diff --git a/.gitattributes b/.gitattributes index ac89d3eaeee07822d2a2c888b9cfde6d53263fd9..53ab3bf829d864b1f610db196f2af68936523375 100644 --- a/.gitattributes +++ b/.gitattributes @@ -4503,3 +4503,6 @@ external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/mpl-data/fonts/ttf/ external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/mpl-data/fonts/ttf/DejaVuSans.ttf filter=lfs diff=lfs merge=lfs -text external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-Bold.ttf filter=lfs diff=lfs merge=lfs -text external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-BoldOblique.ttf filter=lfs diff=lfs merge=lfs -text +external/alphageometry/.venv-ag/Lib/site-packages/scipy.libs/libscipy_openblas-f07f5a5d207a3a47104dca54d6d0c86a.dll filter=lfs diff=lfs merge=lfs -text +external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/propack_test_data.npz filter=lfs diff=lfs merge=lfs -text +external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/lib/npymath.lib filter=lfs diff=lfs merge=lfs -text diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/CheckTactic.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/CheckTactic.lean new file mode 100644 index 0000000000000000000000000000000000000000..d346ec3508bd6f64634ee5903e08c1a38446b3f1 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/CheckTactic.lean @@ -0,0 +1,24 @@ +/- +Copyright (c) 2024 Lean FRO. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Joe Hendrix +-/ +prelude +import Lean.Meta.Basic + +namespace Lean.Meta.CheckTactic + +def mkCheckGoalType (val type : Expr) : MetaM Expr := do + let lvl ← mkFreshLevelMVar + pure <| mkApp2 (mkConst ``CheckGoalType [lvl]) type val + +def matchCheckGoalType (stx : Syntax) (goalType : Expr) : MetaM (Expr × Expr × Level) := do + let u ← mkFreshLevelMVar + let type ← mkFreshExprMVar (some (.sort u)) + let val ← mkFreshExprMVar (some type) + let extType := mkAppN (.const ``CheckGoalType [u]) #[type, val] + if !(← isDefEq goalType extType) then + throwErrorAt stx "Goal{indentExpr goalType}\nis expected to match {indentExpr extType}" + pure (val, type, u) + +end Lean.Meta.CheckTactic diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MVarRenaming.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MVarRenaming.lean new file mode 100644 index 0000000000000000000000000000000000000000..f177c8dc87b5685449e048136607f58269b1c9e0 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MVarRenaming.lean @@ -0,0 +1,36 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Util.ReplaceExpr + +namespace Lean.Meta + +/-- A mapping from MVarId to MVarId -/ +structure MVarRenaming where + map : MVarIdMap MVarId := {} + +def MVarRenaming.isEmpty (s : MVarRenaming) : Bool := + s.map.isEmpty + +def MVarRenaming.find? (s : MVarRenaming) (mvarId : MVarId) : Option MVarId := + s.map.find? mvarId + +def MVarRenaming.find! (s : MVarRenaming) (mvarId : MVarId) : MVarId := + (s.find? mvarId).get! + +def MVarRenaming.insert (s : MVarRenaming) (mvarId mvarId' : MVarId) : MVarRenaming := + { s with map := s.map.insert mvarId mvarId' } + +def MVarRenaming.apply (s : MVarRenaming) (e : Expr) : Expr := + if !e.hasMVar then e + else if s.map.isEmpty then e + else e.replace fun e => match e with + | Expr.mvar mvarId => match s.map.find? mvarId with + | none => e + | some newMVarId => mkMVar newMVarId + | _ => none + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatchPatternAttr.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatchPatternAttr.lean new file mode 100644 index 0000000000000000000000000000000000000000..d374dab2211905bb78b9b5eff6b032c4a7432222 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatchPatternAttr.lean @@ -0,0 +1,37 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Attributes + +namespace Lean + +/-- +Instructs the pattern matcher to unfold occurrences of this definition. + +By default, only constructors and literals can be used for pattern matching. Using +`@[match_pattern]` allows using other definitions, as long as they eventually reduce to +constructors and literals. + +Example: +``` +@[match_pattern] +def yellowString : String := "yellow" + +def isYellow (color : String) : Bool := + match color with + | yellowString => true + | _ => false +``` +-/ +@[builtin_doc] +builtin_initialize matchPatternAttr : TagAttribute ← + registerTagAttribute `match_pattern "mark that a definition can be used in a pattern (remark: the dependent pattern matching compiler will unfold the definition)" + +@[export lean_has_match_pattern_attribute] +def hasMatchPatternAttribute (env : Environment) (n : Name) : Bool := + matchPatternAttr.hasTag env n + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherApp/Basic.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherApp/Basic.lean new file mode 100644 index 0000000000000000000000000000000000000000..762ec475be45c1e550576b215121c92bcd01554e --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherApp/Basic.lean @@ -0,0 +1,80 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Match.MatcherInfo + +namespace Lean.Meta + +structure MatcherApp where + matcherName : Name + matcherLevels : Array Level + uElimPos? : Option Nat + discrInfos : Array Match.DiscrInfo + params : Array Expr + motive : Expr + discrs : Array Expr + altNumParams : Array Nat + alts : Array Expr + remaining : Array Expr + +/-- +Recognizes if `e` is a matcher application, and destructs it into the `MatcherApp` data structure. + +This can optionally also treat `casesOn` recursor applications as a special case +of matcher applications. +-/ +def matchMatcherApp? [Monad m] [MonadEnv m] [MonadError m] (e : Expr) (alsoCasesOn := false) : + m (Option MatcherApp) := do + if let .const declName declLevels := e.getAppFn then + if let some info ← getMatcherInfo? declName then + let args := e.getAppArgs + if args.size < info.arity then + return none + return some { + matcherName := declName + matcherLevels := declLevels.toArray + uElimPos? := info.uElimPos? + discrInfos := info.discrInfos + params := args.extract 0 info.numParams + motive := args[info.getMotivePos]! + discrs := args[(info.numParams + 1)...(info.numParams + 1 + info.numDiscrs)] + altNumParams := info.altNumParams + alts := args[(info.numParams + 1 + info.numDiscrs)...(info.numParams + 1 + info.numDiscrs + info.numAlts)] + remaining := args[(info.numParams + 1 + info.numDiscrs + info.numAlts)...args.size] + } + + if alsoCasesOn && isCasesOnRecursor (← getEnv) declName then + let indName := declName.getPrefix + let .inductInfo info ← getConstInfo indName | return none + let args := e.getAppArgs + unless args.size >= info.numParams + 1 /- motive -/ + info.numIndices + 1 /- major -/ + info.numCtors do return none + let params := args[*...info.numParams] + let motive := args[info.numParams]! + let discrs := args[(info.numParams + 1)...(info.numParams + 1 + info.numIndices + 1)] + let discrInfos := .replicate (info.numIndices + 1) {} + let alts := args[(info.numParams + 1 + info.numIndices + 1)...(info.numParams + 1 + info.numIndices + 1 + info.numCtors)] + let remaining := args[(info.numParams + 1 + info.numIndices + 1 + info.numCtors)...*] + let uElimPos? := if info.levelParams.length == declLevels.length then none else some 0 + let mut altNumParams := #[] + for ctor in info.ctors do + let .ctorInfo ctorInfo ← getConstInfo ctor | unreachable! + altNumParams := altNumParams.push ctorInfo.numFields + return some { + matcherName := declName + matcherLevels := declLevels.toArray + uElimPos?, discrInfos, params, motive, discrs, alts, remaining, altNumParams + } + + return none + +def MatcherApp.toExpr (matcherApp : MatcherApp) : Expr := + let result := mkAppN (mkConst matcherApp.matcherName matcherApp.matcherLevels.toList) matcherApp.params + let result := mkApp result matcherApp.motive + let result := mkAppN result matcherApp.discrs + let result := mkAppN result matcherApp.alts + mkAppN result matcherApp.remaining + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherApp/Transform.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherApp/Transform.lean new file mode 100644 index 0000000000000000000000000000000000000000..1d7198c0e7a32208045f50a61df0c8547d3cd3fd --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherApp/Transform.lean @@ -0,0 +1,443 @@ +/- +Copyright (c) 2024 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura, Joachim Breitner +-/ + +prelude +import Lean.Meta.Match +import Lean.Meta.InferType +import Lean.Meta.Check +import Lean.Meta.Tactic.Split + +namespace Lean.Meta.MatcherApp + +/-- Auxiliary function for MatcherApp.addArg -/ +private partial def updateAlts (unrefinedArgType : Expr) (typeNew : Expr) (altNumParams : Array Nat) (alts : Array Expr) (refined : Bool) (i : Nat) : MetaM (Array Nat × Array Expr) := do + if h : i < alts.size then + let alt := alts[i] + let numParams := altNumParams[i]! + let typeNew ← whnfD typeNew + match typeNew with + | Expr.forallE _ d b _ => + let (alt, refined) ← forallBoundedTelescope d (some numParams) fun xs d => do + let alt ← try instantiateLambda alt xs catch _ => throwError "unexpected matcher application, insufficient number of parameters in alternative" + forallBoundedTelescope d (some 1) fun x _ => do + let alt ← mkLambdaFVars x alt -- x is the new argument we are adding to the alternative + let refined ← if refined then + pure refined + else + pure <| !(← isDefEq unrefinedArgType (← inferType x[0]!)) + return (← mkLambdaFVars xs alt, refined) + updateAlts unrefinedArgType (b.instantiate1 alt) (altNumParams.set! i (numParams+1)) (alts.set i alt) refined (i+1) + | _ => throwError "unexpected type at MatcherApp.addArg" + else + if refined then + return (altNumParams, alts) + else + throwError "failed to add argument to matcher application, argument type was not refined by `casesOn`" + +/-- +Given +- matcherApp `match_i As (fun xs => motive[xs]) discrs (fun ys_1 => (alt_1 : motive (C_1[ys_1])) ... (fun ys_n => (alt_n : motive (C_n[ys_n]) remaining`, and +- expression `e : B[discrs]`, +Construct the term +`match_i As (fun xs => B[xs] -> motive[xs]) discrs (fun ys_1 (y : B[C_1[ys_1]]) => alt_1) ... (fun ys_n (y : B[C_n[ys_n]]) => alt_n) e remaining`. + +We only abstract discriminants that are fvars. We used to use `kabstract` to abstract all +discriminants from `B[discrs]`, but that changes the type of the arg in ways that make it no +longer compatible with the original recursive function (issue #7322). + +If this is still not great, then we could try to use `kabstract`, but only on the last parameter +of the `arg` (the termination proof obligation). + +This method assumes +- the `matcherApp.motive` is a lambda abstraction where `xs.size == discrs.size` +- each alternative is a lambda abstraction where `ys_i.size == matcherApp.altNumParams[i]` + +This is used in `Lean.Elab.PreDefinition.WF.Fix` when replacing recursive calls with calls to +the argument provided by `fix` to refine type of the local variable used for recursive calls, +which may mention `major`. See there for how to use this function. +-/ +def addArg (matcherApp : MatcherApp) (e : Expr) : MetaM MatcherApp := + lambdaTelescope matcherApp.motive fun motiveArgs motiveBody => do + unless motiveArgs.size == matcherApp.discrs.size do + -- This error can only happen if someone implemented a transformation that rewrites the motive created by `mkMatcher`. + throwError "unexpected matcher application, motive must be lambda expression with #{matcherApp.discrs.size} arguments" + let eType ← inferType e + let eTypeAbst := matcherApp.discrs.size.foldRev (init := eType) fun i _ eTypeAbst => + let discr := matcherApp.discrs[i] + if discr.isFVar then + let motiveArg := motiveArgs[i]! + eTypeAbst.replaceFVar discr motiveArg + else + eTypeAbst + let motiveBody ← mkArrow eTypeAbst motiveBody + let matcherLevels ← match matcherApp.uElimPos? with + | none => pure matcherApp.matcherLevels + | some pos => + let uElim ← getLevel motiveBody + pure <| matcherApp.matcherLevels.set! pos uElim + let motive ← mkLambdaFVars motiveArgs motiveBody + -- Construct `aux` `match_i As (fun xs => B[xs] → motive[xs]) discrs`, and infer its type `auxType`. + -- We use `auxType` to infer the type `B[C_i[ys_i]]` of the new argument in each alternative. + let aux := mkAppN (mkConst matcherApp.matcherName matcherLevels.toList) matcherApp.params + let aux := mkApp aux motive + let aux := mkAppN aux matcherApp.discrs + unless (← isTypeCorrect aux) do + throwError "failed to add argument to matcher application, type error when constructing the new motive" + let auxType ← inferType aux + let (altNumParams, alts) ← updateAlts eType auxType matcherApp.altNumParams matcherApp.alts false 0 + return { matcherApp with + matcherLevels := matcherLevels, + motive := motive, + alts := alts, + altNumParams := altNumParams, + remaining := #[e] ++ matcherApp.remaining + } + +/-- Similar to `MatcherApp.addArg`, but returns `none` on failure. -/ +def addArg? (matcherApp : MatcherApp) (e : Expr) : MetaM (Option MatcherApp) := + try + return some (← matcherApp.addArg e) + catch _ => + return none + + +/-- Given + - matcherApp `match_i As (fun xs => motive[xs]) discrs (fun ys_1 => (alt_1 : motive (C_1[ys_1])) ... (fun ys_n => (alt_n : motive (C_n[ys_n]) remaining`, and + - a expression `B[discrs]` (which may not be a type, e.g. `n : Nat`), + returns the expressions `fun ys_1 ... ys_i => B[C_1[ys_1]] ... B[C_n[ys_n]]`, + + This method assumes + - the `matcherApp.motive` is a lambda abstraction where `xs.size == discrs.size` + - each alternative is a lambda abstraction where `ys_i.size == matcherApp.altNumParams[i]` + + This is similar to `MatcherApp.addArg` when you only have an expression to + refined, and not a type with a value. + + This is used in `Lean.Elab.PreDefinition.WF.GuessFix` when constructing the context of recursive + calls to refine the functions' parameter, which may mention `major`. + See there for how to use this function. +-/ +def refineThrough (matcherApp : MatcherApp) (e : Expr) : MetaM (Array Expr) := + lambdaTelescope matcherApp.motive fun motiveArgs _motiveBody => do + unless motiveArgs.size == matcherApp.discrs.size do + -- This error can only happen if someone implemented a transformation that rewrites the motive created by `mkMatcher`. + throwError "failed to transfer argument through matcher application, motive must be lambda expression with #{matcherApp.discrs.size} arguments" + + let eAbst ← matcherApp.discrs.size.foldRevM (init := e) fun i _ eAbst => do + let motiveArg := motiveArgs[i]! + let discr := matcherApp.discrs[i] + let eTypeAbst ← kabstract eAbst discr + return eTypeAbst.instantiate1 motiveArg + -- Let's create something that’s a `Sort` and mentions `e` + -- (recall that `e` itself possibly isn't a type), + -- by writing `e = e`, so that we can use it as a motive + let eEq ← mkEq eAbst eAbst + + let matcherLevels ← match matcherApp.uElimPos? with + | none => pure matcherApp.matcherLevels + | some pos => + pure <| matcherApp.matcherLevels.set! pos levelZero + let motive ← mkLambdaFVars motiveArgs eEq + let aux := mkAppN (mkConst matcherApp.matcherName matcherLevels.toList) matcherApp.params + let aux := mkApp aux motive + let aux := mkAppN aux matcherApp.discrs + unless (← isTypeCorrect aux) do + throwError "failed to transfer argument through matcher application, type error when constructing the new motive" + let auxType ← inferType aux + forallTelescope auxType fun altAuxs _ => do + let altAuxTys ← altAuxs.mapM (inferType ·) + (Array.zip matcherApp.altNumParams altAuxTys).mapM fun (altNumParams, altAuxTy) => do + forallBoundedTelescope altAuxTy altNumParams fun fvs body => do + unless fvs.size = altNumParams do + throwError "failed to transfer argument through matcher application, alt type must be telescope with #{altNumParams} arguments" + -- extract type from our synthetic equality + let body := body.getArg! 2 + -- and abstract over the parameters of the alternatives, so that we can safely pass the Expr out + mkLambdaFVars fvs body + +/-- A non-failing version of `MatcherApp.refineThrough` -/ +def refineThrough? (matcherApp : MatcherApp) (e : Expr) : + MetaM (Option (Array Expr)) := + try + return some (← matcherApp.refineThrough e) + catch _ => + return none + + +private def withUserNamesImpl {α} (fvars : Array Expr) (names : Array Name) (k : MetaM α) : MetaM α := do + let lctx := (Array.zip fvars names).foldl (init := ← (getLCtx)) fun lctx (fvar, name) => + lctx.setUserName fvar.fvarId! name + withLCtx' lctx k + +/-- +Sets the user name of the FVars in the local context according to the given array of names. + +If they differ in size the shorter size wins. +-/ +def withUserNames {n} [MonadControlT MetaM n] [Monad n] + {α} (fvars : Array Expr) (names : Array Name) (k : n α) : n α := do + mapMetaM (withUserNamesImpl fvars names) k + +/- +`Match.forallAltTelescope` lifted to a monad transformer +(and only passing those arguments that we care about below) +-/ +private def forallAltTelescope' + {n} [Monad n] [MonadControlT MetaM n] + {α} (origAltType : Expr) (numParams numDiscrEqs : Nat) + (k : Array Expr → Array Expr → n α) : n α := do + map2MetaM (fun k => + Match.forallAltVarsTelescope origAltType numParams numDiscrEqs + fun ys args _mask _bodyType => k ys args + ) k + +/-- +Performs a possibly type-changing transformation to a `MatcherApp`. + +* `onParams` is run on each parameter and discriminant +* `onMotive` runs on the body of the motive, and is passed the motive parameters + (one for each `MatcherApp.discrs`) +* `onAlt` runs on each alternative, and is passed the expected type of the alternative, + as inferred from the motive +* `onRemaining` runs on the remaining arguments (and may change their number) + +If `useSplitter` is true, the matcher is replaced with the splitter. +NB: Not all operations on `MatcherApp` can handle one `matcherName` is a splitter. + +If `addEqualities` is true, then equalities connecting the discriminant to the parameters of the +alternative (like in `match h : x with …`) are be added, if not already there. + +This function works even if the type of alternatives do *not* fit the inferred type. This +allows you to post-process the `MatcherApp` with `MatcherApp.inferMatchType`, which will +infer a type, given all the alternatives. +-/ +def transform + {n} [MonadLiftT MetaM n] [MonadControlT MetaM n] [Monad n] [MonadError n] [MonadEnv n] [MonadLog n] + [AddMessageContext n] [MonadOptions n] + (matcherApp : MatcherApp) + (useSplitter := false) + (addEqualities : Bool := false) + (onParams : Expr → n Expr := pure) + (onMotive : Array Expr → Expr → n Expr := fun _ e => pure e) + (onAlt : Nat → Expr → Expr → n Expr := fun _ _ e => pure e) + (onRemaining : Array Expr → n (Array Expr) := pure) : + n MatcherApp := do + + -- We also handle CasesOn applications here, and need to treat them specially in a + -- few places. + -- TODO: Expand MatcherApp with the necessary fields to make this more uniform + -- (in particular, include discrEq and whether there is a splitter) + let isCasesOn := isCasesOnRecursor (← getEnv) matcherApp.matcherName + + let numDiscrEqs ← + if isCasesOn then pure 0 else + match ← getMatcherInfo? matcherApp.matcherName with + | some info => pure info.getNumDiscrEqs + | none => throwError "matcher {matcherApp.matcherName} has no MatchInfo found" + + let params' ← matcherApp.params.mapM onParams + let discrs' ← matcherApp.discrs.mapM onParams + + let (motive', uElim, addHEqualities) ← lambdaTelescope matcherApp.motive fun motiveArgs motiveBody => do + unless motiveArgs.size == matcherApp.discrs.size do + throwError "unexpected matcher application, motive must be lambda expression with #{matcherApp.discrs.size} arguments" + let mut motiveBody' ← onMotive motiveArgs motiveBody + + -- Prepend `(x = e) →` or `(x ≍ e) → ` to the motive when an equality is requested + -- and not already present, and remember whether we added an Eq or a HEq + let mut addHEqualities : Array (Option Bool) := #[] + for arg in motiveArgs, discr in discrs', di in matcherApp.discrInfos do + if addEqualities && di.hName?.isNone then + if ← isProof arg then + addHEqualities := addHEqualities.push none + else + let heq ← mkEqHEq discr arg + motiveBody' ← liftMetaM <| mkArrow heq motiveBody' + addHEqualities := addHEqualities.push heq.isHEq + else + addHEqualities := addHEqualities.push none + + return (← mkLambdaFVars motiveArgs motiveBody', ← getLevel motiveBody', addHEqualities) + + let matcherLevels ← match matcherApp.uElimPos? with + | none => pure matcherApp.matcherLevels + | some pos => pure <| matcherApp.matcherLevels.set! pos uElim + + -- We pass `Eq.refl`s for all the equations we added as extra arguments + -- (and count them along the way) + let mut remaining' := #[] + let mut extraEqualities : Nat := 0 + for discr in discrs'.reverse, b in addHEqualities.reverse do + match b with + | none => pure () + | some is_heq => + remaining' := remaining'.push (← (if is_heq then mkHEqRefl else mkEqRefl) discr) + extraEqualities := extraEqualities + 1 + + if useSplitter && !isCasesOn then + let aux1 := mkAppN (mkConst matcherApp.matcherName matcherLevels.toList) params' + let aux1 := mkApp aux1 motive' + let aux1 := mkAppN aux1 discrs' + unless (← isTypeCorrect aux1) do + prependError m!"failed to transform matcher, type error when constructing new pre-splitter motive:{indentExpr aux1}\nfailed with" do + check aux1 + let origAltTypes ← inferArgumentTypesN matcherApp.alts.size aux1 + + -- We replace the matcher with the splitter + let matchEqns ← Match.getEquationsFor matcherApp.matcherName + let splitter := matchEqns.splitterName + + let aux2 := mkAppN (mkConst splitter matcherLevels.toList) params' + let aux2 := mkApp aux2 motive' + let aux2 := mkAppN aux2 discrs' + unless (← isTypeCorrect aux2) do + prependError m!"failed to transform matcher, type error when constructing splitter motive:{indentExpr aux2}\nfailed with" do + check aux2 + let altTypes ← inferArgumentTypesN matcherApp.alts.size aux2 + + let mut alts' := #[] + for altIdx in [:matcherApp.alts.size], + alt in matcherApp.alts, + numParams in matcherApp.altNumParams, + splitterNumParams in matchEqns.splitterAltNumParams, + origAltType in origAltTypes, + altType in altTypes do + let alt' ← forallAltTelescope' origAltType (numParams - numDiscrEqs) 0 fun ys args => do + let altType ← instantiateForall altType ys + -- The splitter inserts its extra parameters after the first ys.size parameters, before + -- the parameters for the numDiscrEqs + forallBoundedTelescope altType (splitterNumParams - ys.size) fun ys2 altType => do + forallBoundedTelescope altType numDiscrEqs fun ys3 altType => do + forallBoundedTelescope altType extraEqualities fun ys4 altType => do + let alt ← try instantiateLambda alt (args ++ ys3) + catch _ => throwError "unexpected matcher application, insufficient number of parameters in alternative" + let alt' ← onAlt altIdx altType alt + mkLambdaFVars (ys ++ ys2 ++ ys3 ++ ys4) alt' + alts' := alts'.push alt' + + remaining' := remaining' ++ (← onRemaining matcherApp.remaining) + + return { matcherApp with + matcherName := splitter + matcherLevels := matcherLevels + params := params' + motive := motive' + discrs := discrs' + altNumParams := matchEqns.splitterAltNumParams.map (· + extraEqualities) + alts := alts' + remaining := remaining' + } + else + let aux := mkAppN (mkConst matcherApp.matcherName matcherLevels.toList) params' + let aux := mkApp aux motive' + let aux := mkAppN aux discrs' + unless (← isTypeCorrect aux) do + logError m!"failed to transform matcher, type error when constructing new motive:{indentExpr aux}" + check aux + let altTypes ← inferArgumentTypesN matcherApp.alts.size aux + + let mut alts' := #[] + for altIdx in [:matcherApp.alts.size], + alt in matcherApp.alts, + numParams in matcherApp.altNumParams, + altType in altTypes do + let alt' ← forallBoundedTelescope altType numParams fun xs altType => do + forallBoundedTelescope altType extraEqualities fun ys4 altType => do + -- we should try to preserve the variable names in the alternative + let names ← lambdaTelescope alt fun xs _ => xs.mapM (·.fvarId!.getUserName) + withUserNames xs names do + let alt ← instantiateLambda alt xs + let alt' ← onAlt altIdx altType alt + mkLambdaFVars (xs ++ ys4) alt' + alts' := alts'.push alt' + + remaining' := remaining' ++ (← onRemaining matcherApp.remaining) + + return { matcherApp with + matcherLevels := matcherLevels + params := params' + motive := motive' + discrs := discrs' + altNumParams := matcherApp.altNumParams.map (· + extraEqualities) + alts := alts' + remaining := remaining' + } + + + +/-- +Given a `MatcherApp`, replaces the motive with one that is inferred from the actual types of the +alternatives. + +For example, given +``` +(match (motive := Nat → Unit → ?) n with + 0 => 1 + _ => true) () +``` +(for any `?`; the motive’s result type be ignored) will give this type +``` +(match n with + | 0 => Nat + | _ => Bool) +``` + +The given `MatcherApp` must not use a splitter in `matcherName`. +The resulting expression *will* use the splitter corresponding to `matcherName` (this is necessary +for the construction). + +Internally, this needs to reduce the matcher in a given branch; this is done using +`Split.simpMatchTarget`. +-/ +def inferMatchType (matcherApp : MatcherApp) : MetaM MatcherApp := do + -- In matcherApp.motive, replace the (dummy) matcher body with a type + -- derived from the inferred types of the alternatives + let nExtra := matcherApp.remaining.size + matcherApp.transform (useSplitter := true) + (onMotive := fun motiveArgs body => do + let extraParams ← arrowDomainsN nExtra body + let propMotive ← mkLambdaFVars motiveArgs (.sort levelZero) + let propAlts ← matcherApp.alts.mapM fun termAlt => + lambdaTelescope termAlt fun xs termAltBody => do + -- We have alt parameters and parameters corresponding to the extra args + let xs1 := xs[*...(xs.size - nExtra)] + let xs2 := xs[(xs.size - nExtra)...xs.size] + -- logInfo m!"altIH: {xs} => {altIH}" + let altType ← inferType termAltBody + for x in xs2 do + if altType.hasAnyFVar (· == x.fvarId!) then + throwError "Type {altType} of alternative {termAlt} still depends on {x}" + -- logInfo m!"altIH type: {altType}" + mkLambdaFVars xs1 altType + let matcherLevels ← match matcherApp.uElimPos? with + | none => pure matcherApp.matcherLevels + | some pos => pure <| matcherApp.matcherLevels.set! pos levelOne + let typeMatcherApp := { matcherApp with + motive := propMotive + matcherLevels := matcherLevels + discrs := motiveArgs + alts := propAlts + remaining := #[] + } + mkArrowN extraParams typeMatcherApp.toExpr + ) + (onAlt := fun _altIdx expAltType alt => do + let altType ← inferType alt + let eq ← mkEq expAltType altType + let proof ← mkFreshExprSyntheticOpaqueMVar eq + let goal := proof.mvarId! + -- logInfo m!"Goal: {goal}" + let goal ← Split.simpMatchTarget goal + -- logInfo m!"Goal after splitting: {goal}" + try + goal.refl + catch _ => + logInfo m!"Cannot close goal after splitting: {goal}" + goal.admit + mkEqMPR proof alt + ) + +end Lean.Meta.MatcherApp diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherInfo.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherInfo.lean new file mode 100644 index 0000000000000000000000000000000000000000..a59b21553f42e943fc8fa18f4566f0a1f69d75f8 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/MatcherInfo.lean @@ -0,0 +1,134 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Basic + +namespace Lean.Meta +namespace Match + +structure DiscrInfo where + /-- `some h` if the discriminant is annotated with `h:` -/ + hName? : Option Name := none + deriving Inhabited + +/-- +A "matcher" auxiliary declaration has the following structure: +- `numParams` parameters +- motive +- `numDiscrs` discriminators (aka major premises) +- `altNumParams.size` alternatives (aka minor premises) where alternative `i` has `altNumParams[i]` parameters +- `uElimPos?` is `some pos` when the matcher can eliminate in different universe levels, and + `pos` is the position of the universe level parameter that specifies the elimination universe. + It is `none` if the matcher only eliminates into `Prop`. -/ +structure MatcherInfo where + numParams : Nat + numDiscrs : Nat + altNumParams : Array Nat + uElimPos? : Option Nat + /-- + `discrInfos[i] = { hName? := some h }` if the i-th discriminant was annotated with `h :`. + -/ + discrInfos : Array DiscrInfo + +def MatcherInfo.numAlts (info : MatcherInfo) : Nat := + info.altNumParams.size + +def MatcherInfo.arity (info : MatcherInfo) : Nat := + info.numParams + 1 + info.numDiscrs + info.numAlts + +def MatcherInfo.getFirstDiscrPos (info : MatcherInfo) : Nat := + info.numParams + 1 + +def MatcherInfo.getDiscrRange (info : MatcherInfo) : Std.Range := + [info.getFirstDiscrPos : info.getFirstDiscrPos + info.numDiscrs] + +def MatcherInfo.getFirstAltPos (info : MatcherInfo) : Nat := + info.numParams + 1 + info.numDiscrs + +def MatcherInfo.getAltRange (info : MatcherInfo) : Std.Range := + [info.getFirstAltPos : info.getFirstAltPos + info.numAlts] + +def MatcherInfo.getMotivePos (info : MatcherInfo) : Nat := + info.numParams + +def getNumEqsFromDiscrInfos (infos : Array DiscrInfo) : Nat := Id.run do + let mut r := 0 + for info in infos do + if info.hName?.isSome then + r := r + 1 + return r + +def MatcherInfo.getNumDiscrEqs (info : MatcherInfo) : Nat := + getNumEqsFromDiscrInfos info.discrInfos + +namespace Extension + +structure Entry where + name : Name + info : MatcherInfo + +structure State where + map : SMap Name MatcherInfo := {} + +instance : Inhabited State := ⟨{}⟩ + +def State.addEntry (s : State) (e : Entry) : State := { s with map := s.map.insert e.name e.info } +def State.switch (s : State) : State := { s with map := s.map.switch } + +builtin_initialize extension : SimplePersistentEnvExtension Entry State ← + registerSimplePersistentEnvExtension { + addEntryFn := State.addEntry + addImportedFn := fun es => (mkStateFromImportedEntries State.addEntry {} es).switch + asyncMode := .async + } + +def addMatcherInfo (env : Environment) (matcherName : Name) (info : MatcherInfo) : Environment := + let _ : Inhabited Environment := ⟨env⟩ + assert! env.asyncMayContain matcherName + extension.addEntry env { name := matcherName, info := info } + +def getMatcherInfo? (env : Environment) (declName : Name) : Option MatcherInfo := + (extension.findStateAsync env declName).map.find? declName + +end Extension + +def addMatcherInfo [Monad m] [MonadEnv m] (matcherName : Name) (info : MatcherInfo) : m Unit := + modifyEnv fun env => Extension.addMatcherInfo env matcherName info + +end Match + +export Match (MatcherInfo) + +def getMatcherInfoCore? (env : Environment) (declName : Name) : Option MatcherInfo := + Match.Extension.getMatcherInfo? env declName + +def getMatcherInfo? [Monad m] [MonadEnv m] (declName : Name) : m (Option MatcherInfo) := + return getMatcherInfoCore? (← getEnv) declName + +@[export lean_is_matcher] +def isMatcherCore (env : Environment) (declName : Name) : Bool := + getMatcherInfoCore? env declName |>.isSome + +def isMatcher [Monad m] [MonadEnv m] (declName : Name) : m Bool := + return isMatcherCore (← getEnv) declName + +def isMatcherAppCore? (env : Environment) (e : Expr) : Option MatcherInfo := + let fn := e.getAppFn + if fn.isConst then + if let some matcherInfo := getMatcherInfoCore? env fn.constName! then + if e.getAppNumArgs ≥ matcherInfo.arity then some matcherInfo else none + else + none + else + none + +def isMatcherAppCore (env : Environment) (e : Expr) : Bool := + isMatcherAppCore? env e |>.isSome + +def isMatcherApp [Monad m] [MonadEnv m] (e : Expr) : m Bool := + return isMatcherAppCore (← getEnv) e + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/Value.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/Value.lean new file mode 100644 index 0000000000000000000000000000000000000000..9ccd0a390a2c446ad213b30ffcdf92d2b3ba9cfc --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Match/Value.lean @@ -0,0 +1,27 @@ +/- +Copyright (c) 2021 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.LitValues +import Lean.Expr + +namespace Lean.Meta + +/-- Return true is `e` is a term that should be processed by the `match`-compiler using `casesValues` -/ +def isMatchValue (e : Expr) : MetaM Bool := do + let e ← instantiateMVars e + if (← getNatValue? e).isSome then return true + if (← getIntValue? e).isSome then return true + if (← getFinValue? e).isSome then return true + if (← getBitVecValue? e).isSome then return true + if (getStringValue? e).isSome then return true + if (← getCharValue? e).isSome then return true + if (← getUInt8Value? e).isSome then return true + if (← getUInt16Value? e).isSome then return true + if (← getUInt32Value? e).isSome then return true + if (← getUInt64Value? e).isSome then return true + return false + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AC.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AC.lean new file mode 100644 index 0000000000000000000000000000000000000000..15469048c2aee92c9e27baf58cfa0f9a90a2738a --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AC.lean @@ -0,0 +1,7 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Dany Fabian +-/ +prelude +import Lean.Meta.Tactic.AC.Main diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AC/Main.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AC/Main.lean new file mode 100644 index 0000000000000000000000000000000000000000..d496b3f9282abc98b52e6a4f5ea5f978991a6134 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AC/Main.lean @@ -0,0 +1,242 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Dany Fabian +-/ +prelude +import Init.Data.AC +import Lean.Meta.AppBuilder +import Lean.Meta.Tactic.Refl +import Lean.Meta.Tactic.Simp.Main +import Lean.Elab.Tactic.Rewrite + +namespace Lean.Meta.AC +open Lean.Data.AC +open Lean.Elab.Tactic +open Std + +abbrev ACExpr := Lean.Data.AC.Expr + +structure PreContext where + id : Nat + op : Expr + assoc : Expr + comm : Option Expr + idem : Option Expr + deriving Inhabited + +instance : ContextInformation (PreContext × Array Bool) where + isComm ctx := ctx.1.comm.isSome + isIdem ctx := ctx.1.idem.isSome + isNeutral ctx x := ctx.2[x]! + +instance : EvalInformation PreContext ACExpr where + arbitrary _ := Data.AC.Expr.var 0 + evalOp _ := Data.AC.Expr.op + evalVar _ x := Data.AC.Expr.var x + +def getInstance (cls : Name) (exprs : Array Expr) : MetaM (Option Expr) := do + try + let app ← mkAppM cls exprs + trace[Meta.AC] "trying: {indentExpr app}" + let inst ← synthInstance app + trace[Meta.AC] "got instance" + return some inst + catch + | _ => return none + +def preContext (expr : Expr) : MetaM (Option PreContext) := do + if let some assoc := ←getInstance ``Associative #[expr] then + return some + { assoc, + op := expr + id := 0 + comm := ←getInstance ``Commutative #[expr] + idem := ←getInstance ``IdempotentOp #[expr] } + + return none + +inductive PreExpr +| op (lhs rhs : PreExpr) +| var (e : Expr) + +@[match_pattern] def bin (op l r : Expr) := + Expr.app (Expr.app op l) r + +def toACExpr (op l r : Expr) : MetaM (Array Expr × ACExpr) := do + let (preExpr, vars) ← + toPreExpr (mkApp2 op l r) + |>.run ∅ + let vars := vars.toArray.insertionSort Expr.lt + let varMap := vars.foldl (fun xs x => xs.insert x xs.size) (∅ : Std.HashMap Expr Nat) |>.get! + + return (vars, toACExpr varMap preExpr) + where + toPreExpr : Expr → StateT ExprSet MetaM PreExpr + | e@(bin op₂ l r) => do + if ←isDefEq op op₂ then + return PreExpr.op (←toPreExpr l) (←toPreExpr r) + modify fun vars => vars.insert e + return PreExpr.var e + | e => do + modify fun vars => vars.insert e + return PreExpr.var e + + toACExpr (varMap : Expr → Nat) : PreExpr → ACExpr + | PreExpr.op l r => Data.AC.Expr.op (toACExpr varMap l) (toACExpr varMap r) + | PreExpr.var x => Data.AC.Expr.var (varMap x) + +/-- +In order to prevent the kernel trying to reduce the atoms of the expression, we abstract the proof +over them. But `ac_rfl` proofs are not completely abstract in the value of the atoms – it recognizes +neutral elements. So we have to abstract over these proofs as well. +-/ +def abstractAtoms (preContext : PreContext) (atoms : Array Expr) + (k : Array (Expr × Option Expr) → MetaM Expr) : MetaM Expr := do + let α ← inferType atoms[0]! + let u ← getLevel α + let rec go i (acc : Array (Expr × Option Expr)) (vars : Array Expr) (args : Array Expr) := do + if h : i < atoms.size then + withLocalDeclD `x α fun v => do + match (← getInstance ``LawfulIdentity #[preContext.op, atoms[i]]) with + | none => + go (i+1) (acc.push (v, .none)) (vars.push v) (args.push atoms[i]) + | some inst => + withLocalDeclD `inst (mkApp3 (mkConst ``LawfulIdentity [u]) α preContext.op v) fun iv => + go (i+1) (acc.push (v, .some iv)) (vars ++ #[v,iv]) (args ++ #[atoms[i], inst]) + else + let proof ← k acc + let proof ← mkLambdaFVars vars proof + let proof := mkAppN proof args + return proof + go 0 #[] #[] #[] + +def buildNormProof (preContext : PreContext) (l r : Expr) : MetaM (Lean.Expr × Lean.Expr) := do + let (atoms, acExpr) ← toACExpr preContext.op l r + let proof ← abstractAtoms preContext atoms fun varsData => do + let α ← inferType atoms[0]! + let u ← getLevel α + let context ← mkContext α u varsData + let isNeutrals := varsData.map (·.2.isSome) + let vars := varsData.map (·.1) + let acExprNormed := Data.AC.evalList ACExpr preContext $ Data.AC.norm (preContext, isNeutrals) acExpr + let lhs := convert acExpr + let rhs := convert acExprNormed + let proof := mkAppN (mkConst ``Context.eq_of_norm [u]) #[α, context, lhs, rhs, ←mkEqRefl (mkConst ``Bool.true)] + let proofType ← mkEq (convertTarget vars acExpr) (convertTarget vars acExprNormed) + let proof := mkExpectedPropHint proof proofType + return proof + let some (_, _, tgt) := (← inferType proof).eq? | panic! "unexpected proof type" + return (proof, tgt) +where + mkContext (α : Expr) (u : Level) (vars : Array (Expr × Option Expr)) : MetaM Expr := do + let arbitrary := vars[0]!.1 + let plift := mkApp (mkConst ``PLift [.zero]) + let pliftUp := mkApp2 (mkConst ``PLift.up [.zero]) + let noneE tp := mkApp (mkConst ``Option.none [.zero]) (plift tp) + let someE tp v := mkApp2 (mkConst ``Option.some [.zero]) (plift tp) (pliftUp tp v) + let vars ← vars.mapM fun ⟨x, inst?⟩ => + let isNeutral := + let isNeutralClass := mkApp3 (mkConst ``LawfulIdentity [u]) α preContext.op x + match inst? with + | none => noneE isNeutralClass + | some isNeutral => someE isNeutralClass isNeutral + return mkApp4 (mkConst ``Variable.mk [u]) α preContext.op x isNeutral + + let vars := vars.toList + let vars ← mkListLit (mkApp2 (mkConst ``Variable [u]) α preContext.op) vars + + let comm := + let commClass := mkApp2 (mkConst ``Commutative [u]) α preContext.op + match preContext.comm with + | none => noneE commClass + | some comm => someE commClass comm + + let idem := + let idemClass := mkApp2 (mkConst ``IdempotentOp [u]) α preContext.op + match preContext.idem with + | none => noneE idemClass + | some idem => someE idemClass idem + + return mkApp7 (mkConst ``Lean.Data.AC.Context.mk [u]) α preContext.op preContext.assoc comm idem vars arbitrary + + convert : ACExpr → Expr + | .op l r => mkApp2 (mkConst ``Data.AC.Expr.op) (convert l) (convert r) + | .var x => mkApp (mkConst ``Data.AC.Expr.var) $ mkNatLit x + + convertTarget (vars : Array Expr) : ACExpr → Expr + | .op l r => mkApp2 preContext.op (convertTarget vars l) (convertTarget vars r) + | .var x => vars[x]! + +def post (e : Expr) : SimpM Simp.Step := do + let ctx ← Simp.getContext + match e, ctx.parent? with + | bin op₁ l r, some (bin op₂ _ _) => + if ←isDefEq op₁ op₂ then + return Simp.Step.done { expr := e } + match ←preContext op₁ with + | some pc => + let (proof, newTgt) ← buildNormProof pc l r + return Simp.Step.done { expr := newTgt, proof? := proof } + | none => return Simp.Step.done { expr := e } + | bin op l r, _ => + match ←preContext op with + | some pc => + let (proof, newTgt) ← buildNormProof pc l r + return Simp.Step.done { expr := newTgt, proof? := proof } + | none => return Simp.Step.done { expr := e } + | e, _ => return Simp.Step.done { expr := e } + +def rewriteUnnormalized (mvarId : MVarId) : MetaM MVarId := do + let simpCtx ← Simp.mkContext + (simpTheorems := {}) + (congrTheorems := (← getSimpCongrTheorems)) + (config := Simp.neutralConfig) + let tgt ← instantiateMVars (← mvarId.getType) + let (res, _) ← Simp.main tgt simpCtx (methods := { post }) + applySimpResultToTarget mvarId tgt res + +def rewriteUnnormalizedRefl (goal : MVarId) : MetaM Unit := do + (← rewriteUnnormalized goal).refl + +@[builtin_tactic acRfl] def acRflTactic : Lean.Elab.Tactic.Tactic := fun _ => do + let goal ← getMainGoal + goal.withContext <| rewriteUnnormalizedRefl goal + +def acNfHypMeta (goal : MVarId) (fvarId : FVarId) : MetaM (Option MVarId) := do + goal.withContext do + let simpCtx ← Simp.mkContext + (simpTheorems := {}) + (congrTheorems := (← getSimpCongrTheorems)) + (config := Simp.neutralConfig) + let tgt ← instantiateMVars (← fvarId.getType) + let (res, _) ← Simp.main tgt simpCtx (methods := { post }) + return (← applySimpResultToLocalDecl goal fvarId res false).map (·.snd) + +/-- Implementation of the `ac_nf` tactic when operating on the main goal. -/ +def acNfTargetTactic : TacticM Unit := + liftMetaTactic1 fun goal => rewriteUnnormalized goal + +/-- Implementation of the `ac_nf` tactic when operating on a hypothesis. -/ +def acNfHypTactic (fvarId : FVarId) : TacticM Unit := + liftMetaTactic1 fun goal => acNfHypMeta goal fvarId + +@[builtin_tactic acNf0] +def evalNf0 : Tactic := fun stx => do + match stx with + | `(tactic| ac_nf0 $[$loc?]?) => + let loc := if let some loc := loc? then expandLocation loc else Location.targets #[] true + withMainContext do + match loc with + | Location.targets hyps target => + if target then acNfTargetTactic + (← getFVarIds hyps).forM acNfHypTactic + | Location.wildcard => + acNfTargetTactic + (← (← getMainGoal).getNondepPropHyps).forM acNfHypTactic + | _ => Lean.Elab.throwUnsupportedSyntax + +builtin_initialize + registerTraceClass `Meta.AC + +end Lean.Meta.AC diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Acyclic.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Acyclic.lean new file mode 100644 index 0000000000000000000000000000000000000000..2770acec5fa2f2de5640253a1b5de65685a7ab9f --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Acyclic.lean @@ -0,0 +1,60 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.MatchUtil +import Lean.Meta.Tactic.Simp.Main + +namespace Lean.MVarId +open Meta + +private def isTarget (lhs rhs : Expr) : MetaM Bool := do + if !lhs.isFVar || !lhs.occurs rhs then + return false + else + isConstructorApp' rhs + +/-- + Close the given goal if `h` is a proof for an equality such as `as = a :: as`. + Inductive datatypes in Lean are acyclic. +-/ +def acyclic (mvarId : MVarId) (h : Expr) : MetaM Bool := mvarId.withContext do + let type ← whnfD (← inferType h) + trace[Meta.Tactic.acyclic] "type: {type}" + let some (_, lhs, rhs) := type.eq? | return false + if (← isTarget lhs rhs) then + go h lhs rhs + else if (← isTarget rhs lhs) then + go (← mkEqSymm h) rhs lhs + else + return false +where + go (h lhs rhs : Expr) : MetaM Bool := do + try + let sizeOf_lhs ← mkAppM ``sizeOf #[lhs] + let sizeOf_rhs ← mkAppM ``sizeOf #[rhs] + let sizeOfEq ← mkLT sizeOf_lhs sizeOf_rhs + let hlt ← mkFreshExprSyntheticOpaqueMVar sizeOfEq + -- TODO: we only need the `sizeOf` simp theorems + let ctx ← Simp.mkContext + (config := { arith := true }) + (simpTheorems := #[ (← getSimpTheorems) ]) + match (← simpTarget hlt.mvarId! ctx {}).1 with + | some _ => return false + | none => + let heq ← mkCongrArg sizeOf_lhs.appFn! (← mkEqSymm h) + let hlt_self ← mkAppM ``Nat.lt_of_lt_of_eq #[hlt, heq] + let hlt_irrelf ← mkAppM ``Nat.lt_irrefl #[sizeOf_lhs] + mvarId.assign (← mkFalseElim (← mvarId.getType) (mkApp hlt_irrelf hlt_self)) + trace[Meta.Tactic.acyclic] "succeeded" + return true + catch ex => + trace[Meta.Tactic.acyclic] "failed with\n{ex.toMessageData}" + return false + +builtin_initialize + registerTraceClass `Meta.Tactic.acyclic + +end Lean.MVarId diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Apply.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Apply.lean new file mode 100644 index 0000000000000000000000000000000000000000..8ced2dac2cacc48e29f046726adf168ecd04e661 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Apply.lean @@ -0,0 +1,377 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura, Siddhartha Gadgil +-/ +prelude +import Lean.Util.FindMVar +import Lean.Meta.SynthInstance +import Lean.Meta.CollectMVars +import Lean.Meta.Tactic.Util +import Lean.PrettyPrinter + +namespace Lean.Meta +/-- + Compute the number of expected arguments and whether the result type is of the form + (?m ...) where ?m is an unassigned metavariable. +-/ +def getExpectedNumArgsAux (e : Expr) : MetaM (Nat × Bool) := + withDefault <| forallTelescopeReducing e fun xs body => + pure (xs.size, body.getAppFn.isMVar) + +def getExpectedNumArgs (e : Expr) : MetaM Nat := do + let (numArgs, _) ← getExpectedNumArgsAux e + pure numArgs + +private def throwApplyError {α} (mvarId : MVarId) + (eType : Expr) (conclusionType? : Option Expr) (targetType : Expr) + (term? : Option MessageData) : MetaM α := do + throwTacticEx `apply mvarId <| MessageData.ofLazyM (es := #[eType, targetType]) do + let conclusionType := conclusionType?.getD eType + let note := if conclusionType?.isSome then .note m!"The full type of {term?.getD "the term"} is{indentExpr eType}" else m!"" + let (conclusionType, targetType) ← addPPExplicitToExposeDiff conclusionType targetType + let conclusion := if conclusionType?.isNone then "type" else "conclusion" + return m!"could not unify the {conclusion} of {term?.getD "the term"}{indentExpr conclusionType}\n\ + with the goal{indentExpr targetType}{note}" + +def synthAppInstances (tacticName : Name) (mvarId : MVarId) (mvarsNew : Array Expr) (binderInfos : Array BinderInfo) + (synthAssignedInstances : Bool) (allowSynthFailures : Bool) : MetaM Unit := do + let mut todo := #[] + -- Collect metavariables to synthesize + for mvar in mvarsNew, binderInfo in binderInfos do + if binderInfo.isInstImplicit then + if synthAssignedInstances || !(← mvar.mvarId!.isAssigned) then + todo := todo.push mvar + while !todo.isEmpty do + todo ← step todo +where + /-- + Try to synthesize instances for the metavariables `mvars`. + Returns metavariables that still need to be synthesized. + We can view the resulting array as the set of metavariables that we should try again. + This is needed when applying or rewriting with functions with complex instances. + For example, consider `rw [@map_smul]` where `map_smul` is + ``` + map_smul {F : Type u_1} {M : Type u_2} {N : Type u_3} {φ : M → N} + {X : Type u_4} {Y : Type u_5} + [SMul M X] [SMul N Y] [FunLike F X Y] [MulActionSemiHomClass F φ X Y] + (f : F) (c : M) (x : X) : DFunLike.coe f (c • x) = φ c • DFunLike.coe f x + ``` + and `MulActionSemiHomClass` is defined as + ``` + class MulActionSemiHomClass (F : Type _) + {M N : outParam (Type _)} (φ : outParam (M → N)) + (X Y : outParam (Type _)) [SMul M X] [SMul N Y] [FunLike F X Y] : Prop where + ``` + The left-hand-side of the equation does not bind `N`. Thus, `SMul N Y` cannot + be synthesized until we synthesize `MulActionSemiHomClass F φ X Y`. Note that + `N` is an output parameter for `MulActionSemiHomClass`. + -/ + step (mvars : Array Expr) : MetaM (Array Expr) := do + -- `ex?` stores the exception for this first synthesis failure in this step. + let mut ex? := none + -- `true` if we managed to synthesize an instance after we hit a failure. + -- That is, there is a chance we may succeed if we try again. + let mut progressAfterEx := false + -- Metavariables that we failed to synthesize. + let mut postponed := #[] + for mvar in mvars do + let mvarType ← inferType mvar + let mvarVal? ← try + let mvarVal ← synthInstance mvarType + unless postponed.isEmpty do + progressAfterEx := true + pure (some mvarVal) + catch ex => + ex? := some ex + postponed := postponed.push mvar + pure none + if let some mvarVal := mvarVal? then + unless (← isDefEq mvar mvarVal) do + -- There is no point in trying again for this kind of failure + unless allowSynthFailures do + throwTacticEx tacticName mvarId "failed to assign synthesized instance" + if let some ex := ex? then + if progressAfterEx then + return postponed + else + -- There is no point in running `step` again. We should give up (`allowSynthFailures`), + -- or throw the first exception we found in this `step`. + if allowSynthFailures then return #[] else throw ex + else + -- Done. We successfully synthesized all metavariables. + return #[] + +def appendParentTag (mvarId : MVarId) (newMVars : Array Expr) (binderInfos : Array BinderInfo) : MetaM Unit := do + let parentTag ← mvarId.getTag + if newMVars.size == 1 then + -- if there is only one subgoal, we inherit the parent tag + newMVars[0]!.mvarId!.setTag parentTag + else + unless parentTag.isAnonymous do + newMVars.size.forM fun i _ => do + let mvarIdNew := newMVars[i].mvarId! + unless (← mvarIdNew.isAssigned) do + unless binderInfos[i]!.isInstImplicit do + let currTag ← mvarIdNew.getTag + mvarIdNew.setTag (appendTag parentTag currTag) + +/-- +If `synthAssignedInstances` is `true`, then `apply` will synthesize instance implicit arguments +even if they have assigned by `isDefEq`, and then check whether the synthesized value matches the +one inferred. The `congr` tactic sets this flag to false. +-/ +def postprocessAppMVars (tacticName : Name) (mvarId : MVarId) (newMVars : Array Expr) (binderInfos : Array BinderInfo) + (synthAssignedInstances := true) (allowSynthFailures := false) : MetaM Unit := do + synthAppInstances tacticName mvarId newMVars binderInfos synthAssignedInstances allowSynthFailures + -- TODO: default and auto params + appendParentTag mvarId newMVars binderInfos + +private def dependsOnOthers (mvar : Expr) (otherMVars : Array Expr) : MetaM Bool := + otherMVars.anyM fun otherMVar => do + if mvar == otherMVar then + return false + else + let otherMVarType ← inferType otherMVar + return (otherMVarType.findMVar? fun mvarId => mvarId == mvar.mvarId!).isSome + +/-- Partitions the given mvars in to two arrays (non-deps, deps) +according to whether the given mvar depends on other mvars in the array.-/ +private def partitionDependentMVars (mvars : Array Expr) : MetaM (Array MVarId × Array MVarId) := + mvars.foldlM (init := (#[], #[])) fun (nonDeps, deps) mvar => do + let currMVarId := mvar.mvarId! + if (← dependsOnOthers mvar mvars) then + return (nonDeps, deps.push currMVarId) + else + return (nonDeps.push currMVarId, deps) + +private def reorderGoals (mvars : Array Expr) : ApplyNewGoals → MetaM (List MVarId) + | ApplyNewGoals.nonDependentFirst => do + let (nonDeps, deps) ← partitionDependentMVars mvars + return nonDeps.toList ++ deps.toList + | ApplyNewGoals.nonDependentOnly => do + let (nonDeps, _) ← partitionDependentMVars mvars + return nonDeps.toList + | ApplyNewGoals.all => return mvars.toList.map Lean.Expr.mvarId! + +/-- Custom `isDefEq` for the `apply` tactic -/ +private def isDefEqApply (approx : Bool) (a b : Expr) : MetaM Bool := do + if approx then + approxDefEq <| isDefEqGuarded a b + else + isDefEqGuarded a b + +/-- +Close the given goal using `apply e`. +-/ +def _root_.Lean.MVarId.apply (mvarId : MVarId) (e : Expr) (cfg : ApplyConfig := {}) + (term? : Option MessageData := none) : MetaM (List MVarId) := + mvarId.withContext do + mvarId.checkNotAssigned `apply + let targetType ← mvarId.getType + let eType ← inferType e + let (numArgs, hasMVarHead) ← getExpectedNumArgsAux eType + /- + The `apply` tactic adds `_`s to `e`, and some of these `_`s become new goals. + When `hasMVarHead` is `false` we try different numbers, until we find a type compatible with `targetType`. + We used to try only `numArgs-targetTypeNumArgs` when `hasMVarHead = false`, but this is not always correct. + For example, consider the following example + ``` + example {α β} [LE_trans β] (x y z : α → β) (h₀ : x ≤ y) (h₁ : y ≤ z) : x ≤ z := by + apply le_trans + assumption + assumption + ``` + In this example, `targetTypeNumArgs = 1` because `LE` for functions is defined as + ``` + instance {α : Type u} {β : Type v} [LE β] : LE (α → β) where + le f g := ∀ i, f i ≤ g i + ``` + -/ + let rangeNumArgs ← if hasMVarHead then + pure [numArgs : numArgs+1] + else + let targetTypeNumArgs ← getExpectedNumArgs targetType + pure [numArgs - targetTypeNumArgs : numArgs+1] + /- + Auxiliary function for trying to add `n` underscores where `n ∈ [i: rangeNumArgs.stop)` + See comment above + -/ + let rec go (i : Nat) : MetaM (Array Expr × Array BinderInfo) := do + if i < rangeNumArgs.stop then + let s ← saveState + let (newMVars, binderInfos, eType) ← forallMetaTelescopeReducing eType i + if (← isDefEqApply cfg.approx eType targetType) then + return (newMVars, binderInfos) + else + s.restore + go (i+1) + else + + let conclusionType? ← if rangeNumArgs.start = 0 then + pure none + else + let (_, _, r) ← forallMetaTelescopeReducing eType (some rangeNumArgs.start) + pure (some r) + throwApplyError mvarId eType conclusionType? targetType term? + termination_by rangeNumArgs.stop - i + let (newMVars, binderInfos) ← go rangeNumArgs.start + postprocessAppMVars `apply mvarId newMVars binderInfos cfg.synthAssignedInstances cfg.allowSynthFailures + let e ← instantiateMVars e + mvarId.assign (mkAppN e newMVars) + let newMVars ← newMVars.filterM fun mvar => not <$> mvar.mvarId!.isAssigned + let otherMVarIds ← getMVarsNoDelayed e + let newMVarIds ← reorderGoals newMVars cfg.newGoals + let otherMVarIds := otherMVarIds.filter fun mvarId => !newMVarIds.contains mvarId + let result := newMVarIds ++ otherMVarIds.toList + result.forM (·.headBetaType) + return result + +/-- Short-hand for applying a constant to the goal. -/ +def _root_.Lean.MVarId.applyConst (mvar : MVarId) (c : Name) (cfg : ApplyConfig := {}) : MetaM (List MVarId) := do + mvar.apply (← mkConstWithFreshMVarLevels c) cfg (term? := m!"'{.ofConstName c}'") + +/-- Close the given goal using `e`, instantiated with `n` metavariables. -/ +def _root_.Lean.MVarId.applyN (mvarId : MVarId) (e : Expr) (n : Nat) (useApproxDefEq := true) : MetaM (List MVarId) := + mvarId.withContext do + mvarId.checkNotAssigned `apply + let targetType ← mvarId.getType + let eType ← inferType e + let (mvarIds, _, eType) ← forallMetaBoundedTelescope eType n + unless mvarIds.size == n do + throwError "Applied type takes fewer than {n} arguments:\n{indentExpr eType}" + unless (← isDefEqApply useApproxDefEq eType targetType) do + throwError "Type mismatch: target is{indentExpr targetType}\nbut applied expression has \ + type{indentExpr eType}\nafter applying {n} arguments." + mvarId.assign (e.beta mvarIds) + return (mvarIds.map (·.mvarId!)).toList + +end Meta + +open Meta +namespace MVarId + +partial def splitAndCore (mvarId : MVarId) : MetaM (List MVarId) := + mvarId.withContext do + mvarId.checkNotAssigned `splitAnd + let type ← mvarId.getType' + if !type.isAppOfArity ``And 2 then + return [mvarId] + else + let tag ← mvarId.getTag + let rec go (type : Expr) : StateRefT (Array MVarId) MetaM Expr := do + let type ← whnf type + if type.isAppOfArity ``And 2 then + let p₁ := type.appFn!.appArg! + let p₂ := type.appArg! + return mkApp4 (mkConst ``And.intro) p₁ p₂ (← go p₁) (← go p₂) + else + let idx := (← get).size + 1 + let mvar ← mkFreshExprSyntheticOpaqueMVar type (tag ++ (`h).appendIndexAfter idx) + modify fun s => s.push mvar.mvarId! + return mvar + let (val, s) ← go type |>.run #[] + mvarId.assign val + return s.toList + +/-- +Apply `And.intro` as much as possible to goal `mvarId`. +-/ +abbrev splitAnd (mvarId : MVarId) : MetaM (List MVarId) := + splitAndCore mvarId + +def exfalso (mvarId : MVarId) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `exfalso + let target ← instantiateMVars (← mvarId.getType) + let u ← getLevel target + let mvarIdNew ← mkFreshExprSyntheticOpaqueMVar (mkConst ``False) (tag := (← mvarId.getTag)) + mvarId.assign (mkApp2 (mkConst ``False.elim [u]) target mvarIdNew) + return mvarIdNew.mvarId! + +/-- +Apply the `n`-th constructor of the target type, +checking that it is an inductive type, +and that there are the expected number of constructors. +-/ +def nthConstructor + (name : Name) (idx : Nat) (expected? : Option Nat := none) (goal : MVarId) : + MetaM (List MVarId) := do + goal.withContext do + goal.checkNotAssigned name + matchConstInduct (← goal.getType').getAppFn + (fun _ => throwTacticEx name goal "target is not an inductive datatype") + fun ival us => do + if let some e := expected? then unless ival.ctors.length == e do + throwTacticEx name goal + s!"{name} tactic works for inductive types with exactly {e} constructors" + if h : idx < ival.ctors.length then + goal.apply <| mkConst ival.ctors[idx] us + else + throwTacticEx name goal s!"index {idx} out of bounds, only {ival.ctors.length} constructors" + +/-- +Try to convert an `Iff` into an `Eq` by applying `iff_of_eq`. +If successful, returns the new goal, and otherwise returns the original `MVarId`. + +This may be regarded as being a special case of `Lean.MVarId.liftReflToEq`, specifically for `Iff`. +-/ +def iffOfEq (mvarId : MVarId) : MetaM MVarId := do + let res ← observing? do + let [mvarId] ← mvarId.apply (mkConst ``iff_of_eq []) | failure + return mvarId + return res.getD mvarId + +/-- +Try to convert an `Eq` into an `Iff` by applying `propext`. +If successful, then returns then new goal, otherwise returns the original `MVarId`. +-/ +def propext (mvarId : MVarId) : MetaM MVarId := do + let res ← observing? do + -- Avoid applying `propext` if the target is not an equality of `Prop`s. + -- We don't want a unification specializing `Sort*` to `Prop`. + let tgt ← withReducible mvarId.getType' + let some (_, x, _) := tgt.eq? | failure + guard <| ← Meta.isProp x + let [mvarId] ← mvarId.apply (mkConst ``propext []) | failure + return mvarId + return res.getD mvarId + +/-- +Try to close the goal using `proof_irrel_heq`. Returns whether or not it succeeds. + +We need to be somewhat careful not to assign metavariables while doing this, otherwise we might +specialize `Sort _` to `Prop`. +-/ +def proofIrrelHeq (mvarId : MVarId) : MetaM Bool := + mvarId.withContext do + let res ← observing? do + mvarId.checkNotAssigned `proofIrrelHeq + let tgt ← withReducible mvarId.getType' + let some (_, lhs, _, rhs) := tgt.heq? | failure + -- Note: `mkAppM` uses `withNewMCtxDepth`, which prevents `Sort _` from specializing to `Prop` + let pf ← mkAppM ``proof_irrel_heq #[lhs, rhs] + mvarId.assign pf + return true + return res.getD false + +/-- +Try to close the goal using `Subsingleton.elim`. Returns whether or not it succeeds. + +We are careful to apply `Subsingleton.elim` in a way that does not assign any metavariables. +This is to prevent the `Subsingleton Prop` instance from being used as justification to specialize +`Sort _` to `Prop`. +-/ +def subsingletonElim (mvarId : MVarId) : MetaM Bool := + mvarId.withContext do + let res ← observing? do + mvarId.checkNotAssigned `subsingletonElim + let tgt ← withReducible mvarId.getType' + let some (_, lhs, rhs) := tgt.eq? | failure + -- Note: `mkAppM` uses `withNewMCtxDepth`, which prevents `Sort _` from specializing to `Prop` + let pf ← mkAppM ``Subsingleton.elim #[lhs, rhs] + mvarId.assign pf + return true + return res.getD false + +end Lean.MVarId diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Assert.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Assert.lean new file mode 100644 index 0000000000000000000000000000000000000000..b22348025fe0c347212270dd0789f3ae779b085f --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Assert.lean @@ -0,0 +1,145 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.FVarSubst +import Lean.Meta.Tactic.Intro +import Lean.Meta.Tactic.Revert +import Lean.Util.ForEachExpr + +namespace Lean.Meta + +/-- + Convert the given goal `Ctx |- target` into `Ctx |- type -> target`. + It assumes `val` has type `type` -/ +def _root_.Lean.MVarId.assert (mvarId : MVarId) (name : Name) (type : Expr) (val : Expr) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `assert + let tag ← mvarId.getTag + let target ← mvarId.getType + let newType := Lean.mkForall name BinderInfo.default type target + let newMVar ← mkFreshExprSyntheticOpaqueMVar newType tag + mvarId.assign (mkApp newMVar val) + return newMVar.mvarId! + +/-- Add the hypothesis `h : t`, given `v : t`, and return the new `FVarId`. -/ +def _root_.Lean.MVarId.note (g : MVarId) (h : Name) (v : Expr) (t? : Option Expr := .none) : + MetaM (FVarId × MVarId) := do + (← g.assert h (← match t? with | some t => pure t | none => inferType v) v).intro1P + +/-- + Convert the given goal `Ctx |- target` into `Ctx |- let name : type := val; target`. + It assumes `val` has type `type` -/ +def _root_.Lean.MVarId.define (mvarId : MVarId) (name : Name) (type : Expr) (val : Expr) : MetaM MVarId := do + mvarId.withContext do + mvarId.checkNotAssigned `define + let tag ← mvarId.getTag + let target ← mvarId.getType + let newType := Lean.mkLet name type val target + let newMVar ← mkFreshExprSyntheticOpaqueMVar newType tag + mvarId.assign newMVar + return newMVar.mvarId! + +/-- + Convert the given goal `Ctx |- target` into `Ctx |- (hName : type) -> hName = val -> target`. + It assumes `val` has type `type` -/ +def _root_.Lean.MVarId.assertExt (mvarId : MVarId) (name : Name) (type : Expr) (val : Expr) (hName : Name := `h) : MetaM MVarId := do + mvarId.withContext do + mvarId.checkNotAssigned `assert + let tag ← mvarId.getTag + let target ← mvarId.getType + let u ← getLevel type + let hType := mkApp3 (mkConst `Eq [u]) type (mkBVar 0) val + let newType := Lean.mkForall name BinderInfo.default type $ Lean.mkForall hName BinderInfo.default hType target + let newMVar ← mkFreshExprSyntheticOpaqueMVar newType tag + let rflPrf ← mkEqRefl val + mvarId.assign (mkApp2 newMVar val rflPrf) + return newMVar.mvarId! + +structure AssertAfterResult where + fvarId : FVarId + mvarId : MVarId + subst : FVarSubst + +/-- + Convert the given goal `Ctx |- target` into a goal containing `(userName : type)` after the local declaration with if `fvarId`. + It assumes `val` has type `type`, and that `type` is well-formed after `fvarId`. + Note that `val` does not need to be well-formed after `fvarId`. That is, it may contain variables that are defined after `fvarId`. -/ +def _root_.Lean.MVarId.assertAfter (mvarId : MVarId) (fvarId : FVarId) (userName : Name) (type : Expr) (val : Expr) : MetaM AssertAfterResult := do + mvarId.checkNotAssigned `assertAfter + let (fvarIds, mvarId) ← mvarId.revertAfter fvarId + let mvarId ← mvarId.assert userName type val + let (fvarIdNew, mvarId) ← mvarId.intro1P + let (fvarIdsNew, mvarId) ← mvarId.introNP fvarIds.size + let mut subst := {} + for f in fvarIds, fNew in fvarIdsNew do + subst := subst.insert f (mkFVar fNew) + return { fvarId := fvarIdNew, mvarId, subst } + +structure Hypothesis where + userName : Name + type : Expr + value : Expr + /-- The hypothesis' `BinderInfo` -/ + binderInfo : BinderInfo := .default + /-- The hypothesis' `LocalDeclKind` -/ + kind : LocalDeclKind := .default + +/-- + Convert the given goal `Ctx |- target` into `Ctx, (hs[0].userName : hs[0].type) ... |-target`. + It assumes `hs[i].val` has type `hs[i].type`. -/ +def _root_.Lean.MVarId.assertHypotheses (mvarId : MVarId) (hs : Array Hypothesis) : MetaM (Array FVarId × MVarId) := do + if hs.isEmpty then + return (#[], mvarId) + else mvarId.withContext do + mvarId.checkNotAssigned `assertHypotheses + let tag ← mvarId.getTag + let target ← mvarId.getType + let targetNew := hs.foldr (init := target) fun h targetNew => + .forallE h.userName h.type targetNew h.binderInfo + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew tag + let val := hs.foldl (init := mvarNew) fun val h => .app val h.value + mvarId.assign val + let (fvarIds, mvarId) ← mvarNew.mvarId!.introNP hs.size + mvarId.modifyLCtx fun lctx => Id.run do + let mut lctx := lctx + for h : i in [:hs.size] do + let h := hs[i] + if h.kind != .default then + lctx := lctx.setKind fvarIds[i]! h.kind + pure lctx + return (fvarIds, mvarId) + +/-- +Replace hypothesis `hyp` in goal `g` with `proof : typeNew`. +The new hypothesis is given the same user name as the original, +it attempts to avoid reordering hypotheses, and the original is cleared if possible. +-/ +-- adapted from Lean.Meta.replaceLocalDeclCore +def _root_.Lean.MVarId.replace (g : MVarId) (hyp : FVarId) (proof : Expr) (typeNew : Option Expr := none) : + MetaM AssertAfterResult := + g.withContext do + let typeNew ← match typeNew with + | some t => pure t + | none => inferType proof + let ldecl ← hyp.getDecl + -- `typeNew` may contain variables that occur after `hyp`. + -- Thus, we use the auxiliary function `findMaxFVar` to ensure `typeNew` is well-formed + -- at the position we are inserting it. + let (_, ldecl') ← findMaxFVar typeNew |>.run ldecl + let result ← g.assertAfter ldecl'.fvarId ldecl.userName typeNew proof + (return { result with mvarId := ← result.mvarId.clear hyp }) <|> pure result +where + /-- Finds the `LocalDecl` for the FVar in `e` with the highest index. -/ + findMaxFVar (e : Expr) : StateRefT LocalDecl MetaM Unit := + e.forEach' fun e => do + if e.isFVar then + let ldecl' ← e.fvarId!.getDecl + modify fun ldecl => if ldecl'.index > ldecl.index then ldecl' else ldecl + return false + else + return e.hasFVar + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Assumption.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Assumption.lean new file mode 100644 index 0000000000000000000000000000000000000000..b7aa848720419bdc8581279b4c5bb31703028bc0 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Assumption.lean @@ -0,0 +1,34 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Util + +namespace Lean.Meta + +/-- Return a local declaration whose type is definitionally equal to `type`. -/ +def findLocalDeclWithType? (type : Expr) : MetaM (Option FVarId) := do + (← getLCtx).findDeclRevM? fun localDecl => do + if localDecl.isImplementationDetail then + return none + else if (← isDefEq type localDecl.type) then + return some localDecl.fvarId + else + return none + +/-- Return `true` if managed to close goal `mvarId` using an assumption. -/ +def _root_.Lean.MVarId.assumptionCore (mvarId : MVarId) : MetaM Bool := + mvarId.withContext do + mvarId.checkNotAssigned `assumption + match (← findLocalDeclWithType? (← mvarId.getType)) with + | none => return false + | some fvarId => mvarId.assign (mkFVar fvarId); return true + +/-- Close goal `mvarId` using an assumption. Throw error message if failed. -/ +def _root_.Lean.MVarId.assumption (mvarId : MVarId) : MetaM Unit := + unless (← mvarId.assumptionCore) do + throwTacticEx `assumption mvarId + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AuxLemma.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AuxLemma.lean new file mode 100644 index 0000000000000000000000000000000000000000..c45493204f9b21018ecfc57ced0357622f0653d1 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/AuxLemma.lean @@ -0,0 +1,62 @@ +/- +Copyright (c) 2021 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.AddDecl +import Lean.Meta.Basic +import Lean.DefEqAttrib + +namespace Lean.Meta + +structure AuxLemmas where + lemmas : PHashMap Expr (Name × List Name) := {} + deriving Inhabited + +builtin_initialize auxLemmasExt : EnvExtension AuxLemmas ← + registerEnvExtension (pure {}) (asyncMode := .local) -- a mere cache, keep local + +/-- + Helper method for creating auxiliary lemmas in the environment. + + It uses a cache that maps `type` to declaration name. The cache is not stored in `.olean` files. + It is useful to make sure the same auxiliary lemma is not created over and over again in the same + environment branch. For expensive auxiliary lemmas that should be deduplicated even across + different environment branches, consider using `realizeConst` instead. + + This method is useful for tactics (e.g., `simp`) that may perform preprocessing steps to lemmas provided by + users. For example, `simp` preprocessor may convert a lemma into multiple ones. +-/ +def mkAuxLemma (levelParams : List Name) (type : Expr) (value : Expr) (kind? : Option Name := none) + (cache := true) (inferRfl := false) : MetaM Name := do + let env ← getEnv + let s := auxLemmasExt.getState env + let mkNewAuxLemma := do + let auxName ← mkAuxDeclName (kind := kind?.getD `_proof) + let decl := + if env.hasUnsafe type || env.hasUnsafe value then + -- `result` contains unsafe code, thus we cannot use a theorem. + Declaration.defnDecl { + name := auxName + hints := ReducibilityHints.opaque + safety := DefinitionSafety.unsafe + levelParams, type, value + } + else + Declaration.thmDecl { + name := auxName + levelParams, type, value + } + addDecl decl + if inferRfl then + inferDefEqAttr auxName + modifyEnv fun env => auxLemmasExt.modifyState env fun ⟨lemmas⟩ => ⟨lemmas.insert type (auxName, levelParams)⟩ + return auxName + if cache then + if let some (name, levelParams') := s.lemmas.find? type then + if levelParams == levelParams' then + return name + mkNewAuxLemma + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Backtrack.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Backtrack.lean new file mode 100644 index 0000000000000000000000000000000000000000..5e811d4fd93b44eb8634ace66acdcfb356f97827 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Backtrack.lean @@ -0,0 +1,205 @@ +/- +Copyright (c) 2023 Kim Morrison. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +prelude +import Init.Data.List.BasicAux +import Lean.Meta.Iterator +import Lean.Meta.Tactic.IndependentOf + +/-! +# `backtrack` + +A meta-tactic for running backtracking search, given a non-deterministic tactic +`alternatives : MVarId → Nondet MetaM (List MVarId)`. + +`backtrack alternatives goals` will recursively try to solve all goals in `goals`, +and the subgoals generated, backtracking as necessary. + +In its default behaviour, it will either solve all goals, or fail. +A customisable `suspend` hook in `BacktrackConfig` allows suspend a goal (or subgoal), +so that it will be returned instead of processed further. +Other hooks `proc` and `discharge` (described in `BacktrackConfig`) allow running other +tactics before `alternatives`, or if all search branches from a given goal fail. + +Currently only `solveByElim` is implemented in terms of `backtrack`. +-/ + +namespace Lean.Meta.Tactic.Backtrack + +open Lean Meta + +/-- +Configuration structure to control the behaviour of `backtrack`: +* control the maximum depth and behaviour (fail or return subgoals) at the maximum depth, +* and hooks allowing + * modifying intermediate goals before running the external tactic, + * 'suspending' goals, returning them in the result, and + * discharging subgoals if the external tactic fails. +-/ +structure BacktrackConfig where + /-- Maximum recursion depth. -/ + maxDepth : Nat := 6 + /-- An arbitrary procedure which can be used to modify the list of goals + before each attempt to generate alternatives. + Called as `proc goals curr`, where `goals` are the original goals for `backtracking`, + and `curr` are the current goals. + Returning `some l` will replace the current goals with `l` and recurse + (consuming one step of maximum depth). + Returning `none` will proceed to generating alternative without changing goals. + Failure will cause backtracking. + (defaults to `none`) -/ + proc : List MVarId → List MVarId → MetaM (Option (List MVarId)) := fun _ _ => pure none + /-- If `suspend g`, then we do not consider alternatives for `g`, + but return `g` as a new subgoal. (defaults to `false`) -/ + suspend : MVarId → MetaM Bool := fun _ => pure false + /-- `discharge g` is called on goals for which there were no alternatives. + If `none` we return `g` as a new subgoal. + If `some l`, we replace `g` by `l` in the list of active goals, and recurse. + If failure, we backtrack. (defaults to failure) -/ + discharge : MVarId → MetaM (Option (List MVarId)) := fun _ => failure + /-- + If we solve any "independent" goals, don't fail. + See `Lean.MVarId.independent?` for the definition of independence. + -/ + commitIndependentGoals : Bool := false + +namespace Backtrack + +/-- +Pretty print a list of goals. +-/ +private def ppMVarId (g : MVarId) : MetaM Format := ppExpr =<< g.getType + +/-- +Pretty print a list of goals. +-/ +private def ppMVarIds (gs : List MVarId) : MetaM (List Format) := gs.mapM ppMVarId + +/-- Run a monadic function on every element of a list, +returning the list of elements on which the function fails, and the list of successful results. -/ +def tryAllM [Monad m] [Alternative m] (L : List α) (f : α → m β) : m (List α × List β) := do + let R ← L.mapM (fun a => (Sum.inr <$> f a) <|> (pure (Sum.inl a))) + return (R.filterMap (fun s => match s with | .inl a => a | _ => none), + R.filterMap (fun s => match s with | .inr b => b | _ => none)) + +variable (cfg : BacktrackConfig) +variable (trace : Name := .anonymous) +variable (next : MVarId → (List MVarId → MetaM (Option (List MVarId))) → MetaM (List MVarId)) + +/-- +* `n : Nat` steps remaining. +* `curr : List MVarId` the current list of unsolved goals. +* `acc : List MVarId` a list of "suspended" goals, which will be returned as subgoals. +-/ + -- `acc` is intentionally a `List` rather than an `Array` so we can share across branches. +private def run (goals : List MVarId) (n : Nat) (curr acc : List MVarId) : MetaM (List MVarId) := do + match n with + | 0 => do + -- We're out of fuel. + throwError "backtrack exceeded the recursion limit" + | n + 1 => do + -- First, run `cfg.proc`, to see if it wants to modify the goals. + let procResult? ← try + cfg.proc goals curr + catch e => + withTraceNode trace + (return m!"{exceptEmoji ·} BacktrackConfig.proc failed: {e.toMessageData}") do + throw e + match procResult? with + | some curr' => run goals n curr' acc + | none => + match curr with + -- If there are no active goals, return the accumulated goals. + | [] => withTraceNode trace (return m!"{exceptEmoji ·} success!") do + return acc.reverse + | g :: gs => + -- Discard any goals which have already been assigned. + if ← g.isAssigned then + withTraceNode trace (return m!"{exceptEmoji ·} discarding already assigned goal {g}") do + run goals (n+1) gs acc + else + withTraceNode trace + -- Note: the `addMessageContextFull` ensures we show the goal using the mvar context before + -- the `do` block below runs, potentially unifying mvars in the goal. + (return m!"{exceptEmoji ·} working on: {← addMessageContextFull g}") + do + -- Check if we should suspend the search here: + if (← cfg.suspend g) then + withTraceNode trace + (fun _ => return m!"⏸️ suspending search and returning as subgoal") do + run goals (n+1) gs (g :: acc) + else + try + -- We attempt to find an alternative, + -- for which all resulting sub-goals can be discharged using `run n`. + next g (fun r => observing? do run goals n (r ++ gs) acc) + catch _ => + -- No lemmas could be applied: + match (← cfg.discharge g) with + | none => (withTraceNode trace + (fun _ => return m!"⏭️ deemed acceptable, returning as subgoal") do + run goals (n+1) gs (g :: acc)) + | some l => (withTraceNode trace + (fun _ => return m!"⏬ discharger generated new subgoals") do + run goals n (l ++ gs) acc) + +/-- +A wrapper around `run`, which works on "independent" goals separately first, +to reduce backtracking. +-/ +private partial def processIndependentGoals (orig : List MVarId) (goals remaining : List MVarId) : + MetaM (List MVarId) := do + -- Partition the remaining goals into "independent" goals + -- (which should be solvable without affecting the solvability of other goals) + -- and all the others. + let (igs, ogs) ← remaining.partitionM (MVarId.isIndependentOf goals) + if igs.isEmpty then + -- If there are no independent goals, we solve all the goals together via backtracking search. + return (← run cfg trace next orig cfg.maxDepth remaining []) + else + withTraceNode trace + (fun _ => return m!"independent goals {← ppMVarIds igs}," + ++ m!" working on them before {← ppMVarIds ogs}") do + -- Invoke `run` on each of the independent goals separately, + -- gathering the subgoals on which `run` fails, + -- and the new subgoals generated from goals on which it is successful. + let (failed, newSubgoals') ← tryAllM igs fun g => + run cfg trace next orig cfg.maxDepth [g] [] + let newSubgoals := newSubgoals'.flatten + withTraceNode trace + (fun _ => return m!"failed: {← ppMVarIds failed}, new: {← ppMVarIds newSubgoals}") do + -- Update the list of goals with respect to which we need to check independence. + let goals' := (← goals.filterM (fun g => do pure !(← g.isAssigned))) ++ newSubgoals + -- If `commitIndependentGoals` is `true`, we will return the new goals + -- regardless of whether we can make further progress on the other goals. + if cfg.commitIndependentGoals && !newSubgoals.isEmpty then + return newSubgoals ++ failed ++ (← (processIndependentGoals orig goals' ogs <|> pure ogs)) + else if !failed.isEmpty then + -- If `commitIndependentGoals` is `false`, and we failed on any of the independent goals, + -- then overall failure is inevitable so we can stop here. + failure + else + -- Finally, having solved this batch of independent goals, + -- recurse (potentially now finding new independent goals). + return newSubgoals ++ (← processIndependentGoals orig goals' ogs) + +end Backtrack + + +/-- +Attempts to solve the `goals`, by recursively calling `next` on each +subgoal that appears with a callback to reenter backtracking search. + +Further flow control options are available via the `Config` argument. + +Returns a list of subgoals which were "suspended" via the `suspend` or +`discharge` hooks in `Config`. In the default configuration, `backtrack` +will either return an empty list or fail. +-/ +def backtrack (cfg : BacktrackConfig := {}) (trace : Name := .anonymous) + (next : MVarId → MetaM (Meta.Iterator (List MVarId))) + (goals : List MVarId) : MetaM (List MVarId) := do + let resolve g f := do (←next g).firstM f + Backtrack.processIndependentGoals cfg trace resolve goals goals goals diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Cases.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Cases.lean new file mode 100644 index 0000000000000000000000000000000000000000..b9b616151c7d8f68c1a0f7fd58bf9110f7a0bebb --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Cases.lean @@ -0,0 +1,361 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.AppBuilder +import Lean.Meta.Tactic.Induction +import Lean.Meta.Tactic.Injection +import Lean.Meta.Tactic.Assert +import Lean.Meta.Tactic.Subst +import Lean.Meta.Tactic.Acyclic +import Lean.Meta.Tactic.UnifyEq + +namespace Lean.Meta + +private def throwInductiveTypeExpected (type : Expr) : MetaM α := do + throwError "failed to compile pattern matching, inductive type expected{indentExpr type}" + +def getInductiveUniverseAndParams (type : Expr) : MetaM (List Level × Array Expr) := do + let type ← whnfD type + matchConstInduct type.getAppFn (fun _ => throwInductiveTypeExpected type) fun val us => + let Iargs := type.getAppArgs + let params := Iargs.extract 0 val.numParams + pure (us, params) + +private def mkEqAndProof (lhs rhs : Expr) : MetaM (Expr × Expr) := do + let lhsType ← inferType lhs + let rhsType ← inferType rhs + let u ← getLevel lhsType + if (← isDefEq lhsType rhsType) then + pure (mkApp3 (mkConst ``Eq [u]) lhsType lhs rhs, mkApp2 (mkConst ``Eq.refl [u]) lhsType lhs) + else + pure (mkApp4 (mkConst ``HEq [u]) lhsType lhs rhsType rhs, mkApp2 (mkConst ``HEq.refl [u]) lhsType lhs) + +partial def withNewEqs (targets targetsNew : Array Expr) (k : Array Expr → Array Expr → MetaM α) : MetaM α := + let rec loop (i : Nat) (newEqs : Array Expr) (newRefls : Array Expr) := do + if i < targets.size then + let (newEqType, newRefl) ← mkEqAndProof targets[i]! targetsNew[i]! + withLocalDeclD `h newEqType fun newEq => do + loop (i+1) (newEqs.push newEq) (newRefls.push newRefl) + else + k newEqs newRefls + loop 0 #[] #[] + +def generalizeTargetsEq (mvarId : MVarId) (motiveType : Expr) (targets : Array Expr) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `generalizeTargets + let (typeNew, eqRefls) ← + forallTelescopeReducing motiveType fun targetsNew _ => do + unless targetsNew.size ≥ targets.size do + throwError "invalid number of targets #{targets.size}, motive only takes #{targetsNew.size}" + let targetsNewAtomic := targetsNew[*...targets.size] + withNewEqs targets targetsNewAtomic fun eqs eqRefls => do + let typeNew ← mvarId.getType + let typeNew ← mkForallFVars eqs typeNew + let typeNew ← mkForallFVars targetsNewAtomic typeNew + pure (typeNew, eqRefls) + let mvarNew ← mkFreshExprSyntheticOpaqueMVar typeNew (← mvarId.getTag) + mvarId.assign (mkAppN (mkAppN mvarNew targets) eqRefls) + pure mvarNew.mvarId! + +structure GeneralizeIndicesSubgoal where + mvarId : MVarId + indicesFVarIds : Array FVarId + fvarId : FVarId + numEqs : Nat + +/-- +Given a metavariable `mvarId` representing the goal +``` +Ctx |- T +``` +and an expression `e : I A j`, where `I A j` is an inductive datatype where `A` are parameters, +and `j` the indices. Generate the goal +``` +Ctx, j' : J, h' : I A j' |- j == j' -> e == h' -> T +``` +Remark: `(j == j' -> e == h')` is a "telescopic" equality. +Remark: `j` is sequence of terms, and `j'` a sequence of free variables. +The result contains the fields +- `mvarId`: the new goal +- `indicesFVarIds`: `j'` ids +- `fvarId`: `h'` id +- `numEqs`: number of equations in the target + +If `varName?` is not none, it is used to name `h'`. +-/ +def generalizeIndices' (mvarId : MVarId) (e : Expr) (varName? : Option Name := none) : MetaM GeneralizeIndicesSubgoal := + mvarId.withContext do + let lctx ← getLCtx + let localInsts ← getLocalInstances + mvarId.checkNotAssigned `generalizeIndices + let type ← whnfD (← inferType e) + type.withApp fun f args => matchConstInduct f (fun _ => throwTacticEx `generalizeIndices mvarId "inductive type expected") fun val _ => do + unless val.numIndices > 0 do throwTacticEx `generalizeIndices mvarId "indexed inductive type expected" + unless args.size == val.numIndices + val.numParams do throwTacticEx `generalizeIndices mvarId "ill-formed inductive datatype" + let indices := args.extract (args.size - val.numIndices) args.size + let IA := mkAppN f (args.extract 0 val.numParams) -- `I A` + let IAType ← inferType IA + forallTelescopeReducing IAType fun newIndices _ => do + let newType := mkAppN IA newIndices + let varName ← if let some varName := varName? then pure varName else mkFreshUserName `x + withLocalDeclD varName newType fun h' => + withNewEqs indices newIndices fun newEqs newRefls => do + let (newEqType, newRefl) ← mkEqAndProof e h' + let newRefls := newRefls.push newRefl + withLocalDeclD `h newEqType fun newEq => do + let newEqs := newEqs.push newEq + /- auxType `forall (j' : J) (h' : I A j'), j == j' -> h == h' -> target -/ + let target ← mvarId.getType + let tag ← mvarId.getTag + let auxType ← mkForallFVars newEqs target + let auxType ← mkForallFVars #[h'] auxType + let auxType ← mkForallFVars newIndices auxType + let newMVar ← mkFreshExprMVarAt lctx localInsts auxType MetavarKind.syntheticOpaque tag + /- assign mvarId := newMVar indices h refls -/ + mvarId.assign (mkAppN (mkApp (mkAppN newMVar indices) e) newRefls) + let (indicesFVarIds, newMVarId) ← newMVar.mvarId!.introNP newIndices.size + let (fvarId, newMVarId) ← newMVarId.intro1P + return { + mvarId := newMVarId, + indicesFVarIds := indicesFVarIds, + fvarId := fvarId, + numEqs := newEqs.size + } + +/-- +Similar to `generalizeTargets` but customized for the `casesOn` motive. +Given a metavariable `mvarId` representing the +``` +Ctx, h : I A j, D |- T +``` +where `fvarId` is `h`s id, and the type `I A j` is an inductive datatype where `A` are parameters, +and `j` the indices. Generate the goal +``` +Ctx, h : I A j, D, j' : J, h' : I A j' |- j == j' -> h == h' -> T +``` +Remark: `(j == j' -> h == h')` is a "telescopic" equality. +Remark: `j` is sequence of terms, and `j'` a sequence of free variables. +The result contains the fields +- `mvarId`: the new goal +- `indicesFVarIds`: `j'` ids +- `fvarId`: `h'` id +- `numEqs`: number of equations in the target -/ +def generalizeIndices (mvarId : MVarId) (fvarId : FVarId) : MetaM GeneralizeIndicesSubgoal := + mvarId.withContext do + let fvarDecl ← fvarId.getDecl + generalizeIndices' mvarId fvarDecl.toExpr fvarDecl.userName + +structure CasesSubgoal extends InductionSubgoal where + ctorName : Name + +namespace Cases + +structure Context where + inductiveVal : InductiveVal + casesOnVal : DefinitionVal + nminors : Nat := inductiveVal.ctors.length + majorDecl : LocalDecl + majorTypeFn : Expr + majorTypeArgs : Array Expr + majorTypeIndices : Array Expr := majorTypeArgs.extract (majorTypeArgs.size - inductiveVal.numIndices) majorTypeArgs.size + +private def mkCasesContext? (majorFVarId : FVarId) : MetaM (Option Context) := do + let env ← getEnv + if !env.contains `Eq || !env.contains `HEq then + pure none + else + let majorDecl ← majorFVarId.getDecl + let majorType ← whnf majorDecl.type + majorType.withApp fun f args => matchConstInduct f (fun _ => pure none) fun ival _ => + if args.size != ival.numIndices + ival.numParams then pure none + else match env.find? (Name.mkStr ival.name "casesOn") with + | ConstantInfo.defnInfo cval => + return some { + inductiveVal := ival, + casesOnVal := cval, + majorDecl := majorDecl, + majorTypeFn := f, + majorTypeArgs := args + } + | _ => pure none + +/-- +We say the major premise has independent indices IF +1- its type is *not* an indexed inductive family, OR +2- its type is an indexed inductive family, but all indices are distinct free variables, and + all local declarations different from the major and its indices do not depend on the indices. +-/ +private def hasIndepIndices (ctx : Context) : MetaM Bool := do + if ctx.majorTypeIndices.isEmpty then + return true + else if ctx.majorTypeIndices.any fun idx => !idx.isFVar then + /- One of the indices is not a free variable. -/ + return false + else if ctx.majorTypeIndices.size.any fun i _ => i.any fun j _ => ctx.majorTypeIndices[i] == ctx.majorTypeIndices[j] then + /- An index occurs more than once -/ + return false + else + (← getLCtx).allM fun decl => + pure (decl.fvarId == ctx.majorDecl.fvarId) <||> -- decl is the major + pure (ctx.majorTypeIndices.any (fun index => decl.fvarId == index.fvarId!)) <||> -- decl is one of the indices + findLocalDeclDependsOn decl (fun fvarId => ctx.majorTypeIndices.all fun idx => idx.fvarId! != fvarId) -- or does not depend on any index + +private def elimAuxIndices (s₁ : GeneralizeIndicesSubgoal) (s₂ : Array CasesSubgoal) : MetaM (Array CasesSubgoal) := + let indicesFVarIds := s₁.indicesFVarIds + s₂.mapM fun s => do + indicesFVarIds.foldlM (init := s) fun s indexFVarId => + match s.subst.get indexFVarId with + | Expr.fvar indexFVarId' => + (do let mvarId ← s.mvarId.clear indexFVarId'; pure { s with mvarId := mvarId, subst := s.subst.erase indexFVarId }) + <|> + (pure s) + | _ => pure s + +/-- + Convert `s` into an array of `CasesSubgoal`, by attaching the corresponding constructor name, + and adding the substitution `majorFVarId -> ctor_i us params fields` into each subgoal. -/ +private def toCasesSubgoals (s : Array InductionSubgoal) (ctorNames : Array Name) (majorFVarId : FVarId) (us : List Level) (params : Array Expr) + : Array CasesSubgoal := + s.mapIdx fun i s => + let ctorName := ctorNames[i]! + let ctorApp := mkAppN (mkAppN (mkConst ctorName us) params) s.fields + let s := { s with subst := s.subst.insert majorFVarId ctorApp } + { ctorName := ctorName, + toInductionSubgoal := s } + +partial def unifyEqs? (numEqs : Nat) (mvarId : MVarId) (subst : FVarSubst) (caseName? : Option Name := none): MetaM (Option (MVarId × FVarSubst)) := withIncRecDepth do + if numEqs == 0 then + return some (mvarId, subst) + else + let (eqFVarId, mvarId) ← mvarId.intro1 + if let some { mvarId, subst, numNewEqs } ← unifyEq? mvarId eqFVarId subst MVarId.acyclic caseName? then + unifyEqs? (numEqs - 1 + numNewEqs) mvarId subst caseName? + else + return none + +private def unifyCasesEqs (numEqs : Nat) (subgoals : Array CasesSubgoal) : MetaM (Array CasesSubgoal) := + subgoals.foldlM (init := #[]) fun subgoals s => do + match (← unifyEqs? numEqs s.mvarId s.subst s.ctorName) with + | none => pure subgoals + | some (mvarId, subst) => + return subgoals.push { s with + mvarId := mvarId, + subst := subst, + fields := s.fields.map (subst.apply ·) + } + +private def inductionCasesOn (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames) (ctx : Context) + (useNatCasesAuxOn : Bool := false) : MetaM (Array CasesSubgoal) := mvarId.withContext do + let majorType ← inferType (mkFVar majorFVarId) + let (us, params) ← getInductiveUniverseAndParams majorType + let mut casesOn := mkCasesOnName ctx.inductiveVal.name + if useNatCasesAuxOn && ctx.inductiveVal.name == ``Nat && (← getEnv).contains ``Nat.casesAuxOn then + casesOn := ``Nat.casesAuxOn + let ctors := ctx.inductiveVal.ctors.toArray + let s ← mvarId.induction majorFVarId casesOn givenNames + return toCasesSubgoals s ctors majorFVarId us params + +def cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames := #[]) (useNatCasesAuxOn : Bool := false) : MetaM (Array CasesSubgoal) := do + try + mvarId.withContext do + mvarId.checkNotAssigned `cases + let context? ← mkCasesContext? majorFVarId + match context? with + | none => throwTacticEx `cases mvarId "not applicable to the given hypothesis" + | some ctx => + /- Remark: if caller does not need a `FVarSubst` (variable substitution), and `hasIndepIndices ctx` is true, + then we can also use the simple case. This is a minor optimization, and we currently do not even + allow callers to specify whether they want the `FVarSubst` or not. -/ + if ctx.inductiveVal.numIndices == 0 then + -- Simple case + inductionCasesOn mvarId majorFVarId givenNames ctx (useNatCasesAuxOn := useNatCasesAuxOn) + else + let s₁ ← generalizeIndices mvarId majorFVarId + trace[Meta.Tactic.cases] "after generalizeIndices\n{MessageData.ofGoal s₁.mvarId}" + let s₂ ← inductionCasesOn s₁.mvarId s₁.fvarId givenNames ctx + let s₂ ← elimAuxIndices s₁ s₂ + unifyCasesEqs s₁.numEqs s₂ + catch ex => + throwNestedTacticEx `cases ex + +end Cases + +/-- +Apply `casesOn` using the free variable `majorFVarId` as the major premise (aka discriminant). +`givenNames` contains user-facing names for each alternative. + +- `useNatCasesAuxOn` is a temporary hack for the `rcases` family of tactics. + Do not use it, as it is subject to removal. + It enables using `Nat.casesAuxOn` instead of `Nat.casesOn`, + which causes case splits on `n : Nat` to be represented as `0` and `n' + 1` rather than as `Nat.zero` and `Nat.succ n'`. +-/ +def _root_.Lean.MVarId.cases (mvarId : MVarId) (majorFVarId : FVarId) (givenNames : Array AltVarNames := #[]) (useNatCasesAuxOn : Bool := false) : MetaM (Array CasesSubgoal) := + Cases.cases mvarId majorFVarId givenNames (useNatCasesAuxOn := useNatCasesAuxOn) + +/-- +Keep applying `cases` on any hypothesis that satisfies `p`. +-/ +def _root_.Lean.MVarId.casesRec (mvarId : MVarId) (p : LocalDecl → MetaM Bool) : MetaM (List MVarId) := + saturate mvarId fun mvarId => + mvarId.withContext do + for localDecl in (← getLCtx) do + if (← p localDecl) then + let r? ← observing? do + let r ← mvarId.cases localDecl.fvarId + return r.toList.map (·.mvarId) + if r?.isSome then + return r? + return none + +/-- +Applies `cases` (recursively) to any hypothesis of the form `h : p ∧ q`. +-/ +def _root_.Lean.MVarId.casesAnd (mvarId : MVarId) : MetaM MVarId := do + let mvarIds ← mvarId.casesRec fun localDecl => return (← instantiateMVars localDecl.type).isAppOfArity ``And 2 + exactlyOne mvarIds + +/-- +Applies `cases` to any hypothesis of the form `h : r = s`. +-/ +def _root_.Lean.MVarId.substEqs (mvarId : MVarId) : MetaM (Option MVarId) := do + let mvarIds ← mvarId.casesRec fun localDecl => do + let type ← instantiateMVars localDecl.type + return type.isEq || type.isHEq + ensureAtMostOne mvarIds + +/-- Auxiliary structure for storing `byCases` tactic result. -/ +structure ByCasesSubgoal where + mvarId : MVarId + fvarId : FVarId + +private def toByCasesSubgoal (s : CasesSubgoal) : MetaM ByCasesSubgoal := do + let #[Expr.fvar fvarId ..] ← pure s.fields | throwError "'byCases' tactic failed, unexpected new hypothesis" + return { mvarId := s.mvarId, fvarId } + +/-- +Split the goal in two subgoals: one containing the hypothesis `h : p` and another containing `h : ¬ p`. +-/ +def _root_.Lean.MVarId.byCases (mvarId : MVarId) (p : Expr) (hName : Name := `h) : MetaM (ByCasesSubgoal × ByCasesSubgoal) := do + let mvarId ← mvarId.assert `hByCases (mkOr p (mkNot p)) (mkEM p) + let (fvarId, mvarId) ← mvarId.intro1 + let #[s₁, s₂] ← mvarId.cases fvarId #[{ varNames := [hName] }, { varNames := [hName] }] | + throwError "'byCases' tactic failed, unexpected number of subgoals" + return ((← toByCasesSubgoal s₁), (← toByCasesSubgoal s₂)) + +/-- +Given `dec : Decidable p`, split the goal in two subgoals: one containing the hypothesis `h : p` and another containing `h : ¬ p`. +-/ +def _root_.Lean.MVarId.byCasesDec (mvarId : MVarId) (p : Expr) (dec : Expr) (hName : Name := `h) : MetaM (ByCasesSubgoal × ByCasesSubgoal) := do + let mvarId ← mvarId.assert `hByCases (mkApp (mkConst ``Decidable) p) dec + let (fvarId, mvarId) ← mvarId.intro1 + let #[s₁, s₂] ← mvarId.cases fvarId #[{ varNames := [hName] }, { varNames := [hName] }] | + throwError "'byCasesDec' tactic failed, unexpected number of subgoals" + -- We flip `s₁` and `s₂` because `isFalse` is the first constructor. + return ((← toByCasesSubgoal s₂), (← toByCasesSubgoal s₁)) + +builtin_initialize registerTraceClass `Meta.Tactic.cases + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Cleanup.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Cleanup.lean new file mode 100644 index 0000000000000000000000000000000000000000..49d5d71eb936aeb6c9e7671ad5d41adec8334232 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Cleanup.lean @@ -0,0 +1,75 @@ +/- +Copyright (c) 2021 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.CollectFVars +import Lean.Meta.Tactic.Clear + +namespace Lean.Meta + +private partial def cleanupCore (mvarId : MVarId) (toPreserve : Array FVarId) (indirectProps : Bool) : MetaM MVarId := do + mvarId.withContext do + mvarId.checkNotAssigned `cleanup + let used ← collectUsed |>.run' (false, {}) + let mut lctx ← getLCtx + for localDecl in lctx do + unless used.contains localDecl.fvarId do + lctx := lctx.erase localDecl.fvarId + let localInsts := (← getLocalInstances).filter fun inst => used.contains inst.fvar.fvarId! + let mvarNew ← mkFreshExprMVarAt lctx localInsts (← instantiateMVars (← mvarId.getType)) .syntheticOpaque (← mvarId.getTag) + mvarId.assign mvarNew + return mvarNew.mvarId! +where + addUsedFVars (e : Expr) : StateRefT (Bool × FVarIdSet) MetaM Unit := do + let (_, s) ← (← instantiateMVars e).collectFVars |>.run {} + for fvarId in s.fvarSet do + addUsedFVar fvarId + + addDeps (fvarId : FVarId) : StateRefT (Bool × FVarIdSet) MetaM Unit := do + let localDecl ← fvarId.getDecl + addUsedFVars localDecl.type + if let some val := localDecl.value? then + addUsedFVars val + + addUsedFVar (fvarId : FVarId) : StateRefT (Bool × FVarIdSet) MetaM Unit := do + unless (← get).2.contains fvarId do + modify fun (_, s) => (true, s.insert fvarId) + addDeps fvarId + + /-- We include `p` in the used-set, if `p` is a proposition that contains a `x` that is in the used-set. -/ + collectPropsStep : StateRefT (Bool × FVarIdSet) MetaM Unit := do + let usedSet := (← get).2 + for localDecl in (← getLCtx) do + if (← isProp localDecl.type) then + if (← dependsOnPred localDecl.type usedSet.contains) then + addUsedFVar localDecl.fvarId + + collectProps : StateRefT (Bool × FVarIdSet) MetaM Unit := do + modify fun s => (false, s.2) + collectPropsStep + if (← get).1 then + collectProps + + collectUsed : StateRefT (Bool × FVarIdSet) MetaM FVarIdSet := do + addUsedFVars (← instantiateMVars (← mvarId.getType)) + toPreserve.forM addUsedFVar + if indirectProps then collectProps + return (← get).2 + +/-- + Auxiliary tactic for cleaning the local context. It removes local declarations (aka hypotheses) that are *not* relevant. + We say a variable `x` is "relevant" if + - It occurs in the `toPreserve` array, or + - It occurs in the target type, or + - There is a relevant variable `y` that depends on `x`, or + - If `indirectProps` is true, the type of `x` is a proposition and it depends on a relevant variable `y`. + + By default, `toPreserve := #[]` and `indirectProps := true`. These settings are used in the mathlib tactic `extract_goal` + to give the user more control over which variables to include. +-/ +abbrev _root_.Lean.MVarId.cleanup (mvarId : MVarId) (toPreserve : Array FVarId := #[]) (indirectProps : Bool := true) : MetaM MVarId := do + cleanupCore mvarId toPreserve indirectProps + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Clear.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Clear.lean new file mode 100644 index 0000000000000000000000000000000000000000..59c1429acc70ecdb13fadc704160eae7596987ec --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Clear.lean @@ -0,0 +1,73 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Util + +namespace Lean.Meta + +/-- +Erase the given free variable from the goal `mvarId`. +-/ +def _root_.Lean.MVarId.clear (mvarId : MVarId) (fvarId : FVarId) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `clear + let lctx ← getLCtx + unless lctx.contains fvarId do + throwTacticEx `clear mvarId m!"unknown variable '{mkFVar fvarId}'" + let tag ← mvarId.getTag + lctx.forM fun localDecl => do + unless localDecl.fvarId == fvarId do + if (← localDeclDependsOn localDecl fvarId) then + throwTacticEx `clear mvarId m!"variable '{localDecl.toExpr}' depends on '{mkFVar fvarId}'" + let mvarDecl ← mvarId.getDecl + if (← exprDependsOn mvarDecl.type fvarId) then + throwTacticEx `clear mvarId m!"target depends on '{mkFVar fvarId}'" + let lctx := lctx.erase fvarId + let localInsts ← getLocalInstances + let localInsts := match localInsts.findFinIdx? fun localInst => localInst.fvar.fvarId! == fvarId with + | none => localInsts + | some idx => localInsts.eraseIdx idx + let newMVar ← mkFreshExprMVarAt lctx localInsts mvarDecl.type MetavarKind.syntheticOpaque tag + mvarId.assign newMVar + pure newMVar.mvarId! + + +/-- +Try to erase the given free variable from the goal `mvarId`. It is no-op if the free variable +cannot be erased due to forward dependencies. +-/ +def _root_.Lean.MVarId.tryClear (mvarId : MVarId) (fvarId : FVarId) : MetaM MVarId := + mvarId.clear fvarId <|> pure mvarId + +/-- +Try to clear the given fvars from the local context. + +The fvars must be given in the order they appear in the local context. + +See also `tryClearMany'` which takes care of reordering internally, +and returns the cleared hypotheses along with the new goal. +-/ +def _root_.Lean.MVarId.tryClearMany (mvarId : MVarId) (fvarIds : Array FVarId) : MetaM MVarId := do + fvarIds.foldrM (init := mvarId) fun fvarId mvarId => mvarId.tryClear fvarId + +/-- +Try to clear the given fvars from the local context. Returns the new goal and +the hypotheses that were cleared. + +Does not require the `hyps` to be given in the order in which they +appear in the local context. +-/ +def _root_.Lean.MVarId.tryClearMany' (goal : MVarId) (fvarIds : Array FVarId) : + MetaM (MVarId × Array FVarId) := + goal.withContext do + let fvarIds := (← getLCtx).sortFVarsByContextOrder fvarIds + fvarIds.foldrM (init := (goal, Array.mkEmpty fvarIds.size)) + fun h (goal, cleared) => do + let goal' ← goal.tryClear h + let cleared := if goal == goal' then cleared else cleared.push h + return (goal', cleared) + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Congr.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Congr.lean new file mode 100644 index 0000000000000000000000000000000000000000..bc30a4e5b3f3edf82254d9f4d5c98dd8cd06eec3 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Congr.lean @@ -0,0 +1,118 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.CongrTheorems +import Lean.Meta.Tactic.Assert +import Lean.Meta.Tactic.Apply +import Lean.Meta.Tactic.Clear +import Lean.Meta.Tactic.Refl +import Lean.Meta.Tactic.Assumption + +namespace Lean +open Meta + +/-- +Preprocessor before applying congruence theorem. +Tries to close new goals using `Eq.refl`, `HEq.refl`, and `assumption`. +It also tries to apply `heq_of_eq`. +-/ +def MVarId.congrPre (mvarId : MVarId) : MetaM (Option MVarId) := do + let mvarId ← mvarId.heqOfEq + try mvarId.refl; return none catch _ => pure () + try mvarId.hrefl; return none catch _ => pure () + if (← mvarId.assumptionCore) then return none + return some mvarId + +/-- +Asserts the given congruence theorem as fresh hypothesis, and then applies it. +Return the `fvarId` for the new hypothesis and the new subgoals. +-/ +private def applyCongrThm? (mvarId : MVarId) (congrThm : CongrTheorem) : MetaM (List MVarId) := do + let mvarId ← mvarId.assert (← mkFreshUserName `h_congr_thm) congrThm.type congrThm.proof + let (fvarId, mvarId) ← mvarId.intro1P + let mvarIds ← mvarId.apply (mkFVar fvarId) { synthAssignedInstances := false } + mvarIds.mapM fun mvarId => mvarId.tryClear fvarId + +/-- +Try to apply a `simp` congruence theorem. +-/ +def MVarId.congr? (mvarId : MVarId) : MetaM (Option (List MVarId)) := + mvarId.withContext do commitWhenSomeNoEx? do + mvarId.checkNotAssigned `congr + let target ← mvarId.getType' + let some (_, lhs, _) := target.eq? | return none + let lhs := lhs.cleanupAnnotations + unless lhs.isApp do return none + let some congrThm ← mkCongrSimp? lhs.getAppFn (subsingletonInstImplicitRhs := false) | return none + applyCongrThm? mvarId congrThm + +/-- +Try to apply a `hcongr` congruence theorem, and then tries to close resulting goals +using `Eq.refl`, `HEq.refl`, and assumption. +-/ +def MVarId.hcongr? (mvarId : MVarId) : MetaM (Option (List MVarId)) := do + commitWhenSomeNoEx? do + mvarId.checkNotAssigned `congr + let mvarId ← mvarId.eqOfHEq + mvarId.withContext do + let target ← mvarId.getType' + let some (_, lhs, _, _) := target.heq? | return none + let lhs := lhs.cleanupAnnotations + unless lhs.isApp do return none + let congrThm ← mkHCongr lhs.getAppFn + applyCongrThm? mvarId congrThm + +/-- +Try to apply `implies_congr`. +-/ +def MVarId.congrImplies? (mvarId : MVarId) : MetaM (Option (List MVarId)) := + observing? do + let mvarId₁ :: mvarId₂ :: _ ← mvarId.apply (← mkConstWithFreshMVarLevels ``implies_congr) | throwError "unexpected number of goals" + return [mvarId₁, mvarId₂] + +/-- +Given a goal of the form `⊢ f as = f bs`, `⊢ (p → q) = (p' → q')`, or `⊢ f as ≍ f bs`, try to apply congruence. +It takes proof irrelevance into account, and the fact that `Decidable p` is a subsingleton. +-/ +def MVarId.congrCore (mvarId : MVarId) : MetaM (List MVarId) := do + if let some mvarIds ← mvarId.congr? then + pure mvarIds + else if let some mvarIds ← mvarId.hcongr? then + pure mvarIds + else if let some mvarIds ← mvarId.congrImplies? then + pure mvarIds + else + throwTacticEx `congr mvarId "failed to apply congruence" + +/-- +Given a goal of the form `⊢ f as = f bs`, `⊢ (p → q) = (p' → q')`, or `⊢ f as ≍ f bs`, try to apply congruence. +It takes proof irrelevance into account, and the fact that `Decidable p` is a subsingleton. + +* Applies `congr` recursively up to depth `depth`. +* If `closePre := true`, it will attempt to close new goals + using `Eq.refl`, `HEq.refl`, and `assumption` with reducible transparency. +* If `closePost := true`, it will try again on goals on which `congr` failed to make progress + with default transparency. +-/ +def MVarId.congrN (mvarId : MVarId) (depth : Nat := 1000000) (closePre := true) (closePost := true) : MetaM (List MVarId) := do + let (_, s) ← go depth mvarId |>.run #[] + return s.toList +where + post (mvarId : MVarId) : StateRefT (Array MVarId) MetaM Unit := do + if closePost && (← getTransparency) != .reducible then + if let some mvarId ← mvarId.congrPre then + modify (·.push mvarId) + else + modify (·.push mvarId) + + go (n : Nat) (mvarId : MVarId) : StateRefT (Array MVarId) MetaM Unit := do + if let some mvarId ← if closePre then withReducible mvarId.congrPre else pure mvarId then + match n with + | 0 => post mvarId + | n+1 => + let some mvarIds ← observing? (m := MetaM) mvarId.congrCore + | post mvarId + mvarIds.forM (go n) diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Constructor.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Constructor.lean new file mode 100644 index 0000000000000000000000000000000000000000..ac8fb3ae969e5c3ce53021f29c4a332cc8e48bfc --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Constructor.lean @@ -0,0 +1,49 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Check +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Apply + +namespace Lean.Meta + +/-- +When the goal `mvarId` type is an inductive datatype, +`constructor` calls `apply` with the first matching constructor. +-/ +def _root_.Lean.MVarId.constructor (mvarId : MVarId) (cfg : ApplyConfig := {}) : MetaM (List MVarId) := do + mvarId.withContext do + mvarId.checkNotAssigned `constructor + let target ← mvarId.getType' + matchConstInduct target.getAppFn + (fun _ => throwTacticEx `constructor mvarId "target is not an inductive datatype") + fun ival us => do + for ctor in ival.ctors do + try + return ← mvarId.apply (Lean.mkConst ctor us) cfg + catch _ => + pure () + throwTacticEx `constructor mvarId "no applicable constructor found" + +def _root_.Lean.MVarId.existsIntro (mvarId : MVarId) (w : Expr) : MetaM MVarId := do + mvarId.withContext do + mvarId.checkNotAssigned `exists + let target ← mvarId.getType' + matchConstStructure target.getAppFn + (fun _ => throwTacticEx `exists mvarId "target is not an inductive datatype with one constructor") + fun _ us cval => do + if cval.numFields < 2 then + throwTacticEx `exists mvarId "constructor must have at least two fields" + let ctor := mkAppN (Lean.mkConst cval.name us) target.getAppArgs[*...cval.numParams] + let ctorType ← inferType ctor + let (mvars, _, _) ← forallMetaTelescopeReducing ctorType (some (cval.numFields-2)) + let f := mkAppN ctor mvars + checkApp f w + let [mvarId] ← mvarId.apply <| mkApp f w + | throwTacticEx `exists mvarId "unexpected number of subgoals" + pure mvarId + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Contradiction.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Contradiction.lean new file mode 100644 index 0000000000000000000000000000000000000000..4bc2cff31ba7430e5067aca617705c59647252e4 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Contradiction.lean @@ -0,0 +1,232 @@ +/- +Copyright (c) 2021 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.MatchUtil +import Lean.Meta.Tactic.Assumption +import Lean.Meta.Tactic.Cases +import Lean.Meta.Tactic.Apply + +namespace Lean.Meta + +structure Contradiction.Config where + useDecide : Bool := true + /-- Check whether any of the hypotheses is an empty type. -/ + emptyType : Bool := true + /-- When checking for empty types, `searchFuel` specifies the number of goals visited. -/ + searchFuel : Nat := 16 + /-- Support for hypotheses such as + ``` + h : (x y : Nat) (ys : List Nat) → x = 0 → y::ys = [a, b, c] → False + ``` + This kind of hypotheses appear when proving conditional equation theorems for match expressions. -/ + genDiseq : Bool := false + +/-- + Try to close the current goal by looking for a proof of `False` nested in a `False.Elim` application in the target. + Return `true` iff the goal has been closed. +-/ +private def nestedFalseElim (mvarId : MVarId) : MetaM Bool := do + let target ← mvarId.getType + if let some falseElim := target.find? fun e => e.isAppOfArity ``False.elim 2 && !e.appArg!.hasLooseBVars then + let falseProof := falseElim.appArg! + mvarId.assign (← mkFalseElim (← mvarId.getType) falseProof) + return true + else + return false + +-- We only consider inductives with no constructors and indexed families +private def isElimEmptyInductiveCandidate (fvarId : FVarId) : MetaM Bool := do + let type ← whnfD (← fvarId.getType) + matchConstInduct type.getAppFn (fun _ => pure false) fun info _ => do + return info.ctors.length == 0 || info.numIndices > 0 + +namespace ElimEmptyInductive + +abbrev M := StateRefT Nat MetaM + +instance : MonadBacktrack SavedState M where + saveState := Meta.saveState + restoreState s := s.restore + +partial def elim (mvarId : MVarId) (fvarId : FVarId) : M Bool := do + if (← get) == 0 then + trace[Meta.Tactic.contradiction] "elimEmptyInductive out-of-fuel" + return false + modify (· - 1) + -- We only consider inductives with no constructors and indexed families + commitWhen do + let subgoals ← try mvarId.cases fvarId catch ex => trace[Meta.Tactic.contradiction] "{ex.toMessageData}"; return false + trace[Meta.Tactic.contradiction] "elimEmptyInductive, number subgoals: {subgoals.size}" + for subgoal in subgoals do + -- If one of the fields is uninhabited, then we are done + let found ← subgoal.mvarId.withContext do + for field in subgoal.fields do + let field := subgoal.subst.apply field + if field.isFVar then + if (← isElimEmptyInductiveCandidate field.fvarId!) then + if (← elim subgoal.mvarId field.fvarId!) then + return true + return false + unless found do + return false + return true + +end ElimEmptyInductive + +private def elimEmptyInductive (mvarId : MVarId) (fvarId : FVarId) (fuel : Nat) : MetaM Bool := do + mvarId.withContext do + if (← isElimEmptyInductiveCandidate fvarId) then + commitWhen do + ElimEmptyInductive.elim (← mvarId.exfalso) fvarId |>.run' fuel + else + return false + +/-- + See `Simp.isEqnThmHypothesis` +-/ +private abbrev isGenDiseq (e : Expr) : Bool := + Simp.isEqnThmHypothesis e + +/-- + Given `e` s.t. `isGenDiseq e`, generate a bit-mask `mask` s.t. `mask[i] = true` iff + the `i`-th binder is an equality without forward dependencies. + + See `processGenDiseq` +-/ +def mkGenDiseqMask (e : Expr) : Array Bool := + go e #[] +where + go (e : Expr) (acc : Array Bool) : Array Bool := + match e with + | .forallE _ d b _ => go b (acc.push (!b.hasLooseBVar 0 && (d.isEq || d.isHEq))) + | _ => acc + +/-- + Close goal if `localDecl` is a "generalized disequality". Example: + ``` + h : (x y : Nat) (ys : List Nat) → x = 0 → y::ys = [a, b, c] → False + ``` + This kind of hypotheses is created when we generate conditional equations for match expressions. +-/ +private def processGenDiseq (mvarId : MVarId) (localDecl : LocalDecl) : MetaM Bool := do + assert! isGenDiseq localDecl.type + let val? ← withNewMCtxDepth do + let (args, _, _) ← forallMetaTelescope localDecl.type + let mask := mkGenDiseqMask localDecl.type + for arg in args, useRefl in mask do + if useRefl then + /- Remark: we should not try to use `refl` for equalities that have forward dependencies because + they correspond to constructor fields. We did not use to have this extra test, and this method failed + to close the following goal. + ``` + ... + ns' : NEList String + h' : NEList.notUno ns' = true + : ∀ (ns : NEList String) (h : NEList.notUno ns = true), Value.lam (Lambda.mk ns' h') = Value.lam (Lambda.mk ns h) → False + ⊢ h_1 l a = h_2 v + + ``` + -/ + if let some (_, lhs, _) ← matchEq? (← inferType arg) then + unless (← isDefEq arg (← mkEqRefl lhs)) do + return none + if let some (_, lhs, _, _) ← matchHEq? (← inferType arg) then + unless (← isDefEq arg (← mkHEqRefl lhs)) do + return none + let falseProof ← instantiateMVars (mkAppN localDecl.toExpr args) + if (← hasAssignableMVar falseProof) then + return none + return some (← mkFalseElim (← mvarId.getType) falseProof) + if let some val := val? then + mvarId.assign val + return true + else + return false + +/-- +Return `true` if goal `mvarId` has contradictory hypotheses. +See `MVarId.contradiction` for the list of tests performed by this method. +-/ +def _root_.Lean.MVarId.contradictionCore (mvarId : MVarId) (config : Contradiction.Config) : MetaM Bool := do + mvarId.withContext do + mvarId.checkNotAssigned `contradiction + if (← nestedFalseElim mvarId) then + return true + for localDecl in (← getLCtx) do + unless localDecl.isImplementationDetail do + -- (h : ¬ p) (h' : p) + if let some p ← matchNot? localDecl.type then + if let some pFVarId ← findLocalDeclWithType? p then + -- We use `False.elim` because `p`'s type may be Type + mvarId.assign (← mkFalseElim (← mvarId.getType) (mkApp localDecl.toExpr (mkFVar pFVarId))) + return true + -- (h : x ≠ x) + if let some (_, lhs, rhs) ← matchNe? localDecl.type then + if (← isDefEq lhs rhs) then + mvarId.assign (← mkAbsurd (← mvarId.getType) (← mkEqRefl lhs) localDecl.toExpr) + return true + let mut isEq := false + -- (h : ctor₁ ... = ctor₂ ...) + if let some (_, lhs, rhs) ← matchEq? localDecl.type then + isEq := true + if let some lhsCtor ← matchConstructorApp? lhs then + if let some rhsCtor ← matchConstructorApp? rhs then + if lhsCtor.name != rhsCtor.name then + mvarId.assign (← mkNoConfusion (← mvarId.getType) localDecl.toExpr) + return true + let mut isHEq := false + -- (h : ctor₁ ... ≍ ctor₂ ...) + if let some (α, lhs, β, rhs) ← matchHEq? localDecl.type then + isHEq := true + if let some lhsCtor ← matchConstructorApp? lhs then + if let some rhsCtor ← matchConstructorApp? rhs then + if lhsCtor.name != rhsCtor.name then + if (← isDefEq α β) then + mvarId.assign (← mkNoConfusion (← mvarId.getType) (← mkEqOfHEq localDecl.toExpr)) + return true + -- (h : p) s.t. `decide p` evaluates to `false` + if config.useDecide && !localDecl.type.hasFVar then + let type ← instantiateMVars localDecl.type + if !type.hasMVar && !type.hasFVar then + try + let d ← mkDecide localDecl.type + let r ← withDefault <| whnf d + if r.isConstOf ``false then + let hn := mkAppN (mkConst ``of_decide_eq_false) <| d.getAppArgs.push (← mkEqRefl d) + mvarId.assign (← mkAbsurd (← mvarId.getType) localDecl.toExpr hn) + return true + catch _ => + pure () + -- "generalized" diseq + if config.genDiseq && isGenDiseq localDecl.type then + if (← processGenDiseq mvarId localDecl) then + return true + -- (h : ) + if config.emptyType && !isEq && !isHEq then + -- We do not use `elimEmptyInductive` for equality, since `cases h` produces a huge proof + -- when `(h : 10000 = 10001)`. TODO: `cases` add a threshold at `cases` + if (← elimEmptyInductive mvarId localDecl.fvarId config.searchFuel) then + return true + return false + +/-- +Try to close the goal using "contradictions" such as +- Contradictory hypotheses `h₁ : p` and `h₂ : ¬ p`. +- Contradictory disequality `h : x ≠ x`. +- Contradictory equality between different constructors, e.g., `h : List.nil = List.cons x xs`. +- Empty inductive types, e.g., `x : Fin 0`. +- Decidable propositions that evaluate to false, i.e., a hypothesis `h : p` s.t. `decide p` reduces to `false`. + This is only tried if `Config.useDecide = true`. + +Throw exception if goal failed to be closed. +-/ +def _root_.Lean.MVarId.contradiction (mvarId : MVarId) (config : Contradiction.Config := {}) : MetaM Unit := + unless (← mvarId.contradictionCore config) do + throwTacticEx `contradiction mvarId + +builtin_initialize registerTraceClass `Meta.Tactic.contradiction + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Delta.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Delta.lean new file mode 100644 index 0000000000000000000000000000000000000000..e27a280eb67fe125c6779745fd8ec3b2e2e0da3c --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Delta.lean @@ -0,0 +1,43 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Transform +import Lean.Meta.Tactic.Replace + +namespace Lean.Meta + +def delta? (e : Expr) (p : Name → Bool := fun _ => true) : CoreM (Option Expr) := + matchConst e.getAppFn (fun _ => return none) fun fInfo fLvls => do + if p fInfo.name && fInfo.hasValue && fInfo.levelParams.length == fLvls.length then + let f ← instantiateValueLevelParams fInfo fLvls + return some (f.betaRev e.getAppRevArgs (useZeta := true)) + else + return none + +/-- Low-level delta expansion. It is used to implement equation lemmas and elimination principles for recursive definitions. -/ +def deltaExpand (e : Expr) (p : Name → Bool) : CoreM Expr := + Core.transform e fun e => do + match (← delta? e p) with + | some e' => return .visit e' + | none => return .continue + +/-- +Delta expand declarations that satisfy `p` at `mvarId` type. +-/ +def _root_.Lean.MVarId.deltaTarget (mvarId : MVarId) (p : Name → Bool) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `delta + mvarId.change (← deltaExpand (← mvarId.getType) p) (checkDefEq := false) + +/-- +Delta expand declarations that satisfy `p` at `fvarId` type. +-/ +def _root_.Lean.MVarId.deltaLocalDecl (mvarId : MVarId) (fvarId : FVarId) (p : Name → Bool) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `delta + mvarId.changeLocalDecl fvarId (← deltaExpand (← fvarId.getType) p) (checkDefEq := false) + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/ElimInfo.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/ElimInfo.lean new file mode 100644 index 0000000000000000000000000000000000000000..01eedf52d9ec9f4b42885aa18b9df71b53378831 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/ElimInfo.lean @@ -0,0 +1,281 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Basic +import Lean.Meta.Check +import Lean.ScopedEnvExtension + +namespace Lean.Meta + +structure ElimAltInfo where + name : Name + /-- A declaration corresponding to the inductive constructor. + (For custom recursors, the alternatives correspond to parameter names in the + recursor, so we may not have a declaration to point to.) + This is used for go-to-definition on the alternative name. -/ + declName? : Option Name + numFields : Nat + /-- If `provesMotive := true`, then this alternative has `motive` as its conclusion. + Only for those alternatives the `induction` tactic should introduce reverted hypotheses. -/ + provesMotive : Bool + deriving Repr, Inhabited + +/-- +Information about an eliminator as used by `induction` or `cases`. + +Created from an expression by `getElimInfo`. This typically contains level metavariables that +are instantiated as we go (e.g. in `addImplicitTargets`), so this is single use. +-/ +structure ElimInfo where + elimExpr : Expr + elimType : Expr + motivePos : Nat + targetsPos : Array Nat := #[] + altsInfo : Array ElimAltInfo := #[] + /-- + Extra arguments to the motive, after the targets, that are instantiated with + non-atomic expressions in the goal + -/ + numComplexMotiveArgs : Nat + deriving Repr, Inhabited + + +/-- Given the type `t` of an alternative, determines the number of parameters +(.forall and .let)-bound, and whether the conclusion is a `motive`-application. -/ +def altArity (motive : Expr) (n : Nat) : Expr → Nat × Bool + | .forallE _ _ b _ => altArity motive (n+1) b + | .letE _ _ _ b _ => altArity motive (n+1) b + | conclusion => (n, conclusion.getAppFn == motive) + + +def getElimExprInfo (elimExpr : Expr) (baseDeclName? : Option Name := none) : MetaM ElimInfo := do + let elimType ← inferType elimExpr + trace[Elab.induction] "eliminator {indentExpr elimExpr}\nhas type{indentExpr elimType}" + forallTelescopeReducing elimType fun xs type => type.withApp fun motive motiveArgs => do + unless motive.isFVar do + throwError "expected resulting type of eliminator to be an application of one of its parameters (the motive):{indentExpr type}" + let targets := motiveArgs.takeWhile (·.isFVar) + let complexMotiveArgs := motiveArgs[targets.size...*] + let motiveType ← inferType motive + forallTelescopeReducing motiveType fun motiveParams motiveResultType => do + unless motiveParams.size == motiveArgs.size do + throwError "expected {motiveArgs.size} parameters at motive type, got {motiveParams.size}:{indentExpr motiveType}" + unless motiveResultType.isSort do + throwError "motive result type must be a sort{indentExpr motiveType}" + let some motivePos ← pure (xs.idxOf? motive) | + throwError "unexpected eliminator type{indentExpr elimType}" + let targetsPos ← targets.mapM fun target => do + match xs.idxOf? target with + | none => throwError "unexpected eliminator type{indentExpr elimType}" + | some targetPos => pure targetPos + let mut altsInfo := #[] + let env ← getEnv + for h : i in [:xs.size] do + let x := xs[i] + if x != motive && !targets.contains x then + let xDecl ← x.fvarId!.getDecl + if xDecl.binderInfo.isExplicit then + let (numFields, provesMotive) := altArity motive 0 xDecl.type + let name := xDecl.userName + let declName? := do + let base ← baseDeclName? + let altDeclName := base ++ name + if env.contains altDeclName then some altDeclName else none + altsInfo := altsInfo.push { name, declName?, numFields, provesMotive } + pure { elimExpr, elimType, motivePos, targetsPos, altsInfo, numComplexMotiveArgs := complexMotiveArgs.size } + +def getElimInfo (elimName : Name) (baseDeclName? : Option Name := none) : MetaM ElimInfo := do + getElimExprInfo (← mkConstWithFreshMVarLevels elimName) baseDeclName? + +/-- + Eliminators/recursors may have implicit targets. For builtin recursors, all indices are implicit targets. + Given an eliminator and the sequence of explicit targets, this methods returns a new sequence containing + implicit and explicit targets. +-/ +partial def addImplicitTargets (elimInfo : ElimInfo) (targets : Array Expr) : MetaM (Array Expr) := do + let (implicitMVars, targets) ← collect elimInfo.elimType 0 0 #[] #[] + for mvar in implicitMVars do + unless ← mvar.isAssigned do + let name := (←mvar.getDecl).userName + if name.isAnonymous || name.hasMacroScopes then + throwError "failed to infer implicit target" + else + throwError "failed to infer implicit target {(←mvar.getDecl).userName}" + targets.mapM instantiateMVars +where + collect (type : Expr) (argIdx targetIdx : Nat) (implicits : Array MVarId) (targets' : Array Expr) : + MetaM (Array MVarId × Array Expr) := do + match (← whnfD type) with + | Expr.forallE n d b bi => + if elimInfo.targetsPos.contains argIdx then + if bi.isExplicit then + unless targetIdx < targets.size do + throwError "insufficient number of targets for '{elimInfo.elimExpr}'" + let target := targets[targetIdx]! + let targetType ← inferType target + unless (← isDefEq d targetType) do + throwError "target{indentExpr target}\n{← mkHasTypeButIsExpectedMsg targetType d}" + collect (b.instantiate1 target) (argIdx+1) (targetIdx+1) implicits (targets'.push target) + else + let implicitTarget ← mkFreshExprMVar (type? := d) (userName := n) + collect (b.instantiate1 implicitTarget) (argIdx+1) targetIdx (implicits.push implicitTarget.mvarId!) (targets'.push implicitTarget) + else + collect (b.instantiate1 (← mkFreshExprMVar d)) (argIdx+1) targetIdx implicits targets' + | _ => + unless targetIdx = targets.size do + throwError "extra targets for '{elimInfo.elimExpr}'" + return (implicits, targets') + +structure CustomEliminator where + induction : Bool + typeNames : Array Name + elimName : Name -- NB: Do not store the ElimInfo, it can contain MVars + deriving Inhabited, Repr + +structure CustomEliminators where + map : SMap (Bool × Array Name) Name := {} + deriving Inhabited, Repr + +def addCustomEliminatorEntry (es : CustomEliminators) (e : CustomEliminator) : CustomEliminators := + match es with + | { map := map } => { map := map.insert (e.induction, e.typeNames) e.elimName } + +builtin_initialize customEliminatorExt : SimpleScopedEnvExtension CustomEliminator CustomEliminators ← + registerSimpleScopedEnvExtension { + initial := {} + addEntry := addCustomEliminatorEntry + finalizeImport := fun { map := map } => { map := map.switch } + } + +def mkCustomEliminator (elimName : Name) (induction : Bool) : MetaM CustomEliminator := do + let elimInfo ← getElimInfo elimName + let info ← getConstInfo elimName + forallTelescopeReducing info.type fun xs _ => do + let mut typeNames := #[] + for hi : i in [:elimInfo.targetsPos.size] do + let targetPos := elimInfo.targetsPos[i] + let x := xs[targetPos]! + /- Return true if there is another target that depends on `x`. -/ + let isImplicitTarget : MetaM Bool := do + for hj : j in [i+1:elimInfo.targetsPos.size] do + let y := xs[elimInfo.targetsPos[j]]! + let yType ← inferType y + if (← dependsOn yType x.fvarId!) then + return true + return false + /- We should only use explicit targets when creating the key for the `CustomEliminators` map. + See test `tests/lean/run/eliminatorImplicitTargets.lean`. -/ + unless (← isImplicitTarget) do + let xType ← inferType x + let .const typeName .. := xType.getAppFn | throwError "unexpected eliminator target type{indentExpr xType}" + typeNames := typeNames.push typeName + return { induction, typeNames, elimName } + +def addCustomEliminator (declName : Name) (attrKind : AttributeKind) (induction : Bool) : MetaM Unit := do + let e ← mkCustomEliminator declName induction + customEliminatorExt.add e attrKind + +/-- +Registers a custom eliminator for the `induction` tactic. + +Whenever the types of the targets in an `induction` call matches a custom eliminator, it is used +instead of the recursor. This can be useful for redefining the default eliminator to a more useful +one. + +Example: +```lean example +structure Three where + val : Fin 3 + +example (x : Three) (p : Three → Prop) : p x := by + induction x + -- val : Fin 3 ⊢ p ⟨val⟩ + +@[induction_eliminator, elab_as_elim] +def Three.myRec {motive : Three → Sort u} + (zero : motive ⟨0⟩) (one : motive ⟨1⟩) (two : motive ⟨2⟩) : + ∀ x, motive x + | ⟨0⟩ => zero | ⟨1⟩ => one | ⟨2⟩ => two + +example (x : Three) (p : Three → Prop) : p x := by + induction x + -- ⊢ p ⟨0⟩ + -- ⊢ p ⟨1⟩ + -- ⊢ p ⟨2⟩ +``` + +`@[cases_eliminator]` works similarly for the `cases` tactic. +-/ +@[builtin_doc] +builtin_initialize + registerBuiltinAttribute { + name := `induction_eliminator + descr := "custom `rec`-like eliminator for the `induction` tactic" + add := fun declName _ attrKind => do + discard <| addCustomEliminator declName attrKind (induction := true) |>.run {} {} + } + +/-- +Registers a custom eliminator for the `cases` tactic. + +Whenever the types of the targets in an `cases` call matches a custom eliminator, it is used +instead of the `casesOn` eliminator. This can be useful for redefining the default eliminator to a +more useful one. + +Example: +```lean example +structure Three where + val : Fin 3 + +example (x : Three) (p : Three → Prop) : p x := by + cases x + -- val : Fin 3 ⊢ p ⟨val⟩ + +@[cases_eliminator, elab_as_elim] +def Three.myRec {motive : Three → Sort u} + (zero : motive ⟨0⟩) (one : motive ⟨1⟩) (two : motive ⟨2⟩) : + ∀ x, motive x + | ⟨0⟩ => zero | ⟨1⟩ => one | ⟨2⟩ => two + +example (x : Three) (p : Three → Prop) : p x := by + cases x + -- ⊢ p ⟨0⟩ + -- ⊢ p ⟨1⟩ + -- ⊢ p ⟨2⟩ +``` + +`@[induction_eliminator]` works similarly for the `induction` tactic. +-/ +@[builtin_doc] +builtin_initialize + registerBuiltinAttribute { + name := `cases_eliminator + descr := "custom `casesOn`-like eliminator for the `cases` tactic" + add := fun declName _ attrKind => do + discard <| addCustomEliminator declName attrKind (induction := false) |>.run {} {} + } + +/-- Gets all the eliminators defined using the `@[induction_eliminator]` and `@[cases_eliminator]` attributes. -/ +def getCustomEliminators : CoreM CustomEliminators := do + return customEliminatorExt.getState (← getEnv) + +/-- +Gets an eliminator appropriate for the provided array of targets. +If `induction` is `true` then returns a matching eliminator defined using the `@[induction_eliminator]` attribute +and otherwise returns one defined using the `@[cases_eliminator]` attribute. + +The `@[induction_eliminator]` attribute is for the `induction` tactic +and the `@[cases_eliminator]` attribute is for the `cases` tactic. +-/ +def getCustomEliminator? (targets : Array Expr) (induction : Bool) : MetaM (Option Name) := do + let mut key := #[] + for target in targets do + let targetType := (← instantiateMVars (← inferType target)).headBeta + let .const declName .. := targetType.getAppFn | return none + key := key.push declName + return customEliminatorExt.getState (← getEnv) |>.map.find? (induction, key) + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/ExposeNames.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/ExposeNames.lean new file mode 100644 index 0000000000000000000000000000000000000000..c565d7d2927de6177d2eb25aec2b6bd04d456914 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/ExposeNames.lean @@ -0,0 +1,66 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Util + +namespace Lean.Meta + +/-- Returns a copy of the local context in which all declarations have clear, accessible names. -/ +private def getLCtxWithExposedNames : MetaM LocalContext := do + let mut map : Std.HashMap Name FVarId := {} + let mut toRename := #[] + let mut lctx ← getLCtx + for localDecl in lctx do + let userName := localDecl.userName + if userName.hasMacroScopes then + toRename := toRename.push localDecl.fvarId + else + if let some fvarId := map[userName]? then + -- Variable has been shadowed + toRename := toRename.push fvarId + map := map.insert userName localDecl.fvarId + if toRename.isEmpty then + return lctx + let mut next : Std.HashMap Name Nat := {} + -- Remark: Shadowed variables may be inserted later. + toRename := toRename.qsort fun fvarId₁ fvarId₂ => + (lctx.get! fvarId₁).index < (lctx.get! fvarId₂).index + for fvarId in toRename do + let localDecl := lctx.get! fvarId + let mut baseName := localDecl.userName + if baseName.hasMacroScopes then + baseName := baseName.eraseMacroScopes + if baseName == `x || baseName == `a then + if (← isProp localDecl.type) then + baseName := `h + let mut userName := baseName + let mut i := next[baseName]?.getD 0 + repeat + if !map.contains userName then + break + i := i + 1 + userName := baseName.appendIndexAfter i + next := next.insert baseName i + map := map.insert userName fvarId + lctx := lctx.modifyLocalDecl fvarId (·.setUserName userName) + return lctx + +/-- +Creates a new goal whose local context has been "exposed" so that every local declaration has a clear, accessible name. +If no local declarations require renaming, the original goal is returned unchanged. +-/ +def _root_.Lean.MVarId.exposeNames (mvarId : MVarId) : MetaM MVarId := mvarId.withContext do + mvarId.checkNotAssigned `expose_names + let lctx ← getLCtxWithExposedNames + let mvarNew ← mkFreshExprMVarAt lctx (← getLocalInstances) (← mvarId.getType) .syntheticOpaque (← mvarId.getTag) + mvarId.assign mvarNew + return mvarNew.mvarId! + +/-- Creates a temporary local context where all names are exposed, and executes `k` -/ +def withExposedNames (k : MetaM α) : MetaM α := do + withNewMCtxDepth <| withLCtx (← getLCtxWithExposedNames) (← getLocalInstances) k + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Ext.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Ext.lean new file mode 100644 index 0000000000000000000000000000000000000000..050b8d68e40371880454b5a9d3c7d5697da9501b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Ext.lean @@ -0,0 +1,85 @@ +/- +Copyright (c) 2021 Gabriel Ebner. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Gabriel Ebner, Mario Carneiro +-/ +prelude +import Init.Data.Array.InsertionSort +import Lean.Meta.DiscrTree + +namespace Lean.Meta.Ext + +/-! +### Environment extension for `ext` theorems +-/ + +/-- Information about an extensionality theorem, stored in the environment extension. -/ +structure ExtTheorem where + /-- Declaration name of the extensionality theorem. -/ + declName : Name + /-- Priority of the extensionality theorem. -/ + priority : Nat + /-- + Key in the discrimination tree, + for the type in which the extensionality theorem holds. + -/ + keys : Array DiscrTree.Key + deriving Inhabited, Repr, BEq, Hashable + +/-- The state of the `ext` extension environment -/ +structure ExtTheorems where + /-- The tree of `ext` extensions. -/ + tree : DiscrTree ExtTheorem := {} + /-- Erased `ext`s via `attribute [-ext]`. -/ + erased : PHashSet Name := {} + deriving Inhabited + +/-- The environment extension to track `@[ext]` theorems. -/ +builtin_initialize extExtension : + SimpleScopedEnvExtension ExtTheorem ExtTheorems ← + registerSimpleScopedEnvExtension { + addEntry := fun { tree, erased } thm => + { tree := tree.insertCore thm.keys thm, erased := erased.erase thm.declName } + initial := {} + } + +/-- Gets the list of `@[ext]` theorems corresponding to the key `ty`, +ordered from high priority to low. -/ +@[inline] def getExtTheorems (ty : Expr) : MetaM (Array ExtTheorem) := do + let extTheorems := extExtension.getState (← getEnv) + let arr ← extTheorems.tree.getMatch ty + let erasedArr := arr.filter fun thm => !extTheorems.erased.contains thm.declName + -- Using insertion sort because it is stable and the list of matches should be mostly sorted. + -- Most ext theorems have default priority. + return erasedArr.insertionSort (·.priority < ·.priority) |>.reverse + +/-- +Erases a name marked `ext` by adding it to the state's `erased` field and +removing it from the state's list of `Entry`s. + +This is triggered by `attribute [-ext] name`. +-/ +def ExtTheorems.eraseCore (d : ExtTheorems) (declName : Name) : ExtTheorems := + { d with erased := d.erased.insert declName } + +/-- Returns `true` if `d` contains theorem with name `declName`. -/ +def ExtTheorems.contains (d : ExtTheorems) (declName : Name) : Bool := + d.tree.containsValueP (·.declName == declName) && !d.erased.contains declName + +/-- Returns `true` if `declName` is tagged with `[ext]` attribute. -/ +def isExtTheorem (declName : Name) : CoreM Bool := do + let extTheorems := extExtension.getState (← getEnv) + return extTheorems.contains declName + +/-- +Erases a name marked as a `ext` attribute. +Check that it does in fact have the `ext` attribute by making sure it names a `ExtTheorem` +found somewhere in the state's tree, and is not erased. +-/ +def ExtTheorems.erase [Monad m] [MonadError m] (d : ExtTheorems) (declName : Name) : + m ExtTheorems := do + unless d.contains declName do + throwError "'{declName}' does not have [ext] attribute" + return d.eraseCore declName + +end Lean.Meta.Ext diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FVarSubst.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FVarSubst.lean new file mode 100644 index 0000000000000000000000000000000000000000..5c72a63de2319c25037ba22da2b4657b36fb125b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FVarSubst.lean @@ -0,0 +1,84 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura, Mario Carneiro +-/ +prelude +import Lean.Data.AssocList +import Lean.Expr +import Lean.LocalContext +import Lean.Util.ReplaceExpr + +namespace Lean.Meta +/-- + Some tactics substitute hypotheses with expressions. + We track these substitutions using `FVarSubst`. + It is just a mapping from the original FVarId (internal) name + to an expression. The free variables occurring in the expression must + be defined in the new goal. -/ +structure FVarSubst where + map : AssocList FVarId Expr := {} + deriving Inhabited + +namespace FVarSubst + +def empty : FVarSubst := {} + +def isEmpty (s : FVarSubst) : Bool := + s.map.isEmpty + +def contains (s : FVarSubst) (fvarId : FVarId) : Bool := + s.map.contains fvarId + +/-- Add entry `fvarId |-> v` to `s` if `s` does not contain an entry for `fvarId`. -/ +def insert (s : FVarSubst) (fvarId : FVarId) (v : Expr) : FVarSubst := + if s.contains fvarId then s + else + let map := s.map.mapVal fun e => e.replaceFVarId fvarId v; + { map := map.insertNew fvarId v } + +def erase (s : FVarSubst) (fvarId : FVarId) : FVarSubst := + { map := s.map.erase fvarId } + +def find? (s : FVarSubst) (fvarId : FVarId) : Option Expr := + s.map.find? fvarId + +def get (s : FVarSubst) (fvarId : FVarId) : Expr := + match s.map.find? fvarId with + | none => mkFVar fvarId -- it has not been replaced + | some v => v + +/-- Given `e`, for each `(x => v)` in `s` replace `x` with `v` in `e` -/ +def apply (s : FVarSubst) (e : Expr) : Expr := + if s.map.isEmpty then e + else if !e.hasFVar then e + else e.replace fun e => match e with + | Expr.fvar fvarId => match s.map.find? fvarId with + | none => e + | some v => v + | _ => none + +def domain (s : FVarSubst) : List FVarId := + s.map.foldl (init := []) fun r k _ => k :: r + +def any (p : FVarId → Expr → Bool) (s : FVarSubst) : Bool := + s.map.any p + +/-- +Constructs a substitution consisting of `s` followed by `t`. +This satisfies `(s.append t).apply e = t.apply (s.apply e)` +-/ +def append (s t : FVarSubst) : FVarSubst := + s.1.foldl (fun s' k v => s'.insert k (t.apply v)) t + +end FVarSubst +end Meta + +def LocalDecl.applyFVarSubst (s : Meta.FVarSubst) : LocalDecl → LocalDecl + | LocalDecl.cdecl i id n t bi k => LocalDecl.cdecl i id n (s.apply t) bi k + | LocalDecl.ldecl i id n t v nd k => LocalDecl.ldecl i id n (s.apply t) (s.apply v) nd k + +abbrev Expr.applyFVarSubst (s : Meta.FVarSubst) (e : Expr) : Expr := + s.apply e + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunInd.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunInd.lean new file mode 100644 index 0000000000000000000000000000000000000000..696fb6898c24a6b4b55efbb4922b708167074c15 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunInd.lean @@ -0,0 +1,1715 @@ +/- +Copyright (c) 2024 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Joachim Breitner +-/ + +prelude +import Lean.Meta.Basic +import Lean.Meta.Match.MatcherApp.Transform +import Lean.Meta.Check +import Lean.Meta.Tactic.Subst +import Lean.Meta.Injective -- for elimOptParam +import Lean.Meta.ArgsPacker +import Lean.Meta.PProdN +import Lean.Elab.PreDefinition.WF.Eqns +import Lean.Elab.PreDefinition.Structural.Eqns +import Lean.Elab.PreDefinition.Structural.IndGroupInfo +import Lean.Elab.PreDefinition.Structural.FindRecArg +import Lean.Elab.Command +import Lean.Meta.Tactic.ElimInfo +import Lean.Meta.Tactic.FunIndInfo + +/-! +This module contains code to derive, from the definition of a recursive function (structural or +well-founded, possibly mutual), a **functional induction principle** tailored to proofs about that +function(s). + +For example from: +``` +def ackermann : Nat → Nat → Nat + | 0, m => m + 1 + | n+1, 0 => ackermann n 1 + | n+1, m+1 => ackermann n (ackermann (n + 1) m) +``` +we get +``` +ackermann.induct (motive : Nat → Nat → Prop) (case1 : ∀ (m : Nat), motive 0 m) + (case2 : ∀ (n : Nat), motive n 1 → motive (Nat.succ n) 0) + (case3 : ∀ (n m : Nat), motive (n + 1) m → motive n (ackermann (n + 1) m) → motive (Nat.succ n) (Nat.succ m)) + (x x : Nat) : motive x x +``` + +## Specification + +The functional induction principle takes the same fixed parameters as the function, and +the motive takes the same non-fixed parameters as the original function. + +For each branch of the original function, there is a case in the induction principle. +Here "branch" roughly corresponds to tail-call positions: branches of top-level +`if`-`then`-`else` and of `match` expressions. + +For every recursive call in that branch, an induction hypothesis asserting the +motive for the arguments of the recursive call is provided. +If the recursive call is under binders and it, or its proof of termination, +depend on the bound values, then these become assumptions of the inductive +hypothesis. + +Additionally, the local context of the branch (e.g. the condition of an +if-then-else; a let-binding, a have-binding) is provided as assumptions in the +corresponding induction case, if they are likely to be useful (as determined +by `MVarId.cleanup`). + +Mutual recursion is supported and results in multiple motives. + + +## Implementation overview (well-founded recursion) + +For a non-mutual, unary function `foo` (or else for the `_unary` function), we + +1. expect its definition to be of the form + ``` + def foo := fun x₁ … xₙ (y : a) => WellFounded.fix (fun y' oldIH => body) y + ``` + where `xᵢ…` are the fixed parameter prefix and `y` is the varying parameter of + the function. + +2. From this structure we derive the type of the motive, and start assembling the induction + principle: + ``` + def foo.induct := fun x₁ … xₙ (motive : (y : a) → Prop) => + fix (fun y' newIH => T[body]) + ``` + +3. The first phase, transformation `T1[body]` (implemented in `buildInductionBody`) + mirrors the branching structure of `foo`, i.e. replicates `dite` and some matcher applications, + while adjusting their motive. It also unfolds calls to `oldIH` and collects induction hypotheses + in conditions (see below). + + In particular, when translating a `match` it is prepared to recognize the idiom + as introduced by `mkFix` via `Lean.Meta.MatcherApp.addArg?`, which refines the type of `oldIH` + throughout the match. The transformation will replace `oldIH` with `newIH` here. + ``` + T[(match (motive := fun oldIH => …) y with | … => fun oldIH' => body) oldIH] + ==> (match (motive := fun newIH => …) y with | … => fun newIH' => T[body]) newIH + ``` + + In addition, the information gathered from the match is preserved, so that when performing the + proof by induction, the user can reliably enter the right case. To achieve this + + * the matcher is replaced by its splitter, which brings extra assumptions into scope when + patterns are overlapping (using `matcherApp.transform (useSplitter := true)`) + * simple discriminants that are mentioned in the goal (i.e plain parameters) are instantiated + in the goal. + * for discriminants that are not instantiated that way, equalities connecting the discriminant + to the instantiation are added (just as if the user wrote `match h : x with …`) + * also, simple discriminants (`FVars`) are remembered as `toClear`, as they are unlikely to + provide useful context, and are redundant given the context that comes from the pattern match. + +4. When a tail position (no more branching) is found, function `buildInductionCase` assembles the + type of the case: a fresh `MVar` asserts the current goal, unwanted values from the local context + are cleared, and the current `body` is searched for recursive calls using `foldAndCollect`, + which are then asserted as inductive hyptheses in the `MVar`. + +5. The function `foldAndCollect` walks the term and performs two operations: + + * collects the induction hypotheses for the current case (with proofs). + * recovering the recursive calls + + So when it encounters a saturated application of `oldIH arg proof`, it + * returns `f arg` and + * remembers the expression `newIH arg proof : motive x` as an inductive hypothesis. + + Since `arg` and `proof` can contain further recursive calls, they are folded there as well. + This assumes that the termination proof `proof` works nevertheless. + + Again, `foldAndCollect` may encounter the `Lean.Meta.Matcherapp.addArg?` idiom, and again it + threads `newIH` through, replacing the extra argument. The resulting type of this induction + hypothesis is now itself a `match` statement (cf. `Lean.Meta.MatcherApp.inferMatchType`) + + The termination proof of `foo` may have abstracted over some proofs; these proofs must be + transferred, so auxiliary lemmas are unfolded if needed. + +7. After this construction, the MVars introduced by `buildInductionCase` are turned into parameters. + +The resulting term then becomes `foo.induct` at its inferred type. + +## Implementation overview (mutual/non-unary well-founded recursion) + +If `foo` is not unary and/or part of a mutual reduction, then the induction theorem for `foo._unary` +(i.e. the unary non-mutual recursion function produced by the equation compiler) +of the form +``` +foo._unary.induct : {motive : (a ⊗' b) ⊕' c → Prop} → + (case1 : ∀ …, motive (PSum.inl (x,y)) → …) → … → + (x : (a ⊗' b) ⊕' c) → motive x +``` +will first in `unpackMutualInduction` be turned into a joint induction theorem of the form +``` +foo.mutual_induct : {motive1 : a → b → Prop} {motive2 : c → Prop} → + (case1 : ∀ …, motive1 x y → …) → … → + ((x : a) → (y : b) → motive1 x y) ∧ ((z : c) → motive2 z) +``` +where all the `PSum`/`PSigma` encoding has been resolved. This theorem is attached to the +name of the first function in the mutual group, like the `._unary` definition. + +Finally, in `deriveUnpackedInduction`, for each of the functions in the mutual group, a simple +projection yields the final `foo.induct` theorem: +``` +foo.induct : {motive1 : a → b → Prop} {motive2 : c → Prop} → + (case1 : ∀ …, motive1 x y → …) → … → + (x : a) → (y : b) → motive1 x y +``` + +## Implementation overview (structural recursion) + +When handling structural recursion, the overall approach is the same, with the following +differences: + +* Instead of `WellFounded.fix` we look for a `.brecOn` application, using `isBRecOnRecursor` + + Despite its name, this function does *not* recognize the `.brecOn` of inductive *predicates*, + which we also do not support at this point. + +* The elaboration of structurally recursive function can handle extra arguments. We keep the + `motive` parameters in the original order. + +## Unfolding principles + +The code can also create a variant of the induction/cases principles that automatically unfolds +the function application. It's motive abstracts over the function call, so for the ackermann +function one gets + +``` +ackermann.fun_cases_unfolding + (motive : Nat → Nat → Nat → Prop) + (case1 : ∀ (m : Nat), motive 0 m (m + 1)) + (case2 : ∀ (n : Nat), motive n.succ 0 (ackermann n 1)) + (case3 : ∀ (n m : Nat), motive n.succ m.succ (ackermann n (ackermann (n + 1) m))) + (x✝ x✝¹ : Nat) : motive x✝ x✝¹ (ackermann x✝ x✝¹) +``` + +To implement this, in the initial goal `motive x (ackermann x)` of `buildInductionBody` we unfold the +function definition, and then reduce is as we go into match, ite or let expressions, using the +`withRewrittenMotive` function. + +This gives us great control over the reduction, for example to move `let` expressions to the context +simultaneously. + +The combinators passed to `withRewrittenMotive` are forgiving, so when `unfolding := false`, or when +something goes wrong, one still gets a useful induction principle, just maybe with the function +not fully simplified. +-/ + +namespace Lean.Tactic.FunInd + +open Lean Elab Meta + +def lambdaTelescope1 {n} [MonadControlT MetaM n] [MonadError n] [MonadNameGenerator n] [Monad n] {α} (e : Expr) (k : FVarId → Expr → n α) : n α := do + lambdaBoundedTelescope e 1 fun xs body => do + unless xs.size == 1 do + throwError "lambdaTelescope1: expected lambda, got {e}" + k xs[0]!.fvarId! body + +/-- There are multiple variants of this function around in the code base, maybe unify at some point. -/ +private def elimTypeAnnotations (type : Expr) : CoreM Expr := do + Core.transform type fun e => + if e.isOptParam || e.isAutoParam then + return .visit e.appFn!.appArg! + else if e.isOutParam || e.isSemiOutParam then + return .visit e.appArg! + else + return .continue + +/-- +A monad to help collecting inductive hypothesis. + +In `foldAndCollect` it's a writer monad (with variants of the `local` combinator), +and in `buildInductionBody` it is more of a reader monad, with inductive hypotheses +being passed down (hence the `ask` and `branch` combinator). +-/ +abbrev M := StateT (Array Expr) MetaM + +namespace M + +variable {α : Type} + +def run (act : M α) : MetaM (α × Array Expr) := StateT.run act #[] +def eval (act : M α) : MetaM α := do return (← run act).1 +def exec (act : M α) : MetaM (Array Expr) := do return (← run act).2 + +def tell (x : Expr) : M Unit := fun xs => pure ((), xs.push x) + +def localM (f : Array Expr → MetaM (Array Expr)) (act : M α) : M α := fun xs => do + let n := xs.size + let (b, xs') ← act xs + pure (b, xs'[*...n] ++ (← f xs'[n...*])) + +def localMapM (f : Expr → MetaM Expr) (act : M α) : M α := + localM (·.mapM f) act + +def ask : M (Array Expr) := get + +def branch (act : M α) : M α := + localM (fun _ => pure #[]) act + +end M + +/-- +The `foldAndCollect` function performs two operations together: + + * it fold recursive calls: applications (and projectsions) of `oldIH` in `e` correspond to + recursive calls, so this function rewrites that back to recursive calls + * it collects induction hypotheses: after replacing `oldIH` with `newIH`, applications thereof + are valuable as induction hypotheses for the cases. + +For well-founded recursion (unary, non-mutual by construction) the terms are rather simple: they +are `oldIH arg proof`, and can be rewritten to `f arg` resp. `newIH arg proof`. But for +structural recursion this can be a more complicted mix of function applications (due to reflexive +data types or extra function arguments) and `PProd` projections (due to the below construction and +mutual function packing), and the main function argument isn't even present. + +To avoid having to think about this, we apply a nice trick: + +We compositionally replace `oldIH` with `newIH`. This likely changes the result type, so when +re-assembling we have to be supple (mainly around `PProd.fst`/`PProd.snd`). As we re-assemble +the term we check if it has type `motive xs..`. If it has, then know we have just found and +rewritten a recursive call, and this type nicely provides us the arguments `xs`. So at this point +we store the rewritten expression as a new induction hypothesis (using `M.tell`) and rewrite to +`f xs..`, which now again has the same type as the original term, and the further re-assembly should +work. Half this logic is in the `isRecCall` parameter. + +If this process fails we’ll get weird type errors (caught later on). We'll see if we need to +improve the errors, for example by passing down a flag whether we expect the same type (and no +occurrences of `newIH`), or whether we are in “supple mode”, and catch it earlier if the rewriting +fails. +-/ +partial def foldAndCollect (oldIH newIH : FVarId) (isRecCall : Expr → Option Expr) (e : Expr) : M Expr := do + unless e.containsFVar oldIH do + return e + withTraceNode `Meta.FunInd (pure m!"{exceptEmoji ·} foldAndCollect ({mkFVar oldIH} → {mkFVar newIH})::{indentExpr e}") do + + let e' ← id do + if let some matcherApp ← matchMatcherApp? e (alsoCasesOn := true) then + if matcherApp.remaining.size == 1 && matcherApp.remaining[0]!.isFVarOf oldIH then + -- We do different things to the matcher when folding recursive calls and when + -- collecting inductive hypotheses. Therefore we do it separately, + -- dropping to `MetaM` in between, and using `M.eval`/`M.exec` as appropriate + -- We could try to do it in one pass by breaking up the `matcherApp.transform` + -- abstraction. + + -- To collect the IHs, we collect them in each branch, and combine + -- them to a type-leve match + let ihMatcherApp' ← liftM <| matcherApp.transform + (onParams := fun e => M.eval <| foldAndCollect oldIH newIH isRecCall e) + (onMotive := fun xs _body => do + -- Remove the old IH that was added in mkFix + let eType ← newIH.getType + let eTypeAbst ← matcherApp.discrs.size.foldRevM (init := eType) fun i _ eTypeAbst => do + let motiveArg := xs[i]! + let discr := matcherApp.discrs[i] + let eTypeAbst ← kabstract eTypeAbst discr + return eTypeAbst.instantiate1 motiveArg + + -- Will later be overridden with a type that’s itself a match + -- statement and the inferred alt types + let dummyGoal := mkConst ``True [] + mkArrow eTypeAbst dummyGoal) + (onAlt := fun _altIdx altType alt => do + lambdaTelescope1 alt fun oldIH' alt => do + forallBoundedTelescope altType (some 1) fun newIH' _goal' => do + let #[newIH'] := newIH' | unreachable! + let altIHs ← M.exec <| foldAndCollect oldIH' newIH'.fvarId! isRecCall alt + let altIH ← PProdN.mk 0 altIHs + mkLambdaFVars #[newIH'] altIH) + (onRemaining := fun _ => pure #[mkFVar newIH]) + let ihMatcherApp'' ← ihMatcherApp'.inferMatchType + M.tell ihMatcherApp''.toExpr + + -- When folding the calls we don't want to remove the extra arg to the matcher + -- that was introduced in the translation + let matcherApp' ← liftM <| matcherApp.transform + (onParams := fun e => M.eval <| foldAndCollect oldIH newIH isRecCall e) + (onMotive := fun _motiveArgs motiveBody => do + let some (_extra, body) := motiveBody.arrow? | throwError "motive not an arrow" + M.eval (foldAndCollect oldIH newIH isRecCall body)) + (onAlt := fun _altIdx altType alt => do + lambdaTelescope1 alt fun oldIH' alt => do + -- We don't have suitable newIH around here, but we don't care since + -- we just want to fold calls. So lets create a fake one. + -- (We cannot use oldIH as that would run into the sanity checks that we could + -- replace all of them) + withLocalDeclD `fakeNewIH (← inferType (mkFVar oldIH')) fun fakeNewIH => + M.eval (foldAndCollect oldIH' fakeNewIH.fvarId! isRecCall alt)) + (onRemaining := fun _ => pure #[]) + return matcherApp'.toExpr + + + if e.getAppArgs.any (·.isFVarOf oldIH) then + -- Sometimes Fix.lean abstracts over oldIH in a proof definition. + -- So beta-reduce that definition. We need to look through theorems here! + if let some e' ← withTransparency .all do unfoldDefinition? e then + return ← foldAndCollect oldIH newIH isRecCall e' + else + throwError "foldAndCollect: cannot reduce application of {e.getAppFn} in:{indentExpr e} " + + match e with + | .app e1 e2 => + pure <|.app (← foldAndCollect oldIH newIH isRecCall e1) (← foldAndCollect oldIH newIH isRecCall e2) + + | .lam n t body bi => + let t' ← foldAndCollect oldIH newIH isRecCall t + withLocalDecl n bi t' fun x => do + M.localMapM (mkLambdaFVars (usedOnly := true) #[x] ·) do + let body' ← foldAndCollect oldIH newIH isRecCall (body.instantiate1 x) + mkLambdaFVars #[x] body' + + | .forallE n t body bi => + let t' ← foldAndCollect oldIH newIH isRecCall t + withLocalDecl n bi t' fun x => do + M.localMapM (mkLambdaFVars (usedOnly := true) #[x] ·) do + let body' ← foldAndCollect oldIH newIH isRecCall (body.instantiate1 x) + mkForallFVars #[x] body' + + | .letE n t v b nondep => + let t' ← foldAndCollect oldIH newIH isRecCall t + let v' ← foldAndCollect oldIH newIH isRecCall v + withLetDecl n t' v' (nondep := nondep) fun x => do + M.localMapM (mkLetFVars (usedLetOnly := true) (generalizeNondepLet := false) #[x] ·) do + let b' ← foldAndCollect oldIH newIH isRecCall (b.instantiate1 x) + mkLetFVars (generalizeNondepLet := false) #[x] b' + + | .mdata m b => + pure <| .mdata m (← foldAndCollect oldIH newIH isRecCall b) + + -- These projections can be type changing (to And), so re-infer their type arguments + | .proj ``PProd 0 e => mkPProdFstM (← foldAndCollect oldIH newIH isRecCall e) + | .proj ``PProd 1 e => mkPProdSndM (← foldAndCollect oldIH newIH isRecCall e) + + | .proj t i e => + pure <| .proj t i (← foldAndCollect oldIH newIH isRecCall e) + + | .sort .. | .lit .. | .const .. | .mvar .. | .bvar .. => + unreachable! -- cannot contain free variables, so early exit above kicks in + + | .fvar fvar => + assert! fvar == oldIH + pure <| mkFVar newIH + + -- Now see if the type of the expression we are building is a motive application. + -- If it is we want to replace it with the corresponding function application, + -- and remember the expression as a IH to be used in an inductive case. + + if e'.containsFVar oldIH then + throwError "Failed to eliminate {mkFVar oldIH} from:{indentExpr e'}" + + let eType ← whnf (← inferType e') + if let .some call := isRecCall eType then + M.tell (← mkExpectedTypeHint e' eType) + return call + + return e' + +-- Because of term duplications we might encounter the same IH multiple times. +-- We deduplicate them (by type, not proof term) here. +-- This could be improved and catch cases where the same IH is used in different contexts. +-- (Cf. `assignSubsumed` in `WF.Fix`) +def deduplicateIHs (vals : Array Expr) : MetaM (Array Expr) := do + let mut vals' := #[] + let mut types := #[] + for v in vals do + let t ← inferType v + unless types.contains t do + vals' := vals'.push v + types := types.push t + return vals' + +def assertIHs (vals : Array Expr) (mvarid : MVarId) : MetaM MVarId := do + let mut mvarid := mvarid + for v in vals.reverse, i in [0:vals.size] do + mvarid ← mvarid.assert (.mkSimple s!"ih{i+1}") (← inferType v) v + return mvarid + +/-- +Goal cleanup: +Substitutes equations (with `substVar`) to remove superfluous variables, and clears unused +let bindings. + +Substitutes from the outside in so that the inner-bound variable name wins, but does a first pass +looking only at variables with names with macro scope, so that preferably they disappear. + +Careful to only touch the context after the motives (given by the index) as the motive could depend +on anything before, and `substVar` would happily drop equations about these fixed parameters. +-/ +partial def cleanupAfter (mvarId : MVarId) (index : Nat) : MetaM MVarId := do + let mvarId ← go mvarId index true + let mvarId ← go mvarId index false + return mvarId +where + go (mvarId : MVarId) (index : Nat) (firstPass : Bool) : MetaM MVarId := do + if let some mvarId ← cleanupAfter? mvarId index firstPass then + go mvarId index firstPass + else + return mvarId + + allHeqToEq (mvarId : MVarId) (index : Nat) : MetaM MVarId := + mvarId.withContext do + let mut mvarId := mvarId + for localDecl in (← getLCtx) do + if localDecl.index > index then + let (_, mvarId') ← heqToEq mvarId localDecl.fvarId + mvarId := mvarId' + return mvarId + + cleanupAfter? (mvarId : MVarId) (index : Nat) (firstPass : Bool) : MetaM (Option MVarId) := do + mvarId.withContext do + for localDecl in (← getLCtx) do + if localDecl.index > index && (!firstPass || localDecl.userName.hasMacroScopes) then + if localDecl.isLet then + if ← Meta.isProp localDecl.type then + if let some mvarId' ← observing? <| mvarId.clearValue localDecl.fvarId then + return some mvarId' + else + continue + if let some mvarId' ← observing? <| mvarId.clear localDecl.fvarId then + return some mvarId' + if let some mvarId' ← substVar? mvarId localDecl.fvarId then + -- After substituting, some HEq might turn into Eqs, and we want to be able to substitute + -- them as well + let mvarId' ← allHeqToEq mvarId' index + return some mvarId' + return none + + +/-- +Second helper monad collecting the cases as mvars +-/ +abbrev M2 α := StateT (Array MVarId) M α + +def M2.run {α} (act : M2 α) : MetaM (α × Array MVarId) := + M.eval (StateT.run (s := #[]) act) + +def M2.branch {α} (act : M2 α) : M2 α := + controlAt M fun runInBase => M.branch (runInBase act) + + +/-- Base case of `buildInductionBody`: Construct a case for the final induction hypthesis. -/ +def buildInductionCase (oldIH newIH : FVarId) (isRecCall : Expr → Option Expr) (toErase toClear : Array FVarId) + (goal : Expr) (e : Expr) : M2 Expr := do + withTraceNode `Meta.FunInd (pure m!"{exceptEmoji ·} buildInductionCase:{indentExpr e}") do + let _e' ← foldAndCollect oldIH newIH isRecCall e + let IHs : Array Expr ← M.ask + let IHs ← deduplicateIHs IHs + + withErasedFVars toErase do + let mvar ← mkFreshExprSyntheticOpaqueMVar goal (tag := `hyp) + let mvarId := mvar.mvarId! + let mvarId ← assertIHs IHs mvarId + trace[Meta.FunInd] "Goal before cleanup:{mvarId}" + let (mvarId, _) ← mvarId.tryClearMany' toClear + modify (·.push mvarId) + let mvar ← instantiateMVars mvar + pure mvar + +/-- +Like `mkLambdaFVars (usedOnly := true)`, but + + * silently skips expression in `xs` that are not `.isFVar` + * returns a mask (same size as `xs`) indicating which variables have been abstracted + (`true` means was abstracted). + +The result `r` can be applied with `r.beta (maskArray mask args)`. + +We use this when generating the functional induction principle to refine the goal through a `match`, +here `xs` are the discriminants of the `match`. +We do not expect non-trivial discriminants to appear in the goal (and if they do, the user will +get a helpful equality into the context). +-/ +def mkLambdaFVarsMasked (xs : Array Expr) (e : Expr) : MetaM (Array Bool × Expr) := do + let mut e := e + let mut xs := xs + let mut mask := #[] + while ! xs.isEmpty do + let discr := xs.back! + if discr.isFVar && e.containsFVar discr.fvarId! then + e ← mkLambdaFVars #[discr] e + mask := mask.push true + else + mask := mask.push false + xs := xs.pop + return (mask.reverse, e) + +/-- `maskArray mask xs` keeps those `x` where the corresponding entry in `mask` is `true` -/ +def maskArray {α} (mask : Array Bool) (xs : Array α) : Array α := Id.run do + let mut ys := #[] + for b in mask, x in xs do + if b then ys := ys.push x + return ys + +/-- +Inverse of `maskArray`: +``` +zipMaskedArray mask (maskArray (mask.map not) xs) (maskArray mask xs) == xs +``` +-/ +def zipMaskedArray {α} (mask : Array Bool) (xs ys : Array α) : Array α := Id.run do + let mut i := 0 + let mut j := 0 + let mut zs := #[] + for b in mask do + if b then + if h : j < ys.size then + zs := zs.push ys[j] + j := j + 1 + else + panic! "zipMaskedArray: not enough elements in ys" + else + if h : i < xs.size then + zs := zs.push xs[i] + i := i + 1 + else + panic! "zipMaskedArray: not enough elements in xs" + return zs + + +/-- +Applies `rw` to `goal`, passes the rewritten `goal'` to `k` (which should return an expression of +type `goal'`), and wraps that using the proof from `rw`. +-/ +def withRewrittenMotive (goal : Expr) (rw : Expr → MetaM Simp.Result) (k : Expr → M2 Expr) : M2 Expr := do + let r ← rw goal + let e ← k r.expr + r.mkEqMPR e + +def inLastArg (rw : Expr → MetaM Simp.Result) (goal : Expr) : MetaM Simp.Result := do + match goal with + | .app goalFn arg => + let r ← rw arg + Simp.mkCongrArg goalFn r + | _ => + return { expr := goal } + +/-- +If `goal` is of the form `motive a b e`, applies `rw` to `e`, passes the simplified +`goal'` to `k` (which should return an expression of type `goal'`), and rewrites that term +accordingly. +-/ +def withRewrittenMotiveArg (goal : Expr) (rw : Expr → MetaM Simp.Result) (k : Expr → M2 Expr) : M2 Expr := do + withRewrittenMotive goal (inLastArg rw) k + +/-- +Use to write inside the packed motives used for mutual structural recursion. +-/ +partial def inProdLambdaLastArg (rw : Expr → MetaM Simp.Result) (goal : Expr) : MetaM Simp.Result := do + match_expr goal with + | PProd.mk _ _ goal1 goal2 => + let r1 ← inProdLambdaLastArg rw goal1 + let r2 ← inProdLambdaLastArg rw goal2 + let f := goal.appFn!.appFn! + Simp.mkCongr (← Simp.mkCongrArg f r1) r2 + | _ => + lambdaTelescope goal fun xs goal => do + let r ← inLastArg rw goal + r.addLambdas xs + +def rwIfWith (hc : Expr) (e : Expr) : MetaM Simp.Result := do + match_expr e with + | ite@ite α c h t f => + let us := ite.constLevels! + if (← isDefEq c (← inferType hc)) then + return { + expr := t + proof? := (mkAppN (mkConst ``if_pos us) #[c, h, hc, α, t, f]) + } + if (← isDefEq (mkNot c) (← inferType hc)) then + return { + expr := f + proof? := (mkAppN (mkConst ``if_neg us) #[c, h, hc, α, t, f]) + } + return { expr := e} + | dite@dite α c h t f => + let us := dite.constLevels! + if (← isDefEq c (← inferType hc)) then + return { + expr := t.beta #[hc] + proof? := (mkAppN (mkConst ``dif_pos us) #[c, h, hc, α, t, f]) + } + if (← isDefEq (mkNot c) (← inferType hc)) then + return { + expr := f.beta #[hc] + proof? := (mkAppN (mkConst ``dif_neg us) #[c, h, hc, α, t, f]) + } + return { expr := e } + | cond@cond α c t f => + let us := cond.constLevels! + if (← isDefEq (← inferType hc) (← mkEq c (mkConst ``Bool.true))) then + return { + expr := t + proof? := (mkAppN (mkConst ``Bool.cond_pos us) #[α, c, t, f, hc]) + } + if (← isDefEq (← inferType hc) (← mkEq c (mkConst ``Bool.false))) then + return { + expr := f + proof? := (mkAppN (mkConst ``Bool.cond_neg us) #[α, c, t, f, hc]) + } + return { expr := e } + | _ => + return { expr := e } + +def rwLetWith (h : Expr) (e : Expr) : MetaM Simp.Result := do + if e.isLet then + if (← isDefEq e.letValue! h) then + return { expr := e.letBody!.instantiate1 h } + return { expr := e } + +def rwMData (e : Expr) : MetaM Simp.Result := do + return { expr := e.consumeMData } + +def rwFun (names : Array Name) (e : Expr) : MetaM Simp.Result := do + e.withApp fun f xs => do + if let some name := names.find? f.isConstOf then + let some unfoldThm ← getUnfoldEqnFor? name (nonRec := true) + | return { expr := e } + let h := mkAppN (mkConst unfoldThm f.constLevels!) xs + let some (_, _, rhs) := (← inferType h).eq? + | throwError "Not an equality: {h}" + return { expr := rhs, proof? := h } + else + return { expr := e } + +def rwMatcher (altIdx : Nat) (e : Expr) : MetaM Simp.Result := do + if e.isAppOf ``PSum.casesOn || e.isAppOf ``PSigma.casesOn then + let mut e := e + while true do + if let some e' ← reduceRecMatcher? e then + e := e'.headBeta + else + let e' := e.headBeta + if e != e' then + e := e' + else + break + return { expr := e } + else + unless (← isMatcherApp e) do + trace[Meta.FunInd] "Not a matcher application:{indentExpr e}" + return { expr := e } + let matcherDeclName := e.getAppFn.constName! + let eqns ← Match.genMatchCongrEqns matcherDeclName + unless altIdx < eqns.size do + trace[Meta.FunInd] "When trying to reduce arm {altIdx}, only {eqns.size} equations for {.ofConstName matcherDeclName}" + return { expr := e } + let eqnThm := eqns[altIdx]! + try + withTraceNode `Meta.FunInd (pure m!"{exceptEmoji ·} rewriting with {.ofConstName eqnThm} in{indentExpr e}") do + let eqProof := mkAppN (mkConst eqnThm e.getAppFn.constLevels!) e.getAppArgs + let (hyps, _, eqType) ← forallMetaTelescope (← inferType eqProof) + trace[Meta.FunInd] "eqProof has type{indentExpr eqType}" + let proof := mkAppN eqProof hyps + let hyps := hyps.map (·.mvarId!) + let (isHeq, lhs, rhs) ← do + if let some (_, lhs, _, rhs) := eqType.heq? then pure (true, lhs, rhs) else + if let some (_, lhs, rhs) := eqType.eq? then pure (false, lhs, rhs) else + throwError m!"Type of {.ofConstName eqnThm} is not an equality" + if !(← isDefEq e lhs) then + throwError m!"Left-hand side {lhs} of {.ofConstName eqnThm} does not apply to {e}" + /- + Here we instantiate the hypotheses of the congruence equation theorem + There are two sets of hypotheses to instantiate: + - `Eq` or `HEq` that relate the discriminants to the patterns + Solving these should instantiate the pattern variables. + - Overlap hypotheses (`isEqnThmHypothesis`) + With more book keeping we could maybe do this very precisely, knowing exactly + which facts provided by the splitter should go where, but it's tedious. + So for now let's use heuristics and try `assumption` and `rfl`. + -/ + for h in hyps do + unless (← h.isAssigned) do + let hType ← h.getType + if Simp.isEqnThmHypothesis hType then + -- Using unrestricted h.substVars here does not work well; it could + -- even introduce a dependency on the `oldIH` we want to eliminate + h.assumption <|> throwError "Failed to discharge {h}" + else if hType.isEq then + h.assumption <|> h.refl <|> throwError m!"Failed to resolve {h}" + else if hType.isHEq then + h.assumption <|> h.hrefl <|> throwError m!"Failed to resolve {h}" + let unassignedHyps ← hyps.filterM fun h => return !(← h.isAssigned) + unless unassignedHyps.isEmpty do + throwError m!"Not all hypotheses of {.ofConstName eqnThm} could be discharged: {unassignedHyps}" + let rhs ← instantiateMVars rhs + let proof ← instantiateMVars proof + let proof ← if isHeq then + try mkEqOfHEq proof + catch e => throwError m!"Could not un-HEq {proof}:{indentD e.toMessageData} " + else + pure proof + return { + expr := rhs + proof? := proof + } + catch ex => + trace[Meta.FunInd] "Failed to apply {.ofConstName eqnThm}:{indentD ex.toMessageData}" + return { expr := e } + +/-- +Builds an expression of type `goal` by replicating the expression `e` into its tail-call-positions, +where it calls `buildInductionCase`. Collects the cases of the final induction hypothesis +as `MVars` as it goes. +-/ +partial def buildInductionBody (toErase toClear : Array FVarId) (goal : Expr) + (oldIH newIH : FVarId) (isRecCall : Expr → Option Expr) (e : Expr) : M2 Expr := do + withTraceNode `Meta.FunInd + (pure m!"{exceptEmoji ·} buildInductionBody: {oldIH.name} → {newIH.name}\ngoal: {goal}:{indentExpr e}") do + + -- if-then-else cause case split: + match_expr e with + | ite _α c h t f => + let c' ← foldAndCollect oldIH newIH isRecCall c + let h' ← foldAndCollect oldIH newIH isRecCall h + let t' ← withLocalDecl `h .default c' fun h => M2.branch do + let t' ← withRewrittenMotiveArg goal (rwIfWith h) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall t + mkLambdaFVars #[h] t' + let f' ← withLocalDecl `h .default (mkNot c') fun h => M2.branch do + let f' ← withRewrittenMotiveArg goal (rwIfWith h) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall f + mkLambdaFVars #[h] f' + let u ← getLevel goal + return mkApp5 (mkConst ``dite [u]) goal c' h' t' f' + | dite _α c h t f => + let c' ← foldAndCollect oldIH newIH isRecCall c + let h' ← foldAndCollect oldIH newIH isRecCall h + let t' ← withLocalDecl `h .default c' fun h => M2.branch do + let t ← instantiateLambda t #[h] + let t' ← withRewrittenMotiveArg goal (rwIfWith h) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall t + mkLambdaFVars #[h] t' + let f' ← withLocalDecl `h .default (mkNot c') fun h => M2.branch do + let f ← instantiateLambda f #[h] + let f' ← withRewrittenMotiveArg goal (rwIfWith h) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall f + mkLambdaFVars #[h] f' + let u ← getLevel goal + return mkApp5 (mkConst ``dite [u]) goal c' h' t' f' + | cond _α c t f => + let c' ← foldAndCollect oldIH newIH isRecCall c + let t' ← withLocalDecl `h .default (← mkEq c' (mkConst ``Bool.true)) fun h => M2.branch do + let t' ← withRewrittenMotiveArg goal (rwIfWith h) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall t + mkLambdaFVars #[h] t' + let f' ← withLocalDecl `h .default (← mkEq c' (mkConst ``Bool.false)) fun h => M2.branch do + let t' ← withRewrittenMotiveArg goal (rwIfWith h) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall f + mkLambdaFVars #[h] t' + let u ← getLevel goal + return mkApp4 (mkConst ``Bool.dcond [u]) goal c' t' f' + | _ => + + + -- Check for unreachable cases. We look for the kind of expressions that `by contradiction` + -- produces + match_expr e with + | False.elim _ h => do + return ← mkFalseElim goal h + | absurd _ _ h₁ h₂ => do + return ← mkAbsurd goal h₁ h₂ + | _ => pure () + if e.isApp && e.getAppFn.isConst && isNoConfusion (← getEnv) e.getAppFn.constName! then + let arity := (← inferType e.getAppFn).getNumHeadForalls -- crucially not reducing the noConfusionType in the type + let h := e.getArg! (arity - 1) + let hType ← inferType h + -- The following duplicates a bit of code from the contradiction tactic, maybe worth extracting + -- into a common helper at some point + if let some (_, lhs, rhs) ← matchEq? hType then + if let some lhsCtor ← matchConstructorApp? lhs then + if let some rhsCtor ← matchConstructorApp? rhs then + if lhsCtor.name != rhsCtor.name then + return (← mkNoConfusion goal h) + + -- we look in to `PProd.mk`, as it occurs in the mutual structural recursion construction + match_expr goal with + | And goal₁ goal₂ => match_expr e with + | PProd.mk _α _β e₁ e₂ => + let e₁' ← buildInductionBody toErase toClear goal₁ oldIH newIH isRecCall e₁ + let e₂' ← buildInductionBody toErase toClear goal₂ oldIH newIH isRecCall e₂ + return mkApp4 (.const ``And.intro []) goal₁ goal₂ e₁' e₂' + | _ => + throwError "Goal is PProd, but expression is:{indentExpr e}" + | True => + return .const ``True.intro [] + | _ => + + -- match and casesOn application cause case splitting + if let some matcherApp ← matchMatcherApp? e (alsoCasesOn := true) then + -- Calculate motive + let eType ← newIH.getType + let motiveBody ← mkArrow eType goal + let (mask, absMotiveBody) ← mkLambdaFVarsMasked matcherApp.discrs motiveBody + + -- A match that refines the parameter has been modified by `Fix.lean` to refine the IH, + -- so we need to replace that IH + if matcherApp.remaining.size == 1 && matcherApp.remaining[0]!.isFVarOf oldIH then + let matcherApp' ← matcherApp.transform (useSplitter := true) + (addEqualities := true) + (onParams := (foldAndCollect oldIH newIH isRecCall ·)) + (onMotive := fun xs _body => pure (absMotiveBody.beta (maskArray mask xs))) + (onAlt := fun altIdx expAltType alt => M2.branch do + lambdaTelescope1 alt fun oldIH' alt => do + forallBoundedTelescope expAltType (some 1) fun newIH' goal' => do + let #[newIH'] := newIH' | unreachable! + let toErase' := toErase ++ #[oldIH', newIH'.fvarId!] + let toClear' := toClear ++ matcherApp.discrs.filterMap (·.fvarId?) + let alt' ← withRewrittenMotiveArg goal' (rwMatcher altIdx) fun goal'' => do + -- logInfo m!"rwMatcher after {matcherApp.matcherName} on{indentExpr goal'}\nyields{indentExpr goal''}" + buildInductionBody toErase' toClear' goal'' oldIH' newIH'.fvarId! isRecCall alt + mkLambdaFVars #[newIH'] alt') + (onRemaining := fun _ => pure #[.fvar newIH]) + return matcherApp'.toExpr + + -- A match that does not refine the parameter, but that we still want to split into separate + -- cases + if matcherApp.remaining.isEmpty then + -- Calculate motive + let (mask, absMotiveBody) ← mkLambdaFVarsMasked matcherApp.discrs goal + + let matcherApp' ← matcherApp.transform (useSplitter := true) + (addEqualities := true) + (onParams := (foldAndCollect oldIH newIH isRecCall ·)) + (onMotive := fun xs _body => pure (absMotiveBody.beta (maskArray mask xs))) + (onAlt := fun altIdx expAltType alt => M2.branch do + withRewrittenMotiveArg expAltType (rwMatcher altIdx) fun expAltType' => + buildInductionBody toErase toClear expAltType' oldIH newIH isRecCall alt) + return matcherApp'.toExpr + + -- we look through mdata + if e.isMData then + let b ← withRewrittenMotiveArg goal (rwMData) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall e.mdataExpr! + return e.updateMData! b + + if let .letE n t v b _ := e then + let t' ← foldAndCollect oldIH newIH isRecCall t + let v' ← foldAndCollect oldIH newIH isRecCall v + return ← withLetDecl n t' v' fun x => M2.branch do + let b' ← withRewrittenMotiveArg goal (rwLetWith x) fun goal' => + buildInductionBody toErase toClear goal' oldIH newIH isRecCall (b.instantiate1 x) + mkLetFVars #[x] b' + + -- Special case for traversing the PProd’ed bodies in our encoding of structural mutual recursion + if let .lam n t b bi := e then + if goal.isForall then + let t' ← foldAndCollect oldIH newIH isRecCall t + return ← withLocalDecl n bi t' fun x => M2.branch do + let goal' ← instantiateForall goal #[x] + let b' ← buildInductionBody toErase toClear goal' oldIH newIH isRecCall (b.instantiate1 x) + mkLambdaFVars #[x] b' + + liftM <| buildInductionCase oldIH newIH isRecCall toErase toClear goal e + +/-- +Given an expression `e` with metavariables `mvars` +* performs more cleanup: + * removes unused let-expressions after index `index` + * tries to substitute variables after index `index` +* lifts them to the current context by reverting all local declarations after index `index` +* introducing a local variable for each of the meta variable +* assigning that local variable to the mvar +* and finally lambda-abstracting over these new local variables. + +This operation only works if the metavariables are independent from each other. + +The resulting meta variable assignment is no longer valid (mentions out-of-scope +variables), so after this operations, terms that still mention these meta variables must not +be used anymore. + +We are not using `mkLambdaFVars` on mvars directly, nor `abstractMVars`, as these at the moment +do not handle delayed assignments correctly. +-/ +def abstractIndependentMVars (mvars : Array MVarId) (index : Nat) (e : Expr) : MetaM Expr := do + trace[Meta.FunInd] "abstractIndependentMVars, to revert after {index}, original mvars: {mvars}" + let mvars ← mvars.mapM fun mvar => do + let mvar ← cleanupAfter mvar index + mvar.withContext do + let fvarIds := (← getLCtx).foldl (init := #[]) (start := index+1) fun fvarIds decl => fvarIds.push decl.fvarId + let (_, mvar) ← mvar.revert fvarIds + pure mvar + trace[Meta.FunInd] "abstractIndependentMVars, reverted mvars: {mvars}" + let names := Array.ofFn (n := mvars.size) fun ⟨i,_⟩ => .mkSimple s!"case{i+1}" + let types ← mvars.mapM MVarId.getType + Meta.withLocalDeclsDND (names.zip types) fun xs => do + for mvar in mvars, x in xs do + mvar.assign x + mkLambdaFVars xs (← instantiateMVars e) + +/-- +Given a unary definition `foo` defined via `WellFounded.fixF`, derive a suitable induction principle +`foo.induct` for it. See module doc for details. + -/ +def deriveUnaryInduction (unfolding : Bool) (name : Name) : MetaM Name := do + let inductName := getFunInductName (unfolding := unfolding) name + realizeConst name inductName (doRealize inductName) + return inductName +where doRealize (inductName : Name) := do + let info ← getConstInfoDefn name + let varNames ← forallTelescope info.type fun xs _ => xs.mapM (·.fvarId!.getUserName) + + -- Uses of WellFounded.fix can be partially applied. Here we eta-expand the body + -- to make sure that `target` indeed the last parameter + let e := info.value + let e ← lambdaTelescope e fun params body => do + if body.isAppOfArity ``WellFounded.fix 5 then + forallBoundedTelescope (← inferType body) (some 1) fun xs _ => do + unless xs.size = 1 do + throwError "functional induction: Failed to eta-expand{indentExpr e}" + mkLambdaFVars (params ++ xs) (mkAppN body xs) + else + pure e + let (e', paramMask) ← lambdaTelescope e fun params funBody => MatcherApp.withUserNames params varNames do + match_expr funBody with + | fix@WellFounded.fix α _motive rel wf body target => + unless params.back! == target do + throwError "functional induction: expected the target as last parameter{indentExpr e}" + let fixedParamPerms := params.pop + let motiveType ← + if unfolding then + withLocalDeclD `r (← instantiateForall info.type params) fun r => + mkForallFVars #[target, r] (.sort 0) + else + mkForallFVars #[target] (.sort 0) + withLocalDeclD `motive motiveType fun motive => do + let fn := mkAppN (← mkConstWithLevelParams name) fixedParamPerms + let isRecCall : Expr → Option Expr := fun e => + e.withApp fun f xs => + if f.isFVarOf motive.fvarId! && xs.size > 0 then + mkApp fn xs[0]! + else + none + + let motiveArg ← + if unfolding then + let motiveArg := mkApp2 motive target (mkAppN (← mkConstWithLevelParams name) params) + mkLambdaFVars #[target] motiveArg + else + pure motive + let e' := .const ``WellFounded.fix [fix.constLevels![0]!, levelZero] + let e' := mkApp4 e' α motiveArg rel wf + check e' + let (body', mvars) ← M2.run do + forallTelescope (← inferType e').bindingDomain! fun xs goal => do + if xs.size ≠ 2 then + throwError "expected recursor argument to take 2 parameters, got {xs}" else + let targets : Array Expr := xs[*...1] + let genIH := xs[1]! + let extraParams := xs[2...*] + -- open body with the same arg + let body ← instantiateLambda body targets + lambdaTelescope1 body fun oldIH body => do + let body ← instantiateLambda body extraParams + let body' ← withRewrittenMotiveArg goal (rwFun #[name]) fun goal => do + buildInductionBody #[oldIH, genIH.fvarId!] #[] goal oldIH genIH.fvarId! isRecCall body + if body'.containsFVar oldIH then + throwError m!"Did not fully eliminate {mkFVar oldIH} from induction principle body:{indentExpr body}" + mkLambdaFVars (targets.push genIH) (← mkLambdaFVars extraParams body') + let e' := mkApp2 e' body' target + let e' ← mkLambdaFVars #[target] e' + let e' ← abstractIndependentMVars mvars (← motive.fvarId!.getDecl).index e' + let e' ← mkLambdaFVars #[motive] e' + + -- We used to pass (usedOnly := false) below in the hope that the types of the + -- induction principle match the type of the function better. + -- But this leads to avoidable parameters that make functional induction strictly less + -- useful (e.g. when the unused parameter mentions bound variables in the users' goal) + let (paramMask, e') ← mkLambdaFVarsMasked fixedParamPerms e' + let e' ← instantiateMVars e' + return (e', paramMask) + | _ => + if funBody.isAppOf ``WellFounded.fix then + throwError "Function {name} defined via WellFounded.fix with unexpected arity {funBody.getAppNumArgs}:{indentExpr funBody}" + else + throwError "Function {name} not defined via WellFounded.fix:{indentExpr funBody}" + + unless (← isTypeCorrect e') do + logError m!"failed to derive a type-correct induction principle:{indentExpr e'}" + check e' + + let eTyp ← inferType e' + let eTyp ← elimTypeAnnotations eTyp + let eTyp ← letToHave eTyp + -- logInfo m!"eTyp: {eTyp}" + let levelParams := (collectLevelParams {} eTyp).params + -- Prune unused level parameters, preserving the original order + let funUs := info.levelParams.toArray + let usMask := funUs.map (levelParams.contains ·) + let us := maskArray usMask funUs |>.toList + + addDecl <| Declaration.thmDecl + { name := inductName, levelParams := us, type := eTyp, value := e' } + + setFunIndInfo { + funName := name + funIndName := inductName + levelMask := usMask + params := paramMask.map (cond · .param .dropped) ++ #[.target] + } + +/-- +Given a realizer for `foo.mutual_induct`, defines `foo.induct`, `bar.induct` etc. +Used for well-founded and structural recursion. +-/ +def projectMutualInduct (unfolding : Bool) (names : Array Name) (mutualInduct : MetaM Name) (finalizeFirstInd : MetaM Unit) : MetaM Unit := do + for name in names, idx in [:names.size] do + let inductName := getFunInductName (unfolding := unfolding) name + realizeConst names[0]! inductName do + let ci ← getConstInfo (← mutualInduct) + let levelParams := ci.levelParams + let value ← forallTelescope ci.type fun xs _body => do + let value := .const ci.name (levelParams.map mkLevelParam) + let value := mkAppN value xs + let value ← PProdN.projM names.size idx value + mkLambdaFVars xs value + let type ← inferType value + let type ← letToHave type + addDecl <| Declaration.thmDecl { name := inductName, levelParams, type, value } + + if idx == 0 then finalizeFirstInd + +/-- +For a (non-mutual!) definition of `name`, uses the `FunIndInfo` associated with the `unaryInduct` and +derives the one for the n-ary function. +-/ +def setNaryFunIndInfo (unfolding : Bool) (fixedParamPerms : FixedParamPerms) (funName : Name) (unaryInduct : Name) : MetaM Unit := do + assert! fixedParamPerms.perms.size = 1 -- only non-mutual for now + let funIndName := getFunInductName (unfolding := unfolding) funName + unless funIndName = unaryInduct do + let some unaryFunIndInfo ← getFunIndInfoForInduct? unaryInduct + | throwError "Expected {unaryInduct} to have FunIndInfo" + let fixedParamPerm := fixedParamPerms.perms[0]! + let mut params := #[] + let mut j := 0 + for h : i in [:fixedParamPerm.size] do + if fixedParamPerm[i].isSome then + assert! j + 1 < unaryFunIndInfo.params.size + params := params.push unaryFunIndInfo.params[j]! + j := j + 1 + else + params := params.push .target + assert! j + 1 = unaryFunIndInfo.params.size + + setFunIndInfo { unaryFunIndInfo with funName, funIndName, params } + +/-- +In the type of `value`, reduces +* Beta-redexes +* `PSigma.casesOn (PSigma.mk a b) (fun x y => k x y) --> k a b` +* `PSum.casesOn (PSum.inl x) k₁ k₂ --> k₁ x` +* `foo._unary (PSum.inl (PSigma.mk a b)) --> foo a b` +and then wraps `value` in an appropriate type hint. + +(The implementation is repetitive and verbose, and should be cleaned up when convenient.) +-/ +def cleanPackedArgs (eqnInfo : WF.EqnInfo) (value : Expr) : MetaM Expr := do + let type ← inferType value + let cleanType ← Meta.transform type (skipConstInApp := true) (post := fun e => do + -- Need to beta-reduce first + let e' := e.headBeta + if e' != e then + return .visit e' + + e.withApp fun f args => do + -- Look for PSigma redexes + if f.isConstOf ``PSigma.casesOn then + if 5 ≤ args.size then + let scrut := args[3]! + let k := args[4]! + let extra := args[5...*] + if scrut.isAppOfArity ``PSigma.mk 4 then + let #[_, _, x, y] := scrut.getAppArgs | unreachable! + let e' := (k.beta #[x, y]).beta extra + return .visit e' + -- Look for PSigma projection + if f.isConstOf ``PSigma.fst then + if h : 3 ≤ args.size then + let scrut := args[2] + let extra := args[3...*] + if scrut.isAppOfArity ``PSigma.mk 4 then + let #[_, _, x, _y] := scrut.getAppArgs | unreachable! + let e' := x.beta extra + return .visit e' + if f.isConstOf ``PSigma.snd then + if h : 3 ≤ args.size then + let scrut := args[2] + let extra := args[3...*] + if scrut.isAppOfArity ``PSigma.mk 4 then + let #[_, _, _x, y] := scrut.getAppArgs | unreachable! + let e' := y.beta extra + return .visit e' + if f.isProj then + let scrut := e.projExpr! + if scrut.isAppOfArity ``PSigma.mk 4 then + let #[_, _, x, y] := scrut.getAppArgs | unreachable! + let e' := (if e.projIdx! = 0 then x else y).beta args + return .visit e' + + -- Look for PSum redexes + if f.isConstOf ``PSum.casesOn then + if 6 ≤ args.size then + let scrut := args[3]! + let k₁ := args[4]! + let k₂ := args[5]! + let extra := args[6...*] + if scrut.isAppOfArity ``PSum.inl 3 then + let e' := (k₁.beta #[scrut.appArg!]).beta extra + return .visit e' + if scrut.isAppOfArity ``PSum.inr 3 then + let e' := (k₂.beta #[scrut.appArg!]).beta extra + return .visit e' + -- Look for _unary redexes + if f.isConstOf eqnInfo.declNameNonRec then + if h : args.size ≥ eqnInfo.fixedParamPerms.numFixed + 1 then + let xs := args[*...eqnInfo.fixedParamPerms.numFixed] + let packedArg := args[eqnInfo.fixedParamPerms.numFixed] + let extraArgs := args[eqnInfo.fixedParamPerms.numFixed<...*] + let some (funIdx, ys) := eqnInfo.argsPacker.unpack packedArg + | throwError "Unexpected packedArg:{indentExpr packedArg}" + let args' := eqnInfo.fixedParamPerms.perms[funIdx]!.buildArgs xs ys + let e' := .const eqnInfo.declNames[funIdx]! e.getAppFn.constLevels! + let e' := mkAppN e' args' + let e' := mkAppN e' extraArgs + return .continue e' + + return .continue e) + mkExpectedTypeHint value cleanType + +/-- +Retrieves `foo._unary.induct`, where the motive is a `PSigma`/`PSum` type, and +unpacks it into a n-ary and (possibly) joint induction principle. +-/ +def unpackMutualInduction (unfolding : Bool) (eqnInfo : WF.EqnInfo) : MetaM Name := do + let inductName := if eqnInfo.declNames.size > 1 then + getMutualInductName (unfolding := unfolding) eqnInfo.declNames[0]! + else + -- If there is no mutual recursion, we generate the `foo.induct` directly. + getFunInductName (unfolding := unfolding) eqnInfo.declNames[0]! + realizeConst eqnInfo.declNames[0]! inductName (doRealize inductName) + return inductName +where doRealize inductName := do + let unaryInductName ← deriveUnaryInduction (unfolding := unfolding) eqnInfo.declNameNonRec + prependError m!"Cannot unpack functional cases principle {.ofConstName unaryInductName} (please report this issue)" do + let ci ← getConstInfo unaryInductName + let us := ci.levelParams + let value := .const ci.name (us.map mkLevelParam) + let motivePos ← forallTelescope ci.type fun xs concl => concl.withApp fun motive targets => do + unless motive.isFVar && targets.size = (if unfolding then 2 else 1) && targets[0]!.isFVar do + throwError "conclusion {concl} does not look like a packed motive application" + let packedTarget := targets[0]! + unless xs.back! == packedTarget do + throwError "packed target not last argument to {unaryInductName}" + let some motivePos := xs.findIdx? (· == motive) + | throwError "could not find motive {motive} in {xs}" + pure motivePos + let value ← forallBoundedTelescope ci.type motivePos fun params type => do + let value := mkAppN value params + eqnInfo.argsPacker.curryParam value type fun motives value type => + -- Bring the rest into scope + forallTelescope type fun xs _concl => do + let alts := xs.pop + let value := mkAppN value alts + let value ← eqnInfo.argsPacker.curry value + let value ← mkLambdaFVars alts value + let value ← mkLambdaFVars motives value + let value ← mkLambdaFVars params value + check value + let value ← cleanPackedArgs eqnInfo value + return value + + unless ← isTypeCorrect value do + logError m!"final term is type incorrect:{indentExpr value}" + check value + let type ← inferType value + let type ← elimOptParam type + let type ← letToHave type + + addDecl <| Declaration.thmDecl + { name := inductName, levelParams := ci.levelParams, type, value } + + if eqnInfo.argsPacker.numFuncs = 1 then + setNaryFunIndInfo (unfolding := unfolding) eqnInfo.fixedParamPerms eqnInfo.declNames[0]! unaryInductName + +def withLetDecls {α} (name : Name) (ts : Array Expr) (es : Array Expr) (k : Array Expr → MetaM α) : MetaM α := do + assert! es.size = ts.size + go 0 #[] +where + go (i : Nat) (acc : Array Expr) := do + if h : i < es.size then + let n := if es.size = 1 then name else name.appendIndexAfter (i + 1) + withLetDecl n ts[i]! es[i] fun x => go (i+1) (acc.push x) + else + k acc + +/-- +Given a recursive definition `foo` defined via structural recursion, derive `foo.mutual_induct`, +if needed, and `foo.induct` for all functions in the group. +See module doc for details. + -/ +def deriveInductionStructural (unfolding : Bool) (names : Array Name) (fixedParamPerms : FixedParamPerms) : MetaM Name := do + let inductName := + if names.size = 1 then + getFunInductName (unfolding := unfolding) names[0]! + else + getMutualInductName (unfolding := unfolding) names[0]! + realizeConst names[0]! inductName (doRealize inductName) + return inductName +where doRealize inductName := do + let infos ← names.mapM getConstInfoDefn + assert! infos.size > 0 + -- First open up the fixed parameters everywhere + let (e', paramMask, motiveArities) ← fixedParamPerms.perms[0]!.forallTelescope infos[0]!.type fun xs => do + -- Now look at the body of an arbitrary of the functions (they are essentially the same + -- up to the final projections) + let body ← fixedParamPerms.perms[0]!.instantiateLambda infos[0]!.value xs + let body := body.eta + + lambdaTelescope body fun ys body => do + -- The body is of the form (brecOn … ).2.2.1 extra1 extra2 etc; ignore the + -- projection here + let f' := body.getAppFn + let body' := PProdN.stripProjs f' + let f := body'.getAppFn + let args := body'.getAppArgs ++ body.getAppArgs + + let body := PProdN.stripProjs body + let .const brecOnName us := f | + throwError "{infos[0]!.name}: unexpected body:{indentExpr infos[0]!.value}\ninstantiated to{indentExpr body}" + unless isBRecOnRecursor (← getEnv) brecOnName do + throwError "{infos[0]!.name}: expected .brecOn application, found:{indentExpr body}" + + let .str indName _ := brecOnName | unreachable! + let indInfo ← getConstInfoInduct indName + + -- we have a `.brecOn` application, so now figure out the length of the fixed prefix + -- we can use the recInfo for `.rec`, since `.brecOn` has a similar structure + let recInfo ← getConstInfoRec (mkRecName indName) + if args.size < recInfo.numParams + recInfo.numMotives + recInfo.numIndices + 1 + recInfo.numMotives then + throwError "insufficient arguments to .brecOn:{indentExpr body}" + let brecOnArgs : Array Expr := args[*...recInfo.numParams] + let _brecOnMotives : Array Expr := args[recInfo.numParams...(recInfo.numParams + recInfo.numMotives)] + let brecOnTargets : Array Expr := args[(recInfo.numParams + recInfo.numMotives)... + (recInfo.numParams + recInfo.numMotives + recInfo.numIndices + 1)] + let brecOnMinors : Array Expr := args[(recInfo.numParams + recInfo.numMotives + recInfo.numIndices + 1)... + (recInfo.numParams + recInfo.numMotives + recInfo.numIndices + 1 + recInfo.numMotives)] + let brecOnExtras : Array Expr := args[(recInfo.numParams + recInfo.numMotives + recInfo.numIndices + 1 + + recInfo.numMotives)...*] + unless brecOnTargets.all (·.isFVar) do + throwError "the indices and major argument of the brecOn application are not variables:{indentExpr body}" + unless brecOnExtras.all (·.isFVar) do + throwError "the extra arguments to the brecOn application are not variables:{indentExpr body}" + let _ :: indLevels := us | throwError "Too few universe parameters in .brecOn application:{indentExpr body}" + + let group : Structural.IndGroupInst := { Structural.IndGroupInfo.ofInductiveVal indInfo with + levels := indLevels, params := brecOnArgs } + + -- We also need to know the number of indices of each type former, including the auxiliary + -- type formers that do not have IndInfo. We can read it off the motives types of the recursor. + let numTypeFormerTargetss ← do + let aux := mkAppN (.const recInfo.name (0 :: group.levels)) group.params + let motives ← inferArgumentTypesN recInfo.numMotives aux + motives.mapM fun motive => + forallTelescopeReducing motive fun xs _ => pure xs.size + + -- Recreate the recArgInfos. Maybe more robust and simpler to store relevant parts in the EqnInfos? + let recArgInfos ← infos.mapIdxM fun funIdx info => do + let some eqnInfo := Structural.eqnInfoExt.find? (← getEnv) info.name | throwError "{info.name} missing eqnInfo" + let value ← fixedParamPerms.perms[funIdx]!.instantiateLambda info.value xs + let recArgInfo' ← lambdaTelescope value fun ys _ => + let args := fixedParamPerms.perms[funIdx]!.buildArgs xs ys + Structural.getRecArgInfo info.name fixedParamPerms.perms[funIdx]! args eqnInfo.recArgPos + let #[recArgInfo] ← Structural.argsInGroup group xs value #[recArgInfo'] + | throwError "Structural.argsInGroup did not return a recArgInfo" + pure recArgInfo + + let positions : Structural.Positions := .groupAndSort (·.indIdx) recArgInfos (Array.range indInfo.numTypeFormers) + + -- Below we'll need the types of the motive arguments (brecOn argument order) + let brecMotiveTypes ← inferArgumentTypesN recInfo.numMotives (group.brecOn 0 0) + trace[Meta.FunInd] m!"brecMotiveTypes: {brecMotiveTypes}" + assert! brecMotiveTypes.size = positions.size + + -- Remove the varying parameters from the environment + withErasedFVars (ys.map (·.fvarId!)) do + -- The brecOnArgs, brecOnMotives and brecOnMinor should still be valid in this + -- context, and are the parts relevant for every function in the mutual group + + -- Calculate the types of the induction motives (natural argument order) for each function + let motiveTypes ← infos.mapIdxM fun funIdx info => do + let funType ← fixedParamPerms.perms[funIdx]!.instantiateForall info.type xs + forallBoundedTelescope funType (some (fixedParamPerms.perms[funIdx]!.size - xs.size)) fun ys rType => do + if unfolding then + withLocalDeclD `r rType fun r => + mkForallFVars (ys.push r) (.sort 0) + else + mkForallFVars ys (.sort 0) + trace[Meta.FunInd] m!"motiveTypes: {motiveTypes}" + let motiveArities ← motiveTypes.mapM fun motiveType => + forallTelescope motiveType fun ys _ => pure ys.size + let motiveNames := Array.ofFn (n := infos.size) fun ⟨i, _⟩ => + if infos.size = 1 then .mkSimple "motive" else .mkSimple s!"motive_{i+1}" + + withLocalDeclsDND (motiveNames.zip motiveTypes) fun motives => do + + -- Prepare the `isRecCall` that recognizes recursive calls + let isRecCall : Expr → Option Expr := fun e => do + if let .some funIdx := motives.idxOf? e.getAppFn then + if e.getAppNumArgs = motiveArities[funIdx]! then + let info := infos[funIdx]! + let args := if unfolding then e.getAppArgs.pop else e.getAppArgs + let args := fixedParamPerms.perms[funIdx]!.buildArgs xs args + return mkAppN (.const info.name (info.levelParams.map mkLevelParam)) args + .none + + -- Motives with parameters reordered, to put indices and major first, + -- and (when unfolding) the result field instantiated + let mut brecMotives := #[] + for motive in motives, recArgInfo in recArgInfos, info in infos, funIdx in [:motives.size] do + let brecMotive ← forallTelescope (← inferType motive) fun ys _ => do + let ys := if unfolding then ys.pop else ys + let (indicesMajor, rest) := recArgInfo.pickIndicesMajor ys + let motiveArg := mkAppN motive ys + let motiveArg ← if unfolding then + let args := fixedParamPerms.perms[funIdx]!.buildArgs xs ys + let fnCall := mkAppN (.const info.name (info.levelParams.map mkLevelParam)) args + pure <| mkApp motiveArg fnCall + else + pure motiveArg + mkLambdaFVars indicesMajor (← mkForallFVars rest motiveArg) + brecMotives := brecMotives.push brecMotive + + -- We need to pack these motives according to the `positions` assignment. + let packedMotives ← positions.mapMwith PProdN.packLambdas brecMotiveTypes brecMotives + trace[Meta.FunInd] m!"packedMotives: {packedMotives}" + + -- Now we can calculate the expected types of the minor arguments + let minorTypes ← inferArgumentTypesN recInfo.numMotives <| + mkAppN (group.brecOn 0 0) (packedMotives ++ brecOnTargets) + trace[Meta.FunInd] m!"minorTypes: {minorTypes}" + -- So that we can transform them + let (minors', mvars) ← M2.run do + let mut minors' := #[] + for brecOnMinor in brecOnMinors, goal in minorTypes, numTargets in numTypeFormerTargetss do + let minor' ← forallTelescope goal fun xs goal => do + unless xs.size ≥ numTargets do + throwError ".brecOn argument has too few parameters, expected at least {numTargets}: {xs}" + let targets : Array Expr := xs[*...numTargets] + let genIH := xs[numTargets]! + let extraParams := xs[numTargets<...*] + -- open body with the same arg + let body ← instantiateLambda brecOnMinor targets + lambdaTelescope1 body fun oldIH body => do + trace[Meta.FunInd] "replacing {Expr.fvar oldIH} with {genIH}" + let body ← instantiateLambda body extraParams + let body' ← + withRewrittenMotive goal (inProdLambdaLastArg (rwFun names)) fun goal' => + buildInductionBody #[oldIH, genIH.fvarId!] #[] goal' oldIH genIH.fvarId! isRecCall body + if body'.containsFVar oldIH then + throwError m!"Did not fully eliminate {mkFVar oldIH} from induction principle body:{indentExpr body}" + mkLambdaFVars (targets.push genIH) (← mkLambdaFVars extraParams body') + minors' := minors'.push minor' + pure minors' + trace[Meta.FunInd] "processed minors: {minors'}" + + -- Now assemble the mutual_induct theorem + -- Let-bind the transformed minors to avoid code duplication of possibly very large + -- terms when we have mutual induction. + let e' ← withLetDecls `minor minorTypes minors' fun minors' => do + let mut brecOnApps := #[] + for info in infos, recArgInfo in recArgInfos, idx in [:infos.size] do + -- Take care to pick the `ys` from the type, to get the variable names expected + -- by the user, but use the value arity + let arity ← lambdaTelescope (← fixedParamPerms.perms[idx]!.instantiateLambda info.value xs) fun ys _ => pure ys.size + let e ← forallBoundedTelescope (← fixedParamPerms.perms[idx]!.instantiateForall info.type xs) arity fun ys _ => do + let (indicesMajor, rest) := recArgInfo.pickIndicesMajor ys + -- Find where in the function packing we are (TODO: abstract out) + let some indIdx := positions.findIdx? (·.contains idx) | panic! "invalid positions" + let some pos := positions.find? (·.contains idx) | panic! "invalid positions" + let some packIdx := pos.findIdx? (· == idx) | panic! "invalid positions" + let e := group.brecOn 0 indIdx + let e := mkAppN e packedMotives + let e := mkAppN e indicesMajor + let e := mkAppN e minors' + let e ← PProdN.projM pos.size packIdx e + let e := mkAppN e rest + let e ← mkLambdaFVars ys e + trace[Meta.FunInd] "assembled call for {info.name}: {e}" + pure e + brecOnApps := brecOnApps.push e + mkLetFVars minors' (← PProdN.mk 0 brecOnApps) + let e' ← abstractIndependentMVars mvars (← motives.back!.fvarId!.getDecl).index e' + let e' ← mkLambdaFVars motives e' + + -- We used to pass (usedOnly := false) below in the hope that the types of the + -- induction principle match the type of the function better. + -- But this leads to avoidable parameters that make functional induction strictly less + -- useful (e.g. when the unused parameter mentions bound variables in the users' goal) + let (paramMask, e') ← mkLambdaFVarsMasked xs e' + let e' ← instantiateMVars e' + trace[Meta.FunInd] "complete body of mutual induction principle:{indentExpr e'}" + pure (e', paramMask, motiveArities) + + unless (← isTypeCorrect e') do + logError m!"constructed induction principle is not type correct:{indentExpr e'}" + check e' + + let eTyp ← inferType e' + let eTyp ← elimTypeAnnotations eTyp + let eTyp ← letToHave eTyp + -- logInfo m!"eTyp: {eTyp}" + let levelParams := (collectLevelParams {} eTyp).params + -- Prune unused level parameters, preserving the original order + let funUs := infos[0]!.levelParams.toArray + let usMask := funUs.map (levelParams.contains ·) + let us := maskArray usMask funUs |>.toList + + addDecl <| Declaration.thmDecl + { name := inductName, levelParams := us, type := eTyp, value := e' } + + if names.size = 1 then + let mut params := #[] + for h : i in [:fixedParamPerms.perms[0]!.size] do + if let some idx := fixedParamPerms.perms[0]![i] then + if paramMask[idx]! then + params := params.push .param + else + params := params.push .dropped + else + params := params.push .target + + setFunIndInfo { + funName := names[0]! + funIndName := inductName, levelMask := usMask, params := params + } + + +/-- +Given an expression `fun x y z => body`, returns a bit mask of the functinon's arity length +that has `true` whenever that parameter of the function appears as a scrutinee of a `match` in +tail position. These are the parameters that are likely useful as targets of the motive +of the functional cases theorem. All others become parameters or may be dropped. + +-/ +partial def refinedArguments (e : Expr) : MetaM (Array Bool) := do + let (_, mask) ← lambdaTelescope e fun xs body => + let mask0 := Array.replicate xs.size false + go xs body |>.run mask0 + let mut mask := mask + let revDeps ← getParamRevDeps e + assert! revDeps.size = mask.size + for i in [:mask.size] do + if mask[i]! then + for j in revDeps[i]! do + mask := mask.set! j true + pure mask +where + -- NB: we process open terms here. + go (xs : Array Expr) (e : Expr) : StateT (Array Bool) MetaM Unit := do + let e := e.consumeMData + + if e.isLambda then + -- Not strictly tail position, but simplifies the code below and should not make + -- a difference in practice + go xs e.bindingBody! + else if e.isLet then + go xs e.letBody! + else + e.withApp fun f args => do + if f.isConst then + if let some matchInfo ← getMatcherInfo? f.constName! then + for scrut in args[matchInfo.getFirstDiscrPos...matchInfo.getFirstAltPos] do + if let some i := xs.idxOf? scrut then + modify (·.set! i true) + for alt in args[matchInfo.getFirstAltPos...matchInfo.arity] do + go xs alt + if f.isConstOf ``ite || f.isConstOf ``dite then + for arg in args[3...5] do + go xs arg + if f.isConstOf ``cond then + for arg in args[2...4] do + go xs arg + +/-- +For non-recursive (and recursive functions) functions we derive a “functional case splitting theorem”. This is very similar +than the functional induction theorem. It splits the goal, but does not give you inductive hyptheses. + +For these, it is not really clear which parameters should be “targets” of the motive, as there is +no “fixed prefix” to guide this decision. All? None? Some? + +We tried none, but that did not work well. Right now it's all parameters, and it seems to work well. +In the future, we might post-process the theorem (or run the code below iteratively) and remove +targets that are unchanged in each case, so simplify applying the lemma when these “fixed” parameters +are not variables, to avoid having to generalize them. +-/ +def deriveCases (unfolding : Bool) (name : Name) : MetaM Unit := do + let casesName := getFunCasesName (unfolding := unfolding) name + realizeConst name casesName do + prependError m!"Cannot derive functional cases principle (please report this issue)" do + let info ← getConstInfo name + let some unfoldEqnName ← getUnfoldEqnFor? (nonRec := true) name + | throwError "'{name}' does not have an unfold theorem nor a value" + let value ← do + let eqInfo ← getConstInfo unfoldEqnName + forallTelescope eqInfo.type fun xs body => do + let some (_, _, rhs) := body.eq? + | throwError "Type of {unfoldEqnName} not an equality: {body}" + mkLambdaFVars xs rhs + let targetMask ← refinedArguments value + trace[Meta.FunInd] "targetMask: {targetMask}" + + let (paramsMask, e') ← lambdaTelescope value fun xs _ => do + let params := maskArray (targetMask.map not) xs + let targets := maskArray targetMask xs + let motiveType ← + if unfolding then + withLocalDeclD `r (← instantiateForall info.type xs) fun r => + mkForallFVars (targets.push r) (.sort 0) + else + mkForallFVars targets (.sort 0) + -- Remove targets from local context, we want to bring them into scope after the motive + -- so that the index passed to `abstractIndependentMVars` works. + withErasedFVars (targets.map (·.fvarId!)) do + withLocalDeclD `motive motiveType fun motive => do + -- Bring targets freshly into scope again + forallBoundedTelescope motiveType targets.size fun targets _ => do + let (e', mvars) ← M2.run do + let args := zipMaskedArray targetMask params targets + let body := value.beta args + let goal := mkAppN motive targets + let goal ← if unfolding then + pure <| mkApp goal (mkAppN (← mkConstWithLevelParams name) args) + else + pure goal + withRewrittenMotiveArg goal (rwFun #[name]) fun goal => do + -- We bring an unused FVars into scope to pass as `oldIH` and `newIH`. These do not appear anywhere + -- so `buildInductionBody` should just do the right thing + withLocalDeclD `fakeIH (mkConst ``Unit) fun fakeIH => + let isRecCall := fun _ => none + buildInductionBody #[fakeIH.fvarId!] #[] goal fakeIH.fvarId! fakeIH.fvarId! isRecCall body + let e' ← mkLambdaFVars targets e' + let e' ← abstractIndependentMVars mvars (← motive.fvarId!.getDecl).index e' + let e' ← mkLambdaFVars #[motive] e' + mkLambdaFVarsMasked params e' + + mapError (f := (m!"constructed functional cases principle is not type correct:{indentExpr e'}\n{indentD ·}")) do + check e' + + let eTyp ← inferType e' + let eTyp ← elimTypeAnnotations eTyp + let eTyp ← letToHave eTyp + -- logInfo m!"eTyp: {eTyp}" + let levelParams := (collectLevelParams {} eTyp).params + -- Prune unused level parameters, preserving the original order + let funUs := info.levelParams.toArray + let usMask := funUs.map (levelParams.contains ·) + let us := maskArray usMask funUs |>.toList + + addDecl <| Declaration.thmDecl + { name := casesName, levelParams := us, type := eTyp, value := e' } + + -- Calculate paramsKind from targetMask (length = arity) and paramsMask (length = params) + let mut paramKinds := #[] + let mut j := 0 + for isTarget in targetMask do + if isTarget then + paramKinds := paramKinds.push .target + else + assert! j < paramsMask.size + if paramsMask[j]! then + paramKinds := paramKinds.push .param + else + paramKinds := paramKinds.push .dropped + j := j + 1 + + setFunIndInfo { + funName := name + funIndName := casesName + levelMask := usMask + params := paramKinds + } + +/-- +Given a recursively defined function `foo`, derives `foo.induct`. See the module doc for details. +-/ +def deriveInduction (unfolding : Bool) (name : Name) : MetaM Unit := do + prependError m!"Cannot derive functional induction principle (please report this issue)" do + if let some eqnInfo := WF.eqnInfoExt.find? (← getEnv) name then + let unaryInductName ← deriveUnaryInduction unfolding eqnInfo.declNameNonRec + if eqnInfo.declNames.size > 1 then + projectMutualInduct unfolding eqnInfo.declNames (unpackMutualInduction unfolding eqnInfo) do + -- We set the FunIndInfo on the first induction principle, which must happen inside its + -- realization. + if eqnInfo.argsPacker.numFuncs = 1 then + setNaryFunIndInfo (unfolding := unfolding) eqnInfo.fixedParamPerms eqnInfo.declNames[0]! unaryInductName + else + -- (in this case, `unpackMutualInduction` already does `setNaryFunIndInfo`) + let _ ← unpackMutualInduction unfolding eqnInfo + else if let some eqnInfo := Structural.eqnInfoExt.find? (← getEnv) name then + if eqnInfo.declNames.size > 1 then + projectMutualInduct unfolding eqnInfo.declNames (deriveInductionStructural unfolding eqnInfo.declNames eqnInfo.fixedParamPerms) (pure ()) + else + let _ ← deriveInductionStructural unfolding eqnInfo.declNames eqnInfo.fixedParamPerms + else + throwError "constant '{name}' is not structurally or well-founded recursive" + +def isFunInductName (env : Environment) (name : Name) : Bool := Id.run do + let .str p s := name | return false + match s with + | "induct" + | "induct_unfolding" => + unless env.isSafeDefinition p do return false + if let some eqnInfo := WF.eqnInfoExt.find? env p then + return true + if (Structural.eqnInfoExt.find? env p).isSome then return true + return false + | "mutual_induct" + | "mutual_induct_unfolding" => + unless env.isSafeDefinition p do return false + if let some eqnInfo := WF.eqnInfoExt.find? env p then + if h : eqnInfo.declNames.size > 1 then + return eqnInfo.declNames[0] = p + if let some eqnInfo := Structural.eqnInfoExt.find? env p then + if h : eqnInfo.declNames.size > 1 then + return eqnInfo.declNames[0] = p + return false + | _ => return false + + +def isFunCasesName (env : Environment) (name : Name) : Bool := Id.run do + let .str p s := name | return false + match s with + | "fun_cases" + | "fun_cases_unfolding" => + unless env.isSafeDefinition p do return false + return true + | _ => return false + +builtin_initialize + registerReservedNamePredicate isFunInductName + registerReservedNamePredicate isFunCasesName + + registerReservedNameAction fun name => do + if isFunInductName (← getEnv) name then + let .str p s := name | return false + unless (← getEnv).isSafeDefinition p do return false + let unfolding := s.endsWith "_unfolding" + MetaM.run' <| deriveInduction unfolding p + return true + if isFunCasesName (← getEnv) name then + let .str p s := name | return false + unless (← getEnv).isSafeDefinition p do return false + let unfolding := s == "fun_cases_unfolding" + MetaM.run' <| deriveCases unfolding p + return true + return false + +end Lean.Tactic.FunInd + +builtin_initialize + Lean.registerTraceClass `Meta.FunInd diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunIndCollect.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunIndCollect.lean new file mode 100644 index 0000000000000000000000000000000000000000..f9d4685d2aed9ac6f77d7dd57adb77d97f709636 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunIndCollect.lean @@ -0,0 +1,115 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Joachim Breitner +-/ + +prelude +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.FunIndInfo + +/-! +Support for collecting function calls that could be used for `fun_induction` or `fun_cases`. +Used by `fun_induction foo`, and the `Calls` structure is also used by `try?`. +-/ + +namespace Lean.Meta.FunInd + +structure Call where + /- The full function call -/ + expr : Expr + /- Used to avoid adding calls that differ only in dropped arguments -/ + relevantArgs : Expr + deriving Hashable, BEq + +structure SeenCalls where + /-- the full calls -/ + calls : Array Expr + /-- only function name and relevant arguments -/ + seen : Std.HashSet (Name × Array Expr) + +instance : EmptyCollection SeenCalls where + emptyCollection := ⟨#[], {}⟩ + +def SeenCalls.isEmpty (sc : SeenCalls) : Bool := + sc.calls.isEmpty + +def SeenCalls.push (e : Expr) (funIndInfo : FunIndInfo) (args : Array Expr) (calls : SeenCalls) : + MetaM SeenCalls := do + if funIndInfo.params.size != args.size then return calls + let mut keys := #[] + for arg in args, kind in funIndInfo.params do + if kind matches .target then + if !arg.isFVar then return calls + unless kind matches .dropped do + keys := keys.push arg + let key := (funIndInfo.funName, keys) + if calls.seen.contains key then return calls + return { calls := calls.calls.push e, seen := calls.seen.insert key } + +/-- +Which functions have exactly one candidate application. Used by `try?` to determine whether +we can use `fun_induction foo` or need `fun_induction foo x y z`. +-/ +def SeenCalls.uniques (calls : SeenCalls) : NameSet := Id.run do + let mut seen : NameSet := {} + let mut seenTwice : NameSet := {} + for (n, _) in calls.seen do + unless seenTwice.contains n do + if seen.contains n then + seenTwice := seenTwice.insert n + else + seen := seen.insert n + return seen.filter (! seenTwice.contains ·) + +namespace Collector + +abbrev M := ReaderT FunIndInfo <| StateRefT SeenCalls MetaM + +def saveFunInd (e : Expr) (funIndInfo : FunIndInfo) (args : Array Expr) : M Unit := do + set (← (← get).push e funIndInfo args) + +def visitApp (e : Expr) (funIndInfo : FunIndInfo) (args : Array Expr) : M Unit := do + saveFunInd e funIndInfo args + +unsafe abbrev Cache := PtrSet Expr + +unsafe def visit (e : Expr) : StateRefT Cache M Unit := do + unless (← get).contains e do + modify fun s => s.insert e + match e with + | .const _ _ => pure () + | .forallE _ d b _ => visit d; visit b + | .lam _ d b _ => visit d; visit b + | .mdata _ b => visit b + | .letE _ t v b _ => visit t; visit v; visit b + | .app .. => e.withApp fun f args => do + if let .const declName _ := f then + if declName = (← read).funName then + unless e.hasLooseBVars do -- TODO: We can allow them in `.dropped` arguments + visitApp e (← read) args + else + visit f + args.forM visit + | .proj _ _ b => visit b + | _ => return () + +unsafe def main (needle : FunIndInfo) (mvarId : MVarId) : MetaM (Array Expr) := mvarId.withContext do + let (_, s) ← go |>.run mkPtrSet |>.run needle |>.run {} + return s.calls +where + go : StateRefT Cache M Unit := do + for localDecl in (← getLCtx) do + unless localDecl.isAuxDecl do + if let some val := localDecl.value? then + visit val + else + visit localDecl.type + visit (← mvarId.getType) + +end Collector + +def collect (needle : FunIndInfo) (mvarId : MVarId) : MetaM (Array Expr) := do + unsafe Collector.main needle mvarId + +end Lean.Meta.FunInd diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunIndInfo.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunIndInfo.lean new file mode 100644 index 0000000000000000000000000000000000000000..a0e450016fabb7a66d88422fce400ddf67eb19a5 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/FunIndInfo.lean @@ -0,0 +1,86 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Joachim Breitner +-/ + +prelude +import Lean.Meta.Basic +import Lean.ScopedEnvExtension +import Lean.ReservedNameAction + +/-! +This module defines the data structure and environment extension to remember how to map the +function's arguments to the functional induction principle's arguments. +Also used for functional cases. +-/ + +namespace Lean.Meta + +inductive FunIndParamKind where + | dropped + | param + | target +deriving BEq, Repr, Inhabited + +/-- +A `FunIndInfo` indicates how a function's arguments map to the arguments of the functional induction +(resp. cases) theorem. + +The size of `params` also indicates the arity of the function. +-/ +structure FunIndInfo where + funName : Name + funIndName : Name + /-- + `true` means that the corresponding level parameter of the function is also a level param + of the induction principle. + -/ + levelMask : Array Bool + params : Array FunIndParamKind +deriving Inhabited, Repr + +builtin_initialize funIndInfoExt : MapDeclarationExtension FunIndInfo ← mkMapDeclarationExtension + +def getFunInductName (declName : Name) (unfolding : Bool := false) : Name := + if unfolding then + declName ++ `induct_unfolding + else + declName ++ `induct + +def getFunCasesName (declName : Name) (unfolding : Bool := false) : Name := + if unfolding then + declName ++ `fun_cases_unfolding + else + declName ++ `fun_cases + +def getMutualInductName (declName : Name) (unfolding : Bool := false) : Name := + if unfolding then + declName ++ `mutual_induct_unfolding + else + declName ++ `mutual_induct + +def getFunInduct? (unfolding : Bool) (cases : Bool) (declName : Name) : CoreM (Option Name) := do + let .defnInfo _ ← getConstInfo declName | return none + try + let thmName := if cases then + getFunCasesName (unfolding := unfolding) declName + else + getFunInductName (unfolding := unfolding) declName + let result ← realizeGlobalConstNoOverloadCore thmName + return some result + catch _ => + return none + +def setFunIndInfo (funIndInfo : FunIndInfo) : CoreM Unit := do + assert! !(funIndInfoExt.contains (← getEnv) funIndInfo.funIndName) + modifyEnv fun env => funIndInfoExt.insert env funIndInfo.funIndName funIndInfo + +def getFunIndInfoForInduct? (inductName : Name) : CoreM (Option FunIndInfo) := do + return funIndInfoExt.find? (← getEnv) inductName + +def getFunIndInfo? (cases : Bool) (unfolding : Bool) (funName : Name) : CoreM (Option FunIndInfo) := do + let some inductName ← getFunInduct? (cases := cases) (unfolding := unfolding) funName | return none + getFunIndInfoForInduct? inductName + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Generalize.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Generalize.lean new file mode 100644 index 0000000000000000000000000000000000000000..50ae3a10c42c1a222f0ce976dbb1c7a63a5f4d73 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Generalize.lean @@ -0,0 +1,152 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.KAbstract +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Intro +import Lean.Meta.Tactic.FVarSubst +import Lean.Meta.Tactic.Revert + +namespace Lean.Meta + +/-- The `generalize` tactic takes arguments of the form `h : e = x` -/ +structure GeneralizeArg where + expr : Expr + xName? : Option Name := none + hName? : Option Name := none + deriving Inhabited + +/-- +Telescopic `generalize` tactic. It can simultaneously generalize many terms. +It uses `kabstract` to occurrences of the terms that need to be generalized. +-/ +private partial def generalizeCore (mvarId : MVarId) (args : Array GeneralizeArg) + (transparency : TransparencyMode) : MetaM (Array FVarId × MVarId) := + mvarId.withContext do + mvarId.checkNotAssigned `generalize + let tag ← mvarId.getTag + let target ← instantiateMVars (← mvarId.getType) + let rec go (i : Nat) : MetaM Expr := do + if _h : i < args.size then + let arg := args[i] + let e ← instantiateMVars arg.expr + let eType ← instantiateMVars (← inferType e) + let type ← go (i+1) + let xName ← if let some xName := arg.xName? then pure xName else mkFreshUserName `x + return Lean.mkForall xName BinderInfo.default eType + (← withTransparency transparency <| kabstract type e) + else + return target + let targetNew ← go 0 + unless (← isTypeCorrect targetNew) do + throwTacticEx `generalize mvarId m!"result is not type correct{indentExpr targetNew}" + let es := args.map (·.expr) + if !args.any fun arg => arg.hName?.isSome then + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew tag + mvarId.assign (mkAppN mvarNew es) + mvarNew.mvarId!.introNP args.size + else + let (rfls, targetNew) ← forallBoundedTelescope targetNew args.size fun xs type => do + let rec go' (i : Nat) : MetaM (List Expr × Expr) := do + if _h : i < xs.size then + let arg := args[i]! + if let some hName := arg.hName? then + let xType ← inferType xs[i] + let e ← instantiateMVars arg.expr + let eType ← instantiateMVars (← inferType e) + let (hType, r) ← if (← isDefEq xType eType) then + pure (← mkEq e xs[i], ← mkEqRefl e) + else + pure (← mkHEq e xs[i], ← mkHEqRefl e) + let (rs, type) ← go' (i+1) + return (r :: rs, mkForall hName BinderInfo.default hType type) + else + go' (i+1) + else + return ([], type) + let (rfls, type) ← go' 0 + return (rfls, ← mkForallFVars xs type) + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew tag + mvarId.assign (mkAppN (mkAppN mvarNew es) rfls.toArray) + mvarNew.mvarId!.introNP (args.size + rfls.length) + +/- +Remark: we use `TransparencyMode.instances` as the default setting at `generalize` +and `generalizeHyp` to avoid excessive resource usage. + +**Motivation:** +The `kabstract e p` operation is widely used, for instance, in the `generalize` tactic. +It operates by taking an expression `e` and a pattern (i.e., an expression containing metavariables) +and employs keyed-matching to identify and abstract instances of `p` within `e`. +For example, if `e` is `a + (2 * (b + c))` and `p` is `2 * ?m`, the resultant expression +would be `a + #0`, where `#0` represents a loose bound variable. + +This matching process is not merely syntactic; it also considers reduction. It's impractical +to attempt matching each sub-term with `p`; therefore, only sub-terms sharing the same "root" +symbol are evaluated. For instance, with the pattern `2 * ?m`, only sub-terms with the +root `*` are considered. Matching is executed using the definitionally equality test +(i.e., `isDefEq`). + +The `generalize` tactic employs `kabstract` and defaults to standard reducibility. +Hence, the `isDefEq` operations invoked by `kabstract` can become highly resource-intensive +and potentially trigger "max recursion depth reached" errors, as observed in issue #3524. +This issue was isolated by @**Kim Morrison** with the following example: +``` +example (a : Nat) : ((2 ^ 7) + a) - 2 ^ 7 = 0 := by + generalize 0 - 0 = x +``` +In this scenario, `kabstract` triggers a "max recursion depth reached" error while +testing whether `((2 ^ 7) + a) - 2 ^ 7` is definitionally equal to `0 - 0`. +Note that the term `((2 ^ 7) + a) - 2 ^ 7` is not ground. +We believe most users find the error message to be uninformative and unexpected. +To fix this issue, we decided to use `TransparencyMode.instances` as the default setting. + +Kyle Miller has performed the following analysis on the potential impact of the +changes on Mathlib (2024-03-02). + +There seem to be just 130 cases of generalize in Mathlib, and after looking through a +good number of them, they seem to come in just two types: + +- Ones where it looks like reducible+instance transparency should work, where in + particular there is nothing obvious being reduced, and +- Ones that don't make use of the `kabstract` feature at all (it's being used like a + `have` that introduces an equality for rewriting). + +That wasn't a systematic review of generalize though. It's possible changing the +transparency settings would break things, but in my opinion it would be better +if generalize weren't used for unfolding things. +-/ + +@[inherit_doc generalizeCore] +def _root_.Lean.MVarId.generalize (mvarId : MVarId) (args : Array GeneralizeArg) + (transparency := TransparencyMode.instances) : MetaM (Array FVarId × MVarId) := + generalizeCore mvarId args transparency + +/-- +Extension of `generalize` to support generalizing within specified hypotheses. +The `hyps` array contains the list of hypotheses within which to look for occurrences +of the generalizing expressions. +-/ +def _root_.Lean.MVarId.generalizeHyp (mvarId : MVarId) (args : Array GeneralizeArg) (hyps : Array FVarId := #[]) + (fvarSubst : FVarSubst := {}) (transparency := TransparencyMode.instances) : MetaM (FVarSubst × Array FVarId × MVarId) := do + if hyps.isEmpty then + -- trivial case + return (fvarSubst, ← mvarId.generalize args transparency) + let args ← args.mapM fun arg => return { arg with expr := ← instantiateMVars arg.expr } + let hyps ← hyps.filterM fun h => do + let type ← instantiateMVars (← h.getType) + args.anyM fun arg => return (← withTransparency transparency <| kabstract type arg.expr).hasLooseBVars + let (reverted, mvarId) ← mvarId.revert hyps true + let (newVars, mvarId) ← mvarId.generalize args transparency + let (reintros, mvarId) ← mvarId.introNP reverted.size + let fvarSubst := Id.run do + let mut subst : FVarSubst := fvarSubst + for h in reverted, reintro in reintros do + subst := subst.insert h (mkFVar reintro) + pure subst + return (fvarSubst, newVars, mvarId) + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind.lean new file mode 100644 index 0000000000000000000000000000000000000000..230125ed198eaed257388c79e5c7c17bb4e0658f --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind.lean @@ -0,0 +1,89 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Attr +import Lean.Meta.Tactic.Grind.RevertAll +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Util +import Lean.Meta.Tactic.Grind.Cases +import Lean.Meta.Tactic.Grind.Injection +import Lean.Meta.Tactic.Grind.Core +import Lean.Meta.Tactic.Grind.Canon +import Lean.Meta.Tactic.Grind.MarkNestedProofs +import Lean.Meta.Tactic.Grind.Inv +import Lean.Meta.Tactic.Grind.Proof +import Lean.Meta.Tactic.Grind.Propagate +import Lean.Meta.Tactic.Grind.PP +import Lean.Meta.Tactic.Grind.Simp +import Lean.Meta.Tactic.Grind.Ctor +import Lean.Meta.Tactic.Grind.Parser +import Lean.Meta.Tactic.Grind.EMatchTheorem +import Lean.Meta.Tactic.Grind.EMatch +import Lean.Meta.Tactic.Grind.Main +import Lean.Meta.Tactic.Grind.CasesMatch +import Lean.Meta.Tactic.Grind.Arith +import Lean.Meta.Tactic.Grind.Ext +import Lean.Meta.Tactic.Grind.MatchCond +import Lean.Meta.Tactic.Grind.MatchDiscrOnly +import Lean.Meta.Tactic.Grind.Diseq +import Lean.Meta.Tactic.Grind.MBTC +import Lean.Meta.Tactic.Grind.Lookahead +import Lean.Meta.Tactic.Grind.LawfulEqCmp +import Lean.Meta.Tactic.Grind.ReflCmp + +namespace Lean + +/-! Trace options for `grind` users -/ +builtin_initialize registerTraceClass `grind +builtin_initialize registerTraceClass `grind.trace +builtin_initialize registerTraceClass `grind.assert +builtin_initialize registerTraceClass `grind.eqc +builtin_initialize registerTraceClass `grind.internalize +builtin_initialize registerTraceClass `grind.ematch +builtin_initialize registerTraceClass `grind.ematch.pattern +builtin_initialize registerTraceClass `grind.ematch.instance +builtin_initialize registerTraceClass `grind.ematch.instance.assignment +builtin_initialize registerTraceClass `grind.eqResolution +builtin_initialize registerTraceClass `grind.issues +builtin_initialize registerTraceClass `grind.simp +builtin_initialize registerTraceClass `grind.split +builtin_initialize registerTraceClass `grind.split.candidate +builtin_initialize registerTraceClass `grind.split.resolved +builtin_initialize registerTraceClass `grind.beta +builtin_initialize registerTraceClass `grind.mbtc +builtin_initialize registerTraceClass `grind.ext +builtin_initialize registerTraceClass `grind.ext.candidate +builtin_initialize registerTraceClass `grind.lookahead +builtin_initialize registerTraceClass `grind.lookahead.add (inherited := true) +builtin_initialize registerTraceClass `grind.lookahead.try (inherited := true) +builtin_initialize registerTraceClass `grind.lookahead.assert (inherited := true) + +/-! Trace options for `grind` developers -/ +builtin_initialize registerTraceClass `grind.debug +builtin_initialize registerTraceClass `grind.debug.proofs +builtin_initialize registerTraceClass `grind.debug.congr +builtin_initialize registerTraceClass `grind.debug.proof +builtin_initialize registerTraceClass `grind.debug.proj +builtin_initialize registerTraceClass `grind.debug.parent +builtin_initialize registerTraceClass `grind.debug.final +builtin_initialize registerTraceClass `grind.debug.forallPropagator +builtin_initialize registerTraceClass `grind.debug.split +builtin_initialize registerTraceClass `grind.debug.canon +builtin_initialize registerTraceClass `grind.debug.ematch.activate +builtin_initialize registerTraceClass `grind.debug.ematch.pattern +builtin_initialize registerTraceClass `grind.debug.beta +builtin_initialize registerTraceClass `grind.debug.internalize +builtin_initialize registerTraceClass `grind.debug.matchCond +builtin_initialize registerTraceClass `grind.debug.matchCond.lambda +builtin_initialize registerTraceClass `grind.debug.matchCond.proveFalse +builtin_initialize registerTraceClass `grind.debug.mbtc +builtin_initialize registerTraceClass `grind.debug.ematch +builtin_initialize registerTraceClass `grind.debug.proveEq +builtin_initialize registerTraceClass `grind.debug.pushNewFact +builtin_initialize registerTraceClass `grind.debug.appMap +builtin_initialize registerTraceClass `grind.debug.ext + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/AlphaShareCommon.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/AlphaShareCommon.lean new file mode 100644 index 0000000000000000000000000000000000000000..e6ec133e5f68a2627e72dfe044b452b9ffb3407a --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/AlphaShareCommon.lean @@ -0,0 +1,112 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.ExprPtr + +namespace Lean.Meta.Grind + +private def hashChild (e : Expr) : UInt64 := + match e with + | .bvar .. | .mvar .. | .const .. | .fvar .. | .sort .. | .lit .. => + hash e + | .app .. | .letE .. | .forallE .. | .lam .. | .mdata .. | .proj .. => + hashPtrExpr e + +private def alphaHash (e : Expr) : UInt64 := + match e with + | .bvar .. | .mvar .. | .const .. | .fvar .. | .sort .. | .lit .. => + hash e + | .app f a => mixHash (hashChild f) (hashChild a) + | .letE _ _ v b _ => mixHash (hashChild v) (hashChild b) + | .forallE _ d b _ | .lam _ d b _ => mixHash (hashChild d) (hashChild b) + | .mdata _ b => mixHash 13 (hashChild b) + | .proj n i b => mixHash (mixHash (hash n) (hash i)) (hashChild b) + +private def alphaEq (e₁ e₂ : Expr) : Bool := Id.run do + match e₁ with + | .bvar .. | .mvar .. | .const .. | .fvar .. | .sort .. | .lit .. => + e₁ == e₂ + | .app f₁ a₁ => + let .app f₂ a₂ := e₂ | false + isSameExpr f₁ f₂ && isSameExpr a₁ a₂ + | .letE _ _ v₁ b₁ _ => + let .letE _ _ v₂ b₂ _ := e₂ | false + isSameExpr v₁ v₂ && isSameExpr b₁ b₂ + | .forallE _ d₁ b₁ _ => + let .forallE _ d₂ b₂ _ := e₂ | false + isSameExpr d₁ d₂ && isSameExpr b₁ b₂ + | .lam _ d₁ b₁ _ => + let .lam _ d₂ b₂ _ := e₂ | false + isSameExpr d₁ d₂ && isSameExpr b₁ b₂ + | .mdata d₁ b₁ => + let .mdata d₂ b₂ := e₂ | false + return isSameExpr b₁ b₂ && d₁ == d₂ + | .proj n₁ i₁ b₁ => + let .proj n₂ i₂ b₂ := e₂ | false + n₁ == n₂ && i₁ == i₂ && isSameExpr b₁ b₂ + +structure AlphaKey where + expr : Expr + +instance : Hashable AlphaKey where + hash k := alphaHash k.expr + +instance : BEq AlphaKey where + beq k₁ k₂ := alphaEq k₁.expr k₂.expr + +structure AlphaShareCommon.State where + map : PHashMap ExprPtr Expr := {} + set : PHashSet AlphaKey := {} + +abbrev AlphaShareCommonM := StateM AlphaShareCommon.State + +private def save (e : Expr) (r : Expr) : AlphaShareCommonM Expr := do + if let some r := (← get).set.find? { expr := r } then + let r := r.expr + modify fun { set, map } => { + set + map := map.insert { expr := e } r + } + return r + else + modify fun { set, map } => { + set := set.insert { expr := r } + map := map.insert { expr := e } r |>.insert { expr := r } r + } + return r + +private abbrev visit (e : Expr) (k : AlphaShareCommonM Expr) : AlphaShareCommonM Expr := do + if let some r := (← get).map.find? { expr := e } then + return r + else + save e (← k) + +/-- Similar to `shareCommon`, but handles alpha-equivalence. -/ +def shareCommonAlpha (e : Expr) : AlphaShareCommonM Expr := + go e +where + go (e : Expr) : AlphaShareCommonM Expr := do + match e with + | .bvar .. | .mvar .. | .const .. | .fvar .. | .sort .. | .lit .. => + if let some r := (← get).set.find? { expr := e } then + return r.expr + else + modify fun { set, map } => { set := set.insert { expr := e }, map } + return e + | .app f a => + visit e (return mkApp (← go f) (← go a)) + | .letE n t v b nd => + visit e (return mkLet n t (← go v) (← go b) nd) + | .forallE n d b bi => + visit e (return mkForall n bi (← go d) (← go b)) + | .lam n d b bi => + visit e (return mkLambda n bi (← go d) (← go b)) + | .mdata d b => + visit e (return mkMData d (← go b)) + | .proj n i b => + visit e (return mkProj n i (← go b)) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Arith.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Arith.lean new file mode 100644 index 0000000000000000000000000000000000000000..53ddb17309bff1525a417aec8c327e7defdb1951 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Arith.lean @@ -0,0 +1,15 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Arith.Util +import Lean.Meta.Tactic.Grind.Arith.ProofUtil +import Lean.Meta.Tactic.Grind.Arith.Types +import Lean.Meta.Tactic.Grind.Arith.Main +import Lean.Meta.Tactic.Grind.Arith.Offset +import Lean.Meta.Tactic.Grind.Arith.Cutsat +import Lean.Meta.Tactic.Grind.Arith.CommRing +import Lean.Meta.Tactic.Grind.Arith.Linear +import Lean.Meta.Tactic.Grind.Arith.Simproc diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Attr.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Attr.lean new file mode 100644 index 0000000000000000000000000000000000000000..7a2a71ed3e5b2728d39fa0f25dedcc45c6060024 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Attr.lean @@ -0,0 +1,125 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.EMatchTheorem +import Lean.Meta.Tactic.Grind.Cases +import Lean.Meta.Tactic.Grind.ExtAttr + +namespace Lean.Meta.Grind + +inductive AttrKind where + | ematch (k : EMatchTheoremKind) + | cases (eager : Bool) + | intro + | infer + | ext + | symbol (prio : Nat) + +/-- Return theorem kind for `stx` of the form `Attr.grindThmMod` -/ +def getAttrKindCore (stx : Syntax) : CoreM AttrKind := do + match stx with + | `(Parser.Attr.grindMod|=) => return .ematch (.eqLhs false) + | `(Parser.Attr.grindMod|= gen) => return .ematch (.eqLhs true) + | `(Parser.Attr.grindMod|→) => return .ematch .fwd + | `(Parser.Attr.grindMod|←) => return .ematch (.bwd false) + | `(Parser.Attr.grindMod|← gen) => return .ematch (.bwd true) + | `(Parser.Attr.grindMod|=_) => return .ematch (.eqRhs false) + | `(Parser.Attr.grindMod|=_ gen) => return .ematch (.eqRhs true) + | `(Parser.Attr.grindMod|_=_) => return .ematch (.eqBoth false) + | `(Parser.Attr.grindMod|_=_ gen) => return .ematch (.eqBoth true) + | `(Parser.Attr.grindMod|←=) => return .ematch .eqBwd + | `(Parser.Attr.grindMod|=>) => return .ematch .leftRight + | `(Parser.Attr.grindMod|<=) => return .ematch .rightLeft + | `(Parser.Attr.grindMod|usr) => return .ematch .user + | `(Parser.Attr.grindMod|gen) => return .ematch (.default true) + | `(Parser.Attr.grindMod|cases) => return .cases false + | `(Parser.Attr.grindMod|cases eager) => return .cases true + | `(Parser.Attr.grindMod|intro) => return .intro + | `(Parser.Attr.grindMod|ext) => return .ext + | `(Parser.Attr.grindMod|symbol $prio:prio) => + let some prio := prio.raw.isNatLit? | throwErrorAt prio "priority expected" + return .symbol prio + | _ => throwError "unexpected `grind` theorem kind: `{stx}`" + +/-- Return theorem kind for `stx` of the form `(Attr.grindMod)?` -/ +def getAttrKindFromOpt (stx : Syntax) : CoreM AttrKind := do + if stx[1].isNone then + return .infer + else + getAttrKindCore stx[1][0] + +def throwInvalidUsrModifier : CoreM α := + throwError "the modifier `usr` is only relevant in parameters for `grind only`" + +/-- +Auxiliary function for registering `grind` and `grind?` attributes. +The `grind?` is an alias for `grind` which displays patterns using `logInfo`. +It is just a convenience for users. +-/ +private def registerGrindAttr (showInfo : Bool) : IO Unit := + registerBuiltinAttribute { + name := if showInfo then `grind? else `grind + descr := + let header := if showInfo then + "The `[grind?]` attribute is identical to the `[grind]` attribute, but displays inferred pattern information." + else + "The `[grind]` attribute is used to annotate declarations." + header ++ "\ + \ + When applied to an equational theorem, `[grind =]`, `[grind =_]`, or `[grind _=_]`\ + will mark the theorem for use in heuristic instantiations by the `grind` tactic, + using respectively the left-hand side, the right-hand side, or both sides of the theorem.\ + When applied to a function, `[grind =]` automatically annotates the equational theorems associated with that function.\ + When applied to a theorem `[grind ←]` will instantiate the theorem whenever it encounters the conclusion of the theorem + (that is, it will use the theorem for backwards reasoning).\ + When applied to a theorem `[grind →]` will instantiate the theorem whenever it encounters sufficiently many of the propositional hypotheses + (that is, it will use the theorem for forwards reasoning).\ + \ + The attribute `[grind]` by itself will effectively try `[grind ←]` (if the conclusion is sufficient for instantiation) and then `[grind →]`.\ + \ + The `grind` tactic utilizes annotated theorems to add instances of matching patterns into the local context during proof search.\ + For example, if a theorem `@[grind =] theorem foo_idempotent : foo (foo x) = foo x` is annotated,\ + `grind` will add an instance of this theorem to the local context whenever it encounters the pattern `foo (foo x)`." + applicationTime := .afterCompilation + add := fun declName stx attrKind => MetaM.run' do + match (← getAttrKindFromOpt stx) with + | .ematch .user => throwInvalidUsrModifier + | .ematch k => addEMatchAttr declName attrKind k (← getGlobalSymbolPriorities) (showInfo := showInfo) + | .cases eager => addCasesAttr declName eager attrKind + | .intro => + if let some info ← isCasesAttrPredicateCandidate? declName false then + for ctor in info.ctors do + addEMatchAttr ctor attrKind (.default false) (← getGlobalSymbolPriorities) (showInfo := showInfo) + else + throwError "invalid `[grind intro]`, `{declName}` is not an inductive predicate" + | .ext => addExtAttr declName attrKind + | .infer => + if let some declName ← isCasesAttrCandidate? declName false then + addCasesAttr declName false attrKind + if let some info ← isInductivePredicate? declName then + -- If it is an inductive predicate, + -- we also add the constructors (intro rules) as E-matching rules + for ctor in info.ctors do + addEMatchAttr ctor attrKind (.default false) (← getGlobalSymbolPriorities) (showInfo := showInfo) + else + addEMatchAttr declName attrKind (.default false) (← getGlobalSymbolPriorities) (showInfo := showInfo) + | .symbol prio => addSymbolPriorityAttr declName attrKind prio + erase := fun declName => MetaM.run' do + if showInfo then + throwError "`[grind?]` is a helper attribute for displaying inferred patterns, if you want to remove the attribute, consider using `[grind]` instead" + if (← isCasesAttrCandidate declName false) then + eraseCasesAttr declName + else if (← isExtTheorem declName) then + eraseExtAttr declName + else + eraseEMatchAttr declName + } + +builtin_initialize + registerGrindAttr true + registerGrindAttr false + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Beta.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Beta.lean new file mode 100644 index 0000000000000000000000000000000000000000..3dfad9d1c3b991a7c899c915dcbf3a3af58e63db --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Beta.lean @@ -0,0 +1,77 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types + +namespace Lean.Meta.Grind + +/-- Returns all lambda expressions in the equivalence class with root `root`. -/ +def getEqcLambdas (root : ENode) : GoalM (Array Expr) := do + unless root.hasLambdas do return #[] + foldEqc root.self (init := #[]) fun n lams => + if n.self.isLambda then return lams.push n.self else return lams + +/-- +Returns the root of the functions in the equivalence class containing `e`. +That is, if `f a` is in `root`s equivalence class, results contains the root of `f`. +-/ +def getFnRoots (e : Expr) : GoalM (Array Expr) := do + foldEqc e (init := #[]) fun n fns => do + let fn := n.self.getAppFn + let fnRoot := (← getRoot? fn).getD fn + if Option.isNone <| fns.find? (isSameExpr · fnRoot) then + return fns.push fnRoot + else + return fns + +/-- +For each `lam` in `lams` s.t. `lam` and `f` are in the same equivalence class, +propagate `f args = lam args`. +-/ +def propagateBetaEqs (lams : Array Expr) (f : Expr) (args : Array Expr) : GoalM Unit := do + if args.isEmpty then return () + for lam in lams do + let rhs := lam.beta args + unless rhs.isLambda do + let mut gen := Nat.max (← getGeneration lam) (← getGeneration f) + let lhs := mkAppN f args + if (← hasSameType f lam) then + let mut h ← mkEqProof f lam + for arg in args do + gen := Nat.max gen (← getGeneration arg) + h ← mkCongrFun h arg + let eq ← mkEq lhs rhs + trace_goal[grind.beta] "{eq}, using {lam}" + addNewRawFact h eq (gen+1) (.beta lam) + +private def isPropagateBetaTarget (e : Expr) : GoalM Bool := do + let .app f _ := e | return false + go f +where + go (f : Expr) : GoalM Bool := do + if let some root ← getRootENode? f then + return root.hasLambdas + let .app f _ := f | return false + go f + +/-- +Applies beta-reduction for lambdas in `f`s equivalence class. +We use this function while internalizing new applications. +-/ +def propagateBetaForNewApp (e : Expr) : GoalM Unit := do + unless (← isPropagateBetaTarget e) do return () + let mut e := e + let mut args := #[] + repeat + unless args.isEmpty do + if let some root ← getRootENode? e then + if root.hasLambdas then + propagateBetaEqs (← getEqcLambdas root) e args.reverse + let .app f arg := e | return () + e := f + args := args.push arg + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Canon.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Canon.lean new file mode 100644 index 0000000000000000000000000000000000000000..c079a52fa22764a450d02165cfb350c60cd8a7f7 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Canon.lean @@ -0,0 +1,211 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Lean.Meta.Basic +import Lean.Meta.FunInfo +import Lean.Util.FVarSubset +import Lean.Util.PtrSet +import Lean.Util.FVarSubset +import Lean.Meta.Tactic.Grind.Types + +namespace Lean.Meta.Grind +namespace Canon + +/-! +A canonicalizer module for the `grind` tactic. The canonicalizer defined in `Meta/Canonicalizer.lean` is +not suitable for the `grind` tactic. It was designed for tactics such as `omega`, where the goal is +to detect when two structurally different atoms are definitionally equal. + +The `grind` tactic, on the other hand, uses congruence closure. Moreover, types, type formers, proofs, and instances +are considered supporting elements and are not factored into congruence detection. + +This module minimizes the number of `isDefEq` checks by comparing two terms `a` and `b` only if they are instances, +types, or type formers and are the `i`-th arguments of two different `f`-applications. This approach is +sufficient for the congruence closure procedure used by the `grind` tactic. + +To further optimize `isDefEq` checks, instances are compared using `TransparencyMode.instances`, which reduces +the number of constants that need to be unfolded. If diagnostics are enabled, instances are compared using +the default transparency mode too for sanity checking, and discrepancies are reported. +Types and type formers are always checked using default transparency. + +Remark: +The canonicalizer minimizes issues with non-canonical instances and structurally different but definitionally equal types, +but it does not solve all problems. For example, consider a situation where we have `(a : BitVec n)` +and `(b : BitVec m)`, along with instances `inst1 n : Add (BitVec n)` and `inst2 m : Add (BitVec m)` where `inst1` +and `inst2` are structurally different. Now consider the terms `a + a` and `b + b`. After canonicalization, the two +additions will still use structurally different (and definitionally different) instances: `inst1 n` and `inst2 m`. +Furthermore, `grind` will not be able to infer that `a + a ≍ b + b` even if we add the assumptions `n = m` and `a ≍ b`. +-/ + +@[inline] private def get' : GoalM State := + return (← get).canon + +@[inline] private def modify' (f : State → State) : GoalM Unit := + modify fun s => { s with canon := f s.canon } + +/-- +Helper function for `canonElemCore`. It tries `isDefEq a b` with default transparency, but using +at most `canonHeartbeats` heartbeats. It reports an issue if the threshold is reached. +Remark: `parent` is use only to report an issue. +-/ +private def isDefEqBounded (a b : Expr) (parent : Expr) : GoalM Bool := do + withCurrHeartbeats do + let curr := (← getConfig).canonHeartbeats + tryCatchRuntimeEx + (withTheReader Core.Context (fun ctx => { ctx with maxHeartbeats := curr*1000 }) do + withDefault <| isDefEq a b) + fun ex => do + if ex.isRuntime then + reportIssue! "failed to show that{indentExpr a}\nis definitionally equal to{indentExpr b}\nwhile canonicalizing{indentExpr parent}\nusing `{curr}*1000` heartbeats, `(canonHeartbeats := {curr})`" + return false + else + throw ex + +/-- +Helper function for canonicalizing `e` occurring as the `i`th argument of an `f`-application. +If `useIsDefEqBounded` is `true`, we try `isDefEqBounded` before returning false +-/ +def canonElemCore (parent : Expr) (f : Expr) (i : Nat) (e : Expr) (useIsDefEqBounded : Bool) : GoalM Expr := do + let s ← get' + if let some c := s.canon.find? e then + return c + let key := (f, i) + let eType ← inferType e + let cs := s.argMap.find? key |>.getD [] + for (c, cType) in cs do + /- + We first check the types + The following checks are a performance bottleneck. + For example, in the test `grind_ite.lean`, there are many checks of the form: + ``` + w_4 ∈ assign.insert v true → Prop =?= w_1 ∈ assign.insert v false → Prop + ``` + where `grind` unfolds the definition of `DHashMap.insert` and `TreeMap.insert`. + -/ + if (← withDefault <| isDefEq eType cType) then + if (← isDefEq e c) then + -- We used to check `c.fvarsSubset e` because it is not + -- in general safe to replace `e` with `c` if `c` has more free variables than `e`. + -- However, we don't revert previously canonicalized elements in the `grind` tactic. + -- Moreover, we store the canonicalizer state in the `Goal` because we case-split + -- and different locals are added in different branches. + modify' fun s => { s with canon := s.canon.insert e c } + trace_goal[grind.debugn.canon] "found {e} ===> {c}" + return c + if useIsDefEqBounded then + -- If `e` and `c` are not types, we use `isDefEqBounded` + if (← isDefEqBounded e c parent) then + modify' fun s => { s with canon := s.canon.insert e c } + trace_goal[grind.debugn.canon] "found using `isDefEqBounded`: {e} ===> {c}" + return c + trace_goal[grind.debug.canon] "({f}, {i}) ↦ {e}" + modify' fun s => { s with canon := s.canon.insert e e, argMap := s.argMap.insert key ((e, eType)::cs) } + return e + +abbrev canonType (parent f : Expr) (i : Nat) (e : Expr) := withDefault <| canonElemCore parent f i e (useIsDefEqBounded := false) +abbrev canonInst (parent f : Expr) (i : Nat) (e : Expr) := withReducibleAndInstances <| canonElemCore parent f i e (useIsDefEqBounded := true) +abbrev canonImplicit (parent f : Expr) (i : Nat) (e : Expr) := withReducible <| canonElemCore parent f i e (useIsDefEqBounded := true) + +/-- +Return type for the `shouldCanon` function. +-/ +private inductive ShouldCanonResult where + | /- Nested types (and type formers) are canonicalized. -/ + canonType + | /- Nested instances are canonicalized. -/ + canonInst + | /- Implicit argument that is not an instance nor a type. -/ + canonImplicit + | /- + Term is not a proof, type (former), nor an instance. + Thus, it must be recursively visited by the canonizer. + -/ + visit + deriving Inhabited + +instance : Repr ShouldCanonResult where + reprPrec r _ := match r with + | .canonType => "canonType" + | .canonInst => "canonInst" + | .canonImplicit => "canonImplicit" + | .visit => "visit" + +/-- +See comments at `ShouldCanonResult`. +-/ +def shouldCanon (pinfos : Array ParamInfo) (i : Nat) (arg : Expr) : MetaM ShouldCanonResult := do + if h : i < pinfos.size then + let pinfo := pinfos[i] + if pinfo.isInstImplicit then + return .canonInst + else if pinfo.isProp then + return .visit + else if pinfo.isImplicit then + if (← isTypeFormer arg) then + return .canonType + else + return .canonImplicit + if (← isProp arg) then + return .visit + else if (← isTypeFormer arg) then + return .canonType + else + return .visit + +unsafe def canonImpl (e : Expr) : GoalM Expr := do + visit e |>.run' mkPtrMap +where + visit (e : Expr) : StateRefT (PtrMap Expr Expr) GoalM Expr := do + unless e.isApp || e.isForall do return e + -- Check whether it is cached + if let some r := (← get).find? e then + return r + let e' ← match e with + | .app .. => e.withApp fun f args => do + if f.isConstOf ``Lean.Grind.nestedProof && args.size == 2 then + let prop := args[0]! + let prop' ← visit prop + if let some r := (← get').proofCanon.find? prop' then + pure r + else + let e' := if ptrEq prop prop' then e else mkAppN f (args.set! 0 prop') + modify' fun s => { s with proofCanon := s.proofCanon.insert prop' e' } + pure e' + else + let pinfos := (← getFunInfo f).paramInfo + let mut modified := false + let mut args := args.toVector + for h : i in [:args.size] do + let arg := args[i] + trace_goal[grind.debug.canon] "[{repr (← shouldCanon pinfos i arg)}]: {arg} : {← inferType arg}" + let arg' ← match (← shouldCanon pinfos i arg) with + | .canonType => canonType e f i arg + | .canonInst => canonInst e f i arg + | .canonImplicit => canonImplicit e f i (← visit arg) + | .visit => visit arg + unless ptrEq arg arg' do + args := args.set i arg' + modified := true + pure <| if modified then mkAppN f args.toArray else e + | .forallE _ d b _ => + -- Recall that we have `ForallProp.lean`. + let d' ← visit d + -- Remark: users may not want to convert `p → q` into `¬p ∨ q` + let b' ← if b.hasLooseBVars then pure b else visit b + pure <| e.updateForallE! d' b' + | _ => unreachable! + modify fun s => s.insert e e' + return e' + +end Canon + +/-- Canonicalizes nested types, type formers, and instances in `e`. -/ +def canon (e : Expr) : GoalM Expr := do profileitM Exception "grind canon" (← getOptions) do + trace_goal[grind.debug.canon] "{e}" + unsafe Canon.canonImpl e + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Cases.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Cases.lean new file mode 100644 index 0000000000000000000000000000000000000000..a28fa48c2df6f5081760bdaa0cc5b4a48e13dc38 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Cases.lean @@ -0,0 +1,191 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Cases + +namespace Lean.Meta.Grind + +/-- Types that `grind` will case-split on. -/ +structure CasesTypes where + casesMap : PHashMap Name Bool := {} + deriving Inhabited + +structure CasesEntry where + declName : Name + eager : Bool + deriving Inhabited + +/-- +`grind` always case-splits on the following types. Even when using `grind only`. +The goal is to reduce noise in the tactic generated by `grind?` +-/ +private def builtinEagerCases : NameSet := + RBTree.ofList [``And, ``Exists, ``True, ``False, ``Unit, ``Empty] + +/-- +Returns `true` if `declName` is the name of inductive type/predicate that +even `grind only` case splits on. +Remark: we added support for them to reduce the noise in the tactic generated by +`grind?` +-/ +def isBuiltinEagerCases (declName : Name) : Bool := + builtinEagerCases.contains declName + +/-- Returns `true` if `s` contains a `declName`. -/ +def CasesTypes.contains (s : CasesTypes) (declName : Name) : Bool := + s.casesMap.contains declName + +/-- Removes the given declaration from `s`. -/ +def CasesTypes.erase (s : CasesTypes) (declName : Name) : CasesTypes := + { s with casesMap := s.casesMap.erase declName } + +def CasesTypes.insert (s : CasesTypes) (declName : Name) (eager : Bool) : CasesTypes := + { s with casesMap := s.casesMap.insert declName eager } + +def CasesTypes.find? (s : CasesTypes) (declName : Name) : Option Bool := + s.casesMap.find? declName + +def CasesTypes.isEagerSplit (s : CasesTypes) (declName : Name) : Bool := + (s.casesMap.find? declName |>.getD false) || isBuiltinEagerCases declName + +def CasesTypes.isSplit (s : CasesTypes) (declName : Name) : Bool := + (s.casesMap.find? declName |>.isSome) || isBuiltinEagerCases declName + +builtin_initialize casesExt : SimpleScopedEnvExtension CasesEntry CasesTypes ← + registerSimpleScopedEnvExtension { + initial := {} + addEntry := fun s {declName, eager} => s.insert declName eager + } + +def resetCasesExt : CoreM Unit := do + modifyEnv fun env => casesExt.modifyState env fun _ => {} + +def getCasesTypes : CoreM CasesTypes := + return casesExt.getState (← getEnv) + +/-- Returns `true` is `declName` is a builtin split or has been tagged with `[grind]` attribute. -/ +def isSplit (declName : Name) : CoreM Bool := do + return (← getCasesTypes).isSplit declName + +partial def isCasesAttrCandidate? (declName : Name) (eager : Bool) : CoreM (Option Name) := do + match (← getConstInfo declName) with + | .inductInfo info => if !info.isRec || !eager then return some declName else return none + | _ => return none + +def isCasesAttrCandidate (declName : Name) (eager : Bool) : CoreM Bool := do + return (← isCasesAttrCandidate? declName eager).isSome + +def isCasesAttrPredicateCandidate? (declName : Name) (eager : Bool) : MetaM (Option InductiveVal) := do + let some declName ← isCasesAttrCandidate? declName eager | return none + isInductivePredicate? declName + +def validateCasesAttr (declName : Name) (eager : Bool) : CoreM Unit := do + unless (← isCasesAttrCandidate declName eager) do + if eager then + throwError "invalid `[grind cases eager]`, `{declName}` is not a non-recursive inductive datatype or an alias for one" + else + throwError "invalid `[grind cases]`, `{declName}` is not an inductive datatype or an alias for one" + +def addCasesAttr (declName : Name) (eager : Bool) (attrKind : AttributeKind) : CoreM Unit := do + validateCasesAttr declName eager + casesExt.add { declName, eager } attrKind + +def CasesTypes.eraseDecl (s : CasesTypes) (declName : Name) : CoreM CasesTypes := do + if s.contains declName then + return s.erase declName + else + throwError "`{declName}` is not marked with the `[grind]` attribute" + +def ensureNotBuiltinCases (declName : Name) : CoreM Unit := do + if isBuiltinEagerCases declName then + throwError "`{declName}` is marked as a built-in case-split for `grind` and cannot be erased" + +def eraseCasesAttr (declName : Name) : CoreM Unit := do + ensureNotBuiltinCases declName + let s := casesExt.getState (← getEnv) + let s ← s.eraseDecl declName + modifyEnv fun env => casesExt.modifyState env fun _ => s + +/-- +We say a free variable is "simple" to be processed by the cases tactic IF: +- It is the latest and consequently there are no forward dependencies, OR +- It is not a proposion. +-/ +private def isSimpleFVar (e : Expr) : MetaM Bool := do + let .fvar fvarId := e | return false + let decl ← fvarId.getDecl + if decl.index == (← getLCtx).numIndices - 1 then + -- It is the latest free variable, so there are no forward dependencies + return true + else + -- It is pointless to add an auxiliary equality if `e`s type is a proposition + isProp decl.type + +/-- +The `grind` tactic includes an auxiliary `cases` tactic that is not intended for direct use by users. +This method implements it. +This tactic is automatically applied when introducing local declarations with a type tagged with `[grind_cases]`. +It is also used for "case-splitting" on terms during the search. + +It differs from the user-facing Lean `cases` tactic in the following ways: + +- It avoids unnecessary `revert` and `intro` operations. + +- It does not introduce new local declarations for each minor premise. Instead, the `grind` tactic preprocessor is responsible for introducing them. + +- If the major premise type is an indexed family, auxiliary declarations and (heterogeneous) equalities are introduced. + However, these equalities are not resolved using `unifyEqs`. Instead, the `grind` tactic employs union-find and + congruence closure to process these auxiliary equalities. This approach avoids applying substitution to propositions + that have already been internalized by `grind`. +-/ +def cases (mvarId : MVarId) (e : Expr) : MetaM (List MVarId) := mvarId.withContext do + let tag ← mvarId.getTag + let type ← whnf (← inferType e) + let .const declName _ := type.getAppFn | throwInductiveExpected type + let .inductInfo _ ← getConstInfo declName | throwInductiveExpected type + let recursorInfo ← mkRecursorInfo (mkCasesOnName declName) + let k (mvarId : MVarId) (fvarId : FVarId) (indices : Array FVarId) : MetaM (List MVarId) := do + let indicesExpr := indices.map mkFVar + let recursor ← mkRecursorAppPrefix mvarId `grind.cases fvarId recursorInfo indicesExpr + let lctx ← getLCtx + let lctx := lctx.setKind fvarId .implDetail + let lctx := indices.foldl (init := lctx) fun lctx fvarId => lctx.setKind fvarId .implDetail + let localInsts ← getLocalInstances + let mut recursor := mkApp (mkAppN recursor indicesExpr) (mkFVar fvarId) + let mut recursorType ← inferType recursor + let mut mvarIdsNew := #[] + let mut idx := 1 + for _ in [:recursorInfo.numMinors] do + let .forallE _ targetNew recursorTypeNew _ ← whnf recursorType + | throwTacticEx `grind.cases mvarId "unexpected recursor type" + recursorType := recursorTypeNew + let tagNew := if recursorInfo.numMinors > 1 then Name.num tag idx else tag + let mvar ← mkFreshExprMVarAt lctx localInsts targetNew .syntheticOpaque tagNew + recursor := mkApp recursor mvar + mvarIdsNew := mvarIdsNew.push mvar.mvarId! + idx := idx + 1 + mvarId.assign recursor + return mvarIdsNew.toList + if recursorInfo.numIndices > 0 then + let s ← generalizeIndices' mvarId e + s.mvarId.withContext do + k s.mvarId s.fvarId s.indicesFVarIds + else if (← isSimpleFVar e) then + -- We don't need to revert anything. + k mvarId e.fvarId! #[] + else + let mvarId ← if (← isProof e) then + mvarId.assert (← mkFreshUserName `x) type e + else + mvarId.assertExt (← mkFreshUserName `x) type e + let (fvarId, mvarId) ← mvarId.intro1 + mvarId.withContext do k mvarId fvarId #[] + +where + throwInductiveExpected {α} (type : Expr) : MetaM α := do + throwTacticEx `grind.cases mvarId m!"(non-recursive) inductive type expected at {e}{indentExpr type}" + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/CasesMatch.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/CasesMatch.lean new file mode 100644 index 0000000000000000000000000000000000000000..45c9644af2aeadbf97b108fa2cd3879785d9095e --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/CasesMatch.lean @@ -0,0 +1,82 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Cases +import Lean.Meta.Match.MatcherApp +import Lean.Meta.Tactic.Grind.MatchCond + +namespace Lean.Meta.Grind + +/-- Returns `true` if `e` is of the form `∀ ..., _ = _ ... -> False` -/ +def isMatchCondCandidate (e : Expr) : Bool := + -- TODO: see comment at the beginning of `MatchCond.lean`. + Simp.isEqnThmHypothesis e + +/-- +Given a splitter alternative, annotate the terms that are `match`-expression +conditions corresponding to overlapping patterns. +-/ +private def addMatchCondsToAlt (alt : Expr) : Expr := Id.run do + let .forallE _ d b _ := alt + | return alt + let d := if isMatchCondCandidate d then markAsPreMatchCond d else d + return alt.updateForallE! d (addMatchCondsToAlt b) + +/-- +Annotates the `match`-expression conditions in the alternatives in the given +`match` splitter type. +-/ +private def addMatchCondsToSplitter (splitterType : Expr) (numAlts : Nat) : Expr := Id.run do + if numAlts == 0 then return splitterType + let .forallE _ alt b _ := splitterType + | return splitterType + return splitterType.updateForallE! (addMatchCondsToAlt alt) (addMatchCondsToSplitter b (numAlts-1)) + +def casesMatch (mvarId : MVarId) (e : Expr) : MetaM (List MVarId) := mvarId.withContext do + let some app ← matchMatcherApp? e + | throwTacticEx `grind.casesMatch mvarId m!"`match`-expression expected{indentExpr e}" + let (motive, eqRefls) ← mkMotiveAndRefls app + let target ← mvarId.getType + let mut us := app.matcherLevels + if let some i := app.uElimPos? then + us := us.set! i (← getLevel target) + let splitterName := (← Match.getEquationsFor app.matcherName).splitterName + let splitterApp := mkConst splitterName us.toList + let splitterApp := mkAppN splitterApp app.params + let splitterApp := mkApp splitterApp motive + let splitterApp := mkAppN splitterApp app.discrs + let numAlts := app.alts.size + let splitterType ← inferType splitterApp + let splitterType := addMatchCondsToSplitter splitterType app.alts.size + let (mvars, _, _) ← forallMetaBoundedTelescope splitterType numAlts (kind := .syntheticOpaque) + let splitterApp := mkAppN splitterApp mvars + let val := mkAppN splitterApp eqRefls + mvarId.assign val + updateTags mvars + return mvars.toList.map (·.mvarId!) +where + mkMotiveAndRefls (app : MatcherApp) : MetaM (Expr × Array Expr) := do + let dummy := mkSort 0 + let aux := mkApp (mkAppN e.getAppFn app.params) dummy + forallBoundedTelescope (← inferType aux) app.discrs.size fun xs _ => do + withNewEqs app.discrs xs fun eqs eqRefls => do + let type ← mvarId.getType + let type ← mkForallFVars eqs type + let motive ← mkLambdaFVars xs type + return (motive, eqRefls) + + updateTags (mvars : Array Expr) : MetaM Unit := do + let tag ← mvarId.getTag + if mvars.size == 1 then + mvars[0]!.mvarId!.setTag tag + else + let mut idx := 1 + for mvar in mvars do + mvar.mvarId!.setTag (Name.num tag idx) + idx := idx + 1 + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Core.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Core.lean new file mode 100644 index 0000000000000000000000000000000000000000..cac7c6bb298b35a484d6cbba38f31cc3d9c1fa5f --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Core.lean @@ -0,0 +1,431 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Lean.Meta.LitValues +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Inv +import Lean.Meta.Tactic.Grind.PP +import Lean.Meta.Tactic.Grind.Ctor +import Lean.Meta.Tactic.Grind.Util +import Lean.Meta.Tactic.Grind.Beta +import Lean.Meta.Tactic.Grind.Internalize +import Lean.Meta.Tactic.Grind.Simp + +namespace Lean.Meta.Grind + +/-- +The fields `target?` and `proof?` in `e`'s `ENode` are encoding a transitivity proof +from `e` to the root of the equivalence class +This method "inverts" the proof, and makes it to go from the root of the equivalence class to `e`. + +We use this method when merging two equivalence classes. +-/ +private partial def invertTrans (e : Expr) : GoalM Unit := do + go e false none none +where + go (e : Expr) (flippedNew : Bool) (targetNew? : Option Expr) (proofNew? : Option Expr) : GoalM Unit := do + let node ← getENode e + if let some target := node.target? then + go target (!node.flipped) (some e) node.proof? + setENode e { node with + target? := targetNew? + flipped := flippedNew + proof? := proofNew? + } + +/-- +Returns `true` if the parent is relevant for congruence closure. +-/ +private def isCongrRelevant (parent : Expr) : Bool := + parent.isApp || parent.isArrow + +/-- +Removes `root` parents from the congruence table. +This is an auxiliary function performed while merging equivalence classes. +-/ +private def removeParents (root : Expr) : GoalM ParentSet := do + let parents ← getParents root + for parent in parents do + -- Recall that we may have `Expr.forallE` in `parents` because of `ForallProp.lean` + if (← pure (isCongrRelevant parent) <&&> isCongrRoot parent) then + trace_goal[grind.debug.parent] "remove: {parent}" + modify fun s => { s with congrTable := s.congrTable.erase { e := parent } } + return parents + +/-- +Reinserts parents into the congruence table and detect new equalities. +This is an auxiliary function performed while merging equivalence classes. +-/ +private def reinsertParents (parents : ParentSet) : GoalM Unit := do + for parent in parents do + if (← pure (isCongrRelevant parent) <&&> isCongrRoot parent) then + trace_goal[grind.debug.parent] "reinsert: {parent}" + addCongrTable parent + +/-- Closes the goal when `True` and `False` are in the same equivalence class. -/ +private def closeGoalWithTrueEqFalse : GoalM Unit := do + let mvarId := (← get).mvarId + unless (← mvarId.isAssigned) do + let trueEqFalse ← mkEqFalseProof (← getTrueExpr) + let falseProof := mkApp4 (mkConst ``Eq.mp [levelZero]) (← getTrueExpr) (← getFalseExpr) trueEqFalse (mkConst ``True.intro) + closeGoal falseProof + +/-- Closes the goal when `lhs` and `rhs` are both literal values and belong to the same equivalence class. -/ +private def closeGoalWithValuesEq (lhs rhs : Expr) : GoalM Unit := do + let p ← mkEq lhs rhs + let hp ← mkEqProof lhs rhs + let d ← mkDecide p + let pEqFalse := mkApp3 (mkConst ``eq_false_of_decide) p d.appArg! (mkApp2 (mkConst ``Eq.refl [1]) (mkConst ``Bool) (mkConst ``false)) + let falseProof := mkApp4 (mkConst ``Eq.mp [levelZero]) p (← getFalseExpr) pEqFalse hp + closeGoal falseProof + +/-- +Updates the modification time to `gmt` for the parents of `root`. +The modification time is used to decide which terms are considered during e-matching. +-/ +private partial def updateMT (root : Expr) : GoalM Unit := do + let gmt := (← get).ematch.gmt + for parent in (← getParents root) do + let node ← getENode parent + if node.mt < gmt then + setENode parent { node with mt := gmt } + updateMT parent + +/-- +Equalities or disequalities to be propagated to a theory solver **after** +two equivalence classes have been merged. + +Some solvers (e.g. `cutsat`) require the core data structures to satisfy +their invariants. During the merge operations some of these invariants do not hold. +Thus, we first *record* the facts that must be propagated in a `PendingTheoryPropagation` value, +complete the merge, and only then perform the propagation. + +We now use this workflow for *all* theory solvers, even when a particular +solver does not rely on these invariants. This keeps the core +solver-agnostic and lets us modify solvers without further adjustments. +-/ +inductive PendingTheoryPropagation where + | /-- Nothing to propagate. -/ + none + | /-- Propagate the equality `lhs = rhs`. -/ + eq (lhs rhs : Expr) + | /-- Propagate the disequalities in `ps`. -/ + diseqs (ps : ParentSet) + +/-- +Helper function for combining `ENode.offset?` fields and detecting what needs +to be propagated to the offset constraint module. +-/ +private def checkOffsetEq (rhsRoot lhsRoot : ENode) : GoalM PendingTheoryPropagation := do + match lhsRoot.offset? with + | some lhsOffset => + if let some rhsOffset := rhsRoot.offset? then + return .eq lhsOffset rhsOffset + else + -- We have to retrieve the node because other fields have been updated + let rhsRoot ← getENode rhsRoot.self + setENode rhsRoot.self { rhsRoot with offset? := lhsOffset } + return .none + | none => return .none + +def propagateOffset : PendingTheoryPropagation → GoalM Unit + | .eq lhs rhs => Arith.Offset.processNewEq lhs rhs + | _ => return () + +/-- +Helper function for combining `ENode.cutsat?` fields and detecting what needs +to be propagated to the cutsat module. +-/ +private def checkCutsatEq (rhsRoot lhsRoot : ENode) : GoalM PendingTheoryPropagation := do + match lhsRoot.cutsat? with + | some lhsCutsat => + if let some rhsCutsat := rhsRoot.cutsat? then + return .eq lhsCutsat rhsCutsat + else + -- We have to retrieve the node because other fields have been updated + let rhsRoot ← getENode rhsRoot.self + setENode rhsRoot.self { rhsRoot with cutsat? := lhsCutsat } + return .diseqs (← getParents rhsRoot.self) + | none => + if rhsRoot.cutsat?.isSome then + return .diseqs (← getParents lhsRoot.self) + else + return .none + +def propagateCutsat : PendingTheoryPropagation → GoalM Unit + | .eq lhs rhs => Arith.Cutsat.processNewEq lhs rhs + | .diseqs ps => propagateCutsatDiseqs ps + | .none => return () + +/-- +Helper function for combining `ENode.ring?` fields and detecting what needs to be +propagated to the commutative ring module. +-/ +private def checkCommRingEq (rhsRoot lhsRoot : ENode) : GoalM PendingTheoryPropagation := do + match lhsRoot.ring? with + | some lhsRing => + if let some rhsRing := rhsRoot.ring? then + return .eq lhsRing rhsRing + else + -- We have to retrieve the node because other fields have been updated + let rhsRoot ← getENode rhsRoot.self + setENode rhsRoot.self { rhsRoot with ring? := lhsRing } + return .diseqs (← getParents rhsRoot.self) + | none => + if rhsRoot.ring?.isSome then + return .diseqs (← getParents lhsRoot.self) + else + return .none + +def propagateCommRing : PendingTheoryPropagation → GoalM Unit + | .eq lhs rhs => Arith.CommRing.processNewEq lhs rhs + | .diseqs ps => propagateCommRingDiseqs ps + | _ => return () + +/-- +Helper function for combining `ENode.linarith?` fields and detecting what needs to be +propagated to the linarith module. +-/ +private def checkLinarithEq (rhsRoot lhsRoot : ENode) : GoalM PendingTheoryPropagation := do + match lhsRoot.linarith? with + | some lhs => + if let some rhs := rhsRoot.linarith? then + return .eq lhs rhs + else + -- We have to retrieve the node because other fields have been updated + let rhsRoot ← getENode rhsRoot.self + setENode rhsRoot.self { rhsRoot with linarith? := lhs } + return .diseqs (← getParents rhsRoot.self) + | none => + if rhsRoot.linarith?.isSome then + return .diseqs (← getParents lhsRoot.self) + else + return .none + +def propagateLinarith : PendingTheoryPropagation → GoalM Unit + | .eq lhs rhs => Arith.Linear.processNewEq lhs rhs + | .diseqs ps => propagateLinarithDiseqs ps + | _ => return () + +/-- +Tries to apply beta-reductiong using the parent applications of the functions in `fns` with +the lambda expressions in `lams`. +-/ +def propagateBeta (lams : Array Expr) (fns : Array Expr) : GoalM Unit := do + if lams.isEmpty then return () + let lamRoot ← getRoot lams.back! + trace_goal[grind.debug.beta] "fns: {fns}, lams: {lams}" + for fn in fns do + trace_goal[grind.debug.beta] "fn: {fn}, parents: {(← getParents fn).toArray}" + for parent in (← getParents fn) do + let mut args := #[] + let mut curr := parent + trace_goal[grind.debug.beta] "parent: {parent}" + repeat + trace_goal[grind.debug.beta] "curr: {curr}" + if (← isEqv curr lamRoot) then + propagateBetaEqs lams curr args.reverse + let .app f arg := curr + | break + -- Remark: recall that we do not eagerly internalize partial applications. + internalize curr (← getGeneration parent) + args := args.push arg + curr := f + +private partial def addEqStep (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit := do + let lhsNode ← getENode lhs + let rhsNode ← getENode rhs + if isSameExpr lhsNode.root rhsNode.root then + -- `lhs` and `rhs` are already in the same equivalence class. + trace_goal[grind.debug] "{← ppENodeRef lhs} and {← ppENodeRef rhs} are already in the same equivalence class" + return () + trace_goal[grind.eqc] "{← if isHEq then mkHEq lhs rhs else mkEq lhs rhs}" + let lhsRoot ← getENode lhsNode.root + let rhsRoot ← getENode rhsNode.root + let mut valueInconsistency := false + let mut trueEqFalse := false + if lhsRoot.interpreted && rhsRoot.interpreted then + if lhsNode.root.isTrue || rhsNode.root.isTrue then + markAsInconsistent + trueEqFalse := true + else + valueInconsistency := true + if (lhsRoot.interpreted && !rhsRoot.interpreted) + || (lhsRoot.ctor && !rhsRoot.ctor) + || (lhsRoot.size > rhsRoot.size && !rhsRoot.interpreted && !rhsRoot.ctor) then + go rhs lhs rhsNode lhsNode rhsRoot lhsRoot true + else + go lhs rhs lhsNode rhsNode lhsRoot rhsRoot false + if trueEqFalse then + closeGoalWithTrueEqFalse + unless (← isInconsistent) do + if lhsRoot.ctor && rhsRoot.ctor then + propagateCtor lhsRoot.self rhsRoot.self + unless (← isInconsistent) do + if valueInconsistency then + closeGoalWithValuesEq lhsRoot.self rhsRoot.self + trace_goal[grind.debug] "after addEqStep, {← (← get).ppState}" + checkInvariants +where + go (lhs rhs : Expr) (lhsNode rhsNode lhsRoot rhsRoot : ENode) (flipped : Bool) : GoalM Unit := do + trace_goal[grind.debug] "adding {← ppENodeRef lhs} ↦ {← ppENodeRef rhs}" + /- + We have the following `target?/proof?` + `lhs -> ... -> lhsNode.root` + `rhs -> ... -> rhsNode.root` + We want to convert it to + `lhsNode.root -> ... -> lhs -*-> rhs -> ... -> rhsNode.root` + where step `-*->` is justified by `proof` (or `proof.symm` if `flipped := true`) + -/ + invertTrans lhs + setENode lhs { lhsNode with + target? := rhs + proof? := proof + flipped + } + let lams₁ ← getEqcLambdas lhsRoot + let lams₂ ← getEqcLambdas rhsRoot + let fns₁ ← if lams₁.isEmpty then pure #[] else getFnRoots rhsRoot.self + let fns₂ ← if lams₂.isEmpty then pure #[] else getFnRoots lhsRoot.self + let parents ← removeParents lhsRoot.self + updateRoots lhs rhsNode.root + trace_goal[grind.debug] "{← ppENodeRef lhs} new root {← ppENodeRef rhsNode.root}, {← ppENodeRef (← getRoot lhs)}" + reinsertParents parents + /- + Remark: we used to `propagateDown` here, but this was problematic + because it limits what the propagator can do because several invariants do not + hold until we complete all updates. + -/ + -- TODO: improve performance: we only need to collect terms that may propagate. + let toPropagateDown ← getEqc lhs + setENode lhsNode.root { (← getENode lhsRoot.self) with -- We must retrieve `lhsRoot` since it was updated. + next := rhsRoot.next + } + setENode rhsNode.root { rhsRoot with + next := lhsRoot.next + size := rhsRoot.size + lhsRoot.size + hasLambdas := rhsRoot.hasLambdas || lhsRoot.hasLambdas + heqProofs := isHEq || rhsRoot.heqProofs || lhsRoot.heqProofs + } + propagateBeta lams₁ fns₁ + propagateBeta lams₂ fns₂ + let offsetTodo ← checkOffsetEq rhsRoot lhsRoot + let cutsatTodo ← checkCutsatEq rhsRoot lhsRoot + let ringTodo ← checkCommRingEq rhsRoot lhsRoot + let linarithTodo ← checkLinarithEq rhsRoot lhsRoot + resetParentsOf lhsRoot.self + copyParentsTo parents rhsNode.root + unless (← isInconsistent) do + updateMT rhsRoot.self + unless (← isInconsistent) do + for parent in parents do + propagateUp parent + for e in toPropagateDown do + propagateDown e + propagateOffset offsetTodo + propagateCutsat cutsatTodo + propagateCommRing ringTodo + propagateLinarith linarithTodo + updateRoots (lhs : Expr) (rootNew : Expr) : GoalM Unit := do + traverseEqc lhs fun n => + setENode n.self { n with root := rootNew } + +/-- Ensures collection of equations to be processed is empty. -/ +private def resetNewFacts : GoalM Unit := + modify fun s => { s with newFacts := #[] } + +/-- Pops and returns the next equality to be processed. -/ +private def popNextFact? : GoalM (Option NewFact) := do + let r := (← get).newFacts.back? + if r.isSome then + modify fun s => { s with newFacts := s.newFacts.pop } + return r + +private def addEqCore (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit := do + addEqStep lhs rhs proof isHEq + processNewFacts + +/-- Adds a new equality `lhs = rhs`. It assumes `lhs` and `rhs` have already been internalized. -/ +private def addEq (lhs rhs proof : Expr) : GoalM Unit := do + addEqCore lhs rhs proof false + +/-- Adds a new heterogeneous equality `lhs ≍ rhs`. It assumes `lhs` and `rhs` have already been internalized. -/ +private def addHEq (lhs rhs proof : Expr) : GoalM Unit := do + addEqCore lhs rhs proof true + +/-- Save asserted facts for pretty printing goal. -/ +private def storeFact (fact : Expr) : GoalM Unit := do + modify fun s => { s with facts := s.facts.push fact } + +/-- Internalizes `lhs` and `rhs`, and then adds equality `lhs = rhs`. -/ +def addNewEq (lhs rhs proof : Expr) (generation : Nat) : GoalM Unit := do + let eq ← mkEq lhs rhs + storeFact eq + internalize lhs generation eq + internalize rhs generation eq + addEq lhs rhs proof + +private def addFactStep (fact : Expr) (proof : Expr) (generation : Nat) : GoalM Unit := do + storeFact fact + trace_goal[grind.assert] "{fact}" + let_expr Not p := fact + | go fact false + go p true +where + go (p : Expr) (isNeg : Bool) : GoalM Unit := do + match_expr p with + | Eq α lhs rhs => + if α.isProp then + -- It is morally an iff. + -- We do not use the `goEq` optimization because we want to register `p` as a case-split + goFact p isNeg + else + goEq p lhs rhs isNeg false + | HEq _ lhs _ rhs => goEq p lhs rhs isNeg true + | _ => goFact p isNeg + + goFact (p : Expr) (isNeg : Bool) : GoalM Unit := do + internalize p generation + if isNeg then + addEq p (← getFalseExpr) (← mkEqFalse proof) + else + addEq p (← getTrueExpr) (← mkEqTrue proof) + + goEq (p : Expr) (lhs rhs : Expr) (isNeg : Bool) (isHEq : Bool) : GoalM Unit := do + if isNeg then + internalize p generation + addEq p (← getFalseExpr) (← mkEqFalse proof) + else + internalize lhs generation p + internalize rhs generation p + addEqCore lhs rhs proof isHEq + +@[export lean_grind_process_new_facts] +private def processNewFactsImpl : GoalM Unit := do + repeat + if (← isInconsistent) then + resetNewFacts + return () + checkSystem "grind" + let some next := (← popNextFact?) + | return () + match next with + | .eq lhs rhs proof isHEq => addEqStep lhs rhs proof isHEq + | .fact prop proof gen => addFactStep prop proof gen + +/-- Adds a new `fact` justified by the given proof and using the given generation. -/ +def add (fact : Expr) (proof : Expr) (generation := 0) : GoalM Unit := do + if fact.isTrue then return () + if (← isInconsistent) then return () + resetNewFacts + addFactStep fact proof generation + +/-- Adds a new hypothesis. -/ +def addHypothesis (fvarId : FVarId) (generation := 0) : GoalM Unit := do + add (← fvarId.getType) (mkFVar fvarId) generation + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Ctor.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Ctor.lean new file mode 100644 index 0000000000000000000000000000000000000000..1d20f9ad546a2fb59ce6640ec6433a5715248275 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Ctor.lean @@ -0,0 +1,60 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.ProveEq + +namespace Lean.Meta.Grind + +private partial def propagateInjEqs (eqs : Expr) (proof : Expr) : GoalM Unit := do + -- Remark: we must use `shareCommon` before using `pushEq` and `pushHEq`. + -- This is needed because the result type of the injection theorem may allocate + match_expr eqs with + | And left right => + propagateInjEqs left (.proj ``And 0 proof) + propagateInjEqs right (.proj ``And 1 proof) + | Eq _ lhs rhs => + pushEq (← preprocessLight lhs) (← preprocessLight rhs) proof + | HEq _ lhs _ rhs => + pushHEq (← preprocessLight lhs) (← preprocessLight rhs) proof + | _ => + reportIssue! "unexpected injectivity theorem result type{indentExpr eqs}" + return () + +/-- +Given constructors `a` and `b`, propagate equalities if they are the same, +and close goal if they are different. +-/ +def propagateCtor (a b : Expr) : GoalM Unit := do + let aType ← whnfD (← inferType a) + let bType ← whnfD (← inferType b) + unless (← withDefault <| isDefEq aType bType) do + return () + let ctor₁ := a.getAppFn + let ctor₂ := b.getAppFn + if ctor₁ == ctor₂ then + let .const ctorName _ := a.getAppFn | return () + let injDeclName := Name.mkStr ctorName "inj" + unless (← getEnv).contains injDeclName do return () + let info ← getConstInfo injDeclName + let n := info.type.getForallArity + let mask : Array (Option Expr) := .replicate n none + /- + We use `mkExpectedTypeHint` here to ensure that `mkAppOptM` will "fill" the implicit + arguments of `injDeclName` using exactly the fields of `a` and `b`. + There is no guarantee that `inferType (← mkEqProof a b)` is structurally equal to `a = b`. + -/ + let mask := mask.set! (n-1) (some (← mkExpectedTypeHint (← mkEqProof a b) (← mkEq a b))) + let injLemma ← mkAppOptM injDeclName mask + let injLemmaType ← inferType injLemma + propagateInjEqs injLemmaType injLemma + else + let .const declName _ := aType.getAppFn | return () + let noConfusionDeclName := Name.mkStr declName "noConfusion" + unless (← getEnv).contains noConfusionDeclName do return () + closeGoal (← mkNoConfusion (← getFalseExpr) (← mkEqProof a b)) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Diseq.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Diseq.lean new file mode 100644 index 0000000000000000000000000000000000000000..89275f0969a87ab0f0e81977d79f0d80bcea7995 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Diseq.lean @@ -0,0 +1,86 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Lemmas +import Lean.Meta.Tactic.Grind.Types + +namespace Lean.Meta.Grind + +/-- +Returns `some (c = d)` if +- `c = d` and `False` are in the same equivalence class, and +- `a` (`b`) and `c` are in the same equivalence class, and +- `b` (`a`) and `d` are in the same equivalence class. +Otherwise return `none`. + +Remark `a` and `b` are assumed to have the same type. +-/ +private def getDiseqFor? (a b : Expr) : GoalM (Option Expr) := do + /- + In Z3, we use the congruence table to find equalities more efficiently, + but this optimization would be more complicated here because equalities have + the type as an implicit argument, and `grind`s congruence table assumes it is + hash-consed and canonicalized. So, we use the "slower" approach of visiting + parents. + -/ + let aRoot ← getRoot a + let bRoot ← getRoot b + let aParents ← getParents aRoot + let bParents ← getParents bRoot + if aParents.size ≤ bParents.size then + go aParents + else + go bParents +where + go (parents : ParentSet) : GoalM (Option Expr) := do + for parent in parents do + let_expr Eq α c d := parent | continue + if (← isEqFalse parent) then + -- Remark: we expect `hasType` test to seldom fail, but it can happen because of + -- heterogeneous equalities + if (← isEqv a c <&&> isEqv b d <&&> hasType a α) then + return some parent + if (← isEqv a d <&&> isEqv b c <&&> hasType a α) then + return some parent + return none + +/-- +Returns `true` if `a` and `b` are known to be disequal. +See `getDiseqFor?` +-/ +def isDiseq (a b : Expr) : GoalM Bool := do + return (← getDiseqFor? a b).isSome + +/-- +Returns a proof for `true` if `a` and `b` are known to be disequal. +See `getDiseqFor?` +-/ +def mkDiseqProof? (a b : Expr) : GoalM (Option Expr) := do + let some eq ← getDiseqFor? a b | return none + let_expr f@Eq α c d := eq | unreachable! + let u := f.constLevels! + let h ← mkOfEqFalse (← mkEqFalseProof eq) + let (c, d, h) ← if (← isEqv a c <&&> isEqv b d) then + pure (c, d, h) + else + pure (d, c, mkApp4 (mkConst ``Ne.symm u) α c d h) + -- We have `a = c` and `b = d` + let h ← if isSameExpr a c then + pure h + else + pure <| mkApp6 (mkConst ``Grind.ne_of_ne_of_eq_left u) α a c d (← mkEqProof a c) h + -- `h : a ≠ d + if isSameExpr b d then + return h + else + return mkApp6 (mkConst ``Grind.ne_of_ne_of_eq_right u) α b a d (← mkEqProof b d) h + +def mkDiseqProof (a b : Expr) : GoalM Expr := do + let some h ← mkDiseqProof? a b + | throwError "internal `grind` error, failed to build disequality proof for{indentExpr a}\nand{indentExpr b}" + return h + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EMatch.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EMatch.lean new file mode 100644 index 0000000000000000000000000000000000000000..ff36e11b901fa92996d0ee619f4d7823b9480ee8 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EMatch.lean @@ -0,0 +1,577 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Intro +import Lean.Meta.Tactic.Grind.MatchDiscrOnly +import Lean.Meta.Tactic.Grind.MatchCond +import Lean.Meta.Tactic.Grind.Core + +namespace Lean.Meta.Grind +namespace EMatch +/-! This module implements a simple E-matching procedure as a backtracking search. -/ + +/-- We represent an `E-matching` problem as a list of constraints. -/ +inductive Cnstr where + | /-- Matches pattern `pat` with term `e` -/ + «match» (gen? : Option GenPatternInfo) (pat : Expr) (e : Expr) + | /-- Matches offset pattern `pat+k` with term `e` -/ + offset (gen? : Option GenPatternInfo) (pat : Expr) (k : Nat) (e : Expr) + | /-- This constraint is used to encode multi-patterns. -/ + «continue» (pat : Expr) + deriving Inhabited + +/-- +Internal "marker" for representing unassigned elements in the `assignment` field. +This is a small hack to avoid one extra level of indirection by using `Option Expr` at `assignment`. +-/ +private def unassigned : Expr := mkConst (Name.mkSimple "[grind_unassigned]") + +/-- +Internal "marker" for representing equality proofs for generalized patterns. +They must be synthesized after we have the partial assignment. +-/ +private def delayedEqProof : Expr := mkConst (Name.mkSimple "[grind_delayed_eq_proof]") + +private def assignmentToMessageData (assignment : Array Expr) : Array MessageData := + assignment.reverse.map fun e => + if isSameExpr e unassigned then m!"_" else m!"{e}" + +/-- +Choice point for the backtracking search. +The state of the procedure contains a stack of choices. +-/ +structure Choice where + /-- Constraints to be processed. -/ + cnstrs : List Cnstr + /-- Maximum term generation found so far. -/ + gen : Nat + /-- Partial assignment so far. Recall that pattern variables are encoded as de-Bruijn variables. -/ + assignment : Array Expr + deriving Inhabited + +/-- Context for the E-matching monad. -/ +structure Context where + /-- `useMT` is `true` if we are using the mod-time optimization. It is always set to false for new `EMatchTheorem`s. -/ + useMT : Bool := true + /-- `EMatchTheorem` being processed. -/ + thm : EMatchTheorem := default + /-- Initial application used to start E-matching -/ + initApp : Expr := default + deriving Inhabited + +/-- State for the E-matching monad -/ +structure SearchState where + /-- Choices that still have to be processed. -/ + choiceStack : List Choice := [] + deriving Inhabited + +abbrev M := ReaderT Context $ StateRefT SearchState GoalM + +def M.run' (x : M α) : GoalM α := + x {} |>.run' {} + +@[inline] private abbrev withInitApp (e : Expr) (x : M α) : M α := + withReader (fun ctx => { ctx with initApp := e }) x + +/-- +Assigns `bidx := e` in `c`. If `bidx` is already assigned in `c`, we check whether +`e` and `c.assignment[bidx]` are in the same equivalence class. +This function assumes `bidx < c.assignment.size`. +Recall that we initialize the assignment array with the number of theorem parameters. +-/ +private def assign? (c : Choice) (bidx : Nat) (e : Expr) : OptionT GoalM Choice := do + if h : bidx < c.assignment.size then + let v := c.assignment[bidx] + if isSameExpr v unassigned then + return { c with assignment := c.assignment.set bidx e } + else + guard (← isEqv v e) + return c + else + -- `Choice` was not properly initialized + unreachable! + +/-- +Assigns `bidx` with the marker for a delayed equality proof for generalized patterns. +The proof is assigned after we have the complete assignment. +-/ +private def assignDelayedEqProof? (c : Choice) (bidx : Nat) : OptionT GoalM Choice := do + if h : bidx < c.assignment.size then + let v := c.assignment[bidx] + if isSameExpr v unassigned then + return { c with assignment := c.assignment.set bidx delayedEqProof } + else + return c + else + -- `Choice` was not properly initialized + unreachable! + + +private def unassign (c : Choice) (bidx : Nat) : Choice := + { c with assignment := c.assignment.set! bidx unassigned } + +/-- +Returns `true` if the function `pFn` of a pattern is equivalent to the function `eFn`. +Recall that we ignore universe levels in patterns. +-/ +private def eqvFunctions (pFn eFn : Expr) : Bool := + (pFn.isFVar && pFn == eFn) + || (pFn.isConst && eFn.isConstOf pFn.constName!) + +protected def _root_.Lean.Meta.Grind.GenPatternInfo.assign? (genInfo : GenPatternInfo) (c : Choice) (x : Expr) : OptionT GoalM Choice := do + let c ← assign? c genInfo.xIdx x + let c ← assignDelayedEqProof? c genInfo.hIdx + return c + +/-- Matches a pattern argument. See `matchArgs?`. -/ +private def matchArg? (c : Choice) (pArg : Expr) (eArg : Expr) : OptionT GoalM Choice := do + if isPatternDontCare pArg then + return c + else if pArg.isBVar then + assign? c pArg.bvarIdx! eArg + else if let some pArg := groundPattern? pArg then + /- + We need to use `withReducibleAndIntances` because ground patterns are often instances. + Here is an example + ``` + instance : Max Nat where + max := Nat.max -- Redefined the instance + + example (a : Nat) : max a a = a := by + grind + ``` + Possible future improvements: + - When `diagnostics` is true, try with `withDefault` and report issue if it succeeds. + - (minor) Only use `withReducibleAndInstances` if the argument is an implicit instance. + Potential issue: some user write `{_ : Class α}` when the instance can be inferred from + explicit arguments. + -/ + guard (← isEqv pArg eArg <||> withReducibleAndInstances (isDefEq pArg eArg)) + return c + else if let some (pArg, k) := isOffsetPattern? pArg then + assert! Option.isNone <| isOffsetPattern? pArg + assert! !isPatternDontCare pArg + return { c with cnstrs := .offset none pArg k eArg :: c.cnstrs } + else if let some (genInfo, pArg) := isGenPattern? pArg then + if pArg.isBVar then + let c ← assign? c pArg.bvarIdx! eArg + genInfo.assign? c eArg + else if let some pArg := groundPattern? pArg then + guard (← isEqv pArg eArg <||> withReducibleAndInstances (isDefEq pArg eArg)) + genInfo.assign? c eArg + else if let some (pArg, k) := isOffsetPattern? pArg then + return { c with cnstrs := .offset (some genInfo) pArg k eArg :: c.cnstrs } + else + return { c with cnstrs := .match (some genInfo) pArg eArg :: c.cnstrs } + else + return { c with cnstrs := .match none pArg eArg :: c.cnstrs } + +private def Choice.updateGen (c : Choice) (gen : Nat) : Choice := + { c with gen := Nat.max gen c.gen } + +private def pushChoice (c : Choice) : M Unit := + modify fun s => { s with choiceStack := c :: s.choiceStack } + +/-- +Matches arguments of pattern `p` with term `e`. Returns `some` if successful, +and `none` otherwise. It may update `c`s assignment and list of constraints to be +processed. +-/ +private partial def matchArgs? (c : Choice) (p : Expr) (e : Expr) : OptionT GoalM Choice := do + if !p.isApp then return c -- Done + let pArg := p.appArg! + let eArg := e.appArg! + let c ← matchArg? c pArg eArg + matchArgs? c p.appFn! e.appFn! + +/-- Similar to `matchArgs?` but if `p` has fewer arguments than `e`, we match `p` with a prefix of `e`. -/ +private partial def matchArgsPrefix? (c : Choice) (p : Expr) (e : Expr) : OptionT GoalM Choice := do + let pn := p.getAppNumArgs + let en := e.getAppNumArgs + guard (pn <= en) + if pn == en then + matchArgs? c p e + else + matchArgs? c p (e.getAppPrefix pn) + +-- Helper function for `processMatch` and `processGenMatch` +@[inline] private def isCandidate (n : ENode) (pFn : Expr) (pNumArgs : Nat) (maxGeneration : Nat) : Bool := + -- Remark: we use `<` because the instance generation is the maximum term generation + 1 + n.generation < maxGeneration + -- uses heterogeneous equality or is the root of its congruence class + && (n.heqProofs || n.isCongrRoot) + && eqvFunctions pFn n.self.getAppFn + && n.self.getAppNumArgs == pNumArgs + +private def assignGenInfo? (genInfo? : Option GenPatternInfo) (c : Choice) (x : Expr) : OptionT GoalM Choice := do + let some genInfo := genInfo? | return c + genInfo.assign? c x + +/-- +Matches pattern `p` with term `e` with respect to choice `c`. +We traverse the equivalence class of `e` looking for applications compatible with `p`. +For each candidate application, we match the arguments and may update `c`s assignments and constraints. +We add the updated choices to the choice stack. +-/ +private partial def processMatch (c : Choice) (genInfo? : Option GenPatternInfo) (p : Expr) (e : Expr) : M Unit := do + let maxGeneration ← getMaxGeneration + let pFn := p.getAppFn + let numArgs := p.getAppNumArgs + let mut curr := e + repeat + let n ← getENode curr + -- Remark: we use `<` because the instance generation is the maximum term generation + 1 + if isCandidate n pFn numArgs maxGeneration then + if let some c ← assignGenInfo? genInfo? c e |>.run then + if let some c ← matchArgs? c p curr |>.run then + pushChoice (c.updateGen n.generation) + curr ← getNext curr + if isSameExpr curr e then break + +/-- +Matches offset pattern `pArg+k` with term `e` with respect to choice `c`. +-/ +private partial def processOffset (c : Choice) (genInfo? : Option GenPatternInfo) (pArg : Expr) (k : Nat) (e : Expr) : M Unit := do + let maxGeneration ← getMaxGeneration + let mut curr := e + repeat + let n ← getENode curr + if n.generation < maxGeneration then + if let some (eArg, k') ← isOffset? curr |>.run then + if let some c ← assignGenInfo? genInfo? c e |>.run then + if k' < k then + let c := c.updateGen n.generation + pushChoice { c with cnstrs := .offset none pArg (k - k') eArg :: c.cnstrs } + else if k' == k then + if let some c ← matchArg? c pArg eArg |>.run then + pushChoice (c.updateGen n.generation) + else if k' > k then + let eArg' := mkNatAdd eArg (mkNatLit (k' - k)) + let eArg' ← shareCommon (← canon eArg') + internalize eArg' (n.generation+1) + if let some c ← matchArg? c pArg eArg' |>.run then + pushChoice (c.updateGen n.generation) + else if let some k' ← evalNat curr |>.run then + if let some c ← assignGenInfo? genInfo? c e |>.run then + if k' >= k then + let eArg' := mkNatLit (k' - k) + let eArg' ← shareCommon (← canon eArg') + internalize eArg' (n.generation+1) + if let some c ← matchArg? c pArg eArg' |>.run then + pushChoice (c.updateGen n.generation) + curr ← getNext curr + if isSameExpr curr e then break + +/-- +Returns "applications" in the given goal that may match `p`. +We say "applications," because we assume a constant is a zero-ary application. +-/ +private def getAppsOf (p : Expr) : GoalM (Option (List Expr)) := do + if p.isConst then + if (← alreadyInternalized p) then + return some [p] + else + return none + else + return (← get).appMap.find? p.toHeadIndex + +/-- Processes `continue` constraint used to implement multi-patterns. -/ +private def processContinue (c : Choice) (p : Expr) : M Unit := do + let some apps ← getAppsOf p + | return () + let maxGeneration ← getMaxGeneration + for app in apps do + let n ← getENode app + if n.generation < maxGeneration + && (n.heqProofs || n.isCongrRoot) then + if let some c ← matchArgsPrefix? c p app |>.run then + let gen := n.generation + let c := { c with gen := Nat.max gen c.gen } + modify fun s => { s with choiceStack := c :: s.choiceStack } + +/-- +Given a proposition `prop` corresponding to an equational theorem. +Annotate the conditions using `Grind.MatchCond`. See `MatchCond.lean`. +-/ +private partial def annotateEqnTypeConds (prop : Expr) (k : Expr → M Expr := pure) : M Expr := do + if let .forallE n d b bi := prop then + let d := if (← isProp d) then + markAsPreMatchCond d + else + d + withLocalDecl n bi d fun x => do + mkForallFVars #[x] (← annotateEqnTypeConds (b.instantiate1 x) k) + else + k prop + +/-- +Helper function for marking parts of `match`-equation theorem as "do-not-simplify" +`initApp` is the match-expression used to instantiate the `match`-equation. +It also introduce `Grind.MatchCond` at equation conditions. See `MatchCond.lean`. +-/ +private def annotateMatchEqnType (prop : Expr) (initApp : Expr) : M Expr := do + annotateEqnTypeConds prop fun prop => do + let_expr f@Eq α lhs rhs := prop | return prop + -- See comment at `Grind.EqMatch` + let lhs ← markAsSimpMatchDiscrsOnly lhs + return mkApp4 (mkConst ``Grind.EqMatch f.constLevels!) α lhs rhs initApp + +/-- +Stores new theorem instance in the state. +Recall that new instances are internalized later, after a full round of ematching. +-/ +private def addNewInstance (thm : EMatchTheorem) (proof : Expr) (generation : Nat) : M Unit := do + let proof ← instantiateMVars proof + if grind.debug.proofs.get (← getOptions) then + check proof + let mut prop ← inferType proof + let mut proof := proof + if (← isMatchEqLikeDeclName thm.origin.key) then + prop ← annotateMatchEqnType prop (← read).initApp + -- We must add a hint here because `annotateMatchEqnType` introduces `simpMatchDiscrsOnly` and + -- `Grind.PreMatchCond` which are not reducible. + proof := mkExpectedPropHint proof prop + else if (← isEqnThm thm.origin.key) then + prop ← annotateEqnTypeConds prop + -- We must add a hint because `annotateEqnTypeConds` introduces `Grind.PreMatchCond` + -- which is not reducible. + proof := mkExpectedPropHint proof prop + trace_goal[grind.ematch.instance] "{← thm.origin.pp}: {prop}" + addTheoremInstance thm proof prop (generation+1) + +private def synthesizeInsts (mvars : Array Expr) (bis : Array BinderInfo) : OptionT M Unit := do + let thm := (← read).thm + for mvar in mvars, bi in bis do + if bi.isInstImplicit && !(← mvar.mvarId!.isAssigned) then + let type ← inferType mvar + unless (← synthesizeInstanceAndAssign mvar type) do + reportIssue! "failed to synthesize instance when instantiating {← thm.origin.pp}{indentExpr type}" + failure + +private def preprocessGeneralizedPatternRHS (lhs : Expr) (rhs : Expr) (origin : Origin) (expectedType : Expr) : OptionT (StateT Choice M) Expr := do + assert! (← alreadyInternalized lhs) + -- We use `dsimp` here to ensure terms such as `Nat.succ x` are normalized as `x+1`. + let rhs ← preprocessLight (← dsimpCore rhs) + internalize rhs (← getGeneration lhs) + processNewFacts + if (← isEqv lhs rhs) then + return rhs + else + reportIssue! "invalid generalized pattern at `{← origin.pp}`\nwhen processing argument with type{indentExpr expectedType}\nfailed to prove{indentExpr lhs}\nis equal to{indentExpr rhs}" + failure + +private def assignGeneralizedPatternProof (mvarId : MVarId) (eqProof : Expr) (origin : Origin) : OptionT (StateT Choice M) Unit := do + unless (← mvarId.checkedAssign eqProof) do + reportIssue! "invalid generalized pattern at `{← origin.pp}`\nfailed to assign {mkMVar mvarId}\nwith{indentExpr eqProof}" + failure + +private def applyAssignment (mvars : Array Expr) : OptionT (StateT Choice M) Unit := do + let thm := (← read).thm + let numParams := thm.numParams + for h : i in [:mvars.size] do + let bidx := numParams - i - 1 + let mut v := (← get).assignment[bidx]! + if isSameExpr v delayedEqProof then + let mvarId := mvars[i].mvarId! + let mvarIdType ← instantiateMVars (← mvarId.getType) + match_expr mvarIdType with + | Eq α lhs rhs => + let rhs ← preprocessGeneralizedPatternRHS lhs rhs thm.origin mvarIdType + assignGeneralizedPatternProof mvarId (← mkEqProof lhs rhs) thm.origin + | HEq α lhs β rhs => + let rhs ← preprocessGeneralizedPatternRHS lhs rhs thm.origin mvarIdType + assignGeneralizedPatternProof mvarId (← mkHEqProof lhs rhs) thm.origin + | _ => + reportIssue! "invalid generalized pattern at `{← thm.origin.pp}`\nequality type expected{indentExpr mvarIdType}" + failure + else if !isSameExpr v unassigned then + let mvarId := mvars[i].mvarId! + let mvarIdType ← mvarId.getType + let vType ← inferType v + let unassignOrFail : OptionT (StateT Choice M) Unit := do + /- + If there is type error and `vType` is a proposition, we can still instantiate the + theorem by unassigning `v` and using it as an extra hypothesis. + Here is an example to motivate the unassignment. + ``` + example (xs : Array Nat) (w : xs.reverse = xs) (j : Nat) (hj : 0 ≤ j) (hj' : j < xs.size / 2) + : xs[j] = xs[xs.size - 1 - j] := by + grind + ``` + Without the unassignment we get a type error while trying to instantiate the theorem + ``` + theorem getElem_reverse {xs : Array α} {i : Nat} (hi : i < xs.reverse.size) : + (xs.reverse)[i] = xs[xs.size - 1 - i]'(by simp at hi; omega) + ``` + The pattern for this theorem is `xs.reverse[i]`. Note that `hi` occurs there as an implicit argument. + The term `xs[j]` in our goal e-matches the pattern because we have the equality `xs.reverse = xs`. + However, the implicit proof at `xs[j]` has type `j < xs.size` instead of `j < xs.reverse.size`. + -/ + if (← isProp vType) then + modify (unassign · bidx) + else + reportIssue! "type error constructing proof for {← thm.origin.pp}\nwhen assigning metavariable {mvars[i]} with {indentExpr v}\n{← mkHasTypeButIsExpectedMsg vType mvarIdType}" + failure + unless (← withDefault <| isDefEq mvarIdType vType) do + if let some heq ← withoutReportingMVarIssues <| proveEq? vType mvarIdType (abstract := true) then + /- + Some of the `cast`s will appear inside the `Grind.MatchCond` binders, and + we want their proofs to be properly wrapped. + -/ + let heq := mkApp2 (mkConst ``Grind.nestedProof) (← mkEq vType mvarIdType) heq + v ← mkAppM ``cast #[heq, v] + else + unassignOrFail + continue + unless (← mvarId.checkedAssign v) do + unassignOrFail + +/-- +After processing a (multi-)pattern, use the choice assignment to instantiate the proof. +Missing parameters are synthesized using type inference and type class synthesis." +-/ +private partial def instantiateTheorem (c : Choice) : M Unit := withDefault do withNewMCtxDepth do + let thm := (← read).thm + unless (← markTheoremInstance thm.proof c.assignment) do + return () + trace_goal[grind.ematch.instance.assignment] "{← thm.origin.pp}: {assignmentToMessageData c.assignment}" + let proof ← thm.getProofWithFreshMVarLevels + let numParams := thm.numParams + assert! c.assignment.size == numParams + let (mvars, bis, _) ← forallMetaBoundedTelescope (← inferType proof) numParams + if mvars.size != thm.numParams then + reportIssue! "unexpected number of parameters at {← thm.origin.pp}" + return () + let (some _, c) ← applyAssignment mvars |>.run c | return () + let some _ ← synthesizeInsts mvars bis | return () + let proof := mkAppN proof mvars + if (← mvars.allM (·.mvarId!.isAssigned)) then + addNewInstance thm proof c.gen + else + let mvars ← mvars.filterM fun mvar => return !(← mvar.mvarId!.isAssigned) + if let some mvarBad ← mvars.findM? fun mvar => return !(← isProof mvar) then + reportIssue! "failed to instantiate {← thm.origin.pp}, failed to instantiate non propositional argument with type{indentExpr (← inferType mvarBad)}" + let proof ← mkLambdaFVars (binderInfoForMVars := .default) mvars (← instantiateMVars proof) + addNewInstance thm proof c.gen + +/-- Process choice stack until we don't have more choices to be processed. -/ +private def processChoices : M Unit := do + let maxGeneration ← getMaxGeneration + while !(← get).choiceStack.isEmpty do + checkSystem "ematch" + if (← checkMaxInstancesExceeded) then return () + let c ← modifyGet fun s : SearchState => (s.choiceStack.head!, { s with choiceStack := s.choiceStack.tail! }) + if c.gen < maxGeneration then + match c.cnstrs with + | [] => instantiateTheorem c + | .match genInfo? p e :: cnstrs => processMatch { c with cnstrs } genInfo? p e + | .offset genInfo? p k e :: cnstrs => processOffset { c with cnstrs } genInfo? p k e + | .continue p :: cnstrs => processContinue { c with cnstrs } p + +private def main (p : Expr) (cnstrs : List Cnstr) : M Unit := do + let some apps ← getAppsOf p + | return () + let numParams := (← read).thm.numParams + let assignment := .replicate numParams unassigned + let useMT := (← read).useMT + let gmt := (← getThe Goal).ematch.gmt + for app in apps do + if (← checkMaxInstancesExceeded) then return () + let n ← getENode app + if (n.heqProofs || n.isCongrRoot) && + (!useMT || n.mt == gmt) then + withInitApp app do + if let some c ← matchArgsPrefix? { cnstrs, assignment, gen := n.generation } p app |>.run then + modify fun s => { s with choiceStack := [c] } + processChoices + +/-- +Entry point for matching `lhs ←= rhs` patterns. +It traverses disequalities `a = b`, and tries to solve two matching problems: +1- match `lhs` with `a` and `rhs` with `b` +2- match `lhs` with `b` and `rhs` with `a` +-/ +private def matchEqBwdPat (p : Expr) : M Unit := do + let_expr Grind.eqBwdPattern pα plhs prhs := p | return () + let numParams := (← read).thm.numParams + let assignment := .replicate numParams unassigned + let useMT := (← read).useMT + let gmt := (← getThe Goal).ematch.gmt + let false ← getFalseExpr + let mut curr := false + repeat + if (← checkMaxInstancesExceeded) then return () + let n ← getENode curr + if (n.heqProofs || n.isCongrRoot) && + (!useMT || n.mt == gmt) then + let_expr Eq α lhs rhs := n.self | pure () + if let some c₀ ← matchArg? { cnstrs := [], assignment, gen := n.generation } pα α |>.run then + let go (lhs rhs : Expr) : M Unit := do + let some c₁ ← matchArg? c₀ plhs lhs |>.run | return () + let some c₂ ← matchArg? c₁ prhs rhs |>.run | return () + modify fun s => { s with choiceStack := [c₂] } + processChoices + go lhs rhs + go rhs lhs + if isSameExpr n.next false then return () + curr := n.next + +def ematchTheorem (thm : EMatchTheorem) : M Unit := do + if (← checkMaxInstancesExceeded) then return () + withReader (fun ctx => { ctx with thm }) do + let ps := thm.patterns + match ps, (← read).useMT with + | [p], _ => if isEqBwdPattern p then matchEqBwdPat p else main p [] + | p::ps, false => main p (ps.map (.continue ·)) + | _::_, true => tryAll ps [] + | _, _ => unreachable! +where + /-- + When using the mod-time optimization with multi-patterns, + we must start ematching at each different pattern. That is, + if we have `[p₁, p₂, p₃]`, we must execute + - `main p₁ [.continue p₂, .continue p₃]` + - `main p₂ [.continue p₁, .continue p₃]` + - `main p₃ [.continue p₁, .continue p₂]` + -/ + tryAll (ps : List Expr) (cs : List Cnstr) : M Unit := do + match ps with + | [] => return () + | p::ps => + main p (cs.reverse ++ (ps.map (.continue ·))) + tryAll ps (.continue p :: cs) + +def ematchTheorems (thms : PArray EMatchTheorem) : M Unit := do + thms.forM ematchTheorem + +end EMatch + +open EMatch + +/-- Performs one round of E-matching, and returns new instances. -/ +private def ematchCore : GoalM Unit := do + let go (thms newThms : PArray EMatchTheorem) : EMatch.M Unit := do + withReader (fun ctx => { ctx with useMT := true }) <| ematchTheorems thms + withReader (fun ctx => { ctx with useMT := false }) <| ematchTheorems newThms + if (← checkMaxInstancesExceeded <||> checkMaxEmatchExceeded) then + return () + else + go (← get).ematch.thms (← get).ematch.newThms |>.run' + modify fun s => { s with + ematch.thms := s.ematch.thms ++ s.ematch.newThms + ematch.newThms := {} + ematch.gmt := s.ematch.gmt + 1 + ematch.num := s.ematch.num + 1 + } + +/-- Performs one round of E-matching, and returns `true` if new instances were generated. -/ +def ematch : GoalM Bool := do + let numInstances := (← get).ematch.numInstances + ematchCore + return (← get).ematch.numInstances != numInstances + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EMatchTheorem.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EMatchTheorem.lean new file mode 100644 index 0000000000000000000000000000000000000000..e02f3cafddd2bc2a397bb7f091d066d58355456b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EMatchTheorem.lean @@ -0,0 +1,1277 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Init.Grind.Tactics +import Lean.HeadIndex +import Lean.PrettyPrinter +import Lean.Util.FoldConsts +import Lean.Util.CollectFVars +import Lean.Meta.Basic +import Lean.Meta.InferType +import Lean.Meta.Eqns +import Lean.Meta.Match.MatchEqs +import Lean.Meta.Tactic.Grind.Util + +namespace Lean.Meta.Grind + +/-- +`grind` uses symbol priorities when inferring patterns for E-matching. +Symbols not in `map` are assumed to have default priority (i.e., `eval_prio default`). +-/ +structure SymbolPriorities where + map : PHashMap Name Nat := {} + deriving Inhabited + +structure SymbolPriorityEntry where + declName : Name + prio : Nat + deriving Inhabited + +/-- Removes the given declaration from `s`. -/ +def SymbolPriorities.erase (s : SymbolPriorities) (declName : Name) : SymbolPriorities := + { s with map := s.map.erase declName } + +/-- Inserts `declName ↦ prio` into `s`. -/ +def SymbolPriorities.insert (s : SymbolPriorities) (declName : Name) (prio : Nat) : SymbolPriorities := + { s with map := s.map.insert declName prio } + +/-- Returns `declName` priority for E-matching pattern inference in `s`. -/ +def SymbolPriorities.getPrio (s : SymbolPriorities) (declName : Name) : Nat := + if let some prio := s.map.find? declName then + prio + else + eval_prio default + +/-- +Returns `true`, if there is an entry `declName ↦ prio` in `s`. +Recall that symbols not in `s` are assumed to have default priority. +-/ +def SymbolPriorities.contains (s : SymbolPriorities) (declName : Name) : Bool := + s.map.contains declName + +private builtin_initialize symbolPrioExt : SimpleScopedEnvExtension SymbolPriorityEntry SymbolPriorities ← + registerSimpleScopedEnvExtension { + initial := {} + addEntry := fun s {declName, prio} => s.insert declName prio + } + +def resetSymbolPrioExt : CoreM Unit := do + modifyEnv fun env => symbolPrioExt.modifyState env fun _ => {} + +def getGlobalSymbolPriorities : CoreM SymbolPriorities := + return symbolPrioExt.getState (← getEnv) + +/-- Sets `declName` priority to be used during E-matching pattern inference -/ +def addSymbolPriorityAttr (declName : Name) (attrKind : AttributeKind) (prio : Nat) : MetaM Unit := do + symbolPrioExt.add { declName, prio } attrKind + +def mkOffsetPattern (pat : Expr) (k : Nat) : Expr := + mkApp2 (mkConst ``Grind.offset) pat (mkRawNatLit k) + +private def detectOffsets (pat : Expr) : MetaM Expr := do + let pre (e : Expr) := do + if e == pat then + -- We only consider nested offset patterns + return .continue e + else match e with + | .letE .. | .lam .. | .forallE .. => return .done e + | _ => + let some (e, k) ← isOffset? e + | return .continue e + if k == 0 then return .continue e + return .continue <| mkOffsetPattern e k + Core.transform pat (pre := pre) + +def isOffsetPattern? (pat : Expr) : Option (Expr × Nat) := Id.run do + let_expr Grind.offset pat k := pat | none + let .lit (.natVal k) := k | none + return some (pat, k) + +/-- +`detectOffsets` inverse. +This function is used to expand `mkOffsetPattern` occurring in a constant pattern. +-/ +private def expandOffsetPatterns (pat : Expr) : CoreM Expr := do + let pre (e : Expr) := do + match e with + | .letE .. | .lam .. | .forallE .. => return .done e + | _ => + let some (e, k) := isOffsetPattern? e + | return .continue e + if k == 0 then return .continue e + return .continue <| mkNatAdd e (mkNatLit k) + Core.transform pat (pre := pre) + +def mkEqBwdPattern (u : List Level) (α : Expr) (lhs rhs : Expr) : Expr := + mkApp3 (mkConst ``Grind.eqBwdPattern u) α lhs rhs + +def isEqBwdPattern (e : Expr) : Bool := + e.isAppOfArity ``Grind.eqBwdPattern 3 + +def isEqBwdPattern? (e : Expr) : Option (Expr × Expr) := + let_expr Grind.eqBwdPattern _ lhs rhs := e + | none + some (lhs, rhs) + +def mkGenPattern (u : List Level) (α : Expr) (h : Expr) (x : Expr) (val : Expr) : Expr := + mkApp4 (mkConst ``Grind.genPattern u) α h x val + +def mkGenHEqPattern (u : List Level) (α β : Expr) (h : Expr) (x : Expr) (val : Expr) : Expr := + mkApp5 (mkConst ``Grind.genHEqPattern u) α β h x val + +/-- Generalized pattern information. See `Grind.genPattern` gadget. -/ +structure GenPatternInfo where + heq : Bool + hIdx : Nat + xIdx : Nat + deriving Repr + +def isGenPattern? (pat : Expr) : Option (GenPatternInfo × Expr) := + match_expr pat with + | Grind.genPattern _ h x pat => Id.run do + let .bvar hIdx := h | unreachable! + let .bvar xIdx := x | unreachable! + return some ({ heq := false, hIdx, xIdx }, pat ) + | Grind.genHEqPattern _ _ h x pat => Id.run do + let .bvar hIdx := h | unreachable! + let .bvar xIdx := x | unreachable! + return some ({ heq := true, hIdx, xIdx }, pat ) + | _ => none + +/-- Returns `true` if `declName` is the name of a `match`-expression congruence equation. -/ +def isMatchCongrEqDeclName (declName : Name) : CoreM Bool := do + let declName := privateToUserName declName + match declName with + | .str p s => return (← isMatcher p) && Match.isCongrEqnReservedNameSuffix s + | _ => return false + +/-- Returns `true` if `e` is a constant for a `match`-expression congruence equation. -/ +private def isMatchCongrEqConst (e : Expr) : CoreM Bool := do + let .const declName _ := e | return false + isMatchCongrEqDeclName declName + +/-- +Given the type of a `match` congruence equation, annotate the discriminants using +the gadgets `Grind.genPattern` and `Grind.genHEqPattern`. +For example, consider the following `match` congruence theorem type +``` +forall + (motive : Option Nat → Sort u_1) (a✝ : Option Nat) + (h_1 : a✝ = none → motive none) + (h_2 : (val : Nat) → a✝ = some val → motive (some val)) + (val✝ : Nat) + (heq_1 : a✝ = some val✝), + g.match_1 motive a✝ h_1 h_2 ≍ h_2 val✝ heq_1 +``` +This function returns the type +``` +forall + (motive : Option Nat → Sort u_1) (a✝ : Option Nat) + (h_1 : a✝ = none → motive none) + (h_2 : (val : Nat) → a✝ = some val → motive (some val)) + (val✝ : Nat) + (heq_1 : a✝ = some val✝), + g.match_1 motive (Grind.genPattern (Option Nat) heq_1 a✝ (some val✝)) h_1 h_2 + ≍ h_2 val✝ heq_1 +``` +The gadget is used to infer a `generalize` pattern. The following term is used +during E-matching `g.match_1 motive (Grind.genPattern (Option Nat) heq_1 a✝ (some val✝)) h_1 h_2` +when matching `Grind.genPattern (Option Nat) heq_1 a✝ (some val✝)` the matcher uses +`(some val✝)` as the actual pattern, but also assigns `heq_1` `a✝` using the information stored +in the equivalence class. +-/ +private def preprocessMatchCongrEqType (type : Expr) : MetaM Expr := do + forallTelescopeReducing type fun hs resultType => do + let lhs ← match_expr resultType with + | Eq _ lhs _ => pure lhs + | HEq _ lhs _ _ => pure lhs + | _ => return type + let lhsFn := lhs.getAppFn + let .const declName _ := lhsFn | return type + let some matcherInfo ← getMatcherInfo? declName | return type + let range := matcherInfo.getDiscrRange + let mut args := lhs.getAppArgs + for h in hs do + match_expr (← inferType h) with + | f@Eq α lhs rhs => + for i in range do + if lhs == args[i]! then + args := args.set! i (mkGenPattern f.constLevels! α h lhs rhs) + break + | f@HEq α lhs β rhs => + for i in range do + if lhs == args[i]! then + args := args.set! i (mkGenHEqPattern f.constLevels! α β h lhs rhs) + break + | _ => pure () + let lhsNew := mkAppN lhsFn args + let resultTypeFn := resultType.getAppFn + let resultArgs := resultType.getAppArgs + let resultType := mkAppN resultTypeFn (resultArgs.set! 1 lhsNew) + mkForallFVars hs resultType + +/-- +A heuristic procedure for detecting generalized patterns. +For example, given the theorem +``` +theorem Option.pbind_some' {α β} {x : Option α} {a : α} {f : (a : α) → x = some a → Option β} + (h : x = some a) : pbind x f = f a h +``` +In the current implementation, we support only occurrences in the resulting type. +Thus, the following resulting type is generated for the example above: +``` +pbind (Grind.genPattern h x (some a)) f = f a h +``` +-/ +private def detectGeneralizedPatterns? (type : Expr) : MetaM Expr := do + forallTelescopeReducing type fun hs resultType => do + let isTarget? (lhs : Expr) (rhs : Expr) (s : FVarSubst) : Option (FVarId × Expr) := Id.run do + let .fvar fvarId := lhs | return none + if !hs.contains lhs then + return none -- It is a foreign free variable + if rhs.containsFVar fvarId then + return none -- It is not a generalization if `rhs` contains it + if s.contains fvarId then + return none -- Remark: may want to abort instead, it is probably not a generalization + let rhs := s.apply rhs + return some (fvarId, rhs) + let mut s : FVarSubst := {} + for h in hs do + match_expr (← inferType h) with + | f@Eq α lhs rhs => + let some (fvarId, rhs) := isTarget? lhs rhs s | pure () + s := s.insert fvarId <| mkGenPattern f.constLevels! α h lhs rhs + | f@HEq α lhs β rhs => + let some (fvarId, rhs) := isTarget? lhs rhs s | pure () + s := s.insert fvarId <| mkGenHEqPattern f.constLevels! α β h lhs rhs + | _ => pure () + if s.isEmpty then + return type + let resultType' := s.apply resultType + if resultType' == resultType then + return type + mkForallFVars hs resultType' + +/-- +Given the proof for a proposition to be used as an E-matching theorem, +infers its type, and preprocess it to identify generalized patterns. +Recall that we infer these generalized patterns automatically for +`match` congruence equations. +-/ +private def inferEMatchProofType (proof : Expr) (gen : Bool) : MetaM Expr := do + let type ← inferType proof + if (← isMatchCongrEqConst proof) then + preprocessMatchCongrEqType type + else if gen then + detectGeneralizedPatterns? type + else + return type + +-- Configuration for the `grind` normalizer. We want both `zetaDelta` and `zeta` +private def normConfig : Grind.Config := {} +theorem normConfig_zeta : normConfig.zeta = true := rfl +theorem normConfig_zetaDelta : normConfig.zetaDelta = true := rfl + +def preprocessPattern (pat : Expr) (normalizePattern := true) : MetaM Expr := do + let pat ← instantiateMVars pat + let pat ← unfoldReducible pat + let pat ← if normalizePattern then normalize pat normConfig else pure pat + let pat ← detectOffsets pat + let pat ← foldProjs pat + return pat + +inductive Origin where + /-- A global declaration in the environment. -/ + | decl (declName : Name) + /-- A local hypothesis. -/ + | fvar (fvarId : FVarId) + /-- + A proof term provided directly to a call to `grind` where `ref` + is the provided grind argument. The `id` is a unique identifier for the call. + -/ + | stx (id : Name) (ref : Syntax) + /-- It is local, but we don't have a local hypothesis for it. -/ + | local (id : Name) + deriving Inhabited, Repr + +/-- A unique identifier corresponding to the origin. -/ +def Origin.key : Origin → Name + | .decl declName => declName + | .fvar fvarId => fvarId.name + | .stx id _ => id + | .local id => id + +def Origin.pp [Monad m] [MonadEnv m] [MonadError m] (o : Origin) : m MessageData := do + match o with + | .decl declName => return MessageData.ofConst (← mkConstWithLevelParams declName) + | .fvar fvarId => return mkFVar fvarId + | .stx _ ref => return ref + | .local id => return id + +instance : BEq Origin where + beq a b := a.key == b.key + +instance : Hashable Origin where + hash a := hash a.key + +inductive EMatchTheoremKind where + | eqLhs (gen : Bool) + | eqRhs (gen : Bool) + | eqBoth (gen : Bool) + | eqBwd + | fwd + | bwd (gen : Bool) + | leftRight + | rightLeft + | default (gen : Bool) + | user /- pattern specified using `grind_pattern` command -/ + deriving Inhabited, BEq, Repr, Hashable + +def EMatchTheoremKind.isEqLhs : EMatchTheoremKind → Bool + | .eqLhs _ => true + | _ => false + +def EMatchTheoremKind.isDefault : EMatchTheoremKind → Bool + | .default _ => true + | _ => false + +private def EMatchTheoremKind.toAttribute : EMatchTheoremKind → String + | .eqLhs true => "[grind = gen]" + | .eqLhs false => "[grind =]" + | .eqRhs true => "[grind =_ gen]" + | .eqRhs false => "[grind =_]" + | .eqBoth false => "[grind _=_]" + | .eqBoth true => "[grind _=_ gen]" + | .eqBwd => "[grind ←=]" + | .fwd => "[grind →]" + | .bwd false => "[grind ←]" + | .bwd true => "[grind ← gen]" + | .leftRight => "[grind =>]" + | .rightLeft => "[grind <=]" + | .default false => "[grind]" + | .default true => "[grind gen]" + | .user => "[grind]" + +private def EMatchTheoremKind.explainFailure : EMatchTheoremKind → String + | .eqLhs _ => "failed to find pattern in the left-hand side of the theorem's conclusion" + | .eqRhs _ => "failed to find pattern in the right-hand side of the theorem's conclusion" + | .eqBoth _ => unreachable! -- eqBoth is a macro + | .eqBwd => "failed to use theorem's conclusion as a pattern" + | .fwd => "failed to find patterns in the antecedents of the theorem" + | .bwd _ => "failed to find patterns in the theorem's conclusion" + | .leftRight => "failed to find patterns searching from left to right" + | .rightLeft => "failed to find patterns searching from right to left" + | .default _ => "failed to find patterns" + | .user => unreachable! + +/-- A theorem for heuristic instantiation based on E-matching. -/ +structure EMatchTheorem where + /-- + It stores universe parameter names for universe polymorphic proofs. + Recall that it is non-empty only when we elaborate an expression provided by the user. + When `proof` is just a constant, we can use the universe parameter names stored in the declaration. + -/ + levelParams : Array Name + proof : Expr + numParams : Nat + patterns : List Expr + /-- Contains all symbols used in `patterns`. -/ + symbols : List HeadIndex + origin : Origin + /-- The `kind` is used for generating the `patterns`. We save it here to implement `grind?`. -/ + kind : EMatchTheoremKind + deriving Inhabited + +/-- Set of E-matching theorems. -/ +structure EMatchTheorems where + /-- The key is a symbol from `EMatchTheorem.symbols`. -/ + private smap : PHashMap Name (List EMatchTheorem) := {} + /-- Set of theorem ids that have been inserted using `insert`. -/ + private origins : PHashSet Origin := {} + /-- Theorems that have been marked as erased -/ + private erased : PHashSet Origin := {} + /-- Mapping from origin to E-matching theorems associated with this origin. -/ + private omap : PHashMap Origin (List EMatchTheorem) := {} + /-- Priorities for pattern inference -/ + private prios : SymbolPriorities := {} + deriving Inhabited + +/-- +Inserts a `thm` with symbols `[s_1, ..., s_n]` to `s`. +We add `s_1 -> { thm with symbols := [s_2, ..., s_n] }`. +When `grind` internalizes a term containing symbol `s`, we +process all theorems `thm` associated with key `s`. +If their `thm.symbols` is empty, we say they are activated. +Otherwise, we reinsert into `map`. +-/ +def EMatchTheorems.insert (s : EMatchTheorems) (thm : EMatchTheorem) : EMatchTheorems := Id.run do + let .const declName :: syms := thm.symbols + | unreachable! + let thm := { thm with symbols := syms } + let { smap, origins, erased, omap, prios } := s + let origin := thm.origin + let origins := origins.insert origin + let erased := erased.erase origin + let smap := if let some thms := smap.find? declName then + smap.insert declName (thm::thms) + else + smap.insert declName [thm] + let omap := if let some thms := omap.find? origin then + omap.insert origin (thm::thms) + else + omap.insert origin [thm] + return { smap, origins, erased, omap, prios } + +/-- Returns `true` if `s` contains a theorem with the given origin. -/ +def EMatchTheorems.contains (s : EMatchTheorems) (origin : Origin) : Bool := + s.origins.contains origin + +/-- Mark the theorem with the given origin as `erased` -/ +def EMatchTheorems.erase (s : EMatchTheorems) (origin : Origin) : EMatchTheorems := + { s with erased := s.erased.insert origin, origins := s.origins.erase origin } + +/-- Returns true if the theorem has been marked as erased. -/ +def EMatchTheorems.isErased (s : EMatchTheorems) (origin : Origin) : Bool := + s.erased.contains origin + +/-- +Retrieves theorems from `s` associated with the given symbol. See `EMatchTheorem.insert`. +The theorems are removed from `s`. +-/ +@[inline] +def EMatchTheorems.retrieve? (s : EMatchTheorems) (sym : Name) : Option (List EMatchTheorem × EMatchTheorems) := + if let some thms := s.smap.find? sym then + some (thms, { s with smap := s.smap.erase sym }) + else + none + +/-- +Returns theorems associated with the given origin. +-/ +def EMatchTheorems.find (s : EMatchTheorems) (origin : Origin) : List EMatchTheorem := + if let some thms := s.omap.find? origin then + thms + else + [] + +def EMatchTheorem.getProofWithFreshMVarLevels (thm : EMatchTheorem) : MetaM Expr := do + if thm.proof.isConst && thm.levelParams.isEmpty then + let declName := thm.proof.constName! + let info ← getConstInfo declName + if info.levelParams.isEmpty then + return thm.proof + else + mkConstWithFreshMVarLevels declName + else if thm.levelParams.isEmpty then + return thm.proof + else + let us ← thm.levelParams.mapM fun _ => mkFreshLevelMVar + return thm.proof.instantiateLevelParamsArray thm.levelParams us + +private builtin_initialize ematchTheoremsExt : SimpleScopedEnvExtension EMatchTheorem EMatchTheorems ← + registerSimpleScopedEnvExtension { + addEntry := EMatchTheorems.insert + initial := {} + } + +/-- Returns `true` if `declName` has been tagged as an E-match theorem using `[grind]`. -/ +def isEMatchTheorem (declName : Name) : CoreM Bool := do + return ematchTheoremsExt.getState (← getEnv) |>.omap.contains (.decl declName) + +def resetEMatchTheoremsExt : CoreM Unit := do + modifyEnv fun env => ematchTheoremsExt.modifyState env fun _ => {} + +/-- +Auxiliary function to expand a pattern containing forbidden application symbols +into a multi-pattern. + +This function enhances the usability of the `[grind =]` attribute by automatically handling +forbidden pattern symbols. For example, consider the following theorem tagged with this attribute: +``` +getLast?_eq_some_iff {xs : List α} {a : α} : xs.getLast? = some a ↔ ∃ ys, xs = ys ++ [a] +``` +Here, the selected pattern is `xs.getLast? = some a`, but `Eq` is a forbidden pattern symbol. +Instead of producing an error, this function converts the pattern into a multi-pattern, +allowing the attribute to be used conveniently. + +The function recursively expands patterns with forbidden symbols by splitting them +into their sub-components. If the pattern does not contain forbidden symbols, +it is returned as-is. +-/ +partial def splitWhileForbidden (pat : Expr) : List Expr := + match_expr pat with + | Not p => splitWhileForbidden p + | And p₁ p₂ => splitWhileForbidden p₁ ++ splitWhileForbidden p₂ + | Or p₁ p₂ => splitWhileForbidden p₁ ++ splitWhileForbidden p₂ + | Eq _ lhs rhs => splitWhileForbidden lhs ++ splitWhileForbidden rhs + | Iff lhs rhs => splitWhileForbidden lhs ++ splitWhileForbidden rhs + | HEq _ lhs _ rhs => splitWhileForbidden lhs ++ splitWhileForbidden rhs + | _ => [pat] + +private def dontCare := mkConst (Name.mkSimple "[grind_dontcare]") + +def mkGroundPattern (e : Expr) : Expr := + mkAnnotation `grind.ground_pat e + +def groundPattern? (e : Expr) : Option Expr := + annotation? `grind.ground_pat e + +private def isGroundPattern (e : Expr) : Bool := + groundPattern? e |>.isSome + +def isPatternDontCare (e : Expr) : Bool := + e == dontCare + +private def isAtomicPattern (e : Expr) : Bool := + e.isBVar || isPatternDontCare e || isGroundPattern e + +partial def ppPattern (pattern : Expr) : MessageData := Id.run do + if let some e := groundPattern? pattern then + return m!"`[{e}]" + else if isPatternDontCare pattern then + return m!"_" + else match pattern with + | .bvar idx => return m!"#{idx}" + | _ => + if pattern.isAppOfArity ``Grind.offset 2 then + let lhs := ppArg <| pattern.getArg! 0 + let rhs := ppPattern <| pattern.getArg! 1 + return m!"{lhs} + {rhs}" + else + let mut r := m!"{pattern.getAppFn}" + for arg in pattern.getAppArgs do + r := r ++ " " ++ ppArg arg + return r +where + ppArg (arg : Expr) : MessageData := + if isAtomicPattern arg then + ppPattern arg + else + .paren (ppPattern arg) + +namespace NormalizePattern + +structure State where + symbols : Array HeadIndex := #[] + symbolSet : Std.HashSet HeadIndex := {} + bvarsFound : Std.HashSet Nat := {} + +private structure Context where + symPrios : SymbolPriorities + /-- Only symbols with priority `>= minPrio` are considered in patterns. -/ + minPrio : Nat + +abbrev M := ReaderT Context StateRefT State MetaM + +/-- Helper declaration for finding bootstrapping issues. See `isCandidateSymbol`. -/ +private abbrev badForPatterns := [``Eq, ``HEq, ``Iff, ``And, ``Or, ``Not] + +def isCandidateSymbol (declName : Name) (root : Bool) : M Bool := do + let ctx ← read + let prio := ctx.symPrios.getPrio declName + -- Priority 0 are never considered, they are treated as forbidden + if prio == 0 then return false + -- Remark: uncomment the following code to fix bootstrapping issues + -- if declName ∈ badForPatterns then + -- throwError "INSERT `import Init.Grind.Tactics`, otherwise a pattern containing `{declName}` will be used, prio: {prio}" + -- If it is the root symbol, then we check whether `prio ≥ minPrio` + if root then + return prio ≥ ctx.minPrio + else + return true + +private def saveSymbol (h : HeadIndex) : M Unit := do + if let .const declName := h then + if declName == ``Grind.genHEqPattern || declName == ``Grind.genPattern then + -- We do not save gadgets in the list of symbols. + return () + unless (← get).symbolSet.contains h do + modify fun s => { s with symbols := s.symbols.push h, symbolSet := s.symbolSet.insert h } + +private def foundBVar (idx : Nat) : M Bool := + return (← get).bvarsFound.contains idx + +private def saveBVar (idx : Nat) : M Unit := do + modify fun s => { s with bvarsFound := s.bvarsFound.insert idx } + +inductive PatternArgKind where + | /-- Argument is relevant for E-matching. -/ + relevant + | /-- Instance implicit arguments are considered support and handled using `isDefEq`. -/ + instImplicit + | /-- Proofs are ignored during E-matching. Lean is proof irrelevant. -/ + proof + | /-- + Types and type formers are mostly ignored during E-matching, and processed using + `isDefEq`. However, if the argument is of the form `C ..` where `C` is inductive type + we process it as part of the pattern. Suppose we have `as bs : List α`, and a pattern + candidate expression `as ++ bs`, i.e., `@HAppend.hAppend (List α) (List α) (List α) inst as bs`. + If we completely ignore the types, the pattern will just be + ``` + @HAppend.hAppend _ _ _ _ #1 #0 + ``` + This is not ideal because the E-matcher will try it in any goal that contains `++`, + even if it does not even mention lists. + -/ + typeFormer + deriving Repr + +def PatternArgKind.isSupport : PatternArgKind → Bool + | .relevant => false + | _ => true + +/-- +Returns an array `kinds` s.ts `kinds[i]` is the kind of the corresponding argument. + +- a type (that is not a proposition) or type former (which has forward dependencies) or +- a proof, or +- an instance implicit argument + +When `kinds[i].isSupport` is `true`, we say the corresponding argument is a "support" argument. +-/ +def getPatternArgKinds (f : Expr) (numArgs : Nat) : MetaM (Array PatternArgKind) := do + let pinfos := (← getFunInfoNArgs f numArgs).paramInfo + forallBoundedTelescope (← inferType f) numArgs fun xs _ => do + xs.mapIdxM fun idx x => do + if (← isProp x) then + return .relevant + else if (← isProof x) then + return .proof + else if (← isTypeFormer x) then + if h : idx < pinfos.size then + /- + We originally wanted to ignore types and type formers in `grind` and treat them as supporting elements. + Thus, we wanted to always return `.typeFormer`. However, we changed our heuristic because of the following example: + ``` + example {α} (f : α → Type) (a : α) (h : ∀ x, Nonempty (f x)) : Nonempty (f a) := by + grind + ``` + In this example, we are reasoning about types. Therefore, we adjusted the heuristic as follows: + a type or type former is considered a supporting element only if it has forward dependencies. + Note that this is not the case for `Nonempty`. + -/ + if pinfos[idx].hasFwdDeps then return .typeFormer else return .relevant + else + return .typeFormer + else if (← x.fvarId!.getDecl).binderInfo matches .instImplicit then + return .instImplicit + else + return .relevant + +private def getPatternFn? (pattern : Expr) (inSupport : Bool) (root : Bool) (argKind : PatternArgKind) : M (Option Expr) := do + if !pattern.isApp && !pattern.isConst then + return none + else match pattern.getAppFn with + | f@(.const declName _) => + if !(← isCandidateSymbol declName root) then + return none + if declName == ``Grind.genPattern || declName == ``Grind.genHEqPattern then + return some f + if inSupport then + if argKind matches .typeFormer | .relevant then + if (← isInductive declName) then + return some f + return none + return some f + | f@(.fvar _) => + if inSupport then return none else return some f + | _ => + return none + +private partial def go (pattern : Expr) (inSupport : Bool) (root : Bool) : M Expr := do + if let some (e, k) := isOffsetPattern? pattern then + let e ← goArg e inSupport .relevant + if e == dontCare then + return dontCare + else + return mkOffsetPattern e k + let some f ← getPatternFn? pattern inSupport root .relevant + | throwError "invalid pattern, (non-forbidden) application expected{indentExpr pattern}" + assert! f.isConst || f.isFVar + unless f.isConstOf ``Grind.eqBwdPattern do + saveSymbol f.toHeadIndex + let mut args := pattern.getAppArgs.toVector + let patternArgKinds ← getPatternArgKinds f args.size + for h : i in [:args.size] do + let arg := args[i] + let argKind := patternArgKinds[i]?.getD .relevant + args := args.set i (← goArg arg (inSupport || argKind.isSupport) argKind) + return mkAppN f args.toArray +where + goArg (arg : Expr) (inSupport : Bool) (argKind : PatternArgKind) : M Expr := do + if !arg.hasLooseBVars then + if arg.hasMVar then + pure dontCare + else + return mkGroundPattern (← expandOffsetPatterns arg) + else match arg with + | .bvar idx => + if inSupport && (← foundBVar idx) then + pure dontCare + else + saveBVar idx + pure arg + | _ => + if let some _ ← getPatternFn? arg inSupport (root := false) argKind then + go arg inSupport (root := false) + else + pure dontCare + +def main (patterns : List Expr) (symPrios : SymbolPriorities) (minPrio : Nat) : MetaM (List Expr × List HeadIndex × Std.HashSet Nat) := do + let (patterns, s) ← patterns.mapM (go (inSupport := false) (root := true)) { symPrios, minPrio } |>.run {} + return (patterns, s.symbols.toList, s.bvarsFound) + +def normalizePattern (e : Expr) : M Expr := do + go e (inSupport := false) (root := true) + +end NormalizePattern + +/-- +Returns `true` if free variables in `type` are not in `thmVars` or are in `fvarsFound`. +We use this function to check whether `type` is fully instantiated. +-/ +private def checkTypeFVars (thmVars : FVarIdSet) (fvarsFound : FVarIdSet) (type : Expr) : Bool := + let typeFVars := (collectFVars {} type).fvarIds + typeFVars.all fun fvarId => !thmVars.contains fvarId || fvarsFound.contains fvarId + +/-- +Given an type class instance type `instType`, returns true if free variables in input parameters +1- are not in `thmVars`, or +2- are in `fvarsFound`. +Remark: `fvarsFound` is a subset of `thmVars` +-/ +private def canBeSynthesized (thmVars : FVarIdSet) (fvarsFound : FVarIdSet) (instType : Expr) : MetaM Bool := do + forallTelescopeReducing instType fun xs type => type.withApp fun classFn classArgs => do + for x in xs do + unless checkTypeFVars thmVars fvarsFound (← inferType x) do return false + forallBoundedTelescope (← inferType classFn) type.getAppNumArgs fun params _ => do + for param in params, classArg in classArgs do + let paramType ← inferType param + if !paramType.isAppOf ``semiOutParam && !paramType.isAppOf ``outParam then + unless checkTypeFVars thmVars fvarsFound classArg do + return false + return true + +/-- +Auxiliary type for the `checkCoverage` function. +-/ +inductive CheckCoverageResult where + | /-- `checkCoverage` succeeded -/ + ok + | /-- + `checkCoverage` failed because some of the theorem parameters are missing, + `pos` contains their positions + -/ + missing (pos : List Nat) + +/-- +After we process a set of patterns, we obtain the set of de Bruijn indices in these patterns. +We say they are pattern variables. This function checks whether the set of pattern variables is sufficient for +instantiating the theorem with proof `thmProof`. The theorem has `numParams` parameters. +The missing parameters: +1- we may be able to infer them using type inference or type class synthesis, or +2- they are propositions, and may become hypotheses of the instantiated theorem. + +For type class instance parameters, we must check whether the free variables in class input parameters are available. +-/ +private def checkCoverage (thmProof : Expr) (numParams : Nat) (bvarsFound : Std.HashSet Nat) : MetaM CheckCoverageResult := do + if bvarsFound.size == numParams then return .ok + forallBoundedTelescope (← inferType thmProof) numParams fun xs _ => do + assert! numParams == xs.size + let patternVars := bvarsFound.toList.map fun bidx => xs[numParams - bidx - 1]!.fvarId! + -- `xs` as a `FVarIdSet`. + let thmVars : FVarIdSet := RBTree.ofList <| xs.toList.map (·.fvarId!) + -- Collect free variables occurring in `e`, and insert the ones that are in `thmVars` into `fvarsFound` + let update (fvarsFound : FVarIdSet) (e : Expr) : FVarIdSet := + (collectFVars {} e).fvarIds.foldl (init := fvarsFound) fun s fvarId => + if thmVars.contains fvarId then s.insert fvarId else s + -- Theorem variables found so far. We initialize with the variables occurring in patterns + -- Remark: fvarsFound is a subset of thmVars + let mut fvarsFound : FVarIdSet := RBTree.ofList patternVars + for patternVar in patternVars do + let type ← patternVar.getType + fvarsFound := update fvarsFound type + if fvarsFound.size == numParams then return .ok + -- Now, we keep traversing remaining variables and collecting + -- `processed` contains the variables we have already processed. + let mut processed : FVarIdSet := RBTree.ofList patternVars + let mut modified := false + repeat + modified := false + for x in xs do + let fvarId := x.fvarId! + unless processed.contains fvarId do + let xType ← inferType x + if fvarsFound.contains fvarId then + -- Collect free vars in `x`s type and mark as processed + fvarsFound := update fvarsFound xType + processed := processed.insert fvarId + modified := true + else if (← isProp xType) then + -- If `x` is a proposition, and all theorem variables in `x`s type have already been found + -- add it to `fvarsFound` and mark it as processed. + if checkTypeFVars thmVars fvarsFound xType then + fvarsFound := fvarsFound.insert fvarId + processed := processed.insert fvarId + modified := true + else if (← fvarId.getDecl).binderInfo matches .instImplicit then + -- If `x` is instance implicit, check whether + -- we have found all free variables needed to synthesize instance + if (← canBeSynthesized thmVars fvarsFound xType) then + fvarsFound := fvarsFound.insert fvarId + fvarsFound := update fvarsFound xType + processed := processed.insert fvarId + modified := true + if fvarsFound.size == numParams then + return .ok + if !modified then + break + let mut pos := #[] + for h : i in [:xs.size] do + let fvarId := xs[i].fvarId! + unless fvarsFound.contains fvarId do + pos := pos.push i + return .missing pos.toList + +/-- +Given a theorem with proof `proof` and `numParams` parameters, returns a message +containing the parameters at positions `paramPos`. +-/ +private def ppParamsAt (proof : Expr) (numParams : Nat) (paramPos : List Nat) : MetaM MessageData := do + forallBoundedTelescope (← inferType proof) numParams fun xs _ => do + let mut msg := m!"" + let mut first := true + for h : i in [:xs.size] do + if paramPos.contains i then + let x := xs[i] + if first then first := false else msg := msg ++ "\n" + msg := msg ++ m!"{x} : {← inferType x}" + addMessageContextFull msg + +private def logPatternWhen (showInfo : Bool) (origin : Origin) (patterns : List Expr) : MetaM Unit := do + if showInfo then + logInfo m!"{← origin.pp}: {patterns.map ppPattern}" + +/-- +Creates an E-matching theorem for a theorem with proof `proof`, `numParams` parameters, and the given set of patterns. +Pattern variables are represented using de Bruijn indices. +-/ +def mkEMatchTheoremCore (origin : Origin) (levelParams : Array Name) (numParams : Nat) (proof : Expr) + (patterns : List Expr) (kind : EMatchTheoremKind) (showInfo := false) : MetaM EMatchTheorem := do + -- the patterns have already been selected, there is no point in using priorities here + let (patterns, symbols, bvarFound) ← NormalizePattern.main patterns (← getGlobalSymbolPriorities) (minPrio := 1) + if symbols.isEmpty then + throwError "invalid pattern for `{← origin.pp}`{indentD (patterns.map ppPattern)}\nthe pattern does not contain constant symbols for indexing" + trace[grind.ematch.pattern] "{← origin.pp}: {patterns.map ppPattern}" + if let .missing pos ← checkCoverage proof numParams bvarFound then + let pats : MessageData := m!"{patterns.map ppPattern}" + throwError "invalid pattern(s) for `{← origin.pp}`{indentD pats}\nthe following theorem parameters cannot be instantiated:{indentD (← ppParamsAt proof numParams pos)}" + logPatternWhen showInfo origin patterns + return { + proof, patterns, numParams, symbols + levelParams, origin, kind + } + +private def getProofFor (declName : Name) : MetaM Expr := do + let info ← getConstInfo declName + -- For theorems, `isProp` has already been checked at declaration time + unless wasOriginallyTheorem (← getEnv) declName do + unless (← isProp info.type) do + throwError "invalid E-matching theorem `{declName}`, type is not a proposition" + let us := info.levelParams.map mkLevelParam + return mkConst declName us + +/-- +Creates an E-matching theorem for `declName` with `numParams` parameters, and the given set of patterns. +Pattern variables are represented using de Bruijn indices. +-/ +def mkEMatchTheorem (declName : Name) (numParams : Nat) (patterns : List Expr) (kind : EMatchTheoremKind) : MetaM EMatchTheorem := do + mkEMatchTheoremCore (.decl declName) #[] numParams (← getProofFor declName) patterns kind + +/-- +Given a theorem with proof `proof` and type of the form `∀ (a_1 ... a_n), lhs = rhs`, +creates an E-matching pattern for it using `addEMatchTheorem n [lhs]` +If `normalizePattern` is true, it applies the `grind` simplification theorems and simprocs to the pattern. +-/ +def mkEMatchEqTheoremCore (origin : Origin) (levelParams : Array Name) (proof : Expr) (normalizePattern : Bool) + (useLhs : Bool) (gen : Bool) (showInfo := false) : MetaM EMatchTheorem := do + let (numParams, patterns) ← forallTelescopeReducing (← inferEMatchProofType proof gen) fun xs type => do + let (lhs, rhs) ← match_expr type with + | Eq _ lhs rhs => pure (lhs, rhs) + | Iff lhs rhs => pure (lhs, rhs) + | HEq _ lhs _ rhs => pure (lhs, rhs) + | _ => throwError "invalid E-matching equality theorem, conclusion must be an equality{indentExpr type}" + let pat := if useLhs then lhs else rhs + trace[grind.debug.ematch.pattern] "mkEMatchEqTheoremCore: origin: {← origin.pp}, pat: {pat}, useLhs: {useLhs}" + let pat ← preprocessPattern pat normalizePattern + trace[grind.debug.ematch.pattern] "mkEMatchEqTheoremCore: after preprocessing: {pat}, {← normalize pat normConfig}" + let pats := splitWhileForbidden (pat.abstract xs) + return (xs.size, pats) + mkEMatchTheoremCore origin levelParams numParams proof patterns (if useLhs then .eqLhs gen else .eqRhs gen) (showInfo := showInfo) + +def mkEMatchEqBwdTheoremCore (origin : Origin) (levelParams : Array Name) (proof : Expr) (showInfo := false) : MetaM EMatchTheorem := do + let (numParams, patterns) ← forallTelescopeReducing (← inferEMatchProofType proof (gen := false)) fun xs type => do + let_expr f@Eq α lhs rhs := type + | throwError "invalid E-matching `←=` theorem, conclusion must be an equality{indentExpr type}" + let pat ← preprocessPattern (mkEqBwdPattern f.constLevels! α lhs rhs) + return (xs.size, [pat.abstract xs]) + mkEMatchTheoremCore origin levelParams numParams proof patterns .eqBwd (showInfo := showInfo) + +/-- +Given theorem with name `declName` and type of the form `∀ (a_1 ... a_n), lhs = rhs`, +creates an E-matching pattern for it using `addEMatchTheorem n [lhs]` + +If `normalizePattern` is true, it applies the `grind` simplification theorems and simprocs to the +pattern. +-/ +def mkEMatchEqTheorem (declName : Name) (normalizePattern := true) (useLhs : Bool := true) (gen : Bool := false) (showInfo := false) : MetaM EMatchTheorem := do + mkEMatchEqTheoremCore (.decl declName) #[] (← getProofFor declName) normalizePattern useLhs gen (showInfo := showInfo) + +/-- +Adds an E-matching theorem to the environment. +See `mkEMatchTheorem`. +-/ +def addEMatchTheorem (declName : Name) (numParams : Nat) (patterns : List Expr) (kind : EMatchTheoremKind) : MetaM Unit := do + ematchTheoremsExt.add (← mkEMatchTheorem declName numParams patterns kind) + +/-- +Adds an E-matching equality theorem to the environment. +See `mkEMatchEqTheorem`. +-/ +def addEMatchEqTheorem (declName : Name) : MetaM Unit := do + ematchTheoremsExt.add (← mkEMatchEqTheorem declName) + +/-- Returns the E-matching theorems registered in the environment. -/ +def getEMatchTheorems : CoreM EMatchTheorems := + return ematchTheoremsExt.getState (← getEnv) + +/-- Returns the types of `xs` that are propositions. -/ +private def getPropTypes (xs : Array Expr) : MetaM (Array Expr) := + xs.filterMapM fun x => do + let type ← inferType x + if (← isProp type) then return some type else return none + +/-- State for the (pattern) `CollectorM` monad -/ +private structure Collector.State where + /-- Pattern found so far. -/ + patterns : Array Expr := #[] + done : Bool := false + +private structure Collector.Context where + proof : Expr + xs : Array Expr + +/-- Monad for collecting patterns for a theorem. -/ +private abbrev CollectorM := ReaderT Collector.Context $ StateRefT Collector.State NormalizePattern.M + +/-- Similar to `getPatternFn?`, but operates on expressions that do not contain loose de Bruijn variables. -/ +private def isPatternFnCandidate (f : Expr) : CollectorM Bool := do + match f with + | .const declName _ => NormalizePattern.isCandidateSymbol declName (root := true) + | .fvar .. => return !(← read).xs.contains f + | _ => return false + +private def addNewPattern (p : Expr) : CollectorM Unit := do + trace[grind.debug.ematch.pattern] "found pattern: {ppPattern p}" + let bvarsFound := (← getThe NormalizePattern.State).bvarsFound + let done := (← checkCoverage (← read).proof (← read).xs.size bvarsFound) matches .ok + if done then + trace[grind.debug.ematch.pattern] "found full coverage" + modify fun s => { s with patterns := s.patterns.push p, done } + +/-- Collect the pattern (i.e., de Bruijn) variables in the given pattern. -/ +private def collectPatternBVars (p : Expr) : List Nat := + go p |>.run [] |>.2 +where + go (e : Expr) : StateM (List Nat) Unit := do + match e with + | .app f a => go f; go a + | .mdata _ b => go b + | .bvar idx => modify fun s => if s.contains idx then s else idx :: s + | _ => return () + +private def diff (s : List Nat) (found : Std.HashSet Nat) : List Nat := + if found.isEmpty then s else s.filter fun x => !found.contains x + +/-- +Returns `true` if pattern `p` contains a child `c` such that +1- `p` and `c` have the same new pattern variables. We say a pattern variable is new if it is not in `alreadyFound`. +2- `c` is not a support argument. See `NormalizePattern.getPatternSupportMask` for definition. +3- `c` is not an offset pattern. +4- `c` is not a bound variable. +-/ +private def hasChildWithSameNewBVars (p : Expr) + (argKinds : Array NormalizePattern.PatternArgKind) (alreadyFound : Std.HashSet Nat) : CoreM Bool := do + let s := diff (collectPatternBVars p) alreadyFound + for arg in p.getAppArgs, argKind in argKinds do + unless argKind.isSupport do + unless arg.isBVar do + unless isOffsetPattern? arg |>.isSome do + let sArg := diff (collectPatternBVars arg) alreadyFound + if s ⊆ sArg then + return true + return false + +private partial def collect (e : Expr) : CollectorM Unit := do + if (← get).done then return () + match e with + | .app .. => + trace[grind.debug.ematch.pattern] "collect: {e}" + let f := e.getAppFn + let argKinds ← NormalizePattern.getPatternArgKinds f e.getAppNumArgs + if (← isPatternFnCandidate f) then + let saved ← getThe NormalizePattern.State + try + trace[grind.debug.ematch.pattern] "candidate: {e}" + let p := e.abstract (← read).xs + unless p.hasLooseBVars do + trace[grind.debug.ematch.pattern] "skip, does not contain pattern variables" + return () + let p ← NormalizePattern.normalizePattern p + if saved.bvarsFound.size < (← getThe NormalizePattern.State).bvarsFound.size then + unless (← hasChildWithSameNewBVars p argKinds saved.bvarsFound) do + addNewPattern p + return () + trace[grind.debug.ematch.pattern] "skip, no new variables covered" + -- restore state and continue search + set saved + catch ex => + trace[grind.debug.ematch.pattern] "skip, exception during normalization{indentD ex.toMessageData}" + -- restore state and continue search + set saved + let args := e.getAppArgs + for arg in args, argKind in argKinds do + trace[grind.debug.ematch.pattern] "arg: {arg}, support: {argKind.isSupport}" + unless argKind.isSupport do + collect arg + | .forallE _ d b _ => + if (← pure e.isArrow <&&> isProp d <&&> isProp b) then + collect d + collect b + | _ => return () + +private def collectPatterns? (proof : Expr) (xs : Array Expr) (searchPlaces : Array Expr) (symPrios : SymbolPriorities) (minPrio : Nat) + : MetaM (Option (List Expr × List HeadIndex)) := do + let go : CollectorM (Option (List Expr)) := do + for place in searchPlaces do + trace[grind.debug.ematch.pattern] "place: {place}" + let place ← preprocessPattern place + collect place + if (← get).done then + return some ((← get).patterns.toList) + return none + let (some ps, s) ← go { proof, xs } |>.run' {} { symPrios, minPrio } |>.run {} + | return none + return some (ps, s.symbols.toList) + +/-- +Tries to find a ground pattern to activate the theorem. +This is used for theorems such as `theorem evenZ : Even 0`. +This function is only used if `collectPatterns?` returns `none`. + +Remark: only symbols with priority `>= minPrio` are considered. +-/ +private partial def collectGroundPattern? (proof : Expr) (xs : Array Expr) (searchPlaces : Array Expr) (symPrios: SymbolPriorities) (minPrio : Nat) + : MetaM (Option (Expr × List HeadIndex)) := do + unless (← checkCoverage proof xs.size {}) matches .ok do + return none + let go? : CollectorM (Option Expr) := do + for place in searchPlaces do + let place ← preprocessPattern place + if let some r ← visit? place then + return r + return none + let (some p, s) ← go? { proof, xs } |>.run' {} { symPrios, minPrio } |>.run {} + | return none + return some (p, s.symbols.toList) +where + visit? (e : Expr) : CollectorM (Option Expr) := do + match e with + | .app .. => + let f := e.getAppFn + if (← isPatternFnCandidate f) then + let e ← NormalizePattern.normalizePattern e + return some e + else + let args := e.getAppArgs + for arg in args, argKind in (← NormalizePattern.getPatternArgKinds f args.size) do + unless argKind.isSupport do + if let some r ← visit? arg then + return r + return none + | .forallE _ d b _ => + if (← pure e.isArrow <&&> isProp d <&&> isProp b) then + if let some d ← visit? d then return d + visit? b + else + return none + | _ => return none + +def EMatchTheoremKind.gen : EMatchTheoremKind → Bool + | .eqLhs gen => gen + | .eqRhs gen => gen + | .eqBoth gen => gen + | .default gen => gen + | .bwd gen => gen + | .eqBwd | .fwd | .rightLeft + | .leftRight | .user => false + +private def collectUsedPriorities (prios : SymbolPriorities) (searchPlaces : Array Expr) : Array Nat := Id.run do + let mut s : Std.HashSet Nat := {} + for place in searchPlaces do + s := place.foldConsts (init := s) fun declName s => + let prio := prios.getPrio declName + if prio > 0 then s.insert prio else s + let r := s.toArray + if r.isEmpty then + return #[eval_prio default] + else + return r.qsort fun p₁ p₂ => p₁ > p₂ + +/-- +Creates an E-match theorem using the given proof and kind. +If `groundPatterns` is `true`, it accepts patterns without pattern variables. This is useful for +theorems such as `theorem evenZ : Even 0`. For local theorems, we use `groundPatterns := false` +since the theorem is already in the `grind` state and there is nothing to be instantiated. +-/ +def mkEMatchTheoremWithKind? + (origin : Origin) (levelParams : Array Name) (proof : Expr) (kind : EMatchTheoremKind) (symPrios : SymbolPriorities) + (groundPatterns := true) (showInfo := false) : MetaM (Option EMatchTheorem) := do + match kind with + | .eqLhs gen => + return (← mkEMatchEqTheoremCore origin levelParams proof (normalizePattern := true) (useLhs := true) (gen := gen) (showInfo := showInfo)) + | .eqRhs gen => + return (← mkEMatchEqTheoremCore origin levelParams proof (normalizePattern := true) (useLhs := false) (gen := gen) (showInfo := showInfo)) + | .eqBwd => + return (← mkEMatchEqBwdTheoremCore origin levelParams proof (showInfo := showInfo)) + | _ => + pure () + let type ← inferEMatchProofType proof kind.gen + /- + Remark: we should not use `forallTelescopeReducing` (with default reducibility) here + because it may unfold a definition/abstraction, and then select a suboptimal pattern. + Here is an example. Suppose we have + ``` + def State.le (σ₁ σ₂ : State) : Prop := ∀ ⦃x : Var⦄ ⦃v : Val⦄, σ₁.find? x = some v → σ₂.find? x = some v + + infix:50 " ≼ " => State.le + ``` + Then, we write the theorem + ``` + @[grind] theorem State.join_le_left (σ₁ σ₂ : State) : σ₁.join σ₂ ≼ σ₁ := by + ``` + We do not want `State.le` to be unfolded and the abstraction exposed. + + That said, we must still reduce `[reducible]` definitions since `grind` unfolds them. + -/ + withReducible <| forallTelescopeReducing type fun xs type => withDefault do + let searchPlaces ← match kind with + | .fwd => + let ps ← getPropTypes xs + if ps.isEmpty then + throwError "invalid `grind` forward theorem, theorem `{← origin.pp}` does not have propositional hypotheses" + pure ps + | .bwd _ => pure #[type] + | .leftRight => pure <| (← getPropTypes xs).push type + | .rightLeft => pure <| #[type] ++ (← getPropTypes xs).reverse + | .default _ => pure <| #[type] ++ (← getPropTypes xs) + | _ => unreachable! + go xs searchPlaces +where + collect (xs : Array Expr) (searchPlaces : Array Expr) : MetaM (Option (List Expr × List HeadIndex)) := do + let prios := collectUsedPriorities symPrios searchPlaces + for minPrio in prios do + if let some r ← collectPatterns? proof xs searchPlaces symPrios minPrio then + return some r + else if groundPatterns then + if let some (pattern, symbols) ← collectGroundPattern? proof xs searchPlaces symPrios minPrio then + return some ([pattern], symbols) + return none + + go (xs : Array Expr) (searchPlaces : Array Expr) : MetaM (Option EMatchTheorem) := do + let some (patterns, symbols) ← collect xs searchPlaces | return none + let numParams := xs.size + trace[grind.ematch.pattern] "{← origin.pp}: {patterns.map ppPattern}" + logPatternWhen showInfo origin patterns + return some { + proof, patterns, numParams, symbols + levelParams, origin, kind + } + +def mkEMatchTheoremForDecl (declName : Name) (thmKind : EMatchTheoremKind) (prios : SymbolPriorities) (showInfo := false) : MetaM EMatchTheorem := do + let some thm ← mkEMatchTheoremWithKind? (.decl declName) #[] (← getProofFor declName) thmKind prios (showInfo := showInfo) + | throwError "`@{thmKind.toAttribute} theorem {declName}` {thmKind.explainFailure}, consider using different options or the `grind_pattern` command" + return thm + +def mkEMatchEqTheoremsForDef? (declName : Name) (showInfo := false) : MetaM (Option (Array EMatchTheorem)) := do + let some eqns ← getEqnsFor? declName | return none + eqns.mapM fun eqn => do + mkEMatchEqTheorem eqn (normalizePattern := true) (showInfo := showInfo) + +private def addGrindEqAttr (declName : Name) (attrKind : AttributeKind) (thmKind : EMatchTheoremKind) (useLhs := true) (showInfo := false) : MetaM Unit := do + if wasOriginallyTheorem (← getEnv) declName then + ematchTheoremsExt.add (← mkEMatchEqTheorem declName (normalizePattern := true) (useLhs := useLhs) (gen := thmKind.gen) (showInfo := showInfo)) attrKind + else if let some thms ← mkEMatchEqTheoremsForDef? declName (showInfo := showInfo) then + unless useLhs do + throwError "`{declName}` is a definition, you must only use the left-hand side for extracting patterns" + thms.forM (ematchTheoremsExt.add · attrKind) + else + throwError s!"`{thmKind.toAttribute}` attribute can only be applied to equational theorems or function definitions" + +def EMatchTheorems.eraseDecl (s : EMatchTheorems) (declName : Name) : MetaM EMatchTheorems := do + let throwErr {α} : MetaM α := + throwError "`{declName}` is not marked with the `[grind]` attribute" + if !wasOriginallyTheorem (← getEnv) declName then + if let some eqns ← getEqnsFor? declName then + let s := ematchTheoremsExt.getState (← getEnv) + unless eqns.all fun eqn => s.contains (.decl eqn) do + throwErr + return eqns.foldl (init := s) fun s eqn => s.erase (.decl eqn) + else + throwErr + else + unless ematchTheoremsExt.getState (← getEnv) |>.contains (.decl declName) do + throwErr + return s.erase <| .decl declName + +def addEMatchAttr (declName : Name) (attrKind : AttributeKind) (thmKind : EMatchTheoremKind) (prios : SymbolPriorities) (showInfo := false) : MetaM Unit := do + match thmKind with + | .eqLhs _ => + addGrindEqAttr declName attrKind thmKind (useLhs := true) (showInfo := showInfo) + | .eqRhs _ => + addGrindEqAttr declName attrKind thmKind (useLhs := false) (showInfo := showInfo) + | .eqBoth _ => + addGrindEqAttr declName attrKind thmKind (useLhs := true) (showInfo := showInfo) + addGrindEqAttr declName attrKind thmKind (useLhs := false) (showInfo := showInfo) + | _ => + let info ← getConstInfo declName + if !wasOriginallyTheorem (← getEnv) declName && !info.isCtor && !info.isAxiom then + addGrindEqAttr declName attrKind thmKind (showInfo := showInfo) + else + let thm ← mkEMatchTheoremForDecl declName thmKind prios (showInfo := showInfo) + ematchTheoremsExt.add thm attrKind + +def eraseEMatchAttr (declName : Name) : MetaM Unit := do + /- + Remark: consider the following example + ``` + attribute [grind] foo -- ok + attribute [-grind] foo.eqn_2 -- ok + attribute [-grind] foo -- error + ``` + One may argue that the correct behavior should be + ``` + attribute [grind] foo -- ok + attribute [-grind] foo.eqn_2 -- error + attribute [-grind] foo -- ok + ``` + -/ + let s := ematchTheoremsExt.getState (← getEnv) + let s ← s.eraseDecl declName + modifyEnv fun env => ematchTheoremsExt.modifyState env fun _ => s + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EqResolution.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EqResolution.lean new file mode 100644 index 0000000000000000000000000000000000000000..fe954f3025f6b637d66487a4e1f399625d9ee2ff --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/EqResolution.lean @@ -0,0 +1,133 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.AppBuilder +import Lean.Meta.MatchUtil +import Lean.Util.ForEachExpr + +namespace Lean.Meta.Grind +/-! A basic "equality resolution" procedure. -/ + +/-- +Similar to `forallMetaTelescopeReducing`, but if `prop` resulting type is of the form `¬ p`, it will "convert" it so `p → False`, and +will return a hypothesis for it and return `False` as the resulting type. +-/ +private def forallMetaTelescopeReducingAndUnfoldingNot (prop : Expr) : MetaM (Array Expr × Expr) := do + let (ms, _, type) ← forallMetaTelescopeReducing prop + if let some p ← matchNot? type then + let m ← mkFreshExprMVar p + return (ms.push m, mkConst ``False) + return (ms, type) + +structure TopSort.State where + tempMark : Std.HashSet Expr := {} + permMark : Std.HashSet Expr := {} + result : Array Expr := #[] + +abbrev TopSortM := OptionT $ StateT TopSort.State MetaM + +/-- +Sorts metavariables `ms` using topological sort. +There is an "edge" from `m₁` to `m₂` if type of `m₁` contains `m₂`. +We use this function to ensure that after applying equality resolution to +``` +∀ x : Nat, p x a → ∀ y : Nat, p y b → x = y → False +``` +we produce +``` +∀ y, p y a → p y b → False +``` +instead of +``` +p ?y a → ∀ y, p y b → False +``` +Recall that in equality resolution we create a meta-variable for each hypothesis. +Thus, we initially have +``` +?x : Nat, ?h₁ : p ?x a, ?y : Nat, ?h₂ : p ?y b, ?h₃ : ?x = ?y +``` +Then, we resolve `?h₃ : ?x = ?y` as `?y := ?x` and `?h₃ := Eq.refl ?y`. +But `?h₁` occurs before `?y`. We use topological sort to address this situation. +If a cycle is detected, it returns `none`. +-/ +private partial def topsortMVars? (ms : Array Expr) : MetaM (Option (Array Expr)) := do + let (some _, s) ← go.run.run {} | return none + return some s.result +where + go : TopSortM Unit := do + for m in ms do + visit m + + visit (m : Expr) : TopSortM Unit := do + if (← get).permMark.contains m then + return () + if (← get).tempMark.contains m then + failure + modify fun s => { s with tempMark := s.tempMark.insert m } + visitTypeOf m + modify fun s => { s with + result := s.result.push m + permMark := s.permMark.insert m + } + + visitTypeOf (m : Expr) : TopSortM Unit := do + let type ← instantiateMVars (← inferType m) + type.forEach' fun e => do + if e.hasExprMVar then + if e.isMVar && ms.contains e then + visit e + return true + else + return false + +private def eqResCore (prop proof : Expr) : MetaM (Option (Expr × Expr)) := withNewMCtxDepth do + /- + We use `forallMetaTelescopeReducingAndUnfoldingNot` because we want to treat + ``` + ∀ x, ¬ f x = f a + ``` + as + ``` + ∀ x, f x = f a → False + ``` + recall that `¬` is not reducible. + -/ + let (ms, type) ← forallMetaTelescopeReducingAndUnfoldingNot prop + if ms.isEmpty then return none + let mut progress := false + for m in ms do + let type ← inferType m + let_expr Eq _ lhs rhs ← type + | pure () + if (← isDefEq lhs rhs) then + unless (← m.mvarId!.checkedAssign (← mkEqRefl lhs)) do + return none + progress := true + unless progress do + return none + if (← ms.anyM fun m => m.mvarId!.isDelayedAssigned) then + return none + let prop' ← instantiateMVars type + let proof' ← instantiateMVars (mkAppN proof ms) + let ms ← ms.filterM fun m => return !(← m.mvarId!.isAssigned) + let some ms ← topsortMVars? ms | return none + let prop' ← mkForallFVars ms prop' (binderInfoForMVars := .default) + let proof' ← mkLambdaFVars ms proof' + return some (prop', proof') + +/-- +A basic "equality resolution" procedure: Given a proposition `prop` with a proof `proof`, it attempts to resolve equality hypotheses using `isDefEq`. For example, it reduces `∀ x y, f x = f (g y y) → g x y = y` to `∀ y, g (g y y) y = y`, and `∀ (x : Nat), f x ≠ f a` to `False`. +If successful, the result is a pair `(prop', proof)`, where `prop'` is the simplified proposition, +and `proof : prop → prop'` +-/ +def eqResolution (prop : Expr) : MetaM (Option (Expr × Expr)) := + withLocalDeclD `h prop fun h => do + let some (prop', proof') ← eqResCore prop h + | return none + let proof' ← mkLambdaFVars #[h] proof' + return some (prop', proof') + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ExprPtr.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ExprPtr.lean new file mode 100644 index 0000000000000000000000000000000000000000..5d40276b9056348e38385231904f79ace39b2caa --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ExprPtr.lean @@ -0,0 +1,33 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Expr + +namespace Lean.Meta.Grind + +@[inline] def isSameExpr (a b : Expr) : Bool := + -- It is safe to use pointer equality because we hashcons all expressions + -- inserted into the E-graph + unsafe ptrEq a b + +/-- +Key for the `ENodeMap`, `ParentMap`, and other maps and sets. +We use pointer addresses and rely on the fact all internalized expressions +have been hash-consed, i.e., we have applied `shareCommon`. +-/ +structure ExprPtr where + expr : Expr + +abbrev hashPtrExpr (e : Expr) : UInt64 := + unsafe (ptrAddrUnsafe e >>> 3).toUInt64 + +instance : Hashable ExprPtr where + hash k := hashPtrExpr k.expr + +instance : BEq ExprPtr where + beq k₁ k₂ := isSameExpr k₁.expr k₂.expr + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Ext.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Ext.lean new file mode 100644 index 0000000000000000000000000000000000000000..e98dda44f8103581c156a6a362faa6570add093b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Ext.lean @@ -0,0 +1,41 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types + +namespace Lean.Meta.Grind +/-! Extensionality theorems support. -/ + +def instantiateExtTheorem (thm : Ext.ExtTheorem) (e : Expr) : GoalM Unit := withNewMCtxDepth do + unless (← getGeneration e) < (← getMaxGeneration) do return () + let c ← mkConstWithFreshMVarLevels thm.declName + let (mvars, bis, type) ← withDefault <| forallMetaTelescopeReducing (← inferType c) + unless (← isDefEq e type) do + reportIssue! "failed to apply extensionality theorem `{thm.declName}` for {indentExpr e}\nis not definitionally equal to{indentExpr type}" + return () + -- Instantiate type class instances + for mvar in mvars, bi in bis do + if bi.isInstImplicit && !(← mvar.mvarId!.isAssigned) then + let type ← inferType mvar + unless (← synthesizeInstanceAndAssign mvar type) do + reportIssue! "failed to synthesize instance when instantiating extensionality theorem `{thm.declName}` for {indentExpr e}" + return () + -- Remark: `proof c mvars` has type `e` + let proof ← instantiateMVars (mkAppN c mvars) + -- `e` is equal to `False` + let eEqFalse ← mkEqFalseProof e + -- So, we use `Eq.mp` to build a `proof` of `False` + let proof := mkApp4 (mkConst ``Eq.mp [levelZero]) e (← getFalseExpr) eEqFalse proof + let mvars ← mvars.filterM fun mvar => return !(← mvar.mvarId!.isAssigned) + let proof' ← instantiateMVars (← mkLambdaFVars mvars proof) + let prop' ← inferType proof' + if proof'.hasMVar || prop'.hasMVar then + reportIssue! "failed to apply extensionality theorem `{thm.declName}` for {indentExpr e}\nresulting terms contain metavariables" + return () + trace[grind.ext] "{thm.declName}: {prop'}" + addNewRawFact proof' prop' ((← getGeneration e) + 1) (.ext thm.declName) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ExtAttr.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ExtAttr.lean new file mode 100644 index 0000000000000000000000000000000000000000..46512f9582596d1c8051fc43379cbce7623aa6b7 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ExtAttr.lean @@ -0,0 +1,44 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Ext + +namespace Lean.Meta.Grind +/-! Grind extensionality attribute to mark which `[ext]` theorems should be used. -/ + +/-- Extensionality theorems that can be used by `grind` -/ +abbrev ExtTheorems := PHashSet Name + +builtin_initialize extTheoremsExt : SimpleScopedEnvExtension Name ExtTheorems ← + registerSimpleScopedEnvExtension { + initial := {} + addEntry := fun s declName => s.insert declName + } + +def validateExtAttr (declName : Name) : CoreM Unit := do + if !(← Ext.isExtTheorem declName) then + if !(isStructure (← getEnv) declName) then + throwError "invalid `[grind ext]`, `{declName}` is neither tagged with `[ext]` nor is a structure" + +def addExtAttr (declName : Name) (attrKind : AttributeKind) : CoreM Unit := do + validateExtAttr declName + extTheoremsExt.add declName attrKind + +private def eraseDecl (s : ExtTheorems) (declName : Name) : CoreM ExtTheorems := do + if s.contains declName then + return s.erase declName + else + throwError "`{declName}` is not marked with the `[grind ext]` attribute" + +def eraseExtAttr (declName : Name) : CoreM Unit := do + let s := extTheoremsExt.getState (← getEnv) + let s ← eraseDecl s declName + modifyEnv fun env => extTheoremsExt.modifyState env fun _ => s + +def isExtTheorem (declName : Name) : CoreM Bool := do + return extTheoremsExt.getState (← getEnv) |>.contains declName + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ForallProp.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ForallProp.lean new file mode 100644 index 0000000000000000000000000000000000000000..1f9dacf525b7fc5109847891307aab3ad6bf2f54 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/ForallProp.lean @@ -0,0 +1,126 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Lemmas +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Internalize +import Lean.Meta.Tactic.Grind.Simp +import Lean.Meta.Tactic.Grind.EqResolution + +namespace Lean.Meta.Grind +/-- +If `parent` is a projection-application `proj_i c`, +check whether the root of the equivalence class containing `c` is a constructor-application `ctor ... a_i ...`. +If so, internalize the term `proj_i (ctor ... a_i ...)` and add the equality `proj_i (ctor ... a_i ...) = a_i`. +-/ +def propagateForallPropUp (e : Expr) : GoalM Unit := do + let .forallE n p q bi := e | return () + trace_goal[grind.debug.forallPropagator] "{e}" + if !q.hasLooseBVars then + propagateImpliesUp p q + else + unless (← isEqTrue p) do return + trace_goal[grind.debug.forallPropagator] "isEqTrue, {e}" + let h₁ ← mkEqTrueProof p + let qh₁ := q.instantiate1 (mkOfEqTrueCore p h₁) + let r ← preprocess qh₁ + let q := mkLambda n bi p q + let q' := r.expr + internalize q' (← getGeneration e) + trace_goal[grind.debug.forallPropagator] "q': {q'} for{indentExpr e}" + let h₂ ← r.getProof + let h := mkApp5 (mkConst ``Lean.Grind.forall_propagator) p q q' h₁ h₂ + pushEq e q' h +where + propagateImpliesUp (a b : Expr) : GoalM Unit := do + unless (← alreadyInternalized b) do return () + if (← isEqFalse a <&&> isProp b) then + -- a = False → (a → b) = True + pushEqTrue e <| mkApp3 (mkConst ``Grind.imp_eq_of_eq_false_left) a b (← mkEqFalseProof a) + else if (← isEqTrue a <&&> isProp b) then + -- a = True → (a → b) = b + pushEq e b <| mkApp3 (mkConst ``Grind.imp_eq_of_eq_true_left) a b (← mkEqTrueProof a) + else if (← isEqTrue b <&&> isProp a) then + -- b = True → (a → b) = True + pushEqTrue e <| mkApp3 (mkConst ``Grind.imp_eq_of_eq_true_right) a b (← mkEqTrueProof b) + else if (← isEqFalse b <&&> isEqTrue e <&&> isProp a) then + -- (a → b) = True → b = False → a = False + pushEqFalse a <| mkApp4 (mkConst ``Grind.eq_false_of_imp_eq_true) a b (← mkEqTrueProof e) (← mkEqFalseProof b) + +private def isEqTrueHyp? (proof : Expr) : Option FVarId := Id.run do + let_expr eq_true _ p := proof | return none + let .fvar fvarId := p | return none + return some fvarId + +/-- Similar to `mkEMatchTheoremWithKind?`, but swallow any exceptions. -/ +private def mkEMatchTheoremWithKind'? (origin : Origin) (proof : Expr) (kind : EMatchTheoremKind) (prios : SymbolPriorities) : MetaM (Option EMatchTheorem) := do + try + mkEMatchTheoremWithKind? origin #[] proof kind prios (groundPatterns := false) + catch _ => + return none + +private def addLocalEMatchTheorems (e : Expr) : GoalM Unit := do + let proof ← mkEqTrueProof e + let origin ← if let some fvarId := isEqTrueHyp? proof then + pure <| .fvar fvarId + else + let idx ← modifyGet fun s => (s.ematch.nextThmIdx, { s with ematch.nextThmIdx := s.ematch.nextThmIdx + 1 }) + pure <| .local ((`local).appendIndexAfter idx) + let proof := mkOfEqTrueCore e proof + let size := (← get).ematch.newThms.size + let gen ← getGeneration e + -- TODO: we should have a flag for collecting all unary patterns in a local theorem + if let some thm ← mkEMatchTheoremWithKind'? origin proof .leftRight (← getSymbolPriorities) then + activateTheorem thm gen + if let some thm ← mkEMatchTheoremWithKind'? origin proof .rightLeft (← getSymbolPriorities) then + activateTheorem thm gen + if (← get).ematch.newThms.size == size then + if let some thm ← mkEMatchTheoremWithKind'? origin proof (.default false) (← getSymbolPriorities) then + activateTheorem thm gen + if (← get).ematch.newThms.size == size then + reportIssue! "failed to create E-match local theorem for{indentExpr e}" + +def propagateForallPropDown (e : Expr) : GoalM Unit := do + let .forallE n a b bi := e | return () + if (← isEqFalse e) then + if b.hasLooseBVars || !(← isProp a) then + let α := a + let p := b + -- `e` is of the form `∀ x : α, p x` + -- Add fact `∃ x : α, ¬ p x` + let u ← getLevel α + let prop := mkApp2 (mkConst ``Exists [u]) α (mkLambda n bi α (mkNot p)) + let proof := mkApp3 (mkConst ``Grind.of_forall_eq_false [u]) α (mkLambda n bi α p) (← mkEqFalseProof e) + addNewRawFact proof prop (← getGeneration e) (.forallProp e) + else + let h ← mkEqFalseProof e + pushEqTrue a <| mkApp3 (mkConst ``Grind.eq_true_of_imp_eq_false) a b h + pushEqFalse b <| mkApp3 (mkConst ``Grind.eq_false_of_imp_eq_false) a b h + else if (← isEqTrue e) then + if let some (e', h') ← eqResolution e then + trace_goal[grind.eqResolution] "{e}, {e'}" + let h := mkOfEqTrueCore e (← mkEqTrueProof e) + let h' := mkApp h' h + addNewRawFact h' e' (← getGeneration e) (.forallProp e) + else + if b.hasLooseBVars then + addLocalEMatchTheorems e + else + unless (← alreadyInternalized b) do return () + if (← isEqFalse b <&&> isProp a) then + -- (a → b) = True → b = False → a = False + pushEqFalse a <| mkApp4 (mkConst ``Grind.eq_false_of_imp_eq_true) a b (← mkEqTrueProof e) (← mkEqFalseProof b) + +builtin_grind_propagator propagateExistsDown ↓Exists := fun e => do + if (← isEqFalse e) then + let_expr f@Exists α p := e | return () + let u := f.constLevels! + let notP := mkApp (mkConst ``Not) (mkApp p (.bvar 0) |>.headBeta) + let prop := mkForall `x .default α notP + let proof := mkApp3 (mkConst ``forall_not_of_not_exists u) α p (mkOfEqFalseCore e (← mkEqFalseProof e)) + addNewRawFact proof prop (← getGeneration e) (.existsProp e) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Injection.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Injection.lean new file mode 100644 index 0000000000000000000000000000000000000000..65c9778621388da6a95bcc418223e84f1a946cd6 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Injection.lean @@ -0,0 +1,39 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.CtorRecognizer +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Clear + +namespace Lean.Meta.Grind +/-- +The `grind` tactic includes an auxiliary `injection?` tactic that is not intended for direct use by users. +This tactic is automatically applied when introducing non-dependent propositions. +It differs from the user-facing Lean `injection` tactic in the following ways: + +- It does not introduce new propositions. Instead, the `grind` tactic preprocessor is responsible for introducing them. +- It assumes that `fvarId` is the latest local declaration in the current goal. +- It does not handle cases where the constructors are different because the simplifier already reduces such equalities to `False`. +- It does not have support for heterogeneous equality. Recall that the simplifier already reduces them to `Eq` if + the types are definitionally equal. +-/ +def injection? (mvarId : MVarId) (fvarId : FVarId) : MetaM (Option MVarId) := mvarId.withContext do + let decl ← fvarId.getDecl + let_expr Eq _ a b := decl.type | return none + match (← isConstructorAppCore? a), (← isConstructorAppCore? b) with + | some aCtor, some bCtor => + if aCtor.name != bCtor.name then return none -- Simplifier handles this case + let target ← mvarId.getType + let val ← mkNoConfusion target (mkFVar fvarId) + let .forallE _ targetNew _ _ ← whnfD (← inferType val) | return none + let targetNew := targetNew.headBeta + let tag ← mvarId.getTag + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew tag + mvarId.assign (mkApp val mvarNew) + return some (← mvarNew.mvarId!.clear fvarId) + | _, _ => return none + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Internalize.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Internalize.lean new file mode 100644 index 0000000000000000000000000000000000000000..4c588f13ad4f4c9fd8eb74279c6164966665ce51 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Internalize.lean @@ -0,0 +1,422 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Init.Grind.Lemmas +import Lean.Meta.LitValues +import Lean.Meta.Match.MatcherInfo +import Lean.Meta.Match.MatchEqsExt +import Lean.Meta.Match.MatchEqs +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Util +import Lean.Meta.Tactic.Grind.Canon +import Lean.Meta.Tactic.Grind.Beta +import Lean.Meta.Tactic.Grind.MatchCond +import Lean.Meta.Tactic.Grind.Arith.Internalize + +namespace Lean.Meta.Grind +/-- Adds `e` to congruence table. -/ +def addCongrTable (e : Expr) : GoalM Unit := do + if let some { e := e' } := (← get).congrTable.find? { e } then + -- `f` and `g` must have the same type. + -- See paper: Congruence Closure in Intensional Type Theory + if e.isApp then + let f := e.getAppFn + let g := e'.getAppFn + unless isSameExpr f g do + unless (← hasSameType f g) do + reportIssue! "found congruence between{indentExpr e}\nand{indentExpr e'}\nbut functions have different types" + return () + trace_goal[grind.debug.congr] "{e} = {e'}" + pushEqHEq e e' congrPlaceholderProof + let node ← getENode e + setENode e { node with congr := e' } + else + modify fun s => { s with congrTable := s.congrTable.insert { e } } + +/-- +Given an application `e` of the form `f a_1 ... a_n`, +adds entry `f ↦ e` to `appMap`. Recall that `appMap` is a multi-map. +-/ +private def updateAppMap (e : Expr) : GoalM Unit := do + let key := e.toHeadIndex + trace_goal[grind.debug.appMap] "{e} => {repr key}" + modify fun s => { s with + appMap := if let some es := s.appMap.find? key then + s.appMap.insert key (e :: es) + else + s.appMap.insert key [e] + } + saveAppOf key + +private def forbiddenSplitTypes := [``Eq, ``HEq, ``True, ``False] + +/-- Returns `true` if `e` is of the form `@Eq Prop a b` -/ +def isMorallyIff (e : Expr) : Bool := + let_expr Eq α _ _ := e | false + α.isProp + +private def mkDefaultSplitInfo (e : Expr) : GrindM SplitInfo := + return .default e (← readThe Context).splitSource + +private def addDefaultSplitCandidate (e : Expr) : GoalM Unit := do + addSplitCandidate (← mkDefaultSplitInfo e) + +/-- Inserts `e` into the list of case-split candidates if applicable. -/ +private def checkAndAddSplitCandidate (e : Expr) : GoalM Unit := do + match h : e with + | .app .. => + if (← getConfig).splitIte && (isIte e || isDIte e) then + addDefaultSplitCandidate e + return () + if isMorallyIff e then + addDefaultSplitCandidate e + return () + if (← getConfig).splitMatch then + if (← isMatcherApp e) then + if let .reduced _ ← reduceMatcher? e then + -- When instantiating `match`-equations, we add `match`-applications that can be reduced, + -- and consequently don't need to be split. + return () + else + addDefaultSplitCandidate e + return () + let .const declName _ := e.getAppFn | return () + if forbiddenSplitTypes.contains declName then + return () + unless (← isInductivePredicate declName) do + return () + if (← get).split.casesTypes.isSplit declName then + addDefaultSplitCandidate e + else if (← getConfig).splitIndPred then + addDefaultSplitCandidate e + | .fvar .. => + let .const declName _ := (← whnf (← inferType e)).getAppFn | return () + if (← get).split.casesTypes.isSplit declName then + addDefaultSplitCandidate e + | .forallE _ d _ _ => + let currSplitSource := (← readThe Context).splitSource + if (← getConfig).splitImp then + if (← isProp d) then + addSplitCandidate (.imp e (h ▸ rfl) currSplitSource) + else if Arith.isRelevantPred d then + -- TODO: should we keep lookahead after we implement non-chronological backtracking? + if (← getConfig).lookahead then + addLookaheadCandidate (.imp e (h ▸ rfl) currSplitSource) + -- We used to add the `split` only if `lookahead := false`, but it was counterintuitive + -- to make `grind` "stronger" by disabling a feature. + addSplitCandidate (.imp e (h ▸ rfl) currSplitSource) + | _ => pure () + +/-- +If `e` is a `cast`-like term (e.g., `cast h a`), add `e ≍ a` to the to-do list. +It could be an E-matching theorem, but we want to ensure it is always applied since +we want to rely on the fact that `cast h a` and `a` are in the same equivalence class. +-/ +private def pushCastHEqs (e : Expr) : GoalM Unit := do + match_expr e with + | f@cast α β h a => pushHEq e a (mkApp4 (mkConst ``cast_heq f.constLevels!) α β h a) + | f@Eq.rec α a motive v b h => pushHEq e v (mkApp6 (mkConst ``Grind.eqRec_heq f.constLevels!) α a motive v b h) + | f@Eq.ndrec α a motive v b h => pushHEq e v (mkApp6 (mkConst ``Grind.eqNDRec_heq f.constLevels!) α a motive v b h) + | f@Eq.recOn α a motive b h v => pushHEq e v (mkApp6 (mkConst ``Grind.eqRecOn_heq f.constLevels!) α a motive b h v) + | _ => return () + +private def mkENode' (e : Expr) (generation : Nat) : GoalM Unit := + mkENodeCore e (ctor := false) (interpreted := false) (generation := generation) + +/-- Internalizes the nested ground terms in the given pattern. -/ +private partial def internalizePattern (pattern : Expr) (generation : Nat) : GoalM Expr := do + -- Recall that it is important to ensure patterns are maximally shared since + -- we assume that in functions such as `getAppsOf` in `EMatch.lean` + go (← shareCommon pattern) +where + go (pattern : Expr) : GoalM Expr := do + if pattern.isBVar || isPatternDontCare pattern then + return pattern + else if let some e := groundPattern? pattern then + let e ← preprocessLight e + internalize e generation none + return mkGroundPattern e + else pattern.withApp fun f args => do + return mkAppN f (← args.mapM go) + +/-- Internalizes the `MatchCond` gadget. -/ +private def internalizeMatchCond (matchCond : Expr) (generation : Nat) : GoalM Unit := do + mkENode' matchCond generation + let (lhss, e') ← collectMatchCondLhssAndAbstract matchCond + lhss.forM fun lhs => do internalize lhs generation; registerParent matchCond lhs + propagateUp matchCond + internalize e' generation + trace_goal[grind.debug.matchCond.lambda] "(idx := {(← getENode e'.getAppFn).idx}) {e'.getAppFn}" + trace_goal[grind.debug.matchCond.lambda] "auxiliary application{indentExpr e'}" + pushEq matchCond e' (← mkEqRefl matchCond) + internalizeSimpleMatchCondImp +where + /-- + We say `MatchCond` is simple if its argument is an implication such as `x = 0 -> ...` + If that is the case, we also internalize the implication to ensure grind can split on the antecedents. + We added this extra case to make sure the user is not surprised by `grind` failing at + ``` + example (x y : Nat) + : 0 < match x, y with + | 0, 0 => 1 + | _, _ => x + y := by -- x or y must be greater than 0 + grind + ``` + We should try to find a better and more general approach for handling the example above. + -/ + internalizeSimpleMatchCondImp : GoalM Unit := do + let_expr Grind.MatchCond e := matchCond | return () + let .forallE _ d b _ := e | return () + if b.hasLooseBVars then return () + if (← isProp d) then + internalize e generation + pushEq matchCond e (← mkEqRefl matchCond) + +def activateTheorem (thm : EMatchTheorem) (generation : Nat) : GoalM Unit := do + -- Recall that we use the proof as part of the key for a set of instances found so far. + -- We don't want to use structural equality when comparing keys. + let proof ← shareCommon thm.proof + let thm := { thm with proof, patterns := (← thm.patterns.mapM (internalizePattern · generation)) } + trace_goal[grind.ematch] "activated `{thm.origin.key}`, {thm.patterns.map ppPattern}" + modify fun s => { s with ematch.newThms := s.ematch.newThms.push thm } + +/-- +If `Config.matchEqs` is set to `true`, and `f` is `match`-auxiliary function, +adds its equations to `newThms`. +-/ +private def addMatchEqns (f : Expr) (generation : Nat) : GoalM Unit := do + if !(← getConfig).matchEqs then return () + let .const declName _ := f | return () + if !(← isMatcher declName) then return () + if (← get).ematch.matchEqNames.contains declName then return () + modify fun s => { s with ematch.matchEqNames := s.ematch.matchEqNames.insert declName } + -- for eqn in (← Match.getEquationsFor declName).eqnNames do + for eqn in (← Match.genMatchCongrEqns declName) do + -- We disable pattern normalization to prevent the `match`-expression to be reduced. + activateTheorem (← mkEMatchEqTheorem eqn (normalizePattern := false)) generation + +private def activateTheoremPatterns (fName : Name) (generation : Nat) : GoalM Unit := do + if let some (thms, thmMap) := (← get).ematch.thmMap.retrieve? fName then + modify fun s => { s with ematch.thmMap := thmMap } + let appMap := (← get).appMap + for thm in thms do + trace_goal[grind.debug.ematch.activate] "`{fName}` => `{thm.origin.key}`" + unless (← get).ematch.thmMap.isErased thm.origin do + let symbols := thm.symbols.filter fun sym => !appMap.contains sym + let thm := { thm with symbols } + match symbols with + | [] => + trace_goal[grind.debug.ematch.activate] "`{thm.origin.key}`" + activateTheorem thm generation + | _ => + trace_goal[grind.debug.ematch.activate] "reinsert `{thm.origin.key}`" + modify fun s => { s with ematch.thmMap := s.ematch.thmMap.insert thm } + +/-- +If type of `a` is a structure and is tagged with `[grind ext]` attribute, +propagate `a = ⟨a.1, ..., a.n⟩` + +This function subsumes the `propagateUnitLike` function we used in the past. +Recall that the `propagateUnitLike` was added because `isDefEq` implements it, +and consequently the simplifier reduces terms of the form `a = ctor` to `True` using `eq_self`. +This `isDefEq` feature was negatively affecting `grind` until we added an +equivalent one here. For example, when splitting on a `match`-expression +using Unit-like types, equalities about these types were being reduced to `True` +by `simp` (i.e., in the `grind` preprocessor), and `grind` would never see +these facts. +-/ +private def propagateEtaStruct (a : Expr) (generation : Nat) : GoalM Unit := do + unless (← getConfig).etaStruct do return () + let aType ← whnf (← inferType a) + matchConstStructureLike aType.getAppFn (fun _ => return ()) fun inductVal us ctorVal => do + unless a.isAppOf ctorVal.name do + -- TODO: remove ctorVal.numFields after update stage0 + if (← isExtTheorem inductVal.name) || ctorVal.numFields == 0 then + let params := aType.getAppArgs[*...inductVal.numParams] + let mut ctorApp := mkAppN (mkConst ctorVal.name us) params + for j in [: ctorVal.numFields] do + let mut proj ← mkProjFn ctorVal us params j a + if (← isProof proj) then + proj ← markProof proj + ctorApp := mkApp ctorApp proj + ctorApp ← preprocessLight ctorApp + internalize ctorApp generation + let u ← getLevel aType + let expectedProp := mkApp3 (mkConst ``Eq [u]) aType a ctorApp + pushEq a ctorApp <| mkExpectedPropHint (mkApp2 (mkConst ``Eq.refl [u]) aType a) expectedProp + +/-- Returns `true` if we can ignore `ext` for functions occurring as arguments of a `declName`-application. -/ +private def extParentsToIgnore (declName : Name) : Bool := + declName == ``Eq || declName == ``HEq || declName == ``dite || declName == ``ite + || declName == ``Exists || declName == ``Subtype + +/-- +Given a term `arg` that occurs as the argument at position `i` of an `f`-application `parent?`, +we consider `arg` as a candidate for case-splitting. For every other argument `arg'` that also appears +at position `i` in an `f`-application and has the same type as `e`, we add the case-split candidate `arg = arg'`. + +When performing the case split, we consider the following two cases: +- `arg = arg'`, which may introduce a new congruence between the corresponding `f`-applications. +- `¬(arg = arg')`, which may trigger extensionality theorems for the type of `arg`. + +This feature enables `grind` to solve examples such as: +```lean +example (f : (Nat → Nat) → Nat) : a = b → f (fun x => a + x) = f (fun x => b + x) := by + grind +``` +-/ +private def addSplitCandidatesForExt (arg : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := do + let some parent := parent? | return () + unless parent.isApp do return () + let f := parent.getAppFn + if let .const declName _ := f then + if extParentsToIgnore declName then return () + let type ← inferType arg + -- Remark: we currently do not perform function extensionality on functions that produce a type that is not a proposition. + -- We may add an option to enable that in the future. + let u? ← typeFormerTypeLevel type + if u? != .none && u? != some .zero then return () + let mut i := parent.getAppNumArgs + let mut it := parent + repeat + if !it.isApp then return () + i := i - 1 + if isSameExpr arg it.appArg! then + found f i type parent + it := it.appFn! +where + found (f : Expr) (i : Nat) (type : Expr) (parent : Expr) : GoalM Unit := do + trace_goal[grind.debug.ext] "{f}, {i}, {arg}" + let others := (← get).split.argsAt.find? (f, i) |>.getD [] + for other in others do + if (← withDefault <| isDefEq type other.type) then + let eq := mkApp3 (mkConst ``Eq [← getLevel type]) type arg other.arg + let eq ← shareCommon eq + internalize eq generation + trace_goal[grind.ext.candidate] "{eq}" + -- We do not use lookahead here because it is too incomplete. + -- if (← getConfig).lookahead then + -- addLookaheadCandidate (.arg other.app parent i eq) + -- else + let currSplitSource := (← readThe Context).splitSource + addSplitCandidate (.arg other.app parent i eq currSplitSource) + modify fun s => { s with split.argsAt := s.split.argsAt.insert (f, i) ({ arg, type, app := parent } :: others) } + return () + +/-- Applies `addSplitCandidatesForExt` if `funext` is enabled. -/ +private def addSplitCandidatesForFunext (arg : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := do + unless (← getConfig).funext do return () + addSplitCandidatesForExt arg generation parent? + +/-- +Tries to eta-reduce the given expression. +If successful, pushes a new equality between the two terms. +-/ +private def tryEta (e : Expr) (generation : Nat) : GoalM Unit := do + let e' := e.eta + if e != e' then + let e' ← shareCommon e' + internalize e' generation + pushEq e e' (← mkEqRefl e) + +@[export lean_grind_internalize] +private partial def internalizeImpl (e : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := withIncRecDepth do + if (← alreadyInternalized e) then + trace_goal[grind.debug.internalize] "already internalized: {e}" + /- + Even if `e` has already been internalized, we must check whether it has also been internalized in + the satellite solvers. For example, suppose we have already internalized the term `f (a + 1)`. + The `1` in this term is treated as an offset for the offset term `a + 1` by the arithmetic module, and + only nodes for `a` and `a+1` are created. However, an ENode for `1` is created here. + Later, if we try to internalize `f 1`, the arithmetic module must create a node for `1`. + Otherwise, it will not be able to propagate that `a + 1 = 1` when `a = 0` + -/ + Arith.internalize e parent? + else + go + propagateEtaStruct e generation +where + go : GoalM Unit := do + trace_goal[grind.internalize] "{e}" + match e with + | .bvar .. => unreachable! + | .sort .. => return () + | .fvar .. => + mkENode' e generation + checkAndAddSplitCandidate e + | .letE .. => + mkENode' e generation + | .lam .. => + addSplitCandidatesForFunext e generation parent? + mkENode' e generation + tryEta e generation + | .forallE _ d b _ => + mkENode' e generation + internalizeImpl d generation e + registerParent e d + unless b.hasLooseBVars do + internalizeImpl b generation e + registerParent e b + addCongrTable e + if (← isProp d <&&> isProp e) then + propagateUp e + checkAndAddSplitCandidate e + | .lit .. => + mkENode e generation + | .const declName _ => + mkENode e generation + activateTheoremPatterns declName generation + | .mvar .. => + if (← reportMVarInternalization) then + reportIssue! "unexpected metavariable during internalization{indentExpr e}\n`grind` is not supposed to be used in goals containing metavariables." + mkENode' e generation + | .mdata .. => + reportIssue! "unexpected metadata found during internalization{indentExpr e}\n`grind` uses a pre-processing step that eliminates metadata" + mkENode' e generation + | .proj .. => + reportIssue! "unexpected kernel projection term during internalization{indentExpr e}\n`grind` uses a pre-processing step that folds them as projection applications, the pre-processor should have failed to fold this term" + mkENode' e generation + | .app .. => + if (← isLitValue e) then + -- We do not want to internalize the components of a literal value. + mkENode e generation + Arith.internalize e parent? + else if e.isAppOfArity ``Grind.MatchCond 1 then + internalizeMatchCond e generation + else e.withApp fun f args => do + mkENode e generation + updateAppMap e + checkAndAddSplitCandidate e + pushCastHEqs e + addMatchEqns f generation + if f.isConstOf ``Lean.Grind.nestedProof && args.size == 2 then + -- We only internalize the proposition. We can skip the proof because of + -- proof irrelevance + let c := args[0]! + internalizeImpl c generation e + registerParent e c + pushEqTrue c <| mkApp2 (mkConst ``eq_true) c args[1]! + else if f.isConstOf ``ite && args.size == 5 then + let c := args[1]! + internalizeImpl c generation e + registerParent e c + else + if let .const fName _ := f then + activateTheoremPatterns fName generation + else + internalizeImpl f generation e + registerParent e f + for h : i in [: args.size] do + let arg := args[i] + internalize arg generation e + registerParent e arg + addCongrTable e + Arith.internalize e parent? + propagateUp e + propagateBetaForNewApp e + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Intro.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Intro.lean new file mode 100644 index 0000000000000000000000000000000000000000..6427ffa0420728b7f0adf73be5d8ea4d72f07fd1 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Intro.lean @@ -0,0 +1,298 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Lemmas +import Lean.Meta.Tactic.Assert +import Lean.Meta.Tactic.Grind.Simp +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Cases +import Lean.Meta.Tactic.Grind.CasesMatch +import Lean.Meta.Tactic.Grind.Injection +import Lean.Meta.Tactic.Grind.Core +import Lean.Meta.Tactic.Grind.SearchM + +namespace Lean.Meta.Grind + +private inductive IntroResult where + | done (goal : Goal) + | newHyp (fvarId : FVarId) (goal : Goal) + | newDepHyp (goal : Goal) + | newLocal (fvarId : FVarId) (goal : Goal) + deriving Inhabited + +/-- +Returns `true` if `e` is marked with the `alreadyNorm` gadget. +See `alreadyNorm` documentation for additional details +-/ +private def isAlreadyNorm? (e : Expr) : Option Expr := + let_expr Lean.Grind.alreadyNorm c := e | none + some c + +/-- +Similar to `Grind.preprocess`, but does not simplify `e` if +`isMatchCondCandidate` (aka `Simp.isEqnThmHypothesis`) is `true`. +We added this feature because it may be coming from external sources +(e.g., manually applying an function induction principle before invoking `grind`). +-/ +private def preprocessHypothesis (e : Expr) : GoalM Simp.Result := do + if isMatchCondCandidate e then + preprocess (markAsPreMatchCond e) + else if let some c := isAlreadyNorm? e then + let c ← shareCommon (← canon c) + return { expr := c } + else + preprocess e + +/-- +Helper function for `mkCleanName`. +It ensures base name is a simple `Name` and does not have a `_` suffix +-/ +private def mkBaseName (name : Name) (type : Expr) : MetaM Name := do + if let .str _ s := name then + let pos := s.find (· == '_') + unless pos < s.endPos do + return Name.mkSimple s + let suffix := s.extract (pos+'_') s.endPos + unless suffix.isNat do + return Name.mkSimple s + let s := s.extract ⟨0⟩ pos + unless s == "" do + return Name.mkSimple s + if (← isProp type) then return `h else return `x + +private def mkCleanName (name : Name) (type : Expr) : GoalM Name := do + unless (← getConfig).clean do + return name + let mut name := name + if name.hasMacroScopes then + name := name.eraseMacroScopes + if name == `x || name == `a then + if (← isProp type) then + name := `h + if (← get).clean.used.contains name then + let base ← mkBaseName name type + let mut i := if let some i := (← get).clean.next.find? base then i else 1 + repeat + name := base.appendIndexAfter i + i := i + 1 + unless (← get).clean.used.contains name do + break + modify fun s => { s with clean.next := s.clean.next.insert base i } + modify fun s => { s with clean.used := s.clean.used.insert name } + return name + +private def intro1 : GoalM FVarId := do + let target ← (← get).mvarId.getType + let (name, type) ← match target with + | .forallE n d .. => pure (n, d) + | .letE n d .. => pure (n, d) + | _ => throwError "`grind` internal error, binder expected" + let name ← mkCleanName name type + let (fvarId, mvarId) ← (← get).mvarId.intro name + modify fun s => { s with mvarId } + return fvarId + +private partial def introNext (goal : Goal) (generation : Nat) : GrindM IntroResult := do + Prod.fst <$> GoalM.run goal do + let target ← (← get).mvarId.getType + if target.isForall then + let p := target.bindingDomain! + if !(← isProp p) then + let fvarId ← intro1 + return .newLocal fvarId (← get) + else + let mvarId := (← get).mvarId + let tag ← mvarId.getTag + let qBase := target.bindingBody! + let fvarId ← mkFreshFVarId + let fvar := mkFVar fvarId + let r ← preprocessHypothesis p + let lctx := (← getLCtx).mkLocalDecl fvarId (← mkCleanName target.bindingName! r.expr) r.expr target.bindingInfo! + let mut localInsts ← getLocalInstances + if let some className ← isClass? r.expr then + localInsts := localInsts.push { className, fvar } + match r.proof? with + | some he => + if target.isArrow then + let q := qBase + let u ← getLevel q + let mvarNew ← mkFreshExprMVarAt lctx localInsts q .syntheticOpaque tag + let mvarIdNew := mvarNew.mvarId! + mvarIdNew.withContext do + let h ← mkLambdaFVars #[fvar] mvarNew + let hNew := mkAppN (mkConst ``Lean.Grind.intro_with_eq [u]) #[p, r.expr, q, he, h] + mvarId.assign hNew + return .newHyp fvarId { (← get) with mvarId := mvarIdNew } + else + let q := mkLambda target.bindingName! target.bindingInfo! p qBase + let q' := qBase.instantiate1 (mkApp4 (mkConst ``Eq.mpr_prop) p r.expr he fvar) + let u ← getLevel q' + let mvarNew ← mkFreshExprMVarAt lctx localInsts q' .syntheticOpaque tag + let mvarIdNew := mvarNew.mvarId! + mvarIdNew.withContext do + let h ← mkLambdaFVars #[fvar] mvarNew + let hNew := mkAppN (mkConst ``Lean.Grind.intro_with_eq' [u]) #[p, r.expr, q, he, h] + mvarId.assign hNew + return .newHyp fvarId { (← get) with mvarId := mvarIdNew } + | none => + -- `p` and `p'` are definitionally equal + let q := if target.isArrow then qBase else qBase.instantiate1 (mkFVar fvarId) + let mvarNew ← mkFreshExprMVarAt lctx localInsts q .syntheticOpaque tag + let mvarIdNew := mvarNew.mvarId! + mvarIdNew.withContext do + let h ← mkLambdaFVars #[mkFVar fvarId] mvarNew + mvarId.assign h + return .newHyp fvarId { (← get) with mvarId := mvarIdNew } + else if target.isLet then + if (← getConfig).zetaDelta then + let targetNew := expandLet target #[] + let mvarId := (← get).mvarId + mvarId.withContext do + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew (← mvarId.getTag) + mvarId.assign mvarNew + introNext { (← get) with mvarId := mvarNew.mvarId! } generation + else + let fvarId ← intro1 + (← get).mvarId.withContext do + let localDecl ← fvarId.getDecl + if (← isProp localDecl.type) then + -- Add a non-dependent copy + let mvarId ← (← get).mvarId.assert (← mkCleanName `h localDecl.type) localDecl.type (mkFVar fvarId) + return .newDepHyp { (← get) with mvarId } + else + let v := (← fvarId.getDecl).value + let r ← preprocessHypothesis v + let x ← shareCommon (mkFVar fvarId) + addNewEq x r.expr (← r.getProof) generation + return .newLocal fvarId (← get) + else + return .done goal + +private def isEagerCasesCandidate (goal : Goal) (type : Expr) : Bool := Id.run do + let .const declName _ := type.getAppFn | return false + return goal.split.casesTypes.isEagerSplit declName + +/-- Returns `true` if `type` is an inductive type with at most one constructor. -/ +private def isCheapInductive (type : Expr) : CoreM Bool := do + let .const declName _ := type.getAppFn | return false + let .inductInfo info ← getConstInfo declName | return false + return info.numCtors <= 1 + +private def applyInjection? (goal : Goal) (fvarId : FVarId) : MetaM (Option Goal) := do + if let some mvarId ← injection? goal.mvarId fvarId then + return some { goal with mvarId } + else + return none + +private def exfalsoIfNotProp (goal : Goal) : MetaM Goal := goal.mvarId.withContext do + if (← isProp (← goal.mvarId.getType)) then + return goal + else + return { goal with mvarId := (← goal.mvarId.exfalso) } + +private def applyCases? (fvarId : FVarId) (generation : Nat) : SearchM Bool := withCurrGoalContext do + /- + Remark: we used to use `whnfD`. This was a mistake, we don't want to unfold user-defined abstractions. + Example: `a ∣ b` is defined as `∃ x, b = a * x` + -/ + let type ← whnf (← fvarId.getType) + if isEagerCasesCandidate (← getGoal) type then + if (← cheapCasesOnly) then + unless (← isCheapInductive type) do + return false + if let .const declName _ := type.getAppFn then + saveCases declName true + let mvarId ← mkAuxMVarForCurrGoal + let mvarIds ← cases mvarId (mkFVar fvarId) + let goal ← getGoal + let goals := mvarIds.map fun mvarId => { goal with mvarId } + mkChoice (mkMVar mvarId) goals generation + return true + return false + +/-- +Introduce new hypotheses (and apply `by_contra`) until goal is of the form `... ⊢ False` +or is inconsistent. +-/ +def intros (generation : Nat) : SearchM Unit := withCurrGoalContext do + repeat + if (← isInconsistent) then + return () + match (← introNext (← getGoal) generation) with + | .done goal => + let goal ← exfalsoIfNotProp goal + if let some mvarId ← goal.mvarId.byContra? then + setGoal { goal with mvarId } + else + setGoal goal + return () + | .newDepHyp goal => + setGoal goal + | .newLocal fvarId goal => + setGoal goal + discard <| applyCases? fvarId generation + | .newHyp fvarId goal => + if let some goal ← applyInjection? goal fvarId then + setGoal goal + else + setGoal goal + if (← applyCases? fvarId generation) then + pure () + else + addHypothesis fvarId generation + +/-- +Similar to `intros`, but returns `true` if new hypotheses have been added, +and `false` otherwise. +-/ +def intros' (generation : Nat) : SearchM Bool := do + let target ← (← getGoal).mvarId.getType + if target.isFalse then return false + intros generation + return true + +/-- Asserts a new fact `prop` with proof `proof` to the given `goal`. -/ +private def assertAt (proof : Expr) (prop : Expr) (generation : Nat) : SearchM Unit := do + if isEagerCasesCandidate (← getGoal) prop then + let goal ← getGoal + let mvarId ← goal.mvarId.assert (← mkFreshUserName `h) prop proof + setGoal { goal with mvarId } + intros generation + else withCurrGoalContext do + let r ← preprocess prop + let prop' := r.expr + let proof' := mkApp4 (mkConst ``Eq.mp [levelZero]) prop r.expr (← r.getProof) proof + add prop' proof' generation + +/-- +Asserts next fact in the `goal` fact queue. +Returns `true` if the queue was not empty and `false` otherwise. +-/ +def assertNext : SearchM Bool := do + if (← isInconsistent) then return false + let goal ← getGoal + let some (fact, newRawFacts) := goal.newRawFacts.dequeue? + | return false + setGoal { goal with newRawFacts } + withSplitSource fact.splitSource do + -- Remark: we should probably add `withGeneration` + assertAt fact.proof fact.prop fact.generation + return true + +/-- +Asserts all facts in the `goal` fact queue. +Returns `true` if the queue was not empty and `false` otherwise. +-/ +def assertAll : SearchM Bool := do + let mut progress := false + repeat + if (← assertNext) then + progress := true + else + return progress + unreachable! + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Inv.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Inv.lean new file mode 100644 index 0000000000000000000000000000000000000000..66fc8097589439db7387d4ab738262cbd465887b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Inv.lean @@ -0,0 +1,137 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Proof +import Lean.Meta.Tactic.Grind.MatchCond +import Lean.Meta.Tactic.Grind.Arith.Inv + +namespace Lean.Meta.Grind + +/-! +Debugging support code for checking basic invariants. +-/ + +private def checkEqc (root : ENode) : GoalM Unit := do + let mut size := 0 + let mut curr := root.self + repeat + size := size + 1 + -- The root of `curr` must be `root` + assert! isSameExpr (← getRoot curr) root.self + -- Check congruence root + if curr.isApp then + if let some { e } := (← get).congrTable.find? { e := curr } then + if (← hasSameType e.getAppFn curr.getAppFn) then + assert! isSameExpr e (← getCongrRoot curr) + else + assert! (← isCongrRoot curr) + -- If the equivalence class does not have HEq proofs, then the types must be definitionally equal. + unless root.heqProofs do + assert! (← hasSameType curr root.self) + -- Starting at `curr`, following the `target?` field leads to `root`. + let mut n := curr + repeat + if let some target ← getTarget? n then + n := target + else + break + assert! isSameExpr n root.self + -- Go to next element + curr ← getNext curr + if isSameExpr root.self curr then + break + -- The size of the equivalence class is correct. + assert! root.size == size + +def checkChild (e : Expr) (child : Expr) : GoalM Bool := do + let some childRoot ← getRoot? child | return false + return isSameExpr childRoot e + +private def checkMatchCondParent (e : Expr) (parent : Expr) : GoalM Bool := do + let_expr Grind.MatchCond parent ← parent | return false + let mut curr := parent + repeat + let .forallE _ d b _ := curr + | return false + match_expr d with + | Eq _ lhs _ => if (← checkChild e lhs) then return true -- found it + | HEq α lhs _ _ => if (← checkChild e α <||> checkChild e lhs) then return true -- found it + | _ => pure () + curr := b + return false + +private def checkParents (e : Expr) : GoalM Unit := do + if (← isRoot e) then + for parent in (← getParents e) do + if isMatchCond parent then + unless (← checkMatchCondParent e parent) do + throwError "e: {e}, parent: {parent}" + assert! (← checkMatchCondParent e parent) + else + let mut found := false + -- There is an argument `arg` s.t. root of `arg` is `e`. + for arg in parent.getAppArgs do + if (← checkChild e arg) then + found := true + break + -- Recall that we have support for `Expr.forallE` propagation. See `ForallProp.lean`. + if let .forallE _ d b _ := parent then + if (← checkChild e d) then + found := true + unless b.hasLooseBVars do + if (← checkChild e b) then + found := true + unless found do + unless (← checkChild e parent.getAppFn) do + throwError "e: {e}, parent: {parent}" + assert! (← checkChild e parent.getAppFn) + else + -- All the parents are stored in the root of the equivalence class. + assert! (← getParents e).isEmpty + +private def checkPtrEqImpliesStructEq : GoalM Unit := do + let exprs ← getExprs + for h₁ : i in [: exprs.size] do + let e₁ := exprs[i] + for h₂ : j in [i+1 : exprs.size] do + let e₂ := exprs[j] + -- We don't have multiple nodes for the same expression + assert! !isSameExpr e₁ e₂ + -- and the two expressions must not be structurally equal + assert! !Expr.equal e₁ e₂ + +private def checkProofs : GoalM Unit := do + let eqcs ← getEqcs + for eqc in eqcs do + for a in eqc do + for b in eqc do + unless isSameExpr a b do + let p ← mkEqHEqProof a b + trace_goal[grind.debug.proofs] "{a} = {b}" + check p + trace_goal[grind.debug.proofs] "checked: {← inferType p}" + +/-- +Checks basic invariants if `grind.debug` is enabled. +-/ +def checkInvariants (expensive := false) : GoalM Unit := do + if grind.debug.get (← getOptions) then + for e in (← getExprs) do + let node ← getENode e + checkParents node.self + if node.isRoot then + checkEqc node + if expensive then + checkPtrEqImpliesStructEq + Arith.checkInvariants + if expensive && grind.debug.proofs.get (← getOptions) then + checkProofs + +def Goal.checkInvariants (goal : Goal) (expensive := false) : GrindM Unit := + discard <| GoalM.run' goal <| Grind.checkInvariants expensive + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/LawfulEqCmp.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/LawfulEqCmp.lean new file mode 100644 index 0000000000000000000000000000000000000000..6bafda34bd782479ffdb1b7cbfe31b8ee30c2d28 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/LawfulEqCmp.lean @@ -0,0 +1,53 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types + +/-! +Support for type class `LawfulEqCmp`. +-/ +/- +Note: we will have similar support for `Associative` and `Commutative`. In the future, we should have +a mechanism for letting users to install their own handlers. +-/ + +namespace Lean.Meta.Grind + +/-- +If `op` implements `LawfulEqCmp`, then returns the proof term for +`∀ a b, op a b = .eq → a = b` +-/ +def getLawfulEqCmpThm? (op : Expr) : GrindM (Option Expr) := do + if let some thm? := (← get).lawfulEqCmpMap.find? { expr := op } then + return thm? + let thm? ← go? + modify fun s => { s with lawfulEqCmpMap := s.lawfulEqCmpMap.insert { expr := op } thm? } + return thm? +where + go? : MetaM (Option Expr) := do + unless (← getEnv).contains ``Std.LawfulEqCmp do return none + let opType ← whnf (← inferType op) + let .forallE _ α b _ := opType | return none + if b.hasLooseBVars then return none + let .forallE _ α' b _ ← whnf b | return none + unless b.isConstOf ``Ordering do return none + unless (← isDefEq α α') do return none + let u ← getLevel α + let some u ← decLevel? u | return none + let lawfulEqCmp := mkApp2 (mkConst ``Std.LawfulEqCmp [u]) α op + let .some lawfulEqCmpInst ← trySynthInstance lawfulEqCmp | return none + return some <| mkApp3 (mkConst ``Std.LawfulEqCmp.eq_of_compare [u]) α op lawfulEqCmpInst + +def propagateLawfulEqCmp (e : Expr) : GoalM Unit := do + let some op := getBinOp e | return () + let some thm ← getLawfulEqCmpThm? op | return () + let oeq ← getOrderingEqExpr + unless (← isEqv e oeq) do return () + let a := e.appFn!.appArg! + let b := e.appArg! + pushEq a b <| mkApp3 thm a b (← mkEqProof e oeq) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Lookahead.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Lookahead.lean new file mode 100644 index 0000000000000000000000000000000000000000..9c3399a6c5e3e8f9e416fc7139adf26d3195afa9 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Lookahead.lean @@ -0,0 +1,82 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Intro +import Lean.Meta.Tactic.Grind.Arith +import Lean.Meta.Tactic.Grind.Split +import Lean.Meta.Tactic.Grind.EMatch +import Lean.Meta.Tactic.Grind.SearchM + +namespace Lean.Meta.Grind + +private partial def solve (generation : Nat) : SearchM Bool := withIncRecDepth do + unless (← get).choiceStack.isEmpty do + return false -- `splitNext` should have been configured to not create choice points + if (← getGoal).inconsistent then + return true + if (← intros' generation <||> assertAll <||> Arith.check <||> splitNext <||> ematch) then + solve generation + else + return false + +private def tryLookahead (e : Expr) : GoalM Bool := + withTheReader Grind.Context + (fun ctx => { ctx with config.qlia := true, cheapCases := true }) do + -- TODO: if `e` is an arithmetic expression, we can avoid creating an auxiliary goal. + -- We can assert it directly to the arithmetic module. + -- Remark: We can simplify this code because the lookahead only really worked for arithmetic. + trace_goal[grind.lookahead.try] "{e}" + let goal ← get + let proof? ← withoutModifyingMCtx do + let tag ← goal.mvarId.getTag + let target ← mkArrow (mkNot e) (← getFalseExpr) + let mvar ← mkFreshExprSyntheticOpaqueMVar target tag + let goalAux := { goal with mvarId := mvar.mvarId!, newFacts := {} } + let gen ← getGeneration e + let (ok, _) ← (solve gen).run goalAux + if ok then + return some (← instantiateMVars mvar) + else + return none + if let some proof := proof? then + trace[grind.lookahead.assert] "{e}" + pushEqTrue e <| mkApp2 (mkConst ``Grind.of_lookahead) e proof + processNewFacts + return true + else + return false + +def lookahead : GoalM Bool := do + unless (← getConfig).lookahead do + return false + if (← get).split.lookaheads.isEmpty then + return false + let mut postponed := [] + let mut progress := false + let infos := (← get).split.lookaheads + modify fun s => { s with split.lookaheads := [] } + for info in infos do + if (← isInconsistent) then + return true + match (← checkSplitStatus info) with + | .resolved => progress := true + | .ready _ _ true + | .notReady => postponed := info :: postponed + | .ready _ _ false => + if (← tryLookahead info.getExpr) then + progress := true + else + postponed := info :: postponed + if progress then + modify fun s => { s with + split.lookaheads := s.split.lookaheads ++ postponed.reverse + } + return true + else + return false + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MBTC.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MBTC.lean new file mode 100644 index 0000000000000000000000000000000000000000..f21ccdb7e908d23f0edbde3c335ca7b8c8538592 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MBTC.lean @@ -0,0 +1,96 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Canon + +namespace Lean.Meta.Grind + +/-- +Model-based theory combination context. +-/ +structure MBTC.Context where + /-- + `isInterpreted e` returns `true` if `e` is an interpreted symbol in the target theory. + Example: `+` for cutsat + -/ + isInterpreted : Expr → GoalM Bool + /-- + `hasTheoryVar e` returns `true` if `e` has a theory variable in the target theory. + For example, suppose we have the constraint `x + y ≤ 0`, then `x` and `y` have theory + vars in the cutsat procedure. + -/ + hasTheoryVar : Expr → GoalM Bool + /-- + `eqAssignment x y` returns `true` it the theory variables for `x` and `y` have the same + interpretation/assignment in the target theory. For example, suppose we have the + constraint `x + y ≤ 0`, and cutsat satisfied it by assignining both `x` and `y` to + `0`. Then, `eqAssignment x y` must return `true`. + -/ + eqAssignment : Expr → Expr → GoalM Bool + +private structure ArgInfo where + arg : Expr + app : Expr + +private abbrev Map := Std.HashMap (Expr × Nat) (List ArgInfo) +private abbrev Candidates := Std.HashSet SplitInfo +private def mkCandidate (a b : ArgInfo) (i : Nat) : GoalM SplitInfo := do + let (lhs, rhs) := if a.arg.lt b.arg then + (a.arg, b.arg) + else + (b.arg, a.arg) + let eq ← mkEq lhs rhs + let eq ← shareCommon (← canon eq) + return .arg a.app b.app i eq (.mbtc a.app b.app i) + +/-- Model-based theory combination. -/ +def mbtc (ctx : MBTC.Context) : GoalM Bool := do + unless (← getConfig).mbtc do return false + -- It is pointless to run `mbtc` if maximum number of splits has been reached. + if (← checkMaxCaseSplit) then return false + let mut map : Map := {} + let mut candidates : Candidates := {} + for e in (← get).exprs do + if e.isApp && !e.isEq && !e.isHEq then + if (← isCongrRoot e) then + unless (← ctx.isInterpreted e) do + let f := e.getAppFn + let mut i := 0 + for arg in e.getAppArgs do + let some arg ← getRoot? arg | pure () + if (← ctx.hasTheoryVar arg) then + trace[grind.debug.mbtc] "{arg} @ {f}:{i}" + let argInfo : ArgInfo := { arg, app := e } + if let some otherInfos := map[(f, i)]? then + unless otherInfos.any fun info => isSameExpr arg info.arg do + for otherInfo in otherInfos do + if (← ctx.eqAssignment arg otherInfo.arg) then + if (← hasSameType arg otherInfo.arg) then + candidates := candidates.insert (← mkCandidate argInfo otherInfo i) + map := map.insert (f, i) (argInfo :: otherInfos) + else + map := map.insert (f, i) [argInfo] + i := i + 1 + if candidates.isEmpty then + return false + if (← get).split.num > (← getConfig).splits then + reportIssue! "skipping `mbtc`, maximum number of splits has been reached `(splits := {(← getConfig).splits})`" + return false + let result := candidates.toArray.qsort fun c₁ c₂ => c₁.lt c₂ + let result ← result.filterMapM fun info => do + if (← isKnownCaseSplit info) then + return none + let .arg a b _ eq _ := info | return none + internalize eq (Nat.max (← getGeneration a) (← getGeneration b)) + return some info + if result.isEmpty then + return false + for info in result do + addSplitCandidate info + return true + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Main.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Main.lean new file mode 100644 index 0000000000000000000000000000000000000000..868c56acbff0854363063cf31f1dcb544b8ecb89 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Main.lean @@ -0,0 +1,224 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Lemmas +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.ExposeNames +import Lean.Meta.Tactic.Simp.Diagnostics +import Lean.Meta.Tactic.Grind.RevertAll +import Lean.Meta.Tactic.Grind.PropagatorAttr +import Lean.Meta.Tactic.Grind.Proj +import Lean.Meta.Tactic.Grind.ForallProp +import Lean.Meta.Tactic.Grind.Util +import Lean.Meta.Tactic.Grind.Inv +import Lean.Meta.Tactic.Grind.Intro +import Lean.Meta.Tactic.Grind.EMatch +import Lean.Meta.Tactic.Grind.Split +import Lean.Meta.Tactic.Grind.Solve +import Lean.Meta.Tactic.Grind.SimpUtil +import Lean.Meta.Tactic.Grind.Cases +import Lean.Meta.Tactic.Grind.LawfulEqCmp +import Lean.Meta.Tactic.Grind.ReflCmp + +namespace Lean.Meta.Grind + +structure Params where + config : Grind.Config + ematch : EMatchTheorems := {} + symPrios : SymbolPriorities := {} + casesTypes : CasesTypes := {} + extra : PArray EMatchTheorem := {} + norm : Simp.Context + normProcs : Array Simprocs + -- TODO: inductives to split + +def mkParams (config : Grind.Config) : MetaM Params := do + let norm ← Grind.getSimpContext config + let normProcs ← Grind.getSimprocs + let symPrios ← getGlobalSymbolPriorities + return { config, norm, normProcs, symPrios } + +def mkMethods (fallback : Fallback) : CoreM Methods := do + let builtinPropagators ← builtinPropagatorsRef.get + return { + fallback + propagateUp := fun e => do + propagateForallPropUp e + propagateReflCmp e + let .const declName _ := e.getAppFn | return () + propagateProjEq e + if let some prop := builtinPropagators.up[declName]? then + prop e + propagateDown := fun e => do + propagateForallPropDown e + propagateLawfulEqCmp e + let .const declName _ := e.getAppFn | return () + if let some prop := builtinPropagators.down[declName]? then + prop e + } + +-- A `simp` discharger that does not use assumptions. +-- We use it to make sure we don't have to reset the `simp` cache used in `grind`. +private def discharge? (e : Expr) : SimpM (Option Expr) := do + let e := e.cleanupAnnotations + let r ← Simp.simp e + if let some p ← Simp.dischargeRfl r.expr then + return some (mkApp4 (mkConst ``Eq.mpr [levelZero]) e r.expr (← r.getProof) p) + else if r.expr.isTrue then + return some (← mkOfEqTrue (← r.getProof)) + else + return none + +def GrindM.run (x : GrindM α) (params : Params) (fallback : Fallback) : MetaM α := do + let (falseExpr, scState) := shareCommonAlpha (mkConst ``False) {} + let (trueExpr, scState) := shareCommonAlpha (mkConst ``True) scState + let (bfalseExpr, scState) := shareCommonAlpha (mkConst ``Bool.false) scState + let (btrueExpr, scState) := shareCommonAlpha (mkConst ``Bool.true) scState + let (natZExpr, scState) := shareCommonAlpha (mkNatLit 0) scState + let (ordEqExpr, scState) := shareCommonAlpha (mkConst ``Ordering.eq) scState + let (intExpr, scState) := shareCommonAlpha Int.mkType scState + let simprocs := params.normProcs + let simpMethods := Simp.mkMethods simprocs discharge? (wellBehavedDischarge := true) + let simp := params.norm + let config := params.config + let symPrios := params.symPrios + x (← mkMethods fallback).toMethodsRef { config, simpMethods, simp, trueExpr, falseExpr, natZExpr, btrueExpr, bfalseExpr, ordEqExpr, intExpr, symPrios } + |>.run' { scState } + +private def mkCleanState (mvarId : MVarId) (params : Params) : MetaM Clean.State := mvarId.withContext do + unless params.config.clean do return {} + let mut used : PHashSet Name := {} + for localDecl in (← getLCtx) do + used := used.insert localDecl.userName + return { used } + +private def mkGoal (mvarId : MVarId) (params : Params) : GrindM Goal := do + let mvarId ← if params.config.clean then mvarId.exposeNames else pure mvarId + let trueExpr ← getTrueExpr + let falseExpr ← getFalseExpr + let btrueExpr ← getBoolTrueExpr + let bfalseExpr ← getBoolFalseExpr + let natZeroExpr ← getNatZeroExpr + let ordEqExpr ← getOrderingEqExpr + let thmMap := params.ematch + let casesTypes := params.casesTypes + let clean ← mkCleanState mvarId params + GoalM.run' { mvarId, ematch.thmMap := thmMap, split.casesTypes := casesTypes, clean } do + mkENodeCore falseExpr (interpreted := true) (ctor := false) (generation := 0) + mkENodeCore trueExpr (interpreted := true) (ctor := false) (generation := 0) + mkENodeCore btrueExpr (interpreted := false) (ctor := true) (generation := 0) + mkENodeCore bfalseExpr (interpreted := false) (ctor := true) (generation := 0) + mkENodeCore natZeroExpr (interpreted := true) (ctor := false) (generation := 0) + mkENodeCore ordEqExpr (interpreted := false) (ctor := true) (generation := 0) + for thm in params.extra do + activateTheorem thm 0 + +structure Result where + failure? : Option Goal + issues : List MessageData + config : Grind.Config + trace : Trace + counters : Counters + simp : Simp.Stats + splitDiags : PArray SplitDiagInfo + +private def countersToMessageData (header : String) (cls : Name) (data : Array (Name × Nat)) : MetaM MessageData := do + let data := data.qsort fun (d₁, c₁) (d₂, c₂) => if c₁ == c₂ then Name.lt d₁ d₂ else c₁ > c₂ + let data ← data.mapM fun (declName, counter) => + return .trace { cls } m!"{.ofConst (← mkConstWithLevelParams declName)} ↦ {counter}" #[] + return .trace { cls } header data + +private def splitDiagInfoToMessageData (ss : Array SplitDiagInfo) : MetaM MessageData := do + let env ← getEnv + let mctx ← getMCtx + let opts ← getOptions + let cls := `split + let data ← ss.mapM fun { c, lctx, numCases, gen, splitSource } => do + let header := m!"{c}" + return MessageData.withContext { env, mctx, lctx, opts } <| .trace { cls } header #[ + .trace { cls } m!"source: {← splitSource.toMessageData}" #[], + .trace { cls } m!"generation: {gen}" #[], + .trace { cls } m!"# cases: {numCases}" #[] + ] + return .trace { cls } "Case splits" data + +-- Diagnostics information for the whole search +private def mkGlobalDiag (cs : Counters) (simp : Simp.Stats) (ss : PArray SplitDiagInfo) : MetaM (Option MessageData) := do + let thms := cs.thm.toList.toArray.filterMap fun (origin, c) => + match origin with + | .decl declName => some (declName, c) + | _ => none + -- We do not report `cases` applications on builtin types + let cases := cs.case.toList.toArray.filter fun (declName, _) => !isBuiltinEagerCases declName + let mut msgs := #[] + unless thms.isEmpty do + msgs := msgs.push <| (← countersToMessageData "E-Matching instances" `thm thms) + let ss := ss.toArray.filter fun { numCases, .. } => numCases > 1 + unless ss.isEmpty do + msgs := msgs.push <| (← splitDiagInfoToMessageData ss) + unless cases.isEmpty do + msgs := msgs.push <| (← countersToMessageData "Cases instances" `cases cases) + unless cs.apps.isEmpty do + msgs := msgs.push <| (← countersToMessageData "Applications" `app cs.apps.toList.toArray) + let simpMsgs ← Simp.mkDiagMessages simp.diag + unless simpMsgs.isEmpty do + msgs := msgs.push <| .trace { cls := `grind} "Simplifier" simpMsgs + if msgs.isEmpty then + return none + else + return some <| .trace { cls := `grind } "Diagnostics" msgs + +def Result.hasFailed (r : Result) : Bool := + r.failure?.isSome + +def Result.toMessageData (result : Result) : MetaM MessageData := do + let mut msgs ← result.failure?.toList.mapM (goalToMessageData · result.config) + if result.config.verbose then + let mut issues := result.issues + -- We did not find the following very useful in practice. + /- + unless result.skipped.isEmpty do + let m := m!"#{result.skipped.length} other goal(s) were not fully processed due to previous failures, threshold: `(failures := {result.config.failures})`" + issues := .trace { cls := `issue } m #[] :: issues + -/ + unless issues.isEmpty do + msgs := msgs ++ [.trace { cls := `grind } "Issues" issues.reverse.toArray] + if let some msg ← mkGlobalDiag result.counters result.simp result.splitDiags then + msgs := msgs ++ [msg] + return MessageData.joinSep msgs m!"\n" + +private def initCore (mvarId : MVarId) (params : Params) : GrindM Goal := do + let mvarId ← mvarId.abstractMVars + let mvarId ← mvarId.clearAuxDecls + let mvarId ← mvarId.revertAll + let mvarId ← mvarId.unfoldReducible + let mvarId ← mvarId.betaReduce + appendTagSuffix mvarId `grind + mkGoal mvarId params + +def main (mvarId : MVarId) (params : Params) (fallback : Fallback) : MetaM Result := do profileitM Exception "grind" (← getOptions) do + if debug.terminalTacticsAsSorry.get (← getOptions) then + mvarId.admit + return { + failure? := none, issues := [], config := params.config, trace := {}, counters := {}, simp := {}, splitDiags := {} + } + let go : GrindM Result := withReducible do + let goal ← initCore mvarId params + let failure? ← solve goal + let issues := (← get).issues + let trace := (← get).trace + let counters := (← get).counters + let splitDiags := (← get).splitDiags + let simp := { (← get).simp with } + if failure?.isNone then + -- If there are no failures and diagnostics are enabled, we still report the performance counters. + if (← isDiagnosticsEnabled) then + if let some msg ← mkGlobalDiag counters simp splitDiags then + logInfo msg + return { failure?, issues, config := params.config, trace, counters, simp, splitDiags } + go.run params fallback + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MarkNestedProofs.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MarkNestedProofs.lean new file mode 100644 index 0000000000000000000000000000000000000000..b9a9fe42a14c90deed8cbd653ec16389952dfab7 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MarkNestedProofs.lean @@ -0,0 +1,98 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Lean.Util.PtrSet +import Lean.Meta.Transform +import Lean.Meta.Basic +import Lean.Meta.InferType +import Lean.Meta.Tactic.Grind.Util + +namespace Lean.Meta.Grind + +private unsafe def markNestedProofImpl (e : Expr) (visit : Expr → StateRefT (PtrMap Expr Expr) MetaM Expr) + : StateRefT (PtrMap Expr Expr) MetaM Expr := do + let prop ← inferType e + /- + We must unfold reducible constants occurring in `prop` because the congruence closure + module in `grind` assumes they have been expanded. + See `grind_mark_nested_proofs_bug.lean` for an example. + TODO: We may have to normalize `prop` too. + -/ + /- We must also apply beta-reduction to improve the effectiveness of the congruence closure procedure. -/ + let prop ← Core.betaReduce prop + let prop ← unfoldReducible prop + /- We must mask proofs occurring in `prop` too. -/ + let prop ← visit prop + let prop ← eraseIrrelevantMData prop + /- We must fold kernel projections like it is done in the preprocessor. -/ + let prop ← foldProjs prop + let prop ← normalizeLevels prop + return mkApp2 (mkConst ``Lean.Grind.nestedProof) prop e + +unsafe def markNestedProofsImpl (e : Expr) : MetaM Expr := do + visit e |>.run' mkPtrMap +where + visit (e : Expr) : StateRefT (PtrMap Expr Expr) MetaM Expr := do + if (← isProof e) then + if e.isAppOf ``Lean.Grind.nestedProof then + return e -- `e` is already marked + if let some r := (← get).find? e then + return r + let e' ← markNestedProofImpl e visit + modify fun s => s.insert e e' + return e' + -- Remark: we have to process `Expr.proj` since we only + -- fold projections later during term internalization + unless e.isApp || e.isForall || e.isProj || e.isMData do + return e + -- Check whether it is cached + if let some r := (← get).find? e then + return r + let e' ← match e with + | .app .. => e.withApp fun f args => do + let mut modified := false + let mut args := args + for i in [:args.size] do + let arg := args[i]! + let arg' ← visit arg + unless ptrEq arg arg' do + args := args.set! i arg' + modified := true + if modified then + pure <| mkAppN f args + else + pure e + | .proj _ _ b => + pure <| e.updateProj! (← visit b) + | .mdata _ b => + pure <| e.updateMData! (← visit b) + | .forallE _ d b _ => + -- Recall that we have `ForallProp.lean`. + let d' ← visit d + let b' ← if b.hasLooseBVars then pure b else visit b + pure <| e.updateForallE! d' b' + | _ => unreachable! + modify fun s => s.insert e e' + return e' + +/-- +Wrap nested proofs `e` with `Lean.Grind.nestedProof`-applications. +Recall that the congruence closure module has special support for `Lean.Grind.nestedProof`. +-/ +def markNestedProofs (e : Expr) : MetaM Expr := + unsafe markNestedProofsImpl e + +/-- +Given a proof `e`, mark it with `Lean.Grind.nestedProof` +-/ +def markProof (e : Expr) : MetaM Expr := do + if e.isAppOf ``Lean.Grind.nestedProof then + return e -- `e` is already marked + else + unsafe markNestedProofImpl e markNestedProofsImpl.visit |>.run' mkPtrMap + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MatchCond.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MatchCond.lean new file mode 100644 index 0000000000000000000000000000000000000000..31d19342f0ddcae8490120a2d0c577eb3cdf15e8 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MatchCond.lean @@ -0,0 +1,429 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind +import Init.Simproc +import Lean.Meta.Tactic.Contradiction +import Lean.Meta.Tactic.Grind.ProveEq +import Lean.Meta.Tactic.Grind.PropagatorAttr + +namespace Lean.Meta.Grind +/- +Remark: the `simp` module has some support for `MatchCond`, but it is +called `isEqnThmHypothesis` there. This is actually a better name. +TODO: rename this module and functions to `EqThmHyp`. +-/ + +/-! +Support for `match`-expressions with overlapping patterns. +Recall that when a `match`-expression has overlapping patterns, some of its equation theorems are +conditional. Let's consider the following example +``` +inductive S where + | mk1 (n : Nat) + | mk2 (n : Nat) (s : S) + | mk3 (n : Bool) + | mk4 (s1 s2 : S) + +def f (x y : S) := + match x, y with + | .mk1 a, c => a + 2 + | a, .mk2 1 (.mk4 b c) => 3 + | .mk3 a, .mk4 b c => 4 + | a, b => 5 +``` + +The `match`-expression in this example has 4 equations. The second and fourth are conditional. +``` +f.match_1.eq_2 + (motive : S → S → Sort u) (a b c : S) + (h_1 : (a : Nat) → (c : S) → motive (S.mk1 a) c) + (h_2 : (a b c : S) → motive a (S.mk2 1 (b.mk4 c))) + (h_3 : (a : Bool) → (b c : S) → motive (S.mk3 a) (b.mk4 c)) + (h_4 : (a b : S) → motive a b) : + (∀ (a_1 : Nat), a = S.mk1 a_1 → False) → -- <<< Condition stating it is not the first case + f.match_1 motive a (S.mk2 1 (b.mk4 c)) h_1 h_2 h_3 h_4 = h_2 a b c + +f.match_1.eq_4 + (motive : S → S → Sort u) (a b : S) + (h_1 : (a : Nat) → (c : S) → motive (S.mk1 a) c) + (h_2 : (a b c : S) → motive a (S.mk2 1 (b.mk4 c))) + (h_3 : (a : Bool) → (b c : S) → motive (S.mk3 a) (b.mk4 c)) + (h_4 : (a b : S) → motive a b) : + (∀ (a_1 : Nat), a = S.mk1 a_1 → False) → -- <<< Condition stating it is not the first case + (∀ (b_1 c : S), b = S.mk2 1 (b_1.mk4 c) → False) → -- <<< Condition stating it is not the second case + (∀ (a_1 : Bool) (b_1 c : S), a = S.mk3 a_1 → b = b_1.mk4 c → False) → -- -- <<< Condition stating it is not the third case + f.match_1 motive a b h_1 h_2 h_3 h_4 = h_4 a b +``` +In the two equational theorems above, we have the following conditions. +``` +- `(∀ (a_1 : Nat), a = S.mk1 a_1 → False)` +- `(∀ (b_1 c : S), b = S.mk2 1 (b_1.mk4 c) → False)` +- `(∀ (a_1 : Bool) (b_1 c : S), a = S.mk3 a_1 → b = b_1.mk4 c → False)` +``` +When instantiating the equations (and `match`-splitter), we wrap the conditions with the gadget `Grind.PreMatchCond`. +`Grind.PreMatchCond` uses the default reducibility setting and cannot be accidentally reduced by `simp`. +After `simp` is applied, it is replaced with `Grind.MatchCond` which is reducible. +This `Grind.MatchCond` is used for implementing truth-value propagation. +See the propagator `propagateMatchCond` below. +For example, given a condition `C` of the form `Grind.MatchCond (∀ (a : Nat), t = S.mk1 a → False)`, +if `t` is merged with an equivalence class containing `S.mk2 n s`, then `C` is asseted to `true` by `propagateMatchCond`. + +This module also provides auxiliary functions for detecting congruences between `match`-expression conditions. +See function `collectMatchCondLhssAndAbstract`. + +Remark: This note highlights that the representation used for encoding `match`-expressions with +overlapping patterns is far from ideal for the `grind` module which operates with equivalence classes +and does not perform substitutions like `simp`. While modifying how `match`-expressions are encoded in Lean +would require major refactoring and affect many modules, this issue is important to acknowledge. +A different representation could simplify `grind`, but it could add extra complexity to other modules. +-/ + +/-- +Returns `some (α?, lhs, rhs)` if `e` is of the form +- `Eq _ lhs rhs` (`?α := none`), or +- `HEq α lhs _ rhs` (`α? := some α`) +-/ +private def isEqHEq? (e : Expr) : Option (Option Expr × Expr × Expr) := + match_expr e with + | Eq _ lhs rhs => some (none, lhs, rhs) + | HEq α lhs _ rhs => some (some α, lhs, rhs) + | _ => none + +/-- +Given `e` a `match`-expression condition, returns the left-hand side +of the ground equations, and its type if it is the left-hand side of +a heterogeneous equality. +-/ +private def collectMatchCondLhss (e : Expr) : Array (Expr × Option Expr) := Id.run do + let mut r := #[] + let mut e := e + repeat + let .forallE _ d b _ := e | return r + if let some (α?, lhs, _) := isEqHEq? d then + unless lhs.hasLooseBVars do + r := r.push (lhs, α?) + e := b + return r + +/-- +Replaces the left-hand side of an equality (or heterogeneous equality) `e` with `lhsNew`. +-/ +private def replaceLhs? (e : Expr) (lhsNew : Expr) (ty? : Option Expr) : Option Expr := + match_expr e with + | f@Eq α lhs rhs => if lhs.hasLooseBVars then none else some (mkApp3 f α lhsNew rhs) + | f@HEq _ lhs β rhs => if lhs.hasLooseBVars then none else some (mkApp4 f ty?.get! lhsNew β rhs) + | _ => none + +/-- +Given `e` a `match`-expression condition, returns the left-hand side +of the ground equations, **and** function application that abstracts the left-hand sides. +As an example, assume we have a `match`-expression condition `C₁` of the form +``` +Grind.MatchCond (∀ y₁ y₂ y₃, t = .mk₁ y₁ → s = .mk₂ y₂ y₃ → False) +``` +then the result returned by this function is +``` +(#[t, s], (fun x₁ x₂ => (∀ y₁ y₂ y₃, x₁ = .mk₁ y₁ → x₂ = .mk₂ y₂ y₃ → False)) t s) +``` +Note that the returned expression is definitionally equal to `C₁`. +We use this expression to detect whether two different `match`-expression conditions are +congruent. +For example, suppose we also have the `match`-expression `C₂` of the form +``` +Grind.MatchCond (∀ y₁ y₂ y₃, a = .mk₁ y₁ → b = .mk₂ y₂ y₃ → False) +``` +This function would return +``` +(#[a, b], (fun x₁ x₂ => (∀ y₁ y₂ y₃, x₁ = .mk₁ y₁ → x₂ = .mk₂ y₂ y₃ → False)) a b) +``` +Note that the lambda abstraction is identical to the first one. Let's call it `l`. +Thus, we can write the two pairs above as +- `(#[t, s], l t s)` +- `(#[a, b], l a b)` +Moreover, `C₁` is definitionally equal to `l t s`, and `C₂` is definitionally equal to `l a b`. +Then, if `grind` infers that `t = a` and `s = b`, it will detect that `l t s` and `l a b` are +equal by congruence, and consequently `C₁` is equal to `C₂`. + +Gruesome details for heterogenenous equalities. + +When pattern matching on indexing families, the generated conditions often use heterogenenous equalities. Here is an example: +``` +(∀ (x : Vec α 0), n = 0 → as ≍ Vec.nil → bs ≍ x → False) +``` +In this case, it is not sufficient to abstract the left-hand side. We also have +to abstract its type. The following is produced in this case. +``` +(#[n, Vec α n, as, Vec α n, bs], + (fun (x_0 : Nat) (ty_1 : Type u_1) (x_1 : ty_1) (ty_2 : Type u_1) (x_2 : ty_2) => + ∀ (x : Vec α 0), x_0 = 0 → x_1 ≍ Vec.nil → x_2 ≍ x → False) + n (Vec α n) as (Vec α n) bs) +``` +The example makes it clear why this is needed, `as` and `bs` depend on `n`. +Note that we can abstract the type without introducing typer errors because +heterogenenous equality is used for `as` and `bs`. +-/ +def collectMatchCondLhssAndAbstract (matchCond : Expr) : GoalM (Array Expr × Expr) := do + let_expr Grind.MatchCond e := matchCond | return (#[], matchCond) + let lhssαs := collectMatchCondLhss e + let rec go (i : Nat) (xs : Array Expr) (tys : Array (Option Expr)) (tysxs : Array Expr) (args : Array Expr) : GoalM (Array Expr × Expr) := do + if h : i < lhssαs.size then + let (lhs, α?) := lhssαs[i] + if let some α := α? then + withLocalDeclD ((`ty).appendIndexAfter i) (← inferType α) fun ty => + withLocalDeclD ((`x).appendIndexAfter i) ty fun x => + go (i+1) (xs.push x) (tys.push ty) (tysxs.push ty |>.push x) (args.push α |>.push lhs) + else + withLocalDeclD ((`x).appendIndexAfter i) (← inferType lhs) fun x => + go (i+1) (xs.push x) (tys.push none) (tysxs.push x) (args.push lhs) + else + let rec replaceLhss (e : Expr) (i : Nat) : Expr := Id.run do + let .forallE _ d b _ := e | return e + if h : i < xs.size then + if let some dNew := replaceLhs? d xs[i] tys[i]! then + return e.updateForallE! dNew (replaceLhss b (i+1)) + else + return e.updateForallE! d (replaceLhss b i) + else + return e + let eAbst := replaceLhss e 0 + let eLam ← mkLambdaFVars tysxs eAbst + let e' := mkAppN eLam args + return (args, ← shareCommon e') + go 0 #[] #[] #[] #[] + +/-- +Helper function for `isSatisfied`. +See `isSatisfied`. +-/ +private partial def isMatchCondFalseHyp (e : Expr) : GoalM Bool := do + let some (_, lhs, rhs) := isEqHEq? e | return false + isFalse lhs rhs +where + isFalse (lhs rhs : Expr) : GoalM Bool := do + if lhs.hasLooseBVars then return false + let root ← getRootENode lhs + if root.ctor then + let some ctorLhs ← isConstructorApp? root.self | return false + let some ctorRhs ← isConstructorApp? rhs | return false + if ctorLhs.name ≠ ctorRhs.name then return true + let lhsArgs := root.self.getAppArgs + let rhsArgs := rhs.getAppArgs + for i in [ctorLhs.numParams : ctorLhs.numParams + ctorLhs.numFields] do + if (← isFalse lhsArgs[i]! rhsArgs[i]!) then + return true + return false + else if root.interpreted then + if rhs.hasLooseBVars then return false + unless (← isLitValue rhs) do return false + return (← normLitValue root.self) != (← normLitValue rhs) + else + return false + +/-- +Returns `true` if `e` is a `Grind.MatchCond`, and it has been satifisfied. +Recall that we use `Grind.MatchCond` to annotate conditional `match`-equations. +Consider the following example: +``` +inductive S where + | mk1 (n : Nat) + | mk2 (n : Nat) (s : S) + | mk3 (n : Bool) + | mk4 (s1 s2 : S) + +def f (x y : S) := + match x, y with + | .mk1 _, _ => 2 + | _, .mk2 1 (.mk4 _ _) => 3 + | .mk3 _, _ => 4 + | _, _ => 5 +``` +The `match`-expression in the example above has overlapping patterns and +consequently produces conditional `match` equations. Thus, `grind` generates +the following auxiliary `Grind.MatchCond` terms for an application `f a b`: +- `Grind.MatchCond (∀ (n : Nat), a = S.mk1 n → False)` +- `Grind.MatchCond (∀ (s1 s2 : S), b = S.mk2 1 (s1.mk4 s2) → False)` +- `Grind.MatchCond (∀ (n : Bool), a = S.mk3 n → False)` + +`isSatisfied` uses the fact that constructor applications and literal values +are always the root of their equivalence classes. +-/ +private partial def isStatisfied (e : Expr) : GoalM Bool := do + let_expr Grind.MatchCond e ← e | return false + let mut e := e + repeat + let .forallE _ d b _ := e | break + if (← isMatchCondFalseHyp d) then + trace_goal[grind.debug.matchCond] "satifised{indentExpr e}\nthe following equality is false{indentExpr d}" + return true + e := b + return false + +-- TODO: we don't have support for offset equalities + +/-- Constructs a proof for a satisfied `match`-expression condition. -/ +private partial def mkMatchCondProof? (e : Expr) : GoalM (Option Expr) := do + let_expr Grind.MatchCond f ← e | return none + forallTelescopeReducing f fun xs _ => do + for x in xs do + let type ← inferType x + if (← isMatchCondFalseHyp type) then + trace_goal[grind.debug.matchCond] ">>> {type}" + let some h ← go? x | pure () + return some (← mkLambdaFVars xs h) + return none +where + go? (h : Expr) : GoalM (Option Expr) := do + trace_goal[grind.debug.matchCond] "go?: {← inferType h}" + let some (α?, lhs, rhs) := isEqHEq? (← inferType h) + | return none + let target ← (← get).mvarId.getType + -- We use `shareCommon` here because we may accessing a new expression + -- created when we infer the type of the `noConfusion` term below + let lhs ← shareCommon lhs + let some root ← getRootENode? lhs + | reportIssue! "found term that has not been internalized{indentExpr lhs}\nwhile trying to construct a proof for `MatchCond`{indentExpr e}" + return none + let isHEq := α?.isSome + unless (← hasSameType root.self rhs) do + return none + let h ← if isHEq then + mkEqOfHEq (← mkHEqTrans (← mkHEqProof root.self lhs) h) (check := false) + else + mkEqTrans (← mkEqProof root.self lhs) h + if root.ctor then + let some ctorLhs ← isConstructorApp? root.self | return none + let some ctorRhs ← isConstructorApp? rhs | return none + -- See comment on `shareCommon` above. + let h ← mkNoConfusion target h + if ctorLhs.name ≠ ctorRhs.name then + return some h + else + let .forallE _ k _ _ ← whnfD (← inferType h) + | return none + forallTelescopeReducing k fun xs _ => do + for x in xs do + let some hx ← go? x | pure () + return some (mkApp h (← mkLambdaFVars xs hx)) + return none + else if root.interpreted then + if (← normLitValue root.self) != (← normLitValue rhs) then + let hne ← mkDecideProof (mkNot (← mkEq root.self rhs)) + return some (mkApp hne h) + else + return none + else + return none + +/-- +Given a `match`-expression condition `e` that is known to be equal to `True`, +try to close the goal by proving `False`. We use the following to example to illustrate +the purpose of this function. +``` +def f : List Nat → List Nat → Nat + | _, 1 :: _ :: _ => 1 + | _, _ :: _ => 2 + | _, _ => 0 + +example : z = a :: as → y = z → f x y > 0 := by + grind [f.eq_def] +``` +After `grind` unfolds `f`, it case splits on the `match`-expression producing +three subgoals. The first two are easily closed by it. In the third one, +we have the following two `match`-expression conditions stating that we +are **not** in the first and second cases. +``` +Lean.Grind.MatchCond (∀ (head : Nat) (tail : List Nat), x✝² = 1 :: head :: tail → False) +Lean.Grind.MatchCond (∀ (head : Nat) (tail : List Nat), x✝² = head :: tail → False) +``` +Moreover, we have the following equivalence class. +``` +{z, y, x✝², a :: as} +``` +Thus, we can close the goal by using the second `match`-expression condition, +we just have to instantiate `head` and `tail` with `a` and `as` respectively, +and use the fact that `x✝²` is equal to `a :: as`. +-/ +partial def tryToProveFalse (e : Expr) : GoalM Unit := do + trace_goal[grind.debug.matchCond.proveFalse] "{e}" + let_expr Grind.MatchCond body ← e | return () + let proof? ← withNewMCtxDepth do + let (args, _, _) ← forallMetaTelescope body + let mask := mkGenDiseqMask body + for arg in args, target in mask do + if target then + let some (α?, lhs, rhs) := isEqHEq? (← inferType arg) + | return none + let lhs' ← go lhs + trace[grind.debug.matchCond.proveFalse] "{lhs'} =?= {rhs}" + unless (← withDefault <| isDefEq lhs' rhs) do + return none + let isHEq := α?.isSome + let some lhsEqLhs' ← if isHEq then proveHEq? lhs lhs' else proveEq? lhs lhs' + | return none + unless (← isDefEq arg lhsEqLhs') do + return none + let he := mkOfEqTrueCore e (← mkEqTrueProof e) + let falseProof ← instantiateMVars (mkAppN he args) + if (← hasAssignableMVar falseProof) then + return none + trace[grind.debug.matchCond.proveFalse] "{falseProof} : {← inferType falseProof}" + return some falseProof + let some proof := proof? | return () + closeGoal proof +where + /-- + Returns a term that is equal to `e`, but containing constructor applications + and literal values. `e` is the left-hand side of the equations in a `match`-expression + condition. + Remark: we could use the right-hand side to interrupt the recursion. For example, + suppose the equation is `x = ?head :: ?tail`. We only need to show that `x` is equal to + some term of the form `a :: as` to satisfy it. This function may return `a₁ :: b :: bs`, + which still allows us to satisfy the equation, but may have a bigger proof (e.g., + a proof that `as` is equal to `b::bs`) + -/ + go (e : Expr) : GoalM Expr := do + let root ← getRootENode e + if root.ctor then + let ctor := root.self + let some ctorInfo ← isConstructorApp? ctor | return ctor + let mut ctorArgs := ctor.getAppArgs + let mut modified := false + for i in [ctorInfo.numParams : ctorInfo.numParams + ctorInfo.numFields] do + let arg := ctorArgs[i]! + let arg' ← go arg + unless isSameExpr arg arg' do + ctorArgs := ctorArgs.set! i arg' + modified := true + if modified then + shareCommon <| mkAppN ctor.getAppFn ctorArgs + else + return root.self + else if root.interpreted then + return root.self + else + return e + +/-- Propagates `MatchCond` upwards -/ +builtin_grind_propagator propagateMatchCondUp ↑Grind.MatchCond := fun e => do + trace_goal[grind.debug.matchCond] "visiting{indentExpr e}" + if (← isEqTrue e) then + unless (← isStatisfied e) do + tryToProveFalse e + else + if !(← isStatisfied e) then return () + let some h ← mkMatchCondProof? e + | reportIssue! "failed to construct proof for{indentExpr e}"; return () + trace_goal[grind.debug.matchCond] "{← inferType h}" + pushEqTrue e <| mkEqTrueCore e h + +/-- Propagates `MatchCond` downwards -/ +builtin_grind_propagator propagateMatchCondDown ↓Grind.MatchCond := fun e => do + if (← isEqTrue e) then + unless (← isStatisfied e) do + tryToProveFalse e + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MatchDiscrOnly.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MatchDiscrOnly.lean new file mode 100644 index 0000000000000000000000000000000000000000..4641b328e7ea25e1308f8ad4744c5c31b8271c65 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/MatchDiscrOnly.lean @@ -0,0 +1,54 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Init.Simproc +import Lean.Meta.Tactic.Simp.Simproc +import Lean.Meta.Tactic.Simp.Rewrite + +namespace Lean.Meta.Grind +/-- +Returns `Grind.simpMatchDiscrsOnly e`. Recall that `Grind.simpMatchDiscrsOnly` is +a gadget for instructing the `grind` simplifier to only normalize/simplify +the discriminants of a `match`-expression. See `reduceSimpMatchDiscrsOnly`. +-/ +def markAsSimpMatchDiscrsOnly (e : Expr) : MetaM Expr := + mkAppM ``Grind.simpMatchDiscrsOnly #[e] + +def isSimpMatchDiscrsOnly (e : Expr) := + e.isAppOfArity ``Grind.simpMatchDiscrsOnly 2 + +builtin_simproc_decl reduceSimpMatchDiscrsOnly (Grind.simpMatchDiscrsOnly _) := fun e => do + let_expr Grind.simpMatchDiscrsOnly _ m ← e | return .continue + let .const declName _ := m.getAppFn + | return .done { expr := e } + let some info ← getMatcherInfo? declName + | return .done { expr := e } + if let some r ← Simp.simpMatchDiscrs? info m then + return .done { r with expr := (← markAsSimpMatchDiscrsOnly r.expr) } + return .done { expr := e } + +/-- Adds `reduceSimpMatchDiscrsOnly` to `s` -/ +def addSimpMatchDiscrsOnly (s : Simprocs) : CoreM Simprocs := do + s.add ``reduceSimpMatchDiscrsOnly (post := false) + +/-- Erases `Grind.simpMatchDiscrsOnly` annotations. -/ +def eraseSimpMatchDiscrsOnly (e : Expr) : MetaM Simp.Result := do + if e.find? isSimpMatchDiscrsOnly |>.isNone then + return { expr := e } + else + let pre (e : Expr) := do + let_expr Grind.simpMatchDiscrsOnly _ a := e | return .continue e + return .continue a + let e' ← Core.transform e (pre := pre) + /- + `grind` uses the `.reducible` transparency setting, and `Grind.simpMatchDiscrsOnly` is not + reducible. Thus, `e` and `e'` are not definitionally equal in this setting, and we must + add a hint. + -/ + return { expr := e', proof? := mkExpectedPropHint (← mkEqRefl e') (← mkEq e e') } + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/PP.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/PP.lean new file mode 100644 index 0000000000000000000000000000000000000000..084b65bb7f05bea5de0f7d2cf6b55e282b5d4dea --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/PP.lean @@ -0,0 +1,207 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Util +import Init.Grind.PP +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Arith.Model +import Lean.Meta.Tactic.Grind.Arith.CommRing.PP +import Lean.Meta.Tactic.Grind.Arith.Linear.PP + +namespace Lean.Meta.Grind + +/-- Helper function for pretty printing the state for debugging purposes. -/ +def Goal.ppENodeRef (goal : Goal) (e : Expr) : MetaM MessageData := do + let some n := goal.getENode? e | return "_" + let type ← inferType e + let u ← getLevel type + let d := mkApp3 (mkConst ``Grind.node_def [u]) (toExpr n.idx) type e + return m!"{d}" + +@[inherit_doc Goal.ppENodeRef] +def ppENodeRef (e : Expr) : GoalM MessageData := do + (← get).ppENodeRef e + +/-- Helper function for pretty printing the state for debugging purposes. -/ +private def Goal.ppENodeDeclValue (goal : Goal) (e : Expr) : MetaM MessageData := do + if e.isApp && !(← isLitValue e) then + e.withApp fun f args => do + let r ← if f.isConst then + pure m!"{f}" + else + goal.ppENodeRef f + let mut r := r + for arg in args do + r := r ++ " " ++ (← goal.ppENodeRef arg) + return r + else + ppExpr e + +/-- Helper function for pretty printing the state for debugging purposes. -/ +private def Goal.ppENodeDecl (goal : Goal) (e : Expr) : MetaM MessageData := do + let mut r := m!"{← goal.ppENodeRef e} := {← goal.ppENodeDeclValue e}" + let n ← goal.getENode e + unless isSameExpr e n.root do + r := r ++ m!" ↦ {← goal.ppENodeRef n.root}" + if n.interpreted then + r := r ++ ", [val]" + if n.ctor then + r := r ++ ", [ctor]" + if grind.debug.get (← getOptions) then + if let some target := goal.getTarget? e then + r := r ++ m!" ↝ {← goal.ppENodeRef target}" + return r + +/-- Pretty print goal state for debugging purposes. -/ +def Goal.ppState (goal : Goal) : MetaM MessageData := do + let mut r := m!"Goal:" + for e in goal.exprs do + let node ← goal.getENode e + r := r ++ "\n" ++ (← goal.ppENodeDecl node.self) + let eqcs := goal.getEqcs (sort := true) + for eqc in eqcs do + if eqc.length > 1 then + r := r ++ "\n" ++ "{" ++ (MessageData.joinSep (← eqc.mapM goal.ppENodeRef) ", ") ++ "}" + return r + +def ppGoals (goals : List Goal) : MetaM MessageData := do + let mut r := m!"" + for goal in goals do + let m ← goal.ppState + r := r ++ Format.line ++ m + return r + +private abbrev M := ReaderT Goal (StateT (Array MessageData) MetaM) + +private def pushMsg (m : MessageData) : M Unit := + modify fun s => s.push m + +def ppExprArray (cls : Name) (header : String) (es : Array Expr) (clsElem : Name := Name.mkSimple "_") : MessageData := + let es := es.map (toTraceElem · clsElem) + .trace { cls } header es + +private def ppEqcs : M Unit := do + let mut trueEqc? : Option MessageData := none + let mut falseEqc? : Option MessageData := none + let mut otherEqcs : Array MessageData := #[] + let goal ← read + for eqc in goal.getEqcs (sort := true) do + if Option.isSome <| eqc.find? (·.isTrue) then + let eqc := eqc.filter fun e => !e.isTrue + unless eqc.isEmpty do + trueEqc? := ppExprArray `eqc "True propositions" eqc.toArray `prop + else if Option.isSome <| eqc.find? (·.isFalse) then + let eqc := eqc.filter fun e => !e.isFalse + unless eqc.isEmpty do + falseEqc? := ppExprArray `eqc "False propositions" eqc.toArray `prop + else if let e :: _ :: _ := eqc then + -- We may want to add a flag to pretty print equivalence classes of nested proofs + unless (← isProof e) do + otherEqcs := otherEqcs.push <| .trace { cls := `eqc } (.group ("{" ++ (MessageData.joinSep (eqc.map toMessageData) ("," ++ Format.line)) ++ "}")) #[] + if let some trueEqc := trueEqc? then pushMsg trueEqc + if let some falseEqc := falseEqc? then pushMsg falseEqc + unless otherEqcs.isEmpty do + pushMsg <| .trace { cls := `eqc } "Equivalence classes" otherEqcs + +private def ppEMatchTheorem (thm : EMatchTheorem) : MetaM MessageData := do + let m := m!"{← thm.origin.pp}: {thm.patterns.map ppPattern}" + return .trace { cls := `thm } m #[] + +private def ppActiveTheoremPatterns : M Unit := do + let goal ← read + let m ← goal.ematch.thms.toArray.mapM fun thm => ppEMatchTheorem thm + let m := m ++ (← goal.ematch.newThms.toArray.mapM fun thm => ppEMatchTheorem thm) + unless m.isEmpty do + pushMsg <| .trace { cls := `ematch } "E-matching patterns" m + +private def ppOffset : M Unit := do + unless grind.debug.get (← getOptions) do + return () + let goal ← read + let s := goal.arith.offset + let nodes := s.nodes + if nodes.isEmpty then return () + let model ← Arith.Offset.mkModel goal + if model.isEmpty then return () + let mut ms := #[] + for (e, val) in model do + ms := ms.push <| .trace { cls := `assign } m!"{Arith.quoteIfArithTerm e} := {val}" #[] + pushMsg <| .trace { cls := `offset } "Assignment satisfying offset constraints" ms + +private def ppCutsat : M Unit := do + let goal ← read + let s := goal.arith.cutsat + let nodes := s.varMap + if nodes.isEmpty then return () + let model ← Arith.Cutsat.mkModel goal + if model.isEmpty then return () + let mut ms := #[] + for (e, val) in model do + ms := ms.push <| .trace { cls := `assign } m!"{Arith.quoteIfArithTerm e} := {val}" #[] + pushMsg <| .trace { cls := `cutsat } "Assignment satisfying linear constraints" ms + +private def ppCommRing : M Unit := do + let goal ← read + let some msg ← Arith.CommRing.pp? goal | return () + pushMsg msg + +private def ppLinarith : M Unit := do + let goal ← read + let some msg ← Arith.Linear.pp? goal | return () + pushMsg msg + +private def ppThresholds (c : Grind.Config) : M Unit := do + let goal ← read + let maxGen := goal.exprs.foldl (init := 0) fun g e => + if let some n := goal.getENode? e then + Nat.max g n.generation + else + g + let mut msgs := #[] + if goal.ematch.numInstances ≥ c.instances then + msgs := msgs.push <| .trace { cls := `limit } m!"maximum number of instances generated by E-matching has been reached, threshold: `(instances := {c.instances})`" #[] + if goal.ematch.num ≥ c.ematch then + msgs := msgs.push <| .trace { cls := `limit } m!"maximum number of E-matching rounds has been reached, threshold: `(ematch := {c.ematch})`" #[] + if goal.split.num ≥ c.splits then + msgs := msgs.push <| .trace { cls := `limit } m!"maximum number of case-splits has been reached, threshold: `(splits := {c.splits})`" #[] + if maxGen ≥ c.gen then + msgs := msgs.push <| .trace { cls := `limit } m!"maximum term generation has been reached, threshold: `(gen := {c.gen})`" #[] + if goal.arith.ring.steps ≥ c.ringSteps then + msgs := msgs.push <| .trace { cls := `limit } m!"maximum number of ring steps has been reached, threshold: `(ringSteps := {c.ringSteps})`" #[] + unless msgs.isEmpty do + pushMsg <| .trace { cls := `limits } "Thresholds reached" msgs + +private def ppCasesTrace : M Unit := do + let goal ← read + unless goal.split.trace.isEmpty do + let mut msgs := #[] + for { expr, i , num, source } in goal.split.trace.reverse do + msgs := msgs.push <| .trace { cls := `cases } m!"[{i+1}/{num}]: {expr}" #[ + .trace { cls := `cases } m!"source: {← source.toMessageData}" #[] + ] + pushMsg <| .trace { cls := `cases } "Case analyses" msgs + +def goalToMessageData (goal : Goal) (config : Grind.Config) : MetaM MessageData := goal.mvarId.withContext do + if config.verbose then + let (_, m) ← go goal |>.run #[] + let gm := MessageData.trace { cls := `grind, collapsed := false } "Goal diagnostics" m + let r := m!"{.ofGoal goal.mvarId}\n{gm}" + addMessageContextFull r + else + return .ofGoal goal.mvarId +where + go : M Unit := do + pushMsg <| ppExprArray `facts "Asserted facts" goal.facts.toArray `prop + ppEqcs + ppCasesTrace + ppActiveTheoremPatterns + ppOffset + ppCutsat + ppLinarith + ppCommRing + ppThresholds config + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Parser.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Parser.lean new file mode 100644 index 0000000000000000000000000000000000000000..718cb8cc9e3a7e48863d6ab879cd6df3ca8f9690 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Parser.lean @@ -0,0 +1,19 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Parser.Command + +namespace Lean.Parser.Command +/-! +Builtin parsers for `grind` related commands +-/ +@[builtin_command_parser] def grindPattern := leading_parser + "grind_pattern " >> ident >> darrow >> sepBy1 termParser "," + +@[builtin_command_parser] def initGrindNorm := leading_parser + "init_grind_norm " >> many ident >> "| " >> many ident + +end Lean.Parser.Command diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Proj.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Proj.lean new file mode 100644 index 0000000000000000000000000000000000000000..f587b78eace5a56f560f68b07043e7f4abf2ac15 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Proj.lean @@ -0,0 +1,39 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.ProjFns +import Lean.Meta.Tactic.Grind.Types +import Lean.Meta.Tactic.Grind.Internalize + +namespace Lean.Meta.Grind + +/-- +If `parent` is a projection-application `proj_i c`, +check whether the root of the equivalence class containing `c` is a constructor-application `ctor ... a_i ...`. +If so, internalize the term `proj_i (ctor ... a_i ...)` and add the equality `proj_i (ctor ... a_i ...) = a_i`. +-/ +def propagateProjEq (parent : Expr) : GoalM Unit := do + let .const declName _ := parent.getAppFn | return () + let some info ← getProjectionFnInfo? declName | return () + unless info.numParams + 1 == parent.getAppNumArgs do return () + -- It is wasteful to add equation if `parent` is not the root of its congruence class + unless (← isCongrRoot parent) do return () + let arg := parent.appArg! + let ctor ← getRoot arg + unless ctor.isAppOf info.ctorName do return () + let parentNew ← if isSameExpr arg ctor then + pure parent + else + let parentNew ← shareCommon (mkApp parent.appFn! ctor) + internalize parentNew (← getGeneration parent) + pure parentNew + trace_goal[grind.debug.proj] "{parentNew}" + let idx := info.numParams + info.i + unless idx < ctor.getAppNumArgs do return () + let v := ctor.getArg! idx + pushEq parentNew v (← mkEqRefl v) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Proof.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Proof.lean new file mode 100644 index 0000000000000000000000000000000000000000..ca077ad53c044200bb7d15d3fee7ec2f74cc52c2 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Proof.lean @@ -0,0 +1,297 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind.Lemmas +import Lean.Meta.Tactic.Grind.Types + +namespace Lean.Meta.Grind + +private def isEqProof (h : Expr) : MetaM Bool := do + return (← whnfD (← inferType h)).isAppOf ``Eq + +/-- +Given two applications `a` and `b`, finds the common prefix. +Returns `some (f, n)` if `a` and `b` are of the form +- `f a_1 ... a_n` +- `f b_1 ... b_n` + +and `none` otherwise. +-/ +private partial def findCommonPrefix (a b : Expr) : Option (Expr × Nat) := + if isSameExpr a b then + some (a, 0) + else if a.isApp && b.isApp then + if let some (f, n) := findCommonPrefix a.appFn! b.appFn! then + some (f, n+1) + else + none + else + none + +private def flipProof (h : Expr) (flipped : Bool) (heq : Bool) : MetaM Expr := do + let mut h' := h + if (← pure heq <&&> isEqProof h') then + h' ← mkHEqOfEq h' + if flipped then + if heq then mkHEqSymm h' else mkEqSymm h' + else + return h' + +private def mkRefl (a : Expr) (heq : Bool) : MetaM Expr := + if heq then mkHEqRefl a else mkEqRefl a + +private def mkTrans (h₁ h₂ : Expr) (heq : Bool) : MetaM Expr := + if heq then + mkHEqTrans h₁ h₂ + else + mkEqTrans h₁ h₂ + +private def mkTrans' (h₁ : Option Expr) (h₂ : Expr) (heq : Bool) : MetaM Expr := do + let some h₁ := h₁ | return h₂ + mkTrans h₁ h₂ heq + +/-- +Given `h : a ≍ b`, returns a proof `a = b` if `heq == false`. +Otherwise, it returns `h`. +-/ +private def mkEqOfHEqIfNeeded (h : Expr) (heq : Bool) : MetaM Expr := do + if heq then return h else mkEqOfHEq h (check := false) + +/-- +Given `lhs` and `rhs` that are in the same equivalence class, +find the common expression that are in the paths from `lhs` and `rhs` to +the root of their equivalence class. +Recall that this expression must exist since it is the root itself in the +worst case. +-/ +private def findCommon (lhs rhs : Expr) : GoalM Expr := do + let mut visited : RBMap Nat Expr compare := {} + let mut it := lhs + -- Mark elements found following the path from `lhs` to the root. + repeat + let n ← getENode it + visited := visited.insert n.idx n.self + let some target := n.target? | break + it := target + -- Find the marked element from the path from `rhs` to the root. + it := rhs + repeat + let n ← getENode it + if let some common := visited.find? n.idx then + return common + let some target := n.target? | unreachable! -- + it := target + unreachable! + +/-- +Returns `true` if we can construct a congruence proof for `lhs = rhs` using `congrArg`, `congrFun`, and `congr`. +`f` (`g`) is the function of the `lhs` (`rhs`) application. `numArgs` is the number of arguments. +-/ +private partial def isCongrDefaultProofTarget (lhs rhs : Expr) (f g : Expr) (numArgs : Nat) : GoalM Bool := do + unless isSameExpr f g do return false + let info := (← getFunInfo f).paramInfo + let rec loop (lhs rhs : Expr) (i : Nat) : GoalM Bool := do + if lhs.isApp then + let a₁ := lhs.appArg! + let a₂ := rhs.appArg! + let i := i - 1 + unless isSameExpr a₁ a₂ do + if h : i < info.size then + if info[i].hasFwdDeps then + -- Cannot use `congrArg` because there are forward dependencies + return false + else + return false -- Don't have information about argument + loop lhs.appFn! rhs.appFn! i + else + return true + loop lhs rhs numArgs + +mutual + /-- + Given `lhs` and `rhs` proof terms of the form `nestedProof p hp` and `nestedProof q hq`, + constructs a congruence proof for `nestedProof p hp ≍ nestedProof q hq`. + `p` and `q` are in the same equivalence class. + -/ + private partial def mkNestedProofCongr (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do + let p := lhs.appFn!.appArg! + let hp := lhs.appArg! + let q := rhs.appFn!.appArg! + let hq := rhs.appArg! + let h := mkApp5 (mkConst ``Lean.Grind.nestedProof_congr) p q (← mkEqProofCore p q false) hp hq + mkEqOfHEqIfNeeded h heq + + /-- + Constructs a congruence proof for `lhs` and `rhs` using `congr`, `congrFun`, and `congrArg`. + This function assumes `isCongrDefaultProofTarget` returned `true`. + -/ + private partial def mkCongrDefaultProof (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do + let rec loop (lhs rhs : Expr) : GoalM (Option Expr) := do + if lhs.isApp then + let a₁ := lhs.appArg! + let a₂ := rhs.appArg! + if let some proof ← loop lhs.appFn! rhs.appFn! then + if isSameExpr a₁ a₂ then + mkCongrFun proof a₁ + else + mkCongr proof (← mkEqProofCore a₁ a₂ false) + else if isSameExpr a₁ a₂ then + return none -- refl case + else + mkCongrArg lhs.appFn! (← mkEqProofCore a₁ a₂ false) + else + return none + let r := (← loop lhs rhs).get! + if heq then mkHEqOfEq r else return r + + private partial def mkHCongrProofHelper (thm : CongrTheorem) (lhs rhs : Expr) (i : Nat) : GoalM Expr := do + if i > 0 then + let i := i - 1 + let proof ← mkHCongrProofHelper thm lhs.appFn! rhs.appFn! i + let a₁ := lhs.appArg! + let a₂ := rhs.appArg! + let k := thm.argKinds[i]! + return mkApp3 proof a₁ a₂ (← mkEqProofCore a₁ a₂ (k matches .heq)) + else + return thm.proof + + private partial def mkHCongrProof (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do + let f := lhs.getAppFn + let g := rhs.getAppFn + let numArgs := lhs.getAppNumArgs + assert! rhs.getAppNumArgs == numArgs + let fInfo ← getFunInfo f + if fInfo.getArity < numArgs then + -- Function is over-applied. We try to find a common prefix between `lhs` and `rhs` + -- and compute the congruence theorem from there. + let throwUnsupported {α} : GoalM α := + throwError "`grind` currently cannot build congruence proofs for over-applied terms such as{indentExpr lhs}\nand{indentExpr rhs}" + let some (f, numArgs) := findCommonPrefix lhs rhs + | throwUnsupported + let thm ← try mkHCongrWithArity f numArgs catch _ => throwUnsupported + let proof ← mkHCongrProofHelper thm lhs rhs numArgs + mkEqOfHEqIfNeeded proof heq + else + let thm ← mkHCongrWithArity f numArgs + assert! thm.argKinds.size == numArgs + let proof ← mkHCongrProofHelper thm lhs rhs numArgs + if isSameExpr f g then + mkEqOfHEqIfNeeded proof heq + else + /- + `lhs` is of the form `f a_1 ... a_n` + `rhs` is of the form `g b_1 ... b_n` + `proof : f a_1 ... a_n ≍ f b_1 ... b_n` + We construct a proof for `f a_1 ... a_n ≍ g b_1 ... b_n` using `Eq.ndrec` + -/ + let motive ← withLocalDeclD (← mkFreshUserName `x) (← inferType f) fun x => do + mkLambdaFVars #[x] (← mkHEq lhs (mkAppN x rhs.getAppArgs)) + let fEq ← mkEqProofCore f g false + let proof ← mkEqNDRec motive proof fEq + mkEqOfHEqIfNeeded proof heq + + private partial def mkEqCongrProof (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do + let_expr f@Eq α₁ a₁ b₁ := lhs | unreachable! + let_expr Eq α₂ a₂ b₂ := rhs | unreachable! + let us := f.constLevels! + if !isSameExpr α₁ α₂ then + mkHCongrProof lhs rhs heq + else if (← get).hasSameRoot a₁ a₂ && (← get).hasSameRoot b₁ b₂ then + return mkApp7 (mkConst ``Grind.eq_congr us) α₁ a₁ b₁ a₂ b₂ (← mkEqProofCore a₁ a₂ false) (← mkEqProofCore b₁ b₂ false) + else + assert! (← get).hasSameRoot a₁ b₂ && (← get).hasSameRoot b₁ a₂ + return mkApp7 (mkConst ``Grind.eq_congr' us) α₁ a₁ b₁ a₂ b₂ (← mkEqProofCore a₁ b₂ false) (← mkEqProofCore b₁ a₂ false) + + /-- Constructs a congruence proof for `lhs` and `rhs`. -/ + private partial def mkCongrProof (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do + if let .forallE _ p₁ q₁ _ := lhs then + let .forallE _ p₂ q₂ _ := rhs | unreachable! + let u ← withDefault <| getLevel p₁ + let v ← withDefault <| getLevel q₁ + return mkApp6 (mkConst ``implies_congr [u, v]) p₁ p₂ q₁ q₂ (← mkEqProofCore p₁ p₂ false) (← mkEqProofCore q₁ q₂ false) + else + let f := lhs.getAppFn + let g := rhs.getAppFn + let numArgs := lhs.getAppNumArgs + assert! rhs.getAppNumArgs == numArgs + if f.isConstOf ``Lean.Grind.nestedProof && g.isConstOf ``Lean.Grind.nestedProof && numArgs == 2 then + mkNestedProofCongr lhs rhs heq + else if f.isConstOf ``Eq && g.isConstOf ``Eq && numArgs == 3 then + mkEqCongrProof lhs rhs heq + else if (← isCongrDefaultProofTarget lhs rhs f g numArgs) then + mkCongrDefaultProof lhs rhs heq + else + mkHCongrProof lhs rhs heq + + private partial def realizeEqProof (lhs rhs : Expr) (h : Expr) (flipped : Bool) (heq : Bool) : GoalM Expr := do + let h ← if h == congrPlaceholderProof then + mkCongrProof lhs rhs heq + else + flipProof h flipped heq + + /-- Given `acc : lhs₀ = lhs`, returns a proof of `lhs₀ = common`. -/ + private partial def mkProofTo (lhs : Expr) (common : Expr) (acc : Option Expr) (heq : Bool) : GoalM (Option Expr) := do + if isSameExpr lhs common then + return acc + let n ← getENode lhs + let some target := n.target? | unreachable! + let some h := n.proof? | unreachable! + let h ← realizeEqProof lhs target h n.flipped heq + -- h : lhs = target + let acc ← mkTrans' acc h heq + mkProofTo target common (some acc) heq + + /-- Given `lhsEqCommon : lhs = common`, returns a proof for `lhs = rhs`. -/ + private partial def mkProofFrom (rhs : Expr) (common : Expr) (lhsEqCommon? : Option Expr) (heq : Bool) : GoalM (Option Expr) := do + if isSameExpr rhs common then + return lhsEqCommon? + let n ← getENode rhs + let some target := n.target? | unreachable! + let some h := n.proof? | unreachable! + let h ← realizeEqProof target rhs h (!n.flipped) heq + -- `h : target = rhs` + let h' ← mkProofFrom target common lhsEqCommon? heq + -- `h' : lhs = target` + mkTrans' h' h heq + + /-- + Returns a proof of `lhs = rhs` (`lhs ≍ rhs`) if `heq = false` (`heq = true`). + If `heq = false`, this function assumes that `lhs` and `rhs` have the same type. + -/ + private partial def mkEqProofCore (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do + if isSameExpr lhs rhs then + return (← mkRefl lhs heq) + -- The equivalence class contains `HEq` proofs. So, we build a proof using HEq. Otherwise, we use `Eq`. + let heqProofs := (← getRootENode lhs).heqProofs + let n₁ ← getENode lhs + let n₂ ← getENode rhs + assert! isSameExpr n₁.root n₂.root + let common ← findCommon lhs rhs + let lhsEqCommon? ← mkProofTo lhs common none heqProofs + let some lhsEqRhs ← mkProofFrom rhs common lhsEqCommon? heqProofs | unreachable! + if heq == heqProofs then + return lhsEqRhs + else if heq then + mkHEqOfEq lhsEqRhs + else + mkEqOfHEq lhsEqRhs (check := false) + +end + +/-- +Returns a proof that `a = b`. +It assumes `a` and `b` are in the same equivalence class. +-/ +@[export lean_grind_mk_eq_proof] +def mkEqProofImpl (a b : Expr) : GoalM Expr := do + assert! (← hasSameType a b) + mkEqProofCore a b (heq := false) + +@[export lean_grind_mk_heq_proof] +def mkHEqProofImpl (a b : Expr) : GoalM Expr := + mkEqProofCore a b (heq := true) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Propagate.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Propagate.lean new file mode 100644 index 0000000000000000000000000000000000000000..31b5efd3654d7d051c95462ea5dd1926d4813529 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Grind/Propagate.lean @@ -0,0 +1,388 @@ +/- +Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Grind +import Lean.Meta.Tactic.Grind.Proof +import Lean.Meta.Tactic.Grind.PropagatorAttr +import Lean.Meta.Tactic.Grind.Simp +import Lean.Meta.Tactic.Grind.Ext +import Lean.Meta.Tactic.Grind.Internalize + +namespace Lean.Meta.Grind + +/-- +Propagates equalities for a conjunction `a ∧ b` based on the truth values +of its components `a` and `b`. This function checks the truth value of `a` and `b`, +and propagates the following equalities: + +- If `a = True`, propagates `(a ∧ b) = b`. +- If `b = True`, propagates `(a ∧ b) = a`. +- If `a = False`, propagates `(a ∧ b) = False`. +- If `b = False`, propagates `(a ∧ b) = False`. +-/ +builtin_grind_propagator propagateAndUp ↑And := fun e => do + let_expr And a b := e | return () + if (← isEqTrue a) then + -- a = True → (a ∧ b) = b + pushEq e b <| mkApp3 (mkConst ``Grind.and_eq_of_eq_true_left) a b (← mkEqTrueProof a) + else if (← isEqTrue b) then + -- b = True → (a ∧ b) = a + pushEq e a <| mkApp3 (mkConst ``Grind.and_eq_of_eq_true_right) a b (← mkEqTrueProof b) + else if (← isEqFalse a) then + -- a = False → (a ∧ b) = False + pushEqFalse e <| mkApp3 (mkConst ``Grind.and_eq_of_eq_false_left) a b (← mkEqFalseProof a) + else if (← isEqFalse b) then + -- b = False → (a ∧ b) = False + pushEqFalse e <| mkApp3 (mkConst ``Grind.and_eq_of_eq_false_right) a b (← mkEqFalseProof b) + +/-- +Propagates truth values downwards for a conjunction `a ∧ b` when the +expression itself is known to be `True`. +-/ +builtin_grind_propagator propagateAndDown ↓And := fun e => do + if (← isEqTrue e) then + let_expr And a b := e | return () + let h ← mkEqTrueProof e + pushEqTrue a <| mkApp3 (mkConst ``Grind.eq_true_of_and_eq_true_left) a b h + pushEqTrue b <| mkApp3 (mkConst ``Grind.eq_true_of_and_eq_true_right) a b h + +/-- +Propagates equalities for a disjunction `a ∨ b` based on the truth values +of its components `a` and `b`. This function checks the truth value of `a` and `b`, +and propagates the following equalities: + +- If `a = False`, propagates `(a ∨ b) = b`. +- If `b = False`, propagates `(a ∨ b) = a`. +- If `a = True`, propagates `(a ∨ b) = True`. +- If `b = True`, propagates `(a ∨ b) = True`. +-/ +builtin_grind_propagator propagateOrUp ↑Or := fun e => do + let_expr Or a b := e | return () + if (← isEqFalse a) then + -- a = False → (a ∨ b) = b + pushEq e b <| mkApp3 (mkConst ``Grind.or_eq_of_eq_false_left) a b (← mkEqFalseProof a) + else if (← isEqFalse b) then + -- b = False → (a ∨ b) = a + pushEq e a <| mkApp3 (mkConst ``Grind.or_eq_of_eq_false_right) a b (← mkEqFalseProof b) + else if (← isEqTrue a) then + -- a = True → (a ∨ b) = True + pushEqTrue e <| mkApp3 (mkConst ``Grind.or_eq_of_eq_true_left) a b (← mkEqTrueProof a) + else if (← isEqTrue b) then + -- b = True → (a ∧ b) = True + pushEqTrue e <| mkApp3 (mkConst ``Grind.or_eq_of_eq_true_right) a b (← mkEqTrueProof b) + +/-- +Propagates truth values downwards for a disjuction `a ∨ b` when the +expression itself is known to be `False`. +-/ +builtin_grind_propagator propagateOrDown ↓Or := fun e => do + if (← isEqFalse e) then + let_expr Or a b := e | return () + let h ← mkEqFalseProof e + pushEqFalse a <| mkApp3 (mkConst ``Grind.eq_false_of_or_eq_false_left) a b h + pushEqFalse b <| mkApp3 (mkConst ``Grind.eq_false_of_or_eq_false_right) a b h + +/-- +Propagates equalities for a negation `Not a` based on the truth value of `a`. +This function checks the truth value of `a` and propagates the following equalities: + +- If `a = False`, propagates `(Not a) = True`. +- If `a = True`, propagates `(Not a) = False`. +-/ +builtin_grind_propagator propagateNotUp ↑Not := fun e => do + let_expr Not a := e | return () + if (← isEqFalse a) then + -- a = False → (Not a) = True + pushEqTrue e <| mkApp2 (mkConst ``Grind.not_eq_of_eq_false) a (← mkEqFalseProof a) + else if (← isEqTrue a) then + -- a = True → (Not a) = False + pushEqFalse e <| mkApp2 (mkConst ``Grind.not_eq_of_eq_true) a (← mkEqTrueProof a) + else if (← isEqv e a) then + closeGoal <| mkApp2 (mkConst ``Grind.false_of_not_eq_self) a (← mkEqProof e a) + +/-- +Propagates truth values downwards for a negation expression `Not a` based on the truth value of `Not a`. +This function performs the following: + +- If `(Not a) = False`, propagates `a = True`. +- If `(Not a) = True`, propagates `a = False`. +-/ +builtin_grind_propagator propagateNotDown ↓Not := fun e => do + let_expr Not a := e | return () + if (← isEqFalse e) then + pushEqTrue a <| mkApp2 (mkConst ``Grind.eq_true_of_not_eq_false) a (← mkEqFalseProof e) + else if (← isEqTrue e) then + pushEqFalse a <| mkApp2 (mkConst ``Grind.eq_false_of_not_eq_true) a (← mkEqTrueProof e) + else if (← isEqv e a) then + closeGoal <| mkApp2 (mkConst ``Grind.false_of_not_eq_self) a (← mkEqProof e a) + +def propagateBoolDiseq (a : Expr) : GoalM Unit := do + if let some h ← mkDiseqProof? a (← getBoolFalseExpr) then + pushEqBoolTrue a <| mkApp2 (mkConst ``Grind.Bool.eq_true_of_not_eq_false') a h + if let some h ← mkDiseqProof? a (← getBoolTrueExpr) then + pushEqBoolFalse a <| mkApp2 (mkConst ``Grind.Bool.eq_false_of_not_eq_true') a h + +/-- Propagates `Eq` upwards -/ +builtin_grind_propagator propagateEqUp ↑Eq := fun e => do + let_expr Eq α a b := e | return () + if (← isEqTrue a) then + pushEq e b <| mkApp3 (mkConst ``Grind.eq_eq_of_eq_true_left) a b (← mkEqTrueProof a) + else if (← isEqTrue b) then + pushEq e a <| mkApp3 (mkConst ``Grind.eq_eq_of_eq_true_right) a b (← mkEqTrueProof b) + else if (← isEqv a b) then + pushEqTrue e <| mkEqTrueCore e (← mkEqProof a b) + if α.isConstOf ``Bool then + if (← isEqFalse e) then + propagateBoolDiseq a + propagateBoolDiseq b + let aRoot ← getRootENode a + let bRoot ← getRootENode b + if aRoot.ctor && bRoot.ctor && aRoot.self.getAppFn != bRoot.self.getAppFn then + -- ¬a = b + let hne ← withLocalDeclD `h (← mkEq a b) fun h => do + let hf ← mkEqTrans (← mkEqProof aRoot.self a) h + let hf ← mkEqTrans hf (← mkEqProof b bRoot.self) + let hf ← mkNoConfusion (← getFalseExpr) hf + mkLambdaFVars #[h] hf + pushEqFalse e <| mkApp2 (mkConst ``eq_false) e hne + +/-- Propagates `Eq` downwards -/ +builtin_grind_propagator propagateEqDown ↓Eq := fun e => do + if (← isEqTrue e) then + let_expr Eq _ a b := e | return () + pushEq a b <| mkOfEqTrueCore e (← mkEqTrueProof e) + else if (← isEqFalse e) then + let_expr Eq α lhs rhs := e | return () + if α.isConstOf ``Bool then + propagateBoolDiseq lhs + propagateBoolDiseq rhs + propagateCutsatDiseq lhs rhs + propagateCommRingDiseq lhs rhs + propagateLinarithDiseq lhs rhs + let thms ← getExtTheorems α + if !thms.isEmpty then + /- + Heuristic for lists: If `lhs` or `rhs` are constructors we do not apply extensionality theorems. + For example, we don't want to apply the extensionality theorem to things like `xs ≠ []`. + TODO: polish this hackish heuristic later. + -/ + if α.isAppOf ``List && ((← getRootENode lhs).ctor || (← getRootENode rhs).ctor) then + return () + for thm in (← getExtTheorems α) do + instantiateExtTheorem thm e + +private def getLawfulBEqInst? (u : List Level) (α : Expr) (binst : Expr) : MetaM (Option Expr) := do + let lawfulBEq := mkApp2 (mkConst ``LawfulBEq u) α binst + let .some linst ← trySynthInstance lawfulBEq | return none + return some linst + +/- +Note about `BEq.beq` +Given `a b : α` in a context where we have `[BEq α] [LawfulBEq α]` +The normalizer (aka `simp`) fails to normalize `if a == b then ... else ...` to `if a = b then ... else ...` using +``` +theorem beq_iff_eq [BEq α] [LawfulBEq α] {a b : α} : a == b ↔ a = b := + ⟨eq_of_beq, beq_of_eq⟩ +``` +The main issue is that `ite_congr` requires that the resulting proposition to be decidable, +and we don't have `[DecidableEq α]`. Thus, the normalization step fails. +The following propagators for `BEq.beq` ensure `grind` does not assume this normalization +rule has been applied. +-/ + +builtin_grind_propagator propagateBEqUp ↑BEq.beq := fun e => do + /- + `grind` uses the normalization rule `Bool.beq_eq_decide_eq`, but it is only applicable if + the type implements the instances `BEq`, `LawfulBEq`, **and** `DecidableEq α`. + However, we may be in a context where only `BEq` and `LawfulBEq` are available. + Thus, we have added this propagator as a backup. + -/ + let_expr f@BEq.beq α binst a b := e | return () + let u := f.constLevels! + if (← isEqv a b) then + let some linst ← getLawfulBEqInst? u α binst | return () + pushEqBoolTrue e <| mkApp6 (mkConst ``Grind.beq_eq_true_of_eq u) α binst linst a b (← mkEqProof a b) + else if let some h ← mkDiseqProof? a b then + let some linst ← getLawfulBEqInst? u α binst | return () + pushEqBoolFalse e <| mkApp6 (mkConst ``Grind.beq_eq_false_of_diseq u) α binst linst a b h + +builtin_grind_propagator propagateBEqDown ↓BEq.beq := fun e => do + /- See comment above -/ + let_expr f@BEq.beq α binst a b := e | return () + let u := f.constLevels! + if (← isEqBoolTrue e) then + let some linst ← getLawfulBEqInst? u α binst | return () + pushEq a b <| mkApp6 (mkConst ``Grind.eq_of_beq_eq_true u) α binst linst a b (← mkEqProof e (← getBoolTrueExpr)) + else if (← isEqBoolFalse e) then + let some linst ← getLawfulBEqInst? u α binst | return () + let eq ← shareCommon (mkApp3 (mkConst ``Eq [u.head!.succ]) α a b) + internalize eq (← getGeneration a) + pushEqFalse eq <| mkApp6 (mkConst ``Grind.ne_of_beq_eq_false u) α binst linst a b (← mkEqProof e (← getBoolFalseExpr)) + +/-- Propagates `EqMatch` downwards -/ +builtin_grind_propagator propagateEqMatchDown ↓Grind.EqMatch := fun e => do + if (← isEqTrue e) then + let_expr Grind.EqMatch _ a b origin := e | return () + markCaseSplitAsResolved origin + pushEq a b <| mkOfEqTrueCore e (← mkEqTrueProof e) + +/-- Propagates `HEq` downwards -/ +builtin_grind_propagator propagateHEqDown ↓HEq := fun e => do + if (← isEqTrue e) then + let_expr HEq _ a _ b := e | return () + pushHEq a b <| mkOfEqTrueCore e (← mkEqTrueProof e) + +/-- Propagates `HEq` upwards -/ +builtin_grind_propagator propagateHEqUp ↑HEq := fun e => do + let_expr HEq _ a _ b := e | return () + if (← isEqv a b) then + pushEqTrue e <| mkEqTrueCore e (← mkHEqProof a b) + +/-- +Helper function for propagating over-applied `ite` and `dite`-applications. +`h` is a proof for the `e`'s prefix (of size `prefixSize`) that is equal to `rhs`. +`args` contains all arguments of `e`. +`prefixSize <= args.size` +-/ +private def applyCongrFun (e rhs : Expr) (h : Expr) (prefixSize : Nat) (args : Array Expr) : GoalM Unit := do + if prefixSize == args.size then + internalize rhs (← getGeneration e) + pushEq e rhs h + else + go rhs h prefixSize +where + go (rhs : Expr) (h : Expr) (i : Nat) : GoalM Unit := do + if _h : i < args.size then + let arg := args[i] + let rhs' := mkApp rhs arg + let h' ← mkCongrFun h arg + go rhs' h' (i+1) + else + let rhs ← preprocessLight rhs + internalize rhs (← getGeneration e) + pushEq e rhs h + +/-- Propagates `ite` upwards -/ +builtin_grind_propagator propagateIte ↑ite := fun e => do + let numArgs := e.getAppNumArgs + if numArgs < 5 then return () + let c := e.getArg! 1 numArgs + if (← isEqTrue c) then + let f := e.getAppFn + let args := e.getAppArgs + let rhs := args[3]! + let h := mkApp (mkAppRange (mkConst ``ite_cond_eq_true f.constLevels!) 0 5 args) (← mkEqTrueProof c) + applyCongrFun e rhs h 5 args + else if (← isEqFalse c) then + let f := e.getAppFn + let args := e.getAppArgs + let rhs := args[4]! + let h := mkApp (mkAppRange (mkConst ``ite_cond_eq_false f.constLevels!) 0 5 args) (← mkEqFalseProof c) + applyCongrFun e rhs h 5 args + +/-- Propagates `dite` upwards -/ +builtin_grind_propagator propagateDIte ↑dite := fun e => do + let numArgs := e.getAppNumArgs + if numArgs < 5 then return () + let c := e.getArg! 1 numArgs + if (← isEqTrue c) then + let f := e.getAppFn + let args := e.getAppArgs + let h₁ ← mkEqTrueProof c + let ah₁ := mkApp args[3]! (mkOfEqTrueCore c h₁) + let p ← preprocess ah₁ + let r := p.expr + let h₂ ← p.getProof + let h := mkApp3 (mkAppRange (mkConst ``Grind.dite_cond_eq_true' f.constLevels!) 0 5 args) r h₁ h₂ + applyCongrFun e r h 5 args + else if (← isEqFalse c) then + let f := e.getAppFn + let args := e.getAppArgs + let h₁ ← mkEqFalseProof c + let bh₁ := mkApp args[4]! (mkOfEqFalseCore c h₁) + let p ← preprocess bh₁ + let r := p.expr + let h₂ ← p.getProof + let h := mkApp3 (mkAppRange (mkConst ``Grind.dite_cond_eq_false' f.constLevels!) 0 5 args) r h₁ h₂ + applyCongrFun e r h 5 args + +builtin_grind_propagator propagateDecideDown ↓decide := fun e => do + let root ← getRootENode e + unless root.ctor do return () + let_expr decide p h := e | return () + if root.self.isConstOf ``true then + pushEqTrue p <| mkApp3 (mkConst ``Grind.of_decide_eq_true) p h (← mkEqProof e root.self) + else if root.self.isConstOf ``false then + pushEqFalse p <| mkApp3 (mkConst ``Grind.of_decide_eq_false) p h (← mkEqProof e root.self) + +builtin_grind_propagator propagateDecideUp ↑decide := fun e => do + let_expr decide p h := e | return () + if (← isEqTrue p) then + pushEq e (← getBoolTrueExpr) <| mkApp3 (mkConst ``Grind.decide_eq_true) p h (← mkEqTrueProof p) + else if (← isEqFalse p) then + pushEq e (← getBoolFalseExpr) <| mkApp3 (mkConst ``Grind.decide_eq_false) p h (← mkEqFalseProof p) + +/-- `Bool` version of `propagateAndUp` -/ +builtin_grind_propagator propagateBoolAndUp ↑Bool.and := fun e => do + let_expr Bool.and a b := e | return () + if (← isEqBoolTrue a) then + pushEq e b <| mkApp3 (mkConst ``Grind.Bool.and_eq_of_eq_true_left) a b (← mkEqBoolTrueProof a) + else if (← isEqBoolTrue b) then + pushEq e a <| mkApp3 (mkConst ``Grind.Bool.and_eq_of_eq_true_right) a b (← mkEqBoolTrueProof b) + else if (← isEqBoolFalse a) then + pushEqBoolFalse e <| mkApp3 (mkConst ``Grind.Bool.and_eq_of_eq_false_left) a b (← mkEqBoolFalseProof a) + else if (← isEqBoolFalse b) then + pushEqBoolFalse e <| mkApp3 (mkConst ``Grind.Bool.and_eq_of_eq_false_right) a b (← mkEqBoolFalseProof b) + +/-- `Bool` version of `propagateAndDown` -/ +builtin_grind_propagator propagateBoolAndDown ↓Bool.and := fun e => do + if (← isEqBoolTrue e) then + let_expr Bool.and a b := e | return () + let h ← mkEqBoolTrueProof e + pushEqBoolTrue a <| mkApp3 (mkConst ``Grind.Bool.eq_true_of_and_eq_true_left) a b h + pushEqBoolTrue b <| mkApp3 (mkConst ``Grind.Bool.eq_true_of_and_eq_true_right) a b h + +/-- `Bool` version of `propagateOrUp` -/ +builtin_grind_propagator propagateBoolOrUp ↑Bool.or := fun e => do + let_expr Bool.or a b := e | return () + if (← isEqBoolFalse a) then + pushEq e b <| mkApp3 (mkConst ``Grind.Bool.or_eq_of_eq_false_left) a b (← mkEqBoolFalseProof a) + else if (← isEqBoolFalse b) then + pushEq e a <| mkApp3 (mkConst ``Grind.Bool.or_eq_of_eq_false_right) a b (← mkEqBoolFalseProof b) + else if (← isEqBoolTrue a) then + pushEqBoolTrue e <| mkApp3 (mkConst ``Grind.Bool.or_eq_of_eq_true_left) a b (← mkEqBoolTrueProof a) + else if (← isEqBoolTrue b) then + pushEqBoolTrue e <| mkApp3 (mkConst ``Grind.Bool.or_eq_of_eq_true_right) a b (← mkEqBoolTrueProof b) + +/-- `Bool` version of `propagateOrDown` -/ +builtin_grind_propagator propagateBoolOrDown ↓Bool.or := fun e => do + if (← isEqBoolFalse e) then + let_expr Bool.or a b := e | return () + let h ← mkEqBoolFalseProof e + pushEqBoolFalse a <| mkApp3 (mkConst ``Grind.Bool.eq_false_of_or_eq_false_left) a b h + pushEqBoolFalse b <| mkApp3 (mkConst ``Grind.Bool.eq_false_of_or_eq_false_right) a b h + +/-- `Bool` version of `propagateNotUp` -/ +builtin_grind_propagator propagateBoolNotUp ↑Bool.not := fun e => do + let_expr Bool.not a := e | return () + if (← isEqBoolFalse a) then + pushEqBoolTrue e <| mkApp2 (mkConst ``Grind.Bool.not_eq_of_eq_false) a (← mkEqBoolFalseProof a) + else if (← isEqBoolTrue a) then + pushEqBoolFalse e <| mkApp2 (mkConst ``Grind.Bool.not_eq_of_eq_true) a (← mkEqBoolTrueProof a) + else if (← isEqv e a) then + closeGoal <| mkApp2 (mkConst ``Grind.Bool.false_of_not_eq_self) a (← mkEqProof e a) + +/-- `Bool` version of `propagateNotDown` -/ +builtin_grind_propagator propagateBoolNotDown ↓Bool.not := fun e => do + let_expr Bool.not a := e | return () + if (← isEqBoolFalse e) then + pushEqBoolTrue a <| mkApp2 (mkConst ``Grind.Bool.eq_true_of_not_eq_false) a (← mkEqBoolFalseProof e) + else if (← isEqBoolTrue e) then + pushEqBoolFalse a <| mkApp2 (mkConst ``Grind.Bool.eq_false_of_not_eq_true) a (← mkEqBoolTrueProof e) + else if (← isEqv e a) then + closeGoal <| mkApp2 (mkConst ``Grind.Bool.false_of_not_eq_self) a (← mkEqProof e a) + +end Lean.Meta.Grind diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/IndependentOf.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/IndependentOf.lean new file mode 100644 index 0000000000000000000000000000000000000000..c18bb81f492f4eeb43dc7efb39f9320761be597b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/IndependentOf.lean @@ -0,0 +1,58 @@ +/- +Copyright (c) 2022 Mario Carneiro. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Mario Carneiro, Jannis Limperg +-/ +prelude +import Lean.Meta.CollectMVars +import Lean.Meta.Tactic.Util + +namespace Lean.MVarId + +open Meta + +/-- +Check if a goal is "independent" of a list of other goals. +We say a goal is independent of other goals if assigning a value to it +can not change the assignability of the other goals. + +Examples: +* `?m_1 : Type` is not independent of `?m_2 : ?m_1`, + because we could assign `true : Bool` to `?m_2`, + but if we first assign `Nat` to `?m_1` then that is no longer possible. +* `?m_1 : Nat` is not independent of `?m_2 : Fin ?m_1`, + because we could assign `37 : Fin 42` to `?m_2`, + but if we first assign `2` to `?m_1` then that is no longer possible. +* `?m_1 : ?m_2` is not independent of `?m_2 : Type`, because we could assign `Bool` to ?m_2`, + but if we first assign `0 : Nat` to `?m_1` then that is no longer possible. +* Given `P : Prop` and `f : P → Type`, `?m_1 : P` is independent of `?m_2 : f ?m_1` + by proof irrelevance. +* Similarly given `f : Fin 0 → Type`, `?m_1 : Fin 0` is independent of `?m_2 : f ?m_1`, + because `Fin 0` is a subsingleton. +* Finally `?m_1 : Nat` is independent of `?m_2 : α`, + as long as `?m_1` does not appear in `Meta.getMVars α` + (note that `Meta.getMVars` follows delayed assignments). + +This function only calculates a conservative approximation of this condition. +That is, it may return `false` when it should return `true`. +(In particular it returns false whenever the type of `g` contains a metavariable, +regardless of whether this is related to the metavariables in `L`.) +-/ +def isIndependentOf (L : List MVarId) (g : MVarId) : MetaM Bool := g.withContext do + let t ← instantiateMVars (← g.getType) + if t.hasExprMVar then + -- If the goal's type contains other meta-variables, + -- we conservatively say that `g` is not independent. + -- It would be possible to check if `L` depends on these meta-variables. + return false + if (← inferType t).isProp then + -- If the goal is propositional, + -- proof-irrelevance ensures it is independent of any other goals. + return true + if ← g.isSubsingleton then + -- If the goal is a subsingleton, it is independent of any other goals. + return true + -- Finally, we check if the goal `g` appears in the type of any of the goals `L`. + L.allM fun g' => do pure !((← getMVarDependencies g').contains g) + +end Lean.MVarId diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Induction.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Induction.lean new file mode 100644 index 0000000000000000000000000000000000000000..13631015b36d7ade60a262cce88090f915abe781 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Induction.lean @@ -0,0 +1,238 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.RecursorInfo +import Lean.Meta.SynthInstance +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Revert +import Lean.Meta.Tactic.Intro +import Lean.Meta.Tactic.Clear +import Lean.Meta.Tactic.FVarSubst + +namespace Lean.Meta + +private partial def getTargetArity : Expr → Nat + | Expr.mdata _ b => getTargetArity b + | Expr.forallE _ _ b _ => getTargetArity b + 1 + | e => if e.isHeadBetaTarget then getTargetArity e.headBeta else 0 + +private def addRecParams (mvarId : MVarId) (majorTypeArgs : Array Expr) : List (Option Nat) → Expr → MetaM Expr + | [], recursor => pure recursor + | some pos :: rest, recursor => + if h : pos < majorTypeArgs.size then + addRecParams mvarId majorTypeArgs rest (mkApp recursor (majorTypeArgs[pos])) + else + throwTacticEx `induction mvarId "ill-formed recursor" + | none :: rest, recursor => do + let recursorType ← inferType recursor + let recursorType ← whnfForall recursorType + match recursorType with + | Expr.forallE _ d _ _ => do + let param ← try synthInstance d catch _ => throwTacticEx `induction mvarId "failed to generate type class instance parameter" + addRecParams mvarId majorTypeArgs rest (mkApp recursor param) + | _ => + throwTacticEx `induction mvarId "ill-formed recursor" + +structure InductionSubgoal where + mvarId : MVarId + fields : Array Expr := #[] + subst : FVarSubst := {} + deriving Inhabited + +private def getTypeBody (mvarId : MVarId) (type : Expr) (x : Expr) : MetaM Expr := do + let type ← whnfForall type + match type with + | Expr.forallE _ _ b _ => pure $ b.instantiate1 x + | _ => throwTacticEx `induction mvarId "ill-formed recursor" + +structure AltVarNames where + explicit : Bool := false -- true if `@` modifier was used + varNames : List Name := [] + deriving Inhabited + +private partial def finalize + (mvarId : MVarId) (givenNames : Array AltVarNames) (recursorInfo : RecursorInfo) + (reverted : Array FVarId) (major : Expr) (indices : Array Expr) (baseSubst : FVarSubst) (recursor : Expr) + : MetaM (Array InductionSubgoal) := do + let target ← mvarId.getType + let initialArity := getTargetArity target + let recursorType ← inferType recursor + let numMinors := recursorInfo.produceMotive.length + let rec loop (pos : Nat) (minorIdx : Nat) (recursor recursorType : Expr) (consumedMajor : Bool) (subgoals : Array InductionSubgoal) := do + let recursorType ← whnfForall recursorType + if recursorType.isForall && pos < recursorInfo.numArgs then + if pos == recursorInfo.firstIndexPos then + let (recursor, recursorType) ← indices.foldlM (init := (recursor, recursorType)) fun (recursor, recursorType) index => do + let recursor := mkApp recursor index + let recursorType ← getTypeBody mvarId recursorType index + pure (recursor, recursorType) + let recursor := mkApp recursor major + let recursorType ← getTypeBody mvarId recursorType major + loop (pos+1+indices.size) minorIdx recursor recursorType true subgoals + else + -- consume motive + let tag ← mvarId.getTag + if minorIdx ≥ numMinors then throwTacticEx `induction mvarId "ill-formed recursor" + match recursorType with + | Expr.forallE n d _ c => + let d := d.headBeta + -- Remark is givenNames is not empty, then user provided explicit alternatives for each minor premise + if c.isInstImplicit && givenNames.isEmpty then + match (← synthInstance? d) with + | some inst => + let recursor := mkApp recursor inst + let recursorType ← getTypeBody mvarId recursorType inst + loop (pos+1) (minorIdx+1) recursor recursorType consumedMajor subgoals + | none => do + -- Add newSubgoal if type class resolution failed + let mvar ← mkFreshExprSyntheticOpaqueMVar d (tag ++ n) + let recursor := mkApp recursor mvar + let recursorType ← getTypeBody mvarId recursorType mvar + loop (pos+1) (minorIdx+1) recursor recursorType consumedMajor (subgoals.push { mvarId := mvar.mvarId! }) + else + let arity := getTargetArity d + if arity < initialArity then throwTacticEx `induction mvarId "ill-formed recursor" + let nparams := arity - initialArity -- number of fields due to minor premise + let nextra := reverted.size - indices.size - 1 -- extra dependencies that have been reverted + let minorGivenNames := if h : minorIdx < givenNames.size then givenNames[minorIdx] else {} + let mvar ← mkFreshExprSyntheticOpaqueMVar d (tag ++ n) + let recursor := mkApp recursor mvar + let recursorType ← getTypeBody mvarId recursorType mvar + -- Try to clear major premise from new goal + let mvarId' ← mvar.mvarId!.tryClear major.fvarId! + let (fields, mvarId') ← mvarId'.introN nparams minorGivenNames.varNames (useNamesForExplicitOnly := !minorGivenNames.explicit) + let (extra, mvarId') ← mvarId'.introNP nextra + let subst := reverted.size.fold (init := baseSubst) fun i _ (subst : FVarSubst) => + if i < indices.size + 1 then subst + else + let revertedFVarId := reverted[i] + let newFVarId := extra[i - indices.size - 1]! + subst.insert revertedFVarId (mkFVar newFVarId) + let fields := fields.map mkFVar + loop (pos+1) (minorIdx+1) recursor recursorType consumedMajor (subgoals.push { mvarId := mvarId', fields := fields, subst := subst }) + | _ => unreachable! + else + unless consumedMajor do throwTacticEx `induction mvarId "ill-formed recursor" + mvarId.assign recursor + pure subgoals + loop (recursorInfo.paramsPos.length + 1) 0 recursor recursorType false #[] + +private def throwUnexpectedMajorType (tacticName : Name) (mvarId : MVarId) (majorType : Expr) : MetaM α := + throwTacticEx tacticName mvarId m!"unexpected major premise type{indentExpr majorType}" + +/-- +Auxiliary method for implementing `induction`-like tactics. +It retrieves indices from `majorType` using `recursorInfo`. +Remark: `mvarId` and `tacticName` are used to generate error messages. +-/ +def getMajorTypeIndices (mvarId : MVarId) (tacticName : Name) (recursorInfo : RecursorInfo) (majorType : Expr) : MetaM (Array Expr) := do + let majorTypeArgs := majorType.getAppArgs + recursorInfo.indicesPos.toArray.mapM fun idxPos => do + if h : idxPos ≥ majorTypeArgs.size then + throwTacticEx tacticName mvarId m!"major premise type is ill-formed{indentExpr majorType}" + else + let idx := majorTypeArgs[idxPos] + unless idx.isFVar do throwTacticEx tacticName mvarId m!"major premise type index {idx} is not a variable{indentExpr majorType}" + majorTypeArgs.size.forM fun i _ => do + let arg := majorTypeArgs[i] + if i != idxPos && arg == idx then + throwTacticEx tacticName mvarId m!"'{idx}' is an index in major premise, but it occurs more than once{indentExpr majorType}" + if i < idxPos then + if (← exprDependsOn arg idx.fvarId!) then + throwTacticEx tacticName mvarId m!"'{idx}' is an index in major premise, but it occurs in previous arguments{indentExpr majorType}" + -- If arg is also and index and a variable occurring after `idx`, we need to make sure it doesn't depend on `idx`. + -- Note that if `arg` is not a variable, we will fail anyway when we visit it. + if i > idxPos && recursorInfo.indicesPos.contains i && arg.isFVar then + let idxDecl ← idx.fvarId!.getDecl + if (← localDeclDependsOn idxDecl arg.fvarId!) then + throwTacticEx tacticName mvarId m!"'{idx}' is an index in major premise, but it depends on index occurring at position #{i+1}" + return idx + +/-- +Auxiliary method for implementing `induction`-like tactics. +It creates the prefix of a recursor application up-to `motive`. +The motive is computed by abstracting `major` and `indices` at `mvarId.getType`. +It retrieves indices from `majorType` using `recursorInfo`. +Remark: `mvarId` and `tacticName` are used to generate error messages. +-/ +def mkRecursorAppPrefix (mvarId : MVarId) (tacticName : Name) (majorFVarId : FVarId) (recursorInfo : RecursorInfo) (indices : Array Expr) : MetaM Expr := do + let major := mkFVar majorFVarId + let target ← mvarId.getType + let targetLevel ← getLevel target + let targetLevel ← normalizeLevel targetLevel + let majorLocalDecl ← majorFVarId.getDecl + let some majorType ← whnfUntil majorLocalDecl.type recursorInfo.typeName + | throwUnexpectedMajorType tacticName mvarId majorLocalDecl.type + majorType.withApp fun majorTypeFn majorTypeArgs => do + match majorTypeFn with + | .const _ majorTypeFnLevels => do + let majorTypeFnLevels := majorTypeFnLevels.toArray + let (recursorLevels, foundTargetLevel) ← recursorInfo.univLevelPos.foldlM (init := (#[], false)) + fun (recursorLevels, foundTargetLevel) (univPos : RecursorUnivLevelPos) => do + match univPos with + | RecursorUnivLevelPos.motive => pure (recursorLevels.push targetLevel, true) + | RecursorUnivLevelPos.majorType idx => + if h : idx ≥ majorTypeFnLevels.size then + throwTacticEx tacticName mvarId "ill-formed recursor" + else + pure (recursorLevels.push majorTypeFnLevels[idx], foundTargetLevel) + if !foundTargetLevel && !targetLevel.isZero then + throwTacticEx tacticName mvarId m!"recursor '{recursorInfo.recursorName}' can only eliminate into Prop" + let recursor := mkConst recursorInfo.recursorName recursorLevels.toList + let recursor ← addRecParams mvarId majorTypeArgs recursorInfo.paramsPos recursor + -- Compute motive + let motive := target + let motive ← if recursorInfo.depElim then + pure <| mkLambda `x BinderInfo.default (← inferType major) (← motive.abstractM #[major]) + else + pure motive + let motive ← mkLambdaFVars indices motive + return mkApp recursor motive + | _ => + throwTacticEx tacticName mvarId "major premise is not of the form (C ...)" + +def _root_.Lean.MVarId.induction (mvarId : MVarId) (majorFVarId : FVarId) (recursorName : Name) (givenNames : Array AltVarNames := #[]) : MetaM (Array InductionSubgoal) := + mvarId.withContext do + trace[Meta.Tactic.induction] "initial\n{MessageData.ofGoal mvarId}" + mvarId.checkNotAssigned `induction + let majorLocalDecl ← majorFVarId.getDecl + let recursorInfo ← mkRecursorInfo recursorName + let some majorType ← whnfUntil majorLocalDecl.type recursorInfo.typeName + | throwUnexpectedMajorType `induction mvarId majorLocalDecl.type + majorType.withApp fun _ majorTypeArgs => do + recursorInfo.paramsPos.forM fun paramPos? => do + match paramPos? with + | none => pure () + | some paramPos => if paramPos ≥ majorTypeArgs.size then throwTacticEx `induction mvarId m!"major premise type is ill-formed{indentExpr majorType}" + let indices ← getMajorTypeIndices mvarId `induction recursorInfo majorType + let target ← mvarId.getType + if (← pure !recursorInfo.depElim <&&> exprDependsOn target majorFVarId) then + throwTacticEx `induction mvarId m!"recursor '{recursorName}' does not support dependent elimination, but conclusion depends on major premise" + -- Revert indices and major premise preserving variable order + let (reverted, mvarId) ← mvarId.revert ((indices.map Expr.fvarId!).push majorFVarId) true + -- Re-introduce indices and major + let (indices', mvarId) ← mvarId.introNP indices.size + let (majorFVarId', mvarId) ← mvarId.intro1P + -- Create FVarSubst with indices + let baseSubst := Id.run do + let mut subst : FVarSubst := {} + let mut i := 0 + for index in indices do + subst := subst.insert index.fvarId! (mkFVar indices'[i]!) + i := i + 1 + pure subst + trace[Meta.Tactic.induction] "after revert&intro\n{MessageData.ofGoal mvarId}" + -- Update indices and major + let indices := indices'.map mkFVar + let majorFVarId := majorFVarId' + let major := mkFVar majorFVarId + mvarId.withContext do + let recursor ← mkRecursorAppPrefix mvarId `induction majorFVarId recursorInfo indices + finalize mvarId givenNames recursorInfo reverted major indices baseSubst recursor + +builtin_initialize registerTraceClass `Meta.Tactic.induction + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Injection.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Injection.lean new file mode 100644 index 0000000000000000000000000000000000000000..0b61cd29f14684fa2345f9b6cd2079d0f76fb7a6 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Injection.lean @@ -0,0 +1,142 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.AppBuilder +import Lean.Meta.MatchUtil +import Lean.Meta.Tactic.Clear +import Lean.Meta.Tactic.Subst +import Lean.Meta.Tactic.Assert +import Lean.Meta.Tactic.Intro + +namespace Lean.Meta + +def getCtorNumPropFields (ctorInfo : ConstructorVal) : MetaM Nat := do + forallTelescopeReducing ctorInfo.type fun xs _ => do + let mut numProps := 0 + for i in [:ctorInfo.numFields] do + if (← isProp (← inferType xs[ctorInfo.numParams + i]!)) then + numProps := numProps + 1 + return numProps + +inductive InjectionResultCore where + | solved + | subgoal (mvarId : MVarId) (numNewEqs : Nat) + +def injectionCore (mvarId : MVarId) (fvarId : FVarId) : MetaM InjectionResultCore := + mvarId.withContext do + mvarId.checkNotAssigned `injection + let decl ← fvarId.getDecl + let type ← whnf decl.type + let go (type prf : Expr) : MetaM InjectionResultCore := do + match type.eq? with + | none => throwTacticEx `injection mvarId "equality expected" + | some (_, a, b) => + let target ← mvarId.getType + match (← isConstructorApp'? a), (← isConstructorApp'? b) with + | some aCtor, some bCtor => + -- We use the default transparency because `a` and `b` may be builtin literals. + let val ← withTransparency .default <| mkNoConfusion target prf + if aCtor.name != bCtor.name then + mvarId.assign val + return InjectionResultCore.solved + else + let valType ← inferType val + -- We use the default transparency setting here because `a` and `b` may be builtin literals + -- that need to expanded into constructors. + let valType ← whnfD valType + match valType with + | Expr.forallE _ newTarget _ _ => + let newTarget := newTarget.headBeta + let tag ← mvarId.getTag + let newMVar ← mkFreshExprSyntheticOpaqueMVar newTarget tag + mvarId.assign (mkApp val newMVar) + let mvarId ← newMVar.mvarId!.tryClear fvarId + /- Recall that `noConfusion` does not include equalities for + propositions since they are trivial due to proof irrelevance. -/ + let numPropFields ← getCtorNumPropFields aCtor + return InjectionResultCore.subgoal mvarId (aCtor.numFields - numPropFields) + | _ => throwTacticEx `injection mvarId "ill-formed noConfusion auxiliary construction" + | _, _ => throwTacticEx `injection mvarId "equality of constructor applications expected" + let prf := mkFVar fvarId + if let some (α, a, β, b) := type.heq? then + if (← isDefEq α β) then + go (← mkEq a b) (← mkEqOfHEq prf) + else + go type prf + else + go type prf + +inductive InjectionResult where + | solved + | subgoal (mvarId : MVarId) (newEqs : Array FVarId) (remainingNames : List Name) + + +def injectionIntro (mvarId : MVarId) (numEqs : Nat) (newNames : List Name) (tryToClear := true) : MetaM InjectionResult := + let rec go : Nat → MVarId → Array FVarId → List Name → MetaM InjectionResult + | 0, mvarId, fvarIds, remainingNames => + return InjectionResult.subgoal mvarId fvarIds remainingNames + | n+1, mvarId, fvarIds, name::remainingNames => do + let (fvarId, mvarId) ← mvarId.intro name + let (fvarId, mvarId) ← heqToEq mvarId fvarId tryToClear + go n mvarId (fvarIds.push fvarId) remainingNames + | n+1, mvarId, fvarIds, [] => do + let (fvarId, mvarId) ← mvarId.intro1 + let (fvarId, mvarId) ← heqToEq mvarId fvarId tryToClear + go n mvarId (fvarIds.push fvarId) [] + go numEqs mvarId #[] newNames + +def injection (mvarId : MVarId) (fvarId : FVarId) (newNames : List Name := []) : MetaM InjectionResult := do + match (← injectionCore mvarId fvarId) with + | .solved => pure .solved + | .subgoal mvarId numEqs => injectionIntro mvarId numEqs newNames + +inductive InjectionsResult where + /-- `injections` closed the input goal. -/ + | solved + /-- + `injections` produces a new goal `mvarId`. `remainingNames` contains the user-facing names that have not been used. + `forbidden` contains all local declarations to which `injection` has been applied. + Recall that some of these declarations may not have been eliminated from the local context due to forward dependencies, and + we use `forbidden` to avoid non-termination when using `injections` in a loop. + -/ + | subgoal (mvarId : MVarId) (remainingNames : List Name) (forbidden : FVarIdSet) + +/-- +Applies `injection` to local declarations in `mvarId`. It uses `newNames` to name the new local declarations. +`maxDepth` is the maximum recursion depth. Only local declarations that are not in `forbidden` are considered. +Recall that some of local declarations may not have been eliminated from the local context due to forward dependencies, and +we use `forbidden` to avoid non-termination when using `injections` in a loop. +-/ +partial def injections (mvarId : MVarId) (newNames : List Name := []) (maxDepth : Nat := 5) (forbidden : FVarIdSet := {}) : MetaM InjectionsResult := + mvarId.withContext do + let fvarIds := (← getLCtx).getFVarIds + go maxDepth fvarIds.toList mvarId newNames forbidden +where + go (depth : Nat) (fvarIds : List FVarId) (mvarId : MVarId) (newNames : List Name) (forbidden : FVarIdSet) : MetaM InjectionsResult := do + match depth, fvarIds with + | 0, _ => throwTacticEx `injections mvarId "recursion depth exceeded" + | _, [] => return .subgoal mvarId newNames forbidden + | d+1, fvarId :: fvarIds => do + let cont := do + go (d+1) fvarIds mvarId newNames forbidden + if forbidden.contains fvarId then + cont + else if let some (_, lhs, rhs) ← matchEqHEq? (← fvarId.getType) then + let lhs ← whnf lhs + let rhs ← whnf rhs + if lhs.isRawNatLit && rhs.isRawNatLit then + cont + else + try + commitIfNoEx do + match (← injection mvarId fvarId newNames) with + | .solved => return .solved + | .subgoal mvarId newEqs remainingNames => + mvarId.withContext <| go d (newEqs.toList ++ fvarIds) mvarId remainingNames (forbidden.insert fvarId) + catch _ => cont + else cont + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Intro.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Intro.lean new file mode 100644 index 0000000000000000000000000000000000000000..0b4002b05f93cac23dd6c672401259c59965ef5e --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Intro.lean @@ -0,0 +1,183 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Util + +namespace Lean.Meta + +@[inline] private partial def introNImp {σ} (mvarId : MVarId) (n : Nat) (mkName : LocalContext → Name → Bool → σ → MetaM (Name × σ)) (s : σ) + : MetaM (Array FVarId × MVarId) := mvarId.withContext do + mvarId.checkNotAssigned `introN + let mvarType ← mvarId.getType + let lctx ← getLCtx + let rec @[specialize] loop (i : Nat) (lctx : LocalContext) (fvars : Array Expr) (j : Nat) (s : σ) (type : Expr) : MetaM (Array Expr × MVarId) := do + match i, type with + | 0, type => + let type := type.instantiateRevRange j fvars.size fvars + withLCtx' lctx do + withNewLocalInstances fvars j do + let tag ← mvarId.getTag + let type := type.headBeta + let newMVar ← mkFreshExprSyntheticOpaqueMVar type tag + let newVal ← mkLambdaFVars fvars newMVar + mvarId.assign newVal + return (fvars, newMVar.mvarId!) + | i+1, .letE n type val body _ => + let type := type.instantiateRevRange j fvars.size fvars + let type := type.headBeta + let val := val.instantiateRevRange j fvars.size fvars + let fvarId ← mkFreshFVarId + let (n, s) ← mkName lctx n true s + /- + We have both dependent and non-dependent `let` expressions result in dependent `ldecl`s. + It is counterintuitive if `have` expressions are introduced with opaque values, + especially when we run transformations to aggressively turn `let`s into `have`s. + -/ + let lctx := lctx.mkLetDecl fvarId n type val + let fvar := mkFVar fvarId + let fvars := fvars.push fvar + loop i lctx fvars j s body + | i+1, .forallE n type body c => + let type := type.instantiateRevRange j fvars.size fvars + let type := type.headBeta + let fvarId ← mkFreshFVarId + let (n, s) ← mkName lctx n c.isExplicit s + let lctx := lctx.mkLocalDecl fvarId n type c + let fvar := mkFVar fvarId + let fvars := fvars.push fvar + loop i lctx fvars j s body + | i+1, type => + let type := type.instantiateRevRange j fvars.size fvars + withLCtx' lctx do + withNewLocalInstances fvars j do + /- We used to use just `whnf`, but it produces counterintuitive behavior if + - `type` is a metavariable `?m` such that `?m := let x := v; b`, or + - `type` has `MData` or annotations such as `optParam` around a `let`-expression. + + `whnf` instantiates metavariables, and consumes `MData`, but it also expands the `let`. + -/ + let newType := (← instantiateMVars type).cleanupAnnotations + if newType.isForall || newType.isLet then + loop (i+1) lctx fvars fvars.size s newType + else + let newType ← whnf newType + if newType.isForall then + loop (i+1) lctx fvars fvars.size s newType + else + throwTacticEx `introN mvarId "insufficient number of binders" + let (fvars, mvarId) ← loop n lctx #[] 0 s mvarType + return (fvars.map Expr.fvarId!, mvarId) + +register_builtin_option tactic.hygienic : Bool := { + defValue := true + group := "tactic" + descr := "make sure tactics are hygienic" +} + +private def mkFreshBinderNameForTacticCore (lctx : LocalContext) (binderName : Name) (hygienic := true) : MetaM Name := do + if hygienic then + mkFreshUserName binderName + else + return lctx.getUnusedName binderName + +/-- +Similar to `Lean.Core.mkFreshUserName`, but takes into account the `tactic.hygienic` option value. +If `tactic.hygienic = true`, then fresh macro scopes are applied to `binderName`. +If not, then returns an (accessible) name based on `binderName` that is unused in the local context. +-/ +def mkFreshBinderNameForTactic (binderName : Name) : MetaM Name := do + mkFreshBinderNameForTacticCore (← getLCtx) binderName (tactic.hygienic.get (← getOptions)) + +private def mkAuxNameImp (preserveBinderNames : Bool) (hygienic : Bool) (useNamesForExplicitOnly : Bool) + (lctx : LocalContext) (binderName : Name) (isExplicit : Bool) : List Name → MetaM (Name × List Name) + | [] => mkAuxNameWithoutGivenName [] + | n :: rest => do + if useNamesForExplicitOnly && !isExplicit then + mkAuxNameWithoutGivenName (n :: rest) + else if n != Name.mkSimple "_" then + return (n, rest) + else + mkAuxNameWithoutGivenName rest +where + mkAuxNameWithoutGivenName (rest : List Name) : MetaM (Name × List Name) := do + -- Use a nicer binder name than `[anonymous]`. + -- In this case, we make sure the name is hygienic. + let binderName ← if binderName.isAnonymous then mkFreshUserName `a else pure binderName + if preserveBinderNames then + return (binderName, rest) + else + let binderName ← mkFreshBinderNameForTacticCore lctx binderName hygienic + return (binderName, rest) + +def introNCore (mvarId : MVarId) (n : Nat) (givenNames : List Name) (useNamesForExplicitOnly : Bool) (preserveBinderNames : Bool) + : MetaM (Array FVarId × MVarId) := do + let hygienic := tactic.hygienic.get (← getOptions) + if n == 0 then + return (#[], mvarId) + else + introNImp mvarId n (mkAuxNameImp preserveBinderNames hygienic useNamesForExplicitOnly) givenNames + +/-- +Introduce `n` binders in the goal `mvarId`. +-/ +abbrev _root_.Lean.MVarId.introN (mvarId : MVarId) (n : Nat) (givenNames : List Name := []) (useNamesForExplicitOnly := false) : MetaM (Array FVarId × MVarId) := + introNCore mvarId n givenNames (useNamesForExplicitOnly := useNamesForExplicitOnly) (preserveBinderNames := false) + +/-- +Introduce `n` binders in the goal `mvarId`. The new hypotheses are named using the binder names. +The suffix `P` stands for "preserving`. +-/ +abbrev _root_.Lean.MVarId.introNP (mvarId : MVarId) (n : Nat) : MetaM (Array FVarId × MVarId) := + introNCore mvarId n [] (useNamesForExplicitOnly := false) (preserveBinderNames := true) + +/-- +Introduce one binder using `name` as the new hypothesis name. +-/ +def _root_.Lean.MVarId.intro (mvarId : MVarId) (name : Name) : MetaM (FVarId × MVarId) := do + let (fvarIds, mvarId) ← mvarId.introN 1 [name] + return (fvarIds[0]!, mvarId) + +def intro1Core (mvarId : MVarId) (preserveBinderNames : Bool) : MetaM (FVarId × MVarId) := do + let (fvarIds, mvarId) ← introNCore mvarId 1 [] (useNamesForExplicitOnly := false) preserveBinderNames + return (fvarIds[0]!, mvarId) + +/-- Introduce one object from the goal `mvarid`, without preserving the name used in the binder. +Returns a pair made of the newly introduced variable (which will have an inaccessible name) +and the new goal. This will fail if there is nothing to introduce, ie when the goal +does not start with a forall, lambda or let. -/ +abbrev _root_.Lean.MVarId.intro1 (mvarId : MVarId) : MetaM (FVarId × MVarId) := + intro1Core mvarId false + +/-- Introduce one object from the goal `mvarid`, preserving the name used in the binder. +Returns a pair made of the newly introduced variable and the new goal. +This will fail if there is nothing to introduce, ie when the goal +does not start with a forall, lambda or let. -/ +abbrev _root_.Lean.MVarId.intro1P (mvarId : MVarId) : MetaM (FVarId × MVarId) := + intro1Core mvarId true + +/-- +Calculate the number of new hypotheses that would be created by `intros`, +i.e. the number of binders which can be introduced without unfolding definitions. +-/ +partial def getIntrosSize : Expr → Nat + | .forallE _ _ b _ => getIntrosSize b + 1 + | .letE _ _ _ b _ => getIntrosSize b + 1 + | .mdata _ b => getIntrosSize b + | _ => 0 + +/-- +Introduce as many binders as possible without unfolding definitions. +-/ +def _root_.Lean.MVarId.intros (mvarId : MVarId) : MetaM (Array FVarId × MVarId) := do + let type ← mvarId.getType + let type ← instantiateMVars type + let n := getIntrosSize type + if n == 0 then + return (#[], mvarId) + else + mvarId.introN n + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Lets.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Lets.lean new file mode 100644 index 0000000000000000000000000000000000000000..e85228b2ec5117396d5442169ad9576f831a011b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Lets.lean @@ -0,0 +1,447 @@ +/- +Copyright (c) 2024 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kyle Miller +-/ +prelude +import Lean.Meta.Tactic.Replace +import Lean.Meta.LetToHave + +/-! +# Tactics to manipulate `let` expressions +-/ + +open Lean Meta + +namespace Lean.Meta + +/-! +### `let` extraction + +Extracting `let`s means to locate `let`/`have`s in a term and to extract them +from the term, extending the local context with new declarations in the process. +A related process is lifting `lets`, which means to move `let`/`have`s toward the root of a term. +-/ + +namespace ExtractLets + +structure LocalDecl' where + /-- An `ldecl` with `nondep := false`. -/ + decl : LocalDecl + /-- + If true, is a `let`, if false, is a `have`. + Used in `lift` mode. + -/ + isLet : Bool + +structure State where + /-- Names to use for local definitions for the extracted lets. -/ + givenNames : List Name + /-- Saved declarations for the extracted `let`s. -/ + decls : Array LocalDecl' := #[] + /-- Map from `let` values to fvars. To support the `merge` option. -/ + valueMap : ExprStructMap FVarId := {} + deriving Inhabited + +-- The `Bool` in the cache key is whether we are looking at a "top-level" expression. +abbrev M := ReaderT ExtractLetsConfig <| MonadCacheT (Bool × ExprStructEq) Expr <| StateRefT State MetaM + +/-- Returns `true` if `nextName?` would return a name. -/ +def hasNextName : M Bool := do + return !(← read).onlyGivenNames || !(← get).givenNames.isEmpty + +/-- Gets the next name to use for extracted `let`s -/ +def nextName? : M (Option Name) := do + let s ← get + match s.givenNames, (← read).onlyGivenNames with + | n :: ns, _ => set { s with givenNames := ns }; return n + | [] , true => return none + | [] , false => return `_ + +/-- +Generate a name to use for a new local declaration, derived possibly from the given binder name. +Returns `none` iff `hasNextName` is false. +-/ +def nextNameForBinderName? (binderName : Name) : M (Option Name) := do + if let some n ← nextName? then + if n != `_ then + return n + else + if binderName.isAnonymous then + -- Use a nicer binder name than `[anonymous]`. + mkFreshUserName `a + else if (← read).preserveBinderNames || n.hasMacroScopes then + return n + else + mkFreshUserName binderName + else + return none + +/-- +Returns 'true' if `e` does not depend on any of the fvars in `fvars`. +-/ +def extractable (fvars : List Expr) (e : Expr) : Bool := + !e.hasAnyFVar (fvars.contains <| .fvar ·) + +/-- +Returns whether a let-like expression with the given type and value is extractable, +given the list `fvars` of binders that inhibit extraction. +-/ +def isExtractableLet (fvars : List Expr) (n : Name) (t v : Expr) : M (Bool × Name) := do + if (← hasNextName) && extractable fvars t && extractable fvars v then + if let some n ← nextNameForBinderName? n then + return (true, n) + -- In lift mode, we temporarily extract non-extractable lets, but we do not make use of `givenNames` for them. + -- These will be flushed as let/have expressions, and we wish to preserve the original binder name. + if (← read).lift then + return (true, n) + return (false, n) + +/-- +Adds the `decl` to the `decls` list. Assumes that `decl` is an ldecl with `nondep := false`. +-/ +def addDecl (decl : LocalDecl) (isLet : Bool) : M Unit := do + let cfg ← read + modify fun s => { s with + decls := s.decls.push { decl, isLet } + valueMap := if cfg.merge then s.valueMap.insert decl.value decl.fvarId else s.valueMap + } + +/-- +Removes and returns all local declarations that (transitively) depend on `fvar`. +-/ +def flushDecls (fvar : FVarId) : M (Array LocalDecl') := do + let mut fvarSet : FVarIdSet := {} + fvarSet := fvarSet.insert fvar + let mut toSave := #[] + let mut toFlush := #[] + for ldecl in (← get).decls do + if ldecl.decl.type.hasAnyFVar (fvarSet.contains ·) || ldecl.decl.value.hasAnyFVar (fvarSet.contains ·) then + toFlush := toFlush.push ldecl + fvarSet := fvarSet.insert ldecl.decl.fvarId + else + toSave := toSave.push ldecl + modify fun s => { s with decls := toSave } + return toFlush + +/-- +Ensures that the given local declarations are in context. Runs `k` in that context. +-/ +def withEnsuringDeclsInContext [Monad m] [MonadControlT MetaM m] [MonadLCtx m] (decls : Array LocalDecl') (k : m α) : m α := do + let lctx ← getLCtx + let decls := decls |>.filter (!lctx.contains ·.decl.fvarId) |>.map (·.decl) + withExistingLocalDecls decls.toList k + +/-- +Closes all the local declarations in `e`, creating `let` and `have` expressions. +Does not require that any of the declarations are in context. +Assumes that `e` contains no metavariables with local contexts that contain any of these metavariables +(the extraction procedure creates no new metavariables, so this is the case). + +This should *not* be used when closing lets for new goal metavariables, since +1. The goal contains the decls in its local context, violating the assumption. +2. We need to use true `let`s in that case, since tactics may zeta-delta reduce these declarations. +-/ +def mkLetDecls (decls : Array LocalDecl') (e : Expr) : Expr := + decls.foldr (init := e) fun { decl, isLet } e => + Expr.letE decl.userName decl.type decl.value (e.abstract #[decl.toExpr]) (nondep := !isLet) + +/-- +Makes sure the declaration for `fvarId` is marked with `isLet := true`. +Used in `lift + merge` mode to ensure that, after merging, if any version was a `let` then it's a `let` rather than a `have`. +-/ +def ensureIsLet (fvarId : FVarId) : M Unit := do + modify fun s => { s with + decls := s.decls.map fun d => + if d.decl.fvarId == fvarId then { d with isLet := true } else d + } + +/-- +Ensures that the given `fvarId` is in context by adding `decls` from the state. +Simplification: since we are not recording which decls depend on which, but we do know all dependencies +come before a particular decl, we add all the decls up to and including `fvarId`. + +Used for `merge` feature. +-/ +def withDeclInContext (fvarId : FVarId) (k : M α) : M α := do + let decls := (← get).decls + if (← getLCtx).contains fvarId then + -- Is either pre-existing or already added. + k + else if let some idx := decls.findIdx? (·.decl.fvarId == fvarId) then + withEnsuringDeclsInContext decls[*...(idx+1)] k + else + k + +/-- +Initializes the `valueMap` with all the local definitions that aren't implementation details. +Used for `merge` feature when `useContext` is enabled. +-/ +def initializeValueMap : M Unit := do + let lctx ← getLCtx + lctx.forM fun decl => do + if decl.isLet && !decl.isImplementationDetail then + let value ← instantiateMVars decl.value + modify fun s => { s with valueMap := s.valueMap.insert value decl.fvarId } + +/-- +Returns `true` if the expression contains a `let` expression or a `have`. +Its purpose is to be a check for whether a subexpression can be skipped. +-/ +def containsLet (e : Expr) : Bool := + Option.isSome <| e.find? (·.isLet) + +/-- +Extracts lets from `e`. +- `fvars` is an array of all the local variables from going under binders, + used to detect whether an expression is extractable. Extracted `let`s do not have their fvarids in this list. + This is not part of the cache key since it's an optimization and in principle derivable. +- `topLevel` is whether we are still looking at the top-level expression. + The body of an extracted top-level let is also considered to be top-level. + This is part of the cache key since it affects what is extracted. + +Note: the return value may refer to fvars that are not in the current local context, but they are in the `decls` list. +-/ +partial def extractCore (fvars : List Expr) (e : Expr) (topLevel : Bool := false) : M Expr := do + let cfg ← read + if e.isAtomic then + return e + else if !cfg.descend && !topLevel then + return e + else + checkCache (topLevel, (e : ExprStructEq)) fun _ => do + if !containsLet e then + return e + -- Don't honor `proofs := false` or `types := false` for top-level lets, since it's confusing not having them be extracted. + unless topLevel && (e.isLet || e.isMData) do + if !cfg.proofs then + if ← isProof e then + return e + if !cfg.types then + if ← isType e then + return e + let whenDescend (k : M Expr) : M Expr := do if cfg.descend then k else pure e + match e with + | .bvar .. | .fvar .. | .mvar .. | .sort .. | .const .. | .lit .. => unreachable! + | .mdata _ e' => return e.updateMData! (← extractCore fvars e' (topLevel := topLevel)) + | .letE n t v b nondep => extractLetLike (!nondep) n t v b (topLevel := topLevel) + | .app .. => whenDescend do extractApp e.getAppFn e.getAppArgs + | .proj _ _ s => whenDescend do return e.updateProj! (← extractCore fvars s) + | .lam n t b i => whenDescend do extractBinder n t b i (fun t b => e.updateLambda! i t b) + | .forallE n t b i => whenDescend do extractBinder n t b i (fun t b => e.updateForall! i t b) +where + extractBinder (n : Name) (t b : Expr) (i : BinderInfo) (mk : Expr → Expr → Expr) : M Expr := do + let t ← extractCore fvars t + if (← read).underBinder then + withLocalDecl n i t fun x => do + let b ← extractCore (x :: fvars) (b.instantiate1 x) + if (← read).lift then + let toFlush ← flushDecls x.fvarId! + let b := mkLetDecls toFlush b + return mk t (b.abstract #[x]) + else + return mk t (b.abstract #[x]) + else + return mk t b + extractLetLike (isLet : Bool) (n : Name) (t v b : Expr) (topLevel : Bool) : M Expr := do + let cfg ← read + let t ← extractCore fvars t + let v ← extractCore fvars v + if cfg.usedOnly && !b.hasLooseBVars then + return ← extractCore fvars b (topLevel := topLevel) + if cfg.merge then + if let some fvarId := (← get).valueMap.get? v then + if isLet && cfg.lift then ensureIsLet fvarId + return ← withDeclInContext fvarId <| + extractCore fvars (b.instantiate1 (.fvar fvarId)) (topLevel := topLevel) + let (extract, n) ← isExtractableLet fvars n t v + if !extract && (!cfg.underBinder || !cfg.descend) then + return e.updateLetE! t v b + withLetDecl n t v fun x => do + if extract then + addDecl (← x.fvarId!.getDecl) isLet + extractCore fvars (b.instantiate1 x) (topLevel := topLevel) + else + let b ← extractCore (x :: fvars) (b.instantiate1 x) + return e.updateLetE! t v (b.abstract #[x]) + extractApp (f : Expr) (args : Array Expr) : M Expr := do + let cfg ← read + let f' ← extractCore fvars f + if cfg.implicits then + return mkAppN f' (← args.mapM (extractCore fvars)) + else + let (paramInfos, _) ← instantiateForallWithParamInfos (← inferType f) args + let mut args := args + for i in [0:args.size] do + if paramInfos[i]!.binderInfo.isExplicit then + args := args.set! i (← extractCore fvars args[i]!) + return mkAppN f' args + +def extractTopLevel (e : Expr) : M Expr := do + let e ← instantiateMVars e + extractCore [] e (topLevel := true) + +/-- +Main entry point for extracting lets. +-/ +def extract (es : Array Expr) : M (Array Expr) := do + let cfg ← read + if cfg.merge && cfg.useContext then + initializeValueMap + es.mapM extractTopLevel + +end ExtractLets + +/-- +Implementation of the `extractLets` function. +- `es` is an array of terms that are valid in the current local context. +- `k` is a callback that is run in the updated local context with all relevant `let`s extracted + and with the post-extraction expressions, and the remaining names from `givenNames`. +-/ +private def extractLetsImp (es : Array Expr) (givenNames : List Name) + (k : Array FVarId → Array Expr → List Name → MetaM α) (config : ExtractLetsConfig) : MetaM α := do + let (es, st) ← ExtractLets.extract es |>.run config |>.run' {} |>.run { givenNames } + let givenNames' := st.givenNames + let decls := st.decls.map (·.decl) + withExistingLocalDecls decls.toList <| k (decls.map (·.fvarId)) es givenNames' + +/-- +Extracts `let` and `have` expressions into local definitions, +evaluating `k` at the post-extracted expressions and the extracted fvarids, within a context containing those local declarations. +- The `givenNames` is a list of explicit names to use for extracted local declarations. + If a name is `_` (or if there is no provided given name and `config.onlyGivenNames` is true) then uses a hygienic name + based on the existing binder name. +-/ +def extractLets [Monad m] [MonadControlT MetaM m] (es : Array Expr) (givenNames : List Name) + (k : Array FVarId → Array Expr → List Name → m α) (config : ExtractLetsConfig := {}) : m α := + map3MetaM (fun k => extractLetsImp es givenNames k config) k + +/-- +Lifts `let` and `have` expressions in the given expression as far out as possible. +-/ +def liftLets (e : Expr) (config : LiftLetsConfig := {}) : MetaM Expr := do + let (es, st) ← ExtractLets.extract #[e] |>.run { config with onlyGivenNames := true } |>.run' {} |>.run { givenNames := [] } + return ExtractLets.mkLetDecls st.decls es[0]! + +end Lean.Meta + +private def throwMadeNoProgress (tactic : Name) (mvarId : MVarId) : MetaM α := + throwTacticEx tactic mvarId m!"made no progress" + +/-- +Extracts `let` and `have` expressions from the target, +returning `FVarId`s for the extracted let declarations along with the new goal. +- The `givenNames` is a list of explicit names to use for extracted local declarations. + If a name is `_` (or if there is no provided given name and `config.onlyGivenNames` is true) then uses a hygienic name + based on the existing binder name. +-/ +def Lean.MVarId.extractLets (mvarId : MVarId) (givenNames : List Name) (config : ExtractLetsConfig := {}) : + MetaM ((Array FVarId × List Name) × MVarId) := + mvarId.withContext do + mvarId.checkNotAssigned `extract_lets + let ty ← mvarId.getType + Meta.extractLets #[ty] givenNames (config := config) fun fvarIds es givenNames' => do + let ty' := es[0]! + if fvarIds.isEmpty && ty == ty' then + throwMadeNoProgress `extract_lets mvarId + let g ← mkFreshExprSyntheticOpaqueMVar ty' (← mvarId.getTag) + mvarId.assign <| ← mkLetFVars (usedLetOnly := false) (fvarIds.map .fvar) g + return ((fvarIds, givenNames'), g.mvarId!) + +/-- +Like `Lean.MVarId.extractLets` but extracts lets from a local declaration. +If the local declaration has a value, then both its type and value are modified. +-/ +def Lean.MVarId.extractLetsLocalDecl (mvarId : MVarId) (fvarId : FVarId) (givenNames : List Name) (config : ExtractLetsConfig := {}) : + MetaM ((Array FVarId × List Name) × MVarId) := do + mvarId.checkNotAssigned `extract_lets + mvarId.withReverted #[fvarId] fun mvarId fvars => mvarId.withContext do + let finalize (fvarIds : Array FVarId) (givenNames' : List Name) (targetNew : Expr) := do + let g ← mkFreshExprSyntheticOpaqueMVar targetNew (← mvarId.getTag) + mvarId.assign <| ← mkLetFVars (usedLetOnly := false) (fvarIds.map .fvar) g + return ((fvarIds, givenNames'), fvars.map .some, g.mvarId!) + match ← mvarId.getType with + | .forallE n t b i => + Meta.extractLets #[t] givenNames (config := config) fun fvarIds es givenNames' => do + let t' := es[0]! + if fvarIds.isEmpty && t == t' then + throwMadeNoProgress `extract_lets mvarId + finalize fvarIds givenNames' (.forallE n t' b i) + | .letE n t v b ndep => + Meta.extractLets #[t, v] givenNames (config := config) fun fvarIds es givenNames' => do + let t' := es[0]! + let v' := es[1]! + if fvarIds.isEmpty && t == t' && v == v' then + throwMadeNoProgress `extract_lets mvarId + finalize fvarIds givenNames' (.letE n t' v' b ndep) + | _ => throwTacticEx `extract_lets mvarId "unexpected auxiliary target" + +/-- +Lifts `let` and `have` expressions in target as far out as possible. +Throws an exception if nothing is lifted. + +Like `Lean.MVarId.extractLets`, but top-level lets are not added to the local context. +-/ +def Lean.MVarId.liftLets (mvarId : MVarId) (config : LiftLetsConfig := {}) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `lift_lets + let ty ← mvarId.getType + let ty' ← Meta.liftLets ty (config := config) + if ty == ty' then + throwMadeNoProgress `lift_lets mvarId + mvarId.replaceTargetDefEq ty' + +/-- +Like `Lean.MVarId.liftLets` but lifts lets in a local declaration. +If the local declaration has a value, then both its type and value are modified. +-/ +def Lean.MVarId.liftLetsLocalDecl (mvarId : MVarId) (fvarId : FVarId) (config : LiftLetsConfig := {}) : MetaM MVarId := do + mvarId.checkNotAssigned `lift_lets + -- Revert to make sure the resulting type/value refers fvars that come after `fvarId`. + -- (Note: reverting isn't necessary unless both `merge := true` and `useContext := true`.) + Prod.snd <$> mvarId.withReverted #[fvarId] fun mvarId fvars => mvarId.withContext do + let finalize (targetNew : Expr) := do + return ((), fvars.map .some, ← mvarId.replaceTargetDefEq targetNew) + match ← mvarId.getType with + | .forallE n t b i => + let t' ← Meta.liftLets t (config := config) + if t == t' then + throwMadeNoProgress `lift_lets mvarId + finalize (.forallE n t' b i) + | .letE n t v b ndep => + let t' ← Meta.liftLets t (config := config) + let v' ← Meta.liftLets v (config := config) + if t == t' && v == v' then + throwMadeNoProgress `lift_lets mvarId + finalize (.letE n t' v' b ndep) + | _ => throwTacticEx `lift_lets mvarId "unexpected auxiliary target" + +/-! +### Let-to-have transformation + +A meta tactic version of `Lean.Meta.letToHave`. +-/ + +/-- +Transforms lets to haves in the target. Throws an error if no progress is made. +-/ +def Lean.MVarId.letToHave (mvarId : MVarId) (failIfUnchanged := true) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `let_to_have + let ty ← mvarId.getType + let ty' ← Meta.letToHave ty + if failIfUnchanged && ty == ty' then + throwMadeNoProgress `let_to_have mvarId + mvarId.replaceTargetDefEq ty' + +/-- +Transforms lets to haves in the type of `fvarId`. Throws an error if no progress is made. +-/ +def Lean.MVarId.letToHaveLocalDecl (mvarId : MVarId) (fvarId : FVarId) (failIfUnchanged := true) : MetaM MVarId := do + mvarId.withContext do + mvarId.checkNotAssigned `let_to_have + let ty ← fvarId.getType + let ty' ← Meta.letToHave ty + if failIfUnchanged && ty == ty' then + throwMadeNoProgress `let_to_have mvarId + mvarId.replaceLocalDeclDefEq fvarId ty' diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/LibrarySearch.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/LibrarySearch.lean new file mode 100644 index 0000000000000000000000000000000000000000..796b3417e7ddb5171ff1ade37c0df360ac3603d8 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/LibrarySearch.lean @@ -0,0 +1,323 @@ +/- +Copyright (c) 2021-2023 Gabriel Ebner and Lean FRO. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Gabriel Ebner, Joe Hendrix, Kim Morrison +-/ +prelude +import Init.Data.Nat.MinMax +import Lean.Meta.LazyDiscrTree +import Lean.Meta.Tactic.SolveByElim +import Lean.Util.Heartbeats + +/-! +# Library search + +This file defines tactics `exact?` and `apply?`, +(formerly known as `library_search`) +and a term elaborator `exact?%` +that tries to find a lemma +solving the current goal +(subgoals are solved using `solveByElim`). + +``` +example : x < x + 1 := exact?% +example : Nat := by exact? +``` +-/ + + +namespace Lean.Meta.LibrarySearch + +builtin_initialize registerTraceClass `Tactic.librarySearch +builtin_initialize registerTraceClass `Tactic.librarySearch.lemmas + +open SolveByElim + +/-- +Wrapper for calling `Lean.Meta.SolveByElim.solveByElim with +appropriate arguments for library search. +-/ +def solveByElim (required : List Expr) (exfalso : Bool) (goals : List MVarId) (maxDepth : Nat) := do + let cfg : SolveByElimConfig := + { maxDepth, exfalso := exfalso, symm := true, commitIndependentGoals := true, + transparency := ← getTransparency, + -- `constructor` has been observed to significantly slow down `exact?` in Mathlib. + constructor := false } + let ⟨lemmas, ctx⟩ ← SolveByElim.mkAssumptionSet false false [] [] #[] + let cfg := if !required.isEmpty then cfg.requireUsingAll required else cfg + SolveByElim.solveByElim cfg lemmas ctx goals + +/-- +A "modifier" for a declaration. +* `none` indicates the original declaration, +* `mp` indicates that (possibly after binders) the declaration is an `↔`, + and we want to consider the forward direction, +* `mpr` similarly, but for the backward direction. +-/ +inductive DeclMod + | /-- the original declaration -/ none + | /-- the forward direction of an `iff` -/ mp + | /-- the backward direction of an `iff` -/ mpr +deriving DecidableEq, Inhabited, Ord, Hashable + +/-- +LibrarySearch has an extension mechanism for replacing the function used +to find candidate lemmas. +-/ +@[reducible] +def CandidateFinder := Expr → MetaM (Array (Name × DeclMod)) + +open LazyDiscrTree (InitEntry findMatches) + +private def addImport (name : Name) (constInfo : ConstantInfo) : + MetaM (Array (InitEntry (Name × DeclMod))) := + -- Don't report lemmas from metaprogramming namespaces. + if name.isMetaprogramming then return #[] else + forallTelescope constInfo.type fun _ type => do + let e ← InitEntry.fromExpr type (name, DeclMod.none) + let a := #[e] + if e.key == .const ``Iff 2 then + let a := a.push (← e.mkSubEntry 0 (name, DeclMod.mp)) + let a := a.push (← e.mkSubEntry 1 (name, DeclMod.mpr)) + pure a + else + pure a + +/-- Stores import discrimination tree. -/ +private def LibSearchState := IO.Ref (Option (LazyDiscrTree (Name × DeclMod))) + +private builtin_initialize defaultLibSearchState : IO.Ref (Option (LazyDiscrTree (Name × DeclMod))) ← do + IO.mkRef .none + +private instance : Inhabited LibSearchState where + default := defaultLibSearchState + +private builtin_initialize ext : EnvExtension LibSearchState ← + registerEnvExtension (IO.mkRef .none) + +/-- +We drop `.star` and `Eq * * *` from the discriminator trees because +they match too much. +-/ +def droppedKeys : List (List LazyDiscrTree.Key) := [[.star], [.const `Eq 3, .star, .star, .star]] + +/-- +The maximum number of constants an individual task may perform. + +The value was picked because it roughly corresponded to 50ms of work on the +machine this was developed on. Smaller numbers did not seem to improve +performance when importing Std and larger numbers (<10k) seemed to degrade +initialization performance. +-/ +private def constantsPerImportTask : Nat := 6500 + +/-- Create function for finding relevant declarations. -/ +def libSearchFindDecls : Expr → MetaM (Array (Name × DeclMod)) := + findMatches ext addImport + (droppedKeys := droppedKeys) + (constantsPerTask := constantsPerImportTask) + +/-- +Return an action that returns true when the remaining heartbeats is less +than the currently remaining heartbeats * leavePercent / 100. +-/ +def mkHeartbeatCheck (leavePercent : Nat) : MetaM (MetaM Bool) := do + let maxHB ← getMaxHeartbeats + let hbThreshold := (← getRemainingHeartbeats) * leavePercent / 100 + -- Return true if we should stop + pure $ + if maxHB = 0 then + pure false + else do + return (← getRemainingHeartbeats) < hbThreshold + +private def librarySearchEmoji : Except ε (Option α) → String +| .error _ => bombEmoji +| .ok (some _) => crossEmoji +| .ok none => checkEmoji + +/-- +Interleave x y interleaves the elements of x and y until one is empty and then returns +final elements in other list. +-/ +def interleaveWith {α β γ} (f : α → γ) (x : Array α) (g : β → γ) (y : Array β) : Array γ := + Id.run do + let mut res := Array.mkEmpty (x.size + y.size) + let n := min x.size y.size + for h : i in [0:n] do + have p : i < min x.size y.size := h.2.1 + have q : i < x.size := Nat.le_trans p (Nat.min_le_left ..) + have r : i < y.size := Nat.le_trans p (Nat.min_le_right ..) + res := res.push (f x[i]) + res := res.push (g y[i]) + let last := + if x.size > n then + (x.extract n x.size).map f + else + (y.extract n y.size).map g + pure $ res ++ last + +/-- +An exception ID that indicates further speculation on candidate lemmas should stop +and current results should be returned. +-/ +private builtin_initialize abortSpeculationId : InternalExceptionId ← + registerInternalExceptionId `Lean.Meta.LibrarySearch.abortSpeculation + +/-- +Called to abort speculative execution in library search. +-/ +def abortSpeculation [MonadExcept Exception m] : m α := + throw (Exception.internal abortSpeculationId {}) + +/-- Returns true if this is an abort speculation exception. -/ +def isAbortSpeculation : Exception → Bool +| .internal id _ => id == abortSpeculationId +| _ => false + +section LibrarySearch + +/-- +A library search candidate using symmetry includes the goal to solve, the metavar +context for that goal, and the name and orientation of a rule to try using with goal. +-/ +@[reducible] +def Candidate := (MVarId × MetavarContext) × (Name × DeclMod) + +/-- +Run `searchFn` on both the goal and `symm` applied to the goal. +-/ +def librarySearchSymm (searchFn : CandidateFinder) (goal : MVarId) : MetaM (Array Candidate) := do + let tgt ← goal.getType + let l1 ← searchFn tgt + let coreMCtx ← getMCtx + let coreGoalCtx := (goal, coreMCtx) + if let some symmGoal ← observing? goal.applySymm then + let newType ← instantiateMVars (← symmGoal.getType) + let l2 ← searchFn newType + let symmMCtx ← getMCtx + let symmGoalCtx := (symmGoal, symmMCtx) + setMCtx coreMCtx + pure $ interleaveWith (coreGoalCtx, ·) l1 (symmGoalCtx, ·) l2 + else + pure $ l1.map (coreGoalCtx, ·) + +private def emoji (e : Except ε α) := if e.toBool then checkEmoji else crossEmoji + +/-- Create lemma from name and mod. -/ +def mkLibrarySearchLemma (lem : Name) (mod : DeclMod) : MetaM Expr := do + let lem ← mkConstWithFreshMVarLevels lem + match mod with + | .none => pure lem + | .mp => mapForallTelescope (fun e => mkAppM ``Iff.mp #[e]) lem + | .mpr => mapForallTelescope (fun e => mkAppM ``Iff.mpr #[e]) lem + +private def isVar (e : Expr) : Bool := + match e with + | .bvar _ => true + | .fvar _ => true + | .mvar _ => true + | _ => false + +/-- +Tries to apply the given lemma (with symmetry modifier) to the goal, +then tries to close subsequent goals using `solveByElim`. +If `solveByElim` succeeds, `[]` is returned as the list of new subgoals, +otherwise the full list of subgoals is returned. +-/ +private def librarySearchLemma (cfg : ApplyConfig) (act : List MVarId → MetaM (List MVarId)) + (allowFailure : MVarId → MetaM Bool) (cand : Candidate) : MetaM (List MVarId) := do + let ((goal, mctx), (name, mod)) := cand + let ppMod (mod : DeclMod) : MessageData := + match mod with | .none => "" | .mp => " with mp" | .mpr => " with mpr" + withTraceNode `Tactic.librarySearch (return m!"{emoji ·} trying {name}{ppMod mod} ") do + setMCtx mctx + let lem ← mkLibrarySearchLemma name mod + let newGoals ← goal.apply lem cfg + try + act newGoals + catch _ => + if ← allowFailure goal then + pure newGoals + else + failure + +/-- +Sequentially invokes a tactic `act` on each value in candidates on the current state. + +The tactic `act` should return a list of meta-variables that still need to be resolved. +If this list is empty, then no variables remain to be solved, and `tryOnEach` returns +`none` with the environment set so each goal is resolved. + +If the action throws an internal exception with the `abortSpeculationId` id then +further computation is stopped and intermediate results returned. If any other +exception is thrown, then it is silently discarded. +-/ +def tryOnEach + (act : Candidate → MetaM (List MVarId)) + (candidates : Array Candidate) : + MetaM (Option (Array (List MVarId × MetavarContext))) := do + let mut a := #[] + let s ← saveState + for c in candidates do + match ← (tryCatch (Except.ok <$> act c) (pure ∘ Except.error)) with + | .error e => + restoreState s + if isAbortSpeculation e then + break + | .ok remaining => + if remaining.isEmpty then + return none + let ctx ← getMCtx + restoreState s + a := a.push (remaining, ctx) + return (.some a) + +private def librarySearch' (goal : MVarId) + (tactic : List MVarId → MetaM (List MVarId)) + (allowFailure : MVarId → MetaM Bool) + (leavePercentHeartbeats : Nat) : + MetaM (Option (Array (List MVarId × MetavarContext))) := do + withTraceNode `Tactic.librarySearch (return m!"{librarySearchEmoji ·} {← goal.getType}") do + profileitM Exception "librarySearch" (← getOptions) do + -- Create predicate that returns true when running low on heartbeats. + let candidates ← librarySearchSymm libSearchFindDecls goal + let cfg : ApplyConfig := { allowSynthFailures := true } + let shouldAbort ← mkHeartbeatCheck leavePercentHeartbeats + let act := fun cand => do + if ←shouldAbort then + abortSpeculation + librarySearchLemma cfg tactic allowFailure cand + tryOnEach act candidates + +/-- +Tries to solve the goal either by: +* calling `tactic true` +* or applying a library lemma then calling `tactic false` on the resulting goals. + +Typically here `tactic` is `solveByElim`, +with the `Bool` flag indicating whether it may retry with `exfalso`. + +If it successfully closes the goal, returns `none`. +Otherwise, it returns `some a`, where `a : Array (List MVarId × MetavarContext)`, +with an entry for each library lemma which was successfully applied, +containing a list of the subsidiary goals, and the metavariable context after the application. + +(Always succeeds, and the metavariable context stored in the monad is reverted, +unless the goal was completely solved.) + +(Note that if `solveByElim` solves some but not all subsidiary goals, +this is not currently tracked.) +-/ +def librarySearch (goal : MVarId) + (tactic : Bool → List MVarId → MetaM (List MVarId) := + fun initial g => solveByElim [] (maxDepth := 6) (exfalso := initial) g) + (allowFailure : MVarId → MetaM Bool := fun _ => pure true) + (leavePercentHeartbeats : Nat := 10) : + MetaM (Option (Array (List MVarId × MetavarContext))) := do + (tactic true [goal] *> pure none) <|> + librarySearch' goal (tactic false) allowFailure leavePercentHeartbeats + +end LibrarySearch + +end Lean.Meta.LibrarySearch diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/NormCast.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/NormCast.lean new file mode 100644 index 0000000000000000000000000000000000000000..d727d309c84187d6b9837dab3171942596b030ec --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/NormCast.lean @@ -0,0 +1,164 @@ +/- +Copyright (c) 2019 Paul-Nicolas Madelaine. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul-Nicolas Madelaine, Robert Y. Lewis, Mario Carneiro, Gabriel Ebner +-/ +prelude +import Lean.Meta.CongrTheorems +import Lean.Meta.Tactic.Simp.Attr +import Lean.Meta.CoeAttr + +namespace Lean.Meta.NormCast +/-- +`Label` is a type used to classify `norm_cast` lemmas. +* elim lemma: LHS has 0 head coes and ≥ 1 internal coe +* move lemma: LHS has 1 head coe and 0 internal coes, RHS has 0 head coes and ≥ 1 internal coes +* squash lemma: LHS has ≥ 1 head coes and 0 internal coes, RHS has fewer head coes +-/ +inductive Label + /-- elim lemma: LHS has 0 head coes and ≥ 1 internal coe -/ + | elim + /-- move lemma: LHS has 1 head coe and 0 internal coes, + RHS has 0 head coes and ≥ 1 internal coes -/ + | move + /-- squash lemma: LHS has ≥ 1 head coes and 0 internal coes, RHS has fewer head coes -/ + | squash + deriving DecidableEq, Repr, Inhabited + +/-- Assuming `e` is an application, returns the list of subterms that `simp` will rewrite in. -/ +def getSimpArgs (e : Expr) : MetaM (Array Expr) := do + match ← mkCongrSimp? e.getAppFn with + | none => return e.getAppArgs + | some {argKinds, ..} => + let mut args := #[] + for a in e.getAppArgs, k in argKinds do + if k matches .eq then + args := args.push a + return args + +/-- Counts how many coercions are at the head of the expression. -/ +partial def countHeadCoes (e : Expr) : MetaM Nat := do + if let Expr.const fn .. := e.getAppFn then + if let some info ← getCoeFnInfo? fn then + if e.getAppNumArgs >= info.numArgs then + return (← countHeadCoes (e.getArg! info.coercee)) + 1 + return 0 + +/-- Counts how many coercions are inside the expression, including the head ones. -/ +partial def countCoes (e : Expr) : MetaM Nat := + lambdaTelescope e fun _ e => do + if let Expr.const fn .. := e.getAppFn then + if let some info ← getCoeFnInfo? fn then + if e.getAppNumArgs >= info.numArgs then + let mut coes := (← countHeadCoes (e.getArg! info.coercee)) + 1 + for i in [info.numArgs:e.getAppNumArgs] do + coes := coes + (← countCoes (e.getArg! i)) + return coes + return (← (← getSimpArgs e).mapM countCoes).foldl (·+·) 0 + +/-- Counts how many coercions are inside the expression, excluding the head ones. -/ +def countInternalCoes (e : Expr) : MetaM Nat := + return (← countCoes e) - (← countHeadCoes e) + +/-- Classifies a declaration of type `ty` as a `norm_cast` rule. -/ +def classifyType (ty : Expr) : MetaM Label := + forallTelescopeReducing ty fun _ ty => do + let ty ← whnf ty + let (lhs, rhs) ← + if ty.isAppOfArity ``Eq 3 then pure (ty.getArg! 1, ty.getArg! 2) + else if ty.isAppOfArity ``Iff 2 then pure (ty.getArg! 0, ty.getArg! 1) + else throwError "norm_cast: lemma must be = or ↔, but is{indentExpr ty}" + let lhsCoes ← countCoes lhs + if lhsCoes = 0 then + throwError "norm_cast: badly shaped lemma, lhs must contain at least one coe{indentExpr lhs}" + let lhsHeadCoes ← countHeadCoes lhs + let rhsHeadCoes ← countHeadCoes rhs + let rhsInternalCoes ← countInternalCoes rhs + if lhsHeadCoes = 0 then + return Label.elim + else if lhsHeadCoes = 1 then do + unless rhsHeadCoes = 0 do + throwError "norm_cast: badly shaped lemma, rhs can't start with coe{indentExpr rhs}" + if rhsInternalCoes = 0 then + return Label.squash + else + return Label.move + else if rhsHeadCoes < lhsHeadCoes then do + return Label.squash + else do + throwError "\ + norm_cast: badly shaped shaped squash lemma, \ + rhs must have fewer head coes than lhs{indentExpr ty}" + +/-- The `push_cast` simp attribute. -/ +builtin_initialize pushCastExt : SimpExtension ← + registerSimpAttr `push_cast "\ + The `push_cast` simp attribute uses `norm_cast` lemmas \ + to move casts toward the leaf nodes of the expression." + +/-- The `norm_cast` attribute stores a simp set for each of the three types of `norm_cast` lemma. -/ +structure NormCastExtension where + /-- A simp set which lifts coercions to the top level. -/ + up : SimpExtension + /-- A simp set which pushes coercions to the leaves. -/ + down : SimpExtension + /-- A simp set which simplifies transitive coercions. -/ + squash : SimpExtension +deriving Inhabited + +/-- The `norm_cast` extension data. -/ +builtin_initialize normCastExt : NormCastExtension ← pure { + up := ← mkSimpExt (decl_name% ++ `up) + down := ← mkSimpExt (decl_name% ++ `down) + squash := ← mkSimpExt (decl_name% ++ `squash) +} + +/-- `addElim decl` adds `decl` as an `elim` lemma to be used by `norm_cast`. -/ +def addElim (decl : Name) + (kind := AttributeKind.global) (prio := eval_prio default) : MetaM Unit := + addSimpTheorem normCastExt.up decl (post := true) (inv := false) kind prio + +/-- `addMove decl` adds `decl` as a `move` lemma to be used by `norm_cast`. -/ +def addMove (decl : Name) + (kind := AttributeKind.global) (prio := eval_prio default) : MetaM Unit := do + addSimpTheorem pushCastExt decl (post := true) (inv := false) kind prio + addSimpTheorem normCastExt.up decl (post := true) (inv := true) kind prio + addSimpTheorem normCastExt.down decl (post := true) (inv := false) kind prio + +/-- `addSquash decl` adds `decl` as a `squash` lemma to be used by `norm_cast`. -/ +def addSquash (decl : Name) + (kind := AttributeKind.global) (prio := eval_prio default) : MetaM Unit := do + addSimpTheorem pushCastExt decl (post := true) (inv := false) kind prio + addSimpTheorem normCastExt.squash decl (post := true) (inv := false) kind prio + addSimpTheorem normCastExt.down decl (post := true) (inv := false) kind prio + +/-- `addInfer decl` infers the label of `decl` (`elim`, `move`, or `squash`) and arranges for it to +be used by `norm_cast`. + +* elim lemma: LHS has 0 head coes and ≥ 1 internal coe +* move lemma: LHS has 1 head coe and 0 internal coes, RHS has 0 head coes and ≥ 1 internal coes +* squash lemma: LHS has ≥ 1 head coes and 0 internal coes, RHS has fewer head coes +-/ +def addInfer (decl : Name) + (kind := AttributeKind.global) (prio := eval_prio default) : MetaM Unit := do + let ty := (← getConstInfo decl).type + match ← classifyType ty with + | Label.elim => addElim decl kind prio + | Label.squash => addSquash decl kind prio + | Label.move => addMove decl kind prio + +builtin_initialize registerBuiltinAttribute { + name := `norm_cast + descr := "attribute for norm_cast" + add := fun decl stx kind => MetaM.run' do + let `(attr| norm_cast $[$label:normCastLabel]? $[$prio]?) := stx | unreachable! + let prio := (prio.bind (·.1.isNatLit?)).getD (eval_prio default) + match label.bind (·.1.isStrLit?) with + | "elim" => addElim decl kind prio + | "move" => addMove decl kind prio + | "squash" => addSquash decl kind prio + | none => addInfer decl kind prio + | _ => unreachable! +} + +end Lean.Meta.NormCast diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Refl.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Refl.lean new file mode 100644 index 0000000000000000000000000000000000000000..597735c7822f2bd1de3cb4fdff17e6e9793ab215 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Refl.lean @@ -0,0 +1,58 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Reduce +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Apply + +namespace Lean.Meta + +/-- +Close given goal using `Eq.refl`. + +See `Lean.MVarId.applyRfl` for the variant that also consults `@[refl]` lemmas, and which +backs the `rfl` tactic. +-/ +def _root_.Lean.MVarId.refl (mvarId : MVarId) : MetaM Unit := do + mvarId.withContext do + mvarId.checkNotAssigned `refl + let targetType ← mvarId.getType' + unless targetType.isAppOfArity ``Eq 3 do + throwTacticEx `rfl mvarId m!"equality expected{indentExpr targetType}" + let lhs ← instantiateMVars targetType.appFn!.appArg! + let rhs ← instantiateMVars targetType.appArg! + let success ← isDefEq lhs rhs + unless success do + throwTacticEx `rfl mvarId m!"equality lhs{indentExpr lhs}\nis not definitionally equal to rhs{indentExpr rhs}" + let us := targetType.getAppFn.constLevels! + let α := targetType.appFn!.appFn!.appArg! + mvarId.assign (mkApp2 (mkConst ``Eq.refl us) α lhs) + +/-- +Try to apply `heq_of_eq`. If successful, then return new goal, otherwise return `mvarId`. +-/ +def _root_.Lean.MVarId.heqOfEq (mvarId : MVarId) : MetaM MVarId := + mvarId.withContext do + let some [mvarId] ← observing? do mvarId.apply (mkConst ``heq_of_eq [← mkFreshLevelMVar]) | return mvarId + return mvarId + +/-- +Try to apply `eq_of_heq`. If successful, then return new goal, otherwise return `mvarId`. +-/ +def _root_.Lean.MVarId.eqOfHEq (mvarId : MVarId) : MetaM MVarId := + mvarId.withContext do + let some [mvarId] ← observing? do mvarId.apply (mkConst ``eq_of_heq [← mkFreshLevelMVar]) | return mvarId + return mvarId + +/-- +Close given goal using `HEq.refl`. +-/ +def _root_.Lean.MVarId.hrefl (mvarId : MVarId) : MetaM Unit := do + mvarId.withContext do + let some [] ← observing? do mvarId.apply (mkConst ``HEq.refl [← mkFreshLevelMVar]) + | throwTacticEx `hrefl mvarId + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rename.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rename.lean new file mode 100644 index 0000000000000000000000000000000000000000..358c198dffd5f1a60087b607562164162ed5186f --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rename.lean @@ -0,0 +1,21 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Util + +namespace Lean.Meta + +/-- +Rename the user-face naming for the free variable `fvarId` at `mvarId`. +-/ +def _root_.Lean.MVarId.rename (mvarId : MVarId) (fvarId : FVarId) (userNameNew : Name) : MetaM MVarId := mvarId.withContext do + mvarId.checkNotAssigned `rename + let lctxNew := (← getLCtx).setUserName fvarId userNameNew + let mvarNew ← mkFreshExprMVarAt lctxNew (← getLocalInstances) (← mvarId.getType) MetavarKind.syntheticOpaque (← mvarId.getTag) + mvarId.assign mvarNew + return mvarNew.mvarId! + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Repeat.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Repeat.lean new file mode 100644 index 0000000000000000000000000000000000000000..2082e40ca25fff20afecd268f802fc01c4eafe60 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Repeat.lean @@ -0,0 +1,65 @@ +/- +Copyright (c) 2022 Mario Carneiro. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Mario Carneiro, Kim Morrison +-/ +prelude +import Lean.Meta.Basic + +namespace Lean.Meta +/-- +Implementation of `repeat'` and `repeat1'`. + +`repeat'Core f` runs `f` on all of the goals to produce a new list of goals, +then runs `f` again on all of those goals, and repeats until `f` fails on all remaining goals, +or until `maxIters` total calls to `f` have occurred. + +Returns a boolean indicating whether `f` succeeded at least once, and +all the remaining goals (i.e. those on which `f` failed). +-/ +def repeat'Core [Monad m] [MonadExcept ε m] [MonadBacktrack s m] [MonadMCtx m] + (f : MVarId → m (List MVarId)) (goals : List MVarId) (maxIters := 100000) : + m (Bool × List MVarId) := do + let (progress, acc) ← go maxIters false goals [] #[] + pure (progress, (← acc.filterM fun g => not <$> g.isAssigned).toList) +where + /-- Auxiliary for `repeat'Core`. `repeat'Core.go f maxIters progress goals stk acc` evaluates to + essentially `acc.toList ++ repeat' f (goals::stk).join maxIters`: that is, `acc` are goals we will + not revisit, and `(goals::stk).join` is the accumulated todo list of subgoals. -/ + go : Nat → Bool → List MVarId → List (List MVarId) → Array MVarId → m (Bool × Array MVarId) + | _, p, [], [], acc => pure (p, acc) + | n, p, [], goals::stk, acc => go n p goals stk acc + | n, p, g::goals, stk, acc => do + if ← g.isAssigned then + go n p goals stk acc + else + match n with + | 0 => pure <| (p, acc.push g ++ goals |> stk.foldl .appendList) + | n+1 => + match ← observing? (f g) with + | some goals' => go n true goals' (goals::stk) acc + | none => go n p goals stk (acc.push g) +termination_by n _ goals stk => (n, stk, goals) + +/-- +`repeat' f` runs `f` on all of the goals to produce a new list of goals, +then runs `f` again on all of those goals, and repeats until `f` fails on all remaining goals, +or until `maxIters` total calls to `f` have occurred. +Always succeeds (returning the original goals if `f` fails on all of them). +-/ +def repeat' [Monad m] [MonadExcept ε m] [MonadBacktrack s m] [MonadMCtx m] + (f : MVarId → m (List MVarId)) (goals : List MVarId) (maxIters := 100000) : m (List MVarId) := + repeat'Core f goals maxIters <&> (·.2) + +/-- +`repeat1' f` runs `f` on all of the goals to produce a new list of goals, +then runs `f` again on all of those goals, and repeats until `f` fails on all remaining goals, +or until `maxIters` total calls to `f` have occurred. +Fails if `f` does not succeed at least once. +-/ +def repeat1' [Monad m] [MonadError m] [MonadExcept ε m] [MonadBacktrack s m] [MonadMCtx m] + (f : MVarId → m (List MVarId)) (goals : List MVarId) (maxIters := 100000) : m (List MVarId) := do + let (.true, goals) ← repeat'Core f goals maxIters | throwError "repeat1' made no progress" + pure goals + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Replace.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Replace.lean new file mode 100644 index 0000000000000000000000000000000000000000..f3dd25ed3ed189f549abcc7613686077b4fed932 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Replace.lean @@ -0,0 +1,266 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Util.ForEachExpr +import Lean.Elab.InfoTree.Main +import Lean.Meta.AppBuilder +import Lean.Meta.MatchUtil +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Revert +import Lean.Meta.Tactic.Intro +import Lean.Meta.Tactic.Clear +import Lean.Meta.Tactic.Assert + +namespace Lean.Meta + +/-- + Convert the given goal `Ctx |- target` into `Ctx |- targetNew` using an equality proof `eqProof : target = targetNew`. + It assumes `eqProof` has type `target = targetNew` -/ +def _root_.Lean.MVarId.replaceTargetEq (mvarId : MVarId) (targetNew : Expr) (eqProof : Expr) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `replaceTarget + let tag ← mvarId.getTag + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew tag + let target ← mvarId.getType + let u ← getLevel target + let eq ← mkEq target targetNew + let newProof := mkExpectedPropHint eqProof eq + let val := mkAppN (Lean.mkConst `Eq.mpr [u]) #[target, targetNew, newProof, mvarNew] + mvarId.assign val + return mvarNew.mvarId! + +/-- +Converts the given goal `Ctx |- target` into `Ctx |- targetNew`. It assumes the goals are definitionally equal. +We use the proof term +``` +@id target mvarNew +``` +to create a checkpoint. + +If `targetNew` is equal to `target`, then returns `mvarId` unchanged. +Uses `Expr.equal` for the comparison so that it is possible to update binder names, etc., which are user-visible. +-/ +def _root_.Lean.MVarId.replaceTargetDefEq (mvarId : MVarId) (targetNew : Expr) : MetaM MVarId := + mvarId.withContext do + mvarId.checkNotAssigned `change + let target ← mvarId.getType + if Expr.equal target targetNew then + return mvarId + else + let tag ← mvarId.getTag + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew tag + let newVal ← mkExpectedTypeHint mvarNew target + mvarId.assign newVal + return mvarNew.mvarId! + +private def replaceLocalDeclCore (mvarId : MVarId) (fvarId : FVarId) (typeNew : Expr) (eqProof : Expr) : MetaM AssertAfterResult := + mvarId.withContext do + let localDecl ← fvarId.getDecl + let typeNewPr ← mkEqMP eqProof (mkFVar fvarId) + /- `typeNew` may contain variables that occur after `fvarId`. + Thus, we use the auxiliary function `findMaxFVar` to ensure `typeNew` is well-formed at the + position we are inserting it. + We must `instantiateMVars` first to ensure that there is no mvar in `typeNew` which is + assigned to some later-occurring fvar. -/ + let (_, localDecl') ← findMaxFVar (← instantiateMVars typeNew) |>.run localDecl + let result ← mvarId.assertAfter localDecl'.fvarId localDecl.userName typeNew typeNewPr + (do let mvarIdNew ← result.mvarId.clear fvarId + pure { result with mvarId := mvarIdNew }) + <|> pure result +where + findMaxFVar (e : Expr) : StateRefT LocalDecl MetaM Unit := + e.forEach' fun e => do + if e.isFVar then + let localDecl' ← e.fvarId!.getDecl + modify fun localDecl => if localDecl'.index > localDecl.index then localDecl' else localDecl + return false + else + return e.hasFVar + +/-- + Replace type of the local declaration with id `fvarId` with one with the same user-facing name, but with type `typeNew`. + This method assumes `eqProof` is a proof that the type of `fvarId` is equal to `typeNew`. + This tactic actually adds a new declaration and (tries to) clear the old one. + If the old one cannot be cleared, then at least its user-facing name becomes inaccessible. + + The new local declaration is inserted at the soonest point after `fvarId` at which it is + well-formed. That is, if `typeNew` involves declarations which occur later than `fvarId` in the + local context, the new local declaration will be inserted immediately after the latest-occurring + one. Otherwise, it will be inserted immediately after `fvarId`. + + Note: `replaceLocalDecl` should not be used when unassigned pending mvars might be present in + `typeNew`, as these may later be synthesized to fvars which occur after `fvarId` (by e.g. + `Term.withSynthesize` or `Term.synthesizeSyntheticMVars`) . + -/ +abbrev _root_.Lean.MVarId.replaceLocalDecl (mvarId : MVarId) (fvarId : FVarId) (typeNew : Expr) (eqProof : Expr) : MetaM AssertAfterResult := + replaceLocalDeclCore mvarId fvarId typeNew eqProof + +/-- +Replaces the type of `fvarId` at `mvarId` with `typeNew`. +Remark: this method assumes that `typeNew` is definitionally equal to the current type of `fvarId`. + +If `typeNew` is equal to current type of `fvarId`, then returns `mvarId` unchanged. +Uses `Expr.equal` for the comparison so that it is possible to update binder names, etc., which are user-visible. +-/ +def _root_.Lean.MVarId.replaceLocalDeclDefEq (mvarId : MVarId) (fvarId : FVarId) (typeNew : Expr) : MetaM MVarId := do + mvarId.withContext do + if Expr.equal typeNew (← fvarId.getType) then + return mvarId + else + let mvarDecl ← mvarId.getDecl + let lctxNew := (← getLCtx).modifyLocalDecl fvarId (·.setType typeNew) + let mvarNew ← mkFreshExprMVarAt lctxNew (← getLocalInstances) mvarDecl.type mvarDecl.kind mvarDecl.userName + mvarId.assign mvarNew + return mvarNew.mvarId! + +/-- +Replace the target type of `mvarId` with `typeNew`. +If `checkDefEq = false`, this method assumes that `typeNew` is definitionally equal to the current target type. +If `checkDefEq = true`, throw an error if `typeNew` is not definitionally equal to the current target type. +-/ +def _root_.Lean.MVarId.change (mvarId : MVarId) (targetNew : Expr) (checkDefEq := true) : MetaM MVarId := mvarId.withContext do + let target ← mvarId.getType + if checkDefEq then + unless (← isDefEq target targetNew) do + throwTacticEx `change mvarId m!"given type{indentExpr targetNew}\nis not definitionally equal to{indentExpr target}" + mvarId.replaceTargetDefEq targetNew + +/-- +Executes the revert/intro pattern, running the continuation `k` after temporarily reverting +the given local variables from the local context of the metavariable `mvarId`, +and then re-introducing the local variables specified by `k`. + +- `mvarId` is the goal metavariable to operate on. +- `fvarIds` is an array of `fvarIds` to revert in the order specified. + An error is thrown if they cannot be reverted in order. +- `clearAuxDeclsInsteadOfRevert` is configuration passed to `Lean.MVarId.revert`. +- `k` is the continuation run once the local variables have been reverted. + It is provided `mvarId` after the requested variables have been reverted along with the array of reverted variables. + This array always contains `fvarIds`, but it may contain additional variables that were reverted due to dependencies. + The continuation returns a value, a new goal, and an _aliasing array_. + +Once `k` has completed, one variable is introduced per entry in the aliasing array. +* If the entry is `none`, the variable is just introduced. +* If the entry is `some fv` (where `fv` is a variable from `fvarIds`), + the variable is introduced and then recorded as an alias of `fv` in the info tree. + This for example causes the unused variable linter as seeing `fv` and this newly introduced variable as being "the same". + +For example, if `k` leaves all the reverted variables in the same order, +having it return `fvars.map .some` as the aliasing array causes those variables to be re-introduced and aliased +to the original local variables. + +Returns the value returned by `k` along with the resulting goal after variable introduction. + +See `Lean.MVarId.changeLocalDecl` for an example. The motivation is that to work on a local variable, +you need to move it into the goal, alter the goal, and then bring it back into the local context, +all while keeping track of any other local variables that depend on this one. +-/ +def _root_.Lean.MVarId.withReverted (mvarId : MVarId) (fvarIds : Array FVarId) + (k : MVarId → Array FVarId → MetaM (α × Array (Option FVarId) × MVarId)) + (clearAuxDeclsInsteadOfRevert := false) : MetaM (α × MVarId) := do + let (xs, mvarId) ← mvarId.revert fvarIds true clearAuxDeclsInsteadOfRevert + let (r, xs', mvarId) ← k mvarId xs + let (ys, mvarId) ← mvarId.introNP xs'.size + mvarId.withContext do + for x? in xs', y in ys do + if let some x := x? then + Elab.pushInfoLeaf (.ofFVarAliasInfo { id := y, baseId := x, userName := ← y.getUserName }) + return (r, mvarId) + +/-- +Like `Lean.MVarId.withReverted`, but reverts all local variables starting from `fvarId`. +-/ +def _root_.Lean.MVarId.withRevertedFrom (mvarId : MVarId) (fvarId : FVarId) + (k : MVarId → Array FVarId → MetaM (α × Array (Option FVarId) × MVarId)) : MetaM (α × MVarId) := do + let (xs, mvarId) ← mvarId.revertFrom fvarId + let (r, xs', mvarId) ← k mvarId xs + let (ys, mvarId) ← mvarId.introNP xs'.size + mvarId.withContext do + for x? in xs', y in ys do + if let some x := x? then + Elab.pushInfoLeaf (.ofFVarAliasInfo { id := y, baseId := x, userName := ← y.getUserName }) + return (r, mvarId) + +/-- +Replaces the type of the free variable `fvarId` with `typeNew`. + +If `checkDefEq` is `true` then an error is thrown if `typeNew` is not definitionally +equal to the type of `fvarId`. Otherwise this function assumes `typeNew` and the type +of `fvarId` are definitionally equal. + +This function is the same as `Lean.MVarId.changeLocalDecl` but makes sure to push substitution +information into the info tree. +-/ +def _root_.Lean.MVarId.changeLocalDecl (mvarId : MVarId) (fvarId : FVarId) (typeNew : Expr) + (checkDefEq := true) : MetaM MVarId := do + mvarId.checkNotAssigned `changeLocalDecl + let (_, mvarId) ← mvarId.withReverted #[fvarId] fun mvarId fvars => mvarId.withContext do + let check (typeOld : Expr) : MetaM Unit := do + if checkDefEq then + unless ← isDefEq typeNew typeOld do + throwTacticEx `changeLocalDecl mvarId + m!"given type{indentExpr typeNew}\nis not definitionally equal to{indentExpr typeOld}" + let finalize (targetNew : Expr) := do + return ((), fvars.map .some, ← mvarId.replaceTargetDefEq targetNew) + match ← mvarId.getType with + | .forallE n d b bi => do check d; finalize (.forallE n typeNew b bi) + | .letE n t v b ndep => do check t; finalize (.letE n typeNew v b ndep) + | _ => throwTacticEx `changeLocalDecl mvarId "unexpected auxiliary target" + return mvarId + +/-- +Modify `mvarId` target type using `f`. +-/ +def _root_.Lean.MVarId.modifyTarget (mvarId : MVarId) (f : Expr → MetaM Expr) : MetaM MVarId := do + mvarId.withContext do + mvarId.checkNotAssigned `modifyTarget + mvarId.change (← f (← mvarId.getType)) (checkDefEq := false) + +/-- +Modify `mvarId` target type left-hand-side using `f`. +Throw an error if target type is not an equality. +-/ +def _root_.Lean.MVarId.modifyTargetEqLHS (mvarId : MVarId) (f : Expr → MetaM Expr) : MetaM MVarId := do + mvarId.modifyTarget fun target => do + if let some (_, lhs, rhs) ← matchEq? target then + mkEq (← f lhs) rhs + else + throwTacticEx `modifyTargetEqLHS mvarId m!"equality expected{indentExpr target}" + +/-- +Clears the value of the local definition `fvarId`. Ensures that the resulting goal state +is still type correct. Throws an error if it is a local hypothesis without a value. + +Preserves the order of the local context. +-/ +def _root_.Lean.MVarId.clearValue (mvarId : MVarId) (fvarId : FVarId) : MetaM MVarId := do + mvarId.checkNotAssigned `clear_value + let tag ← mvarId.getTag + let (_, mvarId) ← mvarId.withRevertedFrom fvarId fun mvarId' fvars => mvarId'.withContext do + let tgt ← mvarId'.getType + unless tgt.isLet do + mvarId.withContext <| + throwTacticEx `clear_value mvarId m!"hypothesis `{Expr.fvar fvarId}` is not a local definition." + let tgt' := Expr.forallE tgt.letName! tgt.letType! tgt.letBody! .default + /- + Note: `isTypeCorrect` does not instantiate metavariables. Ideally this should not matter, + however delayed assigned metavariables assume that the arguments applied to them have the correct + definitional properties (e.g. an argument might be for the value of a `let` binding). + Instantiating the metavariables before `isTypeCorrect` lets examples like the one reported in + https://github.com/leanprover-community/mathlib4/issues/25053 go through. + (See `tests/lean/run/clear_value.lean`, 25053.) + -/ + let tgt' ← instantiateMVars tgt' + unless ← isTypeCorrect tgt' do + mvarId.withContext <| + throwTacticEx `clear_value mvarId + m!"cannot clear {Expr.fvar fvarId}, the resulting context is not type correct." + let mvarId'' ← mkFreshExprSyntheticOpaqueMVar tgt' tag + mvarId'.assign <| .app mvarId'' tgt.letValue! + return ((), fvars.map .some, mvarId''.mvarId!) + return mvarId + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Revert.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Revert.lean new file mode 100644 index 0000000000000000000000000000000000000000..4e1bab54a840327637de54c823345d3a867ec47e --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Revert.lean @@ -0,0 +1,64 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Clear + +namespace Lean.Meta + +/-- +Revert free variables `fvarIds` at goal `mvarId`. +-/ +def _root_.Lean.MVarId.revert (mvarId : MVarId) (fvarIds : Array FVarId) (preserveOrder : Bool := false) + (clearAuxDeclsInsteadOfRevert := false) : MetaM (Array FVarId × MVarId) := do + if fvarIds.isEmpty then + pure (#[], mvarId) + else mvarId.withContext do + mvarId.checkNotAssigned `revert + unless clearAuxDeclsInsteadOfRevert do + for fvarId in fvarIds do + if (← fvarId.getDecl) |>.isAuxDecl then + throwError "failed to revert {mkFVar fvarId}, it is an auxiliary declaration created to represent recursive definitions" + let fvars := fvarIds.map mkFVar + let toRevert ← collectForwardDeps fvars preserveOrder + /- We should clear any `auxDecl` in `toRevert` -/ + let mut mvarId := mvarId + let mut toRevertNew := #[] + for x in toRevert do + if (← x.fvarId!.getDecl).isAuxDecl then + mvarId ← mvarId.clear x.fvarId! + else + toRevertNew := toRevertNew.push x + let tag ← mvarId.getTag + -- TODO: the following code can be optimized because `MetavarContext.revert` will compute `collectDeps` again. + -- We should factor out the relevant part + + -- Set metavariable kind to natural to make sure `revert` will assign it. + mvarId.setKind .natural + let (e, toRevert) ← + try + liftMkBindingM <| MetavarContext.revert toRevertNew mvarId preserveOrder + finally + mvarId.setKind .syntheticOpaque + let mvar := e.getAppFn + mvar.mvarId!.setKind .syntheticOpaque + mvar.mvarId!.setTag tag + return (toRevert.map Expr.fvarId!, mvar.mvarId!) + +/-- Reverts all local declarations after `fvarId`. -/ +def _root_.Lean.MVarId.revertAfter (mvarId : MVarId) (fvarId : FVarId) : MetaM (Array FVarId × MVarId) := + mvarId.withContext do + let localDecl ← fvarId.getDecl + let fvarIds := (← getLCtx).foldl (init := #[]) (start := localDecl.index+1) fun fvarIds decl => fvarIds.push decl.fvarId + mvarId.revert fvarIds (preserveOrder := true) (clearAuxDeclsInsteadOfRevert := true) + +/-- Reverts all local declarations starting from `fvarId`. -/ +def _root_.Lean.MVarId.revertFrom (mvarId : MVarId) (fvarId : FVarId) : MetaM (Array FVarId × MVarId) := + mvarId.withContext do + let localDecl ← fvarId.getDecl + let fvarIds := (← getLCtx).foldl (init := #[]) (start := localDecl.index) fun fvarIds decl => fvarIds.push decl.fvarId + mvarId.revert fvarIds (preserveOrder := true) (clearAuxDeclsInsteadOfRevert := true) + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rewrite.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rewrite.lean new file mode 100644 index 0000000000000000000000000000000000000000..18a174cba141fb532ec7d30f231f95b2c257e942 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rewrite.lean @@ -0,0 +1,94 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.AppBuilder +import Lean.Meta.MatchUtil +import Lean.Meta.KAbstract +import Lean.Meta.Check +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Apply +import Lean.Meta.BinderNameHint + +namespace Lean.Meta + +structure RewriteResult where + eNew : Expr + eqProof : Expr + mvarIds : List MVarId -- new goals + +/-- +Rewrite goal `mvarId` +-/ +def _root_.Lean.MVarId.rewrite (mvarId : MVarId) (e : Expr) (heq : Expr) + (symm : Bool := false) (config := { : Rewrite.Config }) : MetaM RewriteResult := + mvarId.withContext do + mvarId.checkNotAssigned `rewrite + let heqIn := heq + let heqType ← instantiateMVars (← inferType heq) + let (newMVars, binderInfos, heqType) ← forallMetaTelescopeReducing heqType + let heq := mkAppN heq newMVars + let cont (heq heqType : Expr) : MetaM RewriteResult := do + match (← matchEq? heqType) with + | none => throwTacticEx `rewrite mvarId m!"equality or iff proof expected{indentExpr heqType}" + | some (α, lhs, rhs) => + let cont (heq heqType lhs rhs : Expr) : MetaM RewriteResult := do + if lhs.getAppFn.isMVar then + throwTacticEx `rewrite mvarId m!"pattern is a metavariable{indentExpr lhs}\nfrom equation{indentExpr heqType}" + let e ← instantiateMVars e + let eAbst ← withConfig (fun oldConfig => { config, oldConfig with }) <| kabstract e lhs config.occs + unless eAbst.hasLooseBVars do + throwTacticEx `rewrite mvarId m!"did not find instance of the pattern in the target expression{indentExpr lhs}" + -- construct rewrite proof + let eNew := eAbst.instantiate1 rhs + let eNew ← instantiateMVars eNew + let eNew ← if rhs.hasBinderNameHint then eNew.resolveBinderNameHint else pure eNew + let eType ← inferType e + let motive := Lean.mkLambda `_a BinderInfo.default α eAbst + try + check motive + catch ex => + throwTacticEx `rewrite mvarId m!"\ + motive is not type correct:{indentD motive}\nError: {ex.toMessageData}\ + \n\n\ + Explanation: The rewrite tactic rewrites an expression 'e' using an equality 'a = b' by the following process. \ + First, it looks for all 'a' in 'e'. Second, it tries to abstract these occurrences of 'a' to create a function 'm := fun _a => ...', called the *motive*, \ + with the property that 'm a' is definitionally equal to 'e'. \ + Third, we observe that '{.ofConstName ``congrArg}' implies that 'm a = m b', which can be used with lemmas such as '{.ofConstName ``Eq.mpr}' to change the goal. \ + However, if 'e' depends on specific properties of 'a', then the motive 'm' might not typecheck.\ + \n\n\ + Possible solutions: use rewrite's 'occs' configuration option to limit which occurrences are rewritten, \ + or use 'simp' or 'conv' mode, which have strategies for certain kinds of dependencies \ + (these tactics can handle proofs and '{.ofConstName ``Decidable}' instances whose types depend on the rewritten term, \ + and 'simp' can apply user-defined '@[congr]' theorems as well)." + unless (← withLocalDeclD `_a α fun a => do isDefEq (← inferType (eAbst.instantiate1 a)) eType) do + -- NB: using motive.arrow? would disallow motives where the dependency + -- can be reduced away + throwTacticEx `rewrite mvarId m!"motive is dependent{indentD motive}" + let u1 ← getLevel α + let u2 ← getLevel eType + let eqPrf := mkApp6 (.const ``congrArg [u1, u2]) α eType lhs rhs motive heq + postprocessAppMVars `rewrite mvarId newMVars binderInfos + (synthAssignedInstances := !tactic.skipAssignedInstances.get (← getOptions)) + let newMVarIds ← newMVars.map Expr.mvarId! |>.filterM fun mvarId => not <$> mvarId.isAssigned + let otherMVarIds ← getMVarsNoDelayed heqIn + let otherMVarIds := otherMVarIds.filter (!newMVarIds.contains ·) + let newMVarIds := newMVarIds ++ otherMVarIds + pure { eNew := eNew, eqProof := eqPrf, mvarIds := newMVarIds.toList } + match symm with + | false => cont heq heqType lhs rhs + | true => do + let heq ← mkEqSymm heq + let heqType ← mkEq rhs lhs + cont heq heqType rhs lhs + match heqType.iff? with + | some (lhs, rhs) => + let heqType ← mkEq lhs rhs + let heq := mkApp3 (mkConst `propext) lhs rhs heq + cont heq heqType + | none => + cont heq heqType + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rewrites.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rewrites.lean new file mode 100644 index 0000000000000000000000000000000000000000..7728b37e372a14f1ba4e3e2f8bcaf6994955e920 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rewrites.lean @@ -0,0 +1,342 @@ +/- +Copyright (c) 2023 Kim Morrison. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +prelude +import Lean.Meta.LazyDiscrTree +import Lean.Meta.Tactic.Assumption +import Lean.Meta.Tactic.Rewrite +import Lean.Meta.Tactic.Refl +import Lean.Meta.Tactic.SolveByElim +import Lean.Meta.Tactic.TryThis +import Lean.Util.Heartbeats + +namespace Lean.Meta.Rewrites + +open Lean.Meta.LazyDiscrTree (InitEntry MatchResult) +open Lean.Meta.SolveByElim + +builtin_initialize registerTraceClass `Tactic.rewrites +builtin_initialize registerTraceClass `Tactic.rewrites.lemmas + +/-- Extract the lemma, with arguments, that was used to produce a `RewriteResult`. -/ +-- This assumes that `r.eqProof` was constructed as: +-- `mkApp6 (.const ``congrArg _) α eType lhs rhs motive heq` +-- in `Lean.Meta.Tactic.Rewrite` and we want `heq`. +def rewriteResultLemma (r : RewriteResult) : Option Expr := + if r.eqProof.isAppOfArity ``congrArg 6 then + r.eqProof.getArg! 5 + else + none + +/-- Weight to multiply the "specificity" of a rewrite lemma by when rewriting forwards. -/ +def forwardWeight := 2 +/-- Weight to multiply the "specificity" of a rewrite lemma by when rewriting backwards. -/ +def backwardWeight := 1 + +inductive RwDirection : Type where + | forward : RwDirection + | backward : RwDirection + +private def addImport (name : Name) (constInfo : ConstantInfo) : + MetaM (Array (InitEntry (Name × RwDirection))) := do + if constInfo.isUnsafe then return #[] + if !allowCompletion (←getEnv) name then return #[] + -- We now remove some injectivity lemmas which are not useful to rewrite by. + match name with + | .str _ n => if n = "injEq" ∨ n = "sizeOf_spec" ∨ n.endsWith "_inj" ∨ n.endsWith "_inj'" then return #[] + | _ => pure () + -- Don't report lemmas from metaprogramming namespaces. + if name.isMetaprogramming then return #[] + withNewMCtxDepth do withReducible do + forallTelescopeReducing constInfo.type fun _ type => do + match type.getAppFnArgs with + | (``Eq, #[_, lhs, rhs]) + | (``Iff, #[lhs, rhs]) => do + let a := Array.mkEmpty 2 + let a := a.push (← InitEntry.fromExpr lhs (name, RwDirection.forward)) + let a := a.push (← InitEntry.fromExpr rhs (name, RwDirection.backward)) + pure a + | _ => return #[] + +/-- Select `=` and `↔` local hypotheses. -/ +def localHypotheses (except : List FVarId := []) : MetaM (Array (Expr × Bool × Nat)) := do + let r ← getLocalHyps + let mut result := #[] + for h in r do + if except.contains h.fvarId! then continue + let (_, _, type) ← forallMetaTelescopeReducing (← inferType h) + let type ← whnfR type + match type.getAppFnArgs with + | (``Eq, #[_, _, _]) + | (``Iff, #[_, _]) => do + result := result.push (h, false, forwardWeight) + |>.push (h, true, backwardWeight) + | _ => pure () + return result + +/-- +We drop `.star` and `Eq * * *` from the discriminator trees because +they match too much. +-/ +def droppedKeys : List (List LazyDiscrTree.Key) := [[.star], [.const `Eq 3, .star, .star, .star]] + +def createModuleTreeRef : MetaM (LazyDiscrTree.ModuleDiscrTreeRef (Name × RwDirection)) := + LazyDiscrTree.createModuleTreeRef addImport droppedKeys + +private def ExtState := IO.Ref (Option (LazyDiscrTree (Name × RwDirection))) + +private builtin_initialize ExtState.default : IO.Ref (Option (LazyDiscrTree (Name × RwDirection))) ← do + IO.mkRef .none + +private instance : Inhabited ExtState where + default := ExtState.default + +private builtin_initialize ext : EnvExtension ExtState ← + registerEnvExtension (IO.mkRef .none) + +/-- +The maximum number of constants an individual task may perform. + +The value was picked because it roughly corresponded to 50ms of work on the +machine this was developed on. Smaller numbers did not seem to improve +performance when importing Std and larger numbers (<10k) seemed to degrade +initialization performance. +-/ +private def constantsPerImportTask : Nat := 6500 + +def incPrio : Nat → Name × RwDirection → Name × Bool × Nat +| q, (nm, d) => + match d with + | .forward => (nm, false, 2 * q) + | .backward => (nm, true, q) + +/-- Create function for finding relevant declarations. -/ +def rwFindDecls (moduleRef : LazyDiscrTree.ModuleDiscrTreeRef (Name × RwDirection)) : Expr → MetaM (Array (Name × Bool × Nat)) := + LazyDiscrTree.findMatchesExt moduleRef ext addImport + (droppedKeys := droppedKeys) + (constantsPerTask := constantsPerImportTask) + (adjustResult := incPrio) + +/-- Data structure recording a potential rewrite to report from the `rw?` tactic. -/ +structure RewriteResult where + /-- The lemma we rewrote by. + This is `Expr`, not just a `Name`, as it may be a local hypothesis. -/ + expr : Expr + /-- `True` if we rewrote backwards (i.e. with `rw [← h]`). -/ + symm : Bool + /-- The "weight" of the rewrite. This is calculated based on how specific the rewrite rule was. -/ + weight : Nat + /-- The result from the `rw` tactic. -/ + result : Meta.RewriteResult + /-- The metavariable context after the rewrite. + This needs to be stored as part of the result so we can backtrack the state. -/ + mctx : MetavarContext + rfl? : Bool + +/-- Check to see if this expression (which must be a type) can be closed by `with_reducible rfl`. -/ +def dischargableWithRfl? (mctx : MetavarContext) (e : Expr) : MetaM Bool := do + try + withoutModifyingState <| withMCtx mctx do + -- We use `withReducible` here to follow the behaviour of `rw`. + withReducible (← mkFreshExprMVar e).mvarId!.refl + pure true + catch _e => + pure false + +/-- +Pretty print the result of the rewrite. +-/ +private def RewriteResult.ppResult (r : RewriteResult) : MetaM String := + return (← ppExpr r.result.eNew).pretty + + +/-- Should we try discharging side conditions? If so, using `assumption`, or `solve_by_elim`? -/ +inductive SideConditions +| none +| assumption +| solveByElim + +/-- Shortcut for calling `solveByElim`. -/ +def solveByElim (goals : List MVarId) (depth : Nat := 6) : MetaM PUnit := do + -- There is only a marginal decrease in performance for using the `symm` option for `solveByElim`. + -- (measured via `lake build && time lake env lean test/librarySearch.lean`). + let cfg : SolveByElimConfig := { maxDepth := depth, exfalso := false, symm := true } + let ⟨lemmas, ctx⟩ ← mkAssumptionSet false false [] [] #[] + let [] ← SolveByElim.solveByElim cfg lemmas ctx goals + | failure + +def rwLemma (ctx : MetavarContext) (goal : MVarId) (target : Expr) (side : SideConditions := .solveByElim) + (lem : Expr ⊕ Name) (symm : Bool) (weight : Nat) : MetaM (Option RewriteResult) := + withMCtx ctx do + let some expr ← (match lem with + | .inl hyp => pure (some hyp) + | .inr lem => some <$> mkConstWithFreshMVarLevels lem <|> pure none) + | return none + trace[Tactic.rewrites] m!"considering {if symm then "← " else ""}{expr}" + let some result ← some <$> goal.rewrite target expr symm <|> pure none + | return none + if result.mvarIds.isEmpty then + let mctx ← getMCtx + let rfl? ← dischargableWithRfl? mctx result.eNew + return some { expr, symm, weight, result, mctx, rfl? } + else + -- There are side conditions, which we try to discharge using local hypotheses. + let discharge ← + match side with + | .none => pure false + | .assumption => ((fun _ => true) <$> result.mvarIds.mapM fun m => m.assumption) <|> pure false + | .solveByElim => (solveByElim result.mvarIds >>= fun _ => pure true) <|> pure false + match discharge with + | false => + return none + | true => + -- If we succeed, we need to reconstruct the expression to report that we rewrote by. + let some expr := rewriteResultLemma result | return none + let expr ← instantiateMVars expr + let (expr, symm) := if expr.isAppOfArity ``Eq.symm 4 then + (expr.getArg! 3, true) + else + (expr, false) + let mctx ← getMCtx + let rfl? ← dischargableWithRfl? mctx result.eNew + return some { expr, symm, weight, result, mctx, rfl? } + +/-- +Find keys which match the expression, or some subexpression. + +Note that repeated subexpressions will be visited each time they appear, +making this operation potentially very expensive. +It would be good to solve this problem! + +Implementation: we reverse the results from `getMatch`, +so that we return lemmas matching larger subexpressions first, +and amongst those we return more specific lemmas first. +-/ +partial def getSubexpressionMatches (op : Expr → MetaM (Array α)) (e : Expr) : MetaM (Array α) := do + match e with + | .bvar _ => return #[] + | .forallE _ _ _ _ => + forallTelescope e fun args body => do + args.foldlM (fun acc arg => return acc ++ (← getSubexpressionMatches op (← inferType arg))) + (← getSubexpressionMatches op body).reverse + | .lam _ _ _ _ + | .letE _ _ _ _ _ => + lambdaLetTelescope e (fun args body => do + args.foldlM (fun acc arg => return acc ++ (← getSubexpressionMatches op (← inferType arg))) + (← getSubexpressionMatches op body).reverse) + | _ => + let init := ((← op e).reverse) + e.foldlM (init := init) (fun a f => return a ++ (← getSubexpressionMatches op f)) + +/-- +Find lemmas which can rewrite the goal. + +See also `rewrites` for a more convenient interface. +-/ +def rewriteCandidates (hyps : Array (Expr × Bool × Nat)) + (moduleRef : LazyDiscrTree.ModuleDiscrTreeRef (Name × RwDirection)) + (target : Expr) + (forbidden : NameSet := ∅) : + MetaM (Array ((Expr ⊕ Name) × Bool × Nat)) := do + -- Get all lemmas which could match some subexpression + let candidates ← getSubexpressionMatches (rwFindDecls moduleRef) target + -- Sort them by our preferring weighting + -- (length of discriminant key, doubled for the forward implication) + let candidates := candidates.insertionSort fun (_, _, rp) (_, _, sp) => rp > sp + + -- Now deduplicate. We can't use `Array.deduplicateSorted` as we haven't completely sorted, + -- and in fact want to keep some of the residual ordering from the discrimination tree. + let mut forward : NameSet := ∅ + let mut backward : NameSet := ∅ + let mut deduped := #[] + for (l, s, w) in candidates do + if forbidden.contains l then continue + if s then + if ¬ backward.contains l then + deduped := deduped.push (l, s, w) + backward := backward.insert l + else + if ¬ forward.contains l then + deduped := deduped.push (l, s, w) + forward := forward.insert l + + trace[Tactic.rewrites.lemmas] m!"Candidate rewrite lemmas:\n{deduped}" + + let hyps := hyps.map fun ⟨hyp, symm, weight⟩ => (Sum.inl hyp, symm, weight) + let lemmas := deduped.map fun ⟨lem, symm, weight⟩ => (Sum.inr lem, symm, weight) + pure <| hyps ++ lemmas + +def RewriteResult.newGoal (r : RewriteResult) : Option Expr := + if r.rfl? then + none + else + some r.result.eNew + +open Elab Tactic in +def RewriteResult.addSuggestion (ref : Syntax) (r : RewriteResult) + (checkState? : Option Tactic.SavedState := .none) : TacticM Unit := do + withMCtx r.mctx do + Tactic.TryThis.addRewriteSuggestion ref [(r.expr, r.symm)] + (type? := r.newGoal.toLOption) (origSpan? := ← getRef) + (checkState? := checkState?.getD (← saveState)) + +structure RewriteResultConfig where + stopAtRfl : Bool + max : Nat + minHeartbeats : Nat + goal : MVarId + target : Expr + side : SideConditions := .solveByElim + mctx : MetavarContext + +def takeListAux (cfg : RewriteResultConfig) (seen : Std.HashMap String Unit) (acc : Array RewriteResult) + (xs : List ((Expr ⊕ Name) × Bool × Nat)) : MetaM (Array RewriteResult) := do + let mut seen := seen + let mut acc := acc + for (lem, symm, weight) in xs do + if (← getRemainingHeartbeats) < cfg.minHeartbeats then + return acc + if acc.size ≥ cfg.max then + return acc + let res ← + withoutModifyingState <| withMCtx cfg.mctx do + rwLemma cfg.mctx cfg.goal cfg.target cfg.side lem symm weight + match res with + | none => continue + | some r => + let s ← withoutModifyingState <| withMCtx r.mctx r.ppResult + if seen.contains s then + continue + let rfl? ← dischargableWithRfl? r.mctx r.result.eNew + if cfg.stopAtRfl then + if rfl? then + return #[r] + else + seen := seen.insert s () + acc := acc.push r + else + seen := seen.insert s () + acc := acc.push r + return acc + +/-- Find lemmas which can rewrite the goal. -/ +def findRewrites (hyps : Array (Expr × Bool × Nat)) + (moduleRef : LazyDiscrTree.ModuleDiscrTreeRef (Name × RwDirection)) + (goal : MVarId) (target : Expr) + (forbidden : NameSet := ∅) (side : SideConditions := .solveByElim) + (stopAtRfl : Bool) (max : Nat := 20) + (leavePercentHeartbeats : Nat := 10) : MetaM (List RewriteResult) := do + let mctx ← getMCtx + let candidates ← rewriteCandidates hyps moduleRef target forbidden + let minHeartbeats : Nat ← + if (← getMaxHeartbeats) = 0 then + pure 0 + else + pure <| leavePercentHeartbeats * (← getRemainingHeartbeats) / 100 + let cfg : RewriteResultConfig := + { stopAtRfl, minHeartbeats, max, mctx, goal, target, side } + return (← takeListAux cfg {} (Array.mkEmpty max) candidates.toList).toList + +end Lean.Meta.Rewrites diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rfl.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rfl.lean new file mode 100644 index 0000000000000000000000000000000000000000..9202eb7e5324992c3174e5d8a103ccdee915b0a5 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Rfl.lean @@ -0,0 +1,145 @@ +/- +Copyright (c) 2022 Newell Jensen. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Newell Jensen, Thomas Murrills, Joachim Breitner +-/ +prelude +import Lean.Meta.Tactic.Apply +import Lean.Elab.Tactic.Basic +import Lean.Meta.Tactic.Refl + +/-! +# `rfl` tactic extension for reflexive relations + +This extends the `rfl` tactic so that it works on any reflexive relation, +provided the reflexivity lemma has been marked as `@[refl]`. +-/ + +namespace Lean.Meta.Rfl + +open Lean Meta + +/-- Environment extensions for `refl` lemmas -/ +initialize reflExt : + SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name) ← + registerSimpleScopedEnvExtension { + addEntry := fun dt (n, ks) => dt.insertCore ks n + initial := {} + } + +/-- +Tags reflexivity lemmas to be used by the `rfl` tactic. + +A reflexivity lemma should have the conclusion `r x x` where `r` is an arbitrary relation. + +It is not possible to tag reflexivity lemmas for `=` using this attribute. These are handled as a +special case in the `rfl` tactic. +-/ +@[builtin_doc] +builtin_initialize registerBuiltinAttribute { + name := `refl + descr := "reflexivity relation" + add := fun decl _ kind => MetaM.run' do + let declTy := (← getConstInfo decl).type + let (_, _, targetTy) ← withReducible <| forallMetaTelescopeReducing declTy + let fail := throwError + "@[refl] attribute only applies to lemmas proving x ∼ x, got {declTy}" + let .app (.app rel lhs) rhs := targetTy | fail + if let .app (.const ``Eq [_]) _ := rel then + throwError "@[refl] attribute may not be used on `Eq.refl`." + unless ← withNewMCtxDepth <| isDefEq lhs rhs do fail + let key ← DiscrTree.mkPath rel + reflExt.add (decl, key) kind +} + +open Elab Tactic + +/-- `MetaM` version of the `rfl` tactic. + +This tactic applies to a goal whose target has the form `x ~ x`, where `~` is a reflexive +relation, that is, equality or another relation which has a reflexive lemma tagged with the +attribute [refl]. +-/ +def _root_.Lean.MVarId.applyRfl (goal : MVarId) : MetaM Unit := goal.withContext do + -- NB: uses whnfR, we do not want to unfold the relation itself + let t ← whnfR <|← instantiateMVars <|← goal.getType + if t.getAppNumArgs < 2 then + throwTacticEx `rfl goal "expected goal to be a binary relation" + + -- Special case HEq here as it has a different argument order. + if t.isAppOfArity ``HEq 4 then + let gs ← goal.applyConst ``HEq.refl + unless gs.isEmpty do + throwTacticEx `rfl goal <| MessageData.tagged `Tactic.unsolvedGoals <| m!"unsolved goals\n{ + goalsToMessageData gs}" + return + + let rel := t.appFn!.appFn! + let lhs := t.appFn!.appArg! + let rhs := t.appArg! + + let success ← approxDefEq <| isDefEqGuarded lhs rhs + unless success do + let explanation := MessageData.ofLazyM (es := #[lhs, rhs]) do + let (lhs, rhs) ← addPPExplicitToExposeDiff lhs rhs + return m!"the left-hand side{indentExpr lhs}\nis not definitionally equal to the right-hand side{indentExpr rhs}" + throwTacticEx `rfl goal explanation + + if rel.isAppOfArity `Eq 1 then + -- The common case is equality: just use `Eq.refl` + let us := rel.appFn!.constLevels! + let α := rel.appArg! + goal.assign (mkApp2 (mkConst ``Eq.refl us) α lhs) + else + -- Else search through `@refl` keyed by the relation + -- We change the type to `lhs ~ lhs` so that we don't do the (possibly costly) `lhs =?= rhs` + -- check again. + goal.setType (.app t.appFn! lhs) + let s ← saveState + let mut ex? := none + for lem in ← (reflExt.getState (← getEnv)).getMatch rel do + try + let gs ← goal.apply (← mkConstWithFreshMVarLevels lem) (term? := m!"'{.ofConstName lem}'") + if gs.isEmpty then return () else + throwError MessageData.tagged `Tactic.unsolvedGoals <| m!"unsolved goals\n{ + goalsToMessageData gs}" + catch e => + unless ex?.isSome do + ex? := some (← saveState, e) -- stash the first failure of `apply` + s.restore + if let some (sErr, e) := ex? then + sErr.restore + throw e + else + throwTacticEx `rfl goal m!"no @[refl] lemma registered for relation{indentExpr rel}" + +/-- Helper theorem for `Lean.MVarId.liftReflToEq`. -/ +private theorem rel_of_eq_and_refl {α : Sort _} {R : α → α → Prop} + {x y : α} (hxy : x = y) (h : R x x) : R x y := + hxy ▸ h + +/-- +Convert a goal of the form `x ~ y` into the form `x = y`, where `~` is a reflexive +relation, that is, a relation which has a reflexive lemma tagged with the attribute `@[refl]`. +If this can't be done, returns the original `MVarId`. +-/ +def _root_.Lean.MVarId.liftReflToEq (mvarId : MVarId) : MetaM MVarId := do + mvarId.checkNotAssigned `liftReflToEq + let .app (.app rel _) _ ← withReducible mvarId.getType' | return mvarId + if rel.isAppOf `Eq then + -- No need to lift Eq to Eq + return mvarId + for lem in ← (reflExt.getState (← getEnv)).getMatch rel do + let res ← observing? do + -- First create an equality relating the LHS and RHS + -- and reduce the goal to proving that LHS is related to LHS. + let [mvarIdEq, mvarIdR] ← mvarId.apply (← mkConstWithFreshMVarLevels ``rel_of_eq_and_refl) + | failure + -- Then fill in the proof of the latter by reflexivity. + let [] ← mvarIdR.apply (← mkConstWithFreshMVarLevels lem) | failure + return mvarIdEq + if let some mvarIdEq := res then + return mvarIdEq + return mvarId + +end Lean.Meta.Rfl diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Simp.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Simp.lean new file mode 100644 index 0000000000000000000000000000000000000000..6d27c31c2f187d4cf0bc2f51a4a618dfe32a522f --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Simp.lean @@ -0,0 +1,34 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Simp.SimpTheorems +import Lean.Meta.Tactic.Simp.SimpCongrTheorems +import Lean.Meta.Tactic.Simp.Types +import Lean.Meta.Tactic.Simp.Main +import Lean.Meta.Tactic.Simp.Rewrite +import Lean.Meta.Tactic.Simp.SimpAll +import Lean.Meta.Tactic.Simp.Simproc +import Lean.Meta.Tactic.Simp.BuiltinSimprocs +import Lean.Meta.Tactic.Simp.RegisterCommand +import Lean.Meta.Tactic.Simp.Attr +import Lean.Meta.Tactic.Simp.Diagnostics +import Lean.Meta.Tactic.Simp.Arith + +namespace Lean + +builtin_initialize registerTraceClass `Meta.Tactic.simp +builtin_initialize registerTraceClass `Meta.Tactic.simp.congr (inherited := true) +builtin_initialize registerTraceClass `Meta.Tactic.simp.discharge (inherited := true) +builtin_initialize registerTraceClass `Meta.Tactic.simp.rewrite (inherited := true) +builtin_initialize registerTraceClass `Meta.Tactic.simp.unify (inherited := true) +builtin_initialize registerTraceClass `Meta.Tactic.simp.ground (inherited := true) +builtin_initialize registerTraceClass `Meta.Tactic.simp.loopProtection (inherited := true) +builtin_initialize registerTraceClass `Meta.Tactic.simp.numSteps +builtin_initialize registerTraceClass `Meta.Tactic.simp.heads +builtin_initialize registerTraceClass `Debug.Meta.Tactic.simp +builtin_initialize registerTraceClass `Debug.Meta.Tactic.simp.congr (inherited := true) + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/SolveByElim.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/SolveByElim.lean new file mode 100644 index 0000000000000000000000000000000000000000..f23a4303dfb1e78c78db1b062072fd8a704431c0 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/SolveByElim.lean @@ -0,0 +1,366 @@ +/- +Copyright (c) 2021 Kim Morrison. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison, David Renshaw +-/ +prelude +import Init.Data.Sum +import Lean.LabelAttribute +import Lean.Meta.Tactic.Apply +import Lean.Meta.Tactic.Backtrack +import Lean.Meta.Tactic.Constructor +import Lean.Meta.Tactic.Repeat +import Lean.Meta.Tactic.Symm +import Lean.Elab.Term + +/-! +# `solve_by_elim`, `apply_rules`, and `apply_assumption`. + +`solve_by_elim` takes a collection of facts from the local context or +supplied as arguments by the user, and performs a backtracking +depth-first search by attempting to `apply` these facts to the goal. + +It is a highly configurable tactic, with options to control the +backtracking, to solve multiple goals simultaneously (with backtracking +between goals), or to supply a discharging tactic for unprovable goals. + +`apply_rules` and `apply_assumption` are much simpler tactics which do +not perform backtracking, but are currently implemented in terms of +`solve_by_elim` with backtracking disabled, in order to be able to share +the front-end customisation and parsing of user options. It would be +reasonable to further separate these in future. +-/ + +open Lean Meta Elab Tactic + +namespace Lean.Meta.SolveByElim + +initialize registerTraceClass `Meta.Tactic.solveByElim + +/-- +`applyTactics lemmas goal` will return an iterator that applies the +lemmas to the goal `goal` and returns ones that succeed. + +Providing this to the `backtracking` tactic, +we can perform backtracking search based on applying a list of lemmas. + +``applyTactics (trace := `name)`` will construct trace nodes for ``name` indicating which +calls to `apply` succeeded or failed. +-/ +def applyTactics (cfg : ApplyConfig := {}) (transparency : TransparencyMode := .default) + (lemmas : List Expr) (g : MVarId) : MetaM (Lean.Meta.Iterator (List Lean.MVarId)) := do + pure <| + (← Meta.Iterator.ofList lemmas).filterMapM (fun e => observing? do + withTraceNode `Meta.Tactic.solveByElim (return m!"{exceptEmoji ·} trying to apply: {e}") do + let goals ← withTransparency transparency (g.apply e cfg) + -- When we call `apply` interactively, `Lean.Elab.Tactic.evalApplyLikeTactic` + -- deals with closing new typeclass goals by calling + -- `Lean.Elab.Term.synthesizeSyntheticMVarsNoPostponing`. + -- It seems we can't reuse that machinery down here in `MetaM`, + -- so we just settle for trying to close each subgoal using `inferInstance`. + goals.filterM fun g => try g.inferInstance; pure false catch _ => pure true) + +/-- +`applyFirst lemmas goal` applies the first of the `lemmas` +which can be successfully applied to `goal`, and fails if none apply. + +We use this in `apply_rules` and `apply_assumption` where backtracking is not needed. +-/ +def applyFirst (cfg : ApplyConfig := {}) (transparency : TransparencyMode := .default) + (lemmas : List Expr) (g : MVarId) : MetaM (List MVarId) := do + (← applyTactics cfg transparency lemmas g).head + +open Lean.Meta.Tactic.Backtrack (BacktrackConfig backtrack) + +/-- The default `maxDepth` for `apply_rules` is higher. -/ +structure ApplyRulesConfig extends BacktrackConfig, ApplyConfig where + /-- Transparency mode for calls to `apply`. -/ + transparency : TransparencyMode := .default + /-- Also use symmetric versions (via `@[symm]`) of local hypotheses. -/ + symm : Bool := true + /-- Try proving the goal via `exfalso` if `solve_by_elim` otherwise fails. + This is only used when operating on a single goal. -/ + exfalso : Bool := true + maxDepth := 50 + +/-- +Configuration structure to control the behaviour of `solve_by_elim`: +* transparency mode for calls to `apply` +* whether to use `symm` on hypotheses and `exfalso` on the goal as needed, +* see also `BacktrackConfig` for hooks allowing flow control. +-/ +structure SolveByElimConfig extends ApplyRulesConfig where + /-- Enable backtracking search. -/ + backtracking : Bool := true + maxDepth := 6 + /-- Trying calling `intro` if no lemmas apply. -/ + intro : Bool := true + /-- Try calling `constructor` if no lemmas apply. -/ + constructor : Bool := true + +namespace SolveByElimConfig + +instance : Coe SolveByElimConfig BacktrackConfig := ⟨(·.toApplyRulesConfig.toBacktrackConfig)⟩ + +/-- Create or modify a `Config` which allows a class of goals to be returned as subgoals. -/ +def accept (cfg : SolveByElimConfig := {}) (test : MVarId → MetaM Bool) : SolveByElimConfig := + { cfg with + discharge := fun g => do + if (← test g) then + pure none + else + cfg.discharge g } + +/-- +Create or modify a `Config` which runs a tactic on the main goal. +If that tactic fails, fall back to the `proc` behaviour of `cfg`. +-/ +def mainGoalProc (cfg : SolveByElimConfig := {}) (proc : MVarId → MetaM (List MVarId)) : SolveByElimConfig := + { cfg with + proc := fun orig goals => match goals with + | [] => cfg.proc orig [] + | g :: gs => try + return (← proc g) ++ gs + catch _ => cfg.proc orig goals } + +/-- Create or modify a `Config` which calls `intro` on each goal before applying lemmas. -/ +-- Because `SolveByElim` works on each goal in sequence, even though +-- `mainGoalProc` only applies this operation on the main goal, +-- it is applied to every goal before lemmas are applied. +def intros (cfg : SolveByElimConfig := {}) : SolveByElimConfig := + cfg.mainGoalProc fun g => do pure [(← g.intro1P).2] + +/-- Attempt typeclass inference on each goal, before applying lemmas. -/ +-- Because `SolveByElim` works on each goal in sequence, even though +-- `mainGoalProc` only applies this operation on the main goal, +-- it is applied to every goal before lemmas are applied. +def synthInstance (cfg : SolveByElimConfig := {}) : SolveByElimConfig := + cfg.mainGoalProc fun g => do + g.assign (← Lean.Meta.synthInstance (← g.getType)) + pure [] + +/-- Add a discharging tactic, falling back to the original discharging tactic if it fails. +Return `none` to return the goal as a new subgoal, or `some goals` to replace it. -/ +def withDischarge (cfg : SolveByElimConfig := {}) (discharge : MVarId → MetaM (Option (List MVarId))) : SolveByElimConfig := + { cfg with + discharge := fun g => try discharge g + catch _ => cfg.discharge g } + +/-- Create or modify a `SolveByElimConfig` which calls `intro` on any goal for which no lemma applies. -/ +def introsAfter (cfg : SolveByElimConfig := {}) : SolveByElimConfig := + cfg.withDischarge fun g => do pure [(← g.intro1P).2] + +/-- Call `constructor` when no lemmas apply. -/ +def constructorAfter (cfg : SolveByElimConfig := {}) : SolveByElimConfig := + cfg.withDischarge fun g => g.constructor {newGoals := .all} + +/-- Create or modify a `Config` which +calls `synthInstance` on any goal for which no lemma applies. -/ +def synthInstanceAfter (cfg : SolveByElimConfig := {}) : SolveByElimConfig := + cfg.withDischarge fun g => do + g.assign (← Lean.Meta.synthInstance (← g.getType)) + pure (some []) + +/-- +Create or modify a `Config` which rejects branches for which `test`, +applied to the instantiations of the original goals, fails or returns `false`. +-/ +def testPartialSolutions (cfg : SolveByElimConfig := {}) (test : List Expr → MetaM Bool) : SolveByElimConfig := + { cfg with + proc := fun orig goals => do + guard <| ← test (← orig.mapM fun m => m.withContext do instantiateMVars (.mvar m)) + cfg.proc orig goals } + +/-- +Create or modify a `SolveByElimConfig` which rejects complete solutions for which `test`, +applied to the instantiations of the original goals, fails or returns `false`. +-/ +def testSolutions (cfg : SolveByElimConfig := {}) (test : List Expr → MetaM Bool) : SolveByElimConfig := + cfg.testPartialSolutions fun sols => do + if sols.any Expr.hasMVar then + pure true + else + test sols + +/-- +Create or modify a `Config` which only accept solutions +for which every expression in `use` appears as a subexpression. +-/ +def requireUsingAll (cfg : SolveByElimConfig := {}) (use : List Expr) : SolveByElimConfig := + cfg.testSolutions fun sols => do + pure <| use.all fun e => sols.any fun s => e.occurs s + +/-- +Process the `intro` and `constructor` options by implementing the `discharger` tactic. +-/ +def processOptions (cfg : SolveByElimConfig) : SolveByElimConfig := + let cfg := if cfg.intro then introsAfter { cfg with intro := false } else cfg + let cfg := if cfg.constructor then constructorAfter { cfg with constructor := false } else cfg + cfg + +end SolveByElimConfig + +/-- +Elaborate a list of lemmas and local context. +See `mkAssumptionSet` for an explanation of why this is needed. +-/ +def elabContextLemmas (cfg : SolveByElimConfig) (g : MVarId) + (lemmas : List (TermElabM Expr)) (ctx : SolveByElimConfig → TermElabM (List Expr)) : + MetaM (List Expr) := do + g.withContext (Elab.Term.TermElabM.run' do pure ((← ctx cfg) ++ (← lemmas.mapM id))) + +/-- Returns the list of tactics corresponding to applying the available lemmas to the goal. -/ +def applyLemmas (cfg : SolveByElimConfig) + (lemmas : List (TermElabM Expr)) (ctx : SolveByElimConfig → TermElabM (List Expr)) + (g : MVarId) : MetaM (Meta.Iterator (List MVarId)) := do + let es ← elabContextLemmas cfg g lemmas ctx + applyTactics cfg.toApplyConfig cfg.transparency es g + +/-- Applies the first possible lemma to the goal. -/ +def applyFirstLemma (cfg : SolveByElimConfig) + (lemmas : List (TermElabM Expr)) (ctx : SolveByElimConfig → TermElabM (List Expr)) + (g : MVarId) : MetaM (List MVarId) := do + let es ← elabContextLemmas cfg g lemmas ctx + applyFirst cfg.toApplyConfig cfg.transparency es g + +/-- +Solve a collection of goals by repeatedly applying lemmas, backtracking as necessary. + +Arguments: +* `cfg : SolveByElimConfig` additional configuration options + (options for `apply`, maximum depth, and custom flow control) +* `lemmas : List (TermElabM Expr)` lemmas to apply. + These are thunks in `TermElabM` to avoid stuck metavariables. +* `ctx : TermElabM (List Expr)` monadic function returning the local hypotheses to use. +* `goals : List MVarId` the initial list of goals for `solveByElim` + +Returns a list of suspended goals, if it succeeded on all other subgoals. +By default `cfg.suspend` is `false,` `cfg.discharge` fails, and `cfg.failAtMaxDepth` is `true`, +and so the returned list is always empty. +Custom wrappers (e.g. `apply_assumption` and `apply_rules`) may modify this behaviour. +-/ +def solveByElim (cfg : SolveByElimConfig) + (lemmas : List (TermElabM Expr)) (ctx : SolveByElimConfig → TermElabM (List Expr)) + (goals : List MVarId) : MetaM (List MVarId) := do + let cfg := cfg.processOptions + try + run cfg goals + catch e => do + -- Implementation note: as with `cfg.symm`, this is different from the mathlib3 approach, + -- for (not as severe) performance reasons. + match goals, cfg.exfalso with + | [g], true => + withTraceNode `Meta.Tactic.solveByElim + (fun _ => return m!"⏮️ starting over using `exfalso`") do + run cfg [← g.exfalso] + | _, _ => throw e +where + /-- Run either backtracking search, or repeated application, on the list of goals. -/ + run (cfg : SolveByElimConfig) : List MVarId → MetaM (List MVarId) := + if cfg.backtracking then + backtrack cfg `Meta.Tactic.solveByElim (applyLemmas cfg lemmas ctx) + else + Lean.Meta.repeat1' (maxIters := cfg.maxDepth) (applyFirstLemma cfg lemmas ctx) + +/-- +If `symm` is `true`, then adds in symmetric versions of each hypothesis. +-/ +def saturateSymm (symm : Bool) (hyps : List Expr) : MetaM (List Expr) := do + if symm then + let extraHyps ← hyps.filterMapM fun hyp => try some <$> hyp.applySymm catch _ => pure none + return hyps ++ extraHyps + else + return hyps + +/-- +A `MetaM` analogue of the `apply_rules` user tactic. + +We pass the lemmas as `TermElabM Expr` rather than just `Expr`, +so they can be generated fresh for each application, to avoid stuck metavariables. + +By default it uses all local hypotheses, but you can disable this with `only := true`. +If you need to remove particular local hypotheses, call `solveByElim` directly. +-/ +def _root_.Lean.MVarId.applyRules (cfg : SolveByElimConfig) (lemmas : List (TermElabM Expr)) + (only : Bool := false) (g : MVarId) : MetaM (List MVarId) := do + let ctx (cfg : SolveByElimConfig) : TermElabM (List Expr) := + if only then pure [] + else do saturateSymm cfg.symm (← getLocalHyps).toList + solveByElim { cfg with backtracking := false } lemmas ctx [g] + +open Lean.Parser.Tactic + +/-- +`mkAssumptionSet` builds a collection of lemmas for use in +the backtracking search in `solve_by_elim`. + +* By default, it includes all local hypotheses, along with `rfl`, `trivial`, `congrFun` + and `congrArg`. +* The flag `noDefaults` removes these. +* The flag `star` includes all local hypotheses, but not `rfl`, `trivial`, `congrFun`, + or `congrArg`. (It doesn't make sense to use `star` without `noDefaults`.) +* The argument `add` is the list of terms inside the square brackets that did not have `-` + and can be used to add expressions or local hypotheses +* The argument `remove` is the list of terms inside the square brackets that had a `-`, + and can be used to remove local hypotheses. + (It doesn't make sense to remove expressions which are not local hypotheses, + to remove local hypotheses unless `!noDefaults || star`, + and it does not make sense to use `star` unless you remove at least one local hypothesis.) + +`mkAssumptionSet` returns not a `List expr`, but a `List (TermElabM Expr) × TermElabM (List Expr)`. +There are two separate problems that need to be solved. + +### Stuck metavariables + +Lemmas with implicit arguments would be filled in with metavariables if we created the +`Expr` objects immediately, so instead we return thunks that generate the expressions +on demand. This is the first component, with type `List (TermElabM Expr)`. + +As an example, we have `def rfl : ∀ {α : Sort u} {a : α}, a = a`, which on elaboration will become +`@rfl ?m_1 ?m_2`. + +Because `solve_by_elim` works by repeated application of lemmas against subgoals, +the first time such a lemma is successfully applied, +those metavariables will be unified, and thereafter have fixed values. +This would make it impossible to apply the lemma +a second time with different values of the metavariables. + +See https://github.com/leanprover-community/mathlib/issues/2269 + +### Relevant local hypotheses + +`solve_by_elim*` works with multiple goals, +and we need to use separate sets of local hypotheses for each goal. +The second component of the returned value provides these local hypotheses. +(Essentially using `getLocalHyps`, along with some filtering to remove hypotheses +that have been explicitly removed via `only` or `[-h]`.) + +-/ +-- These `TermElabM`s must be run inside a suitable `g.withContext`, +-- usually using `elabContextLemmas`. +def mkAssumptionSet (noDefaults star : Bool) (add remove : List Term) (use : Array Ident) : + MetaM (List (TermElabM Expr) × (SolveByElimConfig → TermElabM (List Expr))) := do + if star && !noDefaults then + throwError "It doesn't make sense to use `*` without `only`." + + let defaults : List (TermElabM Expr) := + [← `(rfl), ← `(trivial), ← `(congrFun), ← `(congrArg)].map elab' + let labelledLemmas := (← use.mapM (Lean.labelled ·.raw.getId)).flatten.toList + |>.map (liftM <| mkConstWithFreshMVarLevels ·) + let lemmas := if noDefaults then + add.map elab' ++ labelledLemmas + else + add.map elab' ++ labelledLemmas ++ defaults + + if !remove.isEmpty && noDefaults && !star then + throwError "It doesn't make sense to remove local hypotheses when using `only` without `*`." + let locals (cfg : SolveByElimConfig) : TermElabM (List Expr) := + if noDefaults && !star then do pure [] + else do saturateSymm cfg.symm <| (← getLocalHyps).toList.removeAll (← remove.mapM elab') + + return (lemmas, locals) +where + /-- Run `elabTerm`. -/ + elab' (t : Term) : TermElabM Expr := Elab.Term.elabTerm t.raw none + +open Syntax diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Split.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Split.lean new file mode 100644 index 0000000000000000000000000000000000000000..bcf442ef978851982916a1bc64ecb6c082d29eb2 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Split.lean @@ -0,0 +1,343 @@ +/- +Copyright (c) 2021 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Match.MatcherApp.Basic +import Lean.Meta.Tactic.Simp.Main +import Lean.Meta.Tactic.SplitIf +import Lean.Meta.Tactic.Apply +import Lean.Meta.Tactic.Generalize + +namespace Lean.Meta +namespace Split + +def getSimpMatchContext : MetaM Simp.Context := do + Simp.mkContext + (simpTheorems := {}) + (congrTheorems := (← getSimpCongrTheorems)) + (config := { Simp.neutralConfig with dsimp := false, etaStruct := .none, letToHave := true }) + +def simpMatch (e : Expr) : MetaM Simp.Result := do + let discharge? ← SplitIf.mkDischarge? + (·.1) <$> Simp.main e (← getSimpMatchContext) (methods := { pre, discharge? }) +where + pre (e : Expr) : SimpM Simp.Step := do + unless (← isMatcherApp e) do + return Simp.Step.continue + let matcherDeclName := e.getAppFn.constName! + -- First try to reduce matcher + match (← reduceRecMatcher? e) with + | some e' => return Simp.Step.done { expr := e' } + | none => Simp.simpMatchCore matcherDeclName e + +def simpMatchTarget (mvarId : MVarId) : MetaM MVarId := mvarId.withContext do + let target ← instantiateMVars (← mvarId.getType) + let r ← simpMatch target + applySimpResultToTarget mvarId target r + +private def simpMatchCore (matchDeclName : Name) (matchEqDeclName : Name) (e : Expr) : MetaM Simp.Result := do + let discharge? ← SplitIf.mkDischarge? + (·.1) <$> Simp.main e (← getSimpMatchContext) (methods := { pre, discharge? }) +where + pre (e : Expr) : SimpM Simp.Step := do + if e.isAppOf matchDeclName then + -- First try to reduce matcher + match (← reduceRecMatcher? e) with + | some e' => return .done { expr := e' } + | none => + -- Try lemma + let simpTheorem := { + origin := .decl matchEqDeclName + proof := mkConst matchEqDeclName + rfl := (← isRflTheorem matchEqDeclName) + } + match (← withReducible <| Simp.tryTheorem? e simpTheorem) with + | none => return .continue + | some r => return .done r + else + return .continue + +private def simpMatchTargetCore (mvarId : MVarId) (matchDeclName : Name) (matchEqDeclName : Name) : MetaM MVarId := do + mvarId.withContext do + let target ← instantiateMVars (← mvarId.getType) + let r ← simpMatchCore matchDeclName matchEqDeclName target + match r.proof? with + | some proof => mvarId.replaceTargetEq r.expr proof + | none => mvarId.replaceTargetDefEq r.expr + +private partial def withEqs (lhs rhs : Array Expr) (k : Array Expr → Array Expr → MetaM α) : MetaM α := do + go 0 #[] #[] +where + go (i : Nat) (hs : Array Expr) (rfls : Array Expr) : MetaM α := do + if i < lhs.size then + withLocalDeclD (← mkFreshUserName `heq) (← mkEqHEq lhs[i]! rhs[i]!) fun h => do + let rfl ← if (← inferType h).isEq then mkEqRefl lhs[i]! else mkHEqRefl lhs[i]! + go (i+1) (hs.push h) (rfls.push rfl) + else + k hs rfls + +/-- +Internal exception for discriminant generalization failures due to type errors. +-/ +builtin_initialize discrGenExId : InternalExceptionId ← + registerInternalExceptionId `discrGeneralizationFailure + +def isDiscrGenException (ex : Exception) : Bool := + match ex with + | .internal id => id == discrGenExId + | _ => false + +/-- + This method makes sure each discriminant is a free variable. + Return the tuple `(discrsNew, discrEqs, mvarId)`. `discrsNew` in an array representing the new discriminants, `discrEqs` is an array of auxiliary equality hypotheses + that connect the new discriminants to the original terms they represent. + Remark: `discrEqs.size ≤ discrsNew.size` + + Remark: + We should only generalize `discrs` occurrences as `match`-expression discriminants. + For example, given the following goal. + ``` + x : Nat + ⊢ (match g x with + | 0 => 1 + | Nat.succ y => g x) = + 2 * x + 1 + ``` + we should not generalize the `g x` in the rhs of the second alternative, and the two resulting goals + for the `split` tactic should be + ``` + case h_1 + x x✝ : Nat + h✝ : g x = 0 + ⊢ 1 = 2 * x + 1 + + case h_2 + x x✝ y✝ : Nat + h✝ : g x = Nat.succ y✝ + ⊢ g x = 2 * x + 1 + ``` +-/ +private partial def generalizeMatchDiscrs (mvarId : MVarId) (matcherDeclName : Name) (motiveType : Expr) (discrs : Array Expr) : MetaM (Array FVarId × Array FVarId × MVarId) := mvarId.withContext do + if discrs.all (·.isFVar) then + return (discrs.map (·.fvarId!), #[], mvarId) + let some matcherInfo ← getMatcherInfo? matcherDeclName | unreachable! + let numDiscrEqs := matcherInfo.getNumDiscrEqs -- Number of `h : discr = pattern` equations + let (targetNew, rfls) ← + forallTelescope motiveType fun discrVars _ => + withEqs discrs discrVars fun eqs rfls => do + let foundRef ← IO.mkRef false + let rec mkNewTarget (e : Expr) : MetaM Expr := do + let pre (e : Expr) : MetaM TransformStep := do + if !e.isAppOfArity matcherDeclName matcherInfo.arity then + return .continue + let some matcherApp ← matchMatcherApp? e | return .continue + for matcherDiscr in matcherApp.discrs, discr in discrs do + unless matcherDiscr == discr do + trace[split.debug] "discr mismatch {matcherDiscr} != {discr}" + return .continue + let matcherApp := { matcherApp with discrs := discrVars } + foundRef.set true + let mut altsNew := #[] + for h : i in [:matcherApp.alts.size] do + let alt := matcherApp.alts[i] + let altNumParams := matcherApp.altNumParams[i]! + let altNew ← lambdaTelescope alt fun xs body => do + if xs.size < altNumParams || xs.size < numDiscrEqs then + throwError "internal error in `split` tactic: encountered an unexpected `match` expression alternative\nthis error typically occurs when the `match` expression has been constructed using meta-programming." + let body ← mkLambdaFVars xs[altNumParams...*] (← mkNewTarget body) + let ys := xs[*...(altNumParams - numDiscrEqs)] + if numDiscrEqs == 0 then + mkLambdaFVars ys body + else + let altEqs := xs[(altNumParams - numDiscrEqs)...altNumParams] + withNewAltEqs matcherInfo eqs altEqs fun altEqsNew subst => do + let body := body.replaceFVars altEqs subst + mkLambdaFVars (ys++altEqsNew) body + altsNew := altsNew.push altNew + return .done { matcherApp with alts := altsNew }.toExpr + transform (← instantiateMVars e) pre + let targetNew ← mkNewTarget (← mvarId.getType) + unless (← foundRef.get) do + throwError "internal error in `split` tactic: failed to find match-expression discriminants\nthis error typically occurs when the `split` tactic internal functions have been used in a new meta-program" + let targetNew ← mkForallFVars (discrVars ++ eqs) targetNew + unless (← isTypeCorrect targetNew) do + throw <| Exception.internal discrGenExId + return (targetNew, rfls) + let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew (← mvarId.getTag) + trace[split.debug] "targetNew:\n{mvarNew.mvarId!}" + mvarId.assign (mkAppN (mkAppN mvarNew discrs) rfls) + let (discrs', mvarId') ← mvarNew.mvarId!.introNP discrs.size + let (discrEqs, mvarId') ← mvarId'.introNP discrs.size + return (discrs', discrEqs, mvarId') +where + /-- + - `eqs` are free variables `h_eq : discr = discrVar`. `eqs.size == discrs.size` + - `altEqs` are free variables of the form `h_altEq : discr = pattern`. `altEqs.size = numDiscrEqs ≤ discrs.size` + This method executes `k altEqsNew subst` where + - `altEqsNew` are fresh free variables of the form `h_altEqNew : discrVar = pattern` + - `subst` are terms of the form `h_eq.trans h_altEqNew : discr = pattern`. We use `subst` later to replace occurrences of `h_altEq` with `h_eq.trans h_altEqNew`. + -/ + withNewAltEqs (matcherInfo : MatcherInfo) (eqs : Array Expr) (altEqs : Array Expr) (k : Array Expr → Array Expr → MetaM Expr) : MetaM Expr := do + let eqs' := (eqs.zip matcherInfo.discrInfos).filterMap fun (eq, info) => if info.hName?.isNone then none else some eq + -- `eqs'.size == altEqs.size ≤ eqs.size` + let rec go (i : Nat) (altEqsNew : Array Expr) (subst : Array Expr) : MetaM Expr := do + if i < altEqs.size then + let altEqDecl ← getFVarLocalDecl altEqs[i]! + let eq := eqs'[i]! + let eqType ← inferType eq + let altEqType := altEqDecl.type + match eqType.eq?, altEqType.eq? with + | some (_, _, discrVar), some (_, _ /- discr -/, pattern) => + withLocalDeclD altEqDecl.userName (← mkEq discrVar pattern) fun altEqNew => do + go (i+1) (altEqsNew.push altEqNew) (subst.push (← mkEqTrans eq altEqNew)) + | _, _ => + match eqType.heq?, altEqType.heq? with + | some (_, _, _, discrVar), some (_, _ /- discr -/, _, pattern) => + withLocalDeclD altEqDecl.userName (← mkHEq discrVar pattern) fun altEqNew => do + go (i+1) (altEqsNew.push altEqNew) (subst.push (← mkHEqTrans eq altEqNew)) + | _, _ => + throwError "internal error in `split` tactic: encountered unexpected auxiliary equalities created to generalize `match`-expression discriminant\nthis error typically occurs when the `split` tactic internal functions have been used in a new meta-program" + else + k altEqsNew subst + go 0 #[] #[] + +private def substDiscrEqs (mvarId : MVarId) (fvarSubst : FVarSubst) (discrEqs : Array FVarId) : MetaM MVarId := mvarId.withContext do + let mut mvarId := mvarId + let mut fvarSubst := fvarSubst + for fvarId in discrEqs do + if let .fvar fvarId := fvarSubst.apply (mkFVar fvarId) then + let (fvarId, mvarId') ← heqToEq mvarId fvarId + match (← substCore? mvarId' fvarId (symm := false) fvarSubst) with + | some (fvarSubst', mvarId') => mvarId := mvarId'; fvarSubst := fvarSubst' + | none => + match (← substCore? mvarId' fvarId (symm := true) fvarSubst) with + | some (fvarSubst', mvarId') => mvarId := mvarId'; fvarSubst := fvarSubst' + | none => mvarId := mvarId' + return mvarId + +def applyMatchSplitter (mvarId : MVarId) (matcherDeclName : Name) (us : Array Level) (params : Array Expr) (discrs : Array Expr) : MetaM (List MVarId) := do + let some info ← getMatcherInfo? matcherDeclName | throwError "internal error in `split` tactic: `{matcherDeclName}` is not an auxiliary declaration used to encode `match`-expressions\nthis error typically occurs when the `split` tactic internal functions have been used in a new meta-program" + let matchEqns ← Match.getEquationsFor matcherDeclName + -- splitterPre does not have the correct universe elimination level, but this is fine, we only use it to compute the `motiveType`, + -- and we only care about the `motiveType` arguments, and not the resulting `Sort u`. + let splitterPre := mkAppN (mkConst matchEqns.splitterName us.toList) params + let motiveType := (← whnfForall (← inferType splitterPre)).bindingDomain! + trace[split.debug] "applyMatchSplitter\n{mvarId}" + let (discrFVarIds, discrEqs, mvarId) ← generalizeMatchDiscrs mvarId matcherDeclName motiveType discrs + trace[split.debug] "after generalizeMatchDiscrs\n{mvarId}" + let mvarId ← generalizeTargetsEq mvarId motiveType (discrFVarIds.map mkFVar) + mvarId.withContext do trace[split.debug] "discrEqs after generalizeTargetsEq: {discrEqs.map mkFVar}" + trace[split.debug] "after generalize\n{mvarId}" + let numEqs := discrs.size + let (discrFVarIdsNew, mvarId) ← mvarId.introN discrs.size + trace[split.debug] "after introN\n{mvarId}" + let discrsNew := discrFVarIdsNew.map mkFVar + let mvarType ← mvarId.getType + let elimUniv ← mvarId.withContext <| getLevel mvarType + let us ← if let some uElimPos := info.uElimPos? then + pure <| us.set! uElimPos elimUniv + else + unless elimUniv.isZero do + throwError "`split` tactic failed to split a match-expression: the splitter auxiliary theorem `{matchEqns.splitterName}` can only eliminate into `Prop`" + pure us + let splitter := mkAppN (mkConst matchEqns.splitterName us.toList) params + mvarId.withContext do + let motive ← mkLambdaFVars discrsNew mvarType + let splitter := mkAppN (mkApp splitter motive) discrsNew + check splitter + trace[split.debug] "after check splitter" + let mvarIds ← mvarId.applyN splitter matchEqns.size + let (_, mvarIds) ← mvarIds.foldlM (init := (0, [])) fun (i, mvarIds) mvarId => do + let numParams := matchEqns.splitterAltNumParams[i]! + let (_, mvarId) ← mvarId.introN numParams + trace[split.debug] "before unifyEqs\n{mvarId}" + match (← Cases.unifyEqs? (numEqs + info.getNumDiscrEqs) mvarId {}) with + | none => return (i+1, mvarIds) -- case was solved + | some (mvarId, fvarSubst) => + trace[split.debug] "after unifyEqs\n{mvarId}" + let mvarId ← substDiscrEqs mvarId fvarSubst discrEqs + return (i+1, mvarId::mvarIds) + return mvarIds.reverse + +def mkDiscrGenErrorMsg (e : Expr) : MessageData := + m!"`split` tactic failed to generalize discriminant(s) at{indentExpr e}\nresulting expression was not type correct\npossible solution: generalize discriminant(s) manually before using `split`" + +def throwDiscrGenError (e : Expr) : MetaM α := + throwError (mkDiscrGenErrorMsg e) + +def splitMatch (mvarId : MVarId) (e : Expr) : MetaM (List MVarId) := mvarId.withContext do + let some app ← matchMatcherApp? e | throwError "internal error in `split` tactic: match application expected{indentExpr e}\nthis error typically occurs when the `split` tactic internal functions have been used in a new meta-program" + let matchEqns ← Match.getEquationsFor app.matcherName + let mvarIds ← applyMatchSplitter mvarId app.matcherName app.matcherLevels app.params app.discrs + let (_, mvarIds) ← mvarIds.foldlM (init := (0, [])) fun (i, mvarIds) mvarId => do + let mvarId ← simpMatchTargetCore mvarId app.matcherName matchEqns.eqnNames[i]! + return (i+1, mvarId::mvarIds) + return mvarIds.reverse + +end Split + +open Split + +partial def splitTarget? (mvarId : MVarId) (splitIte := true) : MetaM (Option (List MVarId)) := commitWhenSome? do mvarId.withContext do + let target ← instantiateMVars (← mvarId.getType) + let rec go (badCases : ExprSet) : MetaM (Option (List MVarId)) := do + if let some e ← findSplit? target (if splitIte then .both else .match) badCases then + if e.isIte || e.isDIte then + return (← splitIfTarget? mvarId).map fun (s₁, s₂) => [s₁.mvarId, s₂.mvarId] + else + try + splitMatch mvarId e + catch ex => + if isDiscrGenException ex then + trace[split.failure] mkDiscrGenErrorMsg e + else + trace[split.failure] "`split` tactic failed at{indentExpr e}\n{ex.toMessageData}" + go (badCases.insert e) + else + trace[split.debug] "did not find term to split\n{MessageData.ofGoal mvarId}" + return none + go {} + +def splitLocalDecl? (mvarId : MVarId) (fvarId : FVarId) : MetaM (Option (List MVarId)) := commitWhenSome? do + mvarId.withContext do + if let some e ← findSplit? (← instantiateMVars (← inferType (mkFVar fvarId))) then + if e.isIte || e.isDIte then + return (← splitIfLocalDecl? mvarId fvarId).map fun (mvarId₁, mvarId₂) => [mvarId₁, mvarId₂] + else + let result? ← commitWhenSome? do try + let (fvarIds, mvarId) ← mvarId.revert #[fvarId] + let num := fvarIds.size + let mvarIds ← splitMatch mvarId e + let mvarIds ← mvarIds.mapM fun mvarId => return (← mvarId.introNP num).2 + return some mvarIds + catch ex => + if isDiscrGenException ex then + return none + else + throw ex + if result?.isSome then + return result? + -- Generalization failed, if `fvarId` is a let-decl or has forward dependencies, we try to `assert` a copy and try again + let localDecl ← fvarId.getDecl + if (← pure localDecl.isLet <||> exprDependsOn (← mvarId.getType) fvarId <||> fvarId.hasForwardDeps) then + try + let mvarId ← mvarId.assert localDecl.userName localDecl.type localDecl.toExpr + let mvarIds ← splitMatch mvarId e + let mvarIds ← mvarIds.mapM fun mvarId => return (← mvarId.intro1P).2 + return some mvarIds + catch ex => + if isDiscrGenException ex then + throwDiscrGenError e + else + throw ex + throwDiscrGenError e + else + return none + +builtin_initialize + registerTraceClass `split.debug + registerTraceClass `split.failure + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/SplitIf.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/SplitIf.lean new file mode 100644 index 0000000000000000000000000000000000000000..549155c8d10a408b76b76e561983ff0684f6560a --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/SplitIf.lean @@ -0,0 +1,240 @@ +/- +Copyright (c) 2021 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Cases +import Lean.Meta.Tactic.Simp.Main + +namespace Lean.Meta + +inductive SplitKind where + | ite | match | both + +def SplitKind.considerIte : SplitKind → Bool + | .ite | .both => true + | _ => false + +def SplitKind.considerMatch : SplitKind → Bool + | .match | .both => true + | _ => false + +namespace FindSplitImpl + +structure Context where + exceptionSet : ExprSet := {} + kind : SplitKind := .both + +unsafe abbrev FindM := ReaderT Context $ StateT (PtrSet Expr) MetaM + +/-- +Checks whether `e` is a candidate for `split`. +Returns `some e'` if a prefix is a candidate. +Example: suppose `e` is `(if b then f else g) x`, then +the result is `some e'` where `e'` is the subterm `(if b then f else g)` +-/ +private def isCandidate? (env : Environment) (ctx : Context) (e : Expr) : Option Expr := Id.run do + let ret (e : Expr) : Option Expr := + if ctx.exceptionSet.contains e then none else some e + if ctx.kind.considerIte then + if e.isAppOf ``ite || e.isAppOf ``dite then + let numArgs := e.getAppNumArgs + if numArgs >= 5 && !(e.getArg! 1 5).hasLooseBVars then + return ret (e.getBoundedAppFn (numArgs - 5)) + if ctx.kind.considerMatch then + if let some info := isMatcherAppCore? env e then + let args := e.getAppArgs + for i in [info.getFirstDiscrPos : info.getFirstDiscrPos + info.numDiscrs] do + if args[i]!.hasLooseBVars then + return none + return ret (e.getBoundedAppFn (args.size - info.arity)) + return none + +@[inline] unsafe def checkVisited (e : Expr) : OptionT FindM Unit := do + if (← get).contains e then + failure + modify fun s => s.insert e + +unsafe def visit (e : Expr) : OptionT FindM Expr := do + checkVisited e + if let some e := isCandidate? (← getEnv) (← read) e then + return e + else + -- We do not look for split candidates in proofs. + unless e.hasLooseBVars do + if (← isProof e) then + failure + match e with + | .lam _ _ b _ | .proj _ _ b -- We do not look for split candidates in the binder of lambdas. + | .mdata _ b => visit b + | .forallE _ d b _ => visit d <|> visit b -- We want to look for candidates at `A → B` + | .letE _ _ v b _ => visit v <|> visit b + | .app .. => visitApp? e + | _ => failure +where + visitApp? (e : Expr) : FindM (Option Expr) := + e.withApp fun f args => do + -- See comment at `Canonicalizer.lean` regarding the case where + -- `f` has loose bound variables. + let info ← if f.hasLooseBVars then + pure {} + else + getFunInfo f + for u : i in [0:args.size] do + let arg := args[i] + if h : i < info.paramInfo.size then + let info := info.paramInfo[i] + unless info.isProp do + if info.isExplicit then + let some found ← visit arg | pure () + return found + else + let some found ← visit arg | pure () + return found + visit f + +end FindSplitImpl + +/-- Return an `if-then-else` or `match-expr` to split. -/ +partial def findSplit? (e : Expr) (kind : SplitKind := .both) (exceptionSet : ExprSet := {}) : MetaM (Option Expr) := do + go (← instantiateMVars e) +where + go (e : Expr) : MetaM (Option Expr) := do + if let some target ← find? e then + if target.isIte || target.isDIte then + let cond := target.getArg! 1 5 + -- Try to find a nested `if` in `cond` + return (← go cond).getD target + else + return some target + else + return none + + find? (e : Expr) : MetaM (Option Expr) := do + let some candidate ← unsafe FindSplitImpl.visit e { kind, exceptionSet } |>.run' mkPtrSet + | return none + trace[split.debug] "candidate:{indentExpr candidate}" + return some candidate + +/-- Return the condition and decidable instance of an `if` expression to case split. -/ +private partial def findIfToSplit? (e : Expr) : MetaM (Option (Expr × Expr)) := do + if let some iteApp ← findSplit? e .ite then + let cond := iteApp.getArg! 1 5 + let dec := iteApp.getArg! 2 5 + return (cond, dec) + else + return none + +namespace SplitIf + +/-- + Default `Simp.Context` for `simpIf` methods. It contains all congruence theorems, but + just the rewriting rules for reducing `if` expressions. -/ +def getSimpContext : MetaM Simp.Context := do + let mut s : SimpTheorems := {} + s ← s.addConst ``if_pos + s ← s.addConst ``if_neg + s ← s.addConst ``dif_pos + s ← s.addConst ``dif_neg + Simp.mkContext + (simpTheorems := #[s]) + (congrTheorems := (← getSimpCongrTheorems)) + (config := { Simp.neutralConfig with dsimp := false, letToHave := true }) + +/-- + Default `discharge?` function for `simpIf` methods. + It only uses hypotheses from the local context that have `index < numIndices`. + It is effective after a case-split. + + Remark: when `simp` goes inside binders it adds new local declarations to the + local context. We don't want to use these local declarations since the cached result + would depend on them (see issue #3229). `numIndices` is the size of the local + context `decls` field before we start the simplifying the expression. + + Remark: this function is now private, and we should use `mkDischarge?`. +-/ +private def discharge? (numIndices : Nat) (useDecide : Bool) : Simp.Discharge := fun prop => do + let prop ← instantiateMVars prop + trace[Meta.Tactic.splitIf] "discharge? {prop}, {prop.notNot?}" + if useDecide then + let prop ← instantiateMVars prop + if !prop.hasFVar && !prop.hasMVar then + let d ← mkDecide prop + let r ← withDefault <| whnf d + if r.isConstOf ``true then + return some <| mkApp3 (mkConst ``of_decide_eq_true) prop d.appArg! (← mkEqRefl (mkConst ``true)) + (← getLCtx).findDeclRevM? fun localDecl => do + if localDecl.index ≥ numIndices || localDecl.isAuxDecl then + return none + else if (← isDefEq prop localDecl.type) then + return some localDecl.toExpr + else match prop.notNot? with + | none => return none + | some arg => + if (← isDefEq arg localDecl.type) then + return some (mkApp2 (mkConst ``not_not_intro) arg localDecl.toExpr) + else + return none + +def mkDischarge? (useDecide := false) : MetaM Simp.Discharge := + return discharge? (← getLCtx).numIndices useDecide + +def splitIfAt? (mvarId : MVarId) (e : Expr) (hName? : Option Name) : MetaM (Option (ByCasesSubgoal × ByCasesSubgoal)) := mvarId.withContext do + let e ← instantiateMVars e + if let some (cond, decInst) ← findIfToSplit? e then + let hName ← match hName? with + | none => mkFreshUserName `h + | some hName => pure hName + trace[Meta.Tactic.splitIf] "splitting on {decInst}" + return some (← mvarId.byCasesDec cond decInst hName) + else + trace[Meta.Tactic.splitIf] "could not find if to split:{indentExpr e}" + return none + +end SplitIf + +open SplitIf + +def simpIfTarget (mvarId : MVarId) (useDecide := false) : MetaM MVarId := do + let mut ctx ← getSimpContext + if let (some mvarId', _) ← simpTarget mvarId ctx {} (← mvarId.withContext <| mkDischarge? useDecide) (mayCloseGoal := false) then + return mvarId' + else + unreachable! + +def simpIfLocalDecl (mvarId : MVarId) (fvarId : FVarId) : MetaM MVarId := do + let mut ctx ← getSimpContext + if let (some (_, mvarId'), _) ← simpLocalDecl mvarId fvarId ctx {} (← mvarId.withContext <| mkDischarge?) (mayCloseGoal := false) then + return mvarId' + else + unreachable! + +def splitIfTarget? (mvarId : MVarId) (hName? : Option Name := none) : MetaM (Option (ByCasesSubgoal × ByCasesSubgoal)) := commitWhenSome? do + if let some (s₁, s₂) ← splitIfAt? mvarId (← mvarId.getType) hName? then + let mvarId₁ ← simpIfTarget s₁.mvarId + let mvarId₂ ← simpIfTarget s₂.mvarId + if s₁.mvarId == mvarId₁ && s₂.mvarId == mvarId₂ then + trace[split.failure] "`split` tactic failed to simplify target using new hypotheses Goals:\n{mvarId₁}\n{mvarId₂}" + return none + else + return some ({ s₁ with mvarId := mvarId₁ }, { s₂ with mvarId := mvarId₂ }) + else + return none + +def splitIfLocalDecl? (mvarId : MVarId) (fvarId : FVarId) (hName? : Option Name := none) : MetaM (Option (MVarId × MVarId)) := commitWhenSome? do + mvarId.withContext do + if let some (s₁, s₂) ← splitIfAt? mvarId (← inferType (mkFVar fvarId)) hName? then + let mvarId₁ ← simpIfLocalDecl s₁.mvarId fvarId + let mvarId₂ ← simpIfLocalDecl s₂.mvarId fvarId + if s₁.mvarId == mvarId₁ && s₂.mvarId == mvarId₂ then + trace[split.failure] "`split` tactic failed to simplify target using new hypotheses Goals:\n{mvarId₁}\n{mvarId₂}" + return none + else + return some (mvarId₁, mvarId₂) + else + return none + +builtin_initialize registerTraceClass `Meta.Tactic.splitIf + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Subst.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Subst.lean new file mode 100644 index 0000000000000000000000000000000000000000..7749269dd6e66c559293778c1004c3c149660f5b --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Subst.lean @@ -0,0 +1,243 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.AppBuilder +import Lean.Meta.MatchUtil +import Lean.Meta.Tactic.Util +import Lean.Meta.Tactic.Revert +import Lean.Meta.Tactic.Assert +import Lean.Meta.Tactic.Intro +import Lean.Meta.Tactic.Clear +import Lean.Meta.Tactic.FVarSubst + +namespace Lean.Meta + +def substCore (mvarId : MVarId) (hFVarId : FVarId) (symm := false) (fvarSubst : FVarSubst := {}) (clearH := true) (tryToSkip := false) : MetaM (FVarSubst × MVarId) := + mvarId.withContext do + let tag ← mvarId.getTag + mvarId.checkNotAssigned `subst + let hFVarIdOriginal := hFVarId + let hLocalDecl ← hFVarId.getDecl + match (← matchEq? hLocalDecl.type) with + | none => throwTacticEx `subst mvarId "argument must be an equality proof" + | some (_, lhs, rhs) => do + let a ← instantiateMVars <| if symm then rhs else lhs + let b ← instantiateMVars <| if symm then lhs else rhs + match a with + | Expr.fvar aFVarId => do + let aFVarIdOriginal := aFVarId + trace[Meta.Tactic.subst] "substituting {a} (id: {aFVarId.name}) with {b}" + if (← exprDependsOn b aFVarId) then + throwTacticEx `subst mvarId m!"'{a}' occurs at{indentExpr b}" + let (vars, mvarId) ← mvarId.revert #[aFVarId, hFVarId] true + trace[Meta.Tactic.subst] "after revert {MessageData.ofGoal mvarId}" + let (twoVars, mvarId) ← mvarId.introNP 2 + trace[Meta.Tactic.subst] "after intro2 {MessageData.ofGoal mvarId}" + trace[Meta.Tactic.subst] "reverted variables {vars.map (·.name)}" + let aFVarId := twoVars[0]! + let a := mkFVar aFVarId + let hFVarId := twoVars[1]! + let h := mkFVar hFVarId + /- Set skip to true if there is no local variable nor the target depend on the equality -/ + let skip ← if !tryToSkip || vars.size != 2 then + pure false + else + let mvarType ← mvarId.getType + if (← exprDependsOn mvarType aFVarId) then pure false + else if (← exprDependsOn mvarType hFVarId) then pure false + else pure true + if skip then + if clearH then + let mvarId ← mvarId.clear hFVarId + let mvarId ← mvarId.clear aFVarId + pure (fvarSubst, mvarId) + else + pure (fvarSubst, mvarId) + else + mvarId.withContext do + let mvarDecl ← mvarId.getDecl + let type := mvarDecl.type + let hLocalDecl ← hFVarId.getDecl + match (← matchEq? hLocalDecl.type) with + | none => unreachable! + | some (_, lhs, rhs) => do + let b ← instantiateMVars <| if symm then lhs else rhs + let depElim ← exprDependsOn mvarDecl.type hFVarId + let cont (motive : Expr) (newType : Expr) : MetaM (FVarSubst × MVarId) := do + let major ← if symm then pure h else mkEqSymm h + let newMVar ← mkFreshExprSyntheticOpaqueMVar newType tag + let minor := newMVar + let newVal ← if depElim then mkEqRec motive minor major else mkEqNDRec motive minor major + mvarId.assign newVal + let mvarId := newMVar.mvarId! + let mvarId ← if clearH then + let mvarId ← mvarId.clear hFVarId + mvarId.clear aFVarId + else + pure mvarId + let (newFVars, mvarId) ← mvarId.introNP (vars.size - 2) + trace[Meta.Tactic.subst] "after intro rest {vars.size - 2} {MessageData.ofGoal mvarId}" + let fvarSubst ← newFVars.size.foldM (init := fvarSubst) fun i _ (fvarSubst : FVarSubst) => + let var := vars[i+2]! + let newFVar := newFVars[i] + pure $ fvarSubst.insert var (mkFVar newFVar) + let fvarSubst := fvarSubst.insert aFVarIdOriginal (if clearH then b else mkFVar aFVarId) + let fvarSubst := fvarSubst.insert hFVarIdOriginal (mkFVar hFVarId) + pure (fvarSubst, mvarId) + if depElim then do + let newType := type.replaceFVar a b + let reflB ← mkEqRefl b + let newType := newType.replaceFVar h reflB + if symm then + let motive ← mkLambdaFVars #[a, h] type + cont motive newType + else + /- `type` depends on (h : a = b). So, we use the following trick to avoid a type incorrect motive. + 1- Create a new local (hAux : b = a) + 2- Create newType := type [hAux.symm / h] + `newType` is type correct because `h` and `hAux.symm` are definitionally equal by proof irrelevance. + 3- Create motive by abstracting `a` and `hAux` in `newType`. -/ + let hAuxType ← mkEq b a + let motive ← withLocalDeclD `_h hAuxType fun hAux => do + let hAuxSymm ← mkEqSymm hAux + /- replace h in type with hAuxSymm -/ + let newType := type.replaceFVar h hAuxSymm + mkLambdaFVars #[a, hAux] newType + cont motive newType + else + let motive ← mkLambdaFVars #[a] type + let newType := type.replaceFVar a b + cont motive newType + | _ => + let eqMsg := if symm then "(t = x)" else "(x = t)" + throwTacticEx `subst mvarId + m!"invalid equality proof, it is not of the form {eqMsg}{indentExpr hLocalDecl.type}\nafter WHNF, variable expected, but obtained{indentExpr a}" + +/-- + Given `h : @HEq α a α b` in the given goal, produce a new goal where `h : @Eq α a b`. + If `h` is not of the give form, then just return `(h, mvarId)` +-/ +def heqToEq (mvarId : MVarId) (fvarId : FVarId) (tryToClear : Bool := true) : MetaM (FVarId × MVarId) := + mvarId.withContext do + let decl ← fvarId.getDecl + let type ← whnf decl.type + match type.heq? with + | none => pure (fvarId, mvarId) + | some (α, a, β, b) => + if (← isDefEq α β) then + let pr ← mkEqOfHEq (mkFVar fvarId) + let eq ← mkEq a b + let mut mvarId ← mvarId.assert decl.userName eq pr + if tryToClear then + mvarId ← mvarId.tryClear fvarId + let (fvarId, mvarId') ← mvarId.intro1P + return (fvarId, mvarId') + else + return (fvarId, mvarId) + +/-- +Given `x`, try to find an equation of the form `heq : x = rhs` or `heq : lhs = x`, +and runs `substCore` on it. Throws an exception if no such equation is found. +-/ +partial def substVar (mvarId : MVarId) (x : FVarId) : MetaM MVarId := + mvarId.withContext do + let localDecl ← x.getDecl + if localDecl.isLet then + throwTacticEx `subst mvarId m!"variable '{mkFVar x}' is a let-declaration" + let lctx ← getLCtx + let some (fvarId, symm) ← lctx.findDeclM? fun localDecl => do + if localDecl.isImplementationDetail then + return none + else + match (← matchEq? localDecl.type) with + | some (_, lhs, rhs) => + let lhs ← instantiateMVars lhs + let rhs ← instantiateMVars rhs + if rhs.isFVar && rhs.fvarId! == x then + if !(← exprDependsOn lhs x) then + return some (localDecl.fvarId, true) + if lhs.isFVar && lhs.fvarId! == x then + if !(← exprDependsOn rhs x) then + return some (localDecl.fvarId, false) + return none + | _ => return none + | throwTacticEx `subst mvarId m!"did not find equation for eliminating '{mkFVar x}'" + return (← substCore mvarId fvarId (symm := symm) (tryToSkip := true)).2 + +partial def subst (mvarId : MVarId) (h : FVarId) : MetaM MVarId := + mvarId.withContext do + let type ← h.getType + match (← matchEq? type) with + | some _ => substEq mvarId h + | none => match (← matchHEq? type) with + | some _ => + let (h', mvarId') ← heqToEq mvarId h + if mvarId == mvarId' then + substVar mvarId h + else + subst mvarId' h' + | none => substVar mvarId h +where + /-- Give `h : Eq α a b`, try to apply `substCore` -/ + substEq (mvarId : MVarId) (h : FVarId) : MetaM MVarId := mvarId.withContext do + let localDecl ← h.getDecl + let some (_, lhs, rhs) ← matchEq? localDecl.type | unreachable! + let substReduced (newType : Expr) (symm : Bool) : MetaM MVarId := do + let mvarId ← mvarId.assert localDecl.userName newType (mkFVar h) + let (hFVarId', mvarId) ← mvarId.intro1P + let mvarId ← mvarId.clear h + return (← substCore mvarId hFVarId' (symm := symm) (tryToSkip := true)).2 + let rhs' ← whnf rhs + if rhs'.isFVar then + if rhs != rhs' then + substReduced (← mkEq lhs rhs') true + else + return (← substCore mvarId h (symm := true) (tryToSkip := true)).2 + else do + let lhs' ← whnf lhs + if lhs'.isFVar then + if lhs != lhs' then + substReduced (← mkEq lhs' rhs) false + else + return (← substCore mvarId h (symm := false) (tryToSkip := true)).2 + else do + throwTacticEx `subst mvarId m!"invalid equality proof, it is not of the form (x = t) or (t = x){indentExpr localDecl.type}" + +/-- +Given `x`, try to find an equation of the form `heq : x = rhs` or `heq : lhs = x`, +and runs `substCore` on it. Returns `none` if no such equation is found, or if `substCore` fails. +-/ +def substVar? (mvarId : MVarId) (hFVarId : FVarId) : MetaM (Option MVarId) := + observing? (substVar mvarId hFVarId) + +def subst? (mvarId : MVarId) (hFVarId : FVarId) : MetaM (Option MVarId) := + observing? (subst mvarId hFVarId) + +def substCore? (mvarId : MVarId) (hFVarId : FVarId) (symm := false) (fvarSubst : FVarSubst := {}) (clearH := true) (tryToSkip := false) : MetaM (Option (FVarSubst × MVarId)) := + observing? (substCore mvarId hFVarId symm fvarSubst clearH tryToSkip) + +def trySubstVar (mvarId : MVarId) (hFVarId : FVarId) : MetaM MVarId := do + return (← substVar? mvarId hFVarId).getD mvarId + +def trySubst (mvarId : MVarId) (hFVarId : FVarId) : MetaM MVarId := do + return (← subst? mvarId hFVarId).getD mvarId + +def substSomeVar? (mvarId : MVarId) : MetaM (Option MVarId) := mvarId.withContext do + for localDecl in (← getLCtx) do + if let some mvarId ← subst? mvarId localDecl.fvarId then + return some mvarId + return none + +partial def substVars (mvarId : MVarId) : MetaM MVarId := do + if let some mvarId ← substSomeVar? mvarId then + substVars mvarId + else + return mvarId + +builtin_initialize registerTraceClass `Meta.Tactic.subst + +end Meta +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Symm.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Symm.lean new file mode 100644 index 0000000000000000000000000000000000000000..122a6f63b16d107055ae43de4860ba9afd803756 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Symm.lean @@ -0,0 +1,118 @@ +/- +Copyright (c) 2022 Siddhartha Gadgil. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Siddhartha Gadgil, Mario Carneiro, Kim Morrison +-/ +prelude +import Lean.Meta.Reduce +import Lean.Meta.Tactic.Assert + +/-! +# `symm` tactic + +This implements the `symm` tactic, which can apply symmetry theorems to either the goal or a +hypothesis. +-/ + +open Lean Meta + +namespace Lean.Meta.Symm + +/-- Environment extensions for symm lemmas -/ +builtin_initialize symmExt : + SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name) ← + registerSimpleScopedEnvExtension { + addEntry := fun dt (n, ks) => dt.insertCore ks n + initial := {} + } + +/-- +Tags symmetry lemmas to be used by the `symm` tactic. + +A symmetry lemma should be of the form `r x y → r y x` where `r` is an arbitrary relation. +-/ +@[builtin_doc] +builtin_initialize registerBuiltinAttribute { + name := `symm + descr := "symmetric relation" + add := fun decl _ kind => MetaM.run' do + let declTy := (← getConstInfo decl).type + let (xs, _, targetTy) ← withReducible <| forallMetaTelescopeReducing declTy + let fail := throwError + "@[symm] attribute only applies to lemmas proving x ∼ y → y ∼ x, got {declTy}" + let some _ := xs.back? | fail + let targetTy ← reduce targetTy + let .app (.app rel _) _ := targetTy | fail + let key ← withReducible <| DiscrTree.mkPath rel + symmExt.add (decl, key) kind +} + +end Lean.Meta.Symm + +open Lean.Meta.Symm + +namespace Lean.Expr + +/-- Return the symmetry lemmas that match the target type. -/ +def getSymmLems (tgt : Expr) : MetaM (Array Name) := do + let .app (.app rel _) _ := tgt + | throwError "symmetry lemmas only apply to binary relations, not{indentExpr tgt}" + (symmExt.getState (← getEnv)).getMatch rel + +/-- Given a term `e : a ~ b`, construct a term in `b ~ a` using `@[symm]` lemmas. -/ +def applySymm (e : Expr) : MetaM Expr := do + let tgt <- instantiateMVars (← inferType e) + let lems ← getSymmLems tgt + let s ← saveState + let act lem := do + restoreState s + let lem ← mkConstWithFreshMVarLevels lem + let (args, _, body) ← withReducible <| forallMetaTelescopeReducing (← inferType lem) + let .true ← isDefEq args.back! e | failure + mkExpectedTypeHint (mkAppN lem args) (← instantiateMVars body) + lems.toList.firstM act + <|> throwError m!"no applicable symmetry lemma found for {indentExpr tgt}" + +end Lean.Expr + +namespace Lean.MVarId + +/-- +Apply a symmetry lemma (i.e. marked with `@[symm]`) to a metavariable. + +The type of `g` should be of the form `a ~ b`, and is used to index the symm lemmas. +-/ +def applySymm (g : MVarId) : MetaM MVarId := do + let tgt := (← instantiateMVars (← g.getType)).cleanupAnnotations + let lems ← Expr.getSymmLems tgt + let act lem : MetaM MVarId := do + let lem ← mkConstWithFreshMVarLevels lem + let (args, _, body) ← withReducible <| forallMetaTelescopeReducing (← inferType lem) + let .true ← isDefEq (← g.getType) body | failure + g.assign (mkAppN lem args) + let g' := args.back!.mvarId! + g'.setTag (← g.getTag) + pure g' + lems.toList.firstM act + <|> throwError m!"no applicable symmetry lemma found for {indentExpr tgt}" + +/-- Use a symmetry lemma (i.e. marked with `@[symm]`) to replace a hypothesis in a goal. -/ +def applySymmAt (h : FVarId) (g : MVarId) : MetaM MVarId := do + let h' ← (Expr.fvar h).applySymm + pure (← g.replace h h').mvarId + +/-- For every hypothesis `h : a ~ b` where a `@[symm]` lemma is available, +add a hypothesis `h_symm : b ~ a`. -/ +def symmSaturate (g : MVarId) : MetaM MVarId := g.withContext do + let mut g' := g + let hyps ← getLocalHyps + let types ← hyps.mapM inferType + for h in hyps do try + let symm ← h.applySymm + let symmType ← inferType symm + if ¬ (← types.anyM (isDefEq symmType)) then + (_, g') ← g'.note ((← h.fvarId!.getUserName).appendAfter "_symm") symm + catch _ => g' ← pure g' + return g' + +end Lean.MVarId diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Try.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Try.lean new file mode 100644 index 0000000000000000000000000000000000000000..46f20874cbb798e41ebf8d0b13e7f0fcfccd627c --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Try.lean @@ -0,0 +1,19 @@ +/- +Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Try.Collect + +namespace Lean + +builtin_initialize registerTraceClass `try +builtin_initialize registerTraceClass `try.collect +builtin_initialize registerTraceClass `try.collect.funInd + +builtin_initialize registerTraceClass `try.debug +builtin_initialize registerTraceClass `try.debug.funInd +builtin_initialize registerTraceClass `try.debug.chain + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/TryThis.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/TryThis.lean new file mode 100644 index 0000000000000000000000000000000000000000..88a0708125f368c6fd49a5c2027710ad0ea0702e --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/TryThis.lean @@ -0,0 +1,513 @@ +/- +Copyright (c) 2021 Gabriel Ebner. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Gabriel Ebner, Mario Carneiro, Thomas Murrills +-/ +prelude +import Lean.Server.CodeActions +import Lean.Widget.UserWidget +import Lean.Data.Json.Elab +import Lean.Data.Lsp.Utf16 +import Lean.Meta.CollectFVars +import Lean.Meta.Tactic.ExposeNames +import Lean.Meta.TryThis +import Lean.Meta.Hint + +/-! +# "Try this" code action and tactic suggestions + +This implements a mechanism for tactics to print a message saying `Try this: `, +where `` is a link to a replacement tactic. Users can either click on the link +in the suggestion (provided by a widget), or use a code action which applies the suggestion. +-/ +namespace Lean.Meta.Tactic.TryThis + +open Lean Elab Tactic PrettyPrinter Meta Server RequestM + +/-! # Raw widget -/ + +/-- +This is a widget which is placed by `TryThis.addSuggestion` and `TryThis.addSuggestions`. + +When placed by `addSuggestion`, it says `Try this: ` +where `` is a link which will perform the replacement. + +When placed by `addSuggestions`, it says: +``` +Try these: +``` +* `` +* `` +* `` +* ... + +where `` is a link which will perform the replacement. +-/ +@[builtin_widget_module] def tryThisWidget : Widget.Module where + javascript := " +import * as React from 'react'; +import { EditorContext, EnvPosContext } from '@leanprover/infoview'; +const e = React.createElement; +export default function ({ suggestions, range, header, isInline, style }) { + const pos = React.useContext(EnvPosContext) + const editorConnection = React.useContext(EditorContext) + const defStyle = style || { + className: 'link pointer dim', + style: { color: 'var(--vscode-textLink-foreground)' } + } + + // Construct the children of the HTML element for a given suggestion. + function makeSuggestion({ suggestion, preInfo, postInfo, style }) { + function onClick() { + editorConnection.api.applyEdit({ + changes: { [pos.uri]: [{ range, newText: suggestion }] } + }) + } + return [ + preInfo, + e('span', { onClick, title: 'Apply suggestion', ...style || defStyle }, suggestion), + postInfo + ] + } + + // Choose between an inline 'Try this'-like display and a list-based 'Try these'-like display. + let inner = null + if (isInline) { + inner = e('div', { className: 'ml1' }, + e('pre', { className: 'font-code pre-wrap' }, header, makeSuggestion(suggestions[0]))) + } else { + inner = e('div', { className: 'ml1' }, + e('pre', { className: 'font-code pre-wrap' }, header), + e('ul', { style: { paddingInlineStart: '20px' } }, suggestions.map(s => + e('li', { className: 'font-code pre-wrap' }, makeSuggestion(s))))) + } + return e('details', { open: true }, + e('summary', { className: 'mv2 pointer' }, 'Suggestions'), + inner) +}" + +-- Because we can't use `builtin_widget_module` in `Lean.Meta.Hint`, we add the attribute here +attribute [builtin_widget_module] Hint.tryThisDiffWidget + +/-! # Code action -/ + +/-- +This is a code action provider that looks for `TryThisInfo` nodes and supplies a code action to +apply the replacement. +-/ +@[builtin_code_action_provider] def tryThisProvider : CodeActionProvider := fun params snap => do + let doc ← readDoc + pure <| snap.infoTree.foldInfo (init := #[]) fun _ctx info result => Id.run do + let .ofCustomInfo { stx, value } := info | result + let some { range, suggestionTexts, codeActionPrefix? } := + value.get? TryThisInfo | result + let some stxRange := stx.getRange? | result + let stxRange := doc.meta.text.utf8RangeToLspRange stxRange + unless stxRange.start.line ≤ params.range.end.line do return result + unless params.range.start.line ≤ stxRange.end.line do return result + let mut result := result + for h : i in [:suggestionTexts.size] do + let (newText, title?) := suggestionTexts[i] + let title := title?.getD <| (codeActionPrefix?.getD "Try this: ") ++ newText + result := result.push { + eager.title := title + eager.kind? := "quickfix" + -- Only make the first option preferred + eager.isPreferred? := if i = 0 then true else none + eager.edit? := some <| .ofTextEdit doc.versionedIdentifier { range, newText } + } + result + +/-! # Formatting -/ + +/-- Delaborate `e` into syntax suitable for use by `refine`. -/ +def delabToRefinableSyntax (e : Expr) : MetaM Term := + withOptions (pp.mvars.anonymous.set · false) do delab e + +/-- Delaborate `e` into a suggestion suitable for use by `refine`. -/ +def delabToRefinableSuggestion (e : Expr) : MetaM Suggestion := + return { suggestion := ← delabToRefinableSyntax e, messageData? := e } + +/-- Core of `addSuggestion` and `addSuggestions`. Whether we use an inline display for a single +element or a list display is controlled by `isInline`. -/ +private def addSuggestionCore (ref : Syntax) (suggestions : Array Suggestion) + (header : String) (isInline : Bool) (origSpan? : Option Syntax := none) + (style? : Option SuggestionStyle := none) + (codeActionPrefix? : Option String := none) : CoreM Unit := do + if let some range := (origSpan?.getD ref).getRange? then + let { suggestions, info, range } ← processSuggestions ref range suggestions codeActionPrefix? + let suggestions := suggestions.map (·.1) + let json := json% { + suggestions: $suggestions, + range: $range, + header: $header, + isInline: $isInline, + style: $style? + } + pushInfoLeaf info + Widget.savePanelWidgetInfo tryThisWidget.javascriptHash ref (props := return json) + +/-- Add a "try this" suggestion. This has three effects: + +* An info diagnostic is displayed saying `Try this: ` +* A widget is registered, saying `Try this: ` with a link on `` to apply + the suggestion +* A code action is added, which will apply the suggestion. + +The parameters are: +* `ref`: the span of the info diagnostic +* `s`: a `Suggestion`, which contains + * `suggestion`: the replacement text; + * `preInfo?`: an optional string shown immediately after the replacement text in the widget + message (only) + * `postInfo?`: an optional string shown immediately after the replacement text in the widget + message (only) + * `style?`: an optional `Json` object used as the value of the `style` attribute of the + suggestion text's element (not the whole suggestion element). + * `messageData?`: an optional message to display in place of `suggestion` in the info diagnostic + (only). The widget message uses only `suggestion`. If `messageData?` is `none`, we simply use + `suggestion` instead. + * `toCodeActionTitle?`: an optional function `String → String` describing how to transform the + pretty-printed suggestion text into the code action text which appears in the lightbulb menu. + If `none`, we simply prepend `"Try This: "` to the suggestion text. +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `header`: a string that begins the display. By default, it is `"Try this: "`. +* `codeActionPrefix?`: an optional string to be used as the prefix of the replacement text if the + suggestion does not have a custom `toCodeActionTitle?`. If not provided, `"Try this: "` is used. +-/ +def addSuggestion (ref : Syntax) (s : Suggestion) (origSpan? : Option Syntax := none) + (header : String := "Try this: ") (codeActionPrefix? : Option String := none) : MetaM Unit := do + logInfoAt ref m!"{header}{s}" + addSuggestionCore ref #[s] header (isInline := true) origSpan? + (codeActionPrefix? := codeActionPrefix?) + +/-- Add a list of "try this" suggestions as a single "try these" suggestion. This has three effects: + +* An info diagnostic is displayed saying `Try these: ` +* A widget is registered, saying `Try these: ` with a link on each + `` to apply the suggestion +* A code action for each suggestion is added, which will apply the suggestion. + +The parameters are: +* `ref`: the span of the info diagnostic +* `suggestions`: an array of `Suggestion`s, which each contain + * `suggestion`: the replacement text; + * `preInfo?`: an optional string shown immediately after the replacement text in the widget + message (only) + * `postInfo?`: an optional string shown immediately after the replacement text in the widget + message (only) + * `style?`: an optional `Json` object used as the value of the `style` attribute of the + suggestion text's element (not the whole suggestion element). + * `messageData?`: an optional message to display in place of `suggestion` in the info diagnostic + (only). The widget message uses only `suggestion`. If `messageData?` is `none`, we simply use + `suggestion` instead. + * `toCodeActionTitle?`: an optional function `String → String` describing how to transform the + pretty-printed suggestion text into the code action text which appears in the lightbulb menu. + If `none`, we simply prepend `"Try This: "` to the suggestion text. +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `header`: a string that precedes the list. By default, it is `"Try these:"`. +* `style?`: a default style for all suggestions which do not have a custom `style?` set. +* `codeActionPrefix?`: an optional string to be used as the prefix of the replacement text for all + suggestions which do not have a custom `toCodeActionTitle?`. If not provided, `"Try this: "` is + used. +-/ +def addSuggestions (ref : Syntax) (suggestions : Array Suggestion) + (origSpan? : Option Syntax := none) (header : String := "Try these:") + (style? : Option SuggestionStyle := none) + (codeActionPrefix? : Option String := none) : MetaM Unit := do + if suggestions.isEmpty then throwErrorAt ref "no suggestions available" + let msgs := suggestions.map toMessageData + let msgs := msgs.foldl (init := MessageData.nil) (fun msg m => msg ++ m!"\n• " ++ .nest 2 m) + logInfoAt ref m!"{header}{msgs}" + addSuggestionCore ref suggestions header (isInline := false) origSpan? style? codeActionPrefix? + +/-! # Tactic-specific widget hooks -/ +/-- +Evaluates `tac` in `initialState` without recovery or sorrying on elaboration failure. If +`expectedType?` is non-`none`, an error is thrown if the resulting goal type is not equal to the +provided type (up to `Expr` equality modulo metavariable instantiation). +-/ +private def evalTacticWithState (initialState : Tactic.SavedState) (tac : TSyntax `tactic) + (expectedType? : Option Expr := none) : TacticM Unit := do + let currState ← saveState + initialState.restore + try + Term.withoutErrToSorry <| withoutRecover <| evalTactic tac + if let some expectedType := expectedType? then + let type ← (← getMainGoal).getType + let type ← instantiateMVars type + let expectedType ← instantiateMVars expectedType + if type != expectedType then + throwError "tactic did not produce expected goal" + finally + currState.restore + +/-- +Returns a possibly modified version of `tac` and `msg` that succeeds in `initialState`, prepending +`expose_names` if necessary. If `expectedType?` is non-`none`, the tactic is only considered to have +"succeeded" if the resulting goal is equal (up to `Expr` equality modulo metavariable instantiation) +to the provided type. Returns `none` if the tactic fails even with `expose_names`. + +Remark: We cannot determine if a tactic requires `expose_names` merely by inspecting its syntax +(shadowed variables are delaborated without daggers) nor the underlying `Expr` used to produce it +(some inaccessible names may be implicit arguments that do not appear in the delaborated syntax). +-/ +private def mkValidatedTactic (tac : TSyntax `tactic) (msg : MessageData) + (initialState : Tactic.SavedState) (expectedType? : Option Expr := none) : + TacticM (Option (TSyntax `tactic × MessageData)) := do + try + evalTacticWithState initialState tac expectedType? + return some (tac, msg) + catch _ => + -- Note: we must use `(expose_names; _)` and not `· expose_names; _` to avoid generating + -- spurious tactic-abort exceptions, since these tactics may not close the goal + let tac ← `(tactic| (expose_names; $tac)) + try + evalTacticWithState initialState tac expectedType? + return some (tac, m!"(expose_names; {msg})") + catch _ => + return none + +private def mkFailedToMakeTacticMsg (targetKind : MessageData) (invalidTactic : MessageData) : MessageData := + m!"found {targetKind}, but the corresponding tactic failed:{indentD invalidTactic}\n\n\ + It may be possible to correct this proof by adding type annotations, explicitly specifying \ + implicit arguments, or eliminating unnecessary function abstractions." + +/-- +Returns the syntax for an `exact` or `refine` (as indicated by `useRefine`) tactic corresponding to +`e` as well as a `MessageData` representation with hover information. +If `exposeNames` is `true`, prepends the tactic with `expose_names.` Note that the tactic is +always generated within `withExposedNames` to avoid generating unprintable characters. +-/ +private def mkExactSuggestionSyntax (e : Expr) (useRefine : Bool) : + MetaM (TSyntax `tactic × MessageData) := + withOptions (pp.mvars.set · false) <| withExposedNames do + let exprStx ← delabToRefinableSyntax e + let tac ← if useRefine then `(tactic| refine $exprStx) else `(tactic| exact $exprStx) + -- We must add the message context here to account for exposed names + let exprMessage ← addMessageContext <| MessageData.ofExpr e + let tacMessage := if useRefine then m!"refine {exprMessage}" else m!"exact {exprMessage}" + return (tac, tacMessage) + +private def addExactSuggestionCore (addSubgoalsMsg : Bool) (checkState? : Option Tactic.SavedState) (e : Expr) : + TacticM (Suggestion ⊕ MessageData) := + withOptions (pp.mvars.set · false) do + let mvars ← getMVars e + let hasMVars := !mvars.isEmpty + let (suggestion, messageData) ← mkExactSuggestionSyntax e (useRefine := hasMVars) + let some checkState := checkState? | return .inl suggestion + let some (suggestion, messageData) ← mkValidatedTactic suggestion messageData checkState + | let messageData := m!"(expose_names; {messageData})" + return .inr <| mkFailedToMakeTacticMsg m!"a {if hasMVars then "partial " else ""}proof" messageData + let postInfo? ← if !addSubgoalsMsg || mvars.isEmpty then pure none else + let mut str := "\nRemaining subgoals:" + for g in mvars do + -- TODO: use a MessageData.ofExpr instead of rendering to string + let e ← withExposedNames <| PrettyPrinter.ppExpr (← instantiateMVars (← g.getType)) + str := str ++ Format.pretty ("\n⊢ " ++ e) + pure str + return .inl { suggestion := suggestion, postInfo?, messageData? := messageData } + +/-- Add an `exact e` or `refine e` suggestion. + +The parameters are: +* `ref`: the span of the info diagnostic +* `e`: the replacement expression +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `addSubgoalsMsg`: if true (default false), any remaining subgoals will be shown after + `Remaining subgoals:` +* `codeActionPrefix?`: an optional string to be used as the prefix of the replacement text if the + suggestion does not have a custom `toCodeActionTitle?`. If not provided, `"Try this: "` is used. +* `checkState?`: if passed, the tactic state in which the generated tactic will be validated, + inserting `expose_names` if necessary. +* `tacticErrorAsInfo`: if true (default false), if a generated tactic is invalid (e.g., due to a + pretty-printing issue), the resulting error message will be logged as an info message instead of + being thrown as an error. Has no effect if `checkState?` is `none`. +-/ +def addExactSuggestion (ref : Syntax) (e : Expr) + (origSpan? : Option Syntax := none) + (addSubgoalsMsg := false) + (codeActionPrefix? : Option String := none) + (checkState? : Option Tactic.SavedState := none) + (tacticErrorAsInfo := false) : TacticM Unit := do + match (← addExactSuggestionCore addSubgoalsMsg checkState? e) with + | .inl suggestion => + addSuggestion ref suggestion + (origSpan? := origSpan?) (codeActionPrefix? := codeActionPrefix?) + | .inr message => + if tacticErrorAsInfo then logInfo message else throwError message + +/-- Add `exact e` or `refine e` suggestions if they can be successfully generated; for those that +cannot, display messages indicating the invalid generated tactics. + +The parameters are: +* `ref`: the span of the info diagnostic +* `es`: the array of replacement expressions +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `addSubgoalsMsg`: if true (default false), any remaining subgoals will be shown after + `Remaining subgoals:` +* `codeActionPrefix?`: an optional string to be used as the prefix of the replacement text for all + suggestions which do not have a custom `toCodeActionTitle?`. If not provided, `"Try this: "` is + used. +* `checkState?`: if passed, the tactic state in which the generated tactics will be validated, + inserting `expose_names` if necessary. +* `tacticErrorAsInfo`: if true (default true), invalid generated tactics will log info messages + instead of throwing an error. The default behavior differs from `addExactSuggestion` because + throwing an error means that any subsequent suggestions will not be displayed. Has no effect if + `checkState?` is `none`. +-/ +def addExactSuggestions (ref : Syntax) (es : Array Expr) + (origSpan? : Option Syntax := none) (addSubgoalsMsg := false) + (codeActionPrefix? : Option String := none) + (checkState? : Option Tactic.SavedState := none) + (tacticErrorAsInfo := true) : TacticM Unit := do + let suggestionOrMessages ← es.mapM <| addExactSuggestionCore addSubgoalsMsg checkState? + let mut suggestions : Array Suggestion := #[] + let mut messages : Array MessageData := #[] + for suggestionOrMessage in suggestionOrMessages do + match suggestionOrMessage with + | .inl suggestion => suggestions := suggestions.push suggestion + | .inr message => + unless tacticErrorAsInfo do throwError message + messages := messages.push message + addSuggestions ref suggestions (origSpan? := origSpan?) (codeActionPrefix? := codeActionPrefix?) + for message in messages do + logInfo message + +/-- Add a term suggestion. + +The parameters are: +* `ref`: the span of the info diagnostic +* `e`: the replacement expression +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `header`: a string which precedes the suggestion. By default, it's `"Try this: "`. +* `codeActionPrefix?`: an optional string to be used as the prefix of the replacement text if the + suggestion does not have a custom `toCodeActionTitle?`. If not provided, `"Try this: "` is used. +-/ +def addTermSuggestion (ref : Syntax) (e : Expr) + (origSpan? : Option Syntax := none) (header : String := "Try this: ") + (codeActionPrefix? : Option String := none) : MetaM Unit := do + addSuggestion ref (← delabToRefinableSuggestion e) (origSpan? := origSpan?) (header := header) + (codeActionPrefix? := codeActionPrefix?) + +/-- Add term suggestions. + +The parameters are: +* `ref`: the span of the info diagnostic +* `es`: an array of the replacement expressions +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `header`: a string which precedes the list of suggestions. By default, it's `"Try these:"`. +* `codeActionPrefix?`: an optional string to be used as the prefix of the replacement text for all + suggestions which do not have a custom `toCodeActionTitle?`. If not provided, `"Try this: "` is + used. +-/ +def addTermSuggestions (ref : Syntax) (es : Array Expr) + (origSpan? : Option Syntax := none) (header : String := "Try these:") + (codeActionPrefix? : Option String := none) : MetaM Unit := do + addSuggestions ref (← es.mapM delabToRefinableSuggestion) + (origSpan? := origSpan?) (header := header) (codeActionPrefix? := codeActionPrefix?) + +open Lean Elab Elab.Tactic PrettyPrinter Meta + +/-- Add a suggestion for `have h : t := e`. -/ +def addHaveSuggestion (ref : Syntax) (h? : Option Name) (t? : Option Expr) (e : Expr) + (origSpan? : Option Syntax := none) (checkState? : Option Tactic.SavedState := none) : TacticM Unit := do + let prop ← isProp (← inferType e) + -- We construct the tactic and message data separately to facilitate hover info + let mut (tac, msg) ← withExposedNames do + let estx ← delabToRefinableSyntax e + let (tac, msg) ← if let some t := t? then + let tstx ← delabToRefinableSyntax t + if prop then + match h? with + | some h => pure (← `(tactic| have $(mkIdent h):ident : $tstx := $estx), m!"have {h} : {t} := {e}") + | none => pure (← `(tactic| have : $tstx := $estx), m!"have : {t} := {e}") + else + let h := h?.getD `_ + pure (← `(tactic| let $(mkIdent h):ident : $tstx := $estx), m!"let {h} : {t} := {e}") + else + if prop then + match h? with + | some h => pure (← `(tactic| have $(mkIdent h):ident := $estx), m!"have {h} := {e}") + | none => pure (← `(tactic| have := $estx), m!"have := {e}") + else + let h := h?.getD `_ + pure (← `(tactic| let $(mkIdent h):ident := $estx), m!"let {h} := {e}") + pure (tac, ← addMessageContext msg) + if let some checkState := checkState? then + let some (tac', msg') ← mkValidatedTactic tac msg checkState + | logInfo <| mkFailedToMakeTacticMsg "a proof" msg + return + tac := tac' + msg := msg' + addSuggestion ref (s := { suggestion := tac, messageData? := msg }) origSpan? + +open Lean.Parser.Tactic +open Lean.Syntax + +/-- Add a suggestion for `rw [h₁, ← h₂] at loc`. + +Parameters: +* `ref`: the span of the info diagnostic +* `rules`: a list of arguments to `rw`, with the second component `true` if the rewrite is reversed +* `type?`: the goal after the suggested rewrite, `.none` if the rewrite closes the goal, or `.undef` + if the resulting goal is unknown +* `loc?`: the hypothesis at which the rewrite is performed, or `none` if the goal is targeted +* `origSpan?`: a syntax object whose span is the actual text to be replaced by `suggestion`. + If not provided it defaults to `ref`. +* `checkState?`: if passed, the tactic state in which the generated tactic will be validated, + inserting `expose_names` if necessary +-/ +def addRewriteSuggestion (ref : Syntax) (rules : List (Expr × Bool)) + (type? : LOption Expr := .undef) (loc? : Option Expr := none) + (origSpan? : Option Syntax := none) (checkState? : Option Tactic.SavedState := none) : + TacticM Unit := do + let mut (tac, tacMsg, extraMsg, extraStr) ← withExposedNames do + let rulesStx := TSepArray.ofElems <| ← rules.toArray.mapM fun ⟨e, symm⟩ => do + let t ← delabToRefinableSyntax e + if symm then `(rwRule| ← $t:term) else `(rwRule| $t:term) + -- The seemingly superfluous `:=` below is a workaround for issue #2663 + let mut tac := ← do + let loc ← loc?.mapM fun loc => do `(location| at $(← delab loc):term) + `(tactic| rw [$rulesStx,*] $(loc)?) + + -- We don't simply write `let mut tacMsg := m!"{tac}"` here + -- but instead rebuild it, so that there are embedded `Expr`s in the message, + -- thus giving more information in the hovers. + -- Perhaps in future we will have a better way to attach elaboration information to + -- `Syntax` embedded in a `MessageData`. + let toMessageData (e : Expr) : MessageData := if e.isConst then .ofConst e else .ofExpr e + let rulesMsg := MessageData.sbracket <| MessageData.joinSep + (rules.map fun ⟨e, symm⟩ => (if symm then "← " else "") ++ toMessageData e) ", " + let mut tacMsg ← addMessageContext <| + if let some loc := loc? then + m!"rw {rulesMsg} at {loc}" + else + m!"rw {rulesMsg}" + + let (extraMsg, extraStr) ← + match type? with + | .some type => + pure (← addMessageContext m!"\n-- {type}", s!"\n-- {← PrettyPrinter.ppExpr type}") + | .none => pure (m!"\n-- no goals", "\n-- no goals") + | .undef => pure (m!"", "") + return (tac, tacMsg, extraMsg, extraStr) + + if let some checkState := checkState? then + let type? := match type? with + | .some type => some type + | _ => none + let some (tac', tacMsg') ← mkValidatedTactic tac tacMsg checkState type? + | tacMsg := m!"(expose_names; {tacMsg})" + logInfo <| mkFailedToMakeTacticMsg "an applicable rewrite lemma" (tacMsg ++ extraMsg) + return + tac := tac' + tacMsg := tacMsg' + addSuggestion ref (s := { suggestion := tac, postInfo? := extraStr, messageData? := tacMsg ++ extraMsg }) + origSpan? diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Unfold.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Unfold.lean new file mode 100644 index 0000000000000000000000000000000000000000..c6e63acb6fcd76f8c241feef338cc78201d0eef9 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Unfold.lean @@ -0,0 +1,57 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Eqns +import Lean.Meta.Tactic.Delta +import Lean.Meta.Tactic.Simp.Main + +namespace Lean.Meta + +private def getSimpUnfoldContext : MetaM Simp.Context := do + Simp.mkContext + (congrTheorems := (← getSimpCongrTheorems)) + (config := { Simp.neutralConfig with letToHave := true }) + +def unfold (e : Expr) (declName : Name) : MetaM Simp.Result := do + if let some unfoldThm ← getUnfoldEqnFor? declName then + (·.1) <$> Simp.main e (← getSimpUnfoldContext) (methods := { pre := pre unfoldThm }) + else + return { expr := (← deltaExpand e (· == declName)) } +where + pre (unfoldThm : Name) (e : Expr) : SimpM Simp.Step := do + match (← withReducible <| Simp.tryTheorem? e { origin := .decl unfoldThm, proof := mkConst unfoldThm, rfl := (← isRflTheorem unfoldThm) }) with + | none => pure () + | some r => match (← reduceMatcher? r.expr) with + | .reduced e' => return .done { r with expr := e' } + | _ => return .done r + return .continue + +def unfoldTarget (mvarId : MVarId) (declName : Name) : MetaM MVarId := mvarId.withContext do + let target ← instantiateMVars (← mvarId.getType) + let r ← unfold target declName + if r.expr == target then throwError "tactic 'unfold' failed to unfold '{declName}' at{indentExpr target}" + applySimpResultToTarget mvarId target r + +def unfoldLocalDecl (mvarId : MVarId) (fvarId : FVarId) (declName : Name) : MetaM MVarId := mvarId.withContext do + let type ← fvarId.getType + let r ← unfold (← instantiateMVars type) declName + if r.expr == type then throwError "tactic 'unfold' failed to unfold '{declName}' at{indentExpr type}" + let some (_, mvarId) ← applySimpResultToLocalDecl mvarId fvarId r (mayCloseGoal := false) | unreachable! + return mvarId + +def zetaDeltaTarget (mvarId : MVarId) (declFVarId : FVarId) : MetaM MVarId := mvarId.withContext do + let target ← instantiateMVars (← mvarId.getType) + let target' ← Meta.zetaDeltaFVars target #[declFVarId] + if target' == target then throwError "tactic 'unfold' failed to unfold '{Expr.fvar declFVarId}' at{indentExpr target}" + mvarId.replaceTargetDefEq target' + +def zetaDeltaLocalDecl (mvarId : MVarId) (fvarId : FVarId) (declFVarId : FVarId) : MetaM MVarId := mvarId.withContext do + let type ← fvarId.getType + let type' ← Meta.zetaDeltaFVars (← instantiateMVars type) #[declFVarId] + if type' == type then throwError "tactic 'unfold' failed to unfold '{Expr.fvar fvarId}' at{indentExpr type}" + mvarId.replaceLocalDeclDefEq fvarId type' + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/UnifyEq.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/UnifyEq.lean new file mode 100644 index 0000000000000000000000000000000000000000..f535956c8ff692f78816340d8445e4430a05a66c --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/UnifyEq.lean @@ -0,0 +1,137 @@ +/- +Copyright (c) 2022 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Meta.Tactic.Injection + +namespace Lean.Meta + +/-- Convert heterogeneous equality into a homogeneous one -/ +private def heqToEq' (mvarId : MVarId) (eqDecl : LocalDecl) : MetaM MVarId := do + /- Convert heterogeneous equality into a homogeneous one -/ + let prf ← mkEqOfHEq (mkFVar eqDecl.fvarId) + let aEqb ← whnf (← inferType prf) + let mvarId ← mvarId.assert eqDecl.userName aEqb prf + mvarId.clear eqDecl.fvarId + +structure UnifyEqResult where + mvarId : MVarId + subst : FVarSubst + numNewEqs : Nat := 0 + +private def toOffset? (e : Expr) : MetaM (Option (Expr × Nat)) := do + match (← evalNat e) with + | some k => return some (mkNatLit 0, k) + | none => isOffset? e + +/-- + Helper method for methods such as `Cases.unifyEqs?`. + Given the given goal `mvarId` containing the local hypothesis `eqFVarId`, it performs the following operations: + + - If `eqFVarId` is a heterogeneous equality, tries to convert it to a homogeneous one. + - If `eqFVarId` is a homogeneous equality of the form `a = b`, it tries + - If `a` and `b` are definitionally equal, clear it + - Normalize `a` and `b` using the current reducibility setting. + - If `a` (`b`) is a free variable not occurring in `b` (`a`), replace it everywhere. + - If `a` and `b` are distinct constructors, return `none` to indicate that the goal has been closed. + - If `a` and `b` are the same constructor, apply `injection`, the result contains the number of new equalities introduced in the goal. + - It also tries to apply the given `acyclic` method to try to close the goal. + Remark: It is a parameter because `simp` uses `unifyEq?`, and `acyclic` depends on `simp`. +-/ +def unifyEq? (mvarId : MVarId) (eqFVarId : FVarId) (subst : FVarSubst := {}) + (acyclic : MVarId → Expr → MetaM Bool := fun _ _ => return false) + (caseName? : Option Name := none) + : MetaM (Option UnifyEqResult) := do + mvarId.withContext do + let eqDecl ← eqFVarId.getDecl + if eqDecl.type.isHEq then + let mvarId ← heqToEq' mvarId eqDecl + return some { mvarId, subst, numNewEqs := 1 } + else match eqDecl.type.eq? with + | none => throwError "equality expected{indentExpr eqDecl.type}" + | some (_, a, b) => + /- + Remark: we do not check `isDefeq` here because we would fail to substitute equalities + such as `x = t` and `t = x` when `x` and `t` are proofs (proof irrelanvance). + -/ + /- Remark: we use `let rec` here because otherwise the compiler would generate an insane amount of code. + We can remove the `rec` after we fix the eagerly inlining issue in the compiler. -/ + let rec substEq (symm : Bool) := do + /- + Remark: `substCore` fails if the equation is of the form `x = x` + We use `(tryToSkip := true)` to ensure we avoid unnecessary `Eq.rec`s in user code. + Recall that `Eq.rec`s can negatively affect term reduction performance in the kernel and elaborator. + See issue https://github.com/leanprover/lean4/issues/2552 for an example. + -/ + if let some (subst, mvarId) ← observing? (substCore mvarId eqFVarId symm subst (tryToSkip := true)) then + return some { mvarId, subst } + else if (← isDefEq a b) then + /- Skip equality -/ + return some { mvarId := (← mvarId.clear eqFVarId), subst } + else if (← acyclic mvarId (mkFVar eqFVarId)) then + return none -- this alternative has been solved + else + throwError "dependent elimination failed, failed to solve equation{indentExpr eqDecl.type}" + /- Special support for offset equalities -/ + let injectionOffset? (a b : Expr) := do + unless (← getEnv).contains ``Nat.elimOffset do return none + let some (xa, ka) ← toOffset? a | return none + let some (xb, kb) ← toOffset? b | return none + if ka == 0 || kb == 0 then return none -- use default noConfusion + let (x, y, k) ← if ka < kb then + pure (xa, (← mkAdd xb (mkNatLit (kb - ka))), ka) + else if ka = kb then + pure (xa, xb, ka) + else + pure ((← mkAdd xa (mkNatLit (ka - kb))), xb, kb) + let target ← mvarId.getType + let u ← getLevel target + let newTarget ← mkArrow (← mkEq x y) target + let tag ← mvarId.getTag + let newMVar ← mkFreshExprSyntheticOpaqueMVar newTarget tag + let val := mkAppN (mkConst ``Nat.elimOffset [u]) #[target, x, y, mkNatLit k, eqDecl.toExpr, newMVar] + mvarId.assign val + let mvarId ← newMVar.mvarId!.tryClear eqDecl.fvarId + return some mvarId + let rec injection (a b : Expr) := do + if let some mvarId ← injectionOffset? a b then + return some { mvarId, numNewEqs := 1, subst } + if (← isConstructorApp a <&&> isConstructorApp b) then + /- ctor_i ... = ctor_j ... -/ + match (← injectionCore mvarId eqFVarId) with + | InjectionResultCore.solved => return none -- this alternative has been solved + | InjectionResultCore.subgoal mvarId numNewEqs => return some { mvarId, numNewEqs, subst } + else + let a' ← whnf a + let b' ← whnf b + if a' != a || b' != b then + /- Reduced lhs/rhs of current equality -/ + let prf := mkFVar eqFVarId + let aEqb' ← mkEq a' b' + let mvarId ← mvarId.assert eqDecl.userName aEqb' prf + let mvarId ← mvarId.clear eqFVarId + return some { mvarId, subst, numNewEqs := 1 } + else + match caseName? with + | none => throwError "dependent elimination failed, failed to solve equation{indentExpr eqDecl.type}" + | some caseName => throwError "dependent elimination failed, failed to solve equation{indentExpr eqDecl.type}\nat case {.ofConstName caseName}" + let a ← instantiateMVars a + let b ← instantiateMVars b + match a, b with + | Expr.fvar aFVarId, Expr.fvar bFVarId => + /- x = y -/ + let aDecl ← aFVarId.getDecl + let bDecl ← bFVarId.getDecl + substEq (aDecl.index < bDecl.index) + | Expr.fvar .., _ => /- x = t -/ substEq (symm := false) + | _, Expr.fvar .. => /- t = x -/ substEq (symm := true) + | a, b => + if (← isDefEq a b) then + /- Skip equality -/ + return some { mvarId := (← mvarId.clear eqFVarId), subst } + else + injection a b + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Util.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Util.lean new file mode 100644 index 0000000000000000000000000000000000000000..33ed99b10a230a96421105460c33530552765409 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Meta/Tactic/Util.lean @@ -0,0 +1,169 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Util.ForEachExprWhere +import Lean.Meta.Basic +import Lean.Meta.AppBuilder +import Lean.Meta.PPGoal + +namespace Lean.Meta + +register_builtin_option debug.terminalTacticsAsSorry : Bool := { + defValue := false + group := "debug" + descr := "when enabled, terminal tactics such as `grind` and `omega` are replaced with `sorry`. Useful for debugging and fixing bootstrapping issues" +} + +/-- Get the user name of the given metavariable. -/ +def _root_.Lean.MVarId.getTag (mvarId : MVarId) : MetaM Name := + return (← mvarId.getDecl).userName + +def _root_.Lean.MVarId.setTag (mvarId : MVarId) (tag : Name) : MetaM Unit := do + modify fun s => { s with mctx := s.mctx.setMVarUserName mvarId tag } + +def appendTag (tag : Name) (suffix : Name) : Name := + tag.modifyBase (· ++ suffix.eraseMacroScopes) + +def appendTagSuffix (mvarId : MVarId) (suffix : Name) : MetaM Unit := do + let tag ← mvarId.getTag + mvarId.setTag (appendTag tag suffix) + +def mkFreshExprSyntheticOpaqueMVar (type : Expr) (tag : Name := Name.anonymous) : MetaM Expr := + mkFreshExprMVar type MetavarKind.syntheticOpaque tag + +def throwTacticEx (tacticName : Name) (mvarId : MVarId) (msg? : Option MessageData := none) : MetaM α := + match msg? with + | none => throwError "tactic '{tacticName}' failed\n{mvarId}" + | some msg => throwError "tactic '{tacticName}' failed, {msg}\n{mvarId}" + +/-- +Rethrows the error as a nested error with the given tactic name prepended. +If the error was tagged, prepends `nested` to the tag and preserves it. +-/ +def throwNestedTacticEx {α} (tacticName : Name) (ex : Exception) : MetaM α := do + let nestedMsg := ex.toMessageData + let msg := m!"tactic '{tacticName}' failed, nested error:\n{ex.toMessageData}" + let msg := if let .tagged tag _ := nestedMsg then + .tagged (`nested ++ tag) msg + else msg + throwError msg + +/-- Throw a tactic exception with given tactic name if the given metavariable is assigned. -/ +def _root_.Lean.MVarId.checkNotAssigned (mvarId : MVarId) (tacticName : Name) : MetaM Unit := do + if (← mvarId.isAssigned) then + throwTacticEx tacticName mvarId "metavariable has already been assigned" + +/-- Get the type the given metavariable. -/ +def _root_.Lean.MVarId.getType (mvarId : MVarId) : MetaM Expr := + return (← mvarId.getDecl).type + +/-- Get the type the given metavariable after instantiating metavariables and reducing to +weak head normal form. -/ +-- The `instantiateMVars` needs to be on the outside, +-- since `whnf` can unfold local definitions which may introduce metavariables. +-- We don't need an `instantiateMVars` before the `whnf`, since it instantiates as necessary. +def _root_.Lean.MVarId.getType' (mvarId : MVarId) : MetaM Expr := do + instantiateMVars (← whnf (← mvarId.getType)) + +builtin_initialize registerTraceClass `Meta.Tactic + +/-- Assign `mvarId` to `sorryAx` -/ +def _root_.Lean.MVarId.admit (mvarId : MVarId) (synthetic := true) : MetaM Unit := + mvarId.withContext do + mvarId.checkNotAssigned `admit + let mvarType ← mvarId.getType + let val ← mkLabeledSorry mvarType synthetic (unique := true) + mvarId.assign val + +/-- Beta reduce the metavariable type head -/ +def _root_.Lean.MVarId.headBetaType (mvarId : MVarId) : MetaM Unit := do + mvarId.setType (← mvarId.getType).headBeta + +/-- Collect nondependent hypotheses that are propositions. -/ +def _root_.Lean.MVarId.getNondepPropHyps (mvarId : MVarId) : MetaM (Array FVarId) := + let removeDeps (e : Expr) (candidates : FVarIdHashSet) : MetaM FVarIdHashSet := do + let e ← instantiateMVars e + let visit : StateRefT FVarIdHashSet MetaM FVarIdHashSet := do + e.forEachWhere Expr.isFVar fun e => modify fun s => s.erase e.fvarId! + get + visit |>.run' candidates + mvarId.withContext do + let mut candidates : FVarIdHashSet := {} + for localDecl in (← getLCtx) do + unless localDecl.isImplementationDetail do + candidates ← removeDeps localDecl.type candidates + match localDecl.value? with + | none => pure () + | some value => candidates ← removeDeps value candidates + if (← isProp localDecl.type) && !localDecl.hasValue then + candidates := candidates.insert localDecl.fvarId + candidates ← removeDeps (← mvarId.getType) candidates + if candidates.isEmpty then + return #[] + else + let mut result := #[] + for localDecl in (← getLCtx) do + if candidates.contains localDecl.fvarId then + result := result.push localDecl.fvarId + return result + +partial def saturate (mvarId : MVarId) (x : MVarId → MetaM (Option (List MVarId))) : MetaM (List MVarId) := do + let (_, r) ← go mvarId |>.run #[] + return r.toList +where + go (mvarId : MVarId) : StateRefT (Array MVarId) MetaM Unit := + withIncRecDepth do + match (← x mvarId) with + | none => modify fun s => s.push mvarId + | some mvarIds => mvarIds.forM go + +def exactlyOne (mvarIds : List MVarId) (msg : MessageData := "unexpected number of goals") : MetaM MVarId := + match mvarIds with + | [mvarId] => return mvarId + | _ => throwError msg + +def ensureAtMostOne (mvarIds : List MVarId) (msg : MessageData := "unexpected number of goals") : MetaM (Option MVarId) := + match mvarIds with + | [] => return none + | [mvarId] => return some mvarId + | _ => throwError msg + +/-- Return all propositions in the local context. -/ +def getPropHyps : MetaM (Array FVarId) := do + let mut result := #[] + for localDecl in (← getLCtx) do + unless localDecl.isImplementationDetail do + if (← isProp localDecl.type) then + result := result.push localDecl.fvarId + return result + +def _root_.Lean.MVarId.inferInstance (mvarId : MVarId) : MetaM Unit := mvarId.withContext do + mvarId.checkNotAssigned `infer_instance + let synthVal ← synthInstance (← mvarId.getType) + unless (← isDefEq (mkMVar mvarId) synthVal) do + throwTacticEx `infer_instance mvarId "`infer_instance` tactic failed to assign instance" + +inductive TacticResultCNM where + | closed + | noChange + | modified (mvarId : MVarId) + + +/-- Check if a goal is of a subsingleton type. -/ +def _root_.Lean.MVarId.isSubsingleton (g : MVarId) : MetaM Bool := do + try + discard <| synthInstance (← mkAppM ``Subsingleton #[← g.getType]) + return true + catch _ => + return false + +register_builtin_option tactic.skipAssignedInstances : Bool := { + defValue := true + group := "backward compatibility" + descr := "in the `rw` and `simp` tactics, if an instance implicit argument is assigned, do not try to synthesize instance." +} + +end Lean.Meta diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Modifiers.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Modifiers.lean new file mode 100644 index 0000000000000000000000000000000000000000..e726ac335a42e915f9664581d2a82bf12a9415db --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Modifiers.lean @@ -0,0 +1,30 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.EnvExtension +import Lean.PrivateName + +namespace Lean + +builtin_initialize protectedExt : TagDeclarationExtension ← mkTagDeclarationExtension + +def addProtected (env : Environment) (n : Name) : Environment := + protectedExt.tag env n + +def isProtected (env : Environment) (n : Name) : Bool := + protectedExt.isTagged env n + +def mkPrivateName (env : Environment) (n : Name) : Name := + -- If name is already private, remove previous suffix first. We need to ensure the resulting name + -- is private to *this* module. + mkPrivateNameCore env.mainModule <| privateToUserName n + +def isPrivateNameFromImportedModule (env : Environment) (n : Name) : Bool := + match privateToUserName? n with + | some userName => mkPrivateName env userName != n + | _ => false + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/MonadEnv.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/MonadEnv.lean new file mode 100644 index 0000000000000000000000000000000000000000..344ae70b880b94731d3316682ab4d04eb00c913a --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/MonadEnv.lean @@ -0,0 +1,174 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Environment +import Lean.Exception +import Lean.Declaration +import Lean.Log +import Lean.AuxRecursor +import Lean.Compiler.Old + +namespace Lean + +def setEnv [MonadEnv m] (env : Environment) : m Unit := + modifyEnv fun _ => env + +def withEnv [Monad m] [MonadFinally m] [MonadEnv m] (env : Environment) (x : m α) : m α := do + let saved ← getEnv + try + setEnv env + x + finally + setEnv saved + +def isInductiveCore (env : Environment) (declName : Name) : Bool := + env.findAsync? declName matches some { kind := .induct, .. } + +def isInductive [Monad m] [MonadEnv m] (declName : Name) : m Bool := + return isInductiveCore (← getEnv) declName + +def isRecCore (env : Environment) (declName : Name) : Bool := + env.findAsync? declName matches some { kind := .recursor, .. } + +def isRec [Monad m] [MonadEnv m] (declName : Name) : m Bool := + return isRecCore (← getEnv) declName + +@[inline] def withoutModifyingEnv [Monad m] [MonadEnv m] [MonadFinally m] {α : Type} (x : m α) : m α := do + -- Allow `x` to define new declarations even outside the asynchronous prefix (if any) as all + -- results will be discarded anyway. + withEnv (← getEnv).unlockAsync x + +/-- Similar to `withoutModifyingEnv`, but also returns the updated environment -/ +@[inline] def withoutModifyingEnv' [Monad m] [MonadEnv m] [MonadFinally m] {α : Type} (x : m α) : m (α × Environment) := do + let env ← getEnv + try + let a ← x + return (a, ← getEnv) + finally + setEnv env + +@[inline] def matchConst [Monad m] [MonadEnv m] (e : Expr) (failK : Unit → m α) (k : ConstantInfo → List Level → m α) : m α := do + match e with + | Expr.const constName us => do + match (← getEnv).find? constName with + | some cinfo => k cinfo us + | none => failK () + | _ => failK () + +@[inline] def matchConstInduct [Monad m] [MonadEnv m] (e : Expr) (failK : Unit → m α) (k : InductiveVal → List Level → m α) : m α := + matchConst e failK fun cinfo us => + match cinfo with + | ConstantInfo.inductInfo val => k val us + | _ => failK () + +@[inline] def matchConstCtor [Monad m] [MonadEnv m] (e : Expr) (failK : Unit → m α) (k : ConstructorVal → List Level → m α) : m α := + matchConst e failK fun cinfo us => + match cinfo with + | ConstantInfo.ctorInfo val => k val us + | _ => failK () + +@[inline] def matchConstRec [Monad m] [MonadEnv m] (e : Expr) (failK : Unit → m α) (k : RecursorVal → List Level → m α) : m α := + matchConst e failK fun cinfo us => + match cinfo with + | ConstantInfo.recInfo val => k val us + | _ => failK () + +def hasConst [Monad m] [MonadEnv m] (constName : Name) (skipRealize := true) : m Bool := do + return (← getEnv).contains (skipRealize := skipRealize) constName + +def getConstInfo [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m ConstantInfo := do + match (← getEnv).find? constName with + | some info => pure info + | none => throwError "unknown constant '{.ofConstName constName}'" + +def getConstVal [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m ConstantVal := do + match (← getEnv).findConstVal? constName with + | some val => pure val + | none => throwError "unknown constant '{mkConst constName}'" + +def getAsyncConstInfo [Monad m] [MonadEnv m] [MonadError m] (constName : Name) (skipRealize := false) : m AsyncConstantInfo := do + match (← getEnv).findAsync? (skipRealize := skipRealize) constName with + | some val => pure val + | none => throwError "unknown constant '{mkConst constName}'" + +def mkConstWithLevelParams [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m Expr := do + let info ← getConstVal constName + return mkConst constName (info.levelParams.map mkLevelParam) + +def getConstInfoDefn [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m DefinitionVal := do + match (← getConstInfo constName) with + | ConstantInfo.defnInfo v => pure v + | _ => throwError "'{.ofConstName constName}' is not a definition" + +def getConstInfoInduct [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m InductiveVal := do + match (← getConstInfo constName) with + | ConstantInfo.inductInfo v => pure v + | _ => throwError "'{.ofConstName constName}' is not a inductive type" + +def getConstInfoCtor [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m ConstructorVal := do + match (← getConstInfo constName) with + | ConstantInfo.ctorInfo v => pure v + | _ => throwError "'{.ofConstName constName}' is not a constructor" + +def getConstInfoRec [Monad m] [MonadEnv m] [MonadError m] (constName : Name) : m RecursorVal := do + match (← getConstInfo constName) with + | ConstantInfo.recInfo v => pure v + | _ => throwError "'{.ofConstName constName}' is not a recursor" + +/-- +Matches if `e` is a constant that is an inductive type with one constructor. +Such types can be used with primitive projections. +See also `Lean.matchConstStructLike` for a more restrictive version. +-/ +@[inline] def matchConstStructure [Monad m] [MonadEnv m] [MonadError m] (e : Expr) (failK : Unit → m α) (k : InductiveVal → List Level → ConstructorVal → m α) : m α := + matchConstInduct e failK fun ival us => do + match ival.ctors with + | [ctor] => + match (← getConstInfo ctor) with + | ConstantInfo.ctorInfo cval => k ival us cval + | _ => failK () + | _ => failK () + +/-- +Matches if `e` is a constant that is an non-recursive inductive type with no indices and with one constructor. +Such a type satisfies `Lean.isStructureLike`. +See also `Lean.matchConstStructure` for a less restrictive version. +-/ +@[inline] def matchConstStructureLike [Monad m] [MonadEnv m] [MonadError m] (e : Expr) (failK : Unit → m α) (k : InductiveVal → List Level → ConstructorVal → m α) : m α := + matchConstInduct e failK fun ival us => do + if ival.isRec || ival.numIndices != 0 then failK () + else match ival.ctors with + | [ctor] => + match (← getConstInfo ctor) with + | ConstantInfo.ctorInfo cval => k ival us cval + | _ => failK () + | _ => failK () + +unsafe def evalConst [Monad m] [MonadEnv m] [MonadError m] [MonadOptions m] (α) (constName : Name) (checkMeta := true) : m α := do + ofExcept <| (← getEnv).evalConst (checkMeta := checkMeta) α (← getOptions) constName + +unsafe def evalConstCheck [Monad m] [MonadEnv m] [MonadError m] [MonadOptions m] (α) (typeName : Name) (constName : Name) : m α := do + ofExcept <| (← getEnv).evalConstCheck α (← getOptions) typeName constName + +def findModuleOf? [Monad m] [MonadEnv m] [MonadError m] (declName : Name) : m (Option Name) := do + discard <| getConstInfo declName -- ensure declaration exists + match (← getEnv).getModuleIdxFor? declName with + | none => return none + | some modIdx => return some ((← getEnv).allImportedModuleNames[modIdx.toNat]!) + +def isEnumType [Monad m] [MonadEnv m] [MonadError m] (declName : Name) : m Bool := do + if let ConstantInfo.inductInfo info ← getConstInfo declName then + if !info.type.isProp && info.numTypeFormers == 1 && info.numIndices == 0 && info.numParams == 0 + && !info.ctors.isEmpty && !info.isRec && !info.isUnsafe then + info.ctors.allM fun ctorName => do + let ConstantInfo.ctorInfo info ← getConstInfo ctorName | return false + return info.numFields == 0 + else + return false + else + return false + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Namespace.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Namespace.lean new file mode 100644 index 0000000000000000000000000000000000000000..30d34320276b2e94851d263be0ab93667ba1cef0 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Namespace.lean @@ -0,0 +1,44 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.EnvExtension + +namespace Lean + +/-- +Environment extension for tracking all `namespace` declared by users. +-/ +builtin_initialize namespacesExt : SimplePersistentEnvExtension Name NameSSet ← + registerSimplePersistentEnvExtension { + addImportedFn := fun as => + /- + We compute a `HashMap Name Unit` and then convert to `NameSSet` to improve Lean startup time. + Note: we have used `perf` to profile Lean startup cost when processing a file containing just `import Lean`. + 6.18% of the runtime is here. It was 9.31% before the `HashMap` optimization. + -/ + let capacity := as.foldl (init := 0) fun r e => r + e.size + let map : Std.HashMap Name Unit := Std.HashMap.emptyWithCapacity capacity + let map := mkStateFromImportedEntries (fun map name => map.insert name ()) map as + SMap.fromHashMap map |>.switch + addEntryFn := fun s n => s.insert n + -- Namespaces from local helper constants can be disregarded in other environment branches. We + -- do *not* want `getNamespaceSet` to have to wait on all prior branches. + asyncMode := .local + } + +namespace Environment + +/-- Register a new namespace in the environment. -/ +def registerNamespace (env : Environment) (n : Name) : Environment := + if (namespacesExt.getState env).contains n then env else namespacesExt.addEntry env n + +/-- Return `true` if `n` is the name of a namespace in `env`. -/ +def isNamespace (env : Environment) (n : Name) : Bool := + (namespacesExt.getState env).contains n + +/-- Return a set containing all namespaces in `env`. -/ +def getNamespaceSet (env : Environment) : NameSSet := + namespacesExt.getState env diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Parser.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Parser.lean new file mode 100644 index 0000000000000000000000000000000000000000..bd21b4ad10cdd368b7c11086a3873b2fa1d2f254 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Parser.lean @@ -0,0 +1,127 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura, Sebastian Ullrich +-/ +prelude +import Lean.Parser.Basic +import Lean.Parser.Level +import Lean.Parser.Term +import Lean.Parser.Tactic +import Lean.Parser.Command +import Lean.Parser.Module +import Lean.Parser.Syntax +import Lean.Parser.Do +import Lean.Parser.Tactic.Doc + +namespace Lean +namespace Parser +open Lean.PrettyPrinter +open Lean.PrettyPrinter.Parenthesizer +open Lean.PrettyPrinter.Formatter + +builtin_initialize + register_parser_alias "ws" checkWsBefore { stackSz? := some 0 } + register_parser_alias "noWs" checkNoWsBefore { stackSz? := some 0 } + register_parser_alias "linebreak" checkLinebreakBefore { stackSz? := some 0 } + register_parser_alias (kind := numLitKind) "num" numLit + register_parser_alias (kind := strLitKind) "str" strLit + register_parser_alias (kind := charLitKind) "char" charLit + register_parser_alias (kind := nameLitKind) "name" nameLit + register_parser_alias (kind := scientificLitKind) "scientific" scientificLit + register_parser_alias (kind := identKind) ident + register_parser_alias (kind := hygieneInfoKind) hygieneInfo + register_parser_alias "colGt" checkColGt { stackSz? := some 0 } + register_parser_alias "colGe" checkColGe { stackSz? := some 0 } + register_parser_alias "colEq" checkColEq { stackSz? := some 0 } + register_parser_alias "lineEq" checkLineEq { stackSz? := some 0 } + register_parser_alias lookahead { stackSz? := some 0 } + register_parser_alias atomic { stackSz? := none } + register_parser_alias many + register_parser_alias many1 + register_parser_alias manyIndent + register_parser_alias many1Indent + register_parser_alias optional { autoGroupArgs := false } + register_parser_alias withPosition { stackSz? := none } + register_parser_alias withoutPosition { stackSz? := none } + register_parser_alias withoutForbidden { stackSz? := none } + register_parser_alias (kind := interpolatedStrKind) interpolatedStr + register_parser_alias orelse + register_parser_alias andthen { stackSz? := none } + register_parser_alias recover + + registerAlias `notFollowedBy ``notFollowedBy (notFollowedBy · "element") + Parenthesizer.registerAlias `notFollowedBy notFollowedBy.parenthesizer + Formatter.registerAlias `notFollowedBy notFollowedBy.formatter + +end Parser + +namespace PrettyPrinter +namespace Parenthesizer + +-- Close the mutual recursion loop; see corresponding `[extern]` in the parenthesizer. +@[export lean_mk_antiquot_parenthesizer] +def mkAntiquot.parenthesizer (name : String) (kind : SyntaxNodeKind) (anonymous := true) (isPseudoKind := true) : Parenthesizer := + Parser.mkAntiquot.parenthesizer name kind anonymous isPseudoKind + +-- The parenthesizer auto-generated these instances correctly, but tagged them with the wrong kind, since the actual kind +-- (e.g. `ident`) is not equal to the parser name `Lean.Parser.Term.ident`. +@[builtin_parenthesizer ident] def ident.parenthesizer : Parenthesizer := Parser.Term.ident.parenthesizer +@[builtin_parenthesizer num] def numLit.parenthesizer : Parenthesizer := Parser.Term.num.parenthesizer +@[builtin_parenthesizer scientific] def scientificLit.parenthesizer : Parenthesizer := Parser.Term.scientific.parenthesizer +@[builtin_parenthesizer char] def charLit.parenthesizer : Parenthesizer := Parser.Term.char.parenthesizer +@[builtin_parenthesizer str] def strLit.parenthesizer : Parenthesizer := Parser.Term.str.parenthesizer + +open Lean.Parser + +@[export lean_pretty_printer_parenthesizer_interpret_parser_descr] +unsafe def interpretParserDescr : ParserDescr → CoreM Parenthesizer + | ParserDescr.const n => getConstAlias parenthesizerAliasesRef n + | ParserDescr.unary n d => return (← getUnaryAlias parenthesizerAliasesRef n) (← interpretParserDescr d) + | ParserDescr.binary n d₁ d₂ => return (← getBinaryAlias parenthesizerAliasesRef n) (← interpretParserDescr d₁) (← interpretParserDescr d₂) + | ParserDescr.node k prec d => return leadingNode.parenthesizer k prec (← interpretParserDescr d) + | ParserDescr.nodeWithAntiquot n k d => return withAntiquot.parenthesizer (mkAntiquot.parenthesizer' n k (anonymous := true)) <| + node.parenthesizer k (← interpretParserDescr d) + | ParserDescr.sepBy p sep psep trail => return sepBy.parenthesizer (← interpretParserDescr p) sep (← interpretParserDescr psep) trail + | ParserDescr.sepBy1 p sep psep trail => return sepBy1.parenthesizer (← interpretParserDescr p) sep (← interpretParserDescr psep) trail + | ParserDescr.trailingNode k prec lhsPrec d => return trailingNode.parenthesizer k prec lhsPrec (← interpretParserDescr d) + | ParserDescr.symbol tk => return symbol.parenthesizer tk + | ParserDescr.nonReservedSymbol tk includeIdent => return nonReservedSymbol.parenthesizer tk includeIdent + | ParserDescr.parser constName => combinatorParenthesizerAttribute.runDeclFor constName + | ParserDescr.cat catName prec => return categoryParser.parenthesizer catName prec + +end Parenthesizer + +namespace Formatter + +@[export lean_mk_antiquot_formatter] +def mkAntiquot.formatter (name : String) (kind : SyntaxNodeKind) (anonymous := true) (isPseudoKind := true) : Formatter := + Parser.mkAntiquot.formatter name kind anonymous isPseudoKind + +@[builtin_formatter ident] def ident.formatter : Formatter := Parser.Term.ident.formatter +@[builtin_formatter num] def numLit.formatter : Formatter := Parser.Term.num.formatter +@[builtin_formatter scientific] def scientificLit.formatter : Formatter := Parser.Term.scientific.formatter +@[builtin_formatter char] def charLit.formatter : Formatter := Parser.Term.char.formatter +@[builtin_formatter str] def strLit.formatter : Formatter := Parser.Term.str.formatter + +open Lean.Parser + +@[export lean_pretty_printer_formatter_interpret_parser_descr] +unsafe def interpretParserDescr : ParserDescr → CoreM Formatter + | ParserDescr.const n => getConstAlias formatterAliasesRef n + | ParserDescr.unary n d => return (← getUnaryAlias formatterAliasesRef n) (← interpretParserDescr d) + | ParserDescr.binary n d₁ d₂ => return (← getBinaryAlias formatterAliasesRef n) (← interpretParserDescr d₁) (← interpretParserDescr d₂) + | ParserDescr.node k _ d => return node.formatter k (← interpretParserDescr d) + | ParserDescr.nodeWithAntiquot n k d => return withAntiquot.formatter (mkAntiquot.formatter' n k (anonymous := true)) <| + node.formatter k (← interpretParserDescr d) + | ParserDescr.sepBy p sep psep trail => return sepBy.formatter (← interpretParserDescr p) sep (← interpretParserDescr psep) trail + | ParserDescr.sepBy1 p sep psep trail => return sepBy1.formatter (← interpretParserDescr p) sep (← interpretParserDescr psep) trail + | ParserDescr.trailingNode k prec lhsPrec d => return trailingNode.formatter k prec lhsPrec (← interpretParserDescr d) + | ParserDescr.symbol tk => return symbol.formatter tk + | ParserDescr.nonReservedSymbol tk _ => return nonReservedSymbol.formatter tk + | ParserDescr.parser constName => combinatorFormatterAttribute.runDeclFor constName + | ParserDescr.cat catName _ => return categoryParser.formatter catName + +end Formatter +end PrettyPrinter +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ParserCompiler.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ParserCompiler.lean new file mode 100644 index 0000000000000000000000000000000000000000..5330b8741d458c341d5a2cdb32bfc4a0665af200 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ParserCompiler.lean @@ -0,0 +1,158 @@ +/- +Copyright (c) 2020 Sebastian Ullrich. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Sebastian Ullrich +-/ +prelude +import Lean.Meta.ReduceEval +import Lean.Meta.WHNF +import Lean.KeyedDeclsAttribute +import Lean.ParserCompiler.Attribute +import Lean.Parser.Extension + +/-! +Gadgets for compiling parser declarations into other programs, such as pretty printers. +-/ + +namespace Lean +namespace ParserCompiler + +structure Context (α : Type) where + varName : Name + categoryAttr : KeyedDeclsAttribute α + combinatorAttr : CombinatorAttribute + +def Context.tyName {α} (ctx : Context α) : Name := ctx.categoryAttr.defn.valueTypeName + +/-- Replace all references of `Parser` with `tyName` -/ +def replaceParserTy {α} (ctx : Context α) (e : Expr) : Expr := + e.replace fun e => + -- strip `optParam` + let e := if e.isOptParam then e.appFn!.appArg! else e + if e.isConstOf `Lean.Parser.Parser then mkConst ctx.tyName else none + +open Meta Parser in +/-- Takes an expression of type `Parser`, and determines the syntax kind of the root node it produces. -/ +partial def parserNodeKind? (e : Expr) : MetaM (Option Name) := do + let reduceEval? e : MetaM (Option Name) := do + try pure <| some (← reduceEval e) catch _ => pure none + let e ← whnfCore e + if e matches Expr.lam .. | Expr.letE .. then + lambdaLetTelescope (preserveNondepLet := false) e fun _ e => parserNodeKind? e + else if e.isAppOfArity ``leadingNode 3 || e.isAppOfArity ``trailingNode 4 || e.isAppOfArity ``node 2 then + reduceEval? (e.getArg! 0) + else if e.isAppOfArity ``withAntiquot 2 then + parserNodeKind? (e.getArg! 1) + else forallTelescope (← inferType e.getAppFn) fun params _ => do + let lctx ← getLCtx + -- if there is exactly one parameter of type `Parser`, search there + if let #[(_, i)] := params.zipIdx.filter (lctx.getFVar! ·.1 |>.type.isConstOf ``Parser) then + parserNodeKind? (e.getArg! i) + else + return none + +section +open Meta + +variable {α} (ctx : Context α) (builtin : Bool) (force : Bool) in +/-- + Translate an expression of type `Parser` into one of type `tyName`, tagging intermediary constants with + `ctx.combinatorAttr`. If `force` is `false`, refuse to do so for imported constants. -/ +partial def compileParserExpr (e : Expr) : MetaM Expr := do + let e ← whnfCore e + match e with + | .lam .. | .letE .. => mapLambdaLetTelescope (preserveNondepLet := false) e fun _ b => compileParserExpr b + | .fvar .. => return e + | _ => do + let fn := e.getAppFn + let .const c .. := fn | throwError "call of unknown parser at '{e}'" + -- call the translated `p` with (a prefix of) the arguments of `e`, recursing for arguments + -- of type `ty` (i.e. formerly `Parser`) + let mkCall (p : Name) := do + let ty ← inferType (mkConst p) + forallTelescope ty fun params _ => do + let mut p := mkConst p + let args := e.getAppArgs + for i in [:Nat.min params.size args.size] do + let param := params[i]! + let arg := args[i]! + let paramTy ← inferType param + let resultTy ← forallTelescope paramTy fun _ b => pure b + let arg ← if resultTy.isConstOf ctx.tyName then compileParserExpr arg else pure arg + p := mkApp p arg + return p + let env ← getEnv + match ctx.combinatorAttr.getDeclFor? env c with + | some p => mkCall p + | none => + let c' := c ++ ctx.varName + let cinfo ← getConstInfo c + let resultTy ← forallTelescope cinfo.type fun _ b => pure b + if resultTy.isConstOf ``Lean.Parser.TrailingParser || resultTy.isConstOf ``Lean.Parser.Parser then do + -- synthesize a new `[combinatorAttr c]` + let some value ← pure cinfo.value? + | throwError "don't know how to generate {ctx.varName} for non-definition '{e}'" + unless (env.getModuleIdxFor? c).isNone || force do + throwError "refusing to generate code for imported parser declaration '{c}'; use `@[run_parser_attribute_hooks]` on its definition instead." + let value ← compileParserExpr <| replaceParserTy ctx value + let ty ← forallTelescope cinfo.type fun params _ => + params.foldrM (init := mkConst ctx.tyName) fun param ty => do + let paramTy ← replaceParserTy ctx <$> inferType param + return mkForall `_ BinderInfo.default paramTy ty + let decl := Declaration.defnDecl { + name := c', levelParams := [] + type := ty, value := value, hints := ReducibilityHints.opaque, safety := DefinitionSafety.safe + } + addAndCompile decl + modifyEnv (ctx.combinatorAttr.setDeclFor · c c') + if cinfo.type.isConst then + if let some kind ← parserNodeKind? cinfo.value! then + -- If the parser is parameter-less and produces a node of kind `kind`, + -- then tag the compiled definition as `[(builtin)Parenthesizer kind]` + -- (or `[(builtin)Formatter kind]`, resp.) + let attrName := if builtin then ctx.categoryAttr.defn.builtinName else ctx.categoryAttr.defn.name + -- Create syntax node for a simple attribute of the form + -- `def simple := leading_parser ident >> optional (ident <|> priorityParser)` + let stx := mkNode `Lean.Parser.Attr.simple #[mkIdent attrName, mkNullNode #[mkIdent kind]] + Attribute.add c' attrName stx + mkCall c' + else + -- if this is a generic function, e.g. `AndThen.andthen`, it's easier to just unfold it until we are + -- back to parser combinators + let some e' ← unfoldDefinition? e + | throwError "don't know how to generate {ctx.varName} for non-parser combinator '{e}'" + compileParserExpr e' +end + +variable {α} (ctx : Context α) (builtin : Bool) in +def compileEmbeddedParsers : ParserDescr → MetaM Unit + | ParserDescr.const _ => pure () + | ParserDescr.unary _ d => compileEmbeddedParsers d + | ParserDescr.binary _ d₁ d₂ => compileEmbeddedParsers d₁ *> compileEmbeddedParsers d₂ + | ParserDescr.parser constName => discard <| compileParserExpr ctx (mkConst constName) (builtin := builtin) (force := false) + | ParserDescr.node _ _ d => compileEmbeddedParsers d + | ParserDescr.nodeWithAntiquot _ _ d => compileEmbeddedParsers d + | ParserDescr.sepBy p _ psep _ => compileEmbeddedParsers p *> compileEmbeddedParsers psep + | ParserDescr.sepBy1 p _ psep _ => compileEmbeddedParsers p *> compileEmbeddedParsers psep + | ParserDescr.trailingNode _ _ _ d => compileEmbeddedParsers d + | ParserDescr.symbol _ => pure () + | ParserDescr.nonReservedSymbol _ _ => pure () + | ParserDescr.cat _ _ => pure () + +/-- Precondition: `α` must match `ctx.tyName`. -/ +unsafe def registerParserCompiler {α} (ctx : Context α) : IO Unit := do + Parser.registerParserAttributeHook { + postAdd := fun catName constName builtin => do + let info ← getConstInfo constName + if info.type.isConstOf ``Lean.ParserDescr || info.type.isConstOf ``Lean.TrailingParserDescr then + let d ← evalConstCheck ParserDescr `Lean.ParserDescr constName <|> + evalConstCheck TrailingParserDescr `Lean.TrailingParserDescr constName + compileEmbeddedParsers ctx d (builtin := builtin) |>.run' + else + -- `[run_builtin_parser_attribute_hooks]` => force compilation even if imported, do not apply `ctx.categoryAttr`. + let force := catName.isAnonymous + discard (compileParserExpr ctx (mkConst constName) (builtin := builtin) (force := force)).run' + } + +end ParserCompiler +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PremiseSelection.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PremiseSelection.lean new file mode 100644 index 0000000000000000000000000000000000000000..272cc0f422363d02b16c3214b0979dc7c7659248 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PremiseSelection.lean @@ -0,0 +1,135 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +prelude +import Lean.Elab.Command +import Lean.Meta.Eval +import Lean.Meta.CompletionName +import Init.Data.Random + +/-! +# An API for premise selection algorithms. + +This module provides a basic API for premise selection algorithms, +which are used to suggest identifiers that should be introduced in a proof. + +The core interface is the `Selector` type, which is a function from a metavariable +and a configuration to a list of suggestions. + +The `Selector` is registered as an environment extension, and the trivial (no suggestions) implementation +is `Lean.PremiseSelection.empty`. + +Lean does not provide a default premise selector, so this module is intended to be used in conjunction +with a downstream package which registers a premise selector. +-/ + +namespace Lean.PremiseSelection + +/-- +A `Suggestion` is essentially just an identifier and a confidence score that the identifier is relevant. +If the premise selection request included information about the intended use (e.g. in the simplifier, in `grind`, etc.) +the score may be adjusted for that application. + +-/ +structure Suggestion where + name : Name + /-- + The score of the suggestion, as a probability that this suggestion should be used. + -/ + score : Float + +structure Config where + /-- + The maximum number of suggestions to return. + -/ + maxSuggestions : Option Nat := none + /-- + The tactic that is calling the premise selection, e.g. `simp`, `grind`, or `aesop`. + This may be used to adjust the score of the suggestions + -/ + caller : Option Name := none + /-- + A filter on suggestions; only suggestions returning `true` should be returned. + (It can be better to filter on the premise selection side, to ensure that enough suggestions are returned.) + -/ + filter : Name → MetaM Bool := fun _ => pure true + /-- + An optional arbitrary "hint" to the premise selection algorithm. + There is no guarantee that the algorithm will make any use of the hint. + + Potential use cases include a natural language comment provided by the user + (e.g. allowing use of the premise selector as a search engine) + or including context from the current proof and/or file. + + We may later split these use cases into separate fields if necessary. + -/ + hint : Option String := none + +abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion) + +/-- +The trivial premise selector, which returns no suggestions. +-/ +def empty : Selector := fun _ _ => pure #[] + +/-- A random premise selection algorithm, provided solely for testing purposes. -/ +def random (gen : StdGen := ⟨37, 59⟩) : Selector := fun _ cfg => do + IO.stdGenRef.set gen + let env ← getEnv + let max := cfg.maxSuggestions.getD 10 + let consts := env.const2ModIdx.keysArray + let mut suggestions := #[] + while suggestions.size < max do + let i ← IO.rand 0 consts.size + let name := consts[i]! + if ! (`Lean).isPrefixOf name && Lean.Meta.allowCompletion env name then + suggestions := suggestions.push { name := name, score := 1.0 / consts.size.toFloat } + return suggestions + +initialize premiseSelectorExt : EnvExtension (Option Selector) ← + registerEnvExtension (pure none) + +/-- Generate premise suggestions for the given metavariable, using the currently registered premise selector. -/ +def select (m : MVarId) (c : Config := {}) : MetaM (Array Suggestion) := do + let some selector := premiseSelectorExt.getState (← getEnv) | + throwError "No premise selector registered. \ + (Note the Lean does not provide a default premise selector, these must be installed by a downstream library.)" + selector m c + +/-! +Currently the registration mechanism is just global state. +This means that if multiple modules register premise selectors, +the behaviour will be dependent on the order of loading modules. + +We should replace this with a mechanism so that +premise selectors are configured via options in the `lakefile`, and +commands are only used to override in a single declaration or file. +-/ + +/-- Set the current premise selector.-/ +def registerPremiseSelector (selector : Selector) : CoreM Unit := do + modifyEnv fun env => premiseSelectorExt.setState env (some selector) + +open Lean Elab Command in +@[builtin_command_elab setPremiseSelectorCmd, inherit_doc setPremiseSelectorCmd] +def elabSetPremiseSelector : CommandElab + | `(command| set_premise_selector $selector) => do + let selector ← liftTermElabM do + try + let selectorTerm ← Term.elabTermEnsuringType selector (some (Expr.const ``Selector [])) + unsafe Meta.evalExpr Selector (Expr.const ``Selector []) selectorTerm + catch _ => + throwError "Failed to elaborate {selector} as a `MVarId → Config → MetaM (Array Suggestion)`." + liftCoreM (registerPremiseSelector selector) + | _ => throwUnsupportedSyntax + +open Lean.Elab.Tactic in +@[builtin_tactic Lean.Parser.Tactic.suggestPremises] def evalSuggestPremises : Tactic := fun _ => + liftMetaTactic1 fun mvarId => do + let suggestions ← select mvarId + logInfo m!"Premise suggestions: {suggestions.map (·.name)}" + return mvarId + +end Lean.PremiseSelection diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PrettyPrinter.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PrettyPrinter.lean new file mode 100644 index 0000000000000000000000000000000000000000..e975626230884a9659c67a82925f563379027810 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PrettyPrinter.lean @@ -0,0 +1,190 @@ +/- +Copyright (c) 2020 Sebastian Ullrich. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Sebastian Ullrich +-/ +prelude +import Lean.PrettyPrinter.Delaborator +import Lean.PrettyPrinter.Parenthesizer +import Lean.PrettyPrinter.Formatter +import Lean.Parser.Module +import Lean.ParserCompiler +import Lean.Util.NumObjs +import Lean.Util.ShareCommon +namespace Lean + +def PPContext.runCoreM {α : Type} (ppCtx : PPContext) (x : CoreM α) : IO α := + Prod.fst <$> x.toIO { options := ppCtx.opts, currNamespace := ppCtx.currNamespace + openDecls := ppCtx.openDecls + fileName := "", fileMap := default + diag := getDiag ppCtx.opts } + { env := ppCtx.env, ngen := { namePrefix := `_pp_uniq } } + +def PPContext.runMetaM {α : Type} (ppCtx : PPContext) (x : MetaM α) : IO α := + ppCtx.runCoreM <| x.run' { lctx := ppCtx.lctx } { mctx := ppCtx.mctx } + +namespace PrettyPrinter + +def ppCategory (cat : Name) (stx : Syntax) : CoreM Format := do + let opts ← getOptions + let stx := (sanitizeSyntax stx).run' { options := opts } + parenthesizeCategory cat stx >>= formatCategory cat + +def ppTerm (stx : Term) : CoreM Format := ppCategory `term stx + +def ppUsing (e : Expr) (delab : Expr → MetaM Term) : MetaM Format := do + let lctx := (← getLCtx).sanitizeNames.run' { options := (← getOptions) } + Meta.withLCtx' lctx do + ppTerm (← delab e) + +register_builtin_option pp.exprSizes : Bool := { + defValue := false + group := "pp" + descr := "(pretty printer) prefix each embedded expression with its sizes in the format \ + (size disregarding sharing/size with sharing/size with max sharing)" +} + +private def maybePrependExprSizes (e : Expr) (f : Format) : MetaM Format := do + if pp.exprSizes.get (← getOptions) then + return f!"[size {e.sizeWithoutSharing}/{← e.numObjs}/{← (ShareCommon.shareCommon' e).numObjs}] {f}" + else + return f + +def ppExpr (e : Expr) : MetaM Format := do + ppUsing e delab >>= maybePrependExprSizes e + +/-- Return a `fmt` representing pretty-printed `e` together with a map from tags in `fmt` +to `Elab.Info` nodes produced by the delaborator at various subexpressions of `e`. -/ +def ppExprWithInfos (e : Expr) (optsPerPos : Delaborator.OptionsPerPos := {}) (delab := Delaborator.delab) + : MetaM FormatWithInfos := do + let lctx := (← getLCtx).sanitizeNames.run' { options := (← getOptions) } + Meta.withLCtx' lctx do + let (stx, infos) ← delabCore e optsPerPos delab + let fmt ← ppTerm stx >>= maybePrependExprSizes e + return ⟨fmt, infos⟩ + +open Delaborator in +def ppConstNameWithInfos (constName : Name) : MetaM FormatWithInfos := do + if let some info := (← getEnv).find? constName then + let delab := withOptionAtCurrPos `pp.tagAppFns true <| delabConst + PrettyPrinter.ppExprWithInfos (delab := delab) (.const constName <| info.levelParams.map mkLevelParam) + else + -- Still, let's sanitize the name. + let stx := mkIdent constName + let stx := (sanitizeSyntax stx).run' { options := (← getOptions) } + formatCategory `term stx + +@[export lean_pp_expr] +def ppExprLegacy (env : Environment) (mctx : MetavarContext) (lctx : LocalContext) (opts : Options) (e : Expr) : IO Format := + Prod.fst <$> ((withOptions (fun _ => opts) <| ppExpr e).run' { lctx := lctx } { mctx := mctx }).toIO + { fileName := "", fileMap := default } + { env := env } + +def ppTactic (stx : TSyntax `tactic) : CoreM Format := ppCategory `tactic stx + +def ppCommand (stx : Syntax.Command) : CoreM Format := ppCategory `command stx + +def ppModule (stx : TSyntax ``Parser.Module.module) : CoreM Format := do + parenthesize Lean.Parser.Module.module.parenthesizer stx >>= format Lean.Parser.Module.module.formatter + +open Delaborator in +/-- Pretty-prints a declaration `c` as `c.{} : `. -/ +def ppSignature (c : Name) : MetaM FormatWithInfos := do + let decl ← getConstInfo c + let e := Expr.const c (decl.levelParams.map mkLevelParam) + if pp.raw.get (← getOptions) then + return s!"{e} : {decl.type}" + else + let (stx, infos) ← delabCore e (delab := delabConstWithSignature) + return ⟨← ppTerm ⟨stx⟩, infos⟩ -- HACK: not a term + +private partial def noContext : MessageData → MessageData + | MessageData.withContext _ msg => noContext msg + | MessageData.withNamingContext ctx msg => MessageData.withNamingContext ctx (noContext msg) + | MessageData.nest n msg => MessageData.nest n (noContext msg) + | MessageData.group msg => MessageData.group (noContext msg) + | MessageData.compose msg₁ msg₂ => MessageData.compose (noContext msg₁) (noContext msg₂) + | MessageData.tagged tag msg => MessageData.tagged tag (noContext msg) + | MessageData.trace data header children => + MessageData.trace data (noContext header) (children.map noContext) + | msg => msg + +-- strip context (including environments with registered pretty printers) to prevent infinite recursion when pretty printing pretty printer error +private def withoutContext {m} [MonadExcept Exception m] (x : m α) : m α := + tryCatch x fun + | Exception.error ref msg => throw <| Exception.error ref (noContext msg) + | ex => throw ex + +builtin_initialize + ppFnsRef.set { + ppExprWithInfos := fun ctx e => ctx.runMetaM <| withoutContext <| ppExprWithInfos e + ppConstNameWithInfos := fun ctx n => ctx.runMetaM <| withoutContext <| ppConstNameWithInfos n + ppTerm := fun ctx stx => ctx.runCoreM <| withoutContext <| ppTerm stx + ppLevel := fun ctx l => return l.format (mvars := getPPMVarsLevels ctx.opts) + ppGoal := fun ctx mvarId => ctx.runMetaM <| withoutContext <| Meta.ppGoal mvarId + } + +builtin_initialize + registerTraceClass `PrettyPrinter + +@[builtin_init] +unsafe def registerParserCompilers : IO Unit := do + ParserCompiler.registerParserCompiler ⟨`parenthesizer, parenthesizerAttribute, combinatorParenthesizerAttribute⟩ + ParserCompiler.registerParserCompiler ⟨`formatter, formatterAttribute, combinatorFormatterAttribute⟩ + +end PrettyPrinter + +namespace MessageData + +open Lean PrettyPrinter Delaborator + +/-- +Turns a `MetaM FormatWithInfos` into a `MessageData.lazy` which will run the monadic value. +-/ +def ofFormatWithInfosM (fmt : MetaM FormatWithInfos) : MessageData := + .lazy fun ctx => do + match (← ctx.runMetaM fmt |>.toBaseIO) with + | .ok fmt => return .ofFormatWithInfos fmt + | .error ex => return m!"[Error pretty printing: {ex}]" + +/-- +Turns a `MetaM MessageData` into a `MessageData.lazy` which will run the monadic value. +The optional array of expressions is used to set the `hasSyntheticSorry` fields, and should +comprise the expressions that are included in the message data. +-/ +def ofLazyM (f : MetaM MessageData) (es : Array Expr := #[]) : MessageData := + .lazy + (f := fun ppctxt => do + match (← ppctxt.runMetaM f |>.toBaseIO) with + | .ok fmt => return fmt + | .error ex => return m!"[Error pretty printing: {ex}]" + ) + (hasSyntheticSorry := fun mvarctxt => es.any (fun a => + instantiateMVarsCore mvarctxt a |>.1.hasSyntheticSorry + )) + +/-- +Pretty print a const expression using `delabConst` and generate terminfo. +This function avoids inserting `@` if the constant is for a function whose first +argument is implicit, which is what the default `toMessageData` for `Expr` does. +Panics if `e` is not a constant. +-/ +def ofConst (e : Expr) : MessageData := + if e.isConst then + let delab : Delab := withOptionAtCurrPos `pp.tagAppFns true delabConst + .ofFormatWithInfosM (PrettyPrinter.ppExprWithInfos (delab := delab) e) + else + have : Inhabited MessageData := + ⟨m!"[Error pretty printing: expression not a constant]{Format.line}{e}"⟩ + panic! "not a constant" + +/-- Generates `MessageData` for a declaration `c` as `c.{} : `, with terminfo. -/ +def signature (c : Name) : MessageData := + .lazy fun ctx => do + match (← ctx.runMetaM (PrettyPrinter.ppSignature c) |>.toBaseIO) with + | .ok fmt => return .ofFormatWithInfos fmt + | .error ex => return m!"[Error pretty printing signature: {ex}]{Format.line}{c}" + +end MessageData + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PrivateName.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PrivateName.lean new file mode 100644 index 0000000000000000000000000000000000000000..7ddbada78a03378d00ef30dc8a1cdbfe8b12c921 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/PrivateName.lean @@ -0,0 +1,76 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Init.Notation +import Init.Data.Option.Coe + +namespace Lean + +/-! # Private name support. + + Suppose the user marks as declaration `n` as private. Then, we create + the name: `_private..0 ++ n`. + We say `_private..0` is the "private prefix" + + We assume that `n` is a valid user name and does not contain + `Name.num` constructors. Thus, we can easily convert from + private internal name to the user given name. +-/ + +def privateHeader : Name := `_private + +def mkPrivateNameCore (mainModule : Name) (n : Name) : Name := + Name.mkNum (privateHeader ++ mainModule) 0 ++ n + +def isPrivateName : Name → Bool + | n@(.str p _) => n == privateHeader || isPrivateName p + | .num p _ => isPrivateName p + | _ => false + +@[export lean_is_private_name] +def isPrivateNameExport (n : Name) : Bool := + isPrivateName n + +/-- +Return `true` if `n` is of the form `_private..0` +See comment above. +-/ +def isPrivatePrefix (n : Name) : Bool := + match n with + | .num p 0 => go p + | _ => false +where + go (n : Name) : Bool := + n == privateHeader || + match n with + | .str p _ => go p + | _ => false + +private def privateToUserNameAux (n : Name) : Name := + match n with + | .str p s => .str (privateToUserNameAux p) s + | .num p i => if isPrivatePrefix n then .anonymous else .num (privateToUserNameAux p) i + | _ => .anonymous + +@[export lean_private_to_user_name] +def privateToUserName? (n : Name) : Option Name := + if isPrivateName n then privateToUserNameAux n + else none + +def privateToUserName (n : Name) : Name := + if isPrivateName n then privateToUserNameAux n + else n + +private def privatePrefixAux : Name → Name + | .str p _ => privatePrefixAux p + | n => n + +@[export lean_private_prefix] +def privatePrefix? (n : Name) : Option Name := + if isPrivateName n then privatePrefixAux n + else none + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ProjFns.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ProjFns.lean new file mode 100644 index 0000000000000000000000000000000000000000..2cd82d5a15d54a30b9a14f4f77b709c619cdc498 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ProjFns.lean @@ -0,0 +1,65 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.EnvExtension + +namespace Lean + +/-- +Given a structure `S`, Lean automatically creates an auxiliary definition (projection function) +for each field. This structure caches information about these auxiliary definitions. +-/ +structure ProjectionFunctionInfo where + /-- Constructor associated with the auxiliary projection function. -/ + ctorName : Name + /-- Number of parameters in the structure -/ + numParams : Nat + /-- The field index associated with the auxiliary projection function. -/ + i : Nat + /-- `true` if the structure is a class. -/ + fromClass : Bool + deriving Inhabited, Repr + +@[export lean_mk_projection_info] +def mkProjectionInfoEx (ctorName : Name) (numParams : Nat) (i : Nat) (fromClass : Bool) : ProjectionFunctionInfo := + { ctorName, numParams, i, fromClass } +@[export lean_projection_info_from_class] +def ProjectionFunctionInfo.fromClassEx (info : ProjectionFunctionInfo) : Bool := + info.fromClass + +builtin_initialize projectionFnInfoExt : MapDeclarationExtension ProjectionFunctionInfo ← mkMapDeclarationExtension + +@[export lean_add_projection_info] +def addProjectionFnInfo (env : Environment) (projName : Name) (ctorName : Name) (numParams : Nat) (i : Nat) (fromClass : Bool) : Environment := + projectionFnInfoExt.insert env projName { ctorName, numParams, i, fromClass } + +namespace Environment + +@[export lean_get_projection_info] +def getProjectionFnInfo? (env : Environment) (projName : Name) : Option ProjectionFunctionInfo := + projectionFnInfoExt.find? env projName + +def isProjectionFn (env : Environment) (declName : Name) : Bool := + projectionFnInfoExt.contains env declName + +/-- If `projName` is the name of a projection function, return the associated structure name -/ +def getProjectionStructureName? (env : Environment) (projName : Name) : Option Name := + match env.getProjectionFnInfo? projName with + | none => none + | some projInfo => + match env.find? projInfo.ctorName with + | some (ConstantInfo.ctorInfo val) => some val.induct + | _ => none + +end Environment + +def isProjectionFn [MonadEnv m] [Monad m] (declName : Name) : m Bool := + return (← getEnv).isProjectionFn declName + +def getProjectionFnInfo? [MonadEnv m] [Monad m] (declName : Name) : m (Option ProjectionFunctionInfo) := + return (← getEnv).getProjectionFnInfo? declName + +end Lean diff --git a/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ReducibilityAttrs.lean b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ReducibilityAttrs.lean new file mode 100644 index 0000000000000000000000000000000000000000..381d47b80bba2594ebd89e16a2ae9ef1fc740396 --- /dev/null +++ b/backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/ReducibilityAttrs.lean @@ -0,0 +1,194 @@ +/- +Copyright (c) 2019 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +prelude +import Lean.Attributes +import Lean.ScopedEnvExtension + +namespace Lean + +/-- +Reducibility status for a definition. +-/ +inductive ReducibilityStatus where + | reducible | semireducible | irreducible + deriving Inhabited, Repr, BEq + +def ReducibilityStatus.toAttrString : ReducibilityStatus → String + | .reducible => "[reducible]" + | .irreducible => "[irreducible]" + | .semireducible => "[semireducible]" + +builtin_initialize reducibilityCoreExt : PersistentEnvExtension (Name × ReducibilityStatus) (Name × ReducibilityStatus) (NameMap ReducibilityStatus) ← + registerPersistentEnvExtension { + name := `reducibilityCore + mkInitial := pure {} + addImportedFn := fun _ _ => pure {} + addEntryFn := fun (s : NameMap ReducibilityStatus) (p : Name × ReducibilityStatus) => s.insert p.1 p.2 + exportEntriesFn := fun m => + let r : Array (Name × ReducibilityStatus) := m.fold (fun a n p => a.push (n, p)) #[] + r.qsort (fun a b => Name.quickLt a.1 b.1) + statsFn := fun s => "reducibility attribute core extension" ++ Format.line ++ "number of local entries: " ++ format s.size + asyncMode := .async + } + +builtin_initialize reducibilityExtraExt : SimpleScopedEnvExtension (Name × ReducibilityStatus) (SMap Name ReducibilityStatus) ← + registerSimpleScopedEnvExtension { + name := `reducibilityExtra + initial := {} + addEntry := fun d (declName, status) => d.insert declName status + finalizeImport := fun d => d.switch + } + +@[export lean_get_reducibility_status] +def getReducibilityStatusCore (env : Environment) (declName : Name) : ReducibilityStatus := + let m := reducibilityExtraExt.getState env + if let some status := m.find? declName then + status + else match env.getModuleIdxFor? declName with + | some modIdx => + match (reducibilityCoreExt.getModuleEntries env modIdx).binSearch (declName, .semireducible) (fun a b => Name.quickLt a.1 b.1) with + | some (_, status) => status + | none => .semireducible + | none => (reducibilityCoreExt.findStateAsync env declName).find? declName |>.getD .semireducible + +private def setReducibilityStatusCore (env : Environment) (declName : Name) (status : ReducibilityStatus) (attrKind : AttributeKind) (currNamespace : Name) : Environment := + if attrKind matches .global then + match env.getModuleIdxFor? declName with + | some _ => + -- Trying to set the attribute of a declaration defined in an imported module. + reducibilityExtraExt.addEntry env (declName, status) + | none => + let _ : Inhabited Environment := ⟨env⟩ + assert! env.asyncMayContain declName + reducibilityCoreExt.addEntry env (declName, status) + else + -- `scoped` and `local` must be handled by `reducibilityExtraExt` + reducibilityExtraExt.addCore env (declName, status) attrKind currNamespace + +@[export lean_set_reducibility_status] +private def setReducibilityStatusImp (env : Environment) (declName : Name) (status : ReducibilityStatus) : Environment := + setReducibilityStatusCore env declName status .global .anonymous + +/- +TODO: it would be great if we could distinguish between the following two situations + +1- +``` +@[reducible] def foo := ... +``` + +2- +``` +def foo := ... +... +attribute [reducible] foo +``` + +Reason: the second one is problematic if user has add simp theorems or TC instances that include `foo`. +Recall that the discrimination trees unfold `[reducible]` declarations while indexing new entries. +-/ + +register_builtin_option allowUnsafeReducibility : Bool := { + defValue := false + descr := "enables users to modify the reducibility settings for declarations even when such changes are deemed potentially hazardous. For example, `simp` and type class resolution maintain term indices where reducible declarations are expanded." +} + +private def validate (declName : Name) (status : ReducibilityStatus) (attrKind : AttributeKind) : CoreM Unit := do + let suffix := .note "Use `set_option allowUnsafeReducibility true` to override reducibility status validation" + unless allowUnsafeReducibility.get (← getOptions) do + match (← getConstInfo declName) with + | .defnInfo _ => + let statusOld := getReducibilityStatusCore (← getEnv) declName + match attrKind with + | .scoped => + throwError "failed to set reducibility status for `{declName}`, the `scoped` modifier is not recommended for this kind of attribute{suffix}" + | .global => + if (← getEnv).getModuleIdxFor? declName matches some _ then + throwError "failed to set reducibility status, `{declName}` has not been defined in this file, consider using the `local` modifier{suffix}" + match status with + | .reducible => + unless statusOld matches .semireducible do + throwError "failed to set `[reducible]`, `{declName}` is not currently `[semireducible]`, but `{statusOld.toAttrString}`{suffix}" + | .irreducible => + unless statusOld matches .semireducible do + throwError "failed to set `[irreducible]`, `{declName}` is not currently `[semireducible]`, but `{statusOld.toAttrString}`{suffix}" + | .semireducible => + throwError "failed to set `[semireducible]` for `{declName}`, declarations are `[semireducible]` by default{suffix}" + | .local => + match status with + | .reducible => + throwError "failed to set `[local reducible]` for `{declName}`, recall that `[reducible]` affects the term indexing datastructures used by `simp` and type class resolution{suffix}" + | .irreducible => + unless statusOld matches .semireducible do + throwError "failed to set `[local irreducible]`, `{declName}` is currently `{statusOld.toAttrString}`, `[semireducible]` expected{suffix}" + | .semireducible => + unless statusOld matches .irreducible do + throwError "failed to set `[local semireducible]`, `{declName}` is currently `{statusOld.toAttrString}`, `[irreducible]` expected{suffix}" + | _ => throwError "failed to set reducibility status, `{declName}` is not a definition{suffix}" + +private def addAttr (status : ReducibilityStatus) (declName : Name) (stx : Syntax) (attrKind : AttributeKind) : AttrM Unit := do + Attribute.Builtin.ensureNoArgs stx + validate declName status attrKind + let ns ← getCurrNamespace + modifyEnv fun env => setReducibilityStatusCore env declName status attrKind ns + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `irreducible + descr := "irreducible declaration" + add := addAttr .irreducible + applicationTime := .afterTypeChecking + } + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `reducible + descr := "reducible declaration" + add := addAttr .reducible + applicationTime := .afterTypeChecking + } + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `semireducible + descr := "semireducible declaration" + add := addAttr .semireducible + applicationTime := .afterTypeChecking + } + +/-- Return the reducibility attribute for the given declaration. -/ +def getReducibilityStatus [Monad m] [MonadEnv m] (declName : Name) : m ReducibilityStatus := do + return getReducibilityStatusCore (← getEnv) declName + +/-- Set the reducibility attribute for the given declaration. -/ +def setReducibilityStatus [MonadEnv m] (declName : Name) (s : ReducibilityStatus) : m Unit := + modifyEnv fun env => setReducibilityStatusCore env declName s .global .anonymous + +/-- Set the given declaration as `[reducible]` -/ +def setReducibleAttribute [MonadEnv m] (declName : Name) : m Unit := + setReducibilityStatus declName ReducibilityStatus.reducible + +/-- Return `true` if the given declaration has been marked as `[reducible]`. -/ +def isReducible [Monad m] [MonadEnv m] (declName : Name) : m Bool := do + match (← getReducibilityStatus declName) with + | .reducible => return true + | _ => return false + +/-- Return `true` if the given declaration has been marked as `[irreducible]` -/ +def isIrreducible [Monad m] [MonadEnv m] (declName : Name) : m Bool := do + match (← getReducibilityStatus declName) with + | .irreducible => return true + | _ => return false + +/-- Set the given declaration as `[irreducible]` -/ +def setIrreducibleAttribute [MonadEnv m] (declName : Name) : m Unit := + setReducibilityStatus declName ReducibilityStatus.irreducible + + +end Lean diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/__main__.py b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..7fbfc015cc4c367ed8c0ec9e8957ff67f308d602 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/__main__.py @@ -0,0 +1,35 @@ +import sys + + +def main(args=None): + if args is None: + args = sys.argv[1:] + + # TODO Handle library-wide options. Eg.: + # --unicodedata + # --verbose / other logging stuff + + # TODO Allow a way to run arbitrary modules? Useful for setting + # library-wide options and calling another library. Eg.: + # + # $ fonttools --unicodedata=... fontmake ... + # + # This allows for a git-like command where thirdparty commands + # can be added. Should we just try importing the fonttools + # module first and try without if it fails? + + if len(sys.argv) < 2: + sys.argv.append("help") + if sys.argv[1] == "-h" or sys.argv[1] == "--help": + sys.argv[1] = "help" + mod = "fontTools." + sys.argv[1] + sys.argv[1] = sys.argv[0] + " " + sys.argv[1] + del sys.argv[0] + + import runpy + + runpy.run_module(mod, run_name="__main__") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/CFF2ToCFF.py b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/CFF2ToCFF.py new file mode 100644 index 0000000000000000000000000000000000000000..1f5ff74e08b81423a21d558731db96677820206a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/CFF2ToCFF.py @@ -0,0 +1,233 @@ +"""CFF2 to CFF converter.""" + +from fontTools.ttLib import TTFont, newTable +from fontTools.misc.cliTools import makeOutputFileName +from fontTools.misc.psCharStrings import T2StackUseExtractor +from fontTools.cffLib import ( + TopDictIndex, + buildOrder, + buildDefaults, + topDictOperators, + privateDictOperators, + FDSelect, +) +from .transforms import desubroutinizeCharString +from .specializer import specializeProgram +from .width import optimizeWidths +from collections import defaultdict +import logging + + +__all__ = ["convertCFF2ToCFF", "main"] + + +log = logging.getLogger("fontTools.cffLib") + + +def _convertCFF2ToCFF(cff, otFont): + """Converts this object from CFF2 format to CFF format. This conversion + is done 'in-place'. The conversion cannot be reversed. + + The CFF2 font cannot be variable. (TODO Accept those and convert to the + default instance?) + + This assumes a decompiled CFF2 table. (i.e. that the object has been + filled via :meth:`decompile` and e.g. not loaded from XML.)""" + + cff.major = 1 + + topDictData = TopDictIndex(None) + for item in cff.topDictIndex: + # Iterate over, such that all are decompiled + item.cff2GetGlyphOrder = None + topDictData.append(item) + cff.topDictIndex = topDictData + topDict = topDictData[0] + + if hasattr(topDict, "VarStore"): + raise ValueError("Variable CFF2 font cannot be converted to CFF format.") + + opOrder = buildOrder(topDictOperators) + topDict.order = opOrder + for key in topDict.rawDict.keys(): + if key not in opOrder: + del topDict.rawDict[key] + if hasattr(topDict, key): + delattr(topDict, key) + + charStrings = topDict.CharStrings + + fdArray = topDict.FDArray + if not hasattr(topDict, "FDSelect"): + # FDSelect is optional in CFF2, but required in CFF. + fdSelect = topDict.FDSelect = FDSelect() + fdSelect.gidArray = [0] * len(charStrings.charStrings) + + defaults = buildDefaults(privateDictOperators) + order = buildOrder(privateDictOperators) + for fd in fdArray: + fd.setCFF2(False) + privateDict = fd.Private + privateDict.order = order + for key in order: + if key not in privateDict.rawDict and key in defaults: + privateDict.rawDict[key] = defaults[key] + for key in privateDict.rawDict.keys(): + if key not in order: + del privateDict.rawDict[key] + if hasattr(privateDict, key): + delattr(privateDict, key) + + # Add ending operators + for cs in charStrings.values(): + cs.decompile() + cs.program.append("endchar") + for subrSets in [cff.GlobalSubrs] + [ + getattr(fd.Private, "Subrs", []) for fd in fdArray + ]: + for cs in subrSets: + cs.program.append("return") + + # Add (optimal) width to CharStrings that need it. + widths = defaultdict(list) + metrics = otFont["hmtx"].metrics + for glyphName in charStrings.keys(): + cs, fdIndex = charStrings.getItemAndSelector(glyphName) + if fdIndex == None: + fdIndex = 0 + widths[fdIndex].append(metrics[glyphName][0]) + for fdIndex, widthList in widths.items(): + bestDefault, bestNominal = optimizeWidths(widthList) + private = fdArray[fdIndex].Private + private.defaultWidthX = bestDefault + private.nominalWidthX = bestNominal + for glyphName in charStrings.keys(): + cs, fdIndex = charStrings.getItemAndSelector(glyphName) + if fdIndex == None: + fdIndex = 0 + private = fdArray[fdIndex].Private + width = metrics[glyphName][0] + if width != private.defaultWidthX: + cs.program.insert(0, width - private.nominalWidthX) + + # Handle stack use since stack-depth is lower in CFF than in CFF2. + for glyphName in charStrings.keys(): + cs, fdIndex = charStrings.getItemAndSelector(glyphName) + if fdIndex is None: + fdIndex = 0 + private = fdArray[fdIndex].Private + extractor = T2StackUseExtractor( + getattr(private, "Subrs", []), cff.GlobalSubrs, private=private + ) + stackUse = extractor.execute(cs) + if stackUse > 48: # CFF stack depth is 48 + desubroutinizeCharString(cs) + cs.program = specializeProgram(cs.program) + + # Unused subroutines are still in CFF2 (ie. lacking 'return' operator) + # because they were not decompiled when we added the 'return'. + # Moreover, some used subroutines may have become unused after the + # stack-use fixup. So we remove all unused subroutines now. + cff.remove_unused_subroutines() + + mapping = { + name: ("cid" + str(n).zfill(5) if n else ".notdef") + for n, name in enumerate(topDict.charset) + } + topDict.charset = [ + "cid" + str(n).zfill(5) if n else ".notdef" for n in range(len(topDict.charset)) + ] + charStrings.charStrings = { + mapping[name]: v for name, v in charStrings.charStrings.items() + } + + topDict.ROS = ("Adobe", "Identity", 0) + + +def convertCFF2ToCFF(font, *, updatePostTable=True): + if "CFF2" not in font: + raise ValueError("Input font does not contain a CFF2 table.") + cff = font["CFF2"].cff + _convertCFF2ToCFF(cff, font) + del font["CFF2"] + table = font["CFF "] = newTable("CFF ") + table.cff = cff + + if updatePostTable and "post" in font: + # Only version supported for fonts with CFF table is 0x00030000 not 0x20000 + post = font["post"] + if post.formatType == 2.0: + post.formatType = 3.0 + + +def main(args=None): + """Convert CFF2 OTF font to CFF OTF font""" + if args is None: + import sys + + args = sys.argv[1:] + + import argparse + + parser = argparse.ArgumentParser( + "fonttools cffLib.CFF2ToCFF", + description="Convert a non-variable CFF2 font to CFF.", + ) + parser.add_argument( + "input", metavar="INPUT.ttf", help="Input OTF file with CFF table." + ) + parser.add_argument( + "-o", + "--output", + metavar="OUTPUT.ttf", + default=None, + help="Output instance OTF file (default: INPUT-CFF2.ttf).", + ) + parser.add_argument( + "--no-recalc-timestamp", + dest="recalc_timestamp", + action="store_false", + help="Don't set the output font's timestamp to the current time.", + ) + loggingGroup = parser.add_mutually_exclusive_group(required=False) + loggingGroup.add_argument( + "-v", "--verbose", action="store_true", help="Run more verbosely." + ) + loggingGroup.add_argument( + "-q", "--quiet", action="store_true", help="Turn verbosity off." + ) + options = parser.parse_args(args) + + from fontTools import configLogger + + configLogger( + level=("DEBUG" if options.verbose else "ERROR" if options.quiet else "INFO") + ) + + import os + + infile = options.input + if not os.path.isfile(infile): + parser.error("No such file '{}'".format(infile)) + + outfile = ( + makeOutputFileName(infile, overWrite=True, suffix="-CFF") + if not options.output + else options.output + ) + + font = TTFont(infile, recalcTimestamp=options.recalc_timestamp, recalcBBoxes=False) + + convertCFF2ToCFF(font) + + log.info( + "Saving %s", + outfile, + ) + font.save(outfile) + + +if __name__ == "__main__": + import sys + + sys.exit(main(sys.argv[1:])) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/CFFToCFF2.py b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/CFFToCFF2.py new file mode 100644 index 0000000000000000000000000000000000000000..4dd728cad4757b45f6f8709e2ca67b747e39951d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/CFFToCFF2.py @@ -0,0 +1,305 @@ +"""CFF to CFF2 converter.""" + +from fontTools.ttLib import TTFont, newTable +from fontTools.misc.cliTools import makeOutputFileName +from fontTools.misc.psCharStrings import T2WidthExtractor +from fontTools.cffLib import ( + TopDictIndex, + FDArrayIndex, + FontDict, + buildOrder, + topDictOperators, + privateDictOperators, + topDictOperators2, + privateDictOperators2, +) +from io import BytesIO +import logging + +__all__ = ["convertCFFToCFF2", "main"] + + +log = logging.getLogger("fontTools.cffLib") + + +class _NominalWidthUsedError(Exception): + def __add__(self, other): + raise self + + def __radd__(self, other): + raise self + + +def _convertCFFToCFF2(cff, otFont): + """Converts this object from CFF format to CFF2 format. This conversion + is done 'in-place'. The conversion cannot be reversed. + + This assumes a decompiled CFF table. (i.e. that the object has been + filled via :meth:`decompile` and e.g. not loaded from XML.)""" + + # Clean up T2CharStrings + + topDict = cff.topDictIndex[0] + fdArray = topDict.FDArray if hasattr(topDict, "FDArray") else None + charStrings = topDict.CharStrings + globalSubrs = cff.GlobalSubrs + localSubrs = ( + [getattr(fd.Private, "Subrs", []) for fd in fdArray] + if fdArray + else ( + [topDict.Private.Subrs] + if hasattr(topDict, "Private") and hasattr(topDict.Private, "Subrs") + else [] + ) + ) + + for glyphName in charStrings.keys(): + cs, fdIndex = charStrings.getItemAndSelector(glyphName) + cs.decompile() + + # Clean up subroutines first + for subrs in [globalSubrs] + localSubrs: + for subr in subrs: + program = subr.program + i = j = len(program) + try: + i = program.index("return") + except ValueError: + pass + try: + j = program.index("endchar") + except ValueError: + pass + program[min(i, j) :] = [] + + # Clean up glyph charstrings + removeUnusedSubrs = False + nominalWidthXError = _NominalWidthUsedError() + for glyphName in charStrings.keys(): + cs, fdIndex = charStrings.getItemAndSelector(glyphName) + program = cs.program + + thisLocalSubrs = ( + localSubrs[fdIndex] + if fdIndex is not None + else ( + getattr(topDict.Private, "Subrs", []) + if hasattr(topDict, "Private") + else [] + ) + ) + + # Intentionally use custom type for nominalWidthX, such that any + # CharString that has an explicit width encoded will throw back to us. + extractor = T2WidthExtractor( + thisLocalSubrs, + globalSubrs, + nominalWidthXError, + 0, + ) + try: + extractor.execute(cs) + except _NominalWidthUsedError: + # Program has explicit width. We want to drop it, but can't + # just pop the first number since it may be a subroutine call. + # Instead, when seeing that, we embed the subroutine and recurse. + # If this ever happened, we later prune unused subroutines. + while len(program) >= 2 and program[1] in ["callsubr", "callgsubr"]: + removeUnusedSubrs = True + subrNumber = program.pop(0) + assert isinstance(subrNumber, int), subrNumber + op = program.pop(0) + bias = extractor.localBias if op == "callsubr" else extractor.globalBias + subrNumber += bias + subrSet = thisLocalSubrs if op == "callsubr" else globalSubrs + subrProgram = subrSet[subrNumber].program + program[:0] = subrProgram + # Now pop the actual width + assert len(program) >= 1, program + program.pop(0) + + if program and program[-1] == "endchar": + program.pop() + + if removeUnusedSubrs: + cff.remove_unused_subroutines() + + # Upconvert TopDict + + cff.major = 2 + cff2GetGlyphOrder = cff.otFont.getGlyphOrder + topDictData = TopDictIndex(None, cff2GetGlyphOrder) + for item in cff.topDictIndex: + # Iterate over, such that all are decompiled + topDictData.append(item) + cff.topDictIndex = topDictData + topDict = topDictData[0] + if hasattr(topDict, "Private"): + privateDict = topDict.Private + else: + privateDict = None + opOrder = buildOrder(topDictOperators2) + topDict.order = opOrder + topDict.cff2GetGlyphOrder = cff2GetGlyphOrder + + if not hasattr(topDict, "FDArray"): + fdArray = topDict.FDArray = FDArrayIndex() + fdArray.strings = None + fdArray.GlobalSubrs = topDict.GlobalSubrs + topDict.GlobalSubrs.fdArray = fdArray + charStrings = topDict.CharStrings + if charStrings.charStringsAreIndexed: + charStrings.charStringsIndex.fdArray = fdArray + else: + charStrings.fdArray = fdArray + fontDict = FontDict() + fontDict.setCFF2(True) + fdArray.append(fontDict) + fontDict.Private = privateDict + privateOpOrder = buildOrder(privateDictOperators2) + if privateDict is not None: + for entry in privateDictOperators: + key = entry[1] + if key not in privateOpOrder: + if key in privateDict.rawDict: + # print "Removing private dict", key + del privateDict.rawDict[key] + if hasattr(privateDict, key): + delattr(privateDict, key) + # print "Removing privateDict attr", key + else: + # clean up the PrivateDicts in the fdArray + fdArray = topDict.FDArray + privateOpOrder = buildOrder(privateDictOperators2) + for fontDict in fdArray: + fontDict.setCFF2(True) + for key in list(fontDict.rawDict.keys()): + if key not in fontDict.order: + del fontDict.rawDict[key] + if hasattr(fontDict, key): + delattr(fontDict, key) + + privateDict = fontDict.Private + for entry in privateDictOperators: + key = entry[1] + if key not in privateOpOrder: + if key in list(privateDict.rawDict.keys()): + # print "Removing private dict", key + del privateDict.rawDict[key] + if hasattr(privateDict, key): + delattr(privateDict, key) + # print "Removing privateDict attr", key + + # Now delete up the deprecated topDict operators from CFF 1.0 + for entry in topDictOperators: + key = entry[1] + # We seem to need to keep the charset operator for now, + # or we fail to compile with some fonts, like AdditionFont.otf. + # I don't know which kind of CFF font those are. But keeping + # charset seems to work. It will be removed when we save and + # read the font again. + # + # AdditionFont.otf has . + if key == "charset": + continue + if key not in opOrder: + if key in topDict.rawDict: + del topDict.rawDict[key] + if hasattr(topDict, key): + delattr(topDict, key) + + # TODO(behdad): What does the following comment even mean? Both CFF and CFF2 + # use the same T2Charstring class. I *think* what it means is that the CharStrings + # were loaded for CFF1, and we need to reload them for CFF2 to set varstore, etc + # on them. At least that's what I understand. It's probably safe to remove this + # and just set vstore where needed. + # + # See comment above about charset as well. + + # At this point, the Subrs and Charstrings are all still T2Charstring class + # easiest to fix this by compiling, then decompiling again + file = BytesIO() + cff.compile(file, otFont, isCFF2=True) + file.seek(0) + cff.decompile(file, otFont, isCFF2=True) + + +def convertCFFToCFF2(font): + cff = font["CFF "].cff + del font["CFF "] + _convertCFFToCFF2(cff, font) + table = font["CFF2"] = newTable("CFF2") + table.cff = cff + + +def main(args=None): + """Convert CFF OTF font to CFF2 OTF font""" + if args is None: + import sys + + args = sys.argv[1:] + + import argparse + + parser = argparse.ArgumentParser( + "fonttools cffLib.CFFToCFF2", + description="Upgrade a CFF font to CFF2.", + ) + parser.add_argument( + "input", metavar="INPUT.ttf", help="Input OTF file with CFF table." + ) + parser.add_argument( + "-o", + "--output", + metavar="OUTPUT.ttf", + default=None, + help="Output instance OTF file (default: INPUT-CFF2.ttf).", + ) + parser.add_argument( + "--no-recalc-timestamp", + dest="recalc_timestamp", + action="store_false", + help="Don't set the output font's timestamp to the current time.", + ) + loggingGroup = parser.add_mutually_exclusive_group(required=False) + loggingGroup.add_argument( + "-v", "--verbose", action="store_true", help="Run more verbosely." + ) + loggingGroup.add_argument( + "-q", "--quiet", action="store_true", help="Turn verbosity off." + ) + options = parser.parse_args(args) + + from fontTools import configLogger + + configLogger( + level=("DEBUG" if options.verbose else "ERROR" if options.quiet else "INFO") + ) + + import os + + infile = options.input + if not os.path.isfile(infile): + parser.error("No such file '{}'".format(infile)) + + outfile = ( + makeOutputFileName(infile, overWrite=True, suffix="-CFF2") + if not options.output + else options.output + ) + + font = TTFont(infile, recalcTimestamp=options.recalc_timestamp, recalcBBoxes=False) + + convertCFFToCFF2(font) + + log.info( + "Saving %s", + outfile, + ) + font.save(outfile) + + +if __name__ == "__main__": + import sys + + sys.exit(main(sys.argv[1:])) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/specializer.py b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/specializer.py new file mode 100644 index 0000000000000000000000000000000000000000..af3aac442b4cddbb8f1bc43208161a42aec80062 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/specializer.py @@ -0,0 +1,927 @@ +# -*- coding: utf-8 -*- + +"""T2CharString operator specializer and generalizer. + +PostScript glyph drawing operations can be expressed in multiple different +ways. For example, as well as the ``lineto`` operator, there is also a +``hlineto`` operator which draws a horizontal line, removing the need to +specify a ``dx`` coordinate, and a ``vlineto`` operator which draws a +vertical line, removing the need to specify a ``dy`` coordinate. As well +as decompiling :class:`fontTools.misc.psCharStrings.T2CharString` objects +into lists of operations, this module allows for conversion between general +and specific forms of the operation. + +""" + +from fontTools.cffLib import maxStackLimit + + +def stringToProgram(string): + if isinstance(string, str): + string = string.split() + program = [] + for token in string: + try: + token = int(token) + except ValueError: + try: + token = float(token) + except ValueError: + pass + program.append(token) + return program + + +def programToString(program): + return " ".join(str(x) for x in program) + + +def programToCommands(program, getNumRegions=None): + """Takes a T2CharString program list and returns list of commands. + Each command is a two-tuple of commandname,arg-list. The commandname might + be empty string if no commandname shall be emitted (used for glyph width, + hintmask/cntrmask argument, as well as stray arguments at the end of the + program (🤷). + 'getNumRegions' may be None, or a callable object. It must return the + number of regions. 'getNumRegions' takes a single argument, vsindex. It + returns the numRegions for the vsindex. + The Charstring may or may not start with a width value. If the first + non-blend operator has an odd number of arguments, then the first argument is + a width, and is popped off. This is complicated with blend operators, as + there may be more than one before the first hint or moveto operator, and each + one reduces several arguments to just one list argument. We have to sum the + number of arguments that are not part of the blend arguments, and all the + 'numBlends' values. We could instead have said that by definition, if there + is a blend operator, there is no width value, since CFF2 Charstrings don't + have width values. I discussed this with Behdad, and we are allowing for an + initial width value in this case because developers may assemble a CFF2 + charstring from CFF Charstrings, which could have width values. + """ + + seenWidthOp = False + vsIndex = 0 + lenBlendStack = 0 + lastBlendIndex = 0 + commands = [] + stack = [] + it = iter(program) + + for token in it: + if not isinstance(token, str): + stack.append(token) + continue + + if token == "blend": + assert getNumRegions is not None + numSourceFonts = 1 + getNumRegions(vsIndex) + # replace the blend op args on the stack with a single list + # containing all the blend op args. + numBlends = stack[-1] + numBlendArgs = numBlends * numSourceFonts + 1 + # replace first blend op by a list of the blend ops. + stack[-numBlendArgs:] = [stack[-numBlendArgs:]] + lenStack = len(stack) + lenBlendStack += numBlends + lenStack - 1 + lastBlendIndex = lenStack + # if a blend op exists, this is or will be a CFF2 charstring. + continue + + elif token == "vsindex": + vsIndex = stack[-1] + assert type(vsIndex) is int + + elif (not seenWidthOp) and token in { + "hstem", + "hstemhm", + "vstem", + "vstemhm", + "cntrmask", + "hintmask", + "hmoveto", + "vmoveto", + "rmoveto", + "endchar", + }: + seenWidthOp = True + parity = token in {"hmoveto", "vmoveto"} + if lenBlendStack: + # lenBlendStack has the number of args represented by the last blend + # arg and all the preceding args. We need to now add the number of + # args following the last blend arg. + numArgs = lenBlendStack + len(stack[lastBlendIndex:]) + else: + numArgs = len(stack) + if numArgs and (numArgs % 2) ^ parity: + width = stack.pop(0) + commands.append(("", [width])) + + if token in {"hintmask", "cntrmask"}: + if stack: + commands.append(("", stack)) + commands.append((token, [])) + commands.append(("", [next(it)])) + else: + commands.append((token, stack)) + stack = [] + if stack: + commands.append(("", stack)) + return commands + + +def _flattenBlendArgs(args): + token_list = [] + for arg in args: + if isinstance(arg, list): + token_list.extend(arg) + token_list.append("blend") + else: + token_list.append(arg) + return token_list + + +def commandsToProgram(commands): + """Takes a commands list as returned by programToCommands() and converts + it back to a T2CharString program list.""" + program = [] + for op, args in commands: + if any(isinstance(arg, list) for arg in args): + args = _flattenBlendArgs(args) + program.extend(args) + if op: + program.append(op) + return program + + +def _everyN(el, n): + """Group the list el into groups of size n""" + l = len(el) + if l % n != 0: + raise ValueError(el) + for i in range(0, l, n): + yield el[i : i + n] + + +class _GeneralizerDecombinerCommandsMap(object): + @staticmethod + def rmoveto(args): + if len(args) != 2: + raise ValueError(args) + yield ("rmoveto", args) + + @staticmethod + def hmoveto(args): + if len(args) != 1: + raise ValueError(args) + yield ("rmoveto", [args[0], 0]) + + @staticmethod + def vmoveto(args): + if len(args) != 1: + raise ValueError(args) + yield ("rmoveto", [0, args[0]]) + + @staticmethod + def rlineto(args): + if not args: + raise ValueError(args) + for args in _everyN(args, 2): + yield ("rlineto", args) + + @staticmethod + def hlineto(args): + if not args: + raise ValueError(args) + it = iter(args) + try: + while True: + yield ("rlineto", [next(it), 0]) + yield ("rlineto", [0, next(it)]) + except StopIteration: + pass + + @staticmethod + def vlineto(args): + if not args: + raise ValueError(args) + it = iter(args) + try: + while True: + yield ("rlineto", [0, next(it)]) + yield ("rlineto", [next(it), 0]) + except StopIteration: + pass + + @staticmethod + def rrcurveto(args): + if not args: + raise ValueError(args) + for args in _everyN(args, 6): + yield ("rrcurveto", args) + + @staticmethod + def hhcurveto(args): + l = len(args) + if l < 4 or l % 4 > 1: + raise ValueError(args) + if l % 2 == 1: + yield ("rrcurveto", [args[1], args[0], args[2], args[3], args[4], 0]) + args = args[5:] + for args in _everyN(args, 4): + yield ("rrcurveto", [args[0], 0, args[1], args[2], args[3], 0]) + + @staticmethod + def vvcurveto(args): + l = len(args) + if l < 4 or l % 4 > 1: + raise ValueError(args) + if l % 2 == 1: + yield ("rrcurveto", [args[0], args[1], args[2], args[3], 0, args[4]]) + args = args[5:] + for args in _everyN(args, 4): + yield ("rrcurveto", [0, args[0], args[1], args[2], 0, args[3]]) + + @staticmethod + def hvcurveto(args): + l = len(args) + if l < 4 or l % 8 not in {0, 1, 4, 5}: + raise ValueError(args) + last_args = None + if l % 2 == 1: + lastStraight = l % 8 == 5 + args, last_args = args[:-5], args[-5:] + it = _everyN(args, 4) + try: + while True: + args = next(it) + yield ("rrcurveto", [args[0], 0, args[1], args[2], 0, args[3]]) + args = next(it) + yield ("rrcurveto", [0, args[0], args[1], args[2], args[3], 0]) + except StopIteration: + pass + if last_args: + args = last_args + if lastStraight: + yield ("rrcurveto", [args[0], 0, args[1], args[2], args[4], args[3]]) + else: + yield ("rrcurveto", [0, args[0], args[1], args[2], args[3], args[4]]) + + @staticmethod + def vhcurveto(args): + l = len(args) + if l < 4 or l % 8 not in {0, 1, 4, 5}: + raise ValueError(args) + last_args = None + if l % 2 == 1: + lastStraight = l % 8 == 5 + args, last_args = args[:-5], args[-5:] + it = _everyN(args, 4) + try: + while True: + args = next(it) + yield ("rrcurveto", [0, args[0], args[1], args[2], args[3], 0]) + args = next(it) + yield ("rrcurveto", [args[0], 0, args[1], args[2], 0, args[3]]) + except StopIteration: + pass + if last_args: + args = last_args + if lastStraight: + yield ("rrcurveto", [0, args[0], args[1], args[2], args[3], args[4]]) + else: + yield ("rrcurveto", [args[0], 0, args[1], args[2], args[4], args[3]]) + + @staticmethod + def rcurveline(args): + l = len(args) + if l < 8 or l % 6 != 2: + raise ValueError(args) + args, last_args = args[:-2], args[-2:] + for args in _everyN(args, 6): + yield ("rrcurveto", args) + yield ("rlineto", last_args) + + @staticmethod + def rlinecurve(args): + l = len(args) + if l < 8 or l % 2 != 0: + raise ValueError(args) + args, last_args = args[:-6], args[-6:] + for args in _everyN(args, 2): + yield ("rlineto", args) + yield ("rrcurveto", last_args) + + +def _convertBlendOpToArgs(blendList): + # args is list of blend op args. Since we are supporting + # recursive blend op calls, some of these args may also + # be a list of blend op args, and need to be converted before + # we convert the current list. + if any([isinstance(arg, list) for arg in blendList]): + args = [ + i + for e in blendList + for i in (_convertBlendOpToArgs(e) if isinstance(e, list) else [e]) + ] + else: + args = blendList + + # We now know that blendList contains a blend op argument list, even if + # some of the args are lists that each contain a blend op argument list. + # Convert from: + # [default font arg sequence x0,...,xn] + [delta tuple for x0] + ... + [delta tuple for xn] + # to: + # [ [x0] + [delta tuple for x0], + # ..., + # [xn] + [delta tuple for xn] ] + numBlends = args[-1] + # Can't use args.pop() when the args are being used in a nested list + # comprehension. See calling context + args = args[:-1] + + l = len(args) + numRegions = l // numBlends - 1 + if not (numBlends * (numRegions + 1) == l): + raise ValueError(blendList) + + defaultArgs = [[arg] for arg in args[:numBlends]] + deltaArgs = args[numBlends:] + numDeltaValues = len(deltaArgs) + deltaList = [ + deltaArgs[i : i + numRegions] for i in range(0, numDeltaValues, numRegions) + ] + blend_args = [a + b + [1] for a, b in zip(defaultArgs, deltaList)] + return blend_args + + +def generalizeCommands(commands, ignoreErrors=False): + result = [] + mapping = _GeneralizerDecombinerCommandsMap + for op, args in commands: + # First, generalize any blend args in the arg list. + if any([isinstance(arg, list) for arg in args]): + try: + args = [ + n + for arg in args + for n in ( + _convertBlendOpToArgs(arg) if isinstance(arg, list) else [arg] + ) + ] + except ValueError: + if ignoreErrors: + # Store op as data, such that consumers of commands do not have to + # deal with incorrect number of arguments. + result.append(("", args)) + result.append(("", [op])) + else: + raise + + func = getattr(mapping, op, None) + if func is None: + result.append((op, args)) + continue + try: + for command in func(args): + result.append(command) + except ValueError: + if ignoreErrors: + # Store op as data, such that consumers of commands do not have to + # deal with incorrect number of arguments. + result.append(("", args)) + result.append(("", [op])) + else: + raise + return result + + +def generalizeProgram(program, getNumRegions=None, **kwargs): + return commandsToProgram( + generalizeCommands(programToCommands(program, getNumRegions), **kwargs) + ) + + +def _categorizeVector(v): + """ + Takes X,Y vector v and returns one of r, h, v, or 0 depending on which + of X and/or Y are zero, plus tuple of nonzero ones. If both are zero, + it returns a single zero still. + + >>> _categorizeVector((0,0)) + ('0', (0,)) + >>> _categorizeVector((1,0)) + ('h', (1,)) + >>> _categorizeVector((0,2)) + ('v', (2,)) + >>> _categorizeVector((1,2)) + ('r', (1, 2)) + """ + if not v[0]: + if not v[1]: + return "0", v[:1] + else: + return "v", v[1:] + else: + if not v[1]: + return "h", v[:1] + else: + return "r", v + + +def _mergeCategories(a, b): + if a == "0": + return b + if b == "0": + return a + if a == b: + return a + return None + + +def _negateCategory(a): + if a == "h": + return "v" + if a == "v": + return "h" + assert a in "0r" + return a + + +def _convertToBlendCmds(args): + # return a list of blend commands, and + # the remaining non-blended args, if any. + num_args = len(args) + stack_use = 0 + new_args = [] + i = 0 + while i < num_args: + arg = args[i] + i += 1 + if not isinstance(arg, list): + new_args.append(arg) + stack_use += 1 + else: + prev_stack_use = stack_use + # The arg is a tuple of blend values. + # These are each (master 0,delta 1..delta n, 1) + # Combine as many successive tuples as we can, + # up to the max stack limit. + num_sources = len(arg) - 1 + blendlist = [arg] + stack_use += 1 + num_sources # 1 for the num_blends arg + + # if we are here, max stack is the CFF2 max stack. + # I use the CFF2 max stack limit here rather than + # the 'maxstack' chosen by the client, as the default + # maxstack may have been used unintentionally. For all + # the other operators, this just produces a little less + # optimization, but here it puts a hard (and low) limit + # on the number of source fonts that can be used. + # + # Make sure the stack depth does not exceed (maxstack - 1), so + # that subroutinizer can insert subroutine calls at any point. + while ( + (i < num_args) + and isinstance(args[i], list) + and stack_use + num_sources < maxStackLimit + ): + blendlist.append(args[i]) + i += 1 + stack_use += num_sources + # blendList now contains as many single blend tuples as can be + # combined without exceeding the CFF2 stack limit. + num_blends = len(blendlist) + # append the 'num_blends' default font values + blend_args = [] + for arg in blendlist: + blend_args.append(arg[0]) + for arg in blendlist: + assert arg[-1] == 1 + blend_args.extend(arg[1:-1]) + blend_args.append(num_blends) + new_args.append(blend_args) + stack_use = prev_stack_use + num_blends + + return new_args + + +def _addArgs(a, b): + if isinstance(b, list): + if isinstance(a, list): + if len(a) != len(b) or a[-1] != b[-1]: + raise ValueError() + return [_addArgs(va, vb) for va, vb in zip(a[:-1], b[:-1])] + [a[-1]] + else: + a, b = b, a + if isinstance(a, list): + assert a[-1] == 1 + return [_addArgs(a[0], b)] + a[1:] + return a + b + + +def _argsStackUse(args): + stackLen = 0 + maxLen = 0 + for arg in args: + if type(arg) is list: + # Blended arg + maxLen = max(maxLen, stackLen + _argsStackUse(arg)) + stackLen += arg[-1] + else: + stackLen += 1 + return max(stackLen, maxLen) + + +def specializeCommands( + commands, + ignoreErrors=False, + generalizeFirst=True, + preserveTopology=False, + maxstack=48, +): + # We perform several rounds of optimizations. They are carefully ordered and are: + # + # 0. Generalize commands. + # This ensures that they are in our expected simple form, with each line/curve only + # having arguments for one segment, and using the generic form (rlineto/rrcurveto). + # If caller is sure the input is in this form, they can turn off generalization to + # save time. + # + # 1. Combine successive rmoveto operations. + # + # 2. Specialize rmoveto/rlineto/rrcurveto operators into horizontal/vertical variants. + # We specialize into some, made-up, variants as well, which simplifies following + # passes. + # + # 3. Merge or delete redundant operations, to the extent requested. + # OpenType spec declares point numbers in CFF undefined. As such, we happily + # change topology. If client relies on point numbers (in GPOS anchors, or for + # hinting purposes(what?)) they can turn this off. + # + # 4. Peephole optimization to revert back some of the h/v variants back into their + # original "relative" operator (rline/rrcurveto) if that saves a byte. + # + # 5. Combine adjacent operators when possible, minding not to go over max stack size. + # + # 6. Resolve any remaining made-up operators into real operators. + # + # I have convinced myself that this produces optimal bytecode (except for, possibly + # one byte each time maxstack size prohibits combining.) YMMV, but you'd be wrong. :-) + # A dynamic-programming approach can do the same but would be significantly slower. + # + # 7. For any args which are blend lists, convert them to a blend command. + + # 0. Generalize commands. + if generalizeFirst: + commands = generalizeCommands(commands, ignoreErrors=ignoreErrors) + else: + commands = list(commands) # Make copy since we modify in-place later. + + # 1. Combine successive rmoveto operations. + for i in range(len(commands) - 1, 0, -1): + if "rmoveto" == commands[i][0] == commands[i - 1][0]: + v1, v2 = commands[i - 1][1], commands[i][1] + commands[i - 1] = ( + "rmoveto", + [_addArgs(v1[0], v2[0]), _addArgs(v1[1], v2[1])], + ) + del commands[i] + + # 2. Specialize rmoveto/rlineto/rrcurveto operators into horizontal/vertical variants. + # + # We, in fact, specialize into more, made-up, variants that special-case when both + # X and Y components are zero. This simplifies the following optimization passes. + # This case is rare, but OCD does not let me skip it. + # + # After this round, we will have four variants that use the following mnemonics: + # + # - 'r' for relative, ie. non-zero X and non-zero Y, + # - 'h' for horizontal, ie. zero X and non-zero Y, + # - 'v' for vertical, ie. non-zero X and zero Y, + # - '0' for zeros, ie. zero X and zero Y. + # + # The '0' pseudo-operators are not part of the spec, but help simplify the following + # optimization rounds. We resolve them at the end. So, after this, we will have four + # moveto and four lineto variants: + # + # - 0moveto, 0lineto + # - hmoveto, hlineto + # - vmoveto, vlineto + # - rmoveto, rlineto + # + # and sixteen curveto variants. For example, a '0hcurveto' operator means a curve + # dx0,dy0,dx1,dy1,dx2,dy2,dx3,dy3 where dx0, dx1, and dy3 are zero but not dx3. + # An 'rvcurveto' means dx3 is zero but not dx0,dy0,dy3. + # + # There are nine different variants of curves without the '0'. Those nine map exactly + # to the existing curve variants in the spec: rrcurveto, and the four variants hhcurveto, + # vvcurveto, hvcurveto, and vhcurveto each cover two cases, one with an odd number of + # arguments and one without. Eg. an hhcurveto with an extra argument (odd number of + # arguments) is in fact an rhcurveto. The operators in the spec are designed such that + # all four of rhcurveto, rvcurveto, hrcurveto, and vrcurveto are encodable for one curve. + # + # Of the curve types with '0', the 00curveto is equivalent to a lineto variant. The rest + # of the curve types with a 0 need to be encoded as a h or v variant. Ie. a '0' can be + # thought of a "don't care" and can be used as either an 'h' or a 'v'. As such, we always + # encode a number 0 as argument when we use a '0' variant. Later on, we can just substitute + # the '0' with either 'h' or 'v' and it works. + # + # When we get to curve splines however, things become more complicated... XXX finish this. + # There's one more complexity with splines. If one side of the spline is not horizontal or + # vertical (or zero), ie. if it's 'r', then it limits which spline types we can encode. + # Only hhcurveto and vvcurveto operators can encode a spline starting with 'r', and + # only hvcurveto and vhcurveto operators can encode a spline ending with 'r'. + # This limits our merge opportunities later. + # + for i in range(len(commands)): + op, args = commands[i] + + if op in {"rmoveto", "rlineto"}: + c, args = _categorizeVector(args) + commands[i] = c + op[1:], args + continue + + if op == "rrcurveto": + c1, args1 = _categorizeVector(args[:2]) + c2, args2 = _categorizeVector(args[-2:]) + commands[i] = c1 + c2 + "curveto", args1 + args[2:4] + args2 + continue + + # 3. Merge or delete redundant operations, to the extent requested. + # + # TODO + # A 0moveto that comes before all other path operations can be removed. + # though I find conflicting evidence for this. + # + # TODO + # "If hstem and vstem hints are both declared at the beginning of a + # CharString, and this sequence is followed directly by the hintmask or + # cntrmask operators, then the vstem hint operator (or, if applicable, + # the vstemhm operator) need not be included." + # + # "The sequence and form of a CFF2 CharString program may be represented as: + # {hs* vs* cm* hm* mt subpath}? {mt subpath}*" + # + # https://www.microsoft.com/typography/otspec/cff2charstr.htm#section3.1 + # + # For Type2 CharStrings the sequence is: + # w? {hs* vs* cm* hm* mt subpath}? {mt subpath}* endchar" + + # Some other redundancies change topology (point numbers). + if not preserveTopology: + for i in range(len(commands) - 1, -1, -1): + op, args = commands[i] + + # A 00curveto is demoted to a (specialized) lineto. + if op == "00curveto": + assert len(args) == 4 + c, args = _categorizeVector(args[1:3]) + op = c + "lineto" + commands[i] = op, args + # and then... + + # A 0lineto can be deleted. + if op == "0lineto": + del commands[i] + continue + + # Merge adjacent hlineto's and vlineto's. + # In CFF2 charstrings from variable fonts, each + # arg item may be a list of blendable values, one from + # each source font. + if i and op in {"hlineto", "vlineto"} and (op == commands[i - 1][0]): + _, other_args = commands[i - 1] + assert len(args) == 1 and len(other_args) == 1 + try: + new_args = [_addArgs(args[0], other_args[0])] + except ValueError: + continue + commands[i - 1] = (op, new_args) + del commands[i] + continue + + # 4. Peephole optimization to revert back some of the h/v variants back into their + # original "relative" operator (rline/rrcurveto) if that saves a byte. + for i in range(1, len(commands) - 1): + op, args = commands[i] + prv, nxt = commands[i - 1][0], commands[i + 1][0] + + if op in {"0lineto", "hlineto", "vlineto"} and prv == nxt == "rlineto": + assert len(args) == 1 + args = [0, args[0]] if op[0] == "v" else [args[0], 0] + commands[i] = ("rlineto", args) + continue + + if op[2:] == "curveto" and len(args) == 5 and prv == nxt == "rrcurveto": + assert (op[0] == "r") ^ (op[1] == "r") + if op[0] == "v": + pos = 0 + elif op[0] != "r": + pos = 1 + elif op[1] == "v": + pos = 4 + else: + pos = 5 + # Insert, while maintaining the type of args (can be tuple or list). + args = args[:pos] + type(args)((0,)) + args[pos:] + commands[i] = ("rrcurveto", args) + continue + + # 5. Combine adjacent operators when possible, minding not to go over max stack size. + stackUse = _argsStackUse(commands[-1][1]) if commands else 0 + for i in range(len(commands) - 1, 0, -1): + op1, args1 = commands[i - 1] + op2, args2 = commands[i] + new_op = None + + # Merge logic... + if {op1, op2} <= {"rlineto", "rrcurveto"}: + if op1 == op2: + new_op = op1 + else: + l = len(args2) + if op2 == "rrcurveto" and l == 6: + new_op = "rlinecurve" + elif l == 2: + new_op = "rcurveline" + + elif (op1, op2) in {("rlineto", "rlinecurve"), ("rrcurveto", "rcurveline")}: + new_op = op2 + + elif {op1, op2} == {"vlineto", "hlineto"}: + new_op = op1 + + elif "curveto" == op1[2:] == op2[2:]: + d0, d1 = op1[:2] + d2, d3 = op2[:2] + + if d1 == "r" or d2 == "r" or d0 == d3 == "r": + continue + + d = _mergeCategories(d1, d2) + if d is None: + continue + if d0 == "r": + d = _mergeCategories(d, d3) + if d is None: + continue + new_op = "r" + d + "curveto" + elif d3 == "r": + d0 = _mergeCategories(d0, _negateCategory(d)) + if d0 is None: + continue + new_op = d0 + "r" + "curveto" + else: + d0 = _mergeCategories(d0, d3) + if d0 is None: + continue + new_op = d0 + d + "curveto" + + # Make sure the stack depth does not exceed (maxstack - 1), so + # that subroutinizer can insert subroutine calls at any point. + args1StackUse = _argsStackUse(args1) + combinedStackUse = max(args1StackUse, len(args1) + stackUse) + if new_op and combinedStackUse < maxstack: + commands[i - 1] = (new_op, args1 + args2) + del commands[i] + stackUse = combinedStackUse + else: + stackUse = args1StackUse + + # 6. Resolve any remaining made-up operators into real operators. + for i in range(len(commands)): + op, args = commands[i] + + if op in {"0moveto", "0lineto"}: + commands[i] = "h" + op[1:], args + continue + + if op[2:] == "curveto" and op[:2] not in {"rr", "hh", "vv", "vh", "hv"}: + l = len(args) + + op0, op1 = op[:2] + if (op0 == "r") ^ (op1 == "r"): + assert l % 2 == 1 + if op0 == "0": + op0 = "h" + if op1 == "0": + op1 = "h" + if op0 == "r": + op0 = op1 + if op1 == "r": + op1 = _negateCategory(op0) + assert {op0, op1} <= {"h", "v"}, (op0, op1) + + if l % 2: + if op0 != op1: # vhcurveto / hvcurveto + if (op0 == "h") ^ (l % 8 == 1): + # Swap last two args order + args = args[:-2] + args[-1:] + args[-2:-1] + else: # hhcurveto / vvcurveto + if op0 == "h": # hhcurveto + # Swap first two args order + args = args[1:2] + args[:1] + args[2:] + + commands[i] = op0 + op1 + "curveto", args + continue + + # 7. For any series of args which are blend lists, convert the series to a single blend arg. + for i in range(len(commands)): + op, args = commands[i] + if any(isinstance(arg, list) for arg in args): + commands[i] = op, _convertToBlendCmds(args) + + return commands + + +def specializeProgram(program, getNumRegions=None, **kwargs): + return commandsToProgram( + specializeCommands(programToCommands(program, getNumRegions), **kwargs) + ) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) == 1: + import doctest + + sys.exit(doctest.testmod().failed) + + import argparse + + parser = argparse.ArgumentParser( + "fonttools cffLib.specializer", + description="CFF CharString generalizer/specializer", + ) + parser.add_argument("program", metavar="command", nargs="*", help="Commands.") + parser.add_argument( + "--num-regions", + metavar="NumRegions", + nargs="*", + default=None, + help="Number of variable-font regions for blend opertaions.", + ) + parser.add_argument( + "--font", + metavar="FONTFILE", + default=None, + help="CFF2 font to specialize.", + ) + parser.add_argument( + "-o", + "--output-file", + type=str, + help="Output font file name.", + ) + + options = parser.parse_args(sys.argv[1:]) + + if options.program: + getNumRegions = ( + None + if options.num_regions is None + else lambda vsIndex: int( + options.num_regions[0 if vsIndex is None else vsIndex] + ) + ) + + program = stringToProgram(options.program) + print("Program:") + print(programToString(program)) + commands = programToCommands(program, getNumRegions) + print("Commands:") + print(commands) + program2 = commandsToProgram(commands) + print("Program from commands:") + print(programToString(program2)) + assert program == program2 + print("Generalized program:") + print(programToString(generalizeProgram(program, getNumRegions))) + print("Specialized program:") + print(programToString(specializeProgram(program, getNumRegions))) + + if options.font: + from fontTools.ttLib import TTFont + + font = TTFont(options.font) + cff2 = font["CFF2"].cff.topDictIndex[0] + charstrings = cff2.CharStrings + for glyphName in charstrings.keys(): + charstring = charstrings[glyphName] + charstring.decompile() + getNumRegions = charstring.private.getNumRegions + charstring.program = specializeProgram( + charstring.program, getNumRegions, maxstack=maxStackLimit + ) + + if options.output_file is None: + from fontTools.misc.cliTools import makeOutputFileName + + outfile = makeOutputFileName( + options.font, overWrite=True, suffix=".specialized" + ) + else: + outfile = options.output_file + if outfile: + print("Saving", outfile) + font.save(outfile) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/transforms.py b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..56d90cf2e2944467aacedc79e92bcb9fbd373e5d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/fontTools/cffLib/transforms.py @@ -0,0 +1,495 @@ +from fontTools.misc.psCharStrings import ( + SimpleT2Decompiler, + T2WidthExtractor, + calcSubrBias, +) + + +def _uniq_sort(l): + return sorted(set(l)) + + +class StopHintCountEvent(Exception): + pass + + +class _DesubroutinizingT2Decompiler(SimpleT2Decompiler): + stop_hintcount_ops = ( + "op_hintmask", + "op_cntrmask", + "op_rmoveto", + "op_hmoveto", + "op_vmoveto", + ) + + def __init__(self, localSubrs, globalSubrs, private=None): + SimpleT2Decompiler.__init__(self, localSubrs, globalSubrs, private) + + def execute(self, charString): + self.need_hintcount = True # until proven otherwise + for op_name in self.stop_hintcount_ops: + setattr(self, op_name, self.stop_hint_count) + + if hasattr(charString, "_desubroutinized"): + # If a charstring has already been desubroutinized, we will still + # need to execute it if we need to count hints in order to + # compute the byte length for mask arguments, and haven't finished + # counting hints pairs. + if self.need_hintcount and self.callingStack: + try: + SimpleT2Decompiler.execute(self, charString) + except StopHintCountEvent: + del self.callingStack[-1] + return + + charString._patches = [] + SimpleT2Decompiler.execute(self, charString) + desubroutinized = charString.program[:] + for idx, expansion in reversed(charString._patches): + assert idx >= 2 + assert desubroutinized[idx - 1] in [ + "callsubr", + "callgsubr", + ], desubroutinized[idx - 1] + assert type(desubroutinized[idx - 2]) == int + if expansion[-1] == "return": + expansion = expansion[:-1] + desubroutinized[idx - 2 : idx] = expansion + if not self.private.in_cff2: + if "endchar" in desubroutinized: + # Cut off after first endchar + desubroutinized = desubroutinized[ + : desubroutinized.index("endchar") + 1 + ] + + charString._desubroutinized = desubroutinized + del charString._patches + + def op_callsubr(self, index): + subr = self.localSubrs[self.operandStack[-1] + self.localBias] + SimpleT2Decompiler.op_callsubr(self, index) + self.processSubr(index, subr) + + def op_callgsubr(self, index): + subr = self.globalSubrs[self.operandStack[-1] + self.globalBias] + SimpleT2Decompiler.op_callgsubr(self, index) + self.processSubr(index, subr) + + def stop_hint_count(self, *args): + self.need_hintcount = False + for op_name in self.stop_hintcount_ops: + setattr(self, op_name, None) + cs = self.callingStack[-1] + if hasattr(cs, "_desubroutinized"): + raise StopHintCountEvent() + + def op_hintmask(self, index): + SimpleT2Decompiler.op_hintmask(self, index) + if self.need_hintcount: + self.stop_hint_count() + + def processSubr(self, index, subr): + cs = self.callingStack[-1] + if not hasattr(cs, "_desubroutinized"): + cs._patches.append((index, subr._desubroutinized)) + + +def desubroutinizeCharString(cs): + """Desubroutinize a charstring in-place.""" + cs.decompile() + subrs = getattr(cs.private, "Subrs", []) + decompiler = _DesubroutinizingT2Decompiler(subrs, cs.globalSubrs, cs.private) + decompiler.execute(cs) + cs.program = cs._desubroutinized + del cs._desubroutinized + + +def desubroutinize(cff): + for fontName in cff.fontNames: + font = cff[fontName] + cs = font.CharStrings + for c in cs.values(): + desubroutinizeCharString(c) + # Delete all the local subrs + if hasattr(font, "FDArray"): + for fd in font.FDArray: + pd = fd.Private + if hasattr(pd, "Subrs"): + del pd.Subrs + if "Subrs" in pd.rawDict: + del pd.rawDict["Subrs"] + else: + pd = font.Private + if hasattr(pd, "Subrs"): + del pd.Subrs + if "Subrs" in pd.rawDict: + del pd.rawDict["Subrs"] + # as well as the global subrs + cff.GlobalSubrs.clear() + + +class _MarkingT2Decompiler(SimpleT2Decompiler): + def __init__(self, localSubrs, globalSubrs, private): + SimpleT2Decompiler.__init__(self, localSubrs, globalSubrs, private) + for subrs in [localSubrs, globalSubrs]: + if subrs and not hasattr(subrs, "_used"): + subrs._used = set() + + def op_callsubr(self, index): + self.localSubrs._used.add(self.operandStack[-1] + self.localBias) + SimpleT2Decompiler.op_callsubr(self, index) + + def op_callgsubr(self, index): + self.globalSubrs._used.add(self.operandStack[-1] + self.globalBias) + SimpleT2Decompiler.op_callgsubr(self, index) + + +class _DehintingT2Decompiler(T2WidthExtractor): + class Hints(object): + def __init__(self): + # Whether calling this charstring produces any hint stems + # Note that if a charstring starts with hintmask, it will + # have has_hint set to True, because it *might* produce an + # implicit vstem if called under certain conditions. + self.has_hint = False + # Index to start at to drop all hints + self.last_hint = 0 + # Index up to which we know more hints are possible. + # Only relevant if status is 0 or 1. + self.last_checked = 0 + # The status means: + # 0: after dropping hints, this charstring is empty + # 1: after dropping hints, there may be more hints + # continuing after this, or there might be + # other things. Not clear yet. + # 2: no more hints possible after this charstring + self.status = 0 + # Has hintmask instructions; not recursive + self.has_hintmask = False + # List of indices of calls to empty subroutines to remove. + self.deletions = [] + + pass + + def __init__( + self, css, localSubrs, globalSubrs, nominalWidthX, defaultWidthX, private=None + ): + self._css = css + T2WidthExtractor.__init__( + self, localSubrs, globalSubrs, nominalWidthX, defaultWidthX + ) + self.private = private + + def execute(self, charString): + old_hints = charString._hints if hasattr(charString, "_hints") else None + charString._hints = self.Hints() + + T2WidthExtractor.execute(self, charString) + + hints = charString._hints + + if hints.has_hint or hints.has_hintmask: + self._css.add(charString) + + if hints.status != 2: + # Check from last_check, make sure we didn't have any operators. + for i in range(hints.last_checked, len(charString.program) - 1): + if isinstance(charString.program[i], str): + hints.status = 2 + break + else: + hints.status = 1 # There's *something* here + hints.last_checked = len(charString.program) + + if old_hints: + assert hints.__dict__ == old_hints.__dict__ + + def op_callsubr(self, index): + subr = self.localSubrs[self.operandStack[-1] + self.localBias] + T2WidthExtractor.op_callsubr(self, index) + self.processSubr(index, subr) + + def op_callgsubr(self, index): + subr = self.globalSubrs[self.operandStack[-1] + self.globalBias] + T2WidthExtractor.op_callgsubr(self, index) + self.processSubr(index, subr) + + def op_hstem(self, index): + T2WidthExtractor.op_hstem(self, index) + self.processHint(index) + + def op_vstem(self, index): + T2WidthExtractor.op_vstem(self, index) + self.processHint(index) + + def op_hstemhm(self, index): + T2WidthExtractor.op_hstemhm(self, index) + self.processHint(index) + + def op_vstemhm(self, index): + T2WidthExtractor.op_vstemhm(self, index) + self.processHint(index) + + def op_hintmask(self, index): + rv = T2WidthExtractor.op_hintmask(self, index) + self.processHintmask(index) + return rv + + def op_cntrmask(self, index): + rv = T2WidthExtractor.op_cntrmask(self, index) + self.processHintmask(index) + return rv + + def processHintmask(self, index): + cs = self.callingStack[-1] + hints = cs._hints + hints.has_hintmask = True + if hints.status != 2: + # Check from last_check, see if we may be an implicit vstem + for i in range(hints.last_checked, index - 1): + if isinstance(cs.program[i], str): + hints.status = 2 + break + else: + # We are an implicit vstem + hints.has_hint = True + hints.last_hint = index + 1 + hints.status = 0 + hints.last_checked = index + 1 + + def processHint(self, index): + cs = self.callingStack[-1] + hints = cs._hints + hints.has_hint = True + hints.last_hint = index + hints.last_checked = index + + def processSubr(self, index, subr): + cs = self.callingStack[-1] + hints = cs._hints + subr_hints = subr._hints + + # Check from last_check, make sure we didn't have + # any operators. + if hints.status != 2: + for i in range(hints.last_checked, index - 1): + if isinstance(cs.program[i], str): + hints.status = 2 + break + hints.last_checked = index + + if hints.status != 2: + if subr_hints.has_hint: + hints.has_hint = True + + # Decide where to chop off from + if subr_hints.status == 0: + hints.last_hint = index + else: + hints.last_hint = index - 2 # Leave the subr call in + + elif subr_hints.status == 0: + hints.deletions.append(index) + + hints.status = max(hints.status, subr_hints.status) + + +def _cs_subset_subroutines(charstring, subrs, gsubrs): + p = charstring.program + for i in range(1, len(p)): + if p[i] == "callsubr": + assert isinstance(p[i - 1], int) + p[i - 1] = subrs._used.index(p[i - 1] + subrs._old_bias) - subrs._new_bias + elif p[i] == "callgsubr": + assert isinstance(p[i - 1], int) + p[i - 1] = ( + gsubrs._used.index(p[i - 1] + gsubrs._old_bias) - gsubrs._new_bias + ) + + +def _cs_drop_hints(charstring): + hints = charstring._hints + + if hints.deletions: + p = charstring.program + for idx in reversed(hints.deletions): + del p[idx - 2 : idx] + + if hints.has_hint: + assert not hints.deletions or hints.last_hint <= hints.deletions[0] + charstring.program = charstring.program[hints.last_hint :] + if not charstring.program: + # TODO CFF2 no need for endchar. + charstring.program.append("endchar") + if hasattr(charstring, "width"): + # Insert width back if needed + if charstring.width != charstring.private.defaultWidthX: + # For CFF2 charstrings, this should never happen + assert ( + charstring.private.defaultWidthX is not None + ), "CFF2 CharStrings must not have an initial width value" + charstring.program.insert( + 0, charstring.width - charstring.private.nominalWidthX + ) + + if hints.has_hintmask: + i = 0 + p = charstring.program + while i < len(p): + if p[i] in ["hintmask", "cntrmask"]: + assert i + 1 <= len(p) + del p[i : i + 2] + continue + i += 1 + + assert len(charstring.program) + + del charstring._hints + + +def remove_hints(cff, *, removeUnusedSubrs: bool = True): + for fontname in cff.keys(): + font = cff[fontname] + cs = font.CharStrings + # This can be tricky, but doesn't have to. What we do is: + # + # - Run all used glyph charstrings and recurse into subroutines, + # - For each charstring (including subroutines), if it has any + # of the hint stem operators, we mark it as such. + # Upon returning, for each charstring we note all the + # subroutine calls it makes that (recursively) contain a stem, + # - Dropping hinting then consists of the following two ops: + # * Drop the piece of the program in each charstring before the + # last call to a stem op or a stem-calling subroutine, + # * Drop all hintmask operations. + # - It's trickier... A hintmask right after hints and a few numbers + # will act as an implicit vstemhm. As such, we track whether + # we have seen any non-hint operators so far and do the right + # thing, recursively... Good luck understanding that :( + css = set() + for c in cs.values(): + c.decompile() + subrs = getattr(c.private, "Subrs", []) + decompiler = _DehintingT2Decompiler( + css, + subrs, + c.globalSubrs, + c.private.nominalWidthX, + c.private.defaultWidthX, + c.private, + ) + decompiler.execute(c) + c.width = decompiler.width + for charstring in css: + _cs_drop_hints(charstring) + del css + + # Drop font-wide hinting values + all_privs = [] + if hasattr(font, "FDArray"): + all_privs.extend(fd.Private for fd in font.FDArray) + else: + all_privs.append(font.Private) + for priv in all_privs: + for k in [ + "BlueValues", + "OtherBlues", + "FamilyBlues", + "FamilyOtherBlues", + "BlueScale", + "BlueShift", + "BlueFuzz", + "StemSnapH", + "StemSnapV", + "StdHW", + "StdVW", + "ForceBold", + "LanguageGroup", + "ExpansionFactor", + ]: + if hasattr(priv, k): + setattr(priv, k, None) + if removeUnusedSubrs: + remove_unused_subroutines(cff) + + +def _pd_delete_empty_subrs(private_dict): + if hasattr(private_dict, "Subrs") and not private_dict.Subrs: + if "Subrs" in private_dict.rawDict: + del private_dict.rawDict["Subrs"] + del private_dict.Subrs + + +def remove_unused_subroutines(cff): + for fontname in cff.keys(): + font = cff[fontname] + cs = font.CharStrings + # Renumber subroutines to remove unused ones + + # Mark all used subroutines + for c in cs.values(): + subrs = getattr(c.private, "Subrs", []) + decompiler = _MarkingT2Decompiler(subrs, c.globalSubrs, c.private) + decompiler.execute(c) + + all_subrs = [font.GlobalSubrs] + if hasattr(font, "FDArray"): + all_subrs.extend( + fd.Private.Subrs + for fd in font.FDArray + if hasattr(fd.Private, "Subrs") and fd.Private.Subrs + ) + elif hasattr(font.Private, "Subrs") and font.Private.Subrs: + all_subrs.append(font.Private.Subrs) + + subrs = set(subrs) # Remove duplicates + + # Prepare + for subrs in all_subrs: + if not hasattr(subrs, "_used"): + subrs._used = set() + subrs._used = _uniq_sort(subrs._used) + subrs._old_bias = calcSubrBias(subrs) + subrs._new_bias = calcSubrBias(subrs._used) + + # Renumber glyph charstrings + for c in cs.values(): + subrs = getattr(c.private, "Subrs", None) + _cs_subset_subroutines(c, subrs, font.GlobalSubrs) + + # Renumber subroutines themselves + for subrs in all_subrs: + if subrs == font.GlobalSubrs: + if not hasattr(font, "FDArray") and hasattr(font.Private, "Subrs"): + local_subrs = font.Private.Subrs + elif ( + hasattr(font, "FDArray") + and len(font.FDArray) == 1 + and hasattr(font.FDArray[0].Private, "Subrs") + ): + # Technically we shouldn't do this. But I've run into fonts that do it. + local_subrs = font.FDArray[0].Private.Subrs + else: + local_subrs = None + else: + local_subrs = subrs + + subrs.items = [subrs.items[i] for i in subrs._used] + if hasattr(subrs, "file"): + del subrs.file + if hasattr(subrs, "offsets"): + del subrs.offsets + + for subr in subrs.items: + _cs_subset_subroutines(subr, local_subrs, font.GlobalSubrs) + + # Delete local SubrsIndex if empty + if hasattr(font, "FDArray"): + for fd in font.FDArray: + _pd_delete_empty_subrs(fd.Private) + else: + _pd_delete_empty_subrs(font.Private) + + # Cleanup + for subrs in all_subrs: + del subrs._used, subrs._old_bias, subrs._new_bias diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/api.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/api.py new file mode 100644 index 0000000000000000000000000000000000000000..9e9ffddfd0a5a053b092e36733f4c78c3b9f58b2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/api.py @@ -0,0 +1,2975 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""JAX user-facing transformations and utilities. + +The transformations here mostly wrap internal transformations, providing +convenience flags to control behavior and handling Python containers of +arguments and outputs. The Python containers handled are pytrees (see +tree_util.py), which include nested tuples/lists/dicts, where the leaves are +arrays. +""" +from __future__ import annotations + +import collections +from collections.abc import Generator, Hashable, Iterable, Sequence +from functools import partial, lru_cache +import inspect +import math +import typing +from typing import (Any, Callable, Literal, NamedTuple, TypeVar, overload, + cast) +import weakref + +import numpy as np +from contextlib import contextmanager, ExitStack + +from jax._src import linear_util as lu +from jax._src import stages +from jax._src.tree_util import ( + tree_map, tree_flatten, tree_unflatten, tree_structure, tree_transpose, + tree_leaves, Partial, PyTreeDef, all_leaves, keystr, broadcast_prefix, + prefix_errors, generate_key_paths) +from jax._src import api_util +from jax._src import config +from jax._src import core +from jax._src import dispatch +from jax._src import effects +from jax._src import array +from jax._src import basearray +from jax._src import dtypes +from jax._src import sharding_impls +from jax._src import sharding_specs +from jax._src import source_info_util +from jax._src import traceback_util +from jax._src import pjit +from jax._src import xla_bridge as xb +from jax._src.core import eval_jaxpr, ShapedArray, ConcreteArray +from jax._src.api_util import ( + flatten_fun, flatten_fun_nokwargs, flatten_fun_nokwargs2, argnums_partial, + argnums_partial_except, flatten_axes, donation_vector, + rebase_donate_argnums, _ensure_index, _ensure_index_tuple, + shaped_abstractify, apply_flat_fun_nokwargs, check_callable, debug_info, + result_paths, flat_out_axes, debug_info_final, fun_sourceinfo) +from jax._src.lax import lax as lax_internal +from jax._src.lib import jax_jit +from jax._src.lib import xla_client as xc +from jax._src.lib import pmap_lib +from jax._src.sharding import Sharding +from jax._src.sharding_impls import PmapSharding, TransferToMemoryKind +from jax._src.layout import Layout, AutoLayout +from jax._src.traceback_util import api_boundary +from jax._src import tree_util +from jax._src.util import (unzip2, safe_map, safe_zip, wrap_name, wraps, + split_list) +from jax._src import util + +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import partial_eval as pe +from jax._src.interpreters import pxla +from jax._src.interpreters import xla + + +traceback_util.register_exclusion(__file__) + +_dtype = partial(dtypes.dtype, canonicalize=True) + +AxisName = Hashable + +Device = xc.Device + +# These TypeVars are used below to express the fact that function types +# (i.e. call signatures) are invariant under the vmap transformation. +F = TypeVar("F", bound=Callable) +T = TypeVar("T") +U = TypeVar("U") + +map, unsafe_map = safe_map, map +zip, unsafe_zip = safe_zip, zip + + +def _nan_check_posthook(fun, args, kwargs, output): + """Hook function called by the C++ jit/pmap to perform NaN checking.""" + buffers = [] + for leaf in tree_leaves(output): + if hasattr(leaf, "addressable_shards"): + buffers.extend([shard.data for shard in leaf.addressable_shards]) + + try: + dispatch.check_special(pjit.pjit_p.name, buffers) + except FloatingPointError: + # compiled_fun can only raise in this case + assert config.debug_nans.value or config.debug_infs.value + print("Invalid nan value encountered in the output of a C++-jit/pmap " + "function. Calling the de-optimized version.") + fun._cache_miss(*args, **kwargs)[0] # probably won't return + +def _update_debug_special_global(_): + if config._read("jax_debug_nans") or config._read("jax_debug_infs"): + jax_jit.global_state().post_hook = _nan_check_posthook + else: + jax_jit.global_state().post_hook = None + +def _update_debug_special_thread_local(_): + if (getattr(config._thread_local_state, "jax_debug_nans", False) or + getattr(config._thread_local_state, "jax_debug_infs", False)): + jax_jit.thread_local_state().post_hook = _nan_check_posthook + else: + jax_jit.thread_local_state().post_hook = None + +config.debug_nans._add_hooks(_update_debug_special_global, + _update_debug_special_thread_local) +config.debug_infs._add_hooks(_update_debug_special_global, + _update_debug_special_thread_local) + + +float0 = dtypes.float0 + + +def jit( + fun: Callable, + in_shardings=sharding_impls.UNSPECIFIED, + out_shardings=sharding_impls.UNSPECIFIED, + static_argnums: int | Sequence[int] | None = None, + static_argnames: str | Iterable[str] | None = None, + donate_argnums: int | Sequence[int] | None = None, + donate_argnames: str | Iterable[str] | None = None, + keep_unused: bool = False, + device: xc.Device | None = None, + backend: str | None = None, + inline: bool = False, + abstracted_axes: Any | None = None, +) -> pjit.JitWrapped: + """Sets up ``fun`` for just-in-time compilation with XLA. + + Args: + fun: Function to be jitted. ``fun`` should be a pure function. + + The arguments and return value of ``fun`` should be arrays, scalar, or + (nested) standard Python containers (tuple/list/dict) thereof. Positional + arguments indicated by ``static_argnums`` can be any hashable type. Static + arguments are included as part of a compilation cache key, which is why + hash and equality operators must be defined. JAX keeps a weak reference to + ``fun`` for use as a compilation cache key, so the object ``fun`` must be + weakly-referenceable. + in_shardings: optional, a :py:class:`Sharding` or pytree with + :py:class:`Sharding` leaves and structure that is a tree prefix of the + positional arguments tuple to ``fun``. If provided, the positional + arguments passed to ``fun`` must have shardings that are compatible with + ``in_shardings`` or an error is raised, and the compiled computation has + input shardings corresponding to ``in_shardings``. If not provided, the + compiled computation's input shardings are inferred from argument + shardings. + out_shardings: optional, a :py:class:`Sharding` or pytree with + :py:class:`Sharding` leaves and structure that is a tree prefix of the + output of ``fun``. If provided, it has the same effect as applying + corresponding :py:func:`jax.lax.with_sharding_constraint`s to the output + of ``fun``. + static_argnums: optional, an int or collection of ints that specify which + positional arguments to treat as static (trace- and compile-time + constant). + + Static arguments should be hashable, meaning both ``__hash__`` and + ``__eq__`` are implemented, and immutable. Otherwise they can be arbitrary + Python objects. Calling the jitted function with different values for + these constants will trigger recompilation. Arguments that are not + array-like or containers thereof must be marked as static. + + If neither ``static_argnums`` nor ``static_argnames`` is provided, no + arguments are treated as static. If ``static_argnums`` is not provided but + ``static_argnames`` is, or vice versa, JAX uses + :code:`inspect.signature(fun)` to find any positional arguments that + correspond to ``static_argnames`` + (or vice versa). If both ``static_argnums`` and ``static_argnames`` are + provided, ``inspect.signature`` is not used, and only actual + parameters listed in either ``static_argnums`` or ``static_argnames`` will + be treated as static. + static_argnames: optional, a string or collection of strings specifying + which named arguments to treat as static (compile-time constant). See the + comment on ``static_argnums`` for details. If not + provided but ``static_argnums`` is set, the default is based on calling + ``inspect.signature(fun)`` to find corresponding named arguments. + donate_argnums: optional, collection of integers to specify which positional + argument buffers can be overwritten by the computation and marked deleted + in the caller. It is safe to donate argument buffers if you no longer need + them once the computation has started. In some cases XLA can make use of + donated buffers to reduce the amount of memory needed to perform a + computation, for example recycling one of your input buffers to store a + result. You should not reuse buffers that you donate to a computation; JAX + will raise an error if you try to. By default, no argument buffers are + donated. + + If neither ``donate_argnums`` nor ``donate_argnames`` is provided, no + arguments are donated. If ``donate_argnums`` is not provided but + ``donate_argnames`` is, or vice versa, JAX uses + :code:`inspect.signature(fun)` to find any positional arguments that + correspond to ``donate_argnames`` + (or vice versa). If both ``donate_argnums`` and ``donate_argnames`` are + provided, ``inspect.signature`` is not used, and only actual + parameters listed in either ``donate_argnums`` or ``donate_argnames`` will + be donated. + + For more details on buffer donation see the + `FAQ `_. + donate_argnames: optional, a string or collection of strings specifying + which named arguments are donated to the computation. See the + comment on ``donate_argnums`` for details. If not + provided but ``donate_argnums`` is set, the default is based on calling + ``inspect.signature(fun)`` to find corresponding named arguments. + keep_unused: optional boolean. If `False` (the default), arguments that JAX + determines to be unused by `fun` *may* be dropped from resulting compiled + XLA executables. Such arguments will not be transferred to the device nor + provided to the underlying executable. If `True`, unused arguments will + not be pruned. + device: This is an experimental feature and the API is likely to change. + Optional, the Device the jitted function will run on. (Available devices + can be retrieved via :py:func:`jax.devices`.) The default is inherited + from XLA's DeviceAssignment logic and is usually to use + ``jax.devices()[0]``. + backend: This is an experimental feature and the API is likely to change. + Optional, a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or + ``'tpu'``. + inline: Optional boolean. Specify whether this function should be inlined + into enclosing jaxprs. Default False. + + Returns: + A wrapped version of ``fun``, set up for just-in-time compilation. + + Examples: + In the following example, ``selu`` can be compiled into a single fused kernel + by XLA: + + >>> import jax + >>> + >>> @jax.jit + ... def selu(x, alpha=1.67, lmbda=1.05): + ... return lmbda * jax.numpy.where(x > 0, x, alpha * jax.numpy.exp(x) - alpha) + >>> + >>> key = jax.random.key(0) + >>> x = jax.random.normal(key, (10,)) + >>> print(selu(x)) # doctest: +SKIP + [-0.54485 0.27744 -0.29255 -0.91421 -0.62452 -0.24748 + -0.85743 -0.78232 0.76827 0.59566 ] + + To pass arguments such as ``static_argnames`` when decorating a function, a + common pattern is to use :func:`functools.partial`: + + >>> from functools import partial + >>> + >>> @partial(jax.jit, static_argnames=['n']) + ... def g(x, n): + ... for i in range(n): + ... x = x ** 2 + ... return x + >>> + >>> g(jnp.arange(4), 3) + Array([ 0, 1, 256, 6561], dtype=int32) + """ + return pjit.make_jit( + fun, in_shardings, out_shardings, donate_argnums, donate_argnames, + static_argnums, static_argnames, device, backend, abstracted_axes, + keep_unused, inline, use_resource_env=False) + + +@contextmanager +def disable_jit(disable: bool = True): + """Context manager that disables :py:func:`jit` behavior under its dynamic context. + + For debugging it is useful to have a mechanism that disables :py:func:`jit` + everywhere in a dynamic context. Note that this not only disables explicit + uses of :func:`jit` by the user, but will also remove any implicit JIT compilation + used by the JAX library: this includes implicit JIT computation of `body` and + `cond` functions passed to higher-level primitives like :func:`~jax.lax.scan` and + :func:`~jax.lax.while_loop`, JIT used in implementations of :mod:`jax.numpy` functions, + and any other case where :func:`jit` is used within an API's implementation. + Note however that even under `disable_jit`, individual primitive operations + will still be compiled by XLA as in normal eager op-by-op execution. + + Values that have a data dependence on the arguments to a jitted function are + traced and abstracted. For example, an abstract value may be a + :py:class:`ShapedArray` instance, representing the set of all possible arrays + with a given shape and dtype, but not representing one concrete array with + specific values. You might notice those if you use a benign side-effecting + operation in a jitted function, like a print: + + >>> import jax + >>> + >>> @jax.jit + ... def f(x): + ... y = x * 2 + ... print("Value of y is", y) + ... return y + 3 + ... + >>> print(f(jax.numpy.array([1, 2, 3]))) # doctest:+ELLIPSIS + Value of y is Tracedwith + [5 7 9] + + Here ``y`` has been abstracted by :py:func:`jit` to a :py:class:`ShapedArray`, + which represents an array with a fixed shape and type but an arbitrary value. + The value of ``y`` is also traced. If we want to see a concrete value while + debugging, and avoid the tracer too, we can use the :py:func:`disable_jit` + context manager: + + >>> import jax + >>> + >>> with jax.disable_jit(): + ... print(f(jax.numpy.array([1, 2, 3]))) + ... + Value of y is [2 4 6] + [5 7 9] + """ + with config.disable_jit(disable): + yield + + +def xla_computation(fun: Callable, + static_argnums: int | Iterable[int] = (), + axis_env: Sequence[tuple[AxisName, int]] | None = None, + in_parts=None, out_parts=None, + backend: str | None = None, + tuple_args: bool = False, + instantiate_const_outputs: bool | None = None, + return_shape: bool = False, + donate_argnums: int | Iterable[int] = ()) -> Callable: + """Creates a function that produces its XLA computation given example args. + + Args: + fun: Function from which to form XLA computations. + static_argnums: See the :py:func:`jax.jit` docstring. + axis_env: Optional, a sequence of pairs where the first element is an axis + name and the second element is a positive integer representing the size of + the mapped axis with that name. This parameter is useful when lowering + functions that involve parallel communication collectives, and it + specifies the axis name/size environment that would be set up by + applications of :py:func:`jax.pmap`. See the examples below. + in_parts: Optional, how each argument to ``fun`` should be partitioned or + replicated. This is used to specify partitioned XLA computations, see + ``sharded_jit`` for more info. + out_parts: Optional, how each output of ``fun`` should be partitioned or + replicated. This is used to specify partitioned XLA computations, see + ``sharded_jit`` for more info. + backend: This is an experimental feature and the API is likely to change. + Optional, a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or + ``'tpu'``. + tuple_args: Optional bool, defaults to ``False``. If ``True``, the resulting + XLA computation will have a single tuple argument that is unpacked into + the specified function arguments. If `None`, tupling will be enabled when + there are more than 100 arguments, since some platforms have limits on + argument arity. + instantiate_const_outputs: Deprecated argument, does nothing. + return_shape: Optional boolean, defaults to ``False``. If ``True``, the + wrapped function returns a pair where the first element is the XLA + computation and the second element is a pytree with the same structure as + the output of ``fun`` and where the leaves are objects with ``shape``, + ``dtype``, and ``named_shape`` attributes representing the corresponding + types of the output leaves. + donate_argnums: Specify which arguments are "donated" to the computation. + It is safe to donate arguments if you no longer need them once the + computation has finished. In some cases XLA can make use of donated + buffers to reduce the amount of memory needed to perform a computation, + for example recycling one of your input buffers to store a result. You + should not reuse buffers that you donate to a computation, JAX will raise + an error if you try to. + + Returns: + A wrapped version of ``fun`` that when applied to example arguments returns + a built XLA Computation (see xla_client.py), from which representations of + the unoptimized XLA HLO computation can be extracted using methods like + ``as_hlo_text``, ``as_serialized_hlo_module_proto``, and + ``as_hlo_dot_graph``. If the argument ``return_shape`` is ``True``, then the + wrapped function returns a pair where the first element is the XLA + Computation and the second element is a pytree representing the structure, + shapes, dtypes, and named shapes of the output of ``fun``. + + Concrete example arguments are not always necessary. For those arguments not + indicated by ``static_argnums``, any object with ``shape`` and ``dtype`` + attributes is acceptable (excepting namedtuples, which are treated as Python + containers). + + For example: + + >>> import jax + >>> + >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x)) + >>> c = jax.xla_computation(f)(3.) + >>> print(c.as_hlo_text()) # doctest: +SKIP + HloModule xla_computation_f.6 + + ENTRY xla_computation_f.6 { + constant.2 = pred[] constant(false) + parameter.1 = f32[] parameter(0) + cosine.3 = f32[] cosine(parameter.1) + sine.4 = f32[] sine(cosine.3) + ROOT tuple.5 = (f32[]) tuple(sine.4) + } + + + + + Alternatively, the assignment to ``c`` above could be written: + + >>> import types + >>> scalar = types.SimpleNamespace(shape=(), dtype=np.dtype(np.float32)) + >>> c = jax.xla_computation(f)(scalar) + + + Here's an example that involves a parallel collective and axis name: + + >>> def f(x): return x - jax.lax.psum(x, 'i') + >>> c = jax.xla_computation(f, axis_env=[('i', 4)])(2) + >>> print(c.as_hlo_text()) # doctest: +SKIP + HloModule jaxpr_computation.9 + primitive_computation.3 { + parameter.4 = s32[] parameter(0) + parameter.5 = s32[] parameter(1) + ROOT add.6 = s32[] add(parameter.4, parameter.5) + } + ENTRY jaxpr_computation.9 { + tuple.1 = () tuple() + parameter.2 = s32[] parameter(0) + all-reduce.7 = s32[] all-reduce(parameter.2), replica_groups={{0,1,2,3}}, to_apply=primitive_computation.3 + ROOT subtract.8 = s32[] subtract(parameter.2, all-reduce.7) + } + + + + Notice the ``replica_groups`` that were generated. Here's an example that + generates more interesting ``replica_groups``: + + >>> from jax import lax + >>> def g(x): + ... rowsum = lax.psum(x, 'i') + ... colsum = lax.psum(x, 'j') + ... allsum = lax.psum(x, ('i', 'j')) + ... return rowsum, colsum, allsum + ... + >>> axis_env = [('i', 4), ('j', 2)] + >>> c = xla_computation(g, axis_env=axis_env)(5.) + >>> print(c.as_hlo_text()) # doctest: +SKIP + HloModule jaxpr_computation__1.19 + [removed uninteresting text here] + ENTRY jaxpr_computation__1.19 { + tuple.1 = () tuple() + parameter.2 = f32[] parameter(0) + all-reduce.7 = f32[] all-reduce(parameter.2), replica_groups={{0,2,4,6},{1,3,5,7}}, to_apply=primitive_computation__1.3 + all-reduce.12 = f32[] all-reduce(parameter.2), replica_groups={{0,1},{2,3},{4,5},{6,7}}, to_apply=primitive_computation__1.8 + all-reduce.17 = f32[] all-reduce(parameter.2), replica_groups={{0,1,2,3,4,5,6,7}}, to_apply=primitive_computation__1.13 + ROOT tuple.18 = (f32[], f32[], f32[]) tuple(all-reduce.7, all-reduce.12, all-reduce.17) + } + """ + if instantiate_const_outputs is not None: + raise ValueError( + "instantiate_const_outputs has been deprecated. Please use the ahead of" + " time APIs. You can read more here:" + " https://jax.readthedocs.io/en/latest/aot.html") + if in_parts is not None: + raise ValueError( + "in_parts has been deprecated. Please use the ahead of time APIs. You" + " can read more here: https://jax.readthedocs.io/en/latest/aot.html") + if out_parts is not None: + raise ValueError( + "out_parts has been deprecated. Please use the ahead of time APIs. You" + " can read more here: https://jax.readthedocs.io/en/latest/aot.html") + + check_callable(fun) + static_argnums = _ensure_index_tuple(static_argnums) + donate_argnums = _ensure_index_tuple(donate_argnums) + donate_argnums = rebase_donate_argnums(donate_argnums, static_argnums) + + fun_name = getattr(fun, "__name__", "unknown") + + platform = backend if backend is not None else xb.get_backend().platform + + def make_axis_env(nreps): + if axis_env is None: + return sharding_impls.AxisEnv(nreps, (), ()) + else: + nreps = nreps * math.prod(size for name, size in axis_env) + names, sizes = unzip2(axis_env) + return sharding_impls.AxisEnv(nreps, names, sizes) + + @wraps(fun) + @api_boundary + def computation_maker(*args, **kwargs): + if max(static_argnums + donate_argnums, default=-1) >= len(args): + raise ValueError(f"jitted function has {static_argnums=}, {donate_argnums=} but " + f"was called with only {len(args)} positional arguments.") + + f = lu.wrap_init(fun) + f, dyn_args = argnums_partial_except(f, static_argnums, args, allow_invalid=False) + args_flat, in_tree = tree_flatten((dyn_args, kwargs)) + if donate_argnums: + donated_invars = donation_vector(donate_argnums, (), in_tree) + else: + donated_invars = (False,) * len(args_flat) + + jaxtree_fun, out_tree = flatten_fun(f, in_tree) + avals = map(shaped_abstractify, args_flat) + with ExitStack() as stack: + for axis_name, size in axis_env or []: + stack.enter_context(core.extend_axis_env(axis_name, size, None)) + jaxpr, out_avals, consts, () = pe.trace_to_jaxpr_dynamic(jaxtree_fun, avals) + jaxpr = dispatch.apply_outfeed_rewriter(jaxpr) + if axis_env: + jaxpr = core.remove_named_axis_effects( + jaxpr, {axis_name for axis_name, _ in axis_env} + ) + axis_env_ = make_axis_env(dispatch.jaxpr_replicas(jaxpr)) + ordered_effects = list( + effects.ordered_effects.filter_in(jaxpr.effects)) + lowering_result = mlir.lower_jaxpr_to_module( + f"xla_computation_{fun_name}", + core.ClosedJaxpr(jaxpr, consts), + ordered_effects=ordered_effects, + backend_or_name=backend, + platforms=[platform], + axis_context=sharding_impls.ReplicaAxisContext(axis_env_), + name_stack=source_info_util.new_name_stack( + wrap_name(fun_name, "xla_computation")), + donated_args=donated_invars, + arg_shardings=None, + result_shardings=None, + lowering_parameters=mlir.LoweringParameters()) + + m = mlir.module_to_bytecode(lowering_result.module) + built = xc._xla.mlir.mlir_module_to_xla_computation( + m, use_tuple_args=tuple_args, return_tuple=True) + out_shapes_flat = [ + ShapeDtypeStruct(a.shape, a.dtype, a.named_shape) for a in out_avals] + out_shape = tree_unflatten(out_tree(), out_shapes_flat) + for out_aval in out_avals: + if not isinstance(out_aval, ShapedArray): + raise RuntimeError("As we want to propagate the weak_type, we need " + "to get a ShapedArray, otherwise this " + "information is lost") + + if return_shape: + return built, out_shape + else: + return built + + return computation_maker + +def grad(fun: Callable, argnums: int | Sequence[int] = 0, + has_aux: bool = False, holomorphic: bool = False, + allow_int: bool = False, + reduce_axes: Sequence[AxisName] = ()) -> Callable: + """Creates a function that evaluates the gradient of ``fun``. + + Args: + fun: Function to be differentiated. Its arguments at positions specified by + ``argnums`` should be arrays, scalars, or standard Python containers. + Argument arrays in the positions specified by ``argnums`` must be of + inexact (i.e., floating-point or complex) type. It + should return a scalar (which includes arrays with shape ``()`` but not + arrays with shape ``(1,)`` etc.) + argnums: Optional, integer or sequence of integers. Specifies which + positional argument(s) to differentiate with respect to (default 0). + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be + holomorphic. If True, inputs and outputs must be complex. Default False. + allow_int: Optional, bool. Whether to allow differentiating with + respect to integer valued inputs. The gradient of an integer input will + have a trivial vector-space dtype (float0). Default False. + + Returns: + A function with the same arguments as ``fun``, that evaluates the gradient + of ``fun``. If ``argnums`` is an integer then the gradient has the same + shape and type as the positional argument indicated by that integer. If + argnums is a tuple of integers, the gradient is a tuple of values with the + same shapes and types as the corresponding arguments. If ``has_aux`` is True + then a pair of (gradient, auxiliary_data) is returned. + + For example: + + >>> import jax + >>> + >>> grad_tanh = jax.grad(jax.numpy.tanh) + >>> print(grad_tanh(0.2)) + 0.961043 + """ + if reduce_axes: + raise NotImplementedError("reduce_axes argument to grad is deprecated") + del reduce_axes + value_and_grad_f = value_and_grad(fun, argnums, has_aux=has_aux, + holomorphic=holomorphic, + allow_int=allow_int) + + docstr = ("Gradient of {fun} with respect to positional argument(s) " + "{argnums}. Takes the same arguments as {fun} but returns the " + "gradient, which has the same shape as the arguments at " + "positions {argnums}.") + + @wraps(fun, docstr=docstr, argnums=argnums) + @api_boundary + def grad_f(*args, **kwargs): + _, g = value_and_grad_f(*args, **kwargs) + return g + + @wraps(fun, docstr=docstr, argnums=argnums) + @api_boundary + def grad_f_aux(*args, **kwargs): + (_, aux), g = value_and_grad_f(*args, **kwargs) + return g, aux + + return grad_f_aux if has_aux else grad_f + +def value_and_grad(fun: Callable, argnums: int | Sequence[int] = 0, + has_aux: bool = False, holomorphic: bool = False, + allow_int: bool = False, reduce_axes: Sequence[AxisName] = () + ) -> Callable[..., tuple[Any, Any]]: + """Create a function that evaluates both ``fun`` and the gradient of ``fun``. + + Args: + fun: Function to be differentiated. Its arguments at positions specified by + ``argnums`` should be arrays, scalars, or standard Python containers. It + should return a scalar (which includes arrays with shape ``()`` but not + arrays with shape ``(1,)`` etc.) + argnums: Optional, integer or sequence of integers. Specifies which + positional argument(s) to differentiate with respect to (default 0). + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be + holomorphic. If True, inputs and outputs must be complex. Default False. + allow_int: Optional, bool. Whether to allow differentiating with + respect to integer valued inputs. The gradient of an integer input will + have a trivial vector-space dtype (float0). Default False. + + Returns: + A function with the same arguments as ``fun`` that evaluates both ``fun`` + and the gradient of ``fun`` and returns them as a pair (a two-element + tuple). If ``argnums`` is an integer then the gradient has the same shape + and type as the positional argument indicated by that integer. If argnums is + a sequence of integers, the gradient is a tuple of values with the same + shapes and types as the corresponding arguments. If ``has_aux`` is True + then a tuple of ((value, auxiliary_data), gradient) is returned. + """ + if reduce_axes: + raise NotImplementedError("reduce_axes argument to grad is deprecated") + del reduce_axes + + docstr = ("Value and gradient of {fun} with respect to positional " + "argument(s) {argnums}. Takes the same arguments as {fun} but " + "returns a two-element tuple where the first element is the value " + "of {fun} and the second element is the gradient, which has the " + "same shape as the arguments at positions {argnums}.") + + check_callable(fun) + argnums = core.concrete_or_error(_ensure_index, argnums) + + @wraps(fun, docstr=docstr, argnums=argnums) + @api_boundary + def value_and_grad_f(*args, **kwargs): + max_argnum = argnums if isinstance(argnums, int) else max(argnums) + if max_argnum >= len(args): + raise TypeError(f"differentiating with respect to {argnums=} requires at least " + f"{max_argnum + 1} positional arguments to be passed by the caller, " + f"but got only {len(args)} positional arguments.") + + f = lu.wrap_init(fun, kwargs) + f_partial, dyn_args = argnums_partial(f, argnums, args, + require_static_args_hashable=False) + for leaf in tree_leaves(dyn_args): + _check_input_dtype_grad(holomorphic, allow_int, leaf) + if not has_aux: + ans, vjp_py = _vjp(f_partial, *dyn_args) + else: + ans, vjp_py, aux = _vjp( + f_partial, *dyn_args, has_aux=True) + _check_scalar(ans) + tree_map(partial(_check_output_dtype_grad, holomorphic), ans) + g = vjp_py(lax_internal._one(ans)) + g = g[0] if isinstance(argnums, int) else g + if not has_aux: + return ans, g + else: + return (ans, aux), g + + return value_and_grad_f + +def _check_scalar(x): + msg = "Gradient only defined for scalar-output functions. Output {}.".format + try: + aval = core.get_aval(x) + except TypeError as e: + raise TypeError(msg(f"was {x}")) from e + else: + if isinstance(aval, ShapedArray): + if aval.shape != (): + raise TypeError(msg(f"had shape: {aval.shape}")) + else: + raise TypeError(msg(f"had abstract value {aval}")) + +def _check_input_dtype_revderiv(name, holomorphic, allow_int, x): + dispatch.check_arg(x) + aval = core.get_aval(x) + if holomorphic: + if not dtypes.issubdtype(aval.dtype, np.complexfloating): + raise TypeError(f"{name} with holomorphic=True requires inputs with complex dtype, " + f"but got {aval.dtype.name}.") + if (dtypes.issubdtype(aval.dtype, dtypes.extended) or + dtypes.issubdtype(aval.dtype, np.integer) or + dtypes.issubdtype(aval.dtype, np.bool_)): + if not allow_int: + raise TypeError(f"{name} requires real- or complex-valued inputs (input dtype " + f"that is a sub-dtype of np.inexact), but got {aval.dtype.name}. " + "If you want to use Boolean- or integer-valued inputs, use vjp " + "or set allow_int to True.") + elif not dtypes.issubdtype(aval.dtype, np.inexact): + raise TypeError(f"{name} requires numerical-valued inputs (input dtype that is a " + f"sub-dtype of np.bool_ or np.number), but got {aval.dtype.name}.") +_check_input_dtype_grad = partial(_check_input_dtype_revderiv, "grad") + +def _check_output_dtype_revderiv(name, holomorphic, x): + aval = core.get_aval(x) + if dtypes.issubdtype(aval.dtype, dtypes.extended): + raise TypeError( + f"{name} with output element type {aval.dtype.name}") + if holomorphic: + if not dtypes.issubdtype(aval.dtype, np.complexfloating): + raise TypeError(f"{name} with holomorphic=True requires outputs with complex dtype, " + f"but got {aval.dtype.name}.") + elif dtypes.issubdtype(aval.dtype, np.complexfloating): + raise TypeError(f"{name} requires real-valued outputs (output dtype that is " + f"a sub-dtype of np.floating), but got {aval.dtype.name}. " + "For holomorphic differentiation, pass holomorphic=True. " + "For differentiation of non-holomorphic functions involving complex " + "outputs, use jax.vjp directly.") + elif not dtypes.issubdtype(aval.dtype, np.floating): + raise TypeError(f"{name} requires real-valued outputs (output dtype that is " + f"a sub-dtype of np.floating), but got {aval.dtype.name}. " + "For differentiation of functions with integer outputs, use " + "jax.vjp directly.") +_check_output_dtype_grad = partial(_check_output_dtype_revderiv, "grad") + + +def jacfwd(fun: Callable, argnums: int | Sequence[int] = 0, + has_aux: bool = False, holomorphic: bool = False) -> Callable: + """Jacobian of ``fun`` evaluated column-by-column using forward-mode AD. + + Args: + fun: Function whose Jacobian is to be computed. + argnums: Optional, integer or sequence of integers. Specifies which + positional argument(s) to differentiate with respect to (default ``0``). + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be + holomorphic. Default False. + + Returns: + A function with the same arguments as ``fun``, that evaluates the Jacobian of + ``fun`` using forward-mode automatic differentiation. If ``has_aux`` is True + then a pair of (jacobian, auxiliary_data) is returned. + + >>> import jax + >>> import jax.numpy as jnp + >>> + >>> def f(x): + ... return jnp.asarray( + ... [x[0], 5*x[2], 4*x[1]**2 - 2*x[2], x[2] * jnp.sin(x[0])]) + ... + >>> print(jax.jacfwd(f)(jnp.array([1., 2., 3.]))) + [[ 1. 0. 0. ] + [ 0. 0. 5. ] + [ 0. 16. -2. ] + [ 1.6209 0. 0.84147]] + """ + check_callable(fun) + argnums = _ensure_index(argnums) + + docstr = ("Jacobian of {fun} with respect to positional argument(s) " + "{argnums}. Takes the same arguments as {fun} but returns the " + "jacobian of the output with respect to the arguments at " + "positions {argnums}.") + + @wraps(fun, docstr=docstr, argnums=argnums) + def jacfun(*args, **kwargs): + f = lu.wrap_init(fun, kwargs) + f_partial, dyn_args = argnums_partial(f, argnums, args, + require_static_args_hashable=False) + tree_map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args) + if not has_aux: + pushfwd: Callable = partial(_jvp, f_partial, dyn_args) + y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args)) + else: + pushfwd: Callable = partial(_jvp, f_partial, dyn_args, has_aux=True) + y, jac, aux = vmap(pushfwd, out_axes=(None, -1, None))(_std_basis(dyn_args)) + tree_map(partial(_check_output_dtype_jacfwd, holomorphic), y) + example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args + jac_tree = tree_map(partial(_jacfwd_unravel, example_args), y, jac) + if not has_aux: + return jac_tree + else: + return jac_tree, aux + + return jacfun + +def _check_input_dtype_jacfwd(holomorphic: bool, x: Any) -> None: + dispatch.check_arg(x) + aval = core.get_aval(x) + if dtypes.issubdtype(aval.dtype, dtypes.extended): + raise TypeError( + f"jacfwd with input element type {aval.dtype.name}") + if holomorphic: + if not dtypes.issubdtype(aval.dtype, np.complexfloating): + raise TypeError("jacfwd with holomorphic=True requires inputs with complex " + f"dtype, but got {aval.dtype.name}.") + elif not dtypes.issubdtype(aval.dtype, np.floating): + raise TypeError("jacfwd requires real-valued inputs (input dtype that is " + f"a sub-dtype of np.floating), but got {aval.dtype.name}. " + "For holomorphic differentiation, pass holomorphic=True. " + "For differentiation of non-holomorphic functions involving " + "complex inputs or integer inputs, use jax.jvp directly.") + +def _check_output_dtype_jacfwd(holomorphic, x): + aval = core.get_aval(x) + if holomorphic: + if not dtypes.issubdtype(aval.dtype, np.complexfloating): + raise TypeError("jacfwd with holomorphic=True requires outputs with complex dtype, " + f"but got {aval.dtype.name}.") + +def jacrev(fun: Callable, argnums: int | Sequence[int] = 0, + has_aux: bool = False, holomorphic: bool = False, allow_int: bool = False) -> Callable: + """Jacobian of ``fun`` evaluated row-by-row using reverse-mode AD. + + Args: + fun: Function whose Jacobian is to be computed. + argnums: Optional, integer or sequence of integers. Specifies which + positional argument(s) to differentiate with respect to (default ``0``). + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be + holomorphic. Default False. + allow_int: Optional, bool. Whether to allow differentiating with + respect to integer valued inputs. The gradient of an integer input will + have a trivial vector-space dtype (float0). Default False. + + Returns: + A function with the same arguments as ``fun``, that evaluates the Jacobian of + ``fun`` using reverse-mode automatic differentiation. If ``has_aux`` is True + then a pair of (jacobian, auxiliary_data) is returned. + + >>> import jax + >>> import jax.numpy as jnp + >>> + >>> def f(x): + ... return jnp.asarray( + ... [x[0], 5*x[2], 4*x[1]**2 - 2*x[2], x[2] * jnp.sin(x[0])]) + ... + >>> print(jax.jacrev(f)(jnp.array([1., 2., 3.]))) + [[ 1. 0. 0. ] + [ 0. 0. 5. ] + [ 0. 16. -2. ] + [ 1.6209 0. 0.84147]] + """ + check_callable(fun) + + docstr = ("Jacobian of {fun} with respect to positional argument(s) " + "{argnums}. Takes the same arguments as {fun} but returns the " + "jacobian of the output with respect to the arguments at " + "positions {argnums}.") + + @wraps(fun, docstr=docstr, argnums=argnums) + def jacfun(*args, **kwargs): + f = lu.wrap_init(fun, kwargs) + f_partial, dyn_args = argnums_partial(f, argnums, args, + require_static_args_hashable=False) + tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args) + if not has_aux: + y, pullback = _vjp(f_partial, *dyn_args) + else: + y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True) + tree_map(partial(_check_output_dtype_jacrev, holomorphic), y) + jac = vmap(pullback)(_std_basis(y)) + jac = jac[0] if isinstance(argnums, int) else jac + example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args + jac_tree = tree_map(partial(_jacrev_unravel, y), example_args, jac) + jac_tree = tree_transpose(tree_structure(example_args), tree_structure(y), jac_tree) + if not has_aux: + return jac_tree + else: + return jac_tree, aux + + return jacfun +jacobian = jacrev + +_check_input_dtype_jacrev = partial(_check_input_dtype_revderiv, "jacrev") +_check_output_dtype_jacrev = partial(_check_output_dtype_revderiv, "jacrev") + + +def hessian(fun: Callable, argnums: int | Sequence[int] = 0, + has_aux: bool = False, holomorphic: bool = False) -> Callable: + """Hessian of ``fun`` as a dense array. + + Args: + fun: Function whose Hessian is to be computed. Its arguments at positions + specified by ``argnums`` should be arrays, scalars, or standard Python + containers thereof. It should return arrays, scalars, or standard Python + containers thereof. + argnums: Optional, integer or sequence of integers. Specifies which + positional argument(s) to differentiate with respect to (default ``0``). + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be + holomorphic. Default False. + + Returns: + A function with the same arguments as ``fun``, that evaluates the Hessian of + ``fun``. + + >>> import jax + >>> + >>> g = lambda x: x[0]**3 - 2*x[0]*x[1] - x[1]**6 + >>> print(jax.hessian(g)(jax.numpy.array([1., 2.]))) + [[ 6. -2.] + [ -2. -480.]] + + :py:func:`hessian` is a generalization of the usual definition of the Hessian + that supports nested Python containers (i.e. pytrees) as inputs and outputs. + The tree structure of ``jax.hessian(fun)(x)`` is given by forming a tree + product of the structure of ``fun(x)`` with a tree product of two copies of + the structure of ``x``. A tree product of two tree structures is formed by + replacing each leaf of the first tree with a copy of the second. For example: + + >>> import jax.numpy as jnp + >>> f = lambda dct: {"c": jnp.power(dct["a"], dct["b"])} + >>> print(jax.hessian(f)({"a": jnp.arange(2.) + 1., "b": jnp.arange(2.) + 2.})) + {'c': {'a': {'a': Array([[[ 2., 0.], [ 0., 0.]], + [[ 0., 0.], [ 0., 12.]]], dtype=float32), + 'b': Array([[[ 1. , 0. ], [ 0. , 0. ]], + [[ 0. , 0. ], [ 0. , 12.317766]]], dtype=float32)}, + 'b': {'a': Array([[[ 1. , 0. ], [ 0. , 0. ]], + [[ 0. , 0. ], [ 0. , 12.317766]]], dtype=float32), + 'b': Array([[[0. , 0. ], [0. , 0. ]], + [[0. , 0. ], [0. , 3.843624]]], dtype=float32)}}} + + Thus each leaf in the tree structure of ``jax.hessian(fun)(x)`` corresponds to + a leaf of ``fun(x)`` and a pair of leaves of ``x``. For each leaf in + ``jax.hessian(fun)(x)``, if the corresponding array leaf of ``fun(x)`` has + shape ``(out_1, out_2, ...)`` and the corresponding array leaves of ``x`` have + shape ``(in_1_1, in_1_2, ...)`` and ``(in_2_1, in_2_2, ...)`` respectively, + then the Hessian leaf has shape ``(out_1, out_2, ..., in_1_1, in_1_2, ..., + in_2_1, in_2_2, ...)``. In other words, the Python tree structure represents + the block structure of the Hessian, with blocks determined by the input and + output pytrees. + + In particular, an array is produced (with no pytrees involved) when the + function input ``x`` and output ``fun(x)`` are each a single array, as in the + ``g`` example above. If ``fun(x)`` has shape ``(out1, out2, ...)`` and ``x`` + has shape ``(in1, in2, ...)`` then ``jax.hessian(fun)(x)`` has shape + ``(out1, out2, ..., in1, in2, ..., in1, in2, ...)``. To flatten pytrees into + 1D vectors, consider using :py:func:`jax.flatten_util.flatten_pytree`. + """ + return jacfwd(jacrev(fun, argnums, has_aux=has_aux, holomorphic=holomorphic), + argnums, has_aux=has_aux, holomorphic=holomorphic) + +def _std_basis(pytree): + import jax.numpy as jnp + leaves, _ = tree_flatten(pytree) + ndim = sum(map(np.size, leaves)) + dtype = dtypes.result_type(*leaves) + flat_basis = jnp.eye(ndim, dtype=dtype) + return _unravel_array_into_pytree(pytree, 1, None, flat_basis) + +def _jacfwd_unravel(input_pytree, output_pytree_leaf, arr): + return _unravel_array_into_pytree( + input_pytree, -1, output_pytree_leaf, arr) + +def _jacrev_unravel(output_pytree, input_pytree_leaf, arr): + return _unravel_array_into_pytree( + output_pytree, 0, input_pytree_leaf, arr) + +def _possible_downcast(x, example): + if (dtypes.issubdtype(x.dtype, np.complexfloating) and + not dtypes.issubdtype(_dtype(example), np.complexfloating)): + x = x.real + dtype = None if example is None else _dtype(example) + weak_type = None if example is None else dtypes.is_weakly_typed(example) + return lax_internal._convert_element_type(x, dtype, weak_type) + +def _unravel_array_into_pytree(pytree, axis, example, arr): + """Unravel an array into a PyTree with a given structure. + Args: + pytree: The pytree that provides the structure. + axis: The parameter axis is either -1, 0, or 1. It controls the + resulting shapes. + example: If specified, cast the components to the matching dtype/weak_type, + or else use the pytree leaf type if example is None. + arr: The array to be unraveled. + """ + leaves, treedef = tree_flatten(pytree) + axis = axis % arr.ndim + shapes = [arr.shape[:axis] + np.shape(l) + arr.shape[axis+1:] for l in leaves] + parts = _split(arr, np.cumsum(map(np.size, leaves[:-1])), axis) + reshaped_parts = [ + _possible_downcast(np.reshape(x, shape), leaf if example is None else example) + for x, shape, leaf in zip(parts, shapes, leaves)] + return tree_unflatten(treedef, reshaped_parts) + +def _split(x, indices, axis): + if isinstance(x, np.ndarray): + return np.split(x, indices, axis) + else: + return x._split(indices, axis) + + +def vmap(fun: F, + in_axes: int | None | Sequence[Any] = 0, + out_axes: Any = 0, + axis_name: AxisName | None = None, + axis_size: int | None = None, + spmd_axis_name: AxisName | tuple[AxisName, ...] | None = None + ) -> F: + """Vectorizing map. Creates a function which maps ``fun`` over argument axes. + + Args: + fun: Function to be mapped over additional axes. + in_axes: An integer, None, or sequence of values specifying which input + array axes to map over. + + If each positional argument to ``fun`` is an array, then ``in_axes`` can + be an integer, a None, or a tuple of integers and Nones with length equal + to the number of positional arguments to ``fun``. An integer or ``None`` + indicates which array axis to map over for all arguments (with ``None`` + indicating not to map any axis), and a tuple indicates which axis to map + for each corresponding positional argument. Axis integers must be in the + range ``[-ndim, ndim)`` for each array, where ``ndim`` is the number of + dimensions (axes) of the corresponding input array. + + If the positional arguments to ``fun`` are container (pytree) types, ``in_axes`` + must be a sequence with length equal to the number of positional arguments to + ``fun``, and for each argument the corresponding element of ``in_axes`` can + be a container with a matching pytree structure specifying the mapping of its + container elements. In other words, ``in_axes`` must be a container tree prefix + of the positional argument tuple passed to ``fun``. See this link for more detail: + https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees + + Either ``axis_size`` must be provided explicitly, or at least one + positional argument must have ``in_axes`` not None. The sizes of the + mapped input axes for all mapped positional arguments must all be equal. + + Arguments passed as keywords are always mapped over their leading axis + (i.e. axis index 0). + + See below for examples. + + out_axes: An integer, None, or (nested) standard Python container + (tuple/list/dict) thereof indicating where the mapped axis should appear + in the output. All outputs with a mapped axis must have a non-None + ``out_axes`` specification. Axis integers must be in the range ``[-ndim, + ndim)`` for each output array, where ``ndim`` is the number of dimensions + (axes) of the array returned by the :func:`vmap`-ed function, which is one + more than the number of dimensions (axes) of the corresponding array + returned by ``fun``. + axis_name: Optional, a hashable Python object used to identify the mapped + axis so that parallel collectives can be applied. + axis_size: Optional, an integer indicating the size of the axis to be + mapped. If not provided, the mapped axis size is inferred from arguments. + + Returns: + Batched/vectorized version of ``fun`` with arguments that correspond to + those of ``fun``, but with extra array axes at positions indicated by + ``in_axes``, and a return value that corresponds to that of ``fun``, but + with extra array axes at positions indicated by ``out_axes``. + + For example, we can implement a matrix-matrix product using a vector dot + product: + + >>> import jax.numpy as jnp + >>> + >>> vv = lambda x, y: jnp.vdot(x, y) # ([a], [a]) -> [] + >>> mv = vmap(vv, (0, None), 0) # ([b,a], [a]) -> [b] (b is the mapped axis) + >>> mm = vmap(mv, (None, 1), 1) # ([b,a], [a,c]) -> [b,c] (c is the mapped axis) + + Here we use ``[a,b]`` to indicate an array with shape (a,b). Here are some + variants: + + >>> mv1 = vmap(vv, (0, 0), 0) # ([b,a], [b,a]) -> [b] (b is the mapped axis) + >>> mv2 = vmap(vv, (0, 1), 0) # ([b,a], [a,b]) -> [b] (b is the mapped axis) + >>> mm2 = vmap(mv2, (1, 1), 0) # ([b,c,a], [a,c,b]) -> [c,b] (c is the mapped axis) + + Here's an example of using container types in ``in_axes`` to specify which + axes of the container elements to map over: + + >>> A, B, C, D = 2, 3, 4, 5 + >>> x = jnp.ones((A, B)) + >>> y = jnp.ones((B, C)) + >>> z = jnp.ones((C, D)) + >>> def foo(tree_arg): + ... x, (y, z) = tree_arg + ... return jnp.dot(x, jnp.dot(y, z)) + >>> tree = (x, (y, z)) + >>> print(foo(tree)) + [[12. 12. 12. 12. 12.] + [12. 12. 12. 12. 12.]] + >>> from jax import vmap + >>> K = 6 # batch size + >>> x = jnp.ones((K, A, B)) # batch axis in different locations + >>> y = jnp.ones((B, K, C)) + >>> z = jnp.ones((C, D, K)) + >>> tree = (x, (y, z)) + >>> vfoo = vmap(foo, in_axes=((0, (1, 2)),)) + >>> print(vfoo(tree).shape) + (6, 2, 5) + + Here's another example using container types in ``in_axes``, this time a + dictionary, to specify the elements of the container to map over: + + >>> dct = {'a': 0., 'b': jnp.arange(5.)} + >>> x = 1. + >>> def foo(dct, x): + ... return dct['a'] + dct['b'] + x + >>> out = vmap(foo, in_axes=({'a': None, 'b': 0}, None))(dct, x) + >>> print(out) + [1. 2. 3. 4. 5.] + + The results of a vectorized function can be mapped or unmapped. For example, + the function below returns a pair with the first element mapped and the second + unmapped. Only for unmapped results we can specify ``out_axes`` to be ``None`` + (to keep it unmapped). + + >>> print(vmap(lambda x, y: (x + y, y * 2.), in_axes=(0, None), out_axes=(0, None))(jnp.arange(2.), 4.)) + (Array([4., 5.], dtype=float32), 8.0) + + If the ``out_axes`` is specified for an unmapped result, the result is + broadcast across the mapped axis: + + >>> print(vmap(lambda x, y: (x + y, y * 2.), in_axes=(0, None), out_axes=0)(jnp.arange(2.), 4.)) + (Array([4., 5.], dtype=float32), Array([8., 8.], dtype=float32, weak_type=True)) + + If the ``out_axes`` is specified for a mapped result, the result is transposed + accordingly. + + Finally, here's an example using ``axis_name`` together with collectives: + + >>> xs = jnp.arange(3. * 4.).reshape(3, 4) + >>> print(vmap(lambda x: lax.psum(x, 'i'), axis_name='i')(xs)) + [[12. 15. 18. 21.] + [12. 15. 18. 21.] + [12. 15. 18. 21.]] + + See the :py:func:`jax.pmap` docstring for more examples involving collectives. + """ + check_callable(fun) + docstr = ("Vectorized version of {fun}. Takes similar arguments as {fun} " + "but with additional array axes over which {fun} is mapped.") + if fun.__doc__: + docstr += "\n\nOriginal documentation:\n\n" + docstr += fun.__doc__ + + axis_name = core.no_axis_name if axis_name is None else axis_name + if spmd_axis_name is not None and type(spmd_axis_name) is not tuple: + spmd_axis_name = (spmd_axis_name,) + + if isinstance(in_axes, list): + # To be a tree prefix of the positional args tuple, in_axes can never be a + # list: if in_axes is not a leaf, it must be a tuple of trees. However, + # in cases like these users expect tuples and lists to be treated + # essentially interchangeably, so we canonicalize lists to tuples here + # rather than raising an error. https://github.com/google/jax/issues/2367 + in_axes = tuple(in_axes) + + if not (in_axes is None or type(in_axes) in {int, tuple, *batching.spec_types}): + raise TypeError("vmap in_axes must be an int, None, or a tuple of entries corresponding " + f"to the positional arguments passed to the function, but got {in_axes}.") + if not all(type(l) in {int, *batching.spec_types} for l in tree_leaves(in_axes)): + raise TypeError("vmap in_axes must be an int, None, or (nested) container " + f"with those types as leaves, but got {in_axes}.") + if not all(type(l) in {int, *batching.spec_types} for l in tree_leaves(out_axes)): + raise TypeError("vmap out_axes must be an int, None, or (nested) container " + f"with those types as leaves, but got {out_axes}.") + + @wraps(fun, docstr=docstr) + @api_boundary + def vmap_f(*args, **kwargs): + if isinstance(in_axes, tuple) and len(in_axes) != len(args): + raise ValueError("vmap in_axes must be an int, None, or a tuple of entries corresponding " + "to the positional arguments passed to the function, " + f"but got {len(in_axes)=}, {len(args)=}") + args_flat, in_tree = tree_flatten((args, kwargs), is_leaf=batching.is_vmappable) + f = lu.wrap_init(fun) + flat_fun, out_tree = batching.flatten_fun_for_vmap(f, in_tree) + in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True) + axis_size_ = (axis_size if axis_size is not None else + _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, "vmap")) + out_flat = batching.batch( + flat_fun, axis_name, axis_size_, in_axes_flat, + lambda: flatten_axes("vmap out_axes", out_tree(), out_axes), + spmd_axis_name=spmd_axis_name + ).call_wrapped(*args_flat) + return tree_unflatten(out_tree(), out_flat) + + return cast(F, vmap_f) + +def _mapped_axis_size(fn, tree, vals, dims, name): + if not vals: + args, kwargs = tree_unflatten(tree, vals) + raise ValueError( + f"{name} wrapped function must be passed at least one argument " + f"containing an array, got empty *args={args} and **kwargs={kwargs}" + ) + + def _get_axis_size(name: str, shape: tuple[core.AxisSize, ...], axis: int + ) -> core.AxisSize: + try: + return shape[axis] + except (IndexError, TypeError) as e: + min_rank = axis + 1 if axis >= 0 else -axis + # TODO(mattjj): better error message here + raise ValueError( + f"{name} was requested to map its argument along axis {axis}, " + f"which implies that its rank should be at least {min_rank}, " + f"but is only {len(shape)} (its shape is {shape})") from e + + sizes = core.dedup_referents(_get_axis_size(name, np.shape(x), d) + for x, d in zip(vals, dims) if d is not None) + if len(sizes) == 1: + sz, = sizes + return sz + if not sizes: + msg = f"{name} must have at least one non-None value in in_axes" + raise ValueError(msg) + + def _get_argument_type(x): + try: + return shaped_abstractify(x).str_short() + except TypeError: #Catch all for user specified objects that can't be interpreted as a data type + return "unknown" + msg = [f"{name} got inconsistent sizes for array axes to be mapped:\n"] + args, kwargs = tree_unflatten(tree, vals) + try: + ba = inspect.signature(fn).bind(*args, **kwargs) + except (TypeError, ValueError): + ba = None + if ba is None: + args_paths = [f'args{keystr(p)} ' + f'of type {_get_argument_type(x)}' + for p, x in generate_key_paths(args)] + kwargs_paths = [f'kwargs{keystr(p)} ' + f'of type {_get_argument_type(x)}' + for p, x in generate_key_paths(kwargs)] + key_paths = [*args_paths, *kwargs_paths] + else: + key_paths = [f'argument {name}{keystr(p)} ' + f'of type {_get_argument_type(x)}' + for name, arg in ba.arguments.items() + for p, x in generate_key_paths(arg)] + all_sizes = [_get_axis_size(name, np.shape(x), d) if d is not None else None + for x, d in zip(vals, dims)] + size_counts = collections.Counter(s for s in all_sizes if s is not None) + (sz, ct), *other_counts = counts = size_counts.most_common() + def _all_sizes_index(sz): + for i, isz in enumerate(all_sizes): + if core.definitely_equal(isz, sz): return i + assert False, (sz, all_sizes) + + ex, *examples = (key_paths[_all_sizes_index(sz)] for sz, _ in counts) + ax, *axs = (dims[_all_sizes_index(sz)] for sz, _ in counts) + if ct == 1: + msg.append(f" * one axis had size {sz}: axis {ax} of {ex};\n") + else: + msg.append(f" * most axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n") + for ex, ax, (sz, ct) in zip(examples, axs, other_counts): + if ct == 1: + msg.append(f" * one axis had size {sz}: axis {ax} of {ex};\n") + else: + msg.append(f" * some axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n") + raise ValueError(''.join(msg)[:-2]) # remove last semicolon and newline + + +def pmap( + fun: Callable, + axis_name: AxisName | None = None, + *, + in_axes=0, + out_axes=0, + static_broadcasted_argnums: int | Iterable[int] = (), + devices: Sequence[xc.Device] | None = None, # noqa: F811 + backend: str | None = None, + axis_size: int | None = None, + donate_argnums: int | Iterable[int] = (), + global_arg_shapes: tuple[tuple[int, ...], ...] | None = None, + ) -> Any: + """Parallel map with support for collective operations. + + The purpose of :py:func:`pmap` is to express single-program multiple-data + (SPMD) programs. Applying :py:func:`pmap` to a function will compile the + function with XLA (similarly to :py:func:`jit`), then execute it in parallel + on XLA devices, such as multiple GPUs or multiple TPU cores. Semantically it + is comparable to :py:func:`vmap` because both transformations map a function + over array axes, but where :py:func:`vmap` vectorizes functions by pushing the + mapped axis down into primitive operations, :py:func:`pmap` instead replicates + the function and executes each replica on its own XLA device in parallel. + + The mapped axis size must be less than or equal to the number of local XLA + devices available, as returned by :py:func:`jax.local_device_count()` (unless + ``devices`` is specified, see below). For nested :py:func:`pmap` calls, the + product of the mapped axis sizes must be less than or equal to the number of + XLA devices. + + .. note:: + :py:func:`pmap` compiles ``fun``, so while it can be combined with + :py:func:`jit`, it's usually unnecessary. + + :py:func:`pmap` requires that all of the participating devices are identical. + For example, it is not possible to use :py:func:`pmap` to parallelize a + computation across two different models of GPU. It is currently an error for + the same device to participate twice in the same `pmap`. + + **Multi-process platforms:** On multi-process platforms such as TPU pods, + :py:func:`pmap` is designed to be used in SPMD Python programs, where every + process is running the same Python code such that all processes run the same + pmapped function in the same order. Each process should still call the pmapped + function with mapped axis size equal to the number of *local* devices (unless + ``devices`` is specified, see below), and an array of the same leading axis + size will be returned as usual. However, any collective operations in ``fun`` + will be computed over *all* participating devices, including those on other + processes, via device-to-device communication. Conceptually, this can be + thought of as running a pmap over a single array sharded across processes, + where each process "sees" only its local shard of the input and output. The + SPMD model requires that the same multi-process pmaps must be run in the same + order on all devices, but they can be interspersed with arbitrary operations + running in a single process. + + Args: + fun: Function to be mapped over argument axes. Its arguments and return + value should be arrays, scalars, or (nested) standard Python containers + (tuple/list/dict) thereof. Positional arguments indicated by + ``static_broadcasted_argnums`` can be anything at all, provided they are + hashable and have an equality operation defined. + axis_name: Optional, a hashable Python object used to identify the mapped + axis so that parallel collectives can be applied. + in_axes: A non-negative integer, None, or nested Python container thereof + that specifies which axes of positional arguments to map over. Arguments + passed as keywords are always mapped over their leading axis (i.e. axis + index 0). See :py:func:`vmap` for details. + out_axes: A non-negative integer, None, or nested Python container thereof + indicating where the mapped axis should appear in the output. All outputs + with a mapped axis must have a non-None ``out_axes`` specification + (see :py:func:`vmap`). + static_broadcasted_argnums: An int or collection of ints specifying which + positional arguments to treat as static (compile-time constant). + Operations that only depend on static arguments will be constant-folded. + Calling the pmapped function with different values for these constants + will trigger recompilation. If the pmapped function is called with fewer + positional arguments than indicated by ``static_broadcasted_argnums`` then + an error is raised. Each of the static arguments will be broadcasted to + all devices. Arguments that are not arrays or containers thereof must be + marked as static. Defaults to (). + + Static arguments must be hashable, meaning both ``__hash__`` and + ``__eq__`` are implemented, and should be immutable. + + devices: This is an experimental feature and the API is likely to change. + Optional, a sequence of Devices to map over. (Available devices can be + retrieved via jax.devices()). Must be given identically for each process + in multi-process settings (and will therefore include devices across + processes). If specified, the size of the mapped axis must be equal to + the number of devices in the sequence local to the given process. Nested + :py:func:`pmap` s with ``devices`` specified in either the inner or outer + :py:func:`pmap` are not yet supported. + backend: This is an experimental feature and the API is likely to change. + Optional, a string representing the XLA backend. 'cpu', 'gpu', or 'tpu'. + axis_size: Optional; the size of the mapped axis. + donate_argnums: Specify which positional argument buffers are "donated" to + the computation. It is safe to donate argument buffers if you no longer need + them once the computation has finished. In some cases XLA can make use of + donated buffers to reduce the amount of memory needed to perform a + computation, for example recycling one of your input buffers to store a + result. You should not reuse buffers that you donate to a computation, JAX + will raise an error if you try to. + Note that donate_argnums only work for positional arguments, and keyword + arguments will not be donated. + + For more details on buffer donation see the + `FAQ `_. + + Returns: + A parallelized version of ``fun`` with arguments that correspond to those of + ``fun`` but with extra array axes at positions indicated by ``in_axes`` and + with output that has an additional leading array axis (with the same size). + + For example, assuming 8 XLA devices are available, :py:func:`pmap` can be used + as a map along a leading array axis: + + >>> import jax.numpy as jnp + >>> + >>> out = pmap(lambda x: x ** 2)(jnp.arange(8)) # doctest: +SKIP + >>> print(out) # doctest: +SKIP + [0, 1, 4, 9, 16, 25, 36, 49] + + When the leading dimension is smaller than the number of available devices JAX + will simply run on a subset of devices: + + >>> x = jnp.arange(3 * 2 * 2.).reshape((3, 2, 2)) + >>> y = jnp.arange(3 * 2 * 2.).reshape((3, 2, 2)) ** 2 + >>> out = pmap(jnp.dot)(x, y) # doctest: +SKIP + >>> print(out) # doctest: +SKIP + [[[ 4. 9.] + [ 12. 29.]] + [[ 244. 345.] + [ 348. 493.]] + [[ 1412. 1737.] + [ 1740. 2141.]]] + + If your leading dimension is larger than the number of available devices you + will get an error: + + >>> pmap(lambda x: x ** 2)(jnp.arange(9)) # doctest: +SKIP + ValueError: ... requires 9 replicas, but only 8 XLA devices are available + + As with :py:func:`vmap`, using ``None`` in ``in_axes`` indicates that an + argument doesn't have an extra axis and should be broadcasted, rather than + mapped, across the replicas: + + >>> x, y = jnp.arange(2.), 4. + >>> out = pmap(lambda x, y: (x + y, y * 2.), in_axes=(0, None))(x, y) # doctest: +SKIP + >>> print(out) # doctest: +SKIP + ([4., 5.], [8., 8.]) + + Note that :py:func:`pmap` always returns values mapped over their leading axis, + equivalent to using ``out_axes=0`` in :py:func:`vmap`. + + In addition to expressing pure maps, :py:func:`pmap` can also be used to express + parallel single-program multiple-data (SPMD) programs that communicate via + collective operations. For example: + + >>> f = lambda x: x / jax.lax.psum(x, axis_name='i') + >>> out = pmap(f, axis_name='i')(jnp.arange(4.)) # doctest: +SKIP + >>> print(out) # doctest: +SKIP + [ 0. 0.16666667 0.33333334 0.5 ] + >>> print(out.sum()) # doctest: +SKIP + 1.0 + + In this example, ``axis_name`` is a string, but it can be any Python object + with ``__hash__`` and ``__eq__`` defined. + + The argument ``axis_name`` to :py:func:`pmap` names the mapped axis so that + collective operations, like :func:`jax.lax.psum`, can refer to it. Axis names + are important particularly in the case of nested :py:func:`pmap` functions, + where collective operations can operate over distinct axes: + + >>> from functools import partial + >>> import jax + >>> + >>> @partial(pmap, axis_name='rows') + ... @partial(pmap, axis_name='cols') + ... def normalize(x): + ... row_normed = x / jax.lax.psum(x, 'rows') + ... col_normed = x / jax.lax.psum(x, 'cols') + ... doubly_normed = x / jax.lax.psum(x, ('rows', 'cols')) + ... return row_normed, col_normed, doubly_normed + >>> + >>> x = jnp.arange(8.).reshape((4, 2)) + >>> row_normed, col_normed, doubly_normed = normalize(x) # doctest: +SKIP + >>> print(row_normed.sum(0)) # doctest: +SKIP + [ 1. 1.] + >>> print(col_normed.sum(1)) # doctest: +SKIP + [ 1. 1. 1. 1.] + >>> print(doubly_normed.sum((0, 1))) # doctest: +SKIP + 1.0 + + On multi-process platforms, collective operations operate over all devices, + including those on other processes. For example, assuming the following code + runs on two processes with 4 XLA devices each: + + >>> f = lambda x: x + jax.lax.psum(x, axis_name='i') + >>> data = jnp.arange(4) if jax.process_index() == 0 else jnp.arange(4, 8) + >>> out = pmap(f, axis_name='i')(data) # doctest: +SKIP + >>> print(out) # doctest: +SKIP + [28 29 30 31] # on process 0 + [32 33 34 35] # on process 1 + + Each process passes in a different length-4 array, corresponding to its 4 + local devices, and the psum operates over all 8 values. Conceptually, the two + length-4 arrays can be thought of as a sharded length-8 array (in this example + equivalent to jnp.arange(8)) that is mapped over, with the length-8 mapped + axis given name 'i'. The pmap call on each process then returns the + corresponding length-4 output shard. + + The ``devices`` argument can be used to specify exactly which devices are used + to run the parallel computation. For example, again assuming a single process + with 8 devices, the following code defines two parallel computations, one + which runs on the first six devices and one on the remaining two: + + >>> from functools import partial + >>> @partial(pmap, axis_name='i', devices=jax.devices()[:6]) + ... def f1(x): + ... return x / jax.lax.psum(x, axis_name='i') + >>> + >>> @partial(pmap, axis_name='i', devices=jax.devices()[-2:]) + ... def f2(x): + ... return jax.lax.psum(x ** 2, axis_name='i') + >>> + >>> print(f1(jnp.arange(6.))) # doctest: +SKIP + [0. 0.06666667 0.13333333 0.2 0.26666667 0.33333333] + >>> print(f2(jnp.array([2., 3.]))) # doctest: +SKIP + [ 13. 13.] + """ + if global_arg_shapes is not None: + raise ValueError( + "global_arg_shapes only worked with sharded_jit which has long been" + " removed from JAX. Please migrate to pjit and remove global_arg_shapes" + " from pmap.") + + # TODO(yashkatariya): Move this out after shard_map is out of experimental and + # in _src + if config.pmap_shmap_merge.value: + from jax.experimental.shard_map import pmap + return pmap(fun, axis_name, in_axes=in_axes, out_axes=out_axes, + static_broadcasted_argnums=static_broadcasted_argnums, + devices=devices, backend=backend, + axis_size=axis_size, + donate_argnums=donate_argnums) + + return _cpp_pmap( + fun, + axis_name, + in_axes=in_axes, + out_axes=out_axes, + static_broadcasted_argnums=static_broadcasted_argnums, + devices=devices, + backend=backend, + axis_size=axis_size, + donate_argnums=donate_argnums) + + +class PmapCallInfo(NamedTuple): + flat_fun: lu.WrappedFun + in_tree: PyTreeDef + out_tree: Callable[[], PyTreeDef] + flat_args: Sequence[Any] + donated_invars: Sequence[bool] + in_axes_flat: Sequence[int | None] + local_axis_size: int + out_axes_thunk: Callable + devices: Sequence[xc.Device] | None + global_axis_size: int + is_explicit_global_axis_size: bool + + +def _get_global_axis_size(local_axis_size: int, in_devices, backend_name: str, + global_axis_size: int | None): + """Determine global_axis_size for multi-host pmap.""" + # TODO(mattjj,skyewm): revive this check (inner_pmap always False now) + # if xb.process_count() > 1 and global_axis_size is None and inner_pmap: + # raise ValueError("'axis_size' must be specified for nested multi-host pmaps") + if (xb.process_count() == 1 and global_axis_size is not None and + global_axis_size != local_axis_size): + raise ValueError( + f"Specified axis_size {global_axis_size} doesn't match received " + f"axis_size {local_axis_size}.") + + if in_devices is not None and backend_name is None: + backend = xb.get_device_backend(in_devices[0]) + else: + backend = xb.get_backend(backend_name) + + if global_axis_size is None: + if xb.process_count(backend) == 1: + global_axis_size = local_axis_size + elif in_devices: + global_axis_size = len(in_devices) + else: + global_axis_size = local_axis_size * xb.process_count(backend) + assert all( + len(xb.local_devices(pi, backend)) == xb.local_device_count(backend) + for pi in range(xb.process_count(backend))) + return global_axis_size + + +def _prepare_pmap(fun, in_axes, out_axes, static_broadcasted_tuple, + donate_tuple, in_devices, backend_name, + axis_size, args, kwargs): + if in_devices is not None and len(in_devices) == 0: + raise ValueError("'devices' argument to pmap must be non-empty, or None.") + + src = fun_sourceinfo(fun) + signature = api_util.fun_signature(fun) + + dbg = debug_info('pmap', src, signature, args, kwargs, + static_broadcasted_tuple, ()) + + f = lu.wrap_init(fun) + if static_broadcasted_tuple: + if max(static_broadcasted_tuple) >= len(args): + raise ValueError( + f"pmapped function has static_broadcasted_argnums={static_broadcasted_tuple}" + f" but was called with only {len(args)} positional " + f"argument{'s' if len(args) > 1 else ''}. " + "All static broadcasted arguments must be passed positionally.") + dyn_argnums = [i for i in range(len(args)) + if i not in static_broadcasted_tuple] + f, dyn_args = argnums_partial(f, dyn_argnums, args) + + if isinstance(in_axes, tuple): + dyn_in_axes = tuple(in_axes[i] for i in dyn_argnums) + else: + dyn_in_axes = in_axes + else: + dyn_args, dyn_in_axes = args, in_axes + args, in_tree = tree_flatten((dyn_args, kwargs)) + + if donate_tuple and not config.debug_nans.value: + donated_invars = donation_vector(donate_tuple, (), in_tree) + else: + donated_invars = (False,) * len(args) + try: + in_axes_flat = tuple(broadcast_prefix((dyn_in_axes, 0), (dyn_args, kwargs), + is_leaf=lambda x: x is None)) + except ValueError: + e, *_ = prefix_errors((dyn_in_axes, 0), (dyn_args, kwargs)) + ex = e('pmap in_axes') + msg, = ex.args + msg += ("\n\nThe 'full pytree' here is the tuple of arguments passed " + "positionally to the pmapped function, and the value of `in_axes` " + "must be a tree prefix of that tuple. But it was not a prefix.") + if kwargs: + msg += ("\n\nWhen some arguments are passed by keyword to the pmapped " + "function, they are not included in the comparison to `in_axes`. " + "Instead, each argument passed by keyword is mapped over its " + "leading axis. See the description of `in_axes` in the `pmap` " + "docstring: " + "https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html#jax.pmap") + msg += ("\n\nCheck that the value of the `in_axes` argument to `pmap` " + "is a tree prefix of the tuple of arguments passed positionally to " + "the pmapped function.") + raise ValueError(msg) from None + local_axis_size = _mapped_axis_size(fun, in_tree, args, in_axes_flat, "pmap") + + f, res_paths = result_paths(f) + f, out_axes_thunk = flat_out_axes(f, out_axes) + flat_fun, out_tree = flatten_fun(f, in_tree) + flat_fun = debug_info_final(flat_fun, dbg, res_paths) + + is_explicit_global_axis_size = axis_size is not None + global_axis_size = _get_global_axis_size(local_axis_size, in_devices, + backend_name, axis_size) + return PmapCallInfo(flat_fun=flat_fun, + in_tree=in_tree, + out_tree=out_tree, + flat_args=args, + donated_invars=donated_invars, + in_axes_flat=in_axes_flat, + local_axis_size=local_axis_size, + out_axes_thunk=out_axes_thunk, + devices=None if in_devices is None else tuple(in_devices), + global_axis_size=global_axis_size, + is_explicit_global_axis_size=is_explicit_global_axis_size) + + +def _shared_code_pmap(fun, axis_name, static_broadcasted_argnums, + donate_argnums, in_axes, out_axes): + # axis_size is an optional integer representing the global axis size. The + # aggregate size (across all processes) size of the mapped axis must match the + # given value. + check_callable(fun) + axis_name = core._TempAxisName(fun) if axis_name is None else axis_name + static_broadcasted_tuple = _ensure_index_tuple(static_broadcasted_argnums) + donate_tuple = rebase_donate_argnums( + _ensure_index_tuple(donate_argnums), static_broadcasted_tuple) + + if not all(type(l) is int for l in tree_leaves(in_axes)): + raise TypeError("pmap in_axes must be an int, None, or (nested) container " + f"with those types as leaves, but got {in_axes}.") + if not all(type(l) is int for l in tree_leaves(out_axes)): + raise TypeError("pmap out_axes must be an int, None, or (nested) container " + f"with those types as leaves, but got {out_axes}.") + + return axis_name, static_broadcasted_tuple, donate_tuple + + +class _PmapFastpathData(NamedTuple): + version: int # For forward and backward compatibility + xla_executable: xc.LoadedExecutable + in_handler: Any + out_handler: Any + out_pytree_def: Any + # Data needed to handle the inputs. + input_devices: Sequence[xc.Device] + input_indices: Sequence[sharding_specs.Index] + input_array_shardings: Sequence[Any] + # Data needed to build the Array from C++. + out_avals: Sequence[Any] + out_array_shardings: Sequence[Any] + out_committed: Sequence[Any] + + +def _cpp_pmap( + fun: Callable, + axis_name: AxisName | None = None, + *, + in_axes=0, + out_axes=0, + static_broadcasted_argnums: int | Iterable[int] = (), + devices: Sequence[xc.Device] | None = None, # noqa: F811 + backend: str | None = None, + axis_size: int | None = None, + donate_argnums: int | Iterable[int] = (), + ) -> Any: + axis_name, static_broadcasted_tuple, donate_tuple = _shared_code_pmap( + fun, axis_name, static_broadcasted_argnums, donate_argnums, in_axes, + out_axes) + del static_broadcasted_argnums, donate_argnums + + @api_boundary + def cache_miss(*args, **kwargs): + p = _prepare_pmap(fun, in_axes, out_axes, static_broadcasted_tuple, + donate_tuple, devices, backend, + axis_size, args, kwargs) + for arg in p.flat_args: + dispatch.check_arg(arg) + + params = dict( + backend=backend, + axis_name=axis_name, + axis_size=p.local_axis_size, + global_axis_size=p.global_axis_size, + devices=p.devices, + in_axes=p.in_axes_flat, + out_axes_thunk=p.out_axes_thunk, + name=p.flat_fun.__name__, + donated_invars=p.donated_invars, + is_explicit_global_axis_size=p.is_explicit_global_axis_size, + ) + + map_bind_continuation, top_trace, fun_, tracers, params = ( + core.map_bind_with_continuation(pxla.xla_pmap_p, p.flat_fun, + *p.flat_args, **params)) + execute: Callable | None = None + if isinstance(top_trace, core.EvalTrace): + execute = pxla.xla_pmap_impl_lazy(fun_, *tracers, **params) + out = map_bind_continuation(execute(*tracers)) + else: + out = map_bind_continuation( + pxla.xla_pmap_p.process(top_trace, fun_, tracers, params)) + + out_tree, out_flat = p.out_tree, out + out_pytree_def = out_tree() + out = tree_unflatten(out_pytree_def, out_flat) + + ### Decide whether we can support the C++ fast path + use_fastpath = False + if execute is not None and isinstance(execute, pxla.ExecuteReplicated): + execute_replicated = typing.cast(pxla.ExecuteReplicated, execute) + use_fastpath = ( + # TODO(sharadmv): Enable effects in replicated computation + not execute_replicated.has_unordered_effects + and not execute_replicated.has_host_callbacks and + # No tracers in the outputs. + all(isinstance(x, xc.ArrayImpl) for x in out_flat)) + + ### If we can use the fastpath, we return required info to the caller. + if use_fastpath: + execute_replicated = typing.cast(pxla.ExecuteReplicated, execute) + out_handler = execute_replicated.out_handler + in_handler = execute_replicated.in_handler + + out_array_shardings = [out.sharding for out in out_flat] + out_committed = [out._committed for out in out_flat] + fastpath_data = _PmapFastpathData( + version=1, + xla_executable=execute_replicated.xla_executable, + in_handler=in_handler, + out_handler=out_handler, + out_pytree_def=out_pytree_def, + input_devices=in_handler.local_devices, + input_indices=in_handler.input_indices, + input_array_shardings=in_handler.in_shardings, + out_avals=out_handler.out_avals, + out_array_shardings=out_array_shardings, + out_committed=out_committed, + ) + + else: + fastpath_data = None + + return out, fastpath_data + + cpp_mapped_f = pmap_lib.pmap( + fun, cache_miss, static_broadcasted_tuple, + pxla.shard_arg, pytree_registry=tree_util.default_registry) + _pmap_cache_clears.add(cpp_mapped_f) + + pmap_f = wraps(fun)(cpp_mapped_f) + + @api_boundary + def lower(*args, **kwargs): + return trace(*args, **kwargs).lower() + + @api_boundary + def trace(*args, **kwargs): + lowering_parameters = kwargs.pop( + '_experimental_lowering_parameters', mlir.LoweringParameters()) + p = _prepare_pmap( + fun, in_axes, out_axes, static_broadcasted_tuple, donate_tuple, + devices, backend, axis_size, args, kwargs) + abstract_args = list(map(shaped_abstractify, p.flat_args)) + closed_jaxpr, xc_backend, replicas, shards, pci = pxla.get_pmap_jaxpr( + p.flat_fun, backend, axis_name, + axis_size=p.local_axis_size, global_axis_size=p.global_axis_size, + devices=p.devices, + name=p.flat_fun.__name__, + in_axes=p.in_axes_flat, + out_axes_thunk=p.out_axes_thunk, + avals=abstract_args) + lower_callable = partial( + pxla.lower_parallel_callable, p.flat_fun, axis_name, + axis_size=p.local_axis_size, global_axis_size=p.global_axis_size, + devices=p.devices, + name=p.flat_fun.__name__, + in_axes=p.in_axes_flat, + donated_invars=p.donated_invars, + is_explicit_global_axis_size=p.is_explicit_global_axis_size, + avals=abstract_args, + lowering_parameters=lowering_parameters, + closed_jaxpr=closed_jaxpr, + backend=xc_backend, + replicas=replicas, + shards=shards, + pci=pci) + args_info = stages.make_args_info(p.in_tree, abstract_args, donate_tuple) + return stages.Traced(closed_jaxpr, args_info, p.flat_fun.__name__, + p.out_tree(), lower_callable) + + pmap_f.lower = lower + pmap_f.trace = trace + + return pmap_f + +_pmap_cache_clears = weakref.WeakSet() # type: ignore + + +def jvp( + fun: Callable, primals, tangents, has_aux: bool = False + ) -> tuple[Any, ...]: + """Computes a (forward-mode) Jacobian-vector product of ``fun``. + + Args: + fun: Function to be differentiated. Its arguments should be arrays, scalars, + or standard Python containers of arrays or scalars. It should return an + array, scalar, or standard Python container of arrays or scalars. + primals: The primal values at which the Jacobian of ``fun`` should be + evaluated. Should be either a tuple or a list of arguments, + and its length should be equal to the number of positional parameters of + ``fun``. + tangents: The tangent vector for which the Jacobian-vector product should be + evaluated. Should be either a tuple or a list of tangents, with the same + tree structure and array shapes as ``primals``. + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + + Returns: + If ``has_aux`` is ``False``, returns a ``(primals_out, tangents_out)`` pair, + where ``primals_out`` is ``fun(*primals)``, + and ``tangents_out`` is the Jacobian-vector product of + ``function`` evaluated at ``primals`` with ``tangents``. The + ``tangents_out`` value has the same Python tree structure and shapes as + ``primals_out``. If ``has_aux`` is ``True``, returns a + ``(primals_out, tangents_out, aux)`` tuple where ``aux`` + is the auxiliary data returned by ``fun``. + + For example: + + >>> import jax + >>> + >>> primals, tangents = jax.jvp(jax.numpy.sin, (0.1,), (0.2,)) + >>> print(primals) + 0.09983342 + >>> print(tangents) + 0.19900084 + """ + check_callable(fun) + return _jvp(lu.wrap_init(fun), primals, tangents, has_aux=has_aux) + +def _jvp(fun: lu.WrappedFun, primals, tangents, has_aux=False): + """Variant of jvp() that takes an lu.WrappedFun.""" + if (not isinstance(primals, (tuple, list)) or + not isinstance(tangents, (tuple, list))): + raise TypeError("primal and tangent arguments to jax.jvp must be tuples or lists; " + f"found {type(primals).__name__} and {type(tangents).__name__}.") + + ps_flat, tree_def = tree_flatten(primals) + ts_flat, tree_def_2 = tree_flatten(tangents) + if tree_def != tree_def_2: + raise TypeError("primal and tangent arguments to jax.jvp must have the same tree " + f"structure; primals have tree structure {tree_def} whereas tangents have " + f"tree structure {tree_def_2}.") + for p, t in zip(ps_flat, ts_flat): + if core.primal_dtype_to_tangent_dtype(_dtype(p)) != _dtype(t): + raise TypeError("primal and tangent arguments to jax.jvp do not match; " + "dtypes must be equal, or in case of int/bool primal dtype " + "the tangent dtype must be float0." + f"Got primal dtype {_dtype(p)} and so expected tangent dtype " + f"{core.primal_dtype_to_tangent_dtype(_dtype(p))}, but got " + f"tangent dtype {_dtype(t)} instead.") + if np.shape(p) != np.shape(t): + raise ValueError("jvp called with different primal and tangent shapes;" + f"Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}") + + if not has_aux: + flat_fun, out_tree = flatten_fun_nokwargs(fun, tree_def) + out_primals, out_tangents = ad.jvp(flat_fun).call_wrapped(ps_flat, ts_flat) + out_tree = out_tree() + return (tree_unflatten(out_tree, out_primals), + tree_unflatten(out_tree, out_tangents)) + else: + flat_fun, out_aux_trees = flatten_fun_nokwargs2(fun, tree_def) + jvp_fun, aux = ad.jvp(flat_fun, has_aux=True) + out_primals, out_tangents = jvp_fun.call_wrapped(ps_flat, ts_flat) + out_tree, aux_tree = out_aux_trees() + return (tree_unflatten(out_tree, out_primals), + tree_unflatten(out_tree, out_tangents), + tree_unflatten(aux_tree, aux())) + +@overload +def linearize(fun: Callable, *primals, has_aux: Literal[False] = False + ) -> tuple[Any, Callable]: + ... + +@overload +def linearize(fun: Callable, *primals, has_aux: Literal[True] + ) -> tuple[Any, Callable, Any]: + ... + +def linearize(fun: Callable, *primals, has_aux: bool = False + ) -> tuple[Any, Callable] | tuple[Any, Callable, Any]: + """Produces a linear approximation to ``fun`` using :py:func:`jvp` and partial eval. + + Args: + fun: Function to be differentiated. Its arguments should be arrays, scalars, + or standard Python containers of arrays or scalars. It should return an + array, scalar, or standard python container of arrays or scalars. + primals: The primal values at which the Jacobian of ``fun`` should be + evaluated. Should be a tuple of arrays, scalar, or standard Python + container thereof. The length of the tuple is equal to the number of + positional parameters of ``fun``. + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the first + element is considered the output of the mathematical function to be linearized, + and the second is auxiliary data. Default False. + + Returns: + If ``has_aux`` is ``False``, returns a pair where the first element is the value of + ``f(*primals)`` and the second element is a function that evaluates the + (forward-mode) Jacobian-vector product of ``fun`` evaluated at ``primals`` without + re-doing the linearization work. If ``has_aux`` is ``True``, returns a + ``(primals_out, lin_fn, aux)`` tuple where ``aux`` is the auxiliary data returned by + ``fun``. + + In terms of values computed, :py:func:`linearize` behaves much like a curried + :py:func:`jvp`, where these two code blocks compute the same values:: + + y, out_tangent = jax.jvp(f, (x,), (in_tangent,)) + + y, f_jvp = jax.linearize(f, x) + out_tangent = f_jvp(in_tangent) + + However, the difference is that :py:func:`linearize` uses partial evaluation + so that the function ``f`` is not re-linearized on calls to ``f_jvp``. In + general that means the memory usage scales with the size of the computation, + much like in reverse-mode. (Indeed, :py:func:`linearize` has a similar + signature to :py:func:`vjp`!) + + This function is mainly useful if you want to apply ``f_jvp`` multiple times, + i.e. to evaluate a pushforward for many different input tangent vectors at the + same linearization point. Moreover if all the input tangent vectors are known + at once, it can be more efficient to vectorize using :py:func:`vmap`, as in:: + + pushfwd = partial(jvp, f, (x,)) + y, out_tangents = vmap(pushfwd, out_axes=(None, 0))((in_tangents,)) + + By using :py:func:`vmap` and :py:func:`jvp` together like this we avoid the stored-linearization + memory cost that scales with the depth of the computation, which is incurred + by both :py:func:`linearize` and :py:func:`vjp`. + + Here's a more complete example of using :py:func:`linearize`: + + >>> import jax + >>> import jax.numpy as jnp + >>> + >>> def f(x): return 3. * jnp.sin(x) + jnp.cos(x / 2.) + ... + >>> jax.jvp(f, (2.,), (3.,)) + (Array(3.26819, dtype=float32, weak_type=True), Array(-5.00753, dtype=float32, weak_type=True)) + >>> y, f_jvp = jax.linearize(f, 2.) + >>> print(y) + 3.2681944 + >>> print(f_jvp(3.)) + -5.007528 + >>> print(f_jvp(4.)) + -6.676704 + """ + check_callable(fun) + f = lu.wrap_init(fun) + primals_flat, in_tree = tree_flatten(primals) + if has_aux: + jaxtree_fun, out_tree = flatten_fun_nokwargs2(f, in_tree) + else: + jaxtree_fun, out_tree = flatten_fun_nokwargs(f, in_tree) + out_primals, out_pvals, jaxpr, consts, *maybe_aux = ad.linearize( + jaxtree_fun, *primals_flat, has_aux=has_aux) + if has_aux: + out_tree, aux_tree = out_tree() + else: + out_tree = out_tree() + out_primal_py = tree_unflatten(out_tree, out_primals) + primal_avals = list(map(core.get_aval, primals_flat)) + # Ensure that lifted_jvp is a PyTree + lifted_jvp = Partial(partial(_lift_linearized, jaxpr, primal_avals, + (in_tree, out_tree), out_pvals), consts) + if has_aux: + [aux] = maybe_aux + return out_primal_py, lifted_jvp, tree_unflatten(aux_tree, aux) + else: + [] = maybe_aux + return out_primal_py, lifted_jvp + +def _lift_linearized(jaxpr, primal_avals, io_tree, out_pvals, consts, *py_args): + def fun(*tangents): + tangent_avals = list(map(core.get_aval, tangents)) + for primal_aval, tangent_aval in zip(primal_avals, tangent_avals): + if not core.typecompat(primal_aval.at_least_vspace(), tangent_aval): + raise ValueError("linearized function called on tangent values inconsistent with " + "the original primal values: " + f"got {tangent_aval} for primal aval {primal_aval}") + tangents_out = eval_jaxpr(jaxpr, consts, *tangents) + tangents_out_ = iter(tangents_out) + full_out = [pval.get_known() if pval.is_known() else next(tangents_out_) + for pval in out_pvals] + assert next(tangents_out_, None) is None + return full_out + + return apply_flat_fun_nokwargs(fun, io_tree, py_args) + +def _vjp_pullback_wrapper(name, out_primal_avals, io_tree, fun, *py_args_): + if len(py_args_) != 1: + msg = (f"The function returned by `jax.vjp` applied to {name} was called " + f"with {len(py_args_)} arguments, but functions returned by " + "`jax.vjp` must be called with a single argument corresponding to " + f"the single value returned by {name} (even if that returned " + "value is a tuple or other container).\n" + "\n" + "For example, if we have:\n" + "\n" + " def f(x):\n" + " return (x, x)\n" + " _, f_vjp = jax.vjp(f, 1.0)\n" + "\n" + "the function `f` returns a single tuple as output, and so we call " + "`f_vjp` with a single tuple as its argument:\n" + "\n" + " x_bar, = f_vjp((2.0, 2.0))\n" + "\n" + "If we instead call `f_vjp(2.0, 2.0)`, with the values 'splatted " + "out' as arguments rather than in a tuple, this error can arise.") + raise TypeError(msg) + py_args, = py_args_ + in_tree_expected, out_tree = io_tree + args, in_tree = tree_flatten(py_args) + if in_tree != in_tree_expected: + raise ValueError(f"unexpected tree structure of argument to vjp function: " + f"got {in_tree}, but expected to match {in_tree_expected}") + for arg, aval in zip(args, out_primal_avals): + ct_aval = shaped_abstractify(arg) + ct_aval_expected = aval.at_least_vspace() + if (not core.typecompat(ct_aval, ct_aval_expected) and + not _temporary_dtype_exception(ct_aval, ct_aval_expected)): + raise ValueError( + "unexpected JAX type (e.g. shape/dtype) for argument to vjp function: " + f"got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} " + f"because the corresponding output of the function {name} had JAX type " + f"{aval.str_short()}") + ans = fun(*args) + return tree_unflatten(out_tree, ans) + +# TODO(mattjj): see similar function in custom_derivatives.py +def _temporary_dtype_exception(a, a_) -> bool: + if isinstance(a, core.ShapedArray) and isinstance(a_, core.ShapedArray): + return a.shape == a_.shape and a_.dtype == float0 + return False + +@overload +def vjp(fun: Callable[..., T], + *primals: Any, + has_aux: Literal[False] = False, + reduce_axes: Sequence[AxisName] = ()) -> tuple[T, Callable]: + ... + +@overload +def vjp(fun: Callable[..., tuple[T, U]], *primals: Any, + has_aux: Literal[True], + reduce_axes: Sequence[AxisName] = ()) -> tuple[T, Callable, U]: + ... +def vjp( + fun: Callable, *primals, has_aux: bool = False, reduce_axes=() + ) -> tuple[Any, Callable] | tuple[Any, Callable, Any]: + """Compute a (reverse-mode) vector-Jacobian product of ``fun``. + + :py:func:`grad` is implemented as a special case of :py:func:`vjp`. + + Args: + fun: Function to be differentiated. Its arguments should be arrays, scalars, + or standard Python containers of arrays or scalars. It should return an + array, scalar, or standard Python container of arrays or scalars. + primals: A sequence of primal values at which the Jacobian of ``fun`` + should be evaluated. The number of ``primals`` should be equal to the + number of positional parameters of ``fun``. Each primal value should be + an array, a scalar, or a pytree (standard Python containers) thereof. + has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the + first element is considered the output of the mathematical function to be + differentiated and the second element is auxiliary data. Default False. + + Returns: + If ``has_aux`` is ``False``, returns a ``(primals_out, vjpfun)`` pair, where + ``primals_out`` is ``fun(*primals)``. If ``has_aux`` is ``True``, returns a + ``(primals_out, vjpfun, aux)`` tuple where ``aux`` is the auxiliary data + returned by ``fun``. + + ``vjpfun`` is a function from a cotangent vector with the same shape as + ``primals_out`` to a tuple of cotangent vectors with the same number and + shapes as ``primals``, representing the vector-Jacobian product of ``fun`` + evaluated at ``primals``. + + >>> import jax + >>> + >>> def f(x, y): + ... return jax.numpy.sin(x), jax.numpy.cos(y) + ... + >>> primals, f_vjp = jax.vjp(f, 0.5, 1.0) + >>> xbar, ybar = f_vjp((-0.7, 0.3)) + >>> print(xbar) + -0.61430776 + >>> print(ybar) + -0.2524413 + """ + if reduce_axes: + raise NotImplementedError("reduce_axes argument to vjp is deprecated") + del reduce_axes + check_callable(fun) + return _vjp( + lu.wrap_init(fun), *primals, has_aux=has_aux) + +def _vjp(fun: lu.WrappedFun, *primals, has_aux=False): + """Variant of vjp() that takes an lu.WrappedFun.""" + primals_flat, in_tree = tree_flatten(primals) + for arg in primals_flat: dispatch.check_arg(arg) + if not has_aux: + flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree) + out_primals, vjp = ad.vjp(flat_fun, primals_flat) + out_tree = out_tree() + else: + flat_fun, out_aux_trees = flatten_fun_nokwargs2(fun, in_tree) + out_primals, vjp, aux = ad.vjp(flat_fun, primals_flat, has_aux=True) + out_tree, aux_tree = out_aux_trees() + out_primal_avals = map(shaped_abstractify, out_primals) + out_primal_py = tree_unflatten(out_tree, out_primals) + vjp_py = Partial(partial(_vjp_pullback_wrapper, fun.__name__, + out_primal_avals, (out_tree, in_tree)), vjp) + if not has_aux: + return out_primal_py, vjp_py + else: + return out_primal_py, vjp_py, tree_unflatten(aux_tree, aux) + + +def linear_transpose(fun: Callable, *primals, reduce_axes=()) -> Callable: + """Transpose a function that is promised to be linear. + + For linear functions, this transformation is equivalent to :py:func:`vjp`, but + avoids the overhead of computing the forward pass. + + The outputs of the transposed function will always have the exact same dtypes + as ``primals``, even if some values are truncated (e.g., from complex to + float, or from float64 to float32). To avoid truncation, use dtypes in + ``primals`` that match the full range of desired outputs from the transposed + function. Integer dtypes are not supported. + + Args: + fun: the linear function to be transposed. + *primals: a positional argument tuple of arrays, scalars, or (nested) + standard Python containers (tuples, lists, dicts, namedtuples, i.e., + pytrees) of those types used for evaluating the shape/dtype of + ``fun(*primals)``. These arguments may be real scalars/ndarrays, but that + is not required: only the ``shape`` and ``dtype`` attributes are accessed. + See below for an example. (Note that the duck-typed objects cannot be + namedtuples because those are treated as standard Python containers.) + + Returns: + A callable that calculates the transpose of ``fun``. Valid input into this + function must have the same shape/dtypes/structure as the result of + ``fun(*primals)``. Output will be a tuple, with the same + shape/dtypes/structure as ``primals``. + + >>> import jax + >>> import types + >>> + >>> f = lambda x, y: 0.5 * x - 0.5 * y + >>> scalar = types.SimpleNamespace(shape=(), dtype=np.dtype(np.float32)) + >>> f_transpose = jax.linear_transpose(f, scalar, scalar) + >>> f_transpose(1.0) + (Array(0.5, dtype=float32), Array(-0.5, dtype=float32)) + """ + if reduce_axes: + raise NotImplementedError("reduce_axes argument to transpose is deprecated") + del reduce_axes + primals_flat, in_tree = tree_flatten(primals) + flat_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree) + in_avals = map(shaped_abstractify, primals_flat) + in_dtypes = map(dtypes.dtype, in_avals) + + in_pvals = map(pe.PartialVal.unknown, in_avals) + jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals, + instantiate=True) + jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True) + out_avals, _ = unzip2(out_pvals) + out_dtypes = map(dtypes.dtype, out_avals) + if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes) + or all(dtypes.issubdtype(d, np.integer) + for d in in_dtypes + out_dtypes)): + raise TypeError("linear_transpose only supports [float or complex] -> " + "[float or complex], and integer -> integer functions, " + f"but got {in_dtypes} -> {out_dtypes}.") + + @api_boundary + def transposed_fun(const, out_cotangent): + out_cts, out_tree2 = tree_flatten(out_cotangent) + if out_tree() != out_tree2: + raise TypeError("cotangent tree does not match function output, " + f"expected {out_tree()} but got {out_tree2}") + if not all(map(core.typecheck, out_avals, out_cts)): + raise TypeError("cotangent type does not match function output, " + f"expected {out_avals} but got {out_cts}") + dummies = [ad.UndefinedPrimal(a) for a in in_avals] + in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts) + in_cts = map(ad.instantiate_zeros, in_cts) + return tree_unflatten(in_tree, in_cts) + + # Ensure that transposed_fun is a PyTree + return Partial(transposed_fun, const) + + +def _flat_axes_specs(abstracted_axes, *args, **kwargs + ) -> list[pe.AbstractedAxesSpec]: + if kwargs: raise NotImplementedError + def ax_leaf(l): + return (isinstance(l, dict) and all_leaves(l.values()) or + isinstance(l, tuple) and all_leaves(l, lambda x: x is None)) + return broadcast_prefix(abstracted_axes, args, ax_leaf) + + +# TODO(phawkins): for some reason mypy cannot determine these overloads are +# non-overlapping. Pytype is happy with them. +@overload +def make_jaxpr(fun: Callable, # type: ignore + static_argnums: int | Iterable[int] = (), + axis_env: Sequence[tuple[AxisName, int]] | None = None, + return_shape: Literal[False] = ..., + abstracted_axes: Any | None = None, + ) -> Callable[..., core.ClosedJaxpr]: + ... + +@overload +def make_jaxpr(fun: Callable, + static_argnums: int | Iterable[int] = (), + axis_env: Sequence[tuple[AxisName, int]] | None = None, + return_shape: Literal[True] = ..., + abstracted_axes: Any | None = None, + ) -> Callable[..., tuple[core.ClosedJaxpr, Any]]: + ... + +def make_jaxpr(fun: Callable, + static_argnums: int | Iterable[int] = (), + axis_env: Sequence[tuple[AxisName, int]] | None = None, + return_shape: bool = False, + abstracted_axes: Any | None = None, + ) -> Callable[..., (core.ClosedJaxpr | + tuple[core.ClosedJaxpr, Any])]: + """Creates a function that produces its jaxpr given example args. + + Args: + fun: The function whose ``jaxpr`` is to be computed. Its positional + arguments and return value should be arrays, scalars, or standard Python + containers (tuple/list/dict) thereof. + static_argnums: See the :py:func:`jax.jit` docstring. + axis_env: Optional, a sequence of pairs where the first element is an axis + name and the second element is a positive integer representing the size of + the mapped axis with that name. This parameter is useful when lowering + functions that involve parallel communication collectives, and it + specifies the axis name/size environment that would be set up by + applications of :py:func:`jax.pmap`. + return_shape: Optional boolean, defaults to ``False``. If ``True``, the + wrapped function returns a pair where the first element is the + ``ClosedJaxpr`` representation of ``fun`` and the second element is a + pytree with the same structure as the output of ``fun`` and where the + leaves are objects with ``shape``, ``dtype``, and ``named_shape`` + attributes representing the corresponding types of the output leaves. + + Returns: + A wrapped version of ``fun`` that when applied to example arguments returns + a ``ClosedJaxpr`` representation of ``fun`` on those arguments. If the + argument ``return_shape`` is ``True``, then the returned function instead + returns a pair where the first element is the ``ClosedJaxpr`` + representation of ``fun`` and the second element is a pytree representing + the structure, shape, dtypes, and named shapes of the output of ``fun``. + + A ``jaxpr`` is JAX's intermediate representation for program traces. The + ``jaxpr`` language is based on the simply-typed first-order lambda calculus + with let-bindings. :py:func:`make_jaxpr` adapts a function to return its + ``jaxpr``, which we can inspect to understand what JAX is doing internally. + The ``jaxpr`` returned is a trace of ``fun`` abstracted to + :py:class:`ShapedArray` level. Other levels of abstraction exist internally. + + We do not describe the semantics of the ``jaxpr`` language in detail here, but + instead give a few examples. + + >>> import jax + >>> + >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x)) + >>> print(f(3.0)) + -0.83602 + >>> jax.make_jaxpr(f)(3.0) + { lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) } + >>> jax.make_jaxpr(jax.grad(f))(3.0) + { lambda ; a:f32[]. let + b:f32[] = cos a + c:f32[] = sin a + _:f32[] = sin b + d:f32[] = cos b + e:f32[] = mul 1.0 d + f:f32[] = neg e + g:f32[] = mul f c + in (g,) } + """ + try: + hash(fun) + weakref.ref(fun) + except TypeError: + fun = partial(fun) + + @wraps(fun) + @api_boundary + def make_jaxpr_f(*args, **kwargs): + with ExitStack() as stack: + for axis_name, size in axis_env or []: + stack.enter_context(core.extend_axis_env(axis_name, size, None)) + traced = jit(fun, static_argnums=static_argnums, + abstracted_axes=abstracted_axes).trace(*args, **kwargs) + # `jit` converts tracers in consts to args but that breaks the semantics of + # `make_jaxpr`. Hence convert the tracers in args back to consts in jaxpr. + if traced._num_consts: + consts, _ = split_list(traced._args_flat, [traced._num_consts]) + jaxpr_ = pe.convert_invars_to_constvars(traced.jaxpr.jaxpr, + traced._num_consts) + jaxpr = core.ClosedJaxpr(jaxpr_, consts) + else: + jaxpr = traced.jaxpr + if return_shape: + out = [ShapeDtypeStruct(o.shape, o.dtype, getattr(o, 'named_shape', None)) + for o in jaxpr.out_avals] + return jaxpr, tree_unflatten(tree_structure(traced.out_info), out) + return jaxpr + + make_jaxpr_f.__module__ = "jax" + if hasattr(fun, "__qualname__"): + make_jaxpr_f.__qualname__ = f"make_jaxpr({fun.__qualname__})" + if hasattr(fun, "__name__"): + make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})" + return make_jaxpr_f + +def _infer_src_sharding(src, x) -> Sharding | None: + if src is not None: + # TODO(slebedev): This looks like an error and needs investigation. + return src # pytype: disable=bad-return-type + if isinstance(x, array.ArrayImpl): + return x.sharding + elif isinstance(x, core.Tracer): + aval = core.get_aval(x) + if isinstance(aval, ConcreteArray) and isinstance(aval.val, array.ArrayImpl): + return aval.val.sharding + return None + + +# TODO(yashkatariya): Generalize check_compatible_aval (maybe renamed) and use +# that to check if shardings are compatible with the input. +@lru_cache(maxsize=2048) +def _check_sharding(aval, s): + if isinstance(s, Sharding): + if isinstance(aval, core.AbstractToken): + aval = core.token_shaped_array + if not isinstance(s, PmapSharding): + pjit.pjit_check_aval_sharding( + (s,), (aval,), None, "device_put args", allow_uneven_sharding=False) + s.shard_shape(aval.shape) # should raise an Error if incompatible + + +def device_put( + x, + device: None | xc.Device | Sharding | Layout | Any | TransferToMemoryKind = None, + *, src: None | xc.Device | Sharding | Layout | Any | TransferToMemoryKind = None): + """Transfers ``x`` to ``device``. + + Args: + x: An array, scalar, or (nested) standard Python container thereof. + device: The (optional) :py:class:`Device`, :py:class:`Sharding`, or a + (nested) :py:class:`Sharding` in standard Python container (must be a tree + prefix of ``x``), representing the device(s) to which ``x`` should be + transferred. If given, then the result is committed to the device(s). + + Returns: + A copy of ``x`` that resides on ``device``. + + If the ``device`` parameter is ``None``, then this operation behaves like the + identity function if the operand is on any device already, otherwise it + transfers the data to the default device, uncommitted. + + For more details on data placement see the + :ref:`FAQ on data placement `. + + This function is always asynchronous, i.e. returns immediately without + blocking the calling Python thread until any transfers are completed. + """ + with config.explicit_device_put_scope(): + if ((device is None or + isinstance(device, (xc.Device, Sharding, TransferToMemoryKind))) and + (src is None or + isinstance(src, (xc.Device, Sharding, TransferToMemoryKind)))): + def _map(y): + _check_sharding(shaped_abstractify(y), s=device) + return dispatch.device_put_p.bind( + y, device=device, src=_infer_src_sharding(src, y)) + return tree_map(_map, x) + + x_flat, treedef = tree_flatten(x) + device_flat = flatten_axes("device_put device", treedef, device) + src_flat = flatten_axes("device_put source", treedef, src) + out_flat = [] + for xf, d, s in zip(x_flat, device_flat, src_flat): + _check_sharding(shaped_abstractify(xf), d) + out_flat.append(dispatch.device_put_p.bind( + xf, device=d, src=_infer_src_sharding(s, xf))) + return tree_unflatten(treedef, out_flat) + + +def device_put_sharded(shards: Sequence[Any], devices: Sequence[xc.Device]): # noqa: F811 + """Transfer array shards to specified devices and form Array(s). + + Args: + shards: A sequence of arrays, scalars, or (nested) standard Python + containers thereof representing the shards to be stacked together to form + the output. The length of ``shards`` must equal the length of ``devices``. + devices: A sequence of :py:class:`Device` instances representing the devices + to which corresponding shards in ``shards`` will be transferred. + + This function is always asynchronous, i.e. returns immediately. + + Returns: + A Array or (nested) Python container thereof representing the + elements of ``shards`` stacked together, with each shard backed by physical + device memory specified by the corresponding entry in ``devices``. + + Examples: + Passing a list of arrays for ``shards`` results in a sharded array + containing a stacked version of the inputs: + + >>> import jax + >>> devices = jax.local_devices() + >>> x = [jax.numpy.ones(5) for device in devices] + >>> y = jax.device_put_sharded(x, devices) + >>> np.allclose(y, jax.numpy.stack(x)) + True + + Passing a list of nested container objects with arrays at the leaves for + ``shards`` corresponds to stacking the shards at each leaf. This requires + all entries in the list to have the same tree structure: + + >>> x = [(i, jax.numpy.arange(i, i + 4)) for i in range(len(devices))] + >>> y = jax.device_put_sharded(x, devices) + >>> type(y) + + >>> y0 = jax.device_put_sharded([a for a, b in x], devices) + >>> y1 = jax.device_put_sharded([b for a, b in x], devices) + >>> np.allclose(y[0], y0) + True + >>> np.allclose(y[1], y1) + True + + See Also: + - device_put + - device_put_replicated + """ + # TODO(jakevdp): provide a default for devices that considers both local + # devices and pods + if not isinstance(shards, Sequence): + raise TypeError("device_put_sharded `shards` input must be a sequence; " + f"got {type(shards)}") + if len(shards) != len(devices): + raise ValueError(f"len(shards) = {len(shards)} must equal " + f"len(devices) = {len(devices)}.") + + def _device_put_sharded(*xs): + avals = [core.raise_to_shaped(core.get_aval(x)) for x in xs] + if not all(a1 == a2 for a1, a2 in zip(avals[:-1], avals[1:])): + a1, a2 = next((a1, a2) for a1, a2 in zip(avals[:-1], avals[1:]) + if a1 != a2) + raise ValueError("the shards passed to device_put_sharded must have " + f"consistent shape and dtype, but got {a1} and {a2}.") + stacked_aval = avals[0].update(shape=(len(devices),) + avals[0].shape) + sharding_spec = sharding_specs.create_pmap_sharding_spec(stacked_aval.shape) + sharding = PmapSharding(np.array(devices), sharding_spec) + if dtypes.issubdtype(stacked_aval.dtype, dtypes.extended): + return stacked_aval.dtype._rules.device_put_sharded(xs, stacked_aval, sharding, devices) + if config.pmap_no_rank_reduction.value: + ys = [] + for x in xs: + if not isinstance(x, (np.ndarray, basearray.Array)): + x = np.asarray(x) + ys.append(x[None]) + else: + ys = xs + return pxla.batched_device_put(stacked_aval, sharding, ys, list(devices)) + + + with config.explicit_device_put_scope(): + return tree_map(_device_put_sharded, *shards) + + +def device_put_replicated(x: Any, devices: Sequence[xc.Device]): # noqa: F811 + """Transfer array(s) to each specified device and form Array(s). + + Args: + x: an array, scalar, or (nested) standard Python container thereof + representing the array to be replicated to form the output. + devices: A sequence of :py:class:`Device` instances representing the devices + to which ``x`` will be transferred. + + This function is always asynchronous, i.e. returns immediately. + + Returns: + An Array or (nested) Python container thereof representing the + value of ``x`` broadcasted along a new leading axis of size + ``len(devices)``, with each slice along that new leading axis backed by + memory on the device specified by the corresponding entry in ``devices``. + + Examples: + Passing an array: + + >>> import jax + >>> devices = jax.local_devices() + >>> x = jax.numpy.array([1., 2., 3.]) + >>> y = jax.device_put_replicated(x, devices) + >>> np.allclose(y, jax.numpy.stack([x for _ in devices])) + True + + See Also: + - device_put + - device_put_sharded + """ + if not isinstance(devices, Sequence) or not devices: + raise ValueError("`devices` argument to `device_put_replicated must be " + "a non-empty sequence.") + def _device_put_replicated(x): + aval = core.unmapped_aval(len(devices), core.no_axis_name, 0, + core.raise_to_shaped(core.get_aval(x))) + assert isinstance(aval, ShapedArray) + sharding_spec = sharding_specs.create_pmap_sharding_spec(aval.shape) + if config.pmap_no_rank_reduction.value: + if isinstance(x, (np.ndarray, basearray.Array)): + buf = device_put(x[None], devices[0]) + else: + buf = device_put(x, devices[0])[None] + else: + buf = device_put(x, devices[0]) + sharding = PmapSharding(np.array(devices), sharding_spec) + if dtypes.issubdtype(aval.dtype, dtypes.extended): + return aval.dtype._rules.device_put_replicated(buf, aval, sharding, devices) + assert len(xla.aval_to_xla_shapes(aval)) == 1 + return pxla.batched_device_put(aval, sharding, [buf] * len(devices), devices) + + with config.explicit_device_put_scope(): + return tree_map(_device_put_replicated, x) + + +# TODO(mattjj): consider revising +def _device_get(x): + if isinstance(x, core.Tracer): + return x + try: + toarray = x.__array__ + except AttributeError: + return x + else: + return toarray() + +def device_get(x: Any): + """Transfer ``x`` to host. + + If ``x`` is a pytree, then the individual buffers are copied in parallel. + + Args: + x: An array, scalar, Array or (nested) standard Python container thereof + representing the array to be transferred to host. + + Returns: + An array or (nested) Python container thereof representing the + value of ``x``. + + Examples: + Passing a Array: + + >>> import jax + >>> x = jax.numpy.array([1., 2., 3.]) + >>> jax.device_get(x) + array([1., 2., 3.], dtype=float32) + + Passing a scalar (has no effect): + + >>> jax.device_get(1) + 1 + + See Also: + - device_put + - device_put_sharded + - device_put_replicated + """ + with config.explicit_device_get_scope(): + for y in tree_leaves(x): + try: + y.copy_to_host_async() + except AttributeError: + pass + return tree_map(_device_get, x) + + +class ShapeDtypeStruct: + """A container for the shape, dtype, and other static attributes of an array. + + ``ShapeDtypeStruct`` is often used in conjunction with :func:`jax.eval_shape`. + + Args: + shape: a sequence of integers representing an array shape + dtype: a dtype-like object + named_shape: (optional) a dictionary representing a named shape + sharding: (optional) a :class:`jax.Sharding` object + """ + __slots__ = ["shape", "dtype", "named_shape", "sharding", "_dll"] + + def __init__(self, shape, dtype, named_shape=None, sharding=None): + self.shape = tuple(shape) + if dtype is None: + raise ValueError("ShapeDtypeStruct: dtype must be specified.") + self.dtype = dtype if dtypes.issubdtype(dtype, dtypes.extended) else np.dtype(dtype) + if sharding is not None and not isinstance(sharding, (Sharding, Layout)): + raise ValueError( + "sharding should be an instance of `jax.sharding.Sharding` or" + f" `jax.experimental.layout.Layout`. Got {sharding} of type" + f" {type(sharding)}.") + if (isinstance(sharding, Layout) and + isinstance(sharding.device_local_layout, AutoLayout)): + raise TypeError( + "`DeviceLocalLayout.AUTO` cannot be used in place of a device-local" + f" layout in a `ShapeDtypeStruct`. Got {sharding}") + self.sharding = sharding.sharding if isinstance(sharding, Layout) else sharding + self._dll = sharding.device_local_layout if isinstance(sharding, Layout) else None + self.named_shape = {} if named_shape is None else dict(named_shape) + + size = property(lambda self: math.prod(self.shape)) + ndim = property(lambda self: len(self.shape)) + + @property + def layout(self): + return Layout(self._dll, self.sharding) + + def __len__(self): + try: + return self.shape[0] + except IndexError as e: + raise TypeError("len() of unsized object") from e # same as numpy error + + def __repr__(self): + ns = f", named_shape={self.named_shape}" if self.named_shape else "" + sh = f", sharding={self.sharding}" if self.sharding is not None else "" + l = f", layout={self.layout}" if self._dll is not None else "" + return (f"{type(self).__name__}(shape={self.shape}, " + f"dtype={self.dtype.name}{ns}{sh}{l})") + + __str__ = __repr__ + + def __eq__(self, other): + if not isinstance(other, ShapeDtypeStruct): + return False + else: + return ((other.shape, other.dtype, other.named_shape, other.sharding, other.layout) == + (self.shape, self.dtype, self.named_shape, self.sharding, self.layout)) + + def __hash__(self): + # TODO(frostig): avoid the conversion from dict by addressing + # https://github.com/google/jax/issues/8182 + named = frozenset(self.named_shape.items()) + return hash((self.shape, self.dtype, named, self.sharding, self.layout)) + + +core.pytype_aval_mappings[ShapeDtypeStruct] = ( + lambda x: ShapedArray(x.shape, dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True), + weak_type=False, named_shape=x.named_shape)) + + +@api_boundary +def eval_shape(fun: Callable, *args, **kwargs): + """Compute the shape/dtype of ``fun`` without any FLOPs. + + This utility function is useful for performing shape inference. Its + input/output behavior is defined by:: + + def eval_shape(fun, *args, **kwargs): + out = fun(*args, **kwargs) + shape_dtype_struct = lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype) + return jax.tree_util.tree_map(shape_dtype_struct, out) + + But instead of applying ``fun`` directly, which might be expensive, it uses + JAX's abstract interpretation machinery to evaluate the shapes without doing + any FLOPs. + + Using :py:func:`eval_shape` can also catch shape errors, and will raise same + shape errors as evaluating ``fun(*args, **kwargs)``. + + Args: + fun: The function whose output shape should be evaluated. + *args: a positional argument tuple of arrays, scalars, or (nested) standard + Python containers (tuples, lists, dicts, namedtuples, i.e. pytrees) of + those types. Since only the ``shape`` and ``dtype`` attributes are + accessed, one can use :class:`jax.ShapeDtypeStruct` or another container + that duck-types as ndarrays (note however that duck-typed objects cannot + be namedtuples because those are treated as standard Python containers). + **kwargs: a keyword argument dict of arrays, scalars, or (nested) standard + Python containers (pytrees) of those types. As in ``args``, array values + need only be duck-typed to have ``shape`` and ``dtype`` attributes. + + Returns: + out: a nested PyTree containing :class:`jax.ShapeDtypeStruct` objects as leaves. + + For example: + + >>> import jax + >>> import jax.numpy as jnp + >>> + >>> f = lambda A, x: jnp.tanh(jnp.dot(A, x)) + >>> A = jax.ShapeDtypeStruct((2000, 3000), jnp.float32) + >>> x = jax.ShapeDtypeStruct((3000, 1000), jnp.float32) + >>> out = jax.eval_shape(f, A, x) # no FLOPs performed + >>> print(out.shape) + (2000, 1000) + >>> print(out.dtype) + float32 + + All arguments passed via :func:`eval_shape` will be treated as dynamic; + static arguments can be included via closure, for example using :func:`functools.partial`: + + >>> import jax + >>> from jax import lax + >>> from functools import partial + >>> import jax.numpy as jnp + >>> + >>> x = jax.ShapeDtypeStruct((1, 1, 28, 28), jnp.float32) + >>> kernel = jax.ShapeDtypeStruct((32, 1, 3, 3), jnp.float32) + >>> + >>> conv_same = partial(lax.conv_general_dilated, window_strides=(1, 1), padding="SAME") + >>> out = jax.eval_shape(conv_same, x, kernel) + >>> print(out.shape) + (1, 32, 28, 28) + >>> print(out.dtype) + float32 + """ + try: hash(fun) + except TypeError: fun = partial(fun) + return jit(fun).eval_shape(*args, **kwargs) + + +def named_call( + fun: F, + *, + name: str | None = None, +) -> F: + """Adds a user specified name to a function when staging out JAX computations. + + When staging out computations for just-in-time compilation to XLA (or other + backends such as TensorFlow) JAX runs your Python program but by default does + not preserve any of the function names or other metadata associated with it. + This can make debugging the staged out (and/or compiled) representation of + your program complicated because there is limited context information for each + operation being executed. + + `named_call` tells JAX to stage the given function out as a subcomputation + with a specific name. When the staged out program is compiled with XLA these + named subcomputations are preserved and show up in debugging utilities like + the TensorFlow Profiler in TensorBoard. Names are also preserved when staging + out JAX programs to TensorFlow using :func:`experimental.jax2tf.convert`. + + Args: + fun: Function to be wrapped. This can be any Callable. + name: Optional. The prefix to use to name all sub computations created + within the name scope. Use the fun.__name__ if not specified. + + Returns: + A version of `fun` that is wrapped in a name_scope. + """ + if name is None: + name = fun.__name__ + + return source_info_util.extend_name_stack(name)(fun) + + +@contextmanager +def named_scope( + name: str, + ) -> Generator[None, None, None]: + """A context manager that adds a user specified name to the JAX name stack. + + When staging out computations for just-in-time compilation to XLA (or other + backends such as TensorFlow) JAX does not, by default, preserve the names + (or other source metadata) of Python functions it encounters. + This can make debugging the staged out (and/or compiled) representation of + your program complicated because there is limited context information for each + operation being executed. + + ``named_scope`` tells JAX to stage the given function with additional + annotations on the underlying operations. JAX internally keeps track of these + annotations in a name stack. When the staged out program is compiled with XLA + these annotations are preserved and show up in debugging utilities like the + TensorFlow Profiler in TensorBoard. Names are also preserved when staging out + JAX programs to TensorFlow using :func:`experimental.jax2tf.convert`. + + + Args: + name: The prefix to use to name all operations created within the name + scope. + Yields: + Yields ``None``, but enters a context in which `name` will be appended to + the active name stack. + + Examples: + ``named_scope`` can be used as a context manager inside compiled functions: + + >>> import jax + >>> + >>> @jax.jit + ... def layer(w, x): + ... with jax.named_scope("dot_product"): + ... logits = w.dot(x) + ... with jax.named_scope("activation"): + ... return jax.nn.relu(logits) + + It can also be used as a decorator: + + >>> @jax.jit + ... @jax.named_scope("layer") + ... def layer(w, x): + ... logits = w.dot(x) + ... return jax.nn.relu(logits) + """ + if not isinstance(name, str): + raise TypeError("named_scope name argument must be a string.") + with source_info_util.extend_name_stack(name): + yield + +def effects_barrier(): + """Waits until existing functions have completed any side-effects.""" + dispatch.runtime_tokens.block_until_ready() + +def block_until_ready(x): + """ + Tries to call a ``block_until_ready`` method on pytree leaves. + + Args: + x: a pytree, usually with at least some JAX array instances at its leaves. + + Returns: + A pytree with the same structure and values of the input, where the values + of all JAX array leaves are ready. + """ + def try_to_block(x): + try: + return x.block_until_ready() + except AttributeError: + return x + + arrays = [] + for leaf in tree_leaves(x): + if isinstance(leaf, array.ArrayImpl): + arrays.append(leaf) + else: + try_to_block(leaf) + + if not arrays: + # `arrays` will be empty if tree_leaves(x) is empty or all leaves are not + # jax.Array. + pass + elif len(arrays) == 1: + # Fast path for single array. + try_to_block(arrays[0]) + else: + # Optimized for multiple arrays. + xc.batched_block_until_ready(arrays) + + return x + + +def clear_backends(): + """ + Clear all backend clients so that new backend clients can be created later. + """ + xb._clear_backends() + xb.local_devices.cache_clear() + xb.process_count.cache_clear() + dispatch.xla_primitive_callable.cache_clear() + pjit._pjit_lower_cached.cache_clear() + pjit._create_pjit_jaxpr.cache_clear() # pytype: disable=attribute-error + pjit._cpp_pjit_cache.clear() + xc._xla.PjitFunctionCache.clear_all() + +def live_arrays(platform=None): + """Return all live arrays in the backend for `platform`. + + If platform is None, it is the default backend. + """ + return xb.get_backend(platform).live_arrays() + +def clear_caches(): + """Clear all compilation and staging caches. + + This doesn't clear the persistent cache; to disable it (e.g. for benchmarks), + set the jax_enable_compilation_cache config option to False. + """ + # Clear all lu.cache and util.weakref_lru_cache instances (used for staging + # and Python-dispatch compiled executable caches). + lu.clear_all_caches() + util.clear_all_weakref_lru_caches() + + # Clear all C++ compiled executable caches for pjit + pjit._cpp_pjit_cache.clear() + xc._xla.PjitFunctionCache.clear_all() + + # Clear all C++ compiled executable caches for pmap + for fun in _pmap_cache_clears: + fun._cache_clear() + + # Clear particular util.cache instances. + dispatch.xla_primitive_callable.cache_clear() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/api_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/api_util.py new file mode 100644 index 0000000000000000000000000000000000000000..67a8ee8bdfacebfb4c469de481aa5ad3e4ce664c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/api_util.py @@ -0,0 +1,718 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +import inspect +import operator +from functools import partial, lru_cache +from typing import Any, Callable, Type + +import numpy as np + +from jax._src import core +from jax._src import dtypes +from jax._src.abstract_arrays import numpy_scalar_types +from jax._src.core import ShapedArray +from jax._src.tree_util import ( + PyTreeDef, tree_flatten, tree_unflatten, tree_map, + treedef_children, generate_key_paths, keystr, broadcast_prefix, + prefix_errors) +from jax._src.tree_util import _replace_nones +from jax._src import linear_util as lu +from jax._src.linear_util import TracingDebugInfo +from jax._src.util import (safe_map, WrapKwArgs, Hashable, HashableFunction, + Unhashable, safe_zip) +from jax._src import traceback_util +traceback_util.register_exclusion(__file__) + +map = safe_map + +def _ensure_index(x: Any) -> int | tuple[int, ...]: + """Ensure x is either an index or a tuple of indices.""" + x = core.concrete_or_error(None, x, "expected a static index or sequence of indices.") + try: + return operator.index(x) + except TypeError: + return tuple(map(operator.index, x)) + +def _ensure_index_tuple(x: Any) -> tuple[int, ...]: + """Convert x to a tuple of indices.""" + x = core.concrete_or_error(None, x, "expected a static index or sequence of indices.") + try: + return (operator.index(x),) + except TypeError: + return tuple(map(operator.index, x)) + +def _ensure_str(x: str) -> str: + if not isinstance(x, str): + raise TypeError(f"argument is not a string: {x}") + return x + +def _ensure_str_tuple(x: str | Iterable[str]) -> tuple[str, ...]: + """Convert x to a tuple of strings.""" + if isinstance(x, str): + return (x,) + else: + return tuple(map(_ensure_str, x)) + +@lu.transformation_with_aux +def flatten_fun(in_tree, *args_flat): + py_args, py_kwargs = tree_unflatten(in_tree, args_flat) + ans = yield py_args, py_kwargs + yield tree_flatten(ans) + +def apply_flat_fun(fun, io_tree, *py_args): + in_tree_expected, out_tree = io_tree + args, in_tree = tree_flatten((py_args, {})) + if in_tree != in_tree_expected: + raise TypeError(f"Expected {in_tree_expected}, got {in_tree}") + ans = fun(*args) + return tree_unflatten(out_tree, ans) + +@lu.transformation_with_aux +def flatten_fun_nokwargs(in_tree, *args_flat): + py_args = tree_unflatten(in_tree, args_flat) + ans = yield py_args, {} + yield tree_flatten(ans) + +def apply_flat_fun_nokwargs(fun, io_tree, py_args): + in_tree_expected, out_tree = io_tree + args, in_tree = tree_flatten(py_args) + if in_tree != in_tree_expected: + raise TypeError(f"Expected {in_tree_expected}, got {in_tree}") + ans = fun(*args) + return tree_unflatten(out_tree, ans) + +def flattened_fun_in_tree( + fn: lu.WrappedFun + ) -> tuple[PyTreeDef, Callable[[], PyTreeDef], bool] | None: + # This implementation relies on internal details of linear_util.py's + # WrappedFun, but it's for the worthy cause of better user error messages. + # It can fail (i.e. return None) if its WrappedFun argument is not transformed + # with flatten_fun or flatten_fun_nokwargs, which could happen e.g. when + # core.eval_jaxpr encounters a call primitive (though at that point we're just + # round-tripping jaxprs and the user errors in question are impossible). + assert isinstance(flatten_fun, partial) and len(flatten_fun.args) == 1 + assert (isinstance(flatten_fun_nokwargs, partial) and + len(flatten_fun_nokwargs.args) == 1) + flattens = {flatten_fun.args[0], flatten_fun_nokwargs.args[0]} + try: + ((in_tree,), out_tree_store, has_kwargs), = ( + (args, store, f is flatten_fun.args[0]) + for (f, args), store in zip(fn.transforms, fn.stores) if f in flattens) + except ValueError: + return None + else: + return in_tree, lambda: out_tree_store.val, has_kwargs + +@lu.transformation_with_aux +def flatten_fun_nokwargs2(in_tree, *args_flat): + py_args = tree_unflatten(in_tree, args_flat) + pair = yield py_args, {} + if not isinstance(pair, (list, tuple)) or len(pair) != 2: + raise TypeError("expected function with aux output to return a two-element " + f"tuple, but got type {type(pair)} with value {pair!r}") + ans, aux = pair + ans_flat, ans_tree = tree_flatten(ans) + aux_flat, aux_tree = tree_flatten(aux) + yield (ans_flat, aux_flat), (ans_tree, aux_tree) + +class _HashableWithStrictTypeEquality: + """Box object used when comparing static arguments as a jit key. + + Requires exact type equality using `is` and value equality.""" + __slots__ = ["val"] + + def __init__(self, val): + self.val = val + + def __hash__(self): + return hash(self.val) + + def __eq__(self, other): + return type(self.val) is type(other.val) and self.val == other.val + +_POSITIONAL_ARGUMENTS = ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD +) + +def _validate_argnums(sig: inspect.Signature, argnums: tuple[int, ...], argnums_name: str) -> None: + """ + Validate that the argnums are sensible for a given function. + + For functions that accept a variable number of positions arguments + (`f(..., *args)`) all positive argnums are considered valid. + """ + n_pos_args = 0 + for param in sig.parameters.values(): + if param.kind in _POSITIONAL_ARGUMENTS: + n_pos_args += 1 + + elif param.kind is inspect.Parameter.VAR_POSITIONAL: + # We can have any number of positional arguments + return + + if argnums and (-min(argnums) > n_pos_args or max(argnums) >= n_pos_args): + raise ValueError(f"Jitted function has {argnums_name}={argnums}, " + f"but only accepts {n_pos_args} positional arguments.") + +_INVALID_KEYWORD_ARGUMENTS = ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL +) + + +_KEYWORD_ARGUMENTS = ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, +) +def _validate_argnames( + sig: inspect.Signature, argnames: tuple[str, ...], argnames_name: str +) -> None: + """ + Validate that the argnames are sensible for a given function. + + For functions that accept a variable keyword arguments + (`f(..., **kwargs)`) all argnames are considered valid except those + marked as position-only (`f(pos_only, /, ...)`). + """ + var_kwargs = False + valid_kwargs: set[str] = set() + invalid_kwargs: set[str] = set() + for param_name, param in sig.parameters.items(): + if param.kind in _KEYWORD_ARGUMENTS: + valid_kwargs.add(param_name) + + elif param.kind is inspect.Parameter.VAR_KEYWORD: + var_kwargs = True + + elif param.kind in _INVALID_KEYWORD_ARGUMENTS: + invalid_kwargs.add(param_name) + + # Check whether any kwargs are invalid due to position only + if invalid_argnames := (invalid_kwargs & set(argnames)): + raise ValueError(f"Jitted function has invalid argnames {invalid_argnames} " + f"in {argnames_name}. These are positional-only") + + # Takes any kwargs + if var_kwargs: + return + + # Check that all argnames exist on function + if invalid_argnames := (set(argnames) - valid_kwargs): + raise ValueError(f"Jitted function has invalid argnames {invalid_argnames} " + f"in {argnames_name}. Function does not take these args.") + + +def argnums_partial(f, dyn_argnums, args, require_static_args_hashable=True): + dyn_argnums = _ensure_index_tuple(dyn_argnums) + dyn_argnums = _ensure_inbounds(False, len(args), dyn_argnums) + if require_static_args_hashable: + fixed_args = [] + for i, arg in enumerate(args): + if i in dyn_argnums: continue + if not is_hashable(arg): + raise ValueError( + "Non-hashable static arguments are not supported, as this can lead " + f"to unexpected cache-misses. Static argument (index {i}) of type " + f"{type(arg)} for function {f.__name__} is non-hashable.") + fixed_args.append(_HashableWithStrictTypeEquality(arg)) + else: + fixed_args = [Unhashable(arg) for i, arg in enumerate(args) + if i not in dyn_argnums] + dyn_args = tuple(args[i] for i in dyn_argnums) + return _argnums_partial(f, dyn_argnums, tuple(fixed_args)), dyn_args + +def _ensure_inbounds(allow_invalid: bool, num_args: int, argnums: Sequence[int] + ) -> tuple[int, ...]: + """Ensure argnum is within bounds. Also resolves negative argnums.""" + result = [] + for i in argnums: + if i >= num_args and allow_invalid: continue + if not -num_args <= i < num_args: + raise ValueError( + "Positional argument indices, e.g. for `static_argnums`, must have " + "value greater than or equal to -len(args) and less than len(args), " + f"but got value {i} for len(args) == {num_args}.") + result.append(i % num_args) # Resolve negative + return tuple(result) + + +def argnums_partial_except(f: lu.WrappedFun, static_argnums: tuple[int, ...], + args: tuple[Any, ...], *, allow_invalid: bool): + "Version of ``argnums_partial`` that checks hashability of static_argnums." + if not static_argnums: + return f, args + static_argnums = _ensure_inbounds(allow_invalid, len(args), static_argnums) + dyn_argnums = tuple(i for i in range(len(args)) if i not in static_argnums) + dyn_args = tuple(args[i] for i in dyn_argnums) + + fixed_args = [] + for i in static_argnums: + # TODO(shoyer): set allow_invalid=True permanently after static_argnames. + if allow_invalid and i >= len(args): + continue + static_arg = args[i] + if not is_hashable(static_arg): + raise ValueError( + "Non-hashable static arguments are not supported, as this can lead " + f"to unexpected cache-misses. Static argument (index {i}) of type " + f"{type(static_arg)} for function {f.__name__} is non-hashable.") + else: + fixed_args.append(_HashableWithStrictTypeEquality(static_arg)) + + return _argnums_partial(f, dyn_argnums, tuple(fixed_args)), dyn_args + +@lu.transformation +def _argnums_partial(dyn_argnums, fixed_args, *dyn_args, **kwargs): + sentinel = object() + args = [sentinel] * (len(fixed_args) + len(dyn_args)) + for i, arg in zip(dyn_argnums, dyn_args): + args[i] = arg + fixed_args_ = iter(fixed_args) + args = [next(fixed_args_).val if x is sentinel else x for x in args] + assert next(fixed_args_, sentinel) is sentinel + ans = yield args, kwargs + yield ans + + +def argnames_partial_except(f: lu.WrappedFun, static_argnames: tuple[str, ...], + kwargs: dict[str, Any]): + if not static_argnames: + return f, kwargs + dyn_kwargs = {k: v for k, v in kwargs.items() if k not in static_argnames} + + fixed_kwargs: dict[str, Any] = {} + for k, arg in kwargs.items(): + if k not in dyn_kwargs: + try: + hash(arg) + except TypeError: + raise ValueError( + "Non-hashable static arguments are not supported, as this can lead " + f"to unexpected cache-misses. Static argument (name {k}) of type " + f"{type(arg)} for function {f.__name__} is non-hashable.") + else: + fixed_kwargs[k] = Hashable(arg) + + return _argnames_partial(f, WrapKwArgs(fixed_kwargs)), dyn_kwargs + +@lu.transformation +def _argnames_partial(fixed_kwargs: WrapKwArgs, *args, **dyn_kwargs): + kwargs = dict({k: v.val for k, v in fixed_kwargs.val.items()}, **dyn_kwargs) + ans = yield args, kwargs + yield ans + + +@lru_cache(maxsize=4096) +def donation_vector(donate_argnums, donate_argnames, in_tree, + kws: bool = True) -> tuple[bool, ...]: + """Returns a tuple with a boolean value for each leaf in args and kwargs. + + What if a user specifies donate_argnums but calls the function with kwargs + or vice-versa? In that case, in `resolve_argnums` using the signature of the + function, the counterpart (donate_argnames or donate_argnums respectively) is + calculated so when this function is called both donate_argnums and + donate_argnames are available. This allows JAX to donate kwargs when only + donate_argnums is specified and vice-versa. + + When both donate_argnums and donate_argnames are specified, only the args and + kwargs specified are donated. + """ + res: list[bool] = [] + if kws: + args_tree, kwargs_tree = treedef_children(in_tree) + else: + args_tree, kwargs_tree = in_tree, None + for i, arg in enumerate(args_tree.children()): + donate = bool(i in donate_argnums) + res.extend((donate,) * arg.num_leaves) + if kwargs_tree is not None: + for key, val in safe_zip(kwargs_tree.node_data()[1], kwargs_tree.children()): # type: ignore + donate = key in donate_argnames + res.extend((donate,) * val.num_leaves) + return tuple(res) + +def rebase_donate_argnums(donate_argnums, static_argnums) -> tuple[int, ...]: + """Shifts donate to account for static. + + >>> rebase_donate_argnums((3, 4), (0, 1)) + (1, 2) + + Args: + donate_argnums: An iterable of ints. + static_argnums: An iterable of ints. + + Returns: + A tuple of unique, sorted integer values based on donate_argnums with each + element offset to account for static_argnums. + """ + if not (static_argnums or donate_argnums): + return tuple(sorted(donate_argnums)) + + static_argnums = sorted(set(static_argnums)) + donate_argnums = sorted(set(donate_argnums)) + i = j = o = 0 + out = [] + while j < len(donate_argnums): + if i < len(static_argnums) and static_argnums[i] == donate_argnums[j]: + raise ValueError(f"`static_argnums` {static_argnums} and " + f"`donate_argnums` {donate_argnums} cannot intersect.") + + if i < len(static_argnums) and static_argnums[i] < donate_argnums[j]: + o += 1 + i += 1 + else: + out.append(donate_argnums[j] - o) + j += 1 + return tuple(out) + + +def is_hashable(arg): + try: + hash(arg) + return True + except TypeError: + return False + + +def flatten_axes(name, treedef, axis_tree, *, kws=False, tupled_args=False): + # given an axis spec tree axis_tree (a pytree with integers and Nones at the + # leaves, i.e. the Nones are to be considered leaves) that is a tree prefix of + # the given treedef, build a complete axis spec tree with the same structure + # and return the flattened result + # TODO(mattjj,phawkins): improve this implementation + proxy = object() + dummy = tree_unflatten(treedef, [object()] * treedef.num_leaves) + axes = [] + add_leaves = lambda i, x: axes.extend([i] * len(tree_flatten(x)[0])) + try: + tree_map(add_leaves, _replace_nones(proxy, axis_tree), dummy) + except ValueError: + if kws: + # if keyword arguments are included in the tree, we make adapt the error + # message only to be about the positional arguments + treedef, _ = treedef_children(treedef) + axis_tree, _ = axis_tree + hint = "" + if tupled_args: + hint += (f" Note that {name} that are non-trivial pytrees should always be " + f"wrapped in a tuple representing the argument list.") + if len(treedef.children()) == 1: + try: + flatten_axes(name, treedef, (axis_tree,)) + except ValueError: + pass # That's not the issue. + else: + hint += (f" In particular, you're passing in a single argument which " + f"means that {name} might need to be wrapped in " + f"a singleton tuple.") + raise ValueError(f"{name} specification must be a tree prefix of the " + f"corresponding value, got specification {axis_tree} " + f"for value tree {treedef}.{hint}") from None + axes = [None if a is proxy else a for a in axes] + assert len(axes) == treedef.num_leaves + return axes + +def flat_out_axes( + f: lu.WrappedFun, out_spec: Any +) -> tuple[lu.WrappedFun, Callable]: + leaves, treedef = tree_flatten(out_spec) + f, out_axes = _flat_out_axes(f, tuple(leaves), treedef) + return f, HashableFunction(out_axes, closure=(tuple(leaves), treedef)) + +@lu.transformation_with_aux +def _flat_out_axes(leaves, treedef, *args, **kwargs): + ans = yield args, kwargs + spec = tree_unflatten(treedef, leaves) + try: + spec_flat = tuple(broadcast_prefix(spec, ans, is_leaf=lambda x: x is None)) + except ValueError: + e, *_ = prefix_errors(spec, ans) + # TODO(mattjj): currently hardcoded for pmap; generalize to vmap in followup + msg, = e('pmap out_axes').args + msg += ("\n\nThe full pytree is the output of the pmapped function. Ensure " + "that the `out_axes` argument to `pmap` is a pytree prefix of the " + "pmapped function's output.") + raise ValueError(msg) from None + yield ans, spec_flat + +def check_callable(fun): + # In Python 3.10+, the only thing stopping us from supporting staticmethods + # is that we can't take weak references to them, which the C++ JIT requires. + if isinstance(fun, staticmethod): + raise TypeError(f"staticmethod arguments are not supported, got {fun}") + if not callable(fun): + raise TypeError(f"Expected a callable value, got {fun}") + if inspect.isgeneratorfunction(fun): + raise TypeError(f"Expected a function, got a generator function: {fun}") + +_POSITIONAL_OR_KEYWORD = inspect.Parameter.POSITIONAL_OR_KEYWORD + +def infer_argnums_and_argnames( + sig: inspect.Signature, + argnums: int | Iterable[int] | None, + argnames: str | Iterable[str] | None, + ) -> tuple[tuple[int, ...], tuple[str, ...]]: + """Infer missing argnums and argnames for a function with inspect.""" + if argnums is None and argnames is None: + return (), () + + if argnums is not None and argnames is not None: + argnums = _ensure_index_tuple(argnums) + argnames = _ensure_str_tuple(argnames) + return argnums, argnames + + parameters = sig.parameters + if argnums is None: + assert argnames is not None + argnames = _ensure_str_tuple(argnames) + argnums = tuple( + i for i, (k, param) in enumerate(parameters.items()) + if param.kind == _POSITIONAL_OR_KEYWORD and k in argnames + ) + else: + argnums = _ensure_index_tuple(argnums) + argnames = tuple( + k for i, (k, param) in enumerate(parameters.items()) + if param.kind == _POSITIONAL_OR_KEYWORD and i in argnums + ) + + return argnums, argnames + + +def resolve_argnums( + fun: Callable, + signature: inspect.Signature | None, + donate_argnums: int | Sequence[int] | None, + donate_argnames: str | Iterable[str] | None, + static_argnums: int | Sequence[int] | None, + static_argnames: str | Iterable[str] | None, +) -> tuple[tuple[int, ...], tuple[str, ...], tuple[int, ...], tuple[str, ...]]: + """Validates and completes the argnum/argname specification for a jit. + + * fills in any missing pieces (e.g., names given numbers, or vice versa), + * validates the argument names/numbers against the function signature, + * validates that donated and static arguments don't intersect. + * rebases the donated arguments so they index into the dynamic arguments, + (after static arguments have been removed), in the order that parameters + are passed into the compiled function. + """ + if signature is None: + # Some built-in functions don't support signature. + # See: https://github.com/python/cpython/issues/73485 + # In this case no validation is done + static_argnums = () if static_argnums is None else _ensure_index_tuple( + static_argnums) + static_argnames = () if static_argnames is None else _ensure_str_tuple( + static_argnames) + donate_argnums = () if donate_argnums is None else _ensure_index_tuple( + donate_argnums) + if donate_argnames is not None: + raise ValueError(f"Getting the signature of function {fun} failed. " + "Pass donate_argnums instead of donate_argnames.") + assert donate_argnames is None + donate_argnames = () + else: + # Infer argnums and argnames according to docstring + # If nums is None and names is not None, then nums are inferred from the + # names and vice-versa. + static_argnums, static_argnames = infer_argnums_and_argnames( + signature, static_argnums, static_argnames) + donate_argnums, donate_argnames = infer_argnums_and_argnames( + signature, donate_argnums, donate_argnames) + + # Validation + _validate_argnums(signature, static_argnums, "static_argnums") + _validate_argnames(signature, static_argnames, "static_argnames") + _validate_argnums(signature, donate_argnums, "donate_argnums") + _validate_argnames(signature, donate_argnames, "donate_argnames") + + # Compensate for static argnums absorbing args + _assert_no_intersection(static_argnames, donate_argnames) + donate_argnums = rebase_donate_argnums(donate_argnums, static_argnums) + return donate_argnums, donate_argnames, static_argnums, static_argnames + + +def _assert_no_intersection(static_argnames, donate_argnames): + out = set(static_argnames).intersection(set(donate_argnames)) + if out: + raise ValueError( + "static_argnames and donate_argnames cannot intersect. Argument names " + f"{out} appear in both static_argnames and donate_argnames") + + +def _dtype(x): + try: + return dtypes.result_type(x) + except ValueError: + return dtypes.result_type(getattr(x, 'dtype')) + +def _shaped_abstractify_slow(x): + try: + return core.raise_to_shaped( + x if isinstance(x, core.AbstractValue) else core.get_aval(x)) + except TypeError: + pass + + weak_type = getattr(x, 'weak_type', False) + named_shape = getattr(x, 'named_shape', {}) + if hasattr(x, 'dtype'): + dtype = dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True) + else: + raise TypeError( + f"Cannot interpret value of type {type(x)} as an abstract array; it " + "does not have a dtype attribute") + return core.ShapedArray(np.shape(x), dtype, weak_type=weak_type, + named_shape=named_shape) + +# TODO(mattjj,yashkatariya): replace xla.abstractify with this, same behavior +def shaped_abstractify(x): + handler = _shaped_abstractify_handlers.get(type(x), None) + return handler(x) if handler is not None else _shaped_abstractify_slow(x) + +_shaped_abstractify_handlers: dict[Any, Callable[[Any], core.ShapedArray]] = {} + + +def _str_abstractify(x): + raise TypeError(f"Argument '{x}' of type {type(x)} is not a valid JAX type") +_shaped_abstractify_handlers[str] = _str_abstractify + +def _numpy_array_abstractify(x: np.ndarray) -> ShapedArray: + dtype = x.dtype + dtypes.check_valid_dtype(dtype) + return ShapedArray(x.shape, + dtypes.canonicalize_dtype(dtype, allow_extended_dtype=True)) +_shaped_abstractify_handlers[np.ndarray] = _numpy_array_abstractify + +def _np_scalar_abstractify(x: np.generic) -> ShapedArray: + dtype = np.dtype(x) + dtypes.check_valid_dtype(dtype) + return ShapedArray(np.shape(x), + dtypes.canonicalize_dtype(dtype, allow_extended_dtype=True)) +_shaped_abstractify_handlers.update((t, _np_scalar_abstractify) + for t in numpy_scalar_types) + +def _python_scalar_abstractify(x: int | float | complex | bool) -> ShapedArray: + typ = type(x) + dtype = dtypes._scalar_type_to_dtype(typ, x) + return ShapedArray((), dtype, weak_type=typ in dtypes._weak_types) +_shaped_abstractify_handlers.update((t, _python_scalar_abstractify) + for t in dtypes.python_scalar_dtypes) + +# This decorator exists to make it easier to monkey-patch APIs in JAX. +# By default it does nothing, but it can be monkey-patched to do other things. +def api_hook(fun, tag: str): + return fun + + +def debug_info( + traced_for: str, src: str | None, fun_signature: inspect.Signature | None, + args: tuple[Any, ...], kwargs: dict[str, Any], static_argnums: tuple[int, ...], + static_argnames: tuple[str, ...] +) -> TracingDebugInfo | None: + """Try to build trace-time debug info for fun when applied to args/kwargs.""" + arg_names = _arg_names(fun_signature, args, kwargs, static_argnums, + static_argnames) + if src is None or arg_names is None: + return None + return TracingDebugInfo(traced_for, src, arg_names, None) + +def fun_signature(fun: Callable) -> inspect.Signature | None: + try: + return inspect.signature(fun) + except (ValueError, TypeError): + return None + +# TODO(mattjj): make this function internal to this module +def fun_sourceinfo(fun: Callable) -> str | None: + while isinstance(fun, partial): + fun = fun.func + fun = inspect.unwrap(fun) + try: + filename = fun.__code__.co_filename + lineno = fun.__code__.co_firstlineno + return f"{fun.__name__} at {filename}:{lineno}" + except AttributeError: + return None + +def _arg_names(fn_signature, args, kwargs, static_argnums, static_argnames, + ) -> tuple[str, ...] | None: + if fn_signature is None: return None + static = object() + static_argnums_ = _ensure_inbounds(True, len(args), static_argnums) + static_argnames_ = set(static_argnames) + args_ = [static if i in static_argnums_ else x for i, x in enumerate(args)] + kwargs = {k:static if k in static_argnames_ else x for k, x in kwargs.items()} + try: + ba = fn_signature.bind(*args_, **kwargs) + except (ValueError, TypeError): + return None + return tuple(f'{name}{keystr(path)}' for name, x in ba.arguments.items() + for path, l in generate_key_paths(x) if l is not static) + +@lu.transformation_with_aux +def result_paths(*args, **kwargs): + "linear_util transform to get output pytree paths of pre-flattened function." + ans = yield args, kwargs + yield ans, [keystr(path) for path, _ in generate_key_paths(ans)] + +def jaxpr_debug_info(jaxpr: core.Jaxpr, trace_debug: TracingDebugInfo | None, + result_paths: tuple[str, ...] | None = None, + ) -> core.Jaxpr: + """Add debug info to jaxpr, given trace-time debug info and result paths.""" + if trace_debug is None: + return jaxpr + assert (result_paths is not None) ^ (trace_debug.result_paths is not None) + if result_paths is None: + result_paths = trace_debug.result_paths() # type: ignore + debug_info = core.JaxprDebugInfo( + trace_debug.traced_for, trace_debug.func_src_info, + trace_debug.arg_names, tuple(result_paths)) + return jaxpr.replace(debug_info=debug_info) + +def debug_info_final(f: lu.WrappedFun, dbg: TracingDebugInfo | None, + res_paths: Callable[[], tuple[str, ...]]) -> lu.WrappedFun: + "Attach trace-time debug info and result paths lazy thunk to an lu.WrappedFun" + if dbg is None: return f + assert dbg.result_paths is None + res_paths_ = HashableFunction(res_paths, closure=()) + return lu.add_debug_info(f, dbg._replace(result_paths=res_paths_)) + + +def hoist_obj_attrs(f, flat_args): + idxs, objs, flat_args_ = [], [], [] + for i, x in enumerate(flat_args): + if type(x) in _class_with_attrs: + objs.append(_HashableByObjectId(x)) + else: + idxs.append(i) + flat_args_.append(x) + return _argnums_partial(f, tuple(idxs), tuple(objs)), flat_args_ + +class _HashableByObjectId: + __slots__ = ['val'] + def __init__(self, val): + self.val = val + def __hash__(self): + return id(self.val) + def __eq__(self, other): + return self.val is other.val + +def register_class_with_attrs(t: Type) -> None: + _class_with_attrs.add(t) +_class_with_attrs: set[Type] = set() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/array.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/array.py new file mode 100644 index 0000000000000000000000000000000000000000..d2064a16eddf29d46ae14ac56dd14c2b0eb31bb5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/array.py @@ -0,0 +1,1155 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Sequence +import enum +import functools +import math +import operator as op +from typing import Any, Callable, TYPE_CHECKING, cast + +from jax._src import abstract_arrays +from jax._src import api +from jax._src import api_util +from jax._src import basearray +from jax._src import config +from jax._src import core +from jax._src import dispatch +from jax._src import dtypes +from jax._src import errors +from jax._src import profiler +from jax._src import tree_util +from jax._src import xla_bridge +from jax._src.interpreters import mlir +from jax._src.interpreters import pxla +from jax._src.interpreters import xla +from jax._src.layout import AutoLayout, DeviceLocalLayout, Layout +from jax._src.lib import xla_client as xc +from jax._src.lib import xla_extension as xe +from jax._src.sharding import Sharding +from jax._src.sharding_impls import ( + PmapSharding, SingleDeviceSharding, + device_replica_id_map, hashed_index, num_addressable_indices, local_to_global_shape) # pyformat: disable +from jax._src.typing import ArrayLike, DLDeviceType +from jax._src.util import safe_zip, unzip3, use_cpp_class, use_cpp_method +import numpy as np + + +Shape = tuple[int, ...] +Device = xc.Device +Index = tuple[slice, ...] +PRNGKeyArray = Any # TODO(jakevdp): fix cycles and import this. + +def _get_device(a: ArrayImpl) -> Device: + devices = a.sharding._internal_device_list # pytype: disable=attribute-error + assert len(devices) == 1 + return devices[0] + + +class Shard: + """A single data shard of an Array. + + Attributes: + device : Which device this shard resides on. + index : The index into the global array of this shard. + replica_id : Integer id indicating which replica of the global array this + shard is part of. Always 0 for fully sharded data + (i.e. when there’s only 1 replica). + data : The data of this shard. None if ``device`` is non-local. + """ + + def __init__(self, device: Device, sharding: Sharding, global_shape: Shape, + data: None | ArrayImpl | PRNGKeyArray = None): + self._device = device + self._sharding = sharding + self._global_shape = global_shape + self._data = data + + def __repr__(self): + try: + return (f'Shard(device={self.device!r}, index={self.index}, ' + f'replica_id={self.replica_id}, data={self.data})') + except ValueError: + return f'Shard(device={self.device!r}, data={self.data})' + + @functools.cached_property + def index(self) -> Index: + try: + device_indices_map_fn = self._sharding.devices_indices_map + except AttributeError: + raise ValueError('Cannot calculate indices from sharding: ' + f'{self._sharding}. Please create a device to index ' + 'mapping for your sharding.') from None + index = device_indices_map_fn(self._global_shape)[self.device] + assert index is not None + return index + + @functools.cached_property + def replica_id(self) -> int: + return device_replica_id_map(self._sharding, self._global_shape)[self.device] + + @property + def device(self): + return self._device + + @property + def data(self): + return self._data + + +def _reconstruct_array(fun, args, arr_state, aval_state): + """Method to reconstruct a device array from a serialized state.""" + np_value = fun(*args) + np_value.__setstate__(arr_state) + jnp_value = api.device_put(np_value) + jnp_value.aval = jnp_value.aval.update(**aval_state) + return jnp_value + + +@functools.lru_cache(maxsize=4096) +def _cached_index_calc(s, shape): + map_ = s.addressable_devices_indices_map(shape) + seen_h_indices = set() + l = [] + for array_index, index in enumerate(map_.values()): + h_index = hashed_index(index) + if h_index not in seen_h_indices: + seen_h_indices.add(h_index) + l.append((array_index, index)) + return l + + +@functools.lru_cache(maxsize=4096) +def _process_has_full_value_in_mcjax(s, shape): + # Return False for single host as a fast path. + if xla_bridge.process_count() == 1: + return False + + num_unique_indices = len( + {hashed_index(v) for v in s.devices_indices_map(shape).values()}) + num_addressable_unique_indices = len( + {hashed_index(v) for v in s.addressable_devices_indices_map(shape).values()}) + return num_unique_indices == num_addressable_unique_indices + + +class ArrayImpl(basearray.Array): + # TODO(yashkatariya): Add __slots__ here. + + aval: core.ShapedArray + _sharding: Sharding + _arrays: list[ArrayImpl] + _committed: bool + _skip_checks: bool + _npy_value: np.ndarray | None + + @use_cpp_method() + def __init__(self, aval: core.ShapedArray, sharding: Sharding, + arrays: Sequence[ArrayImpl], + committed: bool, _skip_checks: bool = False): + # NOTE: the actual implementation of the constructor is moved to C++. + + self.aval = aval + self._sharding = sharding + self._arrays = [a._arrays[0] for a in arrays] + self._committed = committed + self._npy_value = None + + # Don't rearrange if skip_checks is enabled because this assumes that the + # input buffers are already arranged properly. This usually happens when + # Array's are created as output of a JAX transformation + # (like pjit, xmap, etc). + if not _skip_checks or config.enable_checks.value: + self._check_and_rearrange() + + def _check_and_rearrange(self): + for db in self._arrays: + if db.dtype != self.dtype: + raise ValueError( + "Input buffers to `Array` must have matching dtypes. " + f"Got {db.dtype}, expected {self.dtype} for buffer: {db}") + + device_id_to_buffer = {_get_device(db).id: db for db in self._arrays} + + addressable_dev = self.sharding.addressable_devices + if len(self._arrays) != len(addressable_dev): + raise ValueError( + f"Expected {len(addressable_dev)} per-device arrays " + "(this is how many devices are addressable by the sharding), but " + f"got {len(self._arrays)}") + + array_device_ids = set(device_id_to_buffer.keys()) + addressable_device_ids = {d.id for d in addressable_dev} + # Calculate a symmetric difference because the device ids between sharding + # and _arrays should match. + diff = array_device_ids ^ addressable_device_ids + if diff: + dev_in_sharding_not_in_arrays = addressable_device_ids - array_device_ids + dev_in_arrays_not_in_sharding = array_device_ids - addressable_device_ids + err_msg = ( + "Addressable devices and per-device arrays devices do not match.") + if dev_in_sharding_not_in_arrays: + err_msg += (f" Sharding contains devices {dev_in_sharding_not_in_arrays} " + "that are not present in per-device arrays.") + if dev_in_arrays_not_in_sharding: + err_msg += (f" Per-device arrays contain devices {dev_in_arrays_not_in_sharding} " + "that are not present in the sharding.") + raise ValueError(err_msg) + + ss = self.sharding.shard_shape(self.shape) + for db in self._arrays: + if db.shape != ss: + raise ValueError( + f"Expected shard shape {ss} doesn't match the single device array " + f"shape {db.shape}. Shape of Array is " + f"{self.aval.str_short()} with sharding {self.sharding}") + + # Rearrange arrays based on the device assignment. + addressable_da = self.sharding._addressable_device_assignment + self._arrays = [device_id_to_buffer[device.id] for device in addressable_da] + + @property + def shape(self) -> Shape: + return self.aval.shape + + @property + def dtype(self): + return self.aval.dtype + + @property + def ndim(self): + return len(self.shape) + + @property + def size(self): + return math.prod(self.shape) + + @property + def sharding(self): + return self._sharding + + @property + def weak_type(self): + return self.aval.weak_type + + def __str__(self): + return str(self._value) + + def __len__(self): + try: + return self.shape[0] + except IndexError as err: + raise TypeError("len() of unsized object") from err # same as numpy error + + def __bool__(self): + core.check_bool_conversion(self) + return bool(self._value) + + def __float__(self): + core.check_scalar_conversion(self) + return self._value.__float__() + + def __int__(self): + core.check_scalar_conversion(self) + return self._value.__int__() + + def __complex__(self): + core.check_scalar_conversion(self) + return self._value.__complex__() + + def __hex__(self): + core.check_integer_conversion(self) + return hex(self._value) + + def __oct__(self): + core.check_integer_conversion(self) + return oct(self._value) + + def __index__(self): + core.check_integer_conversion(self) + return op.index(self._value) + + def tobytes(self, order="C"): + return self._value.tobytes(order) + + def tolist(self): + return self._value.tolist() + + def __format__(self, format_spec): + # Simulates behavior of https://github.com/numpy/numpy/pull/9883 + if self.ndim == 0: + return format(self._value[()], format_spec) + else: + return format(self._value, format_spec) + + def __getitem__(self, idx): + from jax._src.lax import lax + from jax._src.numpy import lax_numpy + self._check_if_deleted() + + if isinstance(self.sharding, PmapSharding): + if config.pmap_no_rank_reduction.value: + cidx = idx if isinstance(idx, tuple) else (idx,) + + padded_cidx = tuple( + slice(i, i + 1, None) if isinstance(i, int) else i for i in cidx + ) + (slice(None),) * (len(self.shape) - len(cidx)) + else: + if not isinstance(idx, tuple): + padded_cidx = (idx,) + (slice(None),) * (len(self.shape) - 1) + else: + padded_cidx = idx + (slice(None),) * (len(self.shape) - len(idx)) + + indices = tuple(self.sharding.devices_indices_map(self.shape).values()) + try: + arr_idx = indices.index(padded_cidx) + except ValueError: + arr_idx = None + if arr_idx is not None: + a = self._arrays[arr_idx] + out = ArrayImpl( + a.aval, SingleDeviceSharding(_get_device(a)), [a], committed=False, + _skip_checks=True) + + if config.pmap_no_rank_reduction.value: + # If cidx was the index of a single shard, then it corresponds to one + # shard of the chunked dimension. + dims = tuple(i for i, x in enumerate(cidx) if isinstance(x, int)) + return lax.squeeze(out, dimensions=dims) + else: + return out + + return lax_numpy._rewriting_take(self, idx) + + def __iter__(self): + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") # same as numpy error + else: + assert self.is_fully_replicated or self.is_fully_addressable + if dispatch.is_single_device_sharding(self.sharding) or self.is_fully_replicated: + return (sl for chunk in self._chunk_iter(100) for sl in chunk._unstack()) + elif isinstance(self.sharding, PmapSharding): + return (self[i] for i in range(self.shape[0])) + else: + # TODO(yashkatariya): Don't bounce to host and use `_chunk_iter` path + # here after uneven partitioning support is added. + return (api.device_put(self._value[i]) for i in range(self.shape[0])) + + @property + def is_fully_replicated(self) -> bool: + return self.sharding.is_fully_replicated + + def __repr__(self): + prefix = 'Array(' + if self.aval is not None and self.aval.weak_type: + dtype_str = f'dtype={self.dtype.name}, weak_type=True)' + else: + dtype_str = f'dtype={self.dtype.name})' + + if self.is_fully_addressable or self.is_fully_replicated: + line_width = np.get_printoptions()["linewidth"] + if self.size == 0: + s = f"[], shape={self.shape}" + else: + s = np.array2string(self._value, prefix=prefix, suffix=',', + separator=', ', max_line_width=line_width) + last_line_len = len(s) - s.rfind('\n') + 1 + sep = ' ' + if last_line_len + len(dtype_str) + 1 > line_width: + sep = ' ' * len(prefix) + return f"{prefix}{s},{sep}{dtype_str}" + else: + return f"{prefix}{self.shape}, {dtype_str}" + + @property + def is_fully_addressable(self) -> bool: + """Is this Array fully addressable? + + A jax.Array is fully addressable if the current process can address all of + the devices named in the :class:`Sharding`. ``is_fully_addressable`` is + equivalent to "is_local" in multi-process JAX. + + Note that fully replicated is not equal to fully addressable i.e. + a jax.Array which is fully replicated can span across multiple hosts and is + not fully addressable. + """ + return self.sharding.is_fully_addressable + + def __array__(self, dtype=None, context=None, copy=None): + # copy argument is supported by np.asarray starting in numpy 2.0 + kwds = {} if copy is None else {'copy': copy} + return np.asarray(self._value, dtype=dtype, **kwds) + + def __dlpack__(self, *, stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[DLDeviceType, int] | None = None, + copy: bool | None = None): + from jax._src.dlpack import to_dlpack # pylint: disable=g-import-not-at-top + + device_set = self.sharding.device_set + if len(device_set) > 1: + raise BufferError( + "to_dlpack can only pack a dlpack tensor from an array on a singular " + f"device, but an array with a Sharding over {len(device_set)} devices " + "was provided." + ) + device, = device_set + return to_dlpack(self, stream=stream, + max_version=max_version, + src_device=device, + dl_device=dl_device, + copy=copy) + + def __dlpack_device__(self) -> tuple[enum.Enum, int]: + if len(self._arrays) != 1: + raise BufferError("__dlpack__ only supported for unsharded arrays.") + + from jax._src.dlpack import DLDeviceType # pylint: disable=g-import-not-at-top + + if self.platform() == "cpu": + return DLDeviceType.kDLCPU, 0 + + elif self.platform() == "gpu": + platform_version = _get_device(self).client.platform_version + if "cuda" in platform_version: + dl_device_type = DLDeviceType.kDLCUDA + elif "rocm" in platform_version: + dl_device_type = DLDeviceType.kDLROCM + else: + raise BufferError("Unknown GPU platform for __dlpack__: " + f"{platform_version}") + + local_hardware_id = _get_device(self).local_hardware_id + if local_hardware_id is None: + raise BufferError("Couldn't get local_hardware_id for __dlpack__") + + return dl_device_type, local_hardware_id + + else: + raise BufferError( + "__dlpack__ device only supported for CPU and GPU, got platform: " + f"{self.platform()}" + ) + + def __reduce__(self): + fun, args, arr_state = self._value.__reduce__() + aval_state = {'weak_type': self.aval.weak_type, + 'named_shape': self.aval.named_shape} + return (_reconstruct_array, (fun, args, arr_state, aval_state)) + + @use_cpp_method() + def unsafe_buffer_pointer(self): + if len(self._arrays) != 1: + raise ValueError("unsafe_buffer_pointer() is supported only for unsharded" + " arrays.") + return self._arrays[0].unsafe_buffer_pointer() + + @property + @use_cpp_method() + def __cuda_array_interface__(self): + if len(self._arrays) != 1: + raise ValueError("__cuda_array_interface__() is supported only for " + "unsharded arrays.") + return self._arrays[0].__cuda_array_interface__ # pytype: disable=attribute-error # bind-properties + + @use_cpp_method() + def on_device_size_in_bytes(self): + """Returns the total global on-device size of the array in bytes.""" + arr = self._arrays[0] + per_shard_size = arr.on_device_size_in_bytes() + return per_shard_size * len(self.sharding.device_set) + + def devices(self) -> set[Device]: + self._check_if_deleted() + return self.sharding.device_set + + @property + def device_buffer(self): + raise AttributeError( + "arr.device_buffer has been deprecated. Use arr.addressable_data(0)") + + @property + def device_buffers(self): + raise AttributeError( + "arr.device_buffers has been deprecated. Use [x.data for x in arr.addressable_shards]") + + def addressable_data(self, index: int) -> ArrayImpl: + self._check_if_deleted() + if self.is_fully_replicated: + return self._fully_replicated_shard() + return self._arrays[index] + + @functools.cached_property + def addressable_shards(self) -> Sequence[Shard]: + self._check_if_deleted() + out = [] + for a in self._arrays: + out.append(Shard(_get_device(a), self.sharding, self.shape, a)) + return out + + @property + def layout(self): + # TODO(yashkatariya): Remove the deleted check from here. + if self.is_deleted(): + return Layout(None, self.sharding) + try: + return Layout(DeviceLocalLayout(self._pjrt_layout), self.sharding) + except xe.XlaRuntimeError as e: + msg, *_ = e.args + if type(msg) is str and msg.startswith("UNIMPLEMENTED"): + return Layout(None, self.sharding) + else: + raise + + @property + def global_shards(self) -> Sequence[Shard]: + """Returns list of all `Shard`s of the Array across all devices. + + The result includes shards that are not addressable by the current process. + If a `Shard` is not addressable, then its `data` will be `None`. + """ + self._check_if_deleted() + if self.is_fully_addressable: # pylint: disable=using-constant-test + return self.addressable_shards + + out = [] + device_id_to_buffer = {_get_device(a).id: a for a in self._arrays} + for global_d in self.sharding.device_set: + if device_id_to_buffer.get(global_d.id, None) is not None: + array = device_id_to_buffer[global_d.id] + else: + array = None + out.append(Shard(global_d, self.sharding, self.shape, array)) + return out + + @use_cpp_method() + def delete(self): + if self._arrays is None: + return + for buf in self._arrays: + buf.delete() + self._arrays = None + self._npy_value = None + + @use_cpp_method() + def is_deleted(self): + if self._arrays is None: + return True + # This path is taken when a view of `Array` is created and the original + # Array is deleted. In that case, the buffers the view represents also get + # deleted. + return any(buf.is_deleted() for buf in self._arrays) + + def _check_if_deleted(self): + if self.is_deleted(): + raise RuntimeError( + f"Array has been deleted with shape={self.aval.str_short()}.") + + @use_cpp_method() + def block_until_ready(self): + self._check_if_deleted() + for db in self._arrays: + db.block_until_ready() + return self + + @use_cpp_method() + def _single_device_array_to_np_array(self): + return np.asarray(self._arrays[0]) + + @use_cpp_method() + def _copy_single_device_array_to_host_async(self): + self._arrays[0].copy_to_host_async() + + @profiler.annotate_function + def copy_to_host_async(self): + self._check_if_deleted() + if self._npy_value is None: + if self.is_fully_replicated: + self._copy_single_device_array_to_host_async() + return + for i, _ in _cached_index_calc(self.sharding, self.shape): + self._arrays[i]._copy_single_device_array_to_host_async() + + @property + @functools.partial(profiler.annotate_function, name="np.asarray(jax.Array)") + def _value(self) -> np.ndarray: + self._check_if_deleted() + + if self._npy_value is None: + if self.is_fully_replicated: + self._npy_value = self._single_device_array_to_np_array() + self._npy_value.flags.writeable = False + return cast(np.ndarray, self._npy_value) + + # TODO(yashkatariya): Merge `_process_has_full_value_in_mcjax` with + # is_fully_addressable. + if (not self.is_fully_addressable and + not _process_has_full_value_in_mcjax(self.sharding, self.shape)): + raise RuntimeError( + "Fetching value for `jax.Array` that spans non-addressable" + " (non process local) devices is not possible. You can use" + " `jax.experimental.multihost_utils.process_allgather` to print the" + " global array or use `.addressable_shards` method of jax.Array to" + " inspect the addressable (process local) shards." + ) + + for i, _ in _cached_index_calc(self.sharding, self.shape): + self._arrays[i]._copy_single_device_array_to_host_async() + + npy_value = np.empty(self.shape, self.dtype) + for i, ind in _cached_index_calc(self.sharding, self.shape): + npy_value[ind] = self._arrays[i]._single_device_array_to_np_array() + self._npy_value = npy_value + self._npy_value.flags.writeable = False + # https://docs.python.org/3/library/typing.html#typing.cast + return cast(np.ndarray, self._npy_value) + + +# TODO(b/273265390): ideally we would write this as a decorator on the ArrayImpl +# class, however this triggers a pytype bug. Workaround: apply the decorator +# after the fact. +if not TYPE_CHECKING: + ArrayImpl = use_cpp_class(xc.ArrayImpl)(ArrayImpl) + + +# explicitly set to be unhashable. +setattr(ArrayImpl, "__hash__", None) +setattr(ArrayImpl, "__array_priority__", 100) + +# TODO(yashkatariya): Remove None from callback input type. + +def make_array_from_callback( + shape: Shape, sharding: Sharding | Layout, + data_callback: Callable[[Index | None], ArrayLike]) -> ArrayImpl: + # pyformat: disable + """Returns a ``jax.Array`` via data fetched from ``data_callback``. + + ``data_callback`` is used to fetch the data for each addressable shard of the + returned ``jax.Array``. This function must return concrete arrays, meaning that + ``make_array_from_callback`` has limited compatibility with JAX transformations + like :func:`jit` or :func:`vmap`. + + Args: + shape : Shape of the ``jax.Array``. + sharding: A ``Sharding`` instance which describes how the ``jax.Array`` is + laid out across devices. + data_callback : Callback that takes indices into the global array value as + input and returns the corresponding data of the global array value. + The data can be returned as any array-like object, e.g. a ``numpy.ndarray``. + + Returns: + A ``jax.Array`` via data fetched from ``data_callback``. + + Example: + + >>> import math + >>> from jax.sharding import Mesh + >>> from jax.sharding import PartitionSpec as P + >>> import numpy as np + ... + >>> input_shape = (8, 8) + >>> global_input_data = np.arange(math.prod(input_shape)).reshape(input_shape) + >>> global_mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y')) + >>> inp_sharding = jax.sharding.NamedSharding(global_mesh, P('x', 'y')) + ... + >>> def cb(index): + ... return global_input_data[index] + ... + >>> arr = jax.make_array_from_callback(input_shape, inp_sharding, cb) + >>> arr.addressable_data(0).shape + (4, 2) + """ + # pyformat: enable + dll = sharding.device_local_layout if isinstance(sharding, Layout) else None + if isinstance(dll, AutoLayout): + raise TypeError( + "`DeviceLocalLayout.AUTO` cannot be used in place of a device-local" + f" layout when calling `jax.make_array_from_callback`. Got {sharding}") + sharding = sharding.sharding if isinstance(sharding, Layout) else sharding # type: ignore + if not isinstance(sharding, Sharding): + raise TypeError( + f"sharding should be an instance of `jax.sharding`. Got {sharding} of" + f" type {type(sharding)}") + + if sharding.is_fully_replicated: + devices = list(sharding._internal_device_list.addressable_device_list) # type: ignore + per_device_values = [data_callback((slice(None),) * len(shape))] * len(devices) + else: + device_to_index_map = sharding.addressable_devices_indices_map(shape) + devices = list(device_to_index_map.keys()) + per_device_values = [data_callback(device_to_index_map[device]) + for device in devices] + + if isinstance(per_device_values[0], core.Tracer): + raise errors.UnexpectedTracerError( + "jax.make_array_from_callback cannot be called within a traced context.") + + first_value = xla.canonicalize_dtype(per_device_values[0]) + aval = core.ShapedArray(shape, first_value.dtype, weak_type=False) + + # first value can be numpy array, python scalar, etc. + if (sharding.is_fully_replicated and not isinstance(first_value, ArrayImpl) + and not dtypes.issubdtype(aval.dtype, dtypes.extended) and dll is None): + # Do this check outside because `batched_device_put` won't do these checks + # like ArrayImpl. + if shape != first_value.shape: + raise ValueError( + f"Expected shard shape {shape} doesn't match the single device " + f"array shape {first_value.shape}. Shape of Array is " + f"{aval.str_short()} with sharding {sharding}") + return pxla.batched_device_put( + aval, sharding, per_device_values, devices, committed=True) + + if (isinstance(first_value, ArrayImpl) + and first_value._committed + and sharding.is_fully_replicated + and first_value.is_fully_replicated + and first_value.sharding._device_assignment == tuple(devices) + and (first_value.layout.device_local_layout == + pxla._maybe_get_default_layout(Layout(dll, sharding), None, sharding, aval))): + return first_value + + if dll is not None: + devices = [Layout(dll, SingleDeviceSharding(d)) for d in devices] + arrays = api.device_put(per_device_values, devices) + if dtypes.issubdtype(aval.dtype, dtypes.extended): + return aval.dtype._rules.make_sharded_array(aval, sharding, arrays, + committed=True) + return ArrayImpl(aval, sharding, arrays, committed=True) + + +def make_array_from_process_local_data( + sharding: Sharding, + local_data: np.ndarray, + global_shape: Shape | None = None, +) -> ArrayImpl: + # pyformat: disable + """Creates distributed tensor using the data available in process. + + This function is a common special case of `make_array_from_callback`. It + assumes that the data is available in the process and takes care of the + index wrangling. + + The most common case is when the sharding is sharded across the batch + dimension and each host just loads its corresponding sub-batch. This function + supports more general cases as well, such as mixed multi-host and multi-axis + replication and sharding but you would need to compute the size and the + contents of process-local data correctly to satisfy the sharding constraints. + + In particular, if any two hosts are replicas, host_local_data should be + identical as well. + + The global_shape is optional. If not provided it will be be inferred from + the local_data and sharding, under the assumption that + each host represents only their own data for uniform sharding. If sharding + is non-uniform, (see note below) an exception will be raised. + + Setting global_shape explicitly allows for finer grain control and works with + non-uniform shardings. Each dimension of global_shape must either match + host_local_data, or match the inferred global shape of the sharding (in which + case it is equivalent to setting it to None, but is more explicit). + + For example if dimension `i` is fully sharded then this size would be + `per_device_shape[i] * jax.local_device_count()`. Each device will be mapped + into local slice of `local_data` array. For example, if given process + addresses slices (8, 12) and (24, 28), then these slices will be mapped + into (0, 4) and (4, 8) of the `local_data`. + + For each dimension where global_shapes matches local_shape, each device + will lookup the slice in the local_data. For example if + global_shape == local_data.shape, the local data is assumed to be the + actual target array that will be sharded into device. + + If global_shape is the same as local_data.shape, then the data must + be the same across all hosts. + + Examples: + >>> from jax.sharding import PartitionSpec as P + >>> mesh_rows = 2 + >>> mesh_cols = jax.device_count() // 2 + ... + >>> mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(mesh_rows, mesh_cols), ('x', 'y')) + + >>> sharding = jax.sharding.NamedSharding(mesh, P(('x', 'y'),)) + >>> rows_per_device = 2 + >>> feature_length = 32 + >>> per_device_shape = (rows_per_device, feature_length) + >>> per_host_shape = (rows_per_device * len(mesh.local_devices), feature_length) + >>> per_host_generator = lambda : np.arange(np.prod(per_host_shape)).reshape(per_host_shape) + >>> per_host_data = per_host_generator() # replace with your own per-host data pipeline that outputs numpy arrays + >>> global_shape = (rows_per_device * len(sharding.device_set), ) + per_device_shape[1:] + >>> output_global_array = jax.make_array_from_process_local_data(sharding, per_host_data, global_shape) + ... + >>> assert output_global_array.addressable_data(0).shape == per_device_shape + >>> assert output_global_array.shape == global_shape + + NB: While most shardings are uniform, It is possible to design am exotic + sharding mesh where each process's devices will be arranged in a non-grid + like pattern in some dimensions, or for indices to overlap non-trivially. + Such sharding is called "non-uniform" in those dimensions. In that case, + the global shape along those directions must match local shape as there is + no meaningful way to represent all needed + per-process data in non-overlapping fashion. For example for global_shape 4x4 + if sharding looks like this: + + 0123 + 2103 + 4675 + 4567 + + with 4 processes, containing devices (0,1), (2, 3), (4, 5), (6, 7) respectively. + Then the data for each host look like + + xx.. ..xx .... .... + .xx. x..x .... .... + .... .... x..x .xx. + .... .... xx.. ..xx + + the sharding is uniform on rows (each host requires either rows 1-2, or rows 3-4) + and non-uniform on columns (hosts require overlapping but not matching + set of columns). Thus local data must have the shape 2x4 or 4x4 + for all hosts, even though each host can potentially fit into 2x2 shape. + In this case user must provide global_shape explicitly and for + local_shape=(2, 4), potentially valid global shapes are (2, 4) and (4, 4). + + On the other hand for sharding: + + 0213 x.x. .x.x. .... .... + 0213 x.x. .x.x. .... .... + 4657 .... .... .x.x x.x. + 4657 .... .... .x.x x.x. + + for local_shape=(2, 2) this function can accept a choice of 2x2, 2x4, 4x2 + and 4x4 global shapes. Setting global_shape to None, is equivalent to + setting it to (4, 4) in this case. + + Args: + sharding: sharding of the global tensor. + host_local_data: data on the host to be placed on local devices. Each + dimension should either match global_shape, or match + num_addressable_indices(dim). + global_shape: the target shape of the global tensor. If None, + will infer from host_local_data and sharding. + + Returns: + Tensor that will have sharding=sharding and of shape global_shape. + """ + # pyformat: enable + # TODO(sandler): consider supporting partially specified global_shape or + # making local_to_global_shape available in the api. + local_shape = local_data.shape + if global_shape is None: + global_shape = local_to_global_shape(sharding, local_shape) # type: ignore[assignment] + assert global_shape is not None + if None in global_shape: + raise ValueError( + "Unable to compute global_shape due to non-uniform sharding." + f" Specify global shape directly. Partially computed {global_shape=}." + ) + elif None in global_shape: + raise ValueError(f"{global_shape=} has Nones. This is not supported.") + shard_shape = sharding.shard_shape(global_shape) + full_dim = [] + for i, (data_dim, global_dim) in enumerate( + zip(local_data.shape, global_shape) + ): + full_dim.append(data_dim == global_dim) + if data_dim != global_dim: + process_slice = num_addressable_indices(sharding, i, global_shape) + if process_slice != data_dim: + raise ValueError( + "Invalid host data, each dimension should match either global or " + f"process shape. In dimension {i=}, the process data has {data_dim}" + f"elements. Process addresses {process_slice} elements and " + f"{global_shape=}." + ) + addressable_shards = sharding.addressable_devices_indices_map(global_shape) + slices_for_each_dim: list[list[int]] = [[] for _ in global_shape] + for shard_index in addressable_shards.values(): + assert shard_index is not None + for i, slc in enumerate(shard_index): + slices_for_each_dim[i].append(slc.start or 0) + for i in range(len(global_shape)): + slices_for_each_dim[i] = sorted(set(slices_for_each_dim[i])) + + def local_slice(i, slc): + # Looks up the index of this slice in the list of slices for this dimension. + # This will determine the slice in host_local_data + start = slices_for_each_dim[i].index(slc.start or 0) * shard_shape[i] + end = start + shard_shape[i] + return slice(start, end) + + def cb(index: Index | None) -> ArrayLike: + assert index is not None + data_slice = [ + slc if full_dim[i] else local_slice(i, slc) + for i, slc in enumerate(index) + ] + return local_data[tuple(data_slice)] + + return make_array_from_callback(global_shape, sharding, cb) + + +def make_array_from_single_device_arrays( + shape: Shape, sharding: Sharding, arrays: Sequence[basearray.Array] +) -> ArrayImpl: + r"""Returns a ``jax.Array`` from a sequence of ``jax.Array``\s each on a single device. + Every device in input ``sharding``\'s mesh must have an array in ``arrays``\s. + + Args: + shape : Shape of the output ``jax.Array``. This conveys information already included with + ``sharding`` and ``arrays`` and serves as a double check. + sharding: Sharding: A global Sharding instance which describes how the output jax.Array is laid out across devices. + arrays: Sequence of ``jax.Array``\s that are each single device addressable. ``len(arrays)`` + must equal ``len(sharding.addressable_devices)`` and the shape of each array must be the same. For multiprocess code, + each process will call with a different ``arrays`` argument that corresponds to that processes' data. + These arrays are commonly created via ``jax.device_put``. + + Returns: + A global ``jax.Array``, sharded as ``sharding``, with shape equal to ``shape``, and with per-device + contents matching ``arrays``. + + Examples: + + In this single-process example, we use ``make_array_from_single_device_arrays`` to create an + a global array. + + >>> import math + >>> from jax.sharding import Mesh + >>> from jax.sharding import PartitionSpec as P + >>> import numpy as np + ... + >>> mesh_rows = 2 + >>> mesh_cols = jax.device_count() // 2 + ... + >>> global_shape = (8, 8) + >>> mesh = Mesh(np.array(jax.devices()).reshape(mesh_rows, mesh_cols), ('x', 'y')) + >>> sharding = jax.sharding.NamedSharding(mesh, P('x', 'y')) + >>> inp_data = np.arange(math.prod(global_shape)).reshape(global_shape) + ... + >>> arrays = [ + ... jax.device_put(inp_data[index], d) + ... for d, index in sharding.addressable_devices_indices_map(global_shape).items()] + ... + >>> arr = jax.make_array_from_single_device_arrays(global_shape, sharding, arrays) + >>> assert arr.shape == (8,8) # arr.shape is (8,8) regardless of jax.device_count() + + When using multiple processes, a common data pipeline is to have data parallelism across devices, + with each device receiving at least one example. In this case, the following recipe will use + `make_array_from_single_device_arrays` to create a global jax.Array. + + First, we create the per host data as Numpy arrays. + + >>> sharding = jax.sharding.NamedSharding(mesh, P(('x', 'y'),)) + >>> rows_per_device = 2 + >>> feature_length = 32 + >>> per_device_shape = (rows_per_device, feature_length) + >>> per_host_shape = (rows_per_device * len(mesh.local_devices), feature_length) + >>> per_host_generator = lambda : np.arange(np.prod(per_host_shape)).reshape(per_host_shape) + >>> per_host_data = per_host_generator() # replace with your own per-host data pipeline that outputs numpy arrays + + Second, we put the Numpy data onto the local devices as single device Jax Arrays. Then we call + make_array_from_single_device_arrays to make the global Array. + + >>> global_shape = (rows_per_device * len(sharding.device_set), ) + per_device_shape[1:] + >>> per_device_data = np.split(per_host_data, len(mesh.local_devices), axis = 0) # per device data, but on host + >>> per_device_data_on_device = jax.device_put(per_device_data, mesh.local_devices) # per device data, now on device + >>> output_global_array = jax.make_array_from_single_device_arrays(global_shape, sharding, per_device_data_on_device) + ... + >>> assert output_global_array.addressable_data(0).shape == per_device_shape + >>> assert output_global_array.shape == global_shape + + When using tensor parallelism (equivalent to sharding across both rows and columns in the + above example), the above example doesn't generate the data in the sharding that you plan + to consume it with. The most common fix is to simply load the data in this data parallel sharding + and have the reshard happen automatically within the downstream jitted function. + Depending on your use case, you might prefer to directly load sharded data, something that + ``make_array_from_single_device_arrays`` can do but will depend on your data loading pipeline + also loading in the matching sharding. Loading in a data parallel format is typically + fully satisfactory for data loading for LLM use cases. + + """ + # All input arrays should be committed. Checking it is expensive on + # single-controller systems. + if any(isinstance(arr, core.Tracer) for arr in arrays): + raise ValueError( + "jax.make_array_from_single_device_arrays requires a list of concrete" + f" arrays as input. got types {set(map(type, arrays))}") + aval = core.ShapedArray(shape, arrays[0].dtype, weak_type=False) + if dtypes.issubdtype(aval.dtype, dtypes.extended): + return aval.dtype._rules.make_sharded_array(aval, sharding, arrays, + committed=True) + # TODO(phawkins): ideally the cast() could be checked. + return ArrayImpl(aval, sharding, cast(Sequence[ArrayImpl], arrays), + committed=True) + + +core.pytype_aval_mappings[ArrayImpl] = abstract_arrays.canonical_concrete_aval +xla.pytype_aval_mappings[ArrayImpl] = op.attrgetter('aval') +xla.canonicalize_dtype_handlers[ArrayImpl] = pxla.identity +api_util._shaped_abstractify_handlers[ArrayImpl] = op.attrgetter('aval') +# TODO(jakevdp) replace this with true inheritance at the C++ level. +basearray.Array.register(ArrayImpl) + + +def _array_mlir_constant_handler(val): + return mlir.ir_constants(val._value) +mlir.register_constant_handler(ArrayImpl, _array_mlir_constant_handler) + + +# NOTE(skye): we could refactor to generate _multi_slice parameters directly +# from the input ShardingSpec, rather than the indices. However, this would +# require duplicating the ordering logic of spec_to_indices, which is more +# subtle and more likely to change than the index logic we have to support here. +def as_slice_indices(arr: Any, idx: Index) -> tuple[ + tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """Returns start_indices, limit_indices, removed_dims""" + start_indices = [0] * arr.ndim + limit_indices = list(arr.shape) + removed_dims: list[int] = [] + + tuple_idx = idx if isinstance(idx, tuple) else (idx,) + for dim, sub_idx in enumerate(tuple_idx): + if isinstance(sub_idx, int): + start_indices[dim] = sub_idx + limit_indices[dim] = sub_idx + 1 + removed_dims.append(dim) + elif sub_idx == slice(None): + continue + else: + assert isinstance(sub_idx, slice), sub_idx + assert isinstance(sub_idx.start, int), sub_idx + assert isinstance(sub_idx.stop, int), sub_idx + start_indices[dim] = sub_idx.start + limit_indices[dim] = sub_idx.stop + + return tuple(start_indices), tuple(limit_indices), tuple(removed_dims) + + +def shard_device_array(x, devices, indices, sharding): + start_indices, limit_indices, removed_dims = unzip3( + as_slice_indices(x, idx) for idx in indices) + if sharding.is_fully_replicated: + shards = [x] * len(devices) + else: + shards = x._multi_slice(start_indices, limit_indices, removed_dims) + aval = api_util.shaped_abstractify(x) + return pxla.batched_device_put(aval, sharding, shards, devices) + +def _hashable_index(idx): + return tree_util.tree_map( + lambda x: (x.start, x.stop) if type(x) == slice else x, idx) + + +def shard_sharded_device_array_slow_path(x, devices, indices, sharding): + candidates = defaultdict(list) + bufs = [buf.data for buf in x.addressable_shards] + arr_indices = tuple(x.sharding.devices_indices_map(x.shape).values()) + for buf, idx in safe_zip(bufs, arr_indices): + candidates[_hashable_index(idx)].append(buf) + + bufs = [] + for idx, device in safe_zip(indices, devices): + # Look up all buffers that contain the correct slice of the logical array. + candidates_list = candidates[_hashable_index(idx)] + if not candidates_list: + # This array isn't sharded correctly. Reshard it via host roundtrip. + # TODO(skye): more efficient reshard? + return pxla.shard_arg(x._value, sharding, canonicalize=False) + # Try to find a candidate buffer already on the correct device, + # otherwise copy one of them. + for buf in candidates_list: + if buf.devices() == {device}: + bufs.append(buf) + break + else: + bufs.append(buf) + + return pxla.batched_device_put(x.aval, sharding, bufs, devices) + + +@functools.lru_cache(maxsize=4096) +def _sharding_indices_and_eq(src_sharding, shape, dst_sharding): + src_indices = src_sharding.addressable_devices_indices_map(shape).values() + dst_indices = dst_sharding.addressable_devices_indices_map(shape).values() + return dst_indices, tuple(src_indices) == tuple(dst_indices) + + +def _array_shard_arg(x, sharding): + x._check_if_deleted() + + indices, same_indices = _sharding_indices_and_eq(x.sharding, x.shape, sharding) + if not x.is_fully_addressable: + if same_indices: + return x + else: + raise NotImplementedError( + "Cannot reshard an input that is not fully addressable") + else: + devices = sharding._addressable_device_assignment + if same_indices: + return xc.copy_array_to_devices_with_sharding(x, list(devices), sharding) + # Resharding starts here: + if dispatch.is_single_device_sharding(x.sharding): + return shard_device_array(x, devices, indices, sharding) + else: + return shard_sharded_device_array_slow_path(x, devices, indices, sharding) + +pxla.shard_arg_handlers[ArrayImpl] = _array_shard_arg + + +def _array_global_result_handler(global_aval, out_sharding, committed): + if global_aval.dtype == dtypes.float0: + return lambda _: np.zeros(global_aval.shape, dtypes.float0) + if dtypes.issubdtype(global_aval.dtype, dtypes.extended): + return global_aval.dtype._rules.global_sharded_result_handler( + global_aval, out_sharding, committed) + return xc.array_result_handler( + global_aval, out_sharding, committed=committed, _skip_checks=True + ) +pxla.global_result_handlers[core.ShapedArray] = _array_global_result_handler +pxla.global_result_handlers[core.ConcreteArray] = _array_global_result_handler + +# Only used for Arrays that come out of pmap. +def _array_local_result_handler(aval, sharding, indices): + if aval.dtype == dtypes.float0: + return lambda _: np.zeros(aval.shape, dtypes.float0) + if dtypes.issubdtype(aval.dtype, dtypes.extended): + return aval.dtype._rules.local_sharded_result_handler( + aval, sharding, indices) + return xc.array_result_handler( + aval, sharding, committed=True, _skip_checks=True + ) +pxla.local_result_handlers[core.ShapedArray] = _array_local_result_handler +pxla.local_result_handlers[core.ConcreteArray] = _array_local_result_handler + + +# Token handlers + +def _token_shard_arg(x, sharding): + return _array_shard_arg(x._buf, sharding) +pxla.shard_arg_handlers[core.Token] = _token_shard_arg + + +def _token_global_result_handler(global_aval, out_sharding, committed): + array_handler = _array_global_result_handler( + core.token_shaped_array, out_sharding, committed) + + def wrapper(*args, **kwargs): + out_buf = array_handler(*args, **kwargs) + return core.Token(out_buf) + return wrapper +pxla.global_result_handlers[core.AbstractToken] = _token_global_result_handler diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/basearray.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/basearray.py new file mode 100644 index 0000000000000000000000000000000000000000..5809b9649f26874614a2fee9299d8c84f0e604ba --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/basearray.py @@ -0,0 +1,136 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Note that type annotations for this file are defined in basearray.pyi + +from __future__ import annotations + +import abc +import numpy as np +from typing import Any, Union +from collections.abc import Sequence + +# TODO(jakevdp): fix import cycles and define these. +Shard = Any +Sharding = Any + +# Array is a type annotation for standard JAX arrays and tracers produced by +# core functions in jax.lax and jax.numpy; it is not meant to include +# future non-standard array types like KeyArray and BInt. +class Array(abc.ABC): + """Array base class for JAX + + ``jax.Array`` is the public interface for instance checks and type annotation + of JAX arrays and tracers. Its main applications are in instance checks and + type annotations; for example:: + + x = jnp.arange(5) + isinstance(x, jax.Array) # returns True both inside and outside traced functions. + + def f(x: Array) -> Array: # type annotations are valid for traced and non-traced types. + return x + + ``jax.Array`` should not be used directly for creation of arrays; instead you + should use array creation routines offered in :mod:`jax.numpy`, such as + :func:`jax.numpy.array`, :func:`jax.numpy.zeros`, :func:`jax.numpy.ones`, + :func:`jax.numpy.full`, :func:`jax.numpy.arange`, etc. + """ + # Note: abstract methods for this class are defined dynamically in + # lax_numpy.py + # For the sake of static type analysis, these definitions are mirrored in the + # associated basearray.pyi file. + + __slots__ = ['__weakref__'] + + @property + @abc.abstractmethod + def dtype(self) -> np.dtype: + """The data type (:class:`numpy.dtype`) of the array.""" + + @property + @abc.abstractmethod + def ndim(self) -> int: + """The number of dimensions in the array.""" + + @property + @abc.abstractmethod + def size(self) -> int: + """The total number of elements in the array.""" + + @property + @abc.abstractmethod + def shape(self) -> tuple[int, ...]: + """The shape of the array.""" + + # Documentation for sharding-related methods and properties defined on ArrayImpl: + @abc.abstractmethod + def addressable_data(self, index: int) -> Array: + """Return an array of the addressable data at a particular index.""" + + @property + @abc.abstractmethod + def addressable_shards(self) -> Sequence[Shard]: + """List of addressable shards.""" + + @property + @abc.abstractmethod + def global_shards(self) -> Sequence[Shard]: + """List of global shards.""" + + @property + @abc.abstractmethod + def is_fully_addressable(self) -> bool: + """Is this Array fully addressable? + + A jax.Array is fully addressable if the current process can address all of + the devices named in the :class:`Sharding`. ``is_fully_addressable`` is + equivalent to "is_local" in multi-process JAX. + + Note that fully replicated is not equal to fully addressable i.e. + a jax.Array which is fully replicated can span across multiple hosts and is + not fully addressable. + """ + + @property + @abc.abstractmethod + def is_fully_replicated(self) -> bool: + """Is this Array fully replicated?""" + + @property + @abc.abstractmethod + def sharding(self) -> Sharding: + """The sharding for the array.""" + + +Array.__module__ = "jax" + +# StaticScalar is the Union of all scalar types that can be converted to +# JAX arrays, and are possible to mark as static arguments. +StaticScalar = Union[ + np.bool_, np.number, # NumPy scalar types + bool, int, float, complex, # Python scalar types +] +StaticScalar.__doc__ = "Type annotation for JAX-compatible static scalars." + + +# ArrayLike is a Union of all objects that can be implicitly converted to a +# standard JAX array (i.e. not including future non-standard array types like +# KeyArray and BInt). It's different than np.typing.ArrayLike in that it doesn't +# accept arbitrary sequences, nor does it accept string data. +ArrayLike = Union[ + Array, # JAX array type + np.ndarray, # NumPy array type + StaticScalar, # valid scalars +] +ArrayLike.__doc__ = "Type annotation for JAX array-like objects." diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/basearray.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/basearray.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6f37c16e671508376867c4f28396b9b228e2037e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/basearray.pyi @@ -0,0 +1,259 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import abc +from typing import Any, Callable, Optional, Sequence, Union +import numpy as np + +from jax._src.sharding import Sharding + +Shard = Any + +# TODO: alias this to xla_client.Traceback +Device = Any +Traceback = Any + + +class Array(abc.ABC): + aval: Any + + @property + def dtype(self) -> np.dtype: ... + + @property + def ndim(self) -> int: ... + + @property + def size(self) -> int: ... + + @property + def itemsize(self) -> int: ... + + @property + def shape(self) -> tuple[int, ...]: ... + + def __init__(self, shape, dtype=None, buffer=None, offset=0, strides=None, + order=None): + raise TypeError("jax.numpy.ndarray() should not be instantiated explicitly." + " Use jax.numpy.array, or jax.numpy.zeros instead.") + + def __getitem__(self, key) -> Array: ... + def __setitem__(self, key, value) -> None: ... + def __len__(self) -> int: ... + def __iter__(self) -> Any: ... + def __reversed__(self) -> Any: ... + def __round__(self, ndigits=None) -> Array: ... + + # Comparisons + + # these return bool for object, so ignore override errors. + def __lt__(self, other) -> Array: ... + def __le__(self, other) -> Array: ... + def __eq__(self, other) -> Array: ... # type: ignore[override] + def __ne__(self, other) -> Array: ... # type: ignore[override] + def __gt__(self, other) -> Array: ... + def __ge__(self, other) -> Array: ... + + # Unary arithmetic + + def __neg__(self) -> Array: ... + def __pos__(self) -> Array: ... + def __abs__(self) -> Array: ... + def __invert__(self) -> Array: ... + + # Binary arithmetic + + def __add__(self, other) -> Array: ... + def __sub__(self, other) -> Array: ... + def __mul__(self, other) -> Array: ... + def __matmul__(self, other) -> Array: ... + def __truediv__(self, other) -> Array: ... + def __floordiv__(self, other) -> Array: ... + def __mod__(self, other) -> Array: ... + def __divmod__(self, other) -> Array: ... + def __pow__(self, other) -> Array: ... + def __lshift__(self, other) -> Array: ... + def __rshift__(self, other) -> Array: ... + def __and__(self, other) -> Array: ... + def __xor__(self, other) -> Array: ... + def __or__(self, other) -> Array: ... + + def __radd__(self, other) -> Array: ... + def __rsub__(self, other) -> Array: ... + def __rmul__(self, other) -> Array: ... + def __rmatmul__(self, other) -> Array: ... + def __rtruediv__(self, other) -> Array: ... + def __rfloordiv__(self, other) -> Array: ... + def __rmod__(self, other) -> Array: ... + def __rdivmod__(self, other) -> Array: ... + def __rpow__(self, other) -> Array: ... + def __rlshift__(self, other) -> Array: ... + def __rrshift__(self, other) -> Array: ... + def __rand__(self, other) -> Array: ... + def __rxor__(self, other) -> Array: ... + def __ror__(self, other) -> Array: ... + + def __bool__(self) -> bool: ... + def __complex__(self) -> complex: ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __index__(self) -> int: ... + + def __buffer__(self, flags: int) -> memoryview: ... + def __release_buffer__(self, view: memoryview) -> None: ... + + # np.ndarray methods: + def all(self, axis: Optional[Union[int, Sequence[int]]] = None, out=None, + keepdims=None, *, where: Optional[ArrayLike] = ...) -> Array: ... + def any(self, axis: Optional[Union[int, Sequence[int]]] = None, out=None, + keepdims=None, *, where: Optional[ArrayLike] = ...) -> Array: ... + def argmax(self, axis: Optional[int] = None, out=None, keepdims=None) -> Array: ... + def argmin(self, axis: Optional[int] = None, out=None, keepdims=None) -> Array: ... + def argpartition(self, kth, axis=-1, kind='introselect', order=None) -> Array: ... + def argsort(self, axis: Optional[int] = -1, kind='quicksort', order=None) -> Array: ... + def astype(self, dtype) -> Array: ... + def choose(self, choices, out=None, mode='raise') -> Array: ... + def clip(self, min=None, max=None, out=None) -> Array: ... + def compress(self, condition, axis: Optional[int] = None, out=None) -> Array: ... + def conj(self) -> Array: ... + def conjugate(self) -> Array: ... + def copy(self) -> Array: ... + def cumprod(self, axis: Optional[Union[int, Sequence[int]]] = None, + dtype=None, out=None) -> Array: ... + def cumsum(self, axis: Optional[Union[int, Sequence[int]]] = None, + dtype=None, out=None) -> Array: ... + def diagonal(self, offset=0, axis1: int = 0, axis2: int = 1) -> Array: ... + def dot(self, b, *, precision=None) -> Array: ... + def flatten(self) -> Array: ... + @property + def imag(self) -> Array: ... + def item(self, *args) -> Any: ... + def max(self, axis: Optional[Union[int, Sequence[int]]] = None, out=None, + keepdims=None, initial=None, where=None) -> Array: ... + def mean(self, axis: Optional[Union[int, Sequence[int]]] = None, dtype=None, + out=None, keepdims=False, *, where=None,) -> Array: ... + def min(self, axis: Optional[Union[int, Sequence[int]]] = None, out=None, + keepdims=None, initial=None, where=None) -> Array: ... + @property + def nbytes(self) -> int: ... + def nonzero(self, *, size=None, fill_value=None) -> Array: ... + def prod(self, axis: Optional[Union[int, Sequence[int]]] = None, dtype=None, + out=None, keepdims=None, initial=None, where=None) -> Array: ... + def ptp(self, axis: Optional[Union[int, Sequence[int]]] = None, out=None, + keepdims=False,) -> Array: ... + def ravel(self, order='C') -> Array: ... + @property + def real(self) -> Array: ... + def repeat(self, repeats, axis: Optional[int] = None, *, + total_repeat_length=None) -> Array: ... + def reshape(self, *args, order='C') -> Array: ... + def round(self, decimals=0, out=None) -> Array: ... + def searchsorted(self, v, side='left', sorter=None) -> Array: ... + def sort(self, axis: Optional[int] = -1, kind='quicksort', order=None) -> Array: ... + def squeeze(self, axis: Optional[Union[int, Sequence[int]]] = None) -> Array: ... + def std(self, axis: Optional[Union[int, Sequence[int]]] = None, + dtype=None, out=None, ddof=0, keepdims=False, *, where=None) -> Array: ... + def sum(self, axis: Optional[Union[int, Sequence[int]]] = None, dtype=None, + out=None, keepdims=None, initial=None, where=None) -> Array: ... + def swapaxes(self, axis1: int, axis2: int) -> Array: ... + def take(self, indices, axis: Optional[int] = None, out=None, + mode=None) -> Array: ... + def tobytes(self, order='C') -> bytes: ... + def tolist(self) -> list[Any]: ... + def trace(self, offset=0, axis1: int = 0, axis2: int = 1, dtype=None, + out=None) -> Array: ... + def transpose(self, *args) -> Array: ... + @property + def T(self) -> Array: ... + @property + def mT(self) -> Array: ... + def var(self, axis: Optional[Union[int, Sequence[int]]] = None, + dtype=None, out=None, ddof=0, keepdims=False, *, where=None) -> Array: ... + def view(self, dtype=None, type=None) -> Array: ... + + # Even though we don't always support the NumPy array protocol, e.g., for + # tracer types, for type checking purposes we must declare support so we + # implement the NumPy ArrayLike protocol. + def __array__(self, dtype: Optional[np.dtype] = ..., + copy: Optional[bool] = ...) -> np.ndarray: ... + def __dlpack__(self) -> Any: ... + + # JAX extensions + @property + def at(self) -> _IndexUpdateHelper: ... + @property + def weak_type(self) -> bool: ... + + # Methods defined on ArrayImpl, but not on Tracers + def addressable_data(self, index: int) -> Array: ... + def block_until_ready(self) -> Array: ... + def copy_to_host_async(self) -> None: ... + def delete(self) -> None: ... + def devices(self) -> set[Device]: ... + @property + def sharding(self) -> Sharding: ... + @property + def addressable_shards(self) -> Sequence[Shard]: ... + @property + def global_shards(self) -> Sequence[Shard]: ... + def is_deleted(self) -> bool: ... + @property + def is_fully_addressable(self) -> bool: ... + @property + def is_fully_replicated(self) -> bool: ... + def on_device_size_in_bytes(self) -> int: ... + @property + def traceback(self) -> Traceback: ... + def unsafe_buffer_pointer(self) -> int: ... + + +StaticScalar = Union[ + np.bool_, np.number, # NumPy scalar types + bool, int, float, complex, # Python scalar types +] + +ArrayLike = Union[ + Array, # JAX array type + np.ndarray, # NumPy array type + StaticScalar, # valid scalars +] + + +# TODO: restructure to avoid re-defining this here? +# from jax._src.numpy.lax_numpy import _IndexUpdateHelper + +class _IndexUpdateHelper: + def __getitem__(self, index: Any) -> _IndexUpdateRef: ... + +class _IndexUpdateRef: + def get(self, indices_are_sorted: bool = False, unique_indices: bool = False, + mode: Optional[str] = None, fill_value: Optional[StaticScalar] = None) -> Array: ... + def set(self, values: Any, + indices_are_sorted: bool = False, unique_indices: bool = False, + mode: Optional[str] = None, fill_value: Optional[StaticScalar] = None) -> Array: ... + def add(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def mul(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def multiply(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def divide(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def power(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def min(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def max(self, values: Any, indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... + def apply(self, func: Callable[[ArrayLike], ArrayLike], indices_are_sorted: bool = False, + unique_indices: bool = False, mode: Optional[str] = None) -> Array: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cache_key.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cache_key.py new file mode 100644 index 0000000000000000000000000000000000000000..6fdf0c600b7da49cea6d62d60e65af0346e0e2fe --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cache_key.py @@ -0,0 +1,288 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import hashlib +import io +import logging +import os +import sys +from typing import cast as type_cast + +from jax._src import config +from jax._src.lib import version_str as jaxlib_version_str +from jax._src.lib import xla_client +from jax._src.lib.mlir import ir +from jax._src.lib.mlir import passmanager as pm +import numpy as np + + +logger = logging.getLogger(__name__) + +_extra_flag_prefixes: list[str] = [] + +def add_flag_prefixes(flag_prefixes: list[str]) -> None: + """Add flag prefixes to include in the cache key. Call prior to get(). + """ + global _extra_flag_prefixes + _extra_flag_prefixes += flag_prefixes + + +def clear_flag_prefixes() -> None: + """Clear flag prefixes added by add_flag_prefixes(). + """ + global _extra_flag_prefixes + _extra_flag_prefixes = [] + + +def get_flag_prefixes() -> list[str]: + """Return flag prefixes added by add_flag_prefixes(). + """ + return _extra_flag_prefixes + + +def custom_hook() -> str: + """Custom hook for any addition to the cache key. + + The custom hook will be called everytime get() is called and can be + defined to return a string that will be hashed into the cache key. + """ + return "" + + +def get(module: ir.Module, + devices: np.ndarray, + compile_options: xla_client.CompileOptions, + backend: xla_client.Client, + compression_algorithm: str = "zstandard") -> str: + """Creates a hashed string to use as a key to the compilation cache. + + Creates a cache key that is a hex-encoded string of a unique hash based on + the arguments. The hex-encoded string is 256 characters long. + + Args: + module: the input program + devices: an array of accelerator devices that the program will run on + compile_options: options passed to the XLA compiler + backend: description of the platform (e.g., TPU version) + compression_algorithm: a string representing the compression algorithm used + for the executable before persisting in the cache + + Typical return value example: + 'jit__psum-14ac577cdb2ef6d986078b4054cc9893a9a14a16dbb0d8f37b89167c1f1aacdf' + """ + entries = [ + ("computation", lambda hash_obj: _hash_computation(hash_obj, module)), + ("jax_lib version", + lambda hash_obj: hash_obj.update( + bytes(jaxlib_version_str.encode("utf-8")))), + ("XLA flags", + lambda hash_obj: _hash_xla_flags(hash_obj, get_flag_prefixes())), + ("compile_options", + lambda hash_obj: _hash_serialized_compile_options( + hash_obj, compile_options, + # In case of GPU multi-process tasks we need to strip device + # assignment to use cache key as invariant between processes. + strip_device_assignment=(backend.platform == "gpu"))), + ("accelerator_config", + lambda hash_obj: _hash_accelerator_config(hash_obj, devices, backend)), + ("compression", + lambda hash_obj: _hash_string(hash_obj, compression_algorithm)), + ("custom_hook", lambda hash_obj: _hash_string(hash_obj, custom_hook())), + ] + + hash_obj = hashlib.sha256() + for name, hashfn in entries: + hashfn(hash_obj) + _log_cache_key_hash(hash_obj, name, hashfn) + sym_name = module.operation.attributes['sym_name'] + module_name = ir.StringAttr(sym_name).value + return module_name + "-" + hash_obj.digest().hex() + + +def _log_cache_key_hash(hash_obj, last_serialized: str, hashfn): + if logger.isEnabledFor(logging.DEBUG): + # Log the hash of just this entry + fresh_hash_obj = hashlib.sha256() + hashfn(fresh_hash_obj) + logger.debug( + "get_cache_key hash of serialized %s: %s", + last_serialized, + fresh_hash_obj.digest().hex(), + ) + # Log the cumulative hash + logger.debug( + "get_cache_key hash after serializing %s: %s", + last_serialized, + hash_obj.digest().hex(), + ) + + +def _serialize_ir(m: ir.Module) -> bytes: + output = io.BytesIO() + m.operation.write_bytecode(file=output) + return output.getvalue() + + +def _canonicalize_ir(m_original: ir.Module) -> bytes: + with m_original.context: + m = type_cast(ir.Module, m_original.operation.clone()) + passes = pm.PassManager.parse( + "builtin.module(strip-debuginfo)" + ) + passes.run(m.operation) + return _serialize_ir(m) + + +def _hash_computation(hash_obj, module): + if config.compilation_cache_include_metadata_in_key.value: + canonical_ir = _serialize_ir(module) + else: + canonical_ir = _canonicalize_ir(module) + hash_obj.update(canonical_ir) + + +def _hash_devices(hash_obj, devices: np.ndarray) -> None: + for device in devices.flat: + _hash_string(hash_obj, device.device_kind) + + +def _hash_accelerator_config(hash_obj, accelerators: np.ndarray, backend): + accelerator_devices = [] + for accelerator in accelerators.flat: + accelerator_devices.append(accelerator) + try: + hash_obj.update( + xla_client.get_topology_for_devices(accelerator_devices).serialize() + ) + except xla_client._xla.XlaRuntimeError as ex: + # Fall back for those backends that do not support serialized + # PjRtTopologyDescription as yet. + logger.info("get (_hash_accelerator_config): unable to hash " + "accelerator config, falling back to hashing " + "devices + platform: %s (type %s)", ex, type(ex)) + _hash_devices(hash_obj, accelerators) + _hash_platform(hash_obj, backend) + + +def _hash_serialized_compile_options(hash_obj, compile_options_obj, + strip_device_assignment=False): + # Do not mess with the original CompileOptions object since it is passed to + # the compiler. Create a deep copy for the purpose of cache key generation. + compile_options_copy = copy.deepcopy(compile_options_obj) + + # Certain debug options do not affect the compile result and thus, should not + # be part of the cache key as their inclusion will result in unnecessary cache + # misses. Clear them here by setting bool values to False, ints to 0, and + # strings to empty. The exact values used to clear are not relevant as long + # as the same values are used every time for each field. + debug_options = compile_options_copy.executable_build_options.debug_options + # LINT.IfChange(debug_options) + debug_options.xla_force_host_platform_device_count = 0 + debug_options.xla_dump_to = "" + debug_options.xla_dump_hlo_module_re = "" + debug_options.xla_dump_hlo_pass_re = "" + debug_options.xla_dump_hlo_as_text = False + debug_options.xla_dump_hlo_as_proto = False + debug_options.xla_dump_hlo_as_dot = False + debug_options.xla_dump_hlo_as_url = False + debug_options.xla_dump_hlo_as_html = False + debug_options.xla_dump_fusion_visualization = False + debug_options.xla_dump_hlo_snapshots = False + debug_options.xla_dump_max_hlo_modules = False + debug_options.xla_dump_module_metadata = False + debug_options.xla_dump_compress_protos = False + debug_options.xla_dump_hlo_as_long_text = False + debug_options.xla_dump_disable_metadata = False + debug_options.xla_dump_hlo_pipeline_re = "" + # Optional way to specify the cuda install path to be used by the compiler. + # This could possibly affect the cuda version compiled with, but this should + # already be included in the platform information (and might not be reflected + # by the cuda path regardless, since this only hashes on the directory name + # and not the contents). It can also cause spurious cache misses if the cuda + # path changes across runs despite being the same version, so we clear it + # here. + debug_options.xla_gpu_cuda_data_dir = "" + # LINT.ThenChange(:xla_flags) + + if strip_device_assignment and compile_options_copy.device_assignment: + replica_count = compile_options_copy.device_assignment.replica_count() + computation_count = compile_options_copy.device_assignment.computation_count() + compile_options_copy.device_assignment = xla_client.DeviceAssignment.create( + np.arange(replica_count * computation_count).reshape( + [replica_count, computation_count]) + ) + return hash_obj.update(compile_options_copy.SerializeAsString()) + + +def _hash_platform(hash_obj, backend): + _hash_string(hash_obj, backend.platform) + _hash_string(hash_obj, backend.platform_version) + _hash_string(hash_obj, backend.runtime_type) + + +def _hash_xla_flags(hash_obj, extra_flag_prefixes: list[str]): + # LINT.IfChange(xla_flags) + xla_flags_to_exclude_from_cache_key = [ + "--xla_dump_compress_protos", + "--xla_dump_module_metadata", + "--xla_dump_max_hlo_modules", + "--xla_dump_include_timestamp", + "--xla_dump_hlo_pass_re", + "--xla_dump_hlo_module_re", + "--xla_dump_hlo_snapshots", + "--xla_dump_fusion_visualization", + "--xla_dump_hlo_as_url", + "--xla_dump_hlo_as_proto", + "--xla_dump_hlo_as_text", + "--xla_dump_hlo_as_long_text", + "--xla_dump_hlo_as_html", + "--xla_dump_hlo_as_dot", + "--xla_dump_to", + "--xla_force_host_platform_device_count", + "--xla_dump_disable_metadata", + "--xla_dump_hlo_pipeline_re", + "--xla_tpu_sdc_checker_streamz_metric", + "--xla_tpu_sdc_checker_enable_sdc_event_callbacks", + "--xla_gpu_cuda_data_dir", + ] + # LINT.ThenChange(:debug_options) + + xla_flags = [] + + xla_flags_env_var = os.getenv("XLA_FLAGS") + if xla_flags_env_var: + xla_flags.extend(xla_flags_env_var.split()) + libtpu_init_args_env_var = os.getenv("LIBTPU_INIT_ARGS") + if libtpu_init_args_env_var: + xla_flags.extend(libtpu_init_args_env_var.split()) + + for arg in sys.argv: + if arg.startswith("--xla") or any( + arg.startswith(p) for p in extra_flag_prefixes + ): + xla_flags.append(arg) + + # N.B. all XLA flags that take an argument must use '=' and not a space + # (e.g. --xla_force_host_platform_device_count=8) (I think). + for flag in sorted(xla_flags): + if flag.split("=")[0] in xla_flags_to_exclude_from_cache_key: + logger.debug("Not including XLA flag in cache key: %s", flag) + continue + logger.debug("Including XLA flag in cache key: %s", flag) + _hash_string(hash_obj, flag) + + +def _hash_string(hash_obj, str_var): + hash_obj.update(str_var.encode("utf-8").strip()) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/callback.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/callback.py new file mode 100644 index 0000000000000000000000000000000000000000..8e9ab0b62593ee02aa8daaf0aa165f617421d1c5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/callback.py @@ -0,0 +1,522 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Module for JAX callbacks.""" +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +import functools +import logging +from typing import Any, Callable + +import jax +from jax._src import core +from jax._src import dispatch +from jax._src import dtypes +from jax._src import effects +from jax._src import sharding_impls +from jax._src import tree_util +from jax._src import util +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.lax.control_flow.loops import map as lax_map +from jax._src.lib import xla_client as xc +from jax._src.sharding_impls import SingleDeviceSharding +import numpy as np + +logger = logging.getLogger(__name__) + + +# `pure_callback_p` is the main primitive for staging out Python pure callbacks. +pure_callback_p = core.Primitive("pure_callback") +pure_callback_p.multiple_results = True +dispatch.prim_requires_devices_during_lowering.add(pure_callback_p) + +map, unsafe_map = util.safe_map, map + + +@dataclasses.dataclass(frozen=True) +class _FlatCallback: + """A Python function callable with flat arguments and results. + + An instance of this class is used as a parameter for the callback primitives. + We prefer it to an anonymous flattened function because it produces + equal objects when we call the same Python function with the same argument + structure. + """ + callback_func: Callable[..., Any] + in_tree: tree_util.PyTreeDef # (args, kwargs) pytree for `callback_func`. + + def __call__(self, *flat_args: jax.Array) -> Sequence[jax.Array]: + args, kwargs = tree_util.tree_unflatten(self.in_tree, flat_args) + return tree_util.tree_leaves(self.callback_func(*args, **kwargs)) + + +def pure_callback_impl( + *args, + result_avals, + callback: _FlatCallback, + sharding: SingleDeviceSharding | None, + vectorized: bool, +): + del sharding, vectorized, result_avals + try: + cpu_device, *_ = jax.local_devices(backend="cpu") + except RuntimeError as e: + raise RuntimeError( + "jax.pure_callback failed to find a local CPU device to place the" + " inputs on. Make sure \"cpu\" is listed in --jax_platforms or the" + " JAX_PLATFORMS environment variable." + ) from e + args = jax.device_put(args, cpu_device) + with jax.default_device(cpu_device): + try: + return tree_util.tree_map(np.asarray, callback(*args)) + except BaseException: + logger.exception("jax.pure_callback failed") + raise + + +pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive, + pure_callback_p)) + + +@pure_callback_p.def_abstract_eval +def pure_callback_abstract_eval( + *avals, + callback: _FlatCallback, + result_avals, + sharding: SingleDeviceSharding | None, + vectorized: bool, +): + del avals, callback, sharding, vectorized + return result_avals + + +def pure_callback_jvp_rule(*args, **kwargs): + del args, kwargs + raise ValueError( + "Pure callbacks do not support JVP. " + "Please use `jax.custom_jvp` to use callbacks while taking gradients.") + + +ad.primitive_jvps[pure_callback_p] = pure_callback_jvp_rule + + +def pure_callback_transpose_rule(*args, **kwargs): + del args, kwargs + raise ValueError( + "Pure callbacks do not support transpose. " + "Please use `jax.custom_vjp` to use callbacks while taking gradients.") + +ad.primitive_transposes[pure_callback_p] = pure_callback_transpose_rule + + +def pure_callback_batching_rule( + args, + dims, + *, + callback: _FlatCallback, + sharding: SingleDeviceSharding | None, + vectorized: bool, + result_avals: Sequence[core.ShapedArray], +): + axis_size = next(a.shape[d] for a, d in zip(args, dims) + if d is not batching.not_mapped) + new_args = [arg if dim is batching.not_mapped else + batching.moveaxis(arg, dim, 0) for arg, dim in zip(args, dims)] + if vectorized: + result_avals = tuple( + core.unmapped_aval(axis_size, core.no_axis_name, 0, aval) # type: ignore + for aval in result_avals) + outvals = pure_callback_p.bind( + *new_args, + callback=callback, + sharding=sharding, + vectorized=vectorized, + result_avals=result_avals, + ) + else: + is_batched = [d is not batching.not_mapped for d in dims] + unbatched_args, batched_args = util.partition_list(is_batched, new_args) + def _batch_fun(batched_args): + merged_args = util.merge_lists(is_batched, unbatched_args, batched_args) + return pure_callback_p.bind( + *merged_args, + callback=callback, + sharding=sharding, + result_avals=result_avals, + vectorized=vectorized, + ) + outvals = lax_map(_batch_fun, batched_args) + return tuple(outvals), (0,) * len(outvals) + + +batching.primitive_batchers[pure_callback_p] = pure_callback_batching_rule + + +def _callback_op_sharding(axis_context, sharding: SingleDeviceSharding | None): + if isinstance(axis_context, sharding_impls.SPMDAxisContext): + # If we have fully manual sharding during lowering, that means the JAX + # program has per-device semantics, so we run the callback on each device. + if axis_context.manual_axes != frozenset(axis_context.mesh.axis_names): + raise NotImplementedError( + "callbacks are only supported in spmd computations when all mesh" + " axes are partitioned manually (no partial automatic sharding)." + ) + if sharding is not None: + raise NotImplementedError( + "callbacks do not support specifying sharding inside spmd" + " computations" + ) + op_sharding = xc.OpSharding() + op_sharding.type = xc.OpSharding.Type.MANUAL + return op_sharding + + if isinstance(axis_context, sharding_impls.ShardingContext): + if sharding is not None: + if not isinstance(sharding, SingleDeviceSharding): + raise NotImplementedError( + "pure_callback only supports SingleDeviceSharding, but got" + f" {type(sharding)}" + ) + device = next(iter(sharding.device_set)) + device_assignment = axis_context.device_assignment + if device_assignment is None: + raise AssertionError( + "Please file a bug at https://github.com/google/jax/issues") + try: + device_index = device_assignment.index(device) + except IndexError as e: + raise ValueError( + "Sharding provided to pure_callback specifies a device" + f" {device} that is not in the device assignment" + f" ({device_assignment})") from e + else: + device_index = 0 + + # If we have fully automatic sharding during lowering, that means the JAX + # program has bulk array semantics, so we run the callback with a MAXIMAL + # sharding and hence execute it only once on the full logical value). + op_sharding = xc.OpSharding() + op_sharding.type = xc.OpSharding.Type.MAXIMAL + op_sharding.tile_assignment_dimensions = [1] + op_sharding.tile_assignment_devices = [device_index] + return op_sharding + + # When there's no SPMD partitioning going on, don't annotate a sharding. + return None + + +def pure_callback_lowering( + ctx, *args, callback: _FlatCallback, sharding: SingleDeviceSharding | None, **params +): + def _callback(*flat_args): + return tuple( + pure_callback_impl( + *flat_args, + callback=callback, + sharding=None, # unused. + **params, + ) + ) + + op_sharding = _callback_op_sharding(ctx.module_context.axis_context, sharding) + result, _, _ = mlir.emit_python_callback( + ctx, + _callback, + None, + list(args), + ctx.avals_in, + ctx.avals_out, + has_side_effect=False, + sharding=op_sharding, + ) + return result + + +mlir.register_lowering(pure_callback_p, pure_callback_lowering) + +def _check_shape_dtype(shape_dtype): + dt = np.dtype(shape_dtype.dtype) + if dtypes.canonicalize_dtype(dt) != dt: + raise ValueError( + "result_shape_dtypes cannot specify 64-bit types when `jax_enable_x64` is disabled") + + +def pure_callback( + callback: Callable[..., Any], + result_shape_dtypes: Any, + *args: Any, + sharding: SingleDeviceSharding | None = None, + vectorized: bool = False, + **kwargs: Any, +): + """Calls a pure Python callback. Works under :func:`jit`/:func:`~vmap`/etc. + + For more explanation, see `External Callbacks`_. + + ``pure_callback`` enables calling a Python function in JIT-ed JAX functions. + The input ``callback`` will be passed JAX arrays placed on a local CPU, and + it should also return JAX arrays on CPU. + + The callback is treated as functionally pure, meaning it has no side-effects + and its output value depends only on its argument values. As a consequence, it + is safe to be called multiple times (e.g. when transformed by :func:`~vmap` or + :func:`~pmap`), or not to be called at all when e.g. the output of a + `jit`-decorated function has no data dependence on its value. Pure callbacks + may also be reordered if data-dependence allows. + + When `vmap`-ed the behavior will depend on the value of the + ``vectorized`` keyword argument. When ``vectorized`` is ``True``, the callback + is assumed to obey + ``jax.vmap(callback)(xs) == callback(xs) == jnp.stack([callback(x) for x in xs])``. + Therefore, the callback will be called directly on batched inputs (where the + batch axes are the leading dimensions). Additionally, the callbacks should + return outputs that have corresponding leading batch axes. If not vectorized + ``callback`` will be mapped sequentially across the batched axis. + For example, if ``callback = lambda x, y: np.matmul(x, y)``, then we are free + to set ``vectorized=True`` because the ``np.matmul`` function handles + arbitrary leading batch dimensions. + + Args: + callback: function to execute on the host. The callback is assumed to be a pure + function (i.e. one without side-effects): if an impure function is passed, it + may behave in unexpected ways, particularly under transformation. The callable + will be passed PyTrees of arrays as arguments, and should return a PyTree of + arrays that matches ``result_shape_dtypes``. + result_shape_dtypes: pytree whose leaves have ``shape`` and ``dtype`` attributes, + whose structure matches the expected output of the callback function at runtime. + :class:`jax.ShapeDtypeStruct` is often used to define leaf values. + *args: arguments to be passed to the callback function + sharding: optional sharding that specifies the device from which the callback should + be invoked. + vectorized: boolean specifying whether the callback function can operate in a + vectorized manner. + **kwargs: keyword arguments to be passed to the callback function + + Returns: + result: a pytree of :class:`jax.Array` objects whose structure matches that of + ``result_shape_dtypes``. + + See Also: + - :func:`jax.experimental.io_callback`: callback designed for impure functions. + - :func:`jax.debug.callback`: callback designed for general-purpose debugging. + - :func:`jax.debug.print`: callback designed for printing. + + .. _External Callbacks: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html + """ + flat_args, in_tree = tree_util.tree_flatten((args, kwargs)) + tree_util.tree_map(_check_shape_dtype, result_shape_dtypes) + result_avals = tree_util.tree_map( + lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes) + flat_result_avals, out_tree = tree_util.tree_flatten(result_avals) + out_flat = pure_callback_p.bind( + *flat_args, + callback=_FlatCallback(callback, in_tree), + result_avals=tuple(flat_result_avals), + sharding=sharding, + vectorized=vectorized, + ) + return tree_util.tree_unflatten(out_tree, out_flat) + + +# IO Callback + +io_callback_p = core.Primitive("io_callback") +io_callback_p.multiple_results = True +dispatch.prim_requires_devices_during_lowering.add(io_callback_p) + +class IOEffect(effects.Effect): + __str__ = lambda _: "IO" + +class OrderedIOEffect(effects.Effect): + __str__ = lambda _: "OrderedIO" + +_IOEffect = IOEffect() +_OrderedIOEffect = OrderedIOEffect() +effects.lowerable_effects.add_type(IOEffect) +effects.lowerable_effects.add_type(OrderedIOEffect) +effects.control_flow_allowed_effects.add_type(IOEffect) +effects.control_flow_allowed_effects.add_type(OrderedIOEffect) +effects.ordered_effects.add_type(OrderedIOEffect) +effects.shardable_ordered_effects.add_type(OrderedIOEffect) + + +def io_callback_impl( + *args, + result_avals, + callback: _FlatCallback, + sharding: SingleDeviceSharding | None, + ordered: bool, +): + del result_avals, sharding, ordered + try: + cpu_device, *_ = jax.local_devices(backend="cpu") + except RuntimeError as e: + raise RuntimeError( + "jax.io_callback failed to find a local CPU device to place the" + " inputs on. Make sure \"cpu\" is listed in --jax_platforms or the" + " JAX_PLATFORMS environment variable." + ) from e + args = jax.device_put(args, cpu_device) + with jax.default_device(cpu_device): + try: + return tree_util.tree_map(np.asarray, callback(*args)) + except BaseException: + logger.exception("jax.io_callback failed") + raise + + +io_callback_p.def_impl(functools.partial(dispatch.apply_primitive, + io_callback_p)) + + +@io_callback_p.def_effectful_abstract_eval +def io_callback_abstract_eval( + *avals, + callback: _FlatCallback, + result_avals, + sharding: SingleDeviceSharding | None, + ordered: bool, +): + del avals, sharding, callback + effect = _OrderedIOEffect if ordered else _IOEffect + return result_avals, {effect} + + +def io_callback_jvp_rule(*args, **kwargs): + del args, kwargs + raise ValueError("IO callbacks do not support JVP.") +ad.primitive_jvps[io_callback_p] = io_callback_jvp_rule + + +def io_callback_transpose_rule(*args, **kwargs): + del args, kwargs + raise ValueError("IO callbacks do not support transpose.") +ad.primitive_transposes[io_callback_p] = io_callback_transpose_rule + + +def io_callback_batching_rule( + args, dims, callback, result_avals, sharding, ordered +): + if ordered: + raise ValueError("Cannot `vmap` ordered IO callback.") + is_batched = [d is not batching.not_mapped for d in dims] + new_args = [arg if dim is batching.not_mapped else + batching.moveaxis(arg, dim, 0) for arg, dim in zip(args, dims)] + unbatched_args, batched_args = util.partition_list(is_batched, new_args) + def _batch_fun(batched_args): + merged = util.merge_lists(is_batched, unbatched_args, batched_args) + return io_callback_p.bind(*merged, callback=callback, sharding=sharding, + result_avals=result_avals, ordered=False) + out_vals = lax_map(_batch_fun, batched_args) + return out_vals, (0,) * len(out_vals) +batching.primitive_batchers[io_callback_p] = io_callback_batching_rule + + +def io_callback_lowering(ctx, *args, callback, sharding, ordered, **params): + def _callback(*flat_args): + return tuple( + io_callback_impl( + *flat_args, + callback=callback, + sharding=None, # unused. + ordered=ordered, + **params, + ) + ) + + op_sharding = _callback_op_sharding(ctx.module_context.axis_context, sharding) + if ordered: + token = ctx.tokens_in.get(_OrderedIOEffect)[0] + result, token, _ = mlir.emit_python_callback( + ctx, + _callback, + token, + list(args), + ctx.avals_in, + ctx.avals_out, + has_side_effect=True, + sharding=op_sharding, + ) + ctx.set_tokens_out(mlir.TokenSet({_OrderedIOEffect: (token,)})) + else: + result, token, _ = mlir.emit_python_callback( + ctx, + _callback, + None, + list(args), + ctx.avals_in, + ctx.avals_out, + has_side_effect=True, + sharding=op_sharding, + ) + return result + + +mlir.register_lowering(io_callback_p, io_callback_lowering) + + +def io_callback( + callback: Callable[..., Any], + result_shape_dtypes: Any, + *args: Any, + sharding: SingleDeviceSharding | None = None, + ordered: bool = False, + **kwargs: Any, +): + """Calls an impure Python callback. + + For more explanation, see `External Callbacks`_. + + Args: + callback: function to execute on the host. It is assumed to be an impure function. + If ``callback`` is pure, using :func:`jax.pure_callback` instead may lead to + more efficient execution. + result_shape_dtypes: pytree whose leaves have ``shape`` and ``dtype`` attributes, + whose structure matches the expected output of the callback function at runtime. + :class:`jax.ShapeDtypeStruct` is often used to define leaf values. + *args: arguments to be passed to the callback function + sharding: optional sharding that specifies the device from which the callback should + be invoked. + ordered: boolean specifying whether sequential calls to callback must be ordered. + **kwargs: keyword arguments to be passed to the callback function + + Returns: + result: a pytree of :class:`jax.Array` objects whose structure matches that of + ``result_shape_dtypes``. + + See Also: + - :func:`jax.pure_callback`: callback designed for pure functions. + - :func:`jax.debug.callback`: callback designed for general-purpose debugging. + - :func:`jax.debug.print`: callback designed for printing. + + .. _External Callbacks: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html + """ + flat_args, in_tree = tree_util.tree_flatten((args, kwargs)) + tree_util.tree_map(_check_shape_dtype, result_shape_dtypes) + flat_shape_dtypes, out_tree = tree_util.tree_flatten(result_shape_dtypes) + flat_result_avals = map(lambda x: core.ShapedArray(x.shape, x.dtype), + flat_shape_dtypes) + flat_args = map(core.raise_as_much_as_possible, flat_args) + out_flat = io_callback_p.bind( + *flat_args, + callback=_FlatCallback(callback, in_tree), + result_avals=tuple(flat_result_avals), + sharding=sharding, + ordered=ordered, + ) + return tree_util.tree_unflatten(out_tree, out_flat) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/checkify.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/checkify.py new file mode 100644 index 0000000000000000000000000000000000000000..3727c8364a126f3bbc22ca4bd430d376a1d0585b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/checkify.py @@ -0,0 +1,1305 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +import functools +import itertools as it +from typing import Callable, TypeVar, Any, Union + +import numpy as np + +import jax.numpy as jnp +from jax import dtypes +from jax import lax + +from jax._src import api +from jax._src import linear_util as lu +from jax._src import config +from jax._src import core +from jax._src import custom_derivatives +from jax._src import effects +from jax._src import pjit +from jax._src import sharding_impls +from jax._src import source_info_util +from jax._src import traceback_util +from jax._src import tree_util as jtu +from jax._src.ad_util import SymbolicZero +from jax._src.api_util import flatten_fun +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import partial_eval as pe +from jax._src.tree_util import tree_flatten +from jax._src.tree_util import tree_map +from jax._src.tree_util import tree_unflatten +from jax._src.typing import Array +from jax._src.util import (as_hashable_function, split_list, safe_map, safe_zip, + unzip3, weakref_lru_cache, HashableWrapper) + +source_info_util.register_exclusion(__file__) +traceback_util.register_exclusion(__file__) + +map, unsafe_map = safe_map, map +zip, unsafe_zip = safe_zip, zip + +Bool = Union[bool, Array] +Int = Union[int, Array] +ErrorCategory = type['JaxException'] +Payload = list[Union[np.ndarray, Array]] +PyTreeDef = jtu.PyTreeDef +Out = TypeVar('Out') + +## Utils + +def popattr(obj, attrname): + val = getattr(obj, attrname) + delattr(obj, attrname) + return val + +def setnewattr(obj, name, val): + sentinel = object() + assert getattr(obj, name, sentinel) is sentinel + setattr(obj, name, val) + +# Concrete errors + +class JaxException(Exception): + """Python exception which can contain an error message with JAX run-time info.""" + + def __init__(self, traceback_info): + self.traceback_info = traceback_info + # TODO(lenamartens): re-enable tracebacks when they don't leak tracers. + # self.with_traceback(self.traceback_info) + + def __init_subclass__(cls): + jtu.register_pytree_node_class(cls) + + def tree_flatten(self): + return ([], self.traceback_info) + + @classmethod + def tree_unflatten(cls, metadata, payload): + del payload + return cls(metadata) + + def get_effect_type(self) -> ErrorEffect: + raise NotImplementedError + + +@functools.total_ordering +@dataclasses.dataclass(eq=True, frozen=True) +class ErrorEffect(effects.Effect): + error_type: type[JaxException] + shape_dtypes: tuple[api.ShapeDtypeStruct, ...] + + def __lt__(self, other: ErrorEffect): + shape_dtypes = lambda x: tuple((sd.shape, str(sd.dtype)) # dtype is not comparable + for sd in x.shape_dtypes) + unpack = lambda x: (str(x.error_type), shape_dtypes(x)) + return (unpack(self) < unpack(other)) + +effects.control_flow_allowed_effects.add_type(ErrorEffect) +effects.lowerable_effects.add_type(ErrorEffect) + +class DivisionByZeroError(JaxException): + + def __str__(self): + return 'division by zero' + + def get_effect_type(self): + return ErrorEffect(DivisionByZeroError, ()) + +class NaNError(JaxException): + + def __init__(self, traceback_info, primitive_name): + super().__init__(traceback_info) + self.prim = primitive_name + + def tree_flatten(self): + return ([], (self.traceback_info, self.prim)) + + @classmethod + def tree_unflatten(cls, metadata, _): + return cls(*metadata) + + def get_effect_type(self): + return ErrorEffect(NaNError, ()) + + def __str__(self): + return f'nan generated by primitive: {self.prim}.' + +class OOBError(JaxException): + + def __init__(self, traceback_info, primitive_name, operand_shape, payload): + super().__init__(traceback_info) + self.prim = primitive_name + self.operand_shape = operand_shape + self._payload = payload + + def tree_flatten(self): + return ([self._payload], (self.traceback_info, self.prim, self.operand_shape)) + + @classmethod + def tree_unflatten(cls, metadata, payload): + return cls(*metadata, payload[0]) + + def __str__(self): + return (f'out-of-bounds indexing for array of ' + f'shape {self.operand_shape}: ' + f'index {self._payload[0]} is out of bounds for axis ' + f'{self._payload[1]} with size {self._payload[2]}. ') + + def get_effect_type(self): + return ErrorEffect(OOBError, (api.ShapeDtypeStruct((3,), jnp.int32),)) + +class FailedCheckError(JaxException): + + def __init__(self, traceback_info, fmt_string, *a, **k): + super().__init__(traceback_info) + self.fmt_string = fmt_string + self.args = a + self.kwargs = k + + def tree_flatten(self): + return ((self.args, self.kwargs), # leaves + (self.traceback_info, self.fmt_string)) # treedef + + @classmethod + def tree_unflatten(cls, metadata, payload): + args, kwargs = payload + return cls(*metadata, *args, **kwargs) + + def __str__(self): + return (self.fmt_string.format(*self.args, **self.kwargs) + + ' (`check` failed)') + + def get_effect_type(self): + vals = jtu.tree_leaves((self.args, self.kwargs)) + return ErrorEffect( + FailedCheckError, + tuple(api.ShapeDtypeStruct(x.shape, x.dtype) for x in vals)) + +@dataclasses.dataclass +class BatchedError(JaxException): + error_mapping: dict[tuple[int, ...], JaxException] + + def __post_init__(self): + traceback_info = list(self.error_mapping.values())[0].traceback_info + super().__init__(traceback_info) + + + def __str__(self): + return '\n'.join(f'at mapped index {", ".join(map(str, idx))}: {e}' + for idx, e in self.error_mapping.items()) + + +# Error Value + +@jtu.register_pytree_node_class +@dataclasses.dataclass(frozen=True) +class Error: + _pred: dict[ErrorEffect, Bool] + _code: dict[ErrorEffect, Int] + _metadata: dict[Int, PyTreeDef] # mapping of code to JaxException treedef. + _payload: dict[ErrorEffect, Payload] + + def get(self) -> str | None: + """Returns error message if error happened, None if no error happened.""" + exp = self.get_exception() + if exp is not None: + return str(exp) + return None + + def get_exception(self) -> JaxException | None: + """Returns Python exception if error happened, None if no error happened.""" + if any(map(np.shape, self._pred.values())): + return self._get_batched_exception() + else: + min_code = None + cur_effect = None + for error_effect, code in self._code.items(): + if self._pred[error_effect]: + if min_code is None or code < min_code: + min_code = code + cur_effect = error_effect + + if cur_effect is not None: + return tree_unflatten(self._metadata[int(min_code)], # type: ignore + self._payload[cur_effect]) + return None + + def throw(self): + _check_error(self) + + def __str__(self): + return f'Error({self.get()})' + + # Internal helpers + + def _get_batched_exception(self) -> BatchedError | None: + shape = np.shape(list(self._pred.values())[0]) + error_mapping = {} + for idx in np.ndindex(*shape): + min_code = None + cur_effect = None + for error_effect, code in self._code.items(): + if self._pred[error_effect][idx]: # type: ignore + if min_code is None or code[idx] < min_code: + min_code = code[idx] # type: ignore + cur_effect = error_effect + + if cur_effect is not None: + payload = tree_map(lambda x, i=idx: x[i], self._payload[cur_effect]) + jax_error = tree_unflatten(self._metadata[int(min_code)], payload) # type: ignore + error_mapping[idx] = jax_error + if error_mapping: + return BatchedError(error_mapping) + else: + return None + + def _update(self, effect_type: ErrorEffect, pred, code, metadata, payload): + new_errs = {**self._pred, **{effect_type: pred}} + new_codes = {**self._code, **{effect_type: code}} + new_payload = {**self._payload, **{effect_type: payload}} + new_metadata = {**self._metadata, **metadata} + return Error(new_errs, new_codes, new_metadata, new_payload) + + def _add_placeholder_effects(self, effects: set[ErrorEffect]): + """Fill out Error with `effects` and np.ones arrays of their payloads.""" + new_err = self._pred.copy() + new_code = self._code.copy() + new_payload = self._payload.copy() + for effect in effects: + if effect not in self._pred.keys(): + new_err[effect] = False + new_payload[effect] = list( + tree_map(lambda a: jnp.ones(a.shape, a.dtype), effect.shape_dtypes)) + # The error value associated with this effect will never become True, so + # we don't need to set a meaningful code. + new_code[effect] = -1 + return Error(new_err, new_code, self._metadata, new_payload) + + def _replace(self, *args, **kwargs): + return dataclasses.replace(self, *args, **kwargs) + + # PyTree methods + + def tree_flatten(self): + return ((self._pred, self._code, self._payload), (self._metadata)) + + @classmethod + def tree_unflatten(cls, metadata, data): + pred, code, payload = data + return cls(pred, code, metadata, payload) + +init_error = Error({}, {}, {}, {}) # value used as initial (empty) error. +next_code = it.count(1).__next__ # globally unique ids, could be uuid4 + +def assert_func(error: Error, pred: Bool, new_error: JaxException) -> Error: + code = next_code() + effect_type = new_error.get_effect_type() + new_payload, new_metadata = tree_flatten(new_error) + return update_error(error, pred, code, {code: new_metadata}, new_payload, effect_type) + +def update_error(error, pred, code, metadata, payload, effect_type): + err_of_type = error._pred.get(effect_type, False) + out_err = err_of_type | pred + out_code = lax.select(err_of_type, error._code.get(effect_type, -1), code) + cur_payload = error._payload.get(effect_type, None) + if cur_payload is not None: + out_payload = tree_map(functools.partial(lax.select, err_of_type), cur_payload, payload) + else: + out_payload = payload + return error._update(effect_type, out_err, out_code, metadata, out_payload) + + +## Checkify transformation for plumbing functional error values. + +@lu.transformation_with_aux +def _flatten_and_get_error_metadata_thunk(*invals): + error, out = yield invals, {} + out_vals, out_tree = jtu.tree_flatten((error, out)) + yield out_vals, (out_tree, set(error._pred.keys())) + +def default_checkify_rule(primitive: core.Primitive, error: Error, + enabled_errors, *invals: core.Value, + **params: Any) -> tuple[Error, Sequence[core.Value]]: + """Default rule for primitives in `checkify` interpreter.""" + if 'call_jaxpr' not in params: + # Default non-HOP case: just call primitive and don't update error. + return error, primitive.bind(*invals, **params) + + # Code below handles call- and map-primitives, by recursively calling + # checkify_jaxpr. + err_vals, err_tree = jtu.tree_flatten(error) + num_error_vals = len(err_vals) + if 'donated_invars' in params: + params = dict(params, donated_invars=(*[False]*num_error_vals, + *params['donated_invars'])) + + # call_jaxpr handling + call_jaxpr = params.pop('call_jaxpr') + if isinstance(call_jaxpr, core.ClosedJaxpr): # handle closed_call_p + jaxpr, consts = call_jaxpr.jaxpr, call_jaxpr.consts + else: + jaxpr, consts = call_jaxpr, () + consts_ = tuple(HashableWrapper(c) for c in consts) + partial_checkify = lu.hashable_partial(lu.wrap_init( + checkify_jaxpr_flat_hashable), jaxpr, consts_, enabled_errors, err_tree) + partial_checkify, metadata = _flatten_and_get_error_metadata_thunk( + partial_checkify) + + # map-specific params handling. + if isinstance(primitive, core.MapPrimitive): + # Update `in_axes` and `out_axes_thunk` params for map primitive. + out_val_axes = params.pop('out_axes') + + @as_hashable_function(closure=out_val_axes) + def out_axes_thunk(): + out_err_num = metadata()[0].num_leaves - len(out_val_axes) + return (*(0,)*out_err_num, *out_val_axes) + + params = dict(params, + in_axes=(*(None,)*num_error_vals, *params['in_axes']), + out_axes_thunk=out_axes_thunk) + + all_vals = primitive.bind(partial_checkify, *err_vals, *invals, **params) + + out_tree, _ = metadata() + error, out_vals = tree_unflatten(out_tree, all_vals) + if isinstance(primitive, core.MapPrimitive): + error = _reduce_any_error(error) + return error, out_vals + +def get_shaped_aval(val): + return core.raise_to_shaped(core.get_aval(val)) + +def checkify_jaxpr(jaxpr: core.ClosedJaxpr, enabled_errors, + error: Error, *args) -> tuple[Error, list[core.Value]]: + err_vals, err_tree = jtu.tree_flatten(error) + return checkify_jaxpr_flat(jaxpr.jaxpr, jaxpr.consts, + enabled_errors, err_tree, *err_vals, *args) + +def checkify_jaxpr_flat(jaxpr: core.Jaxpr, consts: Sequence[core.Value], + enabled_errors, err_tree: PyTreeDef, + *args: core.Value) -> tuple[Error, list[Any]]: + env: dict[core.Var, Any] = {} + err_vals, in_args = split_list(args, [err_tree.num_leaves]) + error = jtu.tree_unflatten(err_tree, err_vals) + + last_used = core.last_used(jaxpr) + + def read_env(var: core.Atom): + if isinstance(var, core.Literal): + return var.val + return env[var] + + def write_env(var: core.Var, val: Any): + env[var] = val + + map(write_env, jaxpr.constvars, consts) + map(write_env, jaxpr.invars, in_args) + + # interpreter loop + for eqn in jaxpr.eqns: + invals = map(read_env, eqn.invars) + checkify_rule = error_checks.get( + eqn.primitive, functools.partial(default_checkify_rule, eqn.primitive)) + name_stack = source_info_util.current_name_stack() + eqn.source_info.name_stack + with source_info_util.user_context(eqn.source_info.traceback, + name_stack=name_stack): + error, outvals = checkify_rule(error, enabled_errors, + *invals, **eqn.params) + if eqn.primitive.multiple_results: + map(write_env, eqn.outvars, outvals) + else: + write_env(eqn.outvars[0], outvals) + core.clean_up_dead_vars(eqn, env, last_used) + + return error, map(read_env, jaxpr.outvars) + +def checkify_jaxpr_flat_hashable(jaxpr, hashable_consts, enabled_errors, + err_tree, *args): + consts = tuple(c.x for c in hashable_consts) + return checkify_jaxpr_flat(jaxpr, consts, enabled_errors, err_tree, *args) + +@lu.transformation_with_aux +def flatten_fun_output(*args): + ans = yield args, {} + yield tree_flatten(ans) + + +def _reduce_any_error(error: Error): + out_error = init_error + for error_effect in error._pred.keys(): + errs, codes, payloads = (error._pred[error_effect], + error._code[error_effect], + error._payload[error_effect]) + reduced_idx = jnp.argsort(errs)[-1] + pred, code, payload = tree_map(lambda x, idx=reduced_idx: x[idx], + (errs, codes, payloads)) + out_error = out_error._update(error_effect, pred, code, {}, payload) + + out_error = out_error._replace(_metadata=error._metadata) + return out_error + +## check_p primitive + +check_p = core.Primitive('check') +check_p.multiple_results = True # zero results + + +def _pp_check(eqn, context, settings) -> core.pp.Doc: + annotation = (source_info_util.summarize(eqn.source_info) + if settings.source_info else None) + name_stack_annotation = (f'[{eqn.source_info.name_stack}]' + if settings.name_stack else None) + trimmed_params = sorted((k, v) for (k, v) in eqn.params.items() + if k != "err_tree") + rhs = [core.pp.text(eqn.primitive.name, annotation=name_stack_annotation), + core.pp_kv_pairs(trimmed_params, context, settings), + core.pp.text(" ") + core.pp_vars(eqn.invars, context)] + return core.pp.concat([core.pp.text("", annotation), *rhs]) + +core.pp_eqn_rules[check_p] = _pp_check + +# TODO(lenamartens): inherit from Exception instead of ValueError. +class JaxRuntimeError(ValueError): + pass + +@check_p.def_impl +def check_impl(*args, err_tree, debug): + if debug: + # NOOP (check will only trigger when discharged) + return [] + error = tree_unflatten(err_tree, args) + exc = error.get_exception() + if exc: + filtered_tb = traceback_util.filter_traceback( + exc.traceback_info.as_python_traceback()) + exc.with_traceback(filtered_tb) + raise JaxRuntimeError(str(exc)) from exc + return [] + +@check_p.def_effectful_abstract_eval +def check_abstract_eval(*args, err_tree, debug): + del debug + return [], set(tree_unflatten(err_tree, args)._pred.keys()) + +# TODO(lenamartens) add in-depth error explanation to link to in module docs. +functionalization_error = ValueError( + 'Cannot abstractly evaluate a checkify.check which was not' + ' functionalized. This probably means you tried to stage' + ' (jit/scan/pmap/...) a `check` without functionalizing it' + ' through `checkify.checkify`.' + ) + +def check_lowering_rule(ctx, *args, err_tree, debug): + if debug: + # NOOP (check will only trigger when discharged) + return [] + if not config.xla_runtime_errors.value: + raise functionalization_error + + out_op, _, _ = mlir.emit_python_callback( + ctx, callback=functools.partial(python_err, err_tree), + token=None, + operands=args, + operand_avals=list(ctx.avals_in), + result_avals=list(ctx.avals_out), + has_side_effect=True) + return out_op + +def check_lowering_rule_unsupported(*a, debug, **k): + if debug: + return [] + raise functionalization_error + +def python_err(err_tree, *args): + error = tree_unflatten(err_tree, args) + _check_error(error) + return [] + +mlir.register_lowering(check_p, check_lowering_rule_unsupported, + platform='tpu') +mlir.register_lowering(check_p, check_lowering_rule, + platform='cpu') +mlir.register_lowering(check_p, check_lowering_rule, + platform='gpu') + +def check_batching_rule(batched_args, batch_dims, *, err_tree, debug): + size = next(x.shape[dim] for x, dim in zip(batched_args, batch_dims) + if dim is not batching.not_mapped) + batched_args = (batching.bdim_at_front(a, d, size) + for a, d in zip(batched_args, batch_dims)) + err = tree_unflatten(err_tree, batched_args) + _check_error(err, debug=debug) + return [], [] +batching.primitive_batchers[check_p] = check_batching_rule + +def check_jvp_rule(primals, _, *, err_tree, debug): + # Check primals, discard tangents. + check_p.bind(*primals, err_tree=err_tree, debug=debug) + return [], [] +ad.primitive_jvps[check_p] = check_jvp_rule + +## checkify rules + +ErrorCheckRule = Callable # (Error, FrozenSet[ErrorCategory], *in_vals, **params) -> (Any, Error) +error_checks: dict[core.Primitive, ErrorCheckRule] = {} + + +def get_traceback(): + return source_info_util.current().traceback + +def nan_error_check(prim, error, enabled_errors, *in_vals, **params): + out = prim.bind(*in_vals, **params) + err = check_nans(prim, error, enabled_errors, out) + return err, out + +def check_nans(prim, error, enabled_errors, out): + if NaNError not in enabled_errors: + return error + + def isnan(x): + if jnp.issubdtype(x.dtype, dtypes.prng_key): + return False + return jnp.any(jnp.isnan(x)) + + any_nans = (jnp.any(jnp.array([isnan(x) for x in out])) + if prim.multiple_results else isnan(out)) + return assert_func(error, any_nans, NaNError(get_traceback(), prim.name)) + + +# All primitives which can generate a NaN. +nan_primitives = [lax.acos_p, lax.acosh_p, lax.add_p, lax.asin_p, lax.asinh_p, + lax.atan2_p, lax.atan_p, lax.atanh_p, lax.bessel_i0e_p, + lax.bessel_i1e_p, lax.cbrt_p, lax.conv_general_dilated_p, + lax.cos_p, lax.cosh_p, lax.cumlogsumexp_p, lax.cummax_p, + lax.cummin_p, lax.cumprod_p, lax.cumsum_p, lax.digamma_p, + lax.dot_general_p, lax.erf_inv_p, lax.erf_p, lax.erfc_p, + lax.exp_p, lax.expm1_p, lax.fft_p, lax.igamma_grad_a_p, + lax.igamma_p, lax.igammac_p, lax.integer_pow_p, lax.lgamma_p, + lax.linear_solve_p, lax.log1p_p, lax.log_p, lax.logistic_p, + lax.mul_p, lax.pad_p, lax.pow_p, lax.psum_p, + lax.random_gamma_grad_p, lax.reduce_p, lax.reduce_prod_p, + lax.reduce_sum_p, lax.reduce_window_p, + lax.reduce_window_sum_p, lax.regularized_incomplete_beta_p, + lax.rem_p, lax.rng_uniform_p, lax.rsqrt_p, lax.sin_p, + lax.sinh_p, lax.sqrt_p, lax.sub_p, lax.tan_p, lax.tanh_p] + +for _prim in nan_primitives: + error_checks[_prim] = functools.partial(nan_error_check, _prim) + + +def dynamic_slice_error_check(error, enabled_errors, operand, *start_indices, slice_sizes): + out = lax.dynamic_slice_p.bind(operand, *start_indices, slice_sizes=slice_sizes) + + if OOBError not in enabled_errors: + return error, out + + operand_dims = np.array(operand.shape) + slice_sizes = np.array(slice_sizes) + start_indices = jnp.array(start_indices) + oob_mask = (start_indices < 0) | (start_indices + slice_sizes > operand_dims) + + payload = oob_payload(oob_mask, start_indices, range(operand.ndim), operand.shape) + error = assert_func(error, jnp.any(oob_mask), OOBError(get_traceback(), "dynamic_slice", operand.shape, payload)) + return error, out +error_checks[lax.dynamic_slice_p] = dynamic_slice_error_check + +def dynamic_update_slice_error_check(error, enabled_errors, operand, update, *start_indices): + out = lax.dynamic_update_slice_p.bind(operand, update, *start_indices) + + if OOBError not in enabled_errors: + return error, out + + operand_dims = np.array(operand.shape) + update_dims = np.array(update.shape) + start_indices = jnp.array(start_indices) + oob_mask = (start_indices < 0) | (start_indices + update_dims > operand_dims) + + payload = oob_payload(oob_mask, start_indices, range(operand.ndim), operand.shape) + error = assert_func(error, jnp.any(oob_mask), OOBError(get_traceback(), "dynamic_update_slice", operand.shape, payload)) + return error, out +error_checks[lax.dynamic_update_slice_p] = dynamic_update_slice_error_check + +def gather_error_check(error, enabled_errors, operand, start_indices, *, + dimension_numbers, slice_sizes, unique_indices, + indices_are_sorted, mode, fill_value): + out = lax.gather_p.bind( + operand, start_indices, dimension_numbers=dimension_numbers, + slice_sizes=slice_sizes, unique_indices=unique_indices, + indices_are_sorted=indices_are_sorted, mode=mode, fill_value=fill_value) + + if OOBError not in enabled_errors: + return error, out + + # compare to OOB masking logic in lax._gather_translation_rule + dnums = dimension_numbers + operand_dims = np.array(operand.shape) + num_batch_dims = len(start_indices.shape) - 1 + + upper_bound = operand_dims[np.array(dnums.start_index_map)] + upper_bound -= np.array(slice_sizes)[np.array(dnums.start_index_map)] + upper_bound = jnp.expand_dims(upper_bound, axis=tuple(range(num_batch_dims))) + oob_mask = (start_indices < 0) | (start_indices > upper_bound.astype(start_indices.dtype)) + + payload = oob_payload(oob_mask, start_indices, dnums.start_index_map, operand.shape) + error = assert_func(error, jnp.any(oob_mask), OOBError(get_traceback(), "gather", operand.shape, payload)) + return error, out +error_checks[lax.gather_p] = gather_error_check + +def div_error_check(error, enabled_errors, x, y): + """Checks for division by zero and NaN.""" + if DivisionByZeroError in enabled_errors: + any_zero = jnp.any(jnp.equal(y, 0)) + error = assert_func(error, any_zero, DivisionByZeroError(get_traceback())) + return nan_error_check(lax.div_p, error, enabled_errors, x, y) +error_checks[lax.div_p] = div_error_check + +def oob_payload(oob_mask, indices, dims_map, operand_shape): + # Get first OOB index, axis and axis size so it can be added to the error msg. + flat_idx = jnp.argmin(jnp.logical_not(oob_mask)) + multi_idx = jnp.unravel_index(flat_idx, indices.shape) + oob_axis = jnp.array(dims_map)[multi_idx[-1]] + oob_axis_size = jnp.array(operand_shape)[oob_axis] + oob_index = jnp.ravel(indices)[flat_idx] + payload = jnp.array([oob_index, oob_axis, oob_axis_size], dtype=jnp.int32) + return payload + +def scatter_oob(operand, indices, updates, dnums): + # Ref: see clamping code used in scatter_translation_rule + slice_sizes = [] + pos = 0 + for i in range(len(operand.shape)): + if i in dnums.inserted_window_dims: + slice_sizes.append(1) + else: + slice_sizes.append(updates.shape[dnums.update_window_dims[pos]]) + pos += 1 + + upper_bound = np.array([operand.shape[i] - slice_sizes[i] + for i in dnums.scatter_dims_to_operand_dims], + np.int64) + upper_bound = np.minimum(upper_bound, np.iinfo(indices.dtype).max) + upper_bound = lax.broadcast_in_dim(upper_bound, indices.shape, + (len(indices.shape) - 1,)) + + lower_oob = jnp.less(indices, 0) + upper_oob = jnp.greater(indices, upper_bound.astype(indices.dtype)) + oob_mask = jnp.logical_or(lower_oob, upper_oob) + payload = oob_payload(oob_mask, indices, + dnums.scatter_dims_to_operand_dims, operand.shape) + return jnp.any(oob_mask), payload + +def scatter_error_check(prim, error, enabled_errors, operand, indices, updates, + *, update_jaxpr, update_consts, dimension_numbers, + indices_are_sorted, unique_indices, mode): + """Checks if indices are within bounds and update does not generate NaN.""" + out = prim.bind( + operand, indices, updates, update_jaxpr=update_jaxpr, + update_consts=update_consts, dimension_numbers=dimension_numbers, + indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, + mode=mode) + + if OOBError not in enabled_errors: + return error, out + + out_of_bounds, payload = scatter_oob(operand, indices, updates, dimension_numbers) + oob_error = OOBError(get_traceback(), prim.name, operand.shape, payload) + error = assert_func(error, out_of_bounds, oob_error) + error = check_nans(prim, error, enabled_errors, out) + return error, out +error_checks[lax.scatter_p] = functools.partial(scatter_error_check, lax.scatter_p) +error_checks[lax.scatter_add_p] = functools.partial(scatter_error_check, + lax.scatter_add_p) +error_checks[lax.scatter_mul_p] = functools.partial(scatter_error_check, + lax.scatter_mul_p) +error_checks[lax.scatter_min_p] = functools.partial(scatter_error_check, + lax.scatter_min_p) +error_checks[lax.scatter_max_p] = functools.partial(scatter_error_check, + lax.scatter_max_p) + +# HOP error check rules + +@weakref_lru_cache +def jaxpr_to_checkify_jaxpr( + jaxpr: core.ClosedJaxpr, enabled_errors, err_tree: PyTreeDef, + *flat_err_and_in_vals) -> tuple[core.ClosedJaxpr, PyTreeDef, set[ErrorEffect]]: + checkify_jaxpr_partial = functools.partial(checkify_jaxpr_flat, jaxpr.jaxpr, + jaxpr.consts, enabled_errors, + err_tree) + fun = lu.wrap_init(checkify_jaxpr_partial) + fun, metadata = _flatten_and_get_error_metadata_thunk(fun) + + new_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(fun, flat_err_and_in_vals) + checked_jaxpr = core.ClosedJaxpr(new_jaxpr, consts) + out_tree, error_effects = metadata() + return checked_jaxpr, out_tree, error_effects + +def cond_error_check(error: Error, enabled_errors, index, *ops, branches, linear): + # Get the error-effects out of all branches so the cond can be called with + # a merged error with all these effects. + err_vals, err_tree = jtu.tree_flatten(error) + in_avals = map(get_shaped_aval, [*err_vals, *ops]) + def get_error_effects_from_jaxpr(jxpr): + _, _, effects = jaxpr_to_checkify_jaxpr(jxpr, enabled_errors, err_tree, + *in_avals) + return effects + effects = [get_error_effects_from_jaxpr(jxpr) for jxpr in branches] + merged_error = error._add_placeholder_effects(set().union(*effects)) + err_vals, err_tree = jtu.tree_flatten(merged_error) + new_linear = (*[False] * len(err_vals), *linear) + + # Update branch jaxprs to be checkified jaxprs. + in_avals = map(get_shaped_aval, [*err_vals, *ops]) + new_branches, out_trees, _ = unzip3( + jaxpr_to_checkify_jaxpr( + jxpr, enabled_errors, err_tree, *in_avals) for jxpr in branches) + + err_and_outs = lax.cond_p.bind( + index, *err_vals, *ops, + branches=tuple(new_branches), linear=new_linear) + + # we need to merge metadata across out_trees (a tuple) + err0, out = tree_unflatten(out_trees[0], err_and_outs) + merged_metadata = err0._metadata + for tr in out_trees[1:]: + err, _ = tree_unflatten(tr, err_and_outs) + merged_metadata = {**merged_metadata, **err._metadata} + return err0._replace(_metadata=merged_metadata), out +error_checks[lax.cond_p] = cond_error_check + +def scan_error_check(error, enabled_errors, *in_flat, reverse, length, jaxpr, + num_consts, num_carry, linear, unroll, _split_transpose): + + consts, carry, xs = split_list(in_flat, [num_consts, num_carry]) + xs_mapped = [core.mapped_aval(length, 0, get_shaped_aval(val)) for val in xs] + # Query body effects to create a merged error containing all effects (such + # that in and out carried error are of the same type). + err_vals, err_tree = jtu.tree_flatten(error) + new_in_aval = map(get_shaped_aval, [*err_vals, *consts, *carry]) + xs_mapped + _, _, effects = jaxpr_to_checkify_jaxpr(jaxpr, enabled_errors, + err_tree, *new_in_aval) + + merged_error = error._add_placeholder_effects(effects) + err_vals, err_tree = jtu.tree_flatten(merged_error) + + # Create checked-jaxpr, with the needed pre-processing on the inputs. + new_in_aval = map(get_shaped_aval, [*err_vals, *consts, *carry]) + xs_mapped + checked_jaxpr_, out_tree, _ = jaxpr_to_checkify_jaxpr(jaxpr, enabled_errors, + err_tree, *new_in_aval) + + tomove = ([False] * len(err_vals) + [True] * len(consts) + + [False] * (len(carry) + len(xs))) + checked_jaxpr = pe.move_binders_to_front(checked_jaxpr_, tomove) + new_in_flat = [*consts, *err_vals, *carry, *xs] + new_linear = (*[False] * len(err_vals), *linear) + err_and_out = lax.scan_p.bind( + *new_in_flat, reverse=reverse, length=length, jaxpr=checked_jaxpr, + num_consts=len(consts), num_carry=len(carry)+len(err_vals), + linear=new_linear, unroll=unroll, _split_transpose=_split_transpose) + err, out = tree_unflatten(out_tree, err_and_out) + return err, out + +error_checks[lax.scan_p] = scan_error_check + +def checkify_while_body_jaxpr( + cond_jaxpr: core.ClosedJaxpr, body_jaxpr: core.ClosedJaxpr, + enabled_errors, error: Error, + c_consts_num: int) -> tuple[core.ClosedJaxpr, PyTreeDef, set[ErrorEffect]]: + cond_f = core.jaxpr_as_fun(cond_jaxpr) + body_f = core.jaxpr_as_fun(body_jaxpr) + def new_body_f(*c_consts_and_vals): + c_consts, vals = split_list(c_consts_and_vals, [c_consts_num]) + out = body_f(*vals) + # This checks if the next cond application will error + _ = cond_f(*c_consts, *out) + return out + new_body_f_ = lu.wrap_init(new_body_f) + c_consts_avals = cond_jaxpr.in_avals[:c_consts_num] + jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(new_body_f_, [*c_consts_avals, + *body_jaxpr.in_avals]) + closed_jaxpr = pe.close_jaxpr(jaxpr) + err_vals, err_tree = jtu.tree_flatten(error) + err_vals = map(get_shaped_aval, err_vals) + flat_err_and_in_vals = [*err_vals, *c_consts_avals, *body_jaxpr.in_avals] + jaxpr, out_tree, error_effects = jaxpr_to_checkify_jaxpr( + closed_jaxpr, enabled_errors, err_tree, *flat_err_and_in_vals) + return jaxpr, out_tree, error_effects + + +@weakref_lru_cache +def ignore_error_output_jaxpr(jaxpr, num_error_vals: int): + """Constructs a checked jaxpr which does not output its error value.""" + consts = jaxpr.consts + jaxpr = jaxpr.jaxpr + new_jaxpr = jaxpr.replace(outvars=jaxpr.outvars[num_error_vals:]) + return core.ClosedJaxpr(new_jaxpr, consts) + +def while_loop_error_check(error, enabled_errors, *in_flat, cond_nconsts, + cond_jaxpr, body_nconsts, body_jaxpr): + if cond_jaxpr.out_avals[0].shape: + # TODO(lenamartens, sharadmv): support batched while. + raise ValueError('Checkify does not support batched while-loops ' + '(checkify-of-vmap-of-while). \nHint: if possible, move ' + 'the vmap to the outer level to get ' + 'vmap-of-checkify-of-while.') + + c_consts, b_consts, carry = split_list(in_flat, [cond_nconsts, body_nconsts]) + # Check if the first cond application will error. + error, _ = checkify_jaxpr(cond_jaxpr, enabled_errors, error, *c_consts, *carry) + + _, _, error_effects = checkify_while_body_jaxpr(cond_jaxpr, body_jaxpr, + enabled_errors, error, + cond_nconsts) + # merged error! + error = error._add_placeholder_effects(error_effects) + err_vals, err_tree = jtu.tree_flatten(error) + checked_body_jaxpr_, body_out_tree, _ = checkify_while_body_jaxpr( + cond_jaxpr, body_jaxpr, enabled_errors, error, cond_nconsts) + num_error_vals = len(err_vals) + to_move = ([False] * num_error_vals + [True] * cond_nconsts + + [True] * body_nconsts + [False] * len(carry)) + checked_body_jaxpr = pe.move_binders_to_front(checked_body_jaxpr_, to_move) + + cond_in_flat = [*err_vals, *c_consts, *carry] + cond_in_flat = map(get_shaped_aval, cond_in_flat) + checked_cond_jaxpr, _, _ = jaxpr_to_checkify_jaxpr(cond_jaxpr, enabled_errors, + err_tree, *cond_in_flat) + compat_cond_jaxpr_ = ignore_error_output_jaxpr(checked_cond_jaxpr, num_error_vals) + to_move = [False] * num_error_vals + [True] * cond_nconsts + [False] * len(carry) + compat_cond_jaxpr = pe.move_binders_to_front(compat_cond_jaxpr_, to_move) + + new_in_flat = [*c_consts, *c_consts, *b_consts, *err_vals, *carry] + all_out_vals = lax.while_p.bind( + *new_in_flat, cond_nconsts=cond_nconsts, cond_jaxpr=compat_cond_jaxpr, + body_nconsts=cond_nconsts+body_nconsts, body_jaxpr=checked_body_jaxpr) + # body_out_tree will have all the metadata of cond because it executes a cond! + error, out = tree_unflatten(body_out_tree, all_out_vals) + return error, out +error_checks[lax.while_p] = while_loop_error_check + +def pjit_error_check(error, enabled_errors, *vals_in, jaxpr, + in_shardings, out_shardings, + in_layouts, out_layouts, + resource_env, donated_invars, name, inline, keep_unused): + # jaxpr to checked_jaxpr + err_vals, err_tree = jtu.tree_flatten(error) + new_vals_in = [*err_vals, *vals_in] + in_avals = tuple(map(get_shaped_aval, new_vals_in)) + checked_jaxpr, out_tree, _ = jaxpr_to_checkify_jaxpr(jaxpr, enabled_errors, + err_tree, *in_avals) + + # Update pjit params to account for extra error values. + num_error_vals = len(err_vals) + num_out_error_vals = out_tree.num_leaves - len(out_shardings) + + sharding = sharding_impls.UNSPECIFIED + new_in_shardings = (*[sharding] * num_error_vals, *in_shardings) + new_out_shardings = (*[sharding] * num_out_error_vals, *out_shardings) + new_in_layouts = (*[None] * num_error_vals, *in_layouts) + new_out_layouts = (*[None] * num_out_error_vals, *out_layouts) + new_donated_invars = (*[False] * num_error_vals, *donated_invars) + + err_and_out = pjit.pjit_p.bind( + *new_vals_in, + jaxpr=checked_jaxpr, + in_shardings=new_in_shardings, + out_shardings=new_out_shardings, + in_layouts=new_in_layouts, + out_layouts=new_out_layouts, + resource_env=resource_env, + donated_invars=new_donated_invars, + name=name, + inline=inline, + keep_unused=keep_unused, + ) + return tree_unflatten(out_tree, err_and_out) +error_checks[pjit.pjit_p] = pjit_error_check + +def custom_jvp_call_rule(in_err, enabled_errors, *in_vals, num_consts, + jvp_jaxpr_thunk, call_jaxpr, **params): + # The types to have in mind are: + # jvp : (a -> b) -> (a, T a) -> (b, T b) + # checkify : (a -> b) -> a -> Err b + # jvp-of-checkify : (a -> b) -> (a, T a) -> (Err b, T (Err b)) + # where because Err is a pytree, we necessarily have T (Err b) = Err' (T b) + # where the other Err' components are trivial (of float0 dtype). + # Semantically, we don't add checks to the JVP rule. To check the result of a + # JVP rule, one must instead use checkify-of-jvp. Thus this implementation + # just forwards the input error and code (and trivial tangents) to the output. + err_vals, err_tree = jtu.tree_flatten(in_err) + partial_checkify = lu.wrap_init( + functools.partial(checkify_jaxpr_flat, call_jaxpr.jaxpr, + call_jaxpr.consts, enabled_errors, err_tree)) + partial_checkify, f_metadata = _flatten_and_get_error_metadata_thunk( + partial_checkify) + jvp = lift_jvp(err_tree.num_leaves, num_consts, jvp_jaxpr_thunk) + jvp, jvp_out_tree = flatten_fun_output(jvp) + all_outs = custom_derivatives.custom_jvp_call_p.bind( + partial_checkify, jvp, *err_vals, *in_vals, **params) + fst, out_metadata = lu.merge_linear_aux(f_metadata, jvp_out_tree) + if fst: + err_and_out_tree, _ = out_metadata + out_err, out_vals = tree_unflatten(err_and_out_tree, all_outs) + else: + err_vals, out_vals = split_list(all_outs, [len(err_vals)]) + # forward input error to output + out_err = jtu.tree_unflatten(err_tree, err_vals) + return out_err, out_vals +error_checks[custom_derivatives.custom_jvp_call_p] = custom_jvp_call_rule + +# Compared to custom_derivatives.lift_jvp, we're handling the extra inputs and +# outputs that checkify adds (just forwarding the error data's primal and +# tangent components). The jaxpr in jvp_jaxpr_thunk doesn't expect those. +# TODO(mattjj): can we simplify this, or dedup with custom_derivatives.lift_jvp? +# Adding another layer of lu.transformation was tricky, though maybe doable. +def lift_jvp(num_errs, num_consts, jvp_jaxpr_thunk): + @lu.wrap_init + def jvp(*xs): + n, ragged = divmod(len(xs), 2) + assert not ragged + primals, tangents = xs[num_consts+num_errs:n], xs[n+num_consts+num_errs:] + zeros = [type(t) is SymbolicZero for t in tangents] + jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_thunk(*zeros) + nonzero_tangents = [t for t in tangents if type(t) is not SymbolicZero] + out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents) + out_primals, nz_out_tangents = split_list(out, [len(out_zeros)]) + nz_out_tangents_ = iter(nz_out_tangents) + out_tangents = [SymbolicZero(core.get_aval(p).at_least_vspace()) + if z else next(nz_out_tangents_) + for p, z in zip(out_primals, out_zeros)] + assert next(nz_out_tangents_, None) is None + primal_errs = xs[num_consts:num_consts+num_errs] + tangent_errs = xs[n+num_consts:n+num_consts+num_errs] + return [*primal_errs, *out_primals, *tangent_errs, *out_tangents] + return jvp + +def custom_vjp_call_jaxpr_rule(in_err, enabled_errors, *in_vals, fun_jaxpr, + fwd_jaxpr_thunk, num_consts, bwd, out_trees, + symbolic_zeros): + err_vals, err_tree = jtu.tree_flatten(in_err) + num_errs = err_tree.num_leaves + checkified_fun = lu.wrap_init( + functools.partial(checkify_jaxpr_flat, fun_jaxpr.jaxpr, + fun_jaxpr.consts, enabled_errors, err_tree)) + checkified_fun, fun_metadata = _flatten_and_get_error_metadata_thunk( + checkified_fun) + + @lu.wrap_init + def checkified_fwd(*args): + # TODO(lenamartens, sharadmv): why not checkify here? + xs, zeros = args[::2], args[1::2] + xs, zeros = xs[num_errs:], zeros[num_errs:] + fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk(*zeros) + xs_without_consts = xs[num_consts:] + return core.eval_jaxpr(fwd_jaxpr, fwd_consts, *xs_without_consts) + + bwd_ = lambda *args: (*(None,)*num_errs, *bwd(*args)) + checkified_fwd, fwd_out_tree = flatten_fun_output(checkified_fwd) + all_outs = custom_derivatives.custom_vjp_call_p.bind( + checkified_fun, checkified_fwd, bwd_, *err_vals, *in_vals, out_trees=out_trees, + symbolic_zeros=symbolic_zeros) + fst, out_metadata = lu.merge_linear_aux(fun_metadata, fwd_out_tree) + if fst: + err_and_out_tree, _ = out_metadata + out_err, out_vals = tree_unflatten(err_and_out_tree, all_outs) + else: + out_err, out_vals = in_err, all_outs + return out_err, out_vals +error_checks[custom_derivatives.custom_vjp_call_jaxpr_p] = custom_vjp_call_jaxpr_rule + + +def check_discharge_rule(error, enabled_errors, *args, err_tree, debug): + del debug + new_error = tree_unflatten(err_tree, args) + # Split up new_error into error to be functionalized if it's included in + # enabled_errors (=discharged_error) and an error to be defunctionalized if + # it's not included (=recharged_error) + discharged_error = error + recharged_error = init_error + for error_effect in new_error._pred.keys(): + pred = new_error._pred[error_effect] + code = new_error._code[error_effect] + payload = new_error._payload[error_effect] + if error_effect.error_type in enabled_errors: + discharged_error = update_error(discharged_error, pred, code, {}, payload, + error_effect) + else: + recharged_error = update_error(recharged_error, pred, code, {}, payload, + error_effect) + + discharged_error = discharged_error._replace( + _metadata={**new_error._metadata, **discharged_error._metadata}) + recharged_error = recharged_error._replace(_metadata=new_error._metadata) + # TODO(lenamartens): we actually need to recharge, but this would be a + # breaking API change so leaving for a follow-up. + # check_error(recharged_error) + return discharged_error, [] +error_checks[check_p] = check_discharge_rule + + +## checkify public api + +user_checks = frozenset({FailedCheckError}) +nan_checks = frozenset({NaNError}) +index_checks = frozenset({OOBError}) +div_checks = frozenset({DivisionByZeroError}) +float_checks = nan_checks | div_checks +automatic_checks = float_checks | index_checks +all_checks = automatic_checks | user_checks + + +def checkify(f: Callable[..., Out], + errors: frozenset[ErrorCategory] = user_checks + ) -> Callable[..., tuple[Error, Out]]: + """Functionalize `check` calls in `fun`, and optionally add run-time error checks. + + Run-time errors are either user-added :func:`~check` assertions, or + automatically added checks like NaN checks, depending on the ``errors`` + argument. + + The returned function will return an Error object `err` along with the output + of the original function. ``err.get()`` will either return ``None`` (if no + error occurred) or a string containing an error message. This error message + will correspond to the first error which occurred. ``err.throw()`` will raise + a ValueError with the error message if an error occurred. + + By default only user-added :func:`~check` assertions are enabled. You can + enable automatic checks through the ``errors`` argument. + + The automatic check sets which can be enabled, and when an error is generated: + - ``user_checks``: a :func:`~check` evaluated to False. + - ``nan_checks``: a floating-point operation generated a NaN value + as output. + - ``div_checks``: a division by zero. + - ``index_checks``: an index was out-of-bounds. + + Multiple categories can be enabled together by passing in an error `Set` (eg. + ``errors=nan_checks``). Multiple sets can be re-combined (eg. + ``errors=float_checks|user_checks``) + + Args: + fun: Callable which can contain user checks (see :func:`~check`). + errors: A set of ErrorCategory values which defines the set of enabled + checks. By default only explicit ``checks`` are enabled + (``user_checks``). You can also for example enable NAN and + DIV errors by passing the ``float_checks`` set, or for + example combine multiple sets through set operations + (``float_checks | user_checks``) + Returns: + A function which accepts the same arguments as ``fun`` and returns as output + a pair where the first element is an ``Error`` value, representing the first + failed :func:`~check`, and the second element is the original output of + ``fun``. + + For example: + + >>> import jax + >>> import jax.numpy as jnp + >>> from jax.experimental import checkify + >>> + >>> @jax.jit + ... def f(x): + ... y = jnp.sin(x) + ... return x+y + >>> err, out = checkify.checkify(f, errors=checkify.float_checks)(jnp.inf) + >>> err.throw() # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + jax._src.checkify.JaxRuntimeError: nan generated by primitive: sin + """ + @traceback_util.api_boundary + def checked_fun(*args, **kwargs): + # close over all arguments so they're not turned into abstract values. + in_tree = jtu.tree_structure(((), {})) + closed_f = lambda: f(*args, **kwargs) + # stage: + fun_, out_tree = flatten_fun(lu.wrap_init(closed_f), in_tree) + debug = pe.debug_info(closed_f, in_tree, out_tree, False, 'checkify') + jaxpr_, _, consts, () = pe.trace_to_jaxpr_dynamic(fun_, (), debug) + jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr_)) + # checkify: + error, out_flat = checkify_jaxpr(jaxpr, errors, init_error, *consts) + return error, jtu.tree_unflatten(out_tree(), out_flat) + return checked_fun + +def check(pred: Bool, msg: str, *fmt_args, **fmt_kwargs) -> None: + """Check a predicate, add an error with msg if predicate is False. + + This is an effectful operation, and can't be staged (jitted/scanned/...). + Before staging a function with checks, :func:`~checkify` it! + + Args: + pred: if False, a FailedCheckError error is added. + msg: error message if error is added. Can be a format string. + fmt_args, fmt_kwargs: Positional and keyword formatting arguments for + `msg`, eg.: + ``check(.., "check failed on values {} and {named_arg}", x, named_arg=y)`` + Note that these arguments can be traced values allowing you to add + run-time values to the error message. + Note that tracking these run-time arrays will increase your memory usage, + even if no error happens. + + For example: + + >>> import jax + >>> import jax.numpy as jnp + >>> from jax.experimental import checkify + >>> def f(x): + ... checkify.check(x>0, "{x} needs to be positive!", x=x) + ... return 1/x + >>> checked_f = checkify.checkify(f) + >>> err, out = jax.jit(checked_f)(-3.) + >>> err.throw() # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + jax._src.checkify.JaxRuntimeError: -3. needs to be positive! + + """ + _check(pred, msg, False, *fmt_args, **fmt_kwargs) + +def _check(pred, msg, debug, *fmt_args, **fmt_kwargs): + if not is_scalar_pred(pred): + prim_name = 'debug_check' if debug else 'check' + raise TypeError(f'{prim_name} takes a scalar pred as argument, got {pred}') + for arg in jtu.tree_leaves((fmt_args, fmt_kwargs)): + if not isinstance(arg, (Array, np.ndarray)): + raise TypeError('Formatting arguments to checkify.check need to be ' + 'PyTrees of arrays, but got ' + f'{arg!r} of type {type(arg)}.') + new_error = FailedCheckError(get_traceback(), msg, *fmt_args, **fmt_kwargs) + error = assert_func(init_error, jnp.logical_not(pred), new_error) + _check_error(error, debug=debug) + +def _check_error(error, *, debug=False): + if any(map(np.shape, error._pred.values())): + error = _reduce_any_error(error) + err_args, tree_def = tree_flatten(error) + + return check_p.bind(*err_args, err_tree=tree_def, debug=debug) + + +def is_scalar_pred(pred) -> bool: + return (isinstance(pred, bool) or + isinstance(pred, Array) and pred.shape == () and + pred.dtype == jnp.dtype('bool')) + + +def debug_check(pred: Bool, msg: str, *fmt_args, **fmt_kwargs) -> None: + """Check a predicate when running under checkify, otherwise is a no-op. + + A `debug_check` will only be run if it is transformed by :func:`~checkify`, + otherwise the check will be dropped. + + Args: + pred: if False, a FailedCheckError error is added. + msg: error message if error is added. + fmt_args, fmt_kwargs: Positional and keyword formatting arguments for + `msg`, eg.: + ``debug_check(.., "check failed on values {} and {named}", x, named=y)`` + Note that these arguments can be traced values allowing you to add + run-time values to the error message. + Note that tracking these run-time arrays will increase your memory usage, + even if no error happens. + + For example: + + >>> import jax + >>> import jax.numpy as jnp + >>> from jax.experimental import checkify + >>> def f(x): + ... checkify.debug_check(x!=0, "cannot be zero!") + ... return x + >>> _ = f(0) # running without checkify means no debug_check is run. + >>> checked_f = checkify.checkify(f) + >>> err, out = jax.jit(checked_f)(0) # running with checkify runs debug_check. + >>> err.throw() # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + jax._src.checkify.JaxRuntimeError: cannot be zero! + + """ + _check(pred, msg, True, *fmt_args, **fmt_kwargs) + + +def check_error(error: Error) -> None: + """Raise an Exception if ``error`` represents a failure. Functionalized by :func:`~checkify`. + + The semantics of this function are equivalent to: + + >>> def check_error(err: Error) -> None: + ... err.throw() # can raise ValueError + + But unlike that implementation, ``check_error`` can be functionalized using + the :func:`~checkify` transformation. + + This function is similar to :func:`~check` but with a different signature: whereas + :func:`~check` takes as arguments a boolean predicate and a new error message + string, this function takes an ``Error`` value as argument. Both :func:`~check` + and this function raise a Python Exception on failure (a side-effect), and + thus cannot be staged out by :func:`~jax.jit`, :func:`~jax.pmap`, + :func:`~jax.lax.scan`, etc. Both also can + be functionalized by using :func:`~checkify`. + + But unlike :func:`~check`, this function is like a direct inverse of + :func:`~checkify`: + whereas :func:`~checkify` takes as input a function which + can raise a Python + Exception and produces a new function without that effect but which produces + an ``Error`` value as output, this ``check_error`` function can accept an + ``Error`` value as input and can produce the side-effect of raising an + Exception. That is, while :func:`~checkify` goes from + functionalizable Exception + effect to error value, this ``check_error`` goes from error value to + functionalizable Exception effect. + + ``check_error`` is useful when you want to turn checks represented by an + ``Error`` value (produced by functionalizing ``checks`` via + :func:`~checkify`) back into Python Exceptions. + + Args: + error: Error to check. + + For example, you might want to functionalize part of your program through + checkify, stage out your functionalized code through :func:`~jax.jit`, then + re-inject your error value outside of the :func:`~jax.jit`: + + >>> import jax + >>> from jax.experimental import checkify + >>> def f(x): + ... checkify.check(x>0, "must be positive!") + ... return x + >>> def with_inner_jit(x): + ... checked_f = checkify.checkify(f) + ... # a checkified function can be jitted + ... error, out = jax.jit(checked_f)(x) + ... checkify.check_error(error) + ... return out + >>> _ = with_inner_jit(1) # no failed check + >>> with_inner_jit(-1) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + jax._src.JaxRuntimeError: must be positive! + >>> # can re-checkify + >>> error, _ = checkify.checkify(with_inner_jit)(-1) + """ + if not isinstance(error, Error): + raise TypeError('check_error takes an Error as argument, ' + f'got type {type(error)} instead.') + _check_error(error, debug=False) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cloud_tpu_init.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cloud_tpu_init.py new file mode 100644 index 0000000000000000000000000000000000000000..71827d8f8a9f4ee6a4a341700ea6e68de14d1439 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cloud_tpu_init.py @@ -0,0 +1,82 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from jax import version +from jax._src import config +from jax._src import hardware_utils + +running_in_cloud_tpu_vm: bool = False + + +def maybe_import_libtpu(): + try: + # pylint: disable=import-outside-toplevel + # pytype: disable=import-error + import libtpu + + # pytype: enable=import-error + # pylint: enable=import-outside-toplevel + except ImportError: + return None + else: + return libtpu + + +def jax_force_tpu_init() -> bool: + return 'JAX_FORCE_TPU_INIT' in os.environ + + +def cloud_tpu_init() -> None: + """Automatically sets Cloud TPU topology and other env vars. + + **This must be called before the TPU runtime is loaded, which happens as soon + as JAX's C++ backend is loaded! I.e. call this before xla_bridge or xla_client + is imported.** + + Safe to call in non-Cloud TPU environments. + + Some of these environment variables are used to tell the TPU runtime what kind + of mesh topology to use. It assumes a single-host topology by default, so we + manually set them here to default to the full pod slice if applicable. + + This will not set any env vars if a single topology-related env var is already + set. + """ + global running_in_cloud_tpu_vm + + # Exit early if we're not running on a Cloud TPU VM or libtpu isn't installed. + libtpu_module = maybe_import_libtpu() + num_tpu_chips = hardware_utils.num_available_tpu_chips_and_device_id()[0] + if (libtpu_module is None or num_tpu_chips == 0) and not jax_force_tpu_init(): + return + + running_in_cloud_tpu_vm = True + + os.environ.setdefault('GRPC_VERBOSITY', 'ERROR') + os.environ.setdefault('JAX_PLATFORMS', 'tpu,cpu') + os.environ['TPU_ML_PLATFORM'] = 'JAX' + os.environ['TPU_ML_PLATFORM_VERSION'] = version.__version__ + if hardware_utils.tpu_enhanced_barrier_supported(): + os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS","") + " --xla_tpu_use_enhanced_launch_barrier=true" + + # this makes tensorstore serialization work better on TPU + os.environ.setdefault('TENSORSTORE_CURL_LOW_SPEED_TIME_SECONDS', '60') + os.environ.setdefault('TENSORSTORE_CURL_LOW_SPEED_LIMIT_BYTES', '256') + + if config.jax_pjrt_client_create_options.value is None: + config.update( + 'jax_pjrt_client_create_options', + f'ml_framework_name:JAX;ml_framework_version:{version.__version__}' + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3cc76a2b9d1b6b2f587fd66a59e37abddf4b8f93 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .cluster import ClusterEnv + +# Order of declaration of the cluster environments +# will dictate the order in which they will be checked. +# Therefore, if multiple environments are available and +# the user did not explicitly provide the arguments +# to :func:`jax.distributed.initialize`, the first +# available one from the list will be picked. +from .ompi_cluster import OmpiCluster +from .slurm_cluster import SlurmCluster +from .cloud_tpu_cluster import GkeTpuCluster +from .cloud_tpu_cluster import GceTpuCluster diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/ompi_cluster.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/ompi_cluster.py new file mode 100644 index 0000000000000000000000000000000000000000..9249958e28858acd8d005e7a75c2fc27778ac8b5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/ompi_cluster.py @@ -0,0 +1,60 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import os +import re +from jax._src import clusters + +# OMPI_MCA_orte_hnp_uri exists only when processes are launched via mpirun or mpiexec +_ORTE_URI = 'OMPI_MCA_orte_hnp_uri' +_PROCESS_COUNT = 'OMPI_COMM_WORLD_SIZE' +_PROCESS_ID = 'OMPI_COMM_WORLD_RANK' +_LOCAL_PROCESS_ID = 'OMPI_COMM_WORLD_LOCAL_RANK' + +class OmpiCluster(clusters.ClusterEnv): + @classmethod + def is_env_present(cls) -> bool: + return _ORTE_URI in os.environ + + @classmethod + def get_coordinator_address(cls, timeout_secs: int | None) -> str: + # Examples of orte_uri: + # 1531576320.0;tcp://10.96.0.1,10.148.0.1,10.108.0.1:34911 + # 1314521088.0;tcp6://[fe80::b9b:ac5d:9cf0:b858,2620:10d:c083:150e::3000:2]:43370 + orte_uri = os.environ[_ORTE_URI] + job_id_str = orte_uri.split('.', maxsplit=1)[0] + # The jobid is always a multiple of 2^12, let's divide it by 2^12 + # to reduce likelihood of port conflict between jobs + job_id = int(job_id_str) // 2**12 + # Pick port in ephemeral range [(65535 - 2^12 + 1), 65535] + port = job_id % 2**12 + (65535 - 2**12 + 1) + launcher_ip_match = re.search(r"tcp://(.+?)[,:]|tcp6://\[(.+?)[,\]]", orte_uri) + if launcher_ip_match is None: + raise RuntimeError('Could not parse coordinator IP address from Open MPI environment.') + launcher_ip = next(i for i in launcher_ip_match.groups() if i is not None) + return f'{launcher_ip}:{port}' + + @classmethod + def get_process_count(cls) -> int: + return int(os.environ[_PROCESS_COUNT]) + + @classmethod + def get_process_id(cls) -> int: + return int(os.environ[_PROCESS_ID]) + + @classmethod + def get_local_process_id(cls) -> int | None: + return int(os.environ[_LOCAL_PROCESS_ID]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/slurm_cluster.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/slurm_cluster.py new file mode 100644 index 0000000000000000000000000000000000000000..933036a64eab71f5b1cd31ffb6cd9e3471d84ae7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/clusters/slurm_cluster.py @@ -0,0 +1,63 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import os +from jax._src import clusters + +_JOBID_PARAM = 'SLURM_JOB_ID' +_NODE_LIST = 'SLURM_STEP_NODELIST' +_PROCESS_COUNT = 'SLURM_NTASKS' +_PROCESS_ID = 'SLURM_PROCID' +_LOCAL_PROCESS_ID = 'SLURM_LOCALID' +_NUM_NODES = 'SLURM_STEP_NUM_NODES' + +class SlurmCluster(clusters.ClusterEnv): + @classmethod + def is_env_present(cls) -> bool: + return _JOBID_PARAM in os.environ + + @classmethod + def get_coordinator_address(cls, timeout_secs: int | None) -> str: + # Pick port in ephemeral range [(65535 - 2^12 + 1), 65535] + port = int(os.environ[_JOBID_PARAM]) % 2**12 + (65535 - 2**12 + 1) + + # Parse the first hostname of the job + # If we are looking for 'node001', + # node_list potential formats are 'node001', 'node001,host2', + # 'node[001-0015],host2', and 'node[001,007-015],host2'. + node_list = os.environ[_NODE_LIST] + delims = {',', '['} + ind = next((i for i, ch in enumerate(node_list) if ch in delims), len(node_list)) + if ind == len(node_list) or node_list[ind] == ',': # Formats: 'node001' or 'node001,host2' + return f'{node_list[:ind]}:{port}' + else: # Formats: 'node[001-0015],host2' or 'node[001,007-015],host2' + prefix = node_list[:ind] + suffix = node_list[ind+1:] + delims2 = {',', '-'} + ind2 = next((i for i, ch in enumerate(suffix) if ch in delims2), None) + return f'{prefix}{suffix[:ind2]}:{port}' + + @classmethod + def get_process_count(cls) -> int: + return int(os.environ[_PROCESS_COUNT]) + + @classmethod + def get_process_id(cls) -> int: + return int(os.environ[_PROCESS_ID]) + + @classmethod + def get_local_process_id(cls) -> int | None: + return int(os.environ[_LOCAL_PROCESS_ID]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compilation_cache.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compilation_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..7884cda7c13470b88d56c603bb4787cef765a213 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compilation_cache.py @@ -0,0 +1,296 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +import threading +import warnings +import zlib + +import numpy as np + +# If zstandard is installed, we use zstd compression, otherwise we use zlib. +try: + import zstandard +except ImportError: + zstandard = None + +from jax._src import cache_key +from jax._src.compilation_cache_interface import CacheInterface +from jax._src import config +from jax._src import monitoring +from jax._src.gfile_cache import GFileCache +from jax._src.lib import xla_client +from jax._src.lib.mlir import ir + + +logger = logging.getLogger(__name__) + +_cache: CacheInterface | None = None + +_cache_initialized: bool = False + +_cache_used: bool = False + +# Mutex to protect _cache_initialized and _cache_used. +_cache_initialized_mutex = threading.Lock() + +_UNSUPPORTED_RUNTIMES: set[str] = set() + +def set_once_cache_used(f) -> None: + """One-time setting of _cache_used. + + If _cache_used is False, set it to True and execute the provided function + f. No action if _cache_used is True. This provides a mechanism to execute f + once per task. Note that reset_cache() will reset _cache_used also. + """ + global _cache_used + with _cache_initialized_mutex: + if not _cache_used: + _cache_used = True + if f is not None: + f() + + +def get_file_cache(path: str) -> tuple[CacheInterface, str] | None: + """Returns the file cache and the path to the cache.""" + return GFileCache(path), path + + +def set_cache_dir(path) -> None: + """ + Sets the persistent compilation cache directory. + + After calling this, jit-compiled functions are saved to `path`, so they + do not need be recompiled if the process is restarted or otherwise run again. + This also tells Jax where to look for compiled functions before compiling. + """ + config.config.update("jax_compilation_cache_dir", path) + + +def initialize_cache(path) -> None: + """ + This API is deprecated; use set_cache_dir instead. + + Set the path. To take effect, should be called prior to any calls to + get_executable_and_time() and put_executable_and_time(). + """ + warnings.warn("initialize_cache is deprecated; use set_cache_dir instead", + DeprecationWarning, stacklevel=2) + config.config.update("jax_compilation_cache_dir", path) + + +def default_min_cache_entry_size() -> int: + """Returns the minimum size below which the entry should not be cached.""" + return 0 + + +def _is_cache_enabled() -> bool: + return config.enable_compilation_cache.value + + +def _initialize_cache() -> None: + # Attempt to initialize the cache at most once. + global _cache_initialized + with _cache_initialized_mutex: + if _cache_initialized: + return + _cache_initialized = True + + # Nothing to do if the cache is disabled. + if not _is_cache_enabled(): + logger.debug("_initialize_cache: cache is disabled!") + return + + # Set the minimum cache size entry only if the flag + # --jax_persistent_cache_min_entry_size_bytes has not been set. + if config.persistent_cache_min_entry_size_bytes.value == 0: + config.config.update("jax_persistent_cache_min_entry_size_bytes", + default_min_cache_entry_size()) + + global _cache + assert _cache is None, "The cache has already been initialized!" + path: str | None = config.compilation_cache_dir.value + # If the path is not set, the cache will not be enabled. + if not path: + return + + cache_and_path = get_file_cache(path) + if cache_and_path is None: + logger.debug("_initialize_cache: cache initialization failed!") + else: + _cache, path = cache_and_path + logger.debug("Initialized persistent compilation cache at %s", path) + + +def _get_cache(backend) -> CacheInterface | None: + # TODO(b/289098047): consider making this an API and changing the callers of + # get_executable_and_time() and put_executable_and_time() to call get_cache() + # and passing the result to them. + if backend.runtime_type in _UNSUPPORTED_RUNTIMES: + logger.debug("_get_cache: Unsupported runtime: %s", backend.runtime_type) + return None + if _cache is None: + _initialize_cache() # initialization is done at most once; see above + return _cache + + +def compress_executable(executable): + if zstandard: + compressor = zstandard.ZstdCompressor() + return compressor.compress(executable) + else: + return zlib.compress(executable) + +def decompress_executable(executable): + if zstandard: + decompressor = zstandard.ZstdDecompressor() + return decompressor.decompress(executable) + else: + return zlib.decompress(executable) + + +def is_executable_in_cache(backend, cache_key: str) -> bool: + """Checks if the executable is in the cache.""" + cache = _get_cache(backend) + if cache is None: + return False + + # TODO(patrios): add check cache key method to cache interface. + executable_and_time = cache.get(cache_key) + return executable_and_time is not None + + +def get_executable_and_time( + cache_key: str, compile_options, backend +) -> tuple[xla_client.LoadedExecutable | None, int | None]: + """Returns the cached executable and its compilation time if present, or None + otherwise. + """ + cache = _get_cache(backend) + if cache is None: + logger.debug("get_executable_and_time: cache is disabled/not initialized") + return None, None + executable_and_time = cache.get(cache_key) + if not executable_and_time: + return None, None + + executable_and_time = decompress_executable(executable_and_time) + serialized_executable, compile_time = extract_executable_and_time( + executable_and_time) + xla_executable_deserialized = backend.deserialize_executable( + serialized_executable, compile_options) + return xla_executable_deserialized, compile_time + + +def put_executable_and_time( + cache_key: str, + module_name: str, + executable: xla_client.LoadedExecutable, + backend, + compile_time: int +) -> None: + """Adds the 'executable' and its compilation time to the cache, possibly + evicting older entries. + """ + cache = _get_cache(backend) + if cache is None: + logger.debug("put_executable_and_time: cache is disabled/not initialized") + return + + serialized_executable = backend.serialize_executable(executable) + executable_and_time = combine_executable_and_time( + serialized_executable, compile_time) + executable_and_time = compress_executable(executable_and_time) + + min_entry_size = config.persistent_cache_min_entry_size_bytes.value + entry_size = len(executable_and_time) + if entry_size < min_entry_size: + logger.info( + "Not writing cache entry with key %s since its size (%d bytes) " + "is less than threshold (%d bytes)", + cache_key, + entry_size, + min_entry_size, + ) + else: + logger.info( + "Writing %s to persistent compilation cache with key %s.", + module_name, + cache_key + ) + monitoring.record_event('/jax/compilation_cache/cache_misses') + cache.put(cache_key, executable_and_time) + + +def get_cache_key(module: ir.Module, devices: np.ndarray, compile_options, + backend) -> str: + return cache_key.get(module, devices, compile_options, backend, + "zstandard" if zstandard is not None else "zlib") + + +def is_initialized() -> bool: + """ + Deprecated. + + Return whether the cache is enabled. Initialization can be deferred, so + initialized status is not checked. The name is retained for backwards + compatibility. + """ + warnings.warn("is_initialized is deprecated; do not use", + DeprecationWarning, stacklevel=2) + return _is_cache_enabled() + + +def reset_cache() -> None: + """Get back to pristine, uninitialized state.""" + global _cache + global _cache_initialized + global _cache_used + logger.info("Resetting cache at %s.", + _cache._path if _cache is not None else "") + _cache = None + with _cache_initialized_mutex: + _cache_initialized = False + _cache_used = False + + +def combine_executable_and_time( + serialized_executable: bytes, compile_time: int +) -> bytes: + """Given the serialized executable and the compilation time, produce a cache + entry in the format shown below. + + The cache entry is of the form: + Byte: 0 1 2 3 4 ... + Content: compilation time serialized executable + (big-endian int) + """ + return int(compile_time).to_bytes(4, byteorder='big') + serialized_executable + + +def extract_executable_and_time( + exectuable_and_time: bytes +) -> tuple[bytes, int]: + """Given the cache entry in the format shown below, extract the serialized + executable and the compilation time. + + The cache entry 'executable_and_time' is of the form: + Byte: 0 1 2 3 4 ... + Content: compilation time serialized executable + (big-endian int) + """ + return exectuable_and_time[4:], int.from_bytes( + exectuable_and_time[:4], byteorder='big') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compilation_cache_interface.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compilation_cache_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..95d557c5531eade068171704f12facc328497a89 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compilation_cache_interface.py @@ -0,0 +1,30 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import abstractmethod + +from jax._src import path as pathlib +from jax._src import util + + +class CacheInterface(util.StrictABC): + _path: pathlib.Path + + @abstractmethod + def get(self, key: str): + pass + + @abstractmethod + def put(self, key: str, value: bytes): + pass diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compiler.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compiler.py new file mode 100644 index 0000000000000000000000000000000000000000..7bc1fd66fde0e59b0d387b2359f4630f7e9f826d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compiler.py @@ -0,0 +1,690 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Interface to the compiler + +from __future__ import annotations + +from collections.abc import Sequence +import logging +import os +import tempfile +import time +from typing import Any, Optional +import warnings + +from jax._src import compilation_cache +from jax._src import config as config +from jax._src import distributed +from jax._src import lib +from jax._src import monitoring +from jax._src import profiler +from jax._src import traceback_util +from jax._src.interpreters import mlir +from jax._src.lib import xla_client as xc +from jax._src.lib.mlir import ir +import numpy as np + + +_DISABLE_MOST_OPTIMIZATIONS = config.DEFINE_bool( + 'jax_disable_most_optimizations', + config.bool_env('JAX_DISABLE_MOST_OPTIMIZATIONS', False), + 'Try not to do much optimization work. This can be useful if the cost of ' + 'optimization is greater than that of running a less-optimized program.') + +_COMPILER_DETAILED_LOGGING_MIN_OPS = config.DEFINE_integer( + "jax_compiler_detailed_logging_min_ops", + config.int_env("JAX_COMPILER_DETAILED_LOGGING_MIN_OPS", 10), + help=( + 'How big should a module be in MLIR operations before JAX enables ' + 'detailed compiler logging? The intent of this flag is to suppress ' + 'detailed logging for small/uninteresting computations.' + ), +) + +# The special XLA-AutoFDO profile version that indicates that a profile is not +# available and retrieval should not be attempted. +_NO_PROFILE_DONT_RETRIEVE = -1 + +traceback_util.register_exclusion(__file__) + +CompileOptions = xc.CompileOptions + +logger = logging.getLogger(__name__) + + +# Will be monkeypatched with the function that gets the XLA-AutoFDO profile +# version. The default (-1) takes care of errors. +# TODO(b/289098047): consider refactoring this interface. +def get_latest_profile_version(backend: xc.Client) -> int: + del backend + return -1 + + +def _walk_operations(op, k): + k -= 1 + if k < 0: + return k + for region in op.regions: + for block in region: + for child_op in block: + k = _walk_operations(child_op, k) + if k < 0: + return k + return k + + +def use_detailed_logging(module: ir.Module) -> bool: + """Returns 'true' if detailed logging should be enabled for 'module'.""" + bound = _COMPILER_DETAILED_LOGGING_MIN_OPS.value + return _walk_operations(module.operation, bound) < 0 + + +def get_compile_options( + num_replicas: int, + num_partitions: int, + device_assignment=None, + use_spmd_partitioning: bool = True, + use_auto_spmd_partitioning: bool = False, + auto_spmd_partitioning_mesh_shape: list[int] | None = None, + auto_spmd_partitioning_mesh_ids: list[int] | None = None, + env_options_overrides: dict[str, str] | None = None, + fdo_profile: bytes | None = None, + detailed_logging: bool = True, + backend: xc.Client | None = None, +) -> xc.CompileOptions: + """Returns the compile options to use, as derived from flag values. + + Args: + num_replicas: Number of replicas for which to compile. + num_partitions: Number of partitions for which to compile. + device_assignment: Optional ndarray of jax devices indicating the assignment + of logical replicas to physical devices (default inherited from + xla_client.CompileOptions). Must be consistent with `num_replicas` and + `num_partitions`. + use_spmd_partitioning: boolean indicating whether to enable SPMD or MPMD + partitioning in XLA. + use_auto_spmd_partitioning: boolean indicating whether to automatically + generate XLA shardings for SPMD partitioner. + auto_spmd_partitioning_mesh_shape: device mesh shape used to create + auto_spmd_partitioning search space. + auto_spmd_partitioning_mesh_ids: device ids used to create + auto_spmd_partitioning search space. + env_options_overrides: dict of additional options parsed by the compiler + fdo_profile: Optional profile for feedback-directed optimization passed to + XLA. + detailed_logging: Is this an "interesting" computation about which XLA + would be wise to log compilation information? + backend: the client, if available. + """ + compile_options = xc.CompileOptions() + compile_options.num_replicas = num_replicas + compile_options.num_partitions = num_partitions + build_options = compile_options.executable_build_options + build_options.use_spmd_partitioning = use_spmd_partitioning + build_options.use_auto_spmd_partitioning = use_auto_spmd_partitioning + if fdo_profile is not None: + build_options.fdo_profile = fdo_profile + if use_auto_spmd_partitioning: + build_options.auto_spmd_partitioning_mesh_shape = auto_spmd_partitioning_mesh_shape or [] + build_options.auto_spmd_partitioning_mesh_ids = auto_spmd_partitioning_mesh_ids or [] + if device_assignment is not None: + logger.debug( + 'get_compile_options: num_replicas=%s num_partitions=%s device_assignment=%s', + num_replicas, num_partitions, device_assignment) + device_assignment = np.array(device_assignment) + + # Allow 1D device assignment if num_partitions is 1. + if (device_assignment.ndim == 1) and (num_partitions == 1): + device_assignment = device_assignment[:, None] + + if num_replicas != device_assignment.shape[0]: + msg = 'device_assignment does not match num_replicas: {} vs {}.' + raise ValueError(msg.format(device_assignment, num_replicas)) + + if num_partitions != device_assignment.shape[1]: + msg = 'device_assignment does not match num_partitions: {} vs {}.' + raise ValueError(msg.format(device_assignment, num_partitions)) + + if device_assignment.dtype == object: + device_assignment = np.vectorize(lambda d: d.id, otypes=[int])( + device_assignment) + device_assignment = xc.DeviceAssignment.create(device_assignment) + assert device_assignment.replica_count() == num_replicas + assert device_assignment.computation_count() == num_partitions + compile_options.device_assignment = device_assignment + + if env_options_overrides is not None: + compile_options.env_option_overrides = list(env_options_overrides.items()) + + debug_options = compile_options.executable_build_options.debug_options + if lib.cuda_path is not None: + debug_options.xla_gpu_cuda_data_dir = lib.cuda_path + + if _DISABLE_MOST_OPTIMIZATIONS.value: + debug_options.xla_backend_optimization_level = 0 + debug_options.xla_llvm_disable_expensive_passes = True + debug_options.xla_test_all_input_layouts = False + + # XLA-AutoFDO profile version: precedence order is: + # 1. Whatever --jax_xla_profile_version is set to. + # 2. If --jax_xla_profile_version is not set (i.e., 0), call the function + # set in get_latest_profile_version and use the return value if non-zero. + # If the function returns 0, set -1; this is an error. + # -1 indicates that no attempt should be made to retrieve the latest profile + # later on. + jax_xla_profile_version = config.jax_xla_profile_version.value + if jax_xla_profile_version > 0: + compile_options.profile_version = jax_xla_profile_version + logger.debug("get_compile_options XLA-AutoFDO profile: " + + "using JAX XLA profile version %d from flag", + jax_xla_profile_version) + else: + compile_options.profile_version = _NO_PROFILE_DONT_RETRIEVE + if backend is None: + logging.info("get_compile_options: no backend supplied; " + "disabling XLA-AutoFDO profile") + else: + fdo_profile_version = get_latest_profile_version(backend) + if fdo_profile_version != 0: + compile_options.profile_version = fdo_profile_version + logger.debug("get_compile_options XLA-AutoFDO profile: " + + "using XLA-AutoFDO profile version %d", + fdo_profile_version) + else: + logger.error("get_compile_options XLA-AutoFDO profile: " + + "XLA-AutoFDO profile version is 0; this should not happen") + + debug_options.xla_detailed_logging = detailed_logging + + return compile_options + +@profiler.annotate_function +def backend_compile( + backend: xc.Client, + module: ir.Module, + options: xc.CompileOptions, + host_callbacks: Sequence[Any], +) -> xc.LoadedExecutable: + # Convert ir.Module to a string representation, unless the backend + # explicitly flags the ability to handle a module directly (avoiding the + # overhead of back and forth conversions). + # TODO(slebedev): Change the backend.compile() to accept ir.Module. + built_c: Any + if getattr(backend, "needs_str_ir", True): + built_c = mlir.module_to_bytecode(module) + else: + built_c = module + + # we use a separate function call to ensure that XLA compilation appears + # separately in Python profiling results + if host_callbacks: + return backend.compile(built_c, compile_options=options, + host_callbacks=host_callbacks) + # Some backends don't have `host_callbacks` option yet + # TODO(sharadmv): remove this fallback when all backends allow `compile` + # to take in `host_callbacks` + return backend.compile(built_c, compile_options=options) + +def compile_or_get_cached( + backend: xc.Client, + computation: ir.Module, + devices: np.ndarray, + compile_options: xc.CompileOptions, + host_callbacks: Sequence[Any], + pgle_profiler: profiler.PGLEProfiler | None = None, +) -> xc.LoadedExecutable: + sym_name = computation.operation.attributes['sym_name'] + module_name = ir.StringAttr(sym_name).value + + if dumped_to := mlir.dump_module_to_file(computation, "compile"): + logging.info("Dumped the module to %s.", dumped_to) + + # Persistent compilation cache only implemented on TPU and GPU and the backend + # that supports serialization of executables. + # TODO(skye): add warning when initializing cache on unsupported default platform + supported_platforms = ["tpu", "gpu", "cpu"] + use_compilation_cache = ( + config.enable_compilation_cache.value + and getattr(backend, "supports_executable_serialization", True) + and backend.platform in supported_platforms + ) + + if not use_compilation_cache: + return backend_compile(backend, computation, compile_options, + host_callbacks) + + compilation_cache.set_once_cache_used( + lambda: monitoring.record_event( + "/jax/compilation_cache/tasks_using_cache")) + monitoring.record_event('/jax/compilation_cache/compile_requests_use_cache') + + try: + cache_key = compilation_cache.get_cache_key( + computation, devices, compile_options, backend) + except xc._xla.XlaRuntimeError as ex: + logger.error("compile_or_get_cached: unable to generate cache key, " + "skipping the cache: %s", ex) + return backend_compile(backend, computation, compile_options, + host_callbacks) + + is_multi_process = ( + len({device.process_index for device in devices.flatten()}) > 1) + min_device_process_id = ( + min(devices.flatten(), key=lambda device: device.id).process_index) + + # When PGLE is enabled there might be 3 types of situations: + # 1. PGLE profiled module (the one which was recompiled with FDO profile) is + # in the persistent cache. In this case the module should be returned from + # cache and PGLE should be disabled for this module. Is module is stored in + # the persistent cache under the "pgle_profiled_module_key" which calculated + # with replacing FDO profile with flag which identify that module were PGLE + # profiled. + # 2. PGLE profiled module is not in the persistent cache and the module is + # getting built with an FDO profile. In this case we need to share FDO profile + # with other processes and store the result under the + # "pgle_profiled_module_key" so later in case 1 we will be able to find the + # module. + # 3. PGLE profiled module is not in the persistent cache and the module is + # getting compiled to be PGLEd (FDO profile is empty). In this case we need to + # simply return the non-PGLE profiled module from the persistent cache. + if (config.enable_pgle.value + and config.pgle_profiling_runs.value > 0): + fdo_profile = compile_options.executable_build_options.fdo_profile + compile_options.executable_build_options.fdo_profile = b"pgle profiled" + + pgle_profiled_module_key = compilation_cache.get_cache_key( + computation, devices, compile_options, backend) + compile_options.executable_build_options.fdo_profile = fdo_profile + + if _is_executable_in_cache(backend, pgle_profiled_module_key): + # Load PGLE profiled module from the persistent cache. + cache_key = pgle_profiled_module_key + if pgle_profiler is not None: + pgle_profiler.disable() + elif fdo_profile is not None and len(fdo_profile) > 0: + # Store module under PGLE profiled module cache key. + cache_key = pgle_profiled_module_key + if is_multi_process and distributed.global_state.client is not None: + compile_options.executable_build_options.fdo_profile = _share_fdo_profiles( + computation, devices, compile_options, backend, + distributed.global_state.client, + min_device_process_id + ) + + cache_retrieval_start = time.monotonic() + retrieved_executable, retrieved_compile_time = _cache_read( + module_name, cache_key, compile_options, backend) + cache_retrieval_time = time.monotonic() - cache_retrieval_start + + if retrieved_executable is not None: + assert retrieved_compile_time is not None + logger.debug("Persistent compilation cache hit for '%s'", module_name) + + monitoring.record_event('/jax/compilation_cache/cache_hits') + monitoring.record_event_duration_secs( + '/jax/compilation_cache/compile_time_saved_sec', + retrieved_compile_time - cache_retrieval_time) + + monitoring.record_event_duration_secs( + "/jax/compilation_cache/cache_retrieval_time_sec", cache_retrieval_time) + + return retrieved_executable + elif ( + config.share_binary_between_hosts.value + and is_multi_process + and distributed.global_state.client is not None + # Host callbacks are currently baked into the HLO module so we cant share + # them. + and len(host_callbacks) == 0 + ): + return _compile_and_share_module( + backend, + computation, + compile_options, + host_callbacks, + distributed.global_state.client, + module_name, + cache_key, + min_device_process_id + ) + elif ( + config.share_autotune_config_between_hosts.value + and is_multi_process + and distributed.global_state.client is not None + ): + return _compile_and_write_autotune_config( + backend, + computation, + compile_options, + host_callbacks, + distributed.global_state.client, + module_name, + cache_key, + min_device_process_id + ) + else: + return _compile_and_write_cache( + backend, + computation, + compile_options, + host_callbacks, + module_name, + cache_key, + ) + +# The process that has the lowest device ID should share FDO profile before +# compilation with other processes. +def _share_fdo_profiles( + computation: ir.Module, + devices: np.ndarray, + compile_options: xc.CompileOptions, + backend: xc.Client, + global_client: lib.xla_extension.DistributedRuntimeClient, + min_process_id +) -> Optional[bytes]: + sym_name = computation.operation.attributes['sym_name'] + module_name = ir.StringAttr(sym_name).value + fdo_profile = compile_options.executable_build_options.fdo_profile + if fdo_profile is None or len(fdo_profile) == 0: + return fdo_profile + + compile_options.executable_build_options.fdo_profile = b"" + profile_key = ( + compilation_cache.get_cache_key( + computation, devices, compile_options, backend + ) + + "_fdo_sync" + ) + if profile_key in _share_fdo_profiles.modules_profiles: + return _share_fdo_profiles.modules_profiles[profile_key] + + share_timeout = config.share_binary_between_hosts_timeout_ms.value + if distributed.global_state.process_id == min_process_id: + logger.debug( + "Sharing FDO profile: %s. For module %s. Process %d.", + fdo_profile, + module_name, + min_process_id, + ) + global_client.key_value_set_bytes(profile_key, fdo_profile) + else: + logger.debug( + "Waiting for FDO profile: %s. For module %s. Should be set by process %d.", + fdo_profile, + module_name, + min_process_id, + ) + fdo_profile = global_client.blocking_key_value_get_bytes( + profile_key, share_timeout + ) + + _share_fdo_profiles.modules_profiles[profile_key] = fdo_profile + return fdo_profile + + +_share_fdo_profiles.modules_profiles = {} + + +# The process with the first_process_id should compile the module and write an +# autotune config to the K-V storage. +def _compile_and_write_autotune_config( + backend: xc.Client, + computation: ir.Module, + compile_options: xc.CompileOptions, + host_callbacks: Sequence[Any], + global_client: lib.xla_extension.DistributedRuntimeClient, + module_name: str, + cache_key: str, + first_process_id: int +) -> xc.LoadedExecutable: + share_timeout = config.share_binary_between_hosts_timeout_ms.value + debug_options = compile_options.executable_build_options.debug_options + + if _compile_and_write_autotune_config.autotune_configs_dir is None: + _compile_and_write_autotune_config.autotune_configs_dir = tempfile.mkdtemp() + + autotune_tmp_file = os.path.join( + _compile_and_write_autotune_config.autotune_configs_dir, cache_key + ) + + if os.path.exists(autotune_tmp_file): + logger.debug( + "Compiling module: %s. Use existing autotune config file: %s", + module_name, + autotune_tmp_file, + ) + debug_options.xla_gpu_load_autotune_results_from = autotune_tmp_file + return _compile_and_write_cache( + backend, + computation, + compile_options, + host_callbacks, + module_name, + cache_key, + ) + + if distributed.global_state.process_id == first_process_id: + debug_options.xla_gpu_dump_autotune_results_to = autotune_tmp_file + logger.debug("Process %d compiling and dumping autotune for module: %s", + first_process_id, module_name) + executable = _compile_and_write_cache( + backend, + computation, + compile_options, + host_callbacks, + module_name, + cache_key, + ) + + logger.debug( + "Writing autotune config for module %s to %s", + module_name, + autotune_tmp_file, + ) + with open(autotune_tmp_file, "rb") as f: + autotune_config = f.read() + + autotune_config = compilation_cache.compress_executable(autotune_config) + global_client.key_value_set_bytes(cache_key, autotune_config) + logger.debug( + "Autotune config for module %s with size %d shared by cache_key %s", + module_name, + len(autotune_config), + cache_key, + ) + else: + logger.debug( + "Compiling module %s, waiting for config to be shared by cache_key %s" + "from process %d", + module_name, + cache_key, + first_process_id + ) + autotune_config = global_client.blocking_key_value_get_bytes( + cache_key, share_timeout + ) + + logger.debug( + "Received autotune config for module %s of size %d", + module_name, + len(autotune_config), + ) + autotune_config = compilation_cache.decompress_executable(autotune_config) + with open(autotune_tmp_file, "wb") as f: + f.write(autotune_config) + + logger.debug( + "Compiling module %s, using autotune config from %s", + module_name, + autotune_tmp_file, + ) + debug_options.xla_gpu_load_autotune_results_from = autotune_tmp_file + executable = _compile_and_write_cache( + backend, + computation, + compile_options, + host_callbacks, + module_name, + cache_key, + ) + return executable + +_compile_and_write_autotune_config.autotune_configs_dir = None + +# The process with the first_process_id should compile the module and write it +# to the K-V storage. +def _compile_and_share_module( + backend: xc.Client, + computation: ir.Module, + compile_options: xc.CompileOptions, + host_callbacks: Sequence[Any], + global_client: lib.xla_extension.DistributedRuntimeClient, + module_name: str, + cache_key: str, + first_process_id: int +) -> xc.LoadedExecutable: + share_timeout = config.share_binary_between_hosts_timeout_ms.value + + # TODO: We need a proper eviction protocol here, otherwise all compiled + # modules will pile in memory. + if cache_key in _compile_and_share_module.modules_cache: + return _compile_and_share_module.modules_cache[cache_key] + + if distributed.global_state.process_id == first_process_id: + logger.debug("Process %d compiling and sharing module: %s", + first_process_id, module_name) + executable = _compile_and_write_cache( + backend, + computation, + compile_options, + host_callbacks, + module_name, + cache_key, + ) + serialized_executable = backend.serialize_executable(executable) + serialized_executable = compilation_cache.compress_executable( + serialized_executable + ) + global_client.key_value_set_bytes(cache_key, serialized_executable) + else: + logger.debug("Waiting for module: %s from process %d", module_name, + first_process_id) + serialized_executable = global_client.blocking_key_value_get_bytes( + cache_key, share_timeout + ) + serialized_executable = compilation_cache.decompress_executable( + serialized_executable + ) + executable = backend.deserialize_executable( + serialized_executable, compile_options + ) + + _compile_and_share_module.modules_cache[cache_key] = executable + return executable + +_compile_and_share_module.modules_cache = {} + +def _compile_and_write_cache( + backend: xc.Client, + computation: ir.Module, + compile_options: xc.CompileOptions, + host_callbacks: Sequence[Any], + module_name: str, + cache_key: str, +) -> xc.LoadedExecutable: + start_time = time.monotonic() + executable = backend_compile( + backend, computation, compile_options, host_callbacks + ) + compile_time = time.monotonic() - start_time + _cache_write( + cache_key, compile_time, module_name, backend, executable, host_callbacks + ) + return executable + +def _is_executable_in_cache(backend, cache_key) -> bool: + """Checks if executable is presented in cache on a given key + """ + try: + return compilation_cache.is_executable_in_cache(backend, cache_key) + except Exception as ex: + if config.raise_persistent_cache_errors.value: + raise + warnings.warn( + f"Error reading persistent compilation cache entry for " + f"'{cache_key}': {type(ex).__name__}: {ex}") + return False + + +def _cache_read( + module_name: str, cache_key: str, compile_options: xc.CompileOptions, + backend: xc.Client +) -> tuple[xc.LoadedExecutable | None, int | None]: + """Looks up the `computation` and it's compilation time in the persistent + compilation cache repository. + """ + try: + return compilation_cache.get_executable_and_time( + cache_key, compile_options, backend) + except Exception as ex: + if config.raise_persistent_cache_errors.value: + raise + warnings.warn( + f"Error reading persistent compilation cache entry for " + f"'{module_name}': {type(ex).__name__}: {ex}") + return None, None + + +def _cache_write(cache_key: str, + compile_time_secs: float, + module_name: str, + backend: xc.Client, executable: xc.LoadedExecutable, + host_callbacks: Sequence[Any]) -> None: + """Writes the `serialized_computation` and its compilation time to the + persistent compilation cache repository. + """ + # Only write cache entries from the first process. Otherwise we create + # problems with contention for writes on some filesystems, e.g., GCS. + if distributed.global_state.process_id != 0: + logger.debug("Not writing persistent cache entry since process_id != 0") + return + + if host_callbacks: + logger.debug( + "Not writing persistent cache entry for '%s' because it uses host " + "callbacks (e.g. from jax.debug.print or breakpoint)", module_name) + return + + min_compile_time = config.persistent_cache_min_compile_time_secs.value + if compile_time_secs < min_compile_time: + logger.debug( + "Not writing persistent cache entry for '%s' because it took < %.2f " + "seconds to compile (%.2fs)", module_name, min_compile_time, + compile_time_secs) + return + else: + logger.debug( + "'%s' took at least %.2f seconds to compile (%.2fs)", + module_name, min_compile_time, compile_time_secs) + + try: + compilation_cache.put_executable_and_time( + cache_key, module_name, executable, backend, int(compile_time_secs)) + except Exception as ex: + if config.raise_persistent_cache_errors.value: + raise + warnings.warn( + f"Error writing persistent compilation cache entry for " + f"'{module_name}': {type(ex).__name__}: {ex}") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compute_on.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compute_on.py new file mode 100644 index 0000000000000000000000000000000000000000..25b2be78d287d9521eab10f150f5cb27e843ef07 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/compute_on.py @@ -0,0 +1,55 @@ +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations +import threading +from contextlib import contextmanager + + +class ComputeOnContext(threading.local): + + def __init__(self): + self.stack = [] + +compute_on_context = ComputeOnContext() + + +@contextmanager +def extend_compute_type(c_type: str): + compute_on_context.stack.append(c_type) + try: + if len(set(filter(lambda x: x is not None, set(compute_on_context.stack)))) > 1: + raise NotImplementedError( + 'Nesting `compute_on` with different compute types is not supported' + f' yet. Current stack: {compute_on_context.stack}') + yield compute_on_context.stack[-1] + finally: + compute_on_context.stack.pop() + +def current_compute_type() -> str | None: + return compute_on_context.stack[-1] if compute_on_context.stack else None + +def _check_valid(c_type: str): + if c_type not in {'device_host', 'device'}: + raise ValueError('Invalid compute type received. Current supported values ' + f'are `device_host` and `device`. Got {c_type}') + +@contextmanager +def compute_on(compute_type: str): + if not isinstance(compute_type, str): + raise TypeError("`compute_on`'s compute_type argument must be a string.") + _check_valid(compute_type) + + with extend_compute_type(compute_type): + yield diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/config.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/config.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4c18fa0c715d056cc376f028a24c69b6fa8143 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/config.py @@ -0,0 +1,1640 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Hashable, Iterator +import contextlib +import functools +import itertools +import logging +import os +import sys +import threading +from typing import Any, Callable, Generic, NamedTuple, NoReturn, TypeVar, cast + +from jax._src import lib +from jax._src.lib import jax_jit +from jax._src.lib import transfer_guard_lib +from jax._src.lib import xla_client +from jax._src import logging_config + +logger = logging.getLogger(__name__) + +_T = TypeVar('_T') + + +def bool_env(varname: str, default: bool) -> bool: + """Read an environment variable and interpret it as a boolean. + + True values are (case insensitive): 'y', 'yes', 't', 'true', 'on', and '1'; + false values are 'n', 'no', 'f', 'false', 'off', and '0'. + + Args: + varname: the name of the variable + default: the default boolean value + Raises: ValueError if the environment variable is anything else. + """ + val = os.getenv(varname, str(default)) + val = val.lower() + if val in ('y', 'yes', 't', 'true', 'on', '1'): + return True + elif val in ('n', 'no', 'f', 'false', 'off', '0'): + return False + else: + raise ValueError(f"invalid truth value {val!r} for environment {varname!r}") + +def int_env(varname: str, default: int) -> int: + """Read an environment variable and interpret it as an integer.""" + return int(os.getenv(varname, str(default))) + + +UPGRADE_BOOL_HELP = ( + " This will be enabled by default in future versions of JAX, at which " + "point all uses of the flag will be considered deprecated (following " + "the `API compatibility policy " + "`_).") + +UPGRADE_BOOL_EXTRA_DESC = " (transient)" + + +class Config: + _HAS_DYNAMIC_ATTRIBUTES = True + + def __init__(self): + # There are two kinds of value holders: FlagHolders, which hold global + # flags, and StateContextManagers, which hold state that can be changed + # locally within a thread. A value holder needs a `.value` property and a + # `._set()` method. + self._value_holders = {} + self.meta = {} + self.use_absl = False + self._contextmanager_flags = set() + + def update(self, name, val): + if name not in self._value_holders: + raise AttributeError(f"Unrecognized config option: {name}") + self._value_holders[name]._set(val) + + def read(self, name): + if name in self._contextmanager_flags: + raise AttributeError( + "For flags with a corresponding contextmanager, read their value " + f"via e.g. `config.{name}` rather than `config.FLAGS.{name}`.") + return self._read(name) + + def _read(self, name): + try: + return self._value_holders[name].value + except KeyError: + raise AttributeError(f"Unrecognized config option: {name}") + + @property + def values(self): + return {name: holder.value for name, holder in self._value_holders.items()} + + def add_option(self, name, holder, opt_type, meta_args, meta_kwargs): + if name in self._value_holders: + raise Exception(f"Config option {name} already defined") + self._value_holders[name] = holder + self.meta[name] = (opt_type, meta_args, meta_kwargs) + + def config_with_absl(self): + """Registers absl flags for the JAX configs. + + E.g., for each JAX config defined using define_bool_state(), this method + registers an absl boolean flag, with the same name. + + This is the recommended method to call if you use `app.run(main)` and you + need JAX flags. Example: + + ```python + from absl import app + import jax + ... + + if __name__ == '__main__': + jax.config.config_with_absl() + app.run(main) + ``` + + """ + import absl.flags as absl_FLAGS # noqa: F401 # pytype: disable=import-error + from absl import app, flags as absl_flags # pytype: disable=import-error + + self.use_absl = True + self.absl_flags = absl_flags + absl_defs = { bool: absl_flags.DEFINE_bool, + int: absl_flags.DEFINE_integer, + float: absl_flags.DEFINE_float, + str: absl_flags.DEFINE_string, + 'enum': absl_flags.DEFINE_enum } + + for name, (flag_type, meta_args, meta_kwargs) in self.meta.items(): + holder = self._value_holders[name] + absl_defs[flag_type](name, holder.value, *meta_args, **meta_kwargs) + app.call_after_init(lambda: self.complete_absl_config(absl_flags)) + + def complete_absl_config(self, absl_flags): + # NOTE: avoid calling from outside this module. Instead, use + # `config_with_absl()`, and (in rare cases) `parse_flags_with_absl()`. + for name, holder in self._value_holders.items(): + try: + flag = absl_flags.FLAGS[name] + except KeyError: + # This can happen if a new flag was added after config_with_absl() was + # called, but before complete_absl_config was run. We could in principle + # add code to DEFINE_... to register any newly added flags with ABSL + # if config_with_absl() has already been called, but arguably the user + # should have called config_with_absl() later. + continue + if flag.present: + holder._set(flag.value) + + def parse_flags_with_absl(self): + """Parses command-line args that start with --jax. + + This method should be used only by advanced users. Most users should use + :meth:`config_with_absl` instead. + + This method has serious limitations: e.g., although it parses only the + --jax* command-line args, it runs the validators of all registered absl + flags, even non-JAX ones that have not been set yet; as such, for the + non-JAX flags, the validators run on the default flag values, not on the + values indicated by the command-line args. + """ + global already_configured_with_absl + if not already_configured_with_absl: + # Extract just the --jax... flags (before the first --) from argv. In some + # environments (e.g. ipython/colab) argv might be a mess of things + # parseable by absl and other junk. + jax_argv = itertools.takewhile(lambda a: a != '--', sys.argv) + jax_argv = ['', *(a for a in jax_argv if a.startswith('--jax'))] + + import absl.flags # pytype: disable=import-error + self.config_with_absl() + absl.flags.FLAGS(jax_argv, known_only=True) + self.complete_absl_config(absl.flags) + already_configured_with_absl = True + + +def trace_context(): + """Returns a tuple of configuration values that affect tracing. + + These values are included in the cache key for linear_util.cache. + + Values included in this set should also most likely be included in + the C++ JIT state, which is handled separately. + """ + tls = jax_jit.thread_local_state() + axis_env_state = () + mesh_context_manager = () + context: Any = tls.extra_jit_context + if context and context.axis_env_state is not None: + axis_env_state = context.axis_env_state + if context and context.mesh_context_manager: + mesh_context_manager = context.mesh_context_manager + return (axis_env_state, mesh_context_manager, enable_x64.value, + numpy_rank_promotion.value, default_matmul_precision.value, + dynamic_shapes.value, numpy_dtype_promotion.value, + default_device.value, random_seed_offset.value, + threefry_partitionable.value, + threefry_gpu_kernel_lowering.value, + softmax_custom_jvp.value, + enable_memories.value, + disable_jit.value, + debug_key_reuse.value, + jax_xla_profile_version.value, + # Technically this affects jaxpr->stablehlo lowering, not tracing. + hlo_source_file_canonicalization_regex.value, + pgle_profiling_runs.value, + enable_pgle.value) + +config = Config() + +_read = config._read +update = config.update +parse_flags_with_absl = config.parse_flags_with_absl + + +class NoDefault: pass +no_default = NoDefault() + + +class _Unset: pass +unset = _Unset() + +_thread_local_state = threading.local() + +class _StateContextManager(Generic[_T]): + __slots__ = ( + '_name', '_value', '_update_thread_local_hook', '_update_global_hook', + '_validator', '_default_context_manager_value', '__doc__', '__name__', + ) + + def __init__( + self, + name: str, + default: _T, + help, + update_global_hook: Callable[[_T], None] | None = None, + update_thread_local_hook: Callable[[_T | None], None] | None = None, + validator: Callable[[Any], None] | None = None, + extra_description: str = '', + default_context_manager_value: Any = no_default, + ): + self._name = name + self.__name__ = name[4:] if name.startswith('jax_') else name + self.__doc__ = (f"Context manager for `{name}` config option" + f"{extra_description}.\n\n{help}") + self._update_global_hook = update_global_hook + self._update_thread_local_hook = update_thread_local_hook + self._validator = validator + self._default_context_manager_value = default_context_manager_value + self._set(default) + + def __bool__(self) -> NoReturn: + raise TypeError( + "bool() not supported for instances of type '{0}' " + "(did you mean to use '{0}.value' instead?)".format( + type(self).__name__)) + + def _set(self, value: _T) -> None: + self._value = value + if self._update_global_hook: + self._update_global_hook(value) + + @property + def value(self) -> _T: + val = _thread_local_state.__dict__.get(self._name, unset) + return cast(_T, val) if val is not unset else self._value + + @contextlib.contextmanager + def __call__(self, new_val: Any = no_default): + if new_val is no_default: + if self._default_context_manager_value is not no_default: + new_val = self._default_context_manager_value # default_context_manager_value provided to constructor + else: + # no default_value provided to constructor and no value provided as an + # argument, so we raise an error + raise TypeError(f"Context manager for {self.__name__} config option " + "requires an argument representing the new value for " + "the config option.") + if self._validator: + self._validator(new_val) + prev_val = getattr(_thread_local_state, self._name, unset) + setattr(_thread_local_state, self._name, new_val) + if self._update_thread_local_hook: + self._update_thread_local_hook(new_val) + try: + yield + finally: + if prev_val is unset: + delattr(_thread_local_state, self._name) + if self._update_thread_local_hook: + self._update_thread_local_hook(None) + else: + setattr(_thread_local_state, self._name, prev_val) + if self._update_thread_local_hook: + self._update_thread_local_hook(cast(_T, prev_val)) + + def _add_hooks(self, update_global_hook, update_thread_local_hook): + """Private method that adds hooks to an existing context-manager. + + Used to avoid cyclic import dependencies.""" + self._update_thread_local_hook = update_thread_local_hook + self._update_global_hook = update_global_hook + update_global_hook(self._value) + + +def define_bool_state( + name: str, + default: bool, + help: str, + *, + update_global_hook: Callable[[bool], None] | None = None, + update_thread_local_hook: Callable[[bool | None], None] | None = None, + upgrade: bool = False, + extra_description: str = '', +) -> _StateContextManager[bool]: + """Set up thread-local state and return a contextmanager for managing it. + + This function is a convenience wrapper. It defines a flag, environment + variable, and corresponding thread-local state, which can be managed via the + contextmanager it returns. + + The thread-local state value can be read via the ``config.`` + attribute, where ``config`` is the singleton ``Config`` instance. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + default: boolean, a default value for the option. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + update_global_hook: a optional callback that is called with the updated + value of the global state when it is altered or set initially. + update_thread_local_hook: a optional callback that is called with the + updated value of the thread-local state when it is altered or set + initially. + upgrade: optional indicator that this flag controls a canonical feature + upgrade, so that it is `True` for the incoming functionality, `False` + for the outgoing functionality to be deprecated. + extra_description: string, optional: extra information to add to the + summary description. + + Returns: + A contextmanager to control the thread-local state value. + + Example: + + enable_foo = config.define_bool_state( + name='jax_enable_foo', + default=False, + help='Enable foo.') + + # Now the JAX_ENABLE_FOO shell environment variable and --jax_enable_foo + # command-line flag can be used to control the process-level value of + # the configuration option, in addition to using e.g. + # ``config.update("jax_enable_foo", True)`` directly. We can also use a + # context manager: + + with enable_foo(True): + ... + + The value of the thread-local state or flag can be accessed via + ``config.jax_enable_foo``. Reading it via ``config.FLAGS.jax_enable_foo`` is + an error. + """ + if not isinstance(default, bool): + raise TypeError(f"Default value must be of type bool, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + default = bool_env(name.upper(), default) + name = name.lower() + if upgrade: + help += ' ' + UPGRADE_BOOL_HELP + extra_description += UPGRADE_BOOL_EXTRA_DESC + config._contextmanager_flags.add(name) + + s = _StateContextManager[bool]( + name, default, help, update_global_hook=update_global_hook, + update_thread_local_hook=update_thread_local_hook, + extra_description=extra_description, default_context_manager_value=True) + config.add_option(name, s, bool, meta_args=[], meta_kwargs={"help": help}) + setattr(Config, name, property(lambda _: s.value)) + return s + + +def define_enum_state( + name: str, + enum_values: list[str], + default: str, + help: str, + *, + update_global_hook: Callable[[str], None] | None = None, + update_thread_local_hook: Callable[[str | None], None] | None = None, +) -> _StateContextManager[str]: + """Set up thread-local state and return a contextmanager for managing it. + + See docstring for ``define_bool_state``. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + enum_values: list of strings representing the possible values for the + option. + default: string, default value. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + + Returns: + A contextmanager to control the thread-local state value. + """ + if not isinstance(default, str): + raise TypeError(f"Default value must be of type str, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + name = name.lower() + default = os.getenv(name.upper(), default) + if default not in enum_values: + raise ValueError(f"Invalid value \"{default}\" for JAX flag {name}") + config._contextmanager_flags.add(name) + + def validator(new_val): + if type(new_val) is not str or new_val not in enum_values: + raise ValueError(f"new enum value must be in {enum_values}, " + f"got {new_val} of type {type(new_val)}.") + + s = _StateContextManager[str]( + name, + default, + help, + update_global_hook=update_global_hook, + update_thread_local_hook=update_thread_local_hook, + validator=validator, + ) + config.add_option( + name, s, 'enum', + meta_args=[], + meta_kwargs={"enum_values": enum_values, "help": help} + ) + setattr(Config, name, property(lambda _: s.value)) + return s + + +def define_optional_enum_state( + name: str, + enum_values: list[str], + default: str | None, + help: str, + *, + update_global_hook: Callable[[str | None], None] | None = None, + update_thread_local_hook: Callable[[str | None], None] | None = None, +) -> _StateContextManager[str | None]: + """Set up thread-local state and return a contextmanager for managing it. + + See docstring for ``define_bool_state``. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + enum_values: list of strings representing the possible values for the + option. + default: optional string, default value. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + + Returns: + A contextmanager to control the thread-local state value. + """ + if default is not None and not isinstance(default, str): + raise TypeError(f"Default value must be of type str or None, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + name = name.lower() + default = os.getenv(name.upper(), default) + if default is not None and default not in enum_values: + raise ValueError(f"Invalid value \"{default}\" for JAX flag {name}") + config._contextmanager_flags.add(name) + + def validate(new_val): + if (new_val is not None and + (type(new_val) is not str or new_val not in enum_values)): + raise ValueError(f"new enum value must be None or in {enum_values}, " + f"got {new_val} of type {type(new_val)}.") + + s = _StateContextManager['str | None']( + name, default, help, update_global_hook, update_thread_local_hook, + validate + ) + config.add_option( + name, s, 'enum', + meta_args=[], + meta_kwargs={"enum_values": enum_values, "help": help} + ) + setattr(Config, name, property(lambda _: s.value)) + return s + + +def define_int_state( + name: str, + default: int, + help: str, + *, + update_global_hook: Callable[[int], None] | None = None, + update_thread_local_hook: Callable[[int | None], None] | None = None, +) -> _StateContextManager[int]: + """Set up thread-local state and return a contextmanager for managing it. + + See docstring for ``define_bool_state``. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + default: optional int, default value. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + + Returns: + A contextmanager to control the thread-local state value. + """ + if not isinstance(default, int): + raise TypeError(f"Default value must be of type int, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + name = name.lower() + default_env = os.getenv(name.upper()) + if default_env is not None: + try: + default = int(default_env) + except ValueError: + raise ValueError(f"Invalid value \"{default_env}\" for JAX flag {name}") + config._contextmanager_flags.add(name) + + def validate(new_val): + if new_val is not None and not isinstance(new_val, int): + raise ValueError(f'new int config value must be None or of type int, ' + f'got {new_val} of type {type(new_val)}') + + s = _StateContextManager[int](name, default, help, update_global_hook, + update_thread_local_hook, validate) + config.add_option(name, s, int, meta_args=[], meta_kwargs={"help": help}) + setattr(Config, name, property(lambda _: s.value)) + return s + + +def define_float_state( + name: str, + default: float, + help: str, + *, + update_global_hook: Callable[[float], None] | None = None, + update_thread_local_hook: Callable[[float | None], None] | None = None, +) -> _StateContextManager[float]: + """Set up thread-local state and return a contextmanager for managing it. + + See docstring for ``define_bool_state``. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + default: default value. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + + Returns: + A contextmanager to control the thread-local state value. + """ + if not isinstance(default, float): + raise TypeError(f"Default value must be of type float, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + name = name.lower() + default_env = os.getenv(name.upper()) + if default_env is not None: + try: + default = float(default_env) + except ValueError: + raise ValueError(f"Invalid value \"{default_env}\" for JAX flag {name}") + config._contextmanager_flags.add(name) + + def validate(new_val): + if new_val is not None and not isinstance(new_val, (float, int)): + raise ValueError( + f'new float config value must be None or of type float, ' + f'got {new_val} of type {type(new_val)}') + + s = _StateContextManager[float](name, default, help, update_global_hook, + update_thread_local_hook, validate) + config.add_option(name, s, float, meta_args=[], meta_kwargs={"help": help}) + setattr(Config, name, property(lambda _: s.value)) + return s + + +def define_string_state( + name: str, + default: str, + help: str, + *, + update_global_hook: Callable[[str], None] | None = None, + update_thread_local_hook: Callable[[str | None], None] | None = None, +) -> _StateContextManager[str]: + """Set up thread-local state and return a contextmanager for managing it. + + See docstring for ``define_bool_state``. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + default: string, a default value for the option. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + update_global_hook: an optional callback that is called with the updated + value of the global state when it is altered or set initially. + update_thread_local_hook: an optional callback that is called with the + updated value of the thread-local state when it is altered or set + initially. + + Returns: + A contextmanager to control the thread-local state value. + """ + if not isinstance(default, str): + raise TypeError(f"Default value must be of type str, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + + def validator(new_val): + if not isinstance(new_val, str): + raise TypeError('new string config value must be of type str,' + f' got {new_val} of type {type(new_val)}.') + + return define_string_or_object_state( + name, default, help, + update_global_hook=update_global_hook, + update_thread_local_hook=update_thread_local_hook, + validator=validator) + + +def define_optional_string_state( + name: str, + default: str | None, + help: str, + *, + update_global_hook: Callable[[str], None] | None = None, + update_thread_local_hook: Callable[[str | None], None] | None = None, +) -> _StateContextManager[str | None]: + """Set up thread-local state and return a contextmanager for managing it. + + See docstring for ``define_bool_state``. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + default: optional string, a default value for the option. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + update_global_hook: an optional callback that is called with the updated + value of the global state when it is altered or set initially. + update_thread_local_hook: an optional callback that is called with the + updated value of the thread-local state when it is altered or set + initially. + + Returns: + A contextmanager to control the thread-local state value. + """ + if default is not None and not isinstance(default, str): + raise TypeError(f"Default value must be of type str or None, got {default} " + f"of type {getattr(type(default), '__name__', type(default))}") + + def validator(new_val): + if new_val is not None and not isinstance(new_val, str): + raise ValueError('new string config value must be None or of type str,' + f' got {new_val} of type {type(new_val)}.') + + return define_string_or_object_state( + name, default, help, + update_global_hook=update_global_hook, + update_thread_local_hook=update_thread_local_hook, + validator=validator) + +def define_string_or_object_state( + name: str, + default: Any, + help: str, + *, + update_global_hook: Callable[[Any], None] | None = None, + update_thread_local_hook: Callable[[Any], None] | None = None, + validator: Callable[[Any], None] | None = None, +) -> _StateContextManager[Any]: + """Set up thread-local state and return a contextmanager for managing it. + + Similar to ``define_string_state``, except the context manager will accept + any object, not just a string. Any value passed via command line flag or + environment variable will be treated as a string. + + Args: + name: string, converted to lowercase to define the name of the config + option (and absl flag). It is converted to uppercase to define the + corresponding shell environment variable. + default: string, a default value for the option. + help: string, used to populate the flag help information as well as the + docstring of the returned context manager. + update_global_hook: an optional callback that is called with the updated + value of the global state when it is altered or set initially. + update_thread_local_hook: an optional callback that is called with the + updated value of the thread-local state when it is altered or set + initially. + validator: an optional callback that is called with the new + value on any update, and should raise an error if the new value is + invalid. + + Returns: + A contextmanager to control the thread-local state value. + """ + name = name.lower() + default = os.getenv(name.upper(), default) + config._contextmanager_flags.add(name) + + s = _StateContextManager[Any]( + name, default, help, update_global_hook, update_thread_local_hook, + validator) + setattr(Config, name, property(lambda _: s.value)) + config.add_option(name, s, str, meta_args=[], meta_kwargs={"help": help}) + return s + + +class FlagHolder(Generic[_T]): + __slots__ = ("_name", "value", "_update_hook") + + _name: str + value: _T + _update_hook: Callable[[Any], None] | None + + def __init__(self, name: str, default: _T, + update_hook: Callable[[Any], None] | None = None): + self._name = name + self._update_hook = update_hook + self._set(default) + + def __bool__(self) -> NoReturn: + raise TypeError( + "bool() not supported for instances of type '{0}' " + "(did you mean to use '{0}.value' instead?)".format( + type(self).__name__)) + + def _set(self, value: _T) -> None: + self.value = value + if self._update_hook is not None: + self._update_hook(value) + + +def check_exists(name): + if name not in config._value_holders: + raise AttributeError(f"Unrecognized config option: {name}") + + +def DEFINE_bool(name, default, *args, **kwargs) -> FlagHolder[bool]: + update_hook = kwargs.pop("update_hook", None) + holder = FlagHolder(name, default, update_hook) + config.add_option(name, holder, bool, args, kwargs) + return holder + + +def DEFINE_integer(name, default, *args, **kwargs) -> FlagHolder[int]: + update_hook = kwargs.pop("update_hook", None) + holder = FlagHolder(name, default, update_hook) + config.add_option(name, holder, int, args, kwargs) + return holder + + +def DEFINE_float(name, default, *args, **kwargs) -> FlagHolder[float]: + update_hook = kwargs.pop("update_hook", None) + holder = FlagHolder(name, default, update_hook) + config.add_option(name, holder, float, args, kwargs) + return holder + + +def DEFINE_string(name, default, *args, **kwargs) -> FlagHolder[str]: + update_hook = kwargs.pop("update_hook", None) + holder = FlagHolder(name, default, update_hook) + config.add_option(name, holder, str, args, kwargs) + return holder + + +def DEFINE_enum(name, default, *args, **kwargs) -> FlagHolder[str]: + update_hook = kwargs.pop("update_hook", None) + holder = FlagHolder(name, default, update_hook) + config.add_option(name, holder, 'enum', args, kwargs) + return holder + + +already_configured_with_absl = False + + +# The C++ JIT maintains its own copy of several configuration items as +# a global/thread-local state. These methods allow updates to part of the +# state when a configuration value changes. +class _GlobalExtraJitContext(NamedTuple): + numpy_rank_promotion: str | None = None + numpy_dtype_promotion: str | None = None + default_matmul_precision: Any | None = None + dynamic_shapes: bool = False + random_seed_offset: int = 0 + threefry_partitionable: bool = False + threefry_gpu_kernel_lowering: bool = False + softmax_custom_jvp: bool = False + xla_profile_version: int = 0 + pgle_profiling_runs: int = 0 + enable_pgle: bool = False + + +def _update_global_jit_state(**kw): + gs = jax_jit.global_state() + context = gs.extra_jit_context or _GlobalExtraJitContext() + gs.extra_jit_context = context._replace(**kw) + + +class _ThreadLocalExtraJitContext(NamedTuple): + """A namedtuple containing states to add to the cache key. + + Just in time compilation (for jit, pmap, etc) behavior is configurable through + global and thread-local options, used in the cache key. + + The initialization, which uses both config.py and core.py is done using + `_update_thread_local_jit_state` in core.py to prevent circular imports. + """ + dynamic_trace_state: Any | None = None + axis_env_state: Hashable = () + mesh_context_manager: Hashable = () + + # Values set by _StateContextManager context managers. + # CAUTION: these must be initialized to `None`! The state context manager + # restores these to None on exit. If the object default is not `None`, the + # context manager is not a no-op, which leads to problems with stale state + # (e.g. spurious cache misses in tests). + numpy_rank_promotion: str | None = None + numpy_dtype_promotion: str | None = None + default_matmul_precision: Any | None = None + dynamic_shapes: bool | None = None + random_seed_offset: int | None = None + threefry_partitionable: bool | None = None + threefry_gpu_kernel_lowering: bool | None = None + softmax_custom_jvp: bool | None = None + xla_profile_version: int | None = None + pgle_profiling_runs: int | None = None + enable_pgle: bool | None = None + + +class _ThreadLocalStateCache(threading.local): + """"A thread local cache for _ThreadLocalExtraJitContext + + The extra_jit_context in jax_jit.thread_local_state() may get updated and thus + incurring dispatch overhead for comparing this python object during jit calls. + We want to duduplicate the objects that have the same hash/equality to also + have the same object ID, since the equality check is much faster if the object + IDs match. + """ + def __init__(self): + self.canonicalize = functools.lru_cache(128)(lambda x: x) + + +_thread_local_state_cache = _ThreadLocalStateCache() + + +def update_thread_local_jit_state(**kw): + tls = jax_jit.thread_local_state() + # After xla_client._version >= 70, the thread_local object will necessarily + # be initialized when accessed. The following line can be removed when the + # minimum jaxlib version is past version 70 + context = tls.extra_jit_context or _ThreadLocalExtraJitContext() + tmp = context._replace(**kw) + tls.extra_jit_context = _thread_local_state_cache.canonicalize(tmp) + + +# TODO(b/214340779): remove flag when XLA:CPU is improved. +jax2tf_associative_scan_reductions = define_bool_state( + name='jax2tf_associative_scan_reductions', + default=False, + help=( + 'JAX has two separate lowering rules for the cumulative reduction ' + 'primitives (cumsum, cumprod, cummax, cummin). On CPUs and GPUs it uses ' + 'a lax.associative_scan, while for TPUs it uses the HLO ReduceWindow. ' + 'The latter has a slow implementation on CPUs and GPUs. ' + 'By default, jax2tf uses the TPU lowering. Set this flag to True to ' + 'use the associative scan lowering usage, and only if it makes a difference ' + 'for your application. ' + 'See the jax2tf README.md for more details.' + ) +) + +jax2tf_default_native_serialization = define_bool_state( + name='jax2tf_default_native_serialization', + default=bool_env('JAX2TF_DEFAULT_NATIVE_SERIALIZATION', True), + help=( + 'Sets the default value of the native_serialization parameter to ' + 'jax2tf.convert. Prefer using the parameter instead of the flag, ' + 'the flag may be removed in the future.' + ) +) + +jax_serialization_version = define_int_state( + name='jax_serialization_version', + # Note: bump the default serialization version at least one month after + # we update XlaCallModule to support the new version, so that serialized + # modules are forward compatible with deployed versions of XlaCallModule. + # Version 9 of XlaCallModule is supported since October 27th, 2023. + default=int_env('JAX_SERIALIZATION_VERSION', 9), + help=( + 'The version number to use for native serialization. This must be ' + 'within the range of versions supported by the tf.XlaCallModule ' + 'used in your deployment environment. ' + 'See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions.' + ) +) + +jax_platforms = define_optional_string_state( + name='jax_platforms', + default=None, + help=( + 'Comma-separated list of platform names specifying which platforms jax ' + 'should initialize. If any of the platforms in this list are not successfully ' + 'initialized, an exception will be raised and the program will be aborted. ' + 'The first platform in the list will be the default platform. ' + 'For example, config.jax_platforms=cpu,tpu means that CPU and TPU backends ' + 'will be initialized, and the CPU backend will be used unless otherwise ' + 'specified. If TPU initialization fails, it will raise an exception. ' + 'By default, jax will try to initialize all available ' + 'platforms and will default to GPU or TPU if available, and fallback to CPU ' + 'otherwise.' + )) + +jax_pjrt_client_create_options = define_optional_string_state( + name='jax_pjrt_client_create_options', + default=None, + help=('A set of key-value pairs in the format of "k1:v1;k2:v2" strings ' + 'provided to a device platform pjrt client as extra arguments.')) + +enable_checks = define_bool_state( + name='jax_enable_checks', + default=False, + help='Turn on invariant checking for JAX internals. Makes things slower.') + +debug_key_reuse = define_bool_state( + name='jax_debug_key_reuse', + default=False, + help=('Turn on experimental key reuse checking. With this configuration enabled,' + ' typed PRNG keys (i.e. keys created with jax.random.key()) will have their' + ' usage tracked, and incorrect reuse of a previously-used key will lead to' + ' an error. Currently enabling this leads to a small Python overhead on' + ' every call to a JIT-compiled function with keys as inputs or outputs.')) + +check_tracer_leaks = define_bool_state( + name='jax_check_tracer_leaks', + default=False, + help=('Turn on checking for leaked tracers as soon as a trace completes. ' + 'Enabling leak checking may have performance impacts: some caching ' + 'is disabled, and other overheads may be added. Additionally, be aware ' + 'that some Python debuggers can cause false positives, so it is recommended ' + 'to disable any debuggers while leak checking is enabled.')) +checking_leaks = functools.partial(check_tracer_leaks, True) + +debug_nans = define_bool_state( + name='jax_debug_nans', + default=False, + help=('Add nan checks to every operation. When a nan is detected on the ' + 'output of a jit-compiled computation, call into the un-compiled ' + 'version in an attempt to more precisely identify the operation ' + 'which produced the nan.')) + +debug_infs = define_bool_state( + name='jax_debug_infs', + default=False, + help=('Add inf checks to every operation. When an inf is detected on the ' + 'output of a jit-compiled computation, call into the un-compiled ' + 'version in an attempt to more precisely identify the operation ' + 'which produced the inf.')) + +log_compiles = define_bool_state( + name='jax_log_compiles', + default=False, + help=('Log a message each time `jit` or `pmap` compiles an XLA ' + 'computation. Logging is performed with `logging`. When this ' + 'option is set, the log level is WARNING; otherwise the level is ' + 'DEBUG.')) + +explain_cache_misses = define_bool_state( + name='jax_explain_cache_misses', + default=False, + help=('Each time there is a miss on one of the main caches (e.g. the ' + 'tracing cache), log an explanation.. Logging is performed with ' + '`logging`. When this option is set, the log level is WARNING; ' + 'otherwise the level is DEBUG.')) + +log_checkpoint_residuals = define_bool_state( + name='jax_log_checkpoint_residuals', + default=False, + help=('Log a message every time jax.checkpoint (aka jax.remat) is ' + 'partially evaluated (e.g. for autodiff), printing what residuals ' + 'are saved.')) + +pmap_shmap_merge = define_bool_state( + name='jax_pmap_shmap_merge', + default=False, + upgrade=True, + help='If True, pmap and shard_map API will be merged.') + +def _update_jax_memories_global(val): + lib.jax_jit.global_state().enable_memories = val + +def _update_jax_memories_thread_local(val): + lib.jax_jit.thread_local_state().enable_memories = val + +enable_memories = define_bool_state( + 'jax_enable_memories', + default=False, + upgrade=True, + update_global_hook=_update_jax_memories_global, + update_thread_local_hook=_update_jax_memories_thread_local, + help=("If True, will allow fetching memory kinds available on executable " + "and annotate Shardings with it.")) + +spmd_mode = define_enum_state( + name='jax_spmd_mode', + enum_values=['allow_all', 'allow_jit'], + default='allow_jit', + help=("Decides whether Math on `jax.Array`'s that are not fully addressable " + "(i.e. spans across multiple processes) is allowed. The options are: " + "* allow_jit: Default, `pjit` and `jax.jit` computations are allowed " + " to execute on non-fully addressable `jax.Array`s\n" + "* allow_all: `jnp`, normal math (like `a + b`, etc), `pjit`, " + " `jax.jit` and all other operations are allowed to " + " execute on non-fully addressable `jax.Array`s.")) + + +distributed_debug = define_bool_state( + name='jax_distributed_debug', + default=False, + help=('Enable logging useful for debugging multi-process distributed ' + 'computations. Logging is performed with `logging` at WARNING ' + 'level.')) + +random_seed_offset = define_int_state( + name='jax_random_seed_offset', + default=0, + help=('Offset to all random seeds (e.g. argument to jax.random.key()).'), + update_global_hook=lambda val: _update_global_jit_state( + random_seed_offset=val), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + random_seed_offset=val) +) + +legacy_prng_key = define_enum_state( + name='jax_legacy_prng_key', + enum_values=['allow', 'warn', 'error'], + default='allow', + help=('Specify the behavior when raw PRNG keys are passed to ' + 'jax.random APIs.') +) + +enable_custom_prng = define_bool_state( + name='jax_enable_custom_prng', + default=False, + upgrade=True, + help=('Enables an internal upgrade that allows one to define custom ' + 'pseudo-random number generator implementations.')) + +default_prng_impl = define_enum_state( + name='jax_default_prng_impl', + enum_values=['threefry2x32', 'rbg', 'unsafe_rbg'], + default='threefry2x32', + help=('Select the default PRNG implementation, used when one is not ' + 'explicitly provided at seeding time.')) + +threefry_partitionable = define_bool_state( + name='jax_threefry_partitionable', + default=False, + upgrade=True, + help=('Enables internal threefry PRNG implementation changes that ' + 'render it automatically partitionable in some cases. Without this ' + 'flag, using the standard jax.random pseudo-random number generation ' + 'may result in extraneous communication and/or redundant distributed ' + 'computation. With this flag, the communication overheads disappear ' + 'in some cases.'), + update_global_hook=lambda val: _update_global_jit_state( + threefry_partitionable=val), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + threefry_partitionable=val)) + +threefry_gpu_kernel_lowering = define_bool_state( + name='jax_threefry_gpu_kernel_lowering', + default=False, + help=('On GPU, lower threefry PRNG operations to a kernel implementation. ' + 'This makes compile times faster at a potential runtime memory ' + 'cost.'), + update_global_hook=lambda val: _update_global_jit_state( + threefry_gpu_kernel_lowering=val), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + threefry_gpu_kernel_lowering=val)) + + +softmax_custom_jvp = define_bool_state( + name='jax_softmax_custom_jvp', + default=False, + upgrade=True, + help=('Use a new custom_jvp rule for jax.nn.softmax. The new rule should ' + 'improve memory usage and stability. Set True to use new ' + 'behavior. See https://github.com/google/jax/pull/15677'), + update_global_hook=lambda val: _update_global_jit_state( + softmax_custom_jvp=val), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + softmax_custom_jvp=val)) + + +enable_custom_vjp_by_custom_transpose = define_bool_state( + name='jax_enable_custom_vjp_by_custom_transpose', + default=False, + upgrade=True, + help=('Enables an internal upgrade that implements `jax.custom_vjp` by ' + 'reduction to `jax.custom_jvp` and `jax.custom_transpose`.')) + +raise_persistent_cache_errors = define_bool_state( + name='jax_raise_persistent_cache_errors', + default=False, + help=('If true, exceptions raised when reading or writing to the ' + 'persistent compilation cache will be allowed through, halting ' + 'program execution if not manually caught. If false, exceptions are ' + 'caught and raised as warnings, allowing program execution to ' + 'continue. Defaults to false so cache bugs or intermittent issues ' + 'are non-fatal.')) + +persistent_cache_min_compile_time_secs = define_float_state( + name='jax_persistent_cache_min_compile_time_secs', + default=1., + help=('The minimum compile time of a computation to be written to the ' + 'persistent compilation cache. This threshold can be raised to ' + 'decrease the number of entries written to the cache.')) + +persistent_cache_min_entry_size_bytes = define_int_state( + name='jax_persistent_cache_min_entry_size_bytes', + default=0, + help=('The minimum size (in bytes) of an entry that will be cached in the ' + 'persistent compilation cache: ' + '* -1: disable the size restriction and prevent overrides. ' + '* Leave at default (0) to allow for overrides. The override will ' + ' typically ensure that the minimum size is optimal for the ' + ' filesystem being used for the cache. ' + '* > 0: the actual minimum size desired; no overrides.')) + +compilation_cache_include_metadata_in_key = define_bool_state( + name='jax_compilation_cache_include_metadata_in_key', + default=False, + help=( + 'Include metadata, such as file names and line numbers, in the' + ' compilation cache key. If false, the cache will still get hits even' + ' if functions or files are moved, etc. However, it means that' + ' executables loaded from the cache may have stale metadata, which' + ' may show up in, e.g., profiles.' + ), +) + +hlo_source_file_canonicalization_regex = define_optional_string_state( + name='jax_hlo_source_file_canonicalization_regex', + default=None, + help=('Used to canonicalize the source_path metadata of HLO instructions ' + 'by removing the given regex. If set, re.sub() is called on each ' + 'source_file with the given regex, and all matches are removed. ' + 'This can be used to avoid spurious cache misses when using the ' + 'persistent compilation cache, which includes HLO metadata in the ' + 'cache key.')) + +include_full_tracebacks_in_locations = define_bool_state( + name='jax_include_full_tracebacks_in_locations', + default=True, + help=( + 'Include Python tracebacks in MLIR locations in IR emitted by JAX.' + ), +) + +traceback_in_locations_limit = define_int_state( + name='jax_traceback_in_locations_limit', + default=10, + help=( + 'Limit the number of frames at the Python traceback frames included in ' + 'MLIR locations. If set to the negative value, traceback will not be ' + 'limited.' + ), +) + +share_autotune_config_between_hosts = define_bool_state( + name='jax_share_autotune_config_between_hosts', + default=False, + help=( + 'If set to True, the coordinator process will share autotune configs ' + 'other participants. This will increase overall compilation time, but ' + 'will lead to equal compiled modules in each process. ' + 'If both jax_share_binary_between_hosts and ' + 'jax_share_autotune_config_between_hosts are set, compiled HLO will be ' + "shared when it's possible and autotune config sharing will be used " + 'as a fallback.' + ), +) + +share_binary_between_hosts = define_bool_state( + name='jax_share_binary_between_hosts', + default=False, + help=( + 'If set to True, the compiled module will be shared between hosts ' + 'directly.' + ), +) + +share_binary_between_hosts_timeout_ms = define_int_state( + name='jax_share_binary_between_hosts_timeout_ms', + default=20 * 60 * 1000, + help='Timeout for the compiled module share.', +) + +enable_pgle = define_bool_state( + name='jax_enable_pgle', + default=False, + help=( + 'If set to True and the property jax_pgle_profiling_runs is set to ' + 'greater than 0, the modules will be recompiled after running specified ' + 'number times with collected data provided to the profile guided latency ' + 'estimator.' + ), + update_global_hook=lambda val: _update_global_jit_state(enable_pgle=val), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + enable_pgle=val), +) + +pgle_profiling_runs = define_int_state( + name='jax_pgle_profiling_runs', + default=3, + help=( + 'Amount of times module should be profiled before recompilation when ' + 'PGLE is used.' + ), + update_global_hook=lambda val: _update_global_jit_state( + pgle_profiling_runs=val + ), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + pgle_profiling_runs=val + ), +) + +pgle_aggregation_percentile = define_int_state( + name='jax_pgle_aggregation_percentile', + default=90, + help='Percentile used to aggregate performance data between devices when ' + 'PGLE is used.', +) + +enable_compilation_cache = define_bool_state( + name='jax_enable_compilation_cache', + default=True, + help=('If set to False, the compilation cache will be disabled regardless ' + 'of whether set_cache_dir() was called. If set to True, the ' + 'path could be set to a default value or via a call to ' + 'set_cache_dir().'), +) + +compilation_cache_dir = define_optional_string_state( + name='jax_compilation_cache_dir', + default=None, + help=('Path for the cache. ' + 'Precedence: ' + '1. A call to compilation_cache.set_cache_dir(). ' + '2. The value of this flag set in the command line or by default.'), +) + +default_dtype_bits = define_enum_state( + name='jax_default_dtype_bits', + enum_values=['32', '64'], + default='64', + help=('Specify bit width of default dtypes, either 32-bit or 64-bit. ' + 'This is a temporary flag that will be used during the process ' + 'of deprecating the ``jax_enable_x64`` flag.')) + +numpy_dtype_promotion = define_enum_state( + name='jax_numpy_dtype_promotion', + enum_values=['standard', 'strict'], + default='standard', + help=('Specify the rules used for implicit type promotion in operations ' + 'between arrays. Options are "standard" or "strict"; in strict-mode, ' + 'binary operations between arrays of differing strongly-specified ' + 'dtypes will result in an error.'), + update_global_hook=lambda val: \ + _update_global_jit_state(numpy_dtype_promotion=val), + update_thread_local_hook=lambda val: \ + update_thread_local_jit_state(numpy_dtype_promotion=val)) + +def _update_x64_global(val): + lib.jax_jit.global_state().enable_x64 = val + +def _update_x64_thread_local(val): + lib.jax_jit.thread_local_state().enable_x64 = val + +enable_x64 = define_bool_state( + name='jax_enable_x64', + default=False, + help='Enable 64-bit types to be used', + update_global_hook=_update_x64_global, + update_thread_local_hook=_update_x64_thread_local) + +# TODO(phawkins): remove after fixing users of FLAGS.x64_enabled. +config._contextmanager_flags.remove('jax_enable_x64') + +setattr(Config, "x64_enabled", property(lambda _: enable_x64.value)) + +def _update_default_device_global(val): + lib.jax_jit.global_state().default_device = val + + +def _update_default_device_thread_local(val): + lib.jax_jit.thread_local_state().default_device = val + + +def _validate_default_device(val): + if val is not None and not isinstance(val, xla_client.Device): + # TODO(skyewm): this is a workaround for non-PJRT Device types. Remove when + # all JAX backends use a single C++ device interface. + if 'Device' in str(type(val)): + logger.info( + 'Allowing non-`xla_client.Device` default device: %s, type: %s', + repr(val), type(val)) + return + raise ValueError('jax.default_device must be passed a Device object (e.g. ' + f"`jax.devices('cpu')[0]`), got: {val!r}") + + +# TODO(skye): default_device only accepts devices for now. Make it work with +# platform names as well (e.g. "cpu" to mean the same as jax.devices("cpu")[0]). +default_device = define_string_or_object_state( + name='jax_default_device', + default=None, + help=( + 'Configure the default device for JAX operations. Set to a Device ' + 'object (e.g. ``jax.devices("cpu")[0]``) to use that Device as the ' + 'default device for JAX operations and jit\'d function calls (there is ' + 'no effect on multi-device computations, e.g. pmapped function calls). ' + 'Set to None to use the system default device. See ' + ':ref:`faq-data-placement` for more information on device placement.'), + update_global_hook=_update_default_device_global, + update_thread_local_hook=_update_default_device_thread_local, + validator=_validate_default_device) + +def _update_disable_jit_global(val): + lib.jax_jit.global_state().disable_jit = val + +def _update_disable_jit_thread_local(val): + lib.jax_jit.thread_local_state().disable_jit = val + +disable_jit = define_bool_state( + name='jax_disable_jit', + default=False, + help=('Disable JIT compilation and just call original Python.'), + update_global_hook=_update_disable_jit_global, + update_thread_local_hook=_update_disable_jit_thread_local) + + +numpy_rank_promotion = define_enum_state( + name='jax_numpy_rank_promotion', + enum_values=['allow', 'warn', 'raise'], + default='allow', + help=('Control NumPy-style automatic rank promotion broadcasting ' + '("allow", "warn", or "raise").'), + update_global_hook=lambda val: \ + _update_global_jit_state(numpy_rank_promotion=val), + update_thread_local_hook=lambda val: \ + update_thread_local_jit_state(numpy_rank_promotion=val)) + +default_matmul_precision = define_optional_enum_state( + name='jax_default_matmul_precision', + enum_values=['bfloat16', 'tensorfloat32', 'float32'], + default=None, + help=('Control the default matmul and conv precision for 32bit inputs.\n\n' + + 'Some platforms, like TPU, offer configurable precision levels for ' + 'matrix multiplication and convolution computations, trading off ' + 'accuracy for speed. The precision can be controlled for each ' + 'operation; for example, see the :func:`jax.lax.conv_general_dilated` ' + 'and :func:`jax.lax.dot` docstrings. But it can be useful to control ' + 'the default behavior obtained when an operation is not given a ' + 'specific precision.\n\n' + + 'This option can be used to control the default precision ' + 'level for computations involved in matrix multiplication and ' + 'convolution on 32bit inputs. The levels roughly describe the ' + "precision at which scalar products are computed. The 'bfloat16' " + "option is the fastest and least precise; 'float32' is similar to " + "full float32 precision; 'tensorfloat32' is intermediate.\n\n"), + update_global_hook=lambda val: \ + _update_global_jit_state(default_matmul_precision=val), + update_thread_local_hook=lambda val: \ + update_thread_local_jit_state(default_matmul_precision=val)) + +traceback_filtering = define_enum_state( + name = 'jax_traceback_filtering', + enum_values=["off", "tracebackhide", "remove_frames", "quiet_remove_frames", + "auto"], + default="auto", + help="Controls how JAX filters internal frames out of tracebacks.\n\n" + "Valid values are:\n" + " * \"off\": disables traceback filtering.\n" + " * \"auto\": use \"tracebackhide\" if running under a sufficiently" + " new IPython, or \"remove_frames\" otherwise.\n" + " * \"tracebackhide\": adds \"__tracebackhide__\" annotations to" + " hidden stack frames, which some traceback printers support.\n" + " * \"remove_frames\": removes hidden frames from tracebacks, and adds" + " the unfiltered traceback as a __cause__ of the exception.\n" + " * \"quiet_remove_frames\": removes hidden frames from tracebacks, and adds" + " a brief message (to the __cause__ of the exception) describing that this has" + " happened.\n") + +# This flag is for internal use. +# TODO(tianjianlu): Removes once we always enable cusparse lowering. +# TODO(b/262050896): Set to true after bug is fixed +bcoo_cusparse_lowering = define_bool_state( + name='jax_bcoo_cusparse_lowering', + default=False, + help=('Enables lowering BCOO ops to cuSparse.')) + +# TODO(mattjj): remove this flag when we ensure we only succeed at trace-staging +# if the intended backend can handle lowering the result +dynamic_shapes = define_bool_state( + name='jax_dynamic_shapes', + default=bool(os.getenv('JAX_DYNAMIC_SHAPES', '')), + help=('Enables experimental features for staging out computations with ' + 'dynamic shapes.'), + update_global_hook=lambda val: \ + _update_global_jit_state(dynamic_shapes=val), + update_thread_local_hook=lambda val: \ + update_thread_local_jit_state(dynamic_shapes=val)) + +# This flag is temporary during rollout of the remat barrier. +# TODO(parkers): Remove if there are no complaints. +remat_opt_barrier = define_bool_state( + name='jax_remat_opt_barrier', + default=True, + help=('Enables using optimization-barrier op for lowering remat.')) + +# TODO(sharadmv,mattjj): set default to True, then remove +eager_pmap = define_bool_state( + name='jax_eager_pmap', + default=True, + upgrade=True, + help='Enable eager-mode pmap when jax_disable_jit is activated.') + +# TODO(mattjj): remove once we land mutable array plumbing, or face great shame +custom_vjp_disable_shape_check = define_bool_state( + name='jax_custom_vjp_disable_shape_check', + default=False, + upgrade=True, + help='Disable the check from #19009 to enable some custom_vjp hacks.') + +xla_runtime_errors = define_bool_state( + name='jax_experimental_unsafe_xla_runtime_errors', + default=False, + help=('Enable XLA runtime errors for jax.experimental.checkify.checks ' + 'on CPU and GPU. These errors are async, might get lost and are not ' + 'very readable. But, they crash the computation and enable you ' + 'to write jittable checks without needing to checkify. Does not ' + 'work under pmap/pjit.') +) + +jax_xla_profile_version = define_int_state( + name='jax_xla_profile_version', + default=0, + help=( + 'Optional profile version for XLA compilation. This is meaningful ' + 'only when XLA is configured to support the remote compilation ' + 'profile feature.'), + update_global_hook=lambda val: _update_global_jit_state( + xla_profile_version=val), + update_thread_local_hook=lambda val: update_thread_local_jit_state( + xla_profile_version=val), +) + +@contextlib.contextmanager +def explicit_device_put_scope() -> Iterator[None]: + """Indicates that the current context is an explicit device_put*() call.""" + state = transfer_guard_lib.thread_local_state() + prev = state.explicit_device_put + state.explicit_device_put = True + try: + yield + finally: + state.explicit_device_put = prev + +@contextlib.contextmanager +def explicit_device_get_scope() -> Iterator[None]: + """Indicates that the current context is an explicit device_get() call.""" + state = transfer_guard_lib.thread_local_state() + prev = state.explicit_device_get + state.explicit_device_get = True + try: + yield + finally: + state.explicit_device_get = prev + +def _update_transfer_guard(state, key, val): + """Applies the transfer guard level within transfer_guard_lib.""" + if val is None: + setattr(state, key, None) + elif val == 'allow': + setattr(state, key, transfer_guard_lib.TransferGuardLevel.ALLOW) + elif val == 'log': + setattr(state, key, transfer_guard_lib.TransferGuardLevel.LOG) + elif val == 'disallow': + setattr(state, key, transfer_guard_lib.TransferGuardLevel.DISALLOW) + elif val == 'log_explicit': + setattr(state, key, transfer_guard_lib.TransferGuardLevel.LOG_EXPLICIT) + elif val == 'disallow_explicit': + setattr(state, key, transfer_guard_lib.TransferGuardLevel.DISALLOW_EXPLICIT) + else: + assert False, f'Invalid transfer guard level {val}' + +transfer_guard_host_to_device = define_optional_enum_state( + name='jax_transfer_guard_host_to_device', + enum_values=[ + 'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit' + ], + # The default is applied by transfer_guard_lib. Use None here to avoid + # accidentally overriding --jax_transfer_guard. + default=None, + help=('Select the transfer guard level for host-to-device transfers. ' + 'Default is "allow".'), + update_global_hook=lambda val: _update_transfer_guard( + transfer_guard_lib.global_state(), 'host_to_device', val), + update_thread_local_hook=lambda val: _update_transfer_guard( + transfer_guard_lib.thread_local_state(), 'host_to_device', val)) + +transfer_guard_device_to_device = define_optional_enum_state( + name='jax_transfer_guard_device_to_device', + enum_values=[ + 'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit' + ], + # The default is applied by transfer_guard_lib. Use None here to avoid + # accidentally overriding --jax_transfer_guard. + default=None, + help=('Select the transfer guard level for device-to-device transfers. ' + 'Default is "allow".'), + update_global_hook=lambda val: _update_transfer_guard( + transfer_guard_lib.global_state(), 'device_to_device', val), + update_thread_local_hook=lambda val: _update_transfer_guard( + transfer_guard_lib.thread_local_state(), 'device_to_device', val)) + +transfer_guard_device_to_host = define_optional_enum_state( + name='jax_transfer_guard_device_to_host', + enum_values=[ + 'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit' + ], + # The default is applied by transfer_guard_lib. Use None here to avoid + # accidentally overriding --jax_transfer_guard. + default=None, + help=('Select the transfer guard level for device-to-host transfers. ' + 'Default is "allow".'), + update_global_hook=lambda val: _update_transfer_guard( + transfer_guard_lib.global_state(), 'device_to_host', val), + update_thread_local_hook=lambda val: _update_transfer_guard( + transfer_guard_lib.thread_local_state(), 'device_to_host', val)) + +def _update_all_transfer_guard_global(val): + for name in ('jax_transfer_guard_host_to_device', + 'jax_transfer_guard_device_to_device', + 'jax_transfer_guard_device_to_host'): + config.update(name, val) + +_transfer_guard = define_optional_enum_state( + name='jax_transfer_guard', + enum_values=[ + 'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit' + ], + # The default is applied by transfer_guard_lib. Use None here to avoid + # accidentally overriding --jax_transfer_guard_*. + default=None, + help=('Select the transfer guard level for all transfers. This option is ' + 'set-only; the transfer guard level for a specific direction should ' + 'be read using the per-transfer direction option. ' + 'Default is "allow".'), + update_global_hook=_update_all_transfer_guard_global) + +@contextlib.contextmanager +def transfer_guard(new_val: str) -> Iterator[None]: + """A contextmanager to control the transfer guard level for all transfers. + + For more information, see + https://jax.readthedocs.io/en/latest/transfer_guard.html + + Args: + new_val: The new thread-local transfer guard level for all transfers. + + Yields: + None. + """ + with contextlib.ExitStack() as stack: + stack.enter_context(transfer_guard_host_to_device(new_val)) + stack.enter_context(transfer_guard_device_to_device(new_val)) + stack.enter_context(transfer_guard_device_to_host(new_val)) + stack.enter_context(_transfer_guard(new_val)) + yield + + +def _update_debug_log_modules(module_names_str: str | None): + logging_config.disable_all_debug_logging() + if not module_names_str: + return + module_names = module_names_str.split(',') + for module_name in module_names: + logging_config.enable_debug_logging(module_name) + +# Don't define a context manager since this isn't threadsafe. +define_string_state( + name='jax_debug_log_modules', + default='', + help=('Comma-separated list of module names (e.g. "jax" or ' + '"jax._src.xla_bridge,jax._src.dispatch") to enable debug logging ' + 'for.'), + update_global_hook=_update_debug_log_modules) + +pmap_no_rank_reduction = define_bool_state( + name='jax_pmap_no_rank_reduction', + default=False, + help=( + "If True, pmap shards have a the same rank as their enclosing array." + ) +) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/core.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/core.py new file mode 100644 index 0000000000000000000000000000000000000000..53c2e71f8c9b76c10063c0fc3b064460f9bcefd2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/core.py @@ -0,0 +1,3448 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import collections # noqa: F401 +from collections import Counter, defaultdict, deque, namedtuple +from collections.abc import (Collection, Generator, Hashable, Iterable, + Iterator, Set, Sequence, MutableSet, + MutableMapping) +from contextlib import contextmanager, ExitStack +from dataclasses import dataclass +import functools +from functools import partial, partialmethod, total_ordering +import gc +import inspect +import itertools as it +import math +import operator +import threading +import types +from typing import (Any, Callable, ClassVar, Generic, NamedTuple, TypeVar, + cast, overload, Union) +import warnings +from weakref import ref + +import numpy as np + +from jax._src import dtypes +from jax._src import config +from jax._src import effects +from jax._src import compute_on +from jax._src.errors import ( + ConcretizationTypeError, TracerArrayConversionError, TracerBoolConversionError, + TracerIntegerConversionError, UnexpectedTracerError) +from jax._src import linear_util as lu + +from jax._src import source_info_util +from jax._src.util import (safe_zip, safe_map, curry, tuple_insert, + tuple_delete, as_hashable_function, + HashableFunction, HashableWrapper, weakref_lru_cache, + partition_list, StrictABCMeta) +import jax._src.pretty_printer as pp +from jax._src.lib import jax_jit +from jax._src import traceback_util +from jax._src.typing import Array, DimSize, Shape +from jax._src import typing + +traceback_util.register_exclusion(__file__) + +zip, unsafe_zip = safe_zip, zip +map, unsafe_map = safe_map, map + + +_TRACER_ERROR_NUM_TRACEBACK_FRAMES = config.DEFINE_integer( + 'jax_tracer_error_num_traceback_frames', + config.int_env('JAX_TRACER_ERROR_NUM_TRACEBACK_FRAMES', 5), + help='Set the number of stack frames in JAX tracer error messages.' +) + + +# -------------------- jaxprs -------------------- + +Effect = effects.Effect +Effects = effects.Effects +EffectTypeSet = effects.EffectTypeSet +no_effects: Effects = effects.no_effects + +class JaxprDebugInfo(NamedTuple): + traced_for: str # e.g. 'jit', 'scan', etc + func_src_info: str # e.g. f'{fun.__name__} at {filename}:{lineno}' + arg_names: tuple[str | None, ...] # e.g. ('args[0]', ... ) + result_paths: tuple[str, ...] # e.g. ('[0]', '[1]', ...) + +class Jaxpr: + __slots__ = ['__weakref__', '_constvars', '_invars', '_outvars', '_eqns', + '_effects', '_debug_info'] + + _constvars: list[Var] + _invars: list[Var] + _outvars: list[Atom] + _eqns: list[JaxprEqn] + _effects: Effects + _debug_info: JaxprDebugInfo | None + + @property + def constvars(self) -> list[Var]: + return self._constvars + + @property + def invars(self) -> list[Var]: + return self._invars + + @property + def outvars(self) -> list[Atom]: + return self._outvars + + @property + def eqns(self) -> list[JaxprEqn]: + return self._eqns + + @property + def effects(self) -> Effects: + return self._effects + + @property + def debug_info(self) -> JaxprDebugInfo | None: + return self._debug_info + + def __init__(self, constvars: Sequence[Var], invars: Sequence[Var], + outvars: Sequence[Atom], eqns: Sequence[JaxprEqn], + effects: Effects = no_effects, + debug_info: JaxprDebugInfo | None = None): + """ + Args: + constvars: list of variables introduced for constants. Array constants are + replaced with such variables while scalar constants are kept inline. + invars: list of input variables. Together, `constvars` and `invars` are + the inputs to the Jaxpr. + outvars: list of output atoms. + eqns: list of equations. + effects: set of effects. The effects on a jaxpr are a superset of the + union of the effects for each equation. + debug_info: optional JaxprDebugInfo. + """ + self._constvars = list(constvars) + self._invars = list(invars) + self._outvars = list(outvars) + self._eqns = list(eqns) + self._effects = effects + self._debug_info = debug_info + assert (not debug_info or len(debug_info.arg_names) == len(invars) and + len(debug_info.result_paths) == len(outvars)) + + def __str__(self): + return str(self.pretty_print()) + + __repr__ = __str__ + + def pretty_print(self, *, source_info=False, print_shapes=True, + custom_pp_eqn_rules=True, name_stack=False, + print_effects: bool = False, **kwargs): + doc = pp_toplevel_jaxpr( + self, source_info=source_info, print_shapes=print_shapes, + custom_pp_eqn_rules=custom_pp_eqn_rules, name_stack=name_stack, + print_effects=print_effects) + return doc.format(**kwargs) + + def _repr_pretty_(self, p, cycle): + return p.text(self.pretty_print(use_color=True)) + + def replace(self, **kwargs): + jaxpr = Jaxpr( + constvars=kwargs.pop("constvars", self.constvars), + invars=kwargs.pop("invars", self.invars), + outvars=kwargs.pop("outvars", self.outvars), + eqns=kwargs.pop("eqns", self.eqns), + effects=kwargs.pop("effects", self.effects), + debug_info=kwargs.pop("debug_info", self.debug_info), + ) + if kwargs: + raise ValueError(f"Unknown keyword arguments: {kwargs}") + return jaxpr + + +def join_effects(*effects: Effects) -> Effects: + return set().union(*effects) if effects else no_effects + +def jaxprs_in_params(params) -> Iterator[Jaxpr]: + for val in params.values(): + vals = val if isinstance(val, tuple) else (val,) + for v in vals: + if isinstance(v, Jaxpr): + yield v + elif isinstance(v, ClosedJaxpr): + yield v.jaxpr + + +def subjaxprs(jaxpr: Jaxpr) -> Iterator[Jaxpr]: + """Generator for all subjaxprs found in the params of jaxpr.eqns. + Does not descend recursively into the found subjaxprs. + """ + for eqn in jaxpr.eqns: + yield from jaxprs_in_params(eqn.params) + + +class ClosedJaxpr: + __slots__ = ['__weakref__', '_jaxpr', '_consts'] + + _jaxpr: Jaxpr + _consts: list[Any] + + jaxpr = property(lambda self: self._jaxpr) + consts = property(lambda self: self._consts) + + def __init__(self, jaxpr: Jaxpr, consts: Sequence): + assert len(consts) == len(jaxpr.constvars) + # assert not any(isinstance(c, Tracer) for c in consts) # TODO(mattjj): enable + self._jaxpr = jaxpr + self._consts = list(consts) + + @property + def in_avals(self): + return [v.aval for v in self.jaxpr.invars] + + @property + def out_avals(self): + return [v.aval for v in self.jaxpr.outvars] + + @property + def literals(self): + return self.consts # backwards compatible alias + + @property + def eqns(self): + return self.jaxpr.eqns + + @property + def effects(self) -> Effects: + return self.jaxpr.effects + + def map_jaxpr(self, f): + return ClosedJaxpr(f(self.jaxpr), self.consts) + + def replace(self, *, jaxpr=None, consts=None): + jaxpr = self.jaxpr if jaxpr is None else jaxpr + consts = self.consts if consts is None else consts + return ClosedJaxpr(jaxpr, consts) + + def __str__(self): return str(self.jaxpr) + def __repr__(self): return repr(self.jaxpr) + + def pretty_print(self, *, source_info=False, print_shapes=True, + name_stack=False, custom_pp_eqn_rules=True, + print_effects=False, **kwargs): + return self.jaxpr.pretty_print( + source_info=source_info, + print_shapes=print_shapes, + name_stack=name_stack, + custom_pp_eqn_rules=custom_pp_eqn_rules, + print_effects=print_effects, + **kwargs) + + def _repr_pretty_(self, p, cycle): + return p.text(self.pretty_print(use_color=True)) + +@curry +def jaxpr_as_fun(closed_jaxpr: ClosedJaxpr, *args): + return eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.consts, *args) + + +class JaxprEqnContext: + + def __init__(self, compute_type: str | None, threefry_partitionable: bool): + self.compute_type = compute_type + self.threefry_partitionable = threefry_partitionable + self._managers = [ + (compute_on.extend_compute_type, self.compute_type), + (config.threefry_partitionable.__call__, self.threefry_partitionable), + ] + + @property + @contextmanager + def manager(self): + with ExitStack() as stack: + for manager, val in self._managers: + stack.enter_context(manager(val)) + yield + + def __repr__(self): + return (f"JaxprEqnContext(compute_type={self.compute_type}," + f"threefry_partitionable={self.threefry_partitionable})") + + +class JaxprEqn(NamedTuple): + invars: list[Atom] + outvars: list[Var] + primitive: Primitive + params: dict[str, Any] + effects: Effects + source_info: source_info_util.SourceInfo + ctx: JaxprEqnContext + + def __repr__(self): + return str(pp_eqn(self, JaxprPpContext(), JaxprPpSettings())).rstrip() + + def replace( + self, + invars: list[Atom] | None = None, + outvars: list[Var] | None = None, + primitive: Primitive | None = None, + params: dict[str, Any] | None = None, + effects: Effects | None = None, + source_info: source_info_util.SourceInfo | None = None, + ctx: JaxprEqnContext | None = None + ): + # It is slightly faster to rebuild the tuple directly than to call _replace. + return JaxprEqn( + self.invars if invars is None else invars, + self.outvars if outvars is None else outvars, + self.primitive if primitive is None else primitive, + self.params if params is None else params, + self.effects if effects is None else effects, + self.source_info if source_info is None else source_info, + self.ctx if ctx is None else ctx, + ) + + +# TODO(mattjj): call typecheck rules here, so we don't form bad eqns +def new_jaxpr_eqn(invars, outvars, primitive, params, effects, source_info=None, + ctx=None): + source_info = source_info or source_info_util.new_source_info() + ctx = ctx or JaxprEqnContext(compute_on.current_compute_type(), + config.threefry_partitionable.value) + if config.enable_checks.value: + assert all(isinstance(x, (Var, Literal)) for x in invars) + assert all(isinstance(v, Var) for v in outvars) + return JaxprEqn(invars, outvars, primitive, params, effects, source_info, ctx) + +_var_counter = it.count() + +@total_ordering +class Var: + __slots__ = ["count", "suffix", "aval"] + + count: int + suffix: str + aval: AbstractValue + + def __init__(self, suffix: str, aval: AbstractValue): + self.count = next(_var_counter) + self.suffix = suffix + self.aval = raise_to_shaped(aval) + + # TODO(phawkins, mattjj): remove ordering of variables. JAX itself does not + # care about variable ordering, but the downstream package kfac_jax does. + def __lt__(self, other): + return self.count < other.count + + def __repr__(self): + return f'Var(id={id(self)}){self.suffix}:{self.aval.str_short()}' + + +def gensym(suffix: str = '') -> Callable[[AbstractValue], Var]: + """Produce distinct variables, printed with the optional suffix.""" + return partial(Var, suffix) + +# In a jaxpr, `dropvar` can appear in place of a bound variable to indicate that +# the assignment is dropped, i.e. that an expression's output value will never +# be read. In that sense, `dropvar` is not a variable, but it is convenient to +# treat it as a special case of one. Its `aval` is similarly inexact. +class DropVar(Var): + def __init__(self, aval: AbstractValue): + super().__init__('', aval) + def __repr__(self): return '_' + +class Literal: + __slots__ = ["val", "aval", "hash"] + + val: Any + aval: AbstractValue + hash: int | None + + def __init__(self, val, aval): + self.val = val + self.aval = aval + try: + self.hash = hash(val) + except TypeError: + if type(val) in literalable_types: + try: + self.hash = hash((val.item(), val.dtype)) + except (TypeError, AttributeError, ValueError): + self.hash = None + + __hash__ = None # type: ignore + + def __repr__(self): + if hasattr(self, 'hash'): + return f'{self.val}' + else: + return f'Literal(val={self.val})' + +literalable_types: set[type] = set() + +Atom = Union[Var, Literal] + +class Primitive: + name: str + # set for multi-output primitives. + multiple_results: bool = False + # set for call primitives processed in final style. + call_primitive: bool = False + # set for map primitives processed in final style. + map_primitive: bool = False + + def __init__(self, name: str): + self.name = name + + def __repr__(self): + return f'{self.name}' + + def bind(self, *args, **params): + assert (not config.enable_checks.value or + all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args + return self.bind_with_trace(find_top_trace(args), args, params) + + def bind_with_trace(self, trace, args, params): + with pop_level(trace.level): + out = trace.process_primitive(self, map(trace.full_raise, args), params) + return map(full_lower, out) if self.multiple_results else full_lower(out) + + def def_impl(self, impl): + self.impl = impl + return impl + + def def_abstract_eval(self, abstract_eval): + self.abstract_eval = _effect_free_abstract_eval(abstract_eval) + return abstract_eval + + def def_effectful_abstract_eval(self, effectful_abstract_eval): + self.abstract_eval = effectful_abstract_eval + return effectful_abstract_eval + + def def_custom_bind(self, bind): + self.bind = bind + return bind + + def impl(self, *args, **params): + raise NotImplementedError("Evaluation rule for '{}' not implemented" + .format(self.name)) + + def abstract_eval(self, *args, **params): + raise NotImplementedError("Abstract evaluation for '{}' not implemented" + .format(self.name)) + + def get_bind_params(self, params): + return [], params + + +def _effect_free_abstract_eval(abstract_eval): + def abstract_eval_(*args, **kwargs): + return abstract_eval(*args, **kwargs), no_effects + return abstract_eval_ + +# -------------------- lifting -------------------- + +# TODO(mattjj): replace this approach with a primitive-keyed table of rules +def traverse_jaxpr_params(f, params): + """Applies f to each jaxpr parameter and returns a tuple of returned values.""" + return {name: f(p) + for name, param in params.items() + for p in (param if isinstance(param, (tuple, list)) else [param]) + if type(p) in (Jaxpr, ClosedJaxpr)} + + +def eval_jaxpr(jaxpr: Jaxpr, consts, *args, propagate_source_info=True) -> list[Any]: + def read(v: Atom) -> Any: + return v.val if isinstance(v, Literal) else env[v] + + def write(v: Var, val: Any) -> None: + if config.enable_checks.value and not config.dynamic_shapes.value: + assert typecheck(v.aval, val), (v.aval, val) + env[v] = val + + env: dict[Var, Any] = {} + map(write, jaxpr.constvars, consts) + map(write, jaxpr.invars, args) + lu = last_used(jaxpr) + for eqn in jaxpr.eqns: + subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params) + name_stack = source_info_util.current_name_stack() + eqn.source_info.name_stack + traceback = eqn.source_info.traceback if propagate_source_info else None + with source_info_util.user_context( + traceback, name_stack=name_stack), eqn.ctx.manager: + ans = eqn.primitive.bind(*subfuns, *map(read, eqn.invars), **bind_params) + if eqn.primitive.multiple_results: + map(write, eqn.outvars, ans) + else: + write(eqn.outvars[0], ans) + clean_up_dead_vars(eqn, env, lu) + return map(read, jaxpr.outvars) + + +# -------------------- tracing -------------------- + +TracerType = TypeVar('TracerType', bound='Tracer') + +class Trace(Generic[TracerType]): + __slots__ = ['main', 'level', 'sublevel'] + + main: MainTrace + level: int + sublevel: Sublevel + + def __init__(self, main: MainTrace, sublevel: Sublevel) -> None: + self.main = main + self.level = main.level + self.sublevel = sublevel + + def full_raise(self, val) -> TracerType: + if not isinstance(val, Tracer): + # This check is only applied to non-Tracers, because the hasattr() is + # expensive (Tracer.__getattr__) in the common case that val is a Tracer. + if hasattr(val, "dimension_as_value"): # Used for shape_poly._DimExpr + val = val.dimension_as_value() + if not isinstance(val, Tracer): + return self.pure(val) + else: + return self.pure(val) + val._assert_live() + level = self.level + sublevel = self.sublevel + if val._trace.main is self.main: + if val._trace.sublevel == sublevel: + return cast(TracerType, val) + elif val._trace.sublevel < sublevel: + return self.sublift(val) + else: + raise escaped_tracer_error( + val, f"Can't lift sublevels {val._trace.sublevel} to {sublevel}") + elif val._trace.level < level: + if val._trace.sublevel > sublevel: + raise escaped_tracer_error( + val, f"Incompatible sublevel: {val._trace}, {(level, sublevel)}") + return self.lift(val) + elif val._trace.level > level: + raise escaped_tracer_error( + val, f"Can't lift level {val} to {self}") + else: # val._trace.level == self.level: + raise escaped_tracer_error( + val, f"Different traces at same level: {val}, {self}") + + def pure(self, val) -> TracerType: + raise NotImplementedError("must override") + + def lift(self, tracer) -> TracerType: + raise NotImplementedError("must override") + + def sublift(self, tracer) -> TracerType: + raise NotImplementedError("must override") + + def process_primitive(self, primitive, tracers, params): + raise NotImplementedError("must override") + + def __repr__(self): + return '{}(level={}/{})'.format( + self.__class__.__name__, self.level, self.sublevel) + + def process_call(self, call_primitive, f, tracers, params): + msg = (f"{type(self)} must override process_call to handle call-like " + "primitives") + raise NotImplementedError(msg) + + def process_map(self, map_primitive, f, tracers, params): + msg = (f"{type(self)} must override process_map to handle map-like " + "primitives") + raise NotImplementedError(msg) + + def process_custom_jvp_call(self, primitive, fun, jvp, tracers, *, + symbolic_zeros): + msg = (f"{type(self)} must override process_custom_jvp_call " + "to handle custom_jvp primitives") + raise NotImplementedError(msg) + + def process_custom_transpose(self, prim, call, tracers, **params): + msg = (f"{type(self)} must override process_custom_transpose " + "to handle custom_transpose_call primitives") + raise NotImplementedError(msg) + + def process_custom_vjp_call(self, primitive, fun, fwd, bwd, tracers, + out_trees, symbolic_zeros): + msg = (f"{type(self)} must override process_custom_vjp_call " + "to handle custom_vjp primitives") + raise NotImplementedError(msg) + + +def raise_as_much_as_possible(tracer) -> Tracer: + # Find effective bottom of trace stack (highest dynamic Trace on the stack). + trace_stack = thread_local_state.trace_state.trace_stack.stack + idx = next(i for i, m in enumerate(trace_stack) if m is + thread_local_state.trace_state.trace_stack.dynamic) + + # Only pay attention to effective part of trace stack. + trace_stack = trace_stack[idx:] + + # Lift tracer into everything in the effective stack higher than its level + for trace in trace_stack: + trace = trace.with_cur_sublevel() + if (not isinstance(tracer, Tracer) or tracer._trace.level < trace.level): + tracer = trace.full_raise(tracer) + + return tracer + + +def escaped_tracer_error(tracer, detail=None): + num_frames = _TRACER_ERROR_NUM_TRACEBACK_FRAMES.value + msg = ('Encountered an unexpected tracer. A function transformed by JAX ' + 'had a side effect, allowing for a reference to an intermediate value ' + f'with type {tracer.aval.str_short()} wrapped in a ' + f'{type(tracer).__name__} to escape the scope of the transformation.\n' + 'JAX transformations require that functions explicitly return their ' + 'outputs, and disallow saving intermediate values to global state.') + dbg = getattr(tracer, '_debug_info', None) + if dbg is not None: + msg += ('\nThe function being traced when the value leaked was ' + f'{dbg.func_src_info} traced for {dbg.traced_for}.') + line_info = getattr(tracer, '_line_info', None) + if line_info is not None: + divider = '\n' + '-'*30 + '\n' + msg += divider + msg += ('The leaked intermediate value was created on line ' + f'{source_info_util.summarize(line_info)}. ') + msg += divider + if num_frames > 0: + msg += (f'When the value was created, the final {num_frames} stack ' + 'frames (most recent last) excluding JAX-internal frames were:') + msg += divider + source_info_util.summarize( + line_info, num_frames=num_frames) + divider + msg += ('\nTo catch the leak earlier, try setting the environment variable ' + 'JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context ' + 'manager.') + if detail: + msg += f'Detail: {detail}' + return UnexpectedTracerError(msg) + + +def check_scalar_conversion(arr: Array): + if arr.ndim > 0: + raise TypeError("Only scalar arrays can be converted to Python scalars; " + f"got {arr.ndim=}") + + +def check_integer_conversion(arr: Array): + if not (arr.shape == () and dtypes.issubdtype(arr.dtype, np.integer)): + raise TypeError("Only integer scalar arrays can be converted to a scalar index.") + + +def check_bool_conversion(arr: Array): + if arr.size == 0: + raise ValueError("The truth value of an empty array is ambiguous. Use" + " `array.size > 0` to check that an array is not empty.") + if arr.size > 1: + raise ValueError("The truth value of an array with more than one element" + " is ambiguous. Use a.any() or a.all()") + + +def _aval_property(name): + return property(lambda self: getattr(self.aval, name)) + + +class Tracer(typing.Array, metaclass=StrictABCMeta): + __array_priority__ = 1000 + __slots__ = ['_trace', '_line_info'] + + dtype = _aval_property('dtype') + ndim = _aval_property('ndim') + size = _aval_property('size') + shape = _aval_property('shape') + + def __init__(self, trace: Trace): + self._trace = trace + + def _error_repr(self): + if self.aval is None: + return f"traced array with aval {self.aval}" + return f"traced array with shape {raise_to_shaped(self.aval).str_short()}." + + def __array__(self, *args, **kw): + raise TracerArrayConversionError(self) + + def __dlpack__(self, *args, **kw): + raise ConcretizationTypeError(self, + f"The __dlpack__() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def tolist(self): + raise ConcretizationTypeError(self, + f"The tolist() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def tobytes(self, order="C"): + del order + raise ConcretizationTypeError(self, + f"The tobytes() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def __iter__(self): + return iter(self.aval._iter(self)) + + def __reversed__(self): + return iter(self[::-1]) + + def __len__(self): + return self.aval._len(self) + + @property + def sharding(self): + # This attribute is part of the jax.Array API, but only defined on concrete arrays. + # Raising a ConcretizationTypeError would make sense, but for backward compatibility + # we raise an AttributeError so that hasattr() and getattr() work as expected. + try: + orig_msg = self._origin_msg() + except: + orig_msg = '' + raise AttributeError(self, + f"The 'sharding' attribute is not available on {self._error_repr()}." + f"{orig_msg}") + + @property + def addressable_shards(self): + raise ConcretizationTypeError(self, + f"The 'addressable_shards' attribute is not available on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def at(self): + return self.aval.at.fget(self) + + @property + def aval(self): + raise NotImplementedError("must override") + + def _assert_live(self) -> None: + pass # Override for liveness checking + + def get_referent(self) -> Any: + return self # Override for object equivalence checking + + def __bool__(self): + check_bool_conversion(self) + return self.aval._bool(self) + + def __int__(self): + check_scalar_conversion(self) + return self.aval._int(self) + + def __float__(self): + check_scalar_conversion(self) + return self.aval._float(self) + + def __complex__(self): + check_scalar_conversion(self) + return self.aval._complex(self) + + def __hex__(self): + check_integer_conversion(self) + return self.aval._hex(self) + + def __oct__(self): + check_integer_conversion(self) + return self.aval._oct(self) + + def __index__(self): + check_integer_conversion(self) + raise self.aval._index(self) + + # raises a useful error on attempts to pickle a Tracer. + def __reduce__(self): + raise ConcretizationTypeError( + self, ("The error occurred in the __reduce__ method, which may " + "indicate an attempt to serialize/pickle a traced value.")) + + # raises the better error message from ShapedArray + def __setitem__(self, idx, val): return self.aval._setitem(self, idx, val) + + # NumPy also only looks up special methods on classes. + def __array_module__(self, types): return self.aval._array_module(self, types) + + def __getattr__(self, name): + # if the aval property raises an AttributeError, gets caught here + assert not config.enable_checks.value or name != "aval" + + try: + attr = getattr(self.aval, name) + except AttributeError as err: + raise AttributeError( + f"{self.__class__.__name__} has no attribute {name}" + ) from err + else: + t = type(attr) + if t is aval_property: + return attr.fget(self) + elif t is aval_method: + return types.MethodType(attr.fun, self) + else: + return attr + + def _pretty_print(self): + base = pp.text(f'Traced<{self.aval}>with<{self._trace}>') + contents = [(name, attr._pretty_print() if isinstance(attr, Tracer) + else pp.text(repr(attr))) for name, attr in self._contents()] + if contents: + base = pp.group(pp.nest(2, pp.concat([ + base, pp.text(' with'), pp.brk(), pp.join(pp.brk(), [ + pp.text(f'{name} = ') + pp_payload + for name, pp_payload in contents]) + ]))) + return base + + def __repr__(self): + return self._pretty_print().format() + + def _contents(self): + try: + return [(name, getattr(self, name)) for name in self.__slots__] + except AttributeError: + return () + + def _origin_msg(self) -> str: + return "" + + # Methods that are only valid for materialized arrays + def addressable_data(self, index): + raise ConcretizationTypeError(self, + f"The addressable_data() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def block_until_ready(self): + # Raise AttributeError for backward compatibility with hasattr() and getattr() checks. + raise AttributeError(self, + f"The 'block_until_ready' method is not available on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def copy_to_host_async(self): + # Raise AttributeError for backward compatibility with hasattr() and getattr() checks. + raise AttributeError(self, + f"The 'copy_to_host_async' method is not available on {self._error_repr()}." + f"{self._origin_msg()}") + + def delete(self): + raise ConcretizationTypeError(self, + f"The delete() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def devices(self): + raise ConcretizationTypeError(self, + f"The devices() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def global_shards(self): + raise ConcretizationTypeError(self, + f"The global_shards property was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def is_deleted(self): + raise ConcretizationTypeError(self, + f"The is_deleted() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def is_fully_addressable(self): + raise ConcretizationTypeError(self, + f"The is_fully_addressable property was called on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def is_fully_replicated(self): + raise ConcretizationTypeError(self, + f"The is_fully_replicated property was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def on_device_size_in_bytes(self): + raise ConcretizationTypeError(self, + f"The on_device_size_in_bytes() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + + @property + def traceback(self): + raise ConcretizationTypeError(self, + f"The traceback property was called on {self._error_repr()}." + f"{self._origin_msg()}") + + def unsafe_buffer_pointer(self): + raise ConcretizationTypeError(self, + f"The unsafe_buffer_pointer() method was called on {self._error_repr()}." + f"{self._origin_msg()}") + +# these can be used to set up forwarding of properties and instance methods from +# Tracer instances to the underlying avals +aval_property = namedtuple("aval_property", ["fget"]) +aval_method = namedtuple("aval_method", ["fun"]) + + +class EvalTrace(Trace): + # See comments in https://github.com/google/jax/pull/3370 + def pure(self, x): return x + lift = sublift = pure + + def process_primitive(self, primitive, tracers, params): + if config.debug_key_reuse.value: + # Import here to avoid circular imports + from jax.experimental.key_reuse._core import call_impl_with_key_reuse_checks # pytype: disable=import-error + return call_impl_with_key_reuse_checks(primitive, primitive.impl, *tracers, **params) + else: + return primitive.impl(*tracers, **params) + + def process_call(self, primitive, f, tracers, params): + if config.debug_key_reuse.value: + # Import here to avoid circular imports + from jax.experimental.key_reuse._core import call_impl_with_key_reuse_checks # pytype: disable=import-error + return call_impl_with_key_reuse_checks(primitive, primitive.impl, f, *tracers, **params) + else: + return primitive.impl(f, *tracers, **params) + process_map = process_call + + def process_custom_transpose(self, primitive, call, tracers, **_): + del primitive, _ + with new_sublevel(): + return call.call_wrapped(*tracers) + + def process_custom_jvp_call(self, primitive, fun, jvp, tracers, **_): + del primitive, jvp, _ # Unused. + with new_sublevel(): + return fun.call_wrapped(*tracers) + + def process_custom_vjp_call(self, primitive, fun, fwd, bwd, tracers, **_): # pytype: disable=signature-mismatch + del primitive, fwd, bwd, _ # Unused. + with new_sublevel(): + return fun.call_wrapped(*tracers) + + +class MainTrace: + level: int + trace_type: type[Trace] + payload: dict[str, Any] + + def __init__(self, level, trace_type, **payload) -> None: + self.level = level + self.trace_type = trace_type + self.payload = payload + + def __repr__(self) -> str: + return f"MainTrace({self.level},{self.trace_type.__name__})" + + def __hash__(self) -> int: + return hash((self.level, self.trace_type)) + + def __eq__(self, other: object) -> bool: + return (isinstance(other, MainTrace) and + self.level == other.level and + self.trace_type == other.trace_type and + self.payload == other.payload) + + def with_cur_sublevel(self): + return self.trace_type(self, cur_sublevel(), **self.payload) + +class TraceStack: + # See comments in https://github.com/google/jax/pull/3370 + stack: list[MainTrace] + dynamic: MainTrace + + def __init__(self): + eval_trace = MainTrace(0, EvalTrace) + self.stack = [eval_trace] + self.dynamic = eval_trace + + def next_level(self) -> int: + return len(self.stack) + + def push(self, main_trace: MainTrace) -> None: + self.stack.append(main_trace) + + def pop(self) -> None: + self.stack.pop() + + def __repr__(self) -> str: + stack_str = map(' {}\n'.format, self.stack[::-1]) + return f'Trace stack\n{stack_str}\n{self.dynamic}' + + def copy(self): + new = self.__new__(TraceStack) + new.stack = self.stack[:] + new.dynamic = self.dynamic + return new + + +@total_ordering +class Sublevel: + + def __init__(self, level: int): + self.level = level + + def __repr__(self): + return str(self.level) + + def __eq__(self, other): + return type(other) is Sublevel and self.level == other.level + + def __lt__(self, other): + return type(other) is Sublevel and self.level < other.level + + +AxisEnvFrame = namedtuple('AxisEnvFrame', ['name', 'size', 'main_trace']) +AxisName = Hashable + +no_axis_name = object() + +class TraceState: + trace_stack: TraceStack + substack: list[Sublevel] + axis_env: list[AxisEnvFrame] + + def __init__(self) -> None: + self.trace_stack = TraceStack() + self.substack = [Sublevel(0)] + self.axis_env = [] + + def copy(self): + new = self.__new__(TraceState) + new.trace_stack = self.trace_stack.copy() + new.substack = self.substack[:] + new.axis_env = self.axis_env[:] + return new + + +def _update_thread_local_jit_state(dynamic): + # Copies the MainTrace instance, removing any .debug_info or .jaxpr_stack + # fields that should not be kept alive as part of a cache key. + # TODO(mattjj): split debug_info and jaxpr_stack out of MainTrace. + # TODO(mattjj): add a test that verifies that JIT-ted functions are not kept + # alive by the JIT cache, particularly for nested JIT-ted functions. + copy = MainTrace(dynamic.level, dynamic.trace_type, **dynamic.payload) + config.update_thread_local_jit_state(dynamic_trace_state=copy) + + +# The global state of the tracer is accessed by a thread-local object. +# This allows concurrent tracing in separate threads; passing traced objects +# between threads is forbidden. +class ThreadLocalState(threading.local): + def __init__(self): + self.trace_state = TraceState() + +thread_local_state = ThreadLocalState() + + +def _initialize_jax_jit_thread_local_state(): + """Initializes the C++ thread-local context. + + When the user spawns threads, the C++ `jax_jit.thread_local_state` is None. + The C++ accessor calls this function if it realizes the thread_local_state + is None (which means it's not yet initialized for this thread). + + This function does not live in `config.py`, to prevent circular imports. + """ + tls = jax_jit.thread_local_state() + if tls.extra_jit_context is None: + dynamic = thread_local_state.trace_state.trace_stack.dynamic + copy = MainTrace(dynamic.level, dynamic.trace_type, **dynamic.payload) + config.update_thread_local_jit_state(dynamic_trace_state=copy) + + +jax_jit.set_thread_local_state_initialization_callback( + _initialize_jax_jit_thread_local_state) + +def trace_state_clean() -> bool: + trace_state = thread_local_state.trace_state + return (trace_state.substack == [Sublevel(0)] and + trace_state.axis_env == [] and + trace_state.trace_stack.stack == [MainTrace(0, EvalTrace)] and + trace_state.trace_stack.dynamic == MainTrace(0, EvalTrace)) + +def reset_trace_state() -> bool: + """Resets the global trace state and returns True if it was already clean.""" + if not trace_state_clean(): + thread_local_state.trace_state.__init__() + return False + else: + return True + +def cur_sublevel() -> Sublevel: + return thread_local_state.trace_state.substack[-1] + +TRACER_LEAK_DEBUGGER_WARNING = """\ +JAX check_tracer_leaks behavior can trigger false positives when used with a debugger. +To avoid false positives and silence this warning, you can disable thread tracing using +the following: + + import threading + threading.current_thread().pydev_do_not_trace = True +""" + +def maybe_find_leaked_tracers(x: MainTrace | Sublevel | None + ) -> list[Tracer]: + """Find the leaked tracers holding a reference to the MainTrace or SubLevel. + + It's possible there's none! eg. there's some cases where JAX itself holds a + reference to `x` inside of a lambda closure, and no tracers were leaked + by the user. In this case an empty list is returned. + """ + if not getattr(threading.current_thread(), 'pydev_do_not_trace', True): + warnings.warn(TRACER_LEAK_DEBUGGER_WARNING) + # Trigger garbage collection to filter out unreachable objects that are alive + # only due to cyclical dependencies. (We don't care about unreachable leaked + # tracers since they can't interact with user code and cause a problem.) + gc.collect() + traces = list(filter(lambda x: isinstance(x, Trace), gc.get_referrers(x))) + tracers = list(filter(lambda x: isinstance(x, Tracer), gc.get_referrers(*traces))) + return tracers + +def leaked_tracer_error(name: str, t, tracers: list[Tracer]) -> Exception: + assert tracers + why = partial(_why_alive, {id(tracers)}) + msgs = '\n\n'.join(f'{tracers[i]}{tracers[i]._origin_msg()}{why(tracers[i])}' + for i in range(len(tracers))) + return Exception(f'Leaked {name} {t}. Leaked tracer(s):\n\n{msgs}\n') + +def _why_alive(ignore_ids: set[int], x: Any) -> str: + parents = lambda x: [r for r in gc.get_referrers(x) if id(r) not in ignore_ids] + child, lines, seen = x, [], set() + while (id(child) not in seen and type(child) is not types.ModuleType + and parents(child)): + parent = parents(child)[0] # just pick one parent + + # For namespaces (like modules and class instances) and closures, the + # references may form a simple chain: e.g. instance refers to its own + # __dict__ which refers to child, or function refers to its __closure__ + # which refers to cells which refer to child. In these cases, we can provide + # a more intuitive description by collapsing the chain into a single + # parent->child jump. We do that by setting `parent` here to be a + # grandparent (or great-grandparent) of `child`, and then handling that case + # in _why_alive_container_info. See example: + # https://github.com/google/jax/pull/13022#discussion_r1008456599 + # To prevent this collapsing behavior, just comment out this code block. + if (isinstance(parent, dict) and + getattr(parents(parent)[0], '__dict__', None) is parents(child)[0]): + parent = parents(parent)[0] + elif type(parent) is types.CellType: + parent = parents(parents(parent)[0])[0] + + line = f'<{type(child).__name__} {id(child)}> is referred to by ' + lines.append(line + _why_alive_container_info(parent, id(child))) + seen.add(id(child)) + child = parent + return '\n' + '\n'.join(lines) if lines else '' + +def _why_alive_container_info(container, obj_id) -> str: + name = f'<{type(container).__name__} {id(container)}>' + if type(container) is types.ModuleType: + name = getattr(container, '__name__', name) + if type(container) is types.FunctionType: + name_ = getattr(container, '__name__', '') + closure = inspect.getclosurevars(container) + keys = [k for k, v in dict(closure.nonlocals, **closure.globals).items() + if id(v) == obj_id] + if len(keys) == 1: return f'{name} ({name_}) closed-over variable {keys[0]}' + elif len(keys) > 1: return (f'{name} in closed-over variables ' + + ', '.join(map(repr, keys))) + if hasattr(container, '__dict__'): + keys = [k for k in vars(container) if id(vars(container)[k]) == obj_id] + if len(keys) == 1: return f'{name}.{keys[0]}' + elif len(keys) > 1: return f'{name} in vars ' + ', '.join(map(repr, keys)) + if isinstance(container, (list, tuple)): + idxs = [i for i, x in enumerate(container) if id(x) == obj_id] + if len(idxs) == 1: return f'{name}[{idxs[0]}]' + else: return f'{name} at indices ' + ', '.join(map(str, idxs)) + if isinstance(container, dict): + keys = [k for k in container if id(container[k]) == obj_id] + if len(keys) == 1: return f'{name}[{keys[0]!r}]' + else: return f'{name} at keys ' + ', '.join(map(repr, keys)) + if isinstance(container, types.ModuleType): + return f' named {container.__name__}' + return name + + +@contextmanager +def new_main(trace_type: type[Trace], dynamic: bool = False, + **payload) -> Generator[MainTrace, None, None]: + # See comments in https://github.com/google/jax/pull/3370 + stack = thread_local_state.trace_state.trace_stack + level = stack.next_level() + main = MainTrace(level, trace_type, **payload) + stack.push(main) + if dynamic: + prev_dynamic, stack.dynamic = stack.dynamic, main + _update_thread_local_jit_state(stack.dynamic) + + try: + yield main + finally: + stack.pop() + if dynamic: + stack.dynamic = prev_dynamic + _update_thread_local_jit_state(stack.dynamic) + + if config.check_tracer_leaks.value: + t = ref(main) + del main + if t() is not None: + leaked_tracers = maybe_find_leaked_tracers(t()) + if leaked_tracers: raise leaked_tracer_error("trace", t(), leaked_tracers) + +@contextmanager +def new_dynamic(level: int) -> Generator[None, None, None]: + stack = thread_local_state.trace_state.trace_stack + prev_dynamic, stack.dynamic = stack.dynamic, stack.stack[level] + _update_thread_local_jit_state(stack.dynamic) + try: + yield + finally: + stack.dynamic = prev_dynamic + _update_thread_local_jit_state(stack.dynamic) + +def dynamic_level() -> int: + return thread_local_state.trace_state.trace_stack.dynamic.level + +@contextmanager +def new_base_main(trace_type: type[Trace], + **payload) -> Generator[MainTrace, None, None]: + # See comments in https://github.com/google/jax/pull/3370 + stack = thread_local_state.trace_state.trace_stack + main = MainTrace(0, trace_type, **payload) + prev_dynamic, stack.dynamic = stack.dynamic, main + prev_base, stack.stack[0] = stack.stack[0], main + _update_thread_local_jit_state(stack.dynamic) + try: + yield main + finally: + stack.dynamic = prev_dynamic + stack.stack[0] = prev_base + _update_thread_local_jit_state(stack.dynamic) + + if config.check_tracer_leaks.value: + t = ref(main) + del main + if t() is not None: + leaked_tracers = maybe_find_leaked_tracers(t()) + if leaked_tracers: raise leaked_tracer_error("trace", t(), leaked_tracers) + +@contextmanager +def pop_level(level: int): + if level == 0: + return (yield) + prev, thread_local_state.trace_state.trace_stack.stack = \ + thread_local_state.trace_state.trace_stack.stack, \ + thread_local_state.trace_state.trace_stack.stack[:level] + try: + yield + finally: + thread_local_state.trace_state.trace_stack.stack = prev + +@contextmanager +def ensure_compile_time_eval(): + """Context manager to ensure evaluation at trace/compile time (or error). + + Some JAX APIs like :func:`jax.jit` and :func:`jax.lax.scan` involve staging, + i.e., delaying the evaluation of numerical expressions (like :mod:`jax.numpy` + function applications) so that instead of performing those computations + eagerly while evaluating the corresponding Python expressions, their + computation is carried out separately, e.g. after optimized compilation. But + this delay can be undesirable. For example, numerical values might be needed + to evaluate Python control flow and so their evaluation cannot be delayed. As + another example, it may be beneficial to ensure compile time evaluation (or + "constant folding") for performance reasons. + + This context manager ensures that JAX computations are evaluated eagerly. If + eager evaluation is not possible, a ``ConcretizationTypeError`` is raised. + + Here's a contrived example:: + + import jax + import jax.numpy as jnp + + @jax.jit + def f(x): + with jax.ensure_compile_time_eval(): + y = jnp.sin(3.0) + z = jnp.sin(y) + z_positive = z > 0 + if z_positive: # z_positive is usable in Python control flow + return jnp.sin(x) + else: + return jnp.cos(x) + + Here's a real-world example from https://github.com/google/jax/issues/3974:: + + import jax + import jax.numpy as jnp + from jax import random + + @jax.jit + def jax_fn(x): + with jax.ensure_compile_time_eval(): + y = random.randint(random.key(0), (1000,1000), 0, 100) + y2 = y @ y + x2 = jnp.sum(y2) * x + return x2 + + A similar behavior can often be achieved simply by 'hoisting' the constant + expression out of the corresponding staging API:: + + y = random.randint(random.key(0), (1000,1000), 0, 100) + + @jax.jit + def jax_fn(x): + y2 = y @ y + x2 = jnp.sum(y2)*x + return x2 + + But in some cases it can be more convenient to use this context manager. + """ + with new_base_main(EvalTrace): + yield +eval_context = ensure_compile_time_eval # alias, backward compatibility + +@contextmanager +def new_sublevel() -> Generator[None, None, None]: + sublevel = Sublevel(len(thread_local_state.trace_state.substack)) + thread_local_state.trace_state.substack.append(sublevel) + try: + yield + finally: + thread_local_state.trace_state.substack.pop() + + if config.check_tracer_leaks.value: + t = ref(sublevel) + del sublevel + if t() is not None: + leaked_tracers = maybe_find_leaked_tracers(t()) + if leaked_tracers: + raise leaked_tracer_error("sublevel", t(), leaked_tracers) + +def full_lower(val): + if isinstance(val, Tracer): + return val.full_lower() + else: + return val + + +def _get_trace_level(t: Tracer) -> int: return t._trace.level + + +def find_top_trace(xs) -> Trace: + top_tracer = max((x for x in xs if isinstance(x, Tracer)), + default=None, key=_get_trace_level) + if top_tracer is not None: + top_tracer._assert_live() + top_main = top_tracer._trace.main + else: + top_main = None + dynamic = thread_local_state.trace_state.trace_stack.dynamic + top_main = (dynamic if top_main is None or dynamic.level > top_main.level + else top_main) + return top_main.with_cur_sublevel() + +def get_referent(x: Any) -> Any: + return x.get_referent() if isinstance(x, Tracer) else x + +def same_referent(x: Any, y: Any) -> bool: + return get_referent(x) is get_referent(y) + +def dedup_referents(itr: Iterable[Any]) -> list[Any]: + return list({HashableWrapper(get_referent(x)):x for x in itr}.values()) + +def definitely_equal(x, y): + if isinstance(x, Tracer) or isinstance(y, Tracer): + return same_referent(x, y) + elif x is y: + return True + try: + return x == y + except InconclusiveDimensionOperation: + return False + +# -------------------- abstract values -------------------- + +class AbstractValue: + __slots__: list[str] = [] + + def at_least_vspace(self): + raise NotImplementedError("must override") + + def __repr__(self): + try: + kv_pairs = (f'{k}={v}' for k, v in self.__dict__.items()) + return '{}({})'.format(self.__class__.__name__, ','.join(kv_pairs)) + except AttributeError: + return self.__class__.__name__ + + def strip_weak_type(self) -> AbstractValue: + return self + + def strip_named_shape(self) -> AbstractValue: + return self + + def join(self, other): + raise NotImplementedError("must override") + + def update(self, **kwargs): + raise NotImplementedError("must override") + + def str_short(self, short_dtypes=False): + return str(self) + + +# For type signatures involving dynamic shapes, we use lists of abstract values +# which may contain (reverse) de Bruijn indices in their shapes. +class DBIdx(NamedTuple): + val: int + +@dataclass(frozen=True) +class InDBIdx: + val: int + +@dataclass(frozen=True) +class OutDBIdx: + val: int + +# For annotating input types of callables (i.e. linear_util.WrappedFuns), we use +# a sequence of pairs where the first element of each pair is an AbstractValue +# (possibly containing DBIdx instances in its shape) and the second is a boolean +# indicating whether that argument is explicit (i.e. passed to the callable). +InputType = tuple[tuple[AbstractValue, bool], ...] # DBIdx in shapes + +# For annotating jaxpr output types, we use a sequence of pairs where the first +# element of each pair is an AbstractValue (possibly containing InDBIdx and/or +# OutDBIdx instances in its shape) and the second is a boolean indicating +# whether that argument is explicit (i.e. returned by the callable). +OutputType = tuple[tuple[AbstractValue, bool], ...] # InDBIdx / OutDBIdx shapes + + +def _jaxpr_type_to_callable_annotation(jaxpr: Jaxpr) -> InputType: + idxs = {v: DBIdx(i) for i, v in enumerate((*jaxpr.constvars, *jaxpr.invars))} + out = [(v.aval.update(shape=tuple(idxs.get(d, d) for d in v.aval.shape)) # type: ignore + if type(v.aval) is DShapedArray else v.aval, True) + for v in jaxpr.invars] + return tuple(out) + +class Bot(AbstractValue): pass +bot = Bot() + + +def lattice_join(x: AbstractValue | None, + y: AbstractValue | None) -> AbstractValue: + if x is None: + return cast(AbstractValue, y) + elif y is None: + return cast(AbstractValue, x) + elif isinstance(x, type(y)): + return y.join(x) + elif isinstance(y, type(x)): + return x.join(y) + elif isinstance(x, DShapedArray) and isinstance(y, ShapedArray): + # TODO(mattjj): remove this special case after dynamic shapes are integrated + return x.join(y) + else: + raise TypeError(x, y) + +# For use in typing annotations to denote either a Tracer or a `valid_jaxtype`. +Value = Any + +def valid_jaxtype(x) -> bool: + try: + concrete_aval(x) + except TypeError: + return False + else: + return True + +def check_valid_jaxtype(x): + if not valid_jaxtype(x): + raise TypeError( + f"Value {x!r} of type {type(x)} is not a valid JAX type") + + +def concrete_aval(x): + for typ in type(x).__mro__: + handler = pytype_aval_mappings.get(typ) + if handler: return handler(x) + if hasattr(x, '__jax_array__'): + return concrete_aval(x.__jax_array__()) + raise TypeError(f"Value {x!r} with type {type(x)} is not a valid JAX " + "type") + + +def get_aval(x): + if isinstance(x, Tracer): + return x.aval + else: + return concrete_aval(x) + + +def concretization_function_error(fun, suggest_astype=False): + fname = getattr(fun, "__name__", fun) + fname_context = f"The problem arose with the `{fname}` function. " + if suggest_astype: + fname_context += ("If trying to convert the data type of a value, " + f"try using `x.astype({fun.__name__})` " + f"or `jnp.array(x, {fun.__name__})` instead.") + if fun is bool: + def error(self, arg): + raise TracerBoolConversionError(arg) + elif fun in (hex, oct, operator.index): + def error(self, arg): + raise TracerIntegerConversionError(arg) + else: + def error(self, arg): + raise ConcretizationTypeError(arg, fname_context) + return error + +def concrete_or_error(force: Any, val: Any, context=""): + """Like force(val), but gives the context in the error message.""" + if force is None: + force = lambda x: x + if isinstance(val, Tracer): + if isinstance(val.aval, ConcreteArray): + return force(val.aval.val) + else: + raise ConcretizationTypeError(val, context) + else: + return force(val) + +def concrete_dim_or_error(val: Any, context=""): + """Like concrete_or_error(operator.index), allowing symbolic dimensions.""" + if is_dim(val): + return val + else: + return concrete_or_error(operator.index, val, context=context) + +### Extended dtypes +# +# Extended dtypes are JAX-specific dtypes that allow us to represent logical +# arrays of element types that do not have an obvious direct correspondence +# to ("physical") arrays of basic types in a compiler. In particular, their +# element types differ from those of XLA and NumPy (e.g. int32). These dtypes +# are only known to JAX. Their implementation is determined by: +# a) an object representing the extended dtype, accessible via the `dtype` +# attribute on corresponding JAX arrays and, internally, on avals such +# as ShapedArrays that correspond to such JAX arrays; +# b) a set of rules, available via a private attribute on the extended dtype +# object in (a). +# The rules in (b) tell JAX internals how to ground out the element +# type for interaction with the compiler and runtime, e.g. when lowering +# to the compiler's language. + +@overload +def physical_aval(aval: ShapedArray) -> ShapedArray: ... +@overload +def physical_aval(aval: DShapedArray) -> DShapedArray: ... +@overload # TODO(frostig): remove this case +def physical_aval(aval: AbstractValue) -> AbstractValue: ... + +def physical_aval(aval): + aval_dtype = getattr(aval, 'dtype', None) + if aval_dtype and isinstance(aval_dtype, dtypes.ExtendedDType): + ctor = type(aval) + aval_shape = getattr(aval, 'shape', None) + assert aval_shape is not None, (ctor, aval) + elt_aval = aval_dtype._rules.physical_element_aval(aval_dtype) + assert type(elt_aval) is ShapedArray + return ctor((*aval_shape, *elt_aval.shape), elt_aval.dtype) # pytype: disable=wrong-arg-count + else: + return aval + +def _short_dtype_name(dtype) -> str: + if isinstance(dtype, dtypes.ExtendedDType): + return str(dtype) + else: + return (dtype.name.replace('float', 'f').replace('uint' , 'u') + .replace('int' , 'i').replace('complex', 'c')) + +def _dtype_object(dtype): + return dtype if isinstance(dtype, dtypes.ExtendedDType) else np.dtype(dtype) + +class UnshapedArray(AbstractValue): + __slots__ = ['dtype', 'weak_type'] + array_abstraction_level = 4 + + def __init__(self, dtype, weak_type=False): + self.dtype = _dtype_object(dtype) + self.weak_type = weak_type + + def update(self, dtype=None, weak_type=None): + if dtype is None: + dtype = self.dtype + if weak_type is None: + weak_type = self.weak_type + return UnshapedArray(dtype, weak_type) + + def __eq__(self, other): + return (type(self) is type(other) and self.dtype == other.dtype and + self.weak_type == other.weak_type) + + def __ne__(self, other): + return not self == other + + def __hash__(self): + # can use hash(self.dtype) and rely on the fact that numpy reuses base dtype + # objects, e.g. `np.zeros(3).dtype is np.zeros(4).dtype`, or we can use + # the unique character code via hash(self.dtype.char) + return hash((self.dtype, self.weak_type)) + + def __repr__(self): + return '{}({}{})'.format(self.__class__.__name__, self.str_short(), + ", weak_type=True" if self.weak_type else "") + + _bool = concretization_function_error(bool) + _int = concretization_function_error(int, True) + _float = concretization_function_error(float, True) + _complex = concretization_function_error(complex, True) + _hex = concretization_function_error(hex) + _oct = concretization_function_error(oct) + _index = concretization_function_error(operator.index) + + def at_least_vspace(self) -> AbstractValue: + return UnshapedArray(primal_dtype_to_tangent_dtype(self.dtype), + self.weak_type) + + def join(self, other): + if self.dtype == other.dtype: + if self.weak_type == other.weak_type: + return self + else: + return UnshapedArray(self.dtype, weak_type=False) + else: + raise TypeError(self, other) + + def str_short(self, short_dtypes=False) -> str: + return _short_dtype_name(self.dtype) if short_dtypes else self.dtype.name + + def strip_weak_type(self): + """Returns a copy of the aval with weak_type=False.""" + return self.update(weak_type=False) + + @property + def shape(self): + msg = ("UnshapedArray has no shape. Please open an issue at " + "https://github.com/google/jax/issues because it's unexpected for " + "UnshapedArray instances to ever be produced.") + raise TypeError(msg) + +def _canonicalize_dimension(dim: DimSize) -> DimSize: + # Dimensions are most commonly integral (by far), so we check that first. + try: + return operator.index(dim) + except TypeError as e: + type_error = e + if isinstance(dim, Tracer) and config.dynamic_shapes.value: + if not (dim.ndim == 0 and (dtypes.issubdtype(dim.dtype, np.integer) + or isinstance(dim.dtype, bint))): + raise TypeError(f"Dimensions must be integer scalars; got {dim.ndim=} {dim.dtype=}") + return dim + elif (config.dynamic_shapes.value and isinstance(dim, DArray) and + type(dim._aval.dtype) is bint and not dim._aval.shape): + return dim + elif is_dim(dim): + return dim + else: + raise type_error + +def canonicalize_shape(shape: Shape, context: str="") -> tuple[Any, ...]: + """Canonicalizes and checks for errors in a user-provided shape value. + + Args: + shape: a Python value that represents a shape. + + Returns: + A tuple of canonical dimension values. + """ + try: + return tuple(unsafe_map(_canonicalize_dimension, shape)) + except TypeError: + pass + raise _invalid_shape_error(shape, context) + +def canonicalize_dim(d: DimSize, context: str="") -> DimSize: + """Canonicalizes and checks for errors in a user-provided shape dimension value. + + Args: + f: a Python value that represents a dimension. + + Returns: + A canonical dimension value. + """ + return canonicalize_shape((d,), context)[0] + +def _invalid_shape_error(shape: Shape, context: str=""): + if config.dynamic_shapes.value: + msg = ("Shapes must be 1D sequences of integer scalars, " + f"got {shape}") + else: + msg = ("Shapes must be 1D sequences of concrete values of integer type, " + f"got {shape}.") + if context: + msg += f" {context}." + if not config.dynamic_shapes.value and any( + isinstance(x, Tracer) and isinstance(get_aval(x), ShapedArray) + and not isinstance(get_aval(x), ConcreteArray) for x in shape): + msg += ("\nIf using `jit`, try using `static_argnums` or applying `jit` to " + "smaller subfunctions.") + for x in shape: + if isinstance(x, Tracer) and hasattr(x, "_origin_msg"): + msg += x._origin_msg() + + return TypeError(msg) + +class ShapedArray(UnshapedArray): + __slots__ = ['shape', 'named_shape'] + array_abstraction_level = 2 + + def __init__(self, shape, dtype, weak_type=False, named_shape=None): + self.shape = canonicalize_shape(shape) + self.dtype = _dtype_object(dtype) + self.weak_type = weak_type + self.named_shape = {} if named_shape is None else dict(named_shape) + + def update(self, shape=None, dtype=None, weak_type=None, named_shape=None): + if shape is None: + shape = self.shape + if dtype is None: + dtype = self.dtype + if weak_type is None: + weak_type = self.weak_type + if named_shape is None: + named_shape = self.named_shape + return ShapedArray(shape, dtype, weak_type, named_shape) + + ndim = property(lambda self: len(self.shape)) + size = property(lambda self: + 0 if any(type(d) is int and d == 0 for d in self.shape) + else math.prod(self.shape)) + + broadcast: ClassVar[aval_method | None] = None + transpose: ClassVar[aval_method | None] = None + reshape: ClassVar[aval_method | None] = None + _iter: ClassVar[staticmethod | None] = None + + def __eq__(self, other): + return (type(self) is type(other) + and self.dtype == other.dtype and self.shape == other.shape + and self.weak_type == other.weak_type + and self.named_shape == other.named_shape) + + def __hash__(self): + # can use hash(self.dtype) and rely on the fact that numpy reuses base dtype + # objects, e.g. `np.zeros(3).dtype is np.zeros(4).dtype`, or we can use + # the unique character code via hash(self.dtype.char) + return hash((self.shape, self.dtype, self.weak_type, + tuple(self.named_shape.items()))) + + def at_least_vspace(self): + return ShapedArray(self.shape, primal_dtype_to_tangent_dtype(self.dtype), + self.weak_type, self.named_shape) + + def join(self, other): + if definitely_equal_shape(self.shape, other.shape) and self.dtype == other.dtype: + weak_type = self.weak_type and other.weak_type + named_shape = join_named_shapes(self.named_shape, other.named_shape) + return self.update(weak_type=weak_type, named_shape=named_shape) + elif self.dtype == other.dtype: + return UnshapedArray(self.dtype) + else: + raise TypeError(self, other) + + def str_short(self, short_dtypes=False): + dt_str = _short_dtype_name(self.dtype) if short_dtypes else self.dtype.name + dt_str = dt_str.replace('void', 'float0') + shapestr = ','.join(map(str, self.shape)) + if self.named_shape: + named_shapestr = ','.join(f'{k}:{v}' for k, v in self.named_shape.items()) + return f'{dt_str}[{shapestr};{named_shapestr}]' + else: + return f'{dt_str}[{shapestr}]' + + def strip_named_shape(self): + return self.update(named_shape={}) + + def _len(self, ignored_tracer): + try: + return self.shape[0] + except IndexError as err: + raise TypeError("len() of unsized object") from err # same as numpy error + + +def _forward_to_value(self, fun, ignored_tracer, *args): + return fun(self.val, *args) + + +class ConcreteArray(ShapedArray): + __slots__ = ['val'] + array_abstraction_level = 0 + + def __init__(self, dtype, val, weak_type=None): + super().__init__( + np.shape(val), dtype, + weak_type=dtypes.is_weakly_typed(val) if weak_type is None else weak_type) + dtypes.check_valid_dtype(self.dtype) + # Note: canonicalized self.dtype doesn't necessarily match self.val + assert self.dtype == dtypes.canonicalize_dtype(np.result_type(val)), (val, dtype) + self.val = val + + def update(self, dtype=None, val=None, weak_type=None): + dtype = self.dtype if dtype is None else dtype + val = self.val if val is None else val + weak_type = self.weak_type if weak_type is None else weak_type + return ConcreteArray(dtype, val, weak_type) + + def __eq__(self, other): + if (type(self) is type(other) and self.dtype == other.dtype + and self.shape == other.shape and self.weak_type == other.weak_type): + with eval_context(): # in case self.val is an Array + return (self.val == other.val).all() + else: + return False + + def __hash__(self): + return id(self.val) + + def join(self, other) -> AbstractValue: + if self == other: + return self + elif self.shape == other.shape and self.dtype == other.dtype: + weak_type = self.weak_type and other.weak_type + named_shape = join_named_shapes(self.named_shape, other.named_shape) + return ShapedArray( + self.shape, self.dtype, weak_type=weak_type, named_shape=named_shape) + elif self.dtype == other.dtype: + return UnshapedArray(self.dtype, + weak_type=self.weak_type and other.weak_type) + else: + raise TypeError(self, other) + + def str_short(self, short_dtypes=False) -> str: + dt_str = _short_dtype_name(self.dtype) if short_dtypes else self.dtype.name + return f'{self.val}, dtype={dt_str}' + + _bool = partialmethod(_forward_to_value, bool) + _int = partialmethod(_forward_to_value, int) + _hex = partialmethod(_forward_to_value, hex) + _oct = partialmethod(_forward_to_value, oct) + _index = partialmethod(_forward_to_value, operator.index) + + _float = concretization_function_error(float, True) + _complex = concretization_function_error(complex, True) + +def primal_dtype_to_tangent_dtype(primal_dtype): + if isinstance(primal_dtype, dtypes.ExtendedDType): + return primal_dtype._rules.tangent_dtype(primal_dtype) + elif not dtypes.issubdtype(primal_dtype, np.inexact): + return dtypes.float0 + else: + return primal_dtype + + +# Dynamic shape stuff below here! We keep the abstract values distinct just so +# as not to interfere with any static shape machinery. + +# We have a convention of reusing AbsractValues as types, even though we could +# make a distinction and use abstract values during tracing only. This reuse +# becomes a bit more extreme with DShapedArrays. A DShapedArray's shape +# attribute is a tuple which can contain several different types: int, DArray +# (scalar and with dtype of bint type), Tracer (while tracing), Var (when used +# as jaxpr type annotations), or DBIdx/InDBIdx/OutDBIdx (when used in InputType +# or OutputType). We could reduce this polymorphism if it seems cleaner, though +# it's kind of convenient! +class DShapedArray(UnshapedArray): + __slots__ = ['shape'] + shape: tuple[AxisSize, ...] # noqa: F821 + array_abstraction_level: int = 3 + + def __init__(self, shape, dtype, weak_type=False): + self.shape = shape + self.dtype = dtype + self.weak_type = weak_type + + ndim = property(lambda self: len(self.shape)) + size = property(lambda self: + 0 if any(type(d) is int and d == 0 for d in self.shape) + else math.prod(self.shape)) + + def str_short(self, short_dtypes=False) -> str: + del short_dtypes # ignored + shape = f'{",".join(str(d) for d in self.shape)}' if self.shape else '' + dtype = _short_dtype_name(self.dtype) + return f'{dtype}[{shape}]' + __str__ = __repr__ = str_short + + def update(self, shape=None, dtype=None, weak_type=None): + if shape is None: + shape = self.shape + if dtype is None: + dtype = self.dtype + if weak_type is None: + weak_type = self.weak_type + return DShapedArray(shape, dtype, weak_type) + + def _len(self, tracer): + return self.shape[0] + + def __eq__(self, other): + return (type(self) is type(other) + and self.dtype == other.dtype and self.shape == other.shape + and self.weak_type == other.weak_type) + + def __hash__(self): + return hash((self.shape, self.dtype, self.weak_type)) + + def join(self, other): + if (definitely_equal_shape(self.shape, other.shape) and + self.dtype == other.dtype): + weak_type = self.weak_type and other.weak_type + return self.update(weak_type=weak_type) + elif self.dtype == other.dtype: + return UnshapedArray(self.dtype) + else: + raise TypeError(self, other) + + def at_least_vspace(self): + return DShapedArray(self.shape, primal_dtype_to_tangent_dtype(self.dtype), + self.weak_type) + +class DConcreteArray(DShapedArray): + __slots__ = ['val'] + array_abstraction_level = 1 + def __init__(self, shape, dtype, weak_type, val): + super().__init__(shape, dtype, weak_type) + self.val = val + + +pytype_aval_mappings: dict[type, Callable[[Any], AbstractValue]] = {} + + +class DArray: + _aval: DShapedArray + _data: Any # standard array type + def __init__(self, aval, data): + pad_shape = tuple(d.dtype.bound if type(d) is DArray and + type(d.dtype) is bint else d for d in aval.shape) + assert data.shape == pad_shape + self._aval = aval + self._data = data + shape = property(lambda self: self._aval.shape) + dtype = property(lambda self: self._aval.dtype) + aval = property(lambda self: self._aval) + def __repr__(self) -> str: + if not self.shape and type(self.dtype) is bint: + # special-case scalar bints + return f'{int(self._data)}{{≤{self.dtype.bound}}}' + + dtypestr = _short_dtype_name(self._aval.dtype) + shapestr = ','.join(map(str, self.shape)) + slices = tuple(slice(int(d._data)) if type(d) is DArray and + type(d.dtype) is bint else slice(None) for d in self.shape) + data = self._data[slices] + return f'{dtypestr}[{shapestr}] with value: {data}' + def __hash__(self) -> int: + if not self.shape: + return hash((self._aval, int(self._data))) + raise TypeError("unhashable type: DArray") + def __eq__(self, other): + if isinstance(other, DArray) and self._aval == other._aval: + return self._data == other._data + return False + def __len__(self): + return self.shape[0] + +pytype_aval_mappings[DArray] = \ + lambda x: DConcreteArray(x._aval.shape, x._aval.dtype, x._aval.weak_type, + x._data) + +@dataclass(frozen=True) +class bint(dtypes.ExtendedDType): + bound: int + + @property + def type(self) -> type: + return dtypes.extended + + @property + def name(self) -> str: + return f'bint{{≤{self.bound}}}' + + def __str__(self) -> str: + return self.name + +AxisSize = Union[int, DArray, Tracer, Var, DBIdx, InDBIdx, OutDBIdx] + + +class MutableArray: + _aval: ShapedArray + _buf: Array + def __init__(self, aval, buf): + self._aval = aval + self._buf = buf + aval = property(lambda self: self._aval) + shape = property(lambda self: self._aval.shape) + dtype = property(lambda self: self._aval.dtype) + def __getitem__(self, idx): return get_aval(self)._getitem(self, idx) + def __setitem__(self, idx, x): return get_aval(self)._setitem(self, idx, x) + def __repr__(self) -> str: return 'Mutable' + repr(self[...]) +pytype_aval_mappings[MutableArray] = lambda x: x._aval + +def mutable_array(init_val): + return mutable_array_p.bind(init_val) +mutable_array_p = Primitive('mutable_array') + +class InternalMutableArrayEffect(effects.Effect): + pass +internal_mutable_array_effect = InternalMutableArrayEffect() +effects.control_flow_allowed_effects.add_type(InternalMutableArrayEffect) + +@mutable_array_p.def_effectful_abstract_eval +def mutable_array_abstract_eval(init_aval): + from jax._src.state.types import AbstractRef # pytype: disable=import-error + return AbstractRef(init_aval), {internal_mutable_array_effect} + +@mutable_array_p.def_impl +def _mutable_array_impl(init_val): + from jax._src.state.types import AbstractRef # pytype: disable=import-error + aval = raise_to_shaped(get_aval(init_val)) + return MutableArray(AbstractRef(aval), init_val) + + +class AbstractToken(AbstractValue): + def join(self, other): + if isinstance(other, AbstractToken): + return self + else: + assert False, f"Cannot join {self} with {other}" + def str_short(self, short_dtypes=False): return 'Tok' + def at_least_vspace(self): return self +abstract_token: AbstractToken = AbstractToken() + +# Singleton shaped array used by all abstract tokens when shape/dtype is needed. +token_shaped_array: ShapedArray = ShapedArray((0,), np.dtype(np.bool_)) + +# Concrete token object +class Token: + # The underlying data wrapped by the token, could be used to threaded in and + # out of computations to build up data dependency. + _buf: Array + def __init__(self, buf): + self._buf = buf + def block_until_ready(self): + self._buf.block_until_ready() +pytype_aval_mappings[Token] = lambda _: abstract_token + + +def raise_to_shaped(aval: AbstractValue, weak_type=None): + aval_type = type(aval) + if aval_type is ShapedArray and weak_type is None: + return aval + if weak_type is None: + weak_type = getattr(aval, 'weak_type', False) + for typ in aval_type.__mro__: + handler = raise_to_shaped_mappings.get(typ) + if handler: return handler(aval, weak_type) + raise TypeError(type(aval)) + +raise_to_shaped_mappings : dict[type, Callable] = { + AbstractToken: lambda aval, _: aval, + Bot: lambda aval, _: aval, + UnshapedArray: lambda aval, _: aval, + ShapedArray: lambda aval, weak_type: ShapedArray( + aval.shape, aval.dtype, weak_type, aval.named_shape), + DConcreteArray: lambda aval, weak_type: DShapedArray( + aval.shape, aval.dtype, weak_type), +} + +### Operations on shapes and dimension sizes. + +class InconclusiveDimensionOperation(Exception): + """Raised when we cannot conclusively compute with symbolic dimensions.""" + pass + +def is_symbolic_dim(v: Any) -> bool: + """Checks if a value is a symbolic dimension used for shape polymorphism. + + This should be used very rarely, because symbolic dimensions overload all + operators, and should just work. + """ + return hasattr(v, "dimension_as_value") + +def is_constant_dim(d: DimSize) -> bool: + # Whether the dimension is a static integer constant. + try: + operator.index(d) + return True + except: + return False + +def is_dim(v: Any) -> bool: + return is_symbolic_dim(v) or is_constant_dim(v) + +def is_constant_shape(s: Shape) -> bool: + # Whether the shape is a static constant. + return all(is_constant_dim(d) for d in s) + +def definitely_equal_one_of_dim(d1: DimSize, dlist: Sequence[DimSize]) -> bool: + return any(definitely_equal(d1, d) for d in dlist) + +def definitely_equal_shape(s1: Shape, s2: Shape) -> bool: + """Check that two shapes are guaranteed to be element-wise equal. + + In presence of dynamic shapes may return False even when the shapes may + be equal at runtime. + """ + return (len(s1) == len(s2) and + all(unsafe_map(definitely_equal, s1, s2))) + +def divide_shape_sizes(s1: Shape, s2: Shape) -> DimSize: + """Returns an integer "i" s.t., i * size(s2) == size(s1). + Raises InconclusiveDimensionOperation if there is no such integer.""" + sz1 = math.prod(s1) + sz2 = math.prod(s2) + if definitely_equal(sz1, sz2): # Takes care of sz1 and sz2 being 0 + return 1 + q, r = divmod(sz1, sz2) + if isinstance(r, Tracer) or r != 0: + raise InconclusiveDimensionOperation( + f"Cannot divide evenly the sizes of shapes {tuple(s1)} and {tuple(s2)}. " + f"The remainder {r} should be 0.") + return q + +def cancel_divide_tracers(num, denom): + partition = lambda l: partition_list([isinstance(d, Tracer) for d in l], l) + num, num_tracers = partition(num) + denom, denom_tracers = partition(denom) + if num_tracers or denom_tracers: + factor = _cancel_divide(num_tracers, denom_tracers) + if factor is not None: + size1 = math.prod(num) + size2 = math.prod(denom) + if size1 == size2 or size2 != 0: + return factor * (size1 // size2 if size1 != size2 else 1) + +def _cancel_divide(num, denom): + num = list(num) + for a in denom: + i = next((i for i, b in enumerate(num) if definitely_equal(a, b)), None) + if i is None: + break # couldn't cancel + del num[i] + else: + return math.prod(num) + +def is_empty_shape(s: Shape) -> bool: + return any(definitely_equal(d, 0) for d in s) + +def dilate_dim(d: DimSize, dilation: DimSize) -> DimSize: + """max(0, 1 + dilation * (d - 1)). + + Assumes dilation >= 1. + """ + if definitely_equal(dilation, 1): # fast path + return d + return max_dim(1 + dilation * (d - 1), 0) + +def stride_dim(d: DimSize, window_size: DimSize, window_stride: DimSize) -> DimSize: + """max(0, (d - window_size) // window_stride + 1) + + If d < window_size, returns 0. + We assume window_size >= 1 and window_stride >= 1. + """ + # If d < window_size then (d - window_size) // window_stride < 0 + return max_dim((d - window_size) // window_stride + 1, 0) + +# TODO(necula): Deprecated Jan 2024, to be removed. +def non_negative_dim(d: DimSize) -> DimSize: + """max(d, 0).""" + return max_dim(d, 0) + +def min_dim(d1: DimSize, d2: DimSize) -> DimSize: + """Like min(d1, d2) but for both constant and symbolic dimensions.""" + d1_is_constant = is_constant_dim(d1) + if d1_is_constant and is_constant_dim(d2): + return min(d1, d2) + if d1_is_constant: + return d2.rmin(d1) # type: ignore[union-attr] + else: + return d1.min(d2) # type: ignore[union-attr] + +def max_dim(d1: DimSize, d2: DimSize) -> DimSize: + """Like max(d1, d2) but for both constant and symbolic dimensions.""" + d1_is_constant = is_constant_dim(d1) + if d1_is_constant and is_constant_dim(d2): + return max(d1, d2) + if d1_is_constant: + return d2.rmax(d1) # type: ignore[union-attr] + else: + return d1.max(d2) # type: ignore[union-attr] + +def dimension_as_value(d: DimSize): + """Turns a dimension size into a JAX array. + This is the identity function for constant dimensions. + + Has the same abstract value as Python constants. + """ + if isinstance(d, (int, Tracer, np.int32, np.int64)): return d + # For shape_poly._DimPolynomial + if hasattr(d, "dimension_as_value"): return d.dimension_as_value() + return operator.index(d) + +class SomeTracer: + __slots__ = () + def __repr__(self): return "[dynamic]" + +def replace_tracer_for_error_message(obj): + # TODO(mattjj): Many ideas for improving this. Crawl the stack and see if + # there are user variables whose value is == to this object? Or search + # parameters of functions being transformed, at least? Or at least assign + # short unique ids to them? + if isinstance(obj, Tracer): + return SomeTracer() + else: + return obj + +def evaluate_shape(shape: Shape, dim_vars: Sequence[str], + *dim_values: Array) -> Sequence[Array]: + """Evaluates a shape possibly containing non-constants. + + Args: + shape: the shape to evaluate. + dim_vars: the dimension variables names that may appear in `shape`. + dim_values: the dimension values corresponding to `dim_vars`. + + Returns: + a tuple of JAX values corresponding to `shape`, of type + `dim_value_dtype`. + """ + env = dict(zip(dim_vars, dim_values)) + def eval_one_dim(d: DimSize): + try: + return operator.index(d) + except: + # Is a _DimExpr + return d._evaluate(env) # type: ignore + return tuple(eval_one_dim(d) for d in shape) + +def dim_value_dtype(): + """The dtype to be used for dimension values.""" + return dtypes.canonicalize_dtype(np.int64) + +def dim_constant(ct: int): + dtype = dim_value_dtype() + assert dtype in (np.int32, np.int64) + if dtype == np.int32: + return np.int32(ct) + elif dtype == np.int64: + return np.int64(ct) + +def dim_value_aval() -> AbstractValue: + return ShapedArray((), dim_value_dtype(), weak_type=True) + +# ------------------- Named shapes ------------------- + + +class NamedShape: + def __init__(self, *args, **kwargs): + self.__positional = canonicalize_shape(args) + # TODO: Assert that kwargs match axis env? + self.__named = dict(kwargs) + + @property + def rank(self): + return len(self.__positional) + len(self.__named) + + @property + def positional_rank(self): + return len(self.__positional) + + @property + def named_rank(self): + return len(self.__named) + + @property + def positional(self): + return self.__positional + + @property + def names(self): + return self.__named.keys() + + @property + def named_sizes(self): + return self.__named.values() + + @property + def named_items(self): + return self.__named.items() + + def __getitem__(self, idx): + try: + idx = operator.index(idx) + return self.__positional[idx] + except TypeError: + pass + return self.__named[idx] + + @property + def total(self): + total = 1 + for s in self.__positional: total *= s + for s in self.__named.values(): total *= s + return total + + def __str__(self): + # TODO(mattjj,frostig): revise not to miss commas + if not self.__named: + return str(self.__positional) + return (f"({', '.join(map(str, self.__positional))}{', ' if self.__named else ''}" + f"{', '.join(f'{k}={v}' for k, v in self.__named.items())})") + + def __eq__(self, other): + if isinstance(other, NamedShape): + return (self.__positional, self.__named) == (other.__positional, other.__named) + if isinstance(other, tuple): + return not self.__named and self.__positional == other + return False + + def __hash__(self): + named = frozenset(self.__named.items()) + return hash((self.__positional, named)) + +def join_named_shapes(*named_shapes): + result = {} + for named_shape in named_shapes: + for name, size in named_shape.items(): + if result.setdefault(name, size) != size: + raise TypeError( + f"Axis name {name} used with inconsistent sizes: {result[name]} != {size}") + return result + +# TODO: Make canonicalize_shape return named shapes? +def as_named_shape(shape) -> NamedShape: + if isinstance(shape, NamedShape): + return shape + return NamedShape(*shape) + + +# ------------------- Call ------------------- + +class CallPrimitive(Primitive): + multiple_results = True + call_primitive = True + + def bind(self, fun, *args, **params): + call_bind_continuation, top_trace, fun_, tracers, params = ( + call_bind_with_continuation(self, fun, *args, **params)) + outs = top_trace.process_call(self, fun_, tracers, params) + return call_bind_continuation(outs) + + def get_bind_params(self, params): + new_params = dict(params) + jaxpr = new_params.pop('call_jaxpr') + subfun = lu.hashable_partial(lu.wrap_init(eval_jaxpr), jaxpr, ()) + if config.dynamic_shapes.value: + subfun = lu.annotate(subfun, _jaxpr_type_to_callable_annotation(jaxpr)) + return [subfun], new_params + +def call_bind_with_continuation(primitive: CallPrimitive, fun, *args, **params): + top_trace = find_top_trace(args) + fun_, env_trace_todo = process_env_traces_call( + fun, primitive, top_trace.level, tuple(params.items())) + tracers = map(top_trace.full_raise, args) + fun_ = lu.annotate(fun_, fun.in_type) + + def call_bind_continuation(outs): + return map(full_lower, apply_todos(env_trace_todo(), outs)) + return call_bind_continuation, top_trace, fun_, tracers, params + +@lu.transformation_with_aux +def process_env_traces_call(primitive: CallPrimitive, level: int, + params_tuple: tuple, *args): + outs = yield args, {} + params = dict(params_tuple) + todo = [] + while True: + tracers = [x for x in outs if isinstance(x, Tracer) and x._trace.level > level] + if not tracers: + break + ans = max(tracers, key=_get_trace_level) + trace = ans._trace.main.with_cur_sublevel() + outs = map(trace.full_raise, outs) + outs, cur_todo = trace.post_process_call(primitive, outs, params) + todo.append(cur_todo) + yield outs, tuple(todo) # Ensure the aux output is immutable + +def apply_todos(todos, outs): + todos_list = list(todos) + while todos_list: + outs = map(full_lower, todos_list.pop()(outs)) + return outs + + +def call_impl(f: lu.WrappedFun, *args, **params): + del params # params parameterize the call primitive, not the function + with new_sublevel(): + return f.call_wrapped(*args) + +call_p: CallPrimitive = CallPrimitive('call') +call = call_p.bind +call_p.def_impl(call_impl) + + +class ClosedCallPrimitive(CallPrimitive): + def get_bind_params(self, params): + new_params = dict(params) + jaxpr = new_params.pop('call_jaxpr') + subfun = lu.wrap_init(partial(eval_jaxpr, jaxpr.jaxpr, jaxpr.consts)) + return [subfun], new_params + +closed_call_p: ClosedCallPrimitive = ClosedCallPrimitive('closed_call') +closed_call_p.def_impl(call_impl) +closed_call_p.def_effectful_abstract_eval( + lambda *_, call_jaxpr: (call_jaxpr.out_avals, call_jaxpr.effects)) + + +outfeed_primitives: set[Primitive] = set() +def jaxpr_uses_outfeed(jaxpr: Jaxpr) -> bool: + """Finds if there are outfeed primitives anywhere inside a Jaxpr.""" + return any(primitive_uses_outfeed(eqn.primitive, eqn.params) + for eqn in jaxpr.eqns) + +def _param_uses_outfeed(param): + if type(param) is Jaxpr: + if jaxpr_uses_outfeed(param): + return True + elif type(param) is ClosedJaxpr: + if jaxpr_uses_outfeed(param.jaxpr): + return True + return False + +def primitive_uses_outfeed(prim: Primitive, params: dict) -> bool: + if prim in outfeed_primitives: + return True + for param in params.values(): + if isinstance(param, tuple): + if any(unsafe_map(_param_uses_outfeed, param)): + return True + elif _param_uses_outfeed(param): + return True + return False + +# ------------------- Map ------------------- + +class MapPrimitive(Primitive): + multiple_results = True + map_primitive = True + + def bind(self, fun, *args, **params): + assert len(params['in_axes']) == len(args) + return map_bind(self, fun, *args, **params) + + def process(self, trace, fun, tracers, params): + return trace.process_map(self, fun, tracers, params) + + def post_process(self, trace, out_tracers, params): + return trace.post_process_map(self, out_tracers, params) + + def get_bind_params(self, params): + new_params = dict(params) + jaxpr = new_params.pop('call_jaxpr') + subfun = lu.hashable_partial(lu.wrap_init(eval_jaxpr), jaxpr, ()) + axes = new_params.pop('out_axes') + new_params['out_axes_thunk'] = HashableFunction(lambda: axes, closure=axes) + return [subfun], new_params + + +def map_bind_with_continuation(primitive: MapPrimitive, fun, *args, + out_axes_thunk, **params): + # The new thunk depends deterministically on the old thunk and the wrapped + # function. Any caching already has to include the wrapped function as part + # of the key, so we only use the previous thunk for equality checks. + @as_hashable_function(closure=out_axes_thunk) + def new_out_axes_thunk(): + out_axes = out_axes_thunk() + _, out_axes_transforms = todo_and_xforms() + for t in out_axes_transforms: + out_axes = t(out_axes) + return out_axes + params = dict(params, out_axes_thunk=new_out_axes_thunk) + top_trace = find_top_trace(args) + fun, todo_and_xforms = process_env_traces_map( + fun, primitive, top_trace and top_trace.level, tuple(params.items())) + tracers = map(top_trace.full_raise, args) + + def map_bind_continuation(outs): + env_trace_todo, _ = todo_and_xforms() + return map(full_lower, apply_todos(env_trace_todo, outs)) + + return map_bind_continuation, top_trace, fun, tracers, params + + +def map_bind(primitive: MapPrimitive, fun, *args, **params): + map_bind_continuation, top_trace, fun, tracers, params = ( + map_bind_with_continuation(primitive, fun, *args, **params)) + return map_bind_continuation( + primitive.process(top_trace, fun, tracers, params)) + +@lu.transformation_with_aux +def process_env_traces_map(primitive: MapPrimitive, level: int, + params_tuple: tuple, *args): + outs = yield args, {} + params = dict(params_tuple) + todo = [] + out_axes_transforms = [] + while True: + tracers = [x for x in outs if isinstance(x, Tracer) + and (level is None or x._trace.level > level)] + if not tracers: + break + ans = max(tracers, key=_get_trace_level) + trace = ans._trace.main.with_cur_sublevel() + outs = map(trace.full_raise, outs) + outs, (cur_todo, cur_xform) = primitive.post_process(trace, outs, params) + todo.append(cur_todo) + out_axes_transforms.append(cur_xform) + yield outs, (tuple(todo), tuple(out_axes_transforms)) + + +def mapped_aval(size: AxisSize, axis: int | None, + aval: AbstractValue) -> AbstractValue: + handler, _ = aval_mapping_handlers.get(type(aval), (None, None)) + if handler is not None: + return handler(size, axis, aval) + else: + raise TypeError(f"no mapping handler for {aval} of type {type(aval)}") + +def unmapped_aval(size: AxisSize, axis_name, axis: int | None, + aval: AbstractValue) -> AbstractValue: + _, handler = aval_mapping_handlers.get(type(aval), (None, None)) + if handler is not None: + return handler(size, axis_name, axis, aval) + else: + raise TypeError(f"no unmapping handler for {aval} of type {type(aval)}") + + +def _map_shaped_array( + size: int, axis: int | None, aval: ShapedArray) -> ShapedArray: + assert axis is None or aval.shape[axis] == size + # TODO: Extend the named shape + if axis is None: return aval + return ShapedArray(tuple_delete(aval.shape, axis), aval.dtype, + named_shape=aval.named_shape, weak_type=aval.weak_type) + +def _unmap_shaped_array( + size: int, axis_name: AxisName, axis: int | None, aval: ShapedArray + ) -> ShapedArray: + named_shape = dict(aval.named_shape) + named_shape.pop(axis_name, None) # TODO: make this mandatory + if axis is None: return aval.update(named_shape=named_shape) + elif type(axis) is int: + return ShapedArray(tuple_insert(aval.shape, axis, size), aval.dtype, + named_shape=named_shape, weak_type=aval.weak_type) + else: raise TypeError(axis) + +def _map_dshaped_array( + size: AxisSize, axis: int | None, aval: DShapedArray) -> DShapedArray: + if axis is None: return aval + return DShapedArray(tuple_delete(aval.shape, axis), aval.dtype, + aval.weak_type) + +def _unmap_dshaped_array( + size: AxisSize, axis_name: AxisName, axis: int | None, aval: DShapedArray + ) -> DShapedArray: + if axis is None: return aval + elif type(axis) is int: + return DShapedArray(tuple_insert(aval.shape, axis, size), aval.dtype, + weak_type=aval.weak_type) + else: + raise TypeError(axis) + +AvalMapHandlerPair = tuple[Callable, Callable] +aval_mapping_handlers: dict[type, AvalMapHandlerPair] = { + DShapedArray: (_map_dshaped_array, _unmap_dshaped_array), + ShapedArray: (_map_shaped_array, _unmap_shaped_array), + ConcreteArray: (_map_shaped_array, _unmap_shaped_array), + AbstractToken: (lambda _, __, a: a, lambda _, __, ___, a: a) +} + +@contextmanager +def extend_axis_env(axis_name: AxisName, size: int, tag: Any): + frame = AxisEnvFrame(axis_name, size, tag) + ts = thread_local_state.trace_state + ts.axis_env.append(frame) + config.update_thread_local_jit_state( + axis_env_state=tuple(f for f in ts.axis_env + if f.name is not no_axis_name)) + try: + yield + finally: + ts.axis_env.pop() + config.update_thread_local_jit_state( + axis_env_state=tuple(f for f in ts.axis_env + if f.name is not no_axis_name)) + +@contextmanager +def extend_axis_env_nd(axes: Iterable[tuple[AxisName, int]], tag: Any = None): + frames = [AxisEnvFrame(axis_name, size, tag) for axis_name, size in axes] + ts = thread_local_state.trace_state + ts.axis_env.extend(frames) + config.update_thread_local_jit_state( + axis_env_state=tuple(f for f in ts.axis_env + if f.name is not no_axis_name)) + try: + yield + finally: + for _ in frames: ts.axis_env.pop() + config.update_thread_local_jit_state( + axis_env_state=tuple(f for f in ts.axis_env + if f.name is not no_axis_name)) + + +@contextmanager +def stash_axis_env(): + "Promise that a function or with-suite does not depend implicitly on axis env" + # If the promise is broken, then a NameError about an unbound axis name will + # be raised. + ts = thread_local_state.trace_state + prev_axis_env, ts.axis_env = ts.axis_env, [] + config.update_thread_local_jit_state(axis_env_state=()) + try: + yield + finally: + ts.axis_env = prev_axis_env + config.update_thread_local_jit_state( + axis_env_state=tuple(f for f in ts.axis_env + if f.name is not no_axis_name)) + + +# When a mapped function is given no axis name, we generate a name object based +# on the id of the function object. Collisions aren't important because this +# name can't be used in collectives, as user code never gets a ref to this +# object. We don't want to use the function object itself because that might +# persist references to the function object. +# TODO(mattjj): revisit this unique axis name strategy +@total_ordering +class _TempAxisName: + + def __init__(self, obj): + self.id = id(obj) + + def __repr__(self): + return f'' + + def __hash__(self): + return hash(self.id) + + def __eq__(self, other): + return type(other) is _TempAxisName and self.id == other.id + + def __lt__(self, other): + return type(other) is _TempAxisName and self.id < other.id + + +def axis_frame(axis_name: AxisName, main_trace: MainTrace | None = None + ) -> AxisEnvFrame: + frames = thread_local_state.trace_state.axis_env + for frame in reversed(frames): + if (frame.name == axis_name and + (main_trace is None or frame.main_trace is main_trace)): + return frame + named_axes = [frame.name for frame in reversed(frames) + if not isinstance(frame.name, _TempAxisName)] + raise NameError( + f'unbound axis name: {axis_name}. The following axis names (e.g. defined ' + f'by pmap) are available to collective operations: {named_axes}') + + +@dataclass(frozen=True) +class NamedAxisEffect(effects.Effect): + """A side-effect introducing a new named axis into the current scope.""" + + name: AxisName + + +effects.control_flow_allowed_effects.add_type(NamedAxisEffect) +effects.custom_derivatives_allowed_effects.add_type(NamedAxisEffect) +effects.lowerable_effects.add_type(NamedAxisEffect) +effects.remat_allowed_effects.add_type(NamedAxisEffect) + + +def filter_named_axis_effects( + effects: Effects, names: Collection[AxisName] +) -> Effects: + return {e for e in effects + if not isinstance(e, NamedAxisEffect) or e.name not in names} + + +def remove_named_axis_effects( + jaxpr: Jaxpr, names: Collection[AxisName] +) -> Jaxpr: + if not names or not jaxpr.effects: + return jaxpr + return jaxpr.replace(effects=filter_named_axis_effects(jaxpr.effects, names)) + + +ParamDict = dict[str, Any] +AxisSubst = Callable[[AxisName], tuple[AxisName, ...]] + +class NameGatheringSubst: + def __init__(self): + self.axis_names = set() + def __call__(self, axis_name): + self.axis_names.add(axis_name) + return (axis_name,) + +def used_axis_names(primitive: Primitive, params: ParamDict) -> set[AxisName]: + subst = NameGatheringSubst() + subst_axis_names(primitive, params, subst) + return subst.axis_names + +def subst_axis_names(primitive: Primitive, params: ParamDict, subst: AxisSubst, traverse: bool = True) -> ParamDict: + if primitive in axis_substitution_rules: + return axis_substitution_rules[primitive](params, subst, traverse) + if not traverse: + return params + # Default implementation: substitute names in all jaxpr parameters + if isinstance(primitive, MapPrimitive): + def shadowed_subst(name): + return (name,) if name == params['axis_name'] else subst(name) + else: + shadowed_subst = subst + jaxpr_params = [(n, v) for n, v in params.items() if isinstance(v, (Jaxpr, ClosedJaxpr))] + if not jaxpr_params: + return params + new_params = dict(params) + for name, jaxpr in jaxpr_params: + new_params[name] = subst_axis_names_jaxpr(jaxpr, shadowed_subst) + return new_params + +class DuplicateAxisNameError(Exception): + def __init__(self, var): + self.var = var + self.eqn = None + +def subst_axis_names_effects(effects: Set[Effect], subst: AxisSubst) -> Set[Effect]: + new_effects = set[Effect]() + for e in effects: + if isinstance(e, NamedAxisEffect): + new_effects.update(map(NamedAxisEffect, subst(e.name))) + else: + new_effects.add(e) + return new_effects + +def subst_axis_names_var(v: Var, subst: AxisSubst, var_map: dict[Var, Var]) -> Var: + # Var identity is load-bearing, so we can't have duplicates! + if isinstance(v, DropVar): return v + assert v not in var_map + if not hasattr(v.aval, 'named_shape'): + var_map[v] = v + return v + names = tuple(it.chain.from_iterable(subst(name) for name in v.aval.named_shape)) + named_shape = {name: axis_frame(name).size for name in names} + if len(named_shape) != len(names): + raise DuplicateAxisNameError(v) + new_v = Var(v.suffix, v.aval.update(named_shape=named_shape)) + var_map[v] = new_v + return new_v + +def subst_axis_names_eqn(eqn: JaxprEqn, subst: AxisSubst, var_map: dict[Var, Var]) -> JaxprEqn: + invars: list[Atom] = [v if isinstance(v, Literal) else var_map[v] for v in eqn.invars] + try: + outvars = [subst_axis_names_var(v, subst, var_map) for v in eqn.outvars] + except DuplicateAxisNameError as e: + e.eqn = eqn + raise + params = subst_axis_names(eqn.primitive, eqn.params, subst) + effects = subst_axis_names_effects(eqn.effects, subst) + return eqn.replace(invars=invars, outvars=outvars, params=params, effects=effects) + +def do_subst_axis_names_jaxpr(jaxpr: Jaxpr | ClosedJaxpr, subst: AxisSubst): + consts = None + if isinstance(jaxpr, ClosedJaxpr): + consts = jaxpr.consts + jaxpr = jaxpr.jaxpr + var_map: dict[Var, Var] = {} + invars = [subst_axis_names_var(v, subst, var_map) for v in jaxpr.invars] # type: ignore[union-attr] + constvars = [subst_axis_names_var(v, subst, var_map) for v in jaxpr.constvars] # type: ignore[union-attr] + eqns = [subst_axis_names_eqn(eqn, subst, var_map) for eqn in jaxpr.eqns] + outvars: list[Atom] = [v if isinstance(v, Literal) else var_map[v] for v in jaxpr.outvars] # type: ignore[union-attr] + effects = subst_axis_names_effects(jaxpr.effects, subst) + new_jaxpr = Jaxpr(constvars, invars, outvars, eqns, effects) + if consts is not None: + return ClosedJaxpr(new_jaxpr, consts) + return new_jaxpr + +def used_axis_names_jaxpr(jaxpr: Jaxpr | ClosedJaxpr): + return {e.name for e in jaxpr.effects if isinstance(e, NamedAxisEffect)} + +def subst_axis_names_jaxpr(jaxpr: Jaxpr | ClosedJaxpr, subst: AxisSubst): + if isinstance(subst, NameGatheringSubst): # This is a common case, so we optimize it! + subst.axis_names |= used_axis_names_jaxpr(jaxpr) + return jaxpr + return do_subst_axis_names_jaxpr(jaxpr, subst) + +def replace_jaxpr_effects(jaxpr: ClosedJaxpr, effects: Effects): + return _replace_jaxpr_effects(jaxpr, frozenset(effects)) + +@weakref_lru_cache +def _replace_jaxpr_effects(jaxpr: ClosedJaxpr, effects: frozenset[Effect]): + return jaxpr.replace(jaxpr=jaxpr.jaxpr.replace(effects=set(effects))) + + +axis_substitution_rules: dict[Primitive, Callable[[ParamDict, AxisSubst, bool], ParamDict]] = {} + +# ------------------- AxisPrimitive ------------------- +# Primitives that store axis names in params and want those axis names to +# participate in dispatch should subclass AxisPrimitive. + +class AxisPrimitive(Primitive): + def bind(self, *args, **params): + top_trace = find_top_trace(args) + axis_main = max((axis_frame(a).main_trace for a in used_axis_names(self, params)), + default=None, key=lambda t: getattr(t, 'level', -1)) + top_trace = (top_trace if not axis_main or axis_main.level < top_trace.level + else axis_main.with_cur_sublevel()) + return self.bind_with_trace(top_trace, args, params) + + +# ------------------- Jaxpr checking ------------------- + +def typecheck(aval: AbstractValue, x) -> bool: + return typecompat(aval, get_aval(x)) + +def typecompat(aval_ref: AbstractValue, aval: AbstractValue) -> bool: + """Determine whether `aval` conforms to `aval_ref`. + + Ignores weak_type and named_shape, other than to check that an axis name isn't + used with different sizes. + """ + try: + return typematch(aval_ref, lattice_join(aval_ref, aval)) + except TypeError: + return False + +def typematch(aval1: AbstractValue, aval2: AbstractValue) -> bool: + """Determine whether `aval1` and `aval2` are equivalent. + + Ignores weak_type and named_shape, other than to check that an axis name isn't + used with different sizes. + """ + if aval1 == aval2: return True + # unequal avals may still represent the same type, because type is represented + # by avals at the shaped level, and because weak type tags and (for now) named + # shape components aren't considered part of the type + if isinstance(aval1, ShapedArray) and isinstance(aval2, ShapedArray): + # a bonus check for whether any named axes have inconsistent sizes + join_named_shapes(aval1.named_shape, aval2.named_shape) + return (raise_to_shaped(aval1, weak_type=False).strip_named_shape() == + raise_to_shaped(aval2, weak_type=False).strip_named_shape()) + +class JaxprTypeError(TypeError): pass + +custom_typechecks: dict[Primitive, Callable] = {} + +def _check_closed_call(_, *in_atoms, call_jaxpr): + in_avals = [x.aval for x in in_atoms] + if not all(map(typecompat, call_jaxpr.in_avals, in_avals)): + raise JaxprTypeError("Closed call in_avals mismatch") + return call_jaxpr.out_avals, call_jaxpr.effects +custom_typechecks[closed_call_p] = _check_closed_call + +def check_jaxpr(jaxpr: Jaxpr): + """Checks well-formedness of a jaxpr. + + Specifically, check that: + - variables that are read are bound beforehand + - variables are typed equally throughout a jaxpr + - variable type annotations are compatible with their binding expression + + Raises `JaxprTypeError` if `jaxpr` is determined invalid. Returns `None` + otherwise. + """ + @functools.cache + def ctx_factory(): + ctx = JaxprPpContext() + pp_settings = JaxprPpSettings() + try: pp_jaxpr(jaxpr, ctx, pp_settings) # side-effect on ctx, build variable names + except: pass + return ctx, pp_settings + + try: + _check_jaxpr(ctx_factory, jaxpr) + except JaxprTypeError as e: + ctx, pp_settings = ctx_factory() + if len(e.args) == 2: + msg, eqnidx = e.args + jaxpr_str = str(pp_jaxpr_eqn_range(jaxpr, eqnidx - 10, eqnidx + 10, ctx, + pp_settings)) + else: + msg, = e.args + jaxpr_str = str(pp_jaxpr_eqn_range(jaxpr, 0, 20, ctx, pp_settings)) + msg = "\n\n".join([msg, "while checking jaxpr:", jaxpr_str]) + raise JaxprTypeError(msg) from None + + # Run key reuse checker after validating jaxpr: + if config.debug_key_reuse.value: + # Import here to avoid circular imports + from jax.experimental.key_reuse._core import check_key_reuse_jaxpr # pytype: disable=import-error + check_key_reuse_jaxpr(jaxpr) + + +def _check_jaxpr( + ctx_factory: Callable[[], tuple[JaxprPpContext, JaxprPpSettings]], + jaxpr: Jaxpr + ) -> None: + # Use set of variables to types to check that variables are in scope. + env: set[Var] = set() + + def read(x: Atom) -> Atom: + # Check the type annotation is itself well-typed. + check_type(ctx_factory, env, x.aval) + if isinstance(x, Var): + # Check the variable is in-scope and consistently typed. + if x not in env: + ctx, _ = ctx_factory() + raise JaxprTypeError(f"Variable '{pp_var(x, ctx)}' not defined") + return x + elif isinstance(x, Literal): + # Check that the literal matches its type annotation. + if not typecheck(x.aval, x.val): + ctx, _ = ctx_factory() + raise JaxprTypeError( + f"Literal value {x.val} does not match its type annotation " + f"{pp_aval(x.aval, ctx)}") + return x + else: + assert False, "syntactically invalid jaxpr" + + def write(v: Var, a: AbstractValue) -> None: + assert isinstance(v, Var), "syntactically invalid jaxpr" + # Check the type annotation of the binder is itself well-typed. + check_type(ctx_factory, env, v.aval) + # Check that the variable is not already bound. + if v in env: + ctx, _ = ctx_factory() + raise JaxprTypeError(f"Variable '{pp_var(v, ctx)}' already bound") + # Check that the computed type is consistent with the binder annotation. + if not typematch(v.aval, a): + ctx, _ = ctx_factory() + raise JaxprTypeError( + f"Value for variable '{pp_var(v, ctx)}' inconsistently typed " + f"as {pp_aval(a, ctx)} for let-binder of type {pp_aval(v.aval, ctx)}") + # If the variable is not a DropVar, add it to the environment. + if not isinstance(v, DropVar): + env.add(v) + + # Check type annotations on lambda binders. + for v in it.chain(jaxpr.constvars, jaxpr.invars): + check_type(ctx_factory, env, v.aval) + write(v, v.aval) + + # Check each eqn. + sentinel = object() + in_idx = {v: i for i, v in enumerate(it.chain(jaxpr.constvars, jaxpr.invars))} + for eqn_idx, eqn in enumerate(jaxpr.eqns): + prim = eqn.primitive + try: + in_atoms = map(read, eqn.invars) + in_avals = [x.aval for x in in_atoms] # use in_atoms for dyn shapes + + # Compute the type of the primitive application. + if prim in custom_typechecks: + out_type, eqn_effects = custom_typechecks[prim]( + ctx_factory, *in_atoms, **eqn.params) + elif prim.call_primitive: + out_type, eqn_effects = _check_call(ctx_factory, prim, in_atoms, + eqn.params) + elif prim.map_primitive: + out_type, eqn_effects = _check_map(ctx_factory, prim, in_avals, + eqn.params) + else: + out_type, eqn_effects = check_eqn(prim, in_avals, eqn.params) + + # Check the computed effect type matches the eqn's annotation, and is + # included in the jaxpr's annotation. + if prim is mutable_array_p: + outvar, = eqn.outvars + in_idx[outvar] = None # type: ignore + if eqn.effects != eqn_effects: + raise JaxprTypeError("Inferred effects do not match equation effects. " + f"Equation effects: {eqn.effects}. " + f"Inferred effects: {eqn_effects}") + for eff in eqn.effects: + if isinstance(eff, effects.JaxprInputEffect): + eqn_invar = eqn.invars[eff.input_index] + if (jaxpr_index := in_idx.get(eqn_invar, sentinel)) is sentinel: + raise JaxprTypeError( + "Invalid `JaxprInputEffect`: must correspond to a jaxpr invar") + jaxpr_effect = eff.replace(input_index=jaxpr_index) + if jaxpr_effect not in jaxpr.effects: + raise JaxprTypeError( + "Invalid `JaxprInputEffect`: must be present in jaxpr. " + f"{jaxpr_effect} is not in {jaxpr.effects}.") + elif isinstance(eff, NamedAxisEffect): + # It is valid for a primitive to discharge the named axis effect. + continue + elif eff not in jaxpr.effects: + raise JaxprTypeError("Equation effect not present in jaxpr effects. " + f"Equation effect: {eff}. " + f"Jaxpr effects: {jaxpr.effects}") + + # Check out_type matches the let-binders' annotation (after substitution). + out_type = substitute_vars_in_output_ty(out_type, eqn.invars, eqn.outvars) + map(write, eqn.outvars, out_type) + + except JaxprTypeError as e: + ctx, settings = ctx_factory() + msg, = e.args + src = source_info_util.summarize(eqn.source_info) + msg = "\n\n".join([msg, "in equation:", str(pp.nest(2, pp_eqn(eqn, ctx, settings))), + f"from source: {src}"]) + raise JaxprTypeError(msg, eqn_idx) from None + + # TODO(mattjj): include output type annotation on jaxpr and check it here + map(read, jaxpr.outvars) + +def check_type( + ctx_factory: Callable[[], tuple[JaxprPpContext, JaxprPpSettings]], + env: set[Var], + ty: AbstractValue, + ) -> None: + if isinstance(ty, DShapedArray): + # Check all elements in the shape tuple are well-typed. + for d in ty.shape: + if (isinstance(d, int) or + isinstance(d, DArray) and not d.shape and type(d.dtype) == bint): + continue + elif isinstance(d, Var): + if d not in env: + ctx, _ = ctx_factory() + raise JaxprTypeError(f"unbound axis size: '{pp_var(d, ctx)}'") + if not isinstance(d.aval, (ShapedArray, DShapedArray)): + raise JaxprTypeError(f"axis size with unexpected type annotation: " + f"{d.aval} of type {type(d.aval)}") + if isinstance(d.aval, ShapedArray): + shape, dtype = d.aval.shape, d.aval.dtype + if shape: raise JaxprTypeError(f"axis size nonscalar: {d.aval}") + if not dtypes.issubdtype(dtype, np.integer): + raise JaxprTypeError(f"axis size with non-integer dtype: {d.aval}") + else: + assert isinstance(d.aval, DShapedArray) + shape, dtype = d.aval.shape, d.aval.dtype + if shape: raise JaxprTypeError(f"axis size nonscalar: {d.aval}") + if type(dtype) is not bint: + raise JaxprTypeError( + f"DArray axis size with non-bint dtype: {d.aval}") + else: + raise JaxprTypeError(f"unexpected type in shape: {type(d)}") + else: + return # Except in above case(s), all syntactic forms are valid + +def substitute_vars_in_output_ty( + out_type: Sequence[AbstractValue], # shapes may contain InDBIdx / OutDBIdx + in_atoms: Sequence[Atom], + out_binders: Sequence[Var], + ) -> list[AbstractValue]: # shapes may contain Vars + in_atoms = [x.val if type(x) is Literal else x for x in in_atoms] + result = [] + for aval in out_type: + if type(aval) is DShapedArray: + shape = [in_atoms[d.val] if type(d) is InDBIdx else + out_binders[d.val] if type(d) is OutDBIdx else + d for d in aval.shape] + aval = aval.update(shape=tuple(shape)) + result.append(aval) + return result + +def check_eqn(prim, in_avals, params): + for jaxpr in jaxprs_in_params(params): + check_jaxpr(jaxpr) + + out_avals, effects = prim.abstract_eval(*in_avals, **params) + if not prim.multiple_results: + out_avals = [out_avals] + return out_avals, effects + +def _check_call(ctx_factory, prim, in_atoms, params): + if "call_jaxpr" not in params: + raise JaxprTypeError( + f"Call primitive {prim} missing 'call_jaxpr' parameter") + call_jaxpr = params["call_jaxpr"] + + if len(in_atoms) != len(call_jaxpr.invars): + raise JaxprTypeError(f"Call primitive {prim} with {len(in_atoms)} " + f"operands cannot call jaxpr with " + f"{len(call_jaxpr.invars)} inputs") + + # Check `call_jaxpr` can be applied to in_atoms. + env: dict[Var, Atom] = {} + def substitute(aval: AbstractValue): + if isinstance(aval, DShapedArray): + aval = aval.update(shape=tuple(env.get(d, d) for d in aval.shape)) # type: ignore + return aval + for v, x in zip(call_jaxpr.invars, in_atoms): + if not typecompat(substitute(v.aval), x.aval): + # TODO(mattjj): vars in error message are confusing b/c of Var.__repr__ + raise JaxprTypeError(f"Call primitive {prim} passes operand {x} of type " + f"{x.aval} to jaxpr expecting type " + f"{substitute(v.aval)}") + env[v] = x if type(x) is Var else x.val + + _check_jaxpr(ctx_factory, call_jaxpr) + + invars, outvars = call_jaxpr.invars, call_jaxpr.outvars + in_map : dict[Var, InDBIdx] = {v: InDBIdx(i) for i, v in enumerate( invars)} + out_map: dict[Var, OutDBIdx] = {x: OutDBIdx(i) for i, x in enumerate(outvars) + if type(x) is Var} + out_avals = [x.aval for x in call_jaxpr.outvars] + out_type = [a.update(shape=tuple(in_map.get(d, out_map.get(d)) + if type(d) is Var else d for d in a.shape)) + if type(a) is DShapedArray else a for a in out_avals] + return out_type, call_jaxpr.effects + +def _check_map(ctx_factory, prim, in_avals, params): + if "call_jaxpr" not in params: + raise JaxprTypeError(f"Map primitive {prim} missing 'call_jaxpr' parameter") + call_jaxpr = params["call_jaxpr"] + ordered_effects_ = effects.ordered_effects.filter_in(call_jaxpr.effects) + if ordered_effects_: + raise JaxprTypeError( + f"Map primitive {prim} mapping ordered effects: {ordered_effects_}") + if "axis_size" not in params: + raise JaxprTypeError(f"Map primitive {prim} missing 'axis_size' parameter") + axis_size = params["axis_size"] + if "axis_name" not in params: + raise JaxprTypeError(f"Map primitive {prim} missing 'axis_name' parameter") + axis_name = params["axis_name"] + if "in_axes" not in params: + raise JaxprTypeError(f"Map primitive {prim} missing 'in_axes' parameter") + in_axes = params["in_axes"] + if "out_axes" not in params: + raise JaxprTypeError(f"Map primitive {prim} missing 'out_axes' parameter") + out_axes = params["out_axes"] + + binder_avals = [unmapped_aval(axis_size, axis_name, in_axis, v.aval) + if in_axis is not None else v.aval + for v, in_axis in zip(call_jaxpr.invars, in_axes)] + for binder_aval, in_aval in zip(binder_avals, in_avals): + if not typecompat(binder_aval, in_aval): + raise JaxprTypeError(f"Call primitive {prim} passes operand {in_aval} " + f"to jaxpr expecting {binder_aval}") + + with extend_axis_env(params['axis_name'], axis_size, None): + _check_jaxpr(ctx_factory, call_jaxpr) + + mapped_out_avals = [v.aval for v in call_jaxpr.outvars] + out_avals = [unmapped_aval(axis_size, axis_name, out_axis, aval) + if out_axis is not None else aval + for aval, out_axis in zip(mapped_out_avals, out_axes)] + return out_avals, filter_named_axis_effects(call_jaxpr.effects, {axis_name}) + + +# ------------------- Jaxpr printed representation ------------------- + +def pp_toplevel_jaxpr(jaxpr_to_print, *, source_info=False, print_shapes=True, + custom_pp_eqn_rules=True, name_stack=False, + print_effects: bool = False) -> pp.Doc: + context = JaxprPpContext() + settings = JaxprPpSettings( + source_info=source_info, + print_shapes=print_shapes, + custom_pp_eqn_rules=custom_pp_eqn_rules, + name_stack=name_stack, + print_effects=print_effects) + + # Compute how many times each jaxpr is used. + names = defaultdict[Jaxpr, str](lambda: "jaxpr") + jaxpr_counts = Counter[Jaxpr]() + s = deque([jaxpr_to_print]) + while s: + jaxpr = s.popleft() + jaxpr_counts[jaxpr] += 1 + for eqn in jaxpr.eqns: + # TODO(slebedev): Come up with a more elaborate heuristic for name=. + name = eqn.params.get("name") + if name is None: + s.extend(jaxprs_in_params(eqn.params)) + continue + name = name.strip("<>") # -> lambda + for subjaxpr in jaxprs_in_params(eqn.params): + s.append(subjaxpr) + names.setdefault(subjaxpr, name) + + # Pull jaxprs occurring more than once to the top-level, making sure + # that their names are unique. + docs = [] + name_counts = Counter[str]() + for jaxpr, c in jaxpr_counts.items(): + if c == 1: + continue + name = names[jaxpr] + if (count := name_counts[name]) > 0: + name_counts[name] += 1 + name += str(count) + name_counts[name] += 1 + else: + name_counts[name] += 1 + docs.append(pp_top_level_jaxpr(name, jaxpr, context, settings)) + context.used_names.add(name) + context.top_level_jaxprs[jaxpr] = name + docs.append(pp_jaxpr(jaxpr_to_print, context, settings)) + return pp.concat(docs) + + +class JaxprPpSettings(NamedTuple): + print_shapes: bool = True + source_info: bool = False + name_stack: bool = False + custom_pp_eqn_rules: bool = True + print_effects: bool = False + +def _encode_digits_alphabetic(n: int) -> str: + if n == -1: + return '*' + s = '' + while len(s) == 0 or n: + n, i = n // 26, n % 26 + s = chr(97 + i % 26) + s + return s + +# A JaxprPpContext allows us to globally uniquify variable names within nested +# Jaxprs. +class JaxprPpContext: + var_names: defaultdict[Var, str] + used_names: MutableSet[str] + top_level_jaxprs: MutableMapping[Jaxpr, str] + + def __init__(self) -> None: + self.top_level_jaxprs = {} + self.used_names = set() + fresh_names: Iterator[str] = ( + name + for i in it.count() + if (name := _encode_digits_alphabetic(i)) not in self.used_names + ) + self.var_names = defaultdict(fresh_names.__next__) + + +def pp_var(v: Var | Literal, context: JaxprPpContext) -> str: + if isinstance(v, (Literal, DropVar)): return str(v) + return f"{context.var_names[v]}{v.suffix}" + +def pp_aval(a: AbstractValue, context: JaxprPpContext) -> str: + if isinstance(a, DShapedArray): + shape = [pp_var(d, context) if type(d) is Var else str(d) for d in a.shape] + dtype = _short_dtype_name(a.dtype) + return f'{dtype}[{",".join(shape)}]' + else: + return a.str_short(short_dtypes=True) + +def pp_vars(vs: Sequence[Any], context: JaxprPpContext, + *, separator="", print_shapes: bool = False) -> pp.Doc: + if print_shapes: + return pp.nest(2, pp.group( + pp.join(pp.text(separator) + pp.group(pp.brk()), [ + pp.text(pp_var(v, context)) + + pp.type_annotation(pp.text(":" + pp_aval(v.aval, context))) + for v in vs + ]) + )) + else: + return pp.nest(2, pp.group( + pp.join(pp.text(separator) + pp.group(pp.brk()), + [pp.text(pp_var(v, context)) for v in vs]) + )) + +def pp_kv_pair(k:str, v: Any, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: + if type(v) is tuple and all(isinstance(j, (Jaxpr, ClosedJaxpr)) for j in v): + pp_v = pp_jaxprs(v, context, settings) + elif isinstance(v, Jaxpr): + pp_v = pp_jaxpr(v, context, settings) + elif isinstance(v, ClosedJaxpr): + pp_v = pp_jaxpr(v.jaxpr, context, settings) + else: + pp_v = pp.text(str(v)) + return pp.text(f'{k}=') + pp_v + +def pp_kv_pairs(kv_pairs, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: + if not kv_pairs: + return pp.nil() + return pp.group( + pp.nest(2, pp.concat([ + pp.text("["), pp.brk(""), + pp.join(pp.brk(), [pp_kv_pair(k, v, context, settings) for k, v in kv_pairs]) + ])) + + pp.brk("") + pp.text("]") + ) + +def pp_eqn(eqn: JaxprEqn, context: JaxprPpContext, settings: JaxprPpSettings + ) -> pp.Doc: + rule = (_pp_eqn if not settings.custom_pp_eqn_rules else + pp_eqn_rules.get(eqn.primitive, _pp_eqn)) + doc = rule(eqn, context, settings) # type: ignore[operator] + user_frame = source_info_util.user_frame(eqn.source_info) + return doc if user_frame is None else pp.source_map(doc, user_frame) + +def _pp_eqn(eqn, context, settings, params=None) -> pp.Doc: + annotation = (source_info_util.summarize(eqn.source_info) + if settings.source_info else None) + if params is None: + params = sorted(eqn.params) + name_stack_annotation = f'[{eqn.source_info.name_stack}]' if settings.name_stack else None + lhs = pp_vars(eqn.outvars, context, print_shapes=settings.print_shapes) + rhs = [pp.text(eqn.primitive.name, annotation=name_stack_annotation), + pp_kv_pairs([(p, eqn.params[p]) for p in params], context, settings), + pp.text(" ") + pp_vars(eqn.invars, context)] + if lhs.format(): + return pp.concat([lhs, pp.text(" = ", annotation=annotation), *rhs]) + else: + return pp.concat(rhs) +CustomPpEqnRule = Callable[[JaxprEqn, JaxprPpContext, JaxprPpSettings], pp.Doc] +pp_eqn_rules: dict[Primitive, CustomPpEqnRule] = {} + +def pp_eqns(eqns, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: + return pp.join( + pp.brk("; "), + [pp_eqn(e, context, settings) for e in eqns]) + +def _compact_eqn_should_include(k: str, v: Any) -> bool: + if k == 'branches': return False + if isinstance(v, (Jaxpr, ClosedJaxpr)): return False + if (isinstance(v, tuple) and + any(isinstance(e, (Jaxpr, ClosedJaxpr)) for e in v)): + return False + return True + +def str_eqn_compact(primitive_name: str, params: dict) -> str: + "Compact equation to string conversion used in HLO metadata." + kvs = " ".join(f"{k}={v}" for k, v in params.items() + if _compact_eqn_should_include(k, v)) + return f"{primitive_name}[{kvs}]" if len(kvs) > 0 else primitive_name + +def pp_jaxpr_skeleton(jaxpr, eqns_fn, context: JaxprPpContext, + settings: JaxprPpSettings) -> pp.Doc: + constvars = pp_vars(jaxpr.constvars, context, print_shapes=settings.print_shapes) + invars = pp_vars(jaxpr.invars, context, print_shapes=settings.print_shapes) + eqns = eqns_fn() + outvars = pp.concat([ + pp.text("("), pp_vars(jaxpr.outvars, context, separator=","), + pp.text(")" if len(jaxpr.outvars) != 1 else ",)")]) + if settings.print_effects: + # TODO(sharadmv): render an entire signature here + eff_text = [pp.text(" : { ")] + for i, eff in enumerate(jaxpr.effects): + if i > 0: + eff_text.append(pp.text(", ")) + if isinstance(eff, effects.JaxprInputEffect): + index = eff.input_index + all_vars = [*jaxpr.constvars, *jaxpr.invars] + eff_text.append(pp_effect(eff.replace(input_index=all_vars[index]), + context)) + else: + eff_text.append(pp_effect(eff, context)) + eff_text.append(pp.text(" }")) + else: + eff_text = [] + return pp.group(pp.nest(2, pp.concat([ + pp.text("{ "), pp.keyword(pp.text("lambda ")), + constvars, pp.text("; "), invars, + pp.text(". "), pp.keyword(pp.text("let")), + pp.nest(2, pp.brk() + eqns), pp.brk(), + pp.keyword(pp.text("in ")), outvars, + pp.concat(eff_text) + ])) + pp.text(" }")) + + +def pp_top_level_jaxpr( + name: str, + jaxpr: Jaxpr, + context: JaxprPpContext, + settings: JaxprPpSettings, +) -> pp.Doc: + return pp.concat([ + pp.text("let " + name + " = "), + pp_jaxpr(jaxpr, context, settings), + pp.text(" in"), + pp.brk(), + ]) + + +def pp_jaxpr( + jaxpr: Jaxpr, + context: JaxprPpContext, + settings: JaxprPpSettings, +) -> pp.Doc: + if name := context.top_level_jaxprs.get(jaxpr): + return pp.text(name) + eqns_fn = lambda: pp_eqns(jaxpr.eqns, context, settings) + return pp_jaxpr_skeleton(jaxpr, eqns_fn, context, settings) + + +def pp_jaxprs(jaxprs, context: JaxprPpContext, settings: JaxprPpSettings) -> pp.Doc: + jaxprs = [j.jaxpr if isinstance(j, ClosedJaxpr) else j for j in jaxprs] + return pp.group(pp.nest(2, pp.concat([ + pp.text('('), pp.brk(""), + pp.join(pp.brk(), map(lambda x: pp_jaxpr(x, context, settings), jaxprs))] + )) + pp.brk("") + pp.text(')') + ) + + +def pp_jaxpr_eqn_range(jaxpr: Jaxpr, lo: int, hi: int, context: JaxprPpContext, + settings: JaxprPpSettings) -> pp.Doc: + lo = max(lo, 0) + hi = max(lo, min(hi, len(jaxpr.eqns))) + eqns = jaxpr.eqns[lo:hi] + def eqns_fn(): + pps = [] + if len(eqns) == 0 and len(jaxpr.eqns) != 0: + pps.append(pp.text('...')) + else: + if lo != 0: + pps.append(pp.text('...')) + pps.extend(map((lambda e: pp_eqn(e, context, settings)), eqns)) + if hi != len(jaxpr.eqns): + pps.append(pp.text('...')) + return pp.join(pp.brk("; "), pps) + return pp_jaxpr_skeleton(jaxpr, eqns_fn, context, settings) + +def pp_effect(effect: Effect, context: JaxprPpContext) -> pp.Doc: + if hasattr(effect, "_pretty_print"): + return effect._pretty_print(context) + return pp.text(str(effect)) + +# ------------------- Jaxpr util ------------------- + +def last_used(jaxpr: Jaxpr) -> dict[Var, JaxprEqn | None]: + """Returns a mapping from every var in jaxpr to what equation uses it last.""" + last_used: dict[Var, JaxprEqn | None] = { + v: None for v in jaxpr.outvars if not isinstance(v, Literal)} + for eqn in reversed(jaxpr.eqns): + for v in eqn.invars: + if not isinstance(v, Literal) and v not in last_used: + last_used[v] = eqn + return last_used + +def clean_up_dead_vars(eqn: JaxprEqn, env: dict[Var, Any], + last_used: dict[Var, JaxprEqn | None]): + """Remove all eqn.invars from env if eqn is the last time they were used.""" + for v in {v for v in eqn.invars if not isinstance(v, Literal)}: + if last_used[v] is eqn: + # Delete ref to variable when it is no longer needed by next equations. + del env[v] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cudnn/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cudnn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..862a661e24b9ed93da2e7d199372610dfeab08e9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cudnn/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cudnn/fused_attention_stablehlo.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cudnn/fused_attention_stablehlo.py new file mode 100644 index 0000000000000000000000000000000000000000..a9009da74620ec2ec4db9b5062fe7c21a79a0439 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/cudnn/fused_attention_stablehlo.py @@ -0,0 +1,1004 @@ +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum +from functools import partial, reduce +import operator +from typing import Optional +import json + +import jax +import jax.numpy as jnp +from jax import core +from jax import dtypes +from jax.interpreters import mlir +from jax.interpreters import xla +from jax.interpreters.mlir import ir +from jax.interpreters.mlir import hlo +from jax.experimental.custom_partitioning import custom_partitioning +from jax.sharding import PartitionSpec, NamedSharding + +from jax._src import dispatch +from jax._src.interpreters import batching +from jax._src.lib import cuda_versions + +Array = jnp.ndarray +DType = jnp.dtype +PRNGKey = jnp.ndarray + +class AttentionLayout(Enum): + BTNH = 0 + BNTH = 1 + +class MaskType(Enum): + NO_MASK = 0 + PADDING = 1 + CAUSAL = 2 + PADDING_CAUSAL = 3 + ALIBI = 4 + +def convert_mask_type_to_string(mask_type: MaskType) -> str: + if mask_type == MaskType.NO_MASK: + return "NO_MASK" + elif mask_type == MaskType.PADDING: + return "PADDING" + elif mask_type == MaskType.CAUSAL: + return "CAUSAL" + elif mask_type == MaskType.PADDING_CAUSAL: + return "PADDING_CAUSAL" + elif mask_type == MaskType.ALIBI: + return "ALIBI" + else: + raise ValueError(f"Unexpected mask type: {mask_type}") + +def has_padding(mask_type: MaskType) -> bool: + return mask_type == MaskType.PADDING or mask_type == MaskType.PADDING_CAUSAL + +def should_export_dbias(bias_shape, query_shape, layout) -> bool: + b_B, b_N, _, _ = bias_shape + if layout == AttentionLayout.BNTH.value: + _, q_N, _, _ = query_shape + else: + _, _, q_N, _ = query_shape + return b_B == 1 and b_N == q_N + +def _normalize_layout(layout: str) -> AttentionLayout: + layout_upper = layout.upper() + if layout_upper in ["BSNH", "BNSH", "BTNH", "BNTH"]: + return AttentionLayout[layout_upper.replace("S", "T")] + else: + raise ValueError(f"Unsupported qkv_layout: {layout}") + +def element_type_to_backend_config_type_mapping(dtype): + _element_type_to_backend_config_type_mapping = { + ir.BF16Type.get(): "BF16", + ir.F16Type.get(): "F16", + } + return _element_type_to_backend_config_type_mapping[dtype] + +def default_layouts(*shapes): + return [range(len(shape) - 1, -1, -1) for shape in shapes] + +def create_dot_product_attention_backend_config(batch, + num_heads, + seq_q, + seq_kv, + dtype, + fmha_scale, + seed, + dropout_rate, + mask_type, + layout, + is_bwd): + # Q, K, V: query, key, value in shape of BT(S)NH or BNT(S)H + # P: BMM1 output in shape of BNTS + # O: BMM2 output in the same shape with Q + # BMM1: Q @ K -> P + # BMM2: P @ V -> O + # BMM1Grad1: dP @ Q -> dK + # BMM1Grad2: dP @ K -> dQ + # BMM2Grad1: P @ dO -> dV + # BMM2Grad2: dO @ V -> dP + + cudnn_fmha_backend_config = { + "algorithm": { + "algo_id": "0", + "math_type": "TENSOR_OP_MATH", + "tuning_knobs": {"17": "1", "24": "0"}, + "is_cudnn_frontend": True, + "workspace_size": "0", + }, + "fmha_scale": fmha_scale, + "dropout_rate": dropout_rate, + "intermediate_tensor_shape": { + "element_type": element_type_to_backend_config_type_mapping(dtype), + "dimensions": [str(batch), str(num_heads), str(seq_q), str(seq_kv)], + "tuple_shapes": [], + "layout": { + "dim_level_types": [], + "dim_unique": [], + "dim_ordered": [], + "minor_to_major": ["3", "2", "1", "0"], + "tiles": [], + "element_size_in_bits": "0", + "memory_space": "0", + "index_primitive_type": "PRIMITIVE_TYPE_INVALID", + "pointer_primitive_type": "PRIMITIVE_TYPE_INVALID", + "dynamic_shape_metadata_prefix_bytes": "0", + }, + "is_dynamic_dimension": [False, False, False, False], + }, + "seed": seed, + "is_flash_attention": True, + "mask_type": convert_mask_type_to_string(mask_type), + } + + # We define the contracting and batch dims in the format of + # ((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, + # rhs_batch_dims)). + if layout == AttentionLayout.BNTH.value: + dims = [ + ((3, 3), ((0, 1), (0, 1))), # BMM1: BNTH,BNSH->BNTS + ((3, 2), ((0, 1), (0, 1))), # BMM2: BNTS,BNSH->BNTH + ((2, 2), ((0, 1), (0, 1))), # BMM1_grad_1: BNTS,BNTH->BNSH + ((3, 2), ((0, 1), (0, 1))), # BMM1_grad_2: BNTS,BNSH->BNTH + ((2, 2), ((0, 1), (0, 1))), # BMM2_grad_1: BNTS,BNTH->BNSH + ((3, 3), ((0, 1), (0, 1))), # BMM2_grad_2: BNTH,BNSH->BNTS + ] + else: + dims = [ + ((3, 3), ((0, 2), (0, 2))), # BMM1: BTNH,BSNH->BNTS + ((3, 1), ((0, 1), (0, 2))), # BMM2: BNTS,BSNH->BTNH + ((2, 1), ((0, 1), (0, 2))), # BMM1_grad_1: BNTS,BTNH->BSNH + ((3, 1), ((0, 1), (0, 2))), # BMM1_grad_2: BNTS,BSNH->BTNH + ((2, 1), ((0, 1), (0, 2))), # BMM2_grad_1: BNTS,BTNH->BSNH + ((3, 3), ((0, 2), (0, 2))), # BMM2_grad_2: BTNH,BSNH->BNTS + ] + keys = [ + "bmm1_dot_dimension_numbers", + "bmm2_dot_dimension_numbers", + "bmm1_grad_gemm1_dot_dimension_numbers", + "bmm1_grad_gemm2_dot_dimension_numbers", + "bmm2_grad_gemm1_dot_dimension_numbers", + "bmm2_grad_gemm2_dot_dimension_numbers", + ] + fwd_dot_number = {} + bwd_dot_number = {} + for idx, (key, ((lc, rc), (lb, rb))) in enumerate(zip(keys, dims)): + dims_to_write = fwd_dot_number if idx < 2 else bwd_dot_number + dims_to_write[key] = { + "lhs_contracting_dimensions": [str(lc)], + "rhs_contracting_dimensions": [str(rc)], + "lhs_batch_dimensions": [str(i) for i in lb], + "rhs_batch_dimensions": [str(i) for i in rb], + } + + if is_bwd: + cudnn_fmha_backend_config = {**cudnn_fmha_backend_config, **bwd_dot_number} + else: + cudnn_fmha_backend_config = {**cudnn_fmha_backend_config, **fwd_dot_number} + + backend_config = { + "operation_queue_id":"0", + "wait_on_operation_queues":[], + "cudnn_fmha_backend_config": cudnn_fmha_backend_config + } + backend_config = json.dumps(backend_config) + return backend_config + +# mapping from (is_bwd, has_dropout, has_mask, has_bias) to custom call name +_custom_name_maps = { + # fMHA forward call targets. + (False, False, False): "__cudnn$fmhaSoftmax", + (False, False, True): "__cudnn$fmhaScaleBiasSoftmax", + (False, True, False): "__cudnn$fmhaSoftmaxDropout", + (False, True, True): "__cudnn$fmhaScaleBiasSoftmaxDropout", + # fMHA backward call targets. + (True, False, False): "__cudnn$fmhaSoftmaxBackward", + (True, False, True): "__cudnn$fmhaScaleBiasSoftmaxBackward", + (True, True, False): "__cudnn$fmhaSoftmaxDropoutBackward", + (True, True, True): "__cudnn$fmhaScaleBiasSoftmaxDropoutBackward", +} + +def get_custom_call_name(has_bias, has_dropout, is_bwd): + return _custom_name_maps[(is_bwd, has_dropout, has_bias)] + +def check_layout(query, key, value, bias, q_seqlen, kv_seqlen, layout): + def check_eq(a, b, c, msg): + if not (a == b == c): + raise ValueError(f"{msg} must be same, got {a}, {b}, {b}") + + q_rank, k_rank, v_rank = len(query.shape), len(key.shape), len(value.shape) + if q_rank != 4: + raise ValueError(f"Q must have a rank of 4, got {q_rank}") + check_eq(q_rank, k_rank, v_rank, "QKV rank") + + q_dtype, k_dtype, v_dtype = query.dtype, key.dtype, value.dtype + assert q_dtype in [jnp.float16, jnp.bfloat16], "Q must be fp16 or bf16" + check_eq(q_dtype, k_dtype, v_dtype, "QKV dtype") + + if layout == AttentionLayout.BNTH: + qB, qN, qT, qH = query.shape + kB, kN, kS, kH = key.shape + vB, vN, vS, vH = value.shape + else: + assert layout == AttentionLayout.BTNH + qB, qT, qN, qH = query.shape + kB, kS, kN, kH = key.shape + vB, vS, vN, vH = value.shape + + check_eq(qB, kB, vB, "QKV batch") + check_eq(qH, kH, vH, "QKV dim_per_head") + if kN != vN: + raise ValueError(f"KV must have same number of heads, got {kN} vs {vN}") + if kS != vS: + raise ValueError(f"KV must have same seq length, got {kS} vs {vS}") + + # check bias/q_seqlen/kv_seqlen + if bias is not None: + _, _, bT, bS = bias.shape + if bT != qT or bS != vS: + raise ValueError( + f"Bias must have same seq length as QKV, got {bT} and {bS}") + if q_seqlen is not None: + q_seq_dtype = q_seqlen.dtype + q_seq_rank = len(q_seqlen.shape) + if q_seq_dtype != jnp.int32: + raise ValueError(f"q_seqlen must have int32 datatype, got {q_seq_dtype}") + if q_seq_rank != 1: + raise ValueError(f"q_seqlen must have a rank of 1, got {q_seq_rank}") + q_seq_b = q_seqlen.shape[0] + if q_seq_b != qB: + raise ValueError(f"q_seqlen must have same batch as Q, got {q_seq_b}") + if kv_seqlen is not None: + kv_seq_dtype = kv_seqlen.dtype + kv_seq_rank = len(kv_seqlen.shape) + if kv_seq_dtype != jnp.int32: + raise ValueError( + f"kv_seqlen must have int32 datatype, got {kv_seq_dtype}") + if kv_seq_rank != 1: + raise ValueError(f"kv_seq_rank must have a rank of 1, got {kv_seq_rank}") + kv_seq_b = kv_seqlen.shape[0] + if kv_seq_b != qB: + raise ValueError(f"kv_seqlen must have same batch as Q, got {kv_seq_b}") + +def check_is_flash_attention( + query, key, layout, cudnn_version, has_bias, is_training): + if layout == AttentionLayout.BNTH: + _, _, T, H = query.shape + _, _, S, _ = key.shape + else: + _, T, _, H = query.shape + _, S, _, _ = key.shape + + if not ((H <= 128 and H % 8 == 0) and + (not is_training or not has_bias or T % 2 == 0 and S % 2 == 0)): + # check if flash attention is supported + # for training, for patterns with bias, seqlen should be divisible by 2 + raise NotImplementedError( + f"Unsupported sequence length Q {T}, KV {S} and head dim {H}.") + # check if minimum cudnn version requirement is satisfied + if cudnn_version < 8904: + raise RuntimeError( + "JAX requires cuDNN >= 8.9.4 to use flash cross attention.") + +def check_cudnn_version(): + # check if cuDNN is installed + if cuda_versions is None: + raise RuntimeError("cuDNN is not detected.") + return cuda_versions.cudnn_get_version() + +def check_compute_capability(cc): + if cuda_versions is None: + raise RuntimeError("cuDNN is not detected.") + for i in range(jax.device_count()): + compute_cap = cuda_versions.cuda_compute_capability(i) + if compute_cap not in cc: + raise RuntimeError("Require compute capability in " + str(cc)) + +def _dot_product_attention_fwd( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale, seed, + dropout_rate, variadic_args, mask_type, layout, cudnn_version): + # check if flash attention is supported for this attention pattern + check_is_flash_attention( + query, key, layout, cudnn_version, bias is not None, False) + outputs = _dot_product_attention_fwd_p_wrapper.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale=scale, + seed=seed, dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, is_training=False) + output = outputs[0] + return output + +def _dot_product_attention_fwd_rule( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale, seed, + dropout_rate, variadic_args, mask_type, layout, cudnn_version): + # check if flash attention is supported for this attention pattern + check_is_flash_attention( + query, key, layout, cudnn_version, bias is not None, True) + outputs = _dot_product_attention_fwd_p_wrapper.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale=scale, + seed=seed, dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, is_training=True) + res = (query, key, value, bias, mask, q_seqlen, kv_seqlen, + outputs[1], outputs[0]) + return outputs[0], res + +def _dot_product_attention_bwd_rule( + scale, seed, dropout_rate, variadic_args, mask_type, layout, is_training, + res, grad_output): + (query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, + fwd_output) = res + grads = _dot_product_attention_bwd_p_wrapper.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, + fwd_output, grad_output, scale=scale, seed=seed, + dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, + ) + grads = (*grads,) + (None,) * (7 - len(grads)) + return grads + +def _dot_product_attention_fwd_impl( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale, seed, + dropout_rate, variadic_args, mask_type, layout, is_training): + # args: {Q, K, V, mask*, bias*} + outputs = _dot_product_attention_fwd_p.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale=scale, + seed=seed, dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, is_training=is_training) + return outputs + +def _dot_product_attention_bwd_impl( + query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, fwd_output, + grad_output, scale, seed, dropout_rate, variadic_args, mask_type, layout): + grads = _dot_product_attention_bwd_p.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, + fwd_output, grad_output, scale=scale, seed=seed, + dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout) + return grads + +def _dot_product_attention_fwd_abstract( + query, key, value, bias, mask, q_seqlen, kv_seqlen, *, scale, seed, + dropout_rate, variadic_args, mask_type, layout, is_training): + query_dtype = dtypes.canonicalize_dtype(query.dtype) + if layout == AttentionLayout.BNTH.value: + B, N, T, _ = query.shape + _, _, S, _ = key.shape + else: + B, T, N, _ = query.shape + _, S, _, _ = key.shape + output_shape = query.shape + softmax_stat_shape = (B, N, T) + + if is_training: + return ( + core.ShapedArray(output_shape, query_dtype), # output + core.ShapedArray(softmax_stat_shape, jnp.float32), # softmax_stat + ) + else: + return ( + core.ShapedArray(output_shape, query_dtype), # output + ) + +def _dot_product_attention_bwd_abstract( + query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, fwd_output, + grad_output, *, scale, seed, dropout_rate, variadic_args, mask_type, + layout): + query_dtype = dtypes.canonicalize_dtype(query.dtype) + key_dtype = dtypes.canonicalize_dtype(key.dtype) + value_dtype = dtypes.canonicalize_dtype(value.dtype) + + _, _, has_dbias = variadic_args + if has_dbias: + # cuDNN supports bias for this case + bias_dtype = dtypes.canonicalize_dtype(bias.dtype) + return ( + core.ShapedArray( + query.shape, query_dtype + ), # grad query + core.ShapedArray( + key.shape, key_dtype + ), # grad key + core.ShapedArray( + value.shape, value_dtype + ), # grad value + core.ShapedArray( + bias.shape, bias_dtype + ), # grad bias + ) + else: + return ( + core.ShapedArray( + query.shape, query_dtype + ), # grad query + core.ShapedArray( + key.shape, key_dtype + ), # grad key + core.ShapedArray( + value.shape, value_dtype + ), # grad value + ) + +def _dot_product_attention_fwd_cuda_lowering( + ctx, query, key, value, bias, mask, q_seqlen, kv_seqlen, scale, seed, + dropout_rate, variadic_args, mask_type, layout, is_training): + query_type = ir.RankedTensorType(query.type) + query_shape = query_type.shape + key_type = ir.RankedTensorType(key.type) + key_shape = key_type.shape + + if layout == AttentionLayout.BNTH.value: + B, N, T, H = query_shape + _, _, S, _ = key_shape + output_layout = (3, 2, 1, 0) + output_transpose_perm = mlir.dense_int_array((0, 1, 2, 3)) + else: + B, T, N, H = query_shape + _, S, _, _ = key_shape + output_layout = (3, 1, 2, 0) + output_transpose_perm = mlir.dense_int_array((0, 2, 1, 3)) + + output_shape = (B, N, T, H) + softmax_stat_shape = (B, N, T) + workspace_shape = (0,) + workspace_type = ir.IntegerType.get_unsigned(8) + backend_config = create_dot_product_attention_backend_config( + B, N, T, S, query_type.element_type, scale, seed, dropout_rate, + mask_type, layout, is_bwd=False, + ) + # {Q, K, V, mask*, bias*, q_seqlen*, kv_seqlen*} + # {output, activation*, workspace} + has_dropout = dropout_rate > 0 + has_bias, has_mask, _ = variadic_args + operands = [query, key, value] + if has_mask: + operands.append(mask) + if has_bias: + operands.append(bias) + if has_padding(mask_type): + operands.append(q_seqlen) + operands.append(kv_seqlen) + custom_call_name = get_custom_call_name(has_bias, has_dropout, False) + # create output types and layouts + if is_training: + result_types = [ + ir.RankedTensorType.get(output_shape, query_type.element_type), + ir.RankedTensorType.get(softmax_stat_shape, ir.F32Type.get()), + ir.RankedTensorType.get(workspace_shape, workspace_type), + ] + result_layouts = [output_layout] + default_layouts(softmax_stat_shape, workspace_shape) + else: + result_types = [ + ir.RankedTensorType.get(output_shape, query_type.element_type), + ir.RankedTensorType.get(workspace_shape, workspace_type) + ] + result_layouts = [output_layout] + default_layouts(workspace_shape) + # create custom call here + out = mlir.custom_call( + custom_call_name, + result_types=result_types, + operands=operands, + backend_config=backend_config, + operand_layouts=default_layouts( + *[ir.RankedTensorType(operand.type).shape for operand in operands]), + result_layouts=result_layouts, + ) + # drop workspace memory + # output should be (B, T, N, H) instead of (B, N, T, H) + if is_training: + return [hlo.transpose(out.results[0], output_transpose_perm), out.results[1]] + else: + return [hlo.transpose(out.results[0], output_transpose_perm)] + +def _dot_product_attention_bwd_cuda_lowering( + ctx, query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, + fwd_output, grad_output, scale, seed, dropout_rate, variadic_args, + mask_type, layout): + query_type = ir.RankedTensorType(query.type) + query_shape = query_type.shape + key_type = ir.RankedTensorType(key.type) + key_shape = key_type.shape + value_type = ir.RankedTensorType(value.type) + + if layout == AttentionLayout.BNTH.value: + B, q_N, T, H = query_shape + _, k_N, S, _ = key_shape + grad_layout = (3, 2, 1, 0) + grad_transpose_perm = mlir.dense_int_array((0, 1, 2, 3)) + else: + B, T, q_N, H = query_shape + _, S, k_N, _ = key_shape + grad_layout = (3, 1, 2, 0) + grad_transpose_perm = mlir.dense_int_array((0, 2, 1, 3)) + + workspace_shape = (0,) + workspace_type = ir.IntegerType.get_unsigned(8) + + grad_query_shape = (B, q_N, T, H) + grad_key_shape = (B, k_N, S, H) + grad_value_shape = (B, k_N, S, H) + backend_config = create_dot_product_attention_backend_config( + B, q_N, T, S, query_type.element_type, scale, seed, dropout_rate, + mask_type, layout, is_bwd=True, + ) + # {Q, K, V, activation, dO, mask*, bias*, O, q_seqlen*, kv_seqlen*} + # {dQ, dK, dV, dbias*, workspace} + has_dropout = dropout_rate > 0 + has_bias, has_mask, has_dbias = variadic_args + # create operands + operands = [query, key, value, activation, grad_output] + if has_mask: + operands.append(mask) + if has_bias: + # flash attention requires bias in the bwd for remat + operands.append(bias) + operands.append(fwd_output) + if has_padding(mask_type): + operands.append(q_seqlen) + operands.append(kv_seqlen) + # get custom call name + custom_call_name = get_custom_call_name(has_bias, has_dropout, True) + + # create output types and layouts + # grad_query, grad_key, grad_value + result_types = [ + ir.RankedTensorType.get(grad_query_shape, query_type.element_type), + ir.RankedTensorType.get(grad_key_shape, key_type.element_type), + ir.RankedTensorType.get(grad_value_shape, value_type.element_type), + ] + result_layouts = [grad_layout, grad_layout, grad_layout] + bias_type = ir.RankedTensorType(bias.type) + bias_shape = bias_type.shape + if has_dbias: + # cuDNN supports bias for this case + result_types.append( + ir.RankedTensorType.get(bias_shape, bias_type.element_type)) + result_layouts = result_layouts + default_layouts(bias_shape) + # workspace + result_types.append(ir.RankedTensorType.get(workspace_shape, workspace_type)) + result_layouts = result_layouts + default_layouts(workspace_shape) + out = mlir.custom_call( + custom_call_name, + result_types=result_types, + operands=operands, + backend_config=backend_config, + operand_layouts=default_layouts( + *[ir.RankedTensorType(operand.type).shape for operand in operands]), + result_layouts=result_layouts, + ) + dqkv = (hlo.transpose(out.results[0], grad_transpose_perm), + hlo.transpose(out.results[1], grad_transpose_perm), + hlo.transpose(out.results[2], grad_transpose_perm)) + # Only keep dQ, dK, dV and dBias here + if has_dbias: + return dqkv + (out.results[3],) + else: + return dqkv + +# batcher +def _check_valid_batch_dims(bdims): + for dim in bdims: + if dim not in [0, None]: + raise NotImplementedError( + f"Currently only support batch_dim in [0, None], but got {dim=}") + +def _dot_product_attention_fwd_batcher( + batched_args, batch_dims, *, scale, seed, dropout_rate, variadic_args, + mask_type, layout, is_training): + _check_valid_batch_dims(batch_dims) + query, key, value, bias, mask, q_seqlen, kv_seqlen = batched_args + query_bdim = batch_dims[0] + if is_training: + out_bdims = query_bdim, query_bdim + else: + out_bdims = (query_bdim,) + + if layout == AttentionLayout.BNTH.value: + *Bs, N, T, _ = query.shape + *_, _, S, _ = key.shape + else: + *Bs, T, N, _ = query.shape + *_, S, _, _ = key.shape + B = reduce(operator.mul, Bs) + has_bias, has_mask, _ = variadic_args + # reshape to 4D shape + query = jnp.reshape(query, (B,) + query.shape[-3:]) + key = jnp.reshape(key, (B,) + key.shape[-3:]) + value = jnp.reshape(value, (B,) + key.shape[-3:]) + if has_bias: + bias = jnp.reshape(bias, (B, N, T, S)) + if has_mask: + mask = jnp.reshape(mask, (B, N, T, S)) + if has_padding(mask_type): + q_seqlen = jnp.reshape(q_seqlen, (B, )) + kv_seqlen = jnp.reshape(kv_seqlen, (B, )) + + outputs = _dot_product_attention_fwd_p_wrapper.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale=scale, + seed=seed, dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, is_training=is_training) + + # reshape to original shape + output = outputs[0] + output = jnp.reshape(output, query.shape) + if is_training: + activation = outputs[1] + activation = jnp.reshape(activation, (*Bs, N, T)) + return (output, activation), out_bdims + else: + return (output,), out_bdims + +def _dot_product_attention_bwd_batcher( + batched_args, batch_dims, *, scale, seed, dropout_rate, variadic_args, + mask_type, layout): + _check_valid_batch_dims(batch_dims) + query, key, value, bias, mask, q_seqlen, \ + kv_seqlen, activation, fwd_output, grad_output = batched_args + query_bdim = batch_dims[0] + out_bdims = query_bdim, query_bdim, query_bdim + + if layout == AttentionLayout.BNTH.value: + *Bs, N, T, _ = query.shape + *_, _, S, _ = key.shape + else: + *Bs, T, N, _ = query.shape + *_, S, _, _ = key.shape + B = reduce(operator.mul, Bs) + has_bias, has_mask, has_dbias = variadic_args + # reshape to 4D shape + query = jnp.reshape(query, (B,) + query.shape[-3:]) + key = jnp.reshape(key, (B,) + key.shape[-3:]) + value = jnp.reshape(value, (B,) + key.shape[-3:]) + if has_bias: + bias = jnp.reshape(bias, (B, N, T, S)) + if has_mask: + mask = jnp.reshape(mask, (B, N, T, S)) + if has_padding(mask_type): + q_seqlen = jnp.reshape(q_seqlen, (B, )) + kv_seqlen = jnp.reshape(kv_seqlen, (B, )) + + activation = jnp.reshape(activation, (B, N, T)) + fwd_output = jnp.reshape(fwd_output, (B,) + query.shape[-3:]) + grad_output = jnp.reshape(grad_output, (B,) + query.shape[-3:]) + + grads = _dot_product_attention_bwd_p_wrapper.bind( + query, key, value, bias, mask, q_seqlen, kv_seqlen, activation, + fwd_output, grad_output, scale=scale, seed=seed, + dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, + ) + + grad_query, grad_key, grad_value = grads[:3] + # reshape to original shape + grad_query = jnp.reshape(grad_query, query.shape) + grad_key = jnp.reshape(grad_key, key.shape) + grad_value = jnp.reshape(grad_value, value.shape) + if has_dbias: + grad_bias = grads[3] + grad_bias = jnp.reshape(grad_bias, bias.shape) + return grads + (grad_bias,), out_bdims + (query_bdim,) + return grads, out_bdims + +# custom partitioning +def _get_padded_spec(arg_info): + spec = None if arg_info.sharding is None else arg_info.sharding.spec + ndim = arg_info.ndim + if spec is None: + return (None,) * ndim + assert len(spec) <= ndim + return spec + (None,) * (ndim - len(spec)) + +def _check_qkv_bias_mask_spec( + query_spec, key_spec, value_spec, bias_spec, mask_spec): + # check qkv spec + if not query_spec == key_spec == value_spec: + raise ValueError("Query, key and value should have same sharding.") + *batch_spec, q_seq_spec, num_head_spec, head_spec = query_spec + if q_seq_spec is not None: + raise ValueError("Sharding on sequence dim is not allowed.") + if head_spec is not None: + raise ValueError("Sharding on head dim is not allowed.") + # check bias and mask spec + if bias_spec: + *bias_batch_spec, bias_num_head_spec, bias_q_seq_spec, bias_kv_seq_spec = bias_spec + if bias_batch_spec != batch_spec or bias_num_head_spec != num_head_spec: + raise ValueError( + "Query and bias should have same sharding on batch and num_head dim.") + if bias_q_seq_spec is not None or bias_kv_seq_spec is not None: + raise ValueError("Sharding on bias sequence dim is not allowed.") + if mask_spec: + *mask_batch_spec, mask_num_head_spec, mask_q_seq_spec, mask_kv_seq_spec = mask_spec + if mask_batch_spec != batch_spec or mask_num_head_spec != num_head_spec: + raise ValueError( + "Query and mask should have same sharding on batch and num_head dim.") + if mask_q_seq_spec is not None or mask_kv_seq_spec is not None: + raise ValueError("Sharding on mask sequence dim is not allowed.") + +# fwd custom partition +def _infer_fwd_output_sharding(mesh, arg_shapes, variadic_args, is_training): + # only sharding on batch and num_head dim is allowed + # (*batch, q_seq, num_head, head) + query_spec = _get_padded_spec(arg_shapes[0]) + # (*batch, kv_seq, num_head, head) + key_spec = _get_padded_spec(arg_shapes[1]) + value_spec = _get_padded_spec(arg_shapes[2]) + has_bias, has_mask, _ = variadic_args + bias_spec = _get_padded_spec(arg_shapes[3]) if has_bias else None + mask_spec = _get_padded_spec(arg_shapes[4]) if has_mask else None + _check_qkv_bias_mask_spec( + query_spec, key_spec, value_spec, bias_spec, mask_spec) + # keep out sharding same as query sharding since they have same shape + out_sharding = NamedSharding(mesh, PartitionSpec(*query_spec)) + if is_training: + # activation sharding + *batch_spec, q_seq_spec, num_head_spec, _ = query_spec + activation_sharding = NamedSharding( + mesh, PartitionSpec(*batch_spec, num_head_spec, q_seq_spec, None)) + return [out_sharding, activation_sharding] + return [out_sharding] + +_dot_product_attention_fwd_lower = custom_partitioning( + _dot_product_attention_fwd_impl, static_argnums=(7, 8, 9, 10, 11, 12, 13)) + +def _dot_product_attention_fwd_infer_sharding_from_operands( + scale, seed, dropout_rate, variadic_args, mask_type, layout, is_training, + mesh, arg_shapes, result_shape): + return _infer_fwd_output_sharding(mesh, arg_shapes, variadic_args, is_training) + +def _dot_product_attention_fwd_partition( + scale, seed, dropout_rate, variadic_args, mask_type, layout, is_training, + mesh, arg_shapes, result_shape): + # args sharding + arg_shardings = tuple([arg_i.sharding for arg_i in arg_shapes]) + out_shardings = _infer_fwd_output_sharding( + mesh, arg_shapes, variadic_args, is_training) + impl = partial( + _dot_product_attention_fwd_impl, scale=scale, seed=seed, + dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, + layout=layout, is_training=is_training) + return mesh, impl, out_shardings, arg_shardings + +# bwd custom partition +def _infer_bwd_output_sharding(mesh, arg_shapes, variadic_args): + # (*batch, q_seq, num_head, head) + query_spec = _get_padded_spec(arg_shapes[0]) + # (*batch, kv_seq, num_head, head) + key_spec = _get_padded_spec(arg_shapes[1]) + value_spec = _get_padded_spec(arg_shapes[2]) + has_bias, has_mask, has_dbias = variadic_args + bias_spec = _get_padded_spec(arg_shapes[3]) if has_bias else None + mask_spec = _get_padded_spec(arg_shapes[4]) if has_mask else None + _check_qkv_bias_mask_spec( + query_spec, key_spec, value_spec, bias_spec, mask_spec) + # keep grad query sharding same as query sharding + grad_query_sharding = NamedSharding(mesh, PartitionSpec(*query_spec)) + grad_key_sharding = NamedSharding(mesh, PartitionSpec(*key_spec)) + grad_value_sharding = NamedSharding(mesh, PartitionSpec(*key_spec)) + out_shardings = [grad_query_sharding, grad_key_sharding, grad_value_sharding] + if has_dbias: + grad_bias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) + out_shardings = out_shardings + [grad_bias_sharding] + return out_shardings + +_dot_product_attention_bwd_lower = custom_partitioning( + _dot_product_attention_bwd_impl, static_argnums=(10, 11, 12, 13, 14, 15) +) + +def _dot_product_attention_bwd_infer_sharding_from_operands( + scale, seed, dropout_rate, variadic_args, mask_type, layout, mesh, + arg_shapes, result_shape): + return _infer_bwd_output_sharding(mesh, arg_shapes, variadic_args) + +def _dot_product_attention_bwd_partition( + scale, seed, dropout_rate, variadic_args, mask_type, layout, mesh, + arg_shapes, result_shape): + out_shardings = _infer_bwd_output_sharding(mesh, arg_shapes, variadic_args) + # args sharding + arg_shardings = tuple([arg_i.sharding for arg_i in arg_shapes]) + impl = partial( + _dot_product_attention_bwd_impl, scale=scale, seed=seed, + dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, + layout=layout, + ) + return mesh, impl, out_shardings, arg_shardings + +# Create dot_product_attention_fwd_p for forward operation. +_dot_product_attention_fwd_p = core.Primitive("dot_product_attention_fwd") +_dot_product_attention_fwd_p.multiple_results = True +_dot_product_attention_fwd_p.def_impl( + partial(xla.apply_primitive, _dot_product_attention_fwd_p) +) +_dot_product_attention_fwd_p.def_abstract_eval( + _dot_product_attention_fwd_abstract +) + +mlir.register_lowering( + _dot_product_attention_fwd_p, + _dot_product_attention_fwd_cuda_lowering, + platform="cuda", +) + +_dot_product_attention_fwd_p_wrapper = core.Primitive( + "dot_product_attention_fwd_wrapper" +) +_dot_product_attention_fwd_p_wrapper.multiple_results = True +_dot_product_attention_fwd_p_wrapper.def_impl(_dot_product_attention_fwd_impl) +_dot_product_attention_fwd_p_wrapper.def_abstract_eval( + _dot_product_attention_fwd_abstract +) + +# Create dot_product_attention_bwd_p for backward operation. +_dot_product_attention_bwd_p = core.Primitive("dot_product_attention_bwd") +_dot_product_attention_bwd_p.multiple_results = True +_dot_product_attention_bwd_p.def_impl( + partial(xla.apply_primitive, _dot_product_attention_bwd_p) +) +_dot_product_attention_bwd_p.def_abstract_eval( + _dot_product_attention_bwd_abstract +) + +mlir.register_lowering( + _dot_product_attention_bwd_p, + _dot_product_attention_bwd_cuda_lowering, + platform="cuda", +) + +_dot_product_attention_bwd_p_wrapper = core.Primitive( + "dot_product_attention_bwd_wrapper" +) +_dot_product_attention_bwd_p_wrapper.multiple_results = True +_dot_product_attention_bwd_p_wrapper.def_impl(_dot_product_attention_bwd_impl) +_dot_product_attention_bwd_p_wrapper.def_abstract_eval( + _dot_product_attention_bwd_abstract +) + +batching.primitive_batchers[ + _dot_product_attention_fwd_p_wrapper +] = _dot_product_attention_fwd_batcher +batching.primitive_batchers[ + _dot_product_attention_bwd_p_wrapper +] = _dot_product_attention_bwd_batcher + +_dot_product_attention_fwd_lower.def_partition( + infer_sharding_from_operands=_dot_product_attention_fwd_infer_sharding_from_operands, + partition=_dot_product_attention_fwd_partition) + +mlir.register_lowering(_dot_product_attention_fwd_p_wrapper, + mlir.lower_fun(_dot_product_attention_fwd_lower, multiple_results=True)) + +_dot_product_attention_bwd_lower.def_partition( + infer_sharding_from_operands=_dot_product_attention_bwd_infer_sharding_from_operands, + partition=_dot_product_attention_bwd_partition) + +mlir.register_lowering(_dot_product_attention_bwd_p_wrapper, + mlir.lower_fun(_dot_product_attention_bwd_lower, multiple_results=True)) + +dispatch.prim_requires_devices_during_lowering.add( + _dot_product_attention_fwd_p +) +dispatch.prim_requires_devices_during_lowering.add( + _dot_product_attention_fwd_p_wrapper +) +dispatch.prim_requires_devices_during_lowering.add( + _dot_product_attention_bwd_p +) +dispatch.prim_requires_devices_during_lowering.add( + _dot_product_attention_bwd_p_wrapper +) + +@partial(jax.custom_vjp, nondiff_argnums=(7, 8, 9, 10, 11, 12, 13)) +def _dot_product_attention(query: Array, + key: Array, + value: Array, + bias: Array, + mask: Array, + q_seqlen: Array, + kv_seqlen: Array, + scale: float, + seed: int, + dropout_rate: float, + variadic_args: tuple[bool, ...], + mask_type: bool, + layout: int, + cudnn_version: int): + output = _dot_product_attention_fwd( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale=scale, + seed=seed, dropout_rate=dropout_rate, variadic_args=variadic_args, + mask_type=mask_type, layout=layout, cudnn_version=cudnn_version) + return output + +# _dot_product_attention_fwd must have the same func signature as _dot_product_attention +_dot_product_attention.defvjp(_dot_product_attention_fwd_rule, _dot_product_attention_bwd_rule) + +# User interface +def dot_product_attention(query: Array, + key: Array, + value: Array, + bias: Optional[Array] = None, + mask: Optional[Array] = None, + q_seqlen: Optional[Array] = None, + kv_seqlen: Optional[Array] = None, + *, + scale: float = 1.0, + mask_type: MaskType = MaskType.NO_MASK, + seed: int = 42, + dropout_rate: float = 0., + qkv_layout: str = "BTNH"): + """Computes dot-product attention given query (Q), key (K), and value (V). + + This function serves as the core operation for applying attention + mechanisms as described in the paper [https://arxiv.org/abs/1706.03762]. + Initially, it determines the attention weights by processing Q and K, + subsequently combining the outcomes using K. Throughout this function, we + utilize the following uppercase letters to represent specific parameters of + array: + + B = batch size + S = length of the key/value (source) + T = length of the query (target) + N = number of attention heads + H = dimensions of each attention head. + + The supported layouts for Q, K, V are either BT(S)NH or BNT(S)H, and they must + adhere to the same layout. The output layout remains consistent with Q, + defaulting to BT(S)NH. + + Args: + query: Queries for attention calculation with a shape of BTNH or BNTH. + key: Keys for attention calculation with a shape of BSNH or BNSH. + value: Values to be used in attention with a shape of BSNH or BNSH. + bias: Bias to be added to logits with a shape of BNTS. + mask: Mask used to filter out logits with a shape of BNTS. + q_seqlen: Non padded sequence length of Queries with a shape of B. + kv_seqlen: Non padded sequence length of Keys and Values with a shape of B. + scale: Scale for the query. + dropout_rate: Dropout rate. + qkv_layout: Layout string, with supported formats being BTNH, BNTH, BSNH, + BNSH. + + Returns: + Output of the same shape as the query. + """ + # check if cuDNN is installed + cudnn_version = check_cudnn_version() + # only support Ampere and Hopper for now + check_compute_capability((80, 90)) + layout = _normalize_layout(qkv_layout) + if bias is not None: + # reshape bias to have 4D shape + bias = bias.reshape((1,) * (4 - len(bias.shape)) + bias.shape) + # check if input shape and data type is compatiable + check_layout(query, key, value, bias, q_seqlen, kv_seqlen, layout) + if has_padding(mask_type) and (q_seqlen is None or kv_seqlen is None): + raise ValueError("Require q_seqlen and kv_seqlen to generate padding mask") + has_bias = bias is not None + has_mask = mask is not None + has_dbias = has_bias and \ + should_export_dbias(bias.shape, query.shape, layout) # type: ignore[union-attr] + variadic_args = (has_bias, has_mask, has_dbias) + if bias is None: + bias = jnp.zeros(0, dtype=query.dtype) + if mask is None: + mask = jnp.zeros(0, dtype=query.dtype) + if q_seqlen is None: + q_seqlen = jnp.zeros(0, dtype=query.dtype) + if kv_seqlen is None: + kv_seqlen = jnp.zeros(0, dtype=query.dtype) + output = _dot_product_attention( + query, key, value, bias, mask, q_seqlen, kv_seqlen, scale, seed, + dropout_rate, variadic_args, mask_type, layout.value, cudnn_version + ) + return output diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_api_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_api_util.py new file mode 100644 index 0000000000000000000000000000000000000000..76b421f374033c3e5b873a3a2ccb054457f4b5d7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_api_util.py @@ -0,0 +1,26 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +_custom_wrapper_types = set() + +def register_custom_decorator_type(cls): + _custom_wrapper_types.add(cls) + return cls + +def forward_attr(self_, name): + if name.startswith('def') and type(self_.fun) in _custom_wrapper_types: + return getattr(self_.fun, name) + else: + raise AttributeError diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_batching.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_batching.py new file mode 100644 index 0000000000000000000000000000000000000000..22592b854d143e0110a703e3e17f1e2fad5f2988 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_batching.py @@ -0,0 +1,285 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import functools +import operator +from typing import Callable + +from jax import lax +from jax._src import api +from jax._src import core +from jax._src import custom_api_util +from jax._src import linear_util as lu +from jax._src import source_info_util +from jax._src import traceback_util +from jax._src import tree_util +from jax._src import util +from jax._src.api_util import flatten_fun_nokwargs +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters.batching import not_mapped +from jax._src.interpreters import mlir +from jax._src.interpreters import partial_eval as pe +from jax._src.interpreters import xla +from jax._src.tree_util import (tree_flatten, tree_map, tree_structure, + tree_unflatten, treedef_tuple) + + +source_info_util.register_exclusion(__file__) +traceback_util.register_exclusion(__file__) + + +map, unsafe_map = util.safe_map, map +zip, unsafe_zip = util.safe_zip, zip + + +@custom_api_util.register_custom_decorator_type +class custom_vmap: + fun: Callable + vmap_rule: Callable | None + + def __init__(self, fun: Callable): + functools.update_wrapper(self, fun) + self.fun = fun + self.vmap_rule = None + + __getattr__ = custom_api_util.forward_attr + + def def_vmap(self, vmap_rule: Callable) -> Callable: + self.vmap_rule = vmap_rule + return vmap_rule + + @traceback_util.api_boundary + def __call__(self, *args, **kwargs): + assert not kwargs + args_flat, in_tree = tree_flatten(args) + flat_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(self.fun), in_tree) + in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat] + debug = pe.debug_info(self.fun, in_tree, out_tree, False, "custom_vmap") + jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals, debug) + closed_call = core.ClosedJaxpr(pe.convert_constvars_jaxpr(jaxpr), ()) + in_tree = treedef_tuple((tree_structure(consts), in_tree)) + assert self.vmap_rule is not None + out_flat = custom_vmap_p.bind(*consts, *args_flat, + call=closed_call, + rule=ClosedRule(self.vmap_rule), + in_tree=in_tree, + out_tree=out_tree()) + return tree_unflatten(out_tree(), out_flat) + + +### utils + +# Define a class, instead of making a function closing over `rule`, so +# that we can override __str__ +class ClosedRule: + def __init__(self, rule): + functools.update_wrapper(self, rule) + self.rule = rule + + def __call__(self, axis_size, all_in_batched, *all_args): + _, args = all_args + consts_batched, in_batched = all_in_batched + assert not any(tree_util.tree_leaves(consts_batched)), consts_batched + return call_rule(self.rule, axis_size, in_batched, args) + + def __str__(self): + return str(self.rule) + +def ensure_list(xs): + return xs if type(xs) is list else list(xs) + +def rule_name(rule): + return getattr(rule, '__name__', '') + +def call_rule(rule, axis_size, in_batched, args): + return rule(axis_size, ensure_list(in_batched), *args) + +def check_vmap_rule_trees(rule, original_out_tree, out_tree, out_batched_tree): + if out_tree != out_batched_tree: + raise ValueError( + 'structure of output value and output batching specification returned ' + f'by custom vmap rule ({rule_name(rule)}) do not match.\n' + f'Output values: {out_tree}\n' + f'Batching spec: {out_batched_tree}') + if out_tree != original_out_tree: + raise ValueError( + f'structure of output returned by custom vmap rule ({rule_name(rule)}) ' + 'does not match that of original custom-vmapped function.\n' + f'Original output: {original_out_tree}\n' + f'Rule output: {out_tree}') + +# Like batching.bdim_at_front, but doesn't broadcast if not mapped +def maybe_bdim_at_front(x, bdim): + if bdim is not_mapped: + return x + else: + return util.moveaxis(x, bdim, 0) + +# Like batching.batch except (a) not curried and (b) returns inferred output +# axes instead of accepting and matching a given spec of output axes. Assumes +# `f` is pytree-flattened +def vmap_unrestricted(f: lu.WrappedFun, *args, in_axes, axis_name, axis_size): + f, out_axes = batching.batch_subtrace(f) + f = batching._batch_outer(f, axis_name, axis_size, in_axes, + batching.BatchTrace, None) + outs = f.call_wrapped(*args) + return outs, out_axes() + + +### custom_vmap_p rules + + +def custom_vmap_impl(*args, call, rule, in_tree, out_tree): + del rule, in_tree, out_tree + return core.jaxpr_as_fun(call)(*args) + + +def custom_vmap_batching(args_flat, dims, *, call, rule, in_tree, out_tree): + del call + axis_size, = {x.shape[d] for x, d in zip(args_flat, dims) if d is not None} + args_flat = map(maybe_bdim_at_front, args_flat, dims) + flat_in_batched = [d is not not_mapped for d in dims] + + args = tree_unflatten(in_tree, args_flat) + in_batched = tree_unflatten(in_tree, flat_in_batched) + out, out_batched = call_rule(rule, axis_size, in_batched, args) + flat_outs, tree1 = tree_flatten(out) + flat_out_batched, tree2 = tree_flatten(out_batched) + check_vmap_rule_trees(rule, out_tree, tree1, tree2) + flat_out_dims = [0 if b else not_mapped for b in flat_out_batched] + return flat_outs, flat_out_dims + + +def custom_vmap_abstract_eval(*in_avals, call, **_): + return call.out_avals + + +def custom_vmap_jvp(primals, tangents, *, call, rule, in_tree, out_tree): + def jvp_of_rule_rule(axis_size, in_batched, primals, tangents): + in_batched_ps, in_batched_ts = in_batched + + mutually_batched = tree_map(operator.and_, in_batched_ps, in_batched_ts) + extra_batched_ps = tree_map(lambda pb, tb: 0 if pb and not tb else None, + in_batched_ps, in_batched_ts) + extra_batched_ts = tree_map(lambda pb, tb: 0 if tb and not pb else None, + in_batched_ps, in_batched_ts) + + out_mutually_batched = lu.Store() + flat_ps_ts, tree_ps_ts = tree_flatten((primals, tangents)) + flat_extra_batched_ps_ts, tree_ps_ts2 = tree_flatten( + (extra_batched_ps, extra_batched_ts), + is_leaf=lambda x: x is None) + + # TODO(frostig): assert these also equal: + # treedef_tuple((in_tree, in_tree)) + # once https://github.com/google/jax/issues/9066 is fixed + assert tree_ps_ts == tree_ps_ts2 + del tree_ps_ts2 + + def to_jvp(*primals): + out, out_batched = call_rule(rule, axis_size, mutually_batched, primals) + check_vmap_rule_trees( + rule, out_tree, tree_structure(out), tree_structure(out_batched)) + out_mutually_batched.store(out_batched) + return out + + def to_vmap_over_extra_batched_dims(primals, tangents): + return api.jvp(to_jvp, primals, tangents) + + to_vmap_over_extra_batched_dims_flat, out_tree2 = flatten_fun_nokwargs( + lu.wrap_init(to_vmap_over_extra_batched_dims), + tree_ps_ts) + + flat_out_ps_ts, flat_out_axes = vmap_unrestricted( + to_vmap_over_extra_batched_dims_flat, *flat_ps_ts, + in_axes=flat_extra_batched_ps_ts, + axis_name=core.no_axis_name, axis_size=axis_size) + + n, ragged = divmod(len(flat_out_ps_ts), 2) + assert not ragged + flat_out_ps, flat_out_ts = flat_out_ps_ts[:n], flat_out_ps_ts[n:] + flat_out_axes_p, flat_out_axes_t = flat_out_axes[:n], flat_out_axes[n:] + flat_out_ps = map(maybe_bdim_at_front, flat_out_ps, flat_out_axes_p) + flat_out_extra_batched_ps = [d is not not_mapped for d in flat_out_axes_p] + flat_out_ts = map(maybe_bdim_at_front, flat_out_ts, flat_out_axes_t) + flat_out_extra_batched_ts = [d is not not_mapped for d in flat_out_axes_t] + + out_ps, out_ts = tree_unflatten( + out_tree2(), [*flat_out_ps, *flat_out_ts]) + out_extra_batched_ps, out_extra_batched_ts = tree_unflatten( + out_tree2(), [*flat_out_extra_batched_ps, *flat_out_extra_batched_ts]) + + out_batched_ps = tree_map( + operator.or_, out_mutually_batched.val, out_extra_batched_ps) + out_batched_ts = tree_map( + operator.or_, out_mutually_batched.val, out_extra_batched_ts) + + return (out_ps, out_ts), (out_batched_ps, out_batched_ts) + + tangents = map(ad.instantiate_zeros, tangents) + jvp_call, _ = ad.jvp_jaxpr(call, [True] * len(primals), True) + jvp_in_tree = treedef_tuple((in_tree, in_tree)) + jvp_out_tree = treedef_tuple((out_tree, out_tree)) + outs = custom_vmap_p.bind( + *primals, *tangents, + call=jvp_call, rule=jvp_of_rule_rule, + in_tree=jvp_in_tree, out_tree=jvp_out_tree) + assert len(outs) % 2 == 0, len(outs) + out_primals, out_tangents = util.split_list(outs, [len(outs) // 2]) + return out_primals, out_tangents + + +custom_vmap_p = core.Primitive('custom_vmap_call') +custom_vmap_p.multiple_results = True +custom_vmap_p.def_impl(custom_vmap_impl) +custom_vmap_p.def_abstract_eval(custom_vmap_abstract_eval) +batching.primitive_batchers[custom_vmap_p] = custom_vmap_batching +ad.primitive_jvps[custom_vmap_p] = custom_vmap_jvp +xla.register_initial_style_primitive(custom_vmap_p) +mlir.register_lowering(custom_vmap_p, mlir.lower_fun( + custom_vmap_impl, multiple_results=True)) + + +# -- custom vmap applications + + +def tree_split(mask, tree): + lhs = tree_map(lambda l, x: x if l else None, mask, tree) + rhs = tree_map(lambda l, x: None if l else x, mask, tree) + return lhs, rhs + +def tree_merge(mask, lhs_tree, rhs_tree): + return tree_map(lambda l, x_l, x_r: x_l if l else x_r, + mask, lhs_tree, rhs_tree) + +def sequential_vmap(f): + f = custom_vmap(f) + + @f.def_vmap + def rule(axis_size, in_batched, *args): + del axis_size + + def to_map(mapped_args): + args = tree_merge(in_batched, mapped_args, bcast_args) + return f(*args) + + mapped_args, bcast_args = tree_split(in_batched, list(args)) + out = lax.map(to_map, mapped_args) + out_batched = tree_map(lambda _: True, out) + return out, out_batched + + return f diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_derivatives.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_derivatives.py new file mode 100644 index 0000000000000000000000000000000000000000..459ed1a2d3546ad973662c431ba67d7a7b5c3b27 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_derivatives.py @@ -0,0 +1,1412 @@ +# Copyright 2020 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +from functools import update_wrapper, reduce, partial +import inspect +from typing import Any, Callable, Generic, TypeVar + +from jax._src import config +from jax._src import core +from jax._src import custom_api_util +from jax._src.custom_transpose import custom_transpose +from jax._src import dtypes +from jax._src import effects +from jax._src import linear_util as lu +from jax._src import traceback_util +from jax._src.ad_util import ( + stop_gradient_p, SymbolicZero, Zero, zeros_like_aval) +from jax._src.api_util import argnums_partial, flatten_fun_nokwargs +from jax._src.core import raise_to_shaped +from jax._src.errors import UnexpectedTracerError +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import partial_eval as pe +from jax._src.interpreters import xla +from jax._src.interpreters.batching import not_mapped +from jax._src.lax import lax +from jax._src.tree_util import ( + tree_flatten, tree_unflatten, tree_map, treedef_is_leaf, treedef_tuple, + register_pytree_node_class, tree_leaves, tree_flatten_with_path, keystr, + treedef_children) +from jax._src.util import (cache, safe_zip, safe_map, split_list, Unhashable, + unzip2) + + +traceback_util.register_exclusion(__file__) + +map = safe_map +zip = safe_zip + + +### util + +def _resolve_kwargs(fun, args, kwargs): + if isinstance(fun, partial): + # functools.partial should have an opaque signature. + fun = lambda *args, **kwargs: None + ba = inspect.signature(fun).bind(*args, **kwargs) + ba.apply_defaults() + if ba.kwargs: + raise TypeError("keyword arguments could not be resolved to positions") + else: + return ba.args + +def _initial_style_jaxpr(fun, in_avals): + jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(fun, in_avals) + return jaxpr, consts + +def _close_jaxpr(jaxpr): + return pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr)) + +def _sum_tangents(_, x, *xs): + return reduce(ad.add_tangents, xs, x) + +def _zeros_like_pytree(x): + return tree_map(Zero.from_value, x) + +@partial(partial, tree_map) +def _stop_gradient(x): + if isinstance(x, core.Tracer): + return stop_gradient_p.bind(x) + else: + return x + +# like the api_util.py function, but also grabs output avals for error checking +@lu.transformation_with_aux +def _flatten_fun_nokwargs(in_tree, *args_flat): + py_args = tree_unflatten(in_tree, args_flat) + ans = yield py_args, {} + ans_flat, ans_tree = tree_flatten(ans) + ans_avals = [core.raise_to_shaped(core.get_aval(x)) for x in ans_flat] + yield ans_flat, (ans_tree, ans_avals) + + +### JVPs +ReturnValue = TypeVar('ReturnValue') + +@custom_api_util.register_custom_decorator_type +class custom_jvp(Generic[ReturnValue]): + """Set up a JAX-transformable function for a custom JVP rule definition. + + This class is meant to be used as a function decorator. Instances are + callables that behave similarly to the underlying function to which the + decorator was applied, except when a differentiation transformation (like + :py:func:`jax.jvp` or :py:func:`jax.grad`) is applied, in which case a custom + user-supplied JVP rule function is used instead of tracing into and + performing automatic differentiation of the underlying function's + implementation. + + There are two instance methods available for defining the custom JVP rule: + :py:func:`~jax.custom_jvp.defjvp` for defining a *single* custom JVP rule for + all the function's inputs, and for convenience + :py:func:`~jax.custom_jvp.defjvps`, which wraps + :py:func:`~jax.custom_jvp.defjvp`, and allows you to provide separate + definitions for the partial derivatives of the function w.r.t. each of its + arguments. + + For example:: + + @jax.custom_jvp + def f(x, y): + return jnp.sin(x) * y + + @f.defjvp + def f_jvp(primals, tangents): + x, y = primals + x_dot, y_dot = tangents + primal_out = f(x, y) + tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot + return primal_out, tangent_out + + For a more detailed introduction, see the tutorial_. + + .. _tutorial: https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html + """ + fun: Callable[..., ReturnValue] + nondiff_argnums: tuple[int, ...] + jvp: Callable[..., tuple[ReturnValue, ReturnValue]] | None = None + symbolic_zeros: bool = False + + def __init__(self, + fun: Callable[..., ReturnValue], + nondiff_argnums: tuple[int, ...] = (), + ): + update_wrapper(self, fun) + self.fun = fun + self.nondiff_argnums = nondiff_argnums + + __getattr__ = custom_api_util.forward_attr + + def defjvp(self, + jvp: Callable[..., tuple[ReturnValue, ReturnValue]], + symbolic_zeros: bool = False, + ) -> Callable[..., tuple[ReturnValue, ReturnValue]]: + """Define a custom JVP rule for the function represented by this instance. + + Args: + jvp: a Python callable representing the custom JVP rule. When there are no + ``nondiff_argnums``, the ``jvp`` function should accept two arguments, + where the first is a tuple of primal inputs and the second is a tuple of + tangent inputs. The lengths of both tuples are equal to the number of + parameters of the ``custom_jvp`` function. The ``jvp`` function should + produce as output a pair where the first element is the primal output + and the second element is the tangent output. Elements of the input and + output tuples may be arrays or any nested tuples/lists/dicts thereof. + symbolic_zeros: boolean, indicating whether the rule should be passed + objects representing static symbolic zeros in its tangent argument in + correspondence with unperturbed values; otherwise, only standard JAX + types (e.g. array-likes) are passed. Setting this option to ``True`` + allows a JVP rule to detect whether certain inputs are not involved in + differentiation, but at the cost of needing special handling for these + objects (which e.g. can't be passed into jax.numpy functions). Default + ``False``. + + Returns: + None. + + Example:: + + @jax.custom_jvp + def f(x, y): + return jnp.sin(x) * y + + @f.defjvp + def f_jvp(primals, tangents): + x, y = primals + x_dot, y_dot = tangents + primal_out = f(x, y) + tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot + return primal_out, tangent_out + """ + self.jvp = jvp + self.symbolic_zeros = symbolic_zeros + return jvp + + def defjvps(self, *jvps: Callable[..., ReturnValue] | None): + """Convenience wrapper for defining JVPs for each argument separately. + + This convenience wrapper cannot be used together with ``nondiff_argnums``. + + Args: + *jvps: a sequence of functions, one for each positional argument of the + ``custom_jvp`` function. Each function takes as arguments the tangent + value for the corresponding primal input, the primal output, and the + primal inputs. See the example below. + + Returns: + None. + + Example:: + + @jax.custom_jvp + def f(x, y): + return jnp.sin(x) * y + + f.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y, + lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot) + """ + if self.nondiff_argnums: + raise TypeError("Can't use ``defjvps`` with ``nondiff_argnums``.") + + def jvp(primals, tangents): + primal_out = self(*primals) + zeros = _zeros_like_pytree(primal_out) + all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros + for t, jvp in zip(tangents, jvps)] + tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out) + return primal_out, tangent_out + + self.defjvp(jvp) + + @traceback_util.api_boundary + def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue: # pytype: disable=invalid-annotation + primal_name = getattr(self.fun, '__name__', str(self.fun)) + if not self.jvp: + msg = f"No JVP defined for custom_jvp function {primal_name} using defjvp." + raise AttributeError(msg) + jvp_name = getattr(self.jvp, '__name__', str(self.jvp)) + args = _resolve_kwargs(self.fun, args, kwargs) + if self.nondiff_argnums: + nondiff_argnums = set(self.nondiff_argnums) + args = tuple(_stop_gradient(x) if i in nondiff_argnums else x + for i, x in enumerate(args)) + diff_argnums = [i for i in range(len(args)) if i not in nondiff_argnums] + f_, dyn_args = argnums_partial(lu.wrap_init(self.fun), diff_argnums, args, + require_static_args_hashable=False) + static_args = [args[i] for i in self.nondiff_argnums] + jvp = _add_args(lu.wrap_init(self.jvp), static_args) + else: + f_, dyn_args = lu.wrap_init(self.fun), args + jvp = lu.wrap_init(self.jvp) + args_flat, in_tree = tree_flatten(dyn_args) + flat_fun, out_type1 = _flatten_fun_nokwargs(f_, in_tree) + flat_jvp, out_type2 = _flatten_jvp(jvp, primal_name, jvp_name, in_tree, + out_type1) + out_flat = custom_jvp_call_p.bind(flat_fun, flat_jvp, *args_flat, + symbolic_zeros=self.symbolic_zeros) + _, (out_tree, _) = lu.merge_linear_aux(out_type1, out_type2) + return tree_unflatten(out_tree, out_flat) + +def _add_args(f, extra_args): + return _add_args_(f, tuple(Unhashable(arg) for arg in extra_args)) + +@lu.transformation +def _add_args_(extra_args, *args, **kwargs): + extra_args = tuple(arg.val for arg in extra_args) + all_args = (extra_args + args) + yield (yield all_args, kwargs) + +@partial(lu.transformation_with_aux, use_eq_store=True) +def _flatten_jvp(primal_name, jvp_name, in_tree, maybe_out_type, *args): + primals_in, tangents_in = split_list(args, [len(args) // 2]) + py_primals = tree_unflatten(in_tree, primals_in) + py_tangents = tree_unflatten(in_tree, tangents_in) + pair_out = yield (py_primals, py_tangents), {} + if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2: + msg = (f"Custom JVP rule {jvp_name} for function {primal_name} " + "must produce a pair (list or tuple of length two) representing " + f"primal and tangent outputs, but got {pair_out}.") + raise TypeError(msg) + py_primals_out, py_tangents_out = pair_out + primals_out, out_tree = tree_flatten(py_primals_out) + tangents_out, out_tree2 = tree_flatten(py_tangents_out) + primal_avals = [core.raise_to_shaped(core.get_aval(x)) for x in primals_out] + if out_tree != out_tree2: + msg = (f"Custom JVP rule {jvp_name} for function {primal_name} must " + "produce primal and tangent outputs with equal container (pytree) " + f"structures, but got {out_tree} and {out_tree2} respectively.") + raise TypeError(msg) + # If the primal function already ran, check out_tree agreement. + try: out_type_ = maybe_out_type() + except lu.StoreException: out_type_ = None + if out_type_ is not None: + out_tree_, primal_avals_ = out_type_ + ty_tree = tree_unflatten(out_tree , [a.str_short() for a in primal_avals]) + ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_]) + if out_tree_ != out_tree: + m = (f"Custom JVP rule {jvp_name} for function {primal_name} must " + "produce a pair (list or tuple of length two) " + "where the first element represents the primal output " + "(equal in value to the output of the custom_jvp-decorated function " + f"{primal_name}, " + "and in particular of the same container/pytree structure), but " + "instead the JVP rule output's first element had container/pytree " + "structure:\n" + f""" {str(ty_tree ).replace("'", "")}\n""" + f"while the custom_jvp-decorated function {primal_name} had output " + "container/pytree structure:\n" + f""" {str(ty_tree_).replace("'", "")}.""") + raise TypeError(m) + if not all(map(core.typematch, primal_avals, primal_avals_)): + m = (f"Custom JVP rule {jvp_name} for function {primal_name} must " + "produce a pair (list or tuple of length two) " + "where the first element represents the primal output " + "(equal in value to the output of the custom_jvp-decorated function " + f"{primal_name}, " + "and in particular with leaves of the same shape/dtype), but " + "instead the JVP rule output's first element had shapes/dtypes of:\n" + f""" {str(ty_tree ).replace("'", "")}\n""" + f"while the custom_jvp-decorated function {primal_name} had output " + "shapes/dtypes of:\n" + f""" {str(ty_tree_).replace("'", "")}""") + raise TypeError(m) + # TODO(mattjj): compare primals' tangent types to tangent objects' types + primal_avals_out = [ + raise_to_shaped(core.get_aval(x), weak_type=False).strip_named_shape() + for x in primals_out] + tangent_avals_out = [ + raise_to_shaped(core.get_aval(t), weak_type=False).strip_named_shape() + if type(t) is not SymbolicZero else t.aval.strip_weak_type() + for t in tangents_out] + if primal_avals_out != tangent_avals_out: + if len(primal_avals_out) == 1: + (av1,), (av2,) = primal_avals_out, tangent_avals_out + msg = ("Custom JVP rule must produce primal and tangent outputs with " + "equal shapes and dtypes, but got {} and {} respectively.") + raise TypeError(msg.format(av1.str_short(), av2.str_short())) + else: + msg = ("Custom JVP rule must produce primal and tangent outputs with " + "equal shapes and dtypes, but got:\n{}") + disagreements = ( + f" primal {av1.str_short()} for tangent {av2.str_short()}" + for av1, av2 in zip(primal_avals_out, tangent_avals_out) if av1 != av2) + raise TypeError(msg.format('\n'.join(disagreements))) + yield primals_out + tangents_out, (out_tree, primal_avals) + +class CustomJVPCallPrimitive(core.Primitive): + multiple_results = True + + def bind(self, fun, jvp, *args, symbolic_zeros): + args = map(core.full_lower, args) + top_trace = core.find_top_trace(args) + fun, env_trace_todo1 = process_env_traces( + fun, self, top_trace and top_trace.level, False) + jvp, env_trace_todo2 = process_env_traces( + jvp, self, top_trace and top_trace.level, True) + tracers = map(top_trace.full_raise, args) + outs = top_trace.process_custom_jvp_call(self, fun, jvp, tracers, + symbolic_zeros=symbolic_zeros) + _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2) + return core.apply_todos(env_trace_todo, map(core.full_lower, outs)) + + def impl(self, fun, _, *args): + with core.new_sublevel(): + return fun.call_wrapped(*args) + + def post_process(self, trace, out_tracers, jvp_was_run: bool): + return trace.post_process_custom_jvp_call(out_tracers, jvp_was_run) + + def get_bind_params(self, params): + new_params = dict(params) + call_jaxpr = new_params.pop('call_jaxpr') + num_consts = new_params.pop('num_consts') + jvp_jaxpr_thunk = new_params.pop('jvp_jaxpr_thunk') + fun = lu.wrap_init(core.jaxpr_as_fun(call_jaxpr)) + jvp = lift_jvp(num_consts, jvp_jaxpr_thunk) + return [fun, jvp], new_params + +def lift_jvp(num_consts: int, jvp_jaxpr_thunk: Callable) -> lu.WrappedFun: + @lu.wrap_init + def jvp(*xs): + n, ragged = divmod(len(xs), 2) + assert not ragged + primals, tangents = xs[num_consts:n], xs[n+num_consts:] + zeros = [type(t) is SymbolicZero for t in tangents] + jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_thunk(*zeros) + nonzero_tangents = [t for t in tangents if type(t) is not SymbolicZero] + out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents) + out_primals, nz_out_tangents = split_list(out, [len(out_zeros)]) + nz_out_tangents_ = iter(nz_out_tangents) + out_tangents = [SymbolicZero(core.get_aval(p).at_least_vspace()) + if z else next(nz_out_tangents_) + for p, z in zip(out_primals, out_zeros)] + assert next(nz_out_tangents_, None) is None + return [*out_primals, *out_tangents] + return jvp + +@partial(lu.transformation_with_aux, use_eq_store=True) +def process_env_traces(primitive, level: int, jvp_was_run: bool, *args): + outs = yield args, {} + todo = [] + while True: + tracers = [x for x in outs if isinstance(x, core.Tracer) + and (level is None or x._trace.level > level)] + if tracers: + ans = max(tracers, key=lambda x: x._trace.level) + else: + break + trace = ans._trace.main.with_cur_sublevel() + outs = map(trace.full_raise, outs) + outs, cur_todo = primitive.post_process(trace, outs, jvp_was_run) + todo.append(cur_todo) + yield outs, tuple(todo) # Ensure the aux output is immutable + + +effects.custom_derivatives_allowed_effects.add_type(lax.InOutFeedEffect) + +custom_jvp_call_p = CustomJVPCallPrimitive('custom_jvp_call') + +def _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_thunk, + num_consts, symbolic_zeros): + # TODO(mattjj): could do more checking here... + del in_avals, jvp_jaxpr_thunk, num_consts + disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(call_jaxpr.effects) + if disallowed_effects: + raise NotImplementedError( + f'Effects not supported in `custom_jvp`: {disallowed_effects}') + return call_jaxpr.out_avals, call_jaxpr.effects +core.custom_typechecks[custom_jvp_call_p] = _custom_jvp_call_typecheck + +def _custom_jvp_call_mlir_translation(ctx, *args, call_jaxpr, jvp_jaxpr_thunk, + num_consts, symbolic_zeros): + del jvp_jaxpr_thunk, num_consts, symbolic_zeros + args_ = map(mlir.wrap_singleton_ir_values, args) + consts = mlir._ir_consts(call_jaxpr.consts) + out, tokens = mlir.jaxpr_subcomp(ctx.module_context, call_jaxpr.jaxpr, + ctx.name_stack, ctx.tokens_in, consts, + *args_, dim_var_values=ctx.dim_var_values) + ctx.set_tokens_out(tokens) + return out +mlir.register_lowering(custom_jvp_call_p, _custom_jvp_call_mlir_translation) + +# If a (multi)linear function is defined with a custom jvp, then +# custom_jvp_call_ can appear in jaxprs to be transposed. Since it's already +# been linearized, we can drop the jvp rule. +def _custom_jvp_call_transpose(params, jaxpr, args, ct, _): + del params + return ad.backward_pass(jaxpr.jaxpr, None, jaxpr.consts, args, ct) +ad.primitive_transposes[custom_jvp_call_p] = _custom_jvp_call_transpose + + +### VJPs + +@custom_api_util.register_custom_decorator_type +class custom_vjp(Generic[ReturnValue]): + """Set up a JAX-transformable function for a custom VJP rule definition. + + This class is meant to be used as a function decorator. Instances are + callables that behave similarly to the underlying function to which the + decorator was applied, except when a reverse-mode differentiation + transformation (like :py:func:`jax.grad`) is applied, in which case a custom + user-supplied VJP rule function is used instead of tracing into and performing + automatic differentiation of the underlying function's implementation. There + is a single instance method, :py:func:`~jax.custom_vjp.defvjp`, which may be + used to define the custom VJP rule. + + This decorator precludes the use of forward-mode automatic differentiation. + + For example:: + + @jax.custom_vjp + def f(x, y): + return jnp.sin(x) * y + + def f_fwd(x, y): + return f(x, y), (jnp.cos(x), jnp.sin(x), y) + + def f_bwd(res, g): + cos_x, sin_x, y = res + return (cos_x * g * y, sin_x * g) + + f.defvjp(f_fwd, f_bwd) + + For a more detailed introduction, see the tutorial_. + + .. _tutorial: https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html + """ + + def __init__(self, + fun: Callable[..., ReturnValue], + nondiff_argnums: tuple[int, ...] = ()): + update_wrapper(self, fun) + self.fun = fun + self.nondiff_argnums = nondiff_argnums + self.fwd: Callable[..., tuple[ReturnValue, Any]] | None = None + self.bwd: Callable[..., tuple[Any, ...]] | None = None + self.symbolic_zeros = False + + __getattr__ = custom_api_util.forward_attr + + def defvjp(self, + fwd: Callable[..., tuple[ReturnValue, Any]], + bwd: Callable[..., tuple[Any, ...]], + symbolic_zeros: bool = False, + ) -> None: + """Define a custom VJP rule for the function represented by this instance. + + Args: + fwd: a Python callable representing the forward pass of the custom VJP + rule. When there are no ``nondiff_argnums``, the ``fwd`` function has + the same input signature as the underlying primal function. It should + return as output a pair, where the first element represents the primal + output and the second element represents any "residual" values to store + from the forward pass for use on the backward pass by the function + ``bwd``. Input arguments and elements of the output pair may be arrays + or nested tuples/lists/dicts thereof. + bwd: a Python callable representing the backward pass of the custom VJP + rule. When there are no ``nondiff_argnums``, the ``bwd`` function takes + two arguments, where the first is the "residual" values produced on the + forward pass by ``fwd``, and the second is the output cotangent with the + same structure as the primal function output. The output of ``bwd`` must + be a tuple of length equal to the number of arguments of the primal + function, and the tuple elements may be arrays or nested + tuples/lists/dicts thereof so as to match the structure of the primal + input arguments. + symbolic_zeros: boolean, determining whether to indicate symbolic zeros + to the ``fwd`` and ``bwd`` rules. Enabling this option allows custom + derivative rules to detect when certain inputs, and when certain + output cotangents, are not involved in differentiation. If ``True``: + + * ``fwd`` must accept, in place of each leaf value ``x`` in + the pytree comprising an argument to the original function, + an object (of type + ``jax.custom_derivatives.CustomVJPPrimal``) with two + attributes instead: ``value`` and ``perturbed``. The + ``value`` field is the original primal argument, and + ``perturbed`` is a boolean. The ``perturbed`` bit indicates + whether the argument is involved in differentiation (i.e., + if it is ``False``, then the corresponding Jacobian "column" + is zero). + + * ``bwd`` will be passed objects representing static symbolic zeros in + its cotangent argument in correspondence with unperturbed values; + otherwise, only standard JAX types (e.g. array-likes) are passed. + + Setting this option to ``True`` allows these rules to detect whether + certain inputs and outputs are not involved in differentiation, but at + the cost of special handling. For instance: + + * The signature of ``fwd`` changes, and the objects it is passed cannot + be output from the rule directly. + + * The ``bwd`` rule is passed objects that are not entirely array-like, + and that cannot be passed to most ``jax.numpy`` functions. + + * Any custom pytree nodes involved in the primal function's arguments + must accept, in their unflattening functions, the two-field record + objects that are given as input leaves to the ``fwd`` rule. + + Default ``False``. + + Returns: + None. + + Example:: + + @jax.custom_vjp + def f(x, y): + return jnp.sin(x) * y + + def f_fwd(x, y): + return f(x, y), (jnp.cos(x), jnp.sin(x), y) + + def f_bwd(res, g): + cos_x, sin_x, y = res + return (cos_x * g * y, sin_x * g) + + f.defvjp(f_fwd, f_bwd) + """ + self.fwd = fwd + self.bwd = bwd + self.symbolic_zeros = symbolic_zeros + + @traceback_util.api_boundary + def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue: # pytype: disable=invalid-annotation + primal_name = getattr(self.fun, '__name__', str(self.fun)) + if not self.fwd or not self.bwd: + msg = f"No VJP defined for custom_vjp function {primal_name} using defvjp." + raise AttributeError(msg) + fwd_name = getattr(self.fwd, '__name__', str(self.fwd)) + args = _resolve_kwargs(self.fun, args, kwargs) + if config.enable_custom_vjp_by_custom_transpose.value: + if self.nondiff_argnums: + raise NotImplementedError( + 'nondiff_argnums not implemented for new custom_vjp') + return custom_vjp_by_custom_transpose(self.fun, self.fwd, self.bwd)(*args) + else: + if self.nondiff_argnums: + for i in self.nondiff_argnums: _check_for_tracers(args[i]) + nondiff_argnums = set(self.nondiff_argnums) + dyn_argnums = [i for i in range(len(args)) if i not in nondiff_argnums] + f_, dyn_args = argnums_partial(lu.wrap_init(self.fun), dyn_argnums, + args, require_static_args_hashable=False) + static_args = [args[i] for i in self.nondiff_argnums] + fwd, _ = argnums_partial(lu.wrap_init(self.fwd), dyn_argnums, args, + require_static_args_hashable=False) + bwd = _add_args(lu.wrap_init(self.bwd), static_args) + else: + f_, dyn_args = lu.wrap_init(self.fun), args + fwd, bwd = lu.wrap_init(self.fwd), lu.wrap_init(self.bwd) + args_flat, in_tree = tree_flatten(dyn_args) + in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat] + flat_fun, out_type = _flatten_fun_nokwargs(f_, in_tree) + flat_fwd, out_trees = _flatten_fwd(fwd, self.symbolic_zeros, primal_name, + fwd_name, in_tree, out_type) + flat_bwd = _flatten_bwd(bwd, in_tree, in_avals, out_trees).call_wrapped + out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd, + *args_flat, out_trees=out_trees, + symbolic_zeros=self.symbolic_zeros) + _, (out_tree, _) = lu.merge_linear_aux(out_type, out_trees) + return tree_unflatten(out_tree, out_flat) + +@dataclasses.dataclass +class CustomVJPPrimal: + """Primal to a ``custom_vjp``'s forward rule when ``symbolic_zeros`` is set""" + value: Any + perturbed: bool + +def custom_vjp_primal_tree_values(tree): + """Strips away perturbation information from forward rule arguments. + + This is a helper function for user with the ``symbolic_zeros`` option to + the ``defvjp`` method of a ``custom_vjp``-decorated function. + + In ``symbolic_zeros`` mode, the custom forward rule receives arguments + whose pytree leaves are records with a ``value`` attribute that carries + the primal argument. This is a way to convert such argument trees back to + their original form, replacing each such record with its carried value at + each leaf. + """ + def value(leaf): + if type(leaf) is not CustomVJPPrimal: + raise TypeError(f"unexpected leaf type {type(leaf)}") + return leaf.value + return tree_map(value, tree) + +def _check_for_tracers(x): + for leaf in tree_leaves(x): + if isinstance(leaf, core.Tracer): + msg = ("Found a JAX Tracer object passed as an argument to a custom_vjp " + "function in a position indicated by nondiff_argnums as " + "non-differentiable. Tracers cannot be passed as non-differentiable " + "arguments to custom_vjp functions; instead, nondiff_argnums should " + "only be used for arguments that can't be or contain JAX tracers, " + "e.g. function-valued arguments. In particular, array-valued " + "arguments should typically not be indicated as nondiff_argnums.") + raise UnexpectedTracerError(msg) + +@partial(lu.transformation_with_aux, use_eq_store=True) +def _flatten_fwd(symbolic_zeros, primal_name, fwd_name, in_tree, maybe_out_type, + *args): + if symbolic_zeros: + args = [CustomVJPPrimal(x, z) for x, z in zip(args[::2], args[1::2])] + else: + args = args[::2] + py_args = tree_unflatten(in_tree, args) + pair_out = yield py_args, {} + if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2: + msg = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} " + "must produce a pair (list or tuple of length two) where the first " + "element represents the primal output (equal to those of the " + f"custom_vjp-decorated function {primal_name}) and the " + "second element represents residuals (i.e. values stored from the " + "forward pass for use on the backward pass), but " + f"instead of a pair the fwd rule {fwd_name} produced {pair_out}.") + raise TypeError(msg) + py_primals_out, res = pair_out + primals_out, out_tree = tree_flatten(py_primals_out) + res, res_tree = tree_flatten(res) + primal_avals = [core.raise_to_shaped(core.get_aval(x)) for x in primals_out] + # If the primal function already ran, check out_tree agreement. + try: out_type_ = maybe_out_type() + except lu.StoreException: out_type_ = None + if out_type_ is not None: + out_tree_, primal_avals_ = out_type_ + ty_tree = tree_unflatten(out_tree , [a.str_short() for a in primal_avals]) + ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_]) + if out_tree_ != out_tree: + m = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} " + "must produce a pair (list or tuple of length two) where the first " + "element represents the primal output " + "(equal to the output of the custom_vjp-decorated function " + f"{primal_name}) and the " + "second element represents residuals (i.e. values stored from the " + "forward pass for use on the backward pass), but " + "instead the fwd rule output's first element had container/pytree " + "structure:\n" + f""" {str(ty_tree ).replace("'", "")}\n""" + f"while the custom_vjp-decorated function {primal_name} had output " + "container/pytree structure:\n" + f""" {str(ty_tree_).replace("'", "")}.""") + raise TypeError(m) + if not all(map(core.typematch, primal_avals, primal_avals_)): + m = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} must " + "produce a pair (list or tuple of length two) " + "where the first element represents the primal output " + "(equal to the output of the custom_vjp-decorated function " + f"{primal_name}) and the second element represents residuals " + "(i.e. values stored from the forward pass for use on the " + "backward pass), but " + "instead the fwd rule output's first element had shapes/dtypes of:\n" + f""" {str(ty_tree ).replace("'", "")}\n""" + f"while the custom_vjp-decorated function {primal_name} had output " + "shapes/dtypes of:\n" + f""" {str(ty_tree_).replace("'", "")}""") + raise TypeError(m) + yield (*res, *primals_out), (out_tree, res_tree) + +@lu.transformation +def _flatten_bwd(in_tree, in_avals, out_trees, *args): + out_tree, res_tree = out_trees() + assert len(args) == res_tree.num_leaves + out_tree.num_leaves + res, cts_out = split_list(args, [res_tree.num_leaves]) + py_res = tree_unflatten(res_tree, res) + py_cts_out = tree_unflatten(out_tree, cts_out) + py_cts_in = yield (py_res, py_cts_out), {} + if isinstance(py_cts_in, list) and len(py_cts_in) == len(treedef_children(in_tree)): + py_cts_in = tuple(py_cts_in) + # For each None in py_cts_in, indicating an argument for which the rule + # produces no cotangent, we replace it with a pytree with the structure of the + # corresponding subtree of in_tree and with leaves of a non-pytree sentinel + # object, to be replaced with Nones in the final returned result. + zero = object() # non-pytree sentinel to replace Nones in py_cts_in + dummy = tree_unflatten(in_tree, [object()] * in_tree.num_leaves) + keypaths, _ = unzip2(tree_flatten_with_path(dummy)[0]) + cts_in_flat = [] + def append(x, d): + num_leaves = len(tree_flatten(d)[0]) + if x is None and d is not None: + cts_in_flat.extend([zero] * num_leaves) + elif x is not None: + cts_in_flat.extend([x] * num_leaves) + return x + try: + if not isinstance(py_cts_in, tuple): + raise ValueError + tree_map(append, py_cts_in, dummy, is_leaf=lambda x: x is None) + except ValueError: + _, in_tree2 = tree_flatten(py_cts_in) + msg = ("Custom VJP bwd rule must produce an output with the same container " + "(pytree) structure as the args tuple of the primal function, " + "and in particular must produce a tuple of length equal to the " + "number of arguments to the primal function, but got bwd output " + "structure {} for primal input structure {}.") + raise TypeError(msg.format(in_tree2, in_tree)) from None + results = [] + for kp, a, ct in zip(keypaths, in_avals, cts_in_flat): + if ct is zero or a != a.at_least_vspace(): + results.append(Zero(a.at_least_vspace())) + elif type(ct) is SymbolicZero: + if not core.typecompat(a.at_least_vspace(), a_ := ct.aval): + msg = ("Custom VJP bwd rule produced a SymbolicZero with a shape/dtype " + "that does not match the corresponding input tangent shape/dtype: " + f"at output{keystr(kp)} the SymbolicZero had shape/dtype " + f"{a_.str_short()} while the " + f"corresponding input had shape/dtype {a.str_short()}. " + "Consider just returning a None here instead of a SymbolicZero " + "object.") + raise ValueError(msg) + results.append(Zero(ct.aval)) + else: + if (not core.typecompat(a.at_least_vspace(), a_ := core.get_aval(ct)) + and not (_temporary_dtype_exception(a, a_) or + _temporary_shape_exception(a, a_))): + msg = ("Custom VJP bwd rule must produce an output with the same " + "shape/dtypes as the args tuple of the primal function, but at " + f"output{keystr(kp)} the bwd rule produced an output of " + f"shape/dtype {raise_to_shaped(a_).str_short()} corresponding " + f"to an input of shape/dtype {a.str_short()}.") + raise ValueError(msg) + results.append(ct) + yield results + +# TODO(mattjj): remove both these exceptions to cotangent compatibility check +def _temporary_dtype_exception(a, a_) -> bool: + if isinstance(a, core.ShapedArray) and isinstance(a_, core.ShapedArray): + return (a.shape == a_.shape and + (dtypes.issubdtype(a_.dtype, dtypes.extended) or + dtypes.issubdtype(a.dtype, dtypes.np.inexact))) + return False + +# TODO(mattjj): remove both these exceptions to cotangent compatibility check +def _temporary_shape_exception(a, a_) -> bool: + return config.custom_vjp_disable_shape_check.value + +class CustomVJPCallPrimitive(core.CallPrimitive): + initial_style: core.Primitive + + def bind(self, fun, fwd, bwd, *args, out_trees, symbolic_zeros): + args = map(core.full_lower, args) + top_trace = core.find_top_trace(args) + fun, env_trace_todo1 = process_env_traces( + fun, self, top_trace and top_trace.level, False) + fwd, env_trace_todo2 = process_env_traces_fwd( + fwd, top_trace and top_trace.level, out_trees) + tracers = map(top_trace.full_raise, args) + bwd_ = lambda *args: bwd(*args) + outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd_, tracers, + out_trees=out_trees, + symbolic_zeros=symbolic_zeros) + fst, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2) + if fst: + return core.apply_todos(env_trace_todo, map(core.full_lower, outs)) + else: + env_trace_todo, bwd_transform = env_trace_todo + bwd = _apply_bwd_transform(bwd_transform, bwd) + return core.apply_todos(env_trace_todo, map(core.full_lower, outs)) + + def impl(self, fun, fwd, bwd, *args, out_trees): + del fwd, bwd, out_trees + with core.new_sublevel(): + return fun.call_wrapped(*args) + + def post_process(self, trace, out_tracers, params): + return trace.post_process_custom_vjp_call(out_tracers, params) +custom_vjp_call_p = CustomVJPCallPrimitive('custom_vjp_call') + +@partial(lu.transformation_with_aux, use_eq_store=True) +def process_env_traces_fwd(level: int, out_trees, *args): + outs = yield args, {} + todo = [] + bwd_transforms = [] + while True: + tracers = [x for x in outs if isinstance(x, core.Tracer) + and (level is None or x._trace.level > level)] + if tracers: + ans = max(tracers, key=lambda x: x._trace.level) + else: + break + trace = ans._trace.main.with_cur_sublevel() + outs = map(trace.full_raise, outs) + outs, cur_todo, bwd_xform = trace.post_process_custom_vjp_call_fwd(outs, out_trees) + todo.append(cur_todo) + bwd_transforms.append(bwd_xform) + yield outs, (tuple(todo), tuple(bwd_transforms)) + + +def _apply_bwd_transform(todos, bwd): + todos_list = list(todos) + while todos_list: + bwd = todos_list.pop()(bwd) + return bwd + +def _custom_vjp_call_jaxpr_impl(*args, fun_jaxpr, **_): + return core.jaxpr_as_fun(fun_jaxpr)(*args) + +def _custom_vjp_call_jaxpr_abstract_eval(*_, fun_jaxpr, **__): + disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(fun_jaxpr.effects) + if disallowed_effects: + raise NotImplementedError( + f'Effects not supported in `custom_vjp`: {disallowed_effects}') + return fun_jaxpr.out_avals, fun_jaxpr.effects + +custom_vjp_call_jaxpr_p = core.AxisPrimitive('custom_vjp_call_jaxpr') +custom_vjp_call_jaxpr_p.multiple_results = True +custom_vjp_call_jaxpr_p.def_impl(_custom_vjp_call_jaxpr_impl) +custom_vjp_call_jaxpr_p.def_effectful_abstract_eval(_custom_vjp_call_jaxpr_abstract_eval) +CustomVJPCallPrimitive.initial_style = custom_vjp_call_jaxpr_p + +mlir.register_lowering(custom_vjp_call_jaxpr_p, mlir.lower_fun( + _custom_vjp_call_jaxpr_impl, multiple_results=True)) + +def _custom_vjp_call_jaxpr_jvp( + primals, tangents, *, fun_jaxpr: core.ClosedJaxpr, + fwd_jaxpr_thunk: Callable[..., tuple[core.Jaxpr, Sequence[Any]]], + num_consts: int, bwd: Callable, out_trees: Callable, symbolic_zeros: bool): + _, args = split_list(primals, [num_consts]) + consts_dot, args_dot = split_list(tangents, [num_consts]) + if any(type(t) is not Zero for t in consts_dot): + raise ad.CustomVJPException() + zeros = [type(t) is not Zero for t in args_dot] + fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk(*zeros) # consts can be tracers! + _, res_tree = out_trees() + res_and_primals_out = core.eval_jaxpr(fwd_jaxpr, fwd_consts, *args) + res, primals_out = split_list(res_and_primals_out, [res_tree.num_leaves]) + avals_out = [raise_to_shaped(core.get_aval(x)) for x in primals_out] + args_dot = map(ad.instantiate_zeros, args_dot) + # Cast float0 to zeros with the primal dtype because custom vjp rules don't + # currently handle float0s + args_dot = map(ad.replace_float0s, args, args_dot) + tangents_out = ad.custom_lin_p.bind( + *res, *args_dot, num_res=res_tree.num_leaves, bwd=bwd, + out_avals=avals_out, symbolic_zeros=symbolic_zeros) + tangents_out = map(lax.tie_p.bind, primals_out, tangents_out) + tangents_out = map(ad.recast_to_float0, primals_out, tangents_out) + return primals_out, tangents_out +ad.primitive_jvps[custom_vjp_call_jaxpr_p] = _custom_vjp_call_jaxpr_jvp + +def _custom_vjp_call_jaxpr_vmap( + spmd_axis_name, axis_size, axis_name, main_type, args, in_dims, *, + fun_jaxpr: core.ClosedJaxpr, + fwd_jaxpr_thunk: Callable[..., tuple[core.Jaxpr, Sequence[Any]]], + num_consts: int, bwd: Callable, out_trees: Callable, symbolic_zeros: bool): + args = [batching.moveaxis(x, d, 0) if d is not not_mapped and d != 0 + else x for x, d in zip(args, in_dims)] + + in_batched = [d is not not_mapped for d in in_dims] + _, args_batched = split_list(in_batched, [num_consts]) + batched_fun_jaxpr, out_batched = batching.batch_jaxpr( + fun_jaxpr, axis_size, in_batched, False, axis_name, spmd_axis_name, + main_type) + out_dims1 = [0 if b else not_mapped for b in out_batched] + out_dims2 = [] + + @pe._memoize + def batched_fwd_jaxpr_thunk(*zeros): + fwd_jaxpr = core.ClosedJaxpr(*fwd_jaxpr_thunk(*zeros)) # consts can be tracers + batched_fwd_jaxpr, out_batched = batching.batch_jaxpr( + fwd_jaxpr, axis_size, args_batched, False, axis_name, spmd_axis_name, + main_type) + out_dims2.append([0 if b else not_mapped for b in out_batched]) + return batched_fwd_jaxpr.jaxpr, batched_fwd_jaxpr.consts + + fwd_args_batched = [0 if b else not_mapped for b in args_batched] + fwd_out_dims = lambda: out_dims2[0] + batched_bwd = batching.batch_custom_vjp_bwd( + bwd, axis_name, axis_size, fwd_out_dims, fwd_args_batched, main_type, + spmd_axis_name) + + batched_outs = custom_vjp_call_jaxpr_p.bind( + *args, fun_jaxpr=batched_fun_jaxpr, + fwd_jaxpr_thunk=batched_fwd_jaxpr_thunk, bwd=batched_bwd, + num_consts=num_consts, out_trees=out_trees, symbolic_zeros=symbolic_zeros) + out_dims = out_dims2[0] if out_dims2 else out_dims1 + return batched_outs, out_dims +batching.spmd_axis_primitive_batchers[custom_vjp_call_jaxpr_p] = \ + _custom_vjp_call_jaxpr_vmap +batching.axis_primitive_batchers[custom_vjp_call_jaxpr_p] = partial( + _custom_vjp_call_jaxpr_vmap, None) + +xla.register_initial_style_primitive(custom_vjp_call_jaxpr_p) + +batching.primitive_batchers[ad.custom_lin_p] = ad.raise_custom_vjp_error_on_jvp +mlir.register_lowering(ad.custom_lin_p, ad.raise_custom_vjp_error_on_jvp) + + +def custom_gradient(fun): + """Convenience function for defining custom VJP rules (aka custom gradients). + + While the canonical way to define custom VJP rules is via ``jax.custom_vjp``, + the ``custom_gradient`` convenience wrapper follows TensorFlow's + ``tf.custom_gradient`` API. The difference here is that ``custom_gradient`` + can be used as a decorator on one function that returns both the primal value + (representing the output of the mathematical function to be differentiated) + and the VJP (gradient) function. See + https://www.tensorflow.org/api_docs/python/tf/custom_gradient. + + If the mathematical function to be differentiated has Haskell-like signature + ``a -> b``, then the Python callable ``fun`` should have the signature + ``a -> (b, CT b --o CT a)`` where we use ``CT x`` to denote a cotangent type + for ``x`` and the ``--o`` arrow to denote a linear function. See the example + below. That is, ``fun`` should return a pair where the first element + represents the value of the mathematical function to be differentiated and the + second element is a function to be called on the backward pass of reverse-mode + automatic differentiation (i.e. the "custom gradient" function). + + The function returned as the second element of the output of ``fun`` can close + over intermediate values computed when evaluating the function to be + differentiated. That is, use lexical closure to share work between the forward + pass and the backward pass of reverse-mode automatic differentiation. However, + it cannot perform Python control flow which depends on the values of the + closed-over intermediate values or its cotangent arguments; if the function + includes such control flow, an error is raised. + + Args: + fun: a Python callable specifying both the mathematical function to be + differentiated and its reverse-mode differentiation rule. It should return + a pair consisting of an output value and a Python callable that represents + the custom gradient function. + + Returns: + A Python callable that accepts the same arguments as ``fun`` and returns the + output value specified by the first element of ``fun``'s output pair. + + For example: + + >>> @jax.custom_gradient + ... def f(x): + ... return x ** 2, lambda g: (g * x,) + ... + >>> print(f(3.)) + 9.0 + >>> print(jax.grad(f)(3.)) + 3.0 + + An example with a function on two arguments, so that the VJP function must + return a tuple of length two: + + >>> @jax.custom_gradient + ... def f(x, y): + ... return x * y, lambda g: (g * y, g * x) + ... + >>> print(f(3., 4.)) + 12.0 + >>> print(jax.grad(f, argnums=(0, 1))(3., 4.)) + (Array(4., dtype=float32, weak_type=True), Array(3., dtype=float32, weak_type=True)) + """ + @custom_vjp + def wrapped_fun(*args, **kwargs): + ans, _ = fun(*args, **kwargs) + return ans + + def fwd(*args, **kwargs): + ans, rule = fun(*args, **kwargs) + ans_flat, out_tree = tree_flatten((ans,)) + rule, in_tree = flatten_fun_nokwargs(lu.wrap_init(rule), out_tree) + ans_avals = [core.get_aval(x).at_least_vspace() for x in ans_flat] + jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(rule, ans_avals) + return ans, Residuals(jaxpr, in_tree(), out_tree, consts) + + def bwd(res, cts): + jaxpr, in_tree, out_tree, consts = res + cts_flat, out_tree_ = tree_flatten((cts,)) + if out_tree != out_tree_: raise TypeError(f'{out_tree}\n!=\n{out_tree_}') + cts_out = core.eval_jaxpr(jaxpr, consts, *cts_flat) + cts_out = tree_unflatten(in_tree, cts_out) + if treedef_is_leaf(in_tree): + cts_out = (cts_out,) + return cts_out + + wrapped_fun.defvjp(fwd, bwd) + return wrapped_fun + +@register_pytree_node_class +class Residuals: + def __init__(self, jaxpr, in_tree, out_tree, consts): + self.jaxpr = jaxpr + self.in_tree = in_tree + self.out_tree = out_tree + self.consts = consts + def __iter__(self): + return iter((self.jaxpr, self.in_tree, self.out_tree, self.consts)) + def tree_flatten(self): + return self.consts, (self.jaxpr, self.in_tree, self.out_tree) + @classmethod + def tree_unflatten(cls, aux, consts): + jaxpr, in_tree, out_tree = aux + return cls(jaxpr, in_tree, out_tree, consts) + + +def closure_convert(fun: Callable, *example_args) -> tuple[Callable, list[Any]]: + """Closure conversion utility, for use with higher-order custom derivatives. + + To define custom derivatives such as with ``jax.custom_vjp(f)``, the target + function ``f`` must take, as formal arguments, all values involved in + differentiation. If ``f`` is a higher-order function, in that it accepts as an + argument a Python function ``g``, then values stored away in ``g``'s closure + will not be visible to the custom derivative rules, and attempts at AD + involving these values will fail. One way around this is to convert the + closure by extracting these values, and to pass them as explicit formal + arguments across the custom derivative boundary. This utility carries out that + conversion. More precisely, it closure-converts the function ``fun`` + specialized to the types of the arguments given in ``example_args``. + + When we refer here to "values in the closure" of ``fun``, we do not mean the + values that are captured by Python directly when ``fun`` is defined (e.g. the + Python objects in ``fun.__closure__``, if the attribute exists). Rather, we + mean values encountered during the execution of ``fun`` on ``example_args`` + that determine its output. This may include, for instance, arrays captured + transitively in Python closures, i.e. in the Python closure of functions + called by ``fun``, the closures of the functions that they call, and so forth. + + The function ``fun`` must be a pure function. + + Example usage:: + + def minimize(objective_fn, x0): + converted_fn, aux_args = closure_convert(objective_fn, x0) + return _minimize(converted_fn, x0, *aux_args) + + @partial(custom_vjp, nondiff_argnums=(0,)) + def _minimize(objective_fn, x0, *args): + z = objective_fn(x0, *args) + # ... find minimizer x_opt ... + return x_opt + + def fwd(objective_fn, x0, *args): + y = _minimize(objective_fn, x0, *args) + return y, (y, args) + + def rev(objective_fn, res, g): + y, args = res + y_bar = g + # ... custom reverse-mode AD ... + return x0_bar, *args_bars + + _minimize.defvjp(fwd, rev) + + Args: + fun: Python callable to be converted. Must be a pure function. + example_args: Arrays, scalars, or (nested) standard Python + containers (tuples, lists, dicts, namedtuples, i.e., pytrees) + thereof, used to determine the types of the formal arguments to + ``fun``. This type-specialized form of ``fun`` is the function + that will be closure converted. + + Returns: + A pair comprising (i) a Python callable, accepting the same + arguments as ``fun`` followed by arguments corresponding to the + values hoisted from its closure, and (ii) a list of values hoisted + from the closure. + """ + flat_args, in_tree = tree_flatten(example_args) + in_avals = tuple(map(abstractify, flat_args)) + if config.check_tracer_leaks.value: + return _closure_convert_for_avals.__wrapped__(fun, in_tree, in_avals) + else: + return _closure_convert_for_avals(fun, in_tree, in_avals) + +def _maybe_perturbed(x: Any) -> bool: + # False if x can't represent an AD-perturbed value (i.e. a value + # with a nontrivial tangent attached), up to heuristics, and True otherwise. + # See https://github.com/google/jax/issues/6415 for motivation. + x = core.full_lower(x) + if not isinstance(x, core.Tracer): + # If x is not a Tracer, it can't be perturbed. + return False + elif isinstance(x, pe.DynamicJaxprTracer): + # If x is a DynamicJaxprTracer then we're staging out; differentiation could + # happen later, but some types always have trivial tangents. + vspace = x.aval.at_least_vspace() + return not (vspace is core.abstract_token or + getattr(vspace, 'dtype', None) == dtypes.float0) + elif not isinstance(x, ad.JVPTracer): + # If x is not a JVPTracer, recursively check its contents. + return any(_maybe_perturbed(attr) for name, attr in x._contents()) + else: + return True # We can't be sure! + +@cache() +def _closure_convert_for_avals(fun, in_tree, in_avals): + wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree) + jaxpr, out_pvals, consts, () = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals) + out_tree = out_tree() + + (closure_consts, hoisted_consts), merge = partition_list(_maybe_perturbed, consts) + num_consts = len(hoisted_consts) + + def converted_fun(*args_hconsts): + num_args = len(args_hconsts) - num_consts + args, hoisted_consts = split_list(args_hconsts, [num_args]) + consts = merge(closure_consts, hoisted_consts) + all_args, in_tree2 = tree_flatten(tuple(args)) + assert in_tree == in_tree2 + out_flat = core.eval_jaxpr(jaxpr, consts, *all_args) + return tree_unflatten(out_tree, out_flat) + + return converted_fun, hoisted_consts + +def partition_list(choice, lst): + out = [], [] + which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst] + def merge(l1, l2): + i1, i2 = iter(l1), iter(l2) + return [next(i2 if snd else i1) for snd in which] + return out, merge + +def abstractify(x): + return core.raise_to_shaped(core.get_aval(x)) + + +### Custom transposition + +def linear_call(fun: Callable, fun_transpose: Callable, residual_args, + linear_args): + """Call a linear function, with a custom implementation for its transpose. + + The `Haskell-like type signatures`_ of ``fun`` and ``fun_transpose`` are: + + .. code-block:: haskell + + fun :: r -> a -o b + fun_transpose :: r -> b -o a + + where the ``-o`` arrow indicates a linear function, ``r`` is the + residual input type and ``a`` is the linear input type. + + The functions ``fun`` and ``fun_transpose`` are coupled as + transposes of one another. Specifically, the transpose of a + ``linear_call`` primitive is another ``linear_call`` to + ``fun_transpose``, with ``fun`` as its custom transposition. + + For example: + + >>> def f(r, x): + ... return x / r + + >>> def t(r, t): + ... return t / r + + >>> def div_add(x, denom): + ... return x + linear_call(f, t, denom, x) + + >>> def transpose(f, x_example): + ... def transposed(y): + ... x, = jax.linear_transpose(f, x_example)(y) + ... return x + ... return transposed + + >>> div_add(9., 3.) + Array(12., dtype=float32, weak_type=True) + + >>> transpose(partial(div_add, denom=3.), 1.)(18.) # custom + Array(24., dtype=float32, weak_type=True) + + >>> transpose(lambda x: x + x / 3., 1.)(18.) # reference + Array(24., dtype=float32, weak_type=True) + + The above definition of ``f`` illustrates the purpose of a residual + argument: division is linear in one of its inputs (the dividend + ``x``) but not the other (the divisor ``r``). + + As another example: + + >>> def custom_id(x): + ... def f(_, x): return x + ... def t(_, t): return 7. + ... return linear_call(f, t, (), x) + >>> custom_id(1.) + 1.0 + >>> transpose(custom_id, 1.)(1.) + 7.0 + >>> transpose(transpose(custom_id, 1.), 1.)(1.) + 1.0 + >>> transpose(transpose(transpose(custom_id, 1.), 1.), 1.)(1.) + 7.0 + + Args: + fun: a Python callable specifying a linear function. It should + take two arguments: one of "residual" inputs (type ``r``), + i.e. inputs in which the function is not necessarily linear, and + one of "linear" inputs (type ``a``). It should return output + whose components are linear in the linear input (type ``b``). + fun_transpose: a Python callable specifying a structurally linear + function that is the transpose of ``fun`` with respect to its + linear inputs. Its first argument is the same residual inputs + (``r``) as ``fun``. Its second argument is of type + ``b``. Finally, its output is of type ``a`` and each of its + component are linear in its second argument (the ``b`` inputs). + residual_args: Argument in which ``fun`` and ``fun_transpose`` are + not necessarily linear. Not involved in transposition. + linear_args: Argument in which ``fun`` and ``fun_transpose`` are + linear and with respect to which the two are transposes. + + Returns: + The call result, i.e. ``fun(residual_args, linear_args)``. + + .. _Haskell-like type signatures: https://wiki.haskell.org/Type_signature + """ + operands_res, res_tree = tree_flatten(residual_args) + operands_lin, lin_tree = tree_flatten(linear_args) + + f_in_tree = treedef_tuple((res_tree, lin_tree)) + f, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), f_in_tree) + + res_avals = map(abstractify, operands_res) + lin_avals = map(abstractify, operands_lin) + f_jaxpr, f_consts = _initial_style_jaxpr(f, (*res_avals, *lin_avals)) + f_jaxpr = _close_jaxpr(f_jaxpr) + out_avals = map(core.raise_to_shaped, f_jaxpr.out_avals) + + t_in_tree = treedef_tuple((res_tree, out_tree())) + t, t_out_tree = flatten_fun_nokwargs(lu.wrap_init(fun_transpose), t_in_tree) + + t_jaxpr, t_consts = _initial_style_jaxpr(t, (*res_avals, *out_avals)) + t_jaxpr = _close_jaxpr(t_jaxpr) + + if t_out_tree() != lin_tree: + raise TypeError( + 'transpose output pytree structure must match that of linear inputs, ' + f'got output structure {t_out_tree()} ' + f'and input structure {lin_tree}.') + + out = linear_call_p.bind(*f_consts, *t_consts, *operands_res, *operands_lin, + callee=f_jaxpr, + transpose=t_jaxpr, + num_callee_consts=len(f_consts), + num_transpose_consts=len(t_consts), + num_res=len(operands_res)) + + return tree_unflatten(out_tree(), out) + +def _linear_call_impl(*args, callee, transpose, num_callee_consts, + num_transpose_consts, num_res): + del transpose + consts, _, operands_res, operands_lin = split_list( + args, [num_callee_consts, num_transpose_consts, num_res]) + return core.eval_jaxpr(callee.jaxpr, (), *consts, *operands_res, *operands_lin) + +def _linear_call_transpose_rule(cts, *args, callee, transpose, + num_callee_consts, + num_transpose_consts, num_res): + f_consts, t_consts, operands_res, operands_lin = split_list( + args, [num_callee_consts, num_transpose_consts, num_res]) + _, _, cts_avals = split_list( + transpose.in_avals, [num_transpose_consts, num_res]) + + assert all(ad.is_undefined_primal(x) for x in operands_lin) + assert all(not ad.is_undefined_primal(x) for x in operands_res) + + cts = [zeros_like_aval(a) if type(ct) is Zero else ct + for ct, a in zip(cts, cts_avals)] + + cts_out = linear_call_p.bind(*t_consts, *f_consts, *operands_res, *cts, + callee=transpose, + transpose=callee, + num_callee_consts=len(t_consts), + num_transpose_consts=len(f_consts), + num_res=len(operands_res)) + + return [None] * (num_callee_consts + num_transpose_consts + num_res) + cts_out + +def _linear_call_abstract_eval(*args, **kwargs): + return map(core.raise_to_shaped, kwargs['callee'].out_avals) + +linear_call_p = core.Primitive('linear_call') +linear_call_p.multiple_results = True +linear_call_p.def_impl(_linear_call_impl) +linear_call_p.def_abstract_eval(_linear_call_abstract_eval) +ad.primitive_transposes[linear_call_p] = _linear_call_transpose_rule +xla.register_initial_style_primitive(linear_call_p) +mlir.register_lowering(linear_call_p, mlir.lower_fun( + _linear_call_impl, multiple_results=True)) + + +# A stageable primitive that fails when evaluated +unreachable_p: core.Primitive = core.Primitive('unreachable') +unreachable_p.multiple_results = True + +def unreachable_impl(*_, out_avals, exc_type, message): + del out_avals + raise exc_type(message) + +# Evaluation raises an exception +unreachable_p.def_impl(unreachable_impl) + +# Translation raises an exception +# TODO(frostig,mattjj): We have no good way to translate a function +# that errs. Since MLIR lowering over-approximates concrete evaluation, +# we err on MLIR lowering for the time being. +mlir.register_lowering(unreachable_p, unreachable_impl) + +# Abstract evaluation proceeds without issue, to allow for staging +unreachable_p.def_abstract_eval(lambda *_, out_avals, **__: out_avals) + +def unreachable(*args, out_avals=None, exc_type=TypeError, + message='unreachable'): + """Fail when evaluated concretely (but allow for staging). + + This function allows one to assert an impossibility of + evaluation. It can be used to guarantee that evaluation does not + "reach" a certain point in the sense that it does not execute, but + it can nonetheless be staged out by JAX without error. + + Args: + *args: The arbitrary pytree of arguments to the function. + out_avals: Optional specification of the output types of this + function invocation from the point of view of staging. If + ``None``, these are chosen as equal to types of input arguments. + exc_type: Optional constructor for the Python exception raised if + evaluated. + message: Optional string message for the Python exception raised + if evaluated. + + """ + if out_avals is None: + out_avals = tree_map(core.get_aval, args) + + args_flat, in_tree = tree_flatten(args) + out_avals_flat, out_tree = tree_flatten(out_avals) + out = unreachable_p.bind(*args_flat, out_avals=out_avals_flat, + exc_type=exc_type, message=message) + return tree_unflatten(out_tree, out) + + +disallow_jvp = partial( + unreachable, + exc_type=TypeError, + message="can't apply forward-mode autodiff (jvp) to a custom_vjp function.") + + +def custom_vjp_by_custom_transpose(fun, fwd, bwd): + fun = custom_jvp(fun) + + @fun.defjvp + def jvp(primals, tangents): + outs, residuals = fwd(*primals) + tan_out_types = tree_map(lambda o: core.get_aval(o).at_least_vspace(), outs) + tan_fn = custom_transpose(partial(disallow_jvp, out_avals=tan_out_types)) + tan_fn.def_transpose(bwd) + return outs, tan_fn(tan_out_types, residuals, tangents) + + return fun + + +# TODO(mattjj): remove these stubs, which exist to avoid breaking internal users +custom_jvp_call_jaxpr_p = core.Primitive("custom_jvp_call_jaxpr") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_transpose.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_transpose.py new file mode 100644 index 0000000000000000000000000000000000000000..da132e085086be4221768974cb7d59abee8e5936 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/custom_transpose.py @@ -0,0 +1,231 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import functools +from typing import Any, Callable + +from jax._src import ad_util +from jax._src import api_util +from jax._src import core +from jax._src import custom_api_util +from jax._src import linear_util as lu +from jax._src import source_info_util +from jax._src import traceback_util +from jax._src import util +from jax._src.interpreters import ad +from jax._src.interpreters import mlir +from jax._src.interpreters import partial_eval as pe +from jax._src.interpreters import xla +from jax._src.tree_util import (tree_flatten, tree_leaves, tree_map, + tree_structure, treedef_tuple, tree_unflatten) + + +source_info_util.register_exclusion(__file__) +traceback_util.register_exclusion(__file__) + + +map, unsafe_map = util.safe_map, map +zip, unsafe_zip = util.safe_zip, zip + + +### bespoke linear_util and api_util deviations + +class StoreEqual(lu.Store): + """Stores an unchanging value. Checks empty reads and unequal overwrites.""" + def store(self, val): + if self._val is not lu._EMPTY_STORE_VALUE and val != self._val: + raise lu.StoreException( + f"Store assignment mismatch, from {self._val} to {val}") + self._val = val + +@util.curry +def transformation_with_aux( + gen, fun: lu.WrappedFun, *gen_static_args) -> tuple[lu.WrappedFun, Any]: + out_store = StoreEqual() + out_thunk = lambda: out_store.val + return fun.wrap(gen, gen_static_args, out_store), out_thunk + +flatten_fun_nokwargs = transformation_with_aux( + api_util.flatten_fun_nokwargs.args[0]) + + +### api + +@custom_api_util.register_custom_decorator_type +class custom_transpose: + fun: Callable + transpose: Callable | None = None + + def __init__(self, fun: Callable): + functools.update_wrapper(self, fun) + self.fun = fun + + __getattr__ = custom_api_util.forward_attr + + def def_transpose(self, transpose: Callable): + self.transpose = transpose + return transpose + + @traceback_util.api_boundary + def __call__(self, out_types, res_arg, lin_arg): + _, res_tree = tree_flatten(res_arg) + _, lin_tree = tree_flatten(lin_arg) + args_flat, in_tree = tree_flatten((res_arg, lin_arg)) + + # TODO(frostig,mattjj): check that out_trees match + # TODO(frostig,mattjj): could, and should, we avoid flattening + # self.fun at this point? + + flat_fun, out_tree2 = flatten_fun_nokwargs(lu.wrap_init(self.fun), in_tree) + out_types_flat, out_tree = tree_flatten(out_types) + out_flat = custom_transpose_p.bind(flat_fun, *args_flat, + transpose=self.transpose, + out_types=out_types_flat, + lin_tree=lin_tree, + res_tree=res_tree, + out_tree=out_tree) + return tree_unflatten(out_tree, out_flat) + + +### utils + +def tree_fill(x, treedef): + return tree_unflatten(treedef, [x] * treedef.num_leaves) + +def tree_fill_like(x, tree): + return tree_fill(x, tree_structure(tree)) + +def tree_broadcast(full_treedef, tree, is_leaf=None): + full_tree = tree_fill(0, full_treedef) + return tree_map(tree_fill_like, tree, full_tree, is_leaf=is_leaf) + +def is_treedef_prefix(entire, prefix): + entire = tree_fill(0, entire) + prefix = tree_fill(0, prefix) + try: + tree_map(lambda x, y: x, prefix, entire) + except ValueError: + return False + return True + +def rule_name(rule): + return getattr(rule, '__name__', '') + +def check_transpose_rule_trees(rule, lin_tree, rule_out_tree): + if not is_treedef_prefix(lin_tree, rule_out_tree): + if hasattr(rule, '_transpose_type_error'): + raise rule._transpose_type_error(lin_tree, rule_out_tree) + else: + raise TypeError( + 'structure of custom transpose rule\'s output does not prefix-match ' + 'structure of primal function\'s linear inputs under ' + f'custom transpose rule ({rule_name(rule)}).\n' + f'Transpose rule output: {rule_out_tree}\n' + f'Linear primal inputs: {lin_tree}') + +def make_transpose_from_thunk(thunk, lin_tree): + transpose_jaxpr, transpose_consts = thunk() + transpose_jaxpr = core.ClosedJaxpr( + pe.convert_constvars_jaxpr(transpose_jaxpr), ()) + def transpose(res_arg, ct_out): + args_flat = tree_leaves((res_arg, ct_out)) + ct_ins = core.jaxpr_as_fun(transpose_jaxpr)(*transpose_consts, *args_flat) + return tree_unflatten(lin_tree, ct_ins) + return transpose + + +### custom_transpose primitive and rules + +class CustomTransposePrimitive(core.Primitive): + call_primitive = False + map_primitive = False + multiple_results = True + + def bind(self, call, *args, **params): + # TODO(frostig,mattjj): This doesn't handle closures yet, which is + # a bit involved. Closures are complicated by us binding `call` + # twice in the JVP rule for custom transpose. The `env_trace_todo` + # output by `process_env_traces` due to one of those two bindings + # should be passable to the other, and need to be passed onward + # since the second bind is deferred by partial eval (since it + # typically receives unknowns) + top_trace = core.find_top_trace(args) + tracers = map(top_trace.full_raise, args) + outs = top_trace.process_custom_transpose(self, call, tracers, **params) + return outs + + # TODO(frostig,mattjj): consider keeping `call` as a named parameter + # instead of following this "call primitive" convention. + def get_bind_params(self, params): + assert 'call_jaxpr' in params + assert 'transpose_jaxpr_thunk' in params + new_params = dict(params) + new_params['transpose'] = make_transpose_from_thunk( + new_params.pop('transpose_jaxpr_thunk'), + new_params['lin_tree']) + call = lu.wrap_init(core.jaxpr_as_fun(new_params.pop('call_jaxpr'))) + return [call], new_params + + +# TODO(frostig,mattjj): reinstate checks +def custom_transpose_typecheck(_, *in_atoms, out_types, **params): + del in_atoms, params + return out_types, core.no_effects + + +def custom_transpose_transpose_rule( + cts, *args, out_types, res_tree, lin_tree, out_tree, **params): + + if 'transpose_jaxpr_thunk' in params: + assert 'call_jaxpr' in params + transpose = make_transpose_from_thunk( + params['transpose_jaxpr_thunk'], lin_tree) + else: + assert 'call' in params + transpose = params['transpose'] + + call_in_tree = treedef_tuple((res_tree, lin_tree)) + + # TODO(frostig,mattjj): `lin_arg` indicates the inputs with respect + # to which we are transposing (via `ad.is_undefined_primal`). + # Consider passing this information to the custom transpose rule? + + res_arg, lin_arg = tree_unflatten(call_in_tree, args) + del lin_arg + assert all(not ad.is_undefined_primal(x) for x in tree_leaves(res_arg)) + + cts = [ad_util.zeros_like_aval(ct.aval) if type(ct) is ad_util.Zero else ct + for ct in cts] + ct_out = tree_unflatten(out_tree, cts) + ct_lin = transpose(res_arg, ct_out) + check_transpose_rule_trees(transpose, lin_tree, tree_structure(ct_lin)) + ct_lin_flat, _ = tree_flatten( + tree_broadcast(lin_tree, ct_lin, is_leaf=lambda x: x is None), + is_leaf=lambda x: x is None) + return [None] * len(tree_leaves(res_arg)) + ct_lin_flat + + +def custom_transpose_lowering(*args, call_jaxpr, **params): + return core.jaxpr_as_fun(call_jaxpr)(*args) + + +custom_transpose_p = CustomTransposePrimitive('custom_transpose_call') +core.custom_typechecks[custom_transpose_p] = custom_transpose_typecheck +ad.primitive_transposes[custom_transpose_p] = custom_transpose_transpose_rule +mlir.register_lowering( + custom_transpose_p, + mlir.lower_fun(custom_transpose_lowering, multiple_results=True)) +xla.register_initial_style_primitive(custom_transpose_p) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/debugger/cli_debugger.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/debugger/cli_debugger.py new file mode 100644 index 0000000000000000000000000000000000000000..bf4b387650265a9da8a752881a790ebd18131486 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/debugger/cli_debugger.py @@ -0,0 +1,168 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import cmd +import pprint +import sys +import traceback +from typing import Any, IO + +from jax._src.debugger import core as debugger_core + +DebuggerFrame = debugger_core.DebuggerFrame + +class CliDebugger(cmd.Cmd): + """A text-based debugger.""" + prompt = '(jdb) ' + + def __init__(self, frames: list[DebuggerFrame], thread_id, + stdin: IO[str] | None = None, stdout: IO[str] | None = None, + completekey: str = "tab"): + super().__init__(stdin=stdin, stdout=stdout, completekey=completekey) + self.use_rawinput = stdin is None + self.frames = frames + self.frame_index = 0 + self.thread_id = thread_id + self.intro = 'Entering jdb:' + + def current_frame(self): + return self.frames[self.frame_index] + + def evaluate(self, expr): + env = {} + curr_frame = self.frames[self.frame_index] + env.update(curr_frame.globals) + env.update(curr_frame.locals) + return eval(expr, {}, env) + + def default(self, arg): + """Evaluates an expression.""" + try: + print(repr(self.evaluate(arg)), file=self.stdout) + except: + self._error_message() + + def print_backtrace(self): + backtrace = [] + backtrace.append('Traceback:') + for frame in self.frames[::-1]: + backtrace.append(f' File "{frame.filename}", line {frame.lineno}') + if frame.offset is None: + backtrace.append(' ') + else: + line = frame.source[frame.offset] + backtrace.append(f' {line.strip()}') + print("\n".join(backtrace), file=self.stdout) + + def print_context(self, num_lines=2): + curr_frame = self.frames[self.frame_index] + context = [] + context.append(f'> {curr_frame.filename}({curr_frame.lineno})') + for i, line in enumerate(curr_frame.source): + assert curr_frame.offset is not None + if (curr_frame.offset - 1 - num_lines <= i <= + curr_frame.offset + num_lines): + if i == curr_frame.offset: + context.append(f'-> {line}') + else: + context.append(f' {line}') + print("\n".join(context), file=self.stdout) + + def _error_message(self): + exc_info = sys.exc_info()[:2] + msg = traceback.format_exception_only(*exc_info)[-1].strip() + print('***', msg, file=self.stdout) + + def do_p(self, arg): + """p expression + Evaluates and prints the value of an expression + """ + try: + print(repr(self.evaluate(arg)), file=self.stdout) + except: + self._error_message() + + def do_pp(self, arg): + """pp expression + Evaluates and pretty-prints the value of an expression + """ + try: + print(pprint.pformat(self.evaluate(arg)), file=self.stdout) + except: + self._error_message() + + def do_up(self, _): + """u(p) + Move down a stack frame. + """ + if self.frame_index == len(self.frames) - 1: + print('At topmost frame.', file=self.stdout) + else: + self.frame_index += 1 + self.print_context() + do_u = do_up + + def do_down(self, _): + """d(own) + Move down a stack frame. + """ + if self.frame_index == 0: + print('At bottommost frame.', file=self.stdout) + else: + self.frame_index -= 1 + self.print_context() + do_d = do_down + + def do_list(self, _): + """l(ist) + List source code for the current file. + """ + self.print_context(num_lines=5) + do_l = do_list + + def do_continue(self, _): + """c(ont(inue)) + Continue the program's execution. + """ + return True + do_c = do_cont = do_continue + + def do_quit(self, _): + """q(uit)\n(exit) + Quit the debugger. The program is given an exit command. + """ + sys.exit(0) + do_q = do_EOF = do_exit = do_quit + + def do_where(self, _): + """w(here) + Prints a stack trace with the most recent frame on the bottom. + 'bt' is an alias for this command. + """ + self.print_backtrace() + do_w = do_bt = do_where + + def run(self): + while True: + try: + self.cmdloop() + break + except KeyboardInterrupt: + print('--KeyboardInterrupt--', file=sys.stdout) + +def run_debugger(frames: list[DebuggerFrame], thread_id: int | None, + **kwargs: Any): + CliDebugger(frames, thread_id, **kwargs).run() +debugger_core.register_debugger("cli", run_debugger, -1) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/debugging.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/debugging.py new file mode 100644 index 0000000000000000000000000000000000000000..265dcd0de8434d278665579a02f381987fe48657 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/debugging.py @@ -0,0 +1,614 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Module for JAX debugging primitives and related functionality.""" + +from __future__ import annotations + +import importlib.util +from collections.abc import Sequence +import functools +import logging +import string +import sys +from typing import Any, Callable, Union +import weakref + +import numpy as np + +import jax +import jax.numpy as jnp +from jax import lax + +from jax._src import core +from jax._src import effects +from jax._src import mesh as mesh_lib +from jax._src import sharding_impls +from jax._src import dispatch +from jax._src import tree_util +from jax._src import util +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import partial_eval as pe +from jax._src.lib import xla_client as xc +from jax._src.lib.mlir import ir +from jax._src.lib.mlir.dialects import hlo +from jax._src.sharding import Sharding +from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding + +logger = logging.getLogger(__name__) + +class DebugEffect(effects.Effect): + __str__ = lambda self: "Debug" +debug_effect = DebugEffect() + +class OrderedDebugEffect(effects.Effect): + __str__ = lambda self: "OrderedDebug" +ordered_debug_effect = OrderedDebugEffect() + +effects.ordered_effects.add_type(OrderedDebugEffect) +effects.lowerable_effects.add_type(DebugEffect) +effects.lowerable_effects.add_type(OrderedDebugEffect) +effects.control_flow_allowed_effects.add_type(DebugEffect) +effects.control_flow_allowed_effects.add_type(OrderedDebugEffect) +effects.remat_allowed_effects.add_type(DebugEffect) +effects.remat_allowed_effects.add_type(OrderedDebugEffect) +effects.custom_derivatives_allowed_effects.add_type(DebugEffect) +effects.custom_derivatives_allowed_effects.add_type(OrderedDebugEffect) + +# `debug_callback_p` is the main primitive for staging out Python callbacks. +debug_callback_p = core.Primitive('debug_callback') +debug_callback_p.multiple_results = True + +map, unsafe_map = util.safe_map, map + +@debug_callback_p.def_impl +def debug_callback_impl(*args, callback: Callable[..., Any], + effect: DebugEffect): + del effect + try: + cpu_device, *_ = jax.local_devices(backend="cpu") + except RuntimeError as e: + raise RuntimeError( + "jax.debug.callback failed to find a local CPU device to place the" + " inputs on. Make sure \"cpu\" is listed in --jax_platforms or the" + " JAX_PLATFORMS environment variable." + ) from e + args = jax.device_put(args, cpu_device) + with jax.default_device(cpu_device): + try: + callback(*args) + except BaseException: + logger.exception("jax.debug.callback failed") + raise + return () + +@debug_callback_p.def_effectful_abstract_eval +def debug_callback_abstract_eval(*flat_avals, callback: Callable[..., Any], + effect: DebugEffect): + del flat_avals, callback + return [], {effect} + +def debug_callback_batching_rule(args, dims, **params): + """Unrolls the debug callback across the mapped axis.""" + axis_size = next(x.shape[i] for x, i in zip(args, dims) + if i is not None) + # TODO(sharadmv): implement in terms of rolled loop unstead of + # unrolled. + def get_arg_at_dim(i, dim, arg): + if dim is batching.not_mapped: + # Broadcast unmapped argument + return arg + return lax.index_in_dim(arg, i, axis=dim, keepdims=False) + outs = [] + for i in range(axis_size): + args_idx = map(functools.partial(get_arg_at_dim, i), dims, args) + outs.append(debug_callback_p.bind(*args_idx, **params)) + outs = [jnp.stack(xs) for xs in zip(*outs)] + return outs, (0,) * len(outs) +batching.primitive_batchers[debug_callback_p] = debug_callback_batching_rule + +def debug_callback_jvp_rule(primals, tangents, **params): + return debug_callback_p.bind(*primals, **params), [] +ad.primitive_jvps[debug_callback_p] = debug_callback_jvp_rule + +def debug_callback_transpose_rule(*flat_args, callback: Callable[..., Any], + effect: DebugEffect): + del flat_args, callback, effect + raise ValueError("Transpose doesn't support debugging callbacks.") +ad.primitive_transposes[debug_callback_p] = debug_callback_transpose_rule + +def debug_callback_lowering(ctx, *args, effect, callback, **params): + + axis_context = ctx.module_context.axis_context + if (isinstance(axis_context, sharding_impls.SPMDAxisContext) and + set(axis_context.manual_axes) == set(axis_context.mesh.axis_names)): + # If we have fully manual sharding during lowering, that means the JAX + # program has per-device semantics, so we run the callback on each device. + sharding = xc.OpSharding() + sharding.type = xc.OpSharding.Type.MANUAL + elif isinstance( + axis_context, + (sharding_impls.ShardingContext, sharding_impls.SPMDAxisContext), + ): + # If we have fully automatic sharding during lowering, that means the JAX + # program has bulk array semantics, so we run the callback with a MAXIMAL + # sharding and hence execute it only once on the full logical value). + # If we have partially automatic sharding, we do this too... not sure why! + sharding = xc.OpSharding() + sharding.type = xc.OpSharding.Type.MAXIMAL + sharding.tile_assignment_dimensions = [1] + sharding.tile_assignment_devices = [0] + else: + # When there's no SPMD partitioning going on, don't annotate a sharding. + sharding = None + + def _callback(*flat_args): + debug_callback_p.impl( + *flat_args, effect=effect, callback=callback, **params) + return () + if effects.ordered_effects.contains(effect): + [token] = ctx.tokens_in.get(effect) + result, token, _ = mlir.emit_python_callback( + ctx, _callback, token, list(args), ctx.avals_in, ctx.avals_out, + has_side_effect=True) + ctx.set_tokens_out(mlir.TokenSet({effect: (token,)})) + else: + result, _, _ = mlir.emit_python_callback( + ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, + has_side_effect=True, sharding=sharding) + return result +mlir.register_lowering(debug_callback_p, debug_callback_lowering, + platform="cpu") +mlir.register_lowering( + debug_callback_p, debug_callback_lowering, platform="gpu") +mlir.register_lowering( + debug_callback_p, debug_callback_lowering, platform="tpu") + +def _debug_callback_partial_eval_custom(saveable, unks_in, inst_in, eqn): + # The default behavior for effectful primitives is to not stage them if + # possible. For debug callback, we actually want it to be staged to + # provide more information to the user. This rule bypasses partial_eval's + # regular behavior to do that. Specifically, we will stage the callback + # if: + # 1) the policy says debug_callbacks are not saveable + # 2) the policy says debug_callbacks are saveable BUT all of the input + # values are instantiated. + # The purpose is to call back with as much information as possible while + # avoiding unnecessarily staging out other values. + if any(unks_in): + # The usual case (if we have any unknowns, we need to stage it out) + res = [v for v, inst in zip(eqn.invars, inst_in) if not inst] + return None, eqn, [], [], res + if saveable(debug_callback_p, *[v.aval for v in eqn.invars], **eqn.params): + # The policy is telling us we can save the debug callback. + if all(inst_in): + # If all of the inputs are instantiated, we also stage out the + # debug_callback. + return eqn, eqn, [], [], [] + else: + # If any are not instantiated, we don't do any extra staging to avoid + # affecting the computation. + return eqn, None, [], [], [] + # If we can't save the debug callback (thanks to the policy) we listen to + # the policy and stage out the debug callback. + return eqn, eqn, [], [], [] +pe.partial_eval_jaxpr_custom_rules[debug_callback_p] = ( + _debug_callback_partial_eval_custom) + +def debug_callback(callback: Callable[..., None], *args: Any, + ordered: bool = False, **kwargs: Any) -> None: + """Calls a stageable Python callback. + + For more explanation, see `External Callbacks`_. + + ``jax.debug.callback`` enables you to pass in a Python function that can be called + inside of a staged JAX program. A ``jax.debug.callback`` follows existing JAX + transformation *pure* operational semantics, which are therefore unaware of + side-effects. This means the effect could be dropped, duplicated, or + potentially reordered in the presence of higher-order primitives and + transformations. + + We want this behavior because we'd like ``jax.debug.callback`` to be "innocuous", + i.e. we want these primitives to change the JAX computation as little as + possible while revealing as much about them as possible, such as which parts + of the computation are duplicated or dropped. + + Args: + callback: A Python callable returning None. + *args: The positional arguments to the callback. + ordered: A keyword only argument used to indicate whether or not the + staged out computation will enforce ordering of this callback w.r.t. + other ordered callbacks. + **kwargs: The keyword arguments to the callback. + + Returns: + None + + See Also: + - :func:`jax.experimental.io_callback`: callback designed for impure functions. + - :func:`jax.pure_callback`: callback designed for pure functions. + - :func:`jax.debug.print`: callback designed for printing. + + .. _External Callbacks: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html + """ + if not callable(callback): + raise TypeError("first argument to jax.debug.callback must be callable, " + f"but got an object of type {type(callback)}") + flat_args, in_tree = tree_util.tree_flatten((args, kwargs)) + effect = ordered_debug_effect if ordered else debug_effect + def _flat_callback(*flat_args): + args, kwargs = tree_util.tree_unflatten(in_tree, flat_args) + callback(*args, **kwargs) + return () + debug_callback_p.bind(*flat_args, callback=_flat_callback, effect=effect) + +class _DebugPrintFormatChecker(string.Formatter): + + def check_unused_args(self, used_args, args, kwargs): + unused_args = [arg for i, arg in enumerate(args) if i not in used_args] + unused_kwargs = [k for k in kwargs if k not in used_args] + if unused_args: + raise ValueError( + f"Unused positional arguments to `jax.debug.print`: {unused_args}") + if unused_kwargs: + raise ValueError( + f"Unused keyword arguments to `jax.debug.print`: {unused_kwargs}. " + "You may be passing an f-string (i.e, `f\"{x}\"`) into " + "`jax.debug.print` and instead should pass in a regular string.") + +formatter = _DebugPrintFormatChecker() + +def _format_print_callback(fmt: str, *args, **kwargs): + sys.stdout.write(fmt.format(*args, **kwargs) + "\n") + +def debug_print(fmt: str, *args, ordered: bool = False, **kwargs) -> None: + """Prints values and works in staged out JAX functions. + + This function does *not* work with f-strings because formatting is delayed. + So instead of ``jax.debug.print(f"hello {bar}")``, write + ``jax.debug.print("hello {bar}", bar=bar)``. + + This function is a thin convenience wrapper around :func:`jax.debug.callback`. + The implementation is essentially:: + + def debug_print(fmt: str, *args, **kwargs): + jax.debug.callback( + lambda *args, **kwargs: print(fmt.format(*args, **kwargs)), + *args, **kwargs) + + It may be useful to call :func:`jax.debug.callback` directly instead of this + convenience wrapper. For example, to get debug printing in logs, you might + use :func:`jax.debug.callback` together with ``logging.log``. + + Args: + fmt: A format string, e.g. ``"hello {x}"``, that will be used to format + input arguments, like ``str.format``. See the Python docs on + `string formatting `_ + and `format string syntax `_. + *args: A list of positional arguments to be formatted, as if passed to + ``fmt.format``. + ordered: A keyword only argument used to indicate whether or not the + staged out computation will enforce ordering of this ``jax.debug.print`` + w.r.t. other ordered ``jax.debug.print`` calls. + **kwargs: Additional keyword arguments to be formatted, as if passed to + ``fmt.format``. + """ + # Check that we provide the correct arguments to be formatted + formatter.format(fmt, *args, **kwargs) + + debug_callback(functools.partial(_format_print_callback, fmt), *args, + **kwargs, ordered=ordered) + + +# Sharding visualization + +inspect_sharding_p = core.Primitive("inspect_sharding") +inspect_sharding_p.multiple_results = True +dispatch.prim_requires_devices_during_lowering.add(inspect_sharding_p) + +def _inspect_sharding_impl(value, *, callback): + callback(value.sharding) + return [] +inspect_sharding_p.def_impl(_inspect_sharding_impl) + +def _inspect_sharding_abstract_eval(aval, **_): + del aval + # Effectful abstract avoids DCE + return [], {debug_effect} +inspect_sharding_p.def_effectful_abstract_eval(_inspect_sharding_abstract_eval) + +def _inspect_sharding_batching_rule(args, _, *, callback): + value, = args + inspect_sharding_p.bind(value, callback=callback) + return [], [] +batching.primitive_batchers[inspect_sharding_p] = ( + _inspect_sharding_batching_rule) + +def _inspect_sharding_jvp_rule(primals, _, **params): + return inspect_sharding_p.bind(*primals, **params), [] +ad.primitive_jvps[inspect_sharding_p] = _inspect_sharding_jvp_rule + +sharding_callbacks = weakref.WeakValueDictionary() # type: ignore +_INSPECT_SHARDING_CALL_NAME = "InspectSharding" + +class ShardingCallbackInfo: + def __init__(self, callback, module_context): + self.callback = callback + self.module_context = module_context + +def _inspect_sharding_lowering_rule(ctx: mlir.LoweringRuleContext, value, *, + callback): + + mesh = mesh_lib.thread_resources.env.physical_mesh + axis_context = ctx.module_context.axis_context + + if isinstance(axis_context, sharding_impls.ShardingContext): + devices = axis_context.device_assignment + if devices is None: + raise AssertionError( + 'Please file a bug at https://github.com/google/jax/issues') + elif isinstance(axis_context, sharding_impls.SPMDAxisContext): + devices = axis_context.mesh._flat_devices_tuple + else: + raise NotImplementedError(type(axis_context)) + assert devices is not None + + # If we have a nontrivial parallel computation, we need to wait until the SPMD + # partitioner calls back with the `HloSharding. + def _hlo_sharding_callback(hlo_sharding: xc.HloSharding): + if mesh.empty: + return callback( + sharding_impls._op_sharding_to_pos_sharding(hlo_sharding, devices)) + pspec = parse_flatten_op_sharding(hlo_sharding, mesh)[0].get_partition_spec() + return callback(NamedSharding(mesh, pspec)) + + if len(devices) == 1: + # If we only have one device in our computation, we can construct a + # replicated HloSharding and call it right now. + _hlo_sharding_callback(sharding_impls.get_replicated_hlo_sharding()) + return [] + + key = xc.encode_inspect_sharding_callback(_hlo_sharding_callback) + # We need to make sure `_hlo_sharding_callback` is still alive when the SPMD + # partitioner runs so we keep it alive by attaching it to the executable. # + ctx.module_context.add_keepalive(_hlo_sharding_callback) + + hlo.CustomCallOp([value.type], [value], + call_target_name=ir.StringAttr.get( + _INSPECT_SHARDING_CALL_NAME), + has_side_effect=ir.BoolAttr.get(True), + api_version=mlir.i32_attr(1), + called_computations=ir.ArrayAttr.get([]), + backend_config=ir.StringAttr.get(key), # type: ignore[arg-type] + operand_layouts=None, + result_layouts=None) + return [] +mlir.register_lowering(inspect_sharding_p, _inspect_sharding_lowering_rule) + +def inspect_sharding_prop_user_sharding(sharding, backend_string): + del sharding, backend_string + return [] + +def inspect_sharding_infer_sharding_from_operands(arg_shapes, arg_shardings, + shape, backend_string): + del arg_shapes, shape, backend_string + return arg_shardings[0] + +def _slice_to_chunk_idx(size: int, slc: slice) -> int: + if slc.stop == slc.start == None: + return 0 + slice_size = slc.stop - slc.start + assert slc.start % slice_size == 0 + assert size % slice_size == 0 + return slc.start // slice_size + +def _raise_to_slice(slc: slice | int): + if isinstance(slc, int): + return slice(slc, slc + 1) + return slc + +Color = Union[tuple[float, float, float], str] +ColorMap = Callable[[float], tuple[float, float, float, float]] + +def _canonicalize_color(color: Color) -> str: + if isinstance(color, str): + return color + r, g, b = (int(a * 255) for a in color) + return f"#{r:02X}{g:02X}{b:02X}" + +def _get_text_color(color: str) -> str: + r, g, b = map(lambda x: int(x, 16), (color[1:3], color[3:5], color[5:7])) + if (r * 0.299 + g * 0.587 + b * 0.114) > 186: + return "#000000" + return "#ffffff" + +def make_color_iter(color_map, num_rows, num_cols): + num_colors = num_rows * num_cols + color_values = np.linspace(0, 1, num_colors) + idx = 0 + for _ in range(num_colors): + yield color_map(color_values[idx]) + idx = (idx + num_colors // 2 + bool(num_colors % 2 == 0)) % num_colors + +def visualize_sharding(shape: Sequence[int], sharding: Sharding, *, + use_color: bool = True, scale: float = 1., + min_width: int = 9, max_width: int = 80, + color_map: ColorMap | None = None): + """Visualizes a ``Sharding`` using ``rich``.""" + if not importlib.util.find_spec("rich"): + raise ValueError("`visualize_sharding` requires `rich` to be installed.") + + # These imports are local so that they don't affect JAX import times. + # pytype: disable=import-error + import rich.align + import rich.console + import rich.box + import rich.padding + import rich.style + import rich.table + # pytype: enable=import-error + + if len(shape) > 2 or len(shape) < 1: + raise ValueError( + "`visualize_sharding` only works for shapes with 1 and 2 dimensions.") + console = rich.console.Console(width=max_width) + use_color = use_color and console.color_system is not None + if use_color and not color_map: + try: + import matplotlib as mpl # pytype: disable=import-error + color_map = mpl.colormaps["tab20b"] + except ModuleNotFoundError: + use_color = False + + base_height = int(10 * scale) + aspect_ratio = (shape[1] if len(shape) == 2 else 1) / shape[0] + base_width = int(base_height * aspect_ratio) + height_to_width_ratio = 2.5 + + # Grab the device kind from the first device + device_kind = next(iter(sharding.device_set)).platform.upper() + + device_indices_map = sharding.devices_indices_map(tuple(shape)) + slices: dict[tuple[int, ...], set[int]] = {} + heights: dict[tuple[int, ...], float | None] = {} + widths: dict[tuple[int, ...], float] = {} + + for i, (dev, slcs) in enumerate(device_indices_map.items()): + assert slcs is not None + slcs = tuple(map(_raise_to_slice, slcs)) + chunk_idxs = tuple(map(_slice_to_chunk_idx, shape, slcs)) + if slcs is None: + raise NotImplementedError + if len(slcs) == 2: + vert, horiz = slcs + vert_size = ((vert.stop - vert.start ) if vert.stop is not None + else shape[0]) + horiz_size = ((horiz.stop - horiz.start) if horiz.stop is not None + else shape[1]) + chunk_height = vert_size / shape[0] + chunk_width = horiz_size / shape[1] + heights[chunk_idxs] = chunk_height + widths[chunk_idxs] = chunk_width + else: + # In the 1D case, we set the height to 1. + horiz, = slcs + vert = slice(0, 1, None) + horiz_size = ( + (horiz.stop - horiz.start) if horiz.stop is not None else shape[0]) + chunk_idxs = (0, *chunk_idxs) + heights[chunk_idxs] = None + widths[chunk_idxs] = horiz_size / shape[0] + slices.setdefault(chunk_idxs, set()).add(dev.id) + num_rows = max(a[0] for a in slices.keys()) + 1 + if len(list(slices.keys())[0]) == 1: + num_cols = 1 + else: + num_cols = max(a[1] for a in slices.keys()) + 1 + + color_iter = make_color_iter(color_map, num_rows, num_cols) + table = rich.table.Table(show_header=False, show_lines=not use_color, + padding=0, + highlight=not use_color, pad_edge=False, + box=rich.box.SQUARE if not use_color else None) + for i in range(num_rows): + col = [] + for j in range(num_cols): + entry = f"{device_kind} "+",".join([str(s) for s in sorted(slices[i, j])]) + width, maybe_height = widths[i, j], heights[i, j] + width = int(width * base_width * height_to_width_ratio) + if maybe_height is None: + height = 1 + else: + height = int(maybe_height * base_height) + width = min(max(width, min_width), max_width) + left_padding, remainder = divmod(width - len(entry) - 2, 2) + right_padding = left_padding + remainder + top_padding, remainder = divmod(height - 2, 2) + bottom_padding = top_padding + remainder + if use_color: + color = _canonicalize_color(next(color_iter)[:3]) + text_color = _get_text_color(color) + top_padding += 1 + bottom_padding += 1 + left_padding += 1 + right_padding += 1 + else: + color = None + text_color = None + padding = (top_padding, right_padding, bottom_padding, left_padding) + padding = tuple(max(x, 0) for x in padding) # type: ignore + col.append( + rich.padding.Padding( + rich.align.Align(entry, "center", vertical="middle"), padding, + style=rich.style.Style(bgcolor=color, + color=text_color))) + table.add_row(*col) + console.print(table, end='\n\n') + +def inspect_array_sharding(value, *, callback: Callable[[Sharding], None]): + """Enables inspecting array sharding inside JIT-ted functions. + + This function, when provided with a Pytree of arrays, calls back with each of + their shardings and works in ``pjit``-ted computations, enabling inspecting + the chosen intermediate shardings. + + The policy for when ``callback`` is called is *as early as possible* when the + sharding information is available. This means if ``inspect_array_callback`` is + called without any transformations, the callback will happen immediately + since we have the array and its sharding readily available. Inside of a + ``jax.jit``, the callback will happen at lowering time, meaning you can + trigger the callback using the AOT API (``jit(f).lower(...)``). When inside of + a ``pjit``, the callback happens *at compile time* since the sharding is + determined by XLA. You can trigger the callback by using JAX's AOT API + (``pjit(f).lower(...).compile()``). In all cases, the callback will be + triggered by running the function, since running a function entails lowering + and compiling it first. However, once the function is compiled and cached, + the callback will no longer occur. + + This function is experimental and its behavior may change in the future. + + Args: + value: A Pytree of JAX arrays. + callback: A callable that takes in a ``Sharding`` and doesn't return a value. + + In the following example, we print out the sharding of an intermediate value + in a ``pjit``-ted computation: + + >>> import jax + >>> import jax.numpy as jnp + >>> from jax.experimental.pjit import pjit + >>> from jax.sharding import Mesh, PartitionSpec + >>> + >>> x = jnp.arange(8, dtype=jnp.float32) + >>> def f_(x): + ... x = jnp.sin(x) + ... jax.debug.inspect_array_sharding(x, callback=print) + ... return jnp.square(x) + >>> f = pjit(f_, in_shardings=PartitionSpec('dev'), + ... out_shardings=PartitionSpec('dev')) + >>> with Mesh(jax.devices(), ('dev',)): + ... f.lower(x).compile() # doctest: +SKIP + ... + NamedSharding(mesh={'dev': 8}, partition_spec=PartitionSpec(('dev',),)) + """ + def _inspect(val): + inspect_sharding_p.bind(val, callback=callback) + tree_util.tree_map(_inspect, value) + +def visualize_array_sharding(arr, **kwargs): + """Visualizes an array's sharding.""" + def _visualize(sharding): + return visualize_sharding(arr.shape, sharding, **kwargs) + inspect_array_sharding(arr, callback=_visualize) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/deprecations.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/deprecations.py new file mode 100644 index 0000000000000000000000000000000000000000..0bbae2d140ead4da732778b2822e795f97caaece --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/deprecations.py @@ -0,0 +1,103 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from types import ModuleType +import warnings + +# Module __getattr__ factory that warns if deprecated names are used. +# +# Example usage: +# from jax._src.interpreters.pxla import ( +# Mesh as _deprecated_Mesh, +# ) +# +# _deprecations = { +# # Added Feb 8, 2023: +# "Mesh": ( +# "jax.interpreters.pxla.Mesh is deprecated. Use jax.sharding.Mesh.", +# _deprecated_Mesh, +# ), +# } +# +# from jax._src.deprecations import deprecation_getattr as _deprecation_getattr +# __getattr__ = _deprecation_getattr(__name__, _deprecations) +# del _deprecation_getattr + +# Note that type checkers such as Pytype will not know about the deprecated +# names. If it is desirable that a deprecated name is known to the type checker, +# add: +# import typing +# if typing.TYPE_CHECKING: +# from jax._src.interpreters.pxla import ( +# Mesh as Mesh, +# ) +# del typing +def deprecation_getattr(module, deprecations): + def getattr(name): + if name in deprecations: + message, fn = deprecations[name] + if fn is None: # Is the deprecation accelerated? + raise AttributeError(message) + warnings.warn(message, DeprecationWarning, stacklevel=2) + return fn + raise AttributeError(f"module {module!r} has no attribute {name!r}") + + return getattr + + +def accelerate_getattr_deprecation(module: ModuleType, name: str) -> None: + """Accelerate the deprecation of a module-level attribute. + + Raises an AttributeError instead of a DeprecationWarning upon attribute access. + Used in Google-internal code to implement faster deprecation. + """ + message, _ = module._deprecations[name] + module._deprecations[name] = (message, None) + +# The following mechanism is a separate one, for registering and +# accelerating deprecations that are not imports (for example, deprecations +# of a function argument). +# Maps a globally unique string ID to a DeprecationState, which tracks whether +# the deprecation is accelerated. +# The intent is that non-accelerated deprecations will warn, and accelerated +# deprecations will error. + +@dataclass +class DeprecationState: + accelerated: bool = False + +_registered_deprecations: dict[str, DeprecationState] = {} + + +def register(deprecation_id: str) -> None: + _registered_deprecations[deprecation_id] = DeprecationState() + + +def unregister(deprecation_id: str) -> None: + if deprecation_id not in _registered_deprecations: + raise ValueError(f"{deprecation_id=!r} not registered.") + _registered_deprecations.pop(deprecation_id) + + +def accelerate(deprecation_id: str) -> None: + if deprecation_id not in _registered_deprecations: + raise ValueError(f"{deprecation_id=!r} not registered.") + _registered_deprecations[deprecation_id].accelerated = True + + +def is_accelerated(deprecation_id: str) -> bool: + if deprecation_id not in _registered_deprecations: + raise ValueError(f"{deprecation_id=!r} not registered.") + return _registered_deprecations[deprecation_id].accelerated diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dispatch.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..2836fd9b239127d8b232ac07ebf23f8f9d7a6464 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dispatch.py @@ -0,0 +1,499 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Primitive dispatch and jit dispatch. +from __future__ import annotations + +import atexit +from collections.abc import Iterator, Sequence +import contextlib +from functools import partial +import itertools +import time +from typing import Any, Callable, NamedTuple +import logging +import threading + +import numpy as np + +import jax +from jax._src import basearray +from jax._src import config +from jax._src import core +from jax._src import api +from jax._src import dtypes +from jax._src import source_info_util +from jax._src import traceback_util +from jax._src import util +from jax._src import xla_bridge as xb +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import xla +from jax._src.interpreters import pxla +from jax._src import lib +from jax._src.lib import xla_client as xc +from jax._src.monitoring import record_event_duration_secs +from jax._src.partition_spec import PartitionSpec +from jax._src.sharding import Sharding +from jax._src.sharding_impls import ( + SingleDeviceSharding, NamedSharding, + GSPMDSharding, TransferToMemoryKind, is_single_device_sharding) +from jax._src.layout import Layout, DeviceLocalLayout + + +JAXPR_TRACE_EVENT = "/jax/core/compile/jaxpr_trace_duration" +JAXPR_TO_MLIR_MODULE_EVENT = "/jax/core/compile/jaxpr_to_mlir_module_duration" +BACKEND_COMPILE_EVENT = "/jax/core/compile/backend_compile_duration" + +traceback_util.register_exclusion(__file__) + +xe = xc._xla + +Backend = xe.Client +Device = xc.Device + +CompileOptions = xc.CompileOptions + +map, unsafe_map = util.safe_map, map +zip, unsafe_zip = util.safe_zip, zip + +logger = logging.getLogger(__name__) + +# This flag is set on exit; no logging should be attempted +_on_exit = False + +### op-by-op execution + +def apply_primitive(prim, *args, **params): + """Impl rule that compiles and runs a single primitive 'prim' using XLA.""" + fun = xla_primitive_callable(prim, **params) + # TODO(yashkatariya): Investigate adding is_primitive to jit and never + # triggering the disable jit path instead of messing around with it here. + prev = lib.jax_jit.swap_thread_local_state_disable_jit(False) + try: + outs = fun(*args) + finally: + lib.jax_jit.swap_thread_local_state_disable_jit(prev) + return outs + +@util.cache() +def xla_primitive_callable(prim: core.Primitive, **params): + def prim_fun(*args): + return prim.bind(*args, **params) + prim_fun.__name__ = prim.name + prim_fun.__qualname__ = prim.name + return api.jit(prim_fun) + + +def simple_impl(prim): + prim.def_impl(partial(apply_primitive, prim)) + +RuntimeToken = Any + +class RuntimeTokenSet(threading.local): + """See docstring for effects.py module for the calling convention for tokens.""" + + # For each ordered effect, the token returned by the last dispatched + # computation, sharded over the devices in that computation. + current_tokens: dict[core.Effect, core.Token] + + # For each device, the runtime token returned by the last dispatched + # computation on that device. + output_runtime_tokens: dict[Device, RuntimeToken] + + def __init__(self): + self.current_tokens = {} + self.output_runtime_tokens = {} + + def get_token_input( + self, eff: core.Effect, devices: list[Device] + ) -> core.Token: + tok = self.current_tokens.get(eff, np.zeros(0, np.bool_)) + + if isinstance(tok, core.Token): + # The order of devices may change, so we need to reshard if necessary. + # TODO(yueshengys): This might still be buggy in a multi-process SPMD + # scenario. Revise the logic later. A distributed shutdown barrier inside + # the XLA program may be needed. + return jax.device_put(tok, jax.sharding.PositionalSharding(devices)) + + # We only use replicated sharding for the first time when the token for the + # order effect hasn't been created. + s = jax.sharding.GSPMDSharding.get_replicated(devices) + sharded_tok = core.Token(pxla.shard_args([s], [tok])[0]) + self.current_tokens[eff] = sharded_tok + return sharded_tok + + def set_token_result(self, eff: core.Effect, token: core.Token): + self.current_tokens[eff] = token + + def set_output_runtime_token(self, device: Device, token: RuntimeToken): + # We're free to clobber the previous output token because on each + # device we have a total ordering of computations. Only the token + # from the latest computation matters. + self.output_runtime_tokens[device] = token + + def clear(self): + self.current_tokens = {} + self.output_runtime_tokens = {} + + def block_until_ready(self): + for token in self.current_tokens.values(): + token.block_until_ready() + for token in self.output_runtime_tokens.values(): + token.block_until_ready() + self.clear() + +runtime_tokens: RuntimeTokenSet = RuntimeTokenSet() + +@atexit.register +def wait_for_tokens(): + runtime_tokens.block_until_ready() + + +@contextlib.contextmanager +def log_elapsed_time(fmt: str, fun_name: str, event: str | None = None): + if _on_exit: + yield + else: + log_priority = logging.WARNING if config.log_compiles.value else logging.DEBUG + start_time = time.time() + yield + elapsed_time = time.time() - start_time + if logger.isEnabledFor(log_priority): + logger.log(log_priority, fmt.format( + fun_name=fun_name, elapsed_time=elapsed_time)) + if event is not None: + record_event_duration_secs(event, elapsed_time) + + +def should_tuple_args(num_args: int, platform: str) -> bool: + # CPU and GPU do not need tuples as they use host-side data structures that + # do not have small bounds. + # TPU only needs a tuple for very long lists + if platform == "tpu": + return num_args > 2000 + else: + return False + +def jaxpr_has_primitive(jaxpr: core.Jaxpr, prim_name: str) -> bool: + """Whether there is a primitive given by user anywhere inside a Jaxpr.""" + for eqn in jaxpr.eqns: + if prim_name in eqn.primitive.name: + return True + for subjaxpr in core.subjaxprs(jaxpr): + if jaxpr_has_primitive(subjaxpr, prim_name): + return True + return False + + +# Use this registry with caution. It will void the guarantee that lowering to +# stablehlo is oblivious of physical devices. +prim_requires_devices_during_lowering: set[core.Primitive] = set() + +def jaxpr_has_prim_requiring_devices(jaxpr: core.Jaxpr): + for eqn in jaxpr.eqns: + if eqn.primitive in prim_requires_devices_during_lowering: + return True + for subjaxpr in core.subjaxprs(jaxpr): + if jaxpr_has_prim_requiring_devices(subjaxpr): + return True + return False + + +class SourceInfo(NamedTuple): + source_info: source_info_util.SourceInfo + eqn_name: str + + +def jaxpr_shardings( + jaxpr: core.Jaxpr, +) -> Iterator[tuple[Sharding, SourceInfo]]: + from jax._src import pjit + from jax.experimental import shard_map + + for eqn in jaxpr.eqns: + if eqn.primitive is pjit.sharding_constraint_p: + source_info = SourceInfo(eqn.source_info, eqn.primitive.name) + yield (eqn.params['sharding'], source_info) + elif eqn.primitive is pjit.pjit_p: + source_info = SourceInfo(eqn.source_info, eqn.primitive.name) + yield from ((i, source_info) for i in eqn.params['in_shardings']) + yield from ((o, source_info) for o in eqn.params['out_shardings']) + elif eqn.primitive is shard_map.shard_map_p: + source_info = SourceInfo(eqn.source_info, eqn.primitive.name) + def _names_to_pspec(names): + ndmin = max(names) + 1 if names else 0 + return PartitionSpec(*(names.get(i) for i in range(ndmin))) + yield from ((NamedSharding(eqn.params['mesh'], _names_to_pspec(names)), source_info) + for names in [*eqn.params['in_names'], *eqn.params['out_names']]) + elif eqn.primitive is device_put_p: + s = eqn.params['device'] + if isinstance(s, Sharding) and s.memory_kind is not None: + source_info = SourceInfo(eqn.source_info, eqn.primitive.name) + yield (s, source_info) + for subjaxpr in core.subjaxprs(jaxpr): + yield from jaxpr_shardings(subjaxpr) + + +def jaxpr_has_bints(jaxpr: core.Jaxpr) -> bool: + return (any(type(v.aval.dtype) is core.bint for v in jaxpr.invars + if isinstance(v.aval, core.UnshapedArray)) or + any(_is_bint_axis_size(d) + for j in itertools.chain([jaxpr], core.subjaxprs(jaxpr)) + for e in j.eqns for v in e.outvars + if isinstance(v.aval, core.DShapedArray) for d in v.aval.shape)) + +def _is_bint_axis_size(d: core.AxisSize) -> bool: + if isinstance(d, core.DArray): + assert not d.shape + return type(d.dtype) is core.bint + elif isinstance(d, core.Var): + return (isinstance(d.aval, core.DShapedArray) and + type(d.aval.dtype) is core.bint) + return False + + +# We can optionally set a Jaxpr rewriter that can be applied just before +# compilation. This mechanism is used for compiling id_tap, we can +# remove it once we bring the id_tap implementation into the core. +outfeed_rewriter: Callable[[core.Jaxpr], core.Jaxpr] | None = None +def apply_outfeed_rewriter(jaxpr: core.Jaxpr) -> core.Jaxpr: + if outfeed_rewriter is not None: + return outfeed_rewriter(jaxpr) + else: + return jaxpr + + +def check_arg(arg: Any): + if not (isinstance(arg, core.Tracer) or core.valid_jaxtype(arg)): + raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid " + "JAX type.") + + +def jaxpr_replicas(jaxpr: core.Jaxpr) -> int: + """The number of replicas needed for a jaxpr. + + For a eqn, multiply the `axis_size` with the `jaxpr_replicas` of the + subjaxprs. For a list of eqns, take the maximum number of replicas. + """ + return max(unsafe_map(_eqn_replicas, jaxpr.eqns), default=1) + +# TODO(mattjj): this function assumes that only pmap has a parameter named +# axis_size, and that it corresponds to cross-replica mapping +def _eqn_replicas(eqn: core.JaxprEqn) -> int: + call_jaxpr = eqn.params.get("call_jaxpr") + if call_jaxpr: + return eqn.params.get('axis_size', 1) * jaxpr_replicas(call_jaxpr) + elif eqn.primitive in xla.initial_style_primitives: + return _initial_style_primitive_replicas(eqn.params) + else: + return 1 + +def _initial_style_primitive_replicas(params: dict[str, Any]) -> int: + return max(core.traverse_jaxpr_params(jaxpr_replicas, params).values(), + default=1) + +def needs_check_special() -> bool: + return config.debug_infs.value or config.debug_nans.value + +def check_special(name: str, bufs: Sequence[basearray.Array]) -> None: + if needs_check_special(): + for buf in bufs: + _check_special(name, buf.dtype, buf) + +def _check_special(name: str, dtype: np.dtype, buf: basearray.Array) -> None: + if dtypes.issubdtype(dtype, np.inexact): + if config.debug_nans.value and np.any(np.isnan(np.asarray(buf))): + raise FloatingPointError(f"invalid value (nan) encountered in {name}") + if config.debug_infs.value and np.any(np.isinf(np.asarray(buf))): + raise FloatingPointError(f"invalid value (inf) encountered in {name}") + + +def _put_x(x, s: Sharding, aval: core.AbstractValue, committed: bool): + result_handler = pxla.global_aval_to_result_handler(aval, s, committed) + return result_handler(pxla.shard_arg(x, s)) + +def _override_get_device_assignment(sharding, *args, **kwargs): + da = sharding._device_assignment + return xb.get_device_backend(da[0]), da + +def _identity_fn(x): + return x + +def _mcjax_reshard(x, target_sharding): + from jax._src import api, array + + inp_sharding = x.sharding + + if inp_sharding._device_assignment == target_sharding._device_assignment: + return api.jit(_identity_fn, out_shardings=target_sharding)(x) + + if inp_sharding.device_set != target_sharding.device_set: + inp_ids = [d.id for d in inp_sharding._device_assignment] + inp_plat = inp_sharding._device_assignment[0].platform.upper() + target_ids = [d.id for d in target_sharding._device_assignment] + target_plat = target_sharding._device_assignment[0].platform.upper() + raise ValueError("Input and target sharding should have the same set of " + f"devices. Got input's device set ids: {inp_ids} on " + f"platform {inp_plat} and target sharding's device set " + f"ids: {target_ids} on platform {target_plat}") + + old_hlo_sharding = inp_sharding._to_xla_hlo_sharding(x.ndim) + if old_hlo_sharding.is_replicated(): + new_hlo_sharding = old_hlo_sharding + else: + permute_order = np.vectorize(target_sharding._device_assignment.index, + otypes=[int])(inp_sharding._device_assignment) + # Unfortunately need to fallback to V1 sharding here. + new_op_sharding = old_hlo_sharding.to_proto() + new_op_sharding.iota_reshape_dims = [] + new_op_sharding.iota_transpose_perm = [] + new_op_sharding.tile_assignment_devices = np.take( + old_hlo_sharding.tile_assignment_devices(), permute_order + ) + new_hlo_sharding = xc.HloSharding.from_proto(new_op_sharding) + + new_x = array.make_array_from_single_device_arrays( + x.shape, + GSPMDSharding(target_sharding._device_assignment, new_hlo_sharding), + x._arrays, + ) + + _orig_get_and_check_device_assignment = pxla._get_and_check_device_assignment.fn + pxla._get_and_check_device_assignment.fn = partial( + _override_get_device_assignment, target_sharding) + try: + return api.jit(_identity_fn, out_shardings=target_sharding)(new_x) + finally: + pxla._get_and_check_device_assignment.fn = _orig_get_and_check_device_assignment + + +def _device_put_sharding_impl(x, aval, device): + from jax._src import array + + if isinstance(device, Sharding): + s = device + if getattr(x, 'sharding', None) == s and getattr(x, '_committed', False): + return x + if (not s.is_fully_addressable and + isinstance(x, array.ArrayImpl) and not x.is_fully_addressable): + # This has to be XLACompatible because _mcjax_reshard will run a + # XLA computation. + assert isinstance(s, Sharding) + return _mcjax_reshard(x, s) + if not s.is_fully_addressable: + # TODO(yashkatariya,mattjj): Link to a doc about McJAX and jax.Array. + raise ValueError( + "device_put's second argument must be a Device or a Sharding which" + f" represents addressable devices, but got {s}. You are probably" + " trying to use device_put in multi-controller JAX which is not" + " supported. Please use jax.make_array_from_single_device_arrays API" + " or pass device or Sharding which represents addressable devices.") + return _put_x(x, s, aval, True) + + # Only `Device` exists below. `Sharding` instance is handled above. + if isinstance(x, array.ArrayImpl): + if not x.is_fully_addressable: + raise ValueError( + "device_put's first argument must be a fully addressable array, but " + f"got value with devices {x.devices()}") + if device is None: + return x + elif is_single_device_sharding(x.sharding): + return pxla.batched_device_put(aval, SingleDeviceSharding(device), [x], + [device]) + + sh = SingleDeviceSharding(pxla._get_default_device() + if device is None else device) + return _put_x(x, sh, aval, device is not None) + +def _device_put_impl( + x, + device: Device | Sharding | Layout | None = None, + src: Device | Sharding | Layout | None = None): + if (isinstance(device, TransferToMemoryKind) or + isinstance(src, TransferToMemoryKind)): + raise ValueError( + "TransferToMemoryKind argument to jax.device_put can only be used" + " inside jax.jit. If you are using device_put outside jax.jit, then" + " please provide a concrete Sharding with memory_kind.") + + try: + aval = xla.abstractify(x) + except TypeError as err: + raise TypeError( + f"Argument '{x}' of type {type(x)} is not a valid JAX type") from err + + if isinstance(device, Layout): + l = device + dll = l.device_local_layout + x_dll = x.layout.device_local_layout if hasattr(x, 'layout') else None + if dll is None and l.sharding is None: + return _device_put_sharding_impl(x, aval, l.sharding) + if (not isinstance(l.sharding, Sharding) or + not isinstance(dll, (DeviceLocalLayout, type(None)))): + raise ValueError( + "sharding and device_local_layout in `Layout` instance should be" + f" concrete. Got layout: {l} for input {aval.str_short()}") + if getattr(x, 'layout', None) == l and getattr(x, '_committed', False): + return x + if x_dll is None and dll is None: + return _device_put_sharding_impl(x, aval, l.sharding) + return api.jit(_identity_fn, out_shardings=l)(x) + + return _device_put_sharding_impl(x, aval, device) + + +device_put_p = core.Primitive('device_put') +device_put_p.def_impl(_device_put_impl) +device_put_p.def_abstract_eval(lambda x, device=None, src=None: x) + +def device_put_transpose_rule(ct, _, device, src): + return [device_put_p.bind(ct, device=src, src=device)] +ad.deflinear2(device_put_p, device_put_transpose_rule) +batching.defvectorized(device_put_p) + +def _tpu_gpu_device_put_lowering(ctx, x, *, device, src): + if (isinstance(device, (Sharding, TransferToMemoryKind)) and + device.memory_kind is not None): + aval, = ctx.avals_in + out_aval, = ctx.avals_out + if isinstance(device, Sharding): + x = mlir.wrap_with_sharding_op( + ctx, x, out_aval, device._to_xla_hlo_sharding(aval.ndim).to_proto()) + x = mlir.wrap_with_memory_kind(x, device.memory_kind, out_aval) + return [x] + return [x] +mlir.register_lowering( + device_put_p, _tpu_gpu_device_put_lowering, platform='tpu') +mlir.register_lowering( + device_put_p, _tpu_gpu_device_put_lowering, platform='gpu') + + +def _common_device_put_lowering(ctx, x, *, device, src): + if (isinstance(device, (Sharding, TransferToMemoryKind)) and + device.memory_kind is not None): + raise NotImplementedError( + "Passing memory_kind to device_put via Shardings is not supported on" + f" platforms {ctx.module_context.platforms}") + return [x] +mlir.register_lowering(device_put_p, _common_device_put_lowering) + +def _propagate_mem_kind_dp(xm, device=None, src=None): + if isinstance(device, (Sharding, TransferToMemoryKind)): + return device.memory_kind + return None +pxla.memory_kind_propagate_rule[device_put_p] = _propagate_mem_kind_dp diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/distributed.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..9dc5a9821cdbbf8874fce6b8e140a278df1d6135 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/distributed.py @@ -0,0 +1,208 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import atexit +from collections.abc import Sequence +import logging +import os +from typing import Any + +from jax._src import clusters +from jax._src import config +from jax._src import xla_bridge +from jax._src.lib import xla_extension + +logger = logging.getLogger(__name__) + + +class State: + process_id: int = 0 + num_processes: int = 1 + service: Any | None = None + client: Any | None = None + preemption_sync_manager: Any | None = None + coordinator_address: str | None = None + + def initialize(self, + coordinator_address: str | None = None, + num_processes: int | None = None, + process_id: int | None = None, + local_device_ids: int | Sequence[int] | None = None, + initialization_timeout: int = 300, + coordinator_bind_address: str | None = None): + coordinator_address = (coordinator_address or + os.environ.get('JAX_COORDINATOR_ADDRESS', None)) + if isinstance(local_device_ids, int): + local_device_ids = [local_device_ids] + + (coordinator_address, num_processes, process_id, local_device_ids) = ( + clusters.ClusterEnv.auto_detect_unset_distributed_params( + coordinator_address, + num_processes, + process_id, + local_device_ids, + initialization_timeout, + ) + ) + + if coordinator_address is None: + raise ValueError('coordinator_address should be defined.') + if num_processes is None: + raise ValueError('Number of processes must be defined.') + if process_id is None: + raise ValueError('The process id of the current process must be defined.') + + self.coordinator_address = coordinator_address + + # The default value of [::]:port tells the coordinator to bind to all + # available addresses on the same port as coordinator_address. + default_coordinator_bind_address = '[::]:' + coordinator_address.rsplit(':', 1)[1] + coordinator_bind_address = (coordinator_bind_address or + os.environ.get('JAX_COORDINATOR_BIND_ADDRESS', + default_coordinator_bind_address)) + if coordinator_bind_address is None: + raise ValueError('coordinator_bind_address should be defined.') + + if local_device_ids: + visible_devices = ','.join(str(x) for x in local_device_ids) + logger.info('JAX distributed initialized with visible devices: %s', visible_devices) + config.update("jax_cuda_visible_devices", visible_devices) + config.update("jax_rocm_visible_devices", visible_devices) + + self.process_id = process_id + + if process_id == 0: + if self.service is not None: + raise RuntimeError('distributed.initialize should only be called once.') + logger.info( + 'Starting JAX distributed service on %s', coordinator_bind_address + ) + self.service = xla_extension.get_distributed_runtime_service( + coordinator_bind_address, num_processes) + + self.num_processes = num_processes + + if self.client is not None: + raise RuntimeError('distributed.initialize should only be called once.') + + self.client = xla_extension.get_distributed_runtime_client( + coordinator_address, process_id, init_timeout=initialization_timeout) + logger.info('Connecting to JAX distributed service on %s', coordinator_address) + self.client.connect() + + self.initialize_preemption_sync_manager() + + def shutdown(self): + if self.client: + self.client.shutdown() + self.client = None + if self.service: + self.service.shutdown() + self.service = None + if self.preemption_sync_manager: + self.preemption_sync_manager = None + + def initialize_preemption_sync_manager(self): + if self.preemption_sync_manager is not None: + raise RuntimeError( + 'Preemption sync manager should only be initialized once.') + self.preemption_sync_manager = ( + xla_extension.create_preemption_sync_manager()) + self.preemption_sync_manager.initialize(self.client) + +global_state = State() + + +def initialize(coordinator_address: str | None = None, + num_processes: int | None = None, + process_id: int | None = None, + local_device_ids: int | Sequence[int] | None = None, + initialization_timeout: int = 300, + coordinator_bind_address: str | None = None): + """Initializes the JAX distributed system. + + Calling :func:`~jax.distributed.initialize` prepares JAX for execution on + multi-host GPU and Cloud TPU. :func:`~jax.distributed.initialize` must be + called before performing any JAX computations. + + The JAX distributed system serves a number of roles: + + * It allows JAX processes to discover each other and share topology information, + * It performs health checking, ensuring that all processes shut down if any process dies, and + * It is used for distributed checkpointing. + + If you are using TPU, Slurm, or Open MPI, all arguments are optional: if omitted, they + will be chosen automatically. + + Otherwise, you must provide the ``coordinator_address``, + ``num_processes``, and ``process_id`` arguments to :func:`~jax.distributed.initialize`. + + Args: + coordinator_address: the IP address of process `0` and a port on which that + process should launch a coordinator service. The choice of + port does not matter, so long as the port is available on the coordinator + and all processes agree on the port. + May be ``None`` only on supported environments, in which case it will be chosen automatically. + Note that special addresses like ``localhost`` or ``127.0.0.1`` usually mean that the program + will bind to a local interface and are not suitable when running in a multi-host environment. + num_processes: Number of processes. May be ``None`` only on supported environments, in + which case it will be chosen automatically. + process_id: The ID number of the current process. The ``process_id`` values across + the cluster must be a dense range ``0``, ``1``, ..., ``num_processes - 1``. + May be ``None`` only on supported environments; if ``None`` it will be chosen automatically. + local_device_ids: Restricts the visible devices of the current process to ``local_device_ids``. + If ``None``, defaults to all local devices being visible to the process except when processes + are launched via Slurm and Open MPI on GPUs. In that case, it will default to a single device per process. + initialization_timeout: Time period (in seconds) for which connection will + be retried. If the initialization takes more than the timeout specified, + the initialization will error. Defaults to 300 secs i.e. 5 mins. + coordinator_bind_address: the address and port to which the coordinator service + on process `0` should bind. If this is not specified, the default is to bind to + all available addresses on the same port as ``coordinator_address``. On systems + that have multiple network interfaces per node it may be insufficient to only + have the coordinator service listen on one address/interface. + + Raises: + RuntimeError: If :func:`~jax.distributed.initialize` is called more than once + or if called after the backend is already initialized. + + Example: + + Suppose there are two GPU processes, and process 0 is the designated coordinator + with address ``10.0.0.1:1234``. To initialize the GPU cluster, run the + following commands before anything else. + + On process 0: + + >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=0) # doctest: +SKIP + + On process 1: + + >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=1) # doctest: +SKIP + """ + if xla_bridge.backends_are_initialized(): + raise RuntimeError("jax.distributed.initialize() must be called before " + "any JAX computations are executed.") + global_state.initialize(coordinator_address, num_processes, process_id, + local_device_ids, initialization_timeout, coordinator_bind_address) + atexit.register(shutdown) + + +def shutdown(): + """Shuts down the distributed system. + + Does nothing if the distributed system is not running.""" + global_state.shutdown() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dlpack.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dlpack.py new file mode 100644 index 0000000000000000000000000000000000000000..386123ae61f0c667e22e738d5daa239aa369d8d0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dlpack.py @@ -0,0 +1,278 @@ +# Copyright 2020 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any + +from jax._src.api import device_put +from jax import numpy as jnp +from jax._src import array +from jax._src import xla_bridge +from jax._src.lax.lax import _array_copy +from jax._src.lib import xla_client +from jax._src.typing import Array, DLDeviceType +from jax._src.sharding import Sharding + +DLPACK_VERSION = (0, 8) +MIN_DLPACK_VERSION = (0, 5) + +# A set of dtypes that dlpack supports. +# Note: Make sure to use a "type", not a dtype instance, when looking up this set +# because their hashes are different. +# For example, +# hash(jnp.float32) != hash(jnp.dtype(jnp.float32)) +# hash(jnp.float32) == hash(jnp.dtype(jnp.float32).type) +# TODO(phawkins): Migrate to using dtypes instead of the scalar type objects. +SUPPORTED_DTYPES = frozenset({ + jnp.int8, jnp.int16, jnp.int32, jnp.int64, jnp.uint8, jnp.uint16, + jnp.uint32, jnp.uint64, jnp.float16, jnp.bfloat16, jnp.float32, + jnp.float64, jnp.complex64, jnp.complex128, jnp.bool_}) + + +def _to_dlpack(x: Array, stream: int | Any | None, + src_device: xla_client.Device | None = None, + device: xla_client.Device | None = None, + copy: bool | None = None): + + if src_device is None: + src_device, = x.devices() + if device and (src_device is None or device != src_device): + if copy is not None and not copy: + raise ValueError( + f"Specified {device=} which requires a copy since the source device " + f"is {repr(src_device)}, however copy=False. Set copy=True or " + "copy=None to perform the requested operation." + ) + else: + arr = device_put(x, device) + else: + arr = _array_copy(x) if copy else x + return xla_client._xla.buffer_to_dlpack_managed_tensor( + arr.addressable_data(0), stream=stream + ) + +def to_dlpack(x: Array, stream: int | Any | None = None, + src_device: xla_client.Device | None = None, + dl_device: tuple[DLDeviceType, int] | None = None, + max_version: tuple[int, int] | None = None, + copy : bool | None = None): + """Returns a DLPack tensor that encapsulates a :class:`~jax.Array` ``x``. + + Args: + x: a :class:`~jax.Array`, on either CPU or GPU. + stream: optional platform-dependent stream to wait on until the buffer is + ready. This corresponds to the `stream` argument to ``__dlpack__`` + documented in https://dmlc.github.io/dlpack/latest/python_spec.html. + src_device: either a CPU or GPU :class:`~jax.Device`. + dl_device: a tuple of ``(dl_device_type, local_hardware_id)`` in DLPack + format e.g. as produced by ``__dlpack_device__``. + max_version: the maximum DLPack version that the consumer (i.e. caller of + ``__dlpack__``) supports in the form of a 2-tuple of ``(major, minor)``. + This function is not guaranteed to return a capsule of version + ``max_version``. + copy: a boolean indicating whether or not to copy the input. If + ``copy=True`` then the function must always copy. When + ``copy=False`` then the function must never copy, and must raise an error + when a copy is deemed necessary. If ``copy=None`` then the function must + avoid a copy if possible but may copy if needed. + + Returns: + A DLPack PyCapsule object. + + Note: + While JAX arrays are always immutable, ``DLPackManagedTensor`` buffers + cannot be marked as immutable, and it is possible for processes external + to JAX to mutate them in-place. If a DLPack buffer derived from a JAX array + is mutated, it may lead to undefined behavior when using the associated JAX + array. When JAX eventually supports ``DLManagedTensorVersioned`` + (DLPack 1.0), it will be possible to specify that a buffer is read-only. + """ + if not isinstance(x, array.ArrayImpl): + raise TypeError("Argument to to_dlpack must be a jax.Array, " + f"got {type(x)}") + + device = None + dl_device_type, local_hardware_id = dl_device if dl_device else (None, None) + if dl_device_type: + try: + dl_device_platform = { + DLDeviceType.kDLCPU: "cpu", + DLDeviceType.kDLCUDA: "cuda", + DLDeviceType.kDLROCM: "rocm", + }[dl_device_type] + backend = xla_bridge.get_backend(dl_device_platform) + device = backend.device_from_local_hardware_id(local_hardware_id) + except TypeError: + # https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html + # recommends using BufferError. + raise BufferError( + "The device specification passed to to_dlpack contains an unsupported " + f"device type (DLDeviceType: {dl_device_type})") + + # As new versions are adopted over time, we can maintain some legacy paths + # for compatability mediated through the max_version parameter. + # TODO(micky774): Deprecate default usage of DLPackManagedTensor when XLA + # supports DLManagedTensorVersioned (DLPack version 1.0) and repurpose the + # current _to_dlpack as a legacy path for (0,5) <= max_version < (1,0). + if max_version is None or max_version >= DLPACK_VERSION: + # Latest + return _to_dlpack( + x, stream=stream, + src_device=src_device, + device=device, + copy=copy + ) + elif max_version >= MIN_DLPACK_VERSION: + # Oldest supported + return _to_dlpack( + x, stream=stream, + src_device=src_device, + device=device, + copy=copy + ) + else: + raise BufferError( + f"JAX does not support any version below {MIN_DLPACK_VERSION} but " + f"version ({max_version}) was requested." + ) + +def _place_array(_arr, device, dlpack_device, copy): + if device and dlpack_device != device: + if copy is not None and not copy: + raise ValueError( + f"Specified {device=} which requires a copy since the source device " + f"is {repr(dlpack_device)}, however copy=False. Set copy=True or " + "copy=None to perform the requested operation." + ) + else: + return device_put(_arr, device) + if copy: + return jnp.array(_arr, copy=True) + return _arr + +def _legacy_from_dlpack(dlpack, device: xla_client.Device | None = None, + copy: bool | None = None): + preferred_platform = getattr(device, "platform", None) + if device and preferred_platform == "gpu": + preferred_platform = "cuda" if "cuda" in device.client.platform_version else "rocm" + + cpu_backend = xla_bridge.get_backend("cpu") + gpu_backend = None + + if preferred_platform in {"cuda", "rocm"}: + try: + gpu_backend = xla_bridge.get_backend(preferred_platform) + except RuntimeError: + raise TypeError( + f"A {str.upper(preferred_platform)} device was specified, however no " + f"{str.upper(preferred_platform)} backend was found." + ) + + if preferred_platform is None: + try: + gpu_backend = xla_bridge.get_backend("cuda") + except RuntimeError: + pass + # Try ROCm if CUDA backend not found + if gpu_backend is None: + try: + gpu_backend = xla_bridge.get_backend("rocm") + except RuntimeError: + pass + + _arr = jnp.asarray(xla_client._xla.dlpack_managed_tensor_to_buffer( + dlpack, cpu_backend, gpu_backend)) + dlpack_device, = _arr.devices() + return _place_array(_arr, device, dlpack_device, copy) + +def _from_dlpack(external_array, device: xla_client.Device | None = None, + copy: bool | None = None): + dl_device_type, device_id = external_array.__dlpack_device__() + try: + dl_device_platform = { + DLDeviceType.kDLCPU: "cpu", + DLDeviceType.kDLCUDA: "cuda", + DLDeviceType.kDLROCM: "rocm", + }[dl_device_type] + except TypeError: + # https://dmlc.github.io/dlpack/latest/python_spec.html recommends using + # TypeError. + raise TypeError( + "Array passed to from_dlpack is on unsupported device type " + f"(DLDeviceType: {dl_device_type}, array: {external_array}") + + backend = xla_bridge.get_backend(dl_device_platform) + dlpack_device = backend.device_from_local_hardware_id(device_id) + try: + stream = dlpack_device.get_stream_for_external_ready_events() + except xla_client.XlaRuntimeError as err: + if "UNIMPLEMENTED" in str(err): + stream = None + else: + raise + dlpack = external_array.__dlpack__(stream=stream) + + _arr = jnp.asarray(xla_client._xla.dlpack_managed_tensor_to_buffer( + dlpack, dlpack_device, stream)) + return _place_array(_arr, device, dlpack_device, copy) + +def from_dlpack(external_array, + device: xla_client.Device | Sharding | None = None, + copy: bool | None = None): + """Returns a :class:`~jax.Array` representation of a DLPack tensor. + + The returned :class:`~jax.Array` shares memory with ``external_array`` if no + device transfer or copy was requested. + + Args: + external_array: An array object that has __dlpack__ and __dlpack_device__ + methods, or a DLPack tensor on either CPU or GPU (legacy API). + + device: The (optional) :py:class:`Device`, representing the device on which + the returned array should be placed. If given, then the result is committed + to the device. If unspecified, the resulting array will be unpacked onto the + same device it originated from. Setting ``device`` to a device different from + the source of ``external_array`` will require a copy, meaning ``copy`` must be + set to either ``True`` or ``None``. + + copy: An (optional) boolean, controlling whether or not a copy is performed. + If ``copy=True`` then a copy is always performed, even if unpacked onto the + same device. If ``copy=False`` then the copy is never performed and will raise + an error if necessary. When ``copy=None`` then a copy may be performed if + needed for a device transfer. + + Returns: + A jax.Array + + Note: + While JAX arrays are always immutable, dlpack buffers cannot be marked as + immutable, and it is possible for processes external to JAX to mutate them + in-place. If a jax Array is constructed from a dlpack buffer and the buffer + is later modified in-place, it may lead to undefined behavior when using + the associated JAX array. + """ + if isinstance(device, Sharding): + device_set = device.device_set + if len(device_set) > 1: + raise ValueError( + "from_dlpack can only unpack a dlpack tensor onto a singular device, but " + f"a Sharding with {len(device_set)} devices was provided." + ) + device, = device_set + if hasattr(external_array, "__dlpack__"): + return _from_dlpack(external_array, device, copy) + + # Legacy path + return _legacy_from_dlpack(external_array, device, copy) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dtypes.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dtypes.py new file mode 100644 index 0000000000000000000000000000000000000000..f85e4833e13c3ff8010610e184c40a519cce039d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/dtypes.py @@ -0,0 +1,795 @@ +# Copyright 2019 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Array type functions. +# +# JAX dtypes differ from NumPy in both: +# a) their type promotion rules, and +# b) the set of supported types (e.g., bfloat16), +# so we need our own implementation that deviates from NumPy in places. + +from __future__ import annotations + +import abc +import builtins +import functools +from typing import cast, overload, Any, Literal, Union +import warnings + +import ml_dtypes +import numpy as np + +from jax._src import config +from jax._src.typing import DType, DTypeLike +from jax._src.util import set_module, StrictABC + +from jax._src import traceback_util +traceback_util.register_exclusion(__file__) + +try: + _ml_dtypes_version = tuple(map(int, ml_dtypes.__version__.split('.')[:3])) +except: + pass +else: + if _ml_dtypes_version < (0, 4, 0): + raise ValueError("JAX requires ml_dtypes version 0.4.0 or newer; " + f"installed version is {ml_dtypes.__version__}.") + +export = set_module('jax.dtypes') + +@export +class extended(np.generic): + """Scalar class for extended dtypes. + + This is an abstract class that should never be instantiated, but rather + exists for the sake of `jnp.issubdtype`. + + Examples: + >>> from jax import random + >>> from jax import dtypes + >>> key = random.key(0) + >>> jnp.issubdtype(key.dtype, dtypes.extended) + True + """ + + +@export +class prng_key(extended): + """Scalar class for PRNG Key dtypes. + + This is an abstract class that should never be instantiated, but rather + exists for the sake of `jnp.issubdtype`. + + Examples: + >>> from jax import random + >>> from jax import dtypes + >>> key = random.key(0) + >>> jnp.issubdtype(key.dtype, dtypes.prng_key) + True + """ + + +class ExtendedDType(StrictABC): + """Abstract Base Class for extended dtypes""" + @property + @abc.abstractmethod + def type(self) -> type: ... + + +# fp8 support +float8_e4m3b11fnuz: type[np.generic] = ml_dtypes.float8_e4m3b11fnuz +float8_e4m3fn: type[np.generic] = ml_dtypes.float8_e4m3fn +float8_e4m3fnuz: type[np.generic] = ml_dtypes.float8_e4m3fnuz +float8_e5m2: type[np.generic] = ml_dtypes.float8_e5m2 +float8_e5m2fnuz: type[np.generic] = ml_dtypes.float8_e5m2fnuz + +_float8_e4m3b11fnuz_dtype: np.dtype = np.dtype(float8_e4m3b11fnuz) +_float8_e4m3fn_dtype: np.dtype = np.dtype(float8_e4m3fn) +_float8_e4m3fnuz_dtype: np.dtype = np.dtype(float8_e4m3fnuz) +_float8_e5m2_dtype: np.dtype = np.dtype(float8_e5m2) +_float8_e5m2fnuz_dtype: np.dtype = np.dtype(float8_e5m2fnuz) + +def supports_inf(dtype: DTypeLike) -> bool: + """Return true if the dtype supports infinity, else return False.""" + typ = np.dtype(dtype).type + if typ in {float8_e4m3b11fnuz, float8_e4m3fn, float8_e4m3fnuz, float8_e5m2fnuz}: + return False + return issubdtype(dtype, np.inexact) + +# bfloat16 support +bfloat16: type[np.generic] = ml_dtypes.bfloat16 +_bfloat16_dtype: np.dtype = np.dtype(bfloat16) + +_custom_float_scalar_types = [ + float8_e4m3b11fnuz, + float8_e4m3fn, + float8_e4m3fnuz, + float8_e5m2, + float8_e5m2fnuz, + bfloat16, +] +_custom_float_dtypes = [ + _float8_e4m3b11fnuz_dtype, + _float8_e4m3fn_dtype, + _float8_e4m3fnuz_dtype, + _float8_e5m2_dtype, + _float8_e5m2fnuz_dtype, + _bfloat16_dtype, +] +_float8_dtypes = [ + _float8_e4m3b11fnuz_dtype, + _float8_e4m3fn_dtype, + _float8_e4m3fnuz_dtype, + _float8_e5m2_dtype, + _float8_e5m2fnuz_dtype, +] + +# 4-bit integer support +int4: type[np.generic] = ml_dtypes.int4 +uint4: type[np.generic] = ml_dtypes.uint4 + +_int4_dtype: np.dtype = np.dtype(int4) +_uint4_dtype: np.dtype = np.dtype(uint4) + +_int4_dtypes = [ + _int4_dtype, + _uint4_dtype, +] + +# Default types. +bool_ = np.bool_ +int_: type[Any] +uint: type[Any] +float_: type[Any] +complex_: type[Any] +if config.default_dtype_bits.value == '32': + int_ = np.int32 + uint = np.uint32 + float_ = np.float32 + complex_ = np.complex64 +else: + int_ = np.int64 + uint = np.uint64 + float_ = np.float64 + complex_ = np.complex128 +_default_types: dict[str, type[Any]] = { + 'b': bool_, + 'i': int_, + 'u': uint, + 'f': float_, + 'c': complex_, +} + +# Trivial vectorspace datatype needed for tangent values of int/bool primals +float0: np.dtype = np.dtype([('float0', np.void, 0)]) + +_dtype_to_32bit_dtype: dict[DType, DType] = { + np.dtype('int64'): np.dtype('int32'), + np.dtype('uint64'): np.dtype('uint32'), + np.dtype('float64'): np.dtype('float32'), + np.dtype('complex128'): np.dtype('complex64'), +} + +# Note: we promote narrow types to float32 here for backward compatibility +# with earlier approaches. We might consider revisiting this, or perhaps +# tying the logic more closely to the type promotion lattice. +_dtype_to_inexact: dict[DType, DType] = { + np.dtype(k): np.dtype(v) for k, v in [ + ('bool', 'float32'), + ('uint8', 'float32'), ('int8', 'float32'), + ('uint16', 'float32'), ('int16', 'float32'), + ('uint32', 'float32'), ('int32', 'float32'), + ('uint64', 'float64'), ('int64', 'float64') + ] +} + +def to_numeric_dtype(dtype: DTypeLike) -> DType: + """Promotes a dtype into an numeric dtype, if it is not already one.""" + dtype_ = np.dtype(dtype) + return np.dtype('int32') if dtype_ == np.dtype('bool') else dtype_ + + +def to_inexact_dtype(dtype: DTypeLike) -> DType: + """Promotes a dtype into an inexact dtype, if it is not already one.""" + dtype_ = np.dtype(dtype) + return _dtype_to_inexact.get(dtype_, dtype_) + + +def to_complex_dtype(dtype: DTypeLike) -> DType: + ftype = to_inexact_dtype(dtype) + if ftype in [np.dtype('float64'), np.dtype('complex128')]: + return np.dtype('complex128') + return np.dtype('complex64') + + +@functools.cache +def _canonicalize_dtype(x64_enabled: bool, allow_extended_dtype: bool, dtype: Any) -> DType | ExtendedDType: + if issubdtype(dtype, extended): + if not allow_extended_dtype: + raise ValueError(f"Internal: canonicalize_dtype called on extended dtype {dtype} " + "with allow_extended_dtype=False") + return dtype + try: + dtype_ = np.dtype(dtype) + except TypeError as e: + raise TypeError(f'dtype {dtype!r} not understood') from e + + if x64_enabled: + return dtype_ + else: + return _dtype_to_32bit_dtype.get(dtype_, dtype_) + +@overload +def canonicalize_dtype(dtype: Any, allow_extended_dtype: Literal[False] = False) -> DType: ... + +@overload +def canonicalize_dtype(dtype: Any, allow_extended_dtype: bool = False) -> DType | ExtendedDType: ... + +@export +def canonicalize_dtype(dtype: Any, allow_extended_dtype: bool = False) -> DType | ExtendedDType: + """Convert from a dtype to a canonical dtype based on config.x64_enabled.""" + return _canonicalize_dtype(config.enable_x64.value, allow_extended_dtype, dtype) # pytype: disable=bad-return-type + +# Default dtypes corresponding to Python scalars. +python_scalar_dtypes : dict[type, DType] = { + bool: np.dtype('bool'), + int: np.dtype('int64'), + float: np.dtype('float64'), + complex: np.dtype('complex128'), +} + +@export +def scalar_type_of(x: Any) -> type: + """Return the scalar type associated with a JAX value.""" + typ = dtype(x) + if typ in _custom_float_dtypes: + return float + elif typ in _int4_dtypes: + return int + elif np.issubdtype(typ, np.bool_): + return bool + elif np.issubdtype(typ, np.integer): + return int + elif np.issubdtype(typ, np.floating): + return float + elif np.issubdtype(typ, np.complexfloating): + return complex + else: + raise TypeError(f"Invalid scalar value {x}") + + +def _scalar_type_to_dtype(typ: type, value: Any = None) -> DType: + """Return the numpy dtype for the given scalar type. + + Raises + ------ + OverflowError: if `typ` is `int` and the value is too large for int64. + + Examples + -------- + >>> _scalar_type_to_dtype(int) + dtype('int32') + >>> _scalar_type_to_dtype(float) + dtype('float32') + >>> _scalar_type_to_dtype(complex) + dtype('complex64') + >>> _scalar_type_to_dtype(int) + dtype('int32') + >>> _scalar_type_to_dtype(int, 0) + dtype('int32') + >>> _scalar_type_to_dtype(int, 1 << 63) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + OverflowError: Python int 9223372036854775808 too large to convert to int32 + """ + dtype = canonicalize_dtype(python_scalar_dtypes[typ]) + if typ is int and value is not None: + if value < np.iinfo(dtype).min or value > np.iinfo(dtype).max: + raise OverflowError(f"Python int {value} too large to convert to {dtype}") + return dtype + + +def coerce_to_array(x: Any, dtype: DTypeLike | None = None) -> np.ndarray: + """Coerces a scalar or NumPy array to an np.array. + + Handles Python scalar type promotion according to JAX's rules, not NumPy's + rules. + """ + if dtype is None and type(x) in python_scalar_dtypes: + dtype = _scalar_type_to_dtype(type(x), x) + return np.asarray(x, dtype) + +iinfo = ml_dtypes.iinfo +finfo = ml_dtypes.finfo + +def _issubclass(a: Any, b: Any) -> bool: + """Determines if ``a`` is a subclass of ``b``. + + Similar to issubclass, but returns False instead of an exception if `a` is not + a class. + """ + try: + return issubclass(a, b) + except TypeError: + return False + + +# TODO(jakevdp): consider whether to disallow None here. We allow it +# because np.issubdtype allows it (and treats it as equivalent to float64). +def issubdtype(a: DTypeLike | None, b: DTypeLike | None) -> bool: + """Returns True if first argument is a typecode lower/equal in type hierarchy. + + This is like :func:`numpy.issubdtype`, but can handle dtype extensions such as + :obj:`jax.dtypes.bfloat16` and `jax.dtypes.prng_key`. + """ + # Main departures from np.issubdtype are: + # - "extended" dtypes (like prng key types) are not normal numpy dtypes, so we + # need to handle them specifically. However, their scalar types do conform to + # the numpy scalar type hierarchy. + # - custom dtypes (like bfloat16, int4, etc.) are normal numpy dtypes, but they + # don't conform to the standard numpy type hierarchy (e.g. the bfloat16 scalar + # type is not a subclass of np.floating) so we must also handle these specially. + + # We cannot use the cached version directly for all inputs, because some may be + # unhashable (e.g. custom objects with a dtype attribute). The following check is + # fast and covers the majority of calls to this function within JAX library code. + return _issubdtype_cached( + a if isinstance(a, (type, np.dtype, ExtendedDType)) else np.dtype(a), # type: ignore[arg-type] + b if isinstance(b, (type, np.dtype, ExtendedDType)) else np.dtype(b), # type: ignore[arg-type] + ) + + +@functools.lru_cache(512) # don't use util.memoize because there is no X64 dependence. +def _issubdtype_cached(a: type | np.dtype | ExtendedDType, + b: type | np.dtype | ExtendedDType) -> bool: + # First handle extended dtypes, which require their own logic. + a_is_type = isinstance(a, type) + b_is_type = isinstance(b, type) + if b_is_type and _issubclass(b, extended): + if isinstance(a, ExtendedDType): + return _issubclass(a.type, b) + if a_is_type and _issubclass(a, np.generic): + return _issubclass(a, b) + return _issubclass(np.dtype(a).type, b) + if isinstance(b, ExtendedDType): + return isinstance(a, ExtendedDType) and a == b + if isinstance(a, ExtendedDType): + a = a.type + a_is_type = isinstance(a, type) + + # For all others, normalize inputs to scalar types. + a_sctype = a if a_is_type and _issubclass(a, np.generic) else np.dtype(a).type + b_sctype = b if b_is_type and _issubclass(b, np.generic) else np.dtype(b).type + + # Now do special handling of custom float and int types, as they don't conform + # to the normal scalar type hierarchy. + if a_sctype in _custom_float_scalar_types: + return b_sctype in {a_sctype, np.floating, np.inexact, np.number, np.generic} + if a_sctype == int4: + return b_sctype in {a_sctype, np.signedinteger, np.integer, np.number, np.generic} + if a_sctype == uint4: + return b_sctype in {a_sctype, np.unsignedinteger, np.integer, np.number, np.generic} + + # Otherwise, fall back to numpy.issubdtype + return np.issubdtype(a_sctype, b_sctype) + +can_cast = np.can_cast + +JAXType = Union[type, DType] + +# Enumeration of all valid JAX types in order. +_weak_types: list[JAXType] = [int, float, complex] +_bool_types: list[JAXType] = [np.dtype(bool)] +_signed_types: list[JAXType] +_unsigned_types: list[JAXType] +_int_types: list[JAXType] +_unsigned_types = [ + np.dtype(uint4), + np.dtype('uint8'), + np.dtype('uint16'), + np.dtype('uint32'), + np.dtype('uint64'), +] +_signed_types = [ + np.dtype(int4), + np.dtype('int8'), + np.dtype('int16'), + np.dtype('int32'), + np.dtype('int64'), +] + +_int_types = _unsigned_types + _signed_types + +_float_types: list[JAXType] = [ + *_custom_float_dtypes, + np.dtype('float16'), + np.dtype('float32'), + np.dtype('float64'), +] +_complex_types: list[JAXType] = [ + np.dtype('complex64'), + np.dtype('complex128'), +] +_jax_types = _bool_types + _int_types + _float_types + _complex_types +_jax_dtype_set = {float0, *_bool_types, *_int_types, *_float_types, *_complex_types} + + +_dtype_kinds: dict[str, set] = { + 'bool': {*_bool_types}, + 'signed integer': {*_signed_types}, + 'unsigned integer': {*_unsigned_types}, + 'integral': {*_signed_types, *_unsigned_types}, + 'real floating': {*_float_types}, + 'complex floating': {*_complex_types}, + 'numeric': {*_signed_types, *_unsigned_types, *_float_types, *_complex_types}, +} + + +def isdtype(dtype: DTypeLike, kind: str | DTypeLike | tuple[str | DTypeLike, ...]) -> bool: + """Returns a boolean indicating whether a provided dtype is of a specified kind. + + Args: + dtype : the input dtype + kind : the data type kind. + If ``kind`` is dtype-like, return ``dtype = kind``. + If ``kind`` is a string, then return True if the dtype is in the specified category: + + - ``'bool'``: ``{bool}`` + - ``'signed integer'``: ``{int4, int8, int16, int32, int64}`` + - ``'unsigned integer'``: ``{uint4, uint8, uint16, uint32, uint64}`` + - ``'integral'``: shorthand for ``('signed integer', 'unsigned integer')`` + - ``'real floating'``: ``{float8_*, float16, bfloat16, float32, float64}`` + - ``'complex floating'``: ``{complex64, complex128}`` + - ``'numeric'``: shorthand for ``('integral', 'real floating', 'complex floating')`` + + If ``kind`` is a tuple, then return True if dtype matches any entry of the tuple. + + Returns: + True or False + """ + the_dtype = np.dtype(dtype) + kind_tuple: tuple[str | DTypeLike, ...] = ( + kind if isinstance(kind, tuple) else (kind,) + ) + options: set[DType] = set() + for kind in kind_tuple: + if isinstance(kind, str) and kind in _dtype_kinds: + options.update(_dtype_kinds[kind]) + continue + try: + _dtype = np.dtype(kind) + except TypeError as e: + if isinstance(kind, str): + raise ValueError( + f"Unrecognized {kind=} expected one of {list(_dtype_kinds.keys())}, " + "or a compatible input for jnp.dtype()") + raise TypeError( + f"Expected kind to be a dtype, string, or tuple; got {kind=}" + ) from e + options.add(_dtype) + return the_dtype in options + + +def _jax_type(dtype: DType, weak_type: bool) -> JAXType: + """Return the jax type for a dtype and weak type.""" + if weak_type: + if dtype == bool: + return dtype + if dtype in _custom_float_dtypes: + return float + return type(dtype.type(0).item()) + return dtype + +def _dtype_and_weaktype(value: Any) -> tuple[DType, bool]: + """Return a (dtype, weak_type) tuple for the given input.""" + return dtype(value), any(value is typ for typ in _weak_types) or is_weakly_typed(value) + +def _type_promotion_lattice(jax_numpy_dtype_promotion: str) -> dict[JAXType, list[JAXType]]: + """ + Return the type promotion lattice in the form of a DAG. + This DAG maps each type to its immediately higher type on the lattice. + """ + b1, = _bool_types + uint4, u1, u2, u4, u8, int4, i1, i2, i4, i8 = _int_types + *f1_types, bf, f2, f4, f8 = _float_types + c4, c8 = _complex_types + i_, f_, c_ = _weak_types + if jax_numpy_dtype_promotion == 'standard': + out: dict[JAXType, list[JAXType]] + out = { + b1: [i_], + uint4: [], u1: [i2, u2], u2: [i4, u4], u4: [i8, u8], u8: [f_], + i_: [uint4, int4, u1, i1], + int4: [], i1: [i2], i2: [i4], i4: [i8], i8: [f_], + f_: [*f1_types, bf, f2, c_], + **{t: [] for t in f1_types}, bf: [f4], f2: [f4], f4: [f8, c4], f8: [c8], + c_: [c4], c4: [c8], c8: [], + } + return out + elif jax_numpy_dtype_promotion == 'strict': + return { + i_: [f_] + _int_types, + f_: [c_] + _float_types, + c_: _complex_types, + **{t: [] for t in _jax_types} + } + else: + raise ValueError( + f"Unexpected value of jax_numpy_dtype_promotion={jax_numpy_dtype_promotion!r}") + +def _make_lattice_upper_bounds(jax_numpy_dtype_promotion: str) -> dict[JAXType, set[JAXType]]: + lattice = _type_promotion_lattice(jax_numpy_dtype_promotion) + upper_bounds = {node: {node} for node in lattice} + for n in lattice: + while True: + new_upper_bounds = set().union(*(lattice[b] for b in upper_bounds[n])) + if n in new_upper_bounds: + raise ValueError(f"cycle detected in type promotion lattice for node {n}") + if new_upper_bounds.issubset(upper_bounds[n]): + break + upper_bounds[n] |= new_upper_bounds + return upper_bounds + +_lattice_upper_bounds: dict[str, dict[JAXType, set[JAXType]]] = { + 'standard': _make_lattice_upper_bounds('standard'), + 'strict': _make_lattice_upper_bounds('strict'), +} + +class TypePromotionError(ValueError): + pass + +@functools.lru_cache(512) # don't use util.memoize because there is no X64 dependence. +def _least_upper_bound(jax_numpy_dtype_promotion: str, *nodes: JAXType) -> JAXType: + """Compute the least upper bound of a set of nodes. + + Args: + nodes: sequence of entries from _jax_types + _weak_types + Returns: + the _jax_type representing the least upper bound of the input nodes + on the promotion lattice. + """ + # This function computes the least upper bound of a set of nodes N within a partially + # ordered set defined by the lattice generated above. + # Given a partially ordered set S, let the set of upper bounds of n ∈ S be + # UB(n) ≡ {m ∈ S | n ≤ m} + # Further, for a set of nodes N ⊆ S, let the set of common upper bounds be given by + # CUB(N) ≡ {a ∈ S | ∀ b ∈ N: a ∈ UB(b)} + # Then the least upper bound of N is defined as + # LUB(N) ≡ {c ∈ CUB(N) | ∀ d ∈ CUB(N), c ≤ d} + # The definition of an upper bound implies that c ≤ d if and only if d ∈ UB(c), + # so the LUB can be expressed: + # LUB(N) = {c ∈ CUB(N) | ∀ d ∈ CUB(N): d ∈ UB(c)} + # or, equivalently: + # LUB(N) = {c ∈ CUB(N) | CUB(N) ⊆ UB(c)} + # By definition, LUB(N) has a cardinality of 1 for a partially ordered set. + # Note a potential algorithmic shortcut: from the definition of CUB(N), we have + # ∀ c ∈ N: CUB(N) ⊆ UB(c) + # So if N ∩ CUB(N) is nonempty, if follows that LUB(N) = N ∩ CUB(N). + N = set(nodes) + UB = _lattice_upper_bounds[jax_numpy_dtype_promotion] + try: + bounds = [UB[n] for n in N] + except KeyError: + dtype = next(n for n in N if n not in UB) + raise ValueError(f"{dtype=} is not a valid dtype for JAX type promotion.") + CUB = set.intersection(*bounds) + LUB = (CUB & N) or {c for c in CUB if CUB.issubset(UB[c])} + if len(LUB) == 1: + return LUB.pop() + elif len(LUB) == 0: + if config.numpy_dtype_promotion.value == 'strict': + msg = ( + f"Input dtypes {tuple(str(n) for n in nodes)} have no available implicit dtype " + "promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting " + "inputs to the desired output type, or set jax_numpy_dtype_promotion=standard.") + elif any(n in _float8_dtypes for n in nodes): + msg = ( + f"Input dtypes {tuple(str(n) for n in nodes)} have no available implicit dtype " + "promotion path. To avoid unintended promotion, 8-bit floats do not support " + "implicit promotion. If you'd like your inputs to be promoted to another type, " + "you can do so explicitly using e.g. x.astype('float32')") + elif any(n in _int4_dtypes for n in nodes): + msg = ( + f"Input dtypes {tuple(str(n) for n in nodes)} have no available implicit dtype " + "promotion path. To avoid unintended promotion, 4-bit integers do not support " + "implicit promotion. If you'd like your inputs to be promoted to another type, " + "you can do so explicitly using e.g. x.astype('int32')") + else: + msg = ( + f"Input dtypes {tuple(str(n) for n in nodes)} have no available implicit dtype " + "promotion path. Try explicitly casting inputs to the desired output type.") + raise TypePromotionError(msg) + else: + # If we get here, it means the lattice is ill-formed. + raise TypePromotionError( + f"Internal Type Promotion error: {nodes} do not have a unique least upper bound " + f"on the specified lattice; options are {LUB}. This is an unexpected error in " + "JAX's internal logic; please report it to the JAX maintainers." + ) + +def promote_types(a: DTypeLike, b: DTypeLike) -> DType: + """Returns the type to which a binary operation should cast its arguments. + + For details of JAX's type promotion semantics, see :ref:`type-promotion`. + + Args: + a: a :class:`numpy.dtype` or a dtype specifier. + b: a :class:`numpy.dtype` or a dtype specifier. + + Returns: + A :class:`numpy.dtype` object. + """ + # Note: we deliberately avoid `if a in _weak_types` here because we want to check + # object identity, not object equality, due to the behavior of np.dtype.__eq__ + a_tp = cast(JAXType, a if any(a is t for t in _weak_types) else np.dtype(a)) + b_tp = cast(JAXType, b if any(b is t for t in _weak_types) else np.dtype(b)) + return np.dtype(_least_upper_bound(config.numpy_dtype_promotion.value, a_tp, b_tp)) + +def is_weakly_typed(x: Any) -> bool: + if type(x) in _weak_types: + return True + try: + return x.aval.weak_type + except AttributeError: + return False + +def is_python_scalar(x: Any) -> bool: + try: + return x.aval.weak_type and np.ndim(x) == 0 + except AttributeError: + return type(x) in python_scalar_dtypes + +def check_valid_dtype(dtype: DType) -> None: + if dtype not in _jax_dtype_set: + raise TypeError(f"Dtype {dtype} is not a valid JAX array " + "type. Only arrays of numeric types are supported by JAX.") + +def dtype(x: Any, *, canonicalize: bool = False) -> DType: + """Return the dtype object for a value or type, optionally canonicalized based on X64 mode.""" + if x is None: + raise ValueError(f"Invalid argument to dtype: {x}.") + is_type = isinstance(x, type) + if is_type and x in python_scalar_dtypes: + dt = python_scalar_dtypes[x] + elif type(x) in python_scalar_dtypes: + dt = python_scalar_dtypes[type(x)] + elif is_type and _issubclass(x, np.generic): + return np.dtype(x) + elif issubdtype(getattr(x, 'dtype', None), extended): + dt = x.dtype + else: + try: + dt = np.result_type(x) + except TypeError as err: + raise TypeError(f"Cannot determine dtype of {x}") from err + if dt not in _jax_dtype_set and not issubdtype(dt, extended): + raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array " + "type. Only arrays of numeric types are supported by JAX.") + # TODO(jakevdp): fix return type annotation and remove this ignore. + return canonicalize_dtype(dt, allow_extended_dtype=True) if canonicalize else dt # type: ignore[return-value] + +def _lattice_result_type(*args: Any) -> tuple[DType, bool]: + dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args)) + if len(dtypes) == 1: + out_dtype = dtypes[0] + out_weak_type = weak_types[0] + elif len(set(dtypes)) == 1 and not all(weak_types): + # Trivial promotion case. This allows extended dtypes through. + out_dtype = dtypes[0] + out_weak_type = False + elif all(weak_types) and config.numpy_dtype_promotion.value != 'strict': + # If all inputs are weakly typed, we compute the bound of the strongly-typed + # counterparts and apply the weak type at the end. This avoids returning the + # incorrect result with non-canonical weak types (e.g. weak int16). + # TODO(jakevdp): explore removing this special case. + result_type = _least_upper_bound(config.numpy_dtype_promotion.value, + *{_jax_type(dtype, False) for dtype in dtypes}) + out_dtype = dtype(result_type) + out_weak_type = True + else: + result_type = _least_upper_bound(config.numpy_dtype_promotion.value, + *{_jax_type(d, w) for d, w in zip(dtypes, weak_types)}) + out_dtype = dtype(result_type) + out_weak_type = any(result_type is t for t in _weak_types) + return out_dtype, (out_dtype != bool_) and out_weak_type + +@overload +def result_type(*args: Any, return_weak_type_flag: Literal[True]) -> tuple[DType, bool]: ... + +@overload +def result_type(*args: Any, return_weak_type_flag: Literal[False] = False) -> DType: ... + +@overload +def result_type(*args: Any, return_weak_type_flag: bool = False) -> DType | tuple[DType, bool]: ... + +@export +def result_type(*args: Any, return_weak_type_flag: bool = False) -> DType | tuple[DType, bool]: + """Convenience function to apply JAX argument dtype promotion. + + Args: + return_weak_type_flag : if True, then return a ``(dtype, weak_type)`` tuple. + If False, just return `dtype` + + Returns: + dtype or (dtype, weak_type) depending on the value of the ``return_weak_type`` argument. + """ + if len(args) == 0: + raise ValueError("at least one array or dtype is required") + dtype: DType | ExtendedDType + dtype, weak_type = _lattice_result_type(*(float_ if arg is None else arg for arg in args)) + if weak_type: + dtype = canonicalize_dtype( + _default_types['f' if dtype in _custom_float_dtypes else dtype.kind]) + else: + dtype = canonicalize_dtype(dtype, allow_extended_dtype=True) + # TODO(jakevdp): fix return type annotation and remove this ignore. + return (dtype, weak_type) if return_weak_type_flag else dtype # type: ignore[return-value] + +def check_user_dtype_supported(dtype, fun_name=None): + if issubdtype(dtype, extended): + return + # Avoid using `dtype in [...]` because of numpy dtype equality overloading. + if isinstance(dtype, type) and dtype in {bool, int, float, builtins.complex}: + return + np_dtype = np.dtype(dtype) + is_custom_dtype = np_dtype.type in [*_custom_float_scalar_types, int4, uint4] + if np_dtype.kind not in "biufc" and not is_custom_dtype: + msg = f"JAX only supports number and bool dtypes, got dtype {dtype}" + msg += f" in {fun_name}" if fun_name else "" + raise TypeError(msg) + if dtype is not None and np_dtype != canonicalize_dtype(dtype): + msg = ("Explicitly requested dtype {} {} is not available, " + "and will be truncated to dtype {}. To enable more dtypes, set the " + "jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell " + "environment variable. " + "See https://github.com/google/jax#current-gotchas for more.") + fun_name = f"requested in {fun_name}" if fun_name else "" + truncated_dtype = canonicalize_dtype(dtype).name + warnings.warn(msg.format(dtype, fun_name, truncated_dtype), stacklevel=3) + +def safe_to_cast(input_dtype_or_value: Any, + output_dtype_or_value: Any) -> bool: + """Check if a dtype/value is safe to cast to another dtype/value + + Args: + input_dtype_or_value: a dtype or value (to be passed to result_type) + representing the source dtype. + output_dtype_or_value: a dtype or value (to be passed to result_type) + representing the target dtype. + + Returns: + boolean representing whether the values are safe to cast according to + default type promotion semantics. + + Raises: + TypePromotionError: if the inputs have differing types and no type promotion + path under the current jax_numpy_dtype_promotion setting. + + Examples: + + >>> safe_to_cast('int32', 'float64') + True + >>> safe_to_cast('float64', 'int32') + False + >>> safe_to_cast('float32', 'complex64') + True + >>> safe_to_cast('complex64', 'float64') + False + """ + input_dtype = dtype(input_dtype_or_value, canonicalize=True) + output_dtype = dtype(output_dtype_or_value, canonicalize=True) + if input_dtype == output_dtype: + return True + # We deliberately use output_dtype rather than output_dtype_or_value here: + # this effectively treats the output dtype as always strongly-typed. + return result_type(input_dtype_or_value, output_dtype) == output_dtype diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/earray.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/earray.py new file mode 100644 index 0000000000000000000000000000000000000000..fcf0e9c3c149cb5b7d2b9a5022ec0ac0ecbe2d18 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/earray.py @@ -0,0 +1,111 @@ +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math + +from jax._src import api_util +from jax._src import basearray +from jax._src import core +from jax._src import tree_util +from jax._src import sharding_impls +from jax._src.interpreters import pxla +from jax._src.interpreters import xla +from jax._src.util import safe_zip, safe_map + +map, unsafe_map = safe_map, map +zip, unsafe_zip = safe_zip, zip + +# EArray is an Array that can contain extended dtypes. +class EArray(basearray.Array): + __slots__ = ['aval', '_data'] + __hash__ = None # type: ignore[assignment] + __array_priority__ = 100 + + def __init__(self, aval, data): + self.aval = aval + self._data = data + + def block_until_ready(self): + _ = self._data.block_until_ready() + return self + + def copy_to_host_async(self): + self._data.copy_to_host_async() + + def copy(self): + return EArray(self.aval, self._data.copy()) + + def __repr__(self): + return 'E' + repr(self._data) + + def __iter__(self): + if self.ndim == 0: raise TypeError('iteration over a 0-d array') + raise NotImplementedError + + # forward to aval + shape = property(lambda self: self.aval.shape) # type: ignore[assignment] + dtype = property(lambda self: self.aval.dtype) # type: ignore[assignment] + + # computed from shape and dtype + ndim = property(lambda self: len(self.aval.shape)) # type: ignore[assignment] + size = property(lambda self: math.prod(self.aval.shape)) # type: ignore[assignment] + itemsize = property(lambda self: self.aval.dtype.itemsize) # type: ignore[assignment] + def __len__(self): + if self.ndim == 0: raise TypeError('len() of unsized object') + return self.shape[0] + + # forward to self._data + devices = property(lambda self: self._data.devices) # type: ignore[assignment] + _committed = property(lambda self: self._data._committed) + is_fully_addressable = property(lambda self: self._data.is_fully_addressable) # type: ignore[assignment] + is_fully_replicated = property(lambda self: self._data.is_fully_replicated) # type: ignore[assignment] + delete = property(lambda self: self._data.delete) # type: ignore[assignment] + is_deleted = property(lambda self: self._data.is_deleted) # type: ignore[assignment] + on_device_size_in_bytes = property(lambda self: self._data.on_device_size_in_bytes) # type: ignore[assignment] + unsafe_buffer_pointer = property(lambda self: self._data.unsafe_buffer_pointer) # type: ignore[assignment] + + # defer to extended dtype rules + @property + def sharding(self): + phys_sharding = self._data.sharding + return sharding_impls.logical_sharding(self.aval, phys_sharding) + + # TODO(mattjj): not implemented below here, need more methods from ArrayImpl + + def addressable_data(self, index: int) -> EArray: + raise NotImplementedError + + @property + def addressable_shards(self): + raise NotImplementedError + + @property + def global_shards(self): + raise NotImplementedError + +# TODO(mattjj): _set_array_base_attributes + +def _earray_shard_arg_handler(x, sharding): + arr = x._data + phys_sharding = sharding_impls.physical_sharding(x.aval, sharding) + return pxla.shard_arg_handlers[type(arr)](arr, phys_sharding) +pxla.shard_arg_handlers[EArray] = _earray_shard_arg_handler + +api_util._shaped_abstractify_handlers[EArray] = lambda self: self.aval +core.pytype_aval_mappings[EArray] = lambda x: x.aval +xla.canonicalize_dtype_handlers[EArray] = lambda x: x +tree_util.dispatch_registry.register_node( + EArray, lambda x: ((x._data,), x.aval), lambda a, xs: EArray(a, xs[0])) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/effects.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/effects.py new file mode 100644 index 0000000000000000000000000000000000000000..36528c5feae52b3f42142b977611437cd67204a1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/effects.py @@ -0,0 +1,120 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""JAX effects. + +JAX uses effects to describe computations that may have side-effects. Effects +are associated with JAX primitive instances and Jaxprs. + +A primitive instance with an effect will be protected from dead-code elimination +even if its result is unused. + +A special class of effects are the **ordered** effects +(members of `effects.ordered_effects`). +The lowering of a computation with ordered effects will have one additional +input and one additional output for each ordered effect. These appear before +the regular inputs/outputs, and are of type `i1[0]`. These tokens +are threaded through the instructions with ordered effects to ensure that the +compiler will not eliminate, replicate, or reordered the corresponding +instructions. + +To ensure the ordering across multiple computations we maintain a +per-thread set of the tokens returned by the last dispatched computation. There +is one token per ordered effect, and it may be sharded over the devices +used by the last dispatched computation. Upon dispatching a +new computation with ordered effects we take the current token, we shard it +on the devices for the computation to be dispatched and we pass it as an input. +Then we update the current token to refer to the token output of +the dispatched computation. + +When we have ordered effects, we also use the current token to implement +`jax.barrier` which waits until the current tokens are ready. + +The implementation of `jax.barrier` for unordered effects is a bit different, +because for these effects we do not thread tokens in and out of dispatched +computation. Instead, we use a `RuntimeToken`, which is an object returned when +dispatching a computation and on which we can block until is ready. We store +for each thread the `RuntimeToken` returned by the last dispatched computation. + +For more details, see the design note: +https://jax.readthedocs.io/en/latest/jep/10657-sequencing-effects.html. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Set +from typing import Any + + +class Effect: + """A generic side-effect.""" + +Effects = Set[Effect] + +class JaxprInputEffect(Effect): + """A side-effect associated with the input of a jaxpr. + + Note that the `input_index` includes constvars. + """ + + def __init__(self, input_index: Any): + self.input_index = input_index + + def replace(self, *, input_index: Any | None = None): + if input_index is None: + input_index = self.input_index + return self.__class__(input_index) + + def __eq__(self, other): + if not isinstance(other, JaxprInputEffect): + return NotImplemented + return self.input_index == other.input_index + + def __hash__(self): + return hash((self.__class__, self.input_index)) + + def __repr__(self): + return f"{self.__class__.__name__}({self.input_index})" + +class EffectTypeSet: + + def __init__(self): + self._effect_types: set[type[Effect]] = set() + + def add_type(self, effect_type: type[Effect]): + self._effect_types.add(effect_type) + + def contains(self, eff: Effect) -> bool: + return any(isinstance(eff, eff_type) for eff_type in self._effect_types) + + def filter_in(self, effects: Iterable[Effect]) -> list[Effect]: + return [eff for eff in effects if self.contains(eff)] + + def filter_not_in(self, effects: Iterable[Effect]) -> list[Effect]: + return [eff for eff in effects if not self.contains(eff)] + + +no_effects: Effects = frozenset() +ordered_effects: EffectTypeSet = EffectTypeSet() + +# By default, ordered effects are not allowed in multi-device computations, +# because we cannot ensure a total order. Optionally, an effect can be +# declared as shardable, which means that effects will appear in program order +# but for a given program point we may see several side effects on the +# participating devices, and there is no guarantee of their relative ordering. +shardable_ordered_effects: EffectTypeSet = EffectTypeSet() + +lowerable_effects: EffectTypeSet = EffectTypeSet() +control_flow_allowed_effects: EffectTypeSet = EffectTypeSet() +custom_derivatives_allowed_effects: EffectTypeSet = EffectTypeSet() +remat_allowed_effects: EffectTypeSet = EffectTypeSet() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/environment_info.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/environment_info.py new file mode 100644 index 0000000000000000000000000000000000000000..32d34db254fd17b7bf0bd5e807e4cbb5b20ac26b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/environment_info.py @@ -0,0 +1,63 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import platform +import subprocess +import sys +import textwrap + +from jax import version +from jax._src import lib +from jax._src import xla_bridge +import numpy as np + +def try_nvidia_smi() -> str | None: + try: + return subprocess.check_output(['nvidia-smi']).decode() + except Exception: + return None + + +def print_environment_info(return_string: bool = False) -> str | None: + """Returns a string containing local environment & JAX installation information. + + This is useful information to include when asking a question or filing a bug. + + Args: return_string (bool) : if True, return the string rather than printing + to stdout. + """ + # TODO(jakevdp): should we include other info, e.g. jax.config.values? + python_version = sys.version.replace('\n', ' ') + with np.printoptions(threshold=4, edgeitems=2): + devices_short = str(np.array(xla_bridge.devices())).replace('\n', '') + info = textwrap.dedent( + f"""\ + jax: {version.__version__} + jaxlib: {lib.version_str} + numpy: {np.__version__} + python: {python_version} + jax.devices ({xla_bridge.device_count()} total, {xla_bridge.local_device_count()} local): {devices_short} + process_count: {xla_bridge.process_count()} + platform: {platform.uname()} +""" + ) + nvidia_smi = try_nvidia_smi() + if nvidia_smi: + info += '\n\n$ nvidia-smi\n' + nvidia_smi + if return_string: + return info + else: + return print(info) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/errors.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..590f68ac0b3b340d2b8edd4088293fcb666460a8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/errors.py @@ -0,0 +1,683 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from jax._src import core +from jax._src.util import set_module + +export = set_module('jax.errors') + + +class _JAXErrorMixin: + """Mixin for JAX-specific errors""" + _error_page = 'https://jax.readthedocs.io/en/latest/errors.html' + _module_name = "jax.errors" + + def __init__(self, message: str): + error_page = self._error_page + module_name = self._module_name + class_name = self.__class__.__name__ + error_msg = f'{message}\nSee {error_page}#{module_name}.{class_name}' + # https://github.com/python/mypy/issues/5887 + super().__init__(error_msg) # type: ignore + + +@export +class JAXTypeError(_JAXErrorMixin, TypeError): + pass + + +@export +class JAXIndexError(_JAXErrorMixin, IndexError): + pass + + +@export +class ConcretizationTypeError(JAXTypeError): + """ + This error occurs when a JAX Tracer object is used in a context where a + concrete value is required (see :ref:`faq-different-kinds-of-jax-values` + for more on what a Tracer is). In some situations, it can be easily fixed by + marking problematic values as static; in others, it may indicate that your + program is doing operations that are not directly supported by JAX's JIT + compilation model. + + Examples: + + Traced value where static value is expected + One common cause of this error is using a traced value where a static value + is required. For example: + + >>> from functools import partial + >>> from jax import jit + >>> import jax.numpy as jnp + >>> @jit + ... def func(x, axis): + ... return x.min(axis) + + >>> func(jnp.arange(4), 0) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ConcretizationTypeError: Abstract tracer value encountered where concrete + value is expected: axis argument to jnp.min(). + + This can often be fixed by marking the problematic argument as static:: + + >>> @partial(jit, static_argnums=1) + ... def func(x, axis): + ... return x.min(axis) + + >>> func(jnp.arange(4), 0) + Array(0, dtype=int32) + + Shape depends on Traced Value + Such an error may also arise when a shape in your JIT-compiled computation + depends on the values within a traced quantity. For example:: + + >>> @jit + ... def func(x): + ... return jnp.where(x < 0) + + >>> func(jnp.arange(4)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: + The error arose in jnp.nonzero. + + This is an example of an operation that is incompatible with JAX's JIT + compilation model, which requires array sizes to be known at compile-time. + Here the size of the returned array depends on the contents of `x`, and such + code cannot be JIT compiled. + + In many cases it is possible to work around this by modifying the logic used + in the function; for example here is code with a similar issue:: + + >>> @jit + ... def func(x): + ... indices = jnp.where(x > 1) + ... return x[indices].sum() + + >>> func(jnp.arange(4)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ConcretizationTypeError: Abstract tracer value encountered where concrete + value is expected: The error arose in jnp.nonzero. + + And here is how you might express the same operation in a way that avoids + creation of a dynamically-sized index array:: + + >>> @jit + ... def func(x): + ... return jnp.where(x > 1, x, 0).sum() + + >>> func(jnp.arange(4)) + Array(5, dtype=int32) + + To understand more subtleties having to do with tracers vs. regular values, + and concrete vs. abstract values, you may want to read + :ref:`faq-different-kinds-of-jax-values`. + """ + def __init__(self, tracer: core.Tracer, context: str = ""): + super().__init__( + "Abstract tracer value encountered where concrete value is expected: " + f"{tracer._error_repr()}\n{context}{tracer._origin_msg()}\n") + + +@export +class NonConcreteBooleanIndexError(JAXIndexError): + """ + This error occurs when a program attempts to use non-concrete boolean indices + in a traced indexing operation. Under JIT compilation, JAX arrays must have + static shapes (i.e. shapes that are known at compile-time) and so boolean + masks must be used carefully. Some logic implemented via boolean masking is + simply not possible in a :func:`jax.jit` function; in other cases, the logic + can be re-expressed in a JIT-compatible way, often using the three-argument + version of :func:`~jax.numpy.where`. + + Following are a few examples of when this error might arise. + + Constructing arrays via boolean masking + This most commonly arises when attempting to create an array via a boolean + mask within a JIT context. For example:: + + >>> import jax + >>> import jax.numpy as jnp + + >>> @jax.jit + ... def positive_values(x): + ... return x[x > 0] + + >>> positive_values(jnp.arange(-5, 5)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10]) + + This function is attempting to return only the positive values in the input + array; the size of this returned array cannot be determined at compile-time + unless `x` is marked as static, and so operations like this cannot be + performed under JIT compilation. + + Reexpressible Boolean Logic + Although creating dynamically sized arrays is not supported directly, in + many cases it is possible to re-express the logic of the computation in + terms of a JIT-compatible operation. For example, here is another function + that fails under JIT for the same reason:: + + >>> @jax.jit + ... def sum_of_positive(x): + ... return x[x > 0].sum() + + >>> sum_of_positive(jnp.arange(-5, 5)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10]) + + In this case, however, the problematic array is only an intermediate value, + and we can instead express the same logic in terms of the JIT-compatible + three-argument version of :func:`jax.numpy.where`:: + + >>> @jax.jit + ... def sum_of_positive(x): + ... return jnp.where(x > 0, x, 0).sum() + + >>> sum_of_positive(jnp.arange(-5, 5)) + Array(10, dtype=int32) + + This pattern of replacing boolean masking with three-argument + :func:`~jax.numpy.where` is a common solution to this sort of problem. + + Boolean indexing into JAX arrays + The other situation where this error often arises is when using boolean + indices, such as with :code:`.at[...].set(...)`. Here is a simple example:: + + >>> @jax.jit + ... def manual_clip(x): + ... return x.at[x < 0].set(0) + + >>> manual_clip(jnp.arange(-2, 2)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[4]) + + This function is attempting to set values smaller than zero to a scalar fill + value. As above, this can be addressed by re-expressing the logic in terms + of :func:`~jax.numpy.where`:: + + >>> @jax.jit + ... def manual_clip(x): + ... return jnp.where(x < 0, 0, x) + + >>> manual_clip(jnp.arange(-2, 2)) + Array([0, 0, 0, 1], dtype=int32) + """ + def __init__(self, tracer: core.Tracer): + super().__init__( + f"Array boolean indices must be concrete; got {tracer}\n") + + +@export +class TracerArrayConversionError(JAXTypeError): + """ + This error occurs when a program attempts to convert a JAX Tracer object into + a standard NumPy array (see :ref:`faq-different-kinds-of-jax-values` for more + on what a Tracer is). It typically occurs in one of a few situations. + + Using non-JAX functions in JAX transformations + This error can occur if you attempt to use a non-JAX library like ``numpy`` + or ``scipy`` inside a JAX transformation (:func:`~jax.jit`, :func:`~jax.grad`, + :func:`jax.vmap`, etc.). For example:: + + >>> from jax import jit + >>> import numpy as np + + >>> @jit + ... def func(x): + ... return np.sin(x) + + >>> func(np.arange(4)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerArrayConversionError: The numpy.ndarray conversion method + __array__() was called on traced array with shape int32[4] + + In this case, you can fix the issue by using :func:`jax.numpy.sin` in place of + :func:`numpy.sin`:: + + >>> import jax.numpy as jnp + >>> @jit + ... def func(x): + ... return jnp.sin(x) + + >>> func(jnp.arange(4)) + Array([0. , 0.84147096, 0.9092974 , 0.14112 ], dtype=float32) + + See also `External Callbacks`_ for options for calling back to host-side computations + from transformed JAX code. + + Indexing a numpy array with a tracer + If this error arises on a line that involves array indexing, it may be that + the array being indexed ``x`` is a standard numpy.ndarray while the indices + ``idx`` are traced JAX arrays. For example:: + + >>> x = np.arange(10) + + >>> @jit + ... def func(i): + ... return x[i] + + >>> func(0) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerArrayConversionError: The numpy.ndarray conversion method + __array__() was called on traced array with shape int32[0] + + Depending on the context, you may fix this by converting the numpy array + into a JAX array:: + + >>> @jit + ... def func(i): + ... return jnp.asarray(x)[i] + + >>> func(0) + Array(0, dtype=int32) + + or by declaring the index as a static argument:: + + >>> from functools import partial + >>> @partial(jit, static_argnums=(0,)) + ... def func(i): + ... return x[i] + + >>> func(0) + Array(0, dtype=int32) + + To understand more subtleties having to do with tracers vs. regular values, + and concrete vs. abstract values, you may want to read + :ref:`faq-different-kinds-of-jax-values`. + + .. _External Callbacks: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html + """ + def __init__(self, tracer: core.Tracer): + super().__init__( + "The numpy.ndarray conversion method __array__() was called on " + f"{tracer._error_repr()}{tracer._origin_msg()}") + + +@export +class TracerIntegerConversionError(JAXTypeError): + """ + This error can occur when a JAX Tracer object is used in a context where a + Python integer is expected (see :ref:`faq-different-kinds-of-jax-values` for + more on what a Tracer is). It typically occurs in a few situations. + + Passing a tracer in place of an integer + This error can occur if you attempt to pass a traced value to a function + that requires a static integer argument; for example:: + + >>> from jax import jit + >>> import numpy as np + + >>> @jit + ... def func(x, axis): + ... return np.split(x, 2, axis) + + >>> func(np.arange(4), 0) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerIntegerConversionError: The __index__() method was called on + traced array with shape int32[0] + + When this happens, the solution is often to mark the problematic argument as + static:: + + >>> from functools import partial + >>> @partial(jit, static_argnums=1) + ... def func(x, axis): + ... return np.split(x, 2, axis) + + >>> func(np.arange(10), 0) + [Array([0, 1, 2, 3, 4], dtype=int32), + Array([5, 6, 7, 8, 9], dtype=int32)] + + An alternative is to apply the transformation to a closure that encapsulates + the arguments to be protected, either manually as below or by using + :func:`functools.partial`:: + + >>> jit(lambda arr: np.split(arr, 2, 0))(np.arange(4)) + [Array([0, 1], dtype=int32), Array([2, 3], dtype=int32)] + + **Note a new closure is created at every invocation, which defeats the + compilation caching mechanism, which is why static_argnums is preferred.** + + Indexing a list with a Tracer + This error can occur if you attempt to index a Python list with a traced + quantity. + For example:: + + >>> import jax.numpy as jnp + >>> from jax import jit + + >>> L = [1, 2, 3] + + >>> @jit + ... def func(i): + ... return L[i] + + >>> func(0) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerIntegerConversionError: The __index__() method was called on + traced array with shape int32[0] + + Depending on the context, you can generally fix this either by converting + the list to a JAX array:: + + >>> @jit + ... def func(i): + ... return jnp.array(L)[i] + + >>> func(0) + Array(1, dtype=int32) + + or by declaring the index as a static argument:: + + >>> from functools import partial + >>> @partial(jit, static_argnums=0) + ... def func(i): + ... return L[i] + + >>> func(0) + Array(1, dtype=int32, weak_type=True) + + To understand more subtleties having to do with tracers vs. regular values, + and concrete vs. abstract values, you may want to read + :ref:`faq-different-kinds-of-jax-values`. + """ + def __init__(self, tracer: core.Tracer): + super().__init__( + f"The __index__() method was called on {tracer._error_repr()}" + f"{tracer._origin_msg()}") + + +@export +class TracerBoolConversionError(ConcretizationTypeError): + """ + This error occurs when a traced value in JAX is used in a context where a + boolean value is expected (see :ref:`faq-different-kinds-of-jax-values` + for more on what a Tracer is). + + The boolean cast may be an explicit (e.g. ``bool(x)``) or implicit, through use of + control flow (e.g. ``if x > 0`` or ``while x``), use of Python boolean + operators (e.g. ``z = x and y``, ``z = x or y``, ``z = not x``) or functions + that use them (e.g. ``z = max(x, y)``, ``z = min(x, y)`` etc.). + + In some situations, this problem can be easily fixed by marking traced values as + static; in others, it may indicate that your program is doing operations that are + not directly supported by JAX's JIT compilation model. + + Examples: + + Traced value used in control flow + One case where this often arises is when a traced value is used in + Python control flow. For example:: + + >>> from jax import jit + >>> import jax.numpy as jnp + >>> @jit + ... def func(x, y): + ... return x if x.sum() < y.sum() else y + + >>> func(jnp.ones(4), jnp.zeros(4)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerBoolConversionError: Attempted boolean conversion of JAX Tracer [...] + + We could mark both inputs ``x`` and ``y`` as static, but that would defeat + the purpose of using :func:`jax.jit` here. Another option is to re-express + the if statement in terms of the three-term :func:`jax.numpy.where`:: + + >>> @jit + ... def func(x, y): + ... return jnp.where(x.sum() < y.sum(), x, y) + + >>> func(jnp.ones(4), jnp.zeros(4)) + Array([0., 0., 0., 0.], dtype=float32) + + For more complicated control flow including loops, see + :ref:`lax-control-flow`. + + Control flow on traced values + Another common cause of this error is if you inadvertently trace over a boolean + flag. For example:: + + >>> @jit + ... def func(x, normalize=True): + ... if normalize: + ... return x / x.sum() + ... return x + + >>> func(jnp.arange(5), True) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerBoolConversionError: Attempted boolean conversion of JAX Tracer ... + + Here because the flag ``normalize`` is traced, it cannot be used in Python + control flow. In this situation, the best solution is probably to mark this + value as static:: + + >>> from functools import partial + >>> @partial(jit, static_argnames=['normalize']) + ... def func(x, normalize=True): + ... if normalize: + ... return x / x.sum() + ... return x + + >>> func(jnp.arange(5), True) + Array([0. , 0.1, 0.2, 0.3, 0.4], dtype=float32) + + For more on ``static_argnums``, see the documentation of :func:`jax.jit`. + + Using non-JAX aware functions + Another common cause of this error is using non-JAX aware functions within JAX + code. For example: + + >>> @jit + ... def func(x): + ... return min(x, 0) + + >>> func(2) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TracerBoolConversionError: Attempted boolean conversion of JAX Tracer ... + + In this case, the error occurs because Python's built-in ``min`` function is not + compatible with JAX transforms. This can be fixed by replacing it with + ``jnp.minumum``: + + >>> @jit + ... def func(x): + ... return jnp.minimum(x, 0) + + >>> print(func(2)) + 0 + + To understand more subtleties having to do with tracers vs. regular values, + and concrete vs. abstract values, you may want to read + :ref:`faq-different-kinds-of-jax-values`. + """ + def __init__(self, tracer: core.Tracer): + JAXTypeError.__init__(self, + f"Attempted boolean conversion of {tracer._error_repr()}." + f"{tracer._origin_msg()}") + + +@export +class UnexpectedTracerError(JAXTypeError): + """ + This error occurs when you use a JAX value that has leaked out of a function. + What does it mean to leak a value? If you use a JAX transformation on a + function ``f`` that stores, in some scope outside of ``f``, a reference to + an intermediate value, that value is considered to have been leaked. + Leaking values is a side effect. (Read more about avoiding side effects in + `Pure Functions `_) + + JAX detects leaks when you then use the leaked value in another + operation later on, at which point it raises an ``UnexpectedTracerError``. + To fix this, avoid side effects: if a function computes a value needed + in an outer scope, return that value from the transformed function explicitly. + + Specifically, a ``Tracer`` is JAX's internal representation of a function's + intermediate values during transformations, e.g. within :func:`~jax.jit`, + :func:`~jax.pmap`, :func:`~jax.vmap`, etc. Encountering a ``Tracer`` outside + of a transformation implies a leak. + + Life-cycle of a leaked value + Consider the following example of a transformed function which leaks a value + to an outer scope:: + + >>> from jax import jit + >>> import jax.numpy as jnp + + >>> outs = [] + >>> @jit # 1 + ... def side_effecting(x): + ... y = x + 1 # 3 + ... outs.append(y) # 4 + + >>> x = 1 + >>> side_effecting(x) # 2 + >>> outs[0] + 1 # 5 # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + UnexpectedTracerError: Encountered an unexpected tracer. + + In this example we leak a Traced value from an inner transformed scope to an + outer scope. We get an ``UnexpectedTracerError`` when the leaked value is + used, not when the value is leaked. + + This example also demonstrates the life-cycle of a leaked value: + + 1. A function is transformed (in this case, by :func:`~jax.jit`) + 2. The transformed function is called (initiating an abstract trace of the + function and turning ``x`` into a ``Tracer``) + 3. The intermediate value ``y``, which will later be leaked, is created + (an intermediate value of a traced function is also a ``Tracer``) + 4. The value is leaked (appended to a list in an outer scope, escaping + the function through a side-channel) + 5. The leaked value is used, and an UnexpectedTracerError is raised. + + The UnexpectedTracerError message tries to point to these locations in your + code by including information about each stage. Respectively: + + 1. The name of the transformed function (``side_effecting``) and which + transform kicked off the trace :func:`~jax.jit`). + 2. A reconstructed stack trace of where the leaked Tracer was created, + which includes where the transformed function was called. + (``When the Tracer was created, the final 5 stack frames were...``). + 3. From the reconstructed stack trace, the line of code that created + the leaked Tracer. + 4. The leak location is not included in the error message because it is + difficult to pin down! JAX can only tell you what the leaked value + looks like (what shape it has and where it was created) and what + boundary it was leaked over (the name of the transformation and the + name of the transformed function). + 5. The current error's stack trace points to where the value is used. + + The error can be fixed by the returning the value out of the + transformed function:: + + >>> from jax import jit + >>> import jax.numpy as jnp + + >>> outs = [] + >>> @jit + ... def not_side_effecting(x): + ... y = x+1 + ... return y + + >>> x = 1 + >>> y = not_side_effecting(x) + >>> outs.append(y) + >>> outs[0] + 1 # all good! no longer a leaked value. + Array(3, dtype=int32, weak_type=True) + + Leak checker + As discussed in point 2 and 3 above, JAX shows a reconstructed stack trace + which points to where the leaked value was created. This is because + JAX only raises an error when the leaked value is used, not when the + value is leaked. This is not the most useful place to raise this error, + because you need to know the location where the Tracer was leaked to fix the + error. + + To make this location easier to track down, you can use the leak checker. + When the leak checker is enabled, an error is raised as soon as a ``Tracer`` + is leaked. (To be more exact, it will raise an error when the transformed + function from which the ``Tracer`` is leaked returns) + + To enable the leak checker you can use the ``JAX_CHECK_TRACER_LEAKS`` + environment variable or the ``with jax.checking_leaks()`` context manager. + + .. note:: + Note that this tool is experimental and may report false positives. It + works by disabling some JAX caches, so it will have a negative effect on + performance and should only be used when debugging. + + Example usage:: + + >>> from jax import jit + >>> import jax.numpy as jnp + + >>> outs = [] + >>> @jit + ... def side_effecting(x): + ... y = x+1 + ... outs.append(y) + + >>> x = 1 + >>> with jax.checking_leaks(): + ... y = side_effecting(x) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + Exception: Leaked Trace + + """ + + def __init__(self, msg: str): + super().__init__(msg) + + +@export +class KeyReuseError(JAXTypeError): + """ + This error occurs when a PRNG key is reused in an unsafe manner. + Key reuse is checked only when `jax_debug_key_reuse` is + set to `True`. + + Here is a simple example of code that would lead to such an error:: + + >>> with jax.debug_key_reuse(True): # doctest: +SKIP + ... key = jax.random.key(0) + ... value = jax.random.uniform(key) + ... new_value = jax.random.uniform(key) + ... + --------------------------------------------------------------------------- + KeyReuseError Traceback (most recent call last) + ... + KeyReuseError: Previously-consumed key passed to jit-compiled function at index 0 + + This sort of key reuse is problematic because the JAX PRNG is stateless, and keys + must be manually split; For more information on this see `Sharp Bits: Random Numbers + `_. + """ + pass diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/flatten_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/flatten_util.py new file mode 100644 index 0000000000000000000000000000000000000000..e18ad1f6e7932da9c7c871017f21ef3eeec365e1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/flatten_util.py @@ -0,0 +1,89 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings + +import numpy as np + +from jax import lax +import jax.numpy as jnp + +from jax._src import dtypes +from jax._src.tree_util import tree_flatten, tree_unflatten +from jax._src.util import safe_zip, unzip2, HashablePartial + +zip = safe_zip + + +def ravel_pytree(pytree): + """Ravel (flatten) a pytree of arrays down to a 1D array. + + Args: + pytree: a pytree of arrays and scalars to ravel. + + Returns: + A pair where the first element is a 1D array representing the flattened and + concatenated leaf values, with dtype determined by promoting the dtypes of + leaf values, and the second element is a callable for unflattening a 1D + vector of the same length back to a pytree of the same structure as the + input ``pytree``. If the input pytree is empty (i.e. has no leaves) then as + a convention a 1D empty array of dtype float32 is returned in the first + component of the output. + + For details on dtype promotion, see + https://jax.readthedocs.io/en/latest/type_promotion.html. + + """ + leaves, treedef = tree_flatten(pytree) + flat, unravel_list = _ravel_list(leaves) + return flat, HashablePartial(unravel_pytree, treedef, unravel_list) + +def unravel_pytree(treedef, unravel_list, flat): + return tree_unflatten(treedef, unravel_list(flat)) + +def _ravel_list(lst): + if not lst: return jnp.array([], jnp.float32), lambda _: [] + from_dtypes = tuple(dtypes.dtype(l) for l in lst) + to_dtype = dtypes.result_type(*from_dtypes) + sizes, shapes = unzip2((jnp.size(x), jnp.shape(x)) for x in lst) + indices = tuple(np.cumsum(sizes)) + + if all(dt == to_dtype for dt in from_dtypes): + # Skip any dtype conversion, resulting in a dtype-polymorphic `unravel`. + # See https://github.com/google/jax/issues/7809. + del from_dtypes, to_dtype + raveled = jnp.concatenate([jnp.ravel(e) for e in lst]) + return raveled, HashablePartial(_unravel_list_single_dtype, indices, shapes) + + # When there is more than one distinct input dtype, we perform type + # conversions and produce a dtype-specific unravel function. + ravel = lambda e: jnp.ravel(lax.convert_element_type(e, to_dtype)) + raveled = jnp.concatenate([ravel(e) for e in lst]) + unrav = HashablePartial(_unravel_list, indices, shapes, from_dtypes, to_dtype) + return raveled, unrav + +def _unravel_list_single_dtype(indices, shapes, arr): + chunks = jnp.split(arr, indices[:-1]) + return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)] + +def _unravel_list(indices, shapes, from_dtypes, to_dtype, arr): + arr_dtype = dtypes.dtype(arr) + if arr_dtype != to_dtype: + raise TypeError(f"unravel function given array of dtype {arr_dtype}, " + f"but expected dtype {to_dtype}") + chunks = jnp.split(arr, indices[:-1]) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # ignore complex-to-real cast warning + return [lax.convert_element_type(chunk.reshape(shape), dtype) + for chunk, shape, dtype in zip(chunks, shapes, from_dtypes)] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/gfile_cache.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/gfile_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..301f61cc6bdbe732166d2e49cb72382737941f01 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/gfile_cache.py @@ -0,0 +1,55 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from jax._src import path as pathlib +from jax._src.compilation_cache_interface import CacheInterface + +class GFileCache(CacheInterface): + + def __init__(self, path: str): + """Sets up a cache at 'path'. Cached values may already be present.""" + self._path = pathlib.Path(path) + self._path.mkdir(parents=True, exist_ok=True) + + def get(self, key: str): + """Returns None if 'key' isn't present.""" + if not key: + raise ValueError("key cannot be empty") + path_to_key = self._path / key + if path_to_key.exists(): + return path_to_key.read_bytes() + else: + return None + + def put(self, key: str, value: bytes): + """Adds new cache entry.""" + if not key: + raise ValueError("key cannot be empty") + path_to_new_file = self._path / key + if str(path_to_new_file).startswith('gs://'): + # Writes to gcs are atomic. + path_to_new_file.write_bytes(value) + elif str(path_to_new_file).startswith('file://') or '://' not in str(path_to_new_file): + tmp_path = self._path / f"_temp_{key}" + with open(str(tmp_path), "wb") as f: + f.write(value) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, path_to_new_file) + else: + tmp_path = self._path / f"_temp_{key}" + tmp_path.write_bytes(value) + tmp_path.replace(str(path_to_new_file)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/hardware_utils.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/hardware_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3f6686d05ee8db165a343ac5092bd4458f8c0cf1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/hardware_utils.py @@ -0,0 +1,59 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import pathlib +import glob + +_GOOGLE_PCI_VENDOR_ID = '0x1ae0' +_TPU_PCI_DEVICE_IDS = [ + # TPU v2, v3 + '0x0027', + # TPU v4 + '0x005e', + # TPU v5e + '0x0063', + # Testing only + '0x0056', + '0x0062', +] + +_TPU_ENHANCED_BARRIER_SUPPORTED = [ + # TPU v2, v3 + '0x0027', + # TPU v4 + '0x005e', +] + +def num_available_tpu_chips_and_device_id(): + """Returns the device id and number of TPU chips attached through PCI.""" + num_chips = 0 + device_id = '' + for vendor_path in glob.glob('/sys/bus/pci/devices/*/vendor'): + vendor_id = pathlib.Path(vendor_path).read_text().strip() + if vendor_id != _GOOGLE_PCI_VENDOR_ID: + continue + + device_path = os.path.join(os.path.dirname(vendor_path), 'device') + device_id = pathlib.Path(device_path).read_text().strip() + if device_id in _TPU_PCI_DEVICE_IDS: + num_chips += 1 + + return num_chips, device_id + + +def tpu_enhanced_barrier_supported() -> bool: + """Returns if tpu_enhanced_barrier flag is supported on this TPU version.""" + _, device_id = num_available_tpu_chips_and_device_id() + return device_id in _TPU_ENHANCED_BARRIER_SUPPORTED diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/jaxpr_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/jaxpr_util.py new file mode 100644 index 0000000000000000000000000000000000000000..1358980d04ac34ecf0c8f1ba98cf12e608d266f1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/jaxpr_util.py @@ -0,0 +1,210 @@ +# Copyright 2020 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utilities for the Jaxpr IR.""" + +from __future__ import annotations + +from collections import Counter, defaultdict +import gzip +import itertools +import json +import types +from typing import Any, Callable, Union + +from jax._src import core +from jax._src import util +from jax._src import source_info_util +from jax._src.lib import xla_client + +map, unsafe_map = util.safe_map, map +zip, unsafe_zip = util.safe_zip, zip + + +def all_eqns(jaxpr: core.Jaxpr): + for eqn in jaxpr.eqns: + yield (jaxpr, eqn) + for subjaxpr in core.subjaxprs(jaxpr): + yield from all_eqns(subjaxpr) + +def collect_eqns(jaxpr: core.Jaxpr, key: Callable): + d = defaultdict(list) + for _, eqn in all_eqns(jaxpr): + d[key(eqn)].append(eqn) + return dict(d) + +def histogram(jaxpr: core.Jaxpr, key: Callable, + key_fmt: Callable = lambda x: x): + d = collect_eqns(jaxpr, key) + return {key_fmt(k): len(v) for k, v in d.items()} + +def primitives(jaxpr: core.Jaxpr): + return histogram(jaxpr, lambda eqn: eqn.primitive.name) + +def primitives_by_source(jaxpr: core.Jaxpr): + def key(eqn): + src = source_info_util.summarize(eqn.source_info) + return (eqn.primitive.name, src) + return histogram(jaxpr, key, ' @ '.join) + +def primitives_by_shape(jaxpr: core.Jaxpr): + def shape_fmt(var): + return '*' if isinstance(var, core.DropVar) else var.aval.str_short() + def key(eqn): + return (eqn.primitive.name, ' '.join(map(shape_fmt, eqn.outvars))) + return histogram(jaxpr, key, ' :: '.join) + +def source_locations(jaxpr: core.Jaxpr): + def key(eqn): + return source_info_util.summarize(eqn.source_info) + return histogram(jaxpr, key) + +MaybeEqn = Union[core.JaxprEqn, None] + +def var_defs_and_refs(jaxpr: core.Jaxpr): + defs: dict[core.Var, MaybeEqn] = {} + refs: dict[core.Var, list[MaybeEqn]] = {} + + def read(a: core.Atom, eqn: MaybeEqn): + if not isinstance(a, core.Literal): + assert a in defs, a + assert a in refs, a + refs[a].append(eqn) + + def write(v: core.Var, eqn: MaybeEqn): + assert v not in defs, v + assert v not in refs, v + if not isinstance(v, core.DropVar): + defs[v] = eqn + refs[v] = [] + + for v in jaxpr.constvars: + write(v, None) + for v in jaxpr.invars: + write(v, None) + + for eqn in jaxpr.eqns: + for a in eqn.invars: + read(a, eqn) + for v in eqn.outvars: + write(v, eqn) + + for a in jaxpr.outvars: + read(a, None) + + res = [(v, defs[v], refs[v]) for v in defs] + subs = map(var_defs_and_refs, core.subjaxprs(jaxpr)) + return [(jaxpr, res), *subs] if subs else (jaxpr, res) + +def vars_by_fanout(jaxpr: core.Jaxpr): + def fmt_key(var, eqn): + if eqn is None: + return f'{var} <- invar' + else: + src = source_info_util.summarize(eqn.source_info) + return f'{var} <- {eqn.primitive.name} @ {src}' + + def hist(jaxpr, reads): + return {fmt_key(var, var_def): len(var_refs) + for var, var_def, var_refs in reads} + + return [(j, hist(j, reads)) for j, reads in var_defs_and_refs(jaxpr)] # pytype: disable=bad-unpacking + +def print_histogram(histogram: dict[Any, int]): + count_width = max(len(str(v)) for v in histogram.values()) + count_fmt = '{:>' + str(count_width) + 'd}' + pairs = [(v, k) for k, v in histogram.items()] + for count, name in sorted(pairs, reverse=True): + print(count_fmt.format(count), name) + + +def _pprof_profile( + profile: dict[tuple[xla_client.Traceback | None, core.Primitive], int] +) -> bytes: + """Converts a profile into a compressed pprof protocol buffer. + + The input profile is a map from (traceback, primitive) pairs to counts. + """ + s: defaultdict[str, int] + func: defaultdict[types.CodeType, int] + loc: defaultdict[tuple[types.CodeType, int], int] + + s = defaultdict(itertools.count(1).__next__) + func = defaultdict(itertools.count(1).__next__) + loc = defaultdict(itertools.count(1).__next__) + s[""] = 0 + primitive_key = s["primitive"] + samples = [] + for (tb, primitive), count in profile.items(): + if tb is None: + frames = [] + else: + raw_frames = zip(*tb.raw_frames()) + frames = [loc[(code, lasti)] for code, lasti in raw_frames + if source_info_util.is_user_filename(code.co_filename)] + samples.append({ + "location_id": frames, + "value": [count], + "label": [{ + "key": primitive_key, + "str": s[primitive.name] + }] + }) + + locations = [ + {"id": loc_id, + "line": [{"function_id": func[code], + "line": xla_client.Traceback.code_addr2line(code, lasti)}]} + for (code, lasti), loc_id in loc.items() + ] + functions = [ + {"id": func_id, + "name": s[code.co_name], + "system_name": s[code.co_name], + "filename": s[code.co_filename], + "start_line": code.co_firstlineno} + for code, func_id in func.items() + ] + sample_type = [{"type": s["equations"], "unit": s["count"]}] + # This is the JSON encoding of a pprof profile protocol buffer. See: + # https://github.com/google/pprof/blob/master/proto/profile.proto for a + # description of the format. + json_profile = json.dumps({ + "string_table": list(s.keys()), + "location": locations, + "function": functions, + "sample_type": sample_type, + "sample": samples, + }) + return gzip.compress(xla_client._xla.json_to_pprof_profile(json_profile)) + + +def pprof_equation_profile(jaxpr: core.Jaxpr) -> bytes: + """Generates a pprof profile that maps jaxpr equations to Python stack traces. + + By visualizing the profile using pprof, one can identify Python code that is + responsible for yielding large numbers of jaxpr equations. + + Args: + jaxpr: a Jaxpr. + + Returns: + A gzip-compressed pprof Profile protocol buffer, suitable for passing to + pprof tool for visualization. + """ + d = Counter( + (eqn.source_info.traceback, eqn.primitive) + for _, eqn in all_eqns(jaxpr) + ) + return _pprof_profile(d) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/lax_reference.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/lax_reference.py new file mode 100644 index 0000000000000000000000000000000000000000..81209f6ed34a3e5bb07129d0406e0466132d2c2f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/lax_reference.py @@ -0,0 +1,498 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import builtins +import collections +import itertools + +import numpy as np +import opt_einsum +import scipy.special + +from jax._src import dtypes +from jax._src import util + +_slice = builtins.slice +_max = builtins.max +_min = builtins.min +_map = builtins.map + +neg = np.negative +sign = np.sign +floor = np.floor +ceil = np.ceil + +def round(x): + return np.trunc( + x + np.copysign(np.nextafter(np.array(.5, dtype=x.dtype), + np.array(0., dtype=x.dtype), + dtype=x.dtype), x)).astype(x.dtype) + +nextafter = np.nextafter + +is_finite = np.isfinite + +exp = np.exp +exp2 = np.exp2 +expm1 = np.expm1 +log = np.log +log1p = np.log1p +tanh = np.tanh +sin = np.sin +cos = np.cos +atan2 = np.arctan2 + +sqrt = np.sqrt +rsqrt = lambda x: np.ones_like(x) / np.sqrt(x) +cbrt = np.cbrt + +square = np.square +reciprocal = np.reciprocal +tan = np.tan +asin = np.arcsin +acos = np.arccos +atan = np.arctan +sinh = np.sinh +cosh = np.cosh +asinh = np.arcsinh +acosh = np.arccosh +atanh = np.arctanh + +def logistic(x): return (1 / (1 + np.exp(-x))).astype(x.dtype) +def betainc(a, b, x): return scipy.special.betainc(a, b, x).astype(x.dtype) +def lgamma(x): return scipy.special.gammaln(x).astype(x.dtype) +def digamma(x): return scipy.special.digamma(x).astype(x.dtype) +igamma = scipy.special.gammainc +igammac = scipy.special.gammaincc +def erf(x): return scipy.special.erf(x).astype(x.dtype) +def erfc(x): return scipy.special.erfc(x).astype(x.dtype) +def erf_inv(x): return scipy.special.erfinv(x).astype(x.dtype) + +def bessel_i0e(x): return scipy.special.i0e(x).astype(x.dtype) +def bessel_i1e(x): return scipy.special.i1e(x).astype(x.dtype) + +real = np.real +imag = np.imag + +def conj(x): + return np.conj(x) + np.complex64(0) + +def complex(x, y): + return x + np.complex64(1j) * y + +abs = np.absolute +pow = np.power + +bitwise_not = np.bitwise_not +bitwise_and = np.bitwise_and +bitwise_or = np.bitwise_or +bitwise_xor = np.bitwise_xor + +add = np.add +sub = np.subtract +mul = np.multiply + +def div(lhs, rhs): + if dtypes.issubdtype(dtypes.result_type(lhs), np.integer): + quotient = np.floor_divide(lhs, rhs) + select = np.logical_and(np.sign(lhs) != np.sign(rhs), + np.remainder(lhs, rhs) != 0) + return np.where(select, quotient + 1, quotient) + else: + return np.divide(lhs, rhs) + +def rem(lhs, rhs): + return np.sign(lhs) * np.remainder(np.abs(lhs), np.abs(rhs)) + +max = np.maximum +min = np.minimum + +shift_left = np.left_shift +shift_right_arithmetic = np.right_shift +# TODO shift_right_logical + +def population_count(x): + assert np.issubdtype(x.dtype, np.integer) + dtype = x.dtype + iinfo = np.iinfo(x.dtype) + if np.iinfo(x.dtype).bits < 32: + assert iinfo.kind in ('i', 'u') + x = x.astype(np.uint32 if iinfo.kind == 'u' else np.int32) + if iinfo.kind == 'i': + x = x.view(f"uint{np.iinfo(x.dtype).bits}") + assert x.dtype in (np.uint32, np.uint64) + m = [ + np.uint64(0x5555555555555555), # binary: 0101... + np.uint64(0x3333333333333333), # binary: 00110011.. + np.uint64(0x0f0f0f0f0f0f0f0f), # binary: 4 zeros, 4 ones ... + np.uint64(0x00ff00ff00ff00ff), # binary: 8 zeros, 8 ones ... + np.uint64(0x0000ffff0000ffff), # binary: 16 zeros, 16 ones ... + np.uint64(0x00000000ffffffff), # binary: 32 zeros, 32 ones + ] + + if x.dtype == np.uint32: + m = list(map(np.uint32, m[:-1])) + + x = (x & m[0]) + ((x >> 1) & m[0]) # put count of each 2 bits into those 2 bits + x = (x & m[1]) + ((x >> 2) & m[1]) # put count of each 4 bits into those 4 bits + x = (x & m[2]) + ((x >> 4) & m[2]) # put count of each 8 bits into those 8 bits + x = (x & m[3]) + ((x >> 8) & m[3]) # put count of each 16 bits into those 16 bits + x = (x & m[4]) + ((x >> 16) & m[4]) # put count of each 32 bits into those 32 bits + if x.dtype == np.uint64: + x = (x & m[5]) + ((x >> 32) & m[5]) # put count of each 64 bits into those 64 bits + return x.astype(dtype) + +def clz(x): + assert np.issubdtype(x.dtype, np.integer) + nbits = np.iinfo(x.dtype).bits + mask = (2 ** np.arange(nbits, dtype=x.dtype))[::-1] + bits = (x[..., None] & mask).astype(np.bool_) + out = np.argmax(bits, axis=-1).astype(x.dtype) + out[x == 0] = nbits + return out + +eq = np.equal +ne = np.not_equal +ge = np.greater_equal +gt = np.greater +le = np.less_equal +lt = np.less + +def convert_element_type(operand, dtype): + return np.asarray(operand, dtype=dtype) + +def bitcast_convert_type(operand, dtype): + return np.asarray(operand).view(dtype) + +def clamp(min, operand, max): + return np.clip(operand, np.clip(min, None, max), max).astype(operand.dtype) + +def concatenate(operands, dimension): + return np.concatenate(operands, axis=dimension) + +def conv(lhs, rhs, window_strides, padding): + pads = padtype_to_pads(lhs.shape[2:], rhs.shape[2:], window_strides, padding) + return _conv(lhs, rhs, window_strides, pads) + +def conv_with_general_padding( + lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation): + return _conv(_dilate(lhs, lhs_dilation), _dilate(rhs, rhs_dilation), + window_strides, padding) + +def conv_general_dilated(lhs, rhs, window_strides, padding, lhs_dilation, + rhs_dilation, dimension_numbers): + lhs_perm, rhs_perm, out_perm = _conv_general_permutations(dimension_numbers) + if isinstance(padding, str): + padding = padtype_to_pads(np.take(lhs.shape, lhs_perm)[2:], + np.take(rhs.shape, rhs_perm)[2:], + window_strides, padding) + trans_lhs = transpose(lhs, lhs_perm) + trans_rhs = transpose(rhs, rhs_perm) + out = conv_with_general_padding(trans_lhs, trans_rhs, window_strides, padding, + lhs_dilation, rhs_dilation) + return transpose(out, np.argsort(out_perm)) + +dot = np.dot + +def dot_general(lhs, rhs, dimension_numbers): + (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers + new_id = itertools.count() + lhs_axis_ids = [next(new_id) for _ in lhs.shape] + rhs_axis_ids = [next(new_id) for _ in rhs.shape] + lhs_out_axis_ids = lhs_axis_ids[:] + rhs_out_axis_ids = rhs_axis_ids[:] + + for lhs_axis, rhs_axis in zip(lhs_contracting, rhs_contracting): + shared_id = next(new_id) + lhs_axis_ids[lhs_axis] = shared_id + rhs_axis_ids[rhs_axis] = shared_id + lhs_out_axis_ids[lhs_axis] = None + rhs_out_axis_ids[rhs_axis] = None + + batch_ids = [] + for lhs_axis, rhs_axis in zip(lhs_batch, rhs_batch): + shared_id = next(new_id) + lhs_axis_ids[lhs_axis] = shared_id + rhs_axis_ids[rhs_axis] = shared_id + lhs_out_axis_ids[lhs_axis] = None + rhs_out_axis_ids[rhs_axis] = None + batch_ids.append(shared_id) + + not_none = lambda x: x is not None + out_axis_ids = filter(not_none, + batch_ids + lhs_out_axis_ids + rhs_out_axis_ids) + assert lhs.dtype == rhs.dtype + dtype = np.float32 if lhs.dtype == dtypes.bfloat16 else None + out = np.einsum(lhs, lhs_axis_ids, rhs, rhs_axis_ids, out_axis_ids, + dtype=dtype) + return out.astype(dtypes.bfloat16) if lhs.dtype == dtypes.bfloat16 else out + +def ragged_dot( + lhs, + rhs, + group_sizes, +): + """Reference ragged dot implementation.""" + m, lk = lhs.shape + group_count, rk, n = rhs.shape + assert lk == rk + assert group_count == group_sizes.shape[0] + assert lhs.dtype == rhs.dtype + + out = np.zeros((m, n), dtype=lhs.dtype) + result_iota = np.expand_dims(np.arange(out.shape[0]), list(range(1, out.ndim))) + start = 0 + for i, size in enumerate(group_sizes): + out += np.where( + np.logical_and(start <= result_iota, result_iota < (start + size)), + np.einsum( + "nk,km->nm", + lhs, + rhs[i, :, :], + dtype=np.float32 if lhs.dtype == dtypes.bfloat16 else None, + ), + np.zeros(out.shape, dtype=out.dtype), + ) + start += size + return out.astype(dtypes.bfloat16) if lhs.dtype == dtypes.bfloat16 else out + +def broadcast(operand, sizes): + return np.broadcast_to(operand, sizes + np.shape(operand)) + +def broadcast_in_dim(operand, shape, broadcast_dimensions): + in_reshape = np.ones(len(shape), dtype=np.int32) + for i, bd in enumerate(broadcast_dimensions): + in_reshape[bd] = operand.shape[i] + return np.broadcast_to(np.reshape(operand, in_reshape), shape) + +sum = np.sum + +squeeze = np.squeeze + +def reshape(operand, new_sizes, dimensions=None): + if dimensions is None: + dimensions = range(len(np.shape(operand))) + return np.reshape(np.transpose(operand, dimensions), new_sizes) + +def pad(operand, padding_value, padding_config): + # https://www.tensorflow.org/xla/operation_semantics#pad + lo, hi, interior = util.unzip3(padding_config) + # Handle first the positive edge padding and interior + lo_pos, hi_pos = np.clip(lo, 0, None), np.clip(hi, 0, None) + outshape = np.add(np.add(np.add(lo_pos, hi_pos), operand.shape), + np.multiply(interior, np.subtract(operand.shape, 1))) + out = np.full(outshape, padding_value, operand.dtype) + lhs_slices = tuple(_slice(l if l > 0 else 0, -h if h > 0 else None, step) + for l, h, step in zip(lo_pos, hi_pos, np.add(1, interior))) + out[lhs_slices] = operand + trim_slices = tuple(_slice(-l if l < 0 else 0, h if h < 0 else None) + for l, h in zip(lo, hi)) + return out[trim_slices] + +def rev(operand, dimensions): + dimensions = frozenset(dimensions) + indexer = (_slice(None, None, -1) if d in dimensions else _slice(None) + for d in range(np.ndim(operand))) + return operand[tuple(indexer)] + +select = np.where + +def slice(operand, start_indices, limit_indices, strides=None): # pylint: disable=redefined-builtin + if strides is None: + strides = np.ones(len(start_indices)).astype(int) + slices = tuple(_map(_slice, start_indices, limit_indices, strides)) + return operand[slices] + +def dynamic_slice(operand, start_indices, slice_sizes): + out = np.zeros(slice_sizes, dtype=operand.dtype) + idx = tuple(_slice(start, start+size) + for start, size in zip(start_indices, slice_sizes)) + section = operand[idx] + out[tuple(_slice(None, stop) for stop in section.shape)] = section + return out + +def dynamic_update_slice(operand, update, start_indices): + slices = tuple(_map(_slice, start_indices, np.add(start_indices, update.shape))) + updated_operand = np.copy(operand) + updated_operand[slices] = update + return updated_operand + +transpose = np.transpose + +def reduce(operand, init_value, computation, dimensions): # pylint: disable=redefined-builtin + reducer = _make_reducer(computation, init_value) + return reducer(operand, tuple(dimensions)).astype(np.asarray(operand).dtype) + +def reduce_window(operand, init_value, computation, window_dimensions, + window_strides, padding, base_dilation): + op, dims, strides = operand, window_dimensions, window_strides + if isinstance(padding, str): + pads = padtype_to_pads(op.shape, dims, strides, padding) + else: + pads = padding + op = op.reshape((1, 1) + op.shape) + if base_dilation: + op = _dilate(op, base_dilation, init_value) + view = _conv_view(op, (1, 1) + dims, strides, pads, + pad_value=init_value)[0] + view = view.reshape(view.shape[1:1+len(dims)] + (-1,)) + reducer = _make_reducer(computation, init_value) + return reducer(view, axis=-1) + +# TODO(mattjj): select_and_scatter + +sort = np.sort + +def sort_key_val(keys, values, dimension=-1): + idxs = list(np.ix_(*[np.arange(d) for d in keys.shape])) + idxs[dimension] = np.argsort(keys, axis=dimension) + return keys[tuple(idxs)], values[tuple(idxs)] + +### conv util + +def _conv(lhs, rhs, window_strides, pads): + view, view_axes, rhs_axes, out_axes = _conv_view( + lhs, rhs.shape, window_strides, pads, 0.) + return opt_einsum.contract( + view, view_axes, rhs, rhs_axes, out_axes, use_blas=True) + +def padtype_to_pads(in_shape, filter_shape, window_strides, padding): + if padding.upper() == 'SAME' or padding.upper() == 'SAME_LOWER': + out_shape = np.ceil(np.true_divide(in_shape, window_strides)).astype(int) + pad_sizes = [_max((out_size - 1) * stride + filter_size - in_size, 0) + for out_size, stride, filter_size, in_size + in zip(out_shape, window_strides, filter_shape, in_shape)] + if padding.upper() == 'SAME': + return [ + (pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes + ] + else: + return [ + (pad_size - pad_size // 2, pad_size // 2) for pad_size in pad_sizes + ] + else: + return [(0, 0)] * len(in_shape) + +def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value): + """Compute the view (and its axes) of a convolution or window reduction.""" + if (_min(lhs.ndim, len(rhs_shape)) < 2 or lhs.ndim != len(rhs_shape) + or lhs.shape[1] != rhs_shape[1]): + raise ValueError('Dimension mismatch') + if len(window_strides) != len(rhs_shape) - 2: + raise ValueError('Wrong number of strides for spatial dimensions') + if len(pads) != len(rhs_shape) - 2: + raise ValueError('Wrong number of pads for spatial dimensions') + + lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value) + in_shape = lhs.shape[2:] + filter_shape = rhs_shape[2:] + dim = len(filter_shape) # number of 'spatial' dimensions in convolution + + out_strides = np.multiply(window_strides, lhs.strides[2:]) + view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:] + + out_shape = np.floor_divide( + np.subtract(in_shape, filter_shape), window_strides) + 1 + view_shape = lhs.shape[:1] + tuple(out_shape) + rhs_shape[1:] + + view = np.lib.stride_tricks.as_strided(lhs, view_shape, view_strides) + + view_axes = list(range(view.ndim)) + sum_axes = view_axes[-dim-1:] + rhs_axes = [view.ndim] + sum_axes + out_axes = [0, view.ndim] + list(range(1, dim+1)) + + return view, view_axes, rhs_axes, out_axes + +def _pad(arr, pads, pad_value): + out = np.pad(arr, np.maximum(0, pads), mode='constant', + constant_values=pad_value).astype(arr.dtype) + slices = tuple(_slice(abs(lo) if lo < 0 else 0, hi % dim if hi < 0 else None) + for (lo, hi), dim in zip(pads, np.shape(arr))) + return out[slices] + +def _dilate(operand, factors, fill_value=0): + # this logic is like lax.pad, but with two leading dimensions, no edge + # padding, and factors are at least 1 (interior padding is at least 0) + outspace = np.add(operand.shape[2:], + np.multiply(np.subtract(factors, 1), + np.subtract(operand.shape[2:], 1))) + out = np.full(operand.shape[:2] + tuple(outspace), fill_value, operand.dtype) + lhs_slices = tuple(_slice(None, None, step) for step in factors) + out[(_slice(None),) * 2 + lhs_slices] = operand + return out + +def _conv_general_permutations(dimension_numbers): + lhs_spec, rhs_spec, out_spec = dimension_numbers + rhs_perm = ((rhs_spec.index('O'), rhs_spec.index('I')) + + tuple(i for i, c in enumerate(rhs_spec) if c not in {'O', 'I'})) + lhs_perm = ((lhs_spec.index('N'), lhs_spec.index('C')) + + tuple(sorted((i for i, c in enumerate(lhs_spec) + if c not in {'N', 'C'}), + key=lambda i: rhs_spec.index(lhs_spec[i])))) + out_perm = ((out_spec.index('N'), out_spec.index('C')) + + tuple(sorted((i for i, c in enumerate(out_spec) + if c not in {'N', 'C'}), + key=lambda i: rhs_spec.index(out_spec[i])))) + return lhs_perm, rhs_perm, out_perm + +### reduce util + +def _make_reducer(py_binop, init_val): + """Make a reducer function given a Python binop and an initial value.""" + # It's tempting to use np.ufunc.reduce (even with a ufunc generated by + # np.frompyfunc(py_binop)), but this may not agree with custom init_val. + # We make an attempt to uncover an underlying numpy ufunc (which might be + # wrapped by autograd or lax) and check its identity against init_val. + monoid_record = _monoids.get(getattr(py_binop, '__name__')) + if monoid_record: + reducer, monoid_identity = monoid_record + if init_val == monoid_identity(dtypes.result_type(init_val)): + return reducer + return _reducer_from_pyfunc(py_binop, init_val) + +def _get_max_identity(dt): + return -np.inf if dtypes.issubdtype(dt, np.floating) else np.iinfo(dt).min + +def _get_min_identity(dt): + return np.inf if dtypes.issubdtype(dt, np.floating) else np.iinfo(dt).max + +def _identity_getter(op): + return lambda dtype: np.asarray(op.identity, dtype=dtype) + +MonoidRecord = collections.namedtuple('MonoidRecord', ['reducer', 'identity']) +_monoids = { + 'max': MonoidRecord(np.maximum.reduce, _get_max_identity), + 'min': MonoidRecord(np.minimum.reduce, _get_min_identity), + 'add': MonoidRecord(np.add.reduce, _identity_getter(np.add)), + 'mul': MonoidRecord(np.multiply.reduce, _identity_getter(np.multiply)), + 'multiply': MonoidRecord(np.multiply.reduce, + _identity_getter(np.multiply)), + 'logical_and': MonoidRecord(np.logical_and.reduce, + _identity_getter(np.logical_and)), + 'logical_or': MonoidRecord(np.logical_or.reduce, + _identity_getter(np.logical_or)), +} + +def _reducer_from_pyfunc(py_binop, init_val): + def reducer(operand, axis=0): + axis = range(np.ndim(operand)) if axis is None else axis + result = np.full(np.delete(np.shape(operand), axis), init_val, + dtype=np.asarray(operand).dtype) + for idx, _ in np.ndenumerate(operand): + out_idx = tuple(np.delete(idx, axis)) + result[out_idx] = py_binop(result[out_idx], operand[idx]) + return result + return reducer diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/layout.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/layout.py new file mode 100644 index 0000000000000000000000000000000000000000..2071794a09fbae2291140a9a8b8604b4d8d61e4c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/layout.py @@ -0,0 +1,102 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the ific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Union + +from jax._src.sharding import Sharding +from jax._src.sharding_impls import AUTO as AutoSharding, is_auto +from jax._src.lib import xla_client as xc + + +class AutoLayout: + + def __repr__(self): + return "AUTO" + + +class DeviceLocalLayout: + layout: xc.PjRtLayout + + AUTO = AutoLayout() + + def __init__(self, layout: xc.PjRtLayout): + self._layout = layout + self._layout_str = str(self._layout) + + def __repr__(self): + return f'DeviceLocalLayout({self._layout_str})' + + def __hash__(self): + return hash(self._layout) + + def __eq__(self, other): + if not isinstance(other, DeviceLocalLayout): + return False + return self._layout == other._layout + + def _to_xla_layout(self) -> str: + return self._layout_str + + +LayoutOptions = Union[DeviceLocalLayout, None, AutoLayout] +ShardingOptions = Union[Sharding, None, AutoSharding] + + +class Layout: + __slots__ = ['device_local_layout', 'sharding'] + + def __init__(self, device_local_layout: LayoutOptions = None, + sharding: ShardingOptions = None): + # If layout is concrete and sharding is not, error. + if (isinstance(device_local_layout, DeviceLocalLayout) and + (sharding is None or is_auto(sharding))): + raise ValueError( + 'Sharding has to be concrete when layout is of type' + f' {type(device_local_layout)}. Please pass a' + ' `jax.sharding.NamedSharding`, `jax.sharding.PositionalSharding` or' + ' `jax.sharding.SingleDeviceSharding` to the sharding argument. Got' + f' sharding {sharding}' + ) + if not isinstance( + device_local_layout, (DeviceLocalLayout, type(None), AutoLayout)): + raise TypeError( + 'Invalid value received for the device_local_layout argument.' + ' Expected values are `None`, `DeviceLocalLayout.AUTO` or an' + f' instance of `DeviceLocalLayout`. Got {device_local_layout} of' + f' type {type(device_local_layout)}' + ) + if not isinstance( + sharding, (Sharding, type(None), AutoSharding)): + raise TypeError( + 'Invalid value received for the sharding argument. Expected values' + ' are `None`, `pjit.AUTO` or an instance of `jax.Sharding`. Got' + f' {sharding} of type {type(sharding)}') + + self.device_local_layout = device_local_layout + self.sharding = sharding + + def __repr__(self): + return (f'Layout(device_local_layout={self.device_local_layout},' + f' sharding={self.sharding})') + + def __hash__(self): + return hash((self.device_local_layout, self.sharding)) + + def __eq__(self, other): + if not isinstance(other, Layout): + return False + return (self.device_local_layout == other.device_local_layout and + self.sharding == other.sharding) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/lazy_loader.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/lazy_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..6041c77c65c0d594b8cf91d9d8c6e80067f5f850 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/lazy_loader.py @@ -0,0 +1,45 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A LazyLoader class.""" + +from collections.abc import Sequence +import importlib +from typing import Any, Callable + + +def attach(package_name: str, submodules: Sequence[str]) -> tuple[ + Callable[[str], Any], + Callable[[], list[str]], + list[str], +]: + """Lazily loads submodules of a package. + + Example use: + ``` + __getattr__, __dir__, __all__ = lazy_loader.attach(__name__, ["sub1", "sub2"]) + ``` + """ + + __all__: list[str] = list(submodules) + + def __getattr__(name: str) -> Any: + if name in submodules: + return importlib.import_module(f"{package_name}.{name}") + raise AttributeError(f"module '{package_name}' has no attribute '{name}") + + def __dir__() -> list[str]: + return __all__ + + return __getattr__, __dir__, __all__ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/linear_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/linear_util.py new file mode 100644 index 0000000000000000000000000000000000000000..8f431d160153b008572d28a8a24bd0ddf8fe557d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/linear_util.py @@ -0,0 +1,405 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utilities for defining functions composed with transformations. + +For example, + + from jax._src import linear_util as lu + + wf = lu.wrap_init(f) # Produce a WrappedFun for applying transformations on `f` + +A `WrappedFun` object represents a function `f`, together with a sequence of +nested transformations that are to be applied to the positional and keyword +arguments at call time and function return values at return time. +A transformation can take some static positional arguments that are given +at the wrapping time, and may also return some auxiliary output: + + wf, aux_out_thunk = trans1(wf, static_arg) + +We can call the transformed function. First, the transformation is applied +to the dynamic args and keyword args to produce new dynamic and keyword args. +Then the underlying function is called and the transformation is applied to +the results. +If there are multiple transformations, they form a stack. The arguments are +transformed first with the last applied transformation; the results are +transformed first with the first applied transformation. + + res = wf.call_wrapped(dynamic_args, kwargs) + # Now `aux_out_thunk()` is the auxiliary output. + +A transformation is written as a generator function that takes zero or more +static positional arguments (given when the transformation is instantiated), +along with positional and keyword arguments to be transformed. +The generator will yield twice: + + @lu.transformation_with_aux + def trans1(static_arg, *dynamic_args, **kwargs): + ... + # First yield: pair of transformed (args, kwargs). Get back the results. + results = yield (new_dynamic_args, new_kwargs) + ... + # Second yield: pair of (transformed results, and auxiliary output) + yield new_results, auxiliary_output + + +`WrappedFun` objects explicitly represent the set of transformations so that +they can be used as dictionary keys for memoization. `WrappedFun` objects +compare as equal only if they compute the same function. The static and the +dynamic positional arguments for the generators, and also the auxiliary output +data must be immutable, because it will be stored in function memoization tables. +""" +from __future__ import annotations + +from functools import partial +from typing import Any, Callable, NamedTuple +import weakref + +from jax._src import config +from jax._src import core +from jax._src import traceback_util +from jax._src.tree_util import tree_map +from jax._src.util import curry + + +traceback_util.register_exclusion(__file__) + + +class StoreException(Exception): pass + + +class EmptyStoreValue: pass +_EMPTY_STORE_VALUE = EmptyStoreValue() + +class Store: + """Storage for a value, with checks for overwriting or reading empty store.""" + __slots__ = ("_val",) + + def __init__(self): + self._val = _EMPTY_STORE_VALUE + + def store(self, val): + if self._val is not _EMPTY_STORE_VALUE: + raise StoreException("Store occupied") + self._val = val + + def reset(self): + # This should only be called in exceptional circumstances (e.g. debugging). + self._val = _EMPTY_STORE_VALUE + + @property + def val(self): + if not self: + raise StoreException("Store empty") + return self._val + + def __nonzero__(self): + return self._val is not _EMPTY_STORE_VALUE + + __bool__ = __nonzero__ + +class EqualStore: + __slots__ = ('_store',) + + def __init__(self): + self._store = Store() + + @property + def val(self): + return self._store.val + + def store(self, val): + try: + self._store.store(val) + except StoreException as e: + try: + okay = bool(self._store._val == val) + except: + raise e from None + else: + if not okay: + raise StoreException("Store occupied with not-equal value") from None + + def reset(self): + self._store.reset() + + +class WrappedFun: + """Represents a function `f` to which `transforms` are to be applied. + + Args: + f: the function to be transformed. + transforms: a list of `(gen, gen_static_args)` tuples representing + transformations to apply to `f.` Here `gen` is a generator function and + `gen_static_args` is a tuple of static arguments for the generator. See + description at the start of this module for the expected behavior of the + generator. + stores: a list of out_store for the auxiliary output of the `transforms`. + params: extra parameters to pass as keyword arguments to `f`, along with the + transformed keyword arguments. + """ + __slots__ = ("f", "transforms", "stores", "params", "in_type", "debug_info") + + def __init__(self, f, transforms, stores, params, in_type, debug_info): + self.f = f + self.transforms = transforms + self.stores = stores + self.params = params + self.in_type = in_type + self.debug_info = debug_info + + @property + def __name__(self): + return getattr(self.f, '__name__', '') + + def wrap(self, gen, gen_static_args, out_store) -> WrappedFun: + """Add another transform and its store.""" + return WrappedFun(self.f, ((gen, gen_static_args),) + self.transforms, + (out_store,) + self.stores, self.params, None, None) + + def populate_stores(self, stores): + """Copy the values from the `stores` into `self.stores`.""" + for self_store, other_store in zip(self.stores, stores): + if self_store is not None: + self_store.store(other_store.val) + + def call_wrapped(self, *args, **kwargs): + """Calls the underlying function, applying the transforms. + + The positional `args` and keyword `kwargs` are passed to the first + transformation generator. + """ + stack = [] + for (gen, gen_static_args), out_store in zip(self.transforms, self.stores): + gen = gen(*(gen_static_args + tuple(args)), **kwargs) + args, kwargs = next(gen) + stack.append((gen, out_store)) + gen = gen_static_args = out_store = None + + try: + ans = self.f(*args, **dict(self.params, **kwargs)) + except: + # Some transformations yield from inside context managers, so we have to + # interrupt them before reraising the exception. Otherwise they will only + # get garbage-collected at some later time, running their cleanup tasks + # only after this exception is handled, which can corrupt the global + # state. + while stack: + stack.pop()[0].close() + raise + + args = kwargs = None + while stack: + gen, out_store = stack.pop() + try: + ans = gen.send(ans) + except: + # As above does for the first half of the transformation, exceptions + # raised in the second half of the transformation also require us to + # clean up references here. + while stack: + stack.pop()[0].close() + raise + if out_store is not None: + ans, side = ans + out_store.store(side) + + return ans + + def __repr__(self): + def transform_to_str(x): + i, (gen, args) = x + return f"{i} : {fun_name(gen)} {fun_name(args)}" + transformation_stack = map(transform_to_str, enumerate(self.transforms)) + return "Wrapped function:\n" + '\n'.join(transformation_stack) + '\nCore: ' + fun_name(self.f) + '\n' + + def __hash__(self): + return hash((self.f, self.transforms, self.params, self.in_type, + self.debug_info)) + + def __eq__(self, other): + return (self.f == other.f and self.transforms == other.transforms and + self.params == other.params and self.in_type == other.in_type and + self.debug_info == other.debug_info) + +@curry +def transformation(gen, fun: WrappedFun, *gen_static_args) -> WrappedFun: + """Adds one more transformation to a WrappedFun. + + Args: + gen: the transformation generator function + fun: a WrappedFun on which to apply the transformation + gen_static_args: static args for the generator function + """ + return fun.wrap(gen, gen_static_args, None) + +@curry +def transformation_with_aux(gen, fun: WrappedFun, *gen_static_args, + use_eq_store=False) -> tuple[WrappedFun, Any]: + """Adds one more transformation with auxiliary output to a WrappedFun.""" + out_store = Store() if not use_eq_store else EqualStore() + out_thunk = lambda: out_store.val + return fun.wrap(gen, gen_static_args, out_store), out_thunk + +def fun_name(f): + try: + return f.__name__ + except: + return str(f) + +def wrap_init(f, params=None) -> WrappedFun: + """Wraps function `f` as a `WrappedFun`, suitable for transformation.""" + params = () if params is None else tuple(sorted(params.items())) + return WrappedFun(f, (), (), params, None, None) + + +def annotate(f: WrappedFun, in_type: core.InputType | None) -> WrappedFun: + assert f.in_type is None + if in_type is None: + return f + _check_input_type(in_type) + return WrappedFun(f.f, f.transforms, f.stores, f.params, in_type, f.debug_info) + +def _check_input_type(in_type: core.InputType) -> None: + # Check that in_type is syntactically well-formed + assert type(in_type) is tuple and all(type(e) is tuple for e in in_type) + assert all(isinstance(a, core.AbstractValue) and type(b) is bool + and not isinstance(a, core.ConcreteArray) for a, b in in_type) + + def valid_size(d) -> bool: + if isinstance(d, core.DBIdx) and type(d.val) is int and d.val >= 0: + return True + return (isinstance(d, (int, core.DBIdx, core.DArray)) and + (not isinstance(d, core.DArray) or type(d) is core.bint and not d.shape)) + assert all(valid_size(d) for a, _ in in_type if type(a) is core.DShapedArray + for d in a.shape) + + # Check that all DBIdx point to positions to the left of the input on which + # they appear. + assert all(d.val < i for i, (aval, _) in enumerate(in_type) + if isinstance(aval, core.DShapedArray) for d in aval.shape + if isinstance(d, core.DBIdx)) + + # Check that all implicit arguments have at least one DBIdx pointing to them. + provided = [e for _, e in in_type] + for aval, _ in in_type: + if type(aval) is core.DShapedArray: + for d in aval.shape: + if isinstance(d, core.DBIdx): + provided[d.val] = True + assert all(provided) + + +class TracingDebugInfo(NamedTuple): + # Packages up trace/staging-time debug info about a func and its parameters, + # formed just before staging to a jaxpr and read in trace-time error messages. + # TODO(mattjj): delete partial_eval.DebugInfo, replace all uses with this cls + traced_for: str # e.g. 'jit', 'scan', etc + func_src_info: str # e.g. f'{fun.__name__} at {filename}:{lineno}' + arg_names: tuple[str, ...] # e.g. ('args[0]', ... ) + result_paths: Callable[[], tuple[str, ...]] | None + +def add_debug_info(f: WrappedFun, debug_info: TracingDebugInfo | None + ) -> WrappedFun: + """Produce a new WrappedFun with debug_info attached.""" + assert f.debug_info is None + if debug_info is None: + return f + return WrappedFun(f.f, f.transforms, f.stores, f.params, f.in_type, debug_info) + + +def cache(call: Callable, *, explain: Callable | None = None): + """Memoization decorator for functions taking a WrappedFun as first argument. + + Args: + call: a Python callable that takes a WrappedFun as its first argument. The + underlying transforms and params on the WrappedFun are used as part of the + memoization cache key. + + Returns: + A memoized version of ``call``. + """ + fun_caches: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + def memoized_fun(fun: WrappedFun, *args): + cache = fun_caches.setdefault(fun.f, new_cache := {}) # type: ignore + if config.check_tracer_leaks.value: + key = (_copy_main_traces(fun.transforms), fun.params, fun.in_type, args, + config.enable_x64.value, config.default_device.value, + config.trace_context()) + else: + key = (fun.transforms, fun.params, fun.in_type, args, config.enable_x64.value, + config.default_device.value, config.trace_context()) + result = cache.get(key, None) + if result is not None: + ans, stores = result + fun.populate_stores(stores) + else: + ans = call(fun, *args) + if explain and config.explain_cache_misses.value: + explain(fun.f, cache is new_cache, cache, key, ans) + cache[key] = (ans, fun.stores) + + return ans + + def _evict_function(f): + fun_caches.pop(f, None) + + memoized_fun.cache_clear = fun_caches.clear # type: ignore + memoized_fun.evict_function = _evict_function # type: ignore + + cache_clearing_funs.add(memoized_fun.cache_clear) + + return memoized_fun + +cache_clearing_funs = weakref.WeakSet() # type: ignore + +def clear_all_caches(): + global cache_clearing_funs + for clear in cache_clearing_funs: + clear() + +@partial(partial, tree_map) +def _copy_main_traces(x): + if isinstance(x, core.MainTrace): + return core.MainTrace(x.level, x.trace_type, **x.payload) + else: + return x + + +@transformation +def hashable_partial(*args): + yield (yield args, {}) + + +def merge_linear_aux(aux1, aux2): + try: + out1 = aux1() + except StoreException: + # store 1 was not occupied, so store 2 better be + try: + out2 = aux2() + except StoreException: + raise StoreException("neither store occupied") from None + else: + return False, out2 + else: + # store 1 was occupied, so let's check store 2 is not occupied + try: + out2 = aux2() + except StoreException: + return True, out1 + else: + raise StoreException("both stores occupied") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/logging_config.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/logging_config.py new file mode 100644 index 0000000000000000000000000000000000000000..d2f9d9c8fb1f6835e2ed3790aeddaf03c4342def --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/logging_config.py @@ -0,0 +1,54 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import sys + +_debug_handler = logging.StreamHandler(sys.stderr) +_debug_handler.setLevel(logging.DEBUG) +# Example log message: +# DEBUG:2023-06-07 00:14:40,280:jax._src.xla_bridge:590: Initializing backend 'cpu' +_debug_handler.setFormatter(logging.Formatter( + "{levelname}:{asctime}:{name}:{lineno}: {message}", style='{')) + +_debug_enabled_loggers = [] + + +def enable_debug_logging(logger_name): + """Makes the specified logger log everything to stderr. + + Also adds more useful debug information to the log messages, e.g. the time. + + Args: + logger_name: the name of the logger, e.g. "jax._src.xla_bridge". + """ + logger = logging.getLogger(logger_name) + logger.addHandler(_debug_handler) + logger.setLevel(logging.DEBUG) + _debug_enabled_loggers.append(logger) + + +def disable_all_debug_logging(): + """Disables all debug logging enabled via `enable_debug_logging`. + + The default logging behavior will still be in effect, i.e. WARNING and above + will be logged to stderr without extra message formatting. + """ + for logger in _debug_enabled_loggers: + logger.removeHandler(_debug_handler) + # Assume that the default non-debug log level is always WARNING. In theory + # we could keep track of what it was set to before. This shouldn't make a + # difference if not other handlers are attached, but set it back in case + # something else gets attached (e.g. absl logger) and for consistency. + logger.setLevel(logging.WARNING) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/maps.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/maps.py new file mode 100644 index 0000000000000000000000000000000000000000..77896d7f972dbec7b0134249afa380a1913a7fb6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/maps.py @@ -0,0 +1,1876 @@ +# Copyright 2020 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections import OrderedDict, abc +from collections.abc import Iterable, Sequence, Mapping +import contextlib +from functools import wraps, partial, partialmethod, lru_cache +import itertools as it +import math +from typing import Callable, Any, NamedTuple, Union, cast as type_cast + +import numpy as np + +from jax import lax +from jax import numpy as jnp + +from jax._src import config +from jax._src import core +from jax._src import dispatch +from jax._src import effects +from jax._src import mesh as mesh_lib +from jax._src import linear_util as lu +from jax._src import op_shardings +from jax._src import sharding_impls +from jax._src import source_info_util +from jax._src import stages +from jax._src import traceback_util +from jax._src.api_util import (flatten_fun_nokwargs, flatten_axes, + _ensure_index_tuple, donation_vector, + shaped_abstractify, check_callable) +from jax._src.array import ArrayImpl +from jax._src.errors import JAXTypeError +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import xla +from jax._src.interpreters import partial_eval as pe +from jax._src.interpreters.partial_eval import ( + trace_to_subjaxpr_dynamic, DynamicJaxprTracer, + convert_constvars_jaxpr, new_jaxpr_eqn) +from jax._src.interpreters import pxla +from jax._src.pjit import (sharding_constraint_p, get_unconstrained_dims, + GSPMDSharding) +from jax._src.sharding_impls import ( + ArrayMapping, NamedSharding, ParsedPartitionSpec, + array_mapping_to_axis_resources, UNSPECIFIED) +from jax._src.tree_util import (tree_flatten, tree_unflatten, all_leaves, + tree_map, treedef_tuple) +from jax._src.util import (safe_map, safe_zip, HashableFunction, unzip2, unzip3, + as_hashable_function, distributed_debug_log, + tuple_insert, moveaxis, split_list, wrap_name, + merge_lists, partition_list, fun_name) + +source_info_util.register_exclusion(__file__) +traceback_util.register_exclusion(__file__) + +map, unsafe_map = safe_map, map +zip = safe_zip + + +class FrozenDict(abc.Mapping): + def __init__(self, *args, **kwargs): + self.contents = dict(*args, **kwargs) + + def __iter__(self): + return iter(self.contents) + + def __len__(self): + return len(self.contents) + + def __getitem__(self, name): + return self.contents[name] + + def __eq__(self, other): + return isinstance(other, FrozenDict) and self.contents == other.contents + + def __hash__(self): + return hash(tuple(self.contents.items())) + + def __repr__(self): + return f"FrozenDict({self.contents})" + +# Multi-dimensional generalized map + +AxisName = core.AxisName +ResourceAxisName = mesh_lib.ResourceAxisName # Different name just for documentation purposes +Mesh = mesh_lib.Mesh +MeshAxisName = mesh_lib.MeshAxisName +ResourceEnv = mesh_lib.ResourceEnv +EMPTY_ENV = mesh_lib.EMPTY_ENV +thread_resources = mesh_lib.thread_resources + + +class SerialLoop: + """Create an anonymous serial loop resource for use in a single xmap axis. + + A use of :py:class:`SerialLoop` in :py:func:`xmap`'s ``axis_resources`` + extends the resource environment with a new serial loop with a unique + unspecified name, that will only be used to partition the axis that + used a given instance. + + This is unlike :py:func:`serial_loop`, which makes it possible to iterate + jointly over chunks of multiple axes (with the usual requirement that they + do not coincide in a named shape of any value in the program). + + Example:: + + # Processes `x` in a vectorized way, but in 20 micro-batches. + xmap(f, in_axes=['i'], out_axes=[i], axis_resources={'i': SerialLoop(20)})(x) + + # Computes the result in a vectorized way, but in 400 micro-batches, + # once for each coordinate (0, 0) <= (i, j) < (20, 20). Each `SerialLoop` + # creates a fresh anonymous loop. + xmap(h, in_axes=(['i'], ['j']), out_axes=['i', 'j'], + axis_resources={'i': SerialLoop(20), 'j': SerialLoop(20)})(x, y) + """ + length: int + + def __init__(self, length): + self.length = length + + def __eq__(self, other): + return self.length == other.length + + def __hash__(self): + return hash(self.length) + + +@contextlib.contextmanager +def serial_loop(name: ResourceAxisName, length: int): + """Define a serial loop resource to be available in scope of this context manager. + + This is similar to :py:class:`Mesh` in that it extends the resource + environment with a resource called ``name``. But, any use of this resource + axis in ``axis_resources`` argument of :py:func:`xmap` will cause the + body of :py:func:`xmap` to get executed ``length`` times with each execution + only processing only a slice of inputs mapped along logical axes assigned + to this resource. + + This is especially useful in that it makes it possible to lower the memory + usage compared to :py:func:`vmap`, because it will avoid simultaneous + materialization of intermediate values for every point in the batch. + + Note that collectives over loop axes are not supported, so they are less + versatile than physical mesh axes. + + Args: + name: Name of the loop in the resource environment. + length: Number of iterations. + + Example:: + + >>> x = jnp.linspace(0, jnp.pi, 4) + ... + >>> with serial_loop('l', len(x)): + ... out = xmap( + ... lambda x: jnp.sin(x) * 5, # This will be called 4 times with different + ... # slices of x. + ... in_axes=['i'], out_axes=['i'], + ... axis_resources={'i': 'l'})(x) + >>> out.shape + (4,) + """ + old_env: ResourceEnv = getattr(thread_resources, "env", EMPTY_ENV) + thread_resources.env = old_env.with_extra_loop(mesh_lib.Loop(name, length)) + try: + yield + finally: + thread_resources.env = old_env + + +_next_resource_id = 0 +class _UniqueResourceName: + def __init__(self, uid, tag=None): + self.uid = uid + self.tag = tag + def __eq__(self, other): + return type(other) is _UniqueResourceName and self.uid == other.uid + def __hash__(self): + return hash(self.uid) + def __repr__(self): + return f"" + +def fresh_resource_name(tag=None): + global _next_resource_id + try: + return _UniqueResourceName(_next_resource_id, tag) + finally: + _next_resource_id += 1 + + +# This is really a Dict[AxisName, int], but we don't define a +# pytree instance for it, so that it is treated as a leaf. +class AxisNamePos(FrozenDict): + user_repr: str + expected_rank: int | None = None + + def __init__(self, *args, user_repr, **kwargs): + super().__init__(*args, **kwargs) + self.user_repr = user_repr + +class AxisNamePosWithRank(AxisNamePos): + def __init__(self, *args, expected_rank, **kwargs): + super().__init__(*args, **kwargs) + self.expected_rank = expected_rank + + +# str(...) == 'Ellipsis' which is really annoying +class DotDotDotRepr: + def __repr__(self): return '...' + + +def _parse_entry(arg_name, entry): + # Dictionaries mapping axis names to positional axes + if isinstance(entry, dict) and all(isinstance(v, int) for v in entry.keys()): + result = AxisNamePos(((name, axis) for axis, name in entry.items()), + user_repr=str(entry)) + num_mapped_dims = len(entry) + # Non-empty lists or tuples that optionally terminate with an ellipsis + elif isinstance(entry, (tuple, list)): + if entry and entry[-1] == ...: + constr = AxisNamePos + entry = entry[:-1] + tail = [DotDotDotRepr()] if isinstance(entry, list) else (DotDotDotRepr(),) + user_repr = str(entry + tail) + else: + constr = partial(AxisNamePosWithRank, expected_rank=len(entry)) + user_repr = str(entry) + result = constr(((name, axis) for axis, name in enumerate(entry) + if name is not None), + user_repr=user_repr) + num_mapped_dims = sum(name is not None for name in entry) + else: + raise TypeError(f"""\ +Value mapping specification in xmap {arg_name} pytree can be either: +- lists of axis names (possibly ending with the ellipsis object: ...) +- dictionaries that map positional axes (integers) to axis names (e.g. {{2: 'name'}}) +but got: {entry}""") + if len(result) != num_mapped_dims: + raise ValueError(f"Named axes should be unique within each {arg_name} argument " + f"specification, but one them is: {entry}") + for axis in result.values(): + if axis < 0: + raise ValueError(f"xmap doesn't support negative axes in {arg_name}") + return result + +def _is_axes_leaf(entry): + if isinstance(entry, dict) and all_leaves(entry.values()): + return True + # NOTE: `None`s are not considered leaves by `all_leaves` + if isinstance(entry, (tuple, list)) and all_leaves(v for v in entry if v is not None): + return True + return False + +def _prepare_axes(axes, arg_name): + entries, treedef = tree_flatten(axes, is_leaf=_is_axes_leaf) + entries = map(partial(_parse_entry, arg_name), entries) + return tree_unflatten(treedef, entries), entries, treedef + +Resource = Union[ResourceAxisName, SerialLoop] +ResourceSet = Union[Resource, tuple[Resource, ...]] + +# TODO: Some syntactic sugar to make the API more usable in a single-axis case? +# TODO: Are the resource axes scoped lexically or dynamically? Dynamically for now! +def xmap(fun: Callable, + in_axes, + out_axes, + *, + axis_sizes: Mapping[AxisName, int] | None = None, + axis_resources: Mapping[AxisName, ResourceSet] | None = None, + donate_argnums: int | Sequence[int] = (), + backend: str | None = None) -> stages.Wrapped: + """Assign a positional signature to a program that uses named array axes. + + .. warning:: + xmap is deprecated and will be removed in a future release. Use + :py:func:`~jax.shard_map` or :py:func:`~jax.vmap` with the + ``spmd_axis_name`` argument for expressing SPMD device-parallel + computations. Please file an issue on https://github.com/google/jax/issues + if neither are suitable for your use case. + + .. warning:: + This is an experimental feature and the details can change at + any time. Use at your own risk! + + .. warning:: + This docstring is aspirational. Not all features of the named axis + programming model have been implemented just yet. + + The usual programming model of JAX (or really NumPy) associates each array + with two pieces of metadata describing its type: the element type (``dtype``) + and the ``shape``. :py:func:`xmap` extends this model by adding support for + *named axes*. In particular, each array used in a function wrapped by + :py:func:`xmap` can additionally have a non-empty ``named_shape`` attribute, + which can be used to query the set of named axes (introduced by + :py:func:`xmap`) appearing in that value along with their shapes. + Furthermore, in most places where positional axis indices are allowed (for + example the `axes` arguments in :py:func:`sum`), bound axis names are also + accepted. The :py:func:`einsum` language is extended inside :py:func:`xmap` + to additionally allow contractions that involve named axes. Broadcasting of + named axes happens *by name*, i.e. all axes with equal names are expected to + have equal shapes in all arguments of a broadcasting operation, while the + result has a (set) union of all named axes. The positional semantics of the + program remain unchanged, and broadcasting still implicitly right-aligns + positional axes for unification. For an extended description of the + :py:func:`xmap` programming model, please refer to the :py:func:`xmap` + tutorial notebook in main JAX documentation. + + Note that since all top-level JAX expressions are interpreted in the NumPy + programming model, :py:func:`xmap` can also be seen as an adapter that + converts a function that uses named axes (including in arguments and returned + values) into one that takes and returns values that only have positional + axes. + + The default lowering strategy of :py:func:`xmap` converts all named axes into + positional axes, working similarly to multiple applications of + :py:func:`vmap`. However, this behavior can be further customized by the + ``axis_resources`` argument. When specified, each axis introduced by + :py:func:`xmap` can be assigned to one or more *resource axes*. Those include + the axes of the hardware mesh, as defined by the :py:class:`Mesh` context + manager. Each value that has a named axis in its ``named_shape`` will be + partitioned over all mesh axes that axis is assigned to. Hence, + :py:func:`xmap` can be seen as an alternative to :py:func:`pmap` that also + exposes a way to automatically partition the computation over multiple + devices. + + .. warning:: + While it is possible to assign multiple axis names to a single resource axis, + care has to be taken to ensure that none of those named axes co-occur in a + ``named_shape`` of any value in the named program. At the moment this is + **completely unchecked** and will result in **undefined behavior**. The + final release of :py:func:`xmap` will enforce this invariant, but it is a + work in progress. + + Note that you do not have to worry about any of this for as long as no + resource axis is repeated in ``axis_resources.values()``. + + Note that any assignment of ``axis_resources`` doesn't ever change the + results of the computation, but only how it is carried out (e.g. how many + devices are used). This makes it easy to try out various ways of + partitioning a single program in many distributed scenarios (both small- and + large-scale), to maximize the performance. As such, :py:func:`xmap` can be + seen as a way to seamlessly interpolate between :py:func:`vmap` and + :py:func:`pmap`-style execution. + + Args: + fun: Function that uses named axes. Its arguments and return + value should be arrays, scalars, or (nested) standard Python containers + (tuple/list/dict) thereof (in general: valid pytrees). + in_axes: A Python object with the same container (pytree) structure as the + signature of arguments to ``fun``, but with a positional-to-named axis + mapping in place of every array argument. The valid positional-to-named + mappings are: (1) a ``Dict[int, AxisName]`` specifying that a positional + dimensions given by dictionary keys are to be converted to named axes + of given names (2) a list of axis names that ends with the Ellipsis object + (``...``) in which case a number of leading positional axes of the argument + will be converted into named axes inside the function. Note that ``in_axes`` + can also be a prefix of the argument container structure, in which case the + mapping is repeated for all arrays in the collapsed subtree. + out_axes: A Python object with the same container (pytree) structure as the + returns of ``fun``, but with a positional-to-named axis mapping in place + of every returned array. The valid positional-to-named mappings are the same + as in ``in_axes``. Note that ``out_axes`` can also be a prefix of the return + container structure, in which case the mapping is repeated for all arrays + in the collapsed subtree. + axis_sizes: A dict mapping axis names to their sizes. All axes defined by xmap + have to appear either in ``in_axes`` or ``axis_sizes``. Sizes of axes + that appear in ``in_axes`` are inferred from arguments whenever possible. + In multi-host scenarios, the user-specified sizes are expected to be the + global axis sizes (and might not match the expected size of local inputs). + axis_resources: A dictionary mapping the axes introduced in this + :py:func:`xmap` to one or more resource axes. Any array that has in its + shape an axis with some resources assigned will be partitioned over the + resources associated with the respective resource axes. + donate_argnums: Specify which argument buffers are "donated" to the computation. + It is safe to donate argument buffers if you no longer need them once the + computation has finished. In some cases XLA can make use of donated + buffers to reduce the amount of memory needed to perform a computation, + for example recycling one of your input buffers to store a result. You + should not reuse buffers that you donate to a computation, JAX will raise + an error if you try to. + + For more details on buffer donation see the `FAQ `_. + + backend: This is an experimental feature and the API is likely to change. + Optional, a string representing the XLA backend. 'cpu', 'gpu', or 'tpu'. + + Returns: + A version of ``fun`` that takes in arrays with positional axes in place of + named axes bound in this :py:func:`xmap` call, and results with all named + axes converted to positional axes. If ``axis_resources`` is specified, + ``fun`` can additionally execute in parallel on multiple devices. + + For example, :py:func:`xmap` makes it very easy to convert a function that + computes the vector inner product (such as :py:func:`jax.numpy.vdot`) into + one that computes a matrix multiplication: + + >>> import jax.numpy as jnp + >>> x = jnp.arange(10).reshape((2, 5)) + >>> xmap(jnp.vdot, + ... in_axes=({0: 'left'}, {1: 'right'}), + ... out_axes=['left', 'right', ...])(x, x.T) + Array([[ 30, 80], + [ 80, 255]], dtype=int32) + + Note that the contraction in the program is performed over the positional axes, + while named axes are just a convenient way to achieve batching. While this + might seem like a silly example at first, it might turn out to be useful in + practice, since with conjunction with ``axis_resources`` this makes it possible + to implement a distributed matrix-multiplication in just a few lines of code:: + + devices = np.array(jax.devices())[:4].reshape((2, 2)) + with Mesh(devices, ('x', 'y')): # declare a 2D mesh with axes 'x' and 'y' + distributed_out = xmap( + jnp.vdot, + in_axes=({0: 'left'}, {1: 'right'}), + out_axes=['left', 'right', ...], + axis_resources={'left': 'x', 'right': 'y'})(x, x.T) + + Still, the above examples are quite simple. After all, the xmapped + computation was a simple NumPy function that didn't use the axis names at all! + So, let's explore a slightly larger example which is linear regression:: + + def regression_loss(x, y, w, b): + # Contract over in_features. Batch and out_features are present in + # both inputs and output, so they don't need to be mentioned + y_pred = jnp.einsum('{in_features},{in_features}->{}', x, w) + b + error = jnp.sum((y - y_pred) ** 2, axis='out_features') + return jnp.mean(error, axis='batch') + + xmap(regression_loss, + in_axes=(['batch', 'in_features', ...], + ['batch', 'out_features', ...], + ['in_features', 'out_features', ...], + ['out_features', ...]), + out_axes={}) # Loss is reduced over all axes, including batch! + + .. note:: + When using ``axis_resources`` along with a mesh that is controlled by + multiple JAX hosts, keep in mind that in any given process :py:func:`xmap` + only expects the data slice that corresponds to its local devices to be + specified. This is in line with the current multi-host :py:func:`pmap` + programming model. + """ + check_callable(fun) + + if isinstance(in_axes, list) and not _is_axes_leaf(in_axes): + # To be a tree prefix of the positional args tuple, in_axes can never be a + # list: if in_axes is not a leaf, it must be a tuple of trees. However, + # in cases like these users expect tuples and lists to be treated + # essentially interchangeably, so we canonicalize lists to tuples here + # rather than raising an error. https://github.com/google/jax/issues/2367 + in_axes = tuple(in_axes) + + if in_axes == (): # Allow empty argument lists + in_axes, in_axes_entries = (), [] + else: + in_axes, in_axes_entries, _ = _prepare_axes(in_axes, "in_axes") + if out_axes == (): + raise ValueError("xmapped functions cannot have no return values") + else: + out_axes, out_axes_entries, out_axes_treedef = _prepare_axes(out_axes, "out_axes") + out_axes_entries = tuple(out_axes_entries) # Make entries hashable + + axis_sizes = {} if axis_sizes is None else axis_sizes + axis_resources = {} if axis_resources is None else axis_resources + + axis_sizes_names = set(axis_sizes.keys()) + in_axes_names = set(it.chain(*(spec.keys() for spec in in_axes_entries))) + defined_names = axis_sizes_names | in_axes_names + out_axes_names = set(it.chain(*(spec.keys() for spec in out_axes_entries))) + + anon_serial_loops = [] + def normalize_resource(r) -> ResourceAxisName: + if isinstance(r, SerialLoop): + name = fresh_resource_name() + anon_serial_loops.append((name, r.length)) + return name + return r + + axes_with_resources = set(axis_resources.keys()) + if axes_with_resources - defined_names: + raise ValueError(f"All axes that were assigned resources have to appear in " + f"in_axes or axis_sizes, but the following are missing: " + f"{axes_with_resources - defined_names}") + if out_axes_names - defined_names: + raise ValueError(f"All axis names appearing in out_axes must also appear in " + f"in_axes or axis_sizes, but the following are missing: " + f"{out_axes_names - defined_names}") + + normalized_axis_resources: dict[AxisName, tuple[ResourceAxisName, ...]] = {} + for axis in defined_names: + resources = axis_resources.get(axis, ()) + if not isinstance(resources, tuple): + resources = (resources,) + normalized_axis_resources[axis] = tuple(unsafe_map(normalize_resource, resources)) + frozen_axis_resources = FrozenDict(normalized_axis_resources) + necessary_resources = set(it.chain(*frozen_axis_resources.values())) + + for axis, resources in frozen_axis_resources.items(): + if len(set(resources)) != len(resources): # type: ignore + raise ValueError(f"Resource assignment of a single axis must be a tuple of " + f"distinct resources, but specified {resources} for axis {axis}") + + donate_argnums = _ensure_index_tuple(donate_argnums) + + # A little performance optimization to avoid iterating over all args unnecessarily + has_input_rank_assertions = any(spec.expected_rank is not None for spec in in_axes_entries) + has_output_rank_assertions = any(spec.expected_rank is not None for spec in out_axes_entries) + + def infer_params(*args): + # Putting this outside of fun_mapped would make resources lexically scoped + resource_env = thread_resources.env + available_resources = set(resource_env.shape.keys()) + + if necessary_resources - available_resources: + raise ValueError(f"In-scope resources are insufficient to execute the " + f"xmapped function. The missing resources are: " + f"{necessary_resources - available_resources}") + + args_flat, in_tree = tree_flatten(args) + fun_flat, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree) + if donate_argnums: + donated_invars = donation_vector(donate_argnums, (), in_tree, kws=False) + else: + donated_invars = (False,) * len(args_flat) + in_axes_flat = _flatten_axes("xmap in_axes", in_tree, in_axes, tupled_args=True) + + # Some pytree containers might be unhashable, so we flatten the out_axes + # pytree into a treedef and entries which are guaranteed to be hashable. + out_axes_thunk = HashableFunction( + lambda: tuple(_flatten_axes("xmap out_axes", out_tree(), out_axes, tupled_args=False)), + closure=(out_axes_entries, out_axes_treedef)) + + axis_resource_count = _get_axis_resource_count( + frozen_axis_resources, resource_env) + + for axis, size in axis_sizes.items(): + resources = axis_resource_count[axis] + if size % resources.nglobal != 0: + global_size = "Global size" if resources.distributed else "Size" + raise ValueError(f"{global_size} of axis {axis} ({size}) is not divisible " + f"by the total number of resources assigned to this axis " + f"({frozen_axis_resources[axis]}, {resources.nglobal} in total)") + frozen_global_axis_sizes = _get_axis_sizes( + args_flat, in_axes_flat, axis_sizes, axis_resource_count) + + missing_sizes = defined_names - set(frozen_global_axis_sizes.keys()) + if missing_sizes: + raise ValueError(f"Failed to infer size of axes: {', '.join(unsafe_map(str, missing_sizes))}. " + f"You've probably passed in empty containers in place of arguments that had " + f"those axes in their in_axes. Provide the sizes of missing axes explicitly " + f"via axis_sizes to fix this error.") + + if has_input_rank_assertions: + for arg, spec in zip(args_flat, in_axes_flat): + if spec.expected_rank is not None and spec.expected_rank != arg.ndim: + raise ValueError(f"xmap argument has an in_axes specification of {spec.user_repr}, " + f"which asserts that it should be of rank {spec.expected_rank}, " + f"but the argument has rank {arg.ndim} (and shape {arg.shape})") + + _check_gda_or_array_xmap_partitioning( + frozen_axis_resources, resource_env, frozen_global_axis_sizes, + in_axes_flat, args_flat) + + params = dict( + name=fun_name(fun), + in_axes=tuple(in_axes_flat), + out_axes_thunk=out_axes_thunk, + donated_invars=donated_invars, + global_axis_sizes=frozen_global_axis_sizes, + axis_resources=frozen_axis_resources, + resource_env=resource_env, + backend=backend, + spmd_in_axes=None, + spmd_out_axes_thunk=None) + return fun_flat, args_flat, params, in_tree, out_tree + + def verify_outputs(out_flat, out_tree, params): + if has_output_rank_assertions: + for out, spec in zip(out_flat, params['out_axes_thunk']()): + if spec.expected_rank is not None and spec.expected_rank != out.ndim: + raise ValueError(f"xmap output has an out_axes specification of {spec.user_repr}, " + f"which asserts that it should be of rank {spec.expected_rank}, " + f"but the output has rank {out.ndim} (and shape {out.shape})") + return tree_unflatten(out_tree(), out_flat) + + def decorate_serial(f): + for loop_params in reversed(anon_serial_loops): + f = serial_loop(*loop_params)(f) + return f + + @wraps(fun) + @decorate_serial + def fun_mapped(*args): + tree_map(dispatch.check_arg, args) + fun_flat, args_flat, params, _, out_tree = infer_params(*args) + out_flat = xmap_p.bind(fun_flat, *args_flat, **params) + return verify_outputs(out_flat, out_tree, params) + + @decorate_serial + def lower(*args, **kwargs): + lowering_parameters = kwargs.pop( + '_experimental_lowering_platform', mlir.LoweringParameters()) + fun_flat, args_flat, params, in_tree, out_tree = infer_params(*args) + avals_flat = [shaped_abstractify(arg) for arg in args_flat] + computation = make_xmap_callable( + fun_flat, params['name'], params['in_axes'], params['out_axes_thunk'], + params['donated_invars'], params['global_axis_sizes'], params['axis_resources'], + params['resource_env'], params['backend'], params['spmd_in_axes'], + params['spmd_out_axes_thunk'], + lowering_parameters, *avals_flat) + + in_tree = treedef_tuple([in_tree, tree_flatten({})[1]]) + in_avals = in_tree.unflatten(avals_flat) + return stages.Lowered.from_flat_info( + computation, in_tree, in_avals, donate_argnums, out_tree()) + + fun_mapped.lower = lower + return type_cast(stages.Wrapped, fun_mapped) + +def xmap_impl(fun: lu.WrappedFun, *args, name, in_axes, out_axes_thunk, donated_invars, + global_axis_sizes, axis_resources, resource_env, backend, + spmd_in_axes, spmd_out_axes_thunk): + in_avals = [core.raise_to_shaped(core.get_aval(arg)) for arg in args] + computation = make_xmap_callable( + fun, name, in_axes, out_axes_thunk, donated_invars, global_axis_sizes, + axis_resources, resource_env, backend, + spmd_in_axes, spmd_out_axes_thunk, + mlir.LoweringParameters(), *in_avals) + xmap_callable = computation.compile().unsafe_call + distributed_debug_log(("Running xmapped function", name), + ("python function", fun.f), + ("mesh", resource_env.physical_mesh), + ("abstract args", in_avals)) + return xmap_callable(*args) + +@lu.cache +def make_xmap_callable(fun: lu.WrappedFun, + name, + in_axes, out_axes_thunk, donated_invars, + global_axis_sizes, axis_resources, resource_env, backend, + spmd_in_axes, spmd_out_axes_thunk, + lowering_parameters: mlir.LoweringParameters, + *in_avals): + plan = EvaluationPlan.from_axis_resources( + axis_resources, resource_env, global_axis_sizes) + + # TODO: Making axis substitution final style would allow us to avoid + # tracing to jaxpr here + mapped_in_avals = [_delete_aval_axes(aval, in_axes, global_axis_sizes) + for aval, in_axes in zip(in_avals, in_axes)] + with core.extend_axis_env_nd(global_axis_sizes.items()): + with dispatch.log_elapsed_time( + "Finished tracing + transforming {fun_name} for xmap in {elapsed_time} sec", + fun_name=fun.__name__, event=dispatch.JAXPR_TRACE_EVENT): + jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, mapped_in_avals) + out_axes = out_axes_thunk() + # _check_out_avals_vs_out_axes(out_avals, out_axes, global_axis_sizes) + # NOTE: We don't use avals and all params, so only pass in the relevant parts (too lazy...) + _resource_typing_xmap([], dict(axis_resources=axis_resources, + out_axes=out_axes, + call_jaxpr=jaxpr, + resource_env=resource_env, + name=name), + source_info_util.new_source_info(), resource_env, {}) + jaxpr = plan.subst_axes_with_resources(jaxpr) + use_spmd_lowering = SPMD_LOWERING.value + ensure_fixed_sharding = _ENSURE_FIXED_SHARDING.value + if use_spmd_lowering and ensure_fixed_sharding: + jaxpr = _fix_inferred_spmd_sharding(jaxpr, resource_env) + + f = lu.wrap_init(core.jaxpr_as_fun(core.ClosedJaxpr(jaxpr, consts))) + f = hide_mapped_axes(f, tuple(in_axes), tuple(out_axes)) + f = plan.vectorize_and_loop(f, in_axes, out_axes) + + used_resources = _jaxpr_resources(jaxpr, resource_env) | set(it.chain(*axis_resources.values())) + used_mesh_axes = used_resources & resource_env.physical_resource_axes + if used_mesh_axes: + assert spmd_in_axes is None and spmd_out_axes_thunk is None # No outer xmaps, so should be None + mesh_in_axes, mesh_out_axes = plan.to_mesh_axes(in_axes, out_axes) + mesh = resource_env.physical_mesh + tiling_method: pxla.TilingMethod + if SPMD_LOWERING_MANUAL.value: + manual_mesh_axes = frozenset(it.chain.from_iterable(plan.physical_axis_resources.values())) + tiling_method = pxla.TileManual(manual_mesh_axes) + else: + tiling_method = pxla.TileVectorize() + in_shardings = [NamedSharding(mesh, array_mapping_to_axis_resources(i)) + for i in mesh_in_axes] + out_shardings = [NamedSharding(mesh, array_mapping_to_axis_resources(o)) + for o in mesh_out_axes] + return pxla.lower_mesh_computation( + f, 'xmap', name, mesh, + in_shardings, out_shardings, donated_invars, + use_spmd_lowering, in_avals, + tiling_method=tiling_method, + lowering_parameters=lowering_parameters) + else: + jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(f, in_avals) + return pxla.lower_sharding_computation( + core.ClosedJaxpr(jaxpr, consts), 'jit', name, + (UNSPECIFIED,) * len(in_avals), (UNSPECIFIED,) * len(out_avals), + (None,) * len(in_avals), (None,) * len(out_avals), + donated_invars, keep_unused=True, inline=False, + devices_from_context=None, lowering_parameters=lowering_parameters) + + +class EvaluationPlan(NamedTuple): + """Encapsulates preprocessing common to top-level xmap invocations and its translation rule.""" + resource_env: ResourceEnv + physical_axis_resources: dict[AxisName, tuple[ResourceAxisName, ...]] + loop_axis_resources: dict[AxisName, tuple[ResourceAxisName, ...]] + axis_subst_dict: dict[AxisName, tuple[ResourceAxisName, ...]] + axis_vmap_size: dict[AxisName, int | None] + + @property + def axis_subst(self) -> core.AxisSubst: + return lambda name: self.axis_subst_dict.get(name, (name,)) + + @property + def resource_axis_env(self): + env = dict(self.resource_env.shape) + for axis, size in self.axis_vmap_size.items(): + if size is None: + continue + vmap_axis = self.axis_subst_dict[axis][-1] + env[vmap_axis] = size + return env + + @classmethod + def from_axis_resources(cls, + axis_resources: dict[AxisName, tuple[ResourceAxisName, ...]], + resource_env: ResourceEnv, + global_axis_sizes: dict[AxisName, int]): + physical_axis_resources, loop_axis_resources = _unzip_axis_resources( + axis_resources, resource_env) + axis_resource_count = _get_axis_resource_count( + axis_resources, resource_env) + axis_subst_dict = dict(axis_resources) + axis_vmap_size: dict[AxisName, int | None] = {} + for naxis, raxes in sorted(axis_resources.items(), key=lambda x: str(x[0])): + num_resources = axis_resource_count[naxis] + assert global_axis_sizes[naxis] % num_resources.nglobal == 0 + local_tile_size = global_axis_sizes[naxis] // num_resources.nglobal + # We have to vmap when there are no resources (to handle the axis name!) or + # when every resource gets chunks of values. + if not raxes or local_tile_size > 1: + axis_vmap_size[naxis] = local_tile_size + axis_subst_dict[naxis] += (fresh_resource_name(naxis),) + else: + axis_vmap_size[naxis] = None + return cls(resource_env, + physical_axis_resources, loop_axis_resources, + axis_subst_dict, axis_vmap_size) + + def subst_axes_with_resources(self, jaxpr): + try: + if any(self.loop_axis_resources.values()): + _check_no_loop_collectives(jaxpr, self.loop_axis_resources) + with core.extend_axis_env_nd(self.resource_axis_env.items()): + return core.subst_axis_names_jaxpr(jaxpr, self.axis_subst) + except core.DuplicateAxisNameError: + raise AssertionError("Incomplete resource type-checking? Please open a bug report!") + + def vectorize_and_loop(self, f: lu.WrappedFun, in_axes, out_axes) -> lu.WrappedFun: + vmap_axes = { + naxis: raxes[-1] + for naxis, raxes in self.axis_subst_dict.items() + if self.axis_vmap_size[naxis] is not None + } + for naxis, vaxis in sorted(vmap_axes.items(), key=lambda x: x[1].uid): + local_tile_size = self.axis_vmap_size[naxis] + map_in_axes = tuple(unsafe_map(lambda spec: spec.get(naxis, None), in_axes)) + map_out_axes = tuple(unsafe_map(lambda spec: spec.get(naxis, None), out_axes)) + f = batching.vtile(f, map_in_axes, map_out_axes, tile_size=local_tile_size, axis_name=vaxis) + + used_loops = set(it.chain.from_iterable(self.loop_axis_resources.values())) + if not used_loops: + return f + + if len(used_loops) > 1: + # TODO: Support multiple loops + raise NotImplementedError("Only one loop per xmap is supported") + loop_in_axes = _to_resource_axes(in_axes, self.loop_axis_resources) + loop_out_axes = _to_resource_axes(out_axes, self.loop_axis_resources) + loop_name, = used_loops + loop_length = self.resource_env.shape[loop_name] + def looped_f(*args): + def body(i, _): + # XXX: This call_wrapped is only valid under the assumption that scan + # only ever traces the body once (which it does at the moment). + result = f.call_wrapped( + *(_slice_tile(arg, spec.get(loop_name, None), i, loop_length) + for arg, spec in zip(args, loop_in_axes))) + return i + 1, result + _, stacked_results = lax.scan(body, 0, (), length=loop_length) + return [_merge_leading_axis(sresult, spec.get(loop_name, None)) + for sresult, spec in zip(stacked_results, loop_out_axes)] + return lu.wrap_init(looped_f) + + def to_mesh_axes(self, in_axes, out_axes=None): + """ + Convert in/out_axes parameters ranging over logical dimensions to + in/out_axes that range over the mesh dimensions. + """ + if out_axes is None: + return _to_resource_axes(in_axes, self.physical_axis_resources) + else: + return (_to_resource_axes(in_axes, self.physical_axis_resources), + _to_resource_axes(out_axes, self.physical_axis_resources)) + +# -------- xmap primitive and its transforms -------- + +# xmap has a different set of parameters than pmap, so we make it its own primitive type +class XMapPrimitive(core.MapPrimitive): + def __init__(self): + super().__init__('xmap') + self.def_impl(xmap_impl) + self.def_custom_bind(self.bind) + + def bind(self, fun, *args, in_axes, **params): + assert len(in_axes) == len(args), (in_axes, args) + return core.map_bind(self, fun, *args, in_axes=in_axes, **params) + + def process(self, trace, fun, tracers, params): + return trace.process_xmap(self, fun, tracers, params) + + def post_process(self, trace, out_tracers, params): + post_process = getattr(trace, 'post_process_xmap', None) + if post_process is None: + raise NotImplementedError + return post_process(self, out_tracers, params) + + def get_bind_params(self, params): + new_params = dict(params) + jaxpr = new_params.pop('call_jaxpr') + subfun = lu.hashable_partial(lu.wrap_init(core.eval_jaxpr), jaxpr, ()) + axes = new_params.pop('out_axes') + new_params['out_axes_thunk'] = HashableFunction(lambda: axes, closure=axes) + spmd_axes = new_params.pop('spmd_out_axes') + if spmd_axes is not None: + new_params['spmd_out_axes_thunk'] = \ + HashableFunction(lambda: spmd_axes, closure=spmd_axes) + else: + new_params['spmd_out_axes_thunk'] = None + return [subfun], new_params + +xmap_p = XMapPrimitive() +core.EvalTrace.process_xmap = core.EvalTrace.process_call # type: ignore +def _process_xmap_default(self, call_primitive, f, tracers, params): + raise NotImplementedError(f"{type(self)} must override process_xmap to handle xmap") +core.Trace.process_xmap = _process_xmap_default # type: ignore + +def _xmap_axis_subst(params, subst, traverse): + if 'call_jaxpr' not in params: # TODO(apaszke): This feels sketchy, but I'm not sure why + return params + if not traverse: + return params + def shadowed_subst(name): + return (name,) if name in params['global_axis_sizes'] else subst(name) + with core.extend_axis_env_nd(params['global_axis_sizes'].items()): + new_jaxpr = core.subst_axis_names_jaxpr(params['call_jaxpr'], shadowed_subst) + return dict(params, call_jaxpr=new_jaxpr) +core.axis_substitution_rules[xmap_p] = _xmap_axis_subst + +# NOTE: We don't have to handle spmd_{in|out}_axes here, because +# SPMD batching always gets involved as the last transform before XLA translation +ad.JVPTrace.process_xmap = ad.JVPTrace.process_call # type: ignore +ad.call_param_updaters[xmap_p] = pxla.xla_call_jvp_update_params + +def _xmap_transpose(params, call_jaxpr, args, cts_in, cts_in_avals): + all_args, in_tree_def = tree_flatten(((), args, cts_in)) # empty consts + fun = lu.hashable_partial(lu.wrap_init(ad.backward_pass), call_jaxpr, False) + fun, nz_arg_cts = ad.nonzero_outputs(fun) + fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def) + # Preserve axis for primal arguments, skip tangents (represented as undefined primals). + in_axes, out_axes = params['in_axes'], params['out_axes'] + new_in_axes = (*(axis for axis, x in zip(in_axes, args) if not ad.is_undefined_primal(x)), + *(axis for axis, x in zip(out_axes, cts_in) if type(x) is not ad.Zero)) + # NOTE: This assumes that the output cotangents being zero is a deterministic + # function of which input cotangents were zero. + @as_hashable_function(closure=(in_axes, tuple(type(c) is ad.Zero for c in cts_in))) + def out_axes_thunk(): + return tuple(axis for axis, nz in zip(in_axes, nz_arg_cts()) if nz) + new_params = dict(params, + name=wrap_name(params['name'], 'transpose'), + in_axes=new_in_axes, + out_axes_thunk=out_axes_thunk, + donated_invars=(False,) * len(new_in_axes), + spmd_out_axes_thunk=None) + del new_params['out_axes'] + del new_params['spmd_out_axes'] + out_flat = xmap_p.bind(fun, *all_args, **new_params) + arg_cts = tree_unflatten(out_tree(), out_flat) + + axis_resource_count = _get_axis_resource_count( + params['axis_resources'], params['resource_env']) + local_axis_sizes = { + axis: axis_resource_count[axis].to_local(global_size) + for axis, global_size in params['global_axis_sizes'].items() + } + def unmap_zero(zero, axes): + return ad.Zero(_insert_aval_axes(zero.aval, axes, local_axis_sizes)) + return tuple(unmap_zero(arg_ct, in_axis) if type(arg_ct) is ad.Zero else arg_ct + for arg_ct, in_axis in zip(arg_cts, in_axes)) +ad.primitive_transposes[xmap_p] = _xmap_transpose + + +def _typecheck_xmap( + _, *in_atoms, call_jaxpr, name, in_axes, out_axes, donated_invars, + global_axis_sizes, axis_resources, resource_env, backend, + spmd_in_axes, spmd_out_axes): + in_avals = [x.aval for x in in_atoms] + axis_resource_count = _get_axis_resource_count( + axis_resources, resource_env) + local_axis_sizes = { + axis: axis_resource_count[axis].to_local(global_size) + for axis, global_size in global_axis_sizes.items() + } + binder_in_avals = [_insert_aval_axes(v.aval, a_in_axes, local_axis_sizes) + for v, a_in_axes in zip(call_jaxpr.invars, in_axes)] + for binder_in_aval, in_aval in zip(binder_in_avals, in_avals): + if not core.typecompat(binder_in_aval, in_aval): + raise core.JaxprTypeError( + f"xmap passes operand {in_aval} to jaxpr expecting {binder_in_aval}") + + with core.extend_axis_env_nd(global_axis_sizes.items()): + core._check_jaxpr(lambda: core.JaxprPpContext(), call_jaxpr) + + mapped_out_avals = [v.aval for v in call_jaxpr.outvars] + out_avals = [_insert_aval_axes(a, a_out_axes, local_axis_sizes) + for a, a_out_axes in zip(mapped_out_avals, out_axes)] + effs = core.filter_named_axis_effects(call_jaxpr.effects, global_axis_sizes) + return out_avals, effs +core.custom_typechecks[xmap_p] = _typecheck_xmap + + +def _resource_typing_xmap(avals, + params, + source_info: source_info_util.SourceInfo, + resource_env, + outer_axis_resources): + axis_resources = params['axis_resources'] + inner_axis_resources = dict(outer_axis_resources) + inner_axis_resources.update(axis_resources) + if len(inner_axis_resources) < len(outer_axis_resources) + len(axis_resources): + overlap = set(outer_axis_resources) & set(axis_resources) + raise JAXTypeError( + f"Detected disallowed xmap axis name shadowing at " + f"{source_info_util.summarize(source_info)} " + f"(shadowed axes: {mesh_lib.show_axes(overlap)})") + + if resource_env.physical_mesh != params['resource_env'].physical_mesh: + raise RuntimeError("Changing the physical mesh is not allowed inside xmap.") + + call_jaxpr = params['call_jaxpr'] + pxla.resource_typecheck( + params['call_jaxpr'], resource_env, inner_axis_resources, + lambda: (f"an xmapped function {params['name']} " + + (f"(xmap called at {source_info_util.summarize(source_info)})" + if source_info else ""))) + + for v, axes in zip(call_jaxpr.outvars, params['out_axes']): + broadcast_axes = set(axes) - set(v.aval.named_shape) + used_resources = set(it.chain.from_iterable( + inner_axis_resources[a] for a in v.aval.named_shape)) + for baxis in broadcast_axes: + baxis_resources = set(inner_axis_resources[baxis]) + overlap = baxis_resources & used_resources + if overlap: + resource_to_axis = {} + for axis in v.aval.named_shape: + for raxis in inner_axis_resources[axis]: + resource_to_axis[raxis] = axis + partitioning_axes = {resource_to_axis[raxis] for raxis in overlap} + raise JAXTypeError( + f"One of xmapped function ({params['name']}) outputs is broadcast " + f"along axis `{baxis}` which is assigned to resources " + f"{mesh_lib.show_axes(baxis_resources)}, but the output is already " + f"partitioned along {mesh_lib.show_axes(overlap)}, because its " + f"named shape contains {mesh_lib.show_axes(partitioning_axes)}") +pxla.custom_resource_typing_rules[xmap_p] = _resource_typing_xmap + + +# This is DynamicJaxprTrace.process_map with some very minor modifications +def _dynamic_jaxpr_process_xmap(self, primitive, f, tracers, params): + assert primitive is xmap_p + in_avals = [t.aval for t in tracers] + global_axis_sizes = params['global_axis_sizes'] + mapped_in_avals = [_delete_aval_axes(a, a_in_axes, global_axis_sizes) + for a, a_in_axes in zip(in_avals, params['in_axes'])] + with core.extend_axis_env_nd(global_axis_sizes.items()): + with core.new_sublevel(): + jaxpr, mapped_out_avals, consts, () = trace_to_subjaxpr_dynamic( + f, self.main, mapped_in_avals) + out_axes = params['out_axes_thunk']() + if params['spmd_out_axes_thunk'] is not None: + spmd_out_axes = params['spmd_out_axes_thunk']() + else: + spmd_out_axes = None + axis_resource_count = _get_axis_resource_count( + params['axis_resources'], params['resource_env']) + local_axis_sizes = { + axis: axis_resource_count[axis].to_local(global_size) + for axis, global_size in global_axis_sizes.items() + } + out_avals = [_insert_aval_axes(a, a_out_axes, local_axis_sizes) + for a, a_out_axes in zip(mapped_out_avals, out_axes)] + # _check_out_avals_vs_out_axes(out_avals, out_axes, params['global_axis_sizes']) + source_info = source_info_util.current() + out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals] + invars = map(self.getvar, tracers) + constvars = map(self.getvar, map(self.instantiate_const, consts)) + outvars = map(self.makevar, out_tracers) + new_in_axes = (AxisNamePos(user_repr='{}'),) * len(consts) + params['in_axes'] + if params['spmd_in_axes'] is None: + new_spmd_in_axes = None + else: + new_spmd_in_axes = (None,) * len(consts) + params['spmd_in_axes'] + new_donated_invars = (False,) * len(consts) + params['donated_invars'] + with core.extend_axis_env_nd(global_axis_sizes.items()): + call_jaxpr = convert_constvars_jaxpr(jaxpr) + new_params = dict(params, in_axes=new_in_axes, out_axes=out_axes, + donated_invars=new_donated_invars, + spmd_in_axes=new_spmd_in_axes, + spmd_out_axes=spmd_out_axes, + call_jaxpr=call_jaxpr) + del new_params['out_axes_thunk'] + del new_params['spmd_out_axes_thunk'] + effs = core.filter_named_axis_effects(call_jaxpr.effects, global_axis_sizes) + eqn = new_jaxpr_eqn([*constvars, *invars], outvars, primitive, + new_params, effs, source_info) + self.frame.add_eqn(eqn) + return out_tracers +pe.DynamicJaxprTrace.process_xmap = _dynamic_jaxpr_process_xmap # type: ignore + +def _xmap_partial_eval_custom_params_updater( + unks_in: Sequence[bool], inst_in: Sequence[bool], + kept_outs_known: Sequence[bool], kept_outs_staged: Sequence[bool], + num_res: int, params_known: dict, params_staged: dict + ) -> tuple[dict, dict]: + assert params_known['spmd_in_axes'] is None is params_known['spmd_out_axes'] + assert params_staged['spmd_in_axes'] is None is params_staged['spmd_out_axes'] + + # prune inputs to jaxpr_known according to unks_in + donated_invars_known, _ = pe.partition_list(unks_in, params_known['donated_invars']) + in_axes_known, _ = pe.partition_list(unks_in, params_known['in_axes']) + if num_res == 0: + residual_axes = [] + else: + residual_axes = [ + AxisNamePos(zip(sort_named_shape, range(len(sort_named_shape))), + user_repr=f'') + for named_shape in (v.aval.named_shape for v in params_known['call_jaxpr'].outvars[:-num_res]) + # We sort here to make the iteration order deterministic + for sort_named_shape in [sorted(named_shape, key=str)] + ] + _, out_axes_known = pe.partition_list(kept_outs_known, params_known['out_axes']) + new_params_known = dict(params_known, + in_axes=tuple(in_axes_known), + out_axes=(*out_axes_known, *residual_axes), + donated_invars=tuple(donated_invars_known)) + assert len(new_params_known['in_axes']) == len(params_known['call_jaxpr'].invars) + assert len(new_params_known['out_axes']) == len(params_known['call_jaxpr'].outvars) + + # added num_res new inputs to jaxpr_staged, and pruning according to inst_in + _, donated_invars_staged = pe.partition_list(inst_in, params_staged['donated_invars']) + donated_invars_staged = [False] * num_res + donated_invars_staged + _, in_axes_staged = pe.partition_list(inst_in, params_staged['in_axes']) + in_axes_staged = [*residual_axes, *in_axes_staged] + _, out_axes_staged = pe.partition_list(kept_outs_staged, params_staged['out_axes']) + new_params_staged = dict(params_staged, in_axes=tuple(in_axes_staged), + out_axes=tuple(out_axes_staged), + donated_invars=tuple(donated_invars_staged)) + assert len(new_params_staged['in_axes']) == len(params_staged['call_jaxpr'].invars) + assert len(new_params_staged['out_axes']) == len(params_staged['call_jaxpr'].outvars) + return new_params_known, new_params_staged +pe.partial_eval_jaxpr_custom_rules[xmap_p] = \ + partial(pe.call_partial_eval_custom_rule, 'call_jaxpr', + _xmap_partial_eval_custom_params_updater) + + +@lu.transformation_with_aux +def out_local_named_shapes(local_axes, *args, **kwargs): + ans = yield args, kwargs + ans_axes = [frozenset(a.aval.named_shape) & local_axes for a in ans] + yield ans, ans_axes + + +def _jaxpr_trace_process_xmap(self, primitive, f: lu.WrappedFun, tracers, params): + assert primitive is xmap_p + assert params['spmd_out_axes_thunk'] is params['spmd_in_axes'] is None + in_axes = params['in_axes'] + donated_invars = params['donated_invars'] + global_axis_sizes = params['global_axis_sizes'] + out_axes_thunk = params['out_axes_thunk'] + + # Adjust input tracers' pvals for mapped axes, and unpack. + in_pvals = [t.pval if t.pval.is_known() else + pe.PartialVal.unknown( + _delete_aval_axes(t.pval.get_aval(), axes, global_axis_sizes)) + for t, axes in zip(tracers, in_axes)] + in_knowns, in_avals, in_consts = pe.partition_pvals(in_pvals) + + # Wrap f to perform partial evaluation, and plumb out aux data. + f = pe.trace_to_subjaxpr_nounits(f, self.main, False) + f, aux = pe.partial_eval_wrapper_nounits(f, tuple(in_knowns), tuple(in_avals)) + # Also grab the local named shapes of the output (known and res). + f, out_named_shapes = out_local_named_shapes(f, frozenset(global_axis_sizes)) + + # Adjust params for knowns (donated_invars, in_axes, out_axes_thunk). + out_axes = None # cache this to avoid calling out_axes_thunk() more than once + + @as_hashable_function(closure=out_axes_thunk) + def new_out_axes_thunk(): + nonlocal out_axes + out_axes = out_axes_thunk() + out_knowns, _, _, _ = aux() + _, out_axes_known = partition_list(out_knowns, out_axes) + return (*out_axes_known, *res_axes()) + def res_axes(): + _, _, jaxpr_unknown, _ = aux() + num_res = len(jaxpr_unknown.constvars) + res_named_shapes = out_named_shapes()[-num_res:] if num_res else [] + sorted_named_shapes = [sorted(ns, key=str) for ns in res_named_shapes] + return [AxisNamePos(zip(named_shape, range(len(named_shape))), + user_repr=f'') + for named_shape in sorted_named_shapes] + known_params = dict( + params, in_axes=tuple(a for a, k in zip(in_axes, in_knowns) if k), + donated_invars=tuple(d for d, k in zip(donated_invars, in_knowns) if k), + out_axes_thunk=new_out_axes_thunk) + + # Run the known part. + out = primitive.bind(f, *in_consts, **known_params) + out_knowns, out_avals, jaxpr_unknown, env = aux() + known_outvals, res = split_list(out, [len(out)-len(jaxpr_unknown.constvars)]) + with core.extend_axis_env_nd(global_axis_sizes.items()): + jaxpr_unknown = pe.convert_constvars_jaxpr(jaxpr_unknown) + + # Set up new params. + if out_axes is None: + out_axes = out_axes_thunk() # new_out_axes_thunk may have set during bind + out_axes_unknown = [a for a, k in zip(out_axes, out_knowns) if not k] + unknown_params = dict( + params, call_jaxpr=jaxpr_unknown, out_axes=tuple(out_axes_unknown), + spmd_out_axes=None, + donated_invars=(*(False for _ in res), + *(d for d, k in zip(donated_invars, in_knowns) if not k)), + in_axes=(*res_axes(), *(None for _ in env), + *(a for a, k in zip(in_axes, in_knowns) if not k))) + del unknown_params['out_axes_thunk'] + del unknown_params['spmd_out_axes_thunk'] + # Create input tracers for unknown part. + res_tracers = map(self.new_instantiated_const, res) + env_tracers = map(self.full_raise, env) + unknown_arg_tracers = [t for t in tracers if not t.pval.is_known()] + # Create output tracers for unknown part, adjusting avals. + axis_resource_count = _get_axis_resource_count( + params['axis_resources'], params['resource_env']) + local_axis_sizes = { + ax: axis_resource_count[ax].to_local(global_size) + for ax, global_size in global_axis_sizes.items()} + out_pvals = [pe.PartialVal.unknown(_insert_aval_axes(a, ax, local_axis_sizes)) + for a, ax in zip(out_avals, out_axes_unknown)] + unknown_tracers_out = [pe.JaxprTracer(self, pval, None) for pval in out_pvals] + # Build eqn to be staged out and attach it to unknown output tracers. + eqn = pe.new_eqn_recipe((*res_tracers, *env_tracers, *unknown_arg_tracers), + unknown_tracers_out, primitive, unknown_params, + jaxpr_unknown.effects, source_info_util.current()) + for t in unknown_tracers_out: t.recipe = eqn + return merge_lists(out_knowns, unknown_tracers_out, known_outvals) +pe.JaxprTrace.process_xmap = _jaxpr_trace_process_xmap + +def _batch_trace_update_spmd_axes( + spmd_in_axes, spmd_out_axes_thunk, + axis_name, dims, dims_out_thunk): + """Extends spmd in and out axes with the position of the trace's batch dimension.""" + not_mapped = batching.not_mapped + def insert_spmd_axis(axes, nd): + too_short = nd - len(axes) + if too_short > 0: + axes += (None,) * too_short + return tuple_insert(axes, nd, axis_name) + + if spmd_in_axes is None: + spmd_in_axes = ((),) * len(dims) + new_spmd_in_axes = tuple( + spmd_axes if d is not_mapped else insert_spmd_axis(spmd_axes, d) + for spmd_axes, d in zip(spmd_in_axes, dims)) + + @as_hashable_function(closure=spmd_out_axes_thunk) + def new_spmd_out_axes_thunk(): + dims_out = dims_out_thunk() + if spmd_out_axes_thunk is None: + spmd_out_axes = ((),) * len(dims_out) + else: + spmd_out_axes = spmd_out_axes_thunk() + return tuple( + spmd_out_axes if nd is not_mapped else insert_spmd_axis(spmd_out_axes, nd) + for spmd_out_axes, nd in zip(spmd_out_axes, dims_out)) + + return new_spmd_in_axes, new_spmd_out_axes_thunk + +def _axis_after_insertion(axis, inserted_named_axes): + for inserted_axis in sorted(inserted_named_axes.values()): + if inserted_axis >= axis: + break + axis += 1 + return axis + +def _fmap_dims(axes, f): + return AxisNamePos(((name, f(axis)) for name, axis in axes.items()), + user_repr=axes.user_repr) + +def _batch_trace_process_xmap(self, is_spmd, primitive, f: lu.WrappedFun, tracers, params): + not_mapped = batching.not_mapped + vals, dims = unzip2((t.val, t.batch_dim) for t in tracers) + assert primitive is xmap_p + if not is_spmd and all(dim is not_mapped for dim in dims): + return primitive.bind(f, *vals, **params) + else: + assert len({x.shape[d] for x, d in zip(vals, dims) if d is not not_mapped}) == 1 + new_in_axes = tuple( + _fmap_dims(in_axes, lambda a: a + (d is not not_mapped and d <= a)) + for d, in_axes in zip(dims, params['in_axes'])) + mapped_dims_in = tuple( + d if d is not_mapped else d - sum(a < d for a in in_axis.values()) + for d, in_axis in zip(dims, params['in_axes'])) + f, mapped_dims_out = batching.batch_subtrace(f, self.main, mapped_dims_in) + out_axes_thunk: Callable[[], Sequence[AxisNamePos]] = params['out_axes_thunk'] + dims_out_thunk = lambda: tuple(d if d is not_mapped else _axis_after_insertion(d, out_axes) + for d, out_axes in zip(mapped_dims_out(), out_axes_thunk())) + # NOTE: This assumes that the choice of the dimensions over which outputs + # are batched is entirely dependent on the function and not e.g. on the + # data or its shapes. + @as_hashable_function(closure=out_axes_thunk) + def new_out_axes_thunk(): + return tuple( + out_axes if d is not_mapped else + _fmap_dims(out_axes, lambda a, nd=_axis_after_insertion(d, out_axes): a + (nd <= a)) + for out_axes, d in zip(out_axes_thunk(), mapped_dims_out())) + + if not is_spmd: + assert params['spmd_in_axes'] is None and params['spmd_out_axes_thunk'] is None + new_spmd_in_axes = None + new_spmd_out_axes_thunk = None + else: + new_spmd_in_axes, new_spmd_out_axes_thunk = _batch_trace_update_spmd_axes( + params['spmd_in_axes'], params['spmd_out_axes_thunk'], + self.axis_name, dims, dims_out_thunk) + + new_params = dict(params, + in_axes=new_in_axes, out_axes_thunk=new_out_axes_thunk, + spmd_in_axes=new_spmd_in_axes, + spmd_out_axes_thunk=new_spmd_out_axes_thunk) + vals_out = primitive.bind(f, *vals, **new_params) + dims_out = dims_out_thunk() + return [batching.BatchTracer(self, v, d) for v, d in zip(vals_out, dims_out)] +batching.BatchTrace.process_xmap = partialmethod(_batch_trace_process_xmap, False) # type: ignore +pxla.SPMDBatchTrace.process_xmap = partialmethod(_batch_trace_process_xmap, True) # type: ignore + + +def _batch_trace_post_process_xmap(self, primitive, out_tracers, params): + not_mapped = batching.not_mapped + BT = batching.BatchTracer + vals, dims, srcs = unzip3((t.val, t.batch_dim, t.source_info) for t in out_tracers) + main = self.main + def todo(vals): + trace = main.with_cur_sublevel() + return [BT(trace, v, d if d is not_mapped else _axis_after_insertion(d, oa), s) + for v, d, oa, s in zip(vals, dims, params['out_axes_thunk'](), srcs)] + def out_axes_transform(out_axes): + return tuple(oa if d is not_mapped else + _fmap_dims(oa, lambda a, nd=_axis_after_insertion(d, oa): a + (nd <= a)) + for oa, d in zip(out_axes, dims)) + return vals, (todo, out_axes_transform) +batching.BatchTrace.post_process_xmap = _batch_trace_post_process_xmap + + +# -------- nested xmap handling -------- + +def _xmap_lowering_rule(ctx, *args, **kwargs): + if isinstance(ctx.module_context.axis_context, sharding_impls.SPMDAxisContext): + if SPMD_LOWERING_MANUAL.value: + return _xmap_lowering_rule_spmd_manual(ctx, *args, **kwargs) + else: + return _xmap_lowering_rule_spmd(ctx, *args, **kwargs) + # Here ShardingContext is used in place of ReplicaAxisContext because when + # axis_resources and mesh is not used with xmap, `make_xmap_callable` will + # go via `dispatch.sharded_lowering` path which sets the context to + # ShardingContext. sharding_impls.ShardingContext is not used for SPMD. + elif isinstance(ctx.module_context.axis_context, + (sharding_impls.ReplicaAxisContext, sharding_impls.ShardingContext)): + return _xmap_lowering_rule_replica(ctx, *args, **kwargs) + else: + raise AssertionError("Unrecognized axis context type!") +mlir.register_lowering(xmap_p, _xmap_lowering_rule) + +def _xmap_lowering_rule_replica(ctx, *in_nodes, + call_jaxpr, name, + in_axes, out_axes, donated_invars, + global_axis_sizes, + spmd_in_axes, spmd_out_axes, + axis_resources, resource_env, backend): + mlir.check_backend_matches(backend, ctx.module_context.platforms) + del backend, donated_invars + # The only way for any of those two assertions to be violated is when xmap + # is using the SPMD lowering, but then this rule shouldn't even trigger. + assert spmd_in_axes is None and spmd_out_axes is None + plan = EvaluationPlan.from_axis_resources( + axis_resources, resource_env, global_axis_sizes) + + axis_resource_count = _get_axis_resource_count( + axis_resources, resource_env) + if any(resource_count.distributed for resource_count in axis_resource_count.values()): + raise NotImplementedError + + mesh = resource_env.physical_mesh + mesh_in_axes, mesh_out_axes = plan.to_mesh_axes(in_axes, out_axes) + + local_avals = [pxla.tile_aval_nd( + mesh.shape, aval_mesh_in_axes, + _insert_aval_axes(v.aval, aval_in_axes, global_axis_sizes)) + for v, aval_in_axes, aval_mesh_in_axes + in zip(call_jaxpr.invars, in_axes, mesh_in_axes)] + # We have to substitute before tracing, because we want the vectorized + # axes to be used in the jaxpr. + resource_call_jaxpr = plan.subst_axes_with_resources(call_jaxpr) + f = lu.wrap_init(core.jaxpr_as_fun(core.ClosedJaxpr(resource_call_jaxpr, ()))) + f = hide_mapped_axes(f, tuple(in_axes), tuple(out_axes)) + f = plan.vectorize_and_loop(f, in_axes, out_axes) + # NOTE: We don't extend the resource env with the mesh shape, because those + # resources are already in scope! It's the outermost xmap that introduces + # them! + vectorized_jaxpr, out_avals, consts, () = pe.trace_to_jaxpr_dynamic(f, local_avals) + # _check_out_avals_vs_out_axes(out_avals, out_axes, global_axis_sizes) + const_nodes = [mlir.ir_constants(xla.canonicalize_dtype(x)) for x in consts] + + local_mesh_shape = mesh.local_mesh.shape + tiled_ins = ( + mlir.lower_fun(partial(_tile, in_axes=arg_in_axes, + axis_sizes=local_mesh_shape), + multiple_results=False)( + ctx.replace(primitive=None, + avals_in=[aval], avals_out=None), + in_node)[0] + for v, aval, in_node, arg_in_axes + in zip(call_jaxpr.invars, ctx.avals_in, in_nodes, mesh_in_axes)) + + # NOTE: We don't extend the resource env with the mesh shape, because those + # resources are already in scope! It's the outermost xmap that introduces + # them! + # We in-line here rather than generating a Call HLO as in the xla_call + # translation rule just because the extra tuple stuff is a pain. + name_stack = ctx.name_stack.extend(wrap_name(name, 'xmap')) + if any(effects.ordered_effects.contains(eff) for eff + in vectorized_jaxpr.effects): + raise NotImplementedError('Cannot lower `xmap` with ordered effects.') + tiled_outs, _ = mlir.jaxpr_subcomp(ctx.module_context, vectorized_jaxpr, + name_stack, mlir.TokenSet(), + const_nodes, *tiled_ins, + dim_var_values=ctx.dim_var_values) + + outs = [ + mlir.lower_fun( + partial(_untile, out_axes=ans_out_axes, axis_sizes=local_mesh_shape), + multiple_results=False)( + ctx.replace(primitive=None, + avals_in=[vectorized_outvar.aval], + avals_out=None), tiled_out)[0] + for v, vectorized_outvar, tiled_out, ans_out_axes + in zip(call_jaxpr.outvars, vectorized_jaxpr.outvars, tiled_outs, + mesh_out_axes)] + return outs + + +def _xmap_lowering_rule_spmd(ctx, *global_in_nodes, + call_jaxpr, name, in_axes, out_axes, + donated_invars, global_axis_sizes, spmd_in_axes, + spmd_out_axes, axis_resources, + resource_env, backend): + mlir.check_backend_matches(backend, ctx.module_context.platforms) + del backend, donated_invars + plan = EvaluationPlan.from_axis_resources( + axis_resources, resource_env, global_axis_sizes) + + resource_call_jaxpr = plan.subst_axes_with_resources(call_jaxpr) + f = lu.wrap_init(core.jaxpr_as_fun(core.ClosedJaxpr(resource_call_jaxpr, ()))) + f = hide_mapped_axes(f, in_axes, out_axes) + f = plan.vectorize_and_loop(f, in_axes, out_axes) + mesh_in_axes, mesh_out_axes = plan.to_mesh_axes(in_axes, out_axes) + mesh = resource_env.physical_mesh + f = pxla.vtile_by_mesh(f, mesh, mesh_in_axes, mesh_out_axes) + + # XXX: We modify mesh_in_axes and mesh_out_axes here + def add_spmd_axes( + flat_mesh_axes: Sequence[ArrayMapping], + flat_extra_axes: Sequence[Sequence[Sequence[MeshAxisName]]] | None): + if flat_extra_axes is None: + return + for axes, extra in zip(flat_mesh_axes, flat_extra_axes): + for dim, dim_extra_axis in enumerate(extra): + if dim_extra_axis is None: continue + assert dim_extra_axis not in axes + assert not config.enable_checks.value or all(v != dim for v in axes.values()) + axes[dim_extra_axis] = dim + add_spmd_axes(mesh_in_axes, spmd_in_axes) + add_spmd_axes(mesh_out_axes, spmd_out_axes) + global_in_avals = ctx.avals_in + vectorized_jaxpr, global_out_avals, consts, () = pe.trace_to_jaxpr_dynamic(f, global_in_avals) + + sharded_global_in_nodes = [ + [mlir.wrap_with_sharding_op( + ctx, node, aval, + NamedSharding(mesh, array_mapping_to_axis_resources(aval_axes) + )._to_xla_hlo_sharding(aval.ndim).to_proto())] + if aval_axes else [node] + for node, aval, aval_axes in zip(global_in_nodes, global_in_avals, mesh_in_axes) + ] + const_nodes = [mlir.ir_constants(xla.canonicalize_dtype(x)) for x in consts] + + # We in-line here rather than generating a Call HLO as in the xla_call + # translation rule just because the extra tuple stuff is a pain. + name_stack = ctx.name_stack.extend(wrap_name(name, 'xmap')) + if any(effects.ordered_effects.contains(eff) for eff + in vectorized_jaxpr.effects): + raise NotImplementedError('Cannot lower `xmap` with ordered effects.') + global_out_nodes, _ = mlir.jaxpr_subcomp( + ctx.module_context, vectorized_jaxpr, name_stack, mlir.TokenSet(), + const_nodes, *sharded_global_in_nodes, dim_var_values=ctx.dim_var_values) + + sharded_global_out_nodes = [ + mlir.wrap_with_sharding_op( + ctx, node, aval, + NamedSharding(mesh, array_mapping_to_axis_resources(aval_axes) + )._to_xla_hlo_sharding(aval.ndim).to_proto()) + if aval_axes else node + for (node,), aval, aval_axes in zip(global_out_nodes, global_out_avals, mesh_out_axes) + ] + + return sharded_global_out_nodes + + +def _xmap_lowering_rule_spmd_manual(ctx, *global_in_nodes, + call_jaxpr, name, in_axes, out_axes, + donated_invars, global_axis_sizes, spmd_in_axes, + spmd_out_axes, axis_resources, + resource_env, backend): + del donated_invars + assert spmd_in_axes is None and spmd_out_axes is None + # This first part (up to vtile_manual) is shared with non-MANUAL SPMD rule. + mlir.check_backend_matches(backend, ctx.module_context.platforms) + plan = EvaluationPlan.from_axis_resources( + axis_resources, resource_env, global_axis_sizes) + manual_mesh_axes = frozenset(it.chain.from_iterable(plan.physical_axis_resources.values())) + + resource_call_jaxpr = plan.subst_axes_with_resources(call_jaxpr) + f = lu.wrap_init(core.jaxpr_as_fun(core.ClosedJaxpr(resource_call_jaxpr, ()))) + f = hide_mapped_axes(f, in_axes, out_axes) + f = plan.vectorize_and_loop(f, in_axes, out_axes) + + # NOTE: Sharding constraints are handled entirely by vtile_manual! + mesh_in_axes, mesh_out_axes = plan.to_mesh_axes(in_axes, out_axes) + mesh = resource_env.physical_mesh + f = pxla.vtile_manual(f, tuple(manual_mesh_axes), mesh, mesh_in_axes, mesh_out_axes) + + # NOTE: We don't extend the resource env with the mesh shape, because those + # resources are already in scope! It's the outermost xmap that introduces + # them! + global_in_avals = ctx.avals_in + vectorized_jaxpr, global_out_avals, consts, () = pe.trace_to_jaxpr_dynamic(f, global_in_avals) + const_nodes = [mlir.ir_constants(xla.canonicalize_dtype(x)) for x in consts] + + # We in-line here rather than generating a Call HLO as in the xla_call + # translation rule just because the extra tuple stuff is a pain. + assert isinstance(ctx.module_context.axis_context, + sharding_impls.SPMDAxisContext) + name_stack = ctx.name_stack.extend(wrap_name(name, 'xmap')) + sub_ctx = ctx.module_context.replace( + axis_context=ctx.module_context.axis_context.extend_manual(manual_mesh_axes)) + if any(effects.ordered_effects.contains(eff) for eff + in vectorized_jaxpr.effects): + raise NotImplementedError('Cannot lower `xmap` with ordered effects.') + global_out_nodes, _ = mlir.jaxpr_subcomp( + sub_ctx, vectorized_jaxpr, name_stack, + mlir.TokenSet(), const_nodes, *([n] for n in global_in_nodes), + dim_var_values=ctx.dim_var_values) + + return global_out_nodes + + +def _tile_base_indices(tile_shape, axes, axis_sizes): + zero = np.zeros((), dtype=np.int32) + linear_idxs = [zero] * len(tile_shape) + strides = [1] * len(tile_shape) + for name, axis in reversed(axes.items()): + axis_index = lax.axis_index(name) + stride_c = np.array(strides[axis], np.int32) + if linear_idxs[axis] is zero and strides[axis] == 1: + linear_idxs[axis] = axis_index + else: + linear_idxs[axis] = lax.add(linear_idxs[axis], + lax.mul(axis_index, stride_c)) + strides[axis] *= axis_sizes[name] + return [zero if linear_idx is zero else + lax.mul(linear_idx, np.array(tile_dim_size, np.int32)) + for linear_idx, tile_dim_size in zip(linear_idxs, tile_shape)] + + +def _tile(x, in_axes, axis_sizes): + if not in_axes: + return x + tile_shape = list(x.shape) + for name, axis in in_axes.items(): + axis_size = axis_sizes[name] + assert tile_shape[axis] % axis_size == 0 + tile_shape[axis] //= axis_size + base_idxs = _tile_base_indices(tile_shape, in_axes, axis_sizes) + return lax.dynamic_slice(x, base_idxs, tile_shape) + + +# TODO(b/110096942): more efficient gather +def _untile(x, out_axes, axis_sizes): + tile_shape = list(x.shape) + shape = list(tile_shape) + for name, axis in out_axes.items(): + shape[axis] *= axis_sizes[name] + base_idxs = _tile_base_indices(tile_shape, out_axes, axis_sizes) + + padded = lax.broadcast(np.array(0, x.dtype), shape) + padded = lax.dynamic_update_slice(padded, x, base_idxs) + out = lax.psum(padded, tuple(out_axes.keys())) + return out + + +# -------- helper functions -------- + +def _delete_aval_axes(aval, axes: AxisNamePos, global_axis_sizes): + assert isinstance(aval, core.ShapedArray) + shape = list(aval.shape) + named_shape = dict(aval.named_shape) + for name, dim in sorted(axes.items(), key=lambda x: x[1], reverse=True): + named_shape[name] = global_axis_sizes[name] + del shape[dim] + return aval.update(shape=tuple(shape), named_shape=named_shape) + +def _insert_aval_axes(aval, axes: AxisNamePos, local_axis_sizes): + assert isinstance(aval, core.ShapedArray) + shape = list(aval.shape) + named_shape = dict(aval.named_shape) + for name, dim in sorted(axes.items(), key=lambda x: x[1]): + shape.insert(dim, local_axis_sizes[name]) + named_shape.pop(name, None) # The name might be missing --- it's a broadcast. + return aval.update(shape=tuple(shape), named_shape=named_shape) + + +class ResourceCount(NamedTuple): + nglobal: int + nlocal: int | None + distributed: bool + + def to_local(self, global_size): + return global_size + + +def _get_axis_resource_count( + axis_resources, resource_env) -> dict[ResourceAxisName, ResourceCount]: + global_res_shape = resource_env.shape + local_res_shape = None + + distributed = (False if resource_env.physical_mesh.empty else + resource_env.physical_mesh.size != len(resource_env.physical_mesh.local_devices)) + resource_count_map = {} + for axis, resources in axis_resources.items(): + if local_res_shape is None: + nlocal = None + else: + nlocal = math.prod(map(local_res_shape.get, resources)) + resource_count_map[axis] = ResourceCount( + math.prod(map(global_res_shape.get, resources)), + nlocal, distributed) + return resource_count_map + + +def _get_axis_sizes(args_flat: Iterable[Any], + in_axes_flat: Iterable[AxisNamePos], + global_axis_sizes: dict[AxisName, int], + axis_resource_count: dict[AxisName, ResourceCount]): + global_axis_sizes = dict(global_axis_sizes) + for arg, in_axes in zip(args_flat, in_axes_flat): + for name, dim in in_axes.items(): + try: + dim_size = arg.shape[dim] + except IndexError: + # TODO(apaszke): Handle negative indices. Check for overlap too! + raise ValueError(f"One of xmap arguments has an in_axes specification of " + f"{in_axes.user_repr}, which implies that it has at least " + f"{max(in_axes.values()) + 1} dimensions, but the argument " + f"has rank {arg.ndim}") + global_dim_size = dim_size + if name in global_axis_sizes: + expected_global_dim_size = global_axis_sizes[name] + if global_dim_size != expected_global_dim_size: + raise ValueError(f"The size of axis {name} was previously inferred to be " + f"{expected_global_dim_size}, but found an argument of shape {arg.shape} " + f"with in_axes specification {in_axes.user_repr}. Shape mismatch " + f"occurs in dimension {dim}: {global_dim_size} != {expected_global_dim_size}") + global_axis_sizes[name] = global_dim_size + return FrozenDict(global_axis_sizes) + + +@lu.transformation +def hide_mapped_axes(flat_in_axes, flat_out_axes, *flat_args): + def _squeeze_mapped_axes(arg, axes: AxisNamePos): + for dim in sorted(axes.values(), reverse=True): + arg = arg.squeeze(dim) + return arg + + def _unsqueeze_mapped_axes(out, axes: AxisNamePos): + try: + return jnp.expand_dims(out, tuple(axes.values())) + except ValueError as e: + # Improve the axis out of bounds errors + # TODO(apaszke): Handle negative indices. Check for overlap too! + if e.args[0].startswith('axis') and 'out of bounds' in e.args[0]: + raise ValueError(f"One of xmap outputs has an out_axes specification of " + f"{axes.user_repr}, which requires the result of the xmapped " + f"function to have at least {max(axes.values()) - len(axes) + 1} " + f"positional dimensions, but it only has {out.ndim}") + raise + + squeezed_args = map(_squeeze_mapped_axes, flat_args, flat_in_axes) + flat_outputs = yield squeezed_args, {} + yield map(_unsqueeze_mapped_axes, flat_outputs, flat_out_axes) + + +def _jaxpr_resources(jaxpr, resource_env) -> set[ResourceAxisName]: + if isinstance(jaxpr, core.ClosedJaxpr): + jaxpr = jaxpr.jaxpr + assert isinstance(jaxpr, core.Jaxpr) + used_resources = set() + for eqn in jaxpr.eqns: + if eqn.primitive is xmap_p: + if eqn.params['resource_env'].physical_mesh != resource_env.physical_mesh: + raise RuntimeError("Changing the physical mesh is not allowed inside xmap.") + used_resources |= set(it.chain(*eqn.params['axis_resources'].values())) + updates = core.traverse_jaxpr_params( + partial(_jaxpr_resources, resource_env=resource_env), eqn.params).values() + for update in updates: + used_resources |= update + return used_resources + + +def _to_resource_axes(axes_specs: Sequence[AxisNamePos], + axis_resources: dict[AxisName, tuple[ResourceAxisName, ...]]): + """ + Convert in/out_axes parameters ranging over logical dimensions to + ones that range over resource dimensions. + + Note that values no longer have to be distinct, as multiple resource + axes can tile a single positional axes. This is why the result is + an OrderedDict with an implicit major-to-minor ordering. + """ + return tuple(OrderedDict((resource_axis, pos_axis) + for logical_axis, pos_axis in axes.items() + for resource_axis in axis_resources[logical_axis]) + for axes in axes_specs) + + +def _merge_leading_axis(x, axis: int | None): + if axis is None: + # We assume that the output does not vary along the leading axis + return lax.index_in_dim(x, 0, axis=0, keepdims=False) + else: + x_moved = moveaxis(x, 0, axis) + shape = list(x_moved.shape) + shape[axis:axis + 2] = [shape[axis] * shape[axis + 1]] + return x_moved.reshape(shape) + + +def _slice_tile(x, dim: int | None, i, n: int): + """Selects an `i`th (out of `n`) tiles of `x` along `dim`.""" + if dim is None: return x + (tile_size, rem) = divmod(x.shape[dim], n) + assert rem == 0, "Please open a bug report!" + return lax.dynamic_slice_in_dim(x, i * tile_size, slice_size=tile_size, axis=dim) + + +def _unzip_axis_resources(axis_resources: dict[AxisName, tuple[ResourceAxisName, ...]], + resource_env: ResourceEnv): + """Splits axis_resources into separate dicts for physical and loop resources.""" + physical_axis_resources = {} + loop_axis_resources = {} + loop_resource_axes = resource_env.loop_resource_axes + for axis, raxes in axis_resources.items(): + first_loop = 0 + for raxis in raxes: + if raxis in loop_resource_axes: + break + else: + first_loop += 1 + physical_axis_resources[axis] = raxes[:first_loop] + loop_resources = loop_axis_resources[axis] = raxes[first_loop:] + if not all(name in loop_resource_axes for name in loop_resources): + raise NotImplementedError("Loop resources cannot appear before mesh axes " + "in the resource_axis argument") + return physical_axis_resources, loop_axis_resources + + +def _check_out_avals_vs_out_axes(out_avals: Sequence[core.AbstractValue], + out_axes: Sequence[AxisNamePos], + global_axis_sizes: dict[AxisName, int]): + defined_axes = set(global_axis_sizes) + for aval, axes in zip(out_avals, out_axes): + if not isinstance(aval, core.ShapedArray): + if axes: + raise AssertionError(f"Only array abstract values can have non-empty " + f"out_axes, but {aval} has {axes}") + continue + undeclared_axes = (set(aval.named_shape) - set(axes)) & defined_axes + if undeclared_axes: + undeclared_axes_str = sorted(str(axis) for axis in undeclared_axes) + raise TypeError(f"One of xmap results has an out_axes specification of " + f"{axes.user_repr}, but is actually mapped along more axes " + f"defined by this xmap call: {', '.join(undeclared_axes_str)}") + + +def _check_gda_or_array_xmap_partitioning(axis_resources, resource_env, + global_axis_sizes, in_axes_flat, + args_flat): + @lru_cache + def _check_sharding(in_sharding, xmap_sharding, ndim, arr_flavor): + if (not op_shardings.are_op_shardings_equal( + in_sharding._to_xla_hlo_sharding(ndim), + xmap_sharding._to_xla_hlo_sharding(ndim)) or + in_sharding.memory_kind != xmap_sharding.memory_kind): + raise ValueError( + f"Got an input {arr_flavor} to xmap with different partitioning than " + "specified in xmap. The partitioning must match. " + f"Got {arr_flavor} spec: {in_sharding.spec} and " + f"xmap spec: {xmap_sharding.spec}") + + mesh_in_axes = EvaluationPlan.from_axis_resources( # pytype: disable=wrong-arg-types # always-use-return-annotations + axis_resources, resource_env, global_axis_sizes).to_mesh_axes(in_axes_flat) + for arg, xmap_array_mapping in safe_zip(args_flat, mesh_in_axes): + if isinstance(arg, ArrayImpl): + if not isinstance(arg.sharding, NamedSharding): + continue + mesh = arg.sharding.mesh + if mesh != resource_env.physical_mesh: + raise ValueError("xmap's mesh and Array's mesh should be equal. " + f"Got xmap mesh: {resource_env.physical_mesh},\n" + f"Array mesh: {mesh}") + + s = arg.sharding + xmap_sharding = pxla.create_mesh_pspec_sharding( + mesh, array_mapping_to_axis_resources(xmap_array_mapping)) + # This check is cached because comparing OpSharding is expensive during + # dispatch and if the shardings are the same, then there is no need to + # compare twice. + _check_sharding(s, xmap_sharding, arg.ndim, 'Array') + + +# TODO: We should relax this at least for "constructor primitives" +# such as axis_index or zeros. +def _check_no_loop_collectives(jaxpr, loop_axis_resources): + if isinstance(jaxpr, core.ClosedJaxpr): + jaxpr = jaxpr.jaxpr + def subst_no_loop(name): + if loop_axis_resources.get(name, ()): + raise RuntimeError(f"Named axes with loop resources assigned to them cannot " + f"be referenced inside the xmapped computation (e.g. in " + f"collectives), but `{name}` violates that rule") + return (name,) + for eqn in jaxpr.eqns: + core.subst_axis_names(eqn.primitive, eqn.params, subst_no_loop, traverse=False) + rec = partial(_check_no_loop_collectives, loop_axis_resources=loop_axis_resources) + core.traverse_jaxpr_params(rec, eqn.params) + + +def _fix_inferred_spmd_sharding(jaxpr, resource_env, gen_fresh_name = None): + rec = lambda jaxpr: _fix_inferred_spmd_sharding(jaxpr, resource_env, gen_fresh_name) + if isinstance(jaxpr, core.ClosedJaxpr): + return jaxpr.map_jaxpr(rec) + assert isinstance(jaxpr, core.Jaxpr) + if gen_fresh_name is None: + gen_fresh_name = core.gensym() + new_eqns = [] + for eqn in jaxpr.eqns: + new_jaxpr_params = core.traverse_jaxpr_params(rec, eqn.params) + tmp_outvars = [gen_fresh_name(v.aval) for v in eqn.outvars] + new_eqns.append(eqn.replace( + outvars=tmp_outvars, params=dict(eqn.params, **new_jaxpr_params))) + for outvar, tmpvar in zip(eqn.outvars, tmp_outvars): + mps = NamedSharding._from_parsed_pspec( + resource_env.physical_mesh, ParsedPartitionSpec((), ())) + unconstrained_dims = get_unconstrained_dims(mps) + gspmd_sharding = GSPMDSharding.get_replicated(mps._device_assignment) + new_eqns.append(core.JaxprEqn( + [tmpvar], [outvar], sharding_constraint_p, + dict(resource_env=resource_env, + sharding=gspmd_sharding, + unconstrained_dims=unconstrained_dims), + set(), + eqn.source_info, eqn.ctx)) + return jaxpr.replace(eqns=new_eqns) + +def _flatten_axes(what, tree, axes, tupled_args): + try: + return tuple(flatten_axes(what, tree, axes, tupled_args=tupled_args)) + except ValueError: + pass + # Replace axis_resources with unparsed versions to avoid revealing internal details + flatten_axes(what, tree, tree_map(lambda parsed: NoQuotesStr(parsed.user_repr), axes), + tupled_args=tupled_args) + raise AssertionError("Please open a bug request!") # This should be unreachable + +class NoQuotesStr(str): + __repr__ = str.__str__ + + +# -------- config flags -------- + +def _thread_local_flag_unsupported(_): + raise RuntimeError("thread-local xmap flags not supported!") +def _clear_compilation_cache(_): + make_xmap_callable.cache_clear() # pytype: disable=attribute-error + +def _ensure_spmd_and(f): + def update(v): + if v and not SPMD_LOWERING.value: + raise RuntimeError("This flag requires enabling the experimental_xmap_spmd_lowering flag") + return f(v) + return update + + +SPMD_LOWERING = config.define_bool_state( + name="experimental_xmap_spmd_lowering", + default=False, + help=("When set, multi-device xmap computations will be compiled through " + "the XLA SPMD partitioner instead of explicit cross-replica collectives. " + "Not supported on CPU!"), + update_global_hook=_clear_compilation_cache, + update_thread_local_hook=_thread_local_flag_unsupported) +SPMD_LOWERING_MANUAL = config.define_bool_state( + name="experimental_xmap_spmd_lowering_manual", + default=False, + help=("When set, multi-device xmap computations will be compiled using " + "the MANUAL partitioning feature of the XLA SPMD partitioner instead of " + "sharding constraints on vectorized code. " + "Requires experimental_xmap_spmd_lowering!"), + update_global_hook=_ensure_spmd_and(_clear_compilation_cache), + update_thread_local_hook=_thread_local_flag_unsupported) +_ENSURE_FIXED_SHARDING = config.define_bool_state( + name="experimental_xmap_ensure_fixed_sharding", + default=False, + help=("When set and `experimental_xmap_spmd_lowering` is enabled, the lowering will " + "try to limit the flexibility of the automated SPMD partitioner heuristics " + "by emitting additional sharding annotations for program intermediates."), + update_global_hook=_ensure_spmd_and(_clear_compilation_cache), + update_thread_local_hook=_thread_local_flag_unsupported) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/mesh.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..7339bd6cf7c16345328615c1e8bfa9f8eaec13e3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/mesh.py @@ -0,0 +1,328 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Definitions of Mesh and ResourceEnv.""" + +from __future__ import annotations + +import collections +from collections.abc import Hashable, Sequence +import contextlib +import functools +import math +import threading +from typing import Any, NamedTuple + +import numpy as np + +from jax._src import config as jax_config +from jax._src import xla_bridge as xb +from jax._src import util +from jax._src.lib import xla_client as xc + + +MeshAxisName = Any +ResourceAxisName = Hashable + +class Loop(NamedTuple): + name: ResourceAxisName + length: int + + +def show_axes(axes): + return ", ".join(sorted(f"`{a}`" for a in axes)) + + +class ResourceEnv(NamedTuple): + physical_mesh: Mesh + loops: tuple[Loop, ...] + + def with_mesh(self, mesh: Mesh): + overlap = set(mesh.axis_names) & (self.resource_axes - set(self.physical_mesh.axis_names)) + if overlap: + raise ValueError(f"Cannot update the mesh of the current resource " + f"environment. The new mesh shadows already defined axes " + f"{show_axes(overlap)}") + return self._replace(physical_mesh=mesh) + + def with_extra_loop(self, loop: Loop): + if loop.name in self.resource_axes: + raise ValueError(f"Cannot extend the resource environment with loop named " + f"`{loop.name}`. An axis of this name is already defined!") + return self._replace(loops=self.loops + (loop,)) + + @property + def physical_resource_axes(self) -> set[ResourceAxisName]: + return set(self.physical_mesh.axis_names) + + @property + def loop_resource_axes(self) -> set[ResourceAxisName]: + return {loop.name for loop in self.loops} + + @property + def resource_axes(self) -> set[ResourceAxisName]: + return self.physical_resource_axes | self.loop_resource_axes + + @property + def shape(self): + shape = self.physical_mesh.shape + shape.update(self.loops) + return shape + + @property + def local_shape(self): + shape = self.physical_mesh.local_mesh.shape + shape.update(self.loops) + return shape + + def __repr__(self): + mesh_repr = ", ".join( + f"'{k}': {v}" for k, v in self.physical_mesh.shape.items()) + return f"ResourceEnv(mesh=Mesh({mesh_repr}), {self.loops!r})" + + +@functools.lru_cache(maxsize=128) +def _get_local_mesh(global_mesh: Mesh, process_index: int) -> Mesh: + if global_mesh.empty: + return global_mesh + is_local_device = np.vectorize( + lambda d: d.process_index == process_index, otypes=[bool])(global_mesh.devices) + subcube_indices = [] + # We take the smallest slice of each dimension that doesn't skip any local device. + for axis in range(global_mesh.devices.ndim): + other_axes = util.tuple_delete(tuple(range(global_mesh.devices.ndim)), axis) + # NOTE: This re-reduces over many axes multiple times, so we could definitely + # optimize it, but I hope it won't be a bottleneck anytime soon. + local_slices = is_local_device.any(other_axes, keepdims=False) + nonzero_indices = np.flatnonzero(local_slices) + start, end = int(np.min(nonzero_indices)), int(np.max(nonzero_indices)) + subcube_indices.append(slice(start, end + 1)) + subcube_indices_tuple = tuple(subcube_indices) + # We only end up with all conditions being true if the local devices formed a + # subcube of the full array. This is because we were biased towards taking a + # "hull" spanned by the devices, and in case the local devices don't form a + # subcube that hull will contain non-local devices. + if not is_local_device[subcube_indices_tuple].all(): + raise ValueError( + "When passing host local inputs to pjit or xmap, devices " + "connected to a single host must form a contiguous subcube of the " + "global device mesh") + return Mesh(global_mesh.devices[subcube_indices_tuple], global_mesh.axis_names) + + +_mesh_object_dict = {} # type: ignore + + +class Mesh(contextlib.ContextDecorator): + """Declare the hardware resources available in the scope of this manager. + + In particular, all ``axis_names`` become valid resource names inside the + managed block and can be used e.g. in the ``in_axis_resources`` argument of + :py:func:`jax.experimental.pjit.pjit`. Also see JAX's multi-process programming + model (https://jax.readthedocs.io/en/latest/multi_process.html) + and the Distributed arrays and automatic parallelization tutorial + (https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) + + If you are compiling in multiple threads, make sure that the + ``with Mesh`` context manager is inside the function that the threads will + execute. + + Args: + devices: A NumPy ndarray object containing JAX device objects (as + obtained e.g. from :py:func:`jax.devices`). + axis_names: A sequence of resource axis names to be assigned to the + dimensions of the ``devices`` argument. Its length should match the + rank of ``devices``. + + Example: + + >>> from jax.experimental.pjit import pjit + >>> from jax.sharding import Mesh + >>> from jax.sharding import PartitionSpec as P + >>> import numpy as np + ... + >>> inp = np.arange(16).reshape((8, 2)) + >>> devices = np.array(jax.devices()).reshape(4, 2) + ... + >>> # Declare a 2D mesh with axes `x` and `y`. + >>> global_mesh = Mesh(devices, ('x', 'y')) + >>> # Use the mesh object directly as a context manager. + >>> with global_mesh: + ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) + + >>> # Initialize the Mesh and use the mesh as the context manager. + >>> with Mesh(devices, ('x', 'y')) as global_mesh: + ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) + + >>> # Also you can use it as `with ... as ...`. + >>> global_mesh = Mesh(devices, ('x', 'y')) + >>> with global_mesh as m: + ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) + + >>> # You can also use it as `with Mesh(...)`. + >>> with Mesh(devices, ('x', 'y')): + ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) + """ + + devices: np.ndarray + axis_names: tuple[MeshAxisName, ...] + + def __new__(cls, devices: np.ndarray | Sequence[xc.Device], + axis_names: str | Sequence[MeshAxisName]): + if not isinstance(devices, np.ndarray): + devices = np.array(devices) + if isinstance(axis_names, str): + axis_names = (axis_names,) + axis_names = tuple(axis_names) + + if devices.ndim != len(axis_names): + raise ValueError( + "Mesh requires the ndim of its first argument (`devices`) to equal " + "the length of its second argument (`axis_names`), but got " + f"devices.ndim == {devices.ndim} and " + f"len(axis_names) == {len(axis_names)}.") + + key = (axis_names, devices.shape, tuple(devices.flat)) + val = _mesh_object_dict.get(key, None) + if val is not None: + return val + + self = super().__new__(cls) + self.devices = devices.copy() + self.devices.flags.writeable = False + self.axis_names = axis_names + _mesh_object_dict[key] = self + return self + + def __reduce__(self): + return (type(self), (self.devices, self.axis_names)) + + def __eq__(self, other): + if not isinstance(other, Mesh): + return False + # This is a performance optimization. Comparing thousands of devices + # can be expensive. + if id(self) == id(other): + return True + return (self.axis_names == other.axis_names and + self.devices.shape == other.devices.shape and + self._internal_device_list == other._internal_device_list) + + def __hash__(self): + if not hasattr(self, '_hash'): + self._hash = hash( + (self.axis_names, self._internal_device_list, self.devices.shape)) + return self._hash + + def __setattr__(self, name, value): + if hasattr(self, name): + if getattr(self, name) == value: + # This can to happen if two threads race, for example if two threads + # are trying to hash the same Mesh instance. + return + raise RuntimeError( + f"Cannot reassign attributes ({name}) of immutable mesh objects" + ) + super().__setattr__(name, value) + + def __enter__(self): + new_env = thread_resources.stack[-1].with_mesh(self) + thread_resources.stack.append(new_env) + thread_resources.env = new_env + jax_config.update_thread_local_jit_state( + mesh_context_manager=tuple(t.physical_mesh for t in thread_resources.stack + if not t.physical_mesh.empty)) + return self + + def __exit__(self, exc_type, exc_value, traceback): + thread_resources.stack.pop() + thread_resources.env = thread_resources.stack[-1] + jax_config.update_thread_local_jit_state( + mesh_context_manager=tuple(t.physical_mesh for t in thread_resources.stack + if not t.physical_mesh.empty)) + return False + + @property + def shape(self): + return collections.OrderedDict( + (name, size) + for name, size in util.safe_zip(self.axis_names, self.devices.shape)) + + @property + def size(self): + return math.prod(self.shape.values()) + + @property + def empty(self): + return self.devices.ndim == 0 + + @functools.cached_property + def is_multi_process(self): + return self.devices.size != len(self.local_devices) + + @property + def local_mesh(self): + return self._local_mesh(xb.process_index()) + + def _local_mesh(self, process_index): + return _get_local_mesh(self, process_index) + + @functools.cached_property + def device_ids(self): + assert not self.empty + return np.vectorize(lambda d: d.id, otypes=[int])(self.devices) + + @functools.cached_property + def _local_devices_set(self): + return set(self.local_devices) + + @functools.cached_property + def _flat_devices_tuple(self): + return tuple(self.devices.flat) + + @functools.cached_property + def _internal_device_list(self): + return xc.DeviceList(self._flat_devices_tuple) + + @functools.cached_property + def _flat_devices_set(self): + return set(self.devices.flat) + + def __str__(self): + mesh_str = ", ".join(f"'{k}': {v}" for k, v in self.shape.items()) + return f"Mesh({mesh_str})" + + @functools.cached_property + def _repr(self): + if self.empty: + return "Mesh(device_ids=[], axis_names=())" + return f"Mesh(device_ids={self.device_ids!r}, axis_names={self.axis_names!r})" + + def __repr__(self): + return self._repr + + @functools.cached_property + def local_devices(self): + return [d for d in self.devices.flat + if d.process_index == d.client.process_index()] + + +EMPTY_ENV = ResourceEnv(Mesh(np.empty((), dtype=object), ()), ()) + +class _ThreadResourcesLocalState(threading.local): + + def __init__(self): + self.stack = [EMPTY_ENV] + self.env = self.stack[-1] + +thread_resources = _ThreadResourcesLocalState() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/monitoring.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/monitoring.py new file mode 100644 index 0000000000000000000000000000000000000000..3b291de0061acb2257b69a06bac03dc3a6a5120c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/monitoring.py @@ -0,0 +1,118 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utilities for instrumenting code. + +Code points can be marked as a named event. Every time an event is reached +during program execution, the registered listeners will be invoked. + +A typical listener callback is to send an event to a metrics collector for +aggregation/exporting. +""" + +from __future__ import annotations + +from typing import Protocol + + +class EventListenerWithMetadata(Protocol): + + def __call__(self, event: str, **kwargs: str | int) -> None: + ... + + +class EventDurationListenerWithMetadata(Protocol): + + def __call__(self, event: str, duration_secs: float, + **kwargs: str | int) -> None: + ... + + +_event_listeners: list[EventListenerWithMetadata] = [] +_event_duration_secs_listeners: list[EventDurationListenerWithMetadata] = [] + + +def record_event(event: str, **kwargs: str | int) -> None: + """Record an event. + + If **kwargs are specified, all of the named arguments have to be passed in the + same order across all invocations of this method for the same event. + """ + for callback in _event_listeners: + callback(event, **kwargs) + + +def record_event_duration_secs(event: str, duration: float, + **kwargs: str | int) -> None: + """Record an event duration in seconds (float). + + If **kwargs are specified, all of the named arguments have to be passed in the + same order across all invocations of this method for the same event. + """ + for callback in _event_duration_secs_listeners: + callback(event, duration, **kwargs) + + +def register_event_listener( + callback: EventListenerWithMetadata, +) -> None: + """Register a callback to be invoked during record_event().""" + _event_listeners.append(callback) + + +def register_event_duration_secs_listener( + callback : EventDurationListenerWithMetadata) -> None: + """Register a callback to be invoked during record_event_duration_secs().""" + _event_duration_secs_listeners.append(callback) + +def get_event_duration_listeners() -> list[EventDurationListenerWithMetadata]: + """Get event duration listeners.""" + return list(_event_duration_secs_listeners) + +def get_event_listeners() -> list[EventListenerWithMetadata]: + """Get event listeners.""" + return list(_event_listeners) + +def clear_event_listeners(): + """Clear event listeners.""" + global _event_listeners, _event_duration_secs_listeners + _event_listeners = [] + _event_duration_secs_listeners = [] + +def _unregister_event_duration_listener_by_callback( + callback: EventDurationListenerWithMetadata) -> None: + """Unregister an event duration listener by callback. + + This function is supposed to be called for testing only. + """ + assert callback in _event_duration_secs_listeners + _event_duration_secs_listeners.remove(callback) + +def _unregister_event_duration_listener_by_index(index: int) -> None: + """Unregister an event duration listener by index. + + This function is supposed to be called for testing only. + """ + size = len(_event_duration_secs_listeners) + assert -size <= index < size + del _event_duration_secs_listeners[index] + +def _unregister_event_listener_by_callback( + callback: EventListenerWithMetadata) -> None: + """Unregister an event listener by callback. + + This function is supposed to be called for testing only. + """ + assert callback in _event_listeners + _event_listeners.remove(callback) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/op_shardings.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/op_shardings.py new file mode 100644 index 0000000000000000000000000000000000000000..6b248736ce45fbe3fea170ead89cb3d1352955e5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/op_shardings.py @@ -0,0 +1,111 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Sharding utilities""" + +from __future__ import annotations + +from collections.abc import Sequence +import itertools +from typing import Union + +import numpy as np + +from jax._src.lib import xla_client as xc + + +def get_num_ways_dim_sharded( + hlo_sharding: xc.HloSharding) -> tuple[list[int], int]: + if hlo_sharding.is_replicated(): + return [], 1 + partitions = hlo_sharding.tile_assignment_dimensions() + subgroup_types = hlo_sharding.subgroup_types() + + if subgroup_types == [xc.OpSharding.Type.REPLICATED]: + replicate_on_last_tile_dim = True + else: + replicate_on_last_tile_dim = hlo_sharding.replicate_on_last_tile_dim() + if subgroup_types: + raise NotImplementedError( + "Unhandled OpSharding type. Please open a bug report!") + num_replicas = 1 + if replicate_on_last_tile_dim: + num_replicas = partitions[-1] + partitions = partitions[:-1] + return list(partitions), num_replicas + + +def is_op_sharding_replicated(op: xc.OpSharding | xc.HloSharding) -> bool: + if isinstance(op, xc.OpSharding): + op = xc.HloSharding.from_proto(op) + if op.num_devices() == 1: + return True + return op.is_replicated() + + +def are_op_shardings_equal(op1: xc.OpSharding | xc.HloSharding, + op2: xc.OpSharding | xc.HloSharding) -> bool: + if id(op1) == id(op2): + return True + if is_op_sharding_replicated(op1) and is_op_sharding_replicated(op2): + return True + hc1 = xc.HloSharding.from_proto(op1) if isinstance(op1, xc.OpSharding) else op1 + hc2 = xc.HloSharding.from_proto(op2) if isinstance(op2, xc.OpSharding) else op2 + return hc1 == hc2 + + +_Index = Union[int, slice, tuple[Union[int, slice], ...]] + + +def op_sharding_to_numpy_indices( + hlo_sharding: xc.HloSharding, shape: Sequence[int], + num_devices: int) -> np.ndarray: + indices = np.empty(num_devices, dtype=np.object_) + + # num_devices is required as an argument when hlo_sharding is + # REPLICATED. `jax.device_count()` cannot be used because you can create + # an opsharding with less number of devices than `jax.device_count()`. + if is_op_sharding_replicated(hlo_sharding): + indices.fill((slice(None),) * len(shape)) + return indices + + assert num_devices == hlo_sharding.num_devices() + + partitions, num_replicas = get_num_ways_dim_sharded(hlo_sharding) + assert len(partitions) == len(shape), (len(partitions), len(shape)) + + axis_indices: list[Sequence[_Index]] = [] + for dim, n_shards in zip(shape, partitions): + if n_shards == 1: + axis_indices.append([slice(None)]) + elif n_shards > 1: + shard_size, ragged = divmod(dim, n_shards) + assert not ragged, (dim, n_shards) + axis_indices.append([slice(i * shard_size, (i + 1) * shard_size) + for i in range(n_shards)]) + else: + raise AssertionError('Unrecognized number of shards. Please file a bug!') + + device_it = iter(hlo_sharding.tile_assignment_devices()) + + for i, idxs in enumerate(itertools.product(*axis_indices)): + for _ in range(num_replicas): + indices[next(device_it)] = idxs + return indices + + +def op_sharding_to_indices( + op_sharding: xc.HloSharding, shape: Sequence[int], + num_devices: int) -> tuple[tuple[slice, ...], ...]: + indices = op_sharding_to_numpy_indices(op_sharding, shape, num_devices) + return tuple(indices.flat) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/partition_spec.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/partition_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..18e7d18d931ddf5d8ca73ecf284a93463dd37088 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/partition_spec.py @@ -0,0 +1,50 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +class _UnconstrainedPartitionSingleton: + + def __repr__(self): + return "UNCONSTRAINED" + + +# Unconstrained sentinel value for PartitionSpec, representing a dimension for +# which the user wants XLA to assign the best partitioning. +# TODO(yashkatariya): May rename to AUTO. +_UNCONSTRAINED_PARTITION = _UnconstrainedPartitionSingleton() + + +class PartitionSpec(tuple): + """Tuple describing how to partition an array across a mesh of devices. + + Each element is either ``None``, a string, or a tuple of strings. + See the documentation of :class:`jax.sharding.NamedSharding` for more details. + + This class exists so JAX's pytree utilities can distinguish a partition + specifications from tuples that should be treated as pytrees. + """ + + # A sentinel value representing a dim is unconstrained. + UNCONSTRAINED = _UNCONSTRAINED_PARTITION + + def __init__(self, *partitions): + pass + + def __new__(cls, *partitions): + return tuple.__new__(PartitionSpec, partitions) + + def __repr__(self): + return "PartitionSpec%s" % tuple.__repr__(self) + + def __reduce__(self): + return (PartitionSpec, tuple(self)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/path.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/path.py new file mode 100644 index 0000000000000000000000000000000000000000..1f92dd504282fa7e9b6c1df8e8fd3f09393271fb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/path.py @@ -0,0 +1,33 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import pathlib + +logger = logging.getLogger(__name__) + +try: + import etils.epath as epath +except: + epath = None + +# If etils.epath (aka etils[epath] to pip) is present, we prefer it because it +# can read and write to, e.g., GCS buckets. Otherwise we use the builtin +# pathlib and can only read/write to the local filesystem. +if epath: + logger.debug("etils.epath found. Using etils.epath for file I/O.") + Path = epath.Path +else: + logger.debug("etils.epath was not found. Using pathlib for file I/O.") + Path = pathlib.Path diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pickle_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pickle_util.py new file mode 100644 index 0000000000000000000000000000000000000000..2b82bb66074bca0d802614d97ef70d8b229de5d8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pickle_util.py @@ -0,0 +1,67 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import functools +import io +from typing import Any + +try: + import cloudpickle # type: ignore[import-not-found] +except ImportError: + cloudpickle = None + +from jax._src import profiler + + +@functools.partial(profiler.annotate_function, name='pickle_util.dumps') +def dumps(obj: Any) -> bytes: + """See `pickle.dumps`. Used for serializing host callbacks in jaxlib.""" + if cloudpickle is None: + raise ModuleNotFoundError('No module named "cloudpickle"') + + class Pickler(cloudpickle.CloudPickler): + """Customizes the behavior of cloudpickle.""" + + # Make a copy to avoid modifying cloudpickle for other users. + dispatch_table = cloudpickle.CloudPickler.dispatch_table.copy() + + # Fixes for dataclass internal singleton object serialization. + # Bug: https://github.com/cloudpipe/cloudpickle/issues/386 + # pylint: disable=protected-access + # pytype: disable=module-attr + dispatch_table[dataclasses._FIELD_BASE] = lambda x: f'{x.name}' + dispatch_table[dataclasses._MISSING_TYPE] = lambda _: 'MISSING' + dispatch_table[dataclasses._HAS_DEFAULT_FACTORY_CLASS] = ( + lambda _: '_HAS_DEFAULT_FACTORY' + ) + if hasattr(dataclasses, '_KW_ONLY_TYPE'): + dispatch_table[dataclasses._KW_ONLY_TYPE] = ( + lambda _: '_KW_ONLY_TYPE' + ) # Added in Python 3.10. + # pytype: enable=module-attr + # pylint: enable=protected-access + + with io.BytesIO() as file: + Pickler(file).dump(obj) + return file.getvalue() + + +@functools.partial(profiler.annotate_function, name='pickle_util.loads') +def loads(data: bytes) -> Any: + """See `pickle.loads`.""" + if cloudpickle is None: + raise ModuleNotFoundError('No module named "cloudpickle"') + + return cloudpickle.loads(data) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pjit.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pjit.py new file mode 100644 index 0000000000000000000000000000000000000000..edfff92ba1bd5ef45d866855797538deee969db0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pjit.py @@ -0,0 +1,2591 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Sequence, Iterable +import dataclasses +from functools import partial, lru_cache +import inspect +import itertools as it +import logging +import operator as op +import weakref +from typing import Callable, NamedTuple, Any, Union, Optional +import threading +import warnings + +import numpy as np + +from jax._src import api +from jax._src import ad_util +from jax._src import api_util +from jax._src import config +from jax._src import core +from jax._src import dispatch +from jax._src import dtypes +from jax._src import linear_util as lu +from jax._src import mesh as mesh_lib +from jax._src import op_shardings +from jax._src import profiler +from jax._src import sharding_impls +from jax._src import source_info_util +from jax._src import stages +from jax._src import traceback_util +from jax._src import tree_util +from jax._src import util +from jax._src import xla_bridge as xb +from jax._src.api_util import ( + argnums_partial_except, flatten_axes, flatten_fun, flatten_fun_nokwargs, + donation_vector, shaped_abstractify, check_callable, resolve_argnums, + argnames_partial_except, debug_info, result_paths, jaxpr_debug_info, + hoist_obj_attrs) +from jax._src.errors import JAXTypeError +from jax._src.interpreters import partial_eval as pe +from jax._src.partition_spec import PartitionSpec +from jax._src.interpreters import xla +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import pxla +from jax._src.lib import xla_extension_version +from jax._src.lib.mlir import ir +from jax._src.lib.mlir.dialects import func as func_dialect +from jax._src.lib import xla_client as xc +from jax._src import sharding +from jax._src.sharding_impls import ( + NamedSharding, XLACompatibleSharding, GSPMDSharding, + SingleDeviceSharding, PmapSharding, AUTO, UNSPECIFIED, UnspecifiedValue, + ParsedPartitionSpec, SpecSync, get_single_pspec, is_auto, is_unspecified, + is_unspecified_or_auto, prepare_axis_resources, parse_flatten_op_sharding) +from jax._src.layout import Layout, DeviceLocalLayout, AutoLayout +from jax._src.state import discharge as state_discharge, RefEffect, AbstractRef +from jax._src.traceback_util import api_boundary +from jax._src.tree_util import ( + tree_flatten, tree_unflatten, treedef_is_leaf, tree_structure, tree_leaves, + treedef_children, broadcast_prefix, all_leaves, prefix_errors, keystr, + PyTreeDef, none_leaf_registry as none_lr) +from jax._src.util import ( + HashableFunction, safe_map, safe_zip, wraps, + distributed_debug_log, split_list, weakref_lru_cache, + merge_lists, flatten, unflatten, subs_list, fun_name) + +map, unsafe_map = safe_map, map +zip, unsafe_zip = safe_zip, zip + +traceback_util.register_exclusion(__file__) + +PjitSharding = Union[GSPMDSharding, UnspecifiedValue, AUTO] +PjitShardingMinusUnspecified = Union[GSPMDSharding, AUTO] +MeshSharding = Union[NamedSharding, UnspecifiedValue, AUTO] +MeshShardingMinusUnspecified = Union[NamedSharding, AUTO] + +logger = logging.getLogger(__name__) + + +def _find_arg_mismatch(arg_list, fails, fun_name): + mismatched_args_msg = [] + def mismatch(err): + for name, inp_da, aval in arg_list: + if err.m_type == pxla.MismatchType.ARG_SHARDING and err.da == inp_da: + mismatched_args_msg.append( + f"argument {name} of {fun_name} with shape {aval.str_short()} and " + f"{err._dev_ids_plat_str}") + break + first_err, second_err = fails + mismatch(first_err) + mismatch(second_err) + return mismatched_args_msg + + +def _device_assignment_mismatch_error(fun_name, fails, args_flat, api_name, + arg_names): + arg_list = [] + if arg_names is None: + arg_names = [''] * len(args_flat) + for a, n in zip(args_flat, arg_names): + da = (a.sharding._device_assignment + if getattr(a, 'sharding', None) is not None else None) + arg_list.append((n, da, shaped_abstractify(a))) + + mismatched_args_msg = _find_arg_mismatch(arg_list, fails, fun_name) + + if len(mismatched_args_msg) == 2: + first, second = mismatched_args_msg # pytype: disable=bad-unpacking + extra_msg = f" Got {first} and {second}" + elif len(mismatched_args_msg) == 1: + first, second = fails + # Choose the failure left which is not already covered by ARG_SHARDING. + left = second if first.m_type == pxla.MismatchType.ARG_SHARDING else first + extra_msg = f" Got {mismatched_args_msg[0]} and{left._str(api_name)}" + else: + first, second = fails + extra_msg = f" Got{first._str(api_name)} and{second._str(api_name)}" + msg = (f"Received incompatible devices for {api_name}ted computation.{extra_msg}") + return msg + + +class PjitInfo(NamedTuple): + """Things that we know about a jit instance before it is called. + + In other words, this structure contains arguments to jit()/pjit(), + preprocessed and validated. + """ + fun: Callable + fun_sourceinfo: str | None + fun_signature: inspect.Signature | None + # Shardings, as specified by the user. These can either be UNSPECIFIED or they + # can be a tree (prefix) of shardings or None. + user_specified_in_shardings: bool + in_shardings_treedef: PyTreeDef + in_shardings_leaves: tuple[Any, ...] + out_shardings_treedef: PyTreeDef + out_shardings_leaves: tuple[Any, ...] + in_layouts_treedef: PyTreeDef + in_layouts_leaves: tuple[Any, ...] + out_layouts_treedef: PyTreeDef + out_layouts_leaves: tuple[Any, ...] + static_argnums: tuple[int, ...] + static_argnames: tuple[str, ...] + donate_argnums: tuple[int, ...] + donate_argnames: tuple[str, ...] + device: xc.Device | None + backend: str | None + keep_unused: bool + inline: bool + abstracted_axes: Any | None + has_explicit_sharding: bool + use_resource_env: bool # False for jit, True for pjit + + +def _python_pjit_helper(jit_info, *args, **kwargs): + (args_flat, params, in_avals, _, out_tree, _, arg_names, _, + attrs_tracked) = _infer_params(jit_info, args, kwargs) + + for arg in args_flat: + dispatch.check_arg(arg) + + if attrs_tracked: + init_states = _get_states(attrs_tracked) + args_flat = [*init_states, *args_flat] + + try: + out_flat = pjit_p.bind(*args_flat, **params) + except pxla.DeviceAssignmentMismatchError as e: + fails, = e.args + api_name = 'jit' if params['resource_env'] is None else 'pjit' + fun = jit_info.fun + fun_name = getattr(fun, '__qualname__', getattr(fun, '__name__', str(fun))) + msg = _device_assignment_mismatch_error( + fun_name, fails, args_flat, api_name, arg_names) + raise ValueError(msg) from None + except xla.InvalidInputException as e: + arg_names = [''] * len(args_flat) if arg_names is None else arg_names + # Run canonicalization again to figure out which arg failed. + if params['jaxpr'].consts: + raise TypeError(e.args[0]) from e + else: + for arg, name, aval in zip(args_flat, arg_names, in_avals): + try: + xla.canonicalize_dtype(arg) + except xla.InvalidInputException as _: + # Reraise as TypeError with the new message. + raise TypeError( + f"Argument '{name}' of shape {aval.str_short()} of type" + f' {type(arg)} is not a valid JAX type.') from e + raise AssertionError("Unreachable") from e + + if attrs_tracked: + num_states_out = sum(end_tree.num_leaves for _, end_tree, _ in attrs_tracked) + final_states, out_flat = split_list(out_flat, [num_states_out]) + _set_states(attrs_tracked, final_states) + + outs = tree_unflatten(out_tree, out_flat) + return outs, out_flat, out_tree, args_flat, params['jaxpr'], attrs_tracked + + +def _set_states(attrs_tracked, vals): + from jax.experimental.attrs import jax_setattr + valss = split_list(vals, [td.num_leaves for _, td, _ in attrs_tracked[:-1]]) + for ((_, treedef, (obj, attr)), leaves) in zip(attrs_tracked, valss): + val = tree_unflatten(treedef, leaves) + jax_setattr(obj, attr, val) + +def _get_states(attrs_tracked): + from jax.experimental.attrs import jax_getattr + vals = [] + for treedef, _, (obj, attr) in attrs_tracked: + tree = jax_getattr(obj, attr) + leaves, treedef_ = tree_flatten(tree) + assert treedef == treedef_ + vals.extend(leaves) + return vals + +def _need_to_rebuild_with_fdo(pgle_profiler): + return (pgle_profiler is not None and pgle_profiler.is_enabled() + and not pgle_profiler.is_fdo_consumed()) + +def _get_fastpath_data( + executable, out_tree, args_flat, out_flat, attrs_tracked, effects, + consts, abstracted_axes, pgle_profiler +) -> Optional[pxla.MeshExecutableFastpathData]: + out_reflattened, out_tree = pxla.reflatten_outputs_for_dispatch(out_tree, out_flat) + + use_fastpath = ( + executable is not None + and isinstance(executable, pxla.MeshExecutable) + and isinstance(executable.unsafe_call, pxla.ExecuteReplicated) + # No effects in computation + and not executable.unsafe_call.ordered_effects + and not executable.unsafe_call.has_unordered_effects + and not executable.unsafe_call.has_host_callbacks + and all(isinstance(x, xc.ArrayImpl) for x in out_reflattened) + and abstracted_axes is None + # no attr state effects + and not attrs_tracked + # no ref state effects + and not any(isinstance(e, RefEffect) for e in effects) + # no prng reuse checking + and not (config.debug_key_reuse.value and any( + hasattr(arg, 'dtype') and dtypes.issubdtype(arg.dtype, dtypes.prng_key) + for arg in (*args_flat, *out_flat, *consts))) + and not _need_to_rebuild_with_fdo(pgle_profiler) + ) + + if use_fastpath: + out_avals = [o.aval for o in out_reflattened] + out_committed = [o._committed for o in out_reflattened] + kept_var_bitvec = [i in executable._kept_var_idx + for i in range(len(args_flat))] + in_shardings = [ + sharding_impls.physical_sharding(a, s) + if a is not core.abstract_token and dtypes.issubdtype(a.dtype, dtypes.extended) + else s + for s, a in zip(executable._in_shardings, executable.in_avals) + ] + fastpath_data = pxla.MeshExecutableFastpathData( + executable.xla_executable, out_tree, in_shardings, + executable._out_shardings, out_avals, out_committed, kept_var_bitvec, + executable.unsafe_call.in_handler.local_devices, + executable.unsafe_call.in_handler.input_indices) + else: + fastpath_data = None + return fastpath_data + + +class _MostRecentPjitCallExecutable(threading.local): + def __init__(self): + self.weak_key_dict = weakref.WeakKeyDictionary() + self.weak_pgle_profiler_dict = weakref.WeakKeyDictionary() + +_most_recent_pjit_call_executable = _MostRecentPjitCallExecutable() + + +def _read_most_recent_pjit_call_executable(jaxpr): + return _most_recent_pjit_call_executable.weak_key_dict.get(jaxpr, None) + + +def _read_pgle_profiler(jaxpr): + return _most_recent_pjit_call_executable.weak_pgle_profiler_dict.get( + jaxpr, None + ) + +def _cpp_pjit_evict_fn(self): + self._clear_cache() + _create_pjit_jaxpr.evict_function(self._fun) # pytype: disable=attribute-error + + +# The entries are doubled here from the default 4096 because _pjit_call_impl +# also has a cpp dispatch path and that would double the number of entries in +# the global shared cache. +_cpp_pjit_cache = xc._xla.PjitFunctionCache(capacity=8192) + + +def _get_cpp_global_cache(pjit_has_explicit_sharding): + if pjit_has_explicit_sharding: + return xc._xla.PjitFunctionCache() + else: + return _cpp_pjit_cache + + +def _cpp_pjit(jit_info: PjitInfo): + + @api_boundary + def cache_miss(*args, **kwargs): + outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked = _python_pjit_helper( + jit_info, *args, **kwargs) + executable = _read_most_recent_pjit_call_executable(jaxpr) + pgle_profiler = _read_pgle_profiler(jaxpr) + maybe_fastpath_data = _get_fastpath_data( + executable, out_tree, args_flat, out_flat, attrs_tracked, jaxpr.effects, + jaxpr.consts, jit_info.abstracted_axes, + pgle_profiler) + + if xla_extension_version > 267: + return outs, maybe_fastpath_data, _need_to_rebuild_with_fdo(pgle_profiler) + else: + return outs, maybe_fastpath_data + + fun = jit_info.fun + cpp_pjit_f = xc._xla.pjit( + fun_name(fun), + fun, cache_miss, jit_info.static_argnums, jit_info.static_argnames, + jit_info.donate_argnums, tree_util.dispatch_registry, + pxla.shard_arg, + _get_cpp_global_cache(jit_info.has_explicit_sharding)) + + cpp_pjitted_f = wraps(fun)(cpp_pjit_f) + cpp_pjitted_f._fun = fun + type(cpp_pjitted_f).clear_cache = _cpp_pjit_evict_fn + return cpp_pjitted_f + + +def _pjit_explicit_sharding(in_shardings, out_shardings, device, + backend) -> bool: + in_shardings_flat, _ = tree_flatten(in_shardings) + out_shardings_flat, _ = tree_flatten(out_shardings) + return (device is not None or + backend is not None or + any(not is_unspecified(i) for i in in_shardings_flat) or + any(not is_unspecified(i) for i in out_shardings_flat)) + + +def _split_layout_and_sharding(entries): + entries_flat, treedef = tree_flatten(entries, is_leaf=lambda x: x is None) + layouts, shardings = [], [] + + for e in entries_flat: + if e is None or is_unspecified_or_auto(e): + layouts.append(None) + shardings.append(e) + elif isinstance(e, Layout): + layouts.append(e.device_local_layout) + shardings.append(e.sharding) + elif isinstance(e, (DeviceLocalLayout, AutoLayout)): + raise ValueError( + '`jax.jit` does not accept device-local layouts directly. Create ' + 'a `Layout` instance wrapping this device-local layout and pass ' + f'that to `jit` instead. Got {e}') + else: + layouts.append(None) + shardings.append(e) + + assert len(layouts) == len(shardings) + return tree_unflatten(treedef, layouts), tree_unflatten(treedef, shardings) + + +def _parse_jit_arguments(fun: Callable, in_shardings: Any, out_shardings: Any, + donate_argnums: int | Sequence[int] | None, + donate_argnames: str | Iterable[str] | None, + static_argnums: int | Sequence[int] | None, + static_argnames: str | Iterable[str] | None, + device: xc.Device | None, backend: str | None, + abstracted_axes: Any | None, keep_unused: bool, + inline: bool, use_resource_env: bool) -> PjitInfo: + """Parses the arguments to jit/pjit. + + Performs any preprocessing and validation of the arguments that we can do + ahead of time before the jit()-ed function is invoked. + """ + if abstracted_axes and not config.dynamic_shapes.value: + raise ValueError("abstracted_axes must be used with --jax_dynamic_shapes") + + check_callable(fun) + + if backend is not None or device is not None: + warnings.warn( + 'backend and device argument on jit is deprecated. You can use' + ' `jax.device_put(..., jax.local_devices("cpu")[0])` on the inputs to' + ' the jitted function to get the same behavior.', DeprecationWarning) + if device is not None and backend is not None: + raise ValueError("can't specify both a device and a backend for jit, " + f"got {device=} and {backend=}") + if in_shardings is not None and not is_unspecified(in_shardings): + raise ValueError('If backend or device is specified on jit, then ' + 'in_shardings should not be specified.') + if out_shardings is not None and not is_unspecified(out_shardings): + raise ValueError('If backend or device is specified on jit, then ' + 'out_shardings should not be specified.') + + if isinstance(in_shardings, list): + # To be a tree prefix of the positional args tuple, in_axes can never be a + # list: if in_axes is not a leaf, it must be a tuple of trees. However, + # in cases like these users expect tuples and lists to be treated + # essentially interchangeably, so we canonicalize lists to tuples here + # rather than raising an error. https://github.com/google/jax/issues/2367 + in_shardings = tuple(in_shardings) + + in_layouts, in_shardings = _split_layout_and_sharding(in_shardings) + out_layouts, out_shardings = _split_layout_and_sharding(out_shardings) + + in_shardings = prepare_axis_resources(in_shardings, 'in_shardings') + out_shardings = prepare_axis_resources(out_shardings, 'out_shardings') + + user_specified_in_shardings = (in_shardings is not None and + not is_unspecified(in_shardings)) + + in_shardings_leaves, in_shardings_treedef = none_lr.flatten(in_shardings) + out_shardings_leaves, out_shardings_treedef = none_lr.flatten(out_shardings) + in_layouts_leaves, in_layouts_treedef = none_lr.flatten(in_layouts) + out_layouts_leaves, out_layouts_treedef = none_lr.flatten(out_layouts) + + fun_sourceinfo = api_util.fun_sourceinfo(fun) + fun_signature = api_util.fun_signature(fun) + + donate_argnums, donate_argnames, static_argnums, static_argnames = resolve_argnums( + fun, fun_signature, donate_argnums, donate_argnames, static_argnums, + static_argnames) + + has_explicit_sharding = _pjit_explicit_sharding( + in_shardings, out_shardings, device, backend) + + return PjitInfo( + fun=fun, + fun_sourceinfo=fun_sourceinfo, + fun_signature=fun_signature, + user_specified_in_shardings=user_specified_in_shardings, + in_shardings_treedef=in_shardings_treedef, + in_shardings_leaves=tuple(in_shardings_leaves), + out_shardings_treedef=out_shardings_treedef, + out_shardings_leaves=tuple(out_shardings_leaves), + in_layouts_treedef=in_layouts_treedef, + in_layouts_leaves=tuple(in_layouts_leaves), + out_layouts_treedef=out_layouts_treedef, + out_layouts_leaves=tuple(out_layouts_leaves), + static_argnums=static_argnums, + static_argnames=static_argnames, donate_argnums=donate_argnums, + donate_argnames=donate_argnames, device=device, backend=backend, + keep_unused=keep_unused, inline=inline, + abstracted_axes=abstracted_axes, + has_explicit_sharding=has_explicit_sharding, + use_resource_env=use_resource_env) + + +def _make_jit_wrapper(jit_info: PjitInfo): + + @api_boundary + def lower(*args, **kwargs): + traced = trace(*args, **kwargs) + try: + return traced.lower() + except pxla.DeviceAssignmentMismatchError as e: + fails, = e.args + fun = jit_info.fun + fun_name = getattr(fun, '__qualname__', + getattr(fun, '__name__', str(fun))) + msg = _device_assignment_mismatch_error( + fun_name, fails, traced._args_flat, 'jit', traced._arg_names) + raise ValueError(msg) from None + + @api_boundary + def eval_shape(*args, **kwargs): + _, params, _, _, out_tree, _, _, _, _ = _infer_params(jit_info, args, kwargs) + out_s = [None if is_unspecified(s) else s for s in params['out_shardings']] + # TODO(yashkatariya): Add `Layout` to SDS. + out = [api.ShapeDtypeStruct(x.shape, x.dtype, x.named_shape, sharding=s) + for x, s in zip(params['jaxpr'].out_avals, out_s)] + return tree_unflatten(out_tree, out) + + @api_boundary + def trace(*args, **kwargs) -> stages.Traced: + lowering_parameters = kwargs.pop( + '_experimental_lowering_parameters', mlir.LoweringParameters()) + + (args_flat, params, in_avals, in_tree, out_tree, donated_invars, + arg_names, num_consts, _) = _infer_params(jit_info, args, kwargs) + + donate_argnums = tuple(i for i, d in enumerate(donated_invars) if d) + args_info = stages.make_args_info(in_tree, in_avals, donate_argnums) + lower_callable = partial(_resolve_and_lower, args_flat, **params, + lowering_parameters=lowering_parameters) + return stages.Traced(params['jaxpr'], args_info, params["name"], out_tree, + lower_callable, args_flat, arg_names, num_consts) + + wrapped = _cpp_pjit(jit_info) + wrapped.lower = lower + wrapped.eval_shape = eval_shape + wrapped.trace = trace + return wrapped + + +def make_jit(fun: Callable, in_shardings: Any, out_shardings: Any, + donate_argnums: int | Sequence[int] | None, + donate_argnames: str | Iterable[str] | None, + static_argnums: int | Sequence[int] | None, + static_argnames: str | Iterable[str] | None, + device: xc.Device | None, backend: str | None, + abstracted_axes: Any | None, keep_unused: bool, + inline: bool, use_resource_env: bool) -> Any: + """jit() and pjit() are thin wrappers around this function.""" + jit_info = _parse_jit_arguments( + fun, in_shardings, out_shardings, donate_argnums, donate_argnames, + static_argnums, static_argnames, device, backend, abstracted_axes, + keep_unused, inline, use_resource_env) + return _make_jit_wrapper(jit_info) + + +def _infer_params(jit_info, args, kwargs): + (fun, fun_sourceinfo, fun_signature, user_specified_in_shardings, + in_shardings_treedef, in_shardings_leaves, out_shardings_treedef, + out_shardings_leaves, in_layouts_treedef, in_layouts_leaves, + out_layouts_treedef, out_layouts_leaves, static_argnums, static_argnames, + donate_argnums, donate_argnames, device, backend, keep_unused, inline, + abstracted_axes, _, use_resource_env) = jit_info + + have_kwargs = bool(kwargs) + if have_kwargs and user_specified_in_shardings: + raise ValueError( + "pjit does not support kwargs when in_shardings is specified.") + + if use_resource_env: + # We need to fetch the mesh from inside the wrapped function, because + # meshes are dynamically scoped (i.e., with a context manager). + resource_env = mesh_lib.thread_resources.env + pjit_mesh = resource_env.physical_mesh + jit_name = 'pjit' + if (backend or device) and not pjit_mesh.empty: + raise ValueError( + "Mesh context manager should not be used with jit when backend or " + "device is also specified as an argument to jit.") + else: + resource_env = None + pjit_mesh = None + jit_name = 'jit' + + + axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs) + + dbg = debug_info(jit_name, fun_sourceinfo, fun_signature, args, kwargs, + static_argnums, static_argnames) + f = lu.wrap_init(fun) + f, res_paths = result_paths(f) + f, dyn_args = argnums_partial_except(f, static_argnums, args, allow_invalid=True) + del args + + f, dyn_kwargs = argnames_partial_except(f, static_argnames, kwargs) + explicit_args, in_tree = tree_flatten((dyn_args, dyn_kwargs)) + flat_fun, out_tree = flatten_fun(f, in_tree) + flat_fun, explicit_args = hoist_obj_attrs(flat_fun, explicit_args) + + if (donate_argnums or donate_argnames) and not config.debug_nans.value: + donated_invars = donation_vector(donate_argnums, donate_argnames, in_tree) + else: + donated_invars = (False,) * len(explicit_args) + del donate_argnums, donate_argnames + + # If backend or device is set as an arg on jit, then resolve them to + # in_shardings and out_shardings as if user passed in in_shardings + # and out_shardings. + device_or_backend_set = bool(backend or device) + if device_or_backend_set: + sharding = _create_sharding_with_device_backend(device, backend) + leaves, treedef = tree_flatten(sharding) + in_shardings_leaves = out_shardings_leaves = tuple(leaves) + in_shardings_treedef = out_shardings_treedef = treedef + else: + in_shardings_leaves = tuple( + _create_sharding_for_array(pjit_mesh, x, 'in_shardings', jit_name) + for x in in_shardings_leaves) + out_shardings_leaves = tuple( + _create_sharding_for_array(pjit_mesh, x, 'out_shardings', jit_name) + for x in out_shardings_leaves) + + assert None not in in_shardings_leaves + assert None not in out_shardings_leaves + + if config.dynamic_shapes.value: + in_type = pe.infer_lambda_input_type(axes_specs, explicit_args) + in_avals = tuple(a for a, e in in_type if e) + else: + avals = [] + for i, a in enumerate(explicit_args): + try: + avals.append(shaped_abstractify(a)) + except OverflowError as e: + arg_path = (f"argument path is {dbg.arg_names[i]}" if dbg + else f"flattened argument number is {i}") + raise OverflowError( + "An overflow was encountered while parsing an argument to a jitted " + f"computation, whose {arg_path}." + ) from e + in_type = in_avals = tuple(avals) + + in_shardings_flat, in_layouts_flat = _process_in_axis_resources( + in_shardings_treedef, in_shardings_leaves, + in_layouts_treedef, in_layouts_leaves, + in_avals, in_tree, dbg, device_or_backend_set, have_kwargs) + + jaxpr, consts, out_shardings_flat, out_layouts_flat, attrs_tracked = _pjit_jaxpr( + flat_fun, out_shardings_treedef, out_shardings_leaves, + out_layouts_treedef, out_layouts_leaves, in_type, dbg, + device_or_backend_set, HashableFunction(out_tree, closure=()), + HashableFunction(res_paths, closure=()), inline) + + assert len(explicit_args) == len(in_shardings_flat) == len(in_layouts_flat) + + if config.dynamic_shapes.value: + implicit_args = _extract_implicit_args(in_type, explicit_args) + else: + implicit_args = [] + args_flat = [*implicit_args, *explicit_args] + + num_states_in = sum(init_tree.num_leaves for init_tree, _, _ in attrs_tracked) + num_extra_args = len(implicit_args) + num_states_in + len(consts) + in_shardings_flat = (UNSPECIFIED,) * num_extra_args + in_shardings_flat + in_layouts_flat = (None,) * num_extra_args + in_layouts_flat + donated_invars = (False,) * num_extra_args + donated_invars + assert (len(in_shardings_flat) == len(in_layouts_flat) == + len(donated_invars) == num_states_in + len(consts) + len(args_flat)) + + params = dict( + jaxpr=jaxpr, + in_shardings=in_shardings_flat, + out_shardings=out_shardings_flat, + in_layouts=in_layouts_flat, + out_layouts=out_layouts_flat, + resource_env=resource_env, + donated_invars=donated_invars, + name=fun_name(flat_fun), + keep_unused=keep_unused, + inline=inline, + ) + return (consts + args_flat, params, in_avals, in_tree, out_tree(), + donated_invars, dbg.arg_names if dbg else None, len(consts), + attrs_tracked) + +def _extract_implicit_args( + in_type: Sequence[tuple[core.AbstractValue, bool]], + explicit_args: Sequence[Any] +) -> Sequence[core.Tracer]: + """ + Given an input type and explicitly-passed arguments (per the user-facing API + calling convention), extract implicit axis size arguments from shapes of + explicit arguments (for the trace-time / jaxpr-level calling convention). + """ + # First, using `in_type` construct a list to represent the full argument list, + # leaving the implicit arguments as None placeholders for now. + explicit_args_ = iter(explicit_args) + args = [next(explicit_args_) if expl else None for _, expl in in_type] + assert next(explicit_args_, None) is None + del explicit_args, explicit_args_ + + # Next, populate the implicit arguments using the DBIdxs in `in_type`. + for i, (aval, explicit) in enumerate(in_type): + if not explicit or not isinstance(aval, core.DShapedArray): + continue # can't populate an implicit argument + arg = args[i] + assert arg is not None + for d1, d2 in zip(aval.shape, arg.aval.shape): + if isinstance(d1, core.DBIdx): + if args[d1.val] is None: + args[d1.val] = d2 + assert core.same_referent(args[d1.val], d2) + assert all(x is not None for x in args) + return [x for x, (_, e) in zip(args, in_type) if not e] # pytype: disable=bad-return-type + +def _flat_axes_specs(abstracted_axes, *args, **kwargs + ) -> list[pe.AbstractedAxesSpec] | None: + if abstracted_axes is None: return None + if kwargs: raise NotImplementedError + def ax_leaf(l): + return (isinstance(l, dict) and all_leaves(l.values()) or + isinstance(l, tuple) and all_leaves(l, lambda x: x is None)) + return broadcast_prefix(abstracted_axes, args, ax_leaf) + + +class JitWrapped(stages.Wrapped): + + def eval_shape(self, *args, **kwargs): + """See ``jax.eval_shape``.""" + raise NotImplementedError + + def trace(self, *args, **kwargs) -> stages.Traced: + raise NotImplementedError + + +# in_shardings and out_shardings can't be None as the default value +# because `None` means that the input is fully replicated. +def pjit( + fun: Callable, + in_shardings=UNSPECIFIED, + out_shardings=UNSPECIFIED, + static_argnums: int | Sequence[int] | None = None, + static_argnames: str | Iterable[str] | None = None, + donate_argnums: int | Sequence[int] | None = None, + donate_argnames: str | Iterable[str] | None = None, + keep_unused: bool = False, + device: xc.Device | None = None, + backend: str | None = None, + inline: bool = False, + abstracted_axes: Any | None = None, +) -> JitWrapped: + """Makes ``fun`` compiled and automatically partitioned across multiple devices. + + NOTE: This function is now equivalent to jax.jit please use that instead. + The returned function has semantics equivalent to those of ``fun``, but is + compiled to an XLA computation that runs across multiple devices + (e.g. multiple GPUs or multiple TPU cores). This can be useful if the jitted + version of ``fun`` would not fit in a single device's memory, or to speed up + ``fun`` by running each operation in parallel across multiple devices. + + The partitioning over devices happens automatically based on the + propagation of the input partitioning specified in ``in_shardings`` and + the output partitioning specified in ``out_shardings``. The resources + specified in those two arguments must refer to mesh axes, as defined by + the :py:func:`jax.sharding.Mesh` context manager. Note that the mesh + definition at :func:`~pjit` application time is ignored, and the returned function + will use the mesh definition available at each call site. + + Inputs to a :func:`~pjit`'d function will be automatically partitioned across devices + if they're not already correctly partitioned based on ``in_shardings``. + In some scenarios, ensuring that the inputs are already correctly pre-partitioned + can increase performance. For example, if passing the output of one + :func:`~pjit`'d function to another :func:`~pjit`’d function (or the same + :func:`~pjit`’d function in a loop), make sure the relevant + ``out_shardings`` match the corresponding ``in_shardings``. + + .. note:: + **Multi-process platforms:** On multi-process platforms such as TPU pods, + :func:`~pjit` can be used to run computations across all available devices across + processes. To achieve this, :func:`~pjit` is designed to be used in SPMD Python + programs, where every process is running the same Python code such that all + processes run the same :func:`~pjit`'d function in the same order. + + When running in this configuration, the mesh should contain devices across + all processes. All inputs arguments must be globally shaped. + ``fun`` will still be executed across *all* devices in the mesh, + including those from other processes, and will be given a global view of the + data spread across multiple processes as a single array. + + The SPMD model also requires that the same multi-process :func:`~pjit`'d + functions must be run in the same order on all processes, but they can be + interspersed with arbitrary operations running in a single process. + + Args: + fun: Function to be compiled. Should be a pure function, as side-effects may + only be executed once. Its arguments and return value should be arrays, + scalars, or (nested) standard Python containers (tuple/list/dict) thereof. + Positional arguments indicated by ``static_argnums`` can be anything at + all, provided they are hashable and have an equality operation defined. + Static arguments are included as part of a compilation cache key, which is + why hash and equality operators must be defined. + in_shardings: Pytree of structure matching that of arguments to ``fun``, + with all actual arguments replaced by resource assignment specifications. + It is also valid to specify a pytree prefix (e.g. one value in place of a + whole subtree), in which case the leaves get broadcast to all values in + that subtree. + + The ``in_shardings`` argument is optional. JAX will infer the shardings + from the input :py:class:`jax.Array`'s, and defaults to replicating the input + if the sharding cannot be inferred. + + The valid resource assignment specifications are: + + - :py:class:`Sharding`, which will decide how the value + will be partitioned. With this, using a mesh context manager is not + required. + - :py:obj:`None` is a special case whose semantics are: + - if the mesh context manager is *not* provided, JAX has the freedom to + choose whatever sharding it wants. + For in_shardings, JAX will mark is as replicated but this behavior + can change in the future. + For out_shardings, we will rely on the XLA GSPMD partitioner to + determine the output shardings. + - If the mesh context manager is provided, None will imply that the + value will be replicated on all devices of the mesh. + - For backwards compatibility, in_shardings still supports ingesting + :py:class:`PartitionSpec`. This option can *only* be used with the + mesh context manager. + + - :py:class:`PartitionSpec`, a tuple of length at most equal to the rank + of the partitioned value. Each element can be a :py:obj:`None`, a mesh + axis or a tuple of mesh axes, and specifies the set of resources assigned + to partition the value's dimension matching its position in the spec. + + The size of every dimension has to be a multiple of the total number of + resources assigned to it. + out_shardings: Like ``in_shardings``, but specifies resource + assignment for function outputs. + The ``out_shardings`` argument is optional. If not specified, :py:func:`jax.jit` + will use GSPMD's sharding propagation to determine how to shard the outputs. + static_argnums: An optional int or collection of ints that specify which + positional arguments to treat as static (compile-time constant). + Operations that only depend on static arguments will be constant-folded in + Python (during tracing), and so the corresponding argument values can be + any Python object. + + Static arguments should be hashable, meaning both ``__hash__`` and + ``__eq__`` are implemented, and immutable. Calling the jitted function + with different values for these constants will trigger recompilation. + Arguments that are not arrays or containers thereof must be marked as + static. + + If ``static_argnums`` is not provided, no arguments are treated as static. + static_argnames: An optional string or collection of strings specifying + which named arguments to treat as static (compile-time constant). See the + comment on ``static_argnums`` for details. If not + provided but ``static_argnums`` is set, the default is based on calling + ``inspect.signature(fun)`` to find corresponding named arguments. + donate_argnums: Specify which positional argument buffers are "donated" to + the computation. It is safe to donate argument buffers if you no longer + need them once the computation has finished. In some cases XLA can make + use of donated buffers to reduce the amount of memory needed to perform a + computation, for example recycling one of your input buffers to store a + result. You should not reuse buffers that you donate to a computation, JAX + will raise an error if you try to. By default, no argument buffers are + donated. + + If neither ``donate_argnums`` nor ``donate_argnames`` is provided, no + arguments are donated. If ``donate_argnums`` is not provided but + ``donate_argnames`` is, or vice versa, JAX uses + :code:`inspect.signature(fun)` to find any positional arguments that + correspond to ``donate_argnames`` + (or vice versa). If both ``donate_argnums`` and ``donate_argnames`` are + provided, ``inspect.signature`` is not used, and only actual + parameters listed in either ``donate_argnums`` or ``donate_argnames`` will + be donated. + + For more details on buffer donation see the + `FAQ `_. + donate_argnames: An optional string or collection of strings specifying + which named arguments are donated to the computation. See the + comment on ``donate_argnums`` for details. If not + provided but ``donate_argnums`` is set, the default is based on calling + ``inspect.signature(fun)`` to find corresponding named arguments. + keep_unused: If `False` (the default), arguments that JAX determines to be + unused by `fun` *may* be dropped from resulting compiled XLA executables. + Such arguments will not be transferred to the device nor provided to the + underlying executable. If `True`, unused arguments will not be pruned. + device: This argument is deprecated. Please put your arguments on the + device you want before passing them to jit. + Optional, the Device the jitted function will run on. (Available devices + can be retrieved via :py:func:`jax.devices`.) The default is inherited + from XLA's DeviceAssignment logic and is usually to use + ``jax.devices()[0]``. + backend: This argument is deprecated. Please put your arguments on the + backend you want before passing them to jit. + Optional, a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or + ``'tpu'``. + + Returns: + A wrapped version of ``fun``, set up for just-in-time compilation and + automatically partitioned by the mesh available at each call site. + + For example, a convolution operator can be automatically partitioned over + an arbitrary set of devices by a single :func:`~pjit` application: + + >>> import jax + >>> import jax.numpy as jnp + >>> import numpy as np + >>> from jax.sharding import Mesh, PartitionSpec + >>> from jax.experimental.pjit import pjit + >>> + >>> x = jnp.arange(8, dtype=jnp.float32) + >>> f = pjit(lambda x: jax.numpy.convolve(x, jnp.asarray([0.5, 1.0, 0.5]), 'same'), + ... in_shardings=None, out_shardings=PartitionSpec('devices')) + >>> with Mesh(np.array(jax.devices()), ('devices',)): + ... print(f(x)) # doctest: +SKIP + [ 0.5 2. 4. 6. 8. 10. 12. 10. ] + """ + return make_jit( + fun, in_shardings, out_shardings, donate_argnums, donate_argnames, + static_argnums, static_argnames, device, backend, abstracted_axes, + keep_unused, inline, use_resource_env=True) + + +def hashable_pytree(pytree): + vals, treedef = tree_flatten(pytree) + vals = tuple(vals) + return HashableFunction(lambda: tree_unflatten(treedef, vals), + closure=(treedef, vals)) + + +def _create_sharding_for_array(mesh, x, name, api_name): + if x is None and (mesh is None or mesh.empty): + return UNSPECIFIED + if isinstance(x, sharding.Sharding) or is_unspecified_or_auto(x): + return x + if mesh is None: + msg = ('jax.jit only supports `Sharding`s being passed to' + f' {name}. Looks like you are passing either `PartitionSpec` or `None`' + f' which is not allowed in jax.jit.\n') + if name == 'in_shardings': + msg += (f'Note that {name} argument is optional. JAX will infer the shardings' + " from the input jax.Array's and will default to replicating the" + ' input if the sharding cannot be inferred.') + elif name == 'out_shardings': + msg += (f'Note that {name} is optional. If not specified, jax.jit will' + " use GSPMD's sharding propagation to figure out what the sharding" + ' of the output(s) should be.') + raise RuntimeError(msg) + if mesh.empty: + raise RuntimeError( + f'{api_name} requires a non-empty mesh if you are passing' + f' `PartitionSpec`s or `None` to {name}! Is a mesh defined at the call' + f' site? Alternatively, provide `Sharding`s to {name} and' + ' then the mesh context manager is not required.') + # A nice user error is raised in prepare_axis_resources. + assert x is None or isinstance(x, ParsedPartitionSpec), x + return (pxla.create_mesh_pspec_sharding(mesh, x) + if x is None else pxla.create_mesh_pspec_sharding(mesh, x.user_spec, x)) + + +def _create_sharding_with_device_backend(device, backend): + if device is not None: + assert backend is None + out = SingleDeviceSharding(device) + elif backend is not None: + assert device is None + out = SingleDeviceSharding(xb.get_backend(backend).local_devices()[0]) + else: + raise AssertionError('Unreachable!') + out._device_backend = True + return out + + +def flatten_axis_resources(what, tree, shardings, tupled_args): + try: + return tuple(flatten_axes(what, tree, shardings, tupled_args=tupled_args)) + except ValueError: + pass # Raise a tree prefix error below + + # Tree leaves are always valid prefixes, so if there was a prefix error as + # assumed here, axis_resources must not be a leaf. + assert not treedef_is_leaf(tree_structure(shardings)) + + # Check the type directly rather than using isinstance because of namedtuples. + if tupled_args and (type(shardings) is not tuple or + len(shardings) != len(tree.children())): + # We know axis_resources is meant to be a tuple corresponding to the args + # tuple, but while it is a non-leaf pytree, either it wasn't a tuple or it + # wasn't the right length. + msg = (f"{what} specification must be a tree prefix of the positional " + f"arguments tuple passed to the `pjit`-decorated function. In " + f"particular, {what} must either be a None, a PartitionSpec, or " + f"a tuple of length equal to the number of positional arguments.") + # If `tree` represents an args tuple, then `axis_resources` must be a tuple. + # TODO(mattjj,apaszke): disable implicit list casts, remove 'or list' below + if type(shardings) is not tuple: + msg += f" But {what} is not a tuple: got {type(shardings)} instead." + elif len(shardings) != len(tree.children()): + msg += (f" But {what} is the wrong length: got a tuple or list of length " + f"{len(shardings)} for an args tuple of length " + f"{len(tree.children())}.") + + # As an extra hint, let's check if the user just forgot to wrap + # shardings in a singleton tuple. + if len(tree.children()) == 1: + try: flatten_axes(what, tree, (shardings,)) + except ValueError: pass # That's not the issue. + else: + msg += (f" Given the corresponding argument being " + f"passed, it looks like {what} might need to be wrapped in " + f"a singleton tuple.") + + raise ValueError(msg) + + axis_tree = shardings + + # Because we only have the `tree` treedef and not the full pytree here, + # we construct a dummy tree to compare against. Revise this in callers? + dummy_tree = tree_unflatten(tree, [PytreeLeaf()] * tree.num_leaves) + errors = prefix_errors(axis_tree, dummy_tree) + if errors: + e = errors[0] # Only show information about the first disagreement found. + raise e(what) + + # At this point we've failed to find a tree prefix error. + assert False, "Please open a bug report!" # This should be unreachable. + +class PytreeLeaf: + def __repr__(self): return "pytree leaf" + + +@lru_cache(maxsize=4096) +def _process_in_axis_resources(in_shardings_treedef, in_shardings_leaves, + in_layouts_treedef, in_layouts_leaves, + in_avals, in_tree, debug_info, + device_or_backend_set, kws): + if not kws: + in_tree, _ = treedef_children(in_tree) + + orig_in_shardings = tree_unflatten(in_shardings_treedef, in_shardings_leaves) + # Only do this if original in_shardings are unspecified. If it is AUTO, go + # via flatten_axis_resources. + if is_unspecified(orig_in_shardings): + in_shardings_flat = (orig_in_shardings,) * len(in_avals) + else: + in_shardings_flat = flatten_axis_resources( + "pjit in_shardings", in_tree, orig_in_shardings, tupled_args=True) + + in_layouts = tree_unflatten(in_layouts_treedef, in_layouts_leaves) + if in_layouts is None: + in_layouts_flat = (in_layouts,) * len(in_avals) + else: + in_layouts_flat = flatten_axis_resources( + "pjit in_layouts", in_tree, in_layouts, tupled_args=True) + + # TODO(dougalm,mattjj): enable debug info with attrs_tracked + attrs_tracked = debug_info and len(debug_info.arg_names) != len(in_avals) + if not config.dynamic_shapes.value and not attrs_tracked: + pjit_check_aval_sharding(in_shardings_flat, in_avals, + None if debug_info is None else debug_info.arg_names, + "pjit arguments", allow_uneven_sharding=False) + return in_shardings_flat, in_layouts_flat + +callsites: set[str] = set() + +def explain_tracing_cache_miss( + f: Callable, unseen_f: bool, cache: dict, key: tuple, result: tuple): + if config.check_tracer_leaks.value: return + + def unpack(key): + transforms, (), _, (in_type, _, debug_info, _, inline), *_, ctx = key + # TODO(dougalm,mattjj): enable cache miss explanation with attrs + _, (_, (in_tree,)), *_ = transforms + return in_tree, in_type, debug_info, inline.val, ctx + in_tree, in_type, debug_info, inline, ctx = unpack(key) + if inline: return + + msg: list[str] = [] + p = msg.append + done = lambda: logger.log(logging.WARNING, '\n'.join(msg)) + + callsite = source_info_util.summarize(source_info_util.current()) + p(f"TRACING CACHE MISS at {callsite} because:") + + # have we seen this function before at all? + fun_name = getattr(f, '__qualname__', f) + if debug_info.func_src_info: + _, _, *rest = debug_info.func_src_info.split(' ') + src_info = " defined at " + ' '.join(rest) + else: + src_info = '' + if unseen_f: + p(f" never seen function:\n {fun_name} id={id(f)}{src_info}") + if callsite in callsites: + p(" but seen another function defined on the same line; maybe the function is\n" + " being re-defined repeatedly, preventing caching?") + callsites.add(callsite) + return done() + else: + p(f" for {fun_name}{src_info}") + + seen_keys = map(unpack, cache.keys()) + + # have we maybe switched some args to be kwargs or visa-versa? + args_tree, kwargs_tree = treedef_children(in_tree) + args_kwargs_trees = [treedef_children(k) for k, *_ in seen_keys] + args_kwargs_match = [t for t in args_kwargs_trees + if t == [args_tree, kwargs_tree]] + if not args_kwargs_match: + num_args = len(treedef_children(args_tree)) + _, kwarg_keys = kwargs_tree.node_data() # type: ignore + p(f" never seen passing {num_args} positional args and {len(kwarg_keys)} " + "keyword args with keys:\n" + f" {', '.join(map(repr, kwarg_keys))}") + dont_match = [set(t[1].node_data()[1]) for t in args_kwargs_trees # type: ignore + if t != [args_tree, kwargs_tree]] + close_kwargs = min( + dont_match, key=set(kwarg_keys).symmetric_difference, default=None + ) + if not close_kwargs: + p(" closest seen is passing no keyword args") + else: + p(f" closest seen passes {len(close_kwargs)} keyword args with keys:\n" + f" {', '.join(map(repr, close_kwargs))}") + return done() + + # have we never seen this tracing context before? + ctxs_match = [c for *_, c in seen_keys if c == ctx] + if not ctxs_match: + p(" tracing context doesn't match, e.g. due to config or context manager") + dont_match = [c for *_, c in seen_keys if c != ctx] + closest_ctx = min(dont_match, key=lambda c: sum(map(op.ne, c, ctx))) + idxs = [i for i, (c1, c2) in enumerate(zip(ctx, closest_ctx)) if c1 != c2] + p(" closest seen context tuple differs at positions:\n" + f" {', '.join(map(str, idxs))}\n" + " compare to tuple returned by config._trace_context() in jax/_src/config.py.") + return done() + + # have we never seen this input pytree before? + trees_match = [k for k in seen_keys if k[0] == in_tree] + if not trees_match: + in_tree_str = f':\n {in_tree}' if len(str(in_tree)) < 76 else '' + p(f" never seen input pytree{in_tree_str}") + dont_match = [t for t, *_ in seen_keys if t != in_tree] + closest_tree = min(dont_match, key=lambda t: abs(t.num_leaves - in_tree.num_leaves)) + # TODO(mattjj): make equality_errors not print type name, avoid metaclass + leaf = type('LeafMeta', (type,), dict(__repr__=lambda _: 'leaf'))('Leaf', (), {})() + this_dummy = tree_unflatten(in_tree, [leaf] * in_tree.num_leaves) + close_dummy = tree_unflatten(closest_tree, [leaf] * closest_tree.num_leaves) # type: ignore + errs = list(tree_util.equality_errors(this_dummy, close_dummy)) + p(f" closest seen input pytree has {len(errs)} mismatches, including:") + for path, thing1, thing2, explanation in errs: + fst, *path = path # type: ignore + base = ['args', 'kwargs'][fst.idx] + p(f" * at {base}{keystr(tuple(path))}, seen {thing2} but now given {thing1}," + f" so {explanation}") + return done() + + # have we never seen these input types (eg shapes, dtypes) before? + types_match = [k for k in trees_match if k[1] == in_type] + if not types_match: + if len(in_type) < 5: + in_type_str = ':\n {}'.format(', '.join( + f'{n}: {ty.str_short(short_dtypes=True)}' + for n, ty in zip(debug_info.arg_names, in_type))) + else: + in_type_str = '' + p(f" never seen input type signature{in_type_str}") + dont_match = [t for _, t, *_ in trees_match if t != in_type] + closest_ty = min(dont_match, key=lambda t: sum(map(op.ne, t, in_type))) + num_mismatch = sum(map(op.ne, closest_ty, in_type)) + p(f" closest seen input type signature has {num_mismatch} mismatches, including:") + add_weak_type_hint = False + for name, ty1, ty2 in zip(debug_info.arg_names, closest_ty, in_type): + if ty1 != ty2: + if type(ty1) == type(ty2) == core.ShapedArray: + s1, s2 = ty1.str_short(True), ty2.str_short(True) + if s1 == s2: # weak types don't show up in str_short() + assert ty1.weak_type ^ ty2.weak_type + s1 += f'{{weak_type={ty1.weak_type}}}' + s2 += f'{{weak_type={ty2.weak_type}}}' + add_weak_type_hint = True + else: + s1, s2 = str(ty1), str(ty2) + p(f" * at {name}, seen {s1}, but now given {s2}") + if add_weak_type_hint: + p('where weak_type=True often means a Python builtin numeric value, and ') + p('weak_type=False means a jax.Array.') + p('See https://jax.readthedocs.io/en/latest/type_promotion.html#weak-types') + return done() + + # we think this is unreachable... + p("explanation unavailable! please open an issue at https://github.com/google/jax") + return done() + +@partial(lu.cache, explain=explain_tracing_cache_miss) +def _create_pjit_jaxpr(fun, in_type, attr_data, debug_info, out_paths, ignored_inline): + del ignored_inline # just for explain_cache_miss + with dispatch.log_elapsed_time( + "Finished tracing + transforming {fun_name} for pjit in {elapsed_time} sec", + fun_name=fun.__name__, event=dispatch.JAXPR_TRACE_EVENT): + pe_debug = debug_info and pe.debug_info_final(fun, debug_info.traced_for) + if config.dynamic_shapes.value: + jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic2( + lu.annotate(fun, in_type), debug_info=pe_debug) + attrs_tracked = [] + else: + jaxpr, global_out_avals, consts, attrs_tracked = pe.trace_to_jaxpr_dynamic( + fun, in_type, debug_info=pe_debug) + # assert attr_data is sentinel or attr_data matches attrs_tracked + + # TODO(dougalm,mattjj): enable debug info with attrs_tracked + if not config.dynamic_shapes.value and not attrs_tracked: + jaxpr = jaxpr_debug_info(jaxpr, debug_info, out_paths()) + + if config.debug_key_reuse.value: + # Import here to avoid circular imports + from jax.experimental.key_reuse._core import check_key_reuse_jaxpr + check_key_reuse_jaxpr(jaxpr) + + if any(isinstance(c, core.Tracer) for c in consts): + closed_jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr)) + final_consts = consts + else: + closed_jaxpr = core.ClosedJaxpr(jaxpr, consts) + final_consts = [] + return closed_jaxpr, final_consts, global_out_avals, attrs_tracked + + +@lru_cache(maxsize=4096) +def _check_and_canonicalize_out_shardings( + out_shardings_treedef, out_shardings_leaves, out_layouts_treedef, + out_layouts_leaves, out_tree, out_type, debug_info, device_or_backend_set): + orig_out_shardings = tree_unflatten(out_shardings_treedef, out_shardings_leaves) + if (is_unspecified(orig_out_shardings) or + isinstance(orig_out_shardings, sharding.Sharding)): + out_shardings_flat = (orig_out_shardings,) * len(out_type) + else: + out_shardings_flat = flatten_axis_resources( + "pjit out_shardings", out_tree(), orig_out_shardings, + tupled_args=False) + + out_layouts = tree_unflatten(out_layouts_treedef, out_layouts_leaves) + if out_layouts is None: + out_layouts_flat = (out_layouts,) * len(out_type) + else: + out_layouts_flat = flatten_axis_resources( + "pjit out_layouts", out_tree(), out_layouts, tupled_args=False) + + if not config.dynamic_shapes.value: + pjit_check_aval_sharding( + out_shardings_flat, out_type, + None if debug_info is None else debug_info.result_paths, + "pjit outputs", allow_uneven_sharding=False) + return out_shardings_flat, out_layouts_flat + + +AttrRecord = tuple[object, str, PyTreeDef, list[core.AbstractValue]] +_seen_attrs = weakref.WeakKeyDictionary() # type: ignore + +def seen_attrs_get(fun: lu.WrappedFun, in_type: core.InputType) -> list: + cache = _seen_attrs.setdefault(fun.f, defaultdict(list)) + assert fun.in_type is None or fun.in_type == in_type + return cache[(fun.transforms, fun.params, in_type)] + +def _attr_token(fun, in_type): + from jax.experimental.attrs import jax_getattr + cases = seen_attrs_get(fun, in_type) + for i, records in enumerate(cases): + for obj, attr, treedef, avals in records: + val = jax_getattr(obj, attr) + vals, treedef_ = tree_flatten(val) + avals_ = map(shaped_abstractify, vals) + if treedef != treedef_ or avals != avals_: break + else: + return i + return len(cases) + +def _attr_update(fun, in_type, i, attrs_tracked): + from jax.experimental.attrs import jax_getattr + leaves = lambda obj, attr: tree_leaves(jax_getattr(obj, attr)) + records = [(obj, attr, init_tree, map(shaped_abstractify, leaves(obj, attr))) + for init_tree, _, (obj, attr) in attrs_tracked] + cases = seen_attrs_get(fun, in_type) + if i == len(cases): + cases.append(records) + else: + assert i < len(cases) and cases[i] == records + +def _pjit_jaxpr(fun, out_shardings_treedef, out_shardings_leaves, + out_layouts_treedef, out_layouts_leaves, in_type, debug_info, + device_or_backend_set, out_tree, result_paths, inline): + attr_token = _attr_token(fun, in_type) + jaxpr, final_consts, out_type, attrs_tracked = _create_pjit_jaxpr( + fun, in_type, attr_token, debug_info, result_paths, IgnoreKey(inline)) + _attr_update(fun, in_type, attr_token, attrs_tracked) + canonicalized_out_shardings_flat, out_layouts_flat = _check_and_canonicalize_out_shardings( + out_shardings_treedef, out_shardings_leaves, out_layouts_treedef, + out_layouts_leaves, out_tree, tuple(out_type), + jaxpr.jaxpr.debug_info, device_or_backend_set) + return (jaxpr, final_consts, canonicalized_out_shardings_flat, + out_layouts_flat, attrs_tracked) + + +@dataclasses.dataclass(frozen=True) +class IgnoreKey: + val: Any + def __hash__(self): + return hash(self.__class__) + def __eq__(self, other): + return isinstance(other, IgnoreKey) # ignore self.val! + + +def pjit_check_aval_sharding( + shardings, flat_avals, names: tuple[str, ...] | None, + what_aval: str, allow_uneven_sharding: bool): + new_names = [''] * len(shardings) if names is None else names + for aval, s, name in zip(flat_avals, shardings, new_names): + if is_unspecified_or_auto(s): + continue + name_str = f' with pytree key path {name}' if name else '' + shape = aval.shape + try: + # Sharding interfaces can implement `check_compatible_aval` as an optional + # method to raise a more meaningful error. + if hasattr(s, 'check_compatible_aval'): + s.check_compatible_aval(shape) + else: + s._to_xla_hlo_sharding(len(shape)) + except ValueError as e: + raise ValueError( + f'One of {what_aval}{name_str} is incompatible with its sharding ' + f'annotation {s}: {e}') + # Use the `OpSharding` proto to find out how many ways each dimension of + # the aval is sharded. This approach will work across all + # Sharding. + hlo_sharding = s._to_xla_hlo_sharding(len(shape)) + assert hlo_sharding is not None + num_ways_dim_sharded, _ = op_shardings.get_num_ways_dim_sharded(hlo_sharding) + for i, size in enumerate(num_ways_dim_sharded): + if not allow_uneven_sharding and shape[i] % size != 0: + raise ValueError(f"One of {what_aval}{name_str} was given the sharding " + f"of {s}, which implies that " + f"the global size of its dimension {i} should be " + f"divisible by {size}, but it is equal to {shape[i]} " + f"(full shape: {shape})") + + +# -------------------- pjit rules -------------------- + +pjit_p = core.AxisPrimitive("pjit") +pjit_p.multiple_results = True + + +def _resolve_in_layouts(args, jit_in_layouts, resolved_in_shardings, in_avals): + # If device or backend is set, return the default layout. This is because you + # can pass arrays on cpu (with untiled layouts) to jit with backend='tpu' + # which causes error checks to fail. Returning the default layout allows + # this to exist. It's the same for handling shardings. + if pxla.check_device_backend_on_shardings(resolved_in_shardings): + return (None,) * len(jit_in_layouts) + + resolved_in_layouts = [] + for arg, jit_in_l, rs, aval in safe_zip( + args, jit_in_layouts, resolved_in_shardings, in_avals): + arg_layout, committed = ( + pxla._maybe_get_default_layout(getattr(arg, 'layout', None), jit_in_l, rs, aval), + getattr(arg, '_committed', True)) + # Sharding can be unspecified when array is committed if it's a PmapSharding. + is_pmap_sharding = (is_unspecified(rs) or + isinstance(getattr(arg, 'sharding', None), PmapSharding)) + if jit_in_l is None: + if committed: + if is_pmap_sharding: + resolved_in_layouts.append(None) + else: + resolved_in_layouts.append(arg_layout) + else: + resolved_in_layouts.append(None) + else: + # arg_layout can be None because some backends don't implement the + # required layout methods. Hence `arr.layout` can return + # `Layout(None, sharding)` + if (committed and not is_pmap_sharding and + arg_layout is not None and arg_layout != jit_in_l): + extra_msg = '' + if isinstance(jit_in_l, AutoLayout): + extra_msg = ( + ' The layout given to `jax.jit` is `DeviceLocalLayout.AUTO` but' + ' the corresponding argument passed is a `jax.Array` with a' + ' concrete layout. Consider passing a `jax.ShapeDtypeStruct`' + ' instead of `jax.Array` as an argument to the jitted function ' + ' when using `DeviceLocalLayout.AUTO`.' + ) + raise ValueError('Layout passed to jit does not match the layout ' + 'on the respective arg. ' + f'Got pjit layout: {jit_in_l},\n' + f'arg layout: {arg_layout} for ' + f'arg shape: {shaped_abstractify(arg).str_short()}.' + f'{extra_msg}') + resolved_in_layouts.append(jit_in_l) + return tuple(resolved_in_layouts) + + +def _resolve_in_shardings( + args, pjit_in_shardings: Sequence[PjitSharding], + out_shardings: Sequence[PjitSharding], + pjit_mesh: pxla.Mesh | None, + check_device_assignment: bool = True) -> Sequence[PjitSharding]: + # If True, means that device or backend is set by the user on pjit and it + # has the same semantics as device_put i.e. doesn't matter which device the + # arg is on, reshard it to the device mentioned. So don't do any of the + # checks and just return the pjit_in_shardings directly. `shard_args` will + # handle the resharding. + if pxla.check_device_backend_on_shardings(pjit_in_shardings): + return pjit_in_shardings + + committed_arg_shardings = [] + for a in args: + arg_s = getattr(a, 'sharding', None) + # arg sharding can be None in case of ShapeDtypeStruct. jax.Array does + # not allow None as the sharding. + if arg_s is None: + continue + if xla_extension_version < 270: + if not isinstance(arg_s, XLACompatibleSharding): + raise ValueError(f'One of the argument to pjit got sharding {arg_s} ' + 'which is not a subclass of XLACompatibleSharding.') + # Don't consider PmapSharding inputs as committed. They will get resharded + # unconditionally. + if isinstance(arg_s, PmapSharding): + continue + if getattr(a, '_committed', True): + committed_arg_shardings.append((arg_s, pxla.MismatchType.ARG_SHARDING, None)) + + # Check if the device_assignment across inputs, outputs and arguments is the + # same. + if check_device_assignment: + pxla._get_and_check_device_assignment( + it.chain( + util.stable_unique(committed_arg_shardings), + ((i, pxla.MismatchType.IN_SHARDING, None) + for i in util.stable_unique(pjit_in_shardings)), + ((o, pxla.MismatchType.OUT_SHARDING, None) + for o in util.stable_unique(out_shardings))), + (None if pjit_mesh is None or pjit_mesh.empty else list(pjit_mesh.devices.flat))) + + resolved_in_shardings = [] + for arg, pjit_in_s in zip(args, pjit_in_shardings): + # arg sharding can be None in case of ShapeDtypeStruct. jax.Array does + # not allow None as the sharding. + arg_s, committed = ((arg.sharding, getattr(arg, '_committed', True)) + if hasattr(arg, 'sharding') and arg.sharding is not None + else (UNSPECIFIED, False)) + if is_unspecified(pjit_in_s): + if is_unspecified(arg_s): + resolved_in_shardings.append(arg_s) + else: + if committed: + # If the arg has a PmapSharding, then reshard it unconditionally. + if isinstance(arg_s, PmapSharding): + resolved_in_shardings.append(UNSPECIFIED) + else: + resolved_in_shardings.append(arg_s) + else: + if dispatch.is_single_device_sharding(arg_s): + resolved_in_shardings.append(UNSPECIFIED) + else: + raise NotImplementedError('Having uncommitted Array sharded on ' + 'multiple devices is not supported.') + else: + if (isinstance(arg, np.ndarray) and + not pjit_in_s.is_fully_replicated and # type: ignore + xb.process_count() > 1): + raise ValueError( + 'Passing non-trivial shardings for numpy ' + 'inputs is not allowed. To fix this error, either specify a ' + 'replicated sharding explicitly or use ' + '`jax.experimental.multihost_utils.host_local_array_to_global_array(...)` ' + 'to convert your host local numpy inputs to a jax.Array which you ' + 'can pass to pjit. ' + 'If the numpy input is the same on each process, then you can use ' + '`jax.make_array_from_callback(...) to create a `jax.Array` which ' + 'you can pass to pjit. ' + 'Please see the jax.Array migration guide for more information ' + 'https://jax.readthedocs.io/en/latest/jax_array_migration.html#handling-of-host-local-inputs-to-pjit-like-batch-etc. ' + f'Got arg shape: {arg.shape}, arg value: {arg}') + if not is_unspecified(arg_s): + # jax.jit does not allow resharding across different memory kinds even + # if the argument is uncommitted. Use jax.device_put for those cases, + # either outside or inside jax.jit. + if pjit_in_s.memory_kind != arg_s.memory_kind: # type: ignore + raise ValueError( + 'Memory kinds passed to jax.jit does not match memory kind on the' + f' respective arg. Got pjit memory kind: {pjit_in_s.memory_kind}, ' # type: ignore + f'arg memory kind: {arg_s.memory_kind} for ' # pytype: disable=attribute-error + f'arg shape: {shaped_abstractify(arg).str_short()}') + if (committed and + not isinstance(arg_s, PmapSharding) and + not op_shardings.are_op_shardings_equal( + pjit_in_s._to_xla_hlo_sharding(arg.ndim), # type: ignore + arg_s._to_xla_hlo_sharding(arg.ndim))): + raise ValueError('Sharding passed to pjit does not match the sharding ' + 'on the respective arg. ' + f'Got pjit sharding: {pjit_in_s},\n' + f'arg sharding: {arg_s} for ' + f'arg shape: {shaped_abstractify(arg).str_short()}') + resolved_in_shardings.append(pjit_in_s) + + return tuple(resolved_in_shardings) + + +def _resolve_and_lower( + args, jaxpr, in_shardings, out_shardings, in_layouts, + out_layouts, resource_env, donated_invars, name, keep_unused, inline, + lowering_parameters, pgle_profiler=None): + in_shardings = _resolve_in_shardings( + args, in_shardings, out_shardings, + resource_env.physical_mesh if resource_env is not None else None) + in_layouts = _resolve_in_layouts(args, in_layouts, in_shardings, + jaxpr.in_avals) + lowered = _pjit_lower( + jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, resource_env, + donated_invars, name, keep_unused, inline, + lowering_parameters=lowering_parameters, + pgle_profiler=pgle_profiler) + return lowered + +def _pjit_call_impl_python( + *args, jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, + resource_env, donated_invars, name, keep_unused, inline): + global _most_recent_pjit_call_executable + + compile_options = None + pgle_profiler = None + pgle_profiler_dict = _most_recent_pjit_call_executable.weak_pgle_profiler_dict + if config.enable_pgle.value and config.pgle_profiling_runs.value > 0: + if jaxpr not in pgle_profiler_dict: + pgle_profiler_dict[jaxpr] = profiler.PGLEProfiler( + config.pgle_profiling_runs.value, + config.pgle_aggregation_percentile.value) + + pgle_profiler = pgle_profiler_dict[jaxpr] + # The method below will return FDO profile when module was profiled + # config.jax_pgle_profiling_runs amount of times, otherwise the result will + # be None. + fdo_profile = pgle_profiler.consume_fdo_profile() + if fdo_profile is not None: + compile_options = {'fdo_profile': fdo_profile} + + # TODO(patrios): Do not pass mutable profile session through cached lowering + # chain. Instead we need to move profilers dictionary to pxla module and use + # module as key. Right now we can't do that since there is no way to evict _pjit_lower_cached cache for in PGLE mode. + compiled = _resolve_and_lower( + args, jaxpr=jaxpr, in_shardings=in_shardings, + out_shardings=out_shardings, in_layouts=in_layouts, + out_layouts=out_layouts, resource_env=resource_env, + donated_invars=donated_invars, name=name, keep_unused=keep_unused, + inline=inline, lowering_parameters=mlir.LoweringParameters(), + pgle_profiler=pgle_profiler + ).compile(compile_options) + + _most_recent_pjit_call_executable.weak_key_dict[jaxpr] = compiled + # This check is expensive so only do it if enable_checks is on. + if compiled._auto_spmd_lowering and config.enable_checks.value: + pxla.check_array_xla_sharding_layout_match( + args, compiled._in_shardings, compiled._in_layouts, + jaxpr.jaxpr.debug_info, compiled._kept_var_idx) + if config.distributed_debug.value: + # Defensively only perform fingerprint logic if debug logging is enabled + # NOTE(skyewm): I didn't benchmark this + fingerprint = None + if hasattr(compiled.runtime_executable(), "fingerprint"): + fingerprint = compiled.runtime_executable().fingerprint + if fingerprint is not None: + fingerprint = fingerprint.hex() + distributed_debug_log(("Running pjit'd function", name), + ("in_shardings", in_shardings), + ("out_shardings", out_shardings), + ("in_layouts", in_layouts), + ("out_layouts", out_layouts), + ("abstract args", map(xla.abstractify, args)), + ("fingerprint", fingerprint)) + try: + return compiled.unsafe_call(*args), compiled + except FloatingPointError as e: + assert config.debug_nans.value or config.debug_infs.value # compiled_fun can only raise in this case + + if len(jaxpr.eqns) > 1: + _ = core.jaxpr_as_fun(jaxpr)(*args) # may raise, not return + + # If control reaches this line, we got a NaN on the output of `compiled` + # but not `fun.call_wrapped` on the same arguments. Let's tell the user. + msg = (f"{str(e)}. Because " + "jax_config.debug_nans.value and/or config.jax_debug_infs is set, the " + "de-optimized function (i.e., the function as if the `jit` " + "decorator were removed) was called in an attempt to get a more " + "precise error message. However, the de-optimized function did not " + "produce invalid values during its execution. This behavior can " + "result from `jit` optimizations causing the invalid value to be " + "produced. It may also arise from having nan/inf constants as " + "outputs, like `jax.jit(lambda ...: jax.numpy.nan)(...)`. " + "\n\n" + "It may be possible to avoid the invalid value by removing the " + "`jit` decorator, at the cost of losing optimizations. " + "\n\n" + "If you see this error, consider opening a bug report at " + "https://github.com/google/jax.") + raise FloatingPointError(msg) + + +@weakref_lru_cache +def _get_jaxpr_as_fun(jaxpr, in_shardings, out_shardings, in_layouts, + out_layouts, resource_env, donated_invars, name, + keep_unused, inline): + # The input jaxpr to `_get_jaxpr_as_fun` is under a weakref_lru_cache so + # returning `core.jaxpr_as_fun(jaxpr)` directly creates a strong reference to + # the jaxpr defeating the purpose of weakref_lru_cache. So return a function + # that closes over a weakrefed jaxpr and gets called inside that function. + # This way there won't be a strong reference to the jaxpr from the output + # function. + jaxpr = weakref.ref(jaxpr) + return lambda *args: core.jaxpr_as_fun(jaxpr())(*args) # pylint: disable=unnecessary-lambda + + +def _pjit_call_impl(*args, jaxpr, + in_shardings, out_shardings, in_layouts, out_layouts, + resource_env, + donated_invars, name, keep_unused, inline): + def call_impl_cache_miss(*args_, **kwargs_): + out_flat, compiled = _pjit_call_impl_python( + *args, jaxpr=jaxpr, in_shardings=in_shardings, + out_shardings=out_shardings, in_layouts=in_layouts, + out_layouts=out_layouts, resource_env=resource_env, + donated_invars=donated_invars, name=name, keep_unused=keep_unused, + inline=inline) + pgle_profiler = _read_pgle_profiler(jaxpr) + fastpath_data = _get_fastpath_data( + compiled, tree_structure(out_flat), args, out_flat, [], jaxpr.effects, + jaxpr.consts, None, pgle_profiler) + if xla_extension_version > 267: + return out_flat, fastpath_data, _need_to_rebuild_with_fdo(pgle_profiler) + else: + return out_flat, fastpath_data + + f = _get_jaxpr_as_fun( + jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, + resource_env, donated_invars, name, keep_unused, inline) + donated_argnums = [i for i, d in enumerate(donated_invars) if d] + has_explicit_sharding = _pjit_explicit_sharding( + in_shardings, out_shardings, None, None) + return xc._xla.pjit( + name, f, call_impl_cache_miss, [], [], donated_argnums, + tree_util.dispatch_registry, + pxla.shard_arg, + _get_cpp_global_cache(has_explicit_sharding))(*args) + +pjit_p.def_impl(_pjit_call_impl) + + +def _pjit_lower(*args, **kwargs): + return _pjit_lower_cached(*args, **kwargs) + + +@weakref_lru_cache +def _pjit_lower_cached( + jaxpr: core.ClosedJaxpr, + in_shardings, + out_shardings, + in_layouts: pxla.MaybeLayout, + out_layouts: pxla.MaybeLayout, + resource_env, + donated_invars, + name: str, + keep_unused: bool, + inline: bool, + *, + lowering_parameters: mlir.LoweringParameters, + pgle_profiler: profiler.PGLEProfiler | None): + if resource_env is not None: + pxla.resource_typecheck(jaxpr, resource_env, {}, lambda: "pjit") + + if resource_env is not None: + mesh = resource_env.physical_mesh + api_name = 'pjit' + else: + # resource_env is `None` in the jit wrapper around pjit. + mesh = None + api_name = 'jit' + + # For `pjit(xmap)` cases, it needs to take the `lower_mesh_computation` path + # because `xmap` only supports SPMDAxisContext right now. + if dispatch.jaxpr_has_primitive(jaxpr.jaxpr, 'xmap'): + return pxla.lower_mesh_computation( + jaxpr, api_name, name, mesh, + in_shardings, out_shardings, donated_invars, + True, jaxpr.in_avals, tiling_method=None, + lowering_parameters=lowering_parameters) + else: + return pxla.lower_sharding_computation( + jaxpr, api_name, name, in_shardings, out_shardings, + in_layouts, out_layouts, tuple(donated_invars), + keep_unused=keep_unused, inline=inline, + devices_from_context=( + None if mesh is None or mesh.empty else list(mesh.devices.flat)), + lowering_parameters=lowering_parameters, + pgle_profiler=pgle_profiler) + + +def pjit_staging_rule(trace, *args, **params): + jaxpr, in_fwd, out_shardings, out_layouts = _pjit_forwarding( + params['jaxpr'], params['out_shardings'], params['out_layouts']) + params = dict(params, jaxpr=jaxpr, out_shardings=out_shardings, + out_layouts=out_layouts) + + if (params["inline"] and + all(is_unspecified(i) for i in params["in_shardings"]) and + all(is_unspecified(o) for o in params["out_shardings"]) and + all(i is None for i in params["in_layouts"]) and + all(o is None for o in params["out_layouts"])): + + if config.dynamic_shapes.value: + # Inline jaxpr doesn't handle dynamic shapes when inlining. If dynamic + # shapes are enabled, use eval_jaxpr, which uses the tracing machinery, + # but redundantly performs abstract evaluation again. + out_tracers = core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, *args, + propagate_source_info=False) + else: + out_tracers = pe.inline_jaxpr_into_trace( + trace, jaxpr.jaxpr, jaxpr.consts, *args) + elif config.dynamic_shapes.value: + source_info = source_info_util.current() + out_tracers = [] + for aval in _out_type(jaxpr): + if type(aval) is core.DShapedArray: + shape = [args[d.val] if type(d) is core.InDBIdx else + out_tracers[d.val] if type(d) is core.OutDBIdx else + d for d in aval.shape] + aval = aval.update(shape=tuple(core.get_referent(d) for d in shape)) + out_tracers.append(pe.DynamicJaxprTracer(trace, aval, source_info)) + eqn = core.new_jaxpr_eqn( + map(trace.getvar, args), map(trace.makevar, out_tracers), pjit_p, params, + jaxpr.effects, source_info) + trace.frame.add_eqn(eqn) + elif any(isinstance(c, core.MutableArray) for c in jaxpr.consts): + jaxpr, consts = pxla._move_mutable_consts(jaxpr) + consts = map(trace.instantiate_const, consts) + in_shardings = (*params['in_shardings'],) + (UNSPECIFIED,) * len(consts) + in_layouts = (*params['in_layouts'],) + (None,) * len(consts) + donated_invars = (*params['donated_invars'],) + (False,) * len(consts) + new_params = dict(params, jaxpr=jaxpr, in_shardings=in_shardings, + in_layouts=in_layouts, donated_invars=donated_invars) + out_tracers = trace.default_process_primitive( + pjit_p, (*args, *consts), new_params) + else: + out_tracers = trace.default_process_primitive(pjit_p, args, params) + + out_tracers_ = iter(out_tracers) + out_tracers = [args[f] if type(f) is int else next(out_tracers_) + for f in in_fwd] + assert next(out_tracers_, None) is None + return out_tracers +pe.custom_staging_rules[pjit_p] = pjit_staging_rule + + +def _pjit_forwarding(jaxpr, out_shardings, out_layouts): + in_fwd: list[int | None] = pe._jaxpr_forwarding(jaxpr.jaxpr) + in_fwd = [fwd if is_unspecified(os) and ol is None else None for fwd, os, ol + in zip(in_fwd, out_shardings, out_layouts)] + keep = [f is None for f in in_fwd] + jaxpr = pe.prune_closed_jaxpr_outputs(jaxpr, keep) + out_shardings = [o for o, k in zip(out_shardings, keep) if k] + out_layouts = [o for o, k in zip(out_layouts , keep) if k] + return jaxpr, in_fwd, out_shardings, out_layouts + +def pjit_forwarding_rule(eqn): + jaxpr, in_fwd, out_shardings, out_layouts = _pjit_forwarding( + eqn.params['jaxpr'], eqn.params['out_shardings'], eqn.params['out_layouts']) + new_outvars = [v for v, f in zip(eqn.outvars, in_fwd) if f is None] + new_params = dict(eqn.params, jaxpr=jaxpr, out_shardings=(*out_shardings,), + out_layouts=(*out_layouts,)) + new_eqn = eqn.replace(params=new_params, outvars=new_outvars) + fwd_vars = [eqn.invars[f] if f is not None else None for f in in_fwd] + return fwd_vars, new_eqn +pe.forwarding_rules[pjit_p] = pjit_forwarding_rule + + +# TODO(mattjj): remove/trivialize this when jaxprs have type annotation on them, +# since it's actually not possible in general to infer the type from the term +def _out_type(jaxpr: core.ClosedJaxpr) -> list[core.AbstractValue]: + out = [] + in_idx = {v: i for i, v in enumerate(jaxpr.jaxpr.invars)} + out_idx = {x: i for i, x in enumerate(jaxpr.jaxpr.invars) + if type(x) is core.Var} + for x in jaxpr.jaxpr.outvars: + aval = x.aval + if type(aval) is core.DShapedArray: + shape = [core.InDBIdx(in_idx[d]) if d in in_idx else + core.OutDBIdx(out_idx[d]) if d in out_idx else + d for d in x.aval.shape] + aval = aval.update(shape=tuple(shape)) + out.append(aval) + return out + + +def _pjit_typecheck(ctx_factory, *in_atoms, jaxpr, **params): + return core._check_call(ctx_factory, pjit_p, in_atoms, + dict(params, call_jaxpr=jaxpr.jaxpr)) +core.custom_typechecks[pjit_p] = _pjit_typecheck + + +def _pjit_abstract_eval(*args, jaxpr, **_): + return jaxpr.out_avals, jaxpr.effects +pjit_p.def_effectful_abstract_eval(_pjit_abstract_eval) + + +def _pjit_cached_lower_jaxpr_to_fun(ctx, name, jaxpr, effects, in_shardings, + out_shardings, in_layouts, out_layouts, + api_name): + mod_ctx = ctx.module_context + axis_ctx = ctx.module_context.axis_context + num_devices = None + if isinstance(axis_ctx, sharding_impls.ShardingContext): + num_devices = axis_ctx.num_devices + elif isinstance(axis_ctx, sharding_impls.SPMDAxisContext): + num_devices = axis_ctx.mesh.size + key = (pjit_p, name, jaxpr, effects, num_devices, + pxla.SemanticallyEqualShardings(in_shardings, jaxpr.in_avals), + pxla.SemanticallyEqualShardings(out_shardings, jaxpr.out_avals), + in_layouts, out_layouts, api_name) + + func = mod_ctx.cached_primitive_lowerings.get(key, None) + if func is None: + arg_shardings = [None if is_unspecified(i) else i for i in in_shardings] + result_shardings = [None if is_unspecified(o) else o for o in out_shardings] + # TODO(b/228598865): inlined calls cannot have shardings set directly on the + # inputs or outputs because they are lost during MLIR->HLO conversion. + # using_sharding_annotation=False means we add an identity operation instead. + func = mlir.lower_jaxpr_to_fun( + mod_ctx, name, jaxpr, effects, ctx.name_stack, + arg_shardings=arg_shardings, result_shardings=result_shardings, + use_sharding_annotations=False, api_name=api_name, + arg_layouts=in_layouts, result_layouts=out_layouts) + mod_ctx.cached_primitive_lowerings[key] = func + return func + + +def _pjit_lowering(ctx, *args, name, jaxpr, in_shardings, + out_shardings, in_layouts, out_layouts, resource_env, + donated_invars, keep_unused, inline): + effects = list(ctx.tokens_in.effects()) + output_types = map(mlir.aval_to_ir_types, ctx.avals_out) + output_types = [mlir.token_type()] * len(effects) + output_types + flat_output_types = flatten(output_types) + + func = _pjit_cached_lower_jaxpr_to_fun( + ctx, name, jaxpr, tuple(effects), in_shardings, + out_shardings, in_layouts, out_layouts, + api_name=('jit' if resource_env is None else 'pjit')) + + tokens_in = [ctx.tokens_in.get(eff) for eff in effects] + args = (*ctx.dim_var_values, *tokens_in, *args) + call = func_dialect.CallOp(flat_output_types, + ir.FlatSymbolRefAttr.get(func.name.value), + mlir.flatten_lowering_ir_args(args)) + mlir.wrap_compute_type_in_place(ctx, call) + out_nodes = unflatten(call.results, map(len, output_types)) + tokens, out_nodes = split_list(out_nodes, [len(effects)]) + tokens_out = ctx.tokens_in.update_tokens(mlir.TokenSet(zip(effects, tokens))) + ctx.set_tokens_out(tokens_out) + return out_nodes + +mlir.register_lowering(pjit_p, _pjit_lowering) + + +def _pjit_batcher(insert_axis, spmd_axis_name, + axis_size, axis_name, main_type, + vals_in, dims_in, + jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, + resource_env, donated_invars, name, keep_unused, inline): + segment_lens, dims_in = batching.indirectify_ragged_axes(dims_in) + new_jaxpr, axes_out = batching.batch_jaxpr2( + jaxpr, axis_size, dims_in, axis_name=axis_name, + spmd_axis_name=spmd_axis_name, main_type=main_type) + + # `insert_axis` is set to True only for some `xmap` uses. + new_parts = (axis_name,) if insert_axis else ( + () if spmd_axis_name is None else spmd_axis_name) + + if resource_env is not None: + mesh = resource_env.physical_mesh + else: + mesh = None + + # TODO(axch): prepend with Nones (?) to account for new segment_lens inputs + in_shardings = tuple( + _pjit_batcher_for_sharding(i, axis_in, new_parts, mesh, aval.ndim) + if axis_in is not None else i + for axis_in, i, aval in zip(dims_in, in_shardings, new_jaxpr.in_avals)) + out_shardings = tuple( + _pjit_batcher_for_sharding(o, axis_out, new_parts, mesh, aval.ndim) + if axis_out is not None else o + for axis_out, o, aval in zip(axes_out, out_shardings, new_jaxpr.out_avals)) + # TODO(yashkatariya): Figure out layouts should change under vmap. + if not (all(l is None for l in in_layouts) and + all(l is None for l in out_layouts)): + raise NotImplementedError + + vals_out = pjit_p.bind( + *vals_in, + jaxpr=new_jaxpr, + in_shardings=in_shardings, + out_shardings=out_shardings, + in_layouts=in_layouts, + out_layouts=out_layouts, + resource_env=resource_env, + donated_invars=donated_invars, + name=name, + keep_unused=keep_unused, + inline=inline) + + resolved_axes_out = batching.resolve_ragged_axes_against_inputs_outputs( + vals_in, vals_out, axes_out) + return vals_out, resolved_axes_out + +batching.spmd_axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, False) +batching.axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, False, None) +pxla.spmd_primitive_batchers[pjit_p] = partial(_pjit_batcher, True, None) + +def _pjit_batcher_for_sharding( + s: sharding.Sharding | UnspecifiedValue, + dim: int, val: tuple[str, ...], mesh, ndim: int): + if is_unspecified(s): + return s + hlo_s = s._to_xla_hlo_sharding(ndim) # type: ignore + if not val: + if sharding_impls.is_op_sharding_replicated(hlo_s): + return s + new_op = hlo_s.to_proto().clone() + tad = list(new_op.tile_assignment_dimensions) + tad.insert(dim, 1) + new_op.tile_assignment_dimensions = tad + new_gs = GSPMDSharding( + s._device_assignment, new_op, # type: ignore + _device_list=getattr(s, '_internal_device_list', None)) + return pxla._get_out_sharding_from_orig_sharding([new_gs], [None], s, None)[0] + else: + if isinstance(s, NamedSharding): + mesh = s.mesh + if mesh is None or mesh.empty: + raise ValueError( + 'If you are using xmap or spmd_axis_name parameter of jax.vmap,' + ' please make sure to run your jitted function inside the mesh' + ' context manager. Only `jax.lax.with_sharding_constraint` with' + ' `jax.sharding.NamedSharding` as an input can be transformed with' + ' spmd_axis_name batching rules outside of an explicit mesh context' + f' manager scope{s!r}') + parsed_pspec = parse_flatten_op_sharding(hlo_s, mesh)[0] + parsed_pspec = parsed_pspec.insert_axis_partitions(dim, val) + return NamedSharding._from_parsed_pspec(mesh, parsed_pspec) + + +def _pjit_jvp(primals_in, tangents_in, + jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, + resource_env, donated_invars, name, keep_unused, inline): + if any(isinstance(c, core.MutableArray) for c in jaxpr.consts): + jaxpr, mut_primals = pxla._move_mutable_consts(jaxpr) + mut_tangents = map(ad_util.zeros_like_jaxval, mut_primals) + primals_in = [*primals_in, *mut_primals] + tangents_in = [*tangents_in, *mut_tangents] + in_shardings = (*in_shardings,) + (UNSPECIFIED,) * len(mut_primals) + in_layouts = (*in_layouts,) + (None,) * len(mut_primals) + donated_invars = (*donated_invars,) + (False,) * len(mut_primals) + + tangents_in = [ad_util.zeros_like_aval(a) if isinstance(a, AbstractRef) else x + for x, a in zip(tangents_in, jaxpr.in_avals)] + + is_nz_tangents_in = [type(t) is not ad.Zero for t in tangents_in] + jaxpr_jvp, is_nz_tangents_out = ad.jvp_jaxpr( + jaxpr, is_nz_tangents_in, instantiate=False) + + def _filter_zeros(is_nz_l, l): + return (x for nz, x in zip(is_nz_l, l) if nz) + _filter_zeros_in = partial(_filter_zeros, is_nz_tangents_in) + _filter_zeros_out = partial(_filter_zeros, is_nz_tangents_out) + outputs = pjit_p.bind( + *primals_in, *_filter_zeros_in(tangents_in), + jaxpr=jaxpr_jvp, + in_shardings=(*in_shardings, *_filter_zeros_in(in_shardings)), + out_shardings=(*out_shardings, *_filter_zeros_out(out_shardings)), + in_layouts=(*in_layouts, *_filter_zeros_in(in_layouts)), + out_layouts=(*out_layouts, *_filter_zeros_out(out_layouts)), + resource_env=resource_env, + donated_invars=(*donated_invars, *_filter_zeros_in(donated_invars)), + name=name, + keep_unused=keep_unused, + inline=inline) + + primals_out, tangents_out = split_list(outputs, [len(jaxpr.jaxpr.outvars)]) + assert len(primals_out) == len(jaxpr.jaxpr.outvars) + tangents_out_it = iter(tangents_out) + return primals_out, [next(tangents_out_it) if nz else ad.Zero(aval) + for nz, aval in zip(is_nz_tangents_out, jaxpr.out_avals)] +ad.primitive_jvps[pjit_p] = _pjit_jvp + + +@weakref_lru_cache +def _known_jaxpr_fwd(known_jaxpr: core.ClosedJaxpr, + in_fwd: tuple[int | None, ...]) -> core.ClosedJaxpr: + updated_jaxpr = known_jaxpr.jaxpr.replace( + outvars=[x for x, i in zip(known_jaxpr.jaxpr.outvars, in_fwd) + if i is None]) + return known_jaxpr.replace(jaxpr=updated_jaxpr) + + +def _pjit_partial_eval(trace, *in_tracers, + jaxpr, in_shardings, out_shardings, + in_layouts, out_layouts, resource_env, donated_invars, + name, keep_unused, inline): + in_pvals = [t.pval for t in in_tracers] + + known_ins = tuple(pv.is_known() for pv in in_pvals) + unknown_ins = tuple(not k for k in known_ins) + known_jaxpr, unknown_jaxpr, unknown_outs, res_avals = \ + pe.partial_eval_jaxpr_nounits(jaxpr, unknown_ins, instantiate=False) + unknown_outs = tuple(unknown_outs) + known_outs = tuple(not uk for uk in unknown_outs) + num_residuals = len(res_avals) + res_shardings = (UNSPECIFIED,) * num_residuals + res_layouts = (None,) * num_residuals + + def keep_where(l, should_keep): + return tuple(x for x, keep in zip(l, should_keep) if keep) + + known_out_shardings = keep_where(out_shardings, known_outs) + res_shardings + known_out_layouts = keep_where(out_layouts, known_outs) + res_layouts + + # Input-to-output forwarding: compute which outputs are just forwarded inputs. + num_out_primals = len(known_jaxpr.out_avals) - num_residuals + in_fwd: list[int | None] = pe._jaxpr_forwarding(known_jaxpr.jaxpr) + # Only forward primal outputs when corresponding out_sharding is UNSPECIFIED. + in_fwd_primal, in_fwd_res = split_list(in_fwd, [num_out_primals]) + in_fwd = [ + fwd if is_unspecified(os) and ol is None else None + for os, ol, fwd in zip( + keep_where(out_shardings, known_outs), + keep_where(out_layouts, known_outs), in_fwd_primal) + ] + in_fwd_res + del in_fwd_primal, in_fwd_res + # Prune jaxpr outputs and out_shardings by removing the input-forwards. + keep = [f is None for f in in_fwd] + known_jaxpr = pe.prune_closed_jaxpr_outputs(known_jaxpr, keep) + known_out_shardings = keep_where(known_out_shardings, keep) + known_out_layouts = keep_where(known_out_layouts, keep) + # Update num_out_primals to reflect pruning. + kept_primals, kept_res = split_list(keep, [num_out_primals]) + num_out_primals = sum(kept_primals) + del keep, kept_primals, kept_res + + # Output-to-output forwarding: compute which residuals are just primal outputs + out_vars, res_vars = split_list(known_jaxpr.jaxpr.outvars, [num_out_primals]) + idx_map = {id(v): i for i, v in enumerate(out_vars)} + out_fwd = [None] * num_out_primals + [idx_map.get(id(v)) for v in res_vars] + # Prune jaxpr outputs and out_shardings by removing forwarded residuals. + keep = [f is None for f in out_fwd] + known_jaxpr = pe.prune_closed_jaxpr_outputs(known_jaxpr, keep) + known_out_shardings = keep_where(known_out_shardings, keep) + known_out_layouts = keep_where(known_out_layouts, keep) + del keep + + known_params = dict( + jaxpr=known_jaxpr, in_shardings=keep_where(in_shardings, known_ins), + out_shardings=known_out_shardings, + in_layouts=keep_where(in_layouts, known_ins), + out_layouts=known_out_layouts, resource_env=resource_env, + donated_invars=keep_where(donated_invars, known_ins), + name=name, keep_unused=keep_unused, inline=inline) + assert len(known_params['out_shardings']) == len(known_params['jaxpr'].out_avals) + assert len(known_params['out_layouts']) == len(known_params['jaxpr'].out_avals) + + # Bind known things to pjit_p. + known_inputs = [pv.get_known() for pv in in_pvals if pv.is_known()] + all_known_outs = pjit_p.bind(*known_inputs, **known_params) + # Add back in the output fwds. + all_known_outs = subs_list(out_fwd, all_known_outs, all_known_outs) + # Add back in the input fwds. + all_known_outs = subs_list(in_fwd, known_inputs, all_known_outs) + + known_out_vals, residual_vals = \ + split_list(all_known_outs, [len(all_known_outs) - num_residuals]) + residual_tracers = map(trace.new_instantiated_const, residual_vals) + + # The convention of partial_eval_jaxpr_nounits is to place residual binders at + # the front of the jaxpr produced, so we move them to the back since both the + # jaxpr equation built below and the pjit transpose rule assume a + # residual-inputs-last convention. + unknown_jaxpr = pe.move_binders_to_back( + unknown_jaxpr, [True] * num_residuals + [False] * sum(unknown_ins)) + # Prepare unknown tracers + unknown_params = dict( + jaxpr=unknown_jaxpr, + in_shardings=(keep_where(in_shardings, unknown_ins) + res_shardings), + out_shardings=keep_where(out_shardings, unknown_outs), + in_layouts=(keep_where(in_layouts, unknown_ins) + res_layouts), + out_layouts=keep_where(out_layouts, unknown_outs), + resource_env=resource_env, + donated_invars=(keep_where(donated_invars, unknown_ins) + + (False,) * num_residuals), + name=name, + keep_unused=keep_unused, + inline=inline) + unknown_tracers_in = [t for t in in_tracers if not t.pval.is_known()] + unknown_out_avals = unknown_jaxpr.out_avals + unknown_tracers_out = [ + pe.JaxprTracer(trace, pe.PartialVal.unknown(aval), None) + for aval in unknown_out_avals + ] + eqn = pe.new_eqn_recipe((*unknown_tracers_in, *residual_tracers), + unknown_tracers_out, + pjit_p, + unknown_params, + unknown_jaxpr.effects, + source_info_util.current()) + for t in unknown_tracers_out: t.recipe = eqn + return merge_lists(unknown_outs, known_out_vals, unknown_tracers_out) + +pe.custom_partial_eval_rules[pjit_p] = _pjit_partial_eval + + +def _pjit_partial_eval_custom_params_updater( + unks_in: Sequence[bool], inst_in: Sequence[bool], + kept_outs_known: Sequence[bool], kept_outs_staged: Sequence[bool], + num_res_out: int, num_res_in: int, params_known: dict, params_staged: dict + ) -> tuple[dict, dict]: + # prune inputs to jaxpr_known according to unks_in + donated_invars_known, _ = pe.partition_list(unks_in, params_known['donated_invars']) + in_shardings_known, _ = pe.partition_list(unks_in, params_known['in_shardings']) + _, out_shardings_known = pe.partition_list(kept_outs_known, params_known['out_shardings']) + in_layouts_known, _ = pe.partition_list(unks_in, params_known['in_layouts']) + _, out_layouts_known = pe.partition_list(kept_outs_known, params_known['out_layouts']) + + new_params_known = dict(params_known, + in_shardings=tuple(in_shardings_known), + out_shardings=(*out_shardings_known, + *[UNSPECIFIED] * num_res_out), + in_layouts=tuple(in_layouts_known), + out_layouts=(*out_layouts_known, *[None] * num_res_out), + donated_invars=tuple(donated_invars_known)) + assert len(new_params_known['in_shardings']) == len(params_known['jaxpr'].in_avals) + assert len(new_params_known['out_shardings']) == len(params_known['jaxpr'].out_avals) + assert len(new_params_known['in_layouts']) == len(params_known['jaxpr'].in_avals) + assert len(new_params_known['out_layouts']) == len(params_known['jaxpr'].out_avals) + + # added num_res new inputs to jaxpr_staged, and pruning according to inst_in + _, donated_invars_staged = pe.partition_list(inst_in, params_staged['donated_invars']) + donated_invars_staged = [False] * num_res_in + donated_invars_staged + _, in_shardings_staged = pe.partition_list(inst_in, params_staged['in_shardings']) + in_shardings_staged = [*[UNSPECIFIED] * num_res_in, *in_shardings_staged] + _, out_shardings_staged = pe.partition_list(kept_outs_staged, params_staged['out_shardings']) + _, in_layouts_staged = pe.partition_list(inst_in, params_staged['in_layouts']) + in_layouts_staged = [*[None] * num_res_in, *in_layouts_staged] + _, out_layouts_staged = pe.partition_list(kept_outs_staged, params_staged['out_layouts']) + + new_params_staged = dict(params_staged, + in_shardings=tuple(in_shardings_staged), + out_shardings=tuple(out_shardings_staged), + in_layouts=tuple(in_layouts_staged), + out_layouts=tuple(out_layouts_staged), + donated_invars=tuple(donated_invars_staged)) + assert len(new_params_staged['in_shardings']) == len(params_staged['jaxpr'].in_avals) + assert len(new_params_staged['out_shardings']) == len(params_staged['jaxpr'].out_avals) + assert len(new_params_staged['in_layouts']) == len(params_staged['jaxpr'].in_avals) + assert len(new_params_staged['out_layouts']) == len(params_staged['jaxpr'].out_avals) + return new_params_known, new_params_staged + +pe.partial_eval_jaxpr_custom_rules[pjit_p] = \ + partial(pe.closed_call_partial_eval_custom_rule, 'jaxpr', + _pjit_partial_eval_custom_params_updater) + + +@lu.cache +def _pjit_transpose_trace(fun, in_avals): + transpose_jaxpr, _, consts, attrs_tracked = pe.trace_to_jaxpr_dynamic( + fun, in_avals) + transpose_jaxpr = core.ClosedJaxpr(transpose_jaxpr, consts) + return transpose_jaxpr, attrs_tracked + + +def _pjit_transpose(cts_in, *primals_in, + jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, + resource_env, donated_invars, name, keep_unused, inline): + def prune_type(ty, xs, maybe_zeros): + return tuple(x for x, mz in zip(xs, maybe_zeros) if type(mz) is not ty) + + body = lu.wrap_init(ad.closed_backward_pass) + body = lu.hashable_partial(body, jaxpr, False) + primals_and_nz_cts_in, in_treedef = tree_flatten((primals_in, cts_in)) + body, cts_out_treedef_thunk = flatten_fun_nokwargs(body, in_treedef) + + transpose_in_shardings = ( + *prune_type(ad.UndefinedPrimal, in_shardings, primals_in), + *prune_type(ad.Zero, out_shardings, cts_in) + ) + transpose_in_layouts = ( + *prune_type(ad.UndefinedPrimal, in_layouts, primals_in), + *prune_type(ad.Zero, out_layouts, cts_in) + ) + global_cts_in_avals = tuple(core.raise_to_shaped(core.get_aval(ct)) + for ct in primals_and_nz_cts_in) + + transpose_jaxpr, attrs_tracked = _pjit_transpose_trace( + body, global_cts_in_avals) + cts_out_treedef = cts_out_treedef_thunk() + transpose_out_shardings = prune_type( + ad.Zero, + in_shardings, + tree_unflatten(cts_out_treedef, [object()] * cts_out_treedef.num_leaves)) + transpose_out_layouts = prune_type( + ad.Zero, + in_layouts, + tree_unflatten(cts_out_treedef, [object()] * cts_out_treedef.num_leaves)) + + if attrs_tracked: + init_states = _get_states(attrs_tracked) + primals_and_nz_cts_in = [*init_states, *primals_and_nz_cts_in] + transpose_in_shardings = (UNSPECIFIED,) * len(attrs_tracked) + transpose_in_shardings + transpose_out_shardings = (UNSPECIFIED,) * len(attrs_tracked) + transpose_out_shardings + transpose_in_layouts = (None,) * len(attrs_tracked) + transpose_in_layouts + transpose_out_layouts = (None,) * len(attrs_tracked) + transpose_out_layouts + + nz_cts_out = pjit_p.bind( + *primals_and_nz_cts_in, + jaxpr=transpose_jaxpr, + in_shardings=transpose_in_shardings, + out_shardings=transpose_out_shardings, + in_layouts=transpose_in_layouts, + out_layouts=transpose_out_layouts, + resource_env=resource_env, + donated_invars=(False,) * len(primals_and_nz_cts_in), + name=name, + keep_unused=keep_unused, + inline=inline) + + if attrs_tracked: + final_states, nz_cts_out = split_list(nz_cts_out, [len(init_states)]) + _set_states(attrs_tracked, final_states) + + return tree_unflatten(cts_out_treedef, nz_cts_out) +ad.reducing_transposes[pjit_p] = _pjit_transpose + + +@weakref_lru_cache +def _dce_jaxpr_pjit( + jaxpr: core.ClosedJaxpr, used_outputs: tuple[bool, ...] +) -> tuple[core.ClosedJaxpr, list[bool]]: + new_jaxpr, used_inputs = pe.dce_jaxpr(jaxpr.jaxpr, used_outputs) + return core.ClosedJaxpr(new_jaxpr, jaxpr.consts), used_inputs + + +def dce_jaxpr_pjit_rule(used_outputs: list[bool], eqn: core.JaxprEqn + ) -> tuple[list[bool], core.JaxprEqn | None]: + dced_jaxpr, used_inputs = _dce_jaxpr_pjit( + eqn.params['jaxpr'], tuple(used_outputs)) + + def keep_where(xs, keeps): + return tuple(x for x, keep in zip(xs, keeps) if keep) + + eqn_params = eqn.params + new_params = dict( + eqn_params, + jaxpr=dced_jaxpr, + in_shardings=keep_where(eqn_params["in_shardings"], used_inputs), + out_shardings=keep_where(eqn_params["out_shardings"], used_outputs), + in_layouts=keep_where(eqn_params["in_layouts"], used_inputs), + out_layouts=keep_where(eqn_params["out_layouts"], used_outputs), + donated_invars=keep_where(eqn_params["donated_invars"], used_inputs), + ) + if not any(used_inputs) and not any(used_outputs) and not dced_jaxpr.effects: + return used_inputs, None + else: + new_eqn = core.new_jaxpr_eqn( + [v for v, used in zip(eqn.invars, used_inputs) if used], + [v for v, used in zip(eqn.outvars, used_outputs) if used], + eqn.primitive, new_params, dced_jaxpr.effects, eqn.source_info, eqn.ctx) + return used_inputs, new_eqn + +pe.dce_rules[pjit_p] = dce_jaxpr_pjit_rule + + +def _check_resources_against_named_axes(what, aval, pos_axis_resources, named_axis_resources): + pjit_resources = set( + it.chain.from_iterable([d for d in pos_axis_resources if d is not None])) + aval_resources = set(it.chain.from_iterable( + named_axis_resources[a] for a in aval.named_shape)) + overlap = pjit_resources & aval_resources + if overlap: + raise JAXTypeError( + f"{what} has an axis resources specification of " + f"{pos_axis_resources.unsynced_user_spec(SpecSync.DIM_PERMUTE)} " + f"that uses one or more mesh axes already used by xmap to partition " + f"a named axis appearing in its named_shape (both use mesh axes " + f"{mesh_lib.show_axes(overlap)})") + +def _resource_typing_pjit(avals, params, source_info, resource_env, named_axis_resources): + jaxpr = params["jaxpr"] + what = "pjit input" + if (resource_env is not None and params['resource_env'] is not None and + resource_env.physical_mesh != params['resource_env'].physical_mesh): + raise RuntimeError("Changing the physical mesh is not allowed inside pjit.") + + for aval, s in zip(jaxpr.in_avals, params['in_shardings']): + if is_unspecified(s) or is_auto(s): + continue + elif hasattr(s, '_parsed_pspec'): + parsed_pspec = s._parsed_pspec + else: + if resource_env is not None and not resource_env.physical_mesh.empty: + parsed_pspec = parse_flatten_op_sharding( + s._to_xla_hlo_sharding(aval.ndim), resource_env.physical_mesh)[0] + else: + parsed_pspec = None + if parsed_pspec is not None: + _check_resources_against_named_axes(what, aval, parsed_pspec, + named_axis_resources) + + pxla.resource_typecheck( + jaxpr.jaxpr, resource_env, named_axis_resources, + lambda: (f"a pjit'ed function {params['name']} " + f"(pjit called at {source_info_util.summarize(source_info)})")) + + what = "pjit output" + for aval, s in zip(jaxpr.out_avals, params['out_shardings']): + if is_unspecified(s) or is_auto(s): + continue + elif hasattr(s, '_parsed_pspec'): + parsed_pspec = s._parsed_pspec + else: + if resource_env is not None and not resource_env.physical_mesh.empty: + parsed_pspec = parse_flatten_op_sharding( + s._to_xla_hlo_sharding(aval.ndim), resource_env.physical_mesh)[0] + else: + parsed_pspec = None + if parsed_pspec is not None: + _check_resources_against_named_axes(what, aval, parsed_pspec, + named_axis_resources) + +pxla.custom_resource_typing_rules[pjit_p] = _resource_typing_pjit + + +def _pjit_pp_rule(eqn, context, settings): + params = dict(eqn.params) + del params['inline'] + if not any(params['donated_invars']): + del params['donated_invars'] + if all(is_unspecified(s) for s in params['in_shardings']): + del params['in_shardings'] + if all(is_unspecified(s) for s in params['out_shardings']): + del params['out_shardings'] + if all(l is None for l in params['in_layouts']): + del params['in_layouts'] + if all(l is None for l in params['out_layouts']): + del params['out_layouts'] + if not params['keep_unused']: + del params['keep_unused'] + if (params['resource_env'] is None or + params['resource_env'].physical_mesh.empty): + del params['resource_env'] + + # Move name= to the front to make the resulting equation easier to scan. + del params["name"] + return core._pp_eqn(eqn, context, settings, params=["name"] + sorted(params)) + +core.pp_eqn_rules[pjit_p] = _pjit_pp_rule + + +def _pjit_state_discharge_rule( + in_avals, out_avals, *args, jaxpr, in_shardings, out_shardings, + in_layouts, out_layouts, **params): + if not (all(map(is_unspecified, in_shardings)) and + all(map(is_unspecified, out_shardings))): + raise NotImplementedError + + if not (all(l is None for l in in_layouts) and + all(l is None for l in out_layouts)): + raise NotImplementedError + + jaxpr, consts = jaxpr.jaxpr, jaxpr.consts + num_outs = len(jaxpr.outvars) + discharged_jaxpr, discharged_consts = state_discharge.discharge_state(jaxpr, consts) + discharged_closed_jaxpr = core.ClosedJaxpr(discharged_jaxpr, discharged_consts) + new_in_shardings = (UnspecifiedValue(),) * len(discharged_jaxpr.invars) + new_out_shardings = (UnspecifiedValue(),) * len(discharged_jaxpr.outvars) + new_in_layouts = (None,) * len(discharged_jaxpr.invars) + new_out_layouts = (None,) * len(discharged_jaxpr.outvars) + out_and_ref_vals = pjit_p.bind( + *args, jaxpr=discharged_closed_jaxpr, in_shardings=new_in_shardings, + out_shardings=new_out_shardings, in_layouts=new_in_layouts, + out_layouts=new_out_layouts, **params) + out_vals, ref_vals = split_list(out_and_ref_vals, [num_outs]) + ref_vals_iter = iter(ref_vals) + new_invals = tuple(next(ref_vals_iter) if isinstance(aval, AbstractRef) + else None for aval in in_avals) + sentinel = object() + assert next(ref_vals_iter, sentinel) is sentinel + return new_invals, out_vals +state_discharge.register_discharge_rule(pjit_p)(_pjit_state_discharge_rule) + + +# -------------------- with_sharding_constraint -------------------- + +def with_sharding_constraint(x, shardings): + """Mechanism to constrain the sharding of an Array inside a jitted computation + + This is a strict constraint for the GSPMD partitioner and not a hint. For examples + of how to use this function, see `Distributed arrays and automatic parallelization`_. + + Args: + x: PyTree of jax.Arrays which will have their shardings constrained + shardings: PyTree of sharding specifications. Valid values are the same as for + the ``in_shardings`` argument of :func:`jax.experimental.pjit`. + Returns: + x_with_shardings: PyTree of jax.Arrays with specified sharding constraints. + + .. _Distributed arrays and automatic parallelization: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html + """ + x_flat, tree = tree_flatten(x) + user_shardings = prepare_axis_resources( + shardings, "shardings", allow_unconstrained_dims=True) + del shardings + + user_shardings_flat = tuple( + flatten_axes("with_sharding_constraint shardings", tree, user_shardings)) + del user_shardings + + resource_env = mesh_lib.thread_resources.env + mesh = resource_env.physical_mesh + + shardings_flat = [_create_sharding_for_array(mesh, a, 'shardings', + 'with_sharding_constraint') + for a in user_shardings_flat] + unconstrained_dims = [get_unconstrained_dims(s) + if isinstance(s, NamedSharding) else {} + for s in shardings_flat] + del user_shardings_flat + + pjit_check_aval_sharding( + shardings_flat, x_flat, None, "with_sharding_constraint arguments", + allow_uneven_sharding=True) + + outs = [sharding_constraint_p.bind(xf, sharding=s, + resource_env=resource_env, + unconstrained_dims=ud) + for xf, s, ud in zip(x_flat, shardings_flat, unconstrained_dims)] + return tree_unflatten(tree, outs) + +def _identity_fn(x): return x + +def _sharding_constraint_impl(x, sharding, resource_env, unconstrained_dims): + if hasattr(x, 'sharding') and x.sharding.is_equivalent_to(sharding, x.ndim): + return x + # Run a jit here to raise good errors when device assignment don't match. + return api.jit(_identity_fn, out_shardings=sharding)(x) + + +sharding_constraint_p = core.Primitive("sharding_constraint") +sharding_constraint_p.def_impl(_sharding_constraint_impl) +sharding_constraint_p.def_abstract_eval(lambda x, **_: x) +ad.deflinear2(sharding_constraint_p, + lambda ct, _, **params: (sharding_constraint_p.bind(ct, **params),)) + +def _sharding_constraint_hlo_lowering(ctx, x_node, *, sharding, + resource_env, unconstrained_dims): + aval, = ctx.avals_in + out_aval, = ctx.avals_out + axis_ctx = ctx.module_context.axis_context + # axis_ctx and manual_axes is *only used with xmap* and xmap only works with + # NamedSharding. So update the NamedSharding to have the manual axes. + if isinstance(axis_ctx, sharding_impls.SPMDAxisContext): + mesh = resource_env.physical_mesh + if mesh.empty and isinstance(sharding, NamedSharding): + mesh = sharding.mesh + parsed_pspec = parse_flatten_op_sharding( + sharding._to_xla_hlo_sharding(aval.ndim), mesh)[0] + sharding = NamedSharding._from_parsed_pspec( + mesh, parsed_pspec, _manual_axes=axis_ctx.manual_axes) + return [ + mlir.wrap_with_sharding_op(ctx, + x_node, out_aval, + sharding._to_xla_hlo_sharding(aval.ndim).to_proto(), + unspecified_dims=unconstrained_dims) + ] +mlir.register_lowering(sharding_constraint_p, + _sharding_constraint_hlo_lowering) + + +def _sharding_constraint_batcher(insert_axis, spmd_axis_name, axis_size, + axis_name, main_type, vals_in, dims_in, + sharding, resource_env, unconstrained_dims): + if spmd_axis_name is not None and isinstance(sharding, NamedSharding): + used = {n for ns in sharding.spec + for n in (ns if isinstance(ns, tuple) else (ns,))} + if set(spmd_axis_name) & used: + raise ValueError(f"vmap spmd_axis_name {spmd_axis_name} cannot appear in " + "with_sharding_constraint spec, but got spec " + f"{sharding.spec}") + x, = vals_in + d, = dims_in + # None means unconstrained in ParsedPartitionSpec + new_parts = (axis_name,) if insert_axis else ( + None if spmd_axis_name is None else spmd_axis_name) + unconstrained_dims = {ud + (d <= ud) for ud in unconstrained_dims} + + if new_parts is None: + unconstrained_dims.add(d) + + vmapped_sharding = _pjit_batcher_for_sharding( + sharding, d, new_parts, resource_env.physical_mesh, x.ndim) + if unconstrained_dims and isinstance(vmapped_sharding, NamedSharding): + new_spec = list(vmapped_sharding.spec) + [None] * (x.ndim - len(vmapped_sharding.spec)) + for u in unconstrained_dims: + new_spec[u] = PartitionSpec.UNCONSTRAINED + vmapped_sharding = NamedSharding( + vmapped_sharding.mesh, PartitionSpec(*new_spec)) + + y = sharding_constraint_p.bind( + x, + sharding=vmapped_sharding, + resource_env=resource_env, + unconstrained_dims=unconstrained_dims) + return y, d +batching.spmd_axis_primitive_batchers[sharding_constraint_p] = partial( + _sharding_constraint_batcher, False) +batching.axis_primitive_batchers[sharding_constraint_p] = partial( + _sharding_constraint_batcher, False, None) +pxla.spmd_primitive_batchers[sharding_constraint_p] = partial( + _sharding_constraint_batcher, True, None) + + +def _resource_typing_sharding_constraint(avals, params, source_info, + resource_env, named_axis_resources): + aval, = avals + parsed_pspec = None + if isinstance(params['sharding'], NamedSharding): + parsed_pspec = params['sharding']._parsed_pspec + else: + if not resource_env.physical_mesh.empty: + parsed_pspec = parse_flatten_op_sharding( + params['sharding']._to_xla_hlo_sharding(aval.ndim), + resource_env.physical_mesh)[0] + if parsed_pspec is not None: + _check_resources_against_named_axes( + "with_sharding_constraint input", aval, parsed_pspec, named_axis_resources) + +pxla.custom_resource_typing_rules[sharding_constraint_p] = \ + _resource_typing_sharding_constraint + +# -------------------- helpers -------------------- + +def get_unconstrained_dims(sharding: NamedSharding): + assert sharding._parsed_pspec is not None + return {i for i, axes in enumerate(sharding._parsed_pspec) + if axes is None} + + +def _get_partition_spec( + ppspec: Sequence[ParsedPartitionSpec]) -> Sequence[PartitionSpec]: + return [get_single_pspec(p) for p in ppspec] + + +def get_op_sharding_from_executable( + executable) -> tuple[Sequence[xc.OpSharding], Sequence[xc.OpSharding]]: + in_op_shardings: list[xc.OpSharding] = [] + parameter_shardings_from_xla = executable.get_parameter_shardings() + if parameter_shardings_from_xla is not None: + in_op_shardings = parameter_shardings_from_xla + + out_op_shardings: list[xc.OpSharding] = [] + output_shardings_from_xla = executable.get_output_shardings() + if output_shardings_from_xla is not None: + out_op_shardings = output_shardings_from_xla + + return in_op_shardings, out_op_shardings + + +def _get_ppspec_from_executable( + executable, mesh + ) -> tuple[Sequence[ParsedPartitionSpec], Sequence[ParsedPartitionSpec]]: + input_op_shardings, output_op_sharding = get_op_sharding_from_executable( + executable + ) + in_ppspec: list[ParsedPartitionSpec] = [] + for s in input_op_shardings: + in_ppspec.extend(parse_flatten_op_sharding(s, mesh)) + + out_ppspec: list[ParsedPartitionSpec] = [] + for s in output_op_sharding: + out_ppspec.extend(parse_flatten_op_sharding(s, mesh)) + return in_ppspec, out_ppspec + + +def get_pspec_from_executable( + executable, mesh: pxla.Mesh +) -> tuple[tuple[PartitionSpec, ...], tuple[PartitionSpec, ...]]: + in_ppspec, out_ppspec = _get_ppspec_from_executable(executable, mesh) + out_partition_spec = _get_partition_spec(out_ppspec) + in_partition_spec = _get_partition_spec(in_ppspec) + return tuple(in_partition_spec), tuple(out_partition_spec) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pretty_printer.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pretty_printer.py new file mode 100644 index 0000000000000000000000000000000000000000..0614bb8a8d9b33036d581eb996e0e6e87885ba2d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/pretty_printer.py @@ -0,0 +1,487 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Wadler-Lindig pretty printer. +# +# References: +# Wadler, P., 1998. A prettier printer. Journal of Functional Programming, +# pp.223-244. +# +# Lindig, C. 2000. Strictly Pretty. +# https://lindig.github.io/papers/strictly-pretty-2000.pdf +# +# Hafiz, A. 2021. Strictly Annotated: A Pretty-Printer With Support for +# Annotations. https://ayazhafiz.com/articles/21/strictly-annotated +# + +from __future__ import annotations + +from collections.abc import Sequence +import enum +from functools import partial +import sys +from typing import Any, NamedTuple + +from jax._src import config +from jax._src import util + +try: + import colorama # pytype: disable=import-error +except ImportError: + colorama = None + + +_PPRINT_USE_COLOR = config.DEFINE_bool( + 'jax_pprint_use_color', + config.bool_env('JAX_PPRINT_USE_COLOR', True), + help='Enable jaxpr pretty-printing with colorful syntax highlighting.' +) + +def _can_use_color() -> bool: + try: + # Check if we're in IPython or Colab + ipython = get_ipython() # type: ignore[name-defined] + shell = ipython.__class__.__name__ + if shell == "ZMQInteractiveShell": + # Jupyter Notebook + return True + elif "colab" in str(ipython.__class__): + # Google Colab (external or internal) + return True + except NameError: + pass + # Otherwise check if we're in a terminal + return hasattr(sys.stdout, 'isatty') and sys.stdout.isatty() + +CAN_USE_COLOR = _can_use_color() + +class Doc(util.StrictABC): + __slots__ = () + + def format( + self, width: int = 80, *, use_color: bool | None = None, + annotation_prefix: str = " # ", + source_map: list[list[tuple[int, int, Any]]] | None = None + ) -> str: + """ + Formats a pretty-printer document as a string. + + Args: + source_map: for each line in the output, contains a list of + (start column, end column, source) tuples. Each tuple associates a + region of output text with a source. + """ + if use_color is None: + use_color = CAN_USE_COLOR and _PPRINT_USE_COLOR.value + return _format(self, width, use_color=use_color, + annotation_prefix=annotation_prefix, source_map=source_map) + + def __str__(self): + return self.format() + + def __add__(self, other: Doc) -> Doc: + return concat([self, other]) + +class _NilDoc(Doc): + def __repr__(self): return "nil" + +_nil = _NilDoc() + +class _TextDoc(Doc): + __slots__ = ("text", "annotation") + text: str + annotation: str | None + + def __init__(self, text: str, annotation: str | None = None): + assert isinstance(text, str), text + assert annotation is None or isinstance(annotation, str), annotation + self.text = text + self.annotation = annotation + + def __repr__(self): + if self.annotation is not None: + return f"text(\"{self.text}\", annotation=\"{self.annotation}\")" + else: + return f"text(\"{self.text}\")" + +class _ConcatDoc(Doc): + __slots__ = ("children",) + children: list[Doc] + + def __init__(self, children: Sequence[Doc]): + self.children = list(children) + assert all(isinstance(doc, Doc) for doc in self.children), self.children + + def __repr__(self): return f"concat({self.children})" + +class _BreakDoc(Doc): + __slots__ = ("text",) + text: str + + def __init__(self, text: str): + assert isinstance(text, str), text + self.text = text + + def __repr__(self): return f"break({self.text})" + +class _GroupDoc(Doc): + __slots__ = ("child",) + child: Doc + + def __init__(self, child: Doc): + assert isinstance(child, Doc), child + self.child = child + + def __repr__(self): return f"group({self.child})" + +class _NestDoc(Doc): + __slots__ = ("n", "child",) + n: int + child: Doc + + def __init__(self, n: int, child: Doc): + assert isinstance(child, Doc), child + self.n = n + self.child = child + + def __repr__(self): return f"nest({self.n, self.child})" + + +_NO_SOURCE = object() + +class _SourceMapDoc(Doc): + __slots__ = ("child", "source") + child: Doc + source: Any + + def __init__(self, child: Doc, source: Any): + assert isinstance(child, Doc), child + self.child = child + self.source = source + + def __repr__(self): return f"source({self.child}, {self.source})" + + +Color = enum.Enum("_Color", ["BLACK", "RED", "GREEN", "YELLOW", "BLUE", + "MAGENTA", "CYAN", "WHITE", "RESET"]) +Intensity = enum.Enum("_Intensity", ["DIM", "NORMAL", "BRIGHT"]) + +class _ColorDoc(Doc): + __slots__ = ("foreground", "background", "intensity", "child") + foreground: Color | None + background: Color | None + intensity: Intensity | None + child: Doc + + def __init__(self, child: Doc, *, foreground: Color | None = None, + background: Color | None = None, + intensity: Intensity | None = None): + assert isinstance(child, Doc), child + self.child = child + self.foreground = foreground + self.background = background + self.intensity = intensity + + +_BreakMode = enum.Enum("_BreakMode", ["FLAT", "BREAK"]) + + +# In Lindig's paper fits() and format() are defined recursively. This is a +# non-recursive formulation using an explicit stack, necessary because Python +# doesn't have a tail recursion optimization. + +def _fits(doc: Doc, width: int, agenda: list[tuple[int, _BreakMode, Doc]] + ) -> bool: + while width >= 0 and len(agenda) > 0: + i, m, doc = agenda.pop() + if isinstance(doc, _NilDoc): + pass + elif isinstance(doc, _TextDoc): + width -= len(doc.text) + elif isinstance(doc, _ConcatDoc): + agenda.extend((i, m, d) for d in reversed(doc.children)) + elif isinstance(doc, _BreakDoc): + if m == _BreakMode.BREAK: + return True + width -= len(doc.text) + elif isinstance(doc, _NestDoc): + agenda.append((i + doc.n, m, doc.child)) + elif isinstance(doc, _GroupDoc): + agenda.append((i, _BreakMode.FLAT, doc.child)) + elif isinstance(doc, _ColorDoc) or isinstance(doc, _SourceMapDoc): + agenda.append((i, m, doc.child)) + else: + raise ValueError("Invalid document ", doc) + + return width >= 0 + + +# Annotation layout: A flat group is sparse if there are no breaks between +# annotations. +def _sparse(doc: Doc) -> bool: + agenda = [doc] + num_annotations = 0 + seen_break = False + while len(agenda) > 0: + doc = agenda.pop() + if isinstance(doc, _NilDoc): + pass + elif isinstance(doc, _TextDoc): + if doc.annotation is not None: + if num_annotations >= 1 and seen_break: + return False + num_annotations += 1 + elif isinstance(doc, _ConcatDoc): + agenda.extend(reversed(doc.children)) + elif isinstance(doc, _BreakDoc): + seen_break = True + elif isinstance(doc, _NestDoc): + agenda.append(doc.child) + elif isinstance(doc, _GroupDoc): + agenda.append(doc.child) + elif isinstance(doc, _ColorDoc) or isinstance(doc, _SourceMapDoc): + agenda.append(doc.child) + else: + raise ValueError("Invalid document ", doc) + + return True + +class _ColorState(NamedTuple): + foreground: Color + background: Color + intensity: Intensity + +class _State(NamedTuple): + indent: int + mode: _BreakMode + doc: Doc + color: _ColorState + source_map: Any + +class _Line(NamedTuple): + text: str + width: int + annotations: str | None | list[str] + + +def _update_color(use_color: bool, state: _ColorState, update: _ColorState + ) -> tuple[_ColorState, str]: + if not use_color or colorama is None: + return update, "" + color_str = "" + if state.foreground != update.foreground: + color_str += getattr(colorama.Fore, str(update.foreground.name)) + if state.background != update.background: + color_str += getattr(colorama.Back, str(update.background.name)) + if state.intensity != update.intensity: + color_str += colorama.Style.NORMAL # pytype: disable=unsupported-operands + color_str += getattr(colorama.Style, str(update.intensity.name)) + return update, color_str + + +def _align_annotations(lines): + # TODO: Hafiz also implements a local alignment mode, where groups of lines + # with annotations are aligned together. + maxlen = max(l.width for l in lines) + out = [] + for l in lines: + if len(l.annotations) == 0: + out.append(l._replace(annotations=None)) + elif len(l.annotations) == 1: + out.append(l._replace(text=l.text + " " * (maxlen - l.width), + annotations=l.annotations[0])) + else: + out.append(l._replace(text=l.text + " " * (maxlen - l.width), + annotations=l.annotations[0])) + for a in l.annotations[1:]: + out.append(_Line(text=" " * maxlen, width=l.width, annotations=a)) + return out + + + +def _format( + doc: Doc, width: int, *, use_color: bool, annotation_prefix: str, + source_map: list[list[tuple[int, int, Any]]] | None +) -> str: + lines = [] + default_colors = _ColorState(Color.RESET, Color.RESET, Intensity.NORMAL) + annotation_colors = _ColorState(Color.RESET, Color.RESET, Intensity.DIM) + color_state = default_colors + source_start = 0 # The column at which the current source region starts. + source = _NO_SOURCE # The currently active source region. + line_source_map = [] # Source maps for the current line of text. + agenda = [_State(0, _BreakMode.BREAK, doc, default_colors, source)] + k = 0 + line_text = "" + line_annotations = [] + while len(agenda) > 0: + i, m, doc, color, agenda_source = agenda.pop() + if source_map is not None and agenda_source != source: + pos = len(line_text) + if source_start != pos and source is not _NO_SOURCE: + line_source_map.append((source_start, pos, source)) + source = agenda_source + source_start = pos + if isinstance(doc, _NilDoc): + pass + elif isinstance(doc, _TextDoc): + color_state, color_str = _update_color(use_color, color_state, color) + line_text += color_str + line_text += doc.text + if doc.annotation is not None: + line_annotations.append(doc.annotation) + k += len(doc.text) + elif isinstance(doc, _ConcatDoc): + agenda.extend(_State(i, m, d, color, source) + for d in reversed(doc.children)) + elif isinstance(doc, _BreakDoc): + if m == _BreakMode.BREAK: + if len(line_annotations) > 0: + color_state, color_str = _update_color(use_color, color_state, + annotation_colors) + line_text += color_str + lines.append(_Line(line_text, k, line_annotations)) + if source_map is not None: + pos = len(line_text) + if source_start != pos and source is not _NO_SOURCE: + line_source_map.append((source_start, pos, source)) + source_map.append(line_source_map) + line_source_map = [] + source_start = i + line_text = " " * i + line_annotations = [] + k = i + else: + color_state, color_str = _update_color(use_color, color_state, color) + line_text += color_str + line_text += doc.text + k += len(doc.text) + elif isinstance(doc, _NestDoc): + agenda.append(_State(i + doc.n, m, doc.child, color, source)) + elif isinstance(doc, _GroupDoc): + # In Lindig's paper, _fits is passed the remainder of the document. + # I'm pretty sure that's a bug and we care only if the current group fits! + if (_sparse(doc) + and _fits(doc, width - k, [(i, _BreakMode.FLAT, doc.child)])): + agenda.append(_State(i, _BreakMode.FLAT, doc.child, color, source)) + else: + agenda.append(_State(i, _BreakMode.BREAK, doc.child, color, source)) + elif isinstance(doc, _ColorDoc): + color = _ColorState(doc.foreground or color.foreground, + doc.background or color.background, + doc.intensity or color.intensity) + agenda.append(_State(i, m, doc.child, color, source)) + elif isinstance(doc, _SourceMapDoc): + agenda.append(_State(i, m, doc.child, color, doc.source)) + else: + raise ValueError("Invalid document ", doc) + + if len(line_annotations) > 0: + color_state, color_str = _update_color(use_color, color_state, + annotation_colors) + line_text += color_str + if source_map is not None: + pos = len(line_text) + if source_start != pos and source is not _NO_SOURCE: + line_source_map.append((source_start, pos, source)) + source_map.append(line_source_map) + lines.append(_Line(line_text, k, line_annotations)) + lines = _align_annotations(lines) + out = "\n".join( + l.text if l.annotations is None + else f"{l.text}{annotation_prefix}{l.annotations}" for l in lines) + color_state, color_str = _update_color(use_color, color_state, + default_colors) + return out + color_str + + + + +# Public API. + +def nil() -> Doc: + """An empty document.""" + return _nil + +def text(s: str, annotation: str | None = None) -> Doc: + """Literal text.""" + return _TextDoc(s, annotation) + +def concat(docs: Sequence[Doc]) -> Doc: + """Concatenation of documents.""" + docs = list(docs) + if len(docs) == 1: + return docs[0] + return _ConcatDoc(docs) + +def brk(text: str = " ") -> Doc: + """A break. + + Prints either as a newline or as `text`, depending on the enclosing group. + """ + return _BreakDoc(text) + +def group(doc: Doc) -> Doc: + """Layout alternative groups. + + Prints the group with its breaks as their text (typically spaces) if the + entire group would fit on the line when printed that way. Otherwise, breaks + inside the group as printed as newlines. + """ + return _GroupDoc(doc) + +def nest(n: int, doc: Doc) -> Doc: + """Increases the indentation level by `n`.""" + return _NestDoc(n, doc) + + +def color(doc: Doc, *, foreground: Color | None = None, + background: Color | None = None, + intensity: Intensity | None = None): + """ANSI colors. + + Overrides the foreground/background/intensity of the text for the child doc. + Requires use_colors=True to be set when printing and the `colorama` package + to be installed; otherwise does nothing. + """ + return _ColorDoc(doc, foreground=foreground, background=background, + intensity=intensity) + + +def source_map(doc: Doc, source: Any): + """Source mapping. + + A source map associates a region of the pretty-printer's text output with a + source location that produced it. For the purposes of the pretty printer a + ``source`` may be any object: we require only that we can compare sources for + equality. A text region to source object mapping can be populated as a side + output of the ``format`` method. + """ + return _SourceMapDoc(doc, source) + +type_annotation = partial(color, intensity=Intensity.NORMAL, + foreground=Color.MAGENTA) +keyword = partial(color, intensity=Intensity.BRIGHT, foreground=Color.BLUE) + + +def join(sep: Doc, docs: Sequence[Doc]) -> Doc: + """Concatenates `docs`, separated by `sep`.""" + docs = list(docs) + if len(docs) == 0: + return nil() + xs = [docs[0]] + for doc in docs[1:]: + xs.append(sep) + xs.append(doc) + return concat(xs) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/prng.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/prng.py new file mode 100644 index 0000000000000000000000000000000000000000..4d4a045a31e21edc16e7664743ce019707a8edc6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/prng.py @@ -0,0 +1,1287 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from functools import partial, reduce +import math +import operator as op +from typing import Any, Callable, NamedTuple + +import numpy as np + +import jax +from jax import lax +from jax import numpy as jnp +from jax import tree_util + +from jax._src import api_util +from jax._src import api +from jax._src import config as config +from jax._src import core +from jax._src import dispatch +from jax._src import dtypes +from jax._src import pretty_printer as pp +from jax._src import source_info_util +from jax._src import tree_util as tree_util_internal +from jax._src import typing +from jax._src.api import jit, vmap +from jax._src.dtypes import float0 +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.interpreters import pxla +from jax._src.interpreters import xla +from jax._src.lax import lax as lax_internal +from jax._src.lax import utils as lax_utils +from jax._src.lib import gpu_prng +from jax._src.lib import xla_client as xc +from jax._src.lib.mlir import ir +from jax._src.lib.mlir.dialects import hlo +from jax._src.numpy.array_methods import ( + _array_operators, _set_array_base_attributes, _IndexUpdateHelper) +from jax._src.sharding_impls import ( + NamedSharding, PmapSharding, physical_sharding, logical_sharding) +from jax._src.typing import Array +from jax._src.util import safe_map, safe_zip + +map, unsafe_map = safe_map, map +zip, unsafe_zip = safe_zip, zip + +Device = xc.Device +Shard = Any # TODO(jakevdp): fix circular imports and import Shard +Shape = tuple[int, ...] + +UINT_DTYPES = { + 8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64} + +# -- PRNG implementation interface + +class PRNGImpl(NamedTuple): + """Specifies PRNG key shape and operations. + + A PRNG implementation is determined by a key type ``K`` and a + collection of functions that operate on such keys. The key type + ``K`` is an array type with element type uint32 and shape specified + by ``key_shape``. The type signature of each operations is:: + + seed :: int[] -> K + fold_in :: K -> int[] -> K + split[shape] :: K -> K[*shape] + random_bits[shape, bit_width] :: K -> uint[*shape] + + A PRNG implementation is adapted to an array-like object of keys + ``K`` by the ``PRNGKeyArray`` class, which should be created via the + ``random_seed`` function. + """ + key_shape: Shape + seed: Callable + split: Callable + random_bits: Callable + fold_in: Callable + name: str = '' + tag: str = '?' + + def __hash__(self) -> int: + return hash(self.tag) + + def __str__(self) -> str: + return self.tag + + def pprint(self): + ty = self.__class__.__name__ + return (pp.text(f"{ty} [{self.tag}] {{{self.name}}}:") + + pp.nest(2, pp.group(pp.brk() + pp.join(pp.brk(), [ + pp.text(f"{k} = {v}") for k, v in self._asdict().items() + ])))) + + +prngs = {} + +def register_prng(impl: PRNGImpl): + if impl.name in prngs: + raise ValueError(f'PRNG with name {impl.name} already registered: {impl}') + prngs[impl.name] = impl + + +# -- PRNG key arrays + +def _check_prng_key_data(impl, key_data: typing.Array): + ndim = len(impl.key_shape) + if not all(hasattr(key_data, attr) for attr in ['ndim', 'shape', 'dtype']): + raise TypeError("JAX encountered invalid PRNG key data: expected key_data " + f"to have ndim, shape, and dtype attributes. Got {key_data}") + if key_data.ndim < 1: + raise TypeError("JAX encountered invalid PRNG key data: expected " + f"key_data.ndim >= 1; got ndim={key_data.ndim}") + if key_data.shape[-ndim:] != impl.key_shape: + raise TypeError("JAX encountered invalid PRNG key data: expected key_data.shape to " + f"end with {impl.key_shape}; got shape={key_data.shape} for {impl=}") + if key_data.dtype not in [np.uint32, float0]: + raise TypeError("JAX encountered invalid PRNG key data: expected key_data.dtype = uint32; " + f"got dtype={key_data.dtype}") + + +class PRNGKeyArray(jax.Array): + """An array of PRNG keys backed by an RNG implementation. + + This class lifts the definition of a PRNG, provided in the form of a + ``PRNGImpl``, into an array-like pytree class. Instances of this + class behave like an array whose base elements are keys, hiding the + fact that keys are typically arrays (of ``uint32`` dtype) themselves. + + PRNGKeyArrays are also restricted relative to JAX arrays in that + they do not expose arithmetic operations. They instead expose + wrapper methods around the PRNG implementation functions (``split``, + ``random_bits``, ``fold_in``). + """ + # TODO(jakevdp): potentially add tolist(), tobytes(), + # device_buffer, device_buffers, __cuda_interface__() + + _impl: PRNGImpl + _base_array: typing.Array + _consumed: bool | np.ndarray # Used in jax.experimental.key_reuse. + _source_info: None | source_info_util.SourceInfo = None + + def __init__(self, impl, key_data: Any): + assert not isinstance(key_data, core.Tracer) + _check_prng_key_data(impl, key_data) + self._impl = impl + self._base_array = key_data + self._consumed = False # TODO(jakevdp): default to True here? + + def block_until_ready(self): + _ = self._base_array.block_until_ready() + return self + + def copy_to_host_async(self): + _ = self._base_array.copy_to_host_async() + + @property + def aval(self): + return keys_shaped_array(self._impl, self.shape) + + @property + def shape(self): + return base_arr_shape_to_keys_shape(self._impl, self._base_array.shape) + + @property + def size(self): + return math.prod(self.shape) + + @property + def ndim(self): + return len(self.shape) + + @property + def dtype(self): + return KeyTy(self._impl) + + @property + def itemsize(self): + return self.dtype.itemsize + + _device = property(op.attrgetter('_base_array._device')) + _committed = property(op.attrgetter('_base_array._committed')) + device = property(op.attrgetter('_base_array.device')) + devices = property(op.attrgetter('_base_array.devices')) # type: ignore[assignment] + is_fully_addressable = property(op.attrgetter('_base_array.is_fully_addressable')) # type: ignore[assignment] + is_fully_replicated = property(op.attrgetter('_base_array.is_fully_replicated')) # type: ignore[assignment] + delete = property(op.attrgetter('_base_array.delete')) # type: ignore[assignment] + is_deleted = property(op.attrgetter('_base_array.is_deleted')) # type: ignore[assignment] + on_device_size_in_bytes = property(op.attrgetter('_base_array.on_device_size_in_bytes')) # type: ignore[assignment] + unsafe_buffer_pointer = property(op.attrgetter('_base_array.unsafe_buffer_pointer')) # type: ignore[assignment] + + def addressable_data(self, index: int) -> PRNGKeyArray: + return PRNGKeyArray(self._impl, self._base_array.addressable_data(index)) + + @property + def addressable_shards(self) -> list[Shard]: + return [ + type(s)( + device=s._device, + sharding=s._sharding, + global_shape=s._global_shape, + data=PRNGKeyArray(self._impl, s._data), + ) + for s in self._base_array.addressable_shards + ] + + @property + def global_shards(self) -> list[Shard]: + return [ + type(s)( + device=s._device, + sharding=s._sharding, + global_shape=s._global_shape, + data=PRNGKeyArray(self._impl, s._data), + ) + for s in self._base_array.global_shards + ] + + @property + def sharding(self): + return logical_sharding(self.aval, self._base_array.sharding) + + def _is_scalar(self): + base_ndim = len(self._impl.key_shape) + return self._base_array.ndim == base_ndim + + def __len__(self): + if self._is_scalar(): + raise TypeError('len() of unsized object') + return len(self._base_array) + + def __iter__(self) -> Iterator[PRNGKeyArray]: + if self._is_scalar(): + raise TypeError('iteration over a 0-d key array') + # TODO(frostig): we may want to avoid iteration by slicing because + # a very common use of iteration is `k1, k2 = split(key)`, and + # slicing/indexing may be trickier to track for linearity checking + # purposes. Maybe we can: + # * introduce an unpack primitive+traceable (also allow direct use) + # * unpack upfront into shape[0] many keyarray slices + # * return iter over these unpacked slices + # Whatever we do, we'll want to do it by overriding + # ShapedArray._iter when the element type is KeyTy... + return (PRNGKeyArray(self._impl, k) for k in iter(self._base_array)) + + def __repr__(self): + return (f'Array({self.shape}, dtype={self.dtype.name}) overlaying:\n' + f'{self._base_array}') + + def pprint(self): + pp_keys = pp.text('shape = ') + pp.text(str(self.shape)) + pp_impl = pp.text('impl = ') + self._impl.pprint() + return str(pp.group( + pp.text('PRNGKeyArray:') + + pp.nest(2, pp.brk() + pp_keys + pp.brk() + pp_impl))) + + def copy(self): + out = self.__class__(self._impl, self._base_array.copy()) + out._consumed = self._consumed # TODO(jakevdp): is this correct? + return out + + __hash__ = None # type: ignore[assignment] + __array_priority__ = 100 + + # Overwritten immediately below + @property + def at(self) -> _IndexUpdateHelper: assert False # type: ignore[override] + @property + def T(self) -> PRNGKeyArray: assert False + def __getitem__(self, _) -> PRNGKeyArray: assert False + def flatten(self, *_, **__) -> PRNGKeyArray: assert False + def ravel(self, *_, **__) -> PRNGKeyArray: assert False + def reshape(self, *_, **__) -> PRNGKeyArray: assert False + def squeeze(self, *_, **__) -> PRNGKeyArray: assert False + def swapaxes(self, *_, **__) -> PRNGKeyArray: assert False + def take(self, *_, **__) -> PRNGKeyArray: assert False + def transpose(self, *_, **__) -> PRNGKeyArray: assert False + +_set_array_base_attributes(PRNGKeyArray, include=[ + *(f"__{op}__" for op in _array_operators), + 'at', 'flatten', 'ravel', 'reshape', + 'squeeze', 'swapaxes', 'take', 'transpose', 'T']) + +api_util._shaped_abstractify_handlers[PRNGKeyArray] = op.attrgetter('aval') + +def prngkeyarray_flatten(x): + return (x._base_array,), x._impl + +def prngkeyarray_unflatten(impl, children): + base_array, = children + return PRNGKeyArray(impl, base_array) + +tree_util_internal.dispatch_registry.register_node( + PRNGKeyArray, prngkeyarray_flatten, prngkeyarray_unflatten) + + +# TODO(frostig): remove, rerouting callers directly to random_seed +def seed_with_impl(impl: PRNGImpl, seed: int | typing.ArrayLike) -> PRNGKeyArray: + return random_seed(seed, impl=impl) + + +def keys_shaped_array(impl, shape): + return core.ShapedArray(shape, KeyTy(impl)) + +def base_arr_shape_to_keys_shape(impl, base_arr_shape): + base_ndim = len(impl.key_shape) + return base_arr_shape[:-base_ndim] + + +class KeyTyRules: + + @staticmethod + def full(shape, fill_value, dtype): + physical_shape = (*shape, *dtype._impl.key_shape) + if hasattr(fill_value, 'dtype') and jnp.issubdtype(fill_value.dtype, dtypes.prng_key): + key_data = jnp.broadcast_to(random_unwrap(fill_value), physical_shape) + else: + key_data = lax.full(physical_shape, fill_value, dtype=np.dtype('uint32')) + # TODO(frostig,mattjj,vanderplas,lenamartens): consider this consumed from + # the outset. + return random_wrap(key_data, impl=dtype._impl) + + @staticmethod + def physical_element_aval(dtype) -> core.ShapedArray: + return core.ShapedArray(dtype._impl.key_shape, jnp.dtype('uint32')) + + @staticmethod + def physical_const(val) -> Array: + return val._base_array + + @staticmethod + def result_handler(sticky_device, aval): + def handler(_, buf): + buf.aval = core.ShapedArray(buf.shape, buf.dtype) + return PRNGKeyArray(aval.dtype._impl, buf) + return handler + + @staticmethod + def local_sharded_result_handler(aval, sharding, indices): + phys_aval = core.physical_aval(aval) + key_shape = aval.dtype._impl.key_shape + phys_handler_maker = pxla.local_result_handlers[core.ShapedArray] + + # set up a grounded sharding (with a grounded sharding spec) + if isinstance(sharding, (PmapSharding, NamedSharding)): + phys_sharding = physical_sharding(aval, sharding) + else: + assert False, f'impossible sharding {sharding} in local sharded result handler' + + # set up grounded indices + trailing_inds = [slice(None)] * len(key_shape) + phys_indices = [(*inds, *trailing_inds) for inds in indices] + + # make a physical handler + phys_handler = phys_handler_maker(phys_aval, phys_sharding, phys_indices) + + # set up a handler that calls the physical one and wraps back up + def handler(bufs): + return PRNGKeyArray(aval.dtype._impl, phys_handler(bufs)) + + return handler + + @staticmethod + def global_sharded_result_handler(aval, out_sharding, committed): + phys_aval = core.physical_aval(aval) + phys_handler_maker = pxla.global_result_handlers[core.ShapedArray] + + phys_sharding = physical_sharding(aval, out_sharding) + phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed) + def handler(bufs): + return PRNGKeyArray(aval.dtype._impl, phys_handler(bufs)) + return handler + + @staticmethod + def make_sharded_array(aval, sharding, arrays, committed): + phys_aval = core.physical_aval(aval) + phys_handler_maker = pxla.global_result_handlers[core.ShapedArray] + phys_arrays = [random_unwrap(arr) for arr in arrays] + + phys_sharding = physical_sharding(aval, sharding) + phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed) + phys_result = phys_handler(phys_arrays) + return PRNGKeyArray(aval.dtype._impl, phys_result) + + @staticmethod + def device_put_sharded(vals, aval, sharding, devices): + physical_aval = core.physical_aval(aval) + physical_buffers = tree_util.tree_map(random_unwrap, vals) + phys_sharding = physical_sharding(aval, sharding) + physical_result = pxla.batched_device_put(physical_aval, phys_sharding, + physical_buffers, list(devices)) + return random_wrap(physical_result, impl=aval.dtype._impl) + + @staticmethod + def device_put_replicated(val, aval, sharding, devices): + physical_aval = core.physical_aval(aval) + assert len(xla.aval_to_xla_shapes(physical_aval)) == 1 + physical_buf = random_unwrap(val) + phys_sharding = physical_sharding(aval, sharding) + physical_result = pxla.batched_device_put( + physical_aval, phys_sharding, [physical_buf] * len(devices), devices) + return random_wrap(physical_result, impl=aval.dtype._impl) + + @staticmethod + def tangent_dtype(_): + return dtypes.float0 + + # TODO(mattjj,frostig): even though the key dtype shouldn't appear in + # tangents, our ad.replace_float0s in custom_jvp/vjp means passing in zeros + # like the primal to user rules + @staticmethod + def zero(_): + return np.zeros((), dtypes.float0) + + @staticmethod + def convert_from(key_dtype, other_dtype) -> bool: + return False + + @staticmethod + def convert_to(other_dtype, key_dtype) -> bool: + return False + + +class KeyTy(dtypes.ExtendedDType): + _impl: PRNGImpl # TODO(mattjj,frostig): protocol really + _rules = KeyTyRules + type = dtypes.prng_key + + def __init__(self, impl): + self._impl = impl + + @property + def name(self) -> str: + return f'key<{self._impl.tag}>' + + @property + def itemsize(self) -> int: + return math.prod(self._impl.key_shape) * np.dtype('uint32').itemsize + + def __repr__(self) -> str: + return self.name + + def __eq__(self, other): + return type(other) is KeyTy and self._impl == other._impl + + def __hash__(self) -> int: + return hash((self.__class__, self._impl)) + + +core.pytype_aval_mappings[PRNGKeyArray] = lambda x: x.aval +xla.pytype_aval_mappings[PRNGKeyArray] = lambda x: x.aval + +xla.canonicalize_dtype_handlers[PRNGKeyArray] = lambda x: x + + +def key_array_shard_arg_handler(x: PRNGKeyArray, sharding): + arr = x._base_array + phys_sharding = physical_sharding(x.aval, sharding) + return pxla.shard_arg_handlers[type(arr)](arr, phys_sharding) + + +pxla.shard_arg_handlers[PRNGKeyArray] = key_array_shard_arg_handler + + +def key_array_constant_handler(x): + arr = x._base_array + return mlir.get_constant_handler(type(arr))(arr) +mlir.register_constant_handler(PRNGKeyArray, key_array_constant_handler) + + +# -- primitives + +def iterated_vmap_unary(n, f): + for _ in range(n): + f = api.vmap(f) + return f + +# TODO(frostig): Revise the following two functions? These basically +# undo the singleton dimensions added by `batching.defbroadcasting`. +# It works, but introduces some possibly-redundant squeezes. Can we +# borrow from other broadcasting primitives instead? + +def squeeze_vmap(f, left): + def squeeze_vmap_f(x, y): + if left: + x = jnp.squeeze(x, axis=0) + axes = (None, 0) + else: + y = jnp.squeeze(y, axis=0) + axes = (0, None) + return api.vmap(f, in_axes=axes, out_axes=0)(x, y) + return squeeze_vmap_f + +def iterated_vmap_binary_bcast(shape1, shape2, f): + ndim1, ndim2 = len(shape1), len(shape2) + if ndim1 == ndim2 == 0: + return f + if 0 in [ndim1, ndim2]: + if ndim1 == 0: + return lambda x, y: iterated_vmap_unary(ndim2, lambda y: f(x, y))(y) + else: + return lambda x, y: iterated_vmap_unary(ndim1, lambda x: f(x, y))(x) + assert len(shape1) == len(shape2) + for sz1, sz2 in reversed(zip(shape1, shape2)): + if sz1 == sz2: + f = api.vmap(f, out_axes=0) + else: + assert sz1 == 1 or sz2 == 1, (sz1, sz2) + f = squeeze_vmap(f, sz1 == 1) + return f + + +def random_seed(seeds: int | typing.ArrayLike, impl: PRNGImpl) -> PRNGKeyArray: + # Avoid overflow error in X32 mode by first converting ints to int64. + # This breaks JIT invariance for large ints, but supports the common + # use-case of instantiating with Python hashes in X32 mode. + if isinstance(seeds, int): + seeds_arr = jnp.asarray(np.int64(seeds)) + else: + seeds_arr = jnp.asarray(seeds) + if config.random_seed_offset.value: + seeds_arr += config.random_seed_offset.value + return random_seed_p.bind(seeds_arr, impl=impl) + +random_seed_p = core.Primitive('random_seed') +ad.defjvp_zero(random_seed_p) +batching.defvectorized(random_seed_p) + +@random_seed_p.def_abstract_eval +def random_seed_abstract_eval(seeds_aval, *, impl): + return keys_shaped_array(impl, seeds_aval.shape) + +@random_seed_p.def_impl +def random_seed_impl(seeds, *, impl): + base_arr = random_seed_impl_base(seeds, impl=impl) + return PRNGKeyArray(impl, base_arr) + +def random_seed_impl_base(seeds, *, impl): + seed = iterated_vmap_unary(np.ndim(seeds), impl.seed) + return seed(seeds) + +def random_seed_lowering(ctx, seeds, *, impl): + aval, = ctx.avals_in + seed = iterated_vmap_unary(aval.ndim, impl.seed) + seed_lowering = mlir.lower_fun(seed, multiple_results=False) + return mlir.delegate_lowering( + ctx, seed_lowering, seeds, + avals_out=map(core.physical_aval, ctx.avals_out)) + +mlir.register_lowering(random_seed_p, random_seed_lowering) + + +def random_split(keys, shape: Shape): + return random_split_p.bind(keys, shape=shape) + +random_split_p = core.Primitive('random_split') +ad.defjvp_zero(random_split_p) +batching.defvectorized(random_split_p) + +@random_split_p.def_abstract_eval +def random_split_abstract_eval(keys_aval, *, shape): + return keys_shaped_array(keys_aval.dtype._impl, (*keys_aval.shape, *shape)) + +@random_split_p.def_impl +def random_split_impl(keys, *, shape): + base_arr = random_split_impl_base( + keys._impl, keys._base_array, keys.ndim, shape=shape) + return PRNGKeyArray(keys._impl, base_arr) + +def random_split_impl_base(impl, base_arr, keys_ndim, *, shape): + split = iterated_vmap_unary(keys_ndim, lambda k: impl.split(k, shape)) + return split(base_arr) + +def random_split_lowering(ctx, keys, *, shape): + aval, = ctx.avals_in + impl = aval.dtype._impl + split = iterated_vmap_unary(aval.ndim, lambda k: impl.split(k, shape)) + split_lowering = mlir.lower_fun(split, multiple_results=False) + return mlir.delegate_lowering( + ctx, split_lowering, keys, + avals_in=[core.physical_aval(aval)], + avals_out=map(core.physical_aval, ctx.avals_out)) + +mlir.register_lowering(random_split_p, random_split_lowering) + + +def random_fold_in(keys, msgs): + return random_fold_in_p.bind(keys, jnp.asarray(msgs)) + +random_fold_in_p = core.Primitive('random_fold_in') +ad.defjvp_zero(random_fold_in_p) +batching.defbroadcasting(random_fold_in_p) + +@random_fold_in_p.def_abstract_eval +def random_fold_in_abstract_eval(keys_aval, msgs_aval): + shape = lax_internal.broadcasting_shape_rule( + 'random_fold_in', keys_aval, msgs_aval) + named_shape = lax_utils.standard_named_shape_rule(keys_aval, msgs_aval) + return core.ShapedArray(shape, keys_aval.dtype, named_shape=named_shape) + +@random_fold_in_p.def_impl +def random_fold_in_impl(keys, msgs): + base_arr = random_fold_in_impl_base( + keys._impl, keys._base_array, msgs, keys.shape) + return PRNGKeyArray(keys._impl, base_arr) + +def random_fold_in_impl_base(impl, base_arr, msgs, keys_shape): + fold_in = iterated_vmap_binary_bcast( + keys_shape, np.shape(msgs), impl.fold_in) + return fold_in(base_arr, msgs) + +def random_fold_in_lowering(ctx, keys, msgs): + keys_aval, msgs_aval = ctx.avals_in + impl = keys_aval.dtype._impl + fold_in = iterated_vmap_binary_bcast( + keys_aval.shape, msgs_aval.shape, impl.fold_in) + fold_in_lowering = mlir.lower_fun(fold_in, multiple_results=False) + return mlir.delegate_lowering( + ctx, fold_in_lowering, keys, msgs, + avals_in=[core.physical_aval(keys_aval), msgs_aval], + avals_out=map(core.physical_aval, ctx.avals_out)) + +mlir.register_lowering(random_fold_in_p, random_fold_in_lowering) + + +def random_bits(keys, bit_width, shape): + shape = core.as_named_shape(shape) + for name, size in shape.named_items: + # TODO(frostig,mattjj,apaszke): Is this real_size check necessary, + # and is it meant to raise a user-facing ValueError? Should it be + # an `assert` (or RuntimeError) instead? Why do we check it in + # calls to `random_bits` instead of a more common paralleism path? + real_size = lax.psum(1, name) + if real_size != size: + raise ValueError(f"The shape of axis {name} was specified as {size}, " + f"but it really is {real_size}") + axis_index = lax.axis_index(name) + keys = random_fold_in(keys, axis_index) + return random_bits_p.bind(keys, bit_width=bit_width, shape=shape.positional) + +random_bits_p = core.Primitive('random_bits') +ad.defjvp_zero(random_bits_p) +batching.defvectorized(random_bits_p) + +@random_bits_p.def_abstract_eval +def random_bits_abstract_eval(keys_aval, *, bit_width, shape): + out_shape = (*keys_aval.shape, *shape) + out_dtype = dtypes.dtype(f'uint{bit_width}') + return core.ShapedArray(out_shape, out_dtype) + +@random_bits_p.def_impl +def random_bits_impl(keys, *, bit_width, shape): + return random_bits_impl_base(keys._impl, keys._base_array, keys.ndim, + bit_width=bit_width, shape=shape) + +def random_bits_impl_base(impl, base_arr, keys_ndim, *, bit_width, shape): + bits = iterated_vmap_unary( + keys_ndim, lambda k: impl.random_bits(k, bit_width, shape)) + return bits(base_arr) + +def random_bits_lowering(ctx, keys, *, bit_width, shape): + aval, = ctx.avals_in + impl = aval.dtype._impl + bits = iterated_vmap_unary( + aval.ndim, lambda k: impl.random_bits(k, bit_width, shape)) + bits_lowering = mlir.lower_fun(bits, multiple_results=False) + ctx_new = ctx.replace(avals_in=[core.physical_aval(aval)]) + out = bits_lowering(ctx_new, keys) + ctx.set_tokens_out(ctx_new.tokens_out) + return out + +mlir.register_lowering(random_bits_p, random_bits_lowering) + + +# The following wrap/unwrap primitives are at least a stopgap for +# backwards compatibility, namely when `config.jax_enable_custom_prng` +# is False. We need to convert key arrays to and from underlying +# uint32 base array, and we may need to do so under a jit. For +# example, we want to support: +# +# keys = jax.jit(random.split)(key) +# +# where `key` and `keys` are both acceptably old-style uint32 arrays +# so long as enable_custom_prng is False. The way we handle this is +# that `random.split` adapts the input/output by converting to/from +# key arrays across its call to `random_split`. So we rely on these +# wrap/unwrap casting primitives to allow that conversion under jit. +# +# We may want to keep both around for testing and debugging escape +# hatches. We can rename them `unsafe` for emphasis, and/or issue a +# warning on entry to the traceable. +# +# TODO(frostig): Consider removal once we always enable_custom_prng. + +def random_wrap(base_arr, *, impl): + _check_prng_key_data(impl, base_arr) + return random_wrap_p.bind(base_arr, impl=impl) + +random_wrap_p = core.Primitive('random_wrap') +ad.defjvp_zero(random_wrap_p) + +@random_wrap_p.def_abstract_eval +def random_wrap_abstract_eval(base_arr_aval, *, impl): + shape = base_arr_shape_to_keys_shape(impl, base_arr_aval.shape) + return keys_shaped_array(impl, shape) + +@random_wrap_p.def_impl +def random_wrap_impl(base_arr, *, impl): + return PRNGKeyArray(impl, base_arr) + +def random_wrap_lowering(ctx, base_arr, *, impl): + return [base_arr] + +def random_wrap_batch_rule(batched_args, batch_dims, *, impl): + x, = batched_args + d, = batch_dims + x = batching.bdim_at_front(x, d, 1) + return random_wrap(x, impl=impl), 0 + +mlir.register_lowering(random_wrap_p, random_wrap_lowering) +batching.primitive_batchers[random_wrap_p] = random_wrap_batch_rule + + +def random_unwrap(keys): + if not jnp.issubdtype(keys.dtype, dtypes.prng_key): + raise TypeError(f'random_unwrap takes key array operand, got {keys.dtype=}') + return random_unwrap_p.bind(keys) + +random_unwrap_p = core.Primitive('random_unwrap') +ad.defjvp_zero(random_unwrap_p) +batching.defvectorized(random_unwrap_p) + +@random_unwrap_p.def_abstract_eval +def random_unwrap_abstract_eval(keys_aval): + return core.physical_aval(keys_aval) + +@random_unwrap_p.def_impl +def random_unwrap_impl(keys): + return keys._base_array + +def random_unwrap_lowering(ctx, keys): + return [keys] + +mlir.register_lowering(random_unwrap_p, random_unwrap_lowering) + + +# -- threefry2x32 PRNG implementation + + +def _is_threefry_prng_key(key: typing.Array) -> bool: + try: + return key.shape == (2,) and key.dtype == np.uint32 + except AttributeError: + return False + + +def threefry_seed(seed: typing.Array) -> typing.Array: + """Create a single raw threefry PRNG key from an integer seed. + + Args: + seed: a 64- or 32-bit integer used as the value of the key. + + Returns: + The PRNG key contents, modeled as an array of shape (2,) and dtype + uint32. The key is constructed from a 64-bit seed by effectively + bit-casting to a pair of uint32 values (or from a 32-bit seed by + first padding out with zeros). + """ + return _threefry_seed(seed) + +@partial(jit, inline=True) +def _threefry_seed(seed: typing.Array) -> typing.Array: + if seed.shape: + raise TypeError(f"PRNG key seed must be a scalar; got {seed!r}.") + if not np.issubdtype(seed.dtype, np.integer): + raise TypeError(f"PRNG key seed must be an integer; got {seed!r}") + convert = lambda k: lax.expand_dims(lax.convert_element_type(k, np.uint32), [0]) + k1 = convert( + lax.shift_right_logical(seed, lax_internal._const(seed, 32))) + with config.numpy_dtype_promotion('standard'): + # TODO(jakevdp): in X64 mode, this can generate 64-bit computations for 32-bit + # inputs. We should avoid this. + k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF))) + return lax.concatenate([k1, k2], 0) + + +def _make_rotate_left(dtype): + if not jnp.issubdtype(dtype, np.integer): + raise TypeError("_rotate_left only accepts integer dtypes.") + nbits = np.array(jnp.iinfo(dtype).bits, dtype) + + def _rotate_left(x, d): + if lax.dtype(d) != dtype: + d = lax.convert_element_type(d, dtype) + if lax.dtype(x) != dtype: + x = lax.convert_element_type(x, dtype) + return lax.shift_left(x, d) | lax.shift_right_logical(x, nbits - d) + return _rotate_left + + +### hash function and split + +def _threefry2x32_abstract_eval(*args): + if any(a.dtype != jnp.uint32 for a in args): + raise TypeError("Arguments to threefry2x32 must have uint32 type, got {}" + .format(args)) + if all(isinstance(arg, core.ShapedArray) for arg in args): + shape = lax_internal.broadcasting_shape_rule(*args) + named_shape = core.join_named_shapes(*(a.named_shape for a in args)) + aval = core.ShapedArray(shape, jnp.dtype(jnp.uint32), named_shape=named_shape) + else: + aval = core.UnshapedArray(jnp.dtype(jnp.uint32)) + return (aval,) * 2 + + +rotate_left = _make_rotate_left(np.uint32) + + +def apply_round(v, rot): + v = v[:] + v[0] = v[0] + v[1] + v[1] = rotate_left(v[1], rot) + v[1] = v[0] ^ v[1] + return v + + +def rotate_list(xs): + return xs[1:] + xs[:1] + + +def rolled_loop_step(i, state): + x, ks, rotations = state + for r in rotations[0]: + x = apply_round(x, r) + new_x = [x[0] + ks[0], x[1] + ks[1] + jnp.asarray(i + 1, dtype=np.uint32)] + return new_x, rotate_list(ks), rotate_list(rotations) + + +def _threefry2x32_lowering(key1, key2, x1, x2, use_rolled_loops=True): + """Apply the Threefry 2x32 hash. + + Args: + keypair: a pair of 32bit unsigned integers used for the key. + count: an array of dtype uint32 used for the counts. + + Returns: + An array of dtype uint32 with the same shape as `count`. + """ + x = [x1, x2] + + rotations = [np.array([13, 15, 26, 6], dtype=np.uint32), + np.array([17, 29, 16, 24], dtype=np.uint32)] + ks = [key1, key2, key1 ^ key2 ^ np.uint32(0x1BD11BDA)] + + x[0] = x[0] + ks[0] + x[1] = x[1] + ks[1] + + if use_rolled_loops: + x, _, _ = lax.fori_loop(0, 5, rolled_loop_step, (x, rotate_list(ks), rotations)) + + else: + for r in rotations[0]: + x = apply_round(x, r) + x[0] = x[0] + ks[1] + x[1] = x[1] + ks[2] + np.uint32(1) + + for r in rotations[1]: + x = apply_round(x, r) + x[0] = x[0] + ks[2] + x[1] = x[1] + ks[0] + np.uint32(2) + + for r in rotations[0]: + x = apply_round(x, r) + x[0] = x[0] + ks[0] + x[1] = x[1] + ks[1] + np.uint32(3) + + for r in rotations[1]: + x = apply_round(x, r) + x[0] = x[0] + ks[1] + x[1] = x[1] + ks[2] + np.uint32(4) + + for r in rotations[0]: + x = apply_round(x, r) + x[0] = x[0] + ks[2] + x[1] = x[1] + ks[0] + np.uint32(5) + + return tuple(x) + + +_threefry2x32_lowering_rule = mlir.lower_fun( + partial(_threefry2x32_lowering, use_rolled_loops=False), + multiple_results=True) + +_threefry2x32_cpu_lowering_rule = mlir.lower_fun( + partial(_threefry2x32_lowering, use_rolled_loops=True), + multiple_results=True) + + +def _threefry2x32_gpu_lowering_rule(lowering_func, ctx, k1, k2, x1, x2): + if not config.threefry_gpu_kernel_lowering.value: # back to default lowering + return _threefry2x32_lowering_rule(ctx, k1, k2, x1, x2) + + aval_out, aval_out_2 = ctx.avals_out + assert aval_out == aval_out_2 + k1_aval, k2_aval, x1_aval, x2_aval = ctx.avals_in + rank = len(aval_out.shape) + if 0 in aval_out.shape: + zeros = mlir.full_like_aval(ctx, 0, aval_out) + return [zeros, zeros] + def _broadcast(x, aval): + return mlir.broadcast_in_dim(ctx, x, aval_out, + broadcast_dimensions=range(rank - len(aval.shape), rank)) + + out_len = reduce(op.mul, aval_out.shape, 1) + if not core.is_constant_dim(out_len): + length = mlir.eval_dynamic_shape_as_tensor(ctx, [out_len]) + length = mlir.hlo.convert( + ir.RankedTensorType.get((1,), ir.IntegerType.get_signless(64)), + length) + output_shape = mlir.eval_dynamic_shape_as_tensor(ctx, aval_out.shape) + else: + length = int(out_len) # will be passed statically + output_shape = None + + return lowering_func( + (_broadcast(k1, k1_aval), _broadcast(k2, k2_aval)), + (_broadcast(x1, x1_aval), _broadcast(x2, x2_aval)), length, + output_shape) + +threefry2x32_p = core.Primitive("threefry2x32") +threefry2x32_p.multiple_results = True +threefry2x32_p.def_impl(partial(dispatch.apply_primitive, threefry2x32_p)) +threefry2x32_p.def_abstract_eval(_threefry2x32_abstract_eval) +batching.defbroadcasting(threefry2x32_p) +mlir.register_lowering( + threefry2x32_p, _threefry2x32_lowering_rule) +mlir.register_lowering( + threefry2x32_p, _threefry2x32_cpu_lowering_rule, platform='cpu') +mlir.register_lowering( + threefry2x32_p, + partial(_threefry2x32_gpu_lowering_rule, gpu_prng.cuda_threefry2x32), + platform='cuda') +mlir.register_lowering( + threefry2x32_p, + partial(_threefry2x32_gpu_lowering_rule, gpu_prng.rocm_threefry2x32), + platform='rocm') + + +def iota_2x32_shape(shape): + """Reshaped ``uint64`` iota, as two parallel ``uint32`` arrays. + + Setting aside representation, this function essentially computes the + equivalent of:: + + jax.lax.iota(dtype=np.uint64, size=math.prod(shape)).reshape(shape) + + However: + + * It returns two parallel ``uint32`` arrays instead of one + ``uint64`` array. This renders it invariant under either setting of + the system-wide ``jax_enable_x64`` configuration flag. + + * It lowers in a way such that the compiler's automatic SPMD + partitioner recognizes its partitionability. + + For example:: + + >>> import numpy as np + >>> from jax import lax + >>> from jax._src import prng + + >>> prng.iota_2x32_shape((3, 4)) + [Array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=uint32), + Array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]], dtype=uint32)] + + >>> def reshaped_iota(shape): + ... return lax.iota(size=math.prod(shape), dtype=np.uint32).reshape(shape) + ... + >>> reshaped_iota((3, 4)) + Array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]], dtype=uint32) + + Args: + shape: the output shape + + Returns: + A pair of ``uint32`` arrays ``(counts_hi, counts_lo)``, both of + shape ``shape``, representing the higher-order and lower-order 32 + bits of the 64 bit unsigned iota. + """ + if len(shape) == 0: + return (jnp.zeros((), np.dtype('uint32')),) * 2 + return iota_2x32_shape_p.bind(shape=shape) + +iota_2x32_shape_p = core.Primitive('iota_2x32_shape') +iota_2x32_shape_p.multiple_results = True +iota_2x32_shape_p.def_impl(partial(dispatch.apply_primitive, iota_2x32_shape_p)) + +@iota_2x32_shape_p.def_abstract_eval +def iota_2x32_shape_abstract_eval(*, shape): + return (core.ShapedArray(shape, np.dtype('uint32')),) * 2 + +def bcast_iotas_to_reshaped_iota( + add: Callable[[ir.Value, ir.Value], ir.Value], + mul: Callable[[core.DimSize, ir.Value], ir.Value], + shape: core.Shape, + iotas: Sequence[ir.Value]) -> ir.Value: + strides: core.Shape = (*(np.cumprod(shape[1:][::-1])[::-1]), 1) + return reduce(add, [mul(s, i) for i, s in zip(iotas, strides)]) + +def iota_2x32_shape_lowering(ctx, *, shape): + aval_out, _ = ctx.avals_out + aval_u64 = core.ShapedArray(shape, np.dtype('uint64')) + + def _add(x: ir.Value, y: ir.Value) -> ir.Value: + return mlir.hlo.add(x, y) + + def _mul(x: core.DimSize, y: ir.Value) -> ir.Value: + if core.is_constant_dim(x): + x_const = mlir.ir_constant(np.array(x, np.dtype('uint64'))) + else: + x_shape, = mlir.eval_dynamic_shape(ctx, (x,)) + x_const = hlo.convert( + ir.RankedTensorType.get( + [], + mlir.dtype_to_ir_type(np.dtype('uint64'))), x_shape) + x_bcast = mlir.broadcast_in_dim(ctx, x_const, aval_u64, + broadcast_dimensions=[]) + return mlir.hlo.multiply(x_bcast, y) + + assert len(shape) > 0 + + iotas = [mlir.iota(ctx, aval_u64, dimension=dimension) + for dimension in range(len(shape))] + counts = bcast_iotas_to_reshaped_iota(_add, _mul, shape, iotas) + shift = mlir.ir_constant(np.array(32, np.dtype('uint64'))) + shift = mlir.broadcast_in_dim(ctx, shift, aval_u64, + broadcast_dimensions=[]) + counts_shifted = mlir.hlo.shift_right_logical(counts, shift) + counts_lo = mlir.hlo.convert(mlir.aval_to_ir_type(aval_out), counts) + counts_hi = mlir.hlo.convert(mlir.aval_to_ir_type(aval_out), counts_shifted) + return counts_hi, counts_lo +mlir.register_lowering(iota_2x32_shape_p, iota_2x32_shape_lowering) + + +@partial(jit, inline=True) +def threefry_2x32(keypair, count): + """Apply the Threefry 2x32 hash. + + Args: + keypair: a pair of 32bit unsigned integers used for the key. + count: an array of dtype uint32 used for the counts. + + Returns: + An array of dtype uint32 with the same shape as `count`. + """ + key1, key2 = keypair + if not lax.dtype(key1) == lax.dtype(key2) == lax.dtype(count) == np.uint32: + msg = "threefry_2x32 requires uint32 arguments, got {}" + raise TypeError(msg.format([lax.dtype(x) for x in [key1, key2, count]])) + + odd_size = count.size % 2 + if not isinstance(odd_size, int): + msg = ("jax.random functions have limited support for shape polymorphism. " + "In particular, the product of the known dimensions must be even.") + raise core.InconclusiveDimensionOperation(msg) + + if odd_size: + x = list(jnp.split(jnp.concatenate([count.ravel(), np.uint32([0])]), 2)) + else: + x = list(jnp.split(count.ravel(), 2)) + + x = threefry2x32_p.bind(key1, key2, x[0], x[1]) + out = jnp.concatenate(x) + assert out.dtype == np.uint32 + return lax.reshape(out[:-1] if odd_size else out, count.shape) + + +def threefry_split(key: typing.Array, shape: Shape) -> typing.Array: + shape = tuple(unsafe_map(core.concrete_dim_or_error, shape)) + return _threefry_split(key, shape) + +@partial(jit, static_argnums=(1,)) +def _threefry_split(key, shape) -> typing.Array: + if config.threefry_partitionable.value: + return _threefry_split_foldlike(key, shape) + else: + return _threefry_split_original(key, shape) + +@partial(jit, static_argnums=(1,), inline=True) +def _threefry_split_original(key, shape) -> typing.Array: + num = math.prod(shape) + counts = lax.iota(np.uint32, num * 2) + return lax.reshape(threefry_2x32(key, counts), (*shape, 2)) + +@partial(jit, static_argnums=(1,), inline=True) +def _threefry_split_foldlike(key, shape) -> typing.Array: + k1, k2 = key + counts1, counts2 = iota_2x32_shape(shape) + bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2) + return jnp.stack([bits1, bits2], axis=bits1.ndim) + + +def threefry_fold_in(key: typing.Array, data: typing.Array) -> typing.Array: + assert not data.shape + return _threefry_fold_in(key, jnp.uint32(data)) + +@jit +def _threefry_fold_in(key, data): + return threefry_2x32(key, threefry_seed(data)) + + +def threefry_random_bits(key: typing.Array, bit_width, shape): + """Sample uniform random bits of given width and shape using PRNG key.""" + if not _is_threefry_prng_key(key): + raise TypeError("threefry_random_bits got invalid prng key.") + if bit_width not in (8, 16, 32, 64): + raise TypeError("requires 8-, 16-, 32- or 64-bit field width.") + + if config.threefry_partitionable.value: + return _threefry_random_bits_partitionable(key, bit_width, shape) + else: + return _threefry_random_bits_original(key, bit_width, shape) + +def _threefry_random_bits_partitionable(key: typing.Array, bit_width, shape): + if all(core.is_constant_dim(d) for d in shape) and math.prod(shape) > 2 ** 64: + raise NotImplementedError('random bits array of size exceeding 2 ** 64') + + k1, k2 = key + counts1, counts2 = iota_2x32_shape(shape) + bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2) + + dtype = UINT_DTYPES[bit_width] + if bit_width == 64: + bits_hi = lax.convert_element_type(bits1, dtype) + bits_lo = lax.convert_element_type(bits2, dtype) + return lax.shift_left(bits_hi, dtype(32)) | bits_lo + elif bit_width == 32: + return bits1 ^ bits2 + else: + return lax.convert_element_type(bits1 ^ bits2, dtype) + +@partial(jit, static_argnums=(1, 2), inline=True) +def _threefry_random_bits_original(key: typing.Array, bit_width, shape): + size = math.prod(shape) + # Compute ceil(bit_width * size / 32) in a way that is friendly to shape + # polymorphism + max_count, r = divmod(bit_width * size, 32) + if r > 0: + max_count += 1 + + if core.is_constant_dim(max_count): + nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max) + else: + nblocks, rem = 0, max_count + + if not nblocks: + bits = threefry_2x32(key, lax.iota(np.uint32, rem)) + else: + keys = threefry_split(key, (nblocks + 1,)) + subkeys, last_key = keys[:-1], keys[-1] + blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max)) + last = threefry_2x32(last_key, lax.iota(np.uint32, rem)) + bits = lax.concatenate([blocks.ravel(), last], 0) + + dtype = UINT_DTYPES[bit_width] + if bit_width == 64: + bits = [lax.convert_element_type(x, dtype) for x in jnp.split(bits, 2)] + bits = lax.shift_left(bits[0], dtype(32)) | bits[1] + elif bit_width in [8, 16]: + # this is essentially bits.view(dtype)[:size] + bits = lax.bitwise_and( + np.uint32(np.iinfo(dtype).max), + lax.shift_right_logical( + lax.broadcast(bits, (1,)), + lax.mul( + np.uint32(bit_width), + lax.broadcasted_iota(np.uint32, (32 // bit_width, 1), 0) + ) + ) + ) + bits = lax.reshape(bits, ((max_count * 32 // bit_width),), (1, 0)) + bits = lax.convert_element_type(bits, dtype)[:size] + return lax.reshape(bits, shape) + + +threefry_prng_impl = PRNGImpl( + key_shape=(2,), + seed=threefry_seed, + split=threefry_split, + random_bits=threefry_random_bits, + fold_in=threefry_fold_in, + name='threefry2x32', + tag='fry') + +register_prng(threefry_prng_impl) + + +# -- RngBitGenerator PRNG implementation + +# This code is experimental! +# https://www.tensorflow.org/xla/operation_semantics#rngbitgenerator +# Notice that the RngBitGenerator operations are not guaranteed to be +# stable/deterministic across backends or compiler versions. Correspondingly, we +# reserve the right to change any of these implementations at any time! + +def _rbg_seed(seed: typing.Array) -> typing.Array: + assert not seed.shape + halfkey = threefry_seed(seed) + return jnp.concatenate([halfkey, halfkey]) + +def _rbg_split(key: typing.Array, shape: Shape) -> typing.Array: + if config.threefry_partitionable.value: + _threefry_split = _threefry_split_foldlike + else: + _threefry_split = _threefry_split_original + halfkeys = key.reshape(2, 2) + return vmap( + _threefry_split, (0, None), len(shape))(halfkeys, shape).reshape( + *shape, 4) + +def _rbg_fold_in(key: typing.Array, data: typing.Array) -> typing.Array: + assert not data.shape + return vmap(_threefry_fold_in, (0, None), 0)(key.reshape(2, 2), data).reshape(4) + +def _rbg_random_bits(key: typing.Array, bit_width: int, shape: Sequence[int] + ) -> typing.Array: + if not key.shape == (4,) and key.dtype == jnp.dtype('uint32'): + raise TypeError("_rbg_random_bits got invalid prng key.") + if bit_width not in (8, 16, 32, 64): + raise TypeError("requires 8-, 16-, 32- or 64-bit field width.") + _, bits = lax.rng_bit_generator(key, shape, dtype=UINT_DTYPES[bit_width]) + return bits + +rbg_prng_impl = PRNGImpl( + key_shape=(4,), + seed=_rbg_seed, + split=_rbg_split, + random_bits=_rbg_random_bits, + fold_in=_rbg_fold_in, + name='rbg', + tag='rbg') + +register_prng(rbg_prng_impl) + + +def _unsafe_rbg_split(key: typing.Array, shape: Shape) -> typing.Array: + # treat 10 iterations of random bits as a 'hash function' + num = math.prod(shape) + _, keys = lax.rng_bit_generator(key, (10 * num, 4), dtype='uint32') + return lax.slice_in_dim( + keys, start_index=None, limit_index=None, stride=10).reshape(*shape, 4) + +def _unsafe_rbg_fold_in(key: typing.Array, data: typing.Array) -> typing.Array: + assert not data.shape + _, random_bits = lax.rng_bit_generator(_rbg_seed(data), (10, 4), dtype='uint32') + return key ^ random_bits[-1] + +unsafe_rbg_prng_impl = PRNGImpl( + key_shape=(4,), + seed=_rbg_seed, + split=_unsafe_rbg_split, + random_bits=_rbg_random_bits, + fold_in=_unsafe_rbg_fold_in, + name='unsafe_rbg', + tag='urbg') + +register_prng(unsafe_rbg_prng_impl) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/profiler.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..d761d50e6fa80ba1b030aceb402b2f2ac571b6b4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/profiler.py @@ -0,0 +1,441 @@ +# Copyright 2020 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from contextlib import contextmanager +from functools import wraps +import glob +import gzip +import http.server +import json +import logging +import os +import socketserver +import threading +from typing import Callable, List, Optional, Union, Any + +from jax._src import traceback_util +traceback_util.register_exclusion(__file__) + +from jax._src import xla_bridge +from jax._src.lib import xla_client + +_profiler_server: xla_client.profiler.ProfilerServer | None = None + +logger = logging.getLogger(__name__) + + +def start_server(port: int) -> xla_client.profiler.ProfilerServer: + """Starts the profiler server on port `port`. + + Using the "TensorFlow profiler" feature in `TensorBoard + `_ 2.2 or newer, you can + connect to the profiler server and sample execution traces that show CPU, + GPU, and/or TPU device activity. + """ + global _profiler_server + if _profiler_server is not None: + raise ValueError("Only one profiler server can be active at a time.") + + # Make sure backends are initialized before creating a profiler + # session. Otherwise on Cloud TPU, libtpu may not be initialized before + # creating the tracer, which will cause the TPU tracer initialization to + # fail and no TPU operations will be included in the profile. + # NOTE(skyewm): I'm not sure this is necessary for start_server (is definitely + # is for start_trace), but I'm putting it here to be safe. + xla_bridge.get_backend() + + _profiler_server = xla_client.profiler.start_server(port) + return _profiler_server + + +def stop_server(): + """Stops the running profiler server.""" + global _profiler_server + if _profiler_server is None: + raise ValueError("No active profiler server.") + _profiler_server = None # Should destroy the profiler server + + +class _ProfileState: + def __init__(self): + self.profile_session = None + self.log_dir = None + self.create_perfetto_link = False + self.create_perfetto_trace = False + self.lock = threading.Lock() + + def reset(self): + _profile_state.profile_session = None + _profile_state.create_perfetto_link = False + _profile_state.create_perfetto_trace = False + _profile_state.log_dir = None + + +_profile_state = _ProfileState() + + +def start_trace(log_dir, create_perfetto_link: bool = False, + create_perfetto_trace: bool = False) -> None: + """Starts a profiler trace. + + The trace will capture CPU, GPU, and/or TPU activity, including Python + functions and JAX on-device operations. Use :func:`stop_trace` to end the trace + and save the results to ``log_dir``. + + The resulting trace can be viewed with TensorBoard. Note that TensorBoard + doesn't need to be running when collecting the trace. + + Only one trace may be collected at a time. A RuntimeError will be raised if + :func:`start_trace` is called while another trace is running. + + Args: + log_dir: The directory to save the profiler trace to (usually the + TensorBoard log directory). + create_perfetto_link: A boolean which, if true, creates and prints link to + the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will + block until the link is opened and Perfetto loads the trace. + create_perfetto_trace: A boolean which, if true, additionally dumps a + ``perfetto_trace.json.gz`` file that is compatible for upload with the + Perfetto trace viewer UI (https://ui.perfetto.dev). The file will also be + generated if ``create_perfetto_link`` is true. This could be useful if you + want to generate a Perfetto-compatible trace without blocking the + process. + """ + with _profile_state.lock: + if _profile_state.profile_session is not None: + raise RuntimeError("Profile has already been started. " + "Only one profile may be run at a time.") + # Make sure backends are initialized before creating a profiler + # session. Otherwise on Cloud TPU, libtpu may not be initialized before + # creating the tracer, which will cause the TPU tracer initialization to + # fail and no TPU operations will be included in the profile. + xla_bridge.get_backend() + + _profile_state.profile_session = xla_client.profiler.ProfilerSession() + _profile_state.create_perfetto_link = create_perfetto_link + _profile_state.create_perfetto_trace = ( + create_perfetto_trace or create_perfetto_link) + _profile_state.log_dir = log_dir + + +def _write_perfetto_trace_file(log_dir): + # Navigate to folder with the latest trace dump to find `trace.json.jz` + curr_path = os.path.abspath(log_dir) + root_trace_folder = os.path.join(curr_path, "plugins", "profile") + trace_folders = [os.path.join(root_trace_folder, trace_folder) for + trace_folder in os.listdir(root_trace_folder)] + latest_folder = max(trace_folders, key=os.path.getmtime) + trace_jsons = glob.glob(os.path.join(latest_folder, "*.trace.json.gz")) + if len(trace_jsons) != 1: + raise ValueError(f"Invalid trace folder: {latest_folder}") + trace_json, = trace_jsons + + logger.info("Loading trace.json.gz and removing its metadata...") + # Perfetto doesn't like the `metadata` field in `trace.json` so we remove + # it. + # TODO(sharadmv): speed this up by updating the generated `trace.json` + # to not include metadata if possible. + with gzip.open(trace_json, "rb") as fp: + trace = json.load(fp) + del trace["metadata"] + filename = "perfetto_trace.json.gz" + perfetto_trace = os.path.join(latest_folder, filename) + logger.info("Writing perfetto_trace.json.gz...") + with gzip.open(perfetto_trace, "w") as fp: + fp.write(json.dumps(trace).encode("utf-8")) + return perfetto_trace + +class _PerfettoServer(http.server.SimpleHTTPRequestHandler): + """Handles requests from `ui.perfetto.dev` for the `trace.json`""" + + def end_headers(self): + self.send_header('Access-Control-Allow-Origin', '*') + return super().end_headers() + + def do_GET(self): + self.server.last_request = self.path + return super().do_GET() + + def do_POST(self): + self.send_error(404, "File not found") + +def _host_perfetto_trace_file(path): + # ui.perfetto.dev looks for files hosted on `127.0.0.1:9001`. We set up a + # TCP server that is hosting the `perfetto_trace.json.gz` file. + port = 9001 + orig_directory = os.path.abspath(os.getcwd()) + directory, filename = os.path.split(path) + try: + os.chdir(directory) + socketserver.TCPServer.allow_reuse_address = True + with socketserver.TCPServer(('127.0.0.1', port), _PerfettoServer) as httpd: + url = f"https://ui.perfetto.dev/#!/?url=http://127.0.0.1:{port}/{filename}" + print(f"Open URL in browser: {url}") + + # Once ui.perfetto.dev acquires trace.json from this server we can close + # it down. + while httpd.__dict__.get('last_request') != '/' + filename: + httpd.handle_request() + finally: + os.chdir(orig_directory) + +def stop_trace(): + """Stops the currently-running profiler trace. + + The trace will be saved to the ``log_dir`` passed to the corresponding + :func:`start_trace` call. Raises a RuntimeError if a trace hasn't been started. + """ + with _profile_state.lock: + if _profile_state.profile_session is None: + raise RuntimeError("No profile started") + sess = _profile_state.profile_session + sess.export(sess.stop(), _profile_state.log_dir) + if _profile_state.create_perfetto_trace: + abs_filename = _write_perfetto_trace_file(_profile_state.log_dir) + if _profile_state.create_perfetto_link: + _host_perfetto_trace_file(abs_filename) + _profile_state.reset() + + +def stop_and_get_fdo_profile() -> Union[bytes, str]: + """Stops the currently-running profiler trace and export fdo_profile. + + Currently, this is only supported for GPU. + Raises a RuntimeError if a trace hasn't been started. + """ + with _profile_state.lock: + if _profile_state.profile_session is None: + raise RuntimeError("No profile started") + xspace = _profile_state.profile_session.stop() + fdo_profile = xla_client.profiler.get_fdo_profile(xspace) + _profile_state.reset() + return fdo_profile + + +@contextmanager +def trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False): + """Context manager to take a profiler trace. + + The trace will capture CPU, GPU, and/or TPU activity, including Python + functions and JAX on-device operations. + + The resulting trace can be viewed with TensorBoard. Note that TensorBoard + doesn't need to be running when collecting the trace. + + Only one trace may be collected at a time. A RuntimeError will be raised if a + trace is started while another trace is running. + + Args: + log_dir: The directory to save the profiler trace to (usually the + TensorBoard log directory). + create_perfetto_link: A boolean which, if true, creates and prints link to + the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will + block until the link is opened and Perfetto loads the trace. + create_perfetto_trace: A boolean which, if true, additionally dumps a + ``perfetto_trace.json.gz`` file that is compatible for upload with the + Perfetto trace viewer UI (https://ui.perfetto.dev). The file will also be + generated if ``create_perfetto_link`` is true. This could be useful if you + want to generate a Perfetto-compatible trace without blocking the + process. + """ + start_trace(log_dir, create_perfetto_link, create_perfetto_trace) + try: + yield + finally: + stop_trace() + + +class TraceAnnotation(xla_client.profiler.TraceMe): + """Context manager that generates a trace event in the profiler. + + The trace event spans the duration of the code enclosed by the context. + + For example: + + >>> x = jnp.ones((1000, 1000)) + >>> with jax.profiler.TraceAnnotation("my_label"): + ... result = jnp.dot(x, x.T).block_until_ready() + + This will cause a "my_label" event to show up on the trace timeline if the + event occurs while the process is being traced. + """ + pass + + +class StepTraceAnnotation(TraceAnnotation): + """Context manager that generates a step trace event in the profiler. + + The step trace event spans the duration of the code enclosed by the context. + The profiler will provide the performance analysis for each step trace event. + + For example, it can be used to mark training steps and enable the profiler to + provide the performance analysis per step: + + >>> while global_step < NUM_STEPS: # doctest: +SKIP + ... with jax.profiler.StepTraceAnnotation("train", step_num=global_step): # doctest: +SKIP + ... train_step() # doctest: +SKIP + ... global_step += 1 # doctest: +SKIP + + This will cause a "train xx" event to show up on the trace timeline if the + event occurs while the process is being traced by TensorBoard. In addition, + if using accelerators, the device trace timeline will also show a "train xx" + event. Note that "step_num" can be set as a keyword argument to pass the + global step number to the profiler. + + """ + + def __init__(self, name: str, **kwargs): + super().__init__(name, _r=1, **kwargs) + + +def annotate_function(func: Callable, name: str | None = None, + **decorator_kwargs): + """Decorator that generates a trace event for the execution of a function. + + For example: + + >>> @jax.profiler.annotate_function + ... def f(x): + ... return jnp.dot(x, x.T).block_until_ready() + >>> + >>> result = f(jnp.ones((1000, 1000))) + + This will cause an "f" event to show up on the trace timeline if the + function execution occurs while the process is being traced by TensorBoard. + + Arguments can be passed to the decorator via :py:func:`functools.partial`. + + >>> from functools import partial + + >>> @partial(jax.profiler.annotate_function, name="event_name") + ... def f(x): + ... return jnp.dot(x, x.T).block_until_ready() + + >>> result = f(jnp.ones((1000, 1000))) + """ + + name = name or getattr(func, '__qualname__', None) + name = name or func.__name__ + @wraps(func) + def wrapper(*args, **kwargs): + with TraceAnnotation(name, **decorator_kwargs): + return func(*args, **kwargs) + return wrapper + return wrapper + + +def device_memory_profile(backend: str | None = None) -> bytes: + """Captures a JAX device memory profile as ``pprof``-format protocol buffer. + + A device memory profile is a snapshot of the state of memory, that describes the JAX + :class:`~jax.Array` and executable objects present in memory and their + allocation sites. + + For more information how to use the device memory profiler, see + :doc:`/device_memory_profiling`. + + The profiling system works by instrumenting JAX on-device allocations, + capturing a Python stack trace for each allocation. The instrumentation is + always enabled; :func:`device_memory_profile` provides an API to capture it. + + The output of :func:`device_memory_profile` is a binary protocol buffer that + can be interpreted and visualized by the `pprof tool + `_. + + Args: + backend: optional; the name of the JAX backend for which the device memory + profile should be collected. + + Returns: + A byte string containing a binary `pprof`-format protocol buffer. + """ + return xla_client.heap_profile(xla_bridge.get_backend(backend)) + + +def save_device_memory_profile(filename, backend: str | None = None) -> None: + """Collects a device memory profile and writes it to a file. + + :func:`save_device_memory_profile` is a convenience wrapper around :func:`device_memory_profile` + that saves its output to a ``filename``. See the + :func:`device_memory_profile` documentation for more information. + + Args: + filename: the filename to which the profile should be written. + backend: optional; the name of the JAX backend for which the device memory + profile should be collected. + """ + profile = device_memory_profile(backend) + with open(filename, "wb") as f: + f.write(profile) + + +# Allows to run model with profiler given amount of times. After required amount +# of retries achived client can collect FDO data. +class PGLEProfiler: + + def __init__(self, retries: int, percentile: int): + self.retries: int = retries + self.percentile: int = percentile + self.collected_fdo: str | None = None + self.called_times: int = 0 + self.fdo_profiles: List[Any] = [] + self.current_session: xla_client.profiler.ProfilerSession | None = None + + def consume_fdo_profile(self) -> Optional[str]: + if self.collected_fdo is not None: + return self.collected_fdo + + if not self.is_enabled() or self.called_times != self.retries: + return None + + self.collected_fdo = xla_client.profiler.aggregate_profiled_instructions( + self.fdo_profiles, self.percentile + ) + return self.collected_fdo + + def is_fdo_consumed(self): + return self.collected_fdo is not None + + def disable(self): + self.retries = 0 + + def is_enabled(self): + return self.retries > 0 + + def is_running(self): + return self.current_session is not None + + @classmethod + @contextmanager + def trace(cls, runner: PGLEProfiler | None): + if (runner is None or runner.is_running() + or not runner.is_enabled() or runner.is_fdo_consumed()): + yield + else: + options = xla_client.profiler.ProfileOptions() + options.enable_hlo_proto = True + runner.current_session = xla_client.profiler.ProfilerSession(options) + + try: + yield + finally: + xspace = runner.current_session.stop() + runner.fdo_profiles.append( + xla_client.profiler.get_fdo_profile(xspace) + ) + runner.current_session = None + + runner.called_times += 1 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/public_test_util.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/public_test_util.py new file mode 100644 index 0000000000000000000000000000000000000000..c7cb4ee20e1bf0ae4265462cb869f35c28491b05 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/public_test_util.py @@ -0,0 +1,300 @@ +# Copyright 2022 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from functools import partial +import operator + +from jax._src import api +from jax._src import config +from jax._src import dtypes as _dtypes +from jax._src.tree_util import tree_map, tree_reduce + +import numpy as np + + +# The only functions intended to be exported are these; they should be used via +# jax.test_util. All other functionality appearing here is for internal use only, +# and may be changed or removed at any time and without any deprecation cycle. +__all__ = ['check_grads', 'check_jvp', 'check_vjp'] + + +EPS = 1e-4 + + +def _dtype(x): + if hasattr(x, 'dtype'): + return x.dtype + elif type(x) in _dtypes.python_scalar_dtypes: + return np.dtype(_dtypes.python_scalar_dtypes[type(x)]) + else: + return np.asarray(x).dtype + + +_default_tolerance = { + _dtypes.float0: 0, + np.dtype(np.bool_): 0, + np.dtype(_dtypes.int4): 0, + np.dtype(np.int8): 0, + np.dtype(np.int16): 0, + np.dtype(np.int32): 0, + np.dtype(np.int64): 0, + np.dtype(_dtypes.uint4): 0, + np.dtype(np.uint8): 0, + np.dtype(np.uint16): 0, + np.dtype(np.uint32): 0, + np.dtype(np.uint64): 0, + np.dtype(_dtypes.float8_e4m3b11fnuz): 1e-1, + np.dtype(_dtypes.float8_e4m3fn): 1e-1, + np.dtype(_dtypes.float8_e4m3fnuz): 1e-1, + np.dtype(_dtypes.float8_e5m2): 1e-1, + np.dtype(_dtypes.float8_e5m2fnuz): 1e-1, + np.dtype(_dtypes.bfloat16): 1e-2, + np.dtype(np.float16): 1e-3, + np.dtype(np.float32): 1e-6, + np.dtype(np.float64): 1e-15, + np.dtype(np.complex64): 1e-6, + np.dtype(np.complex128): 1e-15, +} + +def default_tolerance(): + return _default_tolerance + + +default_gradient_tolerance = { + np.dtype(_dtypes.float8_e4m3b11fnuz): 1e-1, + np.dtype(_dtypes.float8_e4m3fn): 1e-1, + np.dtype(_dtypes.float8_e4m3fnuz): 1e-1, + np.dtype(_dtypes.float8_e5m2): 1e-1, + np.dtype(_dtypes.float8_e5m2fnuz): 1e-1, + np.dtype(_dtypes.bfloat16): 1e-1, + np.dtype(np.float16): 1e-2, + np.dtype(np.float32): 2e-3, + np.dtype(np.float64): 1e-5, + np.dtype(np.complex64): 1e-3, + np.dtype(np.complex128): 1e-5, +} + +def is_python_scalar(val): + return not isinstance(val, np.generic) and isinstance(val, (bool, int, float, complex)) + +def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''): + if a.dtype == b.dtype == _dtypes.float0: + np.testing.assert_array_equal(a, b, err_msg=err_msg) + return + + custom_float_dtypes = [ + _dtypes.float8_e4m3b11fnuz, + _dtypes.float8_e4m3fn, + _dtypes.float8_e4m3fnuz, + _dtypes.float8_e5m2, + _dtypes.float8_e5m2fnuz, + _dtypes.bfloat16, + ] + def maybe_upcast(x): + if x.dtype in custom_float_dtypes: + return x.astype(np.float32) + return x + + a = maybe_upcast(a) + b = maybe_upcast(b) + + kw = {} + if atol: kw["atol"] = atol + if rtol: kw["rtol"] = rtol + with np.errstate(invalid='ignore'): + # TODO(phawkins): surprisingly, assert_allclose sometimes reports invalid + # value errors. It should not do that. + np.testing.assert_allclose(a, b, **kw, err_msg=err_msg) + +def tolerance(dtype, tol=None): + tol = {} if tol is None else tol + if not isinstance(tol, dict): + return tol + tol = {np.dtype(key): value for key, value in tol.items()} + dtype = _dtypes.canonicalize_dtype(np.dtype(dtype)) + return tol.get(dtype, default_tolerance()[dtype]) + + +def _assert_numpy_close(a, b, atol=None, rtol=None, err_msg=''): + a, b = np.asarray(a), np.asarray(b) + assert a.shape == b.shape + atol = max(tolerance(a.dtype, atol), tolerance(b.dtype, atol)) + rtol = max(tolerance(a.dtype, rtol), tolerance(b.dtype, rtol)) + _assert_numpy_allclose(a, b, atol=atol * a.size, rtol=rtol * b.size, + err_msg=err_msg) + + +def check_close(xs, ys, atol=None, rtol=None, err_msg=''): + assert_close = partial(_assert_numpy_close, atol=atol, rtol=rtol, + err_msg=err_msg) + tree_map(assert_close, xs, ys) + + +def _check_dtypes_match(xs, ys): + def _assert_dtypes_match(x, y): + if config.enable_x64.value: + assert _dtype(x) == _dtype(y) + else: + assert (_dtypes.canonicalize_dtype(_dtype(x)) == + _dtypes.canonicalize_dtype(_dtype(y))) + tree_map(_assert_dtypes_match, xs, ys) + + +def inner_prod(xs, ys): + def contract(x, y): + return np.real(np.dot(np.conj(x).reshape(-1), y.reshape(-1))) + return tree_reduce(np.add, tree_map(contract, xs, ys)) + + +def _safe_subtract(x, y, *, dtype): + """Subtraction that with `inf - inf == 0` semantics.""" + with np.errstate(invalid='ignore'): + return np.where(np.equal(x, y), np.array(0, dtype), + np.subtract(x, y, dtype=dtype)) + +def _preserve_input_types(f): + def wrapped(*args): + dtype = _dtype(args[0]) + result = np.array(f(*args), dtype=dtype) + if all(is_python_scalar(arg) for arg in args): + result = result.item() + return result + return wrapped + + +add = partial(tree_map, _preserve_input_types(operator.add)) +sub = partial(tree_map, _preserve_input_types(operator.sub)) +safe_sub = partial(tree_map, + lambda x, y: _safe_subtract(x, y, dtype=_dtype(x))) +conj = partial(tree_map, _preserve_input_types(np.conj)) + + +def scalar_mul(xs, a): + def mul(x): + dtype = _dtype(x) + result = np.multiply(x, np.array(a, dtype=dtype), dtype=dtype) + return result.item() if is_python_scalar(x) else result + return tree_map(mul, xs) + + +def rand_like(rng, x): + shape = np.shape(x) + dtype = _dtype(x) + randn = lambda: np.asarray(rng.randn(*shape), dtype=dtype) + if _dtypes.issubdtype(dtype, np.complexfloating): + result = randn() + dtype.type(1.0j) * randn() + else: + result = randn() + return result.item() if is_python_scalar(x) else result + + +def numerical_jvp(f, primals, tangents, eps=EPS): + delta = scalar_mul(tangents, eps) + f_pos = f(*add(primals, delta)) + f_neg = f(*sub(primals, delta)) + return scalar_mul(safe_sub(f_pos, f_neg), 0.5 / eps) + + +def _merge_tolerance(tol, default): + if tol is None: + return default + if not isinstance(tol, dict): + return tol + out = default.copy() + for k, v in tol.items(): + out[np.dtype(k)] = v + return out + + +def check_jvp(f, f_jvp, args, atol=None, rtol=None, eps=EPS, err_msg=''): + atol = _merge_tolerance(atol, default_gradient_tolerance) + rtol = _merge_tolerance(rtol, default_gradient_tolerance) + rng = np.random.RandomState(0) + tangent = tree_map(partial(rand_like, rng), args) + v_out, t_out = f_jvp(args, tangent) + _check_dtypes_match(v_out, t_out) + v_out_expected = f(*args) + _check_dtypes_match(v_out, v_out_expected) + t_out_expected = numerical_jvp(f, args, tangent, eps=eps) + # In principle we should expect exact equality of v_out and v_out_expected, + # but due to nondeterminism especially on GPU (e.g., due to convolution + # autotuning) we only require "close". + check_close(v_out, v_out_expected, atol=atol, rtol=rtol, + err_msg=f'{err_msg} primal' if err_msg else 'primal') + check_close(t_out, t_out_expected, atol=atol, rtol=rtol, + err_msg=f'{err_msg} tangent' if err_msg else 'tangent') + + +def check_vjp(f, f_vjp, args, atol=None, rtol=None, eps=EPS, err_msg=''): + atol = _merge_tolerance(atol, default_gradient_tolerance) + rtol = _merge_tolerance(rtol, default_gradient_tolerance) + _rand_like = partial(rand_like, np.random.RandomState(0)) + v_out, vjpfun = f_vjp(*args) + v_out_expected = f(*args) + check_close(v_out, v_out_expected, atol=atol, rtol=rtol, + err_msg=f'{err_msg} primal' if err_msg else 'primal') + tangent = tree_map(_rand_like, args) + tangent_out = numerical_jvp(f, args, tangent, eps=eps) + cotangent = tree_map(_rand_like, v_out) + cotangent_out = conj(vjpfun(conj(cotangent))) + ip = inner_prod(tangent, cotangent_out) + ip_expected = inner_prod(tangent_out, cotangent) + check_close(ip, ip_expected, atol=atol, rtol=rtol, + err_msg=(f'{err_msg} cotangent projection' + if err_msg else 'cotangent projection')) + + +def check_grads(f, args, order, + modes=("fwd", "rev"), atol=None, rtol=None, eps=None): + """Check gradients from automatic differentiation against finite differences. + + Gradients are only checked in a single randomly chosen direction, which + ensures that the finite difference calculation does not become prohibitively + expensive even for large input/output spaces. + + Args: + f: function to check at ``f(*args)``. + args: tuple of argument values. + order: forward and backwards gradients up to this order are checked. + modes: lists of gradient modes to check ('fwd' and/or 'rev'). + atol: absolute tolerance for gradient equality. + rtol: relative tolerance for gradient equality. + eps: step size used for finite differences. + + Raises: + AssertionError: if gradients do not match. + """ + args = tuple(args) + eps = eps or EPS + + _check_jvp = partial(check_jvp, atol=atol, rtol=rtol, eps=eps) + _check_vjp = partial(check_vjp, atol=atol, rtol=rtol, eps=eps) + + def _check_grads(f, args, order, err_msg=''): + if "fwd" in modes: + fwd_msg = f'JVP of {err_msg}' if err_msg else 'JVP' + _check_jvp(f, partial(api.jvp, f), args, err_msg=fwd_msg) + if order > 1: + _check_grads(partial(api.jvp, f), (args, args), order - 1, fwd_msg) + + if "rev" in modes: + rev_msg = f'VJP of {err_msg}' if err_msg else 'VJP' + _check_vjp(f, partial(api.vjp, f), args, err_msg=rev_msg) + if order > 1: + def f_vjp(*args): + out_primal_py, vjp_py = api.vjp(f, *args) + return vjp_py(out_primal_py) + _check_grads(f_vjp, args, order - 1, rev_msg) + + _check_grads(f, args, order) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/random.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/random.py new file mode 100644 index 0000000000000000000000000000000000000000..6a024f758c579f4f98cfba91735971f462e0a4a4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/random.py @@ -0,0 +1,2632 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Hashable, Sequence +from functools import partial +import math +from operator import index +import typing +from typing import Union +import warnings + +import numpy as np + +import jax.numpy as jnp +from jax import lax +from jax.numpy.linalg import cholesky, svd, eigh + +from jax._src import config +from jax._src import core +from jax._src import dispatch +from jax._src import dtypes +from jax._src import prng +from jax._src import xla_bridge +from jax._src.api import jit, vmap +from jax._src.core import NamedShape +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src.interpreters import mlir +from jax._src.lax import lax as lax_internal +from jax._src.numpy.lax_numpy import _convert_and_clip_integer +from jax._src.numpy.util import _arraylike, check_arraylike, promote_dtypes_inexact +from jax._src.typing import Array, ArrayLike, DTypeLike +from jax._src.util import canonicalize_axis + + +RealArray = ArrayLike +IntegerArray = ArrayLike +# TODO: Import or define these to match +# https://github.com/numpy/numpy/blob/main/numpy/typing/_dtype_like.py. +DTypeLikeInt = DTypeLike +DTypeLikeUInt = DTypeLike +DTypeLikeFloat = DTypeLike +Shape = Sequence[int] + +PRNGImpl = prng.PRNGImpl +KeyArray = Array +KeyArrayLike = ArrayLike + +UINT_DTYPES = prng.UINT_DTYPES + + +### utilities + +_lax_const = lax_internal._const + +def _isnan(x: ArrayLike) -> Array: + return lax.ne(x, x) + + +def _check_prng_key(name: str, key: KeyArrayLike, *, + allow_batched: bool = False) -> tuple[KeyArray, bool]: + if isinstance(key, Array) and dtypes.issubdtype(key.dtype, dtypes.prng_key): + wrapped_key = key + wrapped = False + elif _arraylike(key): + # Call random_wrap here to surface errors for invalid keys. + wrapped_key = prng.random_wrap(key, impl=default_prng_impl()) + wrapped = True + if config.legacy_prng_key.value == 'error': + raise ValueError( + 'Legacy uint32 key array passed as key to jax.random function. ' + 'Please create keys using jax.random.key(). If use of a raw key array ' + 'was intended, set jax_legacy_prng_key="allow".') + elif config.legacy_prng_key.value == 'warn': + warnings.warn( + 'Legacy uint32 key array passed as key to jax.random function. ' + 'Please create keys using jax.random.key(). If use of a raw key array ' + 'was intended, set jax_legacy_prng_key="allow".', stacklevel=2) + elif config.enable_custom_prng.value: + # TODO(jakevdp): possibly remove this warning condition. + warnings.warn( + 'Raw arrays as random keys to jax.random functions are deprecated. ' + 'Assuming valid threefry2x32 key for now.', + FutureWarning) + else: + raise TypeError(f'unexpected PRNG key type {type(key)}') + + if (not allow_batched) and wrapped_key.ndim: + raise ValueError(f"{name} accepts a single key, but was given a key array of" + f" shape {np.shape(key)} != (). Use jax.vmap for batching.") + + return wrapped_key, wrapped + + +def _return_prng_keys(was_wrapped, key): + # TODO(frostig): remove once we always enable_custom_prng + assert jnp.issubdtype(key.dtype, dtypes.prng_key) + if config.enable_custom_prng.value: + return key + else: + return prng.random_unwrap(key) if was_wrapped else key + + +def _random_bits(key: KeyArray, bit_width: int, shape: Shape) -> Array: + assert jnp.issubdtype(key.dtype, dtypes.prng_key) + return prng.random_bits(key, bit_width=bit_width, shape=shape) + + +# TODO(frostig,vanderplas): remove from public API altogether, or at +# least change to return after asserting presence in `prng.prngs` +def default_prng_impl(): + """Get the default PRNG implementation. + + The default implementation is determined by ``config.jax_default_prng_impl``, + which specifies it by name. + """ + impl_name = config.default_prng_impl.value + assert impl_name in prng.prngs, impl_name + return prng.prngs[impl_name] + + +### key operations + +# Wrapper around prng.PRNGImpl meant to hide its attributes from the +# public API. +# TODO(frostig,vanderplas): consider hiding all the attributes of +# PRNGImpl and directly returning it. +class PRNGSpec: + """Specifies a PRNG key implementation.""" + + __slots__ = ['_impl'] + _impl: PRNGImpl + + def __init__(self, impl): + self._impl = impl + + def __str__(self) -> str: return str(self._impl) + def __hash__(self) -> int: return hash(self._impl) + + def __eq__(self, other) -> bool: + return self._impl == other._impl + + +# TODO(frostig,vanderplas): remove PRNGImpl from this union when it's +# no longer in the public API because `default_prng_impl` is gone +PRNGSpecDesc = Union[str, PRNGSpec, PRNGImpl] + + +def resolve_prng_impl(impl_spec: PRNGSpecDesc | None) -> PRNGImpl: + if impl_spec is None: + return default_prng_impl() + if type(impl_spec) is PRNGImpl: + # TODO(frostig,vanderplas): remove this case once we remove + # default_prng_impl (and thus PRNGImpl) from the public API and + # PRNGImpl from jex. We won't need to handle these then, and we + # can remove them from the input type annotation above as well. + return impl_spec + if type(impl_spec) is PRNGSpec: + return impl_spec._impl + if type(impl_spec) is str: + if impl_spec in prng.prngs: + return prng.prngs[impl_spec] + + keys_fmt = ', '.join(f'"{s}"' for s in prng.prngs.keys()) + raise ValueError(f'unrecognized PRNG implementation "{impl_spec}". ' + f'Did you mean one of: {keys_fmt}?') + + t = type(impl_spec) + raise TypeError(f'unrecognized type {t} for specifying PRNG implementation.') + + +def _key(ctor_name: str, seed: int | ArrayLike, + impl_spec: PRNGSpecDesc | None) -> KeyArray: + impl = resolve_prng_impl(impl_spec) + if hasattr(seed, 'dtype') and jnp.issubdtype(seed.dtype, dtypes.prng_key): + raise TypeError( + f"{ctor_name} accepts a scalar seed, but was given a PRNG key.") + if np.ndim(seed): + raise TypeError( + f"{ctor_name} accepts a scalar seed, but was given an array of " + f"shape {np.shape(seed)} != (). Use jax.vmap for batching") + return prng.random_seed(seed, impl=impl) + +def key(seed: int | ArrayLike, *, + impl: PRNGSpecDesc | None = None) -> KeyArray: + """Create a pseudo-random number generator (PRNG) key given an integer seed. + + The result is a scalar array with a key that indicates the default PRNG + implementation, as determined by the optional ``impl`` argument or, + otherwise, by the ``jax_default_prng_impl`` config flag. + + Args: + seed: a 64- or 32-bit integer used as the value of the key. + impl: optional string specifying the PRNG implementation (e.g. + ``'threefry2x32'``) + + Returns: + A scalar PRNG key array, consumable by random functions as well as ``split`` + and ``fold_in``. + """ + return _key('key', seed, impl) + +def PRNGKey(seed: int | ArrayLike, *, + impl: PRNGSpecDesc | None = None) -> KeyArray: + """Create a pseudo-random number generator (PRNG) key given an integer seed. + + The resulting key carries the default PRNG implementation, as + determined by the optional ``impl`` argument or, otherwise, by the + ``jax_default_prng_impl`` config flag. + + Args: + seed: a 64- or 32-bit integer used as the value of the key. + impl: optional string specifying the PRNG implementation (e.g. + ``'threefry2x32'``) + + Returns: + A PRNG key, consumable by random functions as well as ``split`` + and ``fold_in``. + """ + return _return_prng_keys(True, _key('PRNGKey', seed, impl)) + + +def fold_in(key: KeyArrayLike, data: IntegerArray) -> KeyArray: + """Folds in data to a PRNG key to form a new PRNG key. + + Args: + key: a PRNG key (from ``key``, ``split``, ``fold_in``). + data: a 32-bit integer representing data to be folded into the key. + + Returns: + A new PRNG key that is a deterministic function of the inputs and is + statistically safe for producing a stream of new pseudo-random values. + """ + key, wrapped = _check_prng_key("fold_in", key) + if np.ndim(data): + raise TypeError("fold_in accepts a scalar, but was given an array of" + f"shape {np.shape(data)} != (). Use jax.vmap for batching.") + key_out = prng.random_fold_in(key, jnp.uint32(data)) + return _return_prng_keys(wrapped, key_out) + + +def _split(key: KeyArray, num: int | tuple[int, ...] = 2) -> KeyArray: + # Alternative to split() to use within random samplers. + # TODO(frostig): remove and use split(); we no longer need to wait + # to always enable_custom_prng + assert jnp.issubdtype(key.dtype, dtypes.prng_key) + if key.ndim: + raise TypeError("split accepts a single key, but was given a key array of " + f"shape {key.shape} != (). Use jax.vmap for batching.") + shape = tuple(num) if isinstance(num, Sequence) else (num,) + return prng.random_split(key, shape=shape) + +def split(key: KeyArrayLike, num: int | tuple[int, ...] = 2) -> KeyArray: + """Splits a PRNG key into `num` new keys by adding a leading axis. + + Args: + key: a PRNG key (from ``key``, ``split``, ``fold_in``). + num: optional, a positive integer (or tuple of integers) indicating + the number (or shape) of keys to produce. Defaults to 2. + + Returns: + An array-like object of `num` new PRNG keys. + """ + typed_key, wrapped = _check_prng_key("split", key) + return _return_prng_keys(wrapped, _split(typed_key, num)) + + +def _key_impl(keys: KeyArray) -> PRNGImpl: + assert jnp.issubdtype(keys.dtype, dtypes.prng_key) + keys_dtype = typing.cast(prng.KeyTy, keys.dtype) + return keys_dtype._impl + +def key_impl(keys: KeyArrayLike) -> Hashable: + typed_keys, _ = _check_prng_key("key_impl", keys, allow_batched=True) + return PRNGSpec(_key_impl(typed_keys)) + + +def _key_data(keys: KeyArray) -> Array: + assert jnp.issubdtype(keys.dtype, dtypes.prng_key) + return prng.random_unwrap(keys) + +def key_data(keys: KeyArrayLike) -> Array: + """Recover the bits of key data underlying a PRNG key array.""" + keys, _ = _check_prng_key("key_data", keys, allow_batched=True) + return _key_data(keys) + + +def wrap_key_data(key_bits_array: Array, *, + impl: PRNGSpecDesc | None = None): + """Wrap an array of key data bits into a PRNG key array. + + Args: + key_bits_array: a ``uint32`` array with trailing shape corresponding to + the key shape of the PRNG implementation specified by ``impl``. + impl: optional, specifies a PRNG implementation, as in ``random.key``. + + Returns: + A PRNG key array, whose dtype is a subdtype of ``jax.dtypes.prng_key`` + corresponding to ``impl``, and whose shape equals the leading shape + of ``key_bits_array.shape`` up to the key bit dimensions. + """ + impl_obj = resolve_prng_impl(impl) + return prng.random_wrap(key_bits_array, impl=impl_obj) + + +### random samplers + + +def _check_shape(name: str, shape: Shape | NamedShape, *param_shapes) -> None: + shape = core.as_named_shape(shape) + + if param_shapes: + shape_ = lax.broadcast_shapes(shape.positional, *param_shapes) + if shape.positional != shape_: + msg = ("{} parameter shapes must be broadcast-compatible with shape " + "argument, and the result of broadcasting the shapes must equal " + "the shape argument, but got result {} for shape argument {}.") + raise ValueError(msg.format(name, shape_, shape)) + + +def bits(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeUInt | None = None) -> Array: + """Sample uniform bits in the form of unsigned integers. + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default ``()``. + dtype: optional, an unsigned integer dtype for the returned values (default + ``uint64`` if ``jax_enable_x64`` is true, otherwise ``uint32``). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("bits", key) + if dtype is None: + dtype = dtypes.canonicalize_dtype(jnp.uint) + else: + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.unsignedinteger): + raise ValueError("dtype argument to `bits` must be an unsigned int dtype, " + f"got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + bit_width = dtype.itemsize * 8 + return _random_bits(key, bit_width, shape) + + +def uniform(key: KeyArrayLike, + shape: Shape | NamedShape = (), + dtype: DTypeLikeFloat = float, + minval: RealArray = 0., + maxval: RealArray = 1.) -> Array: + """Sample uniform random values in [minval, maxval) with given shape/dtype. + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + minval: optional, a minimum (inclusive) value broadcast-compatible with shape for the range (default 0). + maxval: optional, a maximum (exclusive) value broadcast-compatible with shape for the range (default 1). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("uniform", key) + dtypes.check_user_dtype_supported(dtype) + + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `uniform` must be a float dtype, " + f"got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.as_named_shape(shape) + return _uniform(key, shape, dtype, minval, maxval) + +@partial(jit, static_argnums=(1, 2)) +def _uniform(key, shape, dtype, minval, maxval) -> Array: + _check_shape("uniform", shape) + if not jnp.issubdtype(dtype, np.floating): + raise TypeError("uniform only accepts floating point dtypes.") + + minval = lax.convert_element_type(minval, dtype) + maxval = lax.convert_element_type(maxval, dtype) + minval = lax.broadcast_to_rank(minval, shape.positional_rank) + maxval = lax.broadcast_to_rank(maxval, shape.positional_rank) + + finfo = jnp.finfo(dtype) + nbits, nmant = finfo.bits, finfo.nmant + + if nbits not in (8, 16, 32, 64): + raise TypeError( + f"uniform only accepts 8-, 16-, 32-, or 64-bit dtypesgot {dtype}." + ) + + rng_bits = nbits + if nmant < 8: + rng_bits = 8 + bits = _random_bits(key, rng_bits, shape) + uint_dtype = UINT_DTYPES[nbits] + if rng_bits != nbits: + bits = lax.convert_element_type(bits, uint_dtype) + + # The strategy here is to randomize only the mantissa bits with an exponent of + # 1 (after applying the bias), then shift and scale to the desired range. The + # bit-level transformation we use relies on Numpy and XLA having bit-for-bit + # equivalent float representations, which might not be true on all platforms. + float_bits = lax.bitwise_or( + lax.shift_right_logical(bits, np.array(rng_bits - nmant, uint_dtype)), + np.array(1.0, dtype).view(uint_dtype), + ) + floats = lax.bitcast_convert_type(float_bits, dtype) - np.array(1., dtype) + return lax.max( + minval, + lax.reshape(floats * (maxval - minval) + minval, shape.positional)) + + +def randint(key: KeyArrayLike, + shape: Shape, + minval: IntegerArray, + maxval: IntegerArray, + dtype: DTypeLikeInt = int) -> Array: + """Sample uniform random values in [minval, maxval) with given shape/dtype. + + Args: + key: a PRNG key used as the random key. + shape: a tuple of nonnegative integers representing the shape. + minval: int or array of ints broadcast-compatible with ``shape``, a minimum + (inclusive) value for the range. + maxval: int or array of ints broadcast-compatible with ``shape``, a maximum + (exclusive) value for the range. + dtype: optional, an int dtype for the returned values (default int64 if + jax_enable_x64 is true, otherwise int32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("randint", key) + dtypes.check_user_dtype_supported(dtype) + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _randint(key, shape, minval, maxval, dtype) + +@partial(jit, static_argnums=(1, 4)) +def _randint(key, shape, minval, maxval, dtype) -> Array: + _check_shape("randint", shape, np.shape(minval), np.shape(maxval)) + if not jnp.issubdtype(dtype, np.integer): + raise TypeError(f"randint only accepts integer dtypes, got {dtype}") + + check_arraylike("randint", minval, maxval) + minval = jnp.asarray(minval) + maxval = jnp.asarray(maxval) + if not jnp.issubdtype(minval.dtype, np.integer): + minval = minval.astype(int) + if not jnp.issubdtype(maxval.dtype, np.integer): + maxval = maxval.astype(int) + + # Flag where maxval is greater than the maximum value of dtype + # in order to handle cases like randint(key, shape, 0, 256, 'uint8') + maxval_out_of_range = lax.gt( + maxval, _convert_and_clip_integer(jnp.array(jnp.iinfo(dtype).max, dtype), maxval.dtype)) + + minval = _convert_and_clip_integer(minval, dtype) + maxval = _convert_and_clip_integer(maxval, dtype) + minval = lax.broadcast_to_rank(minval, len(shape)) + maxval = lax.broadcast_to_rank(maxval, len(shape)) + nbits = jnp.iinfo(dtype).bits + + if nbits not in (8, 16, 32, 64): + raise TypeError(f"randint only accepts 8-, 16-, 32-, or 64-bit dtypes, got {dtype}") + + # This algorithm is biased whenever (maxval - minval) is not a power of 2. + # We generate double the number of random bits required by the dtype so as to + # reduce that bias. + k1, k2 = _split(key) + rbits = lambda key: _random_bits(key, nbits, shape) + higher_bits, lower_bits = rbits(k1), rbits(k2) + + unsigned_dtype = UINT_DTYPES[nbits] + span = lax.convert_element_type(maxval - minval, unsigned_dtype) + + # Ensure that span=1 when maxval <= minval, so minval is always returned; + # https://github.com/google/jax/issues/222 + span = lax.select(maxval <= minval, lax.full_like(span, 1), span) + + # When maxval is out of range, the span has to be one larger. + # If span is already the maximum representable value, this will wrap to zero, + # causing remainders below to have no effect, which is the correct semantics. + span = lax.select( + maxval_out_of_range & (maxval > minval), + lax.add(span, _lax_const(span, 1)), + span) + + # To compute a remainder operation on an integer that might have twice as many + # bits as we can represent in the native unsigned dtype, we compute a + # multiplier equal to 2**nbits % span. To avoid overflow, we use the identity: + # (a * b) % N = [(a % N) * (b % N)] % N + multiplier = lax.rem(_lax_const(span, 2 ** (nbits // 2)), span) + multiplier = lax.rem(lax.mul(multiplier, multiplier), span) + + random_offset = lax.add(lax.mul(lax.rem(higher_bits, span), multiplier), + lax.rem(lower_bits, span)) + random_offset = lax.rem(random_offset, span) + return lax.add(minval, lax.convert_element_type(random_offset, dtype)) + + +def shuffle(key: KeyArrayLike, x: ArrayLike, axis: int = 0) -> Array: + """Shuffle the elements of an array uniformly at random along an axis. + + Args: + key: a PRNG key used as the random key. + x: the array to be shuffled. + axis: optional, an int axis along which to shuffle (default 0). + + Returns: + A shuffled version of x. + """ + msg = ("jax.random.shuffle is deprecated and will be removed in a future release. " + "Use jax.random.permutation with independent=True.") + warnings.warn(msg, FutureWarning) + key, _ = _check_prng_key("shuffle", key) + return _shuffle(key, x, axis) + + +def permutation(key: KeyArrayLike, + x: int | ArrayLike, + axis: int = 0, + independent: bool = False) -> Array: + """Returns a randomly permuted array or range. + + Args: + key: a PRNG key used as the random key. + x: int or array. If x is an integer, randomly shuffle np.arange(x). + If x is an array, randomly shuffle its elements. + axis: int, optional. The axis which x is shuffled along. Default is 0. + independent: bool, optional. If set to True, each individual vector along + the given axis is shuffled independently. Default is False. + + Returns: + A shuffled version of x or array range + """ + key, _ = _check_prng_key("permutation", key) + check_arraylike("permutation", x) + axis = canonicalize_axis(axis, np.ndim(x) or 1) + if not np.ndim(x): + if not np.issubdtype(lax.dtype(x), np.integer): + raise TypeError("x must be an integer or at least 1-dimensional") + r = core.concrete_or_error(int, x, 'argument x of jax.random.permutation()') + return _shuffle(key, jnp.arange(r), axis) + if independent or np.ndim(x) == 1: + return _shuffle(key, x, axis) + ind = _shuffle(key, jnp.arange(x.shape[axis]), 0) # type: ignore[union-attr] + return jnp.take(x, ind, axis, unique_indices=True) + + +@partial(jit, static_argnums=(2,)) +def _shuffle(key, x, axis) -> Array: + # On parallel architectures, Fisher-Yates is more expensive than doing + # multiple sorts. This algorithm is based on one developed and analyzed by + # tjablin@. We sort according to randomly-generated 32bit keys, but those keys + # may have collisions. If we repeat the process, using fresh 32bit keys for + # each sort, then whenever all pairs of elements have been assigned distinct + # keys at some iteration (or equivalently when the strings formed by + # concatenating the successive keys for each element are all distinct) then we + # are guaranteed to have a perfect sample (assuming that either the sort is + # stable or that any bias is not value-dependent). Since checking uniqueness + # at runtime may be expensive, we use a heuristic static stop criterion + # developed by tjablin@. See tensorflow/compiler/tf2xla/random_ops.cc for more + # info, and for the original implementation of this algorithm. See also + # Section 2 of http://people.csail.mit.edu/costis/6896sp11/lec5s.pdf for + # another analysis (where the keys are generated one bit at a time). + exponent = 3 # see tjablin@'s analysis for explanation of this parameter + uint32max = jnp.iinfo(np.uint32).max + num_rounds = int(np.ceil(exponent * np.log(max(1, x.size)) / np.log(uint32max))) + + for _ in range(num_rounds): + key, subkey = _split(key) + sort_keys = _random_bits(subkey, 32, x.shape) + _, x = lax.sort_key_val(sort_keys, x, axis) + + return x + + +def choice(key: KeyArrayLike, + a: int | ArrayLike, + shape: Shape = (), + replace: bool = True, + p: RealArray | None = None, + axis: int = 0) -> Array: + """Generates a random sample from a given array. + + .. warning:: + If ``p`` has fewer non-zero elements than the requested number of samples, + as specified in ``shape``, and ``replace=False``, the output of this + function is ill-defined. Please make sure to use appropriate inputs. + + Args: + key: a PRNG key used as the random key. + a : array or int. If an ndarray, a random sample is generated from + its elements. If an int, the random sample is generated as if a were + arange(a). + shape : tuple of ints, optional. Output shape. If the given shape is, + e.g., ``(m, n)``, then ``m * n`` samples are drawn. Default is (), + in which case a single value is returned. + replace : boolean. Whether the sample is with or without replacement. + Default is True. + p : 1-D array-like, The probabilities associated with each entry in a. + If not given the sample assumes a uniform distribution over all + entries in a. + axis: int, optional. The axis along which the selection is performed. + The default, 0, selects by row. + + Returns: + An array of shape `shape` containing samples from `a`. + """ + key, _ = _check_prng_key("choice", key) + if not isinstance(shape, Sequence): + raise TypeError("shape argument of jax.random.choice must be a sequence, " + f"got {shape}") + check_arraylike("choice", a) + arr = jnp.asarray(a) + if arr.ndim == 0: + n_inputs = core.concrete_or_error(int, a, "The error occurred in jax.random.choice()") + else: + axis = canonicalize_axis(axis, arr.ndim) + n_inputs = arr.shape[axis] + n_draws = math.prod(shape) + if n_draws == 0: + return jnp.zeros(shape, dtype=arr.dtype) + if n_inputs <= 0: + raise ValueError("a must be greater than 0 unless no samples are taken") + if not replace and n_draws > n_inputs: + raise ValueError("Cannot take a larger sample than population when 'replace=False'") + + if p is None: + if replace: + ind = randint(key, shape, 0, n_inputs) + result = ind if arr.ndim == 0 else jnp.take(arr, ind, axis) + else: + slices = (slice(None),) * axis + (slice(n_draws),) + result = permutation(key, n_inputs if arr.ndim == 0 else arr, axis)[slices] + else: + check_arraylike("choice", p) + p_arr, = promote_dtypes_inexact(p) + if p_arr.shape != (n_inputs,): + raise ValueError("p must be None or match the shape of a") + if replace: + p_cuml = jnp.cumsum(p_arr) + r = p_cuml[-1] * (1 - uniform(key, shape, dtype=p_cuml.dtype)) + ind = jnp.searchsorted(p_cuml, r).astype(int) + else: + # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/ + g = -gumbel(key, (n_inputs,), dtype=p_arr.dtype) - jnp.log(p_arr) + ind = jnp.argsort(g)[:n_draws] + result = ind if arr.ndim == 0 else jnp.take(arr, ind, axis) + + return result.reshape(shape if arr.ndim == 0 else + np.insert(np.delete(arr.shape, axis), axis, shape)) + + +def normal(key: KeyArrayLike, + shape: Shape | NamedShape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample standard normal random values with given shape and float dtype. + + The values are returned according to the probability density function: + + .. math:: + f(x) = \frac{1}{\sqrt{2\pi}}e^{-x^2/2} + + on the domain :math:`-\infty < x < \infty` + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("normal", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.inexact): + raise ValueError(f"dtype argument to `normal` must be a float or complex dtype, " + f"got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.as_named_shape(shape) + return _normal(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _normal(key, shape, dtype) -> Array: + if dtypes.issubdtype(dtype, np.complexfloating): + sqrt2 = np.array(np.sqrt(2), dtype) + + key_re, key_im = _split(key) + real_dtype = np.array(0, dtype).real.dtype + _re = _normal_real(key_re, shape, real_dtype).astype(dtype) + _im = _normal_real(key_im, shape, real_dtype).astype(dtype) + return (_re + 1j * _im) / sqrt2 + else: + return _normal_real(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _normal_real(key, shape, dtype) -> Array: + _check_shape("normal", shape) + lo = np.nextafter(np.array(-1., dtype), np.array(0., dtype), dtype=dtype) + hi = np.array(1., dtype) + u = uniform(key, shape, dtype, lo, hi) + return lax.mul(np.array(np.sqrt(2), dtype), lax.erf_inv(u)) + + +def multivariate_normal(key: KeyArrayLike, + mean: RealArray, + cov: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat | None = None, + method: str = 'cholesky') -> Array: + r"""Sample multivariate normal random values with given mean and covariance. + + The values are returned according to the probability density function: + + .. math:: + f(x;\mu, \Sigma) = (2\pi)^{-k/2} \det(\Sigma)^{-1}e^{-\frac{1}{2}(x - \mu)^T \Sigma^{-1} (x - \mu)} + + where :math:`k` is the dimension, :math:`\mu` is the mean (given by ``mean``) and + :math:`\Sigma` is the covariance matrix (given by ``cov``). + + Args: + key: a PRNG key used as the random key. + mean: a mean vector of shape ``(..., n)``. + cov: a positive definite covariance matrix of shape ``(..., n, n)``. The + batch shape ``...`` must be broadcast-compatible with that of ``mean``. + shape: optional, a tuple of nonnegative integers specifying the result + batch shape; that is, the prefix of the result shape excluding the last + axis. Must be broadcast-compatible with ``mean.shape[:-1]`` and + ``cov.shape[:-2]``. The default (None) produces a result batch shape by + broadcasting together the batch shapes of ``mean`` and ``cov``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + method: optional, a method to compute the factor of ``cov``. + Must be one of 'svd', 'eigh', and 'cholesky'. Default 'cholesky'. For + singular covariance matrices, use 'svd' or 'eigh'. + Returns: + A random array with the specified dtype and shape given by + ``shape + mean.shape[-1:]`` if ``shape`` is not None, or else + ``broadcast_shapes(mean.shape[:-1], cov.shape[:-2]) + mean.shape[-1:]``. + """ + key, _ = _check_prng_key("multivariate_normal", key) + dtypes.check_user_dtype_supported(dtype) + mean, cov = promote_dtypes_inexact(mean, cov) + if method not in {'svd', 'eigh', 'cholesky'}: + raise ValueError("method must be one of {'svd', 'eigh', 'cholesky'}") + if dtype is None: + dtype = mean.dtype + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `multivariate_normal` must be a float " + f"dtype, got {dtype}") + if shape is not None: + shape = core.canonicalize_shape(shape) + return _multivariate_normal(key, mean, cov, shape, dtype, method) + +@partial(jit, static_argnums=(3, 4, 5)) +def _multivariate_normal(key, mean, cov, shape, dtype, method) -> Array: + if not np.ndim(mean) >= 1: + msg = "multivariate_normal requires mean.ndim >= 1, got mean.ndim == {}" + raise ValueError(msg.format(np.ndim(mean))) + if not np.ndim(cov) >= 2: + msg = "multivariate_normal requires cov.ndim >= 2, got cov.ndim == {}" + raise ValueError(msg.format(np.ndim(cov))) + n = mean.shape[-1] + if np.shape(cov)[-2:] != (n, n): + msg = ("multivariate_normal requires cov.shape == (..., n, n) for n={n}, " + "but got cov.shape == {shape}.") + raise ValueError(msg.format(n=n, shape=np.shape(cov))) + + if shape is None: + shape = lax.broadcast_shapes(mean.shape[:-1], cov.shape[:-2]) + else: + _check_shape("normal", shape, mean.shape[:-1], cov.shape[:-2]) + + if method == 'svd': + (u, s, _) = svd(cov) + factor = u * jnp.sqrt(s[..., None, :]) + elif method == 'eigh': + (w, v) = eigh(cov) + factor = v * jnp.sqrt(w[..., None, :]) + else: # 'cholesky' + factor = cholesky(cov) + normal_samples = normal(key, shape + mean.shape[-1:], dtype) + with config.numpy_rank_promotion('allow'): + result = mean + jnp.einsum('...ij,...j->...i', factor, normal_samples) + return result + + +def truncated_normal(key: KeyArrayLike, + lower: RealArray, + upper: RealArray, + shape: Shape | NamedShape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample truncated standard normal random values with given shape and dtype. + + The values are returned according to the probability density function: + + .. math:: + f(x) \propto e^{-x^2/2} + + on the domain :math:`\rm{lower} < x < \rm{upper}`. + + Args: + key: a PRNG key used as the random key. + lower: a float or array of floats representing the lower bound for + truncation. Must be broadcast-compatible with ``upper``. + upper: a float or array of floats representing the upper bound for + truncation. Must be broadcast-compatible with ``lower``. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``lower`` and ``upper``. The + default (None) produces a result shape by broadcasting ``lower`` and + ``upper``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and shape given by ``shape`` if + ``shape`` is not None, or else by broadcasting ``lower`` and ``upper``. + Returns values in the open interval ``(lower, upper)``. + """ + key, _ = _check_prng_key("truncated_normal", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `truncated_normal` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.as_named_shape(shape) + return _truncated_normal(key, lower, upper, shape, dtype) + +@partial(jit, static_argnums=(3, 4)) +def _truncated_normal(key, lower, upper, shape, dtype) -> Array: + if shape is None: + shape = lax.broadcast_shapes(np.shape(lower), np.shape(upper)) + else: + _check_shape("truncated_normal", shape, np.shape(lower), np.shape(upper)) + + sqrt2 = np.array(np.sqrt(2), dtype) + lower = lax.convert_element_type(lower, dtype) + upper = lax.convert_element_type(upper, dtype) + a = lax.erf(lower / sqrt2) + b = lax.erf(upper / sqrt2) + if not jnp.issubdtype(dtype, np.floating): + raise TypeError("truncated_normal only accepts floating point dtypes.") + u = uniform(key, shape, dtype, minval=a, maxval=b) + out = sqrt2 * lax.erf_inv(u) + # Clamp the value to the open interval (lower, upper) to make sure that + # rounding (or if we chose `a` for `u`) doesn't push us outside of the range. + return jnp.clip( + out, + lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)), + lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype))) + + +def bernoulli(key: KeyArrayLike, + p: RealArray = np.float32(0.5), + shape: Shape | NamedShape | None = None) -> Array: + r"""Sample Bernoulli random values with given shape and mean. + + The values are distributed according to the probability mass function: + + .. math:: + f(k; p) = p^k(1 - p)^{1 - k} + + where :math:`k \in \{0, 1\}` and :math:`0 \le p \le 1`. + + Args: + key: a PRNG key used as the random key. + p: optional, a float or array of floats for the mean of the random + variables. Must be broadcast-compatible with ``shape``. Default 0.5. + shape: optional, a tuple of nonnegative integers representing the result + shape. Must be broadcast-compatible with ``p.shape``. The default (None) + produces a result shape equal to ``p.shape``. + + Returns: + A random array with boolean dtype and shape given by ``shape`` if ``shape`` + is not None, or else ``p.shape``. + """ + key, _ = _check_prng_key("bernoulli", key) + dtype = dtypes.canonicalize_dtype(lax.dtype(p)) + if shape is not None: + shape = core.as_named_shape(shape) + if not jnp.issubdtype(dtype, np.floating): + msg = "bernoulli probability `p` must have a floating dtype, got {}." + raise TypeError(msg.format(dtype)) + p = lax.convert_element_type(p, dtype) + return _bernoulli(key, p, shape) + +@partial(jit, static_argnums=(2,)) +def _bernoulli(key, p, shape) -> Array: + if shape is None: + # TODO: Use the named part of `p` as well + shape = np.shape(p) + else: + _check_shape("bernoulli", shape, np.shape(p)) + + return uniform(key, shape, lax.dtype(p)) < p + + +def beta(key: KeyArrayLike, + a: RealArray, + b: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Beta random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x;a,b) \propto x^{a - 1}(1 - x)^{b - 1} + + on the domain :math:`0 \le x \le 1`. + + Args: + key: a PRNG key used as the random key. + a: a float or array of floats broadcast-compatible with ``shape`` + representing the first parameter "alpha". + b: a float or array of floats broadcast-compatible with ``shape`` + representing the second parameter "beta". + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``a`` and ``b``. The default + (None) produces a result shape by broadcasting ``a`` and ``b``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and shape given by ``shape`` if + ``shape`` is not None, or else by broadcasting ``a`` and ``b``. + """ + key, _ = _check_prng_key("beta", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `beta` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _beta(key, a, b, shape, dtype) + + +def _beta(key, a, b, shape, dtype) -> Array: + if shape is None: + shape = lax.broadcast_shapes(np.shape(a), np.shape(b)) + else: + _check_shape("beta", shape, np.shape(a), np.shape(b)) + + a = lax.convert_element_type(a, dtype) + b = lax.convert_element_type(b, dtype) + key_a, key_b = _split(key) + a = jnp.broadcast_to(a, shape) + b = jnp.broadcast_to(b, shape) + log_gamma_a = loggamma(key_a, a, shape, dtype) + log_gamma_b = loggamma(key_b, b, shape, dtype) + # Compute gamma_a / (gamma_a + gamma_b) without losing precision. + log_max = lax.max(log_gamma_a, log_gamma_b) + gamma_a_scaled = jnp.exp(log_gamma_a - log_max) + gamma_b_scaled = jnp.exp(log_gamma_b - log_max) + return gamma_a_scaled / (gamma_a_scaled + gamma_b_scaled) + + +def cauchy(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Cauchy random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x) \propto \frac{1}{x^2 + 1} + + on the domain :math:`-\infty < x < \infty` + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("cauchy", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `cauchy` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _cauchy(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _cauchy(key, shape, dtype) -> Array: + _check_shape("cauchy", shape) + u = uniform(key, shape, dtype, minval=jnp.finfo(dtype).eps, maxval=1.) + pi = _lax_const(u, np.pi) + return lax.tan(lax.mul(pi, lax.sub(u, _lax_const(u, 0.5)))) + + +def dirichlet(key: KeyArrayLike, + alpha: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Dirichlet random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(\{x_i\}; \{\alpha_i\}) \propto \prod_{i=1}^k x_i^{\alpha_i - 1} + + Where :math:`k` is the dimension, and :math:`\{x_i\}` satisfies + + .. math:: + \sum_{i=1}^k x_i = 1 + + and :math:`0 \le x_i \le 1` for all :math:`x_i`. + + Args: + key: a PRNG key used as the random key. + alpha: an array of shape ``(..., n)`` used as the concentration + parameter of the random variables. + shape: optional, a tuple of nonnegative integers specifying the result + batch shape; that is, the prefix of the result shape excluding the last + element of value ``n``. Must be broadcast-compatible with + ``alpha.shape[:-1]``. The default (None) produces a result shape equal to + ``alpha.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and shape given by + ``shape + (alpha.shape[-1],)`` if ``shape`` is not None, or else + ``alpha.shape``. + """ + key, _ = _check_prng_key("dirichlet", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `dirichlet` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _dirichlet(key, alpha, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _dirichlet(key, alpha, shape, dtype) -> Array: + if not np.ndim(alpha) >= 1: + msg = "dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}" + raise ValueError(msg.format(np.ndim(alpha))) + + if shape is None: + shape = np.shape(alpha)[:-1] + else: + _check_shape("dirichlet", shape, np.shape(alpha)[:-1]) + + alpha = lax.convert_element_type(alpha, dtype) + + # Compute gamma in log space, otherwise small alpha can lead to poor behavior. + log_gamma_samples = loggamma(key, alpha, shape + np.shape(alpha)[-1:], dtype) + return _softmax(log_gamma_samples, -1) + + +def _softmax(x, axis) -> Array: + """Utility to compute the softmax of x along a given axis.""" + if not dtypes.issubdtype(x.dtype, np.floating): + raise TypeError(f"_softmax only accepts floating dtypes, got {x.dtype}") + x_max = jnp.max(x, axis, keepdims=True) + unnormalized = jnp.exp(x - lax.stop_gradient(x_max)) + return unnormalized / unnormalized.sum(axis, keepdims=True) + + +def exponential(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Exponential random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x) = e^{-x} + + on the domain :math:`0 \le x < \infty`. + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("exponential", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `exponential` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _exponential(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _exponential(key, shape, dtype) -> Array: + _check_shape("exponential", shape) + u = uniform(key, shape, dtype) + # taking 1 - u to move the domain of log to (0, 1] instead of [0, 1) + return lax.neg(lax.log1p(lax.neg(u))) + + +def _gamma_one(key: KeyArray, alpha, log_space) -> Array: + # Ref: A simple method for generating gamma variables, George Marsaglia and Wai Wan Tsang + # The algorithm can also be founded in: + # https://en.wikipedia.org/wiki/Gamma_distribution#Generating_gamma-distributed_random_variables + zero = _lax_const(alpha, 0) + one = _lax_const(alpha, 1) + minus_one = _lax_const(alpha, -1) + one_over_two = _lax_const(alpha, 0.5) + one_over_three = _lax_const(alpha, 1. / 3.) + squeeze_const = _lax_const(alpha, 0.0331) + dtype = lax.dtype(alpha) + + # for alpha < 1, we boost alpha to alpha + 1 and get a sample according to + # Gamma(alpha) ~ Gamma(alpha+1) * Uniform()^(1 / alpha) + # When alpha is very small, this boost can be problematic because it may result + # in floating point underflow; for this reason we compute it in log space if + # specified by the `log_space` argument: + # log[Gamma(alpha)] ~ log[Gamma(alpha + 1)] + log[Uniform()] / alpha + # Note that log[Uniform()] ~ -Exponential(), but to avoid problems at x=0 + # exponential is computed in terms of log[1 - Uniform()]; we must account for this + # so that log-space and non-log-space samples match. + boost_mask = lax.ge(alpha, one) + alpha_orig = alpha + alpha = lax.select(boost_mask, alpha, lax.add(alpha, one)) + + d = lax.sub(alpha, one_over_three) + c = lax.div(one_over_three, lax.sqrt(d)) + + def _cond_fn(kXVU): + _, X, V, U = kXVU + # TODO: use lax.cond when its batching rule is supported + # The reason is to avoid evaluating second condition which involves log+log + # if the first condition is satisfied + cond = lax.bitwise_and(lax.ge(U, lax.sub(one, lax.mul(squeeze_const, lax.mul(X, X)))), + lax.ge(lax.log(U), lax.add(lax.mul(X, one_over_two), + lax.mul(d, lax.add(lax.sub(one, V), + lax.log(V)))))) + return cond + + def _body_fn(kXVU): + def _next_kxv(kxv): + key = kxv[0] + key, subkey = _split(key) + x = normal(subkey, (), dtype=dtype) + v = lax.add(one, lax.mul(x, c)) + return key, x, v + + key = kXVU[0] + key, x_key, U_key = _split(key, 3) + _, x, v = lax.while_loop(lambda kxv: lax.le(kxv[2], zero), _next_kxv, (x_key, zero, minus_one)) + X = lax.mul(x, x) + V = lax.mul(lax.mul(v, v), v) + U = uniform(U_key, (), dtype=dtype) + return key, X, V, U + + # initial state is chosen such that _cond_fn will return True + key, subkey = _split(key) + _, _, V, _ = lax.while_loop(_cond_fn, _body_fn, (key, zero, one, _lax_const(alpha, 2))) + if log_space: + log_samples = lax.neg(exponential(subkey, (), dtype=dtype)) + log_boost = lax.select(boost_mask | (log_samples == 0), zero, lax.mul(log_samples, lax.div(one, alpha_orig))) + return lax.add(lax.add(lax.log(d), lax.log(V)), log_boost) + else: + samples = 1 - uniform(subkey, (), dtype=dtype) + boost = lax.select(boost_mask, one, lax.pow(samples, lax.div(one, alpha_orig))) + return lax.mul(lax.mul(d, V), boost) + + +def _gamma_grad(sample, a, *, log_space): + samples = jnp.reshape(sample, -1) + alphas = jnp.reshape(a, -1) + if log_space: + # d[log(sample)] = d[sample] / sample + # This requires computing exp(log_sample), which may be zero due to float roundoff. + # In this case, correct it to smallest representable float. + samples = lax.exp(samples) + zero = lax_internal._const(sample, 0) + tiny = lax.full_like(samples, jnp.finfo(samples.dtype).tiny) + samples = lax.select(lax.eq(samples, zero), tiny, samples) + gamma_grad = lambda alpha, sample: lax.random_gamma_grad(alpha, sample) / sample + else: + gamma_grad = lax.random_gamma_grad + if xla_bridge.get_backend().platform == 'cpu': + grads = lax.map(lambda args: gamma_grad(*args), (alphas, samples)) + else: + grads = vmap(gamma_grad)(alphas, samples) + return grads.reshape(np.shape(a)) + +def _gamma_impl(key, a, *, log_space, use_vmap=False): + # split key to match the shape of a + a_shape = jnp.shape(a) + split_count = math.prod(a_shape[key.ndim:]) + keys = key.flatten() + keys = vmap(_split, in_axes=(0, None))(keys, split_count) + keys = keys.flatten() + alphas = a.flatten() + + if use_vmap and _key_impl(key) is prng.threefry_prng_impl: + samples = vmap(partial(_gamma_one, log_space=log_space))(keys, alphas) + else: + samples = lax.map( + lambda args: _gamma_one(*args, log_space=log_space), (keys, alphas)) + + return jnp.reshape(samples, a_shape) + +def _gamma_batching_rule(batched_args, batch_dims, *, log_space): + k, a = batched_args + bk, ba = batch_dims + size = next( + t.shape[i] for t, i in zip(batched_args, batch_dims) if i is not None) + k = batching.bdim_at_front(k, bk, size) + a = batching.bdim_at_front(a, ba, size) + return random_gamma_p.bind(k, a, log_space=log_space), 0 + +random_gamma_p = core.Primitive('random_gamma') +random_gamma_p.def_impl(_gamma_impl) +random_gamma_p.def_abstract_eval(lambda key, a, **_: core.raise_to_shaped(a)) +ad.defjvp2( + random_gamma_p, None, + lambda tangent, ans, key, a, **kwds: tangent * _gamma_grad(ans, a, **kwds)) +mlir.register_lowering(random_gamma_p, mlir.lower_fun( + partial(_gamma_impl, use_vmap=True), + multiple_results=False)) +mlir.register_lowering(random_gamma_p, mlir.lower_fun( + partial(_gamma_impl, use_vmap=False), + multiple_results=False), platform='cpu') +batching.primitive_batchers[random_gamma_p] = _gamma_batching_rule + +def gamma(key: KeyArrayLike, + a: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Gamma random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x;a) \propto x^{a - 1} e^{-x} + + on the domain :math:`0 \le x < \infty`, with :math:`a > 0`. + + This is the standard gamma density, with a unit scale/rate parameter. + Dividing the sample output by the rate is equivalent to sampling from + *gamma(a, rate)*, and multiplying the sample output by the scale is equivalent + to sampling from *gamma(a, scale)*. + + Args: + key: a PRNG key used as the random key. + a: a float or array of floats broadcast-compatible with ``shape`` + representing the parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``a``. The default (None) + produces a result shape equal to ``a.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``a.shape``. + + See Also: + loggamma : sample gamma values in log-space, which can provide improved + accuracy for small values of ``a``. + """ + key, _ = _check_prng_key("gamma", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `gamma` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _gamma(key, a, shape=shape, dtype=dtype) + + +def loggamma(key: KeyArrayLike, + a: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + """Sample log-gamma random values with given shape and float dtype. + + This function is implemented such that the following will hold for a + dtype-appropriate tolerance:: + + np.testing.assert_allclose(jnp.exp(loggamma(*args)), gamma(*args), rtol=rtol) + + The benefit of log-gamma is that for samples very close to zero (which occur frequently + when `a << 1`) sampling in log space provides better precision. + + Args: + key: a PRNG key used as the random key. + a: a float or array of floats broadcast-compatible with ``shape`` + representing the parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``a``. The default (None) + produces a result shape equal to ``a.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``a.shape``. + + See Also: + gamma : standard gamma sampler. + """ + key, _ = _check_prng_key("loggamma", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `gamma` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _gamma(key, a, shape=shape, dtype=dtype, log_space=True) + + +@partial(jit, static_argnames=('shape', 'dtype', 'log_space')) +def _gamma(key, a, shape, dtype, log_space=False) -> Array: + if shape is None: + shape = np.shape(a) + else: + _check_shape("gamma", shape, np.shape(a)) + + a = lax.convert_element_type(a, dtype) + if np.shape(a) != shape: + a = jnp.broadcast_to(a, shape) + return random_gamma_p.bind(key, a, log_space=log_space) + + +@partial(jit, static_argnums=(2, 3, 4)) +def _poisson_knuth(key, lam, shape, dtype, max_iters) -> Array: + # Knuth's algorithm for generating Poisson random variates. + # Reference: + # https://en.wikipedia.org/wiki/Poisson_distribution#Generating_Poisson-distributed_random_variables + + def body_fn(carry): + i, k, rng, log_prod = carry + rng, subkey = _split(rng) + k = lax.select(log_prod > -lam, k + 1, k) + u = uniform(subkey, shape, np.float32) + return i + 1, k, rng, log_prod + jnp.log(u) + + def cond_fn(carry): + i, log_prod = carry[0], carry[3] + return (log_prod > -lam).any() & (i < max_iters) + + k_init = lax.full_like(lam, 0, dtype, shape) + log_rate_init = lax.full_like(lam, 0, np.float32, shape) + k = lax.while_loop(cond_fn, body_fn, (0, k_init, key, log_rate_init))[1] + return (k - 1).astype(dtype) + + +@partial(jit, static_argnums=(2, 3, 4)) +def _poisson_rejection(key, lam, shape, dtype, max_iters) -> Array: + # Transformed rejection due to Hormann. + # Reference: + # http://citeseer.ist.psu.edu/viewdoc/citations;jsessionid=1BEB35946CC807879F55D42512E5490C?doi=10.1.1.48.3054. + log_lam = lax.log(lam) + b = 0.931 + 2.53 * lax.sqrt(lam) + a = -0.059 + 0.02483 * b + inv_alpha = 1.1239 + 1.1328 / (b - 3.4) + v_r = 0.9277 - 3.6224 / (b - 2) + + def body_fn(carry): + i, k_out, accepted, key = carry + key, subkey_0, subkey_1 = _split(key, 3) + + u = uniform(subkey_0, shape, lam.dtype) - 0.5 + v = uniform(subkey_1, shape, lam.dtype) + u_shifted = 0.5 - abs(u) + + k = lax.floor((2 * a / u_shifted + b) * u + lam + 0.43) + s = lax.log(v * inv_alpha / (a / (u_shifted * u_shifted) + b)) + t = -lam + k * log_lam - lax.lgamma(k + 1) + + accept1 = (u_shifted >= 0.07) & (v <= v_r) + reject = (k < 0) | ((u_shifted < 0.013) & (v > u_shifted)) + accept2 = s <= t + accept = accept1 | (~reject & accept2) + + k_out = lax.select(accept, k, k_out) + accepted |= accept + + return i + 1, k_out, accepted, key + + def cond_fn(carry): + i, k_out, accepted, key = carry + return (~accepted).any() & (i < max_iters) + + k_init = lax.full_like(lam, -1, lam.dtype, shape) + accepted = lax.full_like(lam, False, jnp.bool_, shape) + k = lax.while_loop(cond_fn, body_fn, (0, k_init, accepted, key))[1] + return k.astype(dtype) + + +@partial(jit, static_argnums=(2, 3)) +def _poisson(key, lam, shape, dtype) -> Array: + # The implementation matches TensorFlow and NumPy: + # https://github.com/tensorflow/tensorflow/blob/v2.2.0-rc3/tensorflow/core/kernels/random_poisson_op.cc + # https://github.com/numpy/numpy/blob/v1.18.3/numpy/random/src/distributions/distributions.c#L574 + # For lambda < 10, we use the Knuth algorithm; otherwise, we use transformed + # rejection sampling. + use_knuth = _isnan(lam) | (lam < 10) + lam_knuth = lax.select(use_knuth, lam, lax.full_like(lam, 0.0)) + # The acceptance probability for rejection sampling maxes out at 89% as + # λ -> ∞, so pick some arbitrary large value. + lam_rejection = lax.select(use_knuth, lax.full_like(lam, 1e5), lam) + max_iters = dtype.type(jnp.iinfo(dtype).max) # insanely conservative + result = lax.select( + use_knuth, + _poisson_knuth(key, lam_knuth, shape, dtype, max_iters), + _poisson_rejection(key, lam_rejection, shape, dtype, max_iters), + ) + return lax.select(lam == 0, jnp.zeros_like(result), result) + + +def poisson(key: KeyArrayLike, + lam: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeInt = int) -> Array: + r"""Sample Poisson random values with given shape and integer dtype. + + The values are distributed according to the probability mass function: + + .. math:: + f(k; \lambda) = \frac{\lambda^k e^{-\lambda}}{k!} + + Where `k` is a non-negative integer and :math:`\lambda > 0`. + + Args: + key: a PRNG key used as the random key. + lam: rate parameter (mean of the distribution), must be >= 0. Must be broadcast-compatible with ``shape`` + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (None) produces a result shape equal to ``lam.shape``. + dtype: optional, a integer dtype for the returned values (default int64 if + jax_enable_x64 is true, otherwise int32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape is not None, or else by ``lam.shape``. + """ + key, _ = _check_prng_key("poisson", key) + dtypes.check_user_dtype_supported(dtype) + # TODO(frostig): generalize underlying poisson implementation and + # remove this check + keys_dtype = typing.cast(prng.KeyTy, key.dtype) + key_impl = keys_dtype._impl + if key_impl is not prng.threefry_prng_impl: + raise NotImplementedError( + '`poisson` is only implemented for the threefry2x32 RNG, ' + f'not {key_impl}') + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + else: + shape = np.shape(lam) + lam = jnp.broadcast_to(lam, shape) + lam = lax.convert_element_type(lam, np.float32) + return _poisson(key, lam, shape, dtype) + + +def gumbel(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + """Sample Gumbel random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x) = e^{-(x + e^{-x})} + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("gumbel", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `gumbel` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _gumbel(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _gumbel(key, shape, dtype) -> Array: + _check_shape("gumbel", shape) + return -jnp.log(-jnp.log( + uniform(key, shape, dtype, minval=jnp.finfo(dtype).tiny, maxval=1.))) + + +def categorical(key: KeyArrayLike, + logits: RealArray, + axis: int = -1, + shape: Shape | None = None) -> Array: + """Sample random values from categorical distributions. + + Args: + key: a PRNG key used as the random key. + logits: Unnormalized log probabilities of the categorical distribution(s) to sample from, + so that `softmax(logits, axis)` gives the corresponding probabilities. + axis: Axis along which logits belong to the same categorical distribution. + shape: Optional, a tuple of nonnegative integers representing the result shape. + Must be broadcast-compatible with ``np.delete(logits.shape, axis)``. + The default (None) produces a result shape equal to ``np.delete(logits.shape, axis)``. + + Returns: + A random array with int dtype and shape given by ``shape`` if ``shape`` + is not None, or else ``np.delete(logits.shape, axis)``. + """ + key, _ = _check_prng_key("categorical", key) + check_arraylike("categorical", logits) + logits_arr = jnp.asarray(logits) + + if axis >= 0: + axis -= len(logits_arr.shape) + + batch_shape = tuple(np.delete(logits_arr.shape, axis)) + if shape is None: + shape = batch_shape + else: + shape = tuple(shape) + _check_shape("categorical", shape, batch_shape) + + shape_prefix = shape[:len(shape)-len(batch_shape)] + logits_shape = list(shape[len(shape) - len(batch_shape):]) + logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis]) + return jnp.argmax( + gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype) + + lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))), + axis=axis) + + +def laplace(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Laplace random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x) = \frac{1}{2}e^{-|x|} + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("laplace", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `laplace` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _laplace(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _laplace(key, shape, dtype) -> Array: + _check_shape("laplace", shape) + u = uniform( + key, shape, dtype, minval=-1. + jnp.finfo(dtype).epsneg, maxval=1.) + return lax.mul(lax.sign(u), lax.log1p(lax.neg(lax.abs(u)))) + + +def logistic(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample logistic random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x) = \frac{e^{-x}}{(1 + e^{-x})^2} + + Args: + key: a PRNG key used as the random key. + shape: optional, a tuple of nonnegative integers representing the result + shape. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("logistic", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `logistic` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _logistic(key, shape, dtype) + +@partial(jit, static_argnums=(1, 2)) +def _logistic(key, shape, dtype): + _check_shape("logistic", shape) + x = uniform(key, shape, dtype, minval=jnp.finfo(dtype).eps, maxval=1.) + return lax.log(lax.div(x, lax.sub(_lax_const(x, 1), x))) + + +def pareto(key: KeyArrayLike, + b: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Pareto random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x; b) = b / x^{b + 1} + + on the domain :math:`1 \le x < \infty` with :math:`b > 0` + + Args: + key: a PRNG key used as the random key. + b: a float or array of floats broadcast-compatible with ``shape`` + representing the parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``b``. The default (None) + produces a result shape equal to ``b.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``b.shape``. + """ + key, _ = _check_prng_key("pareto", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `pareto` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _pareto(key, b, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _pareto(key, b, shape, dtype) -> Array: + if shape is None: + shape = np.shape(b) + else: + _check_shape("pareto", shape) + + b = lax.convert_element_type(b, dtype) + e = exponential(key, shape, dtype) + return lax.exp(e / b) + + +def t(key: KeyArrayLike, + df: RealArray, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Student's t random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(t; \nu) \propto \left(1 + \frac{t^2}{\nu}\right)^{-(\nu + 1)/2} + + Where :math:`\nu > 0` is the degrees of freedom, given by the parameter ``df``. + + Args: + key: a PRNG key used as the random key. + df: a float or array of floats broadcast-compatible with ``shape`` + representing the degrees of freedom parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``df``. The default (None) + produces a result shape equal to ``df.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``df.shape``. + """ + key, _ = _check_prng_key("t", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `t` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _t(key, df, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _t(key, df, shape, dtype) -> Array: + if shape is None: + shape = np.shape(df) + else: + _check_shape("t", shape, np.shape(df)) + + df = lax.convert_element_type(df, dtype) + key_n, key_g = _split(key) + n = normal(key_n, shape, dtype) + two = _lax_const(n, 2) + half_df = lax.div(df, two) + g = gamma(key_g, half_df, shape, dtype) + return n * jnp.sqrt(half_df / g) + + +def chisquare(key: KeyArrayLike, + df: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Chisquare random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x; \nu) \propto x^{\nu/2 - 1}e^{-x/2} + + on the domain :math:`0 < x < \infty`, where :math:`\nu > 0` represents the + degrees of freedom, given by the parameter ``df``. + + Args: + key: a PRNG key used as the random key. + df: a float or array of floats broadcast-compatible with ``shape`` + representing the parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``df``. The default (None) + produces a result shape equal to ``df.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``df.shape``. + """ + key, _ = _check_prng_key("chisquare", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError("dtype argument to `chisquare` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _chisquare(key, df, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _chisquare(key, df, shape, dtype) -> Array: + if shape is None: + shape = np.shape(df) + else: + _check_shape("chisquare", shape, np.shape(df)) + df = lax.convert_element_type(df, dtype) + two = _lax_const(df, 2) + half_df = lax.div(df, two) + log_g = loggamma(key, a=half_df, shape=shape, dtype=dtype) + chi2 = lax.mul(jnp.exp(log_g), two) + return chi2 + + +def f(key: KeyArrayLike, + dfnum: RealArray, + dfden: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample F-distribution random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x; \nu_1, \nu_2) \propto x^{\nu_1/2 - 1}\left(1 + \frac{\nu_1}{\nu_2}x\right)^{ + -(\nu_1 + \nu_2) / 2} + + on the domain :math:`0 < x < \infty`. Here :math:`\nu_1` is the degrees of + freedom of the numerator (``dfnum``), and :math:`\nu_2` is the degrees of + freedom of the denominator (``dfden``). + + Args: + key: a PRNG key used as the random key. + dfnum: a float or array of floats broadcast-compatible with ``shape`` + representing the numerator's ``df`` of the distribution. + dfden: a float or array of floats broadcast-compatible with ``shape`` + representing the denominator's ``df`` of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``dfnum`` and ``dfden``. + The default (None) produces a result shape equal to ``dfnum.shape``, + and ``dfden.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``df.shape``. + """ + key, _ = _check_prng_key("f", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError("dtype argument to `f` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _f(key, dfnum, dfden, shape, dtype) + +@partial(jit, static_argnums=(3, 4)) +def _f(key, dfnum, dfden, shape, dtype) -> Array: + if shape is None: + shape = lax.broadcast_shapes(np.shape(dfden), np.shape(dfnum)) + else: + _check_shape("f", shape, np.shape(dfden), np.shape(dfnum)) + dfden = lax.convert_element_type(dfden, dtype) + dfnum = lax.convert_element_type(dfnum, dtype) + key_dfd, key_dfn = _split(key) + chi2_dfn = chisquare(key_dfn, dfnum, shape, dtype) + chi2_dfd = chisquare(key_dfd, dfden, shape, dtype) + # broadcast dfden and dfnum to do div operation + dfden = jnp.broadcast_to(dfden, shape) + dfnum = jnp.broadcast_to(dfnum, shape) + num = lax.div(chi2_dfn, dfnum) + den = lax.div(chi2_dfd ,dfden) + f = lax.div(num, den) + return f + + +def rademacher(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeInt = int) -> Array: + r"""Sample from a Rademacher distribution. + + The values are distributed according to the probability mass function: + + .. math:: + f(k) = \frac{1}{2}(\delta(k - 1) + \delta(k + 1)) + + on the domain :math:`k \in \{-1, 1\}`, where :math:`\delta(x)` is the dirac delta function. + + Args: + key: a PRNG key. + shape: The shape of the returned samples. Default (). + dtype: The type used for samples. + + Returns: + A jnp.array of samples, of shape `shape`. Each element in the output has + a 50% change of being 1 or -1. + + """ + key, _ = _check_prng_key("rademacher", key) + dtypes.check_user_dtype_supported(dtype) + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _rademacher(key, shape, dtype) + + +@partial(jit, static_argnums=(1, 2)) +def _rademacher(key, shape, dtype) -> Array: + bernoulli_samples = bernoulli(key=key, p=0.5, shape=shape).astype(dtype) + return (2 * bernoulli_samples - 1).astype(dtype) + + +def maxwell(key: KeyArrayLike, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample from a one sided Maxwell distribution. + + The values are distributed according to the probability density function: + + .. math:: + f(x) \propto x^2 e^{-x^2 / 2} + + on the domain :math:`0 \le x < \infty`. + + Args: + key: a PRNG key. + shape: The shape of the returned samples. + dtype: The type used for samples. + + Returns: + A jnp.array of samples, of shape `shape`. + + """ + # Generate samples using: + # sqrt(X^2 + Y^2 + Z^2), X,Y,Z ~N(0,1) + key, _ = _check_prng_key("maxwell", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `maxwell` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _maxwell(key, shape, dtype) + + +@partial(jit, static_argnums=(1, 2)) +def _maxwell(key, shape, dtype) -> Array: + shape = shape + (3,) + norm_rvs = normal(key=key, shape=shape, dtype=dtype) + return jnp.linalg.norm(norm_rvs, axis=-1) + + +def double_sided_maxwell(key: KeyArrayLike, + loc: RealArray, + scale: RealArray, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample from a double sided Maxwell distribution. + + The values are distributed according to the probability density function: + + .. math:: + f(x;\mu,\sigma) \propto z^2 e^{-z^2 / 2} + + where :math:`z = (x - \mu) / \sigma`, with the center :math:`\mu` specified by + ``loc`` and the scale :math:`\sigma` specified by ``scale``. + + Args: + key: a PRNG key. + loc: The location parameter of the distribution. + scale: The scale parameter of the distribution. + shape: The shape added to the parameters loc and scale broadcastable shape. + dtype: The type used for samples. + + Returns: + A jnp.array of samples. + + """ + key, _ = _check_prng_key("double_sided_maxwell", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `double_sided_maxwell` must be a float" + f" dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _double_sided_maxwell(key, loc, scale, shape, dtype) + + +@partial(jit, static_argnums=(3, 4)) +def _double_sided_maxwell(key, loc, scale, shape, dtype) -> Array: + params_shapes = lax.broadcast_shapes(np.shape(loc), np.shape(scale)) + if not shape: + shape = params_shapes + + shape = shape + params_shapes + maxwell_key, rademacher_key = _split(key) + maxwell_rvs = maxwell(maxwell_key, shape=shape, dtype=dtype) + # Generate random signs for the symmetric variates. + random_sign = rademacher(rademacher_key, shape=shape, dtype=dtype) + assert random_sign.shape == maxwell_rvs.shape + + return random_sign * maxwell_rvs * scale + loc + + +def weibull_min(key: KeyArrayLike, + scale: RealArray, + concentration: RealArray, + shape: Shape = (), + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample from a Weibull distribution. + + The values are distributed according to the probability density function: + + .. math:: + f(x;\sigma,c) \propto x^{c - 1} \exp(-(x / \sigma)^c) + + on the domain :math:`0 < x < \infty`, where :math:`c > 0` is the concentration + parameter, and :math:`\sigma > 0` is the scale parameter. + + Args: + key: a PRNG key. + scale: The scale parameter of the distribution. + concentration: The concentration parameter of the distribution. + shape: The shape added to the parameters loc and scale broadcastable shape. + dtype: The type used for samples. + + Returns: + A jnp.array of samples. + + """ + key, _ = _check_prng_key("weibull_min", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError(f"dtype argument to `weibull_min` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + shape = core.canonicalize_shape(shape) + return _weibull_min(key, scale, concentration, shape, dtype) + + +@partial(jit, static_argnums=(3, 4)) +def _weibull_min(key, scale, concentration, shape, dtype) -> Array: + random_uniform = uniform( + key=key, shape=shape, minval=0, maxval=1, dtype=dtype) + + # Inverse weibull CDF. + return jnp.power(-jnp.log1p(-random_uniform), 1.0/concentration) * scale + + +def orthogonal( + key: KeyArrayLike, + n: int, + shape: Shape = (), + dtype: DTypeLikeFloat = float +) -> Array: + """Sample uniformly from the orthogonal group O(n). + + If the dtype is complex, sample uniformly from the unitary group U(n). + + Args: + key: a PRNG key used as the random key. + n: an integer indicating the resulting dimension. + shape: optional, the batch dimensions of the result. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array of shape `(*shape, n, n)` and specified dtype. + """ + key, _ = _check_prng_key("orthogonal", key) + dtypes.check_user_dtype_supported(dtype) + _check_shape("orthogonal", shape) + n = core.concrete_or_error(index, n, "The error occurred in jax.random.orthogonal()") + z = normal(key, (*shape, n, n), dtype) + q, r = jnp.linalg.qr(z) + d = jnp.diagonal(r, 0, -2, -1) + return lax.mul(q, lax.expand_dims(lax.div(d, abs(d).astype(d.dtype)), [-2])) + +def generalized_normal( + key: KeyArrayLike, + p: float, + shape: Shape = (), + dtype: DTypeLikeFloat = float +) -> Array: + r"""Sample from the generalized normal distribution. + + The values are returned according to the probability density function: + + .. math:: + f(x;p) \propto e^{-|x|^p} + + on the domain :math:`-\infty < x < \infty`, where :math:`p > 0` is the + shape parameter. + + Args: + key: a PRNG key used as the random key. + p: a float representing the shape parameter. + shape: optional, the batch dimensions of the result. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified shape and dtype. + """ + key, _ = _check_prng_key("generalized_normal", key) + dtypes.check_user_dtype_supported(dtype) + _check_shape("generalized_normal", shape) + keys = split(key) + g = gamma(keys[0], 1/p, shape, dtype) + r = rademacher(keys[1], shape, dtype) + return r * g ** (1 / p) + +def ball( + key: KeyArrayLike, + d: int, + p: float = 2, + shape: Shape = (), + dtype: DTypeLikeFloat = float +): + """Sample uniformly from the unit Lp ball. + + Reference: https://arxiv.org/abs/math/0503650. + + Args: + key: a PRNG key used as the random key. + d: a nonnegative int representing the dimensionality of the ball. + p: a float representing the p parameter of the Lp norm. + shape: optional, the batch dimensions of the result. Default (). + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array of shape `(*shape, d)` and specified dtype. + """ + key, _ = _check_prng_key("ball", key) + dtypes.check_user_dtype_supported(dtype) + _check_shape("ball", shape) + d = core.concrete_or_error(index, d, "The error occurred in jax.random.ball()") + k1, k2 = split(key) + g = generalized_normal(k1, p, (*shape, d), dtype) + e = exponential(k2, shape, dtype) + return g / (((jnp.abs(g) ** p).sum(-1) + e) ** (1 / p))[..., None] + + +def rayleigh(key: KeyArrayLike, + scale: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Rayleigh random values with given shape and float dtype. + + The values are returned according to the probability density function: + + .. math:: + f(x;\sigma) \propto xe^{-x^2/(2\sigma^2)} + + on the domain :math:`-\infty < x < \infty`, and where :math:`\sigma > 0` is the scale + parameter of the distribution. + + Args: + key: a PRNG key used as the random key. + scale: a float or array of floats broadcast-compatible with ``shape`` + representing the parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``scale``. The default (None) + produces a result shape equal to ``scale.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``scale.shape``. + """ + key, _ = _check_prng_key("rayleigh", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError("dtype argument to `rayleigh` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _rayleigh(key, scale, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _rayleigh(key, scale, shape, dtype) -> Array: + if shape is None: + shape = np.shape(scale) + else: + _check_shape("rayleigh", shape, np.shape(scale)) + u = uniform(key, shape, dtype) + scale = scale.astype(dtype) + scale = jnp.broadcast_to(scale, shape) + log_u = lax.log(u) + n_two = _lax_const(scale, -2) + sqrt_u = lax.sqrt(lax.mul(log_u, n_two)) + ray = lax.mul(scale, sqrt_u) + return ray + +def wald(key: KeyArrayLike, + mean: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Wald random values with given shape and float dtype. + + The values are returned according to the probability density function: + + .. math:: + f(x;\mu) = \frac{1}{\sqrt{2\pi x^3}} \exp\left(-\frac{(x - \mu)^2}{2\mu^2 x}\right) + + on the domain :math:`-\infty < x < \infty`, and where :math:`\mu > 0` is the location + parameter of the distribution. + + + Args: + key: a PRNG key used as the random key. + mean: a float or array of floats broadcast-compatible with ``shape`` + representing the mean parameter of the distribution. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``mean``. The default + (None) produces a result shape equal to ``np.shape(mean)``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``mean.shape``. + """ + key, _ = _check_prng_key("wald", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError("dtype argument to `wald` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _wald(key, mean, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _wald(key, mean, shape, dtype) -> Array: + if shape is None: + shape = np.shape(mean) + else: + _check_shape("wald", shape, np.shape(mean)) + k1, k2 = _split(key, 2) + mean = mean.astype(dtype) + mean = jnp.broadcast_to(mean, shape) + v = normal(k1, shape, dtype) + z = uniform(k2, shape, dtype) + y = lax.integer_pow(v, 2) + y_sq = lax.integer_pow(y, 2) + mean_sq = lax.integer_pow(mean, 2) + sqrt_term = lax.sqrt(4 * mean * y + mean_sq * y_sq) + x = mean + mean_sq * y / 2 - mean / 2 * sqrt_term + w = lax.select(lax.le(z, mean / (mean + x)), x, mean_sq / x) + return w + +def geometric(key: KeyArrayLike, + p: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeInt = int) -> Array: + r"""Sample Geometric random values with given shape and float dtype. + + The values are returned according to the probability mass function: + + .. math:: + f(k;p) = p(1-p)^{k-1} + + on the domain :math:`0 < p < 1`. + + Args: + key: a PRNG key used as the random key. + p: a float or array of floats broadcast-compatible with ``shape`` + representing the probability of success of an individual trial. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``p``. The default + (None) produces a result shape equal to ``np.shape(p)``. + dtype: optional, a int dtype for the returned values (default int64 if + jax_enable_x64 is true, otherwise int32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``p.shape``. + """ + key, _ = _check_prng_key("geometric", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.integer): + raise ValueError("dtype argument to `geometric` must be an int " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _geometric(key, p, shape, dtype) + +@partial(jit, static_argnums=(2, 3)) +def _geometric(key, p, shape, dtype) -> Array: + if shape is None: + shape = np.shape(p) + else: + _check_shape("geometric", shape, np.shape(p)) + check_arraylike("geometric", p) + p, = promote_dtypes_inexact(p) + u = uniform(key, shape, p.dtype) + log_u = lax.log(u) + log_one_minus_p = lax.log1p(-p) + log_one_minus_p = jnp.broadcast_to(log_one_minus_p, shape) + g = lax.floor(lax.div(log_u, log_one_minus_p)) + 1 + return g.astype(dtype) + + +def triangular(key: KeyArrayLike, + left: RealArray, + mode: RealArray, + right: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r"""Sample Triangular random values with given shape and float dtype. + + The values are returned according to the probability density function: + + .. math:: + f(x; a, b, c) = \frac{2}{c-a} \left\{ \begin{array}{ll} \frac{x-a}{b-a} & a \leq x \leq b \\ \frac{c-x}{c-b} & b \leq x \leq c \end{array} \right. + + on the domain :math:`a \leq x \leq c`. + + Args: + key: a PRNG key used as the random key. + left: a float or array of floats broadcast-compatible with ``shape`` + representing the lower limit parameter of the distribution. + mode: a float or array of floats broadcast-compatible with ``shape`` + representing the peak value parameter of the distribution, value must + fulfill the condition ``left <= mode <= right``. + right: a float or array of floats broadcast-compatible with ``shape`` + representing the upper limit parameter of the distribution, must be + larger than ``left``. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``left``,``mode`` and ``right``. + The default (None) produces a result shape equal to ``left.shape``, ``mode.shape`` + and ``right.shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape`` if + ``shape`` is not None, or else by ``left.shape``, ``mode.shape`` and ``right.shape``. + """ + key, _ = _check_prng_key("triangular", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError("dtype argument to `triangular` must be a float " + f"dtype, got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _triangular(key, left, mode, right, shape, dtype) + +@partial(jit, static_argnums=(4, 5), inline=True) +def _triangular(key, left, mode, right, shape, dtype) -> Array: + # https://en.wikipedia.org/wiki/Triangular_distribution#Generating_triangular-distributed_random_variates + if shape is None: + shape = lax.broadcast_shapes(np.shape(left), np.shape(mode), np.shape(right)) + else: + _check_shape("triangular", shape, np.shape(left), np.shape(mode), np.shape(right)) + left = jnp.broadcast_to(left, shape) + mode = jnp.broadcast_to(mode, shape) + right = jnp.broadcast_to(right, shape) + fc = (mode - left) / (right - left) + u = uniform(key, shape, dtype) + out1 = left + lax.sqrt(u * (right - left) * (mode - left)) + out2 = right - lax.sqrt((1 - u) * (right - left) * (right - mode)) + tri = lax.select(u < fc, out1, out2) + return tri + + +def lognormal(key: KeyArrayLike, + sigma: RealArray = np.float32(1), + shape: Shape | None = None, + dtype: DTypeLikeFloat = float) -> Array: + r""" Sample lognormal random values with given shape and float dtype. + + The values are distributed according to the probability density function: + + .. math:: + f(x) = \frac{1}{x\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(\log x)^2}{2\sigma^2}\right) + + on the domain :math:`x > 0`. + + Args: + key: a PRNG key used as the random key. + sigma: a float or array of floats broadcast-compatible with ``shape`` representing + the standard deviation of the underlying normal distribution. Default 1. + shape: optional, a tuple of nonnegative integers specifying the result + shape. The default (None) produces a result shape equal to ``()``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by ``shape``. + """ + key, _ = _check_prng_key("lognormal", key) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.inexact): + raise ValueError(f"dtype argument to `lognormal` must be a float or complex dtype, " + f"got {dtype}") + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _lognormal(key, sigma, shape, dtype) + +@partial(jit, static_argnums=(2, 3), inline=True) +def _lognormal(key, sigma, shape, dtype) -> Array: + if shape is None: + shape = np.shape(sigma) + else: + _check_shape("triangular", shape, np.shape(sigma)) + sigma = jnp.broadcast_to(sigma, shape) + scaled_norm = normal(key, shape, dtype) * sigma + return lax.exp(scaled_norm) + + +def _stirling_approx_tail(k): + stirling_tail_vals = jnp.array( + [ + 0.0810614667953272, + 0.0413406959554092, + 0.0276779256849983, + 0.02079067210376509, + 0.0166446911898211, + 0.0138761288230707, + 0.0118967099458917, + 0.0104112652619720, + 0.00925546218271273, + 0.00833056343336287, + ], + dtype=k.dtype, + ) + use_tail_values = k <= 9 + k = lax.clamp(0.0, k, 9.0) + kp1sq = (k + 1) * (k + 1) + approx = (1.0 / 12 - (1.0 / 360 - 1.0 / 1260 / kp1sq) / kp1sq) / (k + 1) + k = jnp.floor(k) + return lax.select(use_tail_values, stirling_tail_vals[jnp.int32(k)], approx) + + +@partial(jit, static_argnums=(3, 4, 5), inline=True) +def _binomial_inversion(key, count, prob, shape, dtype, max_iters): + if config.enable_checks.value: + assert jnp.issubdtype(prob.dtype, jnp.floating) + + log1minusprob = jnp.log1p(-prob) + + def body_fn(carry): + i, num_geom, geom_sum, key = carry + subkey, key = split(key) + num_geom_out = lax.select(geom_sum <= count, num_geom + 1, num_geom) + u = uniform(subkey, shape, prob.dtype) + geom = jnp.ceil(jnp.log(u) / log1minusprob) + geom_sum = geom_sum + geom + return i + 1, num_geom_out, geom_sum, key + + def cond_fn(carry): + i, geom_sum = carry[0], carry[2] + return (geom_sum <= count).any() & (i < max_iters) + + num_geom_init = lax.full_like(prob, 0, prob.dtype, shape) + geom_sum_init = lax.full_like(prob, 0, prob.dtype, shape) + carry = (0, num_geom_init, geom_sum_init, key) + k = lax.while_loop(cond_fn, body_fn, carry)[1] + return (k - 1).astype(dtype) + + +@partial(jit, static_argnums=(3, 4, 5), inline=True) +def _btrs(key, count, prob, shape, dtype, max_iters): + # transforman-rejection algorithm + # https://www.tandfonline.com/doi/abs/10.1080/00949659308811496 + stddev = jnp.sqrt(count * prob * (1 - prob)) + b = 1.15 + 2.53 * stddev + a = -0.0873 + 0.0248 * b + 0.01 * prob + c = count * prob + 0.5 + v_r = 0.92 - 4.2 / b + r = prob / (1 - prob) + alpha = (2.83 + 5.1 / b) * stddev + m = jnp.floor((count + 1) * prob) + + def body_fn(carry): + i, k_out, accepted, key = carry + key, subkey_0, subkey_1 = split(key, 3) + u = uniform(subkey_0, shape, prob.dtype) + v = uniform(subkey_1, shape, prob.dtype) + u = u - 0.5 + us = 0.5 - jnp.abs(u) + accept1 = (us >= 0.07) & (v <= v_r) + k = jnp.floor((2 * a / us + b) * u + c) + reject = (k < 0) | (k > count) + v = jnp.log(v * alpha / (a / (us * us) + b)) + ub = ( + (m + 0.5) * jnp.log((m + 1) / (r * (count - m + 1))) + + (count + 1) * jnp.log((count - m + 1) / (count - k + 1)) + + (k + 0.5) * jnp.log(r * (count - k + 1) / (k + 1)) + + _stirling_approx_tail(m) + + _stirling_approx_tail(count - m) + - _stirling_approx_tail(k) + - _stirling_approx_tail(count - k) + ) + accept2 = v <= ub + accept = accept1 | (~reject & accept2) + k_out = lax.select(accept, k, k_out) + accepted |= accept + return i + 1, k_out, accepted, key + + def cond_fn(carry): + i, accepted = carry[0], carry[2] + return (~accepted).any() & (i < max_iters) + + k_init = lax.full_like(prob, -1, prob.dtype, shape) + carry = (0, k_init, jnp.full(shape, False, jnp.bool_), key) + return lax.while_loop(cond_fn, body_fn, carry)[1].astype(dtype) + + +@partial(jit, static_argnums=(3, 4), inline=True) +def _binomial(key, count, prob, shape, dtype) -> Array: + # The implementation matches TensorFlow and TensorFlow Probability: + # https://github.com/tensorflow/tensorflow/blob/v2.2.0-rc3/tensorflow/core/kernels/random_binomial_op.cc + # and tensorflow_probability.substrates.jax.distributions.Binomial + # For n * p < 10, we use the binomial inverse algorithm; otherwise btrs. + if shape is None: + shape = jnp.broadcast_shapes(jnp.shape(count), jnp.shape(prob)) + else: + _check_shape("binomial", shape, np.shape(count), np.shape(prob)) + (prob,) = promote_dtypes_inexact(prob) + count = lax.convert_element_type(count, prob.dtype) + count = jnp.broadcast_to(count, shape) + prob = jnp.broadcast_to(prob, shape) + p_lt_half = prob < 0.5 + q = lax.select(p_lt_half, prob, 1.0 - prob) + count_nan_or_neg = _isnan(count) | (count < 0.0) + count_inf = jnp.isinf(count) + q_is_nan = _isnan(q) + q_l_0 = q < 0.0 + q = lax.select(q_is_nan | q_l_0, lax.full_like(q, 0.01), q) + use_inversion = count_nan_or_neg | (count * q <= 10.0) + + # consistent with np.random.binomial behavior for float count input + count = jnp.floor(count) + + count_inv = lax.select(use_inversion, count, lax.full_like(count, 0.0)) + count_btrs = lax.select(use_inversion, lax.full_like(count, 1e4), count) + q_btrs = lax.select(use_inversion, lax.full_like(q, 0.5), q) + max_iters = dtype.type(jnp.finfo(dtype).max) + samples = lax.select( + use_inversion, + _binomial_inversion(key, count_inv, q, shape, dtype, max_iters), + _btrs(key, count_btrs, q_btrs, shape, dtype, max_iters), + ) + # ensure nan q always leads to nan output and nan or neg count leads to nan + # as discussed in https://github.com/google/jax/pull/16134#pullrequestreview-1446642709 + invalid = (q_l_0 | q_is_nan | count_nan_or_neg) + samples = lax.select( + invalid, + jnp.full_like(samples, jnp.nan, dtype), + samples, + ) + + # +inf count leads to inf + samples = lax.select( + count_inf & (~invalid), + jnp.full_like(samples, jnp.inf, dtype), + samples, + ) + + samples = lax.select( + p_lt_half | count_nan_or_neg | q_is_nan | count_inf, + samples, + count.astype(dtype) - samples, + ) + return samples + + +def binomial( + key: KeyArray, + n: RealArray, + p: RealArray, + shape: Shape | None = None, + dtype: DTypeLikeFloat = float, +) -> Array: + r"""Sample Binomial random values with given shape and float dtype. + + The values are returned according to the probability mass function: + + .. math:: + f(k;n,p) = \binom{n}{k}p^k(1-p)^{n-k} + + on the domain :math:`0 < p < 1`, and where :math:`n` is a nonnegative integer + representing the number of trials and :math:`p` is a float representing the + probability of success of an individual trial. + + Args: + key: a PRNG key used as the random key. + n: a float or array of floats broadcast-compatible with ``shape`` + representing the number of trials. + p: a float or array of floats broadcast-compatible with ``shape`` + representing the probability of success of an individual trial. + shape: optional, a tuple of nonnegative integers specifying the result + shape. Must be broadcast-compatible with ``n`` and ``p``. + The default (None) produces a result shape equal to ``np.broadcast(n, p).shape``. + dtype: optional, a float dtype for the returned values (default float64 if + jax_enable_x64 is true, otherwise float32). + + Returns: + A random array with the specified dtype and with shape given by + ``np.broadcast(n, p).shape``. + """ + key, _ = _check_prng_key("binomial", key) + check_arraylike("binomial", n, p) + dtypes.check_user_dtype_supported(dtype) + if not dtypes.issubdtype(dtype, np.floating): + raise ValueError( + "dtype argument to `binomial` must be a float dtype, got {dtype}" + ) + dtype = dtypes.canonicalize_dtype(dtype) + if shape is not None: + shape = core.canonicalize_shape(shape) + return _binomial(key, n, p, shape, dtype) + + +# Functions related to key reuse checking +random_clone_p = core.Primitive("random_clone") +dispatch.simple_impl(random_clone_p) +random_clone_p.def_abstract_eval(lambda x: x) +batching.defvectorized(random_clone_p) +mlir.register_lowering(random_clone_p, lambda _, k: [k]) + +def clone(key): + """Clone a key for reuse + + Outside the context of key reuse checking (see :mod:`jax.experimental.key_reuse`) + this function operates as an identity. + + Example: + + >>> import jax + >>> key = jax.random.key(0) + >>> data = jax.random.uniform(key) + >>> cloned_key = jax.random.clone(key) + >>> same_data = jax.random.uniform(cloned_key) + >>> assert data == same_data + """ + return random_clone_p.bind(key) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/shard_alike.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/shard_alike.py new file mode 100644 index 0000000000000000000000000000000000000000..2361eaf6426dfccb8c273feeffee49e54fba1bec --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/shard_alike.py @@ -0,0 +1,109 @@ +# Copyright 2023 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from functools import partial +import itertools + +from jax._src import core +from jax._src.interpreters import ad +from jax._src.interpreters import mlir +from jax._src.dispatch import apply_primitive +from jax._src.tree_util import tree_flatten, tree_unflatten +from jax._src.interpreters import batching +from jax._src.util import safe_zip +from jax._src.lib import xla_client as xc +from jax._src.api_util import shaped_abstractify +from jax._src.lib.mlir import ir + +_next_shard_group_id = itertools.count() + +def shard_alike(x, y): + """Shards x and y alike.""" + x_flat, x_tree = tree_flatten(x) + y_flat, y_tree = tree_flatten(y) + + if x_tree != y_tree: + raise ValueError('Trees should be equal. ' + f'Got x_tree: {x_tree}, y_tree: {y_tree}') + + for x_, y_ in safe_zip(x_flat, y_flat): + x_aval = shaped_abstractify(x_) + y_aval = shaped_abstractify(y_) + if x_aval.shape != y_aval.shape: + raise ValueError( + 'The leaves shapes of `x` and `y` should match. Got `x` leaf shape:' + f' {x_aval.shape} and `y` leaf shape: {y_aval.shape}. File an issue at' + ' https://github.com/google/jax/issues if you want this feature.') + + outs = [shard_alike_p.bind(x_, y_) for x_, y_ in safe_zip(x_flat, y_flat)] + x_out_flat, y_out_flat = zip(*outs) + return tree_unflatten(x_tree, x_out_flat), tree_unflatten(y_tree, y_out_flat) + + +shard_alike_p = core.Primitive('shard_alike') +shard_alike_p.multiple_results = True +shard_alike_p.def_impl(partial(apply_primitive, shard_alike_p)) +shard_alike_p.def_abstract_eval(lambda x, y: (x, y)) + +def shard_alike_transpose(ct, **kwargs): + x_ct, y_ct = ct + if type(x_ct) is ad.Zero or type(y_ct) is ad.Zero: + return x_ct, y_ct + else: + return shard_alike(x_ct, y_ct) +ad.deflinear(shard_alike_p, shard_alike_transpose) + + +def _shard_alike_batcher(batched_args, batch_dims): + x, y = batched_args + xd, yd = batch_dims + if xd == yd: + return shard_alike(x, y), (xd, yd) + elif xd is batching.not_mapped: + x = batching.broadcast(x, y.shape[yd], yd) + return shard_alike(x, y), (yd, yd) + elif yd is batching.not_mapped: + y = batching.broadcast(y, x.shape[xd], xd) + return shard_alike(x, y), (xd, xd) + else: + y = batching.moveaxis(y, yd, xd) + return shard_alike(x, y), (xd, xd) +batching.primitive_batchers[shard_alike_p] = _shard_alike_batcher + + +def _group_shard( + ctx, + x: ir.Value, + y: ir.Value, + x_aval_out: core.AbstractValue, + y_aval_out: core.AbstractValue, +) -> tuple[ir.Value, ir.Value]: + shard_group_id = next(_next_shard_group_id) + + unknown_op_sharding = xc.OpSharding() + unknown_op_sharding.type = xc.OpSharding.Type.UNKNOWN + unknown_op_sharding.is_shard_group = True + unknown_op_sharding.shard_group_id = shard_group_id + unknown_op_sharding.shard_group_type = xc.OpSharding.ShardGroupType.AS + + x = mlir.wrap_with_sharding_op(ctx, x, x_aval_out, unknown_op_sharding, + has_side_effect=True) + y = mlir.wrap_with_sharding_op(ctx, y, y_aval_out, unknown_op_sharding, + has_side_effect=True) + return x, y + + +def shard_alike_lowering(ctx, x, y): + return _group_shard(ctx, x, y, *ctx.avals_out) +mlir.register_lowering(shard_alike_p, shard_alike_lowering) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding.py new file mode 100644 index 0000000000000000000000000000000000000000..29f2b67b16ae911ba0930ee7fbc83c1dd1dbe4fe --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding.py @@ -0,0 +1,197 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import functools + +from jax._src.util import safe_zip, use_cpp_class +from jax._src import xla_bridge as xb +from jax._src.lib import xla_client as xc +from jax._src.op_shardings import ( + are_op_shardings_equal, get_num_ways_dim_sharded, is_op_sharding_replicated, + op_sharding_to_indices) + +Shape = tuple[int, ...] +Device = xc.Device +Index = tuple[slice, ...] +XLADeviceAssignment = Sequence[Device] + + +@functools.lru_cache(maxsize=4096) +def _addressable_devices_indices_map( + sharding: Sharding, global_shape: Shape) -> Mapping[Device, Index | None]: + global_map = sharding.devices_indices_map(global_shape) + if sharding.is_fully_addressable: + return global_map + if hasattr(sharding, '_internal_device_list'): + return {d: global_map[d] + for d in sharding._internal_device_list.addressable_device_list} + return {d: ind for d, ind in global_map.items() + if d.process_index == d.client.process_index()} + +@functools.lru_cache(maxsize=4096) +def common_devices_indices_map(s, global_shape: Shape) -> Mapping[Device, Index]: + s.shard_shape(global_shape) # raises a good error message + hlo_sharding = s._to_xla_hlo_sharding(len(global_shape)) + indices = op_sharding_to_indices(hlo_sharding, global_shape, + len(s._device_assignment)) + return dict(safe_zip(s._device_assignment, indices)) + + +@functools.lru_cache(maxsize=4096) +def _common_shard_shape(self, global_shape: Shape) -> Shape: + hlo_sharding = self._to_xla_hlo_sharding(len(global_shape)) + if is_op_sharding_replicated(hlo_sharding): + return global_shape + partitions, _ = get_num_ways_dim_sharded(hlo_sharding) + assert len(partitions) == len(global_shape), (len(partitions), len(global_shape)) + out = [] + for dim, (s, p) in enumerate(safe_zip(global_shape, partitions)): + try: + quotient, remainder = divmod(s, p) + except TypeError: + # TODO Figure out how to partition dynamic shapes + raise NotImplementedError + if remainder != 0: + raise ValueError( + f"Sharding {self} implies that array axis {dim} is partitioned " + f"{p} times, but the dimension size is {s} " + f"(full shape: {global_shape}, " + f"per-dimension tiling factors: {partitions} should evenly divide " + "the shape)") + out.append(quotient) + return tuple(out) + + +@use_cpp_class(xc.Sharding) +class Sharding: + """Describes how a :class:`jax.Array` is laid out across devices. + """ + + # Abstract methods below that subclasses should implement. + @property + def device_set(self) -> set[Device]: + """The set of devices that this :class:`Sharding` spans. + + In multi-controller JAX, the set of devices is global, i.e., includes + non-addressable devices from other processes. + """ + raise NotImplementedError('Subclasses should implement this method.') + + @property + def is_fully_replicated(self) -> bool: + """Is this sharding fully replicated? + + A sharding is fully replicated if each device has a complete copy of the + entire data. + """ + raise NotImplementedError('Subclasses should implement this method.') + + @property + def is_fully_addressable(self) -> bool: + """Is this sharding fully addressable? + + A sharding is fully addressable if the current process can address all of + the devices named in the :class:`Sharding`. ``is_fully_addressable`` is + equivalent to "is_local" in multi-process JAX. + """ + raise NotImplementedError('Subclasses should implement this method.') + + @property + def memory_kind(self) -> str | None: + """Returns the memory kind of the sharding.""" + raise NotImplementedError('Subclasses should implement this method.') + + def with_memory_kind(self, kind: str) -> Sharding: + """Returns a new Sharding instance with the specified memory kind.""" + raise NotImplementedError('Subclasses should implement this method') + + @property + def _device_assignment(self) -> XLADeviceAssignment: + raise NotImplementedError('Subclasses should implement this method.') + + def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding: + raise NotImplementedError('Subclasses should implement this method.') + + + ############################################################################# + # Default implementations below that all subclasses will inherit. + + @functools.cached_property + def addressable_devices(self) -> set[Device]: + """The set of devices in the :class:`Sharding` that are addressable by the + current process. + """ + # Add a fast path for single controller runtimes. + if xb.process_count() == 1: + return self.device_set + return {d for d in self.device_set + if d.process_index == d.client.process_index()} + + def addressable_devices_indices_map( + self, global_shape: Shape) -> Mapping[Device, Index | None]: + """A mapping from addressable devices to the slice of array data each contains. + + ``addressable_devices_indices_map`` contains that part of + ``device_indices_map`` that applies to the addressable devices. + """ + return _addressable_devices_indices_map(self, global_shape) + + def devices_indices_map(self, global_shape: Shape) -> Mapping[Device, Index]: + """Returns a mapping from devices to the array slices each contains. + + The mapping includes all global devices, i.e., including + non-addressable devices from other processes. + """ + return common_devices_indices_map(self, global_shape) + + @functools.cached_property + def _addressable_device_assignment(self) -> XLADeviceAssignment: + if self.is_fully_addressable: + return self._device_assignment + if hasattr(self, '_internal_device_list'): + return tuple(self._internal_device_list.addressable_device_list) + return tuple(d for d in self._device_assignment + if d.process_index == d.client.process_index()) + + def shard_shape(self, global_shape: Shape) -> Shape: + """Returns the shape of the data on each device. + + The shard shape returned by this function is calculated from + ``global_shape`` and the properties of the sharding. + """ + return _common_shard_shape(self, global_shape) + + def is_equivalent_to(self: Sharding, other: Sharding, ndim: int) -> bool: + """Returns ``True`` if two shardings are equivalent. + + Two shardings are equivalent if they place the same logical array shards on + the same devices. + + For example, a :class:`NamedSharding` may be equivalent + to a :class:`PositionalSharding` if both place the same shards of the array + on the same devices. + """ + try: + return (are_op_shardings_equal(self._to_xla_hlo_sharding(ndim), + other._to_xla_hlo_sharding(ndim)) + and self._internal_device_list == other._internal_device_list and # type: ignore + self.memory_kind == other.memory_kind) + # NotImplementedError is raised by PmapSharding because it can't lower + # to OpSharding. So if `other` is a PmapSharding, default to a strict + # equality check. + except NotImplementedError: + return self == other diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding_impls.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding_impls.py new file mode 100644 index 0000000000000000000000000000000000000000..3e98327b49b48cdf1a33a7b2cc688df8fbe90f4f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding_impls.py @@ -0,0 +1,1641 @@ +# Copyright 2021 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import collections +from collections import OrderedDict +from collections.abc import Mapping, Sequence +import dataclasses +import enum +import functools +import itertools +import math +from typing import Any, NamedTuple, Union, cast + +from jax._src import mesh as mesh_lib +from jax._src import sharding +from jax._src import sharding_specs +from jax._src import tree_util +from jax._src import util +from jax._src import xla_bridge +from jax._src import core +from jax._src.lib import xla_client as xc +from jax._src.lib import xla_extension_version +from jax._src.op_shardings import ( + are_op_shardings_equal, get_num_ways_dim_sharded, is_op_sharding_replicated) +from jax._src.partition_spec import PartitionSpec +from jax._src.util import safe_map, safe_zip, use_cpp_class, use_cpp_method +import numpy as np + + +Shape = tuple[int, ...] +Device = xc.Device +Index = tuple[slice, ...] +XLADeviceAssignment = tuple[Device, ...] +# TODO(yashkatariya): Remove this after 3 months of deprecation. +XLACompatibleSharding = sharding.Sharding + +@dataclasses.dataclass(frozen=True) +class TransferToMemoryKind: + memory_kind: str + + +@functools.lru_cache +def _check_mesh_resource_axis(mesh, parsed_pspec, _manual_axes): + try: + for p in parsed_pspec: + if p is not None: + for r in p: + mesh.shape[r] + if r in _manual_axes: + raise ValueError( + f"Axis: {r} of {parsed_pspec.get_partition_spec()} " + f"is also found in manual_axes: {_manual_axes}.") from None + except KeyError as e: + raise ValueError(f"Resource axis: {e.args[0]} of {parsed_pspec.user_spec} is " + "undefined.") from None + + +def hashed_index(x) -> int: + # This works for both `pjit`/`xmap` indices and `pmap` indices (which might + # have an integer instead of a slice). + assert all(v.step is None for v in x if isinstance(v, slice)) + return hash(tuple((v.start, v.stop) if isinstance(v, slice) else v for v in x)) + + +@functools.lru_cache(maxsize=4096) +def device_replica_id_map(sharding, global_shape: Shape) -> Mapping[Device, int]: + try: + device_indices_map_fn = sharding.devices_indices_map + except AttributeError: + raise ValueError( + f'Cannot calculate replica ids from sharding: {sharding}. Please ' + 'create a device to index mapping for your sharding from which replica ' + 'ids will be calculated.') from None + + index_to_replica: dict[int, int] = collections.Counter() + out = {} + for device, index in device_indices_map_fn(global_shape).items(): + h_index = hashed_index(index) + replica_id = index_to_replica[h_index] + index_to_replica[h_index] += 1 + out[device] = replica_id + return out + + +@functools.lru_cache(maxsize=4096) +def named_sharding_to_xla_hlo_sharding( + self, num_dimensions: int) -> xc.HloSharding: + mesh_shape = self.mesh.shape + array_mapping = get_array_mapping(self._parsed_pspec) + mesh_axis_pos = {name: i for i, name in enumerate(self.mesh.axis_names)} + + special_axes = {} + if self._manual_axes: + axis_names = self.mesh.axis_names + for manual_axis in self._manual_axes: + special_axes[axis_names.index(manual_axis)] = xc.OpSharding.Type.MANUAL + + replicated_mesh_axes = [] + for i, (axis_name, axis_val) in enumerate(mesh_shape.items()): + if axis_name not in array_mapping: # type: ignore + replicated_mesh_axes.append((i, axis_val)) + + if len(replicated_mesh_axes) == len(mesh_shape) and not special_axes: + return xc.HloSharding.replicate() + + mesh_permutation = [] + new_mesh_shape = [1] * num_dimensions + for name, pos in sorted(array_mapping.items(), key=lambda x: x[1]): # type: ignore + new_mesh_shape[pos] *= mesh_shape[name] + mesh_permutation.append(mesh_axis_pos[name]) + + last_tile_dims = [] + if replicated_mesh_axes: + axes_by_type = collections.defaultdict(list) + size_by_type = collections.defaultdict(lambda: 1) # type: ignore + assert {x[0] for x in replicated_mesh_axes}.issuperset(set(special_axes.keys())) + for i, size in replicated_mesh_axes: + ty = special_axes.get(i, xc.OpSharding.Type.REPLICATED) + axes_by_type[ty].append(i) + size_by_type[ty] *= size + for ty, axes in sorted(axes_by_type.items(), key=lambda x: x[0].value): + last_tile_dims.append(ty) + new_mesh_shape.append(size_by_type[ty]) + mesh_permutation.extend(axes) + + # Explanation of the parameters of `HloSharding.iota_tile`. + # This is the HloShardingV2 format: + # * dims: How many ways each dimension is sharded. + # Replicated/Manual dims are added added at the end + # * reshape_dims: This is the just the shape of the mesh. + # * transpose_perm: This is the order in which mesh axes in PartitionSpec + # appear relative to mesh.axis_names order. + # * subgroup_types: List of type of OpSharding. Type can be REPLICATED and MANUAL. + # Let's see an example: + # Consider input_shape=(8, 4, 2, 2), mesh={'a': 2, 'b': 2, 'c': 2, 'd': 2} + # and partition_spec=P(None, ('d', 'b'), 'c'). + # Arguments to iota_tile will be: + # dims = [1, 4, 2, 1, 2] # 'a' is replicated hence `2` is at the end. + # reshape_dims = [2, 2, 2, 2] + # transpose_perm = [3, 1, 2, 0] # 'a' is replicated hence 0 is at the end + # subgroup_types = [xc.OpSharding.Type.REPLICATED] + return xc.HloSharding.iota_tile( + dims=new_mesh_shape, reshape_dims=tuple(self.mesh.shape.values()), + transpose_perm=mesh_permutation, subgroup_types=last_tile_dims) + + +@use_cpp_class(xc.NamedSharding) +class NamedSharding(sharding.Sharding): + r"""A :class:`NamedSharding` expresses sharding using named axes. + + A :class:`NamedSharding` is a pair of a :class:`Mesh` of devices and + :class:`PartitionSpec` which describes how to shard an array across that + mesh. + + A :class:`Mesh` is a multidimensional NumPy array of JAX devices, + where each axis of the mesh has a name, e.g. ``'x'`` or ``'y'``. + + A :class:`PartitionSpec` is a tuple, whose elements can be a ``None``, + a mesh axis, or a tuple of mesh axes. Each element describes how an input + dimension is partitioned across zero or more mesh dimensions. For example, + ``PartitionSpec('x', 'y')`` says that the first dimension of data + is sharded across ``x`` axis of the mesh, and the second dimension is sharded + across ``y`` axis of the mesh. + + The Distributed arrays and automatic parallelization + (https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#namedsharding-gives-a-way-to-express-shardings-with-names) + tutorial has more details and diagrams that explain how + :class:`Mesh` and :class:`PartitionSpec` are used. + + Args: + mesh: A :class:`jax.sharding.Mesh` object. + spec: A :class:`jax.sharding.PartitionSpec` object. + + Example: + + >>> from jax.sharding import Mesh + >>> from jax.sharding import PartitionSpec as P + >>> mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y')) + >>> spec = P('x', 'y') + >>> named_sharding = jax.sharding.NamedSharding(mesh, spec) + """ + + mesh: mesh_lib.Mesh + spec: PartitionSpec + _memory_kind: str | None + _parsed_pspec: ParsedPartitionSpec + _manual_axes: frozenset[MeshAxisName] + + @use_cpp_method() + def __init__( + self, mesh: mesh_lib.Mesh, spec: PartitionSpec, *, + memory_kind: str | None = None, _parsed_pspec=None, + _manual_axes=frozenset()): + self.mesh = mesh + self.spec = spec + self._memory_kind = memory_kind + self._manual_axes = _manual_axes + self._parsed_pspec = preprocess(self.mesh, self.spec, _parsed_pspec) + + def __repr__(self): + mesh_repr = ", ".join(f"'{k}': {v}" for k, v in self.mesh.shape.items()) + mem = '' if self.memory_kind is None else f', memory_kind={self.memory_kind}' + return f'NamedSharding(mesh=Mesh({mesh_repr}), spec={self.spec}{mem})' + + def __reduce__(self): + return (type(self), (self.mesh, self.spec), + {'memory_kind': self.memory_kind, + '_manual_axes': self._manual_axes}) + + @property + def memory_kind(self) -> str | None: + return self._memory_kind + + def __hash__(self): + if not hasattr(self, '_hash'): + self._hash = hash( + (self.mesh, self.memory_kind, self._parsed_pspec, self._manual_axes)) + return self._hash + + def __eq__(self, other): + if not isinstance(other, NamedSharding): + return False + if self is other: + return True + if (self._parsed_pspec != other._parsed_pspec + or self.memory_kind != other.memory_kind + or self._manual_axes != other._manual_axes): + return False + return self.mesh is other.mesh or self.mesh == other.mesh + + def check_compatible_aval(self, aval_shape: Shape) -> None: + assert self._parsed_pspec is not None + if len(aval_shape) < len(self._parsed_pspec): + extra_msg = (' For scalars the PartitionSpec should be P()' + if len(aval_shape) == 0 else '') + raise ValueError( + f"Sharding {self} is only valid for values of rank at least " + f"{len(self._parsed_pspec)}, but was applied to a value of rank " + f"{len(aval_shape)}.{extra_msg}") + + @classmethod + def _from_parsed_pspec( + cls, mesh, parsed_pspec, *, memory_kind=None, _manual_axes=frozenset() + ): + return cls(mesh, parsed_pspec.get_partition_spec(), + memory_kind=memory_kind, _parsed_pspec=parsed_pspec, + _manual_axes=_manual_axes) + + @property + def device_set(self) -> set[Device]: + return self.mesh._flat_devices_set + + @property + def _device_assignment(self) -> XLADeviceAssignment: + return self.mesh._flat_devices_tuple + + @property + def is_fully_addressable(self) -> bool: + # Speed up `is_fully_addressable` since there is a high chance that the + # mesh across multiple NamedSharding objects will be the same. + return not self.mesh.is_multi_process + + @property + def addressable_devices(self) -> set[Device]: + # Override addressable devices because there is a high chance that the mesh + # across multiple NamedSharding objects will be the same. + return self.mesh._local_devices_set + + @functools.cached_property + def is_fully_replicated(self) -> bool: + if self.mesh.size == 1: + return True + array_mapping = cast(ParsedPartitionSpec, get_array_mapping(self._parsed_pspec)) + mesh_shape = self.mesh.shape + num_partitions = 1 + for name in array_mapping: + num_partitions *= mesh_shape[name] + return num_partitions == 1 + + def with_memory_kind(self, kind: str) -> NamedSharding: + return NamedSharding(self.mesh, self.spec, memory_kind=kind) + + def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding: + return named_sharding_to_xla_hlo_sharding(self, num_dimensions) + + +@functools.lru_cache +def get_replicated_hlo_sharding(): + return xc.HloSharding.replicate() + + +@use_cpp_class(xc.SingleDeviceSharding) +class SingleDeviceSharding(sharding.Sharding): + """A :class:`Sharding` that places its data on a single device. + + Args: + device: A single :py:class:`Device`. + + Example: + + >>> single_device_sharding = jax.sharding.SingleDeviceSharding( + ... jax.devices()[0]) + """ + + _device: Device + _memory_kind: str | None + + @use_cpp_method() + def __init__(self, device: Device, *, memory_kind: str | None = None): + self._device = device + self._memory_kind = memory_kind + + def __reduce__(self): + return type(self), (self._device,), {'memory_kind': self._memory_kind} + + def __repr__(self): + mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}' + return f"SingleDeviceSharding(device={self._device!r}{mem})" + + def __hash__(self): + if not hasattr(self, '_hash'): + self._hash = hash((self._device, self.memory_kind)) + return self._hash + + def __eq__(self, other): + if not isinstance(other, SingleDeviceSharding): + return False + if self is other: + return True + return (self._device == other._device and + self.memory_kind == other.memory_kind) + + @property + def device_set(self) -> set[Device]: + return {self._device} + + @property + def memory_kind(self) -> str | None: + return self._memory_kind + + def with_memory_kind(self, kind: str) -> SingleDeviceSharding: + return SingleDeviceSharding(self._device, memory_kind=kind) + + def devices_indices_map(self, global_shape: Shape) -> Mapping[Device, Index]: # type: ignore + return {self._device: (slice(None),) * len(global_shape)} + + @property + def _device_assignment(self) -> XLADeviceAssignment: + return (self._device,) + + def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding: + return get_replicated_hlo_sharding() + + @property + def is_fully_replicated(self) -> bool: + return True + + @property + def is_fully_addressable(self) -> bool: + return True + + +@functools.lru_cache(maxsize=4096) +def pmap_sharding_devices_indices_map( + self, global_shape: Shape) -> Mapping[Device, Index]: + self.shard_shape(global_shape) # raises a good error message + indices = sharding_specs.spec_to_indices(global_shape, self.sharding_spec) + return dict(safe_zip(self.devices.flat, indices)) # type: ignore[arg-type] + + +@use_cpp_class(xc.PmapSharding) +class PmapSharding(sharding.Sharding): + """Describes a sharding used by :func:`jax.pmap`.""" + devices: np.ndarray + sharding_spec: sharding_specs.ShardingSpec + _internal_device_list: xc.DeviceList + + @use_cpp_method() + def __init__(self, devices: Sequence[Device] | np.ndarray, + sharding_spec: sharding_specs.ShardingSpec): + self.devices = np.asarray(devices) + # The sharding spec should be pmap's sharding spec. + self.sharding_spec = sharding_spec + + def __reduce__(self): + return (type(self), (self.devices, self.sharding_spec), + {'memory_kind': self.memory_kind}) + + def __eq__(self, other): + if not isinstance(other, PmapSharding): + return False + if self is other: + return True + return (self.sharding_spec == other.sharding_spec and + self.devices.shape == other.devices.shape and + self._internal_device_list == other._internal_device_list) + + def __hash__(self): + if not hasattr(self, '_hash'): + self._hash = hash((self._internal_device_list, self.sharding_spec)) + return self._hash + + def __str__(self): + device_ids = [d.id for d in self.devices.flat] + return (f'PmapSharding(sharding_spec={self.sharding_spec}, ' + f'{device_ids=}, ' + f'device_platform={self.devices.flat[0].platform.upper()}, ' + f'device_shape={self.devices.shape})') + + def __repr__(self): + return (f'PmapSharding(sharding_spec={self.sharding_spec}, ' + f'devices={self.devices})') + + def is_equivalent_to(self: PmapSharding, other: PmapSharding, # type: ignore + ndim: int) -> bool: + return self == other + + # TODO(yashkatariya): Expose `sharded_dim_size` in the API if required. + @classmethod + def default(cls, shape: Shape, sharded_dim: int = 0, + devices: Sequence[xc.Device] | None = None) -> PmapSharding: + """Creates a :class:`PmapSharding` which matches the default placement + used by :func:`jax.pmap`. + + Args: + shape: The shape of the input array. + sharded_dim: Dimension the input array is sharded on. Defaults to 0. + devices: Optional sequence of devices to use. If omitted, the implicit + device order used by pmap is used, which is the order of + :func:`jax.local_devices`. + """ + # The dtype doesn't matter here. Its only used for creating the + # sharding_spec. + sharding_spec = sharding_specs.create_pmap_sharding_spec( + tuple(shape), sharded_dim) + + num_ways_sharded = None + for s in sharding_spec.sharding: + if isinstance(s, sharding_specs.Unstacked): + assert num_ways_sharded is None + num_ways_sharded = s.size + elif isinstance(s, sharding_specs.Chunked): + assert num_ways_sharded is None + if len(s.chunks) == 1: + num_ways_sharded = s.chunks[0] + else: + raise NotImplementedError( + 'Multiple chunks in Chunked dimension not supported.') + + if num_ways_sharded is None: + raise NotImplementedError( + '`None` to sharded_dim is not supported. Please file a jax ' + 'issue if you need this feature.') + + if devices is None: + pmap_devices: np.ndarray = np.array( + xla_bridge.local_devices()[:num_ways_sharded]) + else: + pmap_devices = np.array(devices) + return cls(pmap_devices, sharding_spec) + + @functools.cached_property + def device_set(self) -> set[Device]: + return set(self.devices.flat) + + def devices_indices_map(self, global_shape: Shape) -> Mapping[Device, Index]: + return pmap_sharding_devices_indices_map(self, global_shape) + + @functools.cached_property + def _device_assignment(self) -> XLADeviceAssignment: + return tuple(self.devices.flat) + + @property + def memory_kind(self) -> str | None: + try: + return self._internal_device_list.default_memory_kind + except: + return None + + def with_memory_kind(self, kind: str): + raise NotImplementedError("pmap does not support memories.") + + def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding: + raise NotImplementedError("pmap doesn't use OpSharding.") + + @functools.cached_property + def is_fully_replicated(self) -> bool: + for s in self.sharding_spec.sharding: + if isinstance(s, (sharding_specs.Unstacked, sharding_specs.Chunked)): + return False + return True + + @functools.cached_property + def is_fully_addressable(self) -> bool: + return self._internal_device_list.is_fully_addressable + + def shard_shape(self, global_shape: Shape) -> Shape: + sharded_dim = None + sharded_dim_size = None + for i, s in enumerate(self.sharding_spec.sharding): + if isinstance(s, sharding_specs.Unstacked): + sharded_dim = i + sharded_dim_size = s.size + sharded_shape = util.tuple_delete(global_shape, sharded_dim) + break + elif isinstance(s, sharding_specs.Chunked): + sharded_dim = i + assert len(s.chunks) == 1, s.chunks + sharded_dim_size = s.chunks[0] + sharded_shape = util.tuple_update(global_shape, sharded_dim, 1) + break + if sharded_dim is None: + return global_shape + if global_shape[sharded_dim] != sharded_dim_size: + raise ValueError( + f'The sharded dimension must be equal to the number of ' + f'devices passed to PmapSharding. Got sharded dimension {sharded_dim} ' + f'with value {global_shape[sharded_dim]} in shape {global_shape} and ' + f'the number of devices={len(self._device_assignment)}') + return sharded_shape + + +def _op_sharding_to_pos_sharding( + op_sharding: xc.OpSharding | xc.HloSharding, + device_assignment: Sequence[xc.Device], + memory_kind: str | None = None) -> PositionalSharding: + if isinstance(op_sharding, xc.OpSharding): + op_sharding = xc.HloSharding.from_proto(op_sharding) + + if op_sharding.is_replicated(): + return PositionalSharding( + device_assignment, memory_kind=memory_kind).replicate() + + if len(op_sharding.subgroup_types()) > 1: + raise NotImplementedError( + 'Unhandled HloSharding type. Please open a bug report!' + ) + + name = device_assignment[0].platform.upper() + ids = np.array( + [DeviceIdSet(name, i) for i in op_sharding.tile_assignment_devices()] + ) + p = PositionalSharding._remake(tuple(device_assignment), ids, + memory_kind=memory_kind) + p = p.reshape(op_sharding.tile_assignment_dimensions()) + if op_sharding.replicate_on_last_tile_dim(): + p = p.replicate(-1, keepdims=False) + return p + + +@functools.lru_cache(maxsize=4096) +def _positional_sharding_to_xla_hlo_sharding( + self, num_dimensions: int) -> xc.HloSharding: + if self.shape == (1,) * self.ndim: + return get_replicated_hlo_sharding() + + pbuf = xc.OpSharding() + shape = self.shape[self.ndim - num_dimensions:] # 'rank promotion' of val + set_size, = {len(device_set) for device_set in self._ids.flat} + pbuf.type = xc.OpSharding.Type.OTHER + if set_size > 1: + pbuf.last_tile_dims = [xc.OpSharding.Type.REPLICATED] + pbuf.tile_assignment_dimensions = (*shape, set_size) + else: + pbuf.tile_assignment_dimensions = shape + pbuf.tile_assignment_devices = [i for ids in self._ids.flat for i in ids] + product_of_dims = math.prod(pbuf.tile_assignment_dimensions) + num_devices = len(pbuf.tile_assignment_devices) + assert product_of_dims == num_devices, (product_of_dims, num_devices) + return xc.HloSharding.from_proto(pbuf) + + +class PositionalSharding(sharding.Sharding): + _devices: tuple[xc.Device, ...] + _memory_kind: str | None + _ids: np.ndarray # dtype DeviceIdSet + + def __init__(self, devices: Sequence[xc.Device] | np.ndarray, + *, memory_kind: str | None = None): + super().__init__() + if not isinstance(devices, np.ndarray): + devices = np.array(devices, dtype='object') + if not devices.size: + raise ValueError(f"{self.__class__.__name__}.__init__ requires at least " + f"one device, got {devices}") + self._devices = tuple(devices.flat) + self._memory_kind = memory_kind + name = self._devices[0].platform.upper() + self._ids = np.array([DeviceIdSet(name, i) for i in range(devices.size)], + dtype='object').reshape(devices.shape) + self._internal_device_list = xc.DeviceList(self._devices) + self._memory_kind = xc.check_and_canonicalize_memory_kind( + self._memory_kind, self._internal_device_list) + + @property + def shape(self): + return self._ids.shape + + @property + def ndim(self): + return self._ids.ndim + + def __repr__(self) -> str: + cls_name = self.__class__.__name__ + ids = self._ids.copy() + platform_name = self._devices[0].platform.upper() + for idx, x in np.ndenumerate(ids): + ids[idx] = DeviceIdSet(platform_name, *(self._devices[i].id for i in x)) + body = np.array2string(ids, prefix=cls_name + '(', suffix=')', + max_line_width=100) + mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}' + return f'{cls_name}({body}{mem}, shape={self.shape})' + + def reshape(self, *shape) -> PositionalSharding: + return self._remake(self._devices, self._ids.reshape(*shape)) + + def transpose(self, *axes) -> PositionalSharding: + return self._remake(self._devices, self._ids.transpose(*axes)) + T = property(transpose) + + def replicate(self, axis=None, keepdims=True) -> PositionalSharding: + new_ids = self._ids.sum(axis=axis, keepdims=keepdims) # union + return self._remake(self._devices, new_ids) + + def check_compatible_aval(self, aval_shape: Shape) -> None: + if len(aval_shape) != len(self.shape): + raise ValueError( + f"Sharding {self} is only valid for values of rank " + f"{len(self.shape)}, but was applied to a value of rank " + f"{len(aval_shape)}") + + @classmethod + def _remake( + cls, devices: tuple[xc.Device, ...], ids: np.ndarray, + *, memory_kind: str | None = None) -> PositionalSharding: + self = cls.__new__(cls) + self._devices = devices + self._ids = ids + self._internal_device_list = xc.DeviceList(self._devices) + self._memory_kind = xc.check_and_canonicalize_memory_kind( + memory_kind, self._internal_device_list) + return self + + # Hashable + + def __hash__(self) -> int: + if not hasattr(self, '_hash'): + self._hash = hash((self._internal_device_list, self.memory_kind)) + return self._hash + + def __eq__(self, other) -> bool: + if not isinstance(other, PositionalSharding): + return False + if self is other: + return True + all_ids_equal = np.array_equal(self._ids,other._ids) + mem_kind_equal = self.memory_kind == other.memory_kind + if self._devices is other._devices and mem_kind_equal and all_ids_equal: + return True + return (mem_kind_equal and all_ids_equal and + self._internal_device_list == other._internal_device_list) + + # Sharding interface + + @functools.cached_property + def device_set(self) -> set[xc.Device]: + return set(self._devices) + + @property + def memory_kind(self) -> str | None: + return self._memory_kind + + def with_memory_kind(self, kind: str) -> PositionalSharding: + return PositionalSharding(self._devices, memory_kind=kind) + + @functools.cached_property + def is_fully_replicated(self) -> bool: + return self.shape == (1,) * self.ndim + + # sharding.Sharding interface + + @property + def _device_assignment(self) -> XLADeviceAssignment: + return self._devices + + def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding: + return _positional_sharding_to_xla_hlo_sharding(self, num_dimensions) + + @functools.cached_property + def is_fully_addressable(self) -> bool: + return self._internal_device_list.is_fully_addressable + + +class DeviceIdSet: + _name: str + _ids: frozenset[int] + def __init__(self, name, *ids): + self._name = name + self._ids = frozenset(ids) + + def __iter__(self): + return iter(sorted(self._ids)) + + def __add__(self, other) -> DeviceIdSet: + assert isinstance(other, DeviceIdSet) + return DeviceIdSet(self._name, *(self._ids | other._ids)) + + def __len__(self) -> int: + return len(self._ids) + + def __repr__(self) -> str: + ids = ', '.join(safe_map(str, sorted(self._ids))) + return f'{{{self._name} {ids}}}' + + def __hash__(self) -> int: + return hash((self._name, self._ids)) + + def __eq__(self, other) -> bool: + return (isinstance(other, DeviceIdSet) and self._name == other._name and + self._ids == other._ids) + + +@use_cpp_class(xc.GSPMDSharding) +class GSPMDSharding(sharding.Sharding): + _devices: tuple[Device, ...] + _hlo_sharding: xc.HloSharding + _memory_kind: str | None + _device_list: xc.DeviceList | None + _internal_device_list: xc.DeviceList + + @use_cpp_method() + def __init__(self, devices: Sequence[Device], + op_sharding: xc.OpSharding | xc.HloSharding, + *, memory_kind: str | None = None, + _device_list: xc.DeviceList | None = None): + self._devices = tuple(devices) + if isinstance(op_sharding, xc.OpSharding): + self._hlo_sharding = xc.HloSharding.from_proto(op_sharding) + else: + self._hlo_sharding = op_sharding + self._memory_kind = memory_kind + + def __reduce__(self): + return (type(self), (self._devices, self._hlo_sharding.to_proto()), + {'memory_kind': self._memory_kind}) + + @functools.cached_property + def _hlo_sharding_hash(self): + if self.is_fully_replicated: + return hash(get_replicated_hlo_sharding()) + return hash(self._hlo_sharding) + + def __eq__(self, other): + if not isinstance(other, GSPMDSharding): + return False + if self is other: + return True + return (are_op_shardings_equal(self._hlo_sharding, other._hlo_sharding) + and self.memory_kind == other.memory_kind + and self._internal_device_list == other._internal_device_list) + + def __hash__(self): + if not hasattr(self, '_hash'): + self._hash = hash((self._internal_device_list, self._hlo_sharding_hash, + self.memory_kind)) + return self._hash + + def __repr__(self): + mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}' + return f'GSPMDSharding({self._hlo_sharding!r}{mem})' + + def check_compatible_aval(self, aval_shape: Shape) -> None: + num_ways_dim_sharded, _ = get_num_ways_dim_sharded(self._hlo_sharding) + if len(aval_shape) < len(num_ways_dim_sharded): + raise ValueError( + f"Sharding {self} is only valid for values of rank at least " + f"{len(num_ways_dim_sharded)}, but was applied to a value of rank " + f"{len(aval_shape)}") + + @functools.cached_property + def device_set(self) -> set[Device]: + return set(self._devices) + + @property + def memory_kind(self) -> str | None: + return self._memory_kind + + def with_memory_kind(self, kind: str) -> GSPMDSharding: + return GSPMDSharding(self._devices, self._hlo_sharding, memory_kind=kind) + + @property + def _device_assignment(self) -> XLADeviceAssignment: + return self._devices + + def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding: + return self._hlo_sharding + + @functools.cached_property + def is_fully_replicated(self) -> bool: + return is_op_sharding_replicated(self._hlo_sharding) + + @functools.cached_property + def is_fully_addressable(self) -> bool: + return self._internal_device_list.is_fully_addressable + + @classmethod + def get_replicated(cls, device_assignment, *, memory_kind: str | None = None): + return cls(tuple(device_assignment), get_replicated_hlo_sharding(), + memory_kind=memory_kind) + + +class AUTO: + + def __init__(self, mesh: mesh_lib.Mesh): + self.mesh = mesh + + +def is_auto(x): + return isinstance(x, AUTO) + + +class UnspecifiedValue: + def __repr__(self): + return "UnspecifiedValue" +UNSPECIFIED = UnspecifiedValue() + +def is_unspecified(x): + return isinstance(x, UnspecifiedValue) + +def is_unspecified_or_auto(x): + return is_auto(x) or is_unspecified(x) + + +MeshAxisName = Any + +""" +ArrayMapping specifies how an ndarray should map to mesh axes. + +Note that the ordering is crucial for the cases when this mapping is non-injective +(i.e. when multiple mesh axes map to the same positional axis). Then, the +order of entries of the mapping determines a major-to-minor order on mesh axes, +according to which chunks of the value along the repeated dimension will be assigned. + +For example, consider a mapping {'x': 1, 'y': 1} and a mesh with shape {'x': 2, 'y': 3}. +The second dimension of the value would get chunked into 6 pieces, and assigned to the +mesh in a way that treats 'y' as the fastest changing (minor) dimension. In this case, +that would mean that a flat list of chunks would get assigned to a flattened list of +mesh devices without any modifications. If the mapping was {'y': 1, 'x': 1}, then the +mesh devices ndarray would have to be transposed before flattening and assignment. +""" +ArrayMapping = OrderedDict[MeshAxisName, int] +ArrayMappingOrAutoOrUnspecified = Union[ArrayMapping, AUTO, UnspecifiedValue] + +def array_mapping_to_axis_resources(array_mapping: ArrayMapping): + if not array_mapping: + return PartitionSpec() + max_index = -1 + reverse_map = collections.defaultdict(list) + for axis, index in array_mapping.items(): + reverse_map[index].append(axis) + if index > max_index: + max_index = index + partitions = [] + for i in range(max_index + 1): + axis = reverse_map[i] + if axis: + partitions.append(axis[0] if len(axis) == 1 else tuple(axis)) + else: + partitions.append(None) + return PartitionSpec(*partitions) + +def get_array_mapping( + axis_resources: ParsedPartitionSpec | AUTO | UnspecifiedValue +) -> ArrayMappingOrAutoOrUnspecified: + # TODO(yashkatariya): Use `TypeGuard` on `is_auto` when it is supported. + # Don't use `is_auto` here to satisfy pytype and mypy. + if isinstance(axis_resources, (AUTO, UnspecifiedValue)): + return axis_resources + return OrderedDict((axis, i) + for i, axes in enumerate(axis_resources) + if axes is not None for axis in axes) + + +get_single_pspec = lambda p: array_mapping_to_axis_resources( + cast(ArrayMapping, get_array_mapping(p))) + + +class SpecSync(enum.IntEnum): + """Encodes how much out of sync the real value of partitions is compared to the user specified one. + + We use this to make sure we don't show garbage modified values while claiming + that the users have specified them like that. + """ + OUT_OF_SYNC = 0 # Arbitrary changes, including new axes inserted + DIM_PERMUTE = 1 # Dimensions permuted, but no new sharding axes + IN_SYNC = 2 # Entirely in sync + +class ParsedPartitionSpec: + __slots__ = ('unsafe_user_spec', 'partitions', 'sync') + + def __init__(self, user_spec, partitions, sync=SpecSync.IN_SYNC): + self.unsafe_user_spec = user_spec + # None in partitions represents unconstrained dim. + # TODO(yashkatariya): May use a sentinel value. + self.partitions = tuple(partitions) + self.sync = sync + + @property + def user_spec(self): + return self.unsynced_user_spec(SpecSync.IN_SYNC) + + def get_partition_spec(self) -> PartitionSpec: + if self.sync < SpecSync.IN_SYNC: + return get_single_pspec(self) + else: + if isinstance(self.unsafe_user_spec, PartitionSpec): + return self.unsafe_user_spec + else: + return get_single_pspec(self) + + def unsynced_user_spec(self, min_sync): + if self.sync < min_sync: + raise AssertionError(f"Please open a bug report! ({self.sync} >= {min_sync})") + return self.unsafe_user_spec + + def insert_axis_partitions(self, dim, val): + parts = self.partitions + too_short = dim - len(parts) + if too_short > 0: + parts += ((),) * too_short + new_partitions = util.tuple_insert(parts, dim, val) + new_sync = SpecSync.DIM_PERMUTE if (val == () or val is None) else SpecSync.OUT_OF_SYNC + return ParsedPartitionSpec(self.unsafe_user_spec, new_partitions, sync=new_sync) + + @classmethod + def from_user_input(cls, entry, arg_name, allow_unconstrained_dims=False): + if entry is None: + return cls(entry, ()) + if not isinstance(entry, PartitionSpec): + raise TypeError(f"{arg_name} are expected to be " + f"PartitionSpec instances or None, but got {entry}") + axis_specs = [] + for axis_spec in entry: + if axis_spec is None: + axis_spec = () + elif isinstance(axis_spec, (list, tuple)): + axis_spec = tuple(axis_spec) + elif axis_spec == PartitionSpec.UNCONSTRAINED: + if not allow_unconstrained_dims: + raise ValueError(f"Unconstrained dims are not allowed: {entry}") + axis_spec = None + else: + axis_spec = (axis_spec,) + axis_specs.append(axis_spec) + new_entry = PartitionSpec( + *[tuple(e) if isinstance(e, (list, tuple)) else e for e in entry]) + return cls(new_entry, axis_specs) + + def __hash__(self): + return hash((self.partitions, self.sync)) + + def __eq__(self, other): + return (self.partitions == other.partitions and + self.sync == other.sync) + + def __len__(self): + return len(self.partitions) + + def __getitem__(self, i): + return self.partitions[i] + + def __iter__(self): + return iter(self.partitions) + + def __repr__(self): + return (f"ParsedPartitionSpec(partitions={self.partitions}, " + f"unsafe_user_spec={self.unsafe_user_spec}, " + f"sync={self.sync})") + +class CanonicalizedParsedPartitionSpec(ParsedPartitionSpec): + """ParsedPartitionSpecs that are canonicalized. + + ParsedPartitionSpecs may contain trailing empty tuples, that make them + semantically different in general, and yet in some situations we prefer + to regard them as equivalent. For example, partitions of () and ((),) + cannot be always considered equivalent, since the first one is a valid + spec for a scalar value, while the second is not! However, when either of + those are applied to a 2D array, they both mean that the array is fully + replicated. + + So CanonicalizedParsedPartitionSpecs removes the trailing empty tuples from + partitions. + """ + + def __init__(self, parsed_pspec: ParsedPartitionSpec): + partitions = list(parsed_pspec.partitions) + while partitions and partitions[-1] == (): + partitions.pop() + + super().__init__(parsed_pspec.unsafe_user_spec, partitions, + parsed_pspec.sync) + + def __repr__(self): + return (f"CanonicalizedParsedPartitionSpec(partitions={self.partitions}, " + f"unsafe_user_spec={self.unsafe_user_spec}, " + f"sync={self.sync})") + + +def preprocess(mesh, spec, parsed_pspec, _manual_axes=frozenset()): + # This split exists because you can pass `_parsed_pspec` that has been + # modified from the original. For example: Adding extra dimension to + # axis_resources for vmap handlers. In such cases you need to preserve the + # `sync` attribute of parsed pspecs. + # PartitionSpec is inferred from the parsed pspec in this case. + # TODO(yaskatariya): Remove this and replace this with a normalized + # representation of Parsed Pspec + if parsed_pspec is None: + parsed_pspec = prepare_axis_resources( + PartitionSpec() if spec is None else spec, + "NamedSharding spec", allow_unconstrained_dims=True) + + _check_mesh_resource_axis(mesh, parsed_pspec, _manual_axes) + return parsed_pspec + +# fallback for c++ . +preprocess_with_manual = preprocess + +def prepare_axis_resources(axis_resources, + arg_name, + allow_unconstrained_dims=False): + # PyTrees don't treat None values as leaves, so we use an is_leaf function. + entries, treedef = tree_util.tree_flatten( + axis_resources, is_leaf=lambda x: x is None) + what = f"{arg_name} leaf specifications" + + new_entries = [] + for entry in entries: + if is_unspecified_or_auto(entry) or entry is None: + new_entries.append(entry) + elif isinstance(entry, sharding.Sharding): + if isinstance(entry, PmapSharding): + raise ValueError(f'One of {what} got sharding {entry} which is not ' + 'allowed.') + if xla_extension_version < 270: + if not isinstance(entry, XLACompatibleSharding): + raise ValueError(f'One of {what} got sharding {entry} which is not a ' + 'subclass of XLACompatibleSharding.') + new_entries.append(entry) + else: + new_entries.append(ParsedPartitionSpec.from_user_input( + entry, what, allow_unconstrained_dims=allow_unconstrained_dims)) + + _check_unique_resources(new_entries, arg_name) + return tree_util.tree_unflatten(treedef, new_entries) + + +def _check_unique_resources(axis_resources, arg_name): + for arg_axis_resources in axis_resources: + if not arg_axis_resources: continue + if (is_unspecified_or_auto(arg_axis_resources) or + isinstance(arg_axis_resources, sharding.Sharding)): + continue + constrained_dims = [d for d in arg_axis_resources if d is not None] + resource_counts = collections.Counter( + itertools.chain.from_iterable(constrained_dims)) + if not resource_counts: continue + if resource_counts.most_common(1)[0][1] > 1: + multiple_uses = [r for r, c in resource_counts.items() if c > 1] + if multiple_uses: + raise ValueError(f"A single {arg_name} specification can map every mesh axis " + f"to at most one positional dimension, but {arg_axis_resources.user_spec} " + f"has duplicate entries for {mesh_lib.show_axes(multiple_uses)}") + +# Axis environments + +class AxisEnv(NamedTuple): + """Represents a pmap mesh (only along the replica axes).""" + nreps: int + names: tuple[Any, ...] + sizes: tuple[int, ...] + + +@dataclasses.dataclass(frozen=True) +class SPMDAxisContext: + """A hardware axis context for parallel computations that use the GSPMD partitioner. + + This includes the mesh that will later by used to execute this computation, + as well as a set of mesh axes that are currently (e.g. because the current lowering + is invoked inside an xmap) lowered in the MANUAL sharding mode. + """ + mesh: mesh_lib.Mesh + manual_axes: frozenset[MeshAxisName] = frozenset() + + @property + def axis_env(self): + # All collectives that touch axis_env should remember to set use_global_device_ids + # when this context is enabled! + if self.manual_axes != frozenset(self.mesh.axis_names): + raise NotImplementedError( + "Collectives in manually partitioned computations are only supported " + "when all mesh axes are partitioned manually (no partial automatic sharding). " + "Make sure that you mention all mesh axes in axis_resources!") + return self.unsafe_axis_env + + @property + def unsafe_axis_env(self): + return AxisEnv( + nreps=self.mesh.size, + names=self.mesh.axis_names, + sizes=tuple(self.mesh.shape.values())) + + def extend_manual(self, axes: frozenset[MeshAxisName]) -> SPMDAxisContext: + return SPMDAxisContext(self.mesh, self.manual_axes | axes) + + +@dataclasses.dataclass(frozen=True) +class ReplicaAxisContext: + """A hardware axis context for parallel computations that are partitioned by JAX. + + Unlike in the SPMDAxisContext, this means that JAX might need to emit calls to + explicit collectives. + """ + axis_env: AxisEnv + + +@dataclasses.dataclass(frozen=True) +class ShardingContext: + """A hardware axis context for parallel computations that use the sharding + interface. + + This context also uses the GSPMD partitioner. + """ + num_devices: int + device_assignment: tuple[xc.Device, ...] | None = None + + def __post_init__(self): + if self.device_assignment is not None: + assert isinstance(self.device_assignment, tuple) + assert self.num_devices == len(self.device_assignment) + + # Similar to SPMDContext as ShardingContext also uses the GSPMD partitioner. + @property + def axis_env(self): + return AxisEnv(nreps=1, names=(), sizes=()) + + +# -------------------- XLA OpSharding to PartitionSpec -------------------- +# Note that OpSharding is more expressive than PartitionSpecs, so it's not +# always possible to convert them, but the code below should at least +# support handle all cases when this is possible. + +def strides_for_sizes(sizes): + """Returns an array of strides for major-to-minor sizes.""" + return np.cumprod(sizes[::-1])[::-1] // np.asarray(sizes) + +def unflatten_array(named_sizes, assignment): + """Recovers the ordering of axis names based on a device assignment. + + The device assignments that this function can convert into axis orders + are of the form:: + + np.arange(np.prod(named_sizes.values())).transpose(...).flatten() + + for some transposition ``...``. This is satisfied by all OpSharding assignments + generated from partition specs. + + Arguments: + named_sizes: A dictionary mapping axis names to their sizes. + assignment: A permutation of integers between 0 and the product of all + named sizes. + + Returns: + A major-to-minor list of axis names that corresponds to the given assignment. + """ + named_sizes = {name: size for name, size in named_sizes.items() if size != 1} + sizes = np.fromiter(named_sizes.values(), dtype=np.int64) + strides = strides_for_sizes(sizes) + dims = explode_superdims(sizes, unflatten_superdims(assignment)) + dim_to_name = {(size, stride): name for size, stride, name in zip(sizes, strides, named_sizes)} + return [dim_to_name[d] for d in dims] + +def unflatten_superdims(assignment): + """Unflatten a list of dimension sizes and their strides that generates assignment. + + If this function succeeds for a given ``assignment``, then the following property + should be satisfied:: + + dims_with_strides = unflatten_superdims(assignment) + base_array = np.arange(map(fst, sorted(dims_with_strides, key=snd, reverse=True))) + assignment == base_array.transpose(argsort(dims_with_strides, key=snd, reverse=True)).flatten() + + That is, the returned dimensions list all sizes of the base array (with strides + indicating their initial order). The order of dimensions in the list corresponds + to the permutation that applied to the base array generates the assignment. + """ + def check(cond): + if cond: return + raise NotImplementedError("Failed to convert OpSharding into a ShardingSpec. " + "Please open a bug report!") + flat_assignment = np.asarray(assignment, dtype=np.int64) + check(flat_assignment[0] == 0) + dims = [] + while flat_assignment.size > 1: + stride = flat_assignment[1] + for i in range(len(flat_assignment)): + if flat_assignment[i] != i * stride: break + else: + # After this loop i should point to an "element after the sequence", so + # we have to increment it if the whole array is a strided sequence. + i += 1 + size = i + dims.append((size, stride)) + assert size > 1 # Ensure progress + flat_assignment = flat_assignment[::size] + return dims + +def explode_superdims(sizes, dims): + """Explode superdims to fit a known shape. + + The unflattening process might mistakenly generate too few too large dimensions. + For example, ``unflatten_superdims(np.arange(n))`` always returns ``[(n, 1)]``. + This function takes a list of such contiguous super-dimensions and splits them + into smaller dimensions such that:: + + set(map(fst, explode_superdims(sizes, dims))) == set(sizes) + """ + strides_to_sizes = {stride: size for size, stride in zip(sizes, strides_for_sizes(sizes))} + dims = list(reversed(dims)) + final_dims = [] + for size, stride in dims: + target_size = strides_to_sizes[stride] + new_dims = [] + while size > target_size: + assert target_size > 1 # Ensure progress + assert size % target_size == 0 + new_dims.append((target_size, stride)) + size //= target_size + stride *= target_size + target_size = strides_to_sizes[stride] + assert size == target_size + new_dims.append((size, stride)) + final_dims += reversed(new_dims) + return final_dims + +def parse_flatten_op_sharding(hlo_sharding: xc.OpSharding | xc.HloSharding, + mesh: mesh_lib.Mesh) -> Sequence[ParsedPartitionSpec]: + if isinstance(hlo_sharding, xc.OpSharding): + hlo_sharding = xc.HloSharding.from_proto(hlo_sharding) + if hlo_sharding.tuple_elements(): + out: list[ParsedPartitionSpec] = [] + for s in hlo_sharding.tuple_elements(): + out.extend(parse_flatten_op_sharding(s, mesh)) + return out + elif hlo_sharding.is_replicated(): + return [CanonicalizedParsedPartitionSpec( + ParsedPartitionSpec(PartitionSpec(), ()))] + elif hlo_sharding.is_tiled(): + mesh_shape = mesh.shape + mesh_axis_order = unflatten_array( + mesh.shape, hlo_sharding.tile_assignment_devices() + ) + mesh_axis = iter(mesh_axis_order) + shape = hlo_sharding.tile_assignment_dimensions() + partitions = [] + for dim_size in shape: + dim_partitions = [] + while dim_size > 1: + axis = next(mesh_axis) + axis_size = mesh_shape[axis] + assert dim_size % axis_size == 0 + dim_size //= axis_size + dim_partitions.append(axis) + partitions.append(tuple(dim_partitions)) + if len(hlo_sharding.subgroup_types()) > 1: + raise NotImplementedError( + 'Unhandled HloSharding type. Please open a bug report!' + ) + if hlo_sharding.replicate_on_last_tile_dim(): + partitions = partitions[:-1] + return [CanonicalizedParsedPartitionSpec( + ParsedPartitionSpec('', partitions))] + else: + raise AssertionError("Unhandled OpSharding type. Please open a bug report!") + + +def _slice_as_tuple(s: slice): + assert s.step is None + return (s.start, s.stop) + + +class NonUniformShardingError(ValueError): + """Raised when sharding is not uniform across processes.""" + + +def get_process_index_and_count( + tensor_sharding: sharding.Sharding, + dim: int, + ndims: int, +) -> tuple[int, int]: + """Get current process index and number of unique processes for given dimension. + + This function facilitates mapping of process-level data to individual + devices. Each process can use its index to obtain the data corresponding + to that index. If process level data is sharded on multiple dimensions + this function can be used to build the cross product of indices in + each sharded axis. Processes that need to load the same data will have + the same index. For shardings whose per-process data is not distributed + on a grid, the number of distinct shards will be such that it is possible to + build the target shape while maintaining a "cube" shape of local-process data. + + For example, in case of 4 hosts with sharding distributed like so: + + 1234 + 2143 + + For dim 0 (rows): all processes need to access all rows, so we return (0, 1) + For dim 1 (cols): + process 1 and 2 returns index 0 out of 2 (need cols 0 and 1), + process 3 and 4 returns index 1 out of 2 (need cols 2 and 3). + + On the other hand, for a sharding like: + + 1212 + 3434 + + Dim 0 (rows): process 1 and 2 returns (0, 2), process 3 and 4 returns (1, 2) + Dim 1 (cols): process 1 and 3 returns (0, 2), process 2 and 4 returns (1, 2) + + Note: This function requires sharding to be process uniform in dimension + `dim`: + each process has the same number of addressable indices in that + dimension and all index sets across processes are either disjoint or the same. + + For sharding to be process uniform the addressable shards doesn't need to + form contiguous subtensor, or even a sparse grid and in case of + interleaved high-dimensional tensor it is possible for sharding to be + process uniform only in some dimensions but not others. + + For example: + 1111 and 12 and 1212 and 1212 + 2222 21 2121 1212 + + are all sharding uniform, in both dimensions. However + + 1122 + 2121 + 1121 + 1222 + + is uniform in dimension 0 (both hosts access all rows), but + is not uniform in dimension 1 (host 1 accesses columns: 0, 1, and 3), + while host 2 accesses (0, 1, 2, 3). + + Returns: + A tuple of (index, num_distinct_shards) for the given dimension. + It is guaranteed that `index` will cover 0 to `num_distinct_shards - 1`, + across all processes. + + Raises: + NonUniformShardingError: if the sharding is not process uniform in dimension + `dim`. + """ + # TODO(sandler, yashkatariya): Consider making this function public. + + if ( + tensor_sharding.is_fully_addressable + or tensor_sharding.is_fully_replicated + ): + return (0, 1) + num_devices = len(tensor_sharding.device_set) + # Get device to indices map, we don't care about the concrete + # global shape here, only to get the distribution of shards across the tensor + # using (num_devices, num_devices, ...) This is a universal shape that is + # compatible with any mesh with num_devices. + device_map = tensor_sharding.devices_indices_map((num_devices,) * ndims) + + # Get the slices for 'dim' for all devices. + global_slice = {k: v[dim] for k, v in device_map.items()} + + # Contains mapping from process_index to a set of slices for that process. + process_to_slice = collections.defaultdict(set) + # Contains global set of slices across all processes. + all_slices = set() + + # Compute the set of slices for each process and the global set of slices. + for d, v in global_slice.items(): + key = (v.start, v.stop) + process_to_slice[d.process_index].add(key) + all_slices.add(key) + + # Get the set of slices for the current process which we will use to compute + # the index of the current process. + current_pid = next(iter(tensor_sharding.addressable_devices)).process_index + addressable_slices = frozenset(process_to_slice[current_pid]) + + # Verify that all processes have the same number of slices. + slices_per_process = len(addressable_slices) + if any(len(x) != slices_per_process for x in process_to_slice.values()): + raise NonUniformShardingError( + f'{tensor_sharding=} is non-uniform on {dim=} as some processes have ' + 'different number of slices.' + ) + unique_processes = list({frozenset(x) for x in process_to_slice.values()}) + + # After removing duplicate processes all unique slices should + # cover the dimension exactly once. If they don' it means that + # the sharding is not uniform. + if sum(len(h) for h in unique_processes) != len(all_slices): + raise NonUniformShardingError( + f'{tensor_sharding=} is non-uniform on {dim=}' + ) + return (unique_processes.index(addressable_slices), len(unique_processes)) + + +def local_to_global_shape( + sharding: sharding.Sharding, + local_shape: Shape, +) -> tuple[int | None, ...]: + """Computes the global shape given the per process if possible. + + The returned shape will have the size of the global tensor in that dimension + or None, if it is not computable. The latter can happen when sharding + is not uniform along that dimension, e.g. different hosts require + different shapes, or if different processes have partial data overlap. + + If at most one dimension is sharded the shape is always computable. + Generally, global shape is computable for most practical meshes (including + topology aware such as meshes returned by mesh_utils.create_device_mesh) + + Some examples: Suppose mesh is {'a': 2, 'b': 2, 'c': 2} with 2 devices + per host, 4 hosts total. For different specs we get: + - P(): + global_shape = local_shape + + - P(('a', 'b', 'c'), None): + global_shape = (4 * local_shape[0], local_shape[1]) + Note: per device shape is (local_shape[0] / 2, local_shape[1]) + + - P(('a', 'b'), None) + global_shape = (4 * local_shape[0], local_shape[1]) + # NB: the same global shape as above, since sharding along 'c' dimension + # happens to be within process, and thus doesn't affect the global shape. + # The underlying difference will be in the per *device* shape, which + # would be (local_shape[0], local_shape[1]) in this case. + + - P(None, ('a', 'c')) + global_shape = (local_shape[0], 2 * local_shape[1]) + # Per device shape is (local_shape[0], local_shape[1] / 2) + - P(('a', 'c'), 'b'): + global_shape = (2 * local_shape[0], 2 * local_shape[1]) + # Per device shape is (local_shape[0] / 2, local_shape[1]) + - If devices in the Mesh are randomly permuted: For any partition spec + which shards more than 1 axis: e.g. P('a', ('b', 'c')): + global_shape = (None, None) + + Args: + local_shape: global shape of the tensor. + + Returns: + global_shape with Nones in non-uniform dimensions. + """ + + global_shape : list[int | None] = [None] * len(local_shape) + for i, local_dim in enumerate(local_shape): + try: + _, shard_count = get_process_index_and_count( + sharding, i, ndims=len(local_shape) + ) + global_shape[i] = local_dim * shard_count + except NonUniformShardingError: + global_shape[i] = None + continue + + return tuple(global_shape) + + +def num_addressable_indices( + tensor_sharding: sharding.Sharding, + dim: int, + global_shape: Shape, +) -> int: + """Returns the number of indices for given dimension this host has access to. + + Each host can have multiple number of devices that are spanning + possibly discontiguous slices of data. This function computes the + total number of unique indices for dimension `dim` that any of its + addressable devices hold. + + In most cases the addressable indices form a sparse grid (and in some + cases a subcube), and thus each host will hold the same of number of + indices for each dimension. However, it is possible to design a mesh that + addressable shards form a complicated pattern. In that case, the returned + value is the number of indices that are addressable by at least one device. + + For example, suppose the sharding looks like this: (number indicates + the host index) + + 1221 + 1221 + 0000 + + Then on host 1 and 2, both dim 0 (rows), and dim=1 (cols) will have size 2, + while on host 0, dim 0 will have size 1, and dim 1 will have size 4. + + Args: + tensor_sharding: Sharding of the tensor. + dim: dimension along which to compute the number of addressable indices. + global_shape: global shape of the tensor. + + Returns: + The number of indices for dimension `dim` that this host holds. + """ + # TODO(sandler, yashkatariya): Consider making this function public. + addressables = tensor_sharding.addressable_devices_indices_map(global_shape) + addressables = cast(Mapping[sharding.Device, Index], addressables) + num_unique_slices = len({ + _slice_as_tuple(addressable[dim]) for addressable in addressables.values() + }) + shard_size = tensor_sharding.shard_shape(global_shape)[dim] + return shard_size * num_unique_slices + + +def physical_hlo_sharding(aval, hlo_sharding: xc.HloSharding) -> xc.HloSharding: + elt_aval = aval.dtype._rules.physical_element_aval(aval.dtype) + new_op_sharding = hlo_sharding.to_proto().clone() + partitions, num_replicas = get_num_ways_dim_sharded(hlo_sharding) + suffix = [] if num_replicas == 1 else [num_replicas] + tad = partitions + [1] * elt_aval.ndim + suffix + new_op_sharding.tile_assignment_dimensions = tad + return xc.HloSharding.from_proto(new_op_sharding) + +def is_single_device_sharding(sharding: sharding.Sharding) -> bool: + # Special case PmapSharding here because PmapSharding maps away an axis + # and needs to be handled separately.test_pjit_single_device_sharding_add + return len(sharding.device_set) == 1 and not isinstance(sharding, PmapSharding) + +def make_key_array_phys_sharding(aval, sharding): + if is_single_device_sharding(sharding): + return sharding + elif isinstance(sharding, PmapSharding): + elt_aval = aval.dtype._rules.physical_element_aval(aval.dtype) + trailing_sharding = [sharding_specs.NoSharding()] * elt_aval.ndim + phys_sharding_spec = sharding_specs.ShardingSpec( + sharding=(*sharding.sharding_spec.sharding, *trailing_sharding), + mesh_mapping=sharding.sharding_spec.mesh_mapping) + return PmapSharding(devices=sharding.devices, + sharding_spec=phys_sharding_spec) + elif isinstance(sharding, NamedSharding): + elt_aval = aval.dtype._rules.physical_element_aval(aval.dtype) + trailing_spec = [None] * elt_aval.ndim + return NamedSharding( + sharding.mesh, + PartitionSpec(*sharding.spec, *trailing_spec)) + else: + hlos = sharding._to_xla_hlo_sharding(aval.ndim) + return GSPMDSharding( + sharding._device_assignment, physical_hlo_sharding(aval, hlos)) + + +def physical_sharding( + aval, sharding: sharding.Sharding) -> sharding.Sharding: + return make_key_array_phys_sharding(aval, sharding) + + +def get_logical_gspmd_sharding(aval, phys_sharding): + elt_aval = aval.dtype._rules.physical_element_aval(aval.dtype) + phys_hlo_sharding = phys_sharding._to_xla_hlo_sharding( + aval.ndim + elt_aval.ndim) + partitions, num_replicas = get_num_ways_dim_sharded(phys_hlo_sharding) + suffix = [] if num_replicas == 1 else [num_replicas] + # Create logical sharding by cutting off the replicated trailing dims. + logical_op_sharding = phys_hlo_sharding.to_proto().clone() + tad = partitions[:-elt_aval.ndim] + suffix + logical_op_sharding.tile_assignment_dimensions = tad + return GSPMDSharding(phys_sharding._device_assignment, + xc.HloSharding.from_proto(logical_op_sharding)) + +def check_replicated_trailing_dims(sharding: sharding.Sharding, aval): + if isinstance(sharding, PmapSharding): + return + phys_aval = core.physical_aval(aval) + hlo_s = sharding._to_xla_hlo_sharding(phys_aval.ndim) + partitions, _ = get_num_ways_dim_sharded(hlo_s) + num_trailing_dims = phys_aval.ndim - aval.ndim + if not all(i == 1 for i in partitions[-num_trailing_dims:]): + raise AssertionError( + "The trailing dims of extended dtypes should be replicated. Got" + f" sharding: {sharding}, partitions: {partitions}, " + f"num_trailing_dims: {num_trailing_dims}") + +def logical_sharding(aval, phys_sharding) -> sharding.Sharding: + # The trailing dims should always be replicated. + check_replicated_trailing_dims(phys_sharding, aval) + + if is_single_device_sharding(phys_sharding): + return phys_sharding + elif isinstance(phys_sharding, PmapSharding): + elt_aval = aval.dtype._rules.physical_element_aval(aval.dtype) + logical_sharding_spec = sharding_specs.ShardingSpec( + sharding=phys_sharding.sharding_spec.sharding[:-elt_aval.ndim], + mesh_mapping=phys_sharding.sharding_spec.mesh_mapping) + return PmapSharding(devices=phys_sharding.devices, + sharding_spec=logical_sharding_spec) + elif isinstance(phys_sharding, NamedSharding): + logical_gs = get_logical_gspmd_sharding(aval, phys_sharding) + return _gspmd_to_named_sharding_via_mesh( + logical_gs, phys_sharding.mesh) + else: + return get_logical_gspmd_sharding(aval, phys_sharding) + + +@util.cache() +def create_mesh_pspec_sharding( + mesh: mesh_lib.Mesh, pspec: PartitionSpec | None, parsed_pspec=None, + memory_kind: str | None = None) -> NamedSharding: + if pspec is None: + pspec, parsed_pspec = PartitionSpec(), None + return NamedSharding(mesh, pspec, _parsed_pspec=parsed_pspec, + memory_kind=memory_kind) + + +def _gspmd_to_named_sharding_via_mesh( + out_s: GSPMDSharding, mesh: mesh_lib.Mesh) -> NamedSharding: + parsed_pspec = parse_flatten_op_sharding( + out_s._hlo_sharding, mesh)[0] + return create_mesh_pspec_sharding( + mesh, parsed_pspec.get_partition_spec(), parsed_pspec, + out_s.memory_kind) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding_specs.py b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding_specs.py new file mode 100644 index 0000000000000000000000000000000000000000..7092b51ab8944fac28f681bf9252724b21e2817c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/jax/_src/sharding_specs.py @@ -0,0 +1,206 @@ +# Copyright 2018 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# A ShardingSpec describes at a high level how a logical array is sharded across +# devices (each array sharded with a `PmapSharding` has a ShardingSpec, and +# ShardingSpecs also describe how to shard inputs to a parallel computation). +# spec_to_indices() encodes exactly how a given ShardingSpec is translated to +# device buffers, i.e. how the sharded array is "laid out" across devices. Given +# a sequence of devices, we shard the data across the devices in row-major +# order, with replication treated as an extra inner dimension. +# +# For example, given the logical data array [1, 2, 3, 4], if we were to +# partition this array 4 ways with a replication factor of 2, for a total of 8 +# devices, the data on each device would be: [1, 1], [2, 2], [3, 3], [4, 4]. +# +# This encoding is assumed by various parts of the system, e.g. generating +# replica groups for collective operations. + +from __future__ import annotations + +from collections.abc import Sequence +import itertools +import math +from typing import Union + +import numpy as np + +from jax._src import config +from jax._src import util +from jax._src.lib import pmap_lib + +unsafe_map, map = map, util.safe_map + +NoSharding = pmap_lib.NoSharding +Chunked = pmap_lib.Chunked +Unstacked = pmap_lib.Unstacked + +_UNSHARDED_INSTANCE = NoSharding() + +ShardedAxis = pmap_lib.ShardedAxis +Replicated = pmap_lib.Replicated +MeshDimAssignment = Union[ShardedAxis, Replicated] + +ShardingSpec = pmap_lib.ShardingSpec + + +def _sharding_spec_indices(self, shape: tuple[int, ...]) -> np.ndarray: + """Returns NumPy-style indices corresponding to a sharding spec. + + Args: + shape: The shape of the logical array being sharded. + + Returns: + An ndarray with the same shape as the logical mesh (as derived form + `mesh_mapping`). Each entry is a NumPy-style index selecting the subset of + the data array to be placed on a corresponding device. The indices can be + ints, slice objects with step=1, or tuples of those. + """ + assert len(shape) == len(self.sharding), (shape, self.sharding) + + axis_indices: list[Sequence[Index]] = [] + shard_indices_shape = [] + for dim, sharding in enumerate(self.sharding): + axis_size = shape[dim] + if isinstance(sharding, NoSharding): + axis_indices.append([slice(None)]) + # NOTE: We don't append unsharded dimensions to shard_indices_shape here, + # because they do not appear in the mesh mapping. + elif isinstance(sharding, Unstacked): + assert axis_size == sharding.size, f'{axis_size} != {sharding.size}' + axis_indices.append(range(axis_size)) + shard_indices_shape.append(axis_size) + elif isinstance(sharding, Chunked): + total_chunks = math.prod(sharding.chunks) + shard_size, ragged = divmod(axis_size, total_chunks) + assert not ragged, (axis_size, total_chunks, dim) + axis_indices.append([slice(i * shard_size, (i + 1) * shard_size) + for i in range(total_chunks)]) + shard_indices_shape.extend(sharding.chunks) + else: + util.assert_unreachable(sharding) + + # shard_indices is an ndarray representing the sharded axes of the logical array, + # with each dimension having size equal to the number of shards across the corresponding + # logical array dimension, and each element containing the multi-dimensional index that + # is used to extract the corresponding shard of the logical array. + shard_indices = np.empty([math.prod(shard_indices_shape)], dtype=np.object_) + for i, idxs in enumerate(itertools.product(*axis_indices)): + shard_indices[i] = idxs + shard_indices = shard_indices.reshape(shard_indices_shape) + + # Ensure that each sharded axis is used exactly once in the mesh mapping + num_sharded_dim = len(shard_indices_shape) + sharded_dim_perm = [a.axis for a in self.mesh_mapping if isinstance(a, ShardedAxis)] + assert (set(sharded_dim_perm) == set(range(num_sharded_dim)) and + len(sharded_dim_perm) == num_sharded_dim) + # Replicate/reorder the indices according to the mesh mapping + replica_sizes = tuple(a.replicas for a in self.mesh_mapping if isinstance(a, Replicated)) + replica_dim, sharded_dim = itertools.count(0), iter(sharded_dim_perm) + perm = [next(replica_dim) if isinstance(a, Replicated) else + len(replica_sizes) + next(sharded_dim) + for a in self.mesh_mapping] + return (np.broadcast_to(shard_indices, replica_sizes + shard_indices.shape) + .transpose(perm)) + +def _sharding_spec_repr(self): + return f'ShardingSpec({self.sharding}, {self.mesh_mapping})' + + +ShardingSpec.indices = _sharding_spec_indices +# mypy raises: error: Cannot assign to a method [assignment] +ShardingSpec.__repr__ = _sharding_spec_repr # type: ignore + + +Index = Union[int, slice, tuple[Union[int, slice], ...]] + +def spec_to_indices(shape: Sequence[int], + spec: ShardingSpec) -> tuple[Index, ...]: + """Returns numpy-style indices corresponding to a sharding spec. + + Each index describes a shard of the array. The order of the indices is the + same as the device_buffers of a Array sharded using PmapSharding (i.e. the + data is laid out row-major). + + Args: + shape: The shape of the logical array being sharded. + spec: Describes how the array is sharded and how the shards are assigned to + the logical mesh. + + Returns: + A tuple of length equal to the size of the mesh (inferred as the product of + sharded dimension sizes and all replication factors). Each element is an + int, a slice object with step=1, or a tuple thereof, to be treated as an + index into the full logical array. + """ + return tuple(spec.indices(shape).flat) # type: ignore + + +def pmap_sharding_spec(nrep, axis_size, sharded_shape: Sequence[int], + map_axis: int | None) -> ShardingSpec: + """Sharding spec for arguments or results of a pmap. + Args: + nrep: number of local XLA replicas (product of local axis sizes) + axis_size: local axis size for outer pmap + sharded_aval: the aval of the value inside the outer pmap, an instance of + a ShapedArray. + map_axis: the axis along which the value is mapped in the outer pmap + Returns: + A ShardingSpec. + """ + replication_factor, ragged = divmod(nrep, axis_size) + assert not ragged + pspec = ShardingSpec(sharding=[_UNSHARDED_INSTANCE] * len(sharded_shape), + mesh_mapping=()) + maybe_replicate = () if replication_factor == 1 else (Replicated(replication_factor),) + if map_axis is not None: + sharded_in_axis = sum(not isinstance(s, NoSharding) for s in pspec.sharding[:map_axis]) + def shift_sharded_axis(a: MeshDimAssignment): + if isinstance(a, ShardedAxis) and a.axis >= sharded_in_axis: + return ShardedAxis(a.axis + 1) + return a + # replication_factor represents the product of inner pmaps, so it goes + # after the outer pmapped axis at index 0 + if config.pmap_no_rank_reduction.value: + sharding = util.tuple_update( + pspec.sharding, map_axis, Chunked([axis_size])) + else: + sharding = util.tuple_insert( + pspec.sharding, map_axis, Unstacked(axis_size)) + return ShardingSpec( + sharding=sharding, + mesh_mapping=itertools.chain( + [ShardedAxis(sharded_in_axis)], maybe_replicate, + map(shift_sharded_axis, pspec.mesh_mapping))) + else: + return ShardingSpec( + sharding=pspec.sharding, + mesh_mapping=(Replicated(axis_size),) + maybe_replicate + pspec.mesh_mapping) + + +def create_pmap_sharding_spec(shape: tuple[int, ...], sharded_dim: int = 0, + sharded_dim_size: int | None = None): + if sharded_dim is not None: + if config.pmap_no_rank_reduction.value: + sharded_shape = util.tuple_update(shape, sharded_dim, 1) + else: + sharded_shape = util.tuple_delete(shape, sharded_dim) + if sharded_dim_size is None: + sharded_dim_size = shape[sharded_dim] + else: + assert sharded_dim_size is not None + sharded_shape = shape + + return pmap_sharding_spec(sharded_dim_size, sharded_dim_size, sharded_shape, + sharded_dim) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/INSTALLER b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/LICENSE b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..ec51537db27dd4d9c9ed3cd39fd96485f3cfddea --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/LICENSE @@ -0,0 +1,99 @@ +License agreement for matplotlib versions 1.3.0 and later +========================================================= + +1. This LICENSE AGREEMENT is between the Matplotlib Development Team +("MDT"), and the Individual or Organization ("Licensee") accessing and +otherwise using matplotlib software in source or binary form and its +associated documentation. + +2. Subject to the terms and conditions of this License Agreement, MDT +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that MDT's +License Agreement and MDT's notice of copyright, i.e., "Copyright (c) +2012- Matplotlib Development Team; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib . + +4. MDT is making matplotlib available to Licensee on an "AS +IS" basis. MDT MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, MDT MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. MDT SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between MDT and +Licensee. This License Agreement does not grant permission to use MDT +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib , +Licensee agrees to be bound by the terms and conditions of this License +Agreement. + +License agreement for matplotlib versions prior to 1.3.0 +======================================================== + +1. This LICENSE AGREEMENT is between John D. Hunter ("JDH"), and the +Individual or Organization ("Licensee") accessing and otherwise using +matplotlib software in source or binary form and its associated +documentation. + +2. Subject to the terms and conditions of this License Agreement, JDH +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that JDH's +License Agreement and JDH's notice of copyright, i.e., "Copyright (c) +2002-2011 John D. Hunter; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib. + +4. JDH is making matplotlib available to Licensee on an "AS +IS" basis. JDH MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, JDH MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. JDH SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between JDH and +Licensee. This License Agreement does not grant permission to use JDH +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib, +Licensee agrees to be bound by the terms and conditions of this License +Agreement. \ No newline at end of file diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/METADATA b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..32bc0eaae91deb727e6b94f31931e9544d6e9684 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/METADATA @@ -0,0 +1,215 @@ +Metadata-Version: 2.1 +Name: matplotlib +Version: 3.10.6 +Summary: Python plotting package +Author: John D. Hunter, Michael Droettboom +Author-Email: Unknown +License: License agreement for matplotlib versions 1.3.0 and later + ========================================================= + + 1. This LICENSE AGREEMENT is between the Matplotlib Development Team + ("MDT"), and the Individual or Organization ("Licensee") accessing and + otherwise using matplotlib software in source or binary form and its + associated documentation. + + 2. Subject to the terms and conditions of this License Agreement, MDT + hereby grants Licensee a nonexclusive, royalty-free, world-wide license + to reproduce, analyze, test, perform and/or display publicly, prepare + derivative works, distribute, and otherwise use matplotlib + alone or in any derivative version, provided, however, that MDT's + License Agreement and MDT's notice of copyright, i.e., "Copyright (c) + 2012- Matplotlib Development Team; All Rights Reserved" are retained in + matplotlib alone or in any derivative version prepared by + Licensee. + + 3. In the event Licensee prepares a derivative work that is based on or + incorporates matplotlib or any part thereof, and wants to + make the derivative work available to others as provided herein, then + Licensee hereby agrees to include in any such work a brief summary of + the changes made to matplotlib . + + 4. MDT is making matplotlib available to Licensee on an "AS + IS" basis. MDT MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR + IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, MDT MAKES NO AND + DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS + FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB + WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + + 5. MDT SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR + LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING + MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF + THE POSSIBILITY THEREOF. + + 6. This License Agreement will automatically terminate upon a material + breach of its terms and conditions. + + 7. Nothing in this License Agreement shall be deemed to create any + relationship of agency, partnership, or joint venture between MDT and + Licensee. This License Agreement does not grant permission to use MDT + trademarks or trade name in a trademark sense to endorse or promote + products or services of Licensee, or any third party. + + 8. By copying, installing or otherwise using matplotlib , + Licensee agrees to be bound by the terms and conditions of this License + Agreement. + + License agreement for matplotlib versions prior to 1.3.0 + ======================================================== + + 1. This LICENSE AGREEMENT is between John D. Hunter ("JDH"), and the + Individual or Organization ("Licensee") accessing and otherwise using + matplotlib software in source or binary form and its associated + documentation. + + 2. Subject to the terms and conditions of this License Agreement, JDH + hereby grants Licensee a nonexclusive, royalty-free, world-wide license + to reproduce, analyze, test, perform and/or display publicly, prepare + derivative works, distribute, and otherwise use matplotlib + alone or in any derivative version, provided, however, that JDH's + License Agreement and JDH's notice of copyright, i.e., "Copyright (c) + 2002-2011 John D. Hunter; All Rights Reserved" are retained in + matplotlib alone or in any derivative version prepared by + Licensee. + + 3. In the event Licensee prepares a derivative work that is based on or + incorporates matplotlib or any part thereof, and wants to + make the derivative work available to others as provided herein, then + Licensee hereby agrees to include in any such work a brief summary of + the changes made to matplotlib. + + 4. JDH is making matplotlib available to Licensee on an "AS + IS" basis. JDH MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR + IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, JDH MAKES NO AND + DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS + FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB + WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + + 5. JDH SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR + LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING + MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF + THE POSSIBILITY THEREOF. + + 6. This License Agreement will automatically terminate upon a material + breach of its terms and conditions. + + 7. Nothing in this License Agreement shall be deemed to create any + relationship of agency, partnership, or joint venture between JDH and + Licensee. This License Agreement does not grant permission to use JDH + trademarks or trade name in a trademark sense to endorse or promote + products or services of Licensee, or any third party. + + 8. By copying, installing or otherwise using matplotlib, + Licensee agrees to be bound by the terms and conditions of this License + Agreement. +Classifier: Development Status :: 5 - Production/Stable +Classifier: Framework :: Matplotlib +Classifier: Intended Audience :: Science/Research +Classifier: Intended Audience :: Education +Classifier: License :: OSI Approved :: Python Software Foundation License +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Topic :: Scientific/Engineering :: Visualization +Project-URL: Homepage, https://matplotlib.org +Project-URL: Download, https://matplotlib.org/stable/install/index.html +Project-URL: Documentation, https://matplotlib.org +Project-URL: Source Code, https://github.com/matplotlib/matplotlib +Project-URL: Bug Tracker, https://github.com/matplotlib/matplotlib/issues +Project-URL: Forum, https://discourse.matplotlib.org/ +Project-URL: Donate, https://numfocus.org/donate-to-matplotlib +Requires-Python: >=3.10 +Requires-Dist: contourpy>=1.0.1 +Requires-Dist: cycler>=0.10 +Requires-Dist: fonttools>=4.22.0 +Requires-Dist: kiwisolver>=1.3.1 +Requires-Dist: numpy>=1.23 +Requires-Dist: packaging>=20.0 +Requires-Dist: pillow>=8 +Requires-Dist: pyparsing>=2.3.1 +Requires-Dist: python-dateutil>=2.7 +Provides-Extra: dev +Requires-Dist: meson-python<0.17.0,>=0.13.1; extra == "dev" +Requires-Dist: pybind11!=2.13.3,>=2.13.2; extra == "dev" +Requires-Dist: setuptools_scm>=7; extra == "dev" +Requires-Dist: setuptools>=64; extra == "dev" +Description-Content-Type: text/markdown + +[![PyPi](https://img.shields.io/pypi/v/matplotlib)](https://pypi.org/project/matplotlib/) +[![Conda](https://img.shields.io/conda/vn/conda-forge/matplotlib)](https://anaconda.org/conda-forge/matplotlib) +[![Downloads](https://img.shields.io/pypi/dm/matplotlib)](https://pypi.org/project/matplotlib) +[![NUMFocus](https://img.shields.io/badge/powered%20by-NumFOCUS-orange.svg?style=flat&colorA=E1523D&colorB=007D8A)](https://numfocus.org) + +[![Discourse help forum](https://img.shields.io/badge/help_forum-discourse-blue.svg)](https://discourse.matplotlib.org) +[![Gitter](https://badges.gitter.im/matplotlib/matplotlib.svg)](https://gitter.im/matplotlib/matplotlib) +[![GitHub issues](https://img.shields.io/badge/issue_tracking-github-blue.svg)](https://github.com/matplotlib/matplotlib/issues) +[![Contributing](https://img.shields.io/badge/PR-Welcome-%23FF8300.svg?)](https://matplotlib.org/stable/devel/index.html) + +[![GitHub actions status](https://github.com/matplotlib/matplotlib/workflows/Tests/badge.svg)](https://github.com/matplotlib/matplotlib/actions?query=workflow%3ATests) +[![Azure pipelines status](https://dev.azure.com/matplotlib/matplotlib/_apis/build/status/matplotlib.matplotlib?branchName=main)](https://dev.azure.com/matplotlib/matplotlib/_build/latest?definitionId=1&branchName=main) +[![AppVeyor status](https://ci.appveyor.com/api/projects/status/github/matplotlib/matplotlib?branch=main&svg=true)](https://ci.appveyor.com/project/matplotlib/matplotlib) +[![Codecov status](https://codecov.io/github/matplotlib/matplotlib/badge.svg?branch=main&service=github)](https://app.codecov.io/gh/matplotlib/matplotlib) +[![EffVer Versioning](https://img.shields.io/badge/version_scheme-EffVer-0097a7)](https://jacobtomlinson.dev/effver) + +![Matplotlib logotype](https://matplotlib.org/_static/logo2.svg) + +Matplotlib is a comprehensive library for creating static, animated, and +interactive visualizations in Python. + +Check out our [home page](https://matplotlib.org/) for more information. + +![image](https://matplotlib.org/_static/readme_preview.png) + +Matplotlib produces publication-quality figures in a variety of hardcopy +formats and interactive environments across platforms. Matplotlib can be +used in Python scripts, Python/IPython shells, web application servers, +and various graphical user interface toolkits. + +## Install + +See the [install +documentation](https://matplotlib.org/stable/users/installing/index.html), +which is generated from `/doc/install/index.rst` + +## Contribute + +You've discovered a bug or something else you want to change — excellent! + +You've worked out a way to fix it — even better! + +You want to tell us about it — best of all! + +Start at the [contributing +guide](https://matplotlib.org/devdocs/devel/contribute.html)! + +## Contact + +[Discourse](https://discourse.matplotlib.org/) is the discussion forum +for general questions and discussions and our recommended starting +point. + +Our active mailing lists (which are mirrored on Discourse) are: + +- [Users](https://mail.python.org/mailman/listinfo/matplotlib-users) + mailing list: +- [Announcement](https://mail.python.org/mailman/listinfo/matplotlib-announce) + mailing list: +- [Development](https://mail.python.org/mailman/listinfo/matplotlib-devel) + mailing list: + +[Gitter](https://gitter.im/matplotlib/matplotlib) is for coordinating +development and asking questions directly related to contributing to +matplotlib. + +## Citing Matplotlib + +If Matplotlib contributes to a project that leads to publication, please +acknowledge this by citing Matplotlib. + +[A ready-made citation +entry](https://matplotlib.org/stable/users/project/citing.html) is +available. diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/RECORD b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..9a132b474f346fc7b6e226bc70f42bf2aac0c3e6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/RECORD @@ -0,0 +1,884 @@ +__pycache__/pylab.cpython-310.pyc,, +matplotlib-3.10.6.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +matplotlib-3.10.6.dist-info/LICENSE,sha256=WhqB6jAXKMi7opM9qDLAzWIina8giToCSrPVMkRGjbw,4830 +matplotlib-3.10.6.dist-info/METADATA,sha256=mPi22uTz5Eb6RMhoTcU9tPMAPv3h7xbPjEEb7g_vK7E,11268 +matplotlib-3.10.6.dist-info/RECORD,, +matplotlib-3.10.6.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib-3.10.6.dist-info/WHEEL,sha256=1nIT8bOU3dBEtO1OHNUw1PB7s17JH9tAQ93SLqU9JNM,85 +matplotlib/__init__.py,sha256=xsR9Hs8l15NPOMilXDuOEW1BqgjY7PqC9q6K9Ykhajo,55823 +matplotlib/__init__.pyi,sha256=s0YJ8fkGVioZL9bcBtyYgCq_MbY8LRrQUtmaD3XaIxQ,3439 +matplotlib/__pycache__/__init__.cpython-310.pyc,, +matplotlib/__pycache__/_afm.cpython-310.pyc,, +matplotlib/__pycache__/_animation_data.cpython-310.pyc,, +matplotlib/__pycache__/_blocking_input.cpython-310.pyc,, +matplotlib/__pycache__/_cm.cpython-310.pyc,, +matplotlib/__pycache__/_cm_bivar.cpython-310.pyc,, +matplotlib/__pycache__/_cm_listed.cpython-310.pyc,, +matplotlib/__pycache__/_cm_multivar.cpython-310.pyc,, +matplotlib/__pycache__/_color_data.cpython-310.pyc,, +matplotlib/__pycache__/_constrained_layout.cpython-310.pyc,, +matplotlib/__pycache__/_docstring.cpython-310.pyc,, +matplotlib/__pycache__/_enums.cpython-310.pyc,, +matplotlib/__pycache__/_fontconfig_pattern.cpython-310.pyc,, +matplotlib/__pycache__/_internal_utils.cpython-310.pyc,, +matplotlib/__pycache__/_layoutgrid.cpython-310.pyc,, +matplotlib/__pycache__/_mathtext.cpython-310.pyc,, +matplotlib/__pycache__/_mathtext_data.cpython-310.pyc,, +matplotlib/__pycache__/_pylab_helpers.cpython-310.pyc,, +matplotlib/__pycache__/_text_helpers.cpython-310.pyc,, +matplotlib/__pycache__/_tight_bbox.cpython-310.pyc,, +matplotlib/__pycache__/_tight_layout.cpython-310.pyc,, +matplotlib/__pycache__/_type1font.cpython-310.pyc,, +matplotlib/__pycache__/_version.cpython-310.pyc,, +matplotlib/__pycache__/animation.cpython-310.pyc,, +matplotlib/__pycache__/artist.cpython-310.pyc,, +matplotlib/__pycache__/axis.cpython-310.pyc,, +matplotlib/__pycache__/backend_bases.cpython-310.pyc,, +matplotlib/__pycache__/backend_managers.cpython-310.pyc,, +matplotlib/__pycache__/backend_tools.cpython-310.pyc,, +matplotlib/__pycache__/bezier.cpython-310.pyc,, +matplotlib/__pycache__/category.cpython-310.pyc,, +matplotlib/__pycache__/cbook.cpython-310.pyc,, +matplotlib/__pycache__/cm.cpython-310.pyc,, +matplotlib/__pycache__/collections.cpython-310.pyc,, +matplotlib/__pycache__/colorbar.cpython-310.pyc,, +matplotlib/__pycache__/colorizer.cpython-310.pyc,, +matplotlib/__pycache__/colors.cpython-310.pyc,, +matplotlib/__pycache__/container.cpython-310.pyc,, +matplotlib/__pycache__/contour.cpython-310.pyc,, +matplotlib/__pycache__/dates.cpython-310.pyc,, +matplotlib/__pycache__/dviread.cpython-310.pyc,, +matplotlib/__pycache__/figure.cpython-310.pyc,, +matplotlib/__pycache__/font_manager.cpython-310.pyc,, +matplotlib/__pycache__/gridspec.cpython-310.pyc,, +matplotlib/__pycache__/hatch.cpython-310.pyc,, +matplotlib/__pycache__/image.cpython-310.pyc,, +matplotlib/__pycache__/inset.cpython-310.pyc,, +matplotlib/__pycache__/layout_engine.cpython-310.pyc,, +matplotlib/__pycache__/legend.cpython-310.pyc,, +matplotlib/__pycache__/legend_handler.cpython-310.pyc,, +matplotlib/__pycache__/lines.cpython-310.pyc,, +matplotlib/__pycache__/markers.cpython-310.pyc,, +matplotlib/__pycache__/mathtext.cpython-310.pyc,, +matplotlib/__pycache__/mlab.cpython-310.pyc,, +matplotlib/__pycache__/offsetbox.cpython-310.pyc,, +matplotlib/__pycache__/patches.cpython-310.pyc,, +matplotlib/__pycache__/path.cpython-310.pyc,, +matplotlib/__pycache__/patheffects.cpython-310.pyc,, +matplotlib/__pycache__/pylab.cpython-310.pyc,, +matplotlib/__pycache__/pyplot.cpython-310.pyc,, +matplotlib/__pycache__/quiver.cpython-310.pyc,, +matplotlib/__pycache__/rcsetup.cpython-310.pyc,, +matplotlib/__pycache__/sankey.cpython-310.pyc,, +matplotlib/__pycache__/scale.cpython-310.pyc,, +matplotlib/__pycache__/spines.cpython-310.pyc,, +matplotlib/__pycache__/stackplot.cpython-310.pyc,, +matplotlib/__pycache__/streamplot.cpython-310.pyc,, +matplotlib/__pycache__/table.cpython-310.pyc,, +matplotlib/__pycache__/texmanager.cpython-310.pyc,, +matplotlib/__pycache__/text.cpython-310.pyc,, +matplotlib/__pycache__/textpath.cpython-310.pyc,, +matplotlib/__pycache__/ticker.cpython-310.pyc,, +matplotlib/__pycache__/transforms.cpython-310.pyc,, +matplotlib/__pycache__/typing.cpython-310.pyc,, +matplotlib/__pycache__/units.cpython-310.pyc,, +matplotlib/__pycache__/widgets.cpython-310.pyc,, +matplotlib/_afm.py,sha256=cWe1Ib37T6ZyHbR6_hPuzAjotMmi32y-kDB-i28iyqE,16692 +matplotlib/_animation_data.py,sha256=JJJbbc-fMdPjkbQ7ng9BHL5i91VTDHQVTtEdWOvWBAI,7986 +matplotlib/_api/__init__.py,sha256=7Fs57FtnCmbyr_NwuGWe0EhfXOYrBw22iKMU9onGc_k,13799 +matplotlib/_api/__init__.pyi,sha256=XNL-oGkk1MZPtSXk3rHlC3jLWZxsElmpNpOongKq8qA,2246 +matplotlib/_api/__pycache__/__init__.cpython-310.pyc,, +matplotlib/_api/__pycache__/deprecation.cpython-310.pyc,, +matplotlib/_api/deprecation.py,sha256=XJxksSV8ukS8kSVmsGPS0zgv1ffFa8WNaRLAk31dmOM,20091 +matplotlib/_api/deprecation.pyi,sha256=A8De57amX2GlZSrDYwVYCMxHFb8AsYGrH0k_bCxuJus,2217 +matplotlib/_blocking_input.py,sha256=VHNsxvX2mTx_xBknd30MSicVlRXS4dCDe9hDctbV5rk,1224 +matplotlib/_c_internal_utils.cp310-win_amd64.pyd,sha256=rSJRSnSyzTPyXR8IGNlAcQx8CVmSpaObBdkwIZ3QJ2U,147456 +matplotlib/_c_internal_utils.pyi,sha256=Z3bLs9pMGXrmZjt-4_A-x4321bLP-B54xDbr4PIgUfc,377 +matplotlib/_cm.py,sha256=PuYIAkUpz4u4aiUjvdV5njIfG0J_MQY9pM2Yz3j3KXs,68014 +matplotlib/_cm_bivar.py,sha256=gpmKiSxsWoVGWIifseh0goND7Y7zTWJKsznr_QkEtDg,97461 +matplotlib/_cm_listed.py,sha256=3a02mPUSnOsXkqRFNxKGwIvJCAKY5lezBUrqGcllnvk,135004 +matplotlib/_cm_multivar.py,sha256=0UjNFW7Sytj1t31QD99ebfr-al7NmCpWhKUOyq9VzK8,6630 +matplotlib/_color_data.py,sha256=k-wdTi6ArJxksqBfMT-7Uy2qWz8XX4Th5gsjf32CwmM,34780 +matplotlib/_color_data.pyi,sha256=RdBRk01yuf3jYVlCwG351tIBCxehizkZMnKs9c8gnOw,170 +matplotlib/_constrained_layout.py,sha256=XX_2elqHukF3toeEDupwQPlN_8UEv3-yGKlR0A92n74,31485 +matplotlib/_docstring.py,sha256=u9yJorJidI8k1W8S01SIMqNmh5VOMVc-AqyES66EEdk,4435 +matplotlib/_docstring.pyi,sha256=6ze5DoqZaFy6BQ5Z8vtDiJwSOXhBcpZJlkM5enI9cuE,800 +matplotlib/_enums.py,sha256=euD2sj2FIbQMzIPA4rCrII_y8RVzaMyEBKsSdBQaUb4,6175 +matplotlib/_enums.pyi,sha256=K7j_kDwGOnx37CYnXhwfJ9NJBGeet45KvAj0om05RUs,326 +matplotlib/_fontconfig_pattern.py,sha256=2livocARMbpys8tmicH6wlifcwcNoIstgruSY6DSAfk,4361 +matplotlib/_image.cp310-win_amd64.pyd,sha256=Ow28IJveeATwL9DxS22Q3QAoUVNSMqvOKIpBJzv7F54,348160 +matplotlib/_image.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/_internal_utils.py,sha256=nhK6LLWYW93fBcsFiO09JmqFj2rgHEsGYFOeaC7HRKw,2140 +matplotlib/_layoutgrid.py,sha256=mFB9asZVol2aV9hLBpdntuG5x1isKUrajF-3TUhCgso,21676 +matplotlib/_mathtext.py,sha256=kKFME5pnU2TZR_sD2Dm2KsnR0ldkmIgY3N6lgqYiSII,107848 +matplotlib/_mathtext_data.py,sha256=9y__7jf3bzgOmD2lEYW3QpgsT-z9QDuiwRLVH7Wq8Pw,65067 +matplotlib/_path.cp310-win_amd64.pyd,sha256=aTwAZt4gnb7k87bjoScJv47H57k9ahn967g0ocZu7ro,305152 +matplotlib/_path.pyi,sha256=yznyfzoUogH9vvi0vK68ga4Shlbrn5UBhAnLX8Ght1o,325 +matplotlib/_pylab_helpers.py,sha256=pJERytHDmXo2VP3sH9Qw6NwJP0hefNPQMUXCVntnoHI,4307 +matplotlib/_pylab_helpers.pyi,sha256=7OZKr-OL3ipVt1EDZ6e-tRwHASz-ijYfcIdlPczXhvQ,1012 +matplotlib/_qhull.cp310-win_amd64.pyd,sha256=qabLFf8M7Fi00wvVUMxR2EHvmJJqrPzeOQkOE4DCLlA,537088 +matplotlib/_qhull.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/_text_helpers.py,sha256=nz1pvEp8756PCrn6KIY2M6vVA-beYUNODvvf4NnGF0g,2538 +matplotlib/_tight_bbox.py,sha256=ddJ5ViXulbPLocHr-RkWK27WJSuV2WXUx74jZyL0NOg,2787 +matplotlib/_tight_layout.py,sha256=A3vZKBmxci7w35OO9lxgfUqrRKgTr-N_dUza0M05WxE,12675 +matplotlib/_tri.cp310-win_amd64.pyd,sha256=8_l35zrsrvqcJgHpYWc_XxQv7UbIaTXZCgM-g9x32YM,288256 +matplotlib/_tri.pyi,sha256=bFnu97pIlSboJT3ijQBw4YlZdlR1UzXYtI3DkEPVzeM,1429 +matplotlib/_type1font.py,sha256=c_PXvyOFQ4YCRQJ23snWNkdEhuZ4k0JbLe55zaeoXfQ,28409 +matplotlib/_version.py,sha256=VC9SRUcNAqaFBUs-eTHhFGNx5sumClw4yPTuFOTUDeE,19 +matplotlib/animation.py,sha256=58vUSy_0M2Ve2GO5JcXsCHxCwr7TqaaMed-KNxN6Nno,73114 +matplotlib/animation.pyi,sha256=C8Fgmswh4um-qYgqPZNBEpXVaefRXi6e026pYjonLHg,6566 +matplotlib/artist.py,sha256=hfCGETRkHLfP-i4PmmNsc7syXyFTZi4HYvzyE8q9ak0,63293 +matplotlib/artist.pyi,sha256=9sR0w8Kvi1GaFV6IbUzOw9LvQkPEB8FP_uRvUcI6_hg,7336 +matplotlib/axes/__init__.py,sha256=aHE_zIjphIJkW4_1fyoFuEWTb7gCRewK5jNZWJHEdgM,351 +matplotlib/axes/__init__.pyi,sha256=HP1z2v-PboHQS4dQjvJ7XjUjX-zw6taZRTTB9oVKwYE,303 +matplotlib/axes/__pycache__/__init__.cpython-310.pyc,, +matplotlib/axes/__pycache__/_axes.cpython-310.pyc,, +matplotlib/axes/__pycache__/_base.cpython-310.pyc,, +matplotlib/axes/__pycache__/_secondary_axes.cpython-310.pyc,, +matplotlib/axes/_axes.py,sha256=CwBUFh4fIlJF4Km80apFeH9sPVBi-3OPTOnxqKveZ5Q,353841 +matplotlib/axes/_axes.pyi,sha256=fvlVj8AnKt52EOCgxJ5N_6_lb9n0gmXlOP7WGf2puNI,26116 +matplotlib/axes/_base.py,sha256=eXarXIkk1SiQRPa7Rdu_R6gDgHdm2l60Hh2RPtPTf7E,186445 +matplotlib/axes/_base.pyi,sha256=EX6xYlHGEMRWXoE85w-UfCkN7m990Gg5VFZJ44D0XmE,17058 +matplotlib/axes/_secondary_axes.py,sha256=Sew-LK57mDndUXb-9_VHzs5dJLQ6fh4nstZtfpZjui0,11887 +matplotlib/axes/_secondary_axes.pyi,sha256=GtU55YzLNN7XHMRQcAmcAxcdcN-FzHw1RAJpOiTcZFk,1414 +matplotlib/axis.py,sha256=bsPEyn4eMAnkTCbW4rxSm8ysPNYH2p_ur7RN_ZfZRps,104714 +matplotlib/axis.pyi,sha256=UO1-oCeHyB7Mrxxnuo0Qys2u_k8e6qTmz4B6-rHrWdw,10181 +matplotlib/backend_bases.py,sha256=5EaIYhLC7LdKrk0SYSNBiYW3G16KbbexndFhoOX7_jY,131953 +matplotlib/backend_bases.pyi,sha256=G6MqvmWf-y2SwhuhNF2qN8hmovxAkV1TJMbg4LzU9cM,16270 +matplotlib/backend_managers.py,sha256=RQheCO_cQBlaWsYMbAmswu0UPKU7bmLTI5LEFgotklA,11795 +matplotlib/backend_managers.pyi,sha256=agnuM0wiZRqSgqti2AgbKJijRLvEPNLOrSY8PEwLjFE,2253 +matplotlib/backend_tools.py,sha256=eCx84l8JBYayju5Ry1LUFyKLIrdgxmGUTJw72VNP1yo,33186 +matplotlib/backend_tools.pyi,sha256=lc3W6FcY0XhTb8Z8XAqQ6MBk_SV7WH41t7zjwAlQfmE,4122 +matplotlib/backends/__init__.py,sha256=JowJe-tDrUBMNJTiJATgiEuACpgdxsKnRCYa-nC255A,206 +matplotlib/backends/__pycache__/__init__.cpython-310.pyc,, +matplotlib/backends/__pycache__/_backend_gtk.cpython-310.pyc,, +matplotlib/backends/__pycache__/_backend_pdf_ps.cpython-310.pyc,, +matplotlib/backends/__pycache__/_backend_tk.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk3.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk3agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk3cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk4.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk4agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk4cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_macosx.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_mixed.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_nbagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_pdf.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_pgf.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_ps.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt5.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt5agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt5cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qtagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qtcairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_svg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_template.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_tkagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_tkcairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_webagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_webagg_core.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_wx.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_wxagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_wxcairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/qt_compat.cpython-310.pyc,, +matplotlib/backends/__pycache__/registry.cpython-310.pyc,, +matplotlib/backends/_backend_agg.cp310-win_amd64.pyd,sha256=BvL6mfVDdyOPpgfPqFe0u4sYyh_EGm57sARvSBtn9EM,479744 +matplotlib/backends/_backend_agg.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/backends/_backend_gtk.py,sha256=zNCmMOgjgZGfAIlxDU4ph02XMSSn0wIUoPQ4wsaeeAg,11274 +matplotlib/backends/_backend_pdf_ps.py,sha256=M0kjQpMlBoAim6Pdw8u3FqUIoXz4UwRv4mSzXPJPI44,5968 +matplotlib/backends/_backend_tk.py,sha256=zx9GZYgrcmTwDgUT3rPvU7YgW79sUbJIWQ-SJtEsH5s,44581 +matplotlib/backends/_macosx.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/backends/_tkagg.cp310-win_amd64.pyd,sha256=cI_Y274qy7UNEerHOG-3amSE9S_cbIg7IfjkIWyGCm4,162816 +matplotlib/backends/_tkagg.pyi,sha256=1yKnYKSgNFesrb0ZI5RQCyTO4zONgHW5MrJ9vq4HzyI,379 +matplotlib/backends/backend_agg.py,sha256=Py_rXRbDTP0ZSrNB2O5Tmh5V8G5V9ViXH_YV5PTX9xw,19888 +matplotlib/backends/backend_cairo.py,sha256=Ah5M85Ppa8zYGOtLOizW0gzltjs3b4nD56qFeQR795U,18618 +matplotlib/backends/backend_gtk3.py,sha256=ZHf4CKhDGJvl8J3p6Sx_ZJBTgVetZBbaKWfnXKvNe0c,22277 +matplotlib/backends/backend_gtk3agg.py,sha256=aFJw05L5V-FyFSj_20l_FuC5piPf8fhKzft_m1NMhYc,2463 +matplotlib/backends/backend_gtk3cairo.py,sha256=zPsJzVm750if2LiQ9ybzsbX0rBhkP05XVvR9Lhz65so,1392 +matplotlib/backends/backend_gtk4.py,sha256=lkcW3cEKdASQ4vD6VHBUXvZQZEGvj5uenCVg4fHof_g,23541 +matplotlib/backends/backend_gtk4agg.py,sha256=00i3qpIt9Tcf_S74GOWbeckiPlfVJoQ2pBbhXDMthF0,1262 +matplotlib/backends/backend_gtk4cairo.py,sha256=sqWm3WgfNO8EsBjqzsD4U4cAc0f4q5SufYx7ZJacDf0,1125 +matplotlib/backends/backend_macosx.py,sha256=Gjbu1MKXIZpCM_luMG6U7qDgwotcGSegCvYw-NNWxeQ,7397 +matplotlib/backends/backend_mixed.py,sha256=Gf_2BDjy94FZRXALm0X3xRpK_1l8bZGSYCqnrSjrjME,4696 +matplotlib/backends/backend_nbagg.py,sha256=Au9RHfRufpI0ngT4R0K0CUVtAMFi9Bg-YhDunlj_Lko,8000 +matplotlib/backends/backend_pdf.py,sha256=DfFKdD4-wc7PPpfjXh8MAaaDppXUvLwicaMhXKgRhGw,105320 +matplotlib/backends/backend_pgf.py,sha256=GT7UHkAtHtnZcyV0GxHtUL9TTre9dji7Q6qexyi4AUU,39567 +matplotlib/backends/backend_ps.py,sha256=JRZZ6hsPcCfZmb-HczqyBulgow15yvoaHQe3Zcl1Zd4,51991 +matplotlib/backends/backend_qt.py,sha256=FjwUd9E2Fu2gCUhAVcJjXC--sAJsAoLUEHpMpa-IueE,42312 +matplotlib/backends/backend_qt5.py,sha256=kzfoo2ksEGsiWAa2LGtZYzKvfzqJJWyGOohohcRAu1g,787 +matplotlib/backends/backend_qt5agg.py,sha256=Vh7H8kqWH4X8a3VX2XZ2Vze9srwJavkNHAZxdJUz_bk,352 +matplotlib/backends/backend_qt5cairo.py,sha256=Go2Y0GVkXh1xh6x4F255_e5Xbwwws-OiD1Fc0805E78,292 +matplotlib/backends/backend_qtagg.py,sha256=ZjPtp5wR6tZGjbngPXRdVXYRhiPPrc5C0q2DmtdRkpY,3413 +matplotlib/backends/backend_qtcairo.py,sha256=e3SUG50VGqo68eS_8ebTCVQPa4AaxLxuo1JiWX4TIWg,1770 +matplotlib/backends/backend_svg.py,sha256=RQvPE9pR0xSu55NIMajX1WBqmmXbUWWwdWr2qMG-63k,51000 +matplotlib/backends/backend_template.py,sha256=qhWvXGiPeO2jvH61L3NCRYLReo--d4WJrQYanIaEmiE,8012 +matplotlib/backends/backend_tkagg.py,sha256=z9gB16fq2d-DUNpbeSDDLWaYmc0Jz3cDqNlBKhnQg0c,592 +matplotlib/backends/backend_tkcairo.py,sha256=JaGGXh8Y5FwVZtgryIucN941Olf_Pn6f4Re7Vuxl1-c,845 +matplotlib/backends/backend_webagg.py,sha256=YzNscpsID8yu0XYO3NXiZLMPZnDKMAKdG_TO2Ya3g4k,11022 +matplotlib/backends/backend_webagg_core.py,sha256=RZisxLS_pocuBRr5_lqBnF91flQA5upI8PiP8-Zitl0,18748 +matplotlib/backends/backend_wx.py,sha256=Y3IM1DEtPTeXOqkyLIGmNuKqDReyHMTtHzqTuwdfPhM,51296 +matplotlib/backends/backend_wxagg.py,sha256=tzcwYyW34j4LPfHm9uhuHwepwZIcspi3y8oPC8FJkdk,1468 +matplotlib/backends/backend_wxcairo.py,sha256=TK-m3S0c1WipfKE2IpIPNeE4hoXPjfMvnWAzHpCXpFs,848 +matplotlib/backends/qt_compat.py,sha256=t_4aD4rg5maItXABLtVG6JIfuQwZnHwhMcjtz6AUwQs,5346 +matplotlib/backends/qt_editor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/backends/qt_editor/__pycache__/__init__.cpython-310.pyc,, +matplotlib/backends/qt_editor/__pycache__/_formlayout.cpython-310.pyc,, +matplotlib/backends/qt_editor/__pycache__/figureoptions.cpython-310.pyc,, +matplotlib/backends/qt_editor/_formlayout.py,sha256=QmqqqLO6waqeSGOKjDNUwjvon53Z7yqil5AfqfftDWY,20953 +matplotlib/backends/qt_editor/figureoptions.py,sha256=ljd7fCBiE0kk2ATNTaB5Dg8U6Gn7Lz4Vhi3XX-VbXeE,9851 +matplotlib/backends/registry.py,sha256=M2BHhBKbUJ8E2B2QExfVf8NbAK4ZFhJtmKRelBZZDh4,15480 +matplotlib/backends/web_backend/all_figures.html,sha256=44Y-GvIJbNlqQaKSW3kwVKpTxBSG1WsdYz3ZmYHlUsA,1753 +matplotlib/backends/web_backend/css/boilerplate.css,sha256=qui16QXRnQFNJDbcMasfH6KtN9hLjv8883U9cJmsVCE,2310 +matplotlib/backends/web_backend/css/fbm.css,sha256=wa4vNkNv7fQ_TufjJjecFZEzPMR6W8x6uXJga_wQILw,1456 +matplotlib/backends/web_backend/css/mpl.css,sha256=ruca_aA5kNnP-MZmLkriu8teVP1nIgwcFEpoB16j8Z4,1611 +matplotlib/backends/web_backend/css/page.css,sha256=ca3nO3TaPw7865PN5SlGJBTc2H3rBXQCMaFPywX29y4,1623 +matplotlib/backends/web_backend/ipython_inline_figure.html,sha256=wgSxUh3xpPAxOnZgSMnrhDM5hYncOfWRGgaCUezvedY,1311 +matplotlib/backends/web_backend/js/mpl.js,sha256=FAFf8huEmvymJbTHFUQ5pbQt48NtXmN_DNACP2dCDMw,24432 +matplotlib/backends/web_backend/js/mpl_tornado.js,sha256=Zs2Uzs7YUilG765nYvanCo-IK8HkHDtIum1KAq6bQ_w,302 +matplotlib/backends/web_backend/js/nbagg_mpl.js,sha256=F-By4ZjOSmwpNAkxUUxUk35qCjGlf0B28Z0aOyfpxDM,9514 +matplotlib/backends/web_backend/single_figure.html,sha256=wEBwF602JLHErBUEXiS6jXqmxYAIzHpa3MMFrnev6rs,1357 +matplotlib/bezier.py,sha256=ahJ1BRJoTbcxJoNmF1mCYc96mWs5jSM4DOK58jYaXkU,19049 +matplotlib/bezier.pyi,sha256=itubQt_U1IKnQ508bhy6J8LSpaspvZIeOmJqg7JvOmU,2586 +matplotlib/category.py,sha256=UsZ8rbADAH96YfpjD-N1DLTcS_cQwlszIIQ2ZRaiWm4,7377 +matplotlib/cbook.py,sha256=nF2hNC_ryiJcec58Cnp_0mAGSD0LfsnjbzUuCH4p-Os,80020 +matplotlib/cbook.pyi,sha256=YkPZz_bsTYQJQ_wKx-qbsJYLhy2t2nCyvLNjJyiRtpw,6037 +matplotlib/cm.py,sha256=5xav-BvJ4pWSVdvGql4oYc4EpA-R6OuUP5xjrZCMKvg,10350 +matplotlib/cm.pyi,sha256=Ot5j2zDHYTra4kKuTwxYp7X5TsNwg_cJuYa7BxmQz-A,939 +matplotlib/collections.py,sha256=5AcX3lTsKD1ObSOkLr-irzYPfdJxsNrgw5dAD886vCE,96311 +matplotlib/collections.pyi,sha256=uZ4fPdq34oghdEtS1UkY7P0oOE2oPUzvS0Cr09yktaQ,10775 +matplotlib/colorbar.py,sha256=-PwV3xkczcyPq9Ki_zfpQdP_1EqOFBzoeUSsMlUF2cE,60687 +matplotlib/colorbar.pyi,sha256=C1NzjMWFKm01rPUUr15Yjz1WB0pGakJQeatJO7G0mNs,4966 +matplotlib/colorizer.py,sha256=RzE0zjqB6elpketKidy6n-RRDbYF98speEIOIMfIFb4,25180 +matplotlib/colorizer.pyi,sha256=B_5-rEvNo4DblSifLodf3OVVEc3bTGsIYqqaQdU_4Tw,3308 +matplotlib/colors.py,sha256=h1YV59y4WvstOPC0x8eZzwHYItal7VjjVcyVS3niaSM,137303 +matplotlib/colors.pyi,sha256=XhCi3MtgSozOY5FATtPUHyKO4xMrvi7TvbCAl714Uxc,14908 +matplotlib/container.py,sha256=Y6v4j79gMk8QDYfrdOqbbJHH0BoOIO9enz4dtlaBSJU,4565 +matplotlib/container.pyi,sha256=DdthHVj1-bQQV3VcpMD1uIPLXVUKyKhWMXy8yCniK1I,1805 +matplotlib/contour.py,sha256=wpMtrMZEiDYPRuFbZ592HvQqn4NOBySfM17detKoenY,68391 +matplotlib/contour.pyi,sha256=X4DSH0u1YFUvKxuoAv8YbmVOJTJ-wbHig4lznzhsI0U,5300 +matplotlib/dates.py,sha256=xfcFY3g3AsW9isnEIKFGj4gXq0Xeoxox_3fxNCDQ01c,66306 +matplotlib/dviread.py,sha256=BqQg0fqXD_TmABJPSWDC-H2azt7B53NfOE1SC6mbXiI,42590 +matplotlib/dviread.pyi,sha256=5fps3GiPf1ibKY1TH6__kNEiJ9xhpVgzgait6jo09Hc,2139 +matplotlib/figure.py,sha256=CNPMl0QCxhr3h3WGcd271pE4Xtv4rapRWuUSFSNJSuc,141862 +matplotlib/figure.pyi,sha256=rROP8wvPmeSwTKoJwFaWsCHFpZjfgDk_7HE0wPfZQRg,14914 +matplotlib/font_manager.py,sha256=tSpOlZqiNUdIn3gQCFN_-nOLZ08WXBqj_unP6o1JyIw,57651 +matplotlib/font_manager.pyi,sha256=Jse1IshOdjhXHdjcvTxtos3rgsuECbRJzc264A-ZGus,5052 +matplotlib/ft2font.cp310-win_amd64.pyd,sha256=rnTLhMx9zZaHC-T7uuX3m_JqXbC2Lg2pFAHvqQd1jXw,915968 +matplotlib/ft2font.pyi,sha256=yuDEfdw-7VyFuAHts-sMRm4VKV7IxYresBxEQgNOaPM,9253 +matplotlib/gridspec.py,sha256=s6RkXEk4ORWsMHjXkK98ZvosZ5D-qgenVNk0hIvMwOw,29786 +matplotlib/gridspec.pyi,sha256=3IQi_Q5KdwyPT_J_FWcngJ7iUOOCH4PSm42USm16Ofg,5099 +matplotlib/hatch.py,sha256=tfEM0DQxyLkUE87lHduMNlE55GKINtAZo4Y_YsAmSJo,7453 +matplotlib/hatch.pyi,sha256=OEFkNe8TgrBli_MDMW99Tvhr44S4ySPaYQomjVXpUUA,2098 +matplotlib/image.py,sha256=MLFg06J2xllm1tmgDnlLYw8gTy7__NGnkf4SZP-3mSs,69765 +matplotlib/image.pyi,sha256=R88aGGu6wP0lFKRudaM7iNxGgq4kh9uX9QRaZSorTLQ,7066 +matplotlib/inset.py,sha256=fJV8VMMY07_97iJR3_RO892WAKE7myfV_Gqg-qK99vg,10154 +matplotlib/inset.pyi,sha256=c-IVW14bpqb3qu7bz0EabvNXJnZwJAXEXoQoIcrL8ns,968 +matplotlib/layout_engine.py,sha256=mwRNh13mIQ-wpjQVgBaB_LWKN1AiQTWldg-X9nLmVok,11433 +matplotlib/layout_engine.pyi,sha256=9cKFTFPCvwmKIygTUvrvOq1iWnFYme9Shniiv2QbC74,1788 +matplotlib/legend.py,sha256=5NDPbWqkiIkP1FJ3sLNqmew2jkNjawQf2G_BqZNG_u4,55343 +matplotlib/legend.pyi,sha256=hn-MNF3SPHtSUqIhwVXTebU_Nzk_wIh5iKgf7AEAZRg,5364 +matplotlib/legend_handler.py,sha256=ekhZTT1X9yQd7TNIRdu_vui5FR2CTmI0NWTFJTqF1N4,29872 +matplotlib/legend_handler.pyi,sha256=3VEfeioGIAxhd3mhg4PXETZjCKf4OlXL0jz1MAFGtos,7655 +matplotlib/lines.py,sha256=BG3xhTuzYDFjzf-k5CcQc_V7W4HlQl4CL7RCQQz0EC8,57920 +matplotlib/lines.pyi,sha256=3tG7tD8GZ8YPphw4nmkgS9SfdLV8icW36bv61jb9RlU,6081 +matplotlib/markers.py,sha256=g6ukerZ5n_sD_Enk0eJA4kkTvOOp8AfSGBq60PHot-E,33708 +matplotlib/markers.pyi,sha256=FFFBsvilnbd8-5L04U70kfVoyBwc18w_fZ6DysTj9p0,1678 +matplotlib/mathtext.py,sha256=yTlpUfnKxzrFl7XNjyuC-xMBBnxUufU6GVwQyhQAI2o,5104 +matplotlib/mathtext.pyi,sha256=RCVxYGQ_CJ6wC7v_HkqoouU2PhtcvlJ1ffIyxzAC-so,1045 +matplotlib/mlab.py,sha256=T71p4x0KJNVGrwUuoDdgcjWQucIQTuVrocOC98Zdwsc,30210 +matplotlib/mlab.pyi,sha256=mkR7wbJS9eCQfCFsUWoXnqrAy0kcE3cVqpFZCeOGC_Q,3583 +matplotlib/mpl-data/fonts/afm/cmex10.afm,sha256=blR3ERmrVBV5XKkAnDCj4NMeYVgzH7cXtJ3u59u9GuE,12070 +matplotlib/mpl-data/fonts/afm/cmmi10.afm,sha256=5qwEOpedEo76bDUahyuuF1q0cD84tRrX-VQ4p3MlfBo,10416 +matplotlib/mpl-data/fonts/afm/cmr10.afm,sha256=WDvgC_D3UkGJg9u-J0U6RaT02lF4oz3lQxHtg1r3lYw,10101 +matplotlib/mpl-data/fonts/afm/cmsy10.afm,sha256=AbmzvCVWBceHRfmRfeJ9E6xzOQTFLk0U1zDfpf3_MaM,8295 +matplotlib/mpl-data/fonts/afm/cmtt10.afm,sha256=4ji7_mTpeWMa93o_UHBWPKCnqsBfhJJNllat1lJArP4,6501 +matplotlib/mpl-data/fonts/afm/pagd8a.afm,sha256=jjFrigwkTpYLqa26cpzZvKQNBo-PuF4bmDVqaM4pMWw,17183 +matplotlib/mpl-data/fonts/afm/pagdo8a.afm,sha256=sgNQdeYyx8J-itGw9h31y95aMBiTCRvmNSPTXwwS7xg,17255 +matplotlib/mpl-data/fonts/afm/pagk8a.afm,sha256=ZUtfHPloNqcvGMHMxaKDSlshhOcjwheUx143RwpGdIU,17241 +matplotlib/mpl-data/fonts/afm/pagko8a.afm,sha256=Yj1wBg6Jsqqz1KBfhRoJ3ACR-CMQol8Fj_ZM5NZ1gDk,17346 +matplotlib/mpl-data/fonts/afm/pbkd8a.afm,sha256=Zl5o6J_di9Y5j2EpHtjew-_sfg7-WoeVmO9PzOYSTUc,15157 +matplotlib/mpl-data/fonts/afm/pbkdi8a.afm,sha256=JAOno930iTyfZILMf11vWtiaTgrJcPpP6FRTRhEMMD4,15278 +matplotlib/mpl-data/fonts/afm/pbkl8a.afm,sha256=UJqJjOJ6xQDgDBLX157mKpohIJFVmHM-N6x2-DiGv14,15000 +matplotlib/mpl-data/fonts/afm/pbkli8a.afm,sha256=AWislZ2hDbs0ox_qOWREugsbS8_8lpL48LPMR40qpi0,15181 +matplotlib/mpl-data/fonts/afm/pcrb8a.afm,sha256=6j1TS2Uc7DWSc-8l42TGDc1u0Fg8JspeWfxFayjUwi8,15352 +matplotlib/mpl-data/fonts/afm/pcrbo8a.afm,sha256=smg3mjl9QaBDtQIt06ko5GvaxLsO9QtTvYANuE5hfG0,15422 +matplotlib/mpl-data/fonts/afm/pcrr8a.afm,sha256=7nxFr0Ehz4E5KG_zSE5SZOhxRH8MyfnCbw-7x5wu7tw,15339 +matplotlib/mpl-data/fonts/afm/pcrro8a.afm,sha256=NKEz7XtdFkh9cA8MvY-S3UOZlV2Y_J3tMEWFFxj7QSg,15443 +matplotlib/mpl-data/fonts/afm/phvb8a.afm,sha256=NAx4M4HjL7vANCJbc-tk04Vkol-T0oaXeQ3T2h-XUvM,17155 +matplotlib/mpl-data/fonts/afm/phvb8an.afm,sha256=8e_myD-AQkNF7q9XNLb2m76_lX2TUr3a5wog_LIE1sk,17086 +matplotlib/mpl-data/fonts/afm/phvbo8a.afm,sha256=8fkBRmJ-SWY2YrBg8fFyjJyrJp8daQ6JPO6LvhM8xPI,17230 +matplotlib/mpl-data/fonts/afm/phvbo8an.afm,sha256=aeVRvV4r15BBvxuRJ0MG8ZHuH2HViuIiCYkvuapmkmM,17195 +matplotlib/mpl-data/fonts/afm/phvl8a.afm,sha256=IyMYM-bgl-gI6rG0EuZZ2OLzlxJfGeSh8xqsh0t-eJQ,15627 +matplotlib/mpl-data/fonts/afm/phvlo8a.afm,sha256=s12C-eNnIDHJ_UVbuiprjxBjCiHIbS3Y8ORTC-qTpuI,15729 +matplotlib/mpl-data/fonts/afm/phvr8a.afm,sha256=Kt8KaRidts89EBIK29X2JomDUEDxvroeaJz_RNTi6r4,17839 +matplotlib/mpl-data/fonts/afm/phvr8an.afm,sha256=lL5fAHTRwODl-sB5mH7IfsD1tnnea4yRUK-_Ca2bQHM,17781 +matplotlib/mpl-data/fonts/afm/phvro8a.afm,sha256=3KqK3eejiR4hIFBUynuSX_4lMdE2V2T58xOF8lX-fwc,17919 +matplotlib/mpl-data/fonts/afm/phvro8an.afm,sha256=Vx9rRf3YfasMY7tz-njSxz67xHKk-fNkN7yBi0X2IP0,17877 +matplotlib/mpl-data/fonts/afm/pncb8a.afm,sha256=aoXepTcDQtQa_mspflMJkEFKefzXHoyjz6ioJVI0YNc,16028 +matplotlib/mpl-data/fonts/afm/pncbi8a.afm,sha256=pCWW1MYgy0EmvwaYsaYJaAI_LfrsKmDANHu7Pk0RaiU,17496 +matplotlib/mpl-data/fonts/afm/pncr8a.afm,sha256=0CIB2BLe9r-6_Wl5ObRTTf98UOrezmGQ8ZOuBX5kLks,16665 +matplotlib/mpl-data/fonts/afm/pncri8a.afm,sha256=5R-pLZOnaHNG8pjV6MP3Ai-d2OTQYR_cYCb5zQhzfSU,16920 +matplotlib/mpl-data/fonts/afm/pplb8a.afm,sha256=3EzUbNnXr5Ft5eFLY00W9oWu59rHORgDXUuJaOoKN58,15662 +matplotlib/mpl-data/fonts/afm/pplbi8a.afm,sha256=X_9tVspvrcMer3OS8qvdwjFFqpAXYZneyCL2NHA902g,15810 +matplotlib/mpl-data/fonts/afm/pplr8a.afm,sha256=ijMb497FDJ9nVdVMb21F7W3-cu9sb_9nF0oriFpSn8k,15752 +matplotlib/mpl-data/fonts/afm/pplri8a.afm,sha256=8KITbarcUUMi_hdoRLLmNHtlqs0TtOSKqtPFft7X5nY,15733 +matplotlib/mpl-data/fonts/afm/psyr.afm,sha256=Iyt8ajE4B2Tm34oBj2pKtctIf9kPfq05suQefq8p3Ro,9644 +matplotlib/mpl-data/fonts/afm/ptmb8a.afm,sha256=bL1fA1NC4_nW14Zrnxz4nHlXJb4dzELJPvodqKnYeMg,17983 +matplotlib/mpl-data/fonts/afm/ptmbi8a.afm,sha256=-_Ui6XlKaFTHEnkoS_-1GtIr5VtGa3gFQ2ezLOYHs08,18070 +matplotlib/mpl-data/fonts/afm/ptmr8a.afm,sha256=IEcsWcmzJyjCwkgsw4o6hIMmzlyXUglJat9s1PZNnEU,17942 +matplotlib/mpl-data/fonts/afm/ptmri8a.afm,sha256=49fQMg5fIGguZ7rgc_2styMK55Pv5bPTs7wCzqpcGpk,18068 +matplotlib/mpl-data/fonts/afm/putb8a.afm,sha256=qMaHTdpkrNL-m4DWhjpxJCSmgYkCv1qIzLlFfM0rl40,21532 +matplotlib/mpl-data/fonts/afm/putbi8a.afm,sha256=g7AVJyiTxeMpNk_1cSfmYgM09uNUfPlZyWGv3D1vcAk,21931 +matplotlib/mpl-data/fonts/afm/putr8a.afm,sha256=XYmNC5GQgSVAZKTIYdYeNksE6znNm9GF_0SmQlriqx0,22148 +matplotlib/mpl-data/fonts/afm/putri8a.afm,sha256=i7fVe-iLyLtQxCfAa4IxdxH-ufcHmMk7hbCGG5TxAY4,21891 +matplotlib/mpl-data/fonts/afm/pzcmi8a.afm,sha256=wyuoIWEZOcoXrSl1tPzLkEahik7kGi91JJj-tkFRG4A,16250 +matplotlib/mpl-data/fonts/afm/pzdr.afm,sha256=MyjLAnzKYRdQBfof1W3k_hf30MvqOkqL__G22mQ5xww,9467 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier-Bold.afm,sha256=sIDDI-B82VZ3C0mI_mHFITCZ7PVn37AIYMv1CrHX4sE,15333 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier-BoldOblique.afm,sha256=zg61QobD3YU9UBfCXmvmhBNaFKno-xj8sY0b2RpgfLw,15399 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier-Oblique.afm,sha256=vRQm5j1sTUN4hicT1PcVZ9P9DTTUHhEzfPXqUUzVZhE,15441 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier.afm,sha256=Mdcq2teZEBJrIqVXnsnhee7oZnTs6-P8_292kWGTrw4,15335 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica-Bold.afm,sha256=i2l4gcjuYXoXf28uK7yIVwuf0rnw6J7PwPVQeHj5iPw,69269 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica-BoldOblique.afm,sha256=Um5O6qK11DXLt8uj_0IoWkc84TKqHK3bObSKUswQqvY,69365 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica-Oblique.afm,sha256=hVYDg2b52kqtbVeCzmiv25bW1yYdpkZS-LXlGREN2Rs,74392 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica.afm,sha256=23cvKDD7bQAJB3kdjSahJSTZaUOppznlIO6FXGslyW8,74292 +matplotlib/mpl-data/fonts/pdfcorefonts/Symbol.afm,sha256=P5UaoXr4y0qh4SiMa5uqijDT6ZDr2-jPmj1ayry593E,9740 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-Bold.afm,sha256=cQTmr2LFPwKQE_sGQageMcmFicjye16mKJslsJLHQyE,64251 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-BoldItalic.afm,sha256=pzWOdycm6RqocBWgAVY5Jq0z3Fp7LuqWgLNMx4q6OFw,59642 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-Italic.afm,sha256=bK5puSMpGT_YUILwyJrXoxjfj7XJOdfv5TQ_iKsJRzw,66328 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-Roman.afm,sha256=hhNrUnpazuDDKD1WpraPxqPWCYLrO7D7bMVOg-zI13o,60460 +matplotlib/mpl-data/fonts/pdfcorefonts/ZapfDingbats.afm,sha256=ZuOmt9GcKofjdOq8kqhPhtAIhOwkL2rTJTmZxAjFakA,9527 +matplotlib/mpl-data/fonts/pdfcorefonts/readme.txt,sha256=MRv8ppSITYYAb7lt5EOw9DWWNZIblfxsFhu5TQE7cpI,828 +matplotlib/mpl-data/fonts/ttf/DejaVuSans-Bold.ttf,sha256=sYS4njwQdfIva3FXW2_CDUlys8_TsjMiym_Vltyu8Wc,704128 +matplotlib/mpl-data/fonts/ttf/DejaVuSans-BoldOblique.ttf,sha256=bt8CgxYBhq9FHL7nHnuEXy5Mq_Jku5ks5mjIPCVGXm8,641720 +matplotlib/mpl-data/fonts/ttf/DejaVuSans-Oblique.ttf,sha256=zN90s1DxH9PdV3TeUOXmNGoaXaH1t9X7g1kGZel6UhM,633840 +matplotlib/mpl-data/fonts/ttf/DejaVuSans.ttf,sha256=P99pyr8GBJ6nCgC1kZNA4s4ebQKwzDxLRPtoAb0eDSI,756072 +matplotlib/mpl-data/fonts/ttf/DejaVuSansDisplay.ttf,sha256=ggmdz7paqGjN_CdFGYlSX-MpL3N_s8ngMozpzvWWUvY,25712 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-Bold.ttf,sha256=uq2ppRcv4giGJRr_BDP8OEYZEtXa8HKH577lZiCo2pY,331536 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-BoldOblique.ttf,sha256=ppCBwVx2yCfgonpaf1x0thNchDSZlVSV_6jCDTqYKIs,253116 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-Oblique.ttf,sha256=KAUoE_enCfyJ9S0ZLcmV708P3Fw9e3OknWhJsZFtDNA,251472 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono.ttf,sha256=YC7Ia4lIz82VZIL-ZPlMNshndwFJ7y95HUYT9EO87LM,340240 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif-Bold.ttf,sha256=w3U_Lta8Zz8VhG3EWt2-s7nIcvMvsY_VOiHxvvHtdnY,355692 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif-BoldItalic.ttf,sha256=2T7-x6nS6CZ2jRou6VuVhw4V4pWZqE80hK8d4c7C4YE,347064 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif-Italic.ttf,sha256=PnmU-8VPoQzjNSpC1Uj63X2crbacsRCbydlg9trFfwQ,345612 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif.ttf,sha256=EHJElW6ZYrnpb6zNxVGCXgrgiYrhNzcTPhuSGi_TX_o,379740 +matplotlib/mpl-data/fonts/ttf/DejaVuSerifDisplay.ttf,sha256=KRTzLkfHd8J75Wd6-ufbTeefnkXeb8kJfZlJwjwU99U,14300 +matplotlib/mpl-data/fonts/ttf/LICENSE_DEJAVU,sha256=11k43sCY8G8Kw8AIUwZdlPAgvhw8Yu8dwpdboVtNmw4,4816 +matplotlib/mpl-data/fonts/ttf/LICENSE_STIX,sha256=urPTHf7wf0g2JPL2XycR52BluOcnMnixwHHt4QQcmVk,5476 +matplotlib/mpl-data/fonts/ttf/STIXGeneral.ttf,sha256=FnN4Ax4t3cYhbWeBnJJg6aBv_ExHjk4jy5im_USxg8I,448228 +matplotlib/mpl-data/fonts/ttf/STIXGeneralBol.ttf,sha256=6FM9xwg_o0a9oZM9YOpKg7Z9CUW86vGzVB-CtKDixqA,237360 +matplotlib/mpl-data/fonts/ttf/STIXGeneralBolIta.ttf,sha256=mHiP1LpI37sr0CbA4gokeosGxzcoeWKLemuw1bsJc2w,181152 +matplotlib/mpl-data/fonts/ttf/STIXGeneralItalic.ttf,sha256=bPyzM9IrfDxiO9_UAXTxTIXD1nMcphZsHtyAFA6uhSc,175040 +matplotlib/mpl-data/fonts/ttf/STIXNonUni.ttf,sha256=Ulb34CEzWsSFTRgPDovxmJZOwvyCAXYnbhaqvGU3u1c,59108 +matplotlib/mpl-data/fonts/ttf/STIXNonUniBol.ttf,sha256=XRBqW3jR_8MBdFU0ObhiV7-kXwiBIMs7QVClHcT5tgs,30512 +matplotlib/mpl-data/fonts/ttf/STIXNonUniBolIta.ttf,sha256=pb22DnbDf2yQqizotc3wBDqFGC_g27YcCGJivH9-Le8,41272 +matplotlib/mpl-data/fonts/ttf/STIXNonUniIta.ttf,sha256=BMr9pWiBv2YIZdq04X4c3CgL6NPLUPrl64aV1N4w9Ug,46752 +matplotlib/mpl-data/fonts/ttf/STIXSizFiveSymReg.ttf,sha256=wYuH1gYUpCuusqItRH5kf9p_s6mUD-9X3L5RvRtKSxs,13656 +matplotlib/mpl-data/fonts/ttf/STIXSizFourSymBol.ttf,sha256=yNdvjUoSmsZCULmD7SVq9HabndG9P4dPhboL1JpAf0s,12228 +matplotlib/mpl-data/fonts/ttf/STIXSizFourSymReg.ttf,sha256=-9xVMYL4_1rcO8FiCKrCfR4PaSmKtA42ddLGqwtei1w,15972 +matplotlib/mpl-data/fonts/ttf/STIXSizOneSymBol.ttf,sha256=cYexyo8rZcdqMlpa9fNF5a2IoXLUTZuIvh0JD1Qp0i4,12556 +matplotlib/mpl-data/fonts/ttf/STIXSizOneSymReg.ttf,sha256=0lbHzpndzJmO8S42mlkhsz5NbvJLQCaH5Mcc7QZRDzc,19760 +matplotlib/mpl-data/fonts/ttf/STIXSizThreeSymBol.ttf,sha256=3eBc-VtYbhQU3BnxiypfO6eAzEu8BdDvtIJSFbkS2oY,12192 +matplotlib/mpl-data/fonts/ttf/STIXSizThreeSymReg.ttf,sha256=XFSKCptbESM8uxHtUFSAV2cybwxhSjd8dWVByq6f3w0,15836 +matplotlib/mpl-data/fonts/ttf/STIXSizTwoSymBol.ttf,sha256=MUCYHrA0ZqFiSE_PjIGlJZgMuv79aUgQqE7Dtu3kuo0,12116 +matplotlib/mpl-data/fonts/ttf/STIXSizTwoSymReg.ttf,sha256=_sdxDuEwBDtADpu9CyIXQxV7sIqA2TZVBCUiUjq5UCk,15704 +matplotlib/mpl-data/fonts/ttf/cmb10.ttf,sha256=B0SXtQxD6ldZcYFZH5iT04_BKofpUQT1ZX_CSB9hojo,25680 +matplotlib/mpl-data/fonts/ttf/cmex10.ttf,sha256=ryjwwXByOsd2pxv6WVrKCemNFa5cPVTOGa_VYZyWqQU,21092 +matplotlib/mpl-data/fonts/ttf/cmmi10.ttf,sha256=MJKWW4gR_WpnZXmWZIRRgfwd0TMLk3-RWAjEhdMWI00,32560 +matplotlib/mpl-data/fonts/ttf/cmr10.ttf,sha256=Tdl2GwWMAJ25shRfVe5mF9CTwnPdPWxbPkP_YRD6m_Y,26348 +matplotlib/mpl-data/fonts/ttf/cmss10.ttf,sha256=ffkag9BbLkcexjjLC0NaNgo8eSsJ_EKn2mfpHy55EVo,20376 +matplotlib/mpl-data/fonts/ttf/cmsy10.ttf,sha256=uyJu2TLz8QDNDlL15JEu5VO0G2nnv9uNOFTbDrZgUjI,29396 +matplotlib/mpl-data/fonts/ttf/cmtt10.ttf,sha256=YhHwmuk1mZka_alwwkZp2tGnfiU9kVYk-_IS9wLwcdc,28136 +matplotlib/mpl-data/images/back-symbolic.svg,sha256=Okj_ressZkfe6Ewv_o7GF5toc5qWCeFkQ2cHQ25BdVE,1532 +matplotlib/mpl-data/images/back.pdf,sha256=ZR7CJo_dAeCM-KlaGvskgtHQyRtrPIolc8REOmcoqJk,1623 +matplotlib/mpl-data/images/back.png,sha256=E4dGf4Gnz1xJ1v2tMygHV0YNQgShreDeVApaMb-74mU,380 +matplotlib/mpl-data/images/back.svg,sha256=Okj_ressZkfe6Ewv_o7GF5toc5qWCeFkQ2cHQ25BdVE,1532 +matplotlib/mpl-data/images/back_large.png,sha256=9A6hUSQeszhYONE4ZuH3kvOItM0JfDVu6tkfromCbsQ,620 +matplotlib/mpl-data/images/filesave-symbolic.svg,sha256=dMGXvLSOHPu44kiWgZx-B_My_tLWaP6J6GgxJfL4FW0,2049 +matplotlib/mpl-data/images/filesave.pdf,sha256=P1EPPV2g50WTt8UaX-6kFoTZM1xVqo6S2H6FJ6Zd1ec,1734 +matplotlib/mpl-data/images/filesave.png,sha256=b7ctucrM_F2mG-DycTedG_a_y4pHkx3F-zM7l18GLhk,458 +matplotlib/mpl-data/images/filesave.svg,sha256=dMGXvLSOHPu44kiWgZx-B_My_tLWaP6J6GgxJfL4FW0,2049 +matplotlib/mpl-data/images/filesave_large.png,sha256=LNbRD5KZ3Kf7nbp-stx_a1_6XfGBSWUfDdpgmnzoRvk,720 +matplotlib/mpl-data/images/forward-symbolic.svg,sha256=kOiKq3a4mieMRLVCwQBdOMTRrWG2NOX_5-rbAFHpdmQ,1551 +matplotlib/mpl-data/images/forward.pdf,sha256=KIqIL4YId43LkcOxV_TT5uvz1SP8k5iUNUeJmAElMV8,1630 +matplotlib/mpl-data/images/forward.png,sha256=pKbLepgGiGeyY2TCBl8svjvm7Z4CS3iysFxcq4GR-wk,357 +matplotlib/mpl-data/images/forward.svg,sha256=kOiKq3a4mieMRLVCwQBdOMTRrWG2NOX_5-rbAFHpdmQ,1551 +matplotlib/mpl-data/images/forward_large.png,sha256=36h7m7DZDHql6kkdpNPckyi2LKCe_xhhyavWARz_2kQ,593 +matplotlib/mpl-data/images/hand.pdf,sha256=hspwkNY915KPD7AMWnVQs7LFPOtlcj0VUiLu76dMabQ,4172 +matplotlib/mpl-data/images/hand.png,sha256=2cchRETGKa0hYNKUxnJABwkyYXEBPqJy_VqSPlT0W2Q,979 +matplotlib/mpl-data/images/hand.svg,sha256=hxxBtakaVFA7mpZOGakvo0QUcb2x06rojeS5gnVmyuc,4906 +matplotlib/mpl-data/images/help-symbolic.svg,sha256=XVcFcuzcL3SQ3LjfSbtdLYDjoB5YUkj2jk2Gk8vaZF8,1890 +matplotlib/mpl-data/images/help.pdf,sha256=CeE978IMi0YWznWKjIT1R8IrP4KhZ0S7usPUvreSgcA,1813 +matplotlib/mpl-data/images/help.png,sha256=s4pQrqaQ0py8I7vc9hv3BI3DO_tky-7YBMpaHuBDCBY,472 +matplotlib/mpl-data/images/help.svg,sha256=XVcFcuzcL3SQ3LjfSbtdLYDjoB5YUkj2jk2Gk8vaZF8,1890 +matplotlib/mpl-data/images/help_large.png,sha256=1IwEyWfGRgnoCWM-r9CJHEogTJVD5n1c8LXTK4AJ4RE,747 +matplotlib/mpl-data/images/home-symbolic.svg,sha256=ptrus8h5PZTi9ahYfnaz-uZ8MAHCr72aPeMW48TBR9Q,1911 +matplotlib/mpl-data/images/home.pdf,sha256=e0e0pI-XRtPmvUCW2VTKL1DeYu1pvPmUUeRSgEbWmik,1737 +matplotlib/mpl-data/images/home.png,sha256=IcFdAAUa6_A0qt8IO3I8p4rpXpQgAlJ8ndBECCh7C1w,468 +matplotlib/mpl-data/images/home.svg,sha256=ptrus8h5PZTi9ahYfnaz-uZ8MAHCr72aPeMW48TBR9Q,1911 +matplotlib/mpl-data/images/home_large.png,sha256=uxS2O3tWOHh1iau7CaVV4ermIJaZ007ibm5Z3i8kXYg,790 +matplotlib/mpl-data/images/matplotlib.pdf,sha256=BkSUf-2xoij-eXfpV2t7y1JFKG1zD1gtV6aAg3Xi_wE,22852 +matplotlib/mpl-data/images/matplotlib.png,sha256=w8KLRYVa-voUZXa41hgJauQuoois23f3NFfdc72pUYY,1283 +matplotlib/mpl-data/images/matplotlib.svg,sha256=QiTIcqlQwGaVPtHsEk-vtmJk1wxwZSvijhqBe_b9VCI,62087 +matplotlib/mpl-data/images/matplotlib_large.png,sha256=ElRoue9grUqkZXJngk-nvh4GKfpvJ4gE69WryjCbX5U,3088 +matplotlib/mpl-data/images/move-symbolic.svg,sha256=_uamLnjQ20iwSuKbd8JvTXUFaRq4206MrpFWvtErr8I,2529 +matplotlib/mpl-data/images/move.pdf,sha256=CXk3PGK9WL5t-5J-G2X5Tl-nb6lcErTBS5oUj2St6aU,1867 +matplotlib/mpl-data/images/move.png,sha256=TmjR41IzSzxGbhiUcV64X0zx2BjrxbWH3cSKvnG2vzc,481 +matplotlib/mpl-data/images/move.svg,sha256=_uamLnjQ20iwSuKbd8JvTXUFaRq4206MrpFWvtErr8I,2529 +matplotlib/mpl-data/images/move_large.png,sha256=Skjz2nW_RTA5s_0g88gdq2hrVbm6DOcfYW4Fu42Fn9U,767 +matplotlib/mpl-data/images/qt4_editor_options.pdf,sha256=2qu6GVyBrJvVHxychQoJUiXPYxBylbH2j90QnytXs_w,1568 +matplotlib/mpl-data/images/qt4_editor_options.png,sha256=EryQjQ5hh2dwmIxtzCFiMN1U6Tnd11p1CDfgH5ZHjNM,380 +matplotlib/mpl-data/images/qt4_editor_options.svg,sha256=sdrNIxYT-BLvJ30ASnaRQ5PxF3SB41-pgdaIJT0KqBg,1264 +matplotlib/mpl-data/images/qt4_editor_options_large.png,sha256=-Pd-9Vh5aIr3PZa8O6Ge_BLo41kiEnpmkdDj8a11JkY,619 +matplotlib/mpl-data/images/subplots-symbolic.svg,sha256=Gq4fDSS99Rv5rbR8_nenV6jcY5VsKPARWeH-BZBk9CU,2150 +matplotlib/mpl-data/images/subplots.pdf,sha256=Q0syPMI5EvtgM-CE-YXKOkL9eFUAZnj_X2Ihoj6R4p4,1714 +matplotlib/mpl-data/images/subplots.png,sha256=MUfCItq3_yzb9yRieGOglpn0Y74h8IA7m5i70B63iRc,445 +matplotlib/mpl-data/images/subplots.svg,sha256=Gq4fDSS99Rv5rbR8_nenV6jcY5VsKPARWeH-BZBk9CU,2150 +matplotlib/mpl-data/images/subplots_large.png,sha256=Edu9SwVMQEXJZ5ogU5cyW7VLcwXJdhdf-EtxxmxdkIs,662 +matplotlib/mpl-data/images/zoom_to_rect-symbolic.svg,sha256=uMmdGkO43ZHlezkpieR3_MiqlEc5vROffRDOhY4sxm4,1499 +matplotlib/mpl-data/images/zoom_to_rect.pdf,sha256=SEvPc24gfZRpl-dHv7nx8KkxPyU66Kq4zgQTvGFm9KA,1609 +matplotlib/mpl-data/images/zoom_to_rect.png,sha256=aNz3QZBrIgxu9E-fFfaQweCVNitGuDUFoC27e5NU2L4,530 +matplotlib/mpl-data/images/zoom_to_rect.svg,sha256=uMmdGkO43ZHlezkpieR3_MiqlEc5vROffRDOhY4sxm4,1499 +matplotlib/mpl-data/images/zoom_to_rect_large.png,sha256=V6pkxmm6VwFExdg_PEJWdK37HB7k3cE_corLa7RbUMk,1016 +matplotlib/mpl-data/kpsewhich.lua,sha256=RdyYaBnBLy3NsB5c2R5FGrKu-V-WBcZim24NWilsTfw,139 +matplotlib/mpl-data/matplotlibrc,sha256=2hOw4G5YMRexNcWmcT_dogowmoui_AfYOfVvGM6aPHM,44226 +matplotlib/mpl-data/plot_directive/plot_directive.css,sha256=utSJ1oETz0UG6AC9hU134J_JY78ENijqMZXN0JMBUfk,318 +matplotlib/mpl-data/sample_data/Minduka_Present_Blue_Pack.png,sha256=XnKGiCanpDKalQ5anvo5NZSAeDP7fyflzQAaivuc0IE,13634 +matplotlib/mpl-data/sample_data/README.txt,sha256=ABz19VBKfGewdY39QInG9Qccgn1MTYV3bT5Ph7TCy2Y,128 +matplotlib/mpl-data/sample_data/Stocks.csv,sha256=72878aZNXGxd5wLvFUw_rnj-nfg4gqtrucZji-w830c,67924 +matplotlib/mpl-data/sample_data/axes_grid/bivariate_normal.npy,sha256=DpWZ9udAh6ospYqneEa27D6EkRgORFwHosacZXVu98U,1880 +matplotlib/mpl-data/sample_data/data_x_x2_x3.csv,sha256=A0SU3buOUGhT-NI_6LQ6p70fFSIU3iLFdgzvzrKR6SE,132 +matplotlib/mpl-data/sample_data/eeg.dat,sha256=KGVjFt8ABKz7p6XZirNfcxSTOpGGNuyA8JYErRKLRBc,25600 +matplotlib/mpl-data/sample_data/embedding_in_wx3.xrc,sha256=cUqVw5vDHNSZoaO4J0ebZUf5SrJP36775abs7R9Bclg,2186 +matplotlib/mpl-data/sample_data/goog.npz,sha256=QAkXzzDmtmT3sNqT18dFhg06qQCNqLfxYNLdEuajGLE,22845 +matplotlib/mpl-data/sample_data/grace_hopper.jpg,sha256=qMptc0dlcDsJcoq0f-WfRz2Trjln_CTHwCiMPHrbcTA,61306 +matplotlib/mpl-data/sample_data/jacksboro_fault_dem.npz,sha256=1JP1CjPoKkQgSUxU0fyhU50Xe9wnqxkLxf5ukvYvtjc,174061 +matplotlib/mpl-data/sample_data/logo2.png,sha256=DXNx4FXeyqxHy26AmvNELpwezQLxweLQY9HP7ktKIdc,22279 +matplotlib/mpl-data/sample_data/membrane.dat,sha256=q3lbQpIBpbtXXGNw1eFwkN_PwxdDGqk4L46IE2b0M1c,48000 +matplotlib/mpl-data/sample_data/msft.csv,sha256=GArKb0O3DgKZRsKdJf6lX3rMSf-PCekIiBoLNdgF7Mk,3211 +matplotlib/mpl-data/sample_data/s1045.ima.gz,sha256=MrQk1k9it-ccsk0p_VOTitVmTWCAVaZ6srKvQ2n4uJ4,33229 +matplotlib/mpl-data/sample_data/topobathy.npz,sha256=AkTgMpFwLfRQJNy1ysvE89TLMNct-n_TccSsYcQrT78,45224 +matplotlib/mpl-data/stylelib/Solarize_Light2.mplstyle,sha256=aytOm4eT_SPvs7HC28ZY4GukeN44q-SE0JEMCR8kVOk,1257 +matplotlib/mpl-data/stylelib/_classic_test_patch.mplstyle,sha256=iopHpMaM3im_AK2aiHGuM2DKM5i9Kc84v6NQEoSb10Q,167 +matplotlib/mpl-data/stylelib/_mpl-gallery-nogrid.mplstyle,sha256=1VOL3USqD6iuGQaSynNg1QhyUwvKLnkLyUKdbBMnnqg,489 +matplotlib/mpl-data/stylelib/_mpl-gallery.mplstyle,sha256=MN-q59CiDqHXB8xFKXxzCbJJbJmNDhBe9lDJJAoMTPA,504 +matplotlib/mpl-data/stylelib/bmh.mplstyle,sha256=-KbhaI859BITHIoyUZIfpQDjfckgLAlDAS_ydKsm6mc,712 +matplotlib/mpl-data/stylelib/classic.mplstyle,sha256=1o5b47VD_RIZv3unnG9Gm2tbprTvOeNGXM8hJCmGuYI,24670 +matplotlib/mpl-data/stylelib/dark_background.mplstyle,sha256=GzSBD06jvlRYOqu7D5Z5a5x25l9JnTgtObn7S4D9zug,607 +matplotlib/mpl-data/stylelib/fast.mplstyle,sha256=yTa2YEIIP9xi5V_G0p2vSlxghuhNwjRi9gPECMxyRiM,288 +matplotlib/mpl-data/stylelib/fivethirtyeight.mplstyle,sha256=IcBt2DSDz9CtuyCis9JZtT3Nqoh6LVO-dB66AxObUcA,780 +matplotlib/mpl-data/stylelib/ggplot.mplstyle,sha256=u2oPHMLWFtZcpIjHk2swi2Nrt4NgnEtof5lxcwM0RD0,956 +matplotlib/mpl-data/stylelib/grayscale.mplstyle,sha256=KCLg-pXpns9cnKDXKN2WH6mV41OH-6cbT-5zKQotSdw,526 +matplotlib/mpl-data/stylelib/petroff10.mplstyle,sha256=Qj7pPbHh3L25kdsPt1ypwvKR3dPzuOrFzzQTP3Mfilg,298 +matplotlib/mpl-data/stylelib/seaborn-v0_8-bright.mplstyle,sha256=pDqn3-NUyVLvlfkYs8n8HzNZvmslVMChkeH-HtZuJIc,144 +matplotlib/mpl-data/stylelib/seaborn-v0_8-colorblind.mplstyle,sha256=eCSzFj5_2vR6n5qu1rHE46wvSVGZcdVqz85ov40ZsH8,148 +matplotlib/mpl-data/stylelib/seaborn-v0_8-dark-palette.mplstyle,sha256=p5ABKNQHRG7bk4HXqMQrRBjDlxGAo3RCXHdQmP7g-Ng,142 +matplotlib/mpl-data/stylelib/seaborn-v0_8-dark.mplstyle,sha256=I4xQ75vE5_9X4k0cNDiqhhnF3OcrZ2xlPX8Ll7OCkoE,667 +matplotlib/mpl-data/stylelib/seaborn-v0_8-darkgrid.mplstyle,sha256=2bXOSzS5gmPzRBrRmzVWyhg_7ZaBRQ6t_-O-cRuyZoA,670 +matplotlib/mpl-data/stylelib/seaborn-v0_8-deep.mplstyle,sha256=44dLcXjjRgR-6yaopgGRInaVgz3jk8VJVQTbBIcxRB0,142 +matplotlib/mpl-data/stylelib/seaborn-v0_8-muted.mplstyle,sha256=T4o3jvqKD_ImXDkp66XFOV_xrBVFUolJU34JDFk1Xkk,143 +matplotlib/mpl-data/stylelib/seaborn-v0_8-notebook.mplstyle,sha256=PcvZQbYrDdducrNlavBPmQ1g2minio_9GkUUFRdgtoM,382 +matplotlib/mpl-data/stylelib/seaborn-v0_8-paper.mplstyle,sha256=n0mboUp2C4Usq2j6tNWcu4TZ_YT4-kKgrYO0t-rz1yw,393 +matplotlib/mpl-data/stylelib/seaborn-v0_8-pastel.mplstyle,sha256=8nV8qRpbUrnFZeyE6VcQ1oRuZPLil2W74M2U37DNMOE,144 +matplotlib/mpl-data/stylelib/seaborn-v0_8-poster.mplstyle,sha256=dUaKqTE4MRfUq2rWVXbbou7kzD7Z9PE9Ko8aXLza8JA,403 +matplotlib/mpl-data/stylelib/seaborn-v0_8-talk.mplstyle,sha256=7FnBaBEdWBbncTm6_ER-EQVa_bZgU7dncgez-ez8R74,403 +matplotlib/mpl-data/stylelib/seaborn-v0_8-ticks.mplstyle,sha256=CITZmZFUFp40MK2Oz8tI8a7WRoCizQU9Z4J172YWfWw,665 +matplotlib/mpl-data/stylelib/seaborn-v0_8-white.mplstyle,sha256=WjJ6LEU6rlCwUugToawciAbKP9oERFHr9rfFlUrdTx0,665 +matplotlib/mpl-data/stylelib/seaborn-v0_8-whitegrid.mplstyle,sha256=ec4BjsNzmOvHptcJ3mdPxULF3S1_U1EUocuqfIpw-Nk,664 +matplotlib/mpl-data/stylelib/seaborn-v0_8.mplstyle,sha256=_Xu6qXKzi4b3GymCOB1b1-ykKTQ8xhDliZ8ezHGTiAs,1130 +matplotlib/mpl-data/stylelib/tableau-colorblind10.mplstyle,sha256=BsirZVd1LmPWT4tBIz6loZPjZcInoQrIGfC7rvzqmJw,190 +matplotlib/offsetbox.py,sha256=hcBN7g2UBFLu_WhBg_bDlnHcIdjc6Pj-kHicfhOpTpI,54490 +matplotlib/offsetbox.pyi,sha256=ISf7gRrOhvefNKfutbHIPh6QYS0kAgtA97bYa6N-82o,9909 +matplotlib/patches.py,sha256=uXew19bSegJ1vxVIhZMhmf8f4XXyB7wKIwgAnFPfwnk,165064 +matplotlib/patches.pyi,sha256=d-sCqVXw_ANsvoUpliH7dllMbTy7TcyRm2IV97Ll0nM,22649 +matplotlib/path.py,sha256=QiQRSMjHbKCBpjWyJ3KPIJgK1Eq07qXUYjULaDHgt5g,43269 +matplotlib/path.pyi,sha256=VfIuc2KpkcB28uzRjLvkC3Hq5HwH4Zp0Uyj9H9bwcWU,4808 +matplotlib/patheffects.py,sha256=4paQgPGKc_xdBz_BZ_MaVhKJ2FxZDGeh5zHwzn3evDY,18387 +matplotlib/patheffects.pyi,sha256=7-FhuxGGrer94GtJ1sZ0YxOmK6Nv4oixTmsKb6-ijOg,3664 +matplotlib/projections/__init__.py,sha256=gICKgNfjJ-tioWEFa75BKKJZG6JJjx46_RpEqt5fk94,4438 +matplotlib/projections/__init__.pyi,sha256=D28dSYmwZcSBFBtNDer-QqE_lqXAhHKersuAlvi89jE,673 +matplotlib/projections/__pycache__/__init__.cpython-310.pyc,, +matplotlib/projections/__pycache__/geo.cpython-310.pyc,, +matplotlib/projections/__pycache__/polar.cpython-310.pyc,, +matplotlib/projections/geo.py,sha256=feR2dybCylhk6xe-MzxBr_e8MdbpVc8tBtCQQJ0V2co,17605 +matplotlib/projections/geo.pyi,sha256=vPfhvj7_e0ZnKjyfDUNC89RGCktycJBPnn5D8w0A7N8,3775 +matplotlib/projections/polar.py,sha256=vEy7FB4FnEu7COWyjZT-zYDt4tbTwkW8BvZCr7kqYiQ,57241 +matplotlib/projections/polar.pyi,sha256=IWNSLQIo5cXFBJur5VFVDQBwAtq5n1XFZMnrj85mCBg,6636 +matplotlib/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/pylab.py,sha256=VqUqd2J2-dKtltZtsYP8ePKoX1jsNDpTyqlcmpfC-lo,2369 +matplotlib/pyplot.py,sha256=Dbp8GKxM6B5oCmcjb8lsik2R8QPtPt3JgmgaqeWqkPw,149704 +matplotlib/quiver.py,sha256=L8O6Ke6oH60R9VA_gd9JiA5Fvv5sMPIFyC_uNRrZSo0,48675 +matplotlib/quiver.pyi,sha256=Kq1FvQP-DRX1N9Xlp6XF6OJfKkJ7tc4dYq4CEndxEzg,5640 +matplotlib/rcsetup.py,sha256=87KrQeDFg4sRu4obxF4mc0eGSCuClGyuzXsxoZ_kme4,51606 +matplotlib/rcsetup.pyi,sha256=0VTIhzfKgBcKxOeOdoKq1oa7ZubalOj9Ks0efE0-cms,4337 +matplotlib/sankey.py,sha256=eKA9DrX96-bpoWRdeNnvC73aJfPtCnB78WmMPCXQ9rc,36151 +matplotlib/sankey.pyi,sha256=P2AGkhdg3ZLqXaQdf8yodblrBpdgEQZBLKPL1mbxmFc,1451 +matplotlib/scale.py,sha256=n3JqLU_d3dAKrP1uO5lTipVEaxgiLhDR2sSk8TAXE_M,26924 +matplotlib/scale.pyi,sha256=-ptRptcqiAuzfKwrjgSWWOxFmjRUTOKGwIoWtuBXKgY,5057 +matplotlib/sphinxext/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/sphinxext/__pycache__/__init__.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/figmpl_directive.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/mathmpl.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/plot_directive.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/roles.cpython-310.pyc,, +matplotlib/sphinxext/figmpl_directive.py,sha256=rgoJcGRfLdp4jiOlyBIwG83wcyfzbFVfrlrSCmDHL9g,9308 +matplotlib/sphinxext/mathmpl.py,sha256=CKIocBSV4prRMPR50NL6WxBTULBa1zReoqo7ZtbvCsc,7871 +matplotlib/sphinxext/plot_directive.py,sha256=6GEOX2y0K1Izc98Aojt8pCXS21Ucru6E-T_33lRU92s,32542 +matplotlib/sphinxext/roles.py,sha256=U7P4t06foLF2PsZJrTN4mfaspKGRoYwFZ3ijIhIUVyo,4882 +matplotlib/spines.py,sha256=_9Dc4oIAgDs3MA6rZCwr1scPGCSzbAnmYp3PQzpKFeU,21923 +matplotlib/spines.pyi,sha256=eceCS47bOawgkWdLUMIPt6XX_H-EtWS2t-euz5Bgkbg,2951 +matplotlib/stackplot.py,sha256=a6IUAujDjGpaFqiTzGr85eDD8l6baWLPrM5YNvb6g9I,4997 +matplotlib/stackplot.pyi,sha256=BX2g3-6dJ4RhjjIguo-gwWhvKZNYqFCQYVw5SKvpM4c,561 +matplotlib/streamplot.py,sha256=rXVKYEi4MfSg3HnAaWb8w4qFp95mRf6dosMVqtXZEpE,24011 +matplotlib/streamplot.pyi,sha256=fkdkew8u807qd7WXFd_ZjF7FUbx8okOI8A7u4x_ktvg,2690 +matplotlib/style/__init__.py,sha256=fraQtyBC3TY2ZTsLdxL7zNs9tJYZtje7tbiqEf3M56M,140 +matplotlib/style/__pycache__/__init__.cpython-310.pyc,, +matplotlib/style/__pycache__/core.cpython-310.pyc,, +matplotlib/style/core.py,sha256=RSNH2NO0xK_Y37dbpMWyaWMMjvcMa_NHaKXDa5I4_tA,8362 +matplotlib/style/core.pyi,sha256=mIlyz6eChdMjZ_A7S05iJgcWwrWE2NXg2rhm1C1eYlQ,521 +matplotlib/table.py,sha256=MkF9s8u2X-9_Cei9sJixljb2_I_B_4UOqVx0iReKIU0,27744 +matplotlib/table.pyi,sha256=tcR40hoCWCRLlL-8PtFMa7P5sq8qajDSMFxd25Q90z4,3098 +matplotlib/testing/__init__.py,sha256=pTaqH1qgGNBeWCmvK_vjqFJXJ9dvZXCaKQ-8jHSpaEc,6942 +matplotlib/testing/__init__.pyi,sha256=ffqfetWzyCVrSx7BlnoxCmbgIaZg2x57nDrq9eucRk0,1752 +matplotlib/testing/__pycache__/__init__.cpython-310.pyc,, +matplotlib/testing/__pycache__/_markers.cpython-310.pyc,, +matplotlib/testing/__pycache__/compare.cpython-310.pyc,, +matplotlib/testing/__pycache__/conftest.cpython-310.pyc,, +matplotlib/testing/__pycache__/decorators.cpython-310.pyc,, +matplotlib/testing/__pycache__/exceptions.cpython-310.pyc,, +matplotlib/testing/__pycache__/widgets.cpython-310.pyc,, +matplotlib/testing/_markers.py,sha256=0iNyOi25XLv_gTfSUqiRizdSqJzozePPBMRo72H2Je4,1419 +matplotlib/testing/compare.py,sha256=8dVcuP7FauHY1JZQXBD7fNBfQKb4IoXM6pVeq6zJgBs,19780 +matplotlib/testing/compare.pyi,sha256=xlJ4chgXKe567NUavlu-dalyPr4wAQUbd2Fz6aK_JII,1192 +matplotlib/testing/conftest.py,sha256=CcwfUtIkUOgGBlXXNplQE6dg2U_UBKdH6AajROPJ5uw,6683 +matplotlib/testing/conftest.pyi,sha256=zOF_MM2GjCkNi9nZ8vldV9VI41bD67nVhrIgLVuMNsk,416 +matplotlib/testing/decorators.py,sha256=1Xiyb2exeN0nDEO4he8MW_VKwFE34cB7EtRl6bcGRzg,18022 +matplotlib/testing/decorators.pyi,sha256=0fSpdLBtEH7ZP_trVJ7RPxNtOX9sJ_z-MkNsbUxF8nM,872 +matplotlib/testing/exceptions.py,sha256=72QmjiHG7DwxSvlJf8mei-hRit5AH3NKh0-osBo4YbY,138 +matplotlib/testing/jpl_units/Duration.py,sha256=9FMBu9uj6orCWtf23cf6_9HCFUC50xAHrCzaxATwQfM,3966 +matplotlib/testing/jpl_units/Epoch.py,sha256=-FGxeq-VvCS9GVPwOEE5ind_G4Tl9ztD-gYcW9CWzjo,6100 +matplotlib/testing/jpl_units/EpochConverter.py,sha256=fhWjyP567bzcTU_oNuJJpucoolqS88Nt-yEFg1-3yEk,2944 +matplotlib/testing/jpl_units/StrConverter.py,sha256=codGw9b_Zc-MG_YK4CiyMrnMR8ahR9hw836O2SsV8QI,2865 +matplotlib/testing/jpl_units/UnitDbl.py,sha256=EABjyEK4MVouyvlwi_9KdYDg-qbYY3aLHoUjRw37Fb0,5882 +matplotlib/testing/jpl_units/UnitDblConverter.py,sha256=B8DssrQVyC4mwvSFP78cGL0vCnZgVhDaAbZE-jsXLUg,2828 +matplotlib/testing/jpl_units/UnitDblFormatter.py,sha256=246hgA4_pCfJm-P94hEsxqnTS9t0XlvLC8p1v_bw2pU,657 +matplotlib/testing/jpl_units/__init__.py,sha256=p__9RUwrt2LJ2eoT2JPM-42XLxSJrfA4az3rN5uP6d4,2684 +matplotlib/testing/jpl_units/__pycache__/Duration.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/Epoch.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/EpochConverter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/StrConverter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/UnitDbl.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/UnitDblConverter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/UnitDblFormatter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/__init__.cpython-310.pyc,, +matplotlib/testing/widgets.py,sha256=IaoPHgDguJSHd7t97kSVWUZvqTQ-Tsr9tU9kfnmJrAM,3480 +matplotlib/testing/widgets.pyi,sha256=Ioau7Q2aPRDZLx8hze2DOe3E1vn7QPxePC74WMR7tFc,831 +matplotlib/tests/__init__.py,sha256=XyXveEAxafB87gnbx0jkC0MggzKO8FvORq_6RtJRwo4,366 +matplotlib/tests/__pycache__/__init__.cpython-310.pyc,, +matplotlib/tests/__pycache__/conftest.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_afm.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_agg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_agg_filter.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_animation.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_api.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_arrow_patches.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_artist.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_axes.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_axis.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_bases.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_cairo.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_gtk3.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_inline.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_macosx.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_nbagg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_pdf.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_pgf.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_ps.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_qt.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_registry.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_svg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_template.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_tk.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_tools.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_webagg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backends_interactive.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_basic.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_bbox_tight.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_bezier.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_category.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_cbook.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_collections.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_colorbar.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_colors.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_compare_images.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_constrainedlayout.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_container.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_contour.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_cycles.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_dates.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_datetime.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_determinism.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_doc.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_dviread.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_figure.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_font_manager.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_fontconfig_pattern.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_ft2font.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_getattr.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_gridspec.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_image.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_legend.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_lines.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_marker.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_mathtext.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_matplotlib.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_mlab.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_multivariate_colormaps.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_offsetbox.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_patches.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_path.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_patheffects.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_pickle.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_png.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_polar.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_preprocess_data.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_pyplot.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_quiver.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_rcparams.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_sankey.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_scale.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_simplification.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_skew.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_sphinxext.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_spines.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_streamplot.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_style.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_subplots.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_table.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_testing.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_texmanager.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_text.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_textpath.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_ticker.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_tightlayout.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_transforms.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_triangulation.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_type1font.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_units.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_usetex.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_widgets.cpython-310.pyc,, +matplotlib/tests/conftest.py,sha256=0cKxWVUC25UfQJQcAjFJa73MQE9eSYqKSXGDx7kAq48,138 +matplotlib/tests/test_afm.py,sha256=A7jm2o-QaQH9SSpiFxGtZkbVU0LJZE69jfPv7RczOD4,3701 +matplotlib/tests/test_agg.py,sha256=xAKrTzvT1dNM_49efiWisnGNrJtINOBenYSOtaRXvIg,10884 +matplotlib/tests/test_agg_filter.py,sha256=3c_Smtb4OHEOfdMFCOb2qKhzMXSbNoaUtsJ0pW47Q44,1067 +matplotlib/tests/test_animation.py,sha256=z0FJ0gx15rJB_khEPSg1E-vLGQyTj85Ki-uTjFsMEPE,18357 +matplotlib/tests/test_api.py,sha256=UJU_qoi7Y8Rbz1S89NYKvOkNwv_RBVPNFajrcXixNds,5736 +matplotlib/tests/test_arrow_patches.py,sha256=qgbS9OrlHa4TJihAd2ekW4AupCVgKDd2B5YyF704Pak,6577 +matplotlib/tests/test_artist.py,sha256=oblCTkCPZzufghCp19AKwIJftV69tZwEOjeRLn8soAQ,18599 +matplotlib/tests/test_axes.py,sha256=PbsmnNA9f3GRtSjCEf4HUvZ5onrBEUIASd1bcalhil8,329934 +matplotlib/tests/test_axis.py,sha256=Ljc_Yc7dAFkTuVW0PD0i3vEnjkGOuAehkpcNpWg9A7Y,1446 +matplotlib/tests/test_backend_bases.py,sha256=FwUgearYkC7UgqpRpaKk51o5F62ipDL4kUvFROwGVLo,22775 +matplotlib/tests/test_backend_cairo.py,sha256=O2LTYjsfPn__bKtTz4MGGBodpSshoPkzu0INsc18xmI,1821 +matplotlib/tests/test_backend_gtk3.py,sha256=Zb29s85nbRzwLVmChCA5_LphAeC_yKyP45kCcHxfgPE,2849 +matplotlib/tests/test_backend_inline.py,sha256=b7GrmMCrJfI5kXGJKbxdMM9kWBd-JKUdCkRq-vk8JXo,1608 +matplotlib/tests/test_backend_macosx.py,sha256=s4Kd1fnwYCuamGRLtkOAYtzer7jVOjLaA2LH52hhPqM,2233 +matplotlib/tests/test_backend_nbagg.py,sha256=o6ON7tt_LOW7wzg4StIMYuTgMFSPNWt1FOv2FFMIefk,1459 +matplotlib/tests/test_backend_pdf.py,sha256=CxVYTlz8_AQ7o90yh5cavpSnNBLozIL8UhscRYShUVM,14602 +matplotlib/tests/test_backend_pgf.py,sha256=WlGdZThz5uQ-x8-HqpMcB5IWmYijflNCC-rVNgvz0yw,13028 +matplotlib/tests/test_backend_ps.py,sha256=veHbm4_dpBDcczgomEjQxdz-Sv8wHij3LVSKWST9Bd0,12600 +matplotlib/tests/test_backend_qt.py,sha256=yiXBEbzY9KBWL_cG8o_J-q3AVU-LYTF0znKvwX-3szQ,12878 +matplotlib/tests/test_backend_registry.py,sha256=SiyyRRMT1tIXKxwn7IXq0SaIrYt3DEvBBe8TvNxa644,5950 +matplotlib/tests/test_backend_svg.py,sha256=L73hpBa1P-y978Y3Gpd_eSJAFJMvUvrOqiYaJ-CkQUw,22930 +matplotlib/tests/test_backend_template.py,sha256=uuE7oZ9pSBVDWrPL05B0WgCFsgv6HlXyetuPTfJn6a8,2184 +matplotlib/tests/test_backend_tk.py,sha256=A1h3AfPhEMbG5PLdDr1gbnwt_DlaQOBgFux5QsIpCpM,8850 +matplotlib/tests/test_backend_tools.py,sha256=C-B7NCkyWsQ5KzQEnI5Be16DsAHHZJU9P5v9--wsF-o,501 +matplotlib/tests/test_backend_webagg.py,sha256=4Oc-z7w-mTGETB1x0FQ_gZP9qHfyWh5lwWc9qPkkisc,938 +matplotlib/tests/test_backends_interactive.py,sha256=oo-vRlDTrLQZOm53FfWECF_JSKDDBJ-BKigekZIO8XE,28880 +matplotlib/tests/test_basic.py,sha256=ubAnlE-lFQzMhoBlWYZJKCAsox9Y3jBXS_IMn29Zi84,1141 +matplotlib/tests/test_bbox_tight.py,sha256=5yo6yy3HTUIf2-oGZllieaLynmOupzWcuPpeGmgS5BQ,6314 +matplotlib/tests/test_bezier.py,sha256=IrrDWAV5O6ELMdziaaOeBI7LfyEaTBVsT4B1T9KMzD8,692 +matplotlib/tests/test_category.py,sha256=jWXOAGOR_gpvrjarvJB6Tyot0yOIWaO6d1F9mrNcVPo,12043 +matplotlib/tests/test_cbook.py,sha256=xESpD9VE9dXInGp_QyzRCeGvmsgY2j9IyNMSUBae-WA,33656 +matplotlib/tests/test_collections.py,sha256=C4jbvAq9t7wcsXY-0SH666aP2t8w4dyO8nRnu7TRdYc,48896 +matplotlib/tests/test_colorbar.py,sha256=tDdmWz_LxCbD80ksM6iqrHYqpombyL8QtqpKQSKOS14,46711 +matplotlib/tests/test_colors.py,sha256=-Q5VZfQWI-NSMOa4tusxDgWe705a7uymMHLI8h10C8A,61215 +matplotlib/tests/test_compare_images.py,sha256=NcBoT8HAxtzoR1zZBu0V3MFpGWtdFaDblQ8o3OTO6zM,3260 +matplotlib/tests/test_constrainedlayout.py,sha256=Mx86XCAfw0XEG67dU73GjZZZi38EZMJErwvgrm7-ThU,24270 +matplotlib/tests/test_container.py,sha256=FTewckOd3dJqLOzEUa29Itjqusk7Mx7MK5xPVMhzMmc,694 +matplotlib/tests/test_contour.py,sha256=O_VIhjDgsb0A-brYR-gJ-WUcqzYzH0XLQD4aI_5SocM,30289 +matplotlib/tests/test_cycles.py,sha256=3KyRmWH29WUgvIXUT06tKVDNCfDWqxuxlLueIp-FIl0,5996 +matplotlib/tests/test_dates.py,sha256=jOK-KIAMBRCcnlgpzknpLJb3JvO3xFa5Qe3O5mVTeGM,56533 +matplotlib/tests/test_datetime.py,sha256=n2TYVItNQcUFBLUivuTXBnc0fElIeSmy9hu4ZNMVurc,32635 +matplotlib/tests/test_determinism.py,sha256=harG-hnOpD5HZFKPtHebT0snHu8V6C6BScVFml1Mqyg,7958 +matplotlib/tests/test_doc.py,sha256=K6HhdRcHRUNNC0iIhxBLCq5tOqxqHBSvanHxXg5w5iI,1015 +matplotlib/tests/test_dviread.py,sha256=JeTuA2FMUj1FddxDVBXUtnvZYTgztE-CyRXL_mI20P0,2764 +matplotlib/tests/test_figure.py,sha256=bGyAW9fwJc3gVVul_Vd5f8TEmbrEoDAMWV4Rm5-QXQQ,60289 +matplotlib/tests/test_font_manager.py,sha256=Qkewu4tMylgVciamVwWL8TFoDon8YyzN-9po6usN6Mk,14136 +matplotlib/tests/test_fontconfig_pattern.py,sha256=LSR6lWF_0cVOshkGiflYaTCLcISRIJM6mjXm5QtDjX4,2168 +matplotlib/tests/test_ft2font.py,sha256=lz13lRGmxdQ9eRjQlnigdfQWBwaST1A3_AgEAR5dJA0,40838 +matplotlib/tests/test_getattr.py,sha256=Tl_H1zpwLdSIVutc4vi-QwDCeWPzBGpN31N9ItzTkeQ,1090 +matplotlib/tests/test_gridspec.py,sha256=SYJuJucA_PyQ2Rsov1RaNhadOEWGDcMbQcVFrWIoy3I,1580 +matplotlib/tests/test_image.py,sha256=O9pUxNrGyewPievlC-WDhJ3hkCTsiAq4id6E6hJ22Is,59562 +matplotlib/tests/test_legend.py,sha256=QeMjTc20_IZeVShlxEZQQsv2suFAPqpvD86TBHESyks,55122 +matplotlib/tests/test_lines.py,sha256=KElNi2GlJCtGLhNpf8rl8bkDhMcQ9ittyRdsX4rdSxo,15035 +matplotlib/tests/test_marker.py,sha256=w0WVHoaD-6iybjUSENoVFFdUOOR13943JcE4sgz3qhI,11410 +matplotlib/tests/test_mathtext.py,sha256=2S1i3BLNffEqz1t-mehhWBROBcsdBZsdgrticxW6Xrw,24549 +matplotlib/tests/test_matplotlib.py,sha256=ZQgD9x7UVWFnVDOm9K5iKfobz6BtMtFkgyVRQ5QezYY,3368 +matplotlib/tests/test_mlab.py,sha256=d4qMyogTFMrvlRZEpDs7SjhSmmCnBUMNgSX2bJU6eDk,42269 +matplotlib/tests/test_multivariate_colormaps.py,sha256=wGc08cdUTBDJvhUiwVqTrBS7o19Ei-wQRnWlkny7h5M,20785 +matplotlib/tests/test_offsetbox.py,sha256=VXFfgputVnLodHhK5ofEWqge42trRSvIm2areAbv1ZQ,16656 +matplotlib/tests/test_patches.py,sha256=FeQGNDS0G8DJNWCAQVit2vN4JPN0ARNougGOTCNa3Fk,33592 +matplotlib/tests/test_path.py,sha256=R3oTrr6kS1Nu_q8zeK1uxrrbnQI8NMb5B7yTk8R1eRs,23254 +matplotlib/tests/test_patheffects.py,sha256=q7JA5VbZU5Zsqvc2eGM7YhB9TClkNp3akQCZVsHrNQU,8109 +matplotlib/tests/test_pickle.py,sha256=2VcsJ6QwN9vP4cvkuUCZD7uwKyd1rOCDIaYPEBr3u6E,10014 +matplotlib/tests/test_png.py,sha256=d6u6UkU71T6ULxDVSdMkIT3CUAL46dEuDuNrnjZe5BE,1407 +matplotlib/tests/test_polar.py,sha256=pPXBMmeUIUZhykYyZa5mrrjFL9GJE3vsaTX94cRjstU,18306 +matplotlib/tests/test_preprocess_data.py,sha256=cIVICUi1iahMQS30sqI5IlT2RYJRH2gJ0z60FyuYUFk,11363 +matplotlib/tests/test_pyplot.py,sha256=5-Rv0g037ifl64vWybmlV8y22YjuY5lVX5aPvnzihiA,13910 +matplotlib/tests/test_quiver.py,sha256=5zSF_0DpoDXPBvGWD6iu1VfIUna15npLd2UyCG6_Ksg,11953 +matplotlib/tests/test_rcparams.py,sha256=PLos0asOBQMn0hjrVP_itywcbZjT8og_TaB3tqSFHac,26506 +matplotlib/tests/test_sankey.py,sha256=yg4i-qHT5MljFzFGViOnHbMkumM8bhPwRgoQ5M6CUEs,3900 +matplotlib/tests/test_scale.py,sha256=fqLu88lVy7-JZFNslpgiid11_TIivSkM9DxVnHwZJYo,8429 +matplotlib/tests/test_simplification.py,sha256=IJ1yEz2Y58QtKRRHz41NaXVS4kkRipaE_Un0HgM4K2g,21570 +matplotlib/tests/test_skew.py,sha256=_Mjwfgce6WxLH3dafI7irv2Cw0ANifytCtQXdyOStCk,6349 +matplotlib/tests/test_sphinxext.py,sha256=9hzVcUgMhtfo_6A_c-SvCYWDQY45Oq-lD6Os4PnRg3s,9937 +matplotlib/tests/test_spines.py,sha256=g2CeF0S73XoGA-T04NvYFalIyZvcq7CpbTvN-ikVFn0,5274 +matplotlib/tests/test_streamplot.py,sha256=mSLuZ7E2eEtcI5nKHBZZ_30OYte5awqsXWFWROJYtCQ,5731 +matplotlib/tests/test_style.py,sha256=sd6rMLpPBS1hinR2svI1rXZmnXUnMQk1szxSNBUXtzU,6509 +matplotlib/tests/test_subplots.py,sha256=-YVCEob48FyLUmdInE8vQ2E2b1txTbqChFb1tDUu16s,11082 +matplotlib/tests/test_table.py,sha256=QsrDRpe864LP3zkQBtTciKub5raYz3i3Vrt_Ve3kdzQ,8659 +matplotlib/tests/test_testing.py,sha256=eh-1r4PIXcM7hfSKCNTU899QxRYWhZBg6W7x0TDVUxo,1057 +matplotlib/tests/test_texmanager.py,sha256=j_5ZG02ztIO-m0oXzd31a13YFw_G7JR7Ldt3fwf-phI,2647 +matplotlib/tests/test_text.py,sha256=3RxKrfJJBWKFGSrEvxeCG5rxcvlkIivir8wCj1ls5-M,38473 +matplotlib/tests/test_textpath.py,sha256=WLXvT5OzE6Tew4pr87LH-cgioCzM7srgMNRemiMEC5o,271 +matplotlib/tests/test_ticker.py,sha256=Bj3hiEgXMBXzDi0Lmg8k-Qg4ChPZrQVqGU0Qfh7Sv84,74759 +matplotlib/tests/test_tightlayout.py,sha256=KU0m9BnpbKV-jBdaI4EbDLE5Cf_BLtF0Y9hcFfQ4zlQ,13978 +matplotlib/tests/test_transforms.py,sha256=NfWkPWOKAmY1sq_R1T85ct4CYdFmYUjmjtg2uwOYIIY,48671 +matplotlib/tests/test_triangulation.py,sha256=hstVYhH09BoVl8L0n6vHXBabniZLUu2KYl6wqZn2Xac,55289 +matplotlib/tests/test_type1font.py,sha256=gZZDFi0buFus_aAgOrmtwKqbC36-88UkvM8Afbcngs4,6369 +matplotlib/tests/test_units.py,sha256=0v7JNqjkA8MVnvlOeFoYqQ3sW1KjRY_ijp8BGM8ZW_U,11675 +matplotlib/tests/test_usetex.py,sha256=a-Y6NuyROPHDGP2ELsOZNSVwBUoGtAv1xQZfisl9lSE,6405 +matplotlib/tests/test_widgets.py,sha256=U2FMCCJyKD-pcLbzGzEBAL91-3Fx80wsYdinF7eCOrg,66435 +matplotlib/texmanager.py,sha256=k3Ud7vHbtY0CjXmzHP-ZnEAxg6bHVxo2ALve7_e1rKg,15319 +matplotlib/texmanager.pyi,sha256=di3gbC9muXKTo5VCrW5ye-e19A1mrOl8e8lvI3b904A,1116 +matplotlib/text.py,sha256=xYSkLZiOpew4G1jkLA5pUjtF7qUhOs3yp8IEFMqE6Gg,70856 +matplotlib/text.pyi,sha256=EL_-kiqz88kuPs0hss6IucRrOCsLJ4zaKOYRFV5Wbk8,7019 +matplotlib/textpath.py,sha256=QmrZQ2Mtfyao6qWg-wmEGAbTwkuJ54AnAzFf9NzZHg4,13254 +matplotlib/textpath.pyi,sha256=rqOeTAeQYgm2b2NpetrEX0gMF8PzzW43xS5mNfUA98M,2529 +matplotlib/ticker.py,sha256=YoIJyWVlYk3DIZbBl-w6fzkuVES3j1dN4nyf-SHKVDg,107381 +matplotlib/ticker.pyi,sha256=YvnQDHxJT9H-iomDZpvBCoYkfRRYSot2MR0H0pjZM78,10604 +matplotlib/transforms.py,sha256=ctnOTqgPP4dgHCmSkUayQWL9KkjzWSzQEZALcc9oEzA,99763 +matplotlib/transforms.pyi,sha256=0I70FR4ZWmRoo6w5KjbXRrH0dVDmQyPFbGt_sNTYq10,12102 +matplotlib/tri/__init__.py,sha256=asnfefKRpJv7sGbfddCMybnJInVDPwgph7g0mpoh2u4,820 +matplotlib/tri/__pycache__/__init__.cpython-310.pyc,, +matplotlib/tri/__pycache__/_triangulation.cpython-310.pyc,, +matplotlib/tri/__pycache__/_tricontour.cpython-310.pyc,, +matplotlib/tri/__pycache__/_trifinder.cpython-310.pyc,, +matplotlib/tri/__pycache__/_triinterpolate.cpython-310.pyc,, +matplotlib/tri/__pycache__/_tripcolor.cpython-310.pyc,, +matplotlib/tri/__pycache__/_triplot.cpython-310.pyc,, +matplotlib/tri/__pycache__/_trirefine.cpython-310.pyc,, +matplotlib/tri/__pycache__/_tritools.cpython-310.pyc,, +matplotlib/tri/_triangulation.py,sha256=Ur2lKMOx4NrZxwyi0hBeBnVzicuKaCke0NkrZneSklM,9784 +matplotlib/tri/_triangulation.pyi,sha256=pVw1rvpIcl00p7V7E9GcvJSqQWyoxlZXX_p0_VSxTiY,1017 +matplotlib/tri/_tricontour.py,sha256=yxeH8QgBub1nTq9qbO5Iiiz81bcLXDeMcJHUS2Ahmns,10220 +matplotlib/tri/_tricontour.pyi,sha256=jnsAmVRX0-FOUw9ptUgci9J4T4JQRloKeH8fh8aAi-o,1155 +matplotlib/tri/_trifinder.py,sha256=3gUzJZDIwfdsSJUE8hIKso9e1-UGvynUN9HxaqC1EEc,3522 +matplotlib/tri/_trifinder.pyi,sha256=dXcZucacAS3Ch6nrDBPh2e3LYZLfZ7VwqpBUBb-vMPo,405 +matplotlib/tri/_triinterpolate.py,sha256=4FtyJSoJpHcFxkSkZHZ1aNengVNWqVKF4l78PgCH8O0,62445 +matplotlib/tri/_triinterpolate.pyi,sha256=nR0o0Jm0uPH-6l0ft1WRHqYlq-6o84X6M0aX1UJ9IvE,1044 +matplotlib/tri/_tripcolor.py,sha256=Z6urFv-naDM3_rhv_pkgIZI53h88aOo9qWc3MH7GmBA,6705 +matplotlib/tri/_tripcolor.pyi,sha256=QsA-A2ohj3r_tAElt2-9pzi47JiU01tNlRPDIptqnh4,1781 +matplotlib/tri/_triplot.py,sha256=jlHSz36Z5S18zBKc639PlSqdhfl7jHol8ExlddJuDI4,3102 +matplotlib/tri/_triplot.pyi,sha256=9USU-BfitrcdQE8yWOUlBX59QBNoHCWivDon9JbDQ0k,446 +matplotlib/tri/_trirefine.py,sha256=NG8bsDhZ5EOxMT-MsEWzJm11ZR3_8CAYHlG53IGi0ps,13178 +matplotlib/tri/_trirefine.pyi,sha256=J_PmjbeI6UbLaeecgj1OCvGe_sr9UUsNK9NGBSlQ320,1056 +matplotlib/tri/_tritools.py,sha256=wC9KVE6UqkWVHpyW9FU4hQdqRVRVmJlhaBF1EXsaD8U,10575 +matplotlib/tri/_tritools.pyi,sha256=XWwwvH2nIAmH8k59aRjnLBVQbTwKvd_FzdsRNASCJMw,402 +matplotlib/typing.py,sha256=ZxpQ5jAqiprq3UkqWParWJITgIuOi4ZJ6WEM_q3QOis,2439 +matplotlib/units.py,sha256=7O-llc8k3GpdotUs2tWcEGgoUHHX-Y7o0R7f-1Jve3k,6429 +matplotlib/widgets.py,sha256=zxn9B00f5TEVGe9cIXnQTeG-sGB8s8kBg4IuFHBfyt0,152667 +matplotlib/widgets.pyi,sha256=BpbjCZec7n-uOj_cImu-cB1imVC4oeL8rJUwjxbmN4c,15383 +mpl_toolkits/axes_grid1/__init__.py,sha256=wiuUCQo1g20SW5T3mFOmI9dGCJY6aDmglpQw5DfszEU,371 +mpl_toolkits/axes_grid1/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/anchored_artists.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_divider.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_grid.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_rgb.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_size.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/inset_locator.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/mpl_axes.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/parasite_axes.cpython-310.pyc,, +mpl_toolkits/axes_grid1/anchored_artists.py,sha256=ZO5bIF_29sWrk-lWocKJ_45OQdluf_gu7nLrCtfKtdY,17161 +mpl_toolkits/axes_grid1/axes_divider.py,sha256=JD3jGBAqRHRYb1jablrpq_RFQvYtYfwtIIPJgQpbKPc,21892 +mpl_toolkits/axes_grid1/axes_grid.py,sha256=sjEbq6hu7YFimoOYVIcaccEGXsSljNPaTLfJjFHWVtM,22344 +mpl_toolkits/axes_grid1/axes_rgb.py,sha256=pabgaWJuLTCPw2FlT6Zfy5d0_95CEvaLeosWRTElR98,5227 +mpl_toolkits/axes_grid1/axes_size.py,sha256=UAEQ-t0qeJlQcwbbrY4tK2NuSWheFyUXIkk4BUCauHc,7713 +mpl_toolkits/axes_grid1/inset_locator.py,sha256=_2U8ZAj_x7gjKhPk40VSbzE7L6wO2VY-OoOMfd6xOJs,19649 +mpl_toolkits/axes_grid1/mpl_axes.py,sha256=vFCttnj9JIgY3Mt2eOi-O_FVvdZ6SW_sBtIBFib6bz4,4251 +mpl_toolkits/axes_grid1/parasite_axes.py,sha256=809Uy3bLgXIDIGztefcvS9EoZcu3KfMV1QwLzX8qVT4,9404 +mpl_toolkits/axes_grid1/tests/__init__.py,sha256=sKLxL9jEJBX7eh5OumtXSOnTriPrJUkujTHFtnJVFrM,365 +mpl_toolkits/axes_grid1/tests/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axes_grid1/tests/__pycache__/conftest.cpython-310.pyc,, +mpl_toolkits/axes_grid1/tests/__pycache__/test_axes_grid1.cpython-310.pyc,, +mpl_toolkits/axes_grid1/tests/conftest.py,sha256=zB61sy90X97YJ16mIGiuaEAaBIjBEzRAK_qfSCichQM,147 +mpl_toolkits/axes_grid1/tests/test_axes_grid1.py,sha256=x-W845Cvl2O-eGGhpxKDayFg7-QQXNA94nODLCFPpJc,29093 +mpl_toolkits/axisartist/__init__.py,sha256=RPaNDl22FbmDP7ZRsku1yCqpoNqcclCk0a3rXj3G7fE,631 +mpl_toolkits/axisartist/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/angle_helper.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axes_divider.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axis_artist.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axisline_style.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axislines.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/floating_axes.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/grid_finder.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/grid_helper_curvelinear.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/parasite_axes.cpython-310.pyc,, +mpl_toolkits/axisartist/angle_helper.py,sha256=-mjKpaR1pLMJuoc0sx0_V3bv0iRPMrpS7r_WI0UYrCc,12952 +mpl_toolkits/axisartist/axes_divider.py,sha256=65xSCQ9cHSC3KE7J7HxS4bfsDTAbPmwyz1jJ43BqnBs,122 +mpl_toolkits/axisartist/axis_artist.py,sha256=9FY9yXl8eF5QuBvF8-Vipv_rmdpcOlA3Q48a5v7FoeA,38328 +mpl_toolkits/axisartist/axisline_style.py,sha256=9jbDkXEzMQiDHR-lDYKZEvTADtJwt2qlN1cErVUUdx0,6723 +mpl_toolkits/axisartist/axislines.py,sha256=QxKvChTaRPj0ovvxdrfr3pOzExEIf5svDwxRE7enEXg,16556 +mpl_toolkits/axisartist/floating_axes.py,sha256=kfhWKkmiy8tkXoVHvooTidXIlicU861VncElsUcBhLo,10337 +mpl_toolkits/axisartist/grid_finder.py,sha256=Hi2zwnQilavgrqWKScuSRba_WBPqqrbmErGU0XyAsOo,12265 +mpl_toolkits/axisartist/grid_helper_curvelinear.py,sha256=ofN7pPqEMh3r6bSf9eOHkAZEOrmhiNZV-Yig3jozkC4,12349 +mpl_toolkits/axisartist/parasite_axes.py,sha256=Ydi4-0Lbczr6K7Sz1-fRwK4Tm8KlHrOIumx67Xbo_9c,244 +mpl_toolkits/axisartist/tests/__init__.py,sha256=sKLxL9jEJBX7eh5OumtXSOnTriPrJUkujTHFtnJVFrM,365 +mpl_toolkits/axisartist/tests/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/conftest.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_angle_helper.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_axis_artist.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_axislines.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_floating_axes.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_grid_finder.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_grid_helper_curvelinear.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/conftest.py,sha256=zB61sy90X97YJ16mIGiuaEAaBIjBEzRAK_qfSCichQM,147 +mpl_toolkits/axisartist/tests/test_angle_helper.py,sha256=PwhJwBm2kk4uMyhdO5arQs8IlqSX2vN0hvUzI7YHqrw,5670 +mpl_toolkits/axisartist/tests/test_axis_artist.py,sha256=wt3bicVgUPnBX48-dH0Z6hboHgutIgwVpaGkcUZDeVU,2980 +mpl_toolkits/axisartist/tests/test_axislines.py,sha256=NXegrvEzVWovshta-qjbUKA2tpQcAbjYbfwkf6tKT6Y,4353 +mpl_toolkits/axisartist/tests/test_floating_axes.py,sha256=l24VB1SLrsJZOMMH2jmBny9ETha4AqAM5KokdGOa5Wk,4083 +mpl_toolkits/axisartist/tests/test_grid_finder.py,sha256=cwQLDOdcJbAY2E7dr8595yzuNh1_Yh80r_O8WGT2hMY,1156 +mpl_toolkits/axisartist/tests/test_grid_helper_curvelinear.py,sha256=OhHej0vCfCjJJknT7yIt4OxZd6OMJCXnFoT3pzqUtTo,7216 +mpl_toolkits/mplot3d/__init__.py,sha256=fH9HdMfFMvjbIWqy2gjQnm2m3ae1CvLiuH6LwKHo0kI,49 +mpl_toolkits/mplot3d/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/art3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/axes3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/axis3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/proj3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/art3d.py,sha256=ovygkCV4P7dQ9xIicczbJyyAbaEy2C1HYnzjuV3COus,50548 +mpl_toolkits/mplot3d/axes3d.py,sha256=hCBda2QB6_OYyKEJOG0paPEgDPmLIcrEMkZuKqOUFho,157921 +mpl_toolkits/mplot3d/axis3d.py,sha256=eQPWo2TKbRsPY8JUt5drjBbyRdVKBgeT2EF4HGg51oU,29327 +mpl_toolkits/mplot3d/proj3d.py,sha256=6Hm6WPzeu_wjfeR8afrQ1nCfjS0p3wjvoSIxJWVlS0s,6349 +mpl_toolkits/mplot3d/tests/__init__.py,sha256=sKLxL9jEJBX7eh5OumtXSOnTriPrJUkujTHFtnJVFrM,365 +mpl_toolkits/mplot3d/tests/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/conftest.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/test_art3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/test_axes3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/test_legend3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/conftest.py,sha256=zB61sy90X97YJ16mIGiuaEAaBIjBEzRAK_qfSCichQM,147 +mpl_toolkits/mplot3d/tests/test_art3d.py,sha256=AdFcJf0vz1_b-PIn2cVonPHF1vZw2EOHvmjdgBRa5Yo,3317 +mpl_toolkits/mplot3d/tests/test_axes3d.py,sha256=Xf5M1DipyVfsBBF8dA3AXg5U0Hts194EfAh57VSupE4,92348 +mpl_toolkits/mplot3d/tests/test_legend3d.py,sha256=QGPaoaucJP9KIC68g8zmk4divB_w5PtQc4DMIHMpcA8,4343 +pylab.py,sha256=zUXU0l7e7C5jmDSJbM0GLQxBun3xzuXNf1tuoZYA6Xk,110 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/REQUESTED b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/WHEEL b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..148e77b3f434fe39f597b0af17ecb80a5b72fd10 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib-3.10.6.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: meson +Root-Is-Purelib: false +Tag: cp310-cp310-win_amd64 \ No newline at end of file diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..22b58b62ff8e7f83043c9d99cc6d0dcefb0a93d2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/__init__.py @@ -0,0 +1,391 @@ +""" +Helper functions for managing the Matplotlib API. + +This documentation is only relevant for Matplotlib developers, not for users. + +.. warning:: + + This module and its submodules are for internal use only. Do not use them + in your own code. We may change the API at any time with no warning. + +""" + +import functools +import itertools +import pathlib +import re +import sys +import warnings + +from .deprecation import ( # noqa: F401 + deprecated, warn_deprecated, + rename_parameter, delete_parameter, make_keyword_only, + deprecate_method_override, deprecate_privatize_attribute, + suppress_matplotlib_deprecation_warning, + MatplotlibDeprecationWarning) + + +class classproperty: + """ + Like `property`, but also triggers on access via the class, and it is the + *class* that's passed as argument. + + Examples + -------- + :: + + class C: + @classproperty + def foo(cls): + return cls.__name__ + + assert C.foo == "C" + """ + + def __init__(self, fget, fset=None, fdel=None, doc=None): + self._fget = fget + if fset is not None or fdel is not None: + raise ValueError('classproperty only implements fget.') + self.fset = fset + self.fdel = fdel + # docs are ignored for now + self._doc = doc + + def __get__(self, instance, owner): + return self._fget(owner) + + @property + def fget(self): + return self._fget + + +# In the following check_foo() functions, the first parameter is positional-only to make +# e.g. `_api.check_isinstance([...], types=foo)` work. + +def check_isinstance(types, /, **kwargs): + """ + For each *key, value* pair in *kwargs*, check that *value* is an instance + of one of *types*; if not, raise an appropriate TypeError. + + As a special case, a ``None`` entry in *types* is treated as NoneType. + + Examples + -------- + >>> _api.check_isinstance((SomeClass, None), arg=arg) + """ + none_type = type(None) + types = ((types,) if isinstance(types, type) else + (none_type,) if types is None else + tuple(none_type if tp is None else tp for tp in types)) + + def type_name(tp): + return ("None" if tp is none_type + else tp.__qualname__ if tp.__module__ == "builtins" + else f"{tp.__module__}.{tp.__qualname__}") + + for k, v in kwargs.items(): + if not isinstance(v, types): + names = [*map(type_name, types)] + if "None" in names: # Move it to the end for better wording. + names.remove("None") + names.append("None") + raise TypeError( + "{!r} must be an instance of {}, not a {}".format( + k, + ", ".join(names[:-1]) + " or " + names[-1] + if len(names) > 1 else names[0], + type_name(type(v)))) + + +def check_in_list(values, /, *, _print_supported_values=True, **kwargs): + """ + For each *key, value* pair in *kwargs*, check that *value* is in *values*; + if not, raise an appropriate ValueError. + + Parameters + ---------- + values : iterable + Sequence of values to check on. + _print_supported_values : bool, default: True + Whether to print *values* when raising ValueError. + **kwargs : dict + *key, value* pairs as keyword arguments to find in *values*. + + Raises + ------ + ValueError + If any *value* in *kwargs* is not found in *values*. + + Examples + -------- + >>> _api.check_in_list(["foo", "bar"], arg=arg, other_arg=other_arg) + """ + if not kwargs: + raise TypeError("No argument to check!") + for key, val in kwargs.items(): + if val not in values: + msg = f"{val!r} is not a valid value for {key}" + if _print_supported_values: + msg += f"; supported values are {', '.join(map(repr, values))}" + raise ValueError(msg) + + +def check_shape(shape, /, **kwargs): + """ + For each *key, value* pair in *kwargs*, check that *value* has the shape *shape*; + if not, raise an appropriate ValueError. + + *None* in the shape is treated as a "free" size that can have any length. + e.g. (None, 2) -> (N, 2) + + The values checked must be numpy arrays. + + Examples + -------- + To check for (N, 2) shaped arrays + + >>> _api.check_shape((None, 2), arg=arg, other_arg=other_arg) + """ + for k, v in kwargs.items(): + data_shape = v.shape + + if (len(data_shape) != len(shape) + or any(s != t and t is not None for s, t in zip(data_shape, shape))): + dim_labels = iter(itertools.chain( + 'NMLKJIH', + (f"D{i}" for i in itertools.count()))) + text_shape = ", ".join([str(n) if n is not None else next(dim_labels) + for n in shape[::-1]][::-1]) + if len(shape) == 1: + text_shape += "," + + raise ValueError( + f"{k!r} must be {len(shape)}D with shape ({text_shape}), " + f"but your input has shape {v.shape}" + ) + + +def check_getitem(mapping, /, **kwargs): + """ + *kwargs* must consist of a single *key, value* pair. If *key* is in + *mapping*, return ``mapping[value]``; else, raise an appropriate + ValueError. + + Examples + -------- + >>> _api.check_getitem({"foo": "bar"}, arg=arg) + """ + if len(kwargs) != 1: + raise ValueError("check_getitem takes a single keyword argument") + (k, v), = kwargs.items() + try: + return mapping[v] + except KeyError: + raise ValueError( + f"{v!r} is not a valid value for {k}; supported values are " + f"{', '.join(map(repr, mapping))}") from None + + +def caching_module_getattr(cls): + """ + Helper decorator for implementing module-level ``__getattr__`` as a class. + + This decorator must be used at the module toplevel as follows:: + + @caching_module_getattr + class __getattr__: # The class *must* be named ``__getattr__``. + @property # Only properties are taken into account. + def name(self): ... + + The ``__getattr__`` class will be replaced by a ``__getattr__`` + function such that trying to access ``name`` on the module will + resolve the corresponding property (which may be decorated e.g. with + ``_api.deprecated`` for deprecating module globals). The properties are + all implicitly cached. Moreover, a suitable AttributeError is generated + and raised if no property with the given name exists. + """ + + assert cls.__name__ == "__getattr__" + # Don't accidentally export cls dunders. + props = {name: prop for name, prop in vars(cls).items() + if isinstance(prop, property)} + instance = cls() + + @functools.cache + def __getattr__(name): + if name in props: + return props[name].__get__(instance) + raise AttributeError( + f"module {cls.__module__!r} has no attribute {name!r}") + + return __getattr__ + + +def define_aliases(alias_d, cls=None): + """ + Class decorator for defining property aliases. + + Use as :: + + @_api.define_aliases({"property": ["alias", ...], ...}) + class C: ... + + For each property, if the corresponding ``get_property`` is defined in the + class so far, an alias named ``get_alias`` will be defined; the same will + be done for setters. If neither the getter nor the setter exists, an + exception will be raised. + + The alias map is stored as the ``_alias_map`` attribute on the class and + can be used by `.normalize_kwargs` (which assumes that higher priority + aliases come last). + """ + if cls is None: # Return the actual class decorator. + return functools.partial(define_aliases, alias_d) + + def make_alias(name): # Enforce a closure over *name*. + @functools.wraps(getattr(cls, name)) + def method(self, *args, **kwargs): + return getattr(self, name)(*args, **kwargs) + return method + + for prop, aliases in alias_d.items(): + exists = False + for prefix in ["get_", "set_"]: + if prefix + prop in vars(cls): + exists = True + for alias in aliases: + method = make_alias(prefix + prop) + method.__name__ = prefix + alias + method.__doc__ = f"Alias for `{prefix + prop}`." + setattr(cls, prefix + alias, method) + if not exists: + raise ValueError( + f"Neither getter nor setter exists for {prop!r}") + + def get_aliased_and_aliases(d): + return {*d, *(alias for aliases in d.values() for alias in aliases)} + + preexisting_aliases = getattr(cls, "_alias_map", {}) + conflicting = (get_aliased_and_aliases(preexisting_aliases) + & get_aliased_and_aliases(alias_d)) + if conflicting: + # Need to decide on conflict resolution policy. + raise NotImplementedError( + f"Parent class already defines conflicting aliases: {conflicting}") + cls._alias_map = {**preexisting_aliases, **alias_d} + return cls + + +def select_matching_signature(funcs, *args, **kwargs): + """ + Select and call the function that accepts ``*args, **kwargs``. + + *funcs* is a list of functions which should not raise any exception (other + than `TypeError` if the arguments passed do not match their signature). + + `select_matching_signature` tries to call each of the functions in *funcs* + with ``*args, **kwargs`` (in the order in which they are given). Calls + that fail with a `TypeError` are silently skipped. As soon as a call + succeeds, `select_matching_signature` returns its return value. If no + function accepts ``*args, **kwargs``, then the `TypeError` raised by the + last failing call is re-raised. + + Callers should normally make sure that any ``*args, **kwargs`` can only + bind a single *func* (to avoid any ambiguity), although this is not checked + by `select_matching_signature`. + + Notes + ----- + `select_matching_signature` is intended to help implementing + signature-overloaded functions. In general, such functions should be + avoided, except for back-compatibility concerns. A typical use pattern is + :: + + def my_func(*args, **kwargs): + params = select_matching_signature( + [lambda old1, old2: locals(), lambda new: locals()], + *args, **kwargs) + if "old1" in params: + warn_deprecated(...) + old1, old2 = params.values() # note that locals() is ordered. + else: + new, = params.values() + # do things with params + + which allows *my_func* to be called either with two parameters (*old1* and + *old2*) or a single one (*new*). Note that the new signature is given + last, so that callers get a `TypeError` corresponding to the new signature + if the arguments they passed in do not match any signature. + """ + # Rather than relying on locals() ordering, one could have just used func's + # signature (``bound = inspect.signature(func).bind(*args, **kwargs); + # bound.apply_defaults(); return bound``) but that is significantly slower. + for i, func in enumerate(funcs): + try: + return func(*args, **kwargs) + except TypeError: + if i == len(funcs) - 1: + raise + + +def nargs_error(name, takes, given): + """Generate a TypeError to be raised by function calls with wrong arity.""" + return TypeError(f"{name}() takes {takes} positional arguments but " + f"{given} were given") + + +def kwarg_error(name, kw): + """ + Generate a TypeError to be raised by function calls with wrong kwarg. + + Parameters + ---------- + name : str + The name of the calling function. + kw : str or Iterable[str] + Either the invalid keyword argument name, or an iterable yielding + invalid keyword arguments (e.g., a ``kwargs`` dict). + """ + if not isinstance(kw, str): + kw = next(iter(kw)) + return TypeError(f"{name}() got an unexpected keyword argument '{kw}'") + + +def recursive_subclasses(cls): + """Yield *cls* and direct and indirect subclasses of *cls*.""" + yield cls + for subcls in cls.__subclasses__(): + yield from recursive_subclasses(subcls) + + +def warn_external(message, category=None): + """ + `warnings.warn` wrapper that sets *stacklevel* to "outside Matplotlib". + + The original emitter of the warning can be obtained by patching this + function back to `warnings.warn`, i.e. ``_api.warn_external = + warnings.warn`` (or ``functools.partial(warnings.warn, stacklevel=2)``, + etc.). + """ + kwargs = {} + if sys.version_info[:2] >= (3, 12): + # Go to Python's `site-packages` or `lib` from an editable install. + basedir = pathlib.Path(__file__).parents[2] + kwargs['skip_file_prefixes'] = (str(basedir / 'matplotlib'), + str(basedir / 'mpl_toolkits')) + else: + frame = sys._getframe() + for stacklevel in itertools.count(1): + if frame is None: + # when called in embedded context may hit frame is None + kwargs['stacklevel'] = stacklevel + break + if not re.match(r"\A(matplotlib|mpl_toolkits)(\Z|\.(?!tests\.))", + # Work around sphinx-gallery not setting __name__. + frame.f_globals.get("__name__", "")): + kwargs['stacklevel'] = stacklevel + break + frame = frame.f_back + # preemptively break reference cycle between locals and the frame + del frame + warnings.warn(message, category, **kwargs) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/__init__.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..ea7076feac3c415be6d789e92f87020ea89db7a1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/__init__.pyi @@ -0,0 +1,59 @@ +from collections.abc import Callable, Generator, Iterable, Mapping, Sequence +from typing import Any, TypeVar, overload +from typing_extensions import Self # < Py 3.11 + +from numpy.typing import NDArray + +from .deprecation import ( # noqa: F401, re-exported API + deprecated as deprecated, + warn_deprecated as warn_deprecated, + rename_parameter as rename_parameter, + delete_parameter as delete_parameter, + make_keyword_only as make_keyword_only, + deprecate_method_override as deprecate_method_override, + deprecate_privatize_attribute as deprecate_privatize_attribute, + suppress_matplotlib_deprecation_warning as suppress_matplotlib_deprecation_warning, + MatplotlibDeprecationWarning as MatplotlibDeprecationWarning, +) + +_T = TypeVar("_T") + +class classproperty(Any): + def __init__( + self, + fget: Callable[[_T], Any], + fset: None = ..., + fdel: None = ..., + doc: str | None = None, + ): ... + @overload + def __get__(self, instance: None, owner: None) -> Self: ... + @overload + def __get__(self, instance: object, owner: type[object]) -> Any: ... + @property + def fget(self) -> Callable[[_T], Any]: ... + +def check_isinstance( + types: type | tuple[type | None, ...], /, **kwargs: Any +) -> None: ... +def check_in_list( + values: Sequence[Any], /, *, _print_supported_values: bool = ..., **kwargs: Any +) -> None: ... +def check_shape(shape: tuple[int | None, ...], /, **kwargs: NDArray) -> None: ... +def check_getitem(mapping: Mapping[Any, Any], /, **kwargs: Any) -> Any: ... +def caching_module_getattr(cls: type) -> Callable[[str], Any]: ... +@overload +def define_aliases( + alias_d: dict[str, list[str]], cls: None = ... +) -> Callable[[type[_T]], type[_T]]: ... +@overload +def define_aliases(alias_d: dict[str, list[str]], cls: type[_T]) -> type[_T]: ... +def select_matching_signature( + funcs: list[Callable], *args: Any, **kwargs: Any +) -> Any: ... +def nargs_error(name: str, takes: int | str, given: int) -> TypeError: ... +def kwarg_error(name: str, kw: str | Iterable[str]) -> TypeError: ... +def recursive_subclasses(cls: type) -> Generator[type, None, None]: ... +def warn_external( + message: str | Warning, category: type[Warning] | None = ... +) -> None: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/deprecation.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/deprecation.py new file mode 100644 index 0000000000000000000000000000000000000000..65a754bbb43d424f52e9b11f4580ba5d08dc6e59 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/deprecation.py @@ -0,0 +1,509 @@ +""" +Helper functions for deprecating parts of the Matplotlib API. + +This documentation is only relevant for Matplotlib developers, not for users. + +.. warning:: + + This module is for internal use only. Do not use it in your own code. + We may change the API at any time with no warning. + +""" + +import contextlib +import functools +import inspect +import math +import warnings + + +class MatplotlibDeprecationWarning(DeprecationWarning): + """A class for issuing deprecation warnings for Matplotlib users.""" + + +def _generate_deprecation_warning( + since, message='', name='', alternative='', pending=False, obj_type='', + addendum='', *, removal=''): + if pending: + if removal: + raise ValueError("A pending deprecation cannot have a scheduled removal") + elif removal == '': + macro, meso, *_ = since.split('.') + removal = f'{macro}.{int(meso) + 2}' + if not message: + message = ( + ("The %(name)s %(obj_type)s" if obj_type else "%(name)s") + + (" will be deprecated in a future version" if pending else + (" was deprecated in Matplotlib %(since)s" + + (" and will be removed in %(removal)s" if removal else ""))) + + "." + + (" Use %(alternative)s instead." if alternative else "") + + (" %(addendum)s" if addendum else "")) + warning_cls = PendingDeprecationWarning if pending else MatplotlibDeprecationWarning + return warning_cls(message % dict( + func=name, name=name, obj_type=obj_type, since=since, removal=removal, + alternative=alternative, addendum=addendum)) + + +def warn_deprecated( + since, *, message='', name='', alternative='', pending=False, + obj_type='', addendum='', removal=''): + """ + Display a standardized deprecation. + + Parameters + ---------- + since : str + The release at which this API became deprecated. + message : str, optional + Override the default deprecation message. The ``%(since)s``, + ``%(name)s``, ``%(alternative)s``, ``%(obj_type)s``, ``%(addendum)s``, + and ``%(removal)s`` format specifiers will be replaced by the values + of the respective arguments passed to this function. + name : str, optional + The name of the deprecated object. + alternative : str, optional + An alternative API that the user may use in place of the deprecated + API. The deprecation warning will tell the user about this alternative + if provided. + pending : bool, optional + If True, uses a PendingDeprecationWarning instead of a + DeprecationWarning. Cannot be used together with *removal*. + obj_type : str, optional + The object type being deprecated. + addendum : str, optional + Additional text appended directly to the final message. + removal : str, optional + The expected removal version. With the default (an empty string), a + removal version is automatically computed from *since*. Set to other + Falsy values to not schedule a removal date. Cannot be used together + with *pending*. + + Examples + -------- + :: + + # To warn of the deprecation of "matplotlib.name_of_module" + warn_deprecated('1.4.0', name='matplotlib.name_of_module', + obj_type='module') + """ + warning = _generate_deprecation_warning( + since, message, name, alternative, pending, obj_type, addendum, + removal=removal) + from . import warn_external + warn_external(warning, category=MatplotlibDeprecationWarning) + + +def deprecated(since, *, message='', name='', alternative='', pending=False, + obj_type=None, addendum='', removal=''): + """ + Decorator to mark a function, a class, or a property as deprecated. + + When deprecating a classmethod, a staticmethod, or a property, the + ``@deprecated`` decorator should go *under* ``@classmethod`` and + ``@staticmethod`` (i.e., `deprecated` should directly decorate the + underlying callable), but *over* ``@property``. + + When deprecating a class ``C`` intended to be used as a base class in a + multiple inheritance hierarchy, ``C`` *must* define an ``__init__`` method + (if ``C`` instead inherited its ``__init__`` from its own base class, then + ``@deprecated`` would mess up ``__init__`` inheritance when installing its + own (deprecation-emitting) ``C.__init__``). + + Parameters are the same as for `warn_deprecated`, except that *obj_type* + defaults to 'class' if decorating a class, 'attribute' if decorating a + property, and 'function' otherwise. + + Examples + -------- + :: + + @deprecated('1.4.0') + def the_function_to_deprecate(): + pass + """ + + def deprecate(obj, message=message, name=name, alternative=alternative, + pending=pending, obj_type=obj_type, addendum=addendum): + from matplotlib._api import classproperty + + if isinstance(obj, type): + if obj_type is None: + obj_type = "class" + func = obj.__init__ + name = name or obj.__name__ + old_doc = obj.__doc__ + + def finalize(wrapper, new_doc): + try: + obj.__doc__ = new_doc + except AttributeError: # Can't set on some extension objects. + pass + obj.__init__ = functools.wraps(obj.__init__)(wrapper) + return obj + + elif isinstance(obj, (property, classproperty)): + if obj_type is None: + obj_type = "attribute" + func = None + name = name or obj.fget.__name__ + old_doc = obj.__doc__ + + class _deprecated_property(type(obj)): + def __get__(self, instance, owner=None): + if instance is not None or owner is not None \ + and isinstance(self, classproperty): + emit_warning() + return super().__get__(instance, owner) + + def __set__(self, instance, value): + if instance is not None: + emit_warning() + return super().__set__(instance, value) + + def __delete__(self, instance): + if instance is not None: + emit_warning() + return super().__delete__(instance) + + def __set_name__(self, owner, set_name): + nonlocal name + if name == "": + name = set_name + + def finalize(_, new_doc): + return _deprecated_property( + fget=obj.fget, fset=obj.fset, fdel=obj.fdel, doc=new_doc) + + else: + if obj_type is None: + obj_type = "function" + func = obj + name = name or obj.__name__ + old_doc = func.__doc__ + + def finalize(wrapper, new_doc): + wrapper = functools.wraps(func)(wrapper) + wrapper.__doc__ = new_doc + return wrapper + + def emit_warning(): + warn_deprecated( + since, message=message, name=name, alternative=alternative, + pending=pending, obj_type=obj_type, addendum=addendum, + removal=removal) + + def wrapper(*args, **kwargs): + emit_warning() + return func(*args, **kwargs) + + old_doc = inspect.cleandoc(old_doc or '').strip('\n') + + notes_header = '\nNotes\n-----' + second_arg = ' '.join([t.strip() for t in + (message, f"Use {alternative} instead." + if alternative else "", addendum) if t]) + new_doc = (f"[*Deprecated*] {old_doc}\n" + f"{notes_header if notes_header not in old_doc else ''}\n" + f".. deprecated:: {since}\n" + f" {second_arg}") + + if not old_doc: + # This is to prevent a spurious 'unexpected unindent' warning from + # docutils when the original docstring was blank. + new_doc += r'\ ' + + return finalize(wrapper, new_doc) + + return deprecate + + +class deprecate_privatize_attribute: + """ + Helper to deprecate public access to an attribute (or method). + + This helper should only be used at class scope, as follows:: + + class Foo: + attr = _deprecate_privatize_attribute(*args, **kwargs) + + where *all* parameters are forwarded to `deprecated`. This form makes + ``attr`` a property which forwards read and write access to ``self._attr`` + (same name but with a leading underscore), with a deprecation warning. + Note that the attribute name is derived from *the name this helper is + assigned to*. This helper also works for deprecating methods. + """ + + def __init__(self, *args, **kwargs): + self.deprecator = deprecated(*args, **kwargs) + + def __set_name__(self, owner, name): + setattr(owner, name, self.deprecator( + property(lambda self: getattr(self, f"_{name}"), + lambda self, value: setattr(self, f"_{name}", value)), + name=name)) + + +# Used by _copy_docstring_and_deprecators to redecorate pyplot wrappers and +# boilerplate.py to retrieve original signatures. It may seem natural to store +# this information as an attribute on the wrapper, but if the wrapper gets +# itself functools.wraps()ed, then such attributes are silently propagated to +# the outer wrapper, which is not desired. +DECORATORS = {} + + +def rename_parameter(since, old, new, func=None): + """ + Decorator indicating that parameter *old* of *func* is renamed to *new*. + + The actual implementation of *func* should use *new*, not *old*. If *old* + is passed to *func*, a DeprecationWarning is emitted, and its value is + used, even if *new* is also passed by keyword (this is to simplify pyplot + wrapper functions, which always pass *new* explicitly to the Axes method). + If *new* is also passed but positionally, a TypeError will be raised by the + underlying function during argument binding. + + Examples + -------- + :: + + @_api.rename_parameter("3.1", "bad_name", "good_name") + def func(good_name): ... + """ + + decorator = functools.partial(rename_parameter, since, old, new) + + if func is None: + return decorator + + signature = inspect.signature(func) + assert old not in signature.parameters, ( + f"Matplotlib internal error: {old!r} cannot be a parameter for " + f"{func.__name__}()") + assert new in signature.parameters, ( + f"Matplotlib internal error: {new!r} must be a parameter for " + f"{func.__name__}()") + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if old in kwargs: + warn_deprecated( + since, message=f"The {old!r} parameter of {func.__name__}() " + f"has been renamed {new!r} since Matplotlib {since}; support " + f"for the old name will be dropped in %(removal)s.") + kwargs[new] = kwargs.pop(old) + return func(*args, **kwargs) + + # wrapper() must keep the same documented signature as func(): if we + # instead made both *old* and *new* appear in wrapper()'s signature, they + # would both show up in the pyplot function for an Axes method as well and + # pyplot would explicitly pass both arguments to the Axes method. + + DECORATORS[wrapper] = decorator + return wrapper + + +class _deprecated_parameter_class: + def __repr__(self): + return "" + + +_deprecated_parameter = _deprecated_parameter_class() + + +def delete_parameter(since, name, func=None, **kwargs): + """ + Decorator indicating that parameter *name* of *func* is being deprecated. + + The actual implementation of *func* should keep the *name* parameter in its + signature, or accept a ``**kwargs`` argument (through which *name* would be + passed). + + Parameters that come after the deprecated parameter effectively become + keyword-only (as they cannot be passed positionally without triggering the + DeprecationWarning on the deprecated parameter), and should be marked as + such after the deprecation period has passed and the deprecated parameter + is removed. + + Parameters other than *since*, *name*, and *func* are keyword-only and + forwarded to `.warn_deprecated`. + + Examples + -------- + :: + + @_api.delete_parameter("3.1", "unused") + def func(used_arg, other_arg, unused, more_args): ... + """ + + decorator = functools.partial(delete_parameter, since, name, **kwargs) + + if func is None: + return decorator + + signature = inspect.signature(func) + # Name of `**kwargs` parameter of the decorated function, typically + # "kwargs" if such a parameter exists, or None if the decorated function + # doesn't accept `**kwargs`. + kwargs_name = next((param.name for param in signature.parameters.values() + if param.kind == inspect.Parameter.VAR_KEYWORD), None) + if name in signature.parameters: + kind = signature.parameters[name].kind + is_varargs = kind is inspect.Parameter.VAR_POSITIONAL + is_varkwargs = kind is inspect.Parameter.VAR_KEYWORD + if not is_varargs and not is_varkwargs: + name_idx = ( + # Deprecated parameter can't be passed positionally. + math.inf if kind is inspect.Parameter.KEYWORD_ONLY + # If call site has no more than this number of parameters, the + # deprecated parameter can't have been passed positionally. + else [*signature.parameters].index(name)) + func.__signature__ = signature = signature.replace(parameters=[ + param.replace(default=_deprecated_parameter) + if param.name == name else param + for param in signature.parameters.values()]) + else: + name_idx = -1 # Deprecated parameter can always have been passed. + else: + is_varargs = is_varkwargs = False + # Deprecated parameter can't be passed positionally. + name_idx = math.inf + assert kwargs_name, ( + f"Matplotlib internal error: {name!r} must be a parameter for " + f"{func.__name__}()") + + addendum = kwargs.pop('addendum', None) + + @functools.wraps(func) + def wrapper(*inner_args, **inner_kwargs): + if len(inner_args) <= name_idx and name not in inner_kwargs: + # Early return in the simple, non-deprecated case (much faster than + # calling bind()). + return func(*inner_args, **inner_kwargs) + arguments = signature.bind(*inner_args, **inner_kwargs).arguments + if is_varargs and arguments.get(name): + warn_deprecated( + since, message=f"Additional positional arguments to " + f"{func.__name__}() are deprecated since %(since)s and " + f"support for them will be removed in %(removal)s.") + elif is_varkwargs and arguments.get(name): + warn_deprecated( + since, message=f"Additional keyword arguments to " + f"{func.__name__}() are deprecated since %(since)s and " + f"support for them will be removed in %(removal)s.") + # We cannot just check `name not in arguments` because the pyplot + # wrappers always pass all arguments explicitly. + elif any(name in d and d[name] != _deprecated_parameter + for d in [arguments, arguments.get(kwargs_name, {})]): + deprecation_addendum = ( + f"If any parameter follows {name!r}, they should be passed as " + f"keyword, not positionally.") + warn_deprecated( + since, + name=repr(name), + obj_type=f"parameter of {func.__name__}()", + addendum=(addendum + " " + deprecation_addendum) if addendum + else deprecation_addendum, + **kwargs) + return func(*inner_args, **inner_kwargs) + + DECORATORS[wrapper] = decorator + return wrapper + + +def make_keyword_only(since, name, func=None): + """ + Decorator indicating that passing parameter *name* (or any of the following + ones) positionally to *func* is being deprecated. + + When used on a method that has a pyplot wrapper, this should be the + outermost decorator, so that :file:`boilerplate.py` can access the original + signature. + """ + + decorator = functools.partial(make_keyword_only, since, name) + + if func is None: + return decorator + + signature = inspect.signature(func) + POK = inspect.Parameter.POSITIONAL_OR_KEYWORD + KWO = inspect.Parameter.KEYWORD_ONLY + assert (name in signature.parameters + and signature.parameters[name].kind == POK), ( + f"Matplotlib internal error: {name!r} must be a positional-or-keyword " + f"parameter for {func.__name__}(). If this error happens on a function with a " + f"pyplot wrapper, make sure make_keyword_only() is the outermost decorator.") + names = [*signature.parameters] + name_idx = names.index(name) + kwonly = [name for name in names[name_idx:] + if signature.parameters[name].kind == POK] + + @functools.wraps(func) + def wrapper(*args, **kwargs): + # Don't use signature.bind here, as it would fail when stacked with + # rename_parameter and an "old" argument name is passed in + # (signature.bind would fail, but the actual call would succeed). + if len(args) > name_idx: + warn_deprecated( + since, message="Passing the %(name)s %(obj_type)s " + "positionally is deprecated since Matplotlib %(since)s; the " + "parameter will become keyword-only in %(removal)s.", + name=name, obj_type=f"parameter of {func.__name__}()") + return func(*args, **kwargs) + + # Don't modify *func*'s signature, as boilerplate.py needs it. + wrapper.__signature__ = signature.replace(parameters=[ + param.replace(kind=KWO) if param.name in kwonly else param + for param in signature.parameters.values()]) + DECORATORS[wrapper] = decorator + return wrapper + + +def deprecate_method_override(method, obj, *, allow_empty=False, **kwargs): + """ + Return ``obj.method`` with a deprecation if it was overridden, else None. + + Parameters + ---------- + method + An unbound method, i.e. an expression of the form + ``Class.method_name``. Remember that within the body of a method, one + can always use ``__class__`` to refer to the class that is currently + being defined. + obj + Either an object of the class where *method* is defined, or a subclass + of that class. + allow_empty : bool, default: False + Whether to allow overrides by "empty" methods without emitting a + warning. + **kwargs + Additional parameters passed to `warn_deprecated` to generate the + deprecation warning; must at least include the "since" key. + """ + + def empty(): pass + def empty_with_docstring(): """doc""" + + name = method.__name__ + bound_child = getattr(obj, name) + bound_base = ( + method # If obj is a class, then we need to use unbound methods. + if isinstance(bound_child, type(empty)) and isinstance(obj, type) + else method.__get__(obj)) + if (bound_child != bound_base + and (not allow_empty + or (getattr(getattr(bound_child, "__code__", None), + "co_code", None) + not in [empty.__code__.co_code, + empty_with_docstring.__code__.co_code]))): + warn_deprecated(**{"name": name, "obj_type": "method", **kwargs}) + return bound_child + return None + + +@contextlib.contextmanager +def suppress_matplotlib_deprecation_warning(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", MatplotlibDeprecationWarning) + yield diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/deprecation.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/deprecation.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e050290662d925fd68bd534881512e86db13a4bb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/_api/deprecation.pyi @@ -0,0 +1,75 @@ +from collections.abc import Callable +import contextlib +from typing import Any, Literal, ParamSpec, TypedDict, TypeVar, overload +from typing_extensions import ( + Unpack, # < Py 3.11 +) + +_P = ParamSpec("_P") +_R = TypeVar("_R") +_T = TypeVar("_T") + +class MatplotlibDeprecationWarning(DeprecationWarning): ... + +class DeprecationKwargs(TypedDict, total=False): + message: str + alternative: str + pending: bool + obj_type: str + addendum: str + removal: str | Literal[False] + +class NamedDeprecationKwargs(DeprecationKwargs, total=False): + name: str + +def warn_deprecated(since: str, **kwargs: Unpack[NamedDeprecationKwargs]) -> None: ... +def deprecated( + since: str, **kwargs: Unpack[NamedDeprecationKwargs] +) -> Callable[[_T], _T]: ... + +class deprecate_privatize_attribute(Any): + def __init__(self, since: str, **kwargs: Unpack[NamedDeprecationKwargs]): ... + def __set_name__(self, owner: type[object], name: str) -> None: ... + +DECORATORS: dict[Callable, Callable] = ... + +@overload +def rename_parameter( + since: str, old: str, new: str, func: None = ... +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... +@overload +def rename_parameter( + since: str, old: str, new: str, func: Callable[_P, _R] +) -> Callable[_P, _R]: ... + +class _deprecated_parameter_class: ... + +_deprecated_parameter: _deprecated_parameter_class + +@overload +def delete_parameter( + since: str, name: str, func: None = ..., **kwargs: Unpack[DeprecationKwargs] +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... +@overload +def delete_parameter( + since: str, name: str, func: Callable[_P, _R], **kwargs: Unpack[DeprecationKwargs] +) -> Callable[_P, _R]: ... +@overload +def make_keyword_only( + since: str, name: str, func: None = ... +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... +@overload +def make_keyword_only( + since: str, name: str, func: Callable[_P, _R] +) -> Callable[_P, _R]: ... +def deprecate_method_override( + method: Callable[_P, _R], + obj: object | type, + *, + allow_empty: bool = ..., + since: str, + **kwargs: Unpack[NamedDeprecationKwargs] +) -> Callable[_P, _R]: ... +def suppress_matplotlib_deprecation_warning() -> ( + contextlib.AbstractContextManager[None] +): ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8cce4fe4558d30d6c71ec7a63e347a0fb9bf7ada --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/__init__.py @@ -0,0 +1,10 @@ +from pathlib import Path + + +# Check that the test directories exist. +if not (Path(__file__).parent / 'baseline_images').exists(): + raise OSError( + 'The baseline image directory does not exist. ' + 'This is most likely because the test data is not installed. ' + 'You may need to install matplotlib from source to get the ' + 'test data.') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_simplification.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_simplification.py new file mode 100644 index 0000000000000000000000000000000000000000..58b4d7a9b24fea5ef8c54f51c6f7095ce19e33a3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_simplification.py @@ -0,0 +1,571 @@ +import base64 +import io +import platform + +import numpy as np +from numpy.testing import assert_array_almost_equal, assert_array_equal + +import pytest + +from matplotlib.testing.decorators import ( + check_figures_equal, image_comparison, remove_ticks_and_titles) +import matplotlib.pyplot as plt + +from matplotlib import patches, transforms +from matplotlib.path import Path + + +# NOTE: All of these tests assume that path.simplify is set to True +# (the default) + +@image_comparison(['clipping'], remove_text=True) +def test_clipping(): + t = np.arange(0.0, 2.0, 0.01) + s = np.sin(2*np.pi*t) + + fig, ax = plt.subplots() + ax.plot(t, s, linewidth=1.0) + ax.set_ylim((-0.20, -0.28)) + + +@image_comparison(['overflow'], remove_text=True, + tol=0 if platform.machine() == 'x86_64' else 0.007) +def test_overflow(): + x = np.array([1.0, 2.0, 3.0, 2.0e5]) + y = np.arange(len(x)) + + fig, ax = plt.subplots() + ax.plot(x, y) + ax.set_xlim(2, 6) + + +@image_comparison(['clipping_diamond'], remove_text=True) +def test_diamond(): + x = np.array([0.0, 1.0, 0.0, -1.0, 0.0]) + y = np.array([1.0, 0.0, -1.0, 0.0, 1.0]) + + fig, ax = plt.subplots() + ax.plot(x, y) + ax.set_xlim(-0.6, 0.6) + ax.set_ylim(-0.6, 0.6) + + +def test_clipping_out_of_bounds(): + # Should work on a Path *without* codes. + path = Path([(0, 0), (1, 2), (2, 1)]) + simplified = path.cleaned(clip=(10, 10, 20, 20)) + assert_array_equal(simplified.vertices, [(0, 0)]) + assert simplified.codes == [Path.STOP] + + # Should work on a Path *with* codes, and no curves. + path = Path([(0, 0), (1, 2), (2, 1)], + [Path.MOVETO, Path.LINETO, Path.LINETO]) + simplified = path.cleaned(clip=(10, 10, 20, 20)) + assert_array_equal(simplified.vertices, [(0, 0)]) + assert simplified.codes == [Path.STOP] + + # A Path with curves does not do any clipping yet. + path = Path([(0, 0), (1, 2), (2, 3)], + [Path.MOVETO, Path.CURVE3, Path.CURVE3]) + simplified = path.cleaned() + simplified_clipped = path.cleaned(clip=(10, 10, 20, 20)) + assert_array_equal(simplified.vertices, simplified_clipped.vertices) + assert_array_equal(simplified.codes, simplified_clipped.codes) + + +def test_noise(): + np.random.seed(0) + x = np.random.uniform(size=50000) * 50 + + fig, ax = plt.subplots() + p1 = ax.plot(x, solid_joinstyle='round', linewidth=2.0) + + # Ensure that the path's transform takes the new axes limits into account. + fig.canvas.draw() + path = p1[0].get_path() + transform = p1[0].get_transform() + path = transform.transform_path(path) + simplified = path.cleaned(simplify=True) + + assert simplified.vertices.size == 25512 + + +def test_antiparallel_simplification(): + def _get_simplified(x, y): + fig, ax = plt.subplots() + p1 = ax.plot(x, y) + + path = p1[0].get_path() + transform = p1[0].get_transform() + path = transform.transform_path(path) + simplified = path.cleaned(simplify=True) + simplified = transform.inverted().transform_path(simplified) + + return simplified + + # test ending on a maximum + x = [0, 0, 0, 0, 0, 1] + y = [.5, 1, -1, 1, 2, .5] + + simplified = _get_simplified(x, y) + + assert_array_almost_equal([[0., 0.5], + [0., -1.], + [0., 2.], + [1., 0.5]], + simplified.vertices[:-2, :]) + + # test ending on a minimum + x = [0, 0, 0, 0, 0, 1] + y = [.5, 1, -1, 1, -2, .5] + + simplified = _get_simplified(x, y) + + assert_array_almost_equal([[0., 0.5], + [0., 1.], + [0., -2.], + [1., 0.5]], + simplified.vertices[:-2, :]) + + # test ending in between + x = [0, 0, 0, 0, 0, 1] + y = [.5, 1, -1, 1, 0, .5] + + simplified = _get_simplified(x, y) + + assert_array_almost_equal([[0., 0.5], + [0., 1.], + [0., -1.], + [0., 0.], + [1., 0.5]], + simplified.vertices[:-2, :]) + + # test no anti-parallel ending at max + x = [0, 0, 0, 0, 0, 1] + y = [.5, 1, 2, 1, 3, .5] + + simplified = _get_simplified(x, y) + + assert_array_almost_equal([[0., 0.5], + [0., 3.], + [1., 0.5]], + simplified.vertices[:-2, :]) + + # test no anti-parallel ending in middle + x = [0, 0, 0, 0, 0, 1] + y = [.5, 1, 2, 1, 1, .5] + + simplified = _get_simplified(x, y) + + assert_array_almost_equal([[0., 0.5], + [0., 2.], + [0., 1.], + [1., 0.5]], + simplified.vertices[:-2, :]) + + +# Only consider angles in 0 <= angle <= pi/2, otherwise +# using min/max will get the expected results out of order: +# min/max for simplification code depends on original vector, +# and if angle is outside above range then simplification +# min/max will be opposite from actual min/max. +@pytest.mark.parametrize('angle', [0, np.pi/4, np.pi/3, np.pi/2]) +@pytest.mark.parametrize('offset', [0, .5]) +def test_angled_antiparallel(angle, offset): + scale = 5 + np.random.seed(19680801) + # get 15 random offsets + # TODO: guarantee offset > 0 results in some offsets < 0 + vert_offsets = (np.random.rand(15) - offset) * scale + # always start at 0 so rotation makes sense + vert_offsets[0] = 0 + # always take the first step the same direction + vert_offsets[1] = 1 + # compute points along a diagonal line + x = np.sin(angle) * vert_offsets + y = np.cos(angle) * vert_offsets + + # will check these later + x_max = x[1:].max() + x_min = x[1:].min() + + y_max = y[1:].max() + y_min = y[1:].min() + + if offset > 0: + p_expected = Path([[0, 0], + [x_max, y_max], + [x_min, y_min], + [x[-1], y[-1]], + [0, 0]], + codes=[1, 2, 2, 2, 0]) + + else: + p_expected = Path([[0, 0], + [x_max, y_max], + [x[-1], y[-1]], + [0, 0]], + codes=[1, 2, 2, 0]) + + p = Path(np.vstack([x, y]).T) + p2 = p.cleaned(simplify=True) + + assert_array_almost_equal(p_expected.vertices, + p2.vertices) + assert_array_equal(p_expected.codes, p2.codes) + + +def test_sine_plus_noise(): + np.random.seed(0) + x = (np.sin(np.linspace(0, np.pi * 2.0, 50000)) + + np.random.uniform(size=50000) * 0.01) + + fig, ax = plt.subplots() + p1 = ax.plot(x, solid_joinstyle='round', linewidth=2.0) + + # Ensure that the path's transform takes the new axes limits into account. + fig.canvas.draw() + path = p1[0].get_path() + transform = p1[0].get_transform() + path = transform.transform_path(path) + simplified = path.cleaned(simplify=True) + + assert simplified.vertices.size == 25240 + + +@image_comparison(['simplify_curve'], remove_text=True, tol=0.017) +def test_simplify_curve(): + pp1 = patches.PathPatch( + Path([(0, 0), (1, 0), (1, 1), (np.nan, 1), (0, 0), (2, 0), (2, 2), + (0, 0)], + [Path.MOVETO, Path.CURVE3, Path.CURVE3, Path.CURVE3, Path.CURVE3, + Path.CURVE3, Path.CURVE3, Path.CLOSEPOLY]), + fc="none") + + fig, ax = plt.subplots() + ax.add_patch(pp1) + ax.set_xlim((0, 2)) + ax.set_ylim((0, 2)) + + +@check_figures_equal() +def test_closed_path_nan_removal(fig_test, fig_ref): + ax_test = fig_test.subplots(2, 2).flatten() + ax_ref = fig_ref.subplots(2, 2).flatten() + + # NaN on the first point also removes the last point, because it's closed. + path = Path( + [[-3, np.nan], [3, -3], [3, 3], [-3, 3], [-3, -3]], + [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]) + ax_test[0].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-3, np.nan], [3, -3], [3, 3], [-3, 3], [-3, np.nan]], + [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.LINETO]) + ax_ref[0].add_patch(patches.PathPatch(path, facecolor='none')) + + # NaN on second-last point should not re-close. + path = Path( + [[-2, -2], [2, -2], [2, 2], [-2, np.nan], [-2, -2]], + [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]) + ax_test[0].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-2, -2], [2, -2], [2, 2], [-2, np.nan], [-2, -2]], + [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.LINETO]) + ax_ref[0].add_patch(patches.PathPatch(path, facecolor='none')) + + # Test multiple loops in a single path (with same paths as above). + path = Path( + [[-3, np.nan], [3, -3], [3, 3], [-3, 3], [-3, -3], + [-2, -2], [2, -2], [2, 2], [-2, np.nan], [-2, -2]], + [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY, + Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]) + ax_test[1].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-3, np.nan], [3, -3], [3, 3], [-3, 3], [-3, np.nan], + [-2, -2], [2, -2], [2, 2], [-2, np.nan], [-2, -2]], + [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.LINETO, + Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.LINETO]) + ax_ref[1].add_patch(patches.PathPatch(path, facecolor='none')) + + # NaN in first point of CURVE3 should not re-close, and hide entire curve. + path = Path( + [[-1, -1], [1, -1], [1, np.nan], [0, 1], [-1, 1], [-1, -1]], + [Path.MOVETO, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, + Path.CLOSEPOLY]) + ax_test[2].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-1, -1], [1, -1], [1, np.nan], [0, 1], [-1, 1], [-1, -1]], + [Path.MOVETO, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, + Path.CLOSEPOLY]) + ax_ref[2].add_patch(patches.PathPatch(path, facecolor='none')) + + # NaN in second point of CURVE3 should not re-close, and hide entire curve + # plus next line segment. + path = Path( + [[-3, -3], [3, -3], [3, 0], [0, np.nan], [-3, 3], [-3, -3]], + [Path.MOVETO, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, + Path.LINETO]) + ax_test[2].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-3, -3], [3, -3], [3, 0], [0, np.nan], [-3, 3], [-3, -3]], + [Path.MOVETO, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, + Path.LINETO]) + ax_ref[2].add_patch(patches.PathPatch(path, facecolor='none')) + + # NaN in first point of CURVE4 should not re-close, and hide entire curve. + path = Path( + [[-1, -1], [1, -1], [1, np.nan], [0, 0], [0, 1], [-1, 1], [-1, -1]], + [Path.MOVETO, Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, + Path.LINETO, Path.CLOSEPOLY]) + ax_test[3].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-1, -1], [1, -1], [1, np.nan], [0, 0], [0, 1], [-1, 1], [-1, -1]], + [Path.MOVETO, Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, + Path.LINETO, Path.CLOSEPOLY]) + ax_ref[3].add_patch(patches.PathPatch(path, facecolor='none')) + + # NaN in second point of CURVE4 should not re-close, and hide entire curve. + path = Path( + [[-2, -2], [2, -2], [2, 0], [0, np.nan], [0, 2], [-2, 2], [-2, -2]], + [Path.MOVETO, Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, + Path.LINETO, Path.LINETO]) + ax_test[3].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-2, -2], [2, -2], [2, 0], [0, np.nan], [0, 2], [-2, 2], [-2, -2]], + [Path.MOVETO, Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, + Path.LINETO, Path.LINETO]) + ax_ref[3].add_patch(patches.PathPatch(path, facecolor='none')) + + # NaN in third point of CURVE4 should not re-close, and hide entire curve + # plus next line segment. + path = Path( + [[-3, -3], [3, -3], [3, 0], [0, 0], [0, np.nan], [-3, 3], [-3, -3]], + [Path.MOVETO, Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, + Path.LINETO, Path.LINETO]) + ax_test[3].add_patch(patches.PathPatch(path, facecolor='none')) + path = Path( + [[-3, -3], [3, -3], [3, 0], [0, 0], [0, np.nan], [-3, 3], [-3, -3]], + [Path.MOVETO, Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, + Path.LINETO, Path.LINETO]) + ax_ref[3].add_patch(patches.PathPatch(path, facecolor='none')) + + # Keep everything clean. + for ax in [*ax_test.flat, *ax_ref.flat]: + ax.set(xlim=(-3.5, 3.5), ylim=(-3.5, 3.5)) + remove_ticks_and_titles(fig_test) + remove_ticks_and_titles(fig_ref) + + +@check_figures_equal() +def test_closed_path_clipping(fig_test, fig_ref): + vertices = [] + for roll in range(8): + offset = 0.1 * roll + 0.1 + + # A U-like pattern. + pattern = [ + [-0.5, 1.5], [-0.5, -0.5], [1.5, -0.5], [1.5, 1.5], # Outer square + # With a notch in the top. + [1 - offset / 2, 1.5], [1 - offset / 2, offset], + [offset / 2, offset], [offset / 2, 1.5], + ] + + # Place the initial/final point anywhere in/out of the clipping area. + pattern = np.roll(pattern, roll, axis=0) + pattern = np.concatenate((pattern, pattern[:1, :])) + + vertices.append(pattern) + + # Multiple subpaths are used here to ensure they aren't broken by closed + # loop clipping. + codes = np.full(len(vertices[0]), Path.LINETO) + codes[0] = Path.MOVETO + codes[-1] = Path.CLOSEPOLY + codes = np.tile(codes, len(vertices)) + vertices = np.concatenate(vertices) + + fig_test.set_size_inches((5, 5)) + path = Path(vertices, codes) + fig_test.add_artist(patches.PathPatch(path, facecolor='none')) + + # For reference, we draw the same thing, but unclosed by using a line to + # the last point only. + fig_ref.set_size_inches((5, 5)) + codes = codes.copy() + codes[codes == Path.CLOSEPOLY] = Path.LINETO + path = Path(vertices, codes) + fig_ref.add_artist(patches.PathPatch(path, facecolor='none')) + + +@image_comparison(['hatch_simplify'], remove_text=True) +def test_hatch(): + fig, ax = plt.subplots() + ax.add_patch(plt.Rectangle((0, 0), 1, 1, fill=False, hatch="/")) + ax.set_xlim((0.45, 0.55)) + ax.set_ylim((0.45, 0.55)) + + +@image_comparison(['fft_peaks'], remove_text=True) +def test_fft_peaks(): + fig, ax = plt.subplots() + t = np.arange(65536) + p1 = ax.plot(abs(np.fft.fft(np.sin(2*np.pi*.01*t)*np.blackman(len(t))))) + + # Ensure that the path's transform takes the new axes limits into account. + fig.canvas.draw() + path = p1[0].get_path() + transform = p1[0].get_transform() + path = transform.transform_path(path) + simplified = path.cleaned(simplify=True) + + assert simplified.vertices.size == 36 + + +def test_start_with_moveto(): + # Should be entirely clipped away to a single MOVETO + data = b""" +ZwAAAAku+v9UAQAA+Tj6/z8CAADpQ/r/KAMAANlO+v8QBAAAyVn6//UEAAC6ZPr/2gUAAKpv+v+8 +BgAAm3r6/50HAACLhfr/ewgAAHyQ+v9ZCQAAbZv6/zQKAABepvr/DgsAAE+x+v/lCwAAQLz6/7wM +AAAxx/r/kA0AACPS+v9jDgAAFN36/zQPAAAF6Pr/AxAAAPfy+v/QEAAA6f36/5wRAADbCPv/ZhIA +AMwT+/8uEwAAvh77//UTAACwKfv/uRQAAKM0+/98FQAAlT/7/z0WAACHSvv//RYAAHlV+/+7FwAA +bGD7/3cYAABea/v/MRkAAFF2+//pGQAARIH7/6AaAAA3jPv/VRsAACmX+/8JHAAAHKL7/7ocAAAP +rfv/ah0AAAO4+/8YHgAA9sL7/8QeAADpzfv/bx8AANzY+/8YIAAA0OP7/78gAADD7vv/ZCEAALf5 ++/8IIgAAqwT8/6kiAACeD/z/SiMAAJIa/P/oIwAAhiX8/4QkAAB6MPz/HyUAAG47/P+4JQAAYkb8 +/1AmAABWUfz/5SYAAEpc/P95JwAAPmf8/wsoAAAzcvz/nCgAACd9/P8qKQAAHIj8/7cpAAAQk/z/ +QyoAAAWe/P/MKgAA+aj8/1QrAADus/z/2isAAOO+/P9eLAAA2Mn8/+AsAADM1Pz/YS0AAMHf/P/g +LQAAtur8/10uAACr9fz/2C4AAKEA/f9SLwAAlgv9/8ovAACLFv3/QDAAAIAh/f+1MAAAdSz9/ycx +AABrN/3/mDEAAGBC/f8IMgAAVk39/3UyAABLWP3/4TIAAEFj/f9LMwAANm79/7MzAAAsef3/GjQA +ACKE/f9+NAAAF4/9/+E0AAANmv3/QzUAAAOl/f+iNQAA+a/9/wA2AADvuv3/XDYAAOXF/f+2NgAA +29D9/w83AADR2/3/ZjcAAMfm/f+7NwAAvfH9/w44AACz/P3/XzgAAKkH/v+vOAAAnxL+//04AACW +Hf7/SjkAAIwo/v+UOQAAgjP+/905AAB5Pv7/JDoAAG9J/v9pOgAAZVT+/606AABcX/7/7zoAAFJq +/v8vOwAASXX+/207AAA/gP7/qjsAADaL/v/lOwAALZb+/x48AAAjof7/VTwAABqs/v+LPAAAELf+ +/788AAAHwv7/8TwAAP7M/v8hPQAA9df+/1A9AADr4v7/fT0AAOLt/v+oPQAA2fj+/9E9AADQA/// ++T0AAMYO//8fPgAAvRn//0M+AAC0JP//ZT4AAKsv//+GPgAAojr//6U+AACZRf//wj4AAJBQ///d +PgAAh1v///c+AAB+Zv//Dz8AAHRx//8lPwAAa3z//zk/AABih///TD8AAFmS//9dPwAAUJ3//2w/ +AABHqP//ej8AAD6z//+FPwAANb7//48/AAAsyf//lz8AACPU//+ePwAAGt///6M/AAAR6v//pj8A +AAj1//+nPwAA/////w==""" + + verts = np.frombuffer(base64.decodebytes(data), dtype=' modification_times[3] + assert filecmp.cmp(range_6, plot_file(5)) + + +def test_plot_html_show_source_link(tmp_path): + parent = Path(__file__).parent + shutil.copyfile(parent / 'tinypages/conf.py', tmp_path / 'conf.py') + shutil.copytree(parent / 'tinypages/_static', tmp_path / '_static') + doctree_dir = tmp_path / 'doctrees' + (tmp_path / 'index.rst').write_text(""" +.. plot:: + + plt.plot(range(2)) +""") + # Make sure source scripts are created by default + html_dir1 = tmp_path / '_build' / 'html1' + build_sphinx_html(tmp_path, doctree_dir, html_dir1) + assert len(list(html_dir1.glob("**/index-1.py"))) == 1 + # Make sure source scripts are NOT created when + # plot_html_show_source_link` is False + html_dir2 = tmp_path / '_build' / 'html2' + build_sphinx_html(tmp_path, doctree_dir, html_dir2, + extra_args=['-D', 'plot_html_show_source_link=0']) + assert len(list(html_dir2.glob("**/index-1.py"))) == 0 + + +@pytest.mark.parametrize('plot_html_show_source_link', [0, 1]) +def test_show_source_link_true(tmp_path, plot_html_show_source_link): + # Test that a source link is generated if :show-source-link: is true, + # whether or not plot_html_show_source_link is true. + parent = Path(__file__).parent + shutil.copyfile(parent / 'tinypages/conf.py', tmp_path / 'conf.py') + shutil.copytree(parent / 'tinypages/_static', tmp_path / '_static') + doctree_dir = tmp_path / 'doctrees' + (tmp_path / 'index.rst').write_text(""" +.. plot:: + :show-source-link: true + + plt.plot(range(2)) +""") + html_dir = tmp_path / '_build' / 'html' + build_sphinx_html(tmp_path, doctree_dir, html_dir, extra_args=[ + '-D', f'plot_html_show_source_link={plot_html_show_source_link}']) + assert len(list(html_dir.glob("**/index-1.py"))) == 1 + + +@pytest.mark.parametrize('plot_html_show_source_link', [0, 1]) +def test_show_source_link_false(tmp_path, plot_html_show_source_link): + # Test that a source link is NOT generated if :show-source-link: is false, + # whether or not plot_html_show_source_link is true. + parent = Path(__file__).parent + shutil.copyfile(parent / 'tinypages/conf.py', tmp_path / 'conf.py') + shutil.copytree(parent / 'tinypages/_static', tmp_path / '_static') + doctree_dir = tmp_path / 'doctrees' + (tmp_path / 'index.rst').write_text(""" +.. plot:: + :show-source-link: false + + plt.plot(range(2)) +""") + html_dir = tmp_path / '_build' / 'html' + build_sphinx_html(tmp_path, doctree_dir, html_dir, extra_args=[ + '-D', f'plot_html_show_source_link={plot_html_show_source_link}']) + assert len(list(html_dir.glob("**/index-1.py"))) == 0 + + +def test_srcset_version(tmp_path): + shutil.copytree(Path(__file__).parent / 'tinypages', tmp_path, + dirs_exist_ok=True) + html_dir = tmp_path / '_build' / 'html' + img_dir = html_dir / '_images' + doctree_dir = tmp_path / 'doctrees' + + build_sphinx_html(tmp_path, doctree_dir, html_dir, extra_args=[ + '-D', 'plot_srcset=2x']) + + def plot_file(num, suff=''): + return img_dir / f'some_plots-{num}{suff}.png' + + # check some-plots + for ind in [1, 2, 3, 5, 7, 11, 13, 15, 17]: + assert plot_file(ind).exists() + assert plot_file(ind, suff='.2x').exists() + + assert (img_dir / 'nestedpage-index-1.png').exists() + assert (img_dir / 'nestedpage-index-1.2x.png').exists() + assert (img_dir / 'nestedpage-index-2.png').exists() + assert (img_dir / 'nestedpage-index-2.2x.png').exists() + assert (img_dir / 'nestedpage2-index-1.png').exists() + assert (img_dir / 'nestedpage2-index-1.2x.png').exists() + assert (img_dir / 'nestedpage2-index-2.png').exists() + assert (img_dir / 'nestedpage2-index-2.2x.png').exists() + + # Check html for srcset + + assert ('srcset="_images/some_plots-1.png, _images/some_plots-1.2x.png 2.00x"' + in (html_dir / 'some_plots.html').read_text(encoding='utf-8')) + + st = ('srcset="../_images/nestedpage-index-1.png, ' + '../_images/nestedpage-index-1.2x.png 2.00x"') + assert st in (html_dir / 'nestedpage/index.html').read_text(encoding='utf-8') + + st = ('srcset="../_images/nestedpage2-index-2.png, ' + '../_images/nestedpage2-index-2.2x.png 2.00x"') + assert st in (html_dir / 'nestedpage2/index.html').read_text(encoding='utf-8') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_spines.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_spines.py new file mode 100644 index 0000000000000000000000000000000000000000..99b8461ad8512029a3cb92286e175c9b37eade3b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_spines.py @@ -0,0 +1,168 @@ +import numpy as np +import pytest + +import matplotlib.pyplot as plt +from matplotlib.spines import Spines +from matplotlib.testing.decorators import check_figures_equal, image_comparison + + +def test_spine_class(): + """Test Spines and SpinesProxy in isolation.""" + class SpineMock: + def __init__(self): + self.val = None + + def set(self, **kwargs): + vars(self).update(kwargs) + + def set_val(self, val): + self.val = val + + spines_dict = { + 'left': SpineMock(), + 'right': SpineMock(), + 'top': SpineMock(), + 'bottom': SpineMock(), + } + spines = Spines(**spines_dict) + + assert spines['left'] is spines_dict['left'] + assert spines.left is spines_dict['left'] + + spines[['left', 'right']].set_val('x') + assert spines.left.val == 'x' + assert spines.right.val == 'x' + assert spines.top.val is None + assert spines.bottom.val is None + + spines[:].set_val('y') + assert all(spine.val == 'y' for spine in spines.values()) + + spines[:].set(foo='bar') + assert all(spine.foo == 'bar' for spine in spines.values()) + + with pytest.raises(AttributeError, match='foo'): + spines.foo + with pytest.raises(KeyError, match='foo'): + spines['foo'] + with pytest.raises(KeyError, match='foo, bar'): + spines[['left', 'foo', 'right', 'bar']] + with pytest.raises(ValueError, match='single list'): + spines['left', 'right'] + with pytest.raises(ValueError, match='Spines does not support slicing'): + spines['left':'right'] + with pytest.raises(ValueError, match='Spines does not support slicing'): + spines['top':] + + +@image_comparison(['spines_axes_positions.png']) +def test_spines_axes_positions(): + # SF bug 2852168 + fig = plt.figure() + x = np.linspace(0, 2*np.pi, 100) + y = 2*np.sin(x) + ax = fig.add_subplot(1, 1, 1) + ax.set_title('centered spines') + ax.plot(x, y) + ax.spines.right.set_position(('axes', 0.1)) + ax.yaxis.set_ticks_position('right') + ax.spines.top.set_position(('axes', 0.25)) + ax.xaxis.set_ticks_position('top') + ax.spines.left.set_color('none') + ax.spines.bottom.set_color('none') + + +@image_comparison(['spines_data_positions.png']) +def test_spines_data_positions(): + fig, ax = plt.subplots() + ax.spines.left.set_position(('data', -1.5)) + ax.spines.top.set_position(('data', 0.5)) + ax.spines.right.set_position(('data', -0.5)) + ax.spines.bottom.set_position('zero') + ax.set_xlim([-2, 2]) + ax.set_ylim([-2, 2]) + + +@check_figures_equal(extensions=["png"]) +def test_spine_nonlinear_data_positions(fig_test, fig_ref): + plt.style.use("default") + + ax = fig_test.add_subplot() + ax.set(xscale="log", xlim=(.1, 1)) + # Use position="data" to visually swap the left and right spines, using + # linewidth to distinguish them. The calls to tick_params removes labels + # (for image comparison purposes) and harmonizes tick positions with the + # reference). + ax.spines.left.set_position(("data", 1)) + ax.spines.left.set_linewidth(2) + ax.spines.right.set_position(("data", .1)) + ax.tick_params(axis="y", labelleft=False, direction="in") + + ax = fig_ref.add_subplot() + ax.set(xscale="log", xlim=(.1, 1)) + ax.spines.right.set_linewidth(2) + ax.tick_params(axis="y", labelleft=False, left=False, right=True) + + +@image_comparison(['spines_capstyle.png']) +def test_spines_capstyle(): + # issue 2542 + plt.rc('axes', linewidth=20) + fig, ax = plt.subplots() + ax.set_xticks([]) + ax.set_yticks([]) + + +def test_label_without_ticks(): + fig, ax = plt.subplots() + plt.subplots_adjust(left=0.3, bottom=0.3) + ax.plot(np.arange(10)) + ax.yaxis.set_ticks_position('left') + ax.spines.left.set_position(('outward', 30)) + ax.spines.right.set_visible(False) + ax.set_ylabel('y label') + ax.xaxis.set_ticks_position('bottom') + ax.spines.bottom.set_position(('outward', 30)) + ax.spines.top.set_visible(False) + ax.set_xlabel('x label') + ax.xaxis.set_ticks([]) + ax.yaxis.set_ticks([]) + plt.draw() + + spine = ax.spines.left + spinebbox = spine.get_transform().transform_path( + spine.get_path()).get_extents() + assert ax.yaxis.label.get_position()[0] < spinebbox.xmin, \ + "Y-Axis label not left of the spine" + + spine = ax.spines.bottom + spinebbox = spine.get_transform().transform_path( + spine.get_path()).get_extents() + assert ax.xaxis.label.get_position()[1] < spinebbox.ymin, \ + "X-Axis label not below the spine" + + +@image_comparison(['black_axes.png']) +def test_spines_black_axes(): + # GitHub #18804 + plt.rcParams["savefig.pad_inches"] = 0 + plt.rcParams["savefig.bbox"] = 'tight' + fig = plt.figure(0, figsize=(4, 4)) + ax = fig.add_axes((0, 0, 1, 1)) + ax.set_xticklabels([]) + ax.set_yticklabels([]) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_facecolor((0, 0, 0)) + + +def test_arc_spine_inner_no_axis(): + # Backcompat: smoke test that inner arc spine does not need a registered + # axis in order to be drawn + fig = plt.figure() + ax = fig.add_subplot(projection="polar") + inner_spine = ax.spines["inner"] + inner_spine.register_axis(None) + assert ax.spines["inner"].axis is None + + fig.draw_without_rendering() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_streamplot.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_streamplot.py new file mode 100644 index 0000000000000000000000000000000000000000..92e72c5baeca1d8f6212310d9f0b5b98beb39f82 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_streamplot.py @@ -0,0 +1,169 @@ +import numpy as np +from numpy.testing import assert_array_almost_equal +import pytest +import matplotlib.pyplot as plt +from matplotlib.testing.decorators import image_comparison +import matplotlib.transforms as mtransforms + + +def velocity_field(): + Y, X = np.mgrid[-3:3:100j, -3:3:200j] + U = -1 - X**2 + Y + V = 1 + X - Y**2 + return X, Y, U, V + + +def swirl_velocity_field(): + x = np.linspace(-3., 3., 200) + y = np.linspace(-3., 3., 100) + X, Y = np.meshgrid(x, y) + a = 0.1 + U = np.cos(a) * (-Y) - np.sin(a) * X + V = np.sin(a) * (-Y) + np.cos(a) * X + return x, y, U, V + + +@image_comparison(['streamplot_startpoints'], remove_text=True, style='mpl20', + extensions=['png']) +def test_startpoints(): + X, Y, U, V = velocity_field() + start_x, start_y = np.meshgrid(np.linspace(X.min(), X.max(), 5), + np.linspace(Y.min(), Y.max(), 5)) + start_points = np.column_stack([start_x.ravel(), start_y.ravel()]) + plt.streamplot(X, Y, U, V, start_points=start_points) + plt.plot(start_x, start_y, 'ok') + + +@image_comparison(['streamplot_colormap.png'], remove_text=True, style='mpl20', + tol=0.022) +def test_colormap(): + X, Y, U, V = velocity_field() + plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2, + cmap=plt.cm.autumn) + plt.colorbar() + + +@image_comparison(['streamplot_linewidth'], remove_text=True, style='mpl20', + tol=0.004) +def test_linewidth(): + X, Y, U, V = velocity_field() + speed = np.hypot(U, V) + lw = 5 * speed / speed.max() + ax = plt.figure().subplots() + ax.streamplot(X, Y, U, V, density=[0.5, 1], color='k', linewidth=lw) + + +@image_comparison(['streamplot_masks_and_nans.png'], + remove_text=True, style='mpl20') +def test_masks_and_nans(): + X, Y, U, V = velocity_field() + mask = np.zeros(U.shape, dtype=bool) + mask[40:60, 80:120] = 1 + U[:20, :40] = np.nan + U = np.ma.array(U, mask=mask) + ax = plt.figure().subplots() + with np.errstate(invalid='ignore'): + ax.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues) + + +@image_comparison(['streamplot_maxlength.png'], + remove_text=True, style='mpl20', tol=0.302) +def test_maxlength(): + x, y, U, V = swirl_velocity_field() + ax = plt.figure().subplots() + ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]], + linewidth=2, density=2) + assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3 + # Compatibility for old test image + ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413)) + + +@image_comparison(['streamplot_maxlength_no_broken.png'], + remove_text=True, style='mpl20', tol=0.302) +def test_maxlength_no_broken(): + x, y, U, V = swirl_velocity_field() + ax = plt.figure().subplots() + ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]], + linewidth=2, density=2, broken_streamlines=False) + assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3 + # Compatibility for old test image + ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413)) + + +@image_comparison(['streamplot_direction.png'], + remove_text=True, style='mpl20', tol=0.073) +def test_direction(): + x, y, U, V = swirl_velocity_field() + plt.streamplot(x, y, U, V, integration_direction='backward', + maxlength=1.5, start_points=[[1.5, 0.]], + linewidth=2, density=2) + + +def test_streamplot_limits(): + ax = plt.axes() + x = np.linspace(-5, 10, 20) + y = np.linspace(-2, 4, 10) + y, x = np.meshgrid(y, x) + trans = mtransforms.Affine2D().translate(25, 32) + ax.transData + plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans) + # The calculated bounds are approximately the bounds of the original data, + # this is because the entire path is taken into account when updating the + # datalim. + assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6), + decimal=1) + + +def test_streamplot_grid(): + u = np.ones((2, 2)) + v = np.zeros((2, 2)) + + # Test for same rows and columns + x = np.array([[10, 20], [10, 30]]) + y = np.array([[10, 10], [20, 20]]) + + with pytest.raises(ValueError, match="The rows of 'x' must be equal"): + plt.streamplot(x, y, u, v) + + x = np.array([[10, 20], [10, 20]]) + y = np.array([[10, 10], [20, 30]]) + + with pytest.raises(ValueError, match="The columns of 'y' must be equal"): + plt.streamplot(x, y, u, v) + + x = np.array([[10, 20], [10, 20]]) + y = np.array([[10, 10], [20, 20]]) + plt.streamplot(x, y, u, v) + + # Test for maximum dimensions + x = np.array([0, 10]) + y = np.array([[[0, 10]]]) + + with pytest.raises(ValueError, match="'y' can have at maximum " + "2 dimensions"): + plt.streamplot(x, y, u, v) + + # Test for equal spacing + u = np.ones((3, 3)) + v = np.zeros((3, 3)) + x = np.array([0, 10, 20]) + y = np.array([0, 10, 30]) + + with pytest.raises(ValueError, match="'y' values must be equally spaced"): + plt.streamplot(x, y, u, v) + + # Test for strictly increasing + x = np.array([0, 20, 40]) + y = np.array([0, 20, 10]) + + with pytest.raises(ValueError, match="'y' must be strictly increasing"): + plt.streamplot(x, y, u, v) + + +def test_streamplot_inputs(): # test no exception occurs. + # fully-masked + plt.streamplot(np.arange(3), np.arange(3), + np.full((3, 3), np.nan), np.full((3, 3), np.nan), + color=np.random.rand(3, 3)) + # array-likes + plt.streamplot(range(3), range(3), + np.random.rand(3, 3), np.random.rand(3, 3)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_style.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_style.py new file mode 100644 index 0000000000000000000000000000000000000000..be038965e33da3a0646b5a330a7247e51488fb16 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_style.py @@ -0,0 +1,197 @@ +from contextlib import contextmanager +from pathlib import Path +from tempfile import TemporaryDirectory +import sys + +import numpy as np +import pytest + +import matplotlib as mpl +from matplotlib import pyplot as plt, style +from matplotlib.style.core import USER_LIBRARY_PATHS, STYLE_EXTENSION + + +PARAM = 'image.cmap' +VALUE = 'pink' +DUMMY_SETTINGS = {PARAM: VALUE} + + +@contextmanager +def temp_style(style_name, settings=None): + """Context manager to create a style sheet in a temporary directory.""" + if not settings: + settings = DUMMY_SETTINGS + temp_file = f'{style_name}.{STYLE_EXTENSION}' + try: + with TemporaryDirectory() as tmpdir: + # Write style settings to file in the tmpdir. + Path(tmpdir, temp_file).write_text( + "\n".join(f"{k}: {v}" for k, v in settings.items()), + encoding="utf-8") + # Add tmpdir to style path and reload so we can access this style. + USER_LIBRARY_PATHS.append(tmpdir) + style.reload_library() + yield + finally: + style.reload_library() + + +def test_invalid_rc_warning_includes_filename(caplog): + SETTINGS = {'foo': 'bar'} + basename = 'basename' + with temp_style(basename, SETTINGS): + # style.reload_library() in temp_style() triggers the warning + pass + assert (len(caplog.records) == 1 + and basename in caplog.records[0].getMessage()) + + +def test_available(): + with temp_style('_test_', DUMMY_SETTINGS): + assert '_test_' in style.available + + +def test_use(): + mpl.rcParams[PARAM] = 'gray' + with temp_style('test', DUMMY_SETTINGS): + with style.context('test'): + assert mpl.rcParams[PARAM] == VALUE + + +def test_use_url(tmp_path): + path = tmp_path / 'file' + path.write_text('axes.facecolor: adeade', encoding='utf-8') + with temp_style('test', DUMMY_SETTINGS): + url = ('file:' + + ('///' if sys.platform == 'win32' else '') + + path.resolve().as_posix()) + with style.context(url): + assert mpl.rcParams['axes.facecolor'] == "#adeade" + + +def test_single_path(tmp_path): + mpl.rcParams[PARAM] = 'gray' + path = tmp_path / f'text.{STYLE_EXTENSION}' + path.write_text(f'{PARAM} : {VALUE}', encoding='utf-8') + with style.context(path): + assert mpl.rcParams[PARAM] == VALUE + assert mpl.rcParams[PARAM] == 'gray' + + +def test_context(): + mpl.rcParams[PARAM] = 'gray' + with temp_style('test', DUMMY_SETTINGS): + with style.context('test'): + assert mpl.rcParams[PARAM] == VALUE + # Check that this value is reset after the exiting the context. + assert mpl.rcParams[PARAM] == 'gray' + + +def test_context_with_dict(): + original_value = 'gray' + other_value = 'blue' + mpl.rcParams[PARAM] = original_value + with style.context({PARAM: other_value}): + assert mpl.rcParams[PARAM] == other_value + assert mpl.rcParams[PARAM] == original_value + + +def test_context_with_dict_after_namedstyle(): + # Test dict after style name where dict modifies the same parameter. + original_value = 'gray' + other_value = 'blue' + mpl.rcParams[PARAM] = original_value + with temp_style('test', DUMMY_SETTINGS): + with style.context(['test', {PARAM: other_value}]): + assert mpl.rcParams[PARAM] == other_value + assert mpl.rcParams[PARAM] == original_value + + +def test_context_with_dict_before_namedstyle(): + # Test dict before style name where dict modifies the same parameter. + original_value = 'gray' + other_value = 'blue' + mpl.rcParams[PARAM] = original_value + with temp_style('test', DUMMY_SETTINGS): + with style.context([{PARAM: other_value}, 'test']): + assert mpl.rcParams[PARAM] == VALUE + assert mpl.rcParams[PARAM] == original_value + + +def test_context_with_union_of_dict_and_namedstyle(): + # Test dict after style name where dict modifies the a different parameter. + original_value = 'gray' + other_param = 'text.usetex' + other_value = True + d = {other_param: other_value} + mpl.rcParams[PARAM] = original_value + mpl.rcParams[other_param] = (not other_value) + with temp_style('test', DUMMY_SETTINGS): + with style.context(['test', d]): + assert mpl.rcParams[PARAM] == VALUE + assert mpl.rcParams[other_param] == other_value + assert mpl.rcParams[PARAM] == original_value + assert mpl.rcParams[other_param] == (not other_value) + + +def test_context_with_badparam(): + original_value = 'gray' + other_value = 'blue' + with style.context({PARAM: other_value}): + assert mpl.rcParams[PARAM] == other_value + x = style.context({PARAM: original_value, 'badparam': None}) + with pytest.raises(KeyError): + with x: + pass + assert mpl.rcParams[PARAM] == other_value + + +@pytest.mark.parametrize('equiv_styles', + [('mpl20', 'default'), + ('mpl15', 'classic')], + ids=['mpl20', 'mpl15']) +def test_alias(equiv_styles): + rc_dicts = [] + for sty in equiv_styles: + with style.context(sty): + rc_dicts.append(mpl.rcParams.copy()) + + rc_base = rc_dicts[0] + for nm, rc in zip(equiv_styles[1:], rc_dicts[1:]): + assert rc_base == rc + + +def test_xkcd_no_cm(): + assert mpl.rcParams["path.sketch"] is None + plt.xkcd() + assert mpl.rcParams["path.sketch"] == (1, 100, 2) + np.testing.break_cycles() + assert mpl.rcParams["path.sketch"] == (1, 100, 2) + + +def test_xkcd_cm(): + assert mpl.rcParams["path.sketch"] is None + with plt.xkcd(): + assert mpl.rcParams["path.sketch"] == (1, 100, 2) + assert mpl.rcParams["path.sketch"] is None + + +def test_up_to_date_blacklist(): + assert mpl.style.core.STYLE_BLACKLIST <= {*mpl.rcsetup._validators} + + +def test_style_from_module(tmp_path, monkeypatch): + monkeypatch.syspath_prepend(tmp_path) + monkeypatch.chdir(tmp_path) + pkg_path = tmp_path / "mpl_test_style_pkg" + pkg_path.mkdir() + (pkg_path / "test_style.mplstyle").write_text( + "lines.linewidth: 42", encoding="utf-8") + pkg_path.with_suffix(".mplstyle").write_text( + "lines.linewidth: 84", encoding="utf-8") + mpl.style.use("mpl_test_style_pkg.test_style") + assert mpl.rcParams["lines.linewidth"] == 42 + mpl.style.use("mpl_test_style_pkg.mplstyle") + assert mpl.rcParams["lines.linewidth"] == 84 + mpl.style.use("./mpl_test_style_pkg.mplstyle") + assert mpl.rcParams["lines.linewidth"] == 84 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_subplots.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_subplots.py new file mode 100644 index 0000000000000000000000000000000000000000..aa5debd0b7871f80ac90082389e00be4998d50bd --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_subplots.py @@ -0,0 +1,293 @@ +import itertools +import platform + +import numpy as np +import pytest + +import matplotlib as mpl +from matplotlib.axes import Axes, SubplotBase +import matplotlib.pyplot as plt +from matplotlib.testing.decorators import check_figures_equal, image_comparison + + +def check_shared(axs, x_shared, y_shared): + """ + x_shared and y_shared are n x n boolean matrices; entry (i, j) indicates + whether the x (or y) axes of subplots i and j should be shared. + """ + for (i1, ax1), (i2, ax2), (i3, (name, shared)) in itertools.product( + enumerate(axs), + enumerate(axs), + enumerate(zip("xy", [x_shared, y_shared]))): + if i2 <= i1: + continue + assert axs[0]._shared_axes[name].joined(ax1, ax2) == shared[i1, i2], \ + "axes %i and %i incorrectly %ssharing %s axis" % ( + i1, i2, "not " if shared[i1, i2] else "", name) + + +def check_ticklabel_visible(axs, x_visible, y_visible): + """Check that the x and y ticklabel visibility is as specified.""" + for i, (ax, vx, vy) in enumerate(zip(axs, x_visible, y_visible)): + for l in ax.get_xticklabels() + [ax.xaxis.offsetText]: + assert l.get_visible() == vx, \ + f"Visibility of x axis #{i} is incorrectly {vx}" + for l in ax.get_yticklabels() + [ax.yaxis.offsetText]: + assert l.get_visible() == vy, \ + f"Visibility of y axis #{i} is incorrectly {vy}" + # axis label "visibility" is toggled by label_outer by resetting the + # label to empty, but it can also be empty to start with. + if not vx: + assert ax.get_xlabel() == "" + if not vy: + assert ax.get_ylabel() == "" + + +def check_tick1_visible(axs, x_visible, y_visible): + """ + Check that the x and y tick visibility is as specified. + + Note: This only checks the tick1line, i.e. bottom / left ticks. + """ + for ax, visible, in zip(axs, x_visible): + for tick in ax.xaxis.get_major_ticks(): + assert tick.tick1line.get_visible() == visible + for ax, y_visible, in zip(axs, y_visible): + for tick in ax.yaxis.get_major_ticks(): + assert tick.tick1line.get_visible() == visible + + +def test_shared(): + rdim = (4, 4, 2) + share = { + 'all': np.ones(rdim[:2], dtype=bool), + 'none': np.zeros(rdim[:2], dtype=bool), + 'row': np.array([ + [False, True, False, False], + [True, False, False, False], + [False, False, False, True], + [False, False, True, False]]), + 'col': np.array([ + [False, False, True, False], + [False, False, False, True], + [True, False, False, False], + [False, True, False, False]]), + } + visible = { + 'x': { + 'all': [False, False, True, True], + 'col': [False, False, True, True], + 'row': [True] * 4, + 'none': [True] * 4, + False: [True] * 4, + True: [False, False, True, True], + }, + 'y': { + 'all': [True, False, True, False], + 'col': [True] * 4, + 'row': [True, False, True, False], + 'none': [True] * 4, + False: [True] * 4, + True: [True, False, True, False], + }, + } + share[False] = share['none'] + share[True] = share['all'] + + # test default + f, ((a1, a2), (a3, a4)) = plt.subplots(2, 2) + axs = [a1, a2, a3, a4] + check_shared(axs, share['none'], share['none']) + plt.close(f) + + # test all option combinations + ops = [False, True, 'all', 'none', 'row', 'col', 0, 1] + for xo in ops: + for yo in ops: + f, ((a1, a2), (a3, a4)) = plt.subplots(2, 2, sharex=xo, sharey=yo) + axs = [a1, a2, a3, a4] + check_shared(axs, share[xo], share[yo]) + check_ticklabel_visible(axs, visible['x'][xo], visible['y'][yo]) + plt.close(f) + + +@pytest.mark.parametrize('remove_ticks', [True, False]) +@pytest.mark.parametrize('layout_engine', ['none', 'tight', 'constrained']) +@pytest.mark.parametrize('with_colorbar', [True, False]) +def test_label_outer(remove_ticks, layout_engine, with_colorbar): + fig = plt.figure(layout=layout_engine) + axs = fig.subplots(2, 2, sharex=True, sharey=True) + for ax in axs.flat: + ax.set(xlabel="foo", ylabel="bar") + if with_colorbar: + fig.colorbar(mpl.cm.ScalarMappable(), ax=ax) + ax.label_outer(remove_inner_ticks=remove_ticks) + check_ticklabel_visible( + axs.flat, [False, False, True, True], [True, False, True, False]) + if remove_ticks: + check_tick1_visible( + axs.flat, [False, False, True, True], [True, False, True, False]) + else: + check_tick1_visible( + axs.flat, [True, True, True, True], [True, True, True, True]) + + +def test_label_outer_span(): + fig = plt.figure() + gs = fig.add_gridspec(3, 3) + # +---+---+---+ + # | 1 | | + # +---+---+---+ + # | | | 3 | + # + 2 +---+---+ + # | | 4 | | + # +---+---+---+ + a1 = fig.add_subplot(gs[0, 0:2]) + a2 = fig.add_subplot(gs[1:3, 0]) + a3 = fig.add_subplot(gs[1, 2]) + a4 = fig.add_subplot(gs[2, 1]) + for ax in fig.axes: + ax.label_outer() + check_ticklabel_visible( + fig.axes, [False, True, False, True], [True, True, False, False]) + + +def test_label_outer_non_gridspec(): + ax = plt.axes((0, 0, 1, 1)) + ax.label_outer() # Does nothing. + check_ticklabel_visible([ax], [True], [True]) + + +def test_shared_and_moved(): + # test if sharey is on, but then tick_left is called that labels don't + # re-appear. Seaborn does this just to be sure yaxis is on left... + f, (a1, a2) = plt.subplots(1, 2, sharey=True) + check_ticklabel_visible([a2], [True], [False]) + a2.yaxis.tick_left() + check_ticklabel_visible([a2], [True], [False]) + + f, (a1, a2) = plt.subplots(2, 1, sharex=True) + check_ticklabel_visible([a1], [False], [True]) + a2.xaxis.tick_bottom() + check_ticklabel_visible([a1], [False], [True]) + + +def test_exceptions(): + # TODO should this test more options? + with pytest.raises(ValueError): + plt.subplots(2, 2, sharex='blah') + with pytest.raises(ValueError): + plt.subplots(2, 2, sharey='blah') + + +@image_comparison(['subplots_offset_text.png'], + tol=0 if platform.machine() == 'x86_64' else 0.028) +def test_subplots_offsettext(): + x = np.arange(0, 1e10, 1e9) + y = np.arange(0, 100, 10)+1e4 + fig, axs = plt.subplots(2, 2, sharex='col', sharey='all') + axs[0, 0].plot(x, x) + axs[1, 0].plot(x, x) + axs[0, 1].plot(y, x) + axs[1, 1].plot(y, x) + + +@pytest.mark.parametrize("top", [True, False]) +@pytest.mark.parametrize("bottom", [True, False]) +@pytest.mark.parametrize("left", [True, False]) +@pytest.mark.parametrize("right", [True, False]) +def test_subplots_hide_ticklabels(top, bottom, left, right): + # Ideally, we would also test offset-text visibility (and remove + # test_subplots_offsettext), but currently, setting rcParams fails to move + # the offset texts as well. + with plt.rc_context({"xtick.labeltop": top, "xtick.labelbottom": bottom, + "ytick.labelleft": left, "ytick.labelright": right}): + axs = plt.figure().subplots(3, 3, sharex=True, sharey=True) + for (i, j), ax in np.ndenumerate(axs): + xtop = ax.xaxis._major_tick_kw["label2On"] + xbottom = ax.xaxis._major_tick_kw["label1On"] + yleft = ax.yaxis._major_tick_kw["label1On"] + yright = ax.yaxis._major_tick_kw["label2On"] + assert xtop == (top and i == 0) + assert xbottom == (bottom and i == 2) + assert yleft == (left and j == 0) + assert yright == (right and j == 2) + + +@pytest.mark.parametrize("xlabel_position", ["bottom", "top"]) +@pytest.mark.parametrize("ylabel_position", ["left", "right"]) +def test_subplots_hide_axislabels(xlabel_position, ylabel_position): + axs = plt.figure().subplots(3, 3, sharex=True, sharey=True) + for (i, j), ax in np.ndenumerate(axs): + ax.set(xlabel="foo", ylabel="bar") + ax.xaxis.set_label_position(xlabel_position) + ax.yaxis.set_label_position(ylabel_position) + ax.label_outer() + assert bool(ax.get_xlabel()) == ( + xlabel_position == "bottom" and i == 2 + or xlabel_position == "top" and i == 0) + assert bool(ax.get_ylabel()) == ( + ylabel_position == "left" and j == 0 + or ylabel_position == "right" and j == 2) + + +def test_get_gridspec(): + # ahem, pretty trivial, but... + fig, ax = plt.subplots() + assert ax.get_subplotspec().get_gridspec() == ax.get_gridspec() + + +def test_dont_mutate_kwargs(): + subplot_kw = {'sharex': 'all'} + gridspec_kw = {'width_ratios': [1, 2]} + fig, ax = plt.subplots(1, 2, subplot_kw=subplot_kw, + gridspec_kw=gridspec_kw) + assert subplot_kw == {'sharex': 'all'} + assert gridspec_kw == {'width_ratios': [1, 2]} + + +@pytest.mark.parametrize("width_ratios", [None, [1, 3, 2]]) +@pytest.mark.parametrize("height_ratios", [None, [1, 2]]) +@check_figures_equal(extensions=['png']) +def test_width_and_height_ratios(fig_test, fig_ref, + height_ratios, width_ratios): + fig_test.subplots(2, 3, height_ratios=height_ratios, + width_ratios=width_ratios) + fig_ref.subplots(2, 3, gridspec_kw={ + 'height_ratios': height_ratios, + 'width_ratios': width_ratios}) + + +@pytest.mark.parametrize("width_ratios", [None, [1, 3, 2]]) +@pytest.mark.parametrize("height_ratios", [None, [1, 2]]) +@check_figures_equal(extensions=['png']) +def test_width_and_height_ratios_mosaic(fig_test, fig_ref, + height_ratios, width_ratios): + mosaic_spec = [['A', 'B', 'B'], ['A', 'C', 'D']] + fig_test.subplot_mosaic(mosaic_spec, height_ratios=height_ratios, + width_ratios=width_ratios) + fig_ref.subplot_mosaic(mosaic_spec, gridspec_kw={ + 'height_ratios': height_ratios, + 'width_ratios': width_ratios}) + + +@pytest.mark.parametrize('method,args', [ + ('subplots', (2, 3)), + ('subplot_mosaic', ('abc;def', )) + ] +) +def test_ratio_overlapping_kws(method, args): + with pytest.raises(ValueError, match='height_ratios'): + getattr(plt, method)(*args, height_ratios=[1, 2], + gridspec_kw={'height_ratios': [1, 2]}) + with pytest.raises(ValueError, match='width_ratios'): + getattr(plt, method)(*args, width_ratios=[1, 2, 3], + gridspec_kw={'width_ratios': [1, 2, 3]}) + + +def test_old_subplot_compat(): + fig = plt.figure() + assert isinstance(fig.add_subplot(), SubplotBase) + assert not isinstance(fig.add_axes(rect=[0, 0, 1, 1]), SubplotBase) + with pytest.raises(TypeError): + Axes(fig, [0, 0, 1, 1], rect=[0, 0, 1, 1]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_table.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_table.py new file mode 100644 index 0000000000000000000000000000000000000000..29d1944323fc3e6141b0cc2bac229ae4a82bec35 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_table.py @@ -0,0 +1,283 @@ +import datetime +from unittest.mock import Mock + +import numpy as np + +import matplotlib.pyplot as plt +from matplotlib.path import Path +from matplotlib.table import CustomCell, Table +from matplotlib.testing.decorators import image_comparison, check_figures_equal +from matplotlib.transforms import Bbox +import matplotlib.units as munits + + +def test_non_square(): + # Check that creating a non-square table works + cellcolors = ['b', 'r'] + plt.table(cellColours=cellcolors) + + +@image_comparison(['table_zorder.png'], remove_text=True) +def test_zorder(): + data = [[66386, 174296], + [58230, 381139]] + + colLabels = ('Freeze', 'Wind') + rowLabels = ['%d year' % x for x in (100, 50)] + + cellText = [] + yoff = np.zeros(len(colLabels)) + for row in reversed(data): + yoff += row + cellText.append(['%1.1f' % (x/1000.0) for x in yoff]) + + t = np.linspace(0, 2*np.pi, 100) + plt.plot(t, np.cos(t), lw=4, zorder=2) + + plt.table(cellText=cellText, + rowLabels=rowLabels, + colLabels=colLabels, + loc='center', + zorder=-2, + ) + + plt.table(cellText=cellText, + rowLabels=rowLabels, + colLabels=colLabels, + loc='upper center', + zorder=4, + ) + plt.yticks([]) + + +@image_comparison(['table_labels.png']) +def test_label_colours(): + dim = 3 + + c = np.linspace(0, 1, dim) + colours = plt.cm.RdYlGn(c) + cellText = [['1'] * dim] * dim + + fig = plt.figure() + + ax1 = fig.add_subplot(4, 1, 1) + ax1.axis('off') + ax1.table(cellText=cellText, + rowColours=colours, + loc='best') + + ax2 = fig.add_subplot(4, 1, 2) + ax2.axis('off') + ax2.table(cellText=cellText, + rowColours=colours, + rowLabels=['Header'] * dim, + loc='best') + + ax3 = fig.add_subplot(4, 1, 3) + ax3.axis('off') + ax3.table(cellText=cellText, + colColours=colours, + loc='best') + + ax4 = fig.add_subplot(4, 1, 4) + ax4.axis('off') + ax4.table(cellText=cellText, + colColours=colours, + colLabels=['Header'] * dim, + loc='best') + + +@image_comparison(['table_cell_manipulation.png'], style='mpl20') +def test_diff_cell_table(text_placeholders): + cells = ('horizontal', 'vertical', 'open', 'closed', 'T', 'R', 'B', 'L') + cellText = [['1'] * len(cells)] * 2 + colWidths = [0.1] * len(cells) + + _, axs = plt.subplots(nrows=len(cells), figsize=(4, len(cells)+1), layout='tight') + for ax, cell in zip(axs, cells): + ax.table( + colWidths=colWidths, + cellText=cellText, + loc='center', + edges=cell, + ) + ax.axis('off') + + +def test_customcell(): + types = ('horizontal', 'vertical', 'open', 'closed', 'T', 'R', 'B', 'L') + codes = ( + (Path.MOVETO, Path.LINETO, Path.MOVETO, Path.LINETO, Path.MOVETO), + (Path.MOVETO, Path.MOVETO, Path.LINETO, Path.MOVETO, Path.LINETO), + (Path.MOVETO, Path.MOVETO, Path.MOVETO, Path.MOVETO, Path.MOVETO), + (Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY), + (Path.MOVETO, Path.MOVETO, Path.MOVETO, Path.LINETO, Path.MOVETO), + (Path.MOVETO, Path.MOVETO, Path.LINETO, Path.MOVETO, Path.MOVETO), + (Path.MOVETO, Path.LINETO, Path.MOVETO, Path.MOVETO, Path.MOVETO), + (Path.MOVETO, Path.MOVETO, Path.MOVETO, Path.MOVETO, Path.LINETO), + ) + + for t, c in zip(types, codes): + cell = CustomCell((0, 0), visible_edges=t, width=1, height=1) + code = tuple(s for _, s in cell.get_path().iter_segments()) + assert c == code + + +@image_comparison(['table_auto_column.png']) +def test_auto_column(): + fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1) + + # iterable list input + ax1.axis('off') + tb1 = ax1.table( + cellText=[['Fit Text', 2], + ['very long long text, Longer text than default', 1]], + rowLabels=["A", "B"], + colLabels=["Col1", "Col2"], + loc="center") + tb1.auto_set_font_size(False) + tb1.set_fontsize(12) + tb1.auto_set_column_width([-1, 0, 1]) + + # iterable tuple input + ax2.axis('off') + tb2 = ax2.table( + cellText=[['Fit Text', 2], + ['very long long text, Longer text than default', 1]], + rowLabels=["A", "B"], + colLabels=["Col1", "Col2"], + loc="center") + tb2.auto_set_font_size(False) + tb2.set_fontsize(12) + tb2.auto_set_column_width((-1, 0, 1)) + + # 3 single inputs + ax3.axis('off') + tb3 = ax3.table( + cellText=[['Fit Text', 2], + ['very long long text, Longer text than default', 1]], + rowLabels=["A", "B"], + colLabels=["Col1", "Col2"], + loc="center") + tb3.auto_set_font_size(False) + tb3.set_fontsize(12) + tb3.auto_set_column_width(-1) + tb3.auto_set_column_width(0) + tb3.auto_set_column_width(1) + + # 4 this used to test non-integer iterable input, which did nothing, but only + # remains to avoid re-generating the test image. + ax4.axis('off') + tb4 = ax4.table( + cellText=[['Fit Text', 2], + ['very long long text, Longer text than default', 1]], + rowLabels=["A", "B"], + colLabels=["Col1", "Col2"], + loc="center") + tb4.auto_set_font_size(False) + tb4.set_fontsize(12) + + +def test_table_cells(): + fig, ax = plt.subplots() + table = Table(ax) + + cell = table.add_cell(1, 2, 1, 1) + assert isinstance(cell, CustomCell) + assert cell is table[1, 2] + + cell2 = CustomCell((0, 0), 1, 2, visible_edges=None) + table[2, 1] = cell2 + assert table[2, 1] is cell2 + + # make sure getitem support has not broken + # properties and setp + table.properties() + plt.setp(table) + + +@check_figures_equal(extensions=["png"]) +def test_table_bbox(fig_test, fig_ref): + data = [[2, 3], + [4, 5]] + + col_labels = ('Foo', 'Bar') + row_labels = ('Ada', 'Bob') + + cell_text = [[f"{x}" for x in row] for row in data] + + ax_list = fig_test.subplots() + ax_list.table(cellText=cell_text, + rowLabels=row_labels, + colLabels=col_labels, + loc='center', + bbox=[0.1, 0.2, 0.8, 0.6] + ) + + ax_bbox = fig_ref.subplots() + ax_bbox.table(cellText=cell_text, + rowLabels=row_labels, + colLabels=col_labels, + loc='center', + bbox=Bbox.from_extents(0.1, 0.2, 0.9, 0.8) + ) + + +@check_figures_equal(extensions=['png']) +def test_table_unit(fig_test, fig_ref): + # test that table doesn't participate in unit machinery, instead uses repr/str + + class FakeUnit: + def __init__(self, thing): + pass + def __repr__(self): + return "Hello" + + fake_convertor = munits.ConversionInterface() + # v, u, a = value, unit, axis + fake_convertor.convert = Mock(side_effect=lambda v, u, a: 0) + # not used, here for completeness + fake_convertor.default_units = Mock(side_effect=lambda v, a: None) + fake_convertor.axisinfo = Mock(side_effect=lambda u, a: munits.AxisInfo()) + + munits.registry[FakeUnit] = fake_convertor + + data = [[FakeUnit("yellow"), FakeUnit(42)], + [FakeUnit(datetime.datetime(1968, 8, 1)), FakeUnit(True)]] + + fig_test.subplots().table(data) + fig_ref.subplots().table([["Hello", "Hello"], ["Hello", "Hello"]]) + fig_test.canvas.draw() + fake_convertor.convert.assert_not_called() + + munits.registry.pop(FakeUnit) + assert not munits.registry.get_converter(FakeUnit) + + +def test_table_dataframe(pd): + # Test if Pandas Data Frame can be passed in cellText + + data = { + 'Letter': ['A', 'B', 'C'], + 'Number': [100, 200, 300] + } + + df = pd.DataFrame(data) + fig, ax = plt.subplots() + table = ax.table(df, loc='center') + + for r, (index, row) in enumerate(df.iterrows()): + for c, col in enumerate(df.columns if r == 0 else row.values): + assert table[r if r == 0 else r+1, c].get_text().get_text() == str(col) + + +def test_table_fontsize(): + # Test that the passed fontsize propagates to cells + tableData = [['a', 1], ['b', 2]] + fig, ax = plt.subplots() + test_fontsize = 20 + t = ax.table(cellText=tableData, loc='top', fontsize=test_fontsize) + cell_fontsize = t[(0, 0)].get_fontsize() + assert cell_fontsize == test_fontsize, f"Actual:{test_fontsize},got:{cell_fontsize}" + cell_fontsize = t[(1, 1)].get_fontsize() + assert cell_fontsize == test_fontsize, f"Actual:{test_fontsize},got:{cell_fontsize}" diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_testing.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_testing.py new file mode 100644 index 0000000000000000000000000000000000000000..f13839d6b3b68a6ce7e2f7d21e7ba68088e05f95 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_testing.py @@ -0,0 +1,41 @@ +import warnings + +import pytest + +import matplotlib.pyplot as plt +from matplotlib.testing.decorators import check_figures_equal + + +@pytest.mark.xfail( + strict=True, reason="testing that warnings fail tests" +) +def test_warn_to_fail(): + warnings.warn("This should fail the test") + + +@pytest.mark.parametrize("a", [1]) +@check_figures_equal(extensions=["png"]) +@pytest.mark.parametrize("b", [1]) +def test_parametrize_with_check_figure_equal(a, fig_ref, b, fig_test): + assert a == b + + +def test_wrap_failure(): + with pytest.raises(ValueError, match="^The decorated function"): + @check_figures_equal() + def should_fail(test, ref): + pass + + +@pytest.mark.xfail(raises=RuntimeError, strict=True, + reason='Test for check_figures_equal test creating ' + 'new figures') +@check_figures_equal() +def test_check_figures_equal_extra_fig(fig_test, fig_ref): + plt.figure() + + +@check_figures_equal() +def test_check_figures_equal_closed_fig(fig_test, fig_ref): + fig = plt.figure() + plt.close(fig) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_texmanager.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_texmanager.py new file mode 100644 index 0000000000000000000000000000000000000000..64dcbf46456de679c6c29f104a7e86e6afcd9ac2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_texmanager.py @@ -0,0 +1,75 @@ +import os +from pathlib import Path +import re +import sys + +import pytest + +import matplotlib.pyplot as plt +from matplotlib.testing import subprocess_run_for_testing +from matplotlib.testing._markers import needs_usetex +from matplotlib.texmanager import TexManager + + +def test_fontconfig_preamble(): + """Test that the preamble is included in the source.""" + plt.rcParams['text.usetex'] = True + + src1 = TexManager()._get_tex_source("", fontsize=12) + plt.rcParams['text.latex.preamble'] = '\\usepackage{txfonts}' + src2 = TexManager()._get_tex_source("", fontsize=12) + + assert src1 != src2 + + +@pytest.mark.parametrize( + "rc, preamble, family", [ + ({"font.family": "sans-serif", "font.sans-serif": "helvetica"}, + r"\usepackage{helvet}", r"\sffamily"), + ({"font.family": "serif", "font.serif": "palatino"}, + r"\usepackage{mathpazo}", r"\rmfamily"), + ({"font.family": "cursive", "font.cursive": "zapf chancery"}, + r"\usepackage{chancery}", r"\rmfamily"), + ({"font.family": "monospace", "font.monospace": "courier"}, + r"\usepackage{courier}", r"\ttfamily"), + ({"font.family": "helvetica"}, r"\usepackage{helvet}", r"\sffamily"), + ({"font.family": "palatino"}, r"\usepackage{mathpazo}", r"\rmfamily"), + ({"font.family": "zapf chancery"}, + r"\usepackage{chancery}", r"\rmfamily"), + ({"font.family": "courier"}, r"\usepackage{courier}", r"\ttfamily") + ]) +def test_font_selection(rc, preamble, family): + plt.rcParams.update(rc) + tm = TexManager() + src = Path(tm.make_tex("hello, world", fontsize=12)).read_text() + assert preamble in src + assert [*re.findall(r"\\\w+family", src)] == [family] + + +@needs_usetex +def test_unicode_characters(): + # Smoke test to see that Unicode characters does not cause issues + # See #23019 + plt.rcParams['text.usetex'] = True + fig, ax = plt.subplots() + ax.set_ylabel('\\textit{Velocity (\N{DEGREE SIGN}/sec)}') + ax.set_xlabel('\N{VULGAR FRACTION ONE QUARTER}Öøæ') + fig.canvas.draw() + + # But not all characters. + # Should raise RuntimeError, not UnicodeDecodeError + with pytest.raises(RuntimeError): + ax.set_title('\N{SNOWMAN}') + fig.canvas.draw() + + +@needs_usetex +def test_openin_any_paranoid(): + completed = subprocess_run_for_testing( + [sys.executable, "-c", + 'import matplotlib.pyplot as plt;' + 'plt.rcParams.update({"text.usetex": True});' + 'plt.title("paranoid");' + 'plt.show(block=False);'], + env={**os.environ, 'openin_any': 'p'}, check=True, capture_output=True) + assert completed.stderr == "" diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_text.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_text.py new file mode 100644 index 0000000000000000000000000000000000000000..19262202e5c160c32fe1ebda3d4feabb54233616 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_text.py @@ -0,0 +1,1137 @@ +from datetime import datetime +import io +import warnings + +import numpy as np +from numpy.testing import assert_almost_equal +from packaging.version import parse as parse_version +import pyparsing +import pytest + +import matplotlib as mpl +from matplotlib.backend_bases import MouseEvent +from matplotlib.backends.backend_agg import RendererAgg +from matplotlib.figure import Figure +from matplotlib.font_manager import FontProperties +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +from matplotlib.gridspec import GridSpec +import matplotlib.transforms as mtransforms +from matplotlib.testing.decorators import check_figures_equal, image_comparison +from matplotlib.testing._markers import needs_usetex +from matplotlib.text import Text, Annotation, OffsetFrom + +pyparsing_version = parse_version(pyparsing.__version__) + + +@image_comparison(['font_styles']) +def test_font_styles(): + + def find_matplotlib_font(**kw): + prop = FontProperties(**kw) + path = findfont(prop, directory=mpl.get_data_path()) + return FontProperties(fname=path) + + from matplotlib.font_manager import FontProperties, findfont + warnings.filterwarnings( + 'ignore', + r"findfont: Font family \[u?'Foo'\] not found. Falling back to .", + UserWarning, + module='matplotlib.font_manager') + + fig, ax = plt.subplots() + + normal_font = find_matplotlib_font( + family="sans-serif", + style="normal", + variant="normal", + size=14) + a = ax.annotate( + "Normal Font", + (0.1, 0.1), + xycoords='axes fraction', + fontproperties=normal_font) + assert a.get_fontname() == 'DejaVu Sans' + assert a.get_fontstyle() == 'normal' + assert a.get_fontvariant() == 'normal' + assert a.get_weight() == 'normal' + assert a.get_stretch() == 'normal' + + bold_font = find_matplotlib_font( + family="Foo", + style="normal", + variant="normal", + weight="bold", + stretch=500, + size=14) + ax.annotate( + "Bold Font", + (0.1, 0.2), + xycoords='axes fraction', + fontproperties=bold_font) + + bold_italic_font = find_matplotlib_font( + family="sans serif", + style="italic", + variant="normal", + weight=750, + stretch=500, + size=14) + ax.annotate( + "Bold Italic Font", + (0.1, 0.3), + xycoords='axes fraction', + fontproperties=bold_italic_font) + + light_font = find_matplotlib_font( + family="sans-serif", + style="normal", + variant="normal", + weight=200, + stretch=500, + size=14) + ax.annotate( + "Light Font", + (0.1, 0.4), + xycoords='axes fraction', + fontproperties=light_font) + + condensed_font = find_matplotlib_font( + family="sans-serif", + style="normal", + variant="normal", + weight=500, + stretch=100, + size=14) + ax.annotate( + "Condensed Font", + (0.1, 0.5), + xycoords='axes fraction', + fontproperties=condensed_font) + + ax.set_xticks([]) + ax.set_yticks([]) + + +@image_comparison(['multiline']) +def test_multiline(): + plt.figure() + ax = plt.subplot(1, 1, 1) + ax.set_title("multiline\ntext alignment") + + plt.text( + 0.2, 0.5, "TpTpTp\n$M$\nTpTpTp", size=20, ha="center", va="top") + + plt.text( + 0.5, 0.5, "TpTpTp\n$M^{M^{M^{M}}}$\nTpTpTp", size=20, + ha="center", va="top") + + plt.text( + 0.8, 0.5, "TpTpTp\n$M_{q_{q_{q}}}$\nTpTpTp", size=20, + ha="center", va="top") + + plt.xlim(0, 1) + plt.ylim(0, 0.8) + + ax.set_xticks([]) + ax.set_yticks([]) + + +@image_comparison(['multiline2'], style='mpl20') +def test_multiline2(): + # Remove this line when this test image is regenerated. + plt.rcParams['text.kerning_factor'] = 6 + + fig, ax = plt.subplots() + + ax.set_xlim([0, 1.4]) + ax.set_ylim([0, 2]) + ax.axhline(0.5, color='C2', linewidth=0.3) + sts = ['Line', '2 Lineg\n 2 Lg', '$\\sum_i x $', 'hi $\\sum_i x $\ntest', + 'test\n $\\sum_i x $', '$\\sum_i x $\n $\\sum_i x $'] + renderer = fig.canvas.get_renderer() + + def draw_box(ax, tt): + r = mpatches.Rectangle((0, 0), 1, 1, clip_on=False, + transform=ax.transAxes) + r.set_bounds( + tt.get_window_extent(renderer) + .transformed(ax.transAxes.inverted()) + .bounds) + ax.add_patch(r) + + horal = 'left' + for nn, st in enumerate(sts): + tt = ax.text(0.2 * nn + 0.1, 0.5, st, horizontalalignment=horal, + verticalalignment='bottom') + draw_box(ax, tt) + ax.text(1.2, 0.5, 'Bottom align', color='C2') + + ax.axhline(1.3, color='C2', linewidth=0.3) + for nn, st in enumerate(sts): + tt = ax.text(0.2 * nn + 0.1, 1.3, st, horizontalalignment=horal, + verticalalignment='top') + draw_box(ax, tt) + ax.text(1.2, 1.3, 'Top align', color='C2') + + ax.axhline(1.8, color='C2', linewidth=0.3) + for nn, st in enumerate(sts): + tt = ax.text(0.2 * nn + 0.1, 1.8, st, horizontalalignment=horal, + verticalalignment='baseline') + draw_box(ax, tt) + ax.text(1.2, 1.8, 'Baseline align', color='C2') + + ax.axhline(0.1, color='C2', linewidth=0.3) + for nn, st in enumerate(sts): + tt = ax.text(0.2 * nn + 0.1, 0.1, st, horizontalalignment=horal, + verticalalignment='bottom', rotation=20) + draw_box(ax, tt) + ax.text(1.2, 0.1, 'Bot align, rot20', color='C2') + + +@image_comparison(['antialiased.png'], style='mpl20') +def test_antialiasing(): + mpl.rcParams['text.antialiased'] = False # Passed arguments should override. + + fig = plt.figure(figsize=(5.25, 0.75)) + fig.text(0.3, 0.75, "antialiased", horizontalalignment='center', + verticalalignment='center', antialiased=True) + fig.text(0.3, 0.25, r"$\sqrt{x}$", horizontalalignment='center', + verticalalignment='center', antialiased=True) + + mpl.rcParams['text.antialiased'] = True # Passed arguments should override. + fig.text(0.7, 0.75, "not antialiased", horizontalalignment='center', + verticalalignment='center', antialiased=False) + fig.text(0.7, 0.25, r"$\sqrt{x}$", horizontalalignment='center', + verticalalignment='center', antialiased=False) + + mpl.rcParams['text.antialiased'] = False # Should not affect existing text. + + +def test_afm_kerning(): + fn = mpl.font_manager.findfont("Helvetica", fontext="afm") + with open(fn, 'rb') as fh: + afm = mpl._afm.AFM(fh) + assert afm.string_width_height('VAVAVAVAVAVA') == (7174.0, 718) + + +@image_comparison(['text_contains.png']) +def test_contains(): + fig = plt.figure() + ax = plt.axes() + + mevent = MouseEvent('button_press_event', fig.canvas, 0.5, 0.5, 1, None) + + xs = np.linspace(0.25, 0.75, 30) + ys = np.linspace(0.25, 0.75, 30) + xs, ys = np.meshgrid(xs, ys) + + txt = plt.text( + 0.5, 0.4, 'hello world', ha='center', fontsize=30, rotation=30) + # uncomment to draw the text's bounding box + # txt.set_bbox(dict(edgecolor='black', facecolor='none')) + + # draw the text. This is important, as the contains method can only work + # when a renderer exists. + fig.canvas.draw() + + for x, y in zip(xs.flat, ys.flat): + mevent.x, mevent.y = plt.gca().transAxes.transform([x, y]) + contains, _ = txt.contains(mevent) + color = 'yellow' if contains else 'red' + + # capture the viewLim, plot a point, and reset the viewLim + vl = ax.viewLim.frozen() + ax.plot(x, y, 'o', color=color) + ax.viewLim.set(vl) + + +def test_annotation_contains(): + # Check that Annotation.contains looks at the bboxes of the text and the + # arrow separately, not at the joint bbox. + fig, ax = plt.subplots() + ann = ax.annotate( + "hello", xy=(.4, .4), xytext=(.6, .6), arrowprops={"arrowstyle": "->"}) + fig.canvas.draw() # Needed for the same reason as in test_contains. + event = MouseEvent( + "button_press_event", fig.canvas, *ax.transData.transform((.5, .6))) + assert ann.contains(event) == (False, {}) + + +@pytest.mark.parametrize('err, xycoords, match', ( + (TypeError, print, "xycoords callable must return a BboxBase or Transform, not a"), + (TypeError, [0, 0], r"'xycoords' must be an instance of str, tuple"), + (ValueError, "foo", "'foo' is not a valid coordinate"), + (ValueError, "foo bar", "'foo bar' is not a valid coordinate"), + (ValueError, "offset foo", "xycoords cannot be an offset coordinate"), + (ValueError, "axes foo", "'foo' is not a recognized unit"), +)) +def test_annotate_errors(err, xycoords, match): + fig, ax = plt.subplots() + with pytest.raises(err, match=match): + ax.annotate('xy', (0, 0), xytext=(0.5, 0.5), xycoords=xycoords) + fig.canvas.draw() + + +@image_comparison(['titles']) +def test_titles(): + # left and right side titles + plt.figure() + ax = plt.subplot(1, 1, 1) + ax.set_title("left title", loc="left") + ax.set_title("right title", loc="right") + ax.set_xticks([]) + ax.set_yticks([]) + + +@image_comparison(['text_alignment'], style='mpl20') +def test_alignment(): + plt.figure() + ax = plt.subplot(1, 1, 1) + + x = 0.1 + for rotation in (0, 30): + for alignment in ('top', 'bottom', 'baseline', 'center'): + ax.text( + x, 0.5, alignment + " Tj", va=alignment, rotation=rotation, + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5)) + ax.text( + x, 1.0, r'$\sum_{i=0}^{j}$', va=alignment, rotation=rotation) + x += 0.1 + + ax.plot([0, 1], [0.5, 0.5]) + ax.plot([0, 1], [1.0, 1.0]) + + ax.set_xlim(0, 1) + ax.set_ylim(0, 1.5) + ax.set_xticks([]) + ax.set_yticks([]) + + +@image_comparison(['axes_titles.png']) +def test_axes_titles(): + # Related to issue #3327 + plt.figure() + ax = plt.subplot(1, 1, 1) + ax.set_title('center', loc='center', fontsize=20, fontweight=700) + ax.set_title('left', loc='left', fontsize=12, fontweight=400) + ax.set_title('right', loc='right', fontsize=12, fontweight=400) + + +def test_set_position(): + fig, ax = plt.subplots() + + # test set_position + ann = ax.annotate( + 'test', (0, 0), xytext=(0, 0), textcoords='figure pixels') + fig.canvas.draw() + + init_pos = ann.get_window_extent(fig.canvas.renderer) + shift_val = 15 + ann.set_position((shift_val, shift_val)) + fig.canvas.draw() + post_pos = ann.get_window_extent(fig.canvas.renderer) + + for a, b in zip(init_pos.min, post_pos.min): + assert a + shift_val == b + + # test xyann + ann = ax.annotate( + 'test', (0, 0), xytext=(0, 0), textcoords='figure pixels') + fig.canvas.draw() + + init_pos = ann.get_window_extent(fig.canvas.renderer) + shift_val = 15 + ann.xyann = (shift_val, shift_val) + fig.canvas.draw() + post_pos = ann.get_window_extent(fig.canvas.renderer) + + for a, b in zip(init_pos.min, post_pos.min): + assert a + shift_val == b + + +def test_char_index_at(): + fig = plt.figure() + text = fig.text(0.1, 0.9, "") + + text.set_text("i") + bbox = text.get_window_extent() + size_i = bbox.x1 - bbox.x0 + + text.set_text("m") + bbox = text.get_window_extent() + size_m = bbox.x1 - bbox.x0 + + text.set_text("iiiimmmm") + bbox = text.get_window_extent() + origin = bbox.x0 + + assert text._char_index_at(origin - size_i) == 0 # left of first char + assert text._char_index_at(origin) == 0 + assert text._char_index_at(origin + 0.499*size_i) == 0 + assert text._char_index_at(origin + 0.501*size_i) == 1 + assert text._char_index_at(origin + size_i*3) == 3 + assert text._char_index_at(origin + size_i*4 + size_m*3) == 7 + assert text._char_index_at(origin + size_i*4 + size_m*4) == 8 + assert text._char_index_at(origin + size_i*4 + size_m*10) == 8 + + +@pytest.mark.parametrize('text', ['', 'O'], ids=['empty', 'non-empty']) +def test_non_default_dpi(text): + fig, ax = plt.subplots() + + t1 = ax.text(0.5, 0.5, text, ha='left', va='bottom') + fig.canvas.draw() + dpi = fig.dpi + + bbox1 = t1.get_window_extent() + bbox2 = t1.get_window_extent(dpi=dpi * 10) + np.testing.assert_allclose(bbox2.get_points(), bbox1.get_points() * 10, + rtol=5e-2) + # Text.get_window_extent should not permanently change dpi. + assert fig.dpi == dpi + + +def test_get_rotation_string(): + assert Text(rotation='horizontal').get_rotation() == 0. + assert Text(rotation='vertical').get_rotation() == 90. + + +def test_get_rotation_float(): + for i in [15., 16.70, 77.4]: + assert Text(rotation=i).get_rotation() == i + + +def test_get_rotation_int(): + for i in [67, 16, 41]: + assert Text(rotation=i).get_rotation() == float(i) + + +def test_get_rotation_raises(): + with pytest.raises(ValueError): + Text(rotation='hozirontal') + + +def test_get_rotation_none(): + assert Text(rotation=None).get_rotation() == 0.0 + + +def test_get_rotation_mod360(): + for i, j in zip([360., 377., 720+177.2], [0., 17., 177.2]): + assert_almost_equal(Text(rotation=i).get_rotation(), j) + + +@pytest.mark.parametrize("ha", ["center", "right", "left"]) +@pytest.mark.parametrize("va", ["center", "top", "bottom", + "baseline", "center_baseline"]) +def test_null_rotation_with_rotation_mode(ha, va): + fig, ax = plt.subplots() + kw = dict(rotation=0, va=va, ha=ha) + t0 = ax.text(.5, .5, 'test', rotation_mode='anchor', **kw) + t1 = ax.text(.5, .5, 'test', rotation_mode='default', **kw) + fig.canvas.draw() + assert_almost_equal(t0.get_window_extent(fig.canvas.renderer).get_points(), + t1.get_window_extent(fig.canvas.renderer).get_points()) + + +@image_comparison(['text_bboxclip']) +def test_bbox_clipping(): + plt.text(0.9, 0.2, 'Is bbox clipped?', backgroundcolor='r', clip_on=True) + t = plt.text(0.9, 0.5, 'Is fancy bbox clipped?', clip_on=True) + t.set_bbox({"boxstyle": "round, pad=0.1"}) + + +@image_comparison(['annotation_negative_ax_coords.png']) +def test_annotation_negative_ax_coords(): + fig, ax = plt.subplots() + + ax.annotate('+ pts', + xytext=[30, 20], textcoords='axes points', + xy=[30, 20], xycoords='axes points', fontsize=32) + ax.annotate('- pts', + xytext=[30, -20], textcoords='axes points', + xy=[30, -20], xycoords='axes points', fontsize=32, + va='top') + ax.annotate('+ frac', + xytext=[0.75, 0.05], textcoords='axes fraction', + xy=[0.75, 0.05], xycoords='axes fraction', fontsize=32) + ax.annotate('- frac', + xytext=[0.75, -0.05], textcoords='axes fraction', + xy=[0.75, -0.05], xycoords='axes fraction', fontsize=32, + va='top') + + ax.annotate('+ pixels', + xytext=[160, 25], textcoords='axes pixels', + xy=[160, 25], xycoords='axes pixels', fontsize=32) + ax.annotate('- pixels', + xytext=[160, -25], textcoords='axes pixels', + xy=[160, -25], xycoords='axes pixels', fontsize=32, + va='top') + + +@image_comparison(['annotation_negative_fig_coords.png']) +def test_annotation_negative_fig_coords(): + fig, ax = plt.subplots() + + ax.annotate('+ pts', + xytext=[10, 120], textcoords='figure points', + xy=[10, 120], xycoords='figure points', fontsize=32) + ax.annotate('- pts', + xytext=[-10, 180], textcoords='figure points', + xy=[-10, 180], xycoords='figure points', fontsize=32, + va='top') + ax.annotate('+ frac', + xytext=[0.05, 0.55], textcoords='figure fraction', + xy=[0.05, 0.55], xycoords='figure fraction', fontsize=32) + ax.annotate('- frac', + xytext=[-0.05, 0.5], textcoords='figure fraction', + xy=[-0.05, 0.5], xycoords='figure fraction', fontsize=32, + va='top') + + ax.annotate('+ pixels', + xytext=[50, 50], textcoords='figure pixels', + xy=[50, 50], xycoords='figure pixels', fontsize=32) + ax.annotate('- pixels', + xytext=[-50, 100], textcoords='figure pixels', + xy=[-50, 100], xycoords='figure pixels', fontsize=32, + va='top') + + +def test_text_stale(): + fig, (ax1, ax2) = plt.subplots(1, 2) + plt.draw_all() + assert not ax1.stale + assert not ax2.stale + assert not fig.stale + + txt1 = ax1.text(.5, .5, 'aardvark') + assert ax1.stale + assert txt1.stale + assert fig.stale + + ann1 = ax2.annotate('aardvark', xy=[.5, .5]) + assert ax2.stale + assert ann1.stale + assert fig.stale + + plt.draw_all() + assert not ax1.stale + assert not ax2.stale + assert not fig.stale + + +@image_comparison(['agg_text_clip.png']) +def test_agg_text_clip(): + np.random.seed(1) + fig, (ax1, ax2) = plt.subplots(2) + for x, y in np.random.rand(10, 2): + ax1.text(x, y, "foo", clip_on=True) + ax2.text(x, y, "foo") + + +def test_text_size_binding(): + mpl.rcParams['font.size'] = 10 + fp = mpl.font_manager.FontProperties(size='large') + sz1 = fp.get_size_in_points() + mpl.rcParams['font.size'] = 100 + + assert sz1 == fp.get_size_in_points() + + +@image_comparison(['font_scaling.pdf']) +def test_font_scaling(): + mpl.rcParams['pdf.fonttype'] = 42 + fig, ax = plt.subplots(figsize=(6.4, 12.4)) + ax.xaxis.set_major_locator(plt.NullLocator()) + ax.yaxis.set_major_locator(plt.NullLocator()) + ax.set_ylim(-10, 600) + + for i, fs in enumerate(range(4, 43, 2)): + ax.text(0.1, i*30, f"{fs} pt font size", fontsize=fs) + + +@pytest.mark.parametrize('spacing1, spacing2', [(0.4, 2), (2, 0.4), (2, 2)]) +def test_two_2line_texts(spacing1, spacing2): + text_string = 'line1\nline2' + fig = plt.figure() + renderer = fig.canvas.get_renderer() + + text1 = fig.text(0.25, 0.5, text_string, linespacing=spacing1) + text2 = fig.text(0.25, 0.5, text_string, linespacing=spacing2) + fig.canvas.draw() + + box1 = text1.get_window_extent(renderer=renderer) + box2 = text2.get_window_extent(renderer=renderer) + + # line spacing only affects height + assert box1.width == box2.width + if spacing1 == spacing2: + assert box1.height == box2.height + else: + assert box1.height != box2.height + + +def test_validate_linespacing(): + with pytest.raises(TypeError): + plt.text(.25, .5, "foo", linespacing="abc") + + +def test_nonfinite_pos(): + fig, ax = plt.subplots() + ax.text(0, np.nan, 'nan') + ax.text(np.inf, 0, 'inf') + fig.canvas.draw() + + +def test_hinting_factor_backends(): + plt.rcParams['text.hinting_factor'] = 1 + fig = plt.figure() + t = fig.text(0.5, 0.5, 'some text') + + fig.savefig(io.BytesIO(), format='svg') + expected = t.get_window_extent().intervalx + + fig.savefig(io.BytesIO(), format='png') + # Backends should apply hinting_factor consistently (within 10%). + np.testing.assert_allclose(t.get_window_extent().intervalx, expected, + rtol=0.1) + + +@needs_usetex +def test_usetex_is_copied(): + # Indirectly tests that update_from (which is used to copy tick label + # properties) copies usetex state. + fig = plt.figure() + plt.rcParams["text.usetex"] = False + ax1 = fig.add_subplot(121) + plt.rcParams["text.usetex"] = True + ax2 = fig.add_subplot(122) + fig.canvas.draw() + for ax, usetex in [(ax1, False), (ax2, True)]: + for t in ax.xaxis.majorTicks: + assert t.label1.get_usetex() == usetex + + +@needs_usetex +def test_single_artist_usetex(): + # Check that a single artist marked with usetex does not get passed through + # the mathtext parser at all (for the Agg backend) (the mathtext parser + # currently fails to parse \frac12, requiring \frac{1}{2} instead). + fig = plt.figure() + fig.text(.5, .5, r"$\frac12$", usetex=True) + fig.canvas.draw() + + +@pytest.mark.parametrize("fmt", ["png", "pdf", "svg"]) +def test_single_artist_usenotex(fmt): + # Check that a single artist can be marked as not-usetex even though the + # rcParam is on ("2_2_2" fails if passed to TeX). This currently skips + # postscript output as the ps renderer doesn't support mixing usetex and + # non-usetex. + plt.rcParams["text.usetex"] = True + fig = plt.figure() + fig.text(.5, .5, "2_2_2", usetex=False) + fig.savefig(io.BytesIO(), format=fmt) + + +@image_comparison(['text_as_path_opacity.svg']) +def test_text_as_path_opacity(): + plt.figure() + plt.gca().set_axis_off() + plt.text(0.25, 0.25, 'c', color=(0, 0, 0, 0.5)) + plt.text(0.25, 0.5, 'a', alpha=0.5) + plt.text(0.25, 0.75, 'x', alpha=0.5, color=(0, 0, 0, 1)) + + +@image_comparison(['text_as_text_opacity.svg']) +def test_text_as_text_opacity(): + mpl.rcParams['svg.fonttype'] = 'none' + plt.figure() + plt.gca().set_axis_off() + plt.text(0.25, 0.25, '50% using `color`', color=(0, 0, 0, 0.5)) + plt.text(0.25, 0.5, '50% using `alpha`', alpha=0.5) + plt.text(0.25, 0.75, '50% using `alpha` and 100% `color`', alpha=0.5, + color=(0, 0, 0, 1)) + + +def test_text_repr(): + # smoketest to make sure text repr doesn't error for category + plt.plot(['A', 'B'], [1, 2]) + repr(plt.text(['A'], 0.5, 'Boo')) + + +def test_annotation_update(): + fig, ax = plt.subplots(1, 1) + an = ax.annotate('annotation', xy=(0.5, 0.5)) + extent1 = an.get_window_extent(fig.canvas.get_renderer()) + fig.tight_layout() + extent2 = an.get_window_extent(fig.canvas.get_renderer()) + + assert not np.allclose(extent1.get_points(), extent2.get_points(), + rtol=1e-6) + + +@check_figures_equal(extensions=["png"]) +def test_annotation_units(fig_test, fig_ref): + ax = fig_test.add_subplot() + ax.plot(datetime.now(), 1, "o") # Implicitly set axes extents. + ax.annotate("x", (datetime.now(), 0.5), xycoords=("data", "axes fraction"), + # This used to crash before. + xytext=(0, 0), textcoords="offset points") + ax = fig_ref.add_subplot() + ax.plot(datetime.now(), 1, "o") + ax.annotate("x", (datetime.now(), 0.5), xycoords=("data", "axes fraction")) + + +@image_comparison(['large_subscript_title.png'], style='mpl20') +def test_large_subscript_title(): + # Remove this line when this test image is regenerated. + plt.rcParams['text.kerning_factor'] = 6 + plt.rcParams['axes.titley'] = None + + fig, axs = plt.subplots(1, 2, figsize=(9, 2.5), constrained_layout=True) + ax = axs[0] + ax.set_title(r'$\sum_{i} x_i$') + ax.set_title('New way', loc='left') + ax.set_xticklabels([]) + + ax = axs[1] + ax.set_title(r'$\sum_{i} x_i$', y=1.01) + ax.set_title('Old Way', loc='left') + ax.set_xticklabels([]) + + +@pytest.mark.parametrize( + "x, rotation, halign", + [(0.7, 0, 'left'), + (0.5, 95, 'left'), + (0.3, 0, 'right'), + (0.3, 185, 'left')]) +def test_wrap(x, rotation, halign): + fig = plt.figure(figsize=(18, 18)) + gs = GridSpec(nrows=3, ncols=3, figure=fig) + subfig = fig.add_subfigure(gs[1, 1]) + # we only use the central subfigure, which does not align with any + # figure boundary, to ensure only subfigure boundaries are relevant + s = 'This is a very long text that should be wrapped multiple times.' + text = subfig.text(x, 0.7, s, wrap=True, rotation=rotation, ha=halign) + fig.canvas.draw() + assert text._get_wrapped_text() == ('This is a very long\n' + 'text that should be\n' + 'wrapped multiple\n' + 'times.') + + +def test_mathwrap(): + fig = plt.figure(figsize=(6, 4)) + s = r'This is a very $\overline{\mathrm{long}}$ line of Mathtext.' + text = fig.text(0, 0.5, s, size=40, wrap=True) + fig.canvas.draw() + assert text._get_wrapped_text() == ('This is a very $\\overline{\\mathrm{long}}$\n' + 'line of Mathtext.') + + +def test_get_window_extent_wrapped(): + # Test that a long title that wraps to two lines has the same vertical + # extent as an explicit two line title. + + fig1 = plt.figure(figsize=(3, 3)) + fig1.suptitle("suptitle that is clearly too long in this case", wrap=True) + window_extent_test = fig1._suptitle.get_window_extent() + + fig2 = plt.figure(figsize=(3, 3)) + fig2.suptitle("suptitle that is clearly\ntoo long in this case") + window_extent_ref = fig2._suptitle.get_window_extent() + + assert window_extent_test.y0 == window_extent_ref.y0 + assert window_extent_test.y1 == window_extent_ref.y1 + + +def test_long_word_wrap(): + fig = plt.figure(figsize=(6, 4)) + text = fig.text(9.5, 8, 'Alonglineoftexttowrap', wrap=True) + fig.canvas.draw() + assert text._get_wrapped_text() == 'Alonglineoftexttowrap' + + +def test_wrap_no_wrap(): + fig = plt.figure(figsize=(6, 4)) + text = fig.text(0, 0, 'non wrapped text', wrap=True) + fig.canvas.draw() + assert text._get_wrapped_text() == 'non wrapped text' + + +@check_figures_equal(extensions=["png"]) +def test_buffer_size(fig_test, fig_ref): + # On old versions of the Agg renderer, large non-ascii single-character + # strings (here, "€") would be rendered clipped because the rendering + # buffer would be set by the physical size of the smaller "a" character. + ax = fig_test.add_subplot() + ax.set_yticks([0, 1]) + ax.set_yticklabels(["€", "a"]) + ax.yaxis.majorTicks[1].label1.set_color("w") + ax = fig_ref.add_subplot() + ax.set_yticks([0, 1]) + ax.set_yticklabels(["€", ""]) + + +def test_fontproperties_kwarg_precedence(): + """Test that kwargs take precedence over fontproperties defaults.""" + plt.figure() + text1 = plt.xlabel("value", fontproperties='Times New Roman', size=40.0) + text2 = plt.ylabel("counts", size=40.0, fontproperties='Times New Roman') + assert text1.get_size() == 40.0 + assert text2.get_size() == 40.0 + + +def test_transform_rotates_text(): + ax = plt.gca() + transform = mtransforms.Affine2D().rotate_deg(30) + text = ax.text(0, 0, 'test', transform=transform, + transform_rotates_text=True) + result = text.get_rotation() + assert_almost_equal(result, 30) + + +def test_update_mutate_input(): + inp = dict(fontproperties=FontProperties(weight="bold"), + bbox=None) + cache = dict(inp) + t = Text() + t.update(inp) + assert inp['fontproperties'] == cache['fontproperties'] + assert inp['bbox'] == cache['bbox'] + + +@pytest.mark.parametrize('rotation', ['invalid string', [90]]) +def test_invalid_rotation_values(rotation): + with pytest.raises( + ValueError, + match=("rotation must be 'vertical', 'horizontal' or a number")): + Text(0, 0, 'foo', rotation=rotation) + + +def test_invalid_color(): + with pytest.raises(ValueError): + plt.figtext(.5, .5, "foo", c="foobar") + + +@image_comparison(['text_pdf_kerning.pdf'], style='mpl20') +def test_pdf_kerning(): + plt.figure() + plt.figtext(0.1, 0.5, "ATATATATATATATATATA", size=30) + + +def test_unsupported_script(recwarn): + fig = plt.figure() + t = fig.text(.5, .5, "\N{BENGALI DIGIT ZERO}") + fig.canvas.draw() + assert all(isinstance(warn.message, UserWarning) for warn in recwarn) + assert ( + [warn.message.args for warn in recwarn] == + [(r"Glyph 2534 (\N{BENGALI DIGIT ZERO}) missing from font(s) " + + f"{t.get_fontname()}.",), + (r"Matplotlib currently does not support Bengali natively.",)]) + + +# See gh-26152 for more information on this xfail +@pytest.mark.xfail(pyparsing_version.release == (3, 1, 0), + reason="Error messages are incorrect with pyparsing 3.1.0") +def test_parse_math(): + fig, ax = plt.subplots() + ax.text(0, 0, r"$ \wrong{math} $", parse_math=False) + fig.canvas.draw() + + ax.text(0, 0, r"$ \wrong{math} $", parse_math=True) + with pytest.raises(ValueError, match='Unknown symbol'): + fig.canvas.draw() + + +# See gh-26152 for more information on this xfail +@pytest.mark.xfail(pyparsing_version.release == (3, 1, 0), + reason="Error messages are incorrect with pyparsing 3.1.0") +def test_parse_math_rcparams(): + # Default is True + fig, ax = plt.subplots() + ax.text(0, 0, r"$ \wrong{math} $") + with pytest.raises(ValueError, match='Unknown symbol'): + fig.canvas.draw() + + # Setting rcParams to False + with mpl.rc_context({'text.parse_math': False}): + fig, ax = plt.subplots() + ax.text(0, 0, r"$ \wrong{math} $") + fig.canvas.draw() + + +@image_comparison(['text_pdf_font42_kerning.pdf'], style='mpl20') +def test_pdf_font42_kerning(): + plt.rcParams['pdf.fonttype'] = 42 + plt.figure() + plt.figtext(0.1, 0.5, "ATAVATAVATAVATAVATA", size=30) + + +@image_comparison(['text_pdf_chars_beyond_bmp.pdf'], style='mpl20') +def test_pdf_chars_beyond_bmp(): + plt.rcParams['pdf.fonttype'] = 42 + plt.rcParams['mathtext.fontset'] = 'stixsans' + plt.figure() + plt.figtext(0.1, 0.5, "Mass $m$ \U00010308", size=30) + + +@needs_usetex +def test_metrics_cache(): + mpl.text._get_text_metrics_with_cache_impl.cache_clear() + + fig = plt.figure() + fig.text(.3, .5, "foo\nbar") + fig.text(.3, .5, "foo\nbar", usetex=True) + fig.text(.5, .5, "foo\nbar", usetex=True) + fig.canvas.draw() + renderer = fig._get_renderer() + ys = {} # mapping of strings to where they were drawn in y with draw_tex. + + def call(*args, **kwargs): + renderer, x, y, s, *_ = args + ys.setdefault(s, set()).add(y) + + renderer.draw_tex = call + fig.canvas.draw() + assert [*ys] == ["foo", "bar"] + # Check that both TeX strings were drawn with the same y-position for both + # single-line substrings. Previously, there used to be an incorrect cache + # collision with the non-TeX string (drawn first here) whose metrics would + # get incorrectly reused by the first TeX string. + assert len(ys["foo"]) == len(ys["bar"]) == 1 + + info = mpl.text._get_text_metrics_with_cache_impl.cache_info() + # Every string gets a miss for the first layouting (extents), then a hit + # when drawing, but "foo\nbar" gets two hits as it's drawn twice. + assert info.hits > info.misses + + +def test_annotate_offset_fontsize(): + # Test that offset_fontsize parameter works and uses accurate values + fig, ax = plt.subplots() + text_coords = ['offset points', 'offset fontsize'] + # 10 points should be equal to 1 fontsize unit at fontsize=10 + xy_text = [(10, 10), (1, 1)] + anns = [ax.annotate('test', xy=(0.5, 0.5), + xytext=xy_text[i], + fontsize='10', + xycoords='data', + textcoords=text_coords[i]) for i in range(2)] + points_coords, fontsize_coords = (ann.get_window_extent() for ann in anns) + fig.canvas.draw() + assert str(points_coords) == str(fontsize_coords) + + +def test_get_set_antialiased(): + txt = Text(.5, .5, "foo\nbar") + assert txt._antialiased == mpl.rcParams['text.antialiased'] + assert txt.get_antialiased() == mpl.rcParams['text.antialiased'] + + txt.set_antialiased(True) + assert txt._antialiased is True + assert txt.get_antialiased() == txt._antialiased + + txt.set_antialiased(False) + assert txt._antialiased is False + assert txt.get_antialiased() == txt._antialiased + + +def test_annotation_antialiased(): + annot = Annotation("foo\nbar", (.5, .5), antialiased=True) + assert annot._antialiased is True + assert annot.get_antialiased() == annot._antialiased + + annot2 = Annotation("foo\nbar", (.5, .5), antialiased=False) + assert annot2._antialiased is False + assert annot2.get_antialiased() == annot2._antialiased + + annot3 = Annotation("foo\nbar", (.5, .5), antialiased=False) + annot3.set_antialiased(True) + assert annot3.get_antialiased() is True + assert annot3._antialiased is True + + annot4 = Annotation("foo\nbar", (.5, .5)) + assert annot4._antialiased == mpl.rcParams['text.antialiased'] + + +@check_figures_equal(extensions=["png"]) +def test_annotate_and_offsetfrom_copy_input(fig_test, fig_ref): + # Both approaches place the text (10, 0) pixels away from the center of the line. + ax = fig_test.add_subplot() + l, = ax.plot([0, 2], [0, 2]) + of_xy = np.array([.5, .5]) + ax.annotate("foo", textcoords=OffsetFrom(l, of_xy), xytext=(10, 0), + xy=(0, 0)) # xy is unused. + of_xy[:] = 1 + ax = fig_ref.add_subplot() + l, = ax.plot([0, 2], [0, 2]) + an_xy = np.array([.5, .5]) + ax.annotate("foo", xy=an_xy, xycoords=l, xytext=(10, 0), textcoords="offset points") + an_xy[:] = 2 + + +@check_figures_equal() +def test_text_antialiased_off_default_vs_manual(fig_test, fig_ref): + fig_test.text(0.5, 0.5, '6 inches x 2 inches', + antialiased=False) + + mpl.rcParams['text.antialiased'] = False + fig_ref.text(0.5, 0.5, '6 inches x 2 inches') + + +@check_figures_equal() +def test_text_antialiased_on_default_vs_manual(fig_test, fig_ref): + fig_test.text(0.5, 0.5, '6 inches x 2 inches', antialiased=True) + + mpl.rcParams['text.antialiased'] = True + fig_ref.text(0.5, 0.5, '6 inches x 2 inches') + + +def test_text_annotation_get_window_extent(): + figure = Figure(dpi=100) + renderer = RendererAgg(200, 200, 100) + + # Only text annotation + annotation = Annotation('test', xy=(0, 0), xycoords='figure pixels') + annotation.set_figure(figure) + + text = Text(text='test', x=0, y=0) + text.set_figure(figure) + + bbox = annotation.get_window_extent(renderer=renderer) + + text_bbox = text.get_window_extent(renderer=renderer) + assert bbox.width == text_bbox.width + assert bbox.height == text_bbox.height + + _, _, d = renderer.get_text_width_height_descent( + 'text', annotation._fontproperties, ismath=False) + _, _, lp_d = renderer.get_text_width_height_descent( + 'lp', annotation._fontproperties, ismath=False) + below_line = max(d, lp_d) + + # These numbers are specific to the current implementation of Text + points = bbox.get_points() + assert points[0, 0] == 0.0 + assert points[1, 0] == text_bbox.width + assert points[0, 1] == -below_line + assert points[1, 1] == text_bbox.height - below_line + + +def test_text_with_arrow_annotation_get_window_extent(): + headwidth = 21 + fig, ax = plt.subplots(dpi=100) + txt = ax.text(s='test', x=0, y=0) + ann = ax.annotate( + 'test', + xy=(0.0, 50.0), + xytext=(50.0, 50.0), xycoords='figure pixels', + arrowprops={ + 'facecolor': 'black', 'width': 2, + 'headwidth': headwidth, 'shrink': 0.0}) + + plt.draw() + renderer = fig.canvas.renderer + # bounding box of text + text_bbox = txt.get_window_extent(renderer=renderer) + # bounding box of annotation (text + arrow) + bbox = ann.get_window_extent(renderer=renderer) + # bounding box of arrow + arrow_bbox = ann.arrow_patch.get_window_extent(renderer) + # bounding box of annotation text + ann_txt_bbox = Text.get_window_extent(ann) + + # make sure annotation width is 50 px wider than + # just the text + assert bbox.width == text_bbox.width + 50.0 + # make sure the annotation text bounding box is same size + # as the bounding box of the same string as a Text object + assert ann_txt_bbox.height == text_bbox.height + assert ann_txt_bbox.width == text_bbox.width + # compute the expected bounding box of arrow + text + expected_bbox = mtransforms.Bbox.union([ann_txt_bbox, arrow_bbox]) + assert_almost_equal(bbox.height, expected_bbox.height) + + +def test_arrow_annotation_get_window_extent(): + dpi = 100 + dots_per_point = dpi / 72 + figure = Figure(dpi=dpi) + figure.set_figwidth(2.0) + figure.set_figheight(2.0) + renderer = RendererAgg(200, 200, 100) + + # Text annotation with arrow; arrow dimensions are in points + annotation = Annotation( + '', xy=(0.0, 50.0), xytext=(50.0, 50.0), xycoords='figure pixels', + arrowprops={ + 'facecolor': 'black', 'width': 8, 'headwidth': 10, 'shrink': 0.0}) + annotation.set_figure(figure) + annotation.draw(renderer) + + bbox = annotation.get_window_extent() + points = bbox.get_points() + + assert bbox.width == 50.0 + assert_almost_equal(bbox.height, 10.0 * dots_per_point) + assert points[0, 0] == 0.0 + assert points[0, 1] == 50.0 - 5 * dots_per_point + + +def test_empty_annotation_get_window_extent(): + figure = Figure(dpi=100) + figure.set_figwidth(2.0) + figure.set_figheight(2.0) + renderer = RendererAgg(200, 200, 100) + + # Text annotation with arrow + annotation = Annotation( + '', xy=(0.0, 50.0), xytext=(0.0, 50.0), xycoords='figure pixels') + annotation.set_figure(figure) + annotation.draw(renderer) + + bbox = annotation.get_window_extent() + points = bbox.get_points() + + assert points[0, 0] == 0.0 + assert points[1, 0] == 0.0 + assert points[1, 1] == 50.0 + assert points[0, 1] == 50.0 + + +@image_comparison(baseline_images=['basictext_wrap'], + extensions=['png']) +def test_basic_wrap(): + fig = plt.figure() + plt.axis([0, 10, 0, 10]) + t = "This is a really long string that I'd rather have wrapped so that" \ + " it doesn't go outside of the figure, but if it's long enough it" \ + " will go off the top or bottom!" + plt.text(4, 1, t, ha='left', rotation=15, wrap=True) + plt.text(6, 5, t, ha='left', rotation=15, wrap=True) + plt.text(5, 5, t, ha='right', rotation=-15, wrap=True) + plt.text(5, 10, t, fontsize=18, style='oblique', ha='center', + va='top', wrap=True) + plt.text(3, 4, t, family='serif', style='italic', ha='right', wrap=True) + plt.text(-1, 0, t, ha='left', rotation=-15, wrap=True) + + +@image_comparison(baseline_images=['fonttext_wrap'], + extensions=['png']) +def test_font_wrap(): + fig = plt.figure() + plt.axis([0, 10, 0, 10]) + t = "This is a really long string that I'd rather have wrapped so that" \ + " it doesn't go outside of the figure, but if it's long enough it" \ + " will go off the top or bottom!" + plt.text(4, -1, t, fontsize=18, family='serif', ha='left', rotation=15, + wrap=True) + plt.text(6, 5, t, family='sans serif', ha='left', rotation=15, wrap=True) + plt.text(5, 10, t, weight='heavy', ha='center', va='top', wrap=True) + plt.text(3, 4, t, family='monospace', ha='right', wrap=True) + plt.text(-1, 0, t, fontsize=14, style='italic', ha='left', rotation=-15, + wrap=True) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_textpath.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_textpath.py new file mode 100644 index 0000000000000000000000000000000000000000..e421d2623cadac7931f1f21148bdc0a603c290b1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_textpath.py @@ -0,0 +1,10 @@ +import copy + +from matplotlib.textpath import TextPath + + +def test_copy(): + tp = TextPath((0, 0), ".") + assert copy.deepcopy(tp).vertices is not tp.vertices + assert (copy.deepcopy(tp).vertices == tp.vertices).all() + assert copy.copy(tp).vertices is tp.vertices diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_ticker.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_ticker.py new file mode 100644 index 0000000000000000000000000000000000000000..77c0e917df8a1fd2c7deb450e3dd8b917391f3b8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_ticker.py @@ -0,0 +1,1939 @@ +from contextlib import nullcontext +import itertools +import locale +import logging +import re +from packaging.version import parse as parse_version + +import numpy as np +from numpy.testing import assert_almost_equal, assert_array_equal +import pytest + +import matplotlib as mpl +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker + + +class TestMaxNLocator: + basic_data = [ + (20, 100, np.array([20., 40., 60., 80., 100.])), + (0.001, 0.0001, np.array([0., 0.0002, 0.0004, 0.0006, 0.0008, 0.001])), + (-1e15, 1e15, np.array([-1.0e+15, -5.0e+14, 0e+00, 5e+14, 1.0e+15])), + (0, 0.85e-50, np.arange(6) * 2e-51), + (-0.85e-50, 0, np.arange(-5, 1) * 2e-51), + ] + + integer_data = [ + (-0.1, 1.1, None, np.array([-1, 0, 1, 2])), + (-0.1, 0.95, None, np.array([-0.25, 0, 0.25, 0.5, 0.75, 1.0])), + (1, 55, [1, 1.5, 5, 6, 10], np.array([0, 15, 30, 45, 60])), + ] + + @pytest.mark.parametrize('vmin, vmax, expected', basic_data) + def test_basic(self, vmin, vmax, expected): + loc = mticker.MaxNLocator(nbins=5) + assert_almost_equal(loc.tick_values(vmin, vmax), expected) + + @pytest.mark.parametrize('vmin, vmax, steps, expected', integer_data) + def test_integer(self, vmin, vmax, steps, expected): + loc = mticker.MaxNLocator(nbins=5, integer=True, steps=steps) + assert_almost_equal(loc.tick_values(vmin, vmax), expected) + + @pytest.mark.parametrize('kwargs, errortype, match', [ + ({'foo': 0}, TypeError, + re.escape("set_params() got an unexpected keyword argument 'foo'")), + ({'steps': [2, 1]}, ValueError, "steps argument must be an increasing"), + ({'steps': 2}, ValueError, "steps argument must be an increasing"), + ({'steps': [2, 11]}, ValueError, "steps argument must be an increasing"), + ]) + def test_errors(self, kwargs, errortype, match): + with pytest.raises(errortype, match=match): + mticker.MaxNLocator(**kwargs) + + @pytest.mark.parametrize('steps, result', [ + ([1, 2, 10], [1, 2, 10]), + ([2, 10], [1, 2, 10]), + ([1, 2], [1, 2, 10]), + ([2], [1, 2, 10]), + ]) + def test_padding(self, steps, result): + loc = mticker.MaxNLocator(steps=steps) + assert (loc._steps == result).all() + + +class TestLinearLocator: + def test_basic(self): + loc = mticker.LinearLocator(numticks=3) + test_value = np.array([-0.8, -0.3, 0.2]) + assert_almost_equal(loc.tick_values(-0.8, 0.2), test_value) + + def test_zero_numticks(self): + loc = mticker.LinearLocator(numticks=0) + loc.tick_values(-0.8, 0.2) == [] + + def test_set_params(self): + """ + Create linear locator with presets={}, numticks=2 and change it to + something else. See if change was successful. Should not exception. + """ + loc = mticker.LinearLocator(numticks=2) + loc.set_params(numticks=8, presets={(0, 1): []}) + assert loc.numticks == 8 + assert loc.presets == {(0, 1): []} + + def test_presets(self): + loc = mticker.LinearLocator(presets={(1, 2): [1, 1.25, 1.75], + (0, 2): [0.5, 1.5]}) + assert loc.tick_values(1, 2) == [1, 1.25, 1.75] + assert loc.tick_values(2, 1) == [1, 1.25, 1.75] + assert loc.tick_values(0, 2) == [0.5, 1.5] + assert loc.tick_values(0.0, 2.0) == [0.5, 1.5] + assert (loc.tick_values(0, 1) == np.linspace(0, 1, 11)).all() + + +class TestMultipleLocator: + def test_basic(self): + loc = mticker.MultipleLocator(base=3.147) + test_value = np.array([-9.441, -6.294, -3.147, 0., 3.147, 6.294, + 9.441, 12.588]) + assert_almost_equal(loc.tick_values(-7, 10), test_value) + + def test_basic_with_offset(self): + loc = mticker.MultipleLocator(base=3.147, offset=1.2) + test_value = np.array([-8.241, -5.094, -1.947, 1.2, 4.347, 7.494, + 10.641]) + assert_almost_equal(loc.tick_values(-7, 10), test_value) + + def test_view_limits(self): + """ + Test basic behavior of view limits. + """ + with mpl.rc_context({'axes.autolimit_mode': 'data'}): + loc = mticker.MultipleLocator(base=3.147) + assert_almost_equal(loc.view_limits(-5, 5), (-5, 5)) + + def test_view_limits_round_numbers(self): + """ + Test that everything works properly with 'round_numbers' for auto + limit. + """ + with mpl.rc_context({'axes.autolimit_mode': 'round_numbers'}): + loc = mticker.MultipleLocator(base=3.147) + assert_almost_equal(loc.view_limits(-4, 4), (-6.294, 6.294)) + + def test_view_limits_round_numbers_with_offset(self): + """ + Test that everything works properly with 'round_numbers' for auto + limit. + """ + with mpl.rc_context({'axes.autolimit_mode': 'round_numbers'}): + loc = mticker.MultipleLocator(base=3.147, offset=1.3) + assert_almost_equal(loc.view_limits(-4, 4), (-4.994, 4.447)) + + def test_view_limits_single_bin(self): + """ + Test that 'round_numbers' works properly with a single bin. + """ + with mpl.rc_context({'axes.autolimit_mode': 'round_numbers'}): + loc = mticker.MaxNLocator(nbins=1) + assert_almost_equal(loc.view_limits(-2.3, 2.3), (-4, 4)) + + def test_set_params(self): + """ + Create multiple locator with 0.7 base, and change it to something else. + See if change was successful. + """ + mult = mticker.MultipleLocator(base=0.7) + mult.set_params(base=1.7) + assert mult._edge.step == 1.7 + mult.set_params(offset=3) + assert mult._offset == 3 + + +class TestAutoMinorLocator: + def test_basic(self): + fig, ax = plt.subplots() + ax.set_xlim(0, 1.39) + ax.minorticks_on() + test_value = np.array([0.05, 0.1, 0.15, 0.25, 0.3, 0.35, 0.45, + 0.5, 0.55, 0.65, 0.7, 0.75, 0.85, 0.9, + 0.95, 1.05, 1.1, 1.15, 1.25, 1.3, 1.35]) + assert_almost_equal(ax.xaxis.get_ticklocs(minor=True), test_value) + + # NB: the following values are assuming that *xlim* is [0, 5] + params = [ + (0, 0), # no major tick => no minor tick either + (1, 0) # a single major tick => no minor tick + ] + + def test_first_and_last_minorticks(self): + """ + Test that first and last minor tick appear as expected. + """ + # This test is related to issue #22331 + fig, ax = plt.subplots() + ax.set_xlim(-1.9, 1.9) + ax.xaxis.set_minor_locator(mticker.AutoMinorLocator()) + test_value = np.array([-1.9, -1.8, -1.7, -1.6, -1.4, -1.3, -1.2, -1.1, + -0.9, -0.8, -0.7, -0.6, -0.4, -0.3, -0.2, -0.1, + 0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.1, + 1.2, 1.3, 1.4, 1.6, 1.7, 1.8, 1.9]) + assert_almost_equal(ax.xaxis.get_ticklocs(minor=True), test_value) + + ax.set_xlim(-5, 5) + test_value = np.array([-5.0, -4.5, -3.5, -3.0, -2.5, -1.5, -1.0, -0.5, + 0.5, 1.0, 1.5, 2.5, 3.0, 3.5, 4.5, 5.0]) + assert_almost_equal(ax.xaxis.get_ticklocs(minor=True), test_value) + + @pytest.mark.parametrize('nb_majorticks, expected_nb_minorticks', params) + def test_low_number_of_majorticks( + self, nb_majorticks, expected_nb_minorticks): + # This test is related to issue #8804 + fig, ax = plt.subplots() + xlims = (0, 5) # easier to test the different code paths + ax.set_xlim(*xlims) + ax.set_xticks(np.linspace(xlims[0], xlims[1], nb_majorticks)) + ax.minorticks_on() + ax.xaxis.set_minor_locator(mticker.AutoMinorLocator()) + assert len(ax.xaxis.get_minorticklocs()) == expected_nb_minorticks + + majorstep_minordivisions = [(1, 5), + (2, 4), + (2.5, 5), + (5, 5), + (10, 5)] + + # This test is meant to verify the parameterization for + # test_number_of_minor_ticks + def test_using_all_default_major_steps(self): + with mpl.rc_context({'_internal.classic_mode': False}): + majorsteps = [x[0] for x in self.majorstep_minordivisions] + np.testing.assert_allclose(majorsteps, + mticker.AutoLocator()._steps) + + @pytest.mark.parametrize('major_step, expected_nb_minordivisions', + majorstep_minordivisions) + def test_number_of_minor_ticks( + self, major_step, expected_nb_minordivisions): + fig, ax = plt.subplots() + xlims = (0, major_step) + ax.set_xlim(*xlims) + ax.set_xticks(xlims) + ax.minorticks_on() + ax.xaxis.set_minor_locator(mticker.AutoMinorLocator()) + nb_minor_divisions = len(ax.xaxis.get_minorticklocs()) + 1 + assert nb_minor_divisions == expected_nb_minordivisions + + limits = [(0, 1.39), (0, 0.139), + (0, 0.11e-19), (0, 0.112e-12), + (-2.0e-07, -3.3e-08), (1.20e-06, 1.42e-06), + (-1.34e-06, -1.44e-06), (-8.76e-07, -1.51e-06)] + + reference = [ + [0.05, 0.1, 0.15, 0.25, 0.3, 0.35, 0.45, 0.5, 0.55, 0.65, 0.7, + 0.75, 0.85, 0.9, 0.95, 1.05, 1.1, 1.15, 1.25, 1.3, 1.35], + [0.005, 0.01, 0.015, 0.025, 0.03, 0.035, 0.045, 0.05, 0.055, 0.065, + 0.07, 0.075, 0.085, 0.09, 0.095, 0.105, 0.11, 0.115, 0.125, 0.13, + 0.135], + [5.00e-22, 1.00e-21, 1.50e-21, 2.50e-21, 3.00e-21, 3.50e-21, 4.50e-21, + 5.00e-21, 5.50e-21, 6.50e-21, 7.00e-21, 7.50e-21, 8.50e-21, 9.00e-21, + 9.50e-21, 1.05e-20, 1.10e-20], + [5.00e-15, 1.00e-14, 1.50e-14, 2.50e-14, 3.00e-14, 3.50e-14, 4.50e-14, + 5.00e-14, 5.50e-14, 6.50e-14, 7.00e-14, 7.50e-14, 8.50e-14, 9.00e-14, + 9.50e-14, 1.05e-13, 1.10e-13], + [-1.95e-07, -1.90e-07, -1.85e-07, -1.75e-07, -1.70e-07, -1.65e-07, + -1.55e-07, -1.50e-07, -1.45e-07, -1.35e-07, -1.30e-07, -1.25e-07, + -1.15e-07, -1.10e-07, -1.05e-07, -9.50e-08, -9.00e-08, -8.50e-08, + -7.50e-08, -7.00e-08, -6.50e-08, -5.50e-08, -5.00e-08, -4.50e-08, + -3.50e-08], + [1.21e-06, 1.22e-06, 1.23e-06, 1.24e-06, 1.26e-06, 1.27e-06, 1.28e-06, + 1.29e-06, 1.31e-06, 1.32e-06, 1.33e-06, 1.34e-06, 1.36e-06, 1.37e-06, + 1.38e-06, 1.39e-06, 1.41e-06, 1.42e-06], + [-1.435e-06, -1.430e-06, -1.425e-06, -1.415e-06, -1.410e-06, + -1.405e-06, -1.395e-06, -1.390e-06, -1.385e-06, -1.375e-06, + -1.370e-06, -1.365e-06, -1.355e-06, -1.350e-06, -1.345e-06], + [-1.48e-06, -1.46e-06, -1.44e-06, -1.42e-06, -1.38e-06, -1.36e-06, + -1.34e-06, -1.32e-06, -1.28e-06, -1.26e-06, -1.24e-06, -1.22e-06, + -1.18e-06, -1.16e-06, -1.14e-06, -1.12e-06, -1.08e-06, -1.06e-06, + -1.04e-06, -1.02e-06, -9.80e-07, -9.60e-07, -9.40e-07, -9.20e-07, + -8.80e-07]] + + additional_data = list(zip(limits, reference)) + + @pytest.mark.parametrize('lim, ref', additional_data) + def test_additional(self, lim, ref): + fig, ax = plt.subplots() + + ax.minorticks_on() + ax.grid(True, 'minor', 'y', linewidth=1) + ax.grid(True, 'major', color='k', linewidth=1) + ax.set_ylim(lim) + + assert_almost_equal(ax.yaxis.get_ticklocs(minor=True), ref) + + @pytest.mark.parametrize('use_rcparam', [False, True]) + @pytest.mark.parametrize( + 'lim, ref', [ + ((0, 1.39), + [0.05, 0.1, 0.15, 0.25, 0.3, 0.35, 0.45, 0.5, 0.55, 0.65, 0.7, + 0.75, 0.85, 0.9, 0.95, 1.05, 1.1, 1.15, 1.25, 1.3, 1.35]), + ((0, 0.139), + [0.005, 0.01, 0.015, 0.025, 0.03, 0.035, 0.045, 0.05, 0.055, + 0.065, 0.07, 0.075, 0.085, 0.09, 0.095, 0.105, 0.11, 0.115, + 0.125, 0.13, 0.135]), + ]) + def test_number_of_minor_ticks_auto(self, lim, ref, use_rcparam): + if use_rcparam: + context = {'xtick.minor.ndivs': 'auto', 'ytick.minor.ndivs': 'auto'} + kwargs = {} + else: + context = {} + kwargs = {'n': 'auto'} + + with mpl.rc_context(context): + fig, ax = plt.subplots() + ax.set_xlim(*lim) + ax.set_ylim(*lim) + ax.xaxis.set_minor_locator(mticker.AutoMinorLocator(**kwargs)) + ax.yaxis.set_minor_locator(mticker.AutoMinorLocator(**kwargs)) + assert_almost_equal(ax.xaxis.get_ticklocs(minor=True), ref) + assert_almost_equal(ax.yaxis.get_ticklocs(minor=True), ref) + + @pytest.mark.parametrize('use_rcparam', [False, True]) + @pytest.mark.parametrize( + 'n, lim, ref', [ + (2, (0, 4), [0.5, 1.5, 2.5, 3.5]), + (4, (0, 2), [0.25, 0.5, 0.75, 1.25, 1.5, 1.75]), + (10, (0, 1), [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]), + ]) + def test_number_of_minor_ticks_int(self, n, lim, ref, use_rcparam): + if use_rcparam: + context = {'xtick.minor.ndivs': n, 'ytick.minor.ndivs': n} + kwargs = {} + else: + context = {} + kwargs = {'n': n} + + with mpl.rc_context(context): + fig, ax = plt.subplots() + ax.set_xlim(*lim) + ax.set_ylim(*lim) + ax.xaxis.set_major_locator(mticker.MultipleLocator(1)) + ax.xaxis.set_minor_locator(mticker.AutoMinorLocator(**kwargs)) + ax.yaxis.set_major_locator(mticker.MultipleLocator(1)) + ax.yaxis.set_minor_locator(mticker.AutoMinorLocator(**kwargs)) + assert_almost_equal(ax.xaxis.get_ticklocs(minor=True), ref) + assert_almost_equal(ax.yaxis.get_ticklocs(minor=True), ref) + + +class TestLogLocator: + def test_basic(self): + loc = mticker.LogLocator(numticks=5) + with pytest.raises(ValueError): + loc.tick_values(0, 1000) + + test_value = np.array([1.00000000e-05, 1.00000000e-03, 1.00000000e-01, + 1.00000000e+01, 1.00000000e+03, 1.00000000e+05, + 1.00000000e+07, 1.000000000e+09]) + assert_almost_equal(loc.tick_values(0.001, 1.1e5), test_value) + + loc = mticker.LogLocator(base=2) + test_value = np.array([0.5, 1., 2., 4., 8., 16., 32., 64., 128., 256.]) + assert_almost_equal(loc.tick_values(1, 100), test_value) + + def test_polar_axes(self): + """ + Polar Axes have a different ticking logic. + """ + fig, ax = plt.subplots(subplot_kw={'projection': 'polar'}) + ax.set_yscale('log') + ax.set_ylim(1, 100) + assert_array_equal(ax.get_yticks(), [10, 100, 1000]) + + def test_switch_to_autolocator(self): + loc = mticker.LogLocator(subs="all") + assert_array_equal(loc.tick_values(0.45, 0.55), + [0.44, 0.46, 0.48, 0.5, 0.52, 0.54, 0.56]) + # check that we *skip* 1.0, and 10, because this is a minor locator + loc = mticker.LogLocator(subs=np.arange(2, 10)) + assert 1.0 not in loc.tick_values(0.9, 20.) + assert 10.0 not in loc.tick_values(0.9, 20.) + + def test_set_params(self): + """ + Create log locator with default value, base=10.0, subs=[1.0], + numticks=15 and change it to something else. + See if change was successful. Should not raise exception. + """ + loc = mticker.LogLocator() + loc.set_params(numticks=7, subs=[2.0], base=4) + assert loc.numticks == 7 + assert loc._base == 4 + assert list(loc._subs) == [2.0] + + def test_tick_values_correct(self): + ll = mticker.LogLocator(subs=(1, 2, 5)) + test_value = np.array([1.e-01, 2.e-01, 5.e-01, 1.e+00, 2.e+00, 5.e+00, + 1.e+01, 2.e+01, 5.e+01, 1.e+02, 2.e+02, 5.e+02, + 1.e+03, 2.e+03, 5.e+03, 1.e+04, 2.e+04, 5.e+04, + 1.e+05, 2.e+05, 5.e+05, 1.e+06, 2.e+06, 5.e+06, + 1.e+07, 2.e+07, 5.e+07, 1.e+08, 2.e+08, 5.e+08]) + assert_almost_equal(ll.tick_values(1, 1e7), test_value) + + def test_tick_values_not_empty(self): + mpl.rcParams['_internal.classic_mode'] = False + ll = mticker.LogLocator(subs=(1, 2, 5)) + test_value = np.array([1.e-01, 2.e-01, 5.e-01, 1.e+00, 2.e+00, 5.e+00, + 1.e+01, 2.e+01, 5.e+01, 1.e+02, 2.e+02, 5.e+02, + 1.e+03, 2.e+03, 5.e+03, 1.e+04, 2.e+04, 5.e+04, + 1.e+05, 2.e+05, 5.e+05, 1.e+06, 2.e+06, 5.e+06, + 1.e+07, 2.e+07, 5.e+07, 1.e+08, 2.e+08, 5.e+08, + 1.e+09, 2.e+09, 5.e+09]) + assert_almost_equal(ll.tick_values(1, 1e8), test_value) + + def test_multiple_shared_axes(self): + rng = np.random.default_rng(19680801) + dummy_data = [rng.normal(size=100), [], []] + fig, axes = plt.subplots(len(dummy_data), sharex=True, sharey=True) + + for ax, data in zip(axes.flatten(), dummy_data): + ax.hist(data, bins=10) + ax.set_yscale('log', nonpositive='clip') + + for ax in axes.flatten(): + assert all(ax.get_yticks() == axes[0].get_yticks()) + assert ax.get_ylim() == axes[0].get_ylim() + + +class TestNullLocator: + def test_set_params(self): + """ + Create null locator, and attempt to call set_params() on it. + Should not exception, and should raise a warning. + """ + loc = mticker.NullLocator() + with pytest.warns(UserWarning): + loc.set_params() + + +class _LogitHelper: + @staticmethod + def isclose(x, y): + return (np.isclose(-np.log(1/x-1), -np.log(1/y-1)) + if 0 < x < 1 and 0 < y < 1 else False) + + @staticmethod + def assert_almost_equal(x, y): + ax = np.array(x) + ay = np.array(y) + assert np.all(ax > 0) and np.all(ax < 1) + assert np.all(ay > 0) and np.all(ay < 1) + lx = -np.log(1/ax-1) + ly = -np.log(1/ay-1) + assert_almost_equal(lx, ly) + + +class TestLogitLocator: + ref_basic_limits = [ + (5e-2, 1 - 5e-2), + (5e-3, 1 - 5e-3), + (5e-4, 1 - 5e-4), + (5e-5, 1 - 5e-5), + (5e-6, 1 - 5e-6), + (5e-7, 1 - 5e-7), + (5e-8, 1 - 5e-8), + (5e-9, 1 - 5e-9), + ] + + ref_basic_major_ticks = [ + 1 / (10 ** np.arange(1, 3)), + 1 / (10 ** np.arange(1, 4)), + 1 / (10 ** np.arange(1, 5)), + 1 / (10 ** np.arange(1, 6)), + 1 / (10 ** np.arange(1, 7)), + 1 / (10 ** np.arange(1, 8)), + 1 / (10 ** np.arange(1, 9)), + 1 / (10 ** np.arange(1, 10)), + ] + + ref_maxn_limits = [(0.4, 0.6), (5e-2, 2e-1), (1 - 2e-1, 1 - 5e-2)] + + @pytest.mark.parametrize( + "lims, expected_low_ticks", + zip(ref_basic_limits, ref_basic_major_ticks), + ) + def test_basic_major(self, lims, expected_low_ticks): + """ + Create logit locator with huge number of major, and tests ticks. + """ + expected_ticks = sorted( + [*expected_low_ticks, 0.5, *(1 - expected_low_ticks)] + ) + loc = mticker.LogitLocator(nbins=100) + _LogitHelper.assert_almost_equal( + loc.tick_values(*lims), + expected_ticks + ) + + @pytest.mark.parametrize("lims", ref_maxn_limits) + def test_maxn_major(self, lims): + """ + When the axis is zoomed, the locator must have the same behavior as + MaxNLocator. + """ + loc = mticker.LogitLocator(nbins=100) + maxn_loc = mticker.MaxNLocator(nbins=100, steps=[1, 2, 5, 10]) + for nbins in (4, 8, 16): + loc.set_params(nbins=nbins) + maxn_loc.set_params(nbins=nbins) + ticks = loc.tick_values(*lims) + maxn_ticks = maxn_loc.tick_values(*lims) + assert ticks.shape == maxn_ticks.shape + assert (ticks == maxn_ticks).all() + + @pytest.mark.parametrize("lims", ref_basic_limits + ref_maxn_limits) + def test_nbins_major(self, lims): + """ + Assert logit locator for respecting nbins param. + """ + + basic_needed = int(-np.floor(np.log10(lims[0]))) * 2 + 1 + loc = mticker.LogitLocator(nbins=100) + for nbins in range(basic_needed, 2, -1): + loc.set_params(nbins=nbins) + assert len(loc.tick_values(*lims)) <= nbins + 2 + + @pytest.mark.parametrize( + "lims, expected_low_ticks", + zip(ref_basic_limits, ref_basic_major_ticks), + ) + def test_minor(self, lims, expected_low_ticks): + """ + In large scale, test the presence of minor, + and assert no minor when major are subsampled. + """ + + expected_ticks = sorted( + [*expected_low_ticks, 0.5, *(1 - expected_low_ticks)] + ) + basic_needed = len(expected_ticks) + loc = mticker.LogitLocator(nbins=100) + minor_loc = mticker.LogitLocator(nbins=100, minor=True) + for nbins in range(basic_needed, 2, -1): + loc.set_params(nbins=nbins) + minor_loc.set_params(nbins=nbins) + major_ticks = loc.tick_values(*lims) + minor_ticks = minor_loc.tick_values(*lims) + if len(major_ticks) >= len(expected_ticks): + # no subsample, we must have a lot of minors ticks + assert (len(major_ticks) - 1) * 5 < len(minor_ticks) + else: + # subsample + _LogitHelper.assert_almost_equal( + sorted([*major_ticks, *minor_ticks]), expected_ticks) + + def test_minor_attr(self): + loc = mticker.LogitLocator(nbins=100) + assert not loc.minor + loc.minor = True + assert loc.minor + loc.set_params(minor=False) + assert not loc.minor + + acceptable_vmin_vmax = [ + *(2.5 ** np.arange(-3, 0)), + *(1 - 2.5 ** np.arange(-3, 0)), + ] + + @pytest.mark.parametrize( + "lims", + [ + (a, b) + for (a, b) in itertools.product(acceptable_vmin_vmax, repeat=2) + if a != b + ], + ) + def test_nonsingular_ok(self, lims): + """ + Create logit locator, and test the nonsingular method for acceptable + value + """ + loc = mticker.LogitLocator() + lims2 = loc.nonsingular(*lims) + assert sorted(lims) == sorted(lims2) + + @pytest.mark.parametrize("okval", acceptable_vmin_vmax) + def test_nonsingular_nok(self, okval): + """ + Create logit locator, and test the nonsingular method for non + acceptable value + """ + loc = mticker.LogitLocator() + vmin, vmax = (-1, okval) + vmin2, vmax2 = loc.nonsingular(vmin, vmax) + assert vmax2 == vmax + assert 0 < vmin2 < vmax2 + vmin, vmax = (okval, 2) + vmin2, vmax2 = loc.nonsingular(vmin, vmax) + assert vmin2 == vmin + assert vmin2 < vmax2 < 1 + + +class TestFixedLocator: + def test_set_params(self): + """ + Create fixed locator with 5 nbins, and change it to something else. + See if change was successful. + Should not exception. + """ + fixed = mticker.FixedLocator(range(0, 24), nbins=5) + fixed.set_params(nbins=7) + assert fixed.nbins == 7 + + +class TestIndexLocator: + def test_set_params(self): + """ + Create index locator with 3 base, 4 offset. and change it to something + else. See if change was successful. + Should not exception. + """ + index = mticker.IndexLocator(base=3, offset=4) + index.set_params(base=7, offset=7) + assert index._base == 7 + assert index.offset == 7 + + +class TestSymmetricalLogLocator: + def test_set_params(self): + """ + Create symmetrical log locator with default subs =[1.0] numticks = 15, + and change it to something else. + See if change was successful. + Should not exception. + """ + sym = mticker.SymmetricalLogLocator(base=10, linthresh=1) + sym.set_params(subs=[2.0], numticks=8) + assert sym._subs == [2.0] + assert sym.numticks == 8 + + @pytest.mark.parametrize( + 'vmin, vmax, expected', + [ + (0, 1, [0, 1]), + (-1, 1, [-1, 0, 1]), + ], + ) + def test_values(self, vmin, vmax, expected): + # https://github.com/matplotlib/matplotlib/issues/25945 + sym = mticker.SymmetricalLogLocator(base=10, linthresh=1) + ticks = sym.tick_values(vmin=vmin, vmax=vmax) + assert_array_equal(ticks, expected) + + def test_subs(self): + sym = mticker.SymmetricalLogLocator(base=10, linthresh=1, subs=[2.0, 4.0]) + sym.create_dummy_axis() + sym.axis.set_view_interval(-10, 10) + assert_array_equal(sym(), [-20, -40, -2, -4, 0, 2, 4, 20, 40]) + + def test_extending(self): + sym = mticker.SymmetricalLogLocator(base=10, linthresh=1) + sym.create_dummy_axis() + sym.axis.set_view_interval(8, 9) + assert (sym() == [1.0]).all() + sym.axis.set_view_interval(8, 12) + assert (sym() == [1.0, 10.0]).all() + assert sym.view_limits(10, 10) == (1, 100) + assert sym.view_limits(-10, -10) == (-100, -1) + assert sym.view_limits(0, 0) == (-0.001, 0.001) + + +class TestAsinhLocator: + def test_init(self): + lctr = mticker.AsinhLocator(linear_width=2.718, numticks=19) + assert lctr.linear_width == 2.718 + assert lctr.numticks == 19 + assert lctr.base == 10 + + def test_set_params(self): + lctr = mticker.AsinhLocator(linear_width=5, + numticks=17, symthresh=0.125, + base=4, subs=(2.5, 3.25)) + assert lctr.numticks == 17 + assert lctr.symthresh == 0.125 + assert lctr.base == 4 + assert lctr.subs == (2.5, 3.25) + + lctr.set_params(numticks=23) + assert lctr.numticks == 23 + lctr.set_params(None) + assert lctr.numticks == 23 + + lctr.set_params(symthresh=0.5) + assert lctr.symthresh == 0.5 + lctr.set_params(symthresh=None) + assert lctr.symthresh == 0.5 + + lctr.set_params(base=7) + assert lctr.base == 7 + lctr.set_params(base=None) + assert lctr.base == 7 + + lctr.set_params(subs=(2, 4.125)) + assert lctr.subs == (2, 4.125) + lctr.set_params(subs=None) + assert lctr.subs == (2, 4.125) + lctr.set_params(subs=[]) + assert lctr.subs is None + + def test_linear_values(self): + lctr = mticker.AsinhLocator(linear_width=100, numticks=11, base=0) + + assert_almost_equal(lctr.tick_values(-1, 1), + np.arange(-1, 1.01, 0.2)) + assert_almost_equal(lctr.tick_values(-0.1, 0.1), + np.arange(-0.1, 0.101, 0.02)) + assert_almost_equal(lctr.tick_values(-0.01, 0.01), + np.arange(-0.01, 0.0101, 0.002)) + + def test_wide_values(self): + lctr = mticker.AsinhLocator(linear_width=0.1, numticks=11, base=0) + + assert_almost_equal(lctr.tick_values(-100, 100), + [-100, -20, -5, -1, -0.2, + 0, 0.2, 1, 5, 20, 100]) + assert_almost_equal(lctr.tick_values(-1000, 1000), + [-1000, -100, -20, -3, -0.4, + 0, 0.4, 3, 20, 100, 1000]) + + def test_near_zero(self): + """Check that manually injected zero will supersede nearby tick""" + lctr = mticker.AsinhLocator(linear_width=100, numticks=3, base=0) + + assert_almost_equal(lctr.tick_values(-1.1, 0.9), [-1.0, 0.0, 0.9]) + + def test_fallback(self): + lctr = mticker.AsinhLocator(1.0, numticks=11) + + assert_almost_equal(lctr.tick_values(101, 102), + np.arange(101, 102.01, 0.1)) + + def test_symmetrizing(self): + lctr = mticker.AsinhLocator(linear_width=1, numticks=3, + symthresh=0.25, base=0) + lctr.create_dummy_axis() + + lctr.axis.set_view_interval(-1, 2) + assert_almost_equal(lctr(), [-1, 0, 2]) + + lctr.axis.set_view_interval(-1, 0.9) + assert_almost_equal(lctr(), [-1, 0, 1]) + + lctr.axis.set_view_interval(-0.85, 1.05) + assert_almost_equal(lctr(), [-1, 0, 1]) + + lctr.axis.set_view_interval(1, 1.1) + assert_almost_equal(lctr(), [1, 1.05, 1.1]) + + def test_base_rounding(self): + lctr10 = mticker.AsinhLocator(linear_width=1, numticks=8, + base=10, subs=(1, 3, 5)) + assert_almost_equal(lctr10.tick_values(-110, 110), + [-500, -300, -100, -50, -30, -10, -5, -3, -1, + -0.5, -0.3, -0.1, 0, 0.1, 0.3, 0.5, + 1, 3, 5, 10, 30, 50, 100, 300, 500]) + + lctr5 = mticker.AsinhLocator(linear_width=1, numticks=20, base=5) + assert_almost_equal(lctr5.tick_values(-1050, 1050), + [-625, -125, -25, -5, -1, -0.2, 0, + 0.2, 1, 5, 25, 125, 625]) + + +class TestScalarFormatter: + offset_data = [ + (123, 189, 0), + (-189, -123, 0), + (12341, 12349, 12340), + (-12349, -12341, -12340), + (99999.5, 100010.5, 100000), + (-100010.5, -99999.5, -100000), + (99990.5, 100000.5, 100000), + (-100000.5, -99990.5, -100000), + (1233999, 1234001, 1234000), + (-1234001, -1233999, -1234000), + (1, 1, 1), + (123, 123, 0), + # Test cases courtesy of @WeatherGod + (.4538, .4578, .45), + (3789.12, 3783.1, 3780), + (45124.3, 45831.75, 45000), + (0.000721, 0.0007243, 0.00072), + (12592.82, 12591.43, 12590), + (9., 12., 0), + (900., 1200., 0), + (1900., 1200., 0), + (0.99, 1.01, 1), + (9.99, 10.01, 10), + (99.99, 100.01, 100), + (5.99, 6.01, 6), + (15.99, 16.01, 16), + (-0.452, 0.492, 0), + (-0.492, 0.492, 0), + (12331.4, 12350.5, 12300), + (-12335.3, 12335.3, 0), + ] + + use_offset_data = [True, False] + + useMathText_data = [True, False] + + # (sci_type, scilimits, lim, orderOfMag, fewticks) + scilimits_data = [ + (False, (0, 0), (10.0, 20.0), 0, False), + (True, (-2, 2), (-10, 20), 0, False), + (True, (-2, 2), (-20, 10), 0, False), + (True, (-2, 2), (-110, 120), 2, False), + (True, (-2, 2), (-120, 110), 2, False), + (True, (-2, 2), (-.001, 0.002), -3, False), + (True, (-7, 7), (0.18e10, 0.83e10), 9, True), + (True, (0, 0), (-1e5, 1e5), 5, False), + (True, (6, 6), (-1e5, 1e5), 6, False), + ] + + cursor_data = [ + [0., "0.000"], + [0.0123, "0.012"], + [0.123, "0.123"], + [1.23, "1.230"], + [12.3, "12.300"], + ] + + format_data = [ + (.1, "1e-1"), + (.11, "1.1e-1"), + (1e8, "1e8"), + (1.1e8, "1.1e8"), + ] + + @pytest.mark.parametrize('unicode_minus, result', + [(True, "\N{MINUS SIGN}1"), (False, "-1")]) + def test_unicode_minus(self, unicode_minus, result): + mpl.rcParams['axes.unicode_minus'] = unicode_minus + assert ( + plt.gca().xaxis.get_major_formatter().format_data_short(-1).strip() + == result) + + @pytest.mark.parametrize('left, right, offset', offset_data) + def test_offset_value(self, left, right, offset): + fig, ax = plt.subplots() + formatter = ax.xaxis.get_major_formatter() + + with (pytest.warns(UserWarning, match='Attempting to set identical') + if left == right else nullcontext()): + ax.set_xlim(left, right) + ax.xaxis._update_ticks() + assert formatter.offset == offset + + with (pytest.warns(UserWarning, match='Attempting to set identical') + if left == right else nullcontext()): + ax.set_xlim(right, left) + ax.xaxis._update_ticks() + assert formatter.offset == offset + + @pytest.mark.parametrize('use_offset', use_offset_data) + def test_use_offset(self, use_offset): + with mpl.rc_context({'axes.formatter.useoffset': use_offset}): + tmp_form = mticker.ScalarFormatter() + assert use_offset == tmp_form.get_useOffset() + assert tmp_form.offset == 0 + + @pytest.mark.parametrize('use_math_text', useMathText_data) + def test_useMathText(self, use_math_text): + with mpl.rc_context({'axes.formatter.use_mathtext': use_math_text}): + tmp_form = mticker.ScalarFormatter() + assert use_math_text == tmp_form.get_useMathText() + + def test_set_use_offset_float(self): + tmp_form = mticker.ScalarFormatter() + tmp_form.set_useOffset(0.5) + assert not tmp_form.get_useOffset() + assert tmp_form.offset == 0.5 + + def test_use_locale(self): + conv = locale.localeconv() + sep = conv['thousands_sep'] + if not sep or conv['grouping'][-1:] in ([], [locale.CHAR_MAX]): + pytest.skip('Locale does not apply grouping') # pragma: no cover + + with mpl.rc_context({'axes.formatter.use_locale': True}): + tmp_form = mticker.ScalarFormatter() + assert tmp_form.get_useLocale() + + tmp_form.create_dummy_axis() + tmp_form.axis.set_data_interval(0, 10) + tmp_form.set_locs([1, 2, 3]) + assert sep in tmp_form(1e9) + + @pytest.mark.parametrize( + 'sci_type, scilimits, lim, orderOfMag, fewticks', scilimits_data) + def test_scilimits(self, sci_type, scilimits, lim, orderOfMag, fewticks): + tmp_form = mticker.ScalarFormatter() + tmp_form.set_scientific(sci_type) + tmp_form.set_powerlimits(scilimits) + fig, ax = plt.subplots() + ax.yaxis.set_major_formatter(tmp_form) + ax.set_ylim(*lim) + if fewticks: + ax.yaxis.set_major_locator(mticker.MaxNLocator(4)) + + tmp_form.set_locs(ax.yaxis.get_majorticklocs()) + assert orderOfMag == tmp_form.orderOfMagnitude + + @pytest.mark.parametrize('value, expected', format_data) + def test_format_data(self, value, expected): + mpl.rcParams['axes.unicode_minus'] = False + sf = mticker.ScalarFormatter() + assert sf.format_data(value) == expected + + @pytest.mark.parametrize('data, expected', cursor_data) + def test_cursor_precision(self, data, expected): + fig, ax = plt.subplots() + ax.set_xlim(-1, 1) # Pointing precision of 0.001. + fmt = ax.xaxis.get_major_formatter().format_data_short + assert fmt(data) == expected + + @pytest.mark.parametrize('data, expected', cursor_data) + def test_cursor_dummy_axis(self, data, expected): + # Issue #17624 + sf = mticker.ScalarFormatter() + sf.create_dummy_axis() + sf.axis.set_view_interval(0, 10) + fmt = sf.format_data_short + assert fmt(data) == expected + assert sf.axis.get_tick_space() == 9 + assert sf.axis.get_minpos() == 0 + + def test_mathtext_ticks(self): + mpl.rcParams.update({ + 'font.family': 'serif', + 'font.serif': 'cmr10', + 'axes.formatter.use_mathtext': False + }) + + if parse_version(pytest.__version__).major < 8: + with pytest.warns(UserWarning, match='cmr10 font should ideally'): + fig, ax = plt.subplots() + ax.set_xticks([-1, 0, 1]) + fig.canvas.draw() + else: + with (pytest.warns(UserWarning, match="Glyph 8722"), + pytest.warns(UserWarning, match='cmr10 font should ideally')): + fig, ax = plt.subplots() + ax.set_xticks([-1, 0, 1]) + fig.canvas.draw() + + def test_cmr10_substitutions(self, caplog): + mpl.rcParams.update({ + 'font.family': 'cmr10', + 'mathtext.fontset': 'cm', + 'axes.formatter.use_mathtext': True, + }) + + # Test that it does not log a warning about missing glyphs. + with caplog.at_level(logging.WARNING, logger='matplotlib.mathtext'): + fig, ax = plt.subplots() + ax.plot([-0.03, 0.05], [40, 0.05]) + ax.set_yscale('log') + yticks = [0.02, 0.3, 4, 50] + formatter = mticker.LogFormatterSciNotation() + ax.set_yticks(yticks, map(formatter, yticks)) + fig.canvas.draw() + assert not caplog.text + + def test_empty_locs(self): + sf = mticker.ScalarFormatter() + sf.set_locs([]) + assert sf(0.5) == '' + + +class TestLogFormatterExponent: + param_data = [ + (True, 4, np.arange(-3, 4.0), np.arange(-3, 4.0), + ['-3', '-2', '-1', '0', '1', '2', '3']), + # With labelOnlyBase=False, non-integer powers should be nicely + # formatted. + (False, 10, np.array([0.1, 0.00001, np.pi, 0.2, -0.2, -0.00001]), + range(6), ['0.1', '1e-05', '3.14', '0.2', '-0.2', '-1e-05']), + (False, 50, np.array([3, 5, 12, 42], dtype=float), range(6), + ['3', '5', '12', '42']), + ] + + base_data = [2.0, 5.0, 10.0, np.pi, np.e] + + @pytest.mark.parametrize( + 'labelOnlyBase, exponent, locs, positions, expected', param_data) + @pytest.mark.parametrize('base', base_data) + def test_basic(self, labelOnlyBase, base, exponent, locs, positions, + expected): + formatter = mticker.LogFormatterExponent(base=base, + labelOnlyBase=labelOnlyBase) + formatter.create_dummy_axis() + formatter.axis.set_view_interval(1, base**exponent) + vals = base**locs + labels = [formatter(x, pos) for (x, pos) in zip(vals, positions)] + expected = [label.replace('-', '\N{Minus Sign}') for label in expected] + assert labels == expected + + def test_blank(self): + # Should be a blank string for non-integer powers if labelOnlyBase=True + formatter = mticker.LogFormatterExponent(base=10, labelOnlyBase=True) + formatter.create_dummy_axis() + formatter.axis.set_view_interval(1, 10) + assert formatter(10**0.1) == '' + + +class TestLogFormatterMathtext: + fmt = mticker.LogFormatterMathtext() + test_data = [ + (0, 1, '$\\mathdefault{10^{0}}$'), + (0, 1e-2, '$\\mathdefault{10^{-2}}$'), + (0, 1e2, '$\\mathdefault{10^{2}}$'), + (3, 1, '$\\mathdefault{1}$'), + (3, 1e-2, '$\\mathdefault{0.01}$'), + (3, 1e2, '$\\mathdefault{100}$'), + (3, 1e-3, '$\\mathdefault{10^{-3}}$'), + (3, 1e3, '$\\mathdefault{10^{3}}$'), + ] + + @pytest.mark.parametrize('min_exponent, value, expected', test_data) + def test_min_exponent(self, min_exponent, value, expected): + with mpl.rc_context({'axes.formatter.min_exponent': min_exponent}): + assert self.fmt(value) == expected + + +class TestLogFormatterSciNotation: + test_data = [ + (2, 0.03125, '$\\mathdefault{2^{-5}}$'), + (2, 1, '$\\mathdefault{2^{0}}$'), + (2, 32, '$\\mathdefault{2^{5}}$'), + (2, 0.0375, '$\\mathdefault{1.2\\times2^{-5}}$'), + (2, 1.2, '$\\mathdefault{1.2\\times2^{0}}$'), + (2, 38.4, '$\\mathdefault{1.2\\times2^{5}}$'), + (10, -1, '$\\mathdefault{-10^{0}}$'), + (10, 1e-05, '$\\mathdefault{10^{-5}}$'), + (10, 1, '$\\mathdefault{10^{0}}$'), + (10, 100000, '$\\mathdefault{10^{5}}$'), + (10, 2e-05, '$\\mathdefault{2\\times10^{-5}}$'), + (10, 2, '$\\mathdefault{2\\times10^{0}}$'), + (10, 200000, '$\\mathdefault{2\\times10^{5}}$'), + (10, 5e-05, '$\\mathdefault{5\\times10^{-5}}$'), + (10, 5, '$\\mathdefault{5\\times10^{0}}$'), + (10, 500000, '$\\mathdefault{5\\times10^{5}}$'), + ] + + @mpl.style.context('default') + @pytest.mark.parametrize('base, value, expected', test_data) + def test_basic(self, base, value, expected): + formatter = mticker.LogFormatterSciNotation(base=base) + with mpl.rc_context({'text.usetex': False}): + assert formatter(value) == expected + + +class TestLogFormatter: + pprint_data = [ + (3.141592654e-05, 0.001, '3.142e-5'), + (0.0003141592654, 0.001, '3.142e-4'), + (0.003141592654, 0.001, '3.142e-3'), + (0.03141592654, 0.001, '3.142e-2'), + (0.3141592654, 0.001, '3.142e-1'), + (3.141592654, 0.001, '3.142'), + (31.41592654, 0.001, '3.142e1'), + (314.1592654, 0.001, '3.142e2'), + (3141.592654, 0.001, '3.142e3'), + (31415.92654, 0.001, '3.142e4'), + (314159.2654, 0.001, '3.142e5'), + (1e-05, 0.001, '1e-5'), + (0.0001, 0.001, '1e-4'), + (0.001, 0.001, '1e-3'), + (0.01, 0.001, '1e-2'), + (0.1, 0.001, '1e-1'), + (1, 0.001, '1'), + (10, 0.001, '10'), + (100, 0.001, '100'), + (1000, 0.001, '1000'), + (10000, 0.001, '1e4'), + (100000, 0.001, '1e5'), + (3.141592654e-05, 0.015, '0'), + (0.0003141592654, 0.015, '0'), + (0.003141592654, 0.015, '0.003'), + (0.03141592654, 0.015, '0.031'), + (0.3141592654, 0.015, '0.314'), + (3.141592654, 0.015, '3.142'), + (31.41592654, 0.015, '31.416'), + (314.1592654, 0.015, '314.159'), + (3141.592654, 0.015, '3141.593'), + (31415.92654, 0.015, '31415.927'), + (314159.2654, 0.015, '314159.265'), + (1e-05, 0.015, '0'), + (0.0001, 0.015, '0'), + (0.001, 0.015, '0.001'), + (0.01, 0.015, '0.01'), + (0.1, 0.015, '0.1'), + (1, 0.015, '1'), + (10, 0.015, '10'), + (100, 0.015, '100'), + (1000, 0.015, '1000'), + (10000, 0.015, '10000'), + (100000, 0.015, '100000'), + (3.141592654e-05, 0.5, '0'), + (0.0003141592654, 0.5, '0'), + (0.003141592654, 0.5, '0.003'), + (0.03141592654, 0.5, '0.031'), + (0.3141592654, 0.5, '0.314'), + (3.141592654, 0.5, '3.142'), + (31.41592654, 0.5, '31.416'), + (314.1592654, 0.5, '314.159'), + (3141.592654, 0.5, '3141.593'), + (31415.92654, 0.5, '31415.927'), + (314159.2654, 0.5, '314159.265'), + (1e-05, 0.5, '0'), + (0.0001, 0.5, '0'), + (0.001, 0.5, '0.001'), + (0.01, 0.5, '0.01'), + (0.1, 0.5, '0.1'), + (1, 0.5, '1'), + (10, 0.5, '10'), + (100, 0.5, '100'), + (1000, 0.5, '1000'), + (10000, 0.5, '10000'), + (100000, 0.5, '100000'), + (3.141592654e-05, 5, '0'), + (0.0003141592654, 5, '0'), + (0.003141592654, 5, '0'), + (0.03141592654, 5, '0.03'), + (0.3141592654, 5, '0.31'), + (3.141592654, 5, '3.14'), + (31.41592654, 5, '31.42'), + (314.1592654, 5, '314.16'), + (3141.592654, 5, '3141.59'), + (31415.92654, 5, '31415.93'), + (314159.2654, 5, '314159.27'), + (1e-05, 5, '0'), + (0.0001, 5, '0'), + (0.001, 5, '0'), + (0.01, 5, '0.01'), + (0.1, 5, '0.1'), + (1, 5, '1'), + (10, 5, '10'), + (100, 5, '100'), + (1000, 5, '1000'), + (10000, 5, '10000'), + (100000, 5, '100000'), + (3.141592654e-05, 100, '0'), + (0.0003141592654, 100, '0'), + (0.003141592654, 100, '0'), + (0.03141592654, 100, '0'), + (0.3141592654, 100, '0.3'), + (3.141592654, 100, '3.1'), + (31.41592654, 100, '31.4'), + (314.1592654, 100, '314.2'), + (3141.592654, 100, '3141.6'), + (31415.92654, 100, '31415.9'), + (314159.2654, 100, '314159.3'), + (1e-05, 100, '0'), + (0.0001, 100, '0'), + (0.001, 100, '0'), + (0.01, 100, '0'), + (0.1, 100, '0.1'), + (1, 100, '1'), + (10, 100, '10'), + (100, 100, '100'), + (1000, 100, '1000'), + (10000, 100, '10000'), + (100000, 100, '100000'), + (3.141592654e-05, 1000000.0, '3.1e-5'), + (0.0003141592654, 1000000.0, '3.1e-4'), + (0.003141592654, 1000000.0, '3.1e-3'), + (0.03141592654, 1000000.0, '3.1e-2'), + (0.3141592654, 1000000.0, '3.1e-1'), + (3.141592654, 1000000.0, '3.1'), + (31.41592654, 1000000.0, '3.1e1'), + (314.1592654, 1000000.0, '3.1e2'), + (3141.592654, 1000000.0, '3.1e3'), + (31415.92654, 1000000.0, '3.1e4'), + (314159.2654, 1000000.0, '3.1e5'), + (1e-05, 1000000.0, '1e-5'), + (0.0001, 1000000.0, '1e-4'), + (0.001, 1000000.0, '1e-3'), + (0.01, 1000000.0, '1e-2'), + (0.1, 1000000.0, '1e-1'), + (1, 1000000.0, '1'), + (10, 1000000.0, '10'), + (100, 1000000.0, '100'), + (1000, 1000000.0, '1000'), + (10000, 1000000.0, '1e4'), + (100000, 1000000.0, '1e5'), + ] + + @pytest.mark.parametrize('value, domain, expected', pprint_data) + def test_pprint(self, value, domain, expected): + fmt = mticker.LogFormatter() + label = fmt._pprint_val(value, domain) + assert label == expected + + @pytest.mark.parametrize('value, long, short', [ + (0.0, "0", "0"), + (0, "0", "0"), + (-1.0, "-10^0", "-1"), + (2e-10, "2x10^-10", "2e-10"), + (1e10, "10^10", "1e+10"), + ]) + def test_format_data(self, value, long, short): + fig, ax = plt.subplots() + ax.set_xscale('log') + fmt = ax.xaxis.get_major_formatter() + assert fmt.format_data(value) == long + assert fmt.format_data_short(value) == short + + def _sub_labels(self, axis, subs=()): + """Test whether locator marks subs to be labeled.""" + fmt = axis.get_minor_formatter() + minor_tlocs = axis.get_minorticklocs() + fmt.set_locs(minor_tlocs) + coefs = minor_tlocs / 10**(np.floor(np.log10(minor_tlocs))) + label_expected = [round(c) in subs for c in coefs] + label_test = [fmt(x) != '' for x in minor_tlocs] + assert label_test == label_expected + + @mpl.style.context('default') + def test_sublabel(self): + # test label locator + fig, ax = plt.subplots() + ax.set_xscale('log') + ax.xaxis.set_major_locator(mticker.LogLocator(base=10, subs=[])) + ax.xaxis.set_minor_locator(mticker.LogLocator(base=10, + subs=np.arange(2, 10))) + ax.xaxis.set_major_formatter(mticker.LogFormatter(labelOnlyBase=True)) + ax.xaxis.set_minor_formatter(mticker.LogFormatter(labelOnlyBase=False)) + # axis range above 3 decades, only bases are labeled + ax.set_xlim(1, 1e4) + fmt = ax.xaxis.get_major_formatter() + fmt.set_locs(ax.xaxis.get_majorticklocs()) + show_major_labels = [fmt(x) != '' + for x in ax.xaxis.get_majorticklocs()] + assert np.all(show_major_labels) + self._sub_labels(ax.xaxis, subs=[]) + + # For the next two, if the numdec threshold in LogFormatter.set_locs + # were 3, then the label sub would be 3 for 2-3 decades and (2, 5) + # for 1-2 decades. With a threshold of 1, subs are not labeled. + # axis range at 2 to 3 decades + ax.set_xlim(1, 800) + self._sub_labels(ax.xaxis, subs=[]) + + # axis range at 1 to 2 decades + ax.set_xlim(1, 80) + self._sub_labels(ax.xaxis, subs=[]) + + # axis range at 0.4 to 1 decades, label subs 2, 3, 4, 6 + ax.set_xlim(1, 8) + self._sub_labels(ax.xaxis, subs=[2, 3, 4, 6]) + + # axis range at 0 to 0.4 decades, label all + ax.set_xlim(0.5, 0.9) + self._sub_labels(ax.xaxis, subs=np.arange(2, 10, dtype=int)) + + @pytest.mark.parametrize('val', [1, 10, 100, 1000]) + def test_LogFormatter_call(self, val): + # test _num_to_string method used in __call__ + temp_lf = mticker.LogFormatter() + temp_lf.create_dummy_axis() + temp_lf.axis.set_view_interval(1, 10) + assert temp_lf(val) == str(val) + + @pytest.mark.parametrize('val', [1e-323, 2e-323, 10e-323, 11e-323]) + def test_LogFormatter_call_tiny(self, val): + # test coeff computation in __call__ + temp_lf = mticker.LogFormatter() + temp_lf.create_dummy_axis() + temp_lf.axis.set_view_interval(1, 10) + temp_lf(val) + + +class TestLogitFormatter: + @staticmethod + def logit_deformatter(string): + r""" + Parser to convert string as r'$\mathdefault{1.41\cdot10^{-4}}$' in + float 1.41e-4, as '0.5' or as r'$\mathdefault{\frac{1}{2}}$' in float + 0.5, + """ + match = re.match( + r"[^\d]*" + r"(?P1-)?" + r"(?P\d*\.?\d*)?" + r"(?:\\cdot)?" + r"(?:10\^\{(?P-?\d*)})?" + r"[^\d]*$", + string, + ) + if match: + comp = match["comp"] is not None + mantissa = float(match["mant"]) if match["mant"] else 1 + expo = int(match["expo"]) if match["expo"] is not None else 0 + value = mantissa * 10 ** expo + if match["mant"] or match["expo"] is not None: + if comp: + return 1 - value + return value + match = re.match( + r"[^\d]*\\frac\{(?P\d+)\}\{(?P\d+)\}[^\d]*$", string + ) + if match: + num, deno = float(match["num"]), float(match["deno"]) + return num / deno + raise ValueError("Not formatted by LogitFormatter") + + @pytest.mark.parametrize( + "fx, x", + [ + (r"STUFF0.41OTHERSTUFF", 0.41), + (r"STUFF1.41\cdot10^{-2}OTHERSTUFF", 1.41e-2), + (r"STUFF1-0.41OTHERSTUFF", 1 - 0.41), + (r"STUFF1-1.41\cdot10^{-2}OTHERSTUFF", 1 - 1.41e-2), + (r"STUFF", None), + (r"STUFF12.4e-3OTHERSTUFF", None), + ], + ) + def test_logit_deformater(self, fx, x): + if x is None: + with pytest.raises(ValueError): + TestLogitFormatter.logit_deformatter(fx) + else: + y = TestLogitFormatter.logit_deformatter(fx) + assert _LogitHelper.isclose(x, y) + + decade_test = sorted( + [10 ** (-i) for i in range(1, 10)] + + [1 - 10 ** (-i) for i in range(1, 10)] + + [1 / 2] + ) + + @pytest.mark.parametrize("x", decade_test) + def test_basic(self, x): + """ + Test the formatted value correspond to the value for ideal ticks in + logit space. + """ + formatter = mticker.LogitFormatter(use_overline=False) + formatter.set_locs(self.decade_test) + s = formatter(x) + x2 = TestLogitFormatter.logit_deformatter(s) + assert _LogitHelper.isclose(x, x2) + + @pytest.mark.parametrize("x", (-1, -0.5, -0.1, 1.1, 1.5, 2)) + def test_invalid(self, x): + """ + Test that invalid value are formatted with empty string without + raising exception. + """ + formatter = mticker.LogitFormatter(use_overline=False) + formatter.set_locs(self.decade_test) + s = formatter(x) + assert s == "" + + @pytest.mark.parametrize("x", 1 / (1 + np.exp(-np.linspace(-7, 7, 10)))) + def test_variablelength(self, x): + """ + The format length should change depending on the neighbor labels. + """ + formatter = mticker.LogitFormatter(use_overline=False) + for N in (10, 20, 50, 100, 200, 1000, 2000, 5000, 10000): + if x + 1 / N < 1: + formatter.set_locs([x - 1 / N, x, x + 1 / N]) + sx = formatter(x) + sx1 = formatter(x + 1 / N) + d = ( + TestLogitFormatter.logit_deformatter(sx1) + - TestLogitFormatter.logit_deformatter(sx) + ) + assert 0 < d < 2 / N + + lims_minor_major = [ + (True, (5e-8, 1 - 5e-8), ((25, False), (75, False))), + (True, (5e-5, 1 - 5e-5), ((25, False), (75, True))), + (True, (5e-2, 1 - 5e-2), ((25, True), (75, True))), + (False, (0.75, 0.76, 0.77), ((7, True), (25, True), (75, True))), + ] + + @pytest.mark.parametrize("method, lims, cases", lims_minor_major) + def test_minor_vs_major(self, method, lims, cases): + """ + Test minor/major displays. + """ + + if method: + min_loc = mticker.LogitLocator(minor=True) + ticks = min_loc.tick_values(*lims) + else: + ticks = np.array(lims) + min_form = mticker.LogitFormatter(minor=True) + for threshold, has_minor in cases: + min_form.set_minor_threshold(threshold) + formatted = min_form.format_ticks(ticks) + labelled = [f for f in formatted if len(f) > 0] + if has_minor: + assert len(labelled) > 0, (threshold, has_minor) + else: + assert len(labelled) == 0, (threshold, has_minor) + + def test_minor_number(self): + """ + Test the parameter minor_number + """ + min_loc = mticker.LogitLocator(minor=True) + min_form = mticker.LogitFormatter(minor=True) + ticks = min_loc.tick_values(5e-2, 1 - 5e-2) + for minor_number in (2, 4, 8, 16): + min_form.set_minor_number(minor_number) + formatted = min_form.format_ticks(ticks) + labelled = [f for f in formatted if len(f) > 0] + assert len(labelled) == minor_number + + def test_use_overline(self): + """ + Test the parameter use_overline + """ + x = 1 - 1e-2 + fx1 = r"$\mathdefault{1-10^{-2}}$" + fx2 = r"$\mathdefault{\overline{10^{-2}}}$" + form = mticker.LogitFormatter(use_overline=False) + assert form(x) == fx1 + form.use_overline(True) + assert form(x) == fx2 + form.use_overline(False) + assert form(x) == fx1 + + def test_one_half(self): + """ + Test the parameter one_half + """ + form = mticker.LogitFormatter() + assert r"\frac{1}{2}" in form(1/2) + form.set_one_half("1/2") + assert "1/2" in form(1/2) + form.set_one_half("one half") + assert "one half" in form(1/2) + + @pytest.mark.parametrize("N", (100, 253, 754)) + def test_format_data_short(self, N): + locs = np.linspace(0, 1, N)[1:-1] + form = mticker.LogitFormatter() + for x in locs: + fx = form.format_data_short(x) + if fx.startswith("1-"): + x2 = 1 - float(fx[2:]) + else: + x2 = float(fx) + assert abs(x - x2) < 1 / N + + +class TestFormatStrFormatter: + def test_basic(self): + # test % style formatter + tmp_form = mticker.FormatStrFormatter('%05d') + assert '00002' == tmp_form(2) + + +class TestStrMethodFormatter: + test_data = [ + ('{x:05d}', (2,), False, '00002'), + ('{x:05d}', (2,), True, '00002'), + ('{x:05d}', (-2,), False, '-0002'), + ('{x:05d}', (-2,), True, '\N{MINUS SIGN}0002'), + ('{x:03d}-{pos:02d}', (2, 1), False, '002-01'), + ('{x:03d}-{pos:02d}', (2, 1), True, '002-01'), + ('{x:03d}-{pos:02d}', (-2, 1), False, '-02-01'), + ('{x:03d}-{pos:02d}', (-2, 1), True, '\N{MINUS SIGN}02-01'), + ] + + @pytest.mark.parametrize('format, input, unicode_minus, expected', test_data) + def test_basic(self, format, input, unicode_minus, expected): + with mpl.rc_context({"axes.unicode_minus": unicode_minus}): + fmt = mticker.StrMethodFormatter(format) + assert fmt(*input) == expected + + +class TestEngFormatter: + # (unicode_minus, input, expected) where ''expected'' corresponds to the + # outputs respectively returned when (places=None, places=0, places=2) + # unicode_minus is a boolean value for the rcParam['axes.unicode_minus'] + raw_format_data = [ + (False, -1234.56789, ('-1.23457 k', '-1 k', '-1.23 k')), + (True, -1234.56789, ('\N{MINUS SIGN}1.23457 k', '\N{MINUS SIGN}1 k', + '\N{MINUS SIGN}1.23 k')), + (False, -1.23456789, ('-1.23457', '-1', '-1.23')), + (True, -1.23456789, ('\N{MINUS SIGN}1.23457', '\N{MINUS SIGN}1', + '\N{MINUS SIGN}1.23')), + (False, -0.123456789, ('-123.457 m', '-123 m', '-123.46 m')), + (True, -0.123456789, ('\N{MINUS SIGN}123.457 m', '\N{MINUS SIGN}123 m', + '\N{MINUS SIGN}123.46 m')), + (False, -0.00123456789, ('-1.23457 m', '-1 m', '-1.23 m')), + (True, -0.00123456789, ('\N{MINUS SIGN}1.23457 m', '\N{MINUS SIGN}1 m', + '\N{MINUS SIGN}1.23 m')), + (True, -0.0, ('0', '0', '0.00')), + (True, -0, ('0', '0', '0.00')), + (True, 0, ('0', '0', '0.00')), + (True, 1.23456789e-6, ('1.23457 µ', '1 µ', '1.23 µ')), + (True, 0.123456789, ('123.457 m', '123 m', '123.46 m')), + (True, 0.1, ('100 m', '100 m', '100.00 m')), + (True, 1, ('1', '1', '1.00')), + (True, 1.23456789, ('1.23457', '1', '1.23')), + # places=0: corner-case rounding + (True, 999.9, ('999.9', '1 k', '999.90')), + # corner-case rounding for all + (True, 999.9999, ('1 k', '1 k', '1.00 k')), + # negative corner-case + (False, -999.9999, ('-1 k', '-1 k', '-1.00 k')), + (True, -999.9999, ('\N{MINUS SIGN}1 k', '\N{MINUS SIGN}1 k', + '\N{MINUS SIGN}1.00 k')), + (True, 1000, ('1 k', '1 k', '1.00 k')), + (True, 1001, ('1.001 k', '1 k', '1.00 k')), + (True, 100001, ('100.001 k', '100 k', '100.00 k')), + (True, 987654.321, ('987.654 k', '988 k', '987.65 k')), + # OoR value (> 1000 Q) + (True, 1.23e33, ('1230 Q', '1230 Q', '1230.00 Q')) + ] + + @pytest.mark.parametrize('unicode_minus, input, expected', raw_format_data) + def test_params(self, unicode_minus, input, expected): + """ + Test the formatting of EngFormatter for various values of the 'places' + argument, in several cases: + + 0. without a unit symbol but with a (default) space separator; + 1. with both a unit symbol and a (default) space separator; + 2. with both a unit symbol and some non default separators; + 3. without a unit symbol but with some non default separators. + + Note that cases 2. and 3. are looped over several separator strings. + """ + + plt.rcParams['axes.unicode_minus'] = unicode_minus + UNIT = 's' # seconds + DIGITS = '0123456789' # %timeit showed 10-20% faster search than set + + # Case 0: unit='' (default) and sep=' ' (default). + # 'expected' already corresponds to this reference case. + exp_outputs = expected + formatters = ( + mticker.EngFormatter(), # places=None (default) + mticker.EngFormatter(places=0), + mticker.EngFormatter(places=2) + ) + for _formatter, _exp_output in zip(formatters, exp_outputs): + assert _formatter(input) == _exp_output + + # Case 1: unit=UNIT and sep=' ' (default). + # Append a unit symbol to the reference case. + # Beware of the values in [1, 1000), where there is no prefix! + exp_outputs = (_s + " " + UNIT if _s[-1] in DIGITS # case w/o prefix + else _s + UNIT for _s in expected) + formatters = ( + mticker.EngFormatter(unit=UNIT), # places=None (default) + mticker.EngFormatter(unit=UNIT, places=0), + mticker.EngFormatter(unit=UNIT, places=2) + ) + for _formatter, _exp_output in zip(formatters, exp_outputs): + assert _formatter(input) == _exp_output + + # Test several non default separators: no separator, a narrow + # no-break space (Unicode character) and an extravagant string. + for _sep in ("", "\N{NARROW NO-BREAK SPACE}", "@_@"): + # Case 2: unit=UNIT and sep=_sep. + # Replace the default space separator from the reference case + # with the tested one `_sep` and append a unit symbol to it. + exp_outputs = (_s + _sep + UNIT if _s[-1] in DIGITS # no prefix + else _s.replace(" ", _sep) + UNIT + for _s in expected) + formatters = ( + mticker.EngFormatter(unit=UNIT, sep=_sep), # places=None + mticker.EngFormatter(unit=UNIT, places=0, sep=_sep), + mticker.EngFormatter(unit=UNIT, places=2, sep=_sep) + ) + for _formatter, _exp_output in zip(formatters, exp_outputs): + assert _formatter(input) == _exp_output + + # Case 3: unit='' (default) and sep=_sep. + # Replace the default space separator from the reference case + # with the tested one `_sep`. Reference case is already unitless. + exp_outputs = (_s.replace(" ", _sep) for _s in expected) + formatters = ( + mticker.EngFormatter(sep=_sep), # places=None (default) + mticker.EngFormatter(places=0, sep=_sep), + mticker.EngFormatter(places=2, sep=_sep) + ) + for _formatter, _exp_output in zip(formatters, exp_outputs): + assert _formatter(input) == _exp_output + + +def test_engformatter_usetex_useMathText(): + fig, ax = plt.subplots() + ax.plot([0, 500, 1000], [0, 500, 1000]) + ax.set_xticks([0, 500, 1000]) + for formatter in (mticker.EngFormatter(usetex=True), + mticker.EngFormatter(useMathText=True)): + ax.xaxis.set_major_formatter(formatter) + fig.canvas.draw() + x_tick_label_text = [labl.get_text() for labl in ax.get_xticklabels()] + # Checking if the dollar `$` signs have been inserted around numbers + # in tick labels. + assert x_tick_label_text == ['$0$', '$500$', '$1$ k'] + + +@pytest.mark.parametrize( + 'data_offset, noise, oom_center_desired, oom_noise_desired', [ + (271_490_000_000.0, 10, 9, 0), + (27_149_000_000_000.0, 10_000_000, 12, 6), + (27.149, 0.01, 0, -3), + (2_714.9, 0.01, 3, -3), + (271_490.0, 0.001, 3, -3), + (271.49, 0.001, 0, -3), + # The following sets of parameters demonstrates that when + # oom(data_offset)-1 and oom(noise)-2 equal a standard 3*N oom, we get + # that oom_noise_desired < oom(noise) + (27_149_000_000.0, 100, 9, +3), + (27.149, 1e-07, 0, -6), + (271.49, 0.0001, 0, -3), + (27.149, 0.0001, 0, -3), + # Tests where oom(data_offset) <= oom(noise), those are probably + # covered by the part where formatter.offset != 0 + (27_149.0, 10_000, 0, 3), + (27.149, 10_000, 0, 3), + (27.149, 1_000, 0, 3), + (27.149, 100, 0, 0), + (27.149, 10, 0, 0), + ] +) +def test_engformatter_offset_oom( + data_offset, + noise, + oom_center_desired, + oom_noise_desired +): + UNIT = "eV" + fig, ax = plt.subplots() + ydata = data_offset + np.arange(-5, 7, dtype=float)*noise + ax.plot(ydata) + formatter = mticker.EngFormatter(useOffset=True, unit=UNIT) + # So that offset strings will always have the same size + formatter.ENG_PREFIXES[0] = "_" + ax.yaxis.set_major_formatter(formatter) + fig.canvas.draw() + offset_got = formatter.get_offset() + ticks_got = [labl.get_text() for labl in ax.get_yticklabels()] + # Predicting whether offset should be 0 or not is essentially testing + # ScalarFormatter._compute_offset . This function is pretty complex and it + # would be nice to test it, but this is out of scope for this test which + # only makes sure that offset text and the ticks gets the correct unit + # prefixes and the ticks. + if formatter.offset: + prefix_noise_got = offset_got[2] + prefix_noise_desired = formatter.ENG_PREFIXES[oom_noise_desired] + prefix_center_got = offset_got[-1-len(UNIT)] + prefix_center_desired = formatter.ENG_PREFIXES[oom_center_desired] + assert prefix_noise_desired == prefix_noise_got + assert prefix_center_desired == prefix_center_got + # Make sure the ticks didn't get the UNIT + for tick in ticks_got: + assert UNIT not in tick + else: + assert oom_center_desired == 0 + assert offset_got == "" + # Make sure the ticks contain now the prefixes + for tick in ticks_got: + # 0 is zero on all orders of magnitudes, no matter what is + # oom_noise_desired + prefix_idx = 0 if tick[0] == "0" else oom_noise_desired + assert tick.endswith(formatter.ENG_PREFIXES[prefix_idx] + UNIT) + + +class TestPercentFormatter: + percent_data = [ + # Check explicitly set decimals over different intervals and values + (100, 0, '%', 120, 100, '120%'), + (100, 0, '%', 100, 90, '100%'), + (100, 0, '%', 90, 50, '90%'), + (100, 0, '%', -1.7, 40, '-2%'), + (100, 1, '%', 90.0, 100, '90.0%'), + (100, 1, '%', 80.1, 90, '80.1%'), + (100, 1, '%', 70.23, 50, '70.2%'), + # 60.554 instead of 60.55: see https://bugs.python.org/issue5118 + (100, 1, '%', -60.554, 40, '-60.6%'), + # Check auto decimals over different intervals and values + (100, None, '%', 95, 1, '95.00%'), + (1.0, None, '%', 3, 6, '300%'), + (17.0, None, '%', 1, 8.5, '6%'), + (17.0, None, '%', 1, 8.4, '5.9%'), + (5, None, '%', -100, 0.000001, '-2000.00000%'), + # Check percent symbol + (1.0, 2, None, 1.2, 100, '120.00'), + (75, 3, '', 50, 100, '66.667'), + (42, None, '^^Foobar$$', 21, 12, '50.0^^Foobar$$'), + ] + + percent_ids = [ + # Check explicitly set decimals over different intervals and values + 'decimals=0, x>100%', + 'decimals=0, x=100%', + 'decimals=0, x<100%', + 'decimals=0, x<0%', + 'decimals=1, x>100%', + 'decimals=1, x=100%', + 'decimals=1, x<100%', + 'decimals=1, x<0%', + # Check auto decimals over different intervals and values + 'autodecimal, x<100%, display_range=1', + 'autodecimal, x>100%, display_range=6 (custom xmax test)', + 'autodecimal, x<100%, display_range=8.5 (autodecimal test 1)', + 'autodecimal, x<100%, display_range=8.4 (autodecimal test 2)', + 'autodecimal, x<-100%, display_range=1e-6 (tiny display range)', + # Check percent symbol + 'None as percent symbol', + 'Empty percent symbol', + 'Custom percent symbol', + ] + + latex_data = [ + (False, False, r'50\{t}%'), + (False, True, r'50\\\{t\}\%'), + (True, False, r'50\{t}%'), + (True, True, r'50\{t}%'), + ] + + @pytest.mark.parametrize( + 'xmax, decimals, symbol, x, display_range, expected', + percent_data, ids=percent_ids) + def test_basic(self, xmax, decimals, symbol, + x, display_range, expected): + formatter = mticker.PercentFormatter(xmax, decimals, symbol) + with mpl.rc_context(rc={'text.usetex': False}): + assert formatter.format_pct(x, display_range) == expected + + @pytest.mark.parametrize('is_latex, usetex, expected', latex_data) + def test_latex(self, is_latex, usetex, expected): + fmt = mticker.PercentFormatter(symbol='\\{t}%', is_latex=is_latex) + with mpl.rc_context(rc={'text.usetex': usetex}): + assert fmt.format_pct(50, 100) == expected + + +def _impl_locale_comma(): + try: + locale.setlocale(locale.LC_ALL, 'de_DE.UTF-8') + except locale.Error: + print('SKIP: Locale de_DE.UTF-8 is not supported on this machine') + return + ticks = mticker.ScalarFormatter(useMathText=True, useLocale=True) + fmt = '$\\mathdefault{%1.1f}$' + x = ticks._format_maybe_minus_and_locale(fmt, 0.5) + assert x == '$\\mathdefault{0{,}5}$' + # Do not change , in the format string + fmt = ',$\\mathdefault{,%1.1f},$' + x = ticks._format_maybe_minus_and_locale(fmt, 0.5) + assert x == ',$\\mathdefault{,0{,}5},$' + # Make sure no brackets are added if not using math text + ticks = mticker.ScalarFormatter(useMathText=False, useLocale=True) + fmt = '%1.1f' + x = ticks._format_maybe_minus_and_locale(fmt, 0.5) + assert x == '0,5' + + +def test_locale_comma(): + # On some systems/pytest versions, `pytest.skip` in an exception handler + # does not skip, but is treated as an exception, so directly running this + # test can incorrectly fail instead of skip. + # Instead, run this test in a subprocess, which avoids the problem, and the + # need to fix the locale after. + proc = mpl.testing.subprocess_run_helper(_impl_locale_comma, timeout=60, + extra_env={'MPLBACKEND': 'Agg'}) + skip_msg = next((line[len('SKIP:'):].strip() + for line in proc.stdout.splitlines() + if line.startswith('SKIP:')), + '') + if skip_msg: + pytest.skip(skip_msg) + + +def test_majformatter_type(): + fig, ax = plt.subplots() + with pytest.raises(TypeError): + ax.xaxis.set_major_formatter(mticker.LogLocator()) + + +def test_minformatter_type(): + fig, ax = plt.subplots() + with pytest.raises(TypeError): + ax.xaxis.set_minor_formatter(mticker.LogLocator()) + + +def test_majlocator_type(): + fig, ax = plt.subplots() + with pytest.raises(TypeError): + ax.xaxis.set_major_locator(mticker.LogFormatter()) + + +def test_minlocator_type(): + fig, ax = plt.subplots() + with pytest.raises(TypeError): + ax.xaxis.set_minor_locator(mticker.LogFormatter()) + + +def test_minorticks_rc(): + fig = plt.figure() + + def minorticksubplot(xminor, yminor, i): + rc = {'xtick.minor.visible': xminor, + 'ytick.minor.visible': yminor} + with plt.rc_context(rc=rc): + ax = fig.add_subplot(2, 2, i) + + assert (len(ax.xaxis.get_minor_ticks()) > 0) == xminor + assert (len(ax.yaxis.get_minor_ticks()) > 0) == yminor + + minorticksubplot(False, False, 1) + minorticksubplot(True, False, 2) + minorticksubplot(False, True, 3) + minorticksubplot(True, True, 4) + + +def test_minorticks_toggle(): + """ + Test toggling minor ticks + + Test `.Axis.minorticks_on()` and `.Axis.minorticks_off()`. Testing is + limited to a subset of built-in scales - `'linear'`, `'log'`, `'asinh'` + and `'logit'`. `symlog` scale does not seem to have a working minor + locator and is omitted. In future, this test should cover all scales in + `matplotlib.scale.get_scale_names()`. + """ + fig = plt.figure() + def minortickstoggle(xminor, yminor, scale, i): + ax = fig.add_subplot(2, 2, i) + ax.set_xscale(scale) + ax.set_yscale(scale) + if not xminor and not yminor: + ax.minorticks_off() + if xminor and not yminor: + ax.xaxis.minorticks_on() + ax.yaxis.minorticks_off() + if not xminor and yminor: + ax.xaxis.minorticks_off() + ax.yaxis.minorticks_on() + if xminor and yminor: + ax.minorticks_on() + + assert (len(ax.xaxis.get_minor_ticks()) > 0) == xminor + assert (len(ax.yaxis.get_minor_ticks()) > 0) == yminor + + scales = ['linear', 'log', 'asinh', 'logit'] + for scale in scales: + minortickstoggle(False, False, scale, 1) + minortickstoggle(True, False, scale, 2) + minortickstoggle(False, True, scale, 3) + minortickstoggle(True, True, scale, 4) + fig.clear() + + plt.close(fig) + + +@pytest.mark.parametrize('remove_overlapping_locs, expected_num', + ((True, 6), + (None, 6), # this tests the default + (False, 9))) +def test_remove_overlap(remove_overlapping_locs, expected_num): + t = np.arange("2018-11-03", "2018-11-06", dtype="datetime64") + x = np.ones(len(t)) + + fig, ax = plt.subplots() + ax.plot(t, x) + + ax.xaxis.set_major_locator(mpl.dates.DayLocator()) + ax.xaxis.set_major_formatter(mpl.dates.DateFormatter('\n%a')) + + ax.xaxis.set_minor_locator(mpl.dates.HourLocator((0, 6, 12, 18))) + ax.xaxis.set_minor_formatter(mpl.dates.DateFormatter('%H:%M')) + # force there to be extra ticks + ax.xaxis.get_minor_ticks(15) + if remove_overlapping_locs is not None: + ax.xaxis.remove_overlapping_locs = remove_overlapping_locs + + # check that getter/setter exists + current = ax.xaxis.remove_overlapping_locs + assert (current == ax.xaxis.get_remove_overlapping_locs()) + plt.setp(ax.xaxis, remove_overlapping_locs=current) + new = ax.xaxis.remove_overlapping_locs + assert (new == ax.xaxis.remove_overlapping_locs) + + # check that the accessors filter correctly + # this is the method that does the actual filtering + assert len(ax.xaxis.get_minorticklocs()) == expected_num + # these three are derivative + assert len(ax.xaxis.get_minor_ticks()) == expected_num + assert len(ax.xaxis.get_minorticklabels()) == expected_num + assert len(ax.xaxis.get_minorticklines()) == expected_num*2 + + +@pytest.mark.parametrize('sub', [ + ['hi', 'aardvark'], + np.zeros((2, 2))]) +def test_bad_locator_subs(sub): + ll = mticker.LogLocator() + with pytest.raises(ValueError): + ll.set_params(subs=sub) + + +@pytest.mark.parametrize('numticks', [1, 2, 3, 9]) +@mpl.style.context('default') +def test_small_range_loglocator(numticks): + ll = mticker.LogLocator() + ll.set_params(numticks=numticks) + for top in [5, 7, 9, 11, 15, 50, 100, 1000]: + ticks = ll.tick_values(.5, top) + assert (np.diff(np.log10(ll.tick_values(6, 150))) == 1).all() + + +def test_NullFormatter(): + formatter = mticker.NullFormatter() + assert formatter(1.0) == '' + assert formatter.format_data(1.0) == '' + assert formatter.format_data_short(1.0) == '' + + +@pytest.mark.parametrize('formatter', ( + mticker.FuncFormatter(lambda a: f'val: {a}'), + mticker.FixedFormatter(('foo', 'bar')))) +def test_set_offset_string(formatter): + assert formatter.get_offset() == '' + formatter.set_offset_string('mpl') + assert formatter.get_offset() == 'mpl' + + +def test_minorticks_on_multi_fig(): + """ + Turning on minor gridlines in a multi-Axes Figure + that contains more than one boxplot and shares the x-axis + should not raise an exception. + """ + fig, ax = plt.subplots() + + ax.boxplot(np.arange(10), positions=[0]) + ax.boxplot(np.arange(10), positions=[0]) + ax.boxplot(np.arange(10), positions=[1]) + + ax.grid(which="major") + ax.grid(which="minor") + ax.minorticks_on() + fig.draw_without_rendering() + + assert ax.get_xgridlines() + assert isinstance(ax.xaxis.get_minor_locator(), mpl.ticker.AutoMinorLocator) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_tightlayout.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_tightlayout.py new file mode 100644 index 0000000000000000000000000000000000000000..f6b6d8f644cc50537a7466fa7548a97833033e48 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_tightlayout.py @@ -0,0 +1,430 @@ +import warnings + +import numpy as np +from numpy.testing import assert_array_equal +import pytest + +import matplotlib as mpl +from matplotlib.testing.decorators import image_comparison +import matplotlib.pyplot as plt +from matplotlib.offsetbox import AnchoredOffsetbox, DrawingArea +from matplotlib.patches import Rectangle + + +pytestmark = [ + pytest.mark.usefixtures('text_placeholders') +] + + +def example_plot(ax, fontsize=12): + ax.plot([1, 2]) + ax.locator_params(nbins=3) + ax.set_xlabel('x-label', fontsize=fontsize) + ax.set_ylabel('y-label', fontsize=fontsize) + ax.set_title('Title', fontsize=fontsize) + + +@image_comparison(['tight_layout1'], style='mpl20') +def test_tight_layout1(): + """Test tight_layout for a single subplot.""" + fig, ax = plt.subplots() + example_plot(ax, fontsize=24) + plt.tight_layout() + + +@image_comparison(['tight_layout2'], style='mpl20') +def test_tight_layout2(): + """Test tight_layout for multiple subplots.""" + fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2) + example_plot(ax1) + example_plot(ax2) + example_plot(ax3) + example_plot(ax4) + plt.tight_layout() + + +@image_comparison(['tight_layout3'], style='mpl20') +def test_tight_layout3(): + """Test tight_layout for multiple subplots.""" + ax1 = plt.subplot(221) + ax2 = plt.subplot(223) + ax3 = plt.subplot(122) + example_plot(ax1) + example_plot(ax2) + example_plot(ax3) + plt.tight_layout() + + +@image_comparison(['tight_layout4'], style='mpl20') +def test_tight_layout4(): + """Test tight_layout for subplot2grid.""" + ax1 = plt.subplot2grid((3, 3), (0, 0)) + ax2 = plt.subplot2grid((3, 3), (0, 1), colspan=2) + ax3 = plt.subplot2grid((3, 3), (1, 0), colspan=2, rowspan=2) + ax4 = plt.subplot2grid((3, 3), (1, 2), rowspan=2) + example_plot(ax1) + example_plot(ax2) + example_plot(ax3) + example_plot(ax4) + plt.tight_layout() + + +@image_comparison(['tight_layout5'], style='mpl20') +def test_tight_layout5(): + """Test tight_layout for image.""" + ax = plt.subplot() + arr = np.arange(100).reshape((10, 10)) + ax.imshow(arr, interpolation="none") + plt.tight_layout() + + +@image_comparison(['tight_layout6'], style='mpl20') +def test_tight_layout6(): + """Test tight_layout for gridspec.""" + + # This raises warnings since tight layout cannot + # do this fully automatically. But the test is + # correct since the layout is manually edited + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + fig = plt.figure() + + gs1 = mpl.gridspec.GridSpec(2, 1) + ax1 = fig.add_subplot(gs1[0]) + ax2 = fig.add_subplot(gs1[1]) + + example_plot(ax1) + example_plot(ax2) + + gs1.tight_layout(fig, rect=[0, 0, 0.5, 1]) + + gs2 = mpl.gridspec.GridSpec(3, 1) + + for ss in gs2: + ax = fig.add_subplot(ss) + example_plot(ax) + ax.set_title("") + ax.set_xlabel("") + + ax.set_xlabel("x-label", fontsize=12) + + gs2.tight_layout(fig, rect=[0.5, 0, 1, 1], h_pad=0.45) + + top = min(gs1.top, gs2.top) + bottom = max(gs1.bottom, gs2.bottom) + + gs1.tight_layout(fig, rect=[None, 0 + (bottom-gs1.bottom), + 0.5, 1 - (gs1.top-top)]) + gs2.tight_layout(fig, rect=[0.5, 0 + (bottom-gs2.bottom), + None, 1 - (gs2.top-top)], + h_pad=0.45) + + +@image_comparison(['tight_layout7'], style='mpl20') +def test_tight_layout7(): + # tight layout with left and right titles + fontsize = 24 + fig, ax = plt.subplots() + ax.plot([1, 2]) + ax.locator_params(nbins=3) + ax.set_xlabel('x-label', fontsize=fontsize) + ax.set_ylabel('y-label', fontsize=fontsize) + ax.set_title('Left Title', loc='left', fontsize=fontsize) + ax.set_title('Right Title', loc='right', fontsize=fontsize) + plt.tight_layout() + + +@image_comparison(['tight_layout8'], style='mpl20', tol=0.005) +def test_tight_layout8(): + """Test automatic use of tight_layout.""" + fig = plt.figure() + fig.set_layout_engine(layout='tight', pad=0.1) + ax = fig.add_subplot() + example_plot(ax, fontsize=24) + fig.draw_without_rendering() + + +@image_comparison(['tight_layout9'], style='mpl20') +def test_tight_layout9(): + # Test tight_layout for non-visible subplots + # GH 8244 + f, axarr = plt.subplots(2, 2) + axarr[1][1].set_visible(False) + plt.tight_layout() + + +def test_outward_ticks(): + """Test automatic use of tight_layout.""" + fig = plt.figure() + ax = fig.add_subplot(221) + ax.xaxis.set_tick_params(tickdir='out', length=16, width=3) + ax.yaxis.set_tick_params(tickdir='out', length=16, width=3) + ax.xaxis.set_tick_params( + tickdir='out', length=32, width=3, tick1On=True, which='minor') + ax.yaxis.set_tick_params( + tickdir='out', length=32, width=3, tick1On=True, which='minor') + ax.xaxis.set_ticks([0], minor=True) + ax.yaxis.set_ticks([0], minor=True) + ax = fig.add_subplot(222) + ax.xaxis.set_tick_params(tickdir='in', length=32, width=3) + ax.yaxis.set_tick_params(tickdir='in', length=32, width=3) + ax = fig.add_subplot(223) + ax.xaxis.set_tick_params(tickdir='inout', length=32, width=3) + ax.yaxis.set_tick_params(tickdir='inout', length=32, width=3) + ax = fig.add_subplot(224) + ax.xaxis.set_tick_params(tickdir='out', length=32, width=3) + ax.yaxis.set_tick_params(tickdir='out', length=32, width=3) + plt.tight_layout() + # These values were obtained after visual checking that they correspond + # to a tight layouting that did take the ticks into account. + expected = [ + [[0.092, 0.605], [0.433, 0.933]], + [[0.581, 0.605], [0.922, 0.933]], + [[0.092, 0.138], [0.433, 0.466]], + [[0.581, 0.138], [0.922, 0.466]], + ] + for nn, ax in enumerate(fig.axes): + assert_array_equal(np.round(ax.get_position().get_points(), 3), + expected[nn]) + + +def add_offsetboxes(ax, size=10, margin=.1, color='black'): + """ + Surround ax with OffsetBoxes + """ + m, mp = margin, 1+margin + anchor_points = [(-m, -m), (-m, .5), (-m, mp), + (.5, mp), (mp, mp), (mp, .5), + (mp, -m), (.5, -m)] + for point in anchor_points: + da = DrawingArea(size, size) + background = Rectangle((0, 0), width=size, + height=size, + facecolor=color, + edgecolor='None', + linewidth=0, + antialiased=False) + da.add_artist(background) + + anchored_box = AnchoredOffsetbox( + loc='center', + child=da, + pad=0., + frameon=False, + bbox_to_anchor=point, + bbox_transform=ax.transAxes, + borderpad=0.) + ax.add_artist(anchored_box) + + +def test_tight_layout_offsetboxes(): + # 0. + # - Create 4 subplots + # - Plot a diagonal line on them + # - Use tight_layout + # + # 1. + # - Same 4 subplots + # - Surround each plot with 7 boxes + # - Use tight_layout + # - See that the squares are included in the tight_layout and that the squares do + # not overlap + # + # 2. + # - Make the squares around the Axes invisible + # - See that the invisible squares do not affect the tight_layout + rows = cols = 2 + colors = ['red', 'blue', 'green', 'yellow'] + x = y = [0, 1] + + def _subplots(with_boxes): + fig, axs = plt.subplots(rows, cols) + for ax, color in zip(axs.flat, colors): + ax.plot(x, y, color=color) + if with_boxes: + add_offsetboxes(ax, 20, color=color) + return fig, axs + + # 0. + fig0, axs0 = _subplots(False) + fig0.tight_layout() + + # 1. + fig1, axs1 = _subplots(True) + fig1.tight_layout() + + # The AnchoredOffsetbox should be added to the bounding of the Axes, causing them to + # be smaller than the plain figure. + for ax0, ax1 in zip(axs0.flat, axs1.flat): + bbox0 = ax0.get_position() + bbox1 = ax1.get_position() + assert bbox1.x0 > bbox0.x0 + assert bbox1.x1 < bbox0.x1 + assert bbox1.y0 > bbox0.y0 + assert bbox1.y1 < bbox0.y1 + + # No AnchoredOffsetbox should overlap with another. + bboxes = [] + for ax1 in axs1.flat: + for child in ax1.get_children(): + if not isinstance(child, AnchoredOffsetbox): + continue + bbox = child.get_window_extent() + for other_bbox in bboxes: + assert not bbox.overlaps(other_bbox) + bboxes.append(bbox) + + # 2. + fig2, axs2 = _subplots(True) + for ax in axs2.flat: + for child in ax.get_children(): + if isinstance(child, AnchoredOffsetbox): + child.set_visible(False) + fig2.tight_layout() + # The invisible AnchoredOffsetbox should not count for tight layout, so it should + # look the same as when they were never added. + for ax0, ax2 in zip(axs0.flat, axs2.flat): + bbox0 = ax0.get_position() + bbox2 = ax2.get_position() + assert_array_equal(bbox2.get_points(), bbox0.get_points()) + + +def test_empty_layout(): + """Test that tight layout doesn't cause an error when there are no Axes.""" + fig = plt.gcf() + fig.tight_layout() + + +@pytest.mark.parametrize("label", ["xlabel", "ylabel"]) +def test_verybig_decorators(label): + """Test that no warning emitted when xlabel/ylabel too big.""" + fig, ax = plt.subplots(figsize=(3, 2)) + ax.set(**{label: 'a' * 100}) + + +def test_big_decorators_horizontal(): + """Test that doesn't warn when xlabel too big.""" + fig, axs = plt.subplots(1, 2, figsize=(3, 2)) + axs[0].set_xlabel('a' * 30) + axs[1].set_xlabel('b' * 30) + + +def test_big_decorators_vertical(): + """Test that doesn't warn when ylabel too big.""" + fig, axs = plt.subplots(2, 1, figsize=(3, 2)) + axs[0].set_ylabel('a' * 20) + axs[1].set_ylabel('b' * 20) + + +def test_badsubplotgrid(): + # test that we get warning for mismatched subplot grids, not than an error + plt.subplot2grid((4, 5), (0, 0)) + # this is the bad entry: + plt.subplot2grid((5, 5), (0, 3), colspan=3, rowspan=5) + with pytest.warns(UserWarning): + plt.tight_layout() + + +def test_collapsed(): + # test that if the amount of space required to make all the axes + # decorations fit would mean that the actual Axes would end up with size + # zero (i.e. margins add up to more than the available width) that a call + # to tight_layout will not get applied: + fig, ax = plt.subplots(tight_layout=True) + ax.set_xlim([0, 1]) + ax.set_ylim([0, 1]) + + ax.annotate('BIG LONG STRING', xy=(1.25, 2), xytext=(10.5, 1.75), + annotation_clip=False) + p1 = ax.get_position() + with pytest.warns(UserWarning): + plt.tight_layout() + p2 = ax.get_position() + assert p1.width == p2.width + # test that passing a rect doesn't crash... + with pytest.warns(UserWarning): + plt.tight_layout(rect=[0, 0, 0.8, 0.8]) + + +def test_suptitle(): + fig, ax = plt.subplots(tight_layout=True) + st = fig.suptitle("foo") + t = ax.set_title("bar") + fig.canvas.draw() + assert st.get_window_extent().y0 > t.get_window_extent().y1 + + +@pytest.mark.backend("pdf") +def test_non_agg_renderer(monkeypatch, recwarn): + unpatched_init = mpl.backend_bases.RendererBase.__init__ + + def __init__(self, *args, **kwargs): + # Check that we don't instantiate any other renderer than a pdf + # renderer to perform pdf tight layout. + assert isinstance(self, mpl.backends.backend_pdf.RendererPdf) + unpatched_init(self, *args, **kwargs) + + monkeypatch.setattr(mpl.backend_bases.RendererBase, "__init__", __init__) + fig, ax = plt.subplots() + fig.tight_layout() + + +def test_manual_colorbar(): + # This should warn, but not raise + fig, axes = plt.subplots(1, 2) + pts = axes[1].scatter([0, 1], [0, 1], c=[1, 5]) + ax_rect = axes[1].get_position() + cax = fig.add_axes( + [ax_rect.x1 + 0.005, ax_rect.y0, 0.015, ax_rect.height] + ) + fig.colorbar(pts, cax=cax) + with pytest.warns(UserWarning, match="This figure includes Axes"): + fig.tight_layout() + + +def test_clipped_to_axes(): + # Ensure that _fully_clipped_to_axes() returns True under default + # conditions for all projection types. Axes.get_tightbbox() + # uses this to skip artists in layout calculations. + arr = np.arange(100).reshape((10, 10)) + fig = plt.figure(figsize=(6, 2)) + ax1 = fig.add_subplot(131, projection='rectilinear') + ax2 = fig.add_subplot(132, projection='mollweide') + ax3 = fig.add_subplot(133, projection='polar') + for ax in (ax1, ax2, ax3): + # Default conditions (clipped by ax.bbox or ax.patch) + ax.grid(False) + h, = ax.plot(arr[:, 0]) + m = ax.pcolor(arr) + assert h._fully_clipped_to_axes() + assert m._fully_clipped_to_axes() + # Non-default conditions (not clipped by ax.patch) + rect = Rectangle((0, 0), 0.5, 0.5, transform=ax.transAxes) + h.set_clip_path(rect) + m.set_clip_path(rect.get_path(), rect.get_transform()) + assert not h._fully_clipped_to_axes() + assert not m._fully_clipped_to_axes() + + +def test_tight_pads(): + fig, ax = plt.subplots() + with pytest.warns(PendingDeprecationWarning, + match='will be deprecated'): + fig.set_tight_layout({'pad': 0.15}) + fig.draw_without_rendering() + + +def test_tight_kwargs(): + fig, ax = plt.subplots(tight_layout={'pad': 0.15}) + fig.draw_without_rendering() + + +def test_tight_toggle(): + fig, ax = plt.subplots() + with pytest.warns(PendingDeprecationWarning): + fig.set_tight_layout(True) + assert fig.get_tight_layout() + fig.set_tight_layout(False) + assert not fig.get_tight_layout() + fig.set_tight_layout(True) + assert fig.get_tight_layout() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_transforms.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..96e78b6828f8d061ef61826678a7d5c91df8d884 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_transforms.py @@ -0,0 +1,1131 @@ +import copy + +import numpy as np +from numpy.testing import (assert_allclose, assert_almost_equal, + assert_array_equal, assert_array_almost_equal) +import pytest + +from matplotlib import scale +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +import matplotlib.transforms as mtransforms +from matplotlib.transforms import Affine2D, Bbox, TransformedBbox, _ScaledRotation +from matplotlib.path import Path +from matplotlib.testing.decorators import image_comparison, check_figures_equal +from unittest.mock import MagicMock + + +class TestAffine2D: + single_point = [1.0, 1.0] + multiple_points = [[0.0, 2.0], [3.0, 3.0], [4.0, 0.0]] + pivot = single_point + + def test_init(self): + Affine2D([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + Affine2D(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], int)) + Affine2D(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], float)) + + def test_values(self): + np.random.seed(19680801) + values = np.random.random(6) + assert_array_equal(Affine2D.from_values(*values).to_values(), values) + + def test_modify_inplace(self): + # Some polar transforms require modifying the matrix in place. + trans = Affine2D() + mtx = trans.get_matrix() + mtx[0, 0] = 42 + assert_array_equal(trans.get_matrix(), [[42, 0, 0], [0, 1, 0], [0, 0, 1]]) + + def test_clear(self): + a = Affine2D(np.random.rand(3, 3) + 5) # Anything non-identity. + a.clear() + assert_array_equal(a.get_matrix(), [[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + def test_rotate(self): + r_pi_2 = Affine2D().rotate(np.pi / 2) + r90 = Affine2D().rotate_deg(90) + assert_array_equal(r_pi_2.get_matrix(), r90.get_matrix()) + assert_array_almost_equal(r90.transform(self.single_point), [-1, 1]) + assert_array_almost_equal(r90.transform(self.multiple_points), + [[-2, 0], [-3, 3], [0, 4]]) + + r_pi = Affine2D().rotate(np.pi) + r180 = Affine2D().rotate_deg(180) + assert_array_equal(r_pi.get_matrix(), r180.get_matrix()) + assert_array_almost_equal(r180.transform(self.single_point), [-1, -1]) + assert_array_almost_equal(r180.transform(self.multiple_points), + [[0, -2], [-3, -3], [-4, 0]]) + + r_pi_3_2 = Affine2D().rotate(3 * np.pi / 2) + r270 = Affine2D().rotate_deg(270) + assert_array_equal(r_pi_3_2.get_matrix(), r270.get_matrix()) + assert_array_almost_equal(r270.transform(self.single_point), [1, -1]) + assert_array_almost_equal(r270.transform(self.multiple_points), + [[2, 0], [3, -3], [0, -4]]) + + assert_array_equal((r90 + r90).get_matrix(), r180.get_matrix()) + assert_array_equal((r90 + r180).get_matrix(), r270.get_matrix()) + + def test_rotate_around(self): + r_pi_2 = Affine2D().rotate_around(*self.pivot, np.pi / 2) + r90 = Affine2D().rotate_deg_around(*self.pivot, 90) + assert_array_equal(r_pi_2.get_matrix(), r90.get_matrix()) + assert_array_almost_equal(r90.transform(self.single_point), [1, 1]) + assert_array_almost_equal(r90.transform(self.multiple_points), + [[0, 0], [-1, 3], [2, 4]]) + + r_pi = Affine2D().rotate_around(*self.pivot, np.pi) + r180 = Affine2D().rotate_deg_around(*self.pivot, 180) + assert_array_equal(r_pi.get_matrix(), r180.get_matrix()) + assert_array_almost_equal(r180.transform(self.single_point), [1, 1]) + assert_array_almost_equal(r180.transform(self.multiple_points), + [[2, 0], [-1, -1], [-2, 2]]) + + r_pi_3_2 = Affine2D().rotate_around(*self.pivot, 3 * np.pi / 2) + r270 = Affine2D().rotate_deg_around(*self.pivot, 270) + assert_array_equal(r_pi_3_2.get_matrix(), r270.get_matrix()) + assert_array_almost_equal(r270.transform(self.single_point), [1, 1]) + assert_array_almost_equal(r270.transform(self.multiple_points), + [[2, 2], [3, -1], [0, -2]]) + + assert_array_almost_equal((r90 + r90).get_matrix(), r180.get_matrix()) + assert_array_almost_equal((r90 + r180).get_matrix(), r270.get_matrix()) + + def test_scale(self): + sx = Affine2D().scale(3, 1) + sy = Affine2D().scale(1, -2) + trans = Affine2D().scale(3, -2) + assert_array_equal((sx + sy).get_matrix(), trans.get_matrix()) + assert_array_equal(trans.transform(self.single_point), [3, -2]) + assert_array_equal(trans.transform(self.multiple_points), + [[0, -4], [9, -6], [12, 0]]) + + def test_skew(self): + trans_rad = Affine2D().skew(np.pi / 8, np.pi / 12) + trans_deg = Affine2D().skew_deg(22.5, 15) + assert_array_equal(trans_rad.get_matrix(), trans_deg.get_matrix()) + # Using ~atan(0.5), ~atan(0.25) produces roundish numbers on output. + trans = Affine2D().skew_deg(26.5650512, 14.0362435) + assert_array_almost_equal(trans.transform(self.single_point), [1.5, 1.25]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[1, 2], [4.5, 3.75], [4, 1]]) + + def test_translate(self): + tx = Affine2D().translate(23, 0) + ty = Affine2D().translate(0, 42) + trans = Affine2D().translate(23, 42) + assert_array_equal((tx + ty).get_matrix(), trans.get_matrix()) + assert_array_equal(trans.transform(self.single_point), [24, 43]) + assert_array_equal(trans.transform(self.multiple_points), + [[23, 44], [26, 45], [27, 42]]) + + def test_rotate_plus_other(self): + trans = Affine2D().rotate_deg(90).rotate_deg_around(*self.pivot, 180) + trans_added = (Affine2D().rotate_deg(90) + + Affine2D().rotate_deg_around(*self.pivot, 180)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [3, 1]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[4, 2], [5, -1], [2, -2]]) + + trans = Affine2D().rotate_deg(90).scale(3, -2) + trans_added = Affine2D().rotate_deg(90) + Affine2D().scale(3, -2) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [-3, -2]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[-6, -0], [-9, -6], [0, -8]]) + + trans = (Affine2D().rotate_deg(90) + .skew_deg(26.5650512, 14.0362435)) # ~atan(0.5), ~atan(0.25) + trans_added = (Affine2D().rotate_deg(90) + + Affine2D().skew_deg(26.5650512, 14.0362435)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [-0.5, 0.75]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[-2, -0.5], [-1.5, 2.25], [2, 4]]) + + trans = Affine2D().rotate_deg(90).translate(23, 42) + trans_added = Affine2D().rotate_deg(90) + Affine2D().translate(23, 42) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [22, 43]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[21, 42], [20, 45], [23, 46]]) + + def test_rotate_around_plus_other(self): + trans = Affine2D().rotate_deg_around(*self.pivot, 90).rotate_deg(180) + trans_added = (Affine2D().rotate_deg_around(*self.pivot, 90) + + Affine2D().rotate_deg(180)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [-1, -1]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[0, 0], [1, -3], [-2, -4]]) + + trans = Affine2D().rotate_deg_around(*self.pivot, 90).scale(3, -2) + trans_added = (Affine2D().rotate_deg_around(*self.pivot, 90) + + Affine2D().scale(3, -2)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [3, -2]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[0, 0], [-3, -6], [6, -8]]) + + trans = (Affine2D().rotate_deg_around(*self.pivot, 90) + .skew_deg(26.5650512, 14.0362435)) # ~atan(0.5), ~atan(0.25) + trans_added = (Affine2D().rotate_deg_around(*self.pivot, 90) + + Affine2D().skew_deg(26.5650512, 14.0362435)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [1.5, 1.25]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[0, 0], [0.5, 2.75], [4, 4.5]]) + + trans = Affine2D().rotate_deg_around(*self.pivot, 90).translate(23, 42) + trans_added = (Affine2D().rotate_deg_around(*self.pivot, 90) + + Affine2D().translate(23, 42)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [24, 43]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[23, 42], [22, 45], [25, 46]]) + + def test_scale_plus_other(self): + trans = Affine2D().scale(3, -2).rotate_deg(90) + trans_added = Affine2D().scale(3, -2) + Affine2D().rotate_deg(90) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_equal(trans.transform(self.single_point), [2, 3]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[4, 0], [6, 9], [0, 12]]) + + trans = Affine2D().scale(3, -2).rotate_deg_around(*self.pivot, 90) + trans_added = (Affine2D().scale(3, -2) + + Affine2D().rotate_deg_around(*self.pivot, 90)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_equal(trans.transform(self.single_point), [4, 3]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[6, 0], [8, 9], [2, 12]]) + + trans = (Affine2D().scale(3, -2) + .skew_deg(26.5650512, 14.0362435)) # ~atan(0.5), ~atan(0.25) + trans_added = (Affine2D().scale(3, -2) + + Affine2D().skew_deg(26.5650512, 14.0362435)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [2, -1.25]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[-2, -4], [6, -3.75], [12, 3]]) + + trans = Affine2D().scale(3, -2).translate(23, 42) + trans_added = Affine2D().scale(3, -2) + Affine2D().translate(23, 42) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_equal(trans.transform(self.single_point), [26, 40]) + assert_array_equal(trans.transform(self.multiple_points), + [[23, 38], [32, 36], [35, 42]]) + + def test_skew_plus_other(self): + # Using ~atan(0.5), ~atan(0.25) produces roundish numbers on output. + trans = Affine2D().skew_deg(26.5650512, 14.0362435).rotate_deg(90) + trans_added = (Affine2D().skew_deg(26.5650512, 14.0362435) + + Affine2D().rotate_deg(90)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [-1.25, 1.5]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[-2, 1], [-3.75, 4.5], [-1, 4]]) + + trans = (Affine2D().skew_deg(26.5650512, 14.0362435) + .rotate_deg_around(*self.pivot, 90)) + trans_added = (Affine2D().skew_deg(26.5650512, 14.0362435) + + Affine2D().rotate_deg_around(*self.pivot, 90)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [0.75, 1.5]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[0, 1], [-1.75, 4.5], [1, 4]]) + + trans = Affine2D().skew_deg(26.5650512, 14.0362435).scale(3, -2) + trans_added = (Affine2D().skew_deg(26.5650512, 14.0362435) + + Affine2D().scale(3, -2)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [4.5, -2.5]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[3, -4], [13.5, -7.5], [12, -2]]) + + trans = Affine2D().skew_deg(26.5650512, 14.0362435).translate(23, 42) + trans_added = (Affine2D().skew_deg(26.5650512, 14.0362435) + + Affine2D().translate(23, 42)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [24.5, 43.25]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[24, 44], [27.5, 45.75], [27, 43]]) + + def test_translate_plus_other(self): + trans = Affine2D().translate(23, 42).rotate_deg(90) + trans_added = Affine2D().translate(23, 42) + Affine2D().rotate_deg(90) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [-43, 24]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[-44, 23], [-45, 26], [-42, 27]]) + + trans = Affine2D().translate(23, 42).rotate_deg_around(*self.pivot, 90) + trans_added = (Affine2D().translate(23, 42) + + Affine2D().rotate_deg_around(*self.pivot, 90)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [-41, 24]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[-42, 23], [-43, 26], [-40, 27]]) + + trans = Affine2D().translate(23, 42).scale(3, -2) + trans_added = Affine2D().translate(23, 42) + Affine2D().scale(3, -2) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [72, -86]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[69, -88], [78, -90], [81, -84]]) + + trans = (Affine2D().translate(23, 42) + .skew_deg(26.5650512, 14.0362435)) # ~atan(0.5), ~atan(0.25) + trans_added = (Affine2D().translate(23, 42) + + Affine2D().skew_deg(26.5650512, 14.0362435)) + assert_array_equal(trans.get_matrix(), trans_added.get_matrix()) + assert_array_almost_equal(trans.transform(self.single_point), [45.5, 49]) + assert_array_almost_equal(trans.transform(self.multiple_points), + [[45, 49.75], [48.5, 51.5], [48, 48.75]]) + + def test_invalid_transform(self): + t = mtransforms.Affine2D() + # There are two different exceptions, since the wrong number of + # dimensions is caught when constructing an array_view, and that + # raises a ValueError, and a wrong shape with a possible number + # of dimensions is caught by our CALL_CPP macro, which always + # raises the less precise RuntimeError. + with pytest.raises(ValueError): + t.transform(1) + with pytest.raises(ValueError): + t.transform([[[1]]]) + with pytest.raises(RuntimeError): + t.transform([]) + with pytest.raises(RuntimeError): + t.transform([1]) + with pytest.raises(ValueError): + t.transform([[1]]) + with pytest.raises(ValueError): + t.transform([[1, 2, 3]]) + + def test_copy(self): + a = mtransforms.Affine2D() + b = mtransforms.Affine2D() + s = a + b + # Updating a dependee should invalidate a copy of the dependent. + s.get_matrix() # resolve it. + s1 = copy.copy(s) + assert not s._invalid and not s1._invalid + a.translate(1, 2) + assert s._invalid and s1._invalid + assert (s1.get_matrix() == a.get_matrix()).all() + # Updating a copy of a dependee shouldn't invalidate a dependent. + s.get_matrix() # resolve it. + b1 = copy.copy(b) + b1.translate(3, 4) + assert not s._invalid + assert_array_equal(s.get_matrix(), a.get_matrix()) + + def test_deepcopy(self): + a = mtransforms.Affine2D() + b = mtransforms.Affine2D() + s = a + b + # Updating a dependee shouldn't invalidate a deepcopy of the dependent. + s.get_matrix() # resolve it. + s1 = copy.deepcopy(s) + assert not s._invalid and not s1._invalid + a.translate(1, 2) + assert s._invalid and not s1._invalid + assert_array_equal(s1.get_matrix(), mtransforms.Affine2D().get_matrix()) + # Updating a deepcopy of a dependee shouldn't invalidate a dependent. + s.get_matrix() # resolve it. + b1 = copy.deepcopy(b) + b1.translate(3, 4) + assert not s._invalid + assert_array_equal(s.get_matrix(), a.get_matrix()) + + +class TestAffineDeltaTransform: + def test_invalidate(self): + before = np.array([[1.0, 4.0, 0.0], + [5.0, 1.0, 0.0], + [0.0, 0.0, 1.0]]) + after = np.array([[1.0, 3.0, 0.0], + [5.0, 1.0, 0.0], + [0.0, 0.0, 1.0]]) + + # Translation and skew present + base = mtransforms.Affine2D.from_values(1, 5, 4, 1, 2, 3) + t = mtransforms.AffineDeltaTransform(base) + assert_array_equal(t.get_matrix(), before) + + # Mess with the internal structure of `base` without invalidating + # This should not affect this transform because it's a passthrough: + # it's always invalid + base.get_matrix()[0, 1:] = 3 + assert_array_equal(t.get_matrix(), after) + + # Invalidate the base + base.invalidate() + assert_array_equal(t.get_matrix(), after) + + +def test_non_affine_caching(): + class AssertingNonAffineTransform(mtransforms.Transform): + """ + This transform raises an assertion error when called when it + shouldn't be and ``self.raise_on_transform`` is True. + + """ + input_dims = output_dims = 2 + is_affine = False + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.raise_on_transform = False + self.underlying_transform = mtransforms.Affine2D().scale(10, 10) + + def transform_path_non_affine(self, path): + assert not self.raise_on_transform, \ + 'Invalidated affine part of transform unnecessarily.' + return self.underlying_transform.transform_path(path) + transform_path = transform_path_non_affine + + def transform_non_affine(self, path): + assert not self.raise_on_transform, \ + 'Invalidated affine part of transform unnecessarily.' + return self.underlying_transform.transform(path) + transform = transform_non_affine + + my_trans = AssertingNonAffineTransform() + ax = plt.axes() + plt.plot(np.arange(10), transform=my_trans + ax.transData) + plt.draw() + # enable the transform to raise an exception if it's non-affine transform + # method is triggered again. + my_trans.raise_on_transform = True + ax.transAxes.invalidate() + plt.draw() + + +def test_external_transform_api(): + class ScaledBy: + def __init__(self, scale_factor): + self._scale_factor = scale_factor + + def _as_mpl_transform(self, axes): + return (mtransforms.Affine2D().scale(self._scale_factor) + + axes.transData) + + ax = plt.axes() + line, = plt.plot(np.arange(10), transform=ScaledBy(10)) + ax.set_xlim(0, 100) + ax.set_ylim(0, 100) + # assert that the top transform of the line is the scale transform. + assert_allclose(line.get_transform()._a.get_matrix(), + mtransforms.Affine2D().scale(10).get_matrix()) + + +@image_comparison(['pre_transform_data'], remove_text=True, style='mpl20', + tol=0.05) +def test_pre_transform_plotting(): + # a catch-all for as many as possible plot layouts which handle + # pre-transforming the data NOTE: The axis range is important in this + # plot. It should be x10 what the data suggests it should be + + ax = plt.axes() + times10 = mtransforms.Affine2D().scale(10) + + ax.contourf(np.arange(48).reshape(6, 8), transform=times10 + ax.transData) + + ax.pcolormesh(np.linspace(0, 4, 7), + np.linspace(5.5, 8, 9), + np.arange(48).reshape(8, 6), + transform=times10 + ax.transData) + + ax.scatter(np.linspace(0, 10), np.linspace(10, 0), + transform=times10 + ax.transData) + + x = np.linspace(8, 10, 20) + y = np.linspace(1, 5, 20) + u = 2*np.sin(x) + np.cos(y[:, np.newaxis]) + v = np.sin(x) - np.cos(y[:, np.newaxis]) + + ax.streamplot(x, y, u, v, transform=times10 + ax.transData, + linewidth=np.hypot(u, v)) + + # reduce the vector data down a bit for barb and quiver plotting + x, y = x[::3], y[::3] + u, v = u[::3, ::3], v[::3, ::3] + + ax.quiver(x, y + 5, u, v, transform=times10 + ax.transData) + + ax.barbs(x - 3, y + 5, u**2, v**2, transform=times10 + ax.transData) + + +def test_contour_pre_transform_limits(): + ax = plt.axes() + xs, ys = np.meshgrid(np.linspace(15, 20, 15), np.linspace(12.4, 12.5, 20)) + ax.contourf(xs, ys, np.log(xs * ys), + transform=mtransforms.Affine2D().scale(0.1) + ax.transData) + + expected = np.array([[1.5, 1.24], + [2., 1.25]]) + assert_almost_equal(expected, ax.dataLim.get_points()) + + +def test_pcolor_pre_transform_limits(): + # Based on test_contour_pre_transform_limits() + ax = plt.axes() + xs, ys = np.meshgrid(np.linspace(15, 20, 15), np.linspace(12.4, 12.5, 20)) + ax.pcolor(xs, ys, np.log(xs * ys)[:-1, :-1], + transform=mtransforms.Affine2D().scale(0.1) + ax.transData) + + expected = np.array([[1.5, 1.24], + [2., 1.25]]) + assert_almost_equal(expected, ax.dataLim.get_points()) + + +def test_pcolormesh_pre_transform_limits(): + # Based on test_contour_pre_transform_limits() + ax = plt.axes() + xs, ys = np.meshgrid(np.linspace(15, 20, 15), np.linspace(12.4, 12.5, 20)) + ax.pcolormesh(xs, ys, np.log(xs * ys)[:-1, :-1], + transform=mtransforms.Affine2D().scale(0.1) + ax.transData) + + expected = np.array([[1.5, 1.24], + [2., 1.25]]) + assert_almost_equal(expected, ax.dataLim.get_points()) + + +def test_pcolormesh_gouraud_nans(): + np.random.seed(19680801) + + values = np.linspace(0, 180, 3) + radii = np.linspace(100, 1000, 10) + z, y = np.meshgrid(values, radii) + x = np.radians(np.random.rand(*z.shape) * 100) + + fig = plt.figure() + ax = fig.add_subplot(111, projection="polar") + # Setting the limit to cause clipping of the r values causes NaN to be + # introduced; these should not crash but be ignored as in other path + # operations. + ax.set_rlim(101, 1000) + ax.pcolormesh(x, y, z, shading="gouraud") + + fig.canvas.draw() + + +def test_Affine2D_from_values(): + points = np.array([[0, 0], + [10, 20], + [-1, 0], + ]) + + t = mtransforms.Affine2D.from_values(1, 0, 0, 0, 0, 0) + actual = t.transform(points) + expected = np.array([[0, 0], [10, 0], [-1, 0]]) + assert_almost_equal(actual, expected) + + t = mtransforms.Affine2D.from_values(0, 2, 0, 0, 0, 0) + actual = t.transform(points) + expected = np.array([[0, 0], [0, 20], [0, -2]]) + assert_almost_equal(actual, expected) + + t = mtransforms.Affine2D.from_values(0, 0, 3, 0, 0, 0) + actual = t.transform(points) + expected = np.array([[0, 0], [60, 0], [0, 0]]) + assert_almost_equal(actual, expected) + + t = mtransforms.Affine2D.from_values(0, 0, 0, 4, 0, 0) + actual = t.transform(points) + expected = np.array([[0, 0], [0, 80], [0, 0]]) + assert_almost_equal(actual, expected) + + t = mtransforms.Affine2D.from_values(0, 0, 0, 0, 5, 0) + actual = t.transform(points) + expected = np.array([[5, 0], [5, 0], [5, 0]]) + assert_almost_equal(actual, expected) + + t = mtransforms.Affine2D.from_values(0, 0, 0, 0, 0, 6) + actual = t.transform(points) + expected = np.array([[0, 6], [0, 6], [0, 6]]) + assert_almost_equal(actual, expected) + + +def test_affine_inverted_invalidated(): + # Ensure that the an affine transform is not declared valid on access + point = [1.0, 1.0] + t = mtransforms.Affine2D() + + assert_almost_equal(point, t.transform(t.inverted().transform(point))) + # Change and access the transform + t.translate(1.0, 1.0).get_matrix() + assert_almost_equal(point, t.transform(t.inverted().transform(point))) + + +def test_clipping_of_log(): + # issue 804 + path = Path._create_closed([(0.2, -99), (0.4, -99), (0.4, 20), (0.2, 20)]) + # something like this happens in plotting logarithmic histograms + trans = mtransforms.BlendedGenericTransform( + mtransforms.Affine2D(), scale.LogTransform(10, 'clip')) + tpath = trans.transform_path_non_affine(path) + result = tpath.iter_segments(trans.get_affine(), + clip=(0, 0, 100, 100), + simplify=False) + tpoints, tcodes = zip(*result) + assert_allclose(tcodes, path.codes[:-1]) # No longer closed. + + +class NonAffineForTest(mtransforms.Transform): + """ + A class which looks like a non affine transform, but does whatever + the given transform does (even if it is affine). This is very useful + for testing NonAffine behaviour with a simple Affine transform. + + """ + is_affine = False + output_dims = 2 + input_dims = 2 + + def __init__(self, real_trans, *args, **kwargs): + self.real_trans = real_trans + super().__init__(*args, **kwargs) + + def transform_non_affine(self, values): + return self.real_trans.transform(values) + + def transform_path_non_affine(self, path): + return self.real_trans.transform_path(path) + + +class TestBasicTransform: + def setup_method(self): + + self.ta1 = mtransforms.Affine2D(shorthand_name='ta1').rotate(np.pi / 2) + self.ta2 = mtransforms.Affine2D(shorthand_name='ta2').translate(10, 0) + self.ta3 = mtransforms.Affine2D(shorthand_name='ta3').scale(1, 2) + + self.tn1 = NonAffineForTest(mtransforms.Affine2D().translate(1, 2), + shorthand_name='tn1') + self.tn2 = NonAffineForTest(mtransforms.Affine2D().translate(1, 2), + shorthand_name='tn2') + self.tn3 = NonAffineForTest(mtransforms.Affine2D().translate(1, 2), + shorthand_name='tn3') + + # creates a transform stack which looks like ((A, (N, A)), A) + self.stack1 = (self.ta1 + (self.tn1 + self.ta2)) + self.ta3 + # creates a transform stack which looks like (((A, N), A), A) + self.stack2 = self.ta1 + self.tn1 + self.ta2 + self.ta3 + # creates a transform stack which is a subset of stack2 + self.stack2_subset = self.tn1 + self.ta2 + self.ta3 + + # when in debug, the transform stacks can produce dot images: +# self.stack1.write_graphviz(file('stack1.dot', 'w')) +# self.stack2.write_graphviz(file('stack2.dot', 'w')) +# self.stack2_subset.write_graphviz(file('stack2_subset.dot', 'w')) + + def test_transform_depth(self): + assert self.stack1.depth == 4 + assert self.stack2.depth == 4 + assert self.stack2_subset.depth == 3 + + def test_left_to_right_iteration(self): + stack3 = (self.ta1 + (self.tn1 + (self.ta2 + self.tn2))) + self.ta3 +# stack3.write_graphviz(file('stack3.dot', 'w')) + + target_transforms = [stack3, + (self.tn1 + (self.ta2 + self.tn2)) + self.ta3, + (self.ta2 + self.tn2) + self.ta3, + self.tn2 + self.ta3, + self.ta3, + ] + r = [rh for _, rh in stack3._iter_break_from_left_to_right()] + assert len(r) == len(target_transforms) + + for target_stack, stack in zip(target_transforms, r): + assert target_stack == stack + + def test_transform_shortcuts(self): + assert self.stack1 - self.stack2_subset == self.ta1 + assert self.stack2 - self.stack2_subset == self.ta1 + + assert self.stack2_subset - self.stack2 == self.ta1.inverted() + assert (self.stack2_subset - self.stack2).depth == 1 + + with pytest.raises(ValueError): + self.stack1 - self.stack2 + + aff1 = self.ta1 + (self.ta2 + self.ta3) + aff2 = self.ta2 + self.ta3 + + assert aff1 - aff2 == self.ta1 + assert aff1 - self.ta2 == aff1 + self.ta2.inverted() + + assert self.stack1 - self.ta3 == self.ta1 + (self.tn1 + self.ta2) + assert self.stack2 - self.ta3 == self.ta1 + self.tn1 + self.ta2 + + assert ((self.ta2 + self.ta3) - self.ta3 + self.ta3 == + self.ta2 + self.ta3) + + def test_contains_branch(self): + r1 = (self.ta2 + self.ta1) + r2 = (self.ta2 + self.ta1) + assert r1 == r2 + assert r1 != self.ta1 + assert r1.contains_branch(r2) + assert r1.contains_branch(self.ta1) + assert not r1.contains_branch(self.ta2) + assert not r1.contains_branch(self.ta2 + self.ta2) + + assert r1 == r2 + + assert self.stack1.contains_branch(self.ta3) + assert self.stack2.contains_branch(self.ta3) + + assert self.stack1.contains_branch(self.stack2_subset) + assert self.stack2.contains_branch(self.stack2_subset) + + assert not self.stack2_subset.contains_branch(self.stack1) + assert not self.stack2_subset.contains_branch(self.stack2) + + assert self.stack1.contains_branch(self.ta2 + self.ta3) + assert self.stack2.contains_branch(self.ta2 + self.ta3) + + assert not self.stack1.contains_branch(self.tn1 + self.ta2) + + blend = mtransforms.BlendedGenericTransform(self.tn2, self.stack2) + x, y = blend.contains_branch_seperately(self.stack2_subset) + stack_blend = self.tn3 + blend + sx, sy = stack_blend.contains_branch_seperately(self.stack2_subset) + assert x is sx is False + assert y is sy is True + + def test_affine_simplification(self): + # tests that a transform stack only calls as much is absolutely + # necessary "non-affine" allowing the best possible optimization with + # complex transformation stacks. + points = np.array([[0, 0], [10, 20], [np.nan, 1], [-1, 0]], + dtype=np.float64) + na_pts = self.stack1.transform_non_affine(points) + all_pts = self.stack1.transform(points) + + na_expected = np.array([[1., 2.], [-19., 12.], + [np.nan, np.nan], [1., 1.]], dtype=np.float64) + all_expected = np.array([[11., 4.], [-9., 24.], + [np.nan, np.nan], [11., 2.]], + dtype=np.float64) + + # check we have the expected results from doing the affine part only + assert_array_almost_equal(na_pts, na_expected) + # check we have the expected results from a full transformation + assert_array_almost_equal(all_pts, all_expected) + # check we have the expected results from doing the transformation in + # two steps + assert_array_almost_equal(self.stack1.transform_affine(na_pts), + all_expected) + # check that getting the affine transformation first, then fully + # transforming using that yields the same result as before. + assert_array_almost_equal(self.stack1.get_affine().transform(na_pts), + all_expected) + + # check that the affine part of stack1 & stack2 are equivalent + # (i.e. the optimization is working) + expected_result = (self.ta2 + self.ta3).get_matrix() + result = self.stack1.get_affine().get_matrix() + assert_array_equal(expected_result, result) + + result = self.stack2.get_affine().get_matrix() + assert_array_equal(expected_result, result) + + +class TestTransformPlotInterface: + def test_line_extent_axes_coords(self): + # a simple line in axes coordinates + ax = plt.axes() + ax.plot([0.1, 1.2, 0.8], [0.9, 0.5, 0.8], transform=ax.transAxes) + assert_array_equal(ax.dataLim.get_points(), + np.array([[np.inf, np.inf], + [-np.inf, -np.inf]])) + + def test_line_extent_data_coords(self): + # a simple line in data coordinates + ax = plt.axes() + ax.plot([0.1, 1.2, 0.8], [0.9, 0.5, 0.8], transform=ax.transData) + assert_array_equal(ax.dataLim.get_points(), + np.array([[0.1, 0.5], [1.2, 0.9]])) + + def test_line_extent_compound_coords1(self): + # a simple line in data coordinates in the y component, and in axes + # coordinates in the x + ax = plt.axes() + trans = mtransforms.blended_transform_factory(ax.transAxes, + ax.transData) + ax.plot([0.1, 1.2, 0.8], [35, -5, 18], transform=trans) + assert_array_equal(ax.dataLim.get_points(), + np.array([[np.inf, -5.], + [-np.inf, 35.]])) + + def test_line_extent_predata_transform_coords(self): + # a simple line in (offset + data) coordinates + ax = plt.axes() + trans = mtransforms.Affine2D().scale(10) + ax.transData + ax.plot([0.1, 1.2, 0.8], [35, -5, 18], transform=trans) + assert_array_equal(ax.dataLim.get_points(), + np.array([[1., -50.], [12., 350.]])) + + def test_line_extent_compound_coords2(self): + # a simple line in (offset + data) coordinates in the y component, and + # in axes coordinates in the x + ax = plt.axes() + trans = mtransforms.blended_transform_factory( + ax.transAxes, mtransforms.Affine2D().scale(10) + ax.transData) + ax.plot([0.1, 1.2, 0.8], [35, -5, 18], transform=trans) + assert_array_equal(ax.dataLim.get_points(), + np.array([[np.inf, -50.], [-np.inf, 350.]])) + + def test_line_extents_affine(self): + ax = plt.axes() + offset = mtransforms.Affine2D().translate(10, 10) + plt.plot(np.arange(10), transform=offset + ax.transData) + expected_data_lim = np.array([[0., 0.], [9., 9.]]) + 10 + assert_array_almost_equal(ax.dataLim.get_points(), expected_data_lim) + + def test_line_extents_non_affine(self): + ax = plt.axes() + offset = mtransforms.Affine2D().translate(10, 10) + na_offset = NonAffineForTest(mtransforms.Affine2D().translate(10, 10)) + plt.plot(np.arange(10), transform=offset + na_offset + ax.transData) + expected_data_lim = np.array([[0., 0.], [9., 9.]]) + 20 + assert_array_almost_equal(ax.dataLim.get_points(), expected_data_lim) + + def test_pathc_extents_non_affine(self): + ax = plt.axes() + offset = mtransforms.Affine2D().translate(10, 10) + na_offset = NonAffineForTest(mtransforms.Affine2D().translate(10, 10)) + pth = Path([[0, 0], [0, 10], [10, 10], [10, 0]]) + patch = mpatches.PathPatch(pth, + transform=offset + na_offset + ax.transData) + ax.add_patch(patch) + expected_data_lim = np.array([[0., 0.], [10., 10.]]) + 20 + assert_array_almost_equal(ax.dataLim.get_points(), expected_data_lim) + + def test_pathc_extents_affine(self): + ax = plt.axes() + offset = mtransforms.Affine2D().translate(10, 10) + pth = Path([[0, 0], [0, 10], [10, 10], [10, 0]]) + patch = mpatches.PathPatch(pth, transform=offset + ax.transData) + ax.add_patch(patch) + expected_data_lim = np.array([[0., 0.], [10., 10.]]) + 10 + assert_array_almost_equal(ax.dataLim.get_points(), expected_data_lim) + + def test_line_extents_for_non_affine_transData(self): + ax = plt.axes(projection='polar') + # add 10 to the radius of the data + offset = mtransforms.Affine2D().translate(0, 10) + + plt.plot(np.arange(10), transform=offset + ax.transData) + # the data lim of a polar plot is stored in coordinates + # before a transData transformation, hence the data limits + # are not what is being shown on the actual plot. + expected_data_lim = np.array([[0., 0.], [9., 9.]]) + [0, 10] + assert_array_almost_equal(ax.dataLim.get_points(), expected_data_lim) + + +def assert_bbox_eq(bbox1, bbox2): + assert_array_equal(bbox1.bounds, bbox2.bounds) + + +def test_bbox_frozen_copies_minpos(): + bbox = mtransforms.Bbox.from_extents(0.0, 0.0, 1.0, 1.0, minpos=1.0) + frozen = bbox.frozen() + assert_array_equal(frozen.minpos, bbox.minpos) + + +def test_bbox_intersection(): + bbox_from_ext = mtransforms.Bbox.from_extents + inter = mtransforms.Bbox.intersection + + r1 = bbox_from_ext(0, 0, 1, 1) + r2 = bbox_from_ext(0.5, 0.5, 1.5, 1.5) + r3 = bbox_from_ext(0.5, 0, 0.75, 0.75) + r4 = bbox_from_ext(0.5, 1.5, 1, 2.5) + r5 = bbox_from_ext(1, 1, 2, 2) + + # self intersection -> no change + assert_bbox_eq(inter(r1, r1), r1) + # simple intersection + assert_bbox_eq(inter(r1, r2), bbox_from_ext(0.5, 0.5, 1, 1)) + # r3 contains r2 + assert_bbox_eq(inter(r1, r3), r3) + # no intersection + assert inter(r1, r4) is None + # single point + assert_bbox_eq(inter(r1, r5), bbox_from_ext(1, 1, 1, 1)) + + +def test_bbox_as_strings(): + b = mtransforms.Bbox([[.5, 0], [.75, .75]]) + assert_bbox_eq(b, eval(repr(b), {'Bbox': mtransforms.Bbox})) + asdict = eval(str(b), {'Bbox': dict}) + for k, v in asdict.items(): + assert getattr(b, k) == v + fmt = '.1f' + asdict = eval(format(b, fmt), {'Bbox': dict}) + for k, v in asdict.items(): + assert eval(format(getattr(b, k), fmt)) == v + + +def test_str_transform(): + # The str here should not be considered as "absolutely stable", and may be + # reformatted later; this is just a smoketest for __str__. + assert str(plt.subplot(projection="polar").transData) == """\ +CompositeGenericTransform( + CompositeGenericTransform( + CompositeGenericTransform( + TransformWrapper( + BlendedAffine2D( + IdentityTransform(), + IdentityTransform())), + CompositeAffine2D( + Affine2D().scale(1.0), + Affine2D().scale(1.0))), + PolarTransform( + PolarAxes(0.125,0.1;0.775x0.8), + use_rmin=True, + apply_theta_transforms=False)), + CompositeGenericTransform( + CompositeGenericTransform( + PolarAffine( + TransformWrapper( + BlendedAffine2D( + IdentityTransform(), + IdentityTransform())), + LockableBbox( + Bbox(x0=0.0, y0=0.0, x1=6.283185307179586, y1=1.0), + [[-- --] + [-- --]])), + BboxTransformFrom( + _WedgeBbox( + (0.5, 0.5), + TransformedBbox( + Bbox(x0=0.0, y0=0.0, x1=6.283185307179586, y1=1.0), + CompositeAffine2D( + Affine2D().scale(1.0), + Affine2D().scale(1.0))), + LockableBbox( + Bbox(x0=0.0, y0=0.0, x1=6.283185307179586, y1=1.0), + [[-- --] + [-- --]])))), + BboxTransformTo( + TransformedBbox( + Bbox(x0=0.125, y0=0.09999999999999998, x1=0.9, y1=0.9), + BboxTransformTo( + TransformedBbox( + Bbox(x0=0.0, y0=0.0, x1=8.0, y1=6.0), + Affine2D().scale(80.0)))))))""" + + +def test_transform_single_point(): + t = mtransforms.Affine2D() + r = t.transform_affine((1, 1)) + assert r.shape == (2,) + + +def test_log_transform(): + # Tests that the last line runs without exception (previously the + # transform would fail if one of the axes was logarithmic). + fig, ax = plt.subplots() + ax.set_yscale('log') + ax.transData.transform((1, 1)) + + +def test_nan_overlap(): + a = mtransforms.Bbox([[0, 0], [1, 1]]) + b = mtransforms.Bbox([[0, 0], [1, np.nan]]) + assert not a.overlaps(b) + + +def test_transform_angles(): + t = mtransforms.Affine2D() # Identity transform + angles = np.array([20, 45, 60]) + points = np.array([[0, 0], [1, 1], [2, 2]]) + + # Identity transform does not change angles + new_angles = t.transform_angles(angles, points) + assert_array_almost_equal(angles, new_angles) + + # points missing a 2nd dimension + with pytest.raises(ValueError): + t.transform_angles(angles, points[0:2, 0:1]) + + # Number of angles != Number of points + with pytest.raises(ValueError): + t.transform_angles(angles, points[0:2, :]) + + +def test_nonsingular(): + # test for zero-expansion type cases; other cases may be added later + zero_expansion = np.array([-0.001, 0.001]) + cases = [(0, np.nan), (0, 0), (0, 7.9e-317)] + for args in cases: + out = np.array(mtransforms.nonsingular(*args)) + assert_array_equal(out, zero_expansion) + + +def test_transformed_path(): + points = [(0, 0), (1, 0), (1, 1), (0, 1)] + path = Path(points, closed=True) + + trans = mtransforms.Affine2D() + trans_path = mtransforms.TransformedPath(path, trans) + assert_allclose(trans_path.get_fully_transformed_path().vertices, points) + + # Changing the transform should change the result. + r2 = 1 / np.sqrt(2) + trans.rotate(np.pi / 4) + assert_allclose(trans_path.get_fully_transformed_path().vertices, + [(0, 0), (r2, r2), (0, 2 * r2), (-r2, r2)], + atol=1e-15) + + # Changing the path does not change the result (it's cached). + path.points = [(0, 0)] * 4 + assert_allclose(trans_path.get_fully_transformed_path().vertices, + [(0, 0), (r2, r2), (0, 2 * r2), (-r2, r2)], + atol=1e-15) + + +def test_transformed_patch_path(): + trans = mtransforms.Affine2D() + patch = mpatches.Wedge((0, 0), 1, 45, 135, transform=trans) + + tpatch = mtransforms.TransformedPatchPath(patch) + points = tpatch.get_fully_transformed_path().vertices + + # Changing the transform should change the result. + trans.scale(2) + assert_allclose(tpatch.get_fully_transformed_path().vertices, points * 2) + + # Changing the path should change the result (and cancel out the scaling + # from the transform). + patch.set_radius(0.5) + assert_allclose(tpatch.get_fully_transformed_path().vertices, points) + + +@pytest.mark.parametrize('locked_element', ['x0', 'y0', 'x1', 'y1']) +def test_lockable_bbox(locked_element): + other_elements = ['x0', 'y0', 'x1', 'y1'] + other_elements.remove(locked_element) + + orig = mtransforms.Bbox.unit() + locked = mtransforms.LockableBbox(orig, **{locked_element: 2}) + + # LockableBbox should keep its locked element as specified in __init__. + assert getattr(locked, locked_element) == 2 + assert getattr(locked, 'locked_' + locked_element) == 2 + for elem in other_elements: + assert getattr(locked, elem) == getattr(orig, elem) + + # Changing underlying Bbox should update everything but locked element. + orig.set_points(orig.get_points() + 10) + assert getattr(locked, locked_element) == 2 + assert getattr(locked, 'locked_' + locked_element) == 2 + for elem in other_elements: + assert getattr(locked, elem) == getattr(orig, elem) + + # Unlocking element should revert values back to the underlying Bbox. + setattr(locked, 'locked_' + locked_element, None) + assert getattr(locked, 'locked_' + locked_element) is None + assert np.all(orig.get_points() == locked.get_points()) + + # Relocking an element should change its value, but not others. + setattr(locked, 'locked_' + locked_element, 3) + assert getattr(locked, locked_element) == 3 + assert getattr(locked, 'locked_' + locked_element) == 3 + for elem in other_elements: + assert getattr(locked, elem) == getattr(orig, elem) + + +def test_transformwrapper(): + t = mtransforms.TransformWrapper(mtransforms.Affine2D()) + with pytest.raises(ValueError, match=( + r"The input and output dims of the new child \(1, 1\) " + r"do not match those of current child \(2, 2\)")): + t.set(scale.LogTransform(10)) + + +@check_figures_equal(extensions=["png"]) +def test_scale_swapping(fig_test, fig_ref): + np.random.seed(19680801) + samples = np.random.normal(size=10) + x = np.linspace(-5, 5, 10) + + for fig, log_state in zip([fig_test, fig_ref], [True, False]): + ax = fig.subplots() + ax.hist(samples, log=log_state, density=True) + ax.plot(x, np.exp(-(x**2) / 2) / np.sqrt(2 * np.pi)) + fig.canvas.draw() + ax.set_yscale('linear') + + +def test_offset_copy_errors(): + with pytest.raises(ValueError, + match="'fontsize' is not a valid value for units;" + " supported values are 'dots', 'points', 'inches'"): + mtransforms.offset_copy(None, units='fontsize') + + with pytest.raises(ValueError, + match='For units of inches or points a fig kwarg is needed'): + mtransforms.offset_copy(None, units='inches') + + +def test_transformedbbox_contains(): + bb = TransformedBbox(Bbox.unit(), Affine2D().rotate_deg(30)) + assert bb.contains(.8, .5) + assert bb.contains(-.4, .85) + assert not bb.contains(.9, .5) + bb = TransformedBbox(Bbox.unit(), Affine2D().translate(.25, .5)) + assert bb.contains(1.25, 1.5) + assert not bb.fully_contains(1.25, 1.5) + assert not bb.fully_contains(.1, .1) + + +def test_interval_contains(): + assert mtransforms.interval_contains((0, 1), 0.5) + assert mtransforms.interval_contains((0, 1), 0) + assert mtransforms.interval_contains((0, 1), 1) + assert not mtransforms.interval_contains((0, 1), -1) + assert not mtransforms.interval_contains((0, 1), 2) + assert mtransforms.interval_contains((1, 0), 0.5) + + +def test_interval_contains_open(): + assert mtransforms.interval_contains_open((0, 1), 0.5) + assert not mtransforms.interval_contains_open((0, 1), 0) + assert not mtransforms.interval_contains_open((0, 1), 1) + assert not mtransforms.interval_contains_open((0, 1), -1) + assert not mtransforms.interval_contains_open((0, 1), 2) + assert mtransforms.interval_contains_open((1, 0), 0.5) + + +def test_scaledrotation_initialization(): + """Test that the ScaledRotation object is initialized correctly.""" + theta = 1.0 # Arbitrary theta value for testing + trans_shift = MagicMock() # Mock the trans_shift transformation + scaled_rot = _ScaledRotation(theta, trans_shift) + assert scaled_rot._theta == theta + assert scaled_rot._trans_shift == trans_shift + assert scaled_rot._mtx is None + + +def test_scaledrotation_get_matrix_invalid(): + """Test get_matrix when the matrix is invalid and needs recalculation.""" + theta = np.pi / 2 + trans_shift = MagicMock(transform=MagicMock(return_value=[[theta, 0]])) + scaled_rot = _ScaledRotation(theta, trans_shift) + scaled_rot._invalid = True + matrix = scaled_rot.get_matrix() + trans_shift.transform.assert_called_once_with([[theta, 0]]) + expected_rotation = np.array([[0, -1], + [1, 0]]) + assert matrix is not None + assert_allclose(matrix[:2, :2], expected_rotation, atol=1e-15) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_triangulation.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_triangulation.py new file mode 100644 index 0000000000000000000000000000000000000000..dd91144f240cefa592a0c73387974c7fcd977ba9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_triangulation.py @@ -0,0 +1,1403 @@ +import numpy as np +from numpy.testing import ( + assert_array_equal, assert_array_almost_equal, assert_array_less) +import numpy.ma.testutils as matest +import pytest + +import matplotlib as mpl +import matplotlib.pyplot as plt +import matplotlib.tri as mtri +from matplotlib.path import Path +from matplotlib.testing.decorators import image_comparison, check_figures_equal + + +class TestTriangulationParams: + x = [-1, 0, 1, 0] + y = [0, -1, 0, 1] + triangles = [[0, 1, 2], [0, 2, 3]] + mask = [False, True] + + @pytest.mark.parametrize('args, kwargs, expected', [ + ([x, y], {}, [x, y, None, None]), + ([x, y, triangles], {}, [x, y, triangles, None]), + ([x, y], dict(triangles=triangles), [x, y, triangles, None]), + ([x, y], dict(mask=mask), [x, y, None, mask]), + ([x, y, triangles], dict(mask=mask), [x, y, triangles, mask]), + ([x, y], dict(triangles=triangles, mask=mask), [x, y, triangles, mask]) + ]) + def test_extract_triangulation_params(self, args, kwargs, expected): + other_args = [1, 2] + other_kwargs = {'a': 3, 'b': '4'} + x_, y_, triangles_, mask_, args_, kwargs_ = \ + mtri.Triangulation._extract_triangulation_params( + args + other_args, {**kwargs, **other_kwargs}) + x, y, triangles, mask = expected + assert x_ is x + assert y_ is y + assert_array_equal(triangles_, triangles) + assert mask_ is mask + assert args_ == other_args + assert kwargs_ == other_kwargs + + +def test_extract_triangulation_positional_mask(): + # mask cannot be passed positionally + mask = [True] + args = [[0, 2, 1], [0, 0, 1], [[0, 1, 2]], mask] + x_, y_, triangles_, mask_, args_, kwargs_ = \ + mtri.Triangulation._extract_triangulation_params(args, {}) + assert mask_ is None + assert args_ == [mask] + # the positional mask must be caught downstream because this must pass + # unknown args through + + +def test_triangulation_init(): + x = [-1, 0, 1, 0] + y = [0, -1, 0, 1] + with pytest.raises(ValueError, match="x and y must be equal-length"): + mtri.Triangulation(x, [1, 2]) + with pytest.raises( + ValueError, + match=r"triangles must be a \(N, 3\) int array, but found shape " + r"\(3,\)"): + mtri.Triangulation(x, y, [0, 1, 2]) + with pytest.raises( + ValueError, + match=r"triangles must be a \(N, 3\) int array, not 'other'"): + mtri.Triangulation(x, y, 'other') + with pytest.raises(ValueError, match="found value 99"): + mtri.Triangulation(x, y, [[0, 1, 99]]) + with pytest.raises(ValueError, match="found value -1"): + mtri.Triangulation(x, y, [[0, 1, -1]]) + + +def test_triangulation_set_mask(): + x = [-1, 0, 1, 0] + y = [0, -1, 0, 1] + triangles = [[0, 1, 2], [2, 3, 0]] + triang = mtri.Triangulation(x, y, triangles) + + # Check neighbors, which forces creation of C++ triangulation + assert_array_equal(triang.neighbors, [[-1, -1, 1], [-1, -1, 0]]) + + # Set mask + triang.set_mask([False, True]) + assert_array_equal(triang.mask, [False, True]) + + # Reset mask + triang.set_mask(None) + assert triang.mask is None + + msg = r"mask array must have same length as triangles array" + for mask in ([False, True, False], [False], [True], False, True): + with pytest.raises(ValueError, match=msg): + triang.set_mask(mask) + + +def test_delaunay(): + # No duplicate points, regular grid. + nx = 5 + ny = 4 + x, y = np.meshgrid(np.linspace(0.0, 1.0, nx), np.linspace(0.0, 1.0, ny)) + x = x.ravel() + y = y.ravel() + npoints = nx*ny + ntriangles = 2 * (nx-1) * (ny-1) + nedges = 3*nx*ny - 2*nx - 2*ny + 1 + + # Create delaunay triangulation. + triang = mtri.Triangulation(x, y) + + # The tests in the remainder of this function should be passed by any + # triangulation that does not contain duplicate points. + + # Points - floating point. + assert_array_almost_equal(triang.x, x) + assert_array_almost_equal(triang.y, y) + + # Triangles - integers. + assert len(triang.triangles) == ntriangles + assert np.min(triang.triangles) == 0 + assert np.max(triang.triangles) == npoints-1 + + # Edges - integers. + assert len(triang.edges) == nedges + assert np.min(triang.edges) == 0 + assert np.max(triang.edges) == npoints-1 + + # Neighbors - integers. + # Check that neighbors calculated by C++ triangulation class are the same + # as those returned from delaunay routine. + neighbors = triang.neighbors + triang._neighbors = None + assert_array_equal(triang.neighbors, neighbors) + + # Is each point used in at least one triangle? + assert_array_equal(np.unique(triang.triangles), np.arange(npoints)) + + +def test_delaunay_duplicate_points(): + npoints = 10 + duplicate = 7 + duplicate_of = 3 + + np.random.seed(23) + x = np.random.random(npoints) + y = np.random.random(npoints) + x[duplicate] = x[duplicate_of] + y[duplicate] = y[duplicate_of] + + # Create delaunay triangulation. + triang = mtri.Triangulation(x, y) + + # Duplicate points should be ignored, so the index of the duplicate points + # should not appear in any triangle. + assert_array_equal(np.unique(triang.triangles), + np.delete(np.arange(npoints), duplicate)) + + +def test_delaunay_points_in_line(): + # Cannot triangulate points that are all in a straight line, but check + # that delaunay code fails gracefully. + x = np.linspace(0.0, 10.0, 11) + y = np.linspace(0.0, 10.0, 11) + with pytest.raises(RuntimeError): + mtri.Triangulation(x, y) + + # Add an extra point not on the line and the triangulation is OK. + x = np.append(x, 2.0) + y = np.append(y, 8.0) + mtri.Triangulation(x, y) + + +@pytest.mark.parametrize('x, y', [ + # Triangulation should raise a ValueError if passed less than 3 points. + ([], []), + ([1], [5]), + ([1, 2], [5, 6]), + # Triangulation should also raise a ValueError if passed duplicate points + # such that there are less than 3 unique points. + ([1, 2, 1], [5, 6, 5]), + ([1, 2, 2], [5, 6, 6]), + ([1, 1, 1, 2, 1, 2], [5, 5, 5, 6, 5, 6]), +]) +def test_delaunay_insufficient_points(x, y): + with pytest.raises(ValueError): + mtri.Triangulation(x, y) + + +def test_delaunay_robust(): + # Fails when mtri.Triangulation uses matplotlib.delaunay, works when using + # qhull. + tri_points = np.array([ + [0.8660254037844384, -0.5000000000000004], + [0.7577722283113836, -0.5000000000000004], + [0.6495190528383288, -0.5000000000000003], + [0.5412658773652739, -0.5000000000000003], + [0.811898816047911, -0.40625000000000044], + [0.7036456405748561, -0.4062500000000004], + [0.5953924651018013, -0.40625000000000033]]) + test_points = np.asarray([ + [0.58, -0.46], + [0.65, -0.46], + [0.65, -0.42], + [0.7, -0.48], + [0.7, -0.44], + [0.75, -0.44], + [0.8, -0.48]]) + + # Utility function that indicates if a triangle defined by 3 points + # (xtri, ytri) contains the test point xy. Avoid calling with a point that + # lies on or very near to an edge of the triangle. + def tri_contains_point(xtri, ytri, xy): + tri_points = np.vstack((xtri, ytri)).T + return Path(tri_points).contains_point(xy) + + # Utility function that returns how many triangles of the specified + # triangulation contain the test point xy. Avoid calling with a point that + # lies on or very near to an edge of any triangle in the triangulation. + def tris_contain_point(triang, xy): + return sum(tri_contains_point(triang.x[tri], triang.y[tri], xy) + for tri in triang.triangles) + + # Using matplotlib.delaunay, an invalid triangulation is created with + # overlapping triangles; qhull is OK. + triang = mtri.Triangulation(tri_points[:, 0], tri_points[:, 1]) + for test_point in test_points: + assert tris_contain_point(triang, test_point) == 1 + + # If ignore the first point of tri_points, matplotlib.delaunay throws a + # KeyError when calculating the convex hull; qhull is OK. + triang = mtri.Triangulation(tri_points[1:, 0], tri_points[1:, 1]) + + +@image_comparison(['tripcolor1.png']) +def test_tripcolor(): + x = np.asarray([0, 0.5, 1, 0, 0.5, 1, 0, 0.5, 1, 0.75]) + y = np.asarray([0, 0, 0, 0.5, 0.5, 0.5, 1, 1, 1, 0.75]) + triangles = np.asarray([ + [0, 1, 3], [1, 4, 3], + [1, 2, 4], [2, 5, 4], + [3, 4, 6], [4, 7, 6], + [4, 5, 9], [7, 4, 9], [8, 7, 9], [5, 8, 9]]) + + # Triangulation with same number of points and triangles. + triang = mtri.Triangulation(x, y, triangles) + + Cpoints = x + 0.5*y + + xmid = x[triang.triangles].mean(axis=1) + ymid = y[triang.triangles].mean(axis=1) + Cfaces = 0.5*xmid + ymid + + plt.subplot(121) + plt.tripcolor(triang, Cpoints, edgecolors='k') + plt.title('point colors') + + plt.subplot(122) + plt.tripcolor(triang, facecolors=Cfaces, edgecolors='k') + plt.title('facecolors') + + +def test_tripcolor_color(): + x = [-1, 0, 1, 0] + y = [0, -1, 0, 1] + fig, ax = plt.subplots() + with pytest.raises(TypeError, match=r"tripcolor\(\) missing 1 required "): + ax.tripcolor(x, y) + with pytest.raises(ValueError, match="The length of c must match either"): + ax.tripcolor(x, y, [1, 2, 3]) + with pytest.raises(ValueError, + match="length of facecolors must match .* triangles"): + ax.tripcolor(x, y, facecolors=[1, 2, 3, 4]) + with pytest.raises(ValueError, + match="'gouraud' .* at the points.* not at the faces"): + ax.tripcolor(x, y, facecolors=[1, 2], shading='gouraud') + with pytest.raises(ValueError, + match="'gouraud' .* at the points.* not at the faces"): + ax.tripcolor(x, y, [1, 2], shading='gouraud') # faces + with pytest.raises(TypeError, + match="positional.*'c'.*keyword-only.*'facecolors'"): + ax.tripcolor(x, y, C=[1, 2, 3, 4]) + with pytest.raises(TypeError, match="Unexpected positional parameter"): + ax.tripcolor(x, y, [1, 2], 'unused_positional') + + # smoke test for valid color specifications (via C or facecolors) + ax.tripcolor(x, y, [1, 2, 3, 4]) # edges + ax.tripcolor(x, y, [1, 2, 3, 4], shading='gouraud') # edges + ax.tripcolor(x, y, [1, 2]) # faces + ax.tripcolor(x, y, facecolors=[1, 2]) # faces + + +def test_tripcolor_clim(): + np.random.seed(19680801) + a, b, c = np.random.rand(10), np.random.rand(10), np.random.rand(10) + + ax = plt.figure().add_subplot() + clim = (0.25, 0.75) + norm = ax.tripcolor(a, b, c, clim=clim).norm + assert (norm.vmin, norm.vmax) == clim + + +def test_tripcolor_warnings(): + x = [-1, 0, 1, 0] + y = [0, -1, 0, 1] + c = [0.4, 0.5] + fig, ax = plt.subplots() + # facecolors takes precedence over c + with pytest.warns(UserWarning, match="Positional parameter c .*no effect"): + ax.tripcolor(x, y, c, facecolors=c) + with pytest.warns(UserWarning, match="Positional parameter c .*no effect"): + ax.tripcolor(x, y, 'interpreted as c', facecolors=c) + + +def test_no_modify(): + # Test that Triangulation does not modify triangles array passed to it. + triangles = np.array([[3, 2, 0], [3, 1, 0]], dtype=np.int32) + points = np.array([(0, 0), (0, 1.1), (1, 0), (1, 1)]) + + old_triangles = triangles.copy() + mtri.Triangulation(points[:, 0], points[:, 1], triangles).edges + assert_array_equal(old_triangles, triangles) + + +def test_trifinder(): + # Test points within triangles of masked triangulation. + x, y = np.meshgrid(np.arange(4), np.arange(4)) + x = x.ravel() + y = y.ravel() + triangles = [[0, 1, 4], [1, 5, 4], [1, 2, 5], [2, 6, 5], [2, 3, 6], + [3, 7, 6], [4, 5, 8], [5, 9, 8], [5, 6, 9], [6, 10, 9], + [6, 7, 10], [7, 11, 10], [8, 9, 12], [9, 13, 12], [9, 10, 13], + [10, 14, 13], [10, 11, 14], [11, 15, 14]] + mask = np.zeros(len(triangles)) + mask[8:10] = 1 + triang = mtri.Triangulation(x, y, triangles, mask) + trifinder = triang.get_trifinder() + + xs = [0.25, 1.25, 2.25, 3.25] + ys = [0.25, 1.25, 2.25, 3.25] + xs, ys = np.meshgrid(xs, ys) + xs = xs.ravel() + ys = ys.ravel() + tris = trifinder(xs, ys) + assert_array_equal(tris, [0, 2, 4, -1, 6, -1, 10, -1, + 12, 14, 16, -1, -1, -1, -1, -1]) + tris = trifinder(xs-0.5, ys-0.5) + assert_array_equal(tris, [-1, -1, -1, -1, -1, 1, 3, 5, + -1, 7, -1, 11, -1, 13, 15, 17]) + + # Test points exactly on boundary edges of masked triangulation. + xs = [0.5, 1.5, 2.5, 0.5, 1.5, 2.5, 1.5, 1.5, 0.0, 1.0, 2.0, 3.0] + ys = [0.0, 0.0, 0.0, 3.0, 3.0, 3.0, 1.0, 2.0, 1.5, 1.5, 1.5, 1.5] + tris = trifinder(xs, ys) + assert_array_equal(tris, [0, 2, 4, 13, 15, 17, 3, 14, 6, 7, 10, 11]) + + # Test points exactly on boundary corners of masked triangulation. + xs = [0.0, 3.0] + ys = [0.0, 3.0] + tris = trifinder(xs, ys) + assert_array_equal(tris, [0, 17]) + + # + # Test triangles with horizontal colinear points. These are not valid + # triangulations, but we try to deal with the simplest violations. + # + + # If +ve, triangulation is OK, if -ve triangulation invalid, + # if zero have colinear points but should pass tests anyway. + delta = 0.0 + + x = [1.5, 0, 1, 2, 3, 1.5, 1.5] + y = [-1, 0, 0, 0, 0, delta, 1] + triangles = [[0, 2, 1], [0, 3, 2], [0, 4, 3], [1, 2, 5], [2, 3, 5], + [3, 4, 5], [1, 5, 6], [4, 6, 5]] + triang = mtri.Triangulation(x, y, triangles) + trifinder = triang.get_trifinder() + + xs = [-0.1, 0.4, 0.9, 1.4, 1.9, 2.4, 2.9] + ys = [-0.1, 0.1] + xs, ys = np.meshgrid(xs, ys) + tris = trifinder(xs, ys) + assert_array_equal(tris, [[-1, 0, 0, 1, 1, 2, -1], + [-1, 6, 6, 6, 7, 7, -1]]) + + # + # Test triangles with vertical colinear points. These are not valid + # triangulations, but we try to deal with the simplest violations. + # + + # If +ve, triangulation is OK, if -ve triangulation invalid, + # if zero have colinear points but should pass tests anyway. + delta = 0.0 + + x = [-1, -delta, 0, 0, 0, 0, 1] + y = [1.5, 1.5, 0, 1, 2, 3, 1.5] + triangles = [[0, 1, 2], [0, 1, 5], [1, 2, 3], [1, 3, 4], [1, 4, 5], + [2, 6, 3], [3, 6, 4], [4, 6, 5]] + triang = mtri.Triangulation(x, y, triangles) + trifinder = triang.get_trifinder() + + xs = [-0.1, 0.1] + ys = [-0.1, 0.4, 0.9, 1.4, 1.9, 2.4, 2.9] + xs, ys = np.meshgrid(xs, ys) + tris = trifinder(xs, ys) + assert_array_equal(tris, [[-1, -1], [0, 5], [0, 5], [0, 6], [1, 6], [1, 7], + [-1, -1]]) + + # Test that changing triangulation by setting a mask causes the trifinder + # to be reinitialised. + x = [0, 1, 0, 1] + y = [0, 0, 1, 1] + triangles = [[0, 1, 2], [1, 3, 2]] + triang = mtri.Triangulation(x, y, triangles) + trifinder = triang.get_trifinder() + + xs = [-0.2, 0.2, 0.8, 1.2] + ys = [0.5, 0.5, 0.5, 0.5] + tris = trifinder(xs, ys) + assert_array_equal(tris, [-1, 0, 1, -1]) + + triang.set_mask([1, 0]) + assert trifinder == triang.get_trifinder() + tris = trifinder(xs, ys) + assert_array_equal(tris, [-1, -1, 1, -1]) + + +def test_triinterp(): + # Test points within triangles of masked triangulation. + x, y = np.meshgrid(np.arange(4), np.arange(4)) + x = x.ravel() + y = y.ravel() + z = 1.23*x - 4.79*y + triangles = [[0, 1, 4], [1, 5, 4], [1, 2, 5], [2, 6, 5], [2, 3, 6], + [3, 7, 6], [4, 5, 8], [5, 9, 8], [5, 6, 9], [6, 10, 9], + [6, 7, 10], [7, 11, 10], [8, 9, 12], [9, 13, 12], [9, 10, 13], + [10, 14, 13], [10, 11, 14], [11, 15, 14]] + mask = np.zeros(len(triangles)) + mask[8:10] = 1 + triang = mtri.Triangulation(x, y, triangles, mask) + linear_interp = mtri.LinearTriInterpolator(triang, z) + cubic_min_E = mtri.CubicTriInterpolator(triang, z) + cubic_geom = mtri.CubicTriInterpolator(triang, z, kind='geom') + + xs = np.linspace(0.25, 2.75, 6) + ys = [0.25, 0.75, 2.25, 2.75] + xs, ys = np.meshgrid(xs, ys) # Testing arrays with array.ndim = 2 + for interp in (linear_interp, cubic_min_E, cubic_geom): + zs = interp(xs, ys) + assert_array_almost_equal(zs, (1.23*xs - 4.79*ys)) + + # Test points outside triangulation. + xs = [-0.25, 1.25, 1.75, 3.25] + ys = xs + xs, ys = np.meshgrid(xs, ys) + for interp in (linear_interp, cubic_min_E, cubic_geom): + zs = linear_interp(xs, ys) + assert_array_equal(zs.mask, [[True]*4]*4) + + # Test mixed configuration (outside / inside). + xs = np.linspace(0.25, 1.75, 6) + ys = [0.25, 0.75, 1.25, 1.75] + xs, ys = np.meshgrid(xs, ys) + for interp in (linear_interp, cubic_min_E, cubic_geom): + zs = interp(xs, ys) + matest.assert_array_almost_equal(zs, (1.23*xs - 4.79*ys)) + mask = (xs >= 1) * (xs <= 2) * (ys >= 1) * (ys <= 2) + assert_array_equal(zs.mask, mask) + + # 2nd order patch test: on a grid with an 'arbitrary shaped' triangle, + # patch test shall be exact for quadratic functions and cubic + # interpolator if *kind* = user + (a, b, c) = (1.23, -4.79, 0.6) + + def quad(x, y): + return a*(x-0.5)**2 + b*(y-0.5)**2 + c*x*y + + def gradient_quad(x, y): + return (2*a*(x-0.5) + c*y, 2*b*(y-0.5) + c*x) + + x = np.array([0.2, 0.33367, 0.669, 0., 1., 1., 0.]) + y = np.array([0.3, 0.80755, 0.4335, 0., 0., 1., 1.]) + triangles = np.array([[0, 1, 2], [3, 0, 4], [4, 0, 2], [4, 2, 5], + [1, 5, 2], [6, 5, 1], [6, 1, 0], [6, 0, 3]]) + triang = mtri.Triangulation(x, y, triangles) + z = quad(x, y) + dz = gradient_quad(x, y) + # test points for 2nd order patch test + xs = np.linspace(0., 1., 5) + ys = np.linspace(0., 1., 5) + xs, ys = np.meshgrid(xs, ys) + cubic_user = mtri.CubicTriInterpolator(triang, z, kind='user', dz=dz) + interp_zs = cubic_user(xs, ys) + assert_array_almost_equal(interp_zs, quad(xs, ys)) + (interp_dzsdx, interp_dzsdy) = cubic_user.gradient(x, y) + (dzsdx, dzsdy) = gradient_quad(x, y) + assert_array_almost_equal(interp_dzsdx, dzsdx) + assert_array_almost_equal(interp_dzsdy, dzsdy) + + # Cubic improvement: cubic interpolation shall perform better than linear + # on a sufficiently dense mesh for a quadratic function. + n = 11 + x, y = np.meshgrid(np.linspace(0., 1., n+1), np.linspace(0., 1., n+1)) + x = x.ravel() + y = y.ravel() + z = quad(x, y) + triang = mtri.Triangulation(x, y, triangles=meshgrid_triangles(n+1)) + xs, ys = np.meshgrid(np.linspace(0.1, 0.9, 5), np.linspace(0.1, 0.9, 5)) + xs = xs.ravel() + ys = ys.ravel() + linear_interp = mtri.LinearTriInterpolator(triang, z) + cubic_min_E = mtri.CubicTriInterpolator(triang, z) + cubic_geom = mtri.CubicTriInterpolator(triang, z, kind='geom') + zs = quad(xs, ys) + diff_lin = np.abs(linear_interp(xs, ys) - zs) + for interp in (cubic_min_E, cubic_geom): + diff_cubic = np.abs(interp(xs, ys) - zs) + assert np.max(diff_lin) >= 10 * np.max(diff_cubic) + assert (np.dot(diff_lin, diff_lin) >= + 100 * np.dot(diff_cubic, diff_cubic)) + + +def test_triinterpcubic_C1_continuity(): + # Below the 4 tests which demonstrate C1 continuity of the + # TriCubicInterpolator (testing the cubic shape functions on arbitrary + # triangle): + # + # 1) Testing continuity of function & derivatives at corner for all 9 + # shape functions. Testing also function values at same location. + # 2) Testing C1 continuity along each edge (as gradient is polynomial of + # 2nd order, it is sufficient to test at the middle). + # 3) Testing C1 continuity at triangle barycenter (where the 3 subtriangles + # meet) + # 4) Testing C1 continuity at median 1/3 points (midside between 2 + # subtriangles) + + # Utility test function check_continuity + def check_continuity(interpolator, loc, values=None): + """ + Checks the continuity of interpolator (and its derivatives) near + location loc. Can check the value at loc itself if *values* is + provided. + + *interpolator* TriInterpolator + *loc* location to test (x0, y0) + *values* (optional) array [z0, dzx0, dzy0] to check the value at *loc* + """ + n_star = 24 # Number of continuity points in a boundary of loc + epsilon = 1.e-10 # Distance for loc boundary + k = 100. # Continuity coefficient + (loc_x, loc_y) = loc + star_x = loc_x + epsilon*np.cos(np.linspace(0., 2*np.pi, n_star)) + star_y = loc_y + epsilon*np.sin(np.linspace(0., 2*np.pi, n_star)) + z = interpolator([loc_x], [loc_y])[0] + (dzx, dzy) = interpolator.gradient([loc_x], [loc_y]) + if values is not None: + assert_array_almost_equal(z, values[0]) + assert_array_almost_equal(dzx[0], values[1]) + assert_array_almost_equal(dzy[0], values[2]) + diff_z = interpolator(star_x, star_y) - z + (tab_dzx, tab_dzy) = interpolator.gradient(star_x, star_y) + diff_dzx = tab_dzx - dzx + diff_dzy = tab_dzy - dzy + assert_array_less(diff_z, epsilon*k) + assert_array_less(diff_dzx, epsilon*k) + assert_array_less(diff_dzy, epsilon*k) + + # Drawing arbitrary triangle (a, b, c) inside a unit square. + (ax, ay) = (0.2, 0.3) + (bx, by) = (0.33367, 0.80755) + (cx, cy) = (0.669, 0.4335) + x = np.array([ax, bx, cx, 0., 1., 1., 0.]) + y = np.array([ay, by, cy, 0., 0., 1., 1.]) + triangles = np.array([[0, 1, 2], [3, 0, 4], [4, 0, 2], [4, 2, 5], + [1, 5, 2], [6, 5, 1], [6, 1, 0], [6, 0, 3]]) + triang = mtri.Triangulation(x, y, triangles) + + for idof in range(9): + z = np.zeros(7, dtype=np.float64) + dzx = np.zeros(7, dtype=np.float64) + dzy = np.zeros(7, dtype=np.float64) + values = np.zeros([3, 3], dtype=np.float64) + case = idof//3 + values[case, idof % 3] = 1.0 + if case == 0: + z[idof] = 1.0 + elif case == 1: + dzx[idof % 3] = 1.0 + elif case == 2: + dzy[idof % 3] = 1.0 + interp = mtri.CubicTriInterpolator(triang, z, kind='user', + dz=(dzx, dzy)) + # Test 1) Checking values and continuity at nodes + check_continuity(interp, (ax, ay), values[:, 0]) + check_continuity(interp, (bx, by), values[:, 1]) + check_continuity(interp, (cx, cy), values[:, 2]) + # Test 2) Checking continuity at midside nodes + check_continuity(interp, ((ax+bx)*0.5, (ay+by)*0.5)) + check_continuity(interp, ((ax+cx)*0.5, (ay+cy)*0.5)) + check_continuity(interp, ((cx+bx)*0.5, (cy+by)*0.5)) + # Test 3) Checking continuity at barycenter + check_continuity(interp, ((ax+bx+cx)/3., (ay+by+cy)/3.)) + # Test 4) Checking continuity at median 1/3-point + check_continuity(interp, ((4.*ax+bx+cx)/6., (4.*ay+by+cy)/6.)) + check_continuity(interp, ((ax+4.*bx+cx)/6., (ay+4.*by+cy)/6.)) + check_continuity(interp, ((ax+bx+4.*cx)/6., (ay+by+4.*cy)/6.)) + + +def test_triinterpcubic_cg_solver(): + # Now 3 basic tests of the Sparse CG solver, used for + # TriCubicInterpolator with *kind* = 'min_E' + # 1) A commonly used test involves a 2d Poisson matrix. + def poisson_sparse_matrix(n, m): + """ + Return the sparse, (n*m, n*m) matrix in coo format resulting from the + discretisation of the 2-dimensional Poisson equation according to a + finite difference numerical scheme on a uniform (n, m) grid. + """ + l = m*n + rows = np.concatenate([ + np.arange(l, dtype=np.int32), + np.arange(l-1, dtype=np.int32), np.arange(1, l, dtype=np.int32), + np.arange(l-n, dtype=np.int32), np.arange(n, l, dtype=np.int32)]) + cols = np.concatenate([ + np.arange(l, dtype=np.int32), + np.arange(1, l, dtype=np.int32), np.arange(l-1, dtype=np.int32), + np.arange(n, l, dtype=np.int32), np.arange(l-n, dtype=np.int32)]) + vals = np.concatenate([ + 4*np.ones(l, dtype=np.float64), + -np.ones(l-1, dtype=np.float64), -np.ones(l-1, dtype=np.float64), + -np.ones(l-n, dtype=np.float64), -np.ones(l-n, dtype=np.float64)]) + # In fact +1 and -1 diags have some zeros + vals[l:2*l-1][m-1::m] = 0. + vals[2*l-1:3*l-2][m-1::m] = 0. + return vals, rows, cols, (n*m, n*m) + + # Instantiating a sparse Poisson matrix of size 48 x 48: + (n, m) = (12, 4) + mat = mtri._triinterpolate._Sparse_Matrix_coo(*poisson_sparse_matrix(n, m)) + mat.compress_csc() + mat_dense = mat.to_dense() + # Testing a sparse solve for all 48 basis vector + for itest in range(n*m): + b = np.zeros(n*m, dtype=np.float64) + b[itest] = 1. + x, _ = mtri._triinterpolate._cg(A=mat, b=b, x0=np.zeros(n*m), + tol=1.e-10) + assert_array_almost_equal(np.dot(mat_dense, x), b) + + # 2) Same matrix with inserting 2 rows - cols with null diag terms + # (but still linked with the rest of the matrix by extra-diag terms) + (i_zero, j_zero) = (12, 49) + vals, rows, cols, _ = poisson_sparse_matrix(n, m) + rows = rows + 1*(rows >= i_zero) + 1*(rows >= j_zero) + cols = cols + 1*(cols >= i_zero) + 1*(cols >= j_zero) + # adding extra-diag terms + rows = np.concatenate([rows, [i_zero, i_zero-1, j_zero, j_zero-1]]) + cols = np.concatenate([cols, [i_zero-1, i_zero, j_zero-1, j_zero]]) + vals = np.concatenate([vals, [1., 1., 1., 1.]]) + mat = mtri._triinterpolate._Sparse_Matrix_coo(vals, rows, cols, + (n*m + 2, n*m + 2)) + mat.compress_csc() + mat_dense = mat.to_dense() + # Testing a sparse solve for all 50 basis vec + for itest in range(n*m + 2): + b = np.zeros(n*m + 2, dtype=np.float64) + b[itest] = 1. + x, _ = mtri._triinterpolate._cg(A=mat, b=b, x0=np.ones(n * m + 2), + tol=1.e-10) + assert_array_almost_equal(np.dot(mat_dense, x), b) + + # 3) Now a simple test that summation of duplicate (i.e. with same rows, + # same cols) entries occurs when compressed. + vals = np.ones(17, dtype=np.float64) + rows = np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1], + dtype=np.int32) + cols = np.array([0, 1, 2, 1, 1, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2], + dtype=np.int32) + dim = (3, 3) + mat = mtri._triinterpolate._Sparse_Matrix_coo(vals, rows, cols, dim) + mat.compress_csc() + mat_dense = mat.to_dense() + assert_array_almost_equal(mat_dense, np.array([ + [1., 2., 0.], [2., 1., 5.], [0., 5., 1.]], dtype=np.float64)) + + +def test_triinterpcubic_geom_weights(): + # Tests to check computation of weights for _DOF_estimator_geom: + # The weight sum per triangle can be 1. (in case all angles < 90 degrees) + # or (2*w_i) where w_i = 1-alpha_i/np.pi is the weight of apex i; alpha_i + # is the apex angle > 90 degrees. + (ax, ay) = (0., 1.687) + x = np.array([ax, 0.5*ax, 0., 1.]) + y = np.array([ay, -ay, 0., 0.]) + z = np.zeros(4, dtype=np.float64) + triangles = [[0, 2, 3], [1, 3, 2]] + sum_w = np.zeros([4, 2]) # 4 possibilities; 2 triangles + for theta in np.linspace(0., 2*np.pi, 14): # rotating the figure... + x_rot = np.cos(theta)*x + np.sin(theta)*y + y_rot = -np.sin(theta)*x + np.cos(theta)*y + triang = mtri.Triangulation(x_rot, y_rot, triangles) + cubic_geom = mtri.CubicTriInterpolator(triang, z, kind='geom') + dof_estimator = mtri._triinterpolate._DOF_estimator_geom(cubic_geom) + weights = dof_estimator.compute_geom_weights() + # Testing for the 4 possibilities... + sum_w[0, :] = np.sum(weights, 1) - 1 + for itri in range(3): + sum_w[itri+1, :] = np.sum(weights, 1) - 2*weights[:, itri] + assert_array_almost_equal(np.min(np.abs(sum_w), axis=0), + np.array([0., 0.], dtype=np.float64)) + + +def test_triinterp_colinear(): + # Tests interpolating inside a triangulation with horizontal colinear + # points (refer also to the tests :func:`test_trifinder` ). + # + # These are not valid triangulations, but we try to deal with the + # simplest violations (i. e. those handled by default TriFinder). + # + # Note that the LinearTriInterpolator and the CubicTriInterpolator with + # kind='min_E' or 'geom' still pass a linear patch test. + # We also test interpolation inside a flat triangle, by forcing + # *tri_index* in a call to :meth:`_interpolate_multikeys`. + + # If +ve, triangulation is OK, if -ve triangulation invalid, + # if zero have colinear points but should pass tests anyway. + delta = 0. + + x0 = np.array([1.5, 0, 1, 2, 3, 1.5, 1.5]) + y0 = np.array([-1, 0, 0, 0, 0, delta, 1]) + + # We test different affine transformations of the initial figure; to + # avoid issues related to round-off errors we only use integer + # coefficients (otherwise the Triangulation might become invalid even with + # delta == 0). + transformations = [[1, 0], [0, 1], [1, 1], [1, 2], [-2, -1], [-2, 1]] + for transformation in transformations: + x_rot = transformation[0]*x0 + transformation[1]*y0 + y_rot = -transformation[1]*x0 + transformation[0]*y0 + (x, y) = (x_rot, y_rot) + z = 1.23*x - 4.79*y + triangles = [[0, 2, 1], [0, 3, 2], [0, 4, 3], [1, 2, 5], [2, 3, 5], + [3, 4, 5], [1, 5, 6], [4, 6, 5]] + triang = mtri.Triangulation(x, y, triangles) + xs = np.linspace(np.min(triang.x), np.max(triang.x), 20) + ys = np.linspace(np.min(triang.y), np.max(triang.y), 20) + xs, ys = np.meshgrid(xs, ys) + xs = xs.ravel() + ys = ys.ravel() + mask_out = (triang.get_trifinder()(xs, ys) == -1) + zs_target = np.ma.array(1.23*xs - 4.79*ys, mask=mask_out) + + linear_interp = mtri.LinearTriInterpolator(triang, z) + cubic_min_E = mtri.CubicTriInterpolator(triang, z) + cubic_geom = mtri.CubicTriInterpolator(triang, z, kind='geom') + + for interp in (linear_interp, cubic_min_E, cubic_geom): + zs = interp(xs, ys) + assert_array_almost_equal(zs_target, zs) + + # Testing interpolation inside the flat triangle number 4: [2, 3, 5] + # by imposing *tri_index* in a call to :meth:`_interpolate_multikeys` + itri = 4 + pt1 = triang.triangles[itri, 0] + pt2 = triang.triangles[itri, 1] + xs = np.linspace(triang.x[pt1], triang.x[pt2], 10) + ys = np.linspace(triang.y[pt1], triang.y[pt2], 10) + zs_target = 1.23*xs - 4.79*ys + for interp in (linear_interp, cubic_min_E, cubic_geom): + zs, = interp._interpolate_multikeys( + xs, ys, tri_index=itri*np.ones(10, dtype=np.int32)) + assert_array_almost_equal(zs_target, zs) + + +def test_triinterp_transformations(): + # 1) Testing that the interpolation scheme is invariant by rotation of the + # whole figure. + # Note: This test is non-trivial for a CubicTriInterpolator with + # kind='min_E'. It does fail for a non-isotropic stiffness matrix E of + # :class:`_ReducedHCT_Element` (tested with E=np.diag([1., 1., 1.])), and + # provides a good test for :meth:`get_Kff_and_Ff`of the same class. + # + # 2) Also testing that the interpolation scheme is invariant by expansion + # of the whole figure along one axis. + n_angles = 20 + n_radii = 10 + min_radius = 0.15 + + def z(x, y): + r1 = np.hypot(0.5 - x, 0.5 - y) + theta1 = np.arctan2(0.5 - x, 0.5 - y) + r2 = np.hypot(-x - 0.2, -y - 0.2) + theta2 = np.arctan2(-x - 0.2, -y - 0.2) + z = -(2*(np.exp((r1/10)**2)-1)*30. * np.cos(7.*theta1) + + (np.exp((r2/10)**2)-1)*30. * np.cos(11.*theta2) + + 0.7*(x**2 + y**2)) + return (np.max(z)-z)/(np.max(z)-np.min(z)) + + # First create the x and y coordinates of the points. + radii = np.linspace(min_radius, 0.95, n_radii) + angles = np.linspace(0 + n_angles, 2*np.pi + n_angles, + n_angles, endpoint=False) + angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1) + angles[:, 1::2] += np.pi/n_angles + x0 = (radii*np.cos(angles)).flatten() + y0 = (radii*np.sin(angles)).flatten() + triang0 = mtri.Triangulation(x0, y0) # Delaunay triangulation + z0 = z(x0, y0) + + # Then create the test points + xs0 = np.linspace(-1., 1., 23) + ys0 = np.linspace(-1., 1., 23) + xs0, ys0 = np.meshgrid(xs0, ys0) + xs0 = xs0.ravel() + ys0 = ys0.ravel() + + interp_z0 = {} + for i_angle in range(2): + # Rotating everything + theta = 2*np.pi / n_angles * i_angle + x = np.cos(theta)*x0 + np.sin(theta)*y0 + y = -np.sin(theta)*x0 + np.cos(theta)*y0 + xs = np.cos(theta)*xs0 + np.sin(theta)*ys0 + ys = -np.sin(theta)*xs0 + np.cos(theta)*ys0 + triang = mtri.Triangulation(x, y, triang0.triangles) + linear_interp = mtri.LinearTriInterpolator(triang, z0) + cubic_min_E = mtri.CubicTriInterpolator(triang, z0) + cubic_geom = mtri.CubicTriInterpolator(triang, z0, kind='geom') + dic_interp = {'lin': linear_interp, + 'min_E': cubic_min_E, + 'geom': cubic_geom} + # Testing that the interpolation is invariant by rotation... + for interp_key in ['lin', 'min_E', 'geom']: + interp = dic_interp[interp_key] + if i_angle == 0: + interp_z0[interp_key] = interp(xs0, ys0) # storage + else: + interpz = interp(xs, ys) + matest.assert_array_almost_equal(interpz, + interp_z0[interp_key]) + + scale_factor = 987654.3210 + for scaled_axis in ('x', 'y'): + # Scaling everything (expansion along scaled_axis) + if scaled_axis == 'x': + x = scale_factor * x0 + y = y0 + xs = scale_factor * xs0 + ys = ys0 + else: + x = x0 + y = scale_factor * y0 + xs = xs0 + ys = scale_factor * ys0 + triang = mtri.Triangulation(x, y, triang0.triangles) + linear_interp = mtri.LinearTriInterpolator(triang, z0) + cubic_min_E = mtri.CubicTriInterpolator(triang, z0) + cubic_geom = mtri.CubicTriInterpolator(triang, z0, kind='geom') + dic_interp = {'lin': linear_interp, + 'min_E': cubic_min_E, + 'geom': cubic_geom} + # Test that the interpolation is invariant by expansion along 1 axis... + for interp_key in ['lin', 'min_E', 'geom']: + interpz = dic_interp[interp_key](xs, ys) + matest.assert_array_almost_equal(interpz, interp_z0[interp_key]) + + +@image_comparison(['tri_smooth_contouring.png'], remove_text=True, tol=0.072) +def test_tri_smooth_contouring(): + # Image comparison based on example tricontour_smooth_user. + n_angles = 20 + n_radii = 10 + min_radius = 0.15 + + def z(x, y): + r1 = np.hypot(0.5 - x, 0.5 - y) + theta1 = np.arctan2(0.5 - x, 0.5 - y) + r2 = np.hypot(-x - 0.2, -y - 0.2) + theta2 = np.arctan2(-x - 0.2, -y - 0.2) + z = -(2*(np.exp((r1/10)**2)-1)*30. * np.cos(7.*theta1) + + (np.exp((r2/10)**2)-1)*30. * np.cos(11.*theta2) + + 0.7*(x**2 + y**2)) + return (np.max(z)-z)/(np.max(z)-np.min(z)) + + # First create the x and y coordinates of the points. + radii = np.linspace(min_radius, 0.95, n_radii) + angles = np.linspace(0 + n_angles, 2*np.pi + n_angles, + n_angles, endpoint=False) + angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1) + angles[:, 1::2] += np.pi/n_angles + x0 = (radii*np.cos(angles)).flatten() + y0 = (radii*np.sin(angles)).flatten() + triang0 = mtri.Triangulation(x0, y0) # Delaunay triangulation + z0 = z(x0, y0) + triang0.set_mask(np.hypot(x0[triang0.triangles].mean(axis=1), + y0[triang0.triangles].mean(axis=1)) + < min_radius) + + # Then the plot + refiner = mtri.UniformTriRefiner(triang0) + tri_refi, z_test_refi = refiner.refine_field(z0, subdiv=4) + levels = np.arange(0., 1., 0.025) + plt.triplot(triang0, lw=0.5, color='0.5') + plt.tricontour(tri_refi, z_test_refi, levels=levels, colors="black") + + +@image_comparison(['tri_smooth_gradient.png'], remove_text=True, tol=0.092) +def test_tri_smooth_gradient(): + # Image comparison based on example trigradient_demo. + + def dipole_potential(x, y): + """An electric dipole potential V.""" + r_sq = x**2 + y**2 + theta = np.arctan2(y, x) + z = np.cos(theta)/r_sq + return (np.max(z)-z) / (np.max(z)-np.min(z)) + + # Creating a Triangulation + n_angles = 30 + n_radii = 10 + min_radius = 0.2 + radii = np.linspace(min_radius, 0.95, n_radii) + angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False) + angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1) + angles[:, 1::2] += np.pi/n_angles + x = (radii*np.cos(angles)).flatten() + y = (radii*np.sin(angles)).flatten() + V = dipole_potential(x, y) + triang = mtri.Triangulation(x, y) + triang.set_mask(np.hypot(x[triang.triangles].mean(axis=1), + y[triang.triangles].mean(axis=1)) + < min_radius) + + # Refine data - interpolates the electrical potential V + refiner = mtri.UniformTriRefiner(triang) + tri_refi, z_test_refi = refiner.refine_field(V, subdiv=3) + + # Computes the electrical field (Ex, Ey) as gradient of -V + tci = mtri.CubicTriInterpolator(triang, -V) + Ex, Ey = tci.gradient(triang.x, triang.y) + E_norm = np.hypot(Ex, Ey) + + # Plot the triangulation, the potential iso-contours and the vector field + plt.figure() + plt.gca().set_aspect('equal') + plt.triplot(triang, color='0.8') + + levels = np.arange(0., 1., 0.01) + cmap = mpl.colormaps['hot'] + plt.tricontour(tri_refi, z_test_refi, levels=levels, cmap=cmap, + linewidths=[2.0, 1.0, 1.0, 1.0]) + # Plots direction of the electrical vector field + plt.quiver(triang.x, triang.y, Ex/E_norm, Ey/E_norm, + units='xy', scale=10., zorder=3, color='blue', + width=0.007, headwidth=3., headlength=4.) + # We are leaving ax.use_sticky_margins as True, so the + # view limits are the contour data limits. + + +def test_tritools(): + # Tests TriAnalyzer.scale_factors on masked triangulation + # Tests circle_ratios on equilateral and right-angled triangle. + x = np.array([0., 1., 0.5, 0., 2.]) + y = np.array([0., 0., 0.5*np.sqrt(3.), -1., 1.]) + triangles = np.array([[0, 1, 2], [0, 1, 3], [1, 2, 4]], dtype=np.int32) + mask = np.array([False, False, True], dtype=bool) + triang = mtri.Triangulation(x, y, triangles, mask=mask) + analyser = mtri.TriAnalyzer(triang) + assert_array_almost_equal(analyser.scale_factors, [1, 1/(1+3**.5/2)]) + assert_array_almost_equal( + analyser.circle_ratios(rescale=False), + np.ma.masked_array([0.5, 1./(1.+np.sqrt(2.)), np.nan], mask)) + + # Tests circle ratio of a flat triangle + x = np.array([0., 1., 2.]) + y = np.array([1., 1.+3., 1.+6.]) + triangles = np.array([[0, 1, 2]], dtype=np.int32) + triang = mtri.Triangulation(x, y, triangles) + analyser = mtri.TriAnalyzer(triang) + assert_array_almost_equal(analyser.circle_ratios(), np.array([0.])) + + # Tests TriAnalyzer.get_flat_tri_mask + # Creates a triangulation of [-1, 1] x [-1, 1] with contiguous groups of + # 'flat' triangles at the 4 corners and at the center. Checks that only + # those at the borders are eliminated by TriAnalyzer.get_flat_tri_mask + n = 9 + + def power(x, a): + return np.abs(x)**a*np.sign(x) + + x = np.linspace(-1., 1., n+1) + x, y = np.meshgrid(power(x, 2.), power(x, 0.25)) + x = x.ravel() + y = y.ravel() + + triang = mtri.Triangulation(x, y, triangles=meshgrid_triangles(n+1)) + analyser = mtri.TriAnalyzer(triang) + mask_flat = analyser.get_flat_tri_mask(0.2) + verif_mask = np.zeros(162, dtype=bool) + corners_index = [0, 1, 2, 3, 14, 15, 16, 17, 18, 19, 34, 35, 126, 127, + 142, 143, 144, 145, 146, 147, 158, 159, 160, 161] + verif_mask[corners_index] = True + assert_array_equal(mask_flat, verif_mask) + + # Now including a hole (masked triangle) at the center. The center also + # shall be eliminated by get_flat_tri_mask. + mask = np.zeros(162, dtype=bool) + mask[80] = True + triang.set_mask(mask) + mask_flat = analyser.get_flat_tri_mask(0.2) + center_index = [44, 45, 62, 63, 78, 79, 80, 81, 82, 83, 98, 99, 116, 117] + verif_mask[center_index] = True + assert_array_equal(mask_flat, verif_mask) + + +def test_trirefine(): + # Testing subdiv=2 refinement + n = 3 + subdiv = 2 + x = np.linspace(-1., 1., n+1) + x, y = np.meshgrid(x, x) + x = x.ravel() + y = y.ravel() + mask = np.zeros(2*n**2, dtype=bool) + mask[n**2:] = True + triang = mtri.Triangulation(x, y, triangles=meshgrid_triangles(n+1), + mask=mask) + refiner = mtri.UniformTriRefiner(triang) + refi_triang = refiner.refine_triangulation(subdiv=subdiv) + x_refi = refi_triang.x + y_refi = refi_triang.y + + n_refi = n * subdiv**2 + x_verif = np.linspace(-1., 1., n_refi+1) + x_verif, y_verif = np.meshgrid(x_verif, x_verif) + x_verif = x_verif.ravel() + y_verif = y_verif.ravel() + ind1d = np.isin(np.around(x_verif*(2.5+y_verif), 8), + np.around(x_refi*(2.5+y_refi), 8)) + assert_array_equal(ind1d, True) + + # Testing the mask of the refined triangulation + refi_mask = refi_triang.mask + refi_tri_barycenter_x = np.sum(refi_triang.x[refi_triang.triangles], + axis=1) / 3. + refi_tri_barycenter_y = np.sum(refi_triang.y[refi_triang.triangles], + axis=1) / 3. + tri_finder = triang.get_trifinder() + refi_tri_indices = tri_finder(refi_tri_barycenter_x, + refi_tri_barycenter_y) + refi_tri_mask = triang.mask[refi_tri_indices] + assert_array_equal(refi_mask, refi_tri_mask) + + # Testing that the numbering of triangles does not change the + # interpolation result. + x = np.asarray([0.0, 1.0, 0.0, 1.0]) + y = np.asarray([0.0, 0.0, 1.0, 1.0]) + triang = [mtri.Triangulation(x, y, [[0, 1, 3], [3, 2, 0]]), + mtri.Triangulation(x, y, [[0, 1, 3], [2, 0, 3]])] + z = np.hypot(x - 0.3, y - 0.4) + # Refining the 2 triangulations and reordering the points + xyz_data = [] + for i in range(2): + refiner = mtri.UniformTriRefiner(triang[i]) + refined_triang, refined_z = refiner.refine_field(z, subdiv=1) + xyz = np.dstack((refined_triang.x, refined_triang.y, refined_z))[0] + xyz = xyz[np.lexsort((xyz[:, 1], xyz[:, 0]))] + xyz_data += [xyz] + assert_array_almost_equal(xyz_data[0], xyz_data[1]) + + +@pytest.mark.parametrize('interpolator', + [mtri.LinearTriInterpolator, + mtri.CubicTriInterpolator], + ids=['linear', 'cubic']) +def test_trirefine_masked(interpolator): + # Repeated points means we will have fewer triangles than points, and thus + # get masking. + x, y = np.mgrid[:2, :2] + x = np.repeat(x.flatten(), 2) + y = np.repeat(y.flatten(), 2) + + z = np.zeros_like(x) + tri = mtri.Triangulation(x, y) + refiner = mtri.UniformTriRefiner(tri) + interp = interpolator(tri, z) + refiner.refine_field(z, triinterpolator=interp, subdiv=2) + + +def meshgrid_triangles(n): + """ + Return (2*(N-1)**2, 3) array of triangles to mesh (N, N)-point np.meshgrid. + """ + tri = [] + for i in range(n-1): + for j in range(n-1): + a = i + j*n + b = (i+1) + j*n + c = i + (j+1)*n + d = (i+1) + (j+1)*n + tri += [[a, b, d], [a, d, c]] + return np.array(tri, dtype=np.int32) + + +def test_triplot_return(): + # Check that triplot returns the artists it adds + ax = plt.figure().add_subplot() + triang = mtri.Triangulation( + [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0], + triangles=[[0, 1, 3], [3, 2, 0]]) + assert ax.triplot(triang, "b-") is not None, \ + 'triplot should return the artist it adds' + + +def test_trirefiner_fortran_contiguous_triangles(): + # github issue 4180. Test requires two arrays of triangles that are + # identical except that one is C-contiguous and one is fortran-contiguous. + triangles1 = np.array([[2, 0, 3], [2, 1, 0]]) + assert not np.isfortran(triangles1) + + triangles2 = np.array(triangles1, copy=True, order='F') + assert np.isfortran(triangles2) + + x = np.array([0.39, 0.59, 0.43, 0.32]) + y = np.array([33.99, 34.01, 34.19, 34.18]) + triang1 = mtri.Triangulation(x, y, triangles1) + triang2 = mtri.Triangulation(x, y, triangles2) + + refiner1 = mtri.UniformTriRefiner(triang1) + refiner2 = mtri.UniformTriRefiner(triang2) + + fine_triang1 = refiner1.refine_triangulation(subdiv=1) + fine_triang2 = refiner2.refine_triangulation(subdiv=1) + + assert_array_equal(fine_triang1.triangles, fine_triang2.triangles) + + +def test_qhull_triangle_orientation(): + # github issue 4437. + xi = np.linspace(-2, 2, 100) + x, y = map(np.ravel, np.meshgrid(xi, xi)) + w = (x > y - 1) & (x < -1.95) & (y > -1.2) + x, y = x[w], y[w] + theta = np.radians(25) + x1 = x*np.cos(theta) - y*np.sin(theta) + y1 = x*np.sin(theta) + y*np.cos(theta) + + # Calculate Delaunay triangulation using Qhull. + triang = mtri.Triangulation(x1, y1) + + # Neighbors returned by Qhull. + qhull_neighbors = triang.neighbors + + # Obtain neighbors using own C++ calculation. + triang._neighbors = None + own_neighbors = triang.neighbors + + assert_array_equal(qhull_neighbors, own_neighbors) + + +def test_trianalyzer_mismatched_indices(): + # github issue 4999. + x = np.array([0., 1., 0.5, 0., 2.]) + y = np.array([0., 0., 0.5*np.sqrt(3.), -1., 1.]) + triangles = np.array([[0, 1, 2], [0, 1, 3], [1, 2, 4]], dtype=np.int32) + mask = np.array([False, False, True], dtype=bool) + triang = mtri.Triangulation(x, y, triangles, mask=mask) + analyser = mtri.TriAnalyzer(triang) + # numpy >= 1.10 raises a VisibleDeprecationWarning in the following line + # prior to the fix. + analyser._get_compressed_triangulation() + + +def test_tricontourf_decreasing_levels(): + # github issue 5477. + x = [0.0, 1.0, 1.0] + y = [0.0, 0.0, 1.0] + z = [0.2, 0.4, 0.6] + plt.figure() + with pytest.raises(ValueError): + plt.tricontourf(x, y, z, [1.0, 0.0]) + + +def test_internal_cpp_api() -> None: + # Following github issue 8197. + from matplotlib import _tri # noqa: F401, ensure lazy-loaded module *is* loaded. + + # C++ Triangulation. + with pytest.raises( + TypeError, + match=r'__init__\(\): incompatible constructor arguments.'): + mpl._tri.Triangulation() # type: ignore[call-arg] + + with pytest.raises( + ValueError, match=r'x and y must be 1D arrays of the same length'): + mpl._tri.Triangulation(np.array([]), np.array([1]), np.array([[]]), (), (), (), + False) + + x = np.array([0, 1, 1], dtype=np.float64) + y = np.array([0, 0, 1], dtype=np.float64) + with pytest.raises( + ValueError, + match=r'triangles must be a 2D array of shape \(\?,3\)'): + mpl._tri.Triangulation(x, y, np.array([[0, 1]]), (), (), (), False) + + tris = np.array([[0, 1, 2]], dtype=np.int_) + with pytest.raises( + ValueError, + match=r'mask must be a 1D array with the same length as the ' + r'triangles array'): + mpl._tri.Triangulation(x, y, tris, np.array([0, 1]), (), (), False) + + with pytest.raises( + ValueError, match=r'edges must be a 2D array with shape \(\?,2\)'): + mpl._tri.Triangulation(x, y, tris, (), np.array([[1]]), (), False) + + with pytest.raises( + ValueError, + match=r'neighbors must be a 2D array with the same shape as the ' + r'triangles array'): + mpl._tri.Triangulation(x, y, tris, (), (), np.array([[-1]]), False) + + triang = mpl._tri.Triangulation(x, y, tris, (), (), (), False) + + with pytest.raises( + ValueError, + match=r'z must be a 1D array with the same length as the ' + r'triangulation x and y arrays'): + triang.calculate_plane_coefficients([]) + + for mask in ([0, 1], None): + with pytest.raises( + ValueError, + match=r'mask must be a 1D array with the same length as the ' + r'triangles array'): + triang.set_mask(mask) # type: ignore[arg-type] + + triang.set_mask(np.array([True])) + assert_array_equal(triang.get_edges(), np.empty((0, 2))) + + triang.set_mask(()) # Equivalent to Python Triangulation mask=None + assert_array_equal(triang.get_edges(), [[1, 0], [2, 0], [2, 1]]) + + # C++ TriContourGenerator. + with pytest.raises( + TypeError, + match=r'__init__\(\): incompatible constructor arguments.'): + mpl._tri.TriContourGenerator() # type: ignore[call-arg] + + with pytest.raises( + ValueError, + match=r'z must be a 1D array with the same length as the x and y arrays'): + mpl._tri.TriContourGenerator(triang, np.array([1])) + + z = np.array([0, 1, 2]) + tcg = mpl._tri.TriContourGenerator(triang, z) + + with pytest.raises( + ValueError, match=r'filled contour levels must be increasing'): + tcg.create_filled_contour(1, 0) + + # C++ TrapezoidMapTriFinder. + with pytest.raises( + TypeError, + match=r'__init__\(\): incompatible constructor arguments.'): + mpl._tri.TrapezoidMapTriFinder() # type: ignore[call-arg] + + trifinder = mpl._tri.TrapezoidMapTriFinder(triang) + + with pytest.raises( + ValueError, match=r'x and y must be array-like with same shape'): + trifinder.find_many(np.array([0]), np.array([0, 1])) + + +def test_qhull_large_offset(): + # github issue 8682. + x = np.asarray([0, 1, 0, 1, 0.5]) + y = np.asarray([0, 0, 1, 1, 0.5]) + + offset = 1e10 + triang = mtri.Triangulation(x, y) + triang_offset = mtri.Triangulation(x + offset, y + offset) + assert len(triang.triangles) == len(triang_offset.triangles) + + +def test_tricontour_non_finite_z(): + # github issue 10167. + x = [0, 1, 0, 1] + y = [0, 0, 1, 1] + triang = mtri.Triangulation(x, y) + plt.figure() + + with pytest.raises(ValueError, match='z array must not contain non-finite ' + 'values within the triangulation'): + plt.tricontourf(triang, [0, 1, 2, np.inf]) + + with pytest.raises(ValueError, match='z array must not contain non-finite ' + 'values within the triangulation'): + plt.tricontourf(triang, [0, 1, 2, -np.inf]) + + with pytest.raises(ValueError, match='z array must not contain non-finite ' + 'values within the triangulation'): + plt.tricontourf(triang, [0, 1, 2, np.nan]) + + with pytest.raises(ValueError, match='z must not contain masked points ' + 'within the triangulation'): + plt.tricontourf(triang, np.ma.array([0, 1, 2, 3], mask=[1, 0, 0, 0])) + + +def test_tricontourset_reuse(): + # If TriContourSet returned from one tricontour(f) call is passed as first + # argument to another the underlying C++ contour generator will be reused. + x = [0.0, 0.5, 1.0] + y = [0.0, 1.0, 0.0] + z = [1.0, 2.0, 3.0] + fig, ax = plt.subplots() + tcs1 = ax.tricontourf(x, y, z) + tcs2 = ax.tricontour(x, y, z) + assert tcs2._contour_generator != tcs1._contour_generator + tcs3 = ax.tricontour(tcs1, z) + assert tcs3._contour_generator == tcs1._contour_generator + + +@check_figures_equal(extensions=['png']) +def test_triplot_with_ls(fig_test, fig_ref): + x = [0, 2, 1] + y = [0, 0, 1] + data = [[0, 1, 2]] + fig_test.subplots().triplot(x, y, data, ls='--') + fig_ref.subplots().triplot(x, y, data, linestyle='--') + + +def test_triplot_label(): + x = [0, 2, 1] + y = [0, 0, 1] + data = [[0, 1, 2]] + fig, ax = plt.subplots() + lines, markers = ax.triplot(x, y, data, label='label') + handles, labels = ax.get_legend_handles_labels() + assert labels == ['label'] + assert len(handles) == 1 + assert handles[0] is lines + + +def test_tricontour_path(): + x = [0, 4, 4, 0, 2] + y = [0, 0, 4, 4, 2] + triang = mtri.Triangulation(x, y) + _, ax = plt.subplots() + + # Line strip from boundary to boundary + cs = ax.tricontour(triang, [1, 0, 0, 0, 0], levels=[0.5]) + paths = cs.get_paths() + assert len(paths) == 1 + expected_vertices = [[2, 0], [1, 1], [0, 2]] + assert_array_almost_equal(paths[0].vertices, expected_vertices) + assert_array_equal(paths[0].codes, [1, 2, 2]) + assert_array_almost_equal( + paths[0].to_polygons(closed_only=False), [expected_vertices]) + + # Closed line loop inside domain + cs = ax.tricontour(triang, [0, 0, 0, 0, 1], levels=[0.5]) + paths = cs.get_paths() + assert len(paths) == 1 + expected_vertices = [[3, 1], [3, 3], [1, 3], [1, 1], [3, 1]] + assert_array_almost_equal(paths[0].vertices, expected_vertices) + assert_array_equal(paths[0].codes, [1, 2, 2, 2, 79]) + assert_array_almost_equal(paths[0].to_polygons(), [expected_vertices]) + + +def test_tricontourf_path(): + x = [0, 4, 4, 0, 2] + y = [0, 0, 4, 4, 2] + triang = mtri.Triangulation(x, y) + _, ax = plt.subplots() + + # Polygon inside domain + cs = ax.tricontourf(triang, [0, 0, 0, 0, 1], levels=[0.5, 1.5]) + paths = cs.get_paths() + assert len(paths) == 1 + expected_vertices = [[3, 1], [3, 3], [1, 3], [1, 1], [3, 1]] + assert_array_almost_equal(paths[0].vertices, expected_vertices) + assert_array_equal(paths[0].codes, [1, 2, 2, 2, 79]) + assert_array_almost_equal(paths[0].to_polygons(), [expected_vertices]) + + # Polygon following boundary and inside domain + cs = ax.tricontourf(triang, [1, 0, 0, 0, 0], levels=[0.5, 1.5]) + paths = cs.get_paths() + assert len(paths) == 1 + expected_vertices = [[2, 0], [1, 1], [0, 2], [0, 0], [2, 0]] + assert_array_almost_equal(paths[0].vertices, expected_vertices) + assert_array_equal(paths[0].codes, [1, 2, 2, 2, 79]) + assert_array_almost_equal(paths[0].to_polygons(), [expected_vertices]) + + # Polygon is outer boundary with hole + cs = ax.tricontourf(triang, [0, 0, 0, 0, 1], levels=[-0.5, 0.5]) + paths = cs.get_paths() + assert len(paths) == 1 + expected_vertices = [[0, 0], [4, 0], [4, 4], [0, 4], [0, 0], + [1, 1], [1, 3], [3, 3], [3, 1], [1, 1]] + assert_array_almost_equal(paths[0].vertices, expected_vertices) + assert_array_equal(paths[0].codes, [1, 2, 2, 2, 79, 1, 2, 2, 2, 79]) + assert_array_almost_equal(paths[0].to_polygons(), np.split(expected_vertices, [5])) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_type1font.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_type1font.py new file mode 100644 index 0000000000000000000000000000000000000000..9b8a2d1f07c6fd7124a4bc8349766c01c3a58435 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_type1font.py @@ -0,0 +1,160 @@ +import matplotlib._type1font as t1f +import os.path +import difflib +import pytest + + +def test_Type1Font(): + filename = os.path.join(os.path.dirname(__file__), 'cmr10.pfb') + font = t1f.Type1Font(filename) + slanted = font.transform({'slant': 1}) + condensed = font.transform({'extend': 0.5}) + with open(filename, 'rb') as fd: + rawdata = fd.read() + assert font.parts[0] == rawdata[0x0006:0x10c5] + assert font.parts[1] == rawdata[0x10cb:0x897f] + assert font.parts[2] == rawdata[0x8985:0x8ba6] + assert font.decrypted.startswith(b'dup\n/Private 18 dict dup begin') + assert font.decrypted.endswith(b'mark currentfile closefile\n') + assert slanted.decrypted.startswith(b'dup\n/Private 18 dict dup begin') + assert slanted.decrypted.endswith(b'mark currentfile closefile\n') + assert b'UniqueID 5000793' in font.parts[0] + assert b'UniqueID 5000793' in font.decrypted + assert font._pos['UniqueID'] == [(797, 818), (4483, 4504)] + + len0 = len(font.parts[0]) + for key in font._pos.keys(): + for pos0, pos1 in font._pos[key]: + if pos0 < len0: + data = font.parts[0][pos0:pos1] + else: + data = font.decrypted[pos0-len0:pos1-len0] + assert data.startswith(f'/{key}'.encode('ascii')) + assert {'FontType', 'FontMatrix', 'PaintType', 'ItalicAngle', 'RD' + } < set(font._pos.keys()) + + assert b'UniqueID 5000793' not in slanted.parts[0] + assert b'UniqueID 5000793' not in slanted.decrypted + assert 'UniqueID' not in slanted._pos + assert font.prop['Weight'] == 'Medium' + assert not font.prop['isFixedPitch'] + assert font.prop['ItalicAngle'] == 0 + assert slanted.prop['ItalicAngle'] == -45 + assert font.prop['Encoding'][5] == 'Pi' + assert isinstance(font.prop['CharStrings']['Pi'], bytes) + assert font._abbr['ND'] == 'ND' + + differ = difflib.Differ() + diff = list(differ.compare( + font.parts[0].decode('latin-1').splitlines(), + slanted.parts[0].decode('latin-1').splitlines())) + for line in ( + # Removes UniqueID + '- /UniqueID 5000793 def', + # Changes the font name + '- /FontName /CMR10 def', + '+ /FontName/CMR10_Slant_1000 def', + # Alters FontMatrix + '- /FontMatrix [0.001 0 0 0.001 0 0 ]readonly def', + '+ /FontMatrix [0.001 0 0.001 0.001 0 0] readonly def', + # Alters ItalicAngle + '- /ItalicAngle 0 def', + '+ /ItalicAngle -45.0 def'): + assert line in diff, 'diff to slanted font must contain %s' % line + + diff = list(differ.compare( + font.parts[0].decode('latin-1').splitlines(), + condensed.parts[0].decode('latin-1').splitlines())) + for line in ( + # Removes UniqueID + '- /UniqueID 5000793 def', + # Changes the font name + '- /FontName /CMR10 def', + '+ /FontName/CMR10_Extend_500 def', + # Alters FontMatrix + '- /FontMatrix [0.001 0 0 0.001 0 0 ]readonly def', + '+ /FontMatrix [0.0005 0 0 0.001 0 0] readonly def'): + assert line in diff, 'diff to condensed font must contain %s' % line + + +def test_Type1Font_2(): + filename = os.path.join(os.path.dirname(__file__), + 'Courier10PitchBT-Bold.pfb') + font = t1f.Type1Font(filename) + assert font.prop['Weight'] == 'Bold' + assert font.prop['isFixedPitch'] + assert font.prop['Encoding'][65] == 'A' # the font uses StandardEncoding + (pos0, pos1), = font._pos['Encoding'] + assert font.parts[0][pos0:pos1] == b'/Encoding StandardEncoding' + assert font._abbr['ND'] == '|-' + + +def test_tokenize(): + data = (b'1234/abc false -9.81 Foo <<[0 1 2]<0 1ef a\t>>>\n' + b'(string with(nested\t\\) par)ens\\\\)') + # 1 2 x 2 xx1 + # 1 and 2 are matching parens, x means escaped character + n, w, num, kw, d = 'name', 'whitespace', 'number', 'keyword', 'delimiter' + b, s = 'boolean', 'string' + correct = [ + (num, 1234), (n, 'abc'), (w, ' '), (b, False), (w, ' '), (num, -9.81), + (w, ' '), (kw, 'Foo'), (w, ' '), (d, '<<'), (d, '['), (num, 0), + (w, ' '), (num, 1), (w, ' '), (num, 2), (d, ']'), (s, b'\x01\xef\xa0'), + (d, '>>'), (w, '\n'), (s, 'string with(nested\t) par)ens\\') + ] + correct_no_ws = [x for x in correct if x[0] != w] + + def convert(tokens): + return [(t.kind, t.value()) for t in tokens] + + assert convert(t1f._tokenize(data, False)) == correct + assert convert(t1f._tokenize(data, True)) == correct_no_ws + + def bin_after(n): + tokens = t1f._tokenize(data, True) + result = [] + for _ in range(n): + result.append(next(tokens)) + result.append(tokens.send(10)) + return convert(result) + + for n in range(1, len(correct_no_ws)): + result = bin_after(n) + assert result[:-1] == correct_no_ws[:n] + assert result[-1][0] == 'binary' + assert isinstance(result[-1][1], bytes) + + +def test_tokenize_errors(): + with pytest.raises(ValueError): + list(t1f._tokenize(b'1234 (this (string) is unterminated\\)', True)) + with pytest.raises(ValueError): + list(t1f._tokenize(b'/Foo<01234', True)) + with pytest.raises(ValueError): + list(t1f._tokenize(b'/Foo<01234abcg>/Bar', True)) + + +def test_overprecision(): + # We used to output too many digits in FontMatrix entries and + # ItalicAngle, which could make Type-1 parsers unhappy. + filename = os.path.join(os.path.dirname(__file__), 'cmr10.pfb') + font = t1f.Type1Font(filename) + slanted = font.transform({'slant': .167}) + lines = slanted.parts[0].decode('ascii').splitlines() + matrix, = (line[line.index('[')+1:line.index(']')] + for line in lines if '/FontMatrix' in line) + angle, = (word + for line in lines if '/ItalicAngle' in line + for word in line.split() if word[0] in '-0123456789') + # the following used to include 0.00016700000000000002 + assert matrix == '0.001 0 0.000167 0.001 0 0' + # and here we had -9.48090361795083 + assert angle == '-9.4809' + + +def test_encrypt_decrypt_roundtrip(): + data = b'this is my plaintext \0\1\2\3' + encrypted = t1f.Type1Font._encrypt(data, 'eexec') + decrypted = t1f.Type1Font._decrypt(encrypted, 'eexec') + assert encrypted != decrypted + assert data == decrypted diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_units.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_units.py new file mode 100644 index 0000000000000000000000000000000000000000..cc71f685857e2e6da4677969837bd66c828cc17c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_units.py @@ -0,0 +1,353 @@ +from datetime import datetime, timezone, timedelta +import platform +from unittest.mock import MagicMock + +import matplotlib.pyplot as plt +from matplotlib.testing.decorators import check_figures_equal, image_comparison +import matplotlib.patches as mpatches +import matplotlib.units as munits +from matplotlib.category import StrCategoryConverter, UnitData +from matplotlib.dates import DateConverter +import numpy as np +import pytest + + +# Basic class that wraps numpy array and has units +class Quantity: + def __init__(self, data, units): + self.magnitude = data + self.units = units + + def to(self, new_units): + factors = {('hours', 'seconds'): 3600, ('minutes', 'hours'): 1 / 60, + ('minutes', 'seconds'): 60, ('feet', 'miles'): 1 / 5280., + ('feet', 'inches'): 12, ('miles', 'inches'): 12 * 5280} + if self.units != new_units: + mult = factors[self.units, new_units] + return Quantity(mult * self.magnitude, new_units) + else: + return Quantity(self.magnitude, self.units) + + def __copy__(self): + return Quantity(self.magnitude, self.units) + + def __getattr__(self, attr): + return getattr(self.magnitude, attr) + + def __getitem__(self, item): + if np.iterable(self.magnitude): + return Quantity(self.magnitude[item], self.units) + else: + return Quantity(self.magnitude, self.units) + + def __array__(self): + return np.asarray(self.magnitude) + + +@pytest.fixture +def quantity_converter(): + # Create an instance of the conversion interface and + # mock so we can check methods called + qc = munits.ConversionInterface() + + def convert(value, unit, axis): + if hasattr(value, 'units'): + return value.to(unit).magnitude + elif np.iterable(value): + try: + return [v.to(unit).magnitude for v in value] + except AttributeError: + return [Quantity(v, axis.get_units()).to(unit).magnitude + for v in value] + else: + return Quantity(value, axis.get_units()).to(unit).magnitude + + def default_units(value, axis): + if hasattr(value, 'units'): + return value.units + elif np.iterable(value): + for v in value: + if hasattr(v, 'units'): + return v.units + return None + + qc.convert = MagicMock(side_effect=convert) + qc.axisinfo = MagicMock(side_effect=lambda u, a: + munits.AxisInfo(label=u, default_limits=(0, 100))) + qc.default_units = MagicMock(side_effect=default_units) + return qc + + +# Tests that the conversion machinery works properly for classes that +# work as a facade over numpy arrays (like pint) +@image_comparison(['plot_pint.png'], style='mpl20', + tol=0 if platform.machine() == 'x86_64' else 0.03) +def test_numpy_facade(quantity_converter): + # use former defaults to match existing baseline image + plt.rcParams['axes.formatter.limits'] = -7, 7 + + # Register the class + munits.registry[Quantity] = quantity_converter + + # Simple test + y = Quantity(np.linspace(0, 30), 'miles') + x = Quantity(np.linspace(0, 5), 'hours') + + fig, ax = plt.subplots() + fig.subplots_adjust(left=0.15) # Make space for label + ax.plot(x, y, 'tab:blue') + ax.axhline(Quantity(26400, 'feet'), color='tab:red') + ax.axvline(Quantity(120, 'minutes'), color='tab:green') + ax.yaxis.set_units('inches') + ax.xaxis.set_units('seconds') + + assert quantity_converter.convert.called + assert quantity_converter.axisinfo.called + assert quantity_converter.default_units.called + + +# Tests gh-8908 +@image_comparison(['plot_masked_units.png'], remove_text=True, style='mpl20', + tol=0 if platform.machine() == 'x86_64' else 0.02) +def test_plot_masked_units(): + data = np.linspace(-5, 5) + data_masked = np.ma.array(data, mask=(data > -2) & (data < 2)) + data_masked_units = Quantity(data_masked, 'meters') + + fig, ax = plt.subplots() + ax.plot(data_masked_units) + + +def test_empty_set_limits_with_units(quantity_converter): + # Register the class + munits.registry[Quantity] = quantity_converter + + fig, ax = plt.subplots() + ax.set_xlim(Quantity(-1, 'meters'), Quantity(6, 'meters')) + ax.set_ylim(Quantity(-1, 'hours'), Quantity(16, 'hours')) + + +@image_comparison(['jpl_bar_units.png'], + savefig_kwarg={'dpi': 120}, style='mpl20') +def test_jpl_bar_units(): + import matplotlib.testing.jpl_units as units + units.register() + + day = units.Duration("ET", 24.0 * 60.0 * 60.0) + x = [0 * units.km, 1 * units.km, 2 * units.km] + w = [1 * day, 2 * day, 3 * day] + b = units.Epoch("ET", dt=datetime(2009, 4, 26)) + fig, ax = plt.subplots() + ax.bar(x, w, bottom=b) + ax.set_ylim([b - 1 * day, b + w[-1] + (1.001) * day]) + + +@image_comparison(['jpl_barh_units.png'], + savefig_kwarg={'dpi': 120}, style='mpl20') +def test_jpl_barh_units(): + import matplotlib.testing.jpl_units as units + units.register() + + day = units.Duration("ET", 24.0 * 60.0 * 60.0) + x = [0 * units.km, 1 * units.km, 2 * units.km] + w = [1 * day, 2 * day, 3 * day] + b = units.Epoch("ET", dt=datetime(2009, 4, 26)) + + fig, ax = plt.subplots() + ax.barh(x, w, left=b) + ax.set_xlim([b - 1 * day, b + w[-1] + (1.001) * day]) + + +def test_jpl_datetime_units_consistent(): + import matplotlib.testing.jpl_units as units + units.register() + + dt = datetime(2009, 4, 26) + jpl = units.Epoch("ET", dt=dt) + dt_conv = munits.registry.get_converter(dt).convert(dt, None, None) + jpl_conv = munits.registry.get_converter(jpl).convert(jpl, None, None) + assert dt_conv == jpl_conv + + +def test_empty_arrays(): + # Check that plotting an empty array with a dtype works + plt.scatter(np.array([], dtype='datetime64[ns]'), np.array([])) + + +def test_scatter_element0_masked(): + times = np.arange('2005-02', '2005-03', dtype='datetime64[D]') + y = np.arange(len(times), dtype=float) + y[0] = np.nan + fig, ax = plt.subplots() + ax.scatter(times, y) + fig.canvas.draw() + + +def test_errorbar_mixed_units(): + x = np.arange(10) + y = [datetime(2020, 5, i * 2 + 1) for i in x] + fig, ax = plt.subplots() + ax.errorbar(x, y, timedelta(days=0.5)) + fig.canvas.draw() + + +@check_figures_equal(extensions=["png"]) +def test_subclass(fig_test, fig_ref): + class subdate(datetime): + pass + + fig_test.subplots().plot(subdate(2000, 1, 1), 0, "o") + fig_ref.subplots().plot(datetime(2000, 1, 1), 0, "o") + + +def test_shared_axis_quantity(quantity_converter): + munits.registry[Quantity] = quantity_converter + x = Quantity(np.linspace(0, 1, 10), "hours") + y1 = Quantity(np.linspace(1, 2, 10), "feet") + y2 = Quantity(np.linspace(3, 4, 10), "feet") + fig, (ax1, ax2) = plt.subplots(2, 1, sharex='all', sharey='all') + ax1.plot(x, y1) + ax2.plot(x, y2) + assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "hours" + assert ax2.yaxis.get_units() == ax2.yaxis.get_units() == "feet" + ax1.xaxis.set_units("seconds") + ax2.yaxis.set_units("inches") + assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "seconds" + assert ax1.yaxis.get_units() == ax2.yaxis.get_units() == "inches" + + +def test_shared_axis_datetime(): + # datetime uses dates.DateConverter + y1 = [datetime(2020, i, 1, tzinfo=timezone.utc) for i in range(1, 13)] + y2 = [datetime(2021, i, 1, tzinfo=timezone.utc) for i in range(1, 13)] + fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True) + ax1.plot(y1) + ax2.plot(y2) + ax1.yaxis.set_units(timezone(timedelta(hours=5))) + assert ax2.yaxis.units == timezone(timedelta(hours=5)) + + +def test_shared_axis_categorical(): + # str uses category.StrCategoryConverter + d1 = {"a": 1, "b": 2} + d2 = {"a": 3, "b": 4} + fig, (ax1, ax2) = plt.subplots(1, 2, sharex=True, sharey=True) + ax1.plot(d1.keys(), d1.values()) + ax2.plot(d2.keys(), d2.values()) + ax1.xaxis.set_units(UnitData(["c", "d"])) + assert "c" in ax2.xaxis.get_units()._mapping.keys() + + +def test_explicit_converter(): + d1 = {"a": 1, "b": 2} + str_cat_converter = StrCategoryConverter() + str_cat_converter_2 = StrCategoryConverter() + date_converter = DateConverter() + + # Explicit is set + fig1, ax1 = plt.subplots() + ax1.xaxis.set_converter(str_cat_converter) + assert ax1.xaxis.get_converter() == str_cat_converter + # Explicit not overridden by implicit + ax1.plot(d1.keys(), d1.values()) + assert ax1.xaxis.get_converter() == str_cat_converter + # No error when called twice with equivalent input + ax1.xaxis.set_converter(str_cat_converter) + # Error when explicit called twice + with pytest.raises(RuntimeError): + ax1.xaxis.set_converter(str_cat_converter_2) + + fig2, ax2 = plt.subplots() + ax2.plot(d1.keys(), d1.values()) + + # No error when equivalent type is used + ax2.xaxis.set_converter(str_cat_converter) + + fig3, ax3 = plt.subplots() + ax3.plot(d1.keys(), d1.values()) + + # Warn when implicit overridden + with pytest.warns(): + ax3.xaxis.set_converter(date_converter) + + +def test_empty_default_limits(quantity_converter): + munits.registry[Quantity] = quantity_converter + fig, ax1 = plt.subplots() + ax1.xaxis.update_units(Quantity([10], "miles")) + fig.draw_without_rendering() + assert ax1.get_xlim() == (0, 100) + ax1.yaxis.update_units(Quantity([10], "miles")) + fig.draw_without_rendering() + assert ax1.get_ylim() == (0, 100) + + fig, ax = plt.subplots() + ax.axhline(30) + ax.plot(Quantity(np.arange(0, 3), "miles"), + Quantity(np.arange(0, 6, 2), "feet")) + fig.draw_without_rendering() + assert ax.get_xlim() == (0, 2) + assert ax.get_ylim() == (0, 30) + + fig, ax = plt.subplots() + ax.axvline(30) + ax.plot(Quantity(np.arange(0, 3), "miles"), + Quantity(np.arange(0, 6, 2), "feet")) + fig.draw_without_rendering() + assert ax.get_xlim() == (0, 30) + assert ax.get_ylim() == (0, 4) + + fig, ax = plt.subplots() + ax.xaxis.update_units(Quantity([10], "miles")) + ax.axhline(30) + fig.draw_without_rendering() + assert ax.get_xlim() == (0, 100) + assert ax.get_ylim() == (28.5, 31.5) + + fig, ax = plt.subplots() + ax.yaxis.update_units(Quantity([10], "miles")) + ax.axvline(30) + fig.draw_without_rendering() + assert ax.get_ylim() == (0, 100) + assert ax.get_xlim() == (28.5, 31.5) + + +# test array-like objects... +class Kernel: + def __init__(self, array): + self._array = np.asanyarray(array) + + def __array__(self, dtype=None, copy=None): + if dtype is not None and dtype != self._array.dtype: + if copy is not None and not copy: + raise ValueError( + f"Converting array from {self._array.dtype} to " + f"{dtype} requires a copy" + ) + + arr = np.asarray(self._array, dtype=dtype) + return (arr if not copy else np.copy(arr)) + + @property + def shape(self): + return self._array.shape + + +def test_plot_kernel(): + # just a smoketest that fail + kernel = Kernel([1, 2, 3, 4, 5]) + plt.plot(kernel) + + +def test_connection_patch_units(pd): + # tests that this doesn't raise an error + fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(10, 5)) + x = pd.Timestamp('2017-01-01T12') + ax1.axvline(x) + y = "test test" + ax2.axhline(y) + arr = mpatches.ConnectionPatch((x, 0), (0, y), + coordsA='data', coordsB='data', + axesA=ax1, axesB=ax2) + fig.add_artist(arr) + fig.draw_without_rendering() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_usetex.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_usetex.py new file mode 100644 index 0000000000000000000000000000000000000000..342face4504fbc444c530fa08335cf19b3c98e6f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_usetex.py @@ -0,0 +1,187 @@ +from tempfile import TemporaryFile + +import numpy as np +from packaging.version import parse as parse_version +import pytest + +import matplotlib as mpl +from matplotlib import dviread +from matplotlib.testing import _has_tex_package +from matplotlib.testing.decorators import check_figures_equal, image_comparison +from matplotlib.testing._markers import needs_usetex +import matplotlib.pyplot as plt + + +pytestmark = needs_usetex + + +@image_comparison( + baseline_images=['test_usetex'], + extensions=['pdf', 'png'], + style="mpl20") +def test_usetex(): + mpl.rcParams['text.usetex'] = True + fig, ax = plt.subplots() + kwargs = {"verticalalignment": "baseline", "size": 24, + "bbox": dict(pad=0, edgecolor="k", facecolor="none")} + ax.text(0.2, 0.7, + # the \LaTeX macro exercises character sizing and placement, + # \left[ ... \right\} draw some variable-height characters, + # \sqrt and \frac draw horizontal rules, \mathrm changes the font + r'\LaTeX\ $\left[\int\limits_e^{2e}' + r'\sqrt\frac{\log^3 x}{x}\,\mathrm{d}x \right\}$', + **kwargs) + ax.text(0.2, 0.3, "lg", **kwargs) + ax.text(0.4, 0.3, r"$\frac{1}{2}\pi$", **kwargs) + ax.text(0.6, 0.3, "$p^{3^A}$", **kwargs) + ax.text(0.8, 0.3, "$p_{3_2}$", **kwargs) + for x in {t.get_position()[0] for t in ax.texts}: + ax.axvline(x) + for y in {t.get_position()[1] for t in ax.texts}: + ax.axhline(y) + ax.set_axis_off() + + +@check_figures_equal() +def test_empty(fig_test, fig_ref): + mpl.rcParams['text.usetex'] = True + fig_test.text(.5, .5, "% a comment") + + +@check_figures_equal() +def test_unicode_minus(fig_test, fig_ref): + mpl.rcParams['text.usetex'] = True + fig_test.text(.5, .5, "$-$") + fig_ref.text(.5, .5, "\N{MINUS SIGN}") + + +def test_mathdefault(): + plt.rcParams["axes.formatter.use_mathtext"] = True + fig = plt.figure() + fig.add_subplot().set_xlim(-1, 1) + # Check that \mathdefault commands generated by tickers don't cause + # problems when later switching usetex on. + mpl.rcParams['text.usetex'] = True + fig.canvas.draw() + + +@image_comparison(['eqnarray.png']) +def test_multiline_eqnarray(): + text = ( + r'\begin{eqnarray*}' + r'foo\\' + r'bar\\' + r'baz\\' + r'\end{eqnarray*}' + ) + + fig = plt.figure(figsize=(1, 1)) + fig.text(0.5, 0.5, text, usetex=True, + horizontalalignment='center', verticalalignment='center') + + +@pytest.mark.parametrize("fontsize", [8, 10, 12]) +def test_minus_no_descent(fontsize): + # Test special-casing of minus descent in DviFont._height_depth_of, by + # checking that overdrawing a 1 and a -1 results in an overall height + # equivalent to drawing either of them separately. + mpl.style.use("mpl20") + mpl.rcParams['font.size'] = fontsize + heights = {} + fig = plt.figure() + for vals in [(1,), (-1,), (-1, 1)]: + fig.clear() + for x in vals: + fig.text(.5, .5, f"${x}$", usetex=True) + fig.canvas.draw() + # The following counts the number of non-fully-blank pixel rows. + heights[vals] = ((np.array(fig.canvas.buffer_rgba())[..., 0] != 255) + .any(axis=1).sum()) + assert len({*heights.values()}) == 1 + + +@pytest.mark.parametrize('pkg', ['xcolor', 'chemformula']) +def test_usetex_packages(pkg): + if not _has_tex_package(pkg): + pytest.skip(f'{pkg} is not available') + mpl.rcParams['text.usetex'] = True + + fig = plt.figure() + text = fig.text(0.5, 0.5, "Some text 0123456789") + fig.canvas.draw() + + mpl.rcParams['text.latex.preamble'] = ( + r'\PassOptionsToPackage{dvipsnames}{xcolor}\usepackage{%s}' % pkg) + fig = plt.figure() + text2 = fig.text(0.5, 0.5, "Some text 0123456789") + fig.canvas.draw() + np.testing.assert_array_equal(text2.get_window_extent(), + text.get_window_extent()) + + +@pytest.mark.parametrize( + "preamble", + [r"\usepackage[full]{textcomp}", r"\usepackage{underscore}"], +) +def test_latex_pkg_already_loaded(preamble): + plt.rcParams["text.latex.preamble"] = preamble + fig = plt.figure() + fig.text(.5, .5, "hello, world", usetex=True) + fig.canvas.draw() + + +def test_usetex_with_underscore(): + plt.rcParams["text.usetex"] = True + df = {"a_b": range(5)[::-1], "c": range(5)} + fig, ax = plt.subplots() + ax.plot("c", "a_b", data=df) + ax.legend() + ax.text(0, 0, "foo_bar", usetex=True) + plt.draw() + + +@pytest.mark.flaky(reruns=3) # Tends to hit a TeX cache lock on AppVeyor. +@pytest.mark.parametrize("fmt", ["pdf", "svg"]) +def test_missing_psfont(fmt, monkeypatch): + """An error is raised if a TeX font lacks a Type-1 equivalent""" + monkeypatch.setattr( + dviread.PsfontsMap, '__getitem__', + lambda self, k: dviread.PsFont( + texname=b'texfont', psname=b'Some Font', + effects=None, encoding=None, filename=None)) + mpl.rcParams['text.usetex'] = True + fig, ax = plt.subplots() + ax.text(0.5, 0.5, 'hello') + with TemporaryFile() as tmpfile, pytest.raises(ValueError): + fig.savefig(tmpfile, format=fmt) + + +try: + _old_gs_version = mpl._get_executable_info('gs').version < parse_version('9.55') +except mpl.ExecutableNotFoundError: + _old_gs_version = True + + +@image_comparison(baseline_images=['rotation'], extensions=['eps', 'pdf', 'png', 'svg'], + style='mpl20', tol=3.91 if _old_gs_version else 0) +def test_rotation(): + mpl.rcParams['text.usetex'] = True + + fig = plt.figure() + ax = fig.add_axes([0, 0, 1, 1]) + ax.set(xlim=[-0.5, 5], xticks=[], ylim=[-0.5, 3], yticks=[], frame_on=False) + + text = {val: val[0] for val in ['top', 'center', 'bottom', 'left', 'right']} + text['baseline'] = 'B' + text['center_baseline'] = 'C' + + for i, va in enumerate(['top', 'center', 'bottom', 'baseline', 'center_baseline']): + for j, ha in enumerate(['left', 'center', 'right']): + for k, angle in enumerate([0, 90, 180, 270]): + k //= 2 + x = i + k / 2 + y = j + k / 2 + ax.plot(x, y, '+', c=f'C{k}', markersize=20, markeredgewidth=0.5) + # 'My' checks full height letters plus descenders. + ax.text(x, y, f"$\\mathrm{{My {text[ha]}{text[va]} {angle}}}$", + rotation=angle, horizontalalignment=ha, verticalalignment=va) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_widgets.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_widgets.py new file mode 100644 index 0000000000000000000000000000000000000000..1b42d1ca7c2dc07b086263c3161b428cfb8a275a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tests/test_widgets.py @@ -0,0 +1,1759 @@ +import functools +import io +import operator +from unittest import mock + +from matplotlib.backend_bases import MouseEvent +import matplotlib.colors as mcolors +import matplotlib.widgets as widgets +import matplotlib.pyplot as plt +from matplotlib.testing.decorators import check_figures_equal, image_comparison +from matplotlib.testing.widgets import (click_and_drag, do_event, get_ax, + mock_event, noop) + +import numpy as np +from numpy.testing import assert_allclose + +import pytest + + +@pytest.fixture +def ax(): + return get_ax() + + +def test_save_blitted_widget_as_pdf(): + from matplotlib.widgets import CheckButtons, RadioButtons + from matplotlib.cbook import _get_running_interactive_framework + if _get_running_interactive_framework() not in ['headless', None]: + pytest.xfail("Callback exceptions are not raised otherwise.") + + fig, ax = plt.subplots( + nrows=2, ncols=2, figsize=(5, 2), width_ratios=[1, 2] + ) + default_rb = RadioButtons(ax[0, 0], ['Apples', 'Oranges']) + styled_rb = RadioButtons( + ax[0, 1], ['Apples', 'Oranges'], + label_props={'color': ['red', 'orange'], + 'fontsize': [16, 20]}, + radio_props={'edgecolor': ['red', 'orange'], + 'facecolor': ['mistyrose', 'peachpuff']} + ) + + default_cb = CheckButtons(ax[1, 0], ['Apples', 'Oranges'], + actives=[True, True]) + styled_cb = CheckButtons( + ax[1, 1], ['Apples', 'Oranges'], + actives=[True, True], + label_props={'color': ['red', 'orange'], + 'fontsize': [16, 20]}, + frame_props={'edgecolor': ['red', 'orange'], + 'facecolor': ['mistyrose', 'peachpuff']}, + check_props={'color': ['darkred', 'darkorange']} + ) + + ax[0, 0].set_title('Default') + ax[0, 1].set_title('Stylized') + # force an Agg render + fig.canvas.draw() + # force a pdf save + with io.BytesIO() as result_after: + fig.savefig(result_after, format='pdf') + + +@pytest.mark.parametrize('kwargs', [ + dict(), + dict(useblit=True, button=1), + dict(minspanx=10, minspany=10, spancoords='pixels'), + dict(props=dict(fill=True)), +]) +def test_rectangle_selector(ax, kwargs): + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.RectangleSelector(ax, onselect=onselect, **kwargs) + do_event(tool, 'press', xdata=100, ydata=100, button=1) + do_event(tool, 'onmove', xdata=199, ydata=199, button=1) + + # purposely drag outside of axis for release + do_event(tool, 'release', xdata=250, ydata=250, button=1) + + if kwargs.get('drawtype', None) not in ['line', 'none']: + assert_allclose(tool.geometry, + [[100., 100, 199, 199, 100], + [100, 199, 199, 100, 100]], + err_msg=tool.geometry) + + onselect.assert_called_once() + (epress, erelease), kwargs = onselect.call_args + assert epress.xdata == 100 + assert epress.ydata == 100 + assert erelease.xdata == 199 + assert erelease.ydata == 199 + assert kwargs == {} + + +@pytest.mark.parametrize('spancoords', ['data', 'pixels']) +@pytest.mark.parametrize('minspanx, x1', [[0, 10], [1, 10.5], [1, 11]]) +@pytest.mark.parametrize('minspany, y1', [[0, 10], [1, 10.5], [1, 11]]) +def test_rectangle_minspan(ax, spancoords, minspanx, x1, minspany, y1): + + onselect = mock.Mock(spec=noop, return_value=None) + + x0, y0 = (10, 10) + if spancoords == 'pixels': + minspanx, minspany = (ax.transData.transform((x1, y1)) - + ax.transData.transform((x0, y0))) + + tool = widgets.RectangleSelector(ax, onselect=onselect, interactive=True, + spancoords=spancoords, + minspanx=minspanx, minspany=minspany) + # Too small to create a selector + click_and_drag(tool, start=(x0, x1), end=(y0, y1)) + assert not tool._selection_completed + onselect.assert_not_called() + + click_and_drag(tool, start=(20, 20), end=(30, 30)) + assert tool._selection_completed + onselect.assert_called_once() + + # Too small to create a selector. Should clear existing selector, and + # trigger onselect because there was a preexisting selector + onselect.reset_mock() + click_and_drag(tool, start=(x0, y0), end=(x1, y1)) + assert not tool._selection_completed + onselect.assert_called_once() + (epress, erelease), kwargs = onselect.call_args + assert epress.xdata == x0 + assert epress.ydata == y0 + assert erelease.xdata == x1 + assert erelease.ydata == y1 + assert kwargs == {} + + +@pytest.mark.parametrize('drag_from_anywhere, new_center', + [[True, (60, 75)], + [False, (30, 20)]]) +def test_rectangle_drag(ax, drag_from_anywhere, new_center): + tool = widgets.RectangleSelector(ax, interactive=True, + drag_from_anywhere=drag_from_anywhere) + # Create rectangle + click_and_drag(tool, start=(0, 10), end=(100, 120)) + assert tool.center == (50, 65) + # Drag inside rectangle, but away from centre handle + # + # If drag_from_anywhere == True, this will move the rectangle by (10, 10), + # giving it a new center of (60, 75) + # + # If drag_from_anywhere == False, this will create a new rectangle with + # center (30, 20) + click_and_drag(tool, start=(25, 15), end=(35, 25)) + assert tool.center == new_center + # Check that in both cases, dragging outside the rectangle draws a new + # rectangle + click_and_drag(tool, start=(175, 185), end=(185, 195)) + assert tool.center == (180, 190) + + +def test_rectangle_selector_set_props_handle_props(ax): + tool = widgets.RectangleSelector(ax, interactive=True, + props=dict(facecolor='b', alpha=0.2), + handle_props=dict(alpha=0.5)) + # Create rectangle + click_and_drag(tool, start=(0, 10), end=(100, 120)) + + artist = tool._selection_artist + assert artist.get_facecolor() == mcolors.to_rgba('b', alpha=0.2) + tool.set_props(facecolor='r', alpha=0.3) + assert artist.get_facecolor() == mcolors.to_rgba('r', alpha=0.3) + + for artist in tool._handles_artists: + assert artist.get_markeredgecolor() == 'black' + assert artist.get_alpha() == 0.5 + tool.set_handle_props(markeredgecolor='r', alpha=0.3) + for artist in tool._handles_artists: + assert artist.get_markeredgecolor() == 'r' + assert artist.get_alpha() == 0.3 + + +def test_rectangle_resize(ax): + tool = widgets.RectangleSelector(ax, interactive=True) + # Create rectangle + click_and_drag(tool, start=(0, 10), end=(100, 120)) + assert tool.extents == (0.0, 100.0, 10.0, 120.0) + + # resize NE handle + extents = tool.extents + xdata, ydata = extents[1], extents[3] + xdata_new, ydata_new = xdata + 10, ydata + 5 + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert tool.extents == (extents[0], xdata_new, extents[2], ydata_new) + + # resize E handle + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdata_new, ydata_new = xdata + 10, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert tool.extents == (extents[0], xdata_new, extents[2], extents[3]) + + # resize W handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdata_new, ydata_new = xdata + 15, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert tool.extents == (xdata_new, extents[1], extents[2], extents[3]) + + # resize SW handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + xdata_new, ydata_new = xdata + 20, ydata + 25 + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert tool.extents == (xdata_new, extents[1], ydata_new, extents[3]) + + +def test_rectangle_add_state(ax): + tool = widgets.RectangleSelector(ax, interactive=True) + # Create rectangle + click_and_drag(tool, start=(70, 65), end=(125, 130)) + + with pytest.raises(ValueError): + tool.add_state('unsupported_state') + + with pytest.raises(ValueError): + tool.add_state('clear') + tool.add_state('move') + tool.add_state('square') + tool.add_state('center') + + +@pytest.mark.parametrize('add_state', [True, False]) +def test_rectangle_resize_center(ax, add_state): + tool = widgets.RectangleSelector(ax, interactive=True) + # Create rectangle + click_and_drag(tool, start=(70, 65), end=(125, 130)) + assert tool.extents == (70.0, 125.0, 65.0, 130.0) + + if add_state: + tool.add_state('center') + use_key = None + else: + use_key = 'control' + + # resize NE handle + extents = tool.extents + xdata, ydata = extents[1], extents[3] + xdiff, ydiff = 10, 5 + xdata_new, ydata_new = xdata + xdiff, ydata + ydiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0] - xdiff, xdata_new, + extents[2] - ydiff, ydata_new) + + # resize E handle + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = 10 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0] - xdiff, xdata_new, + extents[2], extents[3]) + + # resize E handle negative diff + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = -20 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0] - xdiff, xdata_new, + extents[2], extents[3]) + + # resize W handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = 15 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (xdata_new, extents[1] - xdiff, + extents[2], extents[3]) + + # resize W handle negative diff + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = -25 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (xdata_new, extents[1] - xdiff, + extents[2], extents[3]) + + # resize SW handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + xdiff, ydiff = 20, 25 + xdata_new, ydata_new = xdata + xdiff, ydata + ydiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (xdata_new, extents[1] - xdiff, + ydata_new, extents[3] - ydiff) + + +@pytest.mark.parametrize('add_state', [True, False]) +def test_rectangle_resize_square(ax, add_state): + tool = widgets.RectangleSelector(ax, interactive=True) + # Create rectangle + click_and_drag(tool, start=(70, 65), end=(120, 115)) + assert tool.extents == (70.0, 120.0, 65.0, 115.0) + + if add_state: + tool.add_state('square') + use_key = None + else: + use_key = 'shift' + + # resize NE handle + extents = tool.extents + xdata, ydata = extents[1], extents[3] + xdiff, ydiff = 10, 5 + xdata_new, ydata_new = xdata + xdiff, ydata + ydiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0], xdata_new, + extents[2], extents[3] + xdiff) + + # resize E handle + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = 10 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0], xdata_new, + extents[2], extents[3] + xdiff) + + # resize E handle negative diff + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = -20 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0], xdata_new, + extents[2], extents[3] + xdiff) + + # resize W handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = 15 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (xdata_new, extents[1], + extents[2], extents[3] - xdiff) + + # resize W handle negative diff + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = -25 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (xdata_new, extents[1], + extents[2], extents[3] - xdiff) + + # resize SW handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + xdiff, ydiff = 20, 25 + xdata_new, ydata_new = xdata + xdiff, ydata + ydiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new), + key=use_key) + assert tool.extents == (extents[0] + ydiff, extents[1], + ydata_new, extents[3]) + + +def test_rectangle_resize_square_center(ax): + tool = widgets.RectangleSelector(ax, interactive=True) + # Create rectangle + click_and_drag(tool, start=(70, 65), end=(120, 115)) + tool.add_state('square') + tool.add_state('center') + assert_allclose(tool.extents, (70.0, 120.0, 65.0, 115.0)) + + # resize NE handle + extents = tool.extents + xdata, ydata = extents[1], extents[3] + xdiff, ydiff = 10, 5 + xdata_new, ydata_new = xdata + xdiff, ydata + ydiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, (extents[0] - xdiff, xdata_new, + extents[2] - xdiff, extents[3] + xdiff)) + + # resize E handle + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = 10 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, (extents[0] - xdiff, xdata_new, + extents[2] - xdiff, extents[3] + xdiff)) + + # resize E handle negative diff + extents = tool.extents + xdata, ydata = extents[1], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = -20 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, (extents[0] - xdiff, xdata_new, + extents[2] - xdiff, extents[3] + xdiff)) + + # resize W handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = 5 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, (xdata_new, extents[1] - xdiff, + extents[2] + xdiff, extents[3] - xdiff)) + + # resize W handle negative diff + extents = tool.extents + xdata, ydata = extents[0], extents[2] + (extents[3] - extents[2]) / 2 + xdiff = -25 + xdata_new, ydata_new = xdata + xdiff, ydata + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, (xdata_new, extents[1] - xdiff, + extents[2] + xdiff, extents[3] - xdiff)) + + # resize SW handle + extents = tool.extents + xdata, ydata = extents[0], extents[2] + xdiff, ydiff = 20, 25 + xdata_new, ydata_new = xdata + xdiff, ydata + ydiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, (extents[0] + ydiff, extents[1] - ydiff, + ydata_new, extents[3] - ydiff)) + + +@pytest.mark.parametrize('selector_class', + [widgets.RectangleSelector, widgets.EllipseSelector]) +def test_rectangle_rotate(ax, selector_class): + tool = selector_class(ax, interactive=True) + # Draw rectangle + click_and_drag(tool, start=(100, 100), end=(130, 140)) + assert tool.extents == (100, 130, 100, 140) + assert len(tool._state) == 0 + + # Rotate anticlockwise using top-right corner + do_event(tool, 'on_key_press', key='r') + assert tool._state == {'rotate'} + assert len(tool._state) == 1 + click_and_drag(tool, start=(130, 140), end=(120, 145)) + do_event(tool, 'on_key_press', key='r') + assert len(tool._state) == 0 + # Extents shouldn't change (as shape of rectangle hasn't changed) + assert tool.extents == (100, 130, 100, 140) + assert_allclose(tool.rotation, 25.56, atol=0.01) + tool.rotation = 45 + assert tool.rotation == 45 + # Corners should move + assert_allclose(tool.corners, + np.array([[118.53, 139.75, 111.46, 90.25], + [95.25, 116.46, 144.75, 123.54]]), atol=0.01) + + # Scale using top-right corner + click_and_drag(tool, start=(110, 145), end=(110, 160)) + assert_allclose(tool.extents, (100, 139.75, 100, 151.82), atol=0.01) + + if selector_class == widgets.RectangleSelector: + with pytest.raises(ValueError): + tool._selection_artist.rotation_point = 'unvalid_value' + + +def test_rectangle_add_remove_set(ax): + tool = widgets.RectangleSelector(ax, interactive=True) + # Draw rectangle + click_and_drag(tool, start=(100, 100), end=(130, 140)) + assert tool.extents == (100, 130, 100, 140) + assert len(tool._state) == 0 + for state in ['rotate', 'square', 'center']: + tool.add_state(state) + assert len(tool._state) == 1 + tool.remove_state(state) + assert len(tool._state) == 0 + + +@pytest.mark.parametrize('use_data_coordinates', [False, True]) +def test_rectangle_resize_square_center_aspect(ax, use_data_coordinates): + ax.set_aspect(0.8) + + tool = widgets.RectangleSelector(ax, interactive=True, + use_data_coordinates=use_data_coordinates) + # Create rectangle + click_and_drag(tool, start=(70, 65), end=(120, 115)) + assert tool.extents == (70.0, 120.0, 65.0, 115.0) + tool.add_state('square') + tool.add_state('center') + + if use_data_coordinates: + # resize E handle + extents = tool.extents + xdata, ydata, width = extents[1], extents[3], extents[1] - extents[0] + xdiff, ycenter = 10, extents[2] + (extents[3] - extents[2]) / 2 + xdata_new, ydata_new = xdata + xdiff, ydata + ychange = width / 2 + xdiff + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, [extents[0] - xdiff, xdata_new, + ycenter - ychange, ycenter + ychange]) + else: + # resize E handle + extents = tool.extents + xdata, ydata = extents[1], extents[3] + xdiff = 10 + xdata_new, ydata_new = xdata + xdiff, ydata + ychange = xdiff * 1 / tool._aspect_ratio_correction + click_and_drag(tool, start=(xdata, ydata), end=(xdata_new, ydata_new)) + assert_allclose(tool.extents, [extents[0] - xdiff, xdata_new, + 46.25, 133.75]) + + +def test_ellipse(ax): + """For ellipse, test out the key modifiers""" + tool = widgets.EllipseSelector(ax, grab_range=10, interactive=True) + tool.extents = (100, 150, 100, 150) + + # drag the rectangle + click_and_drag(tool, start=(125, 125), end=(145, 145)) + assert tool.extents == (120, 170, 120, 170) + + # create from center + click_and_drag(tool, start=(100, 100), end=(125, 125), key='control') + assert tool.extents == (75, 125, 75, 125) + + # create a square + click_and_drag(tool, start=(10, 10), end=(35, 30), key='shift') + extents = [int(e) for e in tool.extents] + assert extents == [10, 35, 10, 35] + + # create a square from center + click_and_drag(tool, start=(100, 100), end=(125, 130), key='ctrl+shift') + extents = [int(e) for e in tool.extents] + assert extents == [70, 130, 70, 130] + + assert tool.geometry.shape == (2, 73) + assert_allclose(tool.geometry[:, 0], [70., 100]) + + +def test_rectangle_handles(ax): + tool = widgets.RectangleSelector(ax, grab_range=10, interactive=True, + handle_props={'markerfacecolor': 'r', + 'markeredgecolor': 'b'}) + tool.extents = (100, 150, 100, 150) + + assert_allclose(tool.corners, ((100, 150, 150, 100), (100, 100, 150, 150))) + assert tool.extents == (100, 150, 100, 150) + assert_allclose(tool.edge_centers, + ((100, 125.0, 150, 125.0), (125.0, 100, 125.0, 150))) + assert tool.extents == (100, 150, 100, 150) + + # grab a corner and move it + click_and_drag(tool, start=(100, 100), end=(120, 120)) + assert tool.extents == (120, 150, 120, 150) + + # grab the center and move it + click_and_drag(tool, start=(132, 132), end=(120, 120)) + assert tool.extents == (108, 138, 108, 138) + + # create a new rectangle + click_and_drag(tool, start=(10, 10), end=(100, 100)) + assert tool.extents == (10, 100, 10, 100) + + # Check that marker_props worked. + assert mcolors.same_color( + tool._corner_handles.artists[0].get_markerfacecolor(), 'r') + assert mcolors.same_color( + tool._corner_handles.artists[0].get_markeredgecolor(), 'b') + + +@pytest.mark.parametrize('interactive', [True, False]) +def test_rectangle_selector_onselect(ax, interactive): + # check when press and release events take place at the same position + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.RectangleSelector(ax, onselect=onselect, interactive=interactive) + # move outside of axis + click_and_drag(tool, start=(100, 110), end=(150, 120)) + + onselect.assert_called_once() + assert tool.extents == (100.0, 150.0, 110.0, 120.0) + + onselect.reset_mock() + click_and_drag(tool, start=(10, 100), end=(10, 100)) + onselect.assert_called_once() + + +@pytest.mark.parametrize('ignore_event_outside', [True, False]) +def test_rectangle_selector_ignore_outside(ax, ignore_event_outside): + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.RectangleSelector(ax, onselect=onselect, + ignore_event_outside=ignore_event_outside) + click_and_drag(tool, start=(100, 110), end=(150, 120)) + onselect.assert_called_once() + assert tool.extents == (100.0, 150.0, 110.0, 120.0) + + onselect.reset_mock() + # Trigger event outside of span + click_and_drag(tool, start=(150, 150), end=(160, 160)) + if ignore_event_outside: + # event have been ignored and span haven't changed. + onselect.assert_not_called() + assert tool.extents == (100.0, 150.0, 110.0, 120.0) + else: + # A new shape is created + onselect.assert_called_once() + assert tool.extents == (150.0, 160.0, 150.0, 160.0) + + +@pytest.mark.parametrize('orientation, onmove_callback, kwargs', [ + ('horizontal', False, dict(minspan=10, useblit=True)), + ('vertical', True, dict(button=1)), + ('horizontal', False, dict(props=dict(fill=True))), + ('horizontal', False, dict(interactive=True)), +]) +def test_span_selector(ax, orientation, onmove_callback, kwargs): + onselect = mock.Mock(spec=noop, return_value=None) + onmove = mock.Mock(spec=noop, return_value=None) + if onmove_callback: + kwargs['onmove_callback'] = onmove + + # While at it, also test that span selectors work in the presence of twin axes on + # top of the axes that contain the selector. Note that we need to unforce the axes + # aspect here, otherwise the twin axes forces the original axes' limits (to respect + # aspect=1) which makes some of the values below go out of bounds. + ax.set_aspect("auto") + tax = ax.twinx() + + tool = widgets.SpanSelector(ax, onselect, orientation, **kwargs) + do_event(tool, 'press', xdata=100, ydata=100, button=1) + # move outside of axis + do_event(tool, 'onmove', xdata=199, ydata=199, button=1) + do_event(tool, 'release', xdata=250, ydata=250, button=1) + + onselect.assert_called_once_with(100, 199) + if onmove_callback: + onmove.assert_called_once_with(100, 199) + + +@pytest.mark.parametrize('interactive', [True, False]) +def test_span_selector_onselect(ax, interactive): + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.SpanSelector(ax, onselect, 'horizontal', + interactive=interactive) + # move outside of axis + click_and_drag(tool, start=(100, 100), end=(150, 100)) + onselect.assert_called_once() + assert tool.extents == (100, 150) + + onselect.reset_mock() + click_and_drag(tool, start=(10, 100), end=(10, 100)) + onselect.assert_called_once() + + +@pytest.mark.parametrize('ignore_event_outside', [True, False]) +def test_span_selector_ignore_outside(ax, ignore_event_outside): + onselect = mock.Mock(spec=noop, return_value=None) + onmove = mock.Mock(spec=noop, return_value=None) + + tool = widgets.SpanSelector(ax, onselect, 'horizontal', + onmove_callback=onmove, + ignore_event_outside=ignore_event_outside) + click_and_drag(tool, start=(100, 100), end=(125, 125)) + onselect.assert_called_once() + onmove.assert_called_once() + assert tool.extents == (100, 125) + + onselect.reset_mock() + onmove.reset_mock() + # Trigger event outside of span + click_and_drag(tool, start=(150, 150), end=(160, 160)) + if ignore_event_outside: + # event have been ignored and span haven't changed. + onselect.assert_not_called() + onmove.assert_not_called() + assert tool.extents == (100, 125) + else: + # A new shape is created + onselect.assert_called_once() + onmove.assert_called_once() + assert tool.extents == (150, 160) + + +@pytest.mark.parametrize('drag_from_anywhere', [True, False]) +def test_span_selector_drag(ax, drag_from_anywhere): + # Create span + tool = widgets.SpanSelector(ax, onselect=noop, direction='horizontal', + interactive=True, + drag_from_anywhere=drag_from_anywhere) + click_and_drag(tool, start=(10, 10), end=(100, 120)) + assert tool.extents == (10, 100) + # Drag inside span + # + # If drag_from_anywhere == True, this will move the span by 10, + # giving new value extents = 20, 110 + # + # If drag_from_anywhere == False, this will create a new span with + # value extents = 25, 35 + click_and_drag(tool, start=(25, 15), end=(35, 25)) + if drag_from_anywhere: + assert tool.extents == (20, 110) + else: + assert tool.extents == (25, 35) + + # Check that in both cases, dragging outside the span draws a new span + click_and_drag(tool, start=(175, 185), end=(185, 195)) + assert tool.extents == (175, 185) + + +def test_span_selector_direction(ax): + tool = widgets.SpanSelector(ax, onselect=noop, direction='horizontal', + interactive=True) + assert tool.direction == 'horizontal' + assert tool._edge_handles.direction == 'horizontal' + + with pytest.raises(ValueError): + tool = widgets.SpanSelector(ax, onselect=noop, + direction='invalid_direction') + + tool.direction = 'vertical' + assert tool.direction == 'vertical' + assert tool._edge_handles.direction == 'vertical' + + with pytest.raises(ValueError): + tool.direction = 'invalid_string' + + +def test_span_selector_set_props_handle_props(ax): + tool = widgets.SpanSelector(ax, onselect=noop, direction='horizontal', + interactive=True, + props=dict(facecolor='b', alpha=0.2), + handle_props=dict(alpha=0.5)) + # Create rectangle + click_and_drag(tool, start=(0, 10), end=(100, 120)) + + artist = tool._selection_artist + assert artist.get_facecolor() == mcolors.to_rgba('b', alpha=0.2) + tool.set_props(facecolor='r', alpha=0.3) + assert artist.get_facecolor() == mcolors.to_rgba('r', alpha=0.3) + + for artist in tool._handles_artists: + assert artist.get_color() == 'b' + assert artist.get_alpha() == 0.5 + tool.set_handle_props(color='r', alpha=0.3) + for artist in tool._handles_artists: + assert artist.get_color() == 'r' + assert artist.get_alpha() == 0.3 + + +@pytest.mark.parametrize('selector', ['span', 'rectangle']) +def test_selector_clear(ax, selector): + kwargs = dict(ax=ax, interactive=True) + if selector == 'span': + Selector = widgets.SpanSelector + kwargs['direction'] = 'horizontal' + kwargs['onselect'] = noop + else: + Selector = widgets.RectangleSelector + + tool = Selector(**kwargs) + click_and_drag(tool, start=(10, 10), end=(100, 120)) + + # press-release event outside the selector to clear the selector + click_and_drag(tool, start=(130, 130), end=(130, 130)) + assert not tool._selection_completed + + kwargs['ignore_event_outside'] = True + tool = Selector(**kwargs) + assert tool.ignore_event_outside + click_and_drag(tool, start=(10, 10), end=(100, 120)) + + # press-release event outside the selector ignored + click_and_drag(tool, start=(130, 130), end=(130, 130)) + assert tool._selection_completed + + do_event(tool, 'on_key_press', key='escape') + assert not tool._selection_completed + + +@pytest.mark.parametrize('selector', ['span', 'rectangle']) +def test_selector_clear_method(ax, selector): + if selector == 'span': + tool = widgets.SpanSelector(ax, onselect=noop, direction='horizontal', + interactive=True, + ignore_event_outside=True) + else: + tool = widgets.RectangleSelector(ax, interactive=True) + click_and_drag(tool, start=(10, 10), end=(100, 120)) + assert tool._selection_completed + assert tool.get_visible() + if selector == 'span': + assert tool.extents == (10, 100) + + tool.clear() + assert not tool._selection_completed + assert not tool.get_visible() + + # Do another cycle of events to make sure we can + click_and_drag(tool, start=(10, 10), end=(50, 120)) + assert tool._selection_completed + assert tool.get_visible() + if selector == 'span': + assert tool.extents == (10, 50) + + +def test_span_selector_add_state(ax): + tool = widgets.SpanSelector(ax, noop, 'horizontal', + interactive=True) + + with pytest.raises(ValueError): + tool.add_state('unsupported_state') + with pytest.raises(ValueError): + tool.add_state('center') + with pytest.raises(ValueError): + tool.add_state('square') + + tool.add_state('move') + + +def test_tool_line_handle(ax): + positions = [20, 30, 50] + tool_line_handle = widgets.ToolLineHandles(ax, positions, 'horizontal', + useblit=False) + + for artist in tool_line_handle.artists: + assert not artist.get_animated() + assert not artist.get_visible() + + tool_line_handle.set_visible(True) + tool_line_handle.set_animated(True) + + for artist in tool_line_handle.artists: + assert artist.get_animated() + assert artist.get_visible() + + assert tool_line_handle.positions == positions + + +@pytest.mark.parametrize('direction', ("horizontal", "vertical")) +def test_span_selector_bound(direction): + fig, ax = plt.subplots(1, 1) + ax.plot([10, 20], [10, 30]) + fig.canvas.draw() + x_bound = ax.get_xbound() + y_bound = ax.get_ybound() + + tool = widgets.SpanSelector(ax, print, direction, interactive=True) + assert ax.get_xbound() == x_bound + assert ax.get_ybound() == y_bound + + bound = x_bound if direction == 'horizontal' else y_bound + assert tool._edge_handles.positions == list(bound) + + press_data = (10.5, 11.5) + move_data = (11, 13) # Updating selector is done in onmove + release_data = move_data + click_and_drag(tool, start=press_data, end=move_data) + + assert ax.get_xbound() == x_bound + assert ax.get_ybound() == y_bound + + index = 0 if direction == 'horizontal' else 1 + handle_positions = [press_data[index], release_data[index]] + assert tool._edge_handles.positions == handle_positions + + +@pytest.mark.backend('QtAgg', skip_on_importerror=True) +def test_span_selector_animated_artists_callback(): + """Check that the animated artists changed in callbacks are updated.""" + x = np.linspace(0, 2 * np.pi, 100) + values = np.sin(x) + + fig, ax = plt.subplots() + ln, = ax.plot(x, values, animated=True) + ln2, = ax.plot([], animated=True) + + # spin the event loop to let the backend process any pending operations + # before drawing artists + # See blitting tutorial + plt.pause(0.1) + ax.draw_artist(ln) + fig.canvas.blit(fig.bbox) + + def mean(vmin, vmax): + # Return mean of values in x between *vmin* and *vmax* + indmin, indmax = np.searchsorted(x, (vmin, vmax)) + v = values[indmin:indmax].mean() + ln2.set_data(x, np.full_like(x, v)) + + span = widgets.SpanSelector(ax, mean, direction='horizontal', + onmove_callback=mean, + interactive=True, + drag_from_anywhere=True, + useblit=True) + + # Add span selector and check that the line is draw after it was updated + # by the callback + press_data = [1, 2] + move_data = [2, 2] + do_event(span, 'press', xdata=press_data[0], ydata=press_data[1], button=1) + do_event(span, 'onmove', xdata=move_data[0], ydata=move_data[1], button=1) + assert span._get_animated_artists() == (ln, ln2) + assert ln.stale is False + assert ln2.stale + assert_allclose(ln2.get_ydata(), 0.9547335049088455) + span.update() + assert ln2.stale is False + + # Change span selector and check that the line is drawn/updated after its + # value was updated by the callback + press_data = [4, 0] + move_data = [5, 2] + release_data = [5, 2] + do_event(span, 'press', xdata=press_data[0], ydata=press_data[1], button=1) + do_event(span, 'onmove', xdata=move_data[0], ydata=move_data[1], button=1) + assert ln.stale is False + assert ln2.stale + assert_allclose(ln2.get_ydata(), -0.9424150707548072) + do_event(span, 'release', xdata=release_data[0], + ydata=release_data[1], button=1) + assert ln2.stale is False + + +def test_snapping_values_span_selector(ax): + def onselect(*args): + pass + + tool = widgets.SpanSelector(ax, onselect, direction='horizontal',) + snap_function = tool._snap + + snap_values = np.linspace(0, 5, 11) + values = np.array([-0.1, 0.1, 0.2, 0.5, 0.6, 0.7, 0.9, 4.76, 5.0, 5.5]) + expect = np.array([00.0, 0.0, 0.0, 0.5, 0.5, 0.5, 1.0, 5.00, 5.0, 5.0]) + values = snap_function(values, snap_values) + assert_allclose(values, expect) + + +def test_span_selector_snap(ax): + def onselect(vmin, vmax): + ax._got_onselect = True + + snap_values = np.arange(50) * 4 + + tool = widgets.SpanSelector(ax, onselect, direction='horizontal', + snap_values=snap_values) + tool.extents = (17, 35) + assert tool.extents == (16, 36) + + tool.snap_values = None + assert tool.snap_values is None + tool.extents = (17, 35) + assert tool.extents == (17, 35) + + +def test_span_selector_extents(ax): + tool = widgets.SpanSelector( + ax, lambda a, b: None, "horizontal", ignore_event_outside=True + ) + tool.extents = (5, 10) + + assert tool.extents == (5, 10) + assert tool._selection_completed + + # Since `ignore_event_outside=True`, this event should be ignored + press_data = (12, 14) + release_data = (20, 14) + click_and_drag(tool, start=press_data, end=release_data) + + assert tool.extents == (5, 10) + + +@pytest.mark.parametrize('kwargs', [ + dict(), + dict(useblit=False, props=dict(color='red')), + dict(useblit=True, button=1), +]) +def test_lasso_selector(ax, kwargs): + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.LassoSelector(ax, onselect=onselect, **kwargs) + do_event(tool, 'press', xdata=100, ydata=100, button=1) + do_event(tool, 'onmove', xdata=125, ydata=125, button=1) + do_event(tool, 'release', xdata=150, ydata=150, button=1) + + onselect.assert_called_once_with([(100, 100), (125, 125), (150, 150)]) + + +def test_lasso_selector_set_props(ax): + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.LassoSelector(ax, onselect=onselect, + props=dict(color='b', alpha=0.2)) + + artist = tool._selection_artist + assert mcolors.same_color(artist.get_color(), 'b') + assert artist.get_alpha() == 0.2 + tool.set_props(color='r', alpha=0.3) + assert mcolors.same_color(artist.get_color(), 'r') + assert artist.get_alpha() == 0.3 + + +def test_lasso_set_props(ax): + onselect = mock.Mock(spec=noop, return_value=None) + tool = widgets.Lasso(ax, (100, 100), onselect) + line = tool.line + assert mcolors.same_color(line.get_color(), 'black') + assert line.get_linestyle() == '-' + assert line.get_lw() == 2 + tool = widgets.Lasso(ax, (100, 100), onselect, props=dict( + linestyle='-', color='darkblue', alpha=0.2, lw=1)) + + line = tool.line + assert mcolors.same_color(line.get_color(), 'darkblue') + assert line.get_alpha() == 0.2 + assert line.get_lw() == 1 + assert line.get_linestyle() == '-' + line.set_color('r') + line.set_alpha(0.3) + assert mcolors.same_color(line.get_color(), 'r') + assert line.get_alpha() == 0.3 + + +def test_CheckButtons(ax): + labels = ('a', 'b', 'c') + check = widgets.CheckButtons(ax, labels, (True, False, True)) + assert check.get_status() == [True, False, True] + check.set_active(0) + assert check.get_status() == [False, False, True] + assert check.get_checked_labels() == ['c'] + check.clear() + assert check.get_status() == [False, False, False] + assert check.get_checked_labels() == [] + + for invalid_index in [-1, len(labels), len(labels)+5]: + with pytest.raises(ValueError): + check.set_active(index=invalid_index) + + for invalid_value in ['invalid', -1]: + with pytest.raises(TypeError): + check.set_active(1, state=invalid_value) + + cid = check.on_clicked(lambda: None) + check.disconnect(cid) + + +@pytest.mark.parametrize("toolbar", ["none", "toolbar2", "toolmanager"]) +def test_TextBox(ax, toolbar): + # Avoid "toolmanager is provisional" warning. + plt.rcParams._set("toolbar", toolbar) + + submit_event = mock.Mock(spec=noop, return_value=None) + text_change_event = mock.Mock(spec=noop, return_value=None) + tool = widgets.TextBox(ax, '') + tool.on_submit(submit_event) + tool.on_text_change(text_change_event) + + assert tool.text == '' + + do_event(tool, '_click') + + tool.set_val('x**2') + + assert tool.text == 'x**2' + assert text_change_event.call_count == 1 + + tool.begin_typing() + tool.stop_typing() + + assert submit_event.call_count == 2 + + do_event(tool, '_click', xdata=.5, ydata=.5) # Ensure the click is in the axes. + do_event(tool, '_keypress', key='+') + do_event(tool, '_keypress', key='5') + + assert text_change_event.call_count == 3 + + +def test_RadioButtons(ax): + radio = widgets.RadioButtons(ax, ('Radio 1', 'Radio 2', 'Radio 3')) + radio.set_active(1) + assert radio.value_selected == 'Radio 2' + assert radio.index_selected == 1 + radio.clear() + assert radio.value_selected == 'Radio 1' + assert radio.index_selected == 0 + + +@image_comparison(['check_radio_buttons.png'], style='mpl20', remove_text=True) +def test_check_radio_buttons_image(): + ax = get_ax() + fig = ax.get_figure(root=False) + fig.subplots_adjust(left=0.3) + + rax1 = fig.add_axes((0.05, 0.7, 0.2, 0.15)) + rb1 = widgets.RadioButtons(rax1, ('Radio 1', 'Radio 2', 'Radio 3')) + + rax2 = fig.add_axes((0.05, 0.5, 0.2, 0.15)) + cb1 = widgets.CheckButtons(rax2, ('Check 1', 'Check 2', 'Check 3'), + (False, True, True)) + + rax3 = fig.add_axes((0.05, 0.3, 0.2, 0.15)) + rb3 = widgets.RadioButtons( + rax3, ('Radio 1', 'Radio 2', 'Radio 3'), + label_props={'fontsize': [8, 12, 16], + 'color': ['red', 'green', 'blue']}, + radio_props={'edgecolor': ['red', 'green', 'blue'], + 'facecolor': ['mistyrose', 'palegreen', 'lightblue']}) + + rax4 = fig.add_axes((0.05, 0.1, 0.2, 0.15)) + cb4 = widgets.CheckButtons( + rax4, ('Check 1', 'Check 2', 'Check 3'), (False, True, True), + label_props={'fontsize': [8, 12, 16], + 'color': ['red', 'green', 'blue']}, + frame_props={'edgecolor': ['red', 'green', 'blue'], + 'facecolor': ['mistyrose', 'palegreen', 'lightblue']}, + check_props={'color': ['red', 'green', 'blue']}) + + +@check_figures_equal(extensions=["png"]) +def test_radio_buttons(fig_test, fig_ref): + widgets.RadioButtons(fig_test.subplots(), ["tea", "coffee"]) + ax = fig_ref.add_subplot(xticks=[], yticks=[]) + ax.scatter([.15, .15], [2/3, 1/3], transform=ax.transAxes, + s=(plt.rcParams["font.size"] / 2) ** 2, c=["C0", "none"]) + ax.text(.25, 2/3, "tea", transform=ax.transAxes, va="center") + ax.text(.25, 1/3, "coffee", transform=ax.transAxes, va="center") + + +@check_figures_equal(extensions=['png']) +def test_radio_buttons_props(fig_test, fig_ref): + label_props = {'color': ['red'], 'fontsize': [24]} + radio_props = {'facecolor': 'green', 'edgecolor': 'blue', 'linewidth': 2} + + widgets.RadioButtons(fig_ref.subplots(), ['tea', 'coffee'], + label_props=label_props, radio_props=radio_props) + + cb = widgets.RadioButtons(fig_test.subplots(), ['tea', 'coffee']) + cb.set_label_props(label_props) + # Setting the label size automatically increases default marker size, so we + # need to do that here as well. + cb.set_radio_props({**radio_props, 's': (24 / 2)**2}) + + +def test_radio_button_active_conflict(ax): + with pytest.warns(UserWarning, + match=r'Both the \*activecolor\* parameter'): + rb = widgets.RadioButtons(ax, ['tea', 'coffee'], activecolor='red', + radio_props={'facecolor': 'green'}) + # *radio_props*' facecolor wins over *activecolor* + assert mcolors.same_color(rb._buttons.get_facecolor(), ['green', 'none']) + + +@check_figures_equal(extensions=['png']) +def test_radio_buttons_activecolor_change(fig_test, fig_ref): + widgets.RadioButtons(fig_ref.subplots(), ['tea', 'coffee'], + activecolor='green') + + # Test property setter. + cb = widgets.RadioButtons(fig_test.subplots(), ['tea', 'coffee'], + activecolor='red') + cb.activecolor = 'green' + + +@check_figures_equal(extensions=["png"]) +def test_check_buttons(fig_test, fig_ref): + widgets.CheckButtons(fig_test.subplots(), ["tea", "coffee"], [True, True]) + ax = fig_ref.add_subplot(xticks=[], yticks=[]) + ax.scatter([.15, .15], [2/3, 1/3], marker='s', transform=ax.transAxes, + s=(plt.rcParams["font.size"] / 2) ** 2, c=["none", "none"]) + ax.scatter([.15, .15], [2/3, 1/3], marker='x', transform=ax.transAxes, + s=(plt.rcParams["font.size"] / 2) ** 2, c=["k", "k"]) + ax.text(.25, 2/3, "tea", transform=ax.transAxes, va="center") + ax.text(.25, 1/3, "coffee", transform=ax.transAxes, va="center") + + +@check_figures_equal(extensions=['png']) +def test_check_button_props(fig_test, fig_ref): + label_props = {'color': ['red'], 'fontsize': [24]} + frame_props = {'facecolor': 'green', 'edgecolor': 'blue', 'linewidth': 2} + check_props = {'facecolor': 'red', 'linewidth': 2} + + widgets.CheckButtons(fig_ref.subplots(), ['tea', 'coffee'], [True, True], + label_props=label_props, frame_props=frame_props, + check_props=check_props) + + cb = widgets.CheckButtons(fig_test.subplots(), ['tea', 'coffee'], + [True, True]) + cb.set_label_props(label_props) + # Setting the label size automatically increases default marker size, so we + # need to do that here as well. + cb.set_frame_props({**frame_props, 's': (24 / 2)**2}) + # FIXME: Axes.scatter promotes facecolor to edgecolor on unfilled markers, + # but Collection.update doesn't do that (it forgot the marker already). + # This means we cannot pass facecolor to both setters directly. + check_props['edgecolor'] = check_props.pop('facecolor') + cb.set_check_props({**check_props, 's': (24 / 2)**2}) + + +def test_slider_slidermin_slidermax_invalid(): + fig, ax = plt.subplots() + # test min/max with floats + with pytest.raises(ValueError): + widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + slidermin=10.0) + with pytest.raises(ValueError): + widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + slidermax=10.0) + + +def test_slider_slidermin_slidermax(): + fig, ax = plt.subplots() + slider_ = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=5.0) + + slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=1.0, slidermin=slider_) + assert slider.val == slider_.val + + slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=10.0, slidermax=slider_) + assert slider.val == slider_.val + + +def test_slider_valmin_valmax(): + fig, ax = plt.subplots() + slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=-10.0) + assert slider.val == slider.valmin + + slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=25.0) + assert slider.val == slider.valmax + + +def test_slider_valstep_snapping(): + fig, ax = plt.subplots() + slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=11.4, valstep=1) + assert slider.val == 11 + + slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0, + valinit=11.4, valstep=[0, 1, 5.5, 19.7]) + assert slider.val == 5.5 + + +def test_slider_horizontal_vertical(): + fig, ax = plt.subplots() + slider = widgets.Slider(ax=ax, label='', valmin=0, valmax=24, + valinit=12, orientation='horizontal') + slider.set_val(10) + assert slider.val == 10 + # check the dimension of the slider patch in axes units + box = slider.poly.get_extents().transformed(ax.transAxes.inverted()) + assert_allclose(box.bounds, [0, .25, 10/24, .5]) + + fig, ax = plt.subplots() + slider = widgets.Slider(ax=ax, label='', valmin=0, valmax=24, + valinit=12, orientation='vertical') + slider.set_val(10) + assert slider.val == 10 + # check the dimension of the slider patch in axes units + box = slider.poly.get_extents().transformed(ax.transAxes.inverted()) + assert_allclose(box.bounds, [.25, 0, .5, 10/24]) + + +def test_slider_reset(): + fig, ax = plt.subplots() + slider = widgets.Slider(ax=ax, label='', valmin=0, valmax=1, valinit=.5) + slider.set_val(0.75) + slider.reset() + assert slider.val == 0.5 + + +@pytest.mark.parametrize("orientation", ["horizontal", "vertical"]) +def test_range_slider(orientation): + if orientation == "vertical": + idx = [1, 0, 3, 2] + else: + idx = [0, 1, 2, 3] + + fig, ax = plt.subplots() + + slider = widgets.RangeSlider( + ax=ax, label="", valmin=0.0, valmax=1.0, orientation=orientation, + valinit=[0.1, 0.34] + ) + box = slider.poly.get_extents().transformed(ax.transAxes.inverted()) + assert_allclose(box.get_points().flatten()[idx], [0.1, 0.25, 0.34, 0.75]) + + # Check initial value is set correctly + assert_allclose(slider.val, (0.1, 0.34)) + + def handle_positions(slider): + if orientation == "vertical": + return [h.get_ydata()[0] for h in slider._handles] + else: + return [h.get_xdata()[0] for h in slider._handles] + + slider.set_val((0.4, 0.6)) + assert_allclose(slider.val, (0.4, 0.6)) + assert_allclose(handle_positions(slider), (0.4, 0.6)) + + box = slider.poly.get_extents().transformed(ax.transAxes.inverted()) + assert_allclose(box.get_points().flatten()[idx], [0.4, .25, 0.6, .75]) + + slider.set_val((0.2, 0.1)) + assert_allclose(slider.val, (0.1, 0.2)) + assert_allclose(handle_positions(slider), (0.1, 0.2)) + + slider.set_val((-1, 10)) + assert_allclose(slider.val, (0, 1)) + assert_allclose(handle_positions(slider), (0, 1)) + + slider.reset() + assert_allclose(slider.val, (0.1, 0.34)) + assert_allclose(handle_positions(slider), (0.1, 0.34)) + + +@pytest.mark.parametrize("orientation", ["horizontal", "vertical"]) +def test_range_slider_same_init_values(orientation): + if orientation == "vertical": + idx = [1, 0, 3, 2] + else: + idx = [0, 1, 2, 3] + + fig, ax = plt.subplots() + + slider = widgets.RangeSlider( + ax=ax, label="", valmin=0.0, valmax=1.0, orientation=orientation, + valinit=[0, 0] + ) + box = slider.poly.get_extents().transformed(ax.transAxes.inverted()) + assert_allclose(box.get_points().flatten()[idx], [0, 0.25, 0, 0.75]) + + +def check_polygon_selector(event_sequence, expected_result, selections_count, + **kwargs): + """ + Helper function to test Polygon Selector. + + Parameters + ---------- + event_sequence : list of tuples (etype, dict()) + A sequence of events to perform. The sequence is a list of tuples + where the first element of the tuple is an etype (e.g., 'onmove', + 'press', etc.), and the second element of the tuple is a dictionary of + the arguments for the event (e.g., xdata=5, key='shift', etc.). + expected_result : list of vertices (xdata, ydata) + The list of vertices that are expected to result from the event + sequence. + selections_count : int + Wait for the tool to call its `onselect` function `selections_count` + times, before comparing the result to the `expected_result` + **kwargs + Keyword arguments are passed to PolygonSelector. + """ + ax = get_ax() + + onselect = mock.Mock(spec=noop, return_value=None) + + tool = widgets.PolygonSelector(ax, onselect=onselect, **kwargs) + + for (etype, event_args) in event_sequence: + do_event(tool, etype, **event_args) + + assert onselect.call_count == selections_count + assert onselect.call_args == ((expected_result, ), {}) + + +def polygon_place_vertex(xdata, ydata): + return [('onmove', dict(xdata=xdata, ydata=ydata)), + ('press', dict(xdata=xdata, ydata=ydata)), + ('release', dict(xdata=xdata, ydata=ydata))] + + +def polygon_remove_vertex(xdata, ydata): + return [('onmove', dict(xdata=xdata, ydata=ydata)), + ('press', dict(xdata=xdata, ydata=ydata, button=3)), + ('release', dict(xdata=xdata, ydata=ydata, button=3))] + + +@pytest.mark.parametrize('draw_bounding_box', [False, True]) +def test_polygon_selector(draw_bounding_box): + check_selector = functools.partial( + check_polygon_selector, draw_bounding_box=draw_bounding_box) + + # Simple polygon + expected_result = [(50, 50), (150, 50), (50, 150)] + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 50), + ] + check_selector(event_sequence, expected_result, 1) + + # Move first vertex before completing the polygon. + expected_result = [(75, 50), (150, 50), (50, 150)] + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + ('on_key_press', dict(key='control')), + ('onmove', dict(xdata=50, ydata=50)), + ('press', dict(xdata=50, ydata=50)), + ('onmove', dict(xdata=75, ydata=50)), + ('release', dict(xdata=75, ydata=50)), + ('on_key_release', dict(key='control')), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(75, 50), + ] + check_selector(event_sequence, expected_result, 1) + + # Move first two vertices at once before completing the polygon. + expected_result = [(50, 75), (150, 75), (50, 150)] + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + ('on_key_press', dict(key='shift')), + ('onmove', dict(xdata=100, ydata=100)), + ('press', dict(xdata=100, ydata=100)), + ('onmove', dict(xdata=100, ydata=125)), + ('release', dict(xdata=100, ydata=125)), + ('on_key_release', dict(key='shift')), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 75), + ] + check_selector(event_sequence, expected_result, 1) + + # Move first vertex after completing the polygon. + expected_result = [(75, 50), (150, 50), (50, 150)] + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 50), + ('onmove', dict(xdata=50, ydata=50)), + ('press', dict(xdata=50, ydata=50)), + ('onmove', dict(xdata=75, ydata=50)), + ('release', dict(xdata=75, ydata=50)), + ] + check_selector(event_sequence, expected_result, 2) + + # Move all vertices after completing the polygon. + expected_result = [(75, 75), (175, 75), (75, 175)] + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 50), + ('on_key_press', dict(key='shift')), + ('onmove', dict(xdata=100, ydata=100)), + ('press', dict(xdata=100, ydata=100)), + ('onmove', dict(xdata=125, ydata=125)), + ('release', dict(xdata=125, ydata=125)), + ('on_key_release', dict(key='shift')), + ] + check_selector(event_sequence, expected_result, 2) + + # Try to move a vertex and move all before placing any vertices. + expected_result = [(50, 50), (150, 50), (50, 150)] + event_sequence = [ + ('on_key_press', dict(key='control')), + ('onmove', dict(xdata=100, ydata=100)), + ('press', dict(xdata=100, ydata=100)), + ('onmove', dict(xdata=125, ydata=125)), + ('release', dict(xdata=125, ydata=125)), + ('on_key_release', dict(key='control')), + ('on_key_press', dict(key='shift')), + ('onmove', dict(xdata=100, ydata=100)), + ('press', dict(xdata=100, ydata=100)), + ('onmove', dict(xdata=125, ydata=125)), + ('release', dict(xdata=125, ydata=125)), + ('on_key_release', dict(key='shift')), + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 50), + ] + check_selector(event_sequence, expected_result, 1) + + # Try to place vertex out-of-bounds, then reset, and start a new polygon. + expected_result = [(50, 50), (150, 50), (50, 150)] + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(250, 50), + ('on_key_press', dict(key='escape')), + ('on_key_release', dict(key='escape')), + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 50), + ] + check_selector(event_sequence, expected_result, 1) + + +@pytest.mark.parametrize('draw_bounding_box', [False, True]) +def test_polygon_selector_set_props_handle_props(ax, draw_bounding_box): + tool = widgets.PolygonSelector(ax, + props=dict(color='b', alpha=0.2), + handle_props=dict(alpha=0.5), + draw_bounding_box=draw_bounding_box) + + event_sequence = [ + *polygon_place_vertex(50, 50), + *polygon_place_vertex(150, 50), + *polygon_place_vertex(50, 150), + *polygon_place_vertex(50, 50), + ] + + for (etype, event_args) in event_sequence: + do_event(tool, etype, **event_args) + + artist = tool._selection_artist + assert artist.get_color() == 'b' + assert artist.get_alpha() == 0.2 + tool.set_props(color='r', alpha=0.3) + assert artist.get_color() == 'r' + assert artist.get_alpha() == 0.3 + + for artist in tool._handles_artists: + assert artist.get_color() == 'b' + assert artist.get_alpha() == 0.5 + tool.set_handle_props(color='r', alpha=0.3) + for artist in tool._handles_artists: + assert artist.get_color() == 'r' + assert artist.get_alpha() == 0.3 + + +@check_figures_equal(extensions=['png']) +def test_rect_visibility(fig_test, fig_ref): + # Check that requesting an invisible selector makes it invisible + ax_test = fig_test.subplots() + _ = fig_ref.subplots() + + tool = widgets.RectangleSelector(ax_test, props={'visible': False}) + tool.extents = (0.2, 0.8, 0.3, 0.7) + + +# Change the order that the extra point is inserted in +@pytest.mark.parametrize('idx', [1, 2, 3]) +@pytest.mark.parametrize('draw_bounding_box', [False, True]) +def test_polygon_selector_remove(idx, draw_bounding_box): + verts = [(50, 50), (150, 50), (50, 150)] + event_sequence = [polygon_place_vertex(*verts[0]), + polygon_place_vertex(*verts[1]), + polygon_place_vertex(*verts[2]), + # Finish the polygon + polygon_place_vertex(*verts[0])] + # Add an extra point + event_sequence.insert(idx, polygon_place_vertex(200, 200)) + # Remove the extra point + event_sequence.append(polygon_remove_vertex(200, 200)) + # Flatten list of lists + event_sequence = functools.reduce(operator.iadd, event_sequence, []) + check_polygon_selector(event_sequence, verts, 2, + draw_bounding_box=draw_bounding_box) + + +@pytest.mark.parametrize('draw_bounding_box', [False, True]) +def test_polygon_selector_remove_first_point(draw_bounding_box): + verts = [(50, 50), (150, 50), (50, 150)] + event_sequence = [ + *polygon_place_vertex(*verts[0]), + *polygon_place_vertex(*verts[1]), + *polygon_place_vertex(*verts[2]), + *polygon_place_vertex(*verts[0]), + *polygon_remove_vertex(*verts[0]), + ] + check_polygon_selector(event_sequence, verts[1:], 2, + draw_bounding_box=draw_bounding_box) + + +@pytest.mark.parametrize('draw_bounding_box', [False, True]) +def test_polygon_selector_redraw(ax, draw_bounding_box): + verts = [(50, 50), (150, 50), (50, 150)] + event_sequence = [ + *polygon_place_vertex(*verts[0]), + *polygon_place_vertex(*verts[1]), + *polygon_place_vertex(*verts[2]), + *polygon_place_vertex(*verts[0]), + # Polygon completed, now remove first two verts. + *polygon_remove_vertex(*verts[1]), + *polygon_remove_vertex(*verts[2]), + # At this point the tool should be reset so we can add more vertices. + *polygon_place_vertex(*verts[1]), + ] + + tool = widgets.PolygonSelector(ax, draw_bounding_box=draw_bounding_box) + for (etype, event_args) in event_sequence: + do_event(tool, etype, **event_args) + # After removing two verts, only one remains, and the + # selector should be automatically resete + assert tool.verts == verts[0:2] + + +@pytest.mark.parametrize('draw_bounding_box', [False, True]) +@check_figures_equal(extensions=['png']) +def test_polygon_selector_verts_setter(fig_test, fig_ref, draw_bounding_box): + verts = [(0.1, 0.4), (0.5, 0.9), (0.3, 0.2)] + ax_test = fig_test.add_subplot() + + tool_test = widgets.PolygonSelector(ax_test, draw_bounding_box=draw_bounding_box) + tool_test.verts = verts + assert tool_test.verts == verts + + ax_ref = fig_ref.add_subplot() + tool_ref = widgets.PolygonSelector(ax_ref, draw_bounding_box=draw_bounding_box) + event_sequence = [ + *polygon_place_vertex(*verts[0]), + *polygon_place_vertex(*verts[1]), + *polygon_place_vertex(*verts[2]), + *polygon_place_vertex(*verts[0]), + ] + for (etype, event_args) in event_sequence: + do_event(tool_ref, etype, **event_args) + + +def test_polygon_selector_box(ax): + # Create a diamond (adjusting axes lims s.t. the diamond lies within axes limits). + ax.set(xlim=(-10, 50), ylim=(-10, 50)) + verts = [(20, 0), (0, 20), (20, 40), (40, 20)] + event_sequence = [ + *polygon_place_vertex(*verts[0]), + *polygon_place_vertex(*verts[1]), + *polygon_place_vertex(*verts[2]), + *polygon_place_vertex(*verts[3]), + *polygon_place_vertex(*verts[0]), + ] + + # Create selector + tool = widgets.PolygonSelector(ax, draw_bounding_box=True) + for (etype, event_args) in event_sequence: + do_event(tool, etype, **event_args) + + # In order to trigger the correct callbacks, trigger events on the canvas + # instead of the individual tools + t = ax.transData + canvas = ax.get_figure(root=True).canvas + + # Scale to half size using the top right corner of the bounding box + MouseEvent( + "button_press_event", canvas, *t.transform((40, 40)), 1)._process() + MouseEvent( + "motion_notify_event", canvas, *t.transform((20, 20)))._process() + MouseEvent( + "button_release_event", canvas, *t.transform((20, 20)), 1)._process() + np.testing.assert_allclose( + tool.verts, [(10, 0), (0, 10), (10, 20), (20, 10)]) + + # Move using the center of the bounding box + MouseEvent( + "button_press_event", canvas, *t.transform((10, 10)), 1)._process() + MouseEvent( + "motion_notify_event", canvas, *t.transform((30, 30)))._process() + MouseEvent( + "button_release_event", canvas, *t.transform((30, 30)), 1)._process() + np.testing.assert_allclose( + tool.verts, [(30, 20), (20, 30), (30, 40), (40, 30)]) + + # Remove a point from the polygon and check that the box extents update + np.testing.assert_allclose( + tool._box.extents, (20.0, 40.0, 20.0, 40.0)) + + MouseEvent( + "button_press_event", canvas, *t.transform((30, 20)), 3)._process() + MouseEvent( + "button_release_event", canvas, *t.transform((30, 20)), 3)._process() + np.testing.assert_allclose( + tool.verts, [(20, 30), (30, 40), (40, 30)]) + np.testing.assert_allclose( + tool._box.extents, (20.0, 40.0, 30.0, 40.0)) + + +def test_polygon_selector_clear_method(ax): + onselect = mock.Mock(spec=noop, return_value=None) + tool = widgets.PolygonSelector(ax, onselect) + + for result in ([(50, 50), (150, 50), (50, 150), (50, 50)], + [(50, 50), (100, 50), (50, 150), (50, 50)]): + for x, y in result: + for etype, event_args in polygon_place_vertex(x, y): + do_event(tool, etype, **event_args) + + artist = tool._selection_artist + + assert tool._selection_completed + assert tool.get_visible() + assert artist.get_visible() + np.testing.assert_equal(artist.get_xydata(), result) + assert onselect.call_args == ((result[:-1],), {}) + + tool.clear() + assert not tool._selection_completed + np.testing.assert_equal(artist.get_xydata(), [(0, 0)]) + + +@pytest.mark.parametrize("horizOn", [False, True]) +@pytest.mark.parametrize("vertOn", [False, True]) +def test_MultiCursor(horizOn, vertOn): + fig = plt.figure() + (ax1, ax3) = fig.subplots(2, sharex=True) + ax2 = plt.figure().subplots() + + # useblit=false to avoid having to draw the figure to cache the renderer + multi = widgets.MultiCursor( + None, (ax1, ax2), useblit=False, horizOn=horizOn, vertOn=vertOn + ) + + # Only two of the axes should have a line drawn on them. + assert len(multi.vlines) == 2 + assert len(multi.hlines) == 2 + + # mock a motion_notify_event + # Can't use `do_event` as that helper requires the widget + # to have a single .ax attribute. + event = mock_event(ax1, xdata=.5, ydata=.25) + multi.onmove(event) + # force a draw + draw event to exercise clear + fig.canvas.draw() + + # the lines in the first two ax should both move + for l in multi.vlines: + assert l.get_xdata() == (.5, .5) + for l in multi.hlines: + assert l.get_ydata() == (.25, .25) + # The relevant lines get turned on after move. + assert len([line for line in multi.vlines if line.get_visible()]) == ( + 2 if vertOn else 0) + assert len([line for line in multi.hlines if line.get_visible()]) == ( + 2 if horizOn else 0) + + # After toggling settings, the opposite lines should be visible after move. + multi.horizOn = not multi.horizOn + multi.vertOn = not multi.vertOn + event = mock_event(ax1, xdata=.5, ydata=.25) + multi.onmove(event) + assert len([line for line in multi.vlines if line.get_visible()]) == ( + 0 if vertOn else 2) + assert len([line for line in multi.hlines if line.get_visible()]) == ( + 0 if horizOn else 2) + + # test a move event in an Axes not part of the MultiCursor + # the lines in ax1 and ax2 should not have moved. + event = mock_event(ax3, xdata=.75, ydata=.75) + multi.onmove(event) + for l in multi.vlines: + assert l.get_xdata() == (.5, .5) + for l in multi.hlines: + assert l.get_ydata() == (.25, .25) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e000831d8a080a4bf6e8af758b1d978755c2bd14 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/__init__.py @@ -0,0 +1,23 @@ +""" +Unstructured triangular grid functions. +""" + +from ._triangulation import Triangulation +from ._tricontour import TriContourSet, tricontour, tricontourf +from ._trifinder import TriFinder, TrapezoidMapTriFinder +from ._triinterpolate import (TriInterpolator, LinearTriInterpolator, + CubicTriInterpolator) +from ._tripcolor import tripcolor +from ._triplot import triplot +from ._trirefine import TriRefiner, UniformTriRefiner +from ._tritools import TriAnalyzer + + +__all__ = ["Triangulation", + "TriContourSet", "tricontour", "tricontourf", + "TriFinder", "TrapezoidMapTriFinder", + "TriInterpolator", "LinearTriInterpolator", "CubicTriInterpolator", + "tripcolor", + "triplot", + "TriRefiner", "UniformTriRefiner", + "TriAnalyzer"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triangulation.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triangulation.py new file mode 100644 index 0000000000000000000000000000000000000000..a07192dfc8cac08a102be6ef3b22047f264a5099 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triangulation.py @@ -0,0 +1,247 @@ +import sys + +import numpy as np + +from matplotlib import _api + + +class Triangulation: + """ + An unstructured triangular grid consisting of npoints points and + ntri triangles. The triangles can either be specified by the user + or automatically generated using a Delaunay triangulation. + + Parameters + ---------- + x, y : (npoints,) array-like + Coordinates of grid points. + triangles : (ntri, 3) array-like of int, optional + For each triangle, the indices of the three points that make + up the triangle, ordered in an anticlockwise manner. If not + specified, the Delaunay triangulation is calculated. + mask : (ntri,) array-like of bool, optional + Which triangles are masked out. + + Attributes + ---------- + triangles : (ntri, 3) array of int + For each triangle, the indices of the three points that make + up the triangle, ordered in an anticlockwise manner. If you want to + take the *mask* into account, use `get_masked_triangles` instead. + mask : (ntri, 3) array of bool or None + Masked out triangles. + is_delaunay : bool + Whether the Triangulation is a calculated Delaunay + triangulation (where *triangles* was not specified) or not. + + Notes + ----- + For a Triangulation to be valid it must not have duplicate points, + triangles formed from colinear points, or overlapping triangles. + """ + def __init__(self, x, y, triangles=None, mask=None): + from matplotlib import _qhull + + self.x = np.asarray(x, dtype=np.float64) + self.y = np.asarray(y, dtype=np.float64) + if self.x.shape != self.y.shape or self.x.ndim != 1: + raise ValueError("x and y must be equal-length 1D arrays, but " + f"found shapes {self.x.shape!r} and " + f"{self.y.shape!r}") + + self.mask = None + self._edges = None + self._neighbors = None + self.is_delaunay = False + + if triangles is None: + # No triangulation specified, so use matplotlib._qhull to obtain + # Delaunay triangulation. + self.triangles, self._neighbors = _qhull.delaunay(x, y, sys.flags.verbose) + self.is_delaunay = True + else: + # Triangulation specified. Copy, since we may correct triangle + # orientation. + try: + self.triangles = np.array(triangles, dtype=np.int32, order='C') + except ValueError as e: + raise ValueError('triangles must be a (N, 3) int array, not ' + f'{triangles!r}') from e + if self.triangles.ndim != 2 or self.triangles.shape[1] != 3: + raise ValueError( + 'triangles must be a (N, 3) int array, but found shape ' + f'{self.triangles.shape!r}') + if self.triangles.max() >= len(self.x): + raise ValueError( + 'triangles are indices into the points and must be in the ' + f'range 0 <= i < {len(self.x)} but found value ' + f'{self.triangles.max()}') + if self.triangles.min() < 0: + raise ValueError( + 'triangles are indices into the points and must be in the ' + f'range 0 <= i < {len(self.x)} but found value ' + f'{self.triangles.min()}') + + # Underlying C++ object is not created until first needed. + self._cpp_triangulation = None + + # Default TriFinder not created until needed. + self._trifinder = None + + self.set_mask(mask) + + def calculate_plane_coefficients(self, z): + """ + Calculate plane equation coefficients for all unmasked triangles from + the point (x, y) coordinates and specified z-array of shape (npoints). + The returned array has shape (npoints, 3) and allows z-value at (x, y) + position in triangle tri to be calculated using + ``z = array[tri, 0] * x + array[tri, 1] * y + array[tri, 2]``. + """ + return self.get_cpp_triangulation().calculate_plane_coefficients(z) + + @property + def edges(self): + """ + Return integer array of shape (nedges, 2) containing all edges of + non-masked triangles. + + Each row defines an edge by its start point index and end point + index. Each edge appears only once, i.e. for an edge between points + *i* and *j*, there will only be either *(i, j)* or *(j, i)*. + """ + if self._edges is None: + self._edges = self.get_cpp_triangulation().get_edges() + return self._edges + + def get_cpp_triangulation(self): + """ + Return the underlying C++ Triangulation object, creating it + if necessary. + """ + from matplotlib import _tri + if self._cpp_triangulation is None: + self._cpp_triangulation = _tri.Triangulation( + # For unset arrays use empty tuple which has size of zero. + self.x, self.y, self.triangles, + self.mask if self.mask is not None else (), + self._edges if self._edges is not None else (), + self._neighbors if self._neighbors is not None else (), + not self.is_delaunay) + return self._cpp_triangulation + + def get_masked_triangles(self): + """ + Return an array of triangles taking the mask into account. + """ + if self.mask is not None: + return self.triangles[~self.mask] + else: + return self.triangles + + @staticmethod + def get_from_args_and_kwargs(*args, **kwargs): + """ + Return a Triangulation object from the args and kwargs, and + the remaining args and kwargs with the consumed values removed. + + There are two alternatives: either the first argument is a + Triangulation object, in which case it is returned, or the args + and kwargs are sufficient to create a new Triangulation to + return. In the latter case, see Triangulation.__init__ for + the possible args and kwargs. + """ + if isinstance(args[0], Triangulation): + triangulation, *args = args + if 'triangles' in kwargs: + _api.warn_external( + "Passing the keyword 'triangles' has no effect when also " + "passing a Triangulation") + if 'mask' in kwargs: + _api.warn_external( + "Passing the keyword 'mask' has no effect when also " + "passing a Triangulation") + else: + x, y, triangles, mask, args, kwargs = \ + Triangulation._extract_triangulation_params(args, kwargs) + triangulation = Triangulation(x, y, triangles, mask) + return triangulation, args, kwargs + + @staticmethod + def _extract_triangulation_params(args, kwargs): + x, y, *args = args + # Check triangles in kwargs then args. + triangles = kwargs.pop('triangles', None) + from_args = False + if triangles is None and args: + triangles = args[0] + from_args = True + if triangles is not None: + try: + triangles = np.asarray(triangles, dtype=np.int32) + except ValueError: + triangles = None + if triangles is not None and (triangles.ndim != 2 or + triangles.shape[1] != 3): + triangles = None + if triangles is not None and from_args: + args = args[1:] # Consumed first item in args. + # Check for mask in kwargs. + mask = kwargs.pop('mask', None) + return x, y, triangles, mask, args, kwargs + + def get_trifinder(self): + """ + Return the default `matplotlib.tri.TriFinder` of this + triangulation, creating it if necessary. This allows the same + TriFinder object to be easily shared. + """ + if self._trifinder is None: + # Default TriFinder class. + from matplotlib.tri._trifinder import TrapezoidMapTriFinder + self._trifinder = TrapezoidMapTriFinder(self) + return self._trifinder + + @property + def neighbors(self): + """ + Return integer array of shape (ntri, 3) containing neighbor triangles. + + For each triangle, the indices of the three triangles that + share the same edges, or -1 if there is no such neighboring + triangle. ``neighbors[i, j]`` is the triangle that is the neighbor + to the edge from point index ``triangles[i, j]`` to point index + ``triangles[i, (j+1)%3]``. + """ + if self._neighbors is None: + self._neighbors = self.get_cpp_triangulation().get_neighbors() + return self._neighbors + + def set_mask(self, mask): + """ + Set or clear the mask array. + + Parameters + ---------- + mask : None or bool array of length ntri + """ + if mask is None: + self.mask = None + else: + self.mask = np.asarray(mask, dtype=bool) + if self.mask.shape != (self.triangles.shape[0],): + raise ValueError('mask array must have same length as ' + 'triangles array') + + # Set mask in C++ Triangulation. + if self._cpp_triangulation is not None: + self._cpp_triangulation.set_mask( + self.mask if self.mask is not None else ()) + + # Clear derived fields so they are recalculated when needed. + self._edges = None + self._neighbors = None + + # Recalculate TriFinder if it exists. + if self._trifinder is not None: + self._trifinder._initialize() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triangulation.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triangulation.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6e00b272eda93554db53b9fc98e9b4f1b94fbebd --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triangulation.pyi @@ -0,0 +1,33 @@ +from matplotlib import _tri +from matplotlib.tri._trifinder import TriFinder + +import numpy as np +from numpy.typing import ArrayLike +from typing import Any + +class Triangulation: + x: np.ndarray + y: np.ndarray + mask: np.ndarray | None + is_delaunay: bool + triangles: np.ndarray + def __init__( + self, + x: ArrayLike, + y: ArrayLike, + triangles: ArrayLike | None = ..., + mask: ArrayLike | None = ..., + ) -> None: ... + def calculate_plane_coefficients(self, z: ArrayLike) -> np.ndarray: ... + @property + def edges(self) -> np.ndarray: ... + def get_cpp_triangulation(self) -> _tri.Triangulation: ... + def get_masked_triangles(self) -> np.ndarray: ... + @staticmethod + def get_from_args_and_kwargs( + *args, **kwargs + ) -> tuple[Triangulation, tuple[Any, ...], dict[str, Any]]: ... + def get_trifinder(self) -> TriFinder: ... + @property + def neighbors(self) -> np.ndarray: ... + def set_mask(self, mask: None | ArrayLike) -> None: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tricontour.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tricontour.py new file mode 100644 index 0000000000000000000000000000000000000000..8250515f3ef8fde7fecb57cc441b7f1ea50b95f0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tricontour.py @@ -0,0 +1,270 @@ +import numpy as np + +from matplotlib import _docstring +from matplotlib.contour import ContourSet +from matplotlib.tri._triangulation import Triangulation + + +@_docstring.interpd +class TriContourSet(ContourSet): + """ + Create and store a set of contour lines or filled regions for + a triangular grid. + + This class is typically not instantiated directly by the user but by + `~.Axes.tricontour` and `~.Axes.tricontourf`. + + %(contour_set_attributes)s + """ + def __init__(self, ax, *args, **kwargs): + """ + Draw triangular grid contour lines or filled regions, + depending on whether keyword arg *filled* is False + (default) or True. + + The first argument of the initializer must be an `~.axes.Axes` + object. The remaining arguments and keyword arguments + are described in the docstring of `~.Axes.tricontour`. + """ + super().__init__(ax, *args, **kwargs) + + def _process_args(self, *args, **kwargs): + """ + Process args and kwargs. + """ + if isinstance(args[0], TriContourSet): + C = args[0]._contour_generator + if self.levels is None: + self.levels = args[0].levels + self.zmin = args[0].zmin + self.zmax = args[0].zmax + self._mins = args[0]._mins + self._maxs = args[0]._maxs + else: + from matplotlib import _tri + tri, z = self._contour_args(args, kwargs) + C = _tri.TriContourGenerator(tri.get_cpp_triangulation(), z) + self._mins = [tri.x.min(), tri.y.min()] + self._maxs = [tri.x.max(), tri.y.max()] + + self._contour_generator = C + return kwargs + + def _contour_args(self, args, kwargs): + tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, + **kwargs) + z, *args = args + z = np.ma.asarray(z) + if z.shape != tri.x.shape: + raise ValueError('z array must have same length as triangulation x' + ' and y arrays') + + # z values must be finite, only need to check points that are included + # in the triangulation. + z_check = z[np.unique(tri.get_masked_triangles())] + if np.ma.is_masked(z_check): + raise ValueError('z must not contain masked points within the ' + 'triangulation') + if not np.isfinite(z_check).all(): + raise ValueError('z array must not contain non-finite values ' + 'within the triangulation') + + z = np.ma.masked_invalid(z, copy=False) + self.zmax = float(z_check.max()) + self.zmin = float(z_check.min()) + if self.logscale and self.zmin <= 0: + func = 'contourf' if self.filled else 'contour' + raise ValueError(f'Cannot {func} log of negative values.') + self._process_contour_level_args(args, z.dtype) + return (tri, z) + + +_docstring.interpd.register(_tricontour_doc=""" +Draw contour %%(type)s on an unstructured triangular grid. + +Call signatures:: + + %%(func)s(triangulation, z, [levels], ...) + %%(func)s(x, y, z, [levels], *, [triangles=triangles], [mask=mask], ...) + +The triangular grid can be specified either by passing a `.Triangulation` +object as the first parameter, or by passing the points *x*, *y* and +optionally the *triangles* and a *mask*. See `.Triangulation` for an +explanation of these parameters. If neither of *triangulation* or +*triangles* are given, the triangulation is calculated on the fly. + +It is possible to pass *triangles* positionally, i.e. +``%%(func)s(x, y, triangles, z, ...)``. However, this is discouraged. For more +clarity, pass *triangles* via keyword argument. + +Parameters +---------- +triangulation : `.Triangulation`, optional + An already created triangular grid. + +x, y, triangles, mask + Parameters defining the triangular grid. See `.Triangulation`. + This is mutually exclusive with specifying *triangulation*. + +z : array-like + The height values over which the contour is drawn. Color-mapping is + controlled by *cmap*, *norm*, *vmin*, and *vmax*. + + .. note:: + All values in *z* must be finite. Hence, nan and inf values must + either be removed or `~.Triangulation.set_mask` be used. + +levels : int or array-like, optional + Determines the number and positions of the contour lines / regions. + + If an int *n*, use `~matplotlib.ticker.MaxNLocator`, which tries to + automatically choose no more than *n+1* "nice" contour levels between + between minimum and maximum numeric values of *Z*. + + If array-like, draw contour lines at the specified levels. The values must + be in increasing order. + +Returns +------- +`~matplotlib.tri.TriContourSet` + +Other Parameters +---------------- +colors : :mpltype:`color` or list of :mpltype:`color`, optional + The colors of the levels, i.e., the contour %%(type)s. + + The sequence is cycled for the levels in ascending order. If the sequence + is shorter than the number of levels, it is repeated. + + As a shortcut, single color strings may be used in place of one-element + lists, i.e. ``'red'`` instead of ``['red']`` to color all levels with the + same color. This shortcut does only work for color strings, not for other + ways of specifying colors. + + By default (value *None*), the colormap specified by *cmap* will be used. + +alpha : float, default: 1 + The alpha blending value, between 0 (transparent) and 1 (opaque). + +%(cmap_doc)s + + This parameter is ignored if *colors* is set. + +%(norm_doc)s + + This parameter is ignored if *colors* is set. + +%(vmin_vmax_doc)s + + If *vmin* or *vmax* are not given, the default color scaling is based on + *levels*. + + This parameter is ignored if *colors* is set. + +origin : {*None*, 'upper', 'lower', 'image'}, default: None + Determines the orientation and exact position of *z* by specifying the + position of ``z[0, 0]``. This is only relevant, if *X*, *Y* are not given. + + - *None*: ``z[0, 0]`` is at X=0, Y=0 in the lower left corner. + - 'lower': ``z[0, 0]`` is at X=0.5, Y=0.5 in the lower left corner. + - 'upper': ``z[0, 0]`` is at X=N+0.5, Y=0.5 in the upper left corner. + - 'image': Use the value from :rc:`image.origin`. + +extent : (x0, x1, y0, y1), optional + If *origin* is not *None*, then *extent* is interpreted as in `.imshow`: it + gives the outer pixel boundaries. In this case, the position of z[0, 0] is + the center of the pixel, not a corner. If *origin* is *None*, then + (*x0*, *y0*) is the position of z[0, 0], and (*x1*, *y1*) is the position + of z[-1, -1]. + + This argument is ignored if *X* and *Y* are specified in the call to + contour. + +locator : ticker.Locator subclass, optional + The locator is used to determine the contour levels if they are not given + explicitly via *levels*. + Defaults to `~.ticker.MaxNLocator`. + +extend : {'neither', 'both', 'min', 'max'}, default: 'neither' + Determines the ``%%(func)s``-coloring of values that are outside the + *levels* range. + + If 'neither', values outside the *levels* range are not colored. If 'min', + 'max' or 'both', color the values below, above or below and above the + *levels* range. + + Values below ``min(levels)`` and above ``max(levels)`` are mapped to the + under/over values of the `.Colormap`. Note that most colormaps do not have + dedicated colors for these by default, so that the over and under values + are the edge values of the colormap. You may want to set these values + explicitly using `.Colormap.set_under` and `.Colormap.set_over`. + + .. note:: + + An existing `.TriContourSet` does not get notified if properties of its + colormap are changed. Therefore, an explicit call to + `.ContourSet.changed()` is needed after modifying the colormap. The + explicit call can be left out, if a colorbar is assigned to the + `.TriContourSet` because it internally calls `.ContourSet.changed()`. + +xunits, yunits : registered units, optional + Override axis units by specifying an instance of a + :class:`matplotlib.units.ConversionInterface`. + +antialiased : bool, optional + Enable antialiasing, overriding the defaults. For + filled contours, the default is *True*. For line contours, + it is taken from :rc:`lines.antialiased`.""" % _docstring.interpd.params) + + +@_docstring.Substitution(func='tricontour', type='lines') +@_docstring.interpd +def tricontour(ax, *args, **kwargs): + """ + %(_tricontour_doc)s + + linewidths : float or array-like, default: :rc:`contour.linewidth` + The line width of the contour lines. + + If a number, all levels will be plotted with this linewidth. + + If a sequence, the levels in ascending order will be plotted with + the linewidths in the order specified. + + If None, this falls back to :rc:`lines.linewidth`. + + linestyles : {*None*, 'solid', 'dashed', 'dashdot', 'dotted'}, optional + If *linestyles* is *None*, the default is 'solid' unless the lines are + monochrome. In that case, negative contours will take their linestyle + from :rc:`contour.negative_linestyle` setting. + + *linestyles* can also be an iterable of the above strings specifying a + set of linestyles to be used. If this iterable is shorter than the + number of contour levels it will be repeated as necessary. + """ + kwargs['filled'] = False + return TriContourSet(ax, *args, **kwargs) + + +@_docstring.Substitution(func='tricontourf', type='regions') +@_docstring.interpd +def tricontourf(ax, *args, **kwargs): + """ + %(_tricontour_doc)s + + hatches : list[str], optional + A list of crosshatch patterns to use on the filled areas. + If None, no hatching will be added to the contour. + + Notes + ----- + `.tricontourf` fills intervals that are closed at the top; that is, for + boundaries *z1* and *z2*, the filled region is:: + + z1 < Z <= z2 + + except for the lowest interval, which is closed on both sides (i.e. it + includes the lowest value). + """ + kwargs['filled'] = True + return TriContourSet(ax, *args, **kwargs) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tricontour.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tricontour.pyi new file mode 100644 index 0000000000000000000000000000000000000000..31929d8661560cce2d63e735c37163baca02c001 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tricontour.pyi @@ -0,0 +1,52 @@ +from matplotlib.axes import Axes +from matplotlib.contour import ContourSet +from matplotlib.tri._triangulation import Triangulation + +from numpy.typing import ArrayLike +from typing import overload + +# TODO: more explicit args/kwargs (for all things in this module)? + +class TriContourSet(ContourSet): + def __init__(self, ax: Axes, *args, **kwargs) -> None: ... + +@overload +def tricontour( + ax: Axes, + triangulation: Triangulation, + z: ArrayLike, + levels: int | ArrayLike = ..., + **kwargs +) -> TriContourSet: ... +@overload +def tricontour( + ax: Axes, + x: ArrayLike, + y: ArrayLike, + z: ArrayLike, + levels: int | ArrayLike = ..., + *, + triangles: ArrayLike = ..., + mask: ArrayLike = ..., + **kwargs +) -> TriContourSet: ... +@overload +def tricontourf( + ax: Axes, + triangulation: Triangulation, + z: ArrayLike, + levels: int | ArrayLike = ..., + **kwargs +) -> TriContourSet: ... +@overload +def tricontourf( + ax: Axes, + x: ArrayLike, + y: ArrayLike, + z: ArrayLike, + levels: int | ArrayLike = ..., + *, + triangles: ArrayLike = ..., + mask: ArrayLike = ..., + **kwargs +) -> TriContourSet: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trifinder.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trifinder.py new file mode 100644 index 0000000000000000000000000000000000000000..852f3d9eebccf1492a98e6bd8f327b87d5e014c8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trifinder.py @@ -0,0 +1,96 @@ +import numpy as np + +from matplotlib import _api +from matplotlib.tri import Triangulation + + +class TriFinder: + """ + Abstract base class for classes used to find the triangles of a + Triangulation in which (x, y) points lie. + + Rather than instantiate an object of a class derived from TriFinder, it is + usually better to use the function `.Triangulation.get_trifinder`. + + Derived classes implement __call__(x, y) where x and y are array-like point + coordinates of the same shape. + """ + + def __init__(self, triangulation): + _api.check_isinstance(Triangulation, triangulation=triangulation) + self._triangulation = triangulation + + def __call__(self, x, y): + raise NotImplementedError + + +class TrapezoidMapTriFinder(TriFinder): + """ + `~matplotlib.tri.TriFinder` class implemented using the trapezoid + map algorithm from the book "Computational Geometry, Algorithms and + Applications", second edition, by M. de Berg, M. van Kreveld, M. Overmars + and O. Schwarzkopf. + + The triangulation must be valid, i.e. it must not have duplicate points, + triangles formed from colinear points, or overlapping triangles. The + algorithm has some tolerance to triangles formed from colinear points, but + this should not be relied upon. + """ + + def __init__(self, triangulation): + from matplotlib import _tri + super().__init__(triangulation) + self._cpp_trifinder = _tri.TrapezoidMapTriFinder( + triangulation.get_cpp_triangulation()) + self._initialize() + + def __call__(self, x, y): + """ + Return an array containing the indices of the triangles in which the + specified *x*, *y* points lie, or -1 for points that do not lie within + a triangle. + + *x*, *y* are array-like x and y coordinates of the same shape and any + number of dimensions. + + Returns integer array with the same shape and *x* and *y*. + """ + x = np.asarray(x, dtype=np.float64) + y = np.asarray(y, dtype=np.float64) + if x.shape != y.shape: + raise ValueError("x and y must be array-like with the same shape") + + # C++ does the heavy lifting, and expects 1D arrays. + indices = (self._cpp_trifinder.find_many(x.ravel(), y.ravel()) + .reshape(x.shape)) + return indices + + def _get_tree_stats(self): + """ + Return a python list containing the statistics about the node tree: + 0: number of nodes (tree size) + 1: number of unique nodes + 2: number of trapezoids (tree leaf nodes) + 3: number of unique trapezoids + 4: maximum parent count (max number of times a node is repeated in + tree) + 5: maximum depth of tree (one more than the maximum number of + comparisons needed to search through the tree) + 6: mean of all trapezoid depths (one more than the average number + of comparisons needed to search through the tree) + """ + return self._cpp_trifinder.get_tree_stats() + + def _initialize(self): + """ + Initialize the underlying C++ object. Can be called multiple times if, + for example, the triangulation is modified. + """ + self._cpp_trifinder.initialize() + + def _print_tree(self): + """ + Print a text representation of the node tree, which is useful for + debugging purposes. + """ + self._cpp_trifinder.print_tree() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trifinder.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trifinder.pyi new file mode 100644 index 0000000000000000000000000000000000000000..41a9990b898826c3a5e455f4e91b5b3af6da83e7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trifinder.pyi @@ -0,0 +1,10 @@ +from matplotlib.tri import Triangulation +from numpy.typing import ArrayLike + +class TriFinder: + def __init__(self, triangulation: Triangulation) -> None: ... + def __call__(self, x: ArrayLike, y: ArrayLike) -> ArrayLike: ... + +class TrapezoidMapTriFinder(TriFinder): + def __init__(self, triangulation: Triangulation) -> None: ... + def __call__(self, x: ArrayLike, y: ArrayLike) -> ArrayLike: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triinterpolate.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triinterpolate.py new file mode 100644 index 0000000000000000000000000000000000000000..90ad6cf3a76c45fa0fa2a607c6acee1cab1c5b05 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triinterpolate.py @@ -0,0 +1,1574 @@ +""" +Interpolation inside triangular grids. +""" + +import numpy as np + +from matplotlib import _api +from matplotlib.tri import Triangulation +from matplotlib.tri._trifinder import TriFinder +from matplotlib.tri._tritools import TriAnalyzer + +__all__ = ('TriInterpolator', 'LinearTriInterpolator', 'CubicTriInterpolator') + + +class TriInterpolator: + """ + Abstract base class for classes used to interpolate on a triangular grid. + + Derived classes implement the following methods: + + - ``__call__(x, y)``, + where x, y are array-like point coordinates of the same shape, and + that returns a masked array of the same shape containing the + interpolated z-values. + + - ``gradient(x, y)``, + where x, y are array-like point coordinates of the same + shape, and that returns a list of 2 masked arrays of the same shape + containing the 2 derivatives of the interpolator (derivatives of + interpolated z values with respect to x and y). + """ + + def __init__(self, triangulation, z, trifinder=None): + _api.check_isinstance(Triangulation, triangulation=triangulation) + self._triangulation = triangulation + + self._z = np.asarray(z) + if self._z.shape != self._triangulation.x.shape: + raise ValueError("z array must have same length as triangulation x" + " and y arrays") + + _api.check_isinstance((TriFinder, None), trifinder=trifinder) + self._trifinder = trifinder or self._triangulation.get_trifinder() + + # Default scaling factors : 1.0 (= no scaling) + # Scaling may be used for interpolations for which the order of + # magnitude of x, y has an impact on the interpolant definition. + # Please refer to :meth:`_interpolate_multikeys` for details. + self._unit_x = 1.0 + self._unit_y = 1.0 + + # Default triangle renumbering: None (= no renumbering) + # Renumbering may be used to avoid unnecessary computations + # if complex calculations are done inside the Interpolator. + # Please refer to :meth:`_interpolate_multikeys` for details. + self._tri_renum = None + + # __call__ and gradient docstrings are shared by all subclasses + # (except, if needed, relevant additions). + # However these methods are only implemented in subclasses to avoid + # confusion in the documentation. + _docstring__call__ = """ + Returns a masked array containing interpolated values at the specified + (x, y) points. + + Parameters + ---------- + x, y : array-like + x and y coordinates of the same shape and any number of + dimensions. + + Returns + ------- + np.ma.array + Masked array of the same shape as *x* and *y*; values corresponding + to (*x*, *y*) points outside of the triangulation are masked out. + + """ + + _docstringgradient = r""" + Returns a list of 2 masked arrays containing interpolated derivatives + at the specified (x, y) points. + + Parameters + ---------- + x, y : array-like + x and y coordinates of the same shape and any number of + dimensions. + + Returns + ------- + dzdx, dzdy : np.ma.array + 2 masked arrays of the same shape as *x* and *y*; values + corresponding to (x, y) points outside of the triangulation + are masked out. + The first returned array contains the values of + :math:`\frac{\partial z}{\partial x}` and the second those of + :math:`\frac{\partial z}{\partial y}`. + + """ + + def _interpolate_multikeys(self, x, y, tri_index=None, + return_keys=('z',)): + """ + Versatile (private) method defined for all TriInterpolators. + + :meth:`_interpolate_multikeys` is a wrapper around method + :meth:`_interpolate_single_key` (to be defined in the child + subclasses). + :meth:`_interpolate_single_key actually performs the interpolation, + but only for 1-dimensional inputs and at valid locations (inside + unmasked triangles of the triangulation). + + The purpose of :meth:`_interpolate_multikeys` is to implement the + following common tasks needed in all subclasses implementations: + + - calculation of containing triangles + - dealing with more than one interpolation request at the same + location (e.g., if the 2 derivatives are requested, it is + unnecessary to compute the containing triangles twice) + - scaling according to self._unit_x, self._unit_y + - dealing with points outside of the grid (with fill value np.nan) + - dealing with multi-dimensional *x*, *y* arrays: flattening for + :meth:`_interpolate_params` call and final reshaping. + + (Note that np.vectorize could do most of those things very well for + you, but it does it by function evaluations over successive tuples of + the input arrays. Therefore, this tends to be more time-consuming than + using optimized numpy functions - e.g., np.dot - which can be used + easily on the flattened inputs, in the child-subclass methods + :meth:`_interpolate_single_key`.) + + It is guaranteed that the calls to :meth:`_interpolate_single_key` + will be done with flattened (1-d) array-like input parameters *x*, *y* + and with flattened, valid `tri_index` arrays (no -1 index allowed). + + Parameters + ---------- + x, y : array-like + x and y coordinates where interpolated values are requested. + tri_index : array-like of int, optional + Array of the containing triangle indices, same shape as + *x* and *y*. Defaults to None. If None, these indices + will be computed by a TriFinder instance. + (Note: For point outside the grid, tri_index[ipt] shall be -1). + return_keys : tuple of keys from {'z', 'dzdx', 'dzdy'} + Defines the interpolation arrays to return, and in which order. + + Returns + ------- + list of arrays + Each array-like contains the expected interpolated values in the + order defined by *return_keys* parameter. + """ + # Flattening and rescaling inputs arrays x, y + # (initial shape is stored for output) + x = np.asarray(x, dtype=np.float64) + y = np.asarray(y, dtype=np.float64) + sh_ret = x.shape + if x.shape != y.shape: + raise ValueError("x and y shall have same shapes." + f" Given: {x.shape} and {y.shape}") + x = np.ravel(x) + y = np.ravel(y) + x_scaled = x/self._unit_x + y_scaled = y/self._unit_y + size_ret = np.size(x_scaled) + + # Computes & ravels the element indexes, extract the valid ones. + if tri_index is None: + tri_index = self._trifinder(x, y) + else: + if tri_index.shape != sh_ret: + raise ValueError( + "tri_index array is provided and shall" + " have same shape as x and y. Given: " + f"{tri_index.shape} and {sh_ret}") + tri_index = np.ravel(tri_index) + + mask_in = (tri_index != -1) + if self._tri_renum is None: + valid_tri_index = tri_index[mask_in] + else: + valid_tri_index = self._tri_renum[tri_index[mask_in]] + valid_x = x_scaled[mask_in] + valid_y = y_scaled[mask_in] + + ret = [] + for return_key in return_keys: + # Find the return index associated with the key. + try: + return_index = {'z': 0, 'dzdx': 1, 'dzdy': 2}[return_key] + except KeyError as err: + raise ValueError("return_keys items shall take values in" + " {'z', 'dzdx', 'dzdy'}") from err + + # Sets the scale factor for f & df components + scale = [1., 1./self._unit_x, 1./self._unit_y][return_index] + + # Computes the interpolation + ret_loc = np.empty(size_ret, dtype=np.float64) + ret_loc[~mask_in] = np.nan + ret_loc[mask_in] = self._interpolate_single_key( + return_key, valid_tri_index, valid_x, valid_y) * scale + ret += [np.ma.masked_invalid(ret_loc.reshape(sh_ret), copy=False)] + + return ret + + def _interpolate_single_key(self, return_key, tri_index, x, y): + """ + Interpolate at points belonging to the triangulation + (inside an unmasked triangles). + + Parameters + ---------- + return_key : {'z', 'dzdx', 'dzdy'} + The requested values (z or its derivatives). + tri_index : 1D int array + Valid triangle index (cannot be -1). + x, y : 1D arrays, same shape as `tri_index` + Valid locations where interpolation is requested. + + Returns + ------- + 1-d array + Returned array of the same size as *tri_index* + """ + raise NotImplementedError("TriInterpolator subclasses" + + "should implement _interpolate_single_key!") + + +class LinearTriInterpolator(TriInterpolator): + """ + Linear interpolator on a triangular grid. + + Each triangle is represented by a plane so that an interpolated value at + point (x, y) lies on the plane of the triangle containing (x, y). + Interpolated values are therefore continuous across the triangulation, but + their first derivatives are discontinuous at edges between triangles. + + Parameters + ---------- + triangulation : `~matplotlib.tri.Triangulation` + The triangulation to interpolate over. + z : (npoints,) array-like + Array of values, defined at grid points, to interpolate between. + trifinder : `~matplotlib.tri.TriFinder`, optional + If this is not specified, the Triangulation's default TriFinder will + be used by calling `.Triangulation.get_trifinder`. + + Methods + ------- + `__call__` (x, y) : Returns interpolated values at (x, y) points. + `gradient` (x, y) : Returns interpolated derivatives at (x, y) points. + + """ + def __init__(self, triangulation, z, trifinder=None): + super().__init__(triangulation, z, trifinder) + + # Store plane coefficients for fast interpolation calculations. + self._plane_coefficients = \ + self._triangulation.calculate_plane_coefficients(self._z) + + def __call__(self, x, y): + return self._interpolate_multikeys(x, y, tri_index=None, + return_keys=('z',))[0] + __call__.__doc__ = TriInterpolator._docstring__call__ + + def gradient(self, x, y): + return self._interpolate_multikeys(x, y, tri_index=None, + return_keys=('dzdx', 'dzdy')) + gradient.__doc__ = TriInterpolator._docstringgradient + + def _interpolate_single_key(self, return_key, tri_index, x, y): + _api.check_in_list(['z', 'dzdx', 'dzdy'], return_key=return_key) + if return_key == 'z': + return (self._plane_coefficients[tri_index, 0]*x + + self._plane_coefficients[tri_index, 1]*y + + self._plane_coefficients[tri_index, 2]) + elif return_key == 'dzdx': + return self._plane_coefficients[tri_index, 0] + else: # 'dzdy' + return self._plane_coefficients[tri_index, 1] + + +class CubicTriInterpolator(TriInterpolator): + r""" + Cubic interpolator on a triangular grid. + + In one-dimension - on a segment - a cubic interpolating function is + defined by the values of the function and its derivative at both ends. + This is almost the same in 2D inside a triangle, except that the values + of the function and its 2 derivatives have to be defined at each triangle + node. + + The CubicTriInterpolator takes the value of the function at each node - + provided by the user - and internally computes the value of the + derivatives, resulting in a smooth interpolation. + (As a special feature, the user can also impose the value of the + derivatives at each node, but this is not supposed to be the common + usage.) + + Parameters + ---------- + triangulation : `~matplotlib.tri.Triangulation` + The triangulation to interpolate over. + z : (npoints,) array-like + Array of values, defined at grid points, to interpolate between. + kind : {'min_E', 'geom', 'user'}, optional + Choice of the smoothing algorithm, in order to compute + the interpolant derivatives (defaults to 'min_E'): + + - if 'min_E': (default) The derivatives at each node is computed + to minimize a bending energy. + - if 'geom': The derivatives at each node is computed as a + weighted average of relevant triangle normals. To be used for + speed optimization (large grids). + - if 'user': The user provides the argument *dz*, no computation + is hence needed. + + trifinder : `~matplotlib.tri.TriFinder`, optional + If not specified, the Triangulation's default TriFinder will + be used by calling `.Triangulation.get_trifinder`. + dz : tuple of array-likes (dzdx, dzdy), optional + Used only if *kind* ='user'. In this case *dz* must be provided as + (dzdx, dzdy) where dzdx, dzdy are arrays of the same shape as *z* and + are the interpolant first derivatives at the *triangulation* points. + + Methods + ------- + `__call__` (x, y) : Returns interpolated values at (x, y) points. + `gradient` (x, y) : Returns interpolated derivatives at (x, y) points. + + Notes + ----- + This note is a bit technical and details how the cubic interpolation is + computed. + + The interpolation is based on a Clough-Tocher subdivision scheme of + the *triangulation* mesh (to make it clearer, each triangle of the + grid will be divided in 3 child-triangles, and on each child triangle + the interpolated function is a cubic polynomial of the 2 coordinates). + This technique originates from FEM (Finite Element Method) analysis; + the element used is a reduced Hsieh-Clough-Tocher (HCT) + element. Its shape functions are described in [1]_. + The assembled function is guaranteed to be C1-smooth, i.e. it is + continuous and its first derivatives are also continuous (this + is easy to show inside the triangles but is also true when crossing the + edges). + + In the default case (*kind* ='min_E'), the interpolant minimizes a + curvature energy on the functional space generated by the HCT element + shape functions - with imposed values but arbitrary derivatives at each + node. The minimized functional is the integral of the so-called total + curvature (implementation based on an algorithm from [2]_ - PCG sparse + solver): + + .. math:: + + E(z) = \frac{1}{2} \int_{\Omega} \left( + \left( \frac{\partial^2{z}}{\partial{x}^2} \right)^2 + + \left( \frac{\partial^2{z}}{\partial{y}^2} \right)^2 + + 2\left( \frac{\partial^2{z}}{\partial{y}\partial{x}} \right)^2 + \right) dx\,dy + + If the case *kind* ='geom' is chosen by the user, a simple geometric + approximation is used (weighted average of the triangle normal + vectors), which could improve speed on very large grids. + + References + ---------- + .. [1] Michel Bernadou, Kamal Hassan, "Basis functions for general + Hsieh-Clough-Tocher triangles, complete or reduced.", + International Journal for Numerical Methods in Engineering, + 17(5):784 - 789. 2.01. + .. [2] C.T. Kelley, "Iterative Methods for Optimization". + + """ + def __init__(self, triangulation, z, kind='min_E', trifinder=None, + dz=None): + super().__init__(triangulation, z, trifinder) + + # Loads the underlying c++ _triangulation. + # (During loading, reordering of triangulation._triangles may occur so + # that all final triangles are now anti-clockwise) + self._triangulation.get_cpp_triangulation() + + # To build the stiffness matrix and avoid zero-energy spurious modes + # we will only store internally the valid (unmasked) triangles and + # the necessary (used) points coordinates. + # 2 renumbering tables need to be computed and stored: + # - a triangle renum table in order to translate the result from a + # TriFinder instance into the internal stored triangle number. + # - a node renum table to overwrite the self._z values into the new + # (used) node numbering. + tri_analyzer = TriAnalyzer(self._triangulation) + (compressed_triangles, compressed_x, compressed_y, tri_renum, + node_renum) = tri_analyzer._get_compressed_triangulation() + self._triangles = compressed_triangles + self._tri_renum = tri_renum + # Taking into account the node renumbering in self._z: + valid_node = (node_renum != -1) + self._z[node_renum[valid_node]] = self._z[valid_node] + + # Computing scale factors + self._unit_x = np.ptp(compressed_x) + self._unit_y = np.ptp(compressed_y) + self._pts = np.column_stack([compressed_x / self._unit_x, + compressed_y / self._unit_y]) + # Computing triangle points + self._tris_pts = self._pts[self._triangles] + # Computing eccentricities + self._eccs = self._compute_tri_eccentricities(self._tris_pts) + # Computing dof estimations for HCT triangle shape function + _api.check_in_list(['user', 'geom', 'min_E'], kind=kind) + self._dof = self._compute_dof(kind, dz=dz) + # Loading HCT element + self._ReferenceElement = _ReducedHCT_Element() + + def __call__(self, x, y): + return self._interpolate_multikeys(x, y, tri_index=None, + return_keys=('z',))[0] + __call__.__doc__ = TriInterpolator._docstring__call__ + + def gradient(self, x, y): + return self._interpolate_multikeys(x, y, tri_index=None, + return_keys=('dzdx', 'dzdy')) + gradient.__doc__ = TriInterpolator._docstringgradient + + def _interpolate_single_key(self, return_key, tri_index, x, y): + _api.check_in_list(['z', 'dzdx', 'dzdy'], return_key=return_key) + tris_pts = self._tris_pts[tri_index] + alpha = self._get_alpha_vec(x, y, tris_pts) + ecc = self._eccs[tri_index] + dof = np.expand_dims(self._dof[tri_index], axis=1) + if return_key == 'z': + return self._ReferenceElement.get_function_values( + alpha, ecc, dof) + else: # 'dzdx', 'dzdy' + J = self._get_jacobian(tris_pts) + dzdx = self._ReferenceElement.get_function_derivatives( + alpha, J, ecc, dof) + if return_key == 'dzdx': + return dzdx[:, 0, 0] + else: + return dzdx[:, 1, 0] + + def _compute_dof(self, kind, dz=None): + """ + Compute and return nodal dofs according to kind. + + Parameters + ---------- + kind : {'min_E', 'geom', 'user'} + Choice of the _DOF_estimator subclass to estimate the gradient. + dz : tuple of array-likes (dzdx, dzdy), optional + Used only if *kind*=user; in this case passed to the + :class:`_DOF_estimator_user`. + + Returns + ------- + array-like, shape (npts, 2) + Estimation of the gradient at triangulation nodes (stored as + degree of freedoms of reduced-HCT triangle elements). + """ + if kind == 'user': + if dz is None: + raise ValueError("For a CubicTriInterpolator with " + "*kind*='user', a valid *dz* " + "argument is expected.") + TE = _DOF_estimator_user(self, dz=dz) + elif kind == 'geom': + TE = _DOF_estimator_geom(self) + else: # 'min_E', checked in __init__ + TE = _DOF_estimator_min_E(self) + return TE.compute_dof_from_df() + + @staticmethod + def _get_alpha_vec(x, y, tris_pts): + """ + Fast (vectorized) function to compute barycentric coordinates alpha. + + Parameters + ---------- + x, y : array-like of dim 1 (shape (nx,)) + Coordinates of the points whose points barycentric coordinates are + requested. + tris_pts : array like of dim 3 (shape: (nx, 3, 2)) + Coordinates of the containing triangles apexes. + + Returns + ------- + array of dim 2 (shape (nx, 3)) + Barycentric coordinates of the points inside the containing + triangles. + """ + ndim = tris_pts.ndim-2 + + a = tris_pts[:, 1, :] - tris_pts[:, 0, :] + b = tris_pts[:, 2, :] - tris_pts[:, 0, :] + abT = np.stack([a, b], axis=-1) + ab = _transpose_vectorized(abT) + OM = np.stack([x, y], axis=1) - tris_pts[:, 0, :] + + metric = ab @ abT + # Here we try to deal with the colinear cases. + # metric_inv is in this case set to the Moore-Penrose pseudo-inverse + # meaning that we will still return a set of valid barycentric + # coordinates. + metric_inv = _pseudo_inv22sym_vectorized(metric) + Covar = ab @ _transpose_vectorized(np.expand_dims(OM, ndim)) + ksi = metric_inv @ Covar + alpha = _to_matrix_vectorized([ + [1-ksi[:, 0, 0]-ksi[:, 1, 0]], [ksi[:, 0, 0]], [ksi[:, 1, 0]]]) + return alpha + + @staticmethod + def _get_jacobian(tris_pts): + """ + Fast (vectorized) function to compute triangle jacobian matrix. + + Parameters + ---------- + tris_pts : array like of dim 3 (shape: (nx, 3, 2)) + Coordinates of the containing triangles apexes. + + Returns + ------- + array of dim 3 (shape (nx, 2, 2)) + Barycentric coordinates of the points inside the containing + triangles. + J[itri, :, :] is the jacobian matrix at apex 0 of the triangle + itri, so that the following (matrix) relationship holds: + [dz/dksi] = [J] x [dz/dx] + with x: global coordinates + ksi: element parametric coordinates in triangle first apex + local basis. + """ + a = np.array(tris_pts[:, 1, :] - tris_pts[:, 0, :]) + b = np.array(tris_pts[:, 2, :] - tris_pts[:, 0, :]) + J = _to_matrix_vectorized([[a[:, 0], a[:, 1]], + [b[:, 0], b[:, 1]]]) + return J + + @staticmethod + def _compute_tri_eccentricities(tris_pts): + """ + Compute triangle eccentricities. + + Parameters + ---------- + tris_pts : array like of dim 3 (shape: (nx, 3, 2)) + Coordinates of the triangles apexes. + + Returns + ------- + array like of dim 2 (shape: (nx, 3)) + The so-called eccentricity parameters [1] needed for HCT triangular + element. + """ + a = np.expand_dims(tris_pts[:, 2, :] - tris_pts[:, 1, :], axis=2) + b = np.expand_dims(tris_pts[:, 0, :] - tris_pts[:, 2, :], axis=2) + c = np.expand_dims(tris_pts[:, 1, :] - tris_pts[:, 0, :], axis=2) + # Do not use np.squeeze, this is dangerous if only one triangle + # in the triangulation... + dot_a = (_transpose_vectorized(a) @ a)[:, 0, 0] + dot_b = (_transpose_vectorized(b) @ b)[:, 0, 0] + dot_c = (_transpose_vectorized(c) @ c)[:, 0, 0] + # Note that this line will raise a warning for dot_a, dot_b or dot_c + # zeros, but we choose not to support triangles with duplicate points. + return _to_matrix_vectorized([[(dot_c-dot_b) / dot_a], + [(dot_a-dot_c) / dot_b], + [(dot_b-dot_a) / dot_c]]) + + +# FEM element used for interpolation and for solving minimisation +# problem (Reduced HCT element) +class _ReducedHCT_Element: + """ + Implementation of reduced HCT triangular element with explicit shape + functions. + + Computes z, dz, d2z and the element stiffness matrix for bending energy: + E(f) = integral( (d2z/dx2 + d2z/dy2)**2 dA) + + *** Reference for the shape functions: *** + [1] Basis functions for general Hsieh-Clough-Tocher _triangles, complete or + reduced. + Michel Bernadou, Kamal Hassan + International Journal for Numerical Methods in Engineering. + 17(5):784 - 789. 2.01 + + *** Element description: *** + 9 dofs: z and dz given at 3 apex + C1 (conform) + + """ + # 1) Loads matrices to generate shape functions as a function of + # triangle eccentricities - based on [1] p.11 ''' + M = np.array([ + [ 0.00, 0.00, 0.00, 4.50, 4.50, 0.00, 0.00, 0.00, 0.00, 0.00], + [-0.25, 0.00, 0.00, 0.50, 1.25, 0.00, 0.00, 0.00, 0.00, 0.00], + [-0.25, 0.00, 0.00, 1.25, 0.50, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.50, 1.00, 0.00, -1.50, 0.00, 3.00, 3.00, 0.00, 0.00, 3.00], + [ 0.00, 0.00, 0.00, -0.25, 0.25, 0.00, 1.00, 0.00, 0.00, 0.50], + [ 0.25, 0.00, 0.00, -0.50, -0.25, 1.00, 0.00, 0.00, 0.00, 1.00], + [ 0.50, 0.00, 1.00, 0.00, -1.50, 0.00, 0.00, 3.00, 3.00, 3.00], + [ 0.25, 0.00, 0.00, -0.25, -0.50, 0.00, 0.00, 0.00, 1.00, 1.00], + [ 0.00, 0.00, 0.00, 0.25, -0.25, 0.00, 0.00, 1.00, 0.00, 0.50]]) + M0 = np.array([ + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [-1.00, 0.00, 0.00, 1.50, 1.50, 0.00, 0.00, 0.00, 0.00, -3.00], + [-0.50, 0.00, 0.00, 0.75, 0.75, 0.00, 0.00, 0.00, 0.00, -1.50], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 1.00, 0.00, 0.00, -1.50, -1.50, 0.00, 0.00, 0.00, 0.00, 3.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.50, 0.00, 0.00, -0.75, -0.75, 0.00, 0.00, 0.00, 0.00, 1.50]]) + M1 = np.array([ + [-0.50, 0.00, 0.00, 1.50, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [-0.25, 0.00, 0.00, 0.75, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.50, 0.00, 0.00, -1.50, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.25, 0.00, 0.00, -0.75, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]]) + M2 = np.array([ + [ 0.50, 0.00, 0.00, 0.00, -1.50, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.25, 0.00, 0.00, 0.00, -0.75, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [-0.50, 0.00, 0.00, 0.00, 1.50, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [-0.25, 0.00, 0.00, 0.00, 0.75, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], + [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]]) + + # 2) Loads matrices to rotate components of gradient & Hessian + # vectors in the reference basis of triangle first apex (a0) + rotate_dV = np.array([[ 1., 0.], [ 0., 1.], + [ 0., 1.], [-1., -1.], + [-1., -1.], [ 1., 0.]]) + + rotate_d2V = np.array([[1., 0., 0.], [0., 1., 0.], [ 0., 0., 1.], + [0., 1., 0.], [1., 1., 1.], [ 0., -2., -1.], + [1., 1., 1.], [1., 0., 0.], [-2., 0., -1.]]) + + # 3) Loads Gauss points & weights on the 3 sub-_triangles for P2 + # exact integral - 3 points on each subtriangles. + # NOTE: as the 2nd derivative is discontinuous , we really need those 9 + # points! + n_gauss = 9 + gauss_pts = np.array([[13./18., 4./18., 1./18.], + [ 4./18., 13./18., 1./18.], + [ 7./18., 7./18., 4./18.], + [ 1./18., 13./18., 4./18.], + [ 1./18., 4./18., 13./18.], + [ 4./18., 7./18., 7./18.], + [ 4./18., 1./18., 13./18.], + [13./18., 1./18., 4./18.], + [ 7./18., 4./18., 7./18.]], dtype=np.float64) + gauss_w = np.ones([9], dtype=np.float64) / 9. + + # 4) Stiffness matrix for curvature energy + E = np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 2.]]) + + # 5) Loads the matrix to compute DOF_rot from tri_J at apex 0 + J0_to_J1 = np.array([[-1., 1.], [-1., 0.]]) + J0_to_J2 = np.array([[ 0., -1.], [ 1., -1.]]) + + def get_function_values(self, alpha, ecc, dofs): + """ + Parameters + ---------- + alpha : is a (N x 3 x 1) array (array of column-matrices) of + barycentric coordinates, + ecc : is a (N x 3 x 1) array (array of column-matrices) of triangle + eccentricities, + dofs : is a (N x 1 x 9) arrays (arrays of row-matrices) of computed + degrees of freedom. + + Returns + ------- + Returns the N-array of interpolated function values. + """ + subtri = np.argmin(alpha, axis=1)[:, 0] + ksi = _roll_vectorized(alpha, -subtri, axis=0) + E = _roll_vectorized(ecc, -subtri, axis=0) + x = ksi[:, 0, 0] + y = ksi[:, 1, 0] + z = ksi[:, 2, 0] + x_sq = x*x + y_sq = y*y + z_sq = z*z + V = _to_matrix_vectorized([ + [x_sq*x], [y_sq*y], [z_sq*z], [x_sq*z], [x_sq*y], [y_sq*x], + [y_sq*z], [z_sq*y], [z_sq*x], [x*y*z]]) + prod = self.M @ V + prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ V) + prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ V) + prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ V) + s = _roll_vectorized(prod, 3*subtri, axis=0) + return (dofs @ s)[:, 0, 0] + + def get_function_derivatives(self, alpha, J, ecc, dofs): + """ + Parameters + ---------- + *alpha* is a (N x 3 x 1) array (array of column-matrices of + barycentric coordinates) + *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at + triangle first apex) + *ecc* is a (N x 3 x 1) array (array of column-matrices of triangle + eccentricities) + *dofs* is a (N x 1 x 9) arrays (arrays of row-matrices) of computed + degrees of freedom. + + Returns + ------- + Returns the values of interpolated function derivatives [dz/dx, dz/dy] + in global coordinates at locations alpha, as a column-matrices of + shape (N x 2 x 1). + """ + subtri = np.argmin(alpha, axis=1)[:, 0] + ksi = _roll_vectorized(alpha, -subtri, axis=0) + E = _roll_vectorized(ecc, -subtri, axis=0) + x = ksi[:, 0, 0] + y = ksi[:, 1, 0] + z = ksi[:, 2, 0] + x_sq = x*x + y_sq = y*y + z_sq = z*z + dV = _to_matrix_vectorized([ + [ -3.*x_sq, -3.*x_sq], + [ 3.*y_sq, 0.], + [ 0., 3.*z_sq], + [ -2.*x*z, -2.*x*z+x_sq], + [-2.*x*y+x_sq, -2.*x*y], + [ 2.*x*y-y_sq, -y_sq], + [ 2.*y*z, y_sq], + [ z_sq, 2.*y*z], + [ -z_sq, 2.*x*z-z_sq], + [ x*z-y*z, x*y-y*z]]) + # Puts back dV in first apex basis + dV = dV @ _extract_submatrices( + self.rotate_dV, subtri, block_size=2, axis=0) + + prod = self.M @ dV + prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ dV) + prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ dV) + prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ dV) + dsdksi = _roll_vectorized(prod, 3*subtri, axis=0) + dfdksi = dofs @ dsdksi + # In global coordinates: + # Here we try to deal with the simplest colinear cases, returning a + # null matrix. + J_inv = _safe_inv22_vectorized(J) + dfdx = J_inv @ _transpose_vectorized(dfdksi) + return dfdx + + def get_function_hessians(self, alpha, J, ecc, dofs): + """ + Parameters + ---------- + *alpha* is a (N x 3 x 1) array (array of column-matrices) of + barycentric coordinates + *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at + triangle first apex) + *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle + eccentricities + *dofs* is a (N x 1 x 9) arrays (arrays of row-matrices) of computed + degrees of freedom. + + Returns + ------- + Returns the values of interpolated function 2nd-derivatives + [d2z/dx2, d2z/dy2, d2z/dxdy] in global coordinates at locations alpha, + as a column-matrices of shape (N x 3 x 1). + """ + d2sdksi2 = self.get_d2Sidksij2(alpha, ecc) + d2fdksi2 = dofs @ d2sdksi2 + H_rot = self.get_Hrot_from_J(J) + d2fdx2 = d2fdksi2 @ H_rot + return _transpose_vectorized(d2fdx2) + + def get_d2Sidksij2(self, alpha, ecc): + """ + Parameters + ---------- + *alpha* is a (N x 3 x 1) array (array of column-matrices) of + barycentric coordinates + *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle + eccentricities + + Returns + ------- + Returns the arrays d2sdksi2 (N x 3 x 1) Hessian of shape functions + expressed in covariant coordinates in first apex basis. + """ + subtri = np.argmin(alpha, axis=1)[:, 0] + ksi = _roll_vectorized(alpha, -subtri, axis=0) + E = _roll_vectorized(ecc, -subtri, axis=0) + x = ksi[:, 0, 0] + y = ksi[:, 1, 0] + z = ksi[:, 2, 0] + d2V = _to_matrix_vectorized([ + [ 6.*x, 6.*x, 6.*x], + [ 6.*y, 0., 0.], + [ 0., 6.*z, 0.], + [ 2.*z, 2.*z-4.*x, 2.*z-2.*x], + [2.*y-4.*x, 2.*y, 2.*y-2.*x], + [2.*x-4.*y, 0., -2.*y], + [ 2.*z, 0., 2.*y], + [ 0., 2.*y, 2.*z], + [ 0., 2.*x-4.*z, -2.*z], + [ -2.*z, -2.*y, x-y-z]]) + # Puts back d2V in first apex basis + d2V = d2V @ _extract_submatrices( + self.rotate_d2V, subtri, block_size=3, axis=0) + prod = self.M @ d2V + prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ d2V) + prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ d2V) + prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ d2V) + d2sdksi2 = _roll_vectorized(prod, 3*subtri, axis=0) + return d2sdksi2 + + def get_bending_matrices(self, J, ecc): + """ + Parameters + ---------- + *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at + triangle first apex) + *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle + eccentricities + + Returns + ------- + Returns the element K matrices for bending energy expressed in + GLOBAL nodal coordinates. + K_ij = integral [ (d2zi/dx2 + d2zi/dy2) * (d2zj/dx2 + d2zj/dy2) dA] + tri_J is needed to rotate dofs from local basis to global basis + """ + n = np.size(ecc, 0) + + # 1) matrix to rotate dofs in global coordinates + J1 = self.J0_to_J1 @ J + J2 = self.J0_to_J2 @ J + DOF_rot = np.zeros([n, 9, 9], dtype=np.float64) + DOF_rot[:, 0, 0] = 1 + DOF_rot[:, 3, 3] = 1 + DOF_rot[:, 6, 6] = 1 + DOF_rot[:, 1:3, 1:3] = J + DOF_rot[:, 4:6, 4:6] = J1 + DOF_rot[:, 7:9, 7:9] = J2 + + # 2) matrix to rotate Hessian in global coordinates. + H_rot, area = self.get_Hrot_from_J(J, return_area=True) + + # 3) Computes stiffness matrix + # Gauss quadrature. + K = np.zeros([n, 9, 9], dtype=np.float64) + weights = self.gauss_w + pts = self.gauss_pts + for igauss in range(self.n_gauss): + alpha = np.tile(pts[igauss, :], n).reshape(n, 3) + alpha = np.expand_dims(alpha, 2) + weight = weights[igauss] + d2Skdksi2 = self.get_d2Sidksij2(alpha, ecc) + d2Skdx2 = d2Skdksi2 @ H_rot + K += weight * (d2Skdx2 @ self.E @ _transpose_vectorized(d2Skdx2)) + + # 4) With nodal (not elem) dofs + K = _transpose_vectorized(DOF_rot) @ K @ DOF_rot + + # 5) Need the area to compute total element energy + return _scalar_vectorized(area, K) + + def get_Hrot_from_J(self, J, return_area=False): + """ + Parameters + ---------- + *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at + triangle first apex) + + Returns + ------- + Returns H_rot used to rotate Hessian from local basis of first apex, + to global coordinates. + if *return_area* is True, returns also the triangle area (0.5*det(J)) + """ + # Here we try to deal with the simplest colinear cases; a null + # energy and area is imposed. + J_inv = _safe_inv22_vectorized(J) + Ji00 = J_inv[:, 0, 0] + Ji11 = J_inv[:, 1, 1] + Ji10 = J_inv[:, 1, 0] + Ji01 = J_inv[:, 0, 1] + H_rot = _to_matrix_vectorized([ + [Ji00*Ji00, Ji10*Ji10, Ji00*Ji10], + [Ji01*Ji01, Ji11*Ji11, Ji01*Ji11], + [2*Ji00*Ji01, 2*Ji11*Ji10, Ji00*Ji11+Ji10*Ji01]]) + if not return_area: + return H_rot + else: + area = 0.5 * (J[:, 0, 0]*J[:, 1, 1] - J[:, 0, 1]*J[:, 1, 0]) + return H_rot, area + + def get_Kff_and_Ff(self, J, ecc, triangles, Uc): + """ + Build K and F for the following elliptic formulation: + minimization of curvature energy with value of function at node + imposed and derivatives 'free'. + + Build the global Kff matrix in cco format. + Build the full Ff vec Ff = - Kfc x Uc. + + Parameters + ---------- + *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at + triangle first apex) + *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle + eccentricities + *triangles* is a (N x 3) array of nodes indexes. + *Uc* is (N x 3) array of imposed displacements at nodes + + Returns + ------- + (Kff_rows, Kff_cols, Kff_vals) Kff matrix in coo format - Duplicate + (row, col) entries must be summed. + Ff: force vector - dim npts * 3 + """ + ntri = np.size(ecc, 0) + vec_range = np.arange(ntri, dtype=np.int32) + c_indices = np.full(ntri, -1, dtype=np.int32) # for unused dofs, -1 + f_dof = [1, 2, 4, 5, 7, 8] + c_dof = [0, 3, 6] + + # vals, rows and cols indices in global dof numbering + f_dof_indices = _to_matrix_vectorized([[ + c_indices, triangles[:, 0]*2, triangles[:, 0]*2+1, + c_indices, triangles[:, 1]*2, triangles[:, 1]*2+1, + c_indices, triangles[:, 2]*2, triangles[:, 2]*2+1]]) + + expand_indices = np.ones([ntri, 9, 1], dtype=np.int32) + f_row_indices = _transpose_vectorized(expand_indices @ f_dof_indices) + f_col_indices = expand_indices @ f_dof_indices + K_elem = self.get_bending_matrices(J, ecc) + + # Extracting sub-matrices + # Explanation & notations: + # * Subscript f denotes 'free' degrees of freedom (i.e. dz/dx, dz/dx) + # * Subscript c denotes 'condensated' (imposed) degrees of freedom + # (i.e. z at all nodes) + # * F = [Ff, Fc] is the force vector + # * U = [Uf, Uc] is the imposed dof vector + # [ Kff Kfc ] + # * K = [ ] is the laplacian stiffness matrix + # [ Kcf Kff ] + # * As F = K x U one gets straightforwardly: Ff = - Kfc x Uc + + # Computing Kff stiffness matrix in sparse coo format + Kff_vals = np.ravel(K_elem[np.ix_(vec_range, f_dof, f_dof)]) + Kff_rows = np.ravel(f_row_indices[np.ix_(vec_range, f_dof, f_dof)]) + Kff_cols = np.ravel(f_col_indices[np.ix_(vec_range, f_dof, f_dof)]) + + # Computing Ff force vector in sparse coo format + Kfc_elem = K_elem[np.ix_(vec_range, f_dof, c_dof)] + Uc_elem = np.expand_dims(Uc, axis=2) + Ff_elem = -(Kfc_elem @ Uc_elem)[:, :, 0] + Ff_indices = f_dof_indices[np.ix_(vec_range, [0], f_dof)][:, 0, :] + + # Extracting Ff force vector in dense format + # We have to sum duplicate indices - using bincount + Ff = np.bincount(np.ravel(Ff_indices), weights=np.ravel(Ff_elem)) + return Kff_rows, Kff_cols, Kff_vals, Ff + + +# :class:_DOF_estimator, _DOF_estimator_user, _DOF_estimator_geom, +# _DOF_estimator_min_E +# Private classes used to compute the degree of freedom of each triangular +# element for the TriCubicInterpolator. +class _DOF_estimator: + """ + Abstract base class for classes used to estimate a function's first + derivatives, and deduce the dofs for a CubicTriInterpolator using a + reduced HCT element formulation. + + Derived classes implement ``compute_df(self, **kwargs)``, returning + ``np.vstack([dfx, dfy]).T`` where ``dfx, dfy`` are the estimation of the 2 + gradient coordinates. + """ + def __init__(self, interpolator, **kwargs): + _api.check_isinstance(CubicTriInterpolator, interpolator=interpolator) + self._pts = interpolator._pts + self._tris_pts = interpolator._tris_pts + self.z = interpolator._z + self._triangles = interpolator._triangles + (self._unit_x, self._unit_y) = (interpolator._unit_x, + interpolator._unit_y) + self.dz = self.compute_dz(**kwargs) + self.compute_dof_from_df() + + def compute_dz(self, **kwargs): + raise NotImplementedError + + def compute_dof_from_df(self): + """ + Compute reduced-HCT elements degrees of freedom, from the gradient. + """ + J = CubicTriInterpolator._get_jacobian(self._tris_pts) + tri_z = self.z[self._triangles] + tri_dz = self.dz[self._triangles] + tri_dof = self.get_dof_vec(tri_z, tri_dz, J) + return tri_dof + + @staticmethod + def get_dof_vec(tri_z, tri_dz, J): + """ + Compute the dof vector of a triangle, from the value of f, df and + of the local Jacobian at each node. + + Parameters + ---------- + tri_z : shape (3,) array + f nodal values. + tri_dz : shape (3, 2) array + df/dx, df/dy nodal values. + J + Jacobian matrix in local basis of apex 0. + + Returns + ------- + dof : shape (9,) array + For each apex ``iapex``:: + + dof[iapex*3+0] = f(Ai) + dof[iapex*3+1] = df(Ai).(AiAi+) + dof[iapex*3+2] = df(Ai).(AiAi-) + """ + npt = tri_z.shape[0] + dof = np.zeros([npt, 9], dtype=np.float64) + J1 = _ReducedHCT_Element.J0_to_J1 @ J + J2 = _ReducedHCT_Element.J0_to_J2 @ J + + col0 = J @ np.expand_dims(tri_dz[:, 0, :], axis=2) + col1 = J1 @ np.expand_dims(tri_dz[:, 1, :], axis=2) + col2 = J2 @ np.expand_dims(tri_dz[:, 2, :], axis=2) + + dfdksi = _to_matrix_vectorized([ + [col0[:, 0, 0], col1[:, 0, 0], col2[:, 0, 0]], + [col0[:, 1, 0], col1[:, 1, 0], col2[:, 1, 0]]]) + dof[:, 0:7:3] = tri_z + dof[:, 1:8:3] = dfdksi[:, 0] + dof[:, 2:9:3] = dfdksi[:, 1] + return dof + + +class _DOF_estimator_user(_DOF_estimator): + """dz is imposed by user; accounts for scaling if any.""" + + def compute_dz(self, dz): + (dzdx, dzdy) = dz + dzdx = dzdx * self._unit_x + dzdy = dzdy * self._unit_y + return np.vstack([dzdx, dzdy]).T + + +class _DOF_estimator_geom(_DOF_estimator): + """Fast 'geometric' approximation, recommended for large arrays.""" + + def compute_dz(self): + """ + self.df is computed as weighted average of _triangles sharing a common + node. On each triangle itri f is first assumed linear (= ~f), which + allows to compute d~f[itri] + Then the following approximation of df nodal values is then proposed: + f[ipt] = SUM ( w[itri] x d~f[itri] , for itri sharing apex ipt) + The weighted coeff. w[itri] are proportional to the angle of the + triangle itri at apex ipt + """ + el_geom_w = self.compute_geom_weights() + el_geom_grad = self.compute_geom_grads() + + # Sum of weights coeffs + w_node_sum = np.bincount(np.ravel(self._triangles), + weights=np.ravel(el_geom_w)) + + # Sum of weighted df = (dfx, dfy) + dfx_el_w = np.empty_like(el_geom_w) + dfy_el_w = np.empty_like(el_geom_w) + for iapex in range(3): + dfx_el_w[:, iapex] = el_geom_w[:, iapex]*el_geom_grad[:, 0] + dfy_el_w[:, iapex] = el_geom_w[:, iapex]*el_geom_grad[:, 1] + dfx_node_sum = np.bincount(np.ravel(self._triangles), + weights=np.ravel(dfx_el_w)) + dfy_node_sum = np.bincount(np.ravel(self._triangles), + weights=np.ravel(dfy_el_w)) + + # Estimation of df + dfx_estim = dfx_node_sum/w_node_sum + dfy_estim = dfy_node_sum/w_node_sum + return np.vstack([dfx_estim, dfy_estim]).T + + def compute_geom_weights(self): + """ + Build the (nelems, 3) weights coeffs of _triangles angles, + renormalized so that np.sum(weights, axis=1) == np.ones(nelems) + """ + weights = np.zeros([np.size(self._triangles, 0), 3]) + tris_pts = self._tris_pts + for ipt in range(3): + p0 = tris_pts[:, ipt % 3, :] + p1 = tris_pts[:, (ipt+1) % 3, :] + p2 = tris_pts[:, (ipt-1) % 3, :] + alpha1 = np.arctan2(p1[:, 1]-p0[:, 1], p1[:, 0]-p0[:, 0]) + alpha2 = np.arctan2(p2[:, 1]-p0[:, 1], p2[:, 0]-p0[:, 0]) + # In the below formula we could take modulo 2. but + # modulo 1. is safer regarding round-off errors (flat triangles). + angle = np.abs(((alpha2-alpha1) / np.pi) % 1) + # Weight proportional to angle up np.pi/2; null weight for + # degenerated cases 0 and np.pi (note that *angle* is normalized + # by np.pi). + weights[:, ipt] = 0.5 - np.abs(angle-0.5) + return weights + + def compute_geom_grads(self): + """ + Compute the (global) gradient component of f assumed linear (~f). + returns array df of shape (nelems, 2) + df[ielem].dM[ielem] = dz[ielem] i.e. df = dz x dM = dM.T^-1 x dz + """ + tris_pts = self._tris_pts + tris_f = self.z[self._triangles] + + dM1 = tris_pts[:, 1, :] - tris_pts[:, 0, :] + dM2 = tris_pts[:, 2, :] - tris_pts[:, 0, :] + dM = np.dstack([dM1, dM2]) + # Here we try to deal with the simplest colinear cases: a null + # gradient is assumed in this case. + dM_inv = _safe_inv22_vectorized(dM) + + dZ1 = tris_f[:, 1] - tris_f[:, 0] + dZ2 = tris_f[:, 2] - tris_f[:, 0] + dZ = np.vstack([dZ1, dZ2]).T + df = np.empty_like(dZ) + + # With np.einsum: could be ej,eji -> ej + df[:, 0] = dZ[:, 0]*dM_inv[:, 0, 0] + dZ[:, 1]*dM_inv[:, 1, 0] + df[:, 1] = dZ[:, 0]*dM_inv[:, 0, 1] + dZ[:, 1]*dM_inv[:, 1, 1] + return df + + +class _DOF_estimator_min_E(_DOF_estimator_geom): + """ + The 'smoothest' approximation, df is computed through global minimization + of the bending energy: + E(f) = integral[(d2z/dx2 + d2z/dy2 + 2 d2z/dxdy)**2 dA] + """ + def __init__(self, Interpolator): + self._eccs = Interpolator._eccs + super().__init__(Interpolator) + + def compute_dz(self): + """ + Elliptic solver for bending energy minimization. + Uses a dedicated 'toy' sparse Jacobi PCG solver. + """ + # Initial guess for iterative PCG solver. + dz_init = super().compute_dz() + Uf0 = np.ravel(dz_init) + + reference_element = _ReducedHCT_Element() + J = CubicTriInterpolator._get_jacobian(self._tris_pts) + eccs = self._eccs + triangles = self._triangles + Uc = self.z[self._triangles] + + # Building stiffness matrix and force vector in coo format + Kff_rows, Kff_cols, Kff_vals, Ff = reference_element.get_Kff_and_Ff( + J, eccs, triangles, Uc) + + # Building sparse matrix and solving minimization problem + # We could use scipy.sparse direct solver; however to avoid this + # external dependency an implementation of a simple PCG solver with + # a simple diagonal Jacobi preconditioner is implemented. + tol = 1.e-10 + n_dof = Ff.shape[0] + Kff_coo = _Sparse_Matrix_coo(Kff_vals, Kff_rows, Kff_cols, + shape=(n_dof, n_dof)) + Kff_coo.compress_csc() + Uf, err = _cg(A=Kff_coo, b=Ff, x0=Uf0, tol=tol) + # If the PCG did not converge, we return the best guess between Uf0 + # and Uf. + err0 = np.linalg.norm(Kff_coo.dot(Uf0) - Ff) + if err0 < err: + # Maybe a good occasion to raise a warning here ? + _api.warn_external("In TriCubicInterpolator initialization, " + "PCG sparse solver did not converge after " + "1000 iterations. `geom` approximation is " + "used instead of `min_E`") + Uf = Uf0 + + # Building dz from Uf + dz = np.empty([self._pts.shape[0], 2], dtype=np.float64) + dz[:, 0] = Uf[::2] + dz[:, 1] = Uf[1::2] + return dz + + +# The following private :class:_Sparse_Matrix_coo and :func:_cg provide +# a PCG sparse solver for (symmetric) elliptic problems. +class _Sparse_Matrix_coo: + def __init__(self, vals, rows, cols, shape): + """ + Create a sparse matrix in coo format. + *vals*: arrays of values of non-null entries of the matrix + *rows*: int arrays of rows of non-null entries of the matrix + *cols*: int arrays of cols of non-null entries of the matrix + *shape*: 2-tuple (n, m) of matrix shape + """ + self.n, self.m = shape + self.vals = np.asarray(vals, dtype=np.float64) + self.rows = np.asarray(rows, dtype=np.int32) + self.cols = np.asarray(cols, dtype=np.int32) + + def dot(self, V): + """ + Dot product of self by a vector *V* in sparse-dense to dense format + *V* dense vector of shape (self.m,). + """ + assert V.shape == (self.m,) + return np.bincount(self.rows, + weights=self.vals*V[self.cols], + minlength=self.m) + + def compress_csc(self): + """ + Compress rows, cols, vals / summing duplicates. Sort for csc format. + """ + _, unique, indices = np.unique( + self.rows + self.n*self.cols, + return_index=True, return_inverse=True) + self.rows = self.rows[unique] + self.cols = self.cols[unique] + self.vals = np.bincount(indices, weights=self.vals) + + def compress_csr(self): + """ + Compress rows, cols, vals / summing duplicates. Sort for csr format. + """ + _, unique, indices = np.unique( + self.m*self.rows + self.cols, + return_index=True, return_inverse=True) + self.rows = self.rows[unique] + self.cols = self.cols[unique] + self.vals = np.bincount(indices, weights=self.vals) + + def to_dense(self): + """ + Return a dense matrix representing self, mainly for debugging purposes. + """ + ret = np.zeros([self.n, self.m], dtype=np.float64) + nvals = self.vals.size + for i in range(nvals): + ret[self.rows[i], self.cols[i]] += self.vals[i] + return ret + + def __str__(self): + return self.to_dense().__str__() + + @property + def diag(self): + """Return the (dense) vector of the diagonal elements.""" + in_diag = (self.rows == self.cols) + diag = np.zeros(min(self.n, self.n), dtype=np.float64) # default 0. + diag[self.rows[in_diag]] = self.vals[in_diag] + return diag + + +def _cg(A, b, x0=None, tol=1.e-10, maxiter=1000): + """ + Use Preconditioned Conjugate Gradient iteration to solve A x = b + A simple Jacobi (diagonal) preconditioner is used. + + Parameters + ---------- + A : _Sparse_Matrix_coo + *A* must have been compressed before by compress_csc or + compress_csr method. + b : array + Right hand side of the linear system. + x0 : array, optional + Starting guess for the solution. Defaults to the zero vector. + tol : float, optional + Tolerance to achieve. The algorithm terminates when the relative + residual is below tol. Default is 1e-10. + maxiter : int, optional + Maximum number of iterations. Iteration will stop after *maxiter* + steps even if the specified tolerance has not been achieved. Defaults + to 1000. + + Returns + ------- + x : array + The converged solution. + err : float + The absolute error np.linalg.norm(A.dot(x) - b) + """ + n = b.size + assert A.n == n + assert A.m == n + b_norm = np.linalg.norm(b) + + # Jacobi pre-conditioner + kvec = A.diag + # For diag elem < 1e-6 we keep 1e-6. + kvec = np.maximum(kvec, 1e-6) + + # Initial guess + if x0 is None: + x = np.zeros(n) + else: + x = x0 + + r = b - A.dot(x) + w = r/kvec + + p = np.zeros(n) + beta = 0.0 + rho = np.dot(r, w) + k = 0 + + # Following C. T. Kelley + while (np.sqrt(abs(rho)) > tol*b_norm) and (k < maxiter): + p = w + beta*p + z = A.dot(p) + alpha = rho/np.dot(p, z) + r = r - alpha*z + w = r/kvec + rhoold = rho + rho = np.dot(r, w) + x = x + alpha*p + beta = rho/rhoold + # err = np.linalg.norm(A.dot(x) - b) # absolute accuracy - not used + k += 1 + err = np.linalg.norm(A.dot(x) - b) + return x, err + + +# The following private functions: +# :func:`_safe_inv22_vectorized` +# :func:`_pseudo_inv22sym_vectorized` +# :func:`_scalar_vectorized` +# :func:`_transpose_vectorized` +# :func:`_roll_vectorized` +# :func:`_to_matrix_vectorized` +# :func:`_extract_submatrices` +# provide fast numpy implementation of some standard operations on arrays of +# matrices - stored as (:, n_rows, n_cols)-shaped np.arrays. + +# Development note: Dealing with pathologic 'flat' triangles in the +# CubicTriInterpolator code and impact on (2, 2)-matrix inversion functions +# :func:`_safe_inv22_vectorized` and :func:`_pseudo_inv22sym_vectorized`. +# +# Goals: +# 1) The CubicTriInterpolator should be able to handle flat or almost flat +# triangles without raising an error, +# 2) These degenerated triangles should have no impact on the automatic dof +# calculation (associated with null weight for the _DOF_estimator_geom and +# with null energy for the _DOF_estimator_min_E), +# 3) Linear patch test should be passed exactly on degenerated meshes, +# 4) Interpolation (with :meth:`_interpolate_single_key` or +# :meth:`_interpolate_multi_key`) shall be correctly handled even *inside* +# the pathologic triangles, to interact correctly with a TriRefiner class. +# +# Difficulties: +# Flat triangles have rank-deficient *J* (so-called jacobian matrix) and +# *metric* (the metric tensor = J x J.T). Computation of the local +# tangent plane is also problematic. +# +# Implementation: +# Most of the time, when computing the inverse of a rank-deficient matrix it +# is safe to simply return the null matrix (which is the implementation in +# :func:`_safe_inv22_vectorized`). This is because of point 2), itself +# enforced by: +# - null area hence null energy in :class:`_DOF_estimator_min_E` +# - angles close or equal to 0 or np.pi hence null weight in +# :class:`_DOF_estimator_geom`. +# Note that the function angle -> weight is continuous and maximum for an +# angle np.pi/2 (refer to :meth:`compute_geom_weights`) +# The exception is the computation of barycentric coordinates, which is done +# by inversion of the *metric* matrix. In this case, we need to compute a set +# of valid coordinates (1 among numerous possibilities), to ensure point 4). +# We benefit here from the symmetry of metric = J x J.T, which makes it easier +# to compute a pseudo-inverse in :func:`_pseudo_inv22sym_vectorized` +def _safe_inv22_vectorized(M): + """ + Inversion of arrays of (2, 2) matrices, returns 0 for rank-deficient + matrices. + + *M* : array of (2, 2) matrices to inverse, shape (n, 2, 2) + """ + _api.check_shape((None, 2, 2), M=M) + M_inv = np.empty_like(M) + prod1 = M[:, 0, 0]*M[:, 1, 1] + delta = prod1 - M[:, 0, 1]*M[:, 1, 0] + + # We set delta_inv to 0. in case of a rank deficient matrix; a + # rank-deficient input matrix *M* will lead to a null matrix in output + rank2 = (np.abs(delta) > 1e-8*np.abs(prod1)) + if np.all(rank2): + # Normal 'optimized' flow. + delta_inv = 1./delta + else: + # 'Pathologic' flow. + delta_inv = np.zeros(M.shape[0]) + delta_inv[rank2] = 1./delta[rank2] + + M_inv[:, 0, 0] = M[:, 1, 1]*delta_inv + M_inv[:, 0, 1] = -M[:, 0, 1]*delta_inv + M_inv[:, 1, 0] = -M[:, 1, 0]*delta_inv + M_inv[:, 1, 1] = M[:, 0, 0]*delta_inv + return M_inv + + +def _pseudo_inv22sym_vectorized(M): + """ + Inversion of arrays of (2, 2) SYMMETRIC matrices; returns the + (Moore-Penrose) pseudo-inverse for rank-deficient matrices. + + In case M is of rank 1, we have M = trace(M) x P where P is the orthogonal + projection on Im(M), and we return trace(M)^-1 x P == M / trace(M)**2 + In case M is of rank 0, we return the null matrix. + + *M* : array of (2, 2) matrices to inverse, shape (n, 2, 2) + """ + _api.check_shape((None, 2, 2), M=M) + M_inv = np.empty_like(M) + prod1 = M[:, 0, 0]*M[:, 1, 1] + delta = prod1 - M[:, 0, 1]*M[:, 1, 0] + rank2 = (np.abs(delta) > 1e-8*np.abs(prod1)) + + if np.all(rank2): + # Normal 'optimized' flow. + M_inv[:, 0, 0] = M[:, 1, 1] / delta + M_inv[:, 0, 1] = -M[:, 0, 1] / delta + M_inv[:, 1, 0] = -M[:, 1, 0] / delta + M_inv[:, 1, 1] = M[:, 0, 0] / delta + else: + # 'Pathologic' flow. + # Here we have to deal with 2 sub-cases + # 1) First sub-case: matrices of rank 2: + delta = delta[rank2] + M_inv[rank2, 0, 0] = M[rank2, 1, 1] / delta + M_inv[rank2, 0, 1] = -M[rank2, 0, 1] / delta + M_inv[rank2, 1, 0] = -M[rank2, 1, 0] / delta + M_inv[rank2, 1, 1] = M[rank2, 0, 0] / delta + # 2) Second sub-case: rank-deficient matrices of rank 0 and 1: + rank01 = ~rank2 + tr = M[rank01, 0, 0] + M[rank01, 1, 1] + tr_zeros = (np.abs(tr) < 1.e-8) + sq_tr_inv = (1.-tr_zeros) / (tr**2+tr_zeros) + # sq_tr_inv = 1. / tr**2 + M_inv[rank01, 0, 0] = M[rank01, 0, 0] * sq_tr_inv + M_inv[rank01, 0, 1] = M[rank01, 0, 1] * sq_tr_inv + M_inv[rank01, 1, 0] = M[rank01, 1, 0] * sq_tr_inv + M_inv[rank01, 1, 1] = M[rank01, 1, 1] * sq_tr_inv + + return M_inv + + +def _scalar_vectorized(scalar, M): + """ + Scalar product between scalars and matrices. + """ + return scalar[:, np.newaxis, np.newaxis]*M + + +def _transpose_vectorized(M): + """ + Transposition of an array of matrices *M*. + """ + return np.transpose(M, [0, 2, 1]) + + +def _roll_vectorized(M, roll_indices, axis): + """ + Roll an array of matrices along *axis* (0: rows, 1: columns) according to + an array of indices *roll_indices*. + """ + assert axis in [0, 1] + ndim = M.ndim + assert ndim == 3 + ndim_roll = roll_indices.ndim + assert ndim_roll == 1 + sh = M.shape + r, c = sh[-2:] + assert sh[0] == roll_indices.shape[0] + vec_indices = np.arange(sh[0], dtype=np.int32) + + # Builds the rolled matrix + M_roll = np.empty_like(M) + if axis == 0: + for ir in range(r): + for ic in range(c): + M_roll[:, ir, ic] = M[vec_indices, (-roll_indices+ir) % r, ic] + else: # 1 + for ir in range(r): + for ic in range(c): + M_roll[:, ir, ic] = M[vec_indices, ir, (-roll_indices+ic) % c] + return M_roll + + +def _to_matrix_vectorized(M): + """ + Build an array of matrices from individuals np.arrays of identical shapes. + + Parameters + ---------- + M + ncols-list of nrows-lists of shape sh. + + Returns + ------- + M_res : np.array of shape (sh, nrow, ncols) + *M_res* satisfies ``M_res[..., i, j] = M[i][j]``. + """ + assert isinstance(M, (tuple, list)) + assert all(isinstance(item, (tuple, list)) for item in M) + c_vec = np.asarray([len(item) for item in M]) + assert np.all(c_vec-c_vec[0] == 0) + r = len(M) + c = c_vec[0] + M00 = np.asarray(M[0][0]) + dt = M00.dtype + sh = [M00.shape[0], r, c] + M_ret = np.empty(sh, dtype=dt) + for irow in range(r): + for icol in range(c): + M_ret[:, irow, icol] = np.asarray(M[irow][icol]) + return M_ret + + +def _extract_submatrices(M, block_indices, block_size, axis): + """ + Extract selected blocks of a matrices *M* depending on parameters + *block_indices* and *block_size*. + + Returns the array of extracted matrices *Mres* so that :: + + M_res[..., ir, :] = M[(block_indices*block_size+ir), :] + """ + assert block_indices.ndim == 1 + assert axis in [0, 1] + + r, c = M.shape + if axis == 0: + sh = [block_indices.shape[0], block_size, c] + else: # 1 + sh = [block_indices.shape[0], r, block_size] + + dt = M.dtype + M_res = np.empty(sh, dtype=dt) + if axis == 0: + for ir in range(block_size): + M_res[:, ir, :] = M[(block_indices*block_size+ir), :] + else: # 1 + for ic in range(block_size): + M_res[:, :, ic] = M[:, (block_indices*block_size+ic)] + + return M_res diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triinterpolate.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triinterpolate.pyi new file mode 100644 index 0000000000000000000000000000000000000000..33b2fd8be4cdac2437ead73c6be972902edfcb4f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triinterpolate.pyi @@ -0,0 +1,32 @@ +from matplotlib.tri import Triangulation, TriFinder + +from typing import Literal +import numpy as np +from numpy.typing import ArrayLike + +class TriInterpolator: + def __init__( + self, + triangulation: Triangulation, + z: ArrayLike, + trifinder: TriFinder | None = ..., + ) -> None: ... + # __call__ and gradient are not actually implemented by the ABC, but are specified as required + def __call__(self, x: ArrayLike, y: ArrayLike) -> np.ma.MaskedArray: ... + def gradient( + self, x: ArrayLike, y: ArrayLike + ) -> tuple[np.ma.MaskedArray, np.ma.MaskedArray]: ... + +class LinearTriInterpolator(TriInterpolator): ... + +class CubicTriInterpolator(TriInterpolator): + def __init__( + self, + triangulation: Triangulation, + z: ArrayLike, + kind: Literal["min_E", "geom", "user"] = ..., + trifinder: TriFinder | None = ..., + dz: tuple[ArrayLike, ArrayLike] | None = ..., + ) -> None: ... + +__all__ = ('TriInterpolator', 'LinearTriInterpolator', 'CubicTriInterpolator') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tripcolor.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tripcolor.py new file mode 100644 index 0000000000000000000000000000000000000000..f3c26b0b25ffa7bd5b449c77f765dd5e66d4801c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tripcolor.py @@ -0,0 +1,167 @@ +import numpy as np + +from matplotlib import _api, _docstring +from matplotlib.collections import PolyCollection, TriMesh +from matplotlib.tri._triangulation import Triangulation + + +@_docstring.interpd +def tripcolor(ax, *args, alpha=1.0, norm=None, cmap=None, vmin=None, + vmax=None, shading='flat', facecolors=None, **kwargs): + """ + Create a pseudocolor plot of an unstructured triangular grid. + + Call signatures:: + + tripcolor(triangulation, c, *, ...) + tripcolor(x, y, c, *, [triangles=triangles], [mask=mask], ...) + + The triangular grid can be specified either by passing a `.Triangulation` + object as the first parameter, or by passing the points *x*, *y* and + optionally the *triangles* and a *mask*. See `.Triangulation` for an + explanation of these parameters. + + It is possible to pass the triangles positionally, i.e. + ``tripcolor(x, y, triangles, c, ...)``. However, this is discouraged. + For more clarity, pass *triangles* via keyword argument. + + If neither of *triangulation* or *triangles* are given, the triangulation + is calculated on the fly. In this case, it does not make sense to provide + colors at the triangle faces via *c* or *facecolors* because there are + multiple possible triangulations for a group of points and you don't know + which triangles will be constructed. + + Parameters + ---------- + triangulation : `.Triangulation` + An already created triangular grid. + x, y, triangles, mask + Parameters defining the triangular grid. See `.Triangulation`. + This is mutually exclusive with specifying *triangulation*. + c : array-like + The color values, either for the points or for the triangles. Which one + is automatically inferred from the length of *c*, i.e. does it match + the number of points or the number of triangles. If there are the same + number of points and triangles in the triangulation it is assumed that + color values are defined at points; to force the use of color values at + triangles use the keyword argument ``facecolors=c`` instead of just + ``c``. + This parameter is position-only. + facecolors : array-like, optional + Can be used alternatively to *c* to specify colors at the triangle + faces. This parameter takes precedence over *c*. + shading : {'flat', 'gouraud'}, default: 'flat' + If 'flat' and the color values *c* are defined at points, the color + values used for each triangle are from the mean c of the triangle's + three points. If *shading* is 'gouraud' then color values must be + defined at points. + %(cmap_doc)s + + %(norm_doc)s + + %(vmin_vmax_doc)s + + %(colorizer_doc)s + + Returns + ------- + `~matplotlib.collections.PolyCollection` or `~matplotlib.collections.TriMesh` + The result depends on *shading*: For ``shading='flat'`` the result is a + `.PolyCollection`, for ``shading='gouraud'`` the result is a `.TriMesh`. + + Other Parameters + ---------------- + **kwargs : `~matplotlib.collections.Collection` properties + + %(Collection:kwdoc)s + """ + _api.check_in_list(['flat', 'gouraud'], shading=shading) + + tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs) + + # Parse the color to be in one of (the other variable will be None): + # - facecolors: if specified at the triangle faces + # - point_colors: if specified at the points + if facecolors is not None: + if args: + _api.warn_external( + "Positional parameter c has no effect when the keyword " + "facecolors is given") + point_colors = None + if len(facecolors) != len(tri.triangles): + raise ValueError("The length of facecolors must match the number " + "of triangles") + else: + # Color from positional parameter c + if not args: + raise TypeError( + "tripcolor() missing 1 required positional argument: 'c'; or " + "1 required keyword-only argument: 'facecolors'") + elif len(args) > 1: + raise TypeError(f"Unexpected positional parameters: {args[1:]!r}") + c = np.asarray(args[0]) + if len(c) == len(tri.x): + # having this before the len(tri.triangles) comparison gives + # precedence to nodes if there are as many nodes as triangles + point_colors = c + facecolors = None + elif len(c) == len(tri.triangles): + point_colors = None + facecolors = c + else: + raise ValueError('The length of c must match either the number ' + 'of points or the number of triangles') + + # Handling of linewidths, shading, edgecolors and antialiased as + # in Axes.pcolor + linewidths = (0.25,) + if 'linewidth' in kwargs: + kwargs['linewidths'] = kwargs.pop('linewidth') + kwargs.setdefault('linewidths', linewidths) + + edgecolors = 'none' + if 'edgecolor' in kwargs: + kwargs['edgecolors'] = kwargs.pop('edgecolor') + ec = kwargs.setdefault('edgecolors', edgecolors) + + if 'antialiased' in kwargs: + kwargs['antialiaseds'] = kwargs.pop('antialiased') + if 'antialiaseds' not in kwargs and ec.lower() == "none": + kwargs['antialiaseds'] = False + + if shading == 'gouraud': + if facecolors is not None: + raise ValueError( + "shading='gouraud' can only be used when the colors " + "are specified at the points, not at the faces.") + collection = TriMesh(tri, alpha=alpha, array=point_colors, + cmap=cmap, norm=norm, **kwargs) + else: # 'flat' + # Vertices of triangles. + maskedTris = tri.get_masked_triangles() + verts = np.stack((tri.x[maskedTris], tri.y[maskedTris]), axis=-1) + + # Color values. + if facecolors is None: + # One color per triangle, the mean of the 3 vertex color values. + colors = point_colors[maskedTris].mean(axis=1) + elif tri.mask is not None: + # Remove color values of masked triangles. + colors = facecolors[~tri.mask] + else: + colors = facecolors + collection = PolyCollection(verts, alpha=alpha, array=colors, + cmap=cmap, norm=norm, **kwargs) + + collection._scale_norm(norm, vmin, vmax) + ax.grid(False) + + minx = tri.x.min() + maxx = tri.x.max() + miny = tri.y.min() + maxy = tri.y.max() + corners = (minx, miny), (maxx, maxy) + ax.update_datalim(corners) + ax.autoscale_view() + ax.add_collection(collection) + return collection diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tripcolor.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tripcolor.pyi new file mode 100644 index 0000000000000000000000000000000000000000..42acffcdc52e0215508a5038ffa292396a457549 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tripcolor.pyi @@ -0,0 +1,71 @@ +from matplotlib.axes import Axes +from matplotlib.collections import PolyCollection, TriMesh +from matplotlib.colors import Normalize, Colormap +from matplotlib.tri._triangulation import Triangulation + +from numpy.typing import ArrayLike + +from typing import overload, Literal + +@overload +def tripcolor( + ax: Axes, + triangulation: Triangulation, + c: ArrayLike = ..., + *, + alpha: float = ..., + norm: str | Normalize | None = ..., + cmap: str | Colormap | None = ..., + vmin: float | None = ..., + vmax: float | None = ..., + shading: Literal["flat"] = ..., + facecolors: ArrayLike | None = ..., + **kwargs +) -> PolyCollection: ... +@overload +def tripcolor( + ax: Axes, + x: ArrayLike, + y: ArrayLike, + c: ArrayLike = ..., + *, + alpha: float = ..., + norm: str | Normalize | None = ..., + cmap: str | Colormap | None = ..., + vmin: float | None = ..., + vmax: float | None = ..., + shading: Literal["flat"] = ..., + facecolors: ArrayLike | None = ..., + **kwargs +) -> PolyCollection: ... +@overload +def tripcolor( + ax: Axes, + triangulation: Triangulation, + c: ArrayLike = ..., + *, + alpha: float = ..., + norm: str | Normalize | None = ..., + cmap: str | Colormap | None = ..., + vmin: float | None = ..., + vmax: float | None = ..., + shading: Literal["gouraud"], + facecolors: ArrayLike | None = ..., + **kwargs +) -> TriMesh: ... +@overload +def tripcolor( + ax: Axes, + x: ArrayLike, + y: ArrayLike, + c: ArrayLike = ..., + *, + alpha: float = ..., + norm: str | Normalize | None = ..., + cmap: str | Colormap | None = ..., + vmin: float | None = ..., + vmax: float | None = ..., + shading: Literal["gouraud"], + facecolors: ArrayLike | None = ..., + **kwargs +) -> TriMesh: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triplot.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triplot.py new file mode 100644 index 0000000000000000000000000000000000000000..6168946b153180f8e6439a01885263eaa017280a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triplot.py @@ -0,0 +1,86 @@ +import numpy as np +from matplotlib.tri._triangulation import Triangulation +import matplotlib.cbook as cbook +import matplotlib.lines as mlines + + +def triplot(ax, *args, **kwargs): + """ + Draw an unstructured triangular grid as lines and/or markers. + + Call signatures:: + + triplot(triangulation, ...) + triplot(x, y, [triangles], *, [mask=mask], ...) + + The triangular grid can be specified either by passing a `.Triangulation` + object as the first parameter, or by passing the points *x*, *y* and + optionally the *triangles* and a *mask*. If neither of *triangulation* or + *triangles* are given, the triangulation is calculated on the fly. + + Parameters + ---------- + triangulation : `.Triangulation` + An already created triangular grid. + x, y, triangles, mask + Parameters defining the triangular grid. See `.Triangulation`. + This is mutually exclusive with specifying *triangulation*. + other_parameters + All other args and kwargs are forwarded to `~.Axes.plot`. + + Returns + ------- + lines : `~matplotlib.lines.Line2D` + The drawn triangles edges. + markers : `~matplotlib.lines.Line2D` + The drawn marker nodes. + """ + import matplotlib.axes + + tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs) + x, y, edges = (tri.x, tri.y, tri.edges) + + # Decode plot format string, e.g., 'ro-' + fmt = args[0] if args else "" + linestyle, marker, color = matplotlib.axes._base._process_plot_format(fmt) + + # Insert plot format string into a copy of kwargs (kwargs values prevail). + kw = cbook.normalize_kwargs(kwargs, mlines.Line2D) + for key, val in zip(('linestyle', 'marker', 'color'), + (linestyle, marker, color)): + if val is not None: + kw.setdefault(key, val) + + # Draw lines without markers. + # Note 1: If we drew markers here, most markers would be drawn more than + # once as they belong to several edges. + # Note 2: We insert nan values in the flattened edges arrays rather than + # plotting directly (triang.x[edges].T, triang.y[edges].T) + # as it considerably speeds-up code execution. + linestyle = kw['linestyle'] + kw_lines = { + **kw, + 'marker': 'None', # No marker to draw. + 'zorder': kw.get('zorder', 1), # Path default zorder is used. + } + if linestyle not in [None, 'None', '', ' ']: + tri_lines_x = np.insert(x[edges], 2, np.nan, axis=1) + tri_lines_y = np.insert(y[edges], 2, np.nan, axis=1) + tri_lines = ax.plot(tri_lines_x.ravel(), tri_lines_y.ravel(), + **kw_lines) + else: + tri_lines = ax.plot([], [], **kw_lines) + + # Draw markers separately. + marker = kw['marker'] + kw_markers = { + **kw, + 'linestyle': 'None', # No line to draw. + } + kw_markers.pop('label', None) + if marker not in [None, 'None', '', ' ']: + tri_markers = ax.plot(x, y, **kw_markers) + else: + tri_markers = ax.plot([], [], **kw_markers) + + return tri_lines + tri_markers diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triplot.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triplot.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6224276afdb86dd85d85f1be24c8c62d73fb963c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_triplot.pyi @@ -0,0 +1,15 @@ +from matplotlib.tri._triangulation import Triangulation +from matplotlib.axes import Axes +from matplotlib.lines import Line2D + +from typing import overload +from numpy.typing import ArrayLike + +@overload +def triplot( + ax: Axes, triangulation: Triangulation, *args, **kwargs +) -> tuple[Line2D, Line2D]: ... +@overload +def triplot( + ax: Axes, x: ArrayLike, y: ArrayLike, triangles: ArrayLike = ..., *args, **kwargs +) -> tuple[Line2D, Line2D]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trirefine.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trirefine.py new file mode 100644 index 0000000000000000000000000000000000000000..7f5110ab9e218e24683acd1489d38bded16c7420 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trirefine.py @@ -0,0 +1,307 @@ +""" +Mesh refinement for triangular grids. +""" + +import numpy as np + +from matplotlib import _api +from matplotlib.tri._triangulation import Triangulation +import matplotlib.tri._triinterpolate + + +class TriRefiner: + """ + Abstract base class for classes implementing mesh refinement. + + A TriRefiner encapsulates a Triangulation object and provides tools for + mesh refinement and interpolation. + + Derived classes must implement: + + - ``refine_triangulation(return_tri_index=False, **kwargs)`` , where + the optional keyword arguments *kwargs* are defined in each + TriRefiner concrete implementation, and which returns: + + - a refined triangulation, + - optionally (depending on *return_tri_index*), for each + point of the refined triangulation: the index of + the initial triangulation triangle to which it belongs. + + - ``refine_field(z, triinterpolator=None, **kwargs)``, where: + + - *z* array of field values (to refine) defined at the base + triangulation nodes, + - *triinterpolator* is an optional `~matplotlib.tri.TriInterpolator`, + - the other optional keyword arguments *kwargs* are defined in + each TriRefiner concrete implementation; + + and which returns (as a tuple) a refined triangular mesh and the + interpolated values of the field at the refined triangulation nodes. + """ + + def __init__(self, triangulation): + _api.check_isinstance(Triangulation, triangulation=triangulation) + self._triangulation = triangulation + + +class UniformTriRefiner(TriRefiner): + """ + Uniform mesh refinement by recursive subdivisions. + + Parameters + ---------- + triangulation : `~matplotlib.tri.Triangulation` + The encapsulated triangulation (to be refined) + """ +# See Also +# -------- +# :class:`~matplotlib.tri.CubicTriInterpolator` and +# :class:`~matplotlib.tri.TriAnalyzer`. +# """ + def __init__(self, triangulation): + super().__init__(triangulation) + + def refine_triangulation(self, return_tri_index=False, subdiv=3): + """ + Compute a uniformly refined triangulation *refi_triangulation* of + the encapsulated :attr:`triangulation`. + + This function refines the encapsulated triangulation by splitting each + father triangle into 4 child sub-triangles built on the edges midside + nodes, recursing *subdiv* times. In the end, each triangle is hence + divided into ``4**subdiv`` child triangles. + + Parameters + ---------- + return_tri_index : bool, default: False + Whether an index table indicating the father triangle index of each + point is returned. + subdiv : int, default: 3 + Recursion level for the subdivision. + Each triangle is divided into ``4**subdiv`` child triangles; + hence, the default results in 64 refined subtriangles for each + triangle of the initial triangulation. + + Returns + ------- + refi_triangulation : `~matplotlib.tri.Triangulation` + The refined triangulation. + found_index : int array + Index of the initial triangulation containing triangle, for each + point of *refi_triangulation*. + Returned only if *return_tri_index* is set to True. + """ + refi_triangulation = self._triangulation + ntri = refi_triangulation.triangles.shape[0] + + # Computes the triangulation ancestors numbers in the reference + # triangulation. + ancestors = np.arange(ntri, dtype=np.int32) + for _ in range(subdiv): + refi_triangulation, ancestors = self._refine_triangulation_once( + refi_triangulation, ancestors) + refi_npts = refi_triangulation.x.shape[0] + refi_triangles = refi_triangulation.triangles + + # Now we compute found_index table if needed + if return_tri_index: + # We have to initialize found_index with -1 because some nodes + # may very well belong to no triangle at all, e.g., in case of + # Delaunay Triangulation with DuplicatePointWarning. + found_index = np.full(refi_npts, -1, dtype=np.int32) + tri_mask = self._triangulation.mask + if tri_mask is None: + found_index[refi_triangles] = np.repeat(ancestors, + 3).reshape(-1, 3) + else: + # There is a subtlety here: we want to avoid whenever possible + # that refined points container is a masked triangle (which + # would result in artifacts in plots). + # So we impose the numbering from masked ancestors first, + # then overwrite it with unmasked ancestor numbers. + ancestor_mask = tri_mask[ancestors] + found_index[refi_triangles[ancestor_mask, :] + ] = np.repeat(ancestors[ancestor_mask], + 3).reshape(-1, 3) + found_index[refi_triangles[~ancestor_mask, :] + ] = np.repeat(ancestors[~ancestor_mask], + 3).reshape(-1, 3) + return refi_triangulation, found_index + else: + return refi_triangulation + + def refine_field(self, z, triinterpolator=None, subdiv=3): + """ + Refine a field defined on the encapsulated triangulation. + + Parameters + ---------- + z : (npoints,) array-like + Values of the field to refine, defined at the nodes of the + encapsulated triangulation. (``n_points`` is the number of points + in the initial triangulation) + triinterpolator : `~matplotlib.tri.TriInterpolator`, optional + Interpolator used for field interpolation. If not specified, + a `~matplotlib.tri.CubicTriInterpolator` will be used. + subdiv : int, default: 3 + Recursion level for the subdivision. + Each triangle is divided into ``4**subdiv`` child triangles. + + Returns + ------- + refi_tri : `~matplotlib.tri.Triangulation` + The returned refined triangulation. + refi_z : 1D array of length: *refi_tri* node count. + The returned interpolated field (at *refi_tri* nodes). + """ + if triinterpolator is None: + interp = matplotlib.tri.CubicTriInterpolator( + self._triangulation, z) + else: + _api.check_isinstance(matplotlib.tri.TriInterpolator, + triinterpolator=triinterpolator) + interp = triinterpolator + + refi_tri, found_index = self.refine_triangulation( + subdiv=subdiv, return_tri_index=True) + refi_z = interp._interpolate_multikeys( + refi_tri.x, refi_tri.y, tri_index=found_index)[0] + return refi_tri, refi_z + + @staticmethod + def _refine_triangulation_once(triangulation, ancestors=None): + """ + Refine a `.Triangulation` by splitting each triangle into 4 + child-masked_triangles built on the edges midside nodes. + + Masked triangles, if present, are also split, but their children + returned masked. + + If *ancestors* is not provided, returns only a new triangulation: + child_triangulation. + + If the array-like key table *ancestor* is given, it shall be of shape + (ntri,) where ntri is the number of *triangulation* masked_triangles. + In this case, the function returns + (child_triangulation, child_ancestors) + child_ancestors is defined so that the 4 child masked_triangles share + the same index as their father: child_ancestors.shape = (4 * ntri,). + """ + + x = triangulation.x + y = triangulation.y + + # According to tri.triangulation doc: + # neighbors[i, j] is the triangle that is the neighbor + # to the edge from point index masked_triangles[i, j] to point + # index masked_triangles[i, (j+1)%3]. + neighbors = triangulation.neighbors + triangles = triangulation.triangles + npts = np.shape(x)[0] + ntri = np.shape(triangles)[0] + if ancestors is not None: + ancestors = np.asarray(ancestors) + if np.shape(ancestors) != (ntri,): + raise ValueError( + "Incompatible shapes provide for " + "triangulation.masked_triangles and ancestors: " + f"{np.shape(triangles)} and {np.shape(ancestors)}") + + # Initiating tables refi_x and refi_y of the refined triangulation + # points + # hint: each apex is shared by 2 masked_triangles except the borders. + borders = np.sum(neighbors == -1) + added_pts = (3*ntri + borders) // 2 + refi_npts = npts + added_pts + refi_x = np.zeros(refi_npts) + refi_y = np.zeros(refi_npts) + + # First part of refi_x, refi_y is just the initial points + refi_x[:npts] = x + refi_y[:npts] = y + + # Second part contains the edge midside nodes. + # Each edge belongs to 1 triangle (if border edge) or is shared by 2 + # masked_triangles (interior edge). + # We first build 2 * ntri arrays of edge starting nodes (edge_elems, + # edge_apexes); we then extract only the masters to avoid overlaps. + # The so-called 'master' is the triangle with biggest index + # The 'slave' is the triangle with lower index + # (can be -1 if border edge) + # For slave and master we will identify the apex pointing to the edge + # start + edge_elems = np.tile(np.arange(ntri, dtype=np.int32), 3) + edge_apexes = np.repeat(np.arange(3, dtype=np.int32), ntri) + edge_neighbors = neighbors[edge_elems, edge_apexes] + mask_masters = (edge_elems > edge_neighbors) + + # Identifying the "masters" and adding to refi_x, refi_y vec + masters = edge_elems[mask_masters] + apex_masters = edge_apexes[mask_masters] + x_add = (x[triangles[masters, apex_masters]] + + x[triangles[masters, (apex_masters+1) % 3]]) * 0.5 + y_add = (y[triangles[masters, apex_masters]] + + y[triangles[masters, (apex_masters+1) % 3]]) * 0.5 + refi_x[npts:] = x_add + refi_y[npts:] = y_add + + # Building the new masked_triangles; each old masked_triangles hosts + # 4 new masked_triangles + # there are 6 pts to identify per 'old' triangle, 3 new_pt_corner and + # 3 new_pt_midside + new_pt_corner = triangles + + # What is the index in refi_x, refi_y of point at middle of apex iapex + # of elem ielem ? + # If ielem is the apex master: simple count, given the way refi_x was + # built. + # If ielem is the apex slave: yet we do not know; but we will soon + # using the neighbors table. + new_pt_midside = np.empty([ntri, 3], dtype=np.int32) + cum_sum = npts + for imid in range(3): + mask_st_loc = (imid == apex_masters) + n_masters_loc = np.sum(mask_st_loc) + elem_masters_loc = masters[mask_st_loc] + new_pt_midside[:, imid][elem_masters_loc] = np.arange( + n_masters_loc, dtype=np.int32) + cum_sum + cum_sum += n_masters_loc + + # Now dealing with slave elems. + # for each slave element we identify the master and then the inode + # once slave_masters is identified, slave_masters_apex is such that: + # neighbors[slaves_masters, slave_masters_apex] == slaves + mask_slaves = np.logical_not(mask_masters) + slaves = edge_elems[mask_slaves] + slaves_masters = edge_neighbors[mask_slaves] + diff_table = np.abs(neighbors[slaves_masters, :] - + np.outer(slaves, np.ones(3, dtype=np.int32))) + slave_masters_apex = np.argmin(diff_table, axis=1) + slaves_apex = edge_apexes[mask_slaves] + new_pt_midside[slaves, slaves_apex] = new_pt_midside[ + slaves_masters, slave_masters_apex] + + # Builds the 4 child masked_triangles + child_triangles = np.empty([ntri*4, 3], dtype=np.int32) + child_triangles[0::4, :] = np.vstack([ + new_pt_corner[:, 0], new_pt_midside[:, 0], + new_pt_midside[:, 2]]).T + child_triangles[1::4, :] = np.vstack([ + new_pt_corner[:, 1], new_pt_midside[:, 1], + new_pt_midside[:, 0]]).T + child_triangles[2::4, :] = np.vstack([ + new_pt_corner[:, 2], new_pt_midside[:, 2], + new_pt_midside[:, 1]]).T + child_triangles[3::4, :] = np.vstack([ + new_pt_midside[:, 0], new_pt_midside[:, 1], + new_pt_midside[:, 2]]).T + child_triangulation = Triangulation(refi_x, refi_y, child_triangles) + + # Builds the child mask + if triangulation.mask is not None: + child_triangulation.set_mask(np.repeat(triangulation.mask, 4)) + + if ancestors is None: + return child_triangulation + else: + return child_triangulation, np.repeat(ancestors, 4) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trirefine.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trirefine.pyi new file mode 100644 index 0000000000000000000000000000000000000000..7c60dc76e2f158325ad2b143abe4228599ee1b83 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_trirefine.pyi @@ -0,0 +1,31 @@ +from typing import Literal, overload + +import numpy as np +from numpy.typing import ArrayLike + +from matplotlib.tri._triangulation import Triangulation +from matplotlib.tri._triinterpolate import TriInterpolator + +class TriRefiner: + def __init__(self, triangulation: Triangulation) -> None: ... + +class UniformTriRefiner(TriRefiner): + def __init__(self, triangulation: Triangulation) -> None: ... + @overload + def refine_triangulation( + self, *, return_tri_index: Literal[True], subdiv: int = ... + ) -> tuple[Triangulation, np.ndarray]: ... + @overload + def refine_triangulation( + self, return_tri_index: Literal[False] = ..., subdiv: int = ... + ) -> Triangulation: ... + @overload + def refine_triangulation( + self, return_tri_index: bool = ..., subdiv: int = ... + ) -> tuple[Triangulation, np.ndarray] | Triangulation: ... + def refine_field( + self, + z: ArrayLike, + triinterpolator: TriInterpolator | None = ..., + subdiv: int = ..., + ) -> tuple[Triangulation, np.ndarray]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tritools.py b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tritools.py new file mode 100644 index 0000000000000000000000000000000000000000..9837309f7e810ba1c9d41aa52a9856cb6e1aa062 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tritools.py @@ -0,0 +1,263 @@ +""" +Tools for triangular grids. +""" + +import numpy as np + +from matplotlib import _api +from matplotlib.tri import Triangulation + + +class TriAnalyzer: + """ + Define basic tools for triangular mesh analysis and improvement. + + A TriAnalyzer encapsulates a `.Triangulation` object and provides basic + tools for mesh analysis and mesh improvement. + + Attributes + ---------- + scale_factors + + Parameters + ---------- + triangulation : `~matplotlib.tri.Triangulation` + The encapsulated triangulation to analyze. + """ + + def __init__(self, triangulation): + _api.check_isinstance(Triangulation, triangulation=triangulation) + self._triangulation = triangulation + + @property + def scale_factors(self): + """ + Factors to rescale the triangulation into a unit square. + + Returns + ------- + (float, float) + Scaling factors (kx, ky) so that the triangulation + ``[triangulation.x * kx, triangulation.y * ky]`` + fits exactly inside a unit square. + """ + compressed_triangles = self._triangulation.get_masked_triangles() + node_used = (np.bincount(np.ravel(compressed_triangles), + minlength=self._triangulation.x.size) != 0) + return (1 / np.ptp(self._triangulation.x[node_used]), + 1 / np.ptp(self._triangulation.y[node_used])) + + def circle_ratios(self, rescale=True): + """ + Return a measure of the triangulation triangles flatness. + + The ratio of the incircle radius over the circumcircle radius is a + widely used indicator of a triangle flatness. + It is always ``<= 0.5`` and ``== 0.5`` only for equilateral + triangles. Circle ratios below 0.01 denote very flat triangles. + + To avoid unduly low values due to a difference of scale between the 2 + axis, the triangular mesh can first be rescaled to fit inside a unit + square with `scale_factors` (Only if *rescale* is True, which is + its default value). + + Parameters + ---------- + rescale : bool, default: True + If True, internally rescale (based on `scale_factors`), so that the + (unmasked) triangles fit exactly inside a unit square mesh. + + Returns + ------- + masked array + Ratio of the incircle radius over the circumcircle radius, for + each 'rescaled' triangle of the encapsulated triangulation. + Values corresponding to masked triangles are masked out. + + """ + # Coords rescaling + if rescale: + (kx, ky) = self.scale_factors + else: + (kx, ky) = (1.0, 1.0) + pts = np.vstack([self._triangulation.x*kx, + self._triangulation.y*ky]).T + tri_pts = pts[self._triangulation.triangles] + # Computes the 3 side lengths + a = tri_pts[:, 1, :] - tri_pts[:, 0, :] + b = tri_pts[:, 2, :] - tri_pts[:, 1, :] + c = tri_pts[:, 0, :] - tri_pts[:, 2, :] + a = np.hypot(a[:, 0], a[:, 1]) + b = np.hypot(b[:, 0], b[:, 1]) + c = np.hypot(c[:, 0], c[:, 1]) + # circumcircle and incircle radii + s = (a+b+c)*0.5 + prod = s*(a+b-s)*(a+c-s)*(b+c-s) + # We have to deal with flat triangles with infinite circum_radius + bool_flat = (prod == 0.) + if np.any(bool_flat): + # Pathologic flow + ntri = tri_pts.shape[0] + circum_radius = np.empty(ntri, dtype=np.float64) + circum_radius[bool_flat] = np.inf + abc = a*b*c + circum_radius[~bool_flat] = abc[~bool_flat] / ( + 4.0*np.sqrt(prod[~bool_flat])) + else: + # Normal optimized flow + circum_radius = (a*b*c) / (4.0*np.sqrt(prod)) + in_radius = (a*b*c) / (4.0*circum_radius*s) + circle_ratio = in_radius/circum_radius + mask = self._triangulation.mask + if mask is None: + return circle_ratio + else: + return np.ma.array(circle_ratio, mask=mask) + + def get_flat_tri_mask(self, min_circle_ratio=0.01, rescale=True): + """ + Eliminate excessively flat border triangles from the triangulation. + + Returns a mask *new_mask* which allows to clean the encapsulated + triangulation from its border-located flat triangles + (according to their :meth:`circle_ratios`). + This mask is meant to be subsequently applied to the triangulation + using `.Triangulation.set_mask`. + *new_mask* is an extension of the initial triangulation mask + in the sense that an initially masked triangle will remain masked. + + The *new_mask* array is computed recursively; at each step flat + triangles are removed only if they share a side with the current mesh + border. Thus, no new holes in the triangulated domain will be created. + + Parameters + ---------- + min_circle_ratio : float, default: 0.01 + Border triangles with incircle/circumcircle radii ratio r/R will + be removed if r/R < *min_circle_ratio*. + rescale : bool, default: True + If True, first, internally rescale (based on `scale_factors`) so + that the (unmasked) triangles fit exactly inside a unit square + mesh. This rescaling accounts for the difference of scale which + might exist between the 2 axis. + + Returns + ------- + array of bool + Mask to apply to encapsulated triangulation. + All the initially masked triangles remain masked in the + *new_mask*. + + Notes + ----- + The rationale behind this function is that a Delaunay + triangulation - of an unstructured set of points - sometimes contains + almost flat triangles at its border, leading to artifacts in plots + (especially for high-resolution contouring). + Masked with computed *new_mask*, the encapsulated + triangulation would contain no more unmasked border triangles + with a circle ratio below *min_circle_ratio*, thus improving the + mesh quality for subsequent plots or interpolation. + """ + # Recursively computes the mask_current_borders, true if a triangle is + # at the border of the mesh OR touching the border through a chain of + # invalid aspect ratio masked_triangles. + ntri = self._triangulation.triangles.shape[0] + mask_bad_ratio = self.circle_ratios(rescale) < min_circle_ratio + + current_mask = self._triangulation.mask + if current_mask is None: + current_mask = np.zeros(ntri, dtype=bool) + valid_neighbors = np.copy(self._triangulation.neighbors) + renum_neighbors = np.arange(ntri, dtype=np.int32) + nadd = -1 + while nadd != 0: + # The active wavefront is the triangles from the border (unmasked + # but with a least 1 neighbor equal to -1 + wavefront = (np.min(valid_neighbors, axis=1) == -1) & ~current_mask + # The element from the active wavefront will be masked if their + # circle ratio is bad. + added_mask = wavefront & mask_bad_ratio + current_mask = added_mask | current_mask + nadd = np.sum(added_mask) + + # now we have to update the tables valid_neighbors + valid_neighbors[added_mask, :] = -1 + renum_neighbors[added_mask] = -1 + valid_neighbors = np.where(valid_neighbors == -1, -1, + renum_neighbors[valid_neighbors]) + + return np.ma.filled(current_mask, True) + + def _get_compressed_triangulation(self): + """ + Compress (if masked) the encapsulated triangulation. + + Returns minimal-length triangles array (*compressed_triangles*) and + coordinates arrays (*compressed_x*, *compressed_y*) that can still + describe the unmasked triangles of the encapsulated triangulation. + + Returns + ------- + compressed_triangles : array-like + the returned compressed triangulation triangles + compressed_x : array-like + the returned compressed triangulation 1st coordinate + compressed_y : array-like + the returned compressed triangulation 2nd coordinate + tri_renum : int array + renumbering table to translate the triangle numbers from the + encapsulated triangulation into the new (compressed) renumbering. + -1 for masked triangles (deleted from *compressed_triangles*). + node_renum : int array + renumbering table to translate the point numbers from the + encapsulated triangulation into the new (compressed) renumbering. + -1 for unused points (i.e. those deleted from *compressed_x* and + *compressed_y*). + + """ + # Valid triangles and renumbering + tri_mask = self._triangulation.mask + compressed_triangles = self._triangulation.get_masked_triangles() + ntri = self._triangulation.triangles.shape[0] + if tri_mask is not None: + tri_renum = self._total_to_compress_renum(~tri_mask) + else: + tri_renum = np.arange(ntri, dtype=np.int32) + + # Valid nodes and renumbering + valid_node = (np.bincount(np.ravel(compressed_triangles), + minlength=self._triangulation.x.size) != 0) + compressed_x = self._triangulation.x[valid_node] + compressed_y = self._triangulation.y[valid_node] + node_renum = self._total_to_compress_renum(valid_node) + + # Now renumbering the valid triangles nodes + compressed_triangles = node_renum[compressed_triangles] + + return (compressed_triangles, compressed_x, compressed_y, tri_renum, + node_renum) + + @staticmethod + def _total_to_compress_renum(valid): + """ + Parameters + ---------- + valid : 1D bool array + Validity mask. + + Returns + ------- + int array + Array so that (`valid_array` being a compressed array + based on a `masked_array` with mask ~*valid*): + + - For all i with valid[i] = True: + valid_array[renum[i]] = masked_array[i] + - For all i with valid[i] = False: + renum[i] = -1 (invalid value) + """ + renum = np.full(np.size(valid), -1, dtype=np.int32) + n_valid = np.sum(valid) + renum[valid] = np.arange(n_valid, dtype=np.int32) + return renum diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tritools.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tritools.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9b5e1bec4b81616c99a00021d91f78845e7cd283 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/matplotlib/tri/_tritools.pyi @@ -0,0 +1,12 @@ +from matplotlib.tri import Triangulation + +import numpy as np + +class TriAnalyzer: + def __init__(self, triangulation: Triangulation) -> None: ... + @property + def scale_factors(self) -> tuple[float, float]: ... + def circle_ratios(self, rescale: bool = ...) -> np.ndarray: ... + def get_flat_tri_mask( + self, min_circle_ratio: float = ..., rescale: bool = ... + ) -> np.ndarray: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5dc01a2e7746161acb241198f902c6f291fec739 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/__init__.py @@ -0,0 +1,180 @@ +""" +Contains the core of NumPy: ndarray, ufuncs, dtypes, etc. + +Please note that this module is private. All functions and objects +are available in the main ``numpy`` namespace - use that instead. + +""" + +import os + +from numpy.version import version as __version__ + + +# disables OpenBLAS affinity setting of the main thread that limits +# python threads or processes to one core +env_added = [] +for envkey in ['OPENBLAS_MAIN_FREE', 'GOTOBLAS_MAIN_FREE']: + if envkey not in os.environ: + os.environ[envkey] = '1' + env_added.append(envkey) + +try: + from . import multiarray +except ImportError as exc: + import sys + msg = """ + +IMPORTANT: PLEASE READ THIS FOR ADVICE ON HOW TO SOLVE THIS ISSUE! + +Importing the numpy C-extensions failed. This error can happen for +many reasons, often due to issues with your setup or how NumPy was +installed. + +We have compiled some common reasons and troubleshooting tips at: + + https://numpy.org/devdocs/user/troubleshooting-importerror.html + +Please note and check the following: + + * The Python version is: Python%d.%d from "%s" + * The NumPy version is: "%s" + +and make sure that they are the versions you expect. +Please carefully study the documentation linked above for further help. + +Original error was: %s +""" % (sys.version_info[0], sys.version_info[1], sys.executable, + __version__, exc) + raise ImportError(msg) +finally: + for envkey in env_added: + del os.environ[envkey] +del envkey +del env_added +del os + +from . import umath + +# Check that multiarray,umath are pure python modules wrapping +# _multiarray_umath and not either of the old c-extension modules +if not (hasattr(multiarray, '_multiarray_umath') and + hasattr(umath, '_multiarray_umath')): + import sys + path = sys.modules['numpy'].__path__ + msg = ("Something is wrong with the numpy installation. " + "While importing we detected an older version of " + "numpy in {}. One method of fixing this is to repeatedly uninstall " + "numpy until none is found, then reinstall this version.") + raise ImportError(msg.format(path)) + +from . import numerictypes as nt +from .numerictypes import sctypes, sctypeDict +multiarray.set_typeDict(nt.sctypeDict) +from . import numeric +from .numeric import * +from . import fromnumeric +from .fromnumeric import * +from .records import record, recarray +# Note: module name memmap is overwritten by a class with same name +from .memmap import * +from . import function_base +from .function_base import * +from . import _machar +from . import getlimits +from .getlimits import * +from . import shape_base +from .shape_base import * +from . import einsumfunc +from .einsumfunc import * +del nt + +from .numeric import absolute as abs + +# do this after everything else, to minimize the chance of this misleadingly +# appearing in an import-time traceback +from . import _add_newdocs +from . import _add_newdocs_scalars +# add these for module-freeze analysis (like PyInstaller) +from . import _dtype_ctypes +from . import _internal +from . import _dtype +from . import _methods + +acos = numeric.arccos +acosh = numeric.arccosh +asin = numeric.arcsin +asinh = numeric.arcsinh +atan = numeric.arctan +atanh = numeric.arctanh +atan2 = numeric.arctan2 +concat = numeric.concatenate +bitwise_left_shift = numeric.left_shift +bitwise_invert = numeric.invert +bitwise_right_shift = numeric.right_shift +permute_dims = numeric.transpose +pow = numeric.power + +__all__ = [ + "abs", "acos", "acosh", "asin", "asinh", "atan", "atanh", "atan2", + "bitwise_invert", "bitwise_left_shift", "bitwise_right_shift", "concat", + "pow", "permute_dims", "memmap", "sctypeDict", "record", "recarray" +] +__all__ += numeric.__all__ +__all__ += function_base.__all__ +__all__ += getlimits.__all__ +__all__ += shape_base.__all__ +__all__ += einsumfunc.__all__ + + +def _ufunc_reduce(func): + # Report the `__name__`. pickle will try to find the module. Note that + # pickle supports for this `__name__` to be a `__qualname__`. It may + # make sense to add a `__qualname__` to ufuncs, to allow this more + # explicitly (Numba has ufuncs as attributes). + # See also: https://github.com/dask/distributed/issues/3450 + return func.__name__ + + +def _DType_reconstruct(scalar_type): + # This is a work-around to pickle type(np.dtype(np.float64)), etc. + # and it should eventually be replaced with a better solution, e.g. when + # DTypes become HeapTypes. + return type(dtype(scalar_type)) + + +def _DType_reduce(DType): + # As types/classes, most DTypes can simply be pickled by their name: + if not DType._legacy or DType.__module__ == "numpy.dtypes": + return DType.__name__ + + # However, user defined legacy dtypes (like rational) do not end up in + # `numpy.dtypes` as module and do not have a public class at all. + # For these, we pickle them by reconstructing them from the scalar type: + scalar_type = DType.type + return _DType_reconstruct, (scalar_type,) + + +def __getattr__(name): + # Deprecated 2022-11-22, NumPy 1.25. + if name == "MachAr": + import warnings + warnings.warn( + "The `np._core.MachAr` is considered private API (NumPy 1.24)", + DeprecationWarning, stacklevel=2, + ) + return _machar.MachAr + raise AttributeError(f"Module {__name__!r} has no attribute {name!r}") + + +import copyreg + +copyreg.pickle(ufunc, _ufunc_reduce) +copyreg.pickle(type(dtype), _DType_reduce, _DType_reconstruct) + +# Unclutter namespace (must keep _*_reconstruct for unpickling) +del copyreg, _ufunc_reduce, _DType_reduce + +from numpy._pytesttester import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/__init__.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..b46d5e3387e60d348d311d03f8689f986762059a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/__init__.pyi @@ -0,0 +1,2 @@ +# NOTE: The `np._core` namespace is deliberately kept empty due to it +# being private diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_multiarray_umath.cp310-win_amd64.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_multiarray_umath.cp310-win_amd64.lib new file mode 100644 index 0000000000000000000000000000000000000000..9a0fbae00843e11e24748520e007f995f6610007 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_multiarray_umath.cp310-win_amd64.lib differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_operand_flag_tests.cp310-win_amd64.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_operand_flag_tests.cp310-win_amd64.lib new file mode 100644 index 0000000000000000000000000000000000000000..0515ba5a6fcdaa0b075419d040dc9f786db8dfad Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_operand_flag_tests.cp310-win_amd64.lib differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_rational_tests.cp310-win_amd64.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_rational_tests.cp310-win_amd64.lib new file mode 100644 index 0000000000000000000000000000000000000000..9728ef1341bd1924cc2770708b9617a844b0030e Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_rational_tests.cp310-win_amd64.lib differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_simd.cp310-win_amd64.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_simd.cp310-win_amd64.lib new file mode 100644 index 0000000000000000000000000000000000000000..8238d2c0850b1a6b966b04274ab192048b7d4889 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_simd.cp310-win_amd64.lib differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_simd.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_simd.pyi new file mode 100644 index 0000000000000000000000000000000000000000..630c1637f516034d410d5478ed889eeda79299a5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_simd.pyi @@ -0,0 +1,25 @@ +from types import ModuleType +from typing import TypedDict, type_check_only + +# NOTE: these 5 are only defined on systems with an intel processor +SSE42: ModuleType | None = ... +FMA3: ModuleType | None = ... +AVX2: ModuleType | None = ... +AVX512F: ModuleType | None = ... +AVX512_SKX: ModuleType | None = ... + +baseline: ModuleType | None = ... + +@type_check_only +class SimdTargets(TypedDict): + SSE42: ModuleType | None + AVX2: ModuleType | None + FMA3: ModuleType | None + AVX512F: ModuleType | None + AVX512_SKX: ModuleType | None + baseline: ModuleType | None + +targets: SimdTargets = ... + +def clear_floatstatus() -> None: ... +def get_floatstatus() -> int: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_string_helpers.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_string_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b69afef93d8731b9056f990ab568b540053417 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_string_helpers.py @@ -0,0 +1,100 @@ +""" +String-handling utilities to avoid locale-dependence. + +Used primarily to generate type name aliases. +""" +# "import string" is costly to import! +# Construct the translation tables directly +# "A" = chr(65), "a" = chr(97) +_all_chars = tuple(map(chr, range(256))) +_ascii_upper = _all_chars[65:65+26] +_ascii_lower = _all_chars[97:97+26] +LOWER_TABLE = _all_chars[:65] + _ascii_lower + _all_chars[65+26:] +UPPER_TABLE = _all_chars[:97] + _ascii_upper + _all_chars[97+26:] + + +def english_lower(s): + """ Apply English case rules to convert ASCII strings to all lower case. + + This is an internal utility function to replace calls to str.lower() such + that we can avoid changing behavior with changing locales. In particular, + Turkish has distinct dotted and dotless variants of the Latin letter "I" in + both lowercase and uppercase. Thus, "I".lower() != "i" in a "tr" locale. + + Parameters + ---------- + s : str + + Returns + ------- + lowered : str + + Examples + -------- + >>> from numpy._core.numerictypes import english_lower + >>> english_lower('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_') + 'abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz0123456789_' + >>> english_lower('') + '' + """ + lowered = s.translate(LOWER_TABLE) + return lowered + + +def english_upper(s): + """ Apply English case rules to convert ASCII strings to all upper case. + + This is an internal utility function to replace calls to str.upper() such + that we can avoid changing behavior with changing locales. In particular, + Turkish has distinct dotted and dotless variants of the Latin letter "I" in + both lowercase and uppercase. Thus, "i".upper() != "I" in a "tr" locale. + + Parameters + ---------- + s : str + + Returns + ------- + uppered : str + + Examples + -------- + >>> from numpy._core.numerictypes import english_upper + >>> english_upper('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_') + 'ABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_' + >>> english_upper('') + '' + """ + uppered = s.translate(UPPER_TABLE) + return uppered + + +def english_capitalize(s): + """ Apply English case rules to convert the first character of an ASCII + string to upper case. + + This is an internal utility function to replace calls to str.capitalize() + such that we can avoid changing behavior with changing locales. + + Parameters + ---------- + s : str + + Returns + ------- + capitalized : str + + Examples + -------- + >>> from numpy._core.numerictypes import english_capitalize + >>> english_capitalize('int8') + 'Int8' + >>> english_capitalize('Int8') + 'Int8' + >>> english_capitalize('') + '' + """ + if s: + return english_upper(s[0]) + s[1:] + else: + return s diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_string_helpers.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_string_helpers.pyi new file mode 100644 index 0000000000000000000000000000000000000000..891c8915a13589d2117c436cdc826654e0eb9bcf --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_string_helpers.pyi @@ -0,0 +1,12 @@ +from typing import Final + +_all_chars: Final[tuple[str, ...]] = ... +_ascii_upper: Final[tuple[str, ...]] = ... +_ascii_lower: Final[tuple[str, ...]] = ... + +LOWER_TABLE: Final[tuple[str, ...]] = ... +UPPER_TABLE: Final[tuple[str, ...]] = ... + +def english_lower(s: str) -> str: ... +def english_upper(s: str) -> str: ... +def english_capitalize(s: str) -> str: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_struct_ufunc_tests.cp310-win_amd64.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_struct_ufunc_tests.cp310-win_amd64.lib new file mode 100644 index 0000000000000000000000000000000000000000..f84f3f71c5bcd0322687e8efe6ee41a9d4c0a678 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_struct_ufunc_tests.cp310-win_amd64.lib differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_type_aliases.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_type_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..f193118765086d71dd13cba69b440f2c70ef4e85 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_type_aliases.py @@ -0,0 +1,119 @@ +""" +Due to compatibility, numpy has a very large number of different naming +conventions for the scalar types (those subclassing from `numpy.generic`). +This file produces a convoluted set of dictionaries mapping names to types, +and sometimes other mappings too. + +.. data:: allTypes + A dictionary of names to types that will be exposed as attributes through + ``np._core.numerictypes.*`` + +.. data:: sctypeDict + Similar to `allTypes`, but maps a broader set of aliases to their types. + +.. data:: sctypes + A dictionary keyed by a "type group" string, providing a list of types + under that group. + +""" + +import numpy._core.multiarray as ma +from numpy._core.multiarray import typeinfo, dtype + +###################################### +# Building `sctypeDict` and `allTypes` +###################################### + +sctypeDict = {} +allTypes = {} +c_names_dict = {} + +_abstract_type_names = { + "generic", "integer", "inexact", "floating", "number", + "flexible", "character", "complexfloating", "unsignedinteger", + "signedinteger" +} + +for _abstract_type_name in _abstract_type_names: + allTypes[_abstract_type_name] = getattr(ma, _abstract_type_name) + +for k, v in typeinfo.items(): + if k.startswith("NPY_") and v not in c_names_dict: + c_names_dict[k[4:]] = v + else: + concrete_type = v.type + allTypes[k] = concrete_type + sctypeDict[k] = concrete_type + +_aliases = { + "double": "float64", + "cdouble": "complex128", + "single": "float32", + "csingle": "complex64", + "half": "float16", + "bool_": "bool", + # Default integer: + "int_": "intp", + "uint": "uintp", +} + +for k, v in _aliases.items(): + sctypeDict[k] = allTypes[v] + allTypes[k] = allTypes[v] + +# extra aliases are added only to `sctypeDict` +# to support dtype name access, such as`np.dtype("float")` +_extra_aliases = { + "float": "float64", + "complex": "complex128", + "object": "object_", + "bytes": "bytes_", + "a": "bytes_", + "int": "int_", + "str": "str_", + "unicode": "str_", +} + +for k, v in _extra_aliases.items(): + sctypeDict[k] = allTypes[v] + +# include extended precision sized aliases +for is_complex, full_name in [(False, "longdouble"), (True, "clongdouble")]: + longdouble_type: type = allTypes[full_name] + + bits: int = dtype(longdouble_type).itemsize * 8 + base_name: str = "complex" if is_complex else "float" + extended_prec_name: str = f"{base_name}{bits}" + if extended_prec_name not in allTypes: + sctypeDict[extended_prec_name] = longdouble_type + allTypes[extended_prec_name] = longdouble_type + + +#################### +# Building `sctypes` +#################### + +sctypes = {"int": set(), "uint": set(), "float": set(), + "complex": set(), "others": set()} + +for type_info in typeinfo.values(): + if type_info.kind in ["M", "m"]: # exclude timedelta and datetime + continue + + concrete_type = type_info.type + + # find proper group for each concrete type + for type_group, abstract_type in [ + ("int", ma.signedinteger), ("uint", ma.unsignedinteger), + ("float", ma.floating), ("complex", ma.complexfloating), + ("others", ma.generic) + ]: + if issubclass(concrete_type, abstract_type): + sctypes[type_group].add(concrete_type) + break + +# sort sctype groups by bitsize +for sctype_key in sctypes.keys(): + sctype_list = list(sctypes[sctype_key]) + sctype_list.sort(key=lambda x: dtype(x).itemsize) + sctypes[sctype_key] = sctype_list diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_type_aliases.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_type_aliases.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6d28b11b13443d05c7b04e0784e12974ae1d4159 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_type_aliases.pyi @@ -0,0 +1,96 @@ +from collections.abc import Collection +from typing import Any, Final, Literal as L, TypeAlias, TypedDict, type_check_only + +import numpy as np + +__all__ = ( + "_abstract_type_names", + "_aliases", + "_extra_aliases", + "allTypes", + "c_names_dict", + "sctypeDict", + "sctypes", +) + +sctypeDict: Final[dict[str, type[np.generic]]] +allTypes: Final[dict[str, type[np.generic]]] + +@type_check_only +class _CNamesDict(TypedDict): + BOOL: np.dtype[np.bool] + HALF: np.dtype[np.half] + FLOAT: np.dtype[np.single] + DOUBLE: np.dtype[np.double] + LONGDOUBLE: np.dtype[np.longdouble] + CFLOAT: np.dtype[np.csingle] + CDOUBLE: np.dtype[np.cdouble] + CLONGDOUBLE: np.dtype[np.clongdouble] + STRING: np.dtype[np.bytes_] + UNICODE: np.dtype[np.str_] + VOID: np.dtype[np.void] + OBJECT: np.dtype[np.object_] + DATETIME: np.dtype[np.datetime64] + TIMEDELTA: np.dtype[np.timedelta64] + BYTE: np.dtype[np.byte] + UBYTE: np.dtype[np.ubyte] + SHORT: np.dtype[np.short] + USHORT: np.dtype[np.ushort] + INT: np.dtype[np.intc] + UINT: np.dtype[np.uintc] + LONG: np.dtype[np.long] + ULONG: np.dtype[np.ulong] + LONGLONG: np.dtype[np.longlong] + ULONGLONG: np.dtype[np.ulonglong] + +c_names_dict: Final[_CNamesDict] + +_AbstractTypeName: TypeAlias = L[ + "generic", + "flexible", + "character", + "number", + "integer", + "inexact", + "unsignedinteger", + "signedinteger", + "floating", + "complexfloating", +] +_abstract_type_names: Final[set[_AbstractTypeName]] + +@type_check_only +class _AliasesType(TypedDict): + double: L["float64"] + cdouble: L["complex128"] + single: L["float32"] + csingle: L["complex64"] + half: L["float16"] + bool_: L["bool"] + int_: L["intp"] + uint: L["intp"] + +_aliases: Final[_AliasesType] + +@type_check_only +class _ExtraAliasesType(TypedDict): + float: L["float64"] + complex: L["complex128"] + object: L["object_"] + bytes: L["bytes_"] + a: L["bytes_"] + int: L["int_"] + str: L["str_"] + unicode: L["str_"] + +_extra_aliases: Final[_ExtraAliasesType] + +@type_check_only +class _SCTypes(TypedDict): + int: Collection[type[np.signedinteger[Any]]] + uint: Collection[type[np.unsignedinteger[Any]]] + float: Collection[type[np.floating[Any]]] + complex: Collection[type[np.complexfloating[Any, Any]]] + others: Collection[type[np.flexible | np.bool | np.object_]] + +sctypes: Final[_SCTypes] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_ufunc_config.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_ufunc_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6383a928a43d051ea4a68b4d1542a9c155df165d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_ufunc_config.py @@ -0,0 +1,483 @@ +""" +Functions for changing global ufunc configuration + +This provides helpers which wrap `_get_extobj_dict` and `_make_extobj`, and +`_extobj_contextvar` from umath. +""" +import contextlib +import contextvars +import functools + +from .._utils import set_module +from .umath import _make_extobj, _get_extobj_dict, _extobj_contextvar + +__all__ = [ + "seterr", "geterr", "setbufsize", "getbufsize", "seterrcall", "geterrcall", + "errstate" +] + + +@set_module('numpy') +def seterr(all=None, divide=None, over=None, under=None, invalid=None): + """ + Set how floating-point errors are handled. + + Note that operations on integer scalar types (such as `int16`) are + handled like floating point, and are affected by these settings. + + Parameters + ---------- + all : {'ignore', 'warn', 'raise', 'call', 'print', 'log'}, optional + Set treatment for all types of floating-point errors at once: + + - ignore: Take no action when the exception occurs. + - warn: Print a :exc:`RuntimeWarning` (via the Python `warnings` + module). + - raise: Raise a :exc:`FloatingPointError`. + - call: Call a function specified using the `seterrcall` function. + - print: Print a warning directly to ``stdout``. + - log: Record error in a Log object specified by `seterrcall`. + + The default is not to change the current behavior. + divide : {'ignore', 'warn', 'raise', 'call', 'print', 'log'}, optional + Treatment for division by zero. + over : {'ignore', 'warn', 'raise', 'call', 'print', 'log'}, optional + Treatment for floating-point overflow. + under : {'ignore', 'warn', 'raise', 'call', 'print', 'log'}, optional + Treatment for floating-point underflow. + invalid : {'ignore', 'warn', 'raise', 'call', 'print', 'log'}, optional + Treatment for invalid floating-point operation. + + Returns + ------- + old_settings : dict + Dictionary containing the old settings. + + See also + -------- + seterrcall : Set a callback function for the 'call' mode. + geterr, geterrcall, errstate + + Notes + ----- + The floating-point exceptions are defined in the IEEE 754 standard [1]_: + + - Division by zero: infinite result obtained from finite numbers. + - Overflow: result too large to be expressed. + - Underflow: result so close to zero that some precision + was lost. + - Invalid operation: result is not an expressible number, typically + indicates that a NaN was produced. + + .. [1] https://en.wikipedia.org/wiki/IEEE_754 + + Examples + -------- + >>> import numpy as np + >>> orig_settings = np.seterr(all='ignore') # seterr to known value + >>> np.int16(32000) * np.int16(3) + np.int16(30464) + >>> np.seterr(over='raise') + {'divide': 'ignore', 'over': 'ignore', 'under': 'ignore', 'invalid': 'ignore'} + >>> old_settings = np.seterr(all='warn', over='raise') + >>> np.int16(32000) * np.int16(3) + Traceback (most recent call last): + File "", line 1, in + FloatingPointError: overflow encountered in scalar multiply + + >>> old_settings = np.seterr(all='print') + >>> np.geterr() + {'divide': 'print', 'over': 'print', 'under': 'print', 'invalid': 'print'} + >>> np.int16(32000) * np.int16(3) + np.int16(30464) + >>> np.seterr(**orig_settings) # restore original + {'divide': 'print', 'over': 'print', 'under': 'print', 'invalid': 'print'} + + """ + + old = _get_extobj_dict() + # The errstate doesn't include call and bufsize, so pop them: + old.pop("call", None) + old.pop("bufsize", None) + + extobj = _make_extobj( + all=all, divide=divide, over=over, under=under, invalid=invalid) + _extobj_contextvar.set(extobj) + return old + + +@set_module('numpy') +def geterr(): + """ + Get the current way of handling floating-point errors. + + Returns + ------- + res : dict + A dictionary with keys "divide", "over", "under", and "invalid", + whose values are from the strings "ignore", "print", "log", "warn", + "raise", and "call". The keys represent possible floating-point + exceptions, and the values define how these exceptions are handled. + + See Also + -------- + geterrcall, seterr, seterrcall + + Notes + ----- + For complete documentation of the types of floating-point exceptions and + treatment options, see `seterr`. + + Examples + -------- + >>> import numpy as np + >>> np.geterr() + {'divide': 'warn', 'over': 'warn', 'under': 'ignore', 'invalid': 'warn'} + >>> np.arange(3.) / np.arange(3.) # doctest: +SKIP + array([nan, 1., 1.]) + RuntimeWarning: invalid value encountered in divide + + >>> oldsettings = np.seterr(all='warn', invalid='raise') + >>> np.geterr() + {'divide': 'warn', 'over': 'warn', 'under': 'warn', 'invalid': 'raise'} + >>> np.arange(3.) / np.arange(3.) + Traceback (most recent call last): + ... + FloatingPointError: invalid value encountered in divide + >>> oldsettings = np.seterr(**oldsettings) # restore original + + """ + res = _get_extobj_dict() + # The "geterr" doesn't include call and bufsize,: + res.pop("call", None) + res.pop("bufsize", None) + return res + + +@set_module('numpy') +def setbufsize(size): + """ + Set the size of the buffer used in ufuncs. + + .. versionchanged:: 2.0 + The scope of setting the buffer is tied to the `numpy.errstate` + context. Exiting a ``with errstate():`` will also restore the bufsize. + + Parameters + ---------- + size : int + Size of buffer. + + Returns + ------- + bufsize : int + Previous size of ufunc buffer in bytes. + + Examples + -------- + When exiting a `numpy.errstate` context manager the bufsize is restored: + + >>> import numpy as np + >>> with np.errstate(): + ... np.setbufsize(4096) + ... print(np.getbufsize()) + ... + 8192 + 4096 + >>> np.getbufsize() + 8192 + + """ + old = _get_extobj_dict()["bufsize"] + extobj = _make_extobj(bufsize=size) + _extobj_contextvar.set(extobj) + return old + + +@set_module('numpy') +def getbufsize(): + """ + Return the size of the buffer used in ufuncs. + + Returns + ------- + getbufsize : int + Size of ufunc buffer in bytes. + + Examples + -------- + >>> import numpy as np + >>> np.getbufsize() + 8192 + + """ + return _get_extobj_dict()["bufsize"] + + +@set_module('numpy') +def seterrcall(func): + """ + Set the floating-point error callback function or log object. + + There are two ways to capture floating-point error messages. The first + is to set the error-handler to 'call', using `seterr`. Then, set + the function to call using this function. + + The second is to set the error-handler to 'log', using `seterr`. + Floating-point errors then trigger a call to the 'write' method of + the provided object. + + Parameters + ---------- + func : callable f(err, flag) or object with write method + Function to call upon floating-point errors ('call'-mode) or + object whose 'write' method is used to log such message ('log'-mode). + + The call function takes two arguments. The first is a string describing + the type of error (such as "divide by zero", "overflow", "underflow", + or "invalid value"), and the second is the status flag. The flag is a + byte, whose four least-significant bits indicate the type of error, one + of "divide", "over", "under", "invalid":: + + [0 0 0 0 divide over under invalid] + + In other words, ``flags = divide + 2*over + 4*under + 8*invalid``. + + If an object is provided, its write method should take one argument, + a string. + + Returns + ------- + h : callable, log instance or None + The old error handler. + + See Also + -------- + seterr, geterr, geterrcall + + Examples + -------- + Callback upon error: + + >>> def err_handler(type, flag): + ... print("Floating point error (%s), with flag %s" % (type, flag)) + ... + + >>> import numpy as np + + >>> orig_handler = np.seterrcall(err_handler) + >>> orig_err = np.seterr(all='call') + + >>> np.array([1, 2, 3]) / 0.0 + Floating point error (divide by zero), with flag 1 + array([inf, inf, inf]) + + >>> np.seterrcall(orig_handler) + + >>> np.seterr(**orig_err) + {'divide': 'call', 'over': 'call', 'under': 'call', 'invalid': 'call'} + + Log error message: + + >>> class Log: + ... def write(self, msg): + ... print("LOG: %s" % msg) + ... + + >>> log = Log() + >>> saved_handler = np.seterrcall(log) + >>> save_err = np.seterr(all='log') + + >>> np.array([1, 2, 3]) / 0.0 + LOG: Warning: divide by zero encountered in divide + array([inf, inf, inf]) + + >>> np.seterrcall(orig_handler) + + >>> np.seterr(**orig_err) + {'divide': 'log', 'over': 'log', 'under': 'log', 'invalid': 'log'} + + """ + old = _get_extobj_dict()["call"] + extobj = _make_extobj(call=func) + _extobj_contextvar.set(extobj) + return old + + +@set_module('numpy') +def geterrcall(): + """ + Return the current callback function used on floating-point errors. + + When the error handling for a floating-point error (one of "divide", + "over", "under", or "invalid") is set to 'call' or 'log', the function + that is called or the log instance that is written to is returned by + `geterrcall`. This function or log instance has been set with + `seterrcall`. + + Returns + ------- + errobj : callable, log instance or None + The current error handler. If no handler was set through `seterrcall`, + ``None`` is returned. + + See Also + -------- + seterrcall, seterr, geterr + + Notes + ----- + For complete documentation of the types of floating-point exceptions and + treatment options, see `seterr`. + + Examples + -------- + >>> import numpy as np + >>> np.geterrcall() # we did not yet set a handler, returns None + + >>> orig_settings = np.seterr(all='call') + >>> def err_handler(type, flag): + ... print("Floating point error (%s), with flag %s" % (type, flag)) + >>> old_handler = np.seterrcall(err_handler) + >>> np.array([1, 2, 3]) / 0.0 + Floating point error (divide by zero), with flag 1 + array([inf, inf, inf]) + + >>> cur_handler = np.geterrcall() + >>> cur_handler is err_handler + True + >>> old_settings = np.seterr(**orig_settings) # restore original + >>> old_handler = np.seterrcall(None) # restore original + + """ + return _get_extobj_dict()["call"] + + +class _unspecified: + pass + + +_Unspecified = _unspecified() + + +@set_module('numpy') +class errstate: + """ + errstate(**kwargs) + + Context manager for floating-point error handling. + + Using an instance of `errstate` as a context manager allows statements in + that context to execute with a known error handling behavior. Upon entering + the context the error handling is set with `seterr` and `seterrcall`, and + upon exiting it is reset to what it was before. + + .. versionchanged:: 1.17.0 + `errstate` is also usable as a function decorator, saving + a level of indentation if an entire function is wrapped. + + .. versionchanged:: 2.0 + `errstate` is now fully thread and asyncio safe, but may not be + entered more than once. + It is not safe to decorate async functions using ``errstate``. + + Parameters + ---------- + kwargs : {divide, over, under, invalid} + Keyword arguments. The valid keywords are the possible floating-point + exceptions. Each keyword should have a string value that defines the + treatment for the particular error. Possible values are + {'ignore', 'warn', 'raise', 'call', 'print', 'log'}. + + See Also + -------- + seterr, geterr, seterrcall, geterrcall + + Notes + ----- + For complete documentation of the types of floating-point exceptions and + treatment options, see `seterr`. + + Examples + -------- + >>> import numpy as np + >>> olderr = np.seterr(all='ignore') # Set error handling to known state. + + >>> np.arange(3) / 0. + array([nan, inf, inf]) + >>> with np.errstate(divide='ignore'): + ... np.arange(3) / 0. + array([nan, inf, inf]) + + >>> np.sqrt(-1) + np.float64(nan) + >>> with np.errstate(invalid='raise'): + ... np.sqrt(-1) + Traceback (most recent call last): + File "", line 2, in + FloatingPointError: invalid value encountered in sqrt + + Outside the context the error handling behavior has not changed: + + >>> np.geterr() + {'divide': 'ignore', 'over': 'ignore', 'under': 'ignore', 'invalid': 'ignore'} + >>> olderr = np.seterr(**olderr) # restore original state + + """ + __slots__ = ( + "_call", "_all", "_divide", "_over", "_under", "_invalid", "_token") + + def __init__(self, *, call=_Unspecified, + all=None, divide=None, over=None, under=None, invalid=None): + self._token = None + self._call = call + self._all = all + self._divide = divide + self._over = over + self._under = under + self._invalid = invalid + + def __enter__(self): + # Note that __call__ duplicates much of this logic + if self._token is not None: + raise TypeError("Cannot enter `np.errstate` twice.") + if self._call is _Unspecified: + extobj = _make_extobj( + all=self._all, divide=self._divide, over=self._over, + under=self._under, invalid=self._invalid) + else: + extobj = _make_extobj( + call=self._call, + all=self._all, divide=self._divide, over=self._over, + under=self._under, invalid=self._invalid) + + self._token = _extobj_contextvar.set(extobj) + + def __exit__(self, *exc_info): + _extobj_contextvar.reset(self._token) + + def __call__(self, func): + # We need to customize `__call__` compared to `ContextDecorator` + # because we must store the token per-thread so cannot store it on + # the instance (we could create a new instance for this). + # This duplicates the code from `__enter__`. + @functools.wraps(func) + def inner(*args, **kwargs): + if self._call is _Unspecified: + extobj = _make_extobj( + all=self._all, divide=self._divide, over=self._over, + under=self._under, invalid=self._invalid) + else: + extobj = _make_extobj( + call=self._call, + all=self._all, divide=self._divide, over=self._over, + under=self._under, invalid=self._invalid) + + _token = _extobj_contextvar.set(extobj) + try: + # Call the original, decorated, function: + return func(*args, **kwargs) + finally: + _extobj_contextvar.reset(_token) + + return inner diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_ufunc_config.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_ufunc_config.pyi new file mode 100644 index 0000000000000000000000000000000000000000..283b7bfff9f641ac2d60365d001a6e085d47588e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_ufunc_config.pyi @@ -0,0 +1,39 @@ +from _typeshed import SupportsWrite +from collections.abc import Callable +from typing import Any, Literal, TypeAlias, TypedDict, type_check_only + +from numpy import errstate as errstate + +_ErrKind: TypeAlias = Literal["ignore", "warn", "raise", "call", "print", "log"] +_ErrFunc: TypeAlias = Callable[[str, int], Any] +_ErrCall: TypeAlias = _ErrFunc | SupportsWrite[str] + +@type_check_only +class _ErrDict(TypedDict): + divide: _ErrKind + over: _ErrKind + under: _ErrKind + invalid: _ErrKind + +@type_check_only +class _ErrDictOptional(TypedDict, total=False): + all: None | _ErrKind + divide: None | _ErrKind + over: None | _ErrKind + under: None | _ErrKind + invalid: None | _ErrKind + +def seterr( + all: None | _ErrKind = ..., + divide: None | _ErrKind = ..., + over: None | _ErrKind = ..., + under: None | _ErrKind = ..., + invalid: None | _ErrKind = ..., +) -> _ErrDict: ... +def geterr() -> _ErrDict: ... +def setbufsize(size: int) -> int: ... +def getbufsize() -> int: ... +def seterrcall(func: _ErrCall | None) -> _ErrCall | None: ... +def geterrcall() -> _ErrCall | None: ... + +# See `numpy/__init__.pyi` for the `errstate` class and `no_nep5_warnings` diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_umath_tests.cp310-win_amd64.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_umath_tests.cp310-win_amd64.lib new file mode 100644 index 0000000000000000000000000000000000000000..5d37a023ffc2900017985a8335fe1ba61c4c59a3 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/_umath_tests.cp310-win_amd64.lib differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/arrayprint.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/arrayprint.py new file mode 100644 index 0000000000000000000000000000000000000000..0b91dc10e84609acef0df1a7f1a423e72dae6f71 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/arrayprint.py @@ -0,0 +1,1756 @@ +"""Array printing function + +$Id: arrayprint.py,v 1.9 2005/09/13 13:58:44 teoliphant Exp $ + +""" +__all__ = ["array2string", "array_str", "array_repr", + "set_printoptions", "get_printoptions", "printoptions", + "format_float_positional", "format_float_scientific"] +__docformat__ = 'restructuredtext' + +# +# Written by Konrad Hinsen +# last revision: 1996-3-13 +# modified by Jim Hugunin 1997-3-3 for repr's and str's (and other details) +# and by Perry Greenfield 2000-4-1 for numarray +# and by Travis Oliphant 2005-8-22 for numpy + + +# Note: Both scalartypes.c.src and arrayprint.py implement strs for numpy +# scalars but for different purposes. scalartypes.c.src has str/reprs for when +# the scalar is printed on its own, while arrayprint.py has strs for when +# scalars are printed inside an ndarray. Only the latter strs are currently +# user-customizable. + +import functools +import numbers +import sys +try: + from _thread import get_ident +except ImportError: + from _dummy_thread import get_ident + +import numpy as np +from . import numerictypes as _nt +from .umath import absolute, isinf, isfinite, isnat +from . import multiarray +from .multiarray import (array, dragon4_positional, dragon4_scientific, + datetime_as_string, datetime_data, ndarray) +from .fromnumeric import any +from .numeric import concatenate, asarray, errstate +from .numerictypes import (longlong, intc, int_, float64, complex128, + flexible) +from .overrides import array_function_dispatch, set_module +from .printoptions import format_options +import operator +import warnings +import contextlib + + +def _make_options_dict(precision=None, threshold=None, edgeitems=None, + linewidth=None, suppress=None, nanstr=None, infstr=None, + sign=None, formatter=None, floatmode=None, legacy=None, + override_repr=None): + """ + Make a dictionary out of the non-None arguments, plus conversion of + *legacy* and sanity checks. + """ + + options = {k: v for k, v in list(locals().items()) if v is not None} + + if suppress is not None: + options['suppress'] = bool(suppress) + + modes = ['fixed', 'unique', 'maxprec', 'maxprec_equal'] + if floatmode not in modes + [None]: + raise ValueError("floatmode option must be one of " + + ", ".join('"{}"'.format(m) for m in modes)) + + if sign not in [None, '-', '+', ' ']: + raise ValueError("sign option must be one of ' ', '+', or '-'") + + if legacy is False: + options['legacy'] = sys.maxsize + elif legacy == False: # noqa: E712 + warnings.warn( + f"Passing `legacy={legacy!r}` is deprecated.", + FutureWarning, stacklevel=3 + ) + options['legacy'] = sys.maxsize + elif legacy == '1.13': + options['legacy'] = 113 + elif legacy == '1.21': + options['legacy'] = 121 + elif legacy == '1.25': + options['legacy'] = 125 + elif legacy == '2.1': + options['legacy'] = 201 + elif legacy is None: + pass # OK, do nothing. + else: + warnings.warn( + "legacy printing option can currently only be '1.13', '1.21', " + "'1.25', '2.1, or `False`", stacklevel=3) + + if threshold is not None: + # forbid the bad threshold arg suggested by stack overflow, gh-12351 + if not isinstance(threshold, numbers.Number): + raise TypeError("threshold must be numeric") + if np.isnan(threshold): + raise ValueError("threshold must be non-NAN, try " + "sys.maxsize for untruncated representation") + + if precision is not None: + # forbid the bad precision arg as suggested by issue #18254 + try: + options['precision'] = operator.index(precision) + except TypeError as e: + raise TypeError('precision must be an integer') from e + + return options + + +@set_module('numpy') +def set_printoptions(precision=None, threshold=None, edgeitems=None, + linewidth=None, suppress=None, nanstr=None, + infstr=None, formatter=None, sign=None, floatmode=None, + *, legacy=None, override_repr=None): + """ + Set printing options. + + These options determine the way floating point numbers, arrays and + other NumPy objects are displayed. + + Parameters + ---------- + precision : int or None, optional + Number of digits of precision for floating point output (default 8). + May be None if `floatmode` is not `fixed`, to print as many digits as + necessary to uniquely specify the value. + threshold : int, optional + Total number of array elements which trigger summarization + rather than full repr (default 1000). + To always use the full repr without summarization, pass `sys.maxsize`. + edgeitems : int, optional + Number of array items in summary at beginning and end of + each dimension (default 3). + linewidth : int, optional + The number of characters per line for the purpose of inserting + line breaks (default 75). + suppress : bool, optional + If True, always print floating point numbers using fixed point + notation, in which case numbers equal to zero in the current precision + will print as zero. If False, then scientific notation is used when + absolute value of the smallest number is < 1e-4 or the ratio of the + maximum absolute value to the minimum is > 1e3. The default is False. + nanstr : str, optional + String representation of floating point not-a-number (default nan). + infstr : str, optional + String representation of floating point infinity (default inf). + sign : string, either '-', '+', or ' ', optional + Controls printing of the sign of floating-point types. If '+', always + print the sign of positive values. If ' ', always prints a space + (whitespace character) in the sign position of positive values. If + '-', omit the sign character of positive values. (default '-') + + .. versionchanged:: 2.0 + The sign parameter can now be an integer type, previously + types were floating-point types. + + formatter : dict of callables, optional + If not None, the keys should indicate the type(s) that the respective + formatting function applies to. Callables should return a string. + Types that are not specified (by their corresponding keys) are handled + by the default formatters. Individual types for which a formatter + can be set are: + + - 'bool' + - 'int' + - 'timedelta' : a `numpy.timedelta64` + - 'datetime' : a `numpy.datetime64` + - 'float' + - 'longfloat' : 128-bit floats + - 'complexfloat' + - 'longcomplexfloat' : composed of two 128-bit floats + - 'numpystr' : types `numpy.bytes_` and `numpy.str_` + - 'object' : `np.object_` arrays + + Other keys that can be used to set a group of types at once are: + + - 'all' : sets all types + - 'int_kind' : sets 'int' + - 'float_kind' : sets 'float' and 'longfloat' + - 'complex_kind' : sets 'complexfloat' and 'longcomplexfloat' + - 'str_kind' : sets 'numpystr' + floatmode : str, optional + Controls the interpretation of the `precision` option for + floating-point types. Can take the following values + (default maxprec_equal): + + * 'fixed': Always print exactly `precision` fractional digits, + even if this would print more or fewer digits than + necessary to specify the value uniquely. + * 'unique': Print the minimum number of fractional digits necessary + to represent each value uniquely. Different elements may + have a different number of digits. The value of the + `precision` option is ignored. + * 'maxprec': Print at most `precision` fractional digits, but if + an element can be uniquely represented with fewer digits + only print it with that many. + * 'maxprec_equal': Print at most `precision` fractional digits, + but if every element in the array can be uniquely + represented with an equal number of fewer digits, use that + many digits for all elements. + legacy : string or `False`, optional + If set to the string ``'1.13'`` enables 1.13 legacy printing mode. This + approximates numpy 1.13 print output by including a space in the sign + position of floats and different behavior for 0d arrays. This also + enables 1.21 legacy printing mode (described below). + + If set to the string ``'1.21'`` enables 1.21 legacy printing mode. This + approximates numpy 1.21 print output of complex structured dtypes + by not inserting spaces after commas that separate fields and after + colons. + + If set to ``'1.25'`` approximates printing of 1.25 which mainly means + that numeric scalars are printed without their type information, e.g. + as ``3.0`` rather than ``np.float64(3.0)``. + + If set to ``'2.1'``, shape information is not given when arrays are + summarized (i.e., multiple elements replaced with ``...``). + + If set to `False`, disables legacy mode. + + Unrecognized strings will be ignored with a warning for forward + compatibility. + + .. versionchanged:: 1.22.0 + .. versionchanged:: 2.2 + + override_repr: callable, optional + If set a passed function will be used for generating arrays' repr. + Other options will be ignored. + + See Also + -------- + get_printoptions, printoptions, array2string + + Notes + ----- + `formatter` is always reset with a call to `set_printoptions`. + + Use `printoptions` as a context manager to set the values temporarily. + + Examples + -------- + Floating point precision can be set: + + >>> import numpy as np + >>> np.set_printoptions(precision=4) + >>> np.array([1.123456789]) + [1.1235] + + Long arrays can be summarised: + + >>> np.set_printoptions(threshold=5) + >>> np.arange(10) + array([0, 1, 2, ..., 7, 8, 9], shape=(10,)) + + Small results can be suppressed: + + >>> eps = np.finfo(float).eps + >>> x = np.arange(4.) + >>> x**2 - (x + eps)**2 + array([-4.9304e-32, -4.4409e-16, 0.0000e+00, 0.0000e+00]) + >>> np.set_printoptions(suppress=True) + >>> x**2 - (x + eps)**2 + array([-0., -0., 0., 0.]) + + A custom formatter can be used to display array elements as desired: + + >>> np.set_printoptions(formatter={'all':lambda x: 'int: '+str(-x)}) + >>> x = np.arange(3) + >>> x + array([int: 0, int: -1, int: -2]) + >>> np.set_printoptions() # formatter gets reset + >>> x + array([0, 1, 2]) + + To put back the default options, you can use: + + >>> np.set_printoptions(edgeitems=3, infstr='inf', + ... linewidth=75, nanstr='nan', precision=8, + ... suppress=False, threshold=1000, formatter=None) + + Also to temporarily override options, use `printoptions` + as a context manager: + + >>> with np.printoptions(precision=2, suppress=True, threshold=5): + ... np.linspace(0, 10, 10) + array([ 0. , 1.11, 2.22, ..., 7.78, 8.89, 10. ], shape=(10,)) + + """ + _set_printoptions(precision, threshold, edgeitems, linewidth, suppress, + nanstr, infstr, formatter, sign, floatmode, + legacy=legacy, override_repr=override_repr) + + +def _set_printoptions(precision=None, threshold=None, edgeitems=None, + linewidth=None, suppress=None, nanstr=None, + infstr=None, formatter=None, sign=None, floatmode=None, + *, legacy=None, override_repr=None): + new_opt = _make_options_dict(precision, threshold, edgeitems, linewidth, + suppress, nanstr, infstr, sign, formatter, + floatmode, legacy) + # formatter and override_repr are always reset + new_opt['formatter'] = formatter + new_opt['override_repr'] = override_repr + + updated_opt = format_options.get() | new_opt + updated_opt.update(new_opt) + + if updated_opt['legacy'] == 113: + updated_opt['sign'] = '-' + + return format_options.set(updated_opt) + + +@set_module('numpy') +def get_printoptions(): + """ + Return the current print options. + + Returns + ------- + print_opts : dict + Dictionary of current print options with keys + + - precision : int + - threshold : int + - edgeitems : int + - linewidth : int + - suppress : bool + - nanstr : str + - infstr : str + - sign : str + - formatter : dict of callables + - floatmode : str + - legacy : str or False + + For a full description of these options, see `set_printoptions`. + + See Also + -------- + set_printoptions, printoptions + + Examples + -------- + >>> import numpy as np + + >>> np.get_printoptions() + {'edgeitems': 3, 'threshold': 1000, ..., 'override_repr': None} + + >>> np.get_printoptions()['linewidth'] + 75 + >>> np.set_printoptions(linewidth=100) + >>> np.get_printoptions()['linewidth'] + 100 + + """ + opts = format_options.get().copy() + opts['legacy'] = { + 113: '1.13', 121: '1.21', 125: '1.25', sys.maxsize: False, + }[opts['legacy']] + return opts + + +def _get_legacy_print_mode(): + """Return the legacy print mode as an int.""" + return format_options.get()['legacy'] + + +@set_module('numpy') +@contextlib.contextmanager +def printoptions(*args, **kwargs): + """Context manager for setting print options. + + Set print options for the scope of the `with` block, and restore the old + options at the end. See `set_printoptions` for the full description of + available options. + + Examples + -------- + >>> import numpy as np + + >>> from numpy.testing import assert_equal + >>> with np.printoptions(precision=2): + ... np.array([2.0]) / 3 + array([0.67]) + + The `as`-clause of the `with`-statement gives the current print options: + + >>> with np.printoptions(precision=2) as opts: + ... assert_equal(opts, np.get_printoptions()) + + See Also + -------- + set_printoptions, get_printoptions + + """ + token = _set_printoptions(*args, **kwargs) + + try: + yield get_printoptions() + finally: + format_options.reset(token) + + +def _leading_trailing(a, edgeitems, index=()): + """ + Keep only the N-D corners (leading and trailing edges) of an array. + + Should be passed a base-class ndarray, since it makes no guarantees about + preserving subclasses. + """ + axis = len(index) + if axis == a.ndim: + return a[index] + + if a.shape[axis] > 2*edgeitems: + return concatenate(( + _leading_trailing(a, edgeitems, index + np.index_exp[:edgeitems]), + _leading_trailing(a, edgeitems, index + np.index_exp[-edgeitems:]) + ), axis=axis) + else: + return _leading_trailing(a, edgeitems, index + np.index_exp[:]) + + +def _object_format(o): + """ Object arrays containing lists should be printed unambiguously """ + if type(o) is list: + fmt = 'list({!r})' + else: + fmt = '{!r}' + return fmt.format(o) + +def repr_format(x): + if isinstance(x, (np.str_, np.bytes_)): + return repr(x.item()) + return repr(x) + +def str_format(x): + if isinstance(x, (np.str_, np.bytes_)): + return str(x.item()) + return str(x) + +def _get_formatdict(data, *, precision, floatmode, suppress, sign, legacy, + formatter, **kwargs): + # note: extra arguments in kwargs are ignored + + # wrapped in lambdas to avoid taking a code path + # with the wrong type of data + formatdict = { + 'bool': lambda: BoolFormat(data), + 'int': lambda: IntegerFormat(data, sign), + 'float': lambda: FloatingFormat( + data, precision, floatmode, suppress, sign, legacy=legacy), + 'longfloat': lambda: FloatingFormat( + data, precision, floatmode, suppress, sign, legacy=legacy), + 'complexfloat': lambda: ComplexFloatingFormat( + data, precision, floatmode, suppress, sign, legacy=legacy), + 'longcomplexfloat': lambda: ComplexFloatingFormat( + data, precision, floatmode, suppress, sign, legacy=legacy), + 'datetime': lambda: DatetimeFormat(data, legacy=legacy), + 'timedelta': lambda: TimedeltaFormat(data), + 'object': lambda: _object_format, + 'void': lambda: str_format, + 'numpystr': lambda: repr_format} + + # we need to wrap values in `formatter` in a lambda, so that the interface + # is the same as the above values. + def indirect(x): + return lambda: x + + if formatter is not None: + fkeys = [k for k in formatter.keys() if formatter[k] is not None] + if 'all' in fkeys: + for key in formatdict.keys(): + formatdict[key] = indirect(formatter['all']) + if 'int_kind' in fkeys: + for key in ['int']: + formatdict[key] = indirect(formatter['int_kind']) + if 'float_kind' in fkeys: + for key in ['float', 'longfloat']: + formatdict[key] = indirect(formatter['float_kind']) + if 'complex_kind' in fkeys: + for key in ['complexfloat', 'longcomplexfloat']: + formatdict[key] = indirect(formatter['complex_kind']) + if 'str_kind' in fkeys: + formatdict['numpystr'] = indirect(formatter['str_kind']) + for key in formatdict.keys(): + if key in fkeys: + formatdict[key] = indirect(formatter[key]) + + return formatdict + +def _get_format_function(data, **options): + """ + find the right formatting function for the dtype_ + """ + dtype_ = data.dtype + dtypeobj = dtype_.type + formatdict = _get_formatdict(data, **options) + if dtypeobj is None: + return formatdict["numpystr"]() + elif issubclass(dtypeobj, _nt.bool): + return formatdict['bool']() + elif issubclass(dtypeobj, _nt.integer): + if issubclass(dtypeobj, _nt.timedelta64): + return formatdict['timedelta']() + else: + return formatdict['int']() + elif issubclass(dtypeobj, _nt.floating): + if issubclass(dtypeobj, _nt.longdouble): + return formatdict['longfloat']() + else: + return formatdict['float']() + elif issubclass(dtypeobj, _nt.complexfloating): + if issubclass(dtypeobj, _nt.clongdouble): + return formatdict['longcomplexfloat']() + else: + return formatdict['complexfloat']() + elif issubclass(dtypeobj, (_nt.str_, _nt.bytes_)): + return formatdict['numpystr']() + elif issubclass(dtypeobj, _nt.datetime64): + return formatdict['datetime']() + elif issubclass(dtypeobj, _nt.object_): + return formatdict['object']() + elif issubclass(dtypeobj, _nt.void): + if dtype_.names is not None: + return StructuredVoidFormat.from_data(data, **options) + else: + return formatdict['void']() + else: + return formatdict['numpystr']() + + +def _recursive_guard(fillvalue='...'): + """ + Like the python 3.2 reprlib.recursive_repr, but forwards *args and **kwargs + + Decorates a function such that if it calls itself with the same first + argument, it returns `fillvalue` instead of recursing. + + Largely copied from reprlib.recursive_repr + """ + + def decorating_function(f): + repr_running = set() + + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + key = id(self), get_ident() + if key in repr_running: + return fillvalue + repr_running.add(key) + try: + return f(self, *args, **kwargs) + finally: + repr_running.discard(key) + + return wrapper + + return decorating_function + + +# gracefully handle recursive calls, when object arrays contain themselves +@_recursive_guard() +def _array2string(a, options, separator=' ', prefix=""): + # The formatter __init__s in _get_format_function cannot deal with + # subclasses yet, and we also need to avoid recursion issues in + # _formatArray with subclasses which return 0d arrays in place of scalars + data = asarray(a) + if a.shape == (): + a = data + + if a.size > options['threshold']: + summary_insert = "..." + data = _leading_trailing(data, options['edgeitems']) + else: + summary_insert = "" + + # find the right formatting function for the array + format_function = _get_format_function(data, **options) + + # skip over "[" + next_line_prefix = " " + # skip over array( + next_line_prefix += " "*len(prefix) + + lst = _formatArray(a, format_function, options['linewidth'], + next_line_prefix, separator, options['edgeitems'], + summary_insert, options['legacy']) + return lst + + +def _array2string_dispatcher( + a, max_line_width=None, precision=None, + suppress_small=None, separator=None, prefix=None, + style=None, formatter=None, threshold=None, + edgeitems=None, sign=None, floatmode=None, suffix=None, + *, legacy=None): + return (a,) + + +@array_function_dispatch(_array2string_dispatcher, module='numpy') +def array2string(a, max_line_width=None, precision=None, + suppress_small=None, separator=' ', prefix="", + style=np._NoValue, formatter=None, threshold=None, + edgeitems=None, sign=None, floatmode=None, suffix="", + *, legacy=None): + """ + Return a string representation of an array. + + Parameters + ---------- + a : ndarray + Input array. + max_line_width : int, optional + Inserts newlines if text is longer than `max_line_width`. + Defaults to ``numpy.get_printoptions()['linewidth']``. + precision : int or None, optional + Floating point precision. + Defaults to ``numpy.get_printoptions()['precision']``. + suppress_small : bool, optional + Represent numbers "very close" to zero as zero; default is False. + Very close is defined by precision: if the precision is 8, e.g., + numbers smaller (in absolute value) than 5e-9 are represented as + zero. + Defaults to ``numpy.get_printoptions()['suppress']``. + separator : str, optional + Inserted between elements. + prefix : str, optional + suffix : str, optional + The length of the prefix and suffix strings are used to respectively + align and wrap the output. An array is typically printed as:: + + prefix + array2string(a) + suffix + + The output is left-padded by the length of the prefix string, and + wrapping is forced at the column ``max_line_width - len(suffix)``. + It should be noted that the content of prefix and suffix strings are + not included in the output. + style : _NoValue, optional + Has no effect, do not use. + + .. deprecated:: 1.14.0 + formatter : dict of callables, optional + If not None, the keys should indicate the type(s) that the respective + formatting function applies to. Callables should return a string. + Types that are not specified (by their corresponding keys) are handled + by the default formatters. Individual types for which a formatter + can be set are: + + - 'bool' + - 'int' + - 'timedelta' : a `numpy.timedelta64` + - 'datetime' : a `numpy.datetime64` + - 'float' + - 'longfloat' : 128-bit floats + - 'complexfloat' + - 'longcomplexfloat' : composed of two 128-bit floats + - 'void' : type `numpy.void` + - 'numpystr' : types `numpy.bytes_` and `numpy.str_` + + Other keys that can be used to set a group of types at once are: + + - 'all' : sets all types + - 'int_kind' : sets 'int' + - 'float_kind' : sets 'float' and 'longfloat' + - 'complex_kind' : sets 'complexfloat' and 'longcomplexfloat' + - 'str_kind' : sets 'numpystr' + threshold : int, optional + Total number of array elements which trigger summarization + rather than full repr. + Defaults to ``numpy.get_printoptions()['threshold']``. + edgeitems : int, optional + Number of array items in summary at beginning and end of + each dimension. + Defaults to ``numpy.get_printoptions()['edgeitems']``. + sign : string, either '-', '+', or ' ', optional + Controls printing of the sign of floating-point types. If '+', always + print the sign of positive values. If ' ', always prints a space + (whitespace character) in the sign position of positive values. If + '-', omit the sign character of positive values. + Defaults to ``numpy.get_printoptions()['sign']``. + + .. versionchanged:: 2.0 + The sign parameter can now be an integer type, previously + types were floating-point types. + + floatmode : str, optional + Controls the interpretation of the `precision` option for + floating-point types. + Defaults to ``numpy.get_printoptions()['floatmode']``. + Can take the following values: + + - 'fixed': Always print exactly `precision` fractional digits, + even if this would print more or fewer digits than + necessary to specify the value uniquely. + - 'unique': Print the minimum number of fractional digits necessary + to represent each value uniquely. Different elements may + have a different number of digits. The value of the + `precision` option is ignored. + - 'maxprec': Print at most `precision` fractional digits, but if + an element can be uniquely represented with fewer digits + only print it with that many. + - 'maxprec_equal': Print at most `precision` fractional digits, + but if every element in the array can be uniquely + represented with an equal number of fewer digits, use that + many digits for all elements. + legacy : string or `False`, optional + If set to the string ``'1.13'`` enables 1.13 legacy printing mode. This + approximates numpy 1.13 print output by including a space in the sign + position of floats and different behavior for 0d arrays. If set to + `False`, disables legacy mode. Unrecognized strings will be ignored + with a warning for forward compatibility. + + Returns + ------- + array_str : str + String representation of the array. + + Raises + ------ + TypeError + if a callable in `formatter` does not return a string. + + See Also + -------- + array_str, array_repr, set_printoptions, get_printoptions + + Notes + ----- + If a formatter is specified for a certain type, the `precision` keyword is + ignored for that type. + + This is a very flexible function; `array_repr` and `array_str` are using + `array2string` internally so keywords with the same name should work + identically in all three functions. + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1e-16,1,2,3]) + >>> np.array2string(x, precision=2, separator=',', + ... suppress_small=True) + '[0.,1.,2.,3.]' + + >>> x = np.arange(3.) + >>> np.array2string(x, formatter={'float_kind':lambda x: "%.2f" % x}) + '[0.00 1.00 2.00]' + + >>> x = np.arange(3) + >>> np.array2string(x, formatter={'int':lambda x: hex(x)}) + '[0x0 0x1 0x2]' + + """ + + overrides = _make_options_dict(precision, threshold, edgeitems, + max_line_width, suppress_small, None, None, + sign, formatter, floatmode, legacy) + options = format_options.get().copy() + options.update(overrides) + + if options['legacy'] <= 113: + if style is np._NoValue: + style = repr + + if a.shape == () and a.dtype.names is None: + return style(a.item()) + elif style is not np._NoValue: + # Deprecation 11-9-2017 v1.14 + warnings.warn("'style' argument is deprecated and no longer functional" + " except in 1.13 'legacy' mode", + DeprecationWarning, stacklevel=2) + + if options['legacy'] > 113: + options['linewidth'] -= len(suffix) + + # treat as a null array if any of shape elements == 0 + if a.size == 0: + return "[]" + + return _array2string(a, options, separator, prefix) + + +def _extendLine(s, line, word, line_width, next_line_prefix, legacy): + needs_wrap = len(line) + len(word) > line_width + if legacy > 113: + # don't wrap lines if it won't help + if len(line) <= len(next_line_prefix): + needs_wrap = False + + if needs_wrap: + s += line.rstrip() + "\n" + line = next_line_prefix + line += word + return s, line + + +def _extendLine_pretty(s, line, word, line_width, next_line_prefix, legacy): + """ + Extends line with nicely formatted (possibly multi-line) string ``word``. + """ + words = word.splitlines() + if len(words) == 1 or legacy <= 113: + return _extendLine(s, line, word, line_width, next_line_prefix, legacy) + + max_word_length = max(len(word) for word in words) + if (len(line) + max_word_length > line_width and + len(line) > len(next_line_prefix)): + s += line.rstrip() + '\n' + line = next_line_prefix + words[0] + indent = next_line_prefix + else: + indent = len(line)*' ' + line += words[0] + + for word in words[1::]: + s += line.rstrip() + '\n' + line = indent + word + + suffix_length = max_word_length - len(words[-1]) + line += suffix_length*' ' + + return s, line + +def _formatArray(a, format_function, line_width, next_line_prefix, + separator, edge_items, summary_insert, legacy): + """formatArray is designed for two modes of operation: + + 1. Full output + + 2. Summarized output + + """ + def recurser(index, hanging_indent, curr_width): + """ + By using this local function, we don't need to recurse with all the + arguments. Since this function is not created recursively, the cost is + not significant + """ + axis = len(index) + axes_left = a.ndim - axis + + if axes_left == 0: + return format_function(a[index]) + + # when recursing, add a space to align with the [ added, and reduce the + # length of the line by 1 + next_hanging_indent = hanging_indent + ' ' + if legacy <= 113: + next_width = curr_width + else: + next_width = curr_width - len(']') + + a_len = a.shape[axis] + show_summary = summary_insert and 2*edge_items < a_len + if show_summary: + leading_items = edge_items + trailing_items = edge_items + else: + leading_items = 0 + trailing_items = a_len + + # stringify the array with the hanging indent on the first line too + s = '' + + # last axis (rows) - wrap elements if they would not fit on one line + if axes_left == 1: + # the length up until the beginning of the separator / bracket + if legacy <= 113: + elem_width = curr_width - len(separator.rstrip()) + else: + elem_width = curr_width - max( + len(separator.rstrip()), len(']') + ) + + line = hanging_indent + for i in range(leading_items): + word = recurser(index + (i,), next_hanging_indent, next_width) + s, line = _extendLine_pretty( + s, line, word, elem_width, hanging_indent, legacy) + line += separator + + if show_summary: + s, line = _extendLine( + s, line, summary_insert, elem_width, hanging_indent, legacy + ) + if legacy <= 113: + line += ", " + else: + line += separator + + for i in range(trailing_items, 1, -1): + word = recurser(index + (-i,), next_hanging_indent, next_width) + s, line = _extendLine_pretty( + s, line, word, elem_width, hanging_indent, legacy) + line += separator + + if legacy <= 113: + # width of the separator is not considered on 1.13 + elem_width = curr_width + word = recurser(index + (-1,), next_hanging_indent, next_width) + s, line = _extendLine_pretty( + s, line, word, elem_width, hanging_indent, legacy) + + s += line + + # other axes - insert newlines between rows + else: + s = '' + line_sep = separator.rstrip() + '\n'*(axes_left - 1) + + for i in range(leading_items): + nested = recurser( + index + (i,), next_hanging_indent, next_width + ) + s += hanging_indent + nested + line_sep + + if show_summary: + if legacy <= 113: + # trailing space, fixed nbr of newlines, + # and fixed separator + s += hanging_indent + summary_insert + ", \n" + else: + s += hanging_indent + summary_insert + line_sep + + for i in range(trailing_items, 1, -1): + nested = recurser(index + (-i,), next_hanging_indent, + next_width) + s += hanging_indent + nested + line_sep + + nested = recurser(index + (-1,), next_hanging_indent, next_width) + s += hanging_indent + nested + + # remove the hanging indent, and wrap in [] + s = '[' + s[len(hanging_indent):] + ']' + return s + + try: + # invoke the recursive part with an initial index and prefix + return recurser(index=(), + hanging_indent=next_line_prefix, + curr_width=line_width) + finally: + # recursive closures have a cyclic reference to themselves, which + # requires gc to collect (gh-10620). To avoid this problem, for + # performance and PyPy friendliness, we break the cycle: + recurser = None + +def _none_or_positive_arg(x, name): + if x is None: + return -1 + if x < 0: + raise ValueError("{} must be >= 0".format(name)) + return x + +class FloatingFormat: + """ Formatter for subtypes of np.floating """ + def __init__(self, data, precision, floatmode, suppress_small, sign=False, + *, legacy=None): + # for backcompatibility, accept bools + if isinstance(sign, bool): + sign = '+' if sign else '-' + + self._legacy = legacy + if self._legacy <= 113: + # when not 0d, legacy does not support '-' + if data.shape != () and sign == '-': + sign = ' ' + + self.floatmode = floatmode + if floatmode == 'unique': + self.precision = None + else: + self.precision = precision + + self.precision = _none_or_positive_arg(self.precision, 'precision') + + self.suppress_small = suppress_small + self.sign = sign + self.exp_format = False + self.large_exponent = False + self.fillFormat(data) + + def fillFormat(self, data): + # only the finite values are used to compute the number of digits + finite_vals = data[isfinite(data)] + + # choose exponential mode based on the non-zero finite values: + abs_non_zero = absolute(finite_vals[finite_vals != 0]) + if len(abs_non_zero) != 0: + max_val = np.max(abs_non_zero) + min_val = np.min(abs_non_zero) + with errstate(over='ignore'): # division can overflow + if max_val >= 1.e8 or (not self.suppress_small and + (min_val < 0.0001 or max_val/min_val > 1000.)): + self.exp_format = True + + # do a first pass of printing all the numbers, to determine sizes + if len(finite_vals) == 0: + self.pad_left = 0 + self.pad_right = 0 + self.trim = '.' + self.exp_size = -1 + self.unique = True + self.min_digits = None + elif self.exp_format: + trim, unique = '.', True + if self.floatmode == 'fixed' or self._legacy <= 113: + trim, unique = 'k', False + strs = (dragon4_scientific(x, precision=self.precision, + unique=unique, trim=trim, sign=self.sign == '+') + for x in finite_vals) + frac_strs, _, exp_strs = zip(*(s.partition('e') for s in strs)) + int_part, frac_part = zip(*(s.split('.') for s in frac_strs)) + self.exp_size = max(len(s) for s in exp_strs) - 1 + + self.trim = 'k' + self.precision = max(len(s) for s in frac_part) + self.min_digits = self.precision + self.unique = unique + + # for back-compat with np 1.13, use 2 spaces & sign and full prec + if self._legacy <= 113: + self.pad_left = 3 + else: + # this should be only 1 or 2. Can be calculated from sign. + self.pad_left = max(len(s) for s in int_part) + # pad_right is only needed for nan length calculation + self.pad_right = self.exp_size + 2 + self.precision + else: + trim, unique = '.', True + if self.floatmode == 'fixed': + trim, unique = 'k', False + strs = (dragon4_positional(x, precision=self.precision, + fractional=True, + unique=unique, trim=trim, + sign=self.sign == '+') + for x in finite_vals) + int_part, frac_part = zip(*(s.split('.') for s in strs)) + if self._legacy <= 113: + self.pad_left = 1 + max(len(s.lstrip('-+')) for s in int_part) + else: + self.pad_left = max(len(s) for s in int_part) + self.pad_right = max(len(s) for s in frac_part) + self.exp_size = -1 + self.unique = unique + + if self.floatmode in ['fixed', 'maxprec_equal']: + self.precision = self.min_digits = self.pad_right + self.trim = 'k' + else: + self.trim = '.' + self.min_digits = 0 + + if self._legacy > 113: + # account for sign = ' ' by adding one to pad_left + if self.sign == ' ' and not any(np.signbit(finite_vals)): + self.pad_left += 1 + + # if there are non-finite values, may need to increase pad_left + if data.size != finite_vals.size: + neginf = self.sign != '-' or any(data[isinf(data)] < 0) + offset = self.pad_right + 1 # +1 for decimal pt + current_options = format_options.get() + self.pad_left = max( + self.pad_left, len(current_options['nanstr']) - offset, + len(current_options['infstr']) + neginf - offset + ) + + def __call__(self, x): + if not np.isfinite(x): + with errstate(invalid='ignore'): + current_options = format_options.get() + if np.isnan(x): + sign = '+' if self.sign == '+' else '' + ret = sign + current_options['nanstr'] + else: # isinf + sign = '-' if x < 0 else '+' if self.sign == '+' else '' + ret = sign + current_options['infstr'] + return ' '*( + self.pad_left + self.pad_right + 1 - len(ret) + ) + ret + + if self.exp_format: + return dragon4_scientific(x, + precision=self.precision, + min_digits=self.min_digits, + unique=self.unique, + trim=self.trim, + sign=self.sign == '+', + pad_left=self.pad_left, + exp_digits=self.exp_size) + else: + return dragon4_positional(x, + precision=self.precision, + min_digits=self.min_digits, + unique=self.unique, + fractional=True, + trim=self.trim, + sign=self.sign == '+', + pad_left=self.pad_left, + pad_right=self.pad_right) + + +@set_module('numpy') +def format_float_scientific(x, precision=None, unique=True, trim='k', + sign=False, pad_left=None, exp_digits=None, + min_digits=None): + """ + Format a floating-point scalar as a decimal string in scientific notation. + + Provides control over rounding, trimming and padding. Uses and assumes + IEEE unbiased rounding. Uses the "Dragon4" algorithm. + + Parameters + ---------- + x : python float or numpy floating scalar + Value to format. + precision : non-negative integer or None, optional + Maximum number of digits to print. May be None if `unique` is + `True`, but must be an integer if unique is `False`. + unique : boolean, optional + If `True`, use a digit-generation strategy which gives the shortest + representation which uniquely identifies the floating-point number from + other values of the same type, by judicious rounding. If `precision` + is given fewer digits than necessary can be printed. If `min_digits` + is given more can be printed, in which cases the last digit is rounded + with unbiased rounding. + If `False`, digits are generated as if printing an infinite-precision + value and stopping after `precision` digits, rounding the remaining + value with unbiased rounding + trim : one of 'k', '.', '0', '-', optional + Controls post-processing trimming of trailing digits, as follows: + + * 'k' : keep trailing zeros, keep decimal point (no trimming) + * '.' : trim all trailing zeros, leave decimal point + * '0' : trim all but the zero before the decimal point. Insert the + zero if it is missing. + * '-' : trim trailing zeros and any trailing decimal point + sign : boolean, optional + Whether to show the sign for positive values. + pad_left : non-negative integer, optional + Pad the left side of the string with whitespace until at least that + many characters are to the left of the decimal point. + exp_digits : non-negative integer, optional + Pad the exponent with zeros until it contains at least this + many digits. If omitted, the exponent will be at least 2 digits. + min_digits : non-negative integer or None, optional + Minimum number of digits to print. This only has an effect for + `unique=True`. In that case more digits than necessary to uniquely + identify the value may be printed and rounded unbiased. + + .. versionadded:: 1.21.0 + + Returns + ------- + rep : string + The string representation of the floating point value + + See Also + -------- + format_float_positional + + Examples + -------- + >>> import numpy as np + >>> np.format_float_scientific(np.float32(np.pi)) + '3.1415927e+00' + >>> s = np.float32(1.23e24) + >>> np.format_float_scientific(s, unique=False, precision=15) + '1.230000071797338e+24' + >>> np.format_float_scientific(s, exp_digits=4) + '1.23e+0024' + """ + precision = _none_or_positive_arg(precision, 'precision') + pad_left = _none_or_positive_arg(pad_left, 'pad_left') + exp_digits = _none_or_positive_arg(exp_digits, 'exp_digits') + min_digits = _none_or_positive_arg(min_digits, 'min_digits') + if min_digits > 0 and precision > 0 and min_digits > precision: + raise ValueError("min_digits must be less than or equal to precision") + return dragon4_scientific(x, precision=precision, unique=unique, + trim=trim, sign=sign, pad_left=pad_left, + exp_digits=exp_digits, min_digits=min_digits) + + +@set_module('numpy') +def format_float_positional(x, precision=None, unique=True, + fractional=True, trim='k', sign=False, + pad_left=None, pad_right=None, min_digits=None): + """ + Format a floating-point scalar as a decimal string in positional notation. + + Provides control over rounding, trimming and padding. Uses and assumes + IEEE unbiased rounding. Uses the "Dragon4" algorithm. + + Parameters + ---------- + x : python float or numpy floating scalar + Value to format. + precision : non-negative integer or None, optional + Maximum number of digits to print. May be None if `unique` is + `True`, but must be an integer if unique is `False`. + unique : boolean, optional + If `True`, use a digit-generation strategy which gives the shortest + representation which uniquely identifies the floating-point number from + other values of the same type, by judicious rounding. If `precision` + is given fewer digits than necessary can be printed, or if `min_digits` + is given more can be printed, in which cases the last digit is rounded + with unbiased rounding. + If `False`, digits are generated as if printing an infinite-precision + value and stopping after `precision` digits, rounding the remaining + value with unbiased rounding + fractional : boolean, optional + If `True`, the cutoffs of `precision` and `min_digits` refer to the + total number of digits after the decimal point, including leading + zeros. + If `False`, `precision` and `min_digits` refer to the total number of + significant digits, before or after the decimal point, ignoring leading + zeros. + trim : one of 'k', '.', '0', '-', optional + Controls post-processing trimming of trailing digits, as follows: + + * 'k' : keep trailing zeros, keep decimal point (no trimming) + * '.' : trim all trailing zeros, leave decimal point + * '0' : trim all but the zero before the decimal point. Insert the + zero if it is missing. + * '-' : trim trailing zeros and any trailing decimal point + sign : boolean, optional + Whether to show the sign for positive values. + pad_left : non-negative integer, optional + Pad the left side of the string with whitespace until at least that + many characters are to the left of the decimal point. + pad_right : non-negative integer, optional + Pad the right side of the string with whitespace until at least that + many characters are to the right of the decimal point. + min_digits : non-negative integer or None, optional + Minimum number of digits to print. Only has an effect if `unique=True` + in which case additional digits past those necessary to uniquely + identify the value may be printed, rounding the last additional digit. + + .. versionadded:: 1.21.0 + + Returns + ------- + rep : string + The string representation of the floating point value + + See Also + -------- + format_float_scientific + + Examples + -------- + >>> import numpy as np + >>> np.format_float_positional(np.float32(np.pi)) + '3.1415927' + >>> np.format_float_positional(np.float16(np.pi)) + '3.14' + >>> np.format_float_positional(np.float16(0.3)) + '0.3' + >>> np.format_float_positional(np.float16(0.3), unique=False, precision=10) + '0.3000488281' + """ + precision = _none_or_positive_arg(precision, 'precision') + pad_left = _none_or_positive_arg(pad_left, 'pad_left') + pad_right = _none_or_positive_arg(pad_right, 'pad_right') + min_digits = _none_or_positive_arg(min_digits, 'min_digits') + if not fractional and precision == 0: + raise ValueError("precision must be greater than 0 if " + "fractional=False") + if min_digits > 0 and precision > 0 and min_digits > precision: + raise ValueError("min_digits must be less than or equal to precision") + return dragon4_positional(x, precision=precision, unique=unique, + fractional=fractional, trim=trim, + sign=sign, pad_left=pad_left, + pad_right=pad_right, min_digits=min_digits) + +class IntegerFormat: + def __init__(self, data, sign='-'): + if data.size > 0: + data_max = np.max(data) + data_min = np.min(data) + data_max_str_len = len(str(data_max)) + if sign == ' ' and data_min < 0: + sign = '-' + if data_max >= 0 and sign in "+ ": + data_max_str_len += 1 + max_str_len = max(data_max_str_len, + len(str(data_min))) + else: + max_str_len = 0 + self.format = f'{{:{sign}{max_str_len}d}}' + + def __call__(self, x): + return self.format.format(x) + +class BoolFormat: + def __init__(self, data, **kwargs): + # add an extra space so " True" and "False" have the same length and + # array elements align nicely when printed, except in 0d arrays + self.truestr = ' True' if data.shape != () else 'True' + + def __call__(self, x): + return self.truestr if x else "False" + + +class ComplexFloatingFormat: + """ Formatter for subtypes of np.complexfloating """ + def __init__(self, x, precision, floatmode, suppress_small, + sign=False, *, legacy=None): + # for backcompatibility, accept bools + if isinstance(sign, bool): + sign = '+' if sign else '-' + + floatmode_real = floatmode_imag = floatmode + if legacy <= 113: + floatmode_real = 'maxprec_equal' + floatmode_imag = 'maxprec' + + self.real_format = FloatingFormat( + x.real, precision, floatmode_real, suppress_small, + sign=sign, legacy=legacy + ) + self.imag_format = FloatingFormat( + x.imag, precision, floatmode_imag, suppress_small, + sign='+', legacy=legacy + ) + + def __call__(self, x): + r = self.real_format(x.real) + i = self.imag_format(x.imag) + + # add the 'j' before the terminal whitespace in i + sp = len(i.rstrip()) + i = i[:sp] + 'j' + i[sp:] + + return r + i + + +class _TimelikeFormat: + def __init__(self, data): + non_nat = data[~isnat(data)] + if len(non_nat) > 0: + # Max str length of non-NaT elements + max_str_len = max(len(self._format_non_nat(np.max(non_nat))), + len(self._format_non_nat(np.min(non_nat)))) + else: + max_str_len = 0 + if len(non_nat) < data.size: + # data contains a NaT + max_str_len = max(max_str_len, 5) + self._format = '%{}s'.format(max_str_len) + self._nat = "'NaT'".rjust(max_str_len) + + def _format_non_nat(self, x): + # override in subclass + raise NotImplementedError + + def __call__(self, x): + if isnat(x): + return self._nat + else: + return self._format % self._format_non_nat(x) + + +class DatetimeFormat(_TimelikeFormat): + def __init__(self, x, unit=None, timezone=None, casting='same_kind', + legacy=False): + # Get the unit from the dtype + if unit is None: + if x.dtype.kind == 'M': + unit = datetime_data(x.dtype)[0] + else: + unit = 's' + + if timezone is None: + timezone = 'naive' + self.timezone = timezone + self.unit = unit + self.casting = casting + self.legacy = legacy + + # must be called after the above are configured + super().__init__(x) + + def __call__(self, x): + if self.legacy <= 113: + return self._format_non_nat(x) + return super().__call__(x) + + def _format_non_nat(self, x): + return "'%s'" % datetime_as_string(x, + unit=self.unit, + timezone=self.timezone, + casting=self.casting) + + +class TimedeltaFormat(_TimelikeFormat): + def _format_non_nat(self, x): + return str(x.astype('i8')) + + +class SubArrayFormat: + def __init__(self, format_function, **options): + self.format_function = format_function + self.threshold = options['threshold'] + self.edge_items = options['edgeitems'] + + def __call__(self, a): + self.summary_insert = "..." if a.size > self.threshold else "" + return self.format_array(a) + + def format_array(self, a): + if np.ndim(a) == 0: + return self.format_function(a) + + if self.summary_insert and a.shape[0] > 2*self.edge_items: + formatted = ( + [self.format_array(a_) for a_ in a[:self.edge_items]] + + [self.summary_insert] + + [self.format_array(a_) for a_ in a[-self.edge_items:]] + ) + else: + formatted = [self.format_array(a_) for a_ in a] + + return "[" + ", ".join(formatted) + "]" + + +class StructuredVoidFormat: + """ + Formatter for structured np.void objects. + + This does not work on structured alias types like + np.dtype(('i4', 'i2,i2')), as alias scalars lose their field information, + and the implementation relies upon np.void.__getitem__. + """ + def __init__(self, format_functions): + self.format_functions = format_functions + + @classmethod + def from_data(cls, data, **options): + """ + This is a second way to initialize StructuredVoidFormat, + using the raw data as input. Added to avoid changing + the signature of __init__. + """ + format_functions = [] + for field_name in data.dtype.names: + format_function = _get_format_function(data[field_name], **options) + if data.dtype[field_name].shape != (): + format_function = SubArrayFormat(format_function, **options) + format_functions.append(format_function) + return cls(format_functions) + + def __call__(self, x): + str_fields = [ + format_function(field) + for field, format_function in zip(x, self.format_functions) + ] + if len(str_fields) == 1: + return "({},)".format(str_fields[0]) + else: + return "({})".format(", ".join(str_fields)) + + +def _void_scalar_to_string(x, is_repr=True): + """ + Implements the repr for structured-void scalars. It is called from the + scalartypes.c.src code, and is placed here because it uses the elementwise + formatters defined above. + """ + options = format_options.get().copy() + + if options["legacy"] <= 125: + return StructuredVoidFormat.from_data(array(x), **options)(x) + + if options.get('formatter') is None: + options['formatter'] = {} + options['formatter'].setdefault('float_kind', str) + val_repr = StructuredVoidFormat.from_data(array(x), **options)(x) + if not is_repr: + return val_repr + cls = type(x) + cls_fqn = cls.__module__.replace("numpy", "np") + "." + cls.__name__ + void_dtype = np.dtype((np.void, x.dtype)) + return f"{cls_fqn}({val_repr}, dtype={void_dtype!s})" + + +_typelessdata = [int_, float64, complex128, _nt.bool] + + +def dtype_is_implied(dtype): + """ + Determine if the given dtype is implied by the representation + of its values. + + Parameters + ---------- + dtype : dtype + Data type + + Returns + ------- + implied : bool + True if the dtype is implied by the representation of its values. + + Examples + -------- + >>> import numpy as np + >>> np._core.arrayprint.dtype_is_implied(int) + True + >>> np.array([1, 2, 3], int) + array([1, 2, 3]) + >>> np._core.arrayprint.dtype_is_implied(np.int8) + False + >>> np.array([1, 2, 3], np.int8) + array([1, 2, 3], dtype=int8) + """ + dtype = np.dtype(dtype) + if format_options.get()['legacy'] <= 113 and dtype.type == np.bool: + return False + + # not just void types can be structured, and names are not part of the repr + if dtype.names is not None: + return False + + # should care about endianness *unless size is 1* (e.g., int8, bool) + if not dtype.isnative: + return False + + return dtype.type in _typelessdata + + +def dtype_short_repr(dtype): + """ + Convert a dtype to a short form which evaluates to the same dtype. + + The intent is roughly that the following holds + + >>> from numpy import * + >>> dt = np.int64([1, 2]).dtype + >>> assert eval(dtype_short_repr(dt)) == dt + """ + if type(dtype).__repr__ != np.dtype.__repr__: + # TODO: Custom repr for user DTypes, logic should likely move. + return repr(dtype) + if dtype.names is not None: + # structured dtypes give a list or tuple repr + return str(dtype) + elif issubclass(dtype.type, flexible): + # handle these separately so they don't give garbage like str256 + return "'%s'" % str(dtype) + + typename = dtype.name + if not dtype.isnative: + # deal with cases like dtype(' 210 + and arr.size > current_options['threshold']): + extras.append(f"shape={arr.shape}") + if not dtype_is_implied(arr.dtype) or arr.size == 0: + extras.append(f"dtype={dtype_short_repr(arr.dtype)}") + + if not extras: + return prefix + lst + ")" + + arr_str = prefix + lst + "," + extra_str = ", ".join(extras) + ")" + # compute whether we should put extras on a new line: Do so if adding the + # extras would extend the last line past max_line_width. + # Note: This line gives the correct result even when rfind returns -1. + last_line_len = len(arr_str) - (arr_str.rfind('\n') + 1) + spacer = " " + if current_options['legacy'] <= 113: + if issubclass(arr.dtype.type, flexible): + spacer = '\n' + ' '*len(prefix) + elif last_line_len + len(extra_str) + 1 > max_line_width: + spacer = '\n' + ' '*len(prefix) + + return arr_str + spacer + extra_str + + +def _array_repr_dispatcher( + arr, max_line_width=None, precision=None, suppress_small=None): + return (arr,) + + +@array_function_dispatch(_array_repr_dispatcher, module='numpy') +def array_repr(arr, max_line_width=None, precision=None, suppress_small=None): + """ + Return the string representation of an array. + + Parameters + ---------- + arr : ndarray + Input array. + max_line_width : int, optional + Inserts newlines if text is longer than `max_line_width`. + Defaults to ``numpy.get_printoptions()['linewidth']``. + precision : int, optional + Floating point precision. + Defaults to ``numpy.get_printoptions()['precision']``. + suppress_small : bool, optional + Represent numbers "very close" to zero as zero; default is False. + Very close is defined by precision: if the precision is 8, e.g., + numbers smaller (in absolute value) than 5e-9 are represented as + zero. + Defaults to ``numpy.get_printoptions()['suppress']``. + + Returns + ------- + string : str + The string representation of an array. + + See Also + -------- + array_str, array2string, set_printoptions + + Examples + -------- + >>> import numpy as np + >>> np.array_repr(np.array([1,2])) + 'array([1, 2])' + >>> np.array_repr(np.ma.array([0.])) + 'MaskedArray([0.])' + >>> np.array_repr(np.array([], np.int32)) + 'array([], dtype=int32)' + + >>> x = np.array([1e-6, 4e-7, 2, 3]) + >>> np.array_repr(x, precision=6, suppress_small=True) + 'array([0.000001, 0. , 2. , 3. ])' + + """ + return _array_repr_implementation( + arr, max_line_width, precision, suppress_small) + + +@_recursive_guard() +def _guarded_repr_or_str(v): + if isinstance(v, bytes): + return repr(v) + return str(v) + + +def _array_str_implementation( + a, max_line_width=None, precision=None, suppress_small=None, + array2string=array2string): + """Internal version of array_str() that allows overriding array2string.""" + if (format_options.get()['legacy'] <= 113 and + a.shape == () and not a.dtype.names): + return str(a.item()) + + # the str of 0d arrays is a special case: It should appear like a scalar, + # so floats are not truncated by `precision`, and strings are not wrapped + # in quotes. So we return the str of the scalar value. + if a.shape == (): + # obtain a scalar and call str on it, avoiding problems for subclasses + # for which indexing with () returns a 0d instead of a scalar by using + # ndarray's getindex. Also guard against recursive 0d object arrays. + return _guarded_repr_or_str(np.ndarray.__getitem__(a, ())) + + return array2string(a, max_line_width, precision, suppress_small, ' ', "") + + +def _array_str_dispatcher( + a, max_line_width=None, precision=None, suppress_small=None): + return (a,) + + +@array_function_dispatch(_array_str_dispatcher, module='numpy') +def array_str(a, max_line_width=None, precision=None, suppress_small=None): + """ + Return a string representation of the data in an array. + + The data in the array is returned as a single string. This function is + similar to `array_repr`, the difference being that `array_repr` also + returns information on the kind of array and its data type. + + Parameters + ---------- + a : ndarray + Input array. + max_line_width : int, optional + Inserts newlines if text is longer than `max_line_width`. + Defaults to ``numpy.get_printoptions()['linewidth']``. + precision : int, optional + Floating point precision. + Defaults to ``numpy.get_printoptions()['precision']``. + suppress_small : bool, optional + Represent numbers "very close" to zero as zero; default is False. + Very close is defined by precision: if the precision is 8, e.g., + numbers smaller (in absolute value) than 5e-9 are represented as + zero. + Defaults to ``numpy.get_printoptions()['suppress']``. + + See Also + -------- + array2string, array_repr, set_printoptions + + Examples + -------- + >>> import numpy as np + >>> np.array_str(np.arange(3)) + '[0 1 2]' + + """ + return _array_str_implementation( + a, max_line_width, precision, suppress_small) + + +# needed if __array_function__ is disabled +_array2string_impl = getattr(array2string, '__wrapped__', array2string) +_default_array_str = functools.partial(_array_str_implementation, + array2string=_array2string_impl) +_default_array_repr = functools.partial(_array_repr_implementation, + array2string=_array2string_impl) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/arrayprint.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/arrayprint.pyi new file mode 100644 index 0000000000000000000000000000000000000000..142d77ff02c5c04f355714901dabb0aaeb2946cc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/arrayprint.pyi @@ -0,0 +1,229 @@ +from collections.abc import Callable + +# Using a private class is by no means ideal, but it is simply a consequence +# of a `contextlib.context` returning an instance of aforementioned class +from contextlib import _GeneratorContextManager +from typing import Any, Final, Literal, SupportsIndex, TypeAlias, TypedDict, overload, type_check_only + +from typing_extensions import deprecated + +import numpy as np +from numpy._globals import _NoValueType +from numpy._typing import NDArray, _CharLike_co, _FloatLike_co + +__all__ = [ + "array2string", + "array_repr", + "array_str", + "format_float_positional", + "format_float_scientific", + "get_printoptions", + "printoptions", + "set_printoptions", +] + +### + +_FloatMode: TypeAlias = Literal["fixed", "unique", "maxprec", "maxprec_equal"] +_LegacyNoStyle: TypeAlias = Literal["1.21", "1.25", "2.1", False] +_Legacy: TypeAlias = Literal["1.13", _LegacyNoStyle] +_Sign: TypeAlias = Literal["-", "+", " "] +_Trim: TypeAlias = Literal["k", ".", "0", "-"] +_ReprFunc: TypeAlias = Callable[[NDArray[Any]], str] + +@type_check_only +class _FormatDict(TypedDict, total=False): + bool: Callable[[np.bool], str] + int: Callable[[np.integer], str] + timedelta: Callable[[np.timedelta64], str] + datetime: Callable[[np.datetime64], str] + float: Callable[[np.floating], str] + longfloat: Callable[[np.longdouble], str] + complexfloat: Callable[[np.complexfloating], str] + longcomplexfloat: Callable[[np.clongdouble], str] + void: Callable[[np.void], str] + numpystr: Callable[[_CharLike_co], str] + object: Callable[[object], str] + all: Callable[[object], str] + int_kind: Callable[[np.integer], str] + float_kind: Callable[[np.floating], str] + complex_kind: Callable[[np.complexfloating], str] + str_kind: Callable[[_CharLike_co], str] + +@type_check_only +class _FormatOptions(TypedDict): + precision: int + threshold: int + edgeitems: int + linewidth: int + suppress: bool + nanstr: str + infstr: str + formatter: _FormatDict | None + sign: _Sign + floatmode: _FloatMode + legacy: _Legacy + +### + +__docformat__: Final = "restructuredtext" # undocumented + +def set_printoptions( + precision: None | SupportsIndex = ..., + threshold: None | int = ..., + edgeitems: None | int = ..., + linewidth: None | int = ..., + suppress: None | bool = ..., + nanstr: None | str = ..., + infstr: None | str = ..., + formatter: None | _FormatDict = ..., + sign: _Sign | None = None, + floatmode: _FloatMode | None = None, + *, + legacy: _Legacy | None = None, + override_repr: _ReprFunc | None = None, +) -> None: ... +def get_printoptions() -> _FormatOptions: ... + +# public numpy export +@overload # no style +def array2string( + a: NDArray[Any], + max_line_width: int | None = None, + precision: SupportsIndex | None = None, + suppress_small: bool | None = None, + separator: str = " ", + prefix: str = "", + style: _NoValueType = ..., + formatter: _FormatDict | None = None, + threshold: int | None = None, + edgeitems: int | None = None, + sign: _Sign | None = None, + floatmode: _FloatMode | None = None, + suffix: str = "", + *, + legacy: _Legacy | None = None, +) -> str: ... +@overload # style= (positional), legacy="1.13" +def array2string( + a: NDArray[Any], + max_line_width: int | None, + precision: SupportsIndex | None, + suppress_small: bool | None, + separator: str, + prefix: str, + style: _ReprFunc, + formatter: _FormatDict | None = None, + threshold: int | None = None, + edgeitems: int | None = None, + sign: _Sign | None = None, + floatmode: _FloatMode | None = None, + suffix: str = "", + *, + legacy: Literal["1.13"], +) -> str: ... +@overload # style= (keyword), legacy="1.13" +def array2string( + a: NDArray[Any], + max_line_width: int | None = None, + precision: SupportsIndex | None = None, + suppress_small: bool | None = None, + separator: str = " ", + prefix: str = "", + *, + style: _ReprFunc, + formatter: _FormatDict | None = None, + threshold: int | None = None, + edgeitems: int | None = None, + sign: _Sign | None = None, + floatmode: _FloatMode | None = None, + suffix: str = "", + legacy: Literal["1.13"], +) -> str: ... +@overload # style= (positional), legacy!="1.13" +@deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode") +def array2string( + a: NDArray[Any], + max_line_width: int | None, + precision: SupportsIndex | None, + suppress_small: bool | None, + separator: str, + prefix: str, + style: _ReprFunc, + formatter: _FormatDict | None = None, + threshold: int | None = None, + edgeitems: int | None = None, + sign: _Sign | None = None, + floatmode: _FloatMode | None = None, + suffix: str = "", + *, + legacy: _LegacyNoStyle | None = None, +) -> str: ... +@overload # style= (keyword), legacy="1.13" +@deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode") +def array2string( + a: NDArray[Any], + max_line_width: int | None = None, + precision: SupportsIndex | None = None, + suppress_small: bool | None = None, + separator: str = " ", + prefix: str = "", + *, + style: _ReprFunc, + formatter: _FormatDict | None = None, + threshold: int | None = None, + edgeitems: int | None = None, + sign: _Sign | None = None, + floatmode: _FloatMode | None = None, + suffix: str = "", + legacy: _LegacyNoStyle | None = None, +) -> str: ... + +def format_float_scientific( + x: _FloatLike_co, + precision: None | int = ..., + unique: bool = ..., + trim: _Trim = "k", + sign: bool = ..., + pad_left: None | int = ..., + exp_digits: None | int = ..., + min_digits: None | int = ..., +) -> str: ... +def format_float_positional( + x: _FloatLike_co, + precision: None | int = ..., + unique: bool = ..., + fractional: bool = ..., + trim: _Trim = "k", + sign: bool = ..., + pad_left: None | int = ..., + pad_right: None | int = ..., + min_digits: None | int = ..., +) -> str: ... +def array_repr( + arr: NDArray[Any], + max_line_width: None | int = ..., + precision: None | SupportsIndex = ..., + suppress_small: None | bool = ..., +) -> str: ... +def array_str( + a: NDArray[Any], + max_line_width: None | int = ..., + precision: None | SupportsIndex = ..., + suppress_small: None | bool = ..., +) -> str: ... +def printoptions( + precision: None | SupportsIndex = ..., + threshold: None | int = ..., + edgeitems: None | int = ..., + linewidth: None | int = ..., + suppress: None | bool = ..., + nanstr: None | str = ..., + infstr: None | str = ..., + formatter: None | _FormatDict = ..., + sign: None | _Sign = None, + floatmode: _FloatMode | None = None, + *, + legacy: _Legacy | None = None, + override_repr: _ReprFunc | None = None, +) -> _GeneratorContextManager[_FormatOptions]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/cversions.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/cversions.py new file mode 100644 index 0000000000000000000000000000000000000000..ef75a6f4f35cd29569c2d806927695f600066066 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/cversions.py @@ -0,0 +1,13 @@ +"""Simple script to compute the api hash of the current API. + +The API has is defined by numpy_api_order and ufunc_api_order. + +""" +from os.path import dirname + +from code_generators.genapi import fullapi_hash +from code_generators.numpy_api import full_api + +if __name__ == '__main__': + curdir = dirname(__file__) + print(fullapi_hash(full_api)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/defchararray.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/defchararray.py new file mode 100644 index 0000000000000000000000000000000000000000..8d8965423750a4b29ff048defbb94e2dee38fbe9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/defchararray.py @@ -0,0 +1,1414 @@ +""" +This module contains a set of functions for vectorized string +operations and methods. + +.. note:: + The `chararray` class exists for backwards compatibility with + Numarray, it is not recommended for new development. Starting from numpy + 1.4, if one needs arrays of strings, it is recommended to use arrays of + `dtype` `object_`, `bytes_` or `str_`, and use the free functions + in the `numpy.char` module for fast vectorized string operations. + +Some methods will only be available if the corresponding string method is +available in your version of Python. + +The preferred alias for `defchararray` is `numpy.char`. + +""" +import functools + +import numpy as np +from .._utils import set_module +from .numerictypes import bytes_, str_, character +from .numeric import ndarray, array as narray, asarray as asnarray +from numpy._core.multiarray import compare_chararrays +from numpy._core import overrides +from numpy.strings import * +from numpy.strings import ( + multiply as strings_multiply, + partition as strings_partition, + rpartition as strings_rpartition, +) +from numpy._core.strings import ( + _split as split, + _rsplit as rsplit, + _splitlines as splitlines, + _join as join, +) + +__all__ = [ + 'equal', 'not_equal', 'greater_equal', 'less_equal', + 'greater', 'less', 'str_len', 'add', 'multiply', 'mod', 'capitalize', + 'center', 'count', 'decode', 'encode', 'endswith', 'expandtabs', + 'find', 'index', 'isalnum', 'isalpha', 'isdigit', 'islower', 'isspace', + 'istitle', 'isupper', 'join', 'ljust', 'lower', 'lstrip', 'partition', + 'replace', 'rfind', 'rindex', 'rjust', 'rpartition', 'rsplit', + 'rstrip', 'split', 'splitlines', 'startswith', 'strip', 'swapcase', + 'title', 'translate', 'upper', 'zfill', 'isnumeric', 'isdecimal', + 'array', 'asarray', 'compare_chararrays', 'chararray' + ] + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy.char') + + +def _binary_op_dispatcher(x1, x2): + return (x1, x2) + + +@array_function_dispatch(_binary_op_dispatcher) +def equal(x1, x2): + """ + Return (x1 == x2) element-wise. + + Unlike `numpy.equal`, this comparison is performed by first + stripping whitespace characters from the end of the string. This + behavior is provided for backward-compatibility with numarray. + + Parameters + ---------- + x1, x2 : array_like of str or unicode + Input arrays of the same shape. + + Returns + ------- + out : ndarray + Output array of bools. + + Examples + -------- + >>> import numpy as np + >>> y = "aa " + >>> x = "aa" + >>> np.char.equal(x, y) + array(True) + + See Also + -------- + not_equal, greater_equal, less_equal, greater, less + """ + return compare_chararrays(x1, x2, '==', True) + + +@array_function_dispatch(_binary_op_dispatcher) +def not_equal(x1, x2): + """ + Return (x1 != x2) element-wise. + + Unlike `numpy.not_equal`, this comparison is performed by first + stripping whitespace characters from the end of the string. This + behavior is provided for backward-compatibility with numarray. + + Parameters + ---------- + x1, x2 : array_like of str or unicode + Input arrays of the same shape. + + Returns + ------- + out : ndarray + Output array of bools. + + See Also + -------- + equal, greater_equal, less_equal, greater, less + + Examples + -------- + >>> import numpy as np + >>> x1 = np.array(['a', 'b', 'c']) + >>> np.char.not_equal(x1, 'b') + array([ True, False, True]) + + """ + return compare_chararrays(x1, x2, '!=', True) + + +@array_function_dispatch(_binary_op_dispatcher) +def greater_equal(x1, x2): + """ + Return (x1 >= x2) element-wise. + + Unlike `numpy.greater_equal`, this comparison is performed by + first stripping whitespace characters from the end of the string. + This behavior is provided for backward-compatibility with + numarray. + + Parameters + ---------- + x1, x2 : array_like of str or unicode + Input arrays of the same shape. + + Returns + ------- + out : ndarray + Output array of bools. + + See Also + -------- + equal, not_equal, less_equal, greater, less + + Examples + -------- + >>> import numpy as np + >>> x1 = np.array(['a', 'b', 'c']) + >>> np.char.greater_equal(x1, 'b') + array([False, True, True]) + + """ + return compare_chararrays(x1, x2, '>=', True) + + +@array_function_dispatch(_binary_op_dispatcher) +def less_equal(x1, x2): + """ + Return (x1 <= x2) element-wise. + + Unlike `numpy.less_equal`, this comparison is performed by first + stripping whitespace characters from the end of the string. This + behavior is provided for backward-compatibility with numarray. + + Parameters + ---------- + x1, x2 : array_like of str or unicode + Input arrays of the same shape. + + Returns + ------- + out : ndarray + Output array of bools. + + See Also + -------- + equal, not_equal, greater_equal, greater, less + + Examples + -------- + >>> import numpy as np + >>> x1 = np.array(['a', 'b', 'c']) + >>> np.char.less_equal(x1, 'b') + array([ True, True, False]) + + """ + return compare_chararrays(x1, x2, '<=', True) + + +@array_function_dispatch(_binary_op_dispatcher) +def greater(x1, x2): + """ + Return (x1 > x2) element-wise. + + Unlike `numpy.greater`, this comparison is performed by first + stripping whitespace characters from the end of the string. This + behavior is provided for backward-compatibility with numarray. + + Parameters + ---------- + x1, x2 : array_like of str or unicode + Input arrays of the same shape. + + Returns + ------- + out : ndarray + Output array of bools. + + See Also + -------- + equal, not_equal, greater_equal, less_equal, less + + Examples + -------- + >>> import numpy as np + >>> x1 = np.array(['a', 'b', 'c']) + >>> np.char.greater(x1, 'b') + array([False, False, True]) + + """ + return compare_chararrays(x1, x2, '>', True) + + +@array_function_dispatch(_binary_op_dispatcher) +def less(x1, x2): + """ + Return (x1 < x2) element-wise. + + Unlike `numpy.greater`, this comparison is performed by first + stripping whitespace characters from the end of the string. This + behavior is provided for backward-compatibility with numarray. + + Parameters + ---------- + x1, x2 : array_like of str or unicode + Input arrays of the same shape. + + Returns + ------- + out : ndarray + Output array of bools. + + See Also + -------- + equal, not_equal, greater_equal, less_equal, greater + + Examples + -------- + >>> import numpy as np + >>> x1 = np.array(['a', 'b', 'c']) + >>> np.char.less(x1, 'b') + array([True, False, False]) + + """ + return compare_chararrays(x1, x2, '<', True) + + +@set_module("numpy.char") +def multiply(a, i): + """ + Return (a * i), that is string multiple concatenation, + element-wise. + + Values in ``i`` of less than 0 are treated as 0 (which yields an + empty string). + + Parameters + ---------- + a : array_like, with `np.bytes_` or `np.str_` dtype + + i : array_like, with any integer dtype + + Returns + ------- + out : ndarray + Output array of str or unicode, depending on input types + + Notes + ----- + This is a thin wrapper around np.strings.multiply that raises + `ValueError` when ``i`` is not an integer. It only + exists for backwards-compatibility. + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["a", "b", "c"]) + >>> np.strings.multiply(a, 3) + array(['aaa', 'bbb', 'ccc'], dtype='>> i = np.array([1, 2, 3]) + >>> np.strings.multiply(a, i) + array(['a', 'bb', 'ccc'], dtype='>> np.strings.multiply(np.array(['a']), i) + array(['a', 'aa', 'aaa'], dtype='>> a = np.array(['a', 'b', 'c', 'd', 'e', 'f']).reshape((2, 3)) + >>> np.strings.multiply(a, 3) + array([['aaa', 'bbb', 'ccc'], + ['ddd', 'eee', 'fff']], dtype='>> np.strings.multiply(a, i) + array([['a', 'bb', 'ccc'], + ['d', 'ee', 'fff']], dtype='>> import numpy as np + >>> x = np.array(["Numpy is nice!"]) + >>> np.char.partition(x, " ") + array([['Numpy', ' ', 'is nice!']], dtype='>> import numpy as np + >>> a = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> np.char.rpartition(a, 'A') + array([['aAaAa', 'A', ''], + [' a', 'A', ' '], + ['abB', 'A', 'Bba']], dtype='= 2`` and ``order='F'``, in which case `strides` + is in "Fortran order". + + Methods + ------- + astype + argsort + copy + count + decode + dump + dumps + encode + endswith + expandtabs + fill + find + flatten + getfield + index + isalnum + isalpha + isdecimal + isdigit + islower + isnumeric + isspace + istitle + isupper + item + join + ljust + lower + lstrip + nonzero + put + ravel + repeat + replace + reshape + resize + rfind + rindex + rjust + rsplit + rstrip + searchsorted + setfield + setflags + sort + split + splitlines + squeeze + startswith + strip + swapaxes + swapcase + take + title + tofile + tolist + tostring + translate + transpose + upper + view + zfill + + Parameters + ---------- + shape : tuple + Shape of the array. + itemsize : int, optional + Length of each array element, in number of characters. Default is 1. + unicode : bool, optional + Are the array elements of type unicode (True) or string (False). + Default is False. + buffer : object exposing the buffer interface or str, optional + Memory address of the start of the array data. Default is None, + in which case a new array is created. + offset : int, optional + Fixed stride displacement from the beginning of an axis? + Default is 0. Needs to be >=0. + strides : array_like of ints, optional + Strides for the array (see `~numpy.ndarray.strides` for + full description). Default is None. + order : {'C', 'F'}, optional + The order in which the array data is stored in memory: 'C' -> + "row major" order (the default), 'F' -> "column major" + (Fortran) order. + + Examples + -------- + >>> import numpy as np + >>> charar = np.char.chararray((3, 3)) + >>> charar[:] = 'a' + >>> charar + chararray([[b'a', b'a', b'a'], + [b'a', b'a', b'a'], + [b'a', b'a', b'a']], dtype='|S1') + + >>> charar = np.char.chararray(charar.shape, itemsize=5) + >>> charar[:] = 'abc' + >>> charar + chararray([[b'abc', b'abc', b'abc'], + [b'abc', b'abc', b'abc'], + [b'abc', b'abc', b'abc']], dtype='|S5') + + """ + def __new__(subtype, shape, itemsize=1, unicode=False, buffer=None, + offset=0, strides=None, order='C'): + if unicode: + dtype = str_ + else: + dtype = bytes_ + + # force itemsize to be a Python int, since using NumPy integer + # types results in itemsize.itemsize being used as the size of + # strings in the new array. + itemsize = int(itemsize) + + if isinstance(buffer, str): + # unicode objects do not have the buffer interface + filler = buffer + buffer = None + else: + filler = None + + if buffer is None: + self = ndarray.__new__(subtype, shape, (dtype, itemsize), + order=order) + else: + self = ndarray.__new__(subtype, shape, (dtype, itemsize), + buffer=buffer, + offset=offset, strides=strides, + order=order) + if filler is not None: + self[...] = filler + + return self + + def __array_wrap__(self, arr, context=None, return_scalar=False): + # When calling a ufunc (and some other functions), we return a + # chararray if the ufunc output is a string-like array, + # or an ndarray otherwise + if arr.dtype.char in "SUbc": + return arr.view(type(self)) + return arr + + def __array_finalize__(self, obj): + # The b is a special case because it is used for reconstructing. + if self.dtype.char not in 'VSUbc': + raise ValueError("Can only create a chararray from string data.") + + def __getitem__(self, obj): + val = ndarray.__getitem__(self, obj) + if isinstance(val, character): + return val.rstrip() + return val + + # IMPLEMENTATION NOTE: Most of the methods of this class are + # direct delegations to the free functions in this module. + # However, those that return an array of strings should instead + # return a chararray, so some extra wrapping is required. + + def __eq__(self, other): + """ + Return (self == other) element-wise. + + See Also + -------- + equal + """ + return equal(self, other) + + def __ne__(self, other): + """ + Return (self != other) element-wise. + + See Also + -------- + not_equal + """ + return not_equal(self, other) + + def __ge__(self, other): + """ + Return (self >= other) element-wise. + + See Also + -------- + greater_equal + """ + return greater_equal(self, other) + + def __le__(self, other): + """ + Return (self <= other) element-wise. + + See Also + -------- + less_equal + """ + return less_equal(self, other) + + def __gt__(self, other): + """ + Return (self > other) element-wise. + + See Also + -------- + greater + """ + return greater(self, other) + + def __lt__(self, other): + """ + Return (self < other) element-wise. + + See Also + -------- + less + """ + return less(self, other) + + def __add__(self, other): + """ + Return (self + other), that is string concatenation, + element-wise for a pair of array_likes of str or unicode. + + See Also + -------- + add + """ + return add(self, other) + + def __radd__(self, other): + """ + Return (other + self), that is string concatenation, + element-wise for a pair of array_likes of `bytes_` or `str_`. + + See Also + -------- + add + """ + return add(other, self) + + def __mul__(self, i): + """ + Return (self * i), that is string multiple concatenation, + element-wise. + + See Also + -------- + multiply + """ + return asarray(multiply(self, i)) + + def __rmul__(self, i): + """ + Return (self * i), that is string multiple concatenation, + element-wise. + + See Also + -------- + multiply + """ + return asarray(multiply(self, i)) + + def __mod__(self, i): + """ + Return (self % i), that is pre-Python 2.6 string formatting + (interpolation), element-wise for a pair of array_likes of `bytes_` + or `str_`. + + See Also + -------- + mod + """ + return asarray(mod(self, i)) + + def __rmod__(self, other): + return NotImplemented + + def argsort(self, axis=-1, kind=None, order=None): + """ + Return the indices that sort the array lexicographically. + + For full documentation see `numpy.argsort`, for which this method is + in fact merely a "thin wrapper." + + Examples + -------- + >>> c = np.array(['a1b c', '1b ca', 'b ca1', 'Ca1b'], 'S5') + >>> c = c.view(np.char.chararray); c + chararray(['a1b c', '1b ca', 'b ca1', 'Ca1b'], + dtype='|S5') + >>> c[c.argsort()] + chararray(['1b ca', 'Ca1b', 'a1b c', 'b ca1'], + dtype='|S5') + + """ + return self.__array__().argsort(axis, kind, order) + argsort.__doc__ = ndarray.argsort.__doc__ + + def capitalize(self): + """ + Return a copy of `self` with only the first character of each element + capitalized. + + See Also + -------- + char.capitalize + + """ + return asarray(capitalize(self)) + + def center(self, width, fillchar=' '): + """ + Return a copy of `self` with its elements centered in a + string of length `width`. + + See Also + -------- + center + """ + return asarray(center(self, width, fillchar)) + + def count(self, sub, start=0, end=None): + """ + Returns an array with the number of non-overlapping occurrences of + substring `sub` in the range [`start`, `end`]. + + See Also + -------- + char.count + + """ + return count(self, sub, start, end) + + def decode(self, encoding=None, errors=None): + """ + Calls ``bytes.decode`` element-wise. + + See Also + -------- + char.decode + + """ + return decode(self, encoding, errors) + + def encode(self, encoding=None, errors=None): + """ + Calls :meth:`str.encode` element-wise. + + See Also + -------- + char.encode + + """ + return encode(self, encoding, errors) + + def endswith(self, suffix, start=0, end=None): + """ + Returns a boolean array which is `True` where the string element + in `self` ends with `suffix`, otherwise `False`. + + See Also + -------- + char.endswith + + """ + return endswith(self, suffix, start, end) + + def expandtabs(self, tabsize=8): + """ + Return a copy of each string element where all tab characters are + replaced by one or more spaces. + + See Also + -------- + char.expandtabs + + """ + return asarray(expandtabs(self, tabsize)) + + def find(self, sub, start=0, end=None): + """ + For each element, return the lowest index in the string where + substring `sub` is found. + + See Also + -------- + char.find + + """ + return find(self, sub, start, end) + + def index(self, sub, start=0, end=None): + """ + Like `find`, but raises :exc:`ValueError` when the substring is not + found. + + See Also + -------- + char.index + + """ + return index(self, sub, start, end) + + def isalnum(self): + """ + Returns true for each element if all characters in the string + are alphanumeric and there is at least one character, false + otherwise. + + See Also + -------- + char.isalnum + + """ + return isalnum(self) + + def isalpha(self): + """ + Returns true for each element if all characters in the string + are alphabetic and there is at least one character, false + otherwise. + + See Also + -------- + char.isalpha + + """ + return isalpha(self) + + def isdigit(self): + """ + Returns true for each element if all characters in the string are + digits and there is at least one character, false otherwise. + + See Also + -------- + char.isdigit + + """ + return isdigit(self) + + def islower(self): + """ + Returns true for each element if all cased characters in the + string are lowercase and there is at least one cased character, + false otherwise. + + See Also + -------- + char.islower + + """ + return islower(self) + + def isspace(self): + """ + Returns true for each element if there are only whitespace + characters in the string and there is at least one character, + false otherwise. + + See Also + -------- + char.isspace + + """ + return isspace(self) + + def istitle(self): + """ + Returns true for each element if the element is a titlecased + string and there is at least one character, false otherwise. + + See Also + -------- + char.istitle + + """ + return istitle(self) + + def isupper(self): + """ + Returns true for each element if all cased characters in the + string are uppercase and there is at least one character, false + otherwise. + + See Also + -------- + char.isupper + + """ + return isupper(self) + + def join(self, seq): + """ + Return a string which is the concatenation of the strings in the + sequence `seq`. + + See Also + -------- + char.join + + """ + return join(self, seq) + + def ljust(self, width, fillchar=' '): + """ + Return an array with the elements of `self` left-justified in a + string of length `width`. + + See Also + -------- + char.ljust + + """ + return asarray(ljust(self, width, fillchar)) + + def lower(self): + """ + Return an array with the elements of `self` converted to + lowercase. + + See Also + -------- + char.lower + + """ + return asarray(lower(self)) + + def lstrip(self, chars=None): + """ + For each element in `self`, return a copy with the leading characters + removed. + + See Also + -------- + char.lstrip + + """ + return lstrip(self, chars) + + def partition(self, sep): + """ + Partition each element in `self` around `sep`. + + See Also + -------- + partition + """ + return asarray(partition(self, sep)) + + def replace(self, old, new, count=None): + """ + For each element in `self`, return a copy of the string with all + occurrences of substring `old` replaced by `new`. + + See Also + -------- + char.replace + + """ + return replace(self, old, new, count if count is not None else -1) + + def rfind(self, sub, start=0, end=None): + """ + For each element in `self`, return the highest index in the string + where substring `sub` is found, such that `sub` is contained + within [`start`, `end`]. + + See Also + -------- + char.rfind + + """ + return rfind(self, sub, start, end) + + def rindex(self, sub, start=0, end=None): + """ + Like `rfind`, but raises :exc:`ValueError` when the substring `sub` is + not found. + + See Also + -------- + char.rindex + + """ + return rindex(self, sub, start, end) + + def rjust(self, width, fillchar=' '): + """ + Return an array with the elements of `self` + right-justified in a string of length `width`. + + See Also + -------- + char.rjust + + """ + return asarray(rjust(self, width, fillchar)) + + def rpartition(self, sep): + """ + Partition each element in `self` around `sep`. + + See Also + -------- + rpartition + """ + return asarray(rpartition(self, sep)) + + def rsplit(self, sep=None, maxsplit=None): + """ + For each element in `self`, return a list of the words in + the string, using `sep` as the delimiter string. + + See Also + -------- + char.rsplit + + """ + return rsplit(self, sep, maxsplit) + + def rstrip(self, chars=None): + """ + For each element in `self`, return a copy with the trailing + characters removed. + + See Also + -------- + char.rstrip + + """ + return rstrip(self, chars) + + def split(self, sep=None, maxsplit=None): + """ + For each element in `self`, return a list of the words in the + string, using `sep` as the delimiter string. + + See Also + -------- + char.split + + """ + return split(self, sep, maxsplit) + + def splitlines(self, keepends=None): + """ + For each element in `self`, return a list of the lines in the + element, breaking at line boundaries. + + See Also + -------- + char.splitlines + + """ + return splitlines(self, keepends) + + def startswith(self, prefix, start=0, end=None): + """ + Returns a boolean array which is `True` where the string element + in `self` starts with `prefix`, otherwise `False`. + + See Also + -------- + char.startswith + + """ + return startswith(self, prefix, start, end) + + def strip(self, chars=None): + """ + For each element in `self`, return a copy with the leading and + trailing characters removed. + + See Also + -------- + char.strip + + """ + return strip(self, chars) + + def swapcase(self): + """ + For each element in `self`, return a copy of the string with + uppercase characters converted to lowercase and vice versa. + + See Also + -------- + char.swapcase + + """ + return asarray(swapcase(self)) + + def title(self): + """ + For each element in `self`, return a titlecased version of the + string: words start with uppercase characters, all remaining cased + characters are lowercase. + + See Also + -------- + char.title + + """ + return asarray(title(self)) + + def translate(self, table, deletechars=None): + """ + For each element in `self`, return a copy of the string where + all characters occurring in the optional argument + `deletechars` are removed, and the remaining characters have + been mapped through the given translation table. + + See Also + -------- + char.translate + + """ + return asarray(translate(self, table, deletechars)) + + def upper(self): + """ + Return an array with the elements of `self` converted to + uppercase. + + See Also + -------- + char.upper + + """ + return asarray(upper(self)) + + def zfill(self, width): + """ + Return the numeric string left-filled with zeros in a string of + length `width`. + + See Also + -------- + char.zfill + + """ + return asarray(zfill(self, width)) + + def isnumeric(self): + """ + For each element in `self`, return True if there are only + numeric characters in the element. + + See Also + -------- + char.isnumeric + + """ + return isnumeric(self) + + def isdecimal(self): + """ + For each element in `self`, return True if there are only + decimal characters in the element. + + See Also + -------- + char.isdecimal + + """ + return isdecimal(self) + + +@set_module("numpy.char") +def array(obj, itemsize=None, copy=True, unicode=None, order=None): + """ + Create a `~numpy.char.chararray`. + + .. note:: + This class is provided for numarray backward-compatibility. + New code (not concerned with numarray compatibility) should use + arrays of type `bytes_` or `str_` and use the free functions + in :mod:`numpy.char` for fast vectorized string operations instead. + + Versus a NumPy array of dtype `bytes_` or `str_`, this + class adds the following functionality: + + 1) values automatically have whitespace removed from the end + when indexed + + 2) comparison operators automatically remove whitespace from the + end when comparing values + + 3) vectorized string operations are provided as methods + (e.g. `chararray.endswith `) + and infix operators (e.g. ``+, *, %``) + + Parameters + ---------- + obj : array of str or unicode-like + + itemsize : int, optional + `itemsize` is the number of characters per scalar in the + resulting array. If `itemsize` is None, and `obj` is an + object array or a Python list, the `itemsize` will be + automatically determined. If `itemsize` is provided and `obj` + is of type str or unicode, then the `obj` string will be + chunked into `itemsize` pieces. + + copy : bool, optional + If true (default), then the object is copied. Otherwise, a copy + will only be made if ``__array__`` returns a copy, if obj is a + nested sequence, or if a copy is needed to satisfy any of the other + requirements (`itemsize`, unicode, `order`, etc.). + + unicode : bool, optional + When true, the resulting `~numpy.char.chararray` can contain Unicode + characters, when false only 8-bit characters. If unicode is + None and `obj` is one of the following: + + - a `~numpy.char.chararray`, + - an ndarray of type :class:`str_` or :class:`bytes_` + - a Python :class:`str` or :class:`bytes` object, + + then the unicode setting of the output array will be + automatically determined. + + order : {'C', 'F', 'A'}, optional + Specify the order of the array. If order is 'C' (default), then the + array will be in C-contiguous order (last-index varies the + fastest). If order is 'F', then the returned array + will be in Fortran-contiguous order (first-index varies the + fastest). If order is 'A', then the returned array may + be in any order (either C-, Fortran-contiguous, or even + discontiguous). + + Examples + -------- + + >>> import numpy as np + >>> char_array = np.char.array(['hello', 'world', 'numpy','array']) + >>> char_array + chararray(['hello', 'world', 'numpy', 'array'], dtype='`) + and infix operators (e.g. ``+``, ``*``, ``%``) + + Parameters + ---------- + obj : array of str or unicode-like + + itemsize : int, optional + `itemsize` is the number of characters per scalar in the + resulting array. If `itemsize` is None, and `obj` is an + object array or a Python list, the `itemsize` will be + automatically determined. If `itemsize` is provided and `obj` + is of type str or unicode, then the `obj` string will be + chunked into `itemsize` pieces. + + unicode : bool, optional + When true, the resulting `~numpy.char.chararray` can contain Unicode + characters, when false only 8-bit characters. If unicode is + None and `obj` is one of the following: + + - a `~numpy.char.chararray`, + - an ndarray of type `str_` or `unicode_` + - a Python str or unicode object, + + then the unicode setting of the output array will be + automatically determined. + + order : {'C', 'F'}, optional + Specify the order of the array. If order is 'C' (default), then the + array will be in C-contiguous order (last-index varies the + fastest). If order is 'F', then the returned array + will be in Fortran-contiguous order (first-index varies the + fastest). + + Examples + -------- + >>> import numpy as np + >>> np.char.asarray(['hello', 'world']) + chararray(['hello', 'world'], dtype=' chararray[_Shape, dtype[bytes_]]: ... + @overload + def __new__( + subtype, + shape: _ShapeLike, + itemsize: SupportsIndex | SupportsInt = ..., + unicode: L[True] = ..., + buffer: _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> chararray[_Shape, dtype[str_]]: ... + + def __array_finalize__(self, obj: object) -> None: ... + def __mul__(self, other: i_co) -> chararray[_Shape, _CharDType_co]: ... + def __rmul__(self, other: i_co) -> chararray[_Shape, _CharDType_co]: ... + def __mod__(self, i: Any) -> chararray[_Shape, _CharDType_co]: ... + + @overload + def __eq__( + self: _CharArray[str_], + other: U_co, + ) -> NDArray[np.bool]: ... + @overload + def __eq__( + self: _CharArray[bytes_], + other: S_co, + ) -> NDArray[np.bool]: ... + + @overload + def __ne__( + self: _CharArray[str_], + other: U_co, + ) -> NDArray[np.bool]: ... + @overload + def __ne__( + self: _CharArray[bytes_], + other: S_co, + ) -> NDArray[np.bool]: ... + + @overload + def __ge__( + self: _CharArray[str_], + other: U_co, + ) -> NDArray[np.bool]: ... + @overload + def __ge__( + self: _CharArray[bytes_], + other: S_co, + ) -> NDArray[np.bool]: ... + + @overload + def __le__( + self: _CharArray[str_], + other: U_co, + ) -> NDArray[np.bool]: ... + @overload + def __le__( + self: _CharArray[bytes_], + other: S_co, + ) -> NDArray[np.bool]: ... + + @overload + def __gt__( + self: _CharArray[str_], + other: U_co, + ) -> NDArray[np.bool]: ... + @overload + def __gt__( + self: _CharArray[bytes_], + other: S_co, + ) -> NDArray[np.bool]: ... + + @overload + def __lt__( + self: _CharArray[str_], + other: U_co, + ) -> NDArray[np.bool]: ... + @overload + def __lt__( + self: _CharArray[bytes_], + other: S_co, + ) -> NDArray[np.bool]: ... + + @overload + def __add__( + self: _CharArray[str_], + other: U_co, + ) -> _CharArray[str_]: ... + @overload + def __add__( + self: _CharArray[bytes_], + other: S_co, + ) -> _CharArray[bytes_]: ... + + @overload + def __radd__( + self: _CharArray[str_], + other: U_co, + ) -> _CharArray[str_]: ... + @overload + def __radd__( + self: _CharArray[bytes_], + other: S_co, + ) -> _CharArray[bytes_]: ... + + @overload + def center( + self: _CharArray[str_], + width: i_co, + fillchar: U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def center( + self: _CharArray[bytes_], + width: i_co, + fillchar: S_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def count( + self: _CharArray[str_], + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + @overload + def count( + self: _CharArray[bytes_], + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + + def decode( + self: _CharArray[bytes_], + encoding: None | str = ..., + errors: None | str = ..., + ) -> _CharArray[str_]: ... + + def encode( + self: _CharArray[str_], + encoding: None | str = ..., + errors: None | str = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def endswith( + self: _CharArray[str_], + suffix: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[np.bool]: ... + @overload + def endswith( + self: _CharArray[bytes_], + suffix: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[np.bool]: ... + + def expandtabs( + self, + tabsize: i_co = ..., + ) -> chararray[_Shape, _CharDType_co]: ... + + @overload + def find( + self: _CharArray[str_], + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + @overload + def find( + self: _CharArray[bytes_], + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + + @overload + def index( + self: _CharArray[str_], + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + @overload + def index( + self: _CharArray[bytes_], + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + + @overload + def join( + self: _CharArray[str_], + seq: U_co, + ) -> _CharArray[str_]: ... + @overload + def join( + self: _CharArray[bytes_], + seq: S_co, + ) -> _CharArray[bytes_]: ... + + @overload + def ljust( + self: _CharArray[str_], + width: i_co, + fillchar: U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def ljust( + self: _CharArray[bytes_], + width: i_co, + fillchar: S_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def lstrip( + self: _CharArray[str_], + chars: None | U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def lstrip( + self: _CharArray[bytes_], + chars: None | S_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def partition( + self: _CharArray[str_], + sep: U_co, + ) -> _CharArray[str_]: ... + @overload + def partition( + self: _CharArray[bytes_], + sep: S_co, + ) -> _CharArray[bytes_]: ... + + @overload + def replace( + self: _CharArray[str_], + old: U_co, + new: U_co, + count: None | i_co = ..., + ) -> _CharArray[str_]: ... + @overload + def replace( + self: _CharArray[bytes_], + old: S_co, + new: S_co, + count: None | i_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def rfind( + self: _CharArray[str_], + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + @overload + def rfind( + self: _CharArray[bytes_], + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + + @overload + def rindex( + self: _CharArray[str_], + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + @overload + def rindex( + self: _CharArray[bytes_], + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[int_]: ... + + @overload + def rjust( + self: _CharArray[str_], + width: i_co, + fillchar: U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def rjust( + self: _CharArray[bytes_], + width: i_co, + fillchar: S_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def rpartition( + self: _CharArray[str_], + sep: U_co, + ) -> _CharArray[str_]: ... + @overload + def rpartition( + self: _CharArray[bytes_], + sep: S_co, + ) -> _CharArray[bytes_]: ... + + @overload + def rsplit( + self: _CharArray[str_], + sep: None | U_co = ..., + maxsplit: None | i_co = ..., + ) -> NDArray[object_]: ... + @overload + def rsplit( + self: _CharArray[bytes_], + sep: None | S_co = ..., + maxsplit: None | i_co = ..., + ) -> NDArray[object_]: ... + + @overload + def rstrip( + self: _CharArray[str_], + chars: None | U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def rstrip( + self: _CharArray[bytes_], + chars: None | S_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def split( + self: _CharArray[str_], + sep: None | U_co = ..., + maxsplit: None | i_co = ..., + ) -> NDArray[object_]: ... + @overload + def split( + self: _CharArray[bytes_], + sep: None | S_co = ..., + maxsplit: None | i_co = ..., + ) -> NDArray[object_]: ... + + def splitlines(self, keepends: None | b_co = ...) -> NDArray[object_]: ... + + @overload + def startswith( + self: _CharArray[str_], + prefix: U_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[np.bool]: ... + @overload + def startswith( + self: _CharArray[bytes_], + prefix: S_co, + start: i_co = ..., + end: None | i_co = ..., + ) -> NDArray[np.bool]: ... + + @overload + def strip( + self: _CharArray[str_], + chars: None | U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def strip( + self: _CharArray[bytes_], + chars: None | S_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def translate( + self: _CharArray[str_], + table: U_co, + deletechars: None | U_co = ..., + ) -> _CharArray[str_]: ... + @overload + def translate( + self: _CharArray[bytes_], + table: S_co, + deletechars: None | S_co = ..., + ) -> _CharArray[bytes_]: ... + + def zfill(self, width: i_co) -> chararray[_Shape, _CharDType_co]: ... + def capitalize(self) -> chararray[_ShapeT_co, _CharDType_co]: ... + def title(self) -> chararray[_ShapeT_co, _CharDType_co]: ... + def swapcase(self) -> chararray[_ShapeT_co, _CharDType_co]: ... + def lower(self) -> chararray[_ShapeT_co, _CharDType_co]: ... + def upper(self) -> chararray[_ShapeT_co, _CharDType_co]: ... + def isalnum(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def isalpha(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def isdigit(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def islower(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def isspace(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def istitle(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def isupper(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def isnumeric(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + def isdecimal(self) -> ndarray[_ShapeT_co, dtype[np.bool]]: ... + + +# Comparison +@overload +def equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def not_equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def not_equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def not_equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def greater_equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def greater_equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def greater_equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def less_equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def less_equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def less_equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def greater(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def greater(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def greater(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def less(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def less(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def less(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def add(x1: U_co, x2: U_co) -> NDArray[np.str_]: ... +@overload +def add(x1: S_co, x2: S_co) -> NDArray[np.bytes_]: ... +@overload +def add(x1: _StringDTypeSupportsArray, x2: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def add(x1: T_co, T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def multiply(a: U_co, i: i_co) -> NDArray[np.str_]: ... +@overload +def multiply(a: S_co, i: i_co) -> NDArray[np.bytes_]: ... +@overload +def multiply(a: _StringDTypeSupportsArray, i: i_co) -> _StringDTypeArray: ... +@overload +def multiply(a: T_co, i: i_co) -> _StringDTypeOrUnicodeArray: ... + + +@overload +def mod(a: U_co, value: Any) -> NDArray[np.str_]: ... +@overload +def mod(a: S_co, value: Any) -> NDArray[np.bytes_]: ... +@overload +def mod(a: _StringDTypeSupportsArray, value: Any) -> _StringDTypeArray: ... +@overload +def mod(a: T_co, value: Any) -> _StringDTypeOrUnicodeArray: ... + +@overload +def capitalize(a: U_co) -> NDArray[str_]: ... +@overload +def capitalize(a: S_co) -> NDArray[bytes_]: ... +@overload +def capitalize(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def capitalize(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def center(a: U_co, width: i_co, fillchar: U_co = ...) -> NDArray[str_]: ... +@overload +def center(a: S_co, width: i_co, fillchar: S_co = ...) -> NDArray[bytes_]: ... +@overload +def center(a: _StringDTypeSupportsArray, width: i_co, fillchar: _StringDTypeSupportsArray = ...) -> _StringDTypeArray: ... +@overload +def center(a: T_co, width: i_co, fillchar: T_co = ...) -> _StringDTypeOrUnicodeArray: ... + +def decode( + a: S_co, + encoding: None | str = ..., + errors: None | str = ..., +) -> NDArray[str_]: ... +def encode( + a: U_co | T_co, + encoding: None | str = ..., + errors: None | str = ..., +) -> NDArray[bytes_]: ... + +@overload +def expandtabs(a: U_co, tabsize: i_co = ...) -> NDArray[str_]: ... +@overload +def expandtabs(a: S_co, tabsize: i_co = ...) -> NDArray[bytes_]: ... +@overload +def expandtabs(a: _StringDTypeSupportsArray, tabsize: i_co = ...) -> _StringDTypeArray: ... +@overload +def expandtabs(a: T_co, tabsize: i_co = ...) -> _StringDTypeOrUnicodeArray: ... + +@overload +def join(sep: U_co, seq: U_co) -> NDArray[str_]: ... +@overload +def join(sep: S_co, seq: S_co) -> NDArray[bytes_]: ... +@overload +def join(sep: _StringDTypeSupportsArray, seq: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def join(sep: T_co, seq: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def ljust(a: U_co, width: i_co, fillchar: U_co = ...) -> NDArray[str_]: ... +@overload +def ljust(a: S_co, width: i_co, fillchar: S_co = ...) -> NDArray[bytes_]: ... +@overload +def ljust(a: _StringDTypeSupportsArray, width: i_co, fillchar: _StringDTypeSupportsArray = ...) -> _StringDTypeArray: ... +@overload +def ljust(a: T_co, width: i_co, fillchar: T_co = ...) -> _StringDTypeOrUnicodeArray: ... + +@overload +def lower(a: U_co) -> NDArray[str_]: ... +@overload +def lower(a: S_co) -> NDArray[bytes_]: ... +@overload +def lower(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def lower(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def lstrip(a: U_co, chars: None | U_co = ...) -> NDArray[str_]: ... +@overload +def lstrip(a: S_co, chars: None | S_co = ...) -> NDArray[bytes_]: ... +@overload +def lstrip(a: _StringDTypeSupportsArray, chars: None | _StringDTypeSupportsArray = ...) -> _StringDTypeArray: ... +@overload +def lstrip(a: T_co, chars: None | T_co = ...) -> _StringDTypeOrUnicodeArray: ... + +@overload +def partition(a: U_co, sep: U_co) -> NDArray[str_]: ... +@overload +def partition(a: S_co, sep: S_co) -> NDArray[bytes_]: ... +@overload +def partition(a: _StringDTypeSupportsArray, sep: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def partition(a: T_co, sep: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def replace( + a: U_co, + old: U_co, + new: U_co, + count: None | i_co = ..., +) -> NDArray[str_]: ... +@overload +def replace( + a: S_co, + old: S_co, + new: S_co, + count: None | i_co = ..., +) -> NDArray[bytes_]: ... +@overload +def replace( + a: _StringDTypeSupportsArray, + old: _StringDTypeSupportsArray, + new: _StringDTypeSupportsArray, + count: i_co = ..., +) -> _StringDTypeArray: ... +@overload +def replace( + a: T_co, + old: T_co, + new: T_co, + count: i_co = ..., +) -> _StringDTypeOrUnicodeArray: ... + +@overload +def rjust( + a: U_co, + width: i_co, + fillchar: U_co = ..., +) -> NDArray[str_]: ... +@overload +def rjust( + a: S_co, + width: i_co, + fillchar: S_co = ..., +) -> NDArray[bytes_]: ... +@overload +def rjust( + a: _StringDTypeSupportsArray, + width: i_co, + fillchar: _StringDTypeSupportsArray = ..., +) -> _StringDTypeArray: ... +@overload +def rjust( + a: T_co, + width: i_co, + fillchar: T_co = ..., +) -> _StringDTypeOrUnicodeArray: ... + +@overload +def rpartition(a: U_co, sep: U_co) -> NDArray[str_]: ... +@overload +def rpartition(a: S_co, sep: S_co) -> NDArray[bytes_]: ... +@overload +def rpartition(a: _StringDTypeSupportsArray, sep: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def rpartition(a: T_co, sep: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def rsplit( + a: U_co, + sep: None | U_co = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... +@overload +def rsplit( + a: S_co, + sep: None | S_co = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... +@overload +def rsplit( + a: _StringDTypeSupportsArray, + sep: None | _StringDTypeSupportsArray = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... +@overload +def rsplit( + a: T_co, + sep: None | T_co = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... + +@overload +def rstrip(a: U_co, chars: None | U_co = ...) -> NDArray[str_]: ... +@overload +def rstrip(a: S_co, chars: None | S_co = ...) -> NDArray[bytes_]: ... +@overload +def rstrip(a: _StringDTypeSupportsArray, chars: None | _StringDTypeSupportsArray = ...) -> _StringDTypeArray: ... +@overload +def rstrip(a: T_co, chars: None | T_co = ...) -> _StringDTypeOrUnicodeArray: ... + +@overload +def split( + a: U_co, + sep: None | U_co = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... +@overload +def split( + a: S_co, + sep: None | S_co = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... +@overload +def split( + a: _StringDTypeSupportsArray, + sep: None | _StringDTypeSupportsArray = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... +@overload +def split( + a: T_co, + sep: None | T_co = ..., + maxsplit: None | i_co = ..., +) -> NDArray[object_]: ... + +def splitlines(a: UST_co, keepends: None | b_co = ...) -> NDArray[np.object_]: ... + +@overload +def strip(a: U_co, chars: None | U_co = ...) -> NDArray[str_]: ... +@overload +def strip(a: S_co, chars: None | S_co = ...) -> NDArray[bytes_]: ... +@overload +def strip(a: _StringDTypeSupportsArray, chars: None | _StringDTypeSupportsArray = ...) -> _StringDTypeArray: ... +@overload +def strip(a: T_co, chars: None | T_co = ...) -> _StringDTypeOrUnicodeArray: ... + +@overload +def swapcase(a: U_co) -> NDArray[str_]: ... +@overload +def swapcase(a: S_co) -> NDArray[bytes_]: ... +@overload +def swapcase(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def swapcase(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def title(a: U_co) -> NDArray[str_]: ... +@overload +def title(a: S_co) -> NDArray[bytes_]: ... +@overload +def title(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def title(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def translate( + a: U_co, + table: str, + deletechars: None | str = ..., +) -> NDArray[str_]: ... +@overload +def translate( + a: S_co, + table: str, + deletechars: None | str = ..., +) -> NDArray[bytes_]: ... +@overload +def translate( + a: _StringDTypeSupportsArray, + table: str, + deletechars: None | str = ..., +) -> _StringDTypeArray: ... +@overload +def translate( + a: T_co, + table: str, + deletechars: None | str = ..., +) -> _StringDTypeOrUnicodeArray: ... + +@overload +def upper(a: U_co) -> NDArray[str_]: ... +@overload +def upper(a: S_co) -> NDArray[bytes_]: ... +@overload +def upper(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def upper(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def zfill(a: U_co, width: i_co) -> NDArray[str_]: ... +@overload +def zfill(a: S_co, width: i_co) -> NDArray[bytes_]: ... +@overload +def zfill(a: _StringDTypeSupportsArray, width: i_co) -> _StringDTypeArray: ... +@overload +def zfill(a: T_co, width: i_co) -> _StringDTypeOrUnicodeArray: ... + +# String information +@overload +def count( + a: U_co, + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def count( + a: S_co, + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def count( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def endswith( + a: U_co, + suffix: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[np.bool]: ... +@overload +def endswith( + a: S_co, + suffix: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[np.bool]: ... +@overload +def endswith( + a: T_co, + suffix: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... + +@overload +def find( + a: U_co, + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def find( + a: S_co, + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def find( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def index( + a: U_co, + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def index( + a: S_co, + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def index( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +def isalpha(a: UST_co) -> NDArray[np.bool]: ... +def isalnum(a: UST_co) -> NDArray[np.bool]: ... +def isdecimal(a: U_co | T_co) -> NDArray[np.bool]: ... +def isdigit(a: UST_co) -> NDArray[np.bool]: ... +def islower(a: UST_co) -> NDArray[np.bool]: ... +def isnumeric(a: U_co | T_co) -> NDArray[np.bool]: ... +def isspace(a: UST_co) -> NDArray[np.bool]: ... +def istitle(a: UST_co) -> NDArray[np.bool]: ... +def isupper(a: UST_co) -> NDArray[np.bool]: ... + +@overload +def rfind( + a: U_co, + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def rfind( + a: S_co, + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def rfind( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def rindex( + a: U_co, + sub: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def rindex( + a: S_co, + sub: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[int_]: ... +@overload +def rindex( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def startswith( + a: U_co, + prefix: U_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[np.bool]: ... +@overload +def startswith( + a: S_co, + prefix: S_co, + start: i_co = ..., + end: None | i_co = ..., +) -> NDArray[np.bool]: ... +@overload +def startswith( + a: T_co, + suffix: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... + +def str_len(A: UST_co) -> NDArray[int_]: ... + +# Overload 1 and 2: str- or bytes-based array-likes +# overload 3: arbitrary object with unicode=False (-> bytes_) +# overload 4: arbitrary object with unicode=True (-> str_) +@overload +def array( + obj: U_co, + itemsize: None | int = ..., + copy: bool = ..., + unicode: L[False] = ..., + order: _OrderKACF = ..., +) -> _CharArray[str_]: ... +@overload +def array( + obj: S_co, + itemsize: None | int = ..., + copy: bool = ..., + unicode: L[False] = ..., + order: _OrderKACF = ..., +) -> _CharArray[bytes_]: ... +@overload +def array( + obj: object, + itemsize: None | int = ..., + copy: bool = ..., + unicode: L[False] = ..., + order: _OrderKACF = ..., +) -> _CharArray[bytes_]: ... +@overload +def array( + obj: object, + itemsize: None | int = ..., + copy: bool = ..., + unicode: L[True] = ..., + order: _OrderKACF = ..., +) -> _CharArray[str_]: ... + +@overload +def asarray( + obj: U_co, + itemsize: None | int = ..., + unicode: L[False] = ..., + order: _OrderKACF = ..., +) -> _CharArray[str_]: ... +@overload +def asarray( + obj: S_co, + itemsize: None | int = ..., + unicode: L[False] = ..., + order: _OrderKACF = ..., +) -> _CharArray[bytes_]: ... +@overload +def asarray( + obj: object, + itemsize: None | int = ..., + unicode: L[False] = ..., + order: _OrderKACF = ..., +) -> _CharArray[bytes_]: ... +@overload +def asarray( + obj: object, + itemsize: None | int = ..., + unicode: L[True] = ..., + order: _OrderKACF = ..., +) -> _CharArray[str_]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/einsumfunc.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/einsumfunc.py new file mode 100644 index 0000000000000000000000000000000000000000..108a1369a90259c0e0d0620e38297d7780653148 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/einsumfunc.py @@ -0,0 +1,1499 @@ +""" +Implementation of optimized einsum. + +""" +import itertools +import operator + +from numpy._core.multiarray import c_einsum +from numpy._core.numeric import asanyarray, tensordot +from numpy._core.overrides import array_function_dispatch + +__all__ = ['einsum', 'einsum_path'] + +# importing string for string.ascii_letters would be too slow +# the first import before caching has been measured to take 800 µs (#23777) +einsum_symbols = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ' +einsum_symbols_set = set(einsum_symbols) + + +def _flop_count(idx_contraction, inner, num_terms, size_dictionary): + """ + Computes the number of FLOPS in the contraction. + + Parameters + ---------- + idx_contraction : iterable + The indices involved in the contraction + inner : bool + Does this contraction require an inner product? + num_terms : int + The number of terms in a contraction + size_dictionary : dict + The size of each of the indices in idx_contraction + + Returns + ------- + flop_count : int + The total number of FLOPS required for the contraction. + + Examples + -------- + + >>> _flop_count('abc', False, 1, {'a': 2, 'b':3, 'c':5}) + 30 + + >>> _flop_count('abc', True, 2, {'a': 2, 'b':3, 'c':5}) + 60 + + """ + + overall_size = _compute_size_by_dict(idx_contraction, size_dictionary) + op_factor = max(1, num_terms - 1) + if inner: + op_factor += 1 + + return overall_size * op_factor + +def _compute_size_by_dict(indices, idx_dict): + """ + Computes the product of the elements in indices based on the dictionary + idx_dict. + + Parameters + ---------- + indices : iterable + Indices to base the product on. + idx_dict : dictionary + Dictionary of index sizes + + Returns + ------- + ret : int + The resulting product. + + Examples + -------- + >>> _compute_size_by_dict('abbc', {'a': 2, 'b':3, 'c':5}) + 90 + + """ + ret = 1 + for i in indices: + ret *= idx_dict[i] + return ret + + +def _find_contraction(positions, input_sets, output_set): + """ + Finds the contraction for a given set of input and output sets. + + Parameters + ---------- + positions : iterable + Integer positions of terms used in the contraction. + input_sets : list + List of sets that represent the lhs side of the einsum subscript + output_set : set + Set that represents the rhs side of the overall einsum subscript + + Returns + ------- + new_result : set + The indices of the resulting contraction + remaining : list + List of sets that have not been contracted, the new set is appended to + the end of this list + idx_removed : set + Indices removed from the entire contraction + idx_contraction : set + The indices used in the current contraction + + Examples + -------- + + # A simple dot product test case + >>> pos = (0, 1) + >>> isets = [set('ab'), set('bc')] + >>> oset = set('ac') + >>> _find_contraction(pos, isets, oset) + ({'a', 'c'}, [{'a', 'c'}], {'b'}, {'a', 'b', 'c'}) + + # A more complex case with additional terms in the contraction + >>> pos = (0, 2) + >>> isets = [set('abd'), set('ac'), set('bdc')] + >>> oset = set('ac') + >>> _find_contraction(pos, isets, oset) + ({'a', 'c'}, [{'a', 'c'}, {'a', 'c'}], {'b', 'd'}, {'a', 'b', 'c', 'd'}) + """ + + idx_contract = set() + idx_remain = output_set.copy() + remaining = [] + for ind, value in enumerate(input_sets): + if ind in positions: + idx_contract |= value + else: + remaining.append(value) + idx_remain |= value + + new_result = idx_remain & idx_contract + idx_removed = (idx_contract - new_result) + remaining.append(new_result) + + return (new_result, remaining, idx_removed, idx_contract) + + +def _optimal_path(input_sets, output_set, idx_dict, memory_limit): + """ + Computes all possible pair contractions, sieves the results based + on ``memory_limit`` and returns the lowest cost path. This algorithm + scales factorial with respect to the elements in the list ``input_sets``. + + Parameters + ---------- + input_sets : list + List of sets that represent the lhs side of the einsum subscript + output_set : set + Set that represents the rhs side of the overall einsum subscript + idx_dict : dictionary + Dictionary of index sizes + memory_limit : int + The maximum number of elements in a temporary array + + Returns + ------- + path : list + The optimal contraction order within the memory limit constraint. + + Examples + -------- + >>> isets = [set('abd'), set('ac'), set('bdc')] + >>> oset = set() + >>> idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4} + >>> _optimal_path(isets, oset, idx_sizes, 5000) + [(0, 2), (0, 1)] + """ + + full_results = [(0, [], input_sets)] + for iteration in range(len(input_sets) - 1): + iter_results = [] + + # Compute all unique pairs + for curr in full_results: + cost, positions, remaining = curr + for con in itertools.combinations( + range(len(input_sets) - iteration), 2 + ): + + # Find the contraction + cont = _find_contraction(con, remaining, output_set) + new_result, new_input_sets, idx_removed, idx_contract = cont + + # Sieve the results based on memory_limit + new_size = _compute_size_by_dict(new_result, idx_dict) + if new_size > memory_limit: + continue + + # Build (total_cost, positions, indices_remaining) + total_cost = cost + _flop_count( + idx_contract, idx_removed, len(con), idx_dict + ) + new_pos = positions + [con] + iter_results.append((total_cost, new_pos, new_input_sets)) + + # Update combinatorial list, if we did not find anything return best + # path + remaining contractions + if iter_results: + full_results = iter_results + else: + path = min(full_results, key=lambda x: x[0])[1] + path += [tuple(range(len(input_sets) - iteration))] + return path + + # If we have not found anything return single einsum contraction + if len(full_results) == 0: + return [tuple(range(len(input_sets)))] + + path = min(full_results, key=lambda x: x[0])[1] + return path + +def _parse_possible_contraction( + positions, input_sets, output_set, idx_dict, + memory_limit, path_cost, naive_cost + ): + """Compute the cost (removed size + flops) and resultant indices for + performing the contraction specified by ``positions``. + + Parameters + ---------- + positions : tuple of int + The locations of the proposed tensors to contract. + input_sets : list of sets + The indices found on each tensors. + output_set : set + The output indices of the expression. + idx_dict : dict + Mapping of each index to its size. + memory_limit : int + The total allowed size for an intermediary tensor. + path_cost : int + The contraction cost so far. + naive_cost : int + The cost of the unoptimized expression. + + Returns + ------- + cost : (int, int) + A tuple containing the size of any indices removed, and the flop cost. + positions : tuple of int + The locations of the proposed tensors to contract. + new_input_sets : list of sets + The resulting new list of indices if this proposed contraction + is performed. + + """ + + # Find the contraction + contract = _find_contraction(positions, input_sets, output_set) + idx_result, new_input_sets, idx_removed, idx_contract = contract + + # Sieve the results based on memory_limit + new_size = _compute_size_by_dict(idx_result, idx_dict) + if new_size > memory_limit: + return None + + # Build sort tuple + old_sizes = ( + _compute_size_by_dict(input_sets[p], idx_dict) for p in positions + ) + removed_size = sum(old_sizes) - new_size + + # NB: removed_size used to be just the size of any removed indices i.e.: + # helpers.compute_size_by_dict(idx_removed, idx_dict) + cost = _flop_count(idx_contract, idx_removed, len(positions), idx_dict) + sort = (-removed_size, cost) + + # Sieve based on total cost as well + if (path_cost + cost) > naive_cost: + return None + + # Add contraction to possible choices + return [sort, positions, new_input_sets] + + +def _update_other_results(results, best): + """Update the positions and provisional input_sets of ``results`` + based on performing the contraction result ``best``. Remove any + involving the tensors contracted. + + Parameters + ---------- + results : list + List of contraction results produced by + ``_parse_possible_contraction``. + best : list + The best contraction of ``results`` i.e. the one that + will be performed. + + Returns + ------- + mod_results : list + The list of modified results, updated with outcome of + ``best`` contraction. + """ + + best_con = best[1] + bx, by = best_con + mod_results = [] + + for cost, (x, y), con_sets in results: + + # Ignore results involving tensors just contracted + if x in best_con or y in best_con: + continue + + # Update the input_sets + del con_sets[by - int(by > x) - int(by > y)] + del con_sets[bx - int(bx > x) - int(bx > y)] + con_sets.insert(-1, best[2][-1]) + + # Update the position indices + mod_con = x - int(x > bx) - int(x > by), y - int(y > bx) - int(y > by) + mod_results.append((cost, mod_con, con_sets)) + + return mod_results + +def _greedy_path(input_sets, output_set, idx_dict, memory_limit): + """ + Finds the path by contracting the best pair until the input list is + exhausted. The best pair is found by minimizing the tuple + ``(-prod(indices_removed), cost)``. What this amounts to is prioritizing + matrix multiplication or inner product operations, then Hadamard like + operations, and finally outer operations. Outer products are limited by + ``memory_limit``. This algorithm scales cubically with respect to the + number of elements in the list ``input_sets``. + + Parameters + ---------- + input_sets : list + List of sets that represent the lhs side of the einsum subscript + output_set : set + Set that represents the rhs side of the overall einsum subscript + idx_dict : dictionary + Dictionary of index sizes + memory_limit : int + The maximum number of elements in a temporary array + + Returns + ------- + path : list + The greedy contraction order within the memory limit constraint. + + Examples + -------- + >>> isets = [set('abd'), set('ac'), set('bdc')] + >>> oset = set() + >>> idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4} + >>> _greedy_path(isets, oset, idx_sizes, 5000) + [(0, 2), (0, 1)] + """ + + # Handle trivial cases that leaked through + if len(input_sets) == 1: + return [(0,)] + elif len(input_sets) == 2: + return [(0, 1)] + + # Build up a naive cost + contract = _find_contraction( + range(len(input_sets)), input_sets, output_set + ) + idx_result, new_input_sets, idx_removed, idx_contract = contract + naive_cost = _flop_count( + idx_contract, idx_removed, len(input_sets), idx_dict + ) + + # Initially iterate over all pairs + comb_iter = itertools.combinations(range(len(input_sets)), 2) + known_contractions = [] + + path_cost = 0 + path = [] + + for iteration in range(len(input_sets) - 1): + + # Iterate over all pairs on the first step, only previously + # found pairs on subsequent steps + for positions in comb_iter: + + # Always initially ignore outer products + if input_sets[positions[0]].isdisjoint(input_sets[positions[1]]): + continue + + result = _parse_possible_contraction( + positions, input_sets, output_set, idx_dict, + memory_limit, path_cost, naive_cost + ) + if result is not None: + known_contractions.append(result) + + # If we do not have a inner contraction, rescan pairs + # including outer products + if len(known_contractions) == 0: + + # Then check the outer products + for positions in itertools.combinations( + range(len(input_sets)), 2 + ): + result = _parse_possible_contraction( + positions, input_sets, output_set, idx_dict, + memory_limit, path_cost, naive_cost + ) + if result is not None: + known_contractions.append(result) + + # If we still did not find any remaining contractions, + # default back to einsum like behavior + if len(known_contractions) == 0: + path.append(tuple(range(len(input_sets)))) + break + + # Sort based on first index + best = min(known_contractions, key=lambda x: x[0]) + + # Now propagate as many unused contractions as possible + # to the next iteration + known_contractions = _update_other_results(known_contractions, best) + + # Next iteration only compute contractions with the new tensor + # All other contractions have been accounted for + input_sets = best[2] + new_tensor_pos = len(input_sets) - 1 + comb_iter = ((i, new_tensor_pos) for i in range(new_tensor_pos)) + + # Update path and total cost + path.append(best[1]) + path_cost += best[0][1] + + return path + + +def _can_dot(inputs, result, idx_removed): + """ + Checks if we can use BLAS (np.tensordot) call and its beneficial to do so. + + Parameters + ---------- + inputs : list of str + Specifies the subscripts for summation. + result : str + Resulting summation. + idx_removed : set + Indices that are removed in the summation + + + Returns + ------- + type : bool + Returns true if BLAS should and can be used, else False + + Notes + ----- + If the operations is BLAS level 1 or 2 and is not already aligned + we default back to einsum as the memory movement to copy is more + costly than the operation itself. + + + Examples + -------- + + # Standard GEMM operation + >>> _can_dot(['ij', 'jk'], 'ik', set('j')) + True + + # Can use the standard BLAS, but requires odd data movement + >>> _can_dot(['ijj', 'jk'], 'ik', set('j')) + False + + # DDOT where the memory is not aligned + >>> _can_dot(['ijk', 'ikj'], '', set('ijk')) + False + + """ + + # All `dot` calls remove indices + if len(idx_removed) == 0: + return False + + # BLAS can only handle two operands + if len(inputs) != 2: + return False + + input_left, input_right = inputs + + for c in set(input_left + input_right): + # can't deal with repeated indices on same input or more than 2 total + nl, nr = input_left.count(c), input_right.count(c) + if (nl > 1) or (nr > 1) or (nl + nr > 2): + return False + + # can't do implicit summation or dimension collapse e.g. + # "ab,bc->c" (implicitly sum over 'a') + # "ab,ca->ca" (take diagonal of 'a') + if nl + nr - 1 == int(c in result): + return False + + # Build a few temporaries + set_left = set(input_left) + set_right = set(input_right) + keep_left = set_left - idx_removed + keep_right = set_right - idx_removed + rs = len(idx_removed) + + # At this point we are a DOT, GEMV, or GEMM operation + + # Handle inner products + + # DDOT with aligned data + if input_left == input_right: + return True + + # DDOT without aligned data (better to use einsum) + if set_left == set_right: + return False + + # Handle the 4 possible (aligned) GEMV or GEMM cases + + # GEMM or GEMV no transpose + if input_left[-rs:] == input_right[:rs]: + return True + + # GEMM or GEMV transpose both + if input_left[:rs] == input_right[-rs:]: + return True + + # GEMM or GEMV transpose right + if input_left[-rs:] == input_right[-rs:]: + return True + + # GEMM or GEMV transpose left + if input_left[:rs] == input_right[:rs]: + return True + + # Einsum is faster than GEMV if we have to copy data + if not keep_left or not keep_right: + return False + + # We are a matrix-matrix product, but we need to copy data + return True + + +def _parse_einsum_input(operands): + """ + A reproduction of einsum c side einsum parsing in python. + + Returns + ------- + input_strings : str + Parsed input strings + output_string : str + Parsed output string + operands : list of array_like + The operands to use in the numpy contraction + + Examples + -------- + The operand list is simplified to reduce printing: + + >>> np.random.seed(123) + >>> a = np.random.rand(4, 4) + >>> b = np.random.rand(4, 4, 4) + >>> _parse_einsum_input(('...a,...a->...', a, b)) + ('za,xza', 'xz', [a, b]) # may vary + + >>> _parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0])) + ('za,xza', 'xz', [a, b]) # may vary + """ + + if len(operands) == 0: + raise ValueError("No input operands") + + if isinstance(operands[0], str): + subscripts = operands[0].replace(" ", "") + operands = [asanyarray(v) for v in operands[1:]] + + # Ensure all characters are valid + for s in subscripts: + if s in '.,->': + continue + if s not in einsum_symbols: + raise ValueError("Character %s is not a valid symbol." % s) + + else: + tmp_operands = list(operands) + operand_list = [] + subscript_list = [] + for p in range(len(operands) // 2): + operand_list.append(tmp_operands.pop(0)) + subscript_list.append(tmp_operands.pop(0)) + + output_list = tmp_operands[-1] if len(tmp_operands) else None + operands = [asanyarray(v) for v in operand_list] + subscripts = "" + last = len(subscript_list) - 1 + for num, sub in enumerate(subscript_list): + for s in sub: + if s is Ellipsis: + subscripts += "..." + else: + try: + s = operator.index(s) + except TypeError as e: + raise TypeError( + "For this input type lists must contain " + "either int or Ellipsis" + ) from e + subscripts += einsum_symbols[s] + if num != last: + subscripts += "," + + if output_list is not None: + subscripts += "->" + for s in output_list: + if s is Ellipsis: + subscripts += "..." + else: + try: + s = operator.index(s) + except TypeError as e: + raise TypeError( + "For this input type lists must contain " + "either int or Ellipsis" + ) from e + subscripts += einsum_symbols[s] + # Check for proper "->" + if ("-" in subscripts) or (">" in subscripts): + invalid = (subscripts.count("-") > 1) or (subscripts.count(">") > 1) + if invalid or (subscripts.count("->") != 1): + raise ValueError("Subscripts can only contain one '->'.") + + # Parse ellipses + if "." in subscripts: + used = subscripts.replace(".", "").replace(",", "").replace("->", "") + unused = list(einsum_symbols_set - set(used)) + ellipse_inds = "".join(unused) + longest = 0 + + if "->" in subscripts: + input_tmp, output_sub = subscripts.split("->") + split_subscripts = input_tmp.split(",") + out_sub = True + else: + split_subscripts = subscripts.split(',') + out_sub = False + + for num, sub in enumerate(split_subscripts): + if "." in sub: + if (sub.count(".") != 3) or (sub.count("...") != 1): + raise ValueError("Invalid Ellipses.") + + # Take into account numerical values + if operands[num].shape == (): + ellipse_count = 0 + else: + ellipse_count = max(operands[num].ndim, 1) + ellipse_count -= (len(sub) - 3) + + if ellipse_count > longest: + longest = ellipse_count + + if ellipse_count < 0: + raise ValueError("Ellipses lengths do not match.") + elif ellipse_count == 0: + split_subscripts[num] = sub.replace('...', '') + else: + rep_inds = ellipse_inds[-ellipse_count:] + split_subscripts[num] = sub.replace('...', rep_inds) + + subscripts = ",".join(split_subscripts) + if longest == 0: + out_ellipse = "" + else: + out_ellipse = ellipse_inds[-longest:] + + if out_sub: + subscripts += "->" + output_sub.replace("...", out_ellipse) + else: + # Special care for outputless ellipses + output_subscript = "" + tmp_subscripts = subscripts.replace(",", "") + for s in sorted(set(tmp_subscripts)): + if s not in (einsum_symbols): + raise ValueError("Character %s is not a valid symbol." % s) + if tmp_subscripts.count(s) == 1: + output_subscript += s + normal_inds = ''.join(sorted(set(output_subscript) - + set(out_ellipse))) + + subscripts += "->" + out_ellipse + normal_inds + + # Build output string if does not exist + if "->" in subscripts: + input_subscripts, output_subscript = subscripts.split("->") + else: + input_subscripts = subscripts + # Build output subscripts + tmp_subscripts = subscripts.replace(",", "") + output_subscript = "" + for s in sorted(set(tmp_subscripts)): + if s not in einsum_symbols: + raise ValueError("Character %s is not a valid symbol." % s) + if tmp_subscripts.count(s) == 1: + output_subscript += s + + # Make sure output subscripts are in the input + for char in output_subscript: + if output_subscript.count(char) != 1: + raise ValueError("Output character %s appeared more than once in " + "the output." % char) + if char not in input_subscripts: + raise ValueError("Output character %s did not appear in the input" + % char) + + # Make sure number operands is equivalent to the number of terms + if len(input_subscripts.split(',')) != len(operands): + raise ValueError("Number of einsum subscripts must be equal to the " + "number of operands.") + + return (input_subscripts, output_subscript, operands) + + +def _einsum_path_dispatcher(*operands, optimize=None, einsum_call=None): + # NOTE: technically, we should only dispatch on array-like arguments, not + # subscripts (given as strings). But separating operands into + # arrays/subscripts is a little tricky/slow (given einsum's two supported + # signatures), so as a practical shortcut we dispatch on everything. + # Strings will be ignored for dispatching since they don't define + # __array_function__. + return operands + + +@array_function_dispatch(_einsum_path_dispatcher, module='numpy') +def einsum_path(*operands, optimize='greedy', einsum_call=False): + """ + einsum_path(subscripts, *operands, optimize='greedy') + + Evaluates the lowest cost contraction order for an einsum expression by + considering the creation of intermediate arrays. + + Parameters + ---------- + subscripts : str + Specifies the subscripts for summation. + *operands : list of array_like + These are the arrays for the operation. + optimize : {bool, list, tuple, 'greedy', 'optimal'} + Choose the type of path. If a tuple is provided, the second argument is + assumed to be the maximum intermediate size created. If only a single + argument is provided the largest input or output array size is used + as a maximum intermediate size. + + * if a list is given that starts with ``einsum_path``, uses this as the + contraction path + * if False no optimization is taken + * if True defaults to the 'greedy' algorithm + * 'optimal' An algorithm that combinatorially explores all possible + ways of contracting the listed tensors and chooses the least costly + path. Scales exponentially with the number of terms in the + contraction. + * 'greedy' An algorithm that chooses the best pair contraction + at each step. Effectively, this algorithm searches the largest inner, + Hadamard, and then outer products at each step. Scales cubically with + the number of terms in the contraction. Equivalent to the 'optimal' + path for most contractions. + + Default is 'greedy'. + + Returns + ------- + path : list of tuples + A list representation of the einsum path. + string_repr : str + A printable representation of the einsum path. + + Notes + ----- + The resulting path indicates which terms of the input contraction should be + contracted first, the result of this contraction is then appended to the + end of the contraction list. This list can then be iterated over until all + intermediate contractions are complete. + + See Also + -------- + einsum, linalg.multi_dot + + Examples + -------- + + We can begin with a chain dot example. In this case, it is optimal to + contract the ``b`` and ``c`` tensors first as represented by the first + element of the path ``(1, 2)``. The resulting tensor is added to the end + of the contraction and the remaining contraction ``(0, 1)`` is then + completed. + + >>> np.random.seed(123) + >>> a = np.random.rand(2, 2) + >>> b = np.random.rand(2, 5) + >>> c = np.random.rand(5, 2) + >>> path_info = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy') + >>> print(path_info[0]) + ['einsum_path', (1, 2), (0, 1)] + >>> print(path_info[1]) + Complete contraction: ij,jk,kl->il # may vary + Naive scaling: 4 + Optimized scaling: 3 + Naive FLOP count: 1.600e+02 + Optimized FLOP count: 5.600e+01 + Theoretical speedup: 2.857 + Largest intermediate: 4.000e+00 elements + ------------------------------------------------------------------------- + scaling current remaining + ------------------------------------------------------------------------- + 3 kl,jk->jl ij,jl->il + 3 jl,ij->il il->il + + + A more complex index transformation example. + + >>> I = np.random.rand(10, 10, 10, 10) + >>> C = np.random.rand(10, 10) + >>> path_info = np.einsum_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C, + ... optimize='greedy') + + >>> print(path_info[0]) + ['einsum_path', (0, 2), (0, 3), (0, 2), (0, 1)] + >>> print(path_info[1]) + Complete contraction: ea,fb,abcd,gc,hd->efgh # may vary + Naive scaling: 8 + Optimized scaling: 5 + Naive FLOP count: 8.000e+08 + Optimized FLOP count: 8.000e+05 + Theoretical speedup: 1000.000 + Largest intermediate: 1.000e+04 elements + -------------------------------------------------------------------------- + scaling current remaining + -------------------------------------------------------------------------- + 5 abcd,ea->bcde fb,gc,hd,bcde->efgh + 5 bcde,fb->cdef gc,hd,cdef->efgh + 5 cdef,gc->defg hd,defg->efgh + 5 defg,hd->efgh efgh->efgh + """ + + # Figure out what the path really is + path_type = optimize + if path_type is True: + path_type = 'greedy' + if path_type is None: + path_type = False + + explicit_einsum_path = False + memory_limit = None + + # No optimization or a named path algorithm + if (path_type is False) or isinstance(path_type, str): + pass + + # Given an explicit path + elif len(path_type) and (path_type[0] == 'einsum_path'): + explicit_einsum_path = True + + # Path tuple with memory limit + elif ((len(path_type) == 2) and isinstance(path_type[0], str) and + isinstance(path_type[1], (int, float))): + memory_limit = int(path_type[1]) + path_type = path_type[0] + + else: + raise TypeError("Did not understand the path: %s" % str(path_type)) + + # Hidden option, only einsum should call this + einsum_call_arg = einsum_call + + # Python side parsing + input_subscripts, output_subscript, operands = ( + _parse_einsum_input(operands) + ) + + # Build a few useful list and sets + input_list = input_subscripts.split(',') + input_sets = [set(x) for x in input_list] + output_set = set(output_subscript) + indices = set(input_subscripts.replace(',', '')) + + # Get length of each unique dimension and ensure all dimensions are correct + dimension_dict = {} + broadcast_indices = [[] for x in range(len(input_list))] + for tnum, term in enumerate(input_list): + sh = operands[tnum].shape + if len(sh) != len(term): + raise ValueError("Einstein sum subscript %s does not contain the " + "correct number of indices for operand %d." + % (input_subscripts[tnum], tnum)) + for cnum, char in enumerate(term): + dim = sh[cnum] + + # Build out broadcast indices + if dim == 1: + broadcast_indices[tnum].append(char) + + if char in dimension_dict.keys(): + # For broadcasting cases we always want the largest dim size + if dimension_dict[char] == 1: + dimension_dict[char] = dim + elif dim not in (1, dimension_dict[char]): + raise ValueError("Size of label '%s' for operand %d (%d) " + "does not match previous terms (%d)." + % (char, tnum, dimension_dict[char], dim)) + else: + dimension_dict[char] = dim + + # Convert broadcast inds to sets + broadcast_indices = [set(x) for x in broadcast_indices] + + # Compute size of each input array plus the output array + size_list = [_compute_size_by_dict(term, dimension_dict) + for term in input_list + [output_subscript]] + max_size = max(size_list) + + if memory_limit is None: + memory_arg = max_size + else: + memory_arg = memory_limit + + # Compute naive cost + # This isn't quite right, need to look into exactly how einsum does this + inner_product = (sum(len(x) for x in input_sets) - len(indices)) > 0 + naive_cost = _flop_count( + indices, inner_product, len(input_list), dimension_dict + ) + + # Compute the path + if explicit_einsum_path: + path = path_type[1:] + elif ( + (path_type is False) + or (len(input_list) in [1, 2]) + or (indices == output_set) + ): + # Nothing to be optimized, leave it to einsum + path = [tuple(range(len(input_list)))] + elif path_type == "greedy": + path = _greedy_path( + input_sets, output_set, dimension_dict, memory_arg + ) + elif path_type == "optimal": + path = _optimal_path( + input_sets, output_set, dimension_dict, memory_arg + ) + else: + raise KeyError("Path name %s not found", path_type) + + cost_list, scale_list, size_list, contraction_list = [], [], [], [] + + # Build contraction tuple (positions, gemm, einsum_str, remaining) + for cnum, contract_inds in enumerate(path): + # Make sure we remove inds from right to left + contract_inds = tuple(sorted(contract_inds, reverse=True)) + + contract = _find_contraction(contract_inds, input_sets, output_set) + out_inds, input_sets, idx_removed, idx_contract = contract + + cost = _flop_count( + idx_contract, idx_removed, len(contract_inds), dimension_dict + ) + cost_list.append(cost) + scale_list.append(len(idx_contract)) + size_list.append(_compute_size_by_dict(out_inds, dimension_dict)) + + bcast = set() + tmp_inputs = [] + for x in contract_inds: + tmp_inputs.append(input_list.pop(x)) + bcast |= broadcast_indices.pop(x) + + new_bcast_inds = bcast - idx_removed + + # If we're broadcasting, nix blas + if not len(idx_removed & bcast): + do_blas = _can_dot(tmp_inputs, out_inds, idx_removed) + else: + do_blas = False + + # Last contraction + if (cnum - len(path)) == -1: + idx_result = output_subscript + else: + sort_result = [(dimension_dict[ind], ind) for ind in out_inds] + idx_result = "".join([x[1] for x in sorted(sort_result)]) + + input_list.append(idx_result) + broadcast_indices.append(new_bcast_inds) + einsum_str = ",".join(tmp_inputs) + "->" + idx_result + + contraction = ( + contract_inds, idx_removed, einsum_str, input_list[:], do_blas + ) + contraction_list.append(contraction) + + opt_cost = sum(cost_list) + 1 + + if len(input_list) != 1: + # Explicit "einsum_path" is usually trusted, but we detect this kind of + # mistake in order to prevent from returning an intermediate value. + raise RuntimeError( + "Invalid einsum_path is specified: {} more operands has to be " + "contracted.".format(len(input_list) - 1)) + + if einsum_call_arg: + return (operands, contraction_list) + + # Return the path along with a nice string representation + overall_contraction = input_subscripts + "->" + output_subscript + header = ("scaling", "current", "remaining") + + speedup = naive_cost / opt_cost + max_i = max(size_list) + + path_print = " Complete contraction: %s\n" % overall_contraction + path_print += " Naive scaling: %d\n" % len(indices) + path_print += " Optimized scaling: %d\n" % max(scale_list) + path_print += " Naive FLOP count: %.3e\n" % naive_cost + path_print += " Optimized FLOP count: %.3e\n" % opt_cost + path_print += " Theoretical speedup: %3.3f\n" % speedup + path_print += " Largest intermediate: %.3e elements\n" % max_i + path_print += "-" * 74 + "\n" + path_print += "%6s %24s %40s\n" % header + path_print += "-" * 74 + + for n, contraction in enumerate(contraction_list): + inds, idx_rm, einsum_str, remaining, blas = contraction + remaining_str = ",".join(remaining) + "->" + output_subscript + path_run = (scale_list[n], einsum_str, remaining_str) + path_print += "\n%4d %24s %40s" % path_run + + path = ['einsum_path'] + path + return (path, path_print) + + +def _einsum_dispatcher(*operands, out=None, optimize=None, **kwargs): + # Arguably we dispatch on more arguments than we really should; see note in + # _einsum_path_dispatcher for why. + yield from operands + yield out + + +# Rewrite einsum to handle different cases +@array_function_dispatch(_einsum_dispatcher, module='numpy') +def einsum(*operands, out=None, optimize=False, **kwargs): + """ + einsum(subscripts, *operands, out=None, dtype=None, order='K', + casting='safe', optimize=False) + + Evaluates the Einstein summation convention on the operands. + + Using the Einstein summation convention, many common multi-dimensional, + linear algebraic array operations can be represented in a simple fashion. + In *implicit* mode `einsum` computes these values. + + In *explicit* mode, `einsum` provides further flexibility to compute + other array operations that might not be considered classical Einstein + summation operations, by disabling, or forcing summation over specified + subscript labels. + + See the notes and examples for clarification. + + Parameters + ---------- + subscripts : str + Specifies the subscripts for summation as comma separated list of + subscript labels. An implicit (classical Einstein summation) + calculation is performed unless the explicit indicator '->' is + included as well as subscript labels of the precise output form. + operands : list of array_like + These are the arrays for the operation. + out : ndarray, optional + If provided, the calculation is done into this array. + dtype : {data-type, None}, optional + If provided, forces the calculation to use the data type specified. + Note that you may have to also give a more liberal `casting` + parameter to allow the conversions. Default is None. + order : {'C', 'F', 'A', 'K'}, optional + Controls the memory layout of the output. 'C' means it should + be C contiguous. 'F' means it should be Fortran contiguous, + 'A' means it should be 'F' if the inputs are all 'F', 'C' otherwise. + 'K' means it should be as close to the layout as the inputs as + is possible, including arbitrarily permuted axes. + Default is 'K'. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Setting this to + 'unsafe' is not recommended, as it can adversely affect accumulations. + + * 'no' means the data types should not be cast at all. + * 'equiv' means only byte-order changes are allowed. + * 'safe' means only casts which can preserve values are allowed. + * 'same_kind' means only safe casts or casts within a kind, + like float64 to float32, are allowed. + * 'unsafe' means any data conversions may be done. + + Default is 'safe'. + optimize : {False, True, 'greedy', 'optimal'}, optional + Controls if intermediate optimization should occur. No optimization + will occur if False and True will default to the 'greedy' algorithm. + Also accepts an explicit contraction list from the ``np.einsum_path`` + function. See ``np.einsum_path`` for more details. Defaults to False. + + Returns + ------- + output : ndarray + The calculation based on the Einstein summation convention. + + See Also + -------- + einsum_path, dot, inner, outer, tensordot, linalg.multi_dot + einsum: + Similar verbose interface is provided by the + `einops `_ package to cover + additional operations: transpose, reshape/flatten, repeat/tile, + squeeze/unsqueeze and reductions. + The `opt_einsum `_ + optimizes contraction order for einsum-like expressions + in backend-agnostic manner. + + Notes + ----- + The Einstein summation convention can be used to compute + many multi-dimensional, linear algebraic array operations. `einsum` + provides a succinct way of representing these. + + A non-exhaustive list of these operations, + which can be computed by `einsum`, is shown below along with examples: + + * Trace of an array, :py:func:`numpy.trace`. + * Return a diagonal, :py:func:`numpy.diag`. + * Array axis summations, :py:func:`numpy.sum`. + * Transpositions and permutations, :py:func:`numpy.transpose`. + * Matrix multiplication and dot product, :py:func:`numpy.matmul` + :py:func:`numpy.dot`. + * Vector inner and outer products, :py:func:`numpy.inner` + :py:func:`numpy.outer`. + * Broadcasting, element-wise and scalar multiplication, + :py:func:`numpy.multiply`. + * Tensor contractions, :py:func:`numpy.tensordot`. + * Chained array operations, in efficient calculation order, + :py:func:`numpy.einsum_path`. + + The subscripts string is a comma-separated list of subscript labels, + where each label refers to a dimension of the corresponding operand. + Whenever a label is repeated it is summed, so ``np.einsum('i,i', a, b)`` + is equivalent to :py:func:`np.inner(a,b) `. If a label + appears only once, it is not summed, so ``np.einsum('i', a)`` + produces a view of ``a`` with no changes. A further example + ``np.einsum('ij,jk', a, b)`` describes traditional matrix multiplication + and is equivalent to :py:func:`np.matmul(a,b) `. + Repeated subscript labels in one operand take the diagonal. + For example, ``np.einsum('ii', a)`` is equivalent to + :py:func:`np.trace(a) `. + + In *implicit mode*, the chosen subscripts are important + since the axes of the output are reordered alphabetically. This + means that ``np.einsum('ij', a)`` doesn't affect a 2D array, while + ``np.einsum('ji', a)`` takes its transpose. Additionally, + ``np.einsum('ij,jk', a, b)`` returns a matrix multiplication, while, + ``np.einsum('ij,jh', a, b)`` returns the transpose of the + multiplication since subscript 'h' precedes subscript 'i'. + + In *explicit mode* the output can be directly controlled by + specifying output subscript labels. This requires the + identifier '->' as well as the list of output subscript labels. + This feature increases the flexibility of the function since + summing can be disabled or forced when required. The call + ``np.einsum('i->', a)`` is like :py:func:`np.sum(a) ` + if ``a`` is a 1-D array, and ``np.einsum('ii->i', a)`` + is like :py:func:`np.diag(a) ` if ``a`` is a square 2-D array. + The difference is that `einsum` does not allow broadcasting by default. + Additionally ``np.einsum('ij,jh->ih', a, b)`` directly specifies the + order of the output subscript labels and therefore returns matrix + multiplication, unlike the example above in implicit mode. + + To enable and control broadcasting, use an ellipsis. Default + NumPy-style broadcasting is done by adding an ellipsis + to the left of each term, like ``np.einsum('...ii->...i', a)``. + ``np.einsum('...i->...', a)`` is like + :py:func:`np.sum(a, axis=-1) ` for array ``a`` of any shape. + To take the trace along the first and last axes, + you can do ``np.einsum('i...i', a)``, or to do a matrix-matrix + product with the left-most indices instead of rightmost, one can do + ``np.einsum('ij...,jk...->ik...', a, b)``. + + When there is only one operand, no axes are summed, and no output + parameter is provided, a view into the operand is returned instead + of a new array. Thus, taking the diagonal as ``np.einsum('ii->i', a)`` + produces a view (changed in version 1.10.0). + + `einsum` also provides an alternative way to provide the subscripts and + operands as ``einsum(op0, sublist0, op1, sublist1, ..., [sublistout])``. + If the output shape is not provided in this format `einsum` will be + calculated in implicit mode, otherwise it will be performed explicitly. + The examples below have corresponding `einsum` calls with the two + parameter methods. + + Views returned from einsum are now writeable whenever the input array + is writeable. For example, ``np.einsum('ijk...->kji...', a)`` will now + have the same effect as :py:func:`np.swapaxes(a, 0, 2) ` + and ``np.einsum('ii->i', a)`` will return a writeable view of the diagonal + of a 2D array. + + Added the ``optimize`` argument which will optimize the contraction order + of an einsum expression. For a contraction with three or more operands + this can greatly increase the computational efficiency at the cost of + a larger memory footprint during computation. + + Typically a 'greedy' algorithm is applied which empirical tests have shown + returns the optimal path in the majority of cases. In some cases 'optimal' + will return the superlative path through a more expensive, exhaustive + search. For iterative calculations it may be advisable to calculate + the optimal path once and reuse that path by supplying it as an argument. + An example is given below. + + See :py:func:`numpy.einsum_path` for more details. + + Examples + -------- + >>> a = np.arange(25).reshape(5,5) + >>> b = np.arange(5) + >>> c = np.arange(6).reshape(2,3) + + Trace of a matrix: + + >>> np.einsum('ii', a) + 60 + >>> np.einsum(a, [0,0]) + 60 + >>> np.trace(a) + 60 + + Extract the diagonal (requires explicit form): + + >>> np.einsum('ii->i', a) + array([ 0, 6, 12, 18, 24]) + >>> np.einsum(a, [0,0], [0]) + array([ 0, 6, 12, 18, 24]) + >>> np.diag(a) + array([ 0, 6, 12, 18, 24]) + + Sum over an axis (requires explicit form): + + >>> np.einsum('ij->i', a) + array([ 10, 35, 60, 85, 110]) + >>> np.einsum(a, [0,1], [0]) + array([ 10, 35, 60, 85, 110]) + >>> np.sum(a, axis=1) + array([ 10, 35, 60, 85, 110]) + + For higher dimensional arrays summing a single axis can be done + with ellipsis: + + >>> np.einsum('...j->...', a) + array([ 10, 35, 60, 85, 110]) + >>> np.einsum(a, [Ellipsis,1], [Ellipsis]) + array([ 10, 35, 60, 85, 110]) + + Compute a matrix transpose, or reorder any number of axes: + + >>> np.einsum('ji', c) + array([[0, 3], + [1, 4], + [2, 5]]) + >>> np.einsum('ij->ji', c) + array([[0, 3], + [1, 4], + [2, 5]]) + >>> np.einsum(c, [1,0]) + array([[0, 3], + [1, 4], + [2, 5]]) + >>> np.transpose(c) + array([[0, 3], + [1, 4], + [2, 5]]) + + Vector inner products: + + >>> np.einsum('i,i', b, b) + 30 + >>> np.einsum(b, [0], b, [0]) + 30 + >>> np.inner(b,b) + 30 + + Matrix vector multiplication: + + >>> np.einsum('ij,j', a, b) + array([ 30, 80, 130, 180, 230]) + >>> np.einsum(a, [0,1], b, [1]) + array([ 30, 80, 130, 180, 230]) + >>> np.dot(a, b) + array([ 30, 80, 130, 180, 230]) + >>> np.einsum('...j,j', a, b) + array([ 30, 80, 130, 180, 230]) + + Broadcasting and scalar multiplication: + + >>> np.einsum('..., ...', 3, c) + array([[ 0, 3, 6], + [ 9, 12, 15]]) + >>> np.einsum(',ij', 3, c) + array([[ 0, 3, 6], + [ 9, 12, 15]]) + >>> np.einsum(3, [Ellipsis], c, [Ellipsis]) + array([[ 0, 3, 6], + [ 9, 12, 15]]) + >>> np.multiply(3, c) + array([[ 0, 3, 6], + [ 9, 12, 15]]) + + Vector outer product: + + >>> np.einsum('i,j', np.arange(2)+1, b) + array([[0, 1, 2, 3, 4], + [0, 2, 4, 6, 8]]) + >>> np.einsum(np.arange(2)+1, [0], b, [1]) + array([[0, 1, 2, 3, 4], + [0, 2, 4, 6, 8]]) + >>> np.outer(np.arange(2)+1, b) + array([[0, 1, 2, 3, 4], + [0, 2, 4, 6, 8]]) + + Tensor contraction: + + >>> a = np.arange(60.).reshape(3,4,5) + >>> b = np.arange(24.).reshape(4,3,2) + >>> np.einsum('ijk,jil->kl', a, b) + array([[4400., 4730.], + [4532., 4874.], + [4664., 5018.], + [4796., 5162.], + [4928., 5306.]]) + >>> np.einsum(a, [0,1,2], b, [1,0,3], [2,3]) + array([[4400., 4730.], + [4532., 4874.], + [4664., 5018.], + [4796., 5162.], + [4928., 5306.]]) + >>> np.tensordot(a,b, axes=([1,0],[0,1])) + array([[4400., 4730.], + [4532., 4874.], + [4664., 5018.], + [4796., 5162.], + [4928., 5306.]]) + + Writeable returned arrays (since version 1.10.0): + + >>> a = np.zeros((3, 3)) + >>> np.einsum('ii->i', a)[:] = 1 + >>> a + array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + + Example of ellipsis use: + + >>> a = np.arange(6).reshape((3,2)) + >>> b = np.arange(12).reshape((4,3)) + >>> np.einsum('ki,jk->ij', a, b) + array([[10, 28, 46, 64], + [13, 40, 67, 94]]) + >>> np.einsum('ki,...k->i...', a, b) + array([[10, 28, 46, 64], + [13, 40, 67, 94]]) + >>> np.einsum('k...,jk', a, b) + array([[10, 28, 46, 64], + [13, 40, 67, 94]]) + + Chained array operations. For more complicated contractions, speed ups + might be achieved by repeatedly computing a 'greedy' path or pre-computing + the 'optimal' path and repeatedly applying it, using an `einsum_path` + insertion (since version 1.12.0). Performance improvements can be + particularly significant with larger arrays: + + >>> a = np.ones(64).reshape(2,4,8) + + Basic `einsum`: ~1520ms (benchmarked on 3.1GHz Intel i5.) + + >>> for iteration in range(500): + ... _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a) + + Sub-optimal `einsum` (due to repeated path calculation time): ~330ms + + >>> for iteration in range(500): + ... _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, + ... optimize='optimal') + + Greedy `einsum` (faster optimal path approximation): ~160ms + + >>> for iteration in range(500): + ... _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy') + + Optimal `einsum` (best usage pattern in some use cases): ~110ms + + >>> path = np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, + ... optimize='optimal')[0] + >>> for iteration in range(500): + ... _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path) + + """ + # Special handling if out is specified + specified_out = out is not None + + # If no optimization, run pure einsum + if optimize is False: + if specified_out: + kwargs['out'] = out + return c_einsum(*operands, **kwargs) + + # Check the kwargs to avoid a more cryptic error later, without having to + # repeat default values here + valid_einsum_kwargs = ['dtype', 'order', 'casting'] + unknown_kwargs = [k for (k, v) in kwargs.items() if + k not in valid_einsum_kwargs] + if len(unknown_kwargs): + raise TypeError("Did not understand the following kwargs: %s" + % unknown_kwargs) + + # Build the contraction list and operand + operands, contraction_list = einsum_path(*operands, optimize=optimize, + einsum_call=True) + + # Handle order kwarg for output array, c_einsum allows mixed case + output_order = kwargs.pop('order', 'K') + if output_order.upper() == 'A': + if all(arr.flags.f_contiguous for arr in operands): + output_order = 'F' + else: + output_order = 'C' + + # Start contraction loop + for num, contraction in enumerate(contraction_list): + inds, idx_rm, einsum_str, remaining, blas = contraction + tmp_operands = [operands.pop(x) for x in inds] + + # Do we need to deal with the output? + handle_out = specified_out and ((num + 1) == len(contraction_list)) + + # Call tensordot if still possible + if blas: + # Checks have already been handled + input_str, results_index = einsum_str.split('->') + input_left, input_right = input_str.split(',') + + tensor_result = input_left + input_right + for s in idx_rm: + tensor_result = tensor_result.replace(s, "") + + # Find indices to contract over + left_pos, right_pos = [], [] + for s in sorted(idx_rm): + left_pos.append(input_left.find(s)) + right_pos.append(input_right.find(s)) + + # Contract! + new_view = tensordot( + *tmp_operands, axes=(tuple(left_pos), tuple(right_pos)) + ) + + # Build a new view if needed + if (tensor_result != results_index) or handle_out: + if handle_out: + kwargs["out"] = out + new_view = c_einsum( + tensor_result + '->' + results_index, new_view, **kwargs + ) + + # Call einsum + else: + # If out was specified + if handle_out: + kwargs["out"] = out + + # Do the contraction + new_view = c_einsum(einsum_str, *tmp_operands, **kwargs) + + # Append new items and dereference what we can + operands.append(new_view) + del tmp_operands, new_view + + if specified_out: + return out + else: + return asanyarray(operands[0], order=output_order) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/einsumfunc.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/einsumfunc.pyi new file mode 100644 index 0000000000000000000000000000000000000000..782e7bcd76a90364a20fa8a2276d82501e2e69a1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/einsumfunc.pyi @@ -0,0 +1,185 @@ +from collections.abc import Sequence +from typing import TypeAlias, TypeVar, Any, overload, Literal + +import numpy as np +from numpy import number, _OrderKACF +from numpy._typing import ( + NDArray, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, + _DTypeLikeBool, + _DTypeLikeUInt, + _DTypeLikeInt, + _DTypeLikeFloat, + _DTypeLikeComplex, + _DTypeLikeComplex_co, + _DTypeLikeObject, +) + +__all__ = ["einsum", "einsum_path"] + +_ArrayType = TypeVar( + "_ArrayType", + bound=NDArray[np.bool | number[Any]], +) + +_OptimizeKind: TypeAlias = bool | Literal["greedy", "optimal"] | Sequence[Any] | None +_CastingSafe: TypeAlias = Literal["no", "equiv", "safe", "same_kind"] +_CastingUnsafe: TypeAlias = Literal["unsafe"] + + +# TODO: Properly handle the `casting`-based combinatorics +# TODO: We need to evaluate the content `__subscripts` in order +# to identify whether or an array or scalar is returned. At a cursory +# glance this seems like something that can quite easily be done with +# a mypy plugin. +# Something like `is_scalar = bool(__subscripts.partition("->")[-1])` +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeBool_co, + out: None = ..., + dtype: None | _DTypeLikeBool = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeUInt_co, + out: None = ..., + dtype: None | _DTypeLikeUInt = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeInt_co, + out: None = ..., + dtype: None | _DTypeLikeInt = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeFloat_co, + out: None = ..., + dtype: None | _DTypeLikeFloat = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeComplex_co, + out: None = ..., + dtype: None | _DTypeLikeComplex = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: Any, + casting: _CastingUnsafe, + dtype: None | _DTypeLikeComplex_co = ..., + out: None = ..., + order: _OrderKACF = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeComplex_co, + out: _ArrayType, + dtype: None | _DTypeLikeComplex_co = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> _ArrayType: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: Any, + out: _ArrayType, + casting: _CastingUnsafe, + dtype: None | _DTypeLikeComplex_co = ..., + order: _OrderKACF = ..., + optimize: _OptimizeKind = ..., +) -> _ArrayType: ... + +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeObject_co, + out: None = ..., + dtype: None | _DTypeLikeObject = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: Any, + casting: _CastingUnsafe, + dtype: None | _DTypeLikeObject = ..., + out: None = ..., + order: _OrderKACF = ..., + optimize: _OptimizeKind = ..., +) -> Any: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeObject_co, + out: _ArrayType, + dtype: None | _DTypeLikeObject = ..., + order: _OrderKACF = ..., + casting: _CastingSafe = ..., + optimize: _OptimizeKind = ..., +) -> _ArrayType: ... +@overload +def einsum( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: Any, + out: _ArrayType, + casting: _CastingUnsafe, + dtype: None | _DTypeLikeObject = ..., + order: _OrderKACF = ..., + optimize: _OptimizeKind = ..., +) -> _ArrayType: ... + +# NOTE: `einsum_call` is a hidden kwarg unavailable for public use. +# It is therefore excluded from the signatures below. +# NOTE: In practice the list consists of a `str` (first element) +# and a variable number of integer tuples. +def einsum_path( + subscripts: str | _ArrayLikeInt_co, + /, + *operands: _ArrayLikeComplex_co | _DTypeLikeObject, + optimize: _OptimizeKind = "greedy", + einsum_call: Literal[False] = False, +) -> tuple[list[Any], str]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/fromnumeric.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/fromnumeric.py new file mode 100644 index 0000000000000000000000000000000000000000..80508356217b725ef6665834932c4f6faf8e392e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/fromnumeric.py @@ -0,0 +1,4269 @@ +"""Module containing non-deprecated functions borrowed from Numeric. + +""" +import functools +import types +import warnings + +import numpy as np +from .._utils import set_module +from . import multiarray as mu +from . import overrides +from . import umath as um +from . import numerictypes as nt +from .multiarray import asarray, array, asanyarray, concatenate +from ._multiarray_umath import _array_converter +from . import _methods + +_dt_ = nt.sctype2char + +# functions that are methods +__all__ = [ + 'all', 'amax', 'amin', 'any', 'argmax', + 'argmin', 'argpartition', 'argsort', 'around', 'choose', 'clip', + 'compress', 'cumprod', 'cumsum', 'cumulative_prod', 'cumulative_sum', + 'diagonal', 'mean', 'max', 'min', 'matrix_transpose', + 'ndim', 'nonzero', 'partition', 'prod', 'ptp', 'put', + 'ravel', 'repeat', 'reshape', 'resize', 'round', + 'searchsorted', 'shape', 'size', 'sort', 'squeeze', + 'std', 'sum', 'swapaxes', 'take', 'trace', 'transpose', 'var', +] + +_gentype = types.GeneratorType +# save away Python sum +_sum_ = sum + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +# functions that are now methods +def _wrapit(obj, method, *args, **kwds): + conv = _array_converter(obj) + # As this already tried the method, subok is maybe quite reasonable here + # but this follows what was done before. TODO: revisit this. + arr, = conv.as_arrays(subok=False) + result = getattr(arr, method)(*args, **kwds) + + return conv.wrap(result, to_scalar=False) + + +def _wrapfunc(obj, method, *args, **kwds): + bound = getattr(obj, method, None) + if bound is None: + return _wrapit(obj, method, *args, **kwds) + + try: + return bound(*args, **kwds) + except TypeError: + # A TypeError occurs if the object does have such a method in its + # class, but its signature is not identical to that of NumPy's. This + # situation has occurred in the case of a downstream library like + # 'pandas'. + # + # Call _wrapit from within the except clause to ensure a potential + # exception has a traceback chain. + return _wrapit(obj, method, *args, **kwds) + + +def _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs): + passkwargs = {k: v for k, v in kwargs.items() + if v is not np._NoValue} + + if type(obj) is not mu.ndarray: + try: + reduction = getattr(obj, method) + except AttributeError: + pass + else: + # This branch is needed for reductions like any which don't + # support a dtype. + if dtype is not None: + return reduction(axis=axis, dtype=dtype, out=out, **passkwargs) + else: + return reduction(axis=axis, out=out, **passkwargs) + + return ufunc.reduce(obj, axis, dtype, out, **passkwargs) + + +def _wrapreduction_any_all(obj, ufunc, method, axis, out, **kwargs): + # Same as above function, but dtype is always bool (but never passed on) + passkwargs = {k: v for k, v in kwargs.items() + if v is not np._NoValue} + + if type(obj) is not mu.ndarray: + try: + reduction = getattr(obj, method) + except AttributeError: + pass + else: + return reduction(axis=axis, out=out, **passkwargs) + + return ufunc.reduce(obj, axis, bool, out, **passkwargs) + + +def _take_dispatcher(a, indices, axis=None, out=None, mode=None): + return (a, out) + + +@array_function_dispatch(_take_dispatcher) +def take(a, indices, axis=None, out=None, mode='raise'): + """ + Take elements from an array along an axis. + + When axis is not None, this function does the same thing as "fancy" + indexing (indexing arrays using arrays); however, it can be easier to use + if you need elements along a given axis. A call such as + ``np.take(arr, indices, axis=3)`` is equivalent to + ``arr[:,:,:,indices,...]``. + + Explained without fancy indexing, this is equivalent to the following use + of `ndindex`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of + indices:: + + Ni, Nk = a.shape[:axis], a.shape[axis+1:] + Nj = indices.shape + for ii in ndindex(Ni): + for jj in ndindex(Nj): + for kk in ndindex(Nk): + out[ii + jj + kk] = a[ii + (indices[jj],) + kk] + + Parameters + ---------- + a : array_like (Ni..., M, Nk...) + The source array. + indices : array_like (Nj...) + The indices of the values to extract. + Also allow scalars for indices. + axis : int, optional + The axis over which to select values. By default, the flattened + input array is used. + out : ndarray, optional (Ni..., Nj..., Nk...) + If provided, the result will be placed in this array. It should + be of the appropriate shape and dtype. Note that `out` is always + buffered if `mode='raise'`; use other modes for better performance. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices will behave. + + * 'raise' -- raise an error (default) + * 'wrap' -- wrap around + * 'clip' -- clip to the range + + 'clip' mode means that all indices that are too large are replaced + by the index that addresses the last element along that axis. Note + that this disables indexing with negative numbers. + + Returns + ------- + out : ndarray (Ni..., Nj..., Nk...) + The returned array has the same type as `a`. + + See Also + -------- + compress : Take elements using a boolean mask + ndarray.take : equivalent method + take_along_axis : Take elements by matching the array and the index arrays + + Notes + ----- + By eliminating the inner loop in the description above, and using `s_` to + build simple slice objects, `take` can be expressed in terms of applying + fancy indexing to each 1-d slice:: + + Ni, Nk = a.shape[:axis], a.shape[axis+1:] + for ii in ndindex(Ni): + for kk in ndindex(Nj): + out[ii + s_[...,] + kk] = a[ii + s_[:,] + kk][indices] + + For this reason, it is equivalent to (but faster than) the following use + of `apply_along_axis`:: + + out = np.apply_along_axis(lambda a_1d: a_1d[indices], axis, a) + + Examples + -------- + >>> import numpy as np + >>> a = [4, 3, 5, 7, 6, 8] + >>> indices = [0, 1, 4] + >>> np.take(a, indices) + array([4, 3, 6]) + + In this example if `a` is an ndarray, "fancy" indexing can be used. + + >>> a = np.array(a) + >>> a[indices] + array([4, 3, 6]) + + If `indices` is not one dimensional, the output also has these dimensions. + + >>> np.take(a, [[0, 1], [2, 3]]) + array([[4, 3], + [5, 7]]) + """ + return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode) + + +def _reshape_dispatcher(a, /, shape=None, order=None, *, newshape=None, + copy=None): + return (a,) + + +@array_function_dispatch(_reshape_dispatcher) +def reshape(a, /, shape=None, order='C', *, newshape=None, copy=None): + """ + Gives a new shape to an array without changing its data. + + Parameters + ---------- + a : array_like + Array to be reshaped. + shape : int or tuple of ints + The new shape should be compatible with the original shape. If + an integer, then the result will be a 1-D array of that length. + One shape dimension can be -1. In this case, the value is + inferred from the length of the array and remaining dimensions. + order : {'C', 'F', 'A'}, optional + Read the elements of ``a`` using this index order, and place the + elements into the reshaped array using this index order. 'C' + means to read / write the elements using C-like index order, + with the last axis index changing fastest, back to the first + axis index changing slowest. 'F' means to read / write the + elements using Fortran-like index order, with the first index + changing fastest, and the last index changing slowest. Note that + the 'C' and 'F' options take no account of the memory layout of + the underlying array, and only refer to the order of indexing. + 'A' means to read / write the elements in Fortran-like index + order if ``a`` is Fortran *contiguous* in memory, C-like order + otherwise. + newshape : int or tuple of ints + .. deprecated:: 2.1 + Replaced by ``shape`` argument. Retained for backward + compatibility. + copy : bool, optional + If ``True``, then the array data is copied. If ``None``, a copy will + only be made if it's required by ``order``. For ``False`` it raises + a ``ValueError`` if a copy cannot be avoided. Default: ``None``. + + Returns + ------- + reshaped_array : ndarray + This will be a new view object if possible; otherwise, it will + be a copy. Note there is no guarantee of the *memory layout* (C- or + Fortran- contiguous) of the returned array. + + See Also + -------- + ndarray.reshape : Equivalent method. + + Notes + ----- + It is not always possible to change the shape of an array without copying + the data. + + The ``order`` keyword gives the index ordering both for *fetching* + the values from ``a``, and then *placing* the values into the output + array. For example, let's say you have an array: + + >>> a = np.arange(6).reshape((3, 2)) + >>> a + array([[0, 1], + [2, 3], + [4, 5]]) + + You can think of reshaping as first raveling the array (using the given + index order), then inserting the elements from the raveled array into the + new array using the same kind of index ordering as was used for the + raveling. + + >>> np.reshape(a, (2, 3)) # C-like index ordering + array([[0, 1, 2], + [3, 4, 5]]) + >>> np.reshape(np.ravel(a), (2, 3)) # equivalent to C ravel then C reshape + array([[0, 1, 2], + [3, 4, 5]]) + >>> np.reshape(a, (2, 3), order='F') # Fortran-like index ordering + array([[0, 4, 3], + [2, 1, 5]]) + >>> np.reshape(np.ravel(a, order='F'), (2, 3), order='F') + array([[0, 4, 3], + [2, 1, 5]]) + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1,2,3], [4,5,6]]) + >>> np.reshape(a, 6) + array([1, 2, 3, 4, 5, 6]) + >>> np.reshape(a, 6, order='F') + array([1, 4, 2, 5, 3, 6]) + + >>> np.reshape(a, (3,-1)) # the unspecified value is inferred to be 2 + array([[1, 2], + [3, 4], + [5, 6]]) + """ + if newshape is None and shape is None: + raise TypeError( + "reshape() missing 1 required positional argument: 'shape'") + if newshape is not None: + if shape is not None: + raise TypeError( + "You cannot specify 'newshape' and 'shape' arguments " + "at the same time.") + # Deprecated in NumPy 2.1, 2024-04-18 + warnings.warn( + "`newshape` keyword argument is deprecated, " + "use `shape=...` or pass shape positionally instead. " + "(deprecated in NumPy 2.1)", + DeprecationWarning, + stacklevel=2, + ) + shape = newshape + if copy is not None: + return _wrapfunc(a, 'reshape', shape, order=order, copy=copy) + return _wrapfunc(a, 'reshape', shape, order=order) + + +def _choose_dispatcher(a, choices, out=None, mode=None): + yield a + yield from choices + yield out + + +@array_function_dispatch(_choose_dispatcher) +def choose(a, choices, out=None, mode='raise'): + """ + Construct an array from an index array and a list of arrays to choose from. + + First of all, if confused or uncertain, definitely look at the Examples - + in its full generality, this function is less simple than it might + seem from the following code description:: + + np.choose(a,c) == np.array([c[a[I]][I] for I in np.ndindex(a.shape)]) + + But this omits some subtleties. Here is a fully general summary: + + Given an "index" array (`a`) of integers and a sequence of ``n`` arrays + (`choices`), `a` and each choice array are first broadcast, as necessary, + to arrays of a common shape; calling these *Ba* and *Bchoices[i], i = + 0,...,n-1* we have that, necessarily, ``Ba.shape == Bchoices[i].shape`` + for each ``i``. Then, a new array with shape ``Ba.shape`` is created as + follows: + + * if ``mode='raise'`` (the default), then, first of all, each element of + ``a`` (and thus ``Ba``) must be in the range ``[0, n-1]``; now, suppose + that ``i`` (in that range) is the value at the ``(j0, j1, ..., jm)`` + position in ``Ba`` - then the value at the same position in the new array + is the value in ``Bchoices[i]`` at that same position; + + * if ``mode='wrap'``, values in `a` (and thus `Ba`) may be any (signed) + integer; modular arithmetic is used to map integers outside the range + `[0, n-1]` back into that range; and then the new array is constructed + as above; + + * if ``mode='clip'``, values in `a` (and thus ``Ba``) may be any (signed) + integer; negative integers are mapped to 0; values greater than ``n-1`` + are mapped to ``n-1``; and then the new array is constructed as above. + + Parameters + ---------- + a : int array + This array must contain integers in ``[0, n-1]``, where ``n`` is the + number of choices, unless ``mode=wrap`` or ``mode=clip``, in which + cases any integers are permissible. + choices : sequence of arrays + Choice arrays. `a` and all of the choices must be broadcastable to the + same shape. If `choices` is itself an array (not recommended), then + its outermost dimension (i.e., the one corresponding to + ``choices.shape[0]``) is taken as defining the "sequence". + out : array, optional + If provided, the result will be inserted into this array. It should + be of the appropriate shape and dtype. Note that `out` is always + buffered if ``mode='raise'``; use other modes for better performance. + mode : {'raise' (default), 'wrap', 'clip'}, optional + Specifies how indices outside ``[0, n-1]`` will be treated: + + * 'raise' : an exception is raised + * 'wrap' : value becomes value mod ``n`` + * 'clip' : values < 0 are mapped to 0, values > n-1 are mapped to n-1 + + Returns + ------- + merged_array : array + The merged result. + + Raises + ------ + ValueError: shape mismatch + If `a` and each choice array are not all broadcastable to the same + shape. + + See Also + -------- + ndarray.choose : equivalent method + numpy.take_along_axis : Preferable if `choices` is an array + + Notes + ----- + To reduce the chance of misinterpretation, even though the following + "abuse" is nominally supported, `choices` should neither be, nor be + thought of as, a single array, i.e., the outermost sequence-like container + should be either a list or a tuple. + + Examples + -------- + + >>> import numpy as np + >>> choices = [[0, 1, 2, 3], [10, 11, 12, 13], + ... [20, 21, 22, 23], [30, 31, 32, 33]] + >>> np.choose([2, 3, 1, 0], choices + ... # the first element of the result will be the first element of the + ... # third (2+1) "array" in choices, namely, 20; the second element + ... # will be the second element of the fourth (3+1) choice array, i.e., + ... # 31, etc. + ... ) + array([20, 31, 12, 3]) + >>> np.choose([2, 4, 1, 0], choices, mode='clip') # 4 goes to 3 (4-1) + array([20, 31, 12, 3]) + >>> # because there are 4 choice arrays + >>> np.choose([2, 4, 1, 0], choices, mode='wrap') # 4 goes to (4 mod 4) + array([20, 1, 12, 3]) + >>> # i.e., 0 + + A couple examples illustrating how choose broadcasts: + + >>> a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]] + >>> choices = [-10, 10] + >>> np.choose(a, choices) + array([[ 10, -10, 10], + [-10, 10, -10], + [ 10, -10, 10]]) + + >>> # With thanks to Anne Archibald + >>> a = np.array([0, 1]).reshape((2,1,1)) + >>> c1 = np.array([1, 2, 3]).reshape((1,3,1)) + >>> c2 = np.array([-1, -2, -3, -4, -5]).reshape((1,1,5)) + >>> np.choose(a, (c1, c2)) # result is 2x3x5, res[0,:,:]=c1, res[1,:,:]=c2 + array([[[ 1, 1, 1, 1, 1], + [ 2, 2, 2, 2, 2], + [ 3, 3, 3, 3, 3]], + [[-1, -2, -3, -4, -5], + [-1, -2, -3, -4, -5], + [-1, -2, -3, -4, -5]]]) + + """ + return _wrapfunc(a, 'choose', choices, out=out, mode=mode) + + +def _repeat_dispatcher(a, repeats, axis=None): + return (a,) + + +@array_function_dispatch(_repeat_dispatcher) +def repeat(a, repeats, axis=None): + """ + Repeat each element of an array after themselves + + Parameters + ---------- + a : array_like + Input array. + repeats : int or array of ints + The number of repetitions for each element. `repeats` is broadcasted + to fit the shape of the given axis. + axis : int, optional + The axis along which to repeat values. By default, use the + flattened input array, and return a flat output array. + + Returns + ------- + repeated_array : ndarray + Output array which has the same shape as `a`, except along + the given axis. + + See Also + -------- + tile : Tile an array. + unique : Find the unique elements of an array. + + Examples + -------- + >>> import numpy as np + >>> np.repeat(3, 4) + array([3, 3, 3, 3]) + >>> x = np.array([[1,2],[3,4]]) + >>> np.repeat(x, 2) + array([1, 1, 2, 2, 3, 3, 4, 4]) + >>> np.repeat(x, 3, axis=1) + array([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 4, 4, 4]]) + >>> np.repeat(x, [1, 2], axis=0) + array([[1, 2], + [3, 4], + [3, 4]]) + + """ + return _wrapfunc(a, 'repeat', repeats, axis=axis) + + +def _put_dispatcher(a, ind, v, mode=None): + return (a, ind, v) + + +@array_function_dispatch(_put_dispatcher) +def put(a, ind, v, mode='raise'): + """ + Replaces specified elements of an array with given values. + + The indexing works on the flattened target array. `put` is roughly + equivalent to: + + :: + + a.flat[ind] = v + + Parameters + ---------- + a : ndarray + Target array. + ind : array_like + Target indices, interpreted as integers. + v : array_like + Values to place in `a` at target indices. If `v` is shorter than + `ind` it will be repeated as necessary. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices will behave. + + * 'raise' -- raise an error (default) + * 'wrap' -- wrap around + * 'clip' -- clip to the range + + 'clip' mode means that all indices that are too large are replaced + by the index that addresses the last element along that axis. Note + that this disables indexing with negative numbers. In 'raise' mode, + if an exception occurs the target array may still be modified. + + See Also + -------- + putmask, place + put_along_axis : Put elements by matching the array and the index arrays + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(5) + >>> np.put(a, [0, 2], [-44, -55]) + >>> a + array([-44, 1, -55, 3, 4]) + + >>> a = np.arange(5) + >>> np.put(a, 22, -5, mode='clip') + >>> a + array([ 0, 1, 2, 3, -5]) + + """ + try: + put = a.put + except AttributeError as e: + raise TypeError("argument 1 must be numpy.ndarray, " + "not {name}".format(name=type(a).__name__)) from e + + return put(ind, v, mode=mode) + + +def _swapaxes_dispatcher(a, axis1, axis2): + return (a,) + + +@array_function_dispatch(_swapaxes_dispatcher) +def swapaxes(a, axis1, axis2): + """ + Interchange two axes of an array. + + Parameters + ---------- + a : array_like + Input array. + axis1 : int + First axis. + axis2 : int + Second axis. + + Returns + ------- + a_swapped : ndarray + For NumPy >= 1.10.0, if `a` is an ndarray, then a view of `a` is + returned; otherwise a new array is created. For earlier NumPy + versions a view of `a` is returned only if the order of the + axes is changed, otherwise the input array is returned. + + Examples + -------- + >>> import numpy as np + >>> x = np.array([[1,2,3]]) + >>> np.swapaxes(x,0,1) + array([[1], + [2], + [3]]) + + >>> x = np.array([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x + array([[[0, 1], + [2, 3]], + [[4, 5], + [6, 7]]]) + + >>> np.swapaxes(x,0,2) + array([[[0, 4], + [2, 6]], + [[1, 5], + [3, 7]]]) + + """ + return _wrapfunc(a, 'swapaxes', axis1, axis2) + + +def _transpose_dispatcher(a, axes=None): + return (a,) + + +@array_function_dispatch(_transpose_dispatcher) +def transpose(a, axes=None): + """ + Returns an array with axes transposed. + + For a 1-D array, this returns an unchanged view of the original array, as a + transposed vector is simply the same vector. + To convert a 1-D array into a 2-D column vector, an additional dimension + must be added, e.g., ``np.atleast_2d(a).T`` achieves this, as does + ``a[:, np.newaxis]``. + For a 2-D array, this is the standard matrix transpose. + For an n-D array, if axes are given, their order indicates how the + axes are permuted (see Examples). If axes are not provided, then + ``transpose(a).shape == a.shape[::-1]``. + + Parameters + ---------- + a : array_like + Input array. + axes : tuple or list of ints, optional + If specified, it must be a tuple or list which contains a permutation + of [0, 1, ..., N-1] where N is the number of axes of `a`. Negative + indices can also be used to specify axes. The i-th axis of the returned + array will correspond to the axis numbered ``axes[i]`` of the input. + If not specified, defaults to ``range(a.ndim)[::-1]``, which reverses + the order of the axes. + + Returns + ------- + p : ndarray + `a` with its axes permuted. A view is returned whenever possible. + + See Also + -------- + ndarray.transpose : Equivalent method. + moveaxis : Move axes of an array to new positions. + argsort : Return the indices that would sort an array. + + Notes + ----- + Use ``transpose(a, argsort(axes))`` to invert the transposition of tensors + when using the `axes` keyword argument. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4]]) + >>> a + array([[1, 2], + [3, 4]]) + >>> np.transpose(a) + array([[1, 3], + [2, 4]]) + + >>> a = np.array([1, 2, 3, 4]) + >>> a + array([1, 2, 3, 4]) + >>> np.transpose(a) + array([1, 2, 3, 4]) + + >>> a = np.ones((1, 2, 3)) + >>> np.transpose(a, (1, 0, 2)).shape + (2, 1, 3) + + >>> a = np.ones((2, 3, 4, 5)) + >>> np.transpose(a).shape + (5, 4, 3, 2) + + >>> a = np.arange(3*4*5).reshape((3, 4, 5)) + >>> np.transpose(a, (-1, 0, -2)).shape + (5, 3, 4) + + """ + return _wrapfunc(a, 'transpose', axes) + + +def _matrix_transpose_dispatcher(x): + return (x,) + +@array_function_dispatch(_matrix_transpose_dispatcher) +def matrix_transpose(x, /): + """ + Transposes a matrix (or a stack of matrices) ``x``. + + This function is Array API compatible. + + Parameters + ---------- + x : array_like + Input array having shape (..., M, N) and whose two innermost + dimensions form ``MxN`` matrices. + + Returns + ------- + out : ndarray + An array containing the transpose for each matrix and having shape + (..., N, M). + + See Also + -------- + transpose : Generic transpose method. + + Examples + -------- + >>> import numpy as np + >>> np.matrix_transpose([[1, 2], [3, 4]]) + array([[1, 3], + [2, 4]]) + + >>> np.matrix_transpose([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + array([[[1, 3], + [2, 4]], + [[5, 7], + [6, 8]]]) + + """ + x = asanyarray(x) + if x.ndim < 2: + raise ValueError( + f"Input array must be at least 2-dimensional, but it is {x.ndim}" + ) + return swapaxes(x, -1, -2) + + +def _partition_dispatcher(a, kth, axis=None, kind=None, order=None): + return (a,) + + +@array_function_dispatch(_partition_dispatcher) +def partition(a, kth, axis=-1, kind='introselect', order=None): + """ + Return a partitioned copy of an array. + + Creates a copy of the array and partially sorts it in such a way that + the value of the element in k-th position is in the position it would be + in a sorted array. In the output array, all elements smaller than the k-th + element are located to the left of this element and all equal or greater + are located to its right. The ordering of the elements in the two + partitions on the either side of the k-th element in the output array is + undefined. + + Parameters + ---------- + a : array_like + Array to be sorted. + kth : int or sequence of ints + Element index to partition by. The k-th value of the element + will be in its final sorted position and all smaller elements + will be moved before it and all equal or greater elements behind + it. The order of all elements in the partitions is undefined. If + provided with a sequence of k-th it will partition all elements + indexed by k-th of them into their sorted position at once. + + .. deprecated:: 1.22.0 + Passing booleans as index is deprecated. + axis : int or None, optional + Axis along which to sort. If None, the array is flattened before + sorting. The default is -1, which sorts along the last axis. + kind : {'introselect'}, optional + Selection algorithm. Default is 'introselect'. + order : str or list of str, optional + When `a` is an array with fields defined, this argument + specifies which fields to compare first, second, etc. A single + field can be specified as a string. Not all fields need be + specified, but unspecified fields will still be used, in the + order in which they come up in the dtype, to break ties. + + Returns + ------- + partitioned_array : ndarray + Array of the same type and shape as `a`. + + See Also + -------- + ndarray.partition : Method to sort an array in-place. + argpartition : Indirect partition. + sort : Full sorting + + Notes + ----- + The various selection algorithms are characterized by their average + speed, worst case performance, work space size, and whether they are + stable. A stable sort keeps items with the same key in the same + relative order. The available algorithms have the following + properties: + + ================= ======= ============= ============ ======= + kind speed worst case work space stable + ================= ======= ============= ============ ======= + 'introselect' 1 O(n) 0 no + ================= ======= ============= ============ ======= + + All the partition algorithms make temporary copies of the data when + partitioning along any but the last axis. Consequently, + partitioning along the last axis is faster and uses less space than + partitioning along any other axis. + + The sort order for complex numbers is lexicographic. If both the + real and imaginary parts are non-nan then the order is determined by + the real parts except when they are equal, in which case the order + is determined by the imaginary parts. + + The sort order of ``np.nan`` is bigger than ``np.inf``. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0]) + >>> p = np.partition(a, 4) + >>> p + array([0, 1, 2, 1, 2, 5, 2, 3, 3, 6, 7, 7, 7, 7]) # may vary + + ``p[4]`` is 2; all elements in ``p[:4]`` are less than or equal + to ``p[4]``, and all elements in ``p[5:]`` are greater than or + equal to ``p[4]``. The partition is:: + + [0, 1, 2, 1], [2], [5, 2, 3, 3, 6, 7, 7, 7, 7] + + The next example shows the use of multiple values passed to `kth`. + + >>> p2 = np.partition(a, (4, 8)) + >>> p2 + array([0, 1, 2, 1, 2, 3, 3, 2, 5, 6, 7, 7, 7, 7]) + + ``p2[4]`` is 2 and ``p2[8]`` is 5. All elements in ``p2[:4]`` + are less than or equal to ``p2[4]``, all elements in ``p2[5:8]`` + are greater than or equal to ``p2[4]`` and less than or equal to + ``p2[8]``, and all elements in ``p2[9:]`` are greater than or + equal to ``p2[8]``. The partition is:: + + [0, 1, 2, 1], [2], [3, 3, 2], [5], [6, 7, 7, 7, 7] + """ + if axis is None: + # flatten returns (1, N) for np.matrix, so always use the last axis + a = asanyarray(a).flatten() + axis = -1 + else: + a = asanyarray(a).copy(order="K") + a.partition(kth, axis=axis, kind=kind, order=order) + return a + + +def _argpartition_dispatcher(a, kth, axis=None, kind=None, order=None): + return (a,) + + +@array_function_dispatch(_argpartition_dispatcher) +def argpartition(a, kth, axis=-1, kind='introselect', order=None): + """ + Perform an indirect partition along the given axis using the + algorithm specified by the `kind` keyword. It returns an array of + indices of the same shape as `a` that index data along the given + axis in partitioned order. + + Parameters + ---------- + a : array_like + Array to sort. + kth : int or sequence of ints + Element index to partition by. The k-th element will be in its + final sorted position and all smaller elements will be moved + before it and all larger elements behind it. The order of all + elements in the partitions is undefined. If provided with a + sequence of k-th it will partition all of them into their sorted + position at once. + + .. deprecated:: 1.22.0 + Passing booleans as index is deprecated. + axis : int or None, optional + Axis along which to sort. The default is -1 (the last axis). If + None, the flattened array is used. + kind : {'introselect'}, optional + Selection algorithm. Default is 'introselect' + order : str or list of str, optional + When `a` is an array with fields defined, this argument + specifies which fields to compare first, second, etc. A single + field can be specified as a string, and not all fields need be + specified, but unspecified fields will still be used, in the + order in which they come up in the dtype, to break ties. + + Returns + ------- + index_array : ndarray, int + Array of indices that partition `a` along the specified axis. + If `a` is one-dimensional, ``a[index_array]`` yields a partitioned `a`. + More generally, ``np.take_along_axis(a, index_array, axis=axis)`` + always yields the partitioned `a`, irrespective of dimensionality. + + See Also + -------- + partition : Describes partition algorithms used. + ndarray.partition : Inplace partition. + argsort : Full indirect sort. + take_along_axis : Apply ``index_array`` from argpartition + to an array as if by calling partition. + + Notes + ----- + The returned indices are not guaranteed to be sorted according to + the values. Furthermore, the default selection algorithm ``introselect`` + is unstable, and hence the returned indices are not guaranteed + to be the earliest/latest occurrence of the element. + + `argpartition` works for real/complex inputs with nan values, + see `partition` for notes on the enhanced sort order and + different selection algorithms. + + Examples + -------- + One dimensional array: + + >>> import numpy as np + >>> x = np.array([3, 4, 2, 1]) + >>> x[np.argpartition(x, 3)] + array([2, 1, 3, 4]) # may vary + >>> x[np.argpartition(x, (1, 3))] + array([1, 2, 3, 4]) # may vary + + >>> x = [3, 4, 2, 1] + >>> np.array(x)[np.argpartition(x, 3)] + array([2, 1, 3, 4]) # may vary + + Multi-dimensional array: + + >>> x = np.array([[3, 4, 2], [1, 3, 1]]) + >>> index_array = np.argpartition(x, kth=1, axis=-1) + >>> # below is the same as np.partition(x, kth=1) + >>> np.take_along_axis(x, index_array, axis=-1) + array([[2, 3, 4], + [1, 1, 3]]) + + """ + return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order) + + +def _sort_dispatcher(a, axis=None, kind=None, order=None, *, stable=None): + return (a,) + + +@array_function_dispatch(_sort_dispatcher) +def sort(a, axis=-1, kind=None, order=None, *, stable=None): + """ + Return a sorted copy of an array. + + Parameters + ---------- + a : array_like + Array to be sorted. + axis : int or None, optional + Axis along which to sort. If None, the array is flattened before + sorting. The default is -1, which sorts along the last axis. + kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, optional + Sorting algorithm. The default is 'quicksort'. Note that both 'stable' + and 'mergesort' use timsort or radix sort under the covers and, + in general, the actual implementation will vary with data type. + The 'mergesort' option is retained for backwards compatibility. + order : str or list of str, optional + When `a` is an array with fields defined, this argument specifies + which fields to compare first, second, etc. A single field can + be specified as a string, and not all fields need be specified, + but unspecified fields will still be used, in the order in which + they come up in the dtype, to break ties. + stable : bool, optional + Sort stability. If ``True``, the returned array will maintain + the relative order of ``a`` values which compare as equal. + If ``False`` or ``None``, this is not guaranteed. Internally, + this option selects ``kind='stable'``. Default: ``None``. + + .. versionadded:: 2.0.0 + + Returns + ------- + sorted_array : ndarray + Array of the same type and shape as `a`. + + See Also + -------- + ndarray.sort : Method to sort an array in-place. + argsort : Indirect sort. + lexsort : Indirect stable sort on multiple keys. + searchsorted : Find elements in a sorted array. + partition : Partial sort. + + Notes + ----- + The various sorting algorithms are characterized by their average speed, + worst case performance, work space size, and whether they are stable. A + stable sort keeps items with the same key in the same relative + order. The four algorithms implemented in NumPy have the following + properties: + + =========== ======= ============= ============ ======== + kind speed worst case work space stable + =========== ======= ============= ============ ======== + 'quicksort' 1 O(n^2) 0 no + 'heapsort' 3 O(n*log(n)) 0 no + 'mergesort' 2 O(n*log(n)) ~n/2 yes + 'timsort' 2 O(n*log(n)) ~n/2 yes + =========== ======= ============= ============ ======== + + .. note:: The datatype determines which of 'mergesort' or 'timsort' + is actually used, even if 'mergesort' is specified. User selection + at a finer scale is not currently available. + + For performance, ``sort`` makes a temporary copy if needed to make the data + `contiguous `_ + in memory along the sort axis. For even better performance and reduced + memory consumption, ensure that the array is already contiguous along the + sort axis. + + The sort order for complex numbers is lexicographic. If both the real + and imaginary parts are non-nan then the order is determined by the + real parts except when they are equal, in which case the order is + determined by the imaginary parts. + + Previous to numpy 1.4.0 sorting real and complex arrays containing nan + values led to undefined behaviour. In numpy versions >= 1.4.0 nan + values are sorted to the end. The extended sort order is: + + * Real: [R, nan] + * Complex: [R + Rj, R + nanj, nan + Rj, nan + nanj] + + where R is a non-nan real value. Complex values with the same nan + placements are sorted according to the non-nan part if it exists. + Non-nan values are sorted as before. + + quicksort has been changed to: + `introsort `_. + When sorting does not make enough progress it switches to + `heapsort `_. + This implementation makes quicksort O(n*log(n)) in the worst case. + + 'stable' automatically chooses the best stable sorting algorithm + for the data type being sorted. + It, along with 'mergesort' is currently mapped to + `timsort `_ + or `radix sort `_ + depending on the data type. + API forward compatibility currently limits the + ability to select the implementation and it is hardwired for the different + data types. + + Timsort is added for better performance on already or nearly + sorted data. On random data timsort is almost identical to + mergesort. It is now used for stable sort while quicksort is still the + default sort if none is chosen. For timsort details, refer to + `CPython listsort.txt + `_ + 'mergesort' and 'stable' are mapped to radix sort for integer data types. + Radix sort is an O(n) sort instead of O(n log n). + + NaT now sorts to the end of arrays for consistency with NaN. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1,4],[3,1]]) + >>> np.sort(a) # sort along the last axis + array([[1, 4], + [1, 3]]) + >>> np.sort(a, axis=None) # sort the flattened array + array([1, 1, 3, 4]) + >>> np.sort(a, axis=0) # sort along the first axis + array([[1, 1], + [3, 4]]) + + Use the `order` keyword to specify a field to use when sorting a + structured array: + + >>> dtype = [('name', 'S10'), ('height', float), ('age', int)] + >>> values = [('Arthur', 1.8, 41), ('Lancelot', 1.9, 38), + ... ('Galahad', 1.7, 38)] + >>> a = np.array(values, dtype=dtype) # create a structured array + >>> np.sort(a, order='height') # doctest: +SKIP + array([('Galahad', 1.7, 38), ('Arthur', 1.8, 41), + ('Lancelot', 1.8999999999999999, 38)], + dtype=[('name', '|S10'), ('height', '>> np.sort(a, order=['age', 'height']) # doctest: +SKIP + array([('Galahad', 1.7, 38), ('Lancelot', 1.8999999999999999, 38), + ('Arthur', 1.8, 41)], + dtype=[('name', '|S10'), ('height', '>> import numpy as np + >>> x = np.array([3, 1, 2]) + >>> np.argsort(x) + array([1, 2, 0]) + + Two-dimensional array: + + >>> x = np.array([[0, 3], [2, 2]]) + >>> x + array([[0, 3], + [2, 2]]) + + >>> ind = np.argsort(x, axis=0) # sorts along first axis (down) + >>> ind + array([[0, 1], + [1, 0]]) + >>> np.take_along_axis(x, ind, axis=0) # same as np.sort(x, axis=0) + array([[0, 2], + [2, 3]]) + + >>> ind = np.argsort(x, axis=1) # sorts along last axis (across) + >>> ind + array([[0, 1], + [0, 1]]) + >>> np.take_along_axis(x, ind, axis=1) # same as np.sort(x, axis=1) + array([[0, 3], + [2, 2]]) + + Indices of the sorted elements of a N-dimensional array: + + >>> ind = np.unravel_index(np.argsort(x, axis=None), x.shape) + >>> ind + (array([0, 1, 1, 0]), array([0, 0, 1, 1])) + >>> x[ind] # same as np.sort(x, axis=None) + array([0, 2, 2, 3]) + + Sorting with keys: + + >>> x = np.array([(1, 0), (0, 1)], dtype=[('x', '>> x + array([(1, 0), (0, 1)], + dtype=[('x', '>> np.argsort(x, order=('x','y')) + array([1, 0]) + + >>> np.argsort(x, order=('y','x')) + array([0, 1]) + + """ + return _wrapfunc( + a, 'argsort', axis=axis, kind=kind, order=order, stable=stable + ) + +def _argmax_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue): + return (a, out) + + +@array_function_dispatch(_argmax_dispatcher) +def argmax(a, axis=None, out=None, *, keepdims=np._NoValue): + """ + Returns the indices of the maximum values along an axis. + + Parameters + ---------- + a : array_like + Input array. + axis : int, optional + By default, the index is into the flattened array, otherwise + along the specified axis. + out : array, optional + If provided, the result will be inserted into this array. It should + be of the appropriate shape and dtype. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the array. + + .. versionadded:: 1.22.0 + + Returns + ------- + index_array : ndarray of ints + Array of indices into the array. It has the same shape as ``a.shape`` + with the dimension along `axis` removed. If `keepdims` is set to True, + then the size of `axis` will be 1 with the resulting array having same + shape as ``a.shape``. + + See Also + -------- + ndarray.argmax, argmin + amax : The maximum value along a given axis. + unravel_index : Convert a flat index into an index tuple. + take_along_axis : Apply ``np.expand_dims(index_array, axis)`` + from argmax to an array as if by calling max. + + Notes + ----- + In case of multiple occurrences of the maximum values, the indices + corresponding to the first occurrence are returned. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(6).reshape(2,3) + 10 + >>> a + array([[10, 11, 12], + [13, 14, 15]]) + >>> np.argmax(a) + 5 + >>> np.argmax(a, axis=0) + array([1, 1, 1]) + >>> np.argmax(a, axis=1) + array([2, 2]) + + Indexes of the maximal elements of a N-dimensional array: + + >>> ind = np.unravel_index(np.argmax(a, axis=None), a.shape) + >>> ind + (1, 2) + >>> a[ind] + 15 + + >>> b = np.arange(6) + >>> b[1] = 5 + >>> b + array([0, 5, 2, 3, 4, 5]) + >>> np.argmax(b) # Only the first occurrence is returned. + 1 + + >>> x = np.array([[4,2,3], [1,0,3]]) + >>> index_array = np.argmax(x, axis=-1) + >>> # Same as np.amax(x, axis=-1, keepdims=True) + >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1) + array([[4], + [3]]) + >>> # Same as np.amax(x, axis=-1) + >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), + ... axis=-1).squeeze(axis=-1) + array([4, 3]) + + Setting `keepdims` to `True`, + + >>> x = np.arange(24).reshape((2, 3, 4)) + >>> res = np.argmax(x, axis=1, keepdims=True) + >>> res.shape + (2, 1, 4) + """ + kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {} + return _wrapfunc(a, 'argmax', axis=axis, out=out, **kwds) + + +def _argmin_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue): + return (a, out) + + +@array_function_dispatch(_argmin_dispatcher) +def argmin(a, axis=None, out=None, *, keepdims=np._NoValue): + """ + Returns the indices of the minimum values along an axis. + + Parameters + ---------- + a : array_like + Input array. + axis : int, optional + By default, the index is into the flattened array, otherwise + along the specified axis. + out : array, optional + If provided, the result will be inserted into this array. It should + be of the appropriate shape and dtype. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the array. + + .. versionadded:: 1.22.0 + + Returns + ------- + index_array : ndarray of ints + Array of indices into the array. It has the same shape as `a.shape` + with the dimension along `axis` removed. If `keepdims` is set to True, + then the size of `axis` will be 1 with the resulting array having same + shape as `a.shape`. + + See Also + -------- + ndarray.argmin, argmax + amin : The minimum value along a given axis. + unravel_index : Convert a flat index into an index tuple. + take_along_axis : Apply ``np.expand_dims(index_array, axis)`` + from argmin to an array as if by calling min. + + Notes + ----- + In case of multiple occurrences of the minimum values, the indices + corresponding to the first occurrence are returned. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(6).reshape(2,3) + 10 + >>> a + array([[10, 11, 12], + [13, 14, 15]]) + >>> np.argmin(a) + 0 + >>> np.argmin(a, axis=0) + array([0, 0, 0]) + >>> np.argmin(a, axis=1) + array([0, 0]) + + Indices of the minimum elements of a N-dimensional array: + + >>> ind = np.unravel_index(np.argmin(a, axis=None), a.shape) + >>> ind + (0, 0) + >>> a[ind] + 10 + + >>> b = np.arange(6) + 10 + >>> b[4] = 10 + >>> b + array([10, 11, 12, 13, 10, 15]) + >>> np.argmin(b) # Only the first occurrence is returned. + 0 + + >>> x = np.array([[4,2,3], [1,0,3]]) + >>> index_array = np.argmin(x, axis=-1) + >>> # Same as np.amin(x, axis=-1, keepdims=True) + >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1) + array([[2], + [0]]) + >>> # Same as np.amax(x, axis=-1) + >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), + ... axis=-1).squeeze(axis=-1) + array([2, 0]) + + Setting `keepdims` to `True`, + + >>> x = np.arange(24).reshape((2, 3, 4)) + >>> res = np.argmin(x, axis=1, keepdims=True) + >>> res.shape + (2, 1, 4) + """ + kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {} + return _wrapfunc(a, 'argmin', axis=axis, out=out, **kwds) + + +def _searchsorted_dispatcher(a, v, side=None, sorter=None): + return (a, v, sorter) + + +@array_function_dispatch(_searchsorted_dispatcher) +def searchsorted(a, v, side='left', sorter=None): + """ + Find indices where elements should be inserted to maintain order. + + Find the indices into a sorted array `a` such that, if the + corresponding elements in `v` were inserted before the indices, the + order of `a` would be preserved. + + Assuming that `a` is sorted: + + ====== ============================ + `side` returned index `i` satisfies + ====== ============================ + left ``a[i-1] < v <= a[i]`` + right ``a[i-1] <= v < a[i]`` + ====== ============================ + + Parameters + ---------- + a : 1-D array_like + Input array. If `sorter` is None, then it must be sorted in + ascending order, otherwise `sorter` must be an array of indices + that sort it. + v : array_like + Values to insert into `a`. + side : {'left', 'right'}, optional + If 'left', the index of the first suitable location found is given. + If 'right', return the last such index. If there is no suitable + index, return either 0 or N (where N is the length of `a`). + sorter : 1-D array_like, optional + Optional array of integer indices that sort array a into ascending + order. They are typically the result of argsort. + + Returns + ------- + indices : int or array of ints + Array of insertion points with the same shape as `v`, + or an integer if `v` is a scalar. + + See Also + -------- + sort : Return a sorted copy of an array. + histogram : Produce histogram from 1-D data. + + Notes + ----- + Binary search is used to find the required insertion points. + + As of NumPy 1.4.0 `searchsorted` works with real/complex arrays containing + `nan` values. The enhanced sort order is documented in `sort`. + + This function uses the same algorithm as the builtin python + `bisect.bisect_left` (``side='left'``) and `bisect.bisect_right` + (``side='right'``) functions, which is also vectorized + in the `v` argument. + + Examples + -------- + >>> import numpy as np + >>> np.searchsorted([11,12,13,14,15], 13) + 2 + >>> np.searchsorted([11,12,13,14,15], 13, side='right') + 3 + >>> np.searchsorted([11,12,13,14,15], [-10, 20, 12, 13]) + array([0, 5, 1, 2]) + + When `sorter` is used, the returned indices refer to the sorted + array of `a` and not `a` itself: + + >>> a = np.array([40, 10, 20, 30]) + >>> sorter = np.argsort(a) + >>> sorter + array([1, 2, 3, 0]) # Indices that would sort the array 'a' + >>> result = np.searchsorted(a, 25, sorter=sorter) + >>> result + 2 + >>> a[sorter[result]] + 30 # The element at index 2 of the sorted array is 30. + """ + return _wrapfunc(a, 'searchsorted', v, side=side, sorter=sorter) + + +def _resize_dispatcher(a, new_shape): + return (a,) + + +@array_function_dispatch(_resize_dispatcher) +def resize(a, new_shape): + """ + Return a new array with the specified shape. + + If the new array is larger than the original array, then the new + array is filled with repeated copies of `a`. Note that this behavior + is different from a.resize(new_shape) which fills with zeros instead + of repeated copies of `a`. + + Parameters + ---------- + a : array_like + Array to be resized. + + new_shape : int or tuple of int + Shape of resized array. + + Returns + ------- + reshaped_array : ndarray + The new array is formed from the data in the old array, repeated + if necessary to fill out the required number of elements. The + data are repeated iterating over the array in C-order. + + See Also + -------- + numpy.reshape : Reshape an array without changing the total size. + numpy.pad : Enlarge and pad an array. + numpy.repeat : Repeat elements of an array. + ndarray.resize : resize an array in-place. + + Notes + ----- + When the total size of the array does not change `~numpy.reshape` should + be used. In most other cases either indexing (to reduce the size) + or padding (to increase the size) may be a more appropriate solution. + + Warning: This functionality does **not** consider axes separately, + i.e. it does not apply interpolation/extrapolation. + It fills the return array with the required number of elements, iterating + over `a` in C-order, disregarding axes (and cycling back from the start if + the new shape is larger). This functionality is therefore not suitable to + resize images, or data where each axis represents a separate and distinct + entity. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[0,1],[2,3]]) + >>> np.resize(a,(2,3)) + array([[0, 1, 2], + [3, 0, 1]]) + >>> np.resize(a,(1,4)) + array([[0, 1, 2, 3]]) + >>> np.resize(a,(2,4)) + array([[0, 1, 2, 3], + [0, 1, 2, 3]]) + + """ + if isinstance(new_shape, (int, nt.integer)): + new_shape = (new_shape,) + + a = ravel(a) + + new_size = 1 + for dim_length in new_shape: + new_size *= dim_length + if dim_length < 0: + raise ValueError( + 'all elements of `new_shape` must be non-negative' + ) + + if a.size == 0 or new_size == 0: + # First case must zero fill. The second would have repeats == 0. + return np.zeros_like(a, shape=new_shape) + + repeats = -(-new_size // a.size) # ceil division + a = concatenate((a,) * repeats)[:new_size] + + return reshape(a, new_shape) + + +def _squeeze_dispatcher(a, axis=None): + return (a,) + + +@array_function_dispatch(_squeeze_dispatcher) +def squeeze(a, axis=None): + """ + Remove axes of length one from `a`. + + Parameters + ---------- + a : array_like + Input data. + axis : None or int or tuple of ints, optional + Selects a subset of the entries of length one in the + shape. If an axis is selected with shape entry greater than + one, an error is raised. + + Returns + ------- + squeezed : ndarray + The input array, but with all or a subset of the + dimensions of length 1 removed. This is always `a` itself + or a view into `a`. Note that if all axes are squeezed, + the result is a 0d array and not a scalar. + + Raises + ------ + ValueError + If `axis` is not None, and an axis being squeezed is not of length 1 + + See Also + -------- + expand_dims : The inverse operation, adding entries of length one + reshape : Insert, remove, and combine dimensions, and resize existing ones + + Examples + -------- + >>> import numpy as np + >>> x = np.array([[[0], [1], [2]]]) + >>> x.shape + (1, 3, 1) + >>> np.squeeze(x).shape + (3,) + >>> np.squeeze(x, axis=0).shape + (3, 1) + >>> np.squeeze(x, axis=1).shape + Traceback (most recent call last): + ... + ValueError: cannot select an axis to squeeze out which has size + not equal to one + >>> np.squeeze(x, axis=2).shape + (1, 3) + >>> x = np.array([[1234]]) + >>> x.shape + (1, 1) + >>> np.squeeze(x) + array(1234) # 0d array + >>> np.squeeze(x).shape + () + >>> np.squeeze(x)[()] + 1234 + + """ + try: + squeeze = a.squeeze + except AttributeError: + return _wrapit(a, 'squeeze', axis=axis) + if axis is None: + return squeeze() + else: + return squeeze(axis=axis) + + +def _diagonal_dispatcher(a, offset=None, axis1=None, axis2=None): + return (a,) + + +@array_function_dispatch(_diagonal_dispatcher) +def diagonal(a, offset=0, axis1=0, axis2=1): + """ + Return specified diagonals. + + If `a` is 2-D, returns the diagonal of `a` with the given offset, + i.e., the collection of elements of the form ``a[i, i+offset]``. If + `a` has more than two dimensions, then the axes specified by `axis1` + and `axis2` are used to determine the 2-D sub-array whose diagonal is + returned. The shape of the resulting array can be determined by + removing `axis1` and `axis2` and appending an index to the right equal + to the size of the resulting diagonals. + + In versions of NumPy prior to 1.7, this function always returned a new, + independent array containing a copy of the values in the diagonal. + + In NumPy 1.7 and 1.8, it continues to return a copy of the diagonal, + but depending on this fact is deprecated. Writing to the resulting + array continues to work as it used to, but a FutureWarning is issued. + + Starting in NumPy 1.9 it returns a read-only view on the original array. + Attempting to write to the resulting array will produce an error. + + In some future release, it will return a read/write view and writing to + the returned array will alter your original array. The returned array + will have the same type as the input array. + + If you don't write to the array returned by this function, then you can + just ignore all of the above. + + If you depend on the current behavior, then we suggest copying the + returned array explicitly, i.e., use ``np.diagonal(a).copy()`` instead + of just ``np.diagonal(a)``. This will work with both past and future + versions of NumPy. + + Parameters + ---------- + a : array_like + Array from which the diagonals are taken. + offset : int, optional + Offset of the diagonal from the main diagonal. Can be positive or + negative. Defaults to main diagonal (0). + axis1 : int, optional + Axis to be used as the first axis of the 2-D sub-arrays from which + the diagonals should be taken. Defaults to first axis (0). + axis2 : int, optional + Axis to be used as the second axis of the 2-D sub-arrays from + which the diagonals should be taken. Defaults to second axis (1). + + Returns + ------- + array_of_diagonals : ndarray + If `a` is 2-D, then a 1-D array containing the diagonal and of the + same type as `a` is returned unless `a` is a `matrix`, in which case + a 1-D array rather than a (2-D) `matrix` is returned in order to + maintain backward compatibility. + + If ``a.ndim > 2``, then the dimensions specified by `axis1` and `axis2` + are removed, and a new axis inserted at the end corresponding to the + diagonal. + + Raises + ------ + ValueError + If the dimension of `a` is less than 2. + + See Also + -------- + diag : MATLAB work-a-like for 1-D and 2-D arrays. + diagflat : Create diagonal arrays. + trace : Sum along diagonals. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(4).reshape(2,2) + >>> a + array([[0, 1], + [2, 3]]) + >>> a.diagonal() + array([0, 3]) + >>> a.diagonal(1) + array([1]) + + A 3-D example: + + >>> a = np.arange(8).reshape(2,2,2); a + array([[[0, 1], + [2, 3]], + [[4, 5], + [6, 7]]]) + >>> a.diagonal(0, # Main diagonals of two arrays created by skipping + ... 0, # across the outer(left)-most axis last and + ... 1) # the "middle" (row) axis first. + array([[0, 6], + [1, 7]]) + + The sub-arrays whose main diagonals we just obtained; note that each + corresponds to fixing the right-most (column) axis, and that the + diagonals are "packed" in rows. + + >>> a[:,:,0] # main diagonal is [0 6] + array([[0, 2], + [4, 6]]) + >>> a[:,:,1] # main diagonal is [1 7] + array([[1, 3], + [5, 7]]) + + The anti-diagonal can be obtained by reversing the order of elements + using either `numpy.flipud` or `numpy.fliplr`. + + >>> a = np.arange(9).reshape(3, 3) + >>> a + array([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> np.fliplr(a).diagonal() # Horizontal flip + array([2, 4, 6]) + >>> np.flipud(a).diagonal() # Vertical flip + array([6, 4, 2]) + + Note that the order in which the diagonal is retrieved varies depending + on the flip function. + """ + if isinstance(a, np.matrix): + # Make diagonal of matrix 1-D to preserve backward compatibility. + return asarray(a).diagonal(offset=offset, axis1=axis1, axis2=axis2) + else: + return asanyarray(a).diagonal(offset=offset, axis1=axis1, axis2=axis2) + + +def _trace_dispatcher( + a, offset=None, axis1=None, axis2=None, dtype=None, out=None): + return (a, out) + + +@array_function_dispatch(_trace_dispatcher) +def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None): + """ + Return the sum along diagonals of the array. + + If `a` is 2-D, the sum along its diagonal with the given offset + is returned, i.e., the sum of elements ``a[i,i+offset]`` for all i. + + If `a` has more than two dimensions, then the axes specified by axis1 and + axis2 are used to determine the 2-D sub-arrays whose traces are returned. + The shape of the resulting array is the same as that of `a` with `axis1` + and `axis2` removed. + + Parameters + ---------- + a : array_like + Input array, from which the diagonals are taken. + offset : int, optional + Offset of the diagonal from the main diagonal. Can be both positive + and negative. Defaults to 0. + axis1, axis2 : int, optional + Axes to be used as the first and second axis of the 2-D sub-arrays + from which the diagonals should be taken. Defaults are the first two + axes of `a`. + dtype : dtype, optional + Determines the data-type of the returned array and of the accumulator + where the elements are summed. If dtype has the value None and `a` is + of integer type of precision less than the default integer + precision, then the default integer precision is used. Otherwise, + the precision is the same as that of `a`. + out : ndarray, optional + Array into which the output is placed. Its type is preserved and + it must be of the right shape to hold the output. + + Returns + ------- + sum_along_diagonals : ndarray + If `a` is 2-D, the sum along the diagonal is returned. If `a` has + larger dimensions, then an array of sums along diagonals is returned. + + See Also + -------- + diag, diagonal, diagflat + + Examples + -------- + >>> import numpy as np + >>> np.trace(np.eye(3)) + 3.0 + >>> a = np.arange(8).reshape((2,2,2)) + >>> np.trace(a) + array([6, 8]) + + >>> a = np.arange(24).reshape((2,2,2,3)) + >>> np.trace(a).shape + (2, 3) + + """ + if isinstance(a, np.matrix): + # Get trace of matrix via an array to preserve backward compatibility. + return asarray(a).trace( + offset=offset, axis1=axis1, axis2=axis2, dtype=dtype, out=out + ) + else: + return asanyarray(a).trace( + offset=offset, axis1=axis1, axis2=axis2, dtype=dtype, out=out + ) + + +def _ravel_dispatcher(a, order=None): + return (a,) + + +@array_function_dispatch(_ravel_dispatcher) +def ravel(a, order='C'): + """Return a contiguous flattened array. + + A 1-D array, containing the elements of the input, is returned. A copy is + made only if needed. + + As of NumPy 1.10, the returned array will have the same type as the input + array. (for example, a masked array will be returned for a masked array + input) + + Parameters + ---------- + a : array_like + Input array. The elements in `a` are read in the order specified by + `order`, and packed as a 1-D array. + order : {'C','F', 'A', 'K'}, optional + + The elements of `a` are read using this index order. 'C' means + to index the elements in row-major, C-style order, + with the last axis index changing fastest, back to the first + axis index changing slowest. 'F' means to index the elements + in column-major, Fortran-style order, with the + first index changing fastest, and the last index changing + slowest. Note that the 'C' and 'F' options take no account of + the memory layout of the underlying array, and only refer to + the order of axis indexing. 'A' means to read the elements in + Fortran-like index order if `a` is Fortran *contiguous* in + memory, C-like order otherwise. 'K' means to read the + elements in the order they occur in memory, except for + reversing the data when strides are negative. By default, 'C' + index order is used. + + Returns + ------- + y : array_like + y is a contiguous 1-D array of the same subtype as `a`, + with shape ``(a.size,)``. + Note that matrices are special cased for backward compatibility, + if `a` is a matrix, then y is a 1-D ndarray. + + See Also + -------- + ndarray.flat : 1-D iterator over an array. + ndarray.flatten : 1-D array copy of the elements of an array + in row-major order. + ndarray.reshape : Change the shape of an array without changing its data. + + Notes + ----- + In row-major, C-style order, in two dimensions, the row index + varies the slowest, and the column index the quickest. This can + be generalized to multiple dimensions, where row-major order + implies that the index along the first axis varies slowest, and + the index along the last quickest. The opposite holds for + column-major, Fortran-style index ordering. + + When a view is desired in as many cases as possible, ``arr.reshape(-1)`` + may be preferable. However, ``ravel`` supports ``K`` in the optional + ``order`` argument while ``reshape`` does not. + + Examples + -------- + It is equivalent to ``reshape(-1, order=order)``. + + >>> import numpy as np + >>> x = np.array([[1, 2, 3], [4, 5, 6]]) + >>> np.ravel(x) + array([1, 2, 3, 4, 5, 6]) + + >>> x.reshape(-1) + array([1, 2, 3, 4, 5, 6]) + + >>> np.ravel(x, order='F') + array([1, 4, 2, 5, 3, 6]) + + When ``order`` is 'A', it will preserve the array's 'C' or 'F' ordering: + + >>> np.ravel(x.T) + array([1, 4, 2, 5, 3, 6]) + >>> np.ravel(x.T, order='A') + array([1, 2, 3, 4, 5, 6]) + + When ``order`` is 'K', it will preserve orderings that are neither 'C' + nor 'F', but won't reverse axes: + + >>> a = np.arange(3)[::-1]; a + array([2, 1, 0]) + >>> a.ravel(order='C') + array([2, 1, 0]) + >>> a.ravel(order='K') + array([2, 1, 0]) + + >>> a = np.arange(12).reshape(2,3,2).swapaxes(1,2); a + array([[[ 0, 2, 4], + [ 1, 3, 5]], + [[ 6, 8, 10], + [ 7, 9, 11]]]) + >>> a.ravel(order='C') + array([ 0, 2, 4, 1, 3, 5, 6, 8, 10, 7, 9, 11]) + >>> a.ravel(order='K') + array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + + """ + if isinstance(a, np.matrix): + return asarray(a).ravel(order=order) + else: + return asanyarray(a).ravel(order=order) + + +def _nonzero_dispatcher(a): + return (a,) + + +@array_function_dispatch(_nonzero_dispatcher) +def nonzero(a): + """ + Return the indices of the elements that are non-zero. + + Returns a tuple of arrays, one for each dimension of `a`, + containing the indices of the non-zero elements in that + dimension. The values in `a` are always tested and returned in + row-major, C-style order. + + To group the indices by element, rather than dimension, use `argwhere`, + which returns a row for each non-zero element. + + .. note:: + + When called on a zero-d array or scalar, ``nonzero(a)`` is treated + as ``nonzero(atleast_1d(a))``. + + .. deprecated:: 1.17.0 + + Use `atleast_1d` explicitly if this behavior is deliberate. + + Parameters + ---------- + a : array_like + Input array. + + Returns + ------- + tuple_of_arrays : tuple + Indices of elements that are non-zero. + + See Also + -------- + flatnonzero : + Return indices that are non-zero in the flattened version of the input + array. + ndarray.nonzero : + Equivalent ndarray method. + count_nonzero : + Counts the number of non-zero elements in the input array. + + Notes + ----- + While the nonzero values can be obtained with ``a[nonzero(a)]``, it is + recommended to use ``x[x.astype(bool)]`` or ``x[x != 0]`` instead, which + will correctly handle 0-d arrays. + + Examples + -------- + >>> import numpy as np + >>> x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]]) + >>> x + array([[3, 0, 0], + [0, 4, 0], + [5, 6, 0]]) + >>> np.nonzero(x) + (array([0, 1, 2, 2]), array([0, 1, 0, 1])) + + >>> x[np.nonzero(x)] + array([3, 4, 5, 6]) + >>> np.transpose(np.nonzero(x)) + array([[0, 0], + [1, 1], + [2, 0], + [2, 1]]) + + A common use for ``nonzero`` is to find the indices of an array, where + a condition is True. Given an array `a`, the condition `a` > 3 is a + boolean array and since False is interpreted as 0, np.nonzero(a > 3) + yields the indices of the `a` where the condition is true. + + >>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> a > 3 + array([[False, False, False], + [ True, True, True], + [ True, True, True]]) + >>> np.nonzero(a > 3) + (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2])) + + Using this result to index `a` is equivalent to using the mask directly: + + >>> a[np.nonzero(a > 3)] + array([4, 5, 6, 7, 8, 9]) + >>> a[a > 3] # prefer this spelling + array([4, 5, 6, 7, 8, 9]) + + ``nonzero`` can also be called as a method of the array. + + >>> (a > 3).nonzero() + (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2])) + + """ + return _wrapfunc(a, 'nonzero') + + +def _shape_dispatcher(a): + return (a,) + + +@array_function_dispatch(_shape_dispatcher) +def shape(a): + """ + Return the shape of an array. + + Parameters + ---------- + a : array_like + Input array. + + Returns + ------- + shape : tuple of ints + The elements of the shape tuple give the lengths of the + corresponding array dimensions. + + See Also + -------- + len : ``len(a)`` is equivalent to ``np.shape(a)[0]`` for N-D arrays with + ``N>=1``. + ndarray.shape : Equivalent array method. + + Examples + -------- + >>> import numpy as np + >>> np.shape(np.eye(3)) + (3, 3) + >>> np.shape([[1, 3]]) + (1, 2) + >>> np.shape([0]) + (1,) + >>> np.shape(0) + () + + >>> a = np.array([(1, 2), (3, 4), (5, 6)], + ... dtype=[('x', 'i4'), ('y', 'i4')]) + >>> np.shape(a) + (3,) + >>> a.shape + (3,) + + """ + try: + result = a.shape + except AttributeError: + result = asarray(a).shape + return result + + +def _compress_dispatcher(condition, a, axis=None, out=None): + return (condition, a, out) + + +@array_function_dispatch(_compress_dispatcher) +def compress(condition, a, axis=None, out=None): + """ + Return selected slices of an array along given axis. + + When working along a given axis, a slice along that axis is returned in + `output` for each index where `condition` evaluates to True. When + working on a 1-D array, `compress` is equivalent to `extract`. + + Parameters + ---------- + condition : 1-D array of bools + Array that selects which entries to return. If len(condition) + is less than the size of `a` along the given axis, then output is + truncated to the length of the condition array. + a : array_like + Array from which to extract a part. + axis : int, optional + Axis along which to take slices. If None (default), work on the + flattened array. + out : ndarray, optional + Output array. Its type is preserved and it must be of the right + shape to hold the output. + + Returns + ------- + compressed_array : ndarray + A copy of `a` without the slices along axis for which `condition` + is false. + + See Also + -------- + take, choose, diag, diagonal, select + ndarray.compress : Equivalent method in ndarray + extract : Equivalent method when working on 1-D arrays + :ref:`ufuncs-output-type` + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4], [5, 6]]) + >>> a + array([[1, 2], + [3, 4], + [5, 6]]) + >>> np.compress([0, 1], a, axis=0) + array([[3, 4]]) + >>> np.compress([False, True, True], a, axis=0) + array([[3, 4], + [5, 6]]) + >>> np.compress([False, True], a, axis=1) + array([[2], + [4], + [6]]) + + Working on the flattened array does not return slices along an axis but + selects elements. + + >>> np.compress([False, True], a) + array([2]) + + """ + return _wrapfunc(a, 'compress', condition, axis=axis, out=out) + + +def _clip_dispatcher(a, a_min=None, a_max=None, out=None, *, min=None, + max=None, **kwargs): + return (a, a_min, a_max, out, min, max) + + +@array_function_dispatch(_clip_dispatcher) +def clip(a, a_min=np._NoValue, a_max=np._NoValue, out=None, *, + min=np._NoValue, max=np._NoValue, **kwargs): + """ + Clip (limit) the values in an array. + + Given an interval, values outside the interval are clipped to + the interval edges. For example, if an interval of ``[0, 1]`` + is specified, values smaller than 0 become 0, and values larger + than 1 become 1. + + Equivalent to but faster than ``np.minimum(a_max, np.maximum(a, a_min))``. + + No check is performed to ensure ``a_min < a_max``. + + Parameters + ---------- + a : array_like + Array containing elements to clip. + a_min, a_max : array_like or None + Minimum and maximum value. If ``None``, clipping is not performed on + the corresponding edge. If both ``a_min`` and ``a_max`` are ``None``, + the elements of the returned array stay the same. Both are broadcasted + against ``a``. + out : ndarray, optional + The results will be placed in this array. It may be the input + array for in-place clipping. `out` must be of the right shape + to hold the output. Its type is preserved. + min, max : array_like or None + Array API compatible alternatives for ``a_min`` and ``a_max`` + arguments. Either ``a_min`` and ``a_max`` or ``min`` and ``max`` + can be passed at the same time. Default: ``None``. + + .. versionadded:: 2.1.0 + **kwargs + For other keyword-only arguments, see the + :ref:`ufunc docs `. + + Returns + ------- + clipped_array : ndarray + An array with the elements of `a`, but where values + < `a_min` are replaced with `a_min`, and those > `a_max` + with `a_max`. + + See Also + -------- + :ref:`ufuncs-output-type` + + Notes + ----- + When `a_min` is greater than `a_max`, `clip` returns an + array in which all values are equal to `a_max`, + as shown in the second example. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(10) + >>> a + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> np.clip(a, 1, 8) + array([1, 1, 2, 3, 4, 5, 6, 7, 8, 8]) + >>> np.clip(a, 8, 1) + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1]) + >>> np.clip(a, 3, 6, out=a) + array([3, 3, 3, 3, 4, 5, 6, 6, 6, 6]) + >>> a + array([3, 3, 3, 3, 4, 5, 6, 6, 6, 6]) + >>> a = np.arange(10) + >>> a + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> np.clip(a, [3, 4, 1, 1, 1, 4, 4, 4, 4, 4], 8) + array([3, 4, 2, 3, 4, 5, 6, 7, 8, 8]) + + """ + if a_min is np._NoValue and a_max is np._NoValue: + a_min = None if min is np._NoValue else min + a_max = None if max is np._NoValue else max + elif a_min is np._NoValue: + raise TypeError("clip() missing 1 required positional " + "argument: 'a_min'") + elif a_max is np._NoValue: + raise TypeError("clip() missing 1 required positional " + "argument: 'a_max'") + elif min is not np._NoValue or max is not np._NoValue: + raise ValueError("Passing `min` or `max` keyword argument when " + "`a_min` and `a_max` are provided is forbidden.") + + return _wrapfunc(a, 'clip', a_min, a_max, out=out, **kwargs) + + +def _sum_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, + initial=None, where=None): + return (a, out) + + +@array_function_dispatch(_sum_dispatcher) +def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + initial=np._NoValue, where=np._NoValue): + """ + Sum of array elements over a given axis. + + Parameters + ---------- + a : array_like + Elements to sum. + axis : None or int or tuple of ints, optional + Axis or axes along which a sum is performed. The default, + axis=None, will sum all of the elements of the input array. If + axis is negative it counts from the last to the first axis. If + axis is a tuple of ints, a sum is performed on all of the axes + specified in the tuple instead of a single axis or all the axes as + before. + dtype : dtype, optional + The type of the returned array and of the accumulator in which the + elements are summed. The dtype of `a` is used by default unless `a` + has an integer dtype of less precision than the default platform + integer. In that case, if `a` is signed then the platform integer + is used while if `a` is unsigned then an unsigned integer of the + same precision as the platform integer is used. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape as the expected output, but the type of the output + values will be cast if necessary. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `sum` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + initial : scalar, optional + Starting value for the sum. See `~numpy.ufunc.reduce` for details. + where : array_like of bool, optional + Elements to include in the sum. See `~numpy.ufunc.reduce` for details. + + Returns + ------- + sum_along_axis : ndarray + An array with the same shape as `a`, with the specified + axis removed. If `a` is a 0-d array, or if `axis` is None, a scalar + is returned. If an output array is specified, a reference to + `out` is returned. + + See Also + -------- + ndarray.sum : Equivalent method. + add: ``numpy.add.reduce`` equivalent function. + cumsum : Cumulative sum of array elements. + trapezoid : Integration of array values using composite trapezoidal rule. + + mean, average + + Notes + ----- + Arithmetic is modular when using integer types, and no error is + raised on overflow. + + The sum of an empty array is the neutral element 0: + + >>> np.sum([]) + 0.0 + + For floating point numbers the numerical precision of sum (and + ``np.add.reduce``) is in general limited by directly adding each number + individually to the result causing rounding errors in every step. + However, often numpy will use a numerically better approach (partial + pairwise summation) leading to improved precision in many use-cases. + This improved precision is always provided when no ``axis`` is given. + When ``axis`` is given, it will depend on which axis is summed. + Technically, to provide the best speed possible, the improved precision + is only used when the summation is along the fast axis in memory. + Note that the exact precision may vary depending on other parameters. + In contrast to NumPy, Python's ``math.fsum`` function uses a slower but + more precise approach to summation. + Especially when summing a large number of lower precision floating point + numbers, such as ``float32``, numerical errors can become significant. + In such cases it can be advisable to use `dtype="float64"` to use a higher + precision for the output. + + Examples + -------- + >>> import numpy as np + >>> np.sum([0.5, 1.5]) + 2.0 + >>> np.sum([0.5, 0.7, 0.2, 1.5], dtype=np.int32) + np.int32(1) + >>> np.sum([[0, 1], [0, 5]]) + 6 + >>> np.sum([[0, 1], [0, 5]], axis=0) + array([0, 6]) + >>> np.sum([[0, 1], [0, 5]], axis=1) + array([1, 5]) + >>> np.sum([[0, 1], [np.nan, 5]], where=[False, True], axis=1) + array([1., 5.]) + + If the accumulator is too small, overflow occurs: + + >>> np.ones(128, dtype=np.int8).sum(dtype=np.int8) + np.int8(-128) + + You can also start the sum with a value other than zero: + + >>> np.sum([10], initial=5) + 15 + """ + if isinstance(a, _gentype): + # 2018-02-25, 1.15.0 + warnings.warn( + "Calling np.sum(generator) is deprecated, and in the future will " + "give a different result. Use np.sum(np.fromiter(generator)) or " + "the python sum builtin instead.", + DeprecationWarning, stacklevel=2 + ) + + res = _sum_(a) + if out is not None: + out[...] = res + return out + return res + + return _wrapreduction( + a, np.add, 'sum', axis, dtype, out, + keepdims=keepdims, initial=initial, where=where + ) + + +def _any_dispatcher(a, axis=None, out=None, keepdims=None, *, + where=np._NoValue): + return (a, where, out) + + +@array_function_dispatch(_any_dispatcher) +def any(a, axis=None, out=None, keepdims=np._NoValue, *, where=np._NoValue): + """ + Test whether any array element along a given axis evaluates to True. + + Returns single boolean if `axis` is ``None`` + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array. + axis : None or int or tuple of ints, optional + Axis or axes along which a logical OR reduction is performed. + The default (``axis=None``) is to perform a logical OR over all + the dimensions of the input array. `axis` may be negative, in + which case it counts from the last to the first axis. If this + is a tuple of ints, a reduction is performed on multiple + axes, instead of a single axis or all the axes as before. + out : ndarray, optional + Alternate output array in which to place the result. It must have + the same shape as the expected output and its type is preserved + (e.g., if it is of type float, then it will remain so, returning + 1.0 for True and 0.0 for False, regardless of the type of `a`). + See :ref:`ufuncs-output-type` for more details. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `any` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + + where : array_like of bool, optional + Elements to include in checking for any `True` values. + See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.20.0 + + Returns + ------- + any : bool or ndarray + A new boolean or `ndarray` is returned unless `out` is specified, + in which case a reference to `out` is returned. + + See Also + -------- + ndarray.any : equivalent method + + all : Test whether all elements along a given axis evaluate to True. + + Notes + ----- + Not a Number (NaN), positive infinity and negative infinity evaluate + to `True` because these are not equal to zero. + + .. versionchanged:: 2.0 + Before NumPy 2.0, ``any`` did not return booleans for object dtype + input arrays. + This behavior is still available via ``np.logical_or.reduce``. + + Examples + -------- + >>> import numpy as np + >>> np.any([[True, False], [True, True]]) + True + + >>> np.any([[True, False, True ], + ... [False, False, False]], axis=0) + array([ True, False, True]) + + >>> np.any([-1, 0, 5]) + True + + >>> np.any([[np.nan], [np.inf]], axis=1, keepdims=True) + array([[ True], + [ True]]) + + >>> np.any([[True, False], [False, False]], where=[[False], [True]]) + False + + >>> a = np.array([[1, 0, 0], + ... [0, 0, 1], + ... [0, 0, 0]]) + >>> np.any(a, axis=0) + array([ True, False, True]) + >>> np.any(a, axis=1) + array([ True, True, False]) + + >>> o=np.array(False) + >>> z=np.any([-1, 4, 5], out=o) + >>> z, o + (array(True), array(True)) + >>> # Check now that z is a reference to o + >>> z is o + True + >>> id(z), id(o) # identity of z and o # doctest: +SKIP + (191614240, 191614240) + + """ + return _wrapreduction_any_all(a, np.logical_or, 'any', axis, out, + keepdims=keepdims, where=where) + + +def _all_dispatcher(a, axis=None, out=None, keepdims=None, *, + where=None): + return (a, where, out) + + +@array_function_dispatch(_all_dispatcher) +def all(a, axis=None, out=None, keepdims=np._NoValue, *, where=np._NoValue): + """ + Test whether all array elements along a given axis evaluate to True. + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array. + axis : None or int or tuple of ints, optional + Axis or axes along which a logical AND reduction is performed. + The default (``axis=None``) is to perform a logical AND over all + the dimensions of the input array. `axis` may be negative, in + which case it counts from the last to the first axis. If this + is a tuple of ints, a reduction is performed on multiple + axes, instead of a single axis or all the axes as before. + out : ndarray, optional + Alternate output array in which to place the result. + It must have the same shape as the expected output and its + type is preserved (e.g., if ``dtype(out)`` is float, the result + will consist of 0.0's and 1.0's). See :ref:`ufuncs-output-type` + for more details. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `all` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + + where : array_like of bool, optional + Elements to include in checking for all `True` values. + See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.20.0 + + Returns + ------- + all : ndarray, bool + A new boolean or array is returned unless `out` is specified, + in which case a reference to `out` is returned. + + See Also + -------- + ndarray.all : equivalent method + + any : Test whether any element along a given axis evaluates to True. + + Notes + ----- + Not a Number (NaN), positive infinity and negative infinity + evaluate to `True` because these are not equal to zero. + + .. versionchanged:: 2.0 + Before NumPy 2.0, ``all`` did not return booleans for object dtype + input arrays. + This behavior is still available via ``np.logical_and.reduce``. + + Examples + -------- + >>> import numpy as np + >>> np.all([[True,False],[True,True]]) + False + + >>> np.all([[True,False],[True,True]], axis=0) + array([ True, False]) + + >>> np.all([-1, 4, 5]) + True + + >>> np.all([1.0, np.nan]) + True + + >>> np.all([[True, True], [False, True]], where=[[True], [False]]) + True + + >>> o=np.array(False) + >>> z=np.all([-1, 4, 5], out=o) + >>> id(z), id(o), z + (28293632, 28293632, array(True)) # may vary + + """ + return _wrapreduction_any_all(a, np.logical_and, 'all', axis, out, + keepdims=keepdims, where=where) + + +def _cumulative_func(x, func, axis, dtype, out, include_initial): + x = np.atleast_1d(x) + x_ndim = x.ndim + if axis is None: + if x_ndim >= 2: + raise ValueError("For arrays which have more than one dimension " + "``axis`` argument is required.") + axis = 0 + + if out is not None and include_initial: + item = [slice(None)] * x_ndim + item[axis] = slice(1, None) + func.accumulate(x, axis=axis, dtype=dtype, out=out[tuple(item)]) + item[axis] = 0 + out[tuple(item)] = func.identity + return out + + res = func.accumulate(x, axis=axis, dtype=dtype, out=out) + if include_initial: + initial_shape = list(x.shape) + initial_shape[axis] = 1 + res = np.concat( + [np.full_like(res, func.identity, shape=initial_shape), res], + axis=axis, + ) + + return res + + +def _cumulative_prod_dispatcher(x, /, *, axis=None, dtype=None, out=None, + include_initial=None): + return (x, out) + + +@array_function_dispatch(_cumulative_prod_dispatcher) +def cumulative_prod(x, /, *, axis=None, dtype=None, out=None, + include_initial=False): + """ + Return the cumulative product of elements along a given axis. + + This function is an Array API compatible alternative to `numpy.cumprod`. + + Parameters + ---------- + x : array_like + Input array. + axis : int, optional + Axis along which the cumulative product is computed. The default + (None) is only allowed for one-dimensional arrays. For arrays + with more than one dimension ``axis`` is required. + dtype : dtype, optional + Type of the returned array, as well as of the accumulator in which + the elements are multiplied. If ``dtype`` is not specified, it + defaults to the dtype of ``x``, unless ``x`` has an integer dtype + with a precision less than that of the default platform integer. + In that case, the default platform integer is used instead. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output + but the type of the resulting values will be cast if necessary. + See :ref:`ufuncs-output-type` for more details. + include_initial : bool, optional + Boolean indicating whether to include the initial value (ones) as + the first value in the output. With ``include_initial=True`` + the shape of the output is different than the shape of the input. + Default: ``False``. + + Returns + ------- + cumulative_prod_along_axis : ndarray + A new array holding the result is returned unless ``out`` is + specified, in which case a reference to ``out`` is returned. The + result has the same shape as ``x`` if ``include_initial=False``. + + Notes + ----- + Arithmetic is modular when using integer types, and no error is + raised on overflow. + + Examples + -------- + >>> a = np.array([1, 2, 3]) + >>> np.cumulative_prod(a) # intermediate results 1, 1*2 + ... # total product 1*2*3 = 6 + array([1, 2, 6]) + >>> a = np.array([1, 2, 3, 4, 5, 6]) + >>> np.cumulative_prod(a, dtype=float) # specify type of output + array([ 1., 2., 6., 24., 120., 720.]) + + The cumulative product for each column (i.e., over the rows) of ``b``: + + >>> b = np.array([[1, 2, 3], [4, 5, 6]]) + >>> np.cumulative_prod(b, axis=0) + array([[ 1, 2, 3], + [ 4, 10, 18]]) + + The cumulative product for each row (i.e. over the columns) of ``b``: + + >>> np.cumulative_prod(b, axis=1) + array([[ 1, 2, 6], + [ 4, 20, 120]]) + + """ + return _cumulative_func(x, um.multiply, axis, dtype, out, include_initial) + + +def _cumulative_sum_dispatcher(x, /, *, axis=None, dtype=None, out=None, + include_initial=None): + return (x, out) + + +@array_function_dispatch(_cumulative_sum_dispatcher) +def cumulative_sum(x, /, *, axis=None, dtype=None, out=None, + include_initial=False): + """ + Return the cumulative sum of the elements along a given axis. + + This function is an Array API compatible alternative to `numpy.cumsum`. + + Parameters + ---------- + x : array_like + Input array. + axis : int, optional + Axis along which the cumulative sum is computed. The default + (None) is only allowed for one-dimensional arrays. For arrays + with more than one dimension ``axis`` is required. + dtype : dtype, optional + Type of the returned array and of the accumulator in which the + elements are summed. If ``dtype`` is not specified, it defaults + to the dtype of ``x``, unless ``x`` has an integer dtype with + a precision less than that of the default platform integer. + In that case, the default platform integer is used. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output + but the type will be cast if necessary. See :ref:`ufuncs-output-type` + for more details. + include_initial : bool, optional + Boolean indicating whether to include the initial value (zeros) as + the first value in the output. With ``include_initial=True`` + the shape of the output is different than the shape of the input. + Default: ``False``. + + Returns + ------- + cumulative_sum_along_axis : ndarray + A new array holding the result is returned unless ``out`` is + specified, in which case a reference to ``out`` is returned. The + result has the same shape as ``x`` if ``include_initial=False``. + + See Also + -------- + sum : Sum array elements. + trapezoid : Integration of array values using composite trapezoidal rule. + diff : Calculate the n-th discrete difference along given axis. + + Notes + ----- + Arithmetic is modular when using integer types, and no error is + raised on overflow. + + ``cumulative_sum(a)[-1]`` may not be equal to ``sum(a)`` for + floating-point values since ``sum`` may use a pairwise summation routine, + reducing the roundoff-error. See `sum` for more information. + + Examples + -------- + >>> a = np.array([1, 2, 3, 4, 5, 6]) + >>> a + array([1, 2, 3, 4, 5, 6]) + >>> np.cumulative_sum(a) + array([ 1, 3, 6, 10, 15, 21]) + >>> np.cumulative_sum(a, dtype=float) # specifies type of output value(s) + array([ 1., 3., 6., 10., 15., 21.]) + + >>> b = np.array([[1, 2, 3], [4, 5, 6]]) + >>> np.cumulative_sum(b,axis=0) # sum over rows for each of the 3 columns + array([[1, 2, 3], + [5, 7, 9]]) + >>> np.cumulative_sum(b,axis=1) # sum over columns for each of the 2 rows + array([[ 1, 3, 6], + [ 4, 9, 15]]) + + ``cumulative_sum(c)[-1]`` may not be equal to ``sum(c)`` + + >>> c = np.array([1, 2e-9, 3e-9] * 1000000) + >>> np.cumulative_sum(c)[-1] + 1000000.0050045159 + >>> c.sum() + 1000000.0050000029 + + """ + return _cumulative_func(x, um.add, axis, dtype, out, include_initial) + + +def _cumsum_dispatcher(a, axis=None, dtype=None, out=None): + return (a, out) + + +@array_function_dispatch(_cumsum_dispatcher) +def cumsum(a, axis=None, dtype=None, out=None): + """ + Return the cumulative sum of the elements along a given axis. + + Parameters + ---------- + a : array_like + Input array. + axis : int, optional + Axis along which the cumulative sum is computed. The default + (None) is to compute the cumsum over the flattened array. + dtype : dtype, optional + Type of the returned array and of the accumulator in which the + elements are summed. If `dtype` is not specified, it defaults + to the dtype of `a`, unless `a` has an integer dtype with a + precision less than that of the default platform integer. In + that case, the default platform integer is used. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output + but the type will be cast if necessary. See :ref:`ufuncs-output-type` + for more details. + + Returns + ------- + cumsum_along_axis : ndarray. + A new array holding the result is returned unless `out` is + specified, in which case a reference to `out` is returned. The + result has the same size as `a`, and the same shape as `a` if + `axis` is not None or `a` is a 1-d array. + + See Also + -------- + cumulative_sum : Array API compatible alternative for ``cumsum``. + sum : Sum array elements. + trapezoid : Integration of array values using composite trapezoidal rule. + diff : Calculate the n-th discrete difference along given axis. + + Notes + ----- + Arithmetic is modular when using integer types, and no error is + raised on overflow. + + ``cumsum(a)[-1]`` may not be equal to ``sum(a)`` for floating-point + values since ``sum`` may use a pairwise summation routine, reducing + the roundoff-error. See `sum` for more information. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1,2,3], [4,5,6]]) + >>> a + array([[1, 2, 3], + [4, 5, 6]]) + >>> np.cumsum(a) + array([ 1, 3, 6, 10, 15, 21]) + >>> np.cumsum(a, dtype=float) # specifies type of output value(s) + array([ 1., 3., 6., 10., 15., 21.]) + + >>> np.cumsum(a,axis=0) # sum over rows for each of the 3 columns + array([[1, 2, 3], + [5, 7, 9]]) + >>> np.cumsum(a,axis=1) # sum over columns for each of the 2 rows + array([[ 1, 3, 6], + [ 4, 9, 15]]) + + ``cumsum(b)[-1]`` may not be equal to ``sum(b)`` + + >>> b = np.array([1, 2e-9, 3e-9] * 1000000) + >>> b.cumsum()[-1] + 1000000.0050045159 + >>> b.sum() + 1000000.0050000029 + + """ + return _wrapfunc(a, 'cumsum', axis=axis, dtype=dtype, out=out) + + +def _ptp_dispatcher(a, axis=None, out=None, keepdims=None): + return (a, out) + + +@array_function_dispatch(_ptp_dispatcher) +def ptp(a, axis=None, out=None, keepdims=np._NoValue): + """ + Range of values (maximum - minimum) along an axis. + + The name of the function comes from the acronym for 'peak to peak'. + + .. warning:: + `ptp` preserves the data type of the array. This means the + return value for an input of signed integers with n bits + (e.g. `numpy.int8`, `numpy.int16`, etc) is also a signed integer + with n bits. In that case, peak-to-peak values greater than + ``2**(n-1)-1`` will be returned as negative values. An example + with a work-around is shown below. + + Parameters + ---------- + a : array_like + Input values. + axis : None or int or tuple of ints, optional + Axis along which to find the peaks. By default, flatten the + array. `axis` may be negative, in + which case it counts from the last to the first axis. + If this is a tuple of ints, a reduction is performed on multiple + axes, instead of a single axis or all the axes as before. + out : array_like + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output, + but the type of the output values will be cast if necessary. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `ptp` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + + Returns + ------- + ptp : ndarray or scalar + The range of a given array - `scalar` if array is one-dimensional + or a new array holding the result along the given axis + + Examples + -------- + >>> import numpy as np + >>> x = np.array([[4, 9, 2, 10], + ... [6, 9, 7, 12]]) + + >>> np.ptp(x, axis=1) + array([8, 6]) + + >>> np.ptp(x, axis=0) + array([2, 0, 5, 2]) + + >>> np.ptp(x) + 10 + + This example shows that a negative value can be returned when + the input is an array of signed integers. + + >>> y = np.array([[1, 127], + ... [0, 127], + ... [-1, 127], + ... [-2, 127]], dtype=np.int8) + >>> np.ptp(y, axis=1) + array([ 126, 127, -128, -127], dtype=int8) + + A work-around is to use the `view()` method to view the result as + unsigned integers with the same bit width: + + >>> np.ptp(y, axis=1).view(np.uint8) + array([126, 127, 128, 129], dtype=uint8) + + """ + kwargs = {} + if keepdims is not np._NoValue: + kwargs['keepdims'] = keepdims + return _methods._ptp(a, axis=axis, out=out, **kwargs) + + +def _max_dispatcher(a, axis=None, out=None, keepdims=None, initial=None, + where=None): + return (a, out) + + +@array_function_dispatch(_max_dispatcher) +@set_module('numpy') +def max(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): + """ + Return the maximum of an array or maximum along an axis. + + Parameters + ---------- + a : array_like + Input data. + axis : None or int or tuple of ints, optional + Axis or axes along which to operate. By default, flattened input is + used. If this is a tuple of ints, the maximum is selected over + multiple axes, instead of a single axis or all the axes as before. + + out : ndarray, optional + Alternative output array in which to place the result. Must + be of the same shape and buffer length as the expected output. + See :ref:`ufuncs-output-type` for more details. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the ``max`` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + + initial : scalar, optional + The minimum value of an output element. Must be present to allow + computation on empty slice. See `~numpy.ufunc.reduce` for details. + + where : array_like of bool, optional + Elements to compare for the maximum. See `~numpy.ufunc.reduce` + for details. + + Returns + ------- + max : ndarray or scalar + Maximum of `a`. If `axis` is None, the result is a scalar value. + If `axis` is an int, the result is an array of dimension + ``a.ndim - 1``. If `axis` is a tuple, the result is an array of + dimension ``a.ndim - len(axis)``. + + See Also + -------- + amin : + The minimum value of an array along a given axis, propagating any NaNs. + nanmax : + The maximum value of an array along a given axis, ignoring any NaNs. + maximum : + Element-wise maximum of two arrays, propagating any NaNs. + fmax : + Element-wise maximum of two arrays, ignoring any NaNs. + argmax : + Return the indices of the maximum values. + + nanmin, minimum, fmin + + Notes + ----- + NaN values are propagated, that is if at least one item is NaN, the + corresponding max value will be NaN as well. To ignore NaN values + (MATLAB behavior), please use nanmax. + + Don't use `~numpy.max` for element-wise comparison of 2 arrays; when + ``a.shape[0]`` is 2, ``maximum(a[0], a[1])`` is faster than + ``max(a, axis=0)``. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(4).reshape((2,2)) + >>> a + array([[0, 1], + [2, 3]]) + >>> np.max(a) # Maximum of the flattened array + 3 + >>> np.max(a, axis=0) # Maxima along the first axis + array([2, 3]) + >>> np.max(a, axis=1) # Maxima along the second axis + array([1, 3]) + >>> np.max(a, where=[False, True], initial=-1, axis=0) + array([-1, 3]) + >>> b = np.arange(5, dtype=float) + >>> b[2] = np.nan + >>> np.max(b) + np.float64(nan) + >>> np.max(b, where=~np.isnan(b), initial=-1) + 4.0 + >>> np.nanmax(b) + 4.0 + + You can use an initial value to compute the maximum of an empty slice, or + to initialize it to a different value: + + >>> np.max([[-50], [10]], axis=-1, initial=0) + array([ 0, 10]) + + Notice that the initial value is used as one of the elements for which the + maximum is determined, unlike for the default argument Python's max + function, which is only used for empty iterables. + + >>> np.max([5], initial=6) + 6 + >>> max([5], default=6) + 5 + """ + return _wrapreduction(a, np.maximum, 'max', axis, None, out, + keepdims=keepdims, initial=initial, where=where) + + +@array_function_dispatch(_max_dispatcher) +def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): + """ + Return the maximum of an array or maximum along an axis. + + `amax` is an alias of `~numpy.max`. + + See Also + -------- + max : alias of this function + ndarray.max : equivalent method + """ + return _wrapreduction(a, np.maximum, 'max', axis, None, out, + keepdims=keepdims, initial=initial, where=where) + + +def _min_dispatcher(a, axis=None, out=None, keepdims=None, initial=None, + where=None): + return (a, out) + + +@array_function_dispatch(_min_dispatcher) +def min(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): + """ + Return the minimum of an array or minimum along an axis. + + Parameters + ---------- + a : array_like + Input data. + axis : None or int or tuple of ints, optional + Axis or axes along which to operate. By default, flattened input is + used. + + If this is a tuple of ints, the minimum is selected over multiple axes, + instead of a single axis or all the axes as before. + out : ndarray, optional + Alternative output array in which to place the result. Must + be of the same shape and buffer length as the expected output. + See :ref:`ufuncs-output-type` for more details. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the ``min`` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + + initial : scalar, optional + The maximum value of an output element. Must be present to allow + computation on empty slice. See `~numpy.ufunc.reduce` for details. + + where : array_like of bool, optional + Elements to compare for the minimum. See `~numpy.ufunc.reduce` + for details. + + Returns + ------- + min : ndarray or scalar + Minimum of `a`. If `axis` is None, the result is a scalar value. + If `axis` is an int, the result is an array of dimension + ``a.ndim - 1``. If `axis` is a tuple, the result is an array of + dimension ``a.ndim - len(axis)``. + + See Also + -------- + amax : + The maximum value of an array along a given axis, propagating any NaNs. + nanmin : + The minimum value of an array along a given axis, ignoring any NaNs. + minimum : + Element-wise minimum of two arrays, propagating any NaNs. + fmin : + Element-wise minimum of two arrays, ignoring any NaNs. + argmin : + Return the indices of the minimum values. + + nanmax, maximum, fmax + + Notes + ----- + NaN values are propagated, that is if at least one item is NaN, the + corresponding min value will be NaN as well. To ignore NaN values + (MATLAB behavior), please use nanmin. + + Don't use `~numpy.min` for element-wise comparison of 2 arrays; when + ``a.shape[0]`` is 2, ``minimum(a[0], a[1])`` is faster than + ``min(a, axis=0)``. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(4).reshape((2,2)) + >>> a + array([[0, 1], + [2, 3]]) + >>> np.min(a) # Minimum of the flattened array + 0 + >>> np.min(a, axis=0) # Minima along the first axis + array([0, 1]) + >>> np.min(a, axis=1) # Minima along the second axis + array([0, 2]) + >>> np.min(a, where=[False, True], initial=10, axis=0) + array([10, 1]) + + >>> b = np.arange(5, dtype=float) + >>> b[2] = np.nan + >>> np.min(b) + np.float64(nan) + >>> np.min(b, where=~np.isnan(b), initial=10) + 0.0 + >>> np.nanmin(b) + 0.0 + + >>> np.min([[-50], [10]], axis=-1, initial=0) + array([-50, 0]) + + Notice that the initial value is used as one of the elements for which the + minimum is determined, unlike for the default argument Python's max + function, which is only used for empty iterables. + + Notice that this isn't the same as Python's ``default`` argument. + + >>> np.min([6], initial=5) + 5 + >>> min([6], default=5) + 6 + """ + return _wrapreduction(a, np.minimum, 'min', axis, None, out, + keepdims=keepdims, initial=initial, where=where) + + +@array_function_dispatch(_min_dispatcher) +def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): + """ + Return the minimum of an array or minimum along an axis. + + `amin` is an alias of `~numpy.min`. + + See Also + -------- + min : alias of this function + ndarray.min : equivalent method + """ + return _wrapreduction(a, np.minimum, 'min', axis, None, out, + keepdims=keepdims, initial=initial, where=where) + + +def _prod_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, + initial=None, where=None): + return (a, out) + + +@array_function_dispatch(_prod_dispatcher) +def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + initial=np._NoValue, where=np._NoValue): + """ + Return the product of array elements over a given axis. + + Parameters + ---------- + a : array_like + Input data. + axis : None or int or tuple of ints, optional + Axis or axes along which a product is performed. The default, + axis=None, will calculate the product of all the elements in the + input array. If axis is negative it counts from the last to the + first axis. + + If axis is a tuple of ints, a product is performed on all of the + axes specified in the tuple instead of a single axis or all the + axes as before. + dtype : dtype, optional + The type of the returned array, as well as of the accumulator in + which the elements are multiplied. The dtype of `a` is used by + default unless `a` has an integer dtype of less precision than the + default platform integer. In that case, if `a` is signed then the + platform integer is used while if `a` is unsigned then an unsigned + integer of the same precision as the platform integer is used. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape as the expected output, but the type of the output + values will be cast if necessary. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `prod` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + initial : scalar, optional + The starting value for this product. See `~numpy.ufunc.reduce` + for details. + where : array_like of bool, optional + Elements to include in the product. See `~numpy.ufunc.reduce` + for details. + + Returns + ------- + product_along_axis : ndarray, see `dtype` parameter above. + An array shaped as `a` but with the specified axis removed. + Returns a reference to `out` if specified. + + See Also + -------- + ndarray.prod : equivalent method + :ref:`ufuncs-output-type` + + Notes + ----- + Arithmetic is modular when using integer types, and no error is + raised on overflow. That means that, on a 32-bit platform: + + >>> x = np.array([536870910, 536870910, 536870910, 536870910]) + >>> np.prod(x) + 16 # may vary + + The product of an empty array is the neutral element 1: + + >>> np.prod([]) + 1.0 + + Examples + -------- + By default, calculate the product of all elements: + + >>> import numpy as np + >>> np.prod([1.,2.]) + 2.0 + + Even when the input array is two-dimensional: + + >>> a = np.array([[1., 2.], [3., 4.]]) + >>> np.prod(a) + 24.0 + + But we can also specify the axis over which to multiply: + + >>> np.prod(a, axis=1) + array([ 2., 12.]) + >>> np.prod(a, axis=0) + array([3., 8.]) + + Or select specific elements to include: + + >>> np.prod([1., np.nan, 3.], where=[True, False, True]) + 3.0 + + If the type of `x` is unsigned, then the output type is + the unsigned platform integer: + + >>> x = np.array([1, 2, 3], dtype=np.uint8) + >>> np.prod(x).dtype == np.uint + True + + If `x` is of a signed integer type, then the output type + is the default platform integer: + + >>> x = np.array([1, 2, 3], dtype=np.int8) + >>> np.prod(x).dtype == int + True + + You can also start the product with a value other than one: + + >>> np.prod([1, 2], initial=5) + 10 + """ + return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out, + keepdims=keepdims, initial=initial, where=where) + + +def _cumprod_dispatcher(a, axis=None, dtype=None, out=None): + return (a, out) + + +@array_function_dispatch(_cumprod_dispatcher) +def cumprod(a, axis=None, dtype=None, out=None): + """ + Return the cumulative product of elements along a given axis. + + Parameters + ---------- + a : array_like + Input array. + axis : int, optional + Axis along which the cumulative product is computed. By default + the input is flattened. + dtype : dtype, optional + Type of the returned array, as well as of the accumulator in which + the elements are multiplied. If *dtype* is not specified, it + defaults to the dtype of `a`, unless `a` has an integer dtype with + a precision less than that of the default platform integer. In + that case, the default platform integer is used instead. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output + but the type of the resulting values will be cast if necessary. + + Returns + ------- + cumprod : ndarray + A new array holding the result is returned unless `out` is + specified, in which case a reference to out is returned. + + See Also + -------- + cumulative_prod : Array API compatible alternative for ``cumprod``. + :ref:`ufuncs-output-type` + + Notes + ----- + Arithmetic is modular when using integer types, and no error is + raised on overflow. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1,2,3]) + >>> np.cumprod(a) # intermediate results 1, 1*2 + ... # total product 1*2*3 = 6 + array([1, 2, 6]) + >>> a = np.array([[1, 2, 3], [4, 5, 6]]) + >>> np.cumprod(a, dtype=float) # specify type of output + array([ 1., 2., 6., 24., 120., 720.]) + + The cumulative product for each column (i.e., over the rows) of `a`: + + >>> np.cumprod(a, axis=0) + array([[ 1, 2, 3], + [ 4, 10, 18]]) + + The cumulative product for each row (i.e. over the columns) of `a`: + + >>> np.cumprod(a,axis=1) + array([[ 1, 2, 6], + [ 4, 20, 120]]) + + """ + return _wrapfunc(a, 'cumprod', axis=axis, dtype=dtype, out=out) + + +def _ndim_dispatcher(a): + return (a,) + + +@array_function_dispatch(_ndim_dispatcher) +def ndim(a): + """ + Return the number of dimensions of an array. + + Parameters + ---------- + a : array_like + Input array. If it is not already an ndarray, a conversion is + attempted. + + Returns + ------- + number_of_dimensions : int + The number of dimensions in `a`. Scalars are zero-dimensional. + + See Also + -------- + ndarray.ndim : equivalent method + shape : dimensions of array + ndarray.shape : dimensions of array + + Examples + -------- + >>> import numpy as np + >>> np.ndim([[1,2,3],[4,5,6]]) + 2 + >>> np.ndim(np.array([[1,2,3],[4,5,6]])) + 2 + >>> np.ndim(1) + 0 + + """ + try: + return a.ndim + except AttributeError: + return asarray(a).ndim + + +def _size_dispatcher(a, axis=None): + return (a,) + + +@array_function_dispatch(_size_dispatcher) +def size(a, axis=None): + """ + Return the number of elements along a given axis. + + Parameters + ---------- + a : array_like + Input data. + axis : int, optional + Axis along which the elements are counted. By default, give + the total number of elements. + + Returns + ------- + element_count : int + Number of elements along the specified axis. + + See Also + -------- + shape : dimensions of array + ndarray.shape : dimensions of array + ndarray.size : number of elements in array + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1,2,3],[4,5,6]]) + >>> np.size(a) + 6 + >>> np.size(a,1) + 3 + >>> np.size(a,0) + 2 + + """ + if axis is None: + try: + return a.size + except AttributeError: + return asarray(a).size + else: + try: + return a.shape[axis] + except AttributeError: + return asarray(a).shape[axis] + + +def _round_dispatcher(a, decimals=None, out=None): + return (a, out) + + +@array_function_dispatch(_round_dispatcher) +def round(a, decimals=0, out=None): + """ + Evenly round to the given number of decimals. + + Parameters + ---------- + a : array_like + Input data. + decimals : int, optional + Number of decimal places to round to (default: 0). If + decimals is negative, it specifies the number of positions to + the left of the decimal point. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape as the expected output, but the type of the output + values will be cast if necessary. See :ref:`ufuncs-output-type` + for more details. + + Returns + ------- + rounded_array : ndarray + An array of the same type as `a`, containing the rounded values. + Unless `out` was specified, a new array is created. A reference to + the result is returned. + + The real and imaginary parts of complex numbers are rounded + separately. The result of rounding a float is a float. + + See Also + -------- + ndarray.round : equivalent method + around : an alias for this function + ceil, fix, floor, rint, trunc + + + Notes + ----- + For values exactly halfway between rounded decimal values, NumPy + rounds to the nearest even value. Thus 1.5 and 2.5 round to 2.0, + -0.5 and 0.5 round to 0.0, etc. + + ``np.round`` uses a fast but sometimes inexact algorithm to round + floating-point datatypes. For positive `decimals` it is equivalent to + ``np.true_divide(np.rint(a * 10**decimals), 10**decimals)``, which has + error due to the inexact representation of decimal fractions in the IEEE + floating point standard [1]_ and errors introduced when scaling by powers + of ten. For instance, note the extra "1" in the following: + + >>> np.round(56294995342131.5, 3) + 56294995342131.51 + + If your goal is to print such values with a fixed number of decimals, it is + preferable to use numpy's float printing routines to limit the number of + printed decimals: + + >>> np.format_float_positional(56294995342131.5, precision=3) + '56294995342131.5' + + The float printing routines use an accurate but much more computationally + demanding algorithm to compute the number of digits after the decimal + point. + + Alternatively, Python's builtin `round` function uses a more accurate + but slower algorithm for 64-bit floating point values: + + >>> round(56294995342131.5, 3) + 56294995342131.5 + >>> np.round(16.055, 2), round(16.055, 2) # equals 16.0549999999999997 + (16.06, 16.05) + + + References + ---------- + .. [1] "Lecture Notes on the Status of IEEE 754", William Kahan, + https://people.eecs.berkeley.edu/~wkahan/ieee754status/IEEE754.PDF + + Examples + -------- + >>> import numpy as np + >>> np.round([0.37, 1.64]) + array([0., 2.]) + >>> np.round([0.37, 1.64], decimals=1) + array([0.4, 1.6]) + >>> np.round([.5, 1.5, 2.5, 3.5, 4.5]) # rounds to nearest even value + array([0., 2., 2., 4., 4.]) + >>> np.round([1,2,3,11], decimals=1) # ndarray of ints is returned + array([ 1, 2, 3, 11]) + >>> np.round([1,2,3,11], decimals=-1) + array([ 0, 0, 0, 10]) + + """ + return _wrapfunc(a, 'round', decimals=decimals, out=out) + + +@array_function_dispatch(_round_dispatcher) +def around(a, decimals=0, out=None): + """ + Round an array to the given number of decimals. + + `around` is an alias of `~numpy.round`. + + See Also + -------- + ndarray.round : equivalent method + round : alias for this function + ceil, fix, floor, rint, trunc + + """ + return _wrapfunc(a, 'round', decimals=decimals, out=out) + + +def _mean_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, *, + where=None): + return (a, where, out) + + +@array_function_dispatch(_mean_dispatcher) +def mean(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, *, + where=np._NoValue): + """ + Compute the arithmetic mean along the specified axis. + + Returns the average of the array elements. The average is taken over + the flattened array by default, otherwise over the specified axis. + `float64` intermediate and return values are used for integer inputs. + + Parameters + ---------- + a : array_like + Array containing numbers whose mean is desired. If `a` is not an + array, a conversion is attempted. + axis : None or int or tuple of ints, optional + Axis or axes along which the means are computed. The default is to + compute the mean of the flattened array. + + If this is a tuple of ints, a mean is performed over multiple axes, + instead of a single axis or all the axes as before. + dtype : data-type, optional + Type to use in computing the mean. For integer inputs, the default + is `float64`; for floating point inputs, it is the same as the + input dtype. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output, but the type will be cast if necessary. + See :ref:`ufuncs-output-type` for more details. + See :ref:`ufuncs-output-type` for more details. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `mean` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + + where : array_like of bool, optional + Elements to include in the mean. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.20.0 + + Returns + ------- + m : ndarray, see dtype parameter above + If `out=None`, returns a new array containing the mean values, + otherwise a reference to the output array is returned. + + See Also + -------- + average : Weighted average + std, var, nanmean, nanstd, nanvar + + Notes + ----- + The arithmetic mean is the sum of the elements along the axis divided + by the number of elements. + + Note that for floating-point input, the mean is computed using the + same precision the input has. Depending on the input data, this can + cause the results to be inaccurate, especially for `float32` (see + example below). Specifying a higher-precision accumulator using the + `dtype` keyword can alleviate this issue. + + By default, `float16` results are computed using `float32` intermediates + for extra precision. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4]]) + >>> np.mean(a) + 2.5 + >>> np.mean(a, axis=0) + array([2., 3.]) + >>> np.mean(a, axis=1) + array([1.5, 3.5]) + + In single precision, `mean` can be inaccurate: + + >>> a = np.zeros((2, 512*512), dtype=np.float32) + >>> a[0, :] = 1.0 + >>> a[1, :] = 0.1 + >>> np.mean(a) + np.float32(0.54999924) + + Computing the mean in float64 is more accurate: + + >>> np.mean(a, dtype=np.float64) + 0.55000000074505806 # may vary + + Computing the mean in timedelta64 is available: + + >>> b = np.array([1, 3], dtype="timedelta64[D]") + >>> np.mean(b) + np.timedelta64(2,'D') + + Specifying a where argument: + + >>> a = np.array([[5, 9, 13], [14, 10, 12], [11, 15, 19]]) + >>> np.mean(a) + 12.0 + >>> np.mean(a, where=[[True], [False], [False]]) + 9.0 + + """ + kwargs = {} + if keepdims is not np._NoValue: + kwargs['keepdims'] = keepdims + if where is not np._NoValue: + kwargs['where'] = where + if type(a) is not mu.ndarray: + try: + mean = a.mean + except AttributeError: + pass + else: + return mean(axis=axis, dtype=dtype, out=out, **kwargs) + + return _methods._mean(a, axis=axis, dtype=dtype, + out=out, **kwargs) + + +def _std_dispatcher(a, axis=None, dtype=None, out=None, ddof=None, + keepdims=None, *, where=None, mean=None, correction=None): + return (a, where, out, mean) + + +@array_function_dispatch(_std_dispatcher) +def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=np._NoValue, *, + where=np._NoValue, mean=np._NoValue, correction=np._NoValue): + r""" + Compute the standard deviation along the specified axis. + + Returns the standard deviation, a measure of the spread of a distribution, + of the array elements. The standard deviation is computed for the + flattened array by default, otherwise over the specified axis. + + Parameters + ---------- + a : array_like + Calculate the standard deviation of these values. + axis : None or int or tuple of ints, optional + Axis or axes along which the standard deviation is computed. The + default is to compute the standard deviation of the flattened array. + If this is a tuple of ints, a standard deviation is performed over + multiple axes, instead of a single axis or all the axes as before. + dtype : dtype, optional + Type to use in computing the standard deviation. For arrays of + integer type the default is float64, for arrays of float types it is + the same as the array type. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape as the expected output but the type (of the calculated + values) will be cast if necessary. + See :ref:`ufuncs-output-type` for more details. + ddof : {int, float}, optional + Means Delta Degrees of Freedom. The divisor used in calculations + is ``N - ddof``, where ``N`` represents the number of elements. + By default `ddof` is zero. See Notes for details about use of `ddof`. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `std` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + where : array_like of bool, optional + Elements to include in the standard deviation. + See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.20.0 + + mean : array_like, optional + Provide the mean to prevent its recalculation. The mean should have + a shape as if it was calculated with ``keepdims=True``. + The axis for the calculation of the mean should be the same as used in + the call to this std function. + + .. versionadded:: 2.0.0 + + correction : {int, float}, optional + Array API compatible name for the ``ddof`` parameter. Only one of them + can be provided at the same time. + + .. versionadded:: 2.0.0 + + Returns + ------- + standard_deviation : ndarray, see dtype parameter above. + If `out` is None, return a new array containing the standard deviation, + otherwise return a reference to the output array. + + See Also + -------- + var, mean, nanmean, nanstd, nanvar + :ref:`ufuncs-output-type` + + Notes + ----- + There are several common variants of the array standard deviation + calculation. Assuming the input `a` is a one-dimensional NumPy array + and ``mean`` is either provided as an argument or computed as + ``a.mean()``, NumPy computes the standard deviation of an array as:: + + N = len(a) + d2 = abs(a - mean)**2 # abs is for complex `a` + var = d2.sum() / (N - ddof) # note use of `ddof` + std = var**0.5 + + Different values of the argument `ddof` are useful in different + contexts. NumPy's default ``ddof=0`` corresponds with the expression: + + .. math:: + + \sqrt{\frac{\sum_i{|a_i - \bar{a}|^2 }}{N}} + + which is sometimes called the "population standard deviation" in the field + of statistics because it applies the definition of standard deviation to + `a` as if `a` were a complete population of possible observations. + + Many other libraries define the standard deviation of an array + differently, e.g.: + + .. math:: + + \sqrt{\frac{\sum_i{|a_i - \bar{a}|^2 }}{N - 1}} + + In statistics, the resulting quantity is sometimes called the "sample + standard deviation" because if `a` is a random sample from a larger + population, this calculation provides the square root of an unbiased + estimate of the variance of the population. The use of :math:`N-1` in the + denominator is often called "Bessel's correction" because it corrects for + bias (toward lower values) in the variance estimate introduced when the + sample mean of `a` is used in place of the true mean of the population. + The resulting estimate of the standard deviation is still biased, but less + than it would have been without the correction. For this quantity, use + ``ddof=1``. + + Note that, for complex numbers, `std` takes the absolute + value before squaring, so that the result is always real and nonnegative. + + For floating-point input, the standard deviation is computed using the same + precision the input has. Depending on the input data, this can cause + the results to be inaccurate, especially for float32 (see example below). + Specifying a higher-accuracy accumulator using the `dtype` keyword can + alleviate this issue. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4]]) + >>> np.std(a) + 1.1180339887498949 # may vary + >>> np.std(a, axis=0) + array([1., 1.]) + >>> np.std(a, axis=1) + array([0.5, 0.5]) + + In single precision, std() can be inaccurate: + + >>> a = np.zeros((2, 512*512), dtype=np.float32) + >>> a[0, :] = 1.0 + >>> a[1, :] = 0.1 + >>> np.std(a) + np.float32(0.45000005) + + Computing the standard deviation in float64 is more accurate: + + >>> np.std(a, dtype=np.float64) + 0.44999999925494177 # may vary + + Specifying a where argument: + + >>> a = np.array([[14, 8, 11, 10], [7, 9, 10, 11], [10, 15, 5, 10]]) + >>> np.std(a) + 2.614064523559687 # may vary + >>> np.std(a, where=[[True], [True], [False]]) + 2.0 + + Using the mean keyword to save computation time: + + >>> import numpy as np + >>> from timeit import timeit + >>> a = np.array([[14, 8, 11, 10], [7, 9, 10, 11], [10, 15, 5, 10]]) + >>> mean = np.mean(a, axis=1, keepdims=True) + >>> + >>> g = globals() + >>> n = 10000 + >>> t1 = timeit("std = np.std(a, axis=1, mean=mean)", globals=g, number=n) + >>> t2 = timeit("std = np.std(a, axis=1)", globals=g, number=n) + >>> print(f'Percentage execution time saved {100*(t2-t1)/t2:.0f}%') + #doctest: +SKIP + Percentage execution time saved 30% + + """ + kwargs = {} + if keepdims is not np._NoValue: + kwargs['keepdims'] = keepdims + if where is not np._NoValue: + kwargs['where'] = where + if mean is not np._NoValue: + kwargs['mean'] = mean + + if correction != np._NoValue: + if ddof != 0: + raise ValueError( + "ddof and correction can't be provided simultaneously." + ) + else: + ddof = correction + + if type(a) is not mu.ndarray: + try: + std = a.std + except AttributeError: + pass + else: + return std(axis=axis, dtype=dtype, out=out, ddof=ddof, **kwargs) + + return _methods._std(a, axis=axis, dtype=dtype, out=out, ddof=ddof, + **kwargs) + + +def _var_dispatcher(a, axis=None, dtype=None, out=None, ddof=None, + keepdims=None, *, where=None, mean=None, correction=None): + return (a, where, out, mean) + + +@array_function_dispatch(_var_dispatcher) +def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=np._NoValue, *, + where=np._NoValue, mean=np._NoValue, correction=np._NoValue): + r""" + Compute the variance along the specified axis. + + Returns the variance of the array elements, a measure of the spread of a + distribution. The variance is computed for the flattened array by + default, otherwise over the specified axis. + + Parameters + ---------- + a : array_like + Array containing numbers whose variance is desired. If `a` is not an + array, a conversion is attempted. + axis : None or int or tuple of ints, optional + Axis or axes along which the variance is computed. The default is to + compute the variance of the flattened array. + If this is a tuple of ints, a variance is performed over multiple axes, + instead of a single axis or all the axes as before. + dtype : data-type, optional + Type to use in computing the variance. For arrays of integer type + the default is `float64`; for arrays of float types it is the same as + the array type. + out : ndarray, optional + Alternate output array in which to place the result. It must have + the same shape as the expected output, but the type is cast if + necessary. + ddof : {int, float}, optional + "Delta Degrees of Freedom": the divisor used in the calculation is + ``N - ddof``, where ``N`` represents the number of elements. By + default `ddof` is zero. See notes for details about use of `ddof`. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + If the default value is passed, then `keepdims` will not be + passed through to the `var` method of sub-classes of + `ndarray`, however any non-default value will be. If the + sub-class' method does not implement `keepdims` any + exceptions will be raised. + where : array_like of bool, optional + Elements to include in the variance. See `~numpy.ufunc.reduce` for + details. + + .. versionadded:: 1.20.0 + + mean : array like, optional + Provide the mean to prevent its recalculation. The mean should have + a shape as if it was calculated with ``keepdims=True``. + The axis for the calculation of the mean should be the same as used in + the call to this var function. + + .. versionadded:: 2.0.0 + + correction : {int, float}, optional + Array API compatible name for the ``ddof`` parameter. Only one of them + can be provided at the same time. + + .. versionadded:: 2.0.0 + + Returns + ------- + variance : ndarray, see dtype parameter above + If ``out=None``, returns a new array containing the variance; + otherwise, a reference to the output array is returned. + + See Also + -------- + std, mean, nanmean, nanstd, nanvar + :ref:`ufuncs-output-type` + + Notes + ----- + There are several common variants of the array variance calculation. + Assuming the input `a` is a one-dimensional NumPy array and ``mean`` is + either provided as an argument or computed as ``a.mean()``, NumPy + computes the variance of an array as:: + + N = len(a) + d2 = abs(a - mean)**2 # abs is for complex `a` + var = d2.sum() / (N - ddof) # note use of `ddof` + + Different values of the argument `ddof` are useful in different + contexts. NumPy's default ``ddof=0`` corresponds with the expression: + + .. math:: + + \frac{\sum_i{|a_i - \bar{a}|^2 }}{N} + + which is sometimes called the "population variance" in the field of + statistics because it applies the definition of variance to `a` as if `a` + were a complete population of possible observations. + + Many other libraries define the variance of an array differently, e.g.: + + .. math:: + + \frac{\sum_i{|a_i - \bar{a}|^2}}{N - 1} + + In statistics, the resulting quantity is sometimes called the "sample + variance" because if `a` is a random sample from a larger population, + this calculation provides an unbiased estimate of the variance of the + population. The use of :math:`N-1` in the denominator is often called + "Bessel's correction" because it corrects for bias (toward lower values) + in the variance estimate introduced when the sample mean of `a` is used + in place of the true mean of the population. For this quantity, use + ``ddof=1``. + + Note that for complex numbers, the absolute value is taken before + squaring, so that the result is always real and nonnegative. + + For floating-point input, the variance is computed using the same + precision the input has. Depending on the input data, this can cause + the results to be inaccurate, especially for `float32` (see example + below). Specifying a higher-accuracy accumulator using the ``dtype`` + keyword can alleviate this issue. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4]]) + >>> np.var(a) + 1.25 + >>> np.var(a, axis=0) + array([1., 1.]) + >>> np.var(a, axis=1) + array([0.25, 0.25]) + + In single precision, var() can be inaccurate: + + >>> a = np.zeros((2, 512*512), dtype=np.float32) + >>> a[0, :] = 1.0 + >>> a[1, :] = 0.1 + >>> np.var(a) + np.float32(0.20250003) + + Computing the variance in float64 is more accurate: + + >>> np.var(a, dtype=np.float64) + 0.20249999932944759 # may vary + >>> ((1-0.55)**2 + (0.1-0.55)**2)/2 + 0.2025 + + Specifying a where argument: + + >>> a = np.array([[14, 8, 11, 10], [7, 9, 10, 11], [10, 15, 5, 10]]) + >>> np.var(a) + 6.833333333333333 # may vary + >>> np.var(a, where=[[True], [True], [False]]) + 4.0 + + Using the mean keyword to save computation time: + + >>> import numpy as np + >>> from timeit import timeit + >>> + >>> a = np.array([[14, 8, 11, 10], [7, 9, 10, 11], [10, 15, 5, 10]]) + >>> mean = np.mean(a, axis=1, keepdims=True) + >>> + >>> g = globals() + >>> n = 10000 + >>> t1 = timeit("var = np.var(a, axis=1, mean=mean)", globals=g, number=n) + >>> t2 = timeit("var = np.var(a, axis=1)", globals=g, number=n) + >>> print(f'Percentage execution time saved {100*(t2-t1)/t2:.0f}%') + #doctest: +SKIP + Percentage execution time saved 32% + + """ + kwargs = {} + if keepdims is not np._NoValue: + kwargs['keepdims'] = keepdims + if where is not np._NoValue: + kwargs['where'] = where + if mean is not np._NoValue: + kwargs['mean'] = mean + + if correction != np._NoValue: + if ddof != 0: + raise ValueError( + "ddof and correction can't be provided simultaneously." + ) + else: + ddof = correction + + if type(a) is not mu.ndarray: + try: + var = a.var + + except AttributeError: + pass + else: + return var(axis=axis, dtype=dtype, out=out, ddof=ddof, **kwargs) + + return _methods._var(a, axis=axis, dtype=dtype, out=out, ddof=ddof, + **kwargs) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/fromnumeric.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/fromnumeric.pyi new file mode 100644 index 0000000000000000000000000000000000000000..31f763c34d768058e5ef84e805d3c21fdb5e37e2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/fromnumeric.pyi @@ -0,0 +1,1733 @@ +# ruff: noqa: ANN401 +from collections.abc import Sequence +from typing import ( + Any, + Literal, + Protocol, + SupportsIndex, + TypeAlias, + TypeVar, + overload, + type_check_only, +) + +from _typeshed import Incomplete +from typing_extensions import Never, deprecated + +import numpy as np +from numpy import ( + number, + uint64, + int_, + int64, + intp, + float16, + floating, + complexfloating, + timedelta64, + object_, + generic, + _AnyShapeType, + _OrderKACF, + _OrderACF, + _ModeKind, + _PartitionKind, + _SortKind, + _SortSide, + _CastingKind, +) +from numpy._globals import _NoValueType +from numpy._typing import ( + DTypeLike, + _DTypeLike, + ArrayLike, + _ArrayLike, + NDArray, + _NestedSequence, + _ShapeLike, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, + _IntLike_co, + _BoolLike_co, + _ComplexLike_co, + _NumberLike_co, + _ScalarLike_co, +) + +__all__ = [ + "all", + "amax", + "amin", + "any", + "argmax", + "argmin", + "argpartition", + "argsort", + "around", + "choose", + "clip", + "compress", + "cumprod", + "cumsum", + "cumulative_prod", + "cumulative_sum", + "diagonal", + "mean", + "max", + "min", + "matrix_transpose", + "ndim", + "nonzero", + "partition", + "prod", + "ptp", + "put", + "ravel", + "repeat", + "reshape", + "resize", + "round", + "searchsorted", + "shape", + "size", + "sort", + "squeeze", + "std", + "sum", + "swapaxes", + "take", + "trace", + "transpose", + "var", +] + +_SCT = TypeVar("_SCT", bound=generic) +_SCT_uifcO = TypeVar("_SCT_uifcO", bound=number[Any] | object_) +_ArrayT = TypeVar("_ArrayT", bound=np.ndarray[Any, Any]) +_ShapeType = TypeVar("_ShapeType", bound=tuple[int, ...]) +_ShapeType_co = TypeVar("_ShapeType_co", bound=tuple[int, ...], covariant=True) + +@type_check_only +class _SupportsShape(Protocol[_ShapeType_co]): + # NOTE: it matters that `self` is positional only + @property + def shape(self, /) -> _ShapeType_co: ... + +# a "sequence" that isn't a string, bytes, bytearray, or memoryview +_T = TypeVar("_T") +_PyArray: TypeAlias = list[_T] | tuple[_T, ...] +# `int` also covers `bool` +_PyScalar: TypeAlias = float | complex | bytes | str + +@overload +def take( + a: _ArrayLike[_SCT], + indices: _IntLike_co, + axis: None = ..., + out: None = ..., + mode: _ModeKind = ..., +) -> _SCT: ... +@overload +def take( + a: ArrayLike, + indices: _IntLike_co, + axis: SupportsIndex | None = ..., + out: None = ..., + mode: _ModeKind = ..., +) -> Any: ... +@overload +def take( + a: _ArrayLike[_SCT], + indices: _ArrayLikeInt_co, + axis: SupportsIndex | None = ..., + out: None = ..., + mode: _ModeKind = ..., +) -> NDArray[_SCT]: ... +@overload +def take( + a: ArrayLike, + indices: _ArrayLikeInt_co, + axis: SupportsIndex | None = ..., + out: None = ..., + mode: _ModeKind = ..., +) -> NDArray[Any]: ... +@overload +def take( + a: ArrayLike, + indices: _ArrayLikeInt_co, + axis: SupportsIndex | None, + out: _ArrayT, + mode: _ModeKind = ..., +) -> _ArrayT: ... +@overload +def take( + a: ArrayLike, + indices: _ArrayLikeInt_co, + axis: SupportsIndex | None = ..., + *, + out: _ArrayT, + mode: _ModeKind = ..., +) -> _ArrayT: ... + +@overload +def reshape( # shape: index + a: _ArrayLike[_SCT], + /, + shape: SupportsIndex, + order: _OrderACF = "C", + *, + copy: bool | None = None, +) -> np.ndarray[tuple[int], np.dtype[_SCT]]: ... +@overload +def reshape( # shape: (int, ...) @ _AnyShapeType + a: _ArrayLike[_SCT], + /, + shape: _AnyShapeType, + order: _OrderACF = "C", + *, + copy: bool | None = None, +) -> np.ndarray[_AnyShapeType, np.dtype[_SCT]]: ... +@overload # shape: Sequence[index] +def reshape( + a: _ArrayLike[_SCT], + /, + shape: Sequence[SupportsIndex], + order: _OrderACF = "C", + *, + copy: bool | None = None, +) -> NDArray[_SCT]: ... +@overload # shape: index +def reshape( + a: ArrayLike, + /, + shape: SupportsIndex, + order: _OrderACF = "C", + *, + copy: bool | None = None, +) -> np.ndarray[tuple[int], np.dtype[Any]]: ... +@overload +def reshape( # shape: (int, ...) @ _AnyShapeType + a: ArrayLike, + /, + shape: _AnyShapeType, + order: _OrderACF = "C", + *, + copy: bool | None = None, +) -> np.ndarray[_AnyShapeType, np.dtype[Any]]: ... +@overload # shape: Sequence[index] +def reshape( + a: ArrayLike, + /, + shape: Sequence[SupportsIndex], + order: _OrderACF = "C", + *, + copy: bool | None = None, +) -> NDArray[Any]: ... +@overload +@deprecated( + "`newshape` keyword argument is deprecated, " + "use `shape=...` or pass shape positionally instead. " + "(deprecated in NumPy 2.1)", +) +def reshape( + a: ArrayLike, + /, + shape: None = None, + order: _OrderACF = "C", + *, + newshape: _ShapeLike, + copy: bool | None = None, +) -> NDArray[Any]: ... + +@overload +def choose( + a: _IntLike_co, + choices: ArrayLike, + out: None = ..., + mode: _ModeKind = ..., +) -> Any: ... +@overload +def choose( + a: _ArrayLikeInt_co, + choices: _ArrayLike[_SCT], + out: None = ..., + mode: _ModeKind = ..., +) -> NDArray[_SCT]: ... +@overload +def choose( + a: _ArrayLikeInt_co, + choices: ArrayLike, + out: None = ..., + mode: _ModeKind = ..., +) -> NDArray[Any]: ... +@overload +def choose( + a: _ArrayLikeInt_co, + choices: ArrayLike, + out: _ArrayT, + mode: _ModeKind = ..., +) -> _ArrayT: ... + +@overload +def repeat( + a: _ArrayLike[_SCT], + repeats: _ArrayLikeInt_co, + axis: SupportsIndex | None = ..., +) -> NDArray[_SCT]: ... +@overload +def repeat( + a: ArrayLike, + repeats: _ArrayLikeInt_co, + axis: SupportsIndex | None = ..., +) -> NDArray[Any]: ... + +def put( + a: NDArray[Any], + ind: _ArrayLikeInt_co, + v: ArrayLike, + mode: _ModeKind = ..., +) -> None: ... + +@overload +def swapaxes( + a: _ArrayLike[_SCT], + axis1: SupportsIndex, + axis2: SupportsIndex, +) -> NDArray[_SCT]: ... +@overload +def swapaxes( + a: ArrayLike, + axis1: SupportsIndex, + axis2: SupportsIndex, +) -> NDArray[Any]: ... + +@overload +def transpose( + a: _ArrayLike[_SCT], + axes: _ShapeLike | None = ... +) -> NDArray[_SCT]: ... +@overload +def transpose( + a: ArrayLike, + axes: _ShapeLike | None = ... +) -> NDArray[Any]: ... + +@overload +def matrix_transpose(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ... +@overload +def matrix_transpose(x: ArrayLike, /) -> NDArray[Any]: ... + +# +@overload +def partition( + a: _ArrayLike[_SCT], + kth: _ArrayLikeInt, + axis: SupportsIndex | None = -1, + kind: _PartitionKind = "introselect", + order: None = None, +) -> NDArray[_SCT]: ... +@overload +def partition( + a: _ArrayLike[np.void], + kth: _ArrayLikeInt, + axis: SupportsIndex | None = -1, + kind: _PartitionKind = "introselect", + order: str | Sequence[str] | None = None, +) -> NDArray[np.void]: ... +@overload +def partition( + a: ArrayLike, + kth: _ArrayLikeInt, + axis: SupportsIndex | None = -1, + kind: _PartitionKind = "introselect", + order: str | Sequence[str] | None = None, +) -> NDArray[Any]: ... + +# +def argpartition( + a: ArrayLike, + kth: _ArrayLikeInt, + axis: SupportsIndex | None = -1, + kind: _PartitionKind = "introselect", + order: str | Sequence[str] | None = None, +) -> NDArray[intp]: ... + +# +@overload +def sort( + a: _ArrayLike[_SCT], + axis: SupportsIndex | None = ..., + kind: _SortKind | None = ..., + order: str | Sequence[str] | None = ..., + *, + stable: bool | None = ..., +) -> NDArray[_SCT]: ... +@overload +def sort( + a: ArrayLike, + axis: SupportsIndex | None = ..., + kind: _SortKind | None = ..., + order: str | Sequence[str] | None = ..., + *, + stable: bool | None = ..., +) -> NDArray[Any]: ... + +def argsort( + a: ArrayLike, + axis: SupportsIndex | None = ..., + kind: _SortKind | None = ..., + order: str | Sequence[str] | None = ..., + *, + stable: bool | None = ..., +) -> NDArray[intp]: ... + +@overload +def argmax( + a: ArrayLike, + axis: None = ..., + out: None = ..., + *, + keepdims: Literal[False] = ..., +) -> intp: ... +@overload +def argmax( + a: ArrayLike, + axis: SupportsIndex | None = ..., + out: None = ..., + *, + keepdims: bool = ..., +) -> Any: ... +@overload +def argmax( + a: ArrayLike, + axis: SupportsIndex | None, + out: _ArrayT, + *, + keepdims: bool = ..., +) -> _ArrayT: ... +@overload +def argmax( + a: ArrayLike, + axis: SupportsIndex | None = ..., + *, + out: _ArrayT, + keepdims: bool = ..., +) -> _ArrayT: ... + +@overload +def argmin( + a: ArrayLike, + axis: None = ..., + out: None = ..., + *, + keepdims: Literal[False] = ..., +) -> intp: ... +@overload +def argmin( + a: ArrayLike, + axis: SupportsIndex | None = ..., + out: None = ..., + *, + keepdims: bool = ..., +) -> Any: ... +@overload +def argmin( + a: ArrayLike, + axis: SupportsIndex | None, + out: _ArrayT, + *, + keepdims: bool = ..., +) -> _ArrayT: ... +@overload +def argmin( + a: ArrayLike, + axis: SupportsIndex | None = ..., + *, + out: _ArrayT, + keepdims: bool = ..., +) -> _ArrayT: ... + +@overload +def searchsorted( + a: ArrayLike, + v: _ScalarLike_co, + side: _SortSide = ..., + sorter: _ArrayLikeInt_co | None = ..., # 1D int array +) -> intp: ... +@overload +def searchsorted( + a: ArrayLike, + v: ArrayLike, + side: _SortSide = ..., + sorter: _ArrayLikeInt_co | None = ..., # 1D int array +) -> NDArray[intp]: ... + +# +@overload +def resize(a: _ArrayLike[_SCT], new_shape: SupportsIndex | tuple[SupportsIndex]) -> np.ndarray[tuple[int], np.dtype[_SCT]]: ... +@overload +def resize(a: _ArrayLike[_SCT], new_shape: _AnyShapeType) -> np.ndarray[_AnyShapeType, np.dtype[_SCT]]: ... +@overload +def resize(a: _ArrayLike[_SCT], new_shape: _ShapeLike) -> NDArray[_SCT]: ... +@overload +def resize(a: ArrayLike, new_shape: SupportsIndex | tuple[SupportsIndex]) -> np.ndarray[tuple[int], np.dtype[Any]]: ... +@overload +def resize(a: ArrayLike, new_shape: _AnyShapeType) -> np.ndarray[_AnyShapeType, np.dtype[Any]]: ... +@overload +def resize(a: ArrayLike, new_shape: _ShapeLike) -> NDArray[Any]: ... + +@overload +def squeeze( + a: _SCT, + axis: _ShapeLike | None = ..., +) -> _SCT: ... +@overload +def squeeze( + a: _ArrayLike[_SCT], + axis: _ShapeLike | None = ..., +) -> NDArray[_SCT]: ... +@overload +def squeeze( + a: ArrayLike, + axis: _ShapeLike | None = ..., +) -> NDArray[Any]: ... + +@overload +def diagonal( + a: _ArrayLike[_SCT], + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., # >= 2D array +) -> NDArray[_SCT]: ... +@overload +def diagonal( + a: ArrayLike, + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., # >= 2D array +) -> NDArray[Any]: ... + +@overload +def trace( + a: ArrayLike, # >= 2D array + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., +) -> Any: ... +@overload +def trace( + a: ArrayLike, # >= 2D array + offset: SupportsIndex, + axis1: SupportsIndex, + axis2: SupportsIndex, + dtype: DTypeLike, + out: _ArrayT, +) -> _ArrayT: ... +@overload +def trace( + a: ArrayLike, # >= 2D array + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + dtype: DTypeLike = ..., + *, + out: _ArrayT, +) -> _ArrayT: ... + +_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_SCT]] + +@overload +def ravel(a: _ArrayLike[_SCT], order: _OrderKACF = "C") -> _Array1D[_SCT]: ... +@overload +def ravel(a: bytes | _NestedSequence[bytes], order: _OrderKACF = "C") -> _Array1D[np.bytes_]: ... +@overload +def ravel(a: str | _NestedSequence[str], order: _OrderKACF = "C") -> _Array1D[np.str_]: ... +@overload +def ravel(a: bool | _NestedSequence[bool], order: _OrderKACF = "C") -> _Array1D[np.bool]: ... +@overload +def ravel(a: int | _NestedSequence[int], order: _OrderKACF = "C") -> _Array1D[np.int_ | np.bool]: ... +@overload +def ravel(a: float | _NestedSequence[float], order: _OrderKACF = "C") -> _Array1D[np.float64 | np.int_ | np.bool]: ... +@overload +def ravel( + a: complex | _NestedSequence[complex], + order: _OrderKACF = "C", +) -> _Array1D[np.complex128 | np.float64 | np.int_ | np.bool]: ... +@overload +def ravel(a: ArrayLike, order: _OrderKACF = "C") -> np.ndarray[tuple[int], np.dtype[Any]]: ... + +def nonzero(a: _ArrayLike[Any]) -> tuple[NDArray[intp], ...]: ... + +# this prevents `Any` from being returned with Pyright +@overload +def shape(a: _SupportsShape[Never]) -> tuple[int, ...]: ... +@overload +def shape(a: _SupportsShape[_ShapeType]) -> _ShapeType: ... +@overload +def shape(a: _PyScalar) -> tuple[()]: ... +# `collections.abc.Sequence` can't be used hesre, since `bytes` and `str` are +# subtypes of it, which would make the return types incompatible. +@overload +def shape(a: _PyArray[_PyScalar]) -> tuple[int]: ... +@overload +def shape(a: _PyArray[_PyArray[_PyScalar]]) -> tuple[int, int]: ... +# this overload will be skipped by typecheckers that don't support PEP 688 +@overload +def shape(a: memoryview | bytearray) -> tuple[int]: ... +@overload +def shape(a: ArrayLike) -> tuple[int, ...]: ... + +@overload +def compress( + condition: _ArrayLikeBool_co, # 1D bool array + a: _ArrayLike[_SCT], + axis: SupportsIndex | None = ..., + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def compress( + condition: _ArrayLikeBool_co, # 1D bool array + a: ArrayLike, + axis: SupportsIndex | None = ..., + out: None = ..., +) -> NDArray[Any]: ... +@overload +def compress( + condition: _ArrayLikeBool_co, # 1D bool array + a: ArrayLike, + axis: SupportsIndex | None, + out: _ArrayT, +) -> _ArrayT: ... +@overload +def compress( + condition: _ArrayLikeBool_co, # 1D bool array + a: ArrayLike, + axis: SupportsIndex | None = ..., + *, + out: _ArrayT, +) -> _ArrayT: ... + +@overload +def clip( + a: _SCT, + a_min: ArrayLike | None, + a_max: ArrayLike | None, + out: None = ..., + *, + min: ArrayLike | None = ..., + max: ArrayLike | None = ..., + dtype: None = ..., + where: _ArrayLikeBool_co | None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + signature: str | tuple[str | None, ...] = ..., + casting: _CastingKind = ..., +) -> _SCT: ... +@overload +def clip( + a: _ScalarLike_co, + a_min: ArrayLike | None, + a_max: ArrayLike | None, + out: None = ..., + *, + min: ArrayLike | None = ..., + max: ArrayLike | None = ..., + dtype: None = ..., + where: _ArrayLikeBool_co | None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + signature: str | tuple[str | None, ...] = ..., + casting: _CastingKind = ..., +) -> Any: ... +@overload +def clip( + a: _ArrayLike[_SCT], + a_min: ArrayLike | None, + a_max: ArrayLike | None, + out: None = ..., + *, + min: ArrayLike | None = ..., + max: ArrayLike | None = ..., + dtype: None = ..., + where: _ArrayLikeBool_co | None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + signature: str | tuple[str | None, ...] = ..., + casting: _CastingKind = ..., +) -> NDArray[_SCT]: ... +@overload +def clip( + a: ArrayLike, + a_min: ArrayLike | None, + a_max: ArrayLike | None, + out: None = ..., + *, + min: ArrayLike | None = ..., + max: ArrayLike | None = ..., + dtype: None = ..., + where: _ArrayLikeBool_co | None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + signature: str | tuple[str | None, ...] = ..., + casting: _CastingKind = ..., +) -> NDArray[Any]: ... +@overload +def clip( + a: ArrayLike, + a_min: ArrayLike | None, + a_max: ArrayLike | None, + out: _ArrayT, + *, + min: ArrayLike | None = ..., + max: ArrayLike | None = ..., + dtype: DTypeLike = ..., + where: _ArrayLikeBool_co | None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + signature: str | tuple[str | None, ...] = ..., + casting: _CastingKind = ..., +) -> _ArrayT: ... +@overload +def clip( + a: ArrayLike, + a_min: ArrayLike | None, + a_max: ArrayLike | None, + out: ArrayLike = ..., + *, + min: ArrayLike | None = ..., + max: ArrayLike | None = ..., + dtype: DTypeLike, + where: _ArrayLikeBool_co | None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + signature: str | tuple[str | None, ...] = ..., + casting: _CastingKind = ..., +) -> Any: ... + +@overload +def sum( + a: _ArrayLike[_SCT], + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def sum( + a: _ArrayLike[_SCT], + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT | NDArray[_SCT]: ... +@overload +def sum( + a: ArrayLike, + axis: None, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def sum( + a: ArrayLike, + axis: None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def sum( + a: ArrayLike, + axis: _ShapeLike | None, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT | NDArray[_SCT]: ... +@overload +def sum( + a: ArrayLike, + axis: _ShapeLike | None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT | NDArray[_SCT]: ... +@overload +def sum( + a: ArrayLike, + axis: _ShapeLike | None = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> Any: ... +@overload +def sum( + a: ArrayLike, + axis: _ShapeLike | None, + dtype: DTypeLike, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... +@overload +def sum( + a: ArrayLike, + axis: _ShapeLike | None = ..., + dtype: DTypeLike = ..., + *, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... + +@overload +def all( + a: ArrayLike, + axis: None = None, + out: None = None, + keepdims: Literal[False, 0] | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> np.bool: ... +@overload +def all( + a: ArrayLike, + axis: int | tuple[int, ...] | None = None, + out: None = None, + keepdims: _BoolLike_co | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> Incomplete: ... +@overload +def all( + a: ArrayLike, + axis: int | tuple[int, ...] | None, + out: _ArrayT, + keepdims: _BoolLike_co | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _ArrayT: ... +@overload +def all( + a: ArrayLike, + axis: int | tuple[int, ...] | None = None, + *, + out: _ArrayT, + keepdims: _BoolLike_co | _NoValueType = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _ArrayT: ... + +@overload +def any( + a: ArrayLike, + axis: None = None, + out: None = None, + keepdims: Literal[False, 0] | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> np.bool: ... +@overload +def any( + a: ArrayLike, + axis: int | tuple[int, ...] | None = None, + out: None = None, + keepdims: _BoolLike_co | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> Incomplete: ... +@overload +def any( + a: ArrayLike, + axis: int | tuple[int, ...] | None, + out: _ArrayT, + keepdims: _BoolLike_co | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _ArrayT: ... +@overload +def any( + a: ArrayLike, + axis: int | tuple[int, ...] | None = None, + *, + out: _ArrayT, + keepdims: _BoolLike_co | _NoValueType = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _ArrayT: ... + +@overload +def cumsum( + a: _ArrayLike[_SCT], + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def cumsum( + a: ArrayLike, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[Any]: ... +@overload +def cumsum( + a: ArrayLike, + axis: SupportsIndex | None, + dtype: _DTypeLike[_SCT], + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def cumsum( + a: ArrayLike, + axis: SupportsIndex | None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def cumsum( + a: ArrayLike, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + out: None = ..., +) -> NDArray[Any]: ... +@overload +def cumsum( + a: ArrayLike, + axis: SupportsIndex | None, + dtype: DTypeLike, + out: _ArrayT, +) -> _ArrayT: ... +@overload +def cumsum( + a: ArrayLike, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + *, + out: _ArrayT, +) -> _ArrayT: ... + +@overload +def cumulative_sum( + x: _ArrayLike[_SCT], + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def cumulative_sum( + x: ArrayLike, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[Any]: ... +@overload +def cumulative_sum( + x: ArrayLike, + /, + *, + axis: SupportsIndex | None = ..., + dtype: _DTypeLike[_SCT], + out: None = ..., + include_initial: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def cumulative_sum( + x: ArrayLike, + /, + *, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[Any]: ... +@overload +def cumulative_sum( + x: ArrayLike, + /, + *, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + out: _ArrayT, + include_initial: bool = ..., +) -> _ArrayT: ... + +@overload +def ptp( + a: _ArrayLike[_SCT], + axis: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., +) -> _SCT: ... +@overload +def ptp( + a: ArrayLike, + axis: _ShapeLike | None = ..., + out: None = ..., + keepdims: bool = ..., +) -> Any: ... +@overload +def ptp( + a: ArrayLike, + axis: _ShapeLike | None, + out: _ArrayT, + keepdims: bool = ..., +) -> _ArrayT: ... +@overload +def ptp( + a: ArrayLike, + axis: _ShapeLike | None = ..., + *, + out: _ArrayT, + keepdims: bool = ..., +) -> _ArrayT: ... + +@overload +def amax( + a: _ArrayLike[_SCT], + axis: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def amax( + a: ArrayLike, + axis: _ShapeLike | None = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> Any: ... +@overload +def amax( + a: ArrayLike, + axis: _ShapeLike | None, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... +@overload +def amax( + a: ArrayLike, + axis: _ShapeLike | None = ..., + *, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... + +@overload +def amin( + a: _ArrayLike[_SCT], + axis: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def amin( + a: ArrayLike, + axis: _ShapeLike | None = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> Any: ... +@overload +def amin( + a: ArrayLike, + axis: _ShapeLike | None, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... +@overload +def amin( + a: ArrayLike, + axis: _ShapeLike | None = ..., + *, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... + +# TODO: `np.prod()``: For object arrays `initial` does not necessarily +# have to be a numerical scalar. +# The only requirement is that it is compatible +# with the `.__mul__()` method(s) of the passed array's elements. + +# Note that the same situation holds for all wrappers around +# `np.ufunc.reduce`, e.g. `np.sum()` (`.__add__()`). +@overload +def prod( + a: _ArrayLikeBool_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> int_: ... +@overload +def prod( + a: _ArrayLikeUInt_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> uint64: ... +@overload +def prod( + a: _ArrayLikeInt_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> int64: ... +@overload +def prod( + a: _ArrayLikeFloat_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> floating[Any]: ... +@overload +def prod( + a: _ArrayLikeComplex_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> complexfloating[Any, Any]: ... +@overload +def prod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: None = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> Any: ... +@overload +def prod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def prod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: Literal[False] = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _SCT: ... +@overload +def prod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike | None = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> Any: ... +@overload +def prod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None, + dtype: DTypeLike | None, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... +@overload +def prod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike | None = ..., + *, + out: _ArrayT, + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., +) -> _ArrayT: ... + +@overload +def cumprod( + a: _ArrayLikeBool_co, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[int_]: ... +@overload +def cumprod( + a: _ArrayLikeUInt_co, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[uint64]: ... +@overload +def cumprod( + a: _ArrayLikeInt_co, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[int64]: ... +@overload +def cumprod( + a: _ArrayLikeFloat_co, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[floating[Any]]: ... +@overload +def cumprod( + a: _ArrayLikeComplex_co, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def cumprod( + a: _ArrayLikeObject_co, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., +) -> NDArray[object_]: ... +@overload +def cumprod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: SupportsIndex | None, + dtype: _DTypeLike[_SCT], + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def cumprod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: SupportsIndex | None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., +) -> NDArray[_SCT]: ... +@overload +def cumprod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + out: None = ..., +) -> NDArray[Any]: ... +@overload +def cumprod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: SupportsIndex | None, + dtype: DTypeLike, + out: _ArrayT, +) -> _ArrayT: ... +@overload +def cumprod( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + *, + out: _ArrayT, +) -> _ArrayT: ... + +@overload +def cumulative_prod( + x: _ArrayLikeBool_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[int_]: ... +@overload +def cumulative_prod( + x: _ArrayLikeUInt_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[uint64]: ... +@overload +def cumulative_prod( + x: _ArrayLikeInt_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[int64]: ... +@overload +def cumulative_prod( + x: _ArrayLikeFloat_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[floating[Any]]: ... +@overload +def cumulative_prod( + x: _ArrayLikeComplex_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def cumulative_prod( + x: _ArrayLikeObject_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: None = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[object_]: ... +@overload +def cumulative_prod( + x: _ArrayLikeComplex_co | _ArrayLikeObject_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: _DTypeLike[_SCT], + out: None = ..., + include_initial: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def cumulative_prod( + x: _ArrayLikeComplex_co | _ArrayLikeObject_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + out: None = ..., + include_initial: bool = ..., +) -> NDArray[Any]: ... +@overload +def cumulative_prod( + x: _ArrayLikeComplex_co | _ArrayLikeObject_co, + /, + *, + axis: SupportsIndex | None = ..., + dtype: DTypeLike = ..., + out: _ArrayT, + include_initial: bool = ..., +) -> _ArrayT: ... + +def ndim(a: ArrayLike) -> int: ... + +def size(a: ArrayLike, axis: int | None = ...) -> int: ... + +@overload +def around( + a: _BoolLike_co, + decimals: SupportsIndex = ..., + out: None = ..., +) -> float16: ... +@overload +def around( + a: _SCT_uifcO, + decimals: SupportsIndex = ..., + out: None = ..., +) -> _SCT_uifcO: ... +@overload +def around( + a: _ComplexLike_co | object_, + decimals: SupportsIndex = ..., + out: None = ..., +) -> Any: ... +@overload +def around( + a: _ArrayLikeBool_co, + decimals: SupportsIndex = ..., + out: None = ..., +) -> NDArray[float16]: ... +@overload +def around( + a: _ArrayLike[_SCT_uifcO], + decimals: SupportsIndex = ..., + out: None = ..., +) -> NDArray[_SCT_uifcO]: ... +@overload +def around( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + decimals: SupportsIndex = ..., + out: None = ..., +) -> NDArray[Any]: ... +@overload +def around( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + decimals: SupportsIndex, + out: _ArrayT, +) -> _ArrayT: ... +@overload +def around( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + decimals: SupportsIndex = ..., + *, + out: _ArrayT, +) -> _ArrayT: ... + +@overload +def mean( + a: _ArrayLikeFloat_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> floating[Any]: ... +@overload +def mean( + a: _ArrayLikeComplex_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> complexfloating[Any]: ... +@overload +def mean( + a: _ArrayLike[np.timedelta64], + axis: None = ..., + dtype: None = ..., + out: None = ..., + keepdims: Literal[False] | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> timedelta64: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None, + dtype: DTypeLike, + out: _ArrayT, + keepdims: bool | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _ArrayT: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike | None = ..., + *, + out: _ArrayT, + keepdims: bool | _NoValueType = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _ArrayT: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: Literal[False] | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _SCT: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: Literal[False] | _NoValueType = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _SCT: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None, + dtype: _DTypeLike[_SCT], + out: None, + keepdims: Literal[True, 1], + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> NDArray[_SCT]: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None, + dtype: _DTypeLike[_SCT], + out: None = ..., + *, + keepdims: bool | _NoValueType = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _SCT | NDArray[_SCT]: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + keepdims: bool | _NoValueType = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> _SCT | NDArray[_SCT]: ... +@overload +def mean( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike | None = ..., + out: None = ..., + keepdims: bool | _NoValueType = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., +) -> Incomplete: ... + +@overload +def std( + a: _ArrayLikeComplex_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + ddof: float = ..., + keepdims: Literal[False] = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> floating[Any]: ... +@overload +def std( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: None = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> Any: ... +@overload +def std( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None, + dtype: _DTypeLike[_SCT], + out: None = ..., + ddof: float = ..., + keepdims: Literal[False] = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _SCT: ... +@overload +def std( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + ddof: float = ..., + keepdims: Literal[False] = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _SCT: ... +@overload +def std( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> Any: ... +@overload +def std( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None, + dtype: DTypeLike, + out: _ArrayT, + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _ArrayT: ... +@overload +def std( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike = ..., + *, + out: _ArrayT, + ddof: float = ..., + keepdims: bool = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _ArrayT: ... + +@overload +def var( + a: _ArrayLikeComplex_co, + axis: None = ..., + dtype: None = ..., + out: None = ..., + ddof: float = ..., + keepdims: Literal[False] = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> floating[Any]: ... +@overload +def var( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: None = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> Any: ... +@overload +def var( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None, + dtype: _DTypeLike[_SCT], + out: None = ..., + ddof: float = ..., + keepdims: Literal[False] = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _SCT: ... +@overload +def var( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: None = ..., + *, + dtype: _DTypeLike[_SCT], + out: None = ..., + ddof: float = ..., + keepdims: Literal[False] = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _SCT: ... +@overload +def var( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> Any: ... +@overload +def var( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None, + dtype: DTypeLike, + out: _ArrayT, + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _ArrayT: ... +@overload +def var( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = ..., + dtype: DTypeLike = ..., + *, + out: _ArrayT, + ddof: float = ..., + keepdims: bool = ..., + where: _ArrayLikeBool_co | _NoValueType = ..., + mean: _ArrayLikeComplex_co | _ArrayLikeObject_co | _NoValueType = ..., + correction: float | _NoValueType = ..., +) -> _ArrayT: ... + +max = amax +min = amin +round = around diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/function_base.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/function_base.py new file mode 100644 index 0000000000000000000000000000000000000000..25724a411a1106ecafc2fab59a5ccd0fe6f1a6a6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/function_base.py @@ -0,0 +1,546 @@ +import functools +import warnings +import operator +import types + +import numpy as np +from . import numeric as _nx +from .numeric import result_type, nan, asanyarray, ndim +from numpy._core.multiarray import add_docstring +from numpy._core._multiarray_umath import _array_converter +from numpy._core import overrides + +__all__ = ['logspace', 'linspace', 'geomspace'] + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None, + dtype=None, axis=None, *, device=None): + return (start, stop) + + +@array_function_dispatch(_linspace_dispatcher) +def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, + axis=0, *, device=None): + """ + Return evenly spaced numbers over a specified interval. + + Returns `num` evenly spaced samples, calculated over the + interval [`start`, `stop`]. + + The endpoint of the interval can optionally be excluded. + + .. versionchanged:: 1.20.0 + Values are rounded towards ``-inf`` instead of ``0`` when an + integer ``dtype`` is specified. The old behavior can + still be obtained with ``np.linspace(start, stop, num).astype(int)`` + + Parameters + ---------- + start : array_like + The starting value of the sequence. + stop : array_like + The end value of the sequence, unless `endpoint` is set to False. + In that case, the sequence consists of all but the last of ``num + 1`` + evenly spaced samples, so that `stop` is excluded. Note that the step + size changes when `endpoint` is False. + num : int, optional + Number of samples to generate. Default is 50. Must be non-negative. + endpoint : bool, optional + If True, `stop` is the last sample. Otherwise, it is not included. + Default is True. + retstep : bool, optional + If True, return (`samples`, `step`), where `step` is the spacing + between samples. + dtype : dtype, optional + The type of the output array. If `dtype` is not given, the data type + is inferred from `start` and `stop`. The inferred dtype will never be + an integer; `float` is chosen even if the arguments would produce an + array of integers. + axis : int, optional + The axis in the result to store the samples. Relevant only if start + or stop are array-like. By default (0), the samples will be along a + new axis inserted at the beginning. Use -1 to get an axis at the end. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + + Returns + ------- + samples : ndarray + There are `num` equally spaced samples in the closed interval + ``[start, stop]`` or the half-open interval ``[start, stop)`` + (depending on whether `endpoint` is True or False). + step : float, optional + Only returned if `retstep` is True + + Size of spacing between samples. + + + See Also + -------- + arange : Similar to `linspace`, but uses a step size (instead of the + number of samples). + geomspace : Similar to `linspace`, but with numbers spaced evenly on a log + scale (a geometric progression). + logspace : Similar to `geomspace`, but with the end points specified as + logarithms. + :ref:`how-to-partition` + + Examples + -------- + >>> import numpy as np + >>> np.linspace(2.0, 3.0, num=5) + array([2. , 2.25, 2.5 , 2.75, 3. ]) + >>> np.linspace(2.0, 3.0, num=5, endpoint=False) + array([2. , 2.2, 2.4, 2.6, 2.8]) + >>> np.linspace(2.0, 3.0, num=5, retstep=True) + (array([2. , 2.25, 2.5 , 2.75, 3. ]), 0.25) + + Graphical illustration: + + >>> import matplotlib.pyplot as plt + >>> N = 8 + >>> y = np.zeros(N) + >>> x1 = np.linspace(0, 10, N, endpoint=True) + >>> x2 = np.linspace(0, 10, N, endpoint=False) + >>> plt.plot(x1, y, 'o') + [] + >>> plt.plot(x2, y + 0.5, 'o') + [] + >>> plt.ylim([-0.5, 1]) + (-0.5, 1) + >>> plt.show() + + """ + num = operator.index(num) + if num < 0: + raise ValueError( + "Number of samples, %s, must be non-negative." % num + ) + div = (num - 1) if endpoint else num + + conv = _array_converter(start, stop) + start, stop = conv.as_arrays() + dt = conv.result_type(ensure_inexact=True) + + if dtype is None: + dtype = dt + integer_dtype = False + else: + integer_dtype = _nx.issubdtype(dtype, _nx.integer) + + # Use `dtype=type(dt)` to enforce a floating point evaluation: + delta = np.subtract(stop, start, dtype=type(dt)) + y = _nx.arange( + 0, num, dtype=dt, device=device + ).reshape((-1,) + (1,) * ndim(delta)) + + # In-place multiplication y *= delta/div is faster, but prevents + # the multiplicant from overriding what class is produced, and thus + # prevents, e.g. use of Quantities, see gh-7142. Hence, we multiply + # in place only for standard scalar types. + if div > 0: + _mult_inplace = _nx.isscalar(delta) + step = delta / div + any_step_zero = ( + step == 0 if _mult_inplace else _nx.asanyarray(step == 0).any()) + if any_step_zero: + # Special handling for denormal numbers, gh-5437 + y /= div + if _mult_inplace: + y *= delta + else: + y = y * delta + else: + if _mult_inplace: + y *= step + else: + y = y * step + else: + # sequences with 0 items or 1 item with endpoint=True (i.e. div <= 0) + # have an undefined step + step = nan + # Multiply with delta to allow possible override of output class. + y = y * delta + + y += start + + if endpoint and num > 1: + y[-1, ...] = stop + + if axis != 0: + y = _nx.moveaxis(y, 0, axis) + + if integer_dtype: + _nx.floor(y, out=y) + + y = conv.wrap(y.astype(dtype, copy=False)) + if retstep: + return y, step + else: + return y + + +def _logspace_dispatcher(start, stop, num=None, endpoint=None, base=None, + dtype=None, axis=None): + return (start, stop, base) + + +@array_function_dispatch(_logspace_dispatcher) +def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, + axis=0): + """ + Return numbers spaced evenly on a log scale. + + In linear space, the sequence starts at ``base ** start`` + (`base` to the power of `start`) and ends with ``base ** stop`` + (see `endpoint` below). + + .. versionchanged:: 1.25.0 + Non-scalar 'base` is now supported + + Parameters + ---------- + start : array_like + ``base ** start`` is the starting value of the sequence. + stop : array_like + ``base ** stop`` is the final value of the sequence, unless `endpoint` + is False. In that case, ``num + 1`` values are spaced over the + interval in log-space, of which all but the last (a sequence of + length `num`) are returned. + num : integer, optional + Number of samples to generate. Default is 50. + endpoint : boolean, optional + If true, `stop` is the last sample. Otherwise, it is not included. + Default is True. + base : array_like, optional + The base of the log space. The step size between the elements in + ``ln(samples) / ln(base)`` (or ``log_base(samples)``) is uniform. + Default is 10.0. + dtype : dtype + The type of the output array. If `dtype` is not given, the data type + is inferred from `start` and `stop`. The inferred type will never be + an integer; `float` is chosen even if the arguments would produce an + array of integers. + axis : int, optional + The axis in the result to store the samples. Relevant only if start, + stop, or base are array-like. By default (0), the samples will be + along a new axis inserted at the beginning. Use -1 to get an axis at + the end. + + Returns + ------- + samples : ndarray + `num` samples, equally spaced on a log scale. + + See Also + -------- + arange : Similar to linspace, with the step size specified instead of the + number of samples. Note that, when used with a float endpoint, the + endpoint may or may not be included. + linspace : Similar to logspace, but with the samples uniformly distributed + in linear space, instead of log space. + geomspace : Similar to logspace, but with endpoints specified directly. + :ref:`how-to-partition` + + Notes + ----- + If base is a scalar, logspace is equivalent to the code + + >>> y = np.linspace(start, stop, num=num, endpoint=endpoint) + ... # doctest: +SKIP + >>> power(base, y).astype(dtype) + ... # doctest: +SKIP + + Examples + -------- + >>> import numpy as np + >>> np.logspace(2.0, 3.0, num=4) + array([ 100. , 215.443469 , 464.15888336, 1000. ]) + >>> np.logspace(2.0, 3.0, num=4, endpoint=False) + array([100. , 177.827941 , 316.22776602, 562.34132519]) + >>> np.logspace(2.0, 3.0, num=4, base=2.0) + array([4. , 5.0396842 , 6.34960421, 8. ]) + >>> np.logspace(2.0, 3.0, num=4, base=[2.0, 3.0], axis=-1) + array([[ 4. , 5.0396842 , 6.34960421, 8. ], + [ 9. , 12.98024613, 18.72075441, 27. ]]) + + Graphical illustration: + + >>> import matplotlib.pyplot as plt + >>> N = 10 + >>> x1 = np.logspace(0.1, 1, N, endpoint=True) + >>> x2 = np.logspace(0.1, 1, N, endpoint=False) + >>> y = np.zeros(N) + >>> plt.plot(x1, y, 'o') + [] + >>> plt.plot(x2, y + 0.5, 'o') + [] + >>> plt.ylim([-0.5, 1]) + (-0.5, 1) + >>> plt.show() + + """ + if not isinstance(base, (float, int)) and np.ndim(base): + # If base is non-scalar, broadcast it with the others, since it + # may influence how axis is interpreted. + ndmax = np.broadcast(start, stop, base).ndim + start, stop, base = ( + np.array(a, copy=None, subok=True, ndmin=ndmax) + for a in (start, stop, base) + ) + base = np.expand_dims(base, axis=axis) + y = linspace(start, stop, num=num, endpoint=endpoint, axis=axis) + if dtype is None: + return _nx.power(base, y) + return _nx.power(base, y).astype(dtype, copy=False) + + +def _geomspace_dispatcher(start, stop, num=None, endpoint=None, dtype=None, + axis=None): + return (start, stop) + + +@array_function_dispatch(_geomspace_dispatcher) +def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0): + """ + Return numbers spaced evenly on a log scale (a geometric progression). + + This is similar to `logspace`, but with endpoints specified directly. + Each output sample is a constant multiple of the previous. + + Parameters + ---------- + start : array_like + The starting value of the sequence. + stop : array_like + The final value of the sequence, unless `endpoint` is False. + In that case, ``num + 1`` values are spaced over the + interval in log-space, of which all but the last (a sequence of + length `num`) are returned. + num : integer, optional + Number of samples to generate. Default is 50. + endpoint : boolean, optional + If true, `stop` is the last sample. Otherwise, it is not included. + Default is True. + dtype : dtype + The type of the output array. If `dtype` is not given, the data type + is inferred from `start` and `stop`. The inferred dtype will never be + an integer; `float` is chosen even if the arguments would produce an + array of integers. + axis : int, optional + The axis in the result to store the samples. Relevant only if start + or stop are array-like. By default (0), the samples will be along a + new axis inserted at the beginning. Use -1 to get an axis at the end. + + Returns + ------- + samples : ndarray + `num` samples, equally spaced on a log scale. + + See Also + -------- + logspace : Similar to geomspace, but with endpoints specified using log + and base. + linspace : Similar to geomspace, but with arithmetic instead of geometric + progression. + arange : Similar to linspace, with the step size specified instead of the + number of samples. + :ref:`how-to-partition` + + Notes + ----- + If the inputs or dtype are complex, the output will follow a logarithmic + spiral in the complex plane. (There are an infinite number of spirals + passing through two points; the output will follow the shortest such path.) + + Examples + -------- + >>> import numpy as np + >>> np.geomspace(1, 1000, num=4) + array([ 1., 10., 100., 1000.]) + >>> np.geomspace(1, 1000, num=3, endpoint=False) + array([ 1., 10., 100.]) + >>> np.geomspace(1, 1000, num=4, endpoint=False) + array([ 1. , 5.62341325, 31.6227766 , 177.827941 ]) + >>> np.geomspace(1, 256, num=9) + array([ 1., 2., 4., 8., 16., 32., 64., 128., 256.]) + + Note that the above may not produce exact integers: + + >>> np.geomspace(1, 256, num=9, dtype=int) + array([ 1, 2, 4, 7, 16, 32, 63, 127, 256]) + >>> np.around(np.geomspace(1, 256, num=9)).astype(int) + array([ 1, 2, 4, 8, 16, 32, 64, 128, 256]) + + Negative, decreasing, and complex inputs are allowed: + + >>> np.geomspace(1000, 1, num=4) + array([1000., 100., 10., 1.]) + >>> np.geomspace(-1000, -1, num=4) + array([-1000., -100., -10., -1.]) + >>> np.geomspace(1j, 1000j, num=4) # Straight line + array([0. +1.j, 0. +10.j, 0. +100.j, 0.+1000.j]) + >>> np.geomspace(-1+0j, 1+0j, num=5) # Circle + array([-1.00000000e+00+1.22464680e-16j, -7.07106781e-01+7.07106781e-01j, + 6.12323400e-17+1.00000000e+00j, 7.07106781e-01+7.07106781e-01j, + 1.00000000e+00+0.00000000e+00j]) + + Graphical illustration of `endpoint` parameter: + + >>> import matplotlib.pyplot as plt + >>> N = 10 + >>> y = np.zeros(N) + >>> plt.semilogx(np.geomspace(1, 1000, N, endpoint=True), y + 1, 'o') + [] + >>> plt.semilogx(np.geomspace(1, 1000, N, endpoint=False), y + 2, 'o') + [] + >>> plt.axis([0.5, 2000, 0, 3]) + [0.5, 2000, 0, 3] + >>> plt.grid(True, color='0.7', linestyle='-', which='both', axis='both') + >>> plt.show() + + """ + start = asanyarray(start) + stop = asanyarray(stop) + if _nx.any(start == 0) or _nx.any(stop == 0): + raise ValueError('Geometric sequence cannot include zero') + + dt = result_type(start, stop, float(num), _nx.zeros((), dtype)) + if dtype is None: + dtype = dt + else: + # complex to dtype('complex128'), for instance + dtype = _nx.dtype(dtype) + + # Promote both arguments to the same dtype in case, for instance, one is + # complex and another is negative and log would produce NaN otherwise. + # Copy since we may change things in-place further down. + start = start.astype(dt, copy=True) + stop = stop.astype(dt, copy=True) + + # Allow negative real values and ensure a consistent result for complex + # (including avoiding negligible real or imaginary parts in output) by + # rotating start to positive real, calculating, then undoing rotation. + out_sign = _nx.sign(start) + start /= out_sign + stop = stop / out_sign + + log_start = _nx.log10(start) + log_stop = _nx.log10(stop) + result = logspace(log_start, log_stop, num=num, + endpoint=endpoint, base=10.0, dtype=dt) + + # Make sure the endpoints match the start and stop arguments. This is + # necessary because np.exp(np.log(x)) is not necessarily equal to x. + if num > 0: + result[0] = start + if num > 1 and endpoint: + result[-1] = stop + + result *= out_sign + + if axis != 0: + result = _nx.moveaxis(result, 0, axis) + + return result.astype(dtype, copy=False) + + +def _needs_add_docstring(obj): + """ + Returns true if the only way to set the docstring of `obj` from python is + via add_docstring. + + This function errs on the side of being overly conservative. + """ + Py_TPFLAGS_HEAPTYPE = 1 << 9 + + if isinstance(obj, (types.FunctionType, types.MethodType, property)): + return False + + if isinstance(obj, type) and obj.__flags__ & Py_TPFLAGS_HEAPTYPE: + return False + + return True + + +def _add_docstring(obj, doc, warn_on_python): + if warn_on_python and not _needs_add_docstring(obj): + warnings.warn( + "add_newdoc was used on a pure-python object {}. " + "Prefer to attach it directly to the source." + .format(obj), + UserWarning, + stacklevel=3) + try: + add_docstring(obj, doc) + except Exception: + pass + + +def add_newdoc(place, obj, doc, warn_on_python=True): + """ + Add documentation to an existing object, typically one defined in C + + The purpose is to allow easier editing of the docstrings without requiring + a re-compile. This exists primarily for internal use within numpy itself. + + Parameters + ---------- + place : str + The absolute name of the module to import from + obj : str or None + The name of the object to add documentation to, typically a class or + function name. + doc : {str, Tuple[str, str], List[Tuple[str, str]]} + If a string, the documentation to apply to `obj` + + If a tuple, then the first element is interpreted as an attribute + of `obj` and the second as the docstring to apply - + ``(method, docstring)`` + + If a list, then each element of the list should be a tuple of length + two - ``[(method1, docstring1), (method2, docstring2), ...]`` + warn_on_python : bool + If True, the default, emit `UserWarning` if this is used to attach + documentation to a pure-python object. + + Notes + ----- + This routine never raises an error if the docstring can't be written, but + will raise an error if the object being documented does not exist. + + This routine cannot modify read-only docstrings, as appear + in new-style classes or built-in functions. Because this + routine never raises an error the caller must check manually + that the docstrings were changed. + + Since this function grabs the ``char *`` from a c-level str object and puts + it into the ``tp_doc`` slot of the type of `obj`, it violates a number of + C-API best-practices, by: + + - modifying a `PyTypeObject` after calling `PyType_Ready` + - calling `Py_INCREF` on the str and losing the reference, so the str + will never be released + + If possible it should be avoided. + """ + new = getattr(__import__(place, globals(), {}, [obj]), obj) + if isinstance(doc, str): + if "${ARRAY_FUNCTION_LIKE}" in doc: + doc = overrides.get_array_function_like_doc(new, doc) + _add_docstring(new, doc.strip(), warn_on_python) + elif isinstance(doc, tuple): + attr, docstring = doc + _add_docstring(getattr(new, attr), docstring.strip(), warn_on_python) + elif isinstance(doc, list): + for attr, docstring in doc: + _add_docstring( + getattr(new, attr), docstring.strip(), warn_on_python + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/function_base.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/function_base.pyi new file mode 100644 index 0000000000000000000000000000000000000000..fc8f1d0bf68f1f632d4c81be2bde54f68d31b1ce --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/function_base.pyi @@ -0,0 +1,235 @@ +from typing import ( + Literal as L, + overload, + Any, + SupportsIndex, + TypeVar, +) + +from numpy import floating, complexfloating, generic +from numpy._typing import ( + NDArray, + DTypeLike, + _DTypeLike, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, +) + +__all__ = ["logspace", "linspace", "geomspace"] + +_SCT = TypeVar("_SCT", bound=generic) + +@overload +def linspace( + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + dtype: None = ..., + axis: SupportsIndex = ..., + *, + device: L["cpu"] | None = ..., +) -> NDArray[floating]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + dtype: None = ..., + axis: SupportsIndex = ..., + *, + device: L["cpu"] | None = ..., +) -> NDArray[complexfloating]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex, + endpoint: bool, + retstep: L[False], + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., + *, + device: L["cpu"] | None = ..., +) -> NDArray[_SCT]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + *, + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., + device: L["cpu"] | None = ..., +) -> NDArray[_SCT]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + dtype: DTypeLike = ..., + axis: SupportsIndex = ..., + *, + device: L["cpu"] | None = ..., +) -> NDArray[Any]: ... +@overload +def linspace( + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + *, + retstep: L[True], + dtype: None = ..., + axis: SupportsIndex = ..., + device: L["cpu"] | None = ..., +) -> tuple[NDArray[floating], floating]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + *, + retstep: L[True], + dtype: None = ..., + axis: SupportsIndex = ..., + device: L["cpu"] | None = ..., +) -> tuple[NDArray[complexfloating], complexfloating]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + *, + retstep: L[True], + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., + device: L["cpu"] | None = ..., +) -> tuple[NDArray[_SCT], _SCT]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + *, + retstep: L[True], + dtype: DTypeLike = ..., + axis: SupportsIndex = ..., + device: L["cpu"] | None = ..., +) -> tuple[NDArray[Any], Any]: ... + +@overload +def logspace( + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeFloat_co = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[floating]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeComplex_co = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[complexfloating]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex, + endpoint: bool, + base: _ArrayLikeComplex_co, + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeComplex_co = ..., + *, + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeComplex_co = ..., + dtype: DTypeLike = ..., + axis: SupportsIndex = ..., +) -> NDArray[Any]: ... + +@overload +def geomspace( + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[floating]: ... +@overload +def geomspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[complexfloating]: ... +@overload +def geomspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex, + endpoint: bool, + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def geomspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + *, + dtype: _DTypeLike[_SCT], + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def geomspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: DTypeLike = ..., + axis: SupportsIndex = ..., +) -> NDArray[Any]: ... + +def add_newdoc( + place: str, + obj: str, + doc: str | tuple[str, str] | list[tuple[str, str]], + warn_on_python: bool = ..., +) -> None: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/getlimits.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/getlimits.py new file mode 100644 index 0000000000000000000000000000000000000000..7e8778e0d2e13477e5126683120c648cd902fb7b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/getlimits.py @@ -0,0 +1,747 @@ +"""Machine limits for Float32 and Float64 and (long double) if available... + +""" +__all__ = ['finfo', 'iinfo'] + +import types +import warnings + +from .._utils import set_module +from ._machar import MachAr +from . import numeric +from . import numerictypes as ntypes +from .numeric import array, inf, nan +from .umath import log10, exp2, nextafter, isnan + + +def _fr0(a): + """fix rank-0 --> rank-1""" + if a.ndim == 0: + a = a.copy() + a.shape = (1,) + return a + + +def _fr1(a): + """fix rank > 0 --> rank-0""" + if a.size == 1: + a = a.copy() + a.shape = () + return a + + +class MachArLike: + """ Object to simulate MachAr instance """ + def __init__(self, ftype, *, eps, epsneg, huge, tiny, + ibeta, smallest_subnormal=None, **kwargs): + self.params = _MACHAR_PARAMS[ftype] + self.ftype = ftype + self.title = self.params['title'] + # Parameter types same as for discovered MachAr object. + if not smallest_subnormal: + self._smallest_subnormal = nextafter( + self.ftype(0), self.ftype(1), dtype=self.ftype) + else: + self._smallest_subnormal = smallest_subnormal + self.epsilon = self.eps = self._float_to_float(eps) + self.epsneg = self._float_to_float(epsneg) + self.xmax = self.huge = self._float_to_float(huge) + self.xmin = self._float_to_float(tiny) + self.smallest_normal = self.tiny = self._float_to_float(tiny) + self.ibeta = self.params['itype'](ibeta) + self.__dict__.update(kwargs) + self.precision = int(-log10(self.eps)) + self.resolution = self._float_to_float( + self._float_conv(10) ** (-self.precision)) + self._str_eps = self._float_to_str(self.eps) + self._str_epsneg = self._float_to_str(self.epsneg) + self._str_xmin = self._float_to_str(self.xmin) + self._str_xmax = self._float_to_str(self.xmax) + self._str_resolution = self._float_to_str(self.resolution) + self._str_smallest_normal = self._float_to_str(self.xmin) + + @property + def smallest_subnormal(self): + """Return the value for the smallest subnormal. + + Returns + ------- + smallest_subnormal : float + value for the smallest subnormal. + + Warns + ----- + UserWarning + If the calculated value for the smallest subnormal is zero. + """ + # Check that the calculated value is not zero, in case it raises a + # warning. + value = self._smallest_subnormal + if self.ftype(0) == value: + warnings.warn( + 'The value of the smallest subnormal for {} type ' + 'is zero.'.format(self.ftype), UserWarning, stacklevel=2) + + return self._float_to_float(value) + + @property + def _str_smallest_subnormal(self): + """Return the string representation of the smallest subnormal.""" + return self._float_to_str(self.smallest_subnormal) + + def _float_to_float(self, value): + """Converts float to float. + + Parameters + ---------- + value : float + value to be converted. + """ + return _fr1(self._float_conv(value)) + + def _float_conv(self, value): + """Converts float to conv. + + Parameters + ---------- + value : float + value to be converted. + """ + return array([value], self.ftype) + + def _float_to_str(self, value): + """Converts float to str. + + Parameters + ---------- + value : float + value to be converted. + """ + return self.params['fmt'] % array(_fr0(value)[0], self.ftype) + + +_convert_to_float = { + ntypes.csingle: ntypes.single, + ntypes.complex128: ntypes.float64, + ntypes.clongdouble: ntypes.longdouble + } + +# Parameters for creating MachAr / MachAr-like objects +_title_fmt = 'numpy {} precision floating point number' +_MACHAR_PARAMS = { + ntypes.double: dict( + itype = ntypes.int64, + fmt = '%24.16e', + title = _title_fmt.format('double')), + ntypes.single: dict( + itype = ntypes.int32, + fmt = '%15.7e', + title = _title_fmt.format('single')), + ntypes.longdouble: dict( + itype = ntypes.longlong, + fmt = '%s', + title = _title_fmt.format('long double')), + ntypes.half: dict( + itype = ntypes.int16, + fmt = '%12.5e', + title = _title_fmt.format('half'))} + +# Key to identify the floating point type. Key is result of +# +# ftype = np.longdouble # or float64, float32, etc. +# v = (ftype(-1.0) / ftype(10.0)) +# v.view(v.dtype.newbyteorder('<')).tobytes() +# +# Uses division to work around deficiencies in strtold on some platforms. +# See: +# https://perl5.git.perl.org/perl.git/blob/3118d7d684b56cbeb702af874f4326683c45f045:/Configure + +_KNOWN_TYPES = {} +def _register_type(machar, bytepat): + _KNOWN_TYPES[bytepat] = machar + + +_float_ma = {} + + +def _register_known_types(): + # Known parameters for float16 + # See docstring of MachAr class for description of parameters. + f16 = ntypes.float16 + float16_ma = MachArLike(f16, + machep=-10, + negep=-11, + minexp=-14, + maxexp=16, + it=10, + iexp=5, + ibeta=2, + irnd=5, + ngrd=0, + eps=exp2(f16(-10)), + epsneg=exp2(f16(-11)), + huge=f16(65504), + tiny=f16(2 ** -14)) + _register_type(float16_ma, b'f\xae') + _float_ma[16] = float16_ma + + # Known parameters for float32 + f32 = ntypes.float32 + float32_ma = MachArLike(f32, + machep=-23, + negep=-24, + minexp=-126, + maxexp=128, + it=23, + iexp=8, + ibeta=2, + irnd=5, + ngrd=0, + eps=exp2(f32(-23)), + epsneg=exp2(f32(-24)), + huge=f32((1 - 2 ** -24) * 2**128), + tiny=exp2(f32(-126))) + _register_type(float32_ma, b'\xcd\xcc\xcc\xbd') + _float_ma[32] = float32_ma + + # Known parameters for float64 + f64 = ntypes.float64 + epsneg_f64 = 2.0 ** -53.0 + tiny_f64 = 2.0 ** -1022.0 + float64_ma = MachArLike(f64, + machep=-52, + negep=-53, + minexp=-1022, + maxexp=1024, + it=52, + iexp=11, + ibeta=2, + irnd=5, + ngrd=0, + eps=2.0 ** -52.0, + epsneg=epsneg_f64, + huge=(1.0 - epsneg_f64) / tiny_f64 * f64(4), + tiny=tiny_f64) + _register_type(float64_ma, b'\x9a\x99\x99\x99\x99\x99\xb9\xbf') + _float_ma[64] = float64_ma + + # Known parameters for IEEE 754 128-bit binary float + ld = ntypes.longdouble + epsneg_f128 = exp2(ld(-113)) + tiny_f128 = exp2(ld(-16382)) + # Ignore runtime error when this is not f128 + with numeric.errstate(all='ignore'): + huge_f128 = (ld(1) - epsneg_f128) / tiny_f128 * ld(4) + float128_ma = MachArLike(ld, + machep=-112, + negep=-113, + minexp=-16382, + maxexp=16384, + it=112, + iexp=15, + ibeta=2, + irnd=5, + ngrd=0, + eps=exp2(ld(-112)), + epsneg=epsneg_f128, + huge=huge_f128, + tiny=tiny_f128) + # IEEE 754 128-bit binary float + _register_type(float128_ma, + b'\x9a\x99\x99\x99\x99\x99\x99\x99\x99\x99\x99\x99\x99\x99\xfb\xbf') + _float_ma[128] = float128_ma + + # Known parameters for float80 (Intel 80-bit extended precision) + epsneg_f80 = exp2(ld(-64)) + tiny_f80 = exp2(ld(-16382)) + # Ignore runtime error when this is not f80 + with numeric.errstate(all='ignore'): + huge_f80 = (ld(1) - epsneg_f80) / tiny_f80 * ld(4) + float80_ma = MachArLike(ld, + machep=-63, + negep=-64, + minexp=-16382, + maxexp=16384, + it=63, + iexp=15, + ibeta=2, + irnd=5, + ngrd=0, + eps=exp2(ld(-63)), + epsneg=epsneg_f80, + huge=huge_f80, + tiny=tiny_f80) + # float80, first 10 bytes containing actual storage + _register_type(float80_ma, b'\xcd\xcc\xcc\xcc\xcc\xcc\xcc\xcc\xfb\xbf') + _float_ma[80] = float80_ma + + # Guessed / known parameters for double double; see: + # https://en.wikipedia.org/wiki/Quadruple-precision_floating-point_format#Double-double_arithmetic + # These numbers have the same exponent range as float64, but extended + # number of digits in the significand. + huge_dd = nextafter(ld(inf), ld(0), dtype=ld) + # As the smallest_normal in double double is so hard to calculate we set + # it to NaN. + smallest_normal_dd = nan + # Leave the same value for the smallest subnormal as double + smallest_subnormal_dd = ld(nextafter(0., 1.)) + float_dd_ma = MachArLike(ld, + machep=-105, + negep=-106, + minexp=-1022, + maxexp=1024, + it=105, + iexp=11, + ibeta=2, + irnd=5, + ngrd=0, + eps=exp2(ld(-105)), + epsneg=exp2(ld(-106)), + huge=huge_dd, + tiny=smallest_normal_dd, + smallest_subnormal=smallest_subnormal_dd) + # double double; low, high order (e.g. PPC 64) + _register_type(float_dd_ma, + b'\x9a\x99\x99\x99\x99\x99Y<\x9a\x99\x99\x99\x99\x99\xb9\xbf') + # double double; high, low order (e.g. PPC 64 le) + _register_type(float_dd_ma, + b'\x9a\x99\x99\x99\x99\x99\xb9\xbf\x9a\x99\x99\x99\x99\x99Y<') + _float_ma['dd'] = float_dd_ma + + +def _get_machar(ftype): + """ Get MachAr instance or MachAr-like instance + + Get parameters for floating point type, by first trying signatures of + various known floating point types, then, if none match, attempting to + identify parameters by analysis. + + Parameters + ---------- + ftype : class + Numpy floating point type class (e.g. ``np.float64``) + + Returns + ------- + ma_like : instance of :class:`MachAr` or :class:`MachArLike` + Object giving floating point parameters for `ftype`. + + Warns + ----- + UserWarning + If the binary signature of the float type is not in the dictionary of + known float types. + """ + params = _MACHAR_PARAMS.get(ftype) + if params is None: + raise ValueError(repr(ftype)) + # Detect known / suspected types + # ftype(-1.0) / ftype(10.0) is better than ftype('-0.1') because stold + # may be deficient + key = (ftype(-1.0) / ftype(10.)) + key = key.view(key.dtype.newbyteorder("<")).tobytes() + ma_like = None + if ftype == ntypes.longdouble: + # Could be 80 bit == 10 byte extended precision, where last bytes can + # be random garbage. + # Comparing first 10 bytes to pattern first to avoid branching on the + # random garbage. + ma_like = _KNOWN_TYPES.get(key[:10]) + if ma_like is None: + # see if the full key is known. + ma_like = _KNOWN_TYPES.get(key) + if ma_like is None and len(key) == 16: + # machine limits could be f80 masquerading as np.float128, + # find all keys with length 16 and make new dict, but make the keys + # only 10 bytes long, the last bytes can be random garbage + _kt = {k[:10]: v for k, v in _KNOWN_TYPES.items() if len(k) == 16} + ma_like = _kt.get(key[:10]) + if ma_like is not None: + return ma_like + # Fall back to parameter discovery + warnings.warn( + f'Signature {key} for {ftype} does not match any known type: ' + 'falling back to type probe function.\n' + 'This warnings indicates broken support for the dtype!', + UserWarning, stacklevel=2) + return _discovered_machar(ftype) + + +def _discovered_machar(ftype): + """ Create MachAr instance with found information on float types + + TODO: MachAr should be retired completely ideally. We currently only + ever use it system with broken longdouble (valgrind, WSL). + """ + params = _MACHAR_PARAMS[ftype] + return MachAr(lambda v: array([v], ftype), + lambda v: _fr0(v.astype(params['itype']))[0], + lambda v: array(_fr0(v)[0], ftype), + lambda v: params['fmt'] % array(_fr0(v)[0], ftype), + params['title']) + + +@set_module('numpy') +class finfo: + """ + finfo(dtype) + + Machine limits for floating point types. + + Attributes + ---------- + bits : int + The number of bits occupied by the type. + dtype : dtype + Returns the dtype for which `finfo` returns information. For complex + input, the returned dtype is the associated ``float*`` dtype for its + real and complex components. + eps : float + The difference between 1.0 and the next smallest representable float + larger than 1.0. For example, for 64-bit binary floats in the IEEE-754 + standard, ``eps = 2**-52``, approximately 2.22e-16. + epsneg : float + The difference between 1.0 and the next smallest representable float + less than 1.0. For example, for 64-bit binary floats in the IEEE-754 + standard, ``epsneg = 2**-53``, approximately 1.11e-16. + iexp : int + The number of bits in the exponent portion of the floating point + representation. + machep : int + The exponent that yields `eps`. + max : floating point number of the appropriate type + The largest representable number. + maxexp : int + The smallest positive power of the base (2) that causes overflow. + min : floating point number of the appropriate type + The smallest representable number, typically ``-max``. + minexp : int + The most negative power of the base (2) consistent with there + being no leading 0's in the mantissa. + negep : int + The exponent that yields `epsneg`. + nexp : int + The number of bits in the exponent including its sign and bias. + nmant : int + The number of bits in the mantissa. + precision : int + The approximate number of decimal digits to which this kind of + float is precise. + resolution : floating point number of the appropriate type + The approximate decimal resolution of this type, i.e., + ``10**-precision``. + tiny : float + An alias for `smallest_normal`, kept for backwards compatibility. + smallest_normal : float + The smallest positive floating point number with 1 as leading bit in + the mantissa following IEEE-754 (see Notes). + smallest_subnormal : float + The smallest positive floating point number with 0 as leading bit in + the mantissa following IEEE-754. + + Parameters + ---------- + dtype : float, dtype, or instance + Kind of floating point or complex floating point + data-type about which to get information. + + See Also + -------- + iinfo : The equivalent for integer data types. + spacing : The distance between a value and the nearest adjacent number + nextafter : The next floating point value after x1 towards x2 + + Notes + ----- + For developers of NumPy: do not instantiate this at the module level. + The initial calculation of these parameters is expensive and negatively + impacts import times. These objects are cached, so calling ``finfo()`` + repeatedly inside your functions is not a problem. + + Note that ``smallest_normal`` is not actually the smallest positive + representable value in a NumPy floating point type. As in the IEEE-754 + standard [1]_, NumPy floating point types make use of subnormal numbers to + fill the gap between 0 and ``smallest_normal``. However, subnormal numbers + may have significantly reduced precision [2]_. + + This function can also be used for complex data types as well. If used, + the output will be the same as the corresponding real float type + (e.g. numpy.finfo(numpy.csingle) is the same as numpy.finfo(numpy.single)). + However, the output is true for the real and imaginary components. + + References + ---------- + .. [1] IEEE Standard for Floating-Point Arithmetic, IEEE Std 754-2008, + pp.1-70, 2008, https://doi.org/10.1109/IEEESTD.2008.4610935 + .. [2] Wikipedia, "Denormal Numbers", + https://en.wikipedia.org/wiki/Denormal_number + + Examples + -------- + >>> import numpy as np + >>> np.finfo(np.float64).dtype + dtype('float64') + >>> np.finfo(np.complex64).dtype + dtype('float32') + + """ + + _finfo_cache = {} + + __class_getitem__ = classmethod(types.GenericAlias) + + def __new__(cls, dtype): + try: + obj = cls._finfo_cache.get(dtype) # most common path + if obj is not None: + return obj + except TypeError: + pass + + if dtype is None: + # Deprecated in NumPy 1.25, 2023-01-16 + warnings.warn( + "finfo() dtype cannot be None. This behavior will " + "raise an error in the future. (Deprecated in NumPy 1.25)", + DeprecationWarning, + stacklevel=2 + ) + + try: + dtype = numeric.dtype(dtype) + except TypeError: + # In case a float instance was given + dtype = numeric.dtype(type(dtype)) + + obj = cls._finfo_cache.get(dtype) + if obj is not None: + return obj + dtypes = [dtype] + newdtype = ntypes.obj2sctype(dtype) + if newdtype is not dtype: + dtypes.append(newdtype) + dtype = newdtype + if not issubclass(dtype, numeric.inexact): + raise ValueError("data type %r not inexact" % (dtype)) + obj = cls._finfo_cache.get(dtype) + if obj is not None: + return obj + if not issubclass(dtype, numeric.floating): + newdtype = _convert_to_float[dtype] + if newdtype is not dtype: + # dtype changed, for example from complex128 to float64 + dtypes.append(newdtype) + dtype = newdtype + + obj = cls._finfo_cache.get(dtype, None) + if obj is not None: + # the original dtype was not in the cache, but the new + # dtype is in the cache. we add the original dtypes to + # the cache and return the result + for dt in dtypes: + cls._finfo_cache[dt] = obj + return obj + obj = object.__new__(cls)._init(dtype) + for dt in dtypes: + cls._finfo_cache[dt] = obj + return obj + + def _init(self, dtype): + self.dtype = numeric.dtype(dtype) + machar = _get_machar(dtype) + + for word in ['precision', 'iexp', + 'maxexp', 'minexp', 'negep', + 'machep']: + setattr(self, word, getattr(machar, word)) + for word in ['resolution', 'epsneg', 'smallest_subnormal']: + setattr(self, word, getattr(machar, word).flat[0]) + self.bits = self.dtype.itemsize * 8 + self.max = machar.huge.flat[0] + self.min = -self.max + self.eps = machar.eps.flat[0] + self.nexp = machar.iexp + self.nmant = machar.it + self._machar = machar + self._str_tiny = machar._str_xmin.strip() + self._str_max = machar._str_xmax.strip() + self._str_epsneg = machar._str_epsneg.strip() + self._str_eps = machar._str_eps.strip() + self._str_resolution = machar._str_resolution.strip() + self._str_smallest_normal = machar._str_smallest_normal.strip() + self._str_smallest_subnormal = machar._str_smallest_subnormal.strip() + return self + + def __str__(self): + fmt = ( + 'Machine parameters for %(dtype)s\n' + '---------------------------------------------------------------\n' + 'precision = %(precision)3s resolution = %(_str_resolution)s\n' + 'machep = %(machep)6s eps = %(_str_eps)s\n' + 'negep = %(negep)6s epsneg = %(_str_epsneg)s\n' + 'minexp = %(minexp)6s tiny = %(_str_tiny)s\n' + 'maxexp = %(maxexp)6s max = %(_str_max)s\n' + 'nexp = %(nexp)6s min = -max\n' + 'smallest_normal = %(_str_smallest_normal)s ' + 'smallest_subnormal = %(_str_smallest_subnormal)s\n' + '---------------------------------------------------------------\n' + ) + return fmt % self.__dict__ + + def __repr__(self): + c = self.__class__.__name__ + d = self.__dict__.copy() + d['klass'] = c + return (("%(klass)s(resolution=%(resolution)s, min=-%(_str_max)s," + " max=%(_str_max)s, dtype=%(dtype)s)") % d) + + @property + def smallest_normal(self): + """Return the value for the smallest normal. + + Returns + ------- + smallest_normal : float + Value for the smallest normal. + + Warns + ----- + UserWarning + If the calculated value for the smallest normal is requested for + double-double. + """ + # This check is necessary because the value for smallest_normal is + # platform dependent for longdouble types. + if isnan(self._machar.smallest_normal.flat[0]): + warnings.warn( + 'The value of smallest normal is undefined for double double', + UserWarning, stacklevel=2) + return self._machar.smallest_normal.flat[0] + + @property + def tiny(self): + """Return the value for tiny, alias of smallest_normal. + + Returns + ------- + tiny : float + Value for the smallest normal, alias of smallest_normal. + + Warns + ----- + UserWarning + If the calculated value for the smallest normal is requested for + double-double. + """ + return self.smallest_normal + + +@set_module('numpy') +class iinfo: + """ + iinfo(type) + + Machine limits for integer types. + + Attributes + ---------- + bits : int + The number of bits occupied by the type. + dtype : dtype + Returns the dtype for which `iinfo` returns information. + min : int + The smallest integer expressible by the type. + max : int + The largest integer expressible by the type. + + Parameters + ---------- + int_type : integer type, dtype, or instance + The kind of integer data type to get information about. + + See Also + -------- + finfo : The equivalent for floating point data types. + + Examples + -------- + With types: + + >>> import numpy as np + >>> ii16 = np.iinfo(np.int16) + >>> ii16.min + -32768 + >>> ii16.max + 32767 + >>> ii32 = np.iinfo(np.int32) + >>> ii32.min + -2147483648 + >>> ii32.max + 2147483647 + + With instances: + + >>> ii32 = np.iinfo(np.int32(10)) + >>> ii32.min + -2147483648 + >>> ii32.max + 2147483647 + + """ + + _min_vals = {} + _max_vals = {} + + __class_getitem__ = classmethod(types.GenericAlias) + + def __init__(self, int_type): + try: + self.dtype = numeric.dtype(int_type) + except TypeError: + self.dtype = numeric.dtype(type(int_type)) + self.kind = self.dtype.kind + self.bits = self.dtype.itemsize * 8 + self.key = "%s%d" % (self.kind, self.bits) + if self.kind not in 'iu': + raise ValueError("Invalid integer data type %r." % (self.kind,)) + + @property + def min(self): + """Minimum value of given dtype.""" + if self.kind == 'u': + return 0 + else: + try: + val = iinfo._min_vals[self.key] + except KeyError: + val = int(-(1 << (self.bits-1))) + iinfo._min_vals[self.key] = val + return val + + @property + def max(self): + """Maximum value of given dtype.""" + try: + val = iinfo._max_vals[self.key] + except KeyError: + if self.kind == 'u': + val = int((1 << self.bits) - 1) + else: + val = int((1 << (self.bits-1)) - 1) + iinfo._max_vals[self.key] = val + return val + + def __str__(self): + """String representation.""" + fmt = ( + 'Machine parameters for %(dtype)s\n' + '---------------------------------------------------------------\n' + 'min = %(min)s\n' + 'max = %(max)s\n' + '---------------------------------------------------------------\n' + ) + return fmt % {'dtype': self.dtype, 'min': self.min, 'max': self.max} + + def __repr__(self): + return "%s(min=%s, max=%s, dtype=%s)" % (self.__class__.__name__, + self.min, self.max, self.dtype) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/getlimits.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/getlimits.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6da4954ed85c008493dfdaf1ded8488143efc674 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/getlimits.pyi @@ -0,0 +1,3 @@ +from numpy import finfo, iinfo + +__all__ = ["finfo", "iinfo"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/arrayobject.h b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/arrayobject.h new file mode 100644 index 0000000000000000000000000000000000000000..d3ca0a64121d2c2185b263f64ab2a6a72a25a218 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/arrayobject.h @@ -0,0 +1,7 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_ARRAYOBJECT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_ARRAYOBJECT_H_ +#define Py_ARRAYOBJECT_H + +#include "ndarrayobject.h" + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_ARRAYOBJECT_H_ */ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/arrayscalars.h b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/arrayscalars.h new file mode 100644 index 0000000000000000000000000000000000000000..35ae1abd6629426af139f35ba93c4356bf3f30ae --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/arrayscalars.h @@ -0,0 +1,196 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_ARRAYSCALARS_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_ARRAYSCALARS_H_ + +#ifndef _MULTIARRAYMODULE +typedef struct { + PyObject_HEAD + npy_bool obval; +} PyBoolScalarObject; +#endif + + +typedef struct { + PyObject_HEAD + signed char obval; +} PyByteScalarObject; + + +typedef struct { + PyObject_HEAD + short obval; +} PyShortScalarObject; + + +typedef struct { + PyObject_HEAD + int obval; +} PyIntScalarObject; + + +typedef struct { + PyObject_HEAD + long obval; +} PyLongScalarObject; + + +typedef struct { + PyObject_HEAD + npy_longlong obval; +} PyLongLongScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned char obval; +} PyUByteScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned short obval; +} PyUShortScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned int obval; +} PyUIntScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned long obval; +} PyULongScalarObject; + + +typedef struct { + PyObject_HEAD + npy_ulonglong obval; +} PyULongLongScalarObject; + + +typedef struct { + PyObject_HEAD + npy_half obval; +} PyHalfScalarObject; + + +typedef struct { + PyObject_HEAD + float obval; +} PyFloatScalarObject; + + +typedef struct { + PyObject_HEAD + double obval; +} PyDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + npy_longdouble obval; +} PyLongDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + npy_cfloat obval; +} PyCFloatScalarObject; + + +typedef struct { + PyObject_HEAD + npy_cdouble obval; +} PyCDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + npy_clongdouble obval; +} PyCLongDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + PyObject * obval; +} PyObjectScalarObject; + +typedef struct { + PyObject_HEAD + npy_datetime obval; + PyArray_DatetimeMetaData obmeta; +} PyDatetimeScalarObject; + +typedef struct { + PyObject_HEAD + npy_timedelta obval; + PyArray_DatetimeMetaData obmeta; +} PyTimedeltaScalarObject; + + +typedef struct { + PyObject_HEAD + char obval; +} PyScalarObject; + +#define PyStringScalarObject PyBytesObject +#ifndef Py_LIMITED_API +typedef struct { + /* note that the PyObject_HEAD macro lives right here */ + PyUnicodeObject base; + Py_UCS4 *obval; + #if NPY_FEATURE_VERSION >= NPY_1_20_API_VERSION + char *buffer_fmt; + #endif +} PyUnicodeScalarObject; +#endif + + +typedef struct { + PyObject_VAR_HEAD + char *obval; +#if defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD + /* Internally use the subclass to allow accessing names/fields */ + _PyArray_LegacyDescr *descr; +#else + PyArray_Descr *descr; +#endif + int flags; + PyObject *base; + #if NPY_FEATURE_VERSION >= NPY_1_20_API_VERSION + void *_buffer_info; /* private buffer info, tagged to allow warning */ + #endif +} PyVoidScalarObject; + +/* Macros + PyScalarObject + PyArrType_Type + are defined in ndarrayobject.h +*/ + +#define PyArrayScalar_False ((PyObject *)(&(_PyArrayScalar_BoolValues[0]))) +#define PyArrayScalar_True ((PyObject *)(&(_PyArrayScalar_BoolValues[1]))) +#define PyArrayScalar_FromLong(i) \ + ((PyObject *)(&(_PyArrayScalar_BoolValues[((i)!=0)]))) +#define PyArrayScalar_RETURN_BOOL_FROM_LONG(i) \ + return Py_INCREF(PyArrayScalar_FromLong(i)), \ + PyArrayScalar_FromLong(i) +#define PyArrayScalar_RETURN_FALSE \ + return Py_INCREF(PyArrayScalar_False), \ + PyArrayScalar_False +#define PyArrayScalar_RETURN_TRUE \ + return Py_INCREF(PyArrayScalar_True), \ + PyArrayScalar_True + +#define PyArrayScalar_New(cls) \ + Py##cls##ArrType_Type.tp_alloc(&Py##cls##ArrType_Type, 0) +#ifndef Py_LIMITED_API +/* For the limited API, use PyArray_ScalarAsCtype instead */ +#define PyArrayScalar_VAL(obj, cls) \ + ((Py##cls##ScalarObject *)obj)->obval +#define PyArrayScalar_ASSIGN(obj, cls, val) \ + PyArrayScalar_VAL(obj, cls) = val +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_ARRAYSCALARS_H_ */ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/dtype_api.h b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/dtype_api.h new file mode 100644 index 0000000000000000000000000000000000000000..2beefb51e1bb149de4439333d39585402441ac2d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/dtype_api.h @@ -0,0 +1,480 @@ +/* + * The public DType API + */ + +#ifndef NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ +#define NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ + +struct PyArrayMethodObject_tag; + +/* + * Largely opaque struct for DType classes (i.e. metaclass instances). + * The internal definition is currently in `ndarraytypes.h` (export is a bit + * more complex because `PyArray_Descr` is a DTypeMeta internally but not + * externally). + */ +#if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD) + +#ifndef Py_LIMITED_API + + typedef struct PyArray_DTypeMeta_tag { + PyHeapTypeObject super; + + /* + * Most DTypes will have a singleton default instance, for the + * parametric legacy DTypes (bytes, string, void, datetime) this + * may be a pointer to the *prototype* instance? + */ + PyArray_Descr *singleton; + /* Copy of the legacy DTypes type number, usually invalid. */ + int type_num; + + /* The type object of the scalar instances (may be NULL?) */ + PyTypeObject *scalar_type; + /* + * DType flags to signal legacy, parametric, or + * abstract. But plenty of space for additional information/flags. + */ + npy_uint64 flags; + + /* + * Use indirection in order to allow a fixed size for this struct. + * A stable ABI size makes creating a static DType less painful + * while also ensuring flexibility for all opaque API (with one + * indirection due the pointer lookup). + */ + void *dt_slots; + /* Allow growing (at the moment also beyond this) */ + void *reserved[3]; + } PyArray_DTypeMeta; + +#else + +typedef PyTypeObject PyArray_DTypeMeta; + +#endif /* Py_LIMITED_API */ + +#endif /* not internal build */ + +/* + * ****************************************************** + * ArrayMethod API (Casting and UFuncs) + * ****************************************************** + */ + + +typedef enum { + /* Flag for whether the GIL is required */ + NPY_METH_REQUIRES_PYAPI = 1 << 0, + /* + * Some functions cannot set floating point error flags, this flag + * gives us the option (not requirement) to skip floating point error + * setup/check. No function should set error flags and ignore them + * since it would interfere with chaining operations (e.g. casting). + */ + NPY_METH_NO_FLOATINGPOINT_ERRORS = 1 << 1, + /* Whether the method supports unaligned access (not runtime) */ + NPY_METH_SUPPORTS_UNALIGNED = 1 << 2, + /* + * Used for reductions to allow reordering the operation. At this point + * assume that if set, it also applies to normal operations though! + */ + NPY_METH_IS_REORDERABLE = 1 << 3, + /* + * Private flag for now for *logic* functions. The logical functions + * `logical_or` and `logical_and` can always cast the inputs to booleans + * "safely" (because that is how the cast to bool is defined). + * @seberg: I am not sure this is the best way to handle this, so its + * private for now (also it is very limited anyway). + * There is one "exception". NA aware dtypes cannot cast to bool + * (hopefully), so the `??->?` loop should error even with this flag. + * But a second NA fallback loop will be necessary. + */ + _NPY_METH_FORCE_CAST_INPUTS = 1 << 17, + + /* All flags which can change at runtime */ + NPY_METH_RUNTIME_FLAGS = ( + NPY_METH_REQUIRES_PYAPI | + NPY_METH_NO_FLOATINGPOINT_ERRORS), +} NPY_ARRAYMETHOD_FLAGS; + + +typedef struct PyArrayMethod_Context_tag { + /* The caller, which is typically the original ufunc. May be NULL */ + PyObject *caller; + /* The method "self". Currently an opaque object. */ + struct PyArrayMethodObject_tag *method; + + /* Operand descriptors, filled in by resolve_descriptors */ + PyArray_Descr *const *descriptors; + /* Structure may grow (this is harmless for DType authors) */ +} PyArrayMethod_Context; + + +/* + * The main object for creating a new ArrayMethod. We use the typical `slots` + * mechanism used by the Python limited API (see below for the slot defs). + */ +typedef struct { + const char *name; + int nin, nout; + NPY_CASTING casting; + NPY_ARRAYMETHOD_FLAGS flags; + PyArray_DTypeMeta **dtypes; + PyType_Slot *slots; +} PyArrayMethod_Spec; + + +/* + * ArrayMethod slots + * ----------------- + * + * SLOTS IDs For the ArrayMethod creation, once fully public, IDs are fixed + * but can be deprecated and arbitrarily extended. + */ +#define _NPY_METH_resolve_descriptors_with_scalars 1 +#define NPY_METH_resolve_descriptors 2 +#define NPY_METH_get_loop 3 +#define NPY_METH_get_reduction_initial 4 +/* specific loops for constructions/default get_loop: */ +#define NPY_METH_strided_loop 5 +#define NPY_METH_contiguous_loop 6 +#define NPY_METH_unaligned_strided_loop 7 +#define NPY_METH_unaligned_contiguous_loop 8 +#define NPY_METH_contiguous_indexed_loop 9 +#define _NPY_METH_static_data 10 + + +/* + * The resolve descriptors function, must be able to handle NULL values for + * all output (but not input) `given_descrs` and fill `loop_descrs`. + * Return -1 on error or 0 if the operation is not possible without an error + * set. (This may still be in flux.) + * Otherwise must return the "casting safety", for normal functions, this is + * almost always "safe" (or even "equivalent"?). + * + * `resolve_descriptors` is optional if all output DTypes are non-parametric. + */ +typedef NPY_CASTING (PyArrayMethod_ResolveDescriptors)( + /* "method" is currently opaque (necessary e.g. to wrap Python) */ + struct PyArrayMethodObject_tag *method, + /* DTypes the method was created for */ + PyArray_DTypeMeta *const *dtypes, + /* Input descriptors (instances). Outputs may be NULL. */ + PyArray_Descr *const *given_descrs, + /* Exact loop descriptors to use, must not hold references on error */ + PyArray_Descr **loop_descrs, + npy_intp *view_offset); + + +/* + * Rarely needed, slightly more powerful version of `resolve_descriptors`. + * See also `PyArrayMethod_ResolveDescriptors` for details on shared arguments. + * + * NOTE: This function is private now as it is unclear how and what to pass + * exactly as additional information to allow dealing with the scalars. + * See also gh-24915. + */ +typedef NPY_CASTING (PyArrayMethod_ResolveDescriptorsWithScalar)( + struct PyArrayMethodObject_tag *method, + PyArray_DTypeMeta *const *dtypes, + /* Unlike above, these can have any DType and we may allow NULL. */ + PyArray_Descr *const *given_descrs, + /* + * Input scalars or NULL. Only ever passed for python scalars. + * WARNING: In some cases, a loop may be explicitly selected and the + * value passed is not available (NULL) or does not have the + * expected type. + */ + PyObject *const *input_scalars, + PyArray_Descr **loop_descrs, + npy_intp *view_offset); + + + +typedef int (PyArrayMethod_StridedLoop)(PyArrayMethod_Context *context, + char *const *data, const npy_intp *dimensions, const npy_intp *strides, + NpyAuxData *transferdata); + + +typedef int (PyArrayMethod_GetLoop)( + PyArrayMethod_Context *context, + int aligned, int move_references, + const npy_intp *strides, + PyArrayMethod_StridedLoop **out_loop, + NpyAuxData **out_transferdata, + NPY_ARRAYMETHOD_FLAGS *flags); + +/** + * Query an ArrayMethod for the initial value for use in reduction. + * + * @param context The arraymethod context, mainly to access the descriptors. + * @param reduction_is_empty Whether the reduction is empty. When it is, the + * value returned may differ. In this case it is a "default" value that + * may differ from the "identity" value normally used. For example: + * - `0.0` is the default for `sum([])`. But `-0.0` is the correct + * identity otherwise as it preserves the sign for `sum([-0.0])`. + * - We use no identity for object, but return the default of `0` and `1` + * for the empty `sum([], dtype=object)` and `prod([], dtype=object)`. + * This allows `np.sum(np.array(["a", "b"], dtype=object))` to work. + * - `-inf` or `INT_MIN` for `max` is an identity, but at least `INT_MIN` + * not a good *default* when there are no items. + * @param initial Pointer to initial data to be filled (if possible) + * + * @returns -1, 0, or 1 indicating error, no initial value, and initial being + * successfully filled. Errors must not be given where 0 is correct, NumPy + * may call this even when not strictly necessary. + */ +typedef int (PyArrayMethod_GetReductionInitial)( + PyArrayMethod_Context *context, npy_bool reduction_is_empty, + void *initial); + +/* + * The following functions are only used by the wrapping array method defined + * in umath/wrapping_array_method.c + */ + + +/* + * The function to convert the given descriptors (passed in to + * `resolve_descriptors`) and translates them for the wrapped loop. + * The new descriptors MUST be viewable with the old ones, `NULL` must be + * supported (for outputs) and should normally be forwarded. + * + * The function must clean up on error. + * + * NOTE: We currently assume that this translation gives "viewable" results. + * I.e. there is no additional casting related to the wrapping process. + * In principle that could be supported, but not sure it is useful. + * This currently also means that e.g. alignment must apply identically + * to the new dtypes. + * + * TODO: Due to the fact that `resolve_descriptors` is also used for `can_cast` + * there is no way to "pass out" the result of this function. This means + * it will be called twice for every ufunc call. + * (I am considering including `auxdata` as an "optional" parameter to + * `resolve_descriptors`, so that it can be filled there if not NULL.) + */ +typedef int (PyArrayMethod_TranslateGivenDescriptors)(int nin, int nout, + PyArray_DTypeMeta *const wrapped_dtypes[], + PyArray_Descr *const given_descrs[], PyArray_Descr *new_descrs[]); + +/** + * The function to convert the actual loop descriptors (as returned by the + * original `resolve_descriptors` function) to the ones the output array + * should use. + * This function must return "viewable" types, it must not mutate them in any + * form that would break the inner-loop logic. Does not need to support NULL. + * + * The function must clean up on error. + * + * @param nin Number of input arguments + * @param nout Number of output arguments + * @param new_dtypes The DTypes of the output (usually probably not needed) + * @param given_descrs Original given_descrs to the resolver, necessary to + * fetch any information related to the new dtypes from the original. + * @param original_descrs The `loop_descrs` returned by the wrapped loop. + * @param loop_descrs The output descriptors, compatible to `original_descrs`. + * + * @returns 0 on success, -1 on failure. + */ +typedef int (PyArrayMethod_TranslateLoopDescriptors)(int nin, int nout, + PyArray_DTypeMeta *const new_dtypes[], PyArray_Descr *const given_descrs[], + PyArray_Descr *original_descrs[], PyArray_Descr *loop_descrs[]); + + + +/* + * A traverse loop working on a single array. This is similar to the general + * strided-loop function. This is designed for loops that need to visit every + * element of a single array. + * + * Currently this is used for array clearing, via the NPY_DT_get_clear_loop + * API hook, and zero-filling, via the NPY_DT_get_fill_zero_loop API hook. + * These are most useful for handling arrays storing embedded references to + * python objects or heap-allocated data. + * + * The `void *traverse_context` is passed in because we may need to pass in + * Interpreter state or similar in the future, but we don't want to pass in + * a full context (with pointers to dtypes, method, caller which all make + * no sense for a traverse function). + * + * We assume for now that this context can be just passed through in the + * the future (for structured dtypes). + * + */ +typedef int (PyArrayMethod_TraverseLoop)( + void *traverse_context, const PyArray_Descr *descr, char *data, + npy_intp size, npy_intp stride, NpyAuxData *auxdata); + + +/* + * Simplified get_loop function specific to dtype traversal + * + * It should set the flags needed for the traversal loop and set out_loop to the + * loop function, which must be a valid PyArrayMethod_TraverseLoop + * pointer. Currently this is used for zero-filling and clearing arrays storing + * embedded references. + * + */ +typedef int (PyArrayMethod_GetTraverseLoop)( + void *traverse_context, const PyArray_Descr *descr, + int aligned, npy_intp fixed_stride, + PyArrayMethod_TraverseLoop **out_loop, NpyAuxData **out_auxdata, + NPY_ARRAYMETHOD_FLAGS *flags); + + +/* + * Type of the C promoter function, which must be wrapped into a + * PyCapsule with name "numpy._ufunc_promoter". + * + * Note that currently the output dtypes are always NULL unless they are + * also part of the signature. This is an implementation detail and could + * change in the future. However, in general promoters should not have a + * need for output dtypes. + * (There are potential use-cases, these are currently unsupported.) + */ +typedef int (PyArrayMethod_PromoterFunction)(PyObject *ufunc, + PyArray_DTypeMeta *const op_dtypes[], PyArray_DTypeMeta *const signature[], + PyArray_DTypeMeta *new_op_dtypes[]); + +/* + * **************************** + * DTYPE API + * **************************** + */ + +#define NPY_DT_ABSTRACT 1 << 1 +#define NPY_DT_PARAMETRIC 1 << 2 +#define NPY_DT_NUMERIC 1 << 3 + +/* + * These correspond to slots in the NPY_DType_Slots struct and must + * be in the same order as the members of that struct. If new slots + * get added or old slots get removed NPY_NUM_DTYPE_SLOTS must also + * be updated + */ + +#define NPY_DT_discover_descr_from_pyobject 1 +// this slot is considered private because its API hasn't been decided +#define _NPY_DT_is_known_scalar_type 2 +#define NPY_DT_default_descr 3 +#define NPY_DT_common_dtype 4 +#define NPY_DT_common_instance 5 +#define NPY_DT_ensure_canonical 6 +#define NPY_DT_setitem 7 +#define NPY_DT_getitem 8 +#define NPY_DT_get_clear_loop 9 +#define NPY_DT_get_fill_zero_loop 10 +#define NPY_DT_finalize_descr 11 + +// These PyArray_ArrFunc slots will be deprecated and replaced eventually +// getitem and setitem can be defined as a performance optimization; +// by default the user dtypes call `legacy_getitem_using_DType` and +// `legacy_setitem_using_DType`, respectively. This functionality is +// only supported for basic NumPy DTypes. + + +// used to separate dtype slots from arrfuncs slots +// intended only for internal use but defined here for clarity +#define _NPY_DT_ARRFUNCS_OFFSET (1 << 10) + +// Cast is disabled +// #define NPY_DT_PyArray_ArrFuncs_cast 0 + _NPY_DT_ARRFUNCS_OFFSET + +#define NPY_DT_PyArray_ArrFuncs_getitem 1 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_setitem 2 + _NPY_DT_ARRFUNCS_OFFSET + +// Copyswap is disabled +// #define NPY_DT_PyArray_ArrFuncs_copyswapn 3 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_copyswap 4 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_compare 5 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_argmax 6 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_dotfunc 7 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_scanfunc 8 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_fromstr 9 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_nonzero 10 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_fill 11 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_fillwithscalar 12 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_sort 13 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_argsort 14 + _NPY_DT_ARRFUNCS_OFFSET + +// Casting related slots are disabled. See +// https://github.com/numpy/numpy/pull/23173#discussion_r1101098163 +// #define NPY_DT_PyArray_ArrFuncs_castdict 15 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_scalarkind 16 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_cancastscalarkindto 17 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_cancastto 18 + _NPY_DT_ARRFUNCS_OFFSET + +// These are deprecated in NumPy 1.19, so are disabled here. +// #define NPY_DT_PyArray_ArrFuncs_fastclip 19 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_fastputmask 20 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_fasttake 21 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_argmin 22 + _NPY_DT_ARRFUNCS_OFFSET + + +// TODO: These slots probably still need some thought, and/or a way to "grow"? +typedef struct { + PyTypeObject *typeobj; /* type of python scalar or NULL */ + int flags; /* flags, including parametric and abstract */ + /* NULL terminated cast definitions. Use NULL for the newly created DType */ + PyArrayMethod_Spec **casts; + PyType_Slot *slots; + /* Baseclass or NULL (will always subclass `np.dtype`) */ + PyTypeObject *baseclass; +} PyArrayDTypeMeta_Spec; + + +typedef PyArray_Descr *(PyArrayDTypeMeta_DiscoverDescrFromPyobject)( + PyArray_DTypeMeta *cls, PyObject *obj); + +/* + * Before making this public, we should decide whether it should pass + * the type, or allow looking at the object. A possible use-case: + * `np.array(np.array([0]), dtype=np.ndarray)` + * Could consider arrays that are not `dtype=ndarray` "scalars". + */ +typedef int (PyArrayDTypeMeta_IsKnownScalarType)( + PyArray_DTypeMeta *cls, PyTypeObject *obj); + +typedef PyArray_Descr *(PyArrayDTypeMeta_DefaultDescriptor)(PyArray_DTypeMeta *cls); +typedef PyArray_DTypeMeta *(PyArrayDTypeMeta_CommonDType)( + PyArray_DTypeMeta *dtype1, PyArray_DTypeMeta *dtype2); + + +/* + * Convenience utility for getting a reference to the DType metaclass associated + * with a dtype instance. + */ +#define NPY_DTYPE(descr) ((PyArray_DTypeMeta *)Py_TYPE(descr)) + +static inline PyArray_DTypeMeta * +NPY_DT_NewRef(PyArray_DTypeMeta *o) { + Py_INCREF((PyObject *)o); + return o; +} + + +typedef PyArray_Descr *(PyArrayDTypeMeta_CommonInstance)( + PyArray_Descr *dtype1, PyArray_Descr *dtype2); +typedef PyArray_Descr *(PyArrayDTypeMeta_EnsureCanonical)(PyArray_Descr *dtype); +/* + * Returns either a new reference to *dtype* or a new descriptor instance + * initialized with the same parameters as *dtype*. The caller cannot know + * which choice a dtype will make. This function is called just before the + * array buffer is created for a newly created array, it is not called for + * views and the descriptor returned by this function is attached to the array. + */ +typedef PyArray_Descr *(PyArrayDTypeMeta_FinalizeDescriptor)(PyArray_Descr *dtype); + +/* + * TODO: These two functions are currently only used for experimental DType + * API support. Their relation should be "reversed": NumPy should + * always use them internally. + * There are open points about "casting safety" though, e.g. setting + * elements is currently always unsafe. + */ +typedef int(PyArrayDTypeMeta_SetItem)(PyArray_Descr *, PyObject *, char *); +typedef PyObject *(PyArrayDTypeMeta_GetItem)(PyArray_Descr *, char *); + +#endif /* NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ */ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/halffloat.h b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/halffloat.h new file mode 100644 index 0000000000000000000000000000000000000000..1365dc0e4f3c3523b32d6d7cbebc0d107627cf40 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/halffloat.h @@ -0,0 +1,70 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_HALFFLOAT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_HALFFLOAT_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * Half-precision routines + */ + +/* Conversions */ +float npy_half_to_float(npy_half h); +double npy_half_to_double(npy_half h); +npy_half npy_float_to_half(float f); +npy_half npy_double_to_half(double d); +/* Comparisons */ +int npy_half_eq(npy_half h1, npy_half h2); +int npy_half_ne(npy_half h1, npy_half h2); +int npy_half_le(npy_half h1, npy_half h2); +int npy_half_lt(npy_half h1, npy_half h2); +int npy_half_ge(npy_half h1, npy_half h2); +int npy_half_gt(npy_half h1, npy_half h2); +/* faster *_nonan variants for when you know h1 and h2 are not NaN */ +int npy_half_eq_nonan(npy_half h1, npy_half h2); +int npy_half_lt_nonan(npy_half h1, npy_half h2); +int npy_half_le_nonan(npy_half h1, npy_half h2); +/* Miscellaneous functions */ +int npy_half_iszero(npy_half h); +int npy_half_isnan(npy_half h); +int npy_half_isinf(npy_half h); +int npy_half_isfinite(npy_half h); +int npy_half_signbit(npy_half h); +npy_half npy_half_copysign(npy_half x, npy_half y); +npy_half npy_half_spacing(npy_half h); +npy_half npy_half_nextafter(npy_half x, npy_half y); +npy_half npy_half_divmod(npy_half x, npy_half y, npy_half *modulus); + +/* + * Half-precision constants + */ + +#define NPY_HALF_ZERO (0x0000u) +#define NPY_HALF_PZERO (0x0000u) +#define NPY_HALF_NZERO (0x8000u) +#define NPY_HALF_ONE (0x3c00u) +#define NPY_HALF_NEGONE (0xbc00u) +#define NPY_HALF_PINF (0x7c00u) +#define NPY_HALF_NINF (0xfc00u) +#define NPY_HALF_NAN (0x7e00u) + +#define NPY_MAX_HALF (0x7bffu) + +/* + * Bit-level conversions + */ + +npy_uint16 npy_floatbits_to_halfbits(npy_uint32 f); +npy_uint16 npy_doublebits_to_halfbits(npy_uint64 d); +npy_uint32 npy_halfbits_to_floatbits(npy_uint16 h); +npy_uint64 npy_halfbits_to_doublebits(npy_uint16 h); + +#ifdef __cplusplus +} +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_HALFFLOAT_H_ */ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/ndarrayobject.h b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/ndarrayobject.h new file mode 100644 index 0000000000000000000000000000000000000000..40f213077dba6abf1f51922e1914516b04b22ddc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/include/numpy/ndarrayobject.h @@ -0,0 +1,304 @@ +/* + * DON'T INCLUDE THIS DIRECTLY. + */ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NDARRAYOBJECT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NDARRAYOBJECT_H_ + +#ifdef __cplusplus +extern "C" { +#endif + +#include +#include "ndarraytypes.h" +#include "dtype_api.h" + +/* Includes the "function" C-API -- these are all stored in a + list of pointers --- one for each file + The two lists are concatenated into one in multiarray. + + They are available as import_array() +*/ + +#include "__multiarray_api.h" + +/* + * Include any definitions which are defined differently for 1.x and 2.x + * (Symbols only available on 2.x are not there, but rather guarded.) + */ +#include "npy_2_compat.h" + +/* C-API that requires previous API to be defined */ + +#define PyArray_DescrCheck(op) PyObject_TypeCheck(op, &PyArrayDescr_Type) + +#define PyArray_Check(op) PyObject_TypeCheck(op, &PyArray_Type) +#define PyArray_CheckExact(op) (((PyObject*)(op))->ob_type == &PyArray_Type) + +#define PyArray_HasArrayInterfaceType(op, type, context, out) \ + ((((out)=PyArray_FromStructInterface(op)) != Py_NotImplemented) || \ + (((out)=PyArray_FromInterface(op)) != Py_NotImplemented) || \ + (((out)=PyArray_FromArrayAttr(op, type, context)) != \ + Py_NotImplemented)) + +#define PyArray_HasArrayInterface(op, out) \ + PyArray_HasArrayInterfaceType(op, NULL, NULL, out) + +#define PyArray_IsZeroDim(op) (PyArray_Check(op) && \ + (PyArray_NDIM((PyArrayObject *)op) == 0)) + +#define PyArray_IsScalar(obj, cls) \ + (PyObject_TypeCheck(obj, &Py##cls##ArrType_Type)) + +#define PyArray_CheckScalar(m) (PyArray_IsScalar(m, Generic) || \ + PyArray_IsZeroDim(m)) +#define PyArray_IsPythonNumber(obj) \ + (PyFloat_Check(obj) || PyComplex_Check(obj) || \ + PyLong_Check(obj) || PyBool_Check(obj)) +#define PyArray_IsIntegerScalar(obj) (PyLong_Check(obj) \ + || PyArray_IsScalar((obj), Integer)) +#define PyArray_IsPythonScalar(obj) \ + (PyArray_IsPythonNumber(obj) || PyBytes_Check(obj) || \ + PyUnicode_Check(obj)) + +#define PyArray_IsAnyScalar(obj) \ + (PyArray_IsScalar(obj, Generic) || PyArray_IsPythonScalar(obj)) + +#define PyArray_CheckAnyScalar(obj) (PyArray_IsPythonScalar(obj) || \ + PyArray_CheckScalar(obj)) + + +#define PyArray_GETCONTIGUOUS(m) (PyArray_ISCONTIGUOUS(m) ? \ + Py_INCREF(m), (m) : \ + (PyArrayObject *)(PyArray_Copy(m))) + +#define PyArray_SAMESHAPE(a1,a2) ((PyArray_NDIM(a1) == PyArray_NDIM(a2)) && \ + PyArray_CompareLists(PyArray_DIMS(a1), \ + PyArray_DIMS(a2), \ + PyArray_NDIM(a1))) + +#define PyArray_SIZE(m) PyArray_MultiplyList(PyArray_DIMS(m), PyArray_NDIM(m)) +#define PyArray_NBYTES(m) (PyArray_ITEMSIZE(m) * PyArray_SIZE(m)) +#define PyArray_FROM_O(m) PyArray_FromAny(m, NULL, 0, 0, 0, NULL) + +#define PyArray_FROM_OF(m,flags) PyArray_CheckFromAny(m, NULL, 0, 0, flags, \ + NULL) + +#define PyArray_FROM_OT(m,type) PyArray_FromAny(m, \ + PyArray_DescrFromType(type), 0, 0, 0, NULL) + +#define PyArray_FROM_OTF(m, type, flags) \ + PyArray_FromAny(m, PyArray_DescrFromType(type), 0, 0, \ + (((flags) & NPY_ARRAY_ENSURECOPY) ? \ + ((flags) | NPY_ARRAY_DEFAULT) : (flags)), NULL) + +#define PyArray_FROMANY(m, type, min, max, flags) \ + PyArray_FromAny(m, PyArray_DescrFromType(type), min, max, \ + (((flags) & NPY_ARRAY_ENSURECOPY) ? \ + (flags) | NPY_ARRAY_DEFAULT : (flags)), NULL) + +#define PyArray_ZEROS(m, dims, type, is_f_order) \ + PyArray_Zeros(m, dims, PyArray_DescrFromType(type), is_f_order) + +#define PyArray_EMPTY(m, dims, type, is_f_order) \ + PyArray_Empty(m, dims, PyArray_DescrFromType(type), is_f_order) + +#define PyArray_FILLWBYTE(obj, val) memset(PyArray_DATA(obj), val, \ + PyArray_NBYTES(obj)) + +#define PyArray_ContiguousFromAny(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_DEFAULT, NULL) + +#define PyArray_EquivArrTypes(a1, a2) \ + PyArray_EquivTypes(PyArray_DESCR(a1), PyArray_DESCR(a2)) + +#define PyArray_EquivByteorders(b1, b2) \ + (((b1) == (b2)) || (PyArray_ISNBO(b1) == PyArray_ISNBO(b2))) + +#define PyArray_SimpleNew(nd, dims, typenum) \ + PyArray_New(&PyArray_Type, nd, dims, typenum, NULL, NULL, 0, 0, NULL) + +#define PyArray_SimpleNewFromData(nd, dims, typenum, data) \ + PyArray_New(&PyArray_Type, nd, dims, typenum, NULL, \ + data, 0, NPY_ARRAY_CARRAY, NULL) + +#define PyArray_SimpleNewFromDescr(nd, dims, descr) \ + PyArray_NewFromDescr(&PyArray_Type, descr, nd, dims, \ + NULL, NULL, 0, NULL) + +#define PyArray_ToScalar(data, arr) \ + PyArray_Scalar(data, PyArray_DESCR(arr), (PyObject *)arr) + + +/* These might be faster without the dereferencing of obj + going on inside -- of course an optimizing compiler should + inline the constants inside a for loop making it a moot point +*/ + +#define PyArray_GETPTR1(obj, i) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0])) + +#define PyArray_GETPTR2(obj, i, j) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0] + \ + (j)*PyArray_STRIDES(obj)[1])) + +#define PyArray_GETPTR3(obj, i, j, k) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0] + \ + (j)*PyArray_STRIDES(obj)[1] + \ + (k)*PyArray_STRIDES(obj)[2])) + +#define PyArray_GETPTR4(obj, i, j, k, l) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0] + \ + (j)*PyArray_STRIDES(obj)[1] + \ + (k)*PyArray_STRIDES(obj)[2] + \ + (l)*PyArray_STRIDES(obj)[3])) + +static inline void +PyArray_DiscardWritebackIfCopy(PyArrayObject *arr) +{ + PyArrayObject_fields *fa = (PyArrayObject_fields *)arr; + if (fa && fa->base) { + if (fa->flags & NPY_ARRAY_WRITEBACKIFCOPY) { + PyArray_ENABLEFLAGS((PyArrayObject*)fa->base, NPY_ARRAY_WRITEABLE); + Py_DECREF(fa->base); + fa->base = NULL; + PyArray_CLEARFLAGS(arr, NPY_ARRAY_WRITEBACKIFCOPY); + } + } +} + +#define PyArray_DESCR_REPLACE(descr) do { \ + PyArray_Descr *_new_; \ + _new_ = PyArray_DescrNew(descr); \ + Py_XDECREF(descr); \ + descr = _new_; \ + } while(0) + +/* Copy should always return contiguous array */ +#define PyArray_Copy(obj) PyArray_NewCopy(obj, NPY_CORDER) + +#define PyArray_FromObject(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_BEHAVED | \ + NPY_ARRAY_ENSUREARRAY, NULL) + +#define PyArray_ContiguousFromObject(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_DEFAULT | \ + NPY_ARRAY_ENSUREARRAY, NULL) + +#define PyArray_CopyFromObject(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_ENSURECOPY | \ + NPY_ARRAY_DEFAULT | \ + NPY_ARRAY_ENSUREARRAY, NULL) + +#define PyArray_Cast(mp, type_num) \ + PyArray_CastToType(mp, PyArray_DescrFromType(type_num), 0) + +#define PyArray_Take(ap, items, axis) \ + PyArray_TakeFrom(ap, items, axis, NULL, NPY_RAISE) + +#define PyArray_Put(ap, items, values) \ + PyArray_PutTo(ap, items, values, NPY_RAISE) + + +/* + Check to see if this key in the dictionary is the "title" + entry of the tuple (i.e. a duplicate dictionary entry in the fields + dict). +*/ + +static inline int +NPY_TITLE_KEY_check(PyObject *key, PyObject *value) +{ + PyObject *title; + if (PyTuple_Size(value) != 3) { + return 0; + } + title = PyTuple_GetItem(value, 2); + if (key == title) { + return 1; + } +#ifdef PYPY_VERSION + /* + * On PyPy, dictionary keys do not always preserve object identity. + * Fall back to comparison by value. + */ + if (PyUnicode_Check(title) && PyUnicode_Check(key)) { + return PyUnicode_Compare(title, key) == 0 ? 1 : 0; + } +#endif + return 0; +} + +/* Macro, for backward compat with "if NPY_TITLE_KEY(key, value) { ..." */ +#define NPY_TITLE_KEY(key, value) (NPY_TITLE_KEY_check((key), (value))) + +#define DEPRECATE(msg) PyErr_WarnEx(PyExc_DeprecationWarning,msg,1) +#define DEPRECATE_FUTUREWARNING(msg) PyErr_WarnEx(PyExc_FutureWarning,msg,1) + + +/* + * These macros and functions unfortunately require runtime version checks + * that are only defined in `npy_2_compat.h`. For that reasons they cannot be + * part of `ndarraytypes.h` which tries to be self contained. + */ + +static inline npy_intp +PyArray_ITEMSIZE(const PyArrayObject *arr) +{ + return PyDataType_ELSIZE(((PyArrayObject_fields *)arr)->descr); +} + +#define PyDataType_HASFIELDS(obj) (PyDataType_ISLEGACY((PyArray_Descr*)(obj)) && PyDataType_NAMES((PyArray_Descr*)(obj)) != NULL) +#define PyDataType_HASSUBARRAY(dtype) (PyDataType_ISLEGACY(dtype) && PyDataType_SUBARRAY(dtype) != NULL) +#define PyDataType_ISUNSIZED(dtype) ((dtype)->elsize == 0 && \ + !PyDataType_HASFIELDS(dtype)) + +#define PyDataType_FLAGCHK(dtype, flag) \ + ((PyDataType_FLAGS(dtype) & (flag)) == (flag)) + +#define PyDataType_REFCHK(dtype) \ + PyDataType_FLAGCHK(dtype, NPY_ITEM_REFCOUNT) + +#define NPY_BEGIN_THREADS_DESCR(dtype) \ + do {if (!(PyDataType_FLAGCHK((dtype), NPY_NEEDS_PYAPI))) \ + NPY_BEGIN_THREADS;} while (0); + +#define NPY_END_THREADS_DESCR(dtype) \ + do {if (!(PyDataType_FLAGCHK((dtype), NPY_NEEDS_PYAPI))) \ + NPY_END_THREADS; } while (0); + +#if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD) +/* The internal copy of this is now defined in `dtypemeta.h` */ +/* + * `PyArray_Scalar` is the same as this function but converts will convert + * most NumPy types to Python scalars. + */ +static inline PyObject * +PyArray_GETITEM(const PyArrayObject *arr, const char *itemptr) +{ + return PyDataType_GetArrFuncs(((PyArrayObject_fields *)arr)->descr)->getitem( + (void *)itemptr, (PyArrayObject *)arr); +} + +/* + * SETITEM should only be used if it is known that the value is a scalar + * and of a type understood by the arrays dtype. + * Use `PyArray_Pack` if the value may be of a different dtype. + */ +static inline int +PyArray_SETITEM(PyArrayObject *arr, char *itemptr, PyObject *v) +{ + return PyDataType_GetArrFuncs(((PyArrayObject_fields *)arr)->descr)->setitem(v, itemptr, arr); +} +#endif /* not internal */ + + +#ifdef __cplusplus +} +#endif + + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NDARRAYOBJECT_H_ */ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/lib/npymath.lib b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/lib/npymath.lib new file mode 100644 index 0000000000000000000000000000000000000000..b994f41abf8754421cd1ffa0a79e944eec34317f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/lib/npymath.lib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3d60e1d7c58dc6e5f772c4887f466738bde84c67259898f6d55a33fd1ef6298c +size 154174 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/memmap.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/memmap.py new file mode 100644 index 0000000000000000000000000000000000000000..1c4b7789b5b4def27fd03625f795c3e48f33b5e6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/memmap.py @@ -0,0 +1,361 @@ +from contextlib import nullcontext +import operator +import numpy as np +from .._utils import set_module +from .numeric import uint8, ndarray, dtype + +__all__ = ['memmap'] + +dtypedescr = dtype +valid_filemodes = ["r", "c", "r+", "w+"] +writeable_filemodes = ["r+", "w+"] + +mode_equivalents = { + "readonly":"r", + "copyonwrite":"c", + "readwrite":"r+", + "write":"w+" + } + + +@set_module('numpy') +class memmap(ndarray): + """Create a memory-map to an array stored in a *binary* file on disk. + + Memory-mapped files are used for accessing small segments of large files + on disk, without reading the entire file into memory. NumPy's + memmap's are array-like objects. This differs from Python's ``mmap`` + module, which uses file-like objects. + + This subclass of ndarray has some unpleasant interactions with + some operations, because it doesn't quite fit properly as a subclass. + An alternative to using this subclass is to create the ``mmap`` + object yourself, then create an ndarray with ndarray.__new__ directly, + passing the object created in its 'buffer=' parameter. + + This class may at some point be turned into a factory function + which returns a view into an mmap buffer. + + Flush the memmap instance to write the changes to the file. Currently there + is no API to close the underlying ``mmap``. It is tricky to ensure the + resource is actually closed, since it may be shared between different + memmap instances. + + + Parameters + ---------- + filename : str, file-like object, or pathlib.Path instance + The file name or file object to be used as the array data buffer. + dtype : data-type, optional + The data-type used to interpret the file contents. + Default is `uint8`. + mode : {'r+', 'r', 'w+', 'c'}, optional + The file is opened in this mode: + + +------+-------------------------------------------------------------+ + | 'r' | Open existing file for reading only. | + +------+-------------------------------------------------------------+ + | 'r+' | Open existing file for reading and writing. | + +------+-------------------------------------------------------------+ + | 'w+' | Create or overwrite existing file for reading and writing. | + | | If ``mode == 'w+'`` then `shape` must also be specified. | + +------+-------------------------------------------------------------+ + | 'c' | Copy-on-write: assignments affect data in memory, but | + | | changes are not saved to disk. The file on disk is | + | | read-only. | + +------+-------------------------------------------------------------+ + + Default is 'r+'. + offset : int, optional + In the file, array data starts at this offset. Since `offset` is + measured in bytes, it should normally be a multiple of the byte-size + of `dtype`. When ``mode != 'r'``, even positive offsets beyond end of + file are valid; The file will be extended to accommodate the + additional data. By default, ``memmap`` will start at the beginning of + the file, even if ``filename`` is a file pointer ``fp`` and + ``fp.tell() != 0``. + shape : int or sequence of ints, optional + The desired shape of the array. If ``mode == 'r'`` and the number + of remaining bytes after `offset` is not a multiple of the byte-size + of `dtype`, you must specify `shape`. By default, the returned array + will be 1-D with the number of elements determined by file size + and data-type. + + .. versionchanged:: 2.0 + The shape parameter can now be any integer sequence type, previously + types were limited to tuple and int. + + order : {'C', 'F'}, optional + Specify the order of the ndarray memory layout: + :term:`row-major`, C-style or :term:`column-major`, + Fortran-style. This only has an effect if the shape is + greater than 1-D. The default order is 'C'. + + Attributes + ---------- + filename : str or pathlib.Path instance + Path to the mapped file. + offset : int + Offset position in the file. + mode : str + File mode. + + Methods + ------- + flush + Flush any changes in memory to file on disk. + When you delete a memmap object, flush is called first to write + changes to disk. + + + See also + -------- + lib.format.open_memmap : Create or load a memory-mapped ``.npy`` file. + + Notes + ----- + The memmap object can be used anywhere an ndarray is accepted. + Given a memmap ``fp``, ``isinstance(fp, numpy.ndarray)`` returns + ``True``. + + Memory-mapped files cannot be larger than 2GB on 32-bit systems. + + When a memmap causes a file to be created or extended beyond its + current size in the filesystem, the contents of the new part are + unspecified. On systems with POSIX filesystem semantics, the extended + part will be filled with zero bytes. + + Examples + -------- + >>> import numpy as np + >>> data = np.arange(12, dtype='float32') + >>> data.resize((3,4)) + + This example uses a temporary file so that doctest doesn't write + files to your directory. You would use a 'normal' filename. + + >>> from tempfile import mkdtemp + >>> import os.path as path + >>> filename = path.join(mkdtemp(), 'newfile.dat') + + Create a memmap with dtype and shape that matches our data: + + >>> fp = np.memmap(filename, dtype='float32', mode='w+', shape=(3,4)) + >>> fp + memmap([[0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.]], dtype=float32) + + Write data to memmap array: + + >>> fp[:] = data[:] + >>> fp + memmap([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], dtype=float32) + + >>> fp.filename == path.abspath(filename) + True + + Flushes memory changes to disk in order to read them back + + >>> fp.flush() + + Load the memmap and verify data was stored: + + >>> newfp = np.memmap(filename, dtype='float32', mode='r', shape=(3,4)) + >>> newfp + memmap([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], dtype=float32) + + Read-only memmap: + + >>> fpr = np.memmap(filename, dtype='float32', mode='r', shape=(3,4)) + >>> fpr.flags.writeable + False + + Copy-on-write memmap: + + >>> fpc = np.memmap(filename, dtype='float32', mode='c', shape=(3,4)) + >>> fpc.flags.writeable + True + + It's possible to assign to copy-on-write array, but values are only + written into the memory copy of the array, and not written to disk: + + >>> fpc + memmap([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], dtype=float32) + >>> fpc[0,:] = 0 + >>> fpc + memmap([[ 0., 0., 0., 0.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], dtype=float32) + + File on disk is unchanged: + + >>> fpr + memmap([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]], dtype=float32) + + Offset into a memmap: + + >>> fpo = np.memmap(filename, dtype='float32', mode='r', offset=16) + >>> fpo + memmap([ 4., 5., 6., 7., 8., 9., 10., 11.], dtype=float32) + + """ + + __array_priority__ = -100.0 + + def __new__(subtype, filename, dtype=uint8, mode='r+', offset=0, + shape=None, order='C'): + # Import here to minimize 'import numpy' overhead + import mmap + import os.path + try: + mode = mode_equivalents[mode] + except KeyError as e: + if mode not in valid_filemodes: + raise ValueError( + "mode must be one of {!r} (got {!r})" + .format(valid_filemodes + list(mode_equivalents.keys()), mode) + ) from None + + if mode == 'w+' and shape is None: + raise ValueError("shape must be given if mode == 'w+'") + + if hasattr(filename, 'read'): + f_ctx = nullcontext(filename) + else: + f_ctx = open( + os.fspath(filename), + ('r' if mode == 'c' else mode)+'b' + ) + + with f_ctx as fid: + fid.seek(0, 2) + flen = fid.tell() + descr = dtypedescr(dtype) + _dbytes = descr.itemsize + + if shape is None: + bytes = flen - offset + if bytes % _dbytes: + raise ValueError("Size of available data is not a " + "multiple of the data-type size.") + size = bytes // _dbytes + shape = (size,) + else: + if type(shape) not in (tuple, list): + try: + shape = [operator.index(shape)] + except TypeError: + pass + shape = tuple(shape) + size = np.intp(1) # avoid default choice of np.int_, which might overflow + for k in shape: + size *= k + + bytes = int(offset + size*_dbytes) + + if mode in ('w+', 'r+'): + # gh-27723 + # if bytes == 0, we write out 1 byte to allow empty memmap. + bytes = max(bytes, 1) + if flen < bytes: + fid.seek(bytes - 1, 0) + fid.write(b'\0') + fid.flush() + + if mode == 'c': + acc = mmap.ACCESS_COPY + elif mode == 'r': + acc = mmap.ACCESS_READ + else: + acc = mmap.ACCESS_WRITE + + start = offset - offset % mmap.ALLOCATIONGRANULARITY + bytes -= start + # bytes == 0 is problematic as in mmap length=0 maps the full file. + # See PR gh-27723 for a more detailed explanation. + if bytes == 0 and start > 0: + bytes += mmap.ALLOCATIONGRANULARITY + start -= mmap.ALLOCATIONGRANULARITY + array_offset = offset - start + mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start) + + self = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm, + offset=array_offset, order=order) + self._mmap = mm + self.offset = offset + self.mode = mode + + if isinstance(filename, os.PathLike): + # special case - if we were constructed with a pathlib.path, + # then filename is a path object, not a string + self.filename = filename.resolve() + elif hasattr(fid, "name") and isinstance(fid.name, str): + # py3 returns int for TemporaryFile().name + self.filename = os.path.abspath(fid.name) + # same as memmap copies (e.g. memmap + 1) + else: + self.filename = None + + return self + + def __array_finalize__(self, obj): + if hasattr(obj, '_mmap') and np.may_share_memory(self, obj): + self._mmap = obj._mmap + self.filename = obj.filename + self.offset = obj.offset + self.mode = obj.mode + else: + self._mmap = None + self.filename = None + self.offset = None + self.mode = None + + def flush(self): + """ + Write any changes in the array to the file on disk. + + For further information, see `memmap`. + + Parameters + ---------- + None + + See Also + -------- + memmap + + """ + if self.base is not None and hasattr(self.base, 'flush'): + self.base.flush() + + def __array_wrap__(self, arr, context=None, return_scalar=False): + arr = super().__array_wrap__(arr, context) + + # Return a memmap if a memmap was given as the output of the + # ufunc. Leave the arr class unchanged if self is not a memmap + # to keep original memmap subclasses behavior + if self is arr or type(self) is not memmap: + return arr + + # Return scalar instead of 0d memmap, e.g. for np.sum with + # axis=None (note that subclasses will not reach here) + if return_scalar: + return arr[()] + + # Return ndarray otherwise + return arr.view(np.ndarray) + + def __getitem__(self, index): + res = super().__getitem__(index) + if type(res) is memmap and res._mmap is None: + return res.view(type=ndarray) + return res diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/memmap.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/memmap.pyi new file mode 100644 index 0000000000000000000000000000000000000000..5e251f66c03139d74ef83520622b2bb0afc8e721 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/memmap.pyi @@ -0,0 +1,3 @@ +from numpy import memmap + +__all__ = ["memmap"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/multiarray.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/multiarray.py new file mode 100644 index 0000000000000000000000000000000000000000..dad1a23756ed636416988cd326e973ca5d0e18d6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/multiarray.py @@ -0,0 +1,1754 @@ +""" +Create the numpy._core.multiarray namespace for backward compatibility. +In v1.16 the multiarray and umath c-extension modules were merged into +a single _multiarray_umath extension module. So we replicate the old +namespace by importing from the extension module. + +""" + +import functools +from . import overrides +from . import _multiarray_umath +from ._multiarray_umath import * # noqa: F403 +# These imports are needed for backward compatibility, +# do not change them. issue gh-15518 +# _get_ndarray_c_version is semi-public, on purpose not added to __all__ +from ._multiarray_umath import ( + _flagdict, from_dlpack, _place, _reconstruct, + _vec_string, _ARRAY_API, _monotonicity, _get_ndarray_c_version, + _get_madvise_hugepage, _set_madvise_hugepage, + ) + +__all__ = [ + '_ARRAY_API', 'ALLOW_THREADS', 'BUFSIZE', 'CLIP', 'DATETIMEUNITS', + 'ITEM_HASOBJECT', 'ITEM_IS_POINTER', 'LIST_PICKLE', 'MAXDIMS', + 'MAY_SHARE_BOUNDS', 'MAY_SHARE_EXACT', 'NEEDS_INIT', 'NEEDS_PYAPI', + 'RAISE', 'USE_GETITEM', 'USE_SETITEM', 'WRAP', + '_flagdict', 'from_dlpack', '_place', '_reconstruct', '_vec_string', + '_monotonicity', 'add_docstring', 'arange', 'array', 'asarray', + 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount', + 'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast', + 'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2', + 'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data', + 'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype', + 'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat', + 'frombuffer', 'fromfile', 'fromiter', 'fromstring', + 'get_handler_name', 'get_handler_version', 'inner', 'interp', + 'interp_complex', 'is_busday', 'lexsort', 'matmul', 'vecdot', + 'may_share_memory', 'min_scalar_type', 'ndarray', 'nditer', 'nested_iters', + 'normalize_axis_index', 'packbits', 'promote_types', 'putmask', + 'ravel_multi_index', 'result_type', 'scalar', 'set_datetimeparse_function', + 'set_typeDict', 'shares_memory', 'typeinfo', + 'unpackbits', 'unravel_index', 'vdot', 'where', 'zeros'] + +# For backward compatibility, make sure pickle imports +# these functions from here +_reconstruct.__module__ = 'numpy._core.multiarray' +scalar.__module__ = 'numpy._core.multiarray' + + +from_dlpack.__module__ = 'numpy' +arange.__module__ = 'numpy' +array.__module__ = 'numpy' +asarray.__module__ = 'numpy' +asanyarray.__module__ = 'numpy' +ascontiguousarray.__module__ = 'numpy' +asfortranarray.__module__ = 'numpy' +datetime_data.__module__ = 'numpy' +empty.__module__ = 'numpy' +frombuffer.__module__ = 'numpy' +fromfile.__module__ = 'numpy' +fromiter.__module__ = 'numpy' +frompyfunc.__module__ = 'numpy' +fromstring.__module__ = 'numpy' +may_share_memory.__module__ = 'numpy' +nested_iters.__module__ = 'numpy' +promote_types.__module__ = 'numpy' +zeros.__module__ = 'numpy' +normalize_axis_index.__module__ = 'numpy.lib.array_utils' +add_docstring.__module__ = 'numpy.lib' +compare_chararrays.__module__ = 'numpy.char' + + +def _override___module__(): + namespace_names = globals() + for ufunc_name in [ + 'absolute', 'arccos', 'arccosh', 'add', 'arcsin', 'arcsinh', 'arctan', + 'arctan2', 'arctanh', 'bitwise_and', 'bitwise_count', 'invert', + 'left_shift', 'bitwise_or', 'right_shift', 'bitwise_xor', 'cbrt', + 'ceil', 'conjugate', 'copysign', 'cos', 'cosh', 'deg2rad', 'degrees', + 'divide', 'divmod', 'equal', 'exp', 'exp2', 'expm1', 'fabs', + 'float_power', 'floor', 'floor_divide', 'fmax', 'fmin', 'fmod', + 'frexp', 'gcd', 'greater', 'greater_equal', 'heaviside', 'hypot', + 'isfinite', 'isinf', 'isnan', 'isnat', 'lcm', 'ldexp', 'less', + 'less_equal', 'log', 'log10', 'log1p', 'log2', 'logaddexp', + 'logaddexp2', 'logical_and', 'logical_not', 'logical_or', + 'logical_xor', 'matmul', 'matvec', 'maximum', 'minimum', 'remainder', + 'modf', 'multiply', 'negative', 'nextafter', 'not_equal', 'positive', + 'power', 'rad2deg', 'radians', 'reciprocal', 'rint', 'sign', 'signbit', + 'sin', 'sinh', 'spacing', 'sqrt', 'square', 'subtract', 'tan', 'tanh', + 'trunc', 'vecdot', 'vecmat', + ]: + ufunc = namespace_names[ufunc_name] + ufunc.__module__ = "numpy" + ufunc.__qualname__ = ufunc_name + + +_override___module__() + + +# We can't verify dispatcher signatures because NumPy's C functions don't +# support introspection. +array_function_from_c_func_and_dispatcher = functools.partial( + overrides.array_function_from_dispatcher, + module='numpy', docs_from_dispatcher=True, verify=False) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.empty_like) +def empty_like( + prototype, dtype=None, order=None, subok=None, shape=None, *, device=None +): + """ + empty_like(prototype, dtype=None, order='K', subok=True, shape=None, *, + device=None) + + Return a new array with the same shape and type as a given array. + + Parameters + ---------- + prototype : array_like + The shape and data-type of `prototype` define these same attributes + of the returned array. + dtype : data-type, optional + Overrides the data type of the result. + order : {'C', 'F', 'A', or 'K'}, optional + Overrides the memory layout of the result. 'C' means C-order, + 'F' means F-order, 'A' means 'F' if `prototype` is Fortran + contiguous, 'C' otherwise. 'K' means match the layout of `prototype` + as closely as possible. + subok : bool, optional. + If True, then the newly created array will use the sub-class + type of `prototype`, otherwise it will be a base-class array. Defaults + to True. + shape : int or sequence of ints, optional. + Overrides the shape of the result. If order='K' and the number of + dimensions is unchanged, will try to keep order, otherwise, + order='C' is implied. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + + Returns + ------- + out : ndarray + Array of uninitialized (arbitrary) data with the same + shape and type as `prototype`. + + See Also + -------- + ones_like : Return an array of ones with shape and type of input. + zeros_like : Return an array of zeros with shape and type of input. + full_like : Return a new array with shape of input filled with value. + empty : Return a new uninitialized array. + + Notes + ----- + Unlike other array creation functions (e.g. `zeros_like`, `ones_like`, + `full_like`), `empty_like` does not initialize the values of the array, + and may therefore be marginally faster. However, the values stored in the + newly allocated array are arbitrary. For reproducible behavior, be sure + to set each element of the array before reading. + + Examples + -------- + >>> import numpy as np + >>> a = ([1,2,3], [4,5,6]) # a is array-like + >>> np.empty_like(a) + array([[-1073741821, -1073741821, 3], # uninitialized + [ 0, 0, -1073741821]]) + >>> a = np.array([[1., 2., 3.],[4.,5.,6.]]) + >>> np.empty_like(a) + array([[ -2.00000715e+000, 1.48219694e-323, -2.00000572e+000], # uninitialized + [ 4.38791518e-305, -2.00000715e+000, 4.17269252e-309]]) + + """ # NOQA + return (prototype,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.concatenate) +def concatenate(arrays, axis=None, out=None, *, dtype=None, casting=None): + """ + concatenate( + (a1, a2, ...), + axis=0, + out=None, + dtype=None, + casting="same_kind" + ) + + Join a sequence of arrays along an existing axis. + + Parameters + ---------- + a1, a2, ... : sequence of array_like + The arrays must have the same shape, except in the dimension + corresponding to `axis` (the first, by default). + axis : int, optional + The axis along which the arrays will be joined. If axis is None, + arrays are flattened before use. Default is 0. + out : ndarray, optional + If provided, the destination to place the result. The shape must be + correct, matching that of what concatenate would have returned if no + out argument were specified. + dtype : str or dtype + If provided, the destination array will have this dtype. Cannot be + provided together with `out`. + + .. versionadded:: 1.20.0 + + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Defaults to 'same_kind'. + For a description of the options, please see :term:`casting`. + + .. versionadded:: 1.20.0 + + Returns + ------- + res : ndarray + The concatenated array. + + See Also + -------- + ma.concatenate : Concatenate function that preserves input masks. + array_split : Split an array into multiple sub-arrays of equal or + near-equal size. + split : Split array into a list of multiple sub-arrays of equal size. + hsplit : Split array into multiple sub-arrays horizontally (column wise). + vsplit : Split array into multiple sub-arrays vertically (row wise). + dsplit : Split array into multiple sub-arrays along the 3rd axis (depth). + stack : Stack a sequence of arrays along a new axis. + block : Assemble arrays from blocks. + hstack : Stack arrays in sequence horizontally (column wise). + vstack : Stack arrays in sequence vertically (row wise). + dstack : Stack arrays in sequence depth wise (along third dimension). + column_stack : Stack 1-D arrays as columns into a 2-D array. + + Notes + ----- + When one or more of the arrays to be concatenated is a MaskedArray, + this function will return a MaskedArray object instead of an ndarray, + but the input masks are *not* preserved. In cases where a MaskedArray + is expected as input, use the ma.concatenate function from the masked + array module instead. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4]]) + >>> b = np.array([[5, 6]]) + >>> np.concatenate((a, b), axis=0) + array([[1, 2], + [3, 4], + [5, 6]]) + >>> np.concatenate((a, b.T), axis=1) + array([[1, 2, 5], + [3, 4, 6]]) + >>> np.concatenate((a, b), axis=None) + array([1, 2, 3, 4, 5, 6]) + + This function will not preserve masking of MaskedArray inputs. + + >>> a = np.ma.arange(3) + >>> a[1] = np.ma.masked + >>> b = np.arange(2, 5) + >>> a + masked_array(data=[0, --, 2], + mask=[False, True, False], + fill_value=999999) + >>> b + array([2, 3, 4]) + >>> np.concatenate([a, b]) + masked_array(data=[0, 1, 2, 2, 3, 4], + mask=False, + fill_value=999999) + >>> np.ma.concatenate([a, b]) + masked_array(data=[0, --, 2, 2, 3, 4], + mask=[False, True, False, False, False, False], + fill_value=999999) + + """ + if out is not None: + # optimize for the typical case where only arrays is provided + arrays = list(arrays) + arrays.append(out) + return arrays + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.inner) +def inner(a, b): + """ + inner(a, b, /) + + Inner product of two arrays. + + Ordinary inner product of vectors for 1-D arrays (without complex + conjugation), in higher dimensions a sum product over the last axes. + + Parameters + ---------- + a, b : array_like + If `a` and `b` are nonscalar, their last dimensions must match. + + Returns + ------- + out : ndarray + If `a` and `b` are both + scalars or both 1-D arrays then a scalar is returned; otherwise + an array is returned. + ``out.shape = (*a.shape[:-1], *b.shape[:-1])`` + + Raises + ------ + ValueError + If both `a` and `b` are nonscalar and their last dimensions have + different sizes. + + See Also + -------- + tensordot : Sum products over arbitrary axes. + dot : Generalised matrix product, using second last dimension of `b`. + vecdot : Vector dot product of two arrays. + einsum : Einstein summation convention. + + Notes + ----- + For vectors (1-D arrays) it computes the ordinary inner-product:: + + np.inner(a, b) = sum(a[:]*b[:]) + + More generally, if ``ndim(a) = r > 0`` and ``ndim(b) = s > 0``:: + + np.inner(a, b) = np.tensordot(a, b, axes=(-1,-1)) + + or explicitly:: + + np.inner(a, b)[i0,...,ir-2,j0,...,js-2] + = sum(a[i0,...,ir-2,:]*b[j0,...,js-2,:]) + + In addition `a` or `b` may be scalars, in which case:: + + np.inner(a,b) = a*b + + Examples + -------- + Ordinary inner product for vectors: + + >>> import numpy as np + >>> a = np.array([1,2,3]) + >>> b = np.array([0,1,0]) + >>> np.inner(a, b) + 2 + + Some multidimensional examples: + + >>> a = np.arange(24).reshape((2,3,4)) + >>> b = np.arange(4) + >>> c = np.inner(a, b) + >>> c.shape + (2, 3) + >>> c + array([[ 14, 38, 62], + [ 86, 110, 134]]) + + >>> a = np.arange(2).reshape((1,1,2)) + >>> b = np.arange(6).reshape((3,2)) + >>> c = np.inner(a, b) + >>> c.shape + (1, 1, 3) + >>> c + array([[[1, 3, 5]]]) + + An example where `b` is a scalar: + + >>> np.inner(np.eye(2), 7) + array([[7., 0.], + [0., 7.]]) + + """ + return (a, b) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.where) +def where(condition, x=None, y=None): + """ + where(condition, [x, y], /) + + Return elements chosen from `x` or `y` depending on `condition`. + + .. note:: + When only `condition` is provided, this function is a shorthand for + ``np.asarray(condition).nonzero()``. Using `nonzero` directly should be + preferred, as it behaves correctly for subclasses. The rest of this + documentation covers only the case where all three arguments are + provided. + + Parameters + ---------- + condition : array_like, bool + Where True, yield `x`, otherwise yield `y`. + x, y : array_like + Values from which to choose. `x`, `y` and `condition` need to be + broadcastable to some shape. + + Returns + ------- + out : ndarray + An array with elements from `x` where `condition` is True, and elements + from `y` elsewhere. + + See Also + -------- + choose + nonzero : The function that is called when x and y are omitted + + Notes + ----- + If all the arrays are 1-D, `where` is equivalent to:: + + [xv if c else yv + for c, xv, yv in zip(condition, x, y)] + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(10) + >>> a + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> np.where(a < 5, a, 10*a) + array([ 0, 1, 2, 3, 4, 50, 60, 70, 80, 90]) + + This can be used on multidimensional arrays too: + + >>> np.where([[True, False], [True, True]], + ... [[1, 2], [3, 4]], + ... [[9, 8], [7, 6]]) + array([[1, 8], + [3, 4]]) + + The shapes of x, y, and the condition are broadcast together: + + >>> x, y = np.ogrid[:3, :4] + >>> np.where(x < y, x, 10 + y) # both x and 10+y are broadcast + array([[10, 0, 0, 0], + [10, 11, 1, 1], + [10, 11, 12, 2]]) + + >>> a = np.array([[0, 1, 2], + ... [0, 2, 4], + ... [0, 3, 6]]) + >>> np.where(a < 4, a, -1) # -1 is broadcast + array([[ 0, 1, 2], + [ 0, 2, -1], + [ 0, 3, -1]]) + """ + return (condition, x, y) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.lexsort) +def lexsort(keys, axis=None): + """ + lexsort(keys, axis=-1) + + Perform an indirect stable sort using a sequence of keys. + + Given multiple sorting keys, lexsort returns an array of integer indices + that describes the sort order by multiple keys. The last key in the + sequence is used for the primary sort order, ties are broken by the + second-to-last key, and so on. + + Parameters + ---------- + keys : (k, m, n, ...) array-like + The `k` keys to be sorted. The *last* key (e.g, the last + row if `keys` is a 2D array) is the primary sort key. + Each element of `keys` along the zeroth axis must be + an array-like object of the same shape. + axis : int, optional + Axis to be indirectly sorted. By default, sort over the last axis + of each sequence. Separate slices along `axis` sorted over + independently; see last example. + + Returns + ------- + indices : (m, n, ...) ndarray of ints + Array of indices that sort the keys along the specified axis. + + See Also + -------- + argsort : Indirect sort. + ndarray.sort : In-place sort. + sort : Return a sorted copy of an array. + + Examples + -------- + Sort names: first by surname, then by name. + + >>> import numpy as np + >>> surnames = ('Hertz', 'Galilei', 'Hertz') + >>> first_names = ('Heinrich', 'Galileo', 'Gustav') + >>> ind = np.lexsort((first_names, surnames)) + >>> ind + array([1, 2, 0]) + + >>> [surnames[i] + ", " + first_names[i] for i in ind] + ['Galilei, Galileo', 'Hertz, Gustav', 'Hertz, Heinrich'] + + Sort according to two numerical keys, first by elements + of ``a``, then breaking ties according to elements of ``b``: + + >>> a = [1, 5, 1, 4, 3, 4, 4] # First sequence + >>> b = [9, 4, 0, 4, 0, 2, 1] # Second sequence + >>> ind = np.lexsort((b, a)) # Sort by `a`, then by `b` + >>> ind + array([2, 0, 4, 6, 5, 3, 1]) + >>> [(a[i], b[i]) for i in ind] + [(1, 0), (1, 9), (3, 0), (4, 1), (4, 2), (4, 4), (5, 4)] + + Compare against `argsort`, which would sort each key independently. + + >>> np.argsort((b, a), kind='stable') + array([[2, 4, 6, 5, 1, 3, 0], + [0, 2, 4, 3, 5, 6, 1]]) + + To sort lexicographically with `argsort`, we would need to provide a + structured array. + + >>> x = np.array([(ai, bi) for ai, bi in zip(a, b)], + ... dtype = np.dtype([('x', int), ('y', int)])) + >>> np.argsort(x) # or np.argsort(x, order=('x', 'y')) + array([2, 0, 4, 6, 5, 3, 1]) + + The zeroth axis of `keys` always corresponds with the sequence of keys, + so 2D arrays are treated just like other sequences of keys. + + >>> arr = np.asarray([b, a]) + >>> ind2 = np.lexsort(arr) + >>> np.testing.assert_equal(ind2, ind) + + Accordingly, the `axis` parameter refers to an axis of *each* key, not of + the `keys` argument itself. For instance, the array ``arr`` is treated as + a sequence of two 1-D keys, so specifying ``axis=0`` is equivalent to + using the default axis, ``axis=-1``. + + >>> np.testing.assert_equal(np.lexsort(arr, axis=0), + ... np.lexsort(arr, axis=-1)) + + For higher-dimensional arrays, the axis parameter begins to matter. The + resulting array has the same shape as each key, and the values are what + we would expect if `lexsort` were performed on corresponding slices + of the keys independently. For instance, + + >>> x = [[1, 2, 3, 4], + ... [4, 3, 2, 1], + ... [2, 1, 4, 3]] + >>> y = [[2, 2, 1, 1], + ... [1, 2, 1, 2], + ... [1, 1, 2, 1]] + >>> np.lexsort((x, y), axis=1) + array([[2, 3, 0, 1], + [2, 0, 3, 1], + [1, 0, 3, 2]]) + + Each row of the result is what we would expect if we were to perform + `lexsort` on the corresponding row of the keys: + + >>> for i in range(3): + ... print(np.lexsort((x[i], y[i]))) + [2 3 0 1] + [2 0 3 1] + [1 0 3 2] + + """ + if isinstance(keys, tuple): + return keys + else: + return (keys,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.can_cast) +def can_cast(from_, to, casting=None): + """ + can_cast(from_, to, casting='safe') + + Returns True if cast between data types can occur according to the + casting rule. + + Parameters + ---------- + from_ : dtype, dtype specifier, NumPy scalar, or array + Data type, NumPy scalar, or array to cast from. + to : dtype or dtype specifier + Data type to cast to. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. + + * 'no' means the data types should not be cast at all. + * 'equiv' means only byte-order changes are allowed. + * 'safe' means only casts which can preserve values are allowed. + * 'same_kind' means only safe casts or casts within a kind, + like float64 to float32, are allowed. + * 'unsafe' means any data conversions may be done. + + Returns + ------- + out : bool + True if cast can occur according to the casting rule. + + Notes + ----- + .. versionchanged:: 2.0 + This function does not support Python scalars anymore and does not + apply any value-based logic for 0-D arrays and NumPy scalars. + + See also + -------- + dtype, result_type + + Examples + -------- + Basic examples + + >>> import numpy as np + >>> np.can_cast(np.int32, np.int64) + True + >>> np.can_cast(np.float64, complex) + True + >>> np.can_cast(complex, float) + False + + >>> np.can_cast('i8', 'f8') + True + >>> np.can_cast('i8', 'f4') + False + >>> np.can_cast('i4', 'S4') + False + + """ + return (from_,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.min_scalar_type) +def min_scalar_type(a): + """ + min_scalar_type(a, /) + + For scalar ``a``, returns the data type with the smallest size + and smallest scalar kind which can hold its value. For non-scalar + array ``a``, returns the vector's dtype unmodified. + + Floating point values are not demoted to integers, + and complex values are not demoted to floats. + + Parameters + ---------- + a : scalar or array_like + The value whose minimal data type is to be found. + + Returns + ------- + out : dtype + The minimal data type. + + See Also + -------- + result_type, promote_types, dtype, can_cast + + Examples + -------- + >>> import numpy as np + >>> np.min_scalar_type(10) + dtype('uint8') + + >>> np.min_scalar_type(-260) + dtype('int16') + + >>> np.min_scalar_type(3.1) + dtype('float16') + + >>> np.min_scalar_type(1e50) + dtype('float64') + + >>> np.min_scalar_type(np.arange(4,dtype='f8')) + dtype('float64') + + """ + return (a,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.result_type) +def result_type(*arrays_and_dtypes): + """ + result_type(*arrays_and_dtypes) + + Returns the type that results from applying the NumPy + type promotion rules to the arguments. + + Type promotion in NumPy works similarly to the rules in languages + like C++, with some slight differences. When both scalars and + arrays are used, the array's type takes precedence and the actual value + of the scalar is taken into account. + + For example, calculating 3*a, where a is an array of 32-bit floats, + intuitively should result in a 32-bit float output. If the 3 is a + 32-bit integer, the NumPy rules indicate it can't convert losslessly + into a 32-bit float, so a 64-bit float should be the result type. + By examining the value of the constant, '3', we see that it fits in + an 8-bit integer, which can be cast losslessly into the 32-bit float. + + Parameters + ---------- + arrays_and_dtypes : list of arrays and dtypes + The operands of some operation whose result type is needed. + + Returns + ------- + out : dtype + The result type. + + See also + -------- + dtype, promote_types, min_scalar_type, can_cast + + Notes + ----- + The specific algorithm used is as follows. + + Categories are determined by first checking which of boolean, + integer (int/uint), or floating point (float/complex) the maximum + kind of all the arrays and the scalars are. + + If there are only scalars or the maximum category of the scalars + is higher than the maximum category of the arrays, + the data types are combined with :func:`promote_types` + to produce the return value. + + Otherwise, `min_scalar_type` is called on each scalar, and + the resulting data types are all combined with :func:`promote_types` + to produce the return value. + + The set of int values is not a subset of the uint values for types + with the same number of bits, something not reflected in + :func:`min_scalar_type`, but handled as a special case in `result_type`. + + Examples + -------- + >>> import numpy as np + >>> np.result_type(3, np.arange(7, dtype='i1')) + dtype('int8') + + >>> np.result_type('i4', 'c8') + dtype('complex128') + + >>> np.result_type(3.0, -2) + dtype('float64') + + """ + return arrays_and_dtypes + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.dot) +def dot(a, b, out=None): + """ + dot(a, b, out=None) + + Dot product of two arrays. Specifically, + + - If both `a` and `b` are 1-D arrays, it is inner product of vectors + (without complex conjugation). + + - If both `a` and `b` are 2-D arrays, it is matrix multiplication, + but using :func:`matmul` or ``a @ b`` is preferred. + + - If either `a` or `b` is 0-D (scalar), it is equivalent to + :func:`multiply` and using ``numpy.multiply(a, b)`` or ``a * b`` is + preferred. + + - If `a` is an N-D array and `b` is a 1-D array, it is a sum product over + the last axis of `a` and `b`. + + - If `a` is an N-D array and `b` is an M-D array (where ``M>=2``), it is a + sum product over the last axis of `a` and the second-to-last axis of + `b`:: + + dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m]) + + It uses an optimized BLAS library when possible (see `numpy.linalg`). + + Parameters + ---------- + a : array_like + First argument. + b : array_like + Second argument. + out : ndarray, optional + Output argument. This must have the exact kind that would be returned + if it was not used. In particular, it must have the right type, must be + C-contiguous, and its dtype must be the dtype that would be returned + for `dot(a,b)`. This is a performance feature. Therefore, if these + conditions are not met, an exception is raised, instead of attempting + to be flexible. + + Returns + ------- + output : ndarray + Returns the dot product of `a` and `b`. If `a` and `b` are both + scalars or both 1-D arrays then a scalar is returned; otherwise + an array is returned. + If `out` is given, then it is returned. + + Raises + ------ + ValueError + If the last dimension of `a` is not the same size as + the second-to-last dimension of `b`. + + See Also + -------- + vdot : Complex-conjugating dot product. + vecdot : Vector dot product of two arrays. + tensordot : Sum products over arbitrary axes. + einsum : Einstein summation convention. + matmul : '@' operator as method with out parameter. + linalg.multi_dot : Chained dot product. + + Examples + -------- + >>> import numpy as np + >>> np.dot(3, 4) + 12 + + Neither argument is complex-conjugated: + + >>> np.dot([2j, 3j], [2j, 3j]) + (-13+0j) + + For 2-D arrays it is the matrix product: + + >>> a = [[1, 0], [0, 1]] + >>> b = [[4, 1], [2, 2]] + >>> np.dot(a, b) + array([[4, 1], + [2, 2]]) + + >>> a = np.arange(3*4*5*6).reshape((3,4,5,6)) + >>> b = np.arange(3*4*5*6)[::-1].reshape((5,4,6,3)) + >>> np.dot(a, b)[2,3,2,1,2,2] + 499128 + >>> sum(a[2,3,2,:] * b[1,2,:,2]) + 499128 + + """ + return (a, b, out) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.vdot) +def vdot(a, b): + r""" + vdot(a, b, /) + + Return the dot product of two vectors. + + The `vdot` function handles complex numbers differently than `dot`: + if the first argument is complex, it is replaced by its complex conjugate + in the dot product calculation. `vdot` also handles multidimensional + arrays differently than `dot`: it does not perform a matrix product, but + flattens the arguments to 1-D arrays before taking a vector dot product. + + Consequently, when the arguments are 2-D arrays of the same shape, this + function effectively returns their + `Frobenius inner product `_ + (also known as the *trace inner product* or the *standard inner product* + on a vector space of matrices). + + Parameters + ---------- + a : array_like + If `a` is complex the complex conjugate is taken before calculation + of the dot product. + b : array_like + Second argument to the dot product. + + Returns + ------- + output : ndarray + Dot product of `a` and `b`. Can be an int, float, or + complex depending on the types of `a` and `b`. + + See Also + -------- + dot : Return the dot product without using the complex conjugate of the + first argument. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1+2j,3+4j]) + >>> b = np.array([5+6j,7+8j]) + >>> np.vdot(a, b) + (70-8j) + >>> np.vdot(b, a) + (70+8j) + + Note that higher-dimensional arrays are flattened! + + >>> a = np.array([[1, 4], [5, 6]]) + >>> b = np.array([[4, 1], [2, 2]]) + >>> np.vdot(a, b) + 30 + >>> np.vdot(b, a) + 30 + >>> 1*4 + 4*1 + 5*2 + 6*2 + 30 + + """ # noqa: E501 + return (a, b) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.bincount) +def bincount(x, weights=None, minlength=None): + """ + bincount(x, /, weights=None, minlength=0) + + Count number of occurrences of each value in array of non-negative ints. + + The number of bins (of size 1) is one larger than the largest value in + `x`. If `minlength` is specified, there will be at least this number + of bins in the output array (though it will be longer if necessary, + depending on the contents of `x`). + Each bin gives the number of occurrences of its index value in `x`. + If `weights` is specified the input array is weighted by it, i.e. if a + value ``n`` is found at position ``i``, ``out[n] += weight[i]`` instead + of ``out[n] += 1``. + + Parameters + ---------- + x : array_like, 1 dimension, nonnegative ints + Input array. + weights : array_like, optional + Weights, array of the same shape as `x`. + minlength : int, optional + A minimum number of bins for the output array. + + Returns + ------- + out : ndarray of ints + The result of binning the input array. + The length of `out` is equal to ``np.amax(x)+1``. + + Raises + ------ + ValueError + If the input is not 1-dimensional, or contains elements with negative + values, or if `minlength` is negative. + TypeError + If the type of the input is float or complex. + + See Also + -------- + histogram, digitize, unique + + Examples + -------- + >>> import numpy as np + >>> np.bincount(np.arange(5)) + array([1, 1, 1, 1, 1]) + >>> np.bincount(np.array([0, 1, 1, 3, 2, 1, 7])) + array([1, 3, 1, 1, 0, 0, 0, 1]) + + >>> x = np.array([0, 1, 1, 3, 2, 1, 7, 23]) + >>> np.bincount(x).size == np.amax(x)+1 + True + + The input array needs to be of integer dtype, otherwise a + TypeError is raised: + + >>> np.bincount(np.arange(5, dtype=float)) + Traceback (most recent call last): + ... + TypeError: Cannot cast array data from dtype('float64') to dtype('int64') + according to the rule 'safe' + + A possible use of ``bincount`` is to perform sums over + variable-size chunks of an array, using the ``weights`` keyword. + + >>> w = np.array([0.3, 0.5, 0.2, 0.7, 1., -0.6]) # weights + >>> x = np.array([0, 1, 1, 2, 2, 2]) + >>> np.bincount(x, weights=w) + array([ 0.3, 0.7, 1.1]) + + """ + return (x, weights) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.ravel_multi_index) +def ravel_multi_index(multi_index, dims, mode=None, order=None): + """ + ravel_multi_index(multi_index, dims, mode='raise', order='C') + + Converts a tuple of index arrays into an array of flat + indices, applying boundary modes to the multi-index. + + Parameters + ---------- + multi_index : tuple of array_like + A tuple of integer arrays, one array for each dimension. + dims : tuple of ints + The shape of array into which the indices from ``multi_index`` apply. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices are handled. Can specify + either one mode or a tuple of modes, one mode per index. + + * 'raise' -- raise an error (default) + * 'wrap' -- wrap around + * 'clip' -- clip to the range + + In 'clip' mode, a negative index which would normally + wrap will clip to 0 instead. + order : {'C', 'F'}, optional + Determines whether the multi-index should be viewed as + indexing in row-major (C-style) or column-major + (Fortran-style) order. + + Returns + ------- + raveled_indices : ndarray + An array of indices into the flattened version of an array + of dimensions ``dims``. + + See Also + -------- + unravel_index + + Examples + -------- + >>> import numpy as np + >>> arr = np.array([[3,6,6],[4,5,1]]) + >>> np.ravel_multi_index(arr, (7,6)) + array([22, 41, 37]) + >>> np.ravel_multi_index(arr, (7,6), order='F') + array([31, 41, 13]) + >>> np.ravel_multi_index(arr, (4,6), mode='clip') + array([22, 23, 19]) + >>> np.ravel_multi_index(arr, (4,4), mode=('clip','wrap')) + array([12, 13, 13]) + + >>> np.ravel_multi_index((3,1,4,1), (6,7,8,9)) + 1621 + """ + return multi_index + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.unravel_index) +def unravel_index(indices, shape=None, order=None): + """ + unravel_index(indices, shape, order='C') + + Converts a flat index or array of flat indices into a tuple + of coordinate arrays. + + Parameters + ---------- + indices : array_like + An integer array whose elements are indices into the flattened + version of an array of dimensions ``shape``. Before version 1.6.0, + this function accepted just one index value. + shape : tuple of ints + The shape of the array to use for unraveling ``indices``. + order : {'C', 'F'}, optional + Determines whether the indices should be viewed as indexing in + row-major (C-style) or column-major (Fortran-style) order. + + Returns + ------- + unraveled_coords : tuple of ndarray + Each array in the tuple has the same shape as the ``indices`` + array. + + See Also + -------- + ravel_multi_index + + Examples + -------- + >>> import numpy as np + >>> np.unravel_index([22, 41, 37], (7,6)) + (array([3, 6, 6]), array([4, 5, 1])) + >>> np.unravel_index([31, 41, 13], (7,6), order='F') + (array([3, 6, 6]), array([4, 5, 1])) + + >>> np.unravel_index(1621, (6,7,8,9)) + (3, 1, 4, 1) + + """ + return (indices,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.copyto) +def copyto(dst, src, casting=None, where=None): + """ + copyto(dst, src, casting='same_kind', where=True) + + Copies values from one array to another, broadcasting as necessary. + + Raises a TypeError if the `casting` rule is violated, and if + `where` is provided, it selects which elements to copy. + + Parameters + ---------- + dst : ndarray + The array into which values are copied. + src : array_like + The array from which values are copied. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur when copying. + + * 'no' means the data types should not be cast at all. + * 'equiv' means only byte-order changes are allowed. + * 'safe' means only casts which can preserve values are allowed. + * 'same_kind' means only safe casts or casts within a kind, + like float64 to float32, are allowed. + * 'unsafe' means any data conversions may be done. + where : array_like of bool, optional + A boolean array which is broadcasted to match the dimensions + of `dst`, and selects elements to copy from `src` to `dst` + wherever it contains the value True. + + Examples + -------- + >>> import numpy as np + >>> A = np.array([4, 5, 6]) + >>> B = [1, 2, 3] + >>> np.copyto(A, B) + >>> A + array([1, 2, 3]) + + >>> A = np.array([[1, 2, 3], [4, 5, 6]]) + >>> B = [[4, 5, 6], [7, 8, 9]] + >>> np.copyto(A, B) + >>> A + array([[4, 5, 6], + [7, 8, 9]]) + + """ + return (dst, src, where) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.putmask) +def putmask(a, /, mask, values): + """ + putmask(a, mask, values) + + Changes elements of an array based on conditional and input values. + + Sets ``a.flat[n] = values[n]`` for each n where ``mask.flat[n]==True``. + + If `values` is not the same size as `a` and `mask` then it will repeat. + This gives behavior different from ``a[mask] = values``. + + Parameters + ---------- + a : ndarray + Target array. + mask : array_like + Boolean mask array. It has to be the same shape as `a`. + values : array_like + Values to put into `a` where `mask` is True. If `values` is smaller + than `a` it will be repeated. + + See Also + -------- + place, put, take, copyto + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(6).reshape(2, 3) + >>> np.putmask(x, x>2, x**2) + >>> x + array([[ 0, 1, 2], + [ 9, 16, 25]]) + + If `values` is smaller than `a` it is repeated: + + >>> x = np.arange(5) + >>> np.putmask(x, x>1, [-33, -44]) + >>> x + array([ 0, 1, -33, -44, -33]) + + """ + return (a, mask, values) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.packbits) +def packbits(a, axis=None, bitorder='big'): + """ + packbits(a, /, axis=None, bitorder='big') + + Packs the elements of a binary-valued array into bits in a uint8 array. + + The result is padded to full bytes by inserting zero bits at the end. + + Parameters + ---------- + a : array_like + An array of integers or booleans whose elements should be packed to + bits. + axis : int, optional + The dimension over which bit-packing is done. + ``None`` implies packing the flattened array. + bitorder : {'big', 'little'}, optional + The order of the input bits. 'big' will mimic bin(val), + ``[0, 0, 0, 0, 0, 0, 1, 1] => 3 = 0b00000011``, 'little' will + reverse the order so ``[1, 1, 0, 0, 0, 0, 0, 0] => 3``. + Defaults to 'big'. + + Returns + ------- + packed : ndarray + Array of type uint8 whose elements represent bits corresponding to the + logical (0 or nonzero) value of the input elements. The shape of + `packed` has the same number of dimensions as the input (unless `axis` + is None, in which case the output is 1-D). + + See Also + -------- + unpackbits: Unpacks elements of a uint8 array into a binary-valued output + array. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[[1,0,1], + ... [0,1,0]], + ... [[1,1,0], + ... [0,0,1]]]) + >>> b = np.packbits(a, axis=-1) + >>> b + array([[[160], + [ 64]], + [[192], + [ 32]]], dtype=uint8) + + Note that in binary 160 = 1010 0000, 64 = 0100 0000, 192 = 1100 0000, + and 32 = 0010 0000. + + """ + return (a,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.unpackbits) +def unpackbits(a, axis=None, count=None, bitorder='big'): + """ + unpackbits(a, /, axis=None, count=None, bitorder='big') + + Unpacks elements of a uint8 array into a binary-valued output array. + + Each element of `a` represents a bit-field that should be unpacked + into a binary-valued output array. The shape of the output array is + either 1-D (if `axis` is ``None``) or the same shape as the input + array with unpacking done along the axis specified. + + Parameters + ---------- + a : ndarray, uint8 type + Input array. + axis : int, optional + The dimension over which bit-unpacking is done. + ``None`` implies unpacking the flattened array. + count : int or None, optional + The number of elements to unpack along `axis`, provided as a way + of undoing the effect of packing a size that is not a multiple + of eight. A non-negative number means to only unpack `count` + bits. A negative number means to trim off that many bits from + the end. ``None`` means to unpack the entire array (the + default). Counts larger than the available number of bits will + add zero padding to the output. Negative counts must not + exceed the available number of bits. + bitorder : {'big', 'little'}, optional + The order of the returned bits. 'big' will mimic bin(val), + ``3 = 0b00000011 => [0, 0, 0, 0, 0, 0, 1, 1]``, 'little' will reverse + the order to ``[1, 1, 0, 0, 0, 0, 0, 0]``. + Defaults to 'big'. + + Returns + ------- + unpacked : ndarray, uint8 type + The elements are binary-valued (0 or 1). + + See Also + -------- + packbits : Packs the elements of a binary-valued array into bits in + a uint8 array. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[2], [7], [23]], dtype=np.uint8) + >>> a + array([[ 2], + [ 7], + [23]], dtype=uint8) + >>> b = np.unpackbits(a, axis=1) + >>> b + array([[0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 0, 1, 1, 1]], dtype=uint8) + >>> c = np.unpackbits(a, axis=1, count=-3) + >>> c + array([[0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 1, 0]], dtype=uint8) + + >>> p = np.packbits(b, axis=0) + >>> np.unpackbits(p, axis=0) + array([[0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 0, 1, 1, 1], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8) + >>> np.array_equal(b, np.unpackbits(p, axis=0, count=b.shape[0])) + True + + """ + return (a,) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.shares_memory) +def shares_memory(a, b, max_work=None): + """ + shares_memory(a, b, /, max_work=None) + + Determine if two arrays share memory. + + .. warning:: + + This function can be exponentially slow for some inputs, unless + `max_work` is set to zero or a positive integer. + If in doubt, use `numpy.may_share_memory` instead. + + Parameters + ---------- + a, b : ndarray + Input arrays + max_work : int, optional + Effort to spend on solving the overlap problem (maximum number + of candidate solutions to consider). The following special + values are recognized: + + max_work=-1 (default) + The problem is solved exactly. In this case, the function returns + True only if there is an element shared between the arrays. Finding + the exact solution may take extremely long in some cases. + max_work=0 + Only the memory bounds of a and b are checked. + This is equivalent to using ``may_share_memory()``. + + Raises + ------ + numpy.exceptions.TooHardError + Exceeded max_work. + + Returns + ------- + out : bool + + See Also + -------- + may_share_memory + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1, 2, 3, 4]) + >>> np.shares_memory(x, np.array([5, 6, 7])) + False + >>> np.shares_memory(x[::2], x) + True + >>> np.shares_memory(x[::2], x[1::2]) + False + + Checking whether two arrays share memory is NP-complete, and + runtime may increase exponentially in the number of + dimensions. Hence, `max_work` should generally be set to a finite + number, as it is possible to construct examples that take + extremely long to run: + + >>> from numpy.lib.stride_tricks import as_strided + >>> x = np.zeros([192163377], dtype=np.int8) + >>> x1 = as_strided( + ... x, strides=(36674, 61119, 85569), shape=(1049, 1049, 1049)) + >>> x2 = as_strided( + ... x[64023025:], strides=(12223, 12224, 1), shape=(1049, 1049, 1)) + >>> np.shares_memory(x1, x2, max_work=1000) + Traceback (most recent call last): + ... + numpy.exceptions.TooHardError: Exceeded max_work + + Running ``np.shares_memory(x1, x2)`` without `max_work` set takes + around 1 minute for this case. It is possible to find problems + that take still significantly longer. + + """ + return (a, b) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.may_share_memory) +def may_share_memory(a, b, max_work=None): + """ + may_share_memory(a, b, /, max_work=None) + + Determine if two arrays might share memory + + A return of True does not necessarily mean that the two arrays + share any element. It just means that they *might*. + + Only the memory bounds of a and b are checked by default. + + Parameters + ---------- + a, b : ndarray + Input arrays + max_work : int, optional + Effort to spend on solving the overlap problem. See + `shares_memory` for details. Default for ``may_share_memory`` + is to do a bounds check. + + Returns + ------- + out : bool + + See Also + -------- + shares_memory + + Examples + -------- + >>> import numpy as np + >>> np.may_share_memory(np.array([1,2]), np.array([5,8,9])) + False + >>> x = np.zeros([3, 4]) + >>> np.may_share_memory(x[:,0], x[:,1]) + True + + """ + return (a, b) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.is_busday) +def is_busday(dates, weekmask=None, holidays=None, busdaycal=None, out=None): + """ + is_busday( + dates, + weekmask='1111100', + holidays=None, + busdaycal=None, + out=None + ) + + Calculates which of the given dates are valid days, and which are not. + + Parameters + ---------- + dates : array_like of datetime64[D] + The array of dates to process. + weekmask : str or array_like of bool, optional + A seven-element array indicating which of Monday through Sunday are + valid days. May be specified as a length-seven list or array, like + [1,1,1,1,1,0,0]; a length-seven string, like '1111100'; or a string + like "Mon Tue Wed Thu Fri", made up of 3-character abbreviations for + weekdays, optionally separated by white space. Valid abbreviations + are: Mon Tue Wed Thu Fri Sat Sun + holidays : array_like of datetime64[D], optional + An array of dates to consider as invalid dates. They may be + specified in any order, and NaT (not-a-time) dates are ignored. + This list is saved in a normalized form that is suited for + fast calculations of valid days. + busdaycal : busdaycalendar, optional + A `busdaycalendar` object which specifies the valid days. If this + parameter is provided, neither weekmask nor holidays may be + provided. + out : array of bool, optional + If provided, this array is filled with the result. + + Returns + ------- + out : array of bool + An array with the same shape as ``dates``, containing True for + each valid day, and False for each invalid day. + + See Also + -------- + busdaycalendar : An object that specifies a custom set of valid days. + busday_offset : Applies an offset counted in valid days. + busday_count : Counts how many valid days are in a half-open date range. + + Examples + -------- + >>> import numpy as np + >>> # The weekdays are Friday, Saturday, and Monday + ... np.is_busday(['2011-07-01', '2011-07-02', '2011-07-18'], + ... holidays=['2011-07-01', '2011-07-04', '2011-07-17']) + array([False, False, True]) + """ + return (dates, weekmask, holidays, out) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.busday_offset) +def busday_offset(dates, offsets, roll=None, weekmask=None, holidays=None, + busdaycal=None, out=None): + """ + busday_offset( + dates, + offsets, + roll='raise', + weekmask='1111100', + holidays=None, + busdaycal=None, + out=None + ) + + First adjusts the date to fall on a valid day according to + the ``roll`` rule, then applies offsets to the given dates + counted in valid days. + + Parameters + ---------- + dates : array_like of datetime64[D] + The array of dates to process. + offsets : array_like of int + The array of offsets, which is broadcast with ``dates``. + roll : {'raise', 'nat', 'forward', 'following', 'backward', 'preceding', \ + 'modifiedfollowing', 'modifiedpreceding'}, optional + How to treat dates that do not fall on a valid day. The default + is 'raise'. + + * 'raise' means to raise an exception for an invalid day. + * 'nat' means to return a NaT (not-a-time) for an invalid day. + * 'forward' and 'following' mean to take the first valid day + later in time. + * 'backward' and 'preceding' mean to take the first valid day + earlier in time. + * 'modifiedfollowing' means to take the first valid day + later in time unless it is across a Month boundary, in which + case to take the first valid day earlier in time. + * 'modifiedpreceding' means to take the first valid day + earlier in time unless it is across a Month boundary, in which + case to take the first valid day later in time. + weekmask : str or array_like of bool, optional + A seven-element array indicating which of Monday through Sunday are + valid days. May be specified as a length-seven list or array, like + [1,1,1,1,1,0,0]; a length-seven string, like '1111100'; or a string + like "Mon Tue Wed Thu Fri", made up of 3-character abbreviations for + weekdays, optionally separated by white space. Valid abbreviations + are: Mon Tue Wed Thu Fri Sat Sun + holidays : array_like of datetime64[D], optional + An array of dates to consider as invalid dates. They may be + specified in any order, and NaT (not-a-time) dates are ignored. + This list is saved in a normalized form that is suited for + fast calculations of valid days. + busdaycal : busdaycalendar, optional + A `busdaycalendar` object which specifies the valid days. If this + parameter is provided, neither weekmask nor holidays may be + provided. + out : array of datetime64[D], optional + If provided, this array is filled with the result. + + Returns + ------- + out : array of datetime64[D] + An array with a shape from broadcasting ``dates`` and ``offsets`` + together, containing the dates with offsets applied. + + See Also + -------- + busdaycalendar : An object that specifies a custom set of valid days. + is_busday : Returns a boolean array indicating valid days. + busday_count : Counts how many valid days are in a half-open date range. + + Examples + -------- + >>> import numpy as np + >>> # First business day in October 2011 (not accounting for holidays) + ... np.busday_offset('2011-10', 0, roll='forward') + np.datetime64('2011-10-03') + >>> # Last business day in February 2012 (not accounting for holidays) + ... np.busday_offset('2012-03', -1, roll='forward') + np.datetime64('2012-02-29') + >>> # Third Wednesday in January 2011 + ... np.busday_offset('2011-01', 2, roll='forward', weekmask='Wed') + np.datetime64('2011-01-19') + >>> # 2012 Mother's Day in Canada and the U.S. + ... np.busday_offset('2012-05', 1, roll='forward', weekmask='Sun') + np.datetime64('2012-05-13') + + >>> # First business day on or after a date + ... np.busday_offset('2011-03-20', 0, roll='forward') + np.datetime64('2011-03-21') + >>> np.busday_offset('2011-03-22', 0, roll='forward') + np.datetime64('2011-03-22') + >>> # First business day after a date + ... np.busday_offset('2011-03-20', 1, roll='backward') + np.datetime64('2011-03-21') + >>> np.busday_offset('2011-03-22', 1, roll='backward') + np.datetime64('2011-03-23') + """ + return (dates, offsets, weekmask, holidays, out) + + +@array_function_from_c_func_and_dispatcher(_multiarray_umath.busday_count) +def busday_count(begindates, enddates, weekmask=None, holidays=None, + busdaycal=None, out=None): + """ + busday_count( + begindates, + enddates, + weekmask='1111100', + holidays=[], + busdaycal=None, + out=None + ) + + Counts the number of valid days between `begindates` and + `enddates`, not including the day of `enddates`. + + If ``enddates`` specifies a date value that is earlier than the + corresponding ``begindates`` date value, the count will be negative. + + Parameters + ---------- + begindates : array_like of datetime64[D] + The array of the first dates for counting. + enddates : array_like of datetime64[D] + The array of the end dates for counting, which are excluded + from the count themselves. + weekmask : str or array_like of bool, optional + A seven-element array indicating which of Monday through Sunday are + valid days. May be specified as a length-seven list or array, like + [1,1,1,1,1,0,0]; a length-seven string, like '1111100'; or a string + like "Mon Tue Wed Thu Fri", made up of 3-character abbreviations for + weekdays, optionally separated by white space. Valid abbreviations + are: Mon Tue Wed Thu Fri Sat Sun + holidays : array_like of datetime64[D], optional + An array of dates to consider as invalid dates. They may be + specified in any order, and NaT (not-a-time) dates are ignored. + This list is saved in a normalized form that is suited for + fast calculations of valid days. + busdaycal : busdaycalendar, optional + A `busdaycalendar` object which specifies the valid days. If this + parameter is provided, neither weekmask nor holidays may be + provided. + out : array of int, optional + If provided, this array is filled with the result. + + Returns + ------- + out : array of int + An array with a shape from broadcasting ``begindates`` and ``enddates`` + together, containing the number of valid days between + the begin and end dates. + + See Also + -------- + busdaycalendar : An object that specifies a custom set of valid days. + is_busday : Returns a boolean array indicating valid days. + busday_offset : Applies an offset counted in valid days. + + Examples + -------- + >>> import numpy as np + >>> # Number of weekdays in January 2011 + ... np.busday_count('2011-01', '2011-02') + 21 + >>> # Number of weekdays in 2011 + >>> np.busday_count('2011', '2012') + 260 + >>> # Number of Saturdays in 2011 + ... np.busday_count('2011', '2012', weekmask='Sat') + 53 + """ + return (begindates, enddates, weekmask, holidays, out) + + +@array_function_from_c_func_and_dispatcher( + _multiarray_umath.datetime_as_string) +def datetime_as_string(arr, unit=None, timezone=None, casting=None): + """ + datetime_as_string(arr, unit=None, timezone='naive', casting='same_kind') + + Convert an array of datetimes into an array of strings. + + Parameters + ---------- + arr : array_like of datetime64 + The array of UTC timestamps to format. + unit : str + One of None, 'auto', or + a :ref:`datetime unit `. + timezone : {'naive', 'UTC', 'local'} or tzinfo + Timezone information to use when displaying the datetime. If 'UTC', + end with a Z to indicate UTC time. If 'local', convert to the local + timezone first, and suffix with a +-#### timezone offset. If a tzinfo + object, then do as with 'local', but use the specified timezone. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'} + Casting to allow when changing between datetime units. + + Returns + ------- + str_arr : ndarray + An array of strings the same shape as `arr`. + + Examples + -------- + >>> import numpy as np + >>> import pytz + >>> d = np.arange('2002-10-27T04:30', 4*60, 60, dtype='M8[m]') + >>> d + array(['2002-10-27T04:30', '2002-10-27T05:30', '2002-10-27T06:30', + '2002-10-27T07:30'], dtype='datetime64[m]') + + Setting the timezone to UTC shows the same information, but with a Z suffix + + >>> np.datetime_as_string(d, timezone='UTC') + array(['2002-10-27T04:30Z', '2002-10-27T05:30Z', '2002-10-27T06:30Z', + '2002-10-27T07:30Z'], dtype='>> np.datetime_as_string(d, timezone=pytz.timezone('US/Eastern')) + array(['2002-10-27T00:30-0400', '2002-10-27T01:30-0400', + '2002-10-27T01:30-0500', '2002-10-27T02:30-0500'], dtype='>> np.datetime_as_string(d, unit='h') + array(['2002-10-27T04', '2002-10-27T05', '2002-10-27T06', '2002-10-27T07'], + dtype='>> np.datetime_as_string(d, unit='s') + array(['2002-10-27T04:30:00', '2002-10-27T05:30:00', '2002-10-27T06:30:00', + '2002-10-27T07:30:00'], dtype='>> np.datetime_as_string(d, unit='h', casting='safe') + Traceback (most recent call last): + ... + TypeError: Cannot create a datetime string as units 'h' from a NumPy + datetime with units 'm' according to the rule 'safe' + """ + return (arr,) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/multiarray.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/multiarray.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e21f12ba9c0ca49bf8f6d41c9d4d1dd7bea7d6a5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/multiarray.pyi @@ -0,0 +1,1355 @@ +# TODO: Sort out any and all missing functions in this namespace +import datetime as dt +from _typeshed import StrOrBytesPath, SupportsLenAndGetItem +from collections.abc import Sequence, Callable, Iterable +from typing import ( + Literal as L, + Any, + TypeAlias, + overload, + TypeVar, + TypedDict, + SupportsIndex, + final, + Final, + Protocol, + ClassVar, + type_check_only, +) +from typing_extensions import CapsuleType, Unpack + +import numpy as np +from numpy import ( # type: ignore[attr-defined] + # Re-exports + busdaycalendar, + broadcast, + correlate, + count_nonzero, + dtype, + einsum as c_einsum, + flatiter, + from_dlpack, + interp, + matmul, + ndarray, + nditer, + vecdot, + + # The rest + ufunc, + str_, + uint8, + intp, + int_, + float64, + timedelta64, + datetime64, + generic, + unsignedinteger, + signedinteger, + floating, + complexfloating, + _AnyShapeType, + _OrderKACF, + _OrderCF, + _CastingKind, + _ModeKind, + _SupportsBuffer, + _SupportsFileMethods, + _CopyMode, + _NDIterFlagsKind, + _NDIterFlagsOp, +) +from numpy.lib._array_utils_impl import normalize_axis_index + +from numpy._typing import ( + # Shapes + _ShapeLike, + + # DTypes + DTypeLike, + _DTypeLike, + _SupportsDType, + + # Arrays + NDArray, + ArrayLike, + _ArrayLike, + _SupportsArrayFunc, + _NestedSequence, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeTD64_co, + _ArrayLikeDT64_co, + _ArrayLikeObject_co, + _ArrayLikeStr_co, + _ArrayLikeBytes_co, + _ScalarLike_co, + _IntLike_co, + _FloatLike_co, + _TD64Like_co, +) +from numpy._typing._ufunc import ( + _2PTuple, + _PyFunc_Nin1_Nout1, + _PyFunc_Nin2_Nout1, + _PyFunc_Nin3P_Nout1, + _PyFunc_Nin1P_Nout2P, +) + +__all__ = [ + "_ARRAY_API", + "ALLOW_THREADS", + "BUFSIZE", + "CLIP", + "DATETIMEUNITS", + "ITEM_HASOBJECT", + "ITEM_IS_POINTER", + "LIST_PICKLE", + "MAXDIMS", + "MAY_SHARE_BOUNDS", + "MAY_SHARE_EXACT", + "NEEDS_INIT", + "NEEDS_PYAPI", + "RAISE", + "USE_GETITEM", + "USE_SETITEM", + "WRAP", + "_flagdict", + "from_dlpack", + "_place", + "_reconstruct", + "_vec_string", + "_monotonicity", + "add_docstring", + "arange", + "array", + "asarray", + "asanyarray", + "ascontiguousarray", + "asfortranarray", + "bincount", + "broadcast", + "busday_count", + "busday_offset", + "busdaycalendar", + "can_cast", + "compare_chararrays", + "concatenate", + "copyto", + "correlate", + "correlate2", + "count_nonzero", + "c_einsum", + "datetime_as_string", + "datetime_data", + "dot", + "dragon4_positional", + "dragon4_scientific", + "dtype", + "empty", + "empty_like", + "error", + "flagsobj", + "flatiter", + "format_longfloat", + "frombuffer", + "fromfile", + "fromiter", + "fromstring", + "get_handler_name", + "get_handler_version", + "inner", + "interp", + "interp_complex", + "is_busday", + "lexsort", + "matmul", + "vecdot", + "may_share_memory", + "min_scalar_type", + "ndarray", + "nditer", + "nested_iters", + "normalize_axis_index", + "packbits", + "promote_types", + "putmask", + "ravel_multi_index", + "result_type", + "scalar", + "set_datetimeparse_function", + "set_typeDict", + "shares_memory", + "typeinfo", + "unpackbits", + "unravel_index", + "vdot", + "where", + "zeros", +] + +_SCT = TypeVar("_SCT", bound=generic) +_DType = TypeVar("_DType", bound=np.dtype[Any]) +_ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any]) +_ArrayType_co = TypeVar( + "_ArrayType_co", + bound=ndarray[Any, Any], + covariant=True, +) +_ReturnType = TypeVar("_ReturnType") +_IDType = TypeVar("_IDType") +_Nin = TypeVar("_Nin", bound=int) +_Nout = TypeVar("_Nout", bound=int) + +_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...]) +_Array: TypeAlias = ndarray[_ShapeT, dtype[_SCT]] +_Array1D: TypeAlias = ndarray[tuple[int], dtype[_SCT]] + +# Valid time units +_UnitKind: TypeAlias = L[ + "Y", + "M", + "D", + "h", + "m", + "s", + "ms", + "us", "μs", + "ns", + "ps", + "fs", + "as", +] +_RollKind: TypeAlias = L[ # `raise` is deliberately excluded + "nat", + "forward", + "following", + "backward", + "preceding", + "modifiedfollowing", + "modifiedpreceding", +] + +@type_check_only +class _SupportsArray(Protocol[_ArrayType_co]): + def __array__(self, /) -> _ArrayType_co: ... + +@type_check_only +class _KwargsEmpty(TypedDict, total=False): + device: None | L["cpu"] + like: None | _SupportsArrayFunc + +@type_check_only +class _ConstructorEmpty(Protocol): + # 1-D shape + @overload + def __call__( + self, + /, + shape: SupportsIndex, + dtype: None = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> _Array1D[float64]: ... + @overload + def __call__( + self, + /, + shape: SupportsIndex, + dtype: _DType | _SupportsDType[_DType], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> ndarray[tuple[int], _DType]: ... + @overload + def __call__( + self, + /, + shape: SupportsIndex, + dtype: type[_SCT], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> _Array1D[_SCT]: ... + @overload + def __call__( + self, + /, + shape: SupportsIndex, + dtype: DTypeLike, + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> _Array1D[Any]: ... + + # known shape + @overload + def __call__( + self, + /, + shape: _AnyShapeType, + dtype: None = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> _Array[_AnyShapeType, float64]: ... + @overload + def __call__( + self, + /, + shape: _AnyShapeType, + dtype: _DType | _SupportsDType[_DType], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> ndarray[_AnyShapeType, _DType]: ... + @overload + def __call__( + self, + /, + shape: _AnyShapeType, + dtype: type[_SCT], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> _Array[_AnyShapeType, _SCT]: ... + @overload + def __call__( + self, + /, + shape: _AnyShapeType, + dtype: DTypeLike, + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> _Array[_AnyShapeType, Any]: ... + + # unknown shape + @overload + def __call__( + self, /, + shape: _ShapeLike, + dtype: None = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> NDArray[float64]: ... + @overload + def __call__( + self, /, + shape: _ShapeLike, + dtype: _DType | _SupportsDType[_DType], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> ndarray[Any, _DType]: ... + @overload + def __call__( + self, /, + shape: _ShapeLike, + dtype: type[_SCT], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> NDArray[_SCT]: ... + @overload + def __call__( + self, /, + shape: _ShapeLike, + dtype: DTypeLike, + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], + ) -> NDArray[Any]: ... + +# using `Final` or `TypeAlias` will break stubtest +error = Exception + +# from ._multiarray_umath +ITEM_HASOBJECT: Final[L[1]] +LIST_PICKLE: Final[L[2]] +ITEM_IS_POINTER: Final[L[4]] +NEEDS_INIT: Final[L[8]] +NEEDS_PYAPI: Final[L[16]] +USE_GETITEM: Final[L[32]] +USE_SETITEM: Final[L[64]] +DATETIMEUNITS: Final[CapsuleType] +_ARRAY_API: Final[CapsuleType] +_flagdict: Final[dict[str, int]] +_monotonicity: Final[Callable[..., object]] +_place: Final[Callable[..., object]] +_reconstruct: Final[Callable[..., object]] +_vec_string: Final[Callable[..., object]] +correlate2: Final[Callable[..., object]] +dragon4_positional: Final[Callable[..., object]] +dragon4_scientific: Final[Callable[..., object]] +interp_complex: Final[Callable[..., object]] +set_datetimeparse_function: Final[Callable[..., object]] +def get_handler_name(a: NDArray[Any] = ..., /) -> str | None: ... +def get_handler_version(a: NDArray[Any] = ..., /) -> int | None: ... +def format_longfloat(x: np.longdouble, precision: int) -> str: ... +def scalar(dtype: _DType, object: bytes | object = ...) -> ndarray[tuple[()], _DType]: ... +def set_typeDict(dict_: dict[str, np.dtype[Any]], /) -> None: ... +typeinfo: Final[dict[str, np.dtype[np.generic]]] + +ALLOW_THREADS: Final[int] # 0 or 1 (system-specific) +BUFSIZE: L[8192] +CLIP: L[0] +WRAP: L[1] +RAISE: L[2] +MAXDIMS: L[32] +MAY_SHARE_BOUNDS: L[0] +MAY_SHARE_EXACT: L[-1] +tracemalloc_domain: L[389047] + +zeros: Final[_ConstructorEmpty] +empty: Final[_ConstructorEmpty] + +@overload +def empty_like( + prototype: _ArrayType, + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> _ArrayType: ... +@overload +def empty_like( + prototype: _ArrayLike[_SCT], + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def empty_like( + prototype: object, + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... +@overload +def empty_like( + prototype: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def empty_like( + prototype: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... + +@overload +def array( + object: _ArrayType, + dtype: None = ..., + *, + copy: None | bool | _CopyMode = ..., + order: _OrderKACF = ..., + subok: L[True], + ndmin: int = ..., + like: None | _SupportsArrayFunc = ..., +) -> _ArrayType: ... +@overload +def array( + object: _SupportsArray[_ArrayType], + dtype: None = ..., + *, + copy: None | bool | _CopyMode = ..., + order: _OrderKACF = ..., + subok: L[True], + ndmin: L[0] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _ArrayType: ... +@overload +def array( + object: _ArrayLike[_SCT], + dtype: None = ..., + *, + copy: None | bool | _CopyMode = ..., + order: _OrderKACF = ..., + subok: bool = ..., + ndmin: int = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def array( + object: object, + dtype: None = ..., + *, + copy: None | bool | _CopyMode = ..., + order: _OrderKACF = ..., + subok: bool = ..., + ndmin: int = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... +@overload +def array( + object: Any, + dtype: _DTypeLike[_SCT], + *, + copy: None | bool | _CopyMode = ..., + order: _OrderKACF = ..., + subok: bool = ..., + ndmin: int = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def array( + object: Any, + dtype: DTypeLike, + *, + copy: None | bool | _CopyMode = ..., + order: _OrderKACF = ..., + subok: bool = ..., + ndmin: int = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def unravel_index( # type: ignore[misc] + indices: _IntLike_co, + shape: _ShapeLike, + order: _OrderCF = ..., +) -> tuple[intp, ...]: ... +@overload +def unravel_index( + indices: _ArrayLikeInt_co, + shape: _ShapeLike, + order: _OrderCF = ..., +) -> tuple[NDArray[intp], ...]: ... + +@overload +def ravel_multi_index( # type: ignore[misc] + multi_index: Sequence[_IntLike_co], + dims: Sequence[SupportsIndex], + mode: _ModeKind | tuple[_ModeKind, ...] = ..., + order: _OrderCF = ..., +) -> intp: ... +@overload +def ravel_multi_index( + multi_index: Sequence[_ArrayLikeInt_co], + dims: Sequence[SupportsIndex], + mode: _ModeKind | tuple[_ModeKind, ...] = ..., + order: _OrderCF = ..., +) -> NDArray[intp]: ... + +# NOTE: Allow any sequence of array-like objects +@overload +def concatenate( # type: ignore[misc] + arrays: _ArrayLike[_SCT], + /, + axis: None | SupportsIndex = ..., + out: None = ..., + *, + dtype: None = ..., + casting: None | _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def concatenate( # type: ignore[misc] + arrays: SupportsLenAndGetItem[ArrayLike], + /, + axis: None | SupportsIndex = ..., + out: None = ..., + *, + dtype: None = ..., + casting: None | _CastingKind = ... +) -> NDArray[Any]: ... +@overload +def concatenate( # type: ignore[misc] + arrays: SupportsLenAndGetItem[ArrayLike], + /, + axis: None | SupportsIndex = ..., + out: None = ..., + *, + dtype: _DTypeLike[_SCT], + casting: None | _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def concatenate( # type: ignore[misc] + arrays: SupportsLenAndGetItem[ArrayLike], + /, + axis: None | SupportsIndex = ..., + out: None = ..., + *, + dtype: DTypeLike, + casting: None | _CastingKind = ... +) -> NDArray[Any]: ... +@overload +def concatenate( + arrays: SupportsLenAndGetItem[ArrayLike], + /, + axis: None | SupportsIndex = ..., + out: _ArrayType = ..., + *, + dtype: DTypeLike = ..., + casting: None | _CastingKind = ... +) -> _ArrayType: ... + +def inner( + a: ArrayLike, + b: ArrayLike, + /, +) -> Any: ... + +@overload +def where( + condition: ArrayLike, + /, +) -> tuple[NDArray[intp], ...]: ... +@overload +def where( + condition: ArrayLike, + x: ArrayLike, + y: ArrayLike, + /, +) -> NDArray[Any]: ... + +def lexsort( + keys: ArrayLike, + axis: None | SupportsIndex = ..., +) -> Any: ... + +def can_cast( + from_: ArrayLike | DTypeLike, + to: DTypeLike, + casting: None | _CastingKind = ..., +) -> bool: ... + +def min_scalar_type( + a: ArrayLike, /, +) -> dtype[Any]: ... + +def result_type( + *arrays_and_dtypes: ArrayLike | DTypeLike, +) -> dtype[Any]: ... + +@overload +def dot(a: ArrayLike, b: ArrayLike, out: None = ...) -> Any: ... +@overload +def dot(a: ArrayLike, b: ArrayLike, out: _ArrayType) -> _ArrayType: ... + +@overload +def vdot(a: _ArrayLikeBool_co, b: _ArrayLikeBool_co, /) -> np.bool: ... # type: ignore[misc] +@overload +def vdot(a: _ArrayLikeUInt_co, b: _ArrayLikeUInt_co, /) -> unsignedinteger[Any]: ... # type: ignore[misc] +@overload +def vdot(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co, /) -> signedinteger[Any]: ... # type: ignore[misc] +@overload +def vdot(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, /) -> floating[Any]: ... # type: ignore[misc] +@overload +def vdot(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, /) -> complexfloating[Any, Any]: ... # type: ignore[misc] +@overload +def vdot(a: _ArrayLikeTD64_co, b: _ArrayLikeTD64_co, /) -> timedelta64: ... +@overload +def vdot(a: _ArrayLikeObject_co, b: Any, /) -> Any: ... +@overload +def vdot(a: Any, b: _ArrayLikeObject_co, /) -> Any: ... + +def bincount( + x: ArrayLike, + /, + weights: None | ArrayLike = ..., + minlength: SupportsIndex = ..., +) -> NDArray[intp]: ... + +def copyto( + dst: NDArray[Any], + src: ArrayLike, + casting: None | _CastingKind = ..., + where: None | _ArrayLikeBool_co = ..., +) -> None: ... + +def putmask( + a: NDArray[Any], + /, + mask: _ArrayLikeBool_co, + values: ArrayLike, +) -> None: ... + +def packbits( + a: _ArrayLikeInt_co, + /, + axis: None | SupportsIndex = ..., + bitorder: L["big", "little"] = ..., +) -> NDArray[uint8]: ... + +def unpackbits( + a: _ArrayLike[uint8], + /, + axis: None | SupportsIndex = ..., + count: None | SupportsIndex = ..., + bitorder: L["big", "little"] = ..., +) -> NDArray[uint8]: ... + +def shares_memory( + a: object, + b: object, + /, + max_work: None | int = ..., +) -> bool: ... + +def may_share_memory( + a: object, + b: object, + /, + max_work: None | int = ..., +) -> bool: ... + +@overload +def asarray( + a: _ArrayLike[_SCT], + dtype: None = ..., + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def asarray( + a: object, + dtype: None = ..., + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... +@overload +def asarray( + a: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def asarray( + a: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def asanyarray( + a: _ArrayType, # Preserve subclass-information + dtype: None = ..., + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> _ArrayType: ... +@overload +def asanyarray( + a: _ArrayLike[_SCT], + dtype: None = ..., + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def asanyarray( + a: object, + dtype: None = ..., + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... +@overload +def asanyarray( + a: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def asanyarray( + a: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., + *, + device: None | L["cpu"] = ..., + copy: None | bool = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def ascontiguousarray( + a: _ArrayLike[_SCT], + dtype: None = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def ascontiguousarray( + a: object, + dtype: None = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... +@overload +def ascontiguousarray( + a: Any, + dtype: _DTypeLike[_SCT], + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def ascontiguousarray( + a: Any, + dtype: DTypeLike, + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def asfortranarray( + a: _ArrayLike[_SCT], + dtype: None = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def asfortranarray( + a: object, + dtype: None = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... +@overload +def asfortranarray( + a: Any, + dtype: _DTypeLike[_SCT], + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def asfortranarray( + a: Any, + dtype: DTypeLike, + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +def promote_types(__type1: DTypeLike, __type2: DTypeLike) -> dtype[Any]: ... + +# `sep` is a de facto mandatory argument, as its default value is deprecated +@overload +def fromstring( + string: str | bytes, + dtype: None = ..., + count: SupportsIndex = ..., + *, + sep: str, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[float64]: ... +@overload +def fromstring( + string: str | bytes, + dtype: _DTypeLike[_SCT], + count: SupportsIndex = ..., + *, + sep: str, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def fromstring( + string: str | bytes, + dtype: DTypeLike, + count: SupportsIndex = ..., + *, + sep: str, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def frompyfunc( # type: ignore[overload-overlap] + func: Callable[[Any], _ReturnType], /, + nin: L[1], + nout: L[1], + *, + identity: None = ..., +) -> _PyFunc_Nin1_Nout1[_ReturnType, None]: ... +@overload +def frompyfunc( # type: ignore[overload-overlap] + func: Callable[[Any], _ReturnType], /, + nin: L[1], + nout: L[1], + *, + identity: _IDType, +) -> _PyFunc_Nin1_Nout1[_ReturnType, _IDType]: ... +@overload +def frompyfunc( # type: ignore[overload-overlap] + func: Callable[[Any, Any], _ReturnType], /, + nin: L[2], + nout: L[1], + *, + identity: None = ..., +) -> _PyFunc_Nin2_Nout1[_ReturnType, None]: ... +@overload +def frompyfunc( # type: ignore[overload-overlap] + func: Callable[[Any, Any], _ReturnType], /, + nin: L[2], + nout: L[1], + *, + identity: _IDType, +) -> _PyFunc_Nin2_Nout1[_ReturnType, _IDType]: ... +@overload +def frompyfunc( # type: ignore[overload-overlap] + func: Callable[..., _ReturnType], /, + nin: _Nin, + nout: L[1], + *, + identity: None = ..., +) -> _PyFunc_Nin3P_Nout1[_ReturnType, None, _Nin]: ... +@overload +def frompyfunc( # type: ignore[overload-overlap] + func: Callable[..., _ReturnType], /, + nin: _Nin, + nout: L[1], + *, + identity: _IDType, +) -> _PyFunc_Nin3P_Nout1[_ReturnType, _IDType, _Nin]: ... +@overload +def frompyfunc( + func: Callable[..., _2PTuple[_ReturnType]], /, + nin: _Nin, + nout: _Nout, + *, + identity: None = ..., +) -> _PyFunc_Nin1P_Nout2P[_ReturnType, None, _Nin, _Nout]: ... +@overload +def frompyfunc( + func: Callable[..., _2PTuple[_ReturnType]], /, + nin: _Nin, + nout: _Nout, + *, + identity: _IDType, +) -> _PyFunc_Nin1P_Nout2P[_ReturnType, _IDType, _Nin, _Nout]: ... +@overload +def frompyfunc( + func: Callable[..., Any], /, + nin: SupportsIndex, + nout: SupportsIndex, + *, + identity: None | object = ..., +) -> ufunc: ... + +@overload +def fromfile( + file: StrOrBytesPath | _SupportsFileMethods, + dtype: None = ..., + count: SupportsIndex = ..., + sep: str = ..., + offset: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[float64]: ... +@overload +def fromfile( + file: StrOrBytesPath | _SupportsFileMethods, + dtype: _DTypeLike[_SCT], + count: SupportsIndex = ..., + sep: str = ..., + offset: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def fromfile( + file: StrOrBytesPath | _SupportsFileMethods, + dtype: DTypeLike, + count: SupportsIndex = ..., + sep: str = ..., + offset: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def fromiter( + iter: Iterable[Any], + dtype: _DTypeLike[_SCT], + count: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def fromiter( + iter: Iterable[Any], + dtype: DTypeLike, + count: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def frombuffer( + buffer: _SupportsBuffer, + dtype: None = ..., + count: SupportsIndex = ..., + offset: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[float64]: ... +@overload +def frombuffer( + buffer: _SupportsBuffer, + dtype: _DTypeLike[_SCT], + count: SupportsIndex = ..., + offset: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def frombuffer( + buffer: _SupportsBuffer, + dtype: DTypeLike, + count: SupportsIndex = ..., + offset: SupportsIndex = ..., + *, + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def arange( # type: ignore[misc] + stop: _IntLike_co, + /, *, + dtype: None = ..., + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[signedinteger]: ... +@overload +def arange( # type: ignore[misc] + start: _IntLike_co, + stop: _IntLike_co, + step: _IntLike_co = ..., + dtype: None = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[signedinteger]: ... +@overload +def arange( # type: ignore[misc] + stop: _FloatLike_co, + /, *, + dtype: None = ..., + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[floating]: ... +@overload +def arange( # type: ignore[misc] + start: _FloatLike_co, + stop: _FloatLike_co, + step: _FloatLike_co = ..., + dtype: None = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[floating]: ... +@overload +def arange( + stop: _TD64Like_co, + /, *, + dtype: None = ..., + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[timedelta64]: ... +@overload +def arange( + start: _TD64Like_co, + stop: _TD64Like_co, + step: _TD64Like_co = ..., + dtype: None = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[timedelta64]: ... +@overload +def arange( # both start and stop must always be specified for datetime64 + start: datetime64, + stop: datetime64, + step: datetime64 = ..., + dtype: None = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[datetime64]: ... +@overload +def arange( + stop: Any, + /, *, + dtype: _DTypeLike[_SCT], + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[_SCT]: ... +@overload +def arange( + start: Any, + stop: Any, + step: Any = ..., + dtype: _DTypeLike[_SCT] = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[_SCT]: ... +@overload +def arange( + stop: Any, /, + *, + dtype: DTypeLike, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[Any]: ... +@overload +def arange( + start: Any, + stop: Any, + step: Any = ..., + dtype: DTypeLike = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> _Array1D[Any]: ... + +def datetime_data( + dtype: str | _DTypeLike[datetime64] | _DTypeLike[timedelta64], /, +) -> tuple[str, int]: ... + +# The datetime functions perform unsafe casts to `datetime64[D]`, +# so a lot of different argument types are allowed here + +@overload +def busday_count( # type: ignore[misc] + begindates: _ScalarLike_co | dt.date, + enddates: _ScalarLike_co | dt.date, + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> int_: ... +@overload +def busday_count( # type: ignore[misc] + begindates: ArrayLike | dt.date | _NestedSequence[dt.date], + enddates: ArrayLike | dt.date | _NestedSequence[dt.date], + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> NDArray[int_]: ... +@overload +def busday_count( + begindates: ArrayLike | dt.date | _NestedSequence[dt.date], + enddates: ArrayLike | dt.date | _NestedSequence[dt.date], + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: _ArrayType = ..., +) -> _ArrayType: ... + +# `roll="raise"` is (more or less?) equivalent to `casting="safe"` +@overload +def busday_offset( # type: ignore[misc] + dates: datetime64 | dt.date, + offsets: _TD64Like_co | dt.timedelta, + roll: L["raise"] = ..., + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> datetime64: ... +@overload +def busday_offset( # type: ignore[misc] + dates: _ArrayLike[datetime64] | dt.date | _NestedSequence[dt.date], + offsets: _ArrayLikeTD64_co | dt.timedelta | _NestedSequence[dt.timedelta], + roll: L["raise"] = ..., + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> NDArray[datetime64]: ... +@overload +def busday_offset( # type: ignore[misc] + dates: _ArrayLike[datetime64] | dt.date | _NestedSequence[dt.date], + offsets: _ArrayLikeTD64_co | dt.timedelta | _NestedSequence[dt.timedelta], + roll: L["raise"] = ..., + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: _ArrayType = ..., +) -> _ArrayType: ... +@overload +def busday_offset( # type: ignore[misc] + dates: _ScalarLike_co | dt.date, + offsets: _ScalarLike_co | dt.timedelta, + roll: _RollKind, + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> datetime64: ... +@overload +def busday_offset( # type: ignore[misc] + dates: ArrayLike | dt.date | _NestedSequence[dt.date], + offsets: ArrayLike | dt.timedelta | _NestedSequence[dt.timedelta], + roll: _RollKind, + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> NDArray[datetime64]: ... +@overload +def busday_offset( + dates: ArrayLike | dt.date | _NestedSequence[dt.date], + offsets: ArrayLike | dt.timedelta | _NestedSequence[dt.timedelta], + roll: _RollKind, + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: _ArrayType = ..., +) -> _ArrayType: ... + +@overload +def is_busday( # type: ignore[misc] + dates: _ScalarLike_co | dt.date, + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> np.bool: ... +@overload +def is_busday( # type: ignore[misc] + dates: ArrayLike | _NestedSequence[dt.date], + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: None = ..., +) -> NDArray[np.bool]: ... +@overload +def is_busday( + dates: ArrayLike | _NestedSequence[dt.date], + weekmask: ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + busdaycal: None | busdaycalendar = ..., + out: _ArrayType = ..., +) -> _ArrayType: ... + +@overload +def datetime_as_string( # type: ignore[misc] + arr: datetime64 | dt.date, + unit: None | L["auto"] | _UnitKind = ..., + timezone: L["naive", "UTC", "local"] | dt.tzinfo = ..., + casting: _CastingKind = ..., +) -> str_: ... +@overload +def datetime_as_string( + arr: _ArrayLikeDT64_co | _NestedSequence[dt.date], + unit: None | L["auto"] | _UnitKind = ..., + timezone: L["naive", "UTC", "local"] | dt.tzinfo = ..., + casting: _CastingKind = ..., +) -> NDArray[str_]: ... + +@overload +def compare_chararrays( + a1: _ArrayLikeStr_co, + a2: _ArrayLikeStr_co, + cmp: L["<", "<=", "==", ">=", ">", "!="], + rstrip: bool, +) -> NDArray[np.bool]: ... +@overload +def compare_chararrays( + a1: _ArrayLikeBytes_co, + a2: _ArrayLikeBytes_co, + cmp: L["<", "<=", "==", ">=", ">", "!="], + rstrip: bool, +) -> NDArray[np.bool]: ... + +def add_docstring(obj: Callable[..., Any], docstring: str, /) -> None: ... + +_GetItemKeys: TypeAlias = L[ + "C", "CONTIGUOUS", "C_CONTIGUOUS", + "F", "FORTRAN", "F_CONTIGUOUS", + "W", "WRITEABLE", + "B", "BEHAVED", + "O", "OWNDATA", + "A", "ALIGNED", + "X", "WRITEBACKIFCOPY", + "CA", "CARRAY", + "FA", "FARRAY", + "FNC", + "FORC", +] +_SetItemKeys: TypeAlias = L[ + "A", "ALIGNED", + "W", "WRITEABLE", + "X", "WRITEBACKIFCOPY", +] + +@final +class flagsobj: + __hash__: ClassVar[None] # type: ignore[assignment] + aligned: bool + # NOTE: deprecated + # updateifcopy: bool + writeable: bool + writebackifcopy: bool + @property + def behaved(self) -> bool: ... + @property + def c_contiguous(self) -> bool: ... + @property + def carray(self) -> bool: ... + @property + def contiguous(self) -> bool: ... + @property + def f_contiguous(self) -> bool: ... + @property + def farray(self) -> bool: ... + @property + def fnc(self) -> bool: ... + @property + def forc(self) -> bool: ... + @property + def fortran(self) -> bool: ... + @property + def num(self) -> int: ... + @property + def owndata(self) -> bool: ... + def __getitem__(self, key: _GetItemKeys) -> bool: ... + def __setitem__(self, key: _SetItemKeys, value: bool) -> None: ... + +def nested_iters( + op: ArrayLike | Sequence[ArrayLike], + axes: Sequence[Sequence[SupportsIndex]], + flags: None | Sequence[_NDIterFlagsKind] = ..., + op_flags: None | Sequence[Sequence[_NDIterFlagsOp]] = ..., + op_dtypes: DTypeLike | Sequence[DTypeLike] = ..., + order: _OrderKACF = ..., + casting: _CastingKind = ..., + buffersize: SupportsIndex = ..., +) -> tuple[nditer, ...]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numeric.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numeric.py new file mode 100644 index 0000000000000000000000000000000000000000..86c7b7c64d5e37ec1585b2d3317381e5bd81120b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numeric.py @@ -0,0 +1,2713 @@ +import functools +import itertools +import operator +import sys +import warnings +import numbers +import builtins +import math + +import numpy as np +from . import multiarray +from . import numerictypes as nt +from .multiarray import ( + ALLOW_THREADS, BUFSIZE, CLIP, MAXDIMS, MAY_SHARE_BOUNDS, MAY_SHARE_EXACT, + RAISE, WRAP, arange, array, asarray, asanyarray, ascontiguousarray, + asfortranarray, broadcast, can_cast, concatenate, copyto, dot, dtype, + empty, empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter, + fromstring, inner, lexsort, matmul, may_share_memory, min_scalar_type, + ndarray, nditer, nested_iters, promote_types, putmask, result_type, + shares_memory, vdot, where, zeros, normalize_axis_index, vecdot +) + +from . import overrides +from . import umath +from . import shape_base +from .overrides import finalize_array_function_like, set_module +from .umath import (multiply, invert, sin, PINF, NAN) +from . import numerictypes +from ..exceptions import AxisError +from ._ufunc_config import errstate + +bitwise_not = invert +ufunc = type(sin) +newaxis = None + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +__all__ = [ + 'newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc', + 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', + 'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype', + 'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where', + 'argwhere', 'copyto', 'concatenate', 'lexsort', 'astype', + 'can_cast', 'promote_types', 'min_scalar_type', + 'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like', + 'correlate', 'convolve', 'inner', 'dot', 'outer', 'vdot', 'roll', + 'rollaxis', 'moveaxis', 'cross', 'tensordot', 'little_endian', + 'fromiter', 'array_equal', 'array_equiv', 'indices', 'fromfunction', + 'isclose', 'isscalar', 'binary_repr', 'base_repr', 'ones', + 'identity', 'allclose', 'putmask', + 'flatnonzero', 'inf', 'nan', 'False_', 'True_', 'bitwise_not', + 'full', 'full_like', 'matmul', 'vecdot', 'shares_memory', + 'may_share_memory'] + + +def _zeros_like_dispatcher( + a, dtype=None, order=None, subok=None, shape=None, *, device=None +): + return (a,) + + +@array_function_dispatch(_zeros_like_dispatcher) +def zeros_like( + a, dtype=None, order='K', subok=True, shape=None, *, device=None +): + """ + Return an array of zeros with the same shape and type as a given array. + + Parameters + ---------- + a : array_like + The shape and data-type of `a` define these same attributes of + the returned array. + dtype : data-type, optional + Overrides the data type of the result. + order : {'C', 'F', 'A', or 'K'}, optional + Overrides the memory layout of the result. 'C' means C-order, + 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous, + 'C' otherwise. 'K' means match the layout of `a` as closely + as possible. + subok : bool, optional. + If True, then the newly created array will use the sub-class + type of `a`, otherwise it will be a base-class array. Defaults + to True. + shape : int or sequence of ints, optional. + Overrides the shape of the result. If order='K' and the number of + dimensions is unchanged, will try to keep order, otherwise, + order='C' is implied. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + + Returns + ------- + out : ndarray + Array of zeros with the same shape and type as `a`. + + See Also + -------- + empty_like : Return an empty array with shape and type of input. + ones_like : Return an array of ones with shape and type of input. + full_like : Return a new array with shape of input filled with value. + zeros : Return a new array setting values to zero. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(6) + >>> x = x.reshape((2, 3)) + >>> x + array([[0, 1, 2], + [3, 4, 5]]) + >>> np.zeros_like(x) + array([[0, 0, 0], + [0, 0, 0]]) + + >>> y = np.arange(3, dtype=float) + >>> y + array([0., 1., 2.]) + >>> np.zeros_like(y) + array([0., 0., 0.]) + + """ + res = empty_like( + a, dtype=dtype, order=order, subok=subok, shape=shape, device=device + ) + # needed instead of a 0 to get same result as zeros for string dtypes + z = zeros(1, dtype=res.dtype) + multiarray.copyto(res, z, casting='unsafe') + return res + + +@finalize_array_function_like +@set_module('numpy') +def ones(shape, dtype=None, order='C', *, device=None, like=None): + """ + Return a new array of given shape and type, filled with ones. + + Parameters + ---------- + shape : int or sequence of ints + Shape of the new array, e.g., ``(2, 3)`` or ``2``. + dtype : data-type, optional + The desired data-type for the array, e.g., `numpy.int8`. Default is + `numpy.float64`. + order : {'C', 'F'}, optional, default: C + Whether to store multi-dimensional data in row-major + (C-style) or column-major (Fortran-style) order in + memory. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + out : ndarray + Array of ones with the given shape, dtype, and order. + + See Also + -------- + ones_like : Return an array of ones with shape and type of input. + empty : Return a new uninitialized array. + zeros : Return a new array setting values to zero. + full : Return a new array of given shape filled with value. + + Examples + -------- + >>> import numpy as np + >>> np.ones(5) + array([1., 1., 1., 1., 1.]) + + >>> np.ones((5,), dtype=int) + array([1, 1, 1, 1, 1]) + + >>> np.ones((2, 1)) + array([[1.], + [1.]]) + + >>> s = (2,2) + >>> np.ones(s) + array([[1., 1.], + [1., 1.]]) + + """ + if like is not None: + return _ones_with_like( + like, shape, dtype=dtype, order=order, device=device + ) + + a = empty(shape, dtype, order, device=device) + multiarray.copyto(a, 1, casting='unsafe') + return a + + +_ones_with_like = array_function_dispatch()(ones) + + +def _ones_like_dispatcher( + a, dtype=None, order=None, subok=None, shape=None, *, device=None +): + return (a,) + + +@array_function_dispatch(_ones_like_dispatcher) +def ones_like( + a, dtype=None, order='K', subok=True, shape=None, *, device=None +): + """ + Return an array of ones with the same shape and type as a given array. + + Parameters + ---------- + a : array_like + The shape and data-type of `a` define these same attributes of + the returned array. + dtype : data-type, optional + Overrides the data type of the result. + order : {'C', 'F', 'A', or 'K'}, optional + Overrides the memory layout of the result. 'C' means C-order, + 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous, + 'C' otherwise. 'K' means match the layout of `a` as closely + as possible. + subok : bool, optional. + If True, then the newly created array will use the sub-class + type of `a`, otherwise it will be a base-class array. Defaults + to True. + shape : int or sequence of ints, optional. + Overrides the shape of the result. If order='K' and the number of + dimensions is unchanged, will try to keep order, otherwise, + order='C' is implied. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + + Returns + ------- + out : ndarray + Array of ones with the same shape and type as `a`. + + See Also + -------- + empty_like : Return an empty array with shape and type of input. + zeros_like : Return an array of zeros with shape and type of input. + full_like : Return a new array with shape of input filled with value. + ones : Return a new array setting values to one. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(6) + >>> x = x.reshape((2, 3)) + >>> x + array([[0, 1, 2], + [3, 4, 5]]) + >>> np.ones_like(x) + array([[1, 1, 1], + [1, 1, 1]]) + + >>> y = np.arange(3, dtype=float) + >>> y + array([0., 1., 2.]) + >>> np.ones_like(y) + array([1., 1., 1.]) + + """ + res = empty_like( + a, dtype=dtype, order=order, subok=subok, shape=shape, device=device + ) + multiarray.copyto(res, 1, casting='unsafe') + return res + + +def _full_dispatcher( + shape, fill_value, dtype=None, order=None, *, device=None, like=None +): + return(like,) + + +@finalize_array_function_like +@set_module('numpy') +def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None): + """ + Return a new array of given shape and type, filled with `fill_value`. + + Parameters + ---------- + shape : int or sequence of ints + Shape of the new array, e.g., ``(2, 3)`` or ``2``. + fill_value : scalar or array_like + Fill value. + dtype : data-type, optional + The desired data-type for the array The default, None, means + ``np.array(fill_value).dtype``. + order : {'C', 'F'}, optional + Whether to store multidimensional data in C- or Fortran-contiguous + (row- or column-wise) order in memory. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + out : ndarray + Array of `fill_value` with the given shape, dtype, and order. + + See Also + -------- + full_like : Return a new array with shape of input filled with value. + empty : Return a new uninitialized array. + ones : Return a new array setting values to one. + zeros : Return a new array setting values to zero. + + Examples + -------- + >>> import numpy as np + >>> np.full((2, 2), np.inf) + array([[inf, inf], + [inf, inf]]) + >>> np.full((2, 2), 10) + array([[10, 10], + [10, 10]]) + + >>> np.full((2, 2), [1, 2]) + array([[1, 2], + [1, 2]]) + + """ + if like is not None: + return _full_with_like( + like, shape, fill_value, dtype=dtype, order=order, device=device + ) + + if dtype is None: + fill_value = asarray(fill_value) + dtype = fill_value.dtype + a = empty(shape, dtype, order, device=device) + multiarray.copyto(a, fill_value, casting='unsafe') + return a + + +_full_with_like = array_function_dispatch()(full) + + +def _full_like_dispatcher( + a, fill_value, dtype=None, order=None, subok=None, shape=None, + *, device=None +): + return (a,) + + +@array_function_dispatch(_full_like_dispatcher) +def full_like( + a, fill_value, dtype=None, order='K', subok=True, shape=None, + *, device=None +): + """ + Return a full array with the same shape and type as a given array. + + Parameters + ---------- + a : array_like + The shape and data-type of `a` define these same attributes of + the returned array. + fill_value : array_like + Fill value. + dtype : data-type, optional + Overrides the data type of the result. + order : {'C', 'F', 'A', or 'K'}, optional + Overrides the memory layout of the result. 'C' means C-order, + 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous, + 'C' otherwise. 'K' means match the layout of `a` as closely + as possible. + subok : bool, optional. + If True, then the newly created array will use the sub-class + type of `a`, otherwise it will be a base-class array. Defaults + to True. + shape : int or sequence of ints, optional. + Overrides the shape of the result. If order='K' and the number of + dimensions is unchanged, will try to keep order, otherwise, + order='C' is implied. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + + Returns + ------- + out : ndarray + Array of `fill_value` with the same shape and type as `a`. + + See Also + -------- + empty_like : Return an empty array with shape and type of input. + ones_like : Return an array of ones with shape and type of input. + zeros_like : Return an array of zeros with shape and type of input. + full : Return a new array of given shape filled with value. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(6, dtype=int) + >>> np.full_like(x, 1) + array([1, 1, 1, 1, 1, 1]) + >>> np.full_like(x, 0.1) + array([0, 0, 0, 0, 0, 0]) + >>> np.full_like(x, 0.1, dtype=np.double) + array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1]) + >>> np.full_like(x, np.nan, dtype=np.double) + array([nan, nan, nan, nan, nan, nan]) + + >>> y = np.arange(6, dtype=np.double) + >>> np.full_like(y, 0.1) + array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1]) + + >>> y = np.zeros([2, 2, 3], dtype=int) + >>> np.full_like(y, [0, 0, 255]) + array([[[ 0, 0, 255], + [ 0, 0, 255]], + [[ 0, 0, 255], + [ 0, 0, 255]]]) + """ + res = empty_like( + a, dtype=dtype, order=order, subok=subok, shape=shape, device=device + ) + multiarray.copyto(res, fill_value, casting='unsafe') + return res + + +def _count_nonzero_dispatcher(a, axis=None, *, keepdims=None): + return (a,) + + +@array_function_dispatch(_count_nonzero_dispatcher) +def count_nonzero(a, axis=None, *, keepdims=False): + """ + Counts the number of non-zero values in the array ``a``. + + The word "non-zero" is in reference to the Python 2.x + built-in method ``__nonzero__()`` (renamed ``__bool__()`` + in Python 3.x) of Python objects that tests an object's + "truthfulness". For example, any number is considered + truthful if it is nonzero, whereas any string is considered + truthful if it is not the empty string. Thus, this function + (recursively) counts how many elements in ``a`` (and in + sub-arrays thereof) have their ``__nonzero__()`` or ``__bool__()`` + method evaluated to ``True``. + + Parameters + ---------- + a : array_like + The array for which to count non-zeros. + axis : int or tuple, optional + Axis or tuple of axes along which to count non-zeros. + Default is None, meaning that non-zeros will be counted + along a flattened version of ``a``. + keepdims : bool, optional + If this is set to True, the axes that are counted are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + Returns + ------- + count : int or array of int + Number of non-zero values in the array along a given axis. + Otherwise, the total number of non-zero values in the array + is returned. + + See Also + -------- + nonzero : Return the coordinates of all the non-zero values. + + Examples + -------- + >>> import numpy as np + >>> np.count_nonzero(np.eye(4)) + 4 + >>> a = np.array([[0, 1, 7, 0], + ... [3, 0, 2, 19]]) + >>> np.count_nonzero(a) + 5 + >>> np.count_nonzero(a, axis=0) + array([1, 1, 2, 1]) + >>> np.count_nonzero(a, axis=1) + array([2, 3]) + >>> np.count_nonzero(a, axis=1, keepdims=True) + array([[2], + [3]]) + """ + if axis is None and not keepdims: + return multiarray.count_nonzero(a) + + a = asanyarray(a) + + # TODO: this works around .astype(bool) not working properly (gh-9847) + if np.issubdtype(a.dtype, np.character): + a_bool = a != a.dtype.type() + else: + a_bool = a.astype(np.bool, copy=False) + + return a_bool.sum(axis=axis, dtype=np.intp, keepdims=keepdims) + + +@set_module('numpy') +def isfortran(a): + """ + Check if the array is Fortran contiguous but *not* C contiguous. + + This function is obsolete. If you only want to check if an array is Fortran + contiguous use ``a.flags.f_contiguous`` instead. + + Parameters + ---------- + a : ndarray + Input array. + + Returns + ------- + isfortran : bool + Returns True if the array is Fortran contiguous but *not* C contiguous. + + + Examples + -------- + + np.array allows to specify whether the array is written in C-contiguous + order (last index varies the fastest), or FORTRAN-contiguous order in + memory (first index varies the fastest). + + >>> import numpy as np + >>> a = np.array([[1, 2, 3], [4, 5, 6]], order='C') + >>> a + array([[1, 2, 3], + [4, 5, 6]]) + >>> np.isfortran(a) + False + + >>> b = np.array([[1, 2, 3], [4, 5, 6]], order='F') + >>> b + array([[1, 2, 3], + [4, 5, 6]]) + >>> np.isfortran(b) + True + + + The transpose of a C-ordered array is a FORTRAN-ordered array. + + >>> a = np.array([[1, 2, 3], [4, 5, 6]], order='C') + >>> a + array([[1, 2, 3], + [4, 5, 6]]) + >>> np.isfortran(a) + False + >>> b = a.T + >>> b + array([[1, 4], + [2, 5], + [3, 6]]) + >>> np.isfortran(b) + True + + C-ordered arrays evaluate as False even if they are also FORTRAN-ordered. + + >>> np.isfortran(np.array([1, 2], order='F')) + False + + """ + return a.flags.fnc + + +def _argwhere_dispatcher(a): + return (a,) + + +@array_function_dispatch(_argwhere_dispatcher) +def argwhere(a): + """ + Find the indices of array elements that are non-zero, grouped by element. + + Parameters + ---------- + a : array_like + Input data. + + Returns + ------- + index_array : (N, a.ndim) ndarray + Indices of elements that are non-zero. Indices are grouped by element. + This array will have shape ``(N, a.ndim)`` where ``N`` is the number of + non-zero items. + + See Also + -------- + where, nonzero + + Notes + ----- + ``np.argwhere(a)`` is almost the same as ``np.transpose(np.nonzero(a))``, + but produces a result of the correct shape for a 0D array. + + The output of ``argwhere`` is not suitable for indexing arrays. + For this purpose use ``nonzero(a)`` instead. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(6).reshape(2,3) + >>> x + array([[0, 1, 2], + [3, 4, 5]]) + >>> np.argwhere(x>1) + array([[0, 2], + [1, 0], + [1, 1], + [1, 2]]) + + """ + # nonzero does not behave well on 0d, so promote to 1d + if np.ndim(a) == 0: + a = shape_base.atleast_1d(a) + # then remove the added dimension + return argwhere(a)[:, :0] + return transpose(nonzero(a)) + + +def _flatnonzero_dispatcher(a): + return (a,) + + +@array_function_dispatch(_flatnonzero_dispatcher) +def flatnonzero(a): + """ + Return indices that are non-zero in the flattened version of a. + + This is equivalent to ``np.nonzero(np.ravel(a))[0]``. + + Parameters + ---------- + a : array_like + Input data. + + Returns + ------- + res : ndarray + Output array, containing the indices of the elements of ``a.ravel()`` + that are non-zero. + + See Also + -------- + nonzero : Return the indices of the non-zero elements of the input array. + ravel : Return a 1-D array containing the elements of the input array. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(-2, 3) + >>> x + array([-2, -1, 0, 1, 2]) + >>> np.flatnonzero(x) + array([0, 1, 3, 4]) + + Use the indices of the non-zero elements as an index array to extract + these elements: + + >>> x.ravel()[np.flatnonzero(x)] + array([-2, -1, 1, 2]) + + """ + return np.nonzero(np.ravel(a))[0] + + +def _correlate_dispatcher(a, v, mode=None): + return (a, v) + + +@array_function_dispatch(_correlate_dispatcher) +def correlate(a, v, mode='valid'): + r""" + Cross-correlation of two 1-dimensional sequences. + + This function computes the correlation as generally defined in signal + processing texts [1]_: + + .. math:: c_k = \sum_n a_{n+k} \cdot \overline{v}_n + + with a and v sequences being zero-padded where necessary and + :math:`\overline v` denoting complex conjugation. + + Parameters + ---------- + a, v : array_like + Input sequences. + mode : {'valid', 'same', 'full'}, optional + Refer to the `convolve` docstring. Note that the default + is 'valid', unlike `convolve`, which uses 'full'. + + Returns + ------- + out : ndarray + Discrete cross-correlation of `a` and `v`. + + See Also + -------- + convolve : Discrete, linear convolution of two one-dimensional sequences. + scipy.signal.correlate : uses FFT which has superior performance + on large arrays. + + Notes + ----- + The definition of correlation above is not unique and sometimes + correlation may be defined differently. Another common definition is [1]_: + + .. math:: c'_k = \sum_n a_{n} \cdot \overline{v_{n+k}} + + which is related to :math:`c_k` by :math:`c'_k = c_{-k}`. + + `numpy.correlate` may perform slowly in large arrays (i.e. n = 1e5) + because it does not use the FFT to compute the convolution; in that case, + `scipy.signal.correlate` might be preferable. + + References + ---------- + .. [1] Wikipedia, "Cross-correlation", + https://en.wikipedia.org/wiki/Cross-correlation + + Examples + -------- + >>> import numpy as np + >>> np.correlate([1, 2, 3], [0, 1, 0.5]) + array([3.5]) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], "same") + array([2. , 3.5, 3. ]) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], "full") + array([0.5, 2. , 3.5, 3. , 0. ]) + + Using complex sequences: + + >>> np.correlate([1+1j, 2, 3-1j], [0, 1, 0.5j], 'full') + array([ 0.5-0.5j, 1.0+0.j , 1.5-1.5j, 3.0-1.j , 0.0+0.j ]) + + Note that you get the time reversed, complex conjugated result + (:math:`\overline{c_{-k}}`) when the two input sequences a and v change + places: + + >>> np.correlate([0, 1, 0.5j], [1+1j, 2, 3-1j], 'full') + array([ 0.0+0.j , 3.0+1.j , 1.5+1.5j, 1.0+0.j , 0.5+0.5j]) + + """ + return multiarray.correlate2(a, v, mode) + + +def _convolve_dispatcher(a, v, mode=None): + return (a, v) + + +@array_function_dispatch(_convolve_dispatcher) +def convolve(a, v, mode='full'): + """ + Returns the discrete, linear convolution of two one-dimensional sequences. + + The convolution operator is often seen in signal processing, where it + models the effect of a linear time-invariant system on a signal [1]_. In + probability theory, the sum of two independent random variables is + distributed according to the convolution of their individual + distributions. + + If `v` is longer than `a`, the arrays are swapped before computation. + + Parameters + ---------- + a : (N,) array_like + First one-dimensional input array. + v : (M,) array_like + Second one-dimensional input array. + mode : {'full', 'valid', 'same'}, optional + 'full': + By default, mode is 'full'. This returns the convolution + at each point of overlap, with an output shape of (N+M-1,). At + the end-points of the convolution, the signals do not overlap + completely, and boundary effects may be seen. + + 'same': + Mode 'same' returns output of length ``max(M, N)``. Boundary + effects are still visible. + + 'valid': + Mode 'valid' returns output of length + ``max(M, N) - min(M, N) + 1``. The convolution product is only given + for points where the signals overlap completely. Values outside + the signal boundary have no effect. + + Returns + ------- + out : ndarray + Discrete, linear convolution of `a` and `v`. + + See Also + -------- + scipy.signal.fftconvolve : Convolve two arrays using the Fast Fourier + Transform. + scipy.linalg.toeplitz : Used to construct the convolution operator. + polymul : Polynomial multiplication. Same output as convolve, but also + accepts poly1d objects as input. + + Notes + ----- + The discrete convolution operation is defined as + + .. math:: (a * v)_n = \\sum_{m = -\\infty}^{\\infty} a_m v_{n - m} + + It can be shown that a convolution :math:`x(t) * y(t)` in time/space + is equivalent to the multiplication :math:`X(f) Y(f)` in the Fourier + domain, after appropriate padding (padding is necessary to prevent + circular convolution). Since multiplication is more efficient (faster) + than convolution, the function `scipy.signal.fftconvolve` exploits the + FFT to calculate the convolution of large data-sets. + + References + ---------- + .. [1] Wikipedia, "Convolution", + https://en.wikipedia.org/wiki/Convolution + + Examples + -------- + Note how the convolution operator flips the second array + before "sliding" the two across one another: + + >>> import numpy as np + >>> np.convolve([1, 2, 3], [0, 1, 0.5]) + array([0. , 1. , 2.5, 4. , 1.5]) + + Only return the middle values of the convolution. + Contains boundary effects, where zeros are taken + into account: + + >>> np.convolve([1,2,3],[0,1,0.5], 'same') + array([1. , 2.5, 4. ]) + + The two arrays are of the same length, so there + is only one position where they completely overlap: + + >>> np.convolve([1,2,3],[0,1,0.5], 'valid') + array([2.5]) + + """ + a, v = array(a, copy=None, ndmin=1), array(v, copy=None, ndmin=1) + if (len(v) > len(a)): + a, v = v, a + if len(a) == 0: + raise ValueError('a cannot be empty') + if len(v) == 0: + raise ValueError('v cannot be empty') + return multiarray.correlate(a, v[::-1], mode) + + +def _outer_dispatcher(a, b, out=None): + return (a, b, out) + + +@array_function_dispatch(_outer_dispatcher) +def outer(a, b, out=None): + """ + Compute the outer product of two vectors. + + Given two vectors `a` and `b` of length ``M`` and ``N``, respectively, + the outer product [1]_ is:: + + [[a_0*b_0 a_0*b_1 ... a_0*b_{N-1} ] + [a_1*b_0 . + [ ... . + [a_{M-1}*b_0 a_{M-1}*b_{N-1} ]] + + Parameters + ---------- + a : (M,) array_like + First input vector. Input is flattened if + not already 1-dimensional. + b : (N,) array_like + Second input vector. Input is flattened if + not already 1-dimensional. + out : (M, N) ndarray, optional + A location where the result is stored + + Returns + ------- + out : (M, N) ndarray + ``out[i, j] = a[i] * b[j]`` + + See also + -------- + inner + einsum : ``einsum('i,j->ij', a.ravel(), b.ravel())`` is the equivalent. + ufunc.outer : A generalization to dimensions other than 1D and other + operations. ``np.multiply.outer(a.ravel(), b.ravel())`` + is the equivalent. + linalg.outer : An Array API compatible variation of ``np.outer``, + which accepts 1-dimensional inputs only. + tensordot : ``np.tensordot(a.ravel(), b.ravel(), axes=((), ()))`` + is the equivalent. + + References + ---------- + .. [1] G. H. Golub and C. F. Van Loan, *Matrix Computations*, 3rd + ed., Baltimore, MD, Johns Hopkins University Press, 1996, + pg. 8. + + Examples + -------- + Make a (*very* coarse) grid for computing a Mandelbrot set: + + >>> import numpy as np + >>> rl = np.outer(np.ones((5,)), np.linspace(-2, 2, 5)) + >>> rl + array([[-2., -1., 0., 1., 2.], + [-2., -1., 0., 1., 2.], + [-2., -1., 0., 1., 2.], + [-2., -1., 0., 1., 2.], + [-2., -1., 0., 1., 2.]]) + >>> im = np.outer(1j*np.linspace(2, -2, 5), np.ones((5,))) + >>> im + array([[0.+2.j, 0.+2.j, 0.+2.j, 0.+2.j, 0.+2.j], + [0.+1.j, 0.+1.j, 0.+1.j, 0.+1.j, 0.+1.j], + [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], + [0.-1.j, 0.-1.j, 0.-1.j, 0.-1.j, 0.-1.j], + [0.-2.j, 0.-2.j, 0.-2.j, 0.-2.j, 0.-2.j]]) + >>> grid = rl + im + >>> grid + array([[-2.+2.j, -1.+2.j, 0.+2.j, 1.+2.j, 2.+2.j], + [-2.+1.j, -1.+1.j, 0.+1.j, 1.+1.j, 2.+1.j], + [-2.+0.j, -1.+0.j, 0.+0.j, 1.+0.j, 2.+0.j], + [-2.-1.j, -1.-1.j, 0.-1.j, 1.-1.j, 2.-1.j], + [-2.-2.j, -1.-2.j, 0.-2.j, 1.-2.j, 2.-2.j]]) + + An example using a "vector" of letters: + + >>> x = np.array(['a', 'b', 'c'], dtype=object) + >>> np.outer(x, [1, 2, 3]) + array([['a', 'aa', 'aaa'], + ['b', 'bb', 'bbb'], + ['c', 'cc', 'ccc']], dtype=object) + + """ + a = asarray(a) + b = asarray(b) + return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis, :], out) + + +def _tensordot_dispatcher(a, b, axes=None): + return (a, b) + + +@array_function_dispatch(_tensordot_dispatcher) +def tensordot(a, b, axes=2): + """ + Compute tensor dot product along specified axes. + + Given two tensors, `a` and `b`, and an array_like object containing + two array_like objects, ``(a_axes, b_axes)``, sum the products of + `a`'s and `b`'s elements (components) over the axes specified by + ``a_axes`` and ``b_axes``. The third argument can be a single non-negative + integer_like scalar, ``N``; if it is such, then the last ``N`` dimensions + of `a` and the first ``N`` dimensions of `b` are summed over. + + Parameters + ---------- + a, b : array_like + Tensors to "dot". + + axes : int or (2,) array_like + * integer_like + If an int N, sum over the last N axes of `a` and the first N axes + of `b` in order. The sizes of the corresponding axes must match. + * (2,) array_like + Or, a list of axes to be summed over, first sequence applying to `a`, + second to `b`. Both elements array_like must be of the same length. + + Returns + ------- + output : ndarray + The tensor dot product of the input. + + See Also + -------- + dot, einsum + + Notes + ----- + Three common use cases are: + * ``axes = 0`` : tensor product :math:`a\\otimes b` + * ``axes = 1`` : tensor dot product :math:`a\\cdot b` + * ``axes = 2`` : (default) tensor double contraction :math:`a:b` + + When `axes` is integer_like, the sequence of axes for evaluation + will be: from the -Nth axis to the -1th axis in `a`, + and from the 0th axis to (N-1)th axis in `b`. + For example, ``axes = 2`` is the equal to + ``axes = [[-2, -1], [0, 1]]``. + When N-1 is smaller than 0, or when -N is larger than -1, + the element of `a` and `b` are defined as the `axes`. + + When there is more than one axis to sum over - and they are not the last + (first) axes of `a` (`b`) - the argument `axes` should consist of + two sequences of the same length, with the first axis to sum over given + first in both sequences, the second axis second, and so forth. + The calculation can be referred to ``numpy.einsum``. + + The shape of the result consists of the non-contracted axes of the + first tensor, followed by the non-contracted axes of the second. + + Examples + -------- + An example on integer_like: + + >>> a_0 = np.array([[1, 2], [3, 4]]) + >>> b_0 = np.array([[5, 6], [7, 8]]) + >>> c_0 = np.tensordot(a_0, b_0, axes=0) + >>> c_0.shape + (2, 2, 2, 2) + >>> c_0 + array([[[[ 5, 6], + [ 7, 8]], + [[10, 12], + [14, 16]]], + [[[15, 18], + [21, 24]], + [[20, 24], + [28, 32]]]]) + + An example on array_like: + + >>> a = np.arange(60.).reshape(3,4,5) + >>> b = np.arange(24.).reshape(4,3,2) + >>> c = np.tensordot(a,b, axes=([1,0],[0,1])) + >>> c.shape + (5, 2) + >>> c + array([[4400., 4730.], + [4532., 4874.], + [4664., 5018.], + [4796., 5162.], + [4928., 5306.]]) + + A slower but equivalent way of computing the same... + + >>> d = np.zeros((5,2)) + >>> for i in range(5): + ... for j in range(2): + ... for k in range(3): + ... for n in range(4): + ... d[i,j] += a[k,n,i] * b[n,k,j] + >>> c == d + array([[ True, True], + [ True, True], + [ True, True], + [ True, True], + [ True, True]]) + + An extended example taking advantage of the overloading of + and \\*: + + >>> a = np.array(range(1, 9)) + >>> a.shape = (2, 2, 2) + >>> A = np.array(('a', 'b', 'c', 'd'), dtype=object) + >>> A.shape = (2, 2) + >>> a; A + array([[[1, 2], + [3, 4]], + [[5, 6], + [7, 8]]]) + array([['a', 'b'], + ['c', 'd']], dtype=object) + + >>> np.tensordot(a, A) # third argument default is 2 for double-contraction + array(['abbcccdddd', 'aaaaabbbbbbcccccccdddddddd'], dtype=object) + + >>> np.tensordot(a, A, 1) + array([[['acc', 'bdd'], + ['aaacccc', 'bbbdddd']], + [['aaaaacccccc', 'bbbbbdddddd'], + ['aaaaaaacccccccc', 'bbbbbbbdddddddd']]], dtype=object) + + >>> np.tensordot(a, A, 0) # tensor product (result too long to incl.) + array([[[[['a', 'b'], + ['c', 'd']], + ... + + >>> np.tensordot(a, A, (0, 1)) + array([[['abbbbb', 'cddddd'], + ['aabbbbbb', 'ccdddddd']], + [['aaabbbbbbb', 'cccddddddd'], + ['aaaabbbbbbbb', 'ccccdddddddd']]], dtype=object) + + >>> np.tensordot(a, A, (2, 1)) + array([[['abb', 'cdd'], + ['aaabbbb', 'cccdddd']], + [['aaaaabbbbbb', 'cccccdddddd'], + ['aaaaaaabbbbbbbb', 'cccccccdddddddd']]], dtype=object) + + >>> np.tensordot(a, A, ((0, 1), (0, 1))) + array(['abbbcccccddddddd', 'aabbbbccccccdddddddd'], dtype=object) + + >>> np.tensordot(a, A, ((2, 1), (1, 0))) + array(['acccbbdddd', 'aaaaacccccccbbbbbbdddddddd'], dtype=object) + + """ + try: + iter(axes) + except Exception: + axes_a = list(range(-axes, 0)) + axes_b = list(range(0, axes)) + else: + axes_a, axes_b = axes + try: + na = len(axes_a) + axes_a = list(axes_a) + except TypeError: + axes_a = [axes_a] + na = 1 + try: + nb = len(axes_b) + axes_b = list(axes_b) + except TypeError: + axes_b = [axes_b] + nb = 1 + + a, b = asarray(a), asarray(b) + as_ = a.shape + nda = a.ndim + bs = b.shape + ndb = b.ndim + equal = True + if na != nb: + equal = False + else: + for k in range(na): + if as_[axes_a[k]] != bs[axes_b[k]]: + equal = False + break + if axes_a[k] < 0: + axes_a[k] += nda + if axes_b[k] < 0: + axes_b[k] += ndb + if not equal: + raise ValueError("shape-mismatch for sum") + + # Move the axes to sum over to the end of "a" + # and to the front of "b" + notin = [k for k in range(nda) if k not in axes_a] + newaxes_a = notin + axes_a + N2 = math.prod(as_[axis] for axis in axes_a) + newshape_a = (math.prod([as_[ax] for ax in notin]), N2) + olda = [as_[axis] for axis in notin] + + notin = [k for k in range(ndb) if k not in axes_b] + newaxes_b = axes_b + notin + N2 = math.prod(bs[axis] for axis in axes_b) + newshape_b = (N2, math.prod([bs[ax] for ax in notin])) + oldb = [bs[axis] for axis in notin] + + at = a.transpose(newaxes_a).reshape(newshape_a) + bt = b.transpose(newaxes_b).reshape(newshape_b) + res = dot(at, bt) + return res.reshape(olda + oldb) + + +def _roll_dispatcher(a, shift, axis=None): + return (a,) + + +@array_function_dispatch(_roll_dispatcher) +def roll(a, shift, axis=None): + """ + Roll array elements along a given axis. + + Elements that roll beyond the last position are re-introduced at + the first. + + Parameters + ---------- + a : array_like + Input array. + shift : int or tuple of ints + The number of places by which elements are shifted. If a tuple, + then `axis` must be a tuple of the same size, and each of the + given axes is shifted by the corresponding number. If an int + while `axis` is a tuple of ints, then the same value is used for + all given axes. + axis : int or tuple of ints, optional + Axis or axes along which elements are shifted. By default, the + array is flattened before shifting, after which the original + shape is restored. + + Returns + ------- + res : ndarray + Output array, with the same shape as `a`. + + See Also + -------- + rollaxis : Roll the specified axis backwards, until it lies in a + given position. + + Notes + ----- + Supports rolling over multiple dimensions simultaneously. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(10) + >>> np.roll(x, 2) + array([8, 9, 0, 1, 2, 3, 4, 5, 6, 7]) + >>> np.roll(x, -2) + array([2, 3, 4, 5, 6, 7, 8, 9, 0, 1]) + + >>> x2 = np.reshape(x, (2, 5)) + >>> x2 + array([[0, 1, 2, 3, 4], + [5, 6, 7, 8, 9]]) + >>> np.roll(x2, 1) + array([[9, 0, 1, 2, 3], + [4, 5, 6, 7, 8]]) + >>> np.roll(x2, -1) + array([[1, 2, 3, 4, 5], + [6, 7, 8, 9, 0]]) + >>> np.roll(x2, 1, axis=0) + array([[5, 6, 7, 8, 9], + [0, 1, 2, 3, 4]]) + >>> np.roll(x2, -1, axis=0) + array([[5, 6, 7, 8, 9], + [0, 1, 2, 3, 4]]) + >>> np.roll(x2, 1, axis=1) + array([[4, 0, 1, 2, 3], + [9, 5, 6, 7, 8]]) + >>> np.roll(x2, -1, axis=1) + array([[1, 2, 3, 4, 0], + [6, 7, 8, 9, 5]]) + >>> np.roll(x2, (1, 1), axis=(1, 0)) + array([[9, 5, 6, 7, 8], + [4, 0, 1, 2, 3]]) + >>> np.roll(x2, (2, 1), axis=(1, 0)) + array([[8, 9, 5, 6, 7], + [3, 4, 0, 1, 2]]) + + """ + a = asanyarray(a) + if axis is None: + return roll(a.ravel(), shift, 0).reshape(a.shape) + + else: + axis = normalize_axis_tuple(axis, a.ndim, allow_duplicate=True) + broadcasted = broadcast(shift, axis) + if broadcasted.ndim > 1: + raise ValueError( + "'shift' and 'axis' should be scalars or 1D sequences") + shifts = {ax: 0 for ax in range(a.ndim)} + for sh, ax in broadcasted: + shifts[ax] += int(sh) + + rolls = [((slice(None), slice(None)),)] * a.ndim + for ax, offset in shifts.items(): + offset %= a.shape[ax] or 1 # If `a` is empty, nothing matters. + if offset: + # (original, result), (original, result) + rolls[ax] = ((slice(None, -offset), slice(offset, None)), + (slice(-offset, None), slice(None, offset))) + + result = empty_like(a) + for indices in itertools.product(*rolls): + arr_index, res_index = zip(*indices) + result[res_index] = a[arr_index] + + return result + + +def _rollaxis_dispatcher(a, axis, start=None): + return (a,) + + +@array_function_dispatch(_rollaxis_dispatcher) +def rollaxis(a, axis, start=0): + """ + Roll the specified axis backwards, until it lies in a given position. + + This function continues to be supported for backward compatibility, but you + should prefer `moveaxis`. The `moveaxis` function was added in NumPy + 1.11. + + Parameters + ---------- + a : ndarray + Input array. + axis : int + The axis to be rolled. The positions of the other axes do not + change relative to one another. + start : int, optional + When ``start <= axis``, the axis is rolled back until it lies in + this position. When ``start > axis``, the axis is rolled until it + lies before this position. The default, 0, results in a "complete" + roll. The following table describes how negative values of ``start`` + are interpreted: + + .. table:: + :align: left + + +-------------------+----------------------+ + | ``start`` | Normalized ``start`` | + +===================+======================+ + | ``-(arr.ndim+1)`` | raise ``AxisError`` | + +-------------------+----------------------+ + | ``-arr.ndim`` | 0 | + +-------------------+----------------------+ + | |vdots| | |vdots| | + +-------------------+----------------------+ + | ``-1`` | ``arr.ndim-1`` | + +-------------------+----------------------+ + | ``0`` | ``0`` | + +-------------------+----------------------+ + | |vdots| | |vdots| | + +-------------------+----------------------+ + | ``arr.ndim`` | ``arr.ndim`` | + +-------------------+----------------------+ + | ``arr.ndim + 1`` | raise ``AxisError`` | + +-------------------+----------------------+ + + .. |vdots| unicode:: U+22EE .. Vertical Ellipsis + + Returns + ------- + res : ndarray + For NumPy >= 1.10.0 a view of `a` is always returned. For earlier + NumPy versions a view of `a` is returned only if the order of the + axes is changed, otherwise the input array is returned. + + See Also + -------- + moveaxis : Move array axes to new positions. + roll : Roll the elements of an array by a number of positions along a + given axis. + + Examples + -------- + >>> import numpy as np + >>> a = np.ones((3,4,5,6)) + >>> np.rollaxis(a, 3, 1).shape + (3, 6, 4, 5) + >>> np.rollaxis(a, 2).shape + (5, 3, 4, 6) + >>> np.rollaxis(a, 1, 4).shape + (3, 5, 6, 4) + + """ + n = a.ndim + axis = normalize_axis_index(axis, n) + if start < 0: + start += n + msg = "'%s' arg requires %d <= %s < %d, but %d was passed in" + if not (0 <= start < n + 1): + raise AxisError(msg % ('start', -n, 'start', n + 1, start)) + if axis < start: + # it's been removed + start -= 1 + if axis == start: + return a[...] + axes = list(range(0, n)) + axes.remove(axis) + axes.insert(start, axis) + return a.transpose(axes) + + +@set_module("numpy.lib.array_utils") +def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False): + """ + Normalizes an axis argument into a tuple of non-negative integer axes. + + This handles shorthands such as ``1`` and converts them to ``(1,)``, + as well as performing the handling of negative indices covered by + `normalize_axis_index`. + + By default, this forbids axes from being specified multiple times. + + Used internally by multi-axis-checking logic. + + Parameters + ---------- + axis : int, iterable of int + The un-normalized index or indices of the axis. + ndim : int + The number of dimensions of the array that `axis` should be normalized + against. + argname : str, optional + A prefix to put before the error message, typically the name of the + argument. + allow_duplicate : bool, optional + If False, the default, disallow an axis from being specified twice. + + Returns + ------- + normalized_axes : tuple of int + The normalized axis index, such that `0 <= normalized_axis < ndim` + + Raises + ------ + AxisError + If any axis provided is out of range + ValueError + If an axis is repeated + + See also + -------- + normalize_axis_index : normalizing a single scalar axis + """ + # Optimization to speed-up the most common cases. + if type(axis) not in (tuple, list): + try: + axis = [operator.index(axis)] + except TypeError: + pass + # Going via an iterator directly is slower than via list comprehension. + axis = tuple([normalize_axis_index(ax, ndim, argname) for ax in axis]) + if not allow_duplicate and len(set(axis)) != len(axis): + if argname: + raise ValueError('repeated axis in `{}` argument'.format(argname)) + else: + raise ValueError('repeated axis') + return axis + + +def _moveaxis_dispatcher(a, source, destination): + return (a,) + + +@array_function_dispatch(_moveaxis_dispatcher) +def moveaxis(a, source, destination): + """ + Move axes of an array to new positions. + + Other axes remain in their original order. + + Parameters + ---------- + a : np.ndarray + The array whose axes should be reordered. + source : int or sequence of int + Original positions of the axes to move. These must be unique. + destination : int or sequence of int + Destination positions for each of the original axes. These must also be + unique. + + Returns + ------- + result : np.ndarray + Array with moved axes. This array is a view of the input array. + + See Also + -------- + transpose : Permute the dimensions of an array. + swapaxes : Interchange two axes of an array. + + Examples + -------- + >>> import numpy as np + >>> x = np.zeros((3, 4, 5)) + >>> np.moveaxis(x, 0, -1).shape + (4, 5, 3) + >>> np.moveaxis(x, -1, 0).shape + (5, 3, 4) + + These all achieve the same result: + + >>> np.transpose(x).shape + (5, 4, 3) + >>> np.swapaxes(x, 0, -1).shape + (5, 4, 3) + >>> np.moveaxis(x, [0, 1], [-1, -2]).shape + (5, 4, 3) + >>> np.moveaxis(x, [0, 1, 2], [-1, -2, -3]).shape + (5, 4, 3) + + """ + try: + # allow duck-array types if they define transpose + transpose = a.transpose + except AttributeError: + a = asarray(a) + transpose = a.transpose + + source = normalize_axis_tuple(source, a.ndim, 'source') + destination = normalize_axis_tuple(destination, a.ndim, 'destination') + if len(source) != len(destination): + raise ValueError('`source` and `destination` arguments must have ' + 'the same number of elements') + + order = [n for n in range(a.ndim) if n not in source] + + for dest, src in sorted(zip(destination, source)): + order.insert(dest, src) + + result = transpose(order) + return result + + +def _cross_dispatcher(a, b, axisa=None, axisb=None, axisc=None, axis=None): + return (a, b) + + +@array_function_dispatch(_cross_dispatcher) +def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None): + """ + Return the cross product of two (arrays of) vectors. + + The cross product of `a` and `b` in :math:`R^3` is a vector perpendicular + to both `a` and `b`. If `a` and `b` are arrays of vectors, the vectors + are defined by the last axis of `a` and `b` by default, and these axes + can have dimensions 2 or 3. Where the dimension of either `a` or `b` is + 2, the third component of the input vector is assumed to be zero and the + cross product calculated accordingly. In cases where both input vectors + have dimension 2, the z-component of the cross product is returned. + + Parameters + ---------- + a : array_like + Components of the first vector(s). + b : array_like + Components of the second vector(s). + axisa : int, optional + Axis of `a` that defines the vector(s). By default, the last axis. + axisb : int, optional + Axis of `b` that defines the vector(s). By default, the last axis. + axisc : int, optional + Axis of `c` containing the cross product vector(s). Ignored if + both input vectors have dimension 2, as the return is scalar. + By default, the last axis. + axis : int, optional + If defined, the axis of `a`, `b` and `c` that defines the vector(s) + and cross product(s). Overrides `axisa`, `axisb` and `axisc`. + + Returns + ------- + c : ndarray + Vector cross product(s). + + Raises + ------ + ValueError + When the dimension of the vector(s) in `a` and/or `b` does not + equal 2 or 3. + + See Also + -------- + inner : Inner product + outer : Outer product. + linalg.cross : An Array API compatible variation of ``np.cross``, + which accepts (arrays of) 3-element vectors only. + ix_ : Construct index arrays. + + Notes + ----- + Supports full broadcasting of the inputs. + + Dimension-2 input arrays were deprecated in 2.0.0. If you do need this + functionality, you can use:: + + def cross2d(x, y): + return x[..., 0] * y[..., 1] - x[..., 1] * y[..., 0] + + Examples + -------- + Vector cross-product. + + >>> import numpy as np + >>> x = [1, 2, 3] + >>> y = [4, 5, 6] + >>> np.cross(x, y) + array([-3, 6, -3]) + + One vector with dimension 2. + + >>> x = [1, 2] + >>> y = [4, 5, 6] + >>> np.cross(x, y) + array([12, -6, -3]) + + Equivalently: + + >>> x = [1, 2, 0] + >>> y = [4, 5, 6] + >>> np.cross(x, y) + array([12, -6, -3]) + + Both vectors with dimension 2. + + >>> x = [1,2] + >>> y = [4,5] + >>> np.cross(x, y) + array(-3) + + Multiple vector cross-products. Note that the direction of the cross + product vector is defined by the *right-hand rule*. + + >>> x = np.array([[1,2,3], [4,5,6]]) + >>> y = np.array([[4,5,6], [1,2,3]]) + >>> np.cross(x, y) + array([[-3, 6, -3], + [ 3, -6, 3]]) + + The orientation of `c` can be changed using the `axisc` keyword. + + >>> np.cross(x, y, axisc=0) + array([[-3, 3], + [ 6, -6], + [-3, 3]]) + + Change the vector definition of `x` and `y` using `axisa` and `axisb`. + + >>> x = np.array([[1,2,3], [4,5,6], [7, 8, 9]]) + >>> y = np.array([[7, 8, 9], [4,5,6], [1,2,3]]) + >>> np.cross(x, y) + array([[ -6, 12, -6], + [ 0, 0, 0], + [ 6, -12, 6]]) + >>> np.cross(x, y, axisa=0, axisb=0) + array([[-24, 48, -24], + [-30, 60, -30], + [-36, 72, -36]]) + + """ + if axis is not None: + axisa, axisb, axisc = (axis,) * 3 + a = asarray(a) + b = asarray(b) + + if (a.ndim < 1) or (b.ndim < 1): + raise ValueError("At least one array has zero dimension") + + # Check axisa and axisb are within bounds + axisa = normalize_axis_index(axisa, a.ndim, msg_prefix='axisa') + axisb = normalize_axis_index(axisb, b.ndim, msg_prefix='axisb') + + # Move working axis to the end of the shape + a = moveaxis(a, axisa, -1) + b = moveaxis(b, axisb, -1) + msg = ("incompatible dimensions for cross product\n" + "(dimension must be 2 or 3)") + if a.shape[-1] not in (2, 3) or b.shape[-1] not in (2, 3): + raise ValueError(msg) + if a.shape[-1] == 2 or b.shape[-1] == 2: + # Deprecated in NumPy 2.0, 2023-09-26 + warnings.warn( + "Arrays of 2-dimensional vectors are deprecated. Use arrays of " + "3-dimensional vectors instead. (deprecated in NumPy 2.0)", + DeprecationWarning, stacklevel=2 + ) + + # Create the output array + shape = broadcast(a[..., 0], b[..., 0]).shape + if a.shape[-1] == 3 or b.shape[-1] == 3: + shape += (3,) + # Check axisc is within bounds + axisc = normalize_axis_index(axisc, len(shape), msg_prefix='axisc') + dtype = promote_types(a.dtype, b.dtype) + cp = empty(shape, dtype) + + # recast arrays as dtype + a = a.astype(dtype) + b = b.astype(dtype) + + # create local aliases for readability + a0 = a[..., 0] + a1 = a[..., 1] + if a.shape[-1] == 3: + a2 = a[..., 2] + b0 = b[..., 0] + b1 = b[..., 1] + if b.shape[-1] == 3: + b2 = b[..., 2] + if cp.ndim != 0 and cp.shape[-1] == 3: + cp0 = cp[..., 0] + cp1 = cp[..., 1] + cp2 = cp[..., 2] + + if a.shape[-1] == 2: + if b.shape[-1] == 2: + # a0 * b1 - a1 * b0 + multiply(a0, b1, out=cp) + cp -= a1 * b0 + return cp + else: + assert b.shape[-1] == 3 + # cp0 = a1 * b2 - 0 (a2 = 0) + # cp1 = 0 - a0 * b2 (a2 = 0) + # cp2 = a0 * b1 - a1 * b0 + multiply(a1, b2, out=cp0) + multiply(a0, b2, out=cp1) + negative(cp1, out=cp1) + multiply(a0, b1, out=cp2) + cp2 -= a1 * b0 + else: + assert a.shape[-1] == 3 + if b.shape[-1] == 3: + # cp0 = a1 * b2 - a2 * b1 + # cp1 = a2 * b0 - a0 * b2 + # cp2 = a0 * b1 - a1 * b0 + multiply(a1, b2, out=cp0) + tmp = array(a2 * b1) + cp0 -= tmp + multiply(a2, b0, out=cp1) + multiply(a0, b2, out=tmp) + cp1 -= tmp + multiply(a0, b1, out=cp2) + multiply(a1, b0, out=tmp) + cp2 -= tmp + else: + assert b.shape[-1] == 2 + # cp0 = 0 - a2 * b1 (b2 = 0) + # cp1 = a2 * b0 - 0 (b2 = 0) + # cp2 = a0 * b1 - a1 * b0 + multiply(a2, b1, out=cp0) + negative(cp0, out=cp0) + multiply(a2, b0, out=cp1) + multiply(a0, b1, out=cp2) + cp2 -= a1 * b0 + + return moveaxis(cp, -1, axisc) + + +little_endian = (sys.byteorder == 'little') + + +@set_module('numpy') +def indices(dimensions, dtype=int, sparse=False): + """ + Return an array representing the indices of a grid. + + Compute an array where the subarrays contain index values 0, 1, ... + varying only along the corresponding axis. + + Parameters + ---------- + dimensions : sequence of ints + The shape of the grid. + dtype : dtype, optional + Data type of the result. + sparse : boolean, optional + Return a sparse representation of the grid instead of a dense + representation. Default is False. + + Returns + ------- + grid : one ndarray or tuple of ndarrays + If sparse is False: + Returns one array of grid indices, + ``grid.shape = (len(dimensions),) + tuple(dimensions)``. + If sparse is True: + Returns a tuple of arrays, with + ``grid[i].shape = (1, ..., 1, dimensions[i], 1, ..., 1)`` with + dimensions[i] in the ith place + + See Also + -------- + mgrid, ogrid, meshgrid + + Notes + ----- + The output shape in the dense case is obtained by prepending the number + of dimensions in front of the tuple of dimensions, i.e. if `dimensions` + is a tuple ``(r0, ..., rN-1)`` of length ``N``, the output shape is + ``(N, r0, ..., rN-1)``. + + The subarrays ``grid[k]`` contains the N-D array of indices along the + ``k-th`` axis. Explicitly:: + + grid[k, i0, i1, ..., iN-1] = ik + + Examples + -------- + >>> import numpy as np + >>> grid = np.indices((2, 3)) + >>> grid.shape + (2, 2, 3) + >>> grid[0] # row indices + array([[0, 0, 0], + [1, 1, 1]]) + >>> grid[1] # column indices + array([[0, 1, 2], + [0, 1, 2]]) + + The indices can be used as an index into an array. + + >>> x = np.arange(20).reshape(5, 4) + >>> row, col = np.indices((2, 3)) + >>> x[row, col] + array([[0, 1, 2], + [4, 5, 6]]) + + Note that it would be more straightforward in the above example to + extract the required elements directly with ``x[:2, :3]``. + + If sparse is set to true, the grid will be returned in a sparse + representation. + + >>> i, j = np.indices((2, 3), sparse=True) + >>> i.shape + (2, 1) + >>> j.shape + (1, 3) + >>> i # row indices + array([[0], + [1]]) + >>> j # column indices + array([[0, 1, 2]]) + + """ + dimensions = tuple(dimensions) + N = len(dimensions) + shape = (1,)*N + if sparse: + res = tuple() + else: + res = empty((N,)+dimensions, dtype=dtype) + for i, dim in enumerate(dimensions): + idx = arange(dim, dtype=dtype).reshape( + shape[:i] + (dim,) + shape[i+1:] + ) + if sparse: + res = res + (idx,) + else: + res[i] = idx + return res + + +@finalize_array_function_like +@set_module('numpy') +def fromfunction(function, shape, *, dtype=float, like=None, **kwargs): + """ + Construct an array by executing a function over each coordinate. + + The resulting array therefore has a value ``fn(x, y, z)`` at + coordinate ``(x, y, z)``. + + Parameters + ---------- + function : callable + The function is called with N parameters, where N is the rank of + `shape`. Each parameter represents the coordinates of the array + varying along a specific axis. For example, if `shape` + were ``(2, 2)``, then the parameters would be + ``array([[0, 0], [1, 1]])`` and ``array([[0, 1], [0, 1]])`` + shape : (N,) tuple of ints + Shape of the output array, which also determines the shape of + the coordinate arrays passed to `function`. + dtype : data-type, optional + Data-type of the coordinate arrays passed to `function`. + By default, `dtype` is float. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + fromfunction : any + The result of the call to `function` is passed back directly. + Therefore the shape of `fromfunction` is completely determined by + `function`. If `function` returns a scalar value, the shape of + `fromfunction` would not match the `shape` parameter. + + See Also + -------- + indices, meshgrid + + Notes + ----- + Keywords other than `dtype` and `like` are passed to `function`. + + Examples + -------- + >>> import numpy as np + >>> np.fromfunction(lambda i, j: i, (2, 2), dtype=float) + array([[0., 0.], + [1., 1.]]) + + >>> np.fromfunction(lambda i, j: j, (2, 2), dtype=float) + array([[0., 1.], + [0., 1.]]) + + >>> np.fromfunction(lambda i, j: i == j, (3, 3), dtype=int) + array([[ True, False, False], + [False, True, False], + [False, False, True]]) + + >>> np.fromfunction(lambda i, j: i + j, (3, 3), dtype=int) + array([[0, 1, 2], + [1, 2, 3], + [2, 3, 4]]) + + """ + if like is not None: + return _fromfunction_with_like( + like, function, shape, dtype=dtype, **kwargs) + + args = indices(shape, dtype=dtype) + return function(*args, **kwargs) + + +_fromfunction_with_like = array_function_dispatch()(fromfunction) + + +def _frombuffer(buf, dtype, shape, order): + return frombuffer(buf, dtype=dtype).reshape(shape, order=order) + + +@set_module('numpy') +def isscalar(element): + """ + Returns True if the type of `element` is a scalar type. + + Parameters + ---------- + element : any + Input argument, can be of any type and shape. + + Returns + ------- + val : bool + True if `element` is a scalar type, False if it is not. + + See Also + -------- + ndim : Get the number of dimensions of an array + + Notes + ----- + If you need a stricter way to identify a *numerical* scalar, use + ``isinstance(x, numbers.Number)``, as that returns ``False`` for most + non-numerical elements such as strings. + + In most cases ``np.ndim(x) == 0`` should be used instead of this function, + as that will also return true for 0d arrays. This is how numpy overloads + functions in the style of the ``dx`` arguments to `gradient` and + the ``bins`` argument to `histogram`. Some key differences: + + +------------------------------------+---------------+-------------------+ + | x |``isscalar(x)``|``np.ndim(x) == 0``| + +====================================+===============+===================+ + | PEP 3141 numeric objects | ``True`` | ``True`` | + | (including builtins) | | | + +------------------------------------+---------------+-------------------+ + | builtin string and buffer objects | ``True`` | ``True`` | + +------------------------------------+---------------+-------------------+ + | other builtin objects, like | ``False`` | ``True`` | + | `pathlib.Path`, `Exception`, | | | + | the result of `re.compile` | | | + +------------------------------------+---------------+-------------------+ + | third-party objects like | ``False`` | ``True`` | + | `matplotlib.figure.Figure` | | | + +------------------------------------+---------------+-------------------+ + | zero-dimensional numpy arrays | ``False`` | ``True`` | + +------------------------------------+---------------+-------------------+ + | other numpy arrays | ``False`` | ``False`` | + +------------------------------------+---------------+-------------------+ + | `list`, `tuple`, and other | ``False`` | ``False`` | + | sequence objects | | | + +------------------------------------+---------------+-------------------+ + + Examples + -------- + >>> import numpy as np + + >>> np.isscalar(3.1) + True + + >>> np.isscalar(np.array(3.1)) + False + + >>> np.isscalar([3.1]) + False + + >>> np.isscalar(False) + True + + >>> np.isscalar('numpy') + True + + NumPy supports PEP 3141 numbers: + + >>> from fractions import Fraction + >>> np.isscalar(Fraction(5, 17)) + True + >>> from numbers import Number + >>> np.isscalar(Number()) + True + + """ + return (isinstance(element, generic) + or type(element) in ScalarType + or isinstance(element, numbers.Number)) + + +@set_module('numpy') +def binary_repr(num, width=None): + """ + Return the binary representation of the input number as a string. + + For negative numbers, if width is not given, a minus sign is added to the + front. If width is given, the two's complement of the number is + returned, with respect to that width. + + In a two's-complement system negative numbers are represented by the two's + complement of the absolute value. This is the most common method of + representing signed integers on computers [1]_. A N-bit two's-complement + system can represent every integer in the range + :math:`-2^{N-1}` to :math:`+2^{N-1}-1`. + + Parameters + ---------- + num : int + Only an integer decimal number can be used. + width : int, optional + The length of the returned string if `num` is positive, or the length + of the two's complement if `num` is negative, provided that `width` is + at least a sufficient number of bits for `num` to be represented in + the designated form. If the `width` value is insufficient, an error is + raised. + + Returns + ------- + bin : str + Binary representation of `num` or two's complement of `num`. + + See Also + -------- + base_repr: Return a string representation of a number in the given base + system. + bin: Python's built-in binary representation generator of an integer. + + Notes + ----- + `binary_repr` is equivalent to using `base_repr` with base 2, but about 25x + faster. + + References + ---------- + .. [1] Wikipedia, "Two's complement", + https://en.wikipedia.org/wiki/Two's_complement + + Examples + -------- + >>> import numpy as np + >>> np.binary_repr(3) + '11' + >>> np.binary_repr(-3) + '-11' + >>> np.binary_repr(3, width=4) + '0011' + + The two's complement is returned when the input number is negative and + width is specified: + + >>> np.binary_repr(-3, width=3) + '101' + >>> np.binary_repr(-3, width=5) + '11101' + + """ + def err_if_insufficient(width, binwidth): + if width is not None and width < binwidth: + raise ValueError( + f"Insufficient bit {width=} provided for {binwidth=}" + ) + + # Ensure that num is a Python integer to avoid overflow or unwanted + # casts to floating point. + num = operator.index(num) + + if num == 0: + return '0' * (width or 1) + + elif num > 0: + binary = bin(num)[2:] + binwidth = len(binary) + outwidth = (binwidth if width is None + else builtins.max(binwidth, width)) + err_if_insufficient(width, binwidth) + return binary.zfill(outwidth) + + else: + if width is None: + return '-' + bin(-num)[2:] + + else: + poswidth = len(bin(-num)[2:]) + + # See gh-8679: remove extra digit + # for numbers at boundaries. + if 2**(poswidth - 1) == -num: + poswidth -= 1 + + twocomp = 2**(poswidth + 1) + num + binary = bin(twocomp)[2:] + binwidth = len(binary) + + outwidth = builtins.max(binwidth, width) + err_if_insufficient(width, binwidth) + return '1' * (outwidth - binwidth) + binary + + +@set_module('numpy') +def base_repr(number, base=2, padding=0): + """ + Return a string representation of a number in the given base system. + + Parameters + ---------- + number : int + The value to convert. Positive and negative values are handled. + base : int, optional + Convert `number` to the `base` number system. The valid range is 2-36, + the default value is 2. + padding : int, optional + Number of zeros padded on the left. Default is 0 (no padding). + + Returns + ------- + out : str + String representation of `number` in `base` system. + + See Also + -------- + binary_repr : Faster version of `base_repr` for base 2. + + Examples + -------- + >>> import numpy as np + >>> np.base_repr(5) + '101' + >>> np.base_repr(6, 5) + '11' + >>> np.base_repr(7, base=5, padding=3) + '00012' + + >>> np.base_repr(10, base=16) + 'A' + >>> np.base_repr(32, base=16) + '20' + + """ + digits = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' + if base > len(digits): + raise ValueError("Bases greater than 36 not handled in base_repr.") + elif base < 2: + raise ValueError("Bases less than 2 not handled in base_repr.") + + num = abs(int(number)) + res = [] + while num: + res.append(digits[num % base]) + num //= base + if padding: + res.append('0' * padding) + if number < 0: + res.append('-') + return ''.join(reversed(res or '0')) + + +# These are all essentially abbreviations +# These might wind up in a special abbreviations module + + +def _maketup(descr, val): + dt = dtype(descr) + # Place val in all scalar tuples: + fields = dt.fields + if fields is None: + return val + else: + res = [_maketup(fields[name][0], val) for name in dt.names] + return tuple(res) + + +@finalize_array_function_like +@set_module('numpy') +def identity(n, dtype=None, *, like=None): + """ + Return the identity array. + + The identity array is a square array with ones on + the main diagonal. + + Parameters + ---------- + n : int + Number of rows (and columns) in `n` x `n` output. + dtype : data-type, optional + Data-type of the output. Defaults to ``float``. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + out : ndarray + `n` x `n` array with its main diagonal set to one, + and all other elements 0. + + Examples + -------- + >>> import numpy as np + >>> np.identity(3) + array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + + """ + if like is not None: + return _identity_with_like(like, n, dtype=dtype) + + from numpy import eye + return eye(n, dtype=dtype, like=like) + + +_identity_with_like = array_function_dispatch()(identity) + + +def _allclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None): + return (a, b, rtol, atol) + + +@array_function_dispatch(_allclose_dispatcher) +def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False): + """ + Returns True if two arrays are element-wise equal within a tolerance. + + The tolerance values are positive, typically very small numbers. The + relative difference (`rtol` * abs(`b`)) and the absolute difference + `atol` are added together to compare against the absolute difference + between `a` and `b`. + + .. warning:: The default `atol` is not appropriate for comparing numbers + with magnitudes much smaller than one (see Notes). + + NaNs are treated as equal if they are in the same place and if + ``equal_nan=True``. Infs are treated as equal if they are in the same + place and of the same sign in both arrays. + + Parameters + ---------- + a, b : array_like + Input arrays to compare. + rtol : array_like + The relative tolerance parameter (see Notes). + atol : array_like + The absolute tolerance parameter (see Notes). + equal_nan : bool + Whether to compare NaN's as equal. If True, NaN's in `a` will be + considered equal to NaN's in `b` in the output array. + + Returns + ------- + allclose : bool + Returns True if the two arrays are equal within the given + tolerance; False otherwise. + + See Also + -------- + isclose, all, any, equal + + Notes + ----- + If the following equation is element-wise True, then allclose returns + True.:: + + absolute(a - b) <= (atol + rtol * absolute(b)) + + The above equation is not symmetric in `a` and `b`, so that + ``allclose(a, b)`` might be different from ``allclose(b, a)`` in + some rare cases. + + The default value of `atol` is not appropriate when the reference value + `b` has magnitude smaller than one. For example, it is unlikely that + ``a = 1e-9`` and ``b = 2e-9`` should be considered "close", yet + ``allclose(1e-9, 2e-9)`` is ``True`` with default settings. Be sure + to select `atol` for the use case at hand, especially for defining the + threshold below which a non-zero value in `a` will be considered "close" + to a very small or zero value in `b`. + + The comparison of `a` and `b` uses standard broadcasting, which + means that `a` and `b` need not have the same shape in order for + ``allclose(a, b)`` to evaluate to True. The same is true for + `equal` but not `array_equal`. + + `allclose` is not defined for non-numeric data types. + `bool` is considered a numeric data-type for this purpose. + + Examples + -------- + >>> import numpy as np + >>> np.allclose([1e10,1e-7], [1.00001e10,1e-8]) + False + + >>> np.allclose([1e10,1e-8], [1.00001e10,1e-9]) + True + + >>> np.allclose([1e10,1e-8], [1.0001e10,1e-9]) + False + + >>> np.allclose([1.0, np.nan], [1.0, np.nan]) + False + + >>> np.allclose([1.0, np.nan], [1.0, np.nan], equal_nan=True) + True + + + """ + res = all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)) + return builtins.bool(res) + + +def _isclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None): + return (a, b, rtol, atol) + + +@array_function_dispatch(_isclose_dispatcher) +def isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False): + """ + Returns a boolean array where two arrays are element-wise equal within a + tolerance. + + The tolerance values are positive, typically very small numbers. The + relative difference (`rtol` * abs(`b`)) and the absolute difference + `atol` are added together to compare against the absolute difference + between `a` and `b`. + + .. warning:: The default `atol` is not appropriate for comparing numbers + with magnitudes much smaller than one (see Notes). + + Parameters + ---------- + a, b : array_like + Input arrays to compare. + rtol : array_like + The relative tolerance parameter (see Notes). + atol : array_like + The absolute tolerance parameter (see Notes). + equal_nan : bool + Whether to compare NaN's as equal. If True, NaN's in `a` will be + considered equal to NaN's in `b` in the output array. + + Returns + ------- + y : array_like + Returns a boolean array of where `a` and `b` are equal within the + given tolerance. If both `a` and `b` are scalars, returns a single + boolean value. + + See Also + -------- + allclose + math.isclose + + Notes + ----- + For finite values, isclose uses the following equation to test whether + two floating point values are equivalent.:: + + absolute(a - b) <= (atol + rtol * absolute(b)) + + Unlike the built-in `math.isclose`, the above equation is not symmetric + in `a` and `b` -- it assumes `b` is the reference value -- so that + `isclose(a, b)` might be different from `isclose(b, a)`. + + The default value of `atol` is not appropriate when the reference value + `b` has magnitude smaller than one. For example, it is unlikely that + ``a = 1e-9`` and ``b = 2e-9`` should be considered "close", yet + ``isclose(1e-9, 2e-9)`` is ``True`` with default settings. Be sure + to select `atol` for the use case at hand, especially for defining the + threshold below which a non-zero value in `a` will be considered "close" + to a very small or zero value in `b`. + + `isclose` is not defined for non-numeric data types. + :class:`bool` is considered a numeric data-type for this purpose. + + Examples + -------- + >>> import numpy as np + >>> np.isclose([1e10,1e-7], [1.00001e10,1e-8]) + array([ True, False]) + + >>> np.isclose([1e10,1e-8], [1.00001e10,1e-9]) + array([ True, True]) + + >>> np.isclose([1e10,1e-8], [1.0001e10,1e-9]) + array([False, True]) + + >>> np.isclose([1.0, np.nan], [1.0, np.nan]) + array([ True, False]) + + >>> np.isclose([1.0, np.nan], [1.0, np.nan], equal_nan=True) + array([ True, True]) + + >>> np.isclose([1e-8, 1e-7], [0.0, 0.0]) + array([ True, False]) + + >>> np.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0) + array([False, False]) + + >>> np.isclose([1e-10, 1e-10], [1e-20, 0.0]) + array([ True, True]) + + >>> np.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], atol=0.0) + array([False, True]) + + """ + # Turn all but python scalars into arrays. + x, y, atol, rtol = ( + a if isinstance(a, (int, float, complex)) else asanyarray(a) + for a in (a, b, atol, rtol)) + + # Make sure y is an inexact type to avoid bad behavior on abs(MIN_INT). + # This will cause casting of x later. Also, make sure to allow subclasses + # (e.g., for numpy.ma). + # NOTE: We explicitly allow timedelta, which used to work. This could + # possibly be deprecated. See also gh-18286. + # timedelta works if `atol` is an integer or also a timedelta. + # Although, the default tolerances are unlikely to be useful + if (dtype := getattr(y, "dtype", None)) is not None and dtype.kind != "m": + dt = multiarray.result_type(y, 1.) + y = asanyarray(y, dtype=dt) + elif isinstance(y, int): + y = float(y) + + with errstate(invalid='ignore'): + result = (less_equal(abs(x-y), atol + rtol * abs(y)) + & isfinite(y) + | (x == y)) + if equal_nan: + result |= isnan(x) & isnan(y) + + return result[()] # Flatten 0d arrays to scalars + + +def _array_equal_dispatcher(a1, a2, equal_nan=None): + return (a1, a2) + + +_no_nan_types = { + # should use np.dtype.BoolDType, but as of writing + # that fails the reloading test. + type(dtype(nt.bool)), + type(dtype(nt.int8)), + type(dtype(nt.int16)), + type(dtype(nt.int32)), + type(dtype(nt.int64)), +} + + +def _dtype_cannot_hold_nan(dtype): + return type(dtype) in _no_nan_types + + +@array_function_dispatch(_array_equal_dispatcher) +def array_equal(a1, a2, equal_nan=False): + """ + True if two arrays have the same shape and elements, False otherwise. + + Parameters + ---------- + a1, a2 : array_like + Input arrays. + equal_nan : bool + Whether to compare NaN's as equal. If the dtype of a1 and a2 is + complex, values will be considered equal if either the real or the + imaginary component of a given value is ``nan``. + + Returns + ------- + b : bool + Returns True if the arrays are equal. + + See Also + -------- + allclose: Returns True if two arrays are element-wise equal within a + tolerance. + array_equiv: Returns True if input arrays are shape consistent and all + elements equal. + + Examples + -------- + >>> import numpy as np + + >>> np.array_equal([1, 2], [1, 2]) + True + + >>> np.array_equal(np.array([1, 2]), np.array([1, 2])) + True + + >>> np.array_equal([1, 2], [1, 2, 3]) + False + + >>> np.array_equal([1, 2], [1, 4]) + False + + >>> a = np.array([1, np.nan]) + >>> np.array_equal(a, a) + False + + >>> np.array_equal(a, a, equal_nan=True) + True + + When ``equal_nan`` is True, complex values with nan components are + considered equal if either the real *or* the imaginary components are nan. + + >>> a = np.array([1 + 1j]) + >>> b = a.copy() + >>> a.real = np.nan + >>> b.imag = np.nan + >>> np.array_equal(a, b, equal_nan=True) + True + """ + try: + a1, a2 = asarray(a1), asarray(a2) + except Exception: + return False + if a1.shape != a2.shape: + return False + if not equal_nan: + return builtins.bool((asanyarray(a1 == a2)).all()) + + if a1 is a2: + # nan will compare equal so an array will compare equal to itself. + return True + + cannot_have_nan = (_dtype_cannot_hold_nan(a1.dtype) + and _dtype_cannot_hold_nan(a2.dtype)) + if cannot_have_nan: + return builtins.bool(asarray(a1 == a2).all()) + + # Handling NaN values if equal_nan is True + a1nan, a2nan = isnan(a1), isnan(a2) + # NaN's occur at different locations + if not (a1nan == a2nan).all(): + return False + # Shapes of a1, a2 and masks are guaranteed to be consistent by this point + return builtins.bool((a1[~a1nan] == a2[~a1nan]).all()) + + +def _array_equiv_dispatcher(a1, a2): + return (a1, a2) + + +@array_function_dispatch(_array_equiv_dispatcher) +def array_equiv(a1, a2): + """ + Returns True if input arrays are shape consistent and all elements equal. + + Shape consistent means they are either the same shape, or one input array + can be broadcasted to create the same shape as the other one. + + Parameters + ---------- + a1, a2 : array_like + Input arrays. + + Returns + ------- + out : bool + True if equivalent, False otherwise. + + Examples + -------- + >>> import numpy as np + >>> np.array_equiv([1, 2], [1, 2]) + True + >>> np.array_equiv([1, 2], [1, 3]) + False + + Showing the shape equivalence: + + >>> np.array_equiv([1, 2], [[1, 2], [1, 2]]) + True + >>> np.array_equiv([1, 2], [[1, 2, 1, 2], [1, 2, 1, 2]]) + False + + >>> np.array_equiv([1, 2], [[1, 2], [1, 3]]) + False + + """ + try: + a1, a2 = asarray(a1), asarray(a2) + except Exception: + return False + try: + multiarray.broadcast(a1, a2) + except Exception: + return False + + return builtins.bool(asanyarray(a1 == a2).all()) + + +def _astype_dispatcher(x, dtype, /, *, copy=None, device=None): + return (x, dtype) + + +@array_function_dispatch(_astype_dispatcher) +def astype(x, dtype, /, *, copy=True, device=None): + """ + Copies an array to a specified data type. + + This function is an Array API compatible alternative to + `numpy.ndarray.astype`. + + Parameters + ---------- + x : ndarray + Input NumPy array to cast. ``array_likes`` are explicitly not + supported here. + dtype : dtype + Data type of the result. + copy : bool, optional + Specifies whether to copy an array when the specified dtype matches + the data type of the input array ``x``. If ``True``, a newly allocated + array must always be returned. If ``False`` and the specified dtype + matches the data type of the input array, the input array must be + returned; otherwise, a newly allocated array must be returned. + Defaults to ``True``. + device : str, optional + The device on which to place the returned array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.1.0 + + Returns + ------- + out : ndarray + An array having the specified data type. + + See Also + -------- + ndarray.astype + + Examples + -------- + >>> import numpy as np + >>> arr = np.array([1, 2, 3]); arr + array([1, 2, 3]) + >>> np.astype(arr, np.float64) + array([1., 2., 3.]) + + Non-copy case: + + >>> arr = np.array([1, 2, 3]) + >>> arr_noncpy = np.astype(arr, arr.dtype, copy=False) + >>> np.shares_memory(arr, arr_noncpy) + True + + """ + if not (isinstance(x, np.ndarray) or isscalar(x)): + raise TypeError( + "Input should be a NumPy array or scalar. " + f"It is a {type(x)} instead." + ) + if device is not None and device != "cpu": + raise ValueError( + 'Device not understood. Only "cpu" is allowed, but received:' + f' {device}' + ) + return x.astype(dtype, copy=copy) + + +inf = PINF +nan = NAN +False_ = nt.bool(False) +True_ = nt.bool(True) + + +def extend_all(module): + existing = set(__all__) + mall = module.__all__ + for a in mall: + if a not in existing: + __all__.append(a) + + +from .umath import * +from .numerictypes import * +from . import fromnumeric +from .fromnumeric import * +from . import arrayprint +from .arrayprint import * +from . import _asarray +from ._asarray import * +from . import _ufunc_config +from ._ufunc_config import * +extend_all(fromnumeric) +extend_all(umath) +extend_all(numerictypes) +extend_all(arrayprint) +extend_all(_asarray) +extend_all(_ufunc_config) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numeric.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numeric.pyi new file mode 100644 index 0000000000000000000000000000000000000000..b6372aca93796026f3940d6fde04cbf46d568dda --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numeric.pyi @@ -0,0 +1,896 @@ +from collections.abc import Callable, Sequence +from typing import ( + Any, + Final, + TypeAlias, + overload, + TypeVar, + Literal as L, + SupportsAbs, + SupportsIndex, + NoReturn, + TypeGuard, +) +from typing_extensions import Unpack + +import numpy as np +from numpy import ( + # re-exports + bitwise_not, + False_, + True_, + broadcast, + dtype, + flatiter, + from_dlpack, + inf, + little_endian, + matmul, + vecdot, + nan, + ndarray, + nditer, + newaxis, + ufunc, + + # other + generic, + unsignedinteger, + signedinteger, + floating, + complexfloating, + int_, + intp, + float64, + timedelta64, + object_, + _AnyShapeType, + _OrderKACF, + _OrderCF, +) +from .fromnumeric import ( + all as all, + any as any, + argpartition as argpartition, + matrix_transpose as matrix_transpose, + mean as mean, +) +from .multiarray import ( + # re-exports + arange, + array, + asarray, + asanyarray, + ascontiguousarray, + asfortranarray, + can_cast, + concatenate, + copyto, + dot, + empty, + empty_like, + frombuffer, + fromfile, + fromiter, + fromstring, + inner, + lexsort, + may_share_memory, + min_scalar_type, + nested_iters, + putmask, + promote_types, + result_type, + shares_memory, + vdot, + where, + zeros, + + # other + _Array, + _ConstructorEmpty, + _KwargsEmpty, +) + +from numpy._typing import ( + ArrayLike, + NDArray, + DTypeLike, + _SupportsDType, + _ShapeLike, + _DTypeLike, + _ArrayLike, + _SupportsArrayFunc, + _ScalarLike_co, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeTD64_co, + _ArrayLikeObject_co, + _ArrayLikeUnknown, + _NestedSequence, +) + +__all__ = [ + "newaxis", + "ndarray", + "flatiter", + "nditer", + "nested_iters", + "ufunc", + "arange", + "array", + "asarray", + "asanyarray", + "ascontiguousarray", + "asfortranarray", + "zeros", + "count_nonzero", + "empty", + "broadcast", + "dtype", + "fromstring", + "fromfile", + "frombuffer", + "from_dlpack", + "where", + "argwhere", + "copyto", + "concatenate", + "lexsort", + "astype", + "can_cast", + "promote_types", + "min_scalar_type", + "result_type", + "isfortran", + "empty_like", + "zeros_like", + "ones_like", + "correlate", + "convolve", + "inner", + "dot", + "outer", + "vdot", + "roll", + "rollaxis", + "moveaxis", + "cross", + "tensordot", + "little_endian", + "fromiter", + "array_equal", + "array_equiv", + "indices", + "fromfunction", + "isclose", + "isscalar", + "binary_repr", + "base_repr", + "ones", + "identity", + "allclose", + "putmask", + "flatnonzero", + "inf", + "nan", + "False_", + "True_", + "bitwise_not", + "full", + "full_like", + "matmul", + "vecdot", + "shares_memory", + "may_share_memory", +] + +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=generic) +_DType = TypeVar("_DType", bound=np.dtype[Any]) +_ArrayType = TypeVar("_ArrayType", bound=np.ndarray[Any, Any]) +_ShapeType = TypeVar("_ShapeType", bound=tuple[int, ...]) + +_CorrelateMode: TypeAlias = L["valid", "same", "full"] + +@overload +def zeros_like( + a: _ArrayType, + dtype: None = ..., + order: _OrderKACF = ..., + subok: L[True] = ..., + shape: None = ..., + *, + device: None | L["cpu"] = ..., +) -> _ArrayType: ... +@overload +def zeros_like( + a: _ArrayLike[_SCT], + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def zeros_like( + a: object, + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... +@overload +def zeros_like( + a: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def zeros_like( + a: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... + +ones: Final[_ConstructorEmpty] + +@overload +def ones_like( + a: _ArrayType, + dtype: None = ..., + order: _OrderKACF = ..., + subok: L[True] = ..., + shape: None = ..., + *, + device: None | L["cpu"] = ..., +) -> _ArrayType: ... +@overload +def ones_like( + a: _ArrayLike[_SCT], + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def ones_like( + a: object, + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... +@overload +def ones_like( + a: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def ones_like( + a: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... + +# TODO: Add overloads for bool, int, float, complex, str, bytes, and memoryview +# 1-D shape +@overload +def full( + shape: SupportsIndex, + fill_value: _SCT, + dtype: None = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> _Array[tuple[int], _SCT]: ... +@overload +def full( + shape: SupportsIndex, + fill_value: Any, + dtype: _DType | _SupportsDType[_DType], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> np.ndarray[tuple[int], _DType]: ... +@overload +def full( + shape: SupportsIndex, + fill_value: Any, + dtype: type[_SCT], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> _Array[tuple[int], _SCT]: ... +@overload +def full( + shape: SupportsIndex, + fill_value: Any, + dtype: None | DTypeLike = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> _Array[tuple[int], Any]: ... +# known shape +@overload +def full( + shape: _AnyShapeType, + fill_value: _SCT, + dtype: None = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> _Array[_AnyShapeType, _SCT]: ... +@overload +def full( + shape: _AnyShapeType, + fill_value: Any, + dtype: _DType | _SupportsDType[_DType], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> np.ndarray[_AnyShapeType, _DType]: ... +@overload +def full( + shape: _AnyShapeType, + fill_value: Any, + dtype: type[_SCT], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> _Array[_AnyShapeType, _SCT]: ... +@overload +def full( + shape: _AnyShapeType, + fill_value: Any, + dtype: None | DTypeLike = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> _Array[_AnyShapeType, Any]: ... +# unknown shape +@overload +def full( + shape: _ShapeLike, + fill_value: _SCT, + dtype: None = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> NDArray[_SCT]: ... +@overload +def full( + shape: _ShapeLike, + fill_value: Any, + dtype: _DType | _SupportsDType[_DType], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> np.ndarray[Any, _DType]: ... +@overload +def full( + shape: _ShapeLike, + fill_value: Any, + dtype: type[_SCT], + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> NDArray[_SCT]: ... +@overload +def full( + shape: _ShapeLike, + fill_value: Any, + dtype: None | DTypeLike = ..., + order: _OrderCF = ..., + **kwargs: Unpack[_KwargsEmpty], +) -> NDArray[Any]: ... + +@overload +def full_like( + a: _ArrayType, + fill_value: Any, + dtype: None = ..., + order: _OrderKACF = ..., + subok: L[True] = ..., + shape: None = ..., + *, + device: None | L["cpu"] = ..., +) -> _ArrayType: ... +@overload +def full_like( + a: _ArrayLike[_SCT], + fill_value: Any, + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def full_like( + a: object, + fill_value: Any, + dtype: None = ..., + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... +@overload +def full_like( + a: Any, + fill_value: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[_SCT]: ... +@overload +def full_like( + a: Any, + fill_value: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., + subok: bool = ..., + shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., +) -> NDArray[Any]: ... + +# +@overload +def count_nonzero(a: ArrayLike, axis: None = None, *, keepdims: L[False] = False) -> int: ... +@overload +def count_nonzero(a: _ScalarLike_co, axis: _ShapeLike | None = None, *, keepdims: L[True]) -> np.intp: ... +@overload +def count_nonzero( + a: NDArray[Any] | _NestedSequence[ArrayLike], axis: _ShapeLike | None = None, *, keepdims: L[True] +) -> NDArray[np.intp]: ... +@overload +def count_nonzero(a: ArrayLike, axis: _ShapeLike | None = None, *, keepdims: bool = False) -> Any: ... + +# +def isfortran(a: NDArray[Any] | generic) -> bool: ... + +def argwhere(a: ArrayLike) -> NDArray[intp]: ... + +def flatnonzero(a: ArrayLike) -> NDArray[intp]: ... + +@overload +def correlate( + a: _ArrayLikeUnknown, + v: _ArrayLikeUnknown, + mode: _CorrelateMode = ..., +) -> NDArray[Any]: ... +@overload +def correlate( + a: _ArrayLikeBool_co, + v: _ArrayLikeBool_co, + mode: _CorrelateMode = ..., +) -> NDArray[np.bool]: ... +@overload +def correlate( + a: _ArrayLikeUInt_co, + v: _ArrayLikeUInt_co, + mode: _CorrelateMode = ..., +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def correlate( + a: _ArrayLikeInt_co, + v: _ArrayLikeInt_co, + mode: _CorrelateMode = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def correlate( + a: _ArrayLikeFloat_co, + v: _ArrayLikeFloat_co, + mode: _CorrelateMode = ..., +) -> NDArray[floating[Any]]: ... +@overload +def correlate( + a: _ArrayLikeComplex_co, + v: _ArrayLikeComplex_co, + mode: _CorrelateMode = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def correlate( + a: _ArrayLikeTD64_co, + v: _ArrayLikeTD64_co, + mode: _CorrelateMode = ..., +) -> NDArray[timedelta64]: ... +@overload +def correlate( + a: _ArrayLikeObject_co, + v: _ArrayLikeObject_co, + mode: _CorrelateMode = ..., +) -> NDArray[object_]: ... + +@overload +def convolve( + a: _ArrayLikeUnknown, + v: _ArrayLikeUnknown, + mode: _CorrelateMode = ..., +) -> NDArray[Any]: ... +@overload +def convolve( + a: _ArrayLikeBool_co, + v: _ArrayLikeBool_co, + mode: _CorrelateMode = ..., +) -> NDArray[np.bool]: ... +@overload +def convolve( + a: _ArrayLikeUInt_co, + v: _ArrayLikeUInt_co, + mode: _CorrelateMode = ..., +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def convolve( + a: _ArrayLikeInt_co, + v: _ArrayLikeInt_co, + mode: _CorrelateMode = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def convolve( + a: _ArrayLikeFloat_co, + v: _ArrayLikeFloat_co, + mode: _CorrelateMode = ..., +) -> NDArray[floating[Any]]: ... +@overload +def convolve( + a: _ArrayLikeComplex_co, + v: _ArrayLikeComplex_co, + mode: _CorrelateMode = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def convolve( + a: _ArrayLikeTD64_co, + v: _ArrayLikeTD64_co, + mode: _CorrelateMode = ..., +) -> NDArray[timedelta64]: ... +@overload +def convolve( + a: _ArrayLikeObject_co, + v: _ArrayLikeObject_co, + mode: _CorrelateMode = ..., +) -> NDArray[object_]: ... + +@overload +def outer( + a: _ArrayLikeUnknown, + b: _ArrayLikeUnknown, + out: None = ..., +) -> NDArray[Any]: ... +@overload +def outer( + a: _ArrayLikeBool_co, + b: _ArrayLikeBool_co, + out: None = ..., +) -> NDArray[np.bool]: ... +@overload +def outer( + a: _ArrayLikeUInt_co, + b: _ArrayLikeUInt_co, + out: None = ..., +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def outer( + a: _ArrayLikeInt_co, + b: _ArrayLikeInt_co, + out: None = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def outer( + a: _ArrayLikeFloat_co, + b: _ArrayLikeFloat_co, + out: None = ..., +) -> NDArray[floating[Any]]: ... +@overload +def outer( + a: _ArrayLikeComplex_co, + b: _ArrayLikeComplex_co, + out: None = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def outer( + a: _ArrayLikeTD64_co, + b: _ArrayLikeTD64_co, + out: None = ..., +) -> NDArray[timedelta64]: ... +@overload +def outer( + a: _ArrayLikeObject_co, + b: _ArrayLikeObject_co, + out: None = ..., +) -> NDArray[object_]: ... +@overload +def outer( + a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + b: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + out: _ArrayType, +) -> _ArrayType: ... + +@overload +def tensordot( + a: _ArrayLikeUnknown, + b: _ArrayLikeUnknown, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[Any]: ... +@overload +def tensordot( + a: _ArrayLikeBool_co, + b: _ArrayLikeBool_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[np.bool]: ... +@overload +def tensordot( + a: _ArrayLikeUInt_co, + b: _ArrayLikeUInt_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def tensordot( + a: _ArrayLikeInt_co, + b: _ArrayLikeInt_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def tensordot( + a: _ArrayLikeFloat_co, + b: _ArrayLikeFloat_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[floating[Any]]: ... +@overload +def tensordot( + a: _ArrayLikeComplex_co, + b: _ArrayLikeComplex_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def tensordot( + a: _ArrayLikeTD64_co, + b: _ArrayLikeTD64_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[timedelta64]: ... +@overload +def tensordot( + a: _ArrayLikeObject_co, + b: _ArrayLikeObject_co, + axes: int | tuple[_ShapeLike, _ShapeLike] = ..., +) -> NDArray[object_]: ... + +@overload +def roll( + a: _ArrayLike[_SCT], + shift: _ShapeLike, + axis: None | _ShapeLike = ..., +) -> NDArray[_SCT]: ... +@overload +def roll( + a: ArrayLike, + shift: _ShapeLike, + axis: None | _ShapeLike = ..., +) -> NDArray[Any]: ... + +def rollaxis( + a: NDArray[_SCT], + axis: int, + start: int = ..., +) -> NDArray[_SCT]: ... + +def moveaxis( + a: NDArray[_SCT], + source: _ShapeLike, + destination: _ShapeLike, +) -> NDArray[_SCT]: ... + +@overload +def cross( + a: _ArrayLikeUnknown, + b: _ArrayLikeUnknown, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NDArray[Any]: ... +@overload +def cross( + a: _ArrayLikeBool_co, + b: _ArrayLikeBool_co, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NoReturn: ... +@overload +def cross( + a: _ArrayLikeUInt_co, + b: _ArrayLikeUInt_co, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def cross( + a: _ArrayLikeInt_co, + b: _ArrayLikeInt_co, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def cross( + a: _ArrayLikeFloat_co, + b: _ArrayLikeFloat_co, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NDArray[floating[Any]]: ... +@overload +def cross( + a: _ArrayLikeComplex_co, + b: _ArrayLikeComplex_co, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def cross( + a: _ArrayLikeObject_co, + b: _ArrayLikeObject_co, + axisa: int = ..., + axisb: int = ..., + axisc: int = ..., + axis: None | int = ..., +) -> NDArray[object_]: ... + +@overload +def indices( + dimensions: Sequence[int], + dtype: type[int] = ..., + sparse: L[False] = ..., +) -> NDArray[int_]: ... +@overload +def indices( + dimensions: Sequence[int], + dtype: type[int] = ..., + sparse: L[True] = ..., +) -> tuple[NDArray[int_], ...]: ... +@overload +def indices( + dimensions: Sequence[int], + dtype: _DTypeLike[_SCT], + sparse: L[False] = ..., +) -> NDArray[_SCT]: ... +@overload +def indices( + dimensions: Sequence[int], + dtype: _DTypeLike[_SCT], + sparse: L[True], +) -> tuple[NDArray[_SCT], ...]: ... +@overload +def indices( + dimensions: Sequence[int], + dtype: DTypeLike, + sparse: L[False] = ..., +) -> NDArray[Any]: ... +@overload +def indices( + dimensions: Sequence[int], + dtype: DTypeLike, + sparse: L[True], +) -> tuple[NDArray[Any], ...]: ... + +def fromfunction( + function: Callable[..., _T], + shape: Sequence[int], + *, + dtype: DTypeLike = ..., + like: _SupportsArrayFunc = ..., + **kwargs: Any, +) -> _T: ... + +def isscalar(element: object) -> TypeGuard[ + generic | bool | int | float | complex | str | bytes | memoryview +]: ... + +def binary_repr(num: SupportsIndex, width: None | int = ...) -> str: ... + +def base_repr( + number: SupportsAbs[float], + base: float = ..., + padding: SupportsIndex = ..., +) -> str: ... + +@overload +def identity( + n: int, + dtype: None = ..., + *, + like: _SupportsArrayFunc = ..., +) -> NDArray[float64]: ... +@overload +def identity( + n: int, + dtype: _DTypeLike[_SCT], + *, + like: _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def identity( + n: int, + dtype: DTypeLike, + *, + like: _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +def allclose( + a: ArrayLike, + b: ArrayLike, + rtol: ArrayLike = ..., + atol: ArrayLike = ..., + equal_nan: bool = ..., +) -> bool: ... + +@overload +def isclose( + a: _ScalarLike_co, + b: _ScalarLike_co, + rtol: ArrayLike = ..., + atol: ArrayLike = ..., + equal_nan: bool = ..., +) -> np.bool: ... +@overload +def isclose( + a: ArrayLike, + b: ArrayLike, + rtol: ArrayLike = ..., + atol: ArrayLike = ..., + equal_nan: bool = ..., +) -> NDArray[np.bool]: ... + +def array_equal(a1: ArrayLike, a2: ArrayLike, equal_nan: bool = ...) -> bool: ... + +def array_equiv(a1: ArrayLike, a2: ArrayLike) -> bool: ... + +@overload +def astype( + x: ndarray[_ShapeType, dtype[Any]], + dtype: _DTypeLike[_SCT], + /, + *, + copy: bool = ..., + device: None | L["cpu"] = ..., +) -> ndarray[_ShapeType, dtype[_SCT]]: ... +@overload +def astype( + x: ndarray[_ShapeType, dtype[Any]], + dtype: DTypeLike, + /, + *, + copy: bool = ..., + device: None | L["cpu"] = ..., +) -> ndarray[_ShapeType, dtype[Any]]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numerictypes.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numerictypes.py new file mode 100644 index 0000000000000000000000000000000000000000..241233ad36e7ca4e74c968785296f8f5036bb720 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numerictypes.py @@ -0,0 +1,629 @@ +""" +numerictypes: Define the numeric type objects + +This module is designed so "from numerictypes import \\*" is safe. +Exported symbols include: + + Dictionary with all registered number types (including aliases): + sctypeDict + + Type objects (not all will be available, depends on platform): + see variable sctypes for which ones you have + + Bit-width names + + int8 int16 int32 int64 int128 + uint8 uint16 uint32 uint64 uint128 + float16 float32 float64 float96 float128 float256 + complex32 complex64 complex128 complex192 complex256 complex512 + datetime64 timedelta64 + + c-based names + + bool + + object_ + + void, str_ + + byte, ubyte, + short, ushort + intc, uintc, + intp, uintp, + int_, uint, + longlong, ulonglong, + + single, csingle, + double, cdouble, + longdouble, clongdouble, + + As part of the type-hierarchy: xx -- is bit-width + + generic + +-> bool (kind=b) + +-> number + | +-> integer + | | +-> signedinteger (intxx) (kind=i) + | | | byte + | | | short + | | | intc + | | | intp + | | | int_ + | | | longlong + | | \\-> unsignedinteger (uintxx) (kind=u) + | | ubyte + | | ushort + | | uintc + | | uintp + | | uint + | | ulonglong + | +-> inexact + | +-> floating (floatxx) (kind=f) + | | half + | | single + | | double + | | longdouble + | \\-> complexfloating (complexxx) (kind=c) + | csingle + | cdouble + | clongdouble + +-> flexible + | +-> character + | | bytes_ (kind=S) + | | str_ (kind=U) + | | + | \\-> void (kind=V) + \\-> object_ (not used much) (kind=O) + +""" +import numbers +import warnings + +from . import multiarray as ma +from .multiarray import ( + ndarray, dtype, datetime_data, datetime_as_string, + busday_offset, busday_count, is_busday, busdaycalendar + ) +from .._utils import set_module + +# we add more at the bottom +__all__ = [ + 'ScalarType', 'typecodes', 'issubdtype', 'datetime_data', + 'datetime_as_string', 'busday_offset', 'busday_count', + 'is_busday', 'busdaycalendar', 'isdtype' +] + +# we don't need all these imports, but we need to keep them for compatibility +# for users using np._core.numerictypes.UPPER_TABLE +from ._string_helpers import ( # noqa: F401 + english_lower, english_upper, english_capitalize, LOWER_TABLE, UPPER_TABLE +) + +from ._type_aliases import ( + sctypeDict, allTypes, sctypes +) +from ._dtype import _kind_name + +# we don't export these for import *, but we do want them accessible +# as numerictypes.bool, etc. +from builtins import bool, int, float, complex, object, str, bytes # noqa: F401, UP029 + + +# We use this later +generic = allTypes['generic'] + +genericTypeRank = ['bool', 'int8', 'uint8', 'int16', 'uint16', + 'int32', 'uint32', 'int64', 'uint64', 'int128', + 'uint128', 'float16', + 'float32', 'float64', 'float80', 'float96', 'float128', + 'float256', + 'complex32', 'complex64', 'complex128', 'complex160', + 'complex192', 'complex256', 'complex512', 'object'] + +@set_module('numpy') +def maximum_sctype(t): + """ + Return the scalar type of highest precision of the same kind as the input. + + .. deprecated:: 2.0 + Use an explicit dtype like int64 or float64 instead. + + Parameters + ---------- + t : dtype or dtype specifier + The input data type. This can be a `dtype` object or an object that + is convertible to a `dtype`. + + Returns + ------- + out : dtype + The highest precision data type of the same kind (`dtype.kind`) as `t`. + + See Also + -------- + obj2sctype, mintypecode, sctype2char + dtype + + Examples + -------- + >>> from numpy._core.numerictypes import maximum_sctype + >>> maximum_sctype(int) + + >>> maximum_sctype(np.uint8) + + >>> maximum_sctype(complex) + # may vary + + >>> maximum_sctype(str) + + + >>> maximum_sctype('i2') + + >>> maximum_sctype('f4') + # may vary + + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`maximum_sctype` is deprecated. Use an explicit dtype like int64 " + "or float64 instead. (deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + g = obj2sctype(t) + if g is None: + return t + t = g + base = _kind_name(dtype(t)) + if base in sctypes: + return sctypes[base][-1] + else: + return t + + +@set_module('numpy') +def issctype(rep): + """ + Determines whether the given object represents a scalar data-type. + + Parameters + ---------- + rep : any + If `rep` is an instance of a scalar dtype, True is returned. If not, + False is returned. + + Returns + ------- + out : bool + Boolean result of check whether `rep` is a scalar dtype. + + See Also + -------- + issubsctype, issubdtype, obj2sctype, sctype2char + + Examples + -------- + >>> from numpy._core.numerictypes import issctype + >>> issctype(np.int32) + True + >>> issctype(list) + False + >>> issctype(1.1) + False + + Strings are also a scalar type: + + >>> issctype(np.dtype('str')) + True + + """ + if not isinstance(rep, (type, dtype)): + return False + try: + res = obj2sctype(rep) + if res and res != object_: + return True + else: + return False + except Exception: + return False + + +@set_module('numpy') +def obj2sctype(rep, default=None): + """ + Return the scalar dtype or NumPy equivalent of Python type of an object. + + Parameters + ---------- + rep : any + The object of which the type is returned. + default : any, optional + If given, this is returned for objects whose types can not be + determined. If not given, None is returned for those objects. + + Returns + ------- + dtype : dtype or Python type + The data type of `rep`. + + See Also + -------- + sctype2char, issctype, issubsctype, issubdtype + + Examples + -------- + >>> from numpy._core.numerictypes import obj2sctype + >>> obj2sctype(np.int32) + + >>> obj2sctype(np.array([1., 2.])) + + >>> obj2sctype(np.array([1.j])) + + + >>> obj2sctype(dict) + + >>> obj2sctype('string') + + >>> obj2sctype(1, default=list) + + + """ + # prevent abstract classes being upcast + if isinstance(rep, type) and issubclass(rep, generic): + return rep + # extract dtype from arrays + if isinstance(rep, ndarray): + return rep.dtype.type + # fall back on dtype to convert + try: + res = dtype(rep) + except Exception: + return default + else: + return res.type + + +@set_module('numpy') +def issubclass_(arg1, arg2): + """ + Determine if a class is a subclass of a second class. + + `issubclass_` is equivalent to the Python built-in ``issubclass``, + except that it returns False instead of raising a TypeError if one + of the arguments is not a class. + + Parameters + ---------- + arg1 : class + Input class. True is returned if `arg1` is a subclass of `arg2`. + arg2 : class or tuple of classes. + Input class. If a tuple of classes, True is returned if `arg1` is a + subclass of any of the tuple elements. + + Returns + ------- + out : bool + Whether `arg1` is a subclass of `arg2` or not. + + See Also + -------- + issubsctype, issubdtype, issctype + + Examples + -------- + >>> np.issubclass_(np.int32, int) + False + >>> np.issubclass_(np.int32, float) + False + >>> np.issubclass_(np.float64, float) + True + + """ + try: + return issubclass(arg1, arg2) + except TypeError: + return False + + +@set_module('numpy') +def issubsctype(arg1, arg2): + """ + Determine if the first argument is a subclass of the second argument. + + Parameters + ---------- + arg1, arg2 : dtype or dtype specifier + Data-types. + + Returns + ------- + out : bool + The result. + + See Also + -------- + issctype, issubdtype, obj2sctype + + Examples + -------- + >>> from numpy._core import issubsctype + >>> issubsctype('S8', str) + False + >>> issubsctype(np.array([1]), int) + True + >>> issubsctype(np.array([1]), float) + False + + """ + return issubclass(obj2sctype(arg1), obj2sctype(arg2)) + + +class _PreprocessDTypeError(Exception): + pass + + +def _preprocess_dtype(dtype): + """ + Preprocess dtype argument by: + 1. fetching type from a data type + 2. verifying that types are built-in NumPy dtypes + """ + if isinstance(dtype, ma.dtype): + dtype = dtype.type + if isinstance(dtype, ndarray) or dtype not in allTypes.values(): + raise _PreprocessDTypeError + return dtype + + +@set_module('numpy') +def isdtype(dtype, kind): + """ + Determine if a provided dtype is of a specified data type ``kind``. + + This function only supports built-in NumPy's data types. + Third-party dtypes are not yet supported. + + Parameters + ---------- + dtype : dtype + The input dtype. + kind : dtype or str or tuple of dtypes/strs. + dtype or dtype kind. Allowed dtype kinds are: + * ``'bool'`` : boolean kind + * ``'signed integer'`` : signed integer data types + * ``'unsigned integer'`` : unsigned integer data types + * ``'integral'`` : integer data types + * ``'real floating'`` : real-valued floating-point data types + * ``'complex floating'`` : complex floating-point data types + * ``'numeric'`` : numeric data types + + Returns + ------- + out : bool + + See Also + -------- + issubdtype + + Examples + -------- + >>> import numpy as np + >>> np.isdtype(np.float32, np.float64) + False + >>> np.isdtype(np.float32, "real floating") + True + >>> np.isdtype(np.complex128, ("real floating", "complex floating")) + True + + """ + try: + dtype = _preprocess_dtype(dtype) + except _PreprocessDTypeError: + raise TypeError( + "dtype argument must be a NumPy dtype, " + f"but it is a {type(dtype)}." + ) from None + + input_kinds = kind if isinstance(kind, tuple) else (kind,) + + processed_kinds = set() + + for kind in input_kinds: + if kind == "bool": + processed_kinds.add(allTypes["bool"]) + elif kind == "signed integer": + processed_kinds.update(sctypes["int"]) + elif kind == "unsigned integer": + processed_kinds.update(sctypes["uint"]) + elif kind == "integral": + processed_kinds.update(sctypes["int"] + sctypes["uint"]) + elif kind == "real floating": + processed_kinds.update(sctypes["float"]) + elif kind == "complex floating": + processed_kinds.update(sctypes["complex"]) + elif kind == "numeric": + processed_kinds.update( + sctypes["int"] + sctypes["uint"] + + sctypes["float"] + sctypes["complex"] + ) + elif isinstance(kind, str): + raise ValueError( + "kind argument is a string, but" + f" {kind!r} is not a known kind name." + ) + else: + try: + kind = _preprocess_dtype(kind) + except _PreprocessDTypeError: + raise TypeError( + "kind argument must be comprised of " + "NumPy dtypes or strings only, " + f"but is a {type(kind)}." + ) from None + processed_kinds.add(kind) + + return dtype in processed_kinds + + +@set_module('numpy') +def issubdtype(arg1, arg2): + r""" + Returns True if first argument is a typecode lower/equal in type hierarchy. + + This is like the builtin :func:`issubclass`, but for `dtype`\ s. + + Parameters + ---------- + arg1, arg2 : dtype_like + `dtype` or object coercible to one + + Returns + ------- + out : bool + + See Also + -------- + :ref:`arrays.scalars` : Overview of the numpy type hierarchy. + + Examples + -------- + `issubdtype` can be used to check the type of arrays: + + >>> ints = np.array([1, 2, 3], dtype=np.int32) + >>> np.issubdtype(ints.dtype, np.integer) + True + >>> np.issubdtype(ints.dtype, np.floating) + False + + >>> floats = np.array([1, 2, 3], dtype=np.float32) + >>> np.issubdtype(floats.dtype, np.integer) + False + >>> np.issubdtype(floats.dtype, np.floating) + True + + Similar types of different sizes are not subdtypes of each other: + + >>> np.issubdtype(np.float64, np.float32) + False + >>> np.issubdtype(np.float32, np.float64) + False + + but both are subtypes of `floating`: + + >>> np.issubdtype(np.float64, np.floating) + True + >>> np.issubdtype(np.float32, np.floating) + True + + For convenience, dtype-like objects are allowed too: + + >>> np.issubdtype('S1', np.bytes_) + True + >>> np.issubdtype('i4', np.signedinteger) + True + + """ + if not issubclass_(arg1, generic): + arg1 = dtype(arg1).type + if not issubclass_(arg2, generic): + arg2 = dtype(arg2).type + + return issubclass(arg1, arg2) + + +@set_module('numpy') +def sctype2char(sctype): + """ + Return the string representation of a scalar dtype. + + Parameters + ---------- + sctype : scalar dtype or object + If a scalar dtype, the corresponding string character is + returned. If an object, `sctype2char` tries to infer its scalar type + and then return the corresponding string character. + + Returns + ------- + typechar : str + The string character corresponding to the scalar type. + + Raises + ------ + ValueError + If `sctype` is an object for which the type can not be inferred. + + See Also + -------- + obj2sctype, issctype, issubsctype, mintypecode + + Examples + -------- + >>> from numpy._core.numerictypes import sctype2char + >>> for sctype in [np.int32, np.double, np.cdouble, np.bytes_, np.ndarray]: + ... print(sctype2char(sctype)) + l # may vary + d + D + S + O + + >>> x = np.array([1., 2-1.j]) + >>> sctype2char(x) + 'D' + >>> sctype2char(list) + 'O' + + """ + sctype = obj2sctype(sctype) + if sctype is None: + raise ValueError("unrecognized type") + if sctype not in sctypeDict.values(): + # for compatibility + raise KeyError(sctype) + return dtype(sctype).char + + +def _scalar_type_key(typ): + """A ``key`` function for `sorted`.""" + dt = dtype(typ) + return (dt.kind.lower(), dt.itemsize) + + +ScalarType = [int, float, complex, bool, bytes, str, memoryview] +ScalarType += sorted(set(sctypeDict.values()), key=_scalar_type_key) +ScalarType = tuple(ScalarType) + + +# Now add the types we've determined to this module +for key in allTypes: + globals()[key] = allTypes[key] + __all__.append(key) + +del key + +typecodes = {'Character': 'c', + 'Integer': 'bhilqnp', + 'UnsignedInteger': 'BHILQNP', + 'Float': 'efdg', + 'Complex': 'FDG', + 'AllInteger': 'bBhHiIlLqQnNpP', + 'AllFloat': 'efdgFDG', + 'Datetime': 'Mm', + 'All': '?bhilqnpBHILQNPefdgFDGSUVOMm'} + +# backwards compatibility --- deprecated name +# Formal deprecation: Numpy 1.20.0, 2020-10-19 (see numpy/__init__.py) +typeDict = sctypeDict + +def _register_types(): + numbers.Integral.register(integer) + numbers.Complex.register(inexact) + numbers.Real.register(floating) + numbers.Number.register(number) + + +_register_types() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numerictypes.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numerictypes.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6a697cfe2d8a70ed09fcaf4802bfc3b2a63bf582 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/numerictypes.pyi @@ -0,0 +1,217 @@ +import builtins +from typing import ( + Any, + Literal as L, + TypedDict, + type_check_only, +) + +import numpy as np +from numpy import ( + dtype, + generic, + bool, + bool_, + uint8, + uint16, + uint32, + uint64, + ubyte, + ushort, + uintc, + ulong, + ulonglong, + uintp, + uint, + int8, + int16, + int32, + int64, + byte, + short, + intc, + long, + longlong, + intp, + int_, + float16, + float32, + float64, + half, + single, + double, + longdouble, + complex64, + complex128, + csingle, + cdouble, + clongdouble, + datetime64, + timedelta64, + object_, + str_, + bytes_, + void, + unsignedinteger, + character, + inexact, + number, + integer, + flexible, + complexfloating, + signedinteger, + floating, +) +from ._type_aliases import sctypeDict # noqa: F401 +from .multiarray import ( + busday_count, + busday_offset, + busdaycalendar, + datetime_as_string, + datetime_data, + is_busday, +) + +from numpy._typing import DTypeLike +from numpy._typing._extended_precision import ( + uint128, + uint256, + int128, + int256, + float80, + float96, + float128, + float256, + complex160, + complex192, + complex256, + complex512, +) + +__all__ = [ + "ScalarType", + "typecodes", + "issubdtype", + "datetime_data", + "datetime_as_string", + "busday_offset", + "busday_count", + "is_busday", + "busdaycalendar", + "isdtype", + "generic", + "unsignedinteger", + "character", + "inexact", + "number", + "integer", + "flexible", + "complexfloating", + "signedinteger", + "floating", + "bool", + "float16", + "float32", + "float64", + "longdouble", + "complex64", + "complex128", + "clongdouble", + "bytes_", + "str_", + "void", + "object_", + "datetime64", + "timedelta64", + "int8", + "byte", + "uint8", + "ubyte", + "int16", + "short", + "uint16", + "ushort", + "int32", + "intc", + "uint32", + "uintc", + "int64", + "long", + "uint64", + "ulong", + "longlong", + "ulonglong", + "intp", + "uintp", + "double", + "cdouble", + "single", + "csingle", + "half", + "bool_", + "int_", + "uint", + "uint128", + "uint256", + "int128", + "int256", + "float80", + "float96", + "float128", + "float256", + "complex160", + "complex192", + "complex256", + "complex512", +] + +@type_check_only +class _TypeCodes(TypedDict): + Character: L['c'] + Integer: L['bhilqnp'] + UnsignedInteger: L['BHILQNP'] + Float: L['efdg'] + Complex: L['FDG'] + AllInteger: L['bBhHiIlLqQnNpP'] + AllFloat: L['efdgFDG'] + Datetime: L['Mm'] + All: L['?bhilqnpBHILQNPefdgFDGSUVOMm'] + +def isdtype(dtype: dtype[Any] | type[Any], kind: DTypeLike | tuple[DTypeLike, ...]) -> builtins.bool: ... + +def issubdtype(arg1: DTypeLike, arg2: DTypeLike) -> builtins.bool: ... + +typecodes: _TypeCodes +ScalarType: tuple[ + type[int], + type[float], + type[complex], + type[builtins.bool], + type[bytes], + type[str], + type[memoryview], + type[np.bool], + type[csingle], + type[cdouble], + type[clongdouble], + type[half], + type[single], + type[double], + type[longdouble], + type[byte], + type[short], + type[intc], + type[long], + type[longlong], + type[timedelta64], + type[datetime64], + type[object_], + type[bytes_], + type[str_], + type[ubyte], + type[ushort], + type[uintc], + type[ulong], + type[ulonglong], + type[void], +] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/overrides.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/overrides.py new file mode 100644 index 0000000000000000000000000000000000000000..1f5039ba083ef9d3fe90912d052a7a0d17742652 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/overrides.py @@ -0,0 +1,181 @@ +"""Implementation of __array_function__ overrides from NEP-18.""" +import collections +import functools + +from .._utils import set_module +from .._utils._inspect import getargspec +from numpy._core._multiarray_umath import ( + add_docstring, _get_implementing_args, _ArrayFunctionDispatcher) + + +ARRAY_FUNCTIONS = set() + +array_function_like_doc = ( + """like : array_like, optional + Reference object to allow the creation of arrays which are not + NumPy arrays. If an array-like passed in as ``like`` supports + the ``__array_function__`` protocol, the result will be defined + by it. In this case, it ensures the creation of an array object + compatible with that passed in via this argument.""" +) + +def get_array_function_like_doc(public_api, docstring_template=""): + ARRAY_FUNCTIONS.add(public_api) + docstring = public_api.__doc__ or docstring_template + return docstring.replace("${ARRAY_FUNCTION_LIKE}", array_function_like_doc) + +def finalize_array_function_like(public_api): + public_api.__doc__ = get_array_function_like_doc(public_api) + return public_api + + +add_docstring( + _ArrayFunctionDispatcher, + """ + Class to wrap functions with checks for __array_function__ overrides. + + All arguments are required, and can only be passed by position. + + Parameters + ---------- + dispatcher : function or None + The dispatcher function that returns a single sequence-like object + of all arguments relevant. It must have the same signature (except + the default values) as the actual implementation. + If ``None``, this is a ``like=`` dispatcher and the + ``_ArrayFunctionDispatcher`` must be called with ``like`` as the + first (additional and positional) argument. + implementation : function + Function that implements the operation on NumPy arrays without + overrides. Arguments passed calling the ``_ArrayFunctionDispatcher`` + will be forwarded to this (and the ``dispatcher``) as if using + ``*args, **kwargs``. + + Attributes + ---------- + _implementation : function + The original implementation passed in. + """) + + +# exposed for testing purposes; used internally by _ArrayFunctionDispatcher +add_docstring( + _get_implementing_args, + """ + Collect arguments on which to call __array_function__. + + Parameters + ---------- + relevant_args : iterable of array-like + Iterable of possibly array-like arguments to check for + __array_function__ methods. + + Returns + ------- + Sequence of arguments with __array_function__ methods, in the order in + which they should be called. + """) + + +ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults') + + +def verify_matching_signatures(implementation, dispatcher): + """Verify that a dispatcher function has the right signature.""" + implementation_spec = ArgSpec(*getargspec(implementation)) + dispatcher_spec = ArgSpec(*getargspec(dispatcher)) + + if (implementation_spec.args != dispatcher_spec.args or + implementation_spec.varargs != dispatcher_spec.varargs or + implementation_spec.keywords != dispatcher_spec.keywords or + (bool(implementation_spec.defaults) != + bool(dispatcher_spec.defaults)) or + (implementation_spec.defaults is not None and + len(implementation_spec.defaults) != + len(dispatcher_spec.defaults))): + raise RuntimeError('implementation and dispatcher for %s have ' + 'different function signatures' % implementation) + + if implementation_spec.defaults is not None: + if dispatcher_spec.defaults != (None,) * len(dispatcher_spec.defaults): + raise RuntimeError('dispatcher functions can only use None for ' + 'default argument values') + + +def array_function_dispatch(dispatcher=None, module=None, verify=True, + docs_from_dispatcher=False): + """Decorator for adding dispatch with the __array_function__ protocol. + + See NEP-18 for example usage. + + Parameters + ---------- + dispatcher : callable or None + Function that when called like ``dispatcher(*args, **kwargs)`` with + arguments from the NumPy function call returns an iterable of + array-like arguments to check for ``__array_function__``. + + If `None`, the first argument is used as the single `like=` argument + and not passed on. A function implementing `like=` must call its + dispatcher with `like` as the first non-keyword argument. + module : str, optional + __module__ attribute to set on new function, e.g., ``module='numpy'``. + By default, module is copied from the decorated function. + verify : bool, optional + If True, verify the that the signature of the dispatcher and decorated + function signatures match exactly: all required and optional arguments + should appear in order with the same names, but the default values for + all optional arguments should be ``None``. Only disable verification + if the dispatcher's signature needs to deviate for some particular + reason, e.g., because the function has a signature like + ``func(*args, **kwargs)``. + docs_from_dispatcher : bool, optional + If True, copy docs from the dispatcher function onto the dispatched + function, rather than from the implementation. This is useful for + functions defined in C, which otherwise don't have docstrings. + + Returns + ------- + Function suitable for decorating the implementation of a NumPy function. + + """ + def decorator(implementation): + if verify: + if dispatcher is not None: + verify_matching_signatures(implementation, dispatcher) + else: + # Using __code__ directly similar to verify_matching_signature + co = implementation.__code__ + last_arg = co.co_argcount + co.co_kwonlyargcount - 1 + last_arg = co.co_varnames[last_arg] + if last_arg != "like" or co.co_kwonlyargcount == 0: + raise RuntimeError( + "__array_function__ expects `like=` to be the last " + "argument and a keyword-only argument. " + f"{implementation} does not seem to comply.") + + if docs_from_dispatcher: + add_docstring(implementation, dispatcher.__doc__) + + public_api = _ArrayFunctionDispatcher(dispatcher, implementation) + public_api = functools.wraps(implementation)(public_api) + + if module is not None: + public_api.__module__ = module + + ARRAY_FUNCTIONS.add(public_api) + + return public_api + + return decorator + + +def array_function_from_dispatcher( + implementation, module=None, verify=True, docs_from_dispatcher=True): + """Like array_function_dispatcher, but with function arguments flipped.""" + + def decorator(dispatcher): + return array_function_dispatch( + dispatcher, module, verify=verify, + docs_from_dispatcher=docs_from_dispatcher)(implementation) + return decorator diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/overrides.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/overrides.pyi new file mode 100644 index 0000000000000000000000000000000000000000..35e4ed36e70fc2ab3760467a1f94f018961fa29d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/overrides.pyi @@ -0,0 +1,50 @@ +from collections.abc import Callable, Iterable +from typing import Any, Final, NamedTuple + +from typing_extensions import ParamSpec, TypeVar + +from numpy._typing import _SupportsArrayFunc + +_T = TypeVar("_T") +_Tss = ParamSpec("_Tss") +_FuncT = TypeVar("_FuncT", bound=Callable[..., object]) + +### + +ARRAY_FUNCTIONS: set[Callable[..., Any]] = ... +array_function_like_doc: Final[str] = ... + +class ArgSpec(NamedTuple): + args: list[str] + varargs: str | None + keywords: str | None + defaults: tuple[Any, ...] + +def get_array_function_like_doc(public_api: Callable[..., Any], docstring_template: str = "") -> str: ... +def finalize_array_function_like(public_api: _FuncT) -> _FuncT: ... + +# +def verify_matching_signatures( + implementation: Callable[_Tss, object], + dispatcher: Callable[_Tss, Iterable[_SupportsArrayFunc]], +) -> None: ... + +# NOTE: This actually returns a `_ArrayFunctionDispatcher` callable wrapper object, with +# the original wrapped callable stored in the `._implementation` attribute. It checks +# for any `__array_function__` of the values of specific arguments that the dispatcher +# specifies. Since the dispatcher only returns an iterable of passed array-like args, +# this overridable behaviour is impossible to annotate. +def array_function_dispatch( + dispatcher: Callable[_Tss, Iterable[_SupportsArrayFunc]] | None = None, + module: str | None = None, + verify: bool = True, + docs_from_dispatcher: bool = False, +) -> Callable[[_FuncT], _FuncT]: ... + +# +def array_function_from_dispatcher( + implementation: Callable[_Tss, _T], + module: str | None = None, + verify: bool = True, + docs_from_dispatcher: bool = True, +) -> Callable[[Callable[_Tss, Iterable[_SupportsArrayFunc]]], Callable[_Tss, _T]]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/printoptions.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/printoptions.py new file mode 100644 index 0000000000000000000000000000000000000000..0e232d38f65d2c670817b870f409371489931ff6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/printoptions.py @@ -0,0 +1,32 @@ +""" +Stores and defines the low-level format_options context variable. + +This is defined in its own file outside of the arrayprint module +so we can import it from C while initializing the multiarray +C module during import without introducing circular dependencies. +""" + +import sys +from contextvars import ContextVar + +__all__ = ["format_options"] + +default_format_options_dict = { + "edgeitems": 3, # repr N leading and trailing items of each dimension + "threshold": 1000, # total items > triggers array summarization + "floatmode": "maxprec", + "precision": 8, # precision of floating point representations + "suppress": False, # suppress printing small floating values in exp format + "linewidth": 75, + "nanstr": "nan", + "infstr": "inf", + "sign": "-", + "formatter": None, + # Internally stored as an int to simplify comparisons; converted from/to + # str/False on the way in/out. + 'legacy': sys.maxsize, + 'override_repr': None, +} + +format_options = ContextVar( + "format_options", default=default_format_options_dict.copy()) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/printoptions.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/printoptions.pyi new file mode 100644 index 0000000000000000000000000000000000000000..1863f41bb374aa9ba3e45e34068253da58d9b91a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/printoptions.pyi @@ -0,0 +1,28 @@ +from collections.abc import Callable +from contextvars import ContextVar +from typing import Any, Final, TypedDict + +from .arrayprint import _FormatDict + +__all__ = ["format_options"] + +### + +class _FormatOptionsDict(TypedDict): + edgeitems: int + threshold: int + floatmode: str + precision: int + suppress: bool + linewidth: int + nanstr: str + infstr: str + sign: str + formatter: _FormatDict | None + legacy: int + override_repr: Callable[[Any], str] | None + +### + +default_format_options_dict: Final[_FormatOptionsDict] = ... +format_options: ContextVar[_FormatOptionsDict] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/records.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/records.py new file mode 100644 index 0000000000000000000000000000000000000000..acb9584dc38898ede605afe10330e6f86cc86a1b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/records.py @@ -0,0 +1,1091 @@ +""" +This module contains a set of functions for record arrays. +""" +import os +import warnings +from collections import Counter +from contextlib import nullcontext + +from .._utils import set_module +from . import numeric as sb +from . import numerictypes as nt +from .arrayprint import _get_legacy_print_mode + +# All of the functions allow formats to be a dtype +__all__ = [ + 'record', 'recarray', 'format_parser', 'fromarrays', 'fromrecords', + 'fromstring', 'fromfile', 'array', 'find_duplicate', +] + + +ndarray = sb.ndarray + +_byteorderconv = {'b': '>', + 'l': '<', + 'n': '=', + 'B': '>', + 'L': '<', + 'N': '=', + 'S': 's', + 's': 's', + '>': '>', + '<': '<', + '=': '=', + '|': '|', + 'I': '|', + 'i': '|'} + +# formats regular expression +# allows multidimensional spec with a tuple syntax in front +# of the letter code '(2,3)f4' and ' ( 2 , 3 ) f4 ' +# are equally allowed + +numfmt = nt.sctypeDict + + +@set_module('numpy.rec') +def find_duplicate(list): + """Find duplication in a list, return a list of duplicated elements""" + return [ + item + for item, counts in Counter(list).items() + if counts > 1 + ] + + +@set_module('numpy.rec') +class format_parser: + """ + Class to convert formats, names, titles description to a dtype. + + After constructing the format_parser object, the dtype attribute is + the converted data-type: + ``dtype = format_parser(formats, names, titles).dtype`` + + Attributes + ---------- + dtype : dtype + The converted data-type. + + Parameters + ---------- + formats : str or list of str + The format description, either specified as a string with + comma-separated format descriptions in the form ``'f8, i4, S5'``, or + a list of format description strings in the form + ``['f8', 'i4', 'S5']``. + names : str or list/tuple of str + The field names, either specified as a comma-separated string in the + form ``'col1, col2, col3'``, or as a list or tuple of strings in the + form ``['col1', 'col2', 'col3']``. + An empty list can be used, in that case default field names + ('f0', 'f1', ...) are used. + titles : sequence + Sequence of title strings. An empty list can be used to leave titles + out. + aligned : bool, optional + If True, align the fields by padding as the C-compiler would. + Default is False. + byteorder : str, optional + If specified, all the fields will be changed to the + provided byte-order. Otherwise, the default byte-order is + used. For all available string specifiers, see `dtype.newbyteorder`. + + See Also + -------- + numpy.dtype, numpy.typename + + Examples + -------- + >>> import numpy as np + >>> np.rec.format_parser(['>> np.rec.format_parser(['f8', 'i4', 'a5'], ['col1', 'col2', 'col3'], + ... []).dtype + dtype([('col1', '>> np.rec.format_parser([' len(titles): + self._titles += [None] * (self._nfields - len(titles)) + + def _createdtype(self, byteorder): + dtype = sb.dtype({ + 'names': self._names, + 'formats': self._f_formats, + 'offsets': self._offsets, + 'titles': self._titles, + }) + if byteorder is not None: + byteorder = _byteorderconv[byteorder[0]] + dtype = dtype.newbyteorder(byteorder) + + self.dtype = dtype + + +class record(nt.void): + """A data-type scalar that allows field access as attribute lookup. + """ + + # manually set name and module so that this class's type shows up + # as numpy.record when printed + __name__ = 'record' + __module__ = 'numpy' + + def __repr__(self): + if _get_legacy_print_mode() <= 113: + return self.__str__() + return super().__repr__() + + def __str__(self): + if _get_legacy_print_mode() <= 113: + return str(self.item()) + return super().__str__() + + def __getattribute__(self, attr): + if attr in ('setfield', 'getfield', 'dtype'): + return nt.void.__getattribute__(self, attr) + try: + return nt.void.__getattribute__(self, attr) + except AttributeError: + pass + fielddict = nt.void.__getattribute__(self, 'dtype').fields + res = fielddict.get(attr, None) + if res: + obj = self.getfield(*res[:2]) + # if it has fields return a record, + # otherwise return the object + try: + dt = obj.dtype + except AttributeError: + #happens if field is Object type + return obj + if dt.names is not None: + return obj.view((self.__class__, obj.dtype)) + return obj + else: + raise AttributeError("'record' object has no " + "attribute '%s'" % attr) + + def __setattr__(self, attr, val): + if attr in ('setfield', 'getfield', 'dtype'): + raise AttributeError("Cannot set '%s' attribute" % attr) + fielddict = nt.void.__getattribute__(self, 'dtype').fields + res = fielddict.get(attr, None) + if res: + return self.setfield(val, *res[:2]) + else: + if getattr(self, attr, None): + return nt.void.__setattr__(self, attr, val) + else: + raise AttributeError("'record' object has no " + "attribute '%s'" % attr) + + def __getitem__(self, indx): + obj = nt.void.__getitem__(self, indx) + + # copy behavior of record.__getattribute__, + if isinstance(obj, nt.void) and obj.dtype.names is not None: + return obj.view((self.__class__, obj.dtype)) + else: + # return a single element + return obj + + def pprint(self): + """Pretty-print all fields.""" + # pretty-print all fields + names = self.dtype.names + maxlen = max(len(name) for name in names) + fmt = '%% %ds: %%s' % maxlen + rows = [fmt % (name, getattr(self, name)) for name in names] + return "\n".join(rows) + +# The recarray is almost identical to a standard array (which supports +# named fields already) The biggest difference is that it can use +# attribute-lookup to find the fields and it is constructed using +# a record. + +# If byteorder is given it forces a particular byteorder on all +# the fields (and any subfields) + + +@set_module("numpy.rec") +class recarray(ndarray): + """Construct an ndarray that allows field access using attributes. + + Arrays may have a data-types containing fields, analogous + to columns in a spread sheet. An example is ``[(x, int), (y, float)]``, + where each entry in the array is a pair of ``(int, float)``. Normally, + these attributes are accessed using dictionary lookups such as ``arr['x']`` + and ``arr['y']``. Record arrays allow the fields to be accessed as members + of the array, using ``arr.x`` and ``arr.y``. + + Parameters + ---------- + shape : tuple + Shape of output array. + dtype : data-type, optional + The desired data-type. By default, the data-type is determined + from `formats`, `names`, `titles`, `aligned` and `byteorder`. + formats : list of data-types, optional + A list containing the data-types for the different columns, e.g. + ``['i4', 'f8', 'i4']``. `formats` does *not* support the new + convention of using types directly, i.e. ``(int, float, int)``. + Note that `formats` must be a list, not a tuple. + Given that `formats` is somewhat limited, we recommend specifying + `dtype` instead. + names : tuple of str, optional + The name of each column, e.g. ``('x', 'y', 'z')``. + buf : buffer, optional + By default, a new array is created of the given shape and data-type. + If `buf` is specified and is an object exposing the buffer interface, + the array will use the memory from the existing buffer. In this case, + the `offset` and `strides` keywords are available. + + Other Parameters + ---------------- + titles : tuple of str, optional + Aliases for column names. For example, if `names` were + ``('x', 'y', 'z')`` and `titles` is + ``('x_coordinate', 'y_coordinate', 'z_coordinate')``, then + ``arr['x']`` is equivalent to both ``arr.x`` and ``arr.x_coordinate``. + byteorder : {'<', '>', '='}, optional + Byte-order for all fields. + aligned : bool, optional + Align the fields in memory as the C-compiler would. + strides : tuple of ints, optional + Buffer (`buf`) is interpreted according to these strides (strides + define how many bytes each array element, row, column, etc. + occupy in memory). + offset : int, optional + Start reading buffer (`buf`) from this offset onwards. + order : {'C', 'F'}, optional + Row-major (C-style) or column-major (Fortran-style) order. + + Returns + ------- + rec : recarray + Empty array of the given shape and type. + + See Also + -------- + numpy.rec.fromrecords : Construct a record array from data. + numpy.record : fundamental data-type for `recarray`. + numpy.rec.format_parser : determine data-type from formats, names, titles. + + Notes + ----- + This constructor can be compared to ``empty``: it creates a new record + array but does not fill it with data. To create a record array from data, + use one of the following methods: + + 1. Create a standard ndarray and convert it to a record array, + using ``arr.view(np.recarray)`` + 2. Use the `buf` keyword. + 3. Use `np.rec.fromrecords`. + + Examples + -------- + Create an array with two fields, ``x`` and ``y``: + + >>> import numpy as np + >>> x = np.array([(1.0, 2), (3.0, 4)], dtype=[('x', '>> x + array([(1., 2), (3., 4)], dtype=[('x', '>> x['x'] + array([1., 3.]) + + View the array as a record array: + + >>> x = x.view(np.recarray) + + >>> x.x + array([1., 3.]) + + >>> x.y + array([2, 4]) + + Create a new, empty record array: + + >>> np.recarray((2,), + ... dtype=[('x', int), ('y', float), ('z', int)]) #doctest: +SKIP + rec.array([(-1073741821, 1.2249118382103472e-301, 24547520), + (3471280, 1.2134086255804012e-316, 0)], + dtype=[('x', ' 0 or self.shape == (0,): + lst = sb.array2string( + self, separator=', ', prefix=prefix, suffix=',') + else: + # show zero-length shape unless it is (0,) + lst = "[], shape=%s" % (repr(self.shape),) + + lf = '\n'+' '*len(prefix) + if _get_legacy_print_mode() <= 113: + lf = ' ' + lf # trailing space + return fmt % (lst, lf, repr_dtype) + + def field(self, attr, val=None): + if isinstance(attr, int): + names = ndarray.__getattribute__(self, 'dtype').names + attr = names[attr] + + fielddict = ndarray.__getattribute__(self, 'dtype').fields + + res = fielddict[attr][:2] + + if val is None: + obj = self.getfield(*res) + if obj.dtype.names is not None: + return obj + return obj.view(ndarray) + else: + return self.setfield(val, *res) + + +def _deprecate_shape_0_as_None(shape): + if shape == 0: + warnings.warn( + "Passing `shape=0` to have the shape be inferred is deprecated, " + "and in future will be equivalent to `shape=(0,)`. To infer " + "the shape and suppress this warning, pass `shape=None` instead.", + FutureWarning, stacklevel=3) + return None + else: + return shape + + +@set_module("numpy.rec") +def fromarrays(arrayList, dtype=None, shape=None, formats=None, + names=None, titles=None, aligned=False, byteorder=None): + """Create a record array from a (flat) list of arrays + + Parameters + ---------- + arrayList : list or tuple + List of array-like objects (such as lists, tuples, + and ndarrays). + dtype : data-type, optional + valid dtype for all arrays + shape : int or tuple of ints, optional + Shape of the resulting array. If not provided, inferred from + ``arrayList[0]``. + formats, names, titles, aligned, byteorder : + If `dtype` is ``None``, these arguments are passed to + `numpy.rec.format_parser` to construct a dtype. See that function for + detailed documentation. + + Returns + ------- + np.recarray + Record array consisting of given arrayList columns. + + Examples + -------- + >>> x1=np.array([1,2,3,4]) + >>> x2=np.array(['a','dd','xyz','12']) + >>> x3=np.array([1.1,2,3,4]) + >>> r = np.rec.fromarrays([x1,x2,x3],names='a,b,c') + >>> print(r[1]) + (2, 'dd', 2.0) # may vary + >>> x1[1]=34 + >>> r.a + array([1, 2, 3, 4]) + + >>> x1 = np.array([1, 2, 3, 4]) + >>> x2 = np.array(['a', 'dd', 'xyz', '12']) + >>> x3 = np.array([1.1, 2, 3,4]) + >>> r = np.rec.fromarrays( + ... [x1, x2, x3], + ... dtype=np.dtype([('a', np.int32), ('b', 'S3'), ('c', np.float32)])) + >>> r + rec.array([(1, b'a', 1.1), (2, b'dd', 2. ), (3, b'xyz', 3. ), + (4, b'12', 4. )], + dtype=[('a', ' 0: + shape = shape[:-nn] + + _array = recarray(shape, descr) + + # populate the record array (makes a copy) + for k, obj in enumerate(arrayList): + nn = descr[k].ndim + testshape = obj.shape[:obj.ndim - nn] + name = _names[k] + if testshape != shape: + raise ValueError(f'array-shape mismatch in array {k} ("{name}")') + + _array[name] = obj + + return _array + + +@set_module("numpy.rec") +def fromrecords(recList, dtype=None, shape=None, formats=None, names=None, + titles=None, aligned=False, byteorder=None): + """Create a recarray from a list of records in text form. + + Parameters + ---------- + recList : sequence + data in the same field may be heterogeneous - they will be promoted + to the highest data type. + dtype : data-type, optional + valid dtype for all arrays + shape : int or tuple of ints, optional + shape of each array. + formats, names, titles, aligned, byteorder : + If `dtype` is ``None``, these arguments are passed to + `numpy.format_parser` to construct a dtype. See that function for + detailed documentation. + + If both `formats` and `dtype` are None, then this will auto-detect + formats. Use list of tuples rather than list of lists for faster + processing. + + Returns + ------- + np.recarray + record array consisting of given recList rows. + + Examples + -------- + >>> r=np.rec.fromrecords([(456,'dbe',1.2),(2,'de',1.3)], + ... names='col1,col2,col3') + >>> print(r[0]) + (456, 'dbe', 1.2) + >>> r.col1 + array([456, 2]) + >>> r.col2 + array(['dbe', 'de'], dtype='>> import pickle + >>> pickle.loads(pickle.dumps(r)) + rec.array([(456, 'dbe', 1.2), ( 2, 'de', 1.3)], + dtype=[('col1', ' 1: + raise ValueError("Can only deal with 1-d array.") + _array = recarray(shape, descr) + for k in range(_array.size): + _array[k] = tuple(recList[k]) + # list of lists instead of list of tuples ? + # 2018-02-07, 1.14.1 + warnings.warn( + "fromrecords expected a list of tuples, may have received a list " + "of lists instead. In the future that will raise an error", + FutureWarning, stacklevel=2) + return _array + else: + if shape is not None and retval.shape != shape: + retval.shape = shape + + res = retval.view(recarray) + + return res + + +@set_module("numpy.rec") +def fromstring(datastring, dtype=None, shape=None, offset=0, formats=None, + names=None, titles=None, aligned=False, byteorder=None): + r"""Create a record array from binary data + + Note that despite the name of this function it does not accept `str` + instances. + + Parameters + ---------- + datastring : bytes-like + Buffer of binary data + dtype : data-type, optional + Valid dtype for all arrays + shape : int or tuple of ints, optional + Shape of each array. + offset : int, optional + Position in the buffer to start reading from. + formats, names, titles, aligned, byteorder : + If `dtype` is ``None``, these arguments are passed to + `numpy.format_parser` to construct a dtype. See that function for + detailed documentation. + + + Returns + ------- + np.recarray + Record array view into the data in datastring. This will be readonly + if `datastring` is readonly. + + See Also + -------- + numpy.frombuffer + + Examples + -------- + >>> a = b'\x01\x02\x03abc' + >>> np.rec.fromstring(a, dtype='u1,u1,u1,S3') + rec.array([(1, 2, 3, b'abc')], + dtype=[('f0', 'u1'), ('f1', 'u1'), ('f2', 'u1'), ('f3', 'S3')]) + + >>> grades_dtype = [('Name', (np.str_, 10)), ('Marks', np.float64), + ... ('GradeLevel', np.int32)] + >>> grades_array = np.array([('Sam', 33.3, 3), ('Mike', 44.4, 5), + ... ('Aadi', 66.6, 6)], dtype=grades_dtype) + >>> np.rec.fromstring(grades_array.tobytes(), dtype=grades_dtype) + rec.array([('Sam', 33.3, 3), ('Mike', 44.4, 5), ('Aadi', 66.6, 6)], + dtype=[('Name', '>> s = '\x01\x02\x03abc' + >>> np.rec.fromstring(s, dtype='u1,u1,u1,S3') + Traceback (most recent call last): + ... + TypeError: a bytes-like object is required, not 'str' + """ + + if dtype is None and formats is None: + raise TypeError("fromstring() needs a 'dtype' or 'formats' argument") + + if dtype is not None: + descr = sb.dtype(dtype) + else: + descr = format_parser(formats, names, titles, aligned, byteorder).dtype + + itemsize = descr.itemsize + + # NumPy 1.19.0, 2020-01-01 + shape = _deprecate_shape_0_as_None(shape) + + if shape in (None, -1): + shape = (len(datastring) - offset) // itemsize + + _array = recarray(shape, descr, buf=datastring, offset=offset) + return _array + +def get_remaining_size(fd): + pos = fd.tell() + try: + fd.seek(0, 2) + return fd.tell() - pos + finally: + fd.seek(pos, 0) + + +@set_module("numpy.rec") +def fromfile(fd, dtype=None, shape=None, offset=0, formats=None, + names=None, titles=None, aligned=False, byteorder=None): + """Create an array from binary file data + + Parameters + ---------- + fd : str or file type + If file is a string or a path-like object then that file is opened, + else it is assumed to be a file object. The file object must + support random access (i.e. it must have tell and seek methods). + dtype : data-type, optional + valid dtype for all arrays + shape : int or tuple of ints, optional + shape of each array. + offset : int, optional + Position in the file to start reading from. + formats, names, titles, aligned, byteorder : + If `dtype` is ``None``, these arguments are passed to + `numpy.format_parser` to construct a dtype. See that function for + detailed documentation + + Returns + ------- + np.recarray + record array consisting of data enclosed in file. + + Examples + -------- + >>> from tempfile import TemporaryFile + >>> a = np.empty(10,dtype='f8,i4,a5') + >>> a[5] = (0.5,10,'abcde') + >>> + >>> fd=TemporaryFile() + >>> a = a.view(a.dtype.newbyteorder('<')) + >>> a.tofile(fd) + >>> + >>> _ = fd.seek(0) + >>> r=np.rec.fromfile(fd, formats='f8,i4,a5', shape=10, + ... byteorder='<') + >>> print(r[5]) + (0.5, 10, b'abcde') + >>> r.shape + (10,) + """ + + if dtype is None and formats is None: + raise TypeError("fromfile() needs a 'dtype' or 'formats' argument") + + # NumPy 1.19.0, 2020-01-01 + shape = _deprecate_shape_0_as_None(shape) + + if shape is None: + shape = (-1,) + elif isinstance(shape, int): + shape = (shape,) + + if hasattr(fd, 'readinto'): + # GH issue 2504. fd supports io.RawIOBase or io.BufferedIOBase + # interface. Example of fd: gzip, BytesIO, BufferedReader + # file already opened + ctx = nullcontext(fd) + else: + # open file + ctx = open(os.fspath(fd), 'rb') + + with ctx as fd: + if offset > 0: + fd.seek(offset, 1) + size = get_remaining_size(fd) + + if dtype is not None: + descr = sb.dtype(dtype) + else: + descr = format_parser( + formats, names, titles, aligned, byteorder + ).dtype + + itemsize = descr.itemsize + + shapeprod = sb.array(shape).prod(dtype=nt.intp) + shapesize = shapeprod * itemsize + if shapesize < 0: + shape = list(shape) + shape[shape.index(-1)] = size // -shapesize + shape = tuple(shape) + shapeprod = sb.array(shape).prod(dtype=nt.intp) + + nbytes = shapeprod * itemsize + + if nbytes > size: + raise ValueError( + "Not enough bytes left in file for specified " + "shape and type." + ) + + # create the array + _array = recarray(shape, descr) + nbytesread = fd.readinto(_array.data) + if nbytesread != nbytes: + raise OSError("Didn't read as many bytes as expected") + + return _array + + +@set_module("numpy.rec") +def array(obj, dtype=None, shape=None, offset=0, strides=None, formats=None, + names=None, titles=None, aligned=False, byteorder=None, copy=True): + """ + Construct a record array from a wide-variety of objects. + + A general-purpose record array constructor that dispatches to the + appropriate `recarray` creation function based on the inputs (see Notes). + + Parameters + ---------- + obj : any + Input object. See Notes for details on how various input types are + treated. + dtype : data-type, optional + Valid dtype for array. + shape : int or tuple of ints, optional + Shape of each array. + offset : int, optional + Position in the file or buffer to start reading from. + strides : tuple of ints, optional + Buffer (`buf`) is interpreted according to these strides (strides + define how many bytes each array element, row, column, etc. + occupy in memory). + formats, names, titles, aligned, byteorder : + If `dtype` is ``None``, these arguments are passed to + `numpy.format_parser` to construct a dtype. See that function for + detailed documentation. + copy : bool, optional + Whether to copy the input object (True), or to use a reference instead. + This option only applies when the input is an ndarray or recarray. + Defaults to True. + + Returns + ------- + np.recarray + Record array created from the specified object. + + Notes + ----- + If `obj` is ``None``, then call the `~numpy.recarray` constructor. If + `obj` is a string, then call the `fromstring` constructor. If `obj` is a + list or a tuple, then if the first object is an `~numpy.ndarray`, call + `fromarrays`, otherwise call `fromrecords`. If `obj` is a + `~numpy.recarray`, then make a copy of the data in the recarray + (if ``copy=True``) and use the new formats, names, and titles. If `obj` + is a file, then call `fromfile`. Finally, if obj is an `ndarray`, then + return ``obj.view(recarray)``, making a copy of the data if ``copy=True``. + + Examples + -------- + >>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> a + array([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + + >>> np.rec.array(a) + rec.array([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]], + dtype=int64) + + >>> b = [(1, 1), (2, 4), (3, 9)] + >>> c = np.rec.array(b, formats = ['i2', 'f2'], names = ('x', 'y')) + >>> c + rec.array([(1, 1.), (2, 4.), (3, 9.)], + dtype=[('x', '>> c.x + array([1, 2, 3], dtype=int16) + + >>> c.y + array([1., 4., 9.], dtype=float16) + + >>> r = np.rec.array(['abc','def'], names=['col1','col2']) + >>> print(r.col1) + abc + + >>> r.col1 + array('abc', dtype='>> r.col2 + array('def', dtype=' object: ... + def tell(self, /) -> int: ... + def readinto(self, buffer: memoryview, /) -> int: ... + +### + +# exported in `numpy.rec` +class record(np.void): + def __getattribute__(self, attr: str) -> Any: ... + def __setattr__(self, attr: str, val: ArrayLike) -> None: ... + def pprint(self) -> str: ... + @overload + def __getitem__(self, key: str | SupportsIndex) -> Any: ... + @overload + def __getitem__(self, key: list[str]) -> record: ... + +# exported in `numpy.rec` +class recarray(np.ndarray[_ShapeT_co, _DTypeT_co]): + __name__: ClassVar[Literal["record"]] = "record" + __module__: Literal["numpy"] = "numpy" + @overload + def __new__( + subtype, + shape: _ShapeLike, + dtype: None = None, + buf: _SupportsBuffer | None = None, + offset: SupportsIndex = 0, + strides: _ShapeLike | None = None, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + byteorder: _ByteOrder | None = None, + aligned: bool = False, + order: _OrderKACF = "C", + ) -> _RecArray[record]: ... + @overload + def __new__( + subtype, + shape: _ShapeLike, + dtype: DTypeLike, + buf: _SupportsBuffer | None = None, + offset: SupportsIndex = 0, + strides: _ShapeLike | None = None, + formats: None = None, + names: None = None, + titles: None = None, + byteorder: None = None, + aligned: Literal[False] = False, + order: _OrderKACF = "C", + ) -> _RecArray[Any]: ... + def __array_finalize__(self, /, obj: object) -> None: ... + def __getattribute__(self, attr: str, /) -> Any: ... + def __setattr__(self, attr: str, val: ArrayLike, /) -> None: ... + + # + @overload + def field(self, /, attr: int | str, val: ArrayLike) -> None: ... + @overload + def field(self, /, attr: int | str, val: None = None) -> Any: ... + +# exported in `numpy.rec` +class format_parser: + dtype: np.dtype[np.void] + def __init__( + self, + /, + formats: DTypeLike, + names: str | Sequence[str] | None, + titles: str | Sequence[str] | None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, + ) -> None: ... + +# exported in `numpy.rec` +@overload +def fromarrays( + arrayList: Iterable[ArrayLike], + dtype: DTypeLike | None = None, + shape: _ShapeLike | None = None, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, +) -> _RecArray[Any]: ... +@overload +def fromarrays( + arrayList: Iterable[ArrayLike], + dtype: None = None, + shape: _ShapeLike | None = None, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, +) -> _RecArray[record]: ... + +@overload +def fromrecords( + recList: _ArrayLikeVoid_co | tuple[object, ...] | _NestedSequence[tuple[object, ...]], + dtype: DTypeLike | None = None, + shape: _ShapeLike | None = None, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, +) -> _RecArray[record]: ... +@overload +def fromrecords( + recList: _ArrayLikeVoid_co | tuple[object, ...] | _NestedSequence[tuple[object, ...]], + dtype: None = None, + shape: _ShapeLike | None = None, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, +) -> _RecArray[record]: ... + +# exported in `numpy.rec` +@overload +def fromstring( + datastring: _SupportsBuffer, + dtype: DTypeLike, + shape: _ShapeLike | None = None, + offset: int = 0, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, +) -> _RecArray[record]: ... +@overload +def fromstring( + datastring: _SupportsBuffer, + dtype: None = None, + shape: _ShapeLike | None = None, + offset: int = 0, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, +) -> _RecArray[record]: ... + +# exported in `numpy.rec` +@overload +def fromfile( + fd: StrOrBytesPath | _SupportsReadInto, + dtype: DTypeLike, + shape: _ShapeLike | None = None, + offset: int = 0, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, +) -> _RecArray[Any]: ... +@overload +def fromfile( + fd: StrOrBytesPath | _SupportsReadInto, + dtype: None = None, + shape: _ShapeLike | None = None, + offset: int = 0, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, +) -> _RecArray[record]: ... + +# exported in `numpy.rec` +@overload +def array( + obj: _SCT | NDArray[_SCT], + dtype: None = None, + shape: _ShapeLike | None = None, + offset: int = 0, + strides: tuple[int, ...] | None = None, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, + copy: bool = True, +) -> _RecArray[_SCT]: ... +@overload +def array( + obj: ArrayLike, + dtype: DTypeLike, + shape: _ShapeLike | None = None, + offset: int = 0, + strides: tuple[int, ...] | None = None, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, + copy: bool = True, +) -> _RecArray[Any]: ... +@overload +def array( + obj: ArrayLike, + dtype: None = None, + shape: _ShapeLike | None = None, + offset: int = 0, + strides: tuple[int, ...] | None = None, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, + copy: bool = True, +) -> _RecArray[record]: ... +@overload +def array( + obj: None, + dtype: DTypeLike, + shape: _ShapeLike, + offset: int = 0, + strides: tuple[int, ...] | None = None, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, + copy: bool = True, +) -> _RecArray[Any]: ... +@overload +def array( + obj: None, + dtype: None = None, + *, + shape: _ShapeLike, + offset: int = 0, + strides: tuple[int, ...] | None = None, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, + copy: bool = True, +) -> _RecArray[record]: ... +@overload +def array( + obj: _SupportsReadInto, + dtype: DTypeLike, + shape: _ShapeLike | None = None, + offset: int = 0, + strides: tuple[int, ...] | None = None, + formats: None = None, + names: None = None, + titles: None = None, + aligned: bool = False, + byteorder: None = None, + copy: bool = True, +) -> _RecArray[Any]: ... +@overload +def array( + obj: _SupportsReadInto, + dtype: None = None, + shape: _ShapeLike | None = None, + offset: int = 0, + strides: tuple[int, ...] | None = None, + *, + formats: DTypeLike, + names: str | Sequence[str] | None = None, + titles: str | Sequence[str] | None = None, + aligned: bool = False, + byteorder: _ByteOrder | None = None, + copy: bool = True, +) -> _RecArray[record]: ... + +# exported in `numpy.rec` +def find_duplicate(list: Iterable[_T]) -> list[_T]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/shape_base.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/shape_base.py new file mode 100644 index 0000000000000000000000000000000000000000..89a2e945c9d90ca8ce29de24287ae6c2ef26243c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/shape_base.py @@ -0,0 +1,1004 @@ +__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack', + 'stack', 'unstack', 'vstack'] + +import functools +import itertools +import operator + +from . import numeric as _nx +from . import overrides +from .multiarray import array, asanyarray, normalize_axis_index +from . import fromnumeric as _from_nx + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +def _atleast_1d_dispatcher(*arys): + return arys + + +@array_function_dispatch(_atleast_1d_dispatcher) +def atleast_1d(*arys): + """ + Convert inputs to arrays with at least one dimension. + + Scalar inputs are converted to 1-dimensional arrays, whilst + higher-dimensional inputs are preserved. + + Parameters + ---------- + arys1, arys2, ... : array_like + One or more input arrays. + + Returns + ------- + ret : ndarray + An array, or tuple of arrays, each with ``a.ndim >= 1``. + Copies are made only if necessary. + + See Also + -------- + atleast_2d, atleast_3d + + Examples + -------- + >>> import numpy as np + >>> np.atleast_1d(1.0) + array([1.]) + + >>> x = np.arange(9.0).reshape(3,3) + >>> np.atleast_1d(x) + array([[0., 1., 2.], + [3., 4., 5.], + [6., 7., 8.]]) + >>> np.atleast_1d(x) is x + True + + >>> np.atleast_1d(1, [3, 4]) + (array([1]), array([3, 4])) + + """ + if len(arys) == 1: + result = asanyarray(arys[0]) + if result.ndim == 0: + result = result.reshape(1) + return result + res = [] + for ary in arys: + result = asanyarray(ary) + if result.ndim == 0: + result = result.reshape(1) + res.append(result) + return tuple(res) + + +def _atleast_2d_dispatcher(*arys): + return arys + + +@array_function_dispatch(_atleast_2d_dispatcher) +def atleast_2d(*arys): + """ + View inputs as arrays with at least two dimensions. + + Parameters + ---------- + arys1, arys2, ... : array_like + One or more array-like sequences. Non-array inputs are converted + to arrays. Arrays that already have two or more dimensions are + preserved. + + Returns + ------- + res, res2, ... : ndarray + An array, or tuple of arrays, each with ``a.ndim >= 2``. + Copies are avoided where possible, and views with two or more + dimensions are returned. + + See Also + -------- + atleast_1d, atleast_3d + + Examples + -------- + >>> import numpy as np + >>> np.atleast_2d(3.0) + array([[3.]]) + + >>> x = np.arange(3.0) + >>> np.atleast_2d(x) + array([[0., 1., 2.]]) + >>> np.atleast_2d(x).base is x + True + + >>> np.atleast_2d(1, [1, 2], [[1, 2]]) + (array([[1]]), array([[1, 2]]), array([[1, 2]])) + + """ + res = [] + for ary in arys: + ary = asanyarray(ary) + if ary.ndim == 0: + result = ary.reshape(1, 1) + elif ary.ndim == 1: + result = ary[_nx.newaxis, :] + else: + result = ary + res.append(result) + if len(res) == 1: + return res[0] + else: + return tuple(res) + + +def _atleast_3d_dispatcher(*arys): + return arys + + +@array_function_dispatch(_atleast_3d_dispatcher) +def atleast_3d(*arys): + """ + View inputs as arrays with at least three dimensions. + + Parameters + ---------- + arys1, arys2, ... : array_like + One or more array-like sequences. Non-array inputs are converted to + arrays. Arrays that already have three or more dimensions are + preserved. + + Returns + ------- + res1, res2, ... : ndarray + An array, or tuple of arrays, each with ``a.ndim >= 3``. Copies are + avoided where possible, and views with three or more dimensions are + returned. For example, a 1-D array of shape ``(N,)`` becomes a view + of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a + view of shape ``(M, N, 1)``. + + See Also + -------- + atleast_1d, atleast_2d + + Examples + -------- + >>> import numpy as np + >>> np.atleast_3d(3.0) + array([[[3.]]]) + + >>> x = np.arange(3.0) + >>> np.atleast_3d(x).shape + (1, 3, 1) + + >>> x = np.arange(12.0).reshape(4,3) + >>> np.atleast_3d(x).shape + (4, 3, 1) + >>> np.atleast_3d(x).base is x.base # x is a reshape, so not base itself + True + + >>> for arr in np.atleast_3d([1, 2], [[1, 2]], [[[1, 2]]]): + ... print(arr, arr.shape) # doctest: +SKIP + ... + [[[1] + [2]]] (1, 2, 1) + [[[1] + [2]]] (1, 2, 1) + [[[1 2]]] (1, 1, 2) + + """ + res = [] + for ary in arys: + ary = asanyarray(ary) + if ary.ndim == 0: + result = ary.reshape(1, 1, 1) + elif ary.ndim == 1: + result = ary[_nx.newaxis, :, _nx.newaxis] + elif ary.ndim == 2: + result = ary[:, :, _nx.newaxis] + else: + result = ary + res.append(result) + if len(res) == 1: + return res[0] + else: + return tuple(res) + + +def _arrays_for_stack_dispatcher(arrays): + if not hasattr(arrays, "__getitem__"): + raise TypeError('arrays to stack must be passed as a "sequence" type ' + 'such as list or tuple.') + + return tuple(arrays) + + +def _vhstack_dispatcher(tup, *, dtype=None, casting=None): + return _arrays_for_stack_dispatcher(tup) + + +@array_function_dispatch(_vhstack_dispatcher) +def vstack(tup, *, dtype=None, casting="same_kind"): + """ + Stack arrays in sequence vertically (row wise). + + This is equivalent to concatenation along the first axis after 1-D arrays + of shape `(N,)` have been reshaped to `(1,N)`. Rebuilds arrays divided by + `vsplit`. + + This function makes most sense for arrays with up to 3 dimensions. For + instance, for pixel-data with a height (first axis), width (second axis), + and r/g/b channels (third axis). The functions `concatenate`, `stack` and + `block` provide more general stacking and concatenation operations. + + Parameters + ---------- + tup : sequence of ndarrays + The arrays must have the same shape along all but the first axis. + 1-D arrays must have the same length. In the case of a single + array_like input, it will be treated as a sequence of arrays; i.e., + each element along the zeroth axis is treated as a separate array. + + dtype : str or dtype + If provided, the destination array will have this dtype. Cannot be + provided together with `out`. + + .. versionadded:: 1.24 + + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Defaults to 'same_kind'. + + .. versionadded:: 1.24 + + Returns + ------- + stacked : ndarray + The array formed by stacking the given arrays, will be at least 2-D. + + See Also + -------- + concatenate : Join a sequence of arrays along an existing axis. + stack : Join a sequence of arrays along a new axis. + block : Assemble an nd-array from nested lists of blocks. + hstack : Stack arrays in sequence horizontally (column wise). + dstack : Stack arrays in sequence depth wise (along third axis). + column_stack : Stack 1-D arrays as columns into a 2-D array. + vsplit : Split an array into multiple sub-arrays vertically (row-wise). + unstack : Split an array into a tuple of sub-arrays along an axis. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1, 2, 3]) + >>> b = np.array([4, 5, 6]) + >>> np.vstack((a,b)) + array([[1, 2, 3], + [4, 5, 6]]) + + >>> a = np.array([[1], [2], [3]]) + >>> b = np.array([[4], [5], [6]]) + >>> np.vstack((a,b)) + array([[1], + [2], + [3], + [4], + [5], + [6]]) + + """ + arrs = atleast_2d(*tup) + if not isinstance(arrs, tuple): + arrs = (arrs,) + return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting) + + +@array_function_dispatch(_vhstack_dispatcher) +def hstack(tup, *, dtype=None, casting="same_kind"): + """ + Stack arrays in sequence horizontally (column wise). + + This is equivalent to concatenation along the second axis, except for 1-D + arrays where it concatenates along the first axis. Rebuilds arrays divided + by `hsplit`. + + This function makes most sense for arrays with up to 3 dimensions. For + instance, for pixel-data with a height (first axis), width (second axis), + and r/g/b channels (third axis). The functions `concatenate`, `stack` and + `block` provide more general stacking and concatenation operations. + + Parameters + ---------- + tup : sequence of ndarrays + The arrays must have the same shape along all but the second axis, + except 1-D arrays which can be any length. In the case of a single + array_like input, it will be treated as a sequence of arrays; i.e., + each element along the zeroth axis is treated as a separate array. + + dtype : str or dtype + If provided, the destination array will have this dtype. Cannot be + provided together with `out`. + + .. versionadded:: 1.24 + + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Defaults to 'same_kind'. + + .. versionadded:: 1.24 + + Returns + ------- + stacked : ndarray + The array formed by stacking the given arrays. + + See Also + -------- + concatenate : Join a sequence of arrays along an existing axis. + stack : Join a sequence of arrays along a new axis. + block : Assemble an nd-array from nested lists of blocks. + vstack : Stack arrays in sequence vertically (row wise). + dstack : Stack arrays in sequence depth wise (along third axis). + column_stack : Stack 1-D arrays as columns into a 2-D array. + hsplit : Split an array into multiple sub-arrays + horizontally (column-wise). + unstack : Split an array into a tuple of sub-arrays along an axis. + + Examples + -------- + >>> import numpy as np + >>> a = np.array((1,2,3)) + >>> b = np.array((4,5,6)) + >>> np.hstack((a,b)) + array([1, 2, 3, 4, 5, 6]) + >>> a = np.array([[1],[2],[3]]) + >>> b = np.array([[4],[5],[6]]) + >>> np.hstack((a,b)) + array([[1, 4], + [2, 5], + [3, 6]]) + + """ + arrs = atleast_1d(*tup) + if not isinstance(arrs, tuple): + arrs = (arrs,) + # As a special case, dimension 0 of 1-dimensional arrays is "horizontal" + if arrs and arrs[0].ndim == 1: + return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting) + else: + return _nx.concatenate(arrs, 1, dtype=dtype, casting=casting) + + +def _stack_dispatcher(arrays, axis=None, out=None, *, + dtype=None, casting=None): + arrays = _arrays_for_stack_dispatcher(arrays) + if out is not None: + # optimize for the typical case where only arrays is provided + arrays = list(arrays) + arrays.append(out) + return arrays + + +@array_function_dispatch(_stack_dispatcher) +def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"): + """ + Join a sequence of arrays along a new axis. + + The ``axis`` parameter specifies the index of the new axis in the + dimensions of the result. For example, if ``axis=0`` it will be the first + dimension and if ``axis=-1`` it will be the last dimension. + + Parameters + ---------- + arrays : sequence of ndarrays + Each array must have the same shape. In the case of a single ndarray + array_like input, it will be treated as a sequence of arrays; i.e., + each element along the zeroth axis is treated as a separate array. + + axis : int, optional + The axis in the result array along which the input arrays are stacked. + + out : ndarray, optional + If provided, the destination to place the result. The shape must be + correct, matching that of what stack would have returned if no + out argument were specified. + + dtype : str or dtype + If provided, the destination array will have this dtype. Cannot be + provided together with `out`. + + .. versionadded:: 1.24 + + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Defaults to 'same_kind'. + + .. versionadded:: 1.24 + + + Returns + ------- + stacked : ndarray + The stacked array has one more dimension than the input arrays. + + See Also + -------- + concatenate : Join a sequence of arrays along an existing axis. + block : Assemble an nd-array from nested lists of blocks. + split : Split array into a list of multiple sub-arrays of equal size. + unstack : Split an array into a tuple of sub-arrays along an axis. + + Examples + -------- + >>> import numpy as np + >>> rng = np.random.default_rng() + >>> arrays = [rng.normal(size=(3,4)) for _ in range(10)] + >>> np.stack(arrays, axis=0).shape + (10, 3, 4) + + >>> np.stack(arrays, axis=1).shape + (3, 10, 4) + + >>> np.stack(arrays, axis=2).shape + (3, 4, 10) + + >>> a = np.array([1, 2, 3]) + >>> b = np.array([4, 5, 6]) + >>> np.stack((a, b)) + array([[1, 2, 3], + [4, 5, 6]]) + + >>> np.stack((a, b), axis=-1) + array([[1, 4], + [2, 5], + [3, 6]]) + + """ + arrays = [asanyarray(arr) for arr in arrays] + if not arrays: + raise ValueError('need at least one array to stack') + + shapes = {arr.shape for arr in arrays} + if len(shapes) != 1: + raise ValueError('all input arrays must have the same shape') + + result_ndim = arrays[0].ndim + 1 + axis = normalize_axis_index(axis, result_ndim) + + sl = (slice(None),) * axis + (_nx.newaxis,) + expanded_arrays = [arr[sl] for arr in arrays] + return _nx.concatenate(expanded_arrays, axis=axis, out=out, + dtype=dtype, casting=casting) + +def _unstack_dispatcher(x, /, *, axis=None): + return (x,) + +@array_function_dispatch(_unstack_dispatcher) +def unstack(x, /, *, axis=0): + """ + Split an array into a sequence of arrays along the given axis. + + The ``axis`` parameter specifies the dimension along which the array will + be split. For example, if ``axis=0`` (the default) it will be the first + dimension and if ``axis=-1`` it will be the last dimension. + + The result is a tuple of arrays split along ``axis``. + + .. versionadded:: 2.1.0 + + Parameters + ---------- + x : ndarray + The array to be unstacked. + axis : int, optional + Axis along which the array will be split. Default: ``0``. + + Returns + ------- + unstacked : tuple of ndarrays + The unstacked arrays. + + See Also + -------- + stack : Join a sequence of arrays along a new axis. + concatenate : Join a sequence of arrays along an existing axis. + block : Assemble an nd-array from nested lists of blocks. + split : Split array into a list of multiple sub-arrays of equal size. + + Notes + ----- + ``unstack`` serves as the reverse operation of :py:func:`stack`, i.e., + ``stack(unstack(x, axis=axis), axis=axis) == x``. + + This function is equivalent to ``tuple(np.moveaxis(x, axis, 0))``, since + iterating on an array iterates along the first axis. + + Examples + -------- + >>> arr = np.arange(24).reshape((2, 3, 4)) + >>> np.unstack(arr) + (array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]]), + array([[12, 13, 14, 15], + [16, 17, 18, 19], + [20, 21, 22, 23]])) + >>> np.unstack(arr, axis=1) + (array([[ 0, 1, 2, 3], + [12, 13, 14, 15]]), + array([[ 4, 5, 6, 7], + [16, 17, 18, 19]]), + array([[ 8, 9, 10, 11], + [20, 21, 22, 23]])) + >>> arr2 = np.stack(np.unstack(arr, axis=1), axis=1) + >>> arr2.shape + (2, 3, 4) + >>> np.all(arr == arr2) + np.True_ + + """ + if x.ndim == 0: + raise ValueError("Input array must be at least 1-d.") + return tuple(_nx.moveaxis(x, axis, 0)) + +# Internal functions to eliminate the overhead of repeated dispatch in one of +# the two possible paths inside np.block. +# Use getattr to protect against __array_function__ being disabled. +_size = getattr(_from_nx.size, '__wrapped__', _from_nx.size) +_ndim = getattr(_from_nx.ndim, '__wrapped__', _from_nx.ndim) +_concatenate = getattr(_from_nx.concatenate, + '__wrapped__', _from_nx.concatenate) + + +def _block_format_index(index): + """ + Convert a list of indices ``[0, 1, 2]`` into ``"arrays[0][1][2]"``. + """ + idx_str = ''.join('[{}]'.format(i) for i in index if i is not None) + return 'arrays' + idx_str + + +def _block_check_depths_match(arrays, parent_index=[]): + """ + Recursive function checking that the depths of nested lists in `arrays` + all match. Mismatch raises a ValueError as described in the block + docstring below. + + The entire index (rather than just the depth) needs to be calculated + for each innermost list, in case an error needs to be raised, so that + the index of the offending list can be printed as part of the error. + + Parameters + ---------- + arrays : nested list of arrays + The arrays to check + parent_index : list of int + The full index of `arrays` within the nested lists passed to + `_block_check_depths_match` at the top of the recursion. + + Returns + ------- + first_index : list of int + The full index of an element from the bottom of the nesting in + `arrays`. If any element at the bottom is an empty list, this will + refer to it, and the last index along the empty axis will be None. + max_arr_ndim : int + The maximum of the ndims of the arrays nested in `arrays`. + final_size: int + The number of elements in the final array. This is used the motivate + the choice of algorithm used using benchmarking wisdom. + + """ + if type(arrays) is tuple: + # not strictly necessary, but saves us from: + # - more than one way to do things - no point treating tuples like + # lists + # - horribly confusing behaviour that results when tuples are + # treated like ndarray + raise TypeError( + '{} is a tuple. ' + 'Only lists can be used to arrange blocks, and np.block does ' + 'not allow implicit conversion from tuple to ndarray.'.format( + _block_format_index(parent_index) + ) + ) + elif type(arrays) is list and len(arrays) > 0: + idxs_ndims = (_block_check_depths_match(arr, parent_index + [i]) + for i, arr in enumerate(arrays)) + + first_index, max_arr_ndim, final_size = next(idxs_ndims) + for index, ndim, size in idxs_ndims: + final_size += size + if ndim > max_arr_ndim: + max_arr_ndim = ndim + if len(index) != len(first_index): + raise ValueError( + "List depths are mismatched. First element was at depth " + "{}, but there is an element at depth {} ({})".format( + len(first_index), + len(index), + _block_format_index(index) + ) + ) + # propagate our flag that indicates an empty list at the bottom + if index[-1] is None: + first_index = index + + return first_index, max_arr_ndim, final_size + elif type(arrays) is list and len(arrays) == 0: + # We've 'bottomed out' on an empty list + return parent_index + [None], 0, 0 + else: + # We've 'bottomed out' - arrays is either a scalar or an array + size = _size(arrays) + return parent_index, _ndim(arrays), size + + +def _atleast_nd(a, ndim): + # Ensures `a` has at least `ndim` dimensions by prepending + # ones to `a.shape` as necessary + return array(a, ndmin=ndim, copy=None, subok=True) + + +def _accumulate(values): + return list(itertools.accumulate(values)) + + +def _concatenate_shapes(shapes, axis): + """Given array shapes, return the resulting shape and slices prefixes. + + These help in nested concatenation. + + Returns + ------- + shape: tuple of int + This tuple satisfies:: + + shape, _ = _concatenate_shapes([arr.shape for shape in arrs], axis) + shape == concatenate(arrs, axis).shape + + slice_prefixes: tuple of (slice(start, end), ) + For a list of arrays being concatenated, this returns the slice + in the larger array at axis that needs to be sliced into. + + For example, the following holds:: + + ret = concatenate([a, b, c], axis) + _, (sl_a, sl_b, sl_c) = concatenate_slices([a, b, c], axis) + + ret[(slice(None),) * axis + sl_a] == a + ret[(slice(None),) * axis + sl_b] == b + ret[(slice(None),) * axis + sl_c] == c + + These are called slice prefixes since they are used in the recursive + blocking algorithm to compute the left-most slices during the + recursion. Therefore, they must be prepended to rest of the slice + that was computed deeper in the recursion. + + These are returned as tuples to ensure that they can quickly be added + to existing slice tuple without creating a new tuple every time. + + """ + # Cache a result that will be reused. + shape_at_axis = [shape[axis] for shape in shapes] + + # Take a shape, any shape + first_shape = shapes[0] + first_shape_pre = first_shape[:axis] + first_shape_post = first_shape[axis+1:] + + if any(shape[:axis] != first_shape_pre or + shape[axis+1:] != first_shape_post for shape in shapes): + raise ValueError( + 'Mismatched array shapes in block along axis {}.'.format(axis)) + + shape = (first_shape_pre + (sum(shape_at_axis),) + first_shape[axis+1:]) + + offsets_at_axis = _accumulate(shape_at_axis) + slice_prefixes = [(slice(start, end),) + for start, end in zip([0] + offsets_at_axis, + offsets_at_axis)] + return shape, slice_prefixes + + +def _block_info_recursion(arrays, max_depth, result_ndim, depth=0): + """ + Returns the shape of the final array, along with a list + of slices and a list of arrays that can be used for assignment inside the + new array + + Parameters + ---------- + arrays : nested list of arrays + The arrays to check + max_depth : list of int + The number of nested lists + result_ndim : int + The number of dimensions in thefinal array. + + Returns + ------- + shape : tuple of int + The shape that the final array will take on. + slices: list of tuple of slices + The slices into the full array required for assignment. These are + required to be prepended with ``(Ellipsis, )`` to obtain to correct + final index. + arrays: list of ndarray + The data to assign to each slice of the full array + + """ + if depth < max_depth: + shapes, slices, arrays = zip( + *[_block_info_recursion(arr, max_depth, result_ndim, depth+1) + for arr in arrays]) + + axis = result_ndim - max_depth + depth + shape, slice_prefixes = _concatenate_shapes(shapes, axis) + + # Prepend the slice prefix and flatten the slices + slices = [slice_prefix + the_slice + for slice_prefix, inner_slices in zip(slice_prefixes, slices) + for the_slice in inner_slices] + + # Flatten the array list + arrays = functools.reduce(operator.add, arrays) + + return shape, slices, arrays + else: + # We've 'bottomed out' - arrays is either a scalar or an array + # type(arrays) is not list + # Return the slice and the array inside a list to be consistent with + # the recursive case. + arr = _atleast_nd(arrays, result_ndim) + return arr.shape, [()], [arr] + + +def _block(arrays, max_depth, result_ndim, depth=0): + """ + Internal implementation of block based on repeated concatenation. + `arrays` is the argument passed to + block. `max_depth` is the depth of nested lists within `arrays` and + `result_ndim` is the greatest of the dimensions of the arrays in + `arrays` and the depth of the lists in `arrays` (see block docstring + for details). + """ + if depth < max_depth: + arrs = [_block(arr, max_depth, result_ndim, depth+1) + for arr in arrays] + return _concatenate(arrs, axis=-(max_depth-depth)) + else: + # We've 'bottomed out' - arrays is either a scalar or an array + # type(arrays) is not list + return _atleast_nd(arrays, result_ndim) + + +def _block_dispatcher(arrays): + # Use type(...) is list to match the behavior of np.block(), which special + # cases list specifically rather than allowing for generic iterables or + # tuple. Also, we know that list.__array_function__ will never exist. + if type(arrays) is list: + for subarrays in arrays: + yield from _block_dispatcher(subarrays) + else: + yield arrays + + +@array_function_dispatch(_block_dispatcher) +def block(arrays): + """ + Assemble an nd-array from nested lists of blocks. + + Blocks in the innermost lists are concatenated (see `concatenate`) along + the last dimension (-1), then these are concatenated along the + second-last dimension (-2), and so on until the outermost list is reached. + + Blocks can be of any dimension, but will not be broadcasted using + the normal rules. Instead, leading axes of size 1 are inserted, + to make ``block.ndim`` the same for all blocks. This is primarily useful + for working with scalars, and means that code like ``np.block([v, 1])`` + is valid, where ``v.ndim == 1``. + + When the nested list is two levels deep, this allows block matrices to be + constructed from their components. + + Parameters + ---------- + arrays : nested list of array_like or scalars (but not tuples) + If passed a single ndarray or scalar (a nested list of depth 0), this + is returned unmodified (and not copied). + + Elements shapes must match along the appropriate axes (without + broadcasting), but leading 1s will be prepended to the shape as + necessary to make the dimensions match. + + Returns + ------- + block_array : ndarray + The array assembled from the given blocks. + + The dimensionality of the output is equal to the greatest of: + + * the dimensionality of all the inputs + * the depth to which the input list is nested + + Raises + ------ + ValueError + * If list depths are mismatched - for instance, ``[[a, b], c]`` is + illegal, and should be spelt ``[[a, b], [c]]`` + * If lists are empty - for instance, ``[[a, b], []]`` + + See Also + -------- + concatenate : Join a sequence of arrays along an existing axis. + stack : Join a sequence of arrays along a new axis. + vstack : Stack arrays in sequence vertically (row wise). + hstack : Stack arrays in sequence horizontally (column wise). + dstack : Stack arrays in sequence depth wise (along third axis). + column_stack : Stack 1-D arrays as columns into a 2-D array. + vsplit : Split an array into multiple sub-arrays vertically (row-wise). + unstack : Split an array into a tuple of sub-arrays along an axis. + + Notes + ----- + When called with only scalars, ``np.block`` is equivalent to an ndarray + call. So ``np.block([[1, 2], [3, 4]])`` is equivalent to + ``np.array([[1, 2], [3, 4]])``. + + This function does not enforce that the blocks lie on a fixed grid. + ``np.block([[a, b], [c, d]])`` is not restricted to arrays of the form:: + + AAAbb + AAAbb + cccDD + + But is also allowed to produce, for some ``a, b, c, d``:: + + AAAbb + AAAbb + cDDDD + + Since concatenation happens along the last axis first, `block` is *not* + capable of producing the following directly:: + + AAAbb + cccbb + cccDD + + Matlab's "square bracket stacking", ``[A, B, ...; p, q, ...]``, is + equivalent to ``np.block([[A, B, ...], [p, q, ...]])``. + + Examples + -------- + The most common use of this function is to build a block matrix: + + >>> import numpy as np + >>> A = np.eye(2) * 2 + >>> B = np.eye(3) * 3 + >>> np.block([ + ... [A, np.zeros((2, 3))], + ... [np.ones((3, 2)), B ] + ... ]) + array([[2., 0., 0., 0., 0.], + [0., 2., 0., 0., 0.], + [1., 1., 3., 0., 0.], + [1., 1., 0., 3., 0.], + [1., 1., 0., 0., 3.]]) + + With a list of depth 1, `block` can be used as `hstack`: + + >>> np.block([1, 2, 3]) # hstack([1, 2, 3]) + array([1, 2, 3]) + + >>> a = np.array([1, 2, 3]) + >>> b = np.array([4, 5, 6]) + >>> np.block([a, b, 10]) # hstack([a, b, 10]) + array([ 1, 2, 3, 4, 5, 6, 10]) + + >>> A = np.ones((2, 2), int) + >>> B = 2 * A + >>> np.block([A, B]) # hstack([A, B]) + array([[1, 1, 2, 2], + [1, 1, 2, 2]]) + + With a list of depth 2, `block` can be used in place of `vstack`: + + >>> a = np.array([1, 2, 3]) + >>> b = np.array([4, 5, 6]) + >>> np.block([[a], [b]]) # vstack([a, b]) + array([[1, 2, 3], + [4, 5, 6]]) + + >>> A = np.ones((2, 2), int) + >>> B = 2 * A + >>> np.block([[A], [B]]) # vstack([A, B]) + array([[1, 1], + [1, 1], + [2, 2], + [2, 2]]) + + It can also be used in place of `atleast_1d` and `atleast_2d`: + + >>> a = np.array(0) + >>> b = np.array([1]) + >>> np.block([a]) # atleast_1d(a) + array([0]) + >>> np.block([b]) # atleast_1d(b) + array([1]) + + >>> np.block([[a]]) # atleast_2d(a) + array([[0]]) + >>> np.block([[b]]) # atleast_2d(b) + array([[1]]) + + + """ + arrays, list_ndim, result_ndim, final_size = _block_setup(arrays) + + # It was found through benchmarking that making an array of final size + # around 256x256 was faster by straight concatenation on a + # i7-7700HQ processor and dual channel ram 2400MHz. + # It didn't seem to matter heavily on the dtype used. + # + # A 2D array using repeated concatenation requires 2 copies of the array. + # + # The fastest algorithm will depend on the ratio of CPU power to memory + # speed. + # One can monitor the results of the benchmark + # https://pv.github.io/numpy-bench/#bench_shape_base.Block2D.time_block2d + # to tune this parameter until a C version of the `_block_info_recursion` + # algorithm is implemented which would likely be faster than the python + # version. + if list_ndim * final_size > (2 * 512 * 512): + return _block_slicing(arrays, list_ndim, result_ndim) + else: + return _block_concatenate(arrays, list_ndim, result_ndim) + + +# These helper functions are mostly used for testing. +# They allow us to write tests that directly call `_block_slicing` +# or `_block_concatenate` without blocking large arrays to force the wisdom +# to trigger the desired path. +def _block_setup(arrays): + """ + Returns + (`arrays`, list_ndim, result_ndim, final_size) + """ + bottom_index, arr_ndim, final_size = _block_check_depths_match(arrays) + list_ndim = len(bottom_index) + if bottom_index and bottom_index[-1] is None: + raise ValueError( + 'List at {} cannot be empty'.format( + _block_format_index(bottom_index) + ) + ) + result_ndim = max(arr_ndim, list_ndim) + return arrays, list_ndim, result_ndim, final_size + + +def _block_slicing(arrays, list_ndim, result_ndim): + shape, slices, arrays = _block_info_recursion( + arrays, list_ndim, result_ndim) + dtype = _nx.result_type(*[arr.dtype for arr in arrays]) + + # Test preferring F only in the case that all input arrays are F + F_order = all(arr.flags['F_CONTIGUOUS'] for arr in arrays) + C_order = all(arr.flags['C_CONTIGUOUS'] for arr in arrays) + order = 'F' if F_order and not C_order else 'C' + result = _nx.empty(shape=shape, dtype=dtype, order=order) + # Note: In a c implementation, the function + # PyArray_CreateMultiSortedStridePerm could be used for more advanced + # guessing of the desired order. + + for the_slice, arr in zip(slices, arrays): + result[(Ellipsis,) + the_slice] = arr + return result + + +def _block_concatenate(arrays, list_ndim, result_ndim): + result = _block(arrays, list_ndim, result_ndim) + if list_ndim == 0: + # Catch an edge case where _block returns a view because + # `arrays` is a single numpy array and not a list of numpy arrays. + # This might copy scalars or lists twice, but this isn't a likely + # usecase for those interested in performance + result = result.copy() + return result diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/shape_base.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/shape_base.pyi new file mode 100644 index 0000000000000000000000000000000000000000..93f1b3ff0a74224d3c0149a7f5cd8c8e0dad2292 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/shape_base.pyi @@ -0,0 +1,175 @@ +from collections.abc import Sequence +from typing import Any, SupportsIndex, TypeVar, overload + +from numpy import _CastingKind, generic +from numpy._typing import ArrayLike, DTypeLike, NDArray, _ArrayLike, _DTypeLike + +__all__ = [ + "atleast_1d", + "atleast_2d", + "atleast_3d", + "block", + "hstack", + "stack", + "unstack", + "vstack", +] + +_SCT = TypeVar("_SCT", bound=generic) +_SCT1 = TypeVar("_SCT1", bound=generic) +_SCT2 = TypeVar("_SCT2", bound=generic) +_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any]) + +### + +@overload +def atleast_1d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ... +@overload +def atleast_1d(a0: _ArrayLike[_SCT1], a1: _ArrayLike[_SCT2], /) -> tuple[NDArray[_SCT1], NDArray[_SCT2]]: ... +@overload +def atleast_1d(a0: _ArrayLike[_SCT], a1: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ... +@overload +def atleast_1d(a0: ArrayLike, /) -> NDArray[Any]: ... +@overload +def atleast_1d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ... +@overload +def atleast_1d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ... + +# +@overload +def atleast_2d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ... +@overload +def atleast_2d(a0: _ArrayLike[_SCT1], a1: _ArrayLike[_SCT2], /) -> tuple[NDArray[_SCT1], NDArray[_SCT2]]: ... +@overload +def atleast_2d(a0: _ArrayLike[_SCT], a1: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ... +@overload +def atleast_2d(a0: ArrayLike, /) -> NDArray[Any]: ... +@overload +def atleast_2d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ... +@overload +def atleast_2d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ... + +# +@overload +def atleast_3d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ... +@overload +def atleast_3d(a0: _ArrayLike[_SCT1], a1: _ArrayLike[_SCT2], /) -> tuple[NDArray[_SCT1], NDArray[_SCT2]]: ... +@overload +def atleast_3d(a0: _ArrayLike[_SCT], a1: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ... +@overload +def atleast_3d(a0: ArrayLike, /) -> NDArray[Any]: ... +@overload +def atleast_3d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ... +@overload +def atleast_3d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ... + +# +@overload +def vstack( + tup: Sequence[_ArrayLike[_SCT]], + *, + dtype: None = ..., + casting: _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def vstack( + tup: Sequence[ArrayLike], + *, + dtype: _DTypeLike[_SCT], + casting: _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def vstack( + tup: Sequence[ArrayLike], + *, + dtype: DTypeLike = ..., + casting: _CastingKind = ... +) -> NDArray[Any]: ... + +@overload +def hstack( + tup: Sequence[_ArrayLike[_SCT]], + *, + dtype: None = ..., + casting: _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def hstack( + tup: Sequence[ArrayLike], + *, + dtype: _DTypeLike[_SCT], + casting: _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def hstack( + tup: Sequence[ArrayLike], + *, + dtype: DTypeLike = ..., + casting: _CastingKind = ... +) -> NDArray[Any]: ... + +@overload +def stack( + arrays: Sequence[_ArrayLike[_SCT]], + axis: SupportsIndex = ..., + out: None = ..., + *, + dtype: None = ..., + casting: _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def stack( + arrays: Sequence[ArrayLike], + axis: SupportsIndex = ..., + out: None = ..., + *, + dtype: _DTypeLike[_SCT], + casting: _CastingKind = ... +) -> NDArray[_SCT]: ... +@overload +def stack( + arrays: Sequence[ArrayLike], + axis: SupportsIndex = ..., + out: None = ..., + *, + dtype: DTypeLike = ..., + casting: _CastingKind = ... +) -> NDArray[Any]: ... +@overload +def stack( + arrays: Sequence[ArrayLike], + axis: SupportsIndex, + out: _ArrayT, + *, + dtype: DTypeLike | None = None, + casting: _CastingKind = "same_kind", +) -> _ArrayT: ... +@overload +def stack( + arrays: Sequence[ArrayLike], + axis: SupportsIndex = 0, + *, + out: _ArrayT, + dtype: DTypeLike | None = None, + casting: _CastingKind = "same_kind", +) -> _ArrayT: ... + +@overload +def unstack( + array: _ArrayLike[_SCT], + /, + *, + axis: int = ..., +) -> tuple[NDArray[_SCT], ...]: ... +@overload +def unstack( + array: ArrayLike, + /, + *, + axis: int = ..., +) -> tuple[NDArray[Any], ...]: ... + +@overload +def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def block(arrays: ArrayLike) -> NDArray[Any]: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/strings.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/strings.py new file mode 100644 index 0000000000000000000000000000000000000000..e0f4e7b2fb97b787504c14b048f14b387dc72811 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/strings.py @@ -0,0 +1,1641 @@ +""" +This module contains a set of functions for vectorized string +operations. +""" + +import sys +import numpy as np +from numpy import ( + equal, not_equal, less, less_equal, greater, greater_equal, + add, multiply as _multiply_ufunc, +) +from numpy._core.multiarray import _vec_string +from numpy._core.overrides import set_module +from numpy._core.umath import ( + isalpha, + isdigit, + isspace, + isalnum, + islower, + isupper, + istitle, + isdecimal, + isnumeric, + str_len, + find as _find_ufunc, + rfind as _rfind_ufunc, + index as _index_ufunc, + rindex as _rindex_ufunc, + count as _count_ufunc, + startswith as _startswith_ufunc, + endswith as _endswith_ufunc, + _lstrip_whitespace, + _lstrip_chars, + _rstrip_whitespace, + _rstrip_chars, + _strip_whitespace, + _strip_chars, + _replace, + _expandtabs_length, + _expandtabs, + _center, + _ljust, + _rjust, + _zfill, + _partition, + _partition_index, + _rpartition, + _rpartition_index, +) + + +def _override___module__(): + for ufunc in [ + isalnum, isalpha, isdecimal, isdigit, islower, isnumeric, isspace, + istitle, isupper, str_len, + ]: + ufunc.__module__ = "numpy.strings" + ufunc.__qualname__ = ufunc.__name__ + + +_override___module__() + + +__all__ = [ + # UFuncs + "equal", "not_equal", "less", "less_equal", "greater", "greater_equal", + "add", "multiply", "isalpha", "isdigit", "isspace", "isalnum", "islower", + "isupper", "istitle", "isdecimal", "isnumeric", "str_len", "find", + "rfind", "index", "rindex", "count", "startswith", "endswith", "lstrip", + "rstrip", "strip", "replace", "expandtabs", "center", "ljust", "rjust", + "zfill", "partition", "rpartition", + + # _vec_string - Will gradually become ufuncs as well + "upper", "lower", "swapcase", "capitalize", "title", + + # _vec_string - Will probably not become ufuncs + "mod", "decode", "encode", "translate", + + # Removed from namespace until behavior has been crystallized + # "join", "split", "rsplit", "splitlines", +] + + +MAX = np.iinfo(np.int64).max + + +def _get_num_chars(a): + """ + Helper function that returns the number of characters per field in + a string or unicode array. This is to abstract out the fact that + for a unicode array this is itemsize / 4. + """ + if issubclass(a.dtype.type, np.str_): + return a.itemsize // 4 + return a.itemsize + + +def _to_bytes_or_str_array(result, output_dtype_like): + """ + Helper function to cast a result back into an array + with the appropriate dtype if an object array must be used + as an intermediary. + """ + output_dtype_like = np.asarray(output_dtype_like) + if result.size == 0: + # Calling asarray & tolist in an empty array would result + # in losing shape information + return result.astype(output_dtype_like.dtype) + ret = np.asarray(result.tolist()) + if isinstance(output_dtype_like.dtype, np.dtypes.StringDType): + return ret.astype(type(output_dtype_like.dtype)) + return ret.astype(type(output_dtype_like.dtype)(_get_num_chars(ret))) + + +def _clean_args(*args): + """ + Helper function for delegating arguments to Python string + functions. + + Many of the Python string operations that have optional arguments + do not use 'None' to indicate a default value. In these cases, + we need to remove all None arguments, and those following them. + """ + newargs = [] + for chk in args: + if chk is None: + break + newargs.append(chk) + return newargs + + +@set_module("numpy.strings") +def multiply(a, i): + """ + Return (a * i), that is string multiple concatenation, + element-wise. + + Values in ``i`` of less than 0 are treated as 0 (which yields an + empty string). + + Parameters + ---------- + a : array_like, with ``StringDType``, ``bytes_`` or ``str_`` dtype + + i : array_like, with any integer dtype + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input types + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["a", "b", "c"]) + >>> np.strings.multiply(a, 3) + array(['aaa', 'bbb', 'ccc'], dtype='>> i = np.array([1, 2, 3]) + >>> np.strings.multiply(a, i) + array(['a', 'bb', 'ccc'], dtype='>> np.strings.multiply(np.array(['a']), i) + array(['a', 'aa', 'aaa'], dtype='>> a = np.array(['a', 'b', 'c', 'd', 'e', 'f']).reshape((2, 3)) + >>> np.strings.multiply(a, 3) + array([['aaa', 'bbb', 'ccc'], + ['ddd', 'eee', 'fff']], dtype='>> np.strings.multiply(a, i) + array([['a', 'bb', 'ccc'], + ['d', 'ee', 'fff']], dtype=' sys.maxsize / np.maximum(i, 1)): + raise MemoryError("repeated string is too long") + + buffersizes = a_len * i + out_dtype = f"{a.dtype.char}{buffersizes.max()}" + out = np.empty_like(a, shape=buffersizes.shape, dtype=out_dtype) + return _multiply_ufunc(a, i, out=out) + + +@set_module("numpy.strings") +def mod(a, values): + """ + Return (a % i), that is pre-Python 2.6 string formatting + (interpolation), element-wise for a pair of array_likes of str + or unicode. + + Parameters + ---------- + a : array_like, with `np.bytes_` or `np.str_` dtype + + values : array_like of values + These values will be element-wise interpolated into the string. + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input types + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["NumPy is a %s library"]) + >>> np.strings.mod(a, values=["Python"]) + array(['NumPy is a Python library'], dtype='>> a = np.array([b'%d bytes', b'%d bits']) + >>> values = np.array([8, 64]) + >>> np.strings.mod(a, values) + array([b'8 bytes', b'64 bits'], dtype='|S7') + + """ + return _to_bytes_or_str_array( + _vec_string(a, np.object_, '__mod__', (values,)), a) + + +@set_module("numpy.strings") +def find(a, sub, start=0, end=None): + """ + For each element, return the lowest index in the string where + substring ``sub`` is found, such that ``sub`` is contained in the + range [``start``, ``end``). + + Parameters + ---------- + a : array_like, with ``StringDType``, ``bytes_`` or ``str_`` dtype + + sub : array_like, with `np.bytes_` or `np.str_` dtype + The substring to search for. + + start, end : array_like, with any integer dtype + The range to look in, interpreted as in slice notation. + + Returns + ------- + y : ndarray + Output array of ints + + See Also + -------- + str.find + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["NumPy is a Python library"]) + >>> np.strings.find(a, "Python") + array([11]) + + """ + end = end if end is not None else MAX + return _find_ufunc(a, sub, start, end) + + +@set_module("numpy.strings") +def rfind(a, sub, start=0, end=None): + """ + For each element, return the highest index in the string where + substring ``sub`` is found, such that ``sub`` is contained in the + range [``start``, ``end``). + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + sub : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + The substring to search for. + + start, end : array_like, with any integer dtype + The range to look in, interpreted as in slice notation. + + Returns + ------- + y : ndarray + Output array of ints + + See Also + -------- + str.rfind + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["Computer Science"]) + >>> np.strings.rfind(a, "Science", start=0, end=None) + array([9]) + >>> np.strings.rfind(a, "Science", start=0, end=8) + array([-1]) + >>> b = np.array(["Computer Science", "Science"]) + >>> np.strings.rfind(b, "Science", start=0, end=None) + array([9, 0]) + + """ + end = end if end is not None else MAX + return _rfind_ufunc(a, sub, start, end) + + +@set_module("numpy.strings") +def index(a, sub, start=0, end=None): + """ + Like `find`, but raises :exc:`ValueError` when the substring is not found. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + sub : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + start, end : array_like, with any integer dtype, optional + + Returns + ------- + out : ndarray + Output array of ints. + + See Also + -------- + find, str.index + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["Computer Science"]) + >>> np.strings.index(a, "Science", start=0, end=None) + array([9]) + + """ + end = end if end is not None else MAX + return _index_ufunc(a, sub, start, end) + + +@set_module("numpy.strings") +def rindex(a, sub, start=0, end=None): + """ + Like `rfind`, but raises :exc:`ValueError` when the substring `sub` is + not found. + + Parameters + ---------- + a : array-like, with `np.bytes_` or `np.str_` dtype + + sub : array-like, with `np.bytes_` or `np.str_` dtype + + start, end : array-like, with any integer dtype, optional + + Returns + ------- + out : ndarray + Output array of ints. + + See Also + -------- + rfind, str.rindex + + Examples + -------- + >>> a = np.array(["Computer Science"]) + >>> np.strings.rindex(a, "Science", start=0, end=None) + array([9]) + + """ + end = end if end is not None else MAX + return _rindex_ufunc(a, sub, start, end) + + +@set_module("numpy.strings") +def count(a, sub, start=0, end=None): + """ + Returns an array with the number of non-overlapping occurrences of + substring ``sub`` in the range [``start``, ``end``). + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + sub : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + The substring to search for. + + start, end : array_like, with any integer dtype + The range to look in, interpreted as in slice notation. + + Returns + ------- + y : ndarray + Output array of ints + + See Also + -------- + str.count + + Examples + -------- + >>> import numpy as np + >>> c = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> c + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> np.strings.count(c, 'A') + array([3, 1, 1]) + >>> np.strings.count(c, 'aA') + array([3, 1, 0]) + >>> np.strings.count(c, 'A', start=1, end=4) + array([2, 1, 1]) + >>> np.strings.count(c, 'A', start=1, end=3) + array([1, 0, 0]) + + """ + end = end if end is not None else MAX + return _count_ufunc(a, sub, start, end) + + +@set_module("numpy.strings") +def startswith(a, prefix, start=0, end=None): + """ + Returns a boolean array which is `True` where the string element + in ``a`` starts with ``prefix``, otherwise `False`. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + prefix : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + start, end : array_like, with any integer dtype + With ``start``, test beginning at that position. With ``end``, + stop comparing at that position. + + Returns + ------- + out : ndarray + Output array of bools + + See Also + -------- + str.startswith + + Examples + -------- + >>> import numpy as np + >>> s = np.array(['foo', 'bar']) + >>> s + array(['foo', 'bar'], dtype='>> np.strings.startswith(s, 'fo') + array([True, False]) + >>> np.strings.startswith(s, 'o', start=1, end=2) + array([True, False]) + + """ + end = end if end is not None else MAX + return _startswith_ufunc(a, prefix, start, end) + + +@set_module("numpy.strings") +def endswith(a, suffix, start=0, end=None): + """ + Returns a boolean array which is `True` where the string element + in ``a`` ends with ``suffix``, otherwise `False`. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + suffix : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + start, end : array_like, with any integer dtype + With ``start``, test beginning at that position. With ``end``, + stop comparing at that position. + + Returns + ------- + out : ndarray + Output array of bools + + See Also + -------- + str.endswith + + Examples + -------- + >>> import numpy as np + >>> s = np.array(['foo', 'bar']) + >>> s + array(['foo', 'bar'], dtype='>> np.strings.endswith(s, 'ar') + array([False, True]) + >>> np.strings.endswith(s, 'a', start=1, end=2) + array([False, True]) + + """ + end = end if end is not None else MAX + return _endswith_ufunc(a, suffix, start, end) + + +@set_module("numpy.strings") +def decode(a, encoding=None, errors=None): + r""" + Calls :meth:`bytes.decode` element-wise. + + The set of available codecs comes from the Python standard library, + and may be extended at runtime. For more information, see the + :mod:`codecs` module. + + Parameters + ---------- + a : array_like, with ``bytes_`` dtype + + encoding : str, optional + The name of an encoding + + errors : str, optional + Specifies how to handle encoding errors + + Returns + ------- + out : ndarray + + See Also + -------- + :py:meth:`bytes.decode` + + Notes + ----- + The type of the result will depend on the encoding specified. + + Examples + -------- + >>> import numpy as np + >>> c = np.array([b'\x81\xc1\x81\xc1\x81\xc1', b'@@\x81\xc1@@', + ... b'\x81\x82\xc2\xc1\xc2\x82\x81']) + >>> c + array([b'\x81\xc1\x81\xc1\x81\xc1', b'@@\x81\xc1@@', + b'\x81\x82\xc2\xc1\xc2\x82\x81'], dtype='|S7') + >>> np.strings.decode(c, encoding='cp037') + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> import numpy as np + >>> a = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> np.strings.encode(a, encoding='cp037') + array([b'\x81\xc1\x81\xc1\x81\xc1', b'@@\x81\xc1@@', + b'\x81\x82\xc2\xc1\xc2\x82\x81'], dtype='|S7') + + """ + return _to_bytes_or_str_array( + _vec_string(a, np.object_, 'encode', _clean_args(encoding, errors)), + np.bytes_(b'')) + + +@set_module("numpy.strings") +def expandtabs(a, tabsize=8): + """ + Return a copy of each string element where all tab characters are + replaced by one or more spaces. + + Calls :meth:`str.expandtabs` element-wise. + + Return a copy of each string element where all tab characters are + replaced by one or more spaces, depending on the current column + and the given `tabsize`. The column number is reset to zero after + each newline occurring in the string. This doesn't understand other + non-printing characters or escape sequences. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + Input array + tabsize : int, optional + Replace tabs with `tabsize` number of spaces. If not given defaults + to 8 spaces. + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input type + + See Also + -------- + str.expandtabs + + Examples + -------- + >>> import numpy as np + >>> a = np.array(['\t\tHello\tworld']) + >>> np.strings.expandtabs(a, tabsize=4) # doctest: +SKIP + array([' Hello world'], dtype='>> import numpy as np + >>> c = np.array(['a1b2','1b2a','b2a1','2a1b']); c + array(['a1b2', '1b2a', 'b2a1', '2a1b'], dtype='>> np.strings.center(c, width=9) + array([' a1b2 ', ' 1b2a ', ' b2a1 ', ' 2a1b '], dtype='>> np.strings.center(c, width=9, fillchar='*') + array(['***a1b2**', '***1b2a**', '***b2a1**', '***2a1b**'], dtype='>> np.strings.center(c, width=1) + array(['a1b2', '1b2a', 'b2a1', '2a1b'], dtype='>> import numpy as np + >>> c = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> np.strings.ljust(c, width=3) + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> np.strings.ljust(c, width=9) + array(['aAaAaA ', ' aA ', 'abBABba '], dtype='>> import numpy as np + >>> a = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> np.strings.rjust(a, width=3) + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> np.strings.rjust(a, width=9) + array([' aAaAaA', ' aA ', ' abBABba'], dtype='>> import numpy as np + >>> np.strings.zfill(['1', '-1', '+1'], 3) + array(['001', '-01', '+01'], dtype='>> import numpy as np + >>> c = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> c + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> np.strings.lstrip(c, 'a') + array(['AaAaA', ' aA ', 'bBABba'], dtype='>> np.strings.lstrip(c, 'A') # leaves c unchanged + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> (np.strings.lstrip(c, ' ') == np.strings.lstrip(c, '')).all() + np.False_ + >>> (np.strings.lstrip(c, ' ') == np.strings.lstrip(c)).all() + np.True_ + + """ + if chars is None: + return _lstrip_whitespace(a) + return _lstrip_chars(a, chars) + + +@set_module("numpy.strings") +def rstrip(a, chars=None): + """ + For each element in `a`, return a copy with the trailing characters + removed. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + chars : scalar with the same dtype as ``a``, optional + The ``chars`` argument is a string specifying the set of + characters to be removed. If ``None``, the ``chars`` + argument defaults to removing whitespace. The ``chars`` argument + is not a prefix or suffix; rather, all combinations of its + values are stripped. + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input types + + See Also + -------- + str.rstrip + + Examples + -------- + >>> import numpy as np + >>> c = np.array(['aAaAaA', 'abBABba']) + >>> c + array(['aAaAaA', 'abBABba'], dtype='>> np.strings.rstrip(c, 'a') + array(['aAaAaA', 'abBABb'], dtype='>> np.strings.rstrip(c, 'A') + array(['aAaAa', 'abBABba'], dtype='>> import numpy as np + >>> c = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> c + array(['aAaAaA', ' aA ', 'abBABba'], dtype='>> np.strings.strip(c) + array(['aAaAaA', 'aA', 'abBABba'], dtype='>> np.strings.strip(c, 'a') + array(['AaAaA', ' aA ', 'bBABb'], dtype='>> np.strings.strip(c, 'A') + array(['aAaAa', ' aA ', 'abBABba'], dtype='>> import numpy as np + >>> c = np.array(['a1b c', '1bca', 'bca1']); c + array(['a1b c', '1bca', 'bca1'], dtype='>> np.strings.upper(c) + array(['A1B C', '1BCA', 'BCA1'], dtype='>> import numpy as np + >>> c = np.array(['A1B C', '1BCA', 'BCA1']); c + array(['A1B C', '1BCA', 'BCA1'], dtype='>> np.strings.lower(c) + array(['a1b c', '1bca', 'bca1'], dtype='>> import numpy as np + >>> c=np.array(['a1B c','1b Ca','b Ca1','cA1b'],'S5'); c + array(['a1B c', '1b Ca', 'b Ca1', 'cA1b'], + dtype='|S5') + >>> np.strings.swapcase(c) + array(['A1b C', '1B cA', 'B cA1', 'Ca1B'], + dtype='|S5') + + """ + a_arr = np.asarray(a) + return _vec_string(a_arr, a_arr.dtype, 'swapcase') + + +@set_module("numpy.strings") +def capitalize(a): + """ + Return a copy of ``a`` with only the first character of each element + capitalized. + + Calls :meth:`str.capitalize` element-wise. + + For byte strings, this method is locale-dependent. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + Input array of strings to capitalize. + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input types + + See Also + -------- + str.capitalize + + Examples + -------- + >>> import numpy as np + >>> c = np.array(['a1b2','1b2a','b2a1','2a1b'],'S4'); c + array(['a1b2', '1b2a', 'b2a1', '2a1b'], + dtype='|S4') + >>> np.strings.capitalize(c) + array(['A1b2', '1b2a', 'B2a1', '2a1b'], + dtype='|S4') + + """ + a_arr = np.asarray(a) + return _vec_string(a_arr, a_arr.dtype, 'capitalize') + + +@set_module("numpy.strings") +def title(a): + """ + Return element-wise title cased version of string or unicode. + + Title case words start with uppercase characters, all remaining cased + characters are lowercase. + + Calls :meth:`str.title` element-wise. + + For 8-bit strings, this method is locale-dependent. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + Input array. + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input types + + See Also + -------- + str.title + + Examples + -------- + >>> import numpy as np + >>> c=np.array(['a1b c','1b ca','b ca1','ca1b'],'S5'); c + array(['a1b c', '1b ca', 'b ca1', 'ca1b'], + dtype='|S5') + >>> np.strings.title(c) + array(['A1B C', '1B Ca', 'B Ca1', 'Ca1B'], + dtype='|S5') + + """ + a_arr = np.asarray(a) + return _vec_string(a_arr, a_arr.dtype, 'title') + + +@set_module("numpy.strings") +def replace(a, old, new, count=-1): + """ + For each element in ``a``, return a copy of the string with + occurrences of substring ``old`` replaced by ``new``. + + Parameters + ---------- + a : array_like, with ``bytes_`` or ``str_`` dtype + + old, new : array_like, with ``bytes_`` or ``str_`` dtype + + count : array_like, with ``int_`` dtype + If the optional argument ``count`` is given, only the first + ``count`` occurrences are replaced. + + Returns + ------- + out : ndarray + Output array of ``StringDType``, ``bytes_`` or ``str_`` dtype, + depending on input types + + See Also + -------- + str.replace + + Examples + -------- + >>> import numpy as np + >>> a = np.array(["That is a mango", "Monkeys eat mangos"]) + >>> np.strings.replace(a, 'mango', 'banana') + array(['That is a banana', 'Monkeys eat bananas'], dtype='>> a = np.array(["The dish is fresh", "This is it"]) + >>> np.strings.replace(a, 'is', 'was') + array(['The dwash was fresh', 'Thwas was it'], dtype='>> import numpy as np + >>> np.strings.join('-', 'osd') # doctest: +SKIP + array('o-s-d', dtype='>> np.strings.join(['-', '.'], ['ghc', 'osd']) # doctest: +SKIP + array(['g-h-c', 'o.s.d'], dtype='>> import numpy as np + >>> x = np.array("Numpy is nice!") + >>> np.strings.split(x, " ") # doctest: +SKIP + array(list(['Numpy', 'is', 'nice!']), dtype=object) # doctest: +SKIP + + >>> np.strings.split(x, " ", 1) # doctest: +SKIP + array(list(['Numpy', 'is nice!']), dtype=object) # doctest: +SKIP + + See Also + -------- + str.split, rsplit + + """ + # This will return an array of lists of different sizes, so we + # leave it as an object array + return _vec_string( + a, np.object_, 'split', [sep] + _clean_args(maxsplit)) + + +def _rsplit(a, sep=None, maxsplit=None): + """ + For each element in `a`, return a list of the words in the + string, using `sep` as the delimiter string. + + Calls :meth:`str.rsplit` element-wise. + + Except for splitting from the right, `rsplit` + behaves like `split`. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + sep : str or unicode, optional + If `sep` is not specified or None, any whitespace string + is a separator. + maxsplit : int, optional + If `maxsplit` is given, at most `maxsplit` splits are done, + the rightmost ones. + + Returns + ------- + out : ndarray + Array of list objects + + See Also + -------- + str.rsplit, split + + Examples + -------- + >>> import numpy as np + >>> a = np.array(['aAaAaA', 'abBABba']) + >>> np.strings.rsplit(a, 'A') # doctest: +SKIP + array([list(['a', 'a', 'a', '']), # doctest: +SKIP + list(['abB', 'Bba'])], dtype=object) # doctest: +SKIP + + """ + # This will return an array of lists of different sizes, so we + # leave it as an object array + return _vec_string( + a, np.object_, 'rsplit', [sep] + _clean_args(maxsplit)) + + +def _splitlines(a, keepends=None): + """ + For each element in `a`, return a list of the lines in the + element, breaking at line boundaries. + + Calls :meth:`str.splitlines` element-wise. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + + keepends : bool, optional + Line breaks are not included in the resulting list unless + keepends is given and true. + + Returns + ------- + out : ndarray + Array of list objects + + See Also + -------- + str.splitlines + + Examples + -------- + >>> np.char.splitlines("first line\\nsecond line") + array(list(['first line', 'second line']), dtype=object) + >>> a = np.array(["first\\nsecond", "third\\nfourth"]) + >>> np.char.splitlines(a) + array([list(['first', 'second']), list(['third', 'fourth'])], dtype=object) + + """ + return _vec_string( + a, np.object_, 'splitlines', _clean_args(keepends)) + + +@set_module("numpy.strings") +def partition(a, sep): + """ + Partition each element in ``a`` around ``sep``. + + For each element in ``a``, split the element at the first + occurrence of ``sep``, and return a 3-tuple containing the part + before the separator, the separator itself, and the part after + the separator. If the separator is not found, the first item of + the tuple will contain the whole string, and the second and third + ones will be the empty string. + + Parameters + ---------- + a : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + Input array + sep : array-like, with ``StringDType``, ``bytes_``, or ``str_`` dtype + Separator to split each string element in ``a``. + + Returns + ------- + out : 3-tuple: + - array with ``StringDType``, ``bytes_`` or ``str_`` dtype with the + part before the separator + - array with ``StringDType``, ``bytes_`` or ``str_`` dtype with the + separator + - array with ``StringDType``, ``bytes_`` or ``str_`` dtype with the + part after the separator + + See Also + -------- + str.partition + + Examples + -------- + >>> import numpy as np + >>> x = np.array(["Numpy is nice!"]) + >>> np.strings.partition(x, " ") + (array(['Numpy'], dtype='>> import numpy as np + >>> a = np.array(['aAaAaA', ' aA ', 'abBABba']) + >>> np.strings.rpartition(a, 'A') + (array(['aAaAa', ' a', 'abB'], dtype='>> import numpy as np + >>> a = np.array(['a1b c', '1bca', 'bca1']) + >>> table = a[0].maketrans('abc', '123') + >>> deletechars = ' ' + >>> np.char.translate(a, table, deletechars) + array(['112 3', '1231', '2311'], dtype=' NDArray[np.bool]: ... +@overload +def equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def not_equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def not_equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def not_equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def greater_equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def greater_equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def greater_equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def less_equal(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def less_equal(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def less_equal(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def greater(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def greater(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def greater(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def less(x1: U_co, x2: U_co) -> NDArray[np.bool]: ... +@overload +def less(x1: S_co, x2: S_co) -> NDArray[np.bool]: ... +@overload +def less(x1: T_co, x2: T_co) -> NDArray[np.bool]: ... + +@overload +def add(x1: U_co, x2: U_co) -> NDArray[np.str_]: ... +@overload +def add(x1: S_co, x2: S_co) -> NDArray[np.bytes_]: ... +@overload +def add(x1: _StringDTypeSupportsArray, x2: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def add(x1: T_co, x2: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def multiply(a: U_co, i: i_co) -> NDArray[np.str_]: ... +@overload +def multiply(a: S_co, i: i_co) -> NDArray[np.bytes_]: ... +@overload +def multiply(a: _StringDTypeSupportsArray, i: i_co) -> _StringDTypeArray: ... +@overload +def multiply(a: T_co, i: i_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def mod(a: U_co, value: object) -> NDArray[np.str_]: ... +@overload +def mod(a: S_co, value: object) -> NDArray[np.bytes_]: ... +@overload +def mod(a: _StringDTypeSupportsArray, value: object) -> _StringDTypeArray: ... +@overload +def mod(a: T_co, value: object) -> _StringDTypeOrUnicodeArray: ... + +def isalpha(x: UST_co) -> NDArray[np.bool]: ... +def isalnum(a: UST_co) -> NDArray[np.bool]: ... +def isdigit(x: UST_co) -> NDArray[np.bool]: ... +def isspace(x: UST_co) -> NDArray[np.bool]: ... +def isdecimal(x: U_co | T_co) -> NDArray[np.bool]: ... +def isnumeric(x: U_co | T_co) -> NDArray[np.bool]: ... +def islower(a: UST_co) -> NDArray[np.bool]: ... +def istitle(a: UST_co) -> NDArray[np.bool]: ... +def isupper(a: UST_co) -> NDArray[np.bool]: ... + +def str_len(x: UST_co) -> NDArray[np.int_]: ... + +@overload +def find( + a: U_co, + sub: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def find( + a: S_co, + sub: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def find( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def rfind( + a: U_co, + sub: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def rfind( + a: S_co, + sub: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def rfind( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def index( + a: U_co, + sub: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def index( + a: S_co, + sub: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def index( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def rindex( + a: U_co, + sub: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def rindex( + a: S_co, + sub: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def rindex( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def count( + a: U_co, + sub: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def count( + a: S_co, + sub: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... +@overload +def count( + a: T_co, + sub: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.int_]: ... + +@overload +def startswith( + a: U_co, + prefix: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... +@overload +def startswith( + a: S_co, + prefix: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... +@overload +def startswith( + a: T_co, + prefix: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... + +@overload +def endswith( + a: U_co, + suffix: U_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... +@overload +def endswith( + a: S_co, + suffix: S_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... +@overload +def endswith( + a: T_co, + suffix: T_co, + start: i_co = ..., + end: i_co | None = ..., +) -> NDArray[np.bool]: ... + +def decode( + a: S_co, + encoding: str | None = None, + errors: str | None = None, +) -> NDArray[np.str_]: ... +def encode( + a: U_co | T_co, + encoding: str | None = None, + errors: str | None = None, +) -> NDArray[np.bytes_]: ... + +@overload +def expandtabs(a: U_co, tabsize: i_co = ...) -> NDArray[np.str_]: ... +@overload +def expandtabs(a: S_co, tabsize: i_co = ...) -> NDArray[np.bytes_]: ... +@overload +def expandtabs(a: _StringDTypeSupportsArray, tabsize: i_co = ...) -> _StringDTypeArray: ... +@overload +def expandtabs(a: T_co, tabsize: i_co = ...) -> _StringDTypeOrUnicodeArray: ... + +@overload +def center(a: U_co, width: i_co, fillchar: UST_co = " ") -> NDArray[np.str_]: ... +@overload +def center(a: S_co, width: i_co, fillchar: UST_co = " ") -> NDArray[np.bytes_]: ... +@overload +def center(a: _StringDTypeSupportsArray, width: i_co, fillchar: UST_co = " ") -> _StringDTypeArray: ... +@overload +def center(a: T_co, width: i_co, fillchar: UST_co = " ") -> _StringDTypeOrUnicodeArray: ... + +@overload +def ljust(a: U_co, width: i_co, fillchar: UST_co = " ") -> NDArray[np.str_]: ... +@overload +def ljust(a: S_co, width: i_co, fillchar: UST_co = " ") -> NDArray[np.bytes_]: ... +@overload +def ljust(a: _StringDTypeSupportsArray, width: i_co, fillchar: UST_co = " ") -> _StringDTypeArray: ... +@overload +def ljust(a: T_co, width: i_co, fillchar: UST_co = " ") -> _StringDTypeOrUnicodeArray: ... + +@overload +def rjust(a: U_co, width: i_co, fillchar: UST_co = " ") -> NDArray[np.str_]: ... +@overload +def rjust(a: S_co, width: i_co, fillchar: UST_co = " ") -> NDArray[np.bytes_]: ... +@overload +def rjust(a: _StringDTypeSupportsArray, width: i_co, fillchar: UST_co = " ") -> _StringDTypeArray: ... +@overload +def rjust(a: T_co, width: i_co, fillchar: UST_co = " ") -> _StringDTypeOrUnicodeArray: ... + +@overload +def lstrip(a: U_co, chars: U_co | None = None) -> NDArray[np.str_]: ... +@overload +def lstrip(a: S_co, chars: S_co | None = None) -> NDArray[np.bytes_]: ... +@overload +def lstrip(a: _StringDTypeSupportsArray, chars: T_co | None = None) -> _StringDTypeArray: ... +@overload +def lstrip(a: T_co, chars: T_co | None = None) -> _StringDTypeOrUnicodeArray: ... + +@overload +def rstrip(a: U_co, chars: U_co | None = None) -> NDArray[np.str_]: ... +@overload +def rstrip(a: S_co, chars: S_co | None = None) -> NDArray[np.bytes_]: ... +@overload +def rstrip(a: _StringDTypeSupportsArray, chars: T_co | None = None) -> _StringDTypeArray: ... +@overload +def rstrip(a: T_co, chars: T_co | None = None) -> _StringDTypeOrUnicodeArray: ... + +@overload +def strip(a: U_co, chars: U_co | None = None) -> NDArray[np.str_]: ... +@overload +def strip(a: S_co, chars: S_co | None = None) -> NDArray[np.bytes_]: ... +@overload +def strip(a: _StringDTypeSupportsArray, chars: T_co | None = None) -> _StringDTypeArray: ... +@overload +def strip(a: T_co, chars: T_co | None = None) -> _StringDTypeOrUnicodeArray: ... + +@overload +def zfill(a: U_co, width: i_co) -> NDArray[np.str_]: ... +@overload +def zfill(a: S_co, width: i_co) -> NDArray[np.bytes_]: ... +@overload +def zfill(a: _StringDTypeSupportsArray, width: i_co) -> _StringDTypeArray: ... +@overload +def zfill(a: T_co, width: i_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def upper(a: U_co) -> NDArray[np.str_]: ... +@overload +def upper(a: S_co) -> NDArray[np.bytes_]: ... +@overload +def upper(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def upper(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def lower(a: U_co) -> NDArray[np.str_]: ... +@overload +def lower(a: S_co) -> NDArray[np.bytes_]: ... +@overload +def lower(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def lower(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def swapcase(a: U_co) -> NDArray[np.str_]: ... +@overload +def swapcase(a: S_co) -> NDArray[np.bytes_]: ... +@overload +def swapcase(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def swapcase(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def capitalize(a: U_co) -> NDArray[np.str_]: ... +@overload +def capitalize(a: S_co) -> NDArray[np.bytes_]: ... +@overload +def capitalize(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def capitalize(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def title(a: U_co) -> NDArray[np.str_]: ... +@overload +def title(a: S_co) -> NDArray[np.bytes_]: ... +@overload +def title(a: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def title(a: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def replace( + a: U_co, + old: U_co, + new: U_co, + count: i_co = ..., +) -> NDArray[np.str_]: ... +@overload +def replace( + a: S_co, + old: S_co, + new: S_co, + count: i_co = ..., +) -> NDArray[np.bytes_]: ... +@overload +def replace( + a: _StringDTypeSupportsArray, + old: _StringDTypeSupportsArray, + new: _StringDTypeSupportsArray, + count: i_co = ..., +) -> _StringDTypeArray: ... +@overload +def replace( + a: T_co, + old: T_co, + new: T_co, + count: i_co = ..., +) -> _StringDTypeOrUnicodeArray: ... + +@overload +def partition(a: U_co, sep: U_co) -> NDArray[np.str_]: ... +@overload +def partition(a: S_co, sep: S_co) -> NDArray[np.bytes_]: ... +@overload +def partition(a: _StringDTypeSupportsArray, sep: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def partition(a: T_co, sep: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def rpartition(a: U_co, sep: U_co) -> NDArray[np.str_]: ... +@overload +def rpartition(a: S_co, sep: S_co) -> NDArray[np.bytes_]: ... +@overload +def rpartition(a: _StringDTypeSupportsArray, sep: _StringDTypeSupportsArray) -> _StringDTypeArray: ... +@overload +def rpartition(a: T_co, sep: T_co) -> _StringDTypeOrUnicodeArray: ... + +@overload +def translate( + a: U_co, + table: str, + deletechars: str | None = None, +) -> NDArray[np.str_]: ... +@overload +def translate( + a: S_co, + table: str, + deletechars: str | None = None, +) -> NDArray[np.bytes_]: ... +@overload +def translate( + a: _StringDTypeSupportsArray, + table: str, + deletechars: str | None = None, +) -> _StringDTypeArray: ... +@overload +def translate( + a: T_co, + table: str, + deletechars: str | None = None, +) -> _StringDTypeOrUnicodeArray: ... diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/_locales.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/_locales.py new file mode 100644 index 0000000000000000000000000000000000000000..b6c3bd5d2a3ef8eb36e71f5e946fe2a002e49243 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/_locales.py @@ -0,0 +1,72 @@ +"""Provide class for testing in French locale + +""" +import sys +import locale + +import pytest + +__ALL__ = ['CommaDecimalPointLocale'] + + +def find_comma_decimal_point_locale(): + """See if platform has a decimal point as comma locale. + + Find a locale that uses a comma instead of a period as the + decimal point. + + Returns + ------- + old_locale: str + Locale when the function was called. + new_locale: {str, None) + First French locale found, None if none found. + + """ + if sys.platform == 'win32': + locales = ['FRENCH'] + else: + locales = ['fr_FR', 'fr_FR.UTF-8', 'fi_FI', 'fi_FI.UTF-8'] + + old_locale = locale.getlocale(locale.LC_NUMERIC) + new_locale = None + try: + for loc in locales: + try: + locale.setlocale(locale.LC_NUMERIC, loc) + new_locale = loc + break + except locale.Error: + pass + finally: + locale.setlocale(locale.LC_NUMERIC, locale=old_locale) + return old_locale, new_locale + + +class CommaDecimalPointLocale: + """Sets LC_NUMERIC to a locale with comma as decimal point. + + Classes derived from this class have setup and teardown methods that run + tests with locale.LC_NUMERIC set to a locale where commas (',') are used as + the decimal point instead of periods ('.'). On exit the locale is restored + to the initial locale. It also serves as context manager with the same + effect. If no such locale is available, the test is skipped. + + """ + (cur_locale, tst_locale) = find_comma_decimal_point_locale() + + def setup_method(self): + if self.tst_locale is None: + pytest.skip("No French locale available") + locale.setlocale(locale.LC_NUMERIC, locale=self.tst_locale) + + def teardown_method(self): + locale.setlocale(locale.LC_NUMERIC, locale=self.cur_locale) + + def __enter__(self): + if self.tst_locale is None: + pytest.skip("No French locale available") + locale.setlocale(locale.LC_NUMERIC, locale=self.tst_locale) + + def __exit__(self, type, value, traceback): + locale.setlocale(locale.LC_NUMERIC, locale=self.cur_locale) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/_natype.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/_natype.py new file mode 100644 index 0000000000000000000000000000000000000000..d5ab47c3564e415126212f0f58c25e12292495fa --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/_natype.py @@ -0,0 +1,198 @@ +# Vendored implementation of pandas.NA, adapted from pandas/_libs/missing.pyx +# +# This is vendored to avoid adding pandas as a test dependency. + +__all__ = ["pd_NA"] + +import numbers + +import numpy as np + +def _create_binary_propagating_op(name, is_divmod=False): + is_cmp = name.strip("_") in ["eq", "ne", "le", "lt", "ge", "gt"] + + def method(self, other): + if ( + other is pd_NA + or isinstance(other, (str, bytes)) + or isinstance(other, (numbers.Number, np.bool)) + or isinstance(other, np.ndarray) + and not other.shape + ): + # Need the other.shape clause to handle NumPy scalars, + # since we do a setitem on `out` below, which + # won't work for NumPy scalars. + if is_divmod: + return pd_NA, pd_NA + else: + return pd_NA + + elif isinstance(other, np.ndarray): + out = np.empty(other.shape, dtype=object) + out[:] = pd_NA + + if is_divmod: + return out, out.copy() + else: + return out + + elif is_cmp and isinstance(other, (np.datetime64, np.timedelta64)): + return pd_NA + + elif isinstance(other, np.datetime64): + if name in ["__sub__", "__rsub__"]: + return pd_NA + + elif isinstance(other, np.timedelta64): + if name in ["__sub__", "__rsub__", "__add__", "__radd__"]: + return pd_NA + + return NotImplemented + + method.__name__ = name + return method + + +def _create_unary_propagating_op(name: str): + def method(self): + return pd_NA + + method.__name__ = name + return method + + +class NAType: + def __repr__(self) -> str: + return "" + + def __format__(self, format_spec) -> str: + try: + return self.__repr__().__format__(format_spec) + except ValueError: + return self.__repr__() + + def __bool__(self): + raise TypeError("boolean value of NA is ambiguous") + + def __hash__(self): + exponent = 31 if is_32bit else 61 + return 2**exponent - 1 + + def __reduce__(self): + return "pd_NA" + + # Binary arithmetic and comparison ops -> propagate + + __add__ = _create_binary_propagating_op("__add__") + __radd__ = _create_binary_propagating_op("__radd__") + __sub__ = _create_binary_propagating_op("__sub__") + __rsub__ = _create_binary_propagating_op("__rsub__") + __mul__ = _create_binary_propagating_op("__mul__") + __rmul__ = _create_binary_propagating_op("__rmul__") + __matmul__ = _create_binary_propagating_op("__matmul__") + __rmatmul__ = _create_binary_propagating_op("__rmatmul__") + __truediv__ = _create_binary_propagating_op("__truediv__") + __rtruediv__ = _create_binary_propagating_op("__rtruediv__") + __floordiv__ = _create_binary_propagating_op("__floordiv__") + __rfloordiv__ = _create_binary_propagating_op("__rfloordiv__") + __mod__ = _create_binary_propagating_op("__mod__") + __rmod__ = _create_binary_propagating_op("__rmod__") + __divmod__ = _create_binary_propagating_op("__divmod__", is_divmod=True) + __rdivmod__ = _create_binary_propagating_op("__rdivmod__", is_divmod=True) + # __lshift__ and __rshift__ are not implemented + + __eq__ = _create_binary_propagating_op("__eq__") + __ne__ = _create_binary_propagating_op("__ne__") + __le__ = _create_binary_propagating_op("__le__") + __lt__ = _create_binary_propagating_op("__lt__") + __gt__ = _create_binary_propagating_op("__gt__") + __ge__ = _create_binary_propagating_op("__ge__") + + # Unary ops + + __neg__ = _create_unary_propagating_op("__neg__") + __pos__ = _create_unary_propagating_op("__pos__") + __abs__ = _create_unary_propagating_op("__abs__") + __invert__ = _create_unary_propagating_op("__invert__") + + # pow has special + def __pow__(self, other): + if other is pd_NA: + return pd_NA + elif isinstance(other, (numbers.Number, np.bool)): + if other == 0: + # returning positive is correct for +/- 0. + return type(other)(1) + else: + return pd_NA + elif util.is_array(other): + return np.where(other == 0, other.dtype.type(1), pd_NA) + + return NotImplemented + + def __rpow__(self, other): + if other is pd_NA: + return pd_NA + elif isinstance(other, (numbers.Number, np.bool)): + if other == 1: + return other + else: + return pd_NA + elif util.is_array(other): + return np.where(other == 1, other, pd_NA) + return NotImplemented + + # Logical ops using Kleene logic + + def __and__(self, other): + if other is False: + return False + elif other is True or other is pd_NA: + return pd_NA + return NotImplemented + + __rand__ = __and__ + + def __or__(self, other): + if other is True: + return True + elif other is False or other is pd_NA: + return pd_NA + return NotImplemented + + __ror__ = __or__ + + def __xor__(self, other): + if other is False or other is True or other is pd_NA: + return pd_NA + return NotImplemented + + __rxor__ = __xor__ + + __array_priority__ = 1000 + _HANDLED_TYPES = (np.ndarray, numbers.Number, str, np.bool) + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + types = self._HANDLED_TYPES + (NAType,) + for x in inputs: + if not isinstance(x, types): + return NotImplemented + + if method != "__call__": + raise ValueError(f"ufunc method '{method}' not supported for NA") + result = maybe_dispatch_ufunc_to_dunder_op( + self, ufunc, method, *inputs, **kwargs + ) + if result is NotImplemented: + # For a NumPy ufunc that's not a binop, like np.logaddexp + index = [i for i, x in enumerate(inputs) if x is pd_NA][0] + result = np.broadcast_arrays(*inputs)[index] + if result.ndim == 0: + result = result.item() + if ufunc.nout > 1: + result = (pd_NA,) * ufunc.nout + + return result + + +pd_NA = NAType() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/astype_copy.pkl b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/astype_copy.pkl new file mode 100644 index 0000000000000000000000000000000000000000..45694ae001c4a103365ff9fd5ae2da0dba3c11f6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/astype_copy.pkl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9564b309cbf3441ff0a6e4468fddaca46230fab34f15c77d87025a455bdf59d9 +size 716 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/generate_umath_validation_data.cpp b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/generate_umath_validation_data.cpp new file mode 100644 index 0000000000000000000000000000000000000000..a63fbe7b7e88ccaebac26e576c26049db7971bf9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/generate_umath_validation_data.cpp @@ -0,0 +1,170 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +struct ufunc { + std::string name; + double (*f32func)(double); + long double (*f64func)(long double); + float f32ulp; + float f64ulp; +}; + +template +T +RandomFloat(T a, T b) +{ + T random = ((T)rand()) / (T)RAND_MAX; + T diff = b - a; + T r = random * diff; + return a + r; +} + +template +void +append_random_array(std::vector &arr, T min, T max, size_t N) +{ + for (size_t ii = 0; ii < N; ++ii) + arr.emplace_back(RandomFloat(min, max)); +} + +template +std::vector +computeTrueVal(const std::vector &in, T2 (*mathfunc)(T2)) +{ + std::vector out; + for (T1 elem : in) { + T2 elem_d = (T2)elem; + T1 out_elem = (T1)mathfunc(elem_d); + out.emplace_back(out_elem); + } + return out; +} + +/* + * FP range: + * [-inf, -maxflt, -1., -minflt, -minden, 0., minden, minflt, 1., maxflt, inf] + */ + +#define MINDEN std::numeric_limits::denorm_min() +#define MINFLT std::numeric_limits::min() +#define MAXFLT std::numeric_limits::max() +#define INF std::numeric_limits::infinity() +#define qNAN std::numeric_limits::quiet_NaN() +#define sNAN std::numeric_limits::signaling_NaN() + +template +std::vector +generate_input_vector(std::string func) +{ + std::vector input = {MINDEN, -MINDEN, MINFLT, -MINFLT, MAXFLT, + -MAXFLT, INF, -INF, qNAN, sNAN, + -1.0, 1.0, 0.0, -0.0}; + + // [-1.0, 1.0] + if ((func == "arcsin") || (func == "arccos") || (func == "arctanh")) { + append_random_array(input, -1.0, 1.0, 700); + } + // (0.0, INF] + else if ((func == "log2") || (func == "log10")) { + append_random_array(input, 0.0, 1.0, 200); + append_random_array(input, MINDEN, MINFLT, 200); + append_random_array(input, MINFLT, 1.0, 200); + append_random_array(input, 1.0, MAXFLT, 200); + } + // (-1.0, INF] + else if (func == "log1p") { + append_random_array(input, -1.0, 1.0, 200); + append_random_array(input, -MINFLT, -MINDEN, 100); + append_random_array(input, -1.0, -MINFLT, 100); + append_random_array(input, MINDEN, MINFLT, 100); + append_random_array(input, MINFLT, 1.0, 100); + append_random_array(input, 1.0, MAXFLT, 100); + } + // [1.0, INF] + else if (func == "arccosh") { + append_random_array(input, 1.0, 2.0, 400); + append_random_array(input, 2.0, MAXFLT, 300); + } + // [-INF, INF] + else { + append_random_array(input, -1.0, 1.0, 100); + append_random_array(input, MINDEN, MINFLT, 100); + append_random_array(input, -MINFLT, -MINDEN, 100); + append_random_array(input, MINFLT, 1.0, 100); + append_random_array(input, -1.0, -MINFLT, 100); + append_random_array(input, 1.0, MAXFLT, 100); + append_random_array(input, -MAXFLT, -100.0, 100); + } + + std::random_shuffle(input.begin(), input.end()); + return input; +} + +int +main() +{ + srand(42); + std::vector umathfunc = { + {"sin", sin, sin, 1.49, 1.00}, + {"cos", cos, cos, 1.49, 1.00}, + {"tan", tan, tan, 3.91, 1.00}, + {"arcsin", asin, asin, 3.12, 1.00}, + {"arccos", acos, acos, 2.1, 1.00}, + {"arctan", atan, atan, 2.3, 1.00}, + {"sinh", sinh, sinh, 1.55, 1.00}, + {"cosh", cosh, cosh, 2.48, 1.00}, + {"tanh", tanh, tanh, 1.38, 2.00}, + {"arcsinh", asinh, asinh, 1.01, 1.00}, + {"arccosh", acosh, acosh, 1.16, 1.00}, + {"arctanh", atanh, atanh, 1.45, 1.00}, + {"cbrt", cbrt, cbrt, 1.94, 2.00}, + //{"exp",exp,exp,3.76,1.00}, + {"exp2", exp2, exp2, 1.01, 1.00}, + {"expm1", expm1, expm1, 2.62, 1.00}, + //{"log",log,log,1.84,1.00}, + {"log10", log10, log10, 3.5, 1.00}, + {"log1p", log1p, log1p, 1.96, 1.0}, + {"log2", log2, log2, 2.12, 1.00}, + }; + + for (int ii = 0; ii < umathfunc.size(); ++ii) { + // ignore sin/cos + if ((umathfunc[ii].name != "sin") && (umathfunc[ii].name != "cos")) { + std::string fileName = + "umath-validation-set-" + umathfunc[ii].name + ".csv"; + std::ofstream txtOut; + txtOut.open(fileName, std::ofstream::trunc); + txtOut << "dtype,input,output,ulperrortol" << std::endl; + + // Single Precision + auto f32in = generate_input_vector(umathfunc[ii].name); + auto f32out = computeTrueVal(f32in, + umathfunc[ii].f32func); + for (int jj = 0; jj < f32in.size(); ++jj) { + txtOut << "np.float32" << std::hex << ",0x" + << *reinterpret_cast(&f32in[jj]) << ",0x" + << *reinterpret_cast(&f32out[jj]) << "," + << ceil(umathfunc[ii].f32ulp) << std::endl; + } + + // Double Precision + auto f64in = generate_input_vector(umathfunc[ii].name); + auto f64out = computeTrueVal( + f64in, umathfunc[ii].f64func); + for (int jj = 0; jj < f64in.size(); ++jj) { + txtOut << "np.float64" << std::hex << ",0x" + << *reinterpret_cast(&f64in[jj]) << ",0x" + << *reinterpret_cast(&f64out[jj]) << "," + << ceil(umathfunc[ii].f64ulp) << std::endl; + } + txtOut.close(); + } + } + return 0; +} diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/recarray_from_file.fits b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/recarray_from_file.fits new file mode 100644 index 0000000000000000000000000000000000000000..ca48ee85153645a7510e201d574e9b119c089dce Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/recarray_from_file.fits differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/umath-validation-set-arccos.csv b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/umath-validation-set-arccos.csv new file mode 100644 index 0000000000000000000000000000000000000000..da33e9c56d7026a73b12232008a6c334c36c0683 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/data/umath-validation-set-arccos.csv @@ -0,0 +1,1429 @@ +dtype,input,output,ulperrortol +np.float32,0xbddd7f50,0x3fd6eec2,3 +np.float32,0xbe32a20c,0x3fdf8182,3 +np.float32,0xbf607c09,0x4028f84f,3 +np.float32,0x3f25d906,0x3f5db544,3 +np.float32,0x3f01cec8,0x3f84febf,3 +np.float32,0x3f1d5c6e,0x3f68a735,3 +np.float32,0xbf0cab89,0x4009c36d,3 +np.float32,0xbf176b40,0x400d0941,3 +np.float32,0x3f3248b2,0x3f4ce6d4,3 +np.float32,0x3f390b48,0x3f434e0d,3 +np.float32,0xbe261698,0x3fddea43,3 +np.float32,0x3f0e1154,0x3f7b848b,3 +np.float32,0xbf379a3c,0x4017b764,3 +np.float32,0xbeda6f2c,0x4000bd62,3 +np.float32,0xbf6a0c3f,0x402e5d5a,3 +np.float32,0x3ef1d700,0x3f8a17b7,3 +np.float32,0xbf6f4f65,0x4031d30d,3 +np.float32,0x3f2c9eee,0x3f54adfd,3 +np.float32,0x3f3cfb18,0x3f3d8a1e,3 +np.float32,0x3ba80800,0x3fc867d2,3 +np.float32,0x3e723b08,0x3faa7e4d,3 +np.float32,0xbf65820f,0x402bb054,3 +np.float32,0xbee64e7a,0x40026410,3 +np.float32,0x3cb15140,0x3fc64a87,3 +np.float32,0x3f193660,0x3f6ddf2a,3 +np.float32,0xbf0e5b52,0x400a44f7,3 +np.float32,0x3ed55f14,0x3f920a4b,3 +np.float32,0x3dd11a80,0x3fbbf85c,3 +np.float32,0xbf4f5c4b,0x4020f4f9,3 +np.float32,0x3f787532,0x3e792e87,3 +np.float32,0x3f40e6ac,0x3f37a74f,3 +np.float32,0x3f1c1318,0x3f6a47b6,3 +np.float32,0xbe3c48d8,0x3fe0bb70,3 +np.float32,0xbe94d4bc,0x3feed08e,3 +np.float32,0xbe5c3688,0x3fe4ce26,3 +np.float32,0xbf6fe026,0x403239cb,3 +np.float32,0x3ea5983c,0x3f9ee7bf,3 +np.float32,0x3f1471e6,0x3f73c5bb,3 +np.float32,0x3f0e2622,0x3f7b6b87,3 +np.float32,0xbf597180,0x40257ad1,3 +np.float32,0xbeb5321c,0x3ff75d34,3 +np.float32,0x3f5afcd2,0x3f0b6012,3 +np.float32,0xbef2ff88,0x40042e14,3 +np.float32,0xbedc747e,0x400104f5,3 +np.float32,0xbee0c2f4,0x40019dfc,3 +np.float32,0xbf152cd8,0x400c57dc,3 +np.float32,0xbf6cf9e2,0x40303bbe,3 +np.float32,0x3ed9cd74,0x3f90d1a1,3 +np.float32,0xbf754406,0x4036767f,3 +np.float32,0x3f59c5c2,0x3f0db42f,3 +np.float32,0x3f2eefd8,0x3f518684,3 +np.float32,0xbf156bf9,0x400c6b49,3 +np.float32,0xbd550790,0x3fcfb8dc,3 +np.float32,0x3ede58fc,0x3f8f8f77,3 +np.float32,0xbf00ac19,0x40063c4b,3 +np.float32,0x3f4d25ba,0x3f24280e,3 +np.float32,0xbe9568be,0x3feef73c,3 +np.float32,0x3f67d154,0x3ee05547,3 +np.float32,0x3f617226,0x3efcb4f4,3 +np.float32,0xbf3ab41a,0x4018d6cc,3 +np.float32,0xbf3186fe,0x401592cd,3 +np.float32,0x3de3ba50,0x3fbacca9,3 +np.float32,0x3e789f98,0x3fa9ab97,3 +np.float32,0x3f016e08,0x3f8536d8,3 +np.float32,0x3e8b618c,0x3fa5c571,3 +np.float32,0x3eff97bc,0x3f8628a9,3 +np.float32,0xbf6729f0,0x402ca32f,3 +np.float32,0xbebec146,0x3ff9eddc,3 +np.float32,0x3ddb2e60,0x3fbb563a,3 +np.float32,0x3caa8e40,0x3fc66595,3 +np.float32,0xbf5973f2,0x40257bfa,3 +np.float32,0xbdd82c70,0x3fd69916,3 +np.float32,0xbedf4c82,0x400169ef,3 +np.float32,0x3ef8f22c,0x3f881184,3 +np.float32,0xbf1d74d4,0x400eedc9,3 +np.float32,0x3f2e10a6,0x3f52b790,3 +np.float32,0xbf08ecc0,0x4008a628,3 +np.float32,0x3ecb7db4,0x3f94be9f,3 +np.float32,0xbf052ded,0x40078bfc,3 +np.float32,0x3f2ee78a,0x3f5191e4,3 +np.float32,0xbf56f4e1,0x40245194,3 +np.float32,0x3f600a3e,0x3f014a25,3 +np.float32,0x3f3836f8,0x3f44808b,3 +np.float32,0x3ecabfbc,0x3f94f25c,3 +np.float32,0x3c70f500,0x3fc72dec,3 +np.float32,0x3f17c444,0x3f6fabf0,3 +np.float32,0xbf4c22a5,0x401f9a09,3 +np.float32,0xbe4205dc,0x3fe1765a,3 +np.float32,0x3ea49138,0x3f9f2d36,3 +np.float32,0xbece0082,0x3ffe106b,3 +np.float32,0xbe387578,0x3fe03eef,3 +np.float32,0xbf2b6466,0x40137a30,3 +np.float32,0xbe9dadb2,0x3ff12204,3 +np.float32,0xbf56b3f2,0x402433bb,3 +np.float32,0xbdf9b4d8,0x3fd8b51f,3 +np.float32,0x3f58a596,0x3f0fd4b4,3 +np.float32,0xbedf5748,0x40016b6e,3 +np.float32,0x3f446442,0x3f32476f,3 +np.float32,0x3f5be886,0x3f099658,3 +np.float32,0x3ea1e44c,0x3f9fe1de,3 +np.float32,0xbf11e9b8,0x400b585f,3 +np.float32,0xbf231f8f,0x4010befb,3 +np.float32,0xbf4395ea,0x401c2dd0,3 +np.float32,0x3e9e7784,0x3fa0c8a6,3 +np.float32,0xbe255184,0x3fddd14c,3 +np.float32,0x3f70d25e,0x3eb13148,3 +np.float32,0x3f220cdc,0x3f62a722,3 +np.float32,0xbd027bf0,0x3fcd23e7,3 +np.float32,0x3e4ef8b8,0x3faf02d2,3 +np.float32,0xbf76fc6b,0x40380728,3 +np.float32,0xbf57e761,0x4024c1cd,3 +np.float32,0x3ed4fc20,0x3f922580,3 +np.float32,0xbf09b64a,0x4008e1db,3 +np.float32,0x3f21ca62,0x3f62fcf5,3 +np.float32,0xbe55f610,0x3fe40170,3 +np.float32,0xbc0def80,0x3fca2bbb,3 +np.float32,0xbebc8764,0x3ff9547b,3 +np.float32,0x3ec1b200,0x3f9766d1,3 +np.float32,0xbf4ee44e,0x4020c1ee,3 +np.float32,0xbea85852,0x3ff3f22a,3 +np.float32,0xbf195c0c,0x400da3d3,3 +np.float32,0xbf754b5d,0x40367ce8,3 +np.float32,0xbdcbfe50,0x3fd5d52b,3 +np.float32,0xbf1adb87,0x400e1be3,3 +np.float32,0xbf6f8491,0x4031f898,3 +np.float32,0xbf6f9ae7,0x4032086e,3 +np.float32,0xbf52b3f0,0x40226790,3 +np.float32,0xbf698452,0x402e09f4,3 +np.float32,0xbf43dc9a,0x401c493a,3 +np.float32,0xbf165f7f,0x400cb664,3 +np.float32,0x3e635468,0x3fac682f,3 +np.float32,0xbe8cf2b6,0x3fecc28a,3 +np.float32,0x7f7fffff,0x7fc00000,3 +np.float32,0xbf4c6513,0x401fb597,3 +np.float32,0xbf02b8f8,0x4006d47e,3 +np.float32,0x3ed3759c,0x3f9290c8,3 +np.float32,0xbf2a7a5f,0x40132b98,3 +np.float32,0xbae65000,0x3fc9496f,3 +np.float32,0x3f65f5ea,0x3ee8ef07,3 +np.float32,0xbe7712fc,0x3fe84106,3 +np.float32,0xbb9ff700,0x3fc9afd2,3 +np.float32,0x3d8d87a0,0x3fc03592,3 +np.float32,0xbefc921c,0x40058c23,3 +np.float32,0xbf286566,0x401279d8,3 +np.float32,0x3f53857e,0x3f192eaf,3 +np.float32,0xbee9b0f4,0x4002dd90,3 +np.float32,0x3f4041f8,0x3f38a14a,3 +np.float32,0x3f54ea96,0x3f16b02d,3 +np.float32,0x3ea50ef8,0x3f9f0c01,3 +np.float32,0xbeaad2dc,0x3ff49a4a,3 +np.float32,0xbec428c8,0x3ffb636f,3 +np.float32,0xbda46178,0x3fd358c7,3 +np.float32,0xbefacfc4,0x40054b7f,3 +np.float32,0xbf7068f9,0x40329c85,3 +np.float32,0x3f70b850,0x3eb1caa7,3 +np.float32,0x7fa00000,0x7fe00000,3 +np.float32,0x80000000,0x3fc90fdb,3 +np.float32,0x3f68d5c8,0x3edb7cf3,3 +np.float32,0x3d9443d0,0x3fbfc98a,3 +np.float32,0xff7fffff,0x7fc00000,3 +np.float32,0xbeee7ba8,0x40038a5e,3 +np.float32,0xbf0aaaba,0x40092a73,3 +np.float32,0x3f36a4e8,0x3f46c0ee,3 +np.float32,0x3ed268e4,0x3f92da82,3 +np.float32,0xbee6002c,0x4002591b,3 +np.float32,0xbe8f2752,0x3fed5576,3 +np.float32,0x3f525912,0x3f1b40e0,3 +np.float32,0xbe8e151e,0x3fed0e16,3 +np.float32,0x1,0x3fc90fdb,3 +np.float32,0x3ee23b84,0x3f8e7ae1,3 +np.float32,0xbf5961ca,0x40257361,3 +np.float32,0x3f6bbca0,0x3ecd14cd,3 +np.float32,0x3e27b230,0x3fb4014d,3 +np.float32,0xbf183bb8,0x400d49fc,3 +np.float32,0x3f57759c,0x3f120b68,3 +np.float32,0xbd6994c0,0x3fd05d84,3 +np.float32,0xbf1dd684,0x400f0cc8,3 +np.float32,0xbececc1c,0x3ffe480a,3 +np.float32,0xbf48855f,0x401e206d,3 +np.float32,0x3f28c922,0x3f59d382,3 +np.float32,0xbf65c094,0x402bd3b0,3 +np.float32,0x3f657d42,0x3eeb11dd,3 +np.float32,0xbed32d4e,0x3fff7b15,3 +np.float32,0xbf31af02,0x4015a0b1,3 +np.float32,0x3d89eb00,0x3fc06f7f,3 +np.float32,0x3dac2830,0x3fbe4a17,3 +np.float32,0x3f7f7cb6,0x3d81a7df,3 +np.float32,0xbedbb570,0x4000ea82,3 +np.float32,0x3db37830,0x3fbdd4a8,3 +np.float32,0xbf376f48,0x4017a7fd,3 +np.float32,0x3f319f12,0x3f4dd2c9,3 +np.float32,0x7fc00000,0x7fc00000,3 +np.float32,0x3f1b4f70,0x3f6b3e31,3 +np.float32,0x3e33c880,0x3fb278d1,3 +np.float32,0x3f2796e0,0x3f5b69bd,3 +np.float32,0x3f4915d6,0x3f2ad4d0,3 +np.float32,0x3e4db120,0x3faf2ca0,3 +np.float32,0x3ef03dd4,0x3f8a8ba9,3 +np.float32,0x3e96ca88,0x3fa2cbf7,3 +np.float32,0xbeb136ce,0x3ff64d2b,3 +np.float32,0xbf2f3938,0x4014c75e,3 +np.float32,0x3f769dde,0x3e8b0d76,3 +np.float32,0x3f67cec8,0x3ee06148,3 +np.float32,0x3f0a1ade,0x3f80204e,3 +np.float32,0x3e4b9718,0x3faf7144,3 +np.float32,0x3cccb480,0x3fc5dcf3,3 +np.float32,0x3caeb740,0x3fc654f0,3 +np.float32,0x3f684e0e,0x3ede0678,3 +np.float32,0x3f0ba93c,0x3f7e6663,3 +np.float32,0xbf12bbc4,0x400b985e,3 +np.float32,0xbf2a8e1a,0x40133235,3 +np.float32,0x3f42029c,0x3f35f5c5,3 +np.float32,0x3eed1728,0x3f8b6f9c,3 +np.float32,0xbe5779ac,0x3fe432fd,3 +np.float32,0x3f6ed8b8,0x3ebc7e4b,3 +np.float32,0x3eea25b0,0x3f8c43c7,3 +np.float32,0x3f1988a4,0x3f6d786b,3 +np.float32,0xbe751674,0x3fe7ff8a,3 +np.float32,0xbe9f7418,0x3ff1997d,3 +np.float32,0x3dca11d0,0x3fbc6979,3 +np.float32,0x3f795226,0x3e6a6cab,3 +np.float32,0xbea780e0,0x3ff3b926,3 +np.float32,0xbed92770,0x4000901e,3 +np.float32,0xbf3e9f8c,0x401a49f8,3 +np.float32,0x3f0f7054,0x3f79ddb2,3 +np.float32,0x3a99d400,0x3fc8e966,3 +np.float32,0xbef082b0,0x4003d3c6,3 +np.float32,0xbf0d0790,0x4009defb,3 +np.float32,0xbf1649da,0x400cafb4,3 +np.float32,0xbea5aca8,0x3ff33d5c,3 +np.float32,0xbf4e1843,0x40206ba1,3 +np.float32,0xbe3d7d5c,0x3fe0e2ad,3 +np.float32,0xbf0e802d,0x400a500e,3 +np.float32,0xbf0de8f0,0x400a2295,3 +np.float32,0xbf3016ba,0x4015137e,3 +np.float32,0x3f36b1ea,0x3f46ae5d,3 +np.float32,0xbd27f170,0x3fce4fc7,3 +np.float32,0x3e96ec54,0x3fa2c31f,3 +np.float32,0x3eb4dfdc,0x3f9ad87d,3 +np.float32,0x3f5cac6c,0x3f0815cc,3 +np.float32,0xbf0489aa,0x40075bf1,3 +np.float32,0x3df010c0,0x3fba05f5,3 +np.float32,0xbf229f4a,0x4010956a,3 +np.float32,0x3f75e474,0x3e905a99,3 +np.float32,0xbcece6a0,0x3fccc397,3 +np.float32,0xbdb41528,0x3fd454e7,3 +np.float32,0x3ec8b2f8,0x3f958118,3 +np.float32,0x3f5eaa70,0x3f041a1d,3 +np.float32,0xbf32e1cc,0x40160b91,3 +np.float32,0xbe8e6026,0x3fed219c,3 +np.float32,0x3e6b3160,0x3fab65e3,3 +np.float32,0x3e6d7460,0x3fab1b81,3 +np.float32,0xbf13fbde,0x400bfa3b,3 +np.float32,0xbe8235ec,0x3fe9f9e3,3 +np.float32,0x3d71c4a0,0x3fc18096,3 +np.float32,0x3eb769d0,0x3f9a2aa0,3 +np.float32,0xbf68cb3b,0x402d99e4,3 +np.float32,0xbd917610,0x3fd22932,3 +np.float32,0x3d3cba60,0x3fc3297f,3 +np.float32,0xbf383cbe,0x4017f1cc,3 +np.float32,0xbeee96d0,0x40038e34,3 +np.float32,0x3ec89cb4,0x3f958725,3 +np.float32,0x3ebf92d8,0x3f97f95f,3 +np.float32,0x3f30f3da,0x3f4ec021,3 +np.float32,0xbd26b560,0x3fce45e4,3 +np.float32,0xbec0eb12,0x3ffa8330,3 +np.float32,0x3f6d592a,0x3ec4a6c1,3 +np.float32,0x3ea6d39c,0x3f9e9463,3 +np.float32,0x3e884184,0x3fa6951e,3 +np.float32,0x3ea566c4,0x3f9ef4d1,3 +np.float32,0x3f0c8f4c,0x3f7d5380,3 +np.float32,0x3f28e1ba,0x3f59b2cb,3 +np.float32,0x3f798538,0x3e66e1c3,3 +np.float32,0xbe2889b8,0x3fde39b8,3 +np.float32,0x3f3da05e,0x3f3c949c,3 +np.float32,0x3f24d700,0x3f5f073e,3 +np.float32,0xbe5b5768,0x3fe4b198,3 +np.float32,0xbed3b03a,0x3fff9f05,3 +np.float32,0x3e8a1c4c,0x3fa619eb,3 +np.float32,0xbf075d24,0x40083030,3 +np.float32,0x3f765648,0x3e8d1f52,3 +np.float32,0xbf70fc5e,0x403308bb,3 +np.float32,0x3f557ae8,0x3f15ab76,3 +np.float32,0x3f02f7ea,0x3f84521c,3 +np.float32,0x3f7ebbde,0x3dcbc5c5,3 +np.float32,0xbefbdfc6,0x40057285,3 +np.float32,0x3ec687ac,0x3f9617d9,3 +np.float32,0x3e4831c8,0x3fafe01b,3 +np.float32,0x3e25cde0,0x3fb43ea8,3 +np.float32,0x3e4f2ab8,0x3faefc70,3 +np.float32,0x3ea60ae4,0x3f9ec973,3 +np.float32,0xbf1ed55f,0x400f5dde,3 +np.float32,0xbf5ad4aa,0x40262479,3 +np.float32,0x3e8b3594,0x3fa5d0de,3 +np.float32,0x3f3a77aa,0x3f413c80,3 +np.float32,0xbf07512b,0x40082ca9,3 +np.float32,0x3f33d990,0x3f4ab5e5,3 +np.float32,0x3f521556,0x3f1bb78f,3 +np.float32,0xbecf6036,0x3ffe7086,3 +np.float32,0x3db91bd0,0x3fbd7a11,3 +np.float32,0x3ef63a74,0x3f88d839,3 +np.float32,0xbf2f1116,0x4014b99c,3 +np.float32,0xbf17fdc0,0x400d36b9,3 +np.float32,0xbe87df2c,0x3feb7117,3 +np.float32,0x80800000,0x3fc90fdb,3 +np.float32,0x3ee24c1c,0x3f8e7641,3 +np.float32,0x3f688dce,0x3edcd644,3 +np.float32,0xbf0f4e1c,0x400a8e1b,3 +np.float32,0x0,0x3fc90fdb,3 +np.float32,0x3f786eba,0x3e7999d4,3 +np.float32,0xbf404f80,0x401aeca8,3 +np.float32,0xbe9ffb6a,0x3ff1bd18,3 +np.float32,0x3f146bfc,0x3f73ccfd,3 +np.float32,0xbe47d630,0x3fe233ee,3 +np.float32,0xbe95847c,0x3feefe7c,3 +np.float32,0xbf135df0,0x400bc9e5,3 +np.float32,0x3ea19f3c,0x3f9ff411,3 +np.float32,0x3f235e20,0x3f60f247,3 +np.float32,0xbec789ec,0x3ffc4def,3 +np.float32,0x3f04b656,0x3f834db6,3 +np.float32,0x3dfaf440,0x3fb95679,3 +np.float32,0xbe4a7f28,0x3fe28abe,3 +np.float32,0x3ed4850c,0x3f92463b,3 +np.float32,0x3ec4ba5c,0x3f9694dd,3 +np.float32,0xbce24ca0,0x3fcc992b,3 +np.float32,0xbf5b7c6e,0x402675a0,3 +np.float32,0xbea3ce2a,0x3ff2bf04,3 +np.float32,0x3db02c60,0x3fbe0998,3 +np.float32,0x3c47b780,0x3fc78069,3 +np.float32,0x3ed33b20,0x3f92a0d5,3 +np.float32,0xbf4556d7,0x401cdcde,3 +np.float32,0xbe1b6e28,0x3fdc90ec,3 +np.float32,0xbf3289b7,0x4015ecd0,3 +np.float32,0x3df3f240,0x3fb9c76d,3 +np.float32,0x3eefa7d0,0x3f8ab61d,3 +np.float32,0xbe945838,0x3feeb006,3 +np.float32,0xbf0b1386,0x400949a3,3 +np.float32,0x3f77e546,0x3e812cc1,3 +np.float32,0x3e804ba0,0x3fa8a480,3 +np.float32,0x3f43dcea,0x3f331a06,3 +np.float32,0x3eb87450,0x3f99e33c,3 +np.float32,0x3e5f4898,0x3facecea,3 +np.float32,0x3f646640,0x3eeff10e,3 +np.float32,0x3f1aa832,0x3f6c1051,3 +np.float32,0xbebf6bfa,0x3ffa1bdc,3 +np.float32,0xbb77f300,0x3fc98bd4,3 +np.float32,0x3f3587fe,0x3f485645,3 +np.float32,0x3ef85f34,0x3f883b8c,3 +np.float32,0x3f50e584,0x3f1dc82c,3 +np.float32,0x3f1d30a8,0x3f68deb0,3 +np.float32,0x3ee75a78,0x3f8d0c86,3 +np.float32,0x3f2c023a,0x3f5581e1,3 +np.float32,0xbf074e34,0x40082bca,3 +np.float32,0xbead71f0,0x3ff54c6d,3 +np.float32,0xbf39ed88,0x40188e69,3 +np.float32,0x3f5d2fe6,0x3f07118b,3 +np.float32,0xbf1f79f8,0x400f9267,3 +np.float32,0x3e900c58,0x3fa48e99,3 +np.float32,0xbf759cb2,0x4036c47b,3 +np.float32,0x3f63329c,0x3ef5359c,3 +np.float32,0xbf5d6755,0x40276709,3 +np.float32,0x3f2ce31c,0x3f54519a,3 +np.float32,0x7f800000,0x7fc00000,3 +np.float32,0x3f1bf50e,0x3f6a6d9a,3 +np.float32,0x3f258334,0x3f5e25d8,3 +np.float32,0xbf661a3f,0x402c06ac,3 +np.float32,0x3d1654c0,0x3fc45cef,3 +np.float32,0xbef14a36,0x4003f009,3 +np.float32,0xbf356051,0x4016ec3a,3 +np.float32,0x3f6ccc42,0x3ec79193,3 +np.float32,0xbf2fe3d6,0x401501f9,3 +np.float32,0x3deedc80,0x3fba195b,3 +np.float32,0x3f2e5a28,0x3f52533e,3 +np.float32,0x3e6b68b8,0x3fab5ec8,3 +np.float32,0x3e458240,0x3fb037b7,3 +np.float32,0xbf24bab0,0x401144cb,3 +np.float32,0x3f600f4c,0x3f013fb2,3 +np.float32,0x3f021a04,0x3f84d316,3 +np.float32,0x3f741732,0x3e9cc948,3 +np.float32,0x3f0788aa,0x3f81a5b0,3 +np.float32,0x3f28802c,0x3f5a347c,3 +np.float32,0x3c9eb400,0x3fc69500,3 +np.float32,0x3e5d11e8,0x3fad357a,3 +np.float32,0x3d921250,0x3fbfecb9,3 +np.float32,0x3f354866,0x3f48b066,3 +np.float32,0xbf72cf43,0x40346d84,3 +np.float32,0x3eecdbb8,0x3f8b805f,3 +np.float32,0xbee585d0,0x400247fd,3 +np.float32,0x3e3607a8,0x3fb22fc6,3 +np.float32,0xbf0cb7d6,0x4009c71c,3 +np.float32,0xbf56b230,0x402432ec,3 +np.float32,0xbf4ced02,0x401fee29,3 +np.float32,0xbf3a325c,0x4018a776,3 +np.float32,0x3ecae8bc,0x3f94e732,3 +np.float32,0xbe48c7e8,0x3fe252bd,3 +np.float32,0xbe175d7c,0x3fdc0d5b,3 +np.float32,0x3ea78dac,0x3f9e632d,3 +np.float32,0xbe7434a8,0x3fe7e279,3 +np.float32,0x3f1f9e02,0x3f65c7b9,3 +np.float32,0xbe150f2c,0x3fdbc2c2,3 +np.float32,0x3ee13480,0x3f8ec423,3 +np.float32,0x3ecb7d54,0x3f94beb9,3 +np.float32,0x3f1cef42,0x3f693181,3 +np.float32,0xbf1ec06a,0x400f5730,3 +np.float32,0xbe112acc,0x3fdb44e8,3 +np.float32,0xbe77b024,0x3fe85545,3 +np.float32,0x3ec86fe0,0x3f959353,3 +np.float32,0x3f36b326,0x3f46ac9a,3 +np.float32,0x3e581a70,0x3fadd829,3 +np.float32,0xbf032c0c,0x4006f5f9,3 +np.float32,0xbf43b1fd,0x401c38b1,3 +np.float32,0x3f3701b4,0x3f463c5c,3 +np.float32,0x3f1a995a,0x3f6c22f1,3 +np.float32,0xbf05de0b,0x4007bf97,3 +np.float32,0x3d4bd960,0x3fc2b063,3 +np.float32,0x3f0e1618,0x3f7b7ed0,3 +np.float32,0x3edfd420,0x3f8f2628,3 +np.float32,0xbf6662fe,0x402c3047,3 +np.float32,0x3ec0690c,0x3f97bf9b,3 +np.float32,0xbeaf4146,0x3ff5c7a0,3 +np.float32,0x3f5e7764,0x3f04816d,3 +np.float32,0xbedd192c,0x40011bc5,3 +np.float32,0x3eb76350,0x3f9a2c5e,3 +np.float32,0xbed8108c,0x400069a5,3 +np.float32,0xbe59f31c,0x3fe48401,3 +np.float32,0xbea3e1e6,0x3ff2c439,3 +np.float32,0x3e26d1f8,0x3fb41db5,3 +np.float32,0x3f3a0a7c,0x3f41dba5,3 +np.float32,0x3ebae068,0x3f993ce4,3 +np.float32,0x3f2d8e30,0x3f536942,3 +np.float32,0xbe838bbe,0x3fea5247,3 +np.float32,0x3ebe4420,0x3f98538f,3 +np.float32,0xbcc59b80,0x3fcc265c,3 +np.float32,0x3eebb5c8,0x3f8bd334,3 +np.float32,0xbafc3400,0x3fc94ee8,3 +np.float32,0xbf63ddc1,0x402ac683,3 +np.float32,0xbeabdf80,0x3ff4e18f,3 +np.float32,0x3ea863f0,0x3f9e2a78,3 +np.float32,0x3f45b292,0x3f303bc1,3 +np.float32,0xbe68aa60,0x3fe666bf,3 +np.float32,0x3eb9de18,0x3f998239,3 +np.float32,0xbf719d85,0x4033815e,3 +np.float32,0x3edef9a8,0x3f8f62db,3 +np.float32,0xbd7781c0,0x3fd0cd1e,3 +np.float32,0x3f0b3b90,0x3f7ee92a,3 +np.float32,0xbe3eb3b4,0x3fe10a27,3 +np.float32,0xbf31a4c4,0x40159d23,3 +np.float32,0x3e929434,0x3fa3e5b0,3 +np.float32,0xbeb1a90e,0x3ff66b9e,3 +np.float32,0xbeba9b5e,0x3ff8d048,3 +np.float32,0xbf272a84,0x4012119e,3 +np.float32,0x3f1ebbd0,0x3f66e889,3 +np.float32,0x3ed3cdc8,0x3f927893,3 +np.float32,0xbf50dfce,0x40219b58,3 +np.float32,0x3f0c02de,0x3f7dfb62,3 +np.float32,0xbf694de3,0x402de8d2,3 +np.float32,0xbeaeb13e,0x3ff5a14f,3 +np.float32,0xbf61aa7a,0x40299702,3 +np.float32,0xbf13d159,0x400bed35,3 +np.float32,0xbeecd034,0x40034e0b,3 +np.float32,0xbe50c2e8,0x3fe35761,3 +np.float32,0x3f714406,0x3eae8e57,3 +np.float32,0xbf1ca486,0x400eabd8,3 +np.float32,0x3f5858cc,0x3f106497,3 +np.float32,0x3f670288,0x3ee41c84,3 +np.float32,0xbf20bd2c,0x400ff9f5,3 +np.float32,0xbe29afd8,0x3fde5eff,3 +np.float32,0xbf635e6a,0x402a80f3,3 +np.float32,0x3e82b7b0,0x3fa80446,3 +np.float32,0x3e982e7c,0x3fa26ece,3 +np.float32,0x3d9f0e00,0x3fbf1c6a,3 +np.float32,0x3e8299b4,0x3fa80c07,3 +np.float32,0xbf0529c1,0x40078ac3,3 +np.float32,0xbf403b8a,0x401ae519,3 +np.float32,0xbe57e09c,0x3fe44027,3 +np.float32,0x3ea1c8f4,0x3f9fe913,3 +np.float32,0xbe216a94,0x3fdd52d0,3 +np.float32,0x3f59c442,0x3f0db709,3 +np.float32,0xbd636260,0x3fd02bdd,3 +np.float32,0xbdbbc788,0x3fd4d08d,3 +np.float32,0x3dd19560,0x3fbbf0a3,3 +np.float32,0x3f060ad4,0x3f828641,3 +np.float32,0x3b102e00,0x3fc8c7c4,3 +np.float32,0x3f42b3b8,0x3f34e5a6,3 +np.float32,0x3f0255ac,0x3f84b071,3 +np.float32,0xbf014898,0x40066996,3 +np.float32,0x3e004dc0,0x3fb8fb51,3 +np.float32,0xbf594ff8,0x40256af2,3 +np.float32,0x3efafddc,0x3f877b80,3 +np.float32,0xbf5f0780,0x40283899,3 +np.float32,0x3ee95e54,0x3f8c7bcc,3 +np.float32,0x3eba2f0c,0x3f996c80,3 +np.float32,0x3f37721c,0x3f459b68,3 +np.float32,0x3e2be780,0x3fb378bf,3 +np.float32,0x3e550270,0x3fae3d69,3 +np.float32,0x3e0f9500,0x3fb70e0a,3 +np.float32,0xbf51974a,0x4021eaf4,3 +np.float32,0x3f393832,0x3f430d05,3 +np.float32,0x3f3df16a,0x3f3c1bd8,3 +np.float32,0xbd662340,0x3fd041ed,3 +np.float32,0x3f7e8418,0x3ddc9fce,3 +np.float32,0xbf392734,0x40184672,3 +np.float32,0x3ee3b278,0x3f8e124e,3 +np.float32,0x3eed4808,0x3f8b61d2,3 +np.float32,0xbf6fccbd,0x40322beb,3 +np.float32,0x3e3ecdd0,0x3fb1123b,3 +np.float32,0x3f4419e0,0x3f32bb45,3 +np.float32,0x3f595e00,0x3f0e7914,3 +np.float32,0xbe8c1486,0x3fec88c6,3 +np.float32,0xbf800000,0x40490fdb,3 +np.float32,0xbdaf5020,0x3fd4084d,3 +np.float32,0xbf407660,0x401afb63,3 +np.float32,0x3f0c3aa8,0x3f7db8b8,3 +np.float32,0xbcdb5980,0x3fcc7d5b,3 +np.float32,0x3f4738d4,0x3f2dd1ed,3 +np.float32,0x3f4d7064,0x3f23ab14,3 +np.float32,0xbeb1d576,0x3ff67774,3 +np.float32,0xbf507166,0x40216bb3,3 +np.float32,0x3e86484c,0x3fa71813,3 +np.float32,0x3f09123e,0x3f80bd35,3 +np.float32,0xbe9abe0e,0x3ff05cb2,3 +np.float32,0x3f3019dc,0x3f4fed21,3 +np.float32,0xbe99e00e,0x3ff0227d,3 +np.float32,0xbf155ec5,0x400c6739,3 +np.float32,0x3f5857ba,0x3f106698,3 +np.float32,0x3edf619c,0x3f8f45fb,3 +np.float32,0xbf5ab76a,0x40261664,3 +np.float32,0x3e54b5a8,0x3fae4738,3 +np.float32,0xbee92772,0x4002ca40,3 +np.float32,0x3f2fd610,0x3f504a7a,3 +np.float32,0xbf38521c,0x4017f97e,3 +np.float32,0xff800000,0x7fc00000,3 +np.float32,0x3e2da348,0x3fb34077,3 +np.float32,0x3f2f85fa,0x3f50b894,3 +np.float32,0x3e88f9c8,0x3fa66551,3 +np.float32,0xbf61e570,0x4029b648,3 +np.float32,0xbeab362c,0x3ff4b4a1,3 +np.float32,0x3ec6c310,0x3f9607bd,3 +np.float32,0x3f0d7bda,0x3f7c3810,3 +np.float32,0xbeba5d36,0x3ff8bf99,3 +np.float32,0x3f4b0554,0x3f27adda,3 +np.float32,0x3f60f5dc,0x3efebfb3,3 +np.float32,0x3f36ce2c,0x3f468603,3 +np.float32,0xbe70afac,0x3fe76e8e,3 +np.float32,0x3f673350,0x3ee339b5,3 +np.float32,0xbe124cf0,0x3fdb698c,3 +np.float32,0xbf1243dc,0x400b73d0,3 +np.float32,0x3f3c8850,0x3f3e3407,3 +np.float32,0x3ea02f24,0x3fa05500,3 +np.float32,0xbeffed34,0x400607db,3 +np.float32,0x3f5c75c2,0x3f08817c,3 +np.float32,0x3f4b2fbe,0x3f27682d,3 +np.float32,0x3ee47c34,0x3f8dd9f9,3 +np.float32,0x3f50d48c,0x3f1de584,3 +np.float32,0x3f12dc5e,0x3f75b628,3 +np.float32,0xbefe7e4a,0x4005d2f4,3 +np.float32,0xbec2e846,0x3ffb0cbc,3 +np.float32,0xbedc3036,0x4000fb80,3 +np.float32,0xbf48aedc,0x401e311f,3 +np.float32,0x3f6e032e,0x3ec11363,3 +np.float32,0xbf60de15,0x40292b72,3 +np.float32,0x3f06585e,0x3f8258ba,3 +np.float32,0x3ef49b98,0x3f894e66,3 +np.float32,0x3cc5fe00,0x3fc5f7cf,3 +np.float32,0xbf7525c5,0x40365c2c,3 +np.float32,0x3f64f9f8,0x3eed5fb2,3 +np.float32,0x3e8849c0,0x3fa692fb,3 +np.float32,0x3e50c878,0x3faec79e,3 +np.float32,0x3ed61530,0x3f91d831,3 +np.float32,0xbf54872e,0x40233724,3 +np.float32,0xbf52ee7f,0x4022815e,3 +np.float32,0xbe708c24,0x3fe769fc,3 +np.float32,0xbf26fc54,0x40120260,3 +np.float32,0x3f226e8a,0x3f6228db,3 +np.float32,0xbef30406,0x40042eb8,3 +np.float32,0x3f5d996c,0x3f063f5f,3 +np.float32,0xbf425f9c,0x401bb618,3 +np.float32,0x3e4bb260,0x3faf6dc9,3 +np.float32,0xbe52d5a4,0x3fe39b29,3 +np.float32,0xbe169cf0,0x3fdbf505,3 +np.float32,0xbedfc422,0x40017a8e,3 +np.float32,0x3d8ffef0,0x3fc00e05,3 +np.float32,0xbf12bdab,0x400b98f2,3 +np.float32,0x3f295d0a,0x3f590e88,3 +np.float32,0x3f49d8e4,0x3f2998aa,3 +np.float32,0xbef914f4,0x40050c12,3 +np.float32,0xbf4ea2b5,0x4020a61e,3 +np.float32,0xbf3a89e5,0x4018c762,3 +np.float32,0x3e8707b4,0x3fa6e67a,3 +np.float32,0x3ac55400,0x3fc8de86,3 +np.float32,0x800000,0x3fc90fdb,3 +np.float32,0xbeb9762c,0x3ff8819b,3 +np.float32,0xbebbe23c,0x3ff92815,3 +np.float32,0xbf598c88,0x402587a1,3 +np.float32,0x3e95d864,0x3fa30b4a,3 +np.float32,0x3f7f6f40,0x3d882486,3 +np.float32,0xbf53658c,0x4022b604,3 +np.float32,0xbf2a35f2,0x401314ad,3 +np.float32,0x3eb14380,0x3f9bcf28,3 +np.float32,0x3f0e0c64,0x3f7b8a7a,3 +np.float32,0x3d349920,0x3fc36a9a,3 +np.float32,0xbec2092c,0x3ffad071,3 +np.float32,0xbe1d08e8,0x3fdcc4e0,3 +np.float32,0xbf008968,0x40063243,3 +np.float32,0xbefad582,0x40054c51,3 +np.float32,0xbe52d010,0x3fe39a72,3 +np.float32,0x3f4afdac,0x3f27ba6b,3 +np.float32,0x3f6c483c,0x3eca4408,3 +np.float32,0xbef3cb68,0x40044b0c,3 +np.float32,0x3e94687c,0x3fa36b6f,3 +np.float32,0xbf64ae5c,0x402b39bb,3 +np.float32,0xbf0022b4,0x40061497,3 +np.float32,0x80000001,0x3fc90fdb,3 +np.float32,0x3f25bcd0,0x3f5dda4b,3 +np.float32,0x3ed91b40,0x3f9102d7,3 +np.float32,0x3f800000,0x0,3 +np.float32,0xbebc6aca,0x3ff94cca,3 +np.float32,0x3f239e9a,0x3f609e7d,3 +np.float32,0xbf7312be,0x4034a305,3 +np.float32,0x3efd16d0,0x3f86e148,3 +np.float32,0x3f52753a,0x3f1b0f72,3 +np.float32,0xbde58960,0x3fd7702c,3 +np.float32,0x3ef88580,0x3f883099,3 +np.float32,0x3eebaefc,0x3f8bd51e,3 +np.float32,0x3e877d2c,0x3fa6c807,3 +np.float32,0x3f1a0324,0x3f6cdf32,3 +np.float32,0xbedfe20a,0x40017eb6,3 +np.float32,0x3f205a3c,0x3f64d69d,3 +np.float32,0xbeed5b7c,0x400361b0,3 +np.float32,0xbf69ba10,0x402e2ad0,3 +np.float32,0x3c4fe200,0x3fc77014,3 +np.float32,0x3f043310,0x3f839a69,3 +np.float32,0xbeaf359a,0x3ff5c485,3 +np.float32,0x3db3f110,0x3fbdcd12,3 +np.float32,0x3e24af88,0x3fb462ed,3 +np.float32,0xbf34e858,0x4016c1c8,3 +np.float32,0x3f3334f2,0x3f4b9cd0,3 +np.float32,0xbf145882,0x400c16a2,3 +np.float32,0xbf541c38,0x40230748,3 +np.float32,0x3eba7e10,0x3f99574b,3 +np.float32,0xbe34c6e0,0x3fdfc731,3 +np.float32,0xbe957abe,0x3feefbf0,3 +np.float32,0xbf595a59,0x40256fdb,3 +np.float32,0xbdedc7b8,0x3fd7f4f0,3 +np.float32,0xbf627c02,0x402a06a9,3 +np.float32,0x3f339b78,0x3f4b0d18,3 +np.float32,0xbf2df6d2,0x40145929,3 +np.float32,0x3f617726,0x3efc9fd8,3 +np.float32,0xbee3a8fc,0x40020561,3 +np.float32,0x3efe9f68,0x3f867043,3 +np.float32,0xbf2c3e76,0x4013c3ba,3 +np.float32,0xbf218f28,0x40103d84,3 +np.float32,0xbf1ea847,0x400f4f7f,3 +np.float32,0x3ded9160,0x3fba2e31,3 +np.float32,0x3bce1b00,0x3fc841bf,3 +np.float32,0xbe90566e,0x3feda46a,3 +np.float32,0xbf5ea2ba,0x4028056b,3 +np.float32,0x3f538e62,0x3f191ee6,3 +np.float32,0xbf59e054,0x4025af74,3 +np.float32,0xbe8c98ba,0x3fecab24,3 +np.float32,0x3ee7bdb0,0x3f8cf0b7,3 +np.float32,0xbf2eb828,0x40149b2b,3 +np.float32,0xbe5eb904,0x3fe52068,3 +np.float32,0xbf16b422,0x400cd08d,3 +np.float32,0x3f1ab9b4,0x3f6bfa58,3 +np.float32,0x3dc23040,0x3fbce82a,3 +np.float32,0xbf29d9e7,0x4012f5e5,3 +np.float32,0xbf38f30a,0x40183393,3 +np.float32,0x3e88e798,0x3fa66a09,3 +np.float32,0x3f1d07e6,0x3f69124f,3 +np.float32,0xbe1d3d34,0x3fdccb7e,3 +np.float32,0xbf1715be,0x400ceec2,3 +np.float32,0x3f7a0eac,0x3e5d11f7,3 +np.float32,0xbe764924,0x3fe82707,3 +np.float32,0xbf01a1f8,0x4006837c,3 +np.float32,0x3f2be730,0x3f55a661,3 +np.float32,0xbf7bb070,0x403d4ce5,3 +np.float32,0xbd602110,0x3fd011c9,3 +np.float32,0x3f5d080c,0x3f07609d,3 +np.float32,0xbda20400,0x3fd332d1,3 +np.float32,0x3f1c62da,0x3f69e308,3 +np.float32,0xbf2c6916,0x4013d223,3 +np.float32,0xbf44f8fd,0x401cb816,3 +np.float32,0x3f4da392,0x3f235539,3 +np.float32,0x3e9e8aa0,0x3fa0c3a0,3 +np.float32,0x3e9633c4,0x3fa2f366,3 +np.float32,0xbf0422ab,0x40073ddd,3 +np.float32,0x3f518386,0x3f1cb603,3 +np.float32,0x3f24307a,0x3f5fe096,3 +np.float32,0xbdfb4220,0x3fd8ce24,3 +np.float32,0x3f179d28,0x3f6fdc7d,3 +np.float32,0xbecc2df0,0x3ffd911e,3 +np.float32,0x3f3dff0c,0x3f3c0782,3 +np.float32,0xbf58c4d8,0x4025295b,3 +np.float32,0xbdcf8438,0x3fd60dd3,3 +np.float32,0xbeeaf1b2,0x40030aa7,3 +np.float32,0xbf298a28,0x4012db45,3 +np.float32,0x3f6c4dec,0x3eca2678,3 +np.float32,0x3f4d1ac8,0x3f243a59,3 +np.float32,0x3f62cdfa,0x3ef6e8f8,3 +np.float32,0xbee8acce,0x4002b909,3 +np.float32,0xbd5f2af0,0x3fd00a15,3 +np.float32,0x3f5fde8e,0x3f01a453,3 +np.float32,0x3e95233c,0x3fa33aa4,3 +np.float32,0x3ecd2a60,0x3f9449be,3 +np.float32,0x3f10aa86,0x3f78619d,3 +np.float32,0x3f3888e8,0x3f440a70,3 +np.float32,0x3eeb5bfc,0x3f8bec7d,3 +np.float32,0xbe12d654,0x3fdb7ae6,3 +np.float32,0x3eca3110,0x3f951931,3 +np.float32,0xbe2d1b7c,0x3fdece05,3 +np.float32,0xbf29e9db,0x4012fb3a,3 +np.float32,0xbf0c50b8,0x4009a845,3 +np.float32,0xbed9f0e4,0x4000abef,3 +np.float64,0x3fd078ec5ba0f1d8,0x3ff4f7c00595a4d3,1 +np.float64,0xbfdbc39743b7872e,0x400027f85bce43b2,1 +np.float64,0xbfacd2707c39a4e0,0x3ffa08ae1075d766,1 +np.float64,0xbfc956890f32ad14,0x3ffc52308e7285fd,1 +np.float64,0xbf939c2298273840,0x3ff9706d18e6ea6b,1 +np.float64,0xbfe0d7048961ae09,0x4000fff4406bd395,1 +np.float64,0xbfe9d19b86f3a337,0x4004139bc683a69f,1 +np.float64,0x3fd35c7f90a6b900,0x3ff437220e9123f8,1 +np.float64,0x3fdddca171bbb944,0x3ff15da61e61ec08,1 +np.float64,0x3feb300de9f6601c,0x3fe1c6fadb68cdca,1 +np.float64,0xbfef1815327e302a,0x400739808fc6f964,1 +np.float64,0xbfe332d78e6665af,0x4001b6c4ef922f7c,1 +np.float64,0xbfedbf4dfb7b7e9c,0x40061cefed62a58b,1 +np.float64,0xbfd8dcc7e3b1b990,0x3fff84307713c2c3,1 +np.float64,0xbfedaf161c7b5e2c,0x400612027c1b2b25,1 +np.float64,0xbfed9bde897b37bd,0x4006053f05bd7d26,1 +np.float64,0xbfe081ebc26103d8,0x4000e70755eb66e0,1 +np.float64,0xbfe0366f9c606cdf,0x4000d11212f29afd,1 +np.float64,0xbfc7c115212f822c,0x3ffc1e8c9d58f7db,1 +np.float64,0x3fd8dd9a78b1bb34,0x3ff2bf8d0f4c9376,1 +np.float64,0xbfe54eff466a9dfe,0x4002655950b611f4,1 +np.float64,0xbfe4aad987e955b3,0x40022efb19882518,1 +np.float64,0x3f70231ca0204600,0x3ff911d834e7abf4,1 +np.float64,0x3fede01d047bc03a,0x3fd773cecbd8561b,1 +np.float64,0xbfd6a00d48ad401a,0x3ffee9fd7051633f,1 +np.float64,0x3fd44f3d50a89e7c,0x3ff3f74dd0fc9c91,1 +np.float64,0x3fe540f0d0ea81e2,0x3feb055a7c7d43d6,1 +np.float64,0xbf3ba2e200374800,0x3ff923b582650c6c,1 +np.float64,0x3fe93b2d3f72765a,0x3fe532fa15331072,1 +np.float64,0x3fee8ce5a17d19cc,0x3fd35666eefbe336,1 +np.float64,0x3fe55d5f8feabac0,0x3feadf3dcfe251d4,1 +np.float64,0xbfd1d2ede8a3a5dc,0x3ffda600041ac884,1 +np.float64,0xbfee41186e7c8231,0x40067a625cc6f64d,1 +np.float64,0x3fe521a8b9ea4352,0x3feb2f1a6c8084e5,1 +np.float64,0x3fc65378ef2ca6f0,0x3ff653dfe81ee9f2,1 +np.float64,0x3fdaba0fbcb57420,0x3ff23d630995c6ba,1 +np.float64,0xbfe6b7441d6d6e88,0x4002e182539a2994,1 +np.float64,0x3fda00b6dcb4016c,0x3ff2703d516f28e7,1 +np.float64,0xbfe8699f01f0d33e,0x400382326920ea9e,1 +np.float64,0xbfef5889367eb112,0x4007832af5983793,1 +np.float64,0x3fefb57c8aff6afa,0x3fc14700ab38dcef,1 +np.float64,0xbfda0dfdaab41bfc,0x3fffd75b6fd497f6,1 +np.float64,0xbfb059c36620b388,0x3ffa27c528b97a42,1 +np.float64,0xbfdd450ab1ba8a16,0x40005dcac6ab50fd,1 +np.float64,0xbfe54d6156ea9ac2,0x400264ce9f3f0fb9,1 +np.float64,0xbfe076e94760edd2,0x4000e3d1374884da,1 +np.float64,0xbfc063286720c650,0x3ffb2fd1d6bff0ef,1 +np.float64,0xbfe24680f2e48d02,0x40016ddfbb5bcc0e,1 +np.float64,0xbfdc9351d2b926a4,0x400044e3756fb765,1 +np.float64,0x3fefb173d8ff62e8,0x3fc1bd5626f80850,1 +np.float64,0x3fe77c117a6ef822,0x3fe7e57089bad2ec,1 +np.float64,0xbfddbcebf7bb79d8,0x40006eadb60406b3,1 +np.float64,0xbfecf6625ff9ecc5,0x40059e6c6961a6db,1 +np.float64,0x3fdc8950b8b912a0,0x3ff1bcfb2e27795b,1 +np.float64,0xbfeb2fa517765f4a,0x4004b00aee3e6888,1 +np.float64,0x3fd0efc88da1df90,0x3ff4d8f7cbd8248a,1 +np.float64,0xbfe6641a2becc834,0x4002c43362c1bd0f,1 +np.float64,0xbfe28aec0fe515d8,0x400182c91d4df039,1 +np.float64,0xbfd5ede8d0abdbd2,0x3ffeba7baef05ae8,1 +np.float64,0xbfbd99702a3b32e0,0x3ffafca21c1053f1,1 +np.float64,0x3f96f043f82de080,0x3ff8c6384d5eb610,1 +np.float64,0xbfe5badbc9eb75b8,0x400289c8cd5873d1,1 +np.float64,0x3fe5c6bf95eb8d80,0x3fea5093e9a3e43e,1 +np.float64,0x3fb1955486232ab0,0x3ff8086d4c3e71d5,1 +np.float64,0xbfea145f397428be,0x4004302237a35871,1 +np.float64,0xbfdabe685db57cd0,0x400003e2e29725fb,1 +np.float64,0xbfefc79758ff8f2f,0x400831814e23bfc8,1 +np.float64,0x3fd7edb66cafdb6c,0x3ff3006c5123bfaf,1 +np.float64,0xbfeaf7644bf5eec8,0x400495a7963ce4ed,1 +np.float64,0x3fdf838d78bf071c,0x3ff0e527eed73800,1 +np.float64,0xbfd1a0165ba3402c,0x3ffd98c5ab76d375,1 +np.float64,0x3fd75b67a9aeb6d0,0x3ff327c8d80b17cf,1 +np.float64,0x3fc2aa9647255530,0x3ff6ca854b157df1,1 +np.float64,0xbfe0957fd4612b00,0x4000ecbf3932becd,1 +np.float64,0x3fda1792c0b42f24,0x3ff269fbb2360487,1 +np.float64,0x3fd480706ca900e0,0x3ff3ea53a6aa3ae8,1 +np.float64,0xbfd0780ed9a0f01e,0x3ffd4bfd544c7d47,1 +np.float64,0x3feeec0cd77dd81a,0x3fd0a8a241fdb441,1 +np.float64,0x3fcfa933e93f5268,0x3ff5223478621a6b,1 +np.float64,0x3fdad2481fb5a490,0x3ff236b86c6b2b49,1 +np.float64,0x3fe03b129de07626,0x3ff09f21fb868451,1 +np.float64,0xbfc01212cd202424,0x3ffb259a07159ae9,1 +np.float64,0x3febdb912df7b722,0x3fe0768e20dac8c9,1 +np.float64,0xbfbf2148763e4290,0x3ffb154c361ce5bf,1 +np.float64,0xbfb1a7eb1e234fd8,0x3ffa3cb37ac4a176,1 +np.float64,0xbfe26ad1ec64d5a4,0x400178f480ecce8d,1 +np.float64,0x3fe6d1cd1b6da39a,0x3fe8dc20ec4dad3b,1 +np.float64,0xbfede0e53dfbc1ca,0x4006340d3bdd7c97,1 +np.float64,0xbfe8fd1bd9f1fa38,0x4003bc3477f93f40,1 +np.float64,0xbfe329d0f26653a2,0x4001b3f345af5648,1 +np.float64,0xbfe4bb20eee97642,0x40023451404d6d08,1 +np.float64,0x3fb574832e2ae900,0x3ff7ca4bed0c7110,1 +np.float64,0xbfdf3c098fbe7814,0x4000a525bb72d659,1 +np.float64,0x3fa453e6d428a7c0,0x3ff87f512bb9b0c6,1 +np.float64,0x3faaec888435d920,0x3ff84a7d9e4def63,1 +np.float64,0xbfcdc240df3b8480,0x3ffce30ece754e7f,1 +np.float64,0xbf8c3220f0386440,0x3ff95a600ae6e157,1 +np.float64,0x3fe806076c700c0e,0x3fe71784a96c76eb,1 +np.float64,0x3fedf9b0e17bf362,0x3fd6e35fc0a7b6c3,1 +np.float64,0xbfe1b48422636908,0x400141bd8ed251bc,1 +np.float64,0xbfe82e2817705c50,0x40036b5a5556d021,1 +np.float64,0xbfc8ef8ff931df20,0x3ffc450ffae7ce58,1 +np.float64,0xbfe919fa94f233f5,0x4003c7cce4697fe8,1 +np.float64,0xbfc3ace4a72759c8,0x3ffb9a197bb22651,1 +np.float64,0x3fe479f71ee8f3ee,0x3fec0bd2f59097aa,1 +np.float64,0xbfeeb54a967d6a95,0x4006da12c83649c5,1 +np.float64,0x3fe5e74ea8ebce9e,0x3fea2407cef0f08c,1 +np.float64,0x3fb382baf2270570,0x3ff7e98213b921ba,1 +np.float64,0xbfdd86fd3cbb0dfa,0x40006712952ddbcf,1 +np.float64,0xbfd250eb52a4a1d6,0x3ffdc6d56253b1cd,1 +np.float64,0x3fea30c4ed74618a,0x3fe3962deba4f30e,1 +np.float64,0x3fc895963d312b30,0x3ff60a5d52fcbccc,1 +np.float64,0x3fe9cc4f6273989e,0x3fe442740942c80f,1 +np.float64,0xbfe8769f5cf0ed3f,0x4003873b4cb5bfce,1 +np.float64,0xbfe382f3726705e7,0x4001cfeb3204d110,1 +np.float64,0x3fbfe9a9163fd350,0x3ff7220bd2b97c8f,1 +np.float64,0xbfca6162bb34c2c4,0x3ffc743f939358f1,1 +np.float64,0x3fe127a014e24f40,0x3ff0147c4bafbc39,1 +np.float64,0x3fee9cdd2a7d39ba,0x3fd2e9ef45ab122f,1 +np.float64,0x3fa9ffb97c33ff80,0x3ff851e69fa3542c,1 +np.float64,0x3fd378f393a6f1e8,0x3ff42faafa77de56,1 +np.float64,0xbfe4df1e1669be3c,0x400240284df1c321,1 +np.float64,0x3fed0ed79bfa1db0,0x3fdba89060aa96fb,1 +np.float64,0x3fdef2ee52bde5dc,0x3ff10e942244f4f1,1 +np.float64,0xbfdab38f3ab5671e,0x40000264d8d5b49b,1 +np.float64,0x3fbe95a96e3d2b50,0x3ff73774cb59ce2d,1 +np.float64,0xbfe945653af28aca,0x4003d9657bf129c2,1 +np.float64,0xbfb18f3f2a231e80,0x3ffa3b27cba23f50,1 +np.float64,0xbfef50bf22fea17e,0x40077998a850082c,1 +np.float64,0xbfc52b8c212a5718,0x3ffbca8d6560a2da,1 +np.float64,0x7ff8000000000000,0x7ff8000000000000,1 +np.float64,0x3fc1e3a02d23c740,0x3ff6e3a5fcac12a4,1 +np.float64,0xbfeb5e4ea5f6bc9d,0x4004c65abef9426f,1 +np.float64,0xbfe425b132684b62,0x400203c29608b00d,1 +np.float64,0xbfbfa1c19e3f4380,0x3ffb1d6367711158,1 +np.float64,0x3fbba2776e3744f0,0x3ff766f6df586fad,1 +np.float64,0xbfb5d0951e2ba128,0x3ffa7f712480b25e,1 +np.float64,0xbfe949fdab7293fb,0x4003db4530a18507,1 +np.float64,0xbfcf13519b3e26a4,0x3ffd0e6f0a6c38ee,1 +np.float64,0x3f91e6d72823cdc0,0x3ff8da5f08909b6e,1 +np.float64,0x3f78a2e360314600,0x3ff909586727caef,1 +np.float64,0xbfe1ae7e8fe35cfd,0x40013fef082caaa3,1 +np.float64,0x3fe97a6dd1f2f4dc,0x3fe4cb4b99863478,1 +np.float64,0xbfcc1e1e69383c3c,0x3ffcad250a949843,1 +np.float64,0x3faccb797c399700,0x3ff83b8066b49330,1 +np.float64,0x3fe7a2647a6f44c8,0x3fe7acceae6ec425,1 +np.float64,0xbfec3bfcf0f877fa,0x4005366af5a7175b,1 +np.float64,0xbfe2310b94646217,0x400167588fceb228,1 +np.float64,0x3feb167372762ce6,0x3fe1f74c0288fad8,1 +np.float64,0xbfb722b4ee2e4568,0x3ffa94a81b94dfca,1 +np.float64,0x3fc58da9712b1b50,0x3ff66cf8f072aa14,1 +np.float64,0xbfe7fff9d6effff4,0x400359d01b8141de,1 +np.float64,0xbfd56691c5aacd24,0x3ffe9686697797e8,1 +np.float64,0x3fe3ab0557e7560a,0x3fed1593959ef8e8,1 +np.float64,0x3fdd458995ba8b14,0x3ff1883d6f22a322,1 +np.float64,0x3fe7bbed2cef77da,0x3fe786d618094cda,1 +np.float64,0x3fa31a30c4263460,0x3ff88920b936fd79,1 +np.float64,0x8010000000000000,0x3ff921fb54442d18,1 +np.float64,0xbfdc5effbdb8be00,0x40003d95fe0dff11,1 +np.float64,0x3febfdad7e77fb5a,0x3fe030b5297dbbdd,1 +np.float64,0x3fe4f3f3b2e9e7e8,0x3feb6bc59eeb2be2,1 +np.float64,0xbfe44469fd6888d4,0x40020daa5488f97a,1 +np.float64,0xbfe19fddb0e33fbc,0x40013b8c902b167b,1 +np.float64,0x3fa36ad17c26d5a0,0x3ff8869b3e828134,1 +np.float64,0x3fcf23e6c93e47d0,0x3ff5336491a65d1e,1 +np.float64,0xffefffffffffffff,0x7ff8000000000000,1 +np.float64,0xbfe375f4cee6ebea,0x4001cbd2ba42e8b5,1 +np.float64,0xbfaef1215c3de240,0x3ffa19ab02081189,1 +np.float64,0xbfec39c59c78738b,0x4005353dc38e3d78,1 +np.float64,0x7ff4000000000000,0x7ffc000000000000,1 +np.float64,0xbfec09bb7b781377,0x40051c0a5754cb3a,1 +np.float64,0x3fe8301f2870603e,0x3fe6d783c5ef0944,1 +np.float64,0xbfed418c987a8319,0x4005cbae1b8693d1,1 +np.float64,0xbfdc16e7adb82dd0,0x4000338b634eaf03,1 +np.float64,0x3fd5d361bdaba6c4,0x3ff390899300a54c,1 +np.float64,0xbff0000000000000,0x400921fb54442d18,1 +np.float64,0x3fd5946232ab28c4,0x3ff3a14767813f29,1 +np.float64,0x3fe833e5fef067cc,0x3fe6d1be720edf2d,1 +np.float64,0x3fedf746a67bee8e,0x3fd6f127fdcadb7b,1 +np.float64,0x3fd90353d3b206a8,0x3ff2b54f7d369ba9,1 +np.float64,0x3fec4b4b72f89696,0x3fdf1b38d2e93532,1 +np.float64,0xbfe9c67596f38ceb,0x40040ee5f524ce03,1 +np.float64,0x3fd350d91aa6a1b4,0x3ff43a303c0da27f,1 +np.float64,0x3fd062603ba0c4c0,0x3ff4fd9514b935d8,1 +np.float64,0xbfe24c075f64980e,0x40016f8e9f2663b3,1 +np.float64,0x3fdaa546eeb54a8c,0x3ff2431a88fef1d5,1 +np.float64,0x3fe92b8151f25702,0x3fe54c67e005cbf9,1 +np.float64,0xbfe1be8b8a637d17,0x400144c078f67c6e,1 +np.float64,0xbfe468a1d7e8d144,0x40021964b118cbf4,1 +np.float64,0xbfdc6de4fab8dbca,0x40003fa9e27893d8,1 +np.float64,0xbfe3c2788ae784f1,0x4001e407ba3aa956,1 +np.float64,0xbfe2bf1542e57e2a,0x400192d4a9072016,1 +np.float64,0xbfe6982f4c6d305e,0x4002d681b1991bbb,1 +np.float64,0x3fdbceb1c4b79d64,0x3ff1f0f117b9d354,1 +np.float64,0x3fdb3705e7b66e0c,0x3ff21af01ca27ace,1 +np.float64,0x3fe3e6358ee7cc6c,0x3fecca4585053983,1 +np.float64,0xbfe16d6a9a62dad5,0x40012c7988aee247,1 +np.float64,0xbfce66e4413ccdc8,0x3ffcf83b08043a0c,1 +np.float64,0xbfeb6cd46876d9a9,0x4004cd61733bfb79,1 +np.float64,0xbfdb1cdd64b639ba,0x400010e6cf087cb7,1 +np.float64,0xbfe09e4e30e13c9c,0x4000ef5277c47721,1 +np.float64,0xbfee88dd127d11ba,0x4006b3cd443643ac,1 +np.float64,0xbf911e06c8223c00,0x3ff966744064fb05,1 +np.float64,0xbfe8f22bc471e458,0x4003b7d5513af295,1 +np.float64,0x3fe3d7329567ae66,0x3fecdd6c241f83ee,1 +np.float64,0x3fc8a9404b315280,0x3ff607dc175edf3f,1 +np.float64,0x3fe7eb80ad6fd702,0x3fe73f8fdb3e6a6c,1 +np.float64,0x3fef0931e37e1264,0x3fcf7fde80a3c5ab,1 +np.float64,0x3fe2ed3c3fe5da78,0x3fee038334cd1860,1 +np.float64,0x3fe251fdb8e4a3fc,0x3feec26dc636ac31,1 +np.float64,0x3feb239436764728,0x3fe1de9462455da7,1 +np.float64,0xbfe63fd7eeec7fb0,0x4002b78cfa3d2fa6,1 +np.float64,0x3fdd639cb5bac738,0x3ff17fc7d92b3eee,1 +np.float64,0x3fd0a7a13fa14f44,0x3ff4eba95c559c84,1 +np.float64,0x3fe804362d70086c,0x3fe71a44cd91ffa4,1 +np.float64,0xbfe0fecf6e61fd9f,0x40010bac8edbdc4f,1 +np.float64,0x3fcb74acfd36e958,0x3ff5ac84437f1b7c,1 +np.float64,0x3fe55053e1eaa0a8,0x3feaf0bf76304c30,1 +np.float64,0x3fc06b508d20d6a0,0x3ff7131da17f3902,1 +np.float64,0x3fdd78750fbaf0ec,0x3ff179e97fbf7f65,1 +np.float64,0x3fe44cb946689972,0x3fec46859b5da6be,1 +np.float64,0xbfeb165a7ff62cb5,0x4004a41c9cc9589e,1 +np.float64,0x3fe01ffb2b603ff6,0x3ff0aed52bf1c3c1,1 +np.float64,0x3f983c60a83078c0,0x3ff8c107805715ab,1 +np.float64,0x3fd8b5ff13b16c00,0x3ff2ca4a837a476a,1 +np.float64,0x3fc80510a1300a20,0x3ff61cc3b4af470b,1 +np.float64,0xbfd3935b06a726b6,0x3ffe1b3a2066f473,1 +np.float64,0xbfdd4a1f31ba943e,0x40005e81979ed445,1 +np.float64,0xbfa76afdd42ed600,0x3ff9dd63ffba72d2,1 +np.float64,0x3fe7e06d496fc0da,0x3fe7503773566707,1 +np.float64,0xbfea5fbfe874bf80,0x40045106af6c538f,1 +np.float64,0x3fee000c487c0018,0x3fd6bef1f8779d88,1 +np.float64,0xbfb39f4ee2273ea0,0x3ffa5c3f2b3888ab,1 +np.float64,0x3feb9247b0772490,0x3fe1092d2905efce,1 +np.float64,0x3fdaa39b4cb54738,0x3ff243901da0da17,1 +np.float64,0x3fcd5b2b493ab658,0x3ff56e262e65b67d,1 +np.float64,0x3fcf82512f3f04a0,0x3ff52738847c55f2,1 +np.float64,0x3fe2af5e0c655ebc,0x3fee4ffab0c82348,1 +np.float64,0xbfec0055d0f800ac,0x4005172d325933e8,1 +np.float64,0x3fe71da9336e3b52,0x3fe86f2e12f6e303,1 +np.float64,0x3fbefab0723df560,0x3ff731188ac716ec,1 +np.float64,0xbfe11dca28623b94,0x400114d3d4ad370d,1 +np.float64,0x3fbcbda8ca397b50,0x3ff755281078abd4,1 +np.float64,0x3fe687c7126d0f8e,0x3fe945099a7855cc,1 +np.float64,0xbfecde510579bca2,0x400590606e244591,1 +np.float64,0xbfd72de681ae5bce,0x3fff0ff797ad1755,1 +np.float64,0xbfe7c0f7386f81ee,0x40034226e0805309,1 +np.float64,0x3fd8d55619b1aaac,0x3ff2c1cb3267b14e,1 +np.float64,0x3fecd7a2ad79af46,0x3fdcabbffeaa279e,1 +np.float64,0x3fee7fb1a8fcff64,0x3fd3ae620286fe19,1 +np.float64,0xbfc5f3a3592be748,0x3ffbe3ed204d9842,1 +np.float64,0x3fec9e5527793caa,0x3fddb00bc8687e4b,1 +np.float64,0x3fc35dc70f26bb90,0x3ff6b3ded7191e33,1 +np.float64,0x3fda91c07ab52380,0x3ff24878848fec8f,1 +np.float64,0xbfe12cde1fe259bc,0x4001194ab99d5134,1 +np.float64,0xbfd35ab736a6b56e,0x3ffe0c5ce8356d16,1 +np.float64,0x3fc9c94123339280,0x3ff5e3239f3ad795,1 +np.float64,0xbfe72f54926e5ea9,0x40030c95d1d02b56,1 +np.float64,0xbfee283186fc5063,0x40066786bd0feb79,1 +np.float64,0xbfe7b383f56f6708,0x40033d23ef0e903d,1 +np.float64,0x3fd6037327ac06e8,0x3ff383bf2f311ddb,1 +np.float64,0x3fe0e344b561c68a,0x3ff03cd90fd4ba65,1 +np.float64,0xbfef0ff54b7e1feb,0x400730fa5fce381e,1 +np.float64,0x3fd269929da4d324,0x3ff476b230136d32,1 +np.float64,0xbfbc5fb9f638bf70,0x3ffae8e63a4e3234,1 +np.float64,0xbfe2e8bc84e5d179,0x40019fb5874f4310,1 +np.float64,0xbfd7017413ae02e8,0x3fff040d843c1531,1 +np.float64,0x3fefd362fa7fa6c6,0x3fbababc3ddbb21d,1 +np.float64,0x3fecb62ed3f96c5e,0x3fdd44ba77ccff94,1 +np.float64,0xbfb16fad5222df58,0x3ffa392d7f02b522,1 +np.float64,0x3fbcf4abc639e950,0x3ff751b23c40e27f,1 +np.float64,0x3fe128adbce2515c,0x3ff013dc91db04b5,1 +np.float64,0x3fa5dd9d842bbb40,0x3ff87300c88d512f,1 +np.float64,0xbfe61efcaf6c3dfa,0x4002ac27117f87c9,1 +np.float64,0x3feffe1233fffc24,0x3f9638d3796a4954,1 +np.float64,0xbfe78548b66f0a92,0x40032c0447b7bfe2,1 +np.float64,0x3fe7bd38416f7a70,0x3fe784e86d6546b6,1 +np.float64,0x3fe0d6bc5961ad78,0x3ff0443899e747ac,1 +np.float64,0xbfd0bb6e47a176dc,0x3ffd5d6dff390d41,1 +np.float64,0xbfec1d16b8f83a2e,0x40052620378d3b78,1 +np.float64,0x3fe9bbec20f377d8,0x3fe45e167c7a3871,1 +np.float64,0xbfeed81d9dfdb03b,0x4006f9dec2db7310,1 +np.float64,0xbfe1e35179e3c6a3,0x40014fd1b1186ac0,1 +np.float64,0xbfc9c7e605338fcc,0x3ffc60a6bd1a7126,1 +np.float64,0x3feec92810fd9250,0x3fd1afde414ab338,1 +np.float64,0xbfeb9f1d90773e3b,0x4004e606b773f5b0,1 +np.float64,0x3fcbabdf6b3757c0,0x3ff5a573866404af,1 +np.float64,0x3fe9f4e1fff3e9c4,0x3fe3fd7b6712dd7b,1 +np.float64,0xbfe6c0175ded802e,0x4002e4a4dc12f3fe,1 +np.float64,0xbfeefc96f37df92e,0x40071d367cd721ff,1 +np.float64,0xbfeaab58dc7556b2,0x400472ce37e31e50,1 +np.float64,0xbfc62668772c4cd0,0x3ffbea5e6c92010a,1 +np.float64,0x3fafe055fc3fc0a0,0x3ff822ce6502519a,1 +np.float64,0x3fd7b648ffaf6c90,0x3ff30f5a42f11418,1 +np.float64,0xbfe934fe827269fd,0x4003d2b9fed9e6ad,1 +np.float64,0xbfe6d691f2edad24,0x4002eca6a4b1797b,1 +np.float64,0x3fc7e62ced2fcc58,0x3ff620b1f44398b7,1 +np.float64,0xbfc89be9f33137d4,0x3ffc3a67a497f59c,1 +np.float64,0xbfe7793d536ef27a,0x40032794bf14dd64,1 +np.float64,0x3fde55a02dbcab40,0x3ff13b5f82d223e4,1 +np.float64,0xbfc8eabd7b31d57c,0x3ffc4472a81cb6d0,1 +np.float64,0x3fddcb5468bb96a8,0x3ff162899c381f2e,1 +np.float64,0xbfec7554d8f8eaaa,0x40055550e18ec463,1 +np.float64,0x3fd0b6e8b6a16dd0,0x3ff4e7b4781a50e3,1 +np.float64,0x3fedaae01b7b55c0,0x3fd8964916cdf53d,1 +np.float64,0x3fe0870f8a610e20,0x3ff072e7db95c2a2,1 +np.float64,0xbfec3e3ce2787c7a,0x4005379d0f6be873,1 +np.float64,0xbfe65502586caa04,0x4002beecff89147f,1 +np.float64,0xbfe0df39a961be74,0x4001025e36d1c061,1 +np.float64,0xbfb5d8edbe2bb1d8,0x3ffa7ff72b7d6a2b,1 +np.float64,0xbfde89574bbd12ae,0x40008ba4cd74544d,1 +np.float64,0xbfe72938f0ee5272,0x40030a5efd1acb6d,1 +np.float64,0xbfcd500d133aa01c,0x3ffcd462f9104689,1 +np.float64,0x3fe0350766606a0e,0x3ff0a2a3664e2c14,1 +np.float64,0xbfc892fb573125f8,0x3ffc3944641cc69d,1 +np.float64,0xbfba7dc7c634fb90,0x3ffaca9a6a0ffe61,1 +np.float64,0xbfeac94478759289,0x40048068a8b83e45,1 +np.float64,0xbfe8f60c1af1ec18,0x4003b961995b6e51,1 +np.float64,0x3fea1c0817743810,0x3fe3ba28c1643cf7,1 +np.float64,0xbfe42a0fefe85420,0x4002052aadd77f01,1 +np.float64,0x3fd2c61c56a58c38,0x3ff45e84cb9a7fa9,1 +np.float64,0xbfd83fb7cdb07f70,0x3fff59ab4790074c,1 +np.float64,0x3fd95e630fb2bcc8,0x3ff29c8bee1335ad,1 +np.float64,0x3feee88f387dd11e,0x3fd0c3ad3ded4094,1 +np.float64,0x3fe061291160c252,0x3ff0890010199bbc,1 +np.float64,0xbfdc7db3b5b8fb68,0x400041dea3759443,1 +np.float64,0x3fee23b320fc4766,0x3fd5ee73d7aa5c56,1 +np.float64,0xbfdc25c590b84b8c,0x4000359cf98a00b4,1 +np.float64,0xbfd63cbfd2ac7980,0x3ffecf7b9cf99b3c,1 +np.float64,0xbfbeb3c29a3d6788,0x3ffb0e66ecc0fc3b,1 +np.float64,0xbfd2f57fd6a5eb00,0x3ffdf1d7c79e1532,1 +np.float64,0xbfab3eda9c367db0,0x3ff9fc0c875f42e9,1 +np.float64,0xbfe12df1c6e25be4,0x4001199c673e698c,1 +np.float64,0x3fef8ab23a7f1564,0x3fc5aff358c59f1c,1 +np.float64,0x3fe562f50feac5ea,0x3fead7bce205f7d9,1 +np.float64,0x3fdc41adbeb8835c,0x3ff1d0f71341b8f2,1 +np.float64,0x3fe2748967e4e912,0x3fee9837f970ff9e,1 +np.float64,0xbfdaa89d57b5513a,0x400000e3889ba4cf,1 +np.float64,0x3fdf2a137dbe5428,0x3ff0fecfbecbbf86,1 +np.float64,0xbfea1fdcd2f43fba,0x4004351974b32163,1 +np.float64,0xbfe34a93a3e69528,0x4001be323946a3e0,1 +np.float64,0x3fe929bacff25376,0x3fe54f47bd7f4cf2,1 +np.float64,0xbfd667fbd6accff8,0x3ffedb04032b3a1a,1 +np.float64,0xbfeb695796f6d2af,0x4004cbb08ec6f525,1 +np.float64,0x3fd204df2ea409c0,0x3ff490f51e6670f5,1 +np.float64,0xbfd89a2757b1344e,0x3fff722127b988c4,1 +np.float64,0xbfd0787187a0f0e4,0x3ffd4c16dbe94f32,1 +np.float64,0x3fd44239bfa88474,0x3ff3fabbfb24b1fa,1 +np.float64,0xbfeb0b3489f61669,0x40049ee33d811d33,1 +np.float64,0x3fdcf04eaab9e09c,0x3ff1a02a29996c4e,1 +np.float64,0x3fd4c51e4fa98a3c,0x3ff3d8302c68fc9a,1 +np.float64,0x3fd1346645a268cc,0x3ff4c72b4970ecaf,1 +np.float64,0x3fd6a89d09ad513c,0x3ff357af6520afac,1 +np.float64,0xbfba0f469a341e90,0x3ffac3a8f41bed23,1 +np.float64,0xbfe13f8ddce27f1c,0x40011ed557719fd6,1 +np.float64,0x3fd43e5e26a87cbc,0x3ff3fbc040fc30dc,1 +np.float64,0x3fe838125a707024,0x3fe6cb5c987248f3,1 +np.float64,0x3fe128c30c625186,0x3ff013cff238dd1b,1 +np.float64,0xbfcd4718833a8e30,0x3ffcd33c96bde6f9,1 +np.float64,0x3fe43fcd08e87f9a,0x3fec573997456ec1,1 +np.float64,0xbfe9a29104734522,0x4003ffd502a1b57f,1 +np.float64,0xbfe4709d7968e13b,0x40021bfc5cd55af4,1 +np.float64,0x3fd21c3925a43874,0x3ff48adf48556cbb,1 +np.float64,0x3fe9a521b2734a44,0x3fe4844fc054e839,1 +np.float64,0xbfdfa6a912bf4d52,0x4000b4730ad8521e,1 +np.float64,0x3fe3740702e6e80e,0x3fed5b106283b6ed,1 +np.float64,0x3fd0a3aa36a14754,0x3ff4ecb02a5e3f49,1 +np.float64,0x3fdcb903d0b97208,0x3ff1afa5d692c5b9,1 +np.float64,0xbfe7d67839efacf0,0x40034a3146abf6f2,1 +np.float64,0x3f9981c6d8330380,0x3ff8bbf1853d7b90,1 +np.float64,0xbfe9d4191673a832,0x400414a9ab453c5d,1 +np.float64,0x3fef0a1e5c7e143c,0x3fcf70b02a54c415,1 +np.float64,0xbfd996dee6b32dbe,0x3fffb6cf707ad8e4,1 +np.float64,0x3fe19bef17e337de,0x3fef9e70d4fcedae,1 +np.float64,0x3fe34a59716694b2,0x3fed8f6d5cfba474,1 +np.float64,0x3fdf27e27cbe4fc4,0x3ff0ff70500e0c7c,1 +np.float64,0xbfe19df87fe33bf1,0x40013afb401de24c,1 +np.float64,0xbfbdfd97ba3bfb30,0x3ffb02ef8c225e57,1 +np.float64,0xbfe3d3417267a683,0x4001e95ed240b0f8,1 +np.float64,0x3fe566498b6acc94,0x3fead342957d4910,1 +np.float64,0x3ff0000000000000,0x0,1 +np.float64,0x3feb329bd8766538,0x3fe1c2225aafe3b4,1 +np.float64,0xbfc19ca703233950,0x3ffb575b5df057b9,1 +np.float64,0x3fe755027d6eaa04,0x3fe81eb99c262e00,1 +np.float64,0xbfe6c2b8306d8570,0x4002e594199f9eec,1 +np.float64,0x3fd69438e6ad2870,0x3ff35d2275ae891d,1 +np.float64,0x3fda3e7285b47ce4,0x3ff25f5573dd47ae,1 +np.float64,0x3fe7928a166f2514,0x3fe7c4490ef4b9a9,1 +np.float64,0xbfd4eb71b9a9d6e4,0x3ffe75e8ccb74be1,1 +np.float64,0xbfcc3a07f1387410,0x3ffcb0b8af914a5b,1 +np.float64,0xbfe6e80225edd004,0x4002f2e26eae8999,1 +np.float64,0xbfb347728a268ee8,0x3ffa56bd526a12db,1 +np.float64,0x3fe5140ead6a281e,0x3feb4132c9140a1c,1 +np.float64,0xbfc147f125228fe4,0x3ffb4cab18b9050f,1 +np.float64,0xbfcb9145b537228c,0x3ffc9b1b6227a8c9,1 +np.float64,0xbfda84ef4bb509de,0x3ffff7f8a674e17d,1 +np.float64,0x3fd2eb6bbfa5d6d8,0x3ff454c225529d7e,1 +np.float64,0x3fe18c95f1e3192c,0x3fefb0cf0efba75a,1 +np.float64,0x3fe78606efef0c0e,0x3fe7d6c3a092d64c,1 +np.float64,0x3fbad5119a35aa20,0x3ff773dffe3ce660,1 +np.float64,0x3fd0cf5903a19eb4,0x3ff4e15fd21fdb42,1 +np.float64,0xbfd85ce90bb0b9d2,0x3fff618ee848e974,1 +np.float64,0x3fe90e11b9f21c24,0x3fe57be62f606f4a,1 +np.float64,0x3fd7a2040faf4408,0x3ff314ce85457ec2,1 +np.float64,0xbfd73fba69ae7f74,0x3fff14bff3504811,1 +np.float64,0x3fa04b4bd42096a0,0x3ff89f9b52f521a2,1 +np.float64,0xbfd7219ce5ae433a,0x3fff0cac0b45cc18,1 +np.float64,0xbfe0cf4661e19e8d,0x4000fdadb14e3c22,1 +np.float64,0x3fd07469fea0e8d4,0x3ff4f8eaa9b2394a,1 +np.float64,0x3f9b05c5d8360b80,0x3ff8b5e10672db5c,1 +np.float64,0x3fe4c25b916984b8,0x3febad29bd0e25e2,1 +np.float64,0xbfde8b4891bd1692,0x40008beb88d5c409,1 +np.float64,0xbfe199a7efe33350,0x400139b089aee21c,1 +np.float64,0x3fecdad25cf9b5a4,0x3fdc9d062867e8c3,1 +np.float64,0xbfe979b277f2f365,0x4003eedb061e25a4,1 +np.float64,0x3fc8c7311f318e60,0x3ff6040b9aeaad9d,1 +np.float64,0x3fd2b605b8a56c0c,0x3ff462b9a955c224,1 +np.float64,0x3fc073b6ad20e770,0x3ff7120e9f2fd63c,1 +np.float64,0xbfec60ede678c1dc,0x40054a3863e24dc2,1 +np.float64,0x3fe225171be44a2e,0x3feef910dca420ea,1 +np.float64,0xbfd7529762aea52e,0x3fff19d00661f650,1 +np.float64,0xbfd781783daf02f0,0x3fff2667b90be461,1 +np.float64,0x3fe3f6ec6d67edd8,0x3fecb4e814a2e33a,1 +np.float64,0x3fece6702df9cce0,0x3fdc6719d92a50d2,1 +np.float64,0xbfb5c602ce2b8c08,0x3ffa7ec761ba856a,1 +np.float64,0xbfd61f0153ac3e02,0x3ffec78e3b1a6c4d,1 +np.float64,0xbfec3462b2f868c5,0x400532630bbd7050,1 +np.float64,0xbfdd248485ba490a,0x400059391c07c1bb,1 +np.float64,0xbfd424921fa84924,0x3ffe416a85d1dcdf,1 +np.float64,0x3fbb23a932364750,0x3ff76eef79209f7f,1 +np.float64,0x3fca248b0f344918,0x3ff5d77c5c1b4e5e,1 +np.float64,0xbfe69af4a4ed35ea,0x4002d77c2e4fbd4e,1 +np.float64,0x3fdafe3cdcb5fc78,0x3ff22a9be6efbbf2,1 +np.float64,0xbfebba3377f77467,0x4004f3836e1fe71a,1 +np.float64,0xbfe650fae06ca1f6,0x4002bd851406377c,1 +np.float64,0x3fda630007b4c600,0x3ff2554f1832bd94,1 +np.float64,0xbfda8107d9b50210,0x3ffff6e6209659f3,1 +np.float64,0x3fea759a02f4eb34,0x3fe31d1a632c9aae,1 +np.float64,0x3fbf88149e3f1030,0x3ff728313aa12ccb,1 +np.float64,0x3f7196d2a0232e00,0x3ff910647e1914c1,1 +np.float64,0x3feeae51d17d5ca4,0x3fd2709698d31f6f,1 +np.float64,0xbfd73cd663ae79ac,0x3fff13f96300b55a,1 +np.float64,0x3fd4fc5f06a9f8c0,0x3ff3c99359854b97,1 +np.float64,0x3fb29f5d6e253ec0,0x3ff7f7c20e396b20,1 +np.float64,0xbfd757c82aaeaf90,0x3fff1b34c6141e98,1 +np.float64,0x3fc56fd4cf2adfa8,0x3ff670c145122909,1 +np.float64,0x3fc609a2f52c1348,0x3ff65d3ef3cade2c,1 +np.float64,0xbfe1de631163bcc6,0x40014e5528fadb73,1 +np.float64,0xbfe7eb4a726fd695,0x40035202f49d95c4,1 +np.float64,0xbfc9223771324470,0x3ffc4b84d5e263b9,1 +np.float64,0x3fee91a8a87d2352,0x3fd3364befde8de6,1 +np.float64,0x3fbc9784fe392f10,0x3ff7578e29f6a1b2,1 +np.float64,0xbfec627c2c78c4f8,0x40054b0ff2cb9c55,1 +np.float64,0xbfb8b406a6316810,0x3ffaadd97062fb8c,1 +np.float64,0xbfecf98384f9f307,0x4005a043d9110d79,1 +np.float64,0xbfe5834bab6b0698,0x400276f114aebee4,1 +np.float64,0xbfd90f391eb21e72,0x3fff91e26a8f48f3,1 +np.float64,0xbfee288ce2fc511a,0x400667cb09aa04b3,1 +np.float64,0x3fd5aa5e32ab54bc,0x3ff39b7080a52214,1 +np.float64,0xbfee7ef907fcfdf2,0x4006ab96a8eba4c5,1 +np.float64,0x3fd6097973ac12f4,0x3ff3822486978bd1,1 +np.float64,0xbfe02d14b8e05a2a,0x4000ce5be53047b1,1 +np.float64,0xbf9c629a6838c540,0x3ff993897728c3f9,1 +np.float64,0xbfee2024667c4049,0x40066188782fb1f0,1 +np.float64,0xbfa42a88fc285510,0x3ff9c35a4bbce104,1 +np.float64,0x3fa407af5c280f60,0x3ff881b360d8eea1,1 +np.float64,0x3fed0ba42cfa1748,0x3fdbb7d55609175f,1 +np.float64,0xbfdd0b5844ba16b0,0x400055b0bb59ebb2,1 +np.float64,0x3fd88d97e6b11b30,0x3ff2d53c1ecb8f8c,1 +np.float64,0xbfeb7a915ef6f523,0x4004d410812eb84c,1 +np.float64,0xbfb5f979ca2bf2f0,0x3ffa8201d73cd4ca,1 +np.float64,0x3fb3b65dd6276cc0,0x3ff7e64576199505,1 +np.float64,0x3fcd47a7793a8f50,0x3ff570a7b672f160,1 +np.float64,0xbfa41dd30c283ba0,0x3ff9c2f488127eb3,1 +np.float64,0x3fe4b1ea1f6963d4,0x3febc2bed7760427,1 +np.float64,0xbfdd0f81d2ba1f04,0x400056463724b768,1 +np.float64,0x3fd15d93f7a2bb28,0x3ff4bc7a24eacfd7,1 +np.float64,0xbfe3213af8e64276,0x4001b14579dfded3,1 +np.float64,0x3fd90dfbeab21bf8,0x3ff2b26a6c2c3bb3,1 +np.float64,0xbfd02d54bca05aaa,0x3ffd38ab3886b203,1 +np.float64,0x3fc218dcad2431b8,0x3ff6dced56d5b417,1 +np.float64,0x3fea5edf71f4bdbe,0x3fe3455ee09f27e6,1 +np.float64,0x3fa74319042e8640,0x3ff867d224545438,1 +np.float64,0x3fd970ad92b2e15c,0x3ff2979084815dc1,1 +np.float64,0x3fce0a4bf73c1498,0x3ff557a4df32df3e,1 +np.float64,0x3fef5c8e10feb91c,0x3fc99ca0eeaaebe4,1 +np.float64,0xbfedae997ffb5d33,0x400611af18f407ab,1 +np.float64,0xbfbcf07d6239e0f8,0x3ffaf201177a2d36,1 +np.float64,0xbfc3c52541278a4c,0x3ffb9d2af0408e4a,1 +np.float64,0x3fe4ef44e4e9de8a,0x3feb71f7331255e5,1 +np.float64,0xbfccd9f5f539b3ec,0x3ffcc53a99339592,1 +np.float64,0xbfda32c745b4658e,0x3fffe16e8727ef89,1 +np.float64,0xbfef54932a7ea926,0x40077e4605e61ca1,1 +np.float64,0x3fe9d4ae3573a95c,0x3fe4344a069a3fd0,1 +np.float64,0x3fda567e73b4acfc,0x3ff258bd77a663c7,1 +np.float64,0xbfd5bcac5eab7958,0x3ffead6379c19c52,1 +np.float64,0xbfee5e56f97cbcae,0x40069131fc54018d,1 +np.float64,0x3fc2d4413925a880,0x3ff6c54163816298,1 +np.float64,0xbfe9ddf6e873bbee,0x400418d8c722f7c5,1 +np.float64,0x3fdaf2a683b5e54c,0x3ff22dcda599d69c,1 +np.float64,0xbfca69789f34d2f0,0x3ffc7547ff10b1a6,1 +np.float64,0x3fed076f62fa0ede,0x3fdbcbda03c1d72a,1 +np.float64,0xbfcb38326f367064,0x3ffc8fb55dadeae5,1 +np.float64,0x3fe1938705e3270e,0x3fefa88130c5adda,1 +np.float64,0x3feaffae3b75ff5c,0x3fe221e3da537c7e,1 +np.float64,0x3fefc94acb7f9296,0x3fbd9a360ace67b4,1 +np.float64,0xbfe8bddeb0f17bbe,0x4003a316685c767e,1 +np.float64,0x3fbe10fbee3c21f0,0x3ff73fceb10650f5,1 +np.float64,0x3fde9126c1bd224c,0x3ff12a742f734d0a,1 +np.float64,0xbfe9686c91f2d0d9,0x4003e7bc6ee77906,1 +np.float64,0xbfb1ba4892237490,0x3ffa3dda064c2509,1 +np.float64,0xbfe2879100e50f22,0x400181c1a5b16f0f,1 +np.float64,0x3fd1cd40b6a39a80,0x3ff49f70e3064e95,1 +np.float64,0xbfc965869132cb0c,0x3ffc5419f3b43701,1 +np.float64,0x3fea7a6f2874f4de,0x3fe31480fb2dd862,1 +np.float64,0x3fc3bc56892778b0,0x3ff6a7e8fa0e8b0e,1 +np.float64,0x3fec1ed451f83da8,0x3fdfd78e564b8ad7,1 +np.float64,0x3feb77d16df6efa2,0x3fe13d083344e45e,1 +np.float64,0xbfe822e7c67045d0,0x400367104a830cf6,1 +np.float64,0x8000000000000001,0x3ff921fb54442d18,1 +np.float64,0xbfd4900918a92012,0x3ffe5dc0e19737b4,1 +np.float64,0x3fed184187fa3084,0x3fdb7b7a39f234f4,1 +np.float64,0x3fecef846179df08,0x3fdc3cb2228c3682,1 +np.float64,0xbfe2d2aed165a55e,0x400198e21c5b861b,1 +np.float64,0x7ff0000000000000,0x7ff8000000000000,1 +np.float64,0xbfee9409a07d2813,0x4006bd358232d073,1 +np.float64,0xbfecedc2baf9db86,0x4005995df566fc21,1 +np.float64,0x3fe6d857396db0ae,0x3fe8d2cb8794aa99,1 +np.float64,0xbf9a579e7834af40,0x3ff98b5cc8021e1c,1 +np.float64,0x3fc664fefb2cca00,0x3ff651a664ccf8fa,1 +np.float64,0xbfe8a7aa0e714f54,0x40039a5b4df938a0,1 +np.float64,0xbfdf27d380be4fa8,0x4000a241074dbae6,1 +np.float64,0x3fe00ddf55e01bbe,0x3ff0b94eb1ea1851,1 +np.float64,0x3feb47edbff68fdc,0x3fe199822d075959,1 +np.float64,0x3fb4993822293270,0x3ff7d80c838186d0,1 +np.float64,0xbfca2cd1473459a4,0x3ffc6d88c8de3d0d,1 +np.float64,0xbfea7d9c7674fb39,0x40045e4559e9e52d,1 +np.float64,0x3fe0dce425e1b9c8,0x3ff04099cab23289,1 +np.float64,0x3fd6bb7e97ad76fc,0x3ff352a30434499c,1 +np.float64,0x3fd4a4f16da949e4,0x3ff3e0b07432c9aa,1 +np.float64,0x8000000000000000,0x3ff921fb54442d18,1 +np.float64,0x3fe688f5b56d11ec,0x3fe9435f63264375,1 +np.float64,0xbfdf5a427ebeb484,0x4000a97a6c5d4abc,1 +np.float64,0xbfd1f3483fa3e690,0x3ffdae6c8a299383,1 +np.float64,0xbfeac920db759242,0x4004805862be51ec,1 +np.float64,0x3fef5bc711feb78e,0x3fc9ac40fba5b93b,1 +np.float64,0x3fe4bd9e12e97b3c,0x3febb363c787d381,1 +np.float64,0x3fef6a59ab7ed4b4,0x3fc880f1324eafce,1 +np.float64,0x3fc07a362120f470,0x3ff7113cf2c672b3,1 +np.float64,0xbfe4d6dbe2e9adb8,0x40023d6f6bea44b7,1 +np.float64,0xbfec2d6a15785ad4,0x40052eb425cc37a2,1 +np.float64,0x3fc90dae05321b60,0x3ff5fb10015d2934,1 +np.float64,0xbfa9239f74324740,0x3ff9eb2d057068ea,1 +np.float64,0xbfeb4fc8baf69f92,0x4004bf5e17fb08a4,1 +np.float64,0x0,0x3ff921fb54442d18,1 +np.float64,0x3faaf1884c35e320,0x3ff84a5591dbe1f3,1 +np.float64,0xbfed842561fb084b,0x4005f5c0a19116ce,1 +np.float64,0xbfc64850c32c90a0,0x3ffbeeac2ee70f9a,1 +np.float64,0x3fd7d879f5afb0f4,0x3ff306254c453436,1 +np.float64,0xbfdabaa586b5754c,0x4000035e6ac83a2b,1 +np.float64,0xbfebfeefa977fddf,0x4005167446fb9faf,1 +np.float64,0xbfe9383462727069,0x4003d407aa6a1577,1 +np.float64,0x3fe108dfb6e211c0,0x3ff026ac924b281d,1 +np.float64,0xbf85096df02a12c0,0x3ff94c0e60a22ede,1 +np.float64,0xbfe3121cd566243a,0x4001ac8f90db5882,1 +np.float64,0xbfd227f62aa44fec,0x3ffdbc26bb175dcc,1 +np.float64,0x3fd931af2cb26360,0x3ff2a8b62dfe003c,1 +np.float64,0xbfd9b794e3b36f2a,0x3fffbfbc89ec013d,1 +np.float64,0x3fc89b2e6f313660,0x3ff609a6e67f15f2,1 +np.float64,0x3fc0b14a8f216298,0x3ff70a4b6905aad2,1 +np.float64,0xbfeda11a657b4235,0x400608b3f9fff574,1 +np.float64,0xbfed2ee9ec7a5dd4,0x4005c040b7c02390,1 +np.float64,0xbfef7819d8fef034,0x4007ac6bf75cf09d,1 +np.float64,0xbfcc4720fb388e40,0x3ffcb2666a00b336,1 +np.float64,0xbfe05dec4be0bbd8,0x4000dc8a25ca3760,1 +np.float64,0x3fb093416e212680,0x3ff81897b6d8b374,1 +np.float64,0xbfc6ab89332d5714,0x3ffbfb4559d143e7,1 +np.float64,0x3fc51948512a3290,0x3ff67bb9df662c0a,1 +np.float64,0x3fed4d94177a9b28,0x3fda76c92f0c0132,1 +np.float64,0x3fdd195fbeba32c0,0x3ff194a5586dd18e,1 +np.float64,0x3fe3f82799e7f050,0x3fecb354c2faf55c,1 +np.float64,0x3fecac2169f95842,0x3fdd7222296cb7a7,1 +np.float64,0x3fe3d3f36fe7a7e6,0x3fece18f45e30dd7,1 +np.float64,0x3fe31ff63d663fec,0x3fedc46c77d30c6a,1 +np.float64,0xbfe3120c83e62419,0x4001ac8a7c4aa742,1 +np.float64,0x3fe7c1a7976f8350,0x3fe77e4a9307c9f8,1 +np.float64,0x3fe226fe9de44dfe,0x3feef6c0f3cb00fa,1 +np.float64,0x3fd5c933baab9268,0x3ff3933e8a37de42,1 +np.float64,0x3feaa98496f5530a,0x3fe2c003832ebf21,1 +np.float64,0xbfc6f80a2f2df014,0x3ffc04fd54cb1317,1 +np.float64,0x3fde5e18d0bcbc30,0x3ff138f7b32a2ca3,1 +np.float64,0xbfe30c8dd566191c,0x4001aad4af935a78,1 +np.float64,0x3fbe8d196e3d1a30,0x3ff737fec8149ecc,1 +np.float64,0x3feaee6731f5dcce,0x3fe241fa42cce22d,1 +np.float64,0x3fef9cc46cff3988,0x3fc3f17b708dbdbb,1 +np.float64,0xbfdb181bdeb63038,0x4000103ecf405602,1 +np.float64,0xbfc58de0ed2b1bc0,0x3ffbd704c14e15cd,1 +np.float64,0xbfee05d5507c0bab,0x40064e480faba6d8,1 +np.float64,0x3fe27d0ffa64fa20,0x3fee8dc71ef79f2c,1 +np.float64,0xbfe4f7ad4c69ef5a,0x400248456cd09a07,1 +np.float64,0xbfe4843e91e9087d,0x4002225f3e139c84,1 +np.float64,0x3fe7158b9c6e2b18,0x3fe87ae845c5ba96,1 +np.float64,0xbfea64316074c863,0x400452fd2bc23a44,1 +np.float64,0xbfc9f3ae4133e75c,0x3ffc663d482afa42,1 +np.float64,0xbfd5e18513abc30a,0x3ffeb72fc76d7071,1 +np.float64,0xbfd52f6438aa5ec8,0x3ffe87e5b18041e5,1 +np.float64,0xbfea970650f52e0d,0x400469a4a6758154,1 +np.float64,0xbfe44321b7e88644,0x40020d404a2141b1,1 +np.float64,0x3fdf5a39bbbeb474,0x3ff0f10453059dbd,1 +np.float64,0xbfa1d4069423a810,0x3ff9b0a2eacd2ce2,1 +np.float64,0xbfc36d16a326da2c,0x3ffb92077d41d26a,1 +np.float64,0x1,0x3ff921fb54442d18,1 +np.float64,0x3feb232a79764654,0x3fe1df5beeb249d0,1 +np.float64,0xbfed2003d5fa4008,0x4005b737c2727583,1 +np.float64,0x3fd5b093a3ab6128,0x3ff399ca2db1d96d,1 +np.float64,0x3fca692c3d34d258,0x3ff5ceb86b79223e,1 +np.float64,0x3fd6bbdf89ad77c0,0x3ff3528916df652d,1 +np.float64,0xbfefdadd46ffb5bb,0x40085ee735e19f19,1 +np.float64,0x3feb69fb2676d3f6,0x3fe157ee0c15691e,1 +np.float64,0x3fe44c931f689926,0x3fec46b6f5e3f265,1 +np.float64,0xbfc43ddbcb287bb8,0x3ffbac71d268d74d,1 +np.float64,0x3fe6e16d43edc2da,0x3fe8c5cf0f0daa66,1 +np.float64,0x3fe489efc76913e0,0x3febf704ca1ac2a6,1 +np.float64,0xbfe590aadceb2156,0x40027b764205cf78,1 +np.float64,0xbf782e8aa0305d00,0x3ff93a29e81928ab,1 +np.float64,0x3fedcb80cffb9702,0x3fd7e5d1f98a418b,1 +np.float64,0x3fe075858060eb0c,0x3ff07d23ab46b60f,1 +np.float64,0x3fe62a68296c54d0,0x3fe9c77f7068043b,1 +np.float64,0x3feff16a3c7fe2d4,0x3fae8e8a739cc67a,1 +np.float64,0xbfd6ed93e3addb28,0x3ffefebab206fa99,1 +np.float64,0x3fe40d8ccf681b1a,0x3fec97e9cd29966d,1 +np.float64,0x3fd6408210ac8104,0x3ff3737a7d374107,1 +np.float64,0x3fec8023b8f90048,0x3fde35ebfb2b3afd,1 +np.float64,0xbfe13babd4627758,0x40011dae5c07c56b,1 +np.float64,0xbfd2183e61a4307c,0x3ffdb80dd747cfbe,1 +np.float64,0x3feae8eb1d75d1d6,0x3fe24c1f6e42ae77,1 +np.float64,0xbfea559b9c74ab37,0x40044c8e5e123b20,1 +np.float64,0xbfd12c9d57a2593a,0x3ffd7ac6222f561c,1 +np.float64,0x3fe32eb697e65d6e,0x3fedb202693875b6,1 +np.float64,0xbfde0808c3bc1012,0x4000794bd8616ea3,1 +np.float64,0x3fe14958a06292b2,0x3ff0007b40ac648a,1 +np.float64,0x3fe3d388a6e7a712,0x3fece21751a6dd7c,1 +np.float64,0x3fe7ad7897ef5af2,0x3fe79c5b3da302a7,1 +np.float64,0x3fec75527e78eaa4,0x3fde655de0cf0508,1 +np.float64,0x3fea920d4c75241a,0x3fe2ea48f031d908,1 +np.float64,0x7fefffffffffffff,0x7ff8000000000000,1 +np.float64,0xbfc17a68cb22f4d0,0x3ffb530925f41aa0,1 +np.float64,0xbfe1c93166e39263,0x400147f3cb435dec,1 +np.float64,0x3feb97c402f72f88,0x3fe0fe5b561bf869,1 +np.float64,0x3fb58ff5162b1ff0,0x3ff7c8933fa969dc,1 +np.float64,0x3fe68e2beded1c58,0x3fe93c075283703b,1 +np.float64,0xbf94564cc828aca0,0x3ff97355e5ee35db,1 +np.float64,0x3fd31061c9a620c4,0x3ff44b150ec96998,1 +np.float64,0xbfc7d0c89f2fa190,0x3ffc208bf4eddc4d,1 +np.float64,0x3fe5736f1d6ae6de,0x3feac18f84992d1e,1 +np.float64,0x3fdb62e480b6c5c8,0x3ff20ecfdc4afe7c,1 +np.float64,0xbfc417228b282e44,0x3ffba78afea35979,1 +np.float64,0x3f8f5ba1303eb780,0x3ff8e343714630ff,1 +np.float64,0x3fe8e99126f1d322,0x3fe5b6511d4c0798,1 +np.float64,0xbfe2ec08a1e5d812,0x4001a0bb28a85875,1 +np.float64,0x3fea3b46cf74768e,0x3fe383dceaa74296,1 +np.float64,0xbfe008b5ed60116c,0x4000c3d62c275d40,1 +np.float64,0xbfcd9f8a4b3b3f14,0x3ffcde98d6484202,1 +np.float64,0xbfdb5fb112b6bf62,0x40001a22137ef1c9,1 +np.float64,0xbfe9079565f20f2b,0x4003c0670c92e401,1 +np.float64,0xbfce250dc53c4a1c,0x3ffcefc2b3dc3332,1 +np.float64,0x3fe9ba85d373750c,0x3fe4607131b28773,1 +np.float64,0x10000000000000,0x3ff921fb54442d18,1 +np.float64,0xbfeb9ef42c773de8,0x4004e5f239203ad8,1 +np.float64,0xbfd6bf457dad7e8a,0x3ffef2563d87b18d,1 +np.float64,0x3fe4de9aa5e9bd36,0x3feb87f97defb04a,1 +np.float64,0x3fedb4f67cfb69ec,0x3fd8603c465bffac,1 +np.float64,0x3fe7b6d9506f6db2,0x3fe78e670c7bdb67,1 +np.float64,0x3fe071717460e2e2,0x3ff07f84472d9cc5,1 +np.float64,0xbfed2e79dbfa5cf4,0x4005bffc6f9ad24f,1 +np.float64,0x3febb8adc377715c,0x3fe0bcebfbd45900,1 +np.float64,0xbfee2cffd87c5a00,0x40066b20a037c478,1 +np.float64,0x3fef7e358d7efc6c,0x3fc6d0ba71a542a8,1 +np.float64,0xbfef027eef7e04fe,0x400723291cb00a7a,1 +np.float64,0x3fac96da34392dc0,0x3ff83d260a936c6a,1 +np.float64,0x3fe9dba94a73b752,0x3fe428736b94885e,1 +np.float64,0x3fed37581efa6eb0,0x3fdae49dcadf1d90,1 +np.float64,0xbfe6e61037edcc20,0x4002f23031b8d522,1 +np.float64,0xbfdea7204dbd4e40,0x40008fe1f37918b7,1 +np.float64,0x3feb9f8edb773f1e,0x3fe0eef20bd4387b,1 +np.float64,0x3feeb0b6ed7d616e,0x3fd25fb3b7a525d6,1 +np.float64,0xbfd7ce9061af9d20,0x3fff3b25d531aa2b,1 +np.float64,0xbfc806b509300d6c,0x3ffc2768743a8360,1 +np.float64,0xbfa283882c250710,0x3ff9b61fda28914a,1 +np.float64,0x3fdec70050bd8e00,0x3ff11b1d769b578f,1 +np.float64,0xbfc858a44930b148,0x3ffc31d6758b4721,1 +np.float64,0x3fdc321150b86424,0x3ff1d5504c3c91e4,1 +np.float64,0x3fd9416870b282d0,0x3ff2a46f3a850f5b,1 +np.float64,0x3fdd756968baead4,0x3ff17ac510a5573f,1 +np.float64,0xbfedfd632cfbfac6,0x400648345a2f89b0,1 +np.float64,0x3fd6874285ad0e84,0x3ff36098ebff763f,1 +np.float64,0x3fe6daacc9edb55a,0x3fe8cf75fae1e35f,1 +np.float64,0x3fe53f19766a7e32,0x3feb07d0e97cd55b,1 +np.float64,0x3fd13cc36ca27988,0x3ff4c4ff801b1faa,1 +np.float64,0x3fe4f21cbce9e43a,0x3feb6e34a72ef529,1 +np.float64,0xbfc21c1cc9243838,0x3ffb67726394ca89,1 +np.float64,0x3fe947a3f2728f48,0x3fe51eae4660e23c,1 +np.float64,0xbfce78cd653cf19c,0x3ffcfa89194b3f5e,1 +np.float64,0x3fe756f049eeade0,0x3fe81be7f2d399e2,1 +np.float64,0xbfcc727cf138e4f8,0x3ffcb7f547841bb0,1 +np.float64,0xbfc2d8d58f25b1ac,0x3ffb7f496cc72458,1 +np.float64,0xbfcfd0e4653fa1c8,0x3ffd26e1309bc80b,1 +np.float64,0xbfe2126c106424d8,0x40015e0e01db6a4a,1 +np.float64,0x3fe580e4306b01c8,0x3feaaf683ce51aa5,1 +np.float64,0x3fcea8a1b93d5140,0x3ff543456c0d28c7,1 +np.float64,0xfff0000000000000,0x7ff8000000000000,1 +np.float64,0xbfd9d5da72b3abb4,0x3fffc8013113f968,1 +np.float64,0xbfe1fdfcea63fbfa,0x400157def2e4808d,1 +np.float64,0xbfc0022e0720045c,0x3ffb239963e7cbf2,1 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test__exceptions.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test__exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..7da56a77eae206f4524357fdcbbed32788101d3c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test__exceptions.py @@ -0,0 +1,89 @@ +""" +Tests of the ._exceptions module. Primarily for exercising the __str__ methods. +""" + +import pickle + +import pytest +import numpy as np +from numpy.exceptions import AxisError + +_ArrayMemoryError = np._core._exceptions._ArrayMemoryError +_UFuncNoLoopError = np._core._exceptions._UFuncNoLoopError + +class TestArrayMemoryError: + def test_pickling(self): + """ Test that _ArrayMemoryError can be pickled """ + error = _ArrayMemoryError((1023,), np.dtype(np.uint8)) + res = pickle.loads(pickle.dumps(error)) + assert res._total_size == error._total_size + + def test_str(self): + e = _ArrayMemoryError((1023,), np.dtype(np.uint8)) + str(e) # not crashing is enough + + # testing these properties is easier than testing the full string repr + def test__size_to_string(self): + """ Test e._size_to_string """ + f = _ArrayMemoryError._size_to_string + Ki = 1024 + assert f(0) == '0 bytes' + assert f(1) == '1 bytes' + assert f(1023) == '1023 bytes' + assert f(Ki) == '1.00 KiB' + assert f(Ki+1) == '1.00 KiB' + assert f(10*Ki) == '10.0 KiB' + assert f(int(999.4*Ki)) == '999. KiB' + assert f(int(1023.4*Ki)) == '1023. KiB' + assert f(int(1023.5*Ki)) == '1.00 MiB' + assert f(Ki*Ki) == '1.00 MiB' + + # 1023.9999 Mib should round to 1 GiB + assert f(int(Ki*Ki*Ki*0.9999)) == '1.00 GiB' + assert f(Ki*Ki*Ki*Ki*Ki*Ki) == '1.00 EiB' + # larger than sys.maxsize, adding larger prefixes isn't going to help + # anyway. + assert f(Ki*Ki*Ki*Ki*Ki*Ki*123456) == '123456. EiB' + + def test__total_size(self): + """ Test e._total_size """ + e = _ArrayMemoryError((1,), np.dtype(np.uint8)) + assert e._total_size == 1 + + e = _ArrayMemoryError((2, 4), np.dtype((np.uint64, 16))) + assert e._total_size == 1024 + + +class TestUFuncNoLoopError: + def test_pickling(self): + """ Test that _UFuncNoLoopError can be pickled """ + assert isinstance(pickle.dumps(_UFuncNoLoopError), bytes) + + +@pytest.mark.parametrize("args", [ + (2, 1, None), + (2, 1, "test_prefix"), + ("test message",), +]) +class TestAxisError: + def test_attr(self, args): + """Validate attribute types.""" + exc = AxisError(*args) + if len(args) == 1: + assert exc.axis is None + assert exc.ndim is None + else: + axis, ndim, *_ = args + assert exc.axis == axis + assert exc.ndim == ndim + + def test_pickling(self, args): + """Test that `AxisError` can be pickled.""" + exc = AxisError(*args) + exc2 = pickle.loads(pickle.dumps(exc)) + + assert type(exc) is type(exc2) + for name in ("axis", "ndim", "args"): + attr1 = getattr(exc, name) + attr2 = getattr(exc2, name) + assert attr1 == attr2, name diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_mem_overlap.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_mem_overlap.py new file mode 100644 index 0000000000000000000000000000000000000000..88ccc36926bede33bc241483b0ca57bf67d05af3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_mem_overlap.py @@ -0,0 +1,933 @@ +import itertools +import pytest + +import numpy as np +from numpy._core._multiarray_tests import solve_diophantine, internal_overlap +from numpy._core import _umath_tests +from numpy.lib.stride_tricks import as_strided +from numpy.testing import ( + assert_, assert_raises, assert_equal, assert_array_equal + ) + + +ndims = 2 +size = 10 +shape = tuple([size] * ndims) + +MAY_SHARE_BOUNDS = 0 +MAY_SHARE_EXACT = -1 + + +def _indices_for_nelems(nelems): + """Returns slices of length nelems, from start onwards, in direction sign.""" + + if nelems == 0: + return [size // 2] # int index + + res = [] + for step in (1, 2): + for sign in (-1, 1): + start = size // 2 - nelems * step * sign // 2 + stop = start + nelems * step * sign + res.append(slice(start, stop, step * sign)) + + return res + + +def _indices_for_axis(): + """Returns (src, dst) pairs of indices.""" + + res = [] + for nelems in (0, 2, 3): + ind = _indices_for_nelems(nelems) + res.extend(itertools.product(ind, ind)) # all assignments of size "nelems" + + return res + + +def _indices(ndims): + """Returns ((axis0_src, axis0_dst), (axis1_src, axis1_dst), ... ) index pairs.""" + + ind = _indices_for_axis() + return itertools.product(ind, repeat=ndims) + + +def _check_assignment(srcidx, dstidx): + """Check assignment arr[dstidx] = arr[srcidx] works.""" + + arr = np.arange(np.prod(shape)).reshape(shape) + + cpy = arr.copy() + + cpy[dstidx] = arr[srcidx] + arr[dstidx] = arr[srcidx] + + assert_(np.all(arr == cpy), + 'assigning arr[%s] = arr[%s]' % (dstidx, srcidx)) + + +def test_overlapping_assignments(): + # Test automatically generated assignments which overlap in memory. + + inds = _indices(ndims) + + for ind in inds: + srcidx = tuple([a[0] for a in ind]) + dstidx = tuple([a[1] for a in ind]) + + _check_assignment(srcidx, dstidx) + + +@pytest.mark.slow +def test_diophantine_fuzz(): + # Fuzz test the diophantine solver + rng = np.random.RandomState(1234) + + max_int = np.iinfo(np.intp).max + + for ndim in range(10): + feasible_count = 0 + infeasible_count = 0 + + min_count = 500//(ndim + 1) + + while min(feasible_count, infeasible_count) < min_count: + # Ensure big and small integer problems + A_max = 1 + rng.randint(0, 11, dtype=np.intp)**6 + U_max = rng.randint(0, 11, dtype=np.intp)**6 + + A_max = min(max_int, A_max) + U_max = min(max_int-1, U_max) + + A = tuple(int(rng.randint(1, A_max+1, dtype=np.intp)) + for j in range(ndim)) + U = tuple(int(rng.randint(0, U_max+2, dtype=np.intp)) + for j in range(ndim)) + + b_ub = min(max_int-2, sum(a*ub for a, ub in zip(A, U))) + b = int(rng.randint(-1, b_ub+2, dtype=np.intp)) + + if ndim == 0 and feasible_count < min_count: + b = 0 + + X = solve_diophantine(A, U, b) + + if X is None: + # Check the simplified decision problem agrees + X_simplified = solve_diophantine(A, U, b, simplify=1) + assert_(X_simplified is None, (A, U, b, X_simplified)) + + # Check no solution exists (provided the problem is + # small enough so that brute force checking doesn't + # take too long) + ranges = tuple(range(0, a*ub+1, a) for a, ub in zip(A, U)) + + size = 1 + for r in ranges: + size *= len(r) + if size < 100000: + assert_(not any(sum(w) == b for w in itertools.product(*ranges))) + infeasible_count += 1 + else: + # Check the simplified decision problem agrees + X_simplified = solve_diophantine(A, U, b, simplify=1) + assert_(X_simplified is not None, (A, U, b, X_simplified)) + + # Check validity + assert_(sum(a*x for a, x in zip(A, X)) == b) + assert_(all(0 <= x <= ub for x, ub in zip(X, U))) + feasible_count += 1 + + +def test_diophantine_overflow(): + # Smoke test integer overflow detection + max_intp = np.iinfo(np.intp).max + max_int64 = np.iinfo(np.int64).max + + if max_int64 <= max_intp: + # Check that the algorithm works internally in 128-bit; + # solving this problem requires large intermediate numbers + A = (max_int64//2, max_int64//2 - 10) + U = (max_int64//2, max_int64//2 - 10) + b = 2*(max_int64//2) - 10 + + assert_equal(solve_diophantine(A, U, b), (1, 1)) + + +def check_may_share_memory_exact(a, b): + got = np.may_share_memory(a, b, max_work=MAY_SHARE_EXACT) + + assert_equal(np.may_share_memory(a, b), + np.may_share_memory(a, b, max_work=MAY_SHARE_BOUNDS)) + + a.fill(0) + b.fill(0) + a.fill(1) + exact = b.any() + + err_msg = "" + if got != exact: + err_msg = " " + "\n ".join([ + "base_a - base_b = %r" % (a.__array_interface__['data'][0] - b.__array_interface__['data'][0],), + "shape_a = %r" % (a.shape,), + "shape_b = %r" % (b.shape,), + "strides_a = %r" % (a.strides,), + "strides_b = %r" % (b.strides,), + "size_a = %r" % (a.size,), + "size_b = %r" % (b.size,) + ]) + + assert_equal(got, exact, err_msg=err_msg) + + +def test_may_share_memory_manual(): + # Manual test cases for may_share_memory + + # Base arrays + xs0 = [ + np.zeros([13, 21, 23, 22], dtype=np.int8), + np.zeros([13, 21, 23*2, 22], dtype=np.int8)[:,:,::2,:] + ] + + # Generate all negative stride combinations + xs = [] + for x in xs0: + for ss in itertools.product(*(([slice(None), slice(None, None, -1)],)*4)): + xp = x[ss] + xs.append(xp) + + for x in xs: + # The default is a simple extent check + assert_(np.may_share_memory(x[:,0,:], x[:,1,:])) + assert_(np.may_share_memory(x[:,0,:], x[:,1,:], max_work=None)) + + # Exact checks + check_may_share_memory_exact(x[:,0,:], x[:,1,:]) + check_may_share_memory_exact(x[:,::7], x[:,3::3]) + + try: + xp = x.ravel() + if xp.flags.owndata: + continue + xp = xp.view(np.int16) + except ValueError: + continue + + # 0-size arrays cannot overlap + check_may_share_memory_exact(x.ravel()[6:6], + xp.reshape(13, 21, 23, 11)[:,::7]) + + # Test itemsize is dealt with + check_may_share_memory_exact(x[:,::7], + xp.reshape(13, 21, 23, 11)) + check_may_share_memory_exact(x[:,::7], + xp.reshape(13, 21, 23, 11)[:,3::3]) + check_may_share_memory_exact(x.ravel()[6:7], + xp.reshape(13, 21, 23, 11)[:,::7]) + + # Check unit size + x = np.zeros([1], dtype=np.int8) + check_may_share_memory_exact(x, x) + check_may_share_memory_exact(x, x.copy()) + + +def iter_random_view_pairs(x, same_steps=True, equal_size=False): + rng = np.random.RandomState(1234) + + if equal_size and same_steps: + raise ValueError + + def random_slice(n, step): + start = rng.randint(0, n+1, dtype=np.intp) + stop = rng.randint(start, n+1, dtype=np.intp) + if rng.randint(0, 2, dtype=np.intp) == 0: + stop, start = start, stop + step *= -1 + return slice(start, stop, step) + + def random_slice_fixed_size(n, step, size): + start = rng.randint(0, n+1 - size*step) + stop = start + (size-1)*step + 1 + if rng.randint(0, 2) == 0: + stop, start = start-1, stop-1 + if stop < 0: + stop = None + step *= -1 + return slice(start, stop, step) + + # First a few regular views + yield x, x + for j in range(1, 7, 3): + yield x[j:], x[:-j] + yield x[...,j:], x[...,:-j] + + # An array with zero stride internal overlap + strides = list(x.strides) + strides[0] = 0 + xp = as_strided(x, shape=x.shape, strides=strides) + yield x, xp + yield xp, xp + + # An array with non-zero stride internal overlap + strides = list(x.strides) + if strides[0] > 1: + strides[0] = 1 + xp = as_strided(x, shape=x.shape, strides=strides) + yield x, xp + yield xp, xp + + # Then discontiguous views + while True: + steps = tuple(rng.randint(1, 11, dtype=np.intp) + if rng.randint(0, 5, dtype=np.intp) == 0 else 1 + for j in range(x.ndim)) + s1 = tuple(random_slice(p, s) for p, s in zip(x.shape, steps)) + + t1 = np.arange(x.ndim) + rng.shuffle(t1) + + if equal_size: + t2 = t1 + else: + t2 = np.arange(x.ndim) + rng.shuffle(t2) + + a = x[s1] + + if equal_size: + if a.size == 0: + continue + + steps2 = tuple(rng.randint(1, max(2, p//(1+pa))) + if rng.randint(0, 5) == 0 else 1 + for p, s, pa in zip(x.shape, s1, a.shape)) + s2 = tuple(random_slice_fixed_size(p, s, pa) + for p, s, pa in zip(x.shape, steps2, a.shape)) + elif same_steps: + steps2 = steps + else: + steps2 = tuple(rng.randint(1, 11, dtype=np.intp) + if rng.randint(0, 5, dtype=np.intp) == 0 else 1 + for j in range(x.ndim)) + + if not equal_size: + s2 = tuple(random_slice(p, s) for p, s in zip(x.shape, steps2)) + + a = a.transpose(t1) + b = x[s2].transpose(t2) + + yield a, b + + +def check_may_share_memory_easy_fuzz(get_max_work, same_steps, min_count): + # Check that overlap problems with common strides are solved with + # little work. + x = np.zeros([17,34,71,97], dtype=np.int16) + + feasible = 0 + infeasible = 0 + + pair_iter = iter_random_view_pairs(x, same_steps) + + while min(feasible, infeasible) < min_count: + a, b = next(pair_iter) + + bounds_overlap = np.may_share_memory(a, b) + may_share_answer = np.may_share_memory(a, b) + easy_answer = np.may_share_memory(a, b, max_work=get_max_work(a, b)) + exact_answer = np.may_share_memory(a, b, max_work=MAY_SHARE_EXACT) + + if easy_answer != exact_answer: + # assert_equal is slow... + assert_equal(easy_answer, exact_answer) + + if may_share_answer != bounds_overlap: + assert_equal(may_share_answer, bounds_overlap) + + if bounds_overlap: + if exact_answer: + feasible += 1 + else: + infeasible += 1 + + +@pytest.mark.slow +def test_may_share_memory_easy_fuzz(): + # Check that overlap problems with common strides are always + # solved with little work. + + check_may_share_memory_easy_fuzz(get_max_work=lambda a, b: 1, + same_steps=True, + min_count=2000) + + +@pytest.mark.slow +def test_may_share_memory_harder_fuzz(): + # Overlap problems with not necessarily common strides take more + # work. + # + # The work bound below can't be reduced much. Harder problems can + # also exist but not be detected here, as the set of problems + # comes from RNG. + + check_may_share_memory_easy_fuzz(get_max_work=lambda a, b: max(a.size, b.size)//2, + same_steps=False, + min_count=2000) + + +def test_shares_memory_api(): + x = np.zeros([4, 5, 6], dtype=np.int8) + + assert_equal(np.shares_memory(x, x), True) + assert_equal(np.shares_memory(x, x.copy()), False) + + a = x[:,::2,::3] + b = x[:,::3,::2] + assert_equal(np.shares_memory(a, b), True) + assert_equal(np.shares_memory(a, b, max_work=None), True) + assert_raises( + np.exceptions.TooHardError, np.shares_memory, a, b, max_work=1 + ) + + +def test_may_share_memory_bad_max_work(): + x = np.zeros([1]) + assert_raises(OverflowError, np.may_share_memory, x, x, max_work=10**100) + assert_raises(OverflowError, np.shares_memory, x, x, max_work=10**100) + + +def test_internal_overlap_diophantine(): + def check(A, U, exists=None): + X = solve_diophantine(A, U, 0, require_ub_nontrivial=1) + + if exists is None: + exists = (X is not None) + + if X is not None: + assert_(sum(a*x for a, x in zip(A, X)) == sum(a*u//2 for a, u in zip(A, U))) + assert_(all(0 <= x <= u for x, u in zip(X, U))) + assert_(any(x != u//2 for x, u in zip(X, U))) + + if exists: + assert_(X is not None, repr(X)) + else: + assert_(X is None, repr(X)) + + # Smoke tests + check((3, 2), (2*2, 3*2), exists=True) + check((3*2, 2), (15*2, (3-1)*2), exists=False) + + +def test_internal_overlap_slices(): + # Slicing an array never generates internal overlap + + x = np.zeros([17,34,71,97], dtype=np.int16) + + rng = np.random.RandomState(1234) + + def random_slice(n, step): + start = rng.randint(0, n+1, dtype=np.intp) + stop = rng.randint(start, n+1, dtype=np.intp) + if rng.randint(0, 2, dtype=np.intp) == 0: + stop, start = start, stop + step *= -1 + return slice(start, stop, step) + + cases = 0 + min_count = 5000 + + while cases < min_count: + steps = tuple(rng.randint(1, 11, dtype=np.intp) + if rng.randint(0, 5, dtype=np.intp) == 0 else 1 + for j in range(x.ndim)) + t1 = np.arange(x.ndim) + rng.shuffle(t1) + s1 = tuple(random_slice(p, s) for p, s in zip(x.shape, steps)) + a = x[s1].transpose(t1) + + assert_(not internal_overlap(a)) + cases += 1 + + +def check_internal_overlap(a, manual_expected=None): + got = internal_overlap(a) + + # Brute-force check + m = set() + ranges = tuple(range(n) for n in a.shape) + for v in itertools.product(*ranges): + offset = sum(s*w for s, w in zip(a.strides, v)) + if offset in m: + expected = True + break + else: + m.add(offset) + else: + expected = False + + # Compare + if got != expected: + assert_equal(got, expected, err_msg=repr((a.strides, a.shape))) + if manual_expected is not None and expected != manual_expected: + assert_equal(expected, manual_expected) + return got + + +def test_internal_overlap_manual(): + # Stride tricks can construct arrays with internal overlap + + # We don't care about memory bounds, the array is not + # read/write accessed + x = np.arange(1).astype(np.int8) + + # Check low-dimensional special cases + + check_internal_overlap(x, False) # 1-dim + check_internal_overlap(x.reshape([]), False) # 0-dim + + a = as_strided(x, strides=(3, 4), shape=(4, 4)) + check_internal_overlap(a, False) + + a = as_strided(x, strides=(3, 4), shape=(5, 4)) + check_internal_overlap(a, True) + + a = as_strided(x, strides=(0,), shape=(0,)) + check_internal_overlap(a, False) + + a = as_strided(x, strides=(0,), shape=(1,)) + check_internal_overlap(a, False) + + a = as_strided(x, strides=(0,), shape=(2,)) + check_internal_overlap(a, True) + + a = as_strided(x, strides=(0, -9993), shape=(87, 22)) + check_internal_overlap(a, True) + + a = as_strided(x, strides=(0, -9993), shape=(1, 22)) + check_internal_overlap(a, False) + + a = as_strided(x, strides=(0, -9993), shape=(0, 22)) + check_internal_overlap(a, False) + + +def test_internal_overlap_fuzz(): + # Fuzz check; the brute-force check is fairly slow + + x = np.arange(1).astype(np.int8) + + overlap = 0 + no_overlap = 0 + min_count = 100 + + rng = np.random.RandomState(1234) + + while min(overlap, no_overlap) < min_count: + ndim = rng.randint(1, 4, dtype=np.intp) + + strides = tuple(rng.randint(-1000, 1000, dtype=np.intp) + for j in range(ndim)) + shape = tuple(rng.randint(1, 30, dtype=np.intp) + for j in range(ndim)) + + a = as_strided(x, strides=strides, shape=shape) + result = check_internal_overlap(a) + + if result: + overlap += 1 + else: + no_overlap += 1 + + +def test_non_ndarray_inputs(): + # Regression check for gh-5604 + + class MyArray: + def __init__(self, data): + self.data = data + + @property + def __array_interface__(self): + return self.data.__array_interface__ + + class MyArray2: + def __init__(self, data): + self.data = data + + def __array__(self, dtype=None, copy=None): + return self.data + + for cls in [MyArray, MyArray2]: + x = np.arange(5) + + assert_(np.may_share_memory(cls(x[::2]), x[1::2])) + assert_(not np.shares_memory(cls(x[::2]), x[1::2])) + + assert_(np.shares_memory(cls(x[1::3]), x[::2])) + assert_(np.may_share_memory(cls(x[1::3]), x[::2])) + + +def view_element_first_byte(x): + """Construct an array viewing the first byte of each element of `x`""" + from numpy.lib._stride_tricks_impl import DummyArray + interface = dict(x.__array_interface__) + interface['typestr'] = '|b1' + interface['descr'] = [('', '|b1')] + return np.asarray(DummyArray(interface, x)) + + +def assert_copy_equivalent(operation, args, out, **kwargs): + """ + Check that operation(*args, out=out) produces results + equivalent to out[...] = operation(*args, out=out.copy()) + """ + + kwargs['out'] = out + kwargs2 = dict(kwargs) + kwargs2['out'] = out.copy() + + out_orig = out.copy() + out[...] = operation(*args, **kwargs2) + expected = out.copy() + out[...] = out_orig + + got = operation(*args, **kwargs).copy() + + if (got != expected).any(): + assert_equal(got, expected) + + +class TestUFunc: + """ + Test ufunc call memory overlap handling + """ + + def check_unary_fuzz(self, operation, get_out_axis_size, dtype=np.int16, + count=5000): + shapes = [7, 13, 8, 21, 29, 32] + + rng = np.random.RandomState(1234) + + for ndim in range(1, 6): + x = rng.randint(0, 2**16, size=shapes[:ndim]).astype(dtype) + + it = iter_random_view_pairs(x, same_steps=False, equal_size=True) + + min_count = count // (ndim + 1)**2 + + overlapping = 0 + while overlapping < min_count: + a, b = next(it) + + a_orig = a.copy() + b_orig = b.copy() + + if get_out_axis_size is None: + assert_copy_equivalent(operation, [a], out=b) + + if np.shares_memory(a, b): + overlapping += 1 + else: + for axis in itertools.chain(range(ndim), [None]): + a[...] = a_orig + b[...] = b_orig + + # Determine size for reduction axis (None if scalar) + outsize, scalarize = get_out_axis_size(a, b, axis) + if outsize == 'skip': + continue + + # Slice b to get an output array of the correct size + sl = [slice(None)] * ndim + if axis is None: + if outsize is None: + sl = [slice(0, 1)] + [0]*(ndim - 1) + else: + sl = [slice(0, outsize)] + [0]*(ndim - 1) + else: + if outsize is None: + k = b.shape[axis]//2 + if ndim == 1: + sl[axis] = slice(k, k + 1) + else: + sl[axis] = k + else: + assert b.shape[axis] >= outsize + sl[axis] = slice(0, outsize) + b_out = b[tuple(sl)] + + if scalarize: + b_out = b_out.reshape([]) + + if np.shares_memory(a, b_out): + overlapping += 1 + + # Check result + assert_copy_equivalent(operation, [a], out=b_out, axis=axis) + + @pytest.mark.slow + def test_unary_ufunc_call_fuzz(self): + self.check_unary_fuzz(np.invert, None, np.int16) + + @pytest.mark.slow + def test_unary_ufunc_call_complex_fuzz(self): + # Complex typically has a smaller alignment than itemsize + self.check_unary_fuzz(np.negative, None, np.complex128, count=500) + + def test_binary_ufunc_accumulate_fuzz(self): + def get_out_axis_size(a, b, axis): + if axis is None: + if a.ndim == 1: + return a.size, False + else: + return 'skip', False # accumulate doesn't support this + else: + return a.shape[axis], False + + self.check_unary_fuzz(np.add.accumulate, get_out_axis_size, + dtype=np.int16, count=500) + + def test_binary_ufunc_reduce_fuzz(self): + def get_out_axis_size(a, b, axis): + return None, (axis is None or a.ndim == 1) + + self.check_unary_fuzz(np.add.reduce, get_out_axis_size, + dtype=np.int16, count=500) + + def test_binary_ufunc_reduceat_fuzz(self): + def get_out_axis_size(a, b, axis): + if axis is None: + if a.ndim == 1: + return a.size, False + else: + return 'skip', False # reduceat doesn't support this + else: + return a.shape[axis], False + + def do_reduceat(a, out, axis): + if axis is None: + size = len(a) + step = size//len(out) + else: + size = a.shape[axis] + step = a.shape[axis] // out.shape[axis] + idx = np.arange(0, size, step) + return np.add.reduceat(a, idx, out=out, axis=axis) + + self.check_unary_fuzz(do_reduceat, get_out_axis_size, + dtype=np.int16, count=500) + + def test_binary_ufunc_reduceat_manual(self): + def check(ufunc, a, ind, out): + c1 = ufunc.reduceat(a.copy(), ind.copy(), out=out.copy()) + c2 = ufunc.reduceat(a, ind, out=out) + assert_array_equal(c1, c2) + + # Exactly same input/output arrays + a = np.arange(10000, dtype=np.int16) + check(np.add, a, a[::-1].copy(), a) + + # Overlap with index + a = np.arange(10000, dtype=np.int16) + check(np.add, a, a[::-1], a) + + @pytest.mark.slow + def test_unary_gufunc_fuzz(self): + shapes = [7, 13, 8, 21, 29, 32] + gufunc = _umath_tests.euclidean_pdist + + rng = np.random.RandomState(1234) + + for ndim in range(2, 6): + x = rng.rand(*shapes[:ndim]) + + it = iter_random_view_pairs(x, same_steps=False, equal_size=True) + + min_count = 500 // (ndim + 1)**2 + + overlapping = 0 + while overlapping < min_count: + a, b = next(it) + + if min(a.shape[-2:]) < 2 or min(b.shape[-2:]) < 2 or a.shape[-1] < 2: + continue + + # Ensure the shapes are so that euclidean_pdist is happy + if b.shape[-1] > b.shape[-2]: + b = b[...,0,:] + else: + b = b[...,:,0] + + n = a.shape[-2] + p = n * (n - 1) // 2 + if p <= b.shape[-1] and p > 0: + b = b[...,:p] + else: + n = max(2, int(np.sqrt(b.shape[-1]))//2) + p = n * (n - 1) // 2 + a = a[...,:n,:] + b = b[...,:p] + + # Call + if np.shares_memory(a, b): + overlapping += 1 + + with np.errstate(over='ignore', invalid='ignore'): + assert_copy_equivalent(gufunc, [a], out=b) + + def test_ufunc_at_manual(self): + def check(ufunc, a, ind, b=None): + a0 = a.copy() + if b is None: + ufunc.at(a0, ind.copy()) + c1 = a0.copy() + ufunc.at(a, ind) + c2 = a.copy() + else: + ufunc.at(a0, ind.copy(), b.copy()) + c1 = a0.copy() + ufunc.at(a, ind, b) + c2 = a.copy() + assert_array_equal(c1, c2) + + # Overlap with index + a = np.arange(10000, dtype=np.int16) + check(np.invert, a[::-1], a) + + # Overlap with second data array + a = np.arange(100, dtype=np.int16) + ind = np.arange(0, 100, 2, dtype=np.int16) + check(np.add, a, ind, a[25:75]) + + def test_unary_ufunc_1d_manual(self): + # Exercise ufunc fast-paths (that avoid creation of an `np.nditer`) + + def check(a, b): + a_orig = a.copy() + b_orig = b.copy() + + b0 = b.copy() + c1 = ufunc(a, out=b0) + c2 = ufunc(a, out=b) + assert_array_equal(c1, c2) + + # Trigger "fancy ufunc loop" code path + mask = view_element_first_byte(b).view(np.bool) + + a[...] = a_orig + b[...] = b_orig + c1 = ufunc(a, out=b.copy(), where=mask.copy()).copy() + + a[...] = a_orig + b[...] = b_orig + c2 = ufunc(a, out=b, where=mask.copy()).copy() + + # Also, mask overlapping with output + a[...] = a_orig + b[...] = b_orig + c3 = ufunc(a, out=b, where=mask).copy() + + assert_array_equal(c1, c2) + assert_array_equal(c1, c3) + + dtypes = [np.int8, np.int16, np.int32, np.int64, np.float32, + np.float64, np.complex64, np.complex128] + dtypes = [np.dtype(x) for x in dtypes] + + for dtype in dtypes: + if np.issubdtype(dtype, np.integer): + ufunc = np.invert + else: + ufunc = np.reciprocal + + n = 1000 + k = 10 + indices = [ + np.index_exp[:n], + np.index_exp[k:k+n], + np.index_exp[n-1::-1], + np.index_exp[k+n-1:k-1:-1], + np.index_exp[:2*n:2], + np.index_exp[k:k+2*n:2], + np.index_exp[2*n-1::-2], + np.index_exp[k+2*n-1:k-1:-2], + ] + + for xi, yi in itertools.product(indices, indices): + v = np.arange(1, 1 + n*2 + k, dtype=dtype) + x = v[xi] + y = v[yi] + + with np.errstate(all='ignore'): + check(x, y) + + # Scalar cases + check(x[:1], y) + check(x[-1:], y) + check(x[:1].reshape([]), y) + check(x[-1:].reshape([]), y) + + def test_unary_ufunc_where_same(self): + # Check behavior at wheremask overlap + ufunc = np.invert + + def check(a, out, mask): + c1 = ufunc(a, out=out.copy(), where=mask.copy()) + c2 = ufunc(a, out=out, where=mask) + assert_array_equal(c1, c2) + + # Check behavior with same input and output arrays + x = np.arange(100).astype(np.bool) + check(x, x, x) + check(x, x.copy(), x) + check(x, x, x.copy()) + + @pytest.mark.slow + def test_binary_ufunc_1d_manual(self): + ufunc = np.add + + def check(a, b, c): + c0 = c.copy() + c1 = ufunc(a, b, out=c0) + c2 = ufunc(a, b, out=c) + assert_array_equal(c1, c2) + + for dtype in [np.int8, np.int16, np.int32, np.int64, + np.float32, np.float64, np.complex64, np.complex128]: + # Check different data dependency orders + + n = 1000 + k = 10 + + indices = [] + for p in [1, 2]: + indices.extend([ + np.index_exp[:p*n:p], + np.index_exp[k:k+p*n:p], + np.index_exp[p*n-1::-p], + np.index_exp[k+p*n-1:k-1:-p], + ]) + + for x, y, z in itertools.product(indices, indices, indices): + v = np.arange(6*n).astype(dtype) + x = v[x] + y = v[y] + z = v[z] + + check(x, y, z) + + # Scalar cases + check(x[:1], y, z) + check(x[-1:], y, z) + check(x[:1].reshape([]), y, z) + check(x[-1:].reshape([]), y, z) + check(x, y[:1], z) + check(x, y[-1:], z) + check(x, y[:1].reshape([]), z) + check(x, y[-1:].reshape([]), z) + + def test_inplace_op_simple_manual(self): + rng = np.random.RandomState(1234) + x = rng.rand(200, 200) # bigger than bufsize + + x += x.T + assert_array_equal(x - x.T, 0) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_mem_policy.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_mem_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..57826c18ccfe95e6e273ddfdde4014a02e82e182 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_mem_policy.py @@ -0,0 +1,449 @@ +import asyncio +import gc +import os +import sys +import threading + +import pytest + +import numpy as np +from numpy.testing import extbuild, assert_warns, IS_WASM, IS_EDITABLE +from numpy._core.multiarray import get_handler_name + + +@pytest.fixture +def get_module(tmp_path): + """ Add a memory policy that returns a false pointer 64 bytes into the + actual allocation, and fill the prefix with some text. Then check at each + memory manipulation that the prefix exists, to make sure all alloc/realloc/ + free/calloc go via the functions here. + """ + if sys.platform.startswith('cygwin'): + pytest.skip('link fails on cygwin') + if IS_WASM: + pytest.skip("Can't build module inside Wasm") + if IS_EDITABLE: + pytest.skip("Can't build module for editable install") + + functions = [ + ("get_default_policy", "METH_NOARGS", """ + Py_INCREF(PyDataMem_DefaultHandler); + return PyDataMem_DefaultHandler; + """), + ("set_secret_data_policy", "METH_NOARGS", """ + PyObject *secret_data = + PyCapsule_New(&secret_data_handler, "mem_handler", NULL); + if (secret_data == NULL) { + return NULL; + } + PyObject *old = PyDataMem_SetHandler(secret_data); + Py_DECREF(secret_data); + return old; + """), + ("set_wrong_capsule_name_data_policy", "METH_NOARGS", """ + PyObject *wrong_name_capsule = + PyCapsule_New(&secret_data_handler, "not_mem_handler", NULL); + if (wrong_name_capsule == NULL) { + return NULL; + } + PyObject *old = PyDataMem_SetHandler(wrong_name_capsule); + Py_DECREF(wrong_name_capsule); + return old; + """), + ("set_old_policy", "METH_O", """ + PyObject *old; + if (args != NULL && PyCapsule_CheckExact(args)) { + old = PyDataMem_SetHandler(args); + } + else { + old = PyDataMem_SetHandler(NULL); + } + return old; + """), + ("get_array", "METH_NOARGS", """ + char *buf = (char *)malloc(20); + npy_intp dims[1]; + dims[0] = 20; + PyArray_Descr *descr = PyArray_DescrNewFromType(NPY_UINT8); + return PyArray_NewFromDescr(&PyArray_Type, descr, 1, dims, NULL, + buf, NPY_ARRAY_WRITEABLE, NULL); + """), + ("set_own", "METH_O", """ + if (!PyArray_Check(args)) { + PyErr_SetString(PyExc_ValueError, + "need an ndarray"); + return NULL; + } + PyArray_ENABLEFLAGS((PyArrayObject*)args, NPY_ARRAY_OWNDATA); + // Maybe try this too? + // PyArray_BASE(PyArrayObject *)args) = NULL; + Py_RETURN_NONE; + """), + ("get_array_with_base", "METH_NOARGS", """ + char *buf = (char *)malloc(20); + npy_intp dims[1]; + dims[0] = 20; + PyArray_Descr *descr = PyArray_DescrNewFromType(NPY_UINT8); + PyObject *arr = PyArray_NewFromDescr(&PyArray_Type, descr, 1, dims, + NULL, buf, + NPY_ARRAY_WRITEABLE, NULL); + if (arr == NULL) return NULL; + PyObject *obj = PyCapsule_New(buf, "buf capsule", + (PyCapsule_Destructor)&warn_on_free); + if (obj == NULL) { + Py_DECREF(arr); + return NULL; + } + if (PyArray_SetBaseObject((PyArrayObject *)arr, obj) < 0) { + Py_DECREF(arr); + Py_DECREF(obj); + return NULL; + } + return arr; + + """), + ] + prologue = ''' + #define NPY_TARGET_VERSION NPY_1_22_API_VERSION + #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION + #include + /* + * This struct allows the dynamic configuration of the allocator funcs + * of the `secret_data_allocator`. It is provided here for + * demonstration purposes, as a valid `ctx` use-case scenario. + */ + typedef struct { + void *(*malloc)(size_t); + void *(*calloc)(size_t, size_t); + void *(*realloc)(void *, size_t); + void (*free)(void *); + } SecretDataAllocatorFuncs; + + NPY_NO_EXPORT void * + shift_alloc(void *ctx, size_t sz) { + SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx; + char *real = (char *)funcs->malloc(sz + 64); + if (real == NULL) { + return NULL; + } + snprintf(real, 64, "originally allocated %ld", (unsigned long)sz); + return (void *)(real + 64); + } + NPY_NO_EXPORT void * + shift_zero(void *ctx, size_t sz, size_t cnt) { + SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx; + char *real = (char *)funcs->calloc(sz + 64, cnt); + if (real == NULL) { + return NULL; + } + snprintf(real, 64, "originally allocated %ld via zero", + (unsigned long)sz); + return (void *)(real + 64); + } + NPY_NO_EXPORT void + shift_free(void *ctx, void * p, npy_uintp sz) { + SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx; + if (p == NULL) { + return ; + } + char *real = (char *)p - 64; + if (strncmp(real, "originally allocated", 20) != 0) { + fprintf(stdout, "uh-oh, unmatched shift_free, " + "no appropriate prefix\\n"); + /* Make C runtime crash by calling free on the wrong address */ + funcs->free((char *)p + 10); + /* funcs->free(real); */ + } + else { + npy_uintp i = (npy_uintp)atoi(real +20); + if (i != sz) { + fprintf(stderr, "uh-oh, unmatched shift_free" + "(ptr, %ld) but allocated %ld\\n", sz, i); + /* This happens in some places, only print */ + funcs->free(real); + } + else { + funcs->free(real); + } + } + } + NPY_NO_EXPORT void * + shift_realloc(void *ctx, void * p, npy_uintp sz) { + SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx; + if (p != NULL) { + char *real = (char *)p - 64; + if (strncmp(real, "originally allocated", 20) != 0) { + fprintf(stdout, "uh-oh, unmatched shift_realloc\\n"); + return realloc(p, sz); + } + return (void *)((char *)funcs->realloc(real, sz + 64) + 64); + } + else { + char *real = (char *)funcs->realloc(p, sz + 64); + if (real == NULL) { + return NULL; + } + snprintf(real, 64, "originally allocated " + "%ld via realloc", (unsigned long)sz); + return (void *)(real + 64); + } + } + /* As an example, we use the standard {m|c|re}alloc/free funcs. */ + static SecretDataAllocatorFuncs secret_data_handler_ctx = { + malloc, + calloc, + realloc, + free + }; + static PyDataMem_Handler secret_data_handler = { + "secret_data_allocator", + 1, + { + &secret_data_handler_ctx, /* ctx */ + shift_alloc, /* malloc */ + shift_zero, /* calloc */ + shift_realloc, /* realloc */ + shift_free /* free */ + } + }; + void warn_on_free(void *capsule) { + PyErr_WarnEx(PyExc_UserWarning, "in warn_on_free", 1); + void * obj = PyCapsule_GetPointer(capsule, + PyCapsule_GetName(capsule)); + free(obj); + }; + ''' + more_init = "import_array();" + try: + import mem_policy + return mem_policy + except ImportError: + pass + # if it does not exist, build and load it + return extbuild.build_and_import_extension('mem_policy', + functions, + prologue=prologue, + include_dirs=[np.get_include()], + build_dir=tmp_path, + more_init=more_init) + + +def test_set_policy(get_module): + + get_handler_name = np._core.multiarray.get_handler_name + get_handler_version = np._core.multiarray.get_handler_version + orig_policy_name = get_handler_name() + + a = np.arange(10).reshape((2, 5)) # a doesn't own its own data + assert get_handler_name(a) is None + assert get_handler_version(a) is None + assert get_handler_name(a.base) == orig_policy_name + assert get_handler_version(a.base) == 1 + + orig_policy = get_module.set_secret_data_policy() + + b = np.arange(10).reshape((2, 5)) # b doesn't own its own data + assert get_handler_name(b) is None + assert get_handler_version(b) is None + assert get_handler_name(b.base) == 'secret_data_allocator' + assert get_handler_version(b.base) == 1 + + if orig_policy_name == 'default_allocator': + get_module.set_old_policy(None) # tests PyDataMem_SetHandler(NULL) + assert get_handler_name() == 'default_allocator' + else: + get_module.set_old_policy(orig_policy) + assert get_handler_name() == orig_policy_name + + with pytest.raises(ValueError, + match="Capsule must be named 'mem_handler'"): + get_module.set_wrong_capsule_name_data_policy() + + +def test_default_policy_singleton(get_module): + get_handler_name = np._core.multiarray.get_handler_name + + # set the policy to default + orig_policy = get_module.set_old_policy(None) + + assert get_handler_name() == 'default_allocator' + + # re-set the policy to default + def_policy_1 = get_module.set_old_policy(None) + + assert get_handler_name() == 'default_allocator' + + # set the policy to original + def_policy_2 = get_module.set_old_policy(orig_policy) + + # since default policy is a singleton, + # these should be the same object + assert def_policy_1 is def_policy_2 is get_module.get_default_policy() + + +def test_policy_propagation(get_module): + # The memory policy goes hand-in-hand with flags.owndata + + class MyArr(np.ndarray): + pass + + get_handler_name = np._core.multiarray.get_handler_name + orig_policy_name = get_handler_name() + a = np.arange(10).view(MyArr).reshape((2, 5)) + assert get_handler_name(a) is None + assert a.flags.owndata is False + + assert get_handler_name(a.base) is None + assert a.base.flags.owndata is False + + assert get_handler_name(a.base.base) == orig_policy_name + assert a.base.base.flags.owndata is True + + +async def concurrent_context1(get_module, orig_policy_name, event): + if orig_policy_name == 'default_allocator': + get_module.set_secret_data_policy() + assert get_handler_name() == 'secret_data_allocator' + else: + get_module.set_old_policy(None) + assert get_handler_name() == 'default_allocator' + event.set() + + +async def concurrent_context2(get_module, orig_policy_name, event): + await event.wait() + # the policy is not affected by changes in parallel contexts + assert get_handler_name() == orig_policy_name + # change policy in the child context + if orig_policy_name == 'default_allocator': + get_module.set_secret_data_policy() + assert get_handler_name() == 'secret_data_allocator' + else: + get_module.set_old_policy(None) + assert get_handler_name() == 'default_allocator' + + +async def async_test_context_locality(get_module): + orig_policy_name = np._core.multiarray.get_handler_name() + + event = asyncio.Event() + # the child contexts inherit the parent policy + concurrent_task1 = asyncio.create_task( + concurrent_context1(get_module, orig_policy_name, event)) + concurrent_task2 = asyncio.create_task( + concurrent_context2(get_module, orig_policy_name, event)) + await concurrent_task1 + await concurrent_task2 + + # the parent context is not affected by child policy changes + assert np._core.multiarray.get_handler_name() == orig_policy_name + + +def test_context_locality(get_module): + if (sys.implementation.name == 'pypy' + and sys.pypy_version_info[:3] < (7, 3, 6)): + pytest.skip('no context-locality support in PyPy < 7.3.6') + asyncio.run(async_test_context_locality(get_module)) + + +def concurrent_thread1(get_module, event): + get_module.set_secret_data_policy() + assert np._core.multiarray.get_handler_name() == 'secret_data_allocator' + event.set() + + +def concurrent_thread2(get_module, event): + event.wait() + # the policy is not affected by changes in parallel threads + assert np._core.multiarray.get_handler_name() == 'default_allocator' + # change policy in the child thread + get_module.set_secret_data_policy() + + +def test_thread_locality(get_module): + orig_policy_name = np._core.multiarray.get_handler_name() + + event = threading.Event() + # the child threads do not inherit the parent policy + concurrent_task1 = threading.Thread(target=concurrent_thread1, + args=(get_module, event)) + concurrent_task2 = threading.Thread(target=concurrent_thread2, + args=(get_module, event)) + concurrent_task1.start() + concurrent_task2.start() + concurrent_task1.join() + concurrent_task2.join() + + # the parent thread is not affected by child policy changes + assert np._core.multiarray.get_handler_name() == orig_policy_name + + +@pytest.mark.skip(reason="too slow, see gh-23975") +def test_new_policy(get_module): + a = np.arange(10) + orig_policy_name = np._core.multiarray.get_handler_name(a) + + orig_policy = get_module.set_secret_data_policy() + + b = np.arange(10) + assert np._core.multiarray.get_handler_name(b) == 'secret_data_allocator' + + # test array manipulation. This is slow + if orig_policy_name == 'default_allocator': + # when the np._core.test tests recurse into this test, the + # policy will be set so this "if" will be false, preventing + # infinite recursion + # + # if needed, debug this by + # - running tests with -- -s (to not capture stdout/stderr + # - setting verbose=2 + # - setting extra_argv=['-vv'] here + assert np._core.test('full', verbose=1, extra_argv=[]) + # also try the ma tests, the pickling test is quite tricky + assert np.ma.test('full', verbose=1, extra_argv=[]) + + get_module.set_old_policy(orig_policy) + + c = np.arange(10) + assert np._core.multiarray.get_handler_name(c) == orig_policy_name + + +@pytest.mark.xfail(sys.implementation.name == "pypy", + reason=("bad interaction between getenv and " + "os.environ inside pytest")) +@pytest.mark.parametrize("policy", ["0", "1", None]) +def test_switch_owner(get_module, policy): + a = get_module.get_array() + assert np._core.multiarray.get_handler_name(a) is None + get_module.set_own(a) + + if policy is None: + # See what we expect to be set based on the env variable + policy = os.getenv("NUMPY_WARN_IF_NO_MEM_POLICY", "0") == "1" + oldval = None + else: + policy = policy == "1" + oldval = np._core._multiarray_umath._set_numpy_warn_if_no_mem_policy( + policy) + try: + # The policy should be NULL, so we have to assume we can call + # "free". A warning is given if the policy == "1" + if policy: + with assert_warns(RuntimeWarning) as w: + del a + gc.collect() + else: + del a + gc.collect() + + finally: + if oldval is not None: + np._core._multiarray_umath._set_numpy_warn_if_no_mem_policy(oldval) + + +def test_owner_is_base(get_module): + a = get_module.get_array_with_base() + with pytest.warns(UserWarning, match='warn_on_free'): + del a + gc.collect() + gc.collect() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_memmap.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_memmap.py new file mode 100644 index 0000000000000000000000000000000000000000..78b79f6ca0051a75f55f9bf0f62f804886f07c5f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_memmap.py @@ -0,0 +1,230 @@ +import sys +import os +import mmap +import pytest +from pathlib import Path +from tempfile import NamedTemporaryFile, TemporaryFile + +from numpy import ( + memmap, sum, average, prod, ndarray, isscalar, add, subtract, multiply) + +from numpy import arange, allclose, asarray +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, suppress_warnings, IS_PYPY, + break_cycles + ) + +class TestMemmap: + def setup_method(self): + self.tmpfp = NamedTemporaryFile(prefix='mmap') + self.shape = (3, 4) + self.dtype = 'float32' + self.data = arange(12, dtype=self.dtype) + self.data.resize(self.shape) + + def teardown_method(self): + self.tmpfp.close() + self.data = None + if IS_PYPY: + break_cycles() + break_cycles() + + def test_roundtrip(self): + # Write data to file + fp = memmap(self.tmpfp, dtype=self.dtype, mode='w+', + shape=self.shape) + fp[:] = self.data[:] + del fp # Test __del__ machinery, which handles cleanup + + # Read data back from file + newfp = memmap(self.tmpfp, dtype=self.dtype, mode='r', + shape=self.shape) + assert_(allclose(self.data, newfp)) + assert_array_equal(self.data, newfp) + assert_equal(newfp.flags.writeable, False) + + def test_open_with_filename(self, tmp_path): + tmpname = tmp_path / 'mmap' + fp = memmap(tmpname, dtype=self.dtype, mode='w+', + shape=self.shape) + fp[:] = self.data[:] + del fp + + def test_unnamed_file(self): + with TemporaryFile() as f: + fp = memmap(f, dtype=self.dtype, shape=self.shape) + del fp + + def test_attributes(self): + offset = 1 + mode = "w+" + fp = memmap(self.tmpfp, dtype=self.dtype, mode=mode, + shape=self.shape, offset=offset) + assert_equal(offset, fp.offset) + assert_equal(mode, fp.mode) + del fp + + def test_filename(self, tmp_path): + tmpname = tmp_path / "mmap" + fp = memmap(tmpname, dtype=self.dtype, mode='w+', + shape=self.shape) + abspath = Path(os.path.abspath(tmpname)) + fp[:] = self.data[:] + assert_equal(abspath, fp.filename) + b = fp[:1] + assert_equal(abspath, b.filename) + del b + del fp + + def test_path(self, tmp_path): + tmpname = tmp_path / "mmap" + fp = memmap(Path(tmpname), dtype=self.dtype, mode='w+', + shape=self.shape) + # os.path.realpath does not resolve symlinks on Windows + # see: https://bugs.python.org/issue9949 + # use Path.resolve, just as memmap class does internally + abspath = str(Path(tmpname).resolve()) + fp[:] = self.data[:] + assert_equal(abspath, str(fp.filename.resolve())) + b = fp[:1] + assert_equal(abspath, str(b.filename.resolve())) + del b + del fp + + def test_filename_fileobj(self): + fp = memmap(self.tmpfp, dtype=self.dtype, mode="w+", + shape=self.shape) + assert_equal(fp.filename, self.tmpfp.name) + + @pytest.mark.skipif(sys.platform == 'gnu0', + reason="Known to fail on hurd") + def test_flush(self): + fp = memmap(self.tmpfp, dtype=self.dtype, mode='w+', + shape=self.shape) + fp[:] = self.data[:] + assert_equal(fp[0], self.data[0]) + fp.flush() + + def test_del(self): + # Make sure a view does not delete the underlying mmap + fp_base = memmap(self.tmpfp, dtype=self.dtype, mode='w+', + shape=self.shape) + fp_base[0] = 5 + fp_view = fp_base[0:1] + assert_equal(fp_view[0], 5) + del fp_view + # Should still be able to access and assign values after + # deleting the view + assert_equal(fp_base[0], 5) + fp_base[0] = 6 + assert_equal(fp_base[0], 6) + + def test_arithmetic_drops_references(self): + fp = memmap(self.tmpfp, dtype=self.dtype, mode='w+', + shape=self.shape) + tmp = (fp + 10) + if isinstance(tmp, memmap): + assert_(tmp._mmap is not fp._mmap) + + def test_indexing_drops_references(self): + fp = memmap(self.tmpfp, dtype=self.dtype, mode='w+', + shape=self.shape) + tmp = fp[(1, 2), (2, 3)] + if isinstance(tmp, memmap): + assert_(tmp._mmap is not fp._mmap) + + def test_slicing_keeps_references(self): + fp = memmap(self.tmpfp, dtype=self.dtype, mode='w+', + shape=self.shape) + assert_(fp[:2, :2]._mmap is fp._mmap) + + def test_view(self): + fp = memmap(self.tmpfp, dtype=self.dtype, shape=self.shape) + new1 = fp.view() + new2 = new1.view() + assert_(new1.base is fp) + assert_(new2.base is fp) + new_array = asarray(fp) + assert_(new_array.base is fp) + + def test_ufunc_return_ndarray(self): + fp = memmap(self.tmpfp, dtype=self.dtype, shape=self.shape) + fp[:] = self.data + + with suppress_warnings() as sup: + sup.filter(FutureWarning, "np.average currently does not preserve") + for unary_op in [sum, average, prod]: + result = unary_op(fp) + assert_(isscalar(result)) + assert_(result.__class__ is self.data[0, 0].__class__) + + assert_(unary_op(fp, axis=0).__class__ is ndarray) + assert_(unary_op(fp, axis=1).__class__ is ndarray) + + for binary_op in [add, subtract, multiply]: + assert_(binary_op(fp, self.data).__class__ is ndarray) + assert_(binary_op(self.data, fp).__class__ is ndarray) + assert_(binary_op(fp, fp).__class__ is ndarray) + + fp += 1 + assert(fp.__class__ is memmap) + add(fp, 1, out=fp) + assert(fp.__class__ is memmap) + + def test_getitem(self): + fp = memmap(self.tmpfp, dtype=self.dtype, shape=self.shape) + fp[:] = self.data + + assert_(fp[1:, :-1].__class__ is memmap) + # Fancy indexing returns a copy that is not memmapped + assert_(fp[[0, 1]].__class__ is ndarray) + + def test_memmap_subclass(self): + class MemmapSubClass(memmap): + pass + + fp = MemmapSubClass(self.tmpfp, dtype=self.dtype, shape=self.shape) + fp[:] = self.data + + # We keep previous behavior for subclasses of memmap, i.e. the + # ufunc and __getitem__ output is never turned into a ndarray + assert_(sum(fp, axis=0).__class__ is MemmapSubClass) + assert_(sum(fp).__class__ is MemmapSubClass) + assert_(fp[1:, :-1].__class__ is MemmapSubClass) + assert(fp[[0, 1]].__class__ is MemmapSubClass) + + def test_mmap_offset_greater_than_allocation_granularity(self): + size = 5 * mmap.ALLOCATIONGRANULARITY + offset = mmap.ALLOCATIONGRANULARITY + 1 + fp = memmap(self.tmpfp, shape=size, mode='w+', offset=offset) + assert_(fp.offset == offset) + + def test_empty_array_with_offset_multiple_of_allocation_granularity(self): + self.tmpfp.write(b'a'*mmap.ALLOCATIONGRANULARITY) + size = 0 + offset = mmap.ALLOCATIONGRANULARITY + fp = memmap(self.tmpfp, shape=size, mode='w+', offset=offset) + assert_equal(fp.offset, offset) + + def test_no_shape(self): + self.tmpfp.write(b'a'*16) + mm = memmap(self.tmpfp, dtype='float64') + assert_equal(mm.shape, (2,)) + + def test_empty_array(self): + # gh-12653 + with pytest.raises(ValueError, match='empty file'): + memmap(self.tmpfp, shape=(0, 4), mode='r') + + # gh-27723 + # empty memmap works with mode in ('w+','r+') + memmap(self.tmpfp, shape=(0, 4), mode='w+') + + # ok now the file is not empty + memmap(self.tmpfp, shape=(0, 4), mode='w+') + + def test_shape_type(self): + memmap(self.tmpfp, shape=3, mode='w+') + memmap(self.tmpfp, shape=self.shape, mode='w+') + memmap(self.tmpfp, shape=list(self.shape), mode='w+') + memmap(self.tmpfp, shape=asarray(self.shape), mode='w+') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_multiarray.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_multiarray.py new file mode 100644 index 0000000000000000000000000000000000000000..955144e266bd60a30aebd21e1c6c0bf77e44b1c7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_multiarray.py @@ -0,0 +1,10381 @@ +from __future__ import annotations + +import collections.abc +import tempfile +import sys +import warnings +import operator +import io +import itertools +import functools +import ctypes +import os +import gc +import re +import weakref +import pytest +from contextlib import contextmanager +import pickle +import pathlib +import builtins +from decimal import Decimal +import mmap + +import numpy as np +import numpy._core._multiarray_tests as _multiarray_tests +from numpy._core._rational_tests import rational +from numpy.exceptions import AxisError, ComplexWarning +from numpy.testing import ( + assert_, assert_raises, assert_warns, assert_equal, assert_almost_equal, + assert_array_equal, assert_raises_regex, assert_array_almost_equal, + assert_allclose, IS_PYPY, IS_WASM, IS_PYSTON, HAS_REFCOUNT, + assert_array_less, runstring, temppath, suppress_warnings, break_cycles, + check_support_sve, assert_array_compare, IS_64BIT + ) +from numpy.testing._private.utils import requires_memory, _no_tracing +from numpy._core.tests._locales import CommaDecimalPointLocale +from numpy.lib.recfunctions import repack_fields +from numpy._core.multiarray import _get_ndarray_c_version, dot + +# Need to test an object that does not fully implement math interface +from datetime import timedelta, datetime + + +def assert_arg_sorted(arr, arg): + # resulting array should be sorted and arg values should be unique + assert_equal(arr[arg], np.sort(arr)) + assert_equal(np.sort(arg), np.arange(len(arg))) + + +def assert_arr_partitioned(kth, k, arr_part): + assert_equal(arr_part[k], kth) + assert_array_compare(operator.__le__, arr_part[:k], kth) + assert_array_compare(operator.__ge__, arr_part[k:], kth) + + +def _aligned_zeros(shape, dtype=float, order="C", align=None): + """ + Allocate a new ndarray with aligned memory. + + The ndarray is guaranteed *not* aligned to twice the requested alignment. + Eg, if align=4, guarantees it is not aligned to 8. If align=None uses + dtype.alignment.""" + dtype = np.dtype(dtype) + if dtype == np.dtype(object): + # Can't do this, fall back to standard allocation (which + # should always be sufficiently aligned) + if align is not None: + raise ValueError("object array alignment not supported") + return np.zeros(shape, dtype=dtype, order=order) + if align is None: + align = dtype.alignment + if not hasattr(shape, '__len__'): + shape = (shape,) + size = functools.reduce(operator.mul, shape) * dtype.itemsize + buf = np.empty(size + 2*align + 1, np.uint8) + + ptr = buf.__array_interface__['data'][0] + offset = ptr % align + if offset != 0: + offset = align - offset + if (ptr % (2*align)) == 0: + offset += align + + # Note: slices producing 0-size arrays do not necessarily change + # data pointer --- so we use and allocate size+1 + buf = buf[offset:offset+size+1][:-1] + buf.fill(0) + data = np.ndarray(shape, dtype, buf, order=order) + return data + + +class TestFlags: + def setup_method(self): + self.a = np.arange(10) + + def test_writeable(self): + mydict = locals() + self.a.flags.writeable = False + assert_raises(ValueError, runstring, 'self.a[0] = 3', mydict) + self.a.flags.writeable = True + self.a[0] = 5 + self.a[0] = 0 + + def test_writeable_any_base(self): + # Ensure that any base being writeable is sufficient to change flag; + # this is especially interesting for arrays from an array interface. + arr = np.arange(10) + + class subclass(np.ndarray): + pass + + # Create subclass so base will not be collapsed, this is OK to change + view1 = arr.view(subclass) + view2 = view1[...] + arr.flags.writeable = False + view2.flags.writeable = False + view2.flags.writeable = True # Can be set to True again. + + arr = np.arange(10) + + class frominterface: + def __init__(self, arr): + self.arr = arr + self.__array_interface__ = arr.__array_interface__ + + view1 = np.asarray(frominterface) + view2 = view1[...] + view2.flags.writeable = False + view2.flags.writeable = True + + view1.flags.writeable = False + view2.flags.writeable = False + with assert_raises(ValueError): + # Must assume not writeable, since only base is not: + view2.flags.writeable = True + + def test_writeable_from_readonly(self): + # gh-9440 - make sure fromstring, from buffer on readonly buffers + # set writeable False + data = b'\x00' * 100 + vals = np.frombuffer(data, 'B') + assert_raises(ValueError, vals.setflags, write=True) + types = np.dtype( [('vals', 'u1'), ('res3', 'S4')] ) + values = np._core.records.fromstring(data, types) + vals = values['vals'] + assert_raises(ValueError, vals.setflags, write=True) + + def test_writeable_from_buffer(self): + data = bytearray(b'\x00' * 100) + vals = np.frombuffer(data, 'B') + assert_(vals.flags.writeable) + vals.setflags(write=False) + assert_(vals.flags.writeable is False) + vals.setflags(write=True) + assert_(vals.flags.writeable) + types = np.dtype( [('vals', 'u1'), ('res3', 'S4')] ) + values = np._core.records.fromstring(data, types) + vals = values['vals'] + assert_(vals.flags.writeable) + vals.setflags(write=False) + assert_(vals.flags.writeable is False) + vals.setflags(write=True) + assert_(vals.flags.writeable) + + @pytest.mark.skipif(IS_PYPY, reason="PyPy always copies") + def test_writeable_pickle(self): + import pickle + # Small arrays will be copied without setting base. + # See condition for using PyArray_SetBaseObject in + # array_setstate. + a = np.arange(1000) + for v in range(pickle.HIGHEST_PROTOCOL): + vals = pickle.loads(pickle.dumps(a, v)) + assert_(vals.flags.writeable) + assert_(isinstance(vals.base, bytes)) + + def test_writeable_from_c_data(self): + # Test that the writeable flag can be changed for an array wrapping + # low level C-data, but not owning its data. + # Also see that this is deprecated to change from python. + from numpy._core._multiarray_tests import get_c_wrapping_array + + arr_writeable = get_c_wrapping_array(True) + assert not arr_writeable.flags.owndata + assert arr_writeable.flags.writeable + view = arr_writeable[...] + + # Toggling the writeable flag works on the view: + view.flags.writeable = False + assert not view.flags.writeable + view.flags.writeable = True + assert view.flags.writeable + # Flag can be unset on the arr_writeable: + arr_writeable.flags.writeable = False + + arr_readonly = get_c_wrapping_array(False) + assert not arr_readonly.flags.owndata + assert not arr_readonly.flags.writeable + + for arr in [arr_writeable, arr_readonly]: + view = arr[...] + view.flags.writeable = False # make sure it is readonly + arr.flags.writeable = False + assert not arr.flags.writeable + + with assert_raises(ValueError): + view.flags.writeable = True + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + with assert_raises(DeprecationWarning): + arr.flags.writeable = True + + with assert_warns(DeprecationWarning): + arr.flags.writeable = True + + def test_warnonwrite(self): + a = np.arange(10) + a.flags._warn_on_write = True + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always') + a[1] = 10 + a[2] = 10 + # only warn once + assert_(len(w) == 1) + + @pytest.mark.parametrize(["flag", "flag_value", "writeable"], + [("writeable", True, True), + # Delete _warn_on_write after deprecation and simplify + # the parameterization: + ("_warn_on_write", True, False), + ("writeable", False, False)]) + def test_readonly_flag_protocols(self, flag, flag_value, writeable): + a = np.arange(10) + setattr(a.flags, flag, flag_value) + + class MyArr: + __array_struct__ = a.__array_struct__ + + assert memoryview(a).readonly is not writeable + assert a.__array_interface__['data'][1] is not writeable + assert np.asarray(MyArr()).flags.writeable is writeable + + def test_otherflags(self): + assert_equal(self.a.flags.carray, True) + assert_equal(self.a.flags['C'], True) + assert_equal(self.a.flags.farray, False) + assert_equal(self.a.flags.behaved, True) + assert_equal(self.a.flags.fnc, False) + assert_equal(self.a.flags.forc, True) + assert_equal(self.a.flags.owndata, True) + assert_equal(self.a.flags.writeable, True) + assert_equal(self.a.flags.aligned, True) + assert_equal(self.a.flags.writebackifcopy, False) + assert_equal(self.a.flags['X'], False) + assert_equal(self.a.flags['WRITEBACKIFCOPY'], False) + + def test_string_align(self): + a = np.zeros(4, dtype=np.dtype('|S4')) + assert_(a.flags.aligned) + # not power of two are accessed byte-wise and thus considered aligned + a = np.zeros(5, dtype=np.dtype('|S4')) + assert_(a.flags.aligned) + + def test_void_align(self): + a = np.zeros(4, dtype=np.dtype([("a", "i4"), ("b", "i4")])) + assert_(a.flags.aligned) + + @pytest.mark.parametrize("row_size", [5, 1 << 16]) + @pytest.mark.parametrize("row_count", [1, 5]) + @pytest.mark.parametrize("ndmin", [0, 1, 2]) + def test_xcontiguous_load_txt(self, row_size, row_count, ndmin): + s = io.StringIO('\n'.join(['1.0 ' * row_size] * row_count)) + a = np.loadtxt(s, ndmin=ndmin) + + assert a.flags.c_contiguous + x = [i for i in a.shape if i != 1] + assert a.flags.f_contiguous == (len(x) <= 1) + + +class TestHash: + # see #3793 + def test_int(self): + for st, ut, s in [(np.int8, np.uint8, 8), + (np.int16, np.uint16, 16), + (np.int32, np.uint32, 32), + (np.int64, np.uint64, 64)]: + for i in range(1, s): + assert_equal(hash(st(-2**i)), hash(-2**i), + err_msg="%r: -2**%d" % (st, i)) + assert_equal(hash(st(2**(i - 1))), hash(2**(i - 1)), + err_msg="%r: 2**%d" % (st, i - 1)) + assert_equal(hash(st(2**i - 1)), hash(2**i - 1), + err_msg="%r: 2**%d - 1" % (st, i)) + + i = max(i - 1, 1) + assert_equal(hash(ut(2**(i - 1))), hash(2**(i - 1)), + err_msg="%r: 2**%d" % (ut, i - 1)) + assert_equal(hash(ut(2**i - 1)), hash(2**i - 1), + err_msg="%r: 2**%d - 1" % (ut, i)) + + +class TestAttributes: + def setup_method(self): + self.one = np.arange(10) + self.two = np.arange(20).reshape(4, 5) + self.three = np.arange(60, dtype=np.float64).reshape(2, 5, 6) + + def test_attributes(self): + assert_equal(self.one.shape, (10,)) + assert_equal(self.two.shape, (4, 5)) + assert_equal(self.three.shape, (2, 5, 6)) + self.three.shape = (10, 3, 2) + assert_equal(self.three.shape, (10, 3, 2)) + self.three.shape = (2, 5, 6) + assert_equal(self.one.strides, (self.one.itemsize,)) + num = self.two.itemsize + assert_equal(self.two.strides, (5*num, num)) + num = self.three.itemsize + assert_equal(self.three.strides, (30*num, 6*num, num)) + assert_equal(self.one.ndim, 1) + assert_equal(self.two.ndim, 2) + assert_equal(self.three.ndim, 3) + num = self.two.itemsize + assert_equal(self.two.size, 20) + assert_equal(self.two.nbytes, 20*num) + assert_equal(self.two.itemsize, self.two.dtype.itemsize) + assert_equal(self.two.base, np.arange(20)) + + def test_dtypeattr(self): + assert_equal(self.one.dtype, np.dtype(np.int_)) + assert_equal(self.three.dtype, np.dtype(np.float64)) + assert_equal(self.one.dtype.char, np.dtype(int).char) + assert self.one.dtype.char in "lq" + assert_equal(self.three.dtype.char, 'd') + assert_(self.three.dtype.str[0] in '<>') + assert_equal(self.one.dtype.str[1], 'i') + assert_equal(self.three.dtype.str[1], 'f') + + def test_int_subclassing(self): + # Regression test for https://github.com/numpy/numpy/pull/3526 + + numpy_int = np.int_(0) + + # int_ doesn't inherit from Python int, because it's not fixed-width + assert_(not isinstance(numpy_int, int)) + + def test_stridesattr(self): + x = self.one + + def make_array(size, offset, strides): + return np.ndarray(size, buffer=x, dtype=int, + offset=offset*x.itemsize, + strides=strides*x.itemsize) + + assert_equal(make_array(4, 4, -1), np.array([4, 3, 2, 1])) + assert_raises(ValueError, make_array, 4, 4, -2) + assert_raises(ValueError, make_array, 4, 2, -1) + assert_raises(ValueError, make_array, 8, 3, 1) + assert_equal(make_array(8, 3, 0), np.array([3]*8)) + # Check behavior reported in gh-2503: + assert_raises(ValueError, make_array, (2, 3), 5, np.array([-2, -3])) + make_array(0, 0, 10) + + def test_set_stridesattr(self): + x = self.one + + def make_array(size, offset, strides): + try: + r = np.ndarray([size], dtype=int, buffer=x, + offset=offset*x.itemsize) + except Exception as e: + raise RuntimeError(e) + r.strides = strides = strides*x.itemsize + return r + + assert_equal(make_array(4, 4, -1), np.array([4, 3, 2, 1])) + assert_equal(make_array(7, 3, 1), np.array([3, 4, 5, 6, 7, 8, 9])) + assert_raises(ValueError, make_array, 4, 4, -2) + assert_raises(ValueError, make_array, 4, 2, -1) + assert_raises(RuntimeError, make_array, 8, 3, 1) + # Check that the true extent of the array is used. + # Test relies on as_strided base not exposing a buffer. + x = np.lib.stride_tricks.as_strided(np.arange(1), (10, 10), (0, 0)) + + def set_strides(arr, strides): + arr.strides = strides + + assert_raises(ValueError, set_strides, x, (10*x.itemsize, x.itemsize)) + + # Test for offset calculations: + x = np.lib.stride_tricks.as_strided(np.arange(10, dtype=np.int8)[-1], + shape=(10,), strides=(-1,)) + assert_raises(ValueError, set_strides, x[::-1], -1) + a = x[::-1] + a.strides = 1 + a[::2].strides = 2 + + # test 0d + arr_0d = np.array(0) + arr_0d.strides = () + assert_raises(TypeError, set_strides, arr_0d, None) + + def test_fill(self): + for t in "?bhilqpBHILQPfdgFDGO": + x = np.empty((3, 2, 1), t) + y = np.empty((3, 2, 1), t) + x.fill(1) + y[...] = 1 + assert_equal(x, y) + + def test_fill_max_uint64(self): + x = np.empty((3, 2, 1), dtype=np.uint64) + y = np.empty((3, 2, 1), dtype=np.uint64) + value = 2**64 - 1 + y[...] = value + x.fill(value) + assert_array_equal(x, y) + + def test_fill_struct_array(self): + # Filling from a scalar + x = np.array([(0, 0.0), (1, 1.0)], dtype='i4,f8') + x.fill(x[0]) + assert_equal(x['f1'][1], x['f1'][0]) + # Filling from a tuple that can be converted + # to a scalar + x = np.zeros(2, dtype=[('a', 'f8'), ('b', 'i4')]) + x.fill((3.5, -2)) + assert_array_equal(x['a'], [3.5, 3.5]) + assert_array_equal(x['b'], [-2, -2]) + + def test_fill_readonly(self): + # gh-22922 + a = np.zeros(11) + a.setflags(write=False) + with pytest.raises(ValueError, match=".*read-only"): + a.fill(0) + + def test_fill_subarrays(self): + # NOTE: + # This is also a regression test for a crash with PYTHONMALLOC=debug + + dtype = np.dtype("2i4')) + assert_(np.dtype([('a', 'i4')])) + + def test_structured_non_void(self): + fields = [('a', 'i8'), ('b', 'f8')]) + assert_equal(a == b, [False, True]) + assert_equal(a != b, [True, False]) + + a = np.array([(5, 42), (10, 1)], dtype=[('a', '>f8'), ('b', 'i8')]) + assert_equal(a == b, [False, True]) + assert_equal(a != b, [True, False]) + + # Including with embedded subarray dtype (although subarray comparison + # itself may still be a bit weird and compare the raw data) + a = np.array([(5, 42), (10, 1)], dtype=[('a', '10>f8'), ('b', '5i8')]) + assert_equal(a == b, [False, True]) + assert_equal(a != b, [True, False]) + + @pytest.mark.parametrize("op", [ + operator.eq, lambda x, y: operator.eq(y, x), + operator.ne, lambda x, y: operator.ne(y, x)]) + def test_void_comparison_failures(self, op): + # In principle, one could decide to return an array of False for some + # if comparisons are impossible. But right now we return TypeError + # when "void" dtype are involved. + x = np.zeros(3, dtype=[('a', 'i1')]) + y = np.zeros(3) + # Cannot compare non-structured to structured: + with pytest.raises(TypeError): + op(x, y) + + # Added title prevents promotion, but casts are OK: + y = np.zeros(3, dtype=[(('title', 'a'), 'i1')]) + assert np.can_cast(y.dtype, x.dtype) + with pytest.raises(TypeError): + op(x, y) + + x = np.zeros(3, dtype="V7") + y = np.zeros(3, dtype="V8") + with pytest.raises(TypeError): + op(x, y) + + def test_casting(self): + # Check that casting a structured array to change its byte order + # works + a = np.array([(1,)], dtype=[('a', 'i4')], casting='unsafe')) + b = a.astype([('a', '>i4')]) + a_tmp = a.byteswap() + a_tmp = a_tmp.view(a_tmp.dtype.newbyteorder()) + assert_equal(b, a_tmp) + assert_equal(a['a'][0], b['a'][0]) + + # Check that equality comparison works on structured arrays if + # they are 'equiv'-castable + a = np.array([(5, 42), (10, 1)], dtype=[('a', '>i4'), ('b', 'f8')]) + assert_(np.can_cast(a.dtype, b.dtype, casting='equiv')) + assert_equal(a == b, [True, True]) + + # Check that 'equiv' casting can change byte order + assert_(np.can_cast(a.dtype, b.dtype, casting='equiv')) + c = a.astype(b.dtype, casting='equiv') + assert_equal(a == c, [True, True]) + + # Check that 'safe' casting can change byte order and up-cast + # fields + t = [('a', 'f8')] + assert_(np.can_cast(a.dtype, t, casting='safe')) + c = a.astype(t, casting='safe') + assert_equal((c == np.array([(5, 42), (10, 1)], dtype=t)), + [True, True]) + + # Check that 'same_kind' casting can change byte order and + # change field widths within a "kind" + t = [('a', 'f4')] + assert_(np.can_cast(a.dtype, t, casting='same_kind')) + c = a.astype(t, casting='same_kind') + assert_equal((c == np.array([(5, 42), (10, 1)], dtype=t)), + [True, True]) + + # Check that casting fails if the casting rule should fail on + # any of the fields + t = [('a', '>i8'), ('b', 'i2'), ('b', 'i8'), ('b', 'i4')] + assert_(not np.can_cast(a.dtype, t, casting=casting)) + t = [('a', '>i4'), ('b', 'i8") + ab = np.array([(1, 2)], dtype=[A, B]) + ba = np.array([(1, 2)], dtype=[B, A]) + assert_raises(TypeError, np.concatenate, ab, ba) + assert_raises(TypeError, np.result_type, ab.dtype, ba.dtype) + assert_raises(TypeError, np.promote_types, ab.dtype, ba.dtype) + + # dtypes with same field names/order but different memory offsets + # and byte-order are promotable to packed nbo. + assert_equal(np.promote_types(ab.dtype, ba[['a', 'b']].dtype), + repack_fields(ab.dtype.newbyteorder('N'))) + + # gh-13667 + # dtypes with different fieldnames but castable field types are castable + assert_equal(np.can_cast(ab.dtype, ba.dtype), True) + assert_equal(ab.astype(ba.dtype).dtype, ba.dtype) + assert_equal(np.can_cast('f8,i8', [('f0', 'f8'), ('f1', 'i8')]), True) + assert_equal(np.can_cast('f8,i8', [('f1', 'f8'), ('f0', 'i8')]), True) + assert_equal(np.can_cast('f8,i8', [('f1', 'i8'), ('f0', 'f8')]), False) + assert_equal(np.can_cast('f8,i8', [('f1', 'i8'), ('f0', 'f8')], + casting='unsafe'), True) + + ab[:] = ba # make sure assignment still works + + # tests of type-promotion of corresponding fields + dt1 = np.dtype([("", "i4")]) + dt2 = np.dtype([("", "i8")]) + assert_equal(np.promote_types(dt1, dt2), np.dtype([('f0', 'i8')])) + assert_equal(np.promote_types(dt2, dt1), np.dtype([('f0', 'i8')])) + assert_raises(TypeError, np.promote_types, dt1, np.dtype([("", "V3")])) + assert_equal(np.promote_types('i4,f8', 'i8,f4'), + np.dtype([('f0', 'i8'), ('f1', 'f8')])) + # test nested case + dt1nest = np.dtype([("", dt1)]) + dt2nest = np.dtype([("", dt2)]) + assert_equal(np.promote_types(dt1nest, dt2nest), + np.dtype([('f0', np.dtype([('f0', 'i8')]))])) + + # note that offsets are lost when promoting: + dt = np.dtype({'names': ['x'], 'formats': ['i4'], 'offsets': [8]}) + a = np.ones(3, dtype=dt) + assert_equal(np.concatenate([a, a]).dtype, np.dtype([('x', 'i4')])) + + @pytest.mark.parametrize("dtype_dict", [ + dict(names=["a", "b"], formats=["i4", "f"], itemsize=100), + dict(names=["a", "b"], formats=["i4", "f"], + offsets=[0, 12])]) + @pytest.mark.parametrize("align", [True, False]) + def test_structured_promotion_packs(self, dtype_dict, align): + # Structured dtypes are packed when promoted (we consider the packed + # form to be "canonical"), so tere is no extra padding. + dtype = np.dtype(dtype_dict, align=align) + # Remove non "canonical" dtype options: + dtype_dict.pop("itemsize", None) + dtype_dict.pop("offsets", None) + expected = np.dtype(dtype_dict, align=align) + + res = np.promote_types(dtype, dtype) + assert res.itemsize == expected.itemsize + assert res.fields == expected.fields + + # But the "expected" one, should just be returned unchanged: + res = np.promote_types(expected, expected) + assert res is expected + + def test_structured_asarray_is_view(self): + # A scalar viewing an array preserves its view even when creating a + # new array. This test documents behaviour, it may not be the best + # desired behaviour. + arr = np.array([1], dtype="i,i") + scalar = arr[0] + assert not scalar.flags.owndata # view into the array + assert np.asarray(scalar).base is scalar + # But never when a dtype is passed in: + assert np.asarray(scalar, dtype=scalar.dtype).base is None + # A scalar which owns its data does not have this property. + # It is not easy to create one, one method is to use pickle: + scalar = pickle.loads(pickle.dumps(scalar)) + assert scalar.flags.owndata + assert np.asarray(scalar).base is None + +class TestBool: + def test_test_interning(self): + a0 = np.bool(0) + b0 = np.bool(False) + assert_(a0 is b0) + a1 = np.bool(1) + b1 = np.bool(True) + assert_(a1 is b1) + assert_(np.array([True])[0] is a1) + assert_(np.array(True)[()] is a1) + + def test_sum(self): + d = np.ones(101, dtype=bool) + assert_equal(d.sum(), d.size) + assert_equal(d[::2].sum(), d[::2].size) + assert_equal(d[::-2].sum(), d[::-2].size) + + d = np.frombuffer(b'\xff\xff' * 100, dtype=bool) + assert_equal(d.sum(), d.size) + assert_equal(d[::2].sum(), d[::2].size) + assert_equal(d[::-2].sum(), d[::-2].size) + + def check_count_nonzero(self, power, length): + powers = [2 ** i for i in range(length)] + for i in range(2**power): + l = [(i & x) != 0 for x in powers] + a = np.array(l, dtype=bool) + c = builtins.sum(l) + assert_equal(np.count_nonzero(a), c) + av = a.view(np.uint8) + av *= 3 + assert_equal(np.count_nonzero(a), c) + av *= 4 + assert_equal(np.count_nonzero(a), c) + av[av != 0] = 0xFF + assert_equal(np.count_nonzero(a), c) + + def test_count_nonzero(self): + # check all 12 bit combinations in a length 17 array + # covers most cases of the 16 byte unrolled code + self.check_count_nonzero(12, 17) + + @pytest.mark.slow + def test_count_nonzero_all(self): + # check all combinations in a length 17 array + # covers all cases of the 16 byte unrolled code + self.check_count_nonzero(17, 17) + + def test_count_nonzero_unaligned(self): + # prevent mistakes as e.g. gh-4060 + for o in range(7): + a = np.zeros((18,), dtype=bool)[o+1:] + a[:o] = True + assert_equal(np.count_nonzero(a), builtins.sum(a.tolist())) + a = np.ones((18,), dtype=bool)[o+1:] + a[:o] = False + assert_equal(np.count_nonzero(a), builtins.sum(a.tolist())) + + def _test_cast_from_flexible(self, dtype): + # empty string -> false + for n in range(3): + v = np.array(b'', (dtype, n)) + assert_equal(bool(v), False) + assert_equal(bool(v[()]), False) + assert_equal(v.astype(bool), False) + assert_(isinstance(v.astype(bool), np.ndarray)) + assert_(v[()].astype(bool) is np.False_) + + # anything else -> true + for n in range(1, 4): + for val in [b'a', b'0', b' ']: + v = np.array(val, (dtype, n)) + assert_equal(bool(v), True) + assert_equal(bool(v[()]), True) + assert_equal(v.astype(bool), True) + assert_(isinstance(v.astype(bool), np.ndarray)) + assert_(v[()].astype(bool) is np.True_) + + def test_cast_from_void(self): + self._test_cast_from_flexible(np.void) + + @pytest.mark.xfail(reason="See gh-9847") + def test_cast_from_unicode(self): + self._test_cast_from_flexible(np.str_) + + @pytest.mark.xfail(reason="See gh-9847") + def test_cast_from_bytes(self): + self._test_cast_from_flexible(np.bytes_) + + +class TestZeroSizeFlexible: + @staticmethod + def _zeros(shape, dtype=str): + dtype = np.dtype(dtype) + if dtype == np.void: + return np.zeros(shape, dtype=(dtype, 0)) + + # not constructable directly + dtype = np.dtype([('x', dtype, 0)]) + return np.zeros(shape, dtype=dtype)['x'] + + def test_create(self): + zs = self._zeros(10, bytes) + assert_equal(zs.itemsize, 0) + zs = self._zeros(10, np.void) + assert_equal(zs.itemsize, 0) + zs = self._zeros(10, str) + assert_equal(zs.itemsize, 0) + + def _test_sort_partition(self, name, kinds, **kwargs): + # Previously, these would all hang + for dt in [bytes, np.void, str]: + zs = self._zeros(10, dt) + sort_method = getattr(zs, name) + sort_func = getattr(np, name) + for kind in kinds: + sort_method(kind=kind, **kwargs) + sort_func(zs, kind=kind, **kwargs) + + def test_sort(self): + self._test_sort_partition('sort', kinds='qhs') + + def test_argsort(self): + self._test_sort_partition('argsort', kinds='qhs') + + def test_partition(self): + self._test_sort_partition('partition', kinds=['introselect'], kth=2) + + def test_argpartition(self): + self._test_sort_partition('argpartition', kinds=['introselect'], kth=2) + + def test_resize(self): + # previously an error + for dt in [bytes, np.void, str]: + zs = self._zeros(10, dt) + zs.resize(25) + zs.resize((10, 10)) + + def test_view(self): + for dt in [bytes, np.void, str]: + zs = self._zeros(10, dt) + + # viewing as itself should be allowed + assert_equal(zs.view(dt).dtype, np.dtype(dt)) + + # viewing as any non-empty type gives an empty result + assert_equal(zs.view((dt, 1)).shape, (0,)) + + def test_dumps(self): + zs = self._zeros(10, int) + assert_equal(zs, pickle.loads(zs.dumps())) + + def test_pickle(self): + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + for dt in [bytes, np.void, str]: + zs = self._zeros(10, dt) + p = pickle.dumps(zs, protocol=proto) + zs2 = pickle.loads(p) + + assert_equal(zs.dtype, zs2.dtype) + + def test_pickle_empty(self): + """Checking if an empty array pickled and un-pickled will not cause a + segmentation fault""" + arr = np.array([]).reshape(999999, 0) + pk_dmp = pickle.dumps(arr) + pk_load = pickle.loads(pk_dmp) + + assert pk_load.size == 0 + + @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5, + reason="requires pickle protocol 5") + def test_pickle_with_buffercallback(self): + array = np.arange(10) + buffers = [] + bytes_string = pickle.dumps(array, buffer_callback=buffers.append, + protocol=5) + array_from_buffer = pickle.loads(bytes_string, buffers=buffers) + # when using pickle protocol 5 with buffer callbacks, + # array_from_buffer is reconstructed from a buffer holding a view + # to the initial array's data, so modifying an element in array + # should modify it in array_from_buffer too. + array[0] = -1 + assert array_from_buffer[0] == -1, array_from_buffer[0] + + +class TestMethods: + + sort_kinds = ['quicksort', 'heapsort', 'stable'] + + def test_all_where(self): + a = np.array([[True, False, True], + [False, False, False], + [True, True, True]]) + wh_full = np.array([[True, False, True], + [False, False, False], + [True, False, True]]) + wh_lower = np.array([[False], + [False], + [True]]) + for _ax in [0, None]: + assert_equal(a.all(axis=_ax, where=wh_lower), + np.all(a[wh_lower[:,0],:], axis=_ax)) + assert_equal(np.all(a, axis=_ax, where=wh_lower), + a[wh_lower[:,0],:].all(axis=_ax)) + + assert_equal(a.all(where=wh_full), True) + assert_equal(np.all(a, where=wh_full), True) + assert_equal(a.all(where=False), True) + assert_equal(np.all(a, where=False), True) + + def test_any_where(self): + a = np.array([[True, False, True], + [False, False, False], + [True, True, True]]) + wh_full = np.array([[False, True, False], + [True, True, True], + [False, False, False]]) + wh_middle = np.array([[False], + [True], + [False]]) + for _ax in [0, None]: + assert_equal(a.any(axis=_ax, where=wh_middle), + np.any(a[wh_middle[:,0],:], axis=_ax)) + assert_equal(np.any(a, axis=_ax, where=wh_middle), + a[wh_middle[:,0],:].any(axis=_ax)) + assert_equal(a.any(where=wh_full), False) + assert_equal(np.any(a, where=wh_full), False) + assert_equal(a.any(where=False), False) + assert_equal(np.any(a, where=False), False) + + @pytest.mark.parametrize("dtype", + ["i8", "U10", "object", "datetime64[ms]"]) + def test_any_and_all_result_dtype(self, dtype): + arr = np.ones(3, dtype=dtype) + assert arr.any().dtype == np.bool + assert arr.all().dtype == np.bool + + def test_any_and_all_object_dtype(self): + # (seberg) Not sure we should even allow dtype here, but it is. + arr = np.ones(3, dtype=object) + # keepdims to prevent getting a scalar. + assert arr.any(dtype=object, keepdims=True).dtype == object + assert arr.all(dtype=object, keepdims=True).dtype == object + + def test_compress(self): + tgt = [[5, 6, 7, 8, 9]] + arr = np.arange(10).reshape(2, 5) + out = arr.compress([0, 1], axis=0) + assert_equal(out, tgt) + + tgt = [[1, 3], [6, 8]] + out = arr.compress([0, 1, 0, 1, 0], axis=1) + assert_equal(out, tgt) + + tgt = [[1], [6]] + arr = np.arange(10).reshape(2, 5) + out = arr.compress([0, 1], axis=1) + assert_equal(out, tgt) + + arr = np.arange(10).reshape(2, 5) + out = arr.compress([0, 1]) + assert_equal(out, 1) + + def test_choose(self): + x = 2*np.ones((3,), dtype=int) + y = 3*np.ones((3,), dtype=int) + x2 = 2*np.ones((2, 3), dtype=int) + y2 = 3*np.ones((2, 3), dtype=int) + ind = np.array([0, 0, 1]) + + A = ind.choose((x, y)) + assert_equal(A, [2, 2, 3]) + + A = ind.choose((x2, y2)) + assert_equal(A, [[2, 2, 3], [2, 2, 3]]) + + A = ind.choose((x, y2)) + assert_equal(A, [[2, 2, 3], [2, 2, 3]]) + + oned = np.ones(1) + # gh-12031, caused SEGFAULT + assert_raises(TypeError, oned.choose,np.void(0), [oned]) + + out = np.array(0) + ret = np.choose(np.array(1), [10, 20, 30], out=out) + assert out is ret + assert_equal(out[()], 20) + + # gh-6272 check overlap on out + x = np.arange(5) + y = np.choose([0,0,0], [x[:3], x[:3], x[:3]], out=x[1:4], mode='wrap') + assert_equal(y, np.array([0, 1, 2])) + + def test_prod(self): + ba = [1, 2, 10, 11, 6, 5, 4] + ba2 = [[1, 2, 3, 4], [5, 6, 7, 9], [10, 3, 4, 5]] + + for ctype in [np.int16, np.uint16, np.int32, np.uint32, + np.float32, np.float64, np.complex64, np.complex128]: + a = np.array(ba, ctype) + a2 = np.array(ba2, ctype) + if ctype in ['1', 'b']: + assert_raises(ArithmeticError, a.prod) + assert_raises(ArithmeticError, a2.prod, axis=1) + else: + assert_equal(a.prod(axis=0), 26400) + assert_array_equal(a2.prod(axis=0), + np.array([50, 36, 84, 180], ctype)) + assert_array_equal(a2.prod(axis=-1), + np.array([24, 1890, 600], ctype)) + + @pytest.mark.parametrize('dtype', [None, object]) + def test_repeat(self, dtype): + m = np.array([1, 2, 3, 4, 5, 6], dtype=dtype) + m_rect = m.reshape((2, 3)) + + A = m.repeat([1, 3, 2, 1, 1, 2]) + assert_equal(A, [1, 2, 2, 2, 3, + 3, 4, 5, 6, 6]) + + A = m.repeat(2) + assert_equal(A, [1, 1, 2, 2, 3, 3, + 4, 4, 5, 5, 6, 6]) + + A = m_rect.repeat([2, 1], axis=0) + assert_equal(A, [[1, 2, 3], + [1, 2, 3], + [4, 5, 6]]) + + A = m_rect.repeat([1, 3, 2], axis=1) + assert_equal(A, [[1, 2, 2, 2, 3, 3], + [4, 5, 5, 5, 6, 6]]) + + A = m_rect.repeat(2, axis=0) + assert_equal(A, [[1, 2, 3], + [1, 2, 3], + [4, 5, 6], + [4, 5, 6]]) + + A = m_rect.repeat(2, axis=1) + assert_equal(A, [[1, 1, 2, 2, 3, 3], + [4, 4, 5, 5, 6, 6]]) + + def test_reshape(self): + arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]) + + tgt = [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]] + assert_equal(arr.reshape(2, 6), tgt) + + tgt = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]] + assert_equal(arr.reshape(3, 4), tgt) + + tgt = [[1, 10, 8, 6], [4, 2, 11, 9], [7, 5, 3, 12]] + assert_equal(arr.reshape((3, 4), order='F'), tgt) + + tgt = [[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]] + assert_equal(arr.T.reshape((3, 4), order='C'), tgt) + + def test_round(self): + def check_round(arr, expected, *round_args): + assert_equal(arr.round(*round_args), expected) + # With output array + out = np.zeros_like(arr) + res = arr.round(*round_args, out=out) + assert_equal(out, expected) + assert out is res + + check_round(np.array([1.2, 1.5]), [1, 2]) + check_round(np.array(1.5), 2) + check_round(np.array([12.2, 15.5]), [10, 20], -1) + check_round(np.array([12.15, 15.51]), [12.2, 15.5], 1) + # Complex rounding + check_round(np.array([4.5 + 1.5j]), [4 + 2j]) + check_round(np.array([12.5 + 15.5j]), [10 + 20j], -1) + + def test_squeeze(self): + a = np.array([[[1], [2], [3]]]) + assert_equal(a.squeeze(), [1, 2, 3]) + assert_equal(a.squeeze(axis=(0,)), [[1], [2], [3]]) + assert_raises(ValueError, a.squeeze, axis=(1,)) + assert_equal(a.squeeze(axis=(2,)), [[1, 2, 3]]) + + def test_transpose(self): + a = np.array([[1, 2], [3, 4]]) + assert_equal(a.transpose(), [[1, 3], [2, 4]]) + assert_raises(ValueError, lambda: a.transpose(0)) + assert_raises(ValueError, lambda: a.transpose(0, 0)) + assert_raises(ValueError, lambda: a.transpose(0, 1, 2)) + + def test_sort(self): + # test ordering for floats and complex containing nans. It is only + # necessary to check the less-than comparison, so sorts that + # only follow the insertion sort path are sufficient. We only + # test doubles and complex doubles as the logic is the same. + + # check doubles + msg = "Test real sort order with nans" + a = np.array([np.nan, 1, 0]) + b = np.sort(a) + assert_equal(b, a[::-1], msg) + # check complex + msg = "Test complex sort order with nans" + a = np.zeros(9, dtype=np.complex128) + a.real += [np.nan, np.nan, np.nan, 1, 0, 1, 1, 0, 0] + a.imag += [np.nan, 1, 0, np.nan, np.nan, 1, 0, 1, 0] + b = np.sort(a) + assert_equal(b, a[::-1], msg) + + with assert_raises_regex( + ValueError, + "kind` and `stable` parameters can't be provided at the same time" + ): + np.sort(a, kind="stable", stable=True) + + # all c scalar sorts use the same code with different types + # so it suffices to run a quick check with one type. The number + # of sorted items must be greater than ~50 to check the actual + # algorithm because quick and merge sort fall over to insertion + # sort for small arrays. + + @pytest.mark.parametrize('dtype', [np.uint8, np.uint16, np.uint32, np.uint64, + np.float16, np.float32, np.float64, + np.longdouble]) + def test_sort_unsigned(self, dtype): + a = np.arange(101, dtype=dtype) + b = a[::-1].copy() + for kind in self.sort_kinds: + msg = "scalar sort, kind=%s" % kind + c = a.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + c = b.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + + @pytest.mark.parametrize('dtype', + [np.int8, np.int16, np.int32, np.int64, np.float16, + np.float32, np.float64, np.longdouble]) + def test_sort_signed(self, dtype): + a = np.arange(-50, 51, dtype=dtype) + b = a[::-1].copy() + for kind in self.sort_kinds: + msg = "scalar sort, kind=%s" % (kind) + c = a.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + c = b.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + + @pytest.mark.parametrize('dtype', [np.float32, np.float64, np.longdouble]) + @pytest.mark.parametrize('part', ['real', 'imag']) + def test_sort_complex(self, part, dtype): + # test complex sorts. These use the same code as the scalars + # but the compare function differs. + cdtype = { + np.single: np.csingle, + np.double: np.cdouble, + np.longdouble: np.clongdouble, + }[dtype] + a = np.arange(-50, 51, dtype=dtype) + b = a[::-1].copy() + ai = (a * (1+1j)).astype(cdtype) + bi = (b * (1+1j)).astype(cdtype) + setattr(ai, part, 1) + setattr(bi, part, 1) + for kind in self.sort_kinds: + msg = "complex sort, %s part == 1, kind=%s" % (part, kind) + c = ai.copy() + c.sort(kind=kind) + assert_equal(c, ai, msg) + c = bi.copy() + c.sort(kind=kind) + assert_equal(c, ai, msg) + + def test_sort_complex_byte_swapping(self): + # test sorting of complex arrays requiring byte-swapping, gh-5441 + for endianness in '<>': + for dt in np.typecodes['Complex']: + arr = np.array([1+3.j, 2+2.j, 3+1.j], dtype=endianness + dt) + c = arr.copy() + c.sort() + msg = 'byte-swapped complex sort, dtype={0}'.format(dt) + assert_equal(c, arr, msg) + + @pytest.mark.parametrize('dtype', [np.bytes_, np.str_]) + def test_sort_string(self, dtype): + # np.array will perform the encoding to bytes for us in the bytes test + a = np.array(['aaaaaaaa' + chr(i) for i in range(101)], dtype=dtype) + b = a[::-1].copy() + for kind in self.sort_kinds: + msg = "kind=%s" % kind + c = a.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + c = b.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + + def test_sort_object(self): + # test object array sorts. + a = np.empty((101,), dtype=object) + a[:] = list(range(101)) + b = a[::-1] + for kind in ['q', 'h', 'm']: + msg = "kind=%s" % kind + c = a.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + c = b.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + + @pytest.mark.parametrize("dt", [ + np.dtype([('f', float), ('i', int)]), + np.dtype([('f', float), ('i', object)])]) + @pytest.mark.parametrize("step", [1, 2]) + def test_sort_structured(self, dt, step): + # test record array sorts. + a = np.array([(i, i) for i in range(101*step)], dtype=dt) + b = a[::-1] + for kind in ['q', 'h', 'm']: + msg = "kind=%s" % kind + c = a.copy()[::step] + indx = c.argsort(kind=kind) + c.sort(kind=kind) + assert_equal(c, a[::step], msg) + assert_equal(a[::step][indx], a[::step], msg) + c = b.copy()[::step] + indx = c.argsort(kind=kind) + c.sort(kind=kind) + assert_equal(c, a[step-1::step], msg) + assert_equal(b[::step][indx], a[step-1::step], msg) + + @pytest.mark.parametrize('dtype', ['datetime64[D]', 'timedelta64[D]']) + def test_sort_time(self, dtype): + # test datetime64 and timedelta64 sorts. + a = np.arange(0, 101, dtype=dtype) + b = a[::-1] + for kind in ['q', 'h', 'm']: + msg = "kind=%s" % kind + c = a.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + c = b.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + + def test_sort_axis(self): + # check axis handling. This should be the same for all type + # specific sorts, so we only check it for one type and one kind + a = np.array([[3, 2], [1, 0]]) + b = np.array([[1, 0], [3, 2]]) + c = np.array([[2, 3], [0, 1]]) + d = a.copy() + d.sort(axis=0) + assert_equal(d, b, "test sort with axis=0") + d = a.copy() + d.sort(axis=1) + assert_equal(d, c, "test sort with axis=1") + d = a.copy() + d.sort() + assert_equal(d, c, "test sort with default axis") + + def test_sort_size_0(self): + # check axis handling for multidimensional empty arrays + a = np.array([]) + a.shape = (3, 2, 1, 0) + for axis in range(-a.ndim, a.ndim): + msg = 'test empty array sort with axis={0}'.format(axis) + assert_equal(np.sort(a, axis=axis), a, msg) + msg = 'test empty array sort with axis=None' + assert_equal(np.sort(a, axis=None), a.ravel(), msg) + + def test_sort_bad_ordering(self): + # test generic class with bogus ordering, + # should not segfault. + class Boom: + def __lt__(self, other): + return True + + a = np.array([Boom()] * 100, dtype=object) + for kind in self.sort_kinds: + msg = "kind=%s" % kind + c = a.copy() + c.sort(kind=kind) + assert_equal(c, a, msg) + + def test_void_sort(self): + # gh-8210 - previously segfaulted + for i in range(4): + rand = np.random.randint(256, size=4000, dtype=np.uint8) + arr = rand.view('V4') + arr[::-1].sort() + + dt = np.dtype([('val', 'i4', (1,))]) + for i in range(4): + rand = np.random.randint(256, size=4000, dtype=np.uint8) + arr = rand.view(dt) + arr[::-1].sort() + + def test_sort_raises(self): + #gh-9404 + arr = np.array([0, datetime.now(), 1], dtype=object) + for kind in self.sort_kinds: + assert_raises(TypeError, arr.sort, kind=kind) + #gh-3879 + class Raiser: + def raises_anything(*args, **kwargs): + raise TypeError("SOMETHING ERRORED") + __eq__ = __ne__ = __lt__ = __gt__ = __ge__ = __le__ = raises_anything + arr = np.array([[Raiser(), n] for n in range(10)]).reshape(-1) + np.random.shuffle(arr) + for kind in self.sort_kinds: + assert_raises(TypeError, arr.sort, kind=kind) + + def test_sort_degraded(self): + # test degraded dataset would take minutes to run with normal qsort + d = np.arange(1000000) + do = d.copy() + x = d + # create a median of 3 killer where each median is the sorted second + # last element of the quicksort partition + while x.size > 3: + mid = x.size // 2 + x[mid], x[-2] = x[-2], x[mid] + x = x[:-2] + + assert_equal(np.sort(d), do) + assert_equal(d[np.argsort(d)], do) + + def test_copy(self): + def assert_fortran(arr): + assert_(arr.flags.fortran) + assert_(arr.flags.f_contiguous) + assert_(not arr.flags.c_contiguous) + + def assert_c(arr): + assert_(not arr.flags.fortran) + assert_(not arr.flags.f_contiguous) + assert_(arr.flags.c_contiguous) + + a = np.empty((2, 2), order='F') + # Test copying a Fortran array + assert_c(a.copy()) + assert_c(a.copy('C')) + assert_fortran(a.copy('F')) + assert_fortran(a.copy('A')) + + # Now test starting with a C array. + a = np.empty((2, 2), order='C') + assert_c(a.copy()) + assert_c(a.copy('C')) + assert_fortran(a.copy('F')) + assert_c(a.copy('A')) + + @pytest.mark.parametrize("dtype", ['O', np.int32, 'i,O']) + def test__deepcopy__(self, dtype): + # Force the entry of NULLs into array + a = np.empty(4, dtype=dtype) + ctypes.memset(a.ctypes.data, 0, a.nbytes) + + # Ensure no error is raised, see gh-21833 + b = a.__deepcopy__({}) + + a[0] = 42 + with pytest.raises(AssertionError): + assert_array_equal(a, b) + + def test__deepcopy__catches_failure(self): + class MyObj: + def __deepcopy__(self, *args, **kwargs): + raise RuntimeError + + arr = np.array([1, MyObj(), 3], dtype='O') + with pytest.raises(RuntimeError): + arr.__deepcopy__({}) + + def test_sort_order(self): + # Test sorting an array with fields + x1 = np.array([21, 32, 14]) + x2 = np.array(['my', 'first', 'name']) + x3 = np.array([3.1, 4.5, 6.2]) + r = np.rec.fromarrays([x1, x2, x3], names='id,word,number') + + r.sort(order=['id']) + assert_equal(r.id, np.array([14, 21, 32])) + assert_equal(r.word, np.array(['name', 'my', 'first'])) + assert_equal(r.number, np.array([6.2, 3.1, 4.5])) + + r.sort(order=['word']) + assert_equal(r.id, np.array([32, 21, 14])) + assert_equal(r.word, np.array(['first', 'my', 'name'])) + assert_equal(r.number, np.array([4.5, 3.1, 6.2])) + + r.sort(order=['number']) + assert_equal(r.id, np.array([21, 32, 14])) + assert_equal(r.word, np.array(['my', 'first', 'name'])) + assert_equal(r.number, np.array([3.1, 4.5, 6.2])) + + assert_raises_regex(ValueError, 'duplicate', + lambda: r.sort(order=['id', 'id'])) + + if sys.byteorder == 'little': + strtype = '>i2' + else: + strtype = '': + for dt in np.typecodes['Complex']: + arr = np.array([1+3.j, 2+2.j, 3+1.j], dtype=endianness + dt) + msg = 'byte-swapped complex argsort, dtype={0}'.format(dt) + assert_equal(arr.argsort(), + np.arange(len(arr), dtype=np.intp), msg) + + # test string argsorts. + s = 'aaaaaaaa' + a = np.array([s + chr(i) for i in range(101)]) + b = a[::-1].copy() + r = np.arange(101) + rr = r[::-1] + for kind in self.sort_kinds: + msg = "string argsort, kind=%s" % kind + assert_equal(a.copy().argsort(kind=kind), r, msg) + assert_equal(b.copy().argsort(kind=kind), rr, msg) + + # test unicode argsorts. + s = 'aaaaaaaa' + a = np.array([s + chr(i) for i in range(101)], dtype=np.str_) + b = a[::-1] + r = np.arange(101) + rr = r[::-1] + for kind in self.sort_kinds: + msg = "unicode argsort, kind=%s" % kind + assert_equal(a.copy().argsort(kind=kind), r, msg) + assert_equal(b.copy().argsort(kind=kind), rr, msg) + + # test object array argsorts. + a = np.empty((101,), dtype=object) + a[:] = list(range(101)) + b = a[::-1] + r = np.arange(101) + rr = r[::-1] + for kind in self.sort_kinds: + msg = "object argsort, kind=%s" % kind + assert_equal(a.copy().argsort(kind=kind), r, msg) + assert_equal(b.copy().argsort(kind=kind), rr, msg) + + # test structured array argsorts. + dt = np.dtype([('f', float), ('i', int)]) + a = np.array([(i, i) for i in range(101)], dtype=dt) + b = a[::-1] + r = np.arange(101) + rr = r[::-1] + for kind in self.sort_kinds: + msg = "structured array argsort, kind=%s" % kind + assert_equal(a.copy().argsort(kind=kind), r, msg) + assert_equal(b.copy().argsort(kind=kind), rr, msg) + + # test datetime64 argsorts. + a = np.arange(0, 101, dtype='datetime64[D]') + b = a[::-1] + r = np.arange(101) + rr = r[::-1] + for kind in ['q', 'h', 'm']: + msg = "datetime64 argsort, kind=%s" % kind + assert_equal(a.copy().argsort(kind=kind), r, msg) + assert_equal(b.copy().argsort(kind=kind), rr, msg) + + # test timedelta64 argsorts. + a = np.arange(0, 101, dtype='timedelta64[D]') + b = a[::-1] + r = np.arange(101) + rr = r[::-1] + for kind in ['q', 'h', 'm']: + msg = "timedelta64 argsort, kind=%s" % kind + assert_equal(a.copy().argsort(kind=kind), r, msg) + assert_equal(b.copy().argsort(kind=kind), rr, msg) + + # check axis handling. This should be the same for all type + # specific argsorts, so we only check it for one type and one kind + a = np.array([[3, 2], [1, 0]]) + b = np.array([[1, 1], [0, 0]]) + c = np.array([[1, 0], [1, 0]]) + assert_equal(a.copy().argsort(axis=0), b) + assert_equal(a.copy().argsort(axis=1), c) + assert_equal(a.copy().argsort(), c) + + # check axis handling for multidimensional empty arrays + a = np.array([]) + a.shape = (3, 2, 1, 0) + for axis in range(-a.ndim, a.ndim): + msg = 'test empty array argsort with axis={0}'.format(axis) + assert_equal(np.argsort(a, axis=axis), + np.zeros_like(a, dtype=np.intp), msg) + msg = 'test empty array argsort with axis=None' + assert_equal(np.argsort(a, axis=None), + np.zeros_like(a.ravel(), dtype=np.intp), msg) + + # check that stable argsorts are stable + r = np.arange(100) + # scalars + a = np.zeros(100) + assert_equal(a.argsort(kind='m'), r) + # complex + a = np.zeros(100, dtype=complex) + assert_equal(a.argsort(kind='m'), r) + # string + a = np.array(['aaaaaaaaa' for i in range(100)]) + assert_equal(a.argsort(kind='m'), r) + # unicode + a = np.array(['aaaaaaaaa' for i in range(100)], dtype=np.str_) + assert_equal(a.argsort(kind='m'), r) + + with assert_raises_regex( + ValueError, + "kind` and `stable` parameters can't be provided at the same time" + ): + np.argsort(a, kind="stable", stable=True) + + def test_sort_unicode_kind(self): + d = np.arange(10) + k = b'\xc3\xa4'.decode("UTF8") + assert_raises(ValueError, d.sort, kind=k) + assert_raises(ValueError, d.argsort, kind=k) + + @pytest.mark.parametrize('a', [ + np.array([0, 1, np.nan], dtype=np.float16), + np.array([0, 1, np.nan], dtype=np.float32), + np.array([0, 1, np.nan]), + ]) + def test_searchsorted_floats(self, a): + # test for floats arrays containing nans. Explicitly test + # half, single, and double precision floats to verify that + # the NaN-handling is correct. + msg = "Test real (%s) searchsorted with nans, side='l'" % a.dtype + b = a.searchsorted(a, side='left') + assert_equal(b, np.arange(3), msg) + msg = "Test real (%s) searchsorted with nans, side='r'" % a.dtype + b = a.searchsorted(a, side='right') + assert_equal(b, np.arange(1, 4), msg) + # check keyword arguments + a.searchsorted(v=1) + x = np.array([0, 1, np.nan], dtype='float32') + y = np.searchsorted(x, x[-1]) + assert_equal(y, 2) + + def test_searchsorted_complex(self): + # test for complex arrays containing nans. + # The search sorted routines use the compare functions for the + # array type, so this checks if that is consistent with the sort + # order. + # check double complex + a = np.zeros(9, dtype=np.complex128) + a.real += [0, 0, 1, 1, 0, 1, np.nan, np.nan, np.nan] + a.imag += [0, 1, 0, 1, np.nan, np.nan, 0, 1, np.nan] + msg = "Test complex searchsorted with nans, side='l'" + b = a.searchsorted(a, side='left') + assert_equal(b, np.arange(9), msg) + msg = "Test complex searchsorted with nans, side='r'" + b = a.searchsorted(a, side='right') + assert_equal(b, np.arange(1, 10), msg) + msg = "Test searchsorted with little endian, side='l'" + a = np.array([0, 128], dtype=' p[:, i]).all(), + msg="%d: %r < %r" % (i, p[:, i], p[:, i + 1:].T)) + for row in range(p.shape[0]): + self.assert_partitioned(p[row], [i]) + self.assert_partitioned(parg[row], [i]) + + p = np.partition(d0, i, axis=0, kind=k) + parg = d0[np.argpartition(d0, i, axis=0, kind=k), + np.arange(d0.shape[1])[None, :]] + aae(p[i, :], np.array([i] * d1.shape[0], dtype=dt)) + # array_less does not seem to work right + at((p[:i, :] <= p[i, :]).all(), + msg="%d: %r <= %r" % (i, p[i, :], p[:i, :])) + at((p[i + 1:, :] > p[i, :]).all(), + msg="%d: %r < %r" % (i, p[i, :], p[:, i + 1:])) + for col in range(p.shape[1]): + self.assert_partitioned(p[:, col], [i]) + self.assert_partitioned(parg[:, col], [i]) + + # check inplace + dc = d.copy() + dc.partition(i, kind=k) + assert_equal(dc, np.partition(d, i, kind=k)) + dc = d0.copy() + dc.partition(i, axis=0, kind=k) + assert_equal(dc, np.partition(d0, i, axis=0, kind=k)) + dc = d1.copy() + dc.partition(i, axis=1, kind=k) + assert_equal(dc, np.partition(d1, i, axis=1, kind=k)) + + def assert_partitioned(self, d, kth): + prev = 0 + for k in np.sort(kth): + assert_array_compare(operator.__le__, d[prev:k], d[k], + err_msg='kth %d' % k) + assert_((d[k:] >= d[k]).all(), + msg="kth %d, %r not greater equal %r" % (k, d[k:], d[k])) + prev = k + 1 + + def test_partition_iterative(self): + d = np.arange(17) + kth = (0, 1, 2, 429, 231) + assert_raises(ValueError, d.partition, kth) + assert_raises(ValueError, d.argpartition, kth) + d = np.arange(10).reshape((2, 5)) + assert_raises(ValueError, d.partition, kth, axis=0) + assert_raises(ValueError, d.partition, kth, axis=1) + assert_raises(ValueError, np.partition, d, kth, axis=1) + assert_raises(ValueError, np.partition, d, kth, axis=None) + + d = np.array([3, 4, 2, 1]) + p = np.partition(d, (0, 3)) + self.assert_partitioned(p, (0, 3)) + self.assert_partitioned(d[np.argpartition(d, (0, 3))], (0, 3)) + + assert_array_equal(p, np.partition(d, (-3, -1))) + assert_array_equal(p, d[np.argpartition(d, (-3, -1))]) + + d = np.arange(17) + np.random.shuffle(d) + d.partition(range(d.size)) + assert_array_equal(np.arange(17), d) + np.random.shuffle(d) + assert_array_equal(np.arange(17), d[d.argpartition(range(d.size))]) + + # test unsorted kth + d = np.arange(17) + np.random.shuffle(d) + keys = np.array([1, 3, 8, -2]) + np.random.shuffle(d) + p = np.partition(d, keys) + self.assert_partitioned(p, keys) + p = d[np.argpartition(d, keys)] + self.assert_partitioned(p, keys) + np.random.shuffle(keys) + assert_array_equal(np.partition(d, keys), p) + assert_array_equal(d[np.argpartition(d, keys)], p) + + # equal kth + d = np.arange(20)[::-1] + self.assert_partitioned(np.partition(d, [5]*4), [5]) + self.assert_partitioned(np.partition(d, [5]*4 + [6, 13]), + [5]*4 + [6, 13]) + self.assert_partitioned(d[np.argpartition(d, [5]*4)], [5]) + self.assert_partitioned(d[np.argpartition(d, [5]*4 + [6, 13])], + [5]*4 + [6, 13]) + + d = np.arange(12) + np.random.shuffle(d) + d1 = np.tile(np.arange(12), (4, 1)) + map(np.random.shuffle, d1) + d0 = np.transpose(d1) + + kth = (1, 6, 7, -1) + p = np.partition(d1, kth, axis=1) + pa = d1[np.arange(d1.shape[0])[:, None], + d1.argpartition(kth, axis=1)] + assert_array_equal(p, pa) + for i in range(d1.shape[0]): + self.assert_partitioned(p[i,:], kth) + p = np.partition(d0, kth, axis=0) + pa = d0[np.argpartition(d0, kth, axis=0), + np.arange(d0.shape[1])[None,:]] + assert_array_equal(p, pa) + for i in range(d0.shape[1]): + self.assert_partitioned(p[:, i], kth) + + def test_partition_cdtype(self): + d = np.array([('Galahad', 1.7, 38), ('Arthur', 1.8, 41), + ('Lancelot', 1.9, 38)], + dtype=[('name', '|S10'), ('height', ' (numpy ufunc, has_in_place_version, preferred_dtype) + ops = { + 'add': (np.add, True, float), + 'sub': (np.subtract, True, float), + 'mul': (np.multiply, True, float), + 'truediv': (np.true_divide, True, float), + 'floordiv': (np.floor_divide, True, float), + 'mod': (np.remainder, True, float), + 'divmod': (np.divmod, False, float), + 'pow': (np.power, True, int), + 'lshift': (np.left_shift, True, int), + 'rshift': (np.right_shift, True, int), + 'and': (np.bitwise_and, True, int), + 'xor': (np.bitwise_xor, True, int), + 'or': (np.bitwise_or, True, int), + 'matmul': (np.matmul, True, float), + # 'ge': (np.less_equal, False), + # 'gt': (np.less, False), + # 'le': (np.greater_equal, False), + # 'lt': (np.greater, False), + # 'eq': (np.equal, False), + # 'ne': (np.not_equal, False), + } + + class Coerced(Exception): + pass + + def array_impl(self): + raise Coerced + + def op_impl(self, other): + return "forward" + + def rop_impl(self, other): + return "reverse" + + def iop_impl(self, other): + return "in-place" + + def array_ufunc_impl(self, ufunc, method, *args, **kwargs): + return ("__array_ufunc__", ufunc, method, args, kwargs) + + # Create an object with the given base, in the given module, with a + # bunch of placeholder __op__ methods, and optionally a + # __array_ufunc__ and __array_priority__. + def make_obj(base, array_priority=False, array_ufunc=False, + alleged_module="__main__"): + class_namespace = {"__array__": array_impl} + if array_priority is not False: + class_namespace["__array_priority__"] = array_priority + for op in ops: + class_namespace["__{0}__".format(op)] = op_impl + class_namespace["__r{0}__".format(op)] = rop_impl + class_namespace["__i{0}__".format(op)] = iop_impl + if array_ufunc is not False: + class_namespace["__array_ufunc__"] = array_ufunc + eval_namespace = {"base": base, + "class_namespace": class_namespace, + "__name__": alleged_module, + } + MyType = eval("type('MyType', (base,), class_namespace)", + eval_namespace) + if issubclass(MyType, np.ndarray): + # Use this range to avoid special case weirdnesses around + # divide-by-0, pow(x, 2), overflow due to pow(big, big), etc. + return np.arange(3, 7).reshape(2, 2).view(MyType) + else: + return MyType() + + def check(obj, binop_override_expected, ufunc_override_expected, + inplace_override_expected, check_scalar=True): + for op, (ufunc, has_inplace, dtype) in ops.items(): + err_msg = ('op: %s, ufunc: %s, has_inplace: %s, dtype: %s' + % (op, ufunc, has_inplace, dtype)) + check_objs = [np.arange(3, 7, dtype=dtype).reshape(2, 2)] + if check_scalar: + check_objs.append(check_objs[0][0]) + for arr in check_objs: + arr_method = getattr(arr, "__{0}__".format(op)) + + def first_out_arg(result): + if op == "divmod": + assert_(isinstance(result, tuple)) + return result[0] + else: + return result + + # arr __op__ obj + if binop_override_expected: + assert_equal(arr_method(obj), NotImplemented, err_msg) + elif ufunc_override_expected: + assert_equal(arr_method(obj)[0], "__array_ufunc__", + err_msg) + else: + if (isinstance(obj, np.ndarray) and + (type(obj).__array_ufunc__ is + np.ndarray.__array_ufunc__)): + # __array__ gets ignored + res = first_out_arg(arr_method(obj)) + assert_(res.__class__ is obj.__class__, err_msg) + else: + assert_raises((TypeError, Coerced), + arr_method, obj, err_msg=err_msg) + # obj __op__ arr + arr_rmethod = getattr(arr, "__r{0}__".format(op)) + if ufunc_override_expected: + res = arr_rmethod(obj) + assert_equal(res[0], "__array_ufunc__", + err_msg=err_msg) + assert_equal(res[1], ufunc, err_msg=err_msg) + else: + if (isinstance(obj, np.ndarray) and + (type(obj).__array_ufunc__ is + np.ndarray.__array_ufunc__)): + # __array__ gets ignored + res = first_out_arg(arr_rmethod(obj)) + assert_(res.__class__ is obj.__class__, err_msg) + else: + # __array_ufunc__ = "asdf" creates a TypeError + assert_raises((TypeError, Coerced), + arr_rmethod, obj, err_msg=err_msg) + + # arr __iop__ obj + # array scalars don't have in-place operators + if has_inplace and isinstance(arr, np.ndarray): + arr_imethod = getattr(arr, "__i{0}__".format(op)) + if inplace_override_expected: + assert_equal(arr_method(obj), NotImplemented, + err_msg=err_msg) + elif ufunc_override_expected: + res = arr_imethod(obj) + assert_equal(res[0], "__array_ufunc__", err_msg) + assert_equal(res[1], ufunc, err_msg) + assert_(type(res[-1]["out"]) is tuple, err_msg) + assert_(res[-1]["out"][0] is arr, err_msg) + else: + if (isinstance(obj, np.ndarray) and + (type(obj).__array_ufunc__ is + np.ndarray.__array_ufunc__)): + # __array__ gets ignored + assert_(arr_imethod(obj) is arr, err_msg) + else: + assert_raises((TypeError, Coerced), + arr_imethod, obj, + err_msg=err_msg) + + op_fn = getattr(operator, op, None) + if op_fn is None: + op_fn = getattr(operator, op + "_", None) + if op_fn is None: + op_fn = getattr(builtins, op) + assert_equal(op_fn(obj, arr), "forward", err_msg) + if not isinstance(obj, np.ndarray): + if binop_override_expected: + assert_equal(op_fn(arr, obj), "reverse", err_msg) + elif ufunc_override_expected: + assert_equal(op_fn(arr, obj)[0], "__array_ufunc__", + err_msg) + if ufunc_override_expected: + assert_equal(ufunc(obj, arr)[0], "__array_ufunc__", + err_msg) + + # No array priority, no array_ufunc -> nothing called + check(make_obj(object), False, False, False) + # Negative array priority, no array_ufunc -> nothing called + # (has to be very negative, because scalar priority is -1000000.0) + check(make_obj(object, array_priority=-2**30), False, False, False) + # Positive array priority, no array_ufunc -> binops and iops only + check(make_obj(object, array_priority=1), True, False, True) + # ndarray ignores array_priority for ndarray subclasses + check(make_obj(np.ndarray, array_priority=1), False, False, False, + check_scalar=False) + # Positive array_priority and array_ufunc -> array_ufunc only + check(make_obj(object, array_priority=1, + array_ufunc=array_ufunc_impl), False, True, False) + check(make_obj(np.ndarray, array_priority=1, + array_ufunc=array_ufunc_impl), False, True, False) + # array_ufunc set to None -> defer binops only + check(make_obj(object, array_ufunc=None), True, False, False) + check(make_obj(np.ndarray, array_ufunc=None), True, False, False, + check_scalar=False) + + @pytest.mark.parametrize("priority", [None, "runtime error"]) + def test_ufunc_binop_bad_array_priority(self, priority): + # Mainly checks that this does not crash. The second array has a lower + # priority than -1 ("error value"). If the __radd__ actually exists, + # bad things can happen (I think via the scalar paths). + # In principle both of these can probably just be errors in the future. + class BadPriority: + @property + def __array_priority__(self): + if priority == "runtime error": + raise RuntimeError("RuntimeError in __array_priority__!") + return priority + + def __radd__(self, other): + return "result" + + class LowPriority(np.ndarray): + __array_priority__ = -1000 + + # Priority failure uses the same as scalars (smaller -1000). So the + # LowPriority wins with 'result' for each element (inner operation). + res = np.arange(3).view(LowPriority) + BadPriority() + assert res.shape == (3,) + assert res[0] == 'result' + + @pytest.mark.parametrize("scalar", [ + np.longdouble(1), np.timedelta64(120, 'm')]) + @pytest.mark.parametrize("op", [operator.add, operator.xor]) + def test_scalar_binop_guarantees_ufunc(self, scalar, op): + # Test that __array_ufunc__ will always cause ufunc use even when + # we have to protect some other calls from recursing (see gh-26904). + class SomeClass: + def __array_ufunc__(self, ufunc, method, *inputs, **kw): + return "result" + + assert SomeClass() + scalar == "result" + assert scalar + SomeClass() == "result" + + def test_ufunc_override_normalize_signature(self): + # gh-5674 + class SomeClass: + def __array_ufunc__(self, ufunc, method, *inputs, **kw): + return kw + + a = SomeClass() + kw = np.add(a, [1]) + assert_('sig' not in kw and 'signature' not in kw) + kw = np.add(a, [1], sig='ii->i') + assert_('sig' not in kw and 'signature' in kw) + assert_equal(kw['signature'], 'ii->i') + kw = np.add(a, [1], signature='ii->i') + assert_('sig' not in kw and 'signature' in kw) + assert_equal(kw['signature'], 'ii->i') + + def test_array_ufunc_index(self): + # Check that index is set appropriately, also if only an output + # is passed on (latter is another regression tests for github bug 4753) + # This also checks implicitly that 'out' is always a tuple. + class CheckIndex: + def __array_ufunc__(self, ufunc, method, *inputs, **kw): + for i, a in enumerate(inputs): + if a is self: + return i + # calls below mean we must be in an output. + for j, a in enumerate(kw['out']): + if a is self: + return (j,) + + a = CheckIndex() + dummy = np.arange(2.) + # 1 input, 1 output + assert_equal(np.sin(a), 0) + assert_equal(np.sin(dummy, a), (0,)) + assert_equal(np.sin(dummy, out=a), (0,)) + assert_equal(np.sin(dummy, out=(a,)), (0,)) + assert_equal(np.sin(a, a), 0) + assert_equal(np.sin(a, out=a), 0) + assert_equal(np.sin(a, out=(a,)), 0) + # 1 input, 2 outputs + assert_equal(np.modf(dummy, a), (0,)) + assert_equal(np.modf(dummy, None, a), (1,)) + assert_equal(np.modf(dummy, dummy, a), (1,)) + assert_equal(np.modf(dummy, out=(a, None)), (0,)) + assert_equal(np.modf(dummy, out=(a, dummy)), (0,)) + assert_equal(np.modf(dummy, out=(None, a)), (1,)) + assert_equal(np.modf(dummy, out=(dummy, a)), (1,)) + assert_equal(np.modf(a, out=(dummy, a)), 0) + with assert_raises(TypeError): + # Out argument must be tuple, since there are multiple outputs + np.modf(dummy, out=a) + + assert_raises(ValueError, np.modf, dummy, out=(a,)) + + # 2 inputs, 1 output + assert_equal(np.add(a, dummy), 0) + assert_equal(np.add(dummy, a), 1) + assert_equal(np.add(dummy, dummy, a), (0,)) + assert_equal(np.add(dummy, a, a), 1) + assert_equal(np.add(dummy, dummy, out=a), (0,)) + assert_equal(np.add(dummy, dummy, out=(a,)), (0,)) + assert_equal(np.add(a, dummy, out=a), 0) + + def test_out_override(self): + # regression test for github bug 4753 + class OutClass(np.ndarray): + def __array_ufunc__(self, ufunc, method, *inputs, **kw): + if 'out' in kw: + tmp_kw = kw.copy() + tmp_kw.pop('out') + func = getattr(ufunc, method) + kw['out'][0][...] = func(*inputs, **tmp_kw) + + A = np.array([0]).view(OutClass) + B = np.array([5]) + C = np.array([6]) + np.multiply(C, B, A) + assert_equal(A[0], 30) + assert_(isinstance(A, OutClass)) + A[0] = 0 + np.multiply(C, B, out=A) + assert_equal(A[0], 30) + assert_(isinstance(A, OutClass)) + + def test_pow_override_with_errors(self): + # regression test for gh-9112 + class PowerOnly(np.ndarray): + def __array_ufunc__(self, ufunc, method, *inputs, **kw): + if ufunc is not np.power: + raise NotImplementedError + return "POWER!" + # explicit cast to float, to ensure the fast power path is taken. + a = np.array(5., dtype=np.float64).view(PowerOnly) + assert_equal(a ** 2.5, "POWER!") + with assert_raises(NotImplementedError): + a ** 0.5 + with assert_raises(NotImplementedError): + a ** 0 + with assert_raises(NotImplementedError): + a ** 1 + with assert_raises(NotImplementedError): + a ** -1 + with assert_raises(NotImplementedError): + a ** 2 + + def test_pow_array_object_dtype(self): + # test pow on arrays of object dtype + class SomeClass: + def __init__(self, num=None): + self.num = num + + # want to ensure a fast pow path is not taken + def __mul__(self, other): + raise AssertionError('__mul__ should not be called') + + def __div__(self, other): + raise AssertionError('__div__ should not be called') + + def __pow__(self, exp): + return SomeClass(num=self.num ** exp) + + def __eq__(self, other): + if isinstance(other, SomeClass): + return self.num == other.num + + __rpow__ = __pow__ + + def pow_for(exp, arr): + return np.array([x ** exp for x in arr]) + + obj_arr = np.array([SomeClass(1), SomeClass(2), SomeClass(3)]) + + assert_equal(obj_arr ** 0.5, pow_for(0.5, obj_arr)) + assert_equal(obj_arr ** 0, pow_for(0, obj_arr)) + assert_equal(obj_arr ** 1, pow_for(1, obj_arr)) + assert_equal(obj_arr ** -1, pow_for(-1, obj_arr)) + assert_equal(obj_arr ** 2, pow_for(2, obj_arr)) + + def test_pos_array_ufunc_override(self): + class A(np.ndarray): + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + return getattr(ufunc, method)(*[i.view(np.ndarray) for + i in inputs], **kwargs) + tst = np.array('foo').view(A) + with assert_raises(TypeError): + +tst + + +class TestTemporaryElide: + # elision is only triggered on relatively large arrays + + def test_extension_incref_elide(self): + # test extension (e.g. cython) calling PyNumber_* slots without + # increasing the reference counts + # + # def incref_elide(a): + # d = input.copy() # refcount 1 + # return d, d + d # PyNumber_Add without increasing refcount + from numpy._core._multiarray_tests import incref_elide + d = np.ones(100000) + orig, res = incref_elide(d) + d + d + # the return original should not be changed to an inplace operation + assert_array_equal(orig, d) + assert_array_equal(res, d + d) + + def test_extension_incref_elide_stack(self): + # scanning if the refcount == 1 object is on the python stack to check + # that we are called directly from python is flawed as object may still + # be above the stack pointer and we have no access to the top of it + # + # def incref_elide_l(d): + # return l[4] + l[4] # PyNumber_Add without increasing refcount + from numpy._core._multiarray_tests import incref_elide_l + # padding with 1 makes sure the object on the stack is not overwritten + l = [1, 1, 1, 1, np.ones(100000)] + res = incref_elide_l(l) + # the return original should not be changed to an inplace operation + assert_array_equal(l[4], np.ones(100000)) + assert_array_equal(res, l[4] + l[4]) + + def test_temporary_with_cast(self): + # check that we don't elide into a temporary which would need casting + d = np.ones(200000, dtype=np.int64) + r = ((d + d) + np.array(2**222, dtype='O')) + assert_equal(r.dtype, np.dtype('O')) + + r = ((d + d) / 2) + assert_equal(r.dtype, np.dtype('f8')) + + r = np.true_divide((d + d), 2) + assert_equal(r.dtype, np.dtype('f8')) + + r = ((d + d) / 2.) + assert_equal(r.dtype, np.dtype('f8')) + + r = ((d + d) // 2) + assert_equal(r.dtype, np.dtype(np.int64)) + + # commutative elision into the astype result + f = np.ones(100000, dtype=np.float32) + assert_equal(((f + f) + f.astype(np.float64)).dtype, np.dtype('f8')) + + # no elision into lower type + d = f.astype(np.float64) + assert_equal(((f + f) + d).dtype, d.dtype) + l = np.ones(100000, dtype=np.longdouble) + assert_equal(((d + d) + l).dtype, l.dtype) + + # test unary abs with different output dtype + for dt in (np.complex64, np.complex128, np.clongdouble): + c = np.ones(100000, dtype=dt) + r = abs(c * 2.0) + assert_equal(r.dtype, np.dtype('f%d' % (c.itemsize // 2))) + + def test_elide_broadcast(self): + # test no elision on broadcast to higher dimension + # only triggers elision code path in debug mode as triggering it in + # normal mode needs 256kb large matching dimension, so a lot of memory + d = np.ones((2000, 1), dtype=int) + b = np.ones((2000), dtype=bool) + r = (1 - d) + b + assert_equal(r, 1) + assert_equal(r.shape, (2000, 2000)) + + def test_elide_scalar(self): + # check inplace op does not create ndarray from scalars + a = np.bool() + assert_(type(~(a & a)) is np.bool) + + def test_elide_scalar_readonly(self): + # The imaginary part of a real array is readonly. This needs to go + # through fast_scalar_power which is only called for powers of + # +1, -1, 0, 0.5, and 2, so use 2. Also need valid refcount for + # elision which can be gotten for the imaginary part of a real + # array. Should not error. + a = np.empty(100000, dtype=np.float64) + a.imag ** 2 + + def test_elide_readonly(self): + # don't try to elide readonly temporaries + r = np.asarray(np.broadcast_to(np.zeros(1), 100000).flat) * 0.0 + assert_equal(r, 0) + + def test_elide_updateifcopy(self): + a = np.ones(2**20)[::2] + b = a.flat.__array__() + 1 + del b + assert_equal(a, 1) + + +class TestCAPI: + def test_IsPythonScalar(self): + from numpy._core._multiarray_tests import IsPythonScalar + assert_(IsPythonScalar(b'foobar')) + assert_(IsPythonScalar(1)) + assert_(IsPythonScalar(2**80)) + assert_(IsPythonScalar(2.)) + assert_(IsPythonScalar("a")) + + @pytest.mark.parametrize("converter", + [_multiarray_tests.run_scalar_intp_converter, + _multiarray_tests.run_scalar_intp_from_sequence]) + def test_intp_sequence_converters(self, converter): + # Test simple values (-1 is special for error return paths) + assert converter(10) == (10,) + assert converter(-1) == (-1,) + # A 0-D array looks a bit like a sequence but must take the integer + # path: + assert converter(np.array(123)) == (123,) + # Test simple sequences (intp_from_sequence only supports length 1): + assert converter((10,)) == (10,) + assert converter(np.array([11])) == (11,) + + @pytest.mark.parametrize("converter", + [_multiarray_tests.run_scalar_intp_converter, + _multiarray_tests.run_scalar_intp_from_sequence]) + @pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") + def test_intp_sequence_converters_errors(self, converter): + with pytest.raises(TypeError, + match="expected a sequence of integers or a single integer, "): + converter(object()) + with pytest.raises(TypeError, + match="expected a sequence of integers or a single integer, " + "got '32.0'"): + converter(32.) + with pytest.raises(TypeError, + match="'float' object cannot be interpreted as an integer"): + converter([32.]) + with pytest.raises(ValueError, + match="Maximum allowed dimension"): + # These converters currently convert overflows to a ValueError + converter(2**64) + + +class TestSubscripting: + def test_test_zero_rank(self): + x = np.array([1, 2, 3]) + assert_(isinstance(x[0], np.int_)) + assert_(type(x[0, ...]) is np.ndarray) + + +class TestPickling: + @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL >= 5, + reason=('this tests the error messages when trying to' + 'protocol 5 although it is not available')) + def test_correct_protocol5_error_message(self): + array = np.arange(10) + + def test_record_array_with_object_dtype(self): + my_object = object() + + arr_with_object = np.array( + [(my_object, 1, 2.0)], + dtype=[('a', object), ('b', int), ('c', float)]) + arr_without_object = np.array( + [('xxx', 1, 2.0)], + dtype=[('a', str), ('b', int), ('c', float)]) + + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + depickled_arr_with_object = pickle.loads( + pickle.dumps(arr_with_object, protocol=proto)) + depickled_arr_without_object = pickle.loads( + pickle.dumps(arr_without_object, protocol=proto)) + + assert_equal(arr_with_object.dtype, + depickled_arr_with_object.dtype) + assert_equal(arr_without_object.dtype, + depickled_arr_without_object.dtype) + + @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5, + reason="requires pickle protocol 5") + def test_f_contiguous_array(self): + f_contiguous_array = np.array([[1, 2, 3], [4, 5, 6]], order='F') + buffers = [] + + # When using pickle protocol 5, Fortran-contiguous arrays can be + # serialized using out-of-band buffers + bytes_string = pickle.dumps(f_contiguous_array, protocol=5, + buffer_callback=buffers.append) + + assert len(buffers) > 0 + + depickled_f_contiguous_array = pickle.loads(bytes_string, + buffers=buffers) + + assert_equal(f_contiguous_array, depickled_f_contiguous_array) + + def test_non_contiguous_array(self): + non_contiguous_array = np.arange(12).reshape(3, 4)[:, :2] + assert not non_contiguous_array.flags.c_contiguous + assert not non_contiguous_array.flags.f_contiguous + + # make sure non-contiguous arrays can be pickled-depickled + # using any protocol + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + depickled_non_contiguous_array = pickle.loads( + pickle.dumps(non_contiguous_array, protocol=proto)) + + assert_equal(non_contiguous_array, depickled_non_contiguous_array) + + def test_roundtrip(self): + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + carray = np.array([[2, 9], [7, 0], [3, 8]]) + DATA = [ + carray, + np.transpose(carray), + np.array([('xxx', 1, 2.0)], dtype=[('a', (str, 3)), ('b', int), + ('c', float)]) + ] + + refs = [weakref.ref(a) for a in DATA] + for a in DATA: + assert_equal( + a, pickle.loads(pickle.dumps(a, protocol=proto)), + err_msg="%r" % a) + del a, DATA, carray + break_cycles() + # check for reference leaks (gh-12793) + for ref in refs: + assert ref() is None + + def _loads(self, obj): + return pickle.loads(obj, encoding='latin1') + + # version 0 pickles, using protocol=2 to pickle + # version 0 doesn't have a version field + def test_version0_int8(self): + s = b"\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x04\x85cnumpy\ndtype\nq\x04U\x02i1K\x00K\x01\x87Rq\x05(U\x01|NNJ\xff\xff\xff\xffJ\xff\xff\xff\xfftb\x89U\x04\x01\x02\x03\x04tb." # noqa + a = np.array([1, 2, 3, 4], dtype=np.int8) + p = self._loads(s) + assert_equal(a, p) + + def test_version0_float32(self): + s = b"\x80\x02cnumpy.core._internal\n_reconstruct\nq\x01cnumpy\nndarray\nq\x02K\x00\x85U\x01b\x87Rq\x03(K\x04\x85cnumpy\ndtype\nq\x04U\x02f4K\x00K\x01\x87Rq\x05(U\x01= g2, [g1[i] >= g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 < g2, [g1[i] < g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 > g2, [g1[i] > g2[i] for i in [0, 1, 2]]) + + def test_mixed(self): + g1 = np.array(["spam", "spa", "spammer", "and eggs"]) + g2 = "spam" + assert_array_equal(g1 == g2, [x == g2 for x in g1]) + assert_array_equal(g1 != g2, [x != g2 for x in g1]) + assert_array_equal(g1 < g2, [x < g2 for x in g1]) + assert_array_equal(g1 > g2, [x > g2 for x in g1]) + assert_array_equal(g1 <= g2, [x <= g2 for x in g1]) + assert_array_equal(g1 >= g2, [x >= g2 for x in g1]) + + def test_unicode(self): + g1 = np.array(["This", "is", "example"]) + g2 = np.array(["This", "was", "example"]) + assert_array_equal(g1 == g2, [g1[i] == g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 != g2, [g1[i] != g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 <= g2, [g1[i] <= g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 >= g2, [g1[i] >= g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 < g2, [g1[i] < g2[i] for i in [0, 1, 2]]) + assert_array_equal(g1 > g2, [g1[i] > g2[i] for i in [0, 1, 2]]) + +class TestArgmaxArgminCommon: + + sizes = [(), (3,), (3, 2), (2, 3), + (3, 3), (2, 3, 4), (4, 3, 2), + (1, 2, 3, 4), (2, 3, 4, 1), + (3, 4, 1, 2), (4, 1, 2, 3), + (64,), (128,), (256,)] + + @pytest.mark.parametrize("size, axis", itertools.chain(*[[(size, axis) + for axis in list(range(-len(size), len(size))) + [None]] + for size in sizes])) + @pytest.mark.parametrize('method', [np.argmax, np.argmin]) + def test_np_argmin_argmax_keepdims(self, size, axis, method): + + arr = np.random.normal(size=size) + + # contiguous arrays + if axis is None: + new_shape = [1 for _ in range(len(size))] + else: + new_shape = list(size) + new_shape[axis] = 1 + new_shape = tuple(new_shape) + + _res_orig = method(arr, axis=axis) + res_orig = _res_orig.reshape(new_shape) + res = method(arr, axis=axis, keepdims=True) + assert_equal(res, res_orig) + assert_(res.shape == new_shape) + outarray = np.empty(res.shape, dtype=res.dtype) + res1 = method(arr, axis=axis, out=outarray, + keepdims=True) + assert_(res1 is outarray) + assert_equal(res, outarray) + + if len(size) > 0: + wrong_shape = list(new_shape) + if axis is not None: + wrong_shape[axis] = 2 + else: + wrong_shape[0] = 2 + wrong_outarray = np.empty(wrong_shape, dtype=res.dtype) + with pytest.raises(ValueError): + method(arr.T, axis=axis, + out=wrong_outarray, keepdims=True) + + # non-contiguous arrays + if axis is None: + new_shape = [1 for _ in range(len(size))] + else: + new_shape = list(size)[::-1] + new_shape[axis] = 1 + new_shape = tuple(new_shape) + + _res_orig = method(arr.T, axis=axis) + res_orig = _res_orig.reshape(new_shape) + res = method(arr.T, axis=axis, keepdims=True) + assert_equal(res, res_orig) + assert_(res.shape == new_shape) + outarray = np.empty(new_shape[::-1], dtype=res.dtype) + outarray = outarray.T + res1 = method(arr.T, axis=axis, out=outarray, + keepdims=True) + assert_(res1 is outarray) + assert_equal(res, outarray) + + if len(size) > 0: + # one dimension lesser for non-zero sized + # array should raise an error + with pytest.raises(ValueError): + method(arr[0], axis=axis, + out=outarray, keepdims=True) + + if len(size) > 0: + wrong_shape = list(new_shape) + if axis is not None: + wrong_shape[axis] = 2 + else: + wrong_shape[0] = 2 + wrong_outarray = np.empty(wrong_shape, dtype=res.dtype) + with pytest.raises(ValueError): + method(arr.T, axis=axis, + out=wrong_outarray, keepdims=True) + + @pytest.mark.parametrize('method', ['max', 'min']) + def test_all(self, method): + a = np.random.normal(0, 1, (4, 5, 6, 7, 8)) + arg_method = getattr(a, 'arg' + method) + val_method = getattr(a, method) + for i in range(a.ndim): + a_maxmin = val_method(i) + aarg_maxmin = arg_method(i) + axes = list(range(a.ndim)) + axes.remove(i) + assert_(np.all(a_maxmin == aarg_maxmin.choose( + *a.transpose(i, *axes)))) + + @pytest.mark.parametrize('method', ['argmax', 'argmin']) + def test_output_shape(self, method): + # see also gh-616 + a = np.ones((10, 5)) + arg_method = getattr(a, method) + # Check some simple shape mismatches + out = np.ones(11, dtype=np.int_) + assert_raises(ValueError, arg_method, -1, out) + + out = np.ones((2, 5), dtype=np.int_) + assert_raises(ValueError, arg_method, -1, out) + + # these could be relaxed possibly (used to allow even the previous) + out = np.ones((1, 10), dtype=np.int_) + assert_raises(ValueError, arg_method, -1, out) + + out = np.ones(10, dtype=np.int_) + arg_method(-1, out=out) + assert_equal(out, arg_method(-1)) + + @pytest.mark.parametrize('ndim', [0, 1]) + @pytest.mark.parametrize('method', ['argmax', 'argmin']) + def test_ret_is_out(self, ndim, method): + a = np.ones((4,) + (256,)*ndim) + arg_method = getattr(a, method) + out = np.empty((256,)*ndim, dtype=np.intp) + ret = arg_method(axis=0, out=out) + assert ret is out + + @pytest.mark.parametrize('np_array, method, idx, val', + [(np.zeros, 'argmax', 5942, "as"), + (np.ones, 'argmin', 6001, "0")]) + def test_unicode(self, np_array, method, idx, val): + d = np_array(6031, dtype='= cmin)) + assert_(np.all(x <= cmax)) + + def _clip_type(self, type_group, array_max, + clip_min, clip_max, inplace=False, + expected_min=None, expected_max=None): + if expected_min is None: + expected_min = clip_min + if expected_max is None: + expected_max = clip_max + + for T in np._core.sctypes[type_group]: + if sys.byteorder == 'little': + byte_orders = ['=', '>'] + else: + byte_orders = ['<', '='] + + for byteorder in byte_orders: + dtype = np.dtype(T).newbyteorder(byteorder) + + x = (np.random.random(1000) * array_max).astype(dtype) + if inplace: + # The tests that call us pass clip_min and clip_max that + # might not fit in the destination dtype. They were written + # assuming the previous unsafe casting, which now must be + # passed explicitly to avoid a warning. + x.clip(clip_min, clip_max, x, casting='unsafe') + else: + x = x.clip(clip_min, clip_max) + byteorder = '=' + + if x.dtype.byteorder == '|': + byteorder = '|' + assert_equal(x.dtype.byteorder, byteorder) + self._check_range(x, expected_min, expected_max) + return x + + def test_basic(self): + for inplace in [False, True]: + self._clip_type( + 'float', 1024, -12.8, 100.2, inplace=inplace) + self._clip_type( + 'float', 1024, 0, 0, inplace=inplace) + + self._clip_type( + 'int', 1024, -120, 100, inplace=inplace) + self._clip_type( + 'int', 1024, 0, 0, inplace=inplace) + + self._clip_type( + 'uint', 1024, 0, 0, inplace=inplace) + self._clip_type( + 'uint', 1024, 10, 100, inplace=inplace) + + @pytest.mark.parametrize("inplace", [False, True]) + def test_int_out_of_range(self, inplace): + # Simple check for out-of-bound integers, also testing the in-place + # path. + x = (np.random.random(1000) * 255).astype("uint8") + out = np.empty_like(x) + res = x.clip(-1, 300, out=out if inplace else None) + assert res is out or not inplace + assert (res == x).all() + + res = x.clip(-1, 50, out=out if inplace else None) + assert res is out or not inplace + assert (res <= 50).all() + assert (res[x <= 50] == x[x <= 50]).all() + + res = x.clip(100, 1000, out=out if inplace else None) + assert res is out or not inplace + assert (res >= 100).all() + assert (res[x >= 100] == x[x >= 100]).all() + + def test_record_array(self): + rec = np.array([(-5, 2.0, 3.0), (5.0, 4.0, 3.0)], + dtype=[('x', '= 3)) + x = val.clip(min=3) + assert_(np.all(x >= 3)) + x = val.clip(max=4) + assert_(np.all(x <= 4)) + + def test_nan(self): + input_arr = np.array([-2., np.nan, 0.5, 3., 0.25, np.nan]) + result = input_arr.clip(-1, 1) + expected = np.array([-1., np.nan, 0.5, 1., 0.25, np.nan]) + assert_array_equal(result, expected) + + +class TestCompress: + def test_axis(self): + tgt = [[5, 6, 7, 8, 9]] + arr = np.arange(10).reshape(2, 5) + out = np.compress([0, 1], arr, axis=0) + assert_equal(out, tgt) + + tgt = [[1, 3], [6, 8]] + out = np.compress([0, 1, 0, 1, 0], arr, axis=1) + assert_equal(out, tgt) + + def test_truncate(self): + tgt = [[1], [6]] + arr = np.arange(10).reshape(2, 5) + out = np.compress([0, 1], arr, axis=1) + assert_equal(out, tgt) + + def test_flatten(self): + arr = np.arange(10).reshape(2, 5) + out = np.compress([0, 1], arr) + assert_equal(out, 1) + + +class TestPutmask: + def tst_basic(self, x, T, mask, val): + np.putmask(x, mask, val) + assert_equal(x[mask], np.array(val, T)) + + def test_ip_types(self): + unchecked_types = [bytes, str, np.void] + + x = np.random.random(1000)*100 + mask = x < 40 + + for val in [-100, 0, 15]: + for types in np._core.sctypes.values(): + for T in types: + if T not in unchecked_types: + if val < 0 and np.dtype(T).kind == "u": + val = np.iinfo(T).max - 99 + self.tst_basic(x.copy().astype(T), T, mask, val) + + # Also test string of a length which uses an untypical length + dt = np.dtype("S3") + self.tst_basic(x.astype(dt), dt.type, mask, dt.type(val)[:3]) + + def test_mask_size(self): + assert_raises(ValueError, np.putmask, np.array([1, 2, 3]), [True], 5) + + @pytest.mark.parametrize('dtype', ('>i4', 'f8'), ('z', '= 2, 3) + + def test_kwargs(self): + x = np.array([0, 0]) + np.putmask(x, [0, 1], [-1, -2]) + assert_array_equal(x, [0, -2]) + + x = np.array([0, 0]) + np.putmask(x, mask=[0, 1], values=[-1, -2]) + assert_array_equal(x, [0, -2]) + + x = np.array([0, 0]) + np.putmask(x, values=[-1, -2], mask=[0, 1]) + assert_array_equal(x, [0, -2]) + + with pytest.raises(TypeError): + np.putmask(a=x, values=[-1, -2], mask=[0, 1]) + + +class TestTake: + def tst_basic(self, x): + ind = list(range(x.shape[0])) + assert_array_equal(x.take(ind, axis=0), x) + + def test_ip_types(self): + unchecked_types = [bytes, str, np.void] + + x = np.random.random(24)*100 + x.shape = 2, 3, 4 + for types in np._core.sctypes.values(): + for T in types: + if T not in unchecked_types: + self.tst_basic(x.copy().astype(T)) + + # Also test string of a length which uses an untypical length + self.tst_basic(x.astype("S3")) + + def test_raise(self): + x = np.random.random(24)*100 + x.shape = 2, 3, 4 + assert_raises(IndexError, x.take, [0, 1, 2], axis=0) + assert_raises(IndexError, x.take, [-3], axis=0) + assert_array_equal(x.take([-1], axis=0)[0], x[1]) + + def test_clip(self): + x = np.random.random(24)*100 + x.shape = 2, 3, 4 + assert_array_equal(x.take([-1], axis=0, mode='clip')[0], x[0]) + assert_array_equal(x.take([2], axis=0, mode='clip')[0], x[1]) + + def test_wrap(self): + x = np.random.random(24)*100 + x.shape = 2, 3, 4 + assert_array_equal(x.take([-1], axis=0, mode='wrap')[0], x[1]) + assert_array_equal(x.take([2], axis=0, mode='wrap')[0], x[0]) + assert_array_equal(x.take([3], axis=0, mode='wrap')[0], x[1]) + + @pytest.mark.parametrize('dtype', ('>i4', 'f8'), ('z', ' 16MB + d = np.zeros(4 * 1024 ** 2) + d.tofile(tmp_filename) + assert_equal(os.path.getsize(tmp_filename), d.nbytes) + assert_array_equal(d, np.fromfile(tmp_filename)) + # check offset + with open(tmp_filename, "r+b") as f: + f.seek(d.nbytes) + d.tofile(f) + assert_equal(os.path.getsize(tmp_filename), d.nbytes * 2) + # check append mode (gh-8329) + open(tmp_filename, "w").close() # delete file contents + with open(tmp_filename, "ab") as f: + d.tofile(f) + assert_array_equal(d, np.fromfile(tmp_filename)) + with open(tmp_filename, "ab") as f: + d.tofile(f) + assert_equal(os.path.getsize(tmp_filename), d.nbytes * 2) + + def test_io_open_buffered_fromfile(self, x, tmp_filename): + # gh-6632 + x.tofile(tmp_filename) + with open(tmp_filename, 'rb', buffering=-1) as f: + y = np.fromfile(f, dtype=x.dtype) + assert_array_equal(y, x.flat) + + def test_file_position_after_fromfile(self, tmp_filename): + # gh-4118 + sizes = [io.DEFAULT_BUFFER_SIZE//8, + io.DEFAULT_BUFFER_SIZE, + io.DEFAULT_BUFFER_SIZE*8] + + for size in sizes: + with open(tmp_filename, 'wb') as f: + f.seek(size-1) + f.write(b'\0') + + for mode in ['rb', 'r+b']: + err_msg = "%d %s" % (size, mode) + + with open(tmp_filename, mode) as f: + f.read(2) + np.fromfile(f, dtype=np.float64, count=1) + pos = f.tell() + assert_equal(pos, 10, err_msg=err_msg) + + def test_file_position_after_tofile(self, tmp_filename): + # gh-4118 + sizes = [io.DEFAULT_BUFFER_SIZE//8, + io.DEFAULT_BUFFER_SIZE, + io.DEFAULT_BUFFER_SIZE*8] + + for size in sizes: + err_msg = "%d" % (size,) + + with open(tmp_filename, 'wb') as f: + f.seek(size-1) + f.write(b'\0') + f.seek(10) + f.write(b'12') + np.array([0], dtype=np.float64).tofile(f) + pos = f.tell() + assert_equal(pos, 10 + 2 + 8, err_msg=err_msg) + + with open(tmp_filename, 'r+b') as f: + f.read(2) + f.seek(0, 1) # seek between read&write required by ANSI C + np.array([0], dtype=np.float64).tofile(f) + pos = f.tell() + assert_equal(pos, 10, err_msg=err_msg) + + def test_load_object_array_fromfile(self, tmp_filename): + # gh-12300 + with open(tmp_filename, 'w') as f: + # Ensure we have a file with consistent contents + pass + + with open(tmp_filename, 'rb') as f: + assert_raises_regex(ValueError, "Cannot read into object array", + np.fromfile, f, dtype=object) + + assert_raises_regex(ValueError, "Cannot read into object array", + np.fromfile, tmp_filename, dtype=object) + + def test_fromfile_offset(self, x, tmp_filename): + with open(tmp_filename, 'wb') as f: + x.tofile(f) + + with open(tmp_filename, 'rb') as f: + y = np.fromfile(f, dtype=x.dtype, offset=0) + assert_array_equal(y, x.flat) + + with open(tmp_filename, 'rb') as f: + count_items = len(x.flat) // 8 + offset_items = len(x.flat) // 4 + offset_bytes = x.dtype.itemsize * offset_items + y = np.fromfile( + f, dtype=x.dtype, count=count_items, offset=offset_bytes + ) + assert_array_equal( + y, x.flat[offset_items:offset_items+count_items] + ) + + # subsequent seeks should stack + offset_bytes = x.dtype.itemsize + z = np.fromfile(f, dtype=x.dtype, offset=offset_bytes) + assert_array_equal(z, x.flat[offset_items+count_items+1:]) + + with open(tmp_filename, 'wb') as f: + x.tofile(f, sep=",") + + with open(tmp_filename, 'rb') as f: + assert_raises_regex( + TypeError, + "'offset' argument only permitted for binary files", + np.fromfile, tmp_filename, dtype=x.dtype, + sep=",", offset=1) + + @pytest.mark.skipif(IS_PYPY, reason="bug in PyPy's PyNumber_AsSsize_t") + def test_fromfile_bad_dup(self, x, tmp_filename): + def dup_str(fd): + return 'abc' + + def dup_bigint(fd): + return 2**68 + + old_dup = os.dup + try: + with open(tmp_filename, 'wb') as f: + x.tofile(f) + for dup, exc in ((dup_str, TypeError), (dup_bigint, OSError)): + os.dup = dup + assert_raises(exc, np.fromfile, f) + finally: + os.dup = old_dup + + def _check_from(self, s, value, filename, **kw): + if 'sep' not in kw: + y = np.frombuffer(s, **kw) + else: + y = np.fromstring(s, **kw) + assert_array_equal(y, value) + + with open(filename, 'wb') as f: + f.write(s) + y = np.fromfile(filename, **kw) + assert_array_equal(y, value) + + @pytest.fixture(params=["period", "comma"]) + def decimal_sep_localization(self, request): + """ + Including this fixture in a test will automatically + execute it with both types of decimal separator. + + So:: + + def test_decimal(decimal_sep_localization): + pass + + is equivalent to the following two tests:: + + def test_decimal_period_separator(): + pass + + def test_decimal_comma_separator(): + with CommaDecimalPointLocale(): + pass + """ + if request.param == "period": + yield + elif request.param == "comma": + with CommaDecimalPointLocale(): + yield + else: + assert False, request.param + + def test_nan(self, tmp_filename, decimal_sep_localization): + self._check_from( + b"nan +nan -nan NaN nan(foo) +NaN(BAR) -NAN(q_u_u_x_)", + [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], + tmp_filename, + sep=' ') + + def test_inf(self, tmp_filename, decimal_sep_localization): + self._check_from( + b"inf +inf -inf infinity -Infinity iNfInItY -inF", + [np.inf, np.inf, -np.inf, np.inf, -np.inf, np.inf, -np.inf], + tmp_filename, + sep=' ') + + def test_numbers(self, tmp_filename, decimal_sep_localization): + self._check_from( + b"1.234 -1.234 .3 .3e55 -123133.1231e+133", + [1.234, -1.234, .3, .3e55, -123133.1231e+133], + tmp_filename, + sep=' ') + + def test_binary(self, tmp_filename): + self._check_from( + b'\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x80@', + np.array([1, 2, 3, 4]), + tmp_filename, + dtype='']) + @pytest.mark.parametrize('dtype', [float, int, complex]) + def test_basic(self, byteorder, dtype): + dt = np.dtype(dtype).newbyteorder(byteorder) + x = (np.random.random((4, 7)) * 5).astype(dt) + buf = x.tobytes() + assert_array_equal(np.frombuffer(buf, dtype=dt), x.flat) + + @pytest.mark.parametrize("obj", [np.arange(10), b"12345678"]) + def test_array_base(self, obj): + # Objects (including NumPy arrays), which do not use the + # `release_buffer` slot should be directly used as a base object. + # See also gh-21612 + new = np.frombuffer(obj) + assert new.base is obj + + def test_empty(self): + assert_array_equal(np.frombuffer(b''), np.array([])) + + @pytest.mark.skipif(IS_PYPY, + reason="PyPy's memoryview currently does not track exports. See: " + "https://foss.heptapod.net/pypy/pypy/-/issues/3724") + def test_mmap_close(self): + # The old buffer protocol was not safe for some things that the new + # one is. But `frombuffer` always used the old one for a long time. + # Checks that it is safe with the new one (using memoryviews) + with tempfile.TemporaryFile(mode='wb') as tmp: + tmp.write(b"asdf") + tmp.flush() + mm = mmap.mmap(tmp.fileno(), 0) + arr = np.frombuffer(mm, dtype=np.uint8) + with pytest.raises(BufferError): + mm.close() # cannot close while array uses the buffer + del arr + mm.close() + +class TestFlat: + def setup_method(self): + a0 = np.arange(20.0) + a = a0.reshape(4, 5) + a0.shape = (4, 5) + a.flags.writeable = False + self.a = a + self.b = a[::2, ::2] + self.a0 = a0 + self.b0 = a0[::2, ::2] + + def test_contiguous(self): + testpassed = False + try: + self.a.flat[12] = 100.0 + except ValueError: + testpassed = True + assert_(testpassed) + assert_(self.a.flat[12] == 12.0) + + def test_discontiguous(self): + testpassed = False + try: + self.b.flat[4] = 100.0 + except ValueError: + testpassed = True + assert_(testpassed) + assert_(self.b.flat[4] == 12.0) + + def test___array__(self): + c = self.a.flat.__array__() + d = self.b.flat.__array__() + e = self.a0.flat.__array__() + f = self.b0.flat.__array__() + + assert_(c.flags.writeable is False) + assert_(d.flags.writeable is False) + assert_(e.flags.writeable is True) + assert_(f.flags.writeable is False) + assert_(c.flags.writebackifcopy is False) + assert_(d.flags.writebackifcopy is False) + assert_(e.flags.writebackifcopy is False) + assert_(f.flags.writebackifcopy is False) + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_refcount(self): + # includes regression test for reference count error gh-13165 + inds = [np.intp(0), np.array([True]*self.a.size), np.array([0]), None] + indtype = np.dtype(np.intp) + rc_indtype = sys.getrefcount(indtype) + for ind in inds: + rc_ind = sys.getrefcount(ind) + for _ in range(100): + try: + self.a.flat[ind] + except IndexError: + pass + assert_(abs(sys.getrefcount(ind) - rc_ind) < 50) + assert_(abs(sys.getrefcount(indtype) - rc_indtype) < 50) + + def test_index_getset(self): + it = np.arange(10).reshape(2, 1, 5).flat + with pytest.raises(AttributeError): + it.index = 10 + + for _ in it: + pass + # Check the value of `.index` is updated correctly (see also gh-19153) + # If the type was incorrect, this would show up on big-endian machines + assert it.index == it.base.size + + def test_maxdims(self): + # The flat iterator and thus attribute is currently unfortunately + # limited to only 32 dimensions (after bumping it to 64 for 2.0) + a = np.ones((1,) * 64) + + with pytest.raises(RuntimeError, + match=".*32 dimensions but the array has 64"): + a.flat + + +class TestResize: + + @_no_tracing + def test_basic(self): + x = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + if IS_PYPY: + x.resize((5, 5), refcheck=False) + else: + x.resize((5, 5)) + assert_array_equal(x.flat[:9], + np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).flat) + assert_array_equal(x[9:].flat, 0) + + def test_check_reference(self): + x = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + y = x + assert_raises(ValueError, x.resize, (5, 1)) + del y # avoid pyflakes unused variable warning. + + @_no_tracing + def test_int_shape(self): + x = np.eye(3) + if IS_PYPY: + x.resize(3, refcheck=False) + else: + x.resize(3) + assert_array_equal(x, np.eye(3)[0,:]) + + def test_none_shape(self): + x = np.eye(3) + x.resize(None) + assert_array_equal(x, np.eye(3)) + x.resize() + assert_array_equal(x, np.eye(3)) + + def test_0d_shape(self): + # to it multiple times to test it does not break alloc cache gh-9216 + for i in range(10): + x = np.empty((1,)) + x.resize(()) + assert_equal(x.shape, ()) + assert_equal(x.size, 1) + x = np.empty(()) + x.resize((1,)) + assert_equal(x.shape, (1,)) + assert_equal(x.size, 1) + + def test_invalid_arguments(self): + assert_raises(TypeError, np.eye(3).resize, 'hi') + assert_raises(ValueError, np.eye(3).resize, -1) + assert_raises(TypeError, np.eye(3).resize, order=1) + assert_raises(TypeError, np.eye(3).resize, refcheck='hi') + + @_no_tracing + def test_freeform_shape(self): + x = np.eye(3) + if IS_PYPY: + x.resize(3, 2, 1, refcheck=False) + else: + x.resize(3, 2, 1) + assert_(x.shape == (3, 2, 1)) + + @_no_tracing + def test_zeros_appended(self): + x = np.eye(3) + if IS_PYPY: + x.resize(2, 3, 3, refcheck=False) + else: + x.resize(2, 3, 3) + assert_array_equal(x[0], np.eye(3)) + assert_array_equal(x[1], np.zeros((3, 3))) + + @_no_tracing + def test_obj_obj(self): + # check memory is initialized on resize, gh-4857 + a = np.ones(10, dtype=[('k', object, 2)]) + if IS_PYPY: + a.resize(15, refcheck=False) + else: + a.resize(15,) + assert_equal(a.shape, (15,)) + assert_array_equal(a['k'][-5:], 0) + assert_array_equal(a['k'][:-5], 1) + + def test_empty_view(self): + # check that sizes containing a zero don't trigger a reallocate for + # already empty arrays + x = np.zeros((10, 0), int) + x_view = x[...] + x_view.resize((0, 10)) + x_view.resize((0, 100)) + + def test_check_weakref(self): + x = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + xref = weakref.ref(x) + assert_raises(ValueError, x.resize, (5, 1)) + del xref # avoid pyflakes unused variable warning. + + +class TestRecord: + def test_field_rename(self): + dt = np.dtype([('f', float), ('i', int)]) + dt.names = ['p', 'q'] + assert_equal(dt.names, ['p', 'q']) + + def test_multiple_field_name_occurrence(self): + def test_dtype_init(): + np.dtype([("A", "f8"), ("B", "f8"), ("A", "f8")]) + + # Error raised when multiple fields have the same name + assert_raises(ValueError, test_dtype_init) + + def test_bytes_fields(self): + # Bytes are not allowed in field names and not recognized in titles + # on Py3 + assert_raises(TypeError, np.dtype, [(b'a', int)]) + assert_raises(TypeError, np.dtype, [(('b', b'a'), int)]) + + dt = np.dtype([((b'a', 'b'), int)]) + assert_raises(TypeError, dt.__getitem__, b'a') + + x = np.array([(1,), (2,), (3,)], dtype=dt) + assert_raises(IndexError, x.__getitem__, b'a') + + y = x[0] + assert_raises(IndexError, y.__getitem__, b'a') + + def test_multiple_field_name_unicode(self): + def test_dtype_unicode(): + np.dtype([("\u20B9", "f8"), ("B", "f8"), ("\u20B9", "f8")]) + + # Error raised when multiple fields have the same name(unicode included) + assert_raises(ValueError, test_dtype_unicode) + + def test_fromarrays_unicode(self): + # A single name string provided to fromarrays() is allowed to be unicode + # on both Python 2 and 3: + x = np._core.records.fromarrays( + [[0], [1]], names='a,b', formats='i4,i4') + assert_equal(x['a'][0], 0) + assert_equal(x['b'][0], 1) + + def test_unicode_order(self): + # Test that we can sort with order as a unicode field name in both Python 2 and + # 3: + name = 'b' + x = np.array([1, 3, 2], dtype=[(name, int)]) + x.sort(order=name) + assert_equal(x['b'], np.array([1, 2, 3])) + + def test_field_names(self): + # Test unicode and 8-bit / byte strings can be used + a = np.zeros((1,), dtype=[('f1', 'i4'), + ('f2', 'i4'), + ('f3', [('sf1', 'i4')])]) + # byte string indexing fails gracefully + assert_raises(IndexError, a.__setitem__, b'f1', 1) + assert_raises(IndexError, a.__getitem__, b'f1') + assert_raises(IndexError, a['f1'].__setitem__, b'sf1', 1) + assert_raises(IndexError, a['f1'].__getitem__, b'sf1') + b = a.copy() + fn1 = str('f1') + b[fn1] = 1 + assert_equal(b[fn1], 1) + fnn = str('not at all') + assert_raises(ValueError, b.__setitem__, fnn, 1) + assert_raises(ValueError, b.__getitem__, fnn) + b[0][fn1] = 2 + assert_equal(b[fn1], 2) + # Subfield + assert_raises(ValueError, b[0].__setitem__, fnn, 1) + assert_raises(ValueError, b[0].__getitem__, fnn) + # Subfield + fn3 = str('f3') + sfn1 = str('sf1') + b[fn3][sfn1] = 1 + assert_equal(b[fn3][sfn1], 1) + assert_raises(ValueError, b[fn3].__setitem__, fnn, 1) + assert_raises(ValueError, b[fn3].__getitem__, fnn) + # multiple subfields + fn2 = str('f2') + b[fn2] = 3 + + assert_equal(b[['f1', 'f2']][0].tolist(), (2, 3)) + assert_equal(b[['f2', 'f1']][0].tolist(), (3, 2)) + assert_equal(b[['f1', 'f3']][0].tolist(), (2, (1,))) + + # non-ascii unicode field indexing is well behaved + assert_raises(ValueError, a.__setitem__, '\u03e0', 1) + assert_raises(ValueError, a.__getitem__, '\u03e0') + + def test_record_hash(self): + a = np.array([(1, 2), (1, 2)], dtype='i1,i2') + a.flags.writeable = False + b = np.array([(1, 2), (3, 4)], dtype=[('num1', 'i1'), ('num2', 'i2')]) + b.flags.writeable = False + c = np.array([(1, 2), (3, 4)], dtype='i1,i2') + c.flags.writeable = False + assert_(hash(a[0]) == hash(a[1])) + assert_(hash(a[0]) == hash(b[0])) + assert_(hash(a[0]) != hash(b[1])) + assert_(hash(c[0]) == hash(a[0]) and c[0] == a[0]) + + def test_record_no_hash(self): + a = np.array([(1, 2), (1, 2)], dtype='i1,i2') + assert_raises(TypeError, hash, a[0]) + + def test_empty_structure_creation(self): + # make sure these do not raise errors (gh-5631) + np.array([()], dtype={'names': [], 'formats': [], + 'offsets': [], 'itemsize': 12}) + np.array([(), (), (), (), ()], dtype={'names': [], 'formats': [], + 'offsets': [], 'itemsize': 12}) + + def test_multifield_indexing_view(self): + a = np.ones(3, dtype=[('a', 'i4'), ('b', 'f4'), ('c', 'u4')]) + v = a[['a', 'c']] + assert_(v.base is a) + assert_(v.dtype == np.dtype({'names': ['a', 'c'], + 'formats': ['i4', 'u4'], + 'offsets': [0, 8]})) + v[:] = (4,5) + assert_equal(a[0].item(), (4, 1, 5)) + +class TestView: + def test_basic(self): + x = np.array([(1, 2, 3, 4), (5, 6, 7, 8)], + dtype=[('r', np.int8), ('g', np.int8), + ('b', np.int8), ('a', np.int8)]) + # We must be specific about the endianness here: + y = x.view(dtype=' 0) + assert_(issubclass(w[0].category, RuntimeWarning)) + + def test_empty(self): + A = np.zeros((0, 3)) + for f in self.funcs: + for axis in [0, None]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_(np.isnan(f(A, axis=axis)).all()) + assert_(len(w) > 0) + assert_(issubclass(w[0].category, RuntimeWarning)) + for axis in [1]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_equal(f(A, axis=axis), np.zeros([])) + + def test_mean_values(self): + for mat in [self.rmat, self.cmat, self.omat]: + for axis in [0, 1]: + tgt = mat.sum(axis=axis) + res = _mean(mat, axis=axis) * mat.shape[axis] + assert_almost_equal(res, tgt) + for axis in [None]: + tgt = mat.sum(axis=axis) + res = _mean(mat, axis=axis) * np.prod(mat.shape) + assert_almost_equal(res, tgt) + + def test_mean_float16(self): + # This fail if the sum inside mean is done in float16 instead + # of float32. + assert_(_mean(np.ones(100000, dtype='float16')) == 1) + + def test_mean_axis_error(self): + # Ensure that AxisError is raised instead of IndexError when axis is + # out of bounds, see gh-15817. + with assert_raises(np.exceptions.AxisError): + np.arange(10).mean(axis=2) + + def test_mean_where(self): + a = np.arange(16).reshape((4, 4)) + wh_full = np.array([[False, True, False, True], + [True, False, True, False], + [True, True, False, False], + [False, False, True, True]]) + wh_partial = np.array([[False], + [True], + [True], + [False]]) + _cases = [(1, True, [1.5, 5.5, 9.5, 13.5]), + (0, wh_full, [6., 5., 10., 9.]), + (1, wh_full, [2., 5., 8.5, 14.5]), + (0, wh_partial, [6., 7., 8., 9.])] + for _ax, _wh, _res in _cases: + assert_allclose(a.mean(axis=_ax, where=_wh), + np.array(_res)) + assert_allclose(np.mean(a, axis=_ax, where=_wh), + np.array(_res)) + + a3d = np.arange(16).reshape((2, 2, 4)) + _wh_partial = np.array([False, True, True, False]) + _res = [[1.5, 5.5], [9.5, 13.5]] + assert_allclose(a3d.mean(axis=2, where=_wh_partial), + np.array(_res)) + assert_allclose(np.mean(a3d, axis=2, where=_wh_partial), + np.array(_res)) + + with pytest.warns(RuntimeWarning) as w: + assert_allclose(a.mean(axis=1, where=wh_partial), + np.array([np.nan, 5.5, 9.5, np.nan])) + with pytest.warns(RuntimeWarning) as w: + assert_equal(a.mean(where=False), np.nan) + with pytest.warns(RuntimeWarning) as w: + assert_equal(np.mean(a, where=False), np.nan) + + def test_var_values(self): + for mat in [self.rmat, self.cmat, self.omat]: + for axis in [0, 1, None]: + msqr = _mean(mat * mat.conj(), axis=axis) + mean = _mean(mat, axis=axis) + tgt = msqr - mean * mean.conjugate() + res = _var(mat, axis=axis) + assert_almost_equal(res, tgt) + + @pytest.mark.parametrize(('complex_dtype', 'ndec'), ( + ('complex64', 6), + ('complex128', 7), + ('clongdouble', 7), + )) + def test_var_complex_values(self, complex_dtype, ndec): + # Test fast-paths for every builtin complex type + for axis in [0, 1, None]: + mat = self.cmat.copy().astype(complex_dtype) + msqr = _mean(mat * mat.conj(), axis=axis) + mean = _mean(mat, axis=axis) + tgt = msqr - mean * mean.conjugate() + res = _var(mat, axis=axis) + assert_almost_equal(res, tgt, decimal=ndec) + + def test_var_dimensions(self): + # _var paths for complex number introduce additions on views that + # increase dimensions. Ensure this generalizes to higher dims + mat = np.stack([self.cmat]*3) + for axis in [0, 1, 2, -1, None]: + msqr = _mean(mat * mat.conj(), axis=axis) + mean = _mean(mat, axis=axis) + tgt = msqr - mean * mean.conjugate() + res = _var(mat, axis=axis) + assert_almost_equal(res, tgt) + + def test_var_complex_byteorder(self): + # Test that var fast-path does not cause failures for complex arrays + # with non-native byteorder + cmat = self.cmat.copy().astype('complex128') + cmat_swapped = cmat.astype(cmat.dtype.newbyteorder()) + assert_almost_equal(cmat.var(), cmat_swapped.var()) + + def test_var_axis_error(self): + # Ensure that AxisError is raised instead of IndexError when axis is + # out of bounds, see gh-15817. + with assert_raises(np.exceptions.AxisError): + np.arange(10).var(axis=2) + + def test_var_where(self): + a = np.arange(25).reshape((5, 5)) + wh_full = np.array([[False, True, False, True, True], + [True, False, True, True, False], + [True, True, False, False, True], + [False, True, True, False, True], + [True, False, True, True, False]]) + wh_partial = np.array([[False], + [True], + [True], + [False], + [True]]) + _cases = [(0, True, [50., 50., 50., 50., 50.]), + (1, True, [2., 2., 2., 2., 2.])] + for _ax, _wh, _res in _cases: + assert_allclose(a.var(axis=_ax, where=_wh), + np.array(_res)) + assert_allclose(np.var(a, axis=_ax, where=_wh), + np.array(_res)) + + a3d = np.arange(16).reshape((2, 2, 4)) + _wh_partial = np.array([False, True, True, False]) + _res = [[0.25, 0.25], [0.25, 0.25]] + assert_allclose(a3d.var(axis=2, where=_wh_partial), + np.array(_res)) + assert_allclose(np.var(a3d, axis=2, where=_wh_partial), + np.array(_res)) + + assert_allclose(np.var(a, axis=1, where=wh_full), + np.var(a[wh_full].reshape((5, 3)), axis=1)) + assert_allclose(np.var(a, axis=0, where=wh_partial), + np.var(a[wh_partial[:,0]], axis=0)) + with pytest.warns(RuntimeWarning) as w: + assert_equal(a.var(where=False), np.nan) + with pytest.warns(RuntimeWarning) as w: + assert_equal(np.var(a, where=False), np.nan) + + def test_std_values(self): + for mat in [self.rmat, self.cmat, self.omat]: + for axis in [0, 1, None]: + tgt = np.sqrt(_var(mat, axis=axis)) + res = _std(mat, axis=axis) + assert_almost_equal(res, tgt) + + def test_std_where(self): + a = np.arange(25).reshape((5,5))[::-1] + whf = np.array([[False, True, False, True, True], + [True, False, True, False, True], + [True, True, False, True, False], + [True, False, True, True, False], + [False, True, False, True, True]]) + whp = np.array([[False], + [False], + [True], + [True], + [False]]) + _cases = [ + (0, True, 7.07106781*np.ones(5)), + (1, True, 1.41421356*np.ones(5)), + (0, whf, + np.array([4.0824829 , 8.16496581, 5., 7.39509973, 8.49836586])), + (0, whp, 2.5*np.ones(5)) + ] + for _ax, _wh, _res in _cases: + assert_allclose(a.std(axis=_ax, where=_wh), _res) + assert_allclose(np.std(a, axis=_ax, where=_wh), _res) + + a3d = np.arange(16).reshape((2, 2, 4)) + _wh_partial = np.array([False, True, True, False]) + _res = [[0.5, 0.5], [0.5, 0.5]] + assert_allclose(a3d.std(axis=2, where=_wh_partial), + np.array(_res)) + assert_allclose(np.std(a3d, axis=2, where=_wh_partial), + np.array(_res)) + + assert_allclose(a.std(axis=1, where=whf), + np.std(a[whf].reshape((5,3)), axis=1)) + assert_allclose(np.std(a, axis=1, where=whf), + (a[whf].reshape((5,3))).std(axis=1)) + assert_allclose(a.std(axis=0, where=whp), + np.std(a[whp[:,0]], axis=0)) + assert_allclose(np.std(a, axis=0, where=whp), + (a[whp[:,0]]).std(axis=0)) + with pytest.warns(RuntimeWarning) as w: + assert_equal(a.std(where=False), np.nan) + with pytest.warns(RuntimeWarning) as w: + assert_equal(np.std(a, where=False), np.nan) + + def test_subclass(self): + class TestArray(np.ndarray): + def __new__(cls, data, info): + result = np.array(data) + result = result.view(cls) + result.info = info + return result + + def __array_finalize__(self, obj): + self.info = getattr(obj, "info", '') + + dat = TestArray([[1, 2, 3, 4], [5, 6, 7, 8]], 'jubba') + res = dat.mean(1) + assert_(res.info == dat.info) + res = dat.std(1) + assert_(res.info == dat.info) + res = dat.var(1) + assert_(res.info == dat.info) + + +class TestVdot: + def test_basic(self): + dt_numeric = np.typecodes['AllFloat'] + np.typecodes['AllInteger'] + dt_complex = np.typecodes['Complex'] + + # test real + a = np.eye(3) + for dt in dt_numeric + 'O': + b = a.astype(dt) + res = np.vdot(b, b) + assert_(np.isscalar(res)) + assert_equal(np.vdot(b, b), 3) + + # test complex + a = np.eye(3) * 1j + for dt in dt_complex + 'O': + b = a.astype(dt) + res = np.vdot(b, b) + assert_(np.isscalar(res)) + assert_equal(np.vdot(b, b), 3) + + # test boolean + b = np.eye(3, dtype=bool) + res = np.vdot(b, b) + assert_(np.isscalar(res)) + assert_equal(np.vdot(b, b), True) + + def test_vdot_array_order(self): + a = np.array([[1, 2], [3, 4]], order='C') + b = np.array([[1, 2], [3, 4]], order='F') + res = np.vdot(a, a) + + # integer arrays are exact + assert_equal(np.vdot(a, b), res) + assert_equal(np.vdot(b, a), res) + assert_equal(np.vdot(b, b), res) + + def test_vdot_uncontiguous(self): + for size in [2, 1000]: + # Different sizes match different branches in vdot. + a = np.zeros((size, 2, 2)) + b = np.zeros((size, 2, 2)) + a[:, 0, 0] = np.arange(size) + b[:, 0, 0] = np.arange(size) + 1 + # Make a and b uncontiguous: + a = a[..., 0] + b = b[..., 0] + + assert_equal(np.vdot(a, b), + np.vdot(a.flatten(), b.flatten())) + assert_equal(np.vdot(a, b.copy()), + np.vdot(a.flatten(), b.flatten())) + assert_equal(np.vdot(a.copy(), b), + np.vdot(a.flatten(), b.flatten())) + assert_equal(np.vdot(a.copy('F'), b), + np.vdot(a.flatten(), b.flatten())) + assert_equal(np.vdot(a, b.copy('F')), + np.vdot(a.flatten(), b.flatten())) + + +class TestDot: + def setup_method(self): + np.random.seed(128) + self.A = np.random.rand(4, 2) + self.b1 = np.random.rand(2, 1) + self.b2 = np.random.rand(2) + self.b3 = np.random.rand(1, 2) + self.b4 = np.random.rand(4) + self.N = 7 + + def test_dotmatmat(self): + A = self.A + res = np.dot(A.transpose(), A) + tgt = np.array([[1.45046013, 0.86323640], + [0.86323640, 0.84934569]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotmatvec(self): + A, b1 = self.A, self.b1 + res = np.dot(A, b1) + tgt = np.array([[0.32114320], [0.04889721], + [0.15696029], [0.33612621]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotmatvec2(self): + A, b2 = self.A, self.b2 + res = np.dot(A, b2) + tgt = np.array([0.29677940, 0.04518649, 0.14468333, 0.31039293]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecmat(self): + A, b4 = self.A, self.b4 + res = np.dot(b4, A) + tgt = np.array([1.23495091, 1.12222648]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecmat2(self): + b3, A = self.b3, self.A + res = np.dot(b3, A.transpose()) + tgt = np.array([[0.58793804, 0.08957460, 0.30605758, 0.62716383]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecmat3(self): + A, b4 = self.A, self.b4 + res = np.dot(A.transpose(), b4) + tgt = np.array([1.23495091, 1.12222648]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecvecouter(self): + b1, b3 = self.b1, self.b3 + res = np.dot(b1, b3) + tgt = np.array([[0.20128610, 0.08400440], [0.07190947, 0.03001058]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecvecinner(self): + b1, b3 = self.b1, self.b3 + res = np.dot(b3, b1) + tgt = np.array([[ 0.23129668]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotcolumnvect1(self): + b1 = np.ones((3, 1)) + b2 = [5.3] + res = np.dot(b1, b2) + tgt = np.array([5.3, 5.3, 5.3]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotcolumnvect2(self): + b1 = np.ones((3, 1)).transpose() + b2 = [6.2] + res = np.dot(b2, b1) + tgt = np.array([6.2, 6.2, 6.2]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecscalar(self): + np.random.seed(100) + b1 = np.random.rand(1, 1) + b2 = np.random.rand(1, 4) + res = np.dot(b1, b2) + tgt = np.array([[0.15126730, 0.23068496, 0.45905553, 0.00256425]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_dotvecscalar2(self): + np.random.seed(100) + b1 = np.random.rand(4, 1) + b2 = np.random.rand(1, 1) + res = np.dot(b1, b2) + tgt = np.array([[0.00256425],[0.00131359],[0.00200324],[ 0.00398638]]) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_all(self): + dims = [(), (1,), (1, 1)] + dout = [(), (1,), (1, 1), (1,), (), (1,), (1, 1), (1,), (1, 1)] + for dim, (dim1, dim2) in zip(dout, itertools.product(dims, dims)): + b1 = np.zeros(dim1) + b2 = np.zeros(dim2) + res = np.dot(b1, b2) + tgt = np.zeros(dim) + assert_(res.shape == tgt.shape) + assert_almost_equal(res, tgt, decimal=self.N) + + def test_vecobject(self): + class Vec: + def __init__(self, sequence=None): + if sequence is None: + sequence = [] + self.array = np.array(sequence) + + def __add__(self, other): + out = Vec() + out.array = self.array + other.array + return out + + def __sub__(self, other): + out = Vec() + out.array = self.array - other.array + return out + + def __mul__(self, other): # with scalar + out = Vec(self.array.copy()) + out.array *= other + return out + + def __rmul__(self, other): + return self*other + + U_non_cont = np.transpose([[1., 1.], [1., 2.]]) + U_cont = np.ascontiguousarray(U_non_cont) + x = np.array([Vec([1., 0.]), Vec([0., 1.])]) + zeros = np.array([Vec([0., 0.]), Vec([0., 0.])]) + zeros_test = np.dot(U_cont, x) - np.dot(U_non_cont, x) + assert_equal(zeros[0].array, zeros_test[0].array) + assert_equal(zeros[1].array, zeros_test[1].array) + + def test_dot_2args(self): + + a = np.array([[1, 2], [3, 4]], dtype=float) + b = np.array([[1, 0], [1, 1]], dtype=float) + c = np.array([[3, 2], [7, 4]], dtype=float) + + d = dot(a, b) + assert_allclose(c, d) + + def test_dot_3args(self): + + np.random.seed(22) + f = np.random.random_sample((1024, 16)) + v = np.random.random_sample((16, 32)) + + r = np.empty((1024, 32)) + for i in range(12): + dot(f, v, r) + if HAS_REFCOUNT: + assert_equal(sys.getrefcount(r), 2) + r2 = dot(f, v, out=None) + assert_array_equal(r2, r) + assert_(r is dot(f, v, out=r)) + + v = v[:, 0].copy() # v.shape == (16,) + r = r[:, 0].copy() # r.shape == (1024,) + r2 = dot(f, v) + assert_(r is dot(f, v, r)) + assert_array_equal(r2, r) + + def test_dot_3args_errors(self): + + np.random.seed(22) + f = np.random.random_sample((1024, 16)) + v = np.random.random_sample((16, 32)) + + r = np.empty((1024, 31)) + assert_raises(ValueError, dot, f, v, r) + + r = np.empty((1024,)) + assert_raises(ValueError, dot, f, v, r) + + r = np.empty((32,)) + assert_raises(ValueError, dot, f, v, r) + + r = np.empty((32, 1024)) + assert_raises(ValueError, dot, f, v, r) + assert_raises(ValueError, dot, f, v, r.T) + + r = np.empty((1024, 64)) + assert_raises(ValueError, dot, f, v, r[:, ::2]) + assert_raises(ValueError, dot, f, v, r[:, :32]) + + r = np.empty((1024, 32), dtype=np.float32) + assert_raises(ValueError, dot, f, v, r) + + r = np.empty((1024, 32), dtype=int) + assert_raises(ValueError, dot, f, v, r) + + def test_dot_out_result(self): + x = np.ones((), dtype=np.float16) + y = np.ones((5,), dtype=np.float16) + z = np.zeros((5,), dtype=np.float16) + res = x.dot(y, out=z) + assert np.array_equal(res, y) + assert np.array_equal(z, y) + + def test_dot_out_aliasing(self): + x = np.ones((), dtype=np.float16) + y = np.ones((5,), dtype=np.float16) + z = np.zeros((5,), dtype=np.float16) + res = x.dot(y, out=z) + z[0] = 2 + assert np.array_equal(res, z) + + def test_dot_array_order(self): + a = np.array([[1, 2], [3, 4]], order='C') + b = np.array([[1, 2], [3, 4]], order='F') + res = np.dot(a, a) + + # integer arrays are exact + assert_equal(np.dot(a, b), res) + assert_equal(np.dot(b, a), res) + assert_equal(np.dot(b, b), res) + + def test_accelerate_framework_sgemv_fix(self): + + def aligned_array(shape, align, dtype, order='C'): + d = dtype(0) + N = np.prod(shape) + tmp = np.zeros(N * d.nbytes + align, dtype=np.uint8) + address = tmp.__array_interface__["data"][0] + for offset in range(align): + if (address + offset) % align == 0: + break + tmp = tmp[offset:offset+N*d.nbytes].view(dtype=dtype) + return tmp.reshape(shape, order=order) + + def as_aligned(arr, align, dtype, order='C'): + aligned = aligned_array(arr.shape, align, dtype, order) + aligned[:] = arr[:] + return aligned + + def assert_dot_close(A, X, desired): + assert_allclose(np.dot(A, X), desired, rtol=1e-5, atol=1e-7) + + m = aligned_array(100, 15, np.float32) + s = aligned_array((100, 100), 15, np.float32) + np.dot(s, m) # this will always segfault if the bug is present + + testdata = itertools.product((15, 32), (10000,), (200, 89), ('C', 'F')) + for align, m, n, a_order in testdata: + # Calculation in double precision + A_d = np.random.rand(m, n) + X_d = np.random.rand(n) + desired = np.dot(A_d, X_d) + # Calculation with aligned single precision + A_f = as_aligned(A_d, align, np.float32, order=a_order) + X_f = as_aligned(X_d, align, np.float32) + assert_dot_close(A_f, X_f, desired) + # Strided A rows + A_d_2 = A_d[::2] + desired = np.dot(A_d_2, X_d) + A_f_2 = A_f[::2] + assert_dot_close(A_f_2, X_f, desired) + # Strided A columns, strided X vector + A_d_22 = A_d_2[:, ::2] + X_d_2 = X_d[::2] + desired = np.dot(A_d_22, X_d_2) + A_f_22 = A_f_2[:, ::2] + X_f_2 = X_f[::2] + assert_dot_close(A_f_22, X_f_2, desired) + # Check the strides are as expected + if a_order == 'F': + assert_equal(A_f_22.strides, (8, 8 * m)) + else: + assert_equal(A_f_22.strides, (8 * n, 8)) + assert_equal(X_f_2.strides, (8,)) + # Strides in A rows + cols only + X_f_2c = as_aligned(X_f_2, align, np.float32) + assert_dot_close(A_f_22, X_f_2c, desired) + # Strides just in A cols + A_d_12 = A_d[:, ::2] + desired = np.dot(A_d_12, X_d_2) + A_f_12 = A_f[:, ::2] + assert_dot_close(A_f_12, X_f_2c, desired) + # Strides in A cols and X + assert_dot_close(A_f_12, X_f_2, desired) + + @pytest.mark.slow + @pytest.mark.parametrize("dtype", [np.float64, np.complex128]) + @requires_memory(free_bytes=18e9) # complex case needs 18GiB+ + def test_huge_vectordot(self, dtype): + # Large vector multiplications are chunked with 32bit BLAS + # Test that the chunking does the right thing, see also gh-22262 + data = np.ones(2**30+100, dtype=dtype) + res = np.dot(data, data) + assert res == 2**30+100 + + def test_dtype_discovery_fails(self): + # See gh-14247, error checking was missing for failed dtype discovery + class BadObject: + def __array__(self, dtype=None, copy=None): + raise TypeError("just this tiny mint leaf") + + with pytest.raises(TypeError): + np.dot(BadObject(), BadObject()) + + with pytest.raises(TypeError): + np.dot(3.0, BadObject()) + + +class MatmulCommon: + """Common tests for '@' operator and numpy.matmul. + + """ + # Should work with these types. Will want to add + # "O" at some point + types = "?bhilqBHILQefdgFDGO" + + def test_exceptions(self): + dims = [ + ((1,), (2,)), # mismatched vector vector + ((2, 1,), (2,)), # mismatched matrix vector + ((2,), (1, 2)), # mismatched vector matrix + ((1, 2), (3, 1)), # mismatched matrix matrix + ((1,), ()), # vector scalar + ((), (1)), # scalar vector + ((1, 1), ()), # matrix scalar + ((), (1, 1)), # scalar matrix + ((2, 2, 1), (3, 1, 2)), # cannot broadcast + ] + + for dt, (dm1, dm2) in itertools.product(self.types, dims): + a = np.ones(dm1, dtype=dt) + b = np.ones(dm2, dtype=dt) + assert_raises(ValueError, self.matmul, a, b) + + def test_shapes(self): + dims = [ + ((1, 1), (2, 1, 1)), # broadcast first argument + ((2, 1, 1), (1, 1)), # broadcast second argument + ((2, 1, 1), (2, 1, 1)), # matrix stack sizes match + ] + + for dt, (dm1, dm2) in itertools.product(self.types, dims): + a = np.ones(dm1, dtype=dt) + b = np.ones(dm2, dtype=dt) + res = self.matmul(a, b) + assert_(res.shape == (2, 1, 1)) + + # vector vector returns scalars. + for dt in self.types: + a = np.ones((2,), dtype=dt) + b = np.ones((2,), dtype=dt) + c = self.matmul(a, b) + assert_(np.array(c).shape == ()) + + def test_result_types(self): + mat = np.ones((1,1)) + vec = np.ones((1,)) + for dt in self.types: + m = mat.astype(dt) + v = vec.astype(dt) + for arg in [(m, v), (v, m), (m, m)]: + res = self.matmul(*arg) + assert_(res.dtype == dt) + + # vector vector returns scalars + if dt != "O": + res = self.matmul(v, v) + assert_(type(res) is np.dtype(dt).type) + + def test_scalar_output(self): + vec1 = np.array([2]) + vec2 = np.array([3, 4]).reshape(1, -1) + tgt = np.array([6, 8]) + for dt in self.types[1:]: + v1 = vec1.astype(dt) + v2 = vec2.astype(dt) + res = self.matmul(v1, v2) + assert_equal(res, tgt) + res = self.matmul(v2.T, v1) + assert_equal(res, tgt) + + # boolean type + vec = np.array([True, True], dtype='?').reshape(1, -1) + res = self.matmul(vec[:, 0], vec) + assert_equal(res, True) + + def test_vector_vector_values(self): + vec1 = np.array([1, 2]) + vec2 = np.array([3, 4]).reshape(-1, 1) + tgt1 = np.array([11]) + tgt2 = np.array([[3, 6], [4, 8]]) + for dt in self.types[1:]: + v1 = vec1.astype(dt) + v2 = vec2.astype(dt) + res = self.matmul(v1, v2) + assert_equal(res, tgt1) + # no broadcast, we must make v1 into a 2d ndarray + res = self.matmul(v2, v1.reshape(1, -1)) + assert_equal(res, tgt2) + + # boolean type + vec = np.array([True, True], dtype='?') + res = self.matmul(vec, vec) + assert_equal(res, True) + + def test_vector_matrix_values(self): + vec = np.array([1, 2]) + mat1 = np.array([[1, 2], [3, 4]]) + mat2 = np.stack([mat1]*2, axis=0) + tgt1 = np.array([7, 10]) + tgt2 = np.stack([tgt1]*2, axis=0) + for dt in self.types[1:]: + v = vec.astype(dt) + m1 = mat1.astype(dt) + m2 = mat2.astype(dt) + res = self.matmul(v, m1) + assert_equal(res, tgt1) + res = self.matmul(v, m2) + assert_equal(res, tgt2) + + # boolean type + vec = np.array([True, False]) + mat1 = np.array([[True, False], [False, True]]) + mat2 = np.stack([mat1]*2, axis=0) + tgt1 = np.array([True, False]) + tgt2 = np.stack([tgt1]*2, axis=0) + + res = self.matmul(vec, mat1) + assert_equal(res, tgt1) + res = self.matmul(vec, mat2) + assert_equal(res, tgt2) + + def test_matrix_vector_values(self): + vec = np.array([1, 2]) + mat1 = np.array([[1, 2], [3, 4]]) + mat2 = np.stack([mat1]*2, axis=0) + tgt1 = np.array([5, 11]) + tgt2 = np.stack([tgt1]*2, axis=0) + for dt in self.types[1:]: + v = vec.astype(dt) + m1 = mat1.astype(dt) + m2 = mat2.astype(dt) + res = self.matmul(m1, v) + assert_equal(res, tgt1) + res = self.matmul(m2, v) + assert_equal(res, tgt2) + + # boolean type + vec = np.array([True, False]) + mat1 = np.array([[True, False], [False, True]]) + mat2 = np.stack([mat1]*2, axis=0) + tgt1 = np.array([True, False]) + tgt2 = np.stack([tgt1]*2, axis=0) + + res = self.matmul(vec, mat1) + assert_equal(res, tgt1) + res = self.matmul(vec, mat2) + assert_equal(res, tgt2) + + def test_matrix_matrix_values(self): + mat1 = np.array([[1, 2], [3, 4]]) + mat2 = np.array([[1, 0], [1, 1]]) + mat12 = np.stack([mat1, mat2], axis=0) + mat21 = np.stack([mat2, mat1], axis=0) + tgt11 = np.array([[7, 10], [15, 22]]) + tgt12 = np.array([[3, 2], [7, 4]]) + tgt21 = np.array([[1, 2], [4, 6]]) + tgt12_21 = np.stack([tgt12, tgt21], axis=0) + tgt11_12 = np.stack((tgt11, tgt12), axis=0) + tgt11_21 = np.stack((tgt11, tgt21), axis=0) + for dt in self.types[1:]: + m1 = mat1.astype(dt) + m2 = mat2.astype(dt) + m12 = mat12.astype(dt) + m21 = mat21.astype(dt) + + # matrix @ matrix + res = self.matmul(m1, m2) + assert_equal(res, tgt12) + res = self.matmul(m2, m1) + assert_equal(res, tgt21) + + # stacked @ matrix + res = self.matmul(m12, m1) + assert_equal(res, tgt11_21) + + # matrix @ stacked + res = self.matmul(m1, m12) + assert_equal(res, tgt11_12) + + # stacked @ stacked + res = self.matmul(m12, m21) + assert_equal(res, tgt12_21) + + # boolean type + m1 = np.array([[1, 1], [0, 0]], dtype=np.bool) + m2 = np.array([[1, 0], [1, 1]], dtype=np.bool) + m12 = np.stack([m1, m2], axis=0) + m21 = np.stack([m2, m1], axis=0) + tgt11 = m1 + tgt12 = m1 + tgt21 = np.array([[1, 1], [1, 1]], dtype=np.bool) + tgt12_21 = np.stack([tgt12, tgt21], axis=0) + tgt11_12 = np.stack((tgt11, tgt12), axis=0) + tgt11_21 = np.stack((tgt11, tgt21), axis=0) + + # matrix @ matrix + res = self.matmul(m1, m2) + assert_equal(res, tgt12) + res = self.matmul(m2, m1) + assert_equal(res, tgt21) + + # stacked @ matrix + res = self.matmul(m12, m1) + assert_equal(res, tgt11_21) + + # matrix @ stacked + res = self.matmul(m1, m12) + assert_equal(res, tgt11_12) + + # stacked @ stacked + res = self.matmul(m12, m21) + assert_equal(res, tgt12_21) + + +class TestMatmul(MatmulCommon): + matmul = np.matmul + + def test_out_arg(self): + a = np.ones((5, 2), dtype=float) + b = np.array([[1, 3], [5, 7]], dtype=float) + tgt = np.dot(a, b) + + # test as positional argument + msg = "out positional argument" + out = np.zeros((5, 2), dtype=float) + self.matmul(a, b, out) + assert_array_equal(out, tgt, err_msg=msg) + + # test as keyword argument + msg = "out keyword argument" + out = np.zeros((5, 2), dtype=float) + self.matmul(a, b, out=out) + assert_array_equal(out, tgt, err_msg=msg) + + # test out with not allowed type cast (safe casting) + msg = "Cannot cast ufunc .* output" + out = np.zeros((5, 2), dtype=np.int32) + assert_raises_regex(TypeError, msg, self.matmul, a, b, out=out) + + # test out with type upcast to complex + out = np.zeros((5, 2), dtype=np.complex128) + c = self.matmul(a, b, out=out) + assert_(c is out) + with suppress_warnings() as sup: + sup.filter(ComplexWarning, '') + c = c.astype(tgt.dtype) + assert_array_equal(c, tgt) + + def test_empty_out(self): + # Check that the output cannot be broadcast, so that it cannot be + # size zero when the outer dimensions (iterator size) has size zero. + arr = np.ones((0, 1, 1)) + out = np.ones((1, 1, 1)) + assert self.matmul(arr, arr).shape == (0, 1, 1) + + with pytest.raises(ValueError, match=r"non-broadcastable"): + self.matmul(arr, arr, out=out) + + def test_out_contiguous(self): + a = np.ones((5, 2), dtype=float) + b = np.array([[1, 3], [5, 7]], dtype=float) + v = np.array([1, 3], dtype=float) + tgt = np.dot(a, b) + tgt_mv = np.dot(a, v) + + # test out non-contiguous + out = np.ones((5, 2, 2), dtype=float) + c = self.matmul(a, b, out=out[..., 0]) + assert c.base is out + assert_array_equal(c, tgt) + c = self.matmul(a, v, out=out[:, 0, 0]) + assert_array_equal(c, tgt_mv) + c = self.matmul(v, a.T, out=out[:, 0, 0]) + assert_array_equal(c, tgt_mv) + + # test out contiguous in only last dim + out = np.ones((10, 2), dtype=float) + c = self.matmul(a, b, out=out[::2, :]) + assert_array_equal(c, tgt) + + # test transposes of out, args + out = np.ones((5, 2), dtype=float) + c = self.matmul(b.T, a.T, out=out.T) + assert_array_equal(out, tgt) + + m1 = np.arange(15.).reshape(5, 3) + m2 = np.arange(21.).reshape(3, 7) + m3 = np.arange(30.).reshape(5, 6)[:, ::2] # non-contiguous + vc = np.arange(10.) + vr = np.arange(6.) + m0 = np.zeros((3, 0)) + @pytest.mark.parametrize('args', ( + # matrix-matrix + (m1, m2), (m2.T, m1.T), (m2.T.copy(), m1.T), (m2.T, m1.T.copy()), + # matrix-matrix-transpose, contiguous and non + (m1, m1.T), (m1.T, m1), (m1, m3.T), (m3, m1.T), + (m3, m3.T), (m3.T, m3), + # matrix-matrix non-contiguous + (m3, m2), (m2.T, m3.T), (m2.T.copy(), m3.T), + # vector-matrix, matrix-vector, contiguous + (m1, vr[:3]), (vc[:5], m1), (m1.T, vc[:5]), (vr[:3], m1.T), + # vector-matrix, matrix-vector, vector non-contiguous + (m1, vr[::2]), (vc[::2], m1), (m1.T, vc[::2]), (vr[::2], m1.T), + # vector-matrix, matrix-vector, matrix non-contiguous + (m3, vr[:3]), (vc[:5], m3), (m3.T, vc[:5]), (vr[:3], m3.T), + # vector-matrix, matrix-vector, both non-contiguous + (m3, vr[::2]), (vc[::2], m3), (m3.T, vc[::2]), (vr[::2], m3.T), + # size == 0 + (m0, m0.T), (m0.T, m0), (m1, m0), (m0.T, m1.T), + )) + def test_dot_equivalent(self, args): + r1 = np.matmul(*args) + r2 = np.dot(*args) + assert_equal(r1, r2) + + r3 = np.matmul(args[0].copy(), args[1].copy()) + assert_equal(r1, r3) + + def test_matmul_object(self): + import fractions + + f = np.vectorize(fractions.Fraction) + def random_ints(): + return np.random.randint(1, 1000, size=(10, 3, 3)) + M1 = f(random_ints(), random_ints()) + M2 = f(random_ints(), random_ints()) + + M3 = self.matmul(M1, M2) + + [N1, N2, N3] = [a.astype(float) for a in [M1, M2, M3]] + + assert_allclose(N3, self.matmul(N1, N2)) + + def test_matmul_object_type_scalar(self): + from fractions import Fraction as F + v = np.array([F(2,3), F(5,7)]) + res = self.matmul(v, v) + assert_(type(res) is F) + + def test_matmul_empty(self): + a = np.empty((3, 0), dtype=object) + b = np.empty((0, 3), dtype=object) + c = np.zeros((3, 3)) + assert_array_equal(np.matmul(a, b), c) + + def test_matmul_exception_multiply(self): + # test that matmul fails if `__mul__` is missing + class add_not_multiply: + def __add__(self, other): + return self + a = np.full((3,3), add_not_multiply()) + with assert_raises(TypeError): + b = np.matmul(a, a) + + def test_matmul_exception_add(self): + # test that matmul fails if `__add__` is missing + class multiply_not_add: + def __mul__(self, other): + return self + a = np.full((3,3), multiply_not_add()) + with assert_raises(TypeError): + b = np.matmul(a, a) + + def test_matmul_bool(self): + # gh-14439 + a = np.array([[1, 0],[1, 1]], dtype=bool) + assert np.max(a.view(np.uint8)) == 1 + b = np.matmul(a, a) + # matmul with boolean output should always be 0, 1 + assert np.max(b.view(np.uint8)) == 1 + + rg = np.random.default_rng(np.random.PCG64(43)) + d = rg.integers(2, size=4*5, dtype=np.int8) + d = d.reshape(4, 5) > 0 + out1 = np.matmul(d, d.reshape(5, 4)) + out2 = np.dot(d, d.reshape(5, 4)) + assert_equal(out1, out2) + + c = np.matmul(np.zeros((2, 0), dtype=bool), np.zeros(0, dtype=bool)) + assert not np.any(c) + + +class TestMatmulOperator(MatmulCommon): + import operator + matmul = operator.matmul + + def test_array_priority_override(self): + + class A: + __array_priority__ = 1000 + + def __matmul__(self, other): + return "A" + + def __rmatmul__(self, other): + return "A" + + a = A() + b = np.ones(2) + assert_equal(self.matmul(a, b), "A") + assert_equal(self.matmul(b, a), "A") + + def test_matmul_raises(self): + assert_raises(TypeError, self.matmul, np.int8(5), np.int8(5)) + assert_raises(TypeError, self.matmul, np.void(b'abc'), np.void(b'abc')) + assert_raises(TypeError, self.matmul, np.arange(10), np.void(b'abc')) + + +class TestMatmulInplace: + DTYPES = {} + for i in MatmulCommon.types: + for j in MatmulCommon.types: + if np.can_cast(j, i): + DTYPES[f"{i}-{j}"] = (np.dtype(i), np.dtype(j)) + + @pytest.mark.parametrize("dtype1,dtype2", DTYPES.values(), ids=DTYPES) + def test_basic(self, dtype1: np.dtype, dtype2: np.dtype) -> None: + a = np.arange(10).reshape(5, 2).astype(dtype1) + a_id = id(a) + b = np.ones((2, 2), dtype=dtype2) + + ref = a @ b + a @= b + + assert id(a) == a_id + assert a.dtype == dtype1 + assert a.shape == (5, 2) + if dtype1.kind in "fc": + np.testing.assert_allclose(a, ref) + else: + np.testing.assert_array_equal(a, ref) + + SHAPES = { + "2d_large": ((10**5, 10), (10, 10)), + "3d_large": ((10**4, 10, 10), (1, 10, 10)), + "1d": ((3,), (3,)), + "2d_1d": ((3, 3), (3,)), + "1d_2d": ((3,), (3, 3)), + "2d_broadcast": ((3, 3), (3, 1)), + "2d_broadcast_reverse": ((1, 3), (3, 3)), + "3d_broadcast1": ((3, 3, 3), (1, 3, 1)), + "3d_broadcast2": ((3, 3, 3), (1, 3, 3)), + "3d_broadcast3": ((3, 3, 3), (3, 3, 1)), + "3d_broadcast_reverse1": ((1, 3, 3), (3, 3, 3)), + "3d_broadcast_reverse2": ((3, 1, 3), (3, 3, 3)), + "3d_broadcast_reverse3": ((1, 1, 3), (3, 3, 3)), + } + + @pytest.mark.parametrize("a_shape,b_shape", SHAPES.values(), ids=SHAPES) + def test_shapes(self, a_shape: tuple[int, ...], b_shape: tuple[int, ...]): + a_size = np.prod(a_shape) + a = np.arange(a_size).reshape(a_shape).astype(np.float64) + a_id = id(a) + + b_size = np.prod(b_shape) + b = np.arange(b_size).reshape(b_shape) + + ref = a @ b + if ref.shape != a_shape: + with pytest.raises(ValueError): + a @= b + return + else: + a @= b + + assert id(a) == a_id + assert a.dtype.type == np.float64 + assert a.shape == a_shape + np.testing.assert_allclose(a, ref) + + +def test_matmul_axes(): + a = np.arange(3*4*5).reshape(3, 4, 5) + c = np.matmul(a, a, axes=[(-2, -1), (-1, -2), (1, 2)]) + assert c.shape == (3, 4, 4) + d = np.matmul(a, a, axes=[(-2, -1), (-1, -2), (0, 1)]) + assert d.shape == (4, 4, 3) + e = np.swapaxes(d, 0, 2) + assert_array_equal(e, c) + f = np.matmul(a, np.arange(3), axes=[(1, 0), (0), (0)]) + assert f.shape == (4, 5) + + +class TestInner: + + def test_inner_type_mismatch(self): + c = 1. + A = np.array((1,1), dtype='i,i') + + assert_raises(TypeError, np.inner, c, A) + assert_raises(TypeError, np.inner, A, c) + + def test_inner_scalar_and_vector(self): + for dt in np.typecodes['AllInteger'] + np.typecodes['AllFloat'] + '?': + sca = np.array(3, dtype=dt)[()] + vec = np.array([1, 2], dtype=dt) + desired = np.array([3, 6], dtype=dt) + assert_equal(np.inner(vec, sca), desired) + assert_equal(np.inner(sca, vec), desired) + + def test_vecself(self): + # Ticket 844. + # Inner product of a vector with itself segfaults or give + # meaningless result + a = np.zeros(shape=(1, 80), dtype=np.float64) + p = np.inner(a, a) + assert_almost_equal(p, 0, decimal=14) + + def test_inner_product_with_various_contiguities(self): + # github issue 6532 + for dt in np.typecodes['AllInteger'] + np.typecodes['AllFloat'] + '?': + # check an inner product involving a matrix transpose + A = np.array([[1, 2], [3, 4]], dtype=dt) + B = np.array([[1, 3], [2, 4]], dtype=dt) + C = np.array([1, 1], dtype=dt) + desired = np.array([4, 6], dtype=dt) + assert_equal(np.inner(A.T, C), desired) + assert_equal(np.inner(C, A.T), desired) + assert_equal(np.inner(B, C), desired) + assert_equal(np.inner(C, B), desired) + # check a matrix product + desired = np.array([[7, 10], [15, 22]], dtype=dt) + assert_equal(np.inner(A, B), desired) + # check the syrk vs. gemm paths + desired = np.array([[5, 11], [11, 25]], dtype=dt) + assert_equal(np.inner(A, A), desired) + assert_equal(np.inner(A, A.copy()), desired) + # check an inner product involving an aliased and reversed view + a = np.arange(5).astype(dt) + b = a[::-1] + desired = np.array(10, dtype=dt).item() + assert_equal(np.inner(b, a), desired) + + def test_3d_tensor(self): + for dt in np.typecodes['AllInteger'] + np.typecodes['AllFloat'] + '?': + a = np.arange(24).reshape(2,3,4).astype(dt) + b = np.arange(24, 48).reshape(2,3,4).astype(dt) + desired = np.array( + [[[[ 158, 182, 206], + [ 230, 254, 278]], + + [[ 566, 654, 742], + [ 830, 918, 1006]], + + [[ 974, 1126, 1278], + [1430, 1582, 1734]]], + + [[[1382, 1598, 1814], + [2030, 2246, 2462]], + + [[1790, 2070, 2350], + [2630, 2910, 3190]], + + [[2198, 2542, 2886], + [3230, 3574, 3918]]]] + ).astype(dt) + assert_equal(np.inner(a, b), desired) + assert_equal(np.inner(b, a).transpose(2,3,0,1), desired) + + +class TestChoose: + def setup_method(self): + self.x = 2*np.ones((3,), dtype=int) + self.y = 3*np.ones((3,), dtype=int) + self.x2 = 2*np.ones((2, 3), dtype=int) + self.y2 = 3*np.ones((2, 3), dtype=int) + self.ind = [0, 0, 1] + + def test_basic(self): + A = np.choose(self.ind, (self.x, self.y)) + assert_equal(A, [2, 2, 3]) + + def test_broadcast1(self): + A = np.choose(self.ind, (self.x2, self.y2)) + assert_equal(A, [[2, 2, 3], [2, 2, 3]]) + + def test_broadcast2(self): + A = np.choose(self.ind, (self.x, self.y2)) + assert_equal(A, [[2, 2, 3], [2, 2, 3]]) + + @pytest.mark.parametrize("ops", + [(1000, np.array([1], dtype=np.uint8)), + (-1, np.array([1], dtype=np.uint8)), + (1., np.float32(3)), + (1., np.array([3], dtype=np.float32))],) + def test_output_dtype(self, ops): + expected_dt = np.result_type(*ops) + assert(np.choose([0], ops).dtype == expected_dt) + + def test_dimension_and_args_limit(self): + # Maxdims for the legacy iterator is 32, but the maximum number + # of arguments is actually larger (a itself also counts here) + a = np.ones((1,) * 32, dtype=np.intp) + res = a.choose([0, a] + [2] * 61) + with pytest.raises(ValueError, + match="Need at least 0 and at most 64 array objects"): + a.choose([0, a] + [2] * 62) + + assert_array_equal(res, a) + # Choose is unfortunately limited to 32 dims as of NumPy 2.0 + a = np.ones((1,) * 60, dtype=np.intp) + with pytest.raises(RuntimeError, + match=".*32 dimensions but the array has 60"): + a.choose([a, a]) + + +class TestRepeat: + def setup_method(self): + self.m = np.array([1, 2, 3, 4, 5, 6]) + self.m_rect = self.m.reshape((2, 3)) + + def test_basic(self): + A = np.repeat(self.m, [1, 3, 2, 1, 1, 2]) + assert_equal(A, [1, 2, 2, 2, 3, + 3, 4, 5, 6, 6]) + + def test_broadcast1(self): + A = np.repeat(self.m, 2) + assert_equal(A, [1, 1, 2, 2, 3, 3, + 4, 4, 5, 5, 6, 6]) + + def test_axis_spec(self): + A = np.repeat(self.m_rect, [2, 1], axis=0) + assert_equal(A, [[1, 2, 3], + [1, 2, 3], + [4, 5, 6]]) + + A = np.repeat(self.m_rect, [1, 3, 2], axis=1) + assert_equal(A, [[1, 2, 2, 2, 3, 3], + [4, 5, 5, 5, 6, 6]]) + + def test_broadcast2(self): + A = np.repeat(self.m_rect, 2, axis=0) + assert_equal(A, [[1, 2, 3], + [1, 2, 3], + [4, 5, 6], + [4, 5, 6]]) + + A = np.repeat(self.m_rect, 2, axis=1) + assert_equal(A, [[1, 1, 2, 2, 3, 3], + [4, 4, 5, 5, 6, 6]]) + + +# TODO: test for multidimensional +NEIGH_MODE = {'zero': 0, 'one': 1, 'constant': 2, 'circular': 3, 'mirror': 4} + + +@pytest.mark.parametrize('dt', [float, Decimal], ids=['float', 'object']) +class TestNeighborhoodIter: + # Simple, 2d tests + def test_simple2d(self, dt): + # Test zero and one padding for simple data type + x = np.array([[0, 1], [2, 3]], dtype=dt) + r = [np.array([[0, 0, 0], [0, 0, 1]], dtype=dt), + np.array([[0, 0, 0], [0, 1, 0]], dtype=dt), + np.array([[0, 0, 1], [0, 2, 3]], dtype=dt), + np.array([[0, 1, 0], [2, 3, 0]], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 0, -1, 1], x[0], NEIGH_MODE['zero']) + assert_array_equal(l, r) + + r = [np.array([[1, 1, 1], [1, 0, 1]], dtype=dt), + np.array([[1, 1, 1], [0, 1, 1]], dtype=dt), + np.array([[1, 0, 1], [1, 2, 3]], dtype=dt), + np.array([[0, 1, 1], [2, 3, 1]], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 0, -1, 1], x[0], NEIGH_MODE['one']) + assert_array_equal(l, r) + + r = [np.array([[4, 4, 4], [4, 0, 1]], dtype=dt), + np.array([[4, 4, 4], [0, 1, 4]], dtype=dt), + np.array([[4, 0, 1], [4, 2, 3]], dtype=dt), + np.array([[0, 1, 4], [2, 3, 4]], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 0, -1, 1], 4, NEIGH_MODE['constant']) + assert_array_equal(l, r) + + # Test with start in the middle + r = [np.array([[4, 0, 1], [4, 2, 3]], dtype=dt), + np.array([[0, 1, 4], [2, 3, 4]], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 0, -1, 1], 4, NEIGH_MODE['constant'], 2) + assert_array_equal(l, r) + + def test_mirror2d(self, dt): + x = np.array([[0, 1], [2, 3]], dtype=dt) + r = [np.array([[0, 0, 1], [0, 0, 1]], dtype=dt), + np.array([[0, 1, 1], [0, 1, 1]], dtype=dt), + np.array([[0, 0, 1], [2, 2, 3]], dtype=dt), + np.array([[0, 1, 1], [2, 3, 3]], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 0, -1, 1], x[0], NEIGH_MODE['mirror']) + assert_array_equal(l, r) + + # Simple, 1d tests + def test_simple(self, dt): + # Test padding with constant values + x = np.linspace(1, 5, 5).astype(dt) + r = [[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 0]] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 1], x[0], NEIGH_MODE['zero']) + assert_array_equal(l, r) + + r = [[1, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 1]] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 1], x[0], NEIGH_MODE['one']) + assert_array_equal(l, r) + + r = [[x[4], 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, x[4]]] + l = _multiarray_tests.test_neighborhood_iterator( + x, [-1, 1], x[4], NEIGH_MODE['constant']) + assert_array_equal(l, r) + + # Test mirror modes + def test_mirror(self, dt): + x = np.linspace(1, 5, 5).astype(dt) + r = np.array([[2, 1, 1, 2, 3], [1, 1, 2, 3, 4], [1, 2, 3, 4, 5], + [2, 3, 4, 5, 5], [3, 4, 5, 5, 4]], dtype=dt) + l = _multiarray_tests.test_neighborhood_iterator( + x, [-2, 2], x[1], NEIGH_MODE['mirror']) + assert_([i.dtype == dt for i in l]) + assert_array_equal(l, r) + + # Circular mode + def test_circular(self, dt): + x = np.linspace(1, 5, 5).astype(dt) + r = np.array([[4, 5, 1, 2, 3], [5, 1, 2, 3, 4], [1, 2, 3, 4, 5], + [2, 3, 4, 5, 1], [3, 4, 5, 1, 2]], dtype=dt) + l = _multiarray_tests.test_neighborhood_iterator( + x, [-2, 2], x[0], NEIGH_MODE['circular']) + assert_array_equal(l, r) + + +# Test stacking neighborhood iterators +class TestStackedNeighborhoodIter: + # Simple, 1d test: stacking 2 constant-padded neigh iterators + def test_simple_const(self): + dt = np.float64 + # Test zero and one padding for simple data type + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([0], dtype=dt), + np.array([0], dtype=dt), + np.array([1], dtype=dt), + np.array([2], dtype=dt), + np.array([3], dtype=dt), + np.array([0], dtype=dt), + np.array([0], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-2, 4], NEIGH_MODE['zero'], [0, 0], NEIGH_MODE['zero']) + assert_array_equal(l, r) + + r = [np.array([1, 0, 1], dtype=dt), + np.array([0, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 0], dtype=dt), + np.array([3, 0, 1], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [-1, 1], NEIGH_MODE['one']) + assert_array_equal(l, r) + + # 2nd simple, 1d test: stacking 2 neigh iterators, mixing const padding and + # mirror padding + def test_simple_mirror(self): + dt = np.float64 + # Stacking zero on top of mirror + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([0, 1, 1], dtype=dt), + np.array([1, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 3], dtype=dt), + np.array([3, 3, 0], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['mirror'], [-1, 1], NEIGH_MODE['zero']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([1, 0, 0], dtype=dt), + np.array([0, 0, 1], dtype=dt), + np.array([0, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 0], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [-2, 0], NEIGH_MODE['mirror']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero: 2nd + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([0, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 0], dtype=dt), + np.array([3, 0, 0], dtype=dt), + np.array([0, 0, 3], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [0, 2], NEIGH_MODE['mirror']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero: 3rd + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([1, 0, 0, 1, 2], dtype=dt), + np.array([0, 0, 1, 2, 3], dtype=dt), + np.array([0, 1, 2, 3, 0], dtype=dt), + np.array([1, 2, 3, 0, 0], dtype=dt), + np.array([2, 3, 0, 0, 3], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [-2, 2], NEIGH_MODE['mirror']) + assert_array_equal(l, r) + + # 3rd simple, 1d test: stacking 2 neigh iterators, mixing const padding and + # circular padding + def test_simple_circular(self): + dt = np.float64 + # Stacking zero on top of mirror + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([0, 3, 1], dtype=dt), + np.array([3, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 1], dtype=dt), + np.array([3, 1, 0], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['circular'], [-1, 1], NEIGH_MODE['zero']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([3, 0, 0], dtype=dt), + np.array([0, 0, 1], dtype=dt), + np.array([0, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 0], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [-2, 0], NEIGH_MODE['circular']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero: 2nd + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([0, 1, 2], dtype=dt), + np.array([1, 2, 3], dtype=dt), + np.array([2, 3, 0], dtype=dt), + np.array([3, 0, 0], dtype=dt), + np.array([0, 0, 1], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [0, 2], NEIGH_MODE['circular']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero: 3rd + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([3, 0, 0, 1, 2], dtype=dt), + np.array([0, 0, 1, 2, 3], dtype=dt), + np.array([0, 1, 2, 3, 0], dtype=dt), + np.array([1, 2, 3, 0, 0], dtype=dt), + np.array([2, 3, 0, 0, 1], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [-1, 3], NEIGH_MODE['zero'], [-2, 2], NEIGH_MODE['circular']) + assert_array_equal(l, r) + + # 4th simple, 1d test: stacking 2 neigh iterators, but with lower iterator + # being strictly within the array + def test_simple_strict_within(self): + dt = np.float64 + # Stacking zero on top of zero, first neighborhood strictly inside the + # array + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([1, 2, 3, 0], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [1, 1], NEIGH_MODE['zero'], [-1, 2], NEIGH_MODE['zero']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero, first neighborhood strictly inside the + # array + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([1, 2, 3, 3], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [1, 1], NEIGH_MODE['zero'], [-1, 2], NEIGH_MODE['mirror']) + assert_array_equal(l, r) + + # Stacking mirror on top of zero, first neighborhood strictly inside the + # array + x = np.array([1, 2, 3], dtype=dt) + r = [np.array([1, 2, 3, 1], dtype=dt)] + l = _multiarray_tests.test_neighborhood_iterator_oob( + x, [1, 1], NEIGH_MODE['zero'], [-1, 2], NEIGH_MODE['circular']) + assert_array_equal(l, r) + +class TestWarnings: + + def test_complex_warning(self): + x = np.array([1, 2]) + y = np.array([1-2j, 1+2j]) + + with warnings.catch_warnings(): + warnings.simplefilter("error", ComplexWarning) + assert_raises(ComplexWarning, x.__setitem__, slice(None), y) + assert_equal(x, [1, 2]) + + +class TestMinScalarType: + + def test_usigned_shortshort(self): + dt = np.min_scalar_type(2**8-1) + wanted = np.dtype('uint8') + assert_equal(wanted, dt) + + def test_usigned_short(self): + dt = np.min_scalar_type(2**16-1) + wanted = np.dtype('uint16') + assert_equal(wanted, dt) + + def test_usigned_int(self): + dt = np.min_scalar_type(2**32-1) + wanted = np.dtype('uint32') + assert_equal(wanted, dt) + + def test_usigned_longlong(self): + dt = np.min_scalar_type(2**63-1) + wanted = np.dtype('uint64') + assert_equal(wanted, dt) + + def test_object(self): + dt = np.min_scalar_type(2**64) + wanted = np.dtype('O') + assert_equal(wanted, dt) + + +from numpy._core._internal import _dtype_from_pep3118 + + +class TestPEP3118Dtype: + def _check(self, spec, wanted): + dt = np.dtype(wanted) + actual = _dtype_from_pep3118(spec) + assert_equal(actual, dt, + err_msg="spec %r != dtype %r" % (spec, wanted)) + + def test_native_padding(self): + align = np.dtype('i').alignment + for j in range(8): + if j == 0: + s = 'bi' + else: + s = 'b%dxi' % j + self._check('@'+s, {'f0': ('i1', 0), + 'f1': ('i', align*(1 + j//align))}) + self._check('='+s, {'f0': ('i1', 0), + 'f1': ('i', 1+j)}) + + def test_native_padding_2(self): + # Native padding should work also for structs and sub-arrays + self._check('x3T{xi}', {'f0': (({'f0': ('i', 4)}, (3,)), 4)}) + self._check('^x3T{xi}', {'f0': (({'f0': ('i', 1)}, (3,)), 1)}) + + def test_trailing_padding(self): + # Trailing padding should be included, *and*, the item size + # should match the alignment if in aligned mode + align = np.dtype('i').alignment + size = np.dtype('i').itemsize + + def aligned(n): + return align*(1 + (n-1)//align) + + base = dict(formats=['i'], names=['f0']) + + self._check('ix', dict(itemsize=aligned(size + 1), **base)) + self._check('ixx', dict(itemsize=aligned(size + 2), **base)) + self._check('ixxx', dict(itemsize=aligned(size + 3), **base)) + self._check('ixxxx', dict(itemsize=aligned(size + 4), **base)) + self._check('i7x', dict(itemsize=aligned(size + 7), **base)) + + self._check('^ix', dict(itemsize=size + 1, **base)) + self._check('^ixx', dict(itemsize=size + 2, **base)) + self._check('^ixxx', dict(itemsize=size + 3, **base)) + self._check('^ixxxx', dict(itemsize=size + 4, **base)) + self._check('^i7x', dict(itemsize=size + 7, **base)) + + def test_native_padding_3(self): + dt = np.dtype( + [('a', 'b'), ('b', 'i'), + ('sub', np.dtype('b,i')), ('c', 'i')], + align=True) + self._check("T{b:a:xxxi:b:T{b:f0:=i:f1:}:sub:xxxi:c:}", dt) + + dt = np.dtype( + [('a', 'b'), ('b', 'i'), ('c', 'b'), ('d', 'b'), + ('e', 'b'), ('sub', np.dtype('b,i', align=True))]) + self._check("T{b:a:=i:b:b:c:b:d:b:e:T{b:f0:xxxi:f1:}:sub:}", dt) + + def test_padding_with_array_inside_struct(self): + dt = np.dtype( + [('a', 'b'), ('b', 'i'), ('c', 'b', (3,)), + ('d', 'i')], + align=True) + self._check("T{b:a:xxxi:b:3b:c:xi:d:}", dt) + + def test_byteorder_inside_struct(self): + # The byte order after @T{=i} should be '=', not '@'. + # Check this by noting the absence of native alignment. + self._check('@T{^i}xi', {'f0': ({'f0': ('i', 0)}, 0), + 'f1': ('i', 5)}) + + def test_intra_padding(self): + # Natively aligned sub-arrays may require some internal padding + align = np.dtype('i').alignment + size = np.dtype('i').itemsize + + def aligned(n): + return (align*(1 + (n-1)//align)) + + self._check('(3)T{ix}', (dict( + names=['f0'], + formats=['i'], + offsets=[0], + itemsize=aligned(size + 1) + ), (3,))) + + def test_char_vs_string(self): + dt = np.dtype('c') + self._check('c', dt) + + dt = np.dtype([('f0', 'S1', (4,)), ('f1', 'S4')]) + self._check('4c4s', dt) + + def test_field_order(self): + # gh-9053 - previously, we relied on dictionary key order + self._check("(0)I:a:f:b:", [('a', 'I', (0,)), ('b', 'f')]) + self._check("(0)I:b:f:a:", [('b', 'I', (0,)), ('a', 'f')]) + + def test_unnamed_fields(self): + self._check('ii', [('f0', 'i'), ('f1', 'i')]) + self._check('ii:f0:', [('f1', 'i'), ('f0', 'i')]) + + self._check('i', 'i') + self._check('i:f0:', [('f0', 'i')]) + + +class TestNewBufferProtocol: + """ Test PEP3118 buffers """ + + def _check_roundtrip(self, obj): + obj = np.asarray(obj) + x = memoryview(obj) + y = np.asarray(x) + y2 = np.array(x) + assert_(not y.flags.owndata) + assert_(y2.flags.owndata) + + assert_equal(y.dtype, obj.dtype) + assert_equal(y.shape, obj.shape) + assert_array_equal(obj, y) + + assert_equal(y2.dtype, obj.dtype) + assert_equal(y2.shape, obj.shape) + assert_array_equal(obj, y2) + + def test_roundtrip(self): + x = np.array([1, 2, 3, 4, 5], dtype='i4') + self._check_roundtrip(x) + + x = np.array([[1, 2], [3, 4]], dtype=np.float64) + self._check_roundtrip(x) + + x = np.zeros((3, 3, 3), dtype=np.float32)[:, 0,:] + self._check_roundtrip(x) + + dt = [('a', 'b'), + ('b', 'h'), + ('c', 'i'), + ('d', 'l'), + ('dx', 'q'), + ('e', 'B'), + ('f', 'H'), + ('g', 'I'), + ('h', 'L'), + ('hx', 'Q'), + ('i', np.single), + ('j', np.double), + ('k', np.longdouble), + ('ix', np.csingle), + ('jx', np.cdouble), + ('kx', np.clongdouble), + ('l', 'S4'), + ('m', 'U4'), + ('n', 'V3'), + ('o', '?'), + ('p', np.half), + ] + x = np.array( + [(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + b'aaaa', 'bbbb', b'xxx', True, 1.0)], + dtype=dt) + self._check_roundtrip(x) + + x = np.array(([[1, 2], [3, 4]],), dtype=[('a', (int, (2, 2)))]) + self._check_roundtrip(x) + + x = np.array([1, 2, 3], dtype='>i2') + self._check_roundtrip(x) + + x = np.array([1, 2, 3], dtype='') + x = np.zeros(4, dtype=dt) + self._check_roundtrip(x) + + def test_roundtrip_scalar(self): + # Issue #4015. + self._check_roundtrip(0) + + def test_invalid_buffer_format(self): + # datetime64 cannot be used fully in a buffer yet + # Should be fixed in the next Numpy major release + dt = np.dtype([('a', 'uint16'), ('b', 'M8[s]')]) + a = np.empty(3, dt) + assert_raises((ValueError, BufferError), memoryview, a) + assert_raises((ValueError, BufferError), memoryview, np.array((3), 'M8[D]')) + + def test_export_simple_1d(self): + x = np.array([1, 2, 3, 4, 5], dtype='i') + y = memoryview(x) + assert_equal(y.format, 'i') + assert_equal(y.shape, (5,)) + assert_equal(y.ndim, 1) + assert_equal(y.strides, (4,)) + assert_equal(y.suboffsets, ()) + assert_equal(y.itemsize, 4) + + def test_export_simple_nd(self): + x = np.array([[1, 2], [3, 4]], dtype=np.float64) + y = memoryview(x) + assert_equal(y.format, 'd') + assert_equal(y.shape, (2, 2)) + assert_equal(y.ndim, 2) + assert_equal(y.strides, (16, 8)) + assert_equal(y.suboffsets, ()) + assert_equal(y.itemsize, 8) + + def test_export_discontiguous(self): + x = np.zeros((3, 3, 3), dtype=np.float32)[:, 0,:] + y = memoryview(x) + assert_equal(y.format, 'f') + assert_equal(y.shape, (3, 3)) + assert_equal(y.ndim, 2) + assert_equal(y.strides, (36, 4)) + assert_equal(y.suboffsets, ()) + assert_equal(y.itemsize, 4) + + def test_export_record(self): + dt = [('a', 'b'), + ('b', 'h'), + ('c', 'i'), + ('d', 'l'), + ('dx', 'q'), + ('e', 'B'), + ('f', 'H'), + ('g', 'I'), + ('h', 'L'), + ('hx', 'Q'), + ('i', np.single), + ('j', np.double), + ('k', np.longdouble), + ('ix', np.csingle), + ('jx', np.cdouble), + ('kx', np.clongdouble), + ('l', 'S4'), + ('m', 'U4'), + ('n', 'V3'), + ('o', '?'), + ('p', np.half), + ] + x = np.array( + [(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + b'aaaa', 'bbbb', b' ', True, 1.0)], + dtype=dt) + y = memoryview(x) + assert_equal(y.shape, (1,)) + assert_equal(y.ndim, 1) + assert_equal(y.suboffsets, ()) + + sz = sum([np.dtype(b).itemsize for a, b in dt]) + if np.dtype('l').itemsize == 4: + assert_equal(y.format, 'T{b:a:=h:b:i:c:l:d:q:dx:B:e:@H:f:=I:g:L:h:Q:hx:f:i:d:j:^g:k:=Zf:ix:Zd:jx:^Zg:kx:4s:l:=4w:m:3x:n:?:o:@e:p:}') + else: + assert_equal(y.format, 'T{b:a:=h:b:i:c:q:d:q:dx:B:e:@H:f:=I:g:Q:h:Q:hx:f:i:d:j:^g:k:=Zf:ix:Zd:jx:^Zg:kx:4s:l:=4w:m:3x:n:?:o:@e:p:}') + assert_equal(y.strides, (sz,)) + assert_equal(y.itemsize, sz) + + def test_export_subarray(self): + x = np.array(([[1, 2], [3, 4]],), dtype=[('a', ('i', (2, 2)))]) + y = memoryview(x) + assert_equal(y.format, 'T{(2,2)i:a:}') + assert_equal(y.shape, ()) + assert_equal(y.ndim, 0) + assert_equal(y.strides, ()) + assert_equal(y.suboffsets, ()) + assert_equal(y.itemsize, 16) + + def test_export_endian(self): + x = np.array([1, 2, 3], dtype='>i') + y = memoryview(x) + if sys.byteorder == 'little': + assert_equal(y.format, '>i') + else: + assert_equal(y.format, 'i') + + x = np.array([1, 2, 3], dtype=' np.array(0, dtype=dt1), "type %s failed" % (dt1,)) + assert_(not 1 < np.array(0, dtype=dt1), "type %s failed" % (dt1,)) + + for dt2 in np.typecodes['AllInteger']: + assert_(np.array(1, dtype=dt1) > np.array(0, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + assert_(not np.array(1, dtype=dt1) < np.array(0, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + + # Unsigned integers + for dt1 in 'BHILQP': + assert_(-1 < np.array(1, dtype=dt1), "type %s failed" % (dt1,)) + assert_(not -1 > np.array(1, dtype=dt1), "type %s failed" % (dt1,)) + assert_(-1 != np.array(1, dtype=dt1), "type %s failed" % (dt1,)) + + # Unsigned vs signed + for dt2 in 'bhilqp': + assert_(np.array(1, dtype=dt1) > np.array(-1, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + assert_(not np.array(1, dtype=dt1) < np.array(-1, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + assert_(np.array(1, dtype=dt1) != np.array(-1, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + + # Signed integers and floats + for dt1 in 'bhlqp' + np.typecodes['Float']: + assert_(1 > np.array(-1, dtype=dt1), "type %s failed" % (dt1,)) + assert_(not 1 < np.array(-1, dtype=dt1), "type %s failed" % (dt1,)) + assert_(-1 == np.array(-1, dtype=dt1), "type %s failed" % (dt1,)) + + for dt2 in 'bhlqp' + np.typecodes['Float']: + assert_(np.array(1, dtype=dt1) > np.array(-1, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + assert_(not np.array(1, dtype=dt1) < np.array(-1, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + assert_(np.array(-1, dtype=dt1) == np.array(-1, dtype=dt2), + "type %s and %s failed" % (dt1, dt2)) + + def test_to_bool_scalar(self): + assert_equal(bool(np.array([False])), False) + assert_equal(bool(np.array([True])), True) + assert_equal(bool(np.array([[42]])), True) + + def test_to_bool_scalar_not_convertible(self): + + class NotConvertible: + def __bool__(self): + raise NotImplementedError + + assert_raises(NotImplementedError, bool, np.array(NotConvertible())) + assert_raises(NotImplementedError, bool, np.array([NotConvertible()])) + if IS_PYSTON: + pytest.skip("Pyston disables recursion checking") + + self_containing = np.array([None]) + self_containing[0] = self_containing + + Error = RecursionError + + assert_raises(Error, bool, self_containing) # previously stack overflow + self_containing[0] = None # resolve circular reference + + def test_to_bool_scalar_size_errors(self): + with pytest.raises(ValueError, match=".*one element is ambiguous"): + bool(np.array([1, 2])) + + with pytest.raises(ValueError, match=".*empty array is ambiguous"): + bool(np.empty((3, 0))) + + with pytest.raises(ValueError, match=".*empty array is ambiguous"): + bool(np.empty((0,))) + + def test_to_int_scalar(self): + # gh-9972 means that these aren't always the same + int_funcs = (int, lambda x: x.__int__()) + for int_func in int_funcs: + assert_equal(int_func(np.array(0)), 0) + with assert_warns(DeprecationWarning): + assert_equal(int_func(np.array([1])), 1) + with assert_warns(DeprecationWarning): + assert_equal(int_func(np.array([[42]])), 42) + assert_raises(TypeError, int_func, np.array([1, 2])) + + # gh-9972 + assert_equal(4, int_func(np.array('4'))) + assert_equal(5, int_func(np.bytes_(b'5'))) + assert_equal(6, int_func(np.str_('6'))) + + # The delegation of int() to __trunc__ was deprecated in + # Python 3.11. + if sys.version_info < (3, 11): + class HasTrunc: + def __trunc__(self): + return 3 + assert_equal(3, int_func(np.array(HasTrunc()))) + with assert_warns(DeprecationWarning): + assert_equal(3, int_func(np.array([HasTrunc()]))) + else: + pass + + class NotConvertible: + def __int__(self): + raise NotImplementedError + assert_raises(NotImplementedError, + int_func, np.array(NotConvertible())) + with assert_warns(DeprecationWarning): + assert_raises(NotImplementedError, + int_func, np.array([NotConvertible()])) + + +class TestWhere: + def test_basic(self): + dts = [bool, np.int16, np.int32, np.int64, np.double, np.complex128, + np.longdouble, np.clongdouble] + for dt in dts: + c = np.ones(53, dtype=bool) + assert_equal(np.where( c, dt(0), dt(1)), dt(0)) + assert_equal(np.where(~c, dt(0), dt(1)), dt(1)) + assert_equal(np.where(True, dt(0), dt(1)), dt(0)) + assert_equal(np.where(False, dt(0), dt(1)), dt(1)) + d = np.ones_like(c).astype(dt) + e = np.zeros_like(d) + r = d.astype(dt) + c[7] = False + r[7] = e[7] + assert_equal(np.where(c, e, e), e) + assert_equal(np.where(c, d, e), r) + assert_equal(np.where(c, d, e[0]), r) + assert_equal(np.where(c, d[0], e), r) + assert_equal(np.where(c[::2], d[::2], e[::2]), r[::2]) + assert_equal(np.where(c[1::2], d[1::2], e[1::2]), r[1::2]) + assert_equal(np.where(c[::3], d[::3], e[::3]), r[::3]) + assert_equal(np.where(c[1::3], d[1::3], e[1::3]), r[1::3]) + assert_equal(np.where(c[::-2], d[::-2], e[::-2]), r[::-2]) + assert_equal(np.where(c[::-3], d[::-3], e[::-3]), r[::-3]) + assert_equal(np.where(c[1::-3], d[1::-3], e[1::-3]), r[1::-3]) + + @pytest.mark.skipif(IS_WASM, reason="no wasm fp exception support") + def test_exotic(self): + # object + assert_array_equal(np.where(True, None, None), np.array(None)) + # zero sized + m = np.array([], dtype=bool).reshape(0, 3) + b = np.array([], dtype=np.float64).reshape(0, 3) + assert_array_equal(np.where(m, 0, b), np.array([]).reshape(0, 3)) + + # object cast + d = np.array([-1.34, -0.16, -0.54, -0.31, -0.08, -0.95, 0.000, 0.313, + 0.547, -0.18, 0.876, 0.236, 1.969, 0.310, 0.699, 1.013, + 1.267, 0.229, -1.39, 0.487]) + nan = float('NaN') + e = np.array(['5z', '0l', nan, 'Wz', nan, nan, 'Xq', 'cs', nan, nan, + 'QN', nan, nan, 'Fd', nan, nan, 'kp', nan, '36', 'i1'], + dtype=object) + m = np.array([0, 0, 1, 0, 1, 1, 0, 0, 1, 1, + 0, 1, 1, 0, 1, 1, 0, 1, 0, 0], dtype=bool) + + r = e[:] + r[np.where(m)] = d[np.where(m)] + assert_array_equal(np.where(m, d, e), r) + + r = e[:] + r[np.where(~m)] = d[np.where(~m)] + assert_array_equal(np.where(m, e, d), r) + + assert_array_equal(np.where(m, e, e), e) + + # minimal dtype result with NaN scalar (e.g required by pandas) + d = np.array([1., 2.], dtype=np.float32) + e = float('NaN') + assert_equal(np.where(True, d, e).dtype, np.float32) + e = float('Infinity') + assert_equal(np.where(True, d, e).dtype, np.float32) + e = float('-Infinity') + assert_equal(np.where(True, d, e).dtype, np.float32) + # With NEP 50 adopted, the float will overflow here: + e = float(1e150) + with pytest.warns(RuntimeWarning, match="overflow"): + res = np.where(True, d, e) + assert res.dtype == np.float32 + + def test_ndim(self): + c = [True, False] + a = np.zeros((2, 25)) + b = np.ones((2, 25)) + r = np.where(np.array(c)[:,np.newaxis], a, b) + assert_array_equal(r[0], a[0]) + assert_array_equal(r[1], b[0]) + + a = a.T + b = b.T + r = np.where(c, a, b) + assert_array_equal(r[:,0], a[:,0]) + assert_array_equal(r[:,1], b[:,0]) + + def test_dtype_mix(self): + c = np.array([False, True, False, False, False, False, True, False, + False, False, True, False]) + a = np.uint32(1) + b = np.array([5., 0., 3., 2., -1., -4., 0., -10., 10., 1., 0., 3.], + dtype=np.float64) + r = np.array([5., 1., 3., 2., -1., -4., 1., -10., 10., 1., 1., 3.], + dtype=np.float64) + assert_equal(np.where(c, a, b), r) + + a = a.astype(np.float32) + b = b.astype(np.int64) + assert_equal(np.where(c, a, b), r) + + # non bool mask + c = c.astype(int) + c[c != 0] = 34242324 + assert_equal(np.where(c, a, b), r) + # invert + tmpmask = c != 0 + c[c == 0] = 41247212 + c[tmpmask] = 0 + assert_equal(np.where(c, b, a), r) + + def test_foreign(self): + c = np.array([False, True, False, False, False, False, True, False, + False, False, True, False]) + r = np.array([5., 1., 3., 2., -1., -4., 1., -10., 10., 1., 1., 3.], + dtype=np.float64) + a = np.ones(1, dtype='>i4') + b = np.array([5., 0., 3., 2., -1., -4., 0., -10., 10., 1., 0., 3.], + dtype=np.float64) + assert_equal(np.where(c, a, b), r) + + b = b.astype('>f8') + assert_equal(np.where(c, a, b), r) + + a = a.astype('i4') + assert_equal(np.where(c, a, b), r) + + def test_error(self): + c = [True, True] + a = np.ones((4, 5)) + b = np.ones((5, 5)) + assert_raises(ValueError, np.where, c, a, a) + assert_raises(ValueError, np.where, c[0], a, b) + + def test_string(self): + # gh-4778 check strings are properly filled with nulls + a = np.array("abc") + b = np.array("x" * 753) + assert_equal(np.where(True, a, b), "abc") + assert_equal(np.where(False, b, a), "abc") + + # check native datatype sized strings + a = np.array("abcd") + b = np.array("x" * 8) + assert_equal(np.where(True, a, b), "abcd") + assert_equal(np.where(False, b, a), "abcd") + + def test_empty_result(self): + # pass empty where result through an assignment which reads the data of + # empty arrays, error detectable with valgrind, see gh-8922 + x = np.zeros((1, 1)) + ibad = np.vstack(np.where(x == 99.)) + assert_array_equal(ibad, + np.atleast_2d(np.array([[],[]], dtype=np.intp))) + + def test_largedim(self): + # invalid read regression gh-9304 + shape = [10, 2, 3, 4, 5, 6] + np.random.seed(2) + array = np.random.rand(*shape) + + for i in range(10): + benchmark = array.nonzero() + result = array.nonzero() + assert_array_equal(benchmark, result) + + def test_kwargs(self): + a = np.zeros(1) + with assert_raises(TypeError): + np.where(a, x=a, y=a) + + +if not IS_PYPY: + # sys.getsizeof() is not valid on PyPy + class TestSizeOf: + + def test_empty_array(self): + x = np.array([]) + assert_(sys.getsizeof(x) > 0) + + def check_array(self, dtype): + elem_size = dtype(0).itemsize + + for length in [10, 50, 100, 500]: + x = np.arange(length, dtype=dtype) + assert_(sys.getsizeof(x) > length * elem_size) + + def test_array_int32(self): + self.check_array(np.int32) + + def test_array_int64(self): + self.check_array(np.int64) + + def test_array_float32(self): + self.check_array(np.float32) + + def test_array_float64(self): + self.check_array(np.float64) + + def test_view(self): + d = np.ones(100) + assert_(sys.getsizeof(d[...]) < sys.getsizeof(d)) + + def test_reshape(self): + d = np.ones(100) + assert_(sys.getsizeof(d) < sys.getsizeof(d.reshape(100, 1, 1).copy())) + + @_no_tracing + def test_resize(self): + d = np.ones(100) + old = sys.getsizeof(d) + d.resize(50) + assert_(old > sys.getsizeof(d)) + d.resize(150) + assert_(old < sys.getsizeof(d)) + + @pytest.mark.parametrize("dtype", ["u4,f4", "u4,O"]) + def test_resize_structured(self, dtype): + a = np.array([(0, 0.0) for i in range(5)], dtype=dtype) + a.resize(1000) + assert_array_equal(a, np.zeros(1000, dtype=dtype)) + + def test_error(self): + d = np.ones(100) + assert_raises(TypeError, d.__sizeof__, "a") + + +class TestHashing: + + def test_arrays_not_hashable(self): + x = np.ones(3) + assert_raises(TypeError, hash, x) + + def test_collections_hashable(self): + x = np.array([]) + assert_(not isinstance(x, collections.abc.Hashable)) + + +class TestArrayPriority: + # This will go away when __array_priority__ is settled, meanwhile + # it serves to check unintended changes. + op = operator + binary_ops = [ + op.pow, op.add, op.sub, op.mul, op.floordiv, op.truediv, op.mod, + op.and_, op.or_, op.xor, op.lshift, op.rshift, op.mod, op.gt, + op.ge, op.lt, op.le, op.ne, op.eq + ] + + class Foo(np.ndarray): + __array_priority__ = 100. + + def __new__(cls, *args, **kwargs): + return np.array(*args, **kwargs).view(cls) + + class Bar(np.ndarray): + __array_priority__ = 101. + + def __new__(cls, *args, **kwargs): + return np.array(*args, **kwargs).view(cls) + + class Other: + __array_priority__ = 1000. + + def _all(self, other): + return self.__class__() + + __add__ = __radd__ = _all + __sub__ = __rsub__ = _all + __mul__ = __rmul__ = _all + __pow__ = __rpow__ = _all + __div__ = __rdiv__ = _all + __mod__ = __rmod__ = _all + __truediv__ = __rtruediv__ = _all + __floordiv__ = __rfloordiv__ = _all + __and__ = __rand__ = _all + __xor__ = __rxor__ = _all + __or__ = __ror__ = _all + __lshift__ = __rlshift__ = _all + __rshift__ = __rrshift__ = _all + __eq__ = _all + __ne__ = _all + __gt__ = _all + __ge__ = _all + __lt__ = _all + __le__ = _all + + def test_ndarray_subclass(self): + a = np.array([1, 2]) + b = self.Bar([1, 2]) + for f in self.binary_ops: + msg = repr(f) + assert_(isinstance(f(a, b), self.Bar), msg) + assert_(isinstance(f(b, a), self.Bar), msg) + + def test_ndarray_other(self): + a = np.array([1, 2]) + b = self.Other() + for f in self.binary_ops: + msg = repr(f) + assert_(isinstance(f(a, b), self.Other), msg) + assert_(isinstance(f(b, a), self.Other), msg) + + def test_subclass_subclass(self): + a = self.Foo([1, 2]) + b = self.Bar([1, 2]) + for f in self.binary_ops: + msg = repr(f) + assert_(isinstance(f(a, b), self.Bar), msg) + assert_(isinstance(f(b, a), self.Bar), msg) + + def test_subclass_other(self): + a = self.Foo([1, 2]) + b = self.Other() + for f in self.binary_ops: + msg = repr(f) + assert_(isinstance(f(a, b), self.Other), msg) + assert_(isinstance(f(b, a), self.Other), msg) + + +class TestBytestringArrayNonzero: + + def test_empty_bstring_array_is_falsey(self): + assert_(not np.array([''], dtype=str)) + + def test_whitespace_bstring_array_is_truthy(self): + a = np.array(['spam'], dtype=str) + a[0] = ' \0\0' + assert_(a) + + def test_all_null_bstring_array_is_falsey(self): + a = np.array(['spam'], dtype=str) + a[0] = '\0\0\0\0' + assert_(not a) + + def test_null_inside_bstring_array_is_truthy(self): + a = np.array(['spam'], dtype=str) + a[0] = ' \0 \0' + assert_(a) + + +class TestUnicodeEncoding: + """ + Tests for encoding related bugs, such as UCS2 vs UCS4, round-tripping + issues, etc + """ + def test_round_trip(self): + """ Tests that GETITEM, SETITEM, and PyArray_Scalar roundtrip """ + # gh-15363 + arr = np.zeros(shape=(), dtype="U1") + for i in range(1, sys.maxunicode + 1): + expected = chr(i) + arr[()] = expected + assert arr[()] == expected + assert arr.item() == expected + + def test_assign_scalar(self): + # gh-3258 + l = np.array(['aa', 'bb']) + l[:] = np.str_('cc') + assert_equal(l, ['cc', 'cc']) + + def test_fill_scalar(self): + # gh-7227 + l = np.array(['aa', 'bb']) + l.fill(np.str_('cc')) + assert_equal(l, ['cc', 'cc']) + + +class TestUnicodeArrayNonzero: + + def test_empty_ustring_array_is_falsey(self): + assert_(not np.array([''], dtype=np.str_)) + + def test_whitespace_ustring_array_is_truthy(self): + a = np.array(['eggs'], dtype=np.str_) + a[0] = ' \0\0' + assert_(a) + + def test_all_null_ustring_array_is_falsey(self): + a = np.array(['eggs'], dtype=np.str_) + a[0] = '\0\0\0\0' + assert_(not a) + + def test_null_inside_ustring_array_is_truthy(self): + a = np.array(['eggs'], dtype=np.str_) + a[0] = ' \0 \0' + assert_(a) + + +class TestFormat: + + def test_0d(self): + a = np.array(np.pi) + assert_equal('{:0.3g}'.format(a), '3.14') + assert_equal('{:0.3g}'.format(a[()]), '3.14') + + def test_1d_no_format(self): + a = np.array([np.pi]) + assert_equal('{}'.format(a), str(a)) + + def test_1d_format(self): + # until gh-5543, ensure that the behaviour matches what it used to be + a = np.array([np.pi]) + assert_raises(TypeError, '{:30}'.format, a) + + +from numpy.testing import IS_PYPY + + +class TestCTypes: + + def test_ctypes_is_available(self): + test_arr = np.array([[1, 2, 3], [4, 5, 6]]) + + assert_equal(ctypes, test_arr.ctypes._ctypes) + assert_equal(tuple(test_arr.ctypes.shape), (2, 3)) + + def test_ctypes_is_not_available(self): + from numpy._core import _internal + _internal.ctypes = None + try: + test_arr = np.array([[1, 2, 3], [4, 5, 6]]) + + assert_(isinstance(test_arr.ctypes._ctypes, + _internal._missing_ctypes)) + assert_equal(tuple(test_arr.ctypes.shape), (2, 3)) + finally: + _internal.ctypes = ctypes + + def _make_readonly(x): + x.flags.writeable = False + return x + + @pytest.mark.parametrize('arr', [ + np.array([1, 2, 3]), + np.array([['one', 'two'], ['three', 'four']]), + np.array((1, 2), dtype='i4,i4'), + np.zeros((2,), dtype= + np.dtype(dict( + formats=['2, [44, 55]) + assert_equal(a, np.array([[0, 44], [1, 55], [2, 44]])) + # hit one of the failing paths + assert_raises(ValueError, np.place, a, a>20, []) + + def test_put_noncontiguous(self): + a = np.arange(6).reshape(2, 3).T # force non-c-contiguous + np.put(a, [0, 2], [44, 55]) + assert_equal(a, np.array([[44, 3], [55, 4], [2, 5]])) + + def test_putmask_noncontiguous(self): + a = np.arange(6).reshape(2, 3).T # force non-c-contiguous + # uses arr_putmask + np.putmask(a, a>2, a**2) + assert_equal(a, np.array([[0, 9], [1, 16], [2, 25]])) + + def test_take_mode_raise(self): + a = np.arange(6, dtype='int') + out = np.empty(2, dtype='int') + np.take(a, [0, 2], out=out, mode='raise') + assert_equal(out, np.array([0, 2])) + + def test_choose_mod_raise(self): + a = np.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]]) + out = np.empty((3,3), dtype='int') + choices = [-10, 10] + np.choose(a, choices, out=out, mode='raise') + assert_equal(out, np.array([[ 10, -10, 10], + [-10, 10, -10], + [ 10, -10, 10]])) + + def test_flatiter__array__(self): + a = np.arange(9).reshape(3, 3) + b = a.T.flat + c = b.__array__() + # triggers the WRITEBACKIFCOPY resolution, assuming refcount semantics + del c + + def test_dot_out(self): + # if HAVE_CBLAS, will use WRITEBACKIFCOPY + a = np.arange(9, dtype=float).reshape(3, 3) + b = np.dot(a, a, out=a) + assert_equal(b, np.array([[15, 18, 21], [42, 54, 66], [69, 90, 111]])) + + def test_view_assign(self): + from numpy._core._multiarray_tests import ( + npy_create_writebackifcopy, npy_resolve + ) + + arr = np.arange(9).reshape(3, 3).T + arr_wb = npy_create_writebackifcopy(arr) + assert_(arr_wb.flags.writebackifcopy) + assert_(arr_wb.base is arr) + arr_wb[...] = -100 + npy_resolve(arr_wb) + # arr changes after resolve, even though we assigned to arr_wb + assert_equal(arr, -100) + # after resolve, the two arrays no longer reference each other + assert_(arr_wb.ctypes.data != 0) + assert_equal(arr_wb.base, None) + # assigning to arr_wb does not get transferred to arr + arr_wb[...] = 100 + assert_equal(arr, -100) + + @pytest.mark.leaks_references( + reason="increments self in dealloc; ignore since deprecated path.") + def test_dealloc_warning(self): + with suppress_warnings() as sup: + sup.record(RuntimeWarning) + arr = np.arange(9).reshape(3, 3) + v = arr.T + _multiarray_tests.npy_abuse_writebackifcopy(v) + assert len(sup.log) == 1 + + def test_view_discard_refcount(self): + from numpy._core._multiarray_tests import ( + npy_create_writebackifcopy, npy_discard + ) + + arr = np.arange(9).reshape(3, 3).T + orig = arr.copy() + if HAS_REFCOUNT: + arr_cnt = sys.getrefcount(arr) + arr_wb = npy_create_writebackifcopy(arr) + assert_(arr_wb.flags.writebackifcopy) + assert_(arr_wb.base is arr) + arr_wb[...] = -100 + npy_discard(arr_wb) + # arr remains unchanged after discard + assert_equal(arr, orig) + # after discard, the two arrays no longer reference each other + assert_(arr_wb.ctypes.data != 0) + assert_equal(arr_wb.base, None) + if HAS_REFCOUNT: + assert_equal(arr_cnt, sys.getrefcount(arr)) + # assigning to arr_wb does not get transferred to arr + arr_wb[...] = 100 + assert_equal(arr, orig) + + +class TestArange: + def test_infinite(self): + assert_raises_regex( + ValueError, "size exceeded", + np.arange, 0, np.inf + ) + + def test_nan_step(self): + assert_raises_regex( + ValueError, "cannot compute length", + np.arange, 0, 1, np.nan + ) + + def test_zero_step(self): + assert_raises(ZeroDivisionError, np.arange, 0, 10, 0) + assert_raises(ZeroDivisionError, np.arange, 0.0, 10.0, 0.0) + + # empty range + assert_raises(ZeroDivisionError, np.arange, 0, 0, 0) + assert_raises(ZeroDivisionError, np.arange, 0.0, 0.0, 0.0) + + def test_require_range(self): + assert_raises(TypeError, np.arange) + assert_raises(TypeError, np.arange, step=3) + assert_raises(TypeError, np.arange, dtype='int64') + assert_raises(TypeError, np.arange, start=4) + + def test_start_stop_kwarg(self): + keyword_stop = np.arange(stop=3) + keyword_zerotostop = np.arange(0, stop=3) + keyword_start_stop = np.arange(start=3, stop=9) + + assert len(keyword_stop) == 3 + assert len(keyword_zerotostop) == 3 + assert len(keyword_start_stop) == 6 + assert_array_equal(keyword_stop, keyword_zerotostop) + + def test_arange_booleans(self): + # Arange makes some sense for booleans and works up to length 2. + # But it is weird since `arange(2, 4, dtype=bool)` works. + # Arguably, much or all of this could be deprecated/removed. + res = np.arange(False, dtype=bool) + assert_array_equal(res, np.array([], dtype="bool")) + + res = np.arange(True, dtype="bool") + assert_array_equal(res, [False]) + + res = np.arange(2, dtype="bool") + assert_array_equal(res, [False, True]) + + # This case is especially weird, but drops out without special case: + res = np.arange(6, 8, dtype="bool") + assert_array_equal(res, [True, True]) + + with pytest.raises(TypeError): + np.arange(3, dtype="bool") + + @pytest.mark.parametrize("dtype", ["S3", "U", "5i"]) + def test_rejects_bad_dtypes(self, dtype): + dtype = np.dtype(dtype) + DType_name = re.escape(str(type(dtype))) + with pytest.raises(TypeError, + match=rf"arange\(\) not supported for inputs .* {DType_name}"): + np.arange(2, dtype=dtype) + + def test_rejects_strings(self): + # Explicitly test error for strings which may call "b" - "a": + DType_name = re.escape(str(type(np.array("a").dtype))) + with pytest.raises(TypeError, + match=rf"arange\(\) not supported for inputs .* {DType_name}"): + np.arange("a", "b") + + def test_byteswapped(self): + res_be = np.arange(1, 1000, dtype=">i4") + res_le = np.arange(1, 1000, dtype="i4" + assert res_le.dtype == " arr2 + + +@pytest.mark.parametrize("op", [ + operator.eq, operator.ne, operator.le, operator.lt, operator.ge, + operator.gt]) +def test_comparisons_forwards_error(op): + class NotArray: + def __array__(self, dtype=None, copy=None): + raise TypeError("run you fools") + + with pytest.raises(TypeError, match="run you fools"): + op(np.arange(2), NotArray()) + + with pytest.raises(TypeError, match="run you fools"): + op(NotArray(), np.arange(2)) + + +def test_richcompare_scalar_boolean_singleton_return(): + # These are currently guaranteed to be the boolean singletons, but maybe + # returning NumPy booleans would also be OK: + assert (np.array(0) == "a") is False + assert (np.array(0) != "a") is True + assert (np.int16(0) == "a") is False + assert (np.int16(0) != "a") is True + + +@pytest.mark.parametrize("op", [ + operator.eq, operator.ne, operator.le, operator.lt, operator.ge, + operator.gt]) +def test_ragged_comparison_fails(op): + # This needs to convert the internal array to True/False, which fails: + a = np.array([1, np.array([1, 2, 3])], dtype=object) + b = np.array([1, np.array([1, 2, 3])], dtype=object) + + with pytest.raises(ValueError, match="The truth value.*ambiguous"): + op(a, b) + + +@pytest.mark.parametrize( + ["fun", "npfun"], + [ + (_multiarray_tests.npy_cabs, np.absolute), + (_multiarray_tests.npy_carg, np.angle) + ] +) +@pytest.mark.parametrize("x", [1, np.inf, -np.inf, np.nan]) +@pytest.mark.parametrize("y", [1, np.inf, -np.inf, np.nan]) +@pytest.mark.parametrize("test_dtype", np.complexfloating.__subclasses__()) +def test_npymath_complex(fun, npfun, x, y, test_dtype): + # Smoketest npymath functions + z = test_dtype(complex(x, y)) + with np.errstate(invalid='ignore'): + # Fallback implementations may emit a warning for +-inf (see gh-24876): + # RuntimeWarning: invalid value encountered in absolute + got = fun(z) + expected = npfun(z) + assert_allclose(got, expected) + + +def test_npymath_real(): + # Smoketest npymath functions + from numpy._core._multiarray_tests import ( + npy_log10, npy_cosh, npy_sinh, npy_tan, npy_tanh) + + funcs = {npy_log10: np.log10, + npy_cosh: np.cosh, + npy_sinh: np.sinh, + npy_tan: np.tan, + npy_tanh: np.tanh} + vals = (1, np.inf, -np.inf, np.nan) + types = (np.float32, np.float64, np.longdouble) + + with np.errstate(all='ignore'): + for fun, npfun in funcs.items(): + for x, t in itertools.product(vals, types): + z = t(x) + got = fun(z) + expected = npfun(z) + assert_allclose(got, expected) + +def test_uintalignment_and_alignment(): + # alignment code needs to satisfy these requirements: + # 1. numpy structs match C struct layout + # 2. ufuncs/casting is safe wrt to aligned access + # 3. copy code is safe wrt to "uint alidned" access + # + # Complex types are the main problem, whose alignment may not be the same + # as their "uint alignment". + # + # This test might only fail on certain platforms, where uint64 alignment is + # not equal to complex64 alignment. The second 2 tests will only fail + # for DEBUG=1. + + d1 = np.dtype('u1,c8', align=True) + d2 = np.dtype('u4,c8', align=True) + d3 = np.dtype({'names': ['a', 'b'], 'formats': ['u1', d1]}, align=True) + + assert_equal(np.zeros(1, dtype=d1)['f1'].flags['ALIGNED'], True) + assert_equal(np.zeros(1, dtype=d2)['f1'].flags['ALIGNED'], True) + assert_equal(np.zeros(1, dtype='u1,c8')['f1'].flags['ALIGNED'], False) + + # check that C struct matches numpy struct size + s = _multiarray_tests.get_struct_alignments() + for d, (alignment, size) in zip([d1,d2,d3], s): + assert_equal(d.alignment, alignment) + assert_equal(d.itemsize, size) + + # check that ufuncs don't complain in debug mode + # (this is probably OK if the aligned flag is true above) + src = np.zeros((2,2), dtype=d1)['f1'] # 4-byte aligned, often + np.exp(src) # assert fails? + + # check that copy code doesn't complain in debug mode + dst = np.zeros((2,2), dtype='c8') + dst[:,1] = src[:,1] # assert in lowlevel_strided_loops fails? + +class TestAlignment: + # adapted from scipy._lib.tests.test__util.test__aligned_zeros + # Checks that unusual memory alignments don't trip up numpy. + + def check(self, shape, dtype, order, align): + err_msg = repr((shape, dtype, order, align)) + x = _aligned_zeros(shape, dtype, order, align=align) + if align is None: + align = np.dtype(dtype).alignment + assert_equal(x.__array_interface__['data'][0] % align, 0) + if hasattr(shape, '__len__'): + assert_equal(x.shape, shape, err_msg) + else: + assert_equal(x.shape, (shape,), err_msg) + assert_equal(x.dtype, dtype) + if order == "C": + assert_(x.flags.c_contiguous, err_msg) + elif order == "F": + if x.size > 0: + assert_(x.flags.f_contiguous, err_msg) + elif order is None: + assert_(x.flags.c_contiguous, err_msg) + else: + raise ValueError + + def test_various_alignments(self): + for align in [1, 2, 3, 4, 8, 12, 16, 32, 64, None]: + for n in [0, 1, 3, 11]: + for order in ["C", "F", None]: + for dtype in list(np.typecodes["All"]) + ['i4,i4,i4']: + if dtype == 'O': + # object dtype can't be misaligned + continue + for shape in [n, (1, 2, 3, n)]: + self.check(shape, np.dtype(dtype), order, align) + + def test_strided_loop_alignments(self): + # particularly test that complex64 and float128 use right alignment + # code-paths, since these are particularly problematic. It is useful to + # turn on USE_DEBUG for this test, so lowlevel-loop asserts are run. + for align in [1, 2, 4, 8, 12, 16, None]: + xf64 = _aligned_zeros(3, np.float64) + + xc64 = _aligned_zeros(3, np.complex64, align=align) + xf128 = _aligned_zeros(3, np.longdouble, align=align) + + # test casting, both to and from misaligned + with suppress_warnings() as sup: + sup.filter(ComplexWarning, "Casting complex values") + xc64.astype('f8') + xf64.astype(np.complex64) + test = xc64 + xf64 + + xf128.astype('f8') + xf64.astype(np.longdouble) + test = xf128 + xf64 + + test = xf128 + xc64 + + # test copy, both to and from misaligned + # contig copy + xf64[:] = xf64.copy() + xc64[:] = xc64.copy() + xf128[:] = xf128.copy() + # strided copy + xf64[::2] = xf64[::2].copy() + xc64[::2] = xc64[::2].copy() + xf128[::2] = xf128[::2].copy() + +def test_getfield(): + a = np.arange(32, dtype='uint16') + if sys.byteorder == 'little': + i = 0 + j = 1 + else: + i = 1 + j = 0 + b = a.getfield('int8', i) + assert_equal(b, a) + b = a.getfield('int8', j) + assert_equal(b, 0) + pytest.raises(ValueError, a.getfield, 'uint8', -1) + pytest.raises(ValueError, a.getfield, 'uint8', 16) + pytest.raises(ValueError, a.getfield, 'uint64', 0) + + +class TestViewDtype: + """ + Verify that making a view of a non-contiguous array works as expected. + """ + def test_smaller_dtype_multiple(self): + # x is non-contiguous + x = np.arange(10, dtype=' rc_a) + assert_(sys.getrefcount(dt) > rc_dt) + # del 'it' + it = None + assert_equal(sys.getrefcount(a), rc_a) + assert_equal(sys.getrefcount(dt), rc_dt) + + # With a copy + a = arange(6, dtype='f4') + dt = np.dtype('f4') + rc_a = sys.getrefcount(a) + rc_dt = sys.getrefcount(dt) + it = nditer(a, [], + [['readwrite']], + op_dtypes=[dt]) + rc2_a = sys.getrefcount(a) + rc2_dt = sys.getrefcount(dt) + it2 = it.copy() + assert_(sys.getrefcount(a) > rc2_a) + if sys.version_info < (3, 13): + # np.dtype('f4') is immortal after Python 3.13 + assert_(sys.getrefcount(dt) > rc2_dt) + it = None + assert_equal(sys.getrefcount(a), rc2_a) + assert_equal(sys.getrefcount(dt), rc2_dt) + it2 = None + assert_equal(sys.getrefcount(a), rc_a) + assert_equal(sys.getrefcount(dt), rc_dt) + + del it2 # avoid pyflakes unused variable warning + +def test_iter_best_order(): + # The iterator should always find the iteration order + # with increasing memory addresses + + # Test the ordering for 1-D to 5-D shapes + for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: + a = arange(np.prod(shape)) + # Test each combination of positive and negative strides + for dirs in range(2**len(shape)): + dirs_index = [slice(None)]*len(shape) + for bit in range(len(shape)): + if ((2**bit) & dirs): + dirs_index[bit] = slice(None, None, -1) + dirs_index = tuple(dirs_index) + + aview = a.reshape(shape)[dirs_index] + # C-order + i = nditer(aview, [], [['readonly']]) + assert_equal(list(i), a) + # Fortran-order + i = nditer(aview.T, [], [['readonly']]) + assert_equal(list(i), a) + # Other order + if len(shape) > 2: + i = nditer(aview.swapaxes(0, 1), [], [['readonly']]) + assert_equal(list(i), a) + +def test_iter_c_order(): + # Test forcing C order + + # Test the ordering for 1-D to 5-D shapes + for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: + a = arange(np.prod(shape)) + # Test each combination of positive and negative strides + for dirs in range(2**len(shape)): + dirs_index = [slice(None)]*len(shape) + for bit in range(len(shape)): + if ((2**bit) & dirs): + dirs_index[bit] = slice(None, None, -1) + dirs_index = tuple(dirs_index) + + aview = a.reshape(shape)[dirs_index] + # C-order + i = nditer(aview, order='C') + assert_equal(list(i), aview.ravel(order='C')) + # Fortran-order + i = nditer(aview.T, order='C') + assert_equal(list(i), aview.T.ravel(order='C')) + # Other order + if len(shape) > 2: + i = nditer(aview.swapaxes(0, 1), order='C') + assert_equal(list(i), + aview.swapaxes(0, 1).ravel(order='C')) + +def test_iter_f_order(): + # Test forcing F order + + # Test the ordering for 1-D to 5-D shapes + for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: + a = arange(np.prod(shape)) + # Test each combination of positive and negative strides + for dirs in range(2**len(shape)): + dirs_index = [slice(None)]*len(shape) + for bit in range(len(shape)): + if ((2**bit) & dirs): + dirs_index[bit] = slice(None, None, -1) + dirs_index = tuple(dirs_index) + + aview = a.reshape(shape)[dirs_index] + # C-order + i = nditer(aview, order='F') + assert_equal(list(i), aview.ravel(order='F')) + # Fortran-order + i = nditer(aview.T, order='F') + assert_equal(list(i), aview.T.ravel(order='F')) + # Other order + if len(shape) > 2: + i = nditer(aview.swapaxes(0, 1), order='F') + assert_equal(list(i), + aview.swapaxes(0, 1).ravel(order='F')) + +def test_iter_c_or_f_order(): + # Test forcing any contiguous (C or F) order + + # Test the ordering for 1-D to 5-D shapes + for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: + a = arange(np.prod(shape)) + # Test each combination of positive and negative strides + for dirs in range(2**len(shape)): + dirs_index = [slice(None)]*len(shape) + for bit in range(len(shape)): + if ((2**bit) & dirs): + dirs_index[bit] = slice(None, None, -1) + dirs_index = tuple(dirs_index) + + aview = a.reshape(shape)[dirs_index] + # C-order + i = nditer(aview, order='A') + assert_equal(list(i), aview.ravel(order='A')) + # Fortran-order + i = nditer(aview.T, order='A') + assert_equal(list(i), aview.T.ravel(order='A')) + # Other order + if len(shape) > 2: + i = nditer(aview.swapaxes(0, 1), order='A') + assert_equal(list(i), + aview.swapaxes(0, 1).ravel(order='A')) + +def test_nditer_multi_index_set(): + # Test the multi_index set + a = np.arange(6).reshape(2, 3) + it = np.nditer(a, flags=['multi_index']) + + # Removes the iteration on two first elements of a[0] + it.multi_index = (0, 2,) + + assert_equal(list(it), [2, 3, 4, 5]) + +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +def test_nditer_multi_index_set_refcount(): + # Test if the reference count on index variable is decreased + + index = 0 + i = np.nditer(np.array([111, 222, 333, 444]), flags=['multi_index']) + + start_count = sys.getrefcount(index) + i.multi_index = (index,) + end_count = sys.getrefcount(index) + + assert_equal(start_count, end_count) + +def test_iter_best_order_multi_index_1d(): + # The multi-indices should be correct with any reordering + + a = arange(4) + # 1D order + i = nditer(a, ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(0,), (1,), (2,), (3,)]) + # 1D reversed order + i = nditer(a[::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(3,), (2,), (1,), (0,)]) + +def test_iter_best_order_multi_index_2d(): + # The multi-indices should be correct with any reordering + + a = arange(6) + # 2D C-order + i = nditer(a.reshape(2, 3), ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]) + # 2D Fortran-order + i = nditer(a.reshape(2, 3).copy(order='F'), ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)]) + # 2D reversed C-order + i = nditer(a.reshape(2, 3)[::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(1, 0), (1, 1), (1, 2), (0, 0), (0, 1), (0, 2)]) + i = nditer(a.reshape(2, 3)[:, ::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(0, 2), (0, 1), (0, 0), (1, 2), (1, 1), (1, 0)]) + i = nditer(a.reshape(2, 3)[::-1, ::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(1, 2), (1, 1), (1, 0), (0, 2), (0, 1), (0, 0)]) + # 2D reversed Fortran-order + i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(1, 0), (0, 0), (1, 1), (0, 1), (1, 2), (0, 2)]) + i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1], + ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(0, 2), (1, 2), (0, 1), (1, 1), (0, 0), (1, 0)]) + i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1], + ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), [(1, 2), (0, 2), (1, 1), (0, 1), (1, 0), (0, 0)]) + +def test_iter_best_order_multi_index_3d(): + # The multi-indices should be correct with any reordering + + a = arange(12) + # 3D C-order + i = nditer(a.reshape(2, 3, 2), ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1), + (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1)]) + # 3D Fortran-order + i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0), + (0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1)]) + # 3D reversed C-order + i = nditer(a.reshape(2, 3, 2)[::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1), + (0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1)]) + i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(0, 2, 0), (0, 2, 1), (0, 1, 0), (0, 1, 1), (0, 0, 0), (0, 0, 1), + (1, 2, 0), (1, 2, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)]) + i = nditer(a.reshape(2, 3, 2)[:, :, ::-1], ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(0, 0, 1), (0, 0, 0), (0, 1, 1), (0, 1, 0), (0, 2, 1), (0, 2, 0), + (1, 0, 1), (1, 0, 0), (1, 1, 1), (1, 1, 0), (1, 2, 1), (1, 2, 0)]) + # 3D reversed Fortran-order + i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1], + ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 2, 0), (0, 2, 0), + (1, 0, 1), (0, 0, 1), (1, 1, 1), (0, 1, 1), (1, 2, 1), (0, 2, 1)]) + i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1], + ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(0, 2, 0), (1, 2, 0), (0, 1, 0), (1, 1, 0), (0, 0, 0), (1, 0, 0), + (0, 2, 1), (1, 2, 1), (0, 1, 1), (1, 1, 1), (0, 0, 1), (1, 0, 1)]) + i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, :, ::-1], + ['multi_index'], [['readonly']]) + assert_equal(iter_multi_index(i), + [(0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1), + (0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0)]) + +def test_iter_best_order_c_index_1d(): + # The C index should be correct with any reordering + + a = arange(4) + # 1D order + i = nditer(a, ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [0, 1, 2, 3]) + # 1D reversed order + i = nditer(a[::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [3, 2, 1, 0]) + +def test_iter_best_order_c_index_2d(): + # The C index should be correct with any reordering + + a = arange(6) + # 2D C-order + i = nditer(a.reshape(2, 3), ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5]) + # 2D Fortran-order + i = nditer(a.reshape(2, 3).copy(order='F'), + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [0, 3, 1, 4, 2, 5]) + # 2D reversed C-order + i = nditer(a.reshape(2, 3)[::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [3, 4, 5, 0, 1, 2]) + i = nditer(a.reshape(2, 3)[:, ::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [2, 1, 0, 5, 4, 3]) + i = nditer(a.reshape(2, 3)[::-1, ::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0]) + # 2D reversed Fortran-order + i = nditer(a.reshape(2, 3).copy(order='F')[::-1], + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [3, 0, 4, 1, 5, 2]) + i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1], + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [2, 5, 1, 4, 0, 3]) + i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1], + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), [5, 2, 4, 1, 3, 0]) + +def test_iter_best_order_c_index_3d(): + # The C index should be correct with any reordering + + a = arange(12) + # 3D C-order + i = nditer(a.reshape(2, 3, 2), ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + # 3D Fortran-order + i = nditer(a.reshape(2, 3, 2).copy(order='F'), + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11]) + # 3D reversed C-order + i = nditer(a.reshape(2, 3, 2)[::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5]) + i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7]) + i = nditer(a.reshape(2, 3, 2)[:, :, ::-1], ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10]) + # 3D reversed Fortran-order + i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1], + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5]) + i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1], + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7]) + i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, :, ::-1], + ['c_index'], [['readonly']]) + assert_equal(iter_indices(i), + [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10]) + +def test_iter_best_order_f_index_1d(): + # The Fortran index should be correct with any reordering + + a = arange(4) + # 1D order + i = nditer(a, ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [0, 1, 2, 3]) + # 1D reversed order + i = nditer(a[::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [3, 2, 1, 0]) + +def test_iter_best_order_f_index_2d(): + # The Fortran index should be correct with any reordering + + a = arange(6) + # 2D C-order + i = nditer(a.reshape(2, 3), ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [0, 2, 4, 1, 3, 5]) + # 2D Fortran-order + i = nditer(a.reshape(2, 3).copy(order='F'), + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5]) + # 2D reversed C-order + i = nditer(a.reshape(2, 3)[::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [1, 3, 5, 0, 2, 4]) + i = nditer(a.reshape(2, 3)[:, ::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [4, 2, 0, 5, 3, 1]) + i = nditer(a.reshape(2, 3)[::-1, ::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [5, 3, 1, 4, 2, 0]) + # 2D reversed Fortran-order + i = nditer(a.reshape(2, 3).copy(order='F')[::-1], + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4]) + i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1], + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1]) + i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1], + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0]) + +def test_iter_best_order_f_index_3d(): + # The Fortran index should be correct with any reordering + + a = arange(12) + # 3D C-order + i = nditer(a.reshape(2, 3, 2), ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11]) + # 3D Fortran-order + i = nditer(a.reshape(2, 3, 2).copy(order='F'), + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + # 3D reversed C-order + i = nditer(a.reshape(2, 3, 2)[::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10]) + i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7]) + i = nditer(a.reshape(2, 3, 2)[:, :, ::-1], ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5]) + # 3D reversed Fortran-order + i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1], + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10]) + i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1], + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7]) + i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, :, ::-1], + ['f_index'], [['readonly']]) + assert_equal(iter_indices(i), + [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5]) + +def test_iter_no_inner_full_coalesce(): + # Check no_inner iterators which coalesce into a single inner loop + + for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: + size = np.prod(shape) + a = arange(size) + # Test each combination of forward and backwards indexing + for dirs in range(2**len(shape)): + dirs_index = [slice(None)]*len(shape) + for bit in range(len(shape)): + if ((2**bit) & dirs): + dirs_index[bit] = slice(None, None, -1) + dirs_index = tuple(dirs_index) + + aview = a.reshape(shape)[dirs_index] + # C-order + i = nditer(aview, ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 1) + assert_equal(i[0].shape, (size,)) + # Fortran-order + i = nditer(aview.T, ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 1) + assert_equal(i[0].shape, (size,)) + # Other order + if len(shape) > 2: + i = nditer(aview.swapaxes(0, 1), + ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 1) + assert_equal(i[0].shape, (size,)) + +def test_iter_no_inner_dim_coalescing(): + # Check no_inner iterators whose dimensions may not coalesce completely + + # Skipping the last element in a dimension prevents coalescing + # with the next-bigger dimension + a = arange(24).reshape(2, 3, 4)[:, :, :-1] + i = nditer(a, ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 2) + assert_equal(i[0].shape, (3,)) + a = arange(24).reshape(2, 3, 4)[:, :-1, :] + i = nditer(a, ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 2) + assert_equal(i[0].shape, (8,)) + a = arange(24).reshape(2, 3, 4)[:-1, :, :] + i = nditer(a, ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 1) + assert_equal(i[0].shape, (12,)) + + # Even with lots of 1-sized dimensions, should still coalesce + a = arange(24).reshape(1, 1, 2, 1, 1, 3, 1, 1, 4, 1, 1) + i = nditer(a, ['external_loop'], [['readonly']]) + assert_equal(i.ndim, 1) + assert_equal(i[0].shape, (24,)) + +def test_iter_dim_coalescing(): + # Check that the correct number of dimensions are coalesced + + # Tracking a multi-index disables coalescing + a = arange(24).reshape(2, 3, 4) + i = nditer(a, ['multi_index'], [['readonly']]) + assert_equal(i.ndim, 3) + + # A tracked index can allow coalescing if it's compatible with the array + a3d = arange(24).reshape(2, 3, 4) + i = nditer(a3d, ['c_index'], [['readonly']]) + assert_equal(i.ndim, 1) + i = nditer(a3d.swapaxes(0, 1), ['c_index'], [['readonly']]) + assert_equal(i.ndim, 3) + i = nditer(a3d.T, ['c_index'], [['readonly']]) + assert_equal(i.ndim, 3) + i = nditer(a3d.T, ['f_index'], [['readonly']]) + assert_equal(i.ndim, 1) + i = nditer(a3d.T.swapaxes(0, 1), ['f_index'], [['readonly']]) + assert_equal(i.ndim, 3) + + # When C or F order is forced, coalescing may still occur + a3d = arange(24).reshape(2, 3, 4) + i = nditer(a3d, order='C') + assert_equal(i.ndim, 1) + i = nditer(a3d.T, order='C') + assert_equal(i.ndim, 3) + i = nditer(a3d, order='F') + assert_equal(i.ndim, 3) + i = nditer(a3d.T, order='F') + assert_equal(i.ndim, 1) + i = nditer(a3d, order='A') + assert_equal(i.ndim, 1) + i = nditer(a3d.T, order='A') + assert_equal(i.ndim, 1) + +def test_iter_broadcasting(): + # Standard NumPy broadcasting rules + + # 1D with scalar + i = nditer([arange(6), np.int32(2)], ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 6) + assert_equal(i.shape, (6,)) + + # 2D with scalar + i = nditer([arange(6).reshape(2, 3), np.int32(2)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 6) + assert_equal(i.shape, (2, 3)) + # 2D with 1D + i = nditer([arange(6).reshape(2, 3), arange(3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 6) + assert_equal(i.shape, (2, 3)) + i = nditer([arange(2).reshape(2, 1), arange(3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 6) + assert_equal(i.shape, (2, 3)) + # 2D with 2D + i = nditer([arange(2).reshape(2, 1), arange(3).reshape(1, 3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 6) + assert_equal(i.shape, (2, 3)) + + # 3D with scalar + i = nditer([np.int32(2), arange(24).reshape(4, 2, 3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + # 3D with 1D + i = nditer([arange(3), arange(24).reshape(4, 2, 3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + i = nditer([arange(3), arange(8).reshape(4, 2, 1)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + # 3D with 2D + i = nditer([arange(6).reshape(2, 3), arange(24).reshape(4, 2, 3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + i = nditer([arange(2).reshape(2, 1), arange(24).reshape(4, 2, 3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + i = nditer([arange(3).reshape(1, 3), arange(8).reshape(4, 2, 1)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + # 3D with 3D + i = nditer([arange(2).reshape(1, 2, 1), arange(3).reshape(1, 1, 3), + arange(4).reshape(4, 1, 1)], + ['multi_index'], [['readonly']]*3) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + i = nditer([arange(6).reshape(1, 2, 3), arange(4).reshape(4, 1, 1)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + i = nditer([arange(24).reshape(4, 2, 3), arange(12).reshape(4, 1, 3)], + ['multi_index'], [['readonly']]*2) + assert_equal(i.itersize, 24) + assert_equal(i.shape, (4, 2, 3)) + +def test_iter_itershape(): + # Check that allocated outputs work with a specified shape + a = np.arange(6, dtype='i2').reshape(2, 3) + i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], + op_axes=[[0, 1, None], None], + itershape=(-1, -1, 4)) + assert_equal(i.operands[1].shape, (2, 3, 4)) + assert_equal(i.operands[1].strides, (24, 8, 2)) + + i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']], + op_axes=[[0, 1, None], None], + itershape=(-1, -1, 4)) + assert_equal(i.operands[1].shape, (3, 2, 4)) + assert_equal(i.operands[1].strides, (8, 24, 2)) + + i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']], + order='F', + op_axes=[[0, 1, None], None], + itershape=(-1, -1, 4)) + assert_equal(i.operands[1].shape, (3, 2, 4)) + assert_equal(i.operands[1].strides, (2, 6, 12)) + + # If we specify 1 in the itershape, it shouldn't allow broadcasting + # of that dimension to a bigger value + assert_raises(ValueError, nditer, [a, None], [], + [['readonly'], ['writeonly', 'allocate']], + op_axes=[[0, 1, None], None], + itershape=(-1, 1, 4)) + # Test bug that for no op_axes but itershape, they are NULLed correctly + i = np.nditer([np.ones(2), None, None], itershape=(2,)) + +def test_iter_broadcasting_errors(): + # Check that errors are thrown for bad broadcasting shapes + + # 1D with 1D + assert_raises(ValueError, nditer, [arange(2), arange(3)], + [], [['readonly']]*2) + # 2D with 1D + assert_raises(ValueError, nditer, + [arange(6).reshape(2, 3), arange(2)], + [], [['readonly']]*2) + # 2D with 2D + assert_raises(ValueError, nditer, + [arange(6).reshape(2, 3), arange(9).reshape(3, 3)], + [], [['readonly']]*2) + assert_raises(ValueError, nditer, + [arange(6).reshape(2, 3), arange(4).reshape(2, 2)], + [], [['readonly']]*2) + # 3D with 3D + assert_raises(ValueError, nditer, + [arange(36).reshape(3, 3, 4), arange(24).reshape(2, 3, 4)], + [], [['readonly']]*2) + assert_raises(ValueError, nditer, + [arange(8).reshape(2, 4, 1), arange(24).reshape(2, 3, 4)], + [], [['readonly']]*2) + + # Verify that the error message mentions the right shapes + try: + nditer([arange(2).reshape(1, 2, 1), + arange(3).reshape(1, 3), + arange(6).reshape(2, 3)], + [], + [['readonly'], ['readonly'], ['writeonly', 'no_broadcast']]) + raise AssertionError('Should have raised a broadcast error') + except ValueError as e: + msg = str(e) + # The message should contain the shape of the 3rd operand + assert_(msg.find('(2,3)') >= 0, + 'Message "%s" doesn\'t contain operand shape (2,3)' % msg) + # The message should contain the broadcast shape + assert_(msg.find('(1,2,3)') >= 0, + 'Message "%s" doesn\'t contain broadcast shape (1,2,3)' % msg) + + try: + nditer([arange(6).reshape(2, 3), arange(2)], + [], + [['readonly'], ['readonly']], + op_axes=[[0, 1], [0, np.newaxis]], + itershape=(4, 3)) + raise AssertionError('Should have raised a broadcast error') + except ValueError as e: + msg = str(e) + # The message should contain "shape->remappedshape" for each operand + assert_(msg.find('(2,3)->(2,3)') >= 0, + 'Message "%s" doesn\'t contain operand shape (2,3)->(2,3)' % msg) + assert_(msg.find('(2,)->(2,newaxis)') >= 0, + ('Message "%s" doesn\'t contain remapped operand shape' + + '(2,)->(2,newaxis)') % msg) + # The message should contain the itershape parameter + assert_(msg.find('(4,3)') >= 0, + 'Message "%s" doesn\'t contain itershape parameter (4,3)' % msg) + + try: + nditer([np.zeros((2, 1, 1)), np.zeros((2,))], + [], + [['writeonly', 'no_broadcast'], ['readonly']]) + raise AssertionError('Should have raised a broadcast error') + except ValueError as e: + msg = str(e) + # The message should contain the shape of the bad operand + assert_(msg.find('(2,1,1)') >= 0, + 'Message "%s" doesn\'t contain operand shape (2,1,1)' % msg) + # The message should contain the broadcast shape + assert_(msg.find('(2,1,2)') >= 0, + 'Message "%s" doesn\'t contain the broadcast shape (2,1,2)' % msg) + +def test_iter_flags_errors(): + # Check that bad combinations of flags produce errors + + a = arange(6) + + # Not enough operands + assert_raises(ValueError, nditer, [], [], []) + # Too many operands + assert_raises(ValueError, nditer, [a]*100, [], [['readonly']]*100) + # Bad global flag + assert_raises(ValueError, nditer, [a], ['bad flag'], [['readonly']]) + # Bad op flag + assert_raises(ValueError, nditer, [a], [], [['readonly', 'bad flag']]) + # Bad order parameter + assert_raises(ValueError, nditer, [a], [], [['readonly']], order='G') + # Bad casting parameter + assert_raises(ValueError, nditer, [a], [], [['readonly']], casting='noon') + # op_flags must match ops + assert_raises(ValueError, nditer, [a]*3, [], [['readonly']]*2) + # Cannot track both a C and an F index + assert_raises(ValueError, nditer, a, + ['c_index', 'f_index'], [['readonly']]) + # Inner iteration and multi-indices/indices are incompatible + assert_raises(ValueError, nditer, a, + ['external_loop', 'multi_index'], [['readonly']]) + assert_raises(ValueError, nditer, a, + ['external_loop', 'c_index'], [['readonly']]) + assert_raises(ValueError, nditer, a, + ['external_loop', 'f_index'], [['readonly']]) + # Must specify exactly one of readwrite/readonly/writeonly per operand + assert_raises(ValueError, nditer, a, [], [[]]) + assert_raises(ValueError, nditer, a, [], [['readonly', 'writeonly']]) + assert_raises(ValueError, nditer, a, [], [['readonly', 'readwrite']]) + assert_raises(ValueError, nditer, a, [], [['writeonly', 'readwrite']]) + assert_raises(ValueError, nditer, a, + [], [['readonly', 'writeonly', 'readwrite']]) + # Python scalars are always readonly + assert_raises(TypeError, nditer, 1.5, [], [['writeonly']]) + assert_raises(TypeError, nditer, 1.5, [], [['readwrite']]) + # Array scalars are always readonly + assert_raises(TypeError, nditer, np.int32(1), [], [['writeonly']]) + assert_raises(TypeError, nditer, np.int32(1), [], [['readwrite']]) + # Check readonly array + a.flags.writeable = False + assert_raises(ValueError, nditer, a, [], [['writeonly']]) + assert_raises(ValueError, nditer, a, [], [['readwrite']]) + a.flags.writeable = True + # Multi-indices available only with the multi_index flag + i = nditer(arange(6), [], [['readonly']]) + assert_raises(ValueError, lambda i: i.multi_index, i) + # Index available only with an index flag + assert_raises(ValueError, lambda i: i.index, i) + # GotoCoords and GotoIndex incompatible with buffering or no_inner + + def assign_multi_index(i): + i.multi_index = (0,) + + def assign_index(i): + i.index = 0 + + def assign_iterindex(i): + i.iterindex = 0 + + def assign_iterrange(i): + i.iterrange = (0, 1) + i = nditer(arange(6), ['external_loop']) + assert_raises(ValueError, assign_multi_index, i) + assert_raises(ValueError, assign_index, i) + assert_raises(ValueError, assign_iterindex, i) + assert_raises(ValueError, assign_iterrange, i) + i = nditer(arange(6), ['buffered']) + assert_raises(ValueError, assign_multi_index, i) + assert_raises(ValueError, assign_index, i) + assert_raises(ValueError, assign_iterrange, i) + # Can't iterate if size is zero + assert_raises(ValueError, nditer, np.array([])) + +def test_iter_slice(): + a, b, c = np.arange(3), np.arange(3), np.arange(3.) + i = nditer([a, b, c], [], ['readwrite']) + with i: + i[0:2] = (3, 3) + assert_equal(a, [3, 1, 2]) + assert_equal(b, [3, 1, 2]) + assert_equal(c, [0, 1, 2]) + i[1] = 12 + assert_equal(i[0:2], [3, 12]) + +def test_iter_assign_mapping(): + a = np.arange(24, dtype='f8').reshape(2, 3, 4).T + it = np.nditer(a, [], [['readwrite', 'updateifcopy']], + casting='same_kind', op_dtypes=[np.dtype('f4')]) + with it: + it.operands[0][...] = 3 + it.operands[0][...] = 14 + assert_equal(a, 14) + it = np.nditer(a, [], [['readwrite', 'updateifcopy']], + casting='same_kind', op_dtypes=[np.dtype('f4')]) + with it: + x = it.operands[0][-1:1] + x[...] = 14 + it.operands[0][...] = -1234 + assert_equal(a, -1234) + # check for no warnings on dealloc + x = None + it = None + +def test_iter_nbo_align_contig(): + # Check that byte order, alignment, and contig changes work + + # Byte order change by requesting a specific dtype + a = np.arange(6, dtype='f4') + au = a.byteswap() + au = au.view(au.dtype.newbyteorder()) + assert_(a.dtype.byteorder != au.dtype.byteorder) + i = nditer(au, [], [['readwrite', 'updateifcopy']], + casting='equiv', + op_dtypes=[np.dtype('f4')]) + with i: + # context manager triggers WRITEBACKIFCOPY on i at exit + assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder) + assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder) + assert_equal(i.operands[0], a) + i.operands[0][:] = 2 + assert_equal(au, [2]*6) + del i # should not raise a warning + # Byte order change by requesting NBO + a = np.arange(6, dtype='f4') + au = a.byteswap() + au = au.view(au.dtype.newbyteorder()) + assert_(a.dtype.byteorder != au.dtype.byteorder) + with nditer(au, [], [['readwrite', 'updateifcopy', 'nbo']], + casting='equiv') as i: + # context manager triggers UPDATEIFCOPY on i at exit + assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder) + assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder) + assert_equal(i.operands[0], a) + i.operands[0][:] = 12345 + i.operands[0][:] = 2 + assert_equal(au, [2]*6) + + # Unaligned input + a = np.zeros((6*4+1,), dtype='i1')[1:] + a.dtype = 'f4' + a[:] = np.arange(6, dtype='f4') + assert_(not a.flags.aligned) + # Without 'aligned', shouldn't copy + i = nditer(a, [], [['readonly']]) + assert_(not i.operands[0].flags.aligned) + assert_equal(i.operands[0], a) + # With 'aligned', should make a copy + with nditer(a, [], [['readwrite', 'updateifcopy', 'aligned']]) as i: + assert_(i.operands[0].flags.aligned) + # context manager triggers UPDATEIFCOPY on i at exit + assert_equal(i.operands[0], a) + i.operands[0][:] = 3 + assert_equal(a, [3]*6) + + # Discontiguous input + a = arange(12) + # If it is contiguous, shouldn't copy + i = nditer(a[:6], [], [['readonly']]) + assert_(i.operands[0].flags.contiguous) + assert_equal(i.operands[0], a[:6]) + # If it isn't contiguous, should buffer + i = nditer(a[::2], ['buffered', 'external_loop'], + [['readonly', 'contig']], + buffersize=10) + assert_(i[0].flags.contiguous) + assert_equal(i[0], a[::2]) + +def test_iter_array_cast(): + # Check that arrays are cast as requested + + # No cast 'f4' -> 'f4' + a = np.arange(6, dtype='f4').reshape(2, 3) + i = nditer(a, [], [['readwrite']], op_dtypes=[np.dtype('f4')]) + with i: + assert_equal(i.operands[0], a) + assert_equal(i.operands[0].dtype, np.dtype('f4')) + + # Byte-order cast ' '>f4' + a = np.arange(6, dtype='f4')]) as i: + assert_equal(i.operands[0], a) + assert_equal(i.operands[0].dtype, np.dtype('>f4')) + + # Safe case 'f4' -> 'f8' + a = np.arange(24, dtype='f4').reshape(2, 3, 4).swapaxes(1, 2) + i = nditer(a, [], [['readonly', 'copy']], + casting='safe', + op_dtypes=[np.dtype('f8')]) + assert_equal(i.operands[0], a) + assert_equal(i.operands[0].dtype, np.dtype('f8')) + # The memory layout of the temporary should match a (a is (48,4,16)) + # except negative strides get flipped to positive strides. + assert_equal(i.operands[0].strides, (96, 8, 32)) + a = a[::-1, :, ::-1] + i = nditer(a, [], [['readonly', 'copy']], + casting='safe', + op_dtypes=[np.dtype('f8')]) + assert_equal(i.operands[0], a) + assert_equal(i.operands[0].dtype, np.dtype('f8')) + assert_equal(i.operands[0].strides, (96, 8, 32)) + + # Same-kind cast 'f8' -> 'f4' -> 'f8' + a = np.arange(24, dtype='f8').reshape(2, 3, 4).T + with nditer(a, [], + [['readwrite', 'updateifcopy']], + casting='same_kind', + op_dtypes=[np.dtype('f4')]) as i: + assert_equal(i.operands[0], a) + assert_equal(i.operands[0].dtype, np.dtype('f4')) + assert_equal(i.operands[0].strides, (4, 16, 48)) + # Check that WRITEBACKIFCOPY is activated at exit + i.operands[0][2, 1, 1] = -12.5 + assert_(a[2, 1, 1] != -12.5) + assert_equal(a[2, 1, 1], -12.5) + + a = np.arange(6, dtype='i4')[::-2] + with nditer(a, [], + [['writeonly', 'updateifcopy']], + casting='unsafe', + op_dtypes=[np.dtype('f4')]) as i: + assert_equal(i.operands[0].dtype, np.dtype('f4')) + # Even though the stride was negative in 'a', it + # becomes positive in the temporary + assert_equal(i.operands[0].strides, (4,)) + i.operands[0][:] = [1, 2, 3] + assert_equal(a, [1, 2, 3]) + +def test_iter_array_cast_errors(): + # Check that invalid casts are caught + + # Need to enable copying for casts to occur + assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], + [['readonly']], op_dtypes=[np.dtype('f8')]) + # Also need to allow casting for casts to occur + assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], + [['readonly', 'copy']], casting='no', + op_dtypes=[np.dtype('f8')]) + assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], + [['readonly', 'copy']], casting='equiv', + op_dtypes=[np.dtype('f8')]) + assert_raises(TypeError, nditer, arange(2, dtype='f8'), [], + [['writeonly', 'updateifcopy']], + casting='no', + op_dtypes=[np.dtype('f4')]) + assert_raises(TypeError, nditer, arange(2, dtype='f8'), [], + [['writeonly', 'updateifcopy']], + casting='equiv', + op_dtypes=[np.dtype('f4')]) + # ' '>f4' should not work with casting='no' + assert_raises(TypeError, nditer, arange(2, dtype='f4')]) + # 'f4' -> 'f8' is a safe cast, but 'f8' -> 'f4' isn't + assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], + [['readwrite', 'updateifcopy']], + casting='safe', + op_dtypes=[np.dtype('f8')]) + assert_raises(TypeError, nditer, arange(2, dtype='f8'), [], + [['readwrite', 'updateifcopy']], + casting='safe', + op_dtypes=[np.dtype('f4')]) + # 'f4' -> 'i4' is neither a safe nor a same-kind cast + assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], + [['readonly', 'copy']], + casting='same_kind', + op_dtypes=[np.dtype('i4')]) + assert_raises(TypeError, nditer, arange(2, dtype='i4'), [], + [['writeonly', 'updateifcopy']], + casting='same_kind', + op_dtypes=[np.dtype('f4')]) + +def test_iter_scalar_cast(): + # Check that scalars are cast as requested + + # No cast 'f4' -> 'f4' + i = nditer(np.float32(2.5), [], [['readonly']], + op_dtypes=[np.dtype('f4')]) + assert_equal(i.dtypes[0], np.dtype('f4')) + assert_equal(i.value.dtype, np.dtype('f4')) + assert_equal(i.value, 2.5) + # Safe cast 'f4' -> 'f8' + i = nditer(np.float32(2.5), [], + [['readonly', 'copy']], + casting='safe', + op_dtypes=[np.dtype('f8')]) + assert_equal(i.dtypes[0], np.dtype('f8')) + assert_equal(i.value.dtype, np.dtype('f8')) + assert_equal(i.value, 2.5) + # Same-kind cast 'f8' -> 'f4' + i = nditer(np.float64(2.5), [], + [['readonly', 'copy']], + casting='same_kind', + op_dtypes=[np.dtype('f4')]) + assert_equal(i.dtypes[0], np.dtype('f4')) + assert_equal(i.value.dtype, np.dtype('f4')) + assert_equal(i.value, 2.5) + # Unsafe cast 'f8' -> 'i4' + i = nditer(np.float64(3.0), [], + [['readonly', 'copy']], + casting='unsafe', + op_dtypes=[np.dtype('i4')]) + assert_equal(i.dtypes[0], np.dtype('i4')) + assert_equal(i.value.dtype, np.dtype('i4')) + assert_equal(i.value, 3) + # Readonly scalars may be cast even without setting COPY or BUFFERED + i = nditer(3, [], [['readonly']], op_dtypes=[np.dtype('f8')]) + assert_equal(i[0].dtype, np.dtype('f8')) + assert_equal(i[0], 3.) + +def test_iter_scalar_cast_errors(): + # Check that invalid casts are caught + + # Need to allow copying/buffering for write casts of scalars to occur + assert_raises(TypeError, nditer, np.float32(2), [], + [['readwrite']], op_dtypes=[np.dtype('f8')]) + assert_raises(TypeError, nditer, 2.5, [], + [['readwrite']], op_dtypes=[np.dtype('f4')]) + # 'f8' -> 'f4' isn't a safe cast if the value would overflow + assert_raises(TypeError, nditer, np.float64(1e60), [], + [['readonly']], + casting='safe', + op_dtypes=[np.dtype('f4')]) + # 'f4' -> 'i4' is neither a safe nor a same-kind cast + assert_raises(TypeError, nditer, np.float32(2), [], + [['readonly']], + casting='same_kind', + op_dtypes=[np.dtype('i4')]) + +def test_iter_object_arrays_basic(): + # Check that object arrays work + + obj = {'a': 3, 'b': 'd'} + a = np.array([[1, 2, 3], None, obj, None], dtype='O') + if HAS_REFCOUNT: + rc = sys.getrefcount(obj) + + # Need to allow references for object arrays + assert_raises(TypeError, nditer, a) + if HAS_REFCOUNT: + assert_equal(sys.getrefcount(obj), rc) + + i = nditer(a, ['refs_ok'], ['readonly']) + vals = [x_[()] for x_ in i] + assert_equal(np.array(vals, dtype='O'), a) + vals, i, x = [None]*3 + if HAS_REFCOUNT: + assert_equal(sys.getrefcount(obj), rc) + + i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'], + ['readonly'], order='C') + assert_(i.iterationneedsapi) + vals = [x_[()] for x_ in i] + assert_equal(np.array(vals, dtype='O'), a.reshape(2, 2).ravel(order='F')) + vals, i, x = [None]*3 + if HAS_REFCOUNT: + assert_equal(sys.getrefcount(obj), rc) + + i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'], + ['readwrite'], order='C') + with i: + for x in i: + x[...] = None + vals, i, x = [None]*3 + if HAS_REFCOUNT: + assert_(sys.getrefcount(obj) == rc-1) + assert_equal(a, np.array([None]*4, dtype='O')) + +def test_iter_object_arrays_conversions(): + # Conversions to/from objects + a = np.arange(6, dtype='O') + i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], + casting='unsafe', op_dtypes='i4') + with i: + for x in i: + x[...] += 1 + assert_equal(a, np.arange(6)+1) + + a = np.arange(6, dtype='i4') + i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], + casting='unsafe', op_dtypes='O') + with i: + for x in i: + x[...] += 1 + assert_equal(a, np.arange(6)+1) + + # Non-contiguous object array + a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'O')]) + a = a['a'] + a[:] = np.arange(6) + i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], + casting='unsafe', op_dtypes='i4') + with i: + for x in i: + x[...] += 1 + assert_equal(a, np.arange(6)+1) + + #Non-contiguous value array + a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'i4')]) + a = a['a'] + a[:] = np.arange(6) + 98172488 + i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], + casting='unsafe', op_dtypes='O') + with i: + ob = i[0][()] + if HAS_REFCOUNT: + rc = sys.getrefcount(ob) + for x in i: + x[...] += 1 + if HAS_REFCOUNT: + assert_(sys.getrefcount(ob) == rc-1) + assert_equal(a, np.arange(6)+98172489) + +def test_iter_common_dtype(): + # Check that the iterator finds a common data type correctly + # (some checks are somewhat duplicate after adopting NEP 50) + + i = nditer([array([3], dtype='f4'), array([0], dtype='f8')], + ['common_dtype'], + [['readonly', 'copy']]*2, + casting='safe') + assert_equal(i.dtypes[0], np.dtype('f8')) + assert_equal(i.dtypes[1], np.dtype('f8')) + i = nditer([array([3], dtype='i4'), array([0], dtype='f4')], + ['common_dtype'], + [['readonly', 'copy']]*2, + casting='safe') + assert_equal(i.dtypes[0], np.dtype('f8')) + assert_equal(i.dtypes[1], np.dtype('f8')) + i = nditer([array([3], dtype='f4'), array(0, dtype='f8')], + ['common_dtype'], + [['readonly', 'copy']]*2, + casting='same_kind') + assert_equal(i.dtypes[0], np.dtype('f8')) + assert_equal(i.dtypes[1], np.dtype('f8')) + i = nditer([array([3], dtype='u4'), array(0, dtype='i4')], + ['common_dtype'], + [['readonly', 'copy']]*2, + casting='safe') + assert_equal(i.dtypes[0], np.dtype('i8')) + assert_equal(i.dtypes[1], np.dtype('i8')) + i = nditer([array([3], dtype='u4'), array(-12, dtype='i4')], + ['common_dtype'], + [['readonly', 'copy']]*2, + casting='safe') + assert_equal(i.dtypes[0], np.dtype('i8')) + assert_equal(i.dtypes[1], np.dtype('i8')) + i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'), + array([2j], dtype='c8'), array([9], dtype='f8')], + ['common_dtype'], + [['readonly', 'copy']]*4, + casting='safe') + assert_equal(i.dtypes[0], np.dtype('c16')) + assert_equal(i.dtypes[1], np.dtype('c16')) + assert_equal(i.dtypes[2], np.dtype('c16')) + assert_equal(i.dtypes[3], np.dtype('c16')) + assert_equal(i.value, (3, -12, 2j, 9)) + + # When allocating outputs, other outputs aren't factored in + i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], [], + [['readonly', 'copy'], + ['writeonly', 'allocate'], + ['writeonly']], + casting='safe') + assert_equal(i.dtypes[0], np.dtype('i4')) + assert_equal(i.dtypes[1], np.dtype('i4')) + assert_equal(i.dtypes[2], np.dtype('c16')) + # But, if common data types are requested, they are + i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], + ['common_dtype'], + [['readonly', 'copy'], + ['writeonly', 'allocate'], + ['writeonly']], + casting='safe') + assert_equal(i.dtypes[0], np.dtype('c16')) + assert_equal(i.dtypes[1], np.dtype('c16')) + assert_equal(i.dtypes[2], np.dtype('c16')) + +def test_iter_copy_if_overlap(): + # Ensure the iterator makes copies on read/write overlap, if requested + + # Copy not needed, 1 op + for flag in ['readonly', 'writeonly', 'readwrite']: + a = arange(10) + i = nditer([a], ['copy_if_overlap'], [[flag]]) + with i: + assert_(i.operands[0] is a) + + # Copy needed, 2 ops, read-write overlap + x = arange(10) + a = x[1:] + b = x[:-1] + with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i: + assert_(not np.shares_memory(*i.operands)) + + # Copy not needed with elementwise, 2 ops, exactly same arrays + x = arange(10) + a = x + b = x + i = nditer([a, b], ['copy_if_overlap'], [['readonly', 'overlap_assume_elementwise'], + ['readwrite', 'overlap_assume_elementwise']]) + with i: + assert_(i.operands[0] is a and i.operands[1] is b) + with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i: + assert_(i.operands[0] is a and not np.shares_memory(i.operands[1], b)) + + # Copy not needed, 2 ops, no overlap + x = arange(10) + a = x[::2] + b = x[1::2] + i = nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']]) + assert_(i.operands[0] is a and i.operands[1] is b) + + # Copy needed, 2 ops, read-write overlap + x = arange(4, dtype=np.int8) + a = x[3:] + b = x.view(np.int32)[:1] + with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']]) as i: + assert_(not np.shares_memory(*i.operands)) + + # Copy needed, 3 ops, read-write overlap + for flag in ['writeonly', 'readwrite']: + x = np.ones([10, 10]) + a = x + b = x.T + c = x + with nditer([a, b, c], ['copy_if_overlap'], + [['readonly'], ['readonly'], [flag]]) as i: + a2, b2, c2 = i.operands + assert_(not np.shares_memory(a2, c2)) + assert_(not np.shares_memory(b2, c2)) + + # Copy not needed, 3 ops, read-only overlap + x = np.ones([10, 10]) + a = x + b = x.T + c = x + i = nditer([a, b, c], ['copy_if_overlap'], + [['readonly'], ['readonly'], ['readonly']]) + a2, b2, c2 = i.operands + assert_(a is a2) + assert_(b is b2) + assert_(c is c2) + + # Copy not needed, 3 ops, read-only overlap + x = np.ones([10, 10]) + a = x + b = np.ones([10, 10]) + c = x.T + i = nditer([a, b, c], ['copy_if_overlap'], + [['readonly'], ['writeonly'], ['readonly']]) + a2, b2, c2 = i.operands + assert_(a is a2) + assert_(b is b2) + assert_(c is c2) + + # Copy not needed, 3 ops, write-only overlap + x = np.arange(7) + a = x[:3] + b = x[3:6] + c = x[4:7] + i = nditer([a, b, c], ['copy_if_overlap'], + [['readonly'], ['writeonly'], ['writeonly']]) + a2, b2, c2 = i.operands + assert_(a is a2) + assert_(b is b2) + assert_(c is c2) + +def test_iter_op_axes(): + # Check that custom axes work + + # Reverse the axes + a = arange(6).reshape(2, 3) + i = nditer([a, a.T], [], [['readonly']]*2, op_axes=[[0, 1], [1, 0]]) + assert_(all([x == y for (x, y) in i])) + a = arange(24).reshape(2, 3, 4) + i = nditer([a.T, a], [], [['readonly']]*2, op_axes=[[2, 1, 0], None]) + assert_(all([x == y for (x, y) in i])) + + # Broadcast 1D to any dimension + a = arange(1, 31).reshape(2, 3, 5) + b = arange(1, 3) + i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [0, -1, -1]]) + assert_equal([x*y for (x, y) in i], (a*b.reshape(2, 1, 1)).ravel()) + b = arange(1, 4) + i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [-1, 0, -1]]) + assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 3, 1)).ravel()) + b = arange(1, 6) + i = nditer([a, b], [], [['readonly']]*2, + op_axes=[None, [np.newaxis, np.newaxis, 0]]) + assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 1, 5)).ravel()) + + # Inner product-style broadcasting + a = arange(24).reshape(2, 3, 4) + b = arange(40).reshape(5, 2, 4) + i = nditer([a, b], ['multi_index'], [['readonly']]*2, + op_axes=[[0, 1, -1, -1], [-1, -1, 0, 1]]) + assert_equal(i.shape, (2, 3, 5, 2)) + + # Matrix product-style broadcasting + a = arange(12).reshape(3, 4) + b = arange(20).reshape(4, 5) + i = nditer([a, b], ['multi_index'], [['readonly']]*2, + op_axes=[[0, -1], [-1, 1]]) + assert_equal(i.shape, (3, 5)) + +def test_iter_op_axes_errors(): + # Check that custom axes throws errors for bad inputs + + # Wrong number of items in op_axes + a = arange(6).reshape(2, 3) + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[0], [1], [0]]) + # Out of bounds items in op_axes + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[2, 1], [0, 1]]) + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[0, 1], [2, -1]]) + # Duplicate items in op_axes + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[0, 0], [0, 1]]) + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[0, 1], [1, 1]]) + + # Different sized arrays in op_axes + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[0, 1], [0, 1, 0]]) + + # Non-broadcastable dimensions in the result + assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, + op_axes=[[0, 1], [1, 0]]) + +def test_iter_copy(): + # Check that copying the iterator works correctly + a = arange(24).reshape(2, 3, 4) + + # Simple iterator + i = nditer(a) + j = i.copy() + assert_equal([x[()] for x in i], [x[()] for x in j]) + + i.iterindex = 3 + j = i.copy() + assert_equal([x[()] for x in i], [x[()] for x in j]) + + # Buffered iterator + i = nditer(a, ['buffered', 'ranged'], order='F', buffersize=3) + j = i.copy() + assert_equal([x[()] for x in i], [x[()] for x in j]) + + i.iterindex = 3 + j = i.copy() + assert_equal([x[()] for x in i], [x[()] for x in j]) + + i.iterrange = (3, 9) + j = i.copy() + assert_equal([x[()] for x in i], [x[()] for x in j]) + + i.iterrange = (2, 18) + next(i) + next(i) + j = i.copy() + assert_equal([x[()] for x in i], [x[()] for x in j]) + + # Casting iterator + with nditer(a, ['buffered'], order='F', casting='unsafe', + op_dtypes='f8', buffersize=5) as i: + j = i.copy() + assert_equal([x[()] for x in j], a.ravel(order='F')) + + a = arange(24, dtype=' unstructured (any to object), and many other + # casts, which cause this to require all steps in the casting machinery + # one level down as well as the iterator copy (which uses NpyAuxData clone) + in_dtype = np.dtype([("a", np.dtype("i,")), + ("b", np.dtype(">i,d,S17,>d,3f,O,i1"))]) + out_dtype = np.dtype([("a", np.dtype("O")), + ("b", np.dtype(">i,>i,S17,>d,>U3,3d,i1,O"))]) + arr = np.ones(1000, dtype=in_dtype) + + it = np.nditer((arr,), ["buffered", "external_loop", "refs_ok"], + op_dtypes=[out_dtype], casting="unsafe") + it_copy = it.copy() + + res1 = next(it) + del it + res2 = next(it_copy) + del it_copy + + expected = arr["a"].astype(out_dtype["a"]) + assert_array_equal(res1["a"], expected) + assert_array_equal(res2["a"], expected) + + for field in in_dtype["b"].names: + # Note that the .base avoids the subarray field + expected = arr["b"][field].astype(out_dtype["b"][field].base) + assert_array_equal(res1["b"][field], expected) + assert_array_equal(res2["b"][field], expected) + + +def test_iter_copy_casts_structured2(): + # Similar to the above, this is a fairly arcane test to cover internals + in_dtype = np.dtype([("a", np.dtype("O,O")), + ("b", np.dtype("5O,3O,(1,)O,(1,)i,(1,)O"))]) + out_dtype = np.dtype([("a", np.dtype("O")), + ("b", np.dtype("O,3i,4O,4O,4i"))]) + + arr = np.ones(1, dtype=in_dtype) + it = np.nditer((arr,), ["buffered", "external_loop", "refs_ok"], + op_dtypes=[out_dtype], casting="unsafe") + it_copy = it.copy() + + res1 = next(it) + del it + res2 = next(it_copy) + del it_copy + + # Array of two structured scalars: + for res in res1, res2: + # Cast to tuple by getitem, which may be weird and changeable?: + assert type(res["a"][0]) == tuple + assert res["a"][0] == (1, 1) + + for res in res1, res2: + assert_array_equal(res["b"]["f0"][0], np.ones(5, dtype=object)) + assert_array_equal(res["b"]["f1"], np.ones((1, 3), dtype="i")) + assert res["b"]["f2"].shape == (1, 4) + assert_array_equal(res["b"]["f2"][0], np.ones(4, dtype=object)) + assert_array_equal(res["b"]["f3"][0], np.ones(4, dtype=object)) + assert_array_equal(res["b"]["f3"][0], np.ones(4, dtype="i")) + + +def test_iter_allocate_output_simple(): + # Check that the iterator will properly allocate outputs + + # Simple case + a = arange(6) + i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')]) + assert_equal(i.operands[1].shape, a.shape) + assert_equal(i.operands[1].dtype, np.dtype('f4')) + +def test_iter_allocate_output_buffered_readwrite(): + # Allocated output with buffering + delay_bufalloc + + a = arange(6) + i = nditer([a, None], ['buffered', 'delay_bufalloc'], + [['readonly'], ['allocate', 'readwrite']]) + with i: + i.operands[1][:] = 1 + i.reset() + for x in i: + x[1][...] += x[0][...] + assert_equal(i.operands[1], a+1) + +def test_iter_allocate_output_itorder(): + # The allocated output should match the iteration order + + # C-order input, best iteration order + a = arange(6, dtype='i4').reshape(2, 3) + i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')]) + assert_equal(i.operands[1].shape, a.shape) + assert_equal(i.operands[1].strides, a.strides) + assert_equal(i.operands[1].dtype, np.dtype('f4')) + # F-order input, best iteration order + a = arange(24, dtype='i4').reshape(2, 3, 4).T + i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')]) + assert_equal(i.operands[1].shape, a.shape) + assert_equal(i.operands[1].strides, a.strides) + assert_equal(i.operands[1].dtype, np.dtype('f4')) + # Non-contiguous input, C iteration order + a = arange(24, dtype='i4').reshape(2, 3, 4).swapaxes(0, 1) + i = nditer([a, None], [], + [['readonly'], ['writeonly', 'allocate']], + order='C', + op_dtypes=[None, np.dtype('f4')]) + assert_equal(i.operands[1].shape, a.shape) + assert_equal(i.operands[1].strides, (32, 16, 4)) + assert_equal(i.operands[1].dtype, np.dtype('f4')) + +def test_iter_allocate_output_opaxes(): + # Specifying op_axes should work + + a = arange(24, dtype='i4').reshape(2, 3, 4) + i = nditer([None, a], [], [['writeonly', 'allocate'], ['readonly']], + op_dtypes=[np.dtype('u4'), None], + op_axes=[[1, 2, 0], None]) + assert_equal(i.operands[0].shape, (4, 2, 3)) + assert_equal(i.operands[0].strides, (4, 48, 16)) + assert_equal(i.operands[0].dtype, np.dtype('u4')) + +def test_iter_allocate_output_types_promotion(): + # Check type promotion of automatic outputs (this was more interesting + # before NEP 50...) + + i = nditer([array([3], dtype='f4'), array([0], dtype='f8'), None], [], + [['readonly']]*2+[['writeonly', 'allocate']]) + assert_equal(i.dtypes[2], np.dtype('f8')) + i = nditer([array([3], dtype='i4'), array([0], dtype='f4'), None], [], + [['readonly']]*2+[['writeonly', 'allocate']]) + assert_equal(i.dtypes[2], np.dtype('f8')) + i = nditer([array([3], dtype='f4'), array(0, dtype='f8'), None], [], + [['readonly']]*2+[['writeonly', 'allocate']]) + assert_equal(i.dtypes[2], np.dtype('f8')) + i = nditer([array([3], dtype='u4'), array(0, dtype='i4'), None], [], + [['readonly']]*2+[['writeonly', 'allocate']]) + assert_equal(i.dtypes[2], np.dtype('i8')) + i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'), None], [], + [['readonly']]*2+[['writeonly', 'allocate']]) + assert_equal(i.dtypes[2], np.dtype('i8')) + +def test_iter_allocate_output_types_byte_order(): + # Verify the rules for byte order changes + + # When there's just one input, the output type exactly matches + a = array([3], dtype='u4') + a = a.view(a.dtype.newbyteorder()) + i = nditer([a, None], [], + [['readonly'], ['writeonly', 'allocate']]) + assert_equal(i.dtypes[0], i.dtypes[1]) + # With two or more inputs, the output type is in native byte order + i = nditer([a, a, None], [], + [['readonly'], ['readonly'], ['writeonly', 'allocate']]) + assert_(i.dtypes[0] != i.dtypes[2]) + assert_equal(i.dtypes[0].newbyteorder('='), i.dtypes[2]) + +def test_iter_allocate_output_types_scalar(): + # If the inputs are all scalars, the output should be a scalar + + i = nditer([None, 1, 2.3, np.float32(12), np.complex128(3)], [], + [['writeonly', 'allocate']] + [['readonly']]*4) + assert_equal(i.operands[0].dtype, np.dtype('complex128')) + assert_equal(i.operands[0].ndim, 0) + +def test_iter_allocate_output_subtype(): + # Make sure that the subtype with priority wins + class MyNDArray(np.ndarray): + __array_priority__ = 15 + + # subclass vs ndarray + a = np.array([[1, 2], [3, 4]]).view(MyNDArray) + b = np.arange(4).reshape(2, 2).T + i = nditer([a, b, None], [], + [['readonly'], ['readonly'], ['writeonly', 'allocate']]) + assert_equal(type(a), type(i.operands[2])) + assert_(type(b) is not type(i.operands[2])) + assert_equal(i.operands[2].shape, (2, 2)) + + # If subtypes are disabled, we should get back an ndarray. + i = nditer([a, b, None], [], + [['readonly'], ['readonly'], + ['writeonly', 'allocate', 'no_subtype']]) + assert_equal(type(b), type(i.operands[2])) + assert_(type(a) is not type(i.operands[2])) + assert_equal(i.operands[2].shape, (2, 2)) + +def test_iter_allocate_output_errors(): + # Check that the iterator will throw errors for bad output allocations + + # Need an input if no output data type is specified + a = arange(6) + assert_raises(TypeError, nditer, [a, None], [], + [['writeonly'], ['writeonly', 'allocate']]) + # Allocated output should be flagged for writing + assert_raises(ValueError, nditer, [a, None], [], + [['readonly'], ['allocate', 'readonly']]) + # Allocated output can't have buffering without delayed bufalloc + assert_raises(ValueError, nditer, [a, None], ['buffered'], + ['allocate', 'readwrite']) + # Must specify dtype if there are no inputs (cannot promote existing ones; + # maybe this should use the 'f4' here, but it does not historically.) + assert_raises(TypeError, nditer, [None, None], [], + [['writeonly', 'allocate'], + ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')]) + # If using op_axes, must specify all the axes + a = arange(24, dtype='i4').reshape(2, 3, 4) + assert_raises(ValueError, nditer, [a, None], [], + [['readonly'], ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')], + op_axes=[None, [0, np.newaxis, 1]]) + # If using op_axes, the axes must be within bounds + assert_raises(ValueError, nditer, [a, None], [], + [['readonly'], ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')], + op_axes=[None, [0, 3, 1]]) + # If using op_axes, there can't be duplicates + assert_raises(ValueError, nditer, [a, None], [], + [['readonly'], ['writeonly', 'allocate']], + op_dtypes=[None, np.dtype('f4')], + op_axes=[None, [0, 2, 1, 0]]) + # Not all axes may be specified if a reduction. If there is a hole + # in op_axes, this is an error. + a = arange(24, dtype='i4').reshape(2, 3, 4) + assert_raises(ValueError, nditer, [a, None], ["reduce_ok"], + [['readonly'], ['readwrite', 'allocate']], + op_dtypes=[None, np.dtype('f4')], + op_axes=[None, [0, np.newaxis, 2]]) + +def test_all_allocated(): + # When no output and no shape is given, `()` is used as shape. + i = np.nditer([None], op_dtypes=["int64"]) + assert i.operands[0].shape == () + assert i.dtypes == (np.dtype("int64"),) + + i = np.nditer([None], op_dtypes=["int64"], itershape=(2, 3, 4)) + assert i.operands[0].shape == (2, 3, 4) + +def test_iter_remove_axis(): + a = arange(24).reshape(2, 3, 4) + + i = nditer(a, ['multi_index']) + i.remove_axis(1) + assert_equal(list(i), a[:, 0, :].ravel()) + + a = a[::-1, :, :] + i = nditer(a, ['multi_index']) + i.remove_axis(0) + assert_equal(list(i), a[0, :, :].ravel()) + +def test_iter_remove_multi_index_inner_loop(): + # Check that removing multi-index support works + + a = arange(24).reshape(2, 3, 4) + + i = nditer(a, ['multi_index']) + assert_equal(i.ndim, 3) + assert_equal(i.shape, (2, 3, 4)) + assert_equal(i.itviews[0].shape, (2, 3, 4)) + + # Removing the multi-index tracking causes all dimensions to coalesce + before = list(i) + i.remove_multi_index() + after = list(i) + + assert_equal(before, after) + assert_equal(i.ndim, 1) + assert_raises(ValueError, lambda i: i.shape, i) + assert_equal(i.itviews[0].shape, (24,)) + + # Removing the inner loop means there's just one iteration + i.reset() + assert_equal(i.itersize, 24) + assert_equal(i[0].shape, tuple()) + i.enable_external_loop() + assert_equal(i.itersize, 24) + assert_equal(i[0].shape, (24,)) + assert_equal(i.value, arange(24)) + +def test_iter_iterindex(): + # Make sure iterindex works + + buffersize = 5 + a = arange(24).reshape(4, 3, 2) + for flags in ([], ['buffered']): + i = nditer(a, flags, buffersize=buffersize) + assert_equal(iter_iterindices(i), list(range(24))) + i.iterindex = 2 + assert_equal(iter_iterindices(i), list(range(2, 24))) + + i = nditer(a, flags, order='F', buffersize=buffersize) + assert_equal(iter_iterindices(i), list(range(24))) + i.iterindex = 5 + assert_equal(iter_iterindices(i), list(range(5, 24))) + + i = nditer(a[::-1], flags, order='F', buffersize=buffersize) + assert_equal(iter_iterindices(i), list(range(24))) + i.iterindex = 9 + assert_equal(iter_iterindices(i), list(range(9, 24))) + + i = nditer(a[::-1, ::-1], flags, order='C', buffersize=buffersize) + assert_equal(iter_iterindices(i), list(range(24))) + i.iterindex = 13 + assert_equal(iter_iterindices(i), list(range(13, 24))) + + i = nditer(a[::1, ::-1], flags, buffersize=buffersize) + assert_equal(iter_iterindices(i), list(range(24))) + i.iterindex = 23 + assert_equal(iter_iterindices(i), list(range(23, 24))) + i.reset() + i.iterindex = 2 + assert_equal(iter_iterindices(i), list(range(2, 24))) + +def test_iter_iterrange(): + # Make sure getting and resetting the iterrange works + + buffersize = 5 + a = arange(24, dtype='i4').reshape(4, 3, 2) + a_fort = a.ravel(order='F') + + i = nditer(a, ['ranged'], ['readonly'], order='F', + buffersize=buffersize) + assert_equal(i.iterrange, (0, 24)) + assert_equal([x[()] for x in i], a_fort) + for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]: + i.iterrange = r + assert_equal(i.iterrange, r) + assert_equal([x[()] for x in i], a_fort[r[0]:r[1]]) + + i = nditer(a, ['ranged', 'buffered'], ['readonly'], order='F', + op_dtypes='f8', buffersize=buffersize) + assert_equal(i.iterrange, (0, 24)) + assert_equal([x[()] for x in i], a_fort) + for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]: + i.iterrange = r + assert_equal(i.iterrange, r) + assert_equal([x[()] for x in i], a_fort[r[0]:r[1]]) + + def get_array(i): + val = np.array([], dtype='f8') + for x in i: + val = np.concatenate((val, x)) + return val + + i = nditer(a, ['ranged', 'buffered', 'external_loop'], + ['readonly'], order='F', + op_dtypes='f8', buffersize=buffersize) + assert_equal(i.iterrange, (0, 24)) + assert_equal(get_array(i), a_fort) + for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]: + i.iterrange = r + assert_equal(i.iterrange, r) + assert_equal(get_array(i), a_fort[r[0]:r[1]]) + +def test_iter_buffering(): + # Test buffering with several buffer sizes and types + arrays = [] + # F-order swapped array + _tmp = np.arange(24, dtype='c16').reshape(2, 3, 4).T + _tmp = _tmp.view(_tmp.dtype.newbyteorder()).byteswap() + arrays.append(_tmp) + # Contiguous 1-dimensional array + arrays.append(np.arange(10, dtype='f4')) + # Unaligned array + a = np.zeros((4*16+1,), dtype='i1')[1:] + a.dtype = 'i4' + a[:] = np.arange(16, dtype='i4') + arrays.append(a) + # 4-D F-order array + arrays.append(np.arange(120, dtype='i4').reshape(5, 3, 2, 4).T) + for a in arrays: + for buffersize in (1, 2, 3, 5, 8, 11, 16, 1024): + vals = [] + i = nditer(a, ['buffered', 'external_loop'], + [['readonly', 'nbo', 'aligned']], + order='C', + casting='equiv', + buffersize=buffersize) + while not i.finished: + assert_(i[0].size <= buffersize) + vals.append(i[0].copy()) + i.iternext() + assert_equal(np.concatenate(vals), a.ravel(order='C')) + +def test_iter_write_buffering(): + # Test that buffering of writes is working + + # F-order swapped array + a = np.arange(24).reshape(2, 3, 4).T + a = a.view(a.dtype.newbyteorder()).byteswap() + i = nditer(a, ['buffered'], + [['readwrite', 'nbo', 'aligned']], + casting='equiv', + order='C', + buffersize=16) + x = 0 + with i: + while not i.finished: + i[0] = x + x += 1 + i.iternext() + assert_equal(a.ravel(order='C'), np.arange(24)) + +def test_iter_buffering_delayed_alloc(): + # Test that delaying buffer allocation works + + a = np.arange(6) + b = np.arange(1, dtype='f4') + i = nditer([a, b], ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok'], + ['readwrite'], + casting='unsafe', + op_dtypes='f4') + assert_(i.has_delayed_bufalloc) + assert_raises(ValueError, lambda i: i.multi_index, i) + assert_raises(ValueError, lambda i: i[0], i) + assert_raises(ValueError, lambda i: i[0:2], i) + + def assign_iter(i): + i[0] = 0 + assert_raises(ValueError, assign_iter, i) + + i.reset() + assert_(not i.has_delayed_bufalloc) + assert_equal(i.multi_index, (0,)) + with i: + assert_equal(i[0], 0) + i[1] = 1 + assert_equal(i[0:2], [0, 1]) + assert_equal([[x[0][()], x[1][()]] for x in i], list(zip(range(6), [1]*6))) + +def test_iter_buffered_cast_simple(): + # Test that buffering can handle a simple cast + + a = np.arange(10, dtype='f4') + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='same_kind', + op_dtypes=[np.dtype('f8')], + buffersize=3) + with i: + for v in i: + v[...] *= 2 + + assert_equal(a, 2*np.arange(10, dtype='f4')) + +def test_iter_buffered_cast_byteswapped(): + # Test that buffering can handle a cast which requires swap->cast->swap + + a = np.arange(10, dtype='f4') + a = a.view(a.dtype.newbyteorder()).byteswap() + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='same_kind', + op_dtypes=[np.dtype('f8').newbyteorder()], + buffersize=3) + with i: + for v in i: + v[...] *= 2 + + assert_equal(a, 2*np.arange(10, dtype='f4')) + + with suppress_warnings() as sup: + sup.filter(np.exceptions.ComplexWarning) + + a = np.arange(10, dtype='f8') + a = a.view(a.dtype.newbyteorder()).byteswap() + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='unsafe', + op_dtypes=[np.dtype('c8').newbyteorder()], + buffersize=3) + with i: + for v in i: + v[...] *= 2 + + assert_equal(a, 2*np.arange(10, dtype='f8')) + +def test_iter_buffered_cast_byteswapped_complex(): + # Test that buffering can handle a cast which requires swap->cast->copy + + a = np.arange(10, dtype='c8') + a = a.view(a.dtype.newbyteorder()).byteswap() + a += 2j + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='same_kind', + op_dtypes=[np.dtype('c16')], + buffersize=3) + with i: + for v in i: + v[...] *= 2 + assert_equal(a, 2*np.arange(10, dtype='c8') + 4j) + + a = np.arange(10, dtype='c8') + a += 2j + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='same_kind', + op_dtypes=[np.dtype('c16').newbyteorder()], + buffersize=3) + with i: + for v in i: + v[...] *= 2 + assert_equal(a, 2*np.arange(10, dtype='c8') + 4j) + + a = np.arange(10, dtype=np.clongdouble) + a = a.view(a.dtype.newbyteorder()).byteswap() + a += 2j + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='same_kind', + op_dtypes=[np.dtype('c16')], + buffersize=3) + with i: + for v in i: + v[...] *= 2 + assert_equal(a, 2*np.arange(10, dtype=np.clongdouble) + 4j) + + a = np.arange(10, dtype=np.longdouble) + a = a.view(a.dtype.newbyteorder()).byteswap() + i = nditer(a, ['buffered', 'external_loop'], + [['readwrite', 'nbo', 'aligned']], + casting='same_kind', + op_dtypes=[np.dtype('f4')], + buffersize=7) + with i: + for v in i: + v[...] *= 2 + assert_equal(a, 2*np.arange(10, dtype=np.longdouble)) + +def test_iter_buffered_cast_structured_type(): + # Tests buffering of structured types + + # simple -> struct type (duplicates the value) + sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')] + a = np.arange(3, dtype='f4') + 0.5 + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt) + vals = [np.array(x) for x in i] + assert_equal(vals[0]['a'], 0.5) + assert_equal(vals[0]['b'], 0) + assert_equal(vals[0]['c'], [[(0.5)]*3]*2) + assert_equal(vals[0]['d'], 0.5) + assert_equal(vals[1]['a'], 1.5) + assert_equal(vals[1]['b'], 1) + assert_equal(vals[1]['c'], [[(1.5)]*3]*2) + assert_equal(vals[1]['d'], 1.5) + assert_equal(vals[0].dtype, np.dtype(sdt)) + + # object -> struct type + sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')] + a = np.zeros((3,), dtype='O') + a[0] = (0.5, 0.5, [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]], 0.5) + a[1] = (1.5, 1.5, [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]], 1.5) + a[2] = (2.5, 2.5, [[2.5, 2.5, 2.5], [2.5, 2.5, 2.5]], 2.5) + if HAS_REFCOUNT: + rc = sys.getrefcount(a[0]) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt) + vals = [x.copy() for x in i] + assert_equal(vals[0]['a'], 0.5) + assert_equal(vals[0]['b'], 0) + assert_equal(vals[0]['c'], [[(0.5)]*3]*2) + assert_equal(vals[0]['d'], 0.5) + assert_equal(vals[1]['a'], 1.5) + assert_equal(vals[1]['b'], 1) + assert_equal(vals[1]['c'], [[(1.5)]*3]*2) + assert_equal(vals[1]['d'], 1.5) + assert_equal(vals[0].dtype, np.dtype(sdt)) + vals, i, x = [None]*3 + if HAS_REFCOUNT: + assert_equal(sys.getrefcount(a[0]), rc) + + # single-field struct type -> simple + sdt = [('a', 'f4')] + a = np.array([(5.5,), (8,)], dtype=sdt) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes='i4') + assert_equal([x_[()] for x_ in i], [5, 8]) + + # make sure multi-field struct type -> simple doesn't work + sdt = [('a', 'f4'), ('b', 'i8'), ('d', 'O')] + a = np.array([(5.5, 7, 'test'), (8, 10, 11)], dtype=sdt) + assert_raises(TypeError, lambda: ( + nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes='i4'))) + + # struct type -> struct type (field-wise copy) + sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')] + sdt2 = [('d', 'u2'), ('a', 'O'), ('b', 'f8')] + a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + assert_equal([np.array(x_) for x_ in i], + [np.array((1, 2, 3), dtype=sdt2), + np.array((4, 5, 6), dtype=sdt2)]) + + +def test_iter_buffered_cast_structured_type_failure_with_cleanup(): + # make sure struct type -> struct type with different + # number of fields fails + sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')] + sdt2 = [('b', 'O'), ('a', 'f8')] + a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1) + + for intent in ["readwrite", "readonly", "writeonly"]: + # This test was initially designed to test an error at a different + # place, but will now raise earlier to to the cast not being possible: + # `assert np.can_cast(a.dtype, sdt2, casting="unsafe")` fails. + # Without a faulty DType, there is probably no reliable + # way to get the initial tested behaviour. + simple_arr = np.array([1, 2], dtype="i,i") # requires clean up + with pytest.raises(TypeError): + nditer((simple_arr, a), ['buffered', 'refs_ok'], [intent, intent], + casting='unsafe', op_dtypes=["f,f", sdt2]) + + +def test_buffered_cast_error_paths(): + with pytest.raises(ValueError): + # The input is cast into an `S3` buffer + np.nditer((np.array("a", dtype="S1"),), op_dtypes=["i"], + casting="unsafe", flags=["buffered"]) + + # The `M8[ns]` is cast into the `S3` output + it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"], + op_flags=["writeonly"], casting="unsafe", flags=["buffered"]) + with pytest.raises(ValueError): + with it: + buf = next(it) + buf[...] = "a" # cannot be converted to int. + +@pytest.mark.skipif(IS_WASM, reason="Cannot start subprocess") +@pytest.mark.skipif(not HAS_REFCOUNT, reason="PyPy seems to not hit this.") +def test_buffered_cast_error_paths_unraisable(): + # The following gives an unraisable error. Pytest sometimes captures that + # (depending python and/or pytest version). So with Python>=3.8 this can + # probably be cleaned out in the future to check for + # pytest.PytestUnraisableExceptionWarning: + code = textwrap.dedent(""" + import numpy as np + + it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"], + op_flags=["writeonly"], casting="unsafe", flags=["buffered"]) + buf = next(it) + buf[...] = "a" + del buf, it # Flushing only happens during deallocate right now. + """) + res = subprocess.check_output([sys.executable, "-c", code], + stderr=subprocess.STDOUT, text=True) + assert "ValueError" in res + + +def test_iter_buffered_cast_subarray(): + # Tests buffering of subarrays + + # one element -> many (copies it to all) + sdt1 = [('a', 'f4')] + sdt2 = [('a', 'f8', (3, 2, 2))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + for x, count in zip(i, list(range(6))): + assert_(np.all(x['a'] == count)) + + # one element -> many -> back (copies it to all) + sdt1 = [('a', 'O', (1, 1))] + sdt2 = [('a', 'O', (3, 2, 2))] + a = np.zeros((6,), dtype=sdt1) + a['a'][:, 0, 0] = np.arange(6) + i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'], + casting='unsafe', + op_dtypes=sdt2) + with i: + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_(np.all(x['a'] == count)) + x['a'][0] += 2 + count += 1 + assert_equal(a['a'], np.arange(6).reshape(6, 1, 1)+2) + + # many -> one element -> back (copies just element 0) + sdt1 = [('a', 'O', (3, 2, 2))] + sdt2 = [('a', 'O', (1,))] + a = np.zeros((6,), dtype=sdt1) + a['a'][:, 0, 0, 0] = np.arange(6) + i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'], + casting='unsafe', + op_dtypes=sdt2) + with i: + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'], count) + x['a'] += 2 + count += 1 + assert_equal(a['a'], np.arange(6).reshape(6, 1, 1, 1)*np.ones((1, 3, 2, 2))+2) + + # many -> one element -> back (copies just element 0) + sdt1 = [('a', 'f8', (3, 2, 2))] + sdt2 = [('a', 'O', (1,))] + a = np.zeros((6,), dtype=sdt1) + a['a'][:, 0, 0, 0] = np.arange(6) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'], count) + count += 1 + + # many -> one element (copies just element 0) + sdt1 = [('a', 'O', (3, 2, 2))] + sdt2 = [('a', 'f4', (1,))] + a = np.zeros((6,), dtype=sdt1) + a['a'][:, 0, 0, 0] = np.arange(6) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'], count) + count += 1 + + # many -> matching shape (straightforward copy) + sdt1 = [('a', 'O', (3, 2, 2))] + sdt2 = [('a', 'f4', (3, 2, 2))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6*3*2*2).reshape(6, 3, 2, 2) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'], a[count]['a']) + count += 1 + + # vector -> smaller vector (truncates) + sdt1 = [('a', 'f8', (6,))] + sdt2 = [('a', 'f4', (2,))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6*6).reshape(6, 6) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'], a[count]['a'][:2]) + count += 1 + + # vector -> bigger vector (pads with zeros) + sdt1 = [('a', 'f8', (2,))] + sdt2 = [('a', 'f4', (6,))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6*2).reshape(6, 2) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'][:2], a[count]['a']) + assert_equal(x['a'][2:], [0, 0, 0, 0]) + count += 1 + + # vector -> matrix (broadcasts) + sdt1 = [('a', 'f8', (2,))] + sdt2 = [('a', 'f4', (2, 2))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6*2).reshape(6, 2) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'][0], a[count]['a']) + assert_equal(x['a'][1], a[count]['a']) + count += 1 + + # vector -> matrix (broadcasts and zero-pads) + sdt1 = [('a', 'f8', (2, 1))] + sdt2 = [('a', 'f4', (3, 2))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6*2).reshape(6, 2, 1) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'][:2, 0], a[count]['a'][:, 0]) + assert_equal(x['a'][:2, 1], a[count]['a'][:, 0]) + assert_equal(x['a'][2, :], [0, 0]) + count += 1 + + # matrix -> matrix (truncates and zero-pads) + sdt1 = [('a', 'f8', (2, 3))] + sdt2 = [('a', 'f4', (3, 2))] + a = np.zeros((6,), dtype=sdt1) + a['a'] = np.arange(6*2*3).reshape(6, 2, 3) + i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], + casting='unsafe', + op_dtypes=sdt2) + assert_equal(i[0].dtype, np.dtype(sdt2)) + count = 0 + for x in i: + assert_equal(x['a'][:2, 0], a[count]['a'][:, 0]) + assert_equal(x['a'][:2, 1], a[count]['a'][:, 1]) + assert_equal(x['a'][2, :], [0, 0]) + count += 1 + +def test_iter_buffering_badwriteback(): + # Writing back from a buffer cannot combine elements + + # a needs write buffering, but had a broadcast dimension + a = np.arange(6).reshape(2, 3, 1) + b = np.arange(12).reshape(2, 3, 2) + assert_raises(ValueError, nditer, [a, b], + ['buffered', 'external_loop'], + [['readwrite'], ['writeonly']], + order='C') + + # But if a is readonly, it's fine + nditer([a, b], ['buffered', 'external_loop'], + [['readonly'], ['writeonly']], + order='C') + + # If a has just one element, it's fine too (constant 0 stride, a reduction) + a = np.arange(1).reshape(1, 1, 1) + nditer([a, b], ['buffered', 'external_loop', 'reduce_ok'], + [['readwrite'], ['writeonly']], + order='C') + + # check that it fails on other dimensions too + a = np.arange(6).reshape(1, 3, 2) + assert_raises(ValueError, nditer, [a, b], + ['buffered', 'external_loop'], + [['readwrite'], ['writeonly']], + order='C') + a = np.arange(4).reshape(2, 1, 2) + assert_raises(ValueError, nditer, [a, b], + ['buffered', 'external_loop'], + [['readwrite'], ['writeonly']], + order='C') + +def test_iter_buffering_string(): + # Safe casting disallows shrinking strings + a = np.array(['abc', 'a', 'abcd'], dtype=np.bytes_) + assert_equal(a.dtype, np.dtype('S4')) + assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'], + op_dtypes='S2') + i = nditer(a, ['buffered'], ['readonly'], op_dtypes='S6') + assert_equal(i[0], b'abc') + assert_equal(i[0].dtype, np.dtype('S6')) + + a = np.array(['abc', 'a', 'abcd'], dtype=np.str_) + assert_equal(a.dtype, np.dtype('U4')) + assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'], + op_dtypes='U2') + i = nditer(a, ['buffered'], ['readonly'], op_dtypes='U6') + assert_equal(i[0], 'abc') + assert_equal(i[0].dtype, np.dtype('U6')) + +def test_iter_buffering_growinner(): + # Test that the inner loop grows when no buffering is needed + a = np.arange(30) + i = nditer(a, ['buffered', 'growinner', 'external_loop'], + buffersize=5) + # Should end up with just one inner loop here + assert_equal(i[0].size, a.size) + + +@pytest.mark.slow +def test_iter_buffered_reduce_reuse(): + # large enough array for all views, including negative strides. + a = np.arange(2*3**5)[3**5:3**5+1] + flags = ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok', 'refs_ok'] + op_flags = [('readonly',), ('readwrite', 'allocate')] + op_axes_list = [[(0, 1, 2), (0, 1, -1)], [(0, 1, 2), (0, -1, -1)]] + # wrong dtype to force buffering + op_dtypes = [float, a.dtype] + + def get_params(): + for xs in range(-3**2, 3**2 + 1): + for ys in range(xs, 3**2 + 1): + for op_axes in op_axes_list: + # last stride is reduced and because of that not + # important for this test, as it is the inner stride. + strides = (xs * a.itemsize, ys * a.itemsize, a.itemsize) + arr = np.lib.stride_tricks.as_strided(a, (3, 3, 3), strides) + + for skip in [0, 1]: + yield arr, op_axes, skip + + for arr, op_axes, skip in get_params(): + nditer2 = np.nditer([arr.copy(), None], + op_axes=op_axes, flags=flags, op_flags=op_flags, + op_dtypes=op_dtypes) + with nditer2: + nditer2.operands[-1][...] = 0 + nditer2.reset() + nditer2.iterindex = skip + + for (a2_in, b2_in) in nditer2: + b2_in += a2_in.astype(np.int_) + + comp_res = nditer2.operands[-1] + + for bufsize in range(0, 3**3): + nditer1 = np.nditer([arr, None], + op_axes=op_axes, flags=flags, op_flags=op_flags, + buffersize=bufsize, op_dtypes=op_dtypes) + with nditer1: + nditer1.operands[-1][...] = 0 + nditer1.reset() + nditer1.iterindex = skip + + for (a1_in, b1_in) in nditer1: + b1_in += a1_in.astype(np.int_) + + res = nditer1.operands[-1] + assert_array_equal(res, comp_res) + + +def test_iter_no_broadcast(): + # Test that the no_broadcast flag works + a = np.arange(24).reshape(2, 3, 4) + b = np.arange(6).reshape(2, 3, 1) + c = np.arange(12).reshape(3, 4) + + nditer([a, b, c], [], + [['readonly', 'no_broadcast'], + ['readonly'], ['readonly']]) + assert_raises(ValueError, nditer, [a, b, c], [], + [['readonly'], ['readonly', 'no_broadcast'], ['readonly']]) + assert_raises(ValueError, nditer, [a, b, c], [], + [['readonly'], ['readonly'], ['readonly', 'no_broadcast']]) + + +class TestIterNested: + + def test_basic(self): + # Test nested iteration basic usage + a = arange(12).reshape(2, 3, 2) + + i, j = np.nested_iters(a, [[0], [1, 2]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]) + + i, j = np.nested_iters(a, [[0, 1], [2]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]]) + + i, j = np.nested_iters(a, [[0, 2], [1]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) + + def test_reorder(self): + # Test nested iteration basic usage + a = arange(12).reshape(2, 3, 2) + + # In 'K' order (default), it gets reordered + i, j = np.nested_iters(a, [[0], [2, 1]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]) + + i, j = np.nested_iters(a, [[1, 0], [2]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]]) + + i, j = np.nested_iters(a, [[2, 0], [1]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) + + # In 'C' order, it doesn't + i, j = np.nested_iters(a, [[0], [2, 1]], order='C') + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 2, 4, 1, 3, 5], [6, 8, 10, 7, 9, 11]]) + + i, j = np.nested_iters(a, [[1, 0], [2]], order='C') + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1], [6, 7], [2, 3], [8, 9], [4, 5], [10, 11]]) + + i, j = np.nested_iters(a, [[2, 0], [1]], order='C') + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 2, 4], [6, 8, 10], [1, 3, 5], [7, 9, 11]]) + + def test_flip_axes(self): + # Test nested iteration with negative axes + a = arange(12).reshape(2, 3, 2)[::-1, ::-1, ::-1] + + # In 'K' order (default), the axes all get flipped + i, j = np.nested_iters(a, [[0], [1, 2]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]) + + i, j = np.nested_iters(a, [[0, 1], [2]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]]) + + i, j = np.nested_iters(a, [[0, 2], [1]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) + + # In 'C' order, flipping axes is disabled + i, j = np.nested_iters(a, [[0], [1, 2]], order='C') + vals = [list(j) for _ in i] + assert_equal(vals, [[11, 10, 9, 8, 7, 6], [5, 4, 3, 2, 1, 0]]) + + i, j = np.nested_iters(a, [[0, 1], [2]], order='C') + vals = [list(j) for _ in i] + assert_equal(vals, [[11, 10], [9, 8], [7, 6], [5, 4], [3, 2], [1, 0]]) + + i, j = np.nested_iters(a, [[0, 2], [1]], order='C') + vals = [list(j) for _ in i] + assert_equal(vals, [[11, 9, 7], [10, 8, 6], [5, 3, 1], [4, 2, 0]]) + + def test_broadcast(self): + # Test nested iteration with broadcasting + a = arange(2).reshape(2, 1) + b = arange(3).reshape(1, 3) + + i, j = np.nested_iters([a, b], [[0], [1]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[[0, 0], [0, 1], [0, 2]], [[1, 0], [1, 1], [1, 2]]]) + + i, j = np.nested_iters([a, b], [[1], [0]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[[0, 0], [1, 0]], [[0, 1], [1, 1]], [[0, 2], [1, 2]]]) + + def test_dtype_copy(self): + # Test nested iteration with a copy to change dtype + + # copy + a = arange(6, dtype='i4').reshape(2, 3) + i, j = np.nested_iters(a, [[0], [1]], + op_flags=['readonly', 'copy'], + op_dtypes='f8') + assert_equal(j[0].dtype, np.dtype('f8')) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1, 2], [3, 4, 5]]) + vals = None + + # writebackifcopy - using context manager + a = arange(6, dtype='f4').reshape(2, 3) + i, j = np.nested_iters(a, [[0], [1]], + op_flags=['readwrite', 'updateifcopy'], + casting='same_kind', + op_dtypes='f8') + with i, j: + assert_equal(j[0].dtype, np.dtype('f8')) + for x in i: + for y in j: + y[...] += 1 + assert_equal(a, [[0, 1, 2], [3, 4, 5]]) + assert_equal(a, [[1, 2, 3], [4, 5, 6]]) + + # writebackifcopy - using close() + a = arange(6, dtype='f4').reshape(2, 3) + i, j = np.nested_iters(a, [[0], [1]], + op_flags=['readwrite', 'updateifcopy'], + casting='same_kind', + op_dtypes='f8') + assert_equal(j[0].dtype, np.dtype('f8')) + for x in i: + for y in j: + y[...] += 1 + assert_equal(a, [[0, 1, 2], [3, 4, 5]]) + i.close() + j.close() + assert_equal(a, [[1, 2, 3], [4, 5, 6]]) + + def test_dtype_buffered(self): + # Test nested iteration with buffering to change dtype + + a = arange(6, dtype='f4').reshape(2, 3) + i, j = np.nested_iters(a, [[0], [1]], + flags=['buffered'], + op_flags=['readwrite'], + casting='same_kind', + op_dtypes='f8') + assert_equal(j[0].dtype, np.dtype('f8')) + for x in i: + for y in j: + y[...] += 1 + assert_equal(a, [[1, 2, 3], [4, 5, 6]]) + + def test_0d(self): + a = np.arange(12).reshape(2, 3, 2) + i, j = np.nested_iters(a, [[], [1, 0, 2]]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]]) + + i, j = np.nested_iters(a, [[1, 0, 2], []]) + vals = [list(j) for _ in i] + assert_equal(vals, [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11]]) + + i, j, k = np.nested_iters(a, [[2, 0], [], [1]]) + vals = [] + for x in i: + for y in j: + vals.append(list(k)) + assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) + + def test_iter_nested_iters_dtype_buffered(self): + # Test nested iteration with buffering to change dtype + + a = arange(6, dtype='f4').reshape(2, 3) + i, j = np.nested_iters(a, [[0], [1]], + flags=['buffered'], + op_flags=['readwrite'], + casting='same_kind', + op_dtypes='f8') + with i, j: + assert_equal(j[0].dtype, np.dtype('f8')) + for x in i: + for y in j: + y[...] += 1 + assert_equal(a, [[1, 2, 3], [4, 5, 6]]) + +def test_iter_reduction_error(): + + a = np.arange(6) + assert_raises(ValueError, nditer, [a, None], [], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[[0], [-1]]) + + a = np.arange(6).reshape(2, 3) + assert_raises(ValueError, nditer, [a, None], ['external_loop'], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[[0, 1], [-1, -1]]) + +def test_iter_reduction(): + # Test doing reductions with the iterator + + a = np.arange(6) + i = nditer([a, None], ['reduce_ok'], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[[0], [-1]]) + # Need to initialize the output operand to the addition unit + with i: + i.operands[1][...] = 0 + # Do the reduction + for x, y in i: + y[...] += x + # Since no axes were specified, should have allocated a scalar + assert_equal(i.operands[1].ndim, 0) + assert_equal(i.operands[1], np.sum(a)) + + a = np.arange(6).reshape(2, 3) + i = nditer([a, None], ['reduce_ok', 'external_loop'], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[[0, 1], [-1, -1]]) + # Need to initialize the output operand to the addition unit + with i: + i.operands[1][...] = 0 + # Reduction shape/strides for the output + assert_equal(i[1].shape, (6,)) + assert_equal(i[1].strides, (0,)) + # Do the reduction + for x, y in i: + # Use a for loop instead of ``y[...] += x`` + # (equivalent to ``y[...] = y[...].copy() + x``), + # because y has zero strides we use for the reduction + for j in range(len(y)): + y[j] += x[j] + # Since no axes were specified, should have allocated a scalar + assert_equal(i.operands[1].ndim, 0) + assert_equal(i.operands[1], np.sum(a)) + + # This is a tricky reduction case for the buffering double loop + # to handle + a = np.ones((2, 3, 5)) + it1 = nditer([a, None], ['reduce_ok', 'external_loop'], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[None, [0, -1, 1]]) + it2 = nditer([a, None], ['reduce_ok', 'external_loop', + 'buffered', 'delay_bufalloc'], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[None, [0, -1, 1]], buffersize=10) + with it1, it2: + it1.operands[1].fill(0) + it2.operands[1].fill(0) + it2.reset() + for x in it1: + x[1][...] += x[0] + for x in it2: + x[1][...] += x[0] + assert_equal(it1.operands[1], it2.operands[1]) + assert_equal(it2.operands[1].sum(), a.size) + +def test_iter_buffering_reduction(): + # Test doing buffered reductions with the iterator + + a = np.arange(6) + b = np.array(0., dtype='f8').byteswap() + b = b.view(b.dtype.newbyteorder()) + i = nditer([a, b], ['reduce_ok', 'buffered'], + [['readonly'], ['readwrite', 'nbo']], + op_axes=[[0], [-1]]) + with i: + assert_equal(i[1].dtype, np.dtype('f8')) + assert_(i[1].dtype != b.dtype) + # Do the reduction + for x, y in i: + y[...] += x + # Since no axes were specified, should have allocated a scalar + assert_equal(b, np.sum(a)) + + a = np.arange(6).reshape(2, 3) + b = np.array([0, 0], dtype='f8').byteswap() + b = b.view(b.dtype.newbyteorder()) + i = nditer([a, b], ['reduce_ok', 'external_loop', 'buffered'], + [['readonly'], ['readwrite', 'nbo']], + op_axes=[[0, 1], [0, -1]]) + # Reduction shape/strides for the output + with i: + assert_equal(i[1].shape, (3,)) + assert_equal(i[1].strides, (0,)) + # Do the reduction + for x, y in i: + # Use a for loop instead of ``y[...] += x`` + # (equivalent to ``y[...] = y[...].copy() + x``), + # because y has zero strides we use for the reduction + for j in range(len(y)): + y[j] += x[j] + assert_equal(b, np.sum(a, axis=1)) + + # Iterator inner double loop was wrong on this one + p = np.arange(2) + 1 + it = np.nditer([p, None], + ['delay_bufalloc', 'reduce_ok', 'buffered', 'external_loop'], + [['readonly'], ['readwrite', 'allocate']], + op_axes=[[-1, 0], [-1, -1]], + itershape=(2, 2)) + with it: + it.operands[1].fill(0) + it.reset() + assert_equal(it[0], [1, 2, 1, 2]) + + # Iterator inner loop should take argument contiguity into account + x = np.ones((7, 13, 8), np.int8)[4:6, 1:11:6, 1:5].transpose(1, 2, 0) + x[...] = np.arange(x.size).reshape(x.shape) + y_base = np.arange(4*4, dtype=np.int8).reshape(4, 4) + y_base_copy = y_base.copy() + y = y_base[::2, :, None] + + it = np.nditer([y, x], + ['buffered', 'external_loop', 'reduce_ok'], + [['readwrite'], ['readonly']]) + with it: + for a, b in it: + a.fill(2) + + assert_equal(y_base[1::2], y_base_copy[1::2]) + assert_equal(y_base[::2], 2) + +def test_iter_buffering_reduction_reuse_reduce_loops(): + # There was a bug triggering reuse of the reduce loop inappropriately, + # which caused processing to happen in unnecessarily small chunks + # and overran the buffer. + + a = np.zeros((2, 7)) + b = np.zeros((1, 7)) + it = np.nditer([a, b], flags=['reduce_ok', 'external_loop', 'buffered'], + op_flags=[['readonly'], ['readwrite']], + buffersize=5) + + with it: + bufsizes = [x.shape[0] for x, y in it] + assert_equal(bufsizes, [5, 2, 5, 2]) + assert_equal(sum(bufsizes), a.size) + +def test_iter_writemasked_badinput(): + a = np.zeros((2, 3)) + b = np.zeros((3,)) + m = np.array([[True, True, False], [False, True, False]]) + m2 = np.array([True, True, False]) + m3 = np.array([0, 1, 1], dtype='u1') + mbad1 = np.array([0, 1, 1], dtype='i1') + mbad2 = np.array([0, 1, 1], dtype='f4') + + # Need an 'arraymask' if any operand is 'writemasked' + assert_raises(ValueError, nditer, [a, m], [], + [['readwrite', 'writemasked'], ['readonly']]) + + # A 'writemasked' operand must not be readonly + assert_raises(ValueError, nditer, [a, m], [], + [['readonly', 'writemasked'], ['readonly', 'arraymask']]) + + # 'writemasked' and 'arraymask' may not be used together + assert_raises(ValueError, nditer, [a, m], [], + [['readonly'], ['readwrite', 'arraymask', 'writemasked']]) + + # 'arraymask' may only be specified once + assert_raises(ValueError, nditer, [a, m, m2], [], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask'], + ['readonly', 'arraymask']]) + + # An 'arraymask' with nothing 'writemasked' also doesn't make sense + assert_raises(ValueError, nditer, [a, m], [], + [['readwrite'], ['readonly', 'arraymask']]) + + # A writemasked reduction requires a similarly smaller mask + assert_raises(ValueError, nditer, [a, b, m], ['reduce_ok'], + [['readonly'], + ['readwrite', 'writemasked'], + ['readonly', 'arraymask']]) + # But this should work with a smaller/equal mask to the reduction operand + np.nditer([a, b, m2], ['reduce_ok'], + [['readonly'], + ['readwrite', 'writemasked'], + ['readonly', 'arraymask']]) + # The arraymask itself cannot be a reduction + assert_raises(ValueError, nditer, [a, b, m2], ['reduce_ok'], + [['readonly'], + ['readwrite', 'writemasked'], + ['readwrite', 'arraymask']]) + + # A uint8 mask is ok too + np.nditer([a, m3], ['buffered'], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']], + op_dtypes=['f4', None], + casting='same_kind') + # An int8 mask isn't ok + assert_raises(TypeError, np.nditer, [a, mbad1], ['buffered'], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']], + op_dtypes=['f4', None], + casting='same_kind') + # A float32 mask isn't ok + assert_raises(TypeError, np.nditer, [a, mbad2], ['buffered'], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']], + op_dtypes=['f4', None], + casting='same_kind') + + +def _is_buffered(iterator): + try: + iterator.itviews + except ValueError: + return True + return False + +@pytest.mark.parametrize("a", + [np.zeros((3,), dtype='f8'), + np.zeros((9876, 3*5), dtype='f8')[::2, :], + np.zeros((4, 312, 124, 3), dtype='f8')[::2, :, ::2, :], + # Also test with the last dimension strided (so it does not fit if + # there is repeated access) + np.zeros((9,), dtype='f8')[::3], + np.zeros((9876, 3*10), dtype='f8')[::2, ::5], + np.zeros((4, 312, 124, 3), dtype='f8')[::2, :, ::2, ::-1]]) +def test_iter_writemasked(a): + # Note, the slicing above is to ensure that nditer cannot combine multiple + # axes into one. The repetition is just to make things a bit more + # interesting. + shape = a.shape + reps = shape[-1] // 3 + msk = np.empty(shape, dtype=bool) + msk[...] = [True, True, False] * reps + + # When buffering is unused, 'writemasked' effectively does nothing. + # It's up to the user of the iterator to obey the requested semantics. + it = np.nditer([a, msk], [], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']]) + with it: + for x, m in it: + x[...] = 1 + # Because we violated the semantics, all the values became 1 + assert_equal(a, np.broadcast_to([1, 1, 1] * reps, shape)) + + # Even if buffering is enabled, we still may be accessing the array + # directly. + it = np.nditer([a, msk], ['buffered'], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']]) + # @seberg: I honestly don't currently understand why a "buffered" iterator + # would end up not using a buffer for the small array here at least when + # "writemasked" is used, that seems confusing... Check by testing for + # actual memory overlap! + is_buffered = True + with it: + for x, m in it: + x[...] = 2.5 + if np.may_share_memory(x, a): + is_buffered = False + + if not is_buffered: + # Because we violated the semantics, all the values became 2.5 + assert_equal(a, np.broadcast_to([2.5, 2.5, 2.5] * reps, shape)) + else: + # For large sizes, the iterator may be buffered: + assert_equal(a, np.broadcast_to([2.5, 2.5, 1] * reps, shape)) + a[...] = 2.5 + + # If buffering will definitely happening, for instance because of + # a cast, only the items selected by the mask will be copied back from + # the buffer. + it = np.nditer([a, msk], ['buffered'], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']], + op_dtypes=['i8', None], + casting='unsafe') + with it: + for x, m in it: + x[...] = 3 + # Even though we violated the semantics, only the selected values + # were copied back + assert_equal(a, np.broadcast_to([3, 3, 2.5] * reps, shape)) + + +@pytest.mark.parametrize(["mask", "mask_axes"], [ + # Allocated operand (only broadcasts with -1) + (None, [-1, 0]), + # Reduction along the first dimension (with and without op_axes) + (np.zeros((1, 4), dtype="bool"), [0, 1]), + (np.zeros((1, 4), dtype="bool"), None), + # Test 0-D and -1 op_axes + (np.zeros(4, dtype="bool"), [-1, 0]), + (np.zeros((), dtype="bool"), [-1, -1]), + (np.zeros((), dtype="bool"), None)]) +def test_iter_writemasked_broadcast_error(mask, mask_axes): + # This assumes that a readwrite mask makes sense. This is likely not the + # case and should simply be deprecated. + arr = np.zeros((3, 4)) + itflags = ["reduce_ok"] + mask_flags = ["arraymask", "readwrite", "allocate"] + a_flags = ["writeonly", "writemasked"] + if mask_axes is None: + op_axes = None + else: + op_axes = [mask_axes, [0, 1]] + + with assert_raises(ValueError): + np.nditer((mask, arr), flags=itflags, op_flags=[mask_flags, a_flags], + op_axes=op_axes) + + +def test_iter_writemasked_decref(): + # force casting (to make it interesting) by using a structured dtype. + arr = np.arange(10000).astype(">i,O") + original = arr.copy() + mask = np.random.randint(0, 2, size=10000).astype(bool) + + it = np.nditer([arr, mask], ['buffered', "refs_ok"], + [['readwrite', 'writemasked'], + ['readonly', 'arraymask']], + op_dtypes=[" string -> longdouble` for the + # conversion. But Python may refuse `str(int)` for huge ints. + # In that case, RuntimeWarning would be correct, but conversion + # fails earlier (seems to happen on 32bit linux, possibly only debug). + if dtype in "gG": + try: + str(too_big_int) + except ValueError: + pytest.skip("`huge_int -> string -> longdouble` failed") + + # Otherwise, we overflow to infinity: + with pytest.warns(RuntimeWarning): + res = scalar_type(1) + too_big_int + assert res.dtype == dtype + assert res == np.inf + + with pytest.warns(RuntimeWarning): + # We force the dtype here, since windows may otherwise pick the + # double instead of the longdouble loop. That leads to slightly + # different results (conversion of the int fails as above). + res = np.add(np.array(1, dtype=dtype), too_big_int, dtype=dtype) + assert res.dtype == dtype + assert res == np.inf + + +@pytest.mark.parametrize("op", [operator.add, operator.pow]) +def test_weak_promotion_scalar_path(op): + # Some additional paths exercising the weak scalars. + + # Integer path: + res = op(np.uint8(3), 5) + assert res == op(3, 5) + assert res.dtype == np.uint8 or res.dtype == bool + + with pytest.raises(OverflowError): + op(np.uint8(3), 1000) + + # Float path: + res = op(np.float32(3), 5.) + assert res == op(3., 5.) + assert res.dtype == np.float32 or res.dtype == bool + + +def test_nep50_complex_promotion(): + with pytest.warns(RuntimeWarning, match=".*overflow"): + res = np.complex64(3) + complex(2**300) + + assert type(res) == np.complex64 + + +def test_nep50_integer_conversion_errors(): + # Implementation for error paths is mostly missing (as of writing) + with pytest.raises(OverflowError, match=".*uint8"): + np.array([1], np.uint8) + 300 + + with pytest.raises(OverflowError, match=".*uint8"): + np.uint8(1) + 300 + + # Error message depends on platform (maybe unsigned int or unsigned long) + with pytest.raises(OverflowError, + match="Python integer -1 out of bounds for uint8"): + np.uint8(1) + -1 + + +def test_nep50_with_axisconcatenator(): + # Concatenate/r_ does not promote, so this has to error: + with pytest.raises(OverflowError): + np.r_[np.arange(5, dtype=np.int8), 255] + + +@pytest.mark.parametrize("ufunc", [np.add, np.power]) +def test_nep50_huge_integers(ufunc): + # Very large integers are complicated, because they go to uint64 or + # object dtype. This tests covers a few possible paths. + with pytest.raises(OverflowError): + ufunc(np.int64(0), 2**63) # 2**63 too large for int64 + + with pytest.raises(OverflowError): + ufunc(np.uint64(0), 2**64) # 2**64 cannot be represented by uint64 + + # However, 2**63 can be represented by the uint64 (and that is used): + res = ufunc(np.uint64(1), 2**63) + + assert res.dtype == np.uint64 + assert res == ufunc(1, 2**63, dtype=object) + + # The following paths fail to warn correctly about the change: + with pytest.raises(OverflowError): + ufunc(np.int64(1), 2**63) # np.array(2**63) would go to uint + + with pytest.raises(OverflowError): + ufunc(np.int64(1), 2**100) # np.array(2**100) would go to object + + # This would go to object and thus a Python float, not a NumPy one: + res = ufunc(1.0, 2**100) + assert isinstance(res, np.float64) + + +def test_nep50_in_concat_and_choose(): + res = np.concatenate([np.float32(1), 1.], axis=None) + assert res.dtype == "float32" + + res = np.choose(1, [np.float32(1), 1.]) + assert res.dtype == "float32" + + +@pytest.mark.parametrize("expected,dtypes,optional_dtypes", [ + (np.float32, [np.float32], + [np.float16, 0.0, np.uint16, np.int16, np.int8, 0]), + (np.complex64, [np.float32, 0j], + [np.float16, 0.0, np.uint16, np.int16, np.int8, 0]), + (np.float32, [np.int16, np.uint16, np.float16], + [np.int8, np.uint8, np.float32, 0., 0]), + (np.int32, [np.int16, np.uint16], + [np.int8, np.uint8, 0, np.bool]), + ]) +@hypothesis.given(data=strategies.data()) +def test_expected_promotion(expected, dtypes, optional_dtypes, data): + # Sample randomly while ensuring "dtypes" is always present: + optional = data.draw(strategies.lists( + strategies.sampled_from(dtypes + optional_dtypes))) + all_dtypes = dtypes + optional + dtypes_sample = data.draw(strategies.permutations(all_dtypes)) + + res = np.result_type(*dtypes_sample) + assert res == expected + + +@pytest.mark.parametrize("sctype", + [np.int8, np.int16, np.int32, np.int64, + np.uint8, np.uint16, np.uint32, np.uint64]) +@pytest.mark.parametrize("other_val", + [-2*100, -1, 0, 9, 10, 11, 2**63, 2*100]) +@pytest.mark.parametrize("comp", + [operator.eq, operator.ne, operator.le, operator.lt, + operator.ge, operator.gt]) +def test_integer_comparison(sctype, other_val, comp): + # Test that comparisons with integers (especially out-of-bound) ones + # works correctly. + val_obj = 10 + val = sctype(val_obj) + # Check that the scalar behaves the same as the python int: + assert comp(10, other_val) == comp(val, other_val) + assert comp(val, other_val) == comp(10, other_val) + # Except for the result type: + assert type(comp(val, other_val)) is np.bool + + # Check that the integer array and object array behave the same: + val_obj = np.array([10, 10], dtype=object) + val = val_obj.astype(sctype) + assert_array_equal(comp(val_obj, other_val), comp(val, other_val)) + assert_array_equal(comp(other_val, val_obj), comp(other_val, val)) + + +@pytest.mark.parametrize("arr", [ + np.ones((100, 100), dtype=np.uint8)[::2], # not trivially iterable + np.ones(20000, dtype=">u4"), # cast and >buffersize + np.ones(100, dtype=">u4"), # fast path compatible with cast +]) +def test_integer_comparison_with_cast(arr): + # Similar to above, but mainly test a few cases that cover the slow path + # the test is limited to unsigned ints and -1 for simplicity. + res = arr >= -1 + assert_array_equal(res, np.ones_like(arr, dtype=bool)) + res = arr < -1 + assert_array_equal(res, np.zeros_like(arr, dtype=bool)) + + +@pytest.mark.parametrize("comp", + [np.equal, np.not_equal, np.less_equal, np.less, + np.greater_equal, np.greater]) +def test_integer_integer_comparison(comp): + # Test that the NumPy comparison ufuncs work with large Python integers + assert comp(2**200, -2**200) == comp(2**200, -2**200, dtype=object) + + +def create_with_scalar(sctype, value): + return sctype(value) + + +def create_with_array(sctype, value): + return np.array([value], dtype=sctype) + + +@pytest.mark.parametrize("sctype", + [np.int8, np.int16, np.int32, np.int64, + np.uint8, np.uint16, np.uint32, np.uint64]) +@pytest.mark.parametrize("create", [create_with_scalar, create_with_array]) +def test_oob_creation(sctype, create): + iinfo = np.iinfo(sctype) + + with pytest.raises(OverflowError): + create(sctype, iinfo.min - 1) + + with pytest.raises(OverflowError): + create(sctype, iinfo.max + 1) + + with pytest.raises(OverflowError): + create(sctype, str(iinfo.min - 1)) + + with pytest.raises(OverflowError): + create(sctype, str(iinfo.max + 1)) + + assert create(sctype, iinfo.min) == iinfo.min + assert create(sctype, iinfo.max) == iinfo.max diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_numeric.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_numeric.py new file mode 100644 index 0000000000000000000000000000000000000000..8285e1529d509bf69446a46a239d0cb806a68605 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_numeric.py @@ -0,0 +1,4210 @@ +import sys +import warnings +import itertools +import platform +import pytest +import math +from decimal import Decimal + +import numpy as np +from numpy._core import umath, sctypes +from numpy._core.numerictypes import obj2sctype +from numpy.exceptions import AxisError +from numpy.random import rand, randint, randn +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_raises_regex, + assert_array_equal, assert_almost_equal, assert_array_almost_equal, + assert_warns, assert_array_max_ulp, HAS_REFCOUNT, IS_WASM + ) +from numpy._core._rational_tests import rational +from numpy import ma + +from hypothesis import given, strategies as st +from hypothesis.extra import numpy as hynp + + +class TestResize: + def test_copies(self): + A = np.array([[1, 2], [3, 4]]) + Ar1 = np.array([[1, 2, 3, 4], [1, 2, 3, 4]]) + assert_equal(np.resize(A, (2, 4)), Ar1) + + Ar2 = np.array([[1, 2], [3, 4], [1, 2], [3, 4]]) + assert_equal(np.resize(A, (4, 2)), Ar2) + + Ar3 = np.array([[1, 2, 3], [4, 1, 2], [3, 4, 1], [2, 3, 4]]) + assert_equal(np.resize(A, (4, 3)), Ar3) + + def test_repeats(self): + A = np.array([1, 2, 3]) + Ar1 = np.array([[1, 2, 3, 1], [2, 3, 1, 2]]) + assert_equal(np.resize(A, (2, 4)), Ar1) + + Ar2 = np.array([[1, 2], [3, 1], [2, 3], [1, 2]]) + assert_equal(np.resize(A, (4, 2)), Ar2) + + Ar3 = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) + assert_equal(np.resize(A, (4, 3)), Ar3) + + def test_zeroresize(self): + A = np.array([[1, 2], [3, 4]]) + Ar = np.resize(A, (0,)) + assert_array_equal(Ar, np.array([])) + assert_equal(A.dtype, Ar.dtype) + + Ar = np.resize(A, (0, 2)) + assert_equal(Ar.shape, (0, 2)) + + Ar = np.resize(A, (2, 0)) + assert_equal(Ar.shape, (2, 0)) + + def test_reshape_from_zero(self): + # See also gh-6740 + A = np.zeros(0, dtype=[('a', np.float32)]) + Ar = np.resize(A, (2, 1)) + assert_array_equal(Ar, np.zeros((2, 1), Ar.dtype)) + assert_equal(A.dtype, Ar.dtype) + + def test_negative_resize(self): + A = np.arange(0, 10, dtype=np.float32) + new_shape = (-10, -1) + with pytest.raises(ValueError, match=r"negative"): + np.resize(A, new_shape=new_shape) + + def test_subclass(self): + class MyArray(np.ndarray): + __array_priority__ = 1. + + my_arr = np.array([1]).view(MyArray) + assert type(np.resize(my_arr, 5)) is MyArray + assert type(np.resize(my_arr, 0)) is MyArray + + my_arr = np.array([]).view(MyArray) + assert type(np.resize(my_arr, 5)) is MyArray + + +class TestNonarrayArgs: + # check that non-array arguments to functions wrap them in arrays + def test_choose(self): + choices = [[0, 1, 2], + [3, 4, 5], + [5, 6, 7]] + tgt = [5, 1, 5] + a = [2, 0, 1] + + out = np.choose(a, choices) + assert_equal(out, tgt) + + def test_clip(self): + arr = [-1, 5, 2, 3, 10, -4, -9] + out = np.clip(arr, 2, 7) + tgt = [2, 5, 2, 3, 7, 2, 2] + assert_equal(out, tgt) + + def test_compress(self): + arr = [[0, 1, 2, 3, 4], + [5, 6, 7, 8, 9]] + tgt = [[5, 6, 7, 8, 9]] + out = np.compress([0, 1], arr, axis=0) + assert_equal(out, tgt) + + def test_count_nonzero(self): + arr = [[0, 1, 7, 0, 0], + [3, 0, 0, 2, 19]] + tgt = np.array([2, 3]) + out = np.count_nonzero(arr, axis=1) + assert_equal(out, tgt) + + def test_diagonal(self): + a = [[0, 1, 2, 3], + [4, 5, 6, 7], + [8, 9, 10, 11]] + out = np.diagonal(a) + tgt = [0, 5, 10] + + assert_equal(out, tgt) + + def test_mean(self): + A = [[1, 2, 3], [4, 5, 6]] + assert_(np.mean(A) == 3.5) + assert_(np.all(np.mean(A, 0) == np.array([2.5, 3.5, 4.5]))) + assert_(np.all(np.mean(A, 1) == np.array([2., 5.]))) + + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_(np.isnan(np.mean([]))) + assert_(w[0].category is RuntimeWarning) + + def test_ptp(self): + a = [3, 4, 5, 10, -3, -5, 6.0] + assert_equal(np.ptp(a, axis=0), 15.0) + + def test_prod(self): + arr = [[1, 2, 3, 4], + [5, 6, 7, 9], + [10, 3, 4, 5]] + tgt = [24, 1890, 600] + + assert_equal(np.prod(arr, axis=-1), tgt) + + def test_ravel(self): + a = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] + tgt = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + assert_equal(np.ravel(a), tgt) + + def test_repeat(self): + a = [1, 2, 3] + tgt = [1, 1, 2, 2, 3, 3] + + out = np.repeat(a, 2) + assert_equal(out, tgt) + + def test_reshape(self): + arr = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] + tgt = [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]] + assert_equal(np.reshape(arr, (2, 6)), tgt) + + def test_reshape_shape_arg(self): + arr = np.arange(12) + shape = (3, 4) + expected = arr.reshape(shape) + + with pytest.raises( + TypeError, + match="You cannot specify 'newshape' and 'shape' " + "arguments at the same time." + ): + np.reshape(arr, shape=shape, newshape=shape) + with pytest.raises( + TypeError, + match=r"reshape\(\) missing 1 required positional " + "argument: 'shape'" + ): + np.reshape(arr) + + assert_equal(np.reshape(arr, shape), expected) + assert_equal(np.reshape(arr, shape, order="C"), expected) + assert_equal(np.reshape(arr, shape, "C"), expected) + assert_equal(np.reshape(arr, shape=shape), expected) + assert_equal(np.reshape(arr, shape=shape, order="C"), expected) + with pytest.warns(DeprecationWarning): + actual = np.reshape(arr, newshape=shape) + assert_equal(actual, expected) + + def test_reshape_copy_arg(self): + arr = np.arange(24).reshape(2, 3, 4) + arr_f_ord = np.array(arr, order="F") + shape = (12, 2) + + assert np.shares_memory(np.reshape(arr, shape), arr) + assert np.shares_memory(np.reshape(arr, shape, order="C"), arr) + assert np.shares_memory( + np.reshape(arr_f_ord, shape, order="F"), arr_f_ord) + assert np.shares_memory(np.reshape(arr, shape, copy=None), arr) + assert np.shares_memory(np.reshape(arr, shape, copy=False), arr) + assert np.shares_memory(arr.reshape(shape, copy=False), arr) + assert not np.shares_memory(np.reshape(arr, shape, copy=True), arr) + assert not np.shares_memory( + np.reshape(arr, shape, order="C", copy=True), arr) + assert not np.shares_memory( + np.reshape(arr, shape, order="F", copy=True), arr) + assert not np.shares_memory( + np.reshape(arr, shape, order="F", copy=None), arr) + + err_msg = "Unable to avoid creating a copy while reshaping." + with pytest.raises(ValueError, match=err_msg): + np.reshape(arr, shape, order="F", copy=False) + with pytest.raises(ValueError, match=err_msg): + np.reshape(arr_f_ord, shape, order="C", copy=False) + + def test_round(self): + arr = [1.56, 72.54, 6.35, 3.25] + tgt = [1.6, 72.5, 6.4, 3.2] + assert_equal(np.around(arr, decimals=1), tgt) + s = np.float64(1.) + assert_(isinstance(s.round(), np.float64)) + assert_equal(s.round(), 1.) + + @pytest.mark.parametrize('dtype', [ + np.int8, np.int16, np.int32, np.int64, + np.uint8, np.uint16, np.uint32, np.uint64, + np.float16, np.float32, np.float64, + ]) + def test_dunder_round(self, dtype): + s = dtype(1) + assert_(isinstance(round(s), int)) + assert_(isinstance(round(s, None), int)) + assert_(isinstance(round(s, ndigits=None), int)) + assert_equal(round(s), 1) + assert_equal(round(s, None), 1) + assert_equal(round(s, ndigits=None), 1) + + @pytest.mark.parametrize('val, ndigits', [ + pytest.param(2**31 - 1, -1, + marks=pytest.mark.skip(reason="Out of range of int32") + ), + (2**31 - 1, 1-math.ceil(math.log10(2**31 - 1))), + (2**31 - 1, -math.ceil(math.log10(2**31 - 1))) + ]) + def test_dunder_round_edgecases(self, val, ndigits): + assert_equal(round(val, ndigits), round(np.int32(val), ndigits)) + + def test_dunder_round_accuracy(self): + f = np.float64(5.1 * 10**73) + assert_(isinstance(round(f, -73), np.float64)) + assert_array_max_ulp(round(f, -73), 5.0 * 10**73) + assert_(isinstance(round(f, ndigits=-73), np.float64)) + assert_array_max_ulp(round(f, ndigits=-73), 5.0 * 10**73) + + i = np.int64(501) + assert_(isinstance(round(i, -2), np.int64)) + assert_array_max_ulp(round(i, -2), 500) + assert_(isinstance(round(i, ndigits=-2), np.int64)) + assert_array_max_ulp(round(i, ndigits=-2), 500) + + @pytest.mark.xfail(raises=AssertionError, reason="gh-15896") + def test_round_py_consistency(self): + f = 5.1 * 10**73 + assert_equal(round(np.float64(f), -73), round(f, -73)) + + def test_searchsorted(self): + arr = [-8, -5, -1, 3, 6, 10] + out = np.searchsorted(arr, 0) + assert_equal(out, 3) + + def test_size(self): + A = [[1, 2, 3], [4, 5, 6]] + assert_(np.size(A) == 6) + assert_(np.size(A, 0) == 2) + assert_(np.size(A, 1) == 3) + + def test_squeeze(self): + A = [[[1, 1, 1], [2, 2, 2], [3, 3, 3]]] + assert_equal(np.squeeze(A).shape, (3, 3)) + assert_equal(np.squeeze(np.zeros((1, 3, 1))).shape, (3,)) + assert_equal(np.squeeze(np.zeros((1, 3, 1)), axis=0).shape, (3, 1)) + assert_equal(np.squeeze(np.zeros((1, 3, 1)), axis=-1).shape, (1, 3)) + assert_equal(np.squeeze(np.zeros((1, 3, 1)), axis=2).shape, (1, 3)) + assert_equal(np.squeeze([np.zeros((3, 1))]).shape, (3,)) + assert_equal(np.squeeze([np.zeros((3, 1))], axis=0).shape, (3, 1)) + assert_equal(np.squeeze([np.zeros((3, 1))], axis=2).shape, (1, 3)) + assert_equal(np.squeeze([np.zeros((3, 1))], axis=-1).shape, (1, 3)) + + def test_std(self): + A = [[1, 2, 3], [4, 5, 6]] + assert_almost_equal(np.std(A), 1.707825127659933) + assert_almost_equal(np.std(A, 0), np.array([1.5, 1.5, 1.5])) + assert_almost_equal(np.std(A, 1), np.array([0.81649658, 0.81649658])) + + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_(np.isnan(np.std([]))) + assert_(w[0].category is RuntimeWarning) + + def test_swapaxes(self): + tgt = [[[0, 4], [2, 6]], [[1, 5], [3, 7]]] + a = [[[0, 1], [2, 3]], [[4, 5], [6, 7]]] + out = np.swapaxes(a, 0, 2) + assert_equal(out, tgt) + + def test_sum(self): + m = [[1, 2, 3], + [4, 5, 6], + [7, 8, 9]] + tgt = [[6], [15], [24]] + out = np.sum(m, axis=1, keepdims=True) + + assert_equal(tgt, out) + + def test_take(self): + tgt = [2, 3, 5] + indices = [1, 2, 4] + a = [1, 2, 3, 4, 5] + + out = np.take(a, indices) + assert_equal(out, tgt) + + pairs = [ + (np.int32, np.int32), (np.int32, np.int64), + (np.int64, np.int32), (np.int64, np.int64) + ] + for array_type, indices_type in pairs: + x = np.array([1, 2, 3, 4, 5], dtype=array_type) + ind = np.array([0, 2, 2, 3], dtype=indices_type) + tgt = np.array([1, 3, 3, 4], dtype=array_type) + out = np.take(x, ind) + assert_equal(out, tgt) + assert_equal(out.dtype, tgt.dtype) + + def test_trace(self): + c = [[1, 2], [3, 4], [5, 6]] + assert_equal(np.trace(c), 5) + + def test_transpose(self): + arr = [[1, 2], [3, 4], [5, 6]] + tgt = [[1, 3, 5], [2, 4, 6]] + assert_equal(np.transpose(arr, (1, 0)), tgt) + assert_equal(np.transpose(arr, (-1, -2)), tgt) + assert_equal(np.matrix_transpose(arr), tgt) + + def test_var(self): + A = [[1, 2, 3], [4, 5, 6]] + assert_almost_equal(np.var(A), 2.9166666666666665) + assert_almost_equal(np.var(A, 0), np.array([2.25, 2.25, 2.25])) + assert_almost_equal(np.var(A, 1), np.array([0.66666667, 0.66666667])) + + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_(np.isnan(np.var([]))) + assert_(w[0].category is RuntimeWarning) + + B = np.array([None, 0]) + B[0] = 1j + assert_almost_equal(np.var(B), 0.25) + + def test_std_with_mean_keyword(self): + # Setting the seed to make the test reproducible + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + mean_out = np.zeros((10, 1, 5)) + std_out = np.zeros((10, 1, 5)) + + mean = np.mean(A, + out=mean_out, + axis=1, + keepdims=True) + + # The returned object should be the object specified during calling + assert mean_out is mean + + std = np.std(A, + out=std_out, + axis=1, + keepdims=True, + mean=mean) + + # The returned object should be the object specified during calling + assert std_out is std + + # Shape of returned mean and std should be same + assert std.shape == mean.shape + assert std.shape == (10, 1, 5) + + # Output should be the same as from the individual algorithms + std_old = np.std(A, axis=1, keepdims=True) + + assert std_old.shape == mean.shape + assert_almost_equal(std, std_old) + + def test_var_with_mean_keyword(self): + # Setting the seed to make the test reproducible + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + mean_out = np.zeros((10, 1, 5)) + var_out = np.zeros((10, 1, 5)) + + mean = np.mean(A, + out=mean_out, + axis=1, + keepdims=True) + + # The returned object should be the object specified during calling + assert mean_out is mean + + var = np.var(A, + out=var_out, + axis=1, + keepdims=True, + mean=mean) + + # The returned object should be the object specified during calling + assert var_out is var + + # Shape of returned mean and var should be same + assert var.shape == mean.shape + assert var.shape == (10, 1, 5) + + # Output should be the same as from the individual algorithms + var_old = np.var(A, axis=1, keepdims=True) + + assert var_old.shape == mean.shape + assert_almost_equal(var, var_old) + + def test_std_with_mean_keyword_keepdims_false(self): + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + mean = np.mean(A, + axis=1, + keepdims=True) + + std = np.std(A, + axis=1, + keepdims=False, + mean=mean) + + # Shape of returned mean and std should be same + assert std.shape == (10, 5) + + # Output should be the same as from the individual algorithms + std_old = np.std(A, axis=1, keepdims=False) + mean_old = np.mean(A, axis=1, keepdims=False) + + assert std_old.shape == mean_old.shape + assert_equal(std, std_old) + + def test_var_with_mean_keyword_keepdims_false(self): + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + mean = np.mean(A, + axis=1, + keepdims=True) + + var = np.var(A, + axis=1, + keepdims=False, + mean=mean) + + # Shape of returned mean and var should be same + assert var.shape == (10, 5) + + # Output should be the same as from the individual algorithms + var_old = np.var(A, axis=1, keepdims=False) + mean_old = np.mean(A, axis=1, keepdims=False) + + assert var_old.shape == mean_old.shape + assert_equal(var, var_old) + + def test_std_with_mean_keyword_where_nontrivial(self): + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + where = A > 0.5 + + mean = np.mean(A, + axis=1, + keepdims=True, + where=where) + + std = np.std(A, + axis=1, + keepdims=False, + mean=mean, + where=where) + + # Shape of returned mean and std should be same + assert std.shape == (10, 5) + + # Output should be the same as from the individual algorithms + std_old = np.std(A, axis=1, where=where) + mean_old = np.mean(A, axis=1, where=where) + + assert std_old.shape == mean_old.shape + assert_equal(std, std_old) + + def test_var_with_mean_keyword_where_nontrivial(self): + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + where = A > 0.5 + + mean = np.mean(A, + axis=1, + keepdims=True, + where=where) + + var = np.var(A, + axis=1, + keepdims=False, + mean=mean, + where=where) + + # Shape of returned mean and var should be same + assert var.shape == (10, 5) + + # Output should be the same as from the individual algorithms + var_old = np.var(A, axis=1, where=where) + mean_old = np.mean(A, axis=1, where=where) + + assert var_old.shape == mean_old.shape + assert_equal(var, var_old) + + def test_std_with_mean_keyword_multiple_axis(self): + # Setting the seed to make the test reproducible + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + axis = (0, 2) + + mean = np.mean(A, + out=None, + axis=axis, + keepdims=True) + + std = np.std(A, + out=None, + axis=axis, + keepdims=False, + mean=mean) + + # Shape of returned mean and std should be same + assert std.shape == (20,) + + # Output should be the same as from the individual algorithms + std_old = np.std(A, axis=axis, keepdims=False) + + assert_almost_equal(std, std_old) + + def test_std_with_mean_keyword_axis_None(self): + # Setting the seed to make the test reproducible + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + + axis = None + + mean = np.mean(A, + out=None, + axis=axis, + keepdims=True) + + std = np.std(A, + out=None, + axis=axis, + keepdims=False, + mean=mean) + + # Shape of returned mean and std should be same + assert std.shape == () + + # Output should be the same as from the individual algorithms + std_old = np.std(A, axis=axis, keepdims=False) + + assert_almost_equal(std, std_old) + + def test_std_with_mean_keyword_keepdims_true_masked(self): + + A = ma.array([[2., 3., 4., 5.], + [1., 2., 3., 4.]], + mask=[[True, False, True, False], + [True, False, True, False]]) + + B = ma.array([[100., 3., 104., 5.], + [101., 2., 103., 4.]], + mask=[[True, False, True, False], + [True, False, True, False]]) + + mean_out = ma.array([[0., 0., 0., 0.]], + mask=[[False, False, False, False]]) + std_out = ma.array([[0., 0., 0., 0.]], + mask=[[False, False, False, False]]) + + axis = 0 + + mean = np.mean(A, out=mean_out, + axis=axis, keepdims=True) + + std = np.std(A, out=std_out, + axis=axis, keepdims=True, + mean=mean) + + # Shape of returned mean and std should be same + assert std.shape == mean.shape + assert std.shape == (1, 4) + + # Output should be the same as from the individual algorithms + std_old = np.std(A, axis=axis, keepdims=True) + mean_old = np.mean(A, axis=axis, keepdims=True) + + assert std_old.shape == mean_old.shape + assert_almost_equal(std, std_old) + assert_almost_equal(mean, mean_old) + + assert mean_out is mean + assert std_out is std + + # masked elements should be ignored + mean_b = np.mean(B, axis=axis, keepdims=True) + std_b = np.std(B, axis=axis, keepdims=True, mean=mean_b) + assert_almost_equal(std, std_b) + assert_almost_equal(mean, mean_b) + + def test_var_with_mean_keyword_keepdims_true_masked(self): + + A = ma.array([[2., 3., 4., 5.], + [1., 2., 3., 4.]], + mask=[[True, False, True, False], + [True, False, True, False]]) + + B = ma.array([[100., 3., 104., 5.], + [101., 2., 103., 4.]], + mask=[[True, False, True, False], + [True, False, True, False]]) + + mean_out = ma.array([[0., 0., 0., 0.]], + mask=[[False, False, False, False]]) + var_out = ma.array([[0., 0., 0., 0.]], + mask=[[False, False, False, False]]) + + axis = 0 + + mean = np.mean(A, out=mean_out, + axis=axis, keepdims=True) + + var = np.var(A, out=var_out, + axis=axis, keepdims=True, + mean=mean) + + # Shape of returned mean and var should be same + assert var.shape == mean.shape + assert var.shape == (1, 4) + + # Output should be the same as from the individual algorithms + var_old = np.var(A, axis=axis, keepdims=True) + mean_old = np.mean(A, axis=axis, keepdims=True) + + assert var_old.shape == mean_old.shape + assert_almost_equal(var, var_old) + assert_almost_equal(mean, mean_old) + + assert mean_out is mean + assert var_out is var + + # masked elements should be ignored + mean_b = np.mean(B, axis=axis, keepdims=True) + var_b = np.var(B, axis=axis, keepdims=True, mean=mean_b) + assert_almost_equal(var, var_b) + assert_almost_equal(mean, mean_b) + + +class TestIsscalar: + def test_isscalar(self): + assert_(np.isscalar(3.1)) + assert_(np.isscalar(np.int16(12345))) + assert_(np.isscalar(False)) + assert_(np.isscalar('numpy')) + assert_(not np.isscalar([3.1])) + assert_(not np.isscalar(None)) + + # PEP 3141 + from fractions import Fraction + assert_(np.isscalar(Fraction(5, 17))) + from numbers import Number + assert_(np.isscalar(Number())) + + +class TestBoolScalar: + def test_logical(self): + f = np.False_ + t = np.True_ + s = "xyz" + assert_((t and s) is s) + assert_((f and s) is f) + + def test_bitwise_or(self): + f = np.False_ + t = np.True_ + assert_((t | t) is t) + assert_((f | t) is t) + assert_((t | f) is t) + assert_((f | f) is f) + + def test_bitwise_and(self): + f = np.False_ + t = np.True_ + assert_((t & t) is t) + assert_((f & t) is f) + assert_((t & f) is f) + assert_((f & f) is f) + + def test_bitwise_xor(self): + f = np.False_ + t = np.True_ + assert_((t ^ t) is f) + assert_((f ^ t) is t) + assert_((t ^ f) is t) + assert_((f ^ f) is f) + + +class TestBoolArray: + def setup_method(self): + # offset for simd tests + self.t = np.array([True] * 41, dtype=bool)[1::] + self.f = np.array([False] * 41, dtype=bool)[1::] + self.o = np.array([False] * 42, dtype=bool)[2::] + self.nm = self.f.copy() + self.im = self.t.copy() + self.nm[3] = True + self.nm[-2] = True + self.im[3] = False + self.im[-2] = False + + def test_all_any(self): + assert_(self.t.all()) + assert_(self.t.any()) + assert_(not self.f.all()) + assert_(not self.f.any()) + assert_(self.nm.any()) + assert_(self.im.any()) + assert_(not self.nm.all()) + assert_(not self.im.all()) + # check bad element in all positions + for i in range(256 - 7): + d = np.array([False] * 256, dtype=bool)[7::] + d[i] = True + assert_(np.any(d)) + e = np.array([True] * 256, dtype=bool)[7::] + e[i] = False + assert_(not np.all(e)) + assert_array_equal(e, ~d) + # big array test for blocked libc loops + for i in list(range(9, 6000, 507)) + [7764, 90021, -10]: + d = np.array([False] * 100043, dtype=bool) + d[i] = True + assert_(np.any(d), msg="%r" % i) + e = np.array([True] * 100043, dtype=bool) + e[i] = False + assert_(not np.all(e), msg="%r" % i) + + def test_logical_not_abs(self): + assert_array_equal(~self.t, self.f) + assert_array_equal(np.abs(~self.t), self.f) + assert_array_equal(np.abs(~self.f), self.t) + assert_array_equal(np.abs(self.f), self.f) + assert_array_equal(~np.abs(self.f), self.t) + assert_array_equal(~np.abs(self.t), self.f) + assert_array_equal(np.abs(~self.nm), self.im) + np.logical_not(self.t, out=self.o) + assert_array_equal(self.o, self.f) + np.abs(self.t, out=self.o) + assert_array_equal(self.o, self.t) + + def test_logical_and_or_xor(self): + assert_array_equal(self.t | self.t, self.t) + assert_array_equal(self.f | self.f, self.f) + assert_array_equal(self.t | self.f, self.t) + assert_array_equal(self.f | self.t, self.t) + np.logical_or(self.t, self.t, out=self.o) + assert_array_equal(self.o, self.t) + assert_array_equal(self.t & self.t, self.t) + assert_array_equal(self.f & self.f, self.f) + assert_array_equal(self.t & self.f, self.f) + assert_array_equal(self.f & self.t, self.f) + np.logical_and(self.t, self.t, out=self.o) + assert_array_equal(self.o, self.t) + assert_array_equal(self.t ^ self.t, self.f) + assert_array_equal(self.f ^ self.f, self.f) + assert_array_equal(self.t ^ self.f, self.t) + assert_array_equal(self.f ^ self.t, self.t) + np.logical_xor(self.t, self.t, out=self.o) + assert_array_equal(self.o, self.f) + + assert_array_equal(self.nm & self.t, self.nm) + assert_array_equal(self.im & self.f, False) + assert_array_equal(self.nm & True, self.nm) + assert_array_equal(self.im & False, self.f) + assert_array_equal(self.nm | self.t, self.t) + assert_array_equal(self.im | self.f, self.im) + assert_array_equal(self.nm | True, self.t) + assert_array_equal(self.im | False, self.im) + assert_array_equal(self.nm ^ self.t, self.im) + assert_array_equal(self.im ^ self.f, self.im) + assert_array_equal(self.nm ^ True, self.im) + assert_array_equal(self.im ^ False, self.im) + + +class TestBoolCmp: + def setup_method(self): + self.f = np.ones(256, dtype=np.float32) + self.ef = np.ones(self.f.size, dtype=bool) + self.d = np.ones(128, dtype=np.float64) + self.ed = np.ones(self.d.size, dtype=bool) + # generate values for all permutation of 256bit simd vectors + s = 0 + for i in range(32): + self.f[s:s+8] = [i & 2**x for x in range(8)] + self.ef[s:s+8] = [(i & 2**x) != 0 for x in range(8)] + s += 8 + s = 0 + for i in range(16): + self.d[s:s+4] = [i & 2**x for x in range(4)] + self.ed[s:s+4] = [(i & 2**x) != 0 for x in range(4)] + s += 4 + + self.nf = self.f.copy() + self.nd = self.d.copy() + self.nf[self.ef] = np.nan + self.nd[self.ed] = np.nan + + self.inff = self.f.copy() + self.infd = self.d.copy() + self.inff[::3][self.ef[::3]] = np.inf + self.infd[::3][self.ed[::3]] = np.inf + self.inff[1::3][self.ef[1::3]] = -np.inf + self.infd[1::3][self.ed[1::3]] = -np.inf + self.inff[2::3][self.ef[2::3]] = np.nan + self.infd[2::3][self.ed[2::3]] = np.nan + self.efnonan = self.ef.copy() + self.efnonan[2::3] = False + self.ednonan = self.ed.copy() + self.ednonan[2::3] = False + + self.signf = self.f.copy() + self.signd = self.d.copy() + self.signf[self.ef] *= -1. + self.signd[self.ed] *= -1. + self.signf[1::6][self.ef[1::6]] = -np.inf + self.signd[1::6][self.ed[1::6]] = -np.inf + # On RISC-V, many operations that produce NaNs, such as converting + # a -NaN from f64 to f32, return a canonical NaN. The canonical + # NaNs are always positive. See section 11.3 NaN Generation and + # Propagation of the RISC-V Unprivileged ISA for more details. + # We disable the float32 sign test on riscv64 for -np.nan as the sign + # of the NaN will be lost when it's converted to a float32. + if platform.machine() != 'riscv64': + self.signf[3::6][self.ef[3::6]] = -np.nan + self.signd[3::6][self.ed[3::6]] = -np.nan + self.signf[4::6][self.ef[4::6]] = -0. + self.signd[4::6][self.ed[4::6]] = -0. + + def test_float(self): + # offset for alignment test + for i in range(4): + assert_array_equal(self.f[i:] > 0, self.ef[i:]) + assert_array_equal(self.f[i:] - 1 >= 0, self.ef[i:]) + assert_array_equal(self.f[i:] == 0, ~self.ef[i:]) + assert_array_equal(-self.f[i:] < 0, self.ef[i:]) + assert_array_equal(-self.f[i:] + 1 <= 0, self.ef[i:]) + r = self.f[i:] != 0 + assert_array_equal(r, self.ef[i:]) + r2 = self.f[i:] != np.zeros_like(self.f[i:]) + r3 = 0 != self.f[i:] + assert_array_equal(r, r2) + assert_array_equal(r, r3) + # check bool == 0x1 + assert_array_equal(r.view(np.int8), r.astype(np.int8)) + assert_array_equal(r2.view(np.int8), r2.astype(np.int8)) + assert_array_equal(r3.view(np.int8), r3.astype(np.int8)) + + # isnan on amd64 takes the same code path + assert_array_equal(np.isnan(self.nf[i:]), self.ef[i:]) + assert_array_equal(np.isfinite(self.nf[i:]), ~self.ef[i:]) + assert_array_equal(np.isfinite(self.inff[i:]), ~self.ef[i:]) + assert_array_equal(np.isinf(self.inff[i:]), self.efnonan[i:]) + assert_array_equal(np.signbit(self.signf[i:]), self.ef[i:]) + + def test_double(self): + # offset for alignment test + for i in range(2): + assert_array_equal(self.d[i:] > 0, self.ed[i:]) + assert_array_equal(self.d[i:] - 1 >= 0, self.ed[i:]) + assert_array_equal(self.d[i:] == 0, ~self.ed[i:]) + assert_array_equal(-self.d[i:] < 0, self.ed[i:]) + assert_array_equal(-self.d[i:] + 1 <= 0, self.ed[i:]) + r = self.d[i:] != 0 + assert_array_equal(r, self.ed[i:]) + r2 = self.d[i:] != np.zeros_like(self.d[i:]) + r3 = 0 != self.d[i:] + assert_array_equal(r, r2) + assert_array_equal(r, r3) + # check bool == 0x1 + assert_array_equal(r.view(np.int8), r.astype(np.int8)) + assert_array_equal(r2.view(np.int8), r2.astype(np.int8)) + assert_array_equal(r3.view(np.int8), r3.astype(np.int8)) + + # isnan on amd64 takes the same code path + assert_array_equal(np.isnan(self.nd[i:]), self.ed[i:]) + assert_array_equal(np.isfinite(self.nd[i:]), ~self.ed[i:]) + assert_array_equal(np.isfinite(self.infd[i:]), ~self.ed[i:]) + assert_array_equal(np.isinf(self.infd[i:]), self.ednonan[i:]) + assert_array_equal(np.signbit(self.signd[i:]), self.ed[i:]) + + +class TestSeterr: + def test_default(self): + err = np.geterr() + assert_equal(err, + dict(divide='warn', + invalid='warn', + over='warn', + under='ignore') + ) + + def test_set(self): + with np.errstate(): + err = np.seterr() + old = np.seterr(divide='print') + assert_(err == old) + new = np.seterr() + assert_(new['divide'] == 'print') + np.seterr(over='raise') + assert_(np.geterr()['over'] == 'raise') + assert_(new['divide'] == 'print') + np.seterr(**old) + assert_(np.geterr() == old) + + @pytest.mark.skipif(IS_WASM, reason="no wasm fp exception support") + @pytest.mark.skipif(platform.machine() == "armv5tel", reason="See gh-413.") + def test_divide_err(self): + with np.errstate(divide='raise'): + with assert_raises(FloatingPointError): + np.array([1.]) / np.array([0.]) + + np.seterr(divide='ignore') + np.array([1.]) / np.array([0.]) + + +class TestFloatExceptions: + def assert_raises_fpe(self, fpeerr, flop, x, y): + ftype = type(x) + try: + flop(x, y) + assert_(False, + "Type %s did not raise fpe error '%s'." % (ftype, fpeerr)) + except FloatingPointError as exc: + assert_(str(exc).find(fpeerr) >= 0, + "Type %s raised wrong fpe error '%s'." % (ftype, exc)) + + def assert_op_raises_fpe(self, fpeerr, flop, sc1, sc2): + # Check that fpe exception is raised. + # + # Given a floating operation `flop` and two scalar values, check that + # the operation raises the floating point exception specified by + # `fpeerr`. Tests all variants with 0-d array scalars as well. + + self.assert_raises_fpe(fpeerr, flop, sc1, sc2) + self.assert_raises_fpe(fpeerr, flop, sc1[()], sc2) + self.assert_raises_fpe(fpeerr, flop, sc1, sc2[()]) + self.assert_raises_fpe(fpeerr, flop, sc1[()], sc2[()]) + + # Test for all real and complex float types + @pytest.mark.skipif(IS_WASM, reason="no wasm fp exception support") + @pytest.mark.parametrize("typecode", np.typecodes["AllFloat"]) + def test_floating_exceptions(self, typecode): + if 'bsd' in sys.platform and typecode in 'gG': + pytest.skip(reason="Fallback impl for (c)longdouble may not raise " + "FPE errors as expected on BSD OSes, " + "see gh-24876, gh-23379") + + # Test basic arithmetic function errors + with np.errstate(all='raise'): + ftype = obj2sctype(typecode) + if np.dtype(ftype).kind == 'f': + # Get some extreme values for the type + fi = np.finfo(ftype) + ft_tiny = fi._machar.tiny + ft_max = fi.max + ft_eps = fi.eps + underflow = 'underflow' + divbyzero = 'divide by zero' + else: + # 'c', complex, corresponding real dtype + rtype = type(ftype(0).real) + fi = np.finfo(rtype) + ft_tiny = ftype(fi._machar.tiny) + ft_max = ftype(fi.max) + ft_eps = ftype(fi.eps) + # The complex types raise different exceptions + underflow = '' + divbyzero = '' + overflow = 'overflow' + invalid = 'invalid' + + # The value of tiny for double double is NaN, so we need to + # pass the assert + if not np.isnan(ft_tiny): + self.assert_raises_fpe(underflow, + lambda a, b: a/b, ft_tiny, ft_max) + self.assert_raises_fpe(underflow, + lambda a, b: a*b, ft_tiny, ft_tiny) + self.assert_raises_fpe(overflow, + lambda a, b: a*b, ft_max, ftype(2)) + self.assert_raises_fpe(overflow, + lambda a, b: a/b, ft_max, ftype(0.5)) + self.assert_raises_fpe(overflow, + lambda a, b: a+b, ft_max, ft_max*ft_eps) + self.assert_raises_fpe(overflow, + lambda a, b: a-b, -ft_max, ft_max*ft_eps) + self.assert_raises_fpe(overflow, + np.power, ftype(2), ftype(2**fi.nexp)) + self.assert_raises_fpe(divbyzero, + lambda a, b: a/b, ftype(1), ftype(0)) + self.assert_raises_fpe( + invalid, lambda a, b: a/b, ftype(np.inf), ftype(np.inf) + ) + self.assert_raises_fpe(invalid, + lambda a, b: a/b, ftype(0), ftype(0)) + self.assert_raises_fpe( + invalid, lambda a, b: a-b, ftype(np.inf), ftype(np.inf) + ) + self.assert_raises_fpe( + invalid, lambda a, b: a+b, ftype(np.inf), ftype(-np.inf) + ) + self.assert_raises_fpe(invalid, + lambda a, b: a*b, ftype(0), ftype(np.inf)) + + @pytest.mark.skipif(IS_WASM, reason="no wasm fp exception support") + def test_warnings(self): + # test warning code path + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + with np.errstate(all="warn"): + np.divide(1, 0.) + assert_equal(len(w), 1) + assert_("divide by zero" in str(w[0].message)) + np.array(1e300) * np.array(1e300) + assert_equal(len(w), 2) + assert_("overflow" in str(w[-1].message)) + np.array(np.inf) - np.array(np.inf) + assert_equal(len(w), 3) + assert_("invalid value" in str(w[-1].message)) + np.array(1e-300) * np.array(1e-300) + assert_equal(len(w), 4) + assert_("underflow" in str(w[-1].message)) + + +class TestTypes: + def check_promotion_cases(self, promote_func): + # tests that the scalars get coerced correctly. + b = np.bool(0) + i8, i16, i32, i64 = np.int8(0), np.int16(0), np.int32(0), np.int64(0) + u8, u16, u32, u64 = np.uint8(0), np.uint16(0), np.uint32(0), np.uint64(0) + f32, f64, fld = np.float32(0), np.float64(0), np.longdouble(0) + c64, c128, cld = np.complex64(0), np.complex128(0), np.clongdouble(0) + + # coercion within the same kind + assert_equal(promote_func(i8, i16), np.dtype(np.int16)) + assert_equal(promote_func(i32, i8), np.dtype(np.int32)) + assert_equal(promote_func(i16, i64), np.dtype(np.int64)) + assert_equal(promote_func(u8, u32), np.dtype(np.uint32)) + assert_equal(promote_func(f32, f64), np.dtype(np.float64)) + assert_equal(promote_func(fld, f32), np.dtype(np.longdouble)) + assert_equal(promote_func(f64, fld), np.dtype(np.longdouble)) + assert_equal(promote_func(c128, c64), np.dtype(np.complex128)) + assert_equal(promote_func(cld, c128), np.dtype(np.clongdouble)) + assert_equal(promote_func(c64, fld), np.dtype(np.clongdouble)) + + # coercion between kinds + assert_equal(promote_func(b, i32), np.dtype(np.int32)) + assert_equal(promote_func(b, u8), np.dtype(np.uint8)) + assert_equal(promote_func(i8, u8), np.dtype(np.int16)) + assert_equal(promote_func(u8, i32), np.dtype(np.int32)) + assert_equal(promote_func(i64, u32), np.dtype(np.int64)) + assert_equal(promote_func(u64, i32), np.dtype(np.float64)) + assert_equal(promote_func(i32, f32), np.dtype(np.float64)) + assert_equal(promote_func(i64, f32), np.dtype(np.float64)) + assert_equal(promote_func(f32, i16), np.dtype(np.float32)) + assert_equal(promote_func(f32, u32), np.dtype(np.float64)) + assert_equal(promote_func(f32, c64), np.dtype(np.complex64)) + assert_equal(promote_func(c128, f32), np.dtype(np.complex128)) + assert_equal(promote_func(cld, f64), np.dtype(np.clongdouble)) + + # coercion between scalars and 1-D arrays + assert_equal(promote_func(np.array([b]), i8), np.dtype(np.int8)) + assert_equal(promote_func(np.array([b]), u8), np.dtype(np.uint8)) + assert_equal(promote_func(np.array([b]), i32), np.dtype(np.int32)) + assert_equal(promote_func(np.array([b]), u32), np.dtype(np.uint32)) + assert_equal(promote_func(np.array([i8]), i64), np.dtype(np.int64)) + # unsigned and signed unfortunately tend to promote to float64: + assert_equal(promote_func(u64, np.array([i32])), np.dtype(np.float64)) + assert_equal(promote_func(i64, np.array([u32])), np.dtype(np.int64)) + assert_equal(promote_func(np.array([u16]), i32), np.dtype(np.int32)) + assert_equal(promote_func(np.int32(-1), np.array([u64])), + np.dtype(np.float64)) + assert_equal(promote_func(f64, np.array([f32])), np.dtype(np.float64)) + assert_equal(promote_func(fld, np.array([f32])), + np.dtype(np.longdouble)) + assert_equal(promote_func(np.array([f64]), fld), + np.dtype(np.longdouble)) + assert_equal(promote_func(fld, np.array([c64])), + np.dtype(np.clongdouble)) + assert_equal(promote_func(c64, np.array([f64])), + np.dtype(np.complex128)) + assert_equal(promote_func(np.complex64(3j), np.array([f64])), + np.dtype(np.complex128)) + assert_equal(promote_func(np.array([f32]), c128), + np.dtype(np.complex128)) + + # coercion between scalars and 1-D arrays, where + # the scalar has greater kind than the array + assert_equal(promote_func(np.array([b]), f64), np.dtype(np.float64)) + assert_equal(promote_func(np.array([b]), i64), np.dtype(np.int64)) + assert_equal(promote_func(np.array([b]), u64), np.dtype(np.uint64)) + assert_equal(promote_func(np.array([i8]), f64), np.dtype(np.float64)) + assert_equal(promote_func(np.array([u16]), f64), np.dtype(np.float64)) + + + def test_coercion(self): + def res_type(a, b): + return np.add(a, b).dtype + + self.check_promotion_cases(res_type) + + # Use-case: float/complex scalar * bool/int8 array + # shouldn't narrow the float/complex type + for a in [np.array([True, False]), np.array([-3, 12], dtype=np.int8)]: + b = 1.234 * a + assert_equal(b.dtype, np.dtype('f8'), "array type %s" % a.dtype) + b = np.longdouble(1.234) * a + assert_equal(b.dtype, np.dtype(np.longdouble), + "array type %s" % a.dtype) + b = np.float64(1.234) * a + assert_equal(b.dtype, np.dtype('f8'), "array type %s" % a.dtype) + b = np.float32(1.234) * a + assert_equal(b.dtype, np.dtype('f4'), "array type %s" % a.dtype) + b = np.float16(1.234) * a + assert_equal(b.dtype, np.dtype('f2'), "array type %s" % a.dtype) + + b = 1.234j * a + assert_equal(b.dtype, np.dtype('c16'), "array type %s" % a.dtype) + b = np.clongdouble(1.234j) * a + assert_equal(b.dtype, np.dtype(np.clongdouble), + "array type %s" % a.dtype) + b = np.complex128(1.234j) * a + assert_equal(b.dtype, np.dtype('c16'), "array type %s" % a.dtype) + b = np.complex64(1.234j) * a + assert_equal(b.dtype, np.dtype('c8'), "array type %s" % a.dtype) + + # The following use-case is problematic, and to resolve its + # tricky side-effects requires more changes. + # + # Use-case: (1-t)*a, where 't' is a boolean array and 'a' is + # a float32, shouldn't promote to float64 + # + # a = np.array([1.0, 1.5], dtype=np.float32) + # t = np.array([True, False]) + # b = t*a + # assert_equal(b, [1.0, 0.0]) + # assert_equal(b.dtype, np.dtype('f4')) + # b = (1-t)*a + # assert_equal(b, [0.0, 1.5]) + # assert_equal(b.dtype, np.dtype('f4')) + # + # Probably ~t (bitwise negation) is more proper to use here, + # but this is arguably less intuitive to understand at a glance, and + # would fail if 't' is actually an integer array instead of boolean: + # + # b = (~t)*a + # assert_equal(b, [0.0, 1.5]) + # assert_equal(b.dtype, np.dtype('f4')) + + def test_result_type(self): + self.check_promotion_cases(np.result_type) + assert_(np.result_type(None) == np.dtype(None)) + + def test_promote_types_endian(self): + # promote_types should always return native-endian types + assert_equal(np.promote_types('i8', '>i8'), np.dtype('i8')) + + assert_equal(np.promote_types('>i8', '>U16'), np.dtype('U21')) + assert_equal(np.promote_types('U16', '>i8'), np.dtype('U21')) + assert_equal(np.promote_types('S5', '>U8'), np.dtype('U8')) + assert_equal(np.promote_types('U8', '>S5'), np.dtype('U8')) + assert_equal(np.promote_types('U8', '>U5'), np.dtype('U8')) + + assert_equal(np.promote_types('M8', '>M8'), np.dtype('M8')) + assert_equal(np.promote_types('m8', '>m8'), np.dtype('m8')) + + def test_can_cast_and_promote_usertypes(self): + # The rational type defines safe casting for signed integers, + # boolean. Rational itself *does* cast safely to double. + # (rational does not actually cast to all signed integers, e.g. + # int64 can be both long and longlong and it registers only the first) + valid_types = ["int8", "int16", "int32", "int64", "bool"] + invalid_types = "BHILQP" + "FDG" + "mM" + "f" + "V" + + rational_dt = np.dtype(rational) + for numpy_dtype in valid_types: + numpy_dtype = np.dtype(numpy_dtype) + assert np.can_cast(numpy_dtype, rational_dt) + assert np.promote_types(numpy_dtype, rational_dt) is rational_dt + + for numpy_dtype in invalid_types: + numpy_dtype = np.dtype(numpy_dtype) + assert not np.can_cast(numpy_dtype, rational_dt) + with pytest.raises(TypeError): + np.promote_types(numpy_dtype, rational_dt) + + double_dt = np.dtype("double") + assert np.can_cast(rational_dt, double_dt) + assert np.promote_types(double_dt, rational_dt) is double_dt + + @pytest.mark.parametrize("swap", ["", "swap"]) + @pytest.mark.parametrize("string_dtype", ["U", "S"]) + def test_promote_types_strings(self, swap, string_dtype): + if swap == "swap": + promote_types = lambda a, b: np.promote_types(b, a) + else: + promote_types = np.promote_types + + S = string_dtype + + # Promote numeric with unsized string: + assert_equal(promote_types('bool', S), np.dtype(S+'5')) + assert_equal(promote_types('b', S), np.dtype(S+'4')) + assert_equal(promote_types('u1', S), np.dtype(S+'3')) + assert_equal(promote_types('u2', S), np.dtype(S+'5')) + assert_equal(promote_types('u4', S), np.dtype(S+'10')) + assert_equal(promote_types('u8', S), np.dtype(S+'20')) + assert_equal(promote_types('i1', S), np.dtype(S+'4')) + assert_equal(promote_types('i2', S), np.dtype(S+'6')) + assert_equal(promote_types('i4', S), np.dtype(S+'11')) + assert_equal(promote_types('i8', S), np.dtype(S+'21')) + # Promote numeric with sized string: + assert_equal(promote_types('bool', S+'1'), np.dtype(S+'5')) + assert_equal(promote_types('bool', S+'30'), np.dtype(S+'30')) + assert_equal(promote_types('b', S+'1'), np.dtype(S+'4')) + assert_equal(promote_types('b', S+'30'), np.dtype(S+'30')) + assert_equal(promote_types('u1', S+'1'), np.dtype(S+'3')) + assert_equal(promote_types('u1', S+'30'), np.dtype(S+'30')) + assert_equal(promote_types('u2', S+'1'), np.dtype(S+'5')) + assert_equal(promote_types('u2', S+'30'), np.dtype(S+'30')) + assert_equal(promote_types('u4', S+'1'), np.dtype(S+'10')) + assert_equal(promote_types('u4', S+'30'), np.dtype(S+'30')) + assert_equal(promote_types('u8', S+'1'), np.dtype(S+'20')) + assert_equal(promote_types('u8', S+'30'), np.dtype(S+'30')) + # Promote with object: + assert_equal(promote_types('O', S+'30'), np.dtype('O')) + + @pytest.mark.parametrize(["dtype1", "dtype2"], + [[np.dtype("V6"), np.dtype("V10")], # mismatch shape + # Mismatching names: + [np.dtype([("name1", "i8")]), np.dtype([("name2", "i8")])], + ]) + def test_invalid_void_promotion(self, dtype1, dtype2): + with pytest.raises(TypeError): + np.promote_types(dtype1, dtype2) + + @pytest.mark.parametrize(["dtype1", "dtype2"], + [[np.dtype("V10"), np.dtype("V10")], + [np.dtype([("name1", "i8")]), + np.dtype([("name1", np.dtype("i8").newbyteorder())])], + [np.dtype("i8,i8"), np.dtype("i8,>i8")], + [np.dtype("i8,i8"), np.dtype("i4,i4")], + ]) + def test_valid_void_promotion(self, dtype1, dtype2): + assert np.promote_types(dtype1, dtype2) == dtype1 + + @pytest.mark.parametrize("dtype", + list(np.typecodes["All"]) + + ["i,i", "10i", "S3", "S100", "U3", "U100", rational]) + def test_promote_identical_types_metadata(self, dtype): + # The same type passed in twice to promote types always + # preserves metadata + metadata = {1: 1} + dtype = np.dtype(dtype, metadata=metadata) + + res = np.promote_types(dtype, dtype) + assert res.metadata == dtype.metadata + + # byte-swapping preserves and makes the dtype native: + dtype = dtype.newbyteorder() + if dtype.isnative: + # The type does not have byte swapping + return + + res = np.promote_types(dtype, dtype) + + # Metadata is (currently) generally lost on byte-swapping (except for + # unicode. + if dtype.char != "U": + assert res.metadata is None + else: + assert res.metadata == metadata + assert res.isnative + + @pytest.mark.slow + @pytest.mark.filterwarnings('ignore:Promotion of numbers:FutureWarning') + @pytest.mark.parametrize(["dtype1", "dtype2"], + itertools.product( + list(np.typecodes["All"]) + + ["i,i", "S3", "S100", "U3", "U100", rational], + repeat=2)) + def test_promote_types_metadata(self, dtype1, dtype2): + """Metadata handling in promotion does not appear formalized + right now in NumPy. This test should thus be considered to + document behaviour, rather than test the correct definition of it. + + This test is very ugly, it was useful for rewriting part of the + promotion, but probably should eventually be replaced/deleted + (i.e. when metadata handling in promotion is better defined). + """ + metadata1 = {1: 1} + metadata2 = {2: 2} + dtype1 = np.dtype(dtype1, metadata=metadata1) + dtype2 = np.dtype(dtype2, metadata=metadata2) + + try: + res = np.promote_types(dtype1, dtype2) + except TypeError: + # Promotion failed, this test only checks metadata + return + + if res.char not in "USV" or res.names is not None or res.shape != (): + # All except string dtypes (and unstructured void) lose metadata + # on promotion (unless both dtypes are identical). + # At some point structured ones did not, but were restrictive. + assert res.metadata is None + elif res == dtype1: + # If one result is the result, it is usually returned unchanged: + assert res is dtype1 + elif res == dtype2: + # dtype1 may have been cast to the same type/kind as dtype2. + # If the resulting dtype is identical we currently pick the cast + # version of dtype1, which lost the metadata: + if np.promote_types(dtype1, dtype2.kind) == dtype2: + res.metadata is None + else: + res.metadata == metadata2 + else: + assert res.metadata is None + + # Try again for byteswapped version + dtype1 = dtype1.newbyteorder() + assert dtype1.metadata == metadata1 + res_bs = np.promote_types(dtype1, dtype2) + assert res_bs == res + assert res_bs.metadata == res.metadata + + def test_can_cast(self): + assert_(np.can_cast(np.int32, np.int64)) + assert_(np.can_cast(np.float64, complex)) + assert_(not np.can_cast(complex, float)) + + assert_(np.can_cast('i8', 'f8')) + assert_(not np.can_cast('i8', 'f4')) + assert_(np.can_cast('i4', 'S11')) + + assert_(np.can_cast('i8', 'i8', 'no')) + assert_(not np.can_cast('i8', 'no')) + + assert_(np.can_cast('i8', 'equiv')) + assert_(not np.can_cast('i8', 'equiv')) + + assert_(np.can_cast('i8', 'safe')) + assert_(not np.can_cast('i4', 'safe')) + + assert_(np.can_cast('i4', 'same_kind')) + assert_(not np.can_cast('u4', 'same_kind')) + + assert_(np.can_cast('u4', 'unsafe')) + + assert_(np.can_cast('bool', 'S5')) + assert_(not np.can_cast('bool', 'S4')) + + assert_(np.can_cast('b', 'S4')) + assert_(not np.can_cast('b', 'S3')) + + assert_(np.can_cast('u1', 'S3')) + assert_(not np.can_cast('u1', 'S2')) + assert_(np.can_cast('u2', 'S5')) + assert_(not np.can_cast('u2', 'S4')) + assert_(np.can_cast('u4', 'S10')) + assert_(not np.can_cast('u4', 'S9')) + assert_(np.can_cast('u8', 'S20')) + assert_(not np.can_cast('u8', 'S19')) + + assert_(np.can_cast('i1', 'S4')) + assert_(not np.can_cast('i1', 'S3')) + assert_(np.can_cast('i2', 'S6')) + assert_(not np.can_cast('i2', 'S5')) + assert_(np.can_cast('i4', 'S11')) + assert_(not np.can_cast('i4', 'S10')) + assert_(np.can_cast('i8', 'S21')) + assert_(not np.can_cast('i8', 'S20')) + + assert_(np.can_cast('bool', 'S5')) + assert_(not np.can_cast('bool', 'S4')) + + assert_(np.can_cast('b', 'U4')) + assert_(not np.can_cast('b', 'U3')) + + assert_(np.can_cast('u1', 'U3')) + assert_(not np.can_cast('u1', 'U2')) + assert_(np.can_cast('u2', 'U5')) + assert_(not np.can_cast('u2', 'U4')) + assert_(np.can_cast('u4', 'U10')) + assert_(not np.can_cast('u4', 'U9')) + assert_(np.can_cast('u8', 'U20')) + assert_(not np.can_cast('u8', 'U19')) + + assert_(np.can_cast('i1', 'U4')) + assert_(not np.can_cast('i1', 'U3')) + assert_(np.can_cast('i2', 'U6')) + assert_(not np.can_cast('i2', 'U5')) + assert_(np.can_cast('i4', 'U11')) + assert_(not np.can_cast('i4', 'U10')) + assert_(np.can_cast('i8', 'U21')) + assert_(not np.can_cast('i8', 'U20')) + + assert_raises(TypeError, np.can_cast, 'i4', None) + assert_raises(TypeError, np.can_cast, None, 'i4') + + # Also test keyword arguments + assert_(np.can_cast(from_=np.int32, to=np.int64)) + + def test_can_cast_simple_to_structured(self): + # Non-structured can only be cast to structured in 'unsafe' mode. + assert_(not np.can_cast('i4', 'i4,i4')) + assert_(not np.can_cast('i4', 'i4,i2')) + assert_(np.can_cast('i4', 'i4,i4', casting='unsafe')) + assert_(np.can_cast('i4', 'i4,i2', casting='unsafe')) + # Even if there is just a single field which is OK. + assert_(not np.can_cast('i2', [('f1', 'i4')])) + assert_(not np.can_cast('i2', [('f1', 'i4')], casting='same_kind')) + assert_(np.can_cast('i2', [('f1', 'i4')], casting='unsafe')) + # It should be the same for recursive structured or subarrays. + assert_(not np.can_cast('i2', [('f1', 'i4,i4')])) + assert_(np.can_cast('i2', [('f1', 'i4,i4')], casting='unsafe')) + assert_(not np.can_cast('i2', [('f1', '(2,3)i4')])) + assert_(np.can_cast('i2', [('f1', '(2,3)i4')], casting='unsafe')) + + def test_can_cast_structured_to_simple(self): + # Need unsafe casting for structured to simple. + assert_(not np.can_cast([('f1', 'i4')], 'i4')) + assert_(np.can_cast([('f1', 'i4')], 'i4', casting='unsafe')) + assert_(np.can_cast([('f1', 'i4')], 'i2', casting='unsafe')) + # Since it is unclear what is being cast, multiple fields to + # single should not work even for unsafe casting. + assert_(not np.can_cast('i4,i4', 'i4', casting='unsafe')) + # But a single field inside a single field is OK. + assert_(not np.can_cast([('f1', [('x', 'i4')])], 'i4')) + assert_(np.can_cast([('f1', [('x', 'i4')])], 'i4', casting='unsafe')) + # And a subarray is fine too - it will just take the first element + # (arguably not very consistently; might also take the first field). + assert_(not np.can_cast([('f0', '(3,)i4')], 'i4')) + assert_(np.can_cast([('f0', '(3,)i4')], 'i4', casting='unsafe')) + # But a structured subarray with multiple fields should fail. + assert_(not np.can_cast([('f0', ('i4,i4'), (2,))], 'i4', + casting='unsafe')) + + def test_can_cast_values(self): + # With NumPy 2 and NEP 50, can_cast errors on Python scalars. We could + # define this as (usually safe) at some point, and already do so + # in `copyto` and ufuncs (but there an error is raised if the integer + # is out of bounds and a warning for out-of-bound floats). + # Raises even for unsafe, previously checked within range (for floats + # that was approximately whether it would overflow to inf). + with pytest.raises(TypeError): + np.can_cast(4, "int8", casting="unsafe") + + with pytest.raises(TypeError): + np.can_cast(4.0, "float64", casting="unsafe") + + with pytest.raises(TypeError): + np.can_cast(4j, "complex128", casting="unsafe") + + + @pytest.mark.parametrize("dtype", + list("?bhilqBHILQefdgFDG") + [rational]) + def test_can_cast_scalars(self, dtype): + # Basic test to ensure that scalars are supported in can-cast + # (does not check behavior exhaustively). + dtype = np.dtype(dtype) + scalar = dtype.type(0) + + assert np.can_cast(scalar, "int64") == np.can_cast(dtype, "int64") + assert np.can_cast(scalar, "float32", casting="unsafe") + + +# Custom exception class to test exception propagation in fromiter +class NIterError(Exception): + pass + + +class TestFromiter: + def makegen(self): + return (x**2 for x in range(24)) + + def test_types(self): + ai32 = np.fromiter(self.makegen(), np.int32) + ai64 = np.fromiter(self.makegen(), np.int64) + af = np.fromiter(self.makegen(), float) + assert_(ai32.dtype == np.dtype(np.int32)) + assert_(ai64.dtype == np.dtype(np.int64)) + assert_(af.dtype == np.dtype(float)) + + def test_lengths(self): + expected = np.array(list(self.makegen())) + a = np.fromiter(self.makegen(), int) + a20 = np.fromiter(self.makegen(), int, 20) + assert_(len(a) == len(expected)) + assert_(len(a20) == 20) + assert_raises(ValueError, np.fromiter, + self.makegen(), int, len(expected) + 10) + + def test_values(self): + expected = np.array(list(self.makegen())) + a = np.fromiter(self.makegen(), int) + a20 = np.fromiter(self.makegen(), int, 20) + assert_(np.all(a == expected, axis=0)) + assert_(np.all(a20 == expected[:20], axis=0)) + + def load_data(self, n, eindex): + # Utility method for the issue 2592 tests. + # Raise an exception at the desired index in the iterator. + for e in range(n): + if e == eindex: + raise NIterError('error at index %s' % eindex) + yield e + + @pytest.mark.parametrize("dtype", [int, object]) + @pytest.mark.parametrize(["count", "error_index"], [(10, 5), (10, 9)]) + def test_2592(self, count, error_index, dtype): + # Test iteration exceptions are correctly raised. The data/generator + # has `count` elements but errors at `error_index` + iterable = self.load_data(count, error_index) + with pytest.raises(NIterError): + np.fromiter(iterable, dtype=dtype, count=count) + + @pytest.mark.parametrize("dtype", ["S", "S0", "V0", "U0"]) + def test_empty_not_structured(self, dtype): + # Note, "S0" could be allowed at some point, so long "S" (without + # any length) is rejected. + with pytest.raises(ValueError, match="Must specify length"): + np.fromiter([], dtype=dtype) + + @pytest.mark.parametrize(["dtype", "data"], + [("d", [1, 2, 3, 4, 5, 6, 7, 8, 9]), + ("O", [1, 2, 3, 4, 5, 6, 7, 8, 9]), + ("i,O", [(1, 2), (5, 4), (2, 3), (9, 8), (6, 7)]), + # subarray dtypes (important because their dimensions end up + # in the result arrays dimension: + ("2i", [(1, 2), (5, 4), (2, 3), (9, 8), (6, 7)]), + (np.dtype(("O", (2, 3))), + [((1, 2, 3), (3, 4, 5)), ((3, 2, 1), (5, 4, 3))])]) + @pytest.mark.parametrize("length_hint", [0, 1]) + def test_growth_and_complicated_dtypes(self, dtype, data, length_hint): + dtype = np.dtype(dtype) + + data = data * 100 # make sure we realloc a bit + + class MyIter: + # Class/example from gh-15789 + def __length_hint__(self): + # only required to be an estimate, this is legal + return length_hint # 0 or 1 + + def __iter__(self): + return iter(data) + + res = np.fromiter(MyIter(), dtype=dtype) + expected = np.array(data, dtype=dtype) + + assert_array_equal(res, expected) + + def test_empty_result(self): + class MyIter: + def __length_hint__(self): + return 10 + + def __iter__(self): + return iter([]) # actual iterator is empty. + + res = np.fromiter(MyIter(), dtype="d") + assert res.shape == (0,) + assert res.dtype == "d" + + def test_too_few_items(self): + msg = "iterator too short: Expected 10 but iterator had only 3 items." + with pytest.raises(ValueError, match=msg): + np.fromiter([1, 2, 3], count=10, dtype=int) + + def test_failed_itemsetting(self): + with pytest.raises(TypeError): + np.fromiter([1, None, 3], dtype=int) + + # The following manages to hit somewhat trickier code paths: + iterable = ((2, 3, 4) for i in range(5)) + with pytest.raises(ValueError): + np.fromiter(iterable, dtype=np.dtype((int, 2))) + +class TestNonzero: + def test_nonzero_trivial(self): + assert_equal(np.count_nonzero(np.array([])), 0) + assert_equal(np.count_nonzero(np.array([], dtype='?')), 0) + assert_equal(np.nonzero(np.array([])), ([],)) + + assert_equal(np.count_nonzero(np.array([0])), 0) + assert_equal(np.count_nonzero(np.array([0], dtype='?')), 0) + assert_equal(np.nonzero(np.array([0])), ([],)) + + assert_equal(np.count_nonzero(np.array([1])), 1) + assert_equal(np.count_nonzero(np.array([1], dtype='?')), 1) + assert_equal(np.nonzero(np.array([1])), ([0],)) + + def test_nonzero_zerodim(self): + err_msg = "Calling nonzero on 0d arrays is not allowed" + with assert_raises_regex(ValueError, err_msg): + np.nonzero(np.array(0)) + with assert_raises_regex(ValueError, err_msg): + np.array(1).nonzero() + + def test_nonzero_onedim(self): + x = np.array([1, 0, 2, -1, 0, 0, 8]) + assert_equal(np.count_nonzero(x), 4) + assert_equal(np.count_nonzero(x), 4) + assert_equal(np.nonzero(x), ([0, 2, 3, 6],)) + + # x = np.array([(1, 2), (0, 0), (1, 1), (-1, 3), (0, 7)], + # dtype=[('a', 'i4'), ('b', 'i2')]) + x = np.array([(1, 2, -5, -3), (0, 0, 2, 7), (1, 1, 0, 1), (-1, 3, 1, 0), (0, 7, 0, 4)], + dtype=[('a', 'i4'), ('b', 'i2'), ('c', 'i1'), ('d', 'i8')]) + assert_equal(np.count_nonzero(x['a']), 3) + assert_equal(np.count_nonzero(x['b']), 4) + assert_equal(np.count_nonzero(x['c']), 3) + assert_equal(np.count_nonzero(x['d']), 4) + assert_equal(np.nonzero(x['a']), ([0, 2, 3],)) + assert_equal(np.nonzero(x['b']), ([0, 2, 3, 4],)) + + def test_nonzero_twodim(self): + x = np.array([[0, 1, 0], [2, 0, 3]]) + assert_equal(np.count_nonzero(x.astype('i1')), 3) + assert_equal(np.count_nonzero(x.astype('i2')), 3) + assert_equal(np.count_nonzero(x.astype('i4')), 3) + assert_equal(np.count_nonzero(x.astype('i8')), 3) + assert_equal(np.nonzero(x), ([0, 1, 1], [1, 0, 2])) + + x = np.eye(3) + assert_equal(np.count_nonzero(x.astype('i1')), 3) + assert_equal(np.count_nonzero(x.astype('i2')), 3) + assert_equal(np.count_nonzero(x.astype('i4')), 3) + assert_equal(np.count_nonzero(x.astype('i8')), 3) + assert_equal(np.nonzero(x), ([0, 1, 2], [0, 1, 2])) + + x = np.array([[(0, 1), (0, 0), (1, 11)], + [(1, 1), (1, 0), (0, 0)], + [(0, 0), (1, 5), (0, 1)]], dtype=[('a', 'f4'), ('b', 'u1')]) + assert_equal(np.count_nonzero(x['a']), 4) + assert_equal(np.count_nonzero(x['b']), 5) + assert_equal(np.nonzero(x['a']), ([0, 1, 1, 2], [2, 0, 1, 1])) + assert_equal(np.nonzero(x['b']), ([0, 0, 1, 2, 2], [0, 2, 0, 1, 2])) + + assert_(not x['a'].T.flags.aligned) + assert_equal(np.count_nonzero(x['a'].T), 4) + assert_equal(np.count_nonzero(x['b'].T), 5) + assert_equal(np.nonzero(x['a'].T), ([0, 1, 1, 2], [1, 1, 2, 0])) + assert_equal(np.nonzero(x['b'].T), ([0, 0, 1, 2, 2], [0, 1, 2, 0, 2])) + + def test_sparse(self): + # test special sparse condition boolean code path + for i in range(20): + c = np.zeros(200, dtype=bool) + c[i::20] = True + assert_equal(np.nonzero(c)[0], np.arange(i, 200 + i, 20)) + + c = np.zeros(400, dtype=bool) + c[10 + i:20 + i] = True + c[20 + i*2] = True + assert_equal(np.nonzero(c)[0], + np.concatenate((np.arange(10 + i, 20 + i), [20 + i*2]))) + + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_nonzero_float_dtypes(self, dtype): + rng = np.random.default_rng(seed=10) + x = ((2**33)*rng.normal(size=100)).astype(dtype) + x[rng.choice(50, size=100)] = 0 + idxs = np.nonzero(x)[0] + assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True) + + @pytest.mark.parametrize('dtype', [bool, np.int8, np.int16, np.int32, np.int64, + np.uint8, np.uint16, np.uint32, np.uint64]) + def test_nonzero_integer_dtypes(self, dtype): + rng = np.random.default_rng(seed=10) + x = rng.integers(0, 255, size=100).astype(dtype) + x[rng.choice(50, size=100)] = 0 + idxs = np.nonzero(x)[0] + assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True) + + def test_return_type(self): + class C(np.ndarray): + pass + + for view in (C, np.ndarray): + for nd in range(1, 4): + shape = tuple(range(2, 2+nd)) + x = np.arange(np.prod(shape)).reshape(shape).view(view) + for nzx in (np.nonzero(x), x.nonzero()): + for nzx_i in nzx: + assert_(type(nzx_i) is np.ndarray) + assert_(nzx_i.flags.writeable) + + def test_count_nonzero_axis(self): + # Basic check of functionality + m = np.array([[0, 1, 7, 0, 0], [3, 0, 0, 2, 19]]) + + expected = np.array([1, 1, 1, 1, 1]) + assert_equal(np.count_nonzero(m, axis=0), expected) + + expected = np.array([2, 3]) + assert_equal(np.count_nonzero(m, axis=1), expected) + + assert_raises(ValueError, np.count_nonzero, m, axis=(1, 1)) + assert_raises(TypeError, np.count_nonzero, m, axis='foo') + assert_raises(AxisError, np.count_nonzero, m, axis=3) + assert_raises(TypeError, np.count_nonzero, + m, axis=np.array([[1], [2]])) + + def test_count_nonzero_axis_all_dtypes(self): + # More thorough test that the axis argument is respected + # for all dtypes and responds correctly when presented with + # either integer or tuple arguments for axis + msg = "Mismatch for dtype: %s" + + def assert_equal_w_dt(a, b, err_msg): + assert_equal(a.dtype, b.dtype, err_msg=err_msg) + assert_equal(a, b, err_msg=err_msg) + + for dt in np.typecodes['All']: + err_msg = msg % (np.dtype(dt).name,) + + if dt != 'V': + if dt != 'M': + m = np.zeros((3, 3), dtype=dt) + n = np.ones(1, dtype=dt) + + m[0, 0] = n[0] + m[1, 0] = n[0] + + else: # np.zeros doesn't work for np.datetime64 + m = np.array(['1970-01-01'] * 9) + m = m.reshape((3, 3)) + + m[0, 0] = '1970-01-12' + m[1, 0] = '1970-01-12' + m = m.astype(dt) + + expected = np.array([2, 0, 0], dtype=np.intp) + assert_equal_w_dt(np.count_nonzero(m, axis=0), + expected, err_msg=err_msg) + + expected = np.array([1, 1, 0], dtype=np.intp) + assert_equal_w_dt(np.count_nonzero(m, axis=1), + expected, err_msg=err_msg) + + expected = np.array(2) + assert_equal(np.count_nonzero(m, axis=(0, 1)), + expected, err_msg=err_msg) + assert_equal(np.count_nonzero(m, axis=None), + expected, err_msg=err_msg) + assert_equal(np.count_nonzero(m), + expected, err_msg=err_msg) + + if dt == 'V': + # There are no 'nonzero' objects for np.void, so the testing + # setup is slightly different for this dtype + m = np.array([np.void(1)] * 6).reshape((2, 3)) + + expected = np.array([0, 0, 0], dtype=np.intp) + assert_equal_w_dt(np.count_nonzero(m, axis=0), + expected, err_msg=err_msg) + + expected = np.array([0, 0], dtype=np.intp) + assert_equal_w_dt(np.count_nonzero(m, axis=1), + expected, err_msg=err_msg) + + expected = np.array(0) + assert_equal(np.count_nonzero(m, axis=(0, 1)), + expected, err_msg=err_msg) + assert_equal(np.count_nonzero(m, axis=None), + expected, err_msg=err_msg) + assert_equal(np.count_nonzero(m), + expected, err_msg=err_msg) + + def test_count_nonzero_axis_consistent(self): + # Check that the axis behaviour for valid axes in + # non-special cases is consistent (and therefore + # correct) by checking it against an integer array + # that is then casted to the generic object dtype + from itertools import combinations, permutations + + axis = (0, 1, 2, 3) + size = (5, 5, 5, 5) + msg = "Mismatch for axis: %s" + + rng = np.random.RandomState(1234) + m = rng.randint(-100, 100, size=size) + n = m.astype(object) + + for length in range(len(axis)): + for combo in combinations(axis, length): + for perm in permutations(combo): + assert_equal( + np.count_nonzero(m, axis=perm), + np.count_nonzero(n, axis=perm), + err_msg=msg % (perm,)) + + def test_countnonzero_axis_empty(self): + a = np.array([[0, 0, 1], [1, 0, 1]]) + assert_equal(np.count_nonzero(a, axis=()), a.astype(bool)) + + def test_countnonzero_keepdims(self): + a = np.array([[0, 0, 1, 0], + [0, 3, 5, 0], + [7, 9, 2, 0]]) + assert_equal(np.count_nonzero(a, axis=0, keepdims=True), + [[1, 2, 3, 0]]) + assert_equal(np.count_nonzero(a, axis=1, keepdims=True), + [[1], [2], [3]]) + assert_equal(np.count_nonzero(a, keepdims=True), + [[6]]) + + def test_array_method(self): + # Tests that the array method + # call to nonzero works + m = np.array([[1, 0, 0], [4, 0, 6]]) + tgt = [[0, 1, 1], [0, 0, 2]] + + assert_equal(m.nonzero(), tgt) + + def test_nonzero_invalid_object(self): + # gh-9295 + a = np.array([np.array([1, 2]), 3], dtype=object) + assert_raises(ValueError, np.nonzero, a) + + class BoolErrors: + def __bool__(self): + raise ValueError("Not allowed") + + assert_raises(ValueError, np.nonzero, np.array([BoolErrors()])) + + def test_nonzero_sideeffect_safety(self): + # gh-13631 + class FalseThenTrue: + _val = False + def __bool__(self): + try: + return self._val + finally: + self._val = True + + class TrueThenFalse: + _val = True + def __bool__(self): + try: + return self._val + finally: + self._val = False + + # result grows on the second pass + a = np.array([True, FalseThenTrue()]) + assert_raises(RuntimeError, np.nonzero, a) + + a = np.array([[True], [FalseThenTrue()]]) + assert_raises(RuntimeError, np.nonzero, a) + + # result shrinks on the second pass + a = np.array([False, TrueThenFalse()]) + assert_raises(RuntimeError, np.nonzero, a) + + a = np.array([[False], [TrueThenFalse()]]) + assert_raises(RuntimeError, np.nonzero, a) + + def test_nonzero_sideffects_structured_void(self): + # Checks that structured void does not mutate alignment flag of + # original array. + arr = np.zeros(5, dtype="i1,i8,i8") # `ones` may short-circuit + assert arr.flags.aligned # structs are considered "aligned" + assert not arr["f2"].flags.aligned + # make sure that nonzero/count_nonzero do not flip the flag: + np.nonzero(arr) + assert arr.flags.aligned + np.count_nonzero(arr) + assert arr.flags.aligned + + def test_nonzero_exception_safe(self): + # gh-13930 + + class ThrowsAfter: + def __init__(self, iters): + self.iters_left = iters + + def __bool__(self): + if self.iters_left == 0: + raise ValueError("called `iters` times") + + self.iters_left -= 1 + return True + + """ + Test that a ValueError is raised instead of a SystemError + + If the __bool__ function is called after the error state is set, + Python (cpython) will raise a SystemError. + """ + + # assert that an exception in first pass is handled correctly + a = np.array([ThrowsAfter(5)]*10) + assert_raises(ValueError, np.nonzero, a) + + # raise exception in second pass for 1-dimensional loop + a = np.array([ThrowsAfter(15)]*10) + assert_raises(ValueError, np.nonzero, a) + + # raise exception in second pass for n-dimensional loop + a = np.array([[ThrowsAfter(15)]]*10) + assert_raises(ValueError, np.nonzero, a) + + @pytest.mark.skipif(IS_WASM, reason="wasm doesn't have threads") + def test_structured_threadsafety(self): + # Nonzero (and some other functions) should be threadsafe for + # structured datatypes, see gh-15387. This test can behave randomly. + from concurrent.futures import ThreadPoolExecutor + + # Create a deeply nested dtype to make a failure more likely: + dt = np.dtype([("", "f8")]) + dt = np.dtype([("", dt)]) + dt = np.dtype([("", dt)] * 2) + # The array should be large enough to likely run into threading issues + arr = np.random.uniform(size=(5000, 4)).view(dt)[:, 0] + def func(arr): + arr.nonzero() + + tpe = ThreadPoolExecutor(max_workers=8) + futures = [tpe.submit(func, arr) for _ in range(10)] + for f in futures: + f.result() + + assert arr.dtype is dt + + +class TestIndex: + def test_boolean(self): + a = rand(3, 5, 8) + V = rand(5, 8) + g1 = randint(0, 5, size=15) + g2 = randint(0, 8, size=15) + V[g1, g2] = -V[g1, g2] + assert_((np.array([a[0][V > 0], a[1][V > 0], a[2][V > 0]]) == a[:, V > 0]).all()) + + def test_boolean_edgecase(self): + a = np.array([], dtype='int32') + b = np.array([], dtype='bool') + c = a[b] + assert_equal(c, []) + assert_equal(c.dtype, np.dtype('int32')) + + +class TestBinaryRepr: + def test_zero(self): + assert_equal(np.binary_repr(0), '0') + + def test_positive(self): + assert_equal(np.binary_repr(10), '1010') + assert_equal(np.binary_repr(12522), + '11000011101010') + assert_equal(np.binary_repr(10736848), + '101000111101010011010000') + + def test_negative(self): + assert_equal(np.binary_repr(-1), '-1') + assert_equal(np.binary_repr(-10), '-1010') + assert_equal(np.binary_repr(-12522), + '-11000011101010') + assert_equal(np.binary_repr(-10736848), + '-101000111101010011010000') + + def test_sufficient_width(self): + assert_equal(np.binary_repr(0, width=5), '00000') + assert_equal(np.binary_repr(10, width=7), '0001010') + assert_equal(np.binary_repr(-5, width=7), '1111011') + + def test_neg_width_boundaries(self): + # see gh-8670 + + # Ensure that the example in the issue does not + # break before proceeding to a more thorough test. + assert_equal(np.binary_repr(-128, width=8), '10000000') + + for width in range(1, 11): + num = -2**(width - 1) + exp = '1' + (width - 1) * '0' + assert_equal(np.binary_repr(num, width=width), exp) + + def test_large_neg_int64(self): + # See gh-14289. + assert_equal(np.binary_repr(np.int64(-2**62), width=64), + '11' + '0'*62) + + +class TestBaseRepr: + def test_base3(self): + assert_equal(np.base_repr(3**5, 3), '100000') + + def test_positive(self): + assert_equal(np.base_repr(12, 10), '12') + assert_equal(np.base_repr(12, 10, 4), '000012') + assert_equal(np.base_repr(12, 4), '30') + assert_equal(np.base_repr(3731624803700888, 36), '10QR0ROFCEW') + + def test_negative(self): + assert_equal(np.base_repr(-12, 10), '-12') + assert_equal(np.base_repr(-12, 10, 4), '-000012') + assert_equal(np.base_repr(-12, 4), '-30') + + def test_base_range(self): + with assert_raises(ValueError): + np.base_repr(1, 1) + with assert_raises(ValueError): + np.base_repr(1, 37) + + def test_minimal_signed_int(self): + assert_equal(np.base_repr(np.int8(-128)), '-10000000') + + +def _test_array_equal_parametrizations(): + """ + we pre-create arrays as we sometime want to pass the same instance + and sometime not. Passing the same instances may not mean the array are + equal, especially when containing None + """ + # those are 0-d arrays, it used to be a special case + # where (e0 == e0).all() would raise + e0 = np.array(0, dtype="int") + e1 = np.array(1, dtype="float") + # x,y, nan_equal, expected_result + yield (e0, e0.copy(), None, True) + yield (e0, e0.copy(), False, True) + yield (e0, e0.copy(), True, True) + + # + yield (e1, e1.copy(), None, True) + yield (e1, e1.copy(), False, True) + yield (e1, e1.copy(), True, True) + + # Non-nanable – those cannot hold nans + a12 = np.array([1, 2]) + a12b = a12.copy() + a123 = np.array([1, 2, 3]) + a13 = np.array([1, 3]) + a34 = np.array([3, 4]) + + aS1 = np.array(["a"], dtype="S1") + aS1b = aS1.copy() + aS1u4 = np.array([("a", 1)], dtype="S1,u4") + aS1u4b = aS1u4.copy() + + yield (a12, a12b, None, True) + yield (a12, a12, None, True) + yield (a12, a123, None, False) + yield (a12, a34, None, False) + yield (a12, a13, None, False) + yield (aS1, aS1b, None, True) + yield (aS1, aS1, None, True) + + # Non-float dtype - equal_nan should have no effect, + yield (a123, a123, None, True) + yield (a123, a123, False, True) + yield (a123, a123, True, True) + yield (a123, a123.copy(), None, True) + yield (a123, a123.copy(), False, True) + yield (a123, a123.copy(), True, True) + yield (a123.astype("float"), a123.astype("float"), None, True) + yield (a123.astype("float"), a123.astype("float"), False, True) + yield (a123.astype("float"), a123.astype("float"), True, True) + + # these can hold None + b1 = np.array([1, 2, np.nan]) + b2 = np.array([1, np.nan, 2]) + b3 = np.array([1, 2, np.inf]) + b4 = np.array(np.nan) + + # instances are the same + yield (b1, b1, None, False) + yield (b1, b1, False, False) + yield (b1, b1, True, True) + + # equal but not same instance + yield (b1, b1.copy(), None, False) + yield (b1, b1.copy(), False, False) + yield (b1, b1.copy(), True, True) + + # same once stripped of Nan + yield (b1, b2, None, False) + yield (b1, b2, False, False) + yield (b1, b2, True, False) + + # nan's not conflated with inf's + yield (b1, b3, None, False) + yield (b1, b3, False, False) + yield (b1, b3, True, False) + + # all Nan + yield (b4, b4, None, False) + yield (b4, b4, False, False) + yield (b4, b4, True, True) + yield (b4, b4.copy(), None, False) + yield (b4, b4.copy(), False, False) + yield (b4, b4.copy(), True, True) + + t1 = b1.astype("timedelta64") + t2 = b2.astype("timedelta64") + + # Timedeltas are particular + yield (t1, t1, None, False) + yield (t1, t1, False, False) + yield (t1, t1, True, True) + + yield (t1, t1.copy(), None, False) + yield (t1, t1.copy(), False, False) + yield (t1, t1.copy(), True, True) + + yield (t1, t2, None, False) + yield (t1, t2, False, False) + yield (t1, t2, True, False) + + # Multi-dimensional array + md1 = np.array([[0, 1], [np.nan, 1]]) + + yield (md1, md1, None, False) + yield (md1, md1, False, False) + yield (md1, md1, True, True) + yield (md1, md1.copy(), None, False) + yield (md1, md1.copy(), False, False) + yield (md1, md1.copy(), True, True) + # both complexes are nan+nan.j but the same instance + cplx1, cplx2 = [np.array([np.nan + np.nan * 1j])] * 2 + + # only real or img are nan. + cplx3, cplx4 = np.complex64(1, np.nan), np.complex64(np.nan, 1) + + # Complex values + yield (cplx1, cplx2, None, False) + yield (cplx1, cplx2, False, False) + yield (cplx1, cplx2, True, True) + + # Complex values, 1+nan, nan+1j + yield (cplx3, cplx4, None, False) + yield (cplx3, cplx4, False, False) + yield (cplx3, cplx4, True, True) + + +class TestArrayComparisons: + @pytest.mark.parametrize( + "bx,by,equal_nan,expected", _test_array_equal_parametrizations() + ) + def test_array_equal_equal_nan(self, bx, by, equal_nan, expected): + """ + This test array_equal for a few combinations: + + - are the two inputs the same object or not (same object may not + be equal if contains NaNs) + - Whether we should consider or not, NaNs, being equal. + + """ + if equal_nan is None: + res = np.array_equal(bx, by) + else: + res = np.array_equal(bx, by, equal_nan=equal_nan) + assert_(res is expected) + assert_(type(res) is bool) + + def test_array_equal_different_scalar_types(self): + # https://github.com/numpy/numpy/issues/27271 + a = np.array("foo") + b = np.array(1) + assert not np.array_equal(a, b) + assert not np.array_equiv(a, b) + + def test_none_compares_elementwise(self): + a = np.array([None, 1, None], dtype=object) + assert_equal(a == None, [True, False, True]) # noqa: E711 + assert_equal(a != None, [False, True, False]) # noqa: E711 + + a = np.ones(3) + assert_equal(a == None, [False, False, False]) # noqa: E711 + assert_equal(a != None, [True, True, True]) # noqa: E711 + + def test_array_equiv(self): + res = np.array_equiv(np.array([1, 2]), np.array([1, 2])) + assert_(res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 2]), np.array([1, 2, 3])) + assert_(not res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 2]), np.array([3, 4])) + assert_(not res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 2]), np.array([1, 3])) + assert_(not res) + assert_(type(res) is bool) + + res = np.array_equiv(np.array([1, 1]), np.array([1])) + assert_(res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 1]), np.array([[1], [1]])) + assert_(res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 2]), np.array([2])) + assert_(not res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 2]), np.array([[1], [2]])) + assert_(not res) + assert_(type(res) is bool) + res = np.array_equiv(np.array([1, 2]), np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + assert_(not res) + assert_(type(res) is bool) + + @pytest.mark.parametrize("dtype", ["V0", "V3", "V10"]) + def test_compare_unstructured_voids(self, dtype): + zeros = np.zeros(3, dtype=dtype) + + assert_array_equal(zeros, zeros) + assert not (zeros != zeros).any() + + if dtype == "V0": + # Can't test != of actually different data + return + + nonzeros = np.array([b"1", b"2", b"3"], dtype=dtype) + + assert not (zeros == nonzeros).any() + assert (zeros != nonzeros).all() + + +def assert_array_strict_equal(x, y): + assert_array_equal(x, y) + # Check flags, 32 bit arches typically don't provide 16 byte alignment + if ((x.dtype.alignment <= 8 or + np.intp().dtype.itemsize != 4) and + sys.platform != 'win32'): + assert_(x.flags == y.flags) + else: + assert_(x.flags.owndata == y.flags.owndata) + assert_(x.flags.writeable == y.flags.writeable) + assert_(x.flags.c_contiguous == y.flags.c_contiguous) + assert_(x.flags.f_contiguous == y.flags.f_contiguous) + assert_(x.flags.writebackifcopy == y.flags.writebackifcopy) + # check endianness + assert_(x.dtype.isnative == y.dtype.isnative) + + +class TestClip: + def setup_method(self): + self.nr = 5 + self.nc = 3 + + def fastclip(self, a, m, M, out=None, **kwargs): + return a.clip(m, M, out=out, **kwargs) + + def clip(self, a, m, M, out=None): + # use a.choose to verify fastclip result + selector = np.less(a, m) + 2*np.greater(a, M) + return selector.choose((a, m, M), out=out) + + # Handy functions + def _generate_data(self, n, m): + return randn(n, m) + + def _generate_data_complex(self, n, m): + return randn(n, m) + 1.j * rand(n, m) + + def _generate_flt_data(self, n, m): + return (randn(n, m)).astype(np.float32) + + def _neg_byteorder(self, a): + a = np.asarray(a) + if sys.byteorder == 'little': + a = a.astype(a.dtype.newbyteorder('>')) + else: + a = a.astype(a.dtype.newbyteorder('<')) + return a + + def _generate_non_native_data(self, n, m): + data = randn(n, m) + data = self._neg_byteorder(data) + assert_(not data.dtype.isnative) + return data + + def _generate_int_data(self, n, m): + return (10 * rand(n, m)).astype(np.int64) + + def _generate_int32_data(self, n, m): + return (10 * rand(n, m)).astype(np.int32) + + # Now the real test cases + + @pytest.mark.parametrize("dtype", '?bhilqpBHILQPefdgFDGO') + def test_ones_pathological(self, dtype): + # for preservation of behavior described in + # gh-12519; amin > amax behavior may still change + # in the future + arr = np.ones(10, dtype=dtype) + expected = np.zeros(10, dtype=dtype) + actual = np.clip(arr, 1, 0) + if dtype == 'O': + assert actual.tolist() == expected.tolist() + else: + assert_equal(actual, expected) + + def test_simple_double(self): + # Test native double input with scalar min/max. + a = self._generate_data(self.nr, self.nc) + m = 0.1 + M = 0.6 + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_simple_int(self): + # Test native int input with scalar min/max. + a = self._generate_int_data(self.nr, self.nc) + a = a.astype(int) + m = -2 + M = 4 + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_array_double(self): + # Test native double input with array min/max. + a = self._generate_data(self.nr, self.nc) + m = np.zeros(a.shape) + M = m + 0.5 + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_simple_nonnative(self): + # Test non native double input with scalar min/max. + # Test native double input with non native double scalar min/max. + a = self._generate_non_native_data(self.nr, self.nc) + m = -0.5 + M = 0.6 + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_equal(ac, act) + + # Test native double input with non native double scalar min/max. + a = self._generate_data(self.nr, self.nc) + m = -0.5 + M = self._neg_byteorder(0.6) + assert_(not M.dtype.isnative) + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_equal(ac, act) + + def test_simple_complex(self): + # Test native complex input with native double scalar min/max. + # Test native input with complex double scalar min/max. + a = 3 * self._generate_data_complex(self.nr, self.nc) + m = -0.5 + M = 1. + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + # Test native input with complex double scalar min/max. + a = 3 * self._generate_data(self.nr, self.nc) + m = -0.5 + 1.j + M = 1. + 2.j + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_clip_complex(self): + # Address Issue gh-5354 for clipping complex arrays + # Test native complex input without explicit min/max + # ie, either min=None or max=None + a = np.ones(10, dtype=complex) + m = a.min() + M = a.max() + am = self.fastclip(a, m, None) + aM = self.fastclip(a, None, M) + assert_array_strict_equal(am, a) + assert_array_strict_equal(aM, a) + + def test_clip_non_contig(self): + # Test clip for non contiguous native input and native scalar min/max. + a = self._generate_data(self.nr * 2, self.nc * 3) + a = a[::2, ::3] + assert_(not a.flags['F_CONTIGUOUS']) + assert_(not a.flags['C_CONTIGUOUS']) + ac = self.fastclip(a, -1.6, 1.7) + act = self.clip(a, -1.6, 1.7) + assert_array_strict_equal(ac, act) + + def test_simple_out(self): + # Test native double input with scalar min/max. + a = self._generate_data(self.nr, self.nc) + m = -0.5 + M = 0.6 + ac = np.zeros(a.shape) + act = np.zeros(a.shape) + self.fastclip(a, m, M, ac) + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + @pytest.mark.parametrize("casting", [None, "unsafe"]) + def test_simple_int32_inout(self, casting): + # Test native int32 input with double min/max and int32 out. + a = self._generate_int32_data(self.nr, self.nc) + m = np.float64(0) + M = np.float64(2) + ac = np.zeros(a.shape, dtype=np.int32) + act = ac.copy() + if casting is None: + with pytest.raises(TypeError): + self.fastclip(a, m, M, ac, casting=casting) + else: + # explicitly passing "unsafe" will silence warning + self.fastclip(a, m, M, ac, casting=casting) + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_simple_int64_out(self): + # Test native int32 input with int32 scalar min/max and int64 out. + a = self._generate_int32_data(self.nr, self.nc) + m = np.int32(-1) + M = np.int32(1) + ac = np.zeros(a.shape, dtype=np.int64) + act = ac.copy() + self.fastclip(a, m, M, ac) + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_simple_int64_inout(self): + # Test native int32 input with double array min/max and int32 out. + a = self._generate_int32_data(self.nr, self.nc) + m = np.zeros(a.shape, np.float64) + M = np.float64(1) + ac = np.zeros(a.shape, dtype=np.int32) + act = ac.copy() + self.fastclip(a, m, M, out=ac, casting="unsafe") + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_simple_int32_out(self): + # Test native double input with scalar min/max and int out. + a = self._generate_data(self.nr, self.nc) + m = -1.0 + M = 2.0 + ac = np.zeros(a.shape, dtype=np.int32) + act = ac.copy() + self.fastclip(a, m, M, out=ac, casting="unsafe") + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_simple_inplace_01(self): + # Test native double input with array min/max in-place. + a = self._generate_data(self.nr, self.nc) + ac = a.copy() + m = np.zeros(a.shape) + M = 1.0 + self.fastclip(a, m, M, a) + self.clip(a, m, M, ac) + assert_array_strict_equal(a, ac) + + def test_simple_inplace_02(self): + # Test native double input with scalar min/max in-place. + a = self._generate_data(self.nr, self.nc) + ac = a.copy() + m = -0.5 + M = 0.6 + self.fastclip(a, m, M, a) + self.clip(ac, m, M, ac) + assert_array_strict_equal(a, ac) + + def test_noncontig_inplace(self): + # Test non contiguous double input with double scalar min/max in-place. + a = self._generate_data(self.nr * 2, self.nc * 3) + a = a[::2, ::3] + assert_(not a.flags['F_CONTIGUOUS']) + assert_(not a.flags['C_CONTIGUOUS']) + ac = a.copy() + m = -0.5 + M = 0.6 + self.fastclip(a, m, M, a) + self.clip(ac, m, M, ac) + assert_array_equal(a, ac) + + def test_type_cast_01(self): + # Test native double input with scalar min/max. + a = self._generate_data(self.nr, self.nc) + m = -0.5 + M = 0.6 + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_02(self): + # Test native int32 input with int32 scalar min/max. + a = self._generate_int_data(self.nr, self.nc) + a = a.astype(np.int32) + m = -2 + M = 4 + ac = self.fastclip(a, m, M) + act = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_03(self): + # Test native int32 input with float64 scalar min/max. + a = self._generate_int32_data(self.nr, self.nc) + m = -2 + M = 4 + ac = self.fastclip(a, np.float64(m), np.float64(M)) + act = self.clip(a, np.float64(m), np.float64(M)) + assert_array_strict_equal(ac, act) + + def test_type_cast_04(self): + # Test native int32 input with float32 scalar min/max. + a = self._generate_int32_data(self.nr, self.nc) + m = np.float32(-2) + M = np.float32(4) + act = self.fastclip(a, m, M) + ac = self.clip(a, m, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_05(self): + # Test native int32 with double arrays min/max. + a = self._generate_int_data(self.nr, self.nc) + m = -0.5 + M = 1. + ac = self.fastclip(a, m * np.zeros(a.shape), M) + act = self.clip(a, m * np.zeros(a.shape), M) + assert_array_strict_equal(ac, act) + + def test_type_cast_06(self): + # Test native with NON native scalar min/max. + a = self._generate_data(self.nr, self.nc) + m = 0.5 + m_s = self._neg_byteorder(m) + M = 1. + act = self.clip(a, m_s, M) + ac = self.fastclip(a, m_s, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_07(self): + # Test NON native with native array min/max. + a = self._generate_data(self.nr, self.nc) + m = -0.5 * np.ones(a.shape) + M = 1. + a_s = self._neg_byteorder(a) + assert_(not a_s.dtype.isnative) + act = a_s.clip(m, M) + ac = self.fastclip(a_s, m, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_08(self): + # Test NON native with native scalar min/max. + a = self._generate_data(self.nr, self.nc) + m = -0.5 + M = 1. + a_s = self._neg_byteorder(a) + assert_(not a_s.dtype.isnative) + ac = self.fastclip(a_s, m, M) + act = a_s.clip(m, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_09(self): + # Test native with NON native array min/max. + a = self._generate_data(self.nr, self.nc) + m = -0.5 * np.ones(a.shape) + M = 1. + m_s = self._neg_byteorder(m) + assert_(not m_s.dtype.isnative) + ac = self.fastclip(a, m_s, M) + act = self.clip(a, m_s, M) + assert_array_strict_equal(ac, act) + + def test_type_cast_10(self): + # Test native int32 with float min/max and float out for output argument. + a = self._generate_int_data(self.nr, self.nc) + b = np.zeros(a.shape, dtype=np.float32) + m = np.float32(-0.5) + M = np.float32(1) + act = self.clip(a, m, M, out=b) + ac = self.fastclip(a, m, M, out=b) + assert_array_strict_equal(ac, act) + + def test_type_cast_11(self): + # Test non native with native scalar, min/max, out non native + a = self._generate_non_native_data(self.nr, self.nc) + b = a.copy() + b = b.astype(b.dtype.newbyteorder('>')) + bt = b.copy() + m = -0.5 + M = 1. + self.fastclip(a, m, M, out=b) + self.clip(a, m, M, out=bt) + assert_array_strict_equal(b, bt) + + def test_type_cast_12(self): + # Test native int32 input and min/max and float out + a = self._generate_int_data(self.nr, self.nc) + b = np.zeros(a.shape, dtype=np.float32) + m = np.int32(0) + M = np.int32(1) + act = self.clip(a, m, M, out=b) + ac = self.fastclip(a, m, M, out=b) + assert_array_strict_equal(ac, act) + + def test_clip_with_out_simple(self): + # Test native double input with scalar min/max + a = self._generate_data(self.nr, self.nc) + m = -0.5 + M = 0.6 + ac = np.zeros(a.shape) + act = np.zeros(a.shape) + self.fastclip(a, m, M, ac) + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_clip_with_out_simple2(self): + # Test native int32 input with double min/max and int32 out + a = self._generate_int32_data(self.nr, self.nc) + m = np.float64(0) + M = np.float64(2) + ac = np.zeros(a.shape, dtype=np.int32) + act = ac.copy() + self.fastclip(a, m, M, out=ac, casting="unsafe") + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_clip_with_out_simple_int32(self): + # Test native int32 input with int32 scalar min/max and int64 out + a = self._generate_int32_data(self.nr, self.nc) + m = np.int32(-1) + M = np.int32(1) + ac = np.zeros(a.shape, dtype=np.int64) + act = ac.copy() + self.fastclip(a, m, M, ac) + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_clip_with_out_array_int32(self): + # Test native int32 input with double array min/max and int32 out + a = self._generate_int32_data(self.nr, self.nc) + m = np.zeros(a.shape, np.float64) + M = np.float64(1) + ac = np.zeros(a.shape, dtype=np.int32) + act = ac.copy() + self.fastclip(a, m, M, out=ac, casting="unsafe") + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_clip_with_out_array_outint32(self): + # Test native double input with scalar min/max and int out + a = self._generate_data(self.nr, self.nc) + m = -1.0 + M = 2.0 + ac = np.zeros(a.shape, dtype=np.int32) + act = ac.copy() + self.fastclip(a, m, M, out=ac, casting="unsafe") + self.clip(a, m, M, act) + assert_array_strict_equal(ac, act) + + def test_clip_with_out_transposed(self): + # Test that the out argument works when transposed + a = np.arange(16).reshape(4, 4) + out = np.empty_like(a).T + a.clip(4, 10, out=out) + expected = self.clip(a, 4, 10) + assert_array_equal(out, expected) + + def test_clip_with_out_memory_overlap(self): + # Test that the out argument works when it has memory overlap + a = np.arange(16).reshape(4, 4) + ac = a.copy() + a[:-1].clip(4, 10, out=a[1:]) + expected = self.clip(ac[:-1], 4, 10) + assert_array_equal(a[1:], expected) + + def test_clip_inplace_array(self): + # Test native double input with array min/max + a = self._generate_data(self.nr, self.nc) + ac = a.copy() + m = np.zeros(a.shape) + M = 1.0 + self.fastclip(a, m, M, a) + self.clip(a, m, M, ac) + assert_array_strict_equal(a, ac) + + def test_clip_inplace_simple(self): + # Test native double input with scalar min/max + a = self._generate_data(self.nr, self.nc) + ac = a.copy() + m = -0.5 + M = 0.6 + self.fastclip(a, m, M, a) + self.clip(a, m, M, ac) + assert_array_strict_equal(a, ac) + + def test_clip_func_takes_out(self): + # Ensure that the clip() function takes an out=argument. + a = self._generate_data(self.nr, self.nc) + ac = a.copy() + m = -0.5 + M = 0.6 + a2 = np.clip(a, m, M, out=a) + self.clip(a, m, M, ac) + assert_array_strict_equal(a2, ac) + assert_(a2 is a) + + def test_clip_nan(self): + d = np.arange(7.) + assert_equal(d.clip(min=np.nan), np.nan) + assert_equal(d.clip(max=np.nan), np.nan) + assert_equal(d.clip(min=np.nan, max=np.nan), np.nan) + assert_equal(d.clip(min=-2, max=np.nan), np.nan) + assert_equal(d.clip(min=np.nan, max=10), np.nan) + + def test_object_clip(self): + a = np.arange(10, dtype=object) + actual = np.clip(a, 1, 5) + expected = np.array([1, 1, 2, 3, 4, 5, 5, 5, 5, 5]) + assert actual.tolist() == expected.tolist() + + def test_clip_all_none(self): + arr = np.arange(10, dtype=object) + assert_equal(np.clip(arr, None, None), arr) + assert_equal(np.clip(arr), arr) + + def test_clip_invalid_casting(self): + a = np.arange(10, dtype=object) + with assert_raises_regex(ValueError, + 'casting must be one of'): + self.fastclip(a, 1, 8, casting="garbage") + + @pytest.mark.parametrize("amin, amax", [ + # two scalars + (1, 0), + # mix scalar and array + (1, np.zeros(10)), + # two arrays + (np.ones(10), np.zeros(10)), + ]) + def test_clip_value_min_max_flip(self, amin, amax): + a = np.arange(10, dtype=np.int64) + # requirement from ufunc_docstrings.py + expected = np.minimum(np.maximum(a, amin), amax) + actual = np.clip(a, amin, amax) + assert_equal(actual, expected) + + @pytest.mark.parametrize("arr, amin, amax, exp", [ + # for a bug in npy_ObjectClip, based on a + # case produced by hypothesis + (np.zeros(10, dtype=object), + 0, + -2**64+1, + np.full(10, -2**64+1, dtype=object)), + # for bugs in NPY_TIMEDELTA_MAX, based on a case + # produced by hypothesis + (np.zeros(10, dtype='m8') - 1, + 0, + 0, + np.zeros(10, dtype='m8')), + ]) + def test_clip_problem_cases(self, arr, amin, amax, exp): + actual = np.clip(arr, amin, amax) + assert_equal(actual, exp) + + @pytest.mark.parametrize("arr, amin, amax", [ + # problematic scalar nan case from hypothesis + (np.zeros(10, dtype=np.int64), + np.array(np.nan), + np.zeros(10, dtype=np.int32)), + ]) + def test_clip_scalar_nan_propagation(self, arr, amin, amax): + # enforcement of scalar nan propagation for comparisons + # called through clip() + expected = np.minimum(np.maximum(arr, amin), amax) + actual = np.clip(arr, amin, amax) + assert_equal(actual, expected) + + @pytest.mark.xfail(reason="propagation doesn't match spec") + @pytest.mark.parametrize("arr, amin, amax", [ + (np.array([1] * 10, dtype='m8'), + np.timedelta64('NaT'), + np.zeros(10, dtype=np.int32)), + ]) + @pytest.mark.filterwarnings("ignore::DeprecationWarning") + def test_NaT_propagation(self, arr, amin, amax): + # NOTE: the expected function spec doesn't + # propagate NaT, but clip() now does + expected = np.minimum(np.maximum(arr, amin), amax) + actual = np.clip(arr, amin, amax) + assert_equal(actual, expected) + + @given( + data=st.data(), + arr=hynp.arrays( + dtype=hynp.integer_dtypes() | hynp.floating_dtypes(), + shape=hynp.array_shapes() + ) + ) + def test_clip_property(self, data, arr): + """A property-based test using Hypothesis. + + This aims for maximum generality: it could in principle generate *any* + valid inputs to np.clip, and in practice generates much more varied + inputs than human testers come up with. + + Because many of the inputs have tricky dependencies - compatible dtypes + and mutually-broadcastable shapes - we use `st.data()` strategy draw + values *inside* the test function, from strategies we construct based + on previous values. An alternative would be to define a custom strategy + with `@st.composite`, but until we have duplicated code inline is fine. + + That accounts for most of the function; the actual test is just three + lines to calculate and compare actual vs expected results! + """ + numeric_dtypes = hynp.integer_dtypes() | hynp.floating_dtypes() + # Generate shapes for the bounds which can be broadcast with each other + # and with the base shape. Below, we might decide to use scalar bounds, + # but it's clearer to generate these shapes unconditionally in advance. + in_shapes, result_shape = data.draw( + hynp.mutually_broadcastable_shapes( + num_shapes=2, base_shape=arr.shape + ) + ) + # Scalar `nan` is deprecated due to the differing behaviour it shows. + s = numeric_dtypes.flatmap( + lambda x: hynp.from_dtype(x, allow_nan=False)) + amin = data.draw(s | hynp.arrays(dtype=numeric_dtypes, + shape=in_shapes[0], elements={"allow_nan": False})) + amax = data.draw(s | hynp.arrays(dtype=numeric_dtypes, + shape=in_shapes[1], elements={"allow_nan": False})) + + # Then calculate our result and expected result and check that they're + # equal! See gh-12519 and gh-19457 for discussion deciding on this + # property and the result_type argument. + result = np.clip(arr, amin, amax) + t = np.result_type(arr, amin, amax) + expected = np.minimum(amax, np.maximum(arr, amin, dtype=t), dtype=t) + assert result.dtype == t + assert_array_equal(result, expected) + + def test_clip_min_max_args(self): + arr = np.arange(5) + + assert_array_equal(np.clip(arr), arr) + assert_array_equal(np.clip(arr, min=2, max=3), np.clip(arr, 2, 3)) + assert_array_equal(np.clip(arr, min=None, max=2), + np.clip(arr, None, 2)) + + with assert_raises_regex(TypeError, "missing 1 required positional " + "argument: 'a_max'"): + np.clip(arr, 2) + with assert_raises_regex(TypeError, "missing 1 required positional " + "argument: 'a_min'"): + np.clip(arr, a_max=2) + msg = ("Passing `min` or `max` keyword argument when `a_min` and " + "`a_max` are provided is forbidden.") + with assert_raises_regex(ValueError, msg): + np.clip(arr, 2, 3, max=3) + with assert_raises_regex(ValueError, msg): + np.clip(arr, 2, 3, min=2) + + @pytest.mark.parametrize("dtype,min,max", [ + ("int32", -2**32-1, 2**32), + ("int32", -2**320, None), + ("int32", None, 2**300), + ("int32", -1000, 2**32), + ("int32", -2**32-1, 1000), + ("uint8", -1, 129), + ]) + def test_out_of_bound_pyints(self, dtype, min, max): + a = np.arange(10000).astype(dtype) + # Check min only + c = np.clip(a, min=min, max=max) + assert not np.may_share_memory(a, c) + assert c.dtype == a.dtype + if min is not None: + assert (c >= min).all() + if max is not None: + assert (c <= max).all() + +class TestAllclose: + rtol = 1e-5 + atol = 1e-8 + + def setup_method(self): + self.olderr = np.seterr(invalid='ignore') + + def teardown_method(self): + np.seterr(**self.olderr) + + def tst_allclose(self, x, y): + assert_(np.allclose(x, y), "%s and %s not close" % (x, y)) + + def tst_not_allclose(self, x, y): + assert_(not np.allclose(x, y), "%s and %s shouldn't be close" % (x, y)) + + def test_ip_allclose(self): + # Parametric test factory. + arr = np.array([100, 1000]) + aran = np.arange(125).reshape((5, 5, 5)) + + atol = self.atol + rtol = self.rtol + + data = [([1, 0], [1, 0]), + ([atol], [0]), + ([1], [1+rtol+atol]), + (arr, arr + arr*rtol), + (arr, arr + arr*rtol + atol*2), + (aran, aran + aran*rtol), + (np.inf, np.inf), + (np.inf, [np.inf])] + + for (x, y) in data: + self.tst_allclose(x, y) + + def test_ip_not_allclose(self): + # Parametric test factory. + aran = np.arange(125).reshape((5, 5, 5)) + + atol = self.atol + rtol = self.rtol + + data = [([np.inf, 0], [1, np.inf]), + ([np.inf, 0], [1, 0]), + ([np.inf, np.inf], [1, np.inf]), + ([np.inf, np.inf], [1, 0]), + ([-np.inf, 0], [np.inf, 0]), + ([np.nan, 0], [np.nan, 0]), + ([atol*2], [0]), + ([1], [1+rtol+atol*2]), + (aran, aran + aran*atol + atol*2), + (np.array([np.inf, 1]), np.array([0, np.inf]))] + + for (x, y) in data: + self.tst_not_allclose(x, y) + + def test_no_parameter_modification(self): + x = np.array([np.inf, 1]) + y = np.array([0, np.inf]) + np.allclose(x, y) + assert_array_equal(x, np.array([np.inf, 1])) + assert_array_equal(y, np.array([0, np.inf])) + + def test_min_int(self): + # Could make problems because of abs(min_int) == min_int + min_int = np.iinfo(np.int_).min + a = np.array([min_int], dtype=np.int_) + assert_(np.allclose(a, a)) + + def test_equalnan(self): + x = np.array([1.0, np.nan]) + assert_(np.allclose(x, x, equal_nan=True)) + + def test_return_class_is_ndarray(self): + # Issue gh-6475 + # Check that allclose does not preserve subtypes + class Foo(np.ndarray): + def __new__(cls, *args, **kwargs): + return np.array(*args, **kwargs).view(cls) + + a = Foo([1]) + assert_(type(np.allclose(a, a)) is bool) + + +class TestIsclose: + rtol = 1e-5 + atol = 1e-8 + + def _setup(self): + atol = self.atol + rtol = self.rtol + arr = np.array([100, 1000]) + aran = np.arange(125).reshape((5, 5, 5)) + + self.all_close_tests = [ + ([1, 0], [1, 0]), + ([atol], [0]), + ([1], [1 + rtol + atol]), + (arr, arr + arr*rtol), + (arr, arr + arr*rtol + atol), + (aran, aran + aran*rtol), + (np.inf, np.inf), + (np.inf, [np.inf]), + ([np.inf, -np.inf], [np.inf, -np.inf]), + ] + self.none_close_tests = [ + ([np.inf, 0], [1, np.inf]), + ([np.inf, -np.inf], [1, 0]), + ([np.inf, np.inf], [1, -np.inf]), + ([np.inf, np.inf], [1, 0]), + ([np.nan, 0], [np.nan, -np.inf]), + ([atol*2], [0]), + ([1], [1 + rtol + atol*2]), + (aran, aran + rtol*1.1*aran + atol*1.1), + (np.array([np.inf, 1]), np.array([0, np.inf])), + ] + self.some_close_tests = [ + ([np.inf, 0], [np.inf, atol*2]), + ([atol, 1, 1e6*(1 + 2*rtol) + atol], [0, np.nan, 1e6]), + (np.arange(3), [0, 1, 2.1]), + (np.nan, [np.nan, np.nan, np.nan]), + ([0], [atol, np.inf, -np.inf, np.nan]), + (0, [atol, np.inf, -np.inf, np.nan]), + ] + self.some_close_results = [ + [True, False], + [True, False, False], + [True, True, False], + [False, False, False], + [True, False, False, False], + [True, False, False, False], + ] + + def test_ip_isclose(self): + self._setup() + tests = self.some_close_tests + results = self.some_close_results + for (x, y), result in zip(tests, results): + assert_array_equal(np.isclose(x, y), result) + + x = np.array([2.1, 2.1, 2.1, 2.1, 5, np.nan]) + y = np.array([2, 2, 2, 2, np.nan, 5]) + atol = [0.11, 0.09, 1e-8, 1e-8, 1, 1] + rtol = [1e-8, 1e-8, 0.06, 0.04, 1, 1] + expected = np.array([True, False, True, False, False, False]) + assert_array_equal(np.isclose(x, y, rtol=rtol, atol=atol), expected) + + message = "operands could not be broadcast together..." + atol = np.array([1e-8, 1e-8]) + with assert_raises(ValueError, msg=message): + np.isclose(x, y, atol=atol) + + rtol = np.array([1e-5, 1e-5]) + with assert_raises(ValueError, msg=message): + np.isclose(x, y, rtol=rtol) + + def test_nep50_isclose(self): + below_one = float(1.-np.finfo('f8').eps) + f32 = np.array(below_one, 'f4') # This is just 1 at float32 precision + assert f32 > np.array(below_one) + # NEP 50 broadcasting of python scalars + assert f32 == below_one + # Test that it works for isclose arguments too (and that those fail if + # one uses a numpy float64). + assert np.isclose(f32, below_one, atol=0, rtol=0) + assert np.isclose(f32, np.float32(0), atol=below_one) + assert np.isclose(f32, 2, atol=0, rtol=below_one/2) + assert not np.isclose(f32, np.float64(below_one), atol=0, rtol=0) + assert not np.isclose(f32, np.float32(0), atol=np.float64(below_one)) + assert not np.isclose(f32, 2, atol=0, rtol=np.float64(below_one/2)) + + def tst_all_isclose(self, x, y): + assert_(np.all(np.isclose(x, y)), "%s and %s not close" % (x, y)) + + def tst_none_isclose(self, x, y): + msg = "%s and %s shouldn't be close" + assert_(not np.any(np.isclose(x, y)), msg % (x, y)) + + def tst_isclose_allclose(self, x, y): + msg = "isclose.all() and allclose aren't same for %s and %s" + msg2 = "isclose and allclose aren't same for %s and %s" + if np.isscalar(x) and np.isscalar(y): + assert_(np.isclose(x, y) == np.allclose(x, y), msg=msg2 % (x, y)) + else: + assert_array_equal(np.isclose(x, y).all(), np.allclose(x, y), msg % (x, y)) + + def test_ip_all_isclose(self): + self._setup() + for (x, y) in self.all_close_tests: + self.tst_all_isclose(x, y) + + x = np.array([2.3, 3.6, 4.4, np.nan]) + y = np.array([2, 3, 4, np.nan]) + atol = [0.31, 0, 0, 1] + rtol = [0, 0.21, 0.11, 1] + assert np.allclose(x, y, atol=atol, rtol=rtol, equal_nan=True) + assert not np.allclose(x, y, atol=0.1, rtol=0.1, equal_nan=True) + + # Show that gh-14330 is resolved + assert np.allclose([1, 2, float('nan')], [1, 2, float('nan')], + atol=[1, 1, 1], equal_nan=True) + + def test_ip_none_isclose(self): + self._setup() + for (x, y) in self.none_close_tests: + self.tst_none_isclose(x, y) + + def test_ip_isclose_allclose(self): + self._setup() + tests = (self.all_close_tests + self.none_close_tests + + self.some_close_tests) + for (x, y) in tests: + self.tst_isclose_allclose(x, y) + + def test_equal_nan(self): + assert_array_equal(np.isclose(np.nan, np.nan, equal_nan=True), [True]) + arr = np.array([1.0, np.nan]) + assert_array_equal(np.isclose(arr, arr, equal_nan=True), [True, True]) + + def test_masked_arrays(self): + # Make sure to test the output type when arguments are interchanged. + + x = np.ma.masked_where([True, True, False], np.arange(3)) + assert_(type(x) is type(np.isclose(2, x))) + assert_(type(x) is type(np.isclose(x, 2))) + + x = np.ma.masked_where([True, True, False], [np.nan, np.inf, np.nan]) + assert_(type(x) is type(np.isclose(np.inf, x))) + assert_(type(x) is type(np.isclose(x, np.inf))) + + x = np.ma.masked_where([True, True, False], [np.nan, np.nan, np.nan]) + y = np.isclose(np.nan, x, equal_nan=True) + assert_(type(x) is type(y)) + # Ensure that the mask isn't modified... + assert_array_equal([True, True, False], y.mask) + y = np.isclose(x, np.nan, equal_nan=True) + assert_(type(x) is type(y)) + # Ensure that the mask isn't modified... + assert_array_equal([True, True, False], y.mask) + + x = np.ma.masked_where([True, True, False], [np.nan, np.nan, np.nan]) + y = np.isclose(x, x, equal_nan=True) + assert_(type(x) is type(y)) + # Ensure that the mask isn't modified... + assert_array_equal([True, True, False], y.mask) + + def test_scalar_return(self): + assert_(np.isscalar(np.isclose(1, 1))) + + def test_no_parameter_modification(self): + x = np.array([np.inf, 1]) + y = np.array([0, np.inf]) + np.isclose(x, y) + assert_array_equal(x, np.array([np.inf, 1])) + assert_array_equal(y, np.array([0, np.inf])) + + def test_non_finite_scalar(self): + # GH7014, when two scalars are compared the output should also be a + # scalar + assert_(np.isclose(np.inf, -np.inf) is np.False_) + assert_(np.isclose(0, np.inf) is np.False_) + assert_(type(np.isclose(0, np.inf)) is np.bool) + + def test_timedelta(self): + # Allclose currently works for timedelta64 as long as `atol` is + # an integer or also a timedelta64 + a = np.array([[1, 2, 3, "NaT"]], dtype="m8[ns]") + assert np.isclose(a, a, atol=0, equal_nan=True).all() + assert np.isclose(a, a, atol=np.timedelta64(1, "ns"), equal_nan=True).all() + assert np.allclose(a, a, atol=0, equal_nan=True) + assert np.allclose(a, a, atol=np.timedelta64(1, "ns"), equal_nan=True) + + +class TestStdVar: + def setup_method(self): + self.A = np.array([1, -1, 1, -1]) + self.real_var = 1 + + def test_basic(self): + assert_almost_equal(np.var(self.A), self.real_var) + assert_almost_equal(np.std(self.A)**2, self.real_var) + + def test_scalars(self): + assert_equal(np.var(1), 0) + assert_equal(np.std(1), 0) + + def test_ddof1(self): + assert_almost_equal(np.var(self.A, ddof=1), + self.real_var * len(self.A) / (len(self.A) - 1)) + assert_almost_equal(np.std(self.A, ddof=1)**2, + self.real_var*len(self.A) / (len(self.A) - 1)) + + def test_ddof2(self): + assert_almost_equal(np.var(self.A, ddof=2), + self.real_var * len(self.A) / (len(self.A) - 2)) + assert_almost_equal(np.std(self.A, ddof=2)**2, + self.real_var * len(self.A) / (len(self.A) - 2)) + + def test_correction(self): + assert_almost_equal( + np.var(self.A, correction=1), np.var(self.A, ddof=1) + ) + assert_almost_equal( + np.std(self.A, correction=1), np.std(self.A, ddof=1) + ) + + err_msg = "ddof and correction can't be provided simultaneously." + + with assert_raises_regex(ValueError, err_msg): + np.var(self.A, ddof=1, correction=0) + + with assert_raises_regex(ValueError, err_msg): + np.std(self.A, ddof=1, correction=1) + + def test_out_scalar(self): + d = np.arange(10) + out = np.array(0.) + r = np.std(d, out=out) + assert_(r is out) + assert_array_equal(r, out) + r = np.var(d, out=out) + assert_(r is out) + assert_array_equal(r, out) + r = np.mean(d, out=out) + assert_(r is out) + assert_array_equal(r, out) + + +class TestStdVarComplex: + def test_basic(self): + A = np.array([1, 1.j, -1, -1.j]) + real_var = 1 + assert_almost_equal(np.var(A), real_var) + assert_almost_equal(np.std(A)**2, real_var) + + def test_scalars(self): + assert_equal(np.var(1j), 0) + assert_equal(np.std(1j), 0) + + +class TestCreationFuncs: + # Test ones, zeros, empty and full. + + def setup_method(self): + dtypes = {np.dtype(tp) for tp in itertools.chain(*sctypes.values())} + # void, bytes, str + variable_sized = {tp for tp in dtypes if tp.str.endswith('0')} + keyfunc = lambda dtype: dtype.str + self.dtypes = sorted(dtypes - variable_sized | + {np.dtype(tp.str.replace("0", str(i))) + for tp in variable_sized for i in range(1, 10)}, + key=keyfunc) + self.dtypes += [type(dt) for dt in sorted(dtypes, key=keyfunc)] + self.orders = {'C': 'c_contiguous', 'F': 'f_contiguous'} + self.ndims = 10 + + def check_function(self, func, fill_value=None): + par = ((0, 1, 2), + range(self.ndims), + self.orders, + self.dtypes) + fill_kwarg = {} + if fill_value is not None: + fill_kwarg = {'fill_value': fill_value} + + for size, ndims, order, dtype in itertools.product(*par): + shape = ndims * [size] + + is_void = dtype is np.dtypes.VoidDType or ( + isinstance(dtype, np.dtype) and dtype.str.startswith('|V')) + + # do not fill void type + if fill_kwarg and is_void: + continue + + arr = func(shape, order=order, dtype=dtype, + **fill_kwarg) + + if isinstance(dtype, np.dtype): + assert_equal(arr.dtype, dtype) + elif isinstance(dtype, type(np.dtype)): + if dtype in (np.dtypes.StrDType, np.dtypes.BytesDType): + dtype_str = np.dtype(dtype.type).str.replace('0', '1') + assert_equal(arr.dtype, np.dtype(dtype_str)) + else: + assert_equal(arr.dtype, np.dtype(dtype.type)) + assert_(getattr(arr.flags, self.orders[order])) + + if fill_value is not None: + if arr.dtype.str.startswith('|S'): + val = str(fill_value) + else: + val = fill_value + assert_equal(arr, dtype.type(val)) + + def test_zeros(self): + self.check_function(np.zeros) + + def test_ones(self): + self.check_function(np.ones) + + def test_empty(self): + self.check_function(np.empty) + + def test_full(self): + self.check_function(np.full, 0) + self.check_function(np.full, 1) + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_for_reference_leak(self): + # Make sure we have an object for reference + dim = 1 + beg = sys.getrefcount(dim) + np.zeros([dim]*10) + assert_(sys.getrefcount(dim) == beg) + np.ones([dim]*10) + assert_(sys.getrefcount(dim) == beg) + np.empty([dim]*10) + assert_(sys.getrefcount(dim) == beg) + np.full([dim]*10, 0) + assert_(sys.getrefcount(dim) == beg) + + +class TestLikeFuncs: + '''Test ones_like, zeros_like, empty_like and full_like''' + + def setup_method(self): + self.data = [ + # Array scalars + (np.array(3.), None), + (np.array(3), 'f8'), + # 1D arrays + (np.arange(6, dtype='f4'), None), + (np.arange(6), 'c16'), + # 2D C-layout arrays + (np.arange(6).reshape(2, 3), None), + (np.arange(6).reshape(3, 2), 'i1'), + # 2D F-layout arrays + (np.arange(6).reshape((2, 3), order='F'), None), + (np.arange(6).reshape((3, 2), order='F'), 'i1'), + # 3D C-layout arrays + (np.arange(24).reshape(2, 3, 4), None), + (np.arange(24).reshape(4, 3, 2), 'f4'), + # 3D F-layout arrays + (np.arange(24).reshape((2, 3, 4), order='F'), None), + (np.arange(24).reshape((4, 3, 2), order='F'), 'f4'), + # 3D non-C/F-layout arrays + (np.arange(24).reshape(2, 3, 4).swapaxes(0, 1), None), + (np.arange(24).reshape(4, 3, 2).swapaxes(0, 1), '?'), + ] + self.shapes = [(), (5,), (5,6,), (5,6,7,)] + + def compare_array_value(self, dz, value, fill_value): + if value is not None: + if fill_value: + # Conversion is close to what np.full_like uses + # but we may want to convert directly in the future + # which may result in errors (where this does not). + z = np.array(value).astype(dz.dtype) + assert_(np.all(dz == z)) + else: + assert_(np.all(dz == value)) + + def check_like_function(self, like_function, value, fill_value=False): + if fill_value: + fill_kwarg = {'fill_value': value} + else: + fill_kwarg = {} + for d, dtype in self.data: + # default (K) order, dtype + dz = like_function(d, dtype=dtype, **fill_kwarg) + assert_equal(dz.shape, d.shape) + assert_equal(np.array(dz.strides)*d.dtype.itemsize, + np.array(d.strides)*dz.dtype.itemsize) + assert_equal(d.flags.c_contiguous, dz.flags.c_contiguous) + assert_equal(d.flags.f_contiguous, dz.flags.f_contiguous) + if dtype is None: + assert_equal(dz.dtype, d.dtype) + else: + assert_equal(dz.dtype, np.dtype(dtype)) + self.compare_array_value(dz, value, fill_value) + + # C order, default dtype + dz = like_function(d, order='C', dtype=dtype, **fill_kwarg) + assert_equal(dz.shape, d.shape) + assert_(dz.flags.c_contiguous) + if dtype is None: + assert_equal(dz.dtype, d.dtype) + else: + assert_equal(dz.dtype, np.dtype(dtype)) + self.compare_array_value(dz, value, fill_value) + + # F order, default dtype + dz = like_function(d, order='F', dtype=dtype, **fill_kwarg) + assert_equal(dz.shape, d.shape) + assert_(dz.flags.f_contiguous) + if dtype is None: + assert_equal(dz.dtype, d.dtype) + else: + assert_equal(dz.dtype, np.dtype(dtype)) + self.compare_array_value(dz, value, fill_value) + + # A order + dz = like_function(d, order='A', dtype=dtype, **fill_kwarg) + assert_equal(dz.shape, d.shape) + if d.flags.f_contiguous: + assert_(dz.flags.f_contiguous) + else: + assert_(dz.flags.c_contiguous) + if dtype is None: + assert_equal(dz.dtype, d.dtype) + else: + assert_equal(dz.dtype, np.dtype(dtype)) + self.compare_array_value(dz, value, fill_value) + + # Test the 'shape' parameter + for s in self.shapes: + for o in 'CFA': + sz = like_function(d, dtype=dtype, shape=s, order=o, + **fill_kwarg) + assert_equal(sz.shape, s) + if dtype is None: + assert_equal(sz.dtype, d.dtype) + else: + assert_equal(sz.dtype, np.dtype(dtype)) + if o == 'C' or (o == 'A' and d.flags.c_contiguous): + assert_(sz.flags.c_contiguous) + elif o == 'F' or (o == 'A' and d.flags.f_contiguous): + assert_(sz.flags.f_contiguous) + self.compare_array_value(sz, value, fill_value) + + if (d.ndim != len(s)): + assert_equal(np.argsort(like_function(d, dtype=dtype, + shape=s, order='K', + **fill_kwarg).strides), + np.argsort(np.empty(s, dtype=dtype, + order='C').strides)) + else: + assert_equal(np.argsort(like_function(d, dtype=dtype, + shape=s, order='K', + **fill_kwarg).strides), + np.argsort(d.strides)) + + # Test the 'subok' parameter + class MyNDArray(np.ndarray): + pass + + a = np.array([[1, 2], [3, 4]]).view(MyNDArray) + + b = like_function(a, **fill_kwarg) + assert_(type(b) is MyNDArray) + + b = like_function(a, subok=False, **fill_kwarg) + assert_(type(b) is not MyNDArray) + + # Test invalid dtype + with assert_raises(TypeError): + a = np.array(b"abc") + like_function(a, dtype="S-1", **fill_kwarg) + + def test_ones_like(self): + self.check_like_function(np.ones_like, 1) + + def test_zeros_like(self): + self.check_like_function(np.zeros_like, 0) + + def test_empty_like(self): + self.check_like_function(np.empty_like, None) + + def test_filled_like(self): + self.check_like_function(np.full_like, 0, True) + self.check_like_function(np.full_like, 1, True) + # Large integers may overflow, but using int64 is OK (casts) + # see also gh-27075 + with pytest.raises(OverflowError): + np.full_like(np.ones(3, dtype=np.int8), 1000) + self.check_like_function(np.full_like, np.int64(1000), True) + self.check_like_function(np.full_like, 123.456, True) + # Inf to integer casts cause invalid-value errors: ignore them. + with np.errstate(invalid="ignore"): + self.check_like_function(np.full_like, np.inf, True) + + @pytest.mark.parametrize('likefunc', [np.empty_like, np.full_like, + np.zeros_like, np.ones_like]) + @pytest.mark.parametrize('dtype', [str, bytes]) + def test_dtype_str_bytes(self, likefunc, dtype): + # Regression test for gh-19860 + a = np.arange(16).reshape(2, 8) + b = a[:, ::2] # Ensure b is not contiguous. + kwargs = {'fill_value': ''} if likefunc == np.full_like else {} + result = likefunc(b, dtype=dtype, **kwargs) + if dtype == str: + assert result.strides == (16, 4) + else: + # dtype is bytes + assert result.strides == (4, 1) + + +class TestCorrelate: + def _setup(self, dt): + self.x = np.array([1, 2, 3, 4, 5], dtype=dt) + self.xs = np.arange(1, 20)[::3] + self.y = np.array([-1, -2, -3], dtype=dt) + self.z1 = np.array([-3., -8., -14., -20., -26., -14., -5.], dtype=dt) + self.z1_4 = np.array([-2., -5., -8., -11., -14., -5.], dtype=dt) + self.z1r = np.array([-15., -22., -22., -16., -10., -4., -1.], dtype=dt) + self.z2 = np.array([-5., -14., -26., -20., -14., -8., -3.], dtype=dt) + self.z2r = np.array([-1., -4., -10., -16., -22., -22., -15.], dtype=dt) + self.zs = np.array([-3., -14., -30., -48., -66., -84., + -102., -54., -19.], dtype=dt) + + def test_float(self): + self._setup(float) + z = np.correlate(self.x, self.y, 'full') + assert_array_almost_equal(z, self.z1) + z = np.correlate(self.x, self.y[:-1], 'full') + assert_array_almost_equal(z, self.z1_4) + z = np.correlate(self.y, self.x, 'full') + assert_array_almost_equal(z, self.z2) + z = np.correlate(self.x[::-1], self.y, 'full') + assert_array_almost_equal(z, self.z1r) + z = np.correlate(self.y, self.x[::-1], 'full') + assert_array_almost_equal(z, self.z2r) + z = np.correlate(self.xs, self.y, 'full') + assert_array_almost_equal(z, self.zs) + + def test_object(self): + self._setup(Decimal) + z = np.correlate(self.x, self.y, 'full') + assert_array_almost_equal(z, self.z1) + z = np.correlate(self.y, self.x, 'full') + assert_array_almost_equal(z, self.z2) + + def test_no_overwrite(self): + d = np.ones(100) + k = np.ones(3) + np.correlate(d, k) + assert_array_equal(d, np.ones(100)) + assert_array_equal(k, np.ones(3)) + + def test_complex(self): + x = np.array([1, 2, 3, 4+1j], dtype=complex) + y = np.array([-1, -2j, 3+1j], dtype=complex) + r_z = np.array([3-1j, 6, 8+1j, 11+5j, -5+8j, -4-1j], dtype=complex) + r_z = r_z[::-1].conjugate() + z = np.correlate(y, x, mode='full') + assert_array_almost_equal(z, r_z) + + def test_zero_size(self): + with pytest.raises(ValueError): + np.correlate(np.array([]), np.ones(1000), mode='full') + with pytest.raises(ValueError): + np.correlate(np.ones(1000), np.array([]), mode='full') + + def test_mode(self): + d = np.ones(100) + k = np.ones(3) + default_mode = np.correlate(d, k, mode='valid') + with assert_warns(DeprecationWarning): + valid_mode = np.correlate(d, k, mode='v') + assert_array_equal(valid_mode, default_mode) + # integer mode + with assert_raises(ValueError): + np.correlate(d, k, mode=-1) + assert_array_equal(np.correlate(d, k, mode=0), valid_mode) + # illegal arguments + with assert_raises(TypeError): + np.correlate(d, k, mode=None) + + +class TestConvolve: + def test_object(self): + d = [1.] * 100 + k = [1.] * 3 + assert_array_almost_equal(np.convolve(d, k)[2:-2], np.full(98, 3)) + + def test_no_overwrite(self): + d = np.ones(100) + k = np.ones(3) + np.convolve(d, k) + assert_array_equal(d, np.ones(100)) + assert_array_equal(k, np.ones(3)) + + def test_mode(self): + d = np.ones(100) + k = np.ones(3) + default_mode = np.convolve(d, k, mode='full') + with assert_warns(DeprecationWarning): + full_mode = np.convolve(d, k, mode='f') + assert_array_equal(full_mode, default_mode) + # integer mode + with assert_raises(ValueError): + np.convolve(d, k, mode=-1) + assert_array_equal(np.convolve(d, k, mode=2), full_mode) + # illegal arguments + with assert_raises(TypeError): + np.convolve(d, k, mode=None) + + +class TestArgwhere: + + @pytest.mark.parametrize('nd', [0, 1, 2]) + def test_nd(self, nd): + # get an nd array with multiple elements in every dimension + x = np.empty((2,)*nd, bool) + + # none + x[...] = False + assert_equal(np.argwhere(x).shape, (0, nd)) + + # only one + x[...] = False + x.flat[0] = True + assert_equal(np.argwhere(x).shape, (1, nd)) + + # all but one + x[...] = True + x.flat[0] = False + assert_equal(np.argwhere(x).shape, (x.size - 1, nd)) + + # all + x[...] = True + assert_equal(np.argwhere(x).shape, (x.size, nd)) + + def test_2D(self): + x = np.arange(6).reshape((2, 3)) + assert_array_equal(np.argwhere(x > 1), + [[0, 2], + [1, 0], + [1, 1], + [1, 2]]) + + def test_list(self): + assert_equal(np.argwhere([4, 0, 2, 1, 3]), [[0], [2], [3], [4]]) + + +class TestRoll: + def test_roll1d(self): + x = np.arange(10) + xr = np.roll(x, 2) + assert_equal(xr, np.array([8, 9, 0, 1, 2, 3, 4, 5, 6, 7])) + + def test_roll2d(self): + x2 = np.reshape(np.arange(10), (2, 5)) + x2r = np.roll(x2, 1) + assert_equal(x2r, np.array([[9, 0, 1, 2, 3], [4, 5, 6, 7, 8]])) + + x2r = np.roll(x2, 1, axis=0) + assert_equal(x2r, np.array([[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]])) + + x2r = np.roll(x2, 1, axis=1) + assert_equal(x2r, np.array([[4, 0, 1, 2, 3], [9, 5, 6, 7, 8]])) + + # Roll multiple axes at once. + x2r = np.roll(x2, 1, axis=(0, 1)) + assert_equal(x2r, np.array([[9, 5, 6, 7, 8], [4, 0, 1, 2, 3]])) + + x2r = np.roll(x2, (1, 0), axis=(0, 1)) + assert_equal(x2r, np.array([[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]])) + + x2r = np.roll(x2, (-1, 0), axis=(0, 1)) + assert_equal(x2r, np.array([[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]])) + + x2r = np.roll(x2, (0, 1), axis=(0, 1)) + assert_equal(x2r, np.array([[4, 0, 1, 2, 3], [9, 5, 6, 7, 8]])) + + x2r = np.roll(x2, (0, -1), axis=(0, 1)) + assert_equal(x2r, np.array([[1, 2, 3, 4, 0], [6, 7, 8, 9, 5]])) + + x2r = np.roll(x2, (1, 1), axis=(0, 1)) + assert_equal(x2r, np.array([[9, 5, 6, 7, 8], [4, 0, 1, 2, 3]])) + + x2r = np.roll(x2, (-1, -1), axis=(0, 1)) + assert_equal(x2r, np.array([[6, 7, 8, 9, 5], [1, 2, 3, 4, 0]])) + + # Roll the same axis multiple times. + x2r = np.roll(x2, 1, axis=(0, 0)) + assert_equal(x2r, np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])) + + x2r = np.roll(x2, 1, axis=(1, 1)) + assert_equal(x2r, np.array([[3, 4, 0, 1, 2], [8, 9, 5, 6, 7]])) + + # Roll more than one turn in either direction. + x2r = np.roll(x2, 6, axis=1) + assert_equal(x2r, np.array([[4, 0, 1, 2, 3], [9, 5, 6, 7, 8]])) + + x2r = np.roll(x2, -4, axis=1) + assert_equal(x2r, np.array([[4, 0, 1, 2, 3], [9, 5, 6, 7, 8]])) + + def test_roll_empty(self): + x = np.array([]) + assert_equal(np.roll(x, 1), np.array([])) + + def test_roll_unsigned_shift(self): + x = np.arange(4) + shift = np.uint16(2) + assert_equal(np.roll(x, shift), np.roll(x, 2)) + + shift = np.uint64(2**63+2) + assert_equal(np.roll(x, shift), np.roll(x, 2)) + + def test_roll_big_int(self): + x = np.arange(4) + assert_equal(np.roll(x, 2**100), x) + + +class TestRollaxis: + + # expected shape indexed by (axis, start) for array of + # shape (1, 2, 3, 4) + tgtshape = {(0, 0): (1, 2, 3, 4), (0, 1): (1, 2, 3, 4), + (0, 2): (2, 1, 3, 4), (0, 3): (2, 3, 1, 4), + (0, 4): (2, 3, 4, 1), + (1, 0): (2, 1, 3, 4), (1, 1): (1, 2, 3, 4), + (1, 2): (1, 2, 3, 4), (1, 3): (1, 3, 2, 4), + (1, 4): (1, 3, 4, 2), + (2, 0): (3, 1, 2, 4), (2, 1): (1, 3, 2, 4), + (2, 2): (1, 2, 3, 4), (2, 3): (1, 2, 3, 4), + (2, 4): (1, 2, 4, 3), + (3, 0): (4, 1, 2, 3), (3, 1): (1, 4, 2, 3), + (3, 2): (1, 2, 4, 3), (3, 3): (1, 2, 3, 4), + (3, 4): (1, 2, 3, 4)} + + def test_exceptions(self): + a = np.arange(1*2*3*4).reshape(1, 2, 3, 4) + assert_raises(AxisError, np.rollaxis, a, -5, 0) + assert_raises(AxisError, np.rollaxis, a, 0, -5) + assert_raises(AxisError, np.rollaxis, a, 4, 0) + assert_raises(AxisError, np.rollaxis, a, 0, 5) + + def test_results(self): + a = np.arange(1*2*3*4).reshape(1, 2, 3, 4).copy() + aind = np.indices(a.shape) + assert_(a.flags['OWNDATA']) + for (i, j) in self.tgtshape: + # positive axis, positive start + res = np.rollaxis(a, axis=i, start=j) + i0, i1, i2, i3 = aind[np.array(res.shape) - 1] + assert_(np.all(res[i0, i1, i2, i3] == a)) + assert_(res.shape == self.tgtshape[(i, j)], str((i,j))) + assert_(not res.flags['OWNDATA']) + + # negative axis, positive start + ip = i + 1 + res = np.rollaxis(a, axis=-ip, start=j) + i0, i1, i2, i3 = aind[np.array(res.shape) - 1] + assert_(np.all(res[i0, i1, i2, i3] == a)) + assert_(res.shape == self.tgtshape[(4 - ip, j)]) + assert_(not res.flags['OWNDATA']) + + # positive axis, negative start + jp = j + 1 if j < 4 else j + res = np.rollaxis(a, axis=i, start=-jp) + i0, i1, i2, i3 = aind[np.array(res.shape) - 1] + assert_(np.all(res[i0, i1, i2, i3] == a)) + assert_(res.shape == self.tgtshape[(i, 4 - jp)]) + assert_(not res.flags['OWNDATA']) + + # negative axis, negative start + ip = i + 1 + jp = j + 1 if j < 4 else j + res = np.rollaxis(a, axis=-ip, start=-jp) + i0, i1, i2, i3 = aind[np.array(res.shape) - 1] + assert_(np.all(res[i0, i1, i2, i3] == a)) + assert_(res.shape == self.tgtshape[(4 - ip, 4 - jp)]) + assert_(not res.flags['OWNDATA']) + + +class TestMoveaxis: + def test_move_to_end(self): + x = np.random.randn(5, 6, 7) + for source, expected in [(0, (6, 7, 5)), + (1, (5, 7, 6)), + (2, (5, 6, 7)), + (-1, (5, 6, 7))]: + actual = np.moveaxis(x, source, -1).shape + assert_(actual, expected) + + def test_move_new_position(self): + x = np.random.randn(1, 2, 3, 4) + for source, destination, expected in [ + (0, 1, (2, 1, 3, 4)), + (1, 2, (1, 3, 2, 4)), + (1, -1, (1, 3, 4, 2)), + ]: + actual = np.moveaxis(x, source, destination).shape + assert_(actual, expected) + + def test_preserve_order(self): + x = np.zeros((1, 2, 3, 4)) + for source, destination in [ + (0, 0), + (3, -1), + (-1, 3), + ([0, -1], [0, -1]), + ([2, 0], [2, 0]), + (range(4), range(4)), + ]: + actual = np.moveaxis(x, source, destination).shape + assert_(actual, (1, 2, 3, 4)) + + def test_move_multiples(self): + x = np.zeros((0, 1, 2, 3)) + for source, destination, expected in [ + ([0, 1], [2, 3], (2, 3, 0, 1)), + ([2, 3], [0, 1], (2, 3, 0, 1)), + ([0, 1, 2], [2, 3, 0], (2, 3, 0, 1)), + ([3, 0], [1, 0], (0, 3, 1, 2)), + ([0, 3], [0, 1], (0, 3, 1, 2)), + ]: + actual = np.moveaxis(x, source, destination).shape + assert_(actual, expected) + + def test_errors(self): + x = np.random.randn(1, 2, 3) + assert_raises_regex(AxisError, 'source.*out of bounds', + np.moveaxis, x, 3, 0) + assert_raises_regex(AxisError, 'source.*out of bounds', + np.moveaxis, x, -4, 0) + assert_raises_regex(AxisError, 'destination.*out of bounds', + np.moveaxis, x, 0, 5) + assert_raises_regex(ValueError, 'repeated axis in `source`', + np.moveaxis, x, [0, 0], [0, 1]) + assert_raises_regex(ValueError, 'repeated axis in `destination`', + np.moveaxis, x, [0, 1], [1, 1]) + assert_raises_regex(ValueError, 'must have the same number', + np.moveaxis, x, 0, [0, 1]) + assert_raises_regex(ValueError, 'must have the same number', + np.moveaxis, x, [0, 1], [0]) + + def test_array_likes(self): + x = np.ma.zeros((1, 2, 3)) + result = np.moveaxis(x, 0, 0) + assert_(x.shape, result.shape) + assert_(isinstance(result, np.ma.MaskedArray)) + + x = [1, 2, 3] + result = np.moveaxis(x, 0, 0) + assert_(x, list(result)) + assert_(isinstance(result, np.ndarray)) + + +class TestCross: + @pytest.mark.filterwarnings( + "ignore:.*2-dimensional vectors.*:DeprecationWarning" + ) + def test_2x2(self): + u = [1, 2] + v = [3, 4] + z = -2 + cp = np.cross(u, v) + assert_equal(cp, z) + cp = np.cross(v, u) + assert_equal(cp, -z) + + @pytest.mark.filterwarnings( + "ignore:.*2-dimensional vectors.*:DeprecationWarning" + ) + def test_2x3(self): + u = [1, 2] + v = [3, 4, 5] + z = np.array([10, -5, -2]) + cp = np.cross(u, v) + assert_equal(cp, z) + cp = np.cross(v, u) + assert_equal(cp, -z) + + def test_3x3(self): + u = [1, 2, 3] + v = [4, 5, 6] + z = np.array([-3, 6, -3]) + cp = np.cross(u, v) + assert_equal(cp, z) + cp = np.cross(v, u) + assert_equal(cp, -z) + + @pytest.mark.filterwarnings( + "ignore:.*2-dimensional vectors.*:DeprecationWarning" + ) + def test_broadcasting(self): + # Ticket #2624 (Trac #2032) + u = np.tile([1, 2], (11, 1)) + v = np.tile([3, 4], (11, 1)) + z = -2 + assert_equal(np.cross(u, v), z) + assert_equal(np.cross(v, u), -z) + assert_equal(np.cross(u, u), 0) + + u = np.tile([1, 2], (11, 1)).T + v = np.tile([3, 4, 5], (11, 1)) + z = np.tile([10, -5, -2], (11, 1)) + assert_equal(np.cross(u, v, axisa=0), z) + assert_equal(np.cross(v, u.T), -z) + assert_equal(np.cross(v, v), 0) + + u = np.tile([1, 2, 3], (11, 1)).T + v = np.tile([3, 4], (11, 1)).T + z = np.tile([-12, 9, -2], (11, 1)) + assert_equal(np.cross(u, v, axisa=0, axisb=0), z) + assert_equal(np.cross(v.T, u.T), -z) + assert_equal(np.cross(u.T, u.T), 0) + + u = np.tile([1, 2, 3], (5, 1)) + v = np.tile([4, 5, 6], (5, 1)).T + z = np.tile([-3, 6, -3], (5, 1)) + assert_equal(np.cross(u, v, axisb=0), z) + assert_equal(np.cross(v.T, u), -z) + assert_equal(np.cross(u, u), 0) + + @pytest.mark.filterwarnings( + "ignore:.*2-dimensional vectors.*:DeprecationWarning" + ) + def test_broadcasting_shapes(self): + u = np.ones((2, 1, 3)) + v = np.ones((5, 3)) + assert_equal(np.cross(u, v).shape, (2, 5, 3)) + u = np.ones((10, 3, 5)) + v = np.ones((2, 5)) + assert_equal(np.cross(u, v, axisa=1, axisb=0).shape, (10, 5, 3)) + assert_raises(AxisError, np.cross, u, v, axisa=1, axisb=2) + assert_raises(AxisError, np.cross, u, v, axisa=3, axisb=0) + u = np.ones((10, 3, 5, 7)) + v = np.ones((5, 7, 2)) + assert_equal(np.cross(u, v, axisa=1, axisc=2).shape, (10, 5, 3, 7)) + assert_raises(AxisError, np.cross, u, v, axisa=-5, axisb=2) + assert_raises(AxisError, np.cross, u, v, axisa=1, axisb=-4) + # gh-5885 + u = np.ones((3, 4, 2)) + for axisc in range(-2, 2): + assert_equal(np.cross(u, u, axisc=axisc).shape, (3, 4)) + + def test_uint8_int32_mixed_dtypes(self): + # regression test for gh-19138 + u = np.array([[195, 8, 9]], np.uint8) + v = np.array([250, 166, 68], np.int32) + z = np.array([[950, 11010, -30370]], dtype=np.int32) + assert_equal(np.cross(v, u), z) + assert_equal(np.cross(u, v), -z) + + @pytest.mark.parametrize("a, b", [(0, [1, 2]), ([1, 2], 3)]) + def test_zero_dimension(self, a, b): + with pytest.raises(ValueError) as exc: + np.cross(a, b) + assert "At least one array has zero dimension" in str(exc.value) + + +def test_outer_out_param(): + arr1 = np.ones((5,)) + arr2 = np.ones((2,)) + arr3 = np.linspace(-2, 2, 5) + out1 = np.ndarray(shape=(5,5)) + out2 = np.ndarray(shape=(2, 5)) + res1 = np.outer(arr1, arr3, out1) + assert_equal(res1, out1) + assert_equal(np.outer(arr2, arr3, out2), out2) + + +class TestIndices: + + def test_simple(self): + [x, y] = np.indices((4, 3)) + assert_array_equal(x, np.array([[0, 0, 0], + [1, 1, 1], + [2, 2, 2], + [3, 3, 3]])) + assert_array_equal(y, np.array([[0, 1, 2], + [0, 1, 2], + [0, 1, 2], + [0, 1, 2]])) + + def test_single_input(self): + [x] = np.indices((4,)) + assert_array_equal(x, np.array([0, 1, 2, 3])) + + [x] = np.indices((4,), sparse=True) + assert_array_equal(x, np.array([0, 1, 2, 3])) + + def test_scalar_input(self): + assert_array_equal([], np.indices(())) + assert_array_equal([], np.indices((), sparse=True)) + assert_array_equal([[]], np.indices((0,))) + assert_array_equal([[]], np.indices((0,), sparse=True)) + + def test_sparse(self): + [x, y] = np.indices((4,3), sparse=True) + assert_array_equal(x, np.array([[0], [1], [2], [3]])) + assert_array_equal(y, np.array([[0, 1, 2]])) + + @pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float32, np.float64]) + @pytest.mark.parametrize("dims", [(), (0,), (4, 3)]) + def test_return_type(self, dtype, dims): + inds = np.indices(dims, dtype=dtype) + assert_(inds.dtype == dtype) + + for arr in np.indices(dims, dtype=dtype, sparse=True): + assert_(arr.dtype == dtype) + + +class TestRequire: + flag_names = ['C', 'C_CONTIGUOUS', 'CONTIGUOUS', + 'F', 'F_CONTIGUOUS', 'FORTRAN', + 'A', 'ALIGNED', + 'W', 'WRITEABLE', + 'O', 'OWNDATA'] + + def generate_all_false(self, dtype): + arr = np.zeros((2, 2), [('junk', 'i1'), ('a', dtype)]) + arr.setflags(write=False) + a = arr['a'] + assert_(not a.flags['C']) + assert_(not a.flags['F']) + assert_(not a.flags['O']) + assert_(not a.flags['W']) + assert_(not a.flags['A']) + return a + + def set_and_check_flag(self, flag, dtype, arr): + if dtype is None: + dtype = arr.dtype + b = np.require(arr, dtype, [flag]) + assert_(b.flags[flag]) + assert_(b.dtype == dtype) + + # a further call to np.require ought to return the same array + # unless OWNDATA is specified. + c = np.require(b, None, [flag]) + if flag[0] != 'O': + assert_(c is b) + else: + assert_(c.flags[flag]) + + def test_require_each(self): + + id = ['f8', 'i4'] + fd = [None, 'f8', 'c16'] + for idtype, fdtype, flag in itertools.product(id, fd, self.flag_names): + a = self.generate_all_false(idtype) + self.set_and_check_flag(flag, fdtype, a) + + def test_unknown_requirement(self): + a = self.generate_all_false('f8') + assert_raises(KeyError, np.require, a, None, 'Q') + + def test_non_array_input(self): + a = np.require([1, 2, 3, 4], 'i4', ['C', 'A', 'O']) + assert_(a.flags['O']) + assert_(a.flags['C']) + assert_(a.flags['A']) + assert_(a.dtype == 'i4') + assert_equal(a, [1, 2, 3, 4]) + + def test_C_and_F_simul(self): + a = self.generate_all_false('f8') + assert_raises(ValueError, np.require, a, None, ['C', 'F']) + + def test_ensure_array(self): + class ArraySubclass(np.ndarray): + pass + + a = ArraySubclass((2, 2)) + b = np.require(a, None, ['E']) + assert_(type(b) is np.ndarray) + + def test_preserve_subtype(self): + class ArraySubclass(np.ndarray): + pass + + for flag in self.flag_names: + a = ArraySubclass((2, 2)) + self.set_and_check_flag(flag, None, a) + + +class TestBroadcast: + def test_broadcast_in_args(self): + # gh-5881 + arrs = [np.empty((6, 7)), np.empty((5, 6, 1)), np.empty((7,)), + np.empty((5, 1, 7))] + mits = [np.broadcast(*arrs), + np.broadcast(np.broadcast(*arrs[:0]), np.broadcast(*arrs[0:])), + np.broadcast(np.broadcast(*arrs[:1]), np.broadcast(*arrs[1:])), + np.broadcast(np.broadcast(*arrs[:2]), np.broadcast(*arrs[2:])), + np.broadcast(arrs[0], np.broadcast(*arrs[1:-1]), arrs[-1])] + for mit in mits: + assert_equal(mit.shape, (5, 6, 7)) + assert_equal(mit.ndim, 3) + assert_equal(mit.nd, 3) + assert_equal(mit.numiter, 4) + for a, ia in zip(arrs, mit.iters): + assert_(a is ia.base) + + def test_broadcast_single_arg(self): + # gh-6899 + arrs = [np.empty((5, 6, 7))] + mit = np.broadcast(*arrs) + assert_equal(mit.shape, (5, 6, 7)) + assert_equal(mit.ndim, 3) + assert_equal(mit.nd, 3) + assert_equal(mit.numiter, 1) + assert_(arrs[0] is mit.iters[0].base) + + def test_number_of_arguments(self): + arr = np.empty((5,)) + for j in range(70): + arrs = [arr] * j + if j > 64: + assert_raises(ValueError, np.broadcast, *arrs) + else: + mit = np.broadcast(*arrs) + assert_equal(mit.numiter, j) + + def test_broadcast_error_kwargs(self): + #gh-13455 + arrs = [np.empty((5, 6, 7))] + mit = np.broadcast(*arrs) + mit2 = np.broadcast(*arrs, **{}) + assert_equal(mit.shape, mit2.shape) + assert_equal(mit.ndim, mit2.ndim) + assert_equal(mit.nd, mit2.nd) + assert_equal(mit.numiter, mit2.numiter) + assert_(mit.iters[0].base is mit2.iters[0].base) + + assert_raises(ValueError, np.broadcast, 1, **{'x': 1}) + + def test_shape_mismatch_error_message(self): + with pytest.raises(ValueError, match=r"arg 0 with shape \(1, 3\) and " + r"arg 2 with shape \(2,\)"): + np.broadcast([[1, 2, 3]], [[4], [5]], [6, 7]) + + +class TestKeepdims: + + class sub_array(np.ndarray): + def sum(self, axis=None, dtype=None, out=None): + return np.ndarray.sum(self, axis, dtype, out, keepdims=True) + + def test_raise(self): + sub_class = self.sub_array + x = np.arange(30).view(sub_class) + assert_raises(TypeError, np.sum, x, keepdims=True) + + +class TestTensordot: + + def test_zero_dimension(self): + # Test resolution to issue #5663 + a = np.ndarray((3,0)) + b = np.ndarray((0,4)) + td = np.tensordot(a, b, (1, 0)) + assert_array_equal(td, np.dot(a, b)) + assert_array_equal(td, np.einsum('ij,jk', a, b)) + + def test_zero_dimensional(self): + # gh-12130 + arr_0d = np.array(1) + ret = np.tensordot(arr_0d, arr_0d, ([], [])) # contracting no axes is well defined + assert_array_equal(ret, arr_0d) + + +class TestAsType: + + def test_astype(self): + data = [[1, 2], [3, 4]] + actual = np.astype( + np.array(data, dtype=np.int64), np.uint32 + ) + expected = np.array(data, dtype=np.uint32) + + assert_array_equal(actual, expected) + assert_equal(actual.dtype, expected.dtype) + + assert np.shares_memory( + actual, np.astype(actual, actual.dtype, copy=False) + ) + + actual = np.astype(np.int64(10), np.float64) + expected = np.float64(10) + assert_equal(actual, expected) + assert_equal(actual.dtype, expected.dtype) + + with pytest.raises(TypeError, match="Input should be a NumPy array"): + np.astype(data, np.float64) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_numerictypes.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_numerictypes.py new file mode 100644 index 0000000000000000000000000000000000000000..310ba82bc5e757c60cf5294ab8d858c4047c8a4d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_numerictypes.py @@ -0,0 +1,620 @@ +import sys +import itertools + +import pytest +import numpy as np +import numpy._core.numerictypes as nt +from numpy._core.numerictypes import ( + issctype, sctype2char, maximum_sctype, sctypes +) +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_raises_regex, IS_PYPY +) + +# This is the structure of the table used for plain objects: +# +# +-+-+-+ +# |x|y|z| +# +-+-+-+ + +# Structure of a plain array description: +Pdescr = [ + ('x', 'i4', (2,)), + ('y', 'f8', (2, 2)), + ('z', 'u1')] + +# A plain list of tuples with values for testing: +PbufferT = [ + # x y z + ([3, 2], [[6., 4.], [6., 4.]], 8), + ([4, 3], [[7., 5.], [7., 5.]], 9), + ] + + +# This is the structure of the table used for nested objects (DON'T PANIC!): +# +# +-+---------------------------------+-----+----------+-+-+ +# |x|Info |color|info |y|z| +# | +-----+--+----------------+----+--+ +----+-----+ | | +# | |value|y2|Info2 |name|z2| |Name|Value| | | +# | | | +----+-----+--+--+ | | | | | | | +# | | | |name|value|y3|z3| | | | | | | | +# +-+-----+--+----+-----+--+--+----+--+-----+----+-----+-+-+ +# + +# The corresponding nested array description: +Ndescr = [ + ('x', 'i4', (2,)), + ('Info', [ + ('value', 'c16'), + ('y2', 'f8'), + ('Info2', [ + ('name', 'S2'), + ('value', 'c16', (2,)), + ('y3', 'f8', (2,)), + ('z3', 'u4', (2,))]), + ('name', 'S2'), + ('z2', 'b1')]), + ('color', 'S2'), + ('info', [ + ('Name', 'U8'), + ('Value', 'c16')]), + ('y', 'f8', (2, 2)), + ('z', 'u1')] + +NbufferT = [ + # x Info color info y z + # value y2 Info2 name z2 Name Value + # name value y3 z3 + ([3, 2], (6j, 6., (b'nn', [6j, 4j], [6., 4.], [1, 2]), b'NN', True), + b'cc', ('NN', 6j), [[6., 4.], [6., 4.]], 8), + ([4, 3], (7j, 7., (b'oo', [7j, 5j], [7., 5.], [2, 1]), b'OO', False), + b'dd', ('OO', 7j), [[7., 5.], [7., 5.]], 9), + ] + + +byteorder = {'little':'<', 'big':'>'}[sys.byteorder] + +def normalize_descr(descr): + "Normalize a description adding the platform byteorder." + + out = [] + for item in descr: + dtype = item[1] + if isinstance(dtype, str): + if dtype[0] not in ['|', '<', '>']: + onebyte = dtype[1:] == "1" + if onebyte or dtype[0] in ['S', 'V', 'b']: + dtype = "|" + dtype + else: + dtype = byteorder + dtype + if len(item) > 2 and np.prod(item[2]) > 1: + nitem = (item[0], dtype, item[2]) + else: + nitem = (item[0], dtype) + out.append(nitem) + elif isinstance(dtype, list): + l = normalize_descr(dtype) + out.append((item[0], l)) + else: + raise ValueError("Expected a str or list and got %s" % + (type(item))) + return out + + +############################################################ +# Creation tests +############################################################ + +class CreateZeros: + """Check the creation of heterogeneous arrays zero-valued""" + + def test_zeros0D(self): + """Check creation of 0-dimensional objects""" + h = np.zeros((), dtype=self._descr) + assert_(normalize_descr(self._descr) == h.dtype.descr) + assert_(h.dtype.fields['x'][0].name[:4] == 'void') + assert_(h.dtype.fields['x'][0].char == 'V') + assert_(h.dtype.fields['x'][0].type == np.void) + # A small check that data is ok + assert_equal(h['z'], np.zeros((), dtype='u1')) + + def test_zerosSD(self): + """Check creation of single-dimensional objects""" + h = np.zeros((2,), dtype=self._descr) + assert_(normalize_descr(self._descr) == h.dtype.descr) + assert_(h.dtype['y'].name[:4] == 'void') + assert_(h.dtype['y'].char == 'V') + assert_(h.dtype['y'].type == np.void) + # A small check that data is ok + assert_equal(h['z'], np.zeros((2,), dtype='u1')) + + def test_zerosMD(self): + """Check creation of multi-dimensional objects""" + h = np.zeros((2, 3), dtype=self._descr) + assert_(normalize_descr(self._descr) == h.dtype.descr) + assert_(h.dtype['z'].name == 'uint8') + assert_(h.dtype['z'].char == 'B') + assert_(h.dtype['z'].type == np.uint8) + # A small check that data is ok + assert_equal(h['z'], np.zeros((2, 3), dtype='u1')) + + +class TestCreateZerosPlain(CreateZeros): + """Check the creation of heterogeneous arrays zero-valued (plain)""" + _descr = Pdescr + +class TestCreateZerosNested(CreateZeros): + """Check the creation of heterogeneous arrays zero-valued (nested)""" + _descr = Ndescr + + +class CreateValues: + """Check the creation of heterogeneous arrays with values""" + + def test_tuple(self): + """Check creation from tuples""" + h = np.array(self._buffer, dtype=self._descr) + assert_(normalize_descr(self._descr) == h.dtype.descr) + if self.multiple_rows: + assert_(h.shape == (2,)) + else: + assert_(h.shape == ()) + + def test_list_of_tuple(self): + """Check creation from list of tuples""" + h = np.array([self._buffer], dtype=self._descr) + assert_(normalize_descr(self._descr) == h.dtype.descr) + if self.multiple_rows: + assert_(h.shape == (1, 2)) + else: + assert_(h.shape == (1,)) + + def test_list_of_list_of_tuple(self): + """Check creation from list of list of tuples""" + h = np.array([[self._buffer]], dtype=self._descr) + assert_(normalize_descr(self._descr) == h.dtype.descr) + if self.multiple_rows: + assert_(h.shape == (1, 1, 2)) + else: + assert_(h.shape == (1, 1)) + + +class TestCreateValuesPlainSingle(CreateValues): + """Check the creation of heterogeneous arrays (plain, single row)""" + _descr = Pdescr + multiple_rows = 0 + _buffer = PbufferT[0] + +class TestCreateValuesPlainMultiple(CreateValues): + """Check the creation of heterogeneous arrays (plain, multiple rows)""" + _descr = Pdescr + multiple_rows = 1 + _buffer = PbufferT + +class TestCreateValuesNestedSingle(CreateValues): + """Check the creation of heterogeneous arrays (nested, single row)""" + _descr = Ndescr + multiple_rows = 0 + _buffer = NbufferT[0] + +class TestCreateValuesNestedMultiple(CreateValues): + """Check the creation of heterogeneous arrays (nested, multiple rows)""" + _descr = Ndescr + multiple_rows = 1 + _buffer = NbufferT + + +############################################################ +# Reading tests +############################################################ + +class ReadValuesPlain: + """Check the reading of values in heterogeneous arrays (plain)""" + + def test_access_fields(self): + h = np.array(self._buffer, dtype=self._descr) + if not self.multiple_rows: + assert_(h.shape == ()) + assert_equal(h['x'], np.array(self._buffer[0], dtype='i4')) + assert_equal(h['y'], np.array(self._buffer[1], dtype='f8')) + assert_equal(h['z'], np.array(self._buffer[2], dtype='u1')) + else: + assert_(len(h) == 2) + assert_equal(h['x'], np.array([self._buffer[0][0], + self._buffer[1][0]], dtype='i4')) + assert_equal(h['y'], np.array([self._buffer[0][1], + self._buffer[1][1]], dtype='f8')) + assert_equal(h['z'], np.array([self._buffer[0][2], + self._buffer[1][2]], dtype='u1')) + + +class TestReadValuesPlainSingle(ReadValuesPlain): + """Check the creation of heterogeneous arrays (plain, single row)""" + _descr = Pdescr + multiple_rows = 0 + _buffer = PbufferT[0] + +class TestReadValuesPlainMultiple(ReadValuesPlain): + """Check the values of heterogeneous arrays (plain, multiple rows)""" + _descr = Pdescr + multiple_rows = 1 + _buffer = PbufferT + +class ReadValuesNested: + """Check the reading of values in heterogeneous arrays (nested)""" + + def test_access_top_fields(self): + """Check reading the top fields of a nested array""" + h = np.array(self._buffer, dtype=self._descr) + if not self.multiple_rows: + assert_(h.shape == ()) + assert_equal(h['x'], np.array(self._buffer[0], dtype='i4')) + assert_equal(h['y'], np.array(self._buffer[4], dtype='f8')) + assert_equal(h['z'], np.array(self._buffer[5], dtype='u1')) + else: + assert_(len(h) == 2) + assert_equal(h['x'], np.array([self._buffer[0][0], + self._buffer[1][0]], dtype='i4')) + assert_equal(h['y'], np.array([self._buffer[0][4], + self._buffer[1][4]], dtype='f8')) + assert_equal(h['z'], np.array([self._buffer[0][5], + self._buffer[1][5]], dtype='u1')) + + def test_nested1_acessors(self): + """Check reading the nested fields of a nested array (1st level)""" + h = np.array(self._buffer, dtype=self._descr) + if not self.multiple_rows: + assert_equal(h['Info']['value'], + np.array(self._buffer[1][0], dtype='c16')) + assert_equal(h['Info']['y2'], + np.array(self._buffer[1][1], dtype='f8')) + assert_equal(h['info']['Name'], + np.array(self._buffer[3][0], dtype='U2')) + assert_equal(h['info']['Value'], + np.array(self._buffer[3][1], dtype='c16')) + else: + assert_equal(h['Info']['value'], + np.array([self._buffer[0][1][0], + self._buffer[1][1][0]], + dtype='c16')) + assert_equal(h['Info']['y2'], + np.array([self._buffer[0][1][1], + self._buffer[1][1][1]], + dtype='f8')) + assert_equal(h['info']['Name'], + np.array([self._buffer[0][3][0], + self._buffer[1][3][0]], + dtype='U2')) + assert_equal(h['info']['Value'], + np.array([self._buffer[0][3][1], + self._buffer[1][3][1]], + dtype='c16')) + + def test_nested2_acessors(self): + """Check reading the nested fields of a nested array (2nd level)""" + h = np.array(self._buffer, dtype=self._descr) + if not self.multiple_rows: + assert_equal(h['Info']['Info2']['value'], + np.array(self._buffer[1][2][1], dtype='c16')) + assert_equal(h['Info']['Info2']['z3'], + np.array(self._buffer[1][2][3], dtype='u4')) + else: + assert_equal(h['Info']['Info2']['value'], + np.array([self._buffer[0][1][2][1], + self._buffer[1][1][2][1]], + dtype='c16')) + assert_equal(h['Info']['Info2']['z3'], + np.array([self._buffer[0][1][2][3], + self._buffer[1][1][2][3]], + dtype='u4')) + + def test_nested1_descriptor(self): + """Check access nested descriptors of a nested array (1st level)""" + h = np.array(self._buffer, dtype=self._descr) + assert_(h.dtype['Info']['value'].name == 'complex128') + assert_(h.dtype['Info']['y2'].name == 'float64') + assert_(h.dtype['info']['Name'].name == 'str256') + assert_(h.dtype['info']['Value'].name == 'complex128') + + def test_nested2_descriptor(self): + """Check access nested descriptors of a nested array (2nd level)""" + h = np.array(self._buffer, dtype=self._descr) + assert_(h.dtype['Info']['Info2']['value'].name == 'void256') + assert_(h.dtype['Info']['Info2']['z3'].name == 'void64') + + +class TestReadValuesNestedSingle(ReadValuesNested): + """Check the values of heterogeneous arrays (nested, single row)""" + _descr = Ndescr + multiple_rows = False + _buffer = NbufferT[0] + +class TestReadValuesNestedMultiple(ReadValuesNested): + """Check the values of heterogeneous arrays (nested, multiple rows)""" + _descr = Ndescr + multiple_rows = True + _buffer = NbufferT + +class TestEmptyField: + def test_assign(self): + a = np.arange(10, dtype=np.float32) + a.dtype = [("int", "<0i4"), ("float", "<2f4")] + assert_(a['int'].shape == (5, 0)) + assert_(a['float'].shape == (5, 2)) + + +class TestMultipleFields: + def setup_method(self): + self.ary = np.array([(1, 2, 3, 4), (5, 6, 7, 8)], dtype='i4,f4,i2,c8') + + def _bad_call(self): + return self.ary['f0', 'f1'] + + def test_no_tuple(self): + assert_raises(IndexError, self._bad_call) + + def test_return(self): + res = self.ary[['f0', 'f2']].tolist() + assert_(res == [(1, 3), (5, 7)]) + + +class TestIsSubDType: + # scalar types can be promoted into dtypes + wrappers = [np.dtype, lambda x: x] + + def test_both_abstract(self): + assert_(np.issubdtype(np.floating, np.inexact)) + assert_(not np.issubdtype(np.inexact, np.floating)) + + def test_same(self): + for cls in (np.float32, np.int32): + for w1, w2 in itertools.product(self.wrappers, repeat=2): + assert_(np.issubdtype(w1(cls), w2(cls))) + + def test_subclass(self): + # note we cannot promote floating to a dtype, as it would turn into a + # concrete type + for w in self.wrappers: + assert_(np.issubdtype(w(np.float32), np.floating)) + assert_(np.issubdtype(w(np.float64), np.floating)) + + def test_subclass_backwards(self): + for w in self.wrappers: + assert_(not np.issubdtype(np.floating, w(np.float32))) + assert_(not np.issubdtype(np.floating, w(np.float64))) + + def test_sibling_class(self): + for w1, w2 in itertools.product(self.wrappers, repeat=2): + assert_(not np.issubdtype(w1(np.float32), w2(np.float64))) + assert_(not np.issubdtype(w1(np.float64), w2(np.float32))) + + def test_nondtype_nonscalartype(self): + # See gh-14619 and gh-9505 which introduced the deprecation to fix + # this. These tests are directly taken from gh-9505 + assert not np.issubdtype(np.float32, 'float64') + assert not np.issubdtype(np.float32, 'f8') + assert not np.issubdtype(np.int32, str) + assert not np.issubdtype(np.int32, 'int64') + assert not np.issubdtype(np.str_, 'void') + # for the following the correct spellings are + # np.integer, np.floating, or np.complexfloating respectively: + assert not np.issubdtype(np.int8, int) # np.int8 is never np.int_ + assert not np.issubdtype(np.float32, float) + assert not np.issubdtype(np.complex64, complex) + assert not np.issubdtype(np.float32, "float") + assert not np.issubdtype(np.float64, "f") + + # Test the same for the correct first datatype and abstract one + # in the case of int, float, complex: + assert np.issubdtype(np.float64, 'float64') + assert np.issubdtype(np.float64, 'f8') + assert np.issubdtype(np.str_, str) + assert np.issubdtype(np.int64, 'int64') + assert np.issubdtype(np.void, 'void') + assert np.issubdtype(np.int8, np.integer) + assert np.issubdtype(np.float32, np.floating) + assert np.issubdtype(np.complex64, np.complexfloating) + assert np.issubdtype(np.float64, "float") + assert np.issubdtype(np.float32, "f") + + +class TestIsDType: + """ + Check correctness of `np.isdtype`. The test considers different argument + configurations: `np.isdtype(dtype, k1)` and `np.isdtype(dtype, (k1, k2))` + with concrete dtypes and dtype groups. + """ + dtype_group_dict = { + "signed integer": sctypes["int"], + "unsigned integer": sctypes["uint"], + "integral": sctypes["int"] + sctypes["uint"], + "real floating": sctypes["float"], + "complex floating": sctypes["complex"], + "numeric": ( + sctypes["int"] + sctypes["uint"] + sctypes["float"] + + sctypes["complex"] + ) + } + + @pytest.mark.parametrize( + "dtype,close_dtype", + [ + (np.int64, np.int32), (np.uint64, np.uint32), + (np.float64, np.float32), (np.complex128, np.complex64) + ] + ) + @pytest.mark.parametrize( + "dtype_group", + [ + None, "signed integer", "unsigned integer", "integral", + "real floating", "complex floating", "numeric" + ] + ) + def test_isdtype(self, dtype, close_dtype, dtype_group): + # First check if same dtypes return `true` and different ones + # give `false` (even if they're close in the dtype hierarchy!) + if dtype_group is None: + assert np.isdtype(dtype, dtype) + assert not np.isdtype(dtype, close_dtype) + assert np.isdtype(dtype, (dtype, close_dtype)) + + # Check that dtype and a dtype group that it belongs to + # return `true`, and `false` otherwise. + elif dtype in self.dtype_group_dict[dtype_group]: + assert np.isdtype(dtype, dtype_group) + assert np.isdtype(dtype, (close_dtype, dtype_group)) + else: + assert not np.isdtype(dtype, dtype_group) + + def test_isdtype_invalid_args(self): + with assert_raises_regex(TypeError, r".*must be a NumPy dtype.*"): + np.isdtype("int64", np.int64) + with assert_raises_regex(TypeError, r".*kind argument must.*"): + np.isdtype(np.int64, 1) + with assert_raises_regex(ValueError, r".*not a known kind name.*"): + np.isdtype(np.int64, "int64") + + def test_sctypes_complete(self): + # issue 26439: int32/intc were masking each other on 32-bit builds + assert np.int32 in sctypes['int'] + assert np.intc in sctypes['int'] + assert np.int64 in sctypes['int'] + assert np.uint32 in sctypes['uint'] + assert np.uintc in sctypes['uint'] + assert np.uint64 in sctypes['uint'] + +class TestSctypeDict: + def test_longdouble(self): + assert_(np._core.sctypeDict['float64'] is not np.longdouble) + assert_(np._core.sctypeDict['complex128'] is not np.clongdouble) + + def test_ulong(self): + assert np._core.sctypeDict['ulong'] is np.ulong + assert np.dtype(np.ulong) is np.dtype("ulong") + assert np.dtype(np.ulong).itemsize == np.dtype(np.long).itemsize + + +@pytest.mark.filterwarnings("ignore:.*maximum_sctype.*:DeprecationWarning") +class TestMaximumSctype: + + # note that parametrizing with sctype['int'] and similar would skip types + # with the same size (gh-11923) + + @pytest.mark.parametrize( + 't', [np.byte, np.short, np.intc, np.long, np.longlong] + ) + def test_int(self, t): + assert_equal(maximum_sctype(t), np._core.sctypes['int'][-1]) + + @pytest.mark.parametrize( + 't', [np.ubyte, np.ushort, np.uintc, np.ulong, np.ulonglong] + ) + def test_uint(self, t): + assert_equal(maximum_sctype(t), np._core.sctypes['uint'][-1]) + + @pytest.mark.parametrize('t', [np.half, np.single, np.double, np.longdouble]) + def test_float(self, t): + assert_equal(maximum_sctype(t), np._core.sctypes['float'][-1]) + + @pytest.mark.parametrize('t', [np.csingle, np.cdouble, np.clongdouble]) + def test_complex(self, t): + assert_equal(maximum_sctype(t), np._core.sctypes['complex'][-1]) + + @pytest.mark.parametrize('t', [np.bool, np.object_, np.str_, np.bytes_, + np.void]) + def test_other(self, t): + assert_equal(maximum_sctype(t), t) + + +class Test_sctype2char: + # This function is old enough that we're really just documenting the quirks + # at this point. + + def test_scalar_type(self): + assert_equal(sctype2char(np.double), 'd') + assert_equal(sctype2char(np.long), 'l') + assert_equal(sctype2char(np.int_), np.array(0).dtype.char) + assert_equal(sctype2char(np.str_), 'U') + assert_equal(sctype2char(np.bytes_), 'S') + + def test_other_type(self): + assert_equal(sctype2char(float), 'd') + assert_equal(sctype2char(list), 'O') + assert_equal(sctype2char(np.ndarray), 'O') + + def test_third_party_scalar_type(self): + from numpy._core._rational_tests import rational + assert_raises(KeyError, sctype2char, rational) + assert_raises(KeyError, sctype2char, rational(1)) + + def test_array_instance(self): + assert_equal(sctype2char(np.array([1.0, 2.0])), 'd') + + def test_abstract_type(self): + assert_raises(KeyError, sctype2char, np.floating) + + def test_non_type(self): + assert_raises(ValueError, sctype2char, 1) + +@pytest.mark.parametrize("rep, expected", [ + (np.int32, True), + (list, False), + (1.1, False), + (str, True), + (np.dtype(np.float64), True), + (np.dtype((np.int16, (3, 4))), True), + (np.dtype([('a', np.int8)]), True), + ]) +def test_issctype(rep, expected): + # ensure proper identification of scalar + # data-types by issctype() + actual = issctype(rep) + assert type(actual) is bool + assert_equal(actual, expected) + + +@pytest.mark.skipif(sys.flags.optimize > 1, + reason="no docstrings present to inspect when PYTHONOPTIMIZE/Py_OptimizeFlag > 1") +@pytest.mark.xfail(IS_PYPY, + reason="PyPy cannot modify tp_doc after PyType_Ready") +class TestDocStrings: + def test_platform_dependent_aliases(self): + if np.int64 is np.int_: + assert_('int64' in np.int_.__doc__) + elif np.int64 is np.longlong: + assert_('int64' in np.longlong.__doc__) + + +class TestScalarTypeNames: + # gh-9799 + + numeric_types = [ + np.byte, np.short, np.intc, np.long, np.longlong, + np.ubyte, np.ushort, np.uintc, np.ulong, np.ulonglong, + np.half, np.single, np.double, np.longdouble, + np.csingle, np.cdouble, np.clongdouble, + ] + + def test_names_are_unique(self): + # none of the above may be aliases for each other + assert len(set(self.numeric_types)) == len(self.numeric_types) + + # names must be unique + names = [t.__name__ for t in self.numeric_types] + assert len(set(names)) == len(names) + + @pytest.mark.parametrize('t', numeric_types) + def test_names_reflect_attributes(self, t): + """ Test that names correspond to where the type is under ``np.`` """ + assert getattr(np, t.__name__) is t + + @pytest.mark.parametrize('t', numeric_types) + def test_names_are_undersood_by_dtype(self, t): + """ Test the dtype constructor maps names back to the type """ + assert np.dtype(t.__name__).type is t + + +class TestBoolDefinition: + def test_bool_definition(self): + assert nt.bool is np.bool diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_overrides.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_overrides.py new file mode 100644 index 0000000000000000000000000000000000000000..f95536c17686d10d539698919e84f680267fa03e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_overrides.py @@ -0,0 +1,797 @@ +import inspect +import sys +import os +import tempfile +from io import StringIO +from unittest import mock +import pickle + +import pytest + +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_raises_regex) +from numpy.testing.overrides import get_overridable_numpy_array_functions +from numpy._core.overrides import ( + _get_implementing_args, array_function_dispatch, + verify_matching_signatures) + +def _return_not_implemented(self, *args, **kwargs): + return NotImplemented + + +# need to define this at the top level to test pickling +@array_function_dispatch(lambda array: (array,)) +def dispatched_one_arg(array): + """Docstring.""" + return 'original' + + +@array_function_dispatch(lambda array1, array2: (array1, array2)) +def dispatched_two_arg(array1, array2): + """Docstring.""" + return 'original' + + +class TestGetImplementingArgs: + + def test_ndarray(self): + array = np.array(1) + + args = _get_implementing_args([array]) + assert_equal(list(args), [array]) + + args = _get_implementing_args([array, array]) + assert_equal(list(args), [array]) + + args = _get_implementing_args([array, 1]) + assert_equal(list(args), [array]) + + args = _get_implementing_args([1, array]) + assert_equal(list(args), [array]) + + def test_ndarray_subclasses(self): + + class OverrideSub(np.ndarray): + __array_function__ = _return_not_implemented + + class NoOverrideSub(np.ndarray): + pass + + array = np.array(1).view(np.ndarray) + override_sub = np.array(1).view(OverrideSub) + no_override_sub = np.array(1).view(NoOverrideSub) + + args = _get_implementing_args([array, override_sub]) + assert_equal(list(args), [override_sub, array]) + + args = _get_implementing_args([array, no_override_sub]) + assert_equal(list(args), [no_override_sub, array]) + + args = _get_implementing_args( + [override_sub, no_override_sub]) + assert_equal(list(args), [override_sub, no_override_sub]) + + def test_ndarray_and_duck_array(self): + + class Other: + __array_function__ = _return_not_implemented + + array = np.array(1) + other = Other() + + args = _get_implementing_args([other, array]) + assert_equal(list(args), [other, array]) + + args = _get_implementing_args([array, other]) + assert_equal(list(args), [array, other]) + + def test_ndarray_subclass_and_duck_array(self): + + class OverrideSub(np.ndarray): + __array_function__ = _return_not_implemented + + class Other: + __array_function__ = _return_not_implemented + + array = np.array(1) + subarray = np.array(1).view(OverrideSub) + other = Other() + + assert_equal(_get_implementing_args([array, subarray, other]), + [subarray, array, other]) + assert_equal(_get_implementing_args([array, other, subarray]), + [subarray, array, other]) + + def test_many_duck_arrays(self): + + class A: + __array_function__ = _return_not_implemented + + class B(A): + __array_function__ = _return_not_implemented + + class C(A): + __array_function__ = _return_not_implemented + + class D: + __array_function__ = _return_not_implemented + + a = A() + b = B() + c = C() + d = D() + + assert_equal(_get_implementing_args([1]), []) + assert_equal(_get_implementing_args([a]), [a]) + assert_equal(_get_implementing_args([a, 1]), [a]) + assert_equal(_get_implementing_args([a, a, a]), [a]) + assert_equal(_get_implementing_args([a, d, a]), [a, d]) + assert_equal(_get_implementing_args([a, b]), [b, a]) + assert_equal(_get_implementing_args([b, a]), [b, a]) + assert_equal(_get_implementing_args([a, b, c]), [b, c, a]) + assert_equal(_get_implementing_args([a, c, b]), [c, b, a]) + + def test_too_many_duck_arrays(self): + namespace = dict(__array_function__=_return_not_implemented) + types = [type('A' + str(i), (object,), namespace) for i in range(65)] + relevant_args = [t() for t in types] + + actual = _get_implementing_args(relevant_args[:64]) + assert_equal(actual, relevant_args[:64]) + + with assert_raises_regex(TypeError, 'distinct argument types'): + _get_implementing_args(relevant_args) + + +class TestNDArrayArrayFunction: + + def test_method(self): + + class Other: + __array_function__ = _return_not_implemented + + class NoOverrideSub(np.ndarray): + pass + + class OverrideSub(np.ndarray): + __array_function__ = _return_not_implemented + + array = np.array([1]) + other = Other() + no_override_sub = array.view(NoOverrideSub) + override_sub = array.view(OverrideSub) + + result = array.__array_function__(func=dispatched_two_arg, + types=(np.ndarray,), + args=(array, 1.), kwargs={}) + assert_equal(result, 'original') + + result = array.__array_function__(func=dispatched_two_arg, + types=(np.ndarray, Other), + args=(array, other), kwargs={}) + assert_(result is NotImplemented) + + result = array.__array_function__(func=dispatched_two_arg, + types=(np.ndarray, NoOverrideSub), + args=(array, no_override_sub), + kwargs={}) + assert_equal(result, 'original') + + result = array.__array_function__(func=dispatched_two_arg, + types=(np.ndarray, OverrideSub), + args=(array, override_sub), + kwargs={}) + assert_equal(result, 'original') + + with assert_raises_regex(TypeError, 'no implementation found'): + np.concatenate((array, other)) + + expected = np.concatenate((array, array)) + result = np.concatenate((array, no_override_sub)) + assert_equal(result, expected.view(NoOverrideSub)) + result = np.concatenate((array, override_sub)) + assert_equal(result, expected.view(OverrideSub)) + + def test_no_wrapper(self): + # Regular numpy functions have wrappers, but do not presume + # all functions do (array creation ones do not): check that + # we just call the function in that case. + array = np.array(1) + func = lambda x: x * 2 + result = array.__array_function__(func=func, types=(np.ndarray,), + args=(array,), kwargs={}) + assert_equal(result, array * 2) + + def test_wrong_arguments(self): + # Check our implementation guards against wrong arguments. + a = np.array([1, 2]) + with pytest.raises(TypeError, match="args must be a tuple"): + a.__array_function__(np.reshape, (np.ndarray,), a, (2, 1)) + with pytest.raises(TypeError, match="kwargs must be a dict"): + a.__array_function__(np.reshape, (np.ndarray,), (a,), (2, 1)) + + def test_wrong_arguments(self): + # Check our implementation guards against wrong arguments. + a = np.array([1, 2]) + with pytest.raises(TypeError, match="args must be a tuple"): + a.__array_function__(np.reshape, (np.ndarray,), a, (2, 1)) + with pytest.raises(TypeError, match="kwargs must be a dict"): + a.__array_function__(np.reshape, (np.ndarray,), (a,), (2, 1)) + + +class TestArrayFunctionDispatch: + + def test_pickle(self): + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + roundtripped = pickle.loads( + pickle.dumps(dispatched_one_arg, protocol=proto)) + assert_(roundtripped is dispatched_one_arg) + + def test_name_and_docstring(self): + assert_equal(dispatched_one_arg.__name__, 'dispatched_one_arg') + if sys.flags.optimize < 2: + assert_equal(dispatched_one_arg.__doc__, 'Docstring.') + + def test_interface(self): + + class MyArray: + def __array_function__(self, func, types, args, kwargs): + return (self, func, types, args, kwargs) + + original = MyArray() + (obj, func, types, args, kwargs) = dispatched_one_arg(original) + assert_(obj is original) + assert_(func is dispatched_one_arg) + assert_equal(set(types), {MyArray}) + # assert_equal uses the overloaded np.iscomplexobj() internally + assert_(args == (original,)) + assert_equal(kwargs, {}) + + def test_not_implemented(self): + + class MyArray: + def __array_function__(self, func, types, args, kwargs): + return NotImplemented + + array = MyArray() + with assert_raises_regex(TypeError, 'no implementation found'): + dispatched_one_arg(array) + + def test_where_dispatch(self): + + class DuckArray: + def __array_function__(self, ufunc, method, *inputs, **kwargs): + return "overridden" + + array = np.array(1) + duck_array = DuckArray() + + result = np.std(array, where=duck_array) + + assert_equal(result, "overridden") + + +class TestVerifyMatchingSignatures: + + def test_verify_matching_signatures(self): + + verify_matching_signatures(lambda x: 0, lambda x: 0) + verify_matching_signatures(lambda x=None: 0, lambda x=None: 0) + verify_matching_signatures(lambda x=1: 0, lambda x=None: 0) + + with assert_raises(RuntimeError): + verify_matching_signatures(lambda a: 0, lambda b: 0) + with assert_raises(RuntimeError): + verify_matching_signatures(lambda x: 0, lambda x=None: 0) + with assert_raises(RuntimeError): + verify_matching_signatures(lambda x=None: 0, lambda y=None: 0) + with assert_raises(RuntimeError): + verify_matching_signatures(lambda x=1: 0, lambda y=1: 0) + + def test_array_function_dispatch(self): + + with assert_raises(RuntimeError): + @array_function_dispatch(lambda x: (x,)) + def f(y): + pass + + # should not raise + @array_function_dispatch(lambda x: (x,), verify=False) + def f(y): + pass + + +def _new_duck_type_and_implements(): + """Create a duck array type and implements functions.""" + HANDLED_FUNCTIONS = {} + + class MyArray: + def __array_function__(self, func, types, args, kwargs): + if func not in HANDLED_FUNCTIONS: + return NotImplemented + if not all(issubclass(t, MyArray) for t in types): + return NotImplemented + return HANDLED_FUNCTIONS[func](*args, **kwargs) + + def implements(numpy_function): + """Register an __array_function__ implementations.""" + def decorator(func): + HANDLED_FUNCTIONS[numpy_function] = func + return func + return decorator + + return (MyArray, implements) + + +class TestArrayFunctionImplementation: + + def test_one_arg(self): + MyArray, implements = _new_duck_type_and_implements() + + @implements(dispatched_one_arg) + def _(array): + return 'myarray' + + assert_equal(dispatched_one_arg(1), 'original') + assert_equal(dispatched_one_arg(MyArray()), 'myarray') + + def test_optional_args(self): + MyArray, implements = _new_duck_type_and_implements() + + @array_function_dispatch(lambda array, option=None: (array,)) + def func_with_option(array, option='default'): + return option + + @implements(func_with_option) + def my_array_func_with_option(array, new_option='myarray'): + return new_option + + # we don't need to implement every option on __array_function__ + # implementations + assert_equal(func_with_option(1), 'default') + assert_equal(func_with_option(1, option='extra'), 'extra') + assert_equal(func_with_option(MyArray()), 'myarray') + with assert_raises(TypeError): + func_with_option(MyArray(), option='extra') + + # but new options on implementations can't be used + result = my_array_func_with_option(MyArray(), new_option='yes') + assert_equal(result, 'yes') + with assert_raises(TypeError): + func_with_option(MyArray(), new_option='no') + + def test_not_implemented(self): + MyArray, implements = _new_duck_type_and_implements() + + @array_function_dispatch(lambda array: (array,), module='my') + def func(array): + return array + + array = np.array(1) + assert_(func(array) is array) + assert_equal(func.__module__, 'my') + + with assert_raises_regex( + TypeError, "no implementation found for 'my.func'"): + func(MyArray()) + + @pytest.mark.parametrize("name", ["concatenate", "mean", "asarray"]) + def test_signature_error_message_simple(self, name): + func = getattr(np, name) + try: + # all of these functions need an argument: + func() + except TypeError as e: + exc = e + + assert exc.args[0].startswith(f"{name}()") + + def test_signature_error_message(self): + # The lambda function will be named "", but the TypeError + # should show the name as "func" + def _dispatcher(): + return () + + @array_function_dispatch(_dispatcher) + def func(): + pass + + try: + func._implementation(bad_arg=3) + except TypeError as e: + expected_exception = e + + try: + func(bad_arg=3) + raise AssertionError("must fail") + except TypeError as exc: + if exc.args[0].startswith("_dispatcher"): + # We replace the qualname currently, but it used `__name__` + # (relevant functions have the same name and qualname anyway) + pytest.skip("Python version is not using __qualname__ for " + "TypeError formatting.") + + assert exc.args == expected_exception.args + + @pytest.mark.parametrize("value", [234, "this func is not replaced"]) + def test_dispatcher_error(self, value): + # If the dispatcher raises an error, we must not attempt to mutate it + error = TypeError(value) + + def dispatcher(): + raise error + + @array_function_dispatch(dispatcher) + def func(): + return 3 + + try: + func() + raise AssertionError("must fail") + except TypeError as exc: + assert exc is error # unmodified exception + + def test_properties(self): + # Check that str and repr are sensible + func = dispatched_two_arg + assert str(func) == str(func._implementation) + repr_no_id = repr(func).split("at ")[0] + repr_no_id_impl = repr(func._implementation).split("at ")[0] + assert repr_no_id == repr_no_id_impl + + @pytest.mark.parametrize("func", [ + lambda x, y: 0, # no like argument + lambda like=None: 0, # not keyword only + lambda *, like=None, a=3: 0, # not last (not that it matters) + ]) + def test_bad_like_sig(self, func): + # We sanity check the signature, and these should fail. + with pytest.raises(RuntimeError): + array_function_dispatch()(func) + + def test_bad_like_passing(self): + # Cover internal sanity check for passing like as first positional arg + def func(*, like=None): + pass + + func_with_like = array_function_dispatch()(func) + with pytest.raises(TypeError): + func_with_like() + with pytest.raises(TypeError): + func_with_like(like=234) + + def test_too_many_args(self): + # Mainly a unit-test to increase coverage + objs = [] + for i in range(80): + class MyArr: + def __array_function__(self, *args, **kwargs): + return NotImplemented + + objs.append(MyArr()) + + def _dispatch(*args): + return args + + @array_function_dispatch(_dispatch) + def func(*args): + pass + + with pytest.raises(TypeError, match="maximum number"): + func(*objs) + + + +class TestNDArrayMethods: + + def test_repr(self): + # gh-12162: should still be defined even if __array_function__ doesn't + # implement np.array_repr() + + class MyArray(np.ndarray): + def __array_function__(*args, **kwargs): + return NotImplemented + + array = np.array(1).view(MyArray) + assert_equal(repr(array), 'MyArray(1)') + assert_equal(str(array), '1') + + +class TestNumPyFunctions: + + def test_set_module(self): + assert_equal(np.sum.__module__, 'numpy') + assert_equal(np.char.equal.__module__, 'numpy.char') + assert_equal(np.fft.fft.__module__, 'numpy.fft') + assert_equal(np.linalg.solve.__module__, 'numpy.linalg') + + def test_inspect_sum(self): + signature = inspect.signature(np.sum) + assert_('axis' in signature.parameters) + + def test_override_sum(self): + MyArray, implements = _new_duck_type_and_implements() + + @implements(np.sum) + def _(array): + return 'yes' + + assert_equal(np.sum(MyArray()), 'yes') + + def test_sum_on_mock_array(self): + + # We need a proxy for mocks because __array_function__ is only looked + # up in the class dict + class ArrayProxy: + def __init__(self, value): + self.value = value + def __array_function__(self, *args, **kwargs): + return self.value.__array_function__(*args, **kwargs) + def __array__(self, *args, **kwargs): + return self.value.__array__(*args, **kwargs) + + proxy = ArrayProxy(mock.Mock(spec=ArrayProxy)) + proxy.value.__array_function__.return_value = 1 + result = np.sum(proxy) + assert_equal(result, 1) + proxy.value.__array_function__.assert_called_once_with( + np.sum, (ArrayProxy,), (proxy,), {}) + proxy.value.__array__.assert_not_called() + + def test_sum_forwarding_implementation(self): + + class MyArray(np.ndarray): + + def sum(self, axis, out): + return 'summed' + + def __array_function__(self, func, types, args, kwargs): + return super().__array_function__(func, types, args, kwargs) + + # note: the internal implementation of np.sum() calls the .sum() method + array = np.array(1).view(MyArray) + assert_equal(np.sum(array), 'summed') + + +class TestArrayLike: + def setup_method(self): + class MyArray: + def __init__(self, function=None): + self.function = function + + def __array_function__(self, func, types, args, kwargs): + assert func is getattr(np, func.__name__) + try: + my_func = getattr(self, func.__name__) + except AttributeError: + return NotImplemented + return my_func(*args, **kwargs) + + self.MyArray = MyArray + + class MyNoArrayFunctionArray: + def __init__(self, function=None): + self.function = function + + self.MyNoArrayFunctionArray = MyNoArrayFunctionArray + + class MySubclass(np.ndarray): + def __array_function__(self, func, types, args, kwargs): + result = super().__array_function__(func, types, args, kwargs) + return result.view(self.__class__) + + self.MySubclass = MySubclass + + def add_method(self, name, arr_class, enable_value_error=False): + def _definition(*args, **kwargs): + # Check that `like=` isn't propagated downstream + assert 'like' not in kwargs + + if enable_value_error and 'value_error' in kwargs: + raise ValueError + + return arr_class(getattr(arr_class, name)) + setattr(arr_class, name, _definition) + + def func_args(*args, **kwargs): + return args, kwargs + + def test_array_like_not_implemented(self): + self.add_method('array', self.MyArray) + + ref = self.MyArray.array() + + with assert_raises_regex(TypeError, 'no implementation found'): + array_like = np.asarray(1, like=ref) + + _array_tests = [ + ('array', *func_args((1,))), + ('asarray', *func_args((1,))), + ('asanyarray', *func_args((1,))), + ('ascontiguousarray', *func_args((2, 3))), + ('asfortranarray', *func_args((2, 3))), + ('require', *func_args((np.arange(6).reshape(2, 3),), + requirements=['A', 'F'])), + ('empty', *func_args((1,))), + ('full', *func_args((1,), 2)), + ('ones', *func_args((1,))), + ('zeros', *func_args((1,))), + ('arange', *func_args(3)), + ('frombuffer', *func_args(b'\x00' * 8, dtype=int)), + ('fromiter', *func_args(range(3), dtype=int)), + ('fromstring', *func_args('1,2', dtype=int, sep=',')), + ('loadtxt', *func_args(lambda: StringIO('0 1\n2 3'))), + ('genfromtxt', *func_args(lambda: StringIO('1,2.1'), + dtype=[('int', 'i8'), ('float', 'f8')], + delimiter=',')), + ] + + + def test_nep35_functions_as_array_functions(self,): + all_array_functions = get_overridable_numpy_array_functions() + like_array_functions_subset = { + getattr(np, func_name) for func_name, *_ in self.__class__._array_tests + } + assert like_array_functions_subset.issubset(all_array_functions) + + nep35_python_functions = { + np.eye, np.fromfunction, np.full, np.genfromtxt, + np.identity, np.loadtxt, np.ones, np.require, np.tri, + } + assert nep35_python_functions.issubset(all_array_functions) + + nep35_C_functions = { + np.arange, np.array, np.asanyarray, np.asarray, + np.ascontiguousarray, np.asfortranarray, np.empty, + np.frombuffer, np.fromfile, np.fromiter, np.fromstring, + np.zeros, + } + assert nep35_C_functions.issubset(all_array_functions) + + @pytest.mark.parametrize('function, args, kwargs', _array_tests) + @pytest.mark.parametrize('numpy_ref', [True, False]) + def test_array_like(self, function, args, kwargs, numpy_ref): + self.add_method('array', self.MyArray) + self.add_method(function, self.MyArray) + np_func = getattr(np, function) + my_func = getattr(self.MyArray, function) + + if numpy_ref is True: + ref = np.array(1) + else: + ref = self.MyArray.array() + + like_args = tuple(a() if callable(a) else a for a in args) + array_like = np_func(*like_args, **kwargs, like=ref) + + if numpy_ref is True: + assert type(array_like) is np.ndarray + + np_args = tuple(a() if callable(a) else a for a in args) + np_arr = np_func(*np_args, **kwargs) + + # Special-case np.empty to ensure values match + if function == "empty": + np_arr.fill(1) + array_like.fill(1) + + assert_equal(array_like, np_arr) + else: + assert type(array_like) is self.MyArray + assert array_like.function is my_func + + @pytest.mark.parametrize('function, args, kwargs', _array_tests) + @pytest.mark.parametrize('ref', [1, [1], "MyNoArrayFunctionArray"]) + def test_no_array_function_like(self, function, args, kwargs, ref): + self.add_method('array', self.MyNoArrayFunctionArray) + self.add_method(function, self.MyNoArrayFunctionArray) + np_func = getattr(np, function) + + # Instantiate ref if it's the MyNoArrayFunctionArray class + if ref == "MyNoArrayFunctionArray": + ref = self.MyNoArrayFunctionArray.array() + + like_args = tuple(a() if callable(a) else a for a in args) + + with assert_raises_regex(TypeError, + 'The `like` argument must be an array-like that implements'): + np_func(*like_args, **kwargs, like=ref) + + @pytest.mark.parametrize('function, args, kwargs', _array_tests) + def test_subclass(self, function, args, kwargs): + ref = np.array(1).view(self.MySubclass) + np_func = getattr(np, function) + like_args = tuple(a() if callable(a) else a for a in args) + array_like = np_func(*like_args, **kwargs, like=ref) + assert type(array_like) is self.MySubclass + if np_func is np.empty: + return + np_args = tuple(a() if callable(a) else a for a in args) + np_arr = np_func(*np_args, **kwargs) + assert_equal(array_like.view(np.ndarray), np_arr) + + @pytest.mark.parametrize('numpy_ref', [True, False]) + def test_array_like_fromfile(self, numpy_ref): + self.add_method('array', self.MyArray) + self.add_method("fromfile", self.MyArray) + + if numpy_ref is True: + ref = np.array(1) + else: + ref = self.MyArray.array() + + data = np.random.random(5) + + with tempfile.TemporaryDirectory() as tmpdir: + fname = os.path.join(tmpdir, "testfile") + data.tofile(fname) + + array_like = np.fromfile(fname, like=ref) + if numpy_ref is True: + assert type(array_like) is np.ndarray + np_res = np.fromfile(fname, like=ref) + assert_equal(np_res, data) + assert_equal(array_like, np_res) + else: + assert type(array_like) is self.MyArray + assert array_like.function is self.MyArray.fromfile + + def test_exception_handling(self): + self.add_method('array', self.MyArray, enable_value_error=True) + + ref = self.MyArray.array() + + with assert_raises(TypeError): + # Raises the error about `value_error` being invalid first + np.array(1, value_error=True, like=ref) + + @pytest.mark.parametrize('function, args, kwargs', _array_tests) + def test_like_as_none(self, function, args, kwargs): + self.add_method('array', self.MyArray) + self.add_method(function, self.MyArray) + np_func = getattr(np, function) + + like_args = tuple(a() if callable(a) else a for a in args) + # required for loadtxt and genfromtxt to init w/o error. + like_args_exp = tuple(a() if callable(a) else a for a in args) + + array_like = np_func(*like_args, **kwargs, like=None) + expected = np_func(*like_args_exp, **kwargs) + # Special-case np.empty to ensure values match + if function == "empty": + array_like.fill(1) + expected.fill(1) + assert_equal(array_like, expected) + + +def test_function_like(): + # We provide a `__get__` implementation, make sure it works + assert type(np.mean) is np._core._multiarray_umath._ArrayFunctionDispatcher + + class MyClass: + def __array__(self, dtype=None, copy=None): + # valid argument to mean: + return np.arange(3) + + func1 = staticmethod(np.mean) + func2 = np.mean + func3 = classmethod(np.mean) + + m = MyClass() + assert m.func1([10]) == 10 + assert m.func2() == 1 # mean of the arange + with pytest.raises(TypeError, match="unsupported operand type"): + # Tries to operate on the class + m.func3() + + # Manual binding also works (the above may shortcut): + bound = np.mean.__get__(m, MyClass) + assert bound() == 1 + + bound = np.mean.__get__(None, MyClass) # unbound actually + assert bound([10]) == 10 + + bound = np.mean.__get__(MyClass) # classmethod + with pytest.raises(TypeError, match="unsupported operand type"): + bound() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_print.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_print.py new file mode 100644 index 0000000000000000000000000000000000000000..68cdb15b4021257171b1b6d5a958a0fc948072b3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_print.py @@ -0,0 +1,202 @@ +import sys + +import pytest + +import numpy as np +from numpy.testing import assert_, assert_equal, IS_MUSL +from numpy._core.tests._locales import CommaDecimalPointLocale + + +from io import StringIO + +_REF = {np.inf: 'inf', -np.inf: '-inf', np.nan: 'nan'} + + +@pytest.mark.parametrize('tp', [np.float32, np.double, np.longdouble]) +def test_float_types(tp): + """ Check formatting. + + This is only for the str function, and only for simple types. + The precision of np.float32 and np.longdouble aren't the same as the + python float precision. + + """ + for x in [0, 1, -1, 1e20]: + assert_equal(str(tp(x)), str(float(x)), + err_msg='Failed str formatting for type %s' % tp) + + if tp(1e16).itemsize > 4: + assert_equal(str(tp(1e16)), str(float('1e16')), + err_msg='Failed str formatting for type %s' % tp) + else: + ref = '1e+16' + assert_equal(str(tp(1e16)), ref, + err_msg='Failed str formatting for type %s' % tp) + + +@pytest.mark.parametrize('tp', [np.float32, np.double, np.longdouble]) +def test_nan_inf_float(tp): + """ Check formatting of nan & inf. + + This is only for the str function, and only for simple types. + The precision of np.float32 and np.longdouble aren't the same as the + python float precision. + + """ + for x in [np.inf, -np.inf, np.nan]: + assert_equal(str(tp(x)), _REF[x], + err_msg='Failed str formatting for type %s' % tp) + + +@pytest.mark.parametrize('tp', [np.complex64, np.cdouble, np.clongdouble]) +def test_complex_types(tp): + """Check formatting of complex types. + + This is only for the str function, and only for simple types. + The precision of np.float32 and np.longdouble aren't the same as the + python float precision. + + """ + for x in [0, 1, -1, 1e20]: + assert_equal(str(tp(x)), str(complex(x)), + err_msg='Failed str formatting for type %s' % tp) + assert_equal(str(tp(x*1j)), str(complex(x*1j)), + err_msg='Failed str formatting for type %s' % tp) + assert_equal(str(tp(x + x*1j)), str(complex(x + x*1j)), + err_msg='Failed str formatting for type %s' % tp) + + if tp(1e16).itemsize > 8: + assert_equal(str(tp(1e16)), str(complex(1e16)), + err_msg='Failed str formatting for type %s' % tp) + else: + ref = '(1e+16+0j)' + assert_equal(str(tp(1e16)), ref, + err_msg='Failed str formatting for type %s' % tp) + + +@pytest.mark.parametrize('dtype', [np.complex64, np.cdouble, np.clongdouble]) +def test_complex_inf_nan(dtype): + """Check inf/nan formatting of complex types.""" + TESTS = { + complex(np.inf, 0): "(inf+0j)", + complex(0, np.inf): "infj", + complex(-np.inf, 0): "(-inf+0j)", + complex(0, -np.inf): "-infj", + complex(np.inf, 1): "(inf+1j)", + complex(1, np.inf): "(1+infj)", + complex(-np.inf, 1): "(-inf+1j)", + complex(1, -np.inf): "(1-infj)", + complex(np.nan, 0): "(nan+0j)", + complex(0, np.nan): "nanj", + complex(-np.nan, 0): "(nan+0j)", + complex(0, -np.nan): "nanj", + complex(np.nan, 1): "(nan+1j)", + complex(1, np.nan): "(1+nanj)", + complex(-np.nan, 1): "(nan+1j)", + complex(1, -np.nan): "(1+nanj)", + } + for c, s in TESTS.items(): + assert_equal(str(dtype(c)), s) + + +# print tests +def _test_redirected_print(x, tp, ref=None): + file = StringIO() + file_tp = StringIO() + stdout = sys.stdout + try: + sys.stdout = file_tp + print(tp(x)) + sys.stdout = file + if ref: + print(ref) + else: + print(x) + finally: + sys.stdout = stdout + + assert_equal(file.getvalue(), file_tp.getvalue(), + err_msg='print failed for type%s' % tp) + + +@pytest.mark.parametrize('tp', [np.float32, np.double, np.longdouble]) +def test_float_type_print(tp): + """Check formatting when using print """ + for x in [0, 1, -1, 1e20]: + _test_redirected_print(float(x), tp) + + for x in [np.inf, -np.inf, np.nan]: + _test_redirected_print(float(x), tp, _REF[x]) + + if tp(1e16).itemsize > 4: + _test_redirected_print(float(1e16), tp) + else: + ref = '1e+16' + _test_redirected_print(float(1e16), tp, ref) + + +@pytest.mark.parametrize('tp', [np.complex64, np.cdouble, np.clongdouble]) +def test_complex_type_print(tp): + """Check formatting when using print """ + # We do not create complex with inf/nan directly because the feature is + # missing in python < 2.6 + for x in [0, 1, -1, 1e20]: + _test_redirected_print(complex(x), tp) + + if tp(1e16).itemsize > 8: + _test_redirected_print(complex(1e16), tp) + else: + ref = '(1e+16+0j)' + _test_redirected_print(complex(1e16), tp, ref) + + _test_redirected_print(complex(np.inf, 1), tp, '(inf+1j)') + _test_redirected_print(complex(-np.inf, 1), tp, '(-inf+1j)') + _test_redirected_print(complex(-np.nan, 1), tp, '(nan+1j)') + + +def test_scalar_format(): + """Test the str.format method with NumPy scalar types""" + tests = [('{0}', True, np.bool), + ('{0}', False, np.bool), + ('{0:d}', 130, np.uint8), + ('{0:d}', 50000, np.uint16), + ('{0:d}', 3000000000, np.uint32), + ('{0:d}', 15000000000000000000, np.uint64), + ('{0:d}', -120, np.int8), + ('{0:d}', -30000, np.int16), + ('{0:d}', -2000000000, np.int32), + ('{0:d}', -7000000000000000000, np.int64), + ('{0:g}', 1.5, np.float16), + ('{0:g}', 1.5, np.float32), + ('{0:g}', 1.5, np.float64), + ('{0:g}', 1.5, np.longdouble), + ('{0:g}', 1.5+0.5j, np.complex64), + ('{0:g}', 1.5+0.5j, np.complex128), + ('{0:g}', 1.5+0.5j, np.clongdouble)] + + for (fmat, val, valtype) in tests: + try: + assert_equal(fmat.format(val), fmat.format(valtype(val)), + "failed with val %s, type %s" % (val, valtype)) + except ValueError as e: + assert_(False, + "format raised exception (fmt='%s', val=%s, type=%s, exc='%s')" % + (fmat, repr(val), repr(valtype), str(e))) + + +# +# Locale tests: scalar types formatting should be independent of the locale +# + +class TestCommaDecimalPointLocale(CommaDecimalPointLocale): + + def test_locale_single(self): + assert_equal(str(np.float32(1.2)), str(float(1.2))) + + def test_locale_double(self): + assert_equal(str(np.double(1.2)), str(float(1.2))) + + @pytest.mark.skipif(IS_MUSL, + reason="test flaky on musllinux") + def test_locale_longdouble(self): + assert_equal(str(np.longdouble('1.2')), str(float(1.2))) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_protocols.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_protocols.py new file mode 100644 index 0000000000000000000000000000000000000000..4c2cbfabba576020cca2d2fb9712b75f46ec1ee1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_protocols.py @@ -0,0 +1,45 @@ +import pytest +import warnings +import numpy as np + + +@pytest.mark.filterwarnings("error") +def test_getattr_warning(): + # issue gh-14735: make sure we clear only getattr errors, and let warnings + # through + class Wrapper: + def __init__(self, array): + self.array = array + + def __len__(self): + return len(self.array) + + def __getitem__(self, item): + return type(self)(self.array[item]) + + def __getattr__(self, name): + if name.startswith("__array_"): + warnings.warn("object got converted", UserWarning, stacklevel=1) + + return getattr(self.array, name) + + def __repr__(self): + return "".format(self=self) + + array = Wrapper(np.arange(10)) + with pytest.raises(UserWarning, match="object got converted"): + np.asarray(array) + + +def test_array_called(): + class Wrapper: + val = '0' * 100 + + def __array__(self, dtype=None, copy=None): + return np.array([self.val], dtype=dtype, copy=copy) + + + wrapped = Wrapper() + arr = np.array(wrapped, dtype=str) + assert arr.dtype == 'U100' + assert arr[0] == Wrapper.val diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_records.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_records.py new file mode 100644 index 0000000000000000000000000000000000000000..98a80b8fd363a0414b16f49b4f192b41343a1f9d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_records.py @@ -0,0 +1,540 @@ +import collections.abc +import textwrap +from io import BytesIO +from os import path +from pathlib import Path +import pickle + +import pytest + +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, assert_array_almost_equal, + assert_raises, temppath, + ) + + +class TestFromrecords: + def test_fromrecords(self): + r = np.rec.fromrecords([[456, 'dbe', 1.2], [2, 'de', 1.3]], + names='col1,col2,col3') + assert_equal(r[0].item(), (456, 'dbe', 1.2)) + assert_equal(r['col1'].dtype.kind, 'i') + assert_equal(r['col2'].dtype.kind, 'U') + assert_equal(r['col2'].dtype.itemsize, 12) + assert_equal(r['col3'].dtype.kind, 'f') + + def test_fromrecords_0len(self): + """ Verify fromrecords works with a 0-length input """ + dtype = [('a', float), ('b', float)] + r = np.rec.fromrecords([], dtype=dtype) + assert_equal(r.shape, (0,)) + + def test_fromrecords_2d(self): + data = [ + [(1, 2), (3, 4), (5, 6)], + [(6, 5), (4, 3), (2, 1)] + ] + expected_a = [[1, 3, 5], [6, 4, 2]] + expected_b = [[2, 4, 6], [5, 3, 1]] + + # try with dtype + r1 = np.rec.fromrecords(data, dtype=[('a', int), ('b', int)]) + assert_equal(r1['a'], expected_a) + assert_equal(r1['b'], expected_b) + + # try with names + r2 = np.rec.fromrecords(data, names=['a', 'b']) + assert_equal(r2['a'], expected_a) + assert_equal(r2['b'], expected_b) + + assert_equal(r1, r2) + + def test_method_array(self): + r = np.rec.array( + b'abcdefg' * 100, formats='i2,S3,i4', shape=3, byteorder='big' + ) + assert_equal(r[1].item(), (25444, b'efg', 1633837924)) + + def test_method_array2(self): + r = np.rec.array( + [ + (1, 11, 'a'), (2, 22, 'b'), (3, 33, 'c'), (4, 44, 'd'), + (5, 55, 'ex'), (6, 66, 'f'), (7, 77, 'g') + ], + formats='u1,f4,S1' + ) + assert_equal(r[1].item(), (2, 22.0, b'b')) + + def test_recarray_slices(self): + r = np.rec.array( + [ + (1, 11, 'a'), (2, 22, 'b'), (3, 33, 'c'), (4, 44, 'd'), + (5, 55, 'ex'), (6, 66, 'f'), (7, 77, 'g') + ], + formats='u1,f4,S1' + ) + assert_equal(r[1::2][1].item(), (4, 44.0, b'd')) + + def test_recarray_fromarrays(self): + x1 = np.array([1, 2, 3, 4]) + x2 = np.array(['a', 'dd', 'xyz', '12']) + x3 = np.array([1.1, 2, 3, 4]) + r = np.rec.fromarrays([x1, x2, x3], names='a,b,c') + assert_equal(r[1].item(), (2, 'dd', 2.0)) + x1[1] = 34 + assert_equal(r.a, np.array([1, 2, 3, 4])) + + def test_recarray_fromfile(self): + data_dir = path.join(path.dirname(__file__), 'data') + filename = path.join(data_dir, 'recarray_from_file.fits') + fd = open(filename, 'rb') + fd.seek(2880 * 2) + r1 = np.rec.fromfile(fd, formats='f8,i4,S5', shape=3, byteorder='big') + fd.seek(2880 * 2) + r2 = np.rec.array(fd, formats='f8,i4,S5', shape=3, byteorder='big') + fd.seek(2880 * 2) + bytes_array = BytesIO() + bytes_array.write(fd.read()) + bytes_array.seek(0) + r3 = np.rec.fromfile( + bytes_array, formats='f8,i4,S5', shape=3, byteorder='big' + ) + fd.close() + assert_equal(r1, r2) + assert_equal(r2, r3) + + def test_recarray_from_obj(self): + count = 10 + a = np.zeros(count, dtype='O') + b = np.zeros(count, dtype='f8') + c = np.zeros(count, dtype='f8') + for i in range(len(a)): + a[i] = list(range(1, 10)) + + mine = np.rec.fromarrays([a, b, c], names='date,data1,data2') + for i in range(len(a)): + assert_(mine.date[i] == list(range(1, 10))) + assert_(mine.data1[i] == 0.0) + assert_(mine.data2[i] == 0.0) + + def test_recarray_repr(self): + a = np.array([(1, 0.1), (2, 0.2)], + dtype=[('foo', ' 2) & (a < 6)) + xb = np.where((b > 2) & (b < 6)) + ya = ((a > 2) & (a < 6)) + yb = ((b > 2) & (b < 6)) + assert_array_almost_equal(xa, ya.nonzero()) + assert_array_almost_equal(xb, yb.nonzero()) + assert_(np.all(a[ya] > 0.5)) + assert_(np.all(b[yb] > 0.5)) + + def test_endian_where(self): + # GitHub issue #369 + net = np.zeros(3, dtype='>f4') + net[1] = 0.00458849 + net[2] = 0.605202 + max_net = net.max() + test = np.where(net <= 0., max_net, net) + correct = np.array([ 0.60520202, 0.00458849, 0.60520202]) + assert_array_almost_equal(test, correct) + + def test_endian_recarray(self): + # Ticket #2185 + dt = np.dtype([ + ('head', '>u4'), + ('data', '>u4', 2), + ]) + buf = np.recarray(1, dtype=dt) + buf[0]['head'] = 1 + buf[0]['data'][:] = [1, 1] + + h = buf[0]['head'] + d = buf[0]['data'][0] + buf[0]['head'] = h + buf[0]['data'][0] = d + assert_(buf[0]['head'] == 1) + + def test_mem_dot(self): + # Ticket #106 + x = np.random.randn(0, 1) + y = np.random.randn(10, 1) + # Dummy array to detect bad memory access: + _z = np.ones(10) + _dummy = np.empty((0, 10)) + z = np.lib.stride_tricks.as_strided(_z, _dummy.shape, _dummy.strides) + np.dot(x, np.transpose(y), out=z) + assert_equal(_z, np.ones(10)) + # Do the same for the built-in dot: + np._core.multiarray.dot(x, np.transpose(y), out=z) + assert_equal(_z, np.ones(10)) + + def test_arange_endian(self): + # Ticket #111 + ref = np.arange(10) + x = np.arange(10, dtype=' 1 and x['two'] > 2) + + def test_method_args(self): + # Make sure methods and functions have same default axis + # keyword and arguments + funcs1 = ['argmax', 'argmin', 'sum', 'any', 'all', 'cumsum', + 'cumprod', 'prod', 'std', 'var', 'mean', + 'round', 'min', 'max', 'argsort', 'sort'] + funcs2 = ['compress', 'take', 'repeat'] + + for func in funcs1: + arr = np.random.rand(8, 7) + arr2 = arr.copy() + res1 = getattr(arr, func)() + res2 = getattr(np, func)(arr2) + if res1 is None: + res1 = arr + + if res1.dtype.kind in 'uib': + assert_((res1 == res2).all(), func) + else: + assert_(abs(res1-res2).max() < 1e-8, func) + + for func in funcs2: + arr1 = np.random.rand(8, 7) + arr2 = np.random.rand(8, 7) + res1 = None + if func == 'compress': + arr1 = arr1.ravel() + res1 = getattr(arr2, func)(arr1) + else: + arr2 = (15*arr2).astype(int).ravel() + if res1 is None: + res1 = getattr(arr1, func)(arr2) + res2 = getattr(np, func)(arr1, arr2) + assert_(abs(res1-res2).max() < 1e-8, func) + + def test_mem_lexsort_strings(self): + # Ticket #298 + lst = ['abc', 'cde', 'fgh'] + np.lexsort((lst,)) + + def test_fancy_index(self): + # Ticket #302 + x = np.array([1, 2])[np.array([0])] + assert_equal(x.shape, (1,)) + + def test_recarray_copy(self): + # Ticket #312 + dt = [('x', np.int16), ('y', np.float64)] + ra = np.array([(1, 2.3)], dtype=dt) + rb = np.rec.array(ra, dtype=dt) + rb['x'] = 2. + assert_(ra['x'] != rb['x']) + + def test_rec_fromarray(self): + # Ticket #322 + x1 = np.array([[1, 2], [3, 4], [5, 6]]) + x2 = np.array(['a', 'dd', 'xyz']) + x3 = np.array([1.1, 2, 3]) + np.rec.fromarrays([x1, x2, x3], formats="(2,)i4,S3,f8") + + def test_object_array_assign(self): + x = np.empty((2, 2), object) + x.flat[2] = (1, 2, 3) + assert_equal(x.flat[2], (1, 2, 3)) + + def test_ndmin_float64(self): + # Ticket #324 + x = np.array([1, 2, 3], dtype=np.float64) + assert_equal(np.array(x, dtype=np.float32, ndmin=2).ndim, 2) + assert_equal(np.array(x, dtype=np.float64, ndmin=2).ndim, 2) + + def test_ndmin_order(self): + # Issue #465 and related checks + assert_(np.array([1, 2], order='C', ndmin=3).flags.c_contiguous) + assert_(np.array([1, 2], order='F', ndmin=3).flags.f_contiguous) + assert_(np.array(np.ones((2, 2), order='F'), ndmin=3).flags.f_contiguous) + assert_(np.array(np.ones((2, 2), order='C'), ndmin=3).flags.c_contiguous) + + def test_mem_axis_minimization(self): + # Ticket #327 + data = np.arange(5) + data = np.add.outer(data, data) + + def test_mem_float_imag(self): + # Ticket #330 + np.float64(1.0).imag + + def test_dtype_tuple(self): + # Ticket #334 + assert_(np.dtype('i4') == np.dtype(('i4', ()))) + + def test_dtype_posttuple(self): + # Ticket #335 + np.dtype([('col1', '()i4')]) + + def test_numeric_carray_compare(self): + # Ticket #341 + assert_equal(np.array(['X'], 'c'), b'X') + + def test_string_array_size(self): + # Ticket #342 + assert_raises(ValueError, + np.array, [['X'], ['X', 'X', 'X']], '|S1') + + def test_dtype_repr(self): + # Ticket #344 + dt1 = np.dtype(('uint32', 2)) + dt2 = np.dtype(('uint32', (2,))) + assert_equal(dt1.__repr__(), dt2.__repr__()) + + def test_reshape_order(self): + # Make sure reshape order works. + a = np.arange(6).reshape(2, 3, order='F') + assert_equal(a, [[0, 2, 4], [1, 3, 5]]) + a = np.array([[1, 2], [3, 4], [5, 6], [7, 8]]) + b = a[:, 1] + assert_equal(b.reshape(2, 2, order='F'), [[2, 6], [4, 8]]) + + def test_reshape_zero_strides(self): + # Issue #380, test reshaping of zero strided arrays + a = np.ones(1) + a = np.lib.stride_tricks.as_strided(a, shape=(5,), strides=(0,)) + assert_(a.reshape(5, 1).strides[0] == 0) + + def test_reshape_zero_size(self): + # GitHub Issue #2700, setting shape failed for 0-sized arrays + a = np.ones((0, 2)) + a.shape = (-1, 2) + + def test_reshape_trailing_ones_strides(self): + # GitHub issue gh-2949, bad strides for trailing ones of new shape + a = np.zeros(12, dtype=np.int32)[::2] # not contiguous + strides_c = (16, 8, 8, 8) + strides_f = (8, 24, 48, 48) + assert_equal(a.reshape(3, 2, 1, 1).strides, strides_c) + assert_equal(a.reshape(3, 2, 1, 1, order='F').strides, strides_f) + assert_equal(np.array(0, dtype=np.int32).reshape(1, 1).strides, (4, 4)) + + def test_repeat_discont(self): + # Ticket #352 + a = np.arange(12).reshape(4, 3)[:, 2] + assert_equal(a.repeat(3), [2, 2, 2, 5, 5, 5, 8, 8, 8, 11, 11, 11]) + + def test_array_index(self): + # Make sure optimization is not called in this case. + a = np.array([1, 2, 3]) + a2 = np.array([[1, 2, 3]]) + assert_equal(a[np.where(a == 3)], a2[np.where(a2 == 3)]) + + def test_object_argmax(self): + a = np.array([1, 2, 3], dtype=object) + assert_(a.argmax() == 2) + + def test_recarray_fields(self): + # Ticket #372 + dt0 = np.dtype([('f0', 'i4'), ('f1', 'i4')]) + dt1 = np.dtype([('f0', 'i8'), ('f1', 'i8')]) + for a in [np.array([(1, 2), (3, 4)], "i4,i4"), + np.rec.array([(1, 2), (3, 4)], "i4,i4"), + np.rec.array([(1, 2), (3, 4)]), + np.rec.fromarrays([(1, 2), (3, 4)], "i4,i4"), + np.rec.fromarrays([(1, 2), (3, 4)])]: + assert_(a.dtype in [dt0, dt1]) + + def test_random_shuffle(self): + # Ticket #374 + a = np.arange(5).reshape((5, 1)) + b = a.copy() + np.random.shuffle(b) + assert_equal(np.sort(b, axis=0), a) + + def test_refcount_vdot(self): + # Changeset #3443 + _assert_valid_refcount(np.vdot) + + def test_startswith(self): + ca = np.char.array(['Hi', 'There']) + assert_equal(ca.startswith('H'), [True, False]) + + def test_noncommutative_reduce_accumulate(self): + # Ticket #413 + tosubtract = np.arange(5) + todivide = np.array([2.0, 0.5, 0.25]) + assert_equal(np.subtract.reduce(tosubtract), -10) + assert_equal(np.divide.reduce(todivide), 16.0) + assert_array_equal(np.subtract.accumulate(tosubtract), + np.array([0, -1, -3, -6, -10])) + assert_array_equal(np.divide.accumulate(todivide), + np.array([2., 4., 16.])) + + def test_convolve_empty(self): + # Convolve should raise an error for empty input array. + assert_raises(ValueError, np.convolve, [], [1]) + assert_raises(ValueError, np.convolve, [1], []) + + def test_multidim_byteswap(self): + # Ticket #449 + r = np.array([(1, (0, 1, 2))], dtype="i2,3i2") + assert_array_equal(r.byteswap(), + np.array([(256, (0, 256, 512))], r.dtype)) + + def test_string_NULL(self): + # Changeset 3557 + assert_equal(np.array("a\x00\x0b\x0c\x00").item(), + 'a\x00\x0b\x0c') + + def test_junk_in_string_fields_of_recarray(self): + # Ticket #483 + r = np.array([[b'abc']], dtype=[('var1', '|S20')]) + assert_(asbytes(r['var1'][0][0]) == b'abc') + + def test_take_output(self): + # Ensure that 'take' honours output parameter. + x = np.arange(12).reshape((3, 4)) + a = np.take(x, [0, 2], axis=1) + b = np.zeros_like(a) + np.take(x, [0, 2], axis=1, out=b) + assert_array_equal(a, b) + + def test_take_object_fail(self): + # Issue gh-3001 + d = 123. + a = np.array([d, 1], dtype=object) + if HAS_REFCOUNT: + ref_d = sys.getrefcount(d) + try: + a.take([0, 100]) + except IndexError: + pass + if HAS_REFCOUNT: + assert_(ref_d == sys.getrefcount(d)) + + def test_array_str_64bit(self): + # Ticket #501 + s = np.array([1, np.nan], dtype=np.float64) + with np.errstate(all='raise'): + np.array_str(s) # Should succeed + + def test_frompyfunc_endian(self): + # Ticket #503 + from math import radians + uradians = np.frompyfunc(radians, 1, 1) + big_endian = np.array([83.4, 83.5], dtype='>f8') + little_endian = np.array([83.4, 83.5], dtype=' object + # casting succeeds + def rs(): + x = np.ones([484, 286]) + y = np.zeros([484, 286]) + x |= y + + assert_raises(TypeError, rs) + + def test_unicode_scalar(self): + # Ticket #600 + x = np.array(["DROND", "DROND1"], dtype="U6") + el = x[1] + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + new = pickle.loads(pickle.dumps(el, protocol=proto)) + assert_equal(new, el) + + def test_arange_non_native_dtype(self): + # Ticket #616 + for T in ('>f4', ' 0)] = v + + assert_raises(IndexError, ia, x, s, np.zeros(9, dtype=float)) + assert_raises(IndexError, ia, x, s, np.zeros(11, dtype=float)) + + # Old special case (different code path): + assert_raises(ValueError, ia, x.flat, s, np.zeros(9, dtype=float)) + assert_raises(ValueError, ia, x.flat, s, np.zeros(11, dtype=float)) + + def test_mem_scalar_indexing(self): + # Ticket #603 + x = np.array([0], dtype=float) + index = np.array(0, dtype=np.int32) + x[index] + + def test_binary_repr_0_width(self): + assert_equal(np.binary_repr(0, width=3), '000') + + def test_fromstring(self): + assert_equal(np.fromstring("12:09:09", dtype=int, sep=":"), + [12, 9, 9]) + + def test_searchsorted_variable_length(self): + x = np.array(['a', 'aa', 'b']) + y = np.array(['d', 'e']) + assert_equal(x.searchsorted(y), [3, 3]) + + def test_string_argsort_with_zeros(self): + # Check argsort for strings containing zeros. + x = np.frombuffer(b"\x00\x02\x00\x01", dtype="|S2") + assert_array_equal(x.argsort(kind='m'), np.array([1, 0])) + assert_array_equal(x.argsort(kind='q'), np.array([1, 0])) + + def test_string_sort_with_zeros(self): + # Check sort for strings containing zeros. + x = np.frombuffer(b"\x00\x02\x00\x01", dtype="|S2") + y = np.frombuffer(b"\x00\x01\x00\x02", dtype="|S2") + assert_array_equal(np.sort(x, kind="q"), y) + + def test_copy_detection_zero_dim(self): + # Ticket #658 + np.indices((0, 3, 4)).T.reshape(-1, 3) + + def test_flat_byteorder(self): + # Ticket #657 + x = np.arange(10) + assert_array_equal(x.astype('>i4'), x.astype('i4').flat[:], x.astype('i4')): + x = np.array([-1, 0, 1], dtype=dt) + assert_equal(x.flat[0].dtype, x[0].dtype) + + def test_copy_detection_corner_case(self): + # Ticket #658 + np.indices((0, 3, 4)).T.reshape(-1, 3) + + def test_object_array_refcounting(self): + # Ticket #633 + if not hasattr(sys, 'getrefcount'): + return + + # NB. this is probably CPython-specific + + cnt = sys.getrefcount + + a = object() + b = object() + c = object() + + cnt0_a = cnt(a) + cnt0_b = cnt(b) + cnt0_c = cnt(c) + + # -- 0d -> 1-d broadcast slice assignment + + arr = np.zeros(5, dtype=np.object_) + + arr[:] = a + assert_equal(cnt(a), cnt0_a + 5) + + arr[:] = b + assert_equal(cnt(a), cnt0_a) + assert_equal(cnt(b), cnt0_b + 5) + + arr[:2] = c + assert_equal(cnt(b), cnt0_b + 3) + assert_equal(cnt(c), cnt0_c + 2) + + del arr + + # -- 1-d -> 2-d broadcast slice assignment + + arr = np.zeros((5, 2), dtype=np.object_) + arr0 = np.zeros(2, dtype=np.object_) + + arr0[0] = a + assert_(cnt(a) == cnt0_a + 1) + arr0[1] = b + assert_(cnt(b) == cnt0_b + 1) + + arr[:, :] = arr0 + assert_(cnt(a) == cnt0_a + 6) + assert_(cnt(b) == cnt0_b + 6) + + arr[:, 0] = None + assert_(cnt(a) == cnt0_a + 1) + + del arr, arr0 + + # -- 2-d copying + flattening + + arr = np.zeros((5, 2), dtype=np.object_) + + arr[:, 0] = a + arr[:, 1] = b + assert_(cnt(a) == cnt0_a + 5) + assert_(cnt(b) == cnt0_b + 5) + + arr2 = arr.copy() + assert_(cnt(a) == cnt0_a + 10) + assert_(cnt(b) == cnt0_b + 10) + + arr2 = arr[:, 0].copy() + assert_(cnt(a) == cnt0_a + 10) + assert_(cnt(b) == cnt0_b + 5) + + arr2 = arr.flatten() + assert_(cnt(a) == cnt0_a + 10) + assert_(cnt(b) == cnt0_b + 10) + + del arr, arr2 + + # -- concatenate, repeat, take, choose + + arr1 = np.zeros((5, 1), dtype=np.object_) + arr2 = np.zeros((5, 1), dtype=np.object_) + + arr1[...] = a + arr2[...] = b + assert_(cnt(a) == cnt0_a + 5) + assert_(cnt(b) == cnt0_b + 5) + + tmp = np.concatenate((arr1, arr2)) + assert_(cnt(a) == cnt0_a + 5 + 5) + assert_(cnt(b) == cnt0_b + 5 + 5) + + tmp = arr1.repeat(3, axis=0) + assert_(cnt(a) == cnt0_a + 5 + 3*5) + + tmp = arr1.take([1, 2, 3], axis=0) + assert_(cnt(a) == cnt0_a + 5 + 3) + + x = np.array([[0], [1], [0], [1], [1]], int) + tmp = x.choose(arr1, arr2) + assert_(cnt(a) == cnt0_a + 5 + 2) + assert_(cnt(b) == cnt0_b + 5 + 3) + + del tmp # Avoid pyflakes unused variable warning + + def test_mem_custom_float_to_array(self): + # Ticket 702 + class MyFloat: + def __float__(self): + return 1.0 + + tmp = np.atleast_1d([MyFloat()]) + tmp.astype(float) # Should succeed + + def test_object_array_refcount_self_assign(self): + # Ticket #711 + class VictimObject: + deleted = False + + def __del__(self): + self.deleted = True + + d = VictimObject() + arr = np.zeros(5, dtype=np.object_) + arr[:] = d + del d + arr[:] = arr # refcount of 'd' might hit zero here + assert_(not arr[0].deleted) + arr[:] = arr # trying to induce a segfault by doing it again... + assert_(not arr[0].deleted) + + def test_mem_fromiter_invalid_dtype_string(self): + x = [1, 2, 3] + assert_raises(ValueError, + np.fromiter, list(x), dtype='S') + + def test_reduce_big_object_array(self): + # Ticket #713 + oldsize = np.setbufsize(10*16) + a = np.array([None]*161, object) + assert_(not np.any(a)) + np.setbufsize(oldsize) + + def test_mem_0d_array_index(self): + # Ticket #714 + np.zeros(10)[np.array(0)] + + def test_nonnative_endian_fill(self): + # Non-native endian arrays were incorrectly filled with scalars + # before r5034. + if sys.byteorder == 'little': + dtype = np.dtype('>i4') + else: + dtype = np.dtype('data contains non-zero floats + x = np.array([123456789e199], dtype=np.float64) + if IS_PYPY: + x.resize((m, 0), refcheck=False) + else: + x.resize((m, 0)) + y = np.array([123456789e199], dtype=np.float64) + if IS_PYPY: + y.resize((0, n), refcheck=False) + else: + y.resize((0, n)) + + # `dot` should just return zero (m, n) matrix + z = np.dot(x, y) + assert_(np.all(z == 0)) + assert_(z.shape == (m, n)) + + def test_zeros(self): + # Regression test for #1061. + # Set a size which cannot fit into a 64 bits signed integer + sz = 2 ** 64 + with assert_raises_regex(ValueError, + 'Maximum allowed dimension exceeded'): + np.empty(sz) + + def test_huge_arange(self): + # Regression test for #1062. + # Set a size which cannot fit into a 64 bits signed integer + sz = 2 ** 64 + with assert_raises_regex(ValueError, + 'Maximum allowed size exceeded'): + np.arange(sz) + assert_(np.size == sz) + + def test_fromiter_bytes(self): + # Ticket #1058 + a = np.fromiter(list(range(10)), dtype='b') + b = np.fromiter(list(range(10)), dtype='B') + assert_(np.all(a == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))) + assert_(np.all(b == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))) + + def test_array_from_sequence_scalar_array(self): + # Ticket #1078: segfaults when creating an array with a sequence of + # 0d arrays. + a = np.array((np.ones(2), np.array(2)), dtype=object) + assert_equal(a.shape, (2,)) + assert_equal(a.dtype, np.dtype(object)) + assert_equal(a[0], np.ones(2)) + assert_equal(a[1], np.array(2)) + + a = np.array(((1,), np.array(1)), dtype=object) + assert_equal(a.shape, (2,)) + assert_equal(a.dtype, np.dtype(object)) + assert_equal(a[0], (1,)) + assert_equal(a[1], np.array(1)) + + def test_array_from_sequence_scalar_array2(self): + # Ticket #1081: weird array with strange input... + t = np.array([np.array([]), np.array(0, object)], dtype=object) + assert_equal(t.shape, (2,)) + assert_equal(t.dtype, np.dtype(object)) + + def test_array_too_big(self): + # Ticket #1080. + assert_raises(ValueError, np.zeros, [975]*7, np.int8) + assert_raises(ValueError, np.zeros, [26244]*5, np.int8) + + def test_dtype_keyerrors_(self): + # Ticket #1106. + dt = np.dtype([('f1', np.uint)]) + assert_raises(KeyError, dt.__getitem__, "f2") + assert_raises(IndexError, dt.__getitem__, 1) + assert_raises(TypeError, dt.__getitem__, 0.0) + + def test_lexsort_buffer_length(self): + # Ticket #1217, don't segfault. + a = np.ones(100, dtype=np.int8) + b = np.ones(100, dtype=np.int32) + i = np.lexsort((a[::-1], b)) + assert_equal(i, np.arange(100, dtype=int)) + + def test_object_array_to_fixed_string(self): + # Ticket #1235. + a = np.array(['abcdefgh', 'ijklmnop'], dtype=np.object_) + b = np.array(a, dtype=(np.str_, 8)) + assert_equal(a, b) + c = np.array(a, dtype=(np.str_, 5)) + assert_equal(c, np.array(['abcde', 'ijklm'])) + d = np.array(a, dtype=(np.str_, 12)) + assert_equal(a, d) + e = np.empty((2, ), dtype=(np.str_, 8)) + e[:] = a[:] + assert_equal(a, e) + + def test_unicode_to_string_cast(self): + # Ticket #1240. + a = np.array([['abc', '\u03a3'], + ['asdf', 'erw']], + dtype='U') + assert_raises(UnicodeEncodeError, np.array, a, 'S4') + + def test_unicode_to_string_cast_error(self): + # gh-15790 + a = np.array(['\x80'] * 129, dtype='U3') + assert_raises(UnicodeEncodeError, np.array, a, 'S') + b = a.reshape(3, 43)[:-1, :-1] + assert_raises(UnicodeEncodeError, np.array, b, 'S') + + def test_mixed_string_byte_array_creation(self): + a = np.array(['1234', b'123']) + assert_(a.itemsize == 16) + a = np.array([b'123', '1234']) + assert_(a.itemsize == 16) + a = np.array(['1234', b'123', '12345']) + assert_(a.itemsize == 20) + a = np.array([b'123', '1234', b'12345']) + assert_(a.itemsize == 20) + a = np.array([b'123', '1234', b'1234']) + assert_(a.itemsize == 16) + + def test_misaligned_objects_segfault(self): + # Ticket #1198 and #1267 + a1 = np.zeros((10,), dtype='O,c') + a2 = np.array(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'], 'S10') + a1['f0'] = a2 + repr(a1) + np.argmax(a1['f0']) + a1['f0'][1] = "FOO" + a1['f0'] = "FOO" + np.array(a1['f0'], dtype='S') + np.nonzero(a1['f0']) + a1.sort() + copy.deepcopy(a1) + + def test_misaligned_scalars_segfault(self): + # Ticket #1267 + s1 = np.array(('a', 'Foo'), dtype='c,O') + s2 = np.array(('b', 'Bar'), dtype='c,O') + s1['f1'] = s2['f1'] + s1['f1'] = 'Baz' + + def test_misaligned_dot_product_objects(self): + # Ticket #1267 + # This didn't require a fix, but it's worth testing anyway, because + # it may fail if .dot stops enforcing the arrays to be BEHAVED + a = np.array([[(1, 'a'), (0, 'a')], [(0, 'a'), (1, 'a')]], dtype='O,c') + b = np.array([[(4, 'a'), (1, 'a')], [(2, 'a'), (2, 'a')]], dtype='O,c') + np.dot(a['f0'], b['f0']) + + def test_byteswap_complex_scalar(self): + # Ticket #1259 and gh-441 + for dtype in [np.dtype('<'+t) for t in np.typecodes['Complex']]: + z = np.array([2.2-1.1j], dtype) + x = z[0] # always native-endian + y = x.byteswap() + if x.dtype.byteorder == z.dtype.byteorder: + # little-endian machine + assert_equal(x, np.frombuffer(y.tobytes(), dtype=dtype.newbyteorder())) + else: + # big-endian machine + assert_equal(x, np.frombuffer(y.tobytes(), dtype=dtype)) + # double check real and imaginary parts: + assert_equal(x.real, y.real.byteswap()) + assert_equal(x.imag, y.imag.byteswap()) + + def test_structured_arrays_with_objects1(self): + # Ticket #1299 + stra = 'aaaa' + strb = 'bbbb' + x = np.array([[(0, stra), (1, strb)]], 'i8,O') + x[x.nonzero()] = x.ravel()[:1] + assert_(x[0, 1] == x[0, 0]) + + @pytest.mark.skipif( + sys.version_info >= (3, 12), + reason="Python 3.12 has immortal refcounts, this test no longer works." + ) + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_structured_arrays_with_objects2(self): + # Ticket #1299 second test + stra = 'aaaa' + strb = 'bbbb' + numb = sys.getrefcount(strb) + numa = sys.getrefcount(stra) + x = np.array([[(0, stra), (1, strb)]], 'i8,O') + x[x.nonzero()] = x.ravel()[:1] + assert_(sys.getrefcount(strb) == numb) + assert_(sys.getrefcount(stra) == numa + 2) + + def test_duplicate_title_and_name(self): + # Ticket #1254 + dtspec = [(('a', 'a'), 'i'), ('b', 'i')] + assert_raises(ValueError, np.dtype, dtspec) + + def test_signed_integer_division_overflow(self): + # Ticket #1317. + def test_type(t): + min = np.array([np.iinfo(t).min]) + min //= -1 + + with np.errstate(over="ignore"): + for t in (np.int8, np.int16, np.int32, np.int64, int): + test_type(t) + + def test_buffer_hashlib(self): + from hashlib import sha256 + + x = np.array([1, 2, 3], dtype=np.dtype('c') + + def test_log1p_compiler_shenanigans(self): + # Check if log1p is behaving on 32 bit intel systems. + assert_(np.isfinite(np.log1p(np.exp2(-53)))) + + def test_fromiter_comparison(self): + a = np.fromiter(list(range(10)), dtype='b') + b = np.fromiter(list(range(10)), dtype='B') + assert_(np.all(a == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))) + assert_(np.all(b == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))) + + def test_fromstring_crash(self): + # Ticket #1345: the following should not cause a crash + with assert_warns(DeprecationWarning): + np.fromstring(b'aa, aa, 1.0', sep=',') + + def test_ticket_1539(self): + dtypes = [x for x in np._core.sctypeDict.values() + if (issubclass(x, np.number) + and not issubclass(x, np.timedelta64))] + a = np.array([], np.bool) # not x[0] because it is unordered + failures = [] + + for x in dtypes: + b = a.astype(x) + for y in dtypes: + c = a.astype(y) + try: + d = np.dot(b, c) + except TypeError: + failures.append((x, y)) + else: + if d != 0: + failures.append((x, y)) + if failures: + raise AssertionError("Failures: %r" % failures) + + def test_ticket_1538(self): + x = np.finfo(np.float32) + for name in 'eps epsneg max min resolution tiny'.split(): + assert_equal(type(getattr(x, name)), np.float32, + err_msg=name) + + def test_ticket_1434(self): + # Check that the out= argument in var and std has an effect + data = np.array(((1, 2, 3), (4, 5, 6), (7, 8, 9))) + out = np.zeros((3,)) + + ret = data.var(axis=1, out=out) + assert_(ret is out) + assert_array_equal(ret, data.var(axis=1)) + + ret = data.std(axis=1, out=out) + assert_(ret is out) + assert_array_equal(ret, data.std(axis=1)) + + def test_complex_nan_maximum(self): + cnan = complex(0, np.nan) + assert_equal(np.maximum(1, cnan), cnan) + + def test_subclass_int_tuple_assignment(self): + # ticket #1563 + class Subclass(np.ndarray): + def __new__(cls, i): + return np.ones((i,)).view(cls) + + x = Subclass(5) + x[(0,)] = 2 # shouldn't raise an exception + assert_equal(x[0], 2) + + def test_ufunc_no_unnecessary_views(self): + # ticket #1548 + class Subclass(np.ndarray): + pass + x = np.array([1, 2, 3]).view(Subclass) + y = np.add(x, x, x) + assert_equal(id(x), id(y)) + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_take_refcount(self): + # ticket #939 + a = np.arange(16, dtype=float) + a.shape = (4, 4) + lut = np.ones((5 + 3, 4), float) + rgba = np.empty(shape=a.shape + (4,), dtype=lut.dtype) + c1 = sys.getrefcount(rgba) + try: + lut.take(a, axis=0, mode='clip', out=rgba) + except TypeError: + pass + c2 = sys.getrefcount(rgba) + assert_equal(c1, c2) + + def test_fromfile_tofile_seeks(self): + # On Python 3, tofile/fromfile used to get (#1610) the Python + # file handle out of sync + f0 = tempfile.NamedTemporaryFile() + f = f0.file + f.write(np.arange(255, dtype='u1').tobytes()) + + f.seek(20) + ret = np.fromfile(f, count=4, dtype='u1') + assert_equal(ret, np.array([20, 21, 22, 23], dtype='u1')) + assert_equal(f.tell(), 24) + + f.seek(40) + np.array([1, 2, 3], dtype='u1').tofile(f) + assert_equal(f.tell(), 43) + + f.seek(40) + data = f.read(3) + assert_equal(data, b"\x01\x02\x03") + + f.seek(80) + f.read(4) + data = np.fromfile(f, dtype='u1', count=4) + assert_equal(data, np.array([84, 85, 86, 87], dtype='u1')) + + f.close() + + def test_complex_scalar_warning(self): + for tp in [np.csingle, np.cdouble, np.clongdouble]: + x = tp(1+2j) + assert_warns(ComplexWarning, float, x) + with suppress_warnings() as sup: + sup.filter(ComplexWarning) + assert_equal(float(x), float(x.real)) + + def test_complex_scalar_complex_cast(self): + for tp in [np.csingle, np.cdouble, np.clongdouble]: + x = tp(1+2j) + assert_equal(complex(x), 1+2j) + + def test_complex_boolean_cast(self): + # Ticket #2218 + for tp in [np.csingle, np.cdouble, np.clongdouble]: + x = np.array([0, 0+0.5j, 0.5+0j], dtype=tp) + assert_equal(x.astype(bool), np.array([0, 1, 1], dtype=bool)) + assert_(np.any(x)) + assert_(np.all(x[1:])) + + def test_uint_int_conversion(self): + x = 2**64 - 1 + assert_equal(int(np.uint64(x)), x) + + def test_duplicate_field_names_assign(self): + ra = np.fromiter(((i*3, i*2) for i in range(10)), dtype='i8,f8') + ra.dtype.names = ('f1', 'f2') + repr(ra) # should not cause a segmentation fault + assert_raises(ValueError, setattr, ra.dtype, 'names', ('f1', 'f1')) + + def test_eq_string_and_object_array(self): + # From e-mail thread "__eq__ with str and object" (Keith Goodman) + a1 = np.array(['a', 'b'], dtype=object) + a2 = np.array(['a', 'c']) + assert_array_equal(a1 == a2, [True, False]) + assert_array_equal(a2 == a1, [True, False]) + + def test_nonzero_byteswap(self): + a = np.array([0x80000000, 0x00000080, 0], dtype=np.uint32) + a.dtype = np.float32 + assert_equal(a.nonzero()[0], [1]) + a = a.byteswap() + a = a.view(a.dtype.newbyteorder()) + assert_equal(a.nonzero()[0], [1]) # [0] if nonzero() ignores swap + + def test_empty_mul(self): + a = np.array([1.]) + a[1:1] *= 2 + assert_equal(a, [1.]) + + def test_array_side_effect(self): + # The second use of itemsize was throwing an exception because in + # ctors.c, discover_itemsize was calling PyObject_Length without + # checking the return code. This failed to get the length of the + # number 2, and the exception hung around until something checked + # PyErr_Occurred() and returned an error. + assert_equal(np.dtype('S10').itemsize, 10) + np.array([['abc', 2], ['long ', '0123456789']], dtype=np.bytes_) + assert_equal(np.dtype('S10').itemsize, 10) + + def test_any_float(self): + # all and any for floats + a = np.array([0.1, 0.9]) + assert_(np.any(a)) + assert_(np.all(a)) + + def test_large_float_sum(self): + a = np.arange(10000, dtype='f') + assert_equal(a.sum(dtype='d'), a.astype('d').sum()) + + def test_ufunc_casting_out(self): + a = np.array(1.0, dtype=np.float32) + b = np.array(1.0, dtype=np.float64) + c = np.array(1.0, dtype=np.float32) + np.add(a, b, out=c) + assert_equal(c, 2.0) + + def test_array_scalar_contiguous(self): + # Array scalars are both C and Fortran contiguous + assert_(np.array(1.0).flags.c_contiguous) + assert_(np.array(1.0).flags.f_contiguous) + assert_(np.array(np.float32(1.0)).flags.c_contiguous) + assert_(np.array(np.float32(1.0)).flags.f_contiguous) + + def test_squeeze_contiguous(self): + # Similar to GitHub issue #387 + a = np.zeros((1, 2)).squeeze() + b = np.zeros((2, 2, 2), order='F')[:, :, ::2].squeeze() + assert_(a.flags.c_contiguous) + assert_(a.flags.f_contiguous) + assert_(b.flags.f_contiguous) + + def test_squeeze_axis_handling(self): + # Issue #10779 + # Ensure proper handling of objects + # that don't support axis specification + # when squeezing + + class OldSqueeze(np.ndarray): + + def __new__(cls, + input_array): + obj = np.asarray(input_array).view(cls) + return obj + + # it is perfectly reasonable that prior + # to numpy version 1.7.0 a subclass of ndarray + # might have been created that did not expect + # squeeze to have an axis argument + # NOTE: this example is somewhat artificial; + # it is designed to simulate an old API + # expectation to guard against regression + def squeeze(self): + return super().squeeze() + + oldsqueeze = OldSqueeze(np.array([[1],[2],[3]])) + + # if no axis argument is specified the old API + # expectation should give the correct result + assert_equal(np.squeeze(oldsqueeze), + np.array([1,2,3])) + + # likewise, axis=None should work perfectly well + # with the old API expectation + assert_equal(np.squeeze(oldsqueeze, axis=None), + np.array([1,2,3])) + + # however, specification of any particular axis + # should raise a TypeError in the context of the + # old API specification, even when using a valid + # axis specification like 1 for this array + with assert_raises(TypeError): + # this would silently succeed for array + # subclasses / objects that did not support + # squeeze axis argument handling before fixing + # Issue #10779 + np.squeeze(oldsqueeze, axis=1) + + # check for the same behavior when using an invalid + # axis specification -- in this case axis=0 does not + # have size 1, but the priority should be to raise + # a TypeError for the axis argument and NOT a + # ValueError for squeezing a non-empty dimension + with assert_raises(TypeError): + np.squeeze(oldsqueeze, axis=0) + + # the new API knows how to handle the axis + # argument and will return a ValueError if + # attempting to squeeze an axis that is not + # of length 1 + with assert_raises(ValueError): + np.squeeze(np.array([[1],[2],[3]]), axis=0) + + def test_reduce_contiguous(self): + # GitHub issue #387 + a = np.add.reduce(np.zeros((2, 1, 2)), (0, 1)) + b = np.add.reduce(np.zeros((2, 1, 2)), 1) + assert_(a.flags.c_contiguous) + assert_(a.flags.f_contiguous) + assert_(b.flags.c_contiguous) + + @pytest.mark.skipif(IS_PYSTON, reason="Pyston disables recursion checking") + def test_object_array_self_reference(self): + # Object arrays with references to themselves can cause problems + a = np.array(0, dtype=object) + a[()] = a + assert_raises(RecursionError, int, a) + assert_raises(RecursionError, float, a) + a[()] = None + + @pytest.mark.skipif(IS_PYSTON, reason="Pyston disables recursion checking") + def test_object_array_circular_reference(self): + # Test the same for a circular reference. + a = np.array(0, dtype=object) + b = np.array(0, dtype=object) + a[()] = b + b[()] = a + assert_raises(RecursionError, int, a) + # NumPy has no tp_traverse currently, so circular references + # cannot be detected. So resolve it: + a[()] = None + + # This was causing a to become like the above + a = np.array(0, dtype=object) + a[...] += 1 + assert_equal(a, 1) + + def test_object_array_nested(self): + # but is fine with a reference to a different array + a = np.array(0, dtype=object) + b = np.array(0, dtype=object) + a[()] = b + assert_equal(int(a), int(0)) + assert_equal(float(a), float(0)) + + def test_object_array_self_copy(self): + # An object array being copied into itself DECREF'ed before INCREF'ing + # causing segmentation faults (gh-3787) + a = np.array(object(), dtype=object) + np.copyto(a, a) + if HAS_REFCOUNT: + assert_(sys.getrefcount(a[()]) == 2) + a[()].__class__ # will segfault if object was deleted + + def test_zerosize_accumulate(self): + "Ticket #1733" + x = np.array([[42, 0]], dtype=np.uint32) + assert_equal(np.add.accumulate(x[:-1, 0]), []) + + def test_objectarray_setfield(self): + # Setfield should not overwrite Object fields with non-Object data + x = np.array([1, 2, 3], dtype=object) + assert_raises(TypeError, x.setfield, 4, np.int32, 0) + + def test_setting_rank0_string(self): + "Ticket #1736" + s1 = b"hello1" + s2 = b"hello2" + a = np.zeros((), dtype="S10") + a[()] = s1 + assert_equal(a, np.array(s1)) + a[()] = np.array(s2) + assert_equal(a, np.array(s2)) + + a = np.zeros((), dtype='f4') + a[()] = 3 + assert_equal(a, np.array(3)) + a[()] = np.array(4) + assert_equal(a, np.array(4)) + + def test_string_astype(self): + "Ticket #1748" + s1 = b'black' + s2 = b'white' + s3 = b'other' + a = np.array([[s1], [s2], [s3]]) + assert_equal(a.dtype, np.dtype('S5')) + b = a.astype(np.dtype('S0')) + assert_equal(b.dtype, np.dtype('S5')) + + def test_ticket_1756(self): + # Ticket #1756 + s = b'0123456789abcdef' + a = np.array([s]*5) + for i in range(1, 17): + a1 = np.array(a, "|S%d" % i) + a2 = np.array([s[:i]]*5) + assert_equal(a1, a2) + + def test_fields_strides(self): + "gh-2355" + r = np.frombuffer(b'abcdefghijklmnop'*4*3, dtype='i4,(2,3)u2') + assert_equal(r[0:3:2]['f1'], r['f1'][0:3:2]) + assert_equal(r[0:3:2]['f1'][0], r[0:3:2][0]['f1']) + assert_equal(r[0:3:2]['f1'][0][()], r[0:3:2][0]['f1'][()]) + assert_equal(r[0:3:2]['f1'][0].strides, r[0:3:2][0]['f1'].strides) + + def test_alignment_update(self): + # Check that alignment flag is updated on stride setting + a = np.arange(10) + assert_(a.flags.aligned) + a.strides = 3 + assert_(not a.flags.aligned) + + def test_ticket_1770(self): + "Should not segfault on python 3k" + import numpy as np + try: + a = np.zeros((1,), dtype=[('f1', 'f')]) + a['f1'] = 1 + a['f2'] = 1 + except ValueError: + pass + except Exception: + raise AssertionError + + def test_ticket_1608(self): + "x.flat shouldn't modify data" + x = np.array([[1, 2], [3, 4]]).T + np.array(x.flat) + assert_equal(x, [[1, 3], [2, 4]]) + + def test_pickle_string_overwrite(self): + import re + + data = np.array([1], dtype='b') + blob = pickle.dumps(data, protocol=1) + data = pickle.loads(blob) + + # Check that loads does not clobber interned strings + s = re.sub("a(.)", "\x01\\1", "a_") + assert_equal(s[0], "\x01") + data[0] = 0x6a + s = re.sub("a(.)", "\x01\\1", "a_") + assert_equal(s[0], "\x01") + + def test_pickle_bytes_overwrite(self): + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + data = np.array([1], dtype='b') + data = pickle.loads(pickle.dumps(data, protocol=proto)) + data[0] = 0x7d + bytestring = "\x01 ".encode('ascii') + assert_equal(bytestring[0:1], '\x01'.encode('ascii')) + + def test_pickle_py2_array_latin1_hack(self): + # Check that unpickling hacks in Py3 that support + # encoding='latin1' work correctly. + + # Python2 output for pickle.dumps(numpy.array([129], dtype='b')) + data = b"cnumpy.core.multiarray\n_reconstruct\np0\n(cnumpy\nndarray\np1\n(I0\ntp2\nS'b'\np3\ntp4\nRp5\n(I1\n(I1\ntp6\ncnumpy\ndtype\np7\n(S'i1'\np8\nI0\nI1\ntp9\nRp10\n(I3\nS'|'\np11\nNNNI-1\nI-1\nI0\ntp12\nbI00\nS'\\x81'\np13\ntp14\nb." # noqa + # This should work: + result = pickle.loads(data, encoding='latin1') + assert_array_equal(result, np.array([129]).astype('b')) + # Should not segfault: + assert_raises(Exception, pickle.loads, data, encoding='koi8-r') + + def test_pickle_py2_scalar_latin1_hack(self): + # Check that scalar unpickling hack in Py3 that supports + # encoding='latin1' work correctly. + + # Python2 output for pickle.dumps(...) + datas = [ + # (original, python2_pickle, koi8r_validity) + (np.str_('\u6bd2'), + b"cnumpy.core.multiarray\nscalar\np0\n(cnumpy\ndtype\np1\n(S'U1'\np2\nI0\nI1\ntp3\nRp4\n(I3\nS'<'\np5\nNNNI4\nI4\nI0\ntp6\nbS'\\xd2k\\x00\\x00'\np7\ntp8\nRp9\n.", # noqa + 'invalid'), + + (np.float64(9e123), + b"cnumpy.core.multiarray\nscalar\np0\n(cnumpy\ndtype\np1\n(S'f8'\np2\nI0\nI1\ntp3\nRp4\n(I3\nS'<'\np5\nNNNI-1\nI-1\nI0\ntp6\nbS'O\\x81\\xb7Z\\xaa:\\xabY'\np7\ntp8\nRp9\n.", # noqa + 'invalid'), + + # different 8-bit code point in KOI8-R vs latin1 + (np.bytes_(b'\x9c'), + b"cnumpy.core.multiarray\nscalar\np0\n(cnumpy\ndtype\np1\n(S'S1'\np2\nI0\nI1\ntp3\nRp4\n(I3\nS'|'\np5\nNNNI1\nI1\nI0\ntp6\nbS'\\x9c'\np7\ntp8\nRp9\n.", # noqa + 'different'), + ] + for original, data, koi8r_validity in datas: + result = pickle.loads(data, encoding='latin1') + assert_equal(result, original) + + # Decoding under non-latin1 encoding (e.g.) KOI8-R can + # produce bad results, but should not segfault. + if koi8r_validity == 'different': + # Unicode code points happen to lie within latin1, + # but are different in koi8-r, resulting to silent + # bogus results + result = pickle.loads(data, encoding='koi8-r') + assert_(result != original) + elif koi8r_validity == 'invalid': + # Unicode code points outside latin1, so results + # to an encoding exception + assert_raises( + ValueError, pickle.loads, data, encoding='koi8-r' + ) + else: + raise ValueError(koi8r_validity) + + def test_structured_type_to_object(self): + a_rec = np.array([(0, 1), (3, 2)], dtype='i4,i8') + a_obj = np.empty((2,), dtype=object) + a_obj[0] = (0, 1) + a_obj[1] = (3, 2) + # astype records -> object + assert_equal(a_rec.astype(object), a_obj) + # '=' records -> object + b = np.empty_like(a_obj) + b[...] = a_rec + assert_equal(b, a_obj) + # '=' object -> records + b = np.empty_like(a_rec) + b[...] = a_obj + assert_equal(b, a_rec) + + def test_assign_obj_listoflists(self): + # Ticket # 1870 + # The inner list should get assigned to the object elements + a = np.zeros(4, dtype=object) + b = a.copy() + a[0] = [1] + a[1] = [2] + a[2] = [3] + a[3] = [4] + b[...] = [[1], [2], [3], [4]] + assert_equal(a, b) + # The first dimension should get broadcast + a = np.zeros((2, 2), dtype=object) + a[...] = [[1, 2]] + assert_equal(a, [[1, 2], [1, 2]]) + + @pytest.mark.slow_pypy + def test_memoryleak(self): + # Ticket #1917 - ensure that array data doesn't leak + for i in range(1000): + # 100MB times 1000 would give 100GB of memory usage if it leaks + a = np.empty((100000000,), dtype='i1') + del a + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_ufunc_reduce_memoryleak(self): + a = np.arange(6) + acnt = sys.getrefcount(a) + np.add.reduce(a) + assert_equal(sys.getrefcount(a), acnt) + + def test_search_sorted_invalid_arguments(self): + # Ticket #2021, should not segfault. + x = np.arange(0, 4, dtype='datetime64[D]') + assert_raises(TypeError, x.searchsorted, 1) + + def test_string_truncation(self): + # Ticket #1990 - Data can be truncated in creation of an array from a + # mixed sequence of numeric values and strings (gh-2583) + for val in [True, 1234, 123.4, complex(1, 234)]: + for tostr, dtype in [(asunicode, "U"), (asbytes, "S")]: + b = np.array([val, tostr('xx')], dtype=dtype) + assert_equal(tostr(b[0]), tostr(val)) + b = np.array([tostr('xx'), val], dtype=dtype) + assert_equal(tostr(b[1]), tostr(val)) + + # test also with longer strings + b = np.array([val, tostr('xxxxxxxxxx')], dtype=dtype) + assert_equal(tostr(b[0]), tostr(val)) + b = np.array([tostr('xxxxxxxxxx'), val], dtype=dtype) + assert_equal(tostr(b[1]), tostr(val)) + + def test_string_truncation_ucs2(self): + # Ticket #2081. Python compiled with two byte unicode + # can lead to truncation if itemsize is not properly + # adjusted for NumPy's four byte unicode. + a = np.array(['abcd']) + assert_equal(a.dtype.itemsize, 16) + + def test_unique_stable(self): + # Ticket #2063 must always choose stable sort for argsort to + # get consistent results + v = np.array(([0]*5 + [1]*6 + [2]*6)*4) + res = np.unique(v, return_index=True) + tgt = (np.array([0, 1, 2]), np.array([ 0, 5, 11])) + assert_equal(res, tgt) + + def test_unicode_alloc_dealloc_match(self): + # Ticket #1578, the mismatch only showed up when running + # python-debug for python versions >= 2.7, and then as + # a core dump and error message. + a = np.array(['abc'], dtype=np.str_)[0] + del a + + def test_refcount_error_in_clip(self): + # Ticket #1588 + a = np.zeros((2,), dtype='>i2').clip(min=0) + x = a + a + # This used to segfault: + y = str(x) + # Check the final string: + assert_(y == "[0 0]") + + def test_searchsorted_wrong_dtype(self): + # Ticket #2189, it used to segfault, so we check that it raises the + # proper exception. + a = np.array([('a', 1)], dtype='S1, int') + assert_raises(TypeError, np.searchsorted, a, 1.2) + # Ticket #2066, similar problem: + dtype = np.rec.format_parser(['i4', 'i4'], [], []) + a = np.recarray((2,), dtype) + a[...] = [(1, 2), (3, 4)] + assert_raises(TypeError, np.searchsorted, a, 1) + + def test_complex64_alignment(self): + # Issue gh-2668 (trac 2076), segfault on sparc due to misalignment + dtt = np.complex64 + arr = np.arange(10, dtype=dtt) + # 2D array + arr2 = np.reshape(arr, (2, 5)) + # Fortran write followed by (C or F) read caused bus error + data_str = arr2.tobytes('F') + data_back = np.ndarray(arr2.shape, + arr2.dtype, + buffer=data_str, + order='F') + assert_array_equal(arr2, data_back) + + def test_structured_count_nonzero(self): + arr = np.array([0, 1]).astype('i4, 2i4')[:1] + count = np.count_nonzero(arr) + assert_equal(count, 0) + + def test_copymodule_preserves_f_contiguity(self): + a = np.empty((2, 2), order='F') + b = copy.copy(a) + c = copy.deepcopy(a) + assert_(b.flags.fortran) + assert_(b.flags.f_contiguous) + assert_(c.flags.fortran) + assert_(c.flags.f_contiguous) + + def test_fortran_order_buffer(self): + import numpy as np + a = np.array([['Hello', 'Foob']], dtype='U5', order='F') + arr = np.ndarray(shape=[1, 2, 5], dtype='U1', buffer=a) + arr2 = np.array([[['H', 'e', 'l', 'l', 'o'], + ['F', 'o', 'o', 'b', '']]]) + assert_array_equal(arr, arr2) + + def test_assign_from_sequence_error(self): + # Ticket #4024. + arr = np.array([1, 2, 3]) + assert_raises(ValueError, arr.__setitem__, slice(None), [9, 9]) + arr.__setitem__(slice(None), [9]) + assert_equal(arr, [9, 9, 9]) + + def test_format_on_flex_array_element(self): + # Ticket #4369. + dt = np.dtype([('date', ' 0: + # unpickling ndarray goes through _frombuffer for protocol 5 + assert b'numpy._core.numeric' in s + else: + assert b'numpy._core.multiarray' in s + + def test_object_casting_errors(self): + # gh-11993 update to ValueError (see gh-16909), since strings can in + # principle be converted to complex, but this string cannot. + arr = np.array(['AAAAA', 18465886.0, 18465886.0], dtype=object) + assert_raises(ValueError, arr.astype, 'c8') + + def test_eff1d_casting(self): + # gh-12711 + x = np.array([1, 2, 4, 7, 0], dtype=np.int16) + res = np.ediff1d(x, to_begin=-99, to_end=np.array([88, 99])) + assert_equal(res, [-99, 1, 2, 3, -7, 88, 99]) + + # The use of safe casting means, that 1<<20 is cast unsafely, an + # error may be better, but currently there is no mechanism for it. + res = np.ediff1d(x, to_begin=(1<<20), to_end=(1<<20)) + assert_equal(res, [0, 1, 2, 3, -7, 0]) + + def test_pickle_datetime64_array(self): + # gh-12745 (would fail with pickle5 installed) + d = np.datetime64('2015-07-04 12:59:59.50', 'ns') + arr = np.array([d]) + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + dumped = pickle.dumps(arr, protocol=proto) + assert_equal(pickle.loads(dumped), arr) + + def test_bad_array_interface(self): + class T: + __array_interface__ = {} + + with assert_raises(ValueError): + np.array([T()]) + + def test_2d__array__shape(self): + class T: + def __array__(self, dtype=None, copy=None): + return np.ndarray(shape=(0,0)) + + # Make sure __array__ is used instead of Sequence methods. + def __iter__(self): + return iter([]) + + def __getitem__(self, idx): + raise AssertionError("__getitem__ was called") + + def __len__(self): + return 0 + + + t = T() + # gh-13659, would raise in broadcasting [x=t for x in result] + arr = np.array([t]) + assert arr.shape == (1, 0, 0) + + @pytest.mark.skipif(sys.maxsize < 2 ** 31 + 1, reason='overflows 32-bit python') + def test_to_ctypes(self): + #gh-14214 + arr = np.zeros((2 ** 31 + 1,), 'b') + assert arr.size * arr.itemsize > 2 ** 31 + c_arr = np.ctypeslib.as_ctypes(arr) + assert_equal(c_arr._length_, arr.size) + + def test_complex_conversion_error(self): + # gh-17068 + with pytest.raises(TypeError, match=r"Unable to convert dtype.*"): + complex(np.array("now", np.datetime64)) + + def test__array_interface__descr(self): + # gh-17068 + dt = np.dtype(dict(names=['a', 'b'], + offsets=[0, 0], + formats=[np.int64, np.int64])) + descr = np.array((1, 1), dtype=dt).__array_interface__['descr'] + assert descr == [('', '|V8')] # instead of [(b'', '|V8')] + + @pytest.mark.skipif(sys.maxsize < 2 ** 31 + 1, reason='overflows 32-bit python') + @requires_memory(free_bytes=9e9) + def test_dot_big_stride(self): + # gh-17111 + # blas stride = stride//itemsize > int32 max + int32_max = np.iinfo(np.int32).max + n = int32_max + 3 + a = np.empty([n], dtype=np.float32) + b = a[::n-1] + b[...] = 1 + assert b.strides[0] > int32_max * b.dtype.itemsize + assert np.dot(b, b) == 2.0 + + def test_frompyfunc_name(self): + # name conversion was failing for python 3 strings + # resulting in the default '?' name. Also test utf-8 + # encoding using non-ascii name. + def cassé(x): + return x + + f = np.frompyfunc(cassé, 1, 1) + assert str(f) == "" + + @pytest.mark.parametrize("operation", [ + 'add', 'subtract', 'multiply', 'floor_divide', + 'conjugate', 'fmod', 'square', 'reciprocal', + 'power', 'absolute', 'negative', 'positive', + 'greater', 'greater_equal', 'less', + 'less_equal', 'equal', 'not_equal', 'logical_and', + 'logical_not', 'logical_or', 'bitwise_and', 'bitwise_or', + 'bitwise_xor', 'invert', 'left_shift', 'right_shift', + 'gcd', 'lcm' + ] + ) + @pytest.mark.parametrize("order", [ + ('b->', 'B->'), + ('h->', 'H->'), + ('i->', 'I->'), + ('l->', 'L->'), + ('q->', 'Q->'), + ] + ) + def test_ufunc_order(self, operation, order): + # gh-18075 + # Ensure signed types before unsigned + def get_idx(string, str_lst): + for i, s in enumerate(str_lst): + if string in s: + return i + raise ValueError(f"{string} not in list") + types = getattr(np, operation).types + assert get_idx(order[0], types) < get_idx(order[1], types), ( + f"Unexpected types order of ufunc in {operation}" + f"for {order}. Possible fix: Use signed before unsigned" + "in generate_umath.py") + + def test_nonbool_logical(self): + # gh-22845 + # create two arrays with bit patterns that do not overlap. + # needs to be large enough to test both SIMD and scalar paths + size = 100 + a = np.frombuffer(b'\x01' * size, dtype=np.bool) + b = np.frombuffer(b'\x80' * size, dtype=np.bool) + expected = np.ones(size, dtype=np.bool) + assert_array_equal(np.logical_and(a, b), expected) + + @pytest.mark.skipif(IS_PYPY, reason="PyPy issue 2742") + def test_gh_23737(self): + with pytest.raises(TypeError, match="not an acceptable base type"): + class Y(np.flexible): + pass + + with pytest.raises(TypeError, match="not an acceptable base type"): + class X(np.flexible, np.ma.core.MaskedArray): + pass + + def test_load_ufunc_pickle(self): + # ufuncs are pickled with a semi-private path in + # numpy.core._multiarray_umath and must be loadable without warning + # despite np.core being deprecated. + test_data = b'\x80\x04\x95(\x00\x00\x00\x00\x00\x00\x00\x8c\x1cnumpy.core._multiarray_umath\x94\x8c\x03add\x94\x93\x94.' # noqa + result = pickle.loads(test_data, encoding='bytes') + assert result is np.add + + def test__array_namespace__(self): + arr = np.arange(2) + + xp = arr.__array_namespace__() + assert xp is np + xp = arr.__array_namespace__(api_version="2021.12") + assert xp is np + xp = arr.__array_namespace__(api_version="2022.12") + assert xp is np + xp = arr.__array_namespace__(api_version="2023.12") + assert xp is np + xp = arr.__array_namespace__(api_version=None) + assert xp is np + + with pytest.raises( + ValueError, + match="Version \"2024.12\" of the Array API Standard " + "is not supported." + ): + arr.__array_namespace__(api_version="2024.12") + + with pytest.raises( + ValueError, + match="Only None and strings are allowed as the Array API version" + ): + arr.__array_namespace__(api_version=2023) + + def test_isin_refcnt_bug(self): + # gh-25295 + for _ in range(1000): + np.isclose(np.int64(2), np.int64(2), atol=1e-15, rtol=1e-300) + + def test_replace_regression(self): + # gh-25513 segfault + carr = np.char.chararray((2,), itemsize=25) + test_strings = [b' 4.52173913043478315E+00', + b' 4.95652173913043548E+00'] + carr[:] = test_strings + out = carr.replace(b"E", b"D") + expected = np.char.chararray((2,), itemsize=25) + expected[:] = [s.replace(b"E", b"D") for s in test_strings] + assert_array_equal(out, expected) + + def test_logspace_base_does_not_determine_dtype(self): + # gh-24957 and cupy/cupy/issues/7946 + start = np.array([0, 2], dtype=np.float16) + stop = np.array([2, 0], dtype=np.float16) + out = np.logspace(start, stop, num=5, axis=1, dtype=np.float32) + expected = np.array([[1., 3.1621094, 10., 31.625, 100.], + [100., 31.625, 10., 3.1621094, 1.]], + dtype=np.float32) + assert_almost_equal(out, expected) + # Check test fails if the calculation is done in float64, as happened + # before when a python float base incorrectly influenced the dtype. + out2 = np.logspace(start, stop, num=5, axis=1, dtype=np.float32, + base=np.array([10.0])) + with pytest.raises(AssertionError, match="not almost equal"): + assert_almost_equal(out2, expected) + + def test_vectorize_fixed_width_string(self): + arr = np.array(["SOme wOrd DŽ ß ᾛ ΣΣ ffi⁵Å Ç Ⅰ"]).astype(np.str_) + f = str.casefold + res = np.vectorize(f, otypes=[arr.dtype])(arr) + assert res.dtype == "U30" + + def test_repeated_square_consistency(self): + # gh-26940 + buf = np.array([-5.171866611150749e-07 + 2.5618634555957426e-07j, + 0, 0, 0, 0, 0]) + # Test buffer with regular and reverse strides + for in_vec in [buf[:3], buf[:3][::-1]]: + expected_res = np.square(in_vec) + # Output vector immediately follows input vector + # to reproduce off-by-one in nomemoverlap check. + for res in [buf[3:], buf[3:][::-1]]: + res = buf[3:] + np.square(in_vec, out=res) + assert_equal(res, expected_res) + + def test_sort_unique_crash(self): + # gh-27037 + for _ in range(4): + vals = np.linspace(0, 1, num=128) + data = np.broadcast_to(vals, (128, 128, 128)) + data = data.transpose(0, 2, 1).copy() + np.unique(data) + + def test_sort_overlap(self): + # gh-27273 + size = 100 + inp = np.linspace(0, size, num=size, dtype=np.intc) + out = np.sort(inp) + assert_equal(inp, out) + + def test_searchsorted_structured(self): + # gh-28190 + x = np.array([(0, 1.)], dtype=[('time', ' None: + cls = np.dtype(code).type + value = cls(str_value) + assert not value.is_integer() + + @pytest.mark.parametrize( + "code", np.typecodes["Float"] + np.typecodes["AllInteger"] + ) + def test_true(self, code: str) -> None: + float_array = np.arange(-5, 5).astype(code) + for value in float_array: + assert value.is_integer() + + @pytest.mark.parametrize("code", np.typecodes["Float"]) + def test_false(self, code: str) -> None: + float_array = np.arange(-5, 5).astype(code) + float_array *= 1.1 + for value in float_array: + if value == 0: + continue + assert not value.is_integer() + + +class TestClassGetItem: + @pytest.mark.parametrize("cls", [ + np.number, + np.integer, + np.inexact, + np.unsignedinteger, + np.signedinteger, + np.floating, + ]) + def test_abc(self, cls: Type[np.number]) -> None: + alias = cls[Any] + assert isinstance(alias, types.GenericAlias) + assert alias.__origin__ is cls + + def test_abc_complexfloating(self) -> None: + alias = np.complexfloating[Any, Any] + assert isinstance(alias, types.GenericAlias) + assert alias.__origin__ is np.complexfloating + + @pytest.mark.parametrize("arg_len", range(4)) + def test_abc_complexfloating_subscript_tuple(self, arg_len: int) -> None: + arg_tup = (Any,) * arg_len + if arg_len in (1, 2): + assert np.complexfloating[arg_tup] + else: + match = f"Too {'few' if arg_len == 0 else 'many'} arguments" + with pytest.raises(TypeError, match=match): + np.complexfloating[arg_tup] + + @pytest.mark.parametrize("cls", [np.generic, np.flexible, np.character]) + def test_abc_non_numeric(self, cls: Type[np.generic]) -> None: + with pytest.raises(TypeError): + cls[Any] + + @pytest.mark.parametrize("code", np.typecodes["All"]) + def test_concrete(self, code: str) -> None: + cls = np.dtype(code).type + with pytest.raises(TypeError): + cls[Any] + + @pytest.mark.parametrize("arg_len", range(4)) + def test_subscript_tuple(self, arg_len: int) -> None: + arg_tup = (Any,) * arg_len + if arg_len == 1: + assert np.number[arg_tup] + else: + with pytest.raises(TypeError): + np.number[arg_tup] + + def test_subscript_scalar(self) -> None: + assert np.number[Any] + + +class TestBitCount: + # derived in part from the cpython test "test_bit_count" + + @pytest.mark.parametrize("itype", sctypes['int']+sctypes['uint']) + def test_small(self, itype): + for a in range(max(np.iinfo(itype).min, 0), 128): + msg = f"Smoke test for {itype}({a}).bit_count()" + assert itype(a).bit_count() == bin(a).count("1"), msg + + def test_bit_count(self): + for exp in [10, 17, 63]: + a = 2**exp + assert np.uint64(a).bit_count() == 1 + assert np.uint64(a - 1).bit_count() == exp + assert np.uint64(a ^ 63).bit_count() == 7 + assert np.uint64((a - 1) ^ 510).bit_count() == exp - 8 + + +class TestDevice: + """ + Test scalar.device attribute and scalar.to_device() method. + """ + scalars = [np.bool(True), np.int64(1), np.uint64(1), np.float64(1.0), + np.complex128(1+1j)] + + @pytest.mark.parametrize("scalar", scalars) + def test_device(self, scalar): + assert scalar.device == "cpu" + + @pytest.mark.parametrize("scalar", scalars) + def test_to_device(self, scalar): + assert scalar.to_device("cpu") is scalar + + @pytest.mark.parametrize("scalar", scalars) + def test___array_namespace__(self, scalar): + assert scalar.__array_namespace__() is np + + +@pytest.mark.parametrize("scalar", [np.bool(True), np.int8(1), np.float64(1)]) +def test_array_wrap(scalar): + # Test scalars array wrap as long as it exists. NumPy itself should + # probably not use it, so it may not be necessary to keep it around. + + arr0d = np.array(3, dtype=np.int8) + # Third argument not passed, None, or True "decays" to scalar. + # (I don't think NumPy would pass `None`, but it seems clear to support) + assert type(scalar.__array_wrap__(arr0d)) is np.int8 + assert type(scalar.__array_wrap__(arr0d, None, None)) is np.int8 + assert type(scalar.__array_wrap__(arr0d, None, True)) is np.int8 + + # Otherwise, result should be the input + assert scalar.__array_wrap__(arr0d, None, False) is arr0d + + # An old bug. A non 0-d array cannot be converted to scalar: + arr1d = np.array([3], dtype=np.int8) + assert scalar.__array_wrap__(arr1d) is arr1d + assert scalar.__array_wrap__(arr1d, None, True) is arr1d diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarbuffer.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarbuffer.py new file mode 100644 index 0000000000000000000000000000000000000000..e8a5c1cedac76e7fcf7f971a2b85047532201ad8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarbuffer.py @@ -0,0 +1,153 @@ +""" +Test scalar buffer interface adheres to PEP 3118 +""" +import numpy as np +from numpy._core._rational_tests import rational +from numpy._core._multiarray_tests import get_buffer_info +import pytest + +from numpy.testing import assert_, assert_equal, assert_raises + +# PEP3118 format strings for native (standard alignment and byteorder) types +scalars_and_codes = [ + (np.bool, '?'), + (np.byte, 'b'), + (np.short, 'h'), + (np.intc, 'i'), + (np.long, 'l'), + (np.longlong, 'q'), + (np.ubyte, 'B'), + (np.ushort, 'H'), + (np.uintc, 'I'), + (np.ulong, 'L'), + (np.ulonglong, 'Q'), + (np.half, 'e'), + (np.single, 'f'), + (np.double, 'd'), + (np.longdouble, 'g'), + (np.csingle, 'Zf'), + (np.cdouble, 'Zd'), + (np.clongdouble, 'Zg'), +] +scalars_only, codes_only = zip(*scalars_and_codes) + + +class TestScalarPEP3118: + + @pytest.mark.parametrize('scalar', scalars_only, ids=codes_only) + def test_scalar_match_array(self, scalar): + x = scalar() + a = np.array([], dtype=np.dtype(scalar)) + mv_x = memoryview(x) + mv_a = memoryview(a) + assert_equal(mv_x.format, mv_a.format) + + @pytest.mark.parametrize('scalar', scalars_only, ids=codes_only) + def test_scalar_dim(self, scalar): + x = scalar() + mv_x = memoryview(x) + assert_equal(mv_x.itemsize, np.dtype(scalar).itemsize) + assert_equal(mv_x.ndim, 0) + assert_equal(mv_x.shape, ()) + assert_equal(mv_x.strides, ()) + assert_equal(mv_x.suboffsets, ()) + + @pytest.mark.parametrize('scalar, code', scalars_and_codes, ids=codes_only) + def test_scalar_code_and_properties(self, scalar, code): + x = scalar() + expected = dict(strides=(), itemsize=x.dtype.itemsize, ndim=0, + shape=(), format=code, readonly=True) + + mv_x = memoryview(x) + assert self._as_dict(mv_x) == expected + + @pytest.mark.parametrize('scalar', scalars_only, ids=codes_only) + def test_scalar_buffers_readonly(self, scalar): + x = scalar() + with pytest.raises(BufferError, match="scalar buffer is readonly"): + get_buffer_info(x, ["WRITABLE"]) + + def test_void_scalar_structured_data(self): + dt = np.dtype([('name', np.str_, 16), ('grades', np.float64, (2,))]) + x = np.array(('ndarray_scalar', (1.2, 3.0)), dtype=dt)[()] + assert_(isinstance(x, np.void)) + mv_x = memoryview(x) + expected_size = 16 * np.dtype((np.str_, 1)).itemsize + expected_size += 2 * np.dtype(np.float64).itemsize + assert_equal(mv_x.itemsize, expected_size) + assert_equal(mv_x.ndim, 0) + assert_equal(mv_x.shape, ()) + assert_equal(mv_x.strides, ()) + assert_equal(mv_x.suboffsets, ()) + + # check scalar format string against ndarray format string + a = np.array([('Sarah', (8.0, 7.0)), ('John', (6.0, 7.0))], dtype=dt) + assert_(isinstance(a, np.ndarray)) + mv_a = memoryview(a) + assert_equal(mv_x.itemsize, mv_a.itemsize) + assert_equal(mv_x.format, mv_a.format) + + # Check that we do not allow writeable buffer export (technically + # we could allow it sometimes here...) + with pytest.raises(BufferError, match="scalar buffer is readonly"): + get_buffer_info(x, ["WRITABLE"]) + + def _as_dict(self, m): + return dict(strides=m.strides, shape=m.shape, itemsize=m.itemsize, + ndim=m.ndim, format=m.format, readonly=m.readonly) + + def test_datetime_memoryview(self): + # gh-11656 + # Values verified with v1.13.3, shape is not () as in test_scalar_dim + + dt1 = np.datetime64('2016-01-01') + dt2 = np.datetime64('2017-01-01') + expected = dict(strides=(1,), itemsize=1, ndim=1, shape=(8,), + format='B', readonly=True) + v = memoryview(dt1) + assert self._as_dict(v) == expected + + v = memoryview(dt2 - dt1) + assert self._as_dict(v) == expected + + dt = np.dtype([('a', 'uint16'), ('b', 'M8[s]')]) + a = np.empty(1, dt) + # Fails to create a PEP 3118 valid buffer + assert_raises((ValueError, BufferError), memoryview, a[0]) + + # Check that we do not allow writeable buffer export + with pytest.raises(BufferError, match="scalar buffer is readonly"): + get_buffer_info(dt1, ["WRITABLE"]) + + @pytest.mark.parametrize('s', [ + pytest.param("\x32\x32", id="ascii"), + pytest.param("\uFE0F\uFE0F", id="basic multilingual"), + pytest.param("\U0001f4bb\U0001f4bb", id="non-BMP"), + ]) + def test_str_ucs4(self, s): + s = np.str_(s) # only our subclass implements the buffer protocol + + # all the same, characters always encode as ucs4 + expected = dict(strides=(), itemsize=8, ndim=0, shape=(), format='2w', + readonly=True) + + v = memoryview(s) + assert self._as_dict(v) == expected + + # integers of the paltform-appropriate endianness + code_points = np.frombuffer(v, dtype='i4') + + assert_equal(code_points, [ord(c) for c in s]) + + # Check that we do not allow writeable buffer export + with pytest.raises(BufferError, match="scalar buffer is readonly"): + get_buffer_info(s, ["WRITABLE"]) + + def test_user_scalar_fails_buffer(self): + r = rational(1) + with assert_raises(TypeError): + memoryview(r) + + # Check that we do not allow writeable buffer export + with pytest.raises(BufferError, match="scalar buffer is readonly"): + get_buffer_info(r, ["WRITABLE"]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarinherit.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarinherit.py new file mode 100644 index 0000000000000000000000000000000000000000..3751c6ca0d3860db06342745fe573bd5a9293230 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarinherit.py @@ -0,0 +1,105 @@ +""" Test printing of scalar types. + +""" +import pytest + +import numpy as np +from numpy.testing import assert_, assert_raises + + +class A: + pass +class B(A, np.float64): + pass + +class C(B): + pass +class D(C, B): + pass + +class B0(np.float64, A): + pass +class C0(B0): + pass + +class HasNew: + def __new__(cls, *args, **kwargs): + return cls, args, kwargs + +class B1(np.float64, HasNew): + pass + + +class TestInherit: + def test_init(self): + x = B(1.0) + assert_(str(x) == '1.0') + y = C(2.0) + assert_(str(y) == '2.0') + z = D(3.0) + assert_(str(z) == '3.0') + + def test_init2(self): + x = B0(1.0) + assert_(str(x) == '1.0') + y = C0(2.0) + assert_(str(y) == '2.0') + + def test_gh_15395(self): + # HasNew is the second base, so `np.float64` should have priority + x = B1(1.0) + assert_(str(x) == '1.0') + + # previously caused RecursionError!? + with pytest.raises(TypeError): + B1(1.0, 2.0) + + def test_int_repr(self): + # Test that integer repr works correctly for subclasses (gh-27106) + class my_int16(np.int16): + pass + + s = repr(my_int16(3)) + assert s == "my_int16(3)" + +class TestCharacter: + def test_char_radd(self): + # GH issue 9620, reached gentype_add and raise TypeError + np_s = np.bytes_('abc') + np_u = np.str_('abc') + s = b'def' + u = 'def' + assert_(np_s.__radd__(np_s) is NotImplemented) + assert_(np_s.__radd__(np_u) is NotImplemented) + assert_(np_s.__radd__(s) is NotImplemented) + assert_(np_s.__radd__(u) is NotImplemented) + assert_(np_u.__radd__(np_s) is NotImplemented) + assert_(np_u.__radd__(np_u) is NotImplemented) + assert_(np_u.__radd__(s) is NotImplemented) + assert_(np_u.__radd__(u) is NotImplemented) + assert_(s + np_s == b'defabc') + assert_(u + np_u == 'defabc') + + class MyStr(str, np.generic): + # would segfault + pass + + with assert_raises(TypeError): + # Previously worked, but gave completely wrong result + ret = s + MyStr('abc') + + class MyBytes(bytes, np.generic): + # would segfault + pass + + ret = s + MyBytes(b'abc') + assert(type(ret) is type(s)) + assert ret == b"defabc" + + def test_char_repeat(self): + np_s = np.bytes_('abc') + np_u = np.str_('abc') + res_s = b'abc' * 5 + res_u = 'abc' * 5 + assert_(np_s * 5 == res_s) + assert_(np_u * 5 == res_u) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarmath.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarmath.py new file mode 100644 index 0000000000000000000000000000000000000000..c21f782a3773141c645f3b1d22e1d000735d118f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarmath.py @@ -0,0 +1,1165 @@ +import contextlib +import sys +import warnings +import itertools +import operator +import platform +from numpy._utils import _pep440 +import pytest +from hypothesis import given, settings +from hypothesis.strategies import sampled_from +from hypothesis.extra import numpy as hynp + +import numpy as np +from numpy.exceptions import ComplexWarning +from numpy._core._rational_tests import rational +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_almost_equal, + assert_array_equal, IS_PYPY, suppress_warnings, _gen_alignment_data, + assert_warns, check_support_sve, + ) + +types = [np.bool, np.byte, np.ubyte, np.short, np.ushort, np.intc, np.uintc, + np.int_, np.uint, np.longlong, np.ulonglong, + np.single, np.double, np.longdouble, np.csingle, + np.cdouble, np.clongdouble] + +floating_types = np.floating.__subclasses__() +complex_floating_types = np.complexfloating.__subclasses__() + +objecty_things = [object(), None, np.array(None, dtype=object)] + +binary_operators_for_scalars = [ + operator.lt, operator.le, operator.eq, operator.ne, operator.ge, + operator.gt, operator.add, operator.floordiv, operator.mod, + operator.mul, operator.pow, operator.sub, operator.truediv +] +binary_operators_for_scalar_ints = binary_operators_for_scalars + [ + operator.xor, operator.or_, operator.and_ +] + + +# This compares scalarmath against ufuncs. + +class TestTypes: + def test_types(self): + for atype in types: + a = atype(1) + assert_(a == 1, "error with %r: got %r" % (atype, a)) + + def test_type_add(self): + # list of types + for k, atype in enumerate(types): + a_scalar = atype(3) + a_array = np.array([3], dtype=atype) + for l, btype in enumerate(types): + b_scalar = btype(1) + b_array = np.array([1], dtype=btype) + c_scalar = a_scalar + b_scalar + c_array = a_array + b_array + # It was comparing the type numbers, but the new ufunc + # function-finding mechanism finds the lowest function + # to which both inputs can be cast - which produces 'l' + # when you do 'q' + 'b'. The old function finding mechanism + # skipped ahead based on the first argument, but that + # does not produce properly symmetric results... + assert_equal(c_scalar.dtype, c_array.dtype, + "error with types (%d/'%c' + %d/'%c')" % + (k, np.dtype(atype).char, l, np.dtype(btype).char)) + + def test_type_create(self): + for k, atype in enumerate(types): + a = np.array([1, 2, 3], atype) + b = atype([1, 2, 3]) + assert_equal(a, b) + + def test_leak(self): + # test leak of scalar objects + # a leak would show up in valgrind as still-reachable of ~2.6MB + for i in range(200000): + np.add(1, 1) + + +def check_ufunc_scalar_equivalence(op, arr1, arr2): + scalar1 = arr1[()] + scalar2 = arr2[()] + assert isinstance(scalar1, np.generic) + assert isinstance(scalar2, np.generic) + + if arr1.dtype.kind == "c" or arr2.dtype.kind == "c": + comp_ops = {operator.ge, operator.gt, operator.le, operator.lt} + if op in comp_ops and (np.isnan(scalar1) or np.isnan(scalar2)): + pytest.xfail("complex comp ufuncs use sort-order, scalars do not.") + if op == operator.pow and arr2.item() in [-1, 0, 0.5, 1, 2]: + # array**scalar special case can have different result dtype + # (Other powers may have issues also, but are not hit here.) + # TODO: It would be nice to resolve this issue. + pytest.skip("array**2 can have incorrect/weird result dtype") + + # ignore fpe's since they may just mismatch for integers anyway. + with warnings.catch_warnings(), np.errstate(all="ignore"): + # Comparisons DeprecationWarnings replacing errors (2022-03): + warnings.simplefilter("error", DeprecationWarning) + try: + res = op(arr1, arr2) + except Exception as e: + with pytest.raises(type(e)): + op(scalar1, scalar2) + else: + scalar_res = op(scalar1, scalar2) + assert_array_equal(scalar_res, res, strict=True) + + +@pytest.mark.slow +@settings(max_examples=10000, deadline=2000) +@given(sampled_from(binary_operators_for_scalars), + hynp.arrays(dtype=hynp.scalar_dtypes(), shape=()), + hynp.arrays(dtype=hynp.scalar_dtypes(), shape=())) +def test_array_scalar_ufunc_equivalence(op, arr1, arr2): + """ + This is a thorough test attempting to cover important promotion paths + and ensuring that arrays and scalars stay as aligned as possible. + However, if it creates troubles, it should maybe just be removed. + """ + check_ufunc_scalar_equivalence(op, arr1, arr2) + + +@pytest.mark.slow +@given(sampled_from(binary_operators_for_scalars), + hynp.scalar_dtypes(), hynp.scalar_dtypes()) +def test_array_scalar_ufunc_dtypes(op, dt1, dt2): + # Same as above, but don't worry about sampling weird values so that we + # do not have to sample as much + arr1 = np.array(2, dtype=dt1) + arr2 = np.array(3, dtype=dt2) # some power do weird things. + + check_ufunc_scalar_equivalence(op, arr1, arr2) + + +@pytest.mark.parametrize("fscalar", [np.float16, np.float32]) +def test_int_float_promotion_truediv(fscalar): + # Promotion for mixed int and float32/float16 must not go to float64 + i = np.int8(1) + f = fscalar(1) + expected = np.result_type(i, f) + assert (i / f).dtype == expected + assert (f / i).dtype == expected + # But normal int / int true division goes to float64: + assert (i / i).dtype == np.dtype("float64") + # For int16, result has to be ast least float32 (takes ufunc path): + assert (np.int16(1) / f).dtype == np.dtype("float32") + + +class TestBaseMath: + @pytest.mark.xfail(check_support_sve(), reason="gh-22982") + def test_blocked(self): + # test alignments offsets for simd instructions + # alignments for vz + 2 * (vs - 1) + 1 + for dt, sz in [(np.float32, 11), (np.float64, 7), (np.int32, 11)]: + for out, inp1, inp2, msg in _gen_alignment_data(dtype=dt, + type='binary', + max_size=sz): + exp1 = np.ones_like(inp1) + inp1[...] = np.ones_like(inp1) + inp2[...] = np.zeros_like(inp2) + assert_almost_equal(np.add(inp1, inp2), exp1, err_msg=msg) + assert_almost_equal(np.add(inp1, 2), exp1 + 2, err_msg=msg) + assert_almost_equal(np.add(1, inp2), exp1, err_msg=msg) + + np.add(inp1, inp2, out=out) + assert_almost_equal(out, exp1, err_msg=msg) + + inp2[...] += np.arange(inp2.size, dtype=dt) + 1 + assert_almost_equal(np.square(inp2), + np.multiply(inp2, inp2), err_msg=msg) + # skip true divide for ints + if dt != np.int32: + assert_almost_equal(np.reciprocal(inp2), + np.divide(1, inp2), err_msg=msg) + + inp1[...] = np.ones_like(inp1) + np.add(inp1, 2, out=out) + assert_almost_equal(out, exp1 + 2, err_msg=msg) + inp2[...] = np.ones_like(inp2) + np.add(2, inp2, out=out) + assert_almost_equal(out, exp1 + 2, err_msg=msg) + + def test_lower_align(self): + # check data that is not aligned to element size + # i.e doubles are aligned to 4 bytes on i386 + d = np.zeros(23 * 8, dtype=np.int8)[4:-4].view(np.float64) + o = np.zeros(23 * 8, dtype=np.int8)[4:-4].view(np.float64) + assert_almost_equal(d + d, d * 2) + np.add(d, d, out=o) + np.add(np.ones_like(d), d, out=o) + np.add(d, np.ones_like(d), out=o) + np.add(np.ones_like(d), d) + np.add(d, np.ones_like(d)) + + +class TestPower: + def test_small_types(self): + for t in [np.int8, np.int16, np.float16]: + a = t(3) + b = a ** 4 + assert_(b == 81, "error with %r: got %r" % (t, b)) + + def test_large_types(self): + for t in [np.int32, np.int64, np.float32, np.float64, np.longdouble]: + a = t(51) + b = a ** 4 + msg = "error with %r: got %r" % (t, b) + if np.issubdtype(t, np.integer): + assert_(b == 6765201, msg) + else: + assert_almost_equal(b, 6765201, err_msg=msg) + + def test_integers_to_negative_integer_power(self): + # Note that the combination of uint64 with a signed integer + # has common type np.float64. The other combinations should all + # raise a ValueError for integer ** negative integer. + exp = [np.array(-1, dt)[()] for dt in 'bhilq'] + + # 1 ** -1 possible special case + base = [np.array(1, dt)[()] for dt in 'bhilqBHILQ'] + for i1, i2 in itertools.product(base, exp): + if i1.dtype != np.uint64: + assert_raises(ValueError, operator.pow, i1, i2) + else: + res = operator.pow(i1, i2) + assert_(res.dtype.type is np.float64) + assert_almost_equal(res, 1.) + + # -1 ** -1 possible special case + base = [np.array(-1, dt)[()] for dt in 'bhilq'] + for i1, i2 in itertools.product(base, exp): + if i1.dtype != np.uint64: + assert_raises(ValueError, operator.pow, i1, i2) + else: + res = operator.pow(i1, i2) + assert_(res.dtype.type is np.float64) + assert_almost_equal(res, -1.) + + # 2 ** -1 perhaps generic + base = [np.array(2, dt)[()] for dt in 'bhilqBHILQ'] + for i1, i2 in itertools.product(base, exp): + if i1.dtype != np.uint64: + assert_raises(ValueError, operator.pow, i1, i2) + else: + res = operator.pow(i1, i2) + assert_(res.dtype.type is np.float64) + assert_almost_equal(res, .5) + + def test_mixed_types(self): + typelist = [np.int8, np.int16, np.float16, + np.float32, np.float64, np.int8, + np.int16, np.int32, np.int64] + for t1 in typelist: + for t2 in typelist: + a = t1(3) + b = t2(2) + result = a**b + msg = ("error with %r and %r:" + "got %r, expected %r") % (t1, t2, result, 9) + if np.issubdtype(np.dtype(result), np.integer): + assert_(result == 9, msg) + else: + assert_almost_equal(result, 9, err_msg=msg) + + def test_modular_power(self): + # modular power is not implemented, so ensure it errors + a = 5 + b = 4 + c = 10 + expected = pow(a, b, c) # noqa: F841 + for t in (np.int32, np.float32, np.complex64): + # note that 3-operand power only dispatches on the first argument + assert_raises(TypeError, operator.pow, t(a), b, c) + assert_raises(TypeError, operator.pow, np.array(t(a)), b, c) + + +def floordiv_and_mod(x, y): + return (x // y, x % y) + + +def _signs(dt): + if dt in np.typecodes['UnsignedInteger']: + return (+1,) + else: + return (+1, -1) + + +class TestModulus: + + def test_modulus_basic(self): + dt = np.typecodes['AllInteger'] + np.typecodes['Float'] + for op in [floordiv_and_mod, divmod]: + for dt1, dt2 in itertools.product(dt, dt): + for sg1, sg2 in itertools.product(_signs(dt1), _signs(dt2)): + fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s' + msg = fmt % (op.__name__, dt1, dt2, sg1, sg2) + a = np.array(sg1*71, dtype=dt1)[()] + b = np.array(sg2*19, dtype=dt2)[()] + div, rem = op(a, b) + assert_equal(div*b + rem, a, err_msg=msg) + if sg2 == -1: + assert_(b < rem <= 0, msg) + else: + assert_(b > rem >= 0, msg) + + def test_float_modulus_exact(self): + # test that float results are exact for small integers. This also + # holds for the same integers scaled by powers of two. + nlst = list(range(-127, 0)) + plst = list(range(1, 128)) + dividend = nlst + [0] + plst + divisor = nlst + plst + arg = list(itertools.product(dividend, divisor)) + tgt = list(divmod(*t) for t in arg) + + a, b = np.array(arg, dtype=int).T + # convert exact integer results from Python to float so that + # signed zero can be used, it is checked. + tgtdiv, tgtrem = np.array(tgt, dtype=float).T + tgtdiv = np.where((tgtdiv == 0.0) & ((b < 0) ^ (a < 0)), -0.0, tgtdiv) + tgtrem = np.where((tgtrem == 0.0) & (b < 0), -0.0, tgtrem) + + for op in [floordiv_and_mod, divmod]: + for dt in np.typecodes['Float']: + msg = 'op: %s, dtype: %s' % (op.__name__, dt) + fa = a.astype(dt) + fb = b.astype(dt) + # use list comprehension so a_ and b_ are scalars + div, rem = zip(*[op(a_, b_) for a_, b_ in zip(fa, fb)]) + assert_equal(div, tgtdiv, err_msg=msg) + assert_equal(rem, tgtrem, err_msg=msg) + + def test_float_modulus_roundoff(self): + # gh-6127 + dt = np.typecodes['Float'] + for op in [floordiv_and_mod, divmod]: + for dt1, dt2 in itertools.product(dt, dt): + for sg1, sg2 in itertools.product((+1, -1), (+1, -1)): + fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s' + msg = fmt % (op.__name__, dt1, dt2, sg1, sg2) + a = np.array(sg1*78*6e-8, dtype=dt1)[()] + b = np.array(sg2*6e-8, dtype=dt2)[()] + div, rem = op(a, b) + # Equal assertion should hold when fmod is used + assert_equal(div*b + rem, a, err_msg=msg) + if sg2 == -1: + assert_(b < rem <= 0, msg) + else: + assert_(b > rem >= 0, msg) + + def test_float_modulus_corner_cases(self): + # Check remainder magnitude. + for dt in np.typecodes['Float']: + b = np.array(1.0, dtype=dt) + a = np.nextafter(np.array(0.0, dtype=dt), -b) + rem = operator.mod(a, b) + assert_(rem <= b, 'dt: %s' % dt) + rem = operator.mod(-a, -b) + assert_(rem >= -b, 'dt: %s' % dt) + + # Check nans, inf + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "invalid value encountered in remainder") + sup.filter(RuntimeWarning, "divide by zero encountered in remainder") + sup.filter(RuntimeWarning, "divide by zero encountered in floor_divide") + sup.filter(RuntimeWarning, "divide by zero encountered in divmod") + sup.filter(RuntimeWarning, "invalid value encountered in divmod") + for dt in np.typecodes['Float']: + fone = np.array(1.0, dtype=dt) + fzer = np.array(0.0, dtype=dt) + finf = np.array(np.inf, dtype=dt) + fnan = np.array(np.nan, dtype=dt) + rem = operator.mod(fone, fzer) + assert_(np.isnan(rem), 'dt: %s' % dt) + # MSVC 2008 returns NaN here, so disable the check. + #rem = operator.mod(fone, finf) + #assert_(rem == fone, 'dt: %s' % dt) + rem = operator.mod(fone, fnan) + assert_(np.isnan(rem), 'dt: %s' % dt) + rem = operator.mod(finf, fone) + assert_(np.isnan(rem), 'dt: %s' % dt) + for op in [floordiv_and_mod, divmod]: + div, mod = op(fone, fzer) + assert_(np.isinf(div)) and assert_(np.isnan(mod)) + + def test_inplace_floordiv_handling(self): + # issue gh-12927 + # this only applies to in-place floordiv //=, because the output type + # promotes to float which does not fit + a = np.array([1, 2], np.int64) + b = np.array([1, 2], np.uint64) + with pytest.raises(TypeError, + match=r"Cannot cast ufunc 'floor_divide' output from"): + a //= b + + +class TestComplexDivision: + def test_zero_division(self): + with np.errstate(all="ignore"): + for t in [np.complex64, np.complex128]: + a = t(0.0) + b = t(1.0) + assert_(np.isinf(b/a)) + b = t(complex(np.inf, np.inf)) + assert_(np.isinf(b/a)) + b = t(complex(np.inf, np.nan)) + assert_(np.isinf(b/a)) + b = t(complex(np.nan, np.inf)) + assert_(np.isinf(b/a)) + b = t(complex(np.nan, np.nan)) + assert_(np.isnan(b/a)) + b = t(0.) + assert_(np.isnan(b/a)) + + def test_signed_zeros(self): + with np.errstate(all="ignore"): + for t in [np.complex64, np.complex128]: + # tupled (numerator, denominator, expected) + # for testing as expected == numerator/denominator + data = ( + (( 0.0,-1.0), ( 0.0, 1.0), (-1.0,-0.0)), + (( 0.0,-1.0), ( 0.0,-1.0), ( 1.0,-0.0)), + (( 0.0,-1.0), (-0.0,-1.0), ( 1.0, 0.0)), + (( 0.0,-1.0), (-0.0, 1.0), (-1.0, 0.0)), + (( 0.0, 1.0), ( 0.0,-1.0), (-1.0, 0.0)), + (( 0.0,-1.0), ( 0.0,-1.0), ( 1.0,-0.0)), + ((-0.0,-1.0), ( 0.0,-1.0), ( 1.0,-0.0)), + ((-0.0, 1.0), ( 0.0,-1.0), (-1.0,-0.0)) + ) + for cases in data: + n = cases[0] + d = cases[1] + ex = cases[2] + result = t(complex(n[0], n[1])) / t(complex(d[0], d[1])) + # check real and imag parts separately to avoid comparison + # in array context, which does not account for signed zeros + assert_equal(result.real, ex[0]) + assert_equal(result.imag, ex[1]) + + def test_branches(self): + with np.errstate(all="ignore"): + for t in [np.complex64, np.complex128]: + # tupled (numerator, denominator, expected) + # for testing as expected == numerator/denominator + data = list() + + # trigger branch: real(fabs(denom)) > imag(fabs(denom)) + # followed by else condition as neither are == 0 + data.append((( 2.0, 1.0), ( 2.0, 1.0), (1.0, 0.0))) + + # trigger branch: real(fabs(denom)) > imag(fabs(denom)) + # followed by if condition as both are == 0 + # is performed in test_zero_division(), so this is skipped + + # trigger else if branch: real(fabs(denom)) < imag(fabs(denom)) + data.append((( 1.0, 2.0), ( 1.0, 2.0), (1.0, 0.0))) + + for cases in data: + n = cases[0] + d = cases[1] + ex = cases[2] + result = t(complex(n[0], n[1])) / t(complex(d[0], d[1])) + # check real and imag parts separately to avoid comparison + # in array context, which does not account for signed zeros + assert_equal(result.real, ex[0]) + assert_equal(result.imag, ex[1]) + + +class TestConversion: + def test_int_from_long(self): + l = [1e6, 1e12, 1e18, -1e6, -1e12, -1e18] + li = [10**6, 10**12, 10**18, -10**6, -10**12, -10**18] + for T in [None, np.float64, np.int64]: + a = np.array(l, dtype=T) + assert_equal([int(_m) for _m in a], li) + + a = np.array(l[:3], dtype=np.uint64) + assert_equal([int(_m) for _m in a], li[:3]) + + def test_iinfo_long_values(self): + for code in 'bBhH': + with pytest.raises(OverflowError): + np.array(np.iinfo(code).max + 1, dtype=code) + + for code in np.typecodes['AllInteger']: + res = np.array(np.iinfo(code).max, dtype=code) + tgt = np.iinfo(code).max + assert_(res == tgt) + + for code in np.typecodes['AllInteger']: + res = np.dtype(code).type(np.iinfo(code).max) + tgt = np.iinfo(code).max + assert_(res == tgt) + + def test_int_raise_behaviour(self): + def overflow_error_func(dtype): + dtype(np.iinfo(dtype).max + 1) + + for code in [np.int_, np.uint, np.longlong, np.ulonglong]: + assert_raises(OverflowError, overflow_error_func, code) + + def test_int_from_infinite_longdouble(self): + # gh-627 + x = np.longdouble(np.inf) + assert_raises(OverflowError, int, x) + with suppress_warnings() as sup: + sup.record(ComplexWarning) + x = np.clongdouble(np.inf) + assert_raises(OverflowError, int, x) + assert_equal(len(sup.log), 1) + + @pytest.mark.skipif(not IS_PYPY, reason="Test is PyPy only (gh-9972)") + def test_int_from_infinite_longdouble___int__(self): + x = np.longdouble(np.inf) + assert_raises(OverflowError, x.__int__) + with suppress_warnings() as sup: + sup.record(ComplexWarning) + x = np.clongdouble(np.inf) + assert_raises(OverflowError, x.__int__) + assert_equal(len(sup.log), 1) + + @pytest.mark.skipif(np.finfo(np.double) == np.finfo(np.longdouble), + reason="long double is same as double") + @pytest.mark.skipif(platform.machine().startswith("ppc"), + reason="IBM double double") + def test_int_from_huge_longdouble(self): + # Produce a longdouble that would overflow a double, + # use exponent that avoids bug in Darwin pow function. + exp = np.finfo(np.double).maxexp - 1 + huge_ld = 2 * 1234 * np.longdouble(2) ** exp + huge_i = 2 * 1234 * 2 ** exp + assert_(huge_ld != np.inf) + assert_equal(int(huge_ld), huge_i) + + def test_int_from_longdouble(self): + x = np.longdouble(1.5) + assert_equal(int(x), 1) + x = np.longdouble(-10.5) + assert_equal(int(x), -10) + + def test_numpy_scalar_relational_operators(self): + # All integer + for dt1 in np.typecodes['AllInteger']: + assert_(1 > np.array(0, dtype=dt1)[()], "type %s failed" % (dt1,)) + assert_(not 1 < np.array(0, dtype=dt1)[()], "type %s failed" % (dt1,)) + + for dt2 in np.typecodes['AllInteger']: + assert_(np.array(1, dtype=dt1)[()] > np.array(0, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + assert_(not np.array(1, dtype=dt1)[()] < np.array(0, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + + #Unsigned integers + for dt1 in 'BHILQP': + assert_(-1 < np.array(1, dtype=dt1)[()], "type %s failed" % (dt1,)) + assert_(not -1 > np.array(1, dtype=dt1)[()], "type %s failed" % (dt1,)) + assert_(-1 != np.array(1, dtype=dt1)[()], "type %s failed" % (dt1,)) + + #unsigned vs signed + for dt2 in 'bhilqp': + assert_(np.array(1, dtype=dt1)[()] > np.array(-1, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + assert_(not np.array(1, dtype=dt1)[()] < np.array(-1, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + assert_(np.array(1, dtype=dt1)[()] != np.array(-1, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + + #Signed integers and floats + for dt1 in 'bhlqp' + np.typecodes['Float']: + assert_(1 > np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,)) + assert_(not 1 < np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,)) + assert_(-1 == np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,)) + + for dt2 in 'bhlqp' + np.typecodes['Float']: + assert_(np.array(1, dtype=dt1)[()] > np.array(-1, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + assert_(not np.array(1, dtype=dt1)[()] < np.array(-1, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + assert_(np.array(-1, dtype=dt1)[()] == np.array(-1, dtype=dt2)[()], + "type %s and %s failed" % (dt1, dt2)) + + def test_scalar_comparison_to_none(self): + # Scalars should just return False and not give a warnings. + # The comparisons are flagged by pep8, ignore that. + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', FutureWarning) + assert_(not np.float32(1) == None) # noqa: E711 + assert_(not np.str_('test') == None) # noqa: E711 + # This is dubious (see below): + assert_(not np.datetime64('NaT') == None) # noqa: E711 + + assert_(np.float32(1) != None) # noqa: E711 + assert_(np.str_('test') != None) # noqa: E711 + # This is dubious (see below): + assert_(np.datetime64('NaT') != None) # noqa: E711 + assert_(len(w) == 0) + + # For documentation purposes, this is why the datetime is dubious. + # At the time of deprecation this was no behaviour change, but + # it has to be considered when the deprecations are done. + assert_(np.equal(np.datetime64('NaT'), None)) + + +#class TestRepr: +# def test_repr(self): +# for t in types: +# val = t(1197346475.0137341) +# val_repr = repr(val) +# val2 = eval(val_repr) +# assert_equal( val, val2 ) + + +class TestRepr: + def _test_type_repr(self, t): + finfo = np.finfo(t) + last_fraction_bit_idx = finfo.nexp + finfo.nmant + last_exponent_bit_idx = finfo.nexp + storage_bytes = np.dtype(t).itemsize*8 + # could add some more types to the list below + for which in ['small denorm', 'small norm']: + # Values from https://en.wikipedia.org/wiki/IEEE_754 + constr = np.array([0x00]*storage_bytes, dtype=np.uint8) + if which == 'small denorm': + byte = last_fraction_bit_idx // 8 + bytebit = 7-(last_fraction_bit_idx % 8) + constr[byte] = 1 << bytebit + elif which == 'small norm': + byte = last_exponent_bit_idx // 8 + bytebit = 7-(last_exponent_bit_idx % 8) + constr[byte] = 1 << bytebit + else: + raise ValueError('hmm') + val = constr.view(t)[0] + val_repr = repr(val) + val2 = t(eval(val_repr)) + if not (val2 == 0 and val < 1e-100): + assert_equal(val, val2) + + def test_float_repr(self): + # long double test cannot work, because eval goes through a python + # float + for t in [np.float32, np.float64]: + self._test_type_repr(t) + + +if not IS_PYPY: + # sys.getsizeof() is not valid on PyPy + class TestSizeOf: + + def test_equal_nbytes(self): + for type in types: + x = type(0) + assert_(sys.getsizeof(x) > x.nbytes) + + def test_error(self): + d = np.float32() + assert_raises(TypeError, d.__sizeof__, "a") + + +class TestMultiply: + def test_seq_repeat(self): + # Test that basic sequences get repeated when multiplied with + # numpy integers. And errors are raised when multiplied with others. + # Some of this behaviour may be controversial and could be open for + # change. + accepted_types = set(np.typecodes["AllInteger"]) + deprecated_types = {'?'} + forbidden_types = ( + set(np.typecodes["All"]) - accepted_types - deprecated_types) + forbidden_types -= {'V'} # can't default-construct void scalars + + for seq_type in (list, tuple): + seq = seq_type([1, 2, 3]) + for numpy_type in accepted_types: + i = np.dtype(numpy_type).type(2) + assert_equal(seq * i, seq * int(i)) + assert_equal(i * seq, int(i) * seq) + + for numpy_type in deprecated_types: + i = np.dtype(numpy_type).type() + assert_equal( + assert_warns(DeprecationWarning, operator.mul, seq, i), + seq * int(i)) + assert_equal( + assert_warns(DeprecationWarning, operator.mul, i, seq), + int(i) * seq) + + for numpy_type in forbidden_types: + i = np.dtype(numpy_type).type() + assert_raises(TypeError, operator.mul, seq, i) + assert_raises(TypeError, operator.mul, i, seq) + + def test_no_seq_repeat_basic_array_like(self): + # Test that an array-like which does not know how to be multiplied + # does not attempt sequence repeat (raise TypeError). + # See also gh-7428. + class ArrayLike: + def __init__(self, arr): + self.arr = arr + + def __array__(self, dtype=None, copy=None): + return self.arr + + # Test for simple ArrayLike above and memoryviews (original report) + for arr_like in (ArrayLike(np.ones(3)), memoryview(np.ones(3))): + assert_array_equal(arr_like * np.float32(3.), np.full(3, 3.)) + assert_array_equal(np.float32(3.) * arr_like, np.full(3, 3.)) + assert_array_equal(arr_like * np.int_(3), np.full(3, 3)) + assert_array_equal(np.int_(3) * arr_like, np.full(3, 3)) + + +class TestNegative: + def test_exceptions(self): + a = np.ones((), dtype=np.bool)[()] + assert_raises(TypeError, operator.neg, a) + + def test_result(self): + types = np.typecodes['AllInteger'] + np.typecodes['AllFloat'] + with suppress_warnings() as sup: + sup.filter(RuntimeWarning) + for dt in types: + a = np.ones((), dtype=dt)[()] + if dt in np.typecodes['UnsignedInteger']: + st = np.dtype(dt).type + max = st(np.iinfo(dt).max) + assert_equal(operator.neg(a), max) + else: + assert_equal(operator.neg(a) + a, 0) + +class TestSubtract: + def test_exceptions(self): + a = np.ones((), dtype=np.bool)[()] + assert_raises(TypeError, operator.sub, a, a) + + def test_result(self): + types = np.typecodes['AllInteger'] + np.typecodes['AllFloat'] + with suppress_warnings() as sup: + sup.filter(RuntimeWarning) + for dt in types: + a = np.ones((), dtype=dt)[()] + assert_equal(operator.sub(a, a), 0) + + +class TestAbs: + def _test_abs_func(self, absfunc, test_dtype): + x = test_dtype(-1.5) + assert_equal(absfunc(x), 1.5) + x = test_dtype(0.0) + res = absfunc(x) + # assert_equal() checks zero signedness + assert_equal(res, 0.0) + x = test_dtype(-0.0) + res = absfunc(x) + assert_equal(res, 0.0) + + x = test_dtype(np.finfo(test_dtype).max) + assert_equal(absfunc(x), x.real) + + with suppress_warnings() as sup: + sup.filter(UserWarning) + x = test_dtype(np.finfo(test_dtype).tiny) + assert_equal(absfunc(x), x.real) + + x = test_dtype(np.finfo(test_dtype).min) + assert_equal(absfunc(x), -x.real) + + @pytest.mark.parametrize("dtype", floating_types + complex_floating_types) + def test_builtin_abs(self, dtype): + if ( + sys.platform == "cygwin" and dtype == np.clongdouble and + ( + _pep440.parse(platform.release().split("-")[0]) + < _pep440.Version("3.3.0") + ) + ): + pytest.xfail( + reason="absl is computed in double precision on cygwin < 3.3" + ) + self._test_abs_func(abs, dtype) + + @pytest.mark.parametrize("dtype", floating_types + complex_floating_types) + def test_numpy_abs(self, dtype): + if ( + sys.platform == "cygwin" and dtype == np.clongdouble and + ( + _pep440.parse(platform.release().split("-")[0]) + < _pep440.Version("3.3.0") + ) + ): + pytest.xfail( + reason="absl is computed in double precision on cygwin < 3.3" + ) + self._test_abs_func(np.abs, dtype) + +class TestBitShifts: + + @pytest.mark.parametrize('type_code', np.typecodes['AllInteger']) + @pytest.mark.parametrize('op', + [operator.rshift, operator.lshift], ids=['>>', '<<']) + def test_shift_all_bits(self, type_code, op): + """Shifts where the shift amount is the width of the type or wider """ + # gh-2449 + dt = np.dtype(type_code) + nbits = dt.itemsize * 8 + for val in [5, -5]: + for shift in [nbits, nbits + 4]: + val_scl = np.array(val).astype(dt)[()] + shift_scl = dt.type(shift) + res_scl = op(val_scl, shift_scl) + if val_scl < 0 and op is operator.rshift: + # sign bit is preserved + assert_equal(res_scl, -1) + else: + assert_equal(res_scl, 0) + + # Result on scalars should be the same as on arrays + val_arr = np.array([val_scl]*32, dtype=dt) + shift_arr = np.array([shift]*32, dtype=dt) + res_arr = op(val_arr, shift_arr) + assert_equal(res_arr, res_scl) + + +class TestHash: + @pytest.mark.parametrize("type_code", np.typecodes['AllInteger']) + def test_integer_hashes(self, type_code): + scalar = np.dtype(type_code).type + for i in range(128): + assert hash(i) == hash(scalar(i)) + + @pytest.mark.parametrize("type_code", np.typecodes['AllFloat']) + def test_float_and_complex_hashes(self, type_code): + scalar = np.dtype(type_code).type + for val in [np.pi, np.inf, 3, 6.]: + numpy_val = scalar(val) + # Cast back to Python, in case the NumPy scalar has less precision + if numpy_val.dtype.kind == 'c': + val = complex(numpy_val) + else: + val = float(numpy_val) + assert val == numpy_val + assert hash(val) == hash(numpy_val) + + if hash(float(np.nan)) != hash(float(np.nan)): + # If Python distinguishes different NaNs we do so too (gh-18833) + assert hash(scalar(np.nan)) != hash(scalar(np.nan)) + + @pytest.mark.parametrize("type_code", np.typecodes['Complex']) + def test_complex_hashes(self, type_code): + # Test some complex valued hashes specifically: + scalar = np.dtype(type_code).type + for val in [np.pi+1j, np.inf-3j, 3j, 6.+1j]: + numpy_val = scalar(val) + assert hash(complex(numpy_val)) == hash(numpy_val) + + +@contextlib.contextmanager +def recursionlimit(n): + o = sys.getrecursionlimit() + try: + sys.setrecursionlimit(n) + yield + finally: + sys.setrecursionlimit(o) + + +@given(sampled_from(objecty_things), + sampled_from(binary_operators_for_scalar_ints), + sampled_from(types + [rational])) +def test_operator_object_left(o, op, type_): + try: + with recursionlimit(200): + op(o, type_(1)) + except TypeError: + pass + + +@given(sampled_from(objecty_things), + sampled_from(binary_operators_for_scalar_ints), + sampled_from(types + [rational])) +def test_operator_object_right(o, op, type_): + try: + with recursionlimit(200): + op(type_(1), o) + except TypeError: + pass + + +@given(sampled_from(binary_operators_for_scalars), + sampled_from(types), + sampled_from(types)) +def test_operator_scalars(op, type1, type2): + try: + op(type1(1), type2(1)) + except TypeError: + pass + + +@pytest.mark.parametrize("op", binary_operators_for_scalars) +@pytest.mark.parametrize("sctype", [np.longdouble, np.clongdouble]) +def test_longdouble_operators_with_obj(sctype, op): + # This is/used to be tricky, because NumPy generally falls back to + # using the ufunc via `np.asarray()`, this effectively might do: + # longdouble + None + # -> asarray(longdouble) + np.array(None, dtype=object) + # -> asarray(longdouble).astype(object) + np.array(None, dtype=object) + # And after getting the scalars in the inner loop: + # -> longdouble + None + # + # That would recurse infinitely. Other scalars return the python object + # on cast, so this type of things works OK. + # + # As of NumPy 2.1, this has been consolidated into the np.generic binops + # and now checks `.item()`. That also allows the below path to work now. + try: + op(sctype(3), None) + except TypeError: + pass + try: + op(None, sctype(3)) + except TypeError: + pass + + +@pytest.mark.parametrize("op", [operator.add, operator.pow, operator.sub]) +@pytest.mark.parametrize("sctype", [np.longdouble, np.clongdouble]) +def test_longdouble_with_arrlike(sctype, op): + # As of NumPy 2.1, longdouble behaves like other types and can coerce + # e.g. lists. (Not necessarily better, but consistent.) + assert_array_equal(op(sctype(3), [1, 2]), op(3, np.array([1, 2]))) + assert_array_equal(op([1, 2], sctype(3)), op(np.array([1, 2]), 3)) + + +@pytest.mark.parametrize("op", binary_operators_for_scalars) +@pytest.mark.parametrize("sctype", [np.longdouble, np.clongdouble]) +@np.errstate(all="ignore") +def test_longdouble_operators_with_large_int(sctype, op): + # (See `test_longdouble_operators_with_obj` for why longdouble is special) + # NEP 50 means that the result is clearly a (c)longdouble here: + if sctype == np.clongdouble and op in [operator.mod, operator.floordiv]: + # The above operators are not support for complex though... + with pytest.raises(TypeError): + op(sctype(3), 2**64) + with pytest.raises(TypeError): + op(sctype(3), 2**64) + else: + assert op(sctype(3), -2**64) == op(sctype(3), sctype(-2**64)) + assert op(2**64, sctype(3)) == op(sctype(2**64), sctype(3)) + + +@pytest.mark.parametrize("dtype", np.typecodes["AllInteger"]) +@pytest.mark.parametrize("operation", [ + lambda min, max: max + max, + lambda min, max: min - max, + lambda min, max: max * max], ids=["+", "-", "*"]) +def test_scalar_integer_operation_overflow(dtype, operation): + st = np.dtype(dtype).type + min = st(np.iinfo(dtype).min) + max = st(np.iinfo(dtype).max) + + with pytest.warns(RuntimeWarning, match="overflow encountered"): + operation(min, max) + + +@pytest.mark.parametrize("dtype", np.typecodes["Integer"]) +@pytest.mark.parametrize("operation", [ + lambda min, neg_1: -min, + lambda min, neg_1: abs(min), + lambda min, neg_1: min * neg_1, + pytest.param(lambda min, neg_1: min // neg_1, + marks=pytest.mark.skip(reason="broken on some platforms"))], + ids=["neg", "abs", "*", "//"]) +def test_scalar_signed_integer_overflow(dtype, operation): + # The minimum signed integer can "overflow" for some additional operations + st = np.dtype(dtype).type + min = st(np.iinfo(dtype).min) + neg_1 = st(-1) + + with pytest.warns(RuntimeWarning, match="overflow encountered"): + operation(min, neg_1) + + +@pytest.mark.parametrize("dtype", np.typecodes["UnsignedInteger"]) +def test_scalar_unsigned_integer_overflow(dtype): + val = np.dtype(dtype).type(8) + with pytest.warns(RuntimeWarning, match="overflow encountered"): + -val + + zero = np.dtype(dtype).type(0) + -zero # does not warn + +@pytest.mark.parametrize("dtype", np.typecodes["AllInteger"]) +@pytest.mark.parametrize("operation", [ + lambda val, zero: val // zero, + lambda val, zero: val % zero, ], ids=["//", "%"]) +def test_scalar_integer_operation_divbyzero(dtype, operation): + st = np.dtype(dtype).type + val = st(100) + zero = st(0) + + with pytest.warns(RuntimeWarning, match="divide by zero"): + operation(val, zero) + + +ops_with_names = [ + ("__lt__", "__gt__", operator.lt, True), + ("__le__", "__ge__", operator.le, True), + ("__eq__", "__eq__", operator.eq, True), + # Note __op__ and __rop__ may be identical here: + ("__ne__", "__ne__", operator.ne, True), + ("__gt__", "__lt__", operator.gt, True), + ("__ge__", "__le__", operator.ge, True), + ("__floordiv__", "__rfloordiv__", operator.floordiv, False), + ("__truediv__", "__rtruediv__", operator.truediv, False), + ("__add__", "__radd__", operator.add, False), + ("__mod__", "__rmod__", operator.mod, False), + ("__mul__", "__rmul__", operator.mul, False), + ("__pow__", "__rpow__", operator.pow, False), + ("__sub__", "__rsub__", operator.sub, False), +] + + +@pytest.mark.parametrize(["__op__", "__rop__", "op", "cmp"], ops_with_names) +@pytest.mark.parametrize("sctype", [np.float32, np.float64, np.longdouble]) +def test_subclass_deferral(sctype, __op__, __rop__, op, cmp): + """ + This test covers scalar subclass deferral. Note that this is exceedingly + complicated, especially since it tends to fall back to the array paths and + these additionally add the "array priority" mechanism. + + The behaviour was modified subtly in 1.22 (to make it closer to how Python + scalars work). Due to its complexity and the fact that subclassing NumPy + scalars is probably a bad idea to begin with. There is probably room + for adjustments here. + """ + class myf_simple1(sctype): + pass + + class myf_simple2(sctype): + pass + + def op_func(self, other): + return __op__ + + def rop_func(self, other): + return __rop__ + + myf_op = type("myf_op", (sctype,), {__op__: op_func, __rop__: rop_func}) + + # inheritance has to override, or this is correctly lost: + res = op(myf_simple1(1), myf_simple2(2)) + assert type(res) == sctype or type(res) == np.bool + assert op(myf_simple1(1), myf_simple2(2)) == op(1, 2) # inherited + + # Two independent subclasses do not really define an order. This could + # be attempted, but we do not since Python's `int` does neither: + assert op(myf_op(1), myf_simple1(2)) == __op__ + assert op(myf_simple1(1), myf_op(2)) == op(1, 2) # inherited + + +def test_longdouble_complex(): + # Simple test to check longdouble and complex combinations, since these + # need to go through promotion, which longdouble needs to be careful about. + x = np.longdouble(1) + assert x + 1j == 1+1j + assert 1j + x == 1+1j + + +@pytest.mark.parametrize(["__op__", "__rop__", "op", "cmp"], ops_with_names) +@pytest.mark.parametrize("subtype", [float, int, complex, np.float16]) +def test_pyscalar_subclasses(subtype, __op__, __rop__, op, cmp): + # This tests that python scalar subclasses behave like a float64 (if they + # don't override it). + # In an earlier version of NEP 50, they behaved like the Python buildins. + def op_func(self, other): + return __op__ + + def rop_func(self, other): + return __rop__ + + # Check that deferring is indicated using `__array_ufunc__`: + myt = type("myt", (subtype,), + {__op__: op_func, __rop__: rop_func, "__array_ufunc__": None}) + + # Just like normally, we should never presume we can modify the float. + assert op(myt(1), np.float64(2)) == __op__ + assert op(np.float64(1), myt(2)) == __rop__ + + if op in {operator.mod, operator.floordiv} and subtype == complex: + return # module is not support for complex. Do not test. + + if __rop__ == __op__: + return + + # When no deferring is indicated, subclasses are handled normally. + myt = type("myt", (subtype,), {__rop__: rop_func}) + behaves_like = lambda x: np.array(subtype(x))[()] + + # Check for float32, as a float subclass float64 may behave differently + res = op(myt(1), np.float16(2)) + expected = op(behaves_like(1), np.float16(2)) + assert res == expected + assert type(res) == type(expected) + res = op(np.float32(2), myt(1)) + expected = op(np.float32(2), behaves_like(1)) + assert res == expected + assert type(res) == type(expected) + + # Same check for longdouble (compare via dtype to accept float64 when + # longdouble has the identical size), which is currently not perfectly + # consistent. + res = op(myt(1), np.longdouble(2)) + expected = op(behaves_like(1), np.longdouble(2)) + assert res == expected + assert np.dtype(type(res)) == np.dtype(type(expected)) + res = op(np.float32(2), myt(1)) + expected = op(np.float32(2), behaves_like(1)) + assert res == expected + assert np.dtype(type(res)) == np.dtype(type(expected)) + + +def test_truediv_int(): + # This should work, as the result is float: + assert np.uint8(3) / 123454 == np.float64(3) / 123454 + + +@pytest.mark.slow +@pytest.mark.parametrize("op", + # TODO: Power is a bit special, but here mostly bools seem to behave oddly + [op for op in binary_operators_for_scalars if op is not operator.pow]) +@pytest.mark.parametrize("sctype", types) +@pytest.mark.parametrize("other_type", [float, int, complex]) +@pytest.mark.parametrize("rop", [True, False]) +def test_scalar_matches_array_op_with_pyscalar(op, sctype, other_type, rop): + # Check that the ufunc path matches by coercing to an array explicitly + val1 = sctype(2) + val2 = other_type(2) + + if rop: + _op = op + op = lambda x, y: _op(y, x) + + try: + res = op(val1, val2) + except TypeError: + try: + expected = op(np.asarray(val1), val2) + raise AssertionError("ufunc didn't raise.") + except TypeError: + return + else: + expected = op(np.asarray(val1), val2) + + # Note that we only check dtype equivalency, as ufuncs may pick the lower + # dtype if they are equivalent. + assert res == expected + if isinstance(val1, float) and other_type is complex and rop: + # Python complex accepts float subclasses, so we don't get a chance + # and the result may be a Python complex (thus, the `np.array()``) + assert np.array(res).dtype == expected.dtype + else: + assert res.dtype == expected.dtype diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarprint.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarprint.py new file mode 100644 index 0000000000000000000000000000000000000000..b109ce7303a355cd9fa0beea2d0178b910a48535 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_scalarprint.py @@ -0,0 +1,423 @@ +""" Test printing of scalar types. + +""" +import code +import platform +import pytest +import sys + +from tempfile import TemporaryFile +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_raises_regex, IS_MUSL) + +class TestRealScalars: + def test_str(self): + svals = [0.0, -0.0, 1, -1, np.inf, -np.inf, np.nan] + styps = [np.float16, np.float32, np.float64, np.longdouble] + wanted = [ + ['0.0', '0.0', '0.0', '0.0' ], + ['-0.0', '-0.0', '-0.0', '-0.0'], + ['1.0', '1.0', '1.0', '1.0' ], + ['-1.0', '-1.0', '-1.0', '-1.0'], + ['inf', 'inf', 'inf', 'inf' ], + ['-inf', '-inf', '-inf', '-inf'], + ['nan', 'nan', 'nan', 'nan']] + + for wants, val in zip(wanted, svals): + for want, styp in zip(wants, styps): + msg = 'for str({}({}))'.format(np.dtype(styp).name, repr(val)) + assert_equal(str(styp(val)), want, err_msg=msg) + + def test_scalar_cutoffs(self): + # test that both the str and repr of np.float64 behaves + # like python floats in python3. + def check(v): + assert_equal(str(np.float64(v)), str(v)) + assert_equal(str(np.float64(v)), repr(v)) + assert_equal(repr(np.float64(v)), f"np.float64({v!r})") + assert_equal(repr(np.float64(v)), f"np.float64({v})") + + # check we use the same number of significant digits + check(1.12345678901234567890) + check(0.0112345678901234567890) + + # check switch from scientific output to positional and back + check(1e-5) + check(1e-4) + check(1e15) + check(1e16) + + def test_py2_float_print(self): + # gh-10753 + # In python2, the python float type implements an obsolete method + # tp_print, which overrides tp_repr and tp_str when using "print" to + # output to a "real file" (ie, not a StringIO). Make sure we don't + # inherit it. + x = np.double(0.1999999999999) + with TemporaryFile('r+t') as f: + print(x, file=f) + f.seek(0) + output = f.read() + assert_equal(output, str(x) + '\n') + # In python2 the value float('0.1999999999999') prints with reduced + # precision as '0.2', but we want numpy's np.double('0.1999999999999') + # to print the unique value, '0.1999999999999'. + + # gh-11031 + # Only in the python2 interactive shell and when stdout is a "real" + # file, the output of the last command is printed to stdout without + # Py_PRINT_RAW (unlike the print statement) so `>>> x` and `>>> print + # x` are potentially different. Make sure they are the same. The only + # way I found to get prompt-like output is using an actual prompt from + # the 'code' module. Again, must use tempfile to get a "real" file. + + # dummy user-input which enters one line and then ctrl-Ds. + def userinput(): + yield 'np.sqrt(2)' + raise EOFError + gen = userinput() + input_func = lambda prompt="": next(gen) + + with TemporaryFile('r+t') as fo, TemporaryFile('r+t') as fe: + orig_stdout, orig_stderr = sys.stdout, sys.stderr + sys.stdout, sys.stderr = fo, fe + + code.interact(local={'np': np}, readfunc=input_func, banner='') + + sys.stdout, sys.stderr = orig_stdout, orig_stderr + + fo.seek(0) + capture = fo.read().strip() + + assert_equal(capture, repr(np.sqrt(2))) + + def test_dragon4(self): + # these tests are adapted from Ryan Juckett's dragon4 implementation, + # see dragon4.c for details. + + fpos32 = lambda x, **k: np.format_float_positional(np.float32(x), **k) + fsci32 = lambda x, **k: np.format_float_scientific(np.float32(x), **k) + fpos64 = lambda x, **k: np.format_float_positional(np.float64(x), **k) + fsci64 = lambda x, **k: np.format_float_scientific(np.float64(x), **k) + + preckwd = lambda prec: {'unique': False, 'precision': prec} + + assert_equal(fpos32('1.0'), "1.") + assert_equal(fsci32('1.0'), "1.e+00") + assert_equal(fpos32('10.234'), "10.234") + assert_equal(fpos32('-10.234'), "-10.234") + assert_equal(fsci32('10.234'), "1.0234e+01") + assert_equal(fsci32('-10.234'), "-1.0234e+01") + assert_equal(fpos32('1000.0'), "1000.") + assert_equal(fpos32('1.0', precision=0), "1.") + assert_equal(fsci32('1.0', precision=0), "1.e+00") + assert_equal(fpos32('10.234', precision=0), "10.") + assert_equal(fpos32('-10.234', precision=0), "-10.") + assert_equal(fsci32('10.234', precision=0), "1.e+01") + assert_equal(fsci32('-10.234', precision=0), "-1.e+01") + assert_equal(fpos32('10.234', precision=2), "10.23") + assert_equal(fsci32('-10.234', precision=2), "-1.02e+01") + assert_equal(fsci64('9.9999999999999995e-08', **preckwd(16)), + '9.9999999999999995e-08') + assert_equal(fsci64('9.8813129168249309e-324', **preckwd(16)), + '9.8813129168249309e-324') + assert_equal(fsci64('9.9999999999999694e-311', **preckwd(16)), + '9.9999999999999694e-311') + + + # test rounding + # 3.1415927410 is closest float32 to np.pi + assert_equal(fpos32('3.14159265358979323846', **preckwd(10)), + "3.1415927410") + assert_equal(fsci32('3.14159265358979323846', **preckwd(10)), + "3.1415927410e+00") + assert_equal(fpos64('3.14159265358979323846', **preckwd(10)), + "3.1415926536") + assert_equal(fsci64('3.14159265358979323846', **preckwd(10)), + "3.1415926536e+00") + # 299792448 is closest float32 to 299792458 + assert_equal(fpos32('299792458.0', **preckwd(5)), "299792448.00000") + assert_equal(fsci32('299792458.0', **preckwd(5)), "2.99792e+08") + assert_equal(fpos64('299792458.0', **preckwd(5)), "299792458.00000") + assert_equal(fsci64('299792458.0', **preckwd(5)), "2.99792e+08") + + assert_equal(fpos32('3.14159265358979323846', **preckwd(25)), + "3.1415927410125732421875000") + assert_equal(fpos64('3.14159265358979323846', **preckwd(50)), + "3.14159265358979311599796346854418516159057617187500") + assert_equal(fpos64('3.14159265358979323846'), "3.141592653589793") + + + # smallest numbers + assert_equal(fpos32(0.5**(126 + 23), unique=False, precision=149), + "0.00000000000000000000000000000000000000000000140129846432" + "4817070923729583289916131280261941876515771757068283889791" + "08268586060148663818836212158203125") + + assert_equal(fpos64(5e-324, unique=False, precision=1074), + "0.00000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000000049406564584124654417656" + "8792868221372365059802614324764425585682500675507270208751" + "8652998363616359923797965646954457177309266567103559397963" + "9877479601078187812630071319031140452784581716784898210368" + "8718636056998730723050006387409153564984387312473397273169" + "6151400317153853980741262385655911710266585566867681870395" + "6031062493194527159149245532930545654440112748012970999954" + "1931989409080416563324524757147869014726780159355238611550" + "1348035264934720193790268107107491703332226844753335720832" + "4319360923828934583680601060115061698097530783422773183292" + "4790498252473077637592724787465608477820373446969953364701" + "7972677717585125660551199131504891101451037862738167250955" + "8373897335989936648099411642057026370902792427675445652290" + "87538682506419718265533447265625") + + # largest numbers + f32x = np.finfo(np.float32).max + assert_equal(fpos32(f32x, **preckwd(0)), + "340282346638528859811704183484516925440.") + assert_equal(fpos64(np.finfo(np.float64).max, **preckwd(0)), + "1797693134862315708145274237317043567980705675258449965989" + "1747680315726078002853876058955863276687817154045895351438" + "2464234321326889464182768467546703537516986049910576551282" + "0762454900903893289440758685084551339423045832369032229481" + "6580855933212334827479782620414472316873817718091929988125" + "0404026184124858368.") + # Warning: In unique mode only the integer digits necessary for + # uniqueness are computed, the rest are 0. + assert_equal(fpos32(f32x), + "340282350000000000000000000000000000000.") + + # Further tests of zero-padding vs rounding in different combinations + # of unique, fractional, precision, min_digits + # precision can only reduce digits, not add them. + # min_digits can only extend digits, not reduce them. + assert_equal(fpos32(f32x, unique=True, fractional=True, precision=0), + "340282350000000000000000000000000000000.") + assert_equal(fpos32(f32x, unique=True, fractional=True, precision=4), + "340282350000000000000000000000000000000.") + assert_equal(fpos32(f32x, unique=True, fractional=True, min_digits=0), + "340282346638528859811704183484516925440.") + assert_equal(fpos32(f32x, unique=True, fractional=True, min_digits=4), + "340282346638528859811704183484516925440.0000") + assert_equal(fpos32(f32x, unique=True, fractional=True, + min_digits=4, precision=4), + "340282346638528859811704183484516925440.0000") + assert_raises(ValueError, fpos32, f32x, unique=True, fractional=False, + precision=0) + assert_equal(fpos32(f32x, unique=True, fractional=False, precision=4), + "340300000000000000000000000000000000000.") + assert_equal(fpos32(f32x, unique=True, fractional=False, precision=20), + "340282350000000000000000000000000000000.") + assert_equal(fpos32(f32x, unique=True, fractional=False, min_digits=4), + "340282350000000000000000000000000000000.") + assert_equal(fpos32(f32x, unique=True, fractional=False, + min_digits=20), + "340282346638528859810000000000000000000.") + assert_equal(fpos32(f32x, unique=True, fractional=False, + min_digits=15), + "340282346638529000000000000000000000000.") + assert_equal(fpos32(f32x, unique=False, fractional=False, precision=4), + "340300000000000000000000000000000000000.") + # test that unique rounding is preserved when precision is supplied + # but no extra digits need to be printed (gh-18609) + a = np.float64.fromhex('-1p-97') + assert_equal(fsci64(a, unique=True), '-6.310887241768095e-30') + assert_equal(fsci64(a, unique=False, precision=15), + '-6.310887241768094e-30') + assert_equal(fsci64(a, unique=True, precision=15), + '-6.310887241768095e-30') + assert_equal(fsci64(a, unique=True, min_digits=15), + '-6.310887241768095e-30') + assert_equal(fsci64(a, unique=True, precision=15, min_digits=15), + '-6.310887241768095e-30') + # adds/remove digits in unique mode with unbiased rnding + assert_equal(fsci64(a, unique=True, precision=14), + '-6.31088724176809e-30') + assert_equal(fsci64(a, unique=True, min_digits=16), + '-6.3108872417680944e-30') + assert_equal(fsci64(a, unique=True, precision=16), + '-6.310887241768095e-30') + assert_equal(fsci64(a, unique=True, min_digits=14), + '-6.310887241768095e-30') + # test min_digits in unique mode with different rounding cases + assert_equal(fsci64('1e120', min_digits=3), '1.000e+120') + assert_equal(fsci64('1e100', min_digits=3), '1.000e+100') + + # test trailing zeros + assert_equal(fpos32('1.0', unique=False, precision=3), "1.000") + assert_equal(fpos64('1.0', unique=False, precision=3), "1.000") + assert_equal(fsci32('1.0', unique=False, precision=3), "1.000e+00") + assert_equal(fsci64('1.0', unique=False, precision=3), "1.000e+00") + assert_equal(fpos32('1.5', unique=False, precision=3), "1.500") + assert_equal(fpos64('1.5', unique=False, precision=3), "1.500") + assert_equal(fsci32('1.5', unique=False, precision=3), "1.500e+00") + assert_equal(fsci64('1.5', unique=False, precision=3), "1.500e+00") + # gh-10713 + assert_equal(fpos64('324', unique=False, precision=5, + fractional=False), "324.00") + + available_float_dtypes = [np.float16, np.float32, np.float64, np.float128]\ + if hasattr(np, 'float128') else [np.float16, np.float32, np.float64] + + @pytest.mark.parametrize("tp", available_float_dtypes) + def test_dragon4_positional_interface(self, tp): + # test is flaky for musllinux on np.float128 + if IS_MUSL and tp == np.float128: + pytest.skip("Skipping flaky test of float128 on musllinux") + + fpos = np.format_float_positional + + # test padding + assert_equal(fpos(tp('1.0'), pad_left=4, pad_right=4), " 1. ") + assert_equal(fpos(tp('-1.0'), pad_left=4, pad_right=4), " -1. ") + assert_equal(fpos(tp('-10.2'), + pad_left=4, pad_right=4), " -10.2 ") + + # test fixed (non-unique) mode + assert_equal(fpos(tp('1.0'), unique=False, precision=4), "1.0000") + + @pytest.mark.parametrize("tp", available_float_dtypes) + def test_dragon4_positional_interface_trim(self, tp): + # test is flaky for musllinux on np.float128 + if IS_MUSL and tp == np.float128: + pytest.skip("Skipping flaky test of float128 on musllinux") + + fpos = np.format_float_positional + # test trimming + # trim of 'k' or '.' only affects non-unique mode, since unique + # mode will not output trailing 0s. + assert_equal(fpos(tp('1.'), unique=False, precision=4, trim='k'), + "1.0000") + + assert_equal(fpos(tp('1.'), unique=False, precision=4, trim='.'), + "1.") + assert_equal(fpos(tp('1.2'), unique=False, precision=4, trim='.'), + "1.2" if tp != np.float16 else "1.2002") + + assert_equal(fpos(tp('1.'), unique=False, precision=4, trim='0'), + "1.0") + assert_equal(fpos(tp('1.2'), unique=False, precision=4, trim='0'), + "1.2" if tp != np.float16 else "1.2002") + assert_equal(fpos(tp('1.'), trim='0'), "1.0") + + assert_equal(fpos(tp('1.'), unique=False, precision=4, trim='-'), + "1") + assert_equal(fpos(tp('1.2'), unique=False, precision=4, trim='-'), + "1.2" if tp != np.float16 else "1.2002") + assert_equal(fpos(tp('1.'), trim='-'), "1") + assert_equal(fpos(tp('1.001'), precision=1, trim='-'), "1") + + @pytest.mark.parametrize("tp", available_float_dtypes) + @pytest.mark.parametrize("pad_val", [10**5, np.iinfo("int32").max]) + def test_dragon4_positional_interface_overflow(self, tp, pad_val): + # test is flaky for musllinux on np.float128 + if IS_MUSL and tp == np.float128: + pytest.skip("Skipping flaky test of float128 on musllinux") + + fpos = np.format_float_positional + + #gh-28068 + with pytest.raises(RuntimeError, + match="Float formating result too large"): + fpos(tp('1.047'), unique=False, precision=pad_val) + + with pytest.raises(RuntimeError, + match="Float formating result too large"): + fpos(tp('1.047'), precision=2, pad_left=pad_val) + + with pytest.raises(RuntimeError, + match="Float formating result too large"): + fpos(tp('1.047'), precision=2, pad_right=pad_val) + + @pytest.mark.parametrize("tp", available_float_dtypes) + def test_dragon4_scientific_interface(self, tp): + # test is flaky for musllinux on np.float128 + if IS_MUSL and tp == np.float128: + pytest.skip("Skipping flaky test of float128 on musllinux") + + fsci = np.format_float_scientific + + # test exp_digits + assert_equal(fsci(tp('1.23e1'), exp_digits=5), "1.23e+00001") + + # test fixed (non-unique) mode + assert_equal(fsci(tp('1.0'), unique=False, precision=4), + "1.0000e+00") + + @pytest.mark.skipif(not platform.machine().startswith("ppc64"), + reason="only applies to ppc float128 values") + def test_ppc64_ibm_double_double128(self): + # check that the precision decreases once we get into the subnormal + # range. Unlike float64, this starts around 1e-292 instead of 1e-308, + # which happens when the first double is normal and the second is + # subnormal. + x = np.float128('2.123123123123123123123123123123123e-286') + got = [str(x/np.float128('2e' + str(i))) for i in range(0,40)] + expected = [ + "1.06156156156156156156156156156157e-286", + "1.06156156156156156156156156156158e-287", + "1.06156156156156156156156156156159e-288", + "1.0615615615615615615615615615616e-289", + "1.06156156156156156156156156156157e-290", + "1.06156156156156156156156156156156e-291", + "1.0615615615615615615615615615616e-292", + "1.0615615615615615615615615615615e-293", + "1.061561561561561561561561561562e-294", + "1.06156156156156156156156156155e-295", + "1.0615615615615615615615615616e-296", + "1.06156156156156156156156156e-297", + "1.06156156156156156156156157e-298", + "1.0615615615615615615615616e-299", + "1.06156156156156156156156e-300", + "1.06156156156156156156155e-301", + "1.0615615615615615615616e-302", + "1.061561561561561561562e-303", + "1.06156156156156156156e-304", + "1.0615615615615615618e-305", + "1.06156156156156156e-306", + "1.06156156156156157e-307", + "1.0615615615615616e-308", + "1.06156156156156e-309", + "1.06156156156157e-310", + "1.0615615615616e-311", + "1.06156156156e-312", + "1.06156156154e-313", + "1.0615615616e-314", + "1.06156156e-315", + "1.06156155e-316", + "1.061562e-317", + "1.06156e-318", + "1.06155e-319", + "1.0617e-320", + "1.06e-321", + "1.04e-322", + "1e-323", + "0.0", + "0.0"] + assert_equal(got, expected) + + # Note: we follow glibc behavior, but it (or gcc) might not be right. + # In particular we can get two values that print the same but are not + # equal: + a = np.float128('2')/np.float128('3') + b = np.float128(str(a)) + assert_equal(str(a), str(b)) + assert_(a != b) + + def float32_roundtrip(self): + # gh-9360 + x = np.float32(1024 - 2**-14) + y = np.float32(1024 - 2**-13) + assert_(repr(x) != repr(y)) + assert_equal(np.float32(repr(x)), x) + assert_equal(np.float32(repr(y)), y) + + def float64_vs_python(self): + # gh-2643, gh-6136, gh-6908 + assert_equal(repr(np.float64(0.1)), repr(0.1)) + assert_(repr(np.float64(0.20000000000000004)) != repr(0.2)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_shape_base.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_shape_base.py new file mode 100644 index 0000000000000000000000000000000000000000..e90c89aa96347945d8ff5316c41b875e3f3ecbda --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_shape_base.py @@ -0,0 +1,860 @@ +import pytest +import numpy as np +from numpy._core import ( + array, arange, atleast_1d, atleast_2d, atleast_3d, block, vstack, hstack, + newaxis, concatenate, stack + ) +from numpy.exceptions import AxisError +from numpy._core.shape_base import (_block_dispatcher, _block_setup, + _block_concatenate, _block_slicing) +from numpy.testing import ( + assert_, assert_raises, assert_array_equal, assert_equal, + assert_raises_regex, assert_warns, IS_PYPY + ) + + +class TestAtleast1d: + def test_0D_array(self): + a = array(1) + b = array(2) + res = [atleast_1d(a), atleast_1d(b)] + desired = [array([1]), array([2])] + assert_array_equal(res, desired) + + def test_1D_array(self): + a = array([1, 2]) + b = array([2, 3]) + res = [atleast_1d(a), atleast_1d(b)] + desired = [array([1, 2]), array([2, 3])] + assert_array_equal(res, desired) + + def test_2D_array(self): + a = array([[1, 2], [1, 2]]) + b = array([[2, 3], [2, 3]]) + res = [atleast_1d(a), atleast_1d(b)] + desired = [a, b] + assert_array_equal(res, desired) + + def test_3D_array(self): + a = array([[1, 2], [1, 2]]) + b = array([[2, 3], [2, 3]]) + a = array([a, a]) + b = array([b, b]) + res = [atleast_1d(a), atleast_1d(b)] + desired = [a, b] + assert_array_equal(res, desired) + + def test_r1array(self): + """ Test to make sure equivalent Travis O's r1array function + """ + assert_(atleast_1d(3).shape == (1,)) + assert_(atleast_1d(3j).shape == (1,)) + assert_(atleast_1d(3.0).shape == (1,)) + assert_(atleast_1d([[2, 3], [4, 5]]).shape == (2, 2)) + + +class TestAtleast2d: + def test_0D_array(self): + a = array(1) + b = array(2) + res = [atleast_2d(a), atleast_2d(b)] + desired = [array([[1]]), array([[2]])] + assert_array_equal(res, desired) + + def test_1D_array(self): + a = array([1, 2]) + b = array([2, 3]) + res = [atleast_2d(a), atleast_2d(b)] + desired = [array([[1, 2]]), array([[2, 3]])] + assert_array_equal(res, desired) + + def test_2D_array(self): + a = array([[1, 2], [1, 2]]) + b = array([[2, 3], [2, 3]]) + res = [atleast_2d(a), atleast_2d(b)] + desired = [a, b] + assert_array_equal(res, desired) + + def test_3D_array(self): + a = array([[1, 2], [1, 2]]) + b = array([[2, 3], [2, 3]]) + a = array([a, a]) + b = array([b, b]) + res = [atleast_2d(a), atleast_2d(b)] + desired = [a, b] + assert_array_equal(res, desired) + + def test_r2array(self): + """ Test to make sure equivalent Travis O's r2array function + """ + assert_(atleast_2d(3).shape == (1, 1)) + assert_(atleast_2d([3j, 1]).shape == (1, 2)) + assert_(atleast_2d([[[3, 1], [4, 5]], [[3, 5], [1, 2]]]).shape == (2, 2, 2)) + + +class TestAtleast3d: + def test_0D_array(self): + a = array(1) + b = array(2) + res = [atleast_3d(a), atleast_3d(b)] + desired = [array([[[1]]]), array([[[2]]])] + assert_array_equal(res, desired) + + def test_1D_array(self): + a = array([1, 2]) + b = array([2, 3]) + res = [atleast_3d(a), atleast_3d(b)] + desired = [array([[[1], [2]]]), array([[[2], [3]]])] + assert_array_equal(res, desired) + + def test_2D_array(self): + a = array([[1, 2], [1, 2]]) + b = array([[2, 3], [2, 3]]) + res = [atleast_3d(a), atleast_3d(b)] + desired = [a[:,:, newaxis], b[:,:, newaxis]] + assert_array_equal(res, desired) + + def test_3D_array(self): + a = array([[1, 2], [1, 2]]) + b = array([[2, 3], [2, 3]]) + a = array([a, a]) + b = array([b, b]) + res = [atleast_3d(a), atleast_3d(b)] + desired = [a, b] + assert_array_equal(res, desired) + + +class TestHstack: + def test_non_iterable(self): + assert_raises(TypeError, hstack, 1) + + def test_empty_input(self): + assert_raises(ValueError, hstack, ()) + + def test_0D_array(self): + a = array(1) + b = array(2) + res = hstack([a, b]) + desired = array([1, 2]) + assert_array_equal(res, desired) + + def test_1D_array(self): + a = array([1]) + b = array([2]) + res = hstack([a, b]) + desired = array([1, 2]) + assert_array_equal(res, desired) + + def test_2D_array(self): + a = array([[1], [2]]) + b = array([[1], [2]]) + res = hstack([a, b]) + desired = array([[1, 1], [2, 2]]) + assert_array_equal(res, desired) + + def test_generator(self): + with pytest.raises(TypeError, match="arrays to stack must be"): + hstack(np.arange(3) for _ in range(2)) + with pytest.raises(TypeError, match="arrays to stack must be"): + hstack((x for x in np.ones((3, 2)))) + + def test_casting_and_dtype(self): + a = np.array([1, 2, 3]) + b = np.array([2.5, 3.5, 4.5]) + res = np.hstack((a, b), casting="unsafe", dtype=np.int64) + expected_res = np.array([1, 2, 3, 2, 3, 4]) + assert_array_equal(res, expected_res) + + def test_casting_and_dtype_type_error(self): + a = np.array([1, 2, 3]) + b = np.array([2.5, 3.5, 4.5]) + with pytest.raises(TypeError): + hstack((a, b), casting="safe", dtype=np.int64) + + +class TestVstack: + def test_non_iterable(self): + assert_raises(TypeError, vstack, 1) + + def test_empty_input(self): + assert_raises(ValueError, vstack, ()) + + def test_0D_array(self): + a = array(1) + b = array(2) + res = vstack([a, b]) + desired = array([[1], [2]]) + assert_array_equal(res, desired) + + def test_1D_array(self): + a = array([1]) + b = array([2]) + res = vstack([a, b]) + desired = array([[1], [2]]) + assert_array_equal(res, desired) + + def test_2D_array(self): + a = array([[1], [2]]) + b = array([[1], [2]]) + res = vstack([a, b]) + desired = array([[1], [2], [1], [2]]) + assert_array_equal(res, desired) + + def test_2D_array2(self): + a = array([1, 2]) + b = array([1, 2]) + res = vstack([a, b]) + desired = array([[1, 2], [1, 2]]) + assert_array_equal(res, desired) + + def test_generator(self): + with pytest.raises(TypeError, match="arrays to stack must be"): + vstack(np.arange(3) for _ in range(2)) + + def test_casting_and_dtype(self): + a = np.array([1, 2, 3]) + b = np.array([2.5, 3.5, 4.5]) + res = np.vstack((a, b), casting="unsafe", dtype=np.int64) + expected_res = np.array([[1, 2, 3], [2, 3, 4]]) + assert_array_equal(res, expected_res) + + def test_casting_and_dtype_type_error(self): + a = np.array([1, 2, 3]) + b = np.array([2.5, 3.5, 4.5]) + with pytest.raises(TypeError): + vstack((a, b), casting="safe", dtype=np.int64) + + + +class TestConcatenate: + def test_returns_copy(self): + a = np.eye(3) + b = np.concatenate([a]) + b[0, 0] = 2 + assert b[0, 0] != a[0, 0] + + def test_exceptions(self): + # test axis must be in bounds + for ndim in [1, 2, 3]: + a = np.ones((1,)*ndim) + np.concatenate((a, a), axis=0) # OK + assert_raises(AxisError, np.concatenate, (a, a), axis=ndim) + assert_raises(AxisError, np.concatenate, (a, a), axis=-(ndim + 1)) + + # Scalars cannot be concatenated + assert_raises(ValueError, concatenate, (0,)) + assert_raises(ValueError, concatenate, (np.array(0),)) + + # dimensionality must match + assert_raises_regex( + ValueError, + r"all the input arrays must have same number of dimensions, but " + r"the array at index 0 has 1 dimension\(s\) and the array at " + r"index 1 has 2 dimension\(s\)", + np.concatenate, (np.zeros(1), np.zeros((1, 1)))) + + # test shapes must match except for concatenation axis + a = np.ones((1, 2, 3)) + b = np.ones((2, 2, 3)) + axis = list(range(3)) + for i in range(3): + np.concatenate((a, b), axis=axis[0]) # OK + assert_raises_regex( + ValueError, + "all the input array dimensions except for the concatenation axis " + "must match exactly, but along dimension {}, the array at " + "index 0 has size 1 and the array at index 1 has size 2" + .format(i), + np.concatenate, (a, b), axis=axis[1]) + assert_raises(ValueError, np.concatenate, (a, b), axis=axis[2]) + a = np.moveaxis(a, -1, 0) + b = np.moveaxis(b, -1, 0) + axis.append(axis.pop(0)) + + # No arrays to concatenate raises ValueError + assert_raises(ValueError, concatenate, ()) + + def test_concatenate_axis_None(self): + a = np.arange(4, dtype=np.float64).reshape((2, 2)) + b = list(range(3)) + c = ['x'] + r = np.concatenate((a, a), axis=None) + assert_equal(r.dtype, a.dtype) + assert_equal(r.ndim, 1) + r = np.concatenate((a, b), axis=None) + assert_equal(r.size, a.size + len(b)) + assert_equal(r.dtype, a.dtype) + r = np.concatenate((a, b, c), axis=None, dtype="U") + d = array(['0.0', '1.0', '2.0', '3.0', + '0', '1', '2', 'x']) + assert_array_equal(r, d) + + out = np.zeros(a.size + len(b)) + r = np.concatenate((a, b), axis=None) + rout = np.concatenate((a, b), axis=None, out=out) + assert_(out is rout) + assert_equal(r, rout) + + def test_large_concatenate_axis_None(self): + # When no axis is given, concatenate uses flattened versions. + # This also had a bug with many arrays (see gh-5979). + x = np.arange(1, 100) + r = np.concatenate(x, None) + assert_array_equal(x, r) + + # Once upon a time, this was the same as `axis=None` now it fails + # (with an unspecified error, as multiple things are wrong here) + with pytest.raises(ValueError): + np.concatenate(x, 100) + + def test_concatenate(self): + # Test concatenate function + # One sequence returns unmodified (but as array) + r4 = list(range(4)) + assert_array_equal(concatenate((r4,)), r4) + # Any sequence + assert_array_equal(concatenate((tuple(r4),)), r4) + assert_array_equal(concatenate((array(r4),)), r4) + # 1D default concatenation + r3 = list(range(3)) + assert_array_equal(concatenate((r4, r3)), r4 + r3) + # Mixed sequence types + assert_array_equal(concatenate((tuple(r4), r3)), r4 + r3) + assert_array_equal(concatenate((array(r4), r3)), r4 + r3) + # Explicit axis specification + assert_array_equal(concatenate((r4, r3), 0), r4 + r3) + # Including negative + assert_array_equal(concatenate((r4, r3), -1), r4 + r3) + # 2D + a23 = array([[10, 11, 12], [13, 14, 15]]) + a13 = array([[0, 1, 2]]) + res = array([[10, 11, 12], [13, 14, 15], [0, 1, 2]]) + assert_array_equal(concatenate((a23, a13)), res) + assert_array_equal(concatenate((a23, a13), 0), res) + assert_array_equal(concatenate((a23.T, a13.T), 1), res.T) + assert_array_equal(concatenate((a23.T, a13.T), -1), res.T) + # Arrays much match shape + assert_raises(ValueError, concatenate, (a23.T, a13.T), 0) + # 3D + res = arange(2 * 3 * 7).reshape((2, 3, 7)) + a0 = res[..., :4] + a1 = res[..., 4:6] + a2 = res[..., 6:] + assert_array_equal(concatenate((a0, a1, a2), 2), res) + assert_array_equal(concatenate((a0, a1, a2), -1), res) + assert_array_equal(concatenate((a0.T, a1.T, a2.T), 0), res.T) + + out = res.copy() + rout = concatenate((a0, a1, a2), 2, out=out) + assert_(out is rout) + assert_equal(res, rout) + + @pytest.mark.skipif(IS_PYPY, reason="PYPY handles sq_concat, nb_add differently than cpython") + def test_operator_concat(self): + import operator + a = array([1, 2]) + b = array([3, 4]) + n = [1,2] + res = array([1, 2, 3, 4]) + assert_raises(TypeError, operator.concat, a, b) + assert_raises(TypeError, operator.concat, a, n) + assert_raises(TypeError, operator.concat, n, a) + assert_raises(TypeError, operator.concat, a, 1) + assert_raises(TypeError, operator.concat, 1, a) + + def test_bad_out_shape(self): + a = array([1, 2]) + b = array([3, 4]) + + assert_raises(ValueError, concatenate, (a, b), out=np.empty(5)) + assert_raises(ValueError, concatenate, (a, b), out=np.empty((4,1))) + assert_raises(ValueError, concatenate, (a, b), out=np.empty((1,4))) + concatenate((a, b), out=np.empty(4)) + + @pytest.mark.parametrize("axis", [None, 0]) + @pytest.mark.parametrize("out_dtype", ["c8", "f4", "f8", ">f8", "i8", "S4"]) + @pytest.mark.parametrize("casting", + ['no', 'equiv', 'safe', 'same_kind', 'unsafe']) + def test_out_and_dtype(self, axis, out_dtype, casting): + # Compare usage of `out=out` with `dtype=out.dtype` + out = np.empty(4, dtype=out_dtype) + to_concat = (array([1.1, 2.2]), array([3.3, 4.4])) + + if not np.can_cast(to_concat[0], out_dtype, casting=casting): + with assert_raises(TypeError): + concatenate(to_concat, out=out, axis=axis, casting=casting) + with assert_raises(TypeError): + concatenate(to_concat, dtype=out.dtype, + axis=axis, casting=casting) + else: + res_out = concatenate(to_concat, out=out, + axis=axis, casting=casting) + res_dtype = concatenate(to_concat, dtype=out.dtype, + axis=axis, casting=casting) + assert res_out is out + assert_array_equal(out, res_dtype) + assert res_dtype.dtype == out_dtype + + with assert_raises(TypeError): + concatenate(to_concat, out=out, dtype=out_dtype, axis=axis) + + @pytest.mark.parametrize("axis", [None, 0]) + @pytest.mark.parametrize("string_dt", ["S", "U", "S0", "U0"]) + @pytest.mark.parametrize("arrs", + [([0.],), ([0.], [1]), ([0], ["string"], [1.])]) + def test_dtype_with_promotion(self, arrs, string_dt, axis): + # Note that U0 and S0 should be deprecated eventually and changed to + # actually give the empty string result (together with `np.array`) + res = np.concatenate(arrs, axis=axis, dtype=string_dt, casting="unsafe") + # The actual dtype should be identical to a cast (of a double array): + assert res.dtype == np.array(1.).astype(string_dt).dtype + + @pytest.mark.parametrize("axis", [None, 0]) + def test_string_dtype_does_not_inspect(self, axis): + with pytest.raises(TypeError): + np.concatenate(([None], [1]), dtype="S", axis=axis) + with pytest.raises(TypeError): + np.concatenate(([None], [1]), dtype="U", axis=axis) + + @pytest.mark.parametrize("axis", [None, 0]) + def test_subarray_error(self, axis): + with pytest.raises(TypeError, match=".*subarray dtype"): + np.concatenate(([1], [1]), dtype="(2,)i", axis=axis) + + +def test_stack(): + # non-iterable input + assert_raises(TypeError, stack, 1) + + # 0d input + for input_ in [(1, 2, 3), + [np.int32(1), np.int32(2), np.int32(3)], + [np.array(1), np.array(2), np.array(3)]]: + assert_array_equal(stack(input_), [1, 2, 3]) + # 1d input examples + a = np.array([1, 2, 3]) + b = np.array([4, 5, 6]) + r1 = array([[1, 2, 3], [4, 5, 6]]) + assert_array_equal(np.stack((a, b)), r1) + assert_array_equal(np.stack((a, b), axis=1), r1.T) + # all input types + assert_array_equal(np.stack([a, b]), r1) + assert_array_equal(np.stack(array([a, b])), r1) + # all shapes for 1d input + arrays = [np.random.randn(3) for _ in range(10)] + axes = [0, 1, -1, -2] + expected_shapes = [(10, 3), (3, 10), (3, 10), (10, 3)] + for axis, expected_shape in zip(axes, expected_shapes): + assert_equal(np.stack(arrays, axis).shape, expected_shape) + assert_raises_regex(AxisError, 'out of bounds', stack, arrays, axis=2) + assert_raises_regex(AxisError, 'out of bounds', stack, arrays, axis=-3) + # all shapes for 2d input + arrays = [np.random.randn(3, 4) for _ in range(10)] + axes = [0, 1, 2, -1, -2, -3] + expected_shapes = [(10, 3, 4), (3, 10, 4), (3, 4, 10), + (3, 4, 10), (3, 10, 4), (10, 3, 4)] + for axis, expected_shape in zip(axes, expected_shapes): + assert_equal(np.stack(arrays, axis).shape, expected_shape) + # empty arrays + assert_(stack([[], [], []]).shape == (3, 0)) + assert_(stack([[], [], []], axis=1).shape == (0, 3)) + # out + out = np.zeros_like(r1) + np.stack((a, b), out=out) + assert_array_equal(out, r1) + # edge cases + assert_raises_regex(ValueError, 'need at least one array', stack, []) + assert_raises_regex(ValueError, 'must have the same shape', + stack, [1, np.arange(3)]) + assert_raises_regex(ValueError, 'must have the same shape', + stack, [np.arange(3), 1]) + assert_raises_regex(ValueError, 'must have the same shape', + stack, [np.arange(3), 1], axis=1) + assert_raises_regex(ValueError, 'must have the same shape', + stack, [np.zeros((3, 3)), np.zeros(3)], axis=1) + assert_raises_regex(ValueError, 'must have the same shape', + stack, [np.arange(2), np.arange(3)]) + + # do not accept generators + with pytest.raises(TypeError, match="arrays to stack must be"): + stack(x for x in range(3)) + + #casting and dtype test + a = np.array([1, 2, 3]) + b = np.array([2.5, 3.5, 4.5]) + res = np.stack((a, b), axis=1, casting="unsafe", dtype=np.int64) + expected_res = np.array([[1, 2], [2, 3], [3, 4]]) + assert_array_equal(res, expected_res) + #casting and dtype with TypeError + with assert_raises(TypeError): + stack((a, b), dtype=np.int64, axis=1, casting="safe") + + +def test_unstack(): + a = np.arange(24).reshape((2, 3, 4)) + + for stacks in [np.unstack(a), + np.unstack(a, axis=0), + np.unstack(a, axis=-3)]: + assert isinstance(stacks, tuple) + assert len(stacks) == 2 + assert_array_equal(stacks[0], a[0]) + assert_array_equal(stacks[1], a[1]) + + for stacks in [np.unstack(a, axis=1), + np.unstack(a, axis=-2)]: + assert isinstance(stacks, tuple) + assert len(stacks) == 3 + assert_array_equal(stacks[0], a[:, 0]) + assert_array_equal(stacks[1], a[:, 1]) + assert_array_equal(stacks[2], a[:, 2]) + + for stacks in [np.unstack(a, axis=2), + np.unstack(a, axis=-1)]: + assert isinstance(stacks, tuple) + assert len(stacks) == 4 + assert_array_equal(stacks[0], a[:, :, 0]) + assert_array_equal(stacks[1], a[:, :, 1]) + assert_array_equal(stacks[2], a[:, :, 2]) + assert_array_equal(stacks[3], a[:, :, 3]) + + assert_raises(ValueError, np.unstack, a, axis=3) + assert_raises(ValueError, np.unstack, a, axis=-4) + assert_raises(ValueError, np.unstack, np.array(0), axis=0) + + +@pytest.mark.parametrize("axis", [0]) +@pytest.mark.parametrize("out_dtype", ["c8", "f4", "f8", ">f8", "i8"]) +@pytest.mark.parametrize("casting", + ['no', 'equiv', 'safe', 'same_kind', 'unsafe']) +def test_stack_out_and_dtype(axis, out_dtype, casting): + to_concat = (array([1, 2]), array([3, 4])) + res = array([[1, 2], [3, 4]]) + out = np.zeros_like(res) + + if not np.can_cast(to_concat[0], out_dtype, casting=casting): + with assert_raises(TypeError): + stack(to_concat, dtype=out_dtype, + axis=axis, casting=casting) + else: + res_out = stack(to_concat, out=out, + axis=axis, casting=casting) + res_dtype = stack(to_concat, dtype=out_dtype, + axis=axis, casting=casting) + assert res_out is out + assert_array_equal(out, res_dtype) + assert res_dtype.dtype == out_dtype + + with assert_raises(TypeError): + stack(to_concat, out=out, dtype=out_dtype, axis=axis) + + +class TestBlock: + @pytest.fixture(params=['block', 'force_concatenate', 'force_slicing']) + def block(self, request): + # blocking small arrays and large arrays go through different paths. + # the algorithm is triggered depending on the number of element + # copies required. + # We define a test fixture that forces most tests to go through + # both code paths. + # Ultimately, this should be removed if a single algorithm is found + # to be faster for both small and large arrays. + def _block_force_concatenate(arrays): + arrays, list_ndim, result_ndim, _ = _block_setup(arrays) + return _block_concatenate(arrays, list_ndim, result_ndim) + + def _block_force_slicing(arrays): + arrays, list_ndim, result_ndim, _ = _block_setup(arrays) + return _block_slicing(arrays, list_ndim, result_ndim) + + if request.param == 'force_concatenate': + return _block_force_concatenate + elif request.param == 'force_slicing': + return _block_force_slicing + elif request.param == 'block': + return block + else: + raise ValueError('Unknown blocking request. There is a typo in the tests.') + + def test_returns_copy(self, block): + a = np.eye(3) + b = block(a) + b[0, 0] = 2 + assert b[0, 0] != a[0, 0] + + def test_block_total_size_estimate(self, block): + _, _, _, total_size = _block_setup([1]) + assert total_size == 1 + + _, _, _, total_size = _block_setup([[1]]) + assert total_size == 1 + + _, _, _, total_size = _block_setup([[1, 1]]) + assert total_size == 2 + + _, _, _, total_size = _block_setup([[1], [1]]) + assert total_size == 2 + + _, _, _, total_size = _block_setup([[1, 2], [3, 4]]) + assert total_size == 4 + + def test_block_simple_row_wise(self, block): + a_2d = np.ones((2, 2)) + b_2d = 2 * a_2d + desired = np.array([[1, 1, 2, 2], + [1, 1, 2, 2]]) + result = block([a_2d, b_2d]) + assert_equal(desired, result) + + def test_block_simple_column_wise(self, block): + a_2d = np.ones((2, 2)) + b_2d = 2 * a_2d + expected = np.array([[1, 1], + [1, 1], + [2, 2], + [2, 2]]) + result = block([[a_2d], [b_2d]]) + assert_equal(expected, result) + + def test_block_with_1d_arrays_row_wise(self, block): + # # # 1-D vectors are treated as row arrays + a = np.array([1, 2, 3]) + b = np.array([2, 3, 4]) + expected = np.array([1, 2, 3, 2, 3, 4]) + result = block([a, b]) + assert_equal(expected, result) + + def test_block_with_1d_arrays_multiple_rows(self, block): + a = np.array([1, 2, 3]) + b = np.array([2, 3, 4]) + expected = np.array([[1, 2, 3, 2, 3, 4], + [1, 2, 3, 2, 3, 4]]) + result = block([[a, b], [a, b]]) + assert_equal(expected, result) + + def test_block_with_1d_arrays_column_wise(self, block): + # # # 1-D vectors are treated as row arrays + a_1d = np.array([1, 2, 3]) + b_1d = np.array([2, 3, 4]) + expected = np.array([[1, 2, 3], + [2, 3, 4]]) + result = block([[a_1d], [b_1d]]) + assert_equal(expected, result) + + def test_block_mixed_1d_and_2d(self, block): + a_2d = np.ones((2, 2)) + b_1d = np.array([2, 2]) + result = block([[a_2d], [b_1d]]) + expected = np.array([[1, 1], + [1, 1], + [2, 2]]) + assert_equal(expected, result) + + def test_block_complicated(self, block): + # a bit more complicated + one_2d = np.array([[1, 1, 1]]) + two_2d = np.array([[2, 2, 2]]) + three_2d = np.array([[3, 3, 3, 3, 3, 3]]) + four_1d = np.array([4, 4, 4, 4, 4, 4]) + five_0d = np.array(5) + six_1d = np.array([6, 6, 6, 6, 6]) + zero_2d = np.zeros((2, 6)) + + expected = np.array([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 3, 3, 3], + [4, 4, 4, 4, 4, 4], + [5, 6, 6, 6, 6, 6], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0]]) + + result = block([[one_2d, two_2d], + [three_2d], + [four_1d], + [five_0d, six_1d], + [zero_2d]]) + assert_equal(result, expected) + + def test_nested(self, block): + one = np.array([1, 1, 1]) + two = np.array([[2, 2, 2], [2, 2, 2], [2, 2, 2]]) + three = np.array([3, 3, 3]) + four = np.array([4, 4, 4]) + five = np.array(5) + six = np.array([6, 6, 6, 6, 6]) + zero = np.zeros((2, 6)) + + result = block([ + [ + block([ + [one], + [three], + [four] + ]), + two + ], + [five, six], + [zero] + ]) + expected = np.array([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 2, 2, 2], + [4, 4, 4, 2, 2, 2], + [5, 6, 6, 6, 6, 6], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0]]) + + assert_equal(result, expected) + + def test_3d(self, block): + a000 = np.ones((2, 2, 2), int) * 1 + + a100 = np.ones((3, 2, 2), int) * 2 + a010 = np.ones((2, 3, 2), int) * 3 + a001 = np.ones((2, 2, 3), int) * 4 + + a011 = np.ones((2, 3, 3), int) * 5 + a101 = np.ones((3, 2, 3), int) * 6 + a110 = np.ones((3, 3, 2), int) * 7 + + a111 = np.ones((3, 3, 3), int) * 8 + + result = block([ + [ + [a000, a001], + [a010, a011], + ], + [ + [a100, a101], + [a110, a111], + ] + ]) + expected = array([[[1, 1, 4, 4, 4], + [1, 1, 4, 4, 4], + [3, 3, 5, 5, 5], + [3, 3, 5, 5, 5], + [3, 3, 5, 5, 5]], + + [[1, 1, 4, 4, 4], + [1, 1, 4, 4, 4], + [3, 3, 5, 5, 5], + [3, 3, 5, 5, 5], + [3, 3, 5, 5, 5]], + + [[2, 2, 6, 6, 6], + [2, 2, 6, 6, 6], + [7, 7, 8, 8, 8], + [7, 7, 8, 8, 8], + [7, 7, 8, 8, 8]], + + [[2, 2, 6, 6, 6], + [2, 2, 6, 6, 6], + [7, 7, 8, 8, 8], + [7, 7, 8, 8, 8], + [7, 7, 8, 8, 8]], + + [[2, 2, 6, 6, 6], + [2, 2, 6, 6, 6], + [7, 7, 8, 8, 8], + [7, 7, 8, 8, 8], + [7, 7, 8, 8, 8]]]) + + assert_array_equal(result, expected) + + def test_block_with_mismatched_shape(self, block): + a = np.array([0, 0]) + b = np.eye(2) + assert_raises(ValueError, block, [a, b]) + assert_raises(ValueError, block, [b, a]) + + to_block = [[np.ones((2,3)), np.ones((2,2))], + [np.ones((2,2)), np.ones((2,2))]] + assert_raises(ValueError, block, to_block) + def test_no_lists(self, block): + assert_equal(block(1), np.array(1)) + assert_equal(block(np.eye(3)), np.eye(3)) + + def test_invalid_nesting(self, block): + msg = 'depths are mismatched' + assert_raises_regex(ValueError, msg, block, [1, [2]]) + assert_raises_regex(ValueError, msg, block, [1, []]) + assert_raises_regex(ValueError, msg, block, [[1], 2]) + assert_raises_regex(ValueError, msg, block, [[], 2]) + assert_raises_regex(ValueError, msg, block, [ + [[1], [2]], + [[3, 4]], + [5] # missing brackets + ]) + + def test_empty_lists(self, block): + assert_raises_regex(ValueError, 'empty', block, []) + assert_raises_regex(ValueError, 'empty', block, [[]]) + assert_raises_regex(ValueError, 'empty', block, [[1], []]) + + def test_tuple(self, block): + assert_raises_regex(TypeError, 'tuple', block, ([1, 2], [3, 4])) + assert_raises_regex(TypeError, 'tuple', block, [(1, 2), (3, 4)]) + + def test_different_ndims(self, block): + a = 1. + b = 2 * np.ones((1, 2)) + c = 3 * np.ones((1, 1, 3)) + + result = block([a, b, c]) + expected = np.array([[[1., 2., 2., 3., 3., 3.]]]) + + assert_equal(result, expected) + + def test_different_ndims_depths(self, block): + a = 1. + b = 2 * np.ones((1, 2)) + c = 3 * np.ones((1, 2, 3)) + + result = block([[a, b], [c]]) + expected = np.array([[[1., 2., 2.], + [3., 3., 3.], + [3., 3., 3.]]]) + + assert_equal(result, expected) + + def test_block_memory_order(self, block): + # 3D + arr_c = np.zeros((3,)*3, order='C') + arr_f = np.zeros((3,)*3, order='F') + + b_c = [[[arr_c, arr_c], + [arr_c, arr_c]], + [[arr_c, arr_c], + [arr_c, arr_c]]] + + b_f = [[[arr_f, arr_f], + [arr_f, arr_f]], + [[arr_f, arr_f], + [arr_f, arr_f]]] + + assert block(b_c).flags['C_CONTIGUOUS'] + assert block(b_f).flags['F_CONTIGUOUS'] + + arr_c = np.zeros((3, 3), order='C') + arr_f = np.zeros((3, 3), order='F') + # 2D + b_c = [[arr_c, arr_c], + [arr_c, arr_c]] + + b_f = [[arr_f, arr_f], + [arr_f, arr_f]] + + assert block(b_c).flags['C_CONTIGUOUS'] + assert block(b_f).flags['F_CONTIGUOUS'] + + +def test_block_dispatcher(): + class ArrayLike: + pass + a = ArrayLike() + b = ArrayLike() + c = ArrayLike() + assert_equal(list(_block_dispatcher(a)), [a]) + assert_equal(list(_block_dispatcher([a])), [a]) + assert_equal(list(_block_dispatcher([a, b])), [a, b]) + assert_equal(list(_block_dispatcher([[a], [b, [c]]])), [a, b, c]) + # don't recurse into non-lists + assert_equal(list(_block_dispatcher((a, b))), [(a, b)]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_simd.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_simd.py new file mode 100644 index 0000000000000000000000000000000000000000..6b8f208ee31c527a560962d924fc9604042db908 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_simd.py @@ -0,0 +1,1335 @@ +# NOTE: Please avoid the use of numpy.testing since NPYV intrinsics +# may be involved in their functionality. +import itertools +import math +import operator +import re +import pytest +from numpy._core._simd import targets, clear_floatstatus, get_floatstatus +from numpy._core._multiarray_umath import __cpu_baseline__ + +def check_floatstatus(divbyzero=False, overflow=False, + underflow=False, invalid=False, + all=False): + #define NPY_FPE_DIVIDEBYZERO 1 + #define NPY_FPE_OVERFLOW 2 + #define NPY_FPE_UNDERFLOW 4 + #define NPY_FPE_INVALID 8 + err = get_floatstatus() + ret = (all or divbyzero) and (err & 1) != 0 + ret |= (all or overflow) and (err & 2) != 0 + ret |= (all or underflow) and (err & 4) != 0 + ret |= (all or invalid) and (err & 8) != 0 + return ret + +class _Test_Utility: + # submodule of the desired SIMD extension, e.g. targets["AVX512F"] + npyv = None + # the current data type suffix e.g. 's8' + sfx = None + # target name can be 'baseline' or one or more of CPU features + target_name = None + + def __getattr__(self, attr): + """ + To call NPV intrinsics without the attribute 'npyv' and + auto suffixing intrinsics according to class attribute 'sfx' + """ + return getattr(self.npyv, attr + "_" + self.sfx) + + def _x2(self, intrin_name): + return getattr(self.npyv, f"{intrin_name}_{self.sfx}x2") + + def _data(self, start=None, count=None, reverse=False): + """ + Create list of consecutive numbers according to number of vector's lanes. + """ + if start is None: + start = 1 + if count is None: + count = self.nlanes + rng = range(start, start + count) + if reverse: + rng = reversed(rng) + if self._is_fp(): + return [x / 1.0 for x in rng] + return list(rng) + + def _is_unsigned(self): + return self.sfx[0] == 'u' + + def _is_signed(self): + return self.sfx[0] == 's' + + def _is_fp(self): + return self.sfx[0] == 'f' + + def _scalar_size(self): + return int(self.sfx[1:]) + + def _int_clip(self, seq): + if self._is_fp(): + return seq + max_int = self._int_max() + min_int = self._int_min() + return [min(max(v, min_int), max_int) for v in seq] + + def _int_max(self): + if self._is_fp(): + return None + max_u = self._to_unsigned(self.setall(-1))[0] + if self._is_signed(): + return max_u // 2 + return max_u + + def _int_min(self): + if self._is_fp(): + return None + if self._is_unsigned(): + return 0 + return -(self._int_max() + 1) + + def _true_mask(self): + max_unsig = getattr(self.npyv, "setall_u" + self.sfx[1:])(-1) + return max_unsig[0] + + def _to_unsigned(self, vector): + if isinstance(vector, (list, tuple)): + return getattr(self.npyv, "load_u" + self.sfx[1:])(vector) + else: + sfx = vector.__name__.replace("npyv_", "") + if sfx[0] == "b": + cvt_intrin = "cvt_u{0}_b{0}" + else: + cvt_intrin = "reinterpret_u{0}_{1}" + return getattr(self.npyv, cvt_intrin.format(sfx[1:], sfx))(vector) + + def _pinfinity(self): + return float("inf") + + def _ninfinity(self): + return -float("inf") + + def _nan(self): + return float("nan") + + def _cpu_features(self): + target = self.target_name + if target == "baseline": + target = __cpu_baseline__ + else: + target = target.split('__') # multi-target separator + return ' '.join(target) + +class _SIMD_BOOL(_Test_Utility): + """ + To test all boolean vector types at once + """ + def _nlanes(self): + return getattr(self.npyv, "nlanes_u" + self.sfx[1:]) + + def _data(self, start=None, count=None, reverse=False): + true_mask = self._true_mask() + rng = range(self._nlanes()) + if reverse: + rng = reversed(rng) + return [true_mask if x % 2 else 0 for x in rng] + + def _load_b(self, data): + len_str = self.sfx[1:] + load = getattr(self.npyv, "load_u" + len_str) + cvt = getattr(self.npyv, f"cvt_b{len_str}_u{len_str}") + return cvt(load(data)) + + def test_operators_logical(self): + """ + Logical operations for boolean types. + Test intrinsics: + npyv_xor_##SFX, npyv_and_##SFX, npyv_or_##SFX, npyv_not_##SFX, + npyv_andc_b8, npvy_orc_b8, nvpy_xnor_b8 + """ + data_a = self._data() + data_b = self._data(reverse=True) + vdata_a = self._load_b(data_a) + vdata_b = self._load_b(data_b) + + data_and = [a & b for a, b in zip(data_a, data_b)] + vand = getattr(self, "and")(vdata_a, vdata_b) + assert vand == data_and + + data_or = [a | b for a, b in zip(data_a, data_b)] + vor = getattr(self, "or")(vdata_a, vdata_b) + assert vor == data_or + + data_xor = [a ^ b for a, b in zip(data_a, data_b)] + vxor = self.xor(vdata_a, vdata_b) + assert vxor == data_xor + + vnot = getattr(self, "not")(vdata_a) + assert vnot == data_b + + # among the boolean types, andc, orc and xnor only support b8 + if self.sfx not in ("b8"): + return + + data_andc = [(a & ~b) & 0xFF for a, b in zip(data_a, data_b)] + vandc = self.andc(vdata_a, vdata_b) + assert data_andc == vandc + + data_orc = [(a | ~b) & 0xFF for a, b in zip(data_a, data_b)] + vorc = self.orc(vdata_a, vdata_b) + assert data_orc == vorc + + data_xnor = [~(a ^ b) & 0xFF for a, b in zip(data_a, data_b)] + vxnor = self.xnor(vdata_a, vdata_b) + assert data_xnor == vxnor + + def test_tobits(self): + data2bits = lambda data: sum([int(x != 0) << i for i, x in enumerate(data, 0)]) + for data in (self._data(), self._data(reverse=True)): + vdata = self._load_b(data) + data_bits = data2bits(data) + tobits = self.tobits(vdata) + bin_tobits = bin(tobits) + assert bin_tobits == bin(data_bits) + + def test_pack(self): + """ + Pack multiple vectors into one + Test intrinsics: + npyv_pack_b8_b16 + npyv_pack_b8_b32 + npyv_pack_b8_b64 + """ + if self.sfx not in ("b16", "b32", "b64"): + return + # create the vectors + data = self._data() + rdata = self._data(reverse=True) + vdata = self._load_b(data) + vrdata = self._load_b(rdata) + pack_simd = getattr(self.npyv, f"pack_b8_{self.sfx}") + # for scalar execution, concatenate the elements of the multiple lists + # into a single list (spack) and then iterate over the elements of + # the created list applying a mask to capture the first byte of them. + if self.sfx == "b16": + spack = [(i & 0xFF) for i in (list(rdata) + list(data))] + vpack = pack_simd(vrdata, vdata) + elif self.sfx == "b32": + spack = [(i & 0xFF) for i in (2*list(rdata) + 2*list(data))] + vpack = pack_simd(vrdata, vrdata, vdata, vdata) + elif self.sfx == "b64": + spack = [(i & 0xFF) for i in (4*list(rdata) + 4*list(data))] + vpack = pack_simd(vrdata, vrdata, vrdata, vrdata, + vdata, vdata, vdata, vdata) + assert vpack == spack + + @pytest.mark.parametrize("intrin", ["any", "all"]) + @pytest.mark.parametrize("data", ( + [-1, 0], + [0, -1], + [-1], + [0] + )) + def test_operators_crosstest(self, intrin, data): + """ + Test intrinsics: + npyv_any_##SFX + npyv_all_##SFX + """ + data_a = self._load_b(data * self._nlanes()) + func = eval(intrin) + intrin = getattr(self, intrin) + desired = func(data_a) + simd = intrin(data_a) + assert not not simd == desired + +class _SIMD_INT(_Test_Utility): + """ + To test all integer vector types at once + """ + def test_operators_shift(self): + if self.sfx in ("u8", "s8"): + return + + data_a = self._data(self._int_max() - self.nlanes) + data_b = self._data(self._int_min(), reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + for count in range(self._scalar_size()): + # load to cast + data_shl_a = self.load([a << count for a in data_a]) + # left shift + shl = self.shl(vdata_a, count) + assert shl == data_shl_a + # load to cast + data_shr_a = self.load([a >> count for a in data_a]) + # right shift + shr = self.shr(vdata_a, count) + assert shr == data_shr_a + + # shift by zero or max or out-range immediate constant is not applicable and illogical + for count in range(1, self._scalar_size()): + # load to cast + data_shl_a = self.load([a << count for a in data_a]) + # left shift by an immediate constant + shli = self.shli(vdata_a, count) + assert shli == data_shl_a + # load to cast + data_shr_a = self.load([a >> count for a in data_a]) + # right shift by an immediate constant + shri = self.shri(vdata_a, count) + assert shri == data_shr_a + + def test_arithmetic_subadd_saturated(self): + if self.sfx in ("u32", "s32", "u64", "s64"): + return + + data_a = self._data(self._int_max() - self.nlanes) + data_b = self._data(self._int_min(), reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + data_adds = self._int_clip([a + b for a, b in zip(data_a, data_b)]) + adds = self.adds(vdata_a, vdata_b) + assert adds == data_adds + + data_subs = self._int_clip([a - b for a, b in zip(data_a, data_b)]) + subs = self.subs(vdata_a, vdata_b) + assert subs == data_subs + + def test_math_max_min(self): + data_a = self._data() + data_b = self._data(self.nlanes) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + data_max = [max(a, b) for a, b in zip(data_a, data_b)] + simd_max = self.max(vdata_a, vdata_b) + assert simd_max == data_max + + data_min = [min(a, b) for a, b in zip(data_a, data_b)] + simd_min = self.min(vdata_a, vdata_b) + assert simd_min == data_min + + @pytest.mark.parametrize("start", [-100, -10000, 0, 100, 10000]) + def test_reduce_max_min(self, start): + """ + Test intrinsics: + npyv_reduce_max_##sfx + npyv_reduce_min_##sfx + """ + vdata_a = self.load(self._data(start)) + assert self.reduce_max(vdata_a) == max(vdata_a) + assert self.reduce_min(vdata_a) == min(vdata_a) + + +class _SIMD_FP32(_Test_Utility): + """ + To only test single precision + """ + def test_conversions(self): + """ + Round to nearest even integer, assume CPU control register is set to rounding. + Test intrinsics: + npyv_round_s32_##SFX + """ + features = self._cpu_features() + if not self.npyv.simd_f64 and re.match(r".*(NEON|ASIMD)", features): + # very costly to emulate nearest even on Armv7 + # instead we round halves to up. e.g. 0.5 -> 1, -0.5 -> -1 + _round = lambda v: int(v + (0.5 if v >= 0 else -0.5)) + else: + _round = round + vdata_a = self.load(self._data()) + vdata_a = self.sub(vdata_a, self.setall(0.5)) + data_round = [_round(x) for x in vdata_a] + vround = self.round_s32(vdata_a) + assert vround == data_round + +class _SIMD_FP64(_Test_Utility): + """ + To only test double precision + """ + def test_conversions(self): + """ + Round to nearest even integer, assume CPU control register is set to rounding. + Test intrinsics: + npyv_round_s32_##SFX + """ + vdata_a = self.load(self._data()) + vdata_a = self.sub(vdata_a, self.setall(0.5)) + vdata_b = self.mul(vdata_a, self.setall(-1.5)) + data_round = [round(x) for x in list(vdata_a) + list(vdata_b)] + vround = self.round_s32(vdata_a, vdata_b) + assert vround == data_round + +class _SIMD_FP(_Test_Utility): + """ + To test all float vector types at once + """ + def test_arithmetic_fused(self): + vdata_a, vdata_b, vdata_c = [self.load(self._data())]*3 + vdata_cx2 = self.add(vdata_c, vdata_c) + # multiply and add, a*b + c + data_fma = self.load([a * b + c for a, b, c in zip(vdata_a, vdata_b, vdata_c)]) + fma = self.muladd(vdata_a, vdata_b, vdata_c) + assert fma == data_fma + # multiply and subtract, a*b - c + fms = self.mulsub(vdata_a, vdata_b, vdata_c) + data_fms = self.sub(data_fma, vdata_cx2) + assert fms == data_fms + # negate multiply and add, -(a*b) + c + nfma = self.nmuladd(vdata_a, vdata_b, vdata_c) + data_nfma = self.sub(vdata_cx2, data_fma) + assert nfma == data_nfma + # negate multiply and subtract, -(a*b) - c + nfms = self.nmulsub(vdata_a, vdata_b, vdata_c) + data_nfms = self.mul(data_fma, self.setall(-1)) + assert nfms == data_nfms + # multiply, add for odd elements and subtract even elements. + # (a * b) -+ c + fmas = list(self.muladdsub(vdata_a, vdata_b, vdata_c)) + assert fmas[0::2] == list(data_fms)[0::2] + assert fmas[1::2] == list(data_fma)[1::2] + + def test_abs(self): + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + data = self._data() + vdata = self.load(self._data()) + + abs_cases = ((-0, 0), (ninf, pinf), (pinf, pinf), (nan, nan)) + for case, desired in abs_cases: + data_abs = [desired]*self.nlanes + vabs = self.abs(self.setall(case)) + assert vabs == pytest.approx(data_abs, nan_ok=True) + + vabs = self.abs(self.mul(vdata, self.setall(-1))) + assert vabs == data + + def test_sqrt(self): + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + data = self._data() + vdata = self.load(self._data()) + + sqrt_cases = ((-0.0, -0.0), (0.0, 0.0), (-1.0, nan), (ninf, nan), (pinf, pinf)) + for case, desired in sqrt_cases: + data_sqrt = [desired]*self.nlanes + sqrt = self.sqrt(self.setall(case)) + assert sqrt == pytest.approx(data_sqrt, nan_ok=True) + + data_sqrt = self.load([math.sqrt(x) for x in data]) # load to truncate precision + sqrt = self.sqrt(vdata) + assert sqrt == data_sqrt + + def test_square(self): + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + data = self._data() + vdata = self.load(self._data()) + # square + square_cases = ((nan, nan), (pinf, pinf), (ninf, pinf)) + for case, desired in square_cases: + data_square = [desired]*self.nlanes + square = self.square(self.setall(case)) + assert square == pytest.approx(data_square, nan_ok=True) + + data_square = [x*x for x in data] + square = self.square(vdata) + assert square == data_square + + @pytest.mark.parametrize("intrin, func", [("ceil", math.ceil), + ("trunc", math.trunc), ("floor", math.floor), ("rint", round)]) + def test_rounding(self, intrin, func): + """ + Test intrinsics: + npyv_rint_##SFX + npyv_ceil_##SFX + npyv_trunc_##SFX + npyv_floor##SFX + """ + intrin_name = intrin + intrin = getattr(self, intrin) + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + # special cases + round_cases = ((nan, nan), (pinf, pinf), (ninf, ninf)) + for case, desired in round_cases: + data_round = [desired]*self.nlanes + _round = intrin(self.setall(case)) + assert _round == pytest.approx(data_round, nan_ok=True) + + for x in range(0, 2**20, 256**2): + for w in (-1.05, -1.10, -1.15, 1.05, 1.10, 1.15): + data = self.load([(x+a)*w for a in range(self.nlanes)]) + data_round = [func(x) for x in data] + _round = intrin(data) + assert _round == data_round + + # test large numbers + for i in ( + 1.1529215045988576e+18, 4.6116860183954304e+18, + 5.902958103546122e+20, 2.3611832414184488e+21 + ): + x = self.setall(i) + y = intrin(x) + data_round = [func(n) for n in x] + assert y == data_round + + # signed zero + if intrin_name == "floor": + data_szero = (-0.0,) + else: + data_szero = (-0.0, -0.25, -0.30, -0.45, -0.5) + + for w in data_szero: + _round = self._to_unsigned(intrin(self.setall(w))) + data_round = self._to_unsigned(self.setall(-0.0)) + assert _round == data_round + + @pytest.mark.parametrize("intrin", [ + "max", "maxp", "maxn", "min", "minp", "minn" + ]) + def test_max_min(self, intrin): + """ + Test intrinsics: + npyv_max_##sfx + npyv_maxp_##sfx + npyv_maxn_##sfx + npyv_min_##sfx + npyv_minp_##sfx + npyv_minn_##sfx + npyv_reduce_max_##sfx + npyv_reduce_maxp_##sfx + npyv_reduce_maxn_##sfx + npyv_reduce_min_##sfx + npyv_reduce_minp_##sfx + npyv_reduce_minn_##sfx + """ + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + chk_nan = {"xp": 1, "np": 1, "nn": 2, "xn": 2}.get(intrin[-2:], 0) + func = eval(intrin[:3]) + reduce_intrin = getattr(self, "reduce_" + intrin) + intrin = getattr(self, intrin) + hf_nlanes = self.nlanes//2 + + cases = ( + ([0.0, -0.0], [-0.0, 0.0]), + ([10, -10], [10, -10]), + ([pinf, 10], [10, ninf]), + ([10, pinf], [ninf, 10]), + ([10, -10], [10, -10]), + ([-10, 10], [-10, 10]) + ) + for op1, op2 in cases: + vdata_a = self.load(op1*hf_nlanes) + vdata_b = self.load(op2*hf_nlanes) + data = func(vdata_a, vdata_b) + simd = intrin(vdata_a, vdata_b) + assert simd == data + data = func(vdata_a) + simd = reduce_intrin(vdata_a) + assert simd == data + + if not chk_nan: + return + if chk_nan == 1: + test_nan = lambda a, b: ( + b if math.isnan(a) else a if math.isnan(b) else b + ) + else: + test_nan = lambda a, b: ( + nan if math.isnan(a) or math.isnan(b) else b + ) + cases = ( + (nan, 10), + (10, nan), + (nan, pinf), + (pinf, nan), + (nan, nan) + ) + for op1, op2 in cases: + vdata_ab = self.load([op1, op2]*hf_nlanes) + data = test_nan(op1, op2) + simd = reduce_intrin(vdata_ab) + assert simd == pytest.approx(data, nan_ok=True) + vdata_a = self.setall(op1) + vdata_b = self.setall(op2) + data = [data] * self.nlanes + simd = intrin(vdata_a, vdata_b) + assert simd == pytest.approx(data, nan_ok=True) + + def test_reciprocal(self): + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + data = self._data() + vdata = self.load(self._data()) + + recip_cases = ((nan, nan), (pinf, 0.0), (ninf, -0.0), (0.0, pinf), (-0.0, ninf)) + for case, desired in recip_cases: + data_recip = [desired]*self.nlanes + recip = self.recip(self.setall(case)) + assert recip == pytest.approx(data_recip, nan_ok=True) + + data_recip = self.load([1/x for x in data]) # load to truncate precision + recip = self.recip(vdata) + assert recip == data_recip + + def test_special_cases(self): + """ + Compare Not NaN. Test intrinsics: + npyv_notnan_##SFX + """ + nnan = self.notnan(self.setall(self._nan())) + assert nnan == [0]*self.nlanes + + @pytest.mark.parametrize("intrin_name", [ + "rint", "trunc", "ceil", "floor" + ]) + def test_unary_invalid_fpexception(self, intrin_name): + intrin = getattr(self, intrin_name) + for d in [float("nan"), float("inf"), -float("inf")]: + v = self.setall(d) + clear_floatstatus() + intrin(v) + assert check_floatstatus(invalid=True) is False + + @pytest.mark.parametrize('py_comp,np_comp', [ + (operator.lt, "cmplt"), + (operator.le, "cmple"), + (operator.gt, "cmpgt"), + (operator.ge, "cmpge"), + (operator.eq, "cmpeq"), + (operator.ne, "cmpneq") + ]) + def test_comparison_with_nan(self, py_comp, np_comp): + pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan() + mask_true = self._true_mask() + + def to_bool(vector): + return [lane == mask_true for lane in vector] + + intrin = getattr(self, np_comp) + cmp_cases = ((0, nan), (nan, 0), (nan, nan), (pinf, nan), + (ninf, nan), (-0.0, +0.0)) + for case_operand1, case_operand2 in cmp_cases: + data_a = [case_operand1]*self.nlanes + data_b = [case_operand2]*self.nlanes + vdata_a = self.setall(case_operand1) + vdata_b = self.setall(case_operand2) + vcmp = to_bool(intrin(vdata_a, vdata_b)) + data_cmp = [py_comp(a, b) for a, b in zip(data_a, data_b)] + assert vcmp == data_cmp + + @pytest.mark.parametrize("intrin", ["any", "all"]) + @pytest.mark.parametrize("data", ( + [float("nan"), 0], + [0, float("nan")], + [float("nan"), 1], + [1, float("nan")], + [float("nan"), float("nan")], + [0.0, -0.0], + [-0.0, 0.0], + [1.0, -0.0] + )) + def test_operators_crosstest(self, intrin, data): + """ + Test intrinsics: + npyv_any_##SFX + npyv_all_##SFX + """ + data_a = self.load(data * self.nlanes) + func = eval(intrin) + intrin = getattr(self, intrin) + desired = func(data_a) + simd = intrin(data_a) + assert not not simd == desired + +class _SIMD_ALL(_Test_Utility): + """ + To test all vector types at once + """ + def test_memory_load(self): + data = self._data() + # unaligned load + load_data = self.load(data) + assert load_data == data + # aligned load + loada_data = self.loada(data) + assert loada_data == data + # stream load + loads_data = self.loads(data) + assert loads_data == data + # load lower part + loadl = self.loadl(data) + loadl_half = list(loadl)[:self.nlanes//2] + data_half = data[:self.nlanes//2] + assert loadl_half == data_half + assert loadl != data # detect overflow + + def test_memory_store(self): + data = self._data() + vdata = self.load(data) + # unaligned store + store = [0] * self.nlanes + self.store(store, vdata) + assert store == data + # aligned store + store_a = [0] * self.nlanes + self.storea(store_a, vdata) + assert store_a == data + # stream store + store_s = [0] * self.nlanes + self.stores(store_s, vdata) + assert store_s == data + # store lower part + store_l = [0] * self.nlanes + self.storel(store_l, vdata) + assert store_l[:self.nlanes//2] == data[:self.nlanes//2] + assert store_l != vdata # detect overflow + # store higher part + store_h = [0] * self.nlanes + self.storeh(store_h, vdata) + assert store_h[:self.nlanes//2] == data[self.nlanes//2:] + assert store_h != vdata # detect overflow + + @pytest.mark.parametrize("intrin, elsizes, scale, fill", [ + ("self.load_tillz, self.load_till", (32, 64), 1, [0xffff]), + ("self.load2_tillz, self.load2_till", (32, 64), 2, [0xffff, 0x7fff]), + ]) + def test_memory_partial_load(self, intrin, elsizes, scale, fill): + if self._scalar_size() not in elsizes: + return + npyv_load_tillz, npyv_load_till = eval(intrin) + data = self._data() + lanes = list(range(1, self.nlanes + 1)) + lanes += [self.nlanes**2, self.nlanes**4] # test out of range + for n in lanes: + load_till = npyv_load_till(data, n, *fill) + load_tillz = npyv_load_tillz(data, n) + n *= scale + data_till = data[:n] + fill * ((self.nlanes-n) // scale) + assert load_till == data_till + data_tillz = data[:n] + [0] * (self.nlanes-n) + assert load_tillz == data_tillz + + @pytest.mark.parametrize("intrin, elsizes, scale", [ + ("self.store_till", (32, 64), 1), + ("self.store2_till", (32, 64), 2), + ]) + def test_memory_partial_store(self, intrin, elsizes, scale): + if self._scalar_size() not in elsizes: + return + npyv_store_till = eval(intrin) + data = self._data() + data_rev = self._data(reverse=True) + vdata = self.load(data) + lanes = list(range(1, self.nlanes + 1)) + lanes += [self.nlanes**2, self.nlanes**4] + for n in lanes: + data_till = data_rev.copy() + data_till[:n*scale] = data[:n*scale] + store_till = self._data(reverse=True) + npyv_store_till(store_till, n, vdata) + assert store_till == data_till + + @pytest.mark.parametrize("intrin, elsizes, scale", [ + ("self.loadn", (32, 64), 1), + ("self.loadn2", (32, 64), 2), + ]) + def test_memory_noncont_load(self, intrin, elsizes, scale): + if self._scalar_size() not in elsizes: + return + npyv_loadn = eval(intrin) + for stride in range(-64, 64): + if stride < 0: + data = self._data(stride, -stride*self.nlanes) + data_stride = list(itertools.chain( + *zip(*[data[-i::stride] for i in range(scale, 0, -1)]) + )) + elif stride == 0: + data = self._data() + data_stride = data[0:scale] * (self.nlanes//scale) + else: + data = self._data(count=stride*self.nlanes) + data_stride = list(itertools.chain( + *zip(*[data[i::stride] for i in range(scale)])) + ) + data_stride = self.load(data_stride) # cast unsigned + loadn = npyv_loadn(data, stride) + assert loadn == data_stride + + @pytest.mark.parametrize("intrin, elsizes, scale, fill", [ + ("self.loadn_tillz, self.loadn_till", (32, 64), 1, [0xffff]), + ("self.loadn2_tillz, self.loadn2_till", (32, 64), 2, [0xffff, 0x7fff]), + ]) + def test_memory_noncont_partial_load(self, intrin, elsizes, scale, fill): + if self._scalar_size() not in elsizes: + return + npyv_loadn_tillz, npyv_loadn_till = eval(intrin) + lanes = list(range(1, self.nlanes + 1)) + lanes += [self.nlanes**2, self.nlanes**4] + for stride in range(-64, 64): + if stride < 0: + data = self._data(stride, -stride*self.nlanes) + data_stride = list(itertools.chain( + *zip(*[data[-i::stride] for i in range(scale, 0, -1)]) + )) + elif stride == 0: + data = self._data() + data_stride = data[0:scale] * (self.nlanes//scale) + else: + data = self._data(count=stride*self.nlanes) + data_stride = list(itertools.chain( + *zip(*[data[i::stride] for i in range(scale)]) + )) + data_stride = list(self.load(data_stride)) # cast unsigned + for n in lanes: + nscale = n * scale + llanes = self.nlanes - nscale + data_stride_till = ( + data_stride[:nscale] + fill * (llanes//scale) + ) + loadn_till = npyv_loadn_till(data, stride, n, *fill) + assert loadn_till == data_stride_till + data_stride_tillz = data_stride[:nscale] + [0] * llanes + loadn_tillz = npyv_loadn_tillz(data, stride, n) + assert loadn_tillz == data_stride_tillz + + @pytest.mark.parametrize("intrin, elsizes, scale", [ + ("self.storen", (32, 64), 1), + ("self.storen2", (32, 64), 2), + ]) + def test_memory_noncont_store(self, intrin, elsizes, scale): + if self._scalar_size() not in elsizes: + return + npyv_storen = eval(intrin) + data = self._data() + vdata = self.load(data) + hlanes = self.nlanes // scale + for stride in range(1, 64): + data_storen = [0xff] * stride * self.nlanes + for s in range(0, hlanes*stride, stride): + i = (s//stride)*scale + data_storen[s:s+scale] = data[i:i+scale] + storen = [0xff] * stride * self.nlanes + storen += [0x7f]*64 + npyv_storen(storen, stride, vdata) + assert storen[:-64] == data_storen + assert storen[-64:] == [0x7f]*64 # detect overflow + + for stride in range(-64, 0): + data_storen = [0xff] * -stride * self.nlanes + for s in range(0, hlanes*stride, stride): + i = (s//stride)*scale + data_storen[s-scale:s or None] = data[i:i+scale] + storen = [0x7f]*64 + storen += [0xff] * -stride * self.nlanes + npyv_storen(storen, stride, vdata) + assert storen[64:] == data_storen + assert storen[:64] == [0x7f]*64 # detect overflow + # stride 0 + data_storen = [0x7f] * self.nlanes + storen = data_storen.copy() + data_storen[0:scale] = data[-scale:] + npyv_storen(storen, 0, vdata) + assert storen == data_storen + + @pytest.mark.parametrize("intrin, elsizes, scale", [ + ("self.storen_till", (32, 64), 1), + ("self.storen2_till", (32, 64), 2), + ]) + def test_memory_noncont_partial_store(self, intrin, elsizes, scale): + if self._scalar_size() not in elsizes: + return + npyv_storen_till = eval(intrin) + data = self._data() + vdata = self.load(data) + lanes = list(range(1, self.nlanes + 1)) + lanes += [self.nlanes**2, self.nlanes**4] + hlanes = self.nlanes // scale + for stride in range(1, 64): + for n in lanes: + data_till = [0xff] * stride * self.nlanes + tdata = data[:n*scale] + [0xff] * (self.nlanes-n*scale) + for s in range(0, hlanes*stride, stride)[:n]: + i = (s//stride)*scale + data_till[s:s+scale] = tdata[i:i+scale] + storen_till = [0xff] * stride * self.nlanes + storen_till += [0x7f]*64 + npyv_storen_till(storen_till, stride, n, vdata) + assert storen_till[:-64] == data_till + assert storen_till[-64:] == [0x7f]*64 # detect overflow + + for stride in range(-64, 0): + for n in lanes: + data_till = [0xff] * -stride * self.nlanes + tdata = data[:n*scale] + [0xff] * (self.nlanes-n*scale) + for s in range(0, hlanes*stride, stride)[:n]: + i = (s//stride)*scale + data_till[s-scale:s or None] = tdata[i:i+scale] + storen_till = [0x7f]*64 + storen_till += [0xff] * -stride * self.nlanes + npyv_storen_till(storen_till, stride, n, vdata) + assert storen_till[64:] == data_till + assert storen_till[:64] == [0x7f]*64 # detect overflow + + # stride 0 + for n in lanes: + data_till = [0x7f] * self.nlanes + storen_till = data_till.copy() + data_till[0:scale] = data[:n*scale][-scale:] + npyv_storen_till(storen_till, 0, n, vdata) + assert storen_till == data_till + + @pytest.mark.parametrize("intrin, table_size, elsize", [ + ("self.lut32", 32, 32), + ("self.lut16", 16, 64) + ]) + def test_lut(self, intrin, table_size, elsize): + """ + Test lookup table intrinsics: + npyv_lut32_##sfx + npyv_lut16_##sfx + """ + if elsize != self._scalar_size(): + return + intrin = eval(intrin) + idx_itrin = getattr(self.npyv, f"setall_u{elsize}") + table = range(0, table_size) + for i in table: + broadi = self.setall(i) + idx = idx_itrin(i) + lut = intrin(table, idx) + assert lut == broadi + + def test_misc(self): + broadcast_zero = self.zero() + assert broadcast_zero == [0] * self.nlanes + for i in range(1, 10): + broadcasti = self.setall(i) + assert broadcasti == [i] * self.nlanes + + data_a, data_b = self._data(), self._data(reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + # py level of npyv_set_* don't support ignoring the extra specified lanes or + # fill non-specified lanes with zero. + vset = self.set(*data_a) + assert vset == data_a + # py level of npyv_setf_* don't support ignoring the extra specified lanes or + # fill non-specified lanes with the specified scalar. + vsetf = self.setf(10, *data_a) + assert vsetf == data_a + + # We're testing the sanity of _simd's type-vector, + # reinterpret* intrinsics itself are tested via compiler + # during the build of _simd module + sfxes = ["u8", "s8", "u16", "s16", "u32", "s32", "u64", "s64"] + if self.npyv.simd_f64: + sfxes.append("f64") + if self.npyv.simd_f32: + sfxes.append("f32") + for sfx in sfxes: + vec_name = getattr(self, "reinterpret_" + sfx)(vdata_a).__name__ + assert vec_name == "npyv_" + sfx + + # select & mask operations + select_a = self.select(self.cmpeq(self.zero(), self.zero()), vdata_a, vdata_b) + assert select_a == data_a + select_b = self.select(self.cmpneq(self.zero(), self.zero()), vdata_a, vdata_b) + assert select_b == data_b + + # test extract elements + assert self.extract0(vdata_b) == vdata_b[0] + + # cleanup intrinsic is only used with AVX for + # zeroing registers to avoid the AVX-SSE transition penalty, + # so nothing to test here + self.npyv.cleanup() + + def test_reorder(self): + data_a, data_b = self._data(), self._data(reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + # lower half part + data_a_lo = data_a[:self.nlanes//2] + data_b_lo = data_b[:self.nlanes//2] + # higher half part + data_a_hi = data_a[self.nlanes//2:] + data_b_hi = data_b[self.nlanes//2:] + # combine two lower parts + combinel = self.combinel(vdata_a, vdata_b) + assert combinel == data_a_lo + data_b_lo + # combine two higher parts + combineh = self.combineh(vdata_a, vdata_b) + assert combineh == data_a_hi + data_b_hi + # combine x2 + combine = self.combine(vdata_a, vdata_b) + assert combine == (data_a_lo + data_b_lo, data_a_hi + data_b_hi) + + # zip(interleave) + data_zipl = self.load([ + v for p in zip(data_a_lo, data_b_lo) for v in p + ]) + data_ziph = self.load([ + v for p in zip(data_a_hi, data_b_hi) for v in p + ]) + vzip = self.zip(vdata_a, vdata_b) + assert vzip == (data_zipl, data_ziph) + vzip = [0]*self.nlanes*2 + self._x2("store")(vzip, (vdata_a, vdata_b)) + assert vzip == list(data_zipl) + list(data_ziph) + + # unzip(deinterleave) + unzip = self.unzip(data_zipl, data_ziph) + assert unzip == (data_a, data_b) + unzip = self._x2("load")(list(data_zipl) + list(data_ziph)) + assert unzip == (data_a, data_b) + + def test_reorder_rev64(self): + # Reverse elements of each 64-bit lane + ssize = self._scalar_size() + if ssize == 64: + return + data_rev64 = [ + y for x in range(0, self.nlanes, 64//ssize) + for y in reversed(range(x, x + 64//ssize)) + ] + rev64 = self.rev64(self.load(range(self.nlanes))) + assert rev64 == data_rev64 + + def test_reorder_permi128(self): + """ + Test permuting elements for each 128-bit lane. + npyv_permi128_##sfx + """ + ssize = self._scalar_size() + if ssize < 32: + return + data = self.load(self._data()) + permn = 128//ssize + permd = permn-1 + nlane128 = self.nlanes//permn + shfl = [0, 1] if ssize == 64 else [0, 2, 4, 6] + for i in range(permn): + indices = [(i >> shf) & permd for shf in shfl] + vperm = self.permi128(data, *indices) + data_vperm = [ + data[j + (e & -permn)] + for e, j in enumerate(indices*nlane128) + ] + assert vperm == data_vperm + + @pytest.mark.parametrize('func, intrin', [ + (operator.lt, "cmplt"), + (operator.le, "cmple"), + (operator.gt, "cmpgt"), + (operator.ge, "cmpge"), + (operator.eq, "cmpeq") + ]) + def test_operators_comparison(self, func, intrin): + if self._is_fp(): + data_a = self._data() + else: + data_a = self._data(self._int_max() - self.nlanes) + data_b = self._data(self._int_min(), reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + intrin = getattr(self, intrin) + + mask_true = self._true_mask() + def to_bool(vector): + return [lane == mask_true for lane in vector] + + data_cmp = [func(a, b) for a, b in zip(data_a, data_b)] + cmp = to_bool(intrin(vdata_a, vdata_b)) + assert cmp == data_cmp + + def test_operators_logical(self): + if self._is_fp(): + data_a = self._data() + else: + data_a = self._data(self._int_max() - self.nlanes) + data_b = self._data(self._int_min(), reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + if self._is_fp(): + data_cast_a = self._to_unsigned(vdata_a) + data_cast_b = self._to_unsigned(vdata_b) + cast, cast_data = self._to_unsigned, self._to_unsigned + else: + data_cast_a, data_cast_b = data_a, data_b + cast, cast_data = lambda a: a, self.load + + data_xor = cast_data([a ^ b for a, b in zip(data_cast_a, data_cast_b)]) + vxor = cast(self.xor(vdata_a, vdata_b)) + assert vxor == data_xor + + data_or = cast_data([a | b for a, b in zip(data_cast_a, data_cast_b)]) + vor = cast(getattr(self, "or")(vdata_a, vdata_b)) + assert vor == data_or + + data_and = cast_data([a & b for a, b in zip(data_cast_a, data_cast_b)]) + vand = cast(getattr(self, "and")(vdata_a, vdata_b)) + assert vand == data_and + + data_not = cast_data([~a for a in data_cast_a]) + vnot = cast(getattr(self, "not")(vdata_a)) + assert vnot == data_not + + if self.sfx not in ("u8"): + return + data_andc = [a & ~b for a, b in zip(data_cast_a, data_cast_b)] + vandc = cast(self.andc(vdata_a, vdata_b)) + assert vandc == data_andc + + @pytest.mark.parametrize("intrin", ["any", "all"]) + @pytest.mark.parametrize("data", ( + [1, 2, 3, 4], + [-1, -2, -3, -4], + [0, 1, 2, 3, 4], + [0x7f, 0x7fff, 0x7fffffff, 0x7fffffffffffffff], + [0, -1, -2, -3, 4], + [0], + [1], + [-1] + )) + def test_operators_crosstest(self, intrin, data): + """ + Test intrinsics: + npyv_any_##SFX + npyv_all_##SFX + """ + data_a = self.load(data * self.nlanes) + func = eval(intrin) + intrin = getattr(self, intrin) + desired = func(data_a) + simd = intrin(data_a) + assert not not simd == desired + + def test_conversion_boolean(self): + bsfx = "b" + self.sfx[1:] + to_boolean = getattr(self.npyv, "cvt_%s_%s" % (bsfx, self.sfx)) + from_boolean = getattr(self.npyv, "cvt_%s_%s" % (self.sfx, bsfx)) + + false_vb = to_boolean(self.setall(0)) + true_vb = self.cmpeq(self.setall(0), self.setall(0)) + assert false_vb != true_vb + + false_vsfx = from_boolean(false_vb) + true_vsfx = from_boolean(true_vb) + assert false_vsfx != true_vsfx + + def test_conversion_expand(self): + """ + Test expand intrinsics: + npyv_expand_u16_u8 + npyv_expand_u32_u16 + """ + if self.sfx not in ("u8", "u16"): + return + totype = self.sfx[0]+str(int(self.sfx[1:])*2) + expand = getattr(self.npyv, f"expand_{totype}_{self.sfx}") + # close enough from the edge to detect any deviation + data = self._data(self._int_max() - self.nlanes) + vdata = self.load(data) + edata = expand(vdata) + # lower half part + data_lo = data[:self.nlanes//2] + # higher half part + data_hi = data[self.nlanes//2:] + assert edata == (data_lo, data_hi) + + def test_arithmetic_subadd(self): + if self._is_fp(): + data_a = self._data() + else: + data_a = self._data(self._int_max() - self.nlanes) + data_b = self._data(self._int_min(), reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + # non-saturated + data_add = self.load([a + b for a, b in zip(data_a, data_b)]) # load to cast + add = self.add(vdata_a, vdata_b) + assert add == data_add + data_sub = self.load([a - b for a, b in zip(data_a, data_b)]) + sub = self.sub(vdata_a, vdata_b) + assert sub == data_sub + + def test_arithmetic_mul(self): + if self.sfx in ("u64", "s64"): + return + + if self._is_fp(): + data_a = self._data() + else: + data_a = self._data(self._int_max() - self.nlanes) + data_b = self._data(self._int_min(), reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + data_mul = self.load([a * b for a, b in zip(data_a, data_b)]) + mul = self.mul(vdata_a, vdata_b) + assert mul == data_mul + + def test_arithmetic_div(self): + if not self._is_fp(): + return + + data_a, data_b = self._data(), self._data(reverse=True) + vdata_a, vdata_b = self.load(data_a), self.load(data_b) + + # load to truncate f64 to precision of f32 + data_div = self.load([a / b for a, b in zip(data_a, data_b)]) + div = self.div(vdata_a, vdata_b) + assert div == data_div + + def test_arithmetic_intdiv(self): + """ + Test integer division intrinsics: + npyv_divisor_##sfx + npyv_divc_##sfx + """ + if self._is_fp(): + return + + int_min = self._int_min() + def trunc_div(a, d): + """ + Divide towards zero works with large integers > 2^53, + and wrap around overflow similar to what C does. + """ + if d == -1 and a == int_min: + return a + sign_a, sign_d = a < 0, d < 0 + if a == 0 or sign_a == sign_d: + return a // d + return (a + sign_d - sign_a) // d + 1 + + data = [1, -int_min] # to test overflow + data += range(0, 2**8, 2**5) + data += range(0, 2**8, 2**5-1) + bsize = self._scalar_size() + if bsize > 8: + data += range(2**8, 2**16, 2**13) + data += range(2**8, 2**16, 2**13-1) + if bsize > 16: + data += range(2**16, 2**32, 2**29) + data += range(2**16, 2**32, 2**29-1) + if bsize > 32: + data += range(2**32, 2**64, 2**61) + data += range(2**32, 2**64, 2**61-1) + # negate + data += [-x for x in data] + for dividend, divisor in itertools.product(data, data): + divisor = self.setall(divisor)[0] # cast + if divisor == 0: + continue + dividend = self.load(self._data(dividend)) + data_divc = [trunc_div(a, divisor) for a in dividend] + divisor_parms = self.divisor(divisor) + divc = self.divc(dividend, divisor_parms) + assert divc == data_divc + + def test_arithmetic_reduce_sum(self): + """ + Test reduce sum intrinsics: + npyv_sum_##sfx + """ + if self.sfx not in ("u32", "u64", "f32", "f64"): + return + # reduce sum + data = self._data() + vdata = self.load(data) + + data_sum = sum(data) + vsum = self.sum(vdata) + assert vsum == data_sum + + def test_arithmetic_reduce_sumup(self): + """ + Test extend reduce sum intrinsics: + npyv_sumup_##sfx + """ + if self.sfx not in ("u8", "u16"): + return + rdata = (0, self.nlanes, self._int_min(), self._int_max()-self.nlanes) + for r in rdata: + data = self._data(r) + vdata = self.load(data) + data_sum = sum(data) + vsum = self.sumup(vdata) + assert vsum == data_sum + + def test_mask_conditional(self): + """ + Conditional addition and subtraction for all supported data types. + Test intrinsics: + npyv_ifadd_##SFX, npyv_ifsub_##SFX + """ + vdata_a = self.load(self._data()) + vdata_b = self.load(self._data(reverse=True)) + true_mask = self.cmpeq(self.zero(), self.zero()) + false_mask = self.cmpneq(self.zero(), self.zero()) + + data_sub = self.sub(vdata_b, vdata_a) + ifsub = self.ifsub(true_mask, vdata_b, vdata_a, vdata_b) + assert ifsub == data_sub + ifsub = self.ifsub(false_mask, vdata_a, vdata_b, vdata_b) + assert ifsub == vdata_b + + data_add = self.add(vdata_b, vdata_a) + ifadd = self.ifadd(true_mask, vdata_b, vdata_a, vdata_b) + assert ifadd == data_add + ifadd = self.ifadd(false_mask, vdata_a, vdata_b, vdata_b) + assert ifadd == vdata_b + + if not self._is_fp(): + return + data_div = self.div(vdata_b, vdata_a) + ifdiv = self.ifdiv(true_mask, vdata_b, vdata_a, vdata_b) + assert ifdiv == data_div + ifdivz = self.ifdivz(true_mask, vdata_b, vdata_a) + assert ifdivz == data_div + ifdiv = self.ifdiv(false_mask, vdata_a, vdata_b, vdata_b) + assert ifdiv == vdata_b + ifdivz = self.ifdivz(false_mask, vdata_a, vdata_b) + assert ifdivz == self.zero() + +bool_sfx = ("b8", "b16", "b32", "b64") +int_sfx = ("u8", "s8", "u16", "s16", "u32", "s32", "u64", "s64") +fp_sfx = ("f32", "f64") +all_sfx = int_sfx + fp_sfx +tests_registry = { + bool_sfx: _SIMD_BOOL, + int_sfx : _SIMD_INT, + fp_sfx : _SIMD_FP, + ("f32",): _SIMD_FP32, + ("f64",): _SIMD_FP64, + all_sfx : _SIMD_ALL +} +for target_name, npyv in targets.items(): + simd_width = npyv.simd if npyv else '' + pretty_name = target_name.split('__') # multi-target separator + if len(pretty_name) > 1: + # multi-target + pretty_name = f"({' '.join(pretty_name)})" + else: + pretty_name = pretty_name[0] + + skip = "" + skip_sfx = dict() + if not npyv: + skip = f"target '{pretty_name}' isn't supported by current machine" + elif not npyv.simd: + skip = f"target '{pretty_name}' isn't supported by NPYV" + else: + if not npyv.simd_f32: + skip_sfx["f32"] = f"target '{pretty_name}' "\ + "doesn't support single-precision" + if not npyv.simd_f64: + skip_sfx["f64"] = f"target '{pretty_name}' doesn't"\ + "support double-precision" + + for sfxes, cls in tests_registry.items(): + for sfx in sfxes: + skip_m = skip_sfx.get(sfx, skip) + inhr = (cls,) + attr = dict(npyv=targets[target_name], sfx=sfx, target_name=target_name) + tcls = type(f"Test{cls.__name__}_{simd_width}_{target_name}_{sfx}", inhr, attr) + if skip_m: + pytest.mark.skip(reason=skip_m)(tcls) + globals()[tcls.__name__] = tcls diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_simd_module.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_simd_module.py new file mode 100644 index 0000000000000000000000000000000000000000..0b9729e5b6401f1108635ae188d5cb340214ee84 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_simd_module.py @@ -0,0 +1,101 @@ +import pytest +from numpy._core._simd import targets +""" +This testing unit only for checking the sanity of common functionality, +therefore all we need is just to take one submodule that represents any +of enabled SIMD extensions to run the test on it and the second submodule +required to run only one check related to the possibility of mixing +the data types among each submodule. +""" +npyvs = [npyv_mod for npyv_mod in targets.values() if npyv_mod and npyv_mod.simd] +npyv, npyv2 = (npyvs + [None, None])[:2] + +unsigned_sfx = ["u8", "u16", "u32", "u64"] +signed_sfx = ["s8", "s16", "s32", "s64"] +fp_sfx = [] +if npyv and npyv.simd_f32: + fp_sfx.append("f32") +if npyv and npyv.simd_f64: + fp_sfx.append("f64") + +int_sfx = unsigned_sfx + signed_sfx +all_sfx = unsigned_sfx + int_sfx + +@pytest.mark.skipif(not npyv, reason="could not find any SIMD extension with NPYV support") +class Test_SIMD_MODULE: + + @pytest.mark.parametrize('sfx', all_sfx) + def test_num_lanes(self, sfx): + nlanes = getattr(npyv, "nlanes_" + sfx) + vector = getattr(npyv, "setall_" + sfx)(1) + assert len(vector) == nlanes + + @pytest.mark.parametrize('sfx', all_sfx) + def test_type_name(self, sfx): + vector = getattr(npyv, "setall_" + sfx)(1) + assert vector.__name__ == "npyv_" + sfx + + def test_raises(self): + a, b = [npyv.setall_u32(1)]*2 + for sfx in all_sfx: + vcb = lambda intrin: getattr(npyv, f"{intrin}_{sfx}") + pytest.raises(TypeError, vcb("add"), a) + pytest.raises(TypeError, vcb("add"), a, b, a) + pytest.raises(TypeError, vcb("setall")) + pytest.raises(TypeError, vcb("setall"), [1]) + pytest.raises(TypeError, vcb("load"), 1) + pytest.raises(ValueError, vcb("load"), [1]) + pytest.raises(ValueError, vcb("store"), [1], getattr(npyv, f"reinterpret_{sfx}_u32")(a)) + + @pytest.mark.skipif(not npyv2, reason=( + "could not find a second SIMD extension with NPYV support" + )) + def test_nomix(self): + # mix among submodules isn't allowed + a = npyv.setall_u32(1) + a2 = npyv2.setall_u32(1) + pytest.raises(TypeError, npyv.add_u32, a2, a2) + pytest.raises(TypeError, npyv2.add_u32, a, a) + + @pytest.mark.parametrize('sfx', unsigned_sfx) + def test_unsigned_overflow(self, sfx): + nlanes = getattr(npyv, "nlanes_" + sfx) + maxu = (1 << int(sfx[1:])) - 1 + maxu_72 = (1 << 72) - 1 + lane = getattr(npyv, "setall_" + sfx)(maxu_72)[0] + assert lane == maxu + lanes = getattr(npyv, "load_" + sfx)([maxu_72] * nlanes) + assert lanes == [maxu] * nlanes + lane = getattr(npyv, "setall_" + sfx)(-1)[0] + assert lane == maxu + lanes = getattr(npyv, "load_" + sfx)([-1] * nlanes) + assert lanes == [maxu] * nlanes + + @pytest.mark.parametrize('sfx', signed_sfx) + def test_signed_overflow(self, sfx): + nlanes = getattr(npyv, "nlanes_" + sfx) + maxs_72 = (1 << 71) - 1 + lane = getattr(npyv, "setall_" + sfx)(maxs_72)[0] + assert lane == -1 + lanes = getattr(npyv, "load_" + sfx)([maxs_72] * nlanes) + assert lanes == [-1] * nlanes + mins_72 = -1 << 71 + lane = getattr(npyv, "setall_" + sfx)(mins_72)[0] + assert lane == 0 + lanes = getattr(npyv, "load_" + sfx)([mins_72] * nlanes) + assert lanes == [0] * nlanes + + def test_truncate_f32(self): + if not npyv.simd_f32: + pytest.skip("F32 isn't support by the SIMD extension") + f32 = npyv.setall_f32(0.1)[0] + assert f32 != 0.1 + assert round(f32, 1) == 0.1 + + def test_compare(self): + data_range = range(0, npyv.nlanes_u32) + vdata = npyv.load_u32(data_range) + assert vdata == list(data_range) + assert vdata == tuple(data_range) + for i in data_range: + assert vdata[i] == data_range[i] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_stringdtype.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_stringdtype.py new file mode 100644 index 0000000000000000000000000000000000000000..fd638126cbd20e897eadaaf3b2d0a1e75f111072 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_stringdtype.py @@ -0,0 +1,1833 @@ +import concurrent.futures +import copy +import itertools +import os +import string +import pickle +import sys +import tempfile + +import numpy as np +import pytest + +from numpy.dtypes import StringDType +from numpy._core.tests._natype import pd_NA +from numpy.testing import assert_array_equal, IS_WASM, IS_PYPY + + +@pytest.fixture +def string_list(): + return ["abc", "def", "ghi" * 10, "A¢☃€ 😊" * 100, "Abc" * 1000, "DEF"] + + +@pytest.fixture +def random_string_list(): + chars = list(string.ascii_letters + string.digits) + chars = np.array(chars, dtype="U1") + ret = np.random.choice(chars, size=100 * 10, replace=True) + return ret.view("U100") + + +@pytest.fixture(params=[True, False]) +def coerce(request): + return request.param + + +@pytest.fixture( + params=["unset", None, pd_NA, np.nan, float("nan"), "__nan__"], + ids=["unset", "None", "pandas.NA", "np.nan", "float('nan')", "string nan"], +) +def na_object(request): + return request.param + + +def get_dtype(na_object, coerce=True): + # explicit is check for pd_NA because != with pd_NA returns pd_NA + if na_object is pd_NA or na_object != "unset": + return StringDType(na_object=na_object, coerce=coerce) + else: + return StringDType(coerce=coerce) + + +@pytest.fixture() +def dtype(na_object, coerce): + return get_dtype(na_object, coerce) + + +# second copy for cast tests to do a cartesian product over dtypes +@pytest.fixture(params=[True, False]) +def coerce2(request): + return request.param + + +@pytest.fixture( + params=["unset", None, pd_NA, np.nan, float("nan"), "__nan__"], + ids=["unset", "None", "pandas.NA", "np.nan", "float('nan')", "string nan"], +) +def na_object2(request): + return request.param + + +@pytest.fixture() +def dtype2(na_object2, coerce2): + # explicit is check for pd_NA because != with pd_NA returns pd_NA + if na_object2 is pd_NA or na_object2 != "unset": + return StringDType(na_object=na_object2, coerce=coerce2) + else: + return StringDType(coerce=coerce2) + + +def test_dtype_creation(): + hashes = set() + dt = StringDType() + assert not hasattr(dt, "na_object") and dt.coerce is True + hashes.add(hash(dt)) + + dt = StringDType(na_object=None) + assert dt.na_object is None and dt.coerce is True + hashes.add(hash(dt)) + + dt = StringDType(coerce=False) + assert not hasattr(dt, "na_object") and dt.coerce is False + hashes.add(hash(dt)) + + dt = StringDType(na_object=None, coerce=False) + assert dt.na_object is None and dt.coerce is False + hashes.add(hash(dt)) + + assert len(hashes) == 4 + + dt = np.dtype("T") + assert dt == StringDType() + assert dt.kind == "T" + assert dt.char == "T" + + hashes.add(hash(dt)) + assert len(hashes) == 4 + + +def test_dtype_equality(dtype): + assert dtype == dtype + for ch in "SU": + assert dtype != np.dtype(ch) + assert dtype != np.dtype(f"{ch}8") + + +def test_dtype_repr(dtype): + if not hasattr(dtype, "na_object") and dtype.coerce: + assert repr(dtype) == "StringDType()" + elif dtype.coerce: + assert repr(dtype) == f"StringDType(na_object={dtype.na_object!r})" + elif not hasattr(dtype, "na_object"): + assert repr(dtype) == "StringDType(coerce=False)" + else: + assert ( + repr(dtype) + == f"StringDType(na_object={dtype.na_object!r}, coerce=False)" + ) + + +def test_create_with_na(dtype): + if not hasattr(dtype, "na_object"): + pytest.skip("does not have an na object") + na_val = dtype.na_object + string_list = ["hello", na_val, "world"] + arr = np.array(string_list, dtype=dtype) + assert str(arr) == "[" + " ".join([repr(s) for s in string_list]) + "]" + assert arr[1] is dtype.na_object + + +@pytest.mark.parametrize("i", list(range(5))) +def test_set_replace_na(i): + # Test strings of various lengths can be set to NaN and then replaced. + s_empty = "" + s_short = "0123456789" + s_medium = "abcdefghijklmnopqrstuvwxyz" + s_long = "-=+" * 100 + strings = [s_medium, s_empty, s_short, s_medium, s_long] + a = np.array(strings, StringDType(na_object=np.nan)) + for s in [a[i], s_medium+s_short, s_short, s_empty, s_long]: + a[i] = np.nan + assert np.isnan(a[i]) + a[i] = s + assert a[i] == s + assert_array_equal(a, strings[:i] + [s] + strings[i+1:]) + + +def test_null_roundtripping(): + data = ["hello\0world", "ABC\0DEF\0\0"] + arr = np.array(data, dtype="T") + assert data[0] == arr[0] + assert data[1] == arr[1] + + +def test_string_too_large_error(): + arr = np.array(["a", "b", "c"], dtype=StringDType()) + with pytest.raises(MemoryError): + arr * (2**63 - 2) + + +@pytest.mark.parametrize( + "data", + [ + ["abc", "def", "ghi"], + ["🤣", "📵", "😰"], + ["🚜", "🙃", "😾"], + ["😹", "🚠", "🚌"], + ], +) +def test_array_creation_utf8(dtype, data): + arr = np.array(data, dtype=dtype) + assert str(arr) == "[" + " ".join(["'" + str(d) + "'" for d in data]) + "]" + assert arr.dtype == dtype + + +@pytest.mark.parametrize( + "data", + [ + [1, 2, 3], + [b"abc", b"def", b"ghi"], + [object, object, object], + ], +) +def test_scalars_string_conversion(data, dtype): + try: + str_vals = [str(d.decode('utf-8')) for d in data] + except AttributeError: + str_vals = [str(d) for d in data] + if dtype.coerce: + assert_array_equal( + np.array(data, dtype=dtype), + np.array(str_vals, dtype=dtype), + ) + else: + with pytest.raises(ValueError): + np.array(data, dtype=dtype) + + +@pytest.mark.parametrize( + ("strings"), + [ + ["this", "is", "an", "array"], + ["€", "", "😊"], + ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"], + ], +) +def test_self_casts(dtype, dtype2, strings): + if hasattr(dtype, "na_object"): + strings = strings + [dtype.na_object] + elif hasattr(dtype2, "na_object"): + strings = strings + [""] + arr = np.array(strings, dtype=dtype) + newarr = arr.astype(dtype2) + + if hasattr(dtype, "na_object") and not hasattr(dtype2, "na_object"): + assert newarr[-1] == str(dtype.na_object) + with pytest.raises(TypeError): + arr.astype(dtype2, casting="safe") + elif hasattr(dtype, "na_object") and hasattr(dtype2, "na_object"): + assert newarr[-1] is dtype2.na_object + arr.astype(dtype2, casting="safe") + elif hasattr(dtype2, "na_object"): + assert newarr[-1] == "" + arr.astype(dtype2, casting="safe") + else: + arr.astype(dtype2, casting="safe") + + if hasattr(dtype, "na_object") and hasattr(dtype2, "na_object"): + na1 = dtype.na_object + na2 = dtype2.na_object + if (na1 is not na2 and + # check for pd_NA first because bool(pd_NA) is an error + ((na1 is pd_NA or na2 is pd_NA) or + # the second check is a NaN check, spelled this way + # to avoid errors from math.isnan and np.isnan + (na1 != na2 and not (na1 != na1 and na2 != na2)))): + with pytest.raises(TypeError): + arr[:-1] == newarr[:-1] + return + assert_array_equal(arr[:-1], newarr[:-1]) + + +@pytest.mark.parametrize( + ("strings"), + [ + ["this", "is", "an", "array"], + ["€", "", "😊"], + ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"], + ], +) +class TestStringLikeCasts: + def test_unicode_casts(self, dtype, strings): + arr = np.array(strings, dtype=np.str_).astype(dtype) + expected = np.array(strings, dtype=dtype) + assert_array_equal(arr, expected) + + arr_as_U8 = expected.astype("U8") + assert_array_equal(arr_as_U8, np.array(strings, dtype="U8")) + assert_array_equal(arr_as_U8.astype(dtype), arr) + arr_as_U3 = expected.astype("U3") + assert_array_equal(arr_as_U3, np.array(strings, dtype="U3")) + assert_array_equal( + arr_as_U3.astype(dtype), + np.array([s[:3] for s in strings], dtype=dtype), + ) + + def test_void_casts(self, dtype, strings): + sarr = np.array(strings, dtype=dtype) + utf8_bytes = [s.encode("utf-8") for s in strings] + void_dtype = f"V{max([len(s) for s in utf8_bytes])}" + varr = np.array(utf8_bytes, dtype=void_dtype) + assert_array_equal(varr, sarr.astype(void_dtype)) + assert_array_equal(varr.astype(dtype), sarr) + + def test_bytes_casts(self, dtype, strings): + sarr = np.array(strings, dtype=dtype) + try: + utf8_bytes = [s.encode("ascii") for s in strings] + bytes_dtype = f"S{max([len(s) for s in utf8_bytes])}" + barr = np.array(utf8_bytes, dtype=bytes_dtype) + assert_array_equal(barr, sarr.astype(bytes_dtype)) + assert_array_equal(barr.astype(dtype), sarr) + if dtype.coerce: + barr = np.array(utf8_bytes, dtype=dtype) + assert_array_equal(barr, sarr) + barr = np.array(utf8_bytes, dtype="O") + assert_array_equal(barr.astype(dtype), sarr) + else: + with pytest.raises(ValueError): + np.array(utf8_bytes, dtype=dtype) + except UnicodeEncodeError: + with pytest.raises(UnicodeEncodeError): + sarr.astype("S20") + + +def test_additional_unicode_cast(random_string_list, dtype): + arr = np.array(random_string_list, dtype=dtype) + # test that this short-circuits correctly + assert_array_equal(arr, arr.astype(arr.dtype)) + # tests the casts via the comparison promoter + assert_array_equal(arr, arr.astype(random_string_list.dtype)) + + +def test_insert_scalar(dtype, string_list): + """Test that inserting a scalar works.""" + arr = np.array(string_list, dtype=dtype) + scalar_instance = "what" + arr[1] = scalar_instance + assert_array_equal( + arr, + np.array(string_list[:1] + ["what"] + string_list[2:], dtype=dtype), + ) + + +comparison_operators = [ + np.equal, + np.not_equal, + np.greater, + np.greater_equal, + np.less, + np.less_equal, +] + + +@pytest.mark.parametrize("op", comparison_operators) +@pytest.mark.parametrize("o_dtype", [np.str_, object, StringDType()]) +def test_comparisons(string_list, dtype, op, o_dtype): + sarr = np.array(string_list, dtype=dtype) + oarr = np.array(string_list, dtype=o_dtype) + + # test that comparison operators work + res = op(sarr, sarr) + ores = op(oarr, oarr) + # test that promotion works as well + orres = op(sarr, oarr) + olres = op(oarr, sarr) + + assert_array_equal(res, ores) + assert_array_equal(res, orres) + assert_array_equal(res, olres) + + # test we get the correct answer for unequal length strings + sarr2 = np.array([s + "2" for s in string_list], dtype=dtype) + oarr2 = np.array([s + "2" for s in string_list], dtype=o_dtype) + + res = op(sarr, sarr2) + ores = op(oarr, oarr2) + olres = op(oarr, sarr2) + orres = op(sarr, oarr2) + + assert_array_equal(res, ores) + assert_array_equal(res, olres) + assert_array_equal(res, orres) + + res = op(sarr2, sarr) + ores = op(oarr2, oarr) + olres = op(oarr2, sarr) + orres = op(sarr2, oarr) + + assert_array_equal(res, ores) + assert_array_equal(res, olres) + assert_array_equal(res, orres) + + +def test_isnan(dtype, string_list): + if not hasattr(dtype, "na_object"): + pytest.skip("no na support") + sarr = np.array(string_list + [dtype.na_object], dtype=dtype) + is_nan = isinstance(dtype.na_object, float) and np.isnan(dtype.na_object) + bool_errors = 0 + try: + bool(dtype.na_object) + except TypeError: + bool_errors = 1 + if is_nan or bool_errors: + # isnan is only true when na_object is a NaN + assert_array_equal( + np.isnan(sarr), + np.array([0] * len(string_list) + [1], dtype=np.bool), + ) + else: + assert not np.any(np.isnan(sarr)) + + +def test_pickle(dtype, string_list): + arr = np.array(string_list, dtype=dtype) + + with tempfile.NamedTemporaryFile("wb", delete=False) as f: + pickle.dump([arr, dtype], f) + + with open(f.name, "rb") as f: + res = pickle.load(f) + + assert_array_equal(res[0], arr) + assert res[1] == dtype + + os.remove(f.name) + + +def test_stdlib_copy(dtype, string_list): + arr = np.array(string_list, dtype=dtype) + + assert_array_equal(copy.copy(arr), arr) + assert_array_equal(copy.deepcopy(arr), arr) + + +@pytest.mark.parametrize( + "strings", + [ + ["left", "right", "leftovers", "righty", "up", "down"], + [ + "left" * 10, + "right" * 10, + "leftovers" * 10, + "righty" * 10, + "up" * 10, + ], + ["🤣🤣", "🤣", "📵", "😰"], + ["🚜", "🙃", "😾"], + ["😹", "🚠", "🚌"], + ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"], + ], +) +def test_sort(dtype, strings): + """Test that sorting matches python's internal sorting.""" + + def test_sort(strings, arr_sorted): + arr = np.array(strings, dtype=dtype) + na_object = getattr(arr.dtype, "na_object", "") + if na_object is None and None in strings: + with pytest.raises( + ValueError, + match="Cannot compare null that is not a nan-like value", + ): + np.argsort(arr) + argsorted = None + elif na_object is pd_NA or na_object != '': + argsorted = None + else: + argsorted = np.argsort(arr) + np.random.default_rng().shuffle(arr) + if na_object is None and None in strings: + with pytest.raises( + ValueError, + match="Cannot compare null that is not a nan-like value", + ): + arr.sort() + else: + arr.sort() + assert np.array_equal(arr, arr_sorted, equal_nan=True) + if argsorted is not None: + assert np.array_equal(argsorted, np.argsort(strings)) + + + # make a copy so we don't mutate the lists in the fixture + strings = strings.copy() + arr_sorted = np.array(sorted(strings), dtype=dtype) + test_sort(strings, arr_sorted) + + if not hasattr(dtype, "na_object"): + return + + # make sure NAs get sorted to the end of the array and string NAs get + # sorted like normal strings + strings.insert(0, dtype.na_object) + strings.insert(2, dtype.na_object) + # can't use append because doing that with NA converts + # the result to object dtype + if not isinstance(dtype.na_object, str): + arr_sorted = np.array( + arr_sorted.tolist() + [dtype.na_object, dtype.na_object], + dtype=dtype, + ) + else: + arr_sorted = np.array(sorted(strings), dtype=dtype) + + test_sort(strings, arr_sorted) + + +@pytest.mark.parametrize( + "strings", + [ + ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"], + ["A¢☃€ 😊", "", " ", " "], + ["", "a", "😸", "ááðfáíóåéë"], + ], +) +def test_nonzero(strings, na_object): + dtype = get_dtype(na_object) + arr = np.array(strings, dtype=dtype) + is_nonzero = np.array( + [i for i, item in enumerate(strings) if len(item) != 0]) + assert_array_equal(arr.nonzero()[0], is_nonzero) + + if na_object is not pd_NA and na_object == 'unset': + return + + strings_with_na = np.array(strings + [na_object], dtype=dtype) + is_nan = np.isnan(np.array([dtype.na_object], dtype=dtype))[0] + + if is_nan: + assert strings_with_na.nonzero()[0][-1] == 4 + else: + assert strings_with_na.nonzero()[0][-1] == 3 + + # check that the casting to bool and nonzero give consistent results + assert_array_equal(strings_with_na[strings_with_na.nonzero()], + strings_with_na[strings_with_na.astype(bool)]) + + +def test_where(string_list, na_object): + dtype = get_dtype(na_object) + a = np.array(string_list, dtype=dtype) + b = a[::-1] + res = np.where([True, False, True, False, True, False], a, b) + assert_array_equal(res, [a[0], b[1], a[2], b[3], a[4], b[5]]) + + +def test_fancy_indexing(string_list): + sarr = np.array(string_list, dtype="T") + assert_array_equal(sarr, sarr[np.arange(sarr.shape[0])]) + + inds = [ + [True, True], + [0, 1], + ..., + np.array([0, 1], dtype='uint8'), + ] + + lops = [ + ['a'*25, 'b'*25], + ['', ''], + ['hello', 'world'], + ['hello', 'world'*25], + ] + + # see gh-27003 and gh-27053 + for ind in inds: + for lop in lops: + a = np.array(lop, dtype="T") + assert_array_equal(a[ind], a) + rop = ['d'*25, 'e'*25] + for b in [rop, np.array(rop, dtype="T")]: + a[ind] = b + assert_array_equal(a, b) + assert a[0] == 'd'*25 + + +def test_creation_functions(): + assert_array_equal(np.zeros(3, dtype="T"), ["", "", ""]) + assert_array_equal(np.empty(3, dtype="T"), ["", "", ""]) + + assert np.zeros(3, dtype="T")[0] == "" + assert np.empty(3, dtype="T")[0] == "" + + +def test_concatenate(string_list): + sarr = np.array(string_list, dtype="T") + sarr_cat = np.array(string_list + string_list, dtype="T") + + assert_array_equal(np.concatenate([sarr], axis=0), sarr) + + +def test_resize_method(string_list): + sarr = np.array(string_list, dtype="T") + if IS_PYPY: + sarr.resize(len(string_list)+3, refcheck=False) + else: + sarr.resize(len(string_list)+3) + assert_array_equal(sarr, np.array(string_list + ['']*3, dtype="T")) + + +def test_create_with_copy_none(string_list): + arr = np.array(string_list, dtype=StringDType()) + # create another stringdtype array with an arena that has a different + # in-memory layout than the first array + arr_rev = np.array(string_list[::-1], dtype=StringDType()) + + # this should create a copy and the resulting array + # shouldn't share an allocator or arena with arr_rev, despite + # explicitly passing arr_rev.dtype + arr_copy = np.array(arr, copy=None, dtype=arr_rev.dtype) + np.testing.assert_array_equal(arr, arr_copy) + assert arr_copy.base is None + + with pytest.raises(ValueError, match="Unable to avoid copy"): + np.array(arr, copy=False, dtype=arr_rev.dtype) + + # because we're using arr's dtype instance, the view is safe + arr_view = np.array(arr, copy=None, dtype=arr.dtype) + np.testing.assert_array_equal(arr, arr) + np.testing.assert_array_equal(arr_view[::-1], arr_rev) + assert arr_view is arr + + +def test_astype_copy_false(): + orig_dt = StringDType() + arr = np.array(["hello", "world"], dtype=StringDType()) + assert not arr.astype(StringDType(coerce=False), copy=False).dtype.coerce + + assert arr.astype(orig_dt, copy=False).dtype is orig_dt + +@pytest.mark.parametrize( + "strings", + [ + ["left", "right", "leftovers", "righty", "up", "down"], + ["🤣🤣", "🤣", "📵", "😰"], + ["🚜", "🙃", "😾"], + ["😹", "🚠", "🚌"], + ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"], + ], +) +def test_argmax(strings): + """Test that argmax/argmin matches what python calculates.""" + arr = np.array(strings, dtype="T") + assert np.argmax(arr) == strings.index(max(strings)) + assert np.argmin(arr) == strings.index(min(strings)) + + +@pytest.mark.parametrize( + "arrfunc,expected", + [ + [np.sort, None], + [np.nonzero, (np.array([], dtype=np.int_),)], + [np.argmax, 0], + [np.argmin, 0], + ], +) +def test_arrfuncs_zeros(arrfunc, expected): + arr = np.zeros(10, dtype="T") + result = arrfunc(arr) + if expected is None: + expected = arr + assert_array_equal(result, expected, strict=True) + + +@pytest.mark.parametrize( + ("strings", "cast_answer", "any_answer", "all_answer"), + [ + [["hello", "world"], [True, True], True, True], + [["", ""], [False, False], False, False], + [["hello", ""], [True, False], True, False], + [["", "world"], [False, True], True, False], + ], +) +def test_cast_to_bool(strings, cast_answer, any_answer, all_answer): + sarr = np.array(strings, dtype="T") + assert_array_equal(sarr.astype("bool"), cast_answer) + + assert np.any(sarr) == any_answer + assert np.all(sarr) == all_answer + + +@pytest.mark.parametrize( + ("strings", "cast_answer"), + [ + [[True, True], ["True", "True"]], + [[False, False], ["False", "False"]], + [[True, False], ["True", "False"]], + [[False, True], ["False", "True"]], + ], +) +def test_cast_from_bool(strings, cast_answer): + barr = np.array(strings, dtype=bool) + assert_array_equal(barr.astype("T"), np.array(cast_answer, dtype="T")) + + +@pytest.mark.parametrize("bitsize", [8, 16, 32, 64]) +@pytest.mark.parametrize("signed", [True, False]) +def test_sized_integer_casts(bitsize, signed): + idtype = f"int{bitsize}" + if signed: + inp = [-(2**p - 1) for p in reversed(range(bitsize - 1))] + inp += [2**p - 1 for p in range(1, bitsize - 1)] + else: + idtype = "u" + idtype + inp = [2**p - 1 for p in range(bitsize)] + ainp = np.array(inp, dtype=idtype) + assert_array_equal(ainp, ainp.astype("T").astype(idtype)) + + # safe casting works + ainp.astype("T", casting="safe") + + with pytest.raises(TypeError): + ainp.astype("T").astype(idtype, casting="safe") + + oob = [str(2**bitsize), str(-(2**bitsize))] + with pytest.raises(OverflowError): + np.array(oob, dtype="T").astype(idtype) + + with pytest.raises(ValueError): + np.array(["1", np.nan, "3"], + dtype=StringDType(na_object=np.nan)).astype(idtype) + + +@pytest.mark.parametrize("typename", ["byte", "short", "int", "longlong"]) +@pytest.mark.parametrize("signed", ["", "u"]) +def test_unsized_integer_casts(typename, signed): + idtype = f"{signed}{typename}" + + inp = [1, 2, 3, 4] + ainp = np.array(inp, dtype=idtype) + assert_array_equal(ainp, ainp.astype("T").astype(idtype)) + + +@pytest.mark.parametrize( + "typename", + [ + pytest.param( + "longdouble", + marks=pytest.mark.xfail( + np.dtypes.LongDoubleDType() != np.dtypes.Float64DType(), + reason="numpy lacks an ld2a implementation", + strict=True, + ), + ), + "float64", + "float32", + "float16", + ], +) +def test_float_casts(typename): + inp = [1.1, 2.8, -3.2, 2.7e4] + ainp = np.array(inp, dtype=typename) + assert_array_equal(ainp, ainp.astype("T").astype(typename)) + + inp = [0.1] + sres = np.array(inp, dtype=typename).astype("T") + res = sres.astype(typename) + assert_array_equal(np.array(inp, dtype=typename), res) + assert sres[0] == "0.1" + + if typename == "longdouble": + # let's not worry about platform-dependent rounding of longdouble + return + + fi = np.finfo(typename) + + inp = [1e-324, fi.smallest_subnormal, -1e-324, -fi.smallest_subnormal] + eres = [0, fi.smallest_subnormal, -0, -fi.smallest_subnormal] + res = np.array(inp, dtype=typename).astype("T").astype(typename) + assert_array_equal(eres, res) + + inp = [2e308, fi.max, -2e308, fi.min] + eres = [np.inf, fi.max, -np.inf, fi.min] + res = np.array(inp, dtype=typename).astype("T").astype(typename) + assert_array_equal(eres, res) + + +@pytest.mark.parametrize( + "typename", + [ + "csingle", + "cdouble", + pytest.param( + "clongdouble", + marks=pytest.mark.xfail( + np.dtypes.CLongDoubleDType() != np.dtypes.Complex128DType(), + reason="numpy lacks an ld2a implementation", + strict=True, + ), + ), + ], +) +def test_cfloat_casts(typename): + inp = [1.1 + 1.1j, 2.8 + 2.8j, -3.2 - 3.2j, 2.7e4 + 2.7e4j] + ainp = np.array(inp, dtype=typename) + assert_array_equal(ainp, ainp.astype("T").astype(typename)) + + inp = [0.1 + 0.1j] + sres = np.array(inp, dtype=typename).astype("T") + res = sres.astype(typename) + assert_array_equal(np.array(inp, dtype=typename), res) + assert sres[0] == "(0.1+0.1j)" + + +def test_take(string_list): + sarr = np.array(string_list, dtype="T") + res = sarr.take(np.arange(len(string_list))) + assert_array_equal(sarr, res) + + # make sure it also works for out + out = np.empty(len(string_list), dtype="T") + out[0] = "hello" + res = sarr.take(np.arange(len(string_list)), out=out) + assert res is out + assert_array_equal(sarr, res) + + +@pytest.mark.parametrize("use_out", [True, False]) +@pytest.mark.parametrize( + "ufunc_name,func", + [ + ("min", min), + ("max", max), + ], +) +def test_ufuncs_minmax(string_list, ufunc_name, func, use_out): + """Test that the min/max ufuncs match Python builtin min/max behavior.""" + arr = np.array(string_list, dtype="T") + uarr = np.array(string_list, dtype=str) + res = np.array(func(string_list), dtype="T") + assert_array_equal(getattr(arr, ufunc_name)(), res) + + ufunc = getattr(np, ufunc_name + "imum") + + if use_out: + res = ufunc(arr, arr, out=arr) + else: + res = ufunc(arr, arr) + + assert_array_equal(uarr, res) + assert_array_equal(getattr(arr, ufunc_name)(), func(string_list)) + + +def test_max_regression(): + arr = np.array(['y', 'y', 'z'], dtype="T") + assert arr.max() == 'z' + + +@pytest.mark.parametrize("use_out", [True, False]) +@pytest.mark.parametrize( + "other_strings", + [ + ["abc", "def" * 500, "ghi" * 16, "🤣" * 100, "📵", "😰"], + ["🚜", "🙃", "😾", "😹", "🚠", "🚌"], + ["🥦", "¨", "⨯", "∰ ", "⨌ ", "⎶ "], + ], +) +def test_ufunc_add(dtype, string_list, other_strings, use_out): + arr1 = np.array(string_list, dtype=dtype) + arr2 = np.array(other_strings, dtype=dtype) + result = np.array([a + b for a, b in zip(arr1, arr2)], dtype=dtype) + + if use_out: + res = np.add(arr1, arr2, out=arr1) + else: + res = np.add(arr1, arr2) + + assert_array_equal(res, result) + + if not hasattr(dtype, "na_object"): + return + + is_nan = isinstance(dtype.na_object, float) and np.isnan(dtype.na_object) + is_str = isinstance(dtype.na_object, str) + bool_errors = 0 + try: + bool(dtype.na_object) + except TypeError: + bool_errors = 1 + + arr1 = np.array([dtype.na_object] + string_list, dtype=dtype) + arr2 = np.array(other_strings + [dtype.na_object], dtype=dtype) + + if is_nan or bool_errors or is_str: + res = np.add(arr1, arr2) + assert_array_equal(res[1:-1], arr1[1:-1] + arr2[1:-1]) + if not is_str: + assert res[0] is dtype.na_object and res[-1] is dtype.na_object + else: + assert res[0] == dtype.na_object + arr2[0] + assert res[-1] == arr1[-1] + dtype.na_object + else: + with pytest.raises(ValueError): + np.add(arr1, arr2) + + +def test_ufunc_add_reduce(dtype): + values = ["a", "this is a long string", "c"] + arr = np.array(values, dtype=dtype) + out = np.empty((), dtype=dtype) + + expected = np.array("".join(values), dtype=dtype) + assert_array_equal(np.add.reduce(arr), expected) + + np.add.reduce(arr, out=out) + assert_array_equal(out, expected) + + +def test_add_promoter(string_list): + arr = np.array(string_list, dtype=StringDType()) + lresult = np.array(["hello" + s for s in string_list], dtype=StringDType()) + rresult = np.array([s + "hello" for s in string_list], dtype=StringDType()) + + for op in ["hello", np.str_("hello"), np.array(["hello"])]: + assert_array_equal(op + arr, lresult) + assert_array_equal(arr + op, rresult) + + # The promoter should be able to handle things if users pass `dtype=` + res = np.add("hello", string_list, dtype=StringDType) + assert res.dtype == StringDType() + + # The promoter should not kick in if users override the input, + # which means arr is cast, this fails because of the unknown length. + with pytest.raises(TypeError, match="cannot cast dtype"): + np.add(arr, "add", signature=("U", "U", None), casting="unsafe") + + # But it must simply reject the following: + with pytest.raises(TypeError, match=".*did not contain a loop"): + np.add(arr, "add", signature=(None, "U", None)) + + with pytest.raises(TypeError, match=".*did not contain a loop"): + np.add("a", "b", signature=("U", "U", StringDType)) + + +def test_add_no_legacy_promote_with_signature(): + # Possibly misplaced, but useful to test with string DType. We check that + # if there is clearly no loop found, a stray `dtype=` doesn't break things + # Regression test for the bad error in gh-26735 + # (If legacy promotion is gone, this can be deleted...) + with pytest.raises(TypeError, match=".*did not contain a loop"): + np.add("3", 6, dtype=StringDType) + + +def test_add_promoter_reduce(): + # Exact TypeError could change, but ensure StringDtype doesn't match + with pytest.raises(TypeError, match="the resolved dtypes are not"): + np.add.reduce(np.array(["a", "b"], dtype="U")) + + # On the other hand, using `dtype=T` in the *ufunc* should work. + np.add.reduce(np.array(["a", "b"], dtype="U"), dtype=np.dtypes.StringDType) + + +def test_multiply_reduce(): + # At the time of writing (NumPy 2.0) this is very limited (and rather + # ridiculous anyway). But it works and actually makes some sense... + # (NumPy does not allow non-scalar initial values) + repeats = np.array([2, 3, 4]) + val = "school-🚌" + res = np.multiply.reduce(repeats, initial=val, dtype=np.dtypes.StringDType) + assert res == val * np.prod(repeats) + + +def test_multiply_two_string_raises(): + arr = np.array(["hello", "world"], dtype="T") + with pytest.raises(np._core._exceptions._UFuncNoLoopError): + np.multiply(arr, arr) + + +@pytest.mark.parametrize("use_out", [True, False]) +@pytest.mark.parametrize("other", [2, [2, 1, 3, 4, 1, 3]]) +@pytest.mark.parametrize( + "other_dtype", + [ + None, + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "short", + "int", + "intp", + "long", + "longlong", + "ushort", + "uint", + "uintp", + "ulong", + "ulonglong", + ], +) +def test_ufunc_multiply(dtype, string_list, other, other_dtype, use_out): + """Test the two-argument ufuncs match python builtin behavior.""" + arr = np.array(string_list, dtype=dtype) + if other_dtype is not None: + other_dtype = np.dtype(other_dtype) + try: + len(other) + result = [s * o for s, o in zip(string_list, other)] + other = np.array(other) + if other_dtype is not None: + other = other.astype(other_dtype) + except TypeError: + if other_dtype is not None: + other = other_dtype.type(other) + result = [s * other for s in string_list] + + if use_out: + arr_cache = arr.copy() + lres = np.multiply(arr, other, out=arr) + assert_array_equal(lres, result) + arr[:] = arr_cache + assert lres is arr + arr *= other + assert_array_equal(arr, result) + arr[:] = arr_cache + rres = np.multiply(other, arr, out=arr) + assert rres is arr + assert_array_equal(rres, result) + else: + lres = arr * other + assert_array_equal(lres, result) + rres = other * arr + assert_array_equal(rres, result) + + if not hasattr(dtype, "na_object"): + return + + is_nan = np.isnan(np.array([dtype.na_object], dtype=dtype))[0] + is_str = isinstance(dtype.na_object, str) + bool_errors = 0 + try: + bool(dtype.na_object) + except TypeError: + bool_errors = 1 + + arr = np.array(string_list + [dtype.na_object], dtype=dtype) + + try: + len(other) + other = np.append(other, 3) + if other_dtype is not None: + other = other.astype(other_dtype) + except TypeError: + pass + + if is_nan or bool_errors or is_str: + for res in [arr * other, other * arr]: + assert_array_equal(res[:-1], result) + if not is_str: + assert res[-1] is dtype.na_object + else: + try: + assert res[-1] == dtype.na_object * other[-1] + except (IndexError, TypeError): + assert res[-1] == dtype.na_object * other + else: + with pytest.raises(TypeError): + arr * other + with pytest.raises(TypeError): + other * arr + + +def test_findlike_promoters(): + r = "Wally" + l = "Where's Wally?" + s = np.int32(3) + e = np.int8(13) + for dtypes in [("T", "U"), ("U", "T")]: + for function, answer in [ + (np.strings.index, 8), + (np.strings.endswith, True), + ]: + assert answer == function( + np.array(l, dtype=dtypes[0]), np.array(r, dtype=dtypes[1]), s, e + ) + + +def test_strip_promoter(): + arg = ["Hello!!!!", "Hello??!!"] + strip_char = "!" + answer = ["Hello", "Hello??"] + for dtypes in [("T", "U"), ("U", "T")]: + result = np.strings.strip( + np.array(arg, dtype=dtypes[0]), + np.array(strip_char, dtype=dtypes[1]) + ) + assert_array_equal(result, answer) + assert result.dtype.char == "T" + + +def test_replace_promoter(): + arg = ["Hello, planet!", "planet, Hello!"] + old = "planet" + new = "world" + answer = ["Hello, world!", "world, Hello!"] + for dtypes in itertools.product("TU", repeat=3): + if dtypes == ("U", "U", "U"): + continue + answer_arr = np.strings.replace( + np.array(arg, dtype=dtypes[0]), + np.array(old, dtype=dtypes[1]), + np.array(new, dtype=dtypes[2]), + ) + assert_array_equal(answer_arr, answer) + assert answer_arr.dtype.char == "T" + + +def test_center_promoter(): + arg = ["Hello", "planet!"] + fillchar = "/" + for dtypes in [("T", "U"), ("U", "T")]: + answer = np.strings.center( + np.array(arg, dtype=dtypes[0]), 9, np.array(fillchar, dtype=dtypes[1]) + ) + assert_array_equal(answer, ["//Hello//", "/planet!/"]) + assert answer.dtype.char == "T" + + +DATETIME_INPUT = [ + np.datetime64("1923-04-14T12:43:12"), + np.datetime64("1994-06-21T14:43:15"), + np.datetime64("2001-10-15T04:10:32"), + np.datetime64("NaT"), + np.datetime64("1995-11-25T16:02:16"), + np.datetime64("2005-01-04T03:14:12"), + np.datetime64("2041-12-03T14:05:03"), +] + + +TIMEDELTA_INPUT = [ + np.timedelta64(12358, "s"), + np.timedelta64(23, "s"), + np.timedelta64(74, "s"), + np.timedelta64("NaT"), + np.timedelta64(23, "s"), + np.timedelta64(73, "s"), + np.timedelta64(7, "s"), +] + + +@pytest.mark.parametrize( + "input_data, input_dtype", + [ + (DATETIME_INPUT, "M8[s]"), + (TIMEDELTA_INPUT, "m8[s]") + ] +) +def test_datetime_timedelta_cast(dtype, input_data, input_dtype): + + a = np.array(input_data, dtype=input_dtype) + + has_na = hasattr(dtype, "na_object") + is_str = isinstance(getattr(dtype, "na_object", None), str) + + if not has_na or is_str: + a = np.delete(a, 3) + + sa = a.astype(dtype) + ra = sa.astype(a.dtype) + + if has_na and not is_str: + assert sa[3] is dtype.na_object + assert np.isnat(ra[3]) + + assert_array_equal(a, ra) + + if has_na and not is_str: + # don't worry about comparing how NaT is converted + sa = np.delete(sa, 3) + a = np.delete(a, 3) + + if input_dtype.startswith("M"): + assert_array_equal(sa, a.astype("U")) + else: + # The timedelta to unicode cast produces strings + # that aren't round-trippable and we don't want to + # reproduce that behavior in stringdtype + assert_array_equal(sa, a.astype("int64").astype("U")) + + +def test_nat_casts(): + s = 'nat' + all_nats = itertools.product(*zip(s.upper(), s.lower())) + all_nats = list(map(''.join, all_nats)) + NaT_dt = np.datetime64('NaT') + NaT_td = np.timedelta64('NaT') + for na_object in [np._NoValue, None, np.nan, 'nat', '']: + # numpy treats empty string and all case combinations of 'nat' as NaT + dtype = StringDType(na_object=na_object) + arr = np.array([''] + all_nats, dtype=dtype) + dt_array = arr.astype('M8[s]') + td_array = arr.astype('m8[s]') + assert_array_equal(dt_array, NaT_dt) + assert_array_equal(td_array, NaT_td) + + if na_object is np._NoValue: + output_object = 'NaT' + else: + output_object = na_object + + for arr in [dt_array, td_array]: + assert_array_equal( + arr.astype(dtype), + np.array([output_object]*arr.size, dtype=dtype)) + + +def test_nat_conversion(): + for nat in [np.datetime64("NaT", "s"), np.timedelta64("NaT", "s")]: + with pytest.raises(ValueError, match="string coercion is disabled"): + np.array(["a", nat], dtype=StringDType(coerce=False)) + + +def test_growing_strings(dtype): + # growing a string leads to a heap allocation, this tests to make sure + # we do that bookkeeping correctly for all possible starting cases + data = [ + "hello", # a short string + "abcdefghijklmnopqestuvwxyz", # a medium heap-allocated string + "hello" * 200, # a long heap-allocated string + ] + + arr = np.array(data, dtype=dtype) + uarr = np.array(data, dtype=str) + + for _ in range(5): + arr = arr + arr + uarr = uarr + uarr + + assert_array_equal(arr, uarr) + + +@pytest.mark.skipif(IS_WASM, reason="no threading support in wasm") +def test_threaded_access_and_mutation(dtype, random_string_list): + # this test uses an RNG and may crash or cause deadlocks if there is a + # threading bug + rng = np.random.default_rng(0x4D3D3D3) + + def func(arr): + rnd = rng.random() + # either write to random locations in the array, compute a ufunc, or + # re-initialize the array + if rnd < 0.25: + num = np.random.randint(0, arr.size) + arr[num] = arr[num] + "hello" + elif rnd < 0.5: + if rnd < 0.375: + np.add(arr, arr) + else: + np.add(arr, arr, out=arr) + elif rnd < 0.75: + if rnd < 0.875: + np.multiply(arr, np.int64(2)) + else: + np.multiply(arr, np.int64(2), out=arr) + else: + arr[:] = random_string_list + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as tpe: + arr = np.array(random_string_list, dtype=dtype) + futures = [tpe.submit(func, arr) for _ in range(500)] + + for f in futures: + f.result() + + +UFUNC_TEST_DATA = [ + "hello" * 10, + "Ae¢☃€ 😊" * 20, + "entry\nwith\nnewlines", + "entry\twith\ttabs", +] + + +@pytest.fixture +def string_array(dtype): + return np.array(UFUNC_TEST_DATA, dtype=dtype) + + +@pytest.fixture +def unicode_array(): + return np.array(UFUNC_TEST_DATA, dtype=np.str_) + + +NAN_PRESERVING_FUNCTIONS = [ + "capitalize", + "expandtabs", + "lower", + "lstrip", + "rstrip", + "splitlines", + "strip", + "swapcase", + "title", + "upper", +] + +BOOL_OUTPUT_FUNCTIONS = [ + "isalnum", + "isalpha", + "isdigit", + "islower", + "isspace", + "istitle", + "isupper", + "isnumeric", + "isdecimal", +] + +UNARY_FUNCTIONS = [ + "str_len", + "capitalize", + "expandtabs", + "isalnum", + "isalpha", + "isdigit", + "islower", + "isspace", + "istitle", + "isupper", + "lower", + "lstrip", + "rstrip", + "splitlines", + "strip", + "swapcase", + "title", + "upper", + "isnumeric", + "isdecimal", + "isalnum", + "islower", + "istitle", + "isupper", +] + +UNIMPLEMENTED_VEC_STRING_FUNCTIONS = [ + "capitalize", + "expandtabs", + "lower", + "splitlines", + "swapcase", + "title", + "upper", +] + +ONLY_IN_NP_CHAR = [ + "join", + "split", + "rsplit", + "splitlines" +] + + +@pytest.mark.parametrize("function_name", UNARY_FUNCTIONS) +def test_unary(string_array, unicode_array, function_name): + if function_name in ONLY_IN_NP_CHAR: + func = getattr(np.char, function_name) + else: + func = getattr(np.strings, function_name) + dtype = string_array.dtype + sres = func(string_array) + ures = func(unicode_array) + if sres.dtype == StringDType(): + ures = ures.astype(StringDType()) + assert_array_equal(sres, ures) + + if not hasattr(dtype, "na_object"): + return + + is_nan = np.isnan(np.array([dtype.na_object], dtype=dtype))[0] + is_str = isinstance(dtype.na_object, str) + na_arr = np.insert(string_array, 0, dtype.na_object) + + if function_name in UNIMPLEMENTED_VEC_STRING_FUNCTIONS: + if not is_str: + # to avoid these errors we'd need to add NA support to _vec_string + with pytest.raises((ValueError, TypeError)): + func(na_arr) + else: + if function_name == "splitlines": + assert func(na_arr)[0] == func(dtype.na_object)[()] + else: + assert func(na_arr)[0] == func(dtype.na_object) + return + if function_name == "str_len" and not is_str: + # str_len always errors for any non-string null, even NA ones because + # it has an integer result + with pytest.raises(ValueError): + func(na_arr) + return + if function_name in BOOL_OUTPUT_FUNCTIONS: + if is_nan: + assert func(na_arr)[0] is np.False_ + elif is_str: + assert func(na_arr)[0] == func(dtype.na_object) + else: + with pytest.raises(ValueError): + func(na_arr) + return + if not (is_nan or is_str): + with pytest.raises(ValueError): + func(na_arr) + return + res = func(na_arr) + if is_nan and function_name in NAN_PRESERVING_FUNCTIONS: + assert res[0] is dtype.na_object + elif is_str: + assert res[0] == func(dtype.na_object) + + +unicode_bug_fail = pytest.mark.xfail( + reason="unicode output width is buggy", strict=True +) + +# None means that the argument is a string array +BINARY_FUNCTIONS = [ + ("add", (None, None)), + ("multiply", (None, 2)), + ("mod", ("format: %s", None)), + ("center", (None, 25)), + ("count", (None, "A")), + ("encode", (None, "UTF-8")), + ("endswith", (None, "lo")), + ("find", (None, "A")), + ("index", (None, "e")), + ("join", ("-", None)), + ("ljust", (None, 12)), + ("lstrip", (None, "A")), + ("partition", (None, "A")), + ("replace", (None, "A", "B")), + ("rfind", (None, "A")), + ("rindex", (None, "e")), + ("rjust", (None, 12)), + ("rsplit", (None, "A")), + ("rstrip", (None, "A")), + ("rpartition", (None, "A")), + ("split", (None, "A")), + ("strip", (None, "A")), + ("startswith", (None, "A")), + ("zfill", (None, 12)), +] + +PASSES_THROUGH_NAN_NULLS = [ + "add", + "center", + "ljust", + "multiply", + "replace", + "rjust", + "strip", + "lstrip", + "rstrip", + "replace" + "zfill", +] + +NULLS_ARE_FALSEY = [ + "startswith", + "endswith", +] + +NULLS_ALWAYS_ERROR = [ + "count", + "find", + "rfind", +] + +SUPPORTS_NULLS = ( + PASSES_THROUGH_NAN_NULLS + + NULLS_ARE_FALSEY + + NULLS_ALWAYS_ERROR +) + + +def call_func(func, args, array, sanitize=True): + if args == (None, None): + return func(array, array) + if args[0] is None: + if sanitize: + san_args = tuple( + np.array(arg, dtype=array.dtype) if isinstance(arg, str) else + arg for arg in args[1:] + ) + else: + san_args = args[1:] + return func(array, *san_args) + if args[1] is None: + return func(args[0], array) + # shouldn't ever happen + assert 0 + + +@pytest.mark.parametrize("function_name, args", BINARY_FUNCTIONS) +def test_binary(string_array, unicode_array, function_name, args): + if function_name in ONLY_IN_NP_CHAR: + func = getattr(np.char, function_name) + else: + func = getattr(np.strings, function_name) + sres = call_func(func, args, string_array) + ures = call_func(func, args, unicode_array, sanitize=False) + if not isinstance(sres, tuple) and sres.dtype == StringDType(): + ures = ures.astype(StringDType()) + assert_array_equal(sres, ures) + + dtype = string_array.dtype + if function_name not in SUPPORTS_NULLS or not hasattr(dtype, "na_object"): + return + + na_arr = np.insert(string_array, 0, dtype.na_object) + is_nan = np.isnan(np.array([dtype.na_object], dtype=dtype))[0] + is_str = isinstance(dtype.na_object, str) + should_error = not (is_nan or is_str) + + if ( + (function_name in NULLS_ALWAYS_ERROR and not is_str) + or (function_name in PASSES_THROUGH_NAN_NULLS and should_error) + or (function_name in NULLS_ARE_FALSEY and should_error) + ): + with pytest.raises((ValueError, TypeError)): + call_func(func, args, na_arr) + return + + res = call_func(func, args, na_arr) + + if is_str: + assert res[0] == call_func(func, args, na_arr[:1]) + elif function_name in NULLS_ARE_FALSEY: + assert res[0] is np.False_ + elif function_name in PASSES_THROUGH_NAN_NULLS: + assert res[0] is dtype.na_object + else: + # shouldn't ever get here + assert 0 + + +@pytest.mark.parametrize("function, expected", [ + (np.strings.find, [[2, -1], [1, -1]]), + (np.strings.startswith, [[False, False], [True, False]])]) +@pytest.mark.parametrize("start, stop", [ + (1, 4), + (np.int8(1), np.int8(4)), + (np.array([1, 1], dtype='u2'), np.array([4, 4], dtype='u2'))]) +def test_non_default_start_stop(function, start, stop, expected): + a = np.array([["--🐍--", "--🦜--"], + ["-🐍---", "-🦜---"]], "T") + indx = function(a, "🐍", start, stop) + assert_array_equal(indx, expected) + + +@pytest.mark.parametrize("count", [2, np.int8(2), np.array([2, 2], 'u2')]) +def test_replace_non_default_repeat(count): + a = np.array(["🐍--", "🦜-🦜-"], "T") + result = np.strings.replace(a, "🦜-", "🦜†", count) + assert_array_equal(result, np.array(["🐍--", "🦜†🦜†"], "T")) + + +def test_strip_ljust_rjust_consistency(string_array, unicode_array): + rjs = np.char.rjust(string_array, 1000) + rju = np.char.rjust(unicode_array, 1000) + + ljs = np.char.ljust(string_array, 1000) + lju = np.char.ljust(unicode_array, 1000) + + assert_array_equal( + np.char.lstrip(rjs), + np.char.lstrip(rju).astype(StringDType()), + ) + + assert_array_equal( + np.char.rstrip(ljs), + np.char.rstrip(lju).astype(StringDType()), + ) + + assert_array_equal( + np.char.strip(ljs), + np.char.strip(lju).astype(StringDType()), + ) + + assert_array_equal( + np.char.strip(rjs), + np.char.strip(rju).astype(StringDType()), + ) + + +def test_unset_na_coercion(): + # a dtype instance with an unset na object is compatible + # with a dtype that has one set + + # this test uses the "add" and "equal" ufunc but all ufuncs that + # accept more than one string argument and produce a string should + # behave this way + # TODO: generalize to more ufuncs + inp = ["hello", "world"] + arr = np.array(inp, dtype=StringDType(na_object=None)) + for op_dtype in [None, StringDType(), StringDType(coerce=False), + StringDType(na_object=None)]: + if op_dtype is None: + op = "2" + else: + op = np.array("2", dtype=op_dtype) + res = arr + op + assert_array_equal(res, ["hello2", "world2"]) + + # dtype instances with distinct explicitly set NA objects are incompatible + for op_dtype in [StringDType(na_object=pd_NA), StringDType(na_object="")]: + op = np.array("2", dtype=op_dtype) + with pytest.raises(TypeError): + arr + op + + # comparisons only consider the na_object + for op_dtype in [None, StringDType(), StringDType(coerce=True), + StringDType(na_object=None)]: + if op_dtype is None: + op = inp + else: + op = np.array(inp, dtype=op_dtype) + assert_array_equal(arr, op) + + for op_dtype in [StringDType(na_object=pd_NA), + StringDType(na_object=np.nan)]: + op = np.array(inp, dtype=op_dtype) + with pytest.raises(TypeError): + arr == op + + +def test_repeat(string_array): + res = string_array.repeat(1000) + # Create an empty array with expanded dimension, and fill it. Then, + # reshape it to the expected result. + expected = np.empty_like(string_array, shape=string_array.shape + (1000,)) + expected[...] = string_array[:, np.newaxis] + expected = expected.reshape(-1) + + assert_array_equal(res, expected, strict=True) + + +@pytest.mark.parametrize("tile", [1, 6, (2, 5)]) +def test_accumulation(string_array, tile): + """Accumulation is odd for StringDType but tests dtypes with references. + """ + # Fill with mostly empty strings to not create absurdly big strings + arr = np.zeros_like(string_array, shape=(100,)) + arr[:len(string_array)] = string_array + arr[-len(string_array):] = string_array + + # Bloat size a bit (get above thresholds and test >1 ndim). + arr = np.tile(string_array, tile) + + res = np.add.accumulate(arr, axis=0) + res_obj = np.add.accumulate(arr.astype(object), axis=0) + assert_array_equal(res, res_obj.astype(arr.dtype), strict=True) + + if arr.ndim > 1: + res = np.add.accumulate(arr, axis=-1) + res_obj = np.add.accumulate(arr.astype(object), axis=-1) + + assert_array_equal(res, res_obj.astype(arr.dtype), strict=True) + + +class TestImplementation: + """Check that strings are stored in the arena when possible. + + This tests implementation details, so should be adjusted if + the implementation changes. + """ + + @classmethod + def setup_class(self): + self.MISSING = 0x80 + self.INITIALIZED = 0x40 + self.OUTSIDE_ARENA = 0x20 + self.LONG = 0x10 + self.dtype = StringDType(na_object=np.nan) + self.sizeofstr = self.dtype.itemsize + sp = self.dtype.itemsize // 2 # pointer size = sizeof(size_t) + # Below, size is not strictly correct, since it really uses + # 7 (or 3) bytes, but good enough for the tests here. + self.view_dtype = np.dtype([ + ('offset', f'u{sp}'), + ('size', f'u{sp // 2}'), + ('xsiz', f'V{sp // 2 - 1}'), + ('size_and_flags', 'u1'), + ] if sys.byteorder == 'little' else [ + ('size_and_flags', 'u1'), + ('xsiz', f'V{sp // 2 - 1}'), + ('size', f'u{sp // 2}'), + ('offset', f'u{sp}'), + ]) + self.s_empty = "" + self.s_short = "01234" + self.s_medium = "abcdefghijklmnopqrstuvwxyz" + self.s_long = "-=+" * 100 + self.a = np.array( + [self.s_empty, self.s_short, self.s_medium, self.s_long], + self.dtype) + + def get_view(self, a): + # Cannot view a StringDType as anything else directly, since + # it has references. So, we use a stride trick hack. + from numpy.lib._stride_tricks_impl import DummyArray + interface = dict(a.__array_interface__) + interface['descr'] = self.view_dtype.descr + interface['typestr'] = self.view_dtype.str + return np.asarray(DummyArray(interface, base=a)) + + def get_flags(self, a): + return self.get_view(a)['size_and_flags'] & 0xf0 + + def is_short(self, a): + return self.get_flags(a) == self.INITIALIZED | self.OUTSIDE_ARENA + + def is_on_heap(self, a): + return self.get_flags(a) == (self.INITIALIZED + | self.OUTSIDE_ARENA + | self.LONG) + + def is_missing(self, a): + return self.get_flags(a) & self.MISSING == self.MISSING + + def in_arena(self, a): + return (self.get_flags(a) & (self.INITIALIZED | self.OUTSIDE_ARENA) + == self.INITIALIZED) + + def test_setup(self): + is_short = self.is_short(self.a) + length = np.strings.str_len(self.a) + assert_array_equal(is_short, (length > 0) & (length <= 15)) + assert_array_equal(self.in_arena(self.a), [False, False, True, True]) + assert_array_equal(self.is_on_heap(self.a), False) + assert_array_equal(self.is_missing(self.a), False) + view = self.get_view(self.a) + sizes = np.where(is_short, view['size_and_flags'] & 0xf, + view['size']) + assert_array_equal(sizes, np.strings.str_len(self.a)) + assert_array_equal(view['xsiz'][2:], + np.void(b'\x00' * (self.sizeofstr // 4 - 1))) + # Check that the medium string uses only 1 byte for its length + # in the arena, while the long string takes 8 (or 4). + offsets = view['offset'] + assert offsets[2] == 1 + assert offsets[3] == 1 + len(self.s_medium) + self.sizeofstr // 2 + + def test_empty(self): + e = np.empty((3,), self.dtype) + assert_array_equal(self.get_flags(e), 0) + assert_array_equal(e, "") + + def test_zeros(self): + z = np.zeros((2,), self.dtype) + assert_array_equal(self.get_flags(z), 0) + assert_array_equal(z, "") + + def test_copy(self): + for c in [self.a.copy(), copy.copy(self.a), copy.deepcopy(self.a)]: + assert_array_equal(self.get_flags(c), self.get_flags(self.a)) + assert_array_equal(c, self.a) + offsets = self.get_view(c)['offset'] + assert offsets[2] == 1 + assert offsets[3] == 1 + len(self.s_medium) + self.sizeofstr // 2 + + def test_arena_use_with_setting(self): + c = np.zeros_like(self.a) + assert_array_equal(self.get_flags(c), 0) + c[:] = self.a + assert_array_equal(self.get_flags(c), self.get_flags(self.a)) + assert_array_equal(c, self.a) + + def test_arena_reuse_with_setting(self): + c = self.a.copy() + c[:] = self.a + assert_array_equal(self.get_flags(c), self.get_flags(self.a)) + assert_array_equal(c, self.a) + + def test_arena_reuse_after_missing(self): + c = self.a.copy() + c[:] = np.nan + assert np.all(self.is_missing(c)) + # Replacing with the original strings, the arena should be reused. + c[:] = self.a + assert_array_equal(self.get_flags(c), self.get_flags(self.a)) + assert_array_equal(c, self.a) + + def test_arena_reuse_after_empty(self): + c = self.a.copy() + c[:] = "" + assert_array_equal(c, "") + # Replacing with the original strings, the arena should be reused. + c[:] = self.a + assert_array_equal(self.get_flags(c), self.get_flags(self.a)) + assert_array_equal(c, self.a) + + def test_arena_reuse_for_shorter(self): + c = self.a.copy() + # A string slightly shorter than the shortest in the arena + # should be used for all strings in the arena. + c[:] = self.s_medium[:-1] + assert_array_equal(c, self.s_medium[:-1]) + # first empty string in original was never initialized, so + # filling it in now leaves it initialized inside the arena. + # second string started as a short string so it can never live + # in the arena. + in_arena = np.array([True, False, True, True]) + assert_array_equal(self.in_arena(c), in_arena) + # But when a short string is replaced, it will go on the heap. + assert_array_equal(self.is_short(c), False) + assert_array_equal(self.is_on_heap(c), ~in_arena) + # We can put the originals back, and they'll still fit, + # and short strings are back as short strings + c[:] = self.a + assert_array_equal(c, self.a) + assert_array_equal(self.in_arena(c), in_arena) + assert_array_equal(self.is_short(c), self.is_short(self.a)) + assert_array_equal(self.is_on_heap(c), False) + + def test_arena_reuse_if_possible(self): + c = self.a.copy() + # A slightly longer string will not fit in the arena for + # the medium string, but will fit for the longer one. + c[:] = self.s_medium + "±" + assert_array_equal(c, self.s_medium + "±") + in_arena_exp = np.strings.str_len(self.a) >= len(self.s_medium) + 1 + # first entry started uninitialized and empty, so filling it leaves + # it in the arena + in_arena_exp[0] = True + assert not np.all(in_arena_exp == self.in_arena(self.a)) + assert_array_equal(self.in_arena(c), in_arena_exp) + assert_array_equal(self.is_short(c), False) + assert_array_equal(self.is_on_heap(c), ~in_arena_exp) + # And once outside arena, it stays outside, since offset is lost. + # But short strings are used again. + c[:] = self.a + is_short_exp = self.is_short(self.a) + assert_array_equal(c, self.a) + assert_array_equal(self.in_arena(c), in_arena_exp) + assert_array_equal(self.is_short(c), is_short_exp) + assert_array_equal(self.is_on_heap(c), ~in_arena_exp & ~is_short_exp) + + def test_arena_no_reuse_after_short(self): + c = self.a.copy() + # If we replace a string with a short string, it cannot + # go into the arena after because the offset is lost. + c[:] = self.s_short + assert_array_equal(c, self.s_short) + assert_array_equal(self.in_arena(c), False) + c[:] = self.a + assert_array_equal(c, self.a) + assert_array_equal(self.in_arena(c), False) + assert_array_equal(self.is_on_heap(c), self.in_arena(self.a)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_strings.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_strings.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c67b37d3266ee842b89e6f1b706f4449e07528 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_strings.py @@ -0,0 +1,1289 @@ +import sys +import pytest + +import operator +import numpy as np + +from numpy.testing import assert_array_equal, assert_raises, IS_PYPY +from numpy.testing._private.utils import requires_memory + +COMPARISONS = [ + (operator.eq, np.equal, "=="), + (operator.ne, np.not_equal, "!="), + (operator.lt, np.less, "<"), + (operator.le, np.less_equal, "<="), + (operator.gt, np.greater, ">"), + (operator.ge, np.greater_equal, ">="), +] + +MAX = np.iinfo(np.int64).max + +IS_PYPY_LT_7_3_16 = IS_PYPY and sys.implementation.version < (7, 3, 16) + +@pytest.mark.parametrize(["op", "ufunc", "sym"], COMPARISONS) +def test_mixed_string_comparison_ufuncs_fail(op, ufunc, sym): + arr_string = np.array(["a", "b"], dtype="S") + arr_unicode = np.array(["a", "c"], dtype="U") + + with pytest.raises(TypeError, match="did not contain a loop"): + ufunc(arr_string, arr_unicode) + + with pytest.raises(TypeError, match="did not contain a loop"): + ufunc(arr_unicode, arr_string) + +@pytest.mark.parametrize(["op", "ufunc", "sym"], COMPARISONS) +def test_mixed_string_comparisons_ufuncs_with_cast(op, ufunc, sym): + arr_string = np.array(["a", "b"], dtype="S") + arr_unicode = np.array(["a", "c"], dtype="U") + + # While there is no loop, manual casting is acceptable: + res1 = ufunc(arr_string, arr_unicode, signature="UU->?", casting="unsafe") + res2 = ufunc(arr_string, arr_unicode, signature="SS->?", casting="unsafe") + + expected = op(arr_string.astype("U"), arr_unicode) + assert_array_equal(res1, expected) + assert_array_equal(res2, expected) + + +@pytest.mark.parametrize(["op", "ufunc", "sym"], COMPARISONS) +@pytest.mark.parametrize("dtypes", [ + ("S2", "S2"), ("S2", "S10"), + ("U1"), (">U1", ">U1"), + ("U10")]) +@pytest.mark.parametrize("aligned", [True, False]) +def test_string_comparisons(op, ufunc, sym, dtypes, aligned): + # ensure native byte-order for the first view to stay within unicode range + native_dt = np.dtype(dtypes[0]).newbyteorder("=") + arr = np.arange(2**15).view(native_dt).astype(dtypes[0]) + if not aligned: + # Make `arr` unaligned: + new = np.zeros(arr.nbytes + 1, dtype=np.uint8)[1:].view(dtypes[0]) + new[...] = arr + arr = new + + arr2 = arr.astype(dtypes[1], copy=True) + np.random.shuffle(arr2) + arr[0] = arr2[0] # make sure one matches + + expected = [op(d1, d2) for d1, d2 in zip(arr.tolist(), arr2.tolist())] + assert_array_equal(op(arr, arr2), expected) + assert_array_equal(ufunc(arr, arr2), expected) + assert_array_equal( + np.char.compare_chararrays(arr, arr2, sym, False), expected + ) + + expected = [op(d2, d1) for d1, d2 in zip(arr.tolist(), arr2.tolist())] + assert_array_equal(op(arr2, arr), expected) + assert_array_equal(ufunc(arr2, arr), expected) + assert_array_equal( + np.char.compare_chararrays(arr2, arr, sym, False), expected + ) + + +@pytest.mark.parametrize(["op", "ufunc", "sym"], COMPARISONS) +@pytest.mark.parametrize("dtypes", [ + ("S2", "S2"), ("S2", "S10"), ("U10")]) +def test_string_comparisons_empty(op, ufunc, sym, dtypes): + arr = np.empty((1, 0, 1, 5), dtype=dtypes[0]) + arr2 = np.empty((100, 1, 0, 1), dtype=dtypes[1]) + + expected = np.empty(np.broadcast_shapes(arr.shape, arr2.shape), dtype=bool) + assert_array_equal(op(arr, arr2), expected) + assert_array_equal(ufunc(arr, arr2), expected) + assert_array_equal( + np.char.compare_chararrays(arr, arr2, sym, False), expected + ) + + +@pytest.mark.parametrize("str_dt", ["S", "U"]) +@pytest.mark.parametrize("float_dt", np.typecodes["AllFloat"]) +def test_float_to_string_cast(str_dt, float_dt): + float_dt = np.dtype(float_dt) + fi = np.finfo(float_dt) + arr = np.array([np.nan, np.inf, -np.inf, fi.max, fi.min], dtype=float_dt) + expected = ["nan", "inf", "-inf", str(fi.max), str(fi.min)] + if float_dt.kind == "c": + expected = [f"({r}+0j)" for r in expected] + + res = arr.astype(str_dt) + assert_array_equal(res, np.array(expected, dtype=str_dt)) + + +@pytest.mark.parametrize("str_dt", "US") +@pytest.mark.parametrize("size", [-1, np.iinfo(np.intc).max]) +def test_string_size_dtype_errors(str_dt, size): + if size > 0: + size = size // np.dtype(f"{str_dt}1").itemsize + 1 + + with pytest.raises(ValueError): + np.dtype((str_dt, size)) + with pytest.raises(TypeError): + np.dtype(f"{str_dt}{size}") + + +@pytest.mark.parametrize("str_dt", "US") +def test_string_size_dtype_large_repr(str_dt): + size = np.iinfo(np.intc).max // np.dtype(f"{str_dt}1").itemsize + size_str = str(size) + + dtype = np.dtype((str_dt, size)) + assert size_str in dtype.str + assert size_str in str(dtype) + assert size_str in repr(dtype) + + +@pytest.mark.slow +@requires_memory(2 * np.iinfo(np.intc).max) +@pytest.mark.parametrize("str_dt", "US") +def test_large_string_coercion_error(str_dt): + very_large = np.iinfo(np.intc).max // np.dtype(f"{str_dt}1").itemsize + try: + large_string = "A" * (very_large + 1) + except Exception: + # We may not be able to create this Python string on 32bit. + pytest.skip("python failed to create huge string") + + class MyStr: + def __str__(self): + return large_string + + try: + # TypeError from NumPy, or OverflowError from 32bit Python. + with pytest.raises((TypeError, OverflowError)): + np.array([large_string], dtype=str_dt) + + # Same as above, but input has to be converted to a string. + with pytest.raises((TypeError, OverflowError)): + np.array([MyStr()], dtype=str_dt) + except MemoryError: + # Catch memory errors, because `requires_memory` would do so. + raise AssertionError("Ops should raise before any large allocation.") + +@pytest.mark.slow +@requires_memory(2 * np.iinfo(np.intc).max) +@pytest.mark.parametrize("str_dt", "US") +def test_large_string_addition_error(str_dt): + very_large = np.iinfo(np.intc).max // np.dtype(f"{str_dt}1").itemsize + + a = np.array(["A" * very_large], dtype=str_dt) + b = np.array("B", dtype=str_dt) + try: + with pytest.raises(TypeError): + np.add(a, b) + with pytest.raises(TypeError): + np.add(a, a) + except MemoryError: + # Catch memory errors, because `requires_memory` would do so. + raise AssertionError("Ops should raise before any large allocation.") + + +def test_large_string_cast(): + very_large = np.iinfo(np.intc).max // 4 + # Could be nice to test very large path, but it makes too many huge + # allocations right now (need non-legacy cast loops for this). + # a = np.array([], dtype=np.dtype(("S", very_large))) + # assert a.astype("U").dtype.itemsize == very_large * 4 + + a = np.array([], dtype=np.dtype(("S", very_large + 1))) + # It is not perfect but OK if this raises a MemoryError during setup + # (this happens due clunky code and/or buffer setup.) + with pytest.raises((TypeError, MemoryError)): + a.astype("U") + + +@pytest.mark.parametrize("dt", ["S", "U", "T"]) +class TestMethods: + + @pytest.mark.parametrize("in1,in2,out", [ + ("", "", ""), + ("abc", "abc", "abcabc"), + ("12345", "12345", "1234512345"), + ("MixedCase", "MixedCase", "MixedCaseMixedCase"), + ("12345 \0 ", "12345 \0 ", "12345 \0 12345 \0 "), + ("UPPER", "UPPER", "UPPERUPPER"), + (["abc", "def"], ["hello", "world"], ["abchello", "defworld"]), + ]) + def test_add(self, in1, in2, out, dt): + in1 = np.array(in1, dtype=dt) + in2 = np.array(in2, dtype=dt) + out = np.array(out, dtype=dt) + assert_array_equal(np.strings.add(in1, in2), out) + + @pytest.mark.parametrize("in1,in2,out", [ + ("abc", 3, "abcabcabc"), + ("abc", 0, ""), + ("abc", -1, ""), + (["abc", "def"], [1, 4], ["abc", "defdefdefdef"]), + ]) + def test_multiply(self, in1, in2, out, dt): + in1 = np.array(in1, dtype=dt) + out = np.array(out, dtype=dt) + assert_array_equal(np.strings.multiply(in1, in2), out) + + def test_multiply_raises(self, dt): + with pytest.raises(TypeError, match="unsupported type"): + np.strings.multiply(np.array("abc", dtype=dt), 3.14) + + with pytest.raises(MemoryError): + np.strings.multiply(np.array("abc", dtype=dt), sys.maxsize) + + @pytest.mark.parametrize("i_dt", [np.int8, np.int16, np.int32, + np.int64, np.int_]) + def test_multiply_integer_dtypes(self, i_dt, dt): + a = np.array("abc", dtype=dt) + i = np.array(3, dtype=i_dt) + res = np.array("abcabcabc", dtype=dt) + assert_array_equal(np.strings.multiply(a, i), res) + + @pytest.mark.parametrize("in_,out", [ + ("", False), + ("a", True), + ("A", True), + ("\n", False), + ("abc", True), + ("aBc123", False), + ("abc\n", False), + (["abc", "aBc123"], [True, False]), + ]) + def test_isalpha(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isalpha(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ('', False), + ('a', True), + ('A', True), + ('\n', False), + ('123abc456', True), + ('a1b3c', True), + ('aBc000 ', False), + ('abc\n', False), + ]) + def test_isalnum(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isalnum(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ("", False), + ("a", False), + ("0", True), + ("012345", True), + ("012345a", False), + (["a", "012345"], [False, True]), + ]) + def test_isdigit(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isdigit(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ("", False), + ("a", False), + ("1", False), + (" ", True), + ("\t", True), + ("\r", True), + ("\n", True), + (" \t\r \n", True), + (" \t\r\na", False), + (["\t1", " \t\r \n"], [False, True]) + ]) + def test_isspace(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isspace(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ('', False), + ('a', True), + ('A', False), + ('\n', False), + ('abc', True), + ('aBc', False), + ('abc\n', True), + ]) + def test_islower(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.islower(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ('', False), + ('a', False), + ('A', True), + ('\n', False), + ('ABC', True), + ('AbC', False), + ('ABC\n', True), + ]) + def test_isupper(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isupper(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ('', False), + ('a', False), + ('A', True), + ('\n', False), + ('A Titlecased Line', True), + ('A\nTitlecased Line', True), + ('A Titlecased, Line', True), + ('Not a capitalized String', False), + ('Not\ta Titlecase String', False), + ('Not--a Titlecase String', False), + ('NOT', False), + ]) + def test_istitle(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.istitle(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ("", 0), + ("abc", 3), + ("12345", 5), + ("MixedCase", 9), + ("12345 \x00 ", 8), + ("UPPER", 5), + (["abc", "12345 \x00 "], [3, 8]), + ]) + def test_str_len(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.str_len(in_), out) + + @pytest.mark.parametrize("a,sub,start,end,out", [ + ("abcdefghiabc", "abc", 0, None, 0), + ("abcdefghiabc", "abc", 1, None, 9), + ("abcdefghiabc", "def", 4, None, -1), + ("abc", "", 0, None, 0), + ("abc", "", 3, None, 3), + ("abc", "", 4, None, -1), + ("rrarrrrrrrrra", "a", 0, None, 2), + ("rrarrrrrrrrra", "a", 4, None, 12), + ("rrarrrrrrrrra", "a", 4, 6, -1), + ("", "", 0, None, 0), + ("", "", 1, 1, -1), + ("", "", MAX, 0, -1), + ("", "xx", 0, None, -1), + ("", "xx", 1, 1, -1), + ("", "xx", MAX, 0, -1), + pytest.param(99*"a" + "b", "b", 0, None, 99, + id="99*a+b-b-0-None-99"), + pytest.param(98*"a" + "ba", "ba", 0, None, 98, + id="98*a+ba-ba-0-None-98"), + pytest.param(100*"a", "b", 0, None, -1, + id="100*a-b-0-None--1"), + pytest.param(30000*"a" + 100*"b", 100*"b", 0, None, 30000, + id="30000*a+100*b-100*b-0-None-30000"), + pytest.param(30000*"a", 100*"b", 0, None, -1, + id="30000*a-100*b-0-None--1"), + pytest.param(15000*"a" + 15000*"b", 15000*"b", 0, None, 15000, + id="15000*a+15000*b-15000*b-0-None-15000"), + pytest.param(15000*"a" + 15000*"b", 15000*"c", 0, None, -1, + id="15000*a+15000*b-15000*c-0-None--1"), + (["abcdefghiabc", "rrarrrrrrrrra"], ["def", "arr"], [0, 3], + None, [3, -1]), + ("Ae¢☃€ 😊" * 2, "😊", 0, None, 6), + ("Ae¢☃€ 😊" * 2, "😊", 7, None, 13), + pytest.param("A" * (2 ** 17), r"[\w]+\Z", 0, None, -1, + id=r"A*2**17-[\w]+\Z-0-None--1"), + ]) + def test_find(self, a, sub, start, end, out, dt): + if "😊" in a and dt == "S": + pytest.skip("Bytes dtype does not support non-ascii input") + a = np.array(a, dtype=dt) + sub = np.array(sub, dtype=dt) + assert_array_equal(np.strings.find(a, sub, start, end), out) + + @pytest.mark.parametrize("a,sub,start,end,out", [ + ("abcdefghiabc", "abc", 0, None, 9), + ("abcdefghiabc", "", 0, None, 12), + ("abcdefghiabc", "abcd", 0, None, 0), + ("abcdefghiabc", "abcz", 0, None, -1), + ("abc", "", 0, None, 3), + ("abc", "", 3, None, 3), + ("abc", "", 4, None, -1), + ("rrarrrrrrrrra", "a", 0, None, 12), + ("rrarrrrrrrrra", "a", 4, None, 12), + ("rrarrrrrrrrra", "a", 4, 6, -1), + (["abcdefghiabc", "rrarrrrrrrrra"], ["abc", "a"], [0, 0], + None, [9, 12]), + ("Ae¢☃€ 😊" * 2, "😊", 0, None, 13), + ("Ae¢☃€ 😊" * 2, "😊", 0, 7, 6), + ]) + def test_rfind(self, a, sub, start, end, out, dt): + if "😊" in a and dt == "S": + pytest.skip("Bytes dtype does not support non-ascii input") + a = np.array(a, dtype=dt) + sub = np.array(sub, dtype=dt) + assert_array_equal(np.strings.rfind(a, sub, start, end), out) + + @pytest.mark.parametrize("a,sub,start,end,out", [ + ("aaa", "a", 0, None, 3), + ("aaa", "b", 0, None, 0), + ("aaa", "a", 1, None, 2), + ("aaa", "a", 10, None, 0), + ("aaa", "a", -1, None, 1), + ("aaa", "a", -10, None, 3), + ("aaa", "a", 0, 1, 1), + ("aaa", "a", 0, 10, 3), + ("aaa", "a", 0, -1, 2), + ("aaa", "a", 0, -10, 0), + ("aaa", "", 1, None, 3), + ("aaa", "", 3, None, 1), + ("aaa", "", 10, None, 0), + ("aaa", "", -1, None, 2), + ("aaa", "", -10, None, 4), + ("aaa", "aaaa", 0, None, 0), + pytest.param(98*"a" + "ba", "ba", 0, None, 1, + id="98*a+ba-ba-0-None-1"), + pytest.param(30000*"a" + 100*"b", 100*"b", 0, None, 1, + id="30000*a+100*b-100*b-0-None-1"), + pytest.param(30000*"a", 100*"b", 0, None, 0, + id="30000*a-100*b-0-None-0"), + pytest.param(30000*"a" + 100*"ab", "ab", 0, None, 100, + id="30000*a+100*ab-ab-0-None-100"), + pytest.param(15000*"a" + 15000*"b", 15000*"b", 0, None, 1, + id="15000*a+15000*b-15000*b-0-None-1"), + pytest.param(15000*"a" + 15000*"b", 15000*"c", 0, None, 0, + id="15000*a+15000*b-15000*c-0-None-0"), + ("", "", 0, None, 1), + ("", "", 1, 1, 0), + ("", "", MAX, 0, 0), + ("", "xx", 0, None, 0), + ("", "xx", 1, 1, 0), + ("", "xx", MAX, 0, 0), + (["aaa", ""], ["a", ""], [0, 0], None, [3, 1]), + ("Ae¢☃€ 😊" * 100, "😊", 0, None, 100), + ]) + def test_count(self, a, sub, start, end, out, dt): + if "😊" in a and dt == "S": + pytest.skip("Bytes dtype does not support non-ascii input") + a = np.array(a, dtype=dt) + sub = np.array(sub, dtype=dt) + assert_array_equal(np.strings.count(a, sub, start, end), out) + + @pytest.mark.parametrize("a,prefix,start,end,out", [ + ("hello", "he", 0, None, True), + ("hello", "hello", 0, None, True), + ("hello", "hello world", 0, None, False), + ("hello", "", 0, None, True), + ("hello", "ello", 0, None, False), + ("hello", "ello", 1, None, True), + ("hello", "o", 4, None, True), + ("hello", "o", 5, None, False), + ("hello", "", 5, None, True), + ("hello", "lo", 6, None, False), + ("helloworld", "lowo", 3, None, True), + ("helloworld", "lowo", 3, 7, True), + ("helloworld", "lowo", 3, 6, False), + ("", "", 0, 1, True), + ("", "", 0, 0, True), + ("", "", 1, 0, False), + ("hello", "he", 0, -1, True), + ("hello", "he", -53, -1, True), + ("hello", "hello", 0, -1, False), + ("hello", "hello world", -1, -10, False), + ("hello", "ello", -5, None, False), + ("hello", "ello", -4, None, True), + ("hello", "o", -2, None, False), + ("hello", "o", -1, None, True), + ("hello", "", -3, -3, True), + ("hello", "lo", -9, None, False), + (["hello", ""], ["he", ""], [0, 0], None, [True, True]), + ]) + def test_startswith(self, a, prefix, start, end, out, dt): + a = np.array(a, dtype=dt) + prefix = np.array(prefix, dtype=dt) + assert_array_equal(np.strings.startswith(a, prefix, start, end), out) + + @pytest.mark.parametrize("a,suffix,start,end,out", [ + ("hello", "lo", 0, None, True), + ("hello", "he", 0, None, False), + ("hello", "", 0, None, True), + ("hello", "hello world", 0, None, False), + ("helloworld", "worl", 0, None, False), + ("helloworld", "worl", 3, 9, True), + ("helloworld", "world", 3, 12, True), + ("helloworld", "lowo", 1, 7, True), + ("helloworld", "lowo", 2, 7, True), + ("helloworld", "lowo", 3, 7, True), + ("helloworld", "lowo", 4, 7, False), + ("helloworld", "lowo", 3, 8, False), + ("ab", "ab", 0, 1, False), + ("ab", "ab", 0, 0, False), + ("", "", 0, 1, True), + ("", "", 0, 0, True), + ("", "", 1, 0, False), + ("hello", "lo", -2, None, True), + ("hello", "he", -2, None, False), + ("hello", "", -3, -3, True), + ("hello", "hello world", -10, -2, False), + ("helloworld", "worl", -6, None, False), + ("helloworld", "worl", -5, -1, True), + ("helloworld", "worl", -5, 9, True), + ("helloworld", "world", -7, 12, True), + ("helloworld", "lowo", -99, -3, True), + ("helloworld", "lowo", -8, -3, True), + ("helloworld", "lowo", -7, -3, True), + ("helloworld", "lowo", 3, -4, False), + ("helloworld", "lowo", -8, -2, False), + (["hello", "helloworld"], ["lo", "worl"], [0, -6], None, + [True, False]), + ]) + def test_endswith(self, a, suffix, start, end, out, dt): + a = np.array(a, dtype=dt) + suffix = np.array(suffix, dtype=dt) + assert_array_equal(np.strings.endswith(a, suffix, start, end), out) + + @pytest.mark.parametrize("a,chars,out", [ + ("", None, ""), + (" hello ", None, "hello "), + ("hello", None, "hello"), + (" \t\n\r\f\vabc \t\n\r\f\v", None, "abc \t\n\r\f\v"), + ([" hello ", "hello"], None, ["hello ", "hello"]), + ("", "", ""), + ("", "xyz", ""), + ("hello", "", "hello"), + ("xyzzyhelloxyzzy", "xyz", "helloxyzzy"), + ("hello", "xyz", "hello"), + ("xyxz", "xyxz", ""), + ("xyxzx", "x", "yxzx"), + (["xyzzyhelloxyzzy", "hello"], ["xyz", "xyz"], + ["helloxyzzy", "hello"]), + (["ba", "ac", "baa", "bba"], "b", ["a", "ac", "aa", "a"]), + ]) + def test_lstrip(self, a, chars, out, dt): + a = np.array(a, dtype=dt) + out = np.array(out, dtype=dt) + if chars is not None: + chars = np.array(chars, dtype=dt) + assert_array_equal(np.strings.lstrip(a, chars), out) + else: + assert_array_equal(np.strings.lstrip(a), out) + + @pytest.mark.parametrize("a,chars,out", [ + ("", None, ""), + (" hello ", None, " hello"), + ("hello", None, "hello"), + (" \t\n\r\f\vabc \t\n\r\f\v", None, " \t\n\r\f\vabc"), + ([" hello ", "hello"], None, [" hello", "hello"]), + ("", "", ""), + ("", "xyz", ""), + ("hello", "", "hello"), + (["hello ", "abcdefghijklmnop"], None, + ["hello", "abcdefghijklmnop"]), + ("xyzzyhelloxyzzy", "xyz", "xyzzyhello"), + ("hello", "xyz", "hello"), + ("xyxz", "xyxz", ""), + (" ", None, ""), + ("xyxzx", "x", "xyxz"), + (["xyzzyhelloxyzzy", "hello"], ["xyz", "xyz"], + ["xyzzyhello", "hello"]), + (["ab", "ac", "aab", "abb"], "b", ["a", "ac", "aa", "a"]), + ]) + def test_rstrip(self, a, chars, out, dt): + a = np.array(a, dtype=dt) + out = np.array(out, dtype=dt) + if chars is not None: + chars = np.array(chars, dtype=dt) + assert_array_equal(np.strings.rstrip(a, chars), out) + else: + assert_array_equal(np.strings.rstrip(a), out) + + @pytest.mark.parametrize("a,chars,out", [ + ("", None, ""), + (" hello ", None, "hello"), + ("hello", None, "hello"), + (" \t\n\r\f\vabc \t\n\r\f\v", None, "abc"), + ([" hello ", "hello"], None, ["hello", "hello"]), + ("", "", ""), + ("", "xyz", ""), + ("hello", "", "hello"), + ("xyzzyhelloxyzzy", "xyz", "hello"), + ("hello", "xyz", "hello"), + ("xyxz", "xyxz", ""), + ("xyxzx", "x", "yxz"), + (["xyzzyhelloxyzzy", "hello"], ["xyz", "xyz"], + ["hello", "hello"]), + (["bab", "ac", "baab", "bbabb"], "b", ["a", "ac", "aa", "a"]), + ]) + def test_strip(self, a, chars, out, dt): + a = np.array(a, dtype=dt) + if chars is not None: + chars = np.array(chars, dtype=dt) + out = np.array(out, dtype=dt) + assert_array_equal(np.strings.strip(a, chars), out) + + @pytest.mark.parametrize("buf,old,new,count,res", [ + ("", "", "", -1, ""), + ("", "", "A", -1, "A"), + ("", "A", "", -1, ""), + ("", "A", "A", -1, ""), + ("", "", "", 100, ""), + ("", "", "A", 100, "A"), + ("A", "", "", -1, "A"), + ("A", "", "*", -1, "*A*"), + ("A", "", "*1", -1, "*1A*1"), + ("A", "", "*-#", -1, "*-#A*-#"), + ("AA", "", "*-", -1, "*-A*-A*-"), + ("AA", "", "*-", -1, "*-A*-A*-"), + ("AA", "", "*-", 4, "*-A*-A*-"), + ("AA", "", "*-", 3, "*-A*-A*-"), + ("AA", "", "*-", 2, "*-A*-A"), + ("AA", "", "*-", 1, "*-AA"), + ("AA", "", "*-", 0, "AA"), + ("A", "A", "", -1, ""), + ("AAA", "A", "", -1, ""), + ("AAA", "A", "", -1, ""), + ("AAA", "A", "", 4, ""), + ("AAA", "A", "", 3, ""), + ("AAA", "A", "", 2, "A"), + ("AAA", "A", "", 1, "AA"), + ("AAA", "A", "", 0, "AAA"), + ("AAAAAAAAAA", "A", "", -1, ""), + ("ABACADA", "A", "", -1, "BCD"), + ("ABACADA", "A", "", -1, "BCD"), + ("ABACADA", "A", "", 5, "BCD"), + ("ABACADA", "A", "", 4, "BCD"), + ("ABACADA", "A", "", 3, "BCDA"), + ("ABACADA", "A", "", 2, "BCADA"), + ("ABACADA", "A", "", 1, "BACADA"), + ("ABACADA", "A", "", 0, "ABACADA"), + ("ABCAD", "A", "", -1, "BCD"), + ("ABCADAA", "A", "", -1, "BCD"), + ("BCD", "A", "", -1, "BCD"), + ("*************", "A", "", -1, "*************"), + ("^"+"A"*1000+"^", "A", "", 999, "^A^"), + ("the", "the", "", -1, ""), + ("theater", "the", "", -1, "ater"), + ("thethe", "the", "", -1, ""), + ("thethethethe", "the", "", -1, ""), + ("theatheatheathea", "the", "", -1, "aaaa"), + ("that", "the", "", -1, "that"), + ("thaet", "the", "", -1, "thaet"), + ("here and there", "the", "", -1, "here and re"), + ("here and there and there", "the", "", -1, "here and re and re"), + ("here and there and there", "the", "", 3, "here and re and re"), + ("here and there and there", "the", "", 2, "here and re and re"), + ("here and there and there", "the", "", 1, "here and re and there"), + ("here and there and there", "the", "", 0, "here and there and there"), + ("here and there and there", "the", "", -1, "here and re and re"), + ("abc", "the", "", -1, "abc"), + ("abcdefg", "the", "", -1, "abcdefg"), + ("bbobob", "bob", "", -1, "bob"), + ("bbobobXbbobob", "bob", "", -1, "bobXbob"), + ("aaaaaaabob", "bob", "", -1, "aaaaaaa"), + ("aaaaaaa", "bob", "", -1, "aaaaaaa"), + ("Who goes there?", "o", "o", -1, "Who goes there?"), + ("Who goes there?", "o", "O", -1, "WhO gOes there?"), + ("Who goes there?", "o", "O", -1, "WhO gOes there?"), + ("Who goes there?", "o", "O", 3, "WhO gOes there?"), + ("Who goes there?", "o", "O", 2, "WhO gOes there?"), + ("Who goes there?", "o", "O", 1, "WhO goes there?"), + ("Who goes there?", "o", "O", 0, "Who goes there?"), + ("Who goes there?", "a", "q", -1, "Who goes there?"), + ("Who goes there?", "W", "w", -1, "who goes there?"), + ("WWho goes there?WW", "W", "w", -1, "wwho goes there?ww"), + ("Who goes there?", "?", "!", -1, "Who goes there!"), + ("Who goes there??", "?", "!", -1, "Who goes there!!"), + ("Who goes there?", ".", "!", -1, "Who goes there?"), + ("This is a tissue", "is", "**", -1, "Th** ** a t**sue"), + ("This is a tissue", "is", "**", -1, "Th** ** a t**sue"), + ("This is a tissue", "is", "**", 4, "Th** ** a t**sue"), + ("This is a tissue", "is", "**", 3, "Th** ** a t**sue"), + ("This is a tissue", "is", "**", 2, "Th** ** a tissue"), + ("This is a tissue", "is", "**", 1, "Th** is a tissue"), + ("This is a tissue", "is", "**", 0, "This is a tissue"), + ("bobob", "bob", "cob", -1, "cobob"), + ("bobobXbobobob", "bob", "cob", -1, "cobobXcobocob"), + ("bobob", "bot", "bot", -1, "bobob"), + ("Reykjavik", "k", "KK", -1, "ReyKKjaviKK"), + ("Reykjavik", "k", "KK", -1, "ReyKKjaviKK"), + ("Reykjavik", "k", "KK", 2, "ReyKKjaviKK"), + ("Reykjavik", "k", "KK", 1, "ReyKKjavik"), + ("Reykjavik", "k", "KK", 0, "Reykjavik"), + ("A.B.C.", ".", "----", -1, "A----B----C----"), + ("Reykjavik", "q", "KK", -1, "Reykjavik"), + ("spam, spam, eggs and spam", "spam", "ham", -1, + "ham, ham, eggs and ham"), + ("spam, spam, eggs and spam", "spam", "ham", -1, + "ham, ham, eggs and ham"), + ("spam, spam, eggs and spam", "spam", "ham", 4, + "ham, ham, eggs and ham"), + ("spam, spam, eggs and spam", "spam", "ham", 3, + "ham, ham, eggs and ham"), + ("spam, spam, eggs and spam", "spam", "ham", 2, + "ham, ham, eggs and spam"), + ("spam, spam, eggs and spam", "spam", "ham", 1, + "ham, spam, eggs and spam"), + ("spam, spam, eggs and spam", "spam", "ham", 0, + "spam, spam, eggs and spam"), + ("bobobob", "bobob", "bob", -1, "bobob"), + ("bobobobXbobobob", "bobob", "bob", -1, "bobobXbobob"), + ("BOBOBOB", "bob", "bobby", -1, "BOBOBOB"), + ("one!two!three!", "!", "@", 1, "one@two!three!"), + ("one!two!three!", "!", "", -1, "onetwothree"), + ("one!two!three!", "!", "@", 2, "one@two@three!"), + ("one!two!three!", "!", "@", 3, "one@two@three@"), + ("one!two!three!", "!", "@", 4, "one@two@three@"), + ("one!two!three!", "!", "@", 0, "one!two!three!"), + ("one!two!three!", "!", "@", -1, "one@two@three@"), + ("one!two!three!", "x", "@", -1, "one!two!three!"), + ("one!two!three!", "x", "@", 2, "one!two!three!"), + ("abc", "", "-", -1, "-a-b-c-"), + ("abc", "", "-", 3, "-a-b-c"), + ("abc", "", "-", 0, "abc"), + ("abc", "ab", "--", 0, "abc"), + ("abc", "xy", "--", -1, "abc"), + (["abbc", "abbd"], "b", "z", [1, 2], ["azbc", "azzd"]), + ]) + def test_replace(self, buf, old, new, count, res, dt): + if "😊" in buf and dt == "S": + pytest.skip("Bytes dtype does not support non-ascii input") + buf = np.array(buf, dtype=dt) + old = np.array(old, dtype=dt) + new = np.array(new, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.replace(buf, old, new, count), res) + + @pytest.mark.parametrize("buf,sub,start,end,res", [ + ("abcdefghiabc", "", 0, None, 0), + ("abcdefghiabc", "def", 0, None, 3), + ("abcdefghiabc", "abc", 0, None, 0), + ("abcdefghiabc", "abc", 1, None, 9), + ]) + def test_index(self, buf, sub, start, end, res, dt): + buf = np.array(buf, dtype=dt) + sub = np.array(sub, dtype=dt) + assert_array_equal(np.strings.index(buf, sub, start, end), res) + + @pytest.mark.parametrize("buf,sub,start,end", [ + ("abcdefghiabc", "hib", 0, None), + ("abcdefghiab", "abc", 1, None), + ("abcdefghi", "ghi", 8, None), + ("abcdefghi", "ghi", -1, None), + ("rrarrrrrrrrra", "a", 4, 6), + ]) + def test_index_raises(self, buf, sub, start, end, dt): + buf = np.array(buf, dtype=dt) + sub = np.array(sub, dtype=dt) + with pytest.raises(ValueError, match="substring not found"): + np.strings.index(buf, sub, start, end) + + @pytest.mark.parametrize("buf,sub,start,end,res", [ + ("abcdefghiabc", "", 0, None, 12), + ("abcdefghiabc", "def", 0, None, 3), + ("abcdefghiabc", "abc", 0, None, 9), + ("abcdefghiabc", "abc", 0, -1, 0), + ]) + def test_rindex(self, buf, sub, start, end, res, dt): + buf = np.array(buf, dtype=dt) + sub = np.array(sub, dtype=dt) + assert_array_equal(np.strings.rindex(buf, sub, start, end), res) + + @pytest.mark.parametrize("buf,sub,start,end", [ + ("abcdefghiabc", "hib", 0, None), + ("defghiabc", "def", 1, None), + ("defghiabc", "abc", 0, -1), + ("abcdefghi", "ghi", 0, 8), + ("abcdefghi", "ghi", 0, -1), + ("rrarrrrrrrrra", "a", 4, 6), + ]) + def test_rindex_raises(self, buf, sub, start, end, dt): + buf = np.array(buf, dtype=dt) + sub = np.array(sub, dtype=dt) + with pytest.raises(ValueError, match="substring not found"): + np.strings.rindex(buf, sub, start, end) + + @pytest.mark.parametrize("buf,tabsize,res", [ + ("abc\rab\tdef\ng\thi", 8, "abc\rab def\ng hi"), + ("abc\rab\tdef\ng\thi", 4, "abc\rab def\ng hi"), + ("abc\r\nab\tdef\ng\thi", 8, "abc\r\nab def\ng hi"), + ("abc\r\nab\tdef\ng\thi", 4, "abc\r\nab def\ng hi"), + ("abc\r\nab\r\ndef\ng\r\nhi", 4, "abc\r\nab\r\ndef\ng\r\nhi"), + (" \ta\n\tb", 1, " a\n b"), + ]) + def test_expandtabs(self, buf, tabsize, res, dt): + buf = np.array(buf, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.expandtabs(buf, tabsize), res) + + def test_expandtabs_raises_overflow(self, dt): + with pytest.raises(OverflowError, match="new string is too long"): + np.strings.expandtabs(np.array("\ta\n\tb", dtype=dt), sys.maxsize) + np.strings.expandtabs(np.array("\ta\n\tb", dtype=dt), 2**61) + + FILL_ERROR = "The fill character must be exactly one character long" + + def test_center_raises_multiple_character_fill(self, dt): + buf = np.array("abc", dtype=dt) + fill = np.array("**", dtype=dt) + with pytest.raises(TypeError, match=self.FILL_ERROR): + np.strings.center(buf, 10, fill) + + def test_ljust_raises_multiple_character_fill(self, dt): + buf = np.array("abc", dtype=dt) + fill = np.array("**", dtype=dt) + with pytest.raises(TypeError, match=self.FILL_ERROR): + np.strings.ljust(buf, 10, fill) + + def test_rjust_raises_multiple_character_fill(self, dt): + buf = np.array("abc", dtype=dt) + fill = np.array("**", dtype=dt) + with pytest.raises(TypeError, match=self.FILL_ERROR): + np.strings.rjust(buf, 10, fill) + + @pytest.mark.parametrize("buf,width,fillchar,res", [ + ('abc', 10, ' ', ' abc '), + ('abc', 6, ' ', ' abc '), + ('abc', 3, ' ', 'abc'), + ('abc', 2, ' ', 'abc'), + ('abc', 10, '*', '***abc****'), + ]) + def test_center(self, buf, width, fillchar, res, dt): + buf = np.array(buf, dtype=dt) + fillchar = np.array(fillchar, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.center(buf, width, fillchar), res) + + @pytest.mark.parametrize("buf,width,fillchar,res", [ + ('abc', 10, ' ', 'abc '), + ('abc', 6, ' ', 'abc '), + ('abc', 3, ' ', 'abc'), + ('abc', 2, ' ', 'abc'), + ('abc', 10, '*', 'abc*******'), + ]) + def test_ljust(self, buf, width, fillchar, res, dt): + buf = np.array(buf, dtype=dt) + fillchar = np.array(fillchar, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.ljust(buf, width, fillchar), res) + + @pytest.mark.parametrize("buf,width,fillchar,res", [ + ('abc', 10, ' ', ' abc'), + ('abc', 6, ' ', ' abc'), + ('abc', 3, ' ', 'abc'), + ('abc', 2, ' ', 'abc'), + ('abc', 10, '*', '*******abc'), + ]) + def test_rjust(self, buf, width, fillchar, res, dt): + buf = np.array(buf, dtype=dt) + fillchar = np.array(fillchar, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.rjust(buf, width, fillchar), res) + + @pytest.mark.parametrize("buf,width,res", [ + ('123', 2, '123'), + ('123', 3, '123'), + ('0123', 4, '0123'), + ('+123', 3, '+123'), + ('+123', 4, '+123'), + ('+123', 5, '+0123'), + ('+0123', 5, '+0123'), + ('-123', 3, '-123'), + ('-123', 4, '-123'), + ('-0123', 5, '-0123'), + ('000', 3, '000'), + ('34', 1, '34'), + ('0034', 4, '0034'), + ]) + def test_zfill(self, buf, width, res, dt): + buf = np.array(buf, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.zfill(buf, width), res) + + @pytest.mark.parametrize("buf,sep,res1,res2,res3", [ + ("this is the partition method", "ti", "this is the par", + "ti", "tion method"), + ("http://www.python.org", "://", "http", "://", "www.python.org"), + ("http://www.python.org", "?", "http://www.python.org", "", ""), + ("http://www.python.org", "http://", "", "http://", "www.python.org"), + ("http://www.python.org", "org", "http://www.python.", "org", ""), + ("http://www.python.org", ["://", "?", "http://", "org"], + ["http", "http://www.python.org", "", "http://www.python."], + ["://", "", "http://", "org"], + ["www.python.org", "", "www.python.org", ""]), + ("mississippi", "ss", "mi", "ss", "issippi"), + ("mississippi", "i", "m", "i", "ssissippi"), + ("mississippi", "w", "mississippi", "", ""), + ]) + def test_partition(self, buf, sep, res1, res2, res3, dt): + buf = np.array(buf, dtype=dt) + sep = np.array(sep, dtype=dt) + res1 = np.array(res1, dtype=dt) + res2 = np.array(res2, dtype=dt) + res3 = np.array(res3, dtype=dt) + act1, act2, act3 = np.strings.partition(buf, sep) + assert_array_equal(act1, res1) + assert_array_equal(act2, res2) + assert_array_equal(act3, res3) + assert_array_equal(act1 + act2 + act3, buf) + + @pytest.mark.parametrize("buf,sep,res1,res2,res3", [ + ("this is the partition method", "ti", "this is the parti", + "ti", "on method"), + ("http://www.python.org", "://", "http", "://", "www.python.org"), + ("http://www.python.org", "?", "", "", "http://www.python.org"), + ("http://www.python.org", "http://", "", "http://", "www.python.org"), + ("http://www.python.org", "org", "http://www.python.", "org", ""), + ("http://www.python.org", ["://", "?", "http://", "org"], + ["http", "", "", "http://www.python."], + ["://", "", "http://", "org"], + ["www.python.org", "http://www.python.org", "www.python.org", ""]), + ("mississippi", "ss", "missi", "ss", "ippi"), + ("mississippi", "i", "mississipp", "i", ""), + ("mississippi", "w", "", "", "mississippi"), + ]) + def test_rpartition(self, buf, sep, res1, res2, res3, dt): + buf = np.array(buf, dtype=dt) + sep = np.array(sep, dtype=dt) + res1 = np.array(res1, dtype=dt) + res2 = np.array(res2, dtype=dt) + res3 = np.array(res3, dtype=dt) + act1, act2, act3 = np.strings.rpartition(buf, sep) + assert_array_equal(act1, res1) + assert_array_equal(act2, res2) + assert_array_equal(act3, res3) + assert_array_equal(act1 + act2 + act3, buf) + + +@pytest.mark.parametrize("dt", ["U", "T"]) +class TestMethodsWithUnicode: + @pytest.mark.parametrize("in_,out", [ + ("", False), + ("a", False), + ("0", True), + ("\u2460", False), # CIRCLED DIGIT 1 + ("\xbc", False), # VULGAR FRACTION ONE QUARTER + ("\u0660", True), # ARABIC_INDIC DIGIT ZERO + ("012345", True), + ("012345a", False), + (["0", "a"], [True, False]), + ]) + def test_isdecimal_unicode(self, in_, out, dt): + buf = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isdecimal(buf), out) + + @pytest.mark.parametrize("in_,out", [ + ("", False), + ("a", False), + ("0", True), + ("\u2460", True), # CIRCLED DIGIT 1 + ("\xbc", True), # VULGAR FRACTION ONE QUARTER + ("\u0660", True), # ARABIC_INDIC DIGIT ZERO + ("012345", True), + ("012345a", False), + (["0", "a"], [True, False]), + ]) + def test_isnumeric_unicode(self, in_, out, dt): + buf = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isnumeric(buf), out) + + @pytest.mark.parametrize("buf,old,new,count,res", [ + ("...\u043c......<", "<", "<", -1, "...\u043c......<"), + ("Ae¢☃€ 😊" * 2, "A", "B", -1, "Be¢☃€ 😊Be¢☃€ 😊"), + ("Ae¢☃€ 😊" * 2, "😊", "B", -1, "Ae¢☃€ BAe¢☃€ B"), + ]) + def test_replace_unicode(self, buf, old, new, count, res, dt): + buf = np.array(buf, dtype=dt) + old = np.array(old, dtype=dt) + new = np.array(new, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.replace(buf, old, new, count), res) + + @pytest.mark.parametrize("in_", [ + '\U00010401', + '\U00010427', + '\U00010429', + '\U0001044E', + '\U0001D7F6', + '\U00011066', + '\U000104A0', + pytest.param('\U0001F107', marks=pytest.mark.xfail( + sys.platform == 'win32' and IS_PYPY_LT_7_3_16, + reason="PYPY bug in Py_UNICODE_ISALNUM", + strict=True)), + ]) + def test_isalnum_unicode(self, in_, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isalnum(in_), True) + + @pytest.mark.parametrize("in_,out", [ + ('\u1FFc', False), + ('\u2167', False), + ('\U00010401', False), + ('\U00010427', False), + ('\U0001F40D', False), + ('\U0001F46F', False), + ('\u2177', True), + pytest.param('\U00010429', True, marks=pytest.mark.xfail( + sys.platform == 'win32' and IS_PYPY_LT_7_3_16, + reason="PYPY bug in Py_UNICODE_ISLOWER", + strict=True)), + ('\U0001044E', True), + ]) + def test_islower_unicode(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.islower(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ('\u1FFc', False), + ('\u2167', True), + ('\U00010401', True), + ('\U00010427', True), + ('\U0001F40D', False), + ('\U0001F46F', False), + ('\u2177', False), + pytest.param('\U00010429', False, marks=pytest.mark.xfail( + sys.platform == 'win32' and IS_PYPY_LT_7_3_16, + reason="PYPY bug in Py_UNICODE_ISUPPER", + strict=True)), + ('\U0001044E', False), + ]) + def test_isupper_unicode(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.isupper(in_), out) + + @pytest.mark.parametrize("in_,out", [ + ('\u1FFc', True), + ('Greek \u1FFcitlecases ...', True), + pytest.param('\U00010401\U00010429', True, marks=pytest.mark.xfail( + sys.platform == 'win32' and IS_PYPY_LT_7_3_16, + reason="PYPY bug in Py_UNICODE_ISISTITLE", + strict=True)), + ('\U00010427\U0001044E', True), + pytest.param('\U00010429', False, marks=pytest.mark.xfail( + sys.platform == 'win32' and IS_PYPY_LT_7_3_16, + reason="PYPY bug in Py_UNICODE_ISISTITLE", + strict=True)), + ('\U0001044E', False), + ('\U0001F40D', False), + ('\U0001F46F', False), + ]) + def test_istitle_unicode(self, in_, out, dt): + in_ = np.array(in_, dtype=dt) + assert_array_equal(np.strings.istitle(in_), out) + + @pytest.mark.parametrize("buf,sub,start,end,res", [ + ("Ae¢☃€ 😊" * 2, "😊", 0, None, 6), + ("Ae¢☃€ 😊" * 2, "😊", 7, None, 13), + ]) + def test_index_unicode(self, buf, sub, start, end, res, dt): + buf = np.array(buf, dtype=dt) + sub = np.array(sub, dtype=dt) + assert_array_equal(np.strings.index(buf, sub, start, end), res) + + def test_index_raises_unicode(self, dt): + with pytest.raises(ValueError, match="substring not found"): + np.strings.index("Ae¢☃€ 😊", "😀") + + @pytest.mark.parametrize("buf,res", [ + ("Ae¢☃€ \t 😊", "Ae¢☃€ 😊"), + ("\t\U0001044E", " \U0001044E"), + ]) + def test_expandtabs(self, buf, res, dt): + buf = np.array(buf, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.expandtabs(buf), res) + + @pytest.mark.parametrize("buf,width,fillchar,res", [ + ('x', 2, '\U0001044E', 'x\U0001044E'), + ('x', 3, '\U0001044E', '\U0001044Ex\U0001044E'), + ('x', 4, '\U0001044E', '\U0001044Ex\U0001044E\U0001044E'), + ]) + def test_center(self, buf, width, fillchar, res, dt): + buf = np.array(buf, dtype=dt) + fillchar = np.array(fillchar, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.center(buf, width, fillchar), res) + + @pytest.mark.parametrize("buf,width,fillchar,res", [ + ('x', 2, '\U0001044E', 'x\U0001044E'), + ('x', 3, '\U0001044E', 'x\U0001044E\U0001044E'), + ('x', 4, '\U0001044E', 'x\U0001044E\U0001044E\U0001044E'), + ]) + def test_ljust(self, buf, width, fillchar, res, dt): + buf = np.array(buf, dtype=dt) + fillchar = np.array(fillchar, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.ljust(buf, width, fillchar), res) + + @pytest.mark.parametrize("buf,width,fillchar,res", [ + ('x', 2, '\U0001044E', '\U0001044Ex'), + ('x', 3, '\U0001044E', '\U0001044E\U0001044Ex'), + ('x', 4, '\U0001044E', '\U0001044E\U0001044E\U0001044Ex'), + ]) + def test_rjust(self, buf, width, fillchar, res, dt): + buf = np.array(buf, dtype=dt) + fillchar = np.array(fillchar, dtype=dt) + res = np.array(res, dtype=dt) + assert_array_equal(np.strings.rjust(buf, width, fillchar), res) + + @pytest.mark.parametrize("buf,sep,res1,res2,res3", [ + ("āāāāĀĀĀĀ", "Ă", "āāāāĀĀĀĀ", "", ""), + ("āāāāĂĀĀĀĀ", "Ă", "āāāā", "Ă", "ĀĀĀĀ"), + ("āāāāĂĂĀĀĀĀ", "ĂĂ", "āāāā", "ĂĂ", "ĀĀĀĀ"), + ("𐌁𐌁𐌁𐌁𐌀𐌀𐌀𐌀", "𐌂", "𐌁𐌁𐌁𐌁𐌀𐌀𐌀𐌀", "", ""), + ("𐌁𐌁𐌁𐌁𐌂𐌀𐌀𐌀𐌀", "𐌂", "𐌁𐌁𐌁𐌁", "𐌂", "𐌀𐌀𐌀𐌀"), + ("𐌁𐌁𐌁𐌁𐌂𐌂𐌀𐌀𐌀𐌀", "𐌂𐌂", "𐌁𐌁𐌁𐌁", "𐌂𐌂", "𐌀𐌀𐌀𐌀"), + ("𐌁𐌁𐌁𐌁𐌂𐌂𐌂𐌂𐌀𐌀𐌀𐌀", "𐌂𐌂𐌂𐌂", "𐌁𐌁𐌁𐌁", "𐌂𐌂𐌂𐌂", "𐌀𐌀𐌀𐌀"), + ]) + def test_partition(self, buf, sep, res1, res2, res3, dt): + buf = np.array(buf, dtype=dt) + sep = np.array(sep, dtype=dt) + res1 = np.array(res1, dtype=dt) + res2 = np.array(res2, dtype=dt) + res3 = np.array(res3, dtype=dt) + act1, act2, act3 = np.strings.partition(buf, sep) + assert_array_equal(act1, res1) + assert_array_equal(act2, res2) + assert_array_equal(act3, res3) + assert_array_equal(act1 + act2 + act3, buf) + + @pytest.mark.parametrize("buf,sep,res1,res2,res3", [ + ("āāāāĀĀĀĀ", "Ă", "", "", "āāāāĀĀĀĀ"), + ("āāāāĂĀĀĀĀ", "Ă", "āāāā", "Ă", "ĀĀĀĀ"), + ("āāāāĂĂĀĀĀĀ", "ĂĂ", "āāāā", "ĂĂ", "ĀĀĀĀ"), + ("𐌁𐌁𐌁𐌁𐌀𐌀𐌀𐌀", "𐌂", "", "", "𐌁𐌁𐌁𐌁𐌀𐌀𐌀𐌀"), + ("𐌁𐌁𐌁𐌁𐌂𐌀𐌀𐌀𐌀", "𐌂", "𐌁𐌁𐌁𐌁", "𐌂", "𐌀𐌀𐌀𐌀"), + ("𐌁𐌁𐌁𐌁𐌂𐌂𐌀𐌀𐌀𐌀", "𐌂𐌂", "𐌁𐌁𐌁𐌁", "𐌂𐌂", "𐌀𐌀𐌀𐌀"), + ]) + def test_rpartition(self, buf, sep, res1, res2, res3, dt): + buf = np.array(buf, dtype=dt) + sep = np.array(sep, dtype=dt) + res1 = np.array(res1, dtype=dt) + res2 = np.array(res2, dtype=dt) + res3 = np.array(res3, dtype=dt) + act1, act2, act3 = np.strings.rpartition(buf, sep) + assert_array_equal(act1, res1) + assert_array_equal(act2, res2) + assert_array_equal(act3, res3) + assert_array_equal(act1 + act2 + act3, buf) + + @pytest.mark.parametrize("method", ["strip", "lstrip", "rstrip"]) + @pytest.mark.parametrize( + "source,strip", + [ + ("λμ", "μ"), + ("λμ", "λ"), + ("λ"*5 + "μ"*2, "μ"), + ("λ" * 5 + "μ" * 2, "λ"), + ("λ" * 5 + "A" + "μ" * 2, "μλ"), + ("λμ" * 5, "μ"), + ("λμ" * 5, "λ"), + ]) + def test_strip_functions_unicode(self, source, strip, method, dt): + src_array = np.array([source], dtype=dt) + + npy_func = getattr(np.strings, method) + py_func = getattr(str, method) + + expected = np.array([py_func(source, strip)], dtype=dt) + actual = npy_func(src_array, strip) + + assert_array_equal(actual, expected) + + +class TestMixedTypeMethods: + def test_center(self): + buf = np.array("😊", dtype="U") + fill = np.array("*", dtype="S") + res = np.array("*😊*", dtype="U") + assert_array_equal(np.strings.center(buf, 3, fill), res) + + buf = np.array("s", dtype="S") + fill = np.array("*", dtype="U") + res = np.array("*s*", dtype="S") + assert_array_equal(np.strings.center(buf, 3, fill), res) + + with pytest.raises(ValueError, match="'ascii' codec can't encode"): + buf = np.array("s", dtype="S") + fill = np.array("😊", dtype="U") + np.strings.center(buf, 3, fill) + + def test_ljust(self): + buf = np.array("😊", dtype="U") + fill = np.array("*", dtype="S") + res = np.array("😊**", dtype="U") + assert_array_equal(np.strings.ljust(buf, 3, fill), res) + + buf = np.array("s", dtype="S") + fill = np.array("*", dtype="U") + res = np.array("s**", dtype="S") + assert_array_equal(np.strings.ljust(buf, 3, fill), res) + + with pytest.raises(ValueError, match="'ascii' codec can't encode"): + buf = np.array("s", dtype="S") + fill = np.array("😊", dtype="U") + np.strings.ljust(buf, 3, fill) + + def test_rjust(self): + buf = np.array("😊", dtype="U") + fill = np.array("*", dtype="S") + res = np.array("**😊", dtype="U") + assert_array_equal(np.strings.rjust(buf, 3, fill), res) + + buf = np.array("s", dtype="S") + fill = np.array("*", dtype="U") + res = np.array("**s", dtype="S") + assert_array_equal(np.strings.rjust(buf, 3, fill), res) + + with pytest.raises(ValueError, match="'ascii' codec can't encode"): + buf = np.array("s", dtype="S") + fill = np.array("😊", dtype="U") + np.strings.rjust(buf, 3, fill) + + +class TestUnicodeOnlyMethodsRaiseWithBytes: + def test_isdecimal_raises(self): + in_ = np.array(b"1") + with assert_raises(TypeError): + np.strings.isdecimal(in_) + + def test_isnumeric_bytes(self): + in_ = np.array(b"1") + with assert_raises(TypeError): + np.strings.isnumeric(in_) + + +def check_itemsize(n_elem, dt): + if dt == "T": + return np.dtype(dt).itemsize + if dt == "S": + return n_elem + if dt == "U": + return n_elem * 4 + +@pytest.mark.parametrize("dt", ["S", "U", "T"]) +class TestReplaceOnArrays: + + def test_replace_count_and_size(self, dt): + a = np.array(["0123456789" * i for i in range(4)], dtype=dt) + r1 = np.strings.replace(a, "5", "ABCDE") + assert r1.dtype.itemsize == check_itemsize(3*10 + 3*4, dt) + r1_res = np.array(["01234ABCDE6789" * i for i in range(4)], dtype=dt) + assert_array_equal(r1, r1_res) + r2 = np.strings.replace(a, "5", "ABCDE", 1) + assert r2.dtype.itemsize == check_itemsize(3*10 + 4, dt) + r3 = np.strings.replace(a, "5", "ABCDE", 0) + assert r3.dtype.itemsize == a.dtype.itemsize + assert_array_equal(r3, a) + # Negative values mean to replace all. + r4 = np.strings.replace(a, "5", "ABCDE", -1) + assert r4.dtype.itemsize == check_itemsize(3*10 + 3*4, dt) + assert_array_equal(r4, r1) + # We can do count on an element-by-element basis. + r5 = np.strings.replace(a, "5", "ABCDE", [-1, -1, -1, 1]) + assert r5.dtype.itemsize == check_itemsize(3*10 + 4, dt) + assert_array_equal(r5, np.array( + ["01234ABCDE6789" * i for i in range(3)] + + ["01234ABCDE6789" + "0123456789" * 2], dtype=dt)) + + def test_replace_broadcasting(self, dt): + a = np.array("0,0,0", dtype=dt) + r1 = np.strings.replace(a, "0", "1", np.arange(3)) + assert r1.dtype == a.dtype + assert_array_equal(r1, np.array(["0,0,0", "1,0,0", "1,1,0"], dtype=dt)) + r2 = np.strings.replace(a, "0", [["1"], ["2"]], np.arange(1, 4)) + assert_array_equal(r2, np.array([["1,0,0", "1,1,0", "1,1,1"], + ["2,0,0", "2,2,0", "2,2,2"]], + dtype=dt)) + r3 = np.strings.replace(a, ["0", "0,0", "0,0,0"], "X") + assert_array_equal(r3, np.array(["X,X,X", "X,0", "X"], dtype=dt)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_ufunc.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_ufunc.py new file mode 100644 index 0000000000000000000000000000000000000000..41c2b787385a16b8d32fa2fe173a1bb02c42fdca --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_ufunc.py @@ -0,0 +1,3201 @@ +import warnings +import itertools +import sys +import ctypes as ct +import pickle + +import pytest +from pytest import param + +import numpy as np +import numpy._core.umath as ncu +import numpy._core._umath_tests as umt +import numpy.linalg._umath_linalg as uml +import numpy._core._operand_flag_tests as opflag_tests +import numpy._core._rational_tests as _rational_tests +from numpy.exceptions import AxisError +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_array_equal, + assert_almost_equal, assert_array_almost_equal, assert_no_warnings, + assert_allclose, HAS_REFCOUNT, suppress_warnings, IS_WASM, IS_PYPY, + ) +from numpy.testing._private.utils import requires_memory + + +UNARY_UFUNCS = [obj for obj in np._core.umath.__dict__.values() + if isinstance(obj, np.ufunc)] +UNARY_OBJECT_UFUNCS = [uf for uf in UNARY_UFUNCS if "O->O" in uf.types] + +# Remove functions that do not support `floats` +UNARY_OBJECT_UFUNCS.remove(np.bitwise_count) + + +class TestUfuncKwargs: + def test_kwarg_exact(self): + assert_raises(TypeError, np.add, 1, 2, castingx='safe') + assert_raises(TypeError, np.add, 1, 2, dtypex=int) + assert_raises(TypeError, np.add, 1, 2, extobjx=[4096]) + assert_raises(TypeError, np.add, 1, 2, outx=None) + assert_raises(TypeError, np.add, 1, 2, sigx='ii->i') + assert_raises(TypeError, np.add, 1, 2, signaturex='ii->i') + assert_raises(TypeError, np.add, 1, 2, subokx=False) + assert_raises(TypeError, np.add, 1, 2, wherex=[True]) + + def test_sig_signature(self): + assert_raises(TypeError, np.add, 1, 2, sig='ii->i', + signature='ii->i') + + def test_sig_dtype(self): + assert_raises(TypeError, np.add, 1, 2, sig='ii->i', + dtype=int) + assert_raises(TypeError, np.add, 1, 2, signature='ii->i', + dtype=int) + + def test_extobj_removed(self): + assert_raises(TypeError, np.add, 1, 2, extobj=[4096]) + + +class TestUfuncGenericLoops: + """Test generic loops. + + The loops to be tested are: + + PyUFunc_ff_f_As_dd_d + PyUFunc_ff_f + PyUFunc_dd_d + PyUFunc_gg_g + PyUFunc_FF_F_As_DD_D + PyUFunc_DD_D + PyUFunc_FF_F + PyUFunc_GG_G + PyUFunc_OO_O + PyUFunc_OO_O_method + PyUFunc_f_f_As_d_d + PyUFunc_d_d + PyUFunc_f_f + PyUFunc_g_g + PyUFunc_F_F_As_D_D + PyUFunc_F_F + PyUFunc_D_D + PyUFunc_G_G + PyUFunc_O_O + PyUFunc_O_O_method + PyUFunc_On_Om + + Where: + + f -- float + d -- double + g -- long double + F -- complex float + D -- complex double + G -- complex long double + O -- python object + + It is difficult to assure that each of these loops is entered from the + Python level as the special cased loops are a moving target and the + corresponding types are architecture dependent. We probably need to + define C level testing ufuncs to get at them. For the time being, I've + just looked at the signatures registered in the build directory to find + relevant functions. + + """ + np_dtypes = [ + (np.single, np.single), (np.single, np.double), + (np.csingle, np.csingle), (np.csingle, np.cdouble), + (np.double, np.double), (np.longdouble, np.longdouble), + (np.cdouble, np.cdouble), (np.clongdouble, np.clongdouble)] + + @pytest.mark.parametrize('input_dtype,output_dtype', np_dtypes) + def test_unary_PyUFunc(self, input_dtype, output_dtype, f=np.exp, x=0, y=1): + xs = np.full(10, input_dtype(x), dtype=output_dtype) + ys = f(xs)[::2] + assert_allclose(ys, y) + assert_equal(ys.dtype, output_dtype) + + def f2(x, y): + return x**y + + @pytest.mark.parametrize('input_dtype,output_dtype', np_dtypes) + def test_binary_PyUFunc(self, input_dtype, output_dtype, f=f2, x=0, y=1): + xs = np.full(10, input_dtype(x), dtype=output_dtype) + ys = f(xs, xs)[::2] + assert_allclose(ys, y) + assert_equal(ys.dtype, output_dtype) + + # class to use in testing object method loops + class foo: + def conjugate(self): + return np.bool(1) + + def logical_xor(self, obj): + return np.bool(1) + + def test_unary_PyUFunc_O_O(self): + x = np.ones(10, dtype=object) + assert_(np.all(np.abs(x) == 1)) + + def test_unary_PyUFunc_O_O_method_simple(self, foo=foo): + x = np.full(10, foo(), dtype=object) + assert_(np.all(np.conjugate(x) == True)) + + def test_binary_PyUFunc_OO_O(self): + x = np.ones(10, dtype=object) + assert_(np.all(np.add(x, x) == 2)) + + def test_binary_PyUFunc_OO_O_method(self, foo=foo): + x = np.full(10, foo(), dtype=object) + assert_(np.all(np.logical_xor(x, x))) + + def test_binary_PyUFunc_On_Om_method(self, foo=foo): + x = np.full((10, 2, 3), foo(), dtype=object) + assert_(np.all(np.logical_xor(x, x))) + + def test_python_complex_conjugate(self): + # The conjugate ufunc should fall back to calling the method: + arr = np.array([1+2j, 3-4j], dtype="O") + assert isinstance(arr[0], complex) + res = np.conjugate(arr) + assert res.dtype == np.dtype("O") + assert_array_equal(res, np.array([1-2j, 3+4j], dtype="O")) + + @pytest.mark.parametrize("ufunc", UNARY_OBJECT_UFUNCS) + def test_unary_PyUFunc_O_O_method_full(self, ufunc): + """Compare the result of the object loop with non-object one""" + val = np.float64(np.pi/4) + + class MyFloat(np.float64): + def __getattr__(self, attr): + try: + return super().__getattr__(attr) + except AttributeError: + return lambda: getattr(np._core.umath, attr)(val) + + # Use 0-D arrays, to ensure the same element call + num_arr = np.array(val, dtype=np.float64) + obj_arr = np.array(MyFloat(val), dtype="O") + + with np.errstate(all="raise"): + try: + res_num = ufunc(num_arr) + except Exception as exc: + with assert_raises(type(exc)): + ufunc(obj_arr) + else: + res_obj = ufunc(obj_arr) + assert_array_almost_equal(res_num.astype("O"), res_obj) + + +def _pickleable_module_global(): + pass + + +class TestUfunc: + def test_pickle(self): + for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): + assert_(pickle.loads(pickle.dumps(np.sin, + protocol=proto)) is np.sin) + + # Check that ufunc not defined in the top level numpy namespace + # such as numpy._core._rational_tests.test_add can also be pickled + res = pickle.loads(pickle.dumps(_rational_tests.test_add, + protocol=proto)) + assert_(res is _rational_tests.test_add) + + def test_pickle_withstring(self): + astring = (b"cnumpy.core\n_ufunc_reconstruct\np0\n" + b"(S'numpy._core.umath'\np1\nS'cos'\np2\ntp3\nRp4\n.") + assert_(pickle.loads(astring) is np.cos) + + @pytest.mark.skipif(IS_PYPY, reason="'is' check does not work on PyPy") + def test_pickle_name_is_qualname(self): + # This tests that a simplification of our ufunc pickle code will + # lead to allowing qualnames as names. Future ufuncs should + # possible add a specific qualname, or a hook into pickling instead + # (dask+numba may benefit). + _pickleable_module_global.ufunc = umt._pickleable_module_global_ufunc + + obj = pickle.loads(pickle.dumps(_pickleable_module_global.ufunc)) + assert obj is umt._pickleable_module_global_ufunc + + def test_reduceat_shifting_sum(self): + L = 6 + x = np.arange(L) + idx = np.array(list(zip(np.arange(L - 2), np.arange(L - 2) + 2))).ravel() + assert_array_equal(np.add.reduceat(x, idx)[::2], [1, 3, 5, 7]) + + def test_all_ufunc(self): + """Try to check presence and results of all ufuncs. + + The list of ufuncs comes from generate_umath.py and is as follows: + + ===== ==== ============= =============== ======================== + done args function types notes + ===== ==== ============= =============== ======================== + n 1 conjugate nums + O + n 1 absolute nums + O complex -> real + n 1 negative nums + O + n 1 sign nums + O -> int + n 1 invert bool + ints + O flts raise an error + n 1 degrees real + M cmplx raise an error + n 1 radians real + M cmplx raise an error + n 1 arccos flts + M + n 1 arccosh flts + M + n 1 arcsin flts + M + n 1 arcsinh flts + M + n 1 arctan flts + M + n 1 arctanh flts + M + n 1 cos flts + M + n 1 sin flts + M + n 1 tan flts + M + n 1 cosh flts + M + n 1 sinh flts + M + n 1 tanh flts + M + n 1 exp flts + M + n 1 expm1 flts + M + n 1 log flts + M + n 1 log10 flts + M + n 1 log1p flts + M + n 1 sqrt flts + M real x < 0 raises error + n 1 ceil real + M + n 1 trunc real + M + n 1 floor real + M + n 1 fabs real + M + n 1 rint flts + M + n 1 isnan flts -> bool + n 1 isinf flts -> bool + n 1 isfinite flts -> bool + n 1 signbit real -> bool + n 1 modf real -> (frac, int) + n 1 logical_not bool + nums + M -> bool + n 2 left_shift ints + O flts raise an error + n 2 right_shift ints + O flts raise an error + n 2 add bool + nums + O boolean + is || + n 2 subtract bool + nums + O boolean - is ^ + n 2 multiply bool + nums + O boolean * is & + n 2 divide nums + O + n 2 floor_divide nums + O + n 2 true_divide nums + O bBhH -> f, iIlLqQ -> d + n 2 fmod nums + M + n 2 power nums + O + n 2 greater bool + nums + O -> bool + n 2 greater_equal bool + nums + O -> bool + n 2 less bool + nums + O -> bool + n 2 less_equal bool + nums + O -> bool + n 2 equal bool + nums + O -> bool + n 2 not_equal bool + nums + O -> bool + n 2 logical_and bool + nums + M -> bool + n 2 logical_or bool + nums + M -> bool + n 2 logical_xor bool + nums + M -> bool + n 2 maximum bool + nums + O + n 2 minimum bool + nums + O + n 2 bitwise_and bool + ints + O flts raise an error + n 2 bitwise_or bool + ints + O flts raise an error + n 2 bitwise_xor bool + ints + O flts raise an error + n 2 arctan2 real + M + n 2 remainder ints + real + O + n 2 hypot real + M + ===== ==== ============= =============== ======================== + + Types other than those listed will be accepted, but they are cast to + the smallest compatible type for which the function is defined. The + casting rules are: + + bool -> int8 -> float32 + ints -> double + + """ + pass + + # from include/numpy/ufuncobject.h + size_inferred = 2 + can_ignore = 4 + def test_signature0(self): + # the arguments to test_signature are: nin, nout, core_signature + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 2, 1, "(i),(i)->()") + assert_equal(enabled, 1) + assert_equal(num_dims, (1, 1, 0)) + assert_equal(ixs, (0, 0)) + assert_equal(flags, (self.size_inferred,)) + assert_equal(sizes, (-1,)) + + def test_signature1(self): + # empty core signature; treat as plain ufunc (with trivial core) + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 2, 1, "(),()->()") + assert_equal(enabled, 0) + assert_equal(num_dims, (0, 0, 0)) + assert_equal(ixs, ()) + assert_equal(flags, ()) + assert_equal(sizes, ()) + + def test_signature2(self): + # more complicated names for variables + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 2, 1, "(i1,i2),(J_1)->(_kAB)") + assert_equal(enabled, 1) + assert_equal(num_dims, (2, 1, 1)) + assert_equal(ixs, (0, 1, 2, 3)) + assert_equal(flags, (self.size_inferred,)*4) + assert_equal(sizes, (-1, -1, -1, -1)) + + def test_signature3(self): + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 2, 1, "(i1, i12), (J_1)->(i12, i2)") + assert_equal(enabled, 1) + assert_equal(num_dims, (2, 1, 2)) + assert_equal(ixs, (0, 1, 2, 1, 3)) + assert_equal(flags, (self.size_inferred,)*4) + assert_equal(sizes, (-1, -1, -1, -1)) + + def test_signature4(self): + # matrix_multiply signature from _umath_tests + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 2, 1, "(n,k),(k,m)->(n,m)") + assert_equal(enabled, 1) + assert_equal(num_dims, (2, 2, 2)) + assert_equal(ixs, (0, 1, 1, 2, 0, 2)) + assert_equal(flags, (self.size_inferred,)*3) + assert_equal(sizes, (-1, -1, -1)) + + def test_signature5(self): + # matmul signature from _umath_tests + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 2, 1, "(n?,k),(k,m?)->(n?,m?)") + assert_equal(enabled, 1) + assert_equal(num_dims, (2, 2, 2)) + assert_equal(ixs, (0, 1, 1, 2, 0, 2)) + assert_equal(flags, (self.size_inferred | self.can_ignore, + self.size_inferred, + self.size_inferred | self.can_ignore)) + assert_equal(sizes, (-1, -1, -1)) + + def test_signature6(self): + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 1, 1, "(3)->()") + assert_equal(enabled, 1) + assert_equal(num_dims, (1, 0)) + assert_equal(ixs, (0,)) + assert_equal(flags, (0,)) + assert_equal(sizes, (3,)) + + def test_signature7(self): + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 3, 1, "(3),(03,3),(n)->(9)") + assert_equal(enabled, 1) + assert_equal(num_dims, (1, 2, 1, 1)) + assert_equal(ixs, (0, 0, 0, 1, 2)) + assert_equal(flags, (0, self.size_inferred, 0)) + assert_equal(sizes, (3, -1, 9)) + + def test_signature8(self): + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 3, 1, "(3?),(3?,3?),(n)->(9)") + assert_equal(enabled, 1) + assert_equal(num_dims, (1, 2, 1, 1)) + assert_equal(ixs, (0, 0, 0, 1, 2)) + assert_equal(flags, (self.can_ignore, self.size_inferred, 0)) + assert_equal(sizes, (3, -1, 9)) + + def test_signature9(self): + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 1, 1, "( 3) -> ( )") + assert_equal(enabled, 1) + assert_equal(num_dims, (1, 0)) + assert_equal(ixs, (0,)) + assert_equal(flags, (0,)) + assert_equal(sizes, (3,)) + + def test_signature10(self): + enabled, num_dims, ixs, flags, sizes = umt.test_signature( + 3, 1, "( 3? ) , (3? , 3?) ,(n )-> ( 9)") + assert_equal(enabled, 1) + assert_equal(num_dims, (1, 2, 1, 1)) + assert_equal(ixs, (0, 0, 0, 1, 2)) + assert_equal(flags, (self.can_ignore, self.size_inferred, 0)) + assert_equal(sizes, (3, -1, 9)) + + def test_signature_failure_extra_parenthesis(self): + with assert_raises(ValueError): + umt.test_signature(2, 1, "((i)),(i)->()") + + def test_signature_failure_mismatching_parenthesis(self): + with assert_raises(ValueError): + umt.test_signature(2, 1, "(i),)i(->()") + + def test_signature_failure_signature_missing_input_arg(self): + with assert_raises(ValueError): + umt.test_signature(2, 1, "(i),->()") + + def test_signature_failure_signature_missing_output_arg(self): + with assert_raises(ValueError): + umt.test_signature(2, 2, "(i),(i)->()") + + def test_get_signature(self): + assert_equal(np.vecdot.signature, "(n),(n)->()") + + def test_forced_sig(self): + a = 0.5*np.arange(3, dtype='f8') + assert_equal(np.add(a, 0.5), [0.5, 1, 1.5]) + with pytest.warns(DeprecationWarning): + assert_equal(np.add(a, 0.5, sig='i', casting='unsafe'), [0, 0, 1]) + assert_equal(np.add(a, 0.5, sig='ii->i', casting='unsafe'), [0, 0, 1]) + with pytest.warns(DeprecationWarning): + assert_equal(np.add(a, 0.5, sig=('i4',), casting='unsafe'), + [0, 0, 1]) + assert_equal(np.add(a, 0.5, sig=('i4', 'i4', 'i4'), + casting='unsafe'), [0, 0, 1]) + + b = np.zeros((3,), dtype='f8') + np.add(a, 0.5, out=b) + assert_equal(b, [0.5, 1, 1.5]) + b[:] = 0 + with pytest.warns(DeprecationWarning): + np.add(a, 0.5, sig='i', out=b, casting='unsafe') + assert_equal(b, [0, 0, 1]) + b[:] = 0 + np.add(a, 0.5, sig='ii->i', out=b, casting='unsafe') + assert_equal(b, [0, 0, 1]) + b[:] = 0 + with pytest.warns(DeprecationWarning): + np.add(a, 0.5, sig=('i4',), out=b, casting='unsafe') + assert_equal(b, [0, 0, 1]) + b[:] = 0 + np.add(a, 0.5, sig=('i4', 'i4', 'i4'), out=b, casting='unsafe') + assert_equal(b, [0, 0, 1]) + + def test_signature_all_None(self): + # signature all None, is an acceptable alternative (since 1.21) + # to not providing a signature. + res1 = np.add([3], [4], sig=(None, None, None)) + res2 = np.add([3], [4]) + assert_array_equal(res1, res2) + res1 = np.maximum([3], [4], sig=(None, None, None)) + res2 = np.maximum([3], [4]) + assert_array_equal(res1, res2) + + with pytest.raises(TypeError): + # special case, that would be deprecated anyway, so errors: + np.add(3, 4, signature=(None,)) + + def test_signature_dtype_type(self): + # Since that will be the normal behaviour (past NumPy 1.21) + # we do support the types already: + float_dtype = type(np.dtype(np.float64)) + np.add(3, 4, signature=(float_dtype, float_dtype, None)) + + @pytest.mark.parametrize("get_kwarg", [ + lambda dt: dict(dtype=dt), + lambda dt: dict(signature=(dt, None, None))]) + def test_signature_dtype_instances_allowed(self, get_kwarg): + # We allow certain dtype instances when there is a clear singleton + # and the given one is equivalent; mainly for backcompat. + int64 = np.dtype("int64") + int64_2 = pickle.loads(pickle.dumps(int64)) + # Relies on pickling behavior, if assert fails just remove test... + assert int64 is not int64_2 + + assert np.add(1, 2, **get_kwarg(int64_2)).dtype == int64 + td = np.timedelta(2, "s") + assert np.add(td, td, **get_kwarg("m8")).dtype == "m8[s]" + + @pytest.mark.parametrize("get_kwarg", [ + param(lambda x: dict(dtype=x), id="dtype"), + param(lambda x: dict(signature=(x, None, None)), id="signature")]) + def test_signature_dtype_instances_allowed(self, get_kwarg): + msg = "The `dtype` and `signature` arguments to ufuncs" + + with pytest.raises(TypeError, match=msg): + np.add(3, 5, **get_kwarg(np.dtype("int64").newbyteorder())) + with pytest.raises(TypeError, match=msg): + np.add(3, 5, **get_kwarg(np.dtype("m8[ns]"))) + with pytest.raises(TypeError, match=msg): + np.add(3, 5, **get_kwarg("m8[ns]")) + + @pytest.mark.parametrize("casting", ["unsafe", "same_kind", "safe"]) + def test_partial_signature_mismatch(self, casting): + # If the second argument matches already, no need to specify it: + res = np.ldexp(np.float32(1.), np.int_(2), dtype="d") + assert res.dtype == "d" + res = np.ldexp(np.float32(1.), np.int_(2), signature=(None, None, "d")) + assert res.dtype == "d" + + # ldexp only has a loop for long input as second argument, overriding + # the output cannot help with that (no matter the casting) + with pytest.raises(TypeError): + np.ldexp(1., np.uint64(3), dtype="d") + with pytest.raises(TypeError): + np.ldexp(1., np.uint64(3), signature=(None, None, "d")) + + def test_partial_signature_mismatch_with_cache(self): + with pytest.raises(TypeError): + np.add(np.float16(1), np.uint64(2), sig=("e", "d", None)) + # Ensure e,d->None is in the dispatching cache (double loop) + np.add(np.float16(1), np.float64(2)) + # The error must still be raised: + with pytest.raises(TypeError): + np.add(np.float16(1), np.uint64(2), sig=("e", "d", None)) + + def test_use_output_signature_for_all_arguments(self): + # Test that providing only `dtype=` or `signature=(None, None, dtype)` + # is sufficient if falling back to a homogeneous signature works. + # In this case, the `intp, intp -> intp` loop is chosen. + res = np.power(1.5, 2.8, dtype=np.intp, casting="unsafe") + assert res == 1 # the cast happens first. + res = np.power(1.5, 2.8, signature=(None, None, np.intp), + casting="unsafe") + assert res == 1 + with pytest.raises(TypeError): + # the unsafe casting would normally cause errors though: + np.power(1.5, 2.8, dtype=np.intp) + + def test_signature_errors(self): + with pytest.raises(TypeError, + match="the signature object to ufunc must be a string or"): + np.add(3, 4, signature=123.) # neither a string nor a tuple + + with pytest.raises(ValueError): + # bad symbols that do not translate to dtypes + np.add(3, 4, signature="%^->#") + + with pytest.raises(ValueError): + np.add(3, 4, signature=b"ii-i") # incomplete and byte string + + with pytest.raises(ValueError): + np.add(3, 4, signature="ii>i") # incomplete string + + with pytest.raises(ValueError): + np.add(3, 4, signature=(None, "f8")) # bad length + + with pytest.raises(UnicodeDecodeError): + np.add(3, 4, signature=b"\xff\xff->i") + + def test_forced_dtype_times(self): + # Signatures only set the type numbers (not the actual loop dtypes) + # so using `M` in a signature/dtype should generally work: + a = np.array(['2010-01-02', '1999-03-14', '1833-03'], dtype='>M8[D]') + np.maximum(a, a, dtype="M") + np.maximum.reduce(a, dtype="M") + + arr = np.arange(10, dtype="m8[s]") + np.add(arr, arr, dtype="m") + np.maximum(arr, arr, dtype="m") + + @pytest.mark.parametrize("ufunc", [np.add, np.sqrt]) + def test_cast_safety(self, ufunc): + """Basic test for the safest casts, because ufuncs inner loops can + indicate a cast-safety as well (which is normally always "no"). + """ + def call_ufunc(arr, **kwargs): + return ufunc(*(arr,) * ufunc.nin, **kwargs) + + arr = np.array([1., 2., 3.], dtype=np.float32) + arr_bs = arr.astype(arr.dtype.newbyteorder()) + expected = call_ufunc(arr) + # Normally, a "no" cast: + res = call_ufunc(arr, casting="no") + assert_array_equal(expected, res) + # Byte-swapping is not allowed with "no" though: + with pytest.raises(TypeError): + call_ufunc(arr_bs, casting="no") + + # But is allowed with "equiv": + res = call_ufunc(arr_bs, casting="equiv") + assert_array_equal(expected, res) + + # Casting to float64 is safe, but not equiv: + with pytest.raises(TypeError): + call_ufunc(arr_bs, dtype=np.float64, casting="equiv") + + # but it is safe cast: + res = call_ufunc(arr_bs, dtype=np.float64, casting="safe") + expected = call_ufunc(arr.astype(np.float64)) # upcast + assert_array_equal(expected, res) + + @pytest.mark.parametrize("ufunc", [np.add, np.equal]) + def test_cast_safety_scalar(self, ufunc): + # We test add and equal, because equal has special scalar handling + # Note that the "equiv" casting behavior should maybe be considered + # a current implementation detail. + with pytest.raises(TypeError): + # this picks an integer loop, which is not safe + ufunc(3., 4., dtype=int, casting="safe") + + with pytest.raises(TypeError): + # We accept python float as float64 but not float32 for equiv. + ufunc(3., 4., dtype="float32", casting="equiv") + + # Special case for object and equal (note that equiv implies safe) + ufunc(3, 4, dtype=object, casting="equiv") + # Picks a double loop for both, first is equiv, second safe: + ufunc(np.array([3.]), 3., casting="equiv") + ufunc(np.array([3.]), 3, casting="safe") + ufunc(np.array([3]), 3, casting="equiv") + + def test_cast_safety_scalar_special(self): + # We allow this (and it succeeds) via object, although the equiv + # part may not be important. + np.equal(np.array([3]), 2**300, casting="equiv") + + def test_true_divide(self): + a = np.array(10) + b = np.array(20) + tgt = np.array(0.5) + + for tc in 'bhilqBHILQefdgFDG': + dt = np.dtype(tc) + aa = a.astype(dt) + bb = b.astype(dt) + + # Check result value and dtype. + for x, y in itertools.product([aa, -aa], [bb, -bb]): + + # Check with no output type specified + if tc in 'FDG': + tgt = complex(x)/complex(y) + else: + tgt = float(x)/float(y) + + res = np.true_divide(x, y) + rtol = max(np.finfo(res).resolution, 1e-15) + assert_allclose(res, tgt, rtol=rtol) + + if tc in 'bhilqBHILQ': + assert_(res.dtype.name == 'float64') + else: + assert_(res.dtype.name == dt.name ) + + # Check with output type specified. This also checks for the + # incorrect casts in issue gh-3484 because the unary '-' does + # not change types, even for unsigned types, Hence casts in the + # ufunc from signed to unsigned and vice versa will lead to + # errors in the values. + for tcout in 'bhilqBHILQ': + dtout = np.dtype(tcout) + assert_raises(TypeError, np.true_divide, x, y, dtype=dtout) + + for tcout in 'efdg': + dtout = np.dtype(tcout) + if tc in 'FDG': + # Casting complex to float is not allowed + assert_raises(TypeError, np.true_divide, x, y, dtype=dtout) + else: + tgt = float(x)/float(y) + rtol = max(np.finfo(dtout).resolution, 1e-15) + # The value of tiny for double double is NaN + with suppress_warnings() as sup: + sup.filter(UserWarning) + if not np.isnan(np.finfo(dtout).tiny): + atol = max(np.finfo(dtout).tiny, 3e-308) + else: + atol = 3e-308 + # Some test values result in invalid for float16 + # and the cast to it may overflow to inf. + with np.errstate(invalid='ignore', over='ignore'): + res = np.true_divide(x, y, dtype=dtout) + if not np.isfinite(res) and tcout == 'e': + continue + assert_allclose(res, tgt, rtol=rtol, atol=atol) + assert_(res.dtype.name == dtout.name) + + for tcout in 'FDG': + dtout = np.dtype(tcout) + tgt = complex(x)/complex(y) + rtol = max(np.finfo(dtout).resolution, 1e-15) + # The value of tiny for double double is NaN + with suppress_warnings() as sup: + sup.filter(UserWarning) + if not np.isnan(np.finfo(dtout).tiny): + atol = max(np.finfo(dtout).tiny, 3e-308) + else: + atol = 3e-308 + res = np.true_divide(x, y, dtype=dtout) + if not np.isfinite(res): + continue + assert_allclose(res, tgt, rtol=rtol, atol=atol) + assert_(res.dtype.name == dtout.name) + + # Check booleans + a = np.ones((), dtype=np.bool) + res = np.true_divide(a, a) + assert_(res == 1.0) + assert_(res.dtype.name == 'float64') + res = np.true_divide(~a, a) + assert_(res == 0.0) + assert_(res.dtype.name == 'float64') + + def test_sum_stability(self): + a = np.ones(500, dtype=np.float32) + assert_almost_equal((a / 10.).sum() - a.size / 10., 0, 4) + + a = np.ones(500, dtype=np.float64) + assert_almost_equal((a / 10.).sum() - a.size / 10., 0, 13) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_sum(self): + for dt in (int, np.float16, np.float32, np.float64, np.longdouble): + for v in (0, 1, 2, 7, 8, 9, 15, 16, 19, 127, + 128, 1024, 1235): + # warning if sum overflows, which it does in float16 + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always", RuntimeWarning) + + tgt = dt(v * (v + 1) / 2) + overflow = not np.isfinite(tgt) + assert_equal(len(w), 1 * overflow) + + d = np.arange(1, v + 1, dtype=dt) + + assert_almost_equal(np.sum(d), tgt) + assert_equal(len(w), 2 * overflow) + + assert_almost_equal(np.sum(d[::-1]), tgt) + assert_equal(len(w), 3 * overflow) + + d = np.ones(500, dtype=dt) + assert_almost_equal(np.sum(d[::2]), 250.) + assert_almost_equal(np.sum(d[1::2]), 250.) + assert_almost_equal(np.sum(d[::3]), 167.) + assert_almost_equal(np.sum(d[1::3]), 167.) + assert_almost_equal(np.sum(d[::-2]), 250.) + assert_almost_equal(np.sum(d[-1::-2]), 250.) + assert_almost_equal(np.sum(d[::-3]), 167.) + assert_almost_equal(np.sum(d[-1::-3]), 167.) + # sum with first reduction entry != 0 + d = np.ones((1,), dtype=dt) + d += d + assert_almost_equal(d, 2.) + + def test_sum_complex(self): + for dt in (np.complex64, np.complex128, np.clongdouble): + for v in (0, 1, 2, 7, 8, 9, 15, 16, 19, 127, + 128, 1024, 1235): + tgt = dt(v * (v + 1) / 2) - dt((v * (v + 1) / 2) * 1j) + d = np.empty(v, dtype=dt) + d.real = np.arange(1, v + 1) + d.imag = -np.arange(1, v + 1) + assert_almost_equal(np.sum(d), tgt) + assert_almost_equal(np.sum(d[::-1]), tgt) + + d = np.ones(500, dtype=dt) + 1j + assert_almost_equal(np.sum(d[::2]), 250. + 250j) + assert_almost_equal(np.sum(d[1::2]), 250. + 250j) + assert_almost_equal(np.sum(d[::3]), 167. + 167j) + assert_almost_equal(np.sum(d[1::3]), 167. + 167j) + assert_almost_equal(np.sum(d[::-2]), 250. + 250j) + assert_almost_equal(np.sum(d[-1::-2]), 250. + 250j) + assert_almost_equal(np.sum(d[::-3]), 167. + 167j) + assert_almost_equal(np.sum(d[-1::-3]), 167. + 167j) + # sum with first reduction entry != 0 + d = np.ones((1,), dtype=dt) + 1j + d += d + assert_almost_equal(d, 2. + 2j) + + def test_sum_initial(self): + # Integer, single axis + assert_equal(np.sum([3], initial=2), 5) + + # Floating point + assert_almost_equal(np.sum([0.2], initial=0.1), 0.3) + + # Multiple non-adjacent axes + assert_equal(np.sum(np.ones((2, 3, 5), dtype=np.int64), axis=(0, 2), initial=2), + [12, 12, 12]) + + def test_sum_where(self): + # More extensive tests done in test_reduction_with_where. + assert_equal(np.sum([[1., 2.], [3., 4.]], where=[True, False]), 4.) + assert_equal(np.sum([[1., 2.], [3., 4.]], axis=0, initial=5., + where=[True, False]), [9., 5.]) + + def test_vecdot(self): + arr1 = np.arange(6).reshape((2, 3)) + arr2 = np.arange(3).reshape((1, 3)) + + actual = np.vecdot(arr1, arr2) + expected = np.array([5, 14]) + + assert_array_equal(actual, expected) + + actual2 = np.vecdot(arr1.T, arr2.T, axis=-2) + assert_array_equal(actual2, expected) + + actual3 = np.vecdot(arr1.astype("object"), arr2) + assert_array_equal(actual3, expected.astype("object")) + + def test_matvec(self): + arr1 = np.arange(6).reshape((2, 3)) + arr2 = np.arange(3).reshape((1, 3)) + + actual = np.matvec(arr1, arr2) + expected = np.array([[5, 14]]) + + assert_array_equal(actual, expected) + + actual2 = np.matvec(arr1.T, arr2.T, axes=[(-1, -2), -2, -1]) + assert_array_equal(actual2, expected) + + actual3 = np.matvec(arr1.astype("object"), arr2) + assert_array_equal(actual3, expected.astype("object")) + + @pytest.mark.parametrize("vec", [ + np.array([[1., 2., 3.], [4., 5., 6.]]), + np.array([[1., 2j, 3.], [4., 5., 6j]]), + np.array([[1., 2., 3.], [4., 5., 6.]], dtype=object), + np.array([[1., 2j, 3.], [4., 5., 6j]], dtype=object)]) + @pytest.mark.parametrize("matrix", [ + None, + np.array([[1.+1j, 0.5, -0.5j], + [0.25, 2j, 0.], + [4., 0., -1j]])]) + def test_vecmatvec_identity(self, matrix, vec): + """Check that (x†A)x equals x†(Ax).""" + mat = matrix if matrix is not None else np.eye(3) + matvec = np.matvec(mat, vec) # Ax + vecmat = np.vecmat(vec, mat) # x†A + if matrix is None: + assert_array_equal(matvec, vec) + assert_array_equal(vecmat.conj(), vec) + assert_array_equal(matvec, (mat @ vec[..., np.newaxis]).squeeze(-1)) + assert_array_equal(vecmat, (vec[..., np.newaxis].mT.conj() + @ mat).squeeze(-2)) + expected = np.einsum('...i,ij,...j', vec.conj(), mat, vec) + vec_matvec = (vec.conj() * matvec).sum(-1) + vecmat_vec = (vecmat * vec).sum(-1) + assert_array_equal(vec_matvec, expected) + assert_array_equal(vecmat_vec, expected) + + @pytest.mark.parametrize("ufunc, shape1, shape2, conj", [ + (np.vecdot, (3,), (3,), True), + (np.vecmat, (3,), (3, 1), True), + (np.matvec, (1, 3), (3,), False), + (np.matmul, (1, 3), (3, 1), False), + ]) + def test_vecdot_matvec_vecmat_complex(self, ufunc, shape1, shape2, conj): + arr1 = np.array([1, 2j, 3]) + arr2 = np.array([1, 2, 3]) + + actual1 = ufunc(arr1.reshape(shape1), arr2.reshape(shape2)) + expected1 = np.array(((arr1.conj() if conj else arr1) * arr2).sum(), + ndmin=min(len(shape1), len(shape2))) + assert_array_equal(actual1, expected1) + # This would fail for conj=True, since matmul omits the conjugate. + if not conj: + assert_array_equal(arr1.reshape(shape1) @ arr2.reshape(shape2), + expected1) + + actual2 = ufunc(arr2.reshape(shape1), arr1.reshape(shape2)) + expected2 = np.array(((arr2.conj() if conj else arr2) * arr1).sum(), + ndmin=min(len(shape1), len(shape2))) + assert_array_equal(actual2, expected2) + + actual3 = ufunc(arr1.reshape(shape1).astype("object"), + arr2.reshape(shape2).astype("object")) + expected3 = expected1.astype(object) + assert_array_equal(actual3, expected3) + + def test_vecdot_subclass(self): + class MySubclass(np.ndarray): + pass + + arr1 = np.arange(6).reshape((2, 3)).view(MySubclass) + arr2 = np.arange(3).reshape((1, 3)).view(MySubclass) + result = np.vecdot(arr1, arr2) + assert isinstance(result, MySubclass) + + def test_vecdot_object_no_conjugate(self): + arr = np.array(["1", "2"], dtype=object) + with pytest.raises(AttributeError, match="conjugate"): + np.vecdot(arr, arr) + + def test_vecdot_object_breaks_outer_loop_on_error(self): + arr1 = np.ones((3, 3)).astype(object) + arr2 = arr1.copy() + arr2[1, 1] = None + out = np.zeros(3).astype(object) + with pytest.raises(TypeError, match=r"\*: 'float' and 'NoneType'"): + np.vecdot(arr1, arr2, out=out) + assert out[0] == 3 + assert out[1] == out[2] == 0 + + def test_broadcast(self): + msg = "broadcast" + a = np.arange(4).reshape((2, 1, 2)) + b = np.arange(4).reshape((1, 2, 2)) + assert_array_equal(np.vecdot(a, b), np.sum(a*b, axis=-1), err_msg=msg) + msg = "extend & broadcast loop dimensions" + b = np.arange(4).reshape((2, 2)) + assert_array_equal(np.vecdot(a, b), np.sum(a*b, axis=-1), err_msg=msg) + # Broadcast in core dimensions should fail + a = np.arange(8).reshape((4, 2)) + b = np.arange(4).reshape((4, 1)) + assert_raises(ValueError, np.vecdot, a, b) + # Extend core dimensions should fail + a = np.arange(8).reshape((4, 2)) + b = np.array(7) + assert_raises(ValueError, np.vecdot, a, b) + # Broadcast should fail + a = np.arange(2).reshape((2, 1, 1)) + b = np.arange(3).reshape((3, 1, 1)) + assert_raises(ValueError, np.vecdot, a, b) + + # Writing to a broadcasted array with overlap should warn, gh-2705 + a = np.arange(2) + b = np.arange(4).reshape((2, 2)) + u, v = np.broadcast_arrays(a, b) + assert_equal(u.strides[0], 0) + x = u + v + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + u += v + assert_equal(len(w), 1) + assert_(x[0, 0] != u[0, 0]) + + # Output reduction should not be allowed. + # See gh-15139 + a = np.arange(6).reshape(3, 2) + b = np.ones(2) + out = np.empty(()) + assert_raises(ValueError, np.vecdot, a, b, out) + out2 = np.empty(3) + c = np.vecdot(a, b, out2) + assert_(c is out2) + + def test_out_broadcasts(self): + # For ufuncs and gufuncs (not for reductions), we currently allow + # the output to cause broadcasting of the input arrays. + # both along dimensions with shape 1 and dimensions which do not + # exist at all in the inputs. + arr = np.arange(3).reshape(1, 3) + out = np.empty((5, 4, 3)) + np.add(arr, arr, out=out) + assert (out == np.arange(3) * 2).all() + + # The same holds for gufuncs (gh-16484) + np.vecdot(arr, arr, out=out) + # the result would be just a scalar `5`, but is broadcast fully: + assert (out == 5).all() + + @pytest.mark.parametrize(["arr", "out"], [ + ([2], np.empty(())), + ([1, 2], np.empty(1)), + (np.ones((4, 3)), np.empty((4, 1)))], + ids=["(1,)->()", "(2,)->(1,)", "(4, 3)->(4, 1)"]) + def test_out_broadcast_errors(self, arr, out): + # Output is (currently) allowed to broadcast inputs, but it cannot be + # smaller than the actual result. + with pytest.raises(ValueError, match="non-broadcastable"): + np.positive(arr, out=out) + + with pytest.raises(ValueError, match="non-broadcastable"): + np.add(np.ones(()), arr, out=out) + + def test_type_cast(self): + msg = "type cast" + a = np.arange(6, dtype='short').reshape((2, 3)) + assert_array_equal(np.vecdot(a, a), np.sum(a*a, axis=-1), + err_msg=msg) + msg = "type cast on one argument" + a = np.arange(6).reshape((2, 3)) + b = a + 0.1 + assert_array_almost_equal(np.vecdot(a, b), np.sum(a*b, axis=-1), + err_msg=msg) + + def test_endian(self): + msg = "big endian" + a = np.arange(6, dtype='>i4').reshape((2, 3)) + assert_array_equal(np.vecdot(a, a), np.sum(a*a, axis=-1), + err_msg=msg) + msg = "little endian" + a = np.arange(6, dtype='()' + a = np.arange(27.).reshape((3, 3, 3)) + b = np.arange(10., 19.).reshape((3, 1, 3)) + # basic tests on inputs (outputs tested below with matrix_multiply). + c = np.vecdot(a, b) + assert_array_equal(c, (a * b).sum(-1)) + # default + c = np.vecdot(a, b, axes=[(-1,), (-1,), ()]) + assert_array_equal(c, (a * b).sum(-1)) + # integers ok for single axis. + c = np.vecdot(a, b, axes=[-1, -1, ()]) + assert_array_equal(c, (a * b).sum(-1)) + # mix fine + c = np.vecdot(a, b, axes=[(-1,), -1, ()]) + assert_array_equal(c, (a * b).sum(-1)) + # can omit last axis. + c = np.vecdot(a, b, axes=[-1, -1]) + assert_array_equal(c, (a * b).sum(-1)) + # can pass in other types of integer (with __index__ protocol) + c = np.vecdot(a, b, axes=[np.int8(-1), np.array(-1, dtype=np.int32)]) + assert_array_equal(c, (a * b).sum(-1)) + # swap some axes + c = np.vecdot(a, b, axes=[0, 0]) + assert_array_equal(c, (a * b).sum(0)) + c = np.vecdot(a, b, axes=[0, 2]) + assert_array_equal(c, (a.transpose(1, 2, 0) * b).sum(-1)) + # Check errors for improperly constructed axes arguments. + # should have list. + assert_raises(TypeError, np.vecdot, a, b, axes=-1) + # needs enough elements + assert_raises(ValueError, np.vecdot, a, b, axes=[-1]) + # should pass in indices. + assert_raises(TypeError, np.vecdot, a, b, axes=[-1.0, -1.0]) + assert_raises(TypeError, np.vecdot, a, b, axes=[(-1.0,), -1]) + assert_raises(TypeError, np.vecdot, a, b, axes=[None, 1]) + # cannot pass an index unless there is only one dimension + # (output is wrong in this case) + assert_raises(AxisError, np.vecdot, a, b, axes=[-1, -1, -1]) + # or pass in generally the wrong number of axes + assert_raises(AxisError, np.vecdot, a, b, axes=[-1, -1, (-1,)]) + assert_raises(AxisError, np.vecdot, a, b, axes=[-1, (-2, -1), ()]) + # axes need to have same length. + assert_raises(ValueError, np.vecdot, a, b, axes=[0, 1]) + + # matrix_multiply signature: '(m,n),(n,p)->(m,p)' + mm = umt.matrix_multiply + a = np.arange(12).reshape((2, 3, 2)) + b = np.arange(8).reshape((2, 2, 2, 1)) + 1 + # Sanity check. + c = mm(a, b) + assert_array_equal(c, np.matmul(a, b)) + # Default axes. + c = mm(a, b, axes=[(-2, -1), (-2, -1), (-2, -1)]) + assert_array_equal(c, np.matmul(a, b)) + # Default with explicit axes. + c = mm(a, b, axes=[(1, 2), (2, 3), (2, 3)]) + assert_array_equal(c, np.matmul(a, b)) + # swap some axes. + c = mm(a, b, axes=[(0, -1), (1, 2), (-2, -1)]) + assert_array_equal(c, np.matmul(a.transpose(1, 0, 2), + b.transpose(0, 3, 1, 2))) + # Default with output array. + c = np.empty((2, 2, 3, 1)) + d = mm(a, b, out=c, axes=[(1, 2), (2, 3), (2, 3)]) + assert_(c is d) + assert_array_equal(c, np.matmul(a, b)) + # Transposed output array + c = np.empty((1, 2, 2, 3)) + d = mm(a, b, out=c, axes=[(-2, -1), (-2, -1), (3, 0)]) + assert_(c is d) + assert_array_equal(c, np.matmul(a, b).transpose(3, 0, 1, 2)) + # Check errors for improperly constructed axes arguments. + # wrong argument + assert_raises(TypeError, mm, a, b, axis=1) + # axes should be list + assert_raises(TypeError, mm, a, b, axes=1) + assert_raises(TypeError, mm, a, b, axes=((-2, -1), (-2, -1), (-2, -1))) + # list needs to have right length + assert_raises(ValueError, mm, a, b, axes=[]) + assert_raises(ValueError, mm, a, b, axes=[(-2, -1)]) + # list should not contain None, or lists + assert_raises(TypeError, mm, a, b, axes=[None, None, None]) + assert_raises(TypeError, + mm, a, b, axes=[[-2, -1], [-2, -1], [-2, -1]]) + assert_raises(TypeError, + mm, a, b, axes=[(-2, -1), (-2, -1), [-2, -1]]) + assert_raises(TypeError, mm, a, b, axes=[(-2, -1), (-2, -1), None]) + # single integers are AxisErrors if more are required + assert_raises(AxisError, mm, a, b, axes=[-1, -1, -1]) + assert_raises(AxisError, mm, a, b, axes=[(-2, -1), (-2, -1), -1]) + # tuples should not have duplicated values + assert_raises(ValueError, mm, a, b, axes=[(-2, -1), (-2, -1), (-2, -2)]) + # arrays should have enough axes. + z = np.zeros((2, 2)) + assert_raises(ValueError, mm, z, z[0]) + assert_raises(ValueError, mm, z, z, out=z[:, 0]) + assert_raises(ValueError, mm, z[1], z, axes=[0, 1]) + assert_raises(ValueError, mm, z, z, out=z[0], axes=[0, 1]) + # Regular ufuncs should not accept axes. + assert_raises(TypeError, np.add, 1., 1., axes=[0]) + # should be able to deal with bad unrelated kwargs. + assert_raises(TypeError, mm, z, z, axes=[0, 1], parrot=True) + + def test_axis_argument(self): + # vecdot signature: '(n),(n)->()' + a = np.arange(27.).reshape((3, 3, 3)) + b = np.arange(10., 19.).reshape((3, 1, 3)) + c = np.vecdot(a, b) + assert_array_equal(c, (a * b).sum(-1)) + c = np.vecdot(a, b, axis=-1) + assert_array_equal(c, (a * b).sum(-1)) + out = np.zeros_like(c) + d = np.vecdot(a, b, axis=-1, out=out) + assert_(d is out) + assert_array_equal(d, c) + c = np.vecdot(a, b, axis=0) + assert_array_equal(c, (a * b).sum(0)) + # Sanity checks on innerwt and cumsum. + a = np.arange(6).reshape((2, 3)) + b = np.arange(10, 16).reshape((2, 3)) + w = np.arange(20, 26).reshape((2, 3)) + assert_array_equal(umt.innerwt(a, b, w, axis=0), + np.sum(a * b * w, axis=0)) + assert_array_equal(umt.cumsum(a, axis=0), np.cumsum(a, axis=0)) + assert_array_equal(umt.cumsum(a, axis=-1), np.cumsum(a, axis=-1)) + out = np.empty_like(a) + b = umt.cumsum(a, out=out, axis=0) + assert_(out is b) + assert_array_equal(b, np.cumsum(a, axis=0)) + b = umt.cumsum(a, out=out, axis=1) + assert_(out is b) + assert_array_equal(b, np.cumsum(a, axis=-1)) + # Check errors. + # Cannot pass in both axis and axes. + assert_raises(TypeError, np.vecdot, a, b, axis=0, axes=[0, 0]) + # Not an integer. + assert_raises(TypeError, np.vecdot, a, b, axis=[0]) + # more than 1 core dimensions. + mm = umt.matrix_multiply + assert_raises(TypeError, mm, a, b, axis=1) + # Output wrong size in axis. + out = np.empty((1, 2, 3), dtype=a.dtype) + assert_raises(ValueError, umt.cumsum, a, out=out, axis=0) + # Regular ufuncs should not accept axis. + assert_raises(TypeError, np.add, 1., 1., axis=0) + + def test_keepdims_argument(self): + # vecdot signature: '(n),(n)->()' + a = np.arange(27.).reshape((3, 3, 3)) + b = np.arange(10., 19.).reshape((3, 1, 3)) + c = np.vecdot(a, b) + assert_array_equal(c, (a * b).sum(-1)) + c = np.vecdot(a, b, keepdims=False) + assert_array_equal(c, (a * b).sum(-1)) + c = np.vecdot(a, b, keepdims=True) + assert_array_equal(c, (a * b).sum(-1, keepdims=True)) + out = np.zeros_like(c) + d = np.vecdot(a, b, keepdims=True, out=out) + assert_(d is out) + assert_array_equal(d, c) + # Now combined with axis and axes. + c = np.vecdot(a, b, axis=-1, keepdims=False) + assert_array_equal(c, (a * b).sum(-1, keepdims=False)) + c = np.vecdot(a, b, axis=-1, keepdims=True) + assert_array_equal(c, (a * b).sum(-1, keepdims=True)) + c = np.vecdot(a, b, axis=0, keepdims=False) + assert_array_equal(c, (a * b).sum(0, keepdims=False)) + c = np.vecdot(a, b, axis=0, keepdims=True) + assert_array_equal(c, (a * b).sum(0, keepdims=True)) + c = np.vecdot(a, b, axes=[(-1,), (-1,), ()], keepdims=False) + assert_array_equal(c, (a * b).sum(-1)) + c = np.vecdot(a, b, axes=[(-1,), (-1,), (-1,)], keepdims=True) + assert_array_equal(c, (a * b).sum(-1, keepdims=True)) + c = np.vecdot(a, b, axes=[0, 0], keepdims=False) + assert_array_equal(c, (a * b).sum(0)) + c = np.vecdot(a, b, axes=[0, 0, 0], keepdims=True) + assert_array_equal(c, (a * b).sum(0, keepdims=True)) + c = np.vecdot(a, b, axes=[0, 2], keepdims=False) + assert_array_equal(c, (a.transpose(1, 2, 0) * b).sum(-1)) + c = np.vecdot(a, b, axes=[0, 2], keepdims=True) + assert_array_equal(c, (a.transpose(1, 2, 0) * b).sum(-1, + keepdims=True)) + c = np.vecdot(a, b, axes=[0, 2, 2], keepdims=True) + assert_array_equal(c, (a.transpose(1, 2, 0) * b).sum(-1, + keepdims=True)) + c = np.vecdot(a, b, axes=[0, 2, 0], keepdims=True) + assert_array_equal(c, (a * b.transpose(2, 0, 1)).sum(0, keepdims=True)) + # Hardly useful, but should work. + c = np.vecdot(a, b, axes=[0, 2, 1], keepdims=True) + assert_array_equal(c, (a.transpose(1, 0, 2) * b.transpose(0, 2, 1)) + .sum(1, keepdims=True)) + # Check with two core dimensions. + a = np.eye(3) * np.arange(4.)[:, np.newaxis, np.newaxis] + expected = uml.det(a) + c = uml.det(a, keepdims=False) + assert_array_equal(c, expected) + c = uml.det(a, keepdims=True) + assert_array_equal(c, expected[:, np.newaxis, np.newaxis]) + a = np.eye(3) * np.arange(4.)[:, np.newaxis, np.newaxis] + expected_s, expected_l = uml.slogdet(a) + cs, cl = uml.slogdet(a, keepdims=False) + assert_array_equal(cs, expected_s) + assert_array_equal(cl, expected_l) + cs, cl = uml.slogdet(a, keepdims=True) + assert_array_equal(cs, expected_s[:, np.newaxis, np.newaxis]) + assert_array_equal(cl, expected_l[:, np.newaxis, np.newaxis]) + # Sanity check on innerwt. + a = np.arange(6).reshape((2, 3)) + b = np.arange(10, 16).reshape((2, 3)) + w = np.arange(20, 26).reshape((2, 3)) + assert_array_equal(umt.innerwt(a, b, w, keepdims=True), + np.sum(a * b * w, axis=-1, keepdims=True)) + assert_array_equal(umt.innerwt(a, b, w, axis=0, keepdims=True), + np.sum(a * b * w, axis=0, keepdims=True)) + # Check errors. + # Not a boolean + assert_raises(TypeError, np.vecdot, a, b, keepdims='true') + # More than 1 core dimension, and core output dimensions. + mm = umt.matrix_multiply + assert_raises(TypeError, mm, a, b, keepdims=True) + assert_raises(TypeError, mm, a, b, keepdims=False) + # Regular ufuncs should not accept keepdims. + assert_raises(TypeError, np.add, 1., 1., keepdims=False) + + def test_innerwt(self): + a = np.arange(6).reshape((2, 3)) + b = np.arange(10, 16).reshape((2, 3)) + w = np.arange(20, 26).reshape((2, 3)) + assert_array_equal(umt.innerwt(a, b, w), np.sum(a*b*w, axis=-1)) + a = np.arange(100, 124).reshape((2, 3, 4)) + b = np.arange(200, 224).reshape((2, 3, 4)) + w = np.arange(300, 324).reshape((2, 3, 4)) + assert_array_equal(umt.innerwt(a, b, w), np.sum(a*b*w, axis=-1)) + + def test_innerwt_empty(self): + """Test generalized ufunc with zero-sized operands""" + a = np.array([], dtype='f8') + b = np.array([], dtype='f8') + w = np.array([], dtype='f8') + assert_array_equal(umt.innerwt(a, b, w), np.sum(a*b*w, axis=-1)) + + def test_cross1d(self): + """Test with fixed-sized signature.""" + a = np.eye(3) + assert_array_equal(umt.cross1d(a, a), np.zeros((3, 3))) + out = np.zeros((3, 3)) + result = umt.cross1d(a[0], a, out) + assert_(result is out) + assert_array_equal(result, np.vstack((np.zeros(3), a[2], -a[1]))) + assert_raises(ValueError, umt.cross1d, np.eye(4), np.eye(4)) + assert_raises(ValueError, umt.cross1d, a, np.arange(4.)) + # Wrong output core dimension. + assert_raises(ValueError, umt.cross1d, a, np.arange(3.), np.zeros((3, 4))) + # Wrong output broadcast dimension (see gh-15139). + assert_raises(ValueError, umt.cross1d, a, np.arange(3.), np.zeros(3)) + + def test_can_ignore_signature(self): + # Comparing the effects of ? in signature: + # matrix_multiply: (m,n),(n,p)->(m,p) # all must be there. + # matmul: (m?,n),(n,p?)->(m?,p?) # allow missing m, p. + mat = np.arange(12).reshape((2, 3, 2)) + single_vec = np.arange(2) + col_vec = single_vec[:, np.newaxis] + col_vec_array = np.arange(8).reshape((2, 2, 2, 1)) + 1 + # matrix @ single column vector with proper dimension + mm_col_vec = umt.matrix_multiply(mat, col_vec) + # matmul does the same thing + matmul_col_vec = umt.matmul(mat, col_vec) + assert_array_equal(matmul_col_vec, mm_col_vec) + # matrix @ vector without dimension making it a column vector. + # matrix multiply fails -> missing core dim. + assert_raises(ValueError, umt.matrix_multiply, mat, single_vec) + # matmul mimicker passes, and returns a vector. + matmul_col = umt.matmul(mat, single_vec) + assert_array_equal(matmul_col, mm_col_vec.squeeze()) + # Now with a column array: same as for column vector, + # broadcasting sensibly. + mm_col_vec = umt.matrix_multiply(mat, col_vec_array) + matmul_col_vec = umt.matmul(mat, col_vec_array) + assert_array_equal(matmul_col_vec, mm_col_vec) + # As above, but for row vector + single_vec = np.arange(3) + row_vec = single_vec[np.newaxis, :] + row_vec_array = np.arange(24).reshape((4, 2, 1, 1, 3)) + 1 + # row vector @ matrix + mm_row_vec = umt.matrix_multiply(row_vec, mat) + matmul_row_vec = umt.matmul(row_vec, mat) + assert_array_equal(matmul_row_vec, mm_row_vec) + # single row vector @ matrix + assert_raises(ValueError, umt.matrix_multiply, single_vec, mat) + matmul_row = umt.matmul(single_vec, mat) + assert_array_equal(matmul_row, mm_row_vec.squeeze()) + # row vector array @ matrix + mm_row_vec = umt.matrix_multiply(row_vec_array, mat) + matmul_row_vec = umt.matmul(row_vec_array, mat) + assert_array_equal(matmul_row_vec, mm_row_vec) + # Now for vector combinations + # row vector @ column vector + col_vec = row_vec.T + col_vec_array = row_vec_array.swapaxes(-2, -1) + mm_row_col_vec = umt.matrix_multiply(row_vec, col_vec) + matmul_row_col_vec = umt.matmul(row_vec, col_vec) + assert_array_equal(matmul_row_col_vec, mm_row_col_vec) + # single row vector @ single col vector + assert_raises(ValueError, umt.matrix_multiply, single_vec, single_vec) + matmul_row_col = umt.matmul(single_vec, single_vec) + assert_array_equal(matmul_row_col, mm_row_col_vec.squeeze()) + # row vector array @ matrix + mm_row_col_array = umt.matrix_multiply(row_vec_array, col_vec_array) + matmul_row_col_array = umt.matmul(row_vec_array, col_vec_array) + assert_array_equal(matmul_row_col_array, mm_row_col_array) + # Finally, check that things are *not* squeezed if one gives an + # output. + out = np.zeros_like(mm_row_col_array) + out = umt.matrix_multiply(row_vec_array, col_vec_array, out=out) + assert_array_equal(out, mm_row_col_array) + out[:] = 0 + out = umt.matmul(row_vec_array, col_vec_array, out=out) + assert_array_equal(out, mm_row_col_array) + # And check one cannot put missing dimensions back. + out = np.zeros_like(mm_row_col_vec) + assert_raises(ValueError, umt.matrix_multiply, single_vec, single_vec, + out) + # But fine for matmul, since it is just a broadcast. + out = umt.matmul(single_vec, single_vec, out) + assert_array_equal(out, mm_row_col_vec.squeeze()) + + def test_matrix_multiply(self): + self.compare_matrix_multiply_results(np.int64) + self.compare_matrix_multiply_results(np.double) + + def test_matrix_multiply_umath_empty(self): + res = umt.matrix_multiply(np.ones((0, 10)), np.ones((10, 0))) + assert_array_equal(res, np.zeros((0, 0))) + res = umt.matrix_multiply(np.ones((10, 0)), np.ones((0, 10))) + assert_array_equal(res, np.zeros((10, 10))) + + def compare_matrix_multiply_results(self, tp): + d1 = np.array(np.random.rand(2, 3, 4), dtype=tp) + d2 = np.array(np.random.rand(2, 3, 4), dtype=tp) + msg = "matrix multiply on type %s" % d1.dtype.name + + def permute_n(n): + if n == 1: + return ([0],) + ret = () + base = permute_n(n-1) + for perm in base: + for i in range(n): + new = perm + [n-1] + new[n-1] = new[i] + new[i] = n-1 + ret += (new,) + return ret + + def slice_n(n): + if n == 0: + return ((),) + ret = () + base = slice_n(n-1) + for sl in base: + ret += (sl+(slice(None),),) + ret += (sl+(slice(0, 1),),) + return ret + + def broadcastable(s1, s2): + return s1 == s2 or s1 == 1 or s2 == 1 + + permute_3 = permute_n(3) + slice_3 = slice_n(3) + ((slice(None, None, -1),)*3,) + + ref = True + for p1 in permute_3: + for p2 in permute_3: + for s1 in slice_3: + for s2 in slice_3: + a1 = d1.transpose(p1)[s1] + a2 = d2.transpose(p2)[s2] + ref = ref and a1.base is not None + ref = ref and a2.base is not None + if (a1.shape[-1] == a2.shape[-2] and + broadcastable(a1.shape[0], a2.shape[0])): + assert_array_almost_equal( + umt.matrix_multiply(a1, a2), + np.sum(a2[..., np.newaxis].swapaxes(-3, -1) * + a1[..., np.newaxis,:], axis=-1), + err_msg=msg + ' %s %s' % (str(a1.shape), + str(a2.shape))) + + assert_equal(ref, True, err_msg="reference check") + + def test_euclidean_pdist(self): + a = np.arange(12, dtype=float).reshape(4, 3) + out = np.empty((a.shape[0] * (a.shape[0] - 1) // 2,), dtype=a.dtype) + umt.euclidean_pdist(a, out) + b = np.sqrt(np.sum((a[:, None] - a)**2, axis=-1)) + b = b[~np.tri(a.shape[0], dtype=bool)] + assert_almost_equal(out, b) + # An output array is required to determine p with signature (n,d)->(p) + assert_raises(ValueError, umt.euclidean_pdist, a) + + def test_cumsum(self): + a = np.arange(10) + result = umt.cumsum(a) + assert_array_equal(result, a.cumsum()) + + def test_object_logical(self): + a = np.array([3, None, True, False, "test", ""], dtype=object) + assert_equal(np.logical_or(a, None), + np.array([x or None for x in a], dtype=object)) + assert_equal(np.logical_or(a, True), + np.array([x or True for x in a], dtype=object)) + assert_equal(np.logical_or(a, 12), + np.array([x or 12 for x in a], dtype=object)) + assert_equal(np.logical_or(a, "blah"), + np.array([x or "blah" for x in a], dtype=object)) + + assert_equal(np.logical_and(a, None), + np.array([x and None for x in a], dtype=object)) + assert_equal(np.logical_and(a, True), + np.array([x and True for x in a], dtype=object)) + assert_equal(np.logical_and(a, 12), + np.array([x and 12 for x in a], dtype=object)) + assert_equal(np.logical_and(a, "blah"), + np.array([x and "blah" for x in a], dtype=object)) + + assert_equal(np.logical_not(a), + np.array([not x for x in a], dtype=object)) + + assert_equal(np.logical_or.reduce(a), 3) + assert_equal(np.logical_and.reduce(a), None) + + def test_object_comparison(self): + class HasComparisons: + def __eq__(self, other): + return '==' + + arr0d = np.array(HasComparisons()) + assert_equal(arr0d == arr0d, True) + assert_equal(np.equal(arr0d, arr0d), True) # normal behavior is a cast + + arr1d = np.array([HasComparisons()]) + assert_equal(arr1d == arr1d, np.array([True])) + assert_equal(np.equal(arr1d, arr1d), np.array([True])) # normal behavior is a cast + assert_equal(np.equal(arr1d, arr1d, dtype=object), np.array(['=='])) + + def test_object_array_reduction(self): + # Reductions on object arrays + a = np.array(['a', 'b', 'c'], dtype=object) + assert_equal(np.sum(a), 'abc') + assert_equal(np.max(a), 'c') + assert_equal(np.min(a), 'a') + a = np.array([True, False, True], dtype=object) + assert_equal(np.sum(a), 2) + assert_equal(np.prod(a), 0) + assert_equal(np.any(a), True) + assert_equal(np.all(a), False) + assert_equal(np.max(a), True) + assert_equal(np.min(a), False) + assert_equal(np.array([[1]], dtype=object).sum(), 1) + assert_equal(np.array([[[1, 2]]], dtype=object).sum((0, 1)), [1, 2]) + assert_equal(np.array([1], dtype=object).sum(initial=1), 2) + assert_equal(np.array([[1], [2, 3]], dtype=object) + .sum(initial=[0], where=[False, True]), [0, 2, 3]) + + def test_object_array_accumulate_inplace(self): + # Checks that in-place accumulates work, see also gh-7402 + arr = np.ones(4, dtype=object) + arr[:] = [[1] for i in range(4)] + # Twice reproduced also for tuples: + np.add.accumulate(arr, out=arr) + np.add.accumulate(arr, out=arr) + assert_array_equal(arr, + np.array([[1]*i for i in [1, 3, 6, 10]], dtype=object), + ) + + # And the same if the axis argument is used + arr = np.ones((2, 4), dtype=object) + arr[0, :] = [[2] for i in range(4)] + np.add.accumulate(arr, out=arr, axis=-1) + np.add.accumulate(arr, out=arr, axis=-1) + assert_array_equal(arr[0, :], + np.array([[2]*i for i in [1, 3, 6, 10]], dtype=object), + ) + + def test_object_array_accumulate_failure(self): + # Typical accumulation on object works as expected: + res = np.add.accumulate(np.array([1, 0, 2], dtype=object)) + assert_array_equal(res, np.array([1, 1, 3], dtype=object)) + # But errors are propagated from the inner-loop if they occur: + with pytest.raises(TypeError): + np.add.accumulate([1, None, 2]) + + def test_object_array_reduceat_inplace(self): + # Checks that in-place reduceats work, see also gh-7465 + arr = np.empty(4, dtype=object) + arr[:] = [[1] for i in range(4)] + out = np.empty(4, dtype=object) + out[:] = [[1] for i in range(4)] + np.add.reduceat(arr, np.arange(4), out=arr) + np.add.reduceat(arr, np.arange(4), out=arr) + assert_array_equal(arr, out) + + # And the same if the axis argument is used + arr = np.ones((2, 4), dtype=object) + arr[0, :] = [[2] for i in range(4)] + out = np.ones((2, 4), dtype=object) + out[0, :] = [[2] for i in range(4)] + np.add.reduceat(arr, np.arange(4), out=arr, axis=-1) + np.add.reduceat(arr, np.arange(4), out=arr, axis=-1) + assert_array_equal(arr, out) + + def test_object_array_reduceat_failure(self): + # Reduceat works as expected when no invalid operation occurs (None is + # not involved in an operation here) + res = np.add.reduceat(np.array([1, None, 2], dtype=object), [1, 2]) + assert_array_equal(res, np.array([None, 2], dtype=object)) + # But errors when None would be involved in an operation: + with pytest.raises(TypeError): + np.add.reduceat([1, None, 2], [0, 2]) + + def test_zerosize_reduction(self): + # Test with default dtype and object dtype + for a in [[], np.array([], dtype=object)]: + assert_equal(np.sum(a), 0) + assert_equal(np.prod(a), 1) + assert_equal(np.any(a), False) + assert_equal(np.all(a), True) + assert_raises(ValueError, np.max, a) + assert_raises(ValueError, np.min, a) + + def test_axis_out_of_bounds(self): + a = np.array([False, False]) + assert_raises(AxisError, a.all, axis=1) + a = np.array([False, False]) + assert_raises(AxisError, a.all, axis=-2) + + a = np.array([False, False]) + assert_raises(AxisError, a.any, axis=1) + a = np.array([False, False]) + assert_raises(AxisError, a.any, axis=-2) + + def test_scalar_reduction(self): + # The functions 'sum', 'prod', etc allow specifying axis=0 + # even for scalars + assert_equal(np.sum(3, axis=0), 3) + assert_equal(np.prod(3.5, axis=0), 3.5) + assert_equal(np.any(True, axis=0), True) + assert_equal(np.all(False, axis=0), False) + assert_equal(np.max(3, axis=0), 3) + assert_equal(np.min(2.5, axis=0), 2.5) + + # Check scalar behaviour for ufuncs without an identity + assert_equal(np.power.reduce(3), 3) + + # Make sure that scalars are coming out from this operation + assert_(type(np.prod(np.float32(2.5), axis=0)) is np.float32) + assert_(type(np.sum(np.float32(2.5), axis=0)) is np.float32) + assert_(type(np.max(np.float32(2.5), axis=0)) is np.float32) + assert_(type(np.min(np.float32(2.5), axis=0)) is np.float32) + + # check if scalars/0-d arrays get cast + assert_(type(np.any(0, axis=0)) is np.bool) + + # assert that 0-d arrays get wrapped + class MyArray(np.ndarray): + pass + a = np.array(1).view(MyArray) + assert_(type(np.any(a)) is MyArray) + + def test_casting_out_param(self): + # Test that it's possible to do casts on output + a = np.ones((200, 100), np.int64) + b = np.ones((200, 100), np.int64) + c = np.ones((200, 100), np.float64) + np.add(a, b, out=c) + assert_equal(c, 2) + + a = np.zeros(65536) + b = np.zeros(65536, dtype=np.float32) + np.subtract(a, 0, out=b) + assert_equal(b, 0) + + def test_where_param(self): + # Test that the where= ufunc parameter works with regular arrays + a = np.arange(7) + b = np.ones(7) + c = np.zeros(7) + np.add(a, b, out=c, where=(a % 2 == 1)) + assert_equal(c, [0, 2, 0, 4, 0, 6, 0]) + + a = np.arange(4).reshape(2, 2) + 2 + np.power(a, [2, 3], out=a, where=[[0, 1], [1, 0]]) + assert_equal(a, [[2, 27], [16, 5]]) + # Broadcasting the where= parameter + np.subtract(a, 2, out=a, where=[True, False]) + assert_equal(a, [[0, 27], [14, 5]]) + + def test_where_param_buffer_output(self): + # This test is temporarily skipped because it requires + # adding masking features to the nditer to work properly + + # With casting on output + a = np.ones(10, np.int64) + b = np.ones(10, np.int64) + c = 1.5 * np.ones(10, np.float64) + np.add(a, b, out=c, where=[1, 0, 0, 1, 0, 0, 1, 1, 1, 0]) + assert_equal(c, [2, 1.5, 1.5, 2, 1.5, 1.5, 2, 2, 2, 1.5]) + + def test_where_param_alloc(self): + # With casting and allocated output + a = np.array([1], dtype=np.int64) + m = np.array([True], dtype=bool) + assert_equal(np.sqrt(a, where=m), [1]) + + # No casting and allocated output + a = np.array([1], dtype=np.float64) + m = np.array([True], dtype=bool) + assert_equal(np.sqrt(a, where=m), [1]) + + def test_where_with_broadcasting(self): + # See gh-17198 + a = np.random.random((5000, 4)) + b = np.random.random((5000, 1)) + + where = a > 0.3 + out = np.full_like(a, 0) + np.less(a, b, where=where, out=out) + b_where = np.broadcast_to(b, a.shape)[where] + assert_array_equal((a[where] < b_where), out[where].astype(bool)) + assert not out[~where].any() # outside mask, out remains all 0 + + @staticmethod + def identityless_reduce_arrs(): + yield np.empty((2, 3, 4), order='C') + yield np.empty((2, 3, 4), order='F') + # Mixed order (reduce order differs outer) + yield np.empty((2, 4, 3), order='C').swapaxes(1, 2) + # Reversed order + yield np.empty((2, 3, 4), order='C')[::-1, ::-1, ::-1] + # Not contiguous + yield np.empty((3, 5, 4), order='C').swapaxes(1, 2)[1:, 1:, 1:] + # Not contiguous and not aligned + a = np.empty((3*4*5*8 + 1,), dtype='i1') + a = a[1:].view(dtype='f8') + a.shape = (3, 4, 5) + a = a[1:, 1:, 1:] + yield a + + @pytest.mark.parametrize("a", identityless_reduce_arrs()) + @pytest.mark.parametrize("pos", [(1, 0, 0), (0, 1, 0), (0, 0, 1)]) + def test_identityless_reduction(self, a, pos): + # np.minimum.reduce is an identityless reduction + a[...] = 1 + a[pos] = 0 + + for axis in [None, (0, 1), (0, 2), (1, 2), 0, 1, 2, ()]: + if axis is None: + axes = np.array([], dtype=np.intp) + else: + axes = np.delete(np.arange(a.ndim), axis) + + expected_pos = tuple(np.array(pos)[axes]) + expected = np.ones(np.array(a.shape)[axes]) + expected[expected_pos] = 0 + + res = np.minimum.reduce(a, axis=axis) + assert_equal(res, expected, strict=True) + + res = np.full_like(res, np.nan) + np.minimum.reduce(a, axis=axis, out=res) + assert_equal(res, expected, strict=True) + + @requires_memory(6 * 1024**3) + @pytest.mark.skipif(sys.maxsize < 2**32, + reason="test array too large for 32bit platform") + def test_identityless_reduction_huge_array(self): + # Regression test for gh-20921 (copying identity incorrectly failed) + arr = np.zeros((2, 2**31), 'uint8') + arr[:, 0] = [1, 3] + arr[:, -1] = [4, 1] + res = np.maximum.reduce(arr, axis=0) + del arr + assert res[0] == 3 + assert res[-1] == 4 + + def test_reduce_identity_depends_on_loop(self): + """ + The type of the result should always depend on the selected loop, not + necessarily the output (only relevant for object arrays). + """ + # For an object loop, the default value 0 with type int is used: + assert type(np.add.reduce([], dtype=object)) is int + out = np.array(None, dtype=object) + # When the loop is float64 but `out` is object this does not happen, + # the result is float64 cast to object (which gives Python `float`). + np.add.reduce([], out=out, dtype=np.float64) + assert type(out[()]) is float + + def test_initial_reduction(self): + # np.minimum.reduce is an identityless reduction + + # For cases like np.maximum(np.abs(...), initial=0) + # More generally, a supremum over non-negative numbers. + assert_equal(np.maximum.reduce([], initial=0), 0) + + # For cases like reduction of an empty array over the reals. + assert_equal(np.minimum.reduce([], initial=np.inf), np.inf) + assert_equal(np.maximum.reduce([], initial=-np.inf), -np.inf) + + # Random tests + assert_equal(np.minimum.reduce([5], initial=4), 4) + assert_equal(np.maximum.reduce([4], initial=5), 5) + assert_equal(np.maximum.reduce([5], initial=4), 5) + assert_equal(np.minimum.reduce([4], initial=5), 4) + + # Check initial=None raises ValueError for both types of ufunc reductions + assert_raises(ValueError, np.minimum.reduce, [], initial=None) + assert_raises(ValueError, np.add.reduce, [], initial=None) + # Also in the somewhat special object case: + with pytest.raises(ValueError): + np.add.reduce([], initial=None, dtype=object) + + # Check that np._NoValue gives default behavior. + assert_equal(np.add.reduce([], initial=np._NoValue), 0) + + # Check that initial kwarg behaves as intended for dtype=object + a = np.array([10], dtype=object) + res = np.add.reduce(a, initial=5) + assert_equal(res, 15) + + def test_empty_reduction_and_identity(self): + arr = np.zeros((0, 5)) + # OK, since the reduction itself is *not* empty, the result is + assert np.true_divide.reduce(arr, axis=1).shape == (0,) + # Not OK, the reduction itself is empty and we have no identity + with pytest.raises(ValueError): + np.true_divide.reduce(arr, axis=0) + + # Test that an empty reduction fails also if the result is empty + arr = np.zeros((0, 0, 5)) + with pytest.raises(ValueError): + np.true_divide.reduce(arr, axis=1) + + # Division reduction makes sense with `initial=1` (empty or not): + res = np.true_divide.reduce(arr, axis=1, initial=1) + assert_array_equal(res, np.ones((0, 5))) + + @pytest.mark.parametrize('axis', (0, 1, None)) + @pytest.mark.parametrize('where', (np.array([False, True, True]), + np.array([[True], [False], [True]]), + np.array([[True, False, False], + [False, True, False], + [False, True, True]]))) + def test_reduction_with_where(self, axis, where): + a = np.arange(9.).reshape(3, 3) + a_copy = a.copy() + a_check = np.zeros_like(a) + np.positive(a, out=a_check, where=where) + + res = np.add.reduce(a, axis=axis, where=where) + check = a_check.sum(axis) + assert_equal(res, check) + # Check we do not overwrite elements of a internally. + assert_array_equal(a, a_copy) + + @pytest.mark.parametrize(('axis', 'where'), + ((0, np.array([True, False, True])), + (1, [True, True, False]), + (None, True))) + @pytest.mark.parametrize('initial', (-np.inf, 5.)) + def test_reduction_with_where_and_initial(self, axis, where, initial): + a = np.arange(9.).reshape(3, 3) + a_copy = a.copy() + a_check = np.full(a.shape, -np.inf) + np.positive(a, out=a_check, where=where) + + res = np.maximum.reduce(a, axis=axis, where=where, initial=initial) + check = a_check.max(axis, initial=initial) + assert_equal(res, check) + + def test_reduction_where_initial_needed(self): + a = np.arange(9.).reshape(3, 3) + m = [False, True, False] + assert_raises(ValueError, np.maximum.reduce, a, where=m) + + def test_identityless_reduction_nonreorderable(self): + a = np.array([[8.0, 2.0, 2.0], [1.0, 0.5, 0.25]]) + + res = np.divide.reduce(a, axis=0) + assert_equal(res, [8.0, 4.0, 8.0]) + + res = np.divide.reduce(a, axis=1) + assert_equal(res, [2.0, 8.0]) + + res = np.divide.reduce(a, axis=()) + assert_equal(res, a) + + assert_raises(ValueError, np.divide.reduce, a, axis=(0, 1)) + + def test_reduce_zero_axis(self): + # If we have a n x m array and do a reduction with axis=1, then we are + # doing n reductions, and each reduction takes an m-element array. For + # a reduction operation without an identity, then: + # n > 0, m > 0: fine + # n = 0, m > 0: fine, doing 0 reductions of m-element arrays + # n > 0, m = 0: can't reduce a 0-element array, ValueError + # n = 0, m = 0: can't reduce a 0-element array, ValueError (for + # consistency with the above case) + # This test doesn't actually look at return values, it just checks to + # make sure that error we get an error in exactly those cases where we + # expect one, and assumes the calculations themselves are done + # correctly. + + def ok(f, *args, **kwargs): + f(*args, **kwargs) + + def err(f, *args, **kwargs): + assert_raises(ValueError, f, *args, **kwargs) + + def t(expect, func, n, m): + expect(func, np.zeros((n, m)), axis=1) + expect(func, np.zeros((m, n)), axis=0) + expect(func, np.zeros((n // 2, n // 2, m)), axis=2) + expect(func, np.zeros((n // 2, m, n // 2)), axis=1) + expect(func, np.zeros((n, m // 2, m // 2)), axis=(1, 2)) + expect(func, np.zeros((m // 2, n, m // 2)), axis=(0, 2)) + expect(func, np.zeros((m // 3, m // 3, m // 3, + n // 2, n // 2)), + axis=(0, 1, 2)) + # Check what happens if the inner (resp. outer) dimensions are a + # mix of zero and non-zero: + expect(func, np.zeros((10, m, n)), axis=(0, 1)) + expect(func, np.zeros((10, n, m)), axis=(0, 2)) + expect(func, np.zeros((m, 10, n)), axis=0) + expect(func, np.zeros((10, m, n)), axis=1) + expect(func, np.zeros((10, n, m)), axis=2) + + # np.maximum is just an arbitrary ufunc with no reduction identity + assert_equal(np.maximum.identity, None) + t(ok, np.maximum.reduce, 30, 30) + t(ok, np.maximum.reduce, 0, 30) + t(err, np.maximum.reduce, 30, 0) + t(err, np.maximum.reduce, 0, 0) + err(np.maximum.reduce, []) + np.maximum.reduce(np.zeros((0, 0)), axis=()) + + # all of the combinations are fine for a reduction that has an + # identity + t(ok, np.add.reduce, 30, 30) + t(ok, np.add.reduce, 0, 30) + t(ok, np.add.reduce, 30, 0) + t(ok, np.add.reduce, 0, 0) + np.add.reduce([]) + np.add.reduce(np.zeros((0, 0)), axis=()) + + # OTOH, accumulate always makes sense for any combination of n and m, + # because it maps an m-element array to an m-element array. These + # tests are simpler because accumulate doesn't accept multiple axes. + for uf in (np.maximum, np.add): + uf.accumulate(np.zeros((30, 0)), axis=0) + uf.accumulate(np.zeros((0, 30)), axis=0) + uf.accumulate(np.zeros((30, 30)), axis=0) + uf.accumulate(np.zeros((0, 0)), axis=0) + + def test_safe_casting(self): + # In old versions of numpy, in-place operations used the 'unsafe' + # casting rules. In versions >= 1.10, 'same_kind' is the + # default and an exception is raised instead of a warning. + # when 'same_kind' is not satisfied. + a = np.array([1, 2, 3], dtype=int) + # Non-in-place addition is fine + assert_array_equal(assert_no_warnings(np.add, a, 1.1), + [2.1, 3.1, 4.1]) + assert_raises(TypeError, np.add, a, 1.1, out=a) + + def add_inplace(a, b): + a += b + + assert_raises(TypeError, add_inplace, a, 1.1) + # Make sure that explicitly overriding the exception is allowed: + assert_no_warnings(np.add, a, 1.1, out=a, casting="unsafe") + assert_array_equal(a, [2, 3, 4]) + + def test_ufunc_custom_out(self): + # Test ufunc with built in input types and custom output type + + a = np.array([0, 1, 2], dtype='i8') + b = np.array([0, 1, 2], dtype='i8') + c = np.empty(3, dtype=_rational_tests.rational) + + # Output must be specified so numpy knows what + # ufunc signature to look for + result = _rational_tests.test_add(a, b, c) + target = np.array([0, 2, 4], dtype=_rational_tests.rational) + assert_equal(result, target) + + # The new resolution means that we can (usually) find custom loops + # as long as they match exactly: + result = _rational_tests.test_add(a, b) + assert_equal(result, target) + + # This works even more generally, so long the default common-dtype + # promoter works out: + result = _rational_tests.test_add(a, b.astype(np.uint16), out=c) + assert_equal(result, target) + + # This scalar path used to go into legacy promotion, but doesn't now: + result = _rational_tests.test_add(a, np.uint16(2)) + target = np.array([2, 3, 4], dtype=_rational_tests.rational) + assert_equal(result, target) + + def test_operand_flags(self): + a = np.arange(16, dtype=int).reshape(4, 4) + b = np.arange(9, dtype=int).reshape(3, 3) + opflag_tests.inplace_add(a[:-1, :-1], b) + assert_equal(a, np.array([[0, 2, 4, 3], [7, 9, 11, 7], + [14, 16, 18, 11], [12, 13, 14, 15]])) + + a = np.array(0) + opflag_tests.inplace_add(a, 3) + assert_equal(a, 3) + opflag_tests.inplace_add(a, [3, 4]) + assert_equal(a, 10) + + def test_struct_ufunc(self): + import numpy._core._struct_ufunc_tests as struct_ufunc + + a = np.array([(1, 2, 3)], dtype='u8,u8,u8') + b = np.array([(1, 2, 3)], dtype='u8,u8,u8') + + result = struct_ufunc.add_triplet(a, b) + assert_equal(result, np.array([(2, 4, 6)], dtype='u8,u8,u8')) + assert_raises(RuntimeError, struct_ufunc.register_fail) + + def test_custom_ufunc(self): + a = np.array( + [_rational_tests.rational(1, 2), + _rational_tests.rational(1, 3), + _rational_tests.rational(1, 4)], + dtype=_rational_tests.rational) + b = np.array( + [_rational_tests.rational(1, 2), + _rational_tests.rational(1, 3), + _rational_tests.rational(1, 4)], + dtype=_rational_tests.rational) + + result = _rational_tests.test_add_rationals(a, b) + expected = np.array( + [_rational_tests.rational(1), + _rational_tests.rational(2, 3), + _rational_tests.rational(1, 2)], + dtype=_rational_tests.rational) + assert_equal(result, expected) + + def test_custom_ufunc_forced_sig(self): + # gh-9351 - looking for a non-first userloop would previously hang + with assert_raises(TypeError): + np.multiply(_rational_tests.rational(1), 1, + signature=(_rational_tests.rational, int, None)) + + def test_custom_array_like(self): + + class MyThing: + __array_priority__ = 1000 + + rmul_count = 0 + getitem_count = 0 + + def __init__(self, shape): + self.shape = shape + + def __len__(self): + return self.shape[0] + + def __getitem__(self, i): + MyThing.getitem_count += 1 + if not isinstance(i, tuple): + i = (i,) + if len(i) > self.ndim: + raise IndexError("boo") + + return MyThing(self.shape[len(i):]) + + def __rmul__(self, other): + MyThing.rmul_count += 1 + return self + + np.float64(5)*MyThing((3, 3)) + assert_(MyThing.rmul_count == 1, MyThing.rmul_count) + assert_(MyThing.getitem_count <= 2, MyThing.getitem_count) + + @pytest.mark.parametrize("a", ( + np.arange(10, dtype=int), + np.arange(10, dtype=_rational_tests.rational), + )) + def test_ufunc_at_basic(self, a): + + aa = a.copy() + np.add.at(aa, [2, 5, 2], 1) + assert_equal(aa, [0, 1, 4, 3, 4, 6, 6, 7, 8, 9]) + + with pytest.raises(ValueError): + # missing second operand + np.add.at(aa, [2, 5, 3]) + + aa = a.copy() + np.negative.at(aa, [2, 5, 3]) + assert_equal(aa, [0, 1, -2, -3, 4, -5, 6, 7, 8, 9]) + + aa = a.copy() + b = np.array([100, 100, 100]) + np.add.at(aa, [2, 5, 2], b) + assert_equal(aa, [0, 1, 202, 3, 4, 105, 6, 7, 8, 9]) + + with pytest.raises(ValueError): + # extraneous second operand + np.negative.at(a, [2, 5, 3], [1, 2, 3]) + + with pytest.raises(ValueError): + # second operand cannot be converted to an array + np.add.at(a, [2, 5, 3], [[1, 2], 1]) + + # ufuncs with indexed loops for performance in ufunc.at + indexed_ufuncs = [np.add, np.subtract, np.multiply, np.floor_divide, + np.maximum, np.minimum, np.fmax, np.fmin] + + @pytest.mark.parametrize( + "typecode", np.typecodes['AllInteger'] + np.typecodes['Float']) + @pytest.mark.parametrize("ufunc", indexed_ufuncs) + def test_ufunc_at_inner_loops(self, typecode, ufunc): + if ufunc is np.divide and typecode in np.typecodes['AllInteger']: + # Avoid divide-by-zero and inf for integer divide + a = np.ones(100, dtype=typecode) + indx = np.random.randint(100, size=30, dtype=np.intp) + vals = np.arange(1, 31, dtype=typecode) + else: + a = np.ones(1000, dtype=typecode) + indx = np.random.randint(1000, size=3000, dtype=np.intp) + vals = np.arange(3000, dtype=typecode) + atag = a.copy() + # Do the calculation twice and compare the answers + with warnings.catch_warnings(record=True) as w_at: + warnings.simplefilter('always') + ufunc.at(a, indx, vals) + with warnings.catch_warnings(record=True) as w_loop: + warnings.simplefilter('always') + for i, v in zip(indx, vals): + # Make sure all the work happens inside the ufunc + # in order to duplicate error/warning handling + ufunc(atag[i], v, out=atag[i:i+1], casting="unsafe") + assert_equal(atag, a) + # If w_loop warned, make sure w_at warned as well + if len(w_loop) > 0: + # + assert len(w_at) > 0 + assert w_at[0].category == w_loop[0].category + assert str(w_at[0].message)[:10] == str(w_loop[0].message)[:10] + + @pytest.mark.parametrize("typecode", np.typecodes['Complex']) + @pytest.mark.parametrize("ufunc", [np.add, np.subtract, np.multiply]) + def test_ufunc_at_inner_loops_complex(self, typecode, ufunc): + a = np.ones(10, dtype=typecode) + indx = np.concatenate([np.ones(6, dtype=np.intp), + np.full(18, 4, dtype=np.intp)]) + value = a.dtype.type(1j) + ufunc.at(a, indx, value) + expected = np.ones_like(a) + if ufunc is np.multiply: + expected[1] = expected[4] = -1 + else: + expected[1] += 6 * (value if ufunc is np.add else -value) + expected[4] += 18 * (value if ufunc is np.add else -value) + + assert_array_equal(a, expected) + + def test_ufunc_at_ellipsis(self): + # Make sure the indexed loop check does not choke on iters + # with subspaces + arr = np.zeros(5) + np.add.at(arr, slice(None), np.ones(5)) + assert_array_equal(arr, np.ones(5)) + + def test_ufunc_at_negative(self): + arr = np.ones(5, dtype=np.int32) + indx = np.arange(5) + umt.indexed_negative.at(arr, indx) + # If it is [-1, -1, -1, -100, 0] then the regular strided loop was used + assert np.all(arr == [-1, -1, -1, -200, -1]) + + def test_ufunc_at_large(self): + # issue gh-23457 + indices = np.zeros(8195, dtype=np.int16) + b = np.zeros(8195, dtype=float) + b[0] = 10 + b[1] = 5 + b[8192:] = 100 + a = np.zeros(1, dtype=float) + np.add.at(a, indices, b) + assert a[0] == b.sum() + + def test_cast_index_fastpath(self): + arr = np.zeros(10) + values = np.ones(100000) + # index must be cast, which may be buffered in chunks: + index = np.zeros(len(values), dtype=np.uint8) + np.add.at(arr, index, values) + assert arr[0] == len(values) + + @pytest.mark.parametrize("value", [ + np.ones(1), np.ones(()), np.float64(1.), 1.]) + def test_ufunc_at_scalar_value_fastpath(self, value): + arr = np.zeros(1000) + # index must be cast, which may be buffered in chunks: + index = np.repeat(np.arange(1000), 2) + np.add.at(arr, index, value) + assert_array_equal(arr, np.full_like(arr, 2 * value)) + + def test_ufunc_at_multiD(self): + a = np.arange(9).reshape(3, 3) + b = np.array([[100, 100, 100], [200, 200, 200], [300, 300, 300]]) + np.add.at(a, (slice(None), [1, 2, 1]), b) + assert_equal(a, [[0, 201, 102], [3, 404, 205], [6, 607, 308]]) + + a = np.arange(27).reshape(3, 3, 3) + b = np.array([100, 200, 300]) + np.add.at(a, (slice(None), slice(None), [1, 2, 1]), b) + assert_equal(a, + [[[0, 401, 202], + [3, 404, 205], + [6, 407, 208]], + + [[9, 410, 211], + [12, 413, 214], + [15, 416, 217]], + + [[18, 419, 220], + [21, 422, 223], + [24, 425, 226]]]) + + a = np.arange(9).reshape(3, 3) + b = np.array([[100, 100, 100], [200, 200, 200], [300, 300, 300]]) + np.add.at(a, ([1, 2, 1], slice(None)), b) + assert_equal(a, [[0, 1, 2], [403, 404, 405], [206, 207, 208]]) + + a = np.arange(27).reshape(3, 3, 3) + b = np.array([100, 200, 300]) + np.add.at(a, (slice(None), [1, 2, 1], slice(None)), b) + assert_equal(a, + [[[0, 1, 2], + [203, 404, 605], + [106, 207, 308]], + + [[9, 10, 11], + [212, 413, 614], + [115, 216, 317]], + + [[18, 19, 20], + [221, 422, 623], + [124, 225, 326]]]) + + a = np.arange(9).reshape(3, 3) + b = np.array([100, 200, 300]) + np.add.at(a, (0, [1, 2, 1]), b) + assert_equal(a, [[0, 401, 202], [3, 4, 5], [6, 7, 8]]) + + a = np.arange(27).reshape(3, 3, 3) + b = np.array([100, 200, 300]) + np.add.at(a, ([1, 2, 1], 0, slice(None)), b) + assert_equal(a, + [[[0, 1, 2], + [3, 4, 5], + [6, 7, 8]], + + [[209, 410, 611], + [12, 13, 14], + [15, 16, 17]], + + [[118, 219, 320], + [21, 22, 23], + [24, 25, 26]]]) + + a = np.arange(27).reshape(3, 3, 3) + b = np.array([100, 200, 300]) + np.add.at(a, (slice(None), slice(None), slice(None)), b) + assert_equal(a, + [[[100, 201, 302], + [103, 204, 305], + [106, 207, 308]], + + [[109, 210, 311], + [112, 213, 314], + [115, 216, 317]], + + [[118, 219, 320], + [121, 222, 323], + [124, 225, 326]]]) + + def test_ufunc_at_0D(self): + a = np.array(0) + np.add.at(a, (), 1) + assert_equal(a, 1) + + assert_raises(IndexError, np.add.at, a, 0, 1) + assert_raises(IndexError, np.add.at, a, [], 1) + + def test_ufunc_at_dtypes(self): + # Test mixed dtypes + a = np.arange(10) + np.power.at(a, [1, 2, 3, 2], 3.5) + assert_equal(a, np.array([0, 1, 4414, 46, 4, 5, 6, 7, 8, 9])) + + def test_ufunc_at_boolean(self): + # Test boolean indexing and boolean ufuncs + a = np.arange(10) + index = a % 2 == 0 + np.equal.at(a, index, [0, 2, 4, 6, 8]) + assert_equal(a, [1, 1, 1, 3, 1, 5, 1, 7, 1, 9]) + + # Test unary operator + a = np.arange(10, dtype='u4') + np.invert.at(a, [2, 5, 2]) + assert_equal(a, [0, 1, 2, 3, 4, 5 ^ 0xffffffff, 6, 7, 8, 9]) + + def test_ufunc_at_advanced(self): + # Test empty subspace + orig = np.arange(4) + a = orig[:, None][:, 0:0] + np.add.at(a, [0, 1], 3) + assert_array_equal(orig, np.arange(4)) + + # Test with swapped byte order + index = np.array([1, 2, 1], np.dtype('i').newbyteorder()) + values = np.array([1, 2, 3, 4], np.dtype('f').newbyteorder()) + np.add.at(values, index, 3) + assert_array_equal(values, [1, 8, 6, 4]) + + # Test exception thrown + values = np.array(['a', 1], dtype=object) + assert_raises(TypeError, np.add.at, values, [0, 1], 1) + assert_array_equal(values, np.array(['a', 1], dtype=object)) + + # Test multiple output ufuncs raise error, gh-5665 + assert_raises(ValueError, np.modf.at, np.arange(10), [1]) + + # Test maximum + a = np.array([1, 2, 3]) + np.maximum.at(a, [0], 0) + assert_equal(a, np.array([1, 2, 3])) + + @pytest.mark.parametrize("dtype", + np.typecodes['AllInteger'] + np.typecodes['Float']) + @pytest.mark.parametrize("ufunc", + [np.add, np.subtract, np.divide, np.minimum, np.maximum]) + def test_at_negative_indexes(self, dtype, ufunc): + a = np.arange(0, 10).astype(dtype) + indxs = np.array([-1, 1, -1, 2]).astype(np.intp) + vals = np.array([1, 5, 2, 10], dtype=a.dtype) + + expected = a.copy() + for i, v in zip(indxs, vals): + expected[i] = ufunc(expected[i], v) + + ufunc.at(a, indxs, vals) + assert_array_equal(a, expected) + assert np.all(indxs == [-1, 1, -1, 2]) + + def test_at_not_none_signature(self): + # Test ufuncs with non-trivial signature raise a TypeError + a = np.ones((2, 2, 2)) + b = np.ones((1, 2, 2)) + assert_raises(TypeError, np.matmul.at, a, [0], b) + + a = np.array([[[1, 2], [3, 4]]]) + assert_raises(TypeError, np.linalg._umath_linalg.det.at, a, [0]) + + def test_at_no_loop_for_op(self): + # str dtype does not have a ufunc loop for np.add + arr = np.ones(10, dtype=str) + with pytest.raises(np._core._exceptions._UFuncNoLoopError): + np.add.at(arr, [0, 1], [0, 1]) + + def test_at_output_casting(self): + arr = np.array([-1]) + np.equal.at(arr, [0], [0]) + assert arr[0] == 0 + + def test_at_broadcast_failure(self): + arr = np.arange(5) + with pytest.raises(ValueError): + np.add.at(arr, [0, 1], [1, 2, 3]) + + + def test_reduce_arguments(self): + f = np.add.reduce + d = np.ones((5,2), dtype=int) + o = np.ones((2,), dtype=d.dtype) + r = o * 5 + assert_equal(f(d), r) + # a, axis=0, dtype=None, out=None, keepdims=False + assert_equal(f(d, axis=0), r) + assert_equal(f(d, 0), r) + assert_equal(f(d, 0, dtype=None), r) + assert_equal(f(d, 0, dtype='i'), r) + assert_equal(f(d, 0, 'i'), r) + assert_equal(f(d, 0, None), r) + assert_equal(f(d, 0, None, out=None), r) + assert_equal(f(d, 0, None, out=o), r) + assert_equal(f(d, 0, None, o), r) + assert_equal(f(d, 0, None, None), r) + assert_equal(f(d, 0, None, None, keepdims=False), r) + assert_equal(f(d, 0, None, None, True), r.reshape((1,) + r.shape)) + assert_equal(f(d, 0, None, None, False, 0), r) + assert_equal(f(d, 0, None, None, False, initial=0), r) + assert_equal(f(d, 0, None, None, False, 0, True), r) + assert_equal(f(d, 0, None, None, False, 0, where=True), r) + # multiple keywords + assert_equal(f(d, axis=0, dtype=None, out=None, keepdims=False), r) + assert_equal(f(d, 0, dtype=None, out=None, keepdims=False), r) + assert_equal(f(d, 0, None, out=None, keepdims=False), r) + assert_equal(f(d, 0, None, out=None, keepdims=False, initial=0, + where=True), r) + + # too little + assert_raises(TypeError, f) + # too much + assert_raises(TypeError, f, d, 0, None, None, False, 0, True, 1) + # invalid axis + assert_raises(TypeError, f, d, "invalid") + assert_raises(TypeError, f, d, axis="invalid") + assert_raises(TypeError, f, d, axis="invalid", dtype=None, + keepdims=True) + # invalid dtype + assert_raises(TypeError, f, d, 0, "invalid") + assert_raises(TypeError, f, d, dtype="invalid") + assert_raises(TypeError, f, d, dtype="invalid", out=None) + # invalid out + assert_raises(TypeError, f, d, 0, None, "invalid") + assert_raises(TypeError, f, d, out="invalid") + assert_raises(TypeError, f, d, out="invalid", dtype=None) + # keepdims boolean, no invalid value + # assert_raises(TypeError, f, d, 0, None, None, "invalid") + # assert_raises(TypeError, f, d, keepdims="invalid", axis=0, dtype=None) + # invalid mix + assert_raises(TypeError, f, d, 0, keepdims="invalid", dtype="invalid", + out=None) + + # invalid keyword + assert_raises(TypeError, f, d, axis=0, dtype=None, invalid=0) + assert_raises(TypeError, f, d, invalid=0) + assert_raises(TypeError, f, d, 0, keepdims=True, invalid="invalid", + out=None) + assert_raises(TypeError, f, d, axis=0, dtype=None, keepdims=True, + out=None, invalid=0) + assert_raises(TypeError, f, d, axis=0, dtype=None, + out=None, invalid=0) + + def test_structured_equal(self): + # https://github.com/numpy/numpy/issues/4855 + + class MyA(np.ndarray): + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + return getattr(ufunc, method)(*(input.view(np.ndarray) + for input in inputs), **kwargs) + a = np.arange(12.).reshape(4,3) + ra = a.view(dtype=('f8,f8,f8')).squeeze() + mra = ra.view(MyA) + + target = np.array([ True, False, False, False], dtype=bool) + assert_equal(np.all(target == (mra == ra[0])), True) + + def test_scalar_equal(self): + # Scalar comparisons should always work, without deprecation warnings. + # even when the ufunc fails. + a = np.array(0.) + b = np.array('a') + assert_(a != b) + assert_(b != a) + assert_(not (a == b)) + assert_(not (b == a)) + + def test_NotImplemented_not_returned(self): + # See gh-5964 and gh-2091. Some of these functions are not operator + # related and were fixed for other reasons in the past. + binary_funcs = [ + np.power, np.add, np.subtract, np.multiply, np.divide, + np.true_divide, np.floor_divide, np.bitwise_and, np.bitwise_or, + np.bitwise_xor, np.left_shift, np.right_shift, np.fmax, + np.fmin, np.fmod, np.hypot, np.logaddexp, np.logaddexp2, + np.maximum, np.minimum, np.mod, + np.greater, np.greater_equal, np.less, np.less_equal, + np.equal, np.not_equal] + + a = np.array('1') + b = 1 + c = np.array([1., 2.]) + for f in binary_funcs: + assert_raises(TypeError, f, a, b) + assert_raises(TypeError, f, c, a) + + @pytest.mark.parametrize("ufunc", + [np.logical_and, np.logical_or]) # logical_xor object loop is bad + @pytest.mark.parametrize("signature", + [(None, None, object), (object, None, None), + (None, object, None)]) + def test_logical_ufuncs_object_signatures(self, ufunc, signature): + a = np.array([True, None, False], dtype=object) + res = ufunc(a, a, signature=signature) + assert res.dtype == object + + @pytest.mark.parametrize("ufunc", + [np.logical_and, np.logical_or, np.logical_xor]) + @pytest.mark.parametrize("signature", + [(bool, None, object), (object, None, bool), + (None, object, bool)]) + def test_logical_ufuncs_mixed_object_signatures(self, ufunc, signature): + # Most mixed signatures fail (except those with bool out, e.g. `OO->?`) + a = np.array([True, None, False]) + with pytest.raises(TypeError): + ufunc(a, a, signature=signature) + + @pytest.mark.parametrize("ufunc", + [np.logical_and, np.logical_or, np.logical_xor]) + def test_logical_ufuncs_support_anything(self, ufunc): + # The logical ufuncs support even input that can't be promoted: + a = np.array(b'1', dtype="V3") + c = np.array([1., 2.]) + assert_array_equal(ufunc(a, c), ufunc([True, True], True)) + assert ufunc.reduce(a) == True + # check that the output has no effect: + out = np.zeros(2, dtype=np.int32) + expected = ufunc([True, True], True).astype(out.dtype) + assert_array_equal(ufunc(a, c, out=out), expected) + out = np.zeros((), dtype=np.int32) + assert ufunc.reduce(a, out=out) == True + # Last check, test reduction when out and a match (the complexity here + # is that the "i,i->?" may seem right, but should not match. + a = np.array([3], dtype="i") + out = np.zeros((), dtype=a.dtype) + assert ufunc.reduce(a, out=out) == 1 + + @pytest.mark.parametrize("ufunc", + [np.logical_and, np.logical_or, np.logical_xor]) + @pytest.mark.parametrize("dtype", ["S", "U"]) + @pytest.mark.parametrize("values", [["1", "hi", "0"], ["", ""]]) + def test_logical_ufuncs_supports_string(self, ufunc, dtype, values): + # note that values are either all true or all false + arr = np.array(values, dtype=dtype) + obj_arr = np.array(values, dtype=object) + res = ufunc(arr, arr) + expected = ufunc(obj_arr, obj_arr, dtype=bool) + + assert_array_equal(res, expected) + + res = ufunc.reduce(arr) + expected = ufunc.reduce(obj_arr, dtype=bool) + assert_array_equal(res, expected) + + @pytest.mark.parametrize("ufunc", + [np.logical_and, np.logical_or, np.logical_xor]) + def test_logical_ufuncs_out_cast_check(self, ufunc): + a = np.array('1') + c = np.array([1., 2.]) + out = a.copy() + with pytest.raises(TypeError): + # It would be safe, but not equiv casting: + ufunc(a, c, out=out, casting="equiv") + + def test_reducelike_byteorder_resolution(self): + # See gh-20699, byte-order changes need some extra care in the type + # resolution to make the following succeed: + arr_be = np.arange(10, dtype=">i8") + arr_le = np.arange(10, dtype="i + if 'O' in typ or '?' in typ: + continue + inp, out = typ.split('->') + args = [np.ones((3, 3), t) for t in inp] + with warnings.catch_warnings(record=True): + warnings.filterwarnings("always") + res = ufunc(*args) + if isinstance(res, tuple): + outs = tuple(out) + assert len(res) == len(outs) + for r, t in zip(res, outs): + assert r.dtype == np.dtype(t) + else: + assert res.dtype == np.dtype(out) + +@pytest.mark.parametrize('ufunc', [getattr(np, x) for x in dir(np) + if isinstance(getattr(np, x), np.ufunc)]) +def test_ufunc_noncontiguous(ufunc): + ''' + Check that contiguous and non-contiguous calls to ufuncs + have the same results for values in range(9) + ''' + for typ in ufunc.types: + # types is a list of strings like ii->i + if any(set('O?mM') & set(typ)): + # bool, object, datetime are too irregular for this simple test + continue + inp, out = typ.split('->') + args_c = [np.empty((6, 6), t) for t in inp] + # non contiguous (2, 3 step on the two dimensions) + args_n = [np.empty((12, 18), t)[::2, ::3] for t in inp] + # alignment != itemsize is possible. So create an array with such + # an odd step manually. + args_o = [] + for t in inp: + orig_dt = np.dtype(t) + off_dt = f"S{orig_dt.alignment}" # offset by alignment + dtype = np.dtype([("_", off_dt), ("t", orig_dt)], align=False) + args_o.append(np.empty((6, 6), dtype=dtype)["t"]) + for a in args_c + args_n + args_o: + a.flat = range(1, 37) + + with warnings.catch_warnings(record=True): + warnings.filterwarnings("always") + res_c = ufunc(*args_c) + res_n = ufunc(*args_n) + res_o = ufunc(*args_o) + if len(out) == 1: + res_c = (res_c,) + res_n = (res_n,) + res_o = (res_o,) + for c_ar, n_ar, o_ar in zip(res_c, res_n, res_o): + dt = c_ar.dtype + if np.issubdtype(dt, np.floating): + # for floating point results allow a small fuss in comparisons + # since different algorithms (libm vs. intrinsics) can be used + # for different input strides + res_eps = np.finfo(dt).eps + tol = 3*res_eps + assert_allclose(res_c, res_n, atol=tol, rtol=tol) + assert_allclose(res_c, res_o, atol=tol, rtol=tol) + else: + assert_equal(c_ar, n_ar) + assert_equal(c_ar, o_ar) + + +@pytest.mark.parametrize('ufunc', [np.sign, np.equal]) +def test_ufunc_warn_with_nan(ufunc): + # issue gh-15127 + # test that calling certain ufuncs with a non-standard `nan` value does not + # emit a warning + # `b` holds a 64 bit signaling nan: the most significant bit of the + # significand is zero. + b = np.array([0x7ff0000000000001], 'i8').view('f8') + assert np.isnan(b) + if ufunc.nin == 1: + ufunc(b) + elif ufunc.nin == 2: + ufunc(b, b.copy()) + else: + raise ValueError('ufunc with more than 2 inputs') + + +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +def test_ufunc_out_casterrors(): + # Tests that casting errors are correctly reported and buffers are + # cleared. + # The following array can be added to itself as an object array, but + # the result cannot be cast to an integer output: + value = 123 # relies on python cache (leak-check will still find it) + arr = np.array([value] * int(ncu.BUFSIZE * 1.5) + + ["string"] + + [value] * int(1.5 * ncu.BUFSIZE), dtype=object) + out = np.ones(len(arr), dtype=np.intp) + + count = sys.getrefcount(value) + with pytest.raises(ValueError): + # Output casting failure: + np.add(arr, arr, out=out, casting="unsafe") + + assert count == sys.getrefcount(value) + # output is unchanged after the error, this shows that the iteration + # was aborted (this is not necessarily defined behaviour) + assert out[-1] == 1 + + with pytest.raises(ValueError): + # Input casting failure: + np.add(arr, arr, out=out, dtype=np.intp, casting="unsafe") + + assert count == sys.getrefcount(value) + # output is unchanged after the error, this shows that the iteration + # was aborted (this is not necessarily defined behaviour) + assert out[-1] == 1 + + +@pytest.mark.parametrize("bad_offset", [0, int(ncu.BUFSIZE * 1.5)]) +def test_ufunc_input_casterrors(bad_offset): + value = 123 + arr = np.array([value] * bad_offset + + ["string"] + + [value] * int(1.5 * ncu.BUFSIZE), dtype=object) + with pytest.raises(ValueError): + # Force cast inputs, but the buffered cast of `arr` to intp fails: + np.add(arr, arr, dtype=np.intp, casting="unsafe") + + +@pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") +@pytest.mark.parametrize("bad_offset", [0, int(ncu.BUFSIZE * 1.5)]) +def test_ufunc_input_floatingpoint_error(bad_offset): + value = 123 + arr = np.array([value] * bad_offset + + [np.nan] + + [value] * int(1.5 * ncu.BUFSIZE)) + with np.errstate(invalid="raise"), pytest.raises(FloatingPointError): + # Force cast inputs, but the buffered cast of `arr` to intp fails: + np.add(arr, arr, dtype=np.intp, casting="unsafe") + + +def test_trivial_loop_invalid_cast(): + # This tests the fast-path "invalid cast", see gh-19904. + with pytest.raises(TypeError, + match="cast ufunc 'add' input 0"): + # the void dtype definitely cannot cast to double: + np.add(np.array(1, "i,i"), 3, signature="dd->d") + + +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +@pytest.mark.parametrize("offset", + [0, ncu.BUFSIZE//2, int(1.5*ncu.BUFSIZE)]) +def test_reduce_casterrors(offset): + # Test reporting of casting errors in reductions, we test various + # offsets to where the casting error will occur, since these may occur + # at different places during the reduction procedure. For example + # the first item may be special. + value = 123 # relies on python cache (leak-check will still find it) + arr = np.array([value] * offset + + ["string"] + + [value] * int(1.5 * ncu.BUFSIZE), dtype=object) + out = np.array(-1, dtype=np.intp) + + count = sys.getrefcount(value) + with pytest.raises(ValueError, match="invalid literal"): + # This is an unsafe cast, but we currently always allow that. + # Note that the double loop is picked, but the cast fails. + # `initial=None` disables the use of an identity here to test failures + # while copying the first values path (not used when identity exists). + np.add.reduce(arr, dtype=np.intp, out=out, initial=None) + assert count == sys.getrefcount(value) + # If an error occurred during casting, the operation is done at most until + # the error occurs (the result of which would be `value * offset`) and -1 + # if the error happened immediately. + # This does not define behaviour, the output is invalid and thus undefined + assert out[()] < value * offset + + +def test_object_reduce_cleanup_on_failure(): + # Test cleanup, including of the initial value (manually provided or not) + with pytest.raises(TypeError): + np.add.reduce([1, 2, None], initial=4) + + with pytest.raises(TypeError): + np.add.reduce([1, 2, None]) + + +@pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") +@pytest.mark.parametrize("method", + [np.add.accumulate, np.add.reduce, + pytest.param(lambda x: np.add.reduceat(x, [0]), id="reduceat"), + pytest.param(lambda x: np.log.at(x, [2]), id="at")]) +def test_ufunc_methods_floaterrors(method): + # adding inf and -inf (or log(-inf) creates an invalid float and warns + arr = np.array([np.inf, 0, -np.inf]) + with np.errstate(all="warn"): + with pytest.warns(RuntimeWarning, match="invalid value"): + method(arr) + + arr = np.array([np.inf, 0, -np.inf]) + with np.errstate(all="raise"): + with pytest.raises(FloatingPointError): + method(arr) + + +def _check_neg_zero(value): + if value != 0.0: + return False + if not np.signbit(value.real): + return False + if value.dtype.kind == "c": + return np.signbit(value.imag) + return True + +@pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) +def test_addition_negative_zero(dtype): + dtype = np.dtype(dtype) + if dtype.kind == "c": + neg_zero = dtype.type(complex(-0.0, -0.0)) + else: + neg_zero = dtype.type(-0.0) + + arr = np.array(neg_zero) + arr2 = np.array(neg_zero) + + assert _check_neg_zero(arr + arr2) + # In-place ops may end up on a different path (reduce path) see gh-21211 + arr += arr2 + assert _check_neg_zero(arr) + + +@pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) +@pytest.mark.parametrize("use_initial", [True, False]) +def test_addition_reduce_negative_zero(dtype, use_initial): + dtype = np.dtype(dtype) + if dtype.kind == "c": + neg_zero = dtype.type(complex(-0.0, -0.0)) + else: + neg_zero = dtype.type(-0.0) + + kwargs = {} + if use_initial: + kwargs["initial"] = neg_zero + else: + pytest.xfail("-0. propagation in sum currently requires initial") + + # Test various length, in case SIMD paths or chunking play a role. + # 150 extends beyond the pairwise blocksize; probably not important. + for i in range(0, 150): + arr = np.array([neg_zero] * i, dtype=dtype) + res = np.sum(arr, **kwargs) + if i > 0 or use_initial: + assert _check_neg_zero(res) + else: + # `sum([])` should probably be 0.0 and not -0.0 like `sum([-0.0])` + assert not np.signbit(res.real) + assert not np.signbit(res.imag) + + +@pytest.mark.parametrize(["dt1", "dt2"], + [("S", "U"), ("U", "S"), ("S", "d"), ("S", "V"), ("U", "l")]) +def test_addition_string_types(dt1, dt2): + arr1 = np.array([1234234], dtype=dt1) + arr2 = np.array([b"423"], dtype=dt2) + with pytest.raises(np._core._exceptions.UFuncTypeError) as exc: + np.add(arr1, arr2) + + +@pytest.mark.parametrize("order1,order2", + [(">", ">"), ("<", "<"), (">", "<"), ("<", ">")]) +def test_addition_unicode_inverse_byte_order(order1, order2): + element = 'abcd' + arr1 = np.array([element], dtype=f"{order1}U4") + arr2 = np.array([element], dtype=f"{order2}U4") + result = arr1 + arr2 + assert result == 2*element + + +@pytest.mark.parametrize("dtype", [np.int8, np.int16, np.int32, np.int64]) +def test_find_non_long_args(dtype): + element = 'abcd' + start = dtype(0) + end = dtype(len(element)) + arr = np.array([element]) + result = np._core.umath.find(arr, "a", start, end) + assert result.dtype == np.dtype("intp") + assert result == 0 + + +def test_find_access_past_buffer(): + # This checks that no read past the string buffer occurs in + # string_fastsearch.h. The buffer class makes sure this is checked. + # To see it in action, you can remove the checks in the buffer and + # this test will produce an 'Invalid read' if run under valgrind. + arr = np.array([b'abcd', b'ebcd']) + result = np._core.umath.find(arr, b'cde', 0, np.iinfo(np.int64).max) + assert np.all(result == -1) + + +class TestLowlevelAPIAccess: + def test_resolve_dtypes_basic(self): + # Basic test for dtype resolution: + i4 = np.dtype("i4") + f4 = np.dtype("f4") + f8 = np.dtype("f8") + + r = np.add.resolve_dtypes((i4, f4, None)) + assert r == (f8, f8, f8) + + # Signature uses the same logic to parse as ufunc (less strict) + # the following is "same-kind" casting so works: + r = np.add.resolve_dtypes(( + i4, i4, None), signature=(None, None, "f4")) + assert r == (f4, f4, f4) + + # Check NEP 50 "weak" promotion also: + r = np.add.resolve_dtypes((f4, int, None)) + assert r == (f4, f4, f4) + + with pytest.raises(TypeError): + np.add.resolve_dtypes((i4, f4, None), casting="no") + + def test_resolve_dtypes_comparison(self): + i4 = np.dtype("i4") + i8 = np.dtype("i8") + b = np.dtype("?") + r = np.equal.resolve_dtypes((i4, i8, None)) + assert r == (i8, i8, b) + + def test_weird_dtypes(self): + S0 = np.dtype("S0") + # S0 is often converted by NumPy to S1, but not here: + r = np.equal.resolve_dtypes((S0, S0, None)) + assert r == (S0, S0, np.dtype(bool)) + + # Subarray dtypes are weird and may not work fully, we preserve them + # leading to a TypeError (currently no equal loop for void/structured) + dts = np.dtype("10i") + with pytest.raises(TypeError): + np.equal.resolve_dtypes((dts, dts, None)) + + def test_resolve_dtypes_reduction(self): + i2 = np.dtype("i2") + default_int_ = np.dtype(np.int_) + # Check special addition resolution: + res = np.add.resolve_dtypes((None, i2, None), reduction=True) + assert res == (default_int_, default_int_, default_int_) + + def test_resolve_dtypes_reduction_no_output(self): + i4 = np.dtype("i4") + with pytest.raises(TypeError): + # May be allowable at some point? + np.add.resolve_dtypes((i4, i4, i4), reduction=True) + + @pytest.mark.parametrize("dtypes", [ + (np.dtype("i"), np.dtype("i")), + (None, np.dtype("i"), np.dtype("f")), + (np.dtype("i"), None, np.dtype("f")), + ("i4", "i4", None)]) + def test_resolve_dtypes_errors(self, dtypes): + with pytest.raises(TypeError): + np.add.resolve_dtypes(dtypes) + + def test_resolve_dtypes_reduction_errors(self): + i2 = np.dtype("i2") + + with pytest.raises(TypeError): + np.add.resolve_dtypes((None, i2, i2)) + + with pytest.raises(TypeError): + np.add.signature((None, None, "i4")) + + @pytest.mark.skipif(not hasattr(ct, "pythonapi"), + reason="`ctypes.pythonapi` required for capsule unpacking.") + def test_loop_access(self): + # This is a basic test for the full strided loop access + data_t = ct.c_char_p * 2 + dim_t = ct.c_ssize_t * 1 + strides_t = ct.c_ssize_t * 2 + strided_loop_t = ct.CFUNCTYPE( + ct.c_int, ct.c_void_p, data_t, dim_t, strides_t, ct.c_void_p) + + class call_info_t(ct.Structure): + _fields_ = [ + ("strided_loop", strided_loop_t), + ("context", ct.c_void_p), + ("auxdata", ct.c_void_p), + ("requires_pyapi", ct.c_byte), + ("no_floatingpoint_errors", ct.c_byte), + ] + + i4 = np.dtype("i4") + dt, call_info_obj = np.negative._resolve_dtypes_and_context((i4, i4)) + assert dt == (i4, i4) # can be used without casting + + # Fill in the rest of the information: + np.negative._get_strided_loop(call_info_obj) + + ct.pythonapi.PyCapsule_GetPointer.restype = ct.c_void_p + call_info = ct.pythonapi.PyCapsule_GetPointer( + ct.py_object(call_info_obj), + ct.c_char_p(b"numpy_1.24_ufunc_call_info")) + + call_info = ct.cast(call_info, ct.POINTER(call_info_t)).contents + + arr = np.arange(10, dtype=i4) + call_info.strided_loop( + call_info.context, + data_t(arr.ctypes.data, arr.ctypes.data), + arr.ctypes.shape, # is a C-array with 10 here + strides_t(arr.ctypes.strides[0], arr.ctypes.strides[0]), + call_info.auxdata) + + # We just directly called the negative inner-loop in-place: + assert_array_equal(arr, -np.arange(10, dtype=i4)) + + @pytest.mark.parametrize("strides", [1, (1, 2, 3), (1, "2")]) + def test__get_strided_loop_errors_bad_strides(self, strides): + i4 = np.dtype("i4") + dt, call_info = np.negative._resolve_dtypes_and_context((i4, i4)) + + with pytest.raises(TypeError, match="fixed_strides.*tuple.*or None"): + np.negative._get_strided_loop(call_info, fixed_strides=strides) + + def test__get_strided_loop_errors_bad_call_info(self): + i4 = np.dtype("i4") + dt, call_info = np.negative._resolve_dtypes_and_context((i4, i4)) + + with pytest.raises(ValueError, match="PyCapsule"): + np.negative._get_strided_loop("not the capsule!") + + with pytest.raises(TypeError, match=".*incompatible context"): + np.add._get_strided_loop(call_info) + + np.negative._get_strided_loop(call_info) + with pytest.raises(TypeError): + # cannot call it a second time: + np.negative._get_strided_loop(call_info) + + def test_long_arrays(self): + t = np.zeros((1029, 917), dtype=np.single) + t[0][0] = 1 + t[28][414] = 1 + tc = np.cos(t) + assert_equal(tc[0][0], tc[28][414]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath.py new file mode 100644 index 0000000000000000000000000000000000000000..737622f1f57ee7f9976384d9bd3b4de1326e5a30 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath.py @@ -0,0 +1,4897 @@ +import platform +import warnings +import fnmatch +import itertools +import pytest +import sys +import operator +from fractions import Fraction +from functools import reduce +from collections import namedtuple + +import numpy._core.umath as ncu +from numpy._core import _umath_tests as ncu_tests, sctypes +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_raises, assert_raises_regex, + assert_array_equal, assert_almost_equal, assert_array_almost_equal, + assert_array_max_ulp, assert_allclose, assert_no_warnings, suppress_warnings, + _gen_alignment_data, assert_array_almost_equal_nulp, IS_WASM, IS_MUSL, + IS_PYPY, HAS_REFCOUNT + ) +from numpy.testing._private.utils import _glibc_older_than + +UFUNCS = [obj for obj in np._core.umath.__dict__.values() + if isinstance(obj, np.ufunc)] + +UFUNCS_UNARY = [ + uf for uf in UFUNCS if uf.nin == 1 +] +UFUNCS_UNARY_FP = [ + uf for uf in UFUNCS_UNARY if 'f->f' in uf.types +] + +UFUNCS_BINARY = [ + uf for uf in UFUNCS if uf.nin == 2 +] +UFUNCS_BINARY_ACC = [ + uf for uf in UFUNCS_BINARY if hasattr(uf, "accumulate") and uf.nout == 1 +] + +def interesting_binop_operands(val1, val2, dtype): + """ + Helper to create "interesting" operands to cover common code paths: + * scalar inputs + * only first "values" is an array (e.g. scalar division fast-paths) + * Longer array (SIMD) placing the value of interest at different positions + * Oddly strided arrays which may not be SIMD compatible + + It does not attempt to cover unaligned access or mixed dtypes. + These are normally handled by the casting/buffering machinery. + + This is not a fixture (currently), since I believe a fixture normally + only yields once? + """ + fill_value = 1 # could be a parameter, but maybe not an optional one? + + arr1 = np.full(10003, dtype=dtype, fill_value=fill_value) + arr2 = np.full(10003, dtype=dtype, fill_value=fill_value) + + arr1[0] = val1 + arr2[0] = val2 + + extractor = lambda res: res + yield arr1[0], arr2[0], extractor, "scalars" + + extractor = lambda res: res + yield arr1[0, ...], arr2[0, ...], extractor, "scalar-arrays" + + # reset array values to fill_value: + arr1[0] = fill_value + arr2[0] = fill_value + + for pos in [0, 1, 2, 3, 4, 5, -1, -2, -3, -4]: + arr1[pos] = val1 + arr2[pos] = val2 + + extractor = lambda res: res[pos] + yield arr1, arr2, extractor, f"off-{pos}" + yield arr1, arr2[pos], extractor, f"off-{pos}-with-scalar" + + arr1[pos] = fill_value + arr2[pos] = fill_value + + for stride in [-1, 113]: + op1 = arr1[::stride] + op2 = arr2[::stride] + op1[10] = val1 + op2[10] = val2 + + extractor = lambda res: res[10] + yield op1, op2, extractor, f"stride-{stride}" + + op1[10] = fill_value + op2[10] = fill_value + + +def on_powerpc(): + """ True if we are running on a Power PC platform.""" + return platform.processor() == 'powerpc' or \ + platform.machine().startswith('ppc') + + +def bad_arcsinh(): + """The blocklisted trig functions are not accurate on aarch64/PPC for + complex256. Rather than dig through the actual problem skip the + test. This should be fixed when we can move past glibc2.17 + which is the version in manylinux2014 + """ + if platform.machine() == 'aarch64': + x = 1.78e-10 + elif on_powerpc(): + x = 2.16e-10 + else: + return False + v1 = np.arcsinh(np.float128(x)) + v2 = np.arcsinh(np.complex256(x)).real + # The eps for float128 is 1-e33, so this is way bigger + return abs((v1 / v2) - 1.0) > 1e-23 + + +class _FilterInvalids: + def setup_method(self): + self.olderr = np.seterr(invalid='ignore') + + def teardown_method(self): + np.seterr(**self.olderr) + + +class TestConstants: + def test_pi(self): + assert_allclose(ncu.pi, 3.141592653589793, 1e-15) + + def test_e(self): + assert_allclose(ncu.e, 2.718281828459045, 1e-15) + + def test_euler_gamma(self): + assert_allclose(ncu.euler_gamma, 0.5772156649015329, 1e-15) + + +class TestOut: + def test_out_subok(self): + for subok in (True, False): + a = np.array(0.5) + o = np.empty(()) + + r = np.add(a, 2, o, subok=subok) + assert_(r is o) + r = np.add(a, 2, out=o, subok=subok) + assert_(r is o) + r = np.add(a, 2, out=(o,), subok=subok) + assert_(r is o) + + d = np.array(5.7) + o1 = np.empty(()) + o2 = np.empty((), dtype=np.int32) + + r1, r2 = np.frexp(d, o1, None, subok=subok) + assert_(r1 is o1) + r1, r2 = np.frexp(d, None, o2, subok=subok) + assert_(r2 is o2) + r1, r2 = np.frexp(d, o1, o2, subok=subok) + assert_(r1 is o1) + assert_(r2 is o2) + + r1, r2 = np.frexp(d, out=(o1, None), subok=subok) + assert_(r1 is o1) + r1, r2 = np.frexp(d, out=(None, o2), subok=subok) + assert_(r2 is o2) + r1, r2 = np.frexp(d, out=(o1, o2), subok=subok) + assert_(r1 is o1) + assert_(r2 is o2) + + with assert_raises(TypeError): + # Out argument must be tuple, since there are multiple outputs. + r1, r2 = np.frexp(d, out=o1, subok=subok) + + assert_raises(TypeError, np.add, a, 2, o, o, subok=subok) + assert_raises(TypeError, np.add, a, 2, o, out=o, subok=subok) + assert_raises(TypeError, np.add, a, 2, None, out=o, subok=subok) + assert_raises(ValueError, np.add, a, 2, out=(o, o), subok=subok) + assert_raises(ValueError, np.add, a, 2, out=(), subok=subok) + assert_raises(TypeError, np.add, a, 2, [], subok=subok) + assert_raises(TypeError, np.add, a, 2, out=[], subok=subok) + assert_raises(TypeError, np.add, a, 2, out=([],), subok=subok) + o.flags.writeable = False + assert_raises(ValueError, np.add, a, 2, o, subok=subok) + assert_raises(ValueError, np.add, a, 2, out=o, subok=subok) + assert_raises(ValueError, np.add, a, 2, out=(o,), subok=subok) + + def test_out_wrap_subok(self): + class ArrayWrap(np.ndarray): + __array_priority__ = 10 + + def __new__(cls, arr): + return np.asarray(arr).view(cls).copy() + + def __array_wrap__(self, arr, context=None, return_scalar=False): + return arr.view(type(self)) + + for subok in (True, False): + a = ArrayWrap([0.5]) + + r = np.add(a, 2, subok=subok) + if subok: + assert_(isinstance(r, ArrayWrap)) + else: + assert_(type(r) == np.ndarray) + + r = np.add(a, 2, None, subok=subok) + if subok: + assert_(isinstance(r, ArrayWrap)) + else: + assert_(type(r) == np.ndarray) + + r = np.add(a, 2, out=None, subok=subok) + if subok: + assert_(isinstance(r, ArrayWrap)) + else: + assert_(type(r) == np.ndarray) + + r = np.add(a, 2, out=(None,), subok=subok) + if subok: + assert_(isinstance(r, ArrayWrap)) + else: + assert_(type(r) == np.ndarray) + + d = ArrayWrap([5.7]) + o1 = np.empty((1,)) + o2 = np.empty((1,), dtype=np.int32) + + r1, r2 = np.frexp(d, o1, subok=subok) + if subok: + assert_(isinstance(r2, ArrayWrap)) + else: + assert_(type(r2) == np.ndarray) + + r1, r2 = np.frexp(d, o1, None, subok=subok) + if subok: + assert_(isinstance(r2, ArrayWrap)) + else: + assert_(type(r2) == np.ndarray) + + r1, r2 = np.frexp(d, None, o2, subok=subok) + if subok: + assert_(isinstance(r1, ArrayWrap)) + else: + assert_(type(r1) == np.ndarray) + + r1, r2 = np.frexp(d, out=(o1, None), subok=subok) + if subok: + assert_(isinstance(r2, ArrayWrap)) + else: + assert_(type(r2) == np.ndarray) + + r1, r2 = np.frexp(d, out=(None, o2), subok=subok) + if subok: + assert_(isinstance(r1, ArrayWrap)) + else: + assert_(type(r1) == np.ndarray) + + with assert_raises(TypeError): + # Out argument must be tuple, since there are multiple outputs. + r1, r2 = np.frexp(d, out=o1, subok=subok) + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_out_wrap_no_leak(self): + # Regression test for gh-26545 + class ArrSubclass(np.ndarray): + pass + + arr = np.arange(10).view(ArrSubclass) + + arr *= 1 + assert sys.getrefcount(arr) == 2 + + +class TestComparisons: + import operator + + @pytest.mark.parametrize('dtype', sctypes['uint'] + sctypes['int'] + + sctypes['float'] + [np.bool]) + @pytest.mark.parametrize('py_comp,np_comp', [ + (operator.lt, np.less), + (operator.le, np.less_equal), + (operator.gt, np.greater), + (operator.ge, np.greater_equal), + (operator.eq, np.equal), + (operator.ne, np.not_equal) + ]) + def test_comparison_functions(self, dtype, py_comp, np_comp): + # Initialize input arrays + if dtype == np.bool: + a = np.random.choice(a=[False, True], size=1000) + b = np.random.choice(a=[False, True], size=1000) + scalar = True + else: + a = np.random.randint(low=1, high=10, size=1000).astype(dtype) + b = np.random.randint(low=1, high=10, size=1000).astype(dtype) + scalar = 5 + np_scalar = np.dtype(dtype).type(scalar) + a_lst = a.tolist() + b_lst = b.tolist() + + # (Binary) Comparison (x1=array, x2=array) + comp_b = np_comp(a, b).view(np.uint8) + comp_b_list = [int(py_comp(x, y)) for x, y in zip(a_lst, b_lst)] + + # (Scalar1) Comparison (x1=scalar, x2=array) + comp_s1 = np_comp(np_scalar, b).view(np.uint8) + comp_s1_list = [int(py_comp(scalar, x)) for x in b_lst] + + # (Scalar2) Comparison (x1=array, x2=scalar) + comp_s2 = np_comp(a, np_scalar).view(np.uint8) + comp_s2_list = [int(py_comp(x, scalar)) for x in a_lst] + + # Sequence: Binary, Scalar1 and Scalar2 + assert_(comp_b.tolist() == comp_b_list, + f"Failed comparison ({py_comp.__name__})") + assert_(comp_s1.tolist() == comp_s1_list, + f"Failed comparison ({py_comp.__name__})") + assert_(comp_s2.tolist() == comp_s2_list, + f"Failed comparison ({py_comp.__name__})") + + def test_ignore_object_identity_in_equal(self): + # Check comparing identical objects whose comparison + # is not a simple boolean, e.g., arrays that are compared elementwise. + a = np.array([np.array([1, 2, 3]), None], dtype=object) + assert_raises(ValueError, np.equal, a, a) + + # Check error raised when comparing identical non-comparable objects. + class FunkyType: + def __eq__(self, other): + raise TypeError("I won't compare") + + a = np.array([FunkyType()]) + assert_raises(TypeError, np.equal, a, a) + + # Check identity doesn't override comparison mismatch. + a = np.array([np.nan], dtype=object) + assert_equal(np.equal(a, a), [False]) + + def test_ignore_object_identity_in_not_equal(self): + # Check comparing identical objects whose comparison + # is not a simple boolean, e.g., arrays that are compared elementwise. + a = np.array([np.array([1, 2, 3]), None], dtype=object) + assert_raises(ValueError, np.not_equal, a, a) + + # Check error raised when comparing identical non-comparable objects. + class FunkyType: + def __ne__(self, other): + raise TypeError("I won't compare") + + a = np.array([FunkyType()]) + assert_raises(TypeError, np.not_equal, a, a) + + # Check identity doesn't override comparison mismatch. + a = np.array([np.nan], dtype=object) + assert_equal(np.not_equal(a, a), [True]) + + def test_error_in_equal_reduce(self): + # gh-20929 + # make sure np.equal.reduce raises a TypeError if an array is passed + # without specifying the dtype + a = np.array([0, 0]) + assert_equal(np.equal.reduce(a, dtype=bool), True) + assert_raises(TypeError, np.equal.reduce, a) + + def test_object_dtype(self): + assert np.equal(1, [1], dtype=object).dtype == object + assert np.equal(1, [1], signature=(None, None, "O")).dtype == object + + def test_object_nonbool_dtype_error(self): + # bool output dtype is fine of course: + assert np.equal(1, [1], dtype=bool).dtype == bool + + # but the following are examples do not have a loop: + with pytest.raises(TypeError, match="No loop matching"): + np.equal(1, 1, dtype=np.int64) + + with pytest.raises(TypeError, match="No loop matching"): + np.equal(1, 1, sig=(None, None, "l")) + + @pytest.mark.parametrize("dtypes", ["qQ", "Qq"]) + @pytest.mark.parametrize('py_comp, np_comp', [ + (operator.lt, np.less), + (operator.le, np.less_equal), + (operator.gt, np.greater), + (operator.ge, np.greater_equal), + (operator.eq, np.equal), + (operator.ne, np.not_equal) + ]) + @pytest.mark.parametrize("vals", [(2**60, 2**60+1), (2**60+1, 2**60)]) + def test_large_integer_direct_comparison( + self, dtypes, py_comp, np_comp, vals): + # Note that float(2**60) + 1 == float(2**60). + a1 = np.array([2**60], dtype=dtypes[0]) + a2 = np.array([2**60 + 1], dtype=dtypes[1]) + expected = py_comp(2**60, 2**60+1) + + assert py_comp(a1, a2) == expected + assert np_comp(a1, a2) == expected + # Also check the scalars: + s1 = a1[0] + s2 = a2[0] + assert isinstance(s1, np.integer) + assert isinstance(s2, np.integer) + # The Python operator here is mainly interesting: + assert py_comp(s1, s2) == expected + assert np_comp(s1, s2) == expected + + @pytest.mark.parametrize("dtype", np.typecodes['UnsignedInteger']) + @pytest.mark.parametrize('py_comp_func, np_comp_func', [ + (operator.lt, np.less), + (operator.le, np.less_equal), + (operator.gt, np.greater), + (operator.ge, np.greater_equal), + (operator.eq, np.equal), + (operator.ne, np.not_equal) + ]) + @pytest.mark.parametrize("flip", [True, False]) + def test_unsigned_signed_direct_comparison( + self, dtype, py_comp_func, np_comp_func, flip): + if flip: + py_comp = lambda x, y: py_comp_func(y, x) + np_comp = lambda x, y: np_comp_func(y, x) + else: + py_comp = py_comp_func + np_comp = np_comp_func + + arr = np.array([np.iinfo(dtype).max], dtype=dtype) + expected = py_comp(int(arr[0]), -1) + + assert py_comp(arr, -1) == expected + assert np_comp(arr, -1) == expected + + scalar = arr[0] + assert isinstance(scalar, np.integer) + # The Python operator here is mainly interesting: + assert py_comp(scalar, -1) == expected + assert np_comp(scalar, -1) == expected + + +class TestAdd: + def test_reduce_alignment(self): + # gh-9876 + # make sure arrays with weird strides work with the optimizations in + # pairwise_sum_@TYPE@. On x86, the 'b' field will count as aligned at a + # 4 byte offset, even though its itemsize is 8. + a = np.zeros(2, dtype=[('a', np.int32), ('b', np.float64)]) + a['a'] = -1 + assert_equal(a['b'].sum(), 0) + + +class TestDivision: + def test_division_int(self): + # int division should follow Python + x = np.array([5, 10, 90, 100, -5, -10, -90, -100, -120]) + if 5 / 10 == 0.5: + assert_equal(x / 100, [0.05, 0.1, 0.9, 1, + -0.05, -0.1, -0.9, -1, -1.2]) + else: + assert_equal(x / 100, [0, 0, 0, 1, -1, -1, -1, -1, -2]) + assert_equal(x // 100, [0, 0, 0, 1, -1, -1, -1, -1, -2]) + assert_equal(x % 100, [5, 10, 90, 0, 95, 90, 10, 0, 80]) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize("dtype,ex_val", itertools.product( + sctypes['int'] + sctypes['uint'], ( + ( + # dividend + "np.array(range(fo.max-lsize, fo.max)).astype(dtype)," + # divisors + "np.arange(lsize).astype(dtype)," + # scalar divisors + "range(15)" + ), + ( + # dividend + "np.arange(fo.min, fo.min+lsize).astype(dtype)," + # divisors + "np.arange(lsize//-2, lsize//2).astype(dtype)," + # scalar divisors + "range(fo.min, fo.min + 15)" + ), ( + # dividend + "np.array(range(fo.max-lsize, fo.max)).astype(dtype)," + # divisors + "np.arange(lsize).astype(dtype)," + # scalar divisors + "[1,3,9,13,neg, fo.min+1, fo.min//2, fo.max//3, fo.max//4]" + ) + ) + )) + def test_division_int_boundary(self, dtype, ex_val): + fo = np.iinfo(dtype) + neg = -1 if fo.min < 0 else 1 + # Large enough to test SIMD loops and remainder elements + lsize = 512 + 7 + a, b, divisors = eval(ex_val) + a_lst, b_lst = a.tolist(), b.tolist() + + c_div = lambda n, d: ( + 0 if d == 0 else ( + fo.min if (n and n == fo.min and d == -1) else n//d + ) + ) + with np.errstate(divide='ignore'): + ac = a.copy() + ac //= b + div_ab = a // b + div_lst = [c_div(x, y) for x, y in zip(a_lst, b_lst)] + + msg = "Integer arrays floor division check (//)" + assert all(div_ab == div_lst), msg + msg_eq = "Integer arrays floor division check (//=)" + assert all(ac == div_lst), msg_eq + + for divisor in divisors: + ac = a.copy() + with np.errstate(divide='ignore', over='ignore'): + div_a = a // divisor + ac //= divisor + div_lst = [c_div(i, divisor) for i in a_lst] + + assert all(div_a == div_lst), msg + assert all(ac == div_lst), msg_eq + + with np.errstate(divide='raise', over='raise'): + if 0 in b: + # Verify overflow case + with pytest.raises(FloatingPointError, + match="divide by zero encountered in floor_divide"): + a // b + else: + a // b + if fo.min and fo.min in a: + with pytest.raises(FloatingPointError, + match='overflow encountered in floor_divide'): + a // -1 + elif fo.min: + a // -1 + with pytest.raises(FloatingPointError, + match="divide by zero encountered in floor_divide"): + a // 0 + with pytest.raises(FloatingPointError, + match="divide by zero encountered in floor_divide"): + ac = a.copy() + ac //= 0 + + np.array([], dtype=dtype) // 0 + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize("dtype,ex_val", itertools.product( + sctypes['int'] + sctypes['uint'], ( + "np.array([fo.max, 1, 2, 1, 1, 2, 3], dtype=dtype)", + "np.array([fo.min, 1, -2, 1, 1, 2, -3]).astype(dtype)", + "np.arange(fo.min, fo.min+(100*10), 10, dtype=dtype)", + "np.array(range(fo.max-(100*7), fo.max, 7)).astype(dtype)", + ) + )) + def test_division_int_reduce(self, dtype, ex_val): + fo = np.iinfo(dtype) + a = eval(ex_val) + lst = a.tolist() + c_div = lambda n, d: ( + 0 if d == 0 or (n and n == fo.min and d == -1) else n//d + ) + + with np.errstate(divide='ignore'): + div_a = np.floor_divide.reduce(a) + div_lst = reduce(c_div, lst) + msg = "Reduce floor integer division check" + assert div_a == div_lst, msg + + with np.errstate(divide='raise', over='raise'): + with pytest.raises(FloatingPointError, + match="divide by zero encountered in reduce"): + np.floor_divide.reduce(np.arange(-100, 100).astype(dtype)) + if fo.min: + with pytest.raises(FloatingPointError, + match='overflow encountered in reduce'): + np.floor_divide.reduce( + np.array([fo.min, 1, -1], dtype=dtype) + ) + + @pytest.mark.parametrize( + "dividend,divisor,quotient", + [(np.timedelta64(2,'Y'), np.timedelta64(2,'M'), 12), + (np.timedelta64(2,'Y'), np.timedelta64(-2,'M'), -12), + (np.timedelta64(-2,'Y'), np.timedelta64(2,'M'), -12), + (np.timedelta64(-2,'Y'), np.timedelta64(-2,'M'), 12), + (np.timedelta64(2,'M'), np.timedelta64(-2,'Y'), -1), + (np.timedelta64(2,'Y'), np.timedelta64(0,'M'), 0), + (np.timedelta64(2,'Y'), 2, np.timedelta64(1,'Y')), + (np.timedelta64(2,'Y'), -2, np.timedelta64(-1,'Y')), + (np.timedelta64(-2,'Y'), 2, np.timedelta64(-1,'Y')), + (np.timedelta64(-2,'Y'), -2, np.timedelta64(1,'Y')), + (np.timedelta64(-2,'Y'), -2, np.timedelta64(1,'Y')), + (np.timedelta64(-2,'Y'), -3, np.timedelta64(0,'Y')), + (np.timedelta64(-2,'Y'), 0, np.timedelta64('Nat','Y')), + ]) + def test_division_int_timedelta(self, dividend, divisor, quotient): + # If either divisor is 0 or quotient is Nat, check for division by 0 + if divisor and (isinstance(quotient, int) or not np.isnat(quotient)): + msg = "Timedelta floor division check" + assert dividend // divisor == quotient, msg + + # Test for arrays as well + msg = "Timedelta arrays floor division check" + dividend_array = np.array([dividend]*5) + quotient_array = np.array([quotient]*5) + assert all(dividend_array // divisor == quotient_array), msg + else: + if IS_WASM: + pytest.skip("fp errors don't work in wasm") + with np.errstate(divide='raise', invalid='raise'): + with pytest.raises(FloatingPointError): + dividend // divisor + + def test_division_complex(self): + # check that implementation is correct + msg = "Complex division implementation check" + x = np.array([1. + 1.*1j, 1. + .5*1j, 1. + 2.*1j], dtype=np.complex128) + assert_almost_equal(x**2/x, x, err_msg=msg) + # check overflow, underflow + msg = "Complex division overflow/underflow check" + x = np.array([1.e+110, 1.e-110], dtype=np.complex128) + y = x**2/x + assert_almost_equal(y/x, [1, 1], err_msg=msg) + + def test_zero_division_complex(self): + with np.errstate(invalid="ignore", divide="ignore"): + x = np.array([0.0], dtype=np.complex128) + y = 1.0/x + assert_(np.isinf(y)[0]) + y = complex(np.inf, np.nan)/x + assert_(np.isinf(y)[0]) + y = complex(np.nan, np.inf)/x + assert_(np.isinf(y)[0]) + y = complex(np.inf, np.inf)/x + assert_(np.isinf(y)[0]) + y = 0.0/x + assert_(np.isnan(y)[0]) + + def test_floor_division_complex(self): + # check that floor division, divmod and remainder raises type errors + x = np.array([.9 + 1j, -.1 + 1j, .9 + .5*1j, .9 + 2.*1j], dtype=np.complex128) + with pytest.raises(TypeError): + x // 7 + with pytest.raises(TypeError): + np.divmod(x, 7) + with pytest.raises(TypeError): + np.remainder(x, 7) + + def test_floor_division_signed_zero(self): + # Check that the sign bit is correctly set when dividing positive and + # negative zero by one. + x = np.zeros(10) + assert_equal(np.signbit(x//1), 0) + assert_equal(np.signbit((-x)//1), 1) + + @pytest.mark.skipif(hasattr(np.__config__, "blas_ssl2_info"), + reason="gh-22982") + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize('dtype', np.typecodes['Float']) + def test_floor_division_errors(self, dtype): + fnan = np.array(np.nan, dtype=dtype) + fone = np.array(1.0, dtype=dtype) + fzer = np.array(0.0, dtype=dtype) + finf = np.array(np.inf, dtype=dtype) + # divide by zero error check + with np.errstate(divide='raise', invalid='ignore'): + assert_raises(FloatingPointError, np.floor_divide, fone, fzer) + with np.errstate(divide='ignore', invalid='raise'): + np.floor_divide(fone, fzer) + + # The following already contain a NaN and should not warn + with np.errstate(all='raise'): + np.floor_divide(fnan, fone) + np.floor_divide(fone, fnan) + np.floor_divide(fnan, fzer) + np.floor_divide(fzer, fnan) + + @pytest.mark.parametrize('dtype', np.typecodes['Float']) + def test_floor_division_corner_cases(self, dtype): + # test corner cases like 1.0//0.0 for errors and return vals + x = np.zeros(10, dtype=dtype) + y = np.ones(10, dtype=dtype) + fnan = np.array(np.nan, dtype=dtype) + fone = np.array(1.0, dtype=dtype) + fzer = np.array(0.0, dtype=dtype) + finf = np.array(np.inf, dtype=dtype) + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "invalid value encountered in floor_divide") + div = np.floor_divide(fnan, fone) + assert(np.isnan(div)), "div: %s" % div + div = np.floor_divide(fone, fnan) + assert(np.isnan(div)), "div: %s" % div + div = np.floor_divide(fnan, fzer) + assert(np.isnan(div)), "div: %s" % div + # verify 1.0//0.0 computations return inf + with np.errstate(divide='ignore'): + z = np.floor_divide(y, x) + assert_(np.isinf(z).all()) + +def floor_divide_and_remainder(x, y): + return (np.floor_divide(x, y), np.remainder(x, y)) + + +def _signs(dt): + if dt in np.typecodes['UnsignedInteger']: + return (+1,) + else: + return (+1, -1) + + +class TestRemainder: + + def test_remainder_basic(self): + dt = np.typecodes['AllInteger'] + np.typecodes['Float'] + for op in [floor_divide_and_remainder, np.divmod]: + for dt1, dt2 in itertools.product(dt, dt): + for sg1, sg2 in itertools.product(_signs(dt1), _signs(dt2)): + fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s' + msg = fmt % (op.__name__, dt1, dt2, sg1, sg2) + a = np.array(sg1*71, dtype=dt1) + b = np.array(sg2*19, dtype=dt2) + div, rem = op(a, b) + assert_equal(div*b + rem, a, err_msg=msg) + if sg2 == -1: + assert_(b < rem <= 0, msg) + else: + assert_(b > rem >= 0, msg) + + def test_float_remainder_exact(self): + # test that float results are exact for small integers. This also + # holds for the same integers scaled by powers of two. + nlst = list(range(-127, 0)) + plst = list(range(1, 128)) + dividend = nlst + [0] + plst + divisor = nlst + plst + arg = list(itertools.product(dividend, divisor)) + tgt = list(divmod(*t) for t in arg) + + a, b = np.array(arg, dtype=int).T + # convert exact integer results from Python to float so that + # signed zero can be used, it is checked. + tgtdiv, tgtrem = np.array(tgt, dtype=float).T + tgtdiv = np.where((tgtdiv == 0.0) & ((b < 0) ^ (a < 0)), -0.0, tgtdiv) + tgtrem = np.where((tgtrem == 0.0) & (b < 0), -0.0, tgtrem) + + for op in [floor_divide_and_remainder, np.divmod]: + for dt in np.typecodes['Float']: + msg = 'op: %s, dtype: %s' % (op.__name__, dt) + fa = a.astype(dt) + fb = b.astype(dt) + div, rem = op(fa, fb) + assert_equal(div, tgtdiv, err_msg=msg) + assert_equal(rem, tgtrem, err_msg=msg) + + def test_float_remainder_roundoff(self): + # gh-6127 + dt = np.typecodes['Float'] + for op in [floor_divide_and_remainder, np.divmod]: + for dt1, dt2 in itertools.product(dt, dt): + for sg1, sg2 in itertools.product((+1, -1), (+1, -1)): + fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s' + msg = fmt % (op.__name__, dt1, dt2, sg1, sg2) + a = np.array(sg1*78*6e-8, dtype=dt1) + b = np.array(sg2*6e-8, dtype=dt2) + div, rem = op(a, b) + # Equal assertion should hold when fmod is used + assert_equal(div*b + rem, a, err_msg=msg) + if sg2 == -1: + assert_(b < rem <= 0, msg) + else: + assert_(b > rem >= 0, msg) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.xfail(sys.platform.startswith("darwin"), + reason="MacOS seems to not give the correct 'invalid' warning for " + "`fmod`. Hopefully, others always do.") + @pytest.mark.parametrize('dtype', np.typecodes['Float']) + def test_float_divmod_errors(self, dtype): + # Check valid errors raised for divmod and remainder + fzero = np.array(0.0, dtype=dtype) + fone = np.array(1.0, dtype=dtype) + finf = np.array(np.inf, dtype=dtype) + fnan = np.array(np.nan, dtype=dtype) + # since divmod is combination of both remainder and divide + # ops it will set both dividebyzero and invalid flags + with np.errstate(divide='raise', invalid='ignore'): + assert_raises(FloatingPointError, np.divmod, fone, fzero) + with np.errstate(divide='ignore', invalid='raise'): + assert_raises(FloatingPointError, np.divmod, fone, fzero) + with np.errstate(invalid='raise'): + assert_raises(FloatingPointError, np.divmod, fzero, fzero) + with np.errstate(invalid='raise'): + assert_raises(FloatingPointError, np.divmod, finf, finf) + with np.errstate(divide='ignore', invalid='raise'): + assert_raises(FloatingPointError, np.divmod, finf, fzero) + with np.errstate(divide='raise', invalid='ignore'): + # inf / 0 does not set any flags, only the modulo creates a NaN + np.divmod(finf, fzero) + + @pytest.mark.skipif(hasattr(np.__config__, "blas_ssl2_info"), + reason="gh-22982") + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.xfail(sys.platform.startswith("darwin"), + reason="MacOS seems to not give the correct 'invalid' warning for " + "`fmod`. Hopefully, others always do.") + @pytest.mark.parametrize('dtype', np.typecodes['Float']) + @pytest.mark.parametrize('fn', [np.fmod, np.remainder]) + def test_float_remainder_errors(self, dtype, fn): + fzero = np.array(0.0, dtype=dtype) + fone = np.array(1.0, dtype=dtype) + finf = np.array(np.inf, dtype=dtype) + fnan = np.array(np.nan, dtype=dtype) + + # The following already contain a NaN and should not warn. + with np.errstate(all='raise'): + with pytest.raises(FloatingPointError, + match="invalid value"): + fn(fone, fzero) + fn(fnan, fzero) + fn(fzero, fnan) + fn(fone, fnan) + fn(fnan, fone) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_float_remainder_overflow(self): + a = np.finfo(np.float64).tiny + with np.errstate(over='ignore', invalid='ignore'): + div, mod = np.divmod(4, a) + np.isinf(div) + assert_(mod == 0) + with np.errstate(over='raise', invalid='ignore'): + assert_raises(FloatingPointError, np.divmod, 4, a) + with np.errstate(invalid='raise', over='ignore'): + assert_raises(FloatingPointError, np.divmod, 4, a) + + def test_float_divmod_corner_cases(self): + # check nan cases + for dt in np.typecodes['Float']: + fnan = np.array(np.nan, dtype=dt) + fone = np.array(1.0, dtype=dt) + fzer = np.array(0.0, dtype=dt) + finf = np.array(np.inf, dtype=dt) + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "invalid value encountered in divmod") + sup.filter(RuntimeWarning, "divide by zero encountered in divmod") + div, rem = np.divmod(fone, fzer) + assert(np.isinf(div)), 'dt: %s, div: %s' % (dt, rem) + assert(np.isnan(rem)), 'dt: %s, rem: %s' % (dt, rem) + div, rem = np.divmod(fzer, fzer) + assert(np.isnan(rem)), 'dt: %s, rem: %s' % (dt, rem) + assert_(np.isnan(div)), 'dt: %s, rem: %s' % (dt, rem) + div, rem = np.divmod(finf, finf) + assert(np.isnan(div)), 'dt: %s, rem: %s' % (dt, rem) + assert(np.isnan(rem)), 'dt: %s, rem: %s' % (dt, rem) + div, rem = np.divmod(finf, fzer) + assert(np.isinf(div)), 'dt: %s, rem: %s' % (dt, rem) + assert(np.isnan(rem)), 'dt: %s, rem: %s' % (dt, rem) + div, rem = np.divmod(fnan, fone) + assert(np.isnan(rem)), "dt: %s, rem: %s" % (dt, rem) + assert(np.isnan(div)), "dt: %s, rem: %s" % (dt, rem) + div, rem = np.divmod(fone, fnan) + assert(np.isnan(rem)), "dt: %s, rem: %s" % (dt, rem) + assert(np.isnan(div)), "dt: %s, rem: %s" % (dt, rem) + div, rem = np.divmod(fnan, fzer) + assert(np.isnan(rem)), "dt: %s, rem: %s" % (dt, rem) + assert(np.isnan(div)), "dt: %s, rem: %s" % (dt, rem) + + def test_float_remainder_corner_cases(self): + # Check remainder magnitude. + for dt in np.typecodes['Float']: + fone = np.array(1.0, dtype=dt) + fzer = np.array(0.0, dtype=dt) + fnan = np.array(np.nan, dtype=dt) + b = np.array(1.0, dtype=dt) + a = np.nextafter(np.array(0.0, dtype=dt), -b) + rem = np.remainder(a, b) + assert_(rem <= b, 'dt: %s' % dt) + rem = np.remainder(-a, -b) + assert_(rem >= -b, 'dt: %s' % dt) + + # Check nans, inf + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "invalid value encountered in remainder") + sup.filter(RuntimeWarning, "invalid value encountered in fmod") + for dt in np.typecodes['Float']: + fone = np.array(1.0, dtype=dt) + fzer = np.array(0.0, dtype=dt) + finf = np.array(np.inf, dtype=dt) + fnan = np.array(np.nan, dtype=dt) + rem = np.remainder(fone, fzer) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + # MSVC 2008 returns NaN here, so disable the check. + #rem = np.remainder(fone, finf) + #assert_(rem == fone, 'dt: %s, rem: %s' % (dt, rem)) + rem = np.remainder(finf, fone) + fmod = np.fmod(finf, fone) + assert_(np.isnan(fmod), 'dt: %s, fmod: %s' % (dt, fmod)) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + rem = np.remainder(finf, finf) + fmod = np.fmod(finf, fone) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + assert_(np.isnan(fmod), 'dt: %s, fmod: %s' % (dt, fmod)) + rem = np.remainder(finf, fzer) + fmod = np.fmod(finf, fzer) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + assert_(np.isnan(fmod), 'dt: %s, fmod: %s' % (dt, fmod)) + rem = np.remainder(fone, fnan) + fmod = np.fmod(fone, fnan) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + assert_(np.isnan(fmod), 'dt: %s, fmod: %s' % (dt, fmod)) + rem = np.remainder(fnan, fzer) + fmod = np.fmod(fnan, fzer) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + assert_(np.isnan(fmod), 'dt: %s, fmod: %s' % (dt, rem)) + rem = np.remainder(fnan, fone) + fmod = np.fmod(fnan, fone) + assert_(np.isnan(rem), 'dt: %s, rem: %s' % (dt, rem)) + assert_(np.isnan(fmod), 'dt: %s, fmod: %s' % (dt, rem)) + + +class TestDivisionIntegerOverflowsAndDivideByZero: + result_type = namedtuple('result_type', + ['nocast', 'casted']) + helper_lambdas = { + 'zero': lambda dtype: 0, + 'min': lambda dtype: np.iinfo(dtype).min, + 'neg_min': lambda dtype: -np.iinfo(dtype).min, + 'min-zero': lambda dtype: (np.iinfo(dtype).min, 0), + 'neg_min-zero': lambda dtype: (-np.iinfo(dtype).min, 0), + } + overflow_results = { + np.remainder: result_type( + helper_lambdas['zero'], helper_lambdas['zero']), + np.fmod: result_type( + helper_lambdas['zero'], helper_lambdas['zero']), + operator.mod: result_type( + helper_lambdas['zero'], helper_lambdas['zero']), + operator.floordiv: result_type( + helper_lambdas['min'], helper_lambdas['neg_min']), + np.floor_divide: result_type( + helper_lambdas['min'], helper_lambdas['neg_min']), + np.divmod: result_type( + helper_lambdas['min-zero'], helper_lambdas['neg_min-zero']) + } + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize("dtype", np.typecodes["Integer"]) + def test_signed_division_overflow(self, dtype): + to_check = interesting_binop_operands(np.iinfo(dtype).min, -1, dtype) + for op1, op2, extractor, operand_identifier in to_check: + with pytest.warns(RuntimeWarning, match="overflow encountered"): + res = op1 // op2 + + assert res.dtype == op1.dtype + assert extractor(res) == np.iinfo(op1.dtype).min + + # Remainder is well defined though, and does not warn: + res = op1 % op2 + assert res.dtype == op1.dtype + assert extractor(res) == 0 + # Check fmod as well: + res = np.fmod(op1, op2) + assert extractor(res) == 0 + + # Divmod warns for the division part: + with pytest.warns(RuntimeWarning, match="overflow encountered"): + res1, res2 = np.divmod(op1, op2) + + assert res1.dtype == res2.dtype == op1.dtype + assert extractor(res1) == np.iinfo(op1.dtype).min + assert extractor(res2) == 0 + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize("dtype", np.typecodes["AllInteger"]) + def test_divide_by_zero(self, dtype): + # Note that the return value cannot be well defined here, but NumPy + # currently uses 0 consistently. This could be changed. + to_check = interesting_binop_operands(1, 0, dtype) + for op1, op2, extractor, operand_identifier in to_check: + with pytest.warns(RuntimeWarning, match="divide by zero"): + res = op1 // op2 + + assert res.dtype == op1.dtype + assert extractor(res) == 0 + + with pytest.warns(RuntimeWarning, match="divide by zero"): + res1, res2 = np.divmod(op1, op2) + + assert res1.dtype == res2.dtype == op1.dtype + assert extractor(res1) == 0 + assert extractor(res2) == 0 + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize("dividend_dtype", sctypes['int']) + @pytest.mark.parametrize("divisor_dtype", sctypes['int']) + @pytest.mark.parametrize("operation", + [np.remainder, np.fmod, np.divmod, np.floor_divide, + operator.mod, operator.floordiv]) + @np.errstate(divide='warn', over='warn') + def test_overflows(self, dividend_dtype, divisor_dtype, operation): + # SIMD tries to perform the operation on as many elements as possible + # that is a multiple of the register's size. We resort to the + # default implementation for the leftover elements. + # We try to cover all paths here. + arrays = [np.array([np.iinfo(dividend_dtype).min]*i, + dtype=dividend_dtype) for i in range(1, 129)] + divisor = np.array([-1], dtype=divisor_dtype) + # If dividend is a larger type than the divisor (`else` case), + # then, result will be a larger type than dividend and will not + # result in an overflow for `divmod` and `floor_divide`. + if np.dtype(dividend_dtype).itemsize >= np.dtype( + divisor_dtype).itemsize and operation in ( + np.divmod, np.floor_divide, operator.floordiv): + with pytest.warns( + RuntimeWarning, + match="overflow encountered in"): + result = operation( + dividend_dtype(np.iinfo(dividend_dtype).min), + divisor_dtype(-1) + ) + assert result == self.overflow_results[operation].nocast( + dividend_dtype) + + # Arrays + for a in arrays: + # In case of divmod, we need to flatten the result + # column first as we get a column vector of quotient and + # remainder and a normal flatten of the expected result. + with pytest.warns( + RuntimeWarning, + match="overflow encountered in"): + result = np.array(operation(a, divisor)).flatten('f') + expected_array = np.array( + [self.overflow_results[operation].nocast( + dividend_dtype)]*len(a)).flatten() + assert_array_equal(result, expected_array) + else: + # Scalars + result = operation( + dividend_dtype(np.iinfo(dividend_dtype).min), + divisor_dtype(-1) + ) + assert result == self.overflow_results[operation].casted( + dividend_dtype) + + # Arrays + for a in arrays: + # See above comment on flatten + result = np.array(operation(a, divisor)).flatten('f') + expected_array = np.array( + [self.overflow_results[operation].casted( + dividend_dtype)]*len(a)).flatten() + assert_array_equal(result, expected_array) + + +class TestCbrt: + def test_cbrt_scalar(self): + assert_almost_equal((np.cbrt(np.float32(-2.5)**3)), -2.5) + + def test_cbrt(self): + x = np.array([1., 2., -3., np.inf, -np.inf]) + assert_almost_equal(np.cbrt(x**3), x) + + assert_(np.isnan(np.cbrt(np.nan))) + assert_equal(np.cbrt(np.inf), np.inf) + assert_equal(np.cbrt(-np.inf), -np.inf) + + +class TestPower: + def test_power_float(self): + x = np.array([1., 2., 3.]) + assert_equal(x**0, [1., 1., 1.]) + assert_equal(x**1, x) + assert_equal(x**2, [1., 4., 9.]) + y = x.copy() + y **= 2 + assert_equal(y, [1., 4., 9.]) + assert_almost_equal(x**(-1), [1., 0.5, 1./3]) + assert_almost_equal(x**(0.5), [1., ncu.sqrt(2), ncu.sqrt(3)]) + + for out, inp, msg in _gen_alignment_data(dtype=np.float32, + type='unary', + max_size=11): + exp = [ncu.sqrt(i) for i in inp] + assert_almost_equal(inp**(0.5), exp, err_msg=msg) + np.sqrt(inp, out=out) + assert_equal(out, exp, err_msg=msg) + + for out, inp, msg in _gen_alignment_data(dtype=np.float64, + type='unary', + max_size=7): + exp = [ncu.sqrt(i) for i in inp] + assert_almost_equal(inp**(0.5), exp, err_msg=msg) + np.sqrt(inp, out=out) + assert_equal(out, exp, err_msg=msg) + + def test_power_complex(self): + x = np.array([1+2j, 2+3j, 3+4j]) + assert_equal(x**0, [1., 1., 1.]) + assert_equal(x**1, x) + assert_almost_equal(x**2, [-3+4j, -5+12j, -7+24j]) + assert_almost_equal(x**3, [(1+2j)**3, (2+3j)**3, (3+4j)**3]) + assert_almost_equal(x**4, [(1+2j)**4, (2+3j)**4, (3+4j)**4]) + assert_almost_equal(x**(-1), [1/(1+2j), 1/(2+3j), 1/(3+4j)]) + assert_almost_equal(x**(-2), [1/(1+2j)**2, 1/(2+3j)**2, 1/(3+4j)**2]) + assert_almost_equal(x**(-3), [(-11+2j)/125, (-46-9j)/2197, + (-117-44j)/15625]) + assert_almost_equal(x**(0.5), [ncu.sqrt(1+2j), ncu.sqrt(2+3j), + ncu.sqrt(3+4j)]) + norm = 1./((x**14)[0]) + assert_almost_equal(x**14 * norm, + [i * norm for i in [-76443+16124j, 23161315+58317492j, + 5583548873 + 2465133864j]]) + + # Ticket #836 + def assert_complex_equal(x, y): + assert_array_equal(x.real, y.real) + assert_array_equal(x.imag, y.imag) + + for z in [complex(0, np.inf), complex(1, np.inf)]: + z = np.array([z], dtype=np.complex128) + with np.errstate(invalid="ignore"): + assert_complex_equal(z**1, z) + assert_complex_equal(z**2, z*z) + assert_complex_equal(z**3, z*z*z) + + def test_power_zero(self): + # ticket #1271 + zero = np.array([0j]) + one = np.array([1+0j]) + cnan = np.array([complex(np.nan, np.nan)]) + # FIXME cinf not tested. + #cinf = np.array([complex(np.inf, 0)]) + + def assert_complex_equal(x, y): + x, y = np.asarray(x), np.asarray(y) + assert_array_equal(x.real, y.real) + assert_array_equal(x.imag, y.imag) + + # positive powers + for p in [0.33, 0.5, 1, 1.5, 2, 3, 4, 5, 6.6]: + assert_complex_equal(np.power(zero, p), zero) + + # zero power + assert_complex_equal(np.power(zero, 0), one) + with np.errstate(invalid="ignore"): + assert_complex_equal(np.power(zero, 0+1j), cnan) + + # negative power + for p in [0.33, 0.5, 1, 1.5, 2, 3, 4, 5, 6.6]: + assert_complex_equal(np.power(zero, -p), cnan) + assert_complex_equal(np.power(zero, -1+0.2j), cnan) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_zero_power_nonzero(self): + # Testing 0^{Non-zero} issue 18378 + zero = np.array([0.0+0.0j]) + cnan = np.array([complex(np.nan, np.nan)]) + + def assert_complex_equal(x, y): + assert_array_equal(x.real, y.real) + assert_array_equal(x.imag, y.imag) + + #Complex powers with positive real part will not generate a warning + assert_complex_equal(np.power(zero, 1+4j), zero) + assert_complex_equal(np.power(zero, 2-3j), zero) + #Testing zero values when real part is greater than zero + assert_complex_equal(np.power(zero, 1+1j), zero) + assert_complex_equal(np.power(zero, 1+0j), zero) + assert_complex_equal(np.power(zero, 1-1j), zero) + #Complex powers will negative real part or 0 (provided imaginary + # part is not zero) will generate a NAN and hence a RUNTIME warning + with pytest.warns(expected_warning=RuntimeWarning) as r: + assert_complex_equal(np.power(zero, -1+1j), cnan) + assert_complex_equal(np.power(zero, -2-3j), cnan) + assert_complex_equal(np.power(zero, -7+0j), cnan) + assert_complex_equal(np.power(zero, 0+1j), cnan) + assert_complex_equal(np.power(zero, 0-1j), cnan) + assert len(r) == 5 + + def test_fast_power(self): + x = np.array([1, 2, 3], np.int16) + res = x**2.0 + assert_((x**2.00001).dtype is res.dtype) + assert_array_equal(res, [1, 4, 9]) + # check the inplace operation on the casted copy doesn't mess with x + assert_(not np.may_share_memory(res, x)) + assert_array_equal(x, [1, 2, 3]) + + # Check that the fast path ignores 1-element not 0-d arrays + res = x ** np.array([[[2]]]) + assert_equal(res.shape, (1, 1, 3)) + + def test_integer_power(self): + a = np.array([15, 15], 'i8') + b = np.power(a, a) + assert_equal(b, [437893890380859375, 437893890380859375]) + + def test_integer_power_with_integer_zero_exponent(self): + dtypes = np.typecodes['Integer'] + for dt in dtypes: + arr = np.arange(-10, 10, dtype=dt) + assert_equal(np.power(arr, 0), np.ones_like(arr)) + + dtypes = np.typecodes['UnsignedInteger'] + for dt in dtypes: + arr = np.arange(10, dtype=dt) + assert_equal(np.power(arr, 0), np.ones_like(arr)) + + def test_integer_power_of_1(self): + dtypes = np.typecodes['AllInteger'] + for dt in dtypes: + arr = np.arange(10, dtype=dt) + assert_equal(np.power(1, arr), np.ones_like(arr)) + + def test_integer_power_of_zero(self): + dtypes = np.typecodes['AllInteger'] + for dt in dtypes: + arr = np.arange(1, 10, dtype=dt) + assert_equal(np.power(0, arr), np.zeros_like(arr)) + + def test_integer_to_negative_power(self): + dtypes = np.typecodes['Integer'] + for dt in dtypes: + a = np.array([0, 1, 2, 3], dtype=dt) + b = np.array([0, 1, 2, -3], dtype=dt) + one = np.array(1, dtype=dt) + minusone = np.array(-1, dtype=dt) + assert_raises(ValueError, np.power, a, b) + assert_raises(ValueError, np.power, a, minusone) + assert_raises(ValueError, np.power, one, b) + assert_raises(ValueError, np.power, one, minusone) + + def test_float_to_inf_power(self): + for dt in [np.float32, np.float64]: + a = np.array([1, 1, 2, 2, -2, -2, np.inf, -np.inf], dt) + b = np.array([np.inf, -np.inf, np.inf, -np.inf, + np.inf, -np.inf, np.inf, -np.inf], dt) + r = np.array([1, 1, np.inf, 0, np.inf, 0, np.inf, 0], dt) + assert_equal(np.power(a, b), r) + + def test_power_fast_paths(self): + # gh-26055 + for dt in [np.float32, np.float64]: + a = np.array([0, 1.1, 2, 12e12, -10., np.inf, -np.inf], dt) + expected = np.array([0.0, 1.21, 4., 1.44e+26, 100, np.inf, np.inf]) + result = np.power(a, 2.) + assert_array_max_ulp(result, expected.astype(dt), maxulp=1) + + a = np.array([0, 1.1, 2, 12e12], dt) + expected = np.sqrt(a).astype(dt) + result = np.power(a, 0.5) + assert_array_max_ulp(result, expected, maxulp=1) + + +class TestFloat_power: + def test_type_conversion(self): + arg_type = '?bhilBHILefdgFDG' + res_type = 'ddddddddddddgDDG' + for dtin, dtout in zip(arg_type, res_type): + msg = "dtin: %s, dtout: %s" % (dtin, dtout) + arg = np.ones(1, dtype=dtin) + res = np.float_power(arg, arg) + assert_(res.dtype.name == np.dtype(dtout).name, msg) + + +class TestLog2: + @pytest.mark.parametrize('dt', ['f', 'd', 'g']) + def test_log2_values(self, dt): + x = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_almost_equal(np.log2(xf), yf) + + @pytest.mark.parametrize("i", range(1, 65)) + def test_log2_ints(self, i): + # a good log2 implementation should provide this, + # might fail on OS with bad libm + v = np.log2(2.**i) + assert_equal(v, float(i), err_msg='at exponent %d' % i) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_log2_special(self): + assert_equal(np.log2(1.), 0.) + assert_equal(np.log2(np.inf), np.inf) + assert_(np.isnan(np.log2(np.nan))) + + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_(np.isnan(np.log2(-1.))) + assert_(np.isnan(np.log2(-np.inf))) + assert_equal(np.log2(0.), -np.inf) + assert_(w[0].category is RuntimeWarning) + assert_(w[1].category is RuntimeWarning) + assert_(w[2].category is RuntimeWarning) + + +class TestExp2: + def test_exp2_values(self): + x = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + for dt in ['f', 'd', 'g']: + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_almost_equal(np.exp2(yf), xf) + + +class TestLogAddExp2(_FilterInvalids): + # Need test for intermediate precisions + def test_logaddexp2_values(self): + x = [1, 2, 3, 4, 5] + y = [5, 4, 3, 2, 1] + z = [6, 6, 6, 6, 6] + for dt, dec_ in zip(['f', 'd', 'g'], [6, 15, 15]): + xf = np.log2(np.array(x, dtype=dt)) + yf = np.log2(np.array(y, dtype=dt)) + zf = np.log2(np.array(z, dtype=dt)) + assert_almost_equal(np.logaddexp2(xf, yf), zf, decimal=dec_) + + def test_logaddexp2_range(self): + x = [1000000, -1000000, 1000200, -1000200] + y = [1000200, -1000200, 1000000, -1000000] + z = [1000200, -1000000, 1000200, -1000000] + for dt in ['f', 'd', 'g']: + logxf = np.array(x, dtype=dt) + logyf = np.array(y, dtype=dt) + logzf = np.array(z, dtype=dt) + assert_almost_equal(np.logaddexp2(logxf, logyf), logzf) + + def test_inf(self): + inf = np.inf + x = [inf, -inf, inf, -inf, inf, 1, -inf, 1] + y = [inf, inf, -inf, -inf, 1, inf, 1, -inf] + z = [inf, inf, inf, -inf, inf, inf, 1, 1] + with np.errstate(invalid='raise'): + for dt in ['f', 'd', 'g']: + logxf = np.array(x, dtype=dt) + logyf = np.array(y, dtype=dt) + logzf = np.array(z, dtype=dt) + assert_equal(np.logaddexp2(logxf, logyf), logzf) + + def test_nan(self): + assert_(np.isnan(np.logaddexp2(np.nan, np.inf))) + assert_(np.isnan(np.logaddexp2(np.inf, np.nan))) + assert_(np.isnan(np.logaddexp2(np.nan, 0))) + assert_(np.isnan(np.logaddexp2(0, np.nan))) + assert_(np.isnan(np.logaddexp2(np.nan, np.nan))) + + def test_reduce(self): + assert_equal(np.logaddexp2.identity, -np.inf) + assert_equal(np.logaddexp2.reduce([]), -np.inf) + assert_equal(np.logaddexp2.reduce([-np.inf]), -np.inf) + assert_equal(np.logaddexp2.reduce([-np.inf, 0]), 0) + + +class TestLog: + def test_log_values(self): + x = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + for dt in ['f', 'd', 'g']: + log2_ = 0.69314718055994530943 + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt)*log2_ + assert_almost_equal(np.log(xf), yf) + + # test aliasing(issue #17761) + x = np.array([2, 0.937500, 3, 0.947500, 1.054697]) + xf = np.log(x) + assert_almost_equal(np.log(x, out=x), xf) + + def test_log_values_maxofdtype(self): + # test log() of max for dtype does not raise + dtypes = [np.float32, np.float64] + # This is failing at least on linux aarch64 (see gh-25460), and on most + # other non x86-64 platforms checking `longdouble` isn't too useful as + # it's an alias for float64. + if platform.machine() == 'x86_64': + dtypes += [np.longdouble] + + for dt in dtypes: + with np.errstate(all='raise'): + x = np.finfo(dt).max + np.log(x) + + def test_log_strides(self): + np.random.seed(42) + strides = np.array([-4,-3,-2,-1,1,2,3,4]) + sizes = np.arange(2,100) + for ii in sizes: + x_f64 = np.float64(np.random.uniform(low=0.01, high=100.0,size=ii)) + x_special = x_f64.copy() + x_special[3:-1:4] = 1.0 + y_true = np.log(x_f64) + y_special = np.log(x_special) + for jj in strides: + assert_array_almost_equal_nulp(np.log(x_f64[::jj]), y_true[::jj], nulp=2) + assert_array_almost_equal_nulp(np.log(x_special[::jj]), y_special[::jj], nulp=2) + + # Reference values were computed with mpmath, with mp.dps = 200. + @pytest.mark.parametrize( + 'z, wref', + [(1 + 1e-12j, 5e-25 + 1e-12j), + (1.000000000000001 + 3e-08j, + 1.5602230246251546e-15 + 2.999999999999996e-08j), + (0.9999995000000417 + 0.0009999998333333417j, + 7.831475869017683e-18 + 0.001j), + (0.9999999999999996 + 2.999999999999999e-08j, + 5.9107901499372034e-18 + 3e-08j), + (0.99995000042 - 0.009999833j, + -7.015159763822903e-15 - 0.009999999665816696j)], + ) + def test_log_precision_float64(self, z, wref): + w = np.log(z) + assert_allclose(w, wref, rtol=1e-15) + + # Reference values were computed with mpmath, with mp.dps = 200. + @pytest.mark.parametrize( + 'z, wref', + [(np.complex64(1.0 + 3e-6j), np.complex64(4.5e-12+3e-06j)), + (np.complex64(1.0 - 2e-5j), np.complex64(1.9999999e-10 - 2e-5j)), + (np.complex64(0.9999999 + 1e-06j), + np.complex64(-1.192088e-07+1.0000001e-06j))], + ) + def test_log_precision_float32(self, z, wref): + w = np.log(z) + assert_allclose(w, wref, rtol=1e-6) + + +class TestExp: + def test_exp_values(self): + x = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + for dt in ['f', 'd', 'g']: + log2_ = 0.69314718055994530943 + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt)*log2_ + assert_almost_equal(np.exp(yf), xf) + + def test_exp_strides(self): + np.random.seed(42) + strides = np.array([-4,-3,-2,-1,1,2,3,4]) + sizes = np.arange(2,100) + for ii in sizes: + x_f64 = np.float64(np.random.uniform(low=0.01, high=709.1,size=ii)) + y_true = np.exp(x_f64) + for jj in strides: + assert_array_almost_equal_nulp(np.exp(x_f64[::jj]), y_true[::jj], nulp=2) + +class TestSpecialFloats: + def test_exp_values(self): + with np.errstate(under='raise', over='raise'): + x = [np.nan, np.nan, np.inf, 0.] + y = [np.nan, -np.nan, np.inf, -np.inf] + for dt in ['e', 'f', 'd', 'g']: + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_equal(np.exp(yf), xf) + + # See: https://github.com/numpy/numpy/issues/19192 + @pytest.mark.xfail( + _glibc_older_than("2.17"), + reason="Older glibc versions may not raise appropriate FP exceptions" + ) + def test_exp_exceptions(self): + with np.errstate(over='raise'): + assert_raises(FloatingPointError, np.exp, np.float16(11.0899)) + assert_raises(FloatingPointError, np.exp, np.float32(100.)) + assert_raises(FloatingPointError, np.exp, np.float32(1E19)) + assert_raises(FloatingPointError, np.exp, np.float64(800.)) + assert_raises(FloatingPointError, np.exp, np.float64(1E19)) + + with np.errstate(under='raise'): + assert_raises(FloatingPointError, np.exp, np.float16(-17.5)) + assert_raises(FloatingPointError, np.exp, np.float32(-1000.)) + assert_raises(FloatingPointError, np.exp, np.float32(-1E19)) + assert_raises(FloatingPointError, np.exp, np.float64(-1000.)) + assert_raises(FloatingPointError, np.exp, np.float64(-1E19)) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_log_values(self): + with np.errstate(all='ignore'): + x = [np.nan, np.nan, np.inf, np.nan, -np.inf, np.nan] + y = [np.nan, -np.nan, np.inf, -np.inf, 0.0, -1.0] + y1p = [np.nan, -np.nan, np.inf, -np.inf, -1.0, -2.0] + for dt in ['e', 'f', 'd', 'g']: + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + yf1p = np.array(y1p, dtype=dt) + assert_equal(np.log(yf), xf) + assert_equal(np.log2(yf), xf) + assert_equal(np.log10(yf), xf) + assert_equal(np.log1p(yf1p), xf) + + with np.errstate(divide='raise'): + for dt in ['e', 'f', 'd']: + assert_raises(FloatingPointError, np.log, + np.array(0.0, dtype=dt)) + assert_raises(FloatingPointError, np.log2, + np.array(0.0, dtype=dt)) + assert_raises(FloatingPointError, np.log10, + np.array(0.0, dtype=dt)) + assert_raises(FloatingPointError, np.log1p, + np.array(-1.0, dtype=dt)) + + with np.errstate(invalid='raise'): + for dt in ['e', 'f', 'd']: + assert_raises(FloatingPointError, np.log, + np.array(-np.inf, dtype=dt)) + assert_raises(FloatingPointError, np.log, + np.array(-1.0, dtype=dt)) + assert_raises(FloatingPointError, np.log2, + np.array(-np.inf, dtype=dt)) + assert_raises(FloatingPointError, np.log2, + np.array(-1.0, dtype=dt)) + assert_raises(FloatingPointError, np.log10, + np.array(-np.inf, dtype=dt)) + assert_raises(FloatingPointError, np.log10, + np.array(-1.0, dtype=dt)) + assert_raises(FloatingPointError, np.log1p, + np.array(-np.inf, dtype=dt)) + assert_raises(FloatingPointError, np.log1p, + np.array(-2.0, dtype=dt)) + + # See https://github.com/numpy/numpy/issues/18005 + with assert_no_warnings(): + a = np.array(1e9, dtype='float32') + np.log(a) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize('dtype', ['e', 'f', 'd', 'g']) + def test_sincos_values(self, dtype): + with np.errstate(all='ignore'): + x = [np.nan, np.nan, np.nan, np.nan] + y = [np.nan, -np.nan, np.inf, -np.inf] + xf = np.array(x, dtype=dtype) + yf = np.array(y, dtype=dtype) + assert_equal(np.sin(yf), xf) + assert_equal(np.cos(yf), xf) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.xfail( + sys.platform.startswith("darwin"), + reason="underflow is triggered for scalar 'sin'" + ) + def test_sincos_underflow(self): + with np.errstate(under='raise'): + underflow_trigger = np.array( + float.fromhex("0x1.f37f47a03f82ap-511"), + dtype=np.float64 + ) + np.sin(underflow_trigger) + np.cos(underflow_trigger) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.parametrize('callable', [np.sin, np.cos]) + @pytest.mark.parametrize('dtype', ['e', 'f', 'd']) + @pytest.mark.parametrize('value', [np.inf, -np.inf]) + def test_sincos_errors(self, callable, dtype, value): + with np.errstate(invalid='raise'): + assert_raises(FloatingPointError, callable, + np.array([value], dtype=dtype)) + + @pytest.mark.parametrize('callable', [np.sin, np.cos]) + @pytest.mark.parametrize('dtype', ['f', 'd']) + @pytest.mark.parametrize('stride', [-1, 1, 2, 4, 5]) + def test_sincos_overlaps(self, callable, dtype, stride): + N = 100 + M = N // abs(stride) + rng = np.random.default_rng(42) + x = rng.standard_normal(N, dtype) + y = callable(x[::stride]) + callable(x[::stride], out=x[:M]) + assert_equal(x[:M], y) + + @pytest.mark.parametrize('dt', ['e', 'f', 'd', 'g']) + def test_sqrt_values(self, dt): + with np.errstate(all='ignore'): + x = [np.nan, np.nan, np.inf, np.nan, 0.] + y = [np.nan, -np.nan, np.inf, -np.inf, 0.] + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_equal(np.sqrt(yf), xf) + + # with np.errstate(invalid='raise'): + # assert_raises( + # FloatingPointError, np.sqrt, np.array(-100., dtype=dt) + # ) + + def test_abs_values(self): + x = [np.nan, np.nan, np.inf, np.inf, 0., 0., 1.0, 1.0] + y = [np.nan, -np.nan, np.inf, -np.inf, 0., -0., -1.0, 1.0] + for dt in ['e', 'f', 'd', 'g']: + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_equal(np.abs(yf), xf) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_square_values(self): + x = [np.nan, np.nan, np.inf, np.inf] + y = [np.nan, -np.nan, np.inf, -np.inf] + with np.errstate(all='ignore'): + for dt in ['e', 'f', 'd', 'g']: + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_equal(np.square(yf), xf) + + with np.errstate(over='raise'): + assert_raises(FloatingPointError, np.square, + np.array(1E3, dtype='e')) + assert_raises(FloatingPointError, np.square, + np.array(1E32, dtype='f')) + assert_raises(FloatingPointError, np.square, + np.array(1E200, dtype='d')) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_reciprocal_values(self): + with np.errstate(all='ignore'): + x = [np.nan, np.nan, 0.0, -0.0, np.inf, -np.inf] + y = [np.nan, -np.nan, np.inf, -np.inf, 0., -0.] + for dt in ['e', 'f', 'd', 'g']: + xf = np.array(x, dtype=dt) + yf = np.array(y, dtype=dt) + assert_equal(np.reciprocal(yf), xf) + + with np.errstate(divide='raise'): + for dt in ['e', 'f', 'd', 'g']: + assert_raises(FloatingPointError, np.reciprocal, + np.array(-0.0, dtype=dt)) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_tan(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan, 0.0, -0.0, np.inf, -np.inf] + out = [np.nan, np.nan, 0.0, -0.0, np.nan, np.nan] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.tan(in_arr), out_arr) + + with np.errstate(invalid='raise'): + for dt in ['e', 'f', 'd']: + assert_raises(FloatingPointError, np.tan, + np.array(np.inf, dtype=dt)) + assert_raises(FloatingPointError, np.tan, + np.array(-np.inf, dtype=dt)) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_arcsincos(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, np.nan, np.nan] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.arcsin(in_arr), out_arr) + assert_equal(np.arccos(in_arr), out_arr) + + for callable in [np.arcsin, np.arccos]: + for value in [np.inf, -np.inf, 2.0, -2.0]: + for dt in ['e', 'f', 'd']: + with np.errstate(invalid='raise'): + assert_raises(FloatingPointError, callable, + np.array(value, dtype=dt)) + + def test_arctan(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan] + out = [np.nan, np.nan] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.arctan(in_arr), out_arr) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_sinh(self): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, np.inf, -np.inf] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.sinh(in_arr), out_arr) + + with np.errstate(over='raise'): + assert_raises(FloatingPointError, np.sinh, + np.array(12.0, dtype='e')) + assert_raises(FloatingPointError, np.sinh, + np.array(120.0, dtype='f')) + assert_raises(FloatingPointError, np.sinh, + np.array(1200.0, dtype='d')) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + @pytest.mark.skipif('bsd' in sys.platform, + reason="fallback implementation may not raise, see gh-2487") + def test_cosh(self): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, np.inf, np.inf] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.cosh(in_arr), out_arr) + + with np.errstate(over='raise'): + assert_raises(FloatingPointError, np.cosh, + np.array(12.0, dtype='e')) + assert_raises(FloatingPointError, np.cosh, + np.array(120.0, dtype='f')) + assert_raises(FloatingPointError, np.cosh, + np.array(1200.0, dtype='d')) + + def test_tanh(self): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, 1.0, -1.0] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_array_max_ulp(np.tanh(in_arr), out_arr, 3) + + def test_arcsinh(self): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, np.inf, -np.inf] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.arcsinh(in_arr), out_arr) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_arccosh(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan, np.inf, -np.inf, 1.0, 0.0] + out = [np.nan, np.nan, np.inf, np.nan, 0.0, np.nan] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.arccosh(in_arr), out_arr) + + for value in [0.0, -np.inf]: + with np.errstate(invalid='raise'): + for dt in ['e', 'f', 'd']: + assert_raises(FloatingPointError, np.arccosh, + np.array(value, dtype=dt)) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_arctanh(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan, np.inf, -np.inf, 1.0, -1.0, 2.0] + out = [np.nan, np.nan, np.nan, np.nan, np.inf, -np.inf, np.nan] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.arctanh(in_arr), out_arr) + + for value in [1.01, np.inf, -np.inf, 1.0, -1.0]: + with np.errstate(invalid='raise', divide='raise'): + for dt in ['e', 'f', 'd']: + assert_raises(FloatingPointError, np.arctanh, + np.array(value, dtype=dt)) + + # Make sure glibc < 2.18 atanh is not used, issue 25087 + assert np.signbit(np.arctanh(-1j).real) + + # See: https://github.com/numpy/numpy/issues/20448 + @pytest.mark.xfail( + _glibc_older_than("2.17"), + reason="Older glibc versions may not raise appropriate FP exceptions" + ) + def test_exp2(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, np.inf, 0.0] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.exp2(in_arr), out_arr) + + for value in [2000.0, -2000.0]: + with np.errstate(over='raise', under='raise'): + for dt in ['e', 'f', 'd']: + assert_raises(FloatingPointError, np.exp2, + np.array(value, dtype=dt)) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work in wasm") + def test_expm1(self): + with np.errstate(all='ignore'): + in_ = [np.nan, -np.nan, np.inf, -np.inf] + out = [np.nan, np.nan, np.inf, -1.0] + for dt in ['e', 'f', 'd']: + in_arr = np.array(in_, dtype=dt) + out_arr = np.array(out, dtype=dt) + assert_equal(np.expm1(in_arr), out_arr) + + for value in [200.0, 2000.0]: + with np.errstate(over='raise'): + for dt in ['e', 'f']: + assert_raises(FloatingPointError, np.expm1, + np.array(value, dtype=dt)) + + # test to ensure no spurious FP exceptions are raised due to SIMD + INF_INVALID_ERR = [ + np.cos, np.sin, np.tan, np.arccos, np.arcsin, np.spacing, np.arctanh + ] + NEG_INVALID_ERR = [ + np.log, np.log2, np.log10, np.log1p, np.sqrt, np.arccosh, + np.arctanh + ] + ONE_INVALID_ERR = [ + np.arctanh, + ] + LTONE_INVALID_ERR = [ + np.arccosh, + ] + BYZERO_ERR = [ + np.log, np.log2, np.log10, np.reciprocal, np.arccosh + ] + + @pytest.mark.parametrize("ufunc", UFUNCS_UNARY_FP) + @pytest.mark.parametrize("dtype", ('e', 'f', 'd')) + @pytest.mark.parametrize("data, escape", ( + ([0.03], LTONE_INVALID_ERR), + ([0.03]*32, LTONE_INVALID_ERR), + # neg + ([-1.0], NEG_INVALID_ERR), + ([-1.0]*32, NEG_INVALID_ERR), + # flat + ([1.0], ONE_INVALID_ERR), + ([1.0]*32, ONE_INVALID_ERR), + # zero + ([0.0], BYZERO_ERR), + ([0.0]*32, BYZERO_ERR), + ([-0.0], BYZERO_ERR), + ([-0.0]*32, BYZERO_ERR), + # nan + ([0.5, 0.5, 0.5, np.nan], LTONE_INVALID_ERR), + ([0.5, 0.5, 0.5, np.nan]*32, LTONE_INVALID_ERR), + ([np.nan, 1.0, 1.0, 1.0], ONE_INVALID_ERR), + ([np.nan, 1.0, 1.0, 1.0]*32, ONE_INVALID_ERR), + ([np.nan], []), + ([np.nan]*32, []), + # inf + ([0.5, 0.5, 0.5, np.inf], INF_INVALID_ERR + LTONE_INVALID_ERR), + ([0.5, 0.5, 0.5, np.inf]*32, INF_INVALID_ERR + LTONE_INVALID_ERR), + ([np.inf, 1.0, 1.0, 1.0], INF_INVALID_ERR), + ([np.inf, 1.0, 1.0, 1.0]*32, INF_INVALID_ERR), + ([np.inf], INF_INVALID_ERR), + ([np.inf]*32, INF_INVALID_ERR), + # ninf + ([0.5, 0.5, 0.5, -np.inf], + NEG_INVALID_ERR + INF_INVALID_ERR + LTONE_INVALID_ERR), + ([0.5, 0.5, 0.5, -np.inf]*32, + NEG_INVALID_ERR + INF_INVALID_ERR + LTONE_INVALID_ERR), + ([-np.inf, 1.0, 1.0, 1.0], NEG_INVALID_ERR + INF_INVALID_ERR), + ([-np.inf, 1.0, 1.0, 1.0]*32, NEG_INVALID_ERR + INF_INVALID_ERR), + ([-np.inf], NEG_INVALID_ERR + INF_INVALID_ERR), + ([-np.inf]*32, NEG_INVALID_ERR + INF_INVALID_ERR), + )) + def test_unary_spurious_fpexception(self, ufunc, dtype, data, escape): + if escape and ufunc in escape: + return + # FIXME: NAN raises FP invalid exception: + # - ceil/float16 on MSVC:32-bit + # - spacing/float16 on almost all platforms + if ufunc in (np.spacing, np.ceil) and dtype == 'e': + return + array = np.array(data, dtype=dtype) + with assert_no_warnings(): + ufunc(array) + + @pytest.mark.parametrize("dtype", ('e', 'f', 'd')) + def test_divide_spurious_fpexception(self, dtype): + dt = np.dtype(dtype) + dt_info = np.finfo(dt) + subnorm = dt_info.smallest_subnormal + # Verify a bug fix caused due to filling the remaining lanes of the + # partially loaded dividend SIMD vector with ones, which leads to + # raising an overflow warning when the divisor is denormal. + # see https://github.com/numpy/numpy/issues/25097 + with assert_no_warnings(): + np.zeros(128 + 1, dtype=dt) / subnorm + +class TestFPClass: + @pytest.mark.parametrize("stride", [-5, -4, -3, -2, -1, 1, + 2, 4, 5, 6, 7, 8, 9, 10]) + def test_fpclass(self, stride): + arr_f64 = np.array([np.nan, -np.nan, np.inf, -np.inf, -1.0, 1.0, -0.0, 0.0, 2.2251e-308, -2.2251e-308], dtype='d') + arr_f32 = np.array([np.nan, -np.nan, np.inf, -np.inf, -1.0, 1.0, -0.0, 0.0, 1.4013e-045, -1.4013e-045], dtype='f') + nan = np.array([True, True, False, False, False, False, False, False, False, False]) + inf = np.array([False, False, True, True, False, False, False, False, False, False]) + sign = np.array([False, True, False, True, True, False, True, False, False, True]) + finite = np.array([False, False, False, False, True, True, True, True, True, True]) + assert_equal(np.isnan(arr_f32[::stride]), nan[::stride]) + assert_equal(np.isnan(arr_f64[::stride]), nan[::stride]) + assert_equal(np.isinf(arr_f32[::stride]), inf[::stride]) + assert_equal(np.isinf(arr_f64[::stride]), inf[::stride]) + if platform.machine() == 'riscv64': + # On RISC-V, many operations that produce NaNs, such as converting + # a -NaN from f64 to f32, return a canonical NaN. The canonical + # NaNs are always positive. See section 11.3 NaN Generation and + # Propagation of the RISC-V Unprivileged ISA for more details. + # We disable the sign test on riscv64 for -np.nan as we + # cannot assume that its sign will be honoured in these tests. + arr_f64_rv = np.copy(arr_f64) + arr_f32_rv = np.copy(arr_f32) + arr_f64_rv[1] = -1.0 + arr_f32_rv[1] = -1.0 + assert_equal(np.signbit(arr_f32_rv[::stride]), sign[::stride]) + assert_equal(np.signbit(arr_f64_rv[::stride]), sign[::stride]) + else: + assert_equal(np.signbit(arr_f32[::stride]), sign[::stride]) + assert_equal(np.signbit(arr_f64[::stride]), sign[::stride]) + assert_equal(np.isfinite(arr_f32[::stride]), finite[::stride]) + assert_equal(np.isfinite(arr_f64[::stride]), finite[::stride]) + + @pytest.mark.parametrize("dtype", ['d', 'f']) + def test_fp_noncontiguous(self, dtype): + data = np.array([np.nan, -np.nan, np.inf, -np.inf, -1.0, + 1.0, -0.0, 0.0, 2.2251e-308, + -2.2251e-308], dtype=dtype) + nan = np.array([True, True, False, False, False, False, + False, False, False, False]) + inf = np.array([False, False, True, True, False, False, + False, False, False, False]) + sign = np.array([False, True, False, True, True, False, + True, False, False, True]) + finite = np.array([False, False, False, False, True, True, + True, True, True, True]) + out = np.ndarray(data.shape, dtype='bool') + ncontig_in = data[1::3] + ncontig_out = out[1::3] + contig_in = np.array(ncontig_in) + + if platform.machine() == 'riscv64': + # Disable the -np.nan signbit tests on riscv64. See comments in + # test_fpclass for more details. + data_rv = np.copy(data) + data_rv[1] = -1.0 + ncontig_sign_in = data_rv[1::3] + contig_sign_in = np.array(ncontig_sign_in) + else: + ncontig_sign_in = ncontig_in + contig_sign_in = contig_in + + assert_equal(ncontig_in.flags.c_contiguous, False) + assert_equal(ncontig_out.flags.c_contiguous, False) + assert_equal(contig_in.flags.c_contiguous, True) + assert_equal(ncontig_sign_in.flags.c_contiguous, False) + assert_equal(contig_sign_in.flags.c_contiguous, True) + # ncontig in, ncontig out + assert_equal(np.isnan(ncontig_in, out=ncontig_out), nan[1::3]) + assert_equal(np.isinf(ncontig_in, out=ncontig_out), inf[1::3]) + assert_equal(np.signbit(ncontig_sign_in, out=ncontig_out), sign[1::3]) + assert_equal(np.isfinite(ncontig_in, out=ncontig_out), finite[1::3]) + # contig in, ncontig out + assert_equal(np.isnan(contig_in, out=ncontig_out), nan[1::3]) + assert_equal(np.isinf(contig_in, out=ncontig_out), inf[1::3]) + assert_equal(np.signbit(contig_sign_in, out=ncontig_out), sign[1::3]) + assert_equal(np.isfinite(contig_in, out=ncontig_out), finite[1::3]) + # ncontig in, contig out + assert_equal(np.isnan(ncontig_in), nan[1::3]) + assert_equal(np.isinf(ncontig_in), inf[1::3]) + assert_equal(np.signbit(ncontig_sign_in), sign[1::3]) + assert_equal(np.isfinite(ncontig_in), finite[1::3]) + # contig in, contig out, nd stride + data_split = np.array(np.array_split(data, 2)) + nan_split = np.array(np.array_split(nan, 2)) + inf_split = np.array(np.array_split(inf, 2)) + sign_split = np.array(np.array_split(sign, 2)) + finite_split = np.array(np.array_split(finite, 2)) + assert_equal(np.isnan(data_split), nan_split) + assert_equal(np.isinf(data_split), inf_split) + if platform.machine() == 'riscv64': + data_split_rv = np.array(np.array_split(data_rv, 2)) + assert_equal(np.signbit(data_split_rv), sign_split) + else: + assert_equal(np.signbit(data_split), sign_split) + assert_equal(np.isfinite(data_split), finite_split) + +class TestLDExp: + @pytest.mark.parametrize("stride", [-4,-2,-1,1,2,4]) + @pytest.mark.parametrize("dtype", ['f', 'd']) + def test_ldexp(self, dtype, stride): + mant = np.array([0.125, 0.25, 0.5, 1., 1., 2., 4., 8.], dtype=dtype) + exp = np.array([3, 2, 1, 0, 0, -1, -2, -3], dtype='i') + out = np.zeros(8, dtype=dtype) + assert_equal(np.ldexp(mant[::stride], exp[::stride], out=out[::stride]), np.ones(8, dtype=dtype)[::stride]) + assert_equal(out[::stride], np.ones(8, dtype=dtype)[::stride]) + +class TestFRExp: + @pytest.mark.parametrize("stride", [-4,-2,-1,1,2,4]) + @pytest.mark.parametrize("dtype", ['f', 'd']) + @pytest.mark.skipif(not sys.platform.startswith('linux'), + reason="np.frexp gives different answers for NAN/INF on windows and linux") + @pytest.mark.xfail(IS_MUSL, reason="gh23049") + def test_frexp(self, dtype, stride): + arr = np.array([np.nan, np.nan, np.inf, -np.inf, 0.0, -0.0, 1.0, -1.0], dtype=dtype) + mant_true = np.array([np.nan, np.nan, np.inf, -np.inf, 0.0, -0.0, 0.5, -0.5], dtype=dtype) + exp_true = np.array([0, 0, 0, 0, 0, 0, 1, 1], dtype='i') + out_mant = np.ones(8, dtype=dtype) + out_exp = 2*np.ones(8, dtype='i') + mant, exp = np.frexp(arr[::stride], out=(out_mant[::stride], out_exp[::stride])) + assert_equal(mant_true[::stride], mant) + assert_equal(exp_true[::stride], exp) + assert_equal(out_mant[::stride], mant_true[::stride]) + assert_equal(out_exp[::stride], exp_true[::stride]) + +# func : [maxulperror, low, high] +avx_ufuncs = {'sqrt' :[1, 0., 100.], + 'absolute' :[0, -100., 100.], + 'reciprocal' :[1, 1., 100.], + 'square' :[1, -100., 100.], + 'rint' :[0, -100., 100.], + 'floor' :[0, -100., 100.], + 'ceil' :[0, -100., 100.], + 'trunc' :[0, -100., 100.]} + +class TestAVXUfuncs: + def test_avx_based_ufunc(self): + strides = np.array([-4,-3,-2,-1,1,2,3,4]) + np.random.seed(42) + for func, prop in avx_ufuncs.items(): + maxulperr = prop[0] + minval = prop[1] + maxval = prop[2] + # various array sizes to ensure masking in AVX is tested + for size in range(1,32): + myfunc = getattr(np, func) + x_f32 = np.random.uniform(low=minval, high=maxval, + size=size).astype(np.float32) + x_f64 = x_f32.astype(np.float64) + x_f128 = x_f32.astype(np.longdouble) + y_true128 = myfunc(x_f128) + if maxulperr == 0: + assert_equal(myfunc(x_f32), y_true128.astype(np.float32)) + assert_equal(myfunc(x_f64), y_true128.astype(np.float64)) + else: + assert_array_max_ulp(myfunc(x_f32), + y_true128.astype(np.float32), + maxulp=maxulperr) + assert_array_max_ulp(myfunc(x_f64), + y_true128.astype(np.float64), + maxulp=maxulperr) + # various strides to test gather instruction + if size > 1: + y_true32 = myfunc(x_f32) + y_true64 = myfunc(x_f64) + for jj in strides: + assert_equal(myfunc(x_f64[::jj]), y_true64[::jj]) + assert_equal(myfunc(x_f32[::jj]), y_true32[::jj]) + +class TestAVXFloat32Transcendental: + def test_exp_float32(self): + np.random.seed(42) + x_f32 = np.float32(np.random.uniform(low=0.0,high=88.1,size=1000000)) + x_f64 = np.float64(x_f32) + assert_array_max_ulp(np.exp(x_f32), np.float32(np.exp(x_f64)), maxulp=3) + + def test_log_float32(self): + np.random.seed(42) + x_f32 = np.float32(np.random.uniform(low=0.0,high=1000,size=1000000)) + x_f64 = np.float64(x_f32) + assert_array_max_ulp(np.log(x_f32), np.float32(np.log(x_f64)), maxulp=4) + + def test_sincos_float32(self): + np.random.seed(42) + N = 1000000 + M = np.int_(N/20) + index = np.random.randint(low=0, high=N, size=M) + x_f32 = np.float32(np.random.uniform(low=-100.,high=100.,size=N)) + if not _glibc_older_than("2.17"): + # test coverage for elements > 117435.992f for which glibc is used + # this is known to be problematic on old glibc, so skip it there + x_f32[index] = np.float32(10E+10*np.random.rand(M)) + x_f64 = np.float64(x_f32) + assert_array_max_ulp(np.sin(x_f32), np.float32(np.sin(x_f64)), maxulp=2) + assert_array_max_ulp(np.cos(x_f32), np.float32(np.cos(x_f64)), maxulp=2) + # test aliasing(issue #17761) + tx_f32 = x_f32.copy() + assert_array_max_ulp(np.sin(x_f32, out=x_f32), np.float32(np.sin(x_f64)), maxulp=2) + assert_array_max_ulp(np.cos(tx_f32, out=tx_f32), np.float32(np.cos(x_f64)), maxulp=2) + + def test_strided_float32(self): + np.random.seed(42) + strides = np.array([-4,-3,-2,-1,1,2,3,4]) + sizes = np.arange(2,100) + for ii in sizes: + x_f32 = np.float32(np.random.uniform(low=0.01,high=88.1,size=ii)) + x_f32_large = x_f32.copy() + x_f32_large[3:-1:4] = 120000.0 + exp_true = np.exp(x_f32) + log_true = np.log(x_f32) + sin_true = np.sin(x_f32_large) + cos_true = np.cos(x_f32_large) + for jj in strides: + assert_array_almost_equal_nulp(np.exp(x_f32[::jj]), exp_true[::jj], nulp=2) + assert_array_almost_equal_nulp(np.log(x_f32[::jj]), log_true[::jj], nulp=2) + assert_array_almost_equal_nulp(np.sin(x_f32_large[::jj]), sin_true[::jj], nulp=2) + assert_array_almost_equal_nulp(np.cos(x_f32_large[::jj]), cos_true[::jj], nulp=2) + +class TestLogAddExp(_FilterInvalids): + def test_logaddexp_values(self): + x = [1, 2, 3, 4, 5] + y = [5, 4, 3, 2, 1] + z = [6, 6, 6, 6, 6] + for dt, dec_ in zip(['f', 'd', 'g'], [6, 15, 15]): + xf = np.log(np.array(x, dtype=dt)) + yf = np.log(np.array(y, dtype=dt)) + zf = np.log(np.array(z, dtype=dt)) + assert_almost_equal(np.logaddexp(xf, yf), zf, decimal=dec_) + + def test_logaddexp_range(self): + x = [1000000, -1000000, 1000200, -1000200] + y = [1000200, -1000200, 1000000, -1000000] + z = [1000200, -1000000, 1000200, -1000000] + for dt in ['f', 'd', 'g']: + logxf = np.array(x, dtype=dt) + logyf = np.array(y, dtype=dt) + logzf = np.array(z, dtype=dt) + assert_almost_equal(np.logaddexp(logxf, logyf), logzf) + + def test_inf(self): + inf = np.inf + x = [inf, -inf, inf, -inf, inf, 1, -inf, 1] + y = [inf, inf, -inf, -inf, 1, inf, 1, -inf] + z = [inf, inf, inf, -inf, inf, inf, 1, 1] + with np.errstate(invalid='raise'): + for dt in ['f', 'd', 'g']: + logxf = np.array(x, dtype=dt) + logyf = np.array(y, dtype=dt) + logzf = np.array(z, dtype=dt) + assert_equal(np.logaddexp(logxf, logyf), logzf) + + def test_nan(self): + assert_(np.isnan(np.logaddexp(np.nan, np.inf))) + assert_(np.isnan(np.logaddexp(np.inf, np.nan))) + assert_(np.isnan(np.logaddexp(np.nan, 0))) + assert_(np.isnan(np.logaddexp(0, np.nan))) + assert_(np.isnan(np.logaddexp(np.nan, np.nan))) + + def test_reduce(self): + assert_equal(np.logaddexp.identity, -np.inf) + assert_equal(np.logaddexp.reduce([]), -np.inf) + + +class TestLog1p: + def test_log1p(self): + assert_almost_equal(ncu.log1p(0.2), ncu.log(1.2)) + assert_almost_equal(ncu.log1p(1e-6), ncu.log(1+1e-6)) + + def test_special(self): + with np.errstate(invalid="ignore", divide="ignore"): + assert_equal(ncu.log1p(np.nan), np.nan) + assert_equal(ncu.log1p(np.inf), np.inf) + assert_equal(ncu.log1p(-1.), -np.inf) + assert_equal(ncu.log1p(-2.), np.nan) + assert_equal(ncu.log1p(-np.inf), np.nan) + + +class TestExpm1: + def test_expm1(self): + assert_almost_equal(ncu.expm1(0.2), ncu.exp(0.2)-1) + assert_almost_equal(ncu.expm1(1e-6), ncu.exp(1e-6)-1) + + def test_special(self): + assert_equal(ncu.expm1(np.inf), np.inf) + assert_equal(ncu.expm1(0.), 0.) + assert_equal(ncu.expm1(-0.), -0.) + assert_equal(ncu.expm1(np.inf), np.inf) + assert_equal(ncu.expm1(-np.inf), -1.) + + def test_complex(self): + x = np.asarray(1e-12) + assert_allclose(x, ncu.expm1(x)) + x = x.astype(np.complex128) + assert_allclose(x, ncu.expm1(x)) + + +class TestHypot: + def test_simple(self): + assert_almost_equal(ncu.hypot(1, 1), ncu.sqrt(2)) + assert_almost_equal(ncu.hypot(0, 0), 0) + + def test_reduce(self): + assert_almost_equal(ncu.hypot.reduce([3.0, 4.0]), 5.0) + assert_almost_equal(ncu.hypot.reduce([3.0, 4.0, 0]), 5.0) + assert_almost_equal(ncu.hypot.reduce([9.0, 12.0, 20.0]), 25.0) + assert_equal(ncu.hypot.reduce([]), 0.0) + + +def assert_hypot_isnan(x, y): + with np.errstate(invalid='ignore'): + assert_(np.isnan(ncu.hypot(x, y)), + "hypot(%s, %s) is %s, not nan" % (x, y, ncu.hypot(x, y))) + + +def assert_hypot_isinf(x, y): + with np.errstate(invalid='ignore'): + assert_(np.isinf(ncu.hypot(x, y)), + "hypot(%s, %s) is %s, not inf" % (x, y, ncu.hypot(x, y))) + + +class TestHypotSpecialValues: + def test_nan_outputs(self): + assert_hypot_isnan(np.nan, np.nan) + assert_hypot_isnan(np.nan, 1) + + def test_nan_outputs2(self): + assert_hypot_isinf(np.nan, np.inf) + assert_hypot_isinf(np.inf, np.nan) + assert_hypot_isinf(np.inf, 0) + assert_hypot_isinf(0, np.inf) + assert_hypot_isinf(np.inf, np.inf) + assert_hypot_isinf(np.inf, 23.0) + + def test_no_fpe(self): + assert_no_warnings(ncu.hypot, np.inf, 0) + + +def assert_arctan2_isnan(x, y): + assert_(np.isnan(ncu.arctan2(x, y)), "arctan(%s, %s) is %s, not nan" % (x, y, ncu.arctan2(x, y))) + + +def assert_arctan2_ispinf(x, y): + assert_((np.isinf(ncu.arctan2(x, y)) and ncu.arctan2(x, y) > 0), "arctan(%s, %s) is %s, not +inf" % (x, y, ncu.arctan2(x, y))) + + +def assert_arctan2_isninf(x, y): + assert_((np.isinf(ncu.arctan2(x, y)) and ncu.arctan2(x, y) < 0), "arctan(%s, %s) is %s, not -inf" % (x, y, ncu.arctan2(x, y))) + + +def assert_arctan2_ispzero(x, y): + assert_((ncu.arctan2(x, y) == 0 and not np.signbit(ncu.arctan2(x, y))), "arctan(%s, %s) is %s, not +0" % (x, y, ncu.arctan2(x, y))) + + +def assert_arctan2_isnzero(x, y): + assert_((ncu.arctan2(x, y) == 0 and np.signbit(ncu.arctan2(x, y))), "arctan(%s, %s) is %s, not -0" % (x, y, ncu.arctan2(x, y))) + + +class TestArctan2SpecialValues: + def test_one_one(self): + # atan2(1, 1) returns pi/4. + assert_almost_equal(ncu.arctan2(1, 1), 0.25 * np.pi) + assert_almost_equal(ncu.arctan2(-1, 1), -0.25 * np.pi) + assert_almost_equal(ncu.arctan2(1, -1), 0.75 * np.pi) + + def test_zero_nzero(self): + # atan2(+-0, -0) returns +-pi. + assert_almost_equal(ncu.arctan2(ncu.PZERO, ncu.NZERO), np.pi) + assert_almost_equal(ncu.arctan2(ncu.NZERO, ncu.NZERO), -np.pi) + + def test_zero_pzero(self): + # atan2(+-0, +0) returns +-0. + assert_arctan2_ispzero(ncu.PZERO, ncu.PZERO) + assert_arctan2_isnzero(ncu.NZERO, ncu.PZERO) + + def test_zero_negative(self): + # atan2(+-0, x) returns +-pi for x < 0. + assert_almost_equal(ncu.arctan2(ncu.PZERO, -1), np.pi) + assert_almost_equal(ncu.arctan2(ncu.NZERO, -1), -np.pi) + + def test_zero_positive(self): + # atan2(+-0, x) returns +-0 for x > 0. + assert_arctan2_ispzero(ncu.PZERO, 1) + assert_arctan2_isnzero(ncu.NZERO, 1) + + def test_positive_zero(self): + # atan2(y, +-0) returns +pi/2 for y > 0. + assert_almost_equal(ncu.arctan2(1, ncu.PZERO), 0.5 * np.pi) + assert_almost_equal(ncu.arctan2(1, ncu.NZERO), 0.5 * np.pi) + + def test_negative_zero(self): + # atan2(y, +-0) returns -pi/2 for y < 0. + assert_almost_equal(ncu.arctan2(-1, ncu.PZERO), -0.5 * np.pi) + assert_almost_equal(ncu.arctan2(-1, ncu.NZERO), -0.5 * np.pi) + + def test_any_ninf(self): + # atan2(+-y, -infinity) returns +-pi for finite y > 0. + assert_almost_equal(ncu.arctan2(1, -np.inf), np.pi) + assert_almost_equal(ncu.arctan2(-1, -np.inf), -np.pi) + + def test_any_pinf(self): + # atan2(+-y, +infinity) returns +-0 for finite y > 0. + assert_arctan2_ispzero(1, np.inf) + assert_arctan2_isnzero(-1, np.inf) + + def test_inf_any(self): + # atan2(+-infinity, x) returns +-pi/2 for finite x. + assert_almost_equal(ncu.arctan2( np.inf, 1), 0.5 * np.pi) + assert_almost_equal(ncu.arctan2(-np.inf, 1), -0.5 * np.pi) + + def test_inf_ninf(self): + # atan2(+-infinity, -infinity) returns +-3*pi/4. + assert_almost_equal(ncu.arctan2( np.inf, -np.inf), 0.75 * np.pi) + assert_almost_equal(ncu.arctan2(-np.inf, -np.inf), -0.75 * np.pi) + + def test_inf_pinf(self): + # atan2(+-infinity, +infinity) returns +-pi/4. + assert_almost_equal(ncu.arctan2( np.inf, np.inf), 0.25 * np.pi) + assert_almost_equal(ncu.arctan2(-np.inf, np.inf), -0.25 * np.pi) + + def test_nan_any(self): + # atan2(nan, x) returns nan for any x, including inf + assert_arctan2_isnan(np.nan, np.inf) + assert_arctan2_isnan(np.inf, np.nan) + assert_arctan2_isnan(np.nan, np.nan) + + +class TestLdexp: + def _check_ldexp(self, tp): + assert_almost_equal(ncu.ldexp(np.array(2., np.float32), + np.array(3, tp)), 16.) + assert_almost_equal(ncu.ldexp(np.array(2., np.float64), + np.array(3, tp)), 16.) + assert_almost_equal(ncu.ldexp(np.array(2., np.longdouble), + np.array(3, tp)), 16.) + + def test_ldexp(self): + # The default Python int type should work + assert_almost_equal(ncu.ldexp(2., 3), 16.) + # The following int types should all be accepted + self._check_ldexp(np.int8) + self._check_ldexp(np.int16) + self._check_ldexp(np.int32) + self._check_ldexp('i') + self._check_ldexp('l') + + def test_ldexp_overflow(self): + # silence warning emitted on overflow + with np.errstate(over="ignore"): + imax = np.iinfo(np.dtype('l')).max + imin = np.iinfo(np.dtype('l')).min + assert_equal(ncu.ldexp(2., imax), np.inf) + assert_equal(ncu.ldexp(2., imin), 0) + + +class TestMaximum(_FilterInvalids): + def test_reduce(self): + dflt = np.typecodes['AllFloat'] + dint = np.typecodes['AllInteger'] + seq1 = np.arange(11) + seq2 = seq1[::-1] + func = np.maximum.reduce + for dt in dint: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 10) + assert_equal(func(tmp2), 10) + for dt in dflt: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 10) + assert_equal(func(tmp2), 10) + tmp1[::2] = np.nan + tmp2[::2] = np.nan + assert_equal(func(tmp1), np.nan) + assert_equal(func(tmp2), np.nan) + + def test_reduce_complex(self): + assert_equal(np.maximum.reduce([1, 2j]), 1) + assert_equal(np.maximum.reduce([1+3j, 2j]), 1+3j) + + def test_float_nans(self): + nan = np.nan + arg1 = np.array([0, nan, nan]) + arg2 = np.array([nan, 0, nan]) + out = np.array([nan, nan, nan]) + assert_equal(np.maximum(arg1, arg2), out) + + def test_object_nans(self): + # Multiple checks to give this a chance to + # fail if cmp is used instead of rich compare. + # Failure cannot be guaranteed. + for i in range(1): + x = np.array(float('nan'), object) + y = 1.0 + z = np.array(float('nan'), object) + assert_(np.maximum(x, y) == 1.0) + assert_(np.maximum(z, y) == 1.0) + + def test_complex_nans(self): + nan = np.nan + for cnan in [complex(nan, 0), complex(0, nan), complex(nan, nan)]: + arg1 = np.array([0, cnan, cnan], dtype=complex) + arg2 = np.array([cnan, 0, cnan], dtype=complex) + out = np.array([nan, nan, nan], dtype=complex) + assert_equal(np.maximum(arg1, arg2), out) + + def test_object_array(self): + arg1 = np.arange(5, dtype=object) + arg2 = arg1 + 1 + assert_equal(np.maximum(arg1, arg2), arg2) + + def test_strided_array(self): + arr1 = np.array([-4.0, 1.0, 10.0, 0.0, np.nan, -np.nan, np.inf, -np.inf]) + arr2 = np.array([-2.0,-1.0, np.nan, 1.0, 0.0, np.nan, 1.0, -3.0]) + maxtrue = np.array([-2.0, 1.0, np.nan, 1.0, np.nan, np.nan, np.inf, -3.0]) + out = np.ones(8) + out_maxtrue = np.array([-2.0, 1.0, 1.0, 10.0, 1.0, 1.0, np.nan, 1.0]) + assert_equal(np.maximum(arr1,arr2), maxtrue) + assert_equal(np.maximum(arr1[::2],arr2[::2]), maxtrue[::2]) + assert_equal(np.maximum(arr1[:4:], arr2[::2]), np.array([-2.0, np.nan, 10.0, 1.0])) + assert_equal(np.maximum(arr1[::3], arr2[:3:]), np.array([-2.0, 0.0, np.nan])) + assert_equal(np.maximum(arr1[:6:2], arr2[::3], out=out[::3]), np.array([-2.0, 10., np.nan])) + assert_equal(out, out_maxtrue) + + def test_precision(self): + dtypes = [np.float16, np.float32, np.float64, np.longdouble] + + for dt in dtypes: + dtmin = np.finfo(dt).min + dtmax = np.finfo(dt).max + d1 = dt(0.1) + d1_next = np.nextafter(d1, np.inf) + + test_cases = [ + # v1 v2 expected + (dtmin, -np.inf, dtmin), + (dtmax, -np.inf, dtmax), + (d1, d1_next, d1_next), + (dtmax, np.nan, np.nan), + ] + + for v1, v2, expected in test_cases: + assert_equal(np.maximum([v1], [v2]), [expected]) + assert_equal(np.maximum.reduce([v1, v2]), expected) + + +class TestMinimum(_FilterInvalids): + def test_reduce(self): + dflt = np.typecodes['AllFloat'] + dint = np.typecodes['AllInteger'] + seq1 = np.arange(11) + seq2 = seq1[::-1] + func = np.minimum.reduce + for dt in dint: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 0) + assert_equal(func(tmp2), 0) + for dt in dflt: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 0) + assert_equal(func(tmp2), 0) + tmp1[::2] = np.nan + tmp2[::2] = np.nan + assert_equal(func(tmp1), np.nan) + assert_equal(func(tmp2), np.nan) + + def test_reduce_complex(self): + assert_equal(np.minimum.reduce([1, 2j]), 2j) + assert_equal(np.minimum.reduce([1+3j, 2j]), 2j) + + def test_float_nans(self): + nan = np.nan + arg1 = np.array([0, nan, nan]) + arg2 = np.array([nan, 0, nan]) + out = np.array([nan, nan, nan]) + assert_equal(np.minimum(arg1, arg2), out) + + def test_object_nans(self): + # Multiple checks to give this a chance to + # fail if cmp is used instead of rich compare. + # Failure cannot be guaranteed. + for i in range(1): + x = np.array(float('nan'), object) + y = 1.0 + z = np.array(float('nan'), object) + assert_(np.minimum(x, y) == 1.0) + assert_(np.minimum(z, y) == 1.0) + + def test_complex_nans(self): + nan = np.nan + for cnan in [complex(nan, 0), complex(0, nan), complex(nan, nan)]: + arg1 = np.array([0, cnan, cnan], dtype=complex) + arg2 = np.array([cnan, 0, cnan], dtype=complex) + out = np.array([nan, nan, nan], dtype=complex) + assert_equal(np.minimum(arg1, arg2), out) + + def test_object_array(self): + arg1 = np.arange(5, dtype=object) + arg2 = arg1 + 1 + assert_equal(np.minimum(arg1, arg2), arg1) + + def test_strided_array(self): + arr1 = np.array([-4.0, 1.0, 10.0, 0.0, np.nan, -np.nan, np.inf, -np.inf]) + arr2 = np.array([-2.0,-1.0, np.nan, 1.0, 0.0, np.nan, 1.0, -3.0]) + mintrue = np.array([-4.0, -1.0, np.nan, 0.0, np.nan, np.nan, 1.0, -np.inf]) + out = np.ones(8) + out_mintrue = np.array([-4.0, 1.0, 1.0, 1.0, 1.0, 1.0, np.nan, 1.0]) + assert_equal(np.minimum(arr1,arr2), mintrue) + assert_equal(np.minimum(arr1[::2],arr2[::2]), mintrue[::2]) + assert_equal(np.minimum(arr1[:4:], arr2[::2]), np.array([-4.0, np.nan, 0.0, 0.0])) + assert_equal(np.minimum(arr1[::3], arr2[:3:]), np.array([-4.0, -1.0, np.nan])) + assert_equal(np.minimum(arr1[:6:2], arr2[::3], out=out[::3]), np.array([-4.0, 1.0, np.nan])) + assert_equal(out, out_mintrue) + + def test_precision(self): + dtypes = [np.float16, np.float32, np.float64, np.longdouble] + + for dt in dtypes: + dtmin = np.finfo(dt).min + dtmax = np.finfo(dt).max + d1 = dt(0.1) + d1_next = np.nextafter(d1, np.inf) + + test_cases = [ + # v1 v2 expected + (dtmin, np.inf, dtmin), + (dtmax, np.inf, dtmax), + (d1, d1_next, d1), + (dtmin, np.nan, np.nan), + ] + + for v1, v2, expected in test_cases: + assert_equal(np.minimum([v1], [v2]), [expected]) + assert_equal(np.minimum.reduce([v1, v2]), expected) + + +class TestFmax(_FilterInvalids): + def test_reduce(self): + dflt = np.typecodes['AllFloat'] + dint = np.typecodes['AllInteger'] + seq1 = np.arange(11) + seq2 = seq1[::-1] + func = np.fmax.reduce + for dt in dint: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 10) + assert_equal(func(tmp2), 10) + for dt in dflt: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 10) + assert_equal(func(tmp2), 10) + tmp1[::2] = np.nan + tmp2[::2] = np.nan + assert_equal(func(tmp1), 9) + assert_equal(func(tmp2), 9) + + def test_reduce_complex(self): + assert_equal(np.fmax.reduce([1, 2j]), 1) + assert_equal(np.fmax.reduce([1+3j, 2j]), 1+3j) + + def test_float_nans(self): + nan = np.nan + arg1 = np.array([0, nan, nan]) + arg2 = np.array([nan, 0, nan]) + out = np.array([0, 0, nan]) + assert_equal(np.fmax(arg1, arg2), out) + + def test_complex_nans(self): + nan = np.nan + for cnan in [complex(nan, 0), complex(0, nan), complex(nan, nan)]: + arg1 = np.array([0, cnan, cnan], dtype=complex) + arg2 = np.array([cnan, 0, cnan], dtype=complex) + out = np.array([0, 0, nan], dtype=complex) + assert_equal(np.fmax(arg1, arg2), out) + + def test_precision(self): + dtypes = [np.float16, np.float32, np.float64, np.longdouble] + + for dt in dtypes: + dtmin = np.finfo(dt).min + dtmax = np.finfo(dt).max + d1 = dt(0.1) + d1_next = np.nextafter(d1, np.inf) + + test_cases = [ + # v1 v2 expected + (dtmin, -np.inf, dtmin), + (dtmax, -np.inf, dtmax), + (d1, d1_next, d1_next), + (dtmax, np.nan, dtmax), + ] + + for v1, v2, expected in test_cases: + assert_equal(np.fmax([v1], [v2]), [expected]) + assert_equal(np.fmax.reduce([v1, v2]), expected) + + +class TestFmin(_FilterInvalids): + def test_reduce(self): + dflt = np.typecodes['AllFloat'] + dint = np.typecodes['AllInteger'] + seq1 = np.arange(11) + seq2 = seq1[::-1] + func = np.fmin.reduce + for dt in dint: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 0) + assert_equal(func(tmp2), 0) + for dt in dflt: + tmp1 = seq1.astype(dt) + tmp2 = seq2.astype(dt) + assert_equal(func(tmp1), 0) + assert_equal(func(tmp2), 0) + tmp1[::2] = np.nan + tmp2[::2] = np.nan + assert_equal(func(tmp1), 1) + assert_equal(func(tmp2), 1) + + def test_reduce_complex(self): + assert_equal(np.fmin.reduce([1, 2j]), 2j) + assert_equal(np.fmin.reduce([1+3j, 2j]), 2j) + + def test_float_nans(self): + nan = np.nan + arg1 = np.array([0, nan, nan]) + arg2 = np.array([nan, 0, nan]) + out = np.array([0, 0, nan]) + assert_equal(np.fmin(arg1, arg2), out) + + def test_complex_nans(self): + nan = np.nan + for cnan in [complex(nan, 0), complex(0, nan), complex(nan, nan)]: + arg1 = np.array([0, cnan, cnan], dtype=complex) + arg2 = np.array([cnan, 0, cnan], dtype=complex) + out = np.array([0, 0, nan], dtype=complex) + assert_equal(np.fmin(arg1, arg2), out) + + def test_precision(self): + dtypes = [np.float16, np.float32, np.float64, np.longdouble] + + for dt in dtypes: + dtmin = np.finfo(dt).min + dtmax = np.finfo(dt).max + d1 = dt(0.1) + d1_next = np.nextafter(d1, np.inf) + + test_cases = [ + # v1 v2 expected + (dtmin, np.inf, dtmin), + (dtmax, np.inf, dtmax), + (d1, d1_next, d1), + (dtmin, np.nan, dtmin), + ] + + for v1, v2, expected in test_cases: + assert_equal(np.fmin([v1], [v2]), [expected]) + assert_equal(np.fmin.reduce([v1, v2]), expected) + + +class TestBool: + def test_exceptions(self): + a = np.ones(1, dtype=np.bool) + assert_raises(TypeError, np.negative, a) + assert_raises(TypeError, np.positive, a) + assert_raises(TypeError, np.subtract, a, a) + + def test_truth_table_logical(self): + # 2, 3 and 4 serves as true values + input1 = [0, 0, 3, 2] + input2 = [0, 4, 0, 2] + + typecodes = (np.typecodes['AllFloat'] + + np.typecodes['AllInteger'] + + '?') # boolean + for dtype in map(np.dtype, typecodes): + arg1 = np.asarray(input1, dtype=dtype) + arg2 = np.asarray(input2, dtype=dtype) + + # OR + out = [False, True, True, True] + for func in (np.logical_or, np.maximum): + assert_equal(func(arg1, arg2).astype(bool), out) + # AND + out = [False, False, False, True] + for func in (np.logical_and, np.minimum): + assert_equal(func(arg1, arg2).astype(bool), out) + # XOR + out = [False, True, True, False] + for func in (np.logical_xor, np.not_equal): + assert_equal(func(arg1, arg2).astype(bool), out) + + def test_truth_table_bitwise(self): + arg1 = [False, False, True, True] + arg2 = [False, True, False, True] + + out = [False, True, True, True] + assert_equal(np.bitwise_or(arg1, arg2), out) + + out = [False, False, False, True] + assert_equal(np.bitwise_and(arg1, arg2), out) + + out = [False, True, True, False] + assert_equal(np.bitwise_xor(arg1, arg2), out) + + def test_reduce(self): + none = np.array([0, 0, 0, 0], bool) + some = np.array([1, 0, 1, 1], bool) + every = np.array([1, 1, 1, 1], bool) + empty = np.array([], bool) + + arrs = [none, some, every, empty] + + for arr in arrs: + assert_equal(np.logical_and.reduce(arr), all(arr)) + + for arr in arrs: + assert_equal(np.logical_or.reduce(arr), any(arr)) + + for arr in arrs: + assert_equal(np.logical_xor.reduce(arr), arr.sum() % 2 == 1) + + +class TestBitwiseUFuncs: + + _all_ints_bits = [ + np.dtype(c).itemsize * 8 for c in np.typecodes["AllInteger"]] + bitwise_types = [ + np.dtype(c) for c in '?' + np.typecodes["AllInteger"] + 'O'] + bitwise_bits = [ + 2, # boolean type + *_all_ints_bits, # All integers + max(_all_ints_bits) + 1, # Object_ type + ] + + def test_values(self): + for dt in self.bitwise_types: + zeros = np.array([0], dtype=dt) + ones = np.array([-1]).astype(dt) + msg = "dt = '%s'" % dt.char + + assert_equal(np.bitwise_not(zeros), ones, err_msg=msg) + assert_equal(np.bitwise_not(ones), zeros, err_msg=msg) + + assert_equal(np.bitwise_or(zeros, zeros), zeros, err_msg=msg) + assert_equal(np.bitwise_or(zeros, ones), ones, err_msg=msg) + assert_equal(np.bitwise_or(ones, zeros), ones, err_msg=msg) + assert_equal(np.bitwise_or(ones, ones), ones, err_msg=msg) + + assert_equal(np.bitwise_xor(zeros, zeros), zeros, err_msg=msg) + assert_equal(np.bitwise_xor(zeros, ones), ones, err_msg=msg) + assert_equal(np.bitwise_xor(ones, zeros), ones, err_msg=msg) + assert_equal(np.bitwise_xor(ones, ones), zeros, err_msg=msg) + + assert_equal(np.bitwise_and(zeros, zeros), zeros, err_msg=msg) + assert_equal(np.bitwise_and(zeros, ones), zeros, err_msg=msg) + assert_equal(np.bitwise_and(ones, zeros), zeros, err_msg=msg) + assert_equal(np.bitwise_and(ones, ones), ones, err_msg=msg) + + def test_types(self): + for dt in self.bitwise_types: + zeros = np.array([0], dtype=dt) + ones = np.array([-1]).astype(dt) + msg = "dt = '%s'" % dt.char + + assert_(np.bitwise_not(zeros).dtype == dt, msg) + assert_(np.bitwise_or(zeros, zeros).dtype == dt, msg) + assert_(np.bitwise_xor(zeros, zeros).dtype == dt, msg) + assert_(np.bitwise_and(zeros, zeros).dtype == dt, msg) + + def test_identity(self): + assert_(np.bitwise_or.identity == 0, 'bitwise_or') + assert_(np.bitwise_xor.identity == 0, 'bitwise_xor') + assert_(np.bitwise_and.identity == -1, 'bitwise_and') + + def test_reduction(self): + binary_funcs = (np.bitwise_or, np.bitwise_xor, np.bitwise_and) + + for dt in self.bitwise_types: + zeros = np.array([0], dtype=dt) + ones = np.array([-1]).astype(dt) + for f in binary_funcs: + msg = "dt: '%s', f: '%s'" % (dt, f) + assert_equal(f.reduce(zeros), zeros, err_msg=msg) + assert_equal(f.reduce(ones), ones, err_msg=msg) + + # Test empty reduction, no object dtype + for dt in self.bitwise_types[:-1]: + # No object array types + empty = np.array([], dtype=dt) + for f in binary_funcs: + msg = "dt: '%s', f: '%s'" % (dt, f) + tgt = np.array(f.identity).astype(dt) + res = f.reduce(empty) + assert_equal(res, tgt, err_msg=msg) + assert_(res.dtype == tgt.dtype, msg) + + # Empty object arrays use the identity. Note that the types may + # differ, the actual type used is determined by the assign_identity + # function and is not the same as the type returned by the identity + # method. + for f in binary_funcs: + msg = "dt: '%s'" % (f,) + empty = np.array([], dtype=object) + tgt = f.identity + res = f.reduce(empty) + assert_equal(res, tgt, err_msg=msg) + + # Non-empty object arrays do not use the identity + for f in binary_funcs: + msg = "dt: '%s'" % (f,) + btype = np.array([True], dtype=object) + assert_(type(f.reduce(btype)) is bool, msg) + + @pytest.mark.parametrize("input_dtype_obj, bitsize", + zip(bitwise_types, bitwise_bits)) + def test_bitwise_count(self, input_dtype_obj, bitsize): + input_dtype = input_dtype_obj.type + + for i in range(1, bitsize): + num = 2**i - 1 + msg = f"bitwise_count for {num}" + assert i == np.bitwise_count(input_dtype(num)), msg + if np.issubdtype( + input_dtype, np.signedinteger) or input_dtype == np.object_: + assert i == np.bitwise_count(input_dtype(-num)), msg + + a = np.array([2**i-1 for i in range(1, bitsize)], dtype=input_dtype) + bitwise_count_a = np.bitwise_count(a) + expected = np.arange(1, bitsize, dtype=input_dtype) + + msg = f"array bitwise_count for {input_dtype}" + assert all(bitwise_count_a == expected), msg + + +class TestInt: + def test_logical_not(self): + x = np.ones(10, dtype=np.int16) + o = np.ones(10 * 2, dtype=bool) + tgt = o.copy() + tgt[::2] = False + os = o[::2] + assert_array_equal(np.logical_not(x, out=os), False) + assert_array_equal(o, tgt) + + +class TestFloatingPoint: + def test_floating_point(self): + assert_equal(ncu.FLOATING_POINT_SUPPORT, 1) + + +class TestDegrees: + def test_degrees(self): + assert_almost_equal(ncu.degrees(np.pi), 180.0) + assert_almost_equal(ncu.degrees(-0.5*np.pi), -90.0) + + +class TestRadians: + def test_radians(self): + assert_almost_equal(ncu.radians(180.0), np.pi) + assert_almost_equal(ncu.radians(-90.0), -0.5*np.pi) + + +class TestHeavside: + def test_heaviside(self): + x = np.array([[-30.0, -0.1, 0.0, 0.2], [7.5, np.nan, np.inf, -np.inf]]) + expectedhalf = np.array([[0.0, 0.0, 0.5, 1.0], [1.0, np.nan, 1.0, 0.0]]) + expected1 = expectedhalf.copy() + expected1[0, 2] = 1 + + h = ncu.heaviside(x, 0.5) + assert_equal(h, expectedhalf) + + h = ncu.heaviside(x, 1.0) + assert_equal(h, expected1) + + x = x.astype(np.float32) + + h = ncu.heaviside(x, np.float32(0.5)) + assert_equal(h, expectedhalf.astype(np.float32)) + + h = ncu.heaviside(x, np.float32(1.0)) + assert_equal(h, expected1.astype(np.float32)) + + +class TestSign: + def test_sign(self): + a = np.array([np.inf, -np.inf, np.nan, 0.0, 3.0, -3.0]) + out = np.zeros(a.shape) + tgt = np.array([1., -1., np.nan, 0.0, 1.0, -1.0]) + + with np.errstate(invalid='ignore'): + res = ncu.sign(a) + assert_equal(res, tgt) + res = ncu.sign(a, out) + assert_equal(res, tgt) + assert_equal(out, tgt) + + def test_sign_complex(self): + a = np.array([ + np.inf, -np.inf, complex(0, np.inf), complex(0, -np.inf), + complex(np.inf, np.inf), complex(np.inf, -np.inf), # nan + np.nan, complex(0, np.nan), complex(np.nan, np.nan), # nan + 0.0, # 0. + 3.0, -3.0, -2j, 3.0+4.0j, -8.0+6.0j + ]) + out = np.zeros(a.shape, a.dtype) + tgt = np.array([ + 1., -1., 1j, -1j, + ] + [complex(np.nan, np.nan)] * 5 + [ + 0.0, + 1.0, -1.0, -1j, 0.6+0.8j, -0.8+0.6j]) + + with np.errstate(invalid='ignore'): + res = ncu.sign(a) + assert_equal(res, tgt) + res = ncu.sign(a, out) + assert_(res is out) + assert_equal(res, tgt) + + def test_sign_dtype_object(self): + # In reference to github issue #6229 + + foo = np.array([-.1, 0, .1]) + a = np.sign(foo.astype(object)) + b = np.sign(foo) + + assert_array_equal(a, b) + + def test_sign_dtype_nan_object(self): + # In reference to github issue #6229 + def test_nan(): + foo = np.array([np.nan]) + # FIXME: a not used + a = np.sign(foo.astype(object)) + + assert_raises(TypeError, test_nan) + +class TestMinMax: + def test_minmax_blocked(self): + # simd tests on max/min, test all alignments, slow but important + # for 2 * vz + 2 * (vs - 1) + 1 (unrolled once) + for dt, sz in [(np.float32, 15), (np.float64, 7)]: + for out, inp, msg in _gen_alignment_data(dtype=dt, type='unary', + max_size=sz): + for i in range(inp.size): + inp[:] = np.arange(inp.size, dtype=dt) + inp[i] = np.nan + emsg = lambda: '%r\n%s' % (inp, msg) + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, + "invalid value encountered in reduce") + assert_(np.isnan(inp.max()), msg=emsg) + assert_(np.isnan(inp.min()), msg=emsg) + + inp[i] = 1e10 + assert_equal(inp.max(), 1e10, err_msg=msg) + inp[i] = -1e10 + assert_equal(inp.min(), -1e10, err_msg=msg) + + def test_lower_align(self): + # check data that is not aligned to element size + # i.e doubles are aligned to 4 bytes on i386 + d = np.zeros(23 * 8, dtype=np.int8)[4:-4].view(np.float64) + assert_equal(d.max(), d[0]) + assert_equal(d.min(), d[0]) + + def test_reduce_reorder(self): + # gh 10370, 11029 Some compilers reorder the call to npy_getfloatstatus + # and put it before the call to an intrinsic function that causes + # invalid status to be set. Also make sure warnings are not emitted + for n in (2, 4, 8, 16, 32): + for dt in (np.float32, np.float16, np.complex64): + for r in np.diagflat(np.array([np.nan] * n, dtype=dt)): + assert_equal(np.min(r), np.nan) + + def test_minimize_no_warns(self): + a = np.minimum(np.nan, 1) + assert_equal(a, np.nan) + + +class TestAbsoluteNegative: + def test_abs_neg_blocked(self): + # simd tests on abs, test all alignments for vz + 2 * (vs - 1) + 1 + for dt, sz in [(np.float32, 11), (np.float64, 5)]: + for out, inp, msg in _gen_alignment_data(dtype=dt, type='unary', + max_size=sz): + tgt = [ncu.absolute(i) for i in inp] + np.absolute(inp, out=out) + assert_equal(out, tgt, err_msg=msg) + assert_((out >= 0).all()) + + tgt = [-1*(i) for i in inp] + np.negative(inp, out=out) + assert_equal(out, tgt, err_msg=msg) + + for v in [np.nan, -np.inf, np.inf]: + for i in range(inp.size): + d = np.arange(inp.size, dtype=dt) + inp[:] = -d + inp[i] = v + d[i] = -v if v == -np.inf else v + assert_array_equal(np.abs(inp), d, err_msg=msg) + np.abs(inp, out=out) + assert_array_equal(out, d, err_msg=msg) + + assert_array_equal(-inp, -1*inp, err_msg=msg) + d = -1 * inp + np.negative(inp, out=out) + assert_array_equal(out, d, err_msg=msg) + + def test_lower_align(self): + # check data that is not aligned to element size + # i.e doubles are aligned to 4 bytes on i386 + d = np.zeros(23 * 8, dtype=np.int8)[4:-4].view(np.float64) + assert_equal(np.abs(d), d) + assert_equal(np.negative(d), -d) + np.negative(d, out=d) + np.negative(np.ones_like(d), out=d) + np.abs(d, out=d) + np.abs(np.ones_like(d), out=d) + + @pytest.mark.parametrize("dtype", ['d', 'f', 'int32', 'int64']) + @pytest.mark.parametrize("big", [True, False]) + def test_noncontiguous(self, dtype, big): + data = np.array([-1.0, 1.0, -0.0, 0.0, 2.2251e-308, -2.5, 2.5, -6, + 6, -2.2251e-308, -8, 10], dtype=dtype) + expect = np.array([1.0, -1.0, 0.0, -0.0, -2.2251e-308, 2.5, -2.5, 6, + -6, 2.2251e-308, 8, -10], dtype=dtype) + if big: + data = np.repeat(data, 10) + expect = np.repeat(expect, 10) + out = np.ndarray(data.shape, dtype=dtype) + ncontig_in = data[1::2] + ncontig_out = out[1::2] + contig_in = np.array(ncontig_in) + # contig in, contig out + assert_array_equal(np.negative(contig_in), expect[1::2]) + # contig in, ncontig out + assert_array_equal(np.negative(contig_in, out=ncontig_out), + expect[1::2]) + # ncontig in, contig out + assert_array_equal(np.negative(ncontig_in), expect[1::2]) + # ncontig in, ncontig out + assert_array_equal(np.negative(ncontig_in, out=ncontig_out), + expect[1::2]) + # contig in, contig out, nd stride + data_split = np.array(np.array_split(data, 2)) + expect_split = np.array(np.array_split(expect, 2)) + assert_equal(np.negative(data_split), expect_split) + + +class TestPositive: + def test_valid(self): + valid_dtypes = [int, float, complex, object] + for dtype in valid_dtypes: + x = np.arange(5, dtype=dtype) + result = np.positive(x) + assert_equal(x, result, err_msg=str(dtype)) + + def test_invalid(self): + with assert_raises(TypeError): + np.positive(True) + with assert_raises(TypeError): + np.positive(np.datetime64('2000-01-01')) + with assert_raises(TypeError): + np.positive(np.array(['foo'], dtype=str)) + with assert_raises(TypeError): + np.positive(np.array(['bar'], dtype=object)) + + +class TestSpecialMethods: + def test_wrap(self): + + class with_wrap: + def __array__(self, dtype=None, copy=None): + return np.zeros(1) + + def __array_wrap__(self, arr, context, return_scalar): + r = with_wrap() + r.arr = arr + r.context = context + return r + + a = with_wrap() + x = ncu.minimum(a, a) + assert_equal(x.arr, np.zeros(1)) + func, args, i = x.context + assert_(func is ncu.minimum) + assert_equal(len(args), 2) + assert_equal(args[0], a) + assert_equal(args[1], a) + assert_equal(i, 0) + + def test_wrap_out(self): + # Calling convention for out should not affect how special methods are + # called + + class StoreArrayPrepareWrap(np.ndarray): + _wrap_args = None + _prepare_args = None + + def __new__(cls): + return np.zeros(()).view(cls) + + def __array_wrap__(self, obj, context, return_scalar): + self._wrap_args = context[1] + return obj + + @property + def args(self): + # We need to ensure these are fetched at the same time, before + # any other ufuncs are called by the assertions + return self._wrap_args + + def __repr__(self): + return "a" # for short test output + + def do_test(f_call, f_expected): + a = StoreArrayPrepareWrap() + + f_call(a) + + w = a.args + expected = f_expected(a) + try: + assert w == expected + except AssertionError as e: + # assert_equal produces truly useless error messages + raise AssertionError("\n".join([ + "Bad arguments passed in ufunc call", + " expected: {}".format(expected), + " __array_wrap__ got: {}".format(w) + ])) + + # method not on the out argument + do_test(lambda a: np.add(a, 0), lambda a: (a, 0)) + do_test(lambda a: np.add(a, 0, None), lambda a: (a, 0)) + do_test(lambda a: np.add(a, 0, out=None), lambda a: (a, 0)) + do_test(lambda a: np.add(a, 0, out=(None,)), lambda a: (a, 0)) + + # method on the out argument + do_test(lambda a: np.add(0, 0, a), lambda a: (0, 0, a)) + do_test(lambda a: np.add(0, 0, out=a), lambda a: (0, 0, a)) + do_test(lambda a: np.add(0, 0, out=(a,)), lambda a: (0, 0, a)) + + # Also check the where mask handling: + do_test(lambda a: np.add(a, 0, where=False), lambda a: (a, 0)) + do_test(lambda a: np.add(0, 0, a, where=False), lambda a: (0, 0, a)) + + def test_wrap_with_iterable(self): + # test fix for bug #1026: + + class with_wrap(np.ndarray): + __array_priority__ = 10 + + def __new__(cls): + return np.asarray(1).view(cls).copy() + + def __array_wrap__(self, arr, context, return_scalar): + return arr.view(type(self)) + + a = with_wrap() + x = ncu.multiply(a, (1, 2, 3)) + assert_(isinstance(x, with_wrap)) + assert_array_equal(x, np.array((1, 2, 3))) + + def test_priority_with_scalar(self): + # test fix for bug #826: + + class A(np.ndarray): + __array_priority__ = 10 + + def __new__(cls): + return np.asarray(1.0, 'float64').view(cls).copy() + + a = A() + x = np.float64(1)*a + assert_(isinstance(x, A)) + assert_array_equal(x, np.array(1)) + + def test_priority(self): + + class A: + def __array__(self, dtype=None, copy=None): + return np.zeros(1) + + def __array_wrap__(self, arr, context, return_scalar): + r = type(self)() + r.arr = arr + r.context = context + return r + + class B(A): + __array_priority__ = 20. + + class C(A): + __array_priority__ = 40. + + x = np.zeros(1) + a = A() + b = B() + c = C() + f = ncu.minimum + assert_(type(f(x, x)) is np.ndarray) + assert_(type(f(x, a)) is A) + assert_(type(f(x, b)) is B) + assert_(type(f(x, c)) is C) + assert_(type(f(a, x)) is A) + assert_(type(f(b, x)) is B) + assert_(type(f(c, x)) is C) + + assert_(type(f(a, a)) is A) + assert_(type(f(a, b)) is B) + assert_(type(f(b, a)) is B) + assert_(type(f(b, b)) is B) + assert_(type(f(b, c)) is C) + assert_(type(f(c, b)) is C) + assert_(type(f(c, c)) is C) + + assert_(type(ncu.exp(a) is A)) + assert_(type(ncu.exp(b) is B)) + assert_(type(ncu.exp(c) is C)) + + def test_failing_wrap(self): + + class A: + def __array__(self, dtype=None, copy=None): + return np.zeros(2) + + def __array_wrap__(self, arr, context, return_scalar): + raise RuntimeError + + a = A() + assert_raises(RuntimeError, ncu.maximum, a, a) + assert_raises(RuntimeError, ncu.maximum.reduce, a) + + def test_failing_out_wrap(self): + + singleton = np.array([1.0]) + + class Ok(np.ndarray): + def __array_wrap__(self, obj, context, return_scalar): + return singleton + + class Bad(np.ndarray): + def __array_wrap__(self, obj, context, return_scalar): + raise RuntimeError + + ok = np.empty(1).view(Ok) + bad = np.empty(1).view(Bad) + # double-free (segfault) of "ok" if "bad" raises an exception + for i in range(10): + assert_raises(RuntimeError, ncu.frexp, 1, ok, bad) + + def test_none_wrap(self): + # Tests that issue #8507 is resolved. Previously, this would segfault + + class A: + def __array__(self, dtype=None, copy=None): + return np.zeros(1) + + def __array_wrap__(self, arr, context=None, return_scalar=False): + return None + + a = A() + assert_equal(ncu.maximum(a, a), None) + + def test_default_prepare(self): + + class with_wrap: + __array_priority__ = 10 + + def __array__(self, dtype=None, copy=None): + return np.zeros(1) + + def __array_wrap__(self, arr, context, return_scalar): + return arr + + a = with_wrap() + x = ncu.minimum(a, a) + assert_equal(x, np.zeros(1)) + assert_equal(type(x), np.ndarray) + + def test_array_too_many_args(self): + + class A: + def __array__(self, dtype, context, copy=None): + return np.zeros(1) + + a = A() + assert_raises_regex(TypeError, '2 required positional', np.sum, a) + + def test_ufunc_override(self): + # check override works even with instance with high priority. + class A: + def __array_ufunc__(self, func, method, *inputs, **kwargs): + return self, func, method, inputs, kwargs + + class MyNDArray(np.ndarray): + __array_priority__ = 100 + + a = A() + b = np.array([1]).view(MyNDArray) + res0 = np.multiply(a, b) + res1 = np.multiply(b, b, out=a) + + # self + assert_equal(res0[0], a) + assert_equal(res1[0], a) + assert_equal(res0[1], np.multiply) + assert_equal(res1[1], np.multiply) + assert_equal(res0[2], '__call__') + assert_equal(res1[2], '__call__') + assert_equal(res0[3], (a, b)) + assert_equal(res1[3], (b, b)) + assert_equal(res0[4], {}) + assert_equal(res1[4], {'out': (a,)}) + + def test_ufunc_override_mro(self): + + # Some multi arg functions for testing. + def tres_mul(a, b, c): + return a * b * c + + def quatro_mul(a, b, c, d): + return a * b * c * d + + # Make these into ufuncs. + three_mul_ufunc = np.frompyfunc(tres_mul, 3, 1) + four_mul_ufunc = np.frompyfunc(quatro_mul, 4, 1) + + class A: + def __array_ufunc__(self, func, method, *inputs, **kwargs): + return "A" + + class ASub(A): + def __array_ufunc__(self, func, method, *inputs, **kwargs): + return "ASub" + + class B: + def __array_ufunc__(self, func, method, *inputs, **kwargs): + return "B" + + class C: + def __init__(self): + self.count = 0 + + def __array_ufunc__(self, func, method, *inputs, **kwargs): + self.count += 1 + return NotImplemented + + class CSub(C): + def __array_ufunc__(self, func, method, *inputs, **kwargs): + self.count += 1 + return NotImplemented + + a = A() + a_sub = ASub() + b = B() + c = C() + + # Standard + res = np.multiply(a, a_sub) + assert_equal(res, "ASub") + res = np.multiply(a_sub, b) + assert_equal(res, "ASub") + + # With 1 NotImplemented + res = np.multiply(c, a) + assert_equal(res, "A") + assert_equal(c.count, 1) + # Check our counter works, so we can trust tests below. + res = np.multiply(c, a) + assert_equal(c.count, 2) + + # Both NotImplemented. + c = C() + c_sub = CSub() + assert_raises(TypeError, np.multiply, c, c_sub) + assert_equal(c.count, 1) + assert_equal(c_sub.count, 1) + c.count = c_sub.count = 0 + assert_raises(TypeError, np.multiply, c_sub, c) + assert_equal(c.count, 1) + assert_equal(c_sub.count, 1) + c.count = 0 + assert_raises(TypeError, np.multiply, c, c) + assert_equal(c.count, 1) + c.count = 0 + assert_raises(TypeError, np.multiply, 2, c) + assert_equal(c.count, 1) + + # Ternary testing. + assert_equal(three_mul_ufunc(a, 1, 2), "A") + assert_equal(three_mul_ufunc(1, a, 2), "A") + assert_equal(three_mul_ufunc(1, 2, a), "A") + + assert_equal(three_mul_ufunc(a, a, 6), "A") + assert_equal(three_mul_ufunc(a, 2, a), "A") + assert_equal(three_mul_ufunc(a, 2, b), "A") + assert_equal(three_mul_ufunc(a, 2, a_sub), "ASub") + assert_equal(three_mul_ufunc(a, a_sub, 3), "ASub") + c.count = 0 + assert_equal(three_mul_ufunc(c, a_sub, 3), "ASub") + assert_equal(c.count, 1) + c.count = 0 + assert_equal(three_mul_ufunc(1, a_sub, c), "ASub") + assert_equal(c.count, 0) + + c.count = 0 + assert_equal(three_mul_ufunc(a, b, c), "A") + assert_equal(c.count, 0) + c_sub.count = 0 + assert_equal(three_mul_ufunc(a, b, c_sub), "A") + assert_equal(c_sub.count, 0) + assert_equal(three_mul_ufunc(1, 2, b), "B") + + assert_raises(TypeError, three_mul_ufunc, 1, 2, c) + assert_raises(TypeError, three_mul_ufunc, c_sub, 2, c) + assert_raises(TypeError, three_mul_ufunc, c_sub, 2, 3) + + # Quaternary testing. + assert_equal(four_mul_ufunc(a, 1, 2, 3), "A") + assert_equal(four_mul_ufunc(1, a, 2, 3), "A") + assert_equal(four_mul_ufunc(1, 1, a, 3), "A") + assert_equal(four_mul_ufunc(1, 1, 2, a), "A") + + assert_equal(four_mul_ufunc(a, b, 2, 3), "A") + assert_equal(four_mul_ufunc(1, a, 2, b), "A") + assert_equal(four_mul_ufunc(b, 1, a, 3), "B") + assert_equal(four_mul_ufunc(a_sub, 1, 2, a), "ASub") + assert_equal(four_mul_ufunc(a, 1, 2, a_sub), "ASub") + + c = C() + c_sub = CSub() + assert_raises(TypeError, four_mul_ufunc, 1, 2, 3, c) + assert_equal(c.count, 1) + c.count = 0 + assert_raises(TypeError, four_mul_ufunc, 1, 2, c_sub, c) + assert_equal(c_sub.count, 1) + assert_equal(c.count, 1) + c2 = C() + c.count = c_sub.count = 0 + assert_raises(TypeError, four_mul_ufunc, 1, c, c_sub, c2) + assert_equal(c_sub.count, 1) + assert_equal(c.count, 1) + assert_equal(c2.count, 0) + c.count = c2.count = c_sub.count = 0 + assert_raises(TypeError, four_mul_ufunc, c2, c, c_sub, c) + assert_equal(c_sub.count, 1) + assert_equal(c.count, 0) + assert_equal(c2.count, 1) + + def test_ufunc_override_methods(self): + + class A: + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + return self, ufunc, method, inputs, kwargs + + # __call__ + a = A() + with assert_raises(TypeError): + np.multiply.__call__(1, a, foo='bar', answer=42) + res = np.multiply.__call__(1, a, subok='bar', where=42) + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], '__call__') + assert_equal(res[3], (1, a)) + assert_equal(res[4], {'subok': 'bar', 'where': 42}) + + # __call__, wrong args + assert_raises(TypeError, np.multiply, a) + assert_raises(TypeError, np.multiply, a, a, a, a) + assert_raises(TypeError, np.multiply, a, a, sig='a', signature='a') + assert_raises(TypeError, ncu_tests.inner1d, a, a, axis=0, axes=[0, 0]) + + # reduce, positional args + res = np.multiply.reduce(a, 'axis0', 'dtype0', 'out0', 'keep0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'reduce') + assert_equal(res[3], (a,)) + assert_equal(res[4], {'dtype':'dtype0', + 'out': ('out0',), + 'keepdims': 'keep0', + 'axis': 'axis0'}) + + # reduce, kwargs + res = np.multiply.reduce(a, axis='axis0', dtype='dtype0', out='out0', + keepdims='keep0', initial='init0', + where='where0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'reduce') + assert_equal(res[3], (a,)) + assert_equal(res[4], {'dtype':'dtype0', + 'out': ('out0',), + 'keepdims': 'keep0', + 'axis': 'axis0', + 'initial': 'init0', + 'where': 'where0'}) + + # reduce, output equal to None removed, but not other explicit ones, + # even if they are at their default value. + res = np.multiply.reduce(a, 0, None, None, False) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False}) + res = np.multiply.reduce(a, out=None, axis=0, keepdims=True) + assert_equal(res[4], {'axis': 0, 'keepdims': True}) + res = np.multiply.reduce(a, None, out=(None,), dtype=None) + assert_equal(res[4], {'axis': None, 'dtype': None}) + res = np.multiply.reduce(a, 0, None, None, False, 2, True) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, + 'initial': 2, 'where': True}) + # np._NoValue ignored for initial + res = np.multiply.reduce(a, 0, None, None, False, + np._NoValue, True) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, + 'where': True}) + # None kept for initial, True for where. + res = np.multiply.reduce(a, 0, None, None, False, None, True) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, + 'initial': None, 'where': True}) + + # reduce, wrong args + assert_raises(ValueError, np.multiply.reduce, a, out=()) + assert_raises(ValueError, np.multiply.reduce, a, out=('out0', 'out1')) + assert_raises(TypeError, np.multiply.reduce, a, 'axis0', axis='axis0') + + # accumulate, pos args + res = np.multiply.accumulate(a, 'axis0', 'dtype0', 'out0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'accumulate') + assert_equal(res[3], (a,)) + assert_equal(res[4], {'dtype':'dtype0', + 'out': ('out0',), + 'axis': 'axis0'}) + + # accumulate, kwargs + res = np.multiply.accumulate(a, axis='axis0', dtype='dtype0', + out='out0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'accumulate') + assert_equal(res[3], (a,)) + assert_equal(res[4], {'dtype':'dtype0', + 'out': ('out0',), + 'axis': 'axis0'}) + + # accumulate, output equal to None removed. + res = np.multiply.accumulate(a, 0, None, None) + assert_equal(res[4], {'axis': 0, 'dtype': None}) + res = np.multiply.accumulate(a, out=None, axis=0, dtype='dtype1') + assert_equal(res[4], {'axis': 0, 'dtype': 'dtype1'}) + res = np.multiply.accumulate(a, None, out=(None,), dtype=None) + assert_equal(res[4], {'axis': None, 'dtype': None}) + + # accumulate, wrong args + assert_raises(ValueError, np.multiply.accumulate, a, out=()) + assert_raises(ValueError, np.multiply.accumulate, a, + out=('out0', 'out1')) + assert_raises(TypeError, np.multiply.accumulate, a, + 'axis0', axis='axis0') + + # reduceat, pos args + res = np.multiply.reduceat(a, [4, 2], 'axis0', 'dtype0', 'out0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'reduceat') + assert_equal(res[3], (a, [4, 2])) + assert_equal(res[4], {'dtype':'dtype0', + 'out': ('out0',), + 'axis': 'axis0'}) + + # reduceat, kwargs + res = np.multiply.reduceat(a, [4, 2], axis='axis0', dtype='dtype0', + out='out0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'reduceat') + assert_equal(res[3], (a, [4, 2])) + assert_equal(res[4], {'dtype':'dtype0', + 'out': ('out0',), + 'axis': 'axis0'}) + + # reduceat, output equal to None removed. + res = np.multiply.reduceat(a, [4, 2], 0, None, None) + assert_equal(res[4], {'axis': 0, 'dtype': None}) + res = np.multiply.reduceat(a, [4, 2], axis=None, out=None, dtype='dt') + assert_equal(res[4], {'axis': None, 'dtype': 'dt'}) + res = np.multiply.reduceat(a, [4, 2], None, None, out=(None,)) + assert_equal(res[4], {'axis': None, 'dtype': None}) + + # reduceat, wrong args + assert_raises(ValueError, np.multiply.reduce, a, [4, 2], out=()) + assert_raises(ValueError, np.multiply.reduce, a, [4, 2], + out=('out0', 'out1')) + assert_raises(TypeError, np.multiply.reduce, a, [4, 2], + 'axis0', axis='axis0') + + # outer + res = np.multiply.outer(a, 42) + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'outer') + assert_equal(res[3], (a, 42)) + assert_equal(res[4], {}) + + # outer, wrong args + assert_raises(TypeError, np.multiply.outer, a) + assert_raises(TypeError, np.multiply.outer, a, a, a, a) + assert_raises(TypeError, np.multiply.outer, a, a, sig='a', signature='a') + + # at + res = np.multiply.at(a, [4, 2], 'b0') + assert_equal(res[0], a) + assert_equal(res[1], np.multiply) + assert_equal(res[2], 'at') + assert_equal(res[3], (a, [4, 2], 'b0')) + + # at, wrong args + assert_raises(TypeError, np.multiply.at, a) + assert_raises(TypeError, np.multiply.at, a, a, a, a) + + def test_ufunc_override_out(self): + + class A: + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + return kwargs + + class B: + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + return kwargs + + a = A() + b = B() + res0 = np.multiply(a, b, 'out_arg') + res1 = np.multiply(a, b, out='out_arg') + res2 = np.multiply(2, b, 'out_arg') + res3 = np.multiply(3, b, out='out_arg') + res4 = np.multiply(a, 4, 'out_arg') + res5 = np.multiply(a, 5, out='out_arg') + + assert_equal(res0['out'][0], 'out_arg') + assert_equal(res1['out'][0], 'out_arg') + assert_equal(res2['out'][0], 'out_arg') + assert_equal(res3['out'][0], 'out_arg') + assert_equal(res4['out'][0], 'out_arg') + assert_equal(res5['out'][0], 'out_arg') + + # ufuncs with multiple output modf and frexp. + res6 = np.modf(a, 'out0', 'out1') + res7 = np.frexp(a, 'out0', 'out1') + assert_equal(res6['out'][0], 'out0') + assert_equal(res6['out'][1], 'out1') + assert_equal(res7['out'][0], 'out0') + assert_equal(res7['out'][1], 'out1') + + # While we're at it, check that default output is never passed on. + assert_(np.sin(a, None) == {}) + assert_(np.sin(a, out=None) == {}) + assert_(np.sin(a, out=(None,)) == {}) + assert_(np.modf(a, None) == {}) + assert_(np.modf(a, None, None) == {}) + assert_(np.modf(a, out=(None, None)) == {}) + with assert_raises(TypeError): + # Out argument must be tuple, since there are multiple outputs. + np.modf(a, out=None) + + # don't give positional and output argument, or too many arguments. + # wrong number of arguments in the tuple is an error too. + assert_raises(TypeError, np.multiply, a, b, 'one', out='two') + assert_raises(TypeError, np.multiply, a, b, 'one', 'two') + assert_raises(ValueError, np.multiply, a, b, out=('one', 'two')) + assert_raises(TypeError, np.multiply, a, out=()) + assert_raises(TypeError, np.modf, a, 'one', out=('two', 'three')) + assert_raises(TypeError, np.modf, a, 'one', 'two', 'three') + assert_raises(ValueError, np.modf, a, out=('one', 'two', 'three')) + assert_raises(ValueError, np.modf, a, out=('one',)) + + def test_ufunc_override_where(self): + + class OverriddenArrayOld(np.ndarray): + + def _unwrap(self, objs): + cls = type(self) + result = [] + for obj in objs: + if isinstance(obj, cls): + obj = np.array(obj) + elif type(obj) != np.ndarray: + return NotImplemented + result.append(obj) + return result + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + + inputs = self._unwrap(inputs) + if inputs is NotImplemented: + return NotImplemented + + kwargs = kwargs.copy() + if "out" in kwargs: + kwargs["out"] = self._unwrap(kwargs["out"]) + if kwargs["out"] is NotImplemented: + return NotImplemented + + r = super().__array_ufunc__(ufunc, method, *inputs, **kwargs) + if r is not NotImplemented: + r = r.view(type(self)) + + return r + + class OverriddenArrayNew(OverriddenArrayOld): + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + + kwargs = kwargs.copy() + if "where" in kwargs: + kwargs["where"] = self._unwrap((kwargs["where"], )) + if kwargs["where"] is NotImplemented: + return NotImplemented + else: + kwargs["where"] = kwargs["where"][0] + + r = super().__array_ufunc__(ufunc, method, *inputs, **kwargs) + if r is not NotImplemented: + r = r.view(type(self)) + + return r + + ufunc = np.negative + + array = np.array([1, 2, 3]) + where = np.array([True, False, True]) + expected = ufunc(array, where=where) + + with pytest.raises(TypeError): + ufunc(array, where=where.view(OverriddenArrayOld)) + + result_1 = ufunc( + array, + where=where.view(OverriddenArrayNew) + ) + assert isinstance(result_1, OverriddenArrayNew) + assert np.all(np.array(result_1) == expected, where=where) + + result_2 = ufunc( + array.view(OverriddenArrayNew), + where=where.view(OverriddenArrayNew) + ) + assert isinstance(result_2, OverriddenArrayNew) + assert np.all(np.array(result_2) == expected, where=where) + + def test_ufunc_override_exception(self): + + class A: + def __array_ufunc__(self, *a, **kwargs): + raise ValueError("oops") + + a = A() + assert_raises(ValueError, np.negative, 1, out=a) + assert_raises(ValueError, np.negative, a) + assert_raises(ValueError, np.divide, 1., a) + + def test_ufunc_override_not_implemented(self): + + class A: + def __array_ufunc__(self, *args, **kwargs): + return NotImplemented + + msg = ("operand type(s) all returned NotImplemented from " + "__array_ufunc__(, '__call__', <*>): 'A'") + with assert_raises_regex(TypeError, fnmatch.translate(msg)): + np.negative(A()) + + msg = ("operand type(s) all returned NotImplemented from " + "__array_ufunc__(, '__call__', <*>, , " + "out=(1,)): 'A', 'object', 'int'") + with assert_raises_regex(TypeError, fnmatch.translate(msg)): + np.add(A(), object(), out=1) + + def test_ufunc_override_disabled(self): + + class OptOut: + __array_ufunc__ = None + + opt_out = OptOut() + + # ufuncs always raise + msg = "operand 'OptOut' does not support ufuncs" + with assert_raises_regex(TypeError, msg): + np.add(opt_out, 1) + with assert_raises_regex(TypeError, msg): + np.add(1, opt_out) + with assert_raises_regex(TypeError, msg): + np.negative(opt_out) + + # opt-outs still hold even when other arguments have pathological + # __array_ufunc__ implementations + + class GreedyArray: + def __array_ufunc__(self, *args, **kwargs): + return self + + greedy = GreedyArray() + assert_(np.negative(greedy) is greedy) + with assert_raises_regex(TypeError, msg): + np.add(greedy, opt_out) + with assert_raises_regex(TypeError, msg): + np.add(greedy, 1, out=opt_out) + + def test_gufunc_override(self): + # gufunc are just ufunc instances, but follow a different path, + # so check __array_ufunc__ overrides them properly. + class A: + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + return self, ufunc, method, inputs, kwargs + + inner1d = ncu_tests.inner1d + a = A() + res = inner1d(a, a) + assert_equal(res[0], a) + assert_equal(res[1], inner1d) + assert_equal(res[2], '__call__') + assert_equal(res[3], (a, a)) + assert_equal(res[4], {}) + + res = inner1d(1, 1, out=a) + assert_equal(res[0], a) + assert_equal(res[1], inner1d) + assert_equal(res[2], '__call__') + assert_equal(res[3], (1, 1)) + assert_equal(res[4], {'out': (a,)}) + + # wrong number of arguments in the tuple is an error too. + assert_raises(TypeError, inner1d, a, out='two') + assert_raises(TypeError, inner1d, a, a, 'one', out='two') + assert_raises(TypeError, inner1d, a, a, 'one', 'two') + assert_raises(ValueError, inner1d, a, a, out=('one', 'two')) + assert_raises(ValueError, inner1d, a, a, out=()) + + def test_ufunc_override_with_super(self): + # NOTE: this class is used in doc/source/user/basics.subclassing.rst + # if you make any changes here, do update it there too. + class A(np.ndarray): + def __array_ufunc__(self, ufunc, method, *inputs, out=None, **kwargs): + args = [] + in_no = [] + for i, input_ in enumerate(inputs): + if isinstance(input_, A): + in_no.append(i) + args.append(input_.view(np.ndarray)) + else: + args.append(input_) + + outputs = out + out_no = [] + if outputs: + out_args = [] + for j, output in enumerate(outputs): + if isinstance(output, A): + out_no.append(j) + out_args.append(output.view(np.ndarray)) + else: + out_args.append(output) + kwargs['out'] = tuple(out_args) + else: + outputs = (None,) * ufunc.nout + + info = {} + if in_no: + info['inputs'] = in_no + if out_no: + info['outputs'] = out_no + + results = super().__array_ufunc__(ufunc, method, + *args, **kwargs) + if results is NotImplemented: + return NotImplemented + + if method == 'at': + if isinstance(inputs[0], A): + inputs[0].info = info + return + + if ufunc.nout == 1: + results = (results,) + + results = tuple((np.asarray(result).view(A) + if output is None else output) + for result, output in zip(results, outputs)) + if results and isinstance(results[0], A): + results[0].info = info + + return results[0] if len(results) == 1 else results + + class B: + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + if any(isinstance(input_, A) for input_ in inputs): + return "A!" + else: + return NotImplemented + + d = np.arange(5.) + # 1 input, 1 output + a = np.arange(5.).view(A) + b = np.sin(a) + check = np.sin(d) + assert_(np.all(check == b)) + assert_equal(b.info, {'inputs': [0]}) + b = np.sin(d, out=(a,)) + assert_(np.all(check == b)) + assert_equal(b.info, {'outputs': [0]}) + assert_(b is a) + a = np.arange(5.).view(A) + b = np.sin(a, out=a) + assert_(np.all(check == b)) + assert_equal(b.info, {'inputs': [0], 'outputs': [0]}) + + # 1 input, 2 outputs + a = np.arange(5.).view(A) + b1, b2 = np.modf(a) + assert_equal(b1.info, {'inputs': [0]}) + b1, b2 = np.modf(d, out=(None, a)) + assert_(b2 is a) + assert_equal(b1.info, {'outputs': [1]}) + a = np.arange(5.).view(A) + b = np.arange(5.).view(A) + c1, c2 = np.modf(a, out=(a, b)) + assert_(c1 is a) + assert_(c2 is b) + assert_equal(c1.info, {'inputs': [0], 'outputs': [0, 1]}) + + # 2 input, 1 output + a = np.arange(5.).view(A) + b = np.arange(5.).view(A) + c = np.add(a, b, out=a) + assert_(c is a) + assert_equal(c.info, {'inputs': [0, 1], 'outputs': [0]}) + # some tests with a non-ndarray subclass + a = np.arange(5.) + b = B() + assert_(a.__array_ufunc__(np.add, '__call__', a, b) is NotImplemented) + assert_(b.__array_ufunc__(np.add, '__call__', a, b) is NotImplemented) + assert_raises(TypeError, np.add, a, b) + a = a.view(A) + assert_(a.__array_ufunc__(np.add, '__call__', a, b) is NotImplemented) + assert_(b.__array_ufunc__(np.add, '__call__', a, b) == "A!") + assert_(np.add(a, b) == "A!") + # regression check for gh-9102 -- tests ufunc.reduce implicitly. + d = np.array([[1, 2, 3], [1, 2, 3]]) + a = d.view(A) + c = a.any() + check = d.any() + assert_equal(c, check) + assert_(c.info, {'inputs': [0]}) + c = a.max() + check = d.max() + assert_equal(c, check) + assert_(c.info, {'inputs': [0]}) + b = np.array(0).view(A) + c = a.max(out=b) + assert_equal(c, check) + assert_(c is b) + assert_(c.info, {'inputs': [0], 'outputs': [0]}) + check = a.max(axis=0) + b = np.zeros_like(check).view(A) + c = a.max(axis=0, out=b) + assert_equal(c, check) + assert_(c is b) + assert_(c.info, {'inputs': [0], 'outputs': [0]}) + # simple explicit tests of reduce, accumulate, reduceat + check = np.add.reduce(d, axis=1) + c = np.add.reduce(a, axis=1) + assert_equal(c, check) + assert_(c.info, {'inputs': [0]}) + b = np.zeros_like(c) + c = np.add.reduce(a, 1, None, b) + assert_equal(c, check) + assert_(c is b) + assert_(c.info, {'inputs': [0], 'outputs': [0]}) + check = np.add.accumulate(d, axis=0) + c = np.add.accumulate(a, axis=0) + assert_equal(c, check) + assert_(c.info, {'inputs': [0]}) + b = np.zeros_like(c) + c = np.add.accumulate(a, 0, None, b) + assert_equal(c, check) + assert_(c is b) + assert_(c.info, {'inputs': [0], 'outputs': [0]}) + indices = [0, 2, 1] + check = np.add.reduceat(d, indices, axis=1) + c = np.add.reduceat(a, indices, axis=1) + assert_equal(c, check) + assert_(c.info, {'inputs': [0]}) + b = np.zeros_like(c) + c = np.add.reduceat(a, indices, 1, None, b) + assert_equal(c, check) + assert_(c is b) + assert_(c.info, {'inputs': [0], 'outputs': [0]}) + # and a few tests for at + d = np.array([[1, 2, 3], [1, 2, 3]]) + check = d.copy() + a = d.copy().view(A) + np.add.at(check, ([0, 1], [0, 2]), 1.) + np.add.at(a, ([0, 1], [0, 2]), 1.) + assert_equal(a, check) + assert_(a.info, {'inputs': [0]}) + b = np.array(1.).view(A) + a = d.copy().view(A) + np.add.at(a, ([0, 1], [0, 2]), b) + assert_equal(a, check) + assert_(a.info, {'inputs': [0, 2]}) + + def test_array_ufunc_direct_call(self): + # This is mainly a regression test for gh-24023 (shouldn't segfault) + a = np.array(1) + with pytest.raises(TypeError): + a.__array_ufunc__() + + # No kwargs means kwargs may be NULL on the C-level + with pytest.raises(TypeError): + a.__array_ufunc__(1, 2) + + # And the same with a valid call: + res = a.__array_ufunc__(np.add, "__call__", a, a) + assert_array_equal(res, a + a) + + def test_ufunc_docstring(self): + original_doc = np.add.__doc__ + new_doc = "new docs" + expected_dict = ( + {} if IS_PYPY else {"__module__": "numpy", "__qualname__": "add"} + ) + + np.add.__doc__ = new_doc + assert np.add.__doc__ == new_doc + assert np.add.__dict__["__doc__"] == new_doc + + del np.add.__doc__ + assert np.add.__doc__ == original_doc + assert np.add.__dict__ == expected_dict + + np.add.__dict__["other"] = 1 + np.add.__dict__["__doc__"] = new_doc + assert np.add.__doc__ == new_doc + + del np.add.__dict__["__doc__"] + assert np.add.__doc__ == original_doc + del np.add.__dict__["other"] + assert np.add.__dict__ == expected_dict + + +class TestChoose: + def test_mixed(self): + c = np.array([True, True]) + a = np.array([True, True]) + assert_equal(np.choose(c, (a, 1)), np.array([1, 1])) + + +class TestRationalFunctions: + def test_lcm(self): + self._test_lcm_inner(np.int16) + self._test_lcm_inner(np.uint16) + + def test_lcm_object(self): + self._test_lcm_inner(np.object_) + + def test_gcd(self): + self._test_gcd_inner(np.int16) + self._test_lcm_inner(np.uint16) + + def test_gcd_object(self): + self._test_gcd_inner(np.object_) + + def _test_lcm_inner(self, dtype): + # basic use + a = np.array([12, 120], dtype=dtype) + b = np.array([20, 200], dtype=dtype) + assert_equal(np.lcm(a, b), [60, 600]) + + if not issubclass(dtype, np.unsignedinteger): + # negatives are ignored + a = np.array([12, -12, 12, -12], dtype=dtype) + b = np.array([20, 20, -20, -20], dtype=dtype) + assert_equal(np.lcm(a, b), [60]*4) + + # reduce + a = np.array([3, 12, 20], dtype=dtype) + assert_equal(np.lcm.reduce([3, 12, 20]), 60) + + # broadcasting, and a test including 0 + a = np.arange(6).astype(dtype) + b = 20 + assert_equal(np.lcm(a, b), [0, 20, 20, 60, 20, 20]) + + def _test_gcd_inner(self, dtype): + # basic use + a = np.array([12, 120], dtype=dtype) + b = np.array([20, 200], dtype=dtype) + assert_equal(np.gcd(a, b), [4, 40]) + + if not issubclass(dtype, np.unsignedinteger): + # negatives are ignored + a = np.array([12, -12, 12, -12], dtype=dtype) + b = np.array([20, 20, -20, -20], dtype=dtype) + assert_equal(np.gcd(a, b), [4]*4) + + # reduce + a = np.array([15, 25, 35], dtype=dtype) + assert_equal(np.gcd.reduce(a), 5) + + # broadcasting, and a test including 0 + a = np.arange(6).astype(dtype) + b = 20 + assert_equal(np.gcd(a, b), [20, 1, 2, 1, 4, 5]) + + def test_lcm_overflow(self): + # verify that we don't overflow when a*b does overflow + big = np.int32(np.iinfo(np.int32).max // 11) + a = 2*big + b = 5*big + assert_equal(np.lcm(a, b), 10*big) + + def test_gcd_overflow(self): + for dtype in (np.int32, np.int64): + # verify that we don't overflow when taking abs(x) + # not relevant for lcm, where the result is unrepresentable anyway + a = dtype(np.iinfo(dtype).min) # negative power of two + q = -(a // 4) + assert_equal(np.gcd(a, q*3), q) + assert_equal(np.gcd(a, -q*3), q) + + def test_decimal(self): + from decimal import Decimal + a = np.array([1, 1, -1, -1]) * Decimal('0.20') + b = np.array([1, -1, 1, -1]) * Decimal('0.12') + + assert_equal(np.gcd(a, b), 4*[Decimal('0.04')]) + assert_equal(np.lcm(a, b), 4*[Decimal('0.60')]) + + def test_float(self): + # not well-defined on float due to rounding errors + assert_raises(TypeError, np.gcd, 0.3, 0.4) + assert_raises(TypeError, np.lcm, 0.3, 0.4) + + def test_huge_integers(self): + # Converting to an array first is a bit different as it means we + # have an explicit object dtype: + assert_equal(np.array(2**200), 2**200) + # Special promotion rules should ensure that this also works for + # two Python integers (even if slow). + # (We do this for comparisons, as the result is always bool and + # we also special case array comparisons with Python integers) + np.equal(2**200, 2**200) + + # But, we cannot do this when it would affect the result dtype: + with pytest.raises(OverflowError): + np.gcd(2**100, 3**100) + + # Asking for `object` explicitly is fine, though: + assert np.gcd(2**100, 3**100, dtype=object) == 1 + + # As of now, the below work, because it is using arrays (which + # will be object arrays) + a = np.array(2**100 * 3**5) + b = np.array([2**100 * 5**7, 2**50 * 3**10]) + assert_equal(np.gcd(a, b), [2**100, 2**50 * 3**5]) + assert_equal(np.lcm(a, b), [2**100 * 3**5 * 5**7, 2**100 * 3**10]) + + def test_inf_and_nan(self): + inf = np.array([np.inf], dtype=np.object_) + assert_raises(ValueError, np.gcd, inf, 1) + assert_raises(ValueError, np.gcd, 1, inf) + assert_raises(ValueError, np.gcd, np.nan, inf) + assert_raises(TypeError, np.gcd, 4, float(np.inf)) + + + +class TestRoundingFunctions: + + def test_object_direct(self): + """ test direct implementation of these magic methods """ + class C: + def __floor__(self): + return 1 + def __ceil__(self): + return 2 + def __trunc__(self): + return 3 + + arr = np.array([C(), C()]) + assert_equal(np.floor(arr), [1, 1]) + assert_equal(np.ceil(arr), [2, 2]) + assert_equal(np.trunc(arr), [3, 3]) + + def test_object_indirect(self): + """ test implementations via __float__ """ + class C: + def __float__(self): + return -2.5 + + arr = np.array([C(), C()]) + assert_equal(np.floor(arr), [-3, -3]) + assert_equal(np.ceil(arr), [-2, -2]) + with pytest.raises(TypeError): + np.trunc(arr) # consistent with math.trunc + + def test_fraction(self): + f = Fraction(-4, 3) + assert_equal(np.floor(f), -2) + assert_equal(np.ceil(f), -1) + assert_equal(np.trunc(f), -1) + + @pytest.mark.parametrize('func', [np.floor, np.ceil, np.trunc]) + @pytest.mark.parametrize('dtype', [np.bool, np.float64, np.float32, + np.int64, np.uint32]) + def test_output_dtype(self, func, dtype): + arr = np.array([-2, 0, 4, 8]).astype(dtype) + result = func(arr) + assert_equal(arr, result) + assert result.dtype == dtype + + +class TestComplexFunctions: + funcs = [np.arcsin, np.arccos, np.arctan, np.arcsinh, np.arccosh, + np.arctanh, np.sin, np.cos, np.tan, np.exp, + np.exp2, np.log, np.sqrt, np.log10, np.log2, + np.log1p] + + def test_it(self): + for f in self.funcs: + if f is np.arccosh: + x = 1.5 + else: + x = .5 + fr = f(x) + fz = f(complex(x)) + assert_almost_equal(fz.real, fr, err_msg='real part %s' % f) + assert_almost_equal(fz.imag, 0., err_msg='imag part %s' % f) + + @pytest.mark.xfail(IS_WASM, reason="doesn't work") + def test_precisions_consistent(self): + z = 1 + 1j + for f in self.funcs: + fcf = f(np.csingle(z)) + fcd = f(np.cdouble(z)) + fcl = f(np.clongdouble(z)) + assert_almost_equal(fcf, fcd, decimal=6, err_msg='fch-fcd %s' % f) + assert_almost_equal(fcl, fcd, decimal=15, err_msg='fch-fcl %s' % f) + + @pytest.mark.xfail(IS_WASM, reason="doesn't work") + def test_branch_cuts(self): + # check branch cuts and continuity on them + _check_branch_cut(np.log, -0.5, 1j, 1, -1, True) + _check_branch_cut(np.log2, -0.5, 1j, 1, -1, True) + _check_branch_cut(np.log10, -0.5, 1j, 1, -1, True) + _check_branch_cut(np.log1p, -1.5, 1j, 1, -1, True) + _check_branch_cut(np.sqrt, -0.5, 1j, 1, -1, True) + + _check_branch_cut(np.arcsin, [ -2, 2], [1j, 1j], 1, -1, True) + _check_branch_cut(np.arccos, [ -2, 2], [1j, 1j], 1, -1, True) + _check_branch_cut(np.arctan, [0-2j, 2j], [1, 1], -1, 1, True) + + _check_branch_cut(np.arcsinh, [0-2j, 2j], [1, 1], -1, 1, True) + _check_branch_cut(np.arccosh, [ -1, 0.5], [1j, 1j], 1, -1, True) + _check_branch_cut(np.arctanh, [ -2, 2], [1j, 1j], 1, -1, True) + + # check against bogus branch cuts: assert continuity between quadrants + _check_branch_cut(np.arcsin, [0-2j, 2j], [ 1, 1], 1, 1) + _check_branch_cut(np.arccos, [0-2j, 2j], [ 1, 1], 1, 1) + _check_branch_cut(np.arctan, [ -2, 2], [1j, 1j], 1, 1) + + _check_branch_cut(np.arcsinh, [ -2, 2, 0], [1j, 1j, 1], 1, 1) + _check_branch_cut(np.arccosh, [0-2j, 2j, 2], [1, 1, 1j], 1, 1) + _check_branch_cut(np.arctanh, [0-2j, 2j, 0], [1, 1, 1j], 1, 1) + + @pytest.mark.xfail(IS_WASM, reason="doesn't work") + def test_branch_cuts_complex64(self): + # check branch cuts and continuity on them + _check_branch_cut(np.log, -0.5, 1j, 1, -1, True, np.complex64) + _check_branch_cut(np.log2, -0.5, 1j, 1, -1, True, np.complex64) + _check_branch_cut(np.log10, -0.5, 1j, 1, -1, True, np.complex64) + _check_branch_cut(np.log1p, -1.5, 1j, 1, -1, True, np.complex64) + _check_branch_cut(np.sqrt, -0.5, 1j, 1, -1, True, np.complex64) + + _check_branch_cut(np.arcsin, [ -2, 2], [1j, 1j], 1, -1, True, np.complex64) + _check_branch_cut(np.arccos, [ -2, 2], [1j, 1j], 1, -1, True, np.complex64) + _check_branch_cut(np.arctan, [0-2j, 2j], [1, 1], -1, 1, True, np.complex64) + + _check_branch_cut(np.arcsinh, [0-2j, 2j], [1, 1], -1, 1, True, np.complex64) + _check_branch_cut(np.arccosh, [ -1, 0.5], [1j, 1j], 1, -1, True, np.complex64) + _check_branch_cut(np.arctanh, [ -2, 2], [1j, 1j], 1, -1, True, np.complex64) + + # check against bogus branch cuts: assert continuity between quadrants + _check_branch_cut(np.arcsin, [0-2j, 2j], [ 1, 1], 1, 1, False, np.complex64) + _check_branch_cut(np.arccos, [0-2j, 2j], [ 1, 1], 1, 1, False, np.complex64) + _check_branch_cut(np.arctan, [ -2, 2], [1j, 1j], 1, 1, False, np.complex64) + + _check_branch_cut(np.arcsinh, [ -2, 2, 0], [1j, 1j, 1], 1, 1, False, np.complex64) + _check_branch_cut(np.arccosh, [0-2j, 2j, 2], [1, 1, 1j], 1, 1, False, np.complex64) + _check_branch_cut(np.arctanh, [0-2j, 2j, 0], [1, 1, 1j], 1, 1, False, np.complex64) + + def test_against_cmath(self): + import cmath + + points = [-1-1j, -1+1j, +1-1j, +1+1j] + name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan', + 'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'} + atol = 4*np.finfo(complex).eps + for func in self.funcs: + fname = func.__name__.split('.')[-1] + cname = name_map.get(fname, fname) + try: + cfunc = getattr(cmath, cname) + except AttributeError: + continue + for p in points: + a = complex(func(np.complex128(p))) + b = cfunc(p) + assert_( + abs(a - b) < atol, + "%s %s: %s; cmath: %s" % (fname, p, a, b) + ) + + @pytest.mark.xfail( + # manylinux2014 uses glibc2.17 + _glibc_older_than("2.18"), + reason="Older glibc versions are imprecise (maybe passes with SIMD?)" + ) + @pytest.mark.xfail(IS_WASM, reason="doesn't work") + @pytest.mark.parametrize('dtype', [ + np.complex64, np.complex128, np.clongdouble + ]) + def test_loss_of_precision(self, dtype): + """Check loss of precision in complex arc* functions""" + if dtype is np.clongdouble and platform.machine() != 'x86_64': + # Failures on musllinux, aarch64, s390x, ppc64le (see gh-17554) + pytest.skip('Only works reliably for x86-64 and recent glibc') + + # Check against known-good functions + + info = np.finfo(dtype) + real_dtype = dtype(0.).real.dtype + eps = info.eps + + def check(x, rtol): + x = x.astype(real_dtype) + + z = x.astype(dtype) + d = np.absolute(np.arcsinh(x)/np.arcsinh(z).real - 1) + assert_(np.all(d < rtol), (np.argmax(d), x[np.argmax(d)], d.max(), + 'arcsinh')) + + z = (1j*x).astype(dtype) + d = np.absolute(np.arcsinh(x)/np.arcsin(z).imag - 1) + assert_(np.all(d < rtol), (np.argmax(d), x[np.argmax(d)], d.max(), + 'arcsin')) + + z = x.astype(dtype) + d = np.absolute(np.arctanh(x)/np.arctanh(z).real - 1) + assert_(np.all(d < rtol), (np.argmax(d), x[np.argmax(d)], d.max(), + 'arctanh')) + + z = (1j*x).astype(dtype) + d = np.absolute(np.arctanh(x)/np.arctan(z).imag - 1) + assert_(np.all(d < rtol), (np.argmax(d), x[np.argmax(d)], d.max(), + 'arctan')) + + # The switchover was chosen as 1e-3; hence there can be up to + # ~eps/1e-3 of relative cancellation error before it + + x_series = np.logspace(-20, -3.001, 200) + x_basic = np.logspace(-2.999, 0, 10, endpoint=False) + + if dtype is np.clongdouble: + if bad_arcsinh(): + pytest.skip("Trig functions of np.clongdouble values known " + "to be inaccurate on aarch64 and PPC for some " + "compilation configurations.") + # It's not guaranteed that the system-provided arc functions + # are accurate down to a few epsilons. (Eg. on Linux 64-bit) + # So, give more leeway for long complex tests here: + check(x_series, 50.0*eps) + else: + check(x_series, 2.1*eps) + check(x_basic, 2.0*eps/1e-3) + + # Check a few points + + z = np.array([1e-5*(1+1j)], dtype=dtype) + p = 9.999999999333333333e-6 + 1.000000000066666666e-5j + d = np.absolute(1-np.arctanh(z)/p) + assert_(np.all(d < 1e-15)) + + p = 1.0000000000333333333e-5 + 9.999999999666666667e-6j + d = np.absolute(1-np.arcsinh(z)/p) + assert_(np.all(d < 1e-15)) + + p = 9.999999999333333333e-6j + 1.000000000066666666e-5 + d = np.absolute(1-np.arctan(z)/p) + assert_(np.all(d < 1e-15)) + + p = 1.0000000000333333333e-5j + 9.999999999666666667e-6 + d = np.absolute(1-np.arcsin(z)/p) + assert_(np.all(d < 1e-15)) + + # Check continuity across switchover points + + def check(func, z0, d=1): + z0 = np.asarray(z0, dtype=dtype) + zp = z0 + abs(z0) * d * eps * 2 + zm = z0 - abs(z0) * d * eps * 2 + assert_(np.all(zp != zm), (zp, zm)) + + # NB: the cancellation error at the switchover is at least eps + good = (abs(func(zp) - func(zm)) < 2*eps) + assert_(np.all(good), (func, z0[~good])) + + for func in (np.arcsinh, np.arcsinh, np.arcsin, np.arctanh, np.arctan): + pts = [rp+1j*ip for rp in (-1e-3, 0, 1e-3) for ip in(-1e-3, 0, 1e-3) + if rp != 0 or ip != 0] + check(func, pts, 1) + check(func, pts, 1j) + check(func, pts, 1+1j) + + @np.errstate(all="ignore") + def test_promotion_corner_cases(self): + for func in self.funcs: + assert func(np.float16(1)).dtype == np.float16 + # Integer to low precision float promotion is a dubious choice: + assert func(np.uint8(1)).dtype == np.float16 + assert func(np.int16(1)).dtype == np.float32 + + +class TestAttributes: + def test_attributes(self): + add = ncu.add + assert_equal(add.__name__, 'add') + assert_(add.ntypes >= 18) # don't fail if types added + assert_('ii->i' in add.types) + assert_equal(add.nin, 2) + assert_equal(add.nout, 1) + assert_equal(add.identity, 0) + + def test_doc(self): + # don't bother checking the long list of kwargs, which are likely to + # change + assert_(ncu.add.__doc__.startswith( + "add(x1, x2, /, out=None, *, where=True")) + assert_(ncu.frexp.__doc__.startswith( + "frexp(x[, out1, out2], / [, out=(None, None)], *, where=True")) + + +class TestSubclass: + + def test_subclass_op(self): + + class simple(np.ndarray): + def __new__(subtype, shape): + self = np.ndarray.__new__(subtype, shape, dtype=object) + self.fill(0) + return self + + a = simple((3, 4)) + assert_equal(a+a, a) + + +class TestFrompyfunc: + + def test_identity(self): + def mul(a, b): + return a * b + + # with identity=value + mul_ufunc = np.frompyfunc(mul, nin=2, nout=1, identity=1) + assert_equal(mul_ufunc.reduce([2, 3, 4]), 24) + assert_equal(mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)), 1) + assert_equal(mul_ufunc.reduce([]), 1) + + # with identity=None (reorderable) + mul_ufunc = np.frompyfunc(mul, nin=2, nout=1, identity=None) + assert_equal(mul_ufunc.reduce([2, 3, 4]), 24) + assert_equal(mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)), 1) + assert_raises(ValueError, lambda: mul_ufunc.reduce([])) + + # with no identity (not reorderable) + mul_ufunc = np.frompyfunc(mul, nin=2, nout=1) + assert_equal(mul_ufunc.reduce([2, 3, 4]), 24) + assert_raises(ValueError, lambda: mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1))) + assert_raises(ValueError, lambda: mul_ufunc.reduce([])) + + +def _check_branch_cut(f, x0, dx, re_sign=1, im_sign=-1, sig_zero_ok=False, + dtype=complex): + """ + Check for a branch cut in a function. + + Assert that `x0` lies on a branch cut of function `f` and `f` is + continuous from the direction `dx`. + + Parameters + ---------- + f : func + Function to check + x0 : array-like + Point on branch cut + dx : array-like + Direction to check continuity in + re_sign, im_sign : {1, -1} + Change of sign of the real or imaginary part expected + sig_zero_ok : bool + Whether to check if the branch cut respects signed zero (if applicable) + dtype : dtype + Dtype to check (should be complex) + + """ + x0 = np.atleast_1d(x0).astype(dtype) + dx = np.atleast_1d(dx).astype(dtype) + + if np.dtype(dtype).char == 'F': + scale = np.finfo(dtype).eps * 1e2 + atol = np.float32(1e-2) + else: + scale = np.finfo(dtype).eps * 1e3 + atol = 1e-4 + + y0 = f(x0) + yp = f(x0 + dx*scale*np.absolute(x0)/np.absolute(dx)) + ym = f(x0 - dx*scale*np.absolute(x0)/np.absolute(dx)) + + assert_(np.all(np.absolute(y0.real - yp.real) < atol), (y0, yp)) + assert_(np.all(np.absolute(y0.imag - yp.imag) < atol), (y0, yp)) + assert_(np.all(np.absolute(y0.real - ym.real*re_sign) < atol), (y0, ym)) + assert_(np.all(np.absolute(y0.imag - ym.imag*im_sign) < atol), (y0, ym)) + + if sig_zero_ok: + # check that signed zeros also work as a displacement + jr = (x0.real == 0) & (dx.real != 0) + ji = (x0.imag == 0) & (dx.imag != 0) + if np.any(jr): + x = x0[jr] + x.real = ncu.NZERO + ym = f(x) + assert_(np.all(np.absolute(y0[jr].real - ym.real*re_sign) < atol), (y0[jr], ym)) + assert_(np.all(np.absolute(y0[jr].imag - ym.imag*im_sign) < atol), (y0[jr], ym)) + + if np.any(ji): + x = x0[ji] + x.imag = ncu.NZERO + ym = f(x) + assert_(np.all(np.absolute(y0[ji].real - ym.real*re_sign) < atol), (y0[ji], ym)) + assert_(np.all(np.absolute(y0[ji].imag - ym.imag*im_sign) < atol), (y0[ji], ym)) + +def test_copysign(): + assert_(np.copysign(1, -1) == -1) + with np.errstate(divide="ignore"): + assert_(1 / np.copysign(0, -1) < 0) + assert_(1 / np.copysign(0, 1) > 0) + assert_(np.signbit(np.copysign(np.nan, -1))) + assert_(not np.signbit(np.copysign(np.nan, 1))) + +def _test_nextafter(t): + one = t(1) + two = t(2) + zero = t(0) + eps = np.finfo(t).eps + assert_(np.nextafter(one, two) - one == eps) + assert_(np.nextafter(one, zero) - one < 0) + assert_(np.isnan(np.nextafter(np.nan, one))) + assert_(np.isnan(np.nextafter(one, np.nan))) + assert_(np.nextafter(one, one) == one) + +def test_nextafter(): + return _test_nextafter(np.float64) + + +def test_nextafterf(): + return _test_nextafter(np.float32) + + +@pytest.mark.skipif(np.finfo(np.double) == np.finfo(np.longdouble), + reason="long double is same as double") +@pytest.mark.xfail(condition=platform.machine().startswith("ppc64"), + reason="IBM double double") +def test_nextafterl(): + return _test_nextafter(np.longdouble) + + +def test_nextafter_0(): + for t, direction in itertools.product(np._core.sctypes['float'], (1, -1)): + # The value of tiny for double double is NaN, so we need to pass the + # assert + with suppress_warnings() as sup: + sup.filter(UserWarning) + if not np.isnan(np.finfo(t).tiny): + tiny = np.finfo(t).tiny + assert_( + 0. < direction * np.nextafter(t(0), t(direction)) < tiny) + assert_equal(np.nextafter(t(0), t(direction)) / t(2.1), direction * 0.0) + +def _test_spacing(t): + one = t(1) + eps = np.finfo(t).eps + nan = t(np.nan) + inf = t(np.inf) + with np.errstate(invalid='ignore'): + assert_equal(np.spacing(one), eps) + assert_(np.isnan(np.spacing(nan))) + assert_(np.isnan(np.spacing(inf))) + assert_(np.isnan(np.spacing(-inf))) + assert_(np.spacing(t(1e30)) != 0) + +def test_spacing(): + return _test_spacing(np.float64) + +def test_spacingf(): + return _test_spacing(np.float32) + + +@pytest.mark.skipif(np.finfo(np.double) == np.finfo(np.longdouble), + reason="long double is same as double") +@pytest.mark.xfail(condition=platform.machine().startswith("ppc64"), + reason="IBM double double") +def test_spacingl(): + return _test_spacing(np.longdouble) + +def test_spacing_gfortran(): + # Reference from this fortran file, built with gfortran 4.3.3 on linux + # 32bits: + # PROGRAM test_spacing + # INTEGER, PARAMETER :: SGL = SELECTED_REAL_KIND(p=6, r=37) + # INTEGER, PARAMETER :: DBL = SELECTED_REAL_KIND(p=13, r=200) + # + # WRITE(*,*) spacing(0.00001_DBL) + # WRITE(*,*) spacing(1.0_DBL) + # WRITE(*,*) spacing(1000._DBL) + # WRITE(*,*) spacing(10500._DBL) + # + # WRITE(*,*) spacing(0.00001_SGL) + # WRITE(*,*) spacing(1.0_SGL) + # WRITE(*,*) spacing(1000._SGL) + # WRITE(*,*) spacing(10500._SGL) + # END PROGRAM + ref = {np.float64: [1.69406589450860068E-021, + 2.22044604925031308E-016, + 1.13686837721616030E-013, + 1.81898940354585648E-012], + np.float32: [9.09494702E-13, + 1.19209290E-07, + 6.10351563E-05, + 9.76562500E-04]} + + for dt, dec_ in zip([np.float32, np.float64], (10, 20)): + x = np.array([1e-5, 1, 1000, 10500], dtype=dt) + assert_array_almost_equal(np.spacing(x), ref[dt], decimal=dec_) + +def test_nextafter_vs_spacing(): + # XXX: spacing does not handle long double yet + for t in [np.float32, np.float64]: + for _f in [1, 1e-5, 1000]: + f = t(_f) + f1 = t(_f + 1) + assert_(np.nextafter(f, f1) - f == np.spacing(f)) + +def test_pos_nan(): + """Check np.nan is a positive nan.""" + assert_(np.signbit(np.nan) == 0) + +def test_reduceat(): + """Test bug in reduceat when structured arrays are not copied.""" + db = np.dtype([('name', 'S11'), ('time', np.int64), ('value', np.float32)]) + a = np.empty([100], dtype=db) + a['name'] = 'Simple' + a['time'] = 10 + a['value'] = 100 + indx = [0, 7, 15, 25] + + h2 = [] + val1 = indx[0] + for val2 in indx[1:]: + h2.append(np.add.reduce(a['value'][val1:val2])) + val1 = val2 + h2.append(np.add.reduce(a['value'][val1:])) + h2 = np.array(h2) + + # test buffered -- this should work + h1 = np.add.reduceat(a['value'], indx) + assert_array_almost_equal(h1, h2) + + # This is when the error occurs. + # test no buffer + np.setbufsize(32) + h1 = np.add.reduceat(a['value'], indx) + np.setbufsize(ncu.UFUNC_BUFSIZE_DEFAULT) + assert_array_almost_equal(h1, h2) + +def test_reduceat_empty(): + """Reduceat should work with empty arrays""" + indices = np.array([], 'i4') + x = np.array([], 'f8') + result = np.add.reduceat(x, indices) + assert_equal(result.dtype, x.dtype) + assert_equal(result.shape, (0,)) + # Another case with a slightly different zero-sized shape + x = np.ones((5, 2)) + result = np.add.reduceat(x, [], axis=0) + assert_equal(result.dtype, x.dtype) + assert_equal(result.shape, (0, 2)) + result = np.add.reduceat(x, [], axis=1) + assert_equal(result.dtype, x.dtype) + assert_equal(result.shape, (5, 0)) + +def test_complex_nan_comparisons(): + nans = [complex(np.nan, 0), complex(0, np.nan), complex(np.nan, np.nan)] + fins = [complex(1, 0), complex(-1, 0), complex(0, 1), complex(0, -1), + complex(1, 1), complex(-1, -1), complex(0, 0)] + + with np.errstate(invalid='ignore'): + for x in nans + fins: + x = np.array([x]) + for y in nans + fins: + y = np.array([y]) + + if np.isfinite(x) and np.isfinite(y): + continue + + assert_equal(x < y, False, err_msg="%r < %r" % (x, y)) + assert_equal(x > y, False, err_msg="%r > %r" % (x, y)) + assert_equal(x <= y, False, err_msg="%r <= %r" % (x, y)) + assert_equal(x >= y, False, err_msg="%r >= %r" % (x, y)) + assert_equal(x == y, False, err_msg="%r == %r" % (x, y)) + + +def test_rint_big_int(): + # np.rint bug for large integer values on Windows 32-bit and MKL + # https://github.com/numpy/numpy/issues/6685 + val = 4607998452777363968 + # This is exactly representable in floating point + assert_equal(val, int(float(val))) + # Rint should not change the value + assert_equal(val, np.rint(val)) + + +@pytest.mark.parametrize('ftype', [np.float32, np.float64]) +def test_memoverlap_accumulate(ftype): + # Reproduces bug https://github.com/numpy/numpy/issues/15597 + arr = np.array([0.61, 0.60, 0.77, 0.41, 0.19], dtype=ftype) + out_max = np.array([0.61, 0.61, 0.77, 0.77, 0.77], dtype=ftype) + out_min = np.array([0.61, 0.60, 0.60, 0.41, 0.19], dtype=ftype) + assert_equal(np.maximum.accumulate(arr), out_max) + assert_equal(np.minimum.accumulate(arr), out_min) + +@pytest.mark.parametrize("ufunc, dtype", [ + (ufunc, t[0]) + for ufunc in UFUNCS_BINARY_ACC + for t in ufunc.types + if t[-1] == '?' and t[0] not in 'DFGMmO' +]) +def test_memoverlap_accumulate_cmp(ufunc, dtype): + if ufunc.signature: + pytest.skip('For generic signatures only') + for size in (2, 8, 32, 64, 128, 256): + arr = np.array([0, 1, 1]*size, dtype=dtype) + acc = ufunc.accumulate(arr, dtype='?') + acc_u8 = acc.view(np.uint8) + exp = np.array(list(itertools.accumulate(arr, ufunc)), dtype=np.uint8) + assert_equal(exp, acc_u8) + +@pytest.mark.parametrize("ufunc, dtype", [ + (ufunc, t[0]) + for ufunc in UFUNCS_BINARY_ACC + for t in ufunc.types + if t[0] == t[1] and t[0] == t[-1] and t[0] not in 'DFGMmO?' +]) +def test_memoverlap_accumulate_symmetric(ufunc, dtype): + if ufunc.signature: + pytest.skip('For generic signatures only') + with np.errstate(all='ignore'): + for size in (2, 8, 32, 64, 128, 256): + arr = np.array([0, 1, 2]*size).astype(dtype) + acc = ufunc.accumulate(arr, dtype=dtype) + exp = np.array(list(itertools.accumulate(arr, ufunc)), dtype=dtype) + assert_equal(exp, acc) + +def test_signaling_nan_exceptions(): + with assert_no_warnings(): + a = np.ndarray(shape=(), dtype='float32', buffer=b'\x00\xe0\xbf\xff') + np.isnan(a) + +@pytest.mark.parametrize("arr", [ + np.arange(2), + np.matrix([0, 1]), + np.matrix([[0, 1], [2, 5]]), + ]) +def test_outer_subclass_preserve(arr): + # for gh-8661 + class foo(np.ndarray): + pass + actual = np.multiply.outer(arr.view(foo), arr.view(foo)) + assert actual.__class__.__name__ == 'foo' + +def test_outer_bad_subclass(): + class BadArr1(np.ndarray): + def __array_finalize__(self, obj): + # The outer call reshapes to 3 dims, try to do a bad reshape. + if self.ndim == 3: + self.shape = self.shape + (1,) + + class BadArr2(np.ndarray): + def __array_finalize__(self, obj): + if isinstance(obj, BadArr2): + # outer inserts 1-sized dims. In that case disturb them. + if self.shape[-1] == 1: + self.shape = self.shape[::-1] + + for cls in [BadArr1, BadArr2]: + arr = np.ones((2, 3)).view(cls) + with assert_raises(TypeError) as a: + # The first array gets reshaped (not the second one) + np.add.outer(arr, [1, 2]) + + # This actually works, since we only see the reshaping error: + arr = np.ones((2, 3)).view(cls) + assert type(np.add.outer([1, 2], arr)) is cls + +def test_outer_exceeds_maxdims(): + deep = np.ones((1,) * 33) + with assert_raises(ValueError): + np.add.outer(deep, deep) + +def test_bad_legacy_ufunc_silent_errors(): + # legacy ufuncs can't report errors and NumPy can't check if the GIL + # is released. So NumPy has to check after the GIL is released just to + # cover all bases. `np.power` uses/used to use this. + arr = np.arange(3).astype(np.float64) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error(arr, arr) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + # not contiguous means the fast-path cannot be taken + non_contig = arr.repeat(20).reshape(-1, 6)[:, ::2] + ncu_tests.always_error(non_contig, arr) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error.outer(arr, arr) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error.reduce(arr) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error.reduceat(arr, [0, 1]) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error.accumulate(arr) + + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error.at(arr, [0, 1, 2], arr) + + +@pytest.mark.parametrize('x1', [np.arange(3.0), [0.0, 1.0, 2.0]]) +def test_bad_legacy_gufunc_silent_errors(x1): + # Verify that an exception raised in a gufunc loop propagates correctly. + # The signature of always_error_gufunc is '(i),()->()'. + with pytest.raises(RuntimeError, match=r"How unexpected :\)!"): + ncu_tests.always_error_gufunc(x1, 0.0) + + +class TestAddDocstring: + @pytest.mark.skipif(sys.flags.optimize == 2, reason="Python running -OO") + @pytest.mark.skipif(IS_PYPY, reason="PyPy does not modify tp_doc") + def test_add_same_docstring(self): + # test for attributes (which are C-level defined) + ncu.add_docstring(np.ndarray.flat, np.ndarray.flat.__doc__) + + # And typical functions: + def func(): + """docstring""" + return + + ncu.add_docstring(func, func.__doc__) + + @pytest.mark.skipif(sys.flags.optimize == 2, reason="Python running -OO") + def test_different_docstring_fails(self): + # test for attributes (which are C-level defined) + with assert_raises(RuntimeError): + ncu.add_docstring(np.ndarray.flat, "different docstring") + + # And typical functions: + def func(): + """docstring""" + return + + with assert_raises(RuntimeError): + ncu.add_docstring(func, "different docstring") + + +class TestAdd_newdoc_ufunc: + @pytest.mark.filterwarnings("ignore:_add_newdoc_ufunc:DeprecationWarning") + def test_ufunc_arg(self): + assert_raises(TypeError, ncu._add_newdoc_ufunc, 2, "blah") + assert_raises(ValueError, ncu._add_newdoc_ufunc, np.add, "blah") + + @pytest.mark.filterwarnings("ignore:_add_newdoc_ufunc:DeprecationWarning") + def test_string_arg(self): + assert_raises(TypeError, ncu._add_newdoc_ufunc, np.add, 3) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath_accuracy.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath_accuracy.py new file mode 100644 index 0000000000000000000000000000000000000000..71967dd4cfff6fdab83206aacdbfee3c0e264756 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath_accuracy.py @@ -0,0 +1,121 @@ +import numpy as np +import os +from os import path +import sys +import pytest +from ctypes import c_longlong, c_double, c_float, c_int, cast, pointer, POINTER +from numpy.testing import assert_array_max_ulp +from numpy.testing._private.utils import _glibc_older_than +from numpy._core._multiarray_umath import __cpu_features__ + +UNARY_UFUNCS = [obj for obj in np._core.umath.__dict__.values() if + isinstance(obj, np.ufunc)] +UNARY_OBJECT_UFUNCS = [uf for uf in UNARY_UFUNCS if "O->O" in uf.types] + +# Remove functions that do not support `floats` +UNARY_OBJECT_UFUNCS.remove(np.invert) +UNARY_OBJECT_UFUNCS.remove(np.bitwise_count) + +IS_AVX = __cpu_features__.get('AVX512F', False) or \ + (__cpu_features__.get('FMA3', False) and __cpu_features__.get('AVX2', False)) + +IS_AVX512FP16 = __cpu_features__.get('AVX512FP16', False) + +# only run on linux with AVX, also avoid old glibc (numpy/numpy#20448). +runtest = (sys.platform.startswith('linux') + and IS_AVX and not _glibc_older_than("2.17")) +platform_skip = pytest.mark.skipif(not runtest, + reason="avoid testing inconsistent platform " + "library implementations") + +# convert string to hex function taken from: +# https://stackoverflow.com/questions/1592158/convert-hex-to-float # +def convert(s, datatype="np.float32"): + i = int(s, 16) # convert from hex to a Python int + if (datatype == "np.float64"): + cp = pointer(c_longlong(i)) # make this into a c long long integer + fp = cast(cp, POINTER(c_double)) # cast the int pointer to a double pointer + else: + cp = pointer(c_int(i)) # make this into a c integer + fp = cast(cp, POINTER(c_float)) # cast the int pointer to a float pointer + + return fp.contents.value # dereference the pointer, get the float + +str_to_float = np.vectorize(convert) + +class TestAccuracy: + @platform_skip + def test_validate_transcendentals(self): + with np.errstate(all='ignore'): + data_dir = path.join(path.dirname(__file__), 'data') + files = os.listdir(data_dir) + files = list(filter(lambda f: f.endswith('.csv'), files)) + for filename in files: + filepath = path.join(data_dir, filename) + with open(filepath) as fid: + file_without_comments = ( + r for r in fid if r[0] not in ('$', '#') + ) + data = np.genfromtxt(file_without_comments, + dtype=('|S39','|S39','|S39',int), + names=('type','input','output','ulperr'), + delimiter=',', + skip_header=1) + npname = path.splitext(filename)[0].split('-')[3] + npfunc = getattr(np, npname) + for datatype in np.unique(data['type']): + data_subset = data[data['type'] == datatype] + inval = np.array(str_to_float(data_subset['input'].astype(str), data_subset['type'].astype(str)), dtype=eval(datatype)) + outval = np.array(str_to_float(data_subset['output'].astype(str), data_subset['type'].astype(str)), dtype=eval(datatype)) + perm = np.random.permutation(len(inval)) + inval = inval[perm] + outval = outval[perm] + maxulperr = data_subset['ulperr'].max() + assert_array_max_ulp(npfunc(inval), outval, maxulperr) + + @pytest.mark.skipif(IS_AVX512FP16, + reason = "SVML FP16 have slightly higher ULP errors") + @pytest.mark.parametrize("ufunc", UNARY_OBJECT_UFUNCS) + def test_validate_fp16_transcendentals(self, ufunc): + with np.errstate(all='ignore'): + arr = np.arange(65536, dtype=np.int16) + datafp16 = np.frombuffer(arr.tobytes(), dtype=np.float16) + datafp32 = datafp16.astype(np.float32) + assert_array_max_ulp(ufunc(datafp16), ufunc(datafp32), + maxulp=1, dtype=np.float16) + + @pytest.mark.skipif(not IS_AVX512FP16, + reason="lower ULP only apply for SVML FP16") + def test_validate_svml_fp16(self): + max_ulp_err = { + "arccos": 2.54, + "arccosh": 2.09, + "arcsin": 3.06, + "arcsinh": 1.51, + "arctan": 2.61, + "arctanh": 1.88, + "cbrt": 1.57, + "cos": 1.43, + "cosh": 1.33, + "exp2": 1.33, + "exp": 1.27, + "expm1": 0.53, + "log": 1.80, + "log10": 1.27, + "log1p": 1.88, + "log2": 1.80, + "sin": 1.88, + "sinh": 2.05, + "tan": 2.26, + "tanh": 3.00, + } + + with np.errstate(all='ignore'): + arr = np.arange(65536, dtype=np.int16) + datafp16 = np.frombuffer(arr.tobytes(), dtype=np.float16) + datafp32 = datafp16.astype(np.float32) + for func in max_ulp_err: + ufunc = getattr(np, func) + ulp = np.ceil(max_ulp_err[func]) + assert_array_max_ulp(ufunc(datafp16), ufunc(datafp32), + maxulp=ulp, dtype=np.float16) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath_complex.py b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath_complex.py new file mode 100644 index 0000000000000000000000000000000000000000..c7f4cec07b2187784f7659ac1e97207783b25d0b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/_core/tests/test_umath_complex.py @@ -0,0 +1,622 @@ +import sys +import platform +import pytest + +import numpy as np +# import the c-extension module directly since _arg is not exported via umath +import numpy._core._multiarray_umath as ncu +from numpy.testing import ( + assert_raises, assert_equal, assert_array_equal, assert_almost_equal, assert_array_max_ulp + ) + +# TODO: branch cuts (use Pauli code) +# TODO: conj 'symmetry' +# TODO: FPU exceptions + +# At least on Windows the results of many complex functions are not conforming +# to the C99 standard. See ticket 1574. +# Ditto for Solaris (ticket 1642) and OS X on PowerPC. +#FIXME: this will probably change when we require full C99 compatibility +with np.errstate(all='ignore'): + functions_seem_flaky = ((np.exp(complex(np.inf, 0)).imag != 0) + or (np.log(complex(ncu.NZERO, 0)).imag != np.pi)) +# TODO: replace with a check on whether platform-provided C99 funcs are used +xfail_complex_tests = (not sys.platform.startswith('linux') or functions_seem_flaky) + +# TODO This can be xfail when the generator functions are got rid of. +platform_skip = pytest.mark.skipif(xfail_complex_tests, + reason="Inadequate C99 complex support") + + + +class TestCexp: + def test_simple(self): + check = check_complex_value + f = np.exp + + check(f, 1, 0, np.exp(1), 0, False) + check(f, 0, 1, np.cos(1), np.sin(1), False) + + ref = np.exp(1) * complex(np.cos(1), np.sin(1)) + check(f, 1, 1, ref.real, ref.imag, False) + + @platform_skip + def test_special_values(self): + # C99: Section G 6.3.1 + + check = check_complex_value + f = np.exp + + # cexp(+-0 + 0i) is 1 + 0i + check(f, ncu.PZERO, 0, 1, 0, False) + check(f, ncu.NZERO, 0, 1, 0, False) + + # cexp(x + infi) is nan + nani for finite x and raises 'invalid' FPU + # exception + check(f, 1, np.inf, np.nan, np.nan) + check(f, -1, np.inf, np.nan, np.nan) + check(f, 0, np.inf, np.nan, np.nan) + + # cexp(inf + 0i) is inf + 0i + check(f, np.inf, 0, np.inf, 0) + + # cexp(-inf + yi) is +0 * (cos(y) + i sin(y)) for finite y + check(f, -np.inf, 1, ncu.PZERO, ncu.PZERO) + check(f, -np.inf, 0.75 * np.pi, ncu.NZERO, ncu.PZERO) + + # cexp(inf + yi) is +inf * (cos(y) + i sin(y)) for finite y + check(f, np.inf, 1, np.inf, np.inf) + check(f, np.inf, 0.75 * np.pi, -np.inf, np.inf) + + # cexp(-inf + inf i) is +-0 +- 0i (signs unspecified) + def _check_ninf_inf(dummy): + msgform = "cexp(-inf, inf) is (%f, %f), expected (+-0, +-0)" + with np.errstate(invalid='ignore'): + z = f(np.array(complex(-np.inf, np.inf))) + if z.real != 0 or z.imag != 0: + raise AssertionError(msgform % (z.real, z.imag)) + + _check_ninf_inf(None) + + # cexp(inf + inf i) is +-inf + NaNi and raised invalid FPU ex. + def _check_inf_inf(dummy): + msgform = "cexp(inf, inf) is (%f, %f), expected (+-inf, nan)" + with np.errstate(invalid='ignore'): + z = f(np.array(complex(np.inf, np.inf))) + if not np.isinf(z.real) or not np.isnan(z.imag): + raise AssertionError(msgform % (z.real, z.imag)) + + _check_inf_inf(None) + + # cexp(-inf + nan i) is +-0 +- 0i + def _check_ninf_nan(dummy): + msgform = "cexp(-inf, nan) is (%f, %f), expected (+-0, +-0)" + with np.errstate(invalid='ignore'): + z = f(np.array(complex(-np.inf, np.nan))) + if z.real != 0 or z.imag != 0: + raise AssertionError(msgform % (z.real, z.imag)) + + _check_ninf_nan(None) + + # cexp(inf + nan i) is +-inf + nan + def _check_inf_nan(dummy): + msgform = "cexp(-inf, nan) is (%f, %f), expected (+-inf, nan)" + with np.errstate(invalid='ignore'): + z = f(np.array(complex(np.inf, np.nan))) + if not np.isinf(z.real) or not np.isnan(z.imag): + raise AssertionError(msgform % (z.real, z.imag)) + + _check_inf_nan(None) + + # cexp(nan + yi) is nan + nani for y != 0 (optional: raises invalid FPU + # ex) + check(f, np.nan, 1, np.nan, np.nan) + check(f, np.nan, -1, np.nan, np.nan) + + check(f, np.nan, np.inf, np.nan, np.nan) + check(f, np.nan, -np.inf, np.nan, np.nan) + + # cexp(nan + nani) is nan + nani + check(f, np.nan, np.nan, np.nan, np.nan) + + # TODO This can be xfail when the generator functions are got rid of. + @pytest.mark.skip(reason="cexp(nan + 0I) is wrong on most platforms") + def test_special_values2(self): + # XXX: most implementations get it wrong here (including glibc <= 2.10) + # cexp(nan + 0i) is nan + 0i + check = check_complex_value + f = np.exp + + check(f, np.nan, 0, np.nan, 0) + +class TestClog: + def test_simple(self): + x = np.array([1+0j, 1+2j]) + y_r = np.log(np.abs(x)) + 1j * np.angle(x) + y = np.log(x) + assert_almost_equal(y, y_r) + + @platform_skip + @pytest.mark.skipif(platform.machine() == "armv5tel", reason="See gh-413.") + def test_special_values(self): + xl = [] + yl = [] + + # From C99 std (Sec 6.3.2) + # XXX: check exceptions raised + # --- raise for invalid fails. + + # clog(-0 + i0) returns -inf + i pi and raises the 'divide-by-zero' + # floating-point exception. + with np.errstate(divide='raise'): + x = np.array([ncu.NZERO], dtype=complex) + y = complex(-np.inf, np.pi) + assert_raises(FloatingPointError, np.log, x) + with np.errstate(divide='ignore'): + assert_almost_equal(np.log(x), y) + + xl.append(x) + yl.append(y) + + # clog(+0 + i0) returns -inf + i0 and raises the 'divide-by-zero' + # floating-point exception. + with np.errstate(divide='raise'): + x = np.array([0], dtype=complex) + y = complex(-np.inf, 0) + assert_raises(FloatingPointError, np.log, x) + with np.errstate(divide='ignore'): + assert_almost_equal(np.log(x), y) + + xl.append(x) + yl.append(y) + + # clog(x + i inf returns +inf + i pi /2, for finite x. + x = np.array([complex(1, np.inf)], dtype=complex) + y = complex(np.inf, 0.5 * np.pi) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + x = np.array([complex(-1, np.inf)], dtype=complex) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(x + iNaN) returns NaN + iNaN and optionally raises the + # 'invalid' floating- point exception, for finite x. + with np.errstate(invalid='raise'): + x = np.array([complex(1., np.nan)], dtype=complex) + y = complex(np.nan, np.nan) + #assert_raises(FloatingPointError, np.log, x) + with np.errstate(invalid='ignore'): + assert_almost_equal(np.log(x), y) + + xl.append(x) + yl.append(y) + + with np.errstate(invalid='raise'): + x = np.array([np.inf + 1j * np.nan], dtype=complex) + #assert_raises(FloatingPointError, np.log, x) + with np.errstate(invalid='ignore'): + assert_almost_equal(np.log(x), y) + + xl.append(x) + yl.append(y) + + # clog(- inf + iy) returns +inf + ipi , for finite positive-signed y. + x = np.array([-np.inf + 1j], dtype=complex) + y = complex(np.inf, np.pi) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(+ inf + iy) returns +inf + i0, for finite positive-signed y. + x = np.array([np.inf + 1j], dtype=complex) + y = complex(np.inf, 0) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(- inf + i inf) returns +inf + i3pi /4. + x = np.array([complex(-np.inf, np.inf)], dtype=complex) + y = complex(np.inf, 0.75 * np.pi) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(+ inf + i inf) returns +inf + ipi /4. + x = np.array([complex(np.inf, np.inf)], dtype=complex) + y = complex(np.inf, 0.25 * np.pi) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(+/- inf + iNaN) returns +inf + iNaN. + x = np.array([complex(np.inf, np.nan)], dtype=complex) + y = complex(np.inf, np.nan) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + x = np.array([complex(-np.inf, np.nan)], dtype=complex) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(NaN + iy) returns NaN + iNaN and optionally raises the + # 'invalid' floating-point exception, for finite y. + x = np.array([complex(np.nan, 1)], dtype=complex) + y = complex(np.nan, np.nan) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(NaN + i inf) returns +inf + iNaN. + x = np.array([complex(np.nan, np.inf)], dtype=complex) + y = complex(np.inf, np.nan) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(NaN + iNaN) returns NaN + iNaN. + x = np.array([complex(np.nan, np.nan)], dtype=complex) + y = complex(np.nan, np.nan) + assert_almost_equal(np.log(x), y) + xl.append(x) + yl.append(y) + + # clog(conj(z)) = conj(clog(z)). + xa = np.array(xl, dtype=complex) + ya = np.array(yl, dtype=complex) + with np.errstate(divide='ignore'): + for i in range(len(xa)): + assert_almost_equal(np.log(xa[i].conj()), ya[i].conj()) + + +class TestCsqrt: + + def test_simple(self): + # sqrt(1) + check_complex_value(np.sqrt, 1, 0, 1, 0) + + # sqrt(1i) + rres = 0.5*np.sqrt(2) + ires = rres + check_complex_value(np.sqrt, 0, 1, rres, ires, False) + + # sqrt(-1) + check_complex_value(np.sqrt, -1, 0, 0, 1) + + def test_simple_conjugate(self): + ref = np.conj(np.sqrt(complex(1, 1))) + + def f(z): + return np.sqrt(np.conj(z)) + + check_complex_value(f, 1, 1, ref.real, ref.imag, False) + + #def test_branch_cut(self): + # _check_branch_cut(f, -1, 0, 1, -1) + + @platform_skip + def test_special_values(self): + # C99: Sec G 6.4.2 + + check = check_complex_value + f = np.sqrt + + # csqrt(+-0 + 0i) is 0 + 0i + check(f, ncu.PZERO, 0, 0, 0) + check(f, ncu.NZERO, 0, 0, 0) + + # csqrt(x + infi) is inf + infi for any x (including NaN) + check(f, 1, np.inf, np.inf, np.inf) + check(f, -1, np.inf, np.inf, np.inf) + + check(f, ncu.PZERO, np.inf, np.inf, np.inf) + check(f, ncu.NZERO, np.inf, np.inf, np.inf) + check(f, np.inf, np.inf, np.inf, np.inf) + check(f, -np.inf, np.inf, np.inf, np.inf) + check(f, -np.nan, np.inf, np.inf, np.inf) + + # csqrt(x + nani) is nan + nani for any finite x + check(f, 1, np.nan, np.nan, np.nan) + check(f, -1, np.nan, np.nan, np.nan) + check(f, 0, np.nan, np.nan, np.nan) + + # csqrt(-inf + yi) is +0 + infi for any finite y > 0 + check(f, -np.inf, 1, ncu.PZERO, np.inf) + + # csqrt(inf + yi) is +inf + 0i for any finite y > 0 + check(f, np.inf, 1, np.inf, ncu.PZERO) + + # csqrt(-inf + nani) is nan +- infi (both +i infi are valid) + def _check_ninf_nan(dummy): + msgform = "csqrt(-inf, nan) is (%f, %f), expected (nan, +-inf)" + z = np.sqrt(np.array(complex(-np.inf, np.nan))) + #Fixme: ugly workaround for isinf bug. + with np.errstate(invalid='ignore'): + if not (np.isnan(z.real) and np.isinf(z.imag)): + raise AssertionError(msgform % (z.real, z.imag)) + + _check_ninf_nan(None) + + # csqrt(+inf + nani) is inf + nani + check(f, np.inf, np.nan, np.inf, np.nan) + + # csqrt(nan + yi) is nan + nani for any finite y (infinite handled in x + # + nani) + check(f, np.nan, 0, np.nan, np.nan) + check(f, np.nan, 1, np.nan, np.nan) + check(f, np.nan, np.nan, np.nan, np.nan) + + # XXX: check for conj(csqrt(z)) == csqrt(conj(z)) (need to fix branch + # cuts first) + +class TestCpow: + def setup_method(self): + self.olderr = np.seterr(invalid='ignore') + + def teardown_method(self): + np.seterr(**self.olderr) + + def test_simple(self): + x = np.array([1+1j, 0+2j, 1+2j, np.inf, np.nan]) + y_r = x ** 2 + y = np.power(x, 2) + assert_almost_equal(y, y_r) + + def test_scalar(self): + x = np.array([1, 1j, 2, 2.5+.37j, np.inf, np.nan]) + y = np.array([1, 1j, -0.5+1.5j, -0.5+1.5j, 2, 3]) + lx = list(range(len(x))) + + # Hardcode the expected `builtins.complex` values, + # as complex exponentiation is broken as of bpo-44698 + p_r = [ + 1+0j, + 0.20787957635076193+0j, + 0.35812203996480685+0.6097119028618724j, + 0.12659112128185032+0.48847676699581527j, + complex(np.inf, np.nan), + complex(np.nan, np.nan), + ] + + n_r = [x[i] ** y[i] for i in lx] + for i in lx: + assert_almost_equal(n_r[i], p_r[i], err_msg='Loop %d\n' % i) + + def test_array(self): + x = np.array([1, 1j, 2, 2.5+.37j, np.inf, np.nan]) + y = np.array([1, 1j, -0.5+1.5j, -0.5+1.5j, 2, 3]) + lx = list(range(len(x))) + + # Hardcode the expected `builtins.complex` values, + # as complex exponentiation is broken as of bpo-44698 + p_r = [ + 1+0j, + 0.20787957635076193+0j, + 0.35812203996480685+0.6097119028618724j, + 0.12659112128185032+0.48847676699581527j, + complex(np.inf, np.nan), + complex(np.nan, np.nan), + ] + + n_r = x ** y + for i in lx: + assert_almost_equal(n_r[i], p_r[i], err_msg='Loop %d\n' % i) + +class TestCabs: + def setup_method(self): + self.olderr = np.seterr(invalid='ignore') + + def teardown_method(self): + np.seterr(**self.olderr) + + def test_simple(self): + x = np.array([1+1j, 0+2j, 1+2j, np.inf, np.nan]) + y_r = np.array([np.sqrt(2.), 2, np.sqrt(5), np.inf, np.nan]) + y = np.abs(x) + assert_almost_equal(y, y_r) + + def test_fabs(self): + # Test that np.abs(x +- 0j) == np.abs(x) (as mandated by C99 for cabs) + x = np.array([1+0j], dtype=complex) + assert_array_equal(np.abs(x), np.real(x)) + + x = np.array([complex(1, ncu.NZERO)], dtype=complex) + assert_array_equal(np.abs(x), np.real(x)) + + x = np.array([complex(np.inf, ncu.NZERO)], dtype=complex) + assert_array_equal(np.abs(x), np.real(x)) + + x = np.array([complex(np.nan, ncu.NZERO)], dtype=complex) + assert_array_equal(np.abs(x), np.real(x)) + + def test_cabs_inf_nan(self): + x, y = [], [] + + # cabs(+-nan + nani) returns nan + x.append(np.nan) + y.append(np.nan) + check_real_value(np.abs, np.nan, np.nan, np.nan) + + x.append(np.nan) + y.append(-np.nan) + check_real_value(np.abs, -np.nan, np.nan, np.nan) + + # According to C99 standard, if exactly one of the real/part is inf and + # the other nan, then cabs should return inf + x.append(np.inf) + y.append(np.nan) + check_real_value(np.abs, np.inf, np.nan, np.inf) + + x.append(-np.inf) + y.append(np.nan) + check_real_value(np.abs, -np.inf, np.nan, np.inf) + + # cabs(conj(z)) == conj(cabs(z)) (= cabs(z)) + def f(a): + return np.abs(np.conj(a)) + + def g(a, b): + return np.abs(complex(a, b)) + + xa = np.array(x, dtype=complex) + assert len(xa) == len(x) == len(y) + for xi, yi in zip(x, y): + ref = g(xi, yi) + check_real_value(f, xi, yi, ref) + +class TestCarg: + def test_simple(self): + check_real_value(ncu._arg, 1, 0, 0, False) + check_real_value(ncu._arg, 0, 1, 0.5*np.pi, False) + + check_real_value(ncu._arg, 1, 1, 0.25*np.pi, False) + check_real_value(ncu._arg, ncu.PZERO, ncu.PZERO, ncu.PZERO) + + # TODO This can be xfail when the generator functions are got rid of. + @pytest.mark.skip( + reason="Complex arithmetic with signed zero fails on most platforms") + def test_zero(self): + # carg(-0 +- 0i) returns +- pi + check_real_value(ncu._arg, ncu.NZERO, ncu.PZERO, np.pi, False) + check_real_value(ncu._arg, ncu.NZERO, ncu.NZERO, -np.pi, False) + + # carg(+0 +- 0i) returns +- 0 + check_real_value(ncu._arg, ncu.PZERO, ncu.PZERO, ncu.PZERO) + check_real_value(ncu._arg, ncu.PZERO, ncu.NZERO, ncu.NZERO) + + # carg(x +- 0i) returns +- 0 for x > 0 + check_real_value(ncu._arg, 1, ncu.PZERO, ncu.PZERO, False) + check_real_value(ncu._arg, 1, ncu.NZERO, ncu.NZERO, False) + + # carg(x +- 0i) returns +- pi for x < 0 + check_real_value(ncu._arg, -1, ncu.PZERO, np.pi, False) + check_real_value(ncu._arg, -1, ncu.NZERO, -np.pi, False) + + # carg(+- 0 + yi) returns pi/2 for y > 0 + check_real_value(ncu._arg, ncu.PZERO, 1, 0.5 * np.pi, False) + check_real_value(ncu._arg, ncu.NZERO, 1, 0.5 * np.pi, False) + + # carg(+- 0 + yi) returns -pi/2 for y < 0 + check_real_value(ncu._arg, ncu.PZERO, -1, 0.5 * np.pi, False) + check_real_value(ncu._arg, ncu.NZERO, -1, -0.5 * np.pi, False) + + #def test_branch_cuts(self): + # _check_branch_cut(ncu._arg, -1, 1j, -1, 1) + + def test_special_values(self): + # carg(-np.inf +- yi) returns +-pi for finite y > 0 + check_real_value(ncu._arg, -np.inf, 1, np.pi, False) + check_real_value(ncu._arg, -np.inf, -1, -np.pi, False) + + # carg(np.inf +- yi) returns +-0 for finite y > 0 + check_real_value(ncu._arg, np.inf, 1, ncu.PZERO, False) + check_real_value(ncu._arg, np.inf, -1, ncu.NZERO, False) + + # carg(x +- np.infi) returns +-pi/2 for finite x + check_real_value(ncu._arg, 1, np.inf, 0.5 * np.pi, False) + check_real_value(ncu._arg, 1, -np.inf, -0.5 * np.pi, False) + + # carg(-np.inf +- np.infi) returns +-3pi/4 + check_real_value(ncu._arg, -np.inf, np.inf, 0.75 * np.pi, False) + check_real_value(ncu._arg, -np.inf, -np.inf, -0.75 * np.pi, False) + + # carg(np.inf +- np.infi) returns +-pi/4 + check_real_value(ncu._arg, np.inf, np.inf, 0.25 * np.pi, False) + check_real_value(ncu._arg, np.inf, -np.inf, -0.25 * np.pi, False) + + # carg(x + yi) returns np.nan if x or y is nan + check_real_value(ncu._arg, np.nan, 0, np.nan, False) + check_real_value(ncu._arg, 0, np.nan, np.nan, False) + + check_real_value(ncu._arg, np.nan, np.inf, np.nan, False) + check_real_value(ncu._arg, np.inf, np.nan, np.nan, False) + + +def check_real_value(f, x1, y1, x, exact=True): + z1 = np.array([complex(x1, y1)]) + if exact: + assert_equal(f(z1), x) + else: + assert_almost_equal(f(z1), x) + + +def check_complex_value(f, x1, y1, x2, y2, exact=True): + z1 = np.array([complex(x1, y1)]) + z2 = complex(x2, y2) + with np.errstate(invalid='ignore'): + if exact: + assert_equal(f(z1), z2) + else: + assert_almost_equal(f(z1), z2) + +class TestSpecialComplexAVX: + @pytest.mark.parametrize("stride", [-4,-2,-1,1,2,4]) + @pytest.mark.parametrize("astype", [np.complex64, np.complex128]) + def test_array(self, stride, astype): + arr = np.array([complex(np.nan , np.nan), + complex(np.nan , np.inf), + complex(np.inf , np.nan), + complex(np.inf , np.inf), + complex(0. , np.inf), + complex(np.inf , 0.), + complex(0. , 0.), + complex(0. , np.nan), + complex(np.nan , 0.)], dtype=astype) + abs_true = np.array([np.nan, np.inf, np.inf, np.inf, np.inf, np.inf, 0., np.nan, np.nan], dtype=arr.real.dtype) + sq_true = np.array([complex(np.nan, np.nan), + complex(np.nan, np.nan), + complex(np.nan, np.nan), + complex(np.nan, np.inf), + complex(-np.inf, np.nan), + complex(np.inf, np.nan), + complex(0., 0.), + complex(np.nan, np.nan), + complex(np.nan, np.nan)], dtype=astype) + with np.errstate(invalid='ignore'): + assert_equal(np.abs(arr[::stride]), abs_true[::stride]) + assert_equal(np.square(arr[::stride]), sq_true[::stride]) + +class TestComplexAbsoluteAVX: + @pytest.mark.parametrize("arraysize", [1,2,3,4,5,6,7,8,9,10,11,13,15,17,18,19]) + @pytest.mark.parametrize("stride", [-4,-3,-2,-1,1,2,3,4]) + @pytest.mark.parametrize("astype", [np.complex64, np.complex128]) + # test to ensure masking and strides work as intended in the AVX implementation + def test_array(self, arraysize, stride, astype): + arr = np.ones(arraysize, dtype=astype) + abs_true = np.ones(arraysize, dtype=arr.real.dtype) + assert_equal(np.abs(arr[::stride]), abs_true[::stride]) + +# Testcase taken as is from https://github.com/numpy/numpy/issues/16660 +class TestComplexAbsoluteMixedDTypes: + @pytest.mark.parametrize("stride", [-4,-3,-2,-1,1,2,3,4]) + @pytest.mark.parametrize("astype", [np.complex64, np.complex128]) + @pytest.mark.parametrize("func", ['abs', 'square', 'conjugate']) + + def test_array(self, stride, astype, func): + dtype = [('template_id', 'U') + uni_arr2 = str_arr.astype(' None: ... + +np.where(AR_b, 1) # E: No overload variant + +np.can_cast(AR_f8, 1) # E: incompatible type + +np.vdot(AR_M, AR_M) # E: incompatible type + +np.copyto(AR_LIKE_f, AR_f8) # E: incompatible type + +np.putmask(AR_LIKE_f, [True, True, False], 1.5) # E: incompatible type + +np.packbits(AR_f8) # E: incompatible type +np.packbits(AR_u1, bitorder=">") # E: incompatible type + +np.unpackbits(AR_i8) # E: incompatible type +np.unpackbits(AR_u1, bitorder=">") # E: incompatible type + +np.shares_memory(1, 1, max_work=i8) # E: incompatible type +np.may_share_memory(1, 1, max_work=i8) # E: incompatible type + +np.arange(M) # E: No overload variant +np.arange(stop=10) # E: No overload variant + +np.datetime_data(int) # E: incompatible type + +np.busday_offset("2012", 10) # E: No overload variant + +np.datetime_as_string("2012") # E: No overload variant + +np.char.compare_chararrays("a", b"a", "==", False) # E: No overload variant + +np.nested_iters([AR_i8, AR_i8]) # E: Missing positional argument +np.nested_iters([AR_i8, AR_i8], 0) # E: incompatible type +np.nested_iters([AR_i8, AR_i8], [0]) # E: incompatible type +np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["test"]) # E: incompatible type +np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["test"]]) # E: incompatible type +np.nested_iters([AR_i8, AR_i8], [[0], [1]], buffersize=1.0) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/ndarray.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/ndarray.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c20c89e5725d86d2072c0818a37e1fcb45d9f741 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/ndarray.pyi @@ -0,0 +1,11 @@ +import numpy as np + +# Ban setting dtype since mutating the type of the array in place +# makes having ndarray be generic over dtype impossible. Generally +# users should use `ndarray.view` in this situation anyway. See +# +# https://github.com/numpy/numpy-stubs/issues/7 +# +# for more context. +float_array = np.array([1.0]) +float_array.dtype = np.bool # E: Property "dtype" defined in "ndarray" is read-only diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/ndarray_misc.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/ndarray_misc.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9e2c8e292db973052d4e441e3370eca078f41000 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/ndarray_misc.pyi @@ -0,0 +1,36 @@ +""" +Tests for miscellaneous (non-magic) ``np.ndarray``/``np.generic`` methods. + +More extensive tests are performed for the methods' +function-based counterpart in `../from_numeric.py`. + +""" + +import numpy as np +import numpy.typing as npt + +f8: np.float64 +AR_f8: npt.NDArray[np.float64] +AR_M: npt.NDArray[np.datetime64] +AR_b: npt.NDArray[np.bool] + +ctypes_obj = AR_f8.ctypes + +f8.argpartition(0) # E: has no attribute +f8.diagonal() # E: has no attribute +f8.dot(1) # E: has no attribute +f8.nonzero() # E: has no attribute +f8.partition(0) # E: has no attribute +f8.put(0, 2) # E: has no attribute +f8.setfield(2, np.float64) # E: has no attribute +f8.sort() # E: has no attribute +f8.trace() # E: has no attribute + +AR_M.__complex__() # E: Invalid self argument +AR_b.__index__() # E: Invalid self argument + +AR_f8[1.5] # E: No overload variant +AR_f8["field_a"] # E: No overload variant +AR_f8[["field_a", "field_b"]] # E: Invalid index type + +AR_f8.__array_finalize__(object()) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/nditer.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/nditer.pyi new file mode 100644 index 0000000000000000000000000000000000000000..eb8a8e2911c831130d02c1fb3735c67c1aa5a9e0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/nditer.pyi @@ -0,0 +1,8 @@ +import numpy as np + +class Test(np.nditer): ... # E: Cannot inherit from final class + +np.nditer([0, 1], flags=["test"]) # E: incompatible type +np.nditer([0, 1], op_flags=[["test"]]) # E: incompatible type +np.nditer([0, 1], itershape=(1.0,)) # E: incompatible type +np.nditer([0, 1], buffersize=1.0) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/nested_sequence.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/nested_sequence.pyi new file mode 100644 index 0000000000000000000000000000000000000000..b6fe43700eaa44c5b4c9b9981950d7fbeaeeffc1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/nested_sequence.pyi @@ -0,0 +1,17 @@ +from collections.abc import Sequence +from numpy._typing import _NestedSequence + +a: Sequence[float] +b: list[complex] +c: tuple[str, ...] +d: int +e: str + +def func(a: _NestedSequence[int]) -> None: + ... + +reveal_type(func(a)) # E: incompatible type +reveal_type(func(b)) # E: incompatible type +reveal_type(func(c)) # E: incompatible type +reveal_type(func(d)) # E: incompatible type +reveal_type(func(e)) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/npyio.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/npyio.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e9da2777305209bdb896cbe1a91750a6eb58b084 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/npyio.pyi @@ -0,0 +1,25 @@ +import pathlib +from typing import IO + +import numpy.typing as npt +import numpy as np + +str_path: str +bytes_path: bytes +pathlib_path: pathlib.Path +str_file: IO[str] +AR_i8: npt.NDArray[np.int64] + +np.load(str_file) # E: incompatible type + +np.save(bytes_path, AR_i8) # E: No overload variant +# https://github.com/python/mypy/issues/16111 +# np.save(str_path, AR_i8, fix_imports=True) # W: deprecated + +np.savez(bytes_path, AR_i8) # E: incompatible type + +np.savez_compressed(bytes_path, AR_i8) # E: incompatible type + +np.loadtxt(bytes_path) # E: incompatible type + +np.fromregex(bytes_path, ".", np.int64) # E: No overload variant diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/numerictypes.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/numerictypes.pyi new file mode 100644 index 0000000000000000000000000000000000000000..dabd5b49d5e09b7a9fc29650e2ba99f784a2633e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/numerictypes.pyi @@ -0,0 +1,5 @@ +import numpy as np + +np.isdtype(1, np.int64) # E: incompatible type + +np.issubdtype(1, np.int64) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/random.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/random.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3370167b41d0598b5bf868c8865915b31d0b3d37 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/random.pyi @@ -0,0 +1,62 @@ +import numpy as np +import numpy.typing as npt + +SEED_FLOAT: float = 457.3 +SEED_ARR_FLOAT: npt.NDArray[np.float64] = np.array([1.0, 2, 3, 4]) +SEED_ARRLIKE_FLOAT: list[float] = [1.0, 2.0, 3.0, 4.0] +SEED_SEED_SEQ: np.random.SeedSequence = np.random.SeedSequence(0) +SEED_STR: str = "String seeding not allowed" + +# default rng +np.random.default_rng(SEED_FLOAT) # E: incompatible type +np.random.default_rng(SEED_ARR_FLOAT) # E: incompatible type +np.random.default_rng(SEED_ARRLIKE_FLOAT) # E: incompatible type +np.random.default_rng(SEED_STR) # E: incompatible type + +# Seed Sequence +np.random.SeedSequence(SEED_FLOAT) # E: incompatible type +np.random.SeedSequence(SEED_ARR_FLOAT) # E: incompatible type +np.random.SeedSequence(SEED_ARRLIKE_FLOAT) # E: incompatible type +np.random.SeedSequence(SEED_SEED_SEQ) # E: incompatible type +np.random.SeedSequence(SEED_STR) # E: incompatible type + +seed_seq: np.random.bit_generator.SeedSequence = np.random.SeedSequence() +seed_seq.spawn(11.5) # E: incompatible type +seed_seq.generate_state(3.14) # E: incompatible type +seed_seq.generate_state(3, np.uint8) # E: incompatible type +seed_seq.generate_state(3, "uint8") # E: incompatible type +seed_seq.generate_state(3, "u1") # E: incompatible type +seed_seq.generate_state(3, np.uint16) # E: incompatible type +seed_seq.generate_state(3, "uint16") # E: incompatible type +seed_seq.generate_state(3, "u2") # E: incompatible type +seed_seq.generate_state(3, np.int32) # E: incompatible type +seed_seq.generate_state(3, "int32") # E: incompatible type +seed_seq.generate_state(3, "i4") # E: incompatible type + +# Bit Generators +np.random.MT19937(SEED_FLOAT) # E: incompatible type +np.random.MT19937(SEED_ARR_FLOAT) # E: incompatible type +np.random.MT19937(SEED_ARRLIKE_FLOAT) # E: incompatible type +np.random.MT19937(SEED_STR) # E: incompatible type + +np.random.PCG64(SEED_FLOAT) # E: incompatible type +np.random.PCG64(SEED_ARR_FLOAT) # E: incompatible type +np.random.PCG64(SEED_ARRLIKE_FLOAT) # E: incompatible type +np.random.PCG64(SEED_STR) # E: incompatible type + +np.random.Philox(SEED_FLOAT) # E: incompatible type +np.random.Philox(SEED_ARR_FLOAT) # E: incompatible type +np.random.Philox(SEED_ARRLIKE_FLOAT) # E: incompatible type +np.random.Philox(SEED_STR) # E: incompatible type + +np.random.SFC64(SEED_FLOAT) # E: incompatible type +np.random.SFC64(SEED_ARR_FLOAT) # E: incompatible type +np.random.SFC64(SEED_ARRLIKE_FLOAT) # E: incompatible type +np.random.SFC64(SEED_STR) # E: incompatible type + +# Generator +np.random.Generator(None) # E: incompatible type +np.random.Generator(12333283902830213) # E: incompatible type +np.random.Generator("OxFEEDF00D") # E: incompatible type +np.random.Generator([123, 234]) # E: incompatible type +np.random.Generator(np.array([123, 234], dtype="u4")) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/rec.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/rec.pyi new file mode 100644 index 0000000000000000000000000000000000000000..0c319f2e3485fd6b77bd6a0d7bcc4c1adc6d070b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/rec.pyi @@ -0,0 +1,17 @@ +import numpy as np +import numpy.typing as npt + +AR_i8: npt.NDArray[np.int64] + +np.rec.fromarrays(1) # E: No overload variant +np.rec.fromarrays([1, 2, 3], dtype=[("f8", "f8")], formats=["f8", "f8"]) # E: No overload variant + +np.rec.fromrecords(AR_i8) # E: incompatible type +np.rec.fromrecords([(1.5,)], dtype=[("f8", "f8")], formats=["f8", "f8"]) # E: No overload variant + +np.rec.fromstring("string", dtype=[("f8", "f8")]) # E: No overload variant +np.rec.fromstring(b"bytes") # E: No overload variant +np.rec.fromstring(b"(1.5,)", dtype=[("f8", "f8")], formats=["f8", "f8"]) # E: No overload variant + +with open("test", "r") as f: + np.rec.fromfile(f, dtype=[("f8", "f8")]) # E: No overload variant diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/scalars.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/scalars.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3fa395ef2ecd6605bdb9aca7093b6b10739245b5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/scalars.pyi @@ -0,0 +1,89 @@ +import sys +import numpy as np + +f2: np.float16 +f8: np.float64 +c8: np.complex64 + +# Construction + +np.float32(3j) # E: incompatible type + +# Technically the following examples are valid NumPy code. But they +# are not considered a best practice, and people who wish to use the +# stubs should instead do +# +# np.array([1.0, 0.0, 0.0], dtype=np.float32) +# np.array([], dtype=np.complex64) +# +# See e.g. the discussion on the mailing list +# +# https://mail.python.org/pipermail/numpy-discussion/2020-April/080566.html +# +# and the issue +# +# https://github.com/numpy/numpy-stubs/issues/41 +# +# for more context. +np.float32([1.0, 0.0, 0.0]) # E: incompatible type +np.complex64([]) # E: incompatible type + +# TODO: protocols (can't check for non-existent protocols w/ __getattr__) + +np.datetime64(0) # E: No overload variant + +class A: + def __float__(self): + return 1.0 + + +np.int8(A()) # E: incompatible type +np.int16(A()) # E: incompatible type +np.int32(A()) # E: incompatible type +np.int64(A()) # E: incompatible type +np.uint8(A()) # E: incompatible type +np.uint16(A()) # E: incompatible type +np.uint32(A()) # E: incompatible type +np.uint64(A()) # E: incompatible type + +np.void("test") # E: No overload variant +np.void("test", dtype=None) # E: No overload variant + +np.generic(1) # E: Cannot instantiate abstract class +np.number(1) # E: Cannot instantiate abstract class +np.integer(1) # E: Cannot instantiate abstract class +np.inexact(1) # E: Cannot instantiate abstract class +np.character("test") # E: Cannot instantiate abstract class +np.flexible(b"test") # E: Cannot instantiate abstract class + +np.float64(value=0.0) # E: Unexpected keyword argument +np.int64(value=0) # E: Unexpected keyword argument +np.uint64(value=0) # E: Unexpected keyword argument +np.complex128(value=0.0j) # E: No overload variant +np.str_(value='bob') # E: No overload variant +np.bytes_(value=b'test') # E: No overload variant +np.void(value=b'test') # E: No overload variant +np.bool(value=True) # E: Unexpected keyword argument +np.datetime64(value="2019") # E: No overload variant +np.timedelta64(value=0) # E: Unexpected keyword argument + +np.bytes_(b"hello", encoding='utf-8') # E: No overload variant +np.str_("hello", encoding='utf-8') # E: No overload variant + +f8.item(1) # E: incompatible type +f8.item((0, 1)) # E: incompatible type +f8.squeeze(axis=1) # E: incompatible type +f8.squeeze(axis=(0, 1)) # E: incompatible type +f8.transpose(1) # E: incompatible type + +def func(a: np.float32) -> None: ... + +func(f2) # E: incompatible type +func(f8) # E: incompatible type + +c8.__getnewargs__() # E: Invalid self argument +f2.__getnewargs__() # E: Invalid self argument +f2.hex() # E: Invalid self argument +np.float16.fromhex("0x0.0p+0") # E: Invalid self argument +f2.__trunc__() # E: Invalid self argument +f2.__getformat__("float") # E: Invalid self argument diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/shape.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/shape.pyi new file mode 100644 index 0000000000000000000000000000000000000000..63faa4514077381402e40cb65dd26a1303a0b896 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/shape.pyi @@ -0,0 +1,6 @@ +from typing import Any +import numpy as np + +# test bounds of _ShapeType_co + +np.ndarray[tuple[str, str], Any] # E: Value of type variable diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/shape_base.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/shape_base.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a7c965a5b52243fa8a9c36d41404640e5560e93e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/shape_base.pyi @@ -0,0 +1,8 @@ +import numpy as np + +class DTypeLike: + dtype: np.dtype[np.int_] + +dtype_like: DTypeLike + +np.expand_dims(dtype_like, (5, 10)) # E: No overload variant diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/stride_tricks.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/stride_tricks.pyi new file mode 100644 index 0000000000000000000000000000000000000000..b213bf226a4179592417b3b890c389e601ef4e0f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/stride_tricks.pyi @@ -0,0 +1,9 @@ +import numpy as np +import numpy.typing as npt + +AR_f8: npt.NDArray[np.float64] + +np.lib.stride_tricks.as_strided(AR_f8, shape=8) # E: No overload variant +np.lib.stride_tricks.as_strided(AR_f8, strides=8) # E: No overload variant + +np.lib.stride_tricks.sliding_window_view(AR_f8, axis=(1,)) # E: No overload variant diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/strings.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/strings.pyi new file mode 100644 index 0000000000000000000000000000000000000000..ce6e6abe81a2e1750dc58ac45ae8285a0de1252a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/strings.pyi @@ -0,0 +1,59 @@ +import numpy as np +import numpy.typing as npt + +AR_U: npt.NDArray[np.str_] +AR_S: npt.NDArray[np.bytes_] + +np.strings.equal(AR_U, AR_S) # E: incompatible type + +np.strings.not_equal(AR_U, AR_S) # E: incompatible type + +np.strings.greater_equal(AR_U, AR_S) # E: incompatible type + +np.strings.less_equal(AR_U, AR_S) # E: incompatible type + +np.strings.greater(AR_U, AR_S) # E: incompatible type + +np.strings.less(AR_U, AR_S) # E: incompatible type + +np.strings.encode(AR_S) # E: incompatible type +np.strings.decode(AR_U) # E: incompatible type + +np.strings.join(AR_U, b"_") # E: incompatible type +np.strings.join(AR_S, "_") # E: incompatible type + +np.strings.lstrip(AR_U, b"a") # E: incompatible type +np.strings.lstrip(AR_S, "a") # E: incompatible type +np.strings.strip(AR_U, b"a") # E: incompatible type +np.strings.strip(AR_S, "a") # E: incompatible type +np.strings.rstrip(AR_U, b"a") # E: incompatible type +np.strings.rstrip(AR_S, "a") # E: incompatible type + +np.strings.partition(AR_U, b"a") # E: incompatible type +np.strings.partition(AR_S, "a") # E: incompatible type +np.strings.rpartition(AR_U, b"a") # E: incompatible type +np.strings.rpartition(AR_S, "a") # E: incompatible type + +np.strings.count(AR_U, b"a", [1, 2, 3], [1, 2, 3]) # E: incompatible type +np.strings.count(AR_S, "a", 0, 9) # E: incompatible type + +np.strings.endswith(AR_U, b"a", [1, 2, 3], [1, 2, 3]) # E: incompatible type +np.strings.endswith(AR_S, "a", 0, 9) # E: incompatible type +np.strings.startswith(AR_U, b"a", [1, 2, 3], [1, 2, 3]) # E: incompatible type +np.strings.startswith(AR_S, "a", 0, 9) # E: incompatible type + +np.strings.find(AR_U, b"a", [1, 2, 3], [1, 2, 3]) # E: incompatible type +np.strings.find(AR_S, "a", 0, 9) # E: incompatible type +np.strings.rfind(AR_U, b"a", [1, 2, 3], [1, 2 , 3]) # E: incompatible type +np.strings.rfind(AR_S, "a", 0, 9) # E: incompatible type + +np.strings.index(AR_U, b"a", start=[1, 2, 3]) # E: incompatible type +np.strings.index(AR_S, "a", end=9) # E: incompatible type +np.strings.rindex(AR_U, b"a", start=[1, 2, 3]) # E: incompatible type +np.strings.rindex(AR_S, "a", end=9) # E: incompatible type + +np.strings.isdecimal(AR_S) # E: incompatible type +np.strings.isnumeric(AR_S) # E: incompatible type + +np.strings.replace(AR_U, b"_", b"-", 10) # E: incompatible type +np.strings.replace(AR_S, "_", "-", 1) # E: incompatible type diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/testing.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/testing.pyi new file mode 100644 index 0000000000000000000000000000000000000000..160c556880fa7ac1537201e3c1f05e45d8800f41 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/fail/testing.pyi @@ -0,0 +1,28 @@ +import numpy as np +import numpy.typing as npt + +AR_U: npt.NDArray[np.str_] + +def func(x: object) -> bool: ... + +np.testing.assert_(True, msg=1) # E: incompatible type +np.testing.build_err_msg(1, "test") # E: incompatible type +np.testing.assert_almost_equal(AR_U, AR_U) # E: incompatible type +np.testing.assert_approx_equal([1, 2, 3], [1, 2, 3]) # E: incompatible type +np.testing.assert_array_almost_equal(AR_U, AR_U) # E: incompatible type +np.testing.assert_array_less(AR_U, AR_U) # E: incompatible type +np.testing.assert_string_equal(b"a", b"a") # E: incompatible type + +np.testing.assert_raises(expected_exception=TypeError, callable=func) # E: No overload variant +np.testing.assert_raises_regex(expected_exception=TypeError, expected_regex="T", callable=func) # E: No overload variant + +np.testing.assert_allclose(AR_U, AR_U) # E: incompatible type +np.testing.assert_array_almost_equal_nulp(AR_U, AR_U) # E: incompatible type +np.testing.assert_array_max_ulp(AR_U, AR_U) # E: incompatible type + +np.testing.assert_warns(RuntimeWarning, func) # E: No overload variant +np.testing.assert_no_warnings(func=func) # E: No overload variant +np.testing.assert_no_warnings(func) # E: Too many arguments +np.testing.assert_no_warnings(func, y=None) # E: No overload variant + +np.testing.assert_no_gc_cycles(func=func) # E: No overload variant diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/polynomial_series.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/polynomial_series.pyi new file mode 100644 index 0000000000000000000000000000000000000000..51ecd4a24d6bf9decd336f096f2d67c5ad5cd78e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/polynomial_series.pyi @@ -0,0 +1,140 @@ +from collections.abc import Sequence +from typing import Any, TypeAlias + +import numpy as np +import numpy.polynomial as npp +import numpy.typing as npt + +from typing_extensions import assert_type + +_ArrFloat1D: TypeAlias = np.ndarray[tuple[int], np.dtype[np.floating[Any]]] +_ArrFloat1D64: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]] +_ArrComplex1D: TypeAlias = np.ndarray[tuple[int], np.dtype[np.complexfloating[Any, Any]]] +_ArrComplex1D128: TypeAlias = np.ndarray[tuple[int], np.dtype[np.complex128]] +_ArrObject1D: TypeAlias = np.ndarray[tuple[int], np.dtype[np.object_]] + +AR_b: npt.NDArray[np.bool] +AR_u4: npt.NDArray[np.uint32] +AR_i8: npt.NDArray[np.int64] +AR_f8: npt.NDArray[np.float64] +AR_c16: npt.NDArray[np.complex128] +AR_O: npt.NDArray[np.object_] + +PS_poly: npp.Polynomial +PS_cheb: npp.Chebyshev + +assert_type(npp.polynomial.polyroots(AR_f8), _ArrFloat1D64) +assert_type(npp.polynomial.polyroots(AR_c16), _ArrComplex1D128) +assert_type(npp.polynomial.polyroots(AR_O), _ArrObject1D) + +assert_type(npp.polynomial.polyfromroots(AR_f8), _ArrFloat1D) +assert_type(npp.polynomial.polyfromroots(AR_c16), _ArrComplex1D) +assert_type(npp.polynomial.polyfromroots(AR_O), _ArrObject1D) + +# assert_type(npp.polynomial.polyadd(AR_b, AR_b), NoReturn) +assert_type(npp.polynomial.polyadd(AR_u4, AR_b), _ArrFloat1D) +assert_type(npp.polynomial.polyadd(AR_i8, AR_i8), _ArrFloat1D) +assert_type(npp.polynomial.polyadd(AR_f8, AR_i8), _ArrFloat1D) +assert_type(npp.polynomial.polyadd(AR_i8, AR_c16), _ArrComplex1D) +assert_type(npp.polynomial.polyadd(AR_O, AR_O), _ArrObject1D) + +assert_type(npp.polynomial.polymulx(AR_u4), _ArrFloat1D) +assert_type(npp.polynomial.polymulx(AR_i8), _ArrFloat1D) +assert_type(npp.polynomial.polymulx(AR_f8), _ArrFloat1D) +assert_type(npp.polynomial.polymulx(AR_c16), _ArrComplex1D) +assert_type(npp.polynomial.polymulx(AR_O), _ArrObject1D) + +assert_type(npp.polynomial.polypow(AR_u4, 2), _ArrFloat1D) +assert_type(npp.polynomial.polypow(AR_i8, 2), _ArrFloat1D) +assert_type(npp.polynomial.polypow(AR_f8, 2), _ArrFloat1D) +assert_type(npp.polynomial.polypow(AR_c16, 2), _ArrComplex1D) +assert_type(npp.polynomial.polypow(AR_O, 2), _ArrObject1D) + +# assert_type(npp.polynomial.polyder(PS_poly), npt.NDArray[np.object_]) +assert_type(npp.polynomial.polyder(AR_f8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyder(AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyder(AR_O, m=2), npt.NDArray[np.object_]) + +# assert_type(npp.polynomial.polyint(PS_poly), npt.NDArray[np.object_]) +assert_type(npp.polynomial.polyint(AR_f8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyint(AR_f8, k=AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyint(AR_O, m=2), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyval(AR_b, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval(AR_u4, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval(AR_i8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval(AR_f8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval(AR_i8, AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyval(AR_O, AR_O), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyval2d(AR_b, AR_b, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval2d(AR_u4, AR_u4, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval2d(AR_i8, AR_i8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval2d(AR_f8, AR_f8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval2d(AR_i8, AR_i8, AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyval2d(AR_O, AR_O, AR_O), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyval3d(AR_b, AR_b, AR_b, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval3d(AR_u4, AR_u4, AR_u4, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval3d(AR_i8, AR_i8, AR_i8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval3d(AR_f8, AR_f8, AR_f8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyval3d(AR_i8, AR_i8, AR_i8, AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyval3d(AR_O, AR_O, AR_O, AR_O), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyvalfromroots(AR_b, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvalfromroots(AR_u4, AR_b), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvalfromroots(AR_i8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvalfromroots(AR_f8, AR_i8), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvalfromroots(AR_i8, AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyvalfromroots(AR_O, AR_O), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyvander(AR_f8, 3), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvander(AR_c16, 3), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyvander(AR_O, 3), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyvander2d(AR_f8, AR_f8, [4, 2]), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvander2d(AR_c16, AR_c16, [4, 2]), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyvander2d(AR_O, AR_O, [4, 2]), npt.NDArray[np.object_]) + +assert_type(npp.polynomial.polyvander3d(AR_f8, AR_f8, AR_f8, [4, 3, 2]), npt.NDArray[np.floating[Any]]) +assert_type(npp.polynomial.polyvander3d(AR_c16, AR_c16, AR_c16, [4, 3, 2]), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(npp.polynomial.polyvander3d(AR_O, AR_O, AR_O, [4, 3, 2]), npt.NDArray[np.object_]) + +assert_type( + npp.polynomial.polyfit(AR_f8, AR_f8, 2), + npt.NDArray[np.floating[Any]], +) +assert_type( + npp.polynomial.polyfit(AR_f8, AR_i8, 1, full=True), + tuple[npt.NDArray[np.floating[Any]], Sequence[np.inexact[Any] | np.int32]], +) +assert_type( + npp.polynomial.polyfit(AR_c16, AR_f8, 2), + npt.NDArray[np.complexfloating[Any, Any]], +) +assert_type( + npp.polynomial.polyfit(AR_f8, AR_c16, 1, full=True)[0], + npt.NDArray[np.complexfloating[Any, Any]], +) + +assert_type(npp.chebyshev.chebgauss(2), tuple[_ArrFloat1D64, _ArrFloat1D64]) + +assert_type(npp.chebyshev.chebweight(AR_f8), npt.NDArray[np.float64]) +assert_type(npp.chebyshev.chebweight(AR_c16), npt.NDArray[np.complex128]) +assert_type(npp.chebyshev.chebweight(AR_O), npt.NDArray[np.object_]) + +assert_type(npp.chebyshev.poly2cheb(AR_f8), _ArrFloat1D) +assert_type(npp.chebyshev.poly2cheb(AR_c16), _ArrComplex1D) +assert_type(npp.chebyshev.poly2cheb(AR_O), _ArrObject1D) + +assert_type(npp.chebyshev.cheb2poly(AR_f8), _ArrFloat1D) +assert_type(npp.chebyshev.cheb2poly(AR_c16), _ArrComplex1D) +assert_type(npp.chebyshev.cheb2poly(AR_O), _ArrObject1D) + +assert_type(npp.chebyshev.chebpts1(6), _ArrFloat1D64) +assert_type(npp.chebyshev.chebpts2(6), _ArrFloat1D64) + +assert_type( + npp.chebyshev.chebinterpolate(np.tanh, 3), + npt.NDArray[np.float64 | np.complex128 | np.object_], +) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/random.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/random.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e01174372be6046221c29e357d39f3bce7e2077d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/random.pyi @@ -0,0 +1,1551 @@ +import threading +from typing import Any +from collections.abc import Sequence + +import numpy as np +import numpy.typing as npt +from numpy.random._generator import Generator +from numpy.random._mt19937 import MT19937 +from numpy.random._pcg64 import PCG64 +from numpy.random._sfc64 import SFC64 +from numpy.random._philox import Philox +from numpy.random.bit_generator import SeedSequence, SeedlessSeedSequence + +from typing_extensions import assert_type + +def_rng = np.random.default_rng() +seed_seq = np.random.SeedSequence() +mt19937 = np.random.MT19937() +pcg64 = np.random.PCG64() +sfc64 = np.random.SFC64() +philox = np.random.Philox() +seedless_seq = SeedlessSeedSequence() + +assert_type(def_rng, Generator) +assert_type(mt19937, MT19937) +assert_type(pcg64, PCG64) +assert_type(sfc64, SFC64) +assert_type(philox, Philox) +assert_type(seed_seq, SeedSequence) +assert_type(seedless_seq, SeedlessSeedSequence) + +mt19937_jumped = mt19937.jumped() +mt19937_jumped3 = mt19937.jumped(3) +mt19937_raw = mt19937.random_raw() +mt19937_raw_arr = mt19937.random_raw(5) + +assert_type(mt19937_jumped, MT19937) +assert_type(mt19937_jumped3, MT19937) +assert_type(mt19937_raw, int) +assert_type(mt19937_raw_arr, npt.NDArray[np.uint64]) +assert_type(mt19937.lock, threading.Lock) + +pcg64_jumped = pcg64.jumped() +pcg64_jumped3 = pcg64.jumped(3) +pcg64_adv = pcg64.advance(3) +pcg64_raw = pcg64.random_raw() +pcg64_raw_arr = pcg64.random_raw(5) + +assert_type(pcg64_jumped, PCG64) +assert_type(pcg64_jumped3, PCG64) +assert_type(pcg64_adv, PCG64) +assert_type(pcg64_raw, int) +assert_type(pcg64_raw_arr, npt.NDArray[np.uint64]) +assert_type(pcg64.lock, threading.Lock) + +philox_jumped = philox.jumped() +philox_jumped3 = philox.jumped(3) +philox_adv = philox.advance(3) +philox_raw = philox.random_raw() +philox_raw_arr = philox.random_raw(5) + +assert_type(philox_jumped, Philox) +assert_type(philox_jumped3, Philox) +assert_type(philox_adv, Philox) +assert_type(philox_raw, int) +assert_type(philox_raw_arr, npt.NDArray[np.uint64]) +assert_type(philox.lock, threading.Lock) + +sfc64_raw = sfc64.random_raw() +sfc64_raw_arr = sfc64.random_raw(5) + +assert_type(sfc64_raw, int) +assert_type(sfc64_raw_arr, npt.NDArray[np.uint64]) +assert_type(sfc64.lock, threading.Lock) + +assert_type(seed_seq.pool, npt.NDArray[np.uint32]) +assert_type(seed_seq.entropy, None | int | Sequence[int]) +assert_type(seed_seq.spawn(1), list[np.random.SeedSequence]) +assert_type(seed_seq.generate_state(8, "uint32"), npt.NDArray[np.uint32 | np.uint64]) +assert_type(seed_seq.generate_state(8, "uint64"), npt.NDArray[np.uint32 | np.uint64]) + + +def_gen: np.random.Generator = np.random.default_rng() + +D_arr_0p1: npt.NDArray[np.float64] = np.array([0.1]) +D_arr_0p5: npt.NDArray[np.float64] = np.array([0.5]) +D_arr_0p9: npt.NDArray[np.float64] = np.array([0.9]) +D_arr_1p5: npt.NDArray[np.float64] = np.array([1.5]) +I_arr_10: npt.NDArray[np.int_] = np.array([10], dtype=np.int_) +I_arr_20: npt.NDArray[np.int_] = np.array([20], dtype=np.int_) +D_arr_like_0p1: list[float] = [0.1] +D_arr_like_0p5: list[float] = [0.5] +D_arr_like_0p9: list[float] = [0.9] +D_arr_like_1p5: list[float] = [1.5] +I_arr_like_10: list[int] = [10] +I_arr_like_20: list[int] = [20] +D_2D_like: list[list[float]] = [[1, 2], [2, 3], [3, 4], [4, 5.1]] +D_2D: npt.NDArray[np.float64] = np.array(D_2D_like) +S_out: npt.NDArray[np.float32] = np.empty(1, dtype=np.float32) +D_out: npt.NDArray[np.float64] = np.empty(1) + +assert_type(def_gen.standard_normal(), float) +assert_type(def_gen.standard_normal(dtype=np.float32), float) +assert_type(def_gen.standard_normal(dtype="float32"), float) +assert_type(def_gen.standard_normal(dtype="double"), float) +assert_type(def_gen.standard_normal(dtype=np.float64), float) +assert_type(def_gen.standard_normal(size=None), float) +assert_type(def_gen.standard_normal(size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_normal(size=1, dtype=np.float32), npt.NDArray[np.float32]) +assert_type(def_gen.standard_normal(size=1, dtype="f4"), npt.NDArray[np.float32]) +assert_type(def_gen.standard_normal(size=1, dtype="float32", out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.standard_normal(dtype=np.float32, out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.standard_normal(size=1, dtype=np.float64), npt.NDArray[np.float64]) +assert_type(def_gen.standard_normal(size=1, dtype="float64"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_normal(size=1, dtype="f8"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_normal(out=D_out), npt.NDArray[np.float64]) +assert_type(def_gen.standard_normal(size=1, dtype="float64"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_normal(size=1, dtype="float64", out=D_out), npt.NDArray[np.float64]) + +assert_type(def_gen.random(), float) +assert_type(def_gen.random(dtype=np.float32), float) +assert_type(def_gen.random(dtype="float32"), float) +assert_type(def_gen.random(dtype="double"), float) +assert_type(def_gen.random(dtype=np.float64), float) +assert_type(def_gen.random(size=None), float) +assert_type(def_gen.random(size=1), npt.NDArray[np.float64]) +assert_type(def_gen.random(size=1, dtype=np.float32), npt.NDArray[np.float32]) +assert_type(def_gen.random(size=1, dtype="f4"), npt.NDArray[np.float32]) +assert_type(def_gen.random(size=1, dtype="float32", out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.random(dtype=np.float32, out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.random(size=1, dtype=np.float64), npt.NDArray[np.float64]) +assert_type(def_gen.random(size=1, dtype="float64"), npt.NDArray[np.float64]) +assert_type(def_gen.random(size=1, dtype="f8"), npt.NDArray[np.float64]) +assert_type(def_gen.random(out=D_out), npt.NDArray[np.float64]) +assert_type(def_gen.random(size=1, dtype="float64"), npt.NDArray[np.float64]) +assert_type(def_gen.random(size=1, dtype="float64", out=D_out), npt.NDArray[np.float64]) + +assert_type(def_gen.standard_cauchy(), float) +assert_type(def_gen.standard_cauchy(size=None), float) +assert_type(def_gen.standard_cauchy(size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.standard_exponential(), float) +assert_type(def_gen.standard_exponential(method="inv"), float) +assert_type(def_gen.standard_exponential(dtype=np.float32), float) +assert_type(def_gen.standard_exponential(dtype="float32"), float) +assert_type(def_gen.standard_exponential(dtype="double"), float) +assert_type(def_gen.standard_exponential(dtype=np.float64), float) +assert_type(def_gen.standard_exponential(size=None), float) +assert_type(def_gen.standard_exponential(size=None, method="inv"), float) +assert_type(def_gen.standard_exponential(size=1, method="inv"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_exponential(size=1, dtype=np.float32), npt.NDArray[np.float32]) +assert_type(def_gen.standard_exponential(size=1, dtype="f4", method="inv"), npt.NDArray[np.float32]) +assert_type(def_gen.standard_exponential(size=1, dtype="float32", out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.standard_exponential(dtype=np.float32, out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.standard_exponential(size=1, dtype=np.float64, method="inv"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_exponential(size=1, dtype="float64"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_exponential(size=1, dtype="f8"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_exponential(out=D_out), npt.NDArray[np.float64]) +assert_type(def_gen.standard_exponential(size=1, dtype="float64"), npt.NDArray[np.float64]) +assert_type(def_gen.standard_exponential(size=1, dtype="float64", out=D_out), npt.NDArray[np.float64]) + +assert_type(def_gen.zipf(1.5), int) +assert_type(def_gen.zipf(1.5, size=None), int) +assert_type(def_gen.zipf(1.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.zipf(D_arr_1p5), npt.NDArray[np.int64]) +assert_type(def_gen.zipf(D_arr_1p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.zipf(D_arr_like_1p5), npt.NDArray[np.int64]) +assert_type(def_gen.zipf(D_arr_like_1p5, size=1), npt.NDArray[np.int64]) + +assert_type(def_gen.weibull(0.5), float) +assert_type(def_gen.weibull(0.5, size=None), float) +assert_type(def_gen.weibull(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.weibull(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.weibull(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.weibull(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.weibull(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.standard_t(0.5), float) +assert_type(def_gen.standard_t(0.5, size=None), float) +assert_type(def_gen.standard_t(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_t(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.standard_t(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_t(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.standard_t(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.poisson(0.5), int) +assert_type(def_gen.poisson(0.5, size=None), int) +assert_type(def_gen.poisson(0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.poisson(D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.poisson(D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.poisson(D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.poisson(D_arr_like_0p5, size=1), npt.NDArray[np.int64]) + +assert_type(def_gen.power(0.5), float) +assert_type(def_gen.power(0.5, size=None), float) +assert_type(def_gen.power(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.power(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.power(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.power(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.power(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.pareto(0.5), float) +assert_type(def_gen.pareto(0.5, size=None), float) +assert_type(def_gen.pareto(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.pareto(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.pareto(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.pareto(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.pareto(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.chisquare(0.5), float) +assert_type(def_gen.chisquare(0.5, size=None), float) +assert_type(def_gen.chisquare(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.chisquare(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.chisquare(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.chisquare(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.chisquare(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.exponential(0.5), float) +assert_type(def_gen.exponential(0.5, size=None), float) +assert_type(def_gen.exponential(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.exponential(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.exponential(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.exponential(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.exponential(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.geometric(0.5), int) +assert_type(def_gen.geometric(0.5, size=None), int) +assert_type(def_gen.geometric(0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.geometric(D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.geometric(D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.geometric(D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.geometric(D_arr_like_0p5, size=1), npt.NDArray[np.int64]) + +assert_type(def_gen.logseries(0.5), int) +assert_type(def_gen.logseries(0.5, size=None), int) +assert_type(def_gen.logseries(0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.logseries(D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.logseries(D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.logseries(D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.logseries(D_arr_like_0p5, size=1), npt.NDArray[np.int64]) + +assert_type(def_gen.rayleigh(0.5), float) +assert_type(def_gen.rayleigh(0.5, size=None), float) +assert_type(def_gen.rayleigh(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.rayleigh(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.rayleigh(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.rayleigh(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.rayleigh(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.standard_gamma(0.5), float) +assert_type(def_gen.standard_gamma(0.5, size=None), float) +assert_type(def_gen.standard_gamma(0.5, dtype="float32"), float) +assert_type(def_gen.standard_gamma(0.5, size=None, dtype="float32"), float) +assert_type(def_gen.standard_gamma(0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_0p5, dtype="f4"), npt.NDArray[np.float32]) +assert_type(def_gen.standard_gamma(0.5, size=1, dtype="float32", out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.standard_gamma(D_arr_0p5, dtype=np.float32, out=S_out), npt.NDArray[np.float32]) +assert_type(def_gen.standard_gamma(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(0.5, out=D_out), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_like_0p5, out=D_out), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.standard_gamma(D_arr_like_0p5, size=1, out=D_out, dtype=np.float64), npt.NDArray[np.float64]) + +assert_type(def_gen.vonmises(0.5, 0.5), float) +assert_type(def_gen.vonmises(0.5, 0.5, size=None), float) +assert_type(def_gen.vonmises(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.vonmises(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.wald(0.5, 0.5), float) +assert_type(def_gen.wald(0.5, 0.5, size=None), float) +assert_type(def_gen.wald(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.wald(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.wald(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.wald(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.wald(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.uniform(0.5, 0.5), float) +assert_type(def_gen.uniform(0.5, 0.5, size=None), float) +assert_type(def_gen.uniform(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.uniform(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.beta(0.5, 0.5), float) +assert_type(def_gen.beta(0.5, 0.5, size=None), float) +assert_type(def_gen.beta(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.beta(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.beta(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.beta(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.beta(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.f(0.5, 0.5), float) +assert_type(def_gen.f(0.5, 0.5, size=None), float) +assert_type(def_gen.f(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.f(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.f(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.f(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.f(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.gamma(0.5, 0.5), float) +assert_type(def_gen.gamma(0.5, 0.5, size=None), float) +assert_type(def_gen.gamma(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gamma(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.gumbel(0.5, 0.5), float) +assert_type(def_gen.gumbel(0.5, 0.5, size=None), float) +assert_type(def_gen.gumbel(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.gumbel(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.laplace(0.5, 0.5), float) +assert_type(def_gen.laplace(0.5, 0.5, size=None), float) +assert_type(def_gen.laplace(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.laplace(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.logistic(0.5, 0.5), float) +assert_type(def_gen.logistic(0.5, 0.5, size=None), float) +assert_type(def_gen.logistic(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.logistic(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.lognormal(0.5, 0.5), float) +assert_type(def_gen.lognormal(0.5, 0.5, size=None), float) +assert_type(def_gen.lognormal(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.lognormal(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.noncentral_chisquare(0.5, 0.5), float) +assert_type(def_gen.noncentral_chisquare(0.5, 0.5, size=None), float) +assert_type(def_gen.noncentral_chisquare(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_chisquare(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.normal(0.5, 0.5), float) +assert_type(def_gen.normal(0.5, 0.5, size=None), float) +assert_type(def_gen.normal(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.normal(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.normal(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(def_gen.normal(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.normal(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.triangular(0.1, 0.5, 0.9), float) +assert_type(def_gen.triangular(0.1, 0.5, 0.9, size=None), float) +assert_type(def_gen.triangular(0.1, 0.5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_0p1, 0.5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(0.1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_0p1, 0.5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(0.1, D_arr_0p5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_like_0p1, 0.5, D_arr_0p9), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(0.5, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_0p1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_like_0p1, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_0p1, D_arr_0p5, D_arr_0p9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.triangular(D_arr_like_0p1, D_arr_like_0p5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.noncentral_f(0.1, 0.5, 0.9), float) +assert_type(def_gen.noncentral_f(0.1, 0.5, 0.9, size=None), float) +assert_type(def_gen.noncentral_f(0.1, 0.5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_0p1, 0.5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(0.1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_0p1, 0.5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(0.1, D_arr_0p5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_like_0p1, 0.5, D_arr_0p9), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(0.5, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_0p1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_like_0p1, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_0p1, D_arr_0p5, D_arr_0p9, size=1), npt.NDArray[np.float64]) +assert_type(def_gen.noncentral_f(D_arr_like_0p1, D_arr_like_0p5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) + +assert_type(def_gen.binomial(10, 0.5), int) +assert_type(def_gen.binomial(10, 0.5, size=None), int) +assert_type(def_gen.binomial(10, 0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_10, 0.5), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(10, D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_10, 0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(10, D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_like_10, 0.5), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(10, D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_10, D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_like_10, D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_10, D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.binomial(I_arr_like_10, D_arr_like_0p5, size=1), npt.NDArray[np.int64]) + +assert_type(def_gen.negative_binomial(10, 0.5), int) +assert_type(def_gen.negative_binomial(10, 0.5, size=None), int) +assert_type(def_gen.negative_binomial(10, 0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_10, 0.5), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(10, D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_10, 0.5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(10, D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_like_10, 0.5), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(10, D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_10, D_arr_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_like_10, D_arr_like_0p5), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_10, D_arr_0p5, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.negative_binomial(I_arr_like_10, D_arr_like_0p5, size=1), npt.NDArray[np.int64]) + +assert_type(def_gen.hypergeometric(20, 20, 10), int) +assert_type(def_gen.hypergeometric(20, 20, 10, size=None), int) +assert_type(def_gen.hypergeometric(20, 20, 10, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_20, 20, 10), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(20, I_arr_20, 10), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_20, 20, I_arr_like_10, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(20, I_arr_20, 10, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_like_20, 20, I_arr_10), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(20, I_arr_like_20, 10), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_20, I_arr_20, 10), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_like_20, I_arr_like_20, 10), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_20, I_arr_20, I_arr_10, size=1), npt.NDArray[np.int64]) +assert_type(def_gen.hypergeometric(I_arr_like_20, I_arr_like_20, I_arr_like_10, size=1), npt.NDArray[np.int64]) + +I_int64_100: npt.NDArray[np.int64] = np.array([100], dtype=np.int64) + +assert_type(def_gen.integers(0, 100), np.int64) +assert_type(def_gen.integers(100), np.int64) +assert_type(def_gen.integers([100]), npt.NDArray[np.int64]) +assert_type(def_gen.integers(0, [100]), npt.NDArray[np.int64]) + +I_bool_low: npt.NDArray[np.bool] = np.array([0], dtype=np.bool) +I_bool_low_like: list[int] = [0] +I_bool_high_open: npt.NDArray[np.bool] = np.array([1], dtype=np.bool) +I_bool_high_closed: npt.NDArray[np.bool] = np.array([1], dtype=np.bool) + +assert_type(def_gen.integers(2, dtype=bool), bool) +assert_type(def_gen.integers(0, 2, dtype=bool), bool) +assert_type(def_gen.integers(1, dtype=bool, endpoint=True), bool) +assert_type(def_gen.integers(0, 1, dtype=bool, endpoint=True), bool) +assert_type(def_gen.integers(I_bool_low_like, 1, dtype=bool, endpoint=True), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_high_open, dtype=bool), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_low, I_bool_high_open, dtype=bool), npt.NDArray[np.bool]) +assert_type(def_gen.integers(0, I_bool_high_open, dtype=bool), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_high_closed, dtype=bool, endpoint=True), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_low, I_bool_high_closed, dtype=bool, endpoint=True), npt.NDArray[np.bool]) +assert_type(def_gen.integers(0, I_bool_high_closed, dtype=bool, endpoint=True), npt.NDArray[np.bool]) + +assert_type(def_gen.integers(2, dtype=np.bool), np.bool) +assert_type(def_gen.integers(0, 2, dtype=np.bool), np.bool) +assert_type(def_gen.integers(1, dtype=np.bool, endpoint=True), np.bool) +assert_type(def_gen.integers(0, 1, dtype=np.bool, endpoint=True), np.bool) +assert_type(def_gen.integers(I_bool_low_like, 1, dtype=np.bool, endpoint=True), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_high_open, dtype=np.bool), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_low, I_bool_high_open, dtype=np.bool), npt.NDArray[np.bool]) +assert_type(def_gen.integers(0, I_bool_high_open, dtype=np.bool), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_high_closed, dtype=np.bool, endpoint=True), npt.NDArray[np.bool]) +assert_type(def_gen.integers(I_bool_low, I_bool_high_closed, dtype=np.bool, endpoint=True), npt.NDArray[np.bool]) +assert_type(def_gen.integers(0, I_bool_high_closed, dtype=np.bool, endpoint=True), npt.NDArray[np.bool]) + +I_u1_low: npt.NDArray[np.uint8] = np.array([0], dtype=np.uint8) +I_u1_low_like: list[int] = [0] +I_u1_high_open: npt.NDArray[np.uint8] = np.array([255], dtype=np.uint8) +I_u1_high_closed: npt.NDArray[np.uint8] = np.array([255], dtype=np.uint8) + +assert_type(def_gen.integers(256, dtype="u1"), np.uint8) +assert_type(def_gen.integers(0, 256, dtype="u1"), np.uint8) +assert_type(def_gen.integers(255, dtype="u1", endpoint=True), np.uint8) +assert_type(def_gen.integers(0, 255, dtype="u1", endpoint=True), np.uint8) +assert_type(def_gen.integers(I_u1_low_like, 255, dtype="u1", endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_high_open, dtype="u1"), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_low, I_u1_high_open, dtype="u1"), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(0, I_u1_high_open, dtype="u1"), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_high_closed, dtype="u1", endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_low, I_u1_high_closed, dtype="u1", endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(0, I_u1_high_closed, dtype="u1", endpoint=True), npt.NDArray[np.uint8]) + +assert_type(def_gen.integers(256, dtype="uint8"), np.uint8) +assert_type(def_gen.integers(0, 256, dtype="uint8"), np.uint8) +assert_type(def_gen.integers(255, dtype="uint8", endpoint=True), np.uint8) +assert_type(def_gen.integers(0, 255, dtype="uint8", endpoint=True), np.uint8) +assert_type(def_gen.integers(I_u1_low_like, 255, dtype="uint8", endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_high_open, dtype="uint8"), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_low, I_u1_high_open, dtype="uint8"), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(0, I_u1_high_open, dtype="uint8"), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_high_closed, dtype="uint8", endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_low, I_u1_high_closed, dtype="uint8", endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(0, I_u1_high_closed, dtype="uint8", endpoint=True), npt.NDArray[np.uint8]) + +assert_type(def_gen.integers(256, dtype=np.uint8), np.uint8) +assert_type(def_gen.integers(0, 256, dtype=np.uint8), np.uint8) +assert_type(def_gen.integers(255, dtype=np.uint8, endpoint=True), np.uint8) +assert_type(def_gen.integers(0, 255, dtype=np.uint8, endpoint=True), np.uint8) +assert_type(def_gen.integers(I_u1_low_like, 255, dtype=np.uint8, endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_high_open, dtype=np.uint8), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_low, I_u1_high_open, dtype=np.uint8), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(0, I_u1_high_open, dtype=np.uint8), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_high_closed, dtype=np.uint8, endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(I_u1_low, I_u1_high_closed, dtype=np.uint8, endpoint=True), npt.NDArray[np.uint8]) +assert_type(def_gen.integers(0, I_u1_high_closed, dtype=np.uint8, endpoint=True), npt.NDArray[np.uint8]) + +I_u2_low: npt.NDArray[np.uint16] = np.array([0], dtype=np.uint16) +I_u2_low_like: list[int] = [0] +I_u2_high_open: npt.NDArray[np.uint16] = np.array([65535], dtype=np.uint16) +I_u2_high_closed: npt.NDArray[np.uint16] = np.array([65535], dtype=np.uint16) + +assert_type(def_gen.integers(65536, dtype="u2"), np.uint16) +assert_type(def_gen.integers(0, 65536, dtype="u2"), np.uint16) +assert_type(def_gen.integers(65535, dtype="u2", endpoint=True), np.uint16) +assert_type(def_gen.integers(0, 65535, dtype="u2", endpoint=True), np.uint16) +assert_type(def_gen.integers(I_u2_low_like, 65535, dtype="u2", endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_high_open, dtype="u2"), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_low, I_u2_high_open, dtype="u2"), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(0, I_u2_high_open, dtype="u2"), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_high_closed, dtype="u2", endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_low, I_u2_high_closed, dtype="u2", endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(0, I_u2_high_closed, dtype="u2", endpoint=True), npt.NDArray[np.uint16]) + +assert_type(def_gen.integers(65536, dtype="uint16"), np.uint16) +assert_type(def_gen.integers(0, 65536, dtype="uint16"), np.uint16) +assert_type(def_gen.integers(65535, dtype="uint16", endpoint=True), np.uint16) +assert_type(def_gen.integers(0, 65535, dtype="uint16", endpoint=True), np.uint16) +assert_type(def_gen.integers(I_u2_low_like, 65535, dtype="uint16", endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_high_open, dtype="uint16"), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_low, I_u2_high_open, dtype="uint16"), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(0, I_u2_high_open, dtype="uint16"), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_high_closed, dtype="uint16", endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_low, I_u2_high_closed, dtype="uint16", endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(0, I_u2_high_closed, dtype="uint16", endpoint=True), npt.NDArray[np.uint16]) + +assert_type(def_gen.integers(65536, dtype=np.uint16), np.uint16) +assert_type(def_gen.integers(0, 65536, dtype=np.uint16), np.uint16) +assert_type(def_gen.integers(65535, dtype=np.uint16, endpoint=True), np.uint16) +assert_type(def_gen.integers(0, 65535, dtype=np.uint16, endpoint=True), np.uint16) +assert_type(def_gen.integers(I_u2_low_like, 65535, dtype=np.uint16, endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_high_open, dtype=np.uint16), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_low, I_u2_high_open, dtype=np.uint16), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(0, I_u2_high_open, dtype=np.uint16), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_high_closed, dtype=np.uint16, endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(I_u2_low, I_u2_high_closed, dtype=np.uint16, endpoint=True), npt.NDArray[np.uint16]) +assert_type(def_gen.integers(0, I_u2_high_closed, dtype=np.uint16, endpoint=True), npt.NDArray[np.uint16]) + +I_u4_low: npt.NDArray[np.uint32] = np.array([0], dtype=np.uint32) +I_u4_low_like: list[int] = [0] +I_u4_high_open: npt.NDArray[np.uint32] = np.array([4294967295], dtype=np.uint32) +I_u4_high_closed: npt.NDArray[np.uint32] = np.array([4294967295], dtype=np.uint32) + +assert_type(def_gen.integers(4294967296, dtype=np.int_), np.int_) +assert_type(def_gen.integers(0, 4294967296, dtype=np.int_), np.int_) +assert_type(def_gen.integers(4294967295, dtype=np.int_, endpoint=True), np.int_) +assert_type(def_gen.integers(0, 4294967295, dtype=np.int_, endpoint=True), np.int_) +assert_type(def_gen.integers(I_u4_low_like, 4294967295, dtype=np.int_, endpoint=True), npt.NDArray[np.int_]) +assert_type(def_gen.integers(I_u4_high_open, dtype=np.int_), npt.NDArray[np.int_]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_open, dtype=np.int_), npt.NDArray[np.int_]) +assert_type(def_gen.integers(0, I_u4_high_open, dtype=np.int_), npt.NDArray[np.int_]) +assert_type(def_gen.integers(I_u4_high_closed, dtype=np.int_, endpoint=True), npt.NDArray[np.int_]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_closed, dtype=np.int_, endpoint=True), npt.NDArray[np.int_]) +assert_type(def_gen.integers(0, I_u4_high_closed, dtype=np.int_, endpoint=True), npt.NDArray[np.int_]) + + +assert_type(def_gen.integers(4294967296, dtype="u4"), np.uint32) +assert_type(def_gen.integers(0, 4294967296, dtype="u4"), np.uint32) +assert_type(def_gen.integers(4294967295, dtype="u4", endpoint=True), np.uint32) +assert_type(def_gen.integers(0, 4294967295, dtype="u4", endpoint=True), np.uint32) +assert_type(def_gen.integers(I_u4_low_like, 4294967295, dtype="u4", endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_high_open, dtype="u4"), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_open, dtype="u4"), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(0, I_u4_high_open, dtype="u4"), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_high_closed, dtype="u4", endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_closed, dtype="u4", endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(0, I_u4_high_closed, dtype="u4", endpoint=True), npt.NDArray[np.uint32]) + +assert_type(def_gen.integers(4294967296, dtype="uint32"), np.uint32) +assert_type(def_gen.integers(0, 4294967296, dtype="uint32"), np.uint32) +assert_type(def_gen.integers(4294967295, dtype="uint32", endpoint=True), np.uint32) +assert_type(def_gen.integers(0, 4294967295, dtype="uint32", endpoint=True), np.uint32) +assert_type(def_gen.integers(I_u4_low_like, 4294967295, dtype="uint32", endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_high_open, dtype="uint32"), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_open, dtype="uint32"), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(0, I_u4_high_open, dtype="uint32"), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_high_closed, dtype="uint32", endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_closed, dtype="uint32", endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(0, I_u4_high_closed, dtype="uint32", endpoint=True), npt.NDArray[np.uint32]) + +assert_type(def_gen.integers(4294967296, dtype=np.uint32), np.uint32) +assert_type(def_gen.integers(0, 4294967296, dtype=np.uint32), np.uint32) +assert_type(def_gen.integers(4294967295, dtype=np.uint32, endpoint=True), np.uint32) +assert_type(def_gen.integers(0, 4294967295, dtype=np.uint32, endpoint=True), np.uint32) +assert_type(def_gen.integers(I_u4_low_like, 4294967295, dtype=np.uint32, endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_high_open, dtype=np.uint32), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_open, dtype=np.uint32), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(0, I_u4_high_open, dtype=np.uint32), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_high_closed, dtype=np.uint32, endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_closed, dtype=np.uint32, endpoint=True), npt.NDArray[np.uint32]) +assert_type(def_gen.integers(0, I_u4_high_closed, dtype=np.uint32, endpoint=True), npt.NDArray[np.uint32]) + +assert_type(def_gen.integers(4294967296, dtype=np.uint), np.uint) +assert_type(def_gen.integers(0, 4294967296, dtype=np.uint), np.uint) +assert_type(def_gen.integers(4294967295, dtype=np.uint, endpoint=True), np.uint) +assert_type(def_gen.integers(0, 4294967295, dtype=np.uint, endpoint=True), np.uint) +assert_type(def_gen.integers(I_u4_low_like, 4294967295, dtype=np.uint, endpoint=True), npt.NDArray[np.uint]) +assert_type(def_gen.integers(I_u4_high_open, dtype=np.uint), npt.NDArray[np.uint]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_open, dtype=np.uint), npt.NDArray[np.uint]) +assert_type(def_gen.integers(0, I_u4_high_open, dtype=np.uint), npt.NDArray[np.uint]) +assert_type(def_gen.integers(I_u4_high_closed, dtype=np.uint, endpoint=True), npt.NDArray[np.uint]) +assert_type(def_gen.integers(I_u4_low, I_u4_high_closed, dtype=np.uint, endpoint=True), npt.NDArray[np.uint]) +assert_type(def_gen.integers(0, I_u4_high_closed, dtype=np.uint, endpoint=True), npt.NDArray[np.uint]) + +I_u8_low: npt.NDArray[np.uint64] = np.array([0], dtype=np.uint64) +I_u8_low_like: list[int] = [0] +I_u8_high_open: npt.NDArray[np.uint64] = np.array([18446744073709551615], dtype=np.uint64) +I_u8_high_closed: npt.NDArray[np.uint64] = np.array([18446744073709551615], dtype=np.uint64) + +assert_type(def_gen.integers(18446744073709551616, dtype="u8"), np.uint64) +assert_type(def_gen.integers(0, 18446744073709551616, dtype="u8"), np.uint64) +assert_type(def_gen.integers(18446744073709551615, dtype="u8", endpoint=True), np.uint64) +assert_type(def_gen.integers(0, 18446744073709551615, dtype="u8", endpoint=True), np.uint64) +assert_type(def_gen.integers(I_u8_low_like, 18446744073709551615, dtype="u8", endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_high_open, dtype="u8"), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_low, I_u8_high_open, dtype="u8"), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(0, I_u8_high_open, dtype="u8"), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_high_closed, dtype="u8", endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_low, I_u8_high_closed, dtype="u8", endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(0, I_u8_high_closed, dtype="u8", endpoint=True), npt.NDArray[np.uint64]) + +assert_type(def_gen.integers(18446744073709551616, dtype="uint64"), np.uint64) +assert_type(def_gen.integers(0, 18446744073709551616, dtype="uint64"), np.uint64) +assert_type(def_gen.integers(18446744073709551615, dtype="uint64", endpoint=True), np.uint64) +assert_type(def_gen.integers(0, 18446744073709551615, dtype="uint64", endpoint=True), np.uint64) +assert_type(def_gen.integers(I_u8_low_like, 18446744073709551615, dtype="uint64", endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_high_open, dtype="uint64"), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_low, I_u8_high_open, dtype="uint64"), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(0, I_u8_high_open, dtype="uint64"), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_high_closed, dtype="uint64", endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_low, I_u8_high_closed, dtype="uint64", endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(0, I_u8_high_closed, dtype="uint64", endpoint=True), npt.NDArray[np.uint64]) + +assert_type(def_gen.integers(18446744073709551616, dtype=np.uint64), np.uint64) +assert_type(def_gen.integers(0, 18446744073709551616, dtype=np.uint64), np.uint64) +assert_type(def_gen.integers(18446744073709551615, dtype=np.uint64, endpoint=True), np.uint64) +assert_type(def_gen.integers(0, 18446744073709551615, dtype=np.uint64, endpoint=True), np.uint64) +assert_type(def_gen.integers(I_u8_low_like, 18446744073709551615, dtype=np.uint64, endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_high_open, dtype=np.uint64), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_low, I_u8_high_open, dtype=np.uint64), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(0, I_u8_high_open, dtype=np.uint64), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_high_closed, dtype=np.uint64, endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(I_u8_low, I_u8_high_closed, dtype=np.uint64, endpoint=True), npt.NDArray[np.uint64]) +assert_type(def_gen.integers(0, I_u8_high_closed, dtype=np.uint64, endpoint=True), npt.NDArray[np.uint64]) + +I_i1_low: npt.NDArray[np.int8] = np.array([-128], dtype=np.int8) +I_i1_low_like: list[int] = [-128] +I_i1_high_open: npt.NDArray[np.int8] = np.array([127], dtype=np.int8) +I_i1_high_closed: npt.NDArray[np.int8] = np.array([127], dtype=np.int8) + +assert_type(def_gen.integers(128, dtype="i1"), np.int8) +assert_type(def_gen.integers(-128, 128, dtype="i1"), np.int8) +assert_type(def_gen.integers(127, dtype="i1", endpoint=True), np.int8) +assert_type(def_gen.integers(-128, 127, dtype="i1", endpoint=True), np.int8) +assert_type(def_gen.integers(I_i1_low_like, 127, dtype="i1", endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_high_open, dtype="i1"), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_low, I_i1_high_open, dtype="i1"), npt.NDArray[np.int8]) +assert_type(def_gen.integers(-128, I_i1_high_open, dtype="i1"), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_high_closed, dtype="i1", endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_low, I_i1_high_closed, dtype="i1", endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(-128, I_i1_high_closed, dtype="i1", endpoint=True), npt.NDArray[np.int8]) + +assert_type(def_gen.integers(128, dtype="int8"), np.int8) +assert_type(def_gen.integers(-128, 128, dtype="int8"), np.int8) +assert_type(def_gen.integers(127, dtype="int8", endpoint=True), np.int8) +assert_type(def_gen.integers(-128, 127, dtype="int8", endpoint=True), np.int8) +assert_type(def_gen.integers(I_i1_low_like, 127, dtype="int8", endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_high_open, dtype="int8"), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_low, I_i1_high_open, dtype="int8"), npt.NDArray[np.int8]) +assert_type(def_gen.integers(-128, I_i1_high_open, dtype="int8"), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_high_closed, dtype="int8", endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_low, I_i1_high_closed, dtype="int8", endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(-128, I_i1_high_closed, dtype="int8", endpoint=True), npt.NDArray[np.int8]) + +assert_type(def_gen.integers(128, dtype=np.int8), np.int8) +assert_type(def_gen.integers(-128, 128, dtype=np.int8), np.int8) +assert_type(def_gen.integers(127, dtype=np.int8, endpoint=True), np.int8) +assert_type(def_gen.integers(-128, 127, dtype=np.int8, endpoint=True), np.int8) +assert_type(def_gen.integers(I_i1_low_like, 127, dtype=np.int8, endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_high_open, dtype=np.int8), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_low, I_i1_high_open, dtype=np.int8), npt.NDArray[np.int8]) +assert_type(def_gen.integers(-128, I_i1_high_open, dtype=np.int8), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_high_closed, dtype=np.int8, endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(I_i1_low, I_i1_high_closed, dtype=np.int8, endpoint=True), npt.NDArray[np.int8]) +assert_type(def_gen.integers(-128, I_i1_high_closed, dtype=np.int8, endpoint=True), npt.NDArray[np.int8]) + +I_i2_low: npt.NDArray[np.int16] = np.array([-32768], dtype=np.int16) +I_i2_low_like: list[int] = [-32768] +I_i2_high_open: npt.NDArray[np.int16] = np.array([32767], dtype=np.int16) +I_i2_high_closed: npt.NDArray[np.int16] = np.array([32767], dtype=np.int16) + +assert_type(def_gen.integers(32768, dtype="i2"), np.int16) +assert_type(def_gen.integers(-32768, 32768, dtype="i2"), np.int16) +assert_type(def_gen.integers(32767, dtype="i2", endpoint=True), np.int16) +assert_type(def_gen.integers(-32768, 32767, dtype="i2", endpoint=True), np.int16) +assert_type(def_gen.integers(I_i2_low_like, 32767, dtype="i2", endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_high_open, dtype="i2"), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_low, I_i2_high_open, dtype="i2"), npt.NDArray[np.int16]) +assert_type(def_gen.integers(-32768, I_i2_high_open, dtype="i2"), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_high_closed, dtype="i2", endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_low, I_i2_high_closed, dtype="i2", endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(-32768, I_i2_high_closed, dtype="i2", endpoint=True), npt.NDArray[np.int16]) + +assert_type(def_gen.integers(32768, dtype="int16"), np.int16) +assert_type(def_gen.integers(-32768, 32768, dtype="int16"), np.int16) +assert_type(def_gen.integers(32767, dtype="int16", endpoint=True), np.int16) +assert_type(def_gen.integers(-32768, 32767, dtype="int16", endpoint=True), np.int16) +assert_type(def_gen.integers(I_i2_low_like, 32767, dtype="int16", endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_high_open, dtype="int16"), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_low, I_i2_high_open, dtype="int16"), npt.NDArray[np.int16]) +assert_type(def_gen.integers(-32768, I_i2_high_open, dtype="int16"), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_high_closed, dtype="int16", endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_low, I_i2_high_closed, dtype="int16", endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(-32768, I_i2_high_closed, dtype="int16", endpoint=True), npt.NDArray[np.int16]) + +assert_type(def_gen.integers(32768, dtype=np.int16), np.int16) +assert_type(def_gen.integers(-32768, 32768, dtype=np.int16), np.int16) +assert_type(def_gen.integers(32767, dtype=np.int16, endpoint=True), np.int16) +assert_type(def_gen.integers(-32768, 32767, dtype=np.int16, endpoint=True), np.int16) +assert_type(def_gen.integers(I_i2_low_like, 32767, dtype=np.int16, endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_high_open, dtype=np.int16), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_low, I_i2_high_open, dtype=np.int16), npt.NDArray[np.int16]) +assert_type(def_gen.integers(-32768, I_i2_high_open, dtype=np.int16), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_high_closed, dtype=np.int16, endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(I_i2_low, I_i2_high_closed, dtype=np.int16, endpoint=True), npt.NDArray[np.int16]) +assert_type(def_gen.integers(-32768, I_i2_high_closed, dtype=np.int16, endpoint=True), npt.NDArray[np.int16]) + +I_i4_low: npt.NDArray[np.int32] = np.array([-2147483648], dtype=np.int32) +I_i4_low_like: list[int] = [-2147483648] +I_i4_high_open: npt.NDArray[np.int32] = np.array([2147483647], dtype=np.int32) +I_i4_high_closed: npt.NDArray[np.int32] = np.array([2147483647], dtype=np.int32) + +assert_type(def_gen.integers(2147483648, dtype="i4"), np.int32) +assert_type(def_gen.integers(-2147483648, 2147483648, dtype="i4"), np.int32) +assert_type(def_gen.integers(2147483647, dtype="i4", endpoint=True), np.int32) +assert_type(def_gen.integers(-2147483648, 2147483647, dtype="i4", endpoint=True), np.int32) +assert_type(def_gen.integers(I_i4_low_like, 2147483647, dtype="i4", endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_high_open, dtype="i4"), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_low, I_i4_high_open, dtype="i4"), npt.NDArray[np.int32]) +assert_type(def_gen.integers(-2147483648, I_i4_high_open, dtype="i4"), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_high_closed, dtype="i4", endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_low, I_i4_high_closed, dtype="i4", endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(-2147483648, I_i4_high_closed, dtype="i4", endpoint=True), npt.NDArray[np.int32]) + +assert_type(def_gen.integers(2147483648, dtype="int32"), np.int32) +assert_type(def_gen.integers(-2147483648, 2147483648, dtype="int32"), np.int32) +assert_type(def_gen.integers(2147483647, dtype="int32", endpoint=True), np.int32) +assert_type(def_gen.integers(-2147483648, 2147483647, dtype="int32", endpoint=True), np.int32) +assert_type(def_gen.integers(I_i4_low_like, 2147483647, dtype="int32", endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_high_open, dtype="int32"), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_low, I_i4_high_open, dtype="int32"), npt.NDArray[np.int32]) +assert_type(def_gen.integers(-2147483648, I_i4_high_open, dtype="int32"), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_high_closed, dtype="int32", endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_low, I_i4_high_closed, dtype="int32", endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(-2147483648, I_i4_high_closed, dtype="int32", endpoint=True), npt.NDArray[np.int32]) + +assert_type(def_gen.integers(2147483648, dtype=np.int32), np.int32) +assert_type(def_gen.integers(-2147483648, 2147483648, dtype=np.int32), np.int32) +assert_type(def_gen.integers(2147483647, dtype=np.int32, endpoint=True), np.int32) +assert_type(def_gen.integers(-2147483648, 2147483647, dtype=np.int32, endpoint=True), np.int32) +assert_type(def_gen.integers(I_i4_low_like, 2147483647, dtype=np.int32, endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_high_open, dtype=np.int32), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_low, I_i4_high_open, dtype=np.int32), npt.NDArray[np.int32]) +assert_type(def_gen.integers(-2147483648, I_i4_high_open, dtype=np.int32), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_high_closed, dtype=np.int32, endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(I_i4_low, I_i4_high_closed, dtype=np.int32, endpoint=True), npt.NDArray[np.int32]) +assert_type(def_gen.integers(-2147483648, I_i4_high_closed, dtype=np.int32, endpoint=True), npt.NDArray[np.int32]) + +I_i8_low: npt.NDArray[np.int64] = np.array([-9223372036854775808], dtype=np.int64) +I_i8_low_like: list[int] = [-9223372036854775808] +I_i8_high_open: npt.NDArray[np.int64] = np.array([9223372036854775807], dtype=np.int64) +I_i8_high_closed: npt.NDArray[np.int64] = np.array([9223372036854775807], dtype=np.int64) + +assert_type(def_gen.integers(9223372036854775808, dtype="i8"), np.int64) +assert_type(def_gen.integers(-9223372036854775808, 9223372036854775808, dtype="i8"), np.int64) +assert_type(def_gen.integers(9223372036854775807, dtype="i8", endpoint=True), np.int64) +assert_type(def_gen.integers(-9223372036854775808, 9223372036854775807, dtype="i8", endpoint=True), np.int64) +assert_type(def_gen.integers(I_i8_low_like, 9223372036854775807, dtype="i8", endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_high_open, dtype="i8"), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_low, I_i8_high_open, dtype="i8"), npt.NDArray[np.int64]) +assert_type(def_gen.integers(-9223372036854775808, I_i8_high_open, dtype="i8"), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_high_closed, dtype="i8", endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_low, I_i8_high_closed, dtype="i8", endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(-9223372036854775808, I_i8_high_closed, dtype="i8", endpoint=True), npt.NDArray[np.int64]) + +assert_type(def_gen.integers(9223372036854775808, dtype="int64"), np.int64) +assert_type(def_gen.integers(-9223372036854775808, 9223372036854775808, dtype="int64"), np.int64) +assert_type(def_gen.integers(9223372036854775807, dtype="int64", endpoint=True), np.int64) +assert_type(def_gen.integers(-9223372036854775808, 9223372036854775807, dtype="int64", endpoint=True), np.int64) +assert_type(def_gen.integers(I_i8_low_like, 9223372036854775807, dtype="int64", endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_high_open, dtype="int64"), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_low, I_i8_high_open, dtype="int64"), npt.NDArray[np.int64]) +assert_type(def_gen.integers(-9223372036854775808, I_i8_high_open, dtype="int64"), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_high_closed, dtype="int64", endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_low, I_i8_high_closed, dtype="int64", endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(-9223372036854775808, I_i8_high_closed, dtype="int64", endpoint=True), npt.NDArray[np.int64]) + +assert_type(def_gen.integers(9223372036854775808, dtype=np.int64), np.int64) +assert_type(def_gen.integers(-9223372036854775808, 9223372036854775808, dtype=np.int64), np.int64) +assert_type(def_gen.integers(9223372036854775807, dtype=np.int64, endpoint=True), np.int64) +assert_type(def_gen.integers(-9223372036854775808, 9223372036854775807, dtype=np.int64, endpoint=True), np.int64) +assert_type(def_gen.integers(I_i8_low_like, 9223372036854775807, dtype=np.int64, endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_high_open, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_low, I_i8_high_open, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(def_gen.integers(-9223372036854775808, I_i8_high_open, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_high_closed, dtype=np.int64, endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(I_i8_low, I_i8_high_closed, dtype=np.int64, endpoint=True), npt.NDArray[np.int64]) +assert_type(def_gen.integers(-9223372036854775808, I_i8_high_closed, dtype=np.int64, endpoint=True), npt.NDArray[np.int64]) + + +assert_type(def_gen.bit_generator, np.random.BitGenerator) + +assert_type(def_gen.bytes(2), bytes) + +assert_type(def_gen.choice(5), int) +assert_type(def_gen.choice(5, 3), npt.NDArray[np.int64]) +assert_type(def_gen.choice(5, 3, replace=True), npt.NDArray[np.int64]) +assert_type(def_gen.choice(5, 3, p=[1 / 5] * 5), npt.NDArray[np.int64]) +assert_type(def_gen.choice(5, 3, p=[1 / 5] * 5, replace=False), npt.NDArray[np.int64]) + +assert_type(def_gen.choice(["pooh", "rabbit", "piglet", "Christopher"]), Any) +assert_type(def_gen.choice(["pooh", "rabbit", "piglet", "Christopher"], 3), npt.NDArray[Any]) +assert_type(def_gen.choice(["pooh", "rabbit", "piglet", "Christopher"], 3, p=[1 / 4] * 4), npt.NDArray[Any]) +assert_type(def_gen.choice(["pooh", "rabbit", "piglet", "Christopher"], 3, replace=True), npt.NDArray[Any]) +assert_type(def_gen.choice(["pooh", "rabbit", "piglet", "Christopher"], 3, replace=False, p=np.array([1 / 8, 1 / 8, 1 / 2, 1 / 4])), npt.NDArray[Any]) + +assert_type(def_gen.dirichlet([0.5, 0.5]), npt.NDArray[np.float64]) +assert_type(def_gen.dirichlet(np.array([0.5, 0.5])), npt.NDArray[np.float64]) +assert_type(def_gen.dirichlet(np.array([0.5, 0.5]), size=3), npt.NDArray[np.float64]) + +assert_type(def_gen.multinomial(20, [1 / 6.0] * 6), npt.NDArray[np.int64]) +assert_type(def_gen.multinomial(20, np.array([0.5, 0.5])), npt.NDArray[np.int64]) +assert_type(def_gen.multinomial(20, [1 / 6.0] * 6, size=2), npt.NDArray[np.int64]) +assert_type(def_gen.multinomial([[10], [20]], [1 / 6.0] * 6, size=(2, 2)), npt.NDArray[np.int64]) +assert_type(def_gen.multinomial(np.array([[10], [20]]), np.array([0.5, 0.5]), size=(2, 2)), npt.NDArray[np.int64]) + +assert_type(def_gen.multivariate_hypergeometric([3, 5, 7], 2), npt.NDArray[np.int64]) +assert_type(def_gen.multivariate_hypergeometric(np.array([3, 5, 7]), 2), npt.NDArray[np.int64]) +assert_type(def_gen.multivariate_hypergeometric(np.array([3, 5, 7]), 2, size=4), npt.NDArray[np.int64]) +assert_type(def_gen.multivariate_hypergeometric(np.array([3, 5, 7]), 2, size=(4, 7)), npt.NDArray[np.int64]) +assert_type(def_gen.multivariate_hypergeometric([3, 5, 7], 2, method="count"), npt.NDArray[np.int64]) +assert_type(def_gen.multivariate_hypergeometric(np.array([3, 5, 7]), 2, method="marginals"), npt.NDArray[np.int64]) + +assert_type(def_gen.multivariate_normal([0.0], [[1.0]]), npt.NDArray[np.float64]) +assert_type(def_gen.multivariate_normal([0.0], np.array([[1.0]])), npt.NDArray[np.float64]) +assert_type(def_gen.multivariate_normal(np.array([0.0]), [[1.0]]), npt.NDArray[np.float64]) +assert_type(def_gen.multivariate_normal([0.0], np.array([[1.0]])), npt.NDArray[np.float64]) + +assert_type(def_gen.permutation(10), npt.NDArray[np.int64]) +assert_type(def_gen.permutation([1, 2, 3, 4]), npt.NDArray[Any]) +assert_type(def_gen.permutation(np.array([1, 2, 3, 4])), npt.NDArray[Any]) +assert_type(def_gen.permutation(D_2D, axis=1), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D_like), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D, axis=1), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D, out=D_2D), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D_like, out=D_2D), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D_like, out=D_2D), npt.NDArray[Any]) +assert_type(def_gen.permuted(D_2D, axis=1, out=D_2D), npt.NDArray[Any]) + +assert_type(def_gen.shuffle(np.arange(10)), None) +assert_type(def_gen.shuffle([1, 2, 3, 4, 5]), None) +assert_type(def_gen.shuffle(D_2D, axis=1), None) + +assert_type(np.random.Generator(pcg64), np.random.Generator) +assert_type(def_gen.__str__(), str) +assert_type(def_gen.__repr__(), str) +assert_type(def_gen.__setstate__(dict(def_gen.bit_generator.state)), None) + +# RandomState +random_st: np.random.RandomState = np.random.RandomState() + +assert_type(random_st.standard_normal(), float) +assert_type(random_st.standard_normal(size=None), float) +assert_type(random_st.standard_normal(size=1), npt.NDArray[np.float64]) + +assert_type(random_st.random(), float) +assert_type(random_st.random(size=None), float) +assert_type(random_st.random(size=1), npt.NDArray[np.float64]) + +assert_type(random_st.standard_cauchy(), float) +assert_type(random_st.standard_cauchy(size=None), float) +assert_type(random_st.standard_cauchy(size=1), npt.NDArray[np.float64]) + +assert_type(random_st.standard_exponential(), float) +assert_type(random_st.standard_exponential(size=None), float) +assert_type(random_st.standard_exponential(size=1), npt.NDArray[np.float64]) + +assert_type(random_st.zipf(1.5), int) +assert_type(random_st.zipf(1.5, size=None), int) +assert_type(random_st.zipf(1.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.zipf(D_arr_1p5), npt.NDArray[np.long]) +assert_type(random_st.zipf(D_arr_1p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.zipf(D_arr_like_1p5), npt.NDArray[np.long]) +assert_type(random_st.zipf(D_arr_like_1p5, size=1), npt.NDArray[np.long]) + +assert_type(random_st.weibull(0.5), float) +assert_type(random_st.weibull(0.5, size=None), float) +assert_type(random_st.weibull(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.weibull(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.weibull(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.weibull(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.weibull(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.standard_t(0.5), float) +assert_type(random_st.standard_t(0.5, size=None), float) +assert_type(random_st.standard_t(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.standard_t(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.standard_t(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.standard_t(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.standard_t(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.poisson(0.5), int) +assert_type(random_st.poisson(0.5, size=None), int) +assert_type(random_st.poisson(0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.poisson(D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.poisson(D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.poisson(D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.poisson(D_arr_like_0p5, size=1), npt.NDArray[np.long]) + +assert_type(random_st.power(0.5), float) +assert_type(random_st.power(0.5, size=None), float) +assert_type(random_st.power(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.power(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.power(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.power(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.power(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.pareto(0.5), float) +assert_type(random_st.pareto(0.5, size=None), float) +assert_type(random_st.pareto(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.pareto(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.pareto(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.pareto(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.pareto(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.chisquare(0.5), float) +assert_type(random_st.chisquare(0.5, size=None), float) +assert_type(random_st.chisquare(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.chisquare(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.chisquare(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.chisquare(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.chisquare(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.exponential(0.5), float) +assert_type(random_st.exponential(0.5, size=None), float) +assert_type(random_st.exponential(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.exponential(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.exponential(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.exponential(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.exponential(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.geometric(0.5), int) +assert_type(random_st.geometric(0.5, size=None), int) +assert_type(random_st.geometric(0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.geometric(D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.geometric(D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.geometric(D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.geometric(D_arr_like_0p5, size=1), npt.NDArray[np.long]) + +assert_type(random_st.logseries(0.5), int) +assert_type(random_st.logseries(0.5, size=None), int) +assert_type(random_st.logseries(0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.logseries(D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.logseries(D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.logseries(D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.logseries(D_arr_like_0p5, size=1), npt.NDArray[np.long]) + +assert_type(random_st.rayleigh(0.5), float) +assert_type(random_st.rayleigh(0.5, size=None), float) +assert_type(random_st.rayleigh(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.rayleigh(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.rayleigh(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.rayleigh(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.rayleigh(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.standard_gamma(0.5), float) +assert_type(random_st.standard_gamma(0.5, size=None), float) +assert_type(random_st.standard_gamma(0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.standard_gamma(D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.standard_gamma(D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.standard_gamma(D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.standard_gamma(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.standard_gamma(D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.vonmises(0.5, 0.5), float) +assert_type(random_st.vonmises(0.5, 0.5, size=None), float) +assert_type(random_st.vonmises(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.vonmises(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.wald(0.5, 0.5), float) +assert_type(random_st.wald(0.5, 0.5, size=None), float) +assert_type(random_st.wald(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.wald(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.wald(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.wald(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.wald(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.uniform(0.5, 0.5), float) +assert_type(random_st.uniform(0.5, 0.5, size=None), float) +assert_type(random_st.uniform(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.uniform(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.uniform(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.uniform(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.uniform(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.beta(0.5, 0.5), float) +assert_type(random_st.beta(0.5, 0.5, size=None), float) +assert_type(random_st.beta(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.beta(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.beta(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.beta(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.beta(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.f(0.5, 0.5), float) +assert_type(random_st.f(0.5, 0.5, size=None), float) +assert_type(random_st.f(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.f(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.f(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.f(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.f(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.gamma(0.5, 0.5), float) +assert_type(random_st.gamma(0.5, 0.5, size=None), float) +assert_type(random_st.gamma(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.gamma(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gamma(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.gamma(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gamma(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.gumbel(0.5, 0.5), float) +assert_type(random_st.gumbel(0.5, 0.5, size=None), float) +assert_type(random_st.gumbel(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.gumbel(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.laplace(0.5, 0.5), float) +assert_type(random_st.laplace(0.5, 0.5, size=None), float) +assert_type(random_st.laplace(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.laplace(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.laplace(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.laplace(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.laplace(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.logistic(0.5, 0.5), float) +assert_type(random_st.logistic(0.5, 0.5, size=None), float) +assert_type(random_st.logistic(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.logistic(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.logistic(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.logistic(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.logistic(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.lognormal(0.5, 0.5), float) +assert_type(random_st.lognormal(0.5, 0.5, size=None), float) +assert_type(random_st.lognormal(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.lognormal(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.noncentral_chisquare(0.5, 0.5), float) +assert_type(random_st.noncentral_chisquare(0.5, 0.5, size=None), float) +assert_type(random_st.noncentral_chisquare(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_chisquare(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.normal(0.5, 0.5), float) +assert_type(random_st.normal(0.5, 0.5, size=None), float) +assert_type(random_st.normal(0.5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.normal(0.5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_0p5, 0.5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.normal(0.5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_like_0p5, 0.5), npt.NDArray[np.float64]) +assert_type(random_st.normal(0.5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_0p5, D_arr_0p5), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_like_0p5, D_arr_like_0p5), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_0p5, D_arr_0p5, size=1), npt.NDArray[np.float64]) +assert_type(random_st.normal(D_arr_like_0p5, D_arr_like_0p5, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.triangular(0.1, 0.5, 0.9), float) +assert_type(random_st.triangular(0.1, 0.5, 0.9, size=None), float) +assert_type(random_st.triangular(0.1, 0.5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_0p1, 0.5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.triangular(0.1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_0p1, 0.5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.triangular(0.1, D_arr_0p5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_like_0p1, 0.5, D_arr_0p9), npt.NDArray[np.float64]) +assert_type(random_st.triangular(0.5, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_0p1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_like_0p1, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_0p1, D_arr_0p5, D_arr_0p9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.triangular(D_arr_like_0p1, D_arr_like_0p5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.noncentral_f(0.1, 0.5, 0.9), float) +assert_type(random_st.noncentral_f(0.1, 0.5, 0.9, size=None), float) +assert_type(random_st.noncentral_f(0.1, 0.5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_0p1, 0.5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(0.1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_0p1, 0.5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(0.1, D_arr_0p5, 0.9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_like_0p1, 0.5, D_arr_0p9), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(0.5, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_0p1, D_arr_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_like_0p1, D_arr_like_0p5, 0.9), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_0p1, D_arr_0p5, D_arr_0p9, size=1), npt.NDArray[np.float64]) +assert_type(random_st.noncentral_f(D_arr_like_0p1, D_arr_like_0p5, D_arr_like_0p9, size=1), npt.NDArray[np.float64]) + +assert_type(random_st.binomial(10, 0.5), int) +assert_type(random_st.binomial(10, 0.5, size=None), int) +assert_type(random_st.binomial(10, 0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_10, 0.5), npt.NDArray[np.long]) +assert_type(random_st.binomial(10, D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_10, 0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.binomial(10, D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_like_10, 0.5), npt.NDArray[np.long]) +assert_type(random_st.binomial(10, D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_10, D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_like_10, D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_10, D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.binomial(I_arr_like_10, D_arr_like_0p5, size=1), npt.NDArray[np.long]) + +assert_type(random_st.negative_binomial(10, 0.5), int) +assert_type(random_st.negative_binomial(10, 0.5, size=None), int) +assert_type(random_st.negative_binomial(10, 0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_10, 0.5), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(10, D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_10, 0.5, size=1), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(10, D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_like_10, 0.5), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(10, D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_10, D_arr_0p5), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_like_10, D_arr_like_0p5), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_10, D_arr_0p5, size=1), npt.NDArray[np.long]) +assert_type(random_st.negative_binomial(I_arr_like_10, D_arr_like_0p5, size=1), npt.NDArray[np.long]) + +assert_type(random_st.hypergeometric(20, 20, 10), int) +assert_type(random_st.hypergeometric(20, 20, 10, size=None), int) +assert_type(random_st.hypergeometric(20, 20, 10, size=1), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_20, 20, 10), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(20, I_arr_20, 10), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_20, 20, I_arr_like_10, size=1), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(20, I_arr_20, 10, size=1), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_like_20, 20, I_arr_10), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(20, I_arr_like_20, 10), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_20, I_arr_20, 10), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_like_20, I_arr_like_20, 10), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_20, I_arr_20, I_arr_10, size=1), npt.NDArray[np.long]) +assert_type(random_st.hypergeometric(I_arr_like_20, I_arr_like_20, I_arr_like_10, size=1), npt.NDArray[np.long]) + +assert_type(random_st.randint(0, 100), int) +assert_type(random_st.randint(100), int) +assert_type(random_st.randint([100]), npt.NDArray[np.long]) +assert_type(random_st.randint(0, [100]), npt.NDArray[np.long]) + +assert_type(random_st.randint(2, dtype=bool), bool) +assert_type(random_st.randint(0, 2, dtype=bool), bool) +assert_type(random_st.randint(I_bool_high_open, dtype=bool), npt.NDArray[np.bool]) +assert_type(random_st.randint(I_bool_low, I_bool_high_open, dtype=bool), npt.NDArray[np.bool]) +assert_type(random_st.randint(0, I_bool_high_open, dtype=bool), npt.NDArray[np.bool]) + +assert_type(random_st.randint(2, dtype=np.bool), np.bool) +assert_type(random_st.randint(0, 2, dtype=np.bool), np.bool) +assert_type(random_st.randint(I_bool_high_open, dtype=np.bool), npt.NDArray[np.bool]) +assert_type(random_st.randint(I_bool_low, I_bool_high_open, dtype=np.bool), npt.NDArray[np.bool]) +assert_type(random_st.randint(0, I_bool_high_open, dtype=np.bool), npt.NDArray[np.bool]) + +assert_type(random_st.randint(256, dtype="u1"), np.uint8) +assert_type(random_st.randint(0, 256, dtype="u1"), np.uint8) +assert_type(random_st.randint(I_u1_high_open, dtype="u1"), npt.NDArray[np.uint8]) +assert_type(random_st.randint(I_u1_low, I_u1_high_open, dtype="u1"), npt.NDArray[np.uint8]) +assert_type(random_st.randint(0, I_u1_high_open, dtype="u1"), npt.NDArray[np.uint8]) + +assert_type(random_st.randint(256, dtype="uint8"), np.uint8) +assert_type(random_st.randint(0, 256, dtype="uint8"), np.uint8) +assert_type(random_st.randint(I_u1_high_open, dtype="uint8"), npt.NDArray[np.uint8]) +assert_type(random_st.randint(I_u1_low, I_u1_high_open, dtype="uint8"), npt.NDArray[np.uint8]) +assert_type(random_st.randint(0, I_u1_high_open, dtype="uint8"), npt.NDArray[np.uint8]) + +assert_type(random_st.randint(256, dtype=np.uint8), np.uint8) +assert_type(random_st.randint(0, 256, dtype=np.uint8), np.uint8) +assert_type(random_st.randint(I_u1_high_open, dtype=np.uint8), npt.NDArray[np.uint8]) +assert_type(random_st.randint(I_u1_low, I_u1_high_open, dtype=np.uint8), npt.NDArray[np.uint8]) +assert_type(random_st.randint(0, I_u1_high_open, dtype=np.uint8), npt.NDArray[np.uint8]) + +assert_type(random_st.randint(65536, dtype="u2"), np.uint16) +assert_type(random_st.randint(0, 65536, dtype="u2"), np.uint16) +assert_type(random_st.randint(I_u2_high_open, dtype="u2"), npt.NDArray[np.uint16]) +assert_type(random_st.randint(I_u2_low, I_u2_high_open, dtype="u2"), npt.NDArray[np.uint16]) +assert_type(random_st.randint(0, I_u2_high_open, dtype="u2"), npt.NDArray[np.uint16]) + +assert_type(random_st.randint(65536, dtype="uint16"), np.uint16) +assert_type(random_st.randint(0, 65536, dtype="uint16"), np.uint16) +assert_type(random_st.randint(I_u2_high_open, dtype="uint16"), npt.NDArray[np.uint16]) +assert_type(random_st.randint(I_u2_low, I_u2_high_open, dtype="uint16"), npt.NDArray[np.uint16]) +assert_type(random_st.randint(0, I_u2_high_open, dtype="uint16"), npt.NDArray[np.uint16]) + +assert_type(random_st.randint(65536, dtype=np.uint16), np.uint16) +assert_type(random_st.randint(0, 65536, dtype=np.uint16), np.uint16) +assert_type(random_st.randint(I_u2_high_open, dtype=np.uint16), npt.NDArray[np.uint16]) +assert_type(random_st.randint(I_u2_low, I_u2_high_open, dtype=np.uint16), npt.NDArray[np.uint16]) +assert_type(random_st.randint(0, I_u2_high_open, dtype=np.uint16), npt.NDArray[np.uint16]) + +assert_type(random_st.randint(4294967296, dtype="u4"), np.uint32) +assert_type(random_st.randint(0, 4294967296, dtype="u4"), np.uint32) +assert_type(random_st.randint(I_u4_high_open, dtype="u4"), npt.NDArray[np.uint32]) +assert_type(random_st.randint(I_u4_low, I_u4_high_open, dtype="u4"), npt.NDArray[np.uint32]) +assert_type(random_st.randint(0, I_u4_high_open, dtype="u4"), npt.NDArray[np.uint32]) + +assert_type(random_st.randint(4294967296, dtype="uint32"), np.uint32) +assert_type(random_st.randint(0, 4294967296, dtype="uint32"), np.uint32) +assert_type(random_st.randint(I_u4_high_open, dtype="uint32"), npt.NDArray[np.uint32]) +assert_type(random_st.randint(I_u4_low, I_u4_high_open, dtype="uint32"), npt.NDArray[np.uint32]) +assert_type(random_st.randint(0, I_u4_high_open, dtype="uint32"), npt.NDArray[np.uint32]) + +assert_type(random_st.randint(4294967296, dtype=np.uint32), np.uint32) +assert_type(random_st.randint(0, 4294967296, dtype=np.uint32), np.uint32) +assert_type(random_st.randint(I_u4_high_open, dtype=np.uint32), npt.NDArray[np.uint32]) +assert_type(random_st.randint(I_u4_low, I_u4_high_open, dtype=np.uint32), npt.NDArray[np.uint32]) +assert_type(random_st.randint(0, I_u4_high_open, dtype=np.uint32), npt.NDArray[np.uint32]) + +assert_type(random_st.randint(4294967296, dtype=np.uint), np.uint) +assert_type(random_st.randint(0, 4294967296, dtype=np.uint), np.uint) +assert_type(random_st.randint(I_u4_high_open, dtype=np.uint), npt.NDArray[np.uint]) +assert_type(random_st.randint(I_u4_low, I_u4_high_open, dtype=np.uint), npt.NDArray[np.uint]) +assert_type(random_st.randint(0, I_u4_high_open, dtype=np.uint), npt.NDArray[np.uint]) + +assert_type(random_st.randint(18446744073709551616, dtype="u8"), np.uint64) +assert_type(random_st.randint(0, 18446744073709551616, dtype="u8"), np.uint64) +assert_type(random_st.randint(I_u8_high_open, dtype="u8"), npt.NDArray[np.uint64]) +assert_type(random_st.randint(I_u8_low, I_u8_high_open, dtype="u8"), npt.NDArray[np.uint64]) +assert_type(random_st.randint(0, I_u8_high_open, dtype="u8"), npt.NDArray[np.uint64]) + +assert_type(random_st.randint(18446744073709551616, dtype="uint64"), np.uint64) +assert_type(random_st.randint(0, 18446744073709551616, dtype="uint64"), np.uint64) +assert_type(random_st.randint(I_u8_high_open, dtype="uint64"), npt.NDArray[np.uint64]) +assert_type(random_st.randint(I_u8_low, I_u8_high_open, dtype="uint64"), npt.NDArray[np.uint64]) +assert_type(random_st.randint(0, I_u8_high_open, dtype="uint64"), npt.NDArray[np.uint64]) + +assert_type(random_st.randint(18446744073709551616, dtype=np.uint64), np.uint64) +assert_type(random_st.randint(0, 18446744073709551616, dtype=np.uint64), np.uint64) +assert_type(random_st.randint(I_u8_high_open, dtype=np.uint64), npt.NDArray[np.uint64]) +assert_type(random_st.randint(I_u8_low, I_u8_high_open, dtype=np.uint64), npt.NDArray[np.uint64]) +assert_type(random_st.randint(0, I_u8_high_open, dtype=np.uint64), npt.NDArray[np.uint64]) + +assert_type(random_st.randint(128, dtype="i1"), np.int8) +assert_type(random_st.randint(-128, 128, dtype="i1"), np.int8) +assert_type(random_st.randint(I_i1_high_open, dtype="i1"), npt.NDArray[np.int8]) +assert_type(random_st.randint(I_i1_low, I_i1_high_open, dtype="i1"), npt.NDArray[np.int8]) +assert_type(random_st.randint(-128, I_i1_high_open, dtype="i1"), npt.NDArray[np.int8]) + +assert_type(random_st.randint(128, dtype="int8"), np.int8) +assert_type(random_st.randint(-128, 128, dtype="int8"), np.int8) +assert_type(random_st.randint(I_i1_high_open, dtype="int8"), npt.NDArray[np.int8]) +assert_type(random_st.randint(I_i1_low, I_i1_high_open, dtype="int8"), npt.NDArray[np.int8]) +assert_type(random_st.randint(-128, I_i1_high_open, dtype="int8"), npt.NDArray[np.int8]) + +assert_type(random_st.randint(128, dtype=np.int8), np.int8) +assert_type(random_st.randint(-128, 128, dtype=np.int8), np.int8) +assert_type(random_st.randint(I_i1_high_open, dtype=np.int8), npt.NDArray[np.int8]) +assert_type(random_st.randint(I_i1_low, I_i1_high_open, dtype=np.int8), npt.NDArray[np.int8]) +assert_type(random_st.randint(-128, I_i1_high_open, dtype=np.int8), npt.NDArray[np.int8]) + +assert_type(random_st.randint(32768, dtype="i2"), np.int16) +assert_type(random_st.randint(-32768, 32768, dtype="i2"), np.int16) +assert_type(random_st.randint(I_i2_high_open, dtype="i2"), npt.NDArray[np.int16]) +assert_type(random_st.randint(I_i2_low, I_i2_high_open, dtype="i2"), npt.NDArray[np.int16]) +assert_type(random_st.randint(-32768, I_i2_high_open, dtype="i2"), npt.NDArray[np.int16]) + +assert_type(random_st.randint(32768, dtype="int16"), np.int16) +assert_type(random_st.randint(-32768, 32768, dtype="int16"), np.int16) +assert_type(random_st.randint(I_i2_high_open, dtype="int16"), npt.NDArray[np.int16]) +assert_type(random_st.randint(I_i2_low, I_i2_high_open, dtype="int16"), npt.NDArray[np.int16]) +assert_type(random_st.randint(-32768, I_i2_high_open, dtype="int16"), npt.NDArray[np.int16]) + +assert_type(random_st.randint(32768, dtype=np.int16), np.int16) +assert_type(random_st.randint(-32768, 32768, dtype=np.int16), np.int16) +assert_type(random_st.randint(I_i2_high_open, dtype=np.int16), npt.NDArray[np.int16]) +assert_type(random_st.randint(I_i2_low, I_i2_high_open, dtype=np.int16), npt.NDArray[np.int16]) +assert_type(random_st.randint(-32768, I_i2_high_open, dtype=np.int16), npt.NDArray[np.int16]) + +assert_type(random_st.randint(2147483648, dtype="i4"), np.int32) +assert_type(random_st.randint(-2147483648, 2147483648, dtype="i4"), np.int32) +assert_type(random_st.randint(I_i4_high_open, dtype="i4"), npt.NDArray[np.int32]) +assert_type(random_st.randint(I_i4_low, I_i4_high_open, dtype="i4"), npt.NDArray[np.int32]) +assert_type(random_st.randint(-2147483648, I_i4_high_open, dtype="i4"), npt.NDArray[np.int32]) + +assert_type(random_st.randint(2147483648, dtype="int32"), np.int32) +assert_type(random_st.randint(-2147483648, 2147483648, dtype="int32"), np.int32) +assert_type(random_st.randint(I_i4_high_open, dtype="int32"), npt.NDArray[np.int32]) +assert_type(random_st.randint(I_i4_low, I_i4_high_open, dtype="int32"), npt.NDArray[np.int32]) +assert_type(random_st.randint(-2147483648, I_i4_high_open, dtype="int32"), npt.NDArray[np.int32]) + +assert_type(random_st.randint(2147483648, dtype=np.int32), np.int32) +assert_type(random_st.randint(-2147483648, 2147483648, dtype=np.int32), np.int32) +assert_type(random_st.randint(I_i4_high_open, dtype=np.int32), npt.NDArray[np.int32]) +assert_type(random_st.randint(I_i4_low, I_i4_high_open, dtype=np.int32), npt.NDArray[np.int32]) +assert_type(random_st.randint(-2147483648, I_i4_high_open, dtype=np.int32), npt.NDArray[np.int32]) + +assert_type(random_st.randint(2147483648, dtype=np.int_), np.int_) +assert_type(random_st.randint(-2147483648, 2147483648, dtype=np.int_), np.int_) +assert_type(random_st.randint(I_i4_high_open, dtype=np.int_), npt.NDArray[np.int_]) +assert_type(random_st.randint(I_i4_low, I_i4_high_open, dtype=np.int_), npt.NDArray[np.int_]) +assert_type(random_st.randint(-2147483648, I_i4_high_open, dtype=np.int_), npt.NDArray[np.int_]) + +assert_type(random_st.randint(9223372036854775808, dtype="i8"), np.int64) +assert_type(random_st.randint(-9223372036854775808, 9223372036854775808, dtype="i8"), np.int64) +assert_type(random_st.randint(I_i8_high_open, dtype="i8"), npt.NDArray[np.int64]) +assert_type(random_st.randint(I_i8_low, I_i8_high_open, dtype="i8"), npt.NDArray[np.int64]) +assert_type(random_st.randint(-9223372036854775808, I_i8_high_open, dtype="i8"), npt.NDArray[np.int64]) + +assert_type(random_st.randint(9223372036854775808, dtype="int64"), np.int64) +assert_type(random_st.randint(-9223372036854775808, 9223372036854775808, dtype="int64"), np.int64) +assert_type(random_st.randint(I_i8_high_open, dtype="int64"), npt.NDArray[np.int64]) +assert_type(random_st.randint(I_i8_low, I_i8_high_open, dtype="int64"), npt.NDArray[np.int64]) +assert_type(random_st.randint(-9223372036854775808, I_i8_high_open, dtype="int64"), npt.NDArray[np.int64]) + +assert_type(random_st.randint(9223372036854775808, dtype=np.int64), np.int64) +assert_type(random_st.randint(-9223372036854775808, 9223372036854775808, dtype=np.int64), np.int64) +assert_type(random_st.randint(I_i8_high_open, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(random_st.randint(I_i8_low, I_i8_high_open, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(random_st.randint(-9223372036854775808, I_i8_high_open, dtype=np.int64), npt.NDArray[np.int64]) + +assert_type(random_st._bit_generator, np.random.BitGenerator) + +assert_type(random_st.bytes(2), bytes) + +assert_type(random_st.choice(5), int) +assert_type(random_st.choice(5, 3), npt.NDArray[np.long]) +assert_type(random_st.choice(5, 3, replace=True), npt.NDArray[np.long]) +assert_type(random_st.choice(5, 3, p=[1 / 5] * 5), npt.NDArray[np.long]) +assert_type(random_st.choice(5, 3, p=[1 / 5] * 5, replace=False), npt.NDArray[np.long]) + +assert_type(random_st.choice(["pooh", "rabbit", "piglet", "Christopher"]), Any) +assert_type(random_st.choice(["pooh", "rabbit", "piglet", "Christopher"], 3), npt.NDArray[Any]) +assert_type(random_st.choice(["pooh", "rabbit", "piglet", "Christopher"], 3, p=[1 / 4] * 4), npt.NDArray[Any]) +assert_type(random_st.choice(["pooh", "rabbit", "piglet", "Christopher"], 3, replace=True), npt.NDArray[Any]) +assert_type(random_st.choice(["pooh", "rabbit", "piglet", "Christopher"], 3, replace=False, p=np.array([1 / 8, 1 / 8, 1 / 2, 1 / 4])), npt.NDArray[Any]) + +assert_type(random_st.dirichlet([0.5, 0.5]), npt.NDArray[np.float64]) +assert_type(random_st.dirichlet(np.array([0.5, 0.5])), npt.NDArray[np.float64]) +assert_type(random_st.dirichlet(np.array([0.5, 0.5]), size=3), npt.NDArray[np.float64]) + +assert_type(random_st.multinomial(20, [1 / 6.0] * 6), npt.NDArray[np.long]) +assert_type(random_st.multinomial(20, np.array([0.5, 0.5])), npt.NDArray[np.long]) +assert_type(random_st.multinomial(20, [1 / 6.0] * 6, size=2), npt.NDArray[np.long]) + +assert_type(random_st.multivariate_normal([0.0], [[1.0]]), npt.NDArray[np.float64]) +assert_type(random_st.multivariate_normal([0.0], np.array([[1.0]])), npt.NDArray[np.float64]) +assert_type(random_st.multivariate_normal(np.array([0.0]), [[1.0]]), npt.NDArray[np.float64]) +assert_type(random_st.multivariate_normal([0.0], np.array([[1.0]])), npt.NDArray[np.float64]) + +assert_type(random_st.permutation(10), npt.NDArray[np.long]) +assert_type(random_st.permutation([1, 2, 3, 4]), npt.NDArray[Any]) +assert_type(random_st.permutation(np.array([1, 2, 3, 4])), npt.NDArray[Any]) +assert_type(random_st.permutation(D_2D), npt.NDArray[Any]) + +assert_type(random_st.shuffle(np.arange(10)), None) +assert_type(random_st.shuffle([1, 2, 3, 4, 5]), None) +assert_type(random_st.shuffle(D_2D), None) + +assert_type(np.random.RandomState(pcg64), np.random.RandomState) +assert_type(np.random.RandomState(0), np.random.RandomState) +assert_type(np.random.RandomState([0, 1, 2]), np.random.RandomState) +assert_type(random_st.__str__(), str) +assert_type(random_st.__repr__(), str) +random_st_state = random_st.__getstate__() +assert_type(random_st_state, dict[str, Any]) +assert_type(random_st.__setstate__(random_st_state), None) +assert_type(random_st.seed(), None) +assert_type(random_st.seed(1), None) +assert_type(random_st.seed([0, 1]), None) +random_st_get_state = random_st.get_state() +assert_type(random_st_state, dict[str, Any]) +random_st_get_state_legacy = random_st.get_state(legacy=True) +assert_type(random_st_get_state_legacy, dict[str, Any] | tuple[str, npt.NDArray[np.uint32], int, int, float]) +assert_type(random_st.set_state(random_st_get_state), None) + +assert_type(random_st.rand(), float) +assert_type(random_st.rand(1), npt.NDArray[np.float64]) +assert_type(random_st.rand(1, 2), npt.NDArray[np.float64]) +assert_type(random_st.randn(), float) +assert_type(random_st.randn(1), npt.NDArray[np.float64]) +assert_type(random_st.randn(1, 2), npt.NDArray[np.float64]) +assert_type(random_st.random_sample(), float) +assert_type(random_st.random_sample(1), npt.NDArray[np.float64]) +assert_type(random_st.random_sample(size=(1, 2)), npt.NDArray[np.float64]) + +assert_type(random_st.tomaxint(), int) +assert_type(random_st.tomaxint(1), npt.NDArray[np.int64]) +assert_type(random_st.tomaxint((1,)), npt.NDArray[np.int64]) + +assert_type(np.random.mtrand.set_bit_generator(pcg64), None) +assert_type(np.random.mtrand.get_bit_generator(), np.random.BitGenerator) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/rec.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/rec.pyi new file mode 100644 index 0000000000000000000000000000000000000000..5a9bf54ae7dba40a90f89dc8bc52c7141f82eaf4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/rec.pyi @@ -0,0 +1,168 @@ +import io +from typing import Any + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +AR_i8: npt.NDArray[np.int64] +REC_AR_V: np.recarray[tuple[int, ...], np.dtype[np.record]] +AR_LIST: list[npt.NDArray[np.int64]] + +record: np.record +file_obj: io.BufferedIOBase + +assert_type(np.rec.format_parser( + formats=[np.float64, np.int64, np.bool], + names=["f8", "i8", "?"], + titles=None, + aligned=True, +), np.rec.format_parser) +assert_type(np.rec.format_parser.dtype, np.dtype[np.void]) + +assert_type(record.field_a, Any) +assert_type(record.field_b, Any) +assert_type(record["field_a"], Any) +assert_type(record["field_b"], Any) +assert_type(record.pprint(), str) +record.field_c = 5 + +assert_type(REC_AR_V.field(0), Any) +assert_type(REC_AR_V.field("field_a"), Any) +assert_type(REC_AR_V.field(0, AR_i8), None) +assert_type(REC_AR_V.field("field_a", AR_i8), None) +assert_type(REC_AR_V["field_a"], npt.NDArray[Any]) +assert_type(REC_AR_V.field_a, Any) +assert_type(REC_AR_V.__array_finalize__(object()), None) + +assert_type( + np.recarray( + shape=(10, 5), + formats=[np.float64, np.int64, np.bool], + order="K", + byteorder="|", + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.recarray( + shape=(10, 5), + dtype=[("f8", np.float64), ("i8", np.int64)], + strides=(5, 5), + ), + np.recarray[Any, np.dtype[Any]], +) + +assert_type(np.rec.fromarrays(AR_LIST), np.recarray[Any, np.dtype[Any]]) +assert_type( + np.rec.fromarrays(AR_LIST, dtype=np.int64), + np.recarray[Any, np.dtype[Any]], +) +assert_type( + np.rec.fromarrays( + AR_LIST, + formats=[np.int64, np.float64], + names=["i8", "f8"] + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.rec.fromrecords((1, 1.5)), + np.recarray[Any, np.dtype[np.record]] +) + +assert_type( + np.rec.fromrecords( + [(1, 1.5)], + dtype=[("i8", np.int64), ("f8", np.float64)], + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.rec.fromrecords( + REC_AR_V, + formats=[np.int64, np.float64], + names=["i8", "f8"] + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.rec.fromstring( + b"(1, 1.5)", + dtype=[("i8", np.int64), ("f8", np.float64)], + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.rec.fromstring( + REC_AR_V, + formats=[np.int64, np.float64], + names=["i8", "f8"] + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type(np.rec.fromfile( + "test_file.txt", + dtype=[("i8", np.int64), ("f8", np.float64)], +), np.recarray[Any, np.dtype[Any]]) + +assert_type( + np.rec.fromfile( + file_obj, + formats=[np.int64, np.float64], + names=["i8", "f8"] + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type(np.rec.array(AR_i8), np.recarray[Any, np.dtype[np.int64]]) + +assert_type( + np.rec.array([(1, 1.5)], dtype=[("i8", np.int64), ("f8", np.float64)]), + np.recarray[Any, np.dtype[Any]], +) + +assert_type( + np.rec.array( + [(1, 1.5)], + formats=[np.int64, np.float64], + names=["i8", "f8"] + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.rec.array( + None, + dtype=np.float64, + shape=(10, 3), + ), + np.recarray[Any, np.dtype[Any]], +) + +assert_type( + np.rec.array( + None, + formats=[np.int64, np.float64], + names=["i8", "f8"], + shape=(10, 3), + ), + np.recarray[Any, np.dtype[np.record]], +) + +assert_type( + np.rec.array(file_obj, dtype=np.float64), + np.recarray[Any, np.dtype[Any]], +) + +assert_type( + np.rec.array(file_obj, formats=[np.int64, np.float64], names=["i8", "f8"]), + np.recarray[Any, np.dtype[np.record]], +) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/scalars.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/scalars.pyi new file mode 100644 index 0000000000000000000000000000000000000000..365c1e7ad1f8ae6bd72e714a8106ac8bb870d277 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/scalars.pyi @@ -0,0 +1,193 @@ +from typing import Any, Literal, TypeAlias +from typing_extensions import Unpack, assert_type + +import numpy as np +import numpy.typing as npt + +_1: TypeAlias = Literal[1] + +b: np.bool +u8: np.uint64 +i8: np.int64 +f8: np.float64 +c8: np.complex64 +c16: np.complex128 +m: np.timedelta64 +U: np.str_ +S: np.bytes_ +V: np.void +O: np.object_ # cannot exists at runtime + +array_nd: np.ndarray[Any, Any] +array_0d: np.ndarray[tuple[()], Any] +array_2d_2x2: np.ndarray[tuple[Literal[2], Literal[2]], Any] + +assert_type(c8.real, np.float32) +assert_type(c8.imag, np.float32) + +assert_type(c8.real.real, np.float32) +assert_type(c8.real.imag, np.float32) + +assert_type(c8.itemsize, int) +assert_type(c8.shape, tuple[()]) +assert_type(c8.strides, tuple[()]) + +assert_type(c8.ndim, Literal[0]) +assert_type(c8.size, Literal[1]) + +assert_type(c8.squeeze(), np.complex64) +assert_type(c8.byteswap(), np.complex64) +assert_type(c8.transpose(), np.complex64) + +assert_type(c8.dtype, np.dtype[np.complex64]) + +assert_type(c8.real, np.float32) +assert_type(c16.imag, np.float64) + +assert_type(np.str_('foo'), np.str_) + +assert_type(V[0], Any) +assert_type(V["field1"], Any) +assert_type(V[["field1", "field2"]], np.void) +V[0] = 5 + +# Aliases +assert_type(np.bool_(), np.bool[Literal[False]]) +assert_type(np.byte(), np.byte) +assert_type(np.short(), np.short) +assert_type(np.intc(), np.intc) +assert_type(np.intp(), np.intp) +assert_type(np.int_(), np.int_) +assert_type(np.long(), np.long) +assert_type(np.longlong(), np.longlong) + +assert_type(np.ubyte(), np.ubyte) +assert_type(np.ushort(), np.ushort) +assert_type(np.uintc(), np.uintc) +assert_type(np.uintp(), np.uintp) +assert_type(np.uint(), np.uint) +assert_type(np.ulong(), np.ulong) +assert_type(np.ulonglong(), np.ulonglong) + +assert_type(np.half(), np.half) +assert_type(np.single(), np.single) +assert_type(np.double(), np.double) +assert_type(np.longdouble(), np.longdouble) + +assert_type(np.csingle(), np.csingle) +assert_type(np.cdouble(), np.cdouble) +assert_type(np.clongdouble(), np.clongdouble) + +assert_type(b.item(), bool) +assert_type(i8.item(), int) +assert_type(u8.item(), int) +assert_type(f8.item(), float) +assert_type(c16.item(), complex) +assert_type(U.item(), str) +assert_type(S.item(), bytes) + +assert_type(b.tolist(), bool) +assert_type(i8.tolist(), int) +assert_type(u8.tolist(), int) +assert_type(f8.tolist(), float) +assert_type(c16.tolist(), complex) +assert_type(U.tolist(), str) +assert_type(S.tolist(), bytes) + +assert_type(b.ravel(), np.ndarray[tuple[int], np.dtype[np.bool]]) +assert_type(i8.ravel(), np.ndarray[tuple[int], np.dtype[np.int64]]) +assert_type(u8.ravel(), np.ndarray[tuple[int], np.dtype[np.uint64]]) +assert_type(f8.ravel(), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(c16.ravel(), np.ndarray[tuple[int], np.dtype[np.complex128]]) +assert_type(U.ravel(), np.ndarray[tuple[int], np.dtype[np.str_]]) +assert_type(S.ravel(), np.ndarray[tuple[int], np.dtype[np.bytes_]]) + +assert_type(b.flatten(), np.ndarray[tuple[int], np.dtype[np.bool]]) +assert_type(i8.flatten(), np.ndarray[tuple[int], np.dtype[np.int64]]) +assert_type(u8.flatten(), np.ndarray[tuple[int], np.dtype[np.uint64]]) +assert_type(f8.flatten(), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(c16.flatten(), np.ndarray[tuple[int], np.dtype[np.complex128]]) +assert_type(U.flatten(), np.ndarray[tuple[int], np.dtype[np.str_]]) +assert_type(S.flatten(), np.ndarray[tuple[int], np.dtype[np.bytes_]]) + +assert_type(b.reshape(()), np.bool) +assert_type(i8.reshape([]), np.int64) +assert_type(b.reshape(1), np.ndarray[tuple[_1], np.dtype[np.bool]]) +assert_type(i8.reshape(-1), np.ndarray[tuple[_1], np.dtype[np.int64]]) +assert_type(u8.reshape(1, 1), np.ndarray[tuple[_1, _1], np.dtype[np.uint64]]) +assert_type(f8.reshape(1, -1), np.ndarray[tuple[_1, _1], np.dtype[np.float64]]) +assert_type(c16.reshape(1, 1, 1), np.ndarray[tuple[_1, _1, _1], np.dtype[np.complex128]]) +assert_type(U.reshape(1, 1, 1, 1), np.ndarray[tuple[_1, _1, _1, _1], np.dtype[np.str_]]) +assert_type( + S.reshape(1, 1, 1, 1, 1), + np.ndarray[ + # len(shape) >= 5 + tuple[_1, _1, _1, _1, _1, Unpack[tuple[_1, ...]]], + np.dtype[np.bytes_], + ], +) + +assert_type(i8.astype(float), Any) +assert_type(i8.astype(np.float64), np.float64) + +assert_type(i8.view(), np.int64) +assert_type(i8.view(np.float64), np.float64) +assert_type(i8.view(float), Any) +assert_type(i8.view(np.float64, np.ndarray), np.float64) + +assert_type(i8.getfield(float), Any) +assert_type(i8.getfield(np.float64), np.float64) +assert_type(i8.getfield(np.float64, 8), np.float64) + +assert_type(f8.as_integer_ratio(), tuple[int, int]) +assert_type(f8.is_integer(), bool) +assert_type(f8.__trunc__(), int) +assert_type(f8.__getformat__("float"), str) +assert_type(f8.hex(), str) +assert_type(np.float64.fromhex("0x0.0p+0"), np.float64) + +assert_type(f8.__getnewargs__(), tuple[float]) +assert_type(c16.__getnewargs__(), tuple[float, float]) + +assert_type(i8.numerator, np.int64) +assert_type(i8.denominator, Literal[1]) +assert_type(u8.numerator, np.uint64) +assert_type(u8.denominator, Literal[1]) +assert_type(m.numerator, np.timedelta64) +assert_type(m.denominator, Literal[1]) + +assert_type(round(i8), int) +assert_type(round(i8, 3), np.int64) +assert_type(round(u8), int) +assert_type(round(u8, 3), np.uint64) +assert_type(round(f8), int) +assert_type(round(f8, 3), np.float64) + +assert_type(f8.__ceil__(), int) +assert_type(f8.__floor__(), int) + +assert_type(i8.is_integer(), Literal[True]) + +assert_type(O.real, np.object_) +assert_type(O.imag, np.object_) +assert_type(int(O), int) +assert_type(float(O), float) +assert_type(complex(O), complex) + +# These fail fail because of a mypy __new__ bug: +# https://github.com/python/mypy/issues/15182 +# According to the typing spec, the following statements are valid, see +# https://typing.readthedocs.io/en/latest/spec/constructors.html#new-method + +# assert_type(np.object_(), None) +# assert_type(np.object_(None), None) +# assert_type(np.object_(array_nd), np.ndarray[Any, np.dtype[np.object_]]) +# assert_type(np.object_([]), npt.NDArray[np.object_]) +# assert_type(np.object_(()), npt.NDArray[np.object_]) +# assert_type(np.object_(range(4)), npt.NDArray[np.object_]) +# assert_type(np.object_(+42), int) +# assert_type(np.object_(1 / 137), float) +# assert_type(np.object_('Developers! ' * (1 << 6)), str) +# assert_type(np.object_(object()), object) +# assert_type(np.object_({False, True, NotADirectoryError}), set[Any]) +# assert_type(np.object_({'spam': 'food', 'ham': 'food'}), dict[str, str]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/shape.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/shape.pyi new file mode 100644 index 0000000000000000000000000000000000000000..d236479e51167addf1ab1fd2f6faa191160800ee --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/shape.pyi @@ -0,0 +1,15 @@ +from typing import Any, NamedTuple + +import numpy as np +from typing_extensions import assert_type + + +# Subtype of tuple[int, int] +class XYGrid(NamedTuple): + x_axis: int + y_axis: int + +arr: np.ndarray[XYGrid, Any] + +# Test shape property matches shape typevar +assert_type(arr.shape, XYGrid) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/shape_base.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/shape_base.pyi new file mode 100644 index 0000000000000000000000000000000000000000..baaa78b049e73f17d36c8deb908dafbe219d8f1f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/shape_base.pyi @@ -0,0 +1,54 @@ +from typing import Any + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +i8: np.int64 +f8: np.float64 + +AR_b: npt.NDArray[np.bool] +AR_i8: npt.NDArray[np.int64] +AR_f8: npt.NDArray[np.float64] + +AR_LIKE_f8: list[float] + +assert_type(np.take_along_axis(AR_f8, AR_i8, axis=1), npt.NDArray[np.float64]) +assert_type(np.take_along_axis(f8, AR_i8, axis=None), npt.NDArray[np.float64]) + +assert_type(np.put_along_axis(AR_f8, AR_i8, "1.0", axis=1), None) + +assert_type(np.expand_dims(AR_i8, 2), npt.NDArray[np.int64]) +assert_type(np.expand_dims(AR_LIKE_f8, 2), npt.NDArray[Any]) + +assert_type(np.column_stack([AR_i8]), npt.NDArray[np.int64]) +assert_type(np.column_stack([AR_LIKE_f8]), npt.NDArray[Any]) + +assert_type(np.dstack([AR_i8]), npt.NDArray[np.int64]) +assert_type(np.dstack([AR_LIKE_f8]), npt.NDArray[Any]) + +assert_type(np.array_split(AR_i8, [3, 5, 6, 10]), list[npt.NDArray[np.int64]]) +assert_type(np.array_split(AR_LIKE_f8, [3, 5, 6, 10]), list[npt.NDArray[Any]]) + +assert_type(np.split(AR_i8, [3, 5, 6, 10]), list[npt.NDArray[np.int64]]) +assert_type(np.split(AR_LIKE_f8, [3, 5, 6, 10]), list[npt.NDArray[Any]]) + +assert_type(np.hsplit(AR_i8, [3, 5, 6, 10]), list[npt.NDArray[np.int64]]) +assert_type(np.hsplit(AR_LIKE_f8, [3, 5, 6, 10]), list[npt.NDArray[Any]]) + +assert_type(np.vsplit(AR_i8, [3, 5, 6, 10]), list[npt.NDArray[np.int64]]) +assert_type(np.vsplit(AR_LIKE_f8, [3, 5, 6, 10]), list[npt.NDArray[Any]]) + +assert_type(np.dsplit(AR_i8, [3, 5, 6, 10]), list[npt.NDArray[np.int64]]) +assert_type(np.dsplit(AR_LIKE_f8, [3, 5, 6, 10]), list[npt.NDArray[Any]]) + +assert_type(np.kron(AR_b, AR_b), npt.NDArray[np.bool]) +assert_type(np.kron(AR_b, AR_i8), npt.NDArray[np.signedinteger[Any]]) +assert_type(np.kron(AR_f8, AR_f8), npt.NDArray[np.floating[Any]]) + +assert_type(np.tile(AR_i8, 5), npt.NDArray[np.int64]) +assert_type(np.tile(AR_LIKE_f8, [2, 2]), npt.NDArray[Any]) + +assert_type(np.unstack(AR_i8, axis=0), tuple[npt.NDArray[np.int64], ...]) +assert_type(np.unstack(AR_LIKE_f8, axis=0), tuple[npt.NDArray[Any], ...]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/stride_tricks.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/stride_tricks.pyi new file mode 100644 index 0000000000000000000000000000000000000000..fda149e5ec7b1fbafdc8bfece9801d987694a3de --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/stride_tricks.pyi @@ -0,0 +1,29 @@ +from typing import Any + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +AR_f8: npt.NDArray[np.float64] +AR_LIKE_f: list[float] +interface_dict: dict[str, Any] + +assert_type(np.lib.stride_tricks.as_strided(AR_f8), npt.NDArray[np.float64]) +assert_type(np.lib.stride_tricks.as_strided(AR_LIKE_f), npt.NDArray[Any]) +assert_type(np.lib.stride_tricks.as_strided(AR_f8, strides=(1, 5)), npt.NDArray[np.float64]) +assert_type(np.lib.stride_tricks.as_strided(AR_f8, shape=[9, 20]), npt.NDArray[np.float64]) + +assert_type(np.lib.stride_tricks.sliding_window_view(AR_f8, 5), npt.NDArray[np.float64]) +assert_type(np.lib.stride_tricks.sliding_window_view(AR_LIKE_f, (1, 5)), npt.NDArray[Any]) +assert_type(np.lib.stride_tricks.sliding_window_view(AR_f8, [9], axis=1), npt.NDArray[np.float64]) + +assert_type(np.broadcast_to(AR_f8, 5), npt.NDArray[np.float64]) +assert_type(np.broadcast_to(AR_LIKE_f, (1, 5)), npt.NDArray[Any]) +assert_type(np.broadcast_to(AR_f8, [4, 6], subok=True), npt.NDArray[np.float64]) + +assert_type(np.broadcast_shapes((1, 2), [3, 1], (3, 2)), tuple[int, ...]) +assert_type(np.broadcast_shapes((6, 7), (5, 6, 1), 7, (5, 1, 7)), tuple[int, ...]) + +assert_type(np.broadcast_arrays(AR_f8, AR_f8), tuple[npt.NDArray[Any], ...]) +assert_type(np.broadcast_arrays(AR_f8, AR_LIKE_f), tuple[npt.NDArray[Any], ...]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/strings.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/strings.pyi new file mode 100644 index 0000000000000000000000000000000000000000..d9e812a0bef6477b073a5a65af87096abe2360df --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/strings.pyi @@ -0,0 +1,193 @@ +import numpy as np +import numpy.typing as npt +import numpy._typing as np_t + +from typing_extensions import assert_type +from typing import TypeAlias + +AR_U: npt.NDArray[np.str_] +AR_S: npt.NDArray[np.bytes_] +AR_T: np.ndarray[np_t._Shape, np.dtypes.StringDType] + +AR_T_alias: TypeAlias = np.ndarray[np_t._Shape, np.dtypes.StringDType] +AR_TU_alias: TypeAlias = AR_T_alias | npt.NDArray[np.str_] + +assert_type(np.strings.equal(AR_U, AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.equal(AR_S, AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.equal(AR_T, AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.not_equal(AR_U, AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.not_equal(AR_S, AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.not_equal(AR_T, AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.greater_equal(AR_U, AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.greater_equal(AR_S, AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.greater_equal(AR_T, AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.less_equal(AR_U, AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.less_equal(AR_S, AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.less_equal(AR_T, AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.greater(AR_U, AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.greater(AR_S, AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.greater(AR_T, AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.less(AR_U, AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.less(AR_S, AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.less(AR_T, AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.add(AR_U, AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.add(AR_S, AR_S), npt.NDArray[np.bytes_]) +assert_type(np.strings.add(AR_T, AR_T), AR_T_alias) + +assert_type(np.strings.multiply(AR_U, 5), npt.NDArray[np.str_]) +assert_type(np.strings.multiply(AR_S, [5, 4, 3]), npt.NDArray[np.bytes_]) +assert_type(np.strings.multiply(AR_T, 5), AR_T_alias) + +assert_type(np.strings.mod(AR_U, "test"), npt.NDArray[np.str_]) +assert_type(np.strings.mod(AR_S, "test"), npt.NDArray[np.bytes_]) +assert_type(np.strings.mod(AR_T, "test"), AR_T_alias) + +assert_type(np.strings.capitalize(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.capitalize(AR_S), npt.NDArray[np.bytes_]) +assert_type(np.strings.capitalize(AR_T), AR_T_alias) + +assert_type(np.strings.center(AR_U, 5), npt.NDArray[np.str_]) +assert_type(np.strings.center(AR_S, [2, 3, 4], b"a"), npt.NDArray[np.bytes_]) +assert_type(np.strings.center(AR_T, 5), AR_T_alias) + +assert_type(np.strings.encode(AR_U), npt.NDArray[np.bytes_]) +assert_type(np.strings.encode(AR_T), npt.NDArray[np.bytes_]) +assert_type(np.strings.decode(AR_S), npt.NDArray[np.str_]) + +assert_type(np.strings.expandtabs(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.expandtabs(AR_S, tabsize=4), npt.NDArray[np.bytes_]) +assert_type(np.strings.expandtabs(AR_T), AR_T_alias) + +assert_type(np.strings.ljust(AR_U, 5), npt.NDArray[np.str_]) +assert_type(np.strings.ljust(AR_S, [4, 3, 1], fillchar=[b"a", b"b", b"c"]), npt.NDArray[np.bytes_]) +assert_type(np.strings.ljust(AR_T, 5), AR_T_alias) +assert_type(np.strings.ljust(AR_T, [4, 2, 1], fillchar=["a", "b", "c"]), AR_T_alias) + +assert_type(np.strings.rjust(AR_U, 5), npt.NDArray[np.str_]) +assert_type(np.strings.rjust(AR_S, [4, 3, 1], fillchar=[b"a", b"b", b"c"]), npt.NDArray[np.bytes_]) +assert_type(np.strings.rjust(AR_T, 5), AR_T_alias) +assert_type(np.strings.rjust(AR_T, [4, 2, 1], fillchar=["a", "b", "c"]), AR_T_alias) + +assert_type(np.strings.lstrip(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.lstrip(AR_S, b"_"), npt.NDArray[np.bytes_]) +assert_type(np.strings.lstrip(AR_T), AR_T_alias) +assert_type(np.strings.lstrip(AR_T, "_"), AR_T_alias) + +assert_type(np.strings.rstrip(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.rstrip(AR_S, b"_"), npt.NDArray[np.bytes_]) +assert_type(np.strings.rstrip(AR_T), AR_T_alias) +assert_type(np.strings.rstrip(AR_T, "_"), AR_T_alias) + +assert_type(np.strings.strip(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.strip(AR_S, b"_"), npt.NDArray[np.bytes_]) +assert_type(np.strings.strip(AR_T), AR_T_alias) +assert_type(np.strings.strip(AR_T, "_"), AR_T_alias) + +assert_type(np.strings.count(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) +assert_type(np.strings.count(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.int_]) +assert_type(np.strings.count(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) +assert_type(np.strings.count(AR_T, ["a", "b", "c"], end=9), npt.NDArray[np.int_]) + +assert_type(np.strings.partition(AR_U, "\n"), npt.NDArray[np.str_]) +assert_type(np.strings.partition(AR_S, [b"a", b"b", b"c"]), npt.NDArray[np.bytes_]) +assert_type(np.strings.partition(AR_T, "\n"), AR_TU_alias) + +assert_type(np.strings.rpartition(AR_U, "\n"), npt.NDArray[np.str_]) +assert_type(np.strings.rpartition(AR_S, [b"a", b"b", b"c"]), npt.NDArray[np.bytes_]) +assert_type(np.strings.rpartition(AR_T, "\n"), AR_TU_alias) + +assert_type(np.strings.replace(AR_U, "_", "-"), npt.NDArray[np.str_]) +assert_type(np.strings.replace(AR_S, [b"_", b""], [b"a", b"b"]), npt.NDArray[np.bytes_]) +assert_type(np.strings.replace(AR_T, "_", "_"), AR_TU_alias) + +assert_type(np.strings.lower(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.lower(AR_S), npt.NDArray[np.bytes_]) +assert_type(np.strings.lower(AR_T), AR_T_alias) + +assert_type(np.strings.upper(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.upper(AR_S), npt.NDArray[np.bytes_]) +assert_type(np.strings.upper(AR_T), AR_T_alias) + +assert_type(np.strings.swapcase(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.swapcase(AR_S), npt.NDArray[np.bytes_]) +assert_type(np.strings.swapcase(AR_T), AR_T_alias) + +assert_type(np.strings.title(AR_U), npt.NDArray[np.str_]) +assert_type(np.strings.title(AR_S), npt.NDArray[np.bytes_]) +assert_type(np.strings.title(AR_T), AR_T_alias) + +assert_type(np.strings.zfill(AR_U, 5), npt.NDArray[np.str_]) +assert_type(np.strings.zfill(AR_S, [2, 3, 4]), npt.NDArray[np.bytes_]) +assert_type(np.strings.zfill(AR_T, 5), AR_T_alias) + +assert_type(np.strings.endswith(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.bool]) +assert_type(np.strings.endswith(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.bool]) +assert_type(np.strings.endswith(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.bool]) + +assert_type(np.strings.startswith(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.bool]) +assert_type(np.strings.startswith(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.bool]) +assert_type(np.strings.startswith(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.bool]) + +assert_type(np.strings.find(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) +assert_type(np.strings.find(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.int_]) +assert_type(np.strings.find(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) + +assert_type(np.strings.rfind(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) +assert_type(np.strings.rfind(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.int_]) +assert_type(np.strings.rfind(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) + +assert_type(np.strings.index(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) +assert_type(np.strings.index(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.int_]) +assert_type(np.strings.index(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) + +assert_type(np.strings.rindex(AR_U, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) +assert_type(np.strings.rindex(AR_S, [b"a", b"b", b"c"], end=9), npt.NDArray[np.int_]) +assert_type(np.strings.rindex(AR_T, "a", start=[1, 2, 3]), npt.NDArray[np.int_]) + +assert_type(np.strings.isalpha(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isalpha(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.isalpha(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.isalnum(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isalnum(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.isalnum(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.isdecimal(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isdecimal(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.isdigit(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isdigit(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.isdigit(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.islower(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.islower(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.islower(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.isnumeric(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isnumeric(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.isspace(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isspace(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.isspace(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.istitle(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.istitle(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.istitle(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.isupper(AR_U), npt.NDArray[np.bool]) +assert_type(np.strings.isupper(AR_S), npt.NDArray[np.bool]) +assert_type(np.strings.isupper(AR_T), npt.NDArray[np.bool]) + +assert_type(np.strings.str_len(AR_U), npt.NDArray[np.int_]) +assert_type(np.strings.str_len(AR_S), npt.NDArray[np.int_]) +assert_type(np.strings.str_len(AR_T), npt.NDArray[np.int_]) + +assert_type(np.strings.translate(AR_U, ""), npt.NDArray[np.str_]) +assert_type(np.strings.translate(AR_S, ""), npt.NDArray[np.bytes_]) +assert_type(np.strings.translate(AR_T, ""), AR_T_alias) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/testing.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/testing.pyi new file mode 100644 index 0000000000000000000000000000000000000000..14fd57d79c938c8fd5fe8426b02160bf0f766c84 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/testing.pyi @@ -0,0 +1,200 @@ +import re +import sys +import warnings +import types +import unittest +import contextlib +from collections.abc import Callable +from typing import Any, TypeVar +from pathlib import Path + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +AR_f8: npt.NDArray[np.float64] +AR_i8: npt.NDArray[np.int64] + +bool_obj: bool +suppress_obj: np.testing.suppress_warnings +FT = TypeVar("FT", bound=Callable[..., Any]) + +def func() -> int: ... + +def func2( + x: npt.NDArray[np.number[Any]], + y: npt.NDArray[np.number[Any]], +) -> npt.NDArray[np.bool]: ... + +assert_type(np.testing.KnownFailureException(), np.testing.KnownFailureException) +assert_type(np.testing.IgnoreException(), np.testing.IgnoreException) + +assert_type( + np.testing.clear_and_catch_warnings(modules=[np.testing]), + np.testing.clear_and_catch_warnings[None], +) +assert_type( + np.testing.clear_and_catch_warnings(True), + np.testing.clear_and_catch_warnings[list[warnings.WarningMessage]], +) +assert_type( + np.testing.clear_and_catch_warnings(False), + np.testing.clear_and_catch_warnings[None], +) +assert_type( + np.testing.clear_and_catch_warnings(bool_obj), + np.testing.clear_and_catch_warnings, +) +assert_type( + np.testing.clear_and_catch_warnings.class_modules, + tuple[types.ModuleType, ...], +) +assert_type( + np.testing.clear_and_catch_warnings.modules, + set[types.ModuleType], +) + +with np.testing.clear_and_catch_warnings(True) as c1: + assert_type(c1, list[warnings.WarningMessage]) +with np.testing.clear_and_catch_warnings() as c2: + assert_type(c2, None) + +assert_type(np.testing.suppress_warnings("once"), np.testing.suppress_warnings) +assert_type(np.testing.suppress_warnings()(func), Callable[[], int]) +assert_type(suppress_obj.filter(RuntimeWarning), None) +assert_type(suppress_obj.record(RuntimeWarning), list[warnings.WarningMessage]) +with suppress_obj as c3: + assert_type(c3, np.testing.suppress_warnings) + +assert_type(np.testing.verbose, int) +assert_type(np.testing.IS_PYPY, bool) +assert_type(np.testing.HAS_REFCOUNT, bool) +assert_type(np.testing.HAS_LAPACK64, bool) + +assert_type(np.testing.assert_(1, msg="test"), None) +assert_type(np.testing.assert_(2, msg=lambda: "test"), None) + +if sys.platform == "win32" or sys.platform == "cygwin": + assert_type(np.testing.memusage(), int) +elif sys.platform == "linux": + assert_type(np.testing.memusage(), None | int) + +assert_type(np.testing.jiffies(), int) + +assert_type(np.testing.build_err_msg([0, 1, 2], "test"), str) +assert_type(np.testing.build_err_msg(range(2), "test", header="header"), str) +assert_type(np.testing.build_err_msg(np.arange(9).reshape(3, 3), "test", verbose=False), str) +assert_type(np.testing.build_err_msg("abc", "test", names=["x", "y"]), str) +assert_type(np.testing.build_err_msg([1.0, 2.0], "test", precision=5), str) + +assert_type(np.testing.assert_equal({1}, {1}), None) +assert_type(np.testing.assert_equal([1, 2, 3], [1, 2, 3], err_msg="fail"), None) +assert_type(np.testing.assert_equal(1, 1.0, verbose=True), None) + +assert_type(np.testing.print_assert_equal('Test XYZ of func xyz', [0, 1], [0, 1]), None) + +assert_type(np.testing.assert_almost_equal(1.0, 1.1), None) +assert_type(np.testing.assert_almost_equal([1, 2, 3], [1, 2, 3], err_msg="fail"), None) +assert_type(np.testing.assert_almost_equal(1, 1.0, verbose=True), None) +assert_type(np.testing.assert_almost_equal(1, 1.0001, decimal=2), None) + +assert_type(np.testing.assert_approx_equal(1.0, 1.1), None) +assert_type(np.testing.assert_approx_equal("1", "2", err_msg="fail"), None) +assert_type(np.testing.assert_approx_equal(1, 1.0, verbose=True), None) +assert_type(np.testing.assert_approx_equal(1, 1.0001, significant=2), None) + +assert_type(np.testing.assert_array_compare(func2, AR_i8, AR_f8, err_msg="test"), None) +assert_type(np.testing.assert_array_compare(func2, AR_i8, AR_f8, verbose=True), None) +assert_type(np.testing.assert_array_compare(func2, AR_i8, AR_f8, header="header"), None) +assert_type(np.testing.assert_array_compare(func2, AR_i8, AR_f8, precision=np.int64()), None) +assert_type(np.testing.assert_array_compare(func2, AR_i8, AR_f8, equal_nan=False), None) +assert_type(np.testing.assert_array_compare(func2, AR_i8, AR_f8, equal_inf=True), None) + +assert_type(np.testing.assert_array_equal(AR_i8, AR_f8), None) +assert_type(np.testing.assert_array_equal(AR_i8, AR_f8, err_msg="test"), None) +assert_type(np.testing.assert_array_equal(AR_i8, AR_f8, verbose=True), None) + +assert_type(np.testing.assert_array_almost_equal(AR_i8, AR_f8), None) +assert_type(np.testing.assert_array_almost_equal(AR_i8, AR_f8, err_msg="test"), None) +assert_type(np.testing.assert_array_almost_equal(AR_i8, AR_f8, verbose=True), None) +assert_type(np.testing.assert_array_almost_equal(AR_i8, AR_f8, decimal=1), None) + +assert_type(np.testing.assert_array_less(AR_i8, AR_f8), None) +assert_type(np.testing.assert_array_less(AR_i8, AR_f8, err_msg="test"), None) +assert_type(np.testing.assert_array_less(AR_i8, AR_f8, verbose=True), None) + +assert_type(np.testing.runstring("1 + 1", {}), Any) +assert_type(np.testing.runstring("int64() + 1", {"int64": np.int64}), Any) + +assert_type(np.testing.assert_string_equal("1", "1"), None) + +assert_type(np.testing.rundocs(), None) +assert_type(np.testing.rundocs("test.py"), None) +assert_type(np.testing.rundocs(Path("test.py"), raise_on_error=True), None) + +def func3(a: int) -> bool: ... + +assert_type( + np.testing.assert_raises(RuntimeWarning), + unittest.case._AssertRaisesContext[RuntimeWarning], +) +assert_type(np.testing.assert_raises(RuntimeWarning, func3, 5), None) + +assert_type( + np.testing.assert_raises_regex(RuntimeWarning, r"test"), + unittest.case._AssertRaisesContext[RuntimeWarning], +) +assert_type(np.testing.assert_raises_regex(RuntimeWarning, b"test", func3, 5), None) +assert_type(np.testing.assert_raises_regex(RuntimeWarning, re.compile(b"test"), func3, 5), None) + +class Test: ... + +def decorate(a: FT) -> FT: + return a + +assert_type(np.testing.decorate_methods(Test, decorate), None) +assert_type(np.testing.decorate_methods(Test, decorate, None), None) +assert_type(np.testing.decorate_methods(Test, decorate, "test"), None) +assert_type(np.testing.decorate_methods(Test, decorate, b"test"), None) +assert_type(np.testing.decorate_methods(Test, decorate, re.compile("test")), None) + +assert_type(np.testing.measure("for i in range(1000): np.sqrt(i**2)"), float) +assert_type(np.testing.measure(b"for i in range(1000): np.sqrt(i**2)", times=5), float) + +assert_type(np.testing.assert_allclose(AR_i8, AR_f8), None) +assert_type(np.testing.assert_allclose(AR_i8, AR_f8, rtol=0.005), None) +assert_type(np.testing.assert_allclose(AR_i8, AR_f8, atol=1), None) +assert_type(np.testing.assert_allclose(AR_i8, AR_f8, equal_nan=True), None) +assert_type(np.testing.assert_allclose(AR_i8, AR_f8, err_msg="err"), None) +assert_type(np.testing.assert_allclose(AR_i8, AR_f8, verbose=False), None) + +assert_type(np.testing.assert_array_almost_equal_nulp(AR_i8, AR_f8, nulp=2), None) + +assert_type(np.testing.assert_array_max_ulp(AR_i8, AR_f8, maxulp=2), npt.NDArray[Any]) +assert_type(np.testing.assert_array_max_ulp(AR_i8, AR_f8, dtype=np.float32), npt.NDArray[Any]) + +assert_type(np.testing.assert_warns(RuntimeWarning), contextlib._GeneratorContextManager[None]) +assert_type(np.testing.assert_warns(RuntimeWarning, func3, 5), bool) + +def func4(a: int, b: str) -> bool: ... + +assert_type(np.testing.assert_no_warnings(), contextlib._GeneratorContextManager[None]) +assert_type(np.testing.assert_no_warnings(func3, 5), bool) +assert_type(np.testing.assert_no_warnings(func4, a=1, b="test"), bool) +assert_type(np.testing.assert_no_warnings(func4, 1, "test"), bool) + +assert_type(np.testing.tempdir("test_dir"), contextlib._GeneratorContextManager[str]) +assert_type(np.testing.tempdir(prefix=b"test"), contextlib._GeneratorContextManager[bytes]) +assert_type(np.testing.tempdir("test_dir", dir=Path("here")), contextlib._GeneratorContextManager[str]) + +assert_type(np.testing.temppath("test_dir", text=True), contextlib._GeneratorContextManager[str]) +assert_type(np.testing.temppath(prefix=b"test"), contextlib._GeneratorContextManager[bytes]) +assert_type(np.testing.temppath("test_dir", dir=Path("here")), contextlib._GeneratorContextManager[str]) + +assert_type(np.testing.assert_no_gc_cycles(), contextlib._GeneratorContextManager[None]) +assert_type(np.testing.assert_no_gc_cycles(func3, 5), None) + +assert_type(np.testing.break_cycles(), None) + +assert_type(np.testing.TestCase(), unittest.case.TestCase) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/twodim_base.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/twodim_base.pyi new file mode 100644 index 0000000000000000000000000000000000000000..ac9b7067d5610d96499cf4cd17ce5045a12ec7ee --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/twodim_base.pyi @@ -0,0 +1,152 @@ +from typing import Any, TypeVar + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +_SCT = TypeVar("_SCT", bound=np.generic) + + +def func1(ar: npt.NDArray[_SCT], a: int) -> npt.NDArray[_SCT]: + pass + + +def func2(ar: npt.NDArray[np.number[Any]], a: str) -> npt.NDArray[np.float64]: + pass + + +AR_b: npt.NDArray[np.bool] +AR_u: npt.NDArray[np.uint64] +AR_i: npt.NDArray[np.int64] +AR_f: npt.NDArray[np.float64] +AR_c: npt.NDArray[np.complex128] +AR_O: npt.NDArray[np.object_] + +AR_LIKE_b: list[bool] +AR_LIKE_c: list[complex] + +assert_type(np.fliplr(AR_b), npt.NDArray[np.bool]) +assert_type(np.fliplr(AR_LIKE_b), npt.NDArray[Any]) + +assert_type(np.flipud(AR_b), npt.NDArray[np.bool]) +assert_type(np.flipud(AR_LIKE_b), npt.NDArray[Any]) + +assert_type(np.eye(10), npt.NDArray[np.float64]) +assert_type(np.eye(10, M=20, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(np.eye(10, k=2, dtype=int), npt.NDArray[Any]) + +assert_type(np.diag(AR_b), npt.NDArray[np.bool]) +assert_type(np.diag(AR_LIKE_b, k=0), npt.NDArray[Any]) + +assert_type(np.diagflat(AR_b), npt.NDArray[np.bool]) +assert_type(np.diagflat(AR_LIKE_b, k=0), npt.NDArray[Any]) + +assert_type(np.tri(10), npt.NDArray[np.float64]) +assert_type(np.tri(10, M=20, dtype=np.int64), npt.NDArray[np.int64]) +assert_type(np.tri(10, k=2, dtype=int), npt.NDArray[Any]) + +assert_type(np.tril(AR_b), npt.NDArray[np.bool]) +assert_type(np.tril(AR_LIKE_b, k=0), npt.NDArray[Any]) + +assert_type(np.triu(AR_b), npt.NDArray[np.bool]) +assert_type(np.triu(AR_LIKE_b, k=0), npt.NDArray[Any]) + +assert_type(np.vander(AR_b), npt.NDArray[np.signedinteger[Any]]) +assert_type(np.vander(AR_u), npt.NDArray[np.signedinteger[Any]]) +assert_type(np.vander(AR_i, N=2), npt.NDArray[np.signedinteger[Any]]) +assert_type(np.vander(AR_f, increasing=True), npt.NDArray[np.floating[Any]]) +assert_type(np.vander(AR_c), npt.NDArray[np.complexfloating[Any, Any]]) +assert_type(np.vander(AR_O), npt.NDArray[np.object_]) + +assert_type( + np.histogram2d(AR_LIKE_c, AR_LIKE_c), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.complex128 | np.float64], + npt.NDArray[np.complex128 | np.float64], + ], +) +assert_type( + np.histogram2d(AR_i, AR_b), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.float64], + npt.NDArray[np.float64], + ], +) +assert_type( + np.histogram2d(AR_f, AR_i), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.float64], + npt.NDArray[np.float64], + ], +) +assert_type( + np.histogram2d(AR_i, AR_f), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.float64], + npt.NDArray[np.float64], + ], +) +assert_type( + np.histogram2d(AR_f, AR_c, weights=AR_LIKE_b), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.complex128], + npt.NDArray[np.complex128], + ], +) +assert_type( + np.histogram2d(AR_f, AR_c, bins=8), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.complex128], + npt.NDArray[np.complex128], + ], +) +assert_type( + np.histogram2d(AR_c, AR_f, bins=(8, 5)), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.complex128], + npt.NDArray[np.complex128], + ], +) +assert_type( + np.histogram2d(AR_c, AR_i, bins=AR_u), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.uint64], + npt.NDArray[np.uint64], + ], +) +assert_type( + np.histogram2d(AR_c, AR_c, bins=(AR_u, AR_u)), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.uint64], + npt.NDArray[np.uint64], + ], +) +assert_type( + np.histogram2d(AR_c, AR_c, bins=(AR_b, 8)), + tuple[ + npt.NDArray[np.float64], + npt.NDArray[np.bool | np.complex128], + npt.NDArray[np.bool | np.complex128], + ], +) + +assert_type(np.mask_indices(10, func1), tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]) +assert_type(np.mask_indices(8, func2, "0"), tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]) + +assert_type(np.tril_indices(10), tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]) + +assert_type(np.tril_indices_from(AR_b), tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]) + +assert_type(np.triu_indices(10), tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]) + +assert_type(np.triu_indices_from(AR_b), tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/type_check.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/type_check.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3212ca1af95f0fccf086ae3cf794f5215060abf9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/type_check.pyi @@ -0,0 +1,79 @@ +from typing import Any, Literal + +import numpy as np +import numpy.typing as npt +from numpy._typing import _16Bit, _32Bit, _64Bit, _128Bit + +from typing_extensions import assert_type + +f8: np.float64 +f: float + +# NOTE: Avoid importing the platform specific `np.float128` type +AR_i8: npt.NDArray[np.int64] +AR_i4: npt.NDArray[np.int32] +AR_f2: npt.NDArray[np.float16] +AR_f8: npt.NDArray[np.float64] +AR_f16: npt.NDArray[np.floating[_128Bit]] +AR_c8: npt.NDArray[np.complex64] +AR_c16: npt.NDArray[np.complex128] + +AR_LIKE_f: list[float] + +class ComplexObj: + real: slice + imag: slice + +assert_type(np.mintypecode(["f8"], typeset="qfQF"), str) + +assert_type(np.real(ComplexObj()), slice) +assert_type(np.real(AR_f8), npt.NDArray[np.float64]) +assert_type(np.real(AR_c16), npt.NDArray[np.float64]) +assert_type(np.real(AR_LIKE_f), npt.NDArray[Any]) + +assert_type(np.imag(ComplexObj()), slice) +assert_type(np.imag(AR_f8), npt.NDArray[np.float64]) +assert_type(np.imag(AR_c16), npt.NDArray[np.float64]) +assert_type(np.imag(AR_LIKE_f), npt.NDArray[Any]) + +assert_type(np.iscomplex(f8), np.bool) +assert_type(np.iscomplex(AR_f8), npt.NDArray[np.bool]) +assert_type(np.iscomplex(AR_LIKE_f), npt.NDArray[np.bool]) + +assert_type(np.isreal(f8), np.bool) +assert_type(np.isreal(AR_f8), npt.NDArray[np.bool]) +assert_type(np.isreal(AR_LIKE_f), npt.NDArray[np.bool]) + +assert_type(np.iscomplexobj(f8), bool) +assert_type(np.isrealobj(f8), bool) + +assert_type(np.nan_to_num(f8), np.float64) +assert_type(np.nan_to_num(f, copy=True), Any) +assert_type(np.nan_to_num(AR_f8, nan=1.5), npt.NDArray[np.float64]) +assert_type(np.nan_to_num(AR_LIKE_f, posinf=9999), npt.NDArray[Any]) + +assert_type(np.real_if_close(AR_f8), npt.NDArray[np.float64]) +assert_type( + np.real_if_close(AR_c16), + npt.NDArray[np.floating[_64Bit]] | npt.NDArray[np.complexfloating[_64Bit, _64Bit]], +) +assert_type(np.real_if_close(AR_c8), npt.NDArray[np.float32] | npt.NDArray[np.complex64]) +assert_type(np.real_if_close(AR_LIKE_f), npt.NDArray[Any]) + +assert_type(np.typename("h"), Literal["short"]) +assert_type(np.typename("B"), Literal["unsigned char"]) +assert_type(np.typename("V"), Literal["void"]) +assert_type(np.typename("S1"), Literal["character"]) + +assert_type(np.common_type(AR_i4), type[np.floating[_64Bit]]) +assert_type(np.common_type(AR_f2), type[np.float16]) +assert_type(np.common_type(AR_f2, AR_i4), type[np.floating[_16Bit | _64Bit]]) +assert_type(np.common_type(AR_f16, AR_i4), type[np.floating[_64Bit | _128Bit]]) +assert_type( + np.common_type(AR_c8, AR_f2), + type[np.complexfloating[_16Bit | _32Bit, _16Bit | _32Bit]], +) +assert_type( + np.common_type(AR_f2, AR_c8, AR_i4), + type[np.complexfloating[_16Bit | _32Bit | _64Bit, _16Bit | _32Bit | _64Bit]], +) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufunc_config.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufunc_config.pyi new file mode 100644 index 0000000000000000000000000000000000000000..236318163615b1f663e73798cd3eec8f6a7dc42d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufunc_config.pyi @@ -0,0 +1,31 @@ +"""Typing tests for `_core._ufunc_config`.""" + +from _typeshed import SupportsWrite +from typing import Any +from collections.abc import Callable + +import numpy as np + +from typing_extensions import assert_type + +def func(a: str, b: int) -> None: ... + +class Write: + def write(self, value: str) -> None: ... + +assert_type(np.seterr(all=None), np._core._ufunc_config._ErrDict) +assert_type(np.seterr(divide="ignore"), np._core._ufunc_config._ErrDict) +assert_type(np.seterr(over="warn"), np._core._ufunc_config._ErrDict) +assert_type(np.seterr(under="call"), np._core._ufunc_config._ErrDict) +assert_type(np.seterr(invalid="raise"), np._core._ufunc_config._ErrDict) +assert_type(np.geterr(), np._core._ufunc_config._ErrDict) + +assert_type(np.setbufsize(4096), int) +assert_type(np.getbufsize(), int) + +assert_type(np.seterrcall(func), Callable[[str, int], Any] | None | SupportsWrite[str]) +assert_type(np.seterrcall(Write()), Callable[[str, int], Any] | None | SupportsWrite[str]) +assert_type(np.geterrcall(), Callable[[str, int], Any] | None | SupportsWrite[str]) + +assert_type(np.errstate(call=func, all="call"), np.errstate) +assert_type(np.errstate(call=Write(), divide="log", over="log"), np.errstate) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufunclike.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufunclike.pyi new file mode 100644 index 0000000000000000000000000000000000000000..8fbfa2287a1bff879cd67ebe40380f81b119608f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufunclike.pyi @@ -0,0 +1,33 @@ +from typing import Any + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +AR_LIKE_b: list[bool] +AR_LIKE_u: list[np.uint32] +AR_LIKE_i: list[int] +AR_LIKE_f: list[float] +AR_LIKE_O: list[np.object_] + +AR_U: npt.NDArray[np.str_] + +assert_type(np.fix(AR_LIKE_b), npt.NDArray[np.floating[Any]]) +assert_type(np.fix(AR_LIKE_u), npt.NDArray[np.floating[Any]]) +assert_type(np.fix(AR_LIKE_i), npt.NDArray[np.floating[Any]]) +assert_type(np.fix(AR_LIKE_f), npt.NDArray[np.floating[Any]]) +assert_type(np.fix(AR_LIKE_O), npt.NDArray[np.object_]) +assert_type(np.fix(AR_LIKE_f, out=AR_U), npt.NDArray[np.str_]) + +assert_type(np.isposinf(AR_LIKE_b), npt.NDArray[np.bool]) +assert_type(np.isposinf(AR_LIKE_u), npt.NDArray[np.bool]) +assert_type(np.isposinf(AR_LIKE_i), npt.NDArray[np.bool]) +assert_type(np.isposinf(AR_LIKE_f), npt.NDArray[np.bool]) +assert_type(np.isposinf(AR_LIKE_f, out=AR_U), npt.NDArray[np.str_]) + +assert_type(np.isneginf(AR_LIKE_b), npt.NDArray[np.bool]) +assert_type(np.isneginf(AR_LIKE_u), npt.NDArray[np.bool]) +assert_type(np.isneginf(AR_LIKE_i), npt.NDArray[np.bool]) +assert_type(np.isneginf(AR_LIKE_f), npt.NDArray[np.bool]) +assert_type(np.isneginf(AR_LIKE_f, out=AR_U), npt.NDArray[np.str_]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufuncs.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufuncs.pyi new file mode 100644 index 0000000000000000000000000000000000000000..77a96d0f528525286c361043b14a7a3c6277437f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/ufuncs.pyi @@ -0,0 +1,125 @@ +from typing import Literal, Any, NoReturn + +import numpy as np +import numpy.typing as npt + +from typing_extensions import assert_type + +i8: np.int64 +f8: np.float64 +AR_f8: npt.NDArray[np.float64] +AR_i8: npt.NDArray[np.int64] + +assert_type(np.absolute.__doc__, str) +assert_type(np.absolute.types, list[str]) + +assert_type(np.absolute.__name__, Literal["absolute"]) +assert_type(np.absolute.__qualname__, Literal["absolute"]) +assert_type(np.absolute.ntypes, Literal[20]) +assert_type(np.absolute.identity, None) +assert_type(np.absolute.nin, Literal[1]) +assert_type(np.absolute.nin, Literal[1]) +assert_type(np.absolute.nout, Literal[1]) +assert_type(np.absolute.nargs, Literal[2]) +assert_type(np.absolute.signature, None) +assert_type(np.absolute(f8), Any) +assert_type(np.absolute(AR_f8), npt.NDArray[Any]) +assert_type(np.absolute.at(AR_f8, AR_i8), None) + +assert_type(np.add.__name__, Literal["add"]) +assert_type(np.add.__qualname__, Literal["add"]) +assert_type(np.add.ntypes, Literal[22]) +assert_type(np.add.identity, Literal[0]) +assert_type(np.add.nin, Literal[2]) +assert_type(np.add.nout, Literal[1]) +assert_type(np.add.nargs, Literal[3]) +assert_type(np.add.signature, None) +assert_type(np.add(f8, f8), Any) +assert_type(np.add(AR_f8, f8), npt.NDArray[Any]) +assert_type(np.add.at(AR_f8, AR_i8, f8), None) +assert_type(np.add.reduce(AR_f8, axis=0), Any) +assert_type(np.add.accumulate(AR_f8), npt.NDArray[Any]) +assert_type(np.add.reduceat(AR_f8, AR_i8), npt.NDArray[Any]) +assert_type(np.add.outer(f8, f8), Any) +assert_type(np.add.outer(AR_f8, f8), npt.NDArray[Any]) + +assert_type(np.frexp.__name__, Literal["frexp"]) +assert_type(np.frexp.__qualname__, Literal["frexp"]) +assert_type(np.frexp.ntypes, Literal[4]) +assert_type(np.frexp.identity, None) +assert_type(np.frexp.nin, Literal[1]) +assert_type(np.frexp.nout, Literal[2]) +assert_type(np.frexp.nargs, Literal[3]) +assert_type(np.frexp.signature, None) +assert_type(np.frexp(f8), tuple[Any, Any]) +assert_type(np.frexp(AR_f8), tuple[npt.NDArray[Any], npt.NDArray[Any]]) + +assert_type(np.divmod.__name__, Literal["divmod"]) +assert_type(np.divmod.__qualname__, Literal["divmod"]) +assert_type(np.divmod.ntypes, Literal[15]) +assert_type(np.divmod.identity, None) +assert_type(np.divmod.nin, Literal[2]) +assert_type(np.divmod.nout, Literal[2]) +assert_type(np.divmod.nargs, Literal[4]) +assert_type(np.divmod.signature, None) +assert_type(np.divmod(f8, f8), tuple[Any, Any]) +assert_type(np.divmod(AR_f8, f8), tuple[npt.NDArray[Any], npt.NDArray[Any]]) + +assert_type(np.matmul.__name__, Literal["matmul"]) +assert_type(np.matmul.__qualname__, Literal["matmul"]) +assert_type(np.matmul.ntypes, Literal[19]) +assert_type(np.matmul.identity, None) +assert_type(np.matmul.nin, Literal[2]) +assert_type(np.matmul.nout, Literal[1]) +assert_type(np.matmul.nargs, Literal[3]) +assert_type(np.matmul.signature, Literal["(n?,k),(k,m?)->(n?,m?)"]) +assert_type(np.matmul.identity, None) +assert_type(np.matmul(AR_f8, AR_f8), Any) +assert_type(np.matmul(AR_f8, AR_f8, axes=[(0, 1), (0, 1), (0, 1)]), Any) + +assert_type(np.vecdot.__name__, Literal["vecdot"]) +assert_type(np.vecdot.__qualname__, Literal["vecdot"]) +assert_type(np.vecdot.ntypes, Literal[19]) +assert_type(np.vecdot.identity, None) +assert_type(np.vecdot.nin, Literal[2]) +assert_type(np.vecdot.nout, Literal[1]) +assert_type(np.vecdot.nargs, Literal[3]) +assert_type(np.vecdot.signature, Literal["(n),(n)->()"]) +assert_type(np.vecdot.identity, None) +assert_type(np.vecdot(AR_f8, AR_f8), Any) + +assert_type(np.bitwise_count.__name__, Literal["bitwise_count"]) +assert_type(np.bitwise_count.__qualname__, Literal["bitwise_count"]) +assert_type(np.bitwise_count.ntypes, Literal[11]) +assert_type(np.bitwise_count.identity, None) +assert_type(np.bitwise_count.nin, Literal[1]) +assert_type(np.bitwise_count.nout, Literal[1]) +assert_type(np.bitwise_count.nargs, Literal[2]) +assert_type(np.bitwise_count.signature, None) +assert_type(np.bitwise_count.identity, None) +assert_type(np.bitwise_count(i8), Any) +assert_type(np.bitwise_count(AR_i8), npt.NDArray[Any]) + +assert_type(np.absolute.outer(), NoReturn) +assert_type(np.frexp.outer(), NoReturn) +assert_type(np.divmod.outer(), NoReturn) +assert_type(np.matmul.outer(), NoReturn) + +assert_type(np.absolute.reduceat(), NoReturn) +assert_type(np.frexp.reduceat(), NoReturn) +assert_type(np.divmod.reduceat(), NoReturn) +assert_type(np.matmul.reduceat(), NoReturn) + +assert_type(np.absolute.reduce(), NoReturn) +assert_type(np.frexp.reduce(), NoReturn) +assert_type(np.divmod.reduce(), NoReturn) +assert_type(np.matmul.reduce(), NoReturn) + +assert_type(np.absolute.accumulate(), NoReturn) +assert_type(np.frexp.accumulate(), NoReturn) +assert_type(np.divmod.accumulate(), NoReturn) +assert_type(np.matmul.accumulate(), NoReturn) + +assert_type(np.frexp.at(), NoReturn) +assert_type(np.divmod.at(), NoReturn) +assert_type(np.matmul.at(), NoReturn) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/warnings_and_errors.pyi b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/warnings_and_errors.pyi new file mode 100644 index 0000000000000000000000000000000000000000..1f1be133be4f2e5d2bfeac6c6109bbb04c5df97c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/numpy/typing/tests/data/reveal/warnings_and_errors.pyi @@ -0,0 +1,11 @@ +import numpy.exceptions as ex + +from typing_extensions import assert_type + +assert_type(ex.ModuleDeprecationWarning(), ex.ModuleDeprecationWarning) +assert_type(ex.VisibleDeprecationWarning(), ex.VisibleDeprecationWarning) +assert_type(ex.ComplexWarning(), ex.ComplexWarning) +assert_type(ex.RankWarning(), ex.RankWarning) +assert_type(ex.TooHardError(), ex.TooHardError) +assert_type(ex.AxisError("test"), ex.AxisError) +assert_type(ex.AxisError(5, 1), ex.AxisError) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-source/setup.cfg b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-source/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-source/setup.py b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-source/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..ce9080640c0cf4a5b42274ac0c91278d6df17643 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-source/setup.py @@ -0,0 +1,7 @@ +import setuptools + +setuptools.setup( + name="my-test-package", + version="1.0", + zip_safe=True, +) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-zip/my-test-package.zip b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-zip/my-test-package.zip new file mode 100644 index 0000000000000000000000000000000000000000..81f9a0170f7aca13a0d1d8e5201c2c186554ce7e Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package-zip/my-test-package.zip differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/PKG-INFO b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..7328e3f7d18d31422d13b4fd63a7ef746956b7db --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/PKG-INFO @@ -0,0 +1,10 @@ +Metadata-Version: 1.0 +Name: my-test-package +Version: 1.0 +Summary: UNKNOWN +Home-page: UNKNOWN +Author: UNKNOWN +Author-email: UNKNOWN +License: UNKNOWN +Description: UNKNOWN +Platform: UNKNOWN diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/SOURCES.txt b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..3c4ee1676d4fec9f10f3c0f58c2d992d97975b85 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/SOURCES.txt @@ -0,0 +1,7 @@ +setup.cfg +setup.py +my_test_package.egg-info/PKG-INFO +my_test_package.egg-info/SOURCES.txt +my_test_package.egg-info/dependency_links.txt +my_test_package.egg-info/top_level.txt +my_test_package.egg-info/zip-safe \ No newline at end of file diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/dependency_links.txt b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/top_level.txt b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/top_level.txt @@ -0,0 +1 @@ + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/zip-safe b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/zip-safe new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_unpacked-egg/my_test_package-1.0-py3.7.egg/EGG-INFO/zip-safe @@ -0,0 +1 @@ + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_zipped-egg/my_test_package-1.0-py3.7.egg b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_zipped-egg/my_test_package-1.0-py3.7.egg new file mode 100644 index 0000000000000000000000000000000000000000..5115b8957da8213d02b718649c8702ac88bf0e72 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/pkg_resources/tests/data/my-test-package_zipped-egg/my_test_package-1.0-py3.7.egg differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/INSTALLER b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/LICENSE b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..1bf98523e33170372cf0702f03e38dadec3c5094 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/LICENSE @@ -0,0 +1,18 @@ +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/METADATA b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..d03671b6130146132d6e6fc8a42256cfbcd3db0e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/METADATA @@ -0,0 +1,125 @@ +Metadata-Version: 2.1 +Name: pyparsing +Version: 3.2.3 +Summary: pyparsing module - Classes and methods to define and execute parsing grammars +Author-email: Paul McGuire +Requires-Python: >=3.9 +Description-Content-Type: text/x-rst +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Information Technology +Classifier: License :: OSI Approved :: MIT License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Software Development :: Compilers +Classifier: Topic :: Text Processing +Classifier: Typing :: Typed +Requires-Dist: railroad-diagrams ; extra == "diagrams" +Requires-Dist: jinja2 ; extra == "diagrams" +Project-URL: Homepage, https://github.com/pyparsing/pyparsing/ +Provides-Extra: diagrams + +PyParsing -- A Python Parsing Module +==================================== + +|Version| |Build Status| |Coverage| |License| |Python Versions| |Snyk Score| + +Introduction +============ + +The pyparsing module is an alternative approach to creating and +executing simple grammars, vs. the traditional lex/yacc approach, or the +use of regular expressions. The pyparsing module provides a library of +classes that client code uses to construct the grammar directly in +Python code. + +*[Since first writing this description of pyparsing in late 2003, this +technique for developing parsers has become more widespread, under the +name Parsing Expression Grammars - PEGs. See more information on PEGs* +`here `__ +*.]* + +Here is a program to parse ``"Hello, World!"`` (or any greeting of the form +``"salutation, addressee!"``): + +.. code:: python + + from pyparsing import Word, alphas + greet = Word(alphas) + "," + Word(alphas) + "!" + hello = "Hello, World!" + print(hello, "->", greet.parse_string(hello)) + +The program outputs the following:: + + Hello, World! -> ['Hello', ',', 'World', '!'] + +The Python representation of the grammar is quite readable, owing to the +self-explanatory class names, and the use of '+', '|' and '^' operator +definitions. + +The parsed results returned from ``parse_string()`` is a collection of type +``ParseResults``, which can be accessed as a +nested list, a dictionary, or an object with named attributes. + +The pyparsing module handles some of the problems that are typically +vexing when writing text parsers: + +- extra or missing whitespace (the above program will also handle ``"Hello,World!"``, ``"Hello , World !"``, etc.) +- quoted strings +- embedded comments + +The examples directory includes a simple SQL parser, simple CORBA IDL +parser, a config file parser, a chemical formula parser, and a four- +function algebraic notation parser, among many others. + +Documentation +============= + +There are many examples in the online docstrings of the classes +and methods in pyparsing. You can find them compiled into `online docs `__. Additional +documentation resources and project info are listed in the online +`GitHub wiki `__. An +entire directory of examples can be found `here `__. + +License +======= + +MIT License. See header of the `pyparsing __init__.py `__ file. + +History +======= + +See `CHANGES `__ file. + +.. |Build Status| image:: https://github.com/pyparsing/pyparsing/actions/workflows/ci.yml/badge.svg + :target: https://github.com/pyparsing/pyparsing/actions/workflows/ci.yml + +.. |Coverage| image:: https://codecov.io/gh/pyparsing/pyparsing/branch/master/graph/badge.svg + :target: https://codecov.io/gh/pyparsing/pyparsing + +.. |Version| image:: https://img.shields.io/pypi/v/pyparsing?style=flat-square + :target: https://pypi.org/project/pyparsing/ + :alt: Version + +.. |License| image:: https://img.shields.io/pypi/l/pyparsing.svg?style=flat-square + :target: https://pypi.org/project/pyparsing/ + :alt: License + +.. |Python Versions| image:: https://img.shields.io/pypi/pyversions/pyparsing.svg?style=flat-square + :target: https://pypi.org/project/python-liquid/ + :alt: Python versions + +.. |Snyk Score| image:: https://snyk.io//advisor/python/pyparsing/badge.svg + :target: https://snyk.io//advisor/python/pyparsing + :alt: pyparsing + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/RECORD b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..1d0c7bd466bbdb5d0f9e27aafb7410e1601bb539 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/RECORD @@ -0,0 +1,32 @@ +pyparsing-3.2.3.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +pyparsing-3.2.3.dist-info/LICENSE,sha256=ENUSChaAWAT_2otojCIL-06POXQbVzIGBNRVowngGXI,1023 +pyparsing-3.2.3.dist-info/METADATA,sha256=DkRwMIwMm2LRS_DidFhg1Osj9kV-4QLQCke98yA5PKA,5042 +pyparsing-3.2.3.dist-info/RECORD,, +pyparsing-3.2.3.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81 +pyparsing/__init__.py,sha256=ZwTdfiWUpx8XVdEseRd8Clj0egtRca5SUCqzKq58aT8,9038 +pyparsing/__pycache__/__init__.cpython-310.pyc,, +pyparsing/__pycache__/actions.cpython-310.pyc,, +pyparsing/__pycache__/common.cpython-310.pyc,, +pyparsing/__pycache__/core.cpython-310.pyc,, +pyparsing/__pycache__/exceptions.cpython-310.pyc,, +pyparsing/__pycache__/helpers.cpython-310.pyc,, +pyparsing/__pycache__/results.cpython-310.pyc,, +pyparsing/__pycache__/testing.cpython-310.pyc,, +pyparsing/__pycache__/unicode.cpython-310.pyc,, +pyparsing/__pycache__/util.cpython-310.pyc,, +pyparsing/actions.py,sha256=c7j1u5A2l2BF4jzQdIX27IsqyzfTK1vZuCWgtIsm1lQ,7265 +pyparsing/common.py,sha256=SyZAEQpwDAPOpUenOltIv-y1N-tzSN3nCEzNYcOpiDw,13673 +pyparsing/core.py,sha256=Ita1roVKwqzfb4IvX88LwZglxuiio7JuW9eoylkhZfU,234772 +pyparsing/diagram/__init__.py,sha256=7VkLQc6p4bnQC3CFL4VPtECih5j-b-LkN57q_rXMkzM,26976 +pyparsing/diagram/__pycache__/__init__.cpython-310.pyc,, +pyparsing/exceptions.py,sha256=zI9ZGsMBqpcDB2qKsrrpYGun1SCwNrvpuNhSb_goBuo,9892 +pyparsing/helpers.py,sha256=TE6jBjkIBb_xjaeUKXPjS-1-RnH_IkH_DaHnaMDRwWM,40155 +pyparsing/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pyparsing/results.py,sha256=208K1xgCyijpE0auARMNqwJXLCPN0Y7YCsdXf3ynZiE,26254 +pyparsing/testing.py,sha256=r-kKXcv_8r1G-4vrprBkJeuQFpr94OSm0jlXdGwAYF0,14682 +pyparsing/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pyparsing/tools/__pycache__/__init__.cpython-310.pyc,, +pyparsing/tools/__pycache__/cvt_pyparsing_pep8_names.cpython-310.pyc,, +pyparsing/tools/cvt_pyparsing_pep8_names.py,sha256=CKvxIBiDJLvmbnVv3mK0tyNUZN3ub3eQ4Z1eO54x8_U,5369 +pyparsing/unicode.py,sha256=doanv7BYQB4EdQRfdXigaiV_SDhZyyaBSwgvVEz6eXc,10614 +pyparsing/util.py,sha256=knfaF8_pjHLy7PL_YdpsM9VaQBrQDqo4Lnh0huTwTl4,14344 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/WHEEL b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..3b5e64b5e6c4a210201d1676a891fd57b15cda99 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing-3.2.3.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: flit 3.9.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d839fd253756c768e0a783984ee0aaf6b4588023 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/__init__.py @@ -0,0 +1,326 @@ +# module pyparsing.py +# +# Copyright (c) 2003-2022 Paul T. McGuire +# +# Permission is hereby granted, free of charge, to any person obtaining +# a copy of this software and associated documentation files (the +# "Software"), to deal in the Software without restriction, including +# without limitation the rights to use, copy, modify, merge, publish, +# distribute, sublicense, and/or sell copies of the Software, and to +# permit persons to whom the Software is furnished to do so, subject to +# the following conditions: +# +# The above copyright notice and this permission notice shall be +# included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +# + +__doc__ = """ +pyparsing module - Classes and methods to define and execute parsing grammars +============================================================================= + +The pyparsing module is an alternative approach to creating and +executing simple grammars, vs. the traditional lex/yacc approach, or the +use of regular expressions. With pyparsing, you don't need to learn +a new syntax for defining grammars or matching expressions - the parsing +module provides a library of classes that you use to construct the +grammar directly in Python. + +Here is a program to parse "Hello, World!" (or any greeting of the form +``", !"``), built up using :class:`Word`, +:class:`Literal`, and :class:`And` elements +(the :meth:`'+'` operators create :class:`And` expressions, +and the strings are auto-converted to :class:`Literal` expressions):: + + from pyparsing import Word, alphas + + # define grammar of a greeting + greet = Word(alphas) + "," + Word(alphas) + "!" + + hello = "Hello, World!" + print(hello, "->", greet.parse_string(hello)) + +The program outputs the following:: + + Hello, World! -> ['Hello', ',', 'World', '!'] + +The Python representation of the grammar is quite readable, owing to the +self-explanatory class names, and the use of :class:`'+'`, +:class:`'|'`, :class:`'^'` and :class:`'&'` operators. + +The :class:`ParseResults` object returned from +:class:`ParserElement.parse_string` can be +accessed as a nested list, a dictionary, or an object with named +attributes. + +The pyparsing module handles some of the problems that are typically +vexing when writing text parsers: + + - extra or missing whitespace (the above program will also handle + "Hello,World!", "Hello , World !", etc.) + - quoted strings + - embedded comments + + +Getting Started - +----------------- +Visit the classes :class:`ParserElement` and :class:`ParseResults` to +see the base classes that most other pyparsing +classes inherit from. Use the docstrings for examples of how to: + + - construct literal match expressions from :class:`Literal` and + :class:`CaselessLiteral` classes + - construct character word-group expressions using the :class:`Word` + class + - see how to create repetitive expressions using :class:`ZeroOrMore` + and :class:`OneOrMore` classes + - use :class:`'+'`, :class:`'|'`, :class:`'^'`, + and :class:`'&'` operators to combine simple expressions into + more complex ones + - associate names with your parsed results using + :class:`ParserElement.set_results_name` + - access the parsed data, which is returned as a :class:`ParseResults` + object + - find some helpful expression short-cuts like :class:`DelimitedList` + and :class:`one_of` + - find more useful common expressions in the :class:`pyparsing_common` + namespace class +""" +from typing import NamedTuple + + +class version_info(NamedTuple): + major: int + minor: int + micro: int + releaselevel: str + serial: int + + @property + def __version__(self): + return ( + f"{self.major}.{self.minor}.{self.micro}" + + ( + f"{'r' if self.releaselevel[0] == 'c' else ''}{self.releaselevel[0]}{self.serial}", + "", + )[self.releaselevel == "final"] + ) + + def __str__(self): + return f"{__name__} {self.__version__} / {__version_time__}" + + def __repr__(self): + return f"{__name__}.{type(self).__name__}({', '.join('{}={!r}'.format(*nv) for nv in zip(self._fields, self))})" + + +__version_info__ = version_info(3, 2, 3, "final", 1) +__version_time__ = "25 Mar 2025 01:38 UTC" +__version__ = __version_info__.__version__ +__versionTime__ = __version_time__ +__author__ = "Paul McGuire " + +from .util import * +from .exceptions import * +from .actions import * +from .core import __diag__, __compat__ +from .results import * +from .core import * +from .core import _builtin_exprs as core_builtin_exprs +from .helpers import * +from .helpers import _builtin_exprs as helper_builtin_exprs + +from .unicode import unicode_set, UnicodeRangeList, pyparsing_unicode as unicode +from .testing import pyparsing_test as testing +from .common import ( + pyparsing_common as common, + _builtin_exprs as common_builtin_exprs, +) + +# Compatibility synonyms +if "pyparsing_unicode" not in globals(): + pyparsing_unicode = unicode # type: ignore[misc] +if "pyparsing_common" not in globals(): + pyparsing_common = common +if "pyparsing_test" not in globals(): + pyparsing_test = testing + +core_builtin_exprs += common_builtin_exprs + helper_builtin_exprs + + +__all__ = [ + "__version__", + "__version_time__", + "__author__", + "__compat__", + "__diag__", + "And", + "AtLineStart", + "AtStringStart", + "CaselessKeyword", + "CaselessLiteral", + "CharsNotIn", + "CloseMatch", + "Combine", + "DelimitedList", + "Dict", + "Each", + "Empty", + "FollowedBy", + "Forward", + "GoToColumn", + "Group", + "IndentedBlock", + "Keyword", + "LineEnd", + "LineStart", + "Literal", + "Located", + "PrecededBy", + "MatchFirst", + "NoMatch", + "NotAny", + "OneOrMore", + "OnlyOnce", + "OpAssoc", + "Opt", + "Optional", + "Or", + "ParseBaseException", + "ParseElementEnhance", + "ParseException", + "ParseExpression", + "ParseFatalException", + "ParseResults", + "ParseSyntaxException", + "ParserElement", + "PositionToken", + "QuotedString", + "RecursiveGrammarException", + "Regex", + "SkipTo", + "StringEnd", + "StringStart", + "Suppress", + "Tag", + "Token", + "TokenConverter", + "White", + "Word", + "WordEnd", + "WordStart", + "ZeroOrMore", + "Char", + "alphanums", + "alphas", + "alphas8bit", + "any_close_tag", + "any_open_tag", + "autoname_elements", + "c_style_comment", + "col", + "common_html_entity", + "condition_as_parse_action", + "counted_array", + "cpp_style_comment", + "dbl_quoted_string", + "dbl_slash_comment", + "delimited_list", + "dict_of", + "empty", + "hexnums", + "html_comment", + "identchars", + "identbodychars", + "infix_notation", + "java_style_comment", + "line", + "line_end", + "line_start", + "lineno", + "make_html_tags", + "make_xml_tags", + "match_only_at_col", + "match_previous_expr", + "match_previous_literal", + "nested_expr", + "null_debug_action", + "nums", + "one_of", + "original_text_for", + "printables", + "punc8bit", + "pyparsing_common", + "pyparsing_test", + "pyparsing_unicode", + "python_style_comment", + "quoted_string", + "remove_quotes", + "replace_with", + "replace_html_entity", + "rest_of_line", + "sgl_quoted_string", + "srange", + "string_end", + "string_start", + "token_map", + "trace_parse_action", + "ungroup", + "unicode_set", + "unicode_string", + "with_attribute", + "with_class", + # pre-PEP8 compatibility names + "__versionTime__", + "anyCloseTag", + "anyOpenTag", + "cStyleComment", + "commonHTMLEntity", + "conditionAsParseAction", + "countedArray", + "cppStyleComment", + "dblQuotedString", + "dblSlashComment", + "delimitedList", + "dictOf", + "htmlComment", + "indentedBlock", + "infixNotation", + "javaStyleComment", + "lineEnd", + "lineStart", + "locatedExpr", + "makeHTMLTags", + "makeXMLTags", + "matchOnlyAtCol", + "matchPreviousExpr", + "matchPreviousLiteral", + "nestedExpr", + "nullDebugAction", + "oneOf", + "opAssoc", + "originalTextFor", + "pythonStyleComment", + "quotedString", + "removeQuotes", + "replaceHTMLEntity", + "replaceWith", + "restOfLine", + "sglQuotedString", + "stringEnd", + "stringStart", + "tokenMap", + "traceParseAction", + "unicodeString", + "withAttribute", + "withClass", + "common", + "unicode", + "testing", +] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/actions.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/actions.py new file mode 100644 index 0000000000000000000000000000000000000000..0153cc7132abe61554521e8dc15ae3d3041de94c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/actions.py @@ -0,0 +1,219 @@ +# actions.py +from __future__ import annotations + +from typing import Union, Callable, Any + +from .exceptions import ParseException +from .util import col, replaced_by_pep8 +from .results import ParseResults + + +ParseAction = Union[ + Callable[[], Any], + Callable[[ParseResults], Any], + Callable[[int, ParseResults], Any], + Callable[[str, int, ParseResults], Any], +] + + +class OnlyOnce: + """ + Wrapper for parse actions, to ensure they are only called once. + Note: parse action signature must include all 3 arguments. + """ + + def __init__(self, method_call: Callable[[str, int, ParseResults], Any]) -> None: + from .core import _trim_arity + + self.callable = _trim_arity(method_call) + self.called = False + + def __call__(self, s: str, l: int, t: ParseResults) -> ParseResults: + if not self.called: + results = self.callable(s, l, t) + self.called = True + return results + raise ParseException(s, l, "OnlyOnce obj called multiple times w/out reset") + + def reset(self): + """ + Allow the associated parse action to be called once more. + """ + + self.called = False + + +def match_only_at_col(n: int) -> ParseAction: + """ + Helper method for defining parse actions that require matching at + a specific column in the input text. + """ + + def verify_col(strg: str, locn: int, toks: ParseResults) -> None: + if col(locn, strg) != n: + raise ParseException(strg, locn, f"matched token not at column {n}") + + return verify_col + + +def replace_with(repl_str: str) -> ParseAction: + """ + Helper method for common parse actions that simply return + a literal value. Especially useful when used with + :class:`transform_string` (). + + Example:: + + num = Word(nums).set_parse_action(lambda toks: int(toks[0])) + na = one_of("N/A NA").set_parse_action(replace_with(math.nan)) + term = na | num + + term[1, ...].parse_string("324 234 N/A 234") # -> [324, 234, nan, 234] + """ + return lambda s, l, t: [repl_str] + + +def remove_quotes(s: str, l: int, t: ParseResults) -> Any: + """ + Helper parse action for removing quotation marks from parsed + quoted strings. + + Example:: + + # by default, quotation marks are included in parsed results + quoted_string.parse_string("'Now is the Winter of our Discontent'") # -> ["'Now is the Winter of our Discontent'"] + + # use remove_quotes to strip quotation marks from parsed results + quoted_string.set_parse_action(remove_quotes) + quoted_string.parse_string("'Now is the Winter of our Discontent'") # -> ["Now is the Winter of our Discontent"] + """ + return t[0][1:-1] + + +def with_attribute(*args: tuple[str, str], **attr_dict) -> ParseAction: + """ + Helper to create a validating parse action to be used with start + tags created with :class:`make_xml_tags` or + :class:`make_html_tags`. Use ``with_attribute`` to qualify + a starting tag with a required attribute value, to avoid false + matches on common tags such as ```` or ``
``. + + Call ``with_attribute`` with a series of attribute names and + values. Specify the list of filter attributes names and values as: + + - keyword arguments, as in ``(align="right")``, or + - as an explicit dict with ``**`` operator, when an attribute + name is also a Python reserved word, as in ``**{"class":"Customer", "align":"right"}`` + - a list of name-value tuples, as in ``(("ns1:class", "Customer"), ("ns2:align", "right"))`` + + For attribute names with a namespace prefix, you must use the second + form. Attribute names are matched insensitive to upper/lower case. + + If just testing for ``class`` (with or without a namespace), use + :class:`with_class`. + + To verify that the attribute exists, but without specifying a value, + pass ``with_attribute.ANY_VALUE`` as the value. + + Example:: + + html = ''' +
+ Some text +
1 4 0 1 0
+
1,3 2,3 1,1
+
this has no type
+
+ ''' + div,div_end = make_html_tags("div") + + # only match div tag having a type attribute with value "grid" + div_grid = div().set_parse_action(with_attribute(type="grid")) + grid_expr = div_grid + SkipTo(div | div_end)("body") + for grid_header in grid_expr.search_string(html): + print(grid_header.body) + + # construct a match with any div tag having a type attribute, regardless of the value + div_any_type = div().set_parse_action(with_attribute(type=with_attribute.ANY_VALUE)) + div_expr = div_any_type + SkipTo(div | div_end)("body") + for div_header in div_expr.search_string(html): + print(div_header.body) + + prints:: + + 1 4 0 1 0 + + 1 4 0 1 0 + 1,3 2,3 1,1 + """ + attrs_list: list[tuple[str, str]] = [] + if args: + attrs_list.extend(args) + else: + attrs_list.extend(attr_dict.items()) + + def pa(s: str, l: int, tokens: ParseResults) -> None: + for attrName, attrValue in attrs_list: + if attrName not in tokens: + raise ParseException(s, l, "no matching attribute " + attrName) + if attrValue != with_attribute.ANY_VALUE and tokens[attrName] != attrValue: # type: ignore [attr-defined] + raise ParseException( + s, + l, + f"attribute {attrName!r} has value {tokens[attrName]!r}, must be {attrValue!r}", + ) + + return pa + + +with_attribute.ANY_VALUE = object() # type: ignore [attr-defined] + + +def with_class(classname: str, namespace: str = "") -> ParseAction: + """ + Simplified version of :class:`with_attribute` when + matching on a div class - made difficult because ``class`` is + a reserved word in Python. + + Example:: + + html = ''' +
+ Some text +
1 4 0 1 0
+
1,3 2,3 1,1
+
this <div> has no class
+
+ + ''' + div,div_end = make_html_tags("div") + div_grid = div().set_parse_action(with_class("grid")) + + grid_expr = div_grid + SkipTo(div | div_end)("body") + for grid_header in grid_expr.search_string(html): + print(grid_header.body) + + div_any_type = div().set_parse_action(with_class(withAttribute.ANY_VALUE)) + div_expr = div_any_type + SkipTo(div | div_end)("body") + for div_header in div_expr.search_string(html): + print(div_header.body) + + prints:: + + 1 4 0 1 0 + + 1 4 0 1 0 + 1,3 2,3 1,1 + """ + classattr = f"{namespace}:class" if namespace else "class" + return with_attribute(**{classattr: classname}) + + +# Compatibility synonyms +# fmt: off +replaceWith = replaced_by_pep8("replaceWith", replace_with) +removeQuotes = replaced_by_pep8("removeQuotes", remove_quotes) +withAttribute = replaced_by_pep8("withAttribute", with_attribute) +withClass = replaced_by_pep8("withClass", with_class) +matchOnlyAtCol = replaced_by_pep8("matchOnlyAtCol", match_only_at_col) +# fmt: on diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/common.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/common.py new file mode 100644 index 0000000000000000000000000000000000000000..e46511086dadb79357441680531d286124da63c7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/common.py @@ -0,0 +1,434 @@ +# common.py +from .core import * +from .helpers import DelimitedList, any_open_tag, any_close_tag +from datetime import datetime + + +# some other useful expressions - using lower-case class name since we are really using this as a namespace +class pyparsing_common: + """Here are some common low-level expressions that may be useful in + jump-starting parser development: + + - numeric forms (:class:`integers`, :class:`reals`, + :class:`scientific notation`) + - common :class:`programming identifiers` + - network addresses (:class:`MAC`, + :class:`IPv4`, :class:`IPv6`) + - ISO8601 :class:`dates` and + :class:`datetime` + - :class:`UUID` + - :class:`comma-separated list` + - :class:`url` + + Parse actions: + + - :class:`convert_to_integer` + - :class:`convert_to_float` + - :class:`convert_to_date` + - :class:`convert_to_datetime` + - :class:`strip_html_tags` + - :class:`upcase_tokens` + - :class:`downcase_tokens` + + Example:: + + pyparsing_common.number.run_tests(''' + # any int or real number, returned as the appropriate type + 100 + -100 + +100 + 3.14159 + 6.02e23 + 1e-12 + ''') + + pyparsing_common.fnumber.run_tests(''' + # any int or real number, returned as float + 100 + -100 + +100 + 3.14159 + 6.02e23 + 1e-12 + ''') + + pyparsing_common.hex_integer.run_tests(''' + # hex numbers + 100 + FF + ''') + + pyparsing_common.fraction.run_tests(''' + # fractions + 1/2 + -3/4 + ''') + + pyparsing_common.mixed_integer.run_tests(''' + # mixed fractions + 1 + 1/2 + -3/4 + 1-3/4 + ''') + + import uuid + pyparsing_common.uuid.set_parse_action(token_map(uuid.UUID)) + pyparsing_common.uuid.run_tests(''' + # uuid + 12345678-1234-5678-1234-567812345678 + ''') + + prints:: + + # any int or real number, returned as the appropriate type + 100 + [100] + + -100 + [-100] + + +100 + [100] + + 3.14159 + [3.14159] + + 6.02e23 + [6.02e+23] + + 1e-12 + [1e-12] + + # any int or real number, returned as float + 100 + [100.0] + + -100 + [-100.0] + + +100 + [100.0] + + 3.14159 + [3.14159] + + 6.02e23 + [6.02e+23] + + 1e-12 + [1e-12] + + # hex numbers + 100 + [256] + + FF + [255] + + # fractions + 1/2 + [0.5] + + -3/4 + [-0.75] + + # mixed fractions + 1 + [1] + + 1/2 + [0.5] + + -3/4 + [-0.75] + + 1-3/4 + [1.75] + + # uuid + 12345678-1234-5678-1234-567812345678 + [UUID('12345678-1234-5678-1234-567812345678')] + """ + + convert_to_integer = token_map(int) + """ + Parse action for converting parsed integers to Python int + """ + + convert_to_float = token_map(float) + """ + Parse action for converting parsed numbers to Python float + """ + + integer = Word(nums).set_name("integer").set_parse_action(convert_to_integer) + """expression that parses an unsigned integer, returns an int""" + + hex_integer = ( + Word(hexnums).set_name("hex integer").set_parse_action(token_map(int, 16)) + ) + """expression that parses a hexadecimal integer, returns an int""" + + signed_integer = ( + Regex(r"[+-]?\d+") + .set_name("signed integer") + .set_parse_action(convert_to_integer) + ) + """expression that parses an integer with optional leading sign, returns an int""" + + fraction = ( + signed_integer().set_parse_action(convert_to_float) + + "/" + + signed_integer().set_parse_action(convert_to_float) + ).set_name("fraction") + """fractional expression of an integer divided by an integer, returns a float""" + fraction.add_parse_action(lambda tt: tt[0] / tt[-1]) + + mixed_integer = ( + fraction | signed_integer + Opt(Opt("-").suppress() + fraction) + ).set_name("fraction or mixed integer-fraction") + """mixed integer of the form 'integer - fraction', with optional leading integer, returns float""" + mixed_integer.add_parse_action(sum) + + real = ( + Regex(r"[+-]?(?:\d+\.\d*|\.\d+)") + .set_name("real number") + .set_parse_action(convert_to_float) + ) + """expression that parses a floating point number and returns a float""" + + sci_real = ( + Regex(r"[+-]?(?:\d+(?:[eE][+-]?\d+)|(?:\d+\.\d*|\.\d+)(?:[eE][+-]?\d+)?)") + .set_name("real number with scientific notation") + .set_parse_action(convert_to_float) + ) + """expression that parses a floating point number with optional + scientific notation and returns a float""" + + # streamlining this expression makes the docs nicer-looking + number = (sci_real | real | signed_integer).set_name("number").streamline() + """any numeric expression, returns the corresponding Python type""" + + fnumber = ( + Regex(r"[+-]?\d+\.?\d*(?:[eE][+-]?\d+)?") + .set_name("fnumber") + .set_parse_action(convert_to_float) + ) + """any int or real number, returned as float""" + + ieee_float = ( + Regex(r"(?i:[+-]?(?:(?:\d+\.?\d*(?:e[+-]?\d+)?)|nan|inf(?:inity)?))") + .set_name("ieee_float") + .set_parse_action(convert_to_float) + ) + """any floating-point literal (int, real number, infinity, or NaN), returned as float""" + + identifier = Word(identchars, identbodychars).set_name("identifier") + """typical code identifier (leading alpha or '_', followed by 0 or more alphas, nums, or '_')""" + + ipv4_address = Regex( + r"(25[0-5]|2[0-4][0-9]|1?[0-9]{1,2})(\.(25[0-5]|2[0-4][0-9]|1?[0-9]{1,2})){3}" + ).set_name("IPv4 address") + "IPv4 address (``0.0.0.0 - 255.255.255.255``)" + + _ipv6_part = Regex(r"[0-9a-fA-F]{1,4}").set_name("hex_integer") + _full_ipv6_address = (_ipv6_part + (":" + _ipv6_part) * 7).set_name( + "full IPv6 address" + ) + _short_ipv6_address = ( + Opt(_ipv6_part + (":" + _ipv6_part) * (0, 6)) + + "::" + + Opt(_ipv6_part + (":" + _ipv6_part) * (0, 6)) + ).set_name("short IPv6 address") + _short_ipv6_address.add_condition( + lambda t: sum(1 for tt in t if pyparsing_common._ipv6_part.matches(tt)) < 8 + ) + _mixed_ipv6_address = ("::ffff:" + ipv4_address).set_name("mixed IPv6 address") + ipv6_address = Combine( + (_full_ipv6_address | _mixed_ipv6_address | _short_ipv6_address).set_name( + "IPv6 address" + ) + ).set_name("IPv6 address") + "IPv6 address (long, short, or mixed form)" + + mac_address = Regex( + r"[0-9a-fA-F]{2}([:.-])[0-9a-fA-F]{2}(?:\1[0-9a-fA-F]{2}){4}" + ).set_name("MAC address") + "MAC address xx:xx:xx:xx:xx (may also have '-' or '.' delimiters)" + + @staticmethod + def convert_to_date(fmt: str = "%Y-%m-%d"): + """ + Helper to create a parse action for converting parsed date string to Python datetime.date + + Params - + - fmt - format to be passed to datetime.strptime (default= ``"%Y-%m-%d"``) + + Example:: + + date_expr = pyparsing_common.iso8601_date.copy() + date_expr.set_parse_action(pyparsing_common.convert_to_date()) + print(date_expr.parse_string("1999-12-31")) + + prints:: + + [datetime.date(1999, 12, 31)] + """ + + def cvt_fn(ss, ll, tt): + try: + return datetime.strptime(tt[0], fmt).date() + except ValueError as ve: + raise ParseException(ss, ll, str(ve)) + + return cvt_fn + + @staticmethod + def convert_to_datetime(fmt: str = "%Y-%m-%dT%H:%M:%S.%f"): + """Helper to create a parse action for converting parsed + datetime string to Python datetime.datetime + + Params - + - fmt - format to be passed to datetime.strptime (default= ``"%Y-%m-%dT%H:%M:%S.%f"``) + + Example:: + + dt_expr = pyparsing_common.iso8601_datetime.copy() + dt_expr.set_parse_action(pyparsing_common.convert_to_datetime()) + print(dt_expr.parse_string("1999-12-31T23:59:59.999")) + + prints:: + + [datetime.datetime(1999, 12, 31, 23, 59, 59, 999000)] + """ + + def cvt_fn(s, l, t): + try: + return datetime.strptime(t[0], fmt) + except ValueError as ve: + raise ParseException(s, l, str(ve)) + + return cvt_fn + + iso8601_date = Regex( + r"(?P\d{4})(?:-(?P\d\d)(?:-(?P\d\d))?)?" + ).set_name("ISO8601 date") + "ISO8601 date (``yyyy-mm-dd``)" + + iso8601_datetime = Regex( + r"(?P\d{4})-(?P\d\d)-(?P\d\d)[T ](?P\d\d):(?P\d\d)(:(?P\d\d(\.\d*)?)?)?(?PZ|[+-]\d\d:?\d\d)?" + ).set_name("ISO8601 datetime") + "ISO8601 datetime (``yyyy-mm-ddThh:mm:ss.s(Z|+-00:00)``) - trailing seconds, milliseconds, and timezone optional; accepts separating ``'T'`` or ``' '``" + + uuid = Regex(r"[0-9a-fA-F]{8}(-[0-9a-fA-F]{4}){3}-[0-9a-fA-F]{12}").set_name("UUID") + "UUID (``xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx``)" + + _html_stripper = any_open_tag.suppress() | any_close_tag.suppress() + + @staticmethod + def strip_html_tags(s: str, l: int, tokens: ParseResults): + """Parse action to remove HTML tags from web page HTML source + + Example:: + + # strip HTML links from normal text + text = 'More info at the pyparsing wiki page' + td, td_end = make_html_tags("TD") + table_text = td + SkipTo(td_end).set_parse_action(pyparsing_common.strip_html_tags)("body") + td_end + print(table_text.parse_string(text).body) + + Prints:: + + More info at the pyparsing wiki page + """ + return pyparsing_common._html_stripper.transform_string(tokens[0]) + + _commasepitem = ( + Combine( + OneOrMore( + ~Literal(",") + + ~LineEnd() + + Word(printables, exclude_chars=",") + + Opt(White(" \t") + ~FollowedBy(LineEnd() | ",")) + ) + ) + .streamline() + .set_name("commaItem") + ) + comma_separated_list = DelimitedList( + Opt(quoted_string.copy() | _commasepitem, default="") + ).set_name("comma separated list") + """Predefined expression of 1 or more printable words or quoted strings, separated by commas.""" + + upcase_tokens = staticmethod(token_map(lambda t: t.upper())) + """Parse action to convert tokens to upper case.""" + + downcase_tokens = staticmethod(token_map(lambda t: t.lower())) + """Parse action to convert tokens to lower case.""" + + # fmt: off + url = Regex( + # https://mathiasbynens.be/demo/url-regex + # https://gist.github.com/dperini/729294 + r"(?P" + + # protocol identifier (optional) + # short syntax // still required + r"(?:(?:(?Phttps?|ftp):)?\/\/)" + + # user:pass BasicAuth (optional) + r"(?:(?P\S+(?::\S*)?)@)?" + + r"(?P" + + # IP address exclusion + # private & local networks + r"(?!(?:10|127)(?:\.\d{1,3}){3})" + + r"(?!(?:169\.254|192\.168)(?:\.\d{1,3}){2})" + + r"(?!172\.(?:1[6-9]|2\d|3[0-1])(?:\.\d{1,3}){2})" + + # IP address dotted notation octets + # excludes loopback network 0.0.0.0 + # excludes reserved space >= 224.0.0.0 + # excludes network & broadcast addresses + # (first & last IP address of each class) + r"(?:[1-9]\d?|1\d\d|2[01]\d|22[0-3])" + + r"(?:\.(?:1?\d{1,2}|2[0-4]\d|25[0-5])){2}" + + r"(?:\.(?:[1-9]\d?|1\d\d|2[0-4]\d|25[0-4]))" + + r"|" + + # host & domain names, may end with dot + # can be replaced by a shortest alternative + # (?![-_])(?:[-\w\u00a1-\uffff]{0,63}[^-_]\.)+ + r"(?:" + + r"(?:" + + r"[a-z0-9\u00a1-\uffff]" + + r"[a-z0-9\u00a1-\uffff_-]{0,62}" + + r")?" + + r"[a-z0-9\u00a1-\uffff]\." + + r")+" + + # TLD identifier name, may end with dot + r"(?:[a-z\u00a1-\uffff]{2,}\.?)" + + r")" + + # port number (optional) + r"(:(?P\d{2,5}))?" + + # resource path (optional) + r"(?P\/[^?# ]*)?" + + # query string (optional) + r"(\?(?P[^#]*))?" + + # fragment (optional) + r"(#(?P\S*))?" + + r")" + ).set_name("url") + """URL (http/https/ftp scheme)""" + # fmt: on + + # pre-PEP8 compatibility names + # fmt: off + convertToInteger = staticmethod(replaced_by_pep8("convertToInteger", convert_to_integer)) + convertToFloat = staticmethod(replaced_by_pep8("convertToFloat", convert_to_float)) + convertToDate = staticmethod(replaced_by_pep8("convertToDate", convert_to_date)) + convertToDatetime = staticmethod(replaced_by_pep8("convertToDatetime", convert_to_datetime)) + stripHTMLTags = staticmethod(replaced_by_pep8("stripHTMLTags", strip_html_tags)) + upcaseTokens = staticmethod(replaced_by_pep8("upcaseTokens", upcase_tokens)) + downcaseTokens = staticmethod(replaced_by_pep8("downcaseTokens", downcase_tokens)) + # fmt: on + + +_builtin_exprs = [ + v for v in vars(pyparsing_common).values() if isinstance(v, ParserElement) +] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/core.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/core.py new file mode 100644 index 0000000000000000000000000000000000000000..86be949ad473a8ab536a15e9b4fdaabd35f225f6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/core.py @@ -0,0 +1,6299 @@ +# +# core.py +# +from __future__ import annotations + +import collections.abc +from collections import deque +import os +import typing +from typing import ( + Any, + Callable, + Generator, + NamedTuple, + Sequence, + TextIO, + Union, + cast, +) +from abc import ABC, abstractmethod +from enum import Enum +import string +import copy +import warnings +import re +import sys +from collections.abc import Iterable +import traceback +import types +from operator import itemgetter +from functools import wraps +from threading import RLock +from pathlib import Path + +from .util import ( + _FifoCache, + _UnboundedCache, + __config_flags, + _collapse_string_to_ranges, + _escape_regex_range_chars, + _flatten, + LRUMemo as _LRUMemo, + UnboundedMemo as _UnboundedMemo, + replaced_by_pep8, +) +from .exceptions import * +from .actions import * +from .results import ParseResults, _ParseResultsWithOffset +from .unicode import pyparsing_unicode + +_MAX_INT = sys.maxsize +str_type: tuple[type, ...] = (str, bytes) + +# +# Copyright (c) 2003-2022 Paul T. McGuire +# +# Permission is hereby granted, free of charge, to any person obtaining +# a copy of this software and associated documentation files (the +# "Software"), to deal in the Software without restriction, including +# without limitation the rights to use, copy, modify, merge, publish, +# distribute, sublicense, and/or sell copies of the Software, and to +# permit persons to whom the Software is furnished to do so, subject to +# the following conditions: +# +# The above copyright notice and this permission notice shall be +# included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +# + +from functools import cached_property + + +class __compat__(__config_flags): + """ + A cross-version compatibility configuration for pyparsing features that will be + released in a future version. By setting values in this configuration to True, + those features can be enabled in prior versions for compatibility development + and testing. + + - ``collect_all_And_tokens`` - flag to enable fix for Issue #63 that fixes erroneous grouping + of results names when an :class:`And` expression is nested within an :class:`Or` or :class:`MatchFirst`; + maintained for compatibility, but setting to ``False`` no longer restores pre-2.3.1 + behavior + """ + + _type_desc = "compatibility" + + collect_all_And_tokens = True + + _all_names = [__ for __ in locals() if not __.startswith("_")] + _fixed_names = """ + collect_all_And_tokens + """.split() + + +class __diag__(__config_flags): + _type_desc = "diagnostic" + + warn_multiple_tokens_in_named_alternation = False + warn_ungrouped_named_tokens_in_collection = False + warn_name_set_on_empty_Forward = False + warn_on_parse_using_empty_Forward = False + warn_on_assignment_to_Forward = False + warn_on_multiple_string_args_to_oneof = False + warn_on_match_first_with_lshift_operator = False + enable_debug_on_named_expressions = False + + _all_names = [__ for __ in locals() if not __.startswith("_")] + _warning_names = [name for name in _all_names if name.startswith("warn")] + _debug_names = [name for name in _all_names if name.startswith("enable_debug")] + + @classmethod + def enable_all_warnings(cls) -> None: + for name in cls._warning_names: + cls.enable(name) + + +class Diagnostics(Enum): + """ + Diagnostic configuration (all default to disabled) + + - ``warn_multiple_tokens_in_named_alternation`` - flag to enable warnings when a results + name is defined on a :class:`MatchFirst` or :class:`Or` expression with one or more :class:`And` subexpressions + - ``warn_ungrouped_named_tokens_in_collection`` - flag to enable warnings when a results + name is defined on a containing expression with ungrouped subexpressions that also + have results names + - ``warn_name_set_on_empty_Forward`` - flag to enable warnings when a :class:`Forward` is defined + with a results name, but has no contents defined + - ``warn_on_parse_using_empty_Forward`` - flag to enable warnings when a :class:`Forward` is + defined in a grammar but has never had an expression attached to it + - ``warn_on_assignment_to_Forward`` - flag to enable warnings when a :class:`Forward` is defined + but is overwritten by assigning using ``'='`` instead of ``'<<='`` or ``'<<'`` + - ``warn_on_multiple_string_args_to_oneof`` - flag to enable warnings when :class:`one_of` is + incorrectly called with multiple str arguments + - ``enable_debug_on_named_expressions`` - flag to auto-enable debug on all subsequent + calls to :class:`ParserElement.set_name` + + Diagnostics are enabled/disabled by calling :class:`enable_diag` and :class:`disable_diag`. + All warnings can be enabled by calling :class:`enable_all_warnings`. + """ + + warn_multiple_tokens_in_named_alternation = 0 + warn_ungrouped_named_tokens_in_collection = 1 + warn_name_set_on_empty_Forward = 2 + warn_on_parse_using_empty_Forward = 3 + warn_on_assignment_to_Forward = 4 + warn_on_multiple_string_args_to_oneof = 5 + warn_on_match_first_with_lshift_operator = 6 + enable_debug_on_named_expressions = 7 + + +def enable_diag(diag_enum: Diagnostics) -> None: + """ + Enable a global pyparsing diagnostic flag (see :class:`Diagnostics`). + """ + __diag__.enable(diag_enum.name) + + +def disable_diag(diag_enum: Diagnostics) -> None: + """ + Disable a global pyparsing diagnostic flag (see :class:`Diagnostics`). + """ + __diag__.disable(diag_enum.name) + + +def enable_all_warnings() -> None: + """ + Enable all global pyparsing diagnostic warnings (see :class:`Diagnostics`). + """ + __diag__.enable_all_warnings() + + +# hide abstract class +del __config_flags + + +def _should_enable_warnings( + cmd_line_warn_options: typing.Iterable[str], warn_env_var: typing.Optional[str] +) -> bool: + enable = bool(warn_env_var) + for warn_opt in cmd_line_warn_options: + w_action, w_message, w_category, w_module, w_line = (warn_opt + "::::").split( + ":" + )[:5] + if not w_action.lower().startswith("i") and ( + not (w_message or w_category or w_module) or w_module == "pyparsing" + ): + enable = True + elif w_action.lower().startswith("i") and w_module in ("pyparsing", ""): + enable = False + return enable + + +if _should_enable_warnings( + sys.warnoptions, os.environ.get("PYPARSINGENABLEALLWARNINGS") +): + enable_all_warnings() + + +# build list of single arg builtins, that can be used as parse actions +# fmt: off +_single_arg_builtins = { + sum, len, sorted, reversed, list, tuple, set, any, all, min, max +} +# fmt: on + +_generatorType = types.GeneratorType +ParseImplReturnType = tuple[int, Any] +PostParseReturnType = Union[ParseResults, Sequence[ParseResults]] + +ParseCondition = Union[ + Callable[[], bool], + Callable[[ParseResults], bool], + Callable[[int, ParseResults], bool], + Callable[[str, int, ParseResults], bool], +] +ParseFailAction = Callable[[str, int, "ParserElement", Exception], None] +DebugStartAction = Callable[[str, int, "ParserElement", bool], None] +DebugSuccessAction = Callable[ + [str, int, int, "ParserElement", ParseResults, bool], None +] +DebugExceptionAction = Callable[[str, int, "ParserElement", Exception, bool], None] + + +alphas: str = string.ascii_uppercase + string.ascii_lowercase +identchars: str = pyparsing_unicode.Latin1.identchars +identbodychars: str = pyparsing_unicode.Latin1.identbodychars +nums: str = "0123456789" +hexnums: str = nums + "ABCDEFabcdef" +alphanums: str = alphas + nums +printables: str = "".join([c for c in string.printable if c not in string.whitespace]) + + +class _ParseActionIndexError(Exception): + """ + Internal wrapper around IndexError so that IndexErrors raised inside + parse actions aren't misinterpreted as IndexErrors raised inside + ParserElement parseImpl methods. + """ + + def __init__(self, msg: str, exc: BaseException) -> None: + self.msg: str = msg + self.exc: BaseException = exc + + +_trim_arity_call_line: traceback.StackSummary = None # type: ignore[assignment] +pa_call_line_synth = () + + +def _trim_arity(func, max_limit=3): + """decorator to trim function calls to match the arity of the target""" + global _trim_arity_call_line, pa_call_line_synth + + if func in _single_arg_builtins: + return lambda s, l, t: func(t) + + limit = 0 + found_arity = False + + # synthesize what would be returned by traceback.extract_stack at the call to + # user's parse action 'func', so that we don't incur call penalty at parse time + + # fmt: off + LINE_DIFF = 9 + # IF ANY CODE CHANGES, EVEN JUST COMMENTS OR BLANK LINES, BETWEEN THE NEXT LINE AND + # THE CALL TO FUNC INSIDE WRAPPER, LINE_DIFF MUST BE MODIFIED!!!! + _trim_arity_call_line = _trim_arity_call_line or traceback.extract_stack(limit=2)[-1] + pa_call_line_synth = pa_call_line_synth or (_trim_arity_call_line[0], _trim_arity_call_line[1] + LINE_DIFF) + + def wrapper(*args): + nonlocal found_arity, limit + if found_arity: + return func(*args[limit:]) + while 1: + try: + ret = func(*args[limit:]) + found_arity = True + return ret + except TypeError as te: + # re-raise TypeErrors if they did not come from our arity testing + if found_arity: + raise + else: + tb = te.__traceback__ + frames = traceback.extract_tb(tb, limit=2) + frame_summary = frames[-1] + trim_arity_type_error = ( + [frame_summary[:2]][-1][:2] == pa_call_line_synth + ) + del tb + + if trim_arity_type_error: + if limit < max_limit: + limit += 1 + continue + + raise + except IndexError as ie: + # wrap IndexErrors inside a _ParseActionIndexError + raise _ParseActionIndexError( + "IndexError raised in parse action", ie + ).with_traceback(None) + # fmt: on + + # copy func name to wrapper for sensible debug output + # (can't use functools.wraps, since that messes with function signature) + func_name = getattr(func, "__name__", getattr(func, "__class__").__name__) + wrapper.__name__ = func_name + wrapper.__doc__ = func.__doc__ + + return wrapper + + +def condition_as_parse_action( + fn: ParseCondition, message: typing.Optional[str] = None, fatal: bool = False +) -> ParseAction: + """ + Function to convert a simple predicate function that returns ``True`` or ``False`` + into a parse action. Can be used in places when a parse action is required + and :class:`ParserElement.add_condition` cannot be used (such as when adding a condition + to an operator level in :class:`infix_notation`). + + Optional keyword arguments: + + - ``message`` - define a custom message to be used in the raised exception + - ``fatal`` - if True, will raise :class:`ParseFatalException` to stop parsing immediately; + otherwise will raise :class:`ParseException` + + """ + msg = message if message is not None else "failed user-defined condition" + exc_type = ParseFatalException if fatal else ParseException + fn = _trim_arity(fn) + + @wraps(fn) + def pa(s, l, t): + if not bool(fn(s, l, t)): + raise exc_type(s, l, msg) + + return pa + + +def _default_start_debug_action( + instring: str, loc: int, expr: ParserElement, cache_hit: bool = False +): + cache_hit_str = "*" if cache_hit else "" + print( + ( + f"{cache_hit_str}Match {expr} at loc {loc}({lineno(loc, instring)},{col(loc, instring)})\n" + f" {line(loc, instring)}\n" + f" {'^':>{col(loc, instring)}}" + ) + ) + + +def _default_success_debug_action( + instring: str, + startloc: int, + endloc: int, + expr: ParserElement, + toks: ParseResults, + cache_hit: bool = False, +): + cache_hit_str = "*" if cache_hit else "" + print(f"{cache_hit_str}Matched {expr} -> {toks.as_list()}") + + +def _default_exception_debug_action( + instring: str, + loc: int, + expr: ParserElement, + exc: Exception, + cache_hit: bool = False, +): + cache_hit_str = "*" if cache_hit else "" + print(f"{cache_hit_str}Match {expr} failed, {type(exc).__name__} raised: {exc}") + + +def null_debug_action(*args): + """'Do-nothing' debug action, to suppress debugging output during parsing.""" + + +class ParserElement(ABC): + """Abstract base level parser element class.""" + + DEFAULT_WHITE_CHARS: str = " \n\t\r" + verbose_stacktrace: bool = False + _literalStringClass: type = None # type: ignore[assignment] + + @staticmethod + def set_default_whitespace_chars(chars: str) -> None: + r""" + Overrides the default whitespace chars + + Example:: + + # default whitespace chars are space, and newline + Word(alphas)[1, ...].parse_string("abc def\nghi jkl") # -> ['abc', 'def', 'ghi', 'jkl'] + + # change to just treat newline as significant + ParserElement.set_default_whitespace_chars(" \t") + Word(alphas)[1, ...].parse_string("abc def\nghi jkl") # -> ['abc', 'def'] + """ + ParserElement.DEFAULT_WHITE_CHARS = chars + + # update whitespace all parse expressions defined in this module + for expr in _builtin_exprs: + if expr.copyDefaultWhiteChars: + expr.whiteChars = set(chars) + + @staticmethod + def inline_literals_using(cls: type) -> None: + """ + Set class to be used for inclusion of string literals into a parser. + + Example:: + + # default literal class used is Literal + integer = Word(nums) + date_str = integer("year") + '/' + integer("month") + '/' + integer("day") + + date_str.parse_string("1999/12/31") # -> ['1999', '/', '12', '/', '31'] + + + # change to Suppress + ParserElement.inline_literals_using(Suppress) + date_str = integer("year") + '/' + integer("month") + '/' + integer("day") + + date_str.parse_string("1999/12/31") # -> ['1999', '12', '31'] + """ + ParserElement._literalStringClass = cls + + @classmethod + def using_each(cls, seq, **class_kwargs): + """ + Yields a sequence of ``class(obj, **class_kwargs)`` for obj in seq. + + Example:: + + LPAR, RPAR, LBRACE, RBRACE, SEMI = Suppress.using_each("(){};") + + """ + yield from (cls(obj, **class_kwargs) for obj in seq) + + class DebugActions(NamedTuple): + debug_try: typing.Optional[DebugStartAction] + debug_match: typing.Optional[DebugSuccessAction] + debug_fail: typing.Optional[DebugExceptionAction] + + def __init__(self, savelist: bool = False) -> None: + self.parseAction: list[ParseAction] = list() + self.failAction: typing.Optional[ParseFailAction] = None + self.customName: str = None # type: ignore[assignment] + self._defaultName: typing.Optional[str] = None + self.resultsName: str = None # type: ignore[assignment] + self.saveAsList = savelist + self.skipWhitespace = True + self.whiteChars = set(ParserElement.DEFAULT_WHITE_CHARS) + self.copyDefaultWhiteChars = True + # used when checking for left-recursion + self._may_return_empty = False + self.keepTabs = False + self.ignoreExprs: list[ParserElement] = list() + self.debug = False + self.streamlined = False + # optimize exception handling for subclasses that don't advance parse index + self.mayIndexError = True + self.errmsg: Union[str, None] = "" + # mark results names as modal (report only last) or cumulative (list all) + self.modalResults = True + # custom debug actions + self.debugActions = self.DebugActions(None, None, None) + # avoid redundant calls to preParse + self.callPreparse = True + self.callDuringTry = False + self.suppress_warnings_: list[Diagnostics] = [] + self.show_in_diagram = True + + @property + def mayReturnEmpty(self): + return self._may_return_empty + + @mayReturnEmpty.setter + def mayReturnEmpty(self, value): + self._may_return_empty = value + + def suppress_warning(self, warning_type: Diagnostics) -> ParserElement: + """ + Suppress warnings emitted for a particular diagnostic on this expression. + + Example:: + + base = pp.Forward() + base.suppress_warning(Diagnostics.warn_on_parse_using_empty_Forward) + + # statement would normally raise a warning, but is now suppressed + print(base.parse_string("x")) + + """ + self.suppress_warnings_.append(warning_type) + return self + + def visit_all(self): + """General-purpose method to yield all expressions and sub-expressions + in a grammar. Typically just for internal use. + """ + to_visit = deque([self]) + seen = set() + while to_visit: + cur = to_visit.popleft() + + # guard against looping forever through recursive grammars + if cur in seen: + continue + seen.add(cur) + + to_visit.extend(cur.recurse()) + yield cur + + def copy(self) -> ParserElement: + """ + Make a copy of this :class:`ParserElement`. Useful for defining + different parse actions for the same parsing pattern, using copies of + the original parse element. + + Example:: + + integer = Word(nums).set_parse_action(lambda toks: int(toks[0])) + integerK = integer.copy().add_parse_action(lambda toks: toks[0] * 1024) + Suppress("K") + integerM = integer.copy().add_parse_action(lambda toks: toks[0] * 1024 * 1024) + Suppress("M") + + print((integerK | integerM | integer)[1, ...].parse_string("5K 100 640K 256M")) + + prints:: + + [5120, 100, 655360, 268435456] + + Equivalent form of ``expr.copy()`` is just ``expr()``:: + + integerM = integer().add_parse_action(lambda toks: toks[0] * 1024 * 1024) + Suppress("M") + """ + cpy = copy.copy(self) + cpy.parseAction = self.parseAction[:] + cpy.ignoreExprs = self.ignoreExprs[:] + if self.copyDefaultWhiteChars: + cpy.whiteChars = set(ParserElement.DEFAULT_WHITE_CHARS) + return cpy + + def set_results_name( + self, name: str, list_all_matches: bool = False, *, listAllMatches: bool = False + ) -> ParserElement: + """ + Define name for referencing matching tokens as a nested attribute + of the returned parse results. + + Normally, results names are assigned as you would assign keys in a dict: + any existing value is overwritten by later values. If it is necessary to + keep all values captured for a particular results name, call ``set_results_name`` + with ``list_all_matches`` = True. + + NOTE: ``set_results_name`` returns a *copy* of the original :class:`ParserElement` object; + this is so that the client can define a basic element, such as an + integer, and reference it in multiple places with different names. + + You can also set results names using the abbreviated syntax, + ``expr("name")`` in place of ``expr.set_results_name("name")`` + - see :class:`__call__`. If ``list_all_matches`` is required, use + ``expr("name*")``. + + Example:: + + integer = Word(nums) + date_str = (integer.set_results_name("year") + '/' + + integer.set_results_name("month") + '/' + + integer.set_results_name("day")) + + # equivalent form: + date_str = integer("year") + '/' + integer("month") + '/' + integer("day") + """ + listAllMatches = listAllMatches or list_all_matches + return self._setResultsName(name, listAllMatches) + + def _setResultsName(self, name, list_all_matches=False) -> ParserElement: + if name is None: + return self + newself = self.copy() + if name.endswith("*"): + name = name[:-1] + list_all_matches = True + newself.resultsName = name + newself.modalResults = not list_all_matches + return newself + + def set_break(self, break_flag: bool = True) -> ParserElement: + """ + Method to invoke the Python pdb debugger when this element is + about to be parsed. Set ``break_flag`` to ``True`` to enable, ``False`` to + disable. + """ + if break_flag: + _parseMethod = self._parse + + def breaker(instring, loc, do_actions=True, callPreParse=True): + # this call to breakpoint() is intentional, not a checkin error + breakpoint() + return _parseMethod(instring, loc, do_actions, callPreParse) + + breaker._originalParseMethod = _parseMethod # type: ignore [attr-defined] + self._parse = breaker # type: ignore [method-assign] + elif hasattr(self._parse, "_originalParseMethod"): + self._parse = self._parse._originalParseMethod # type: ignore [method-assign] + return self + + def set_parse_action(self, *fns: ParseAction, **kwargs: Any) -> ParserElement: + """ + Define one or more actions to perform when successfully matching parse element definition. + + Parse actions can be called to perform data conversions, do extra validation, + update external data structures, or enhance or replace the parsed tokens. + Each parse action ``fn`` is a callable method with 0-3 arguments, called as + ``fn(s, loc, toks)`` , ``fn(loc, toks)`` , ``fn(toks)`` , or just ``fn()`` , where: + + - ``s`` = the original string being parsed (see note below) + - ``loc`` = the location of the matching substring + - ``toks`` = a list of the matched tokens, packaged as a :class:`ParseResults` object + + The parsed tokens are passed to the parse action as ParseResults. They can be + modified in place using list-style append, extend, and pop operations to update + the parsed list elements; and with dictionary-style item set and del operations + to add, update, or remove any named results. If the tokens are modified in place, + it is not necessary to return them with a return statement. + + Parse actions can also completely replace the given tokens, with another ``ParseResults`` + object, or with some entirely different object (common for parse actions that perform data + conversions). A convenient way to build a new parse result is to define the values + using a dict, and then create the return value using :class:`ParseResults.from_dict`. + + If None is passed as the ``fn`` parse action, all previously added parse actions for this + expression are cleared. + + Optional keyword arguments: + + - ``call_during_try`` = (default= ``False``) indicate if parse action should be run during + lookaheads and alternate testing. For parse actions that have side effects, it is + important to only call the parse action once it is determined that it is being + called as part of a successful parse. For parse actions that perform additional + validation, then call_during_try should be passed as True, so that the validation + code is included in the preliminary "try" parses. + + Note: the default parsing behavior is to expand tabs in the input string + before starting the parsing process. See :class:`parse_string` for more + information on parsing strings containing ```` s, and suggested + methods to maintain a consistent view of the parsed string, the parse + location, and line and column positions within the parsed string. + + Example:: + + # parse dates in the form YYYY/MM/DD + + # use parse action to convert toks from str to int at parse time + def convert_to_int(toks): + return int(toks[0]) + + # use a parse action to verify that the date is a valid date + def is_valid_date(instring, loc, toks): + from datetime import date + year, month, day = toks[::2] + try: + date(year, month, day) + except ValueError: + raise ParseException(instring, loc, "invalid date given") + + integer = Word(nums) + date_str = integer + '/' + integer + '/' + integer + + # add parse actions + integer.set_parse_action(convert_to_int) + date_str.set_parse_action(is_valid_date) + + # note that integer fields are now ints, not strings + date_str.run_tests(''' + # successful parse - note that integer fields were converted to ints + 1999/12/31 + + # fail - invalid date + 1999/13/31 + ''') + """ + if list(fns) == [None]: + self.parseAction.clear() + return self + + if not all(callable(fn) for fn in fns): + raise TypeError("parse actions must be callable") + self.parseAction[:] = [_trim_arity(fn) for fn in fns] + self.callDuringTry = kwargs.get( + "call_during_try", kwargs.get("callDuringTry", False) + ) + + return self + + def add_parse_action(self, *fns: ParseAction, **kwargs: Any) -> ParserElement: + """ + Add one or more parse actions to expression's list of parse actions. See :class:`set_parse_action`. + + See examples in :class:`copy`. + """ + self.parseAction += [_trim_arity(fn) for fn in fns] + self.callDuringTry = self.callDuringTry or kwargs.get( + "call_during_try", kwargs.get("callDuringTry", False) + ) + return self + + def add_condition(self, *fns: ParseCondition, **kwargs: Any) -> ParserElement: + """Add a boolean predicate function to expression's list of parse actions. See + :class:`set_parse_action` for function call signatures. Unlike ``set_parse_action``, + functions passed to ``add_condition`` need to return boolean success/fail of the condition. + + Optional keyword arguments: + + - ``message`` = define a custom message to be used in the raised exception + - ``fatal`` = if True, will raise ParseFatalException to stop parsing immediately; otherwise will raise + ParseException + - ``call_during_try`` = boolean to indicate if this method should be called during internal tryParse calls, + default=False + + Example:: + + integer = Word(nums).set_parse_action(lambda toks: int(toks[0])) + year_int = integer.copy() + year_int.add_condition(lambda toks: toks[0] >= 2000, message="Only support years 2000 and later") + date_str = year_int + '/' + integer + '/' + integer + + result = date_str.parse_string("1999/12/31") # -> Exception: Only support years 2000 and later (at char 0), + (line:1, col:1) + """ + for fn in fns: + self.parseAction.append( + condition_as_parse_action( + fn, + message=str(kwargs.get("message")), + fatal=bool(kwargs.get("fatal", False)), + ) + ) + + self.callDuringTry = self.callDuringTry or kwargs.get( + "call_during_try", kwargs.get("callDuringTry", False) + ) + return self + + def set_fail_action(self, fn: ParseFailAction) -> ParserElement: + """ + Define action to perform if parsing fails at this expression. + Fail acton fn is a callable function that takes the arguments + ``fn(s, loc, expr, err)`` where: + + - ``s`` = string being parsed + - ``loc`` = location where expression match was attempted and failed + - ``expr`` = the parse expression that failed + - ``err`` = the exception thrown + + The function returns no value. It may throw :class:`ParseFatalException` + if it is desired to stop parsing immediately.""" + self.failAction = fn + return self + + def _skipIgnorables(self, instring: str, loc: int) -> int: + if not self.ignoreExprs: + return loc + exprsFound = True + ignore_expr_fns = [e._parse for e in self.ignoreExprs] + last_loc = loc + while exprsFound: + exprsFound = False + for ignore_fn in ignore_expr_fns: + try: + while 1: + loc, dummy = ignore_fn(instring, loc) + exprsFound = True + except ParseException: + pass + # check if all ignore exprs matched but didn't actually advance the parse location + if loc == last_loc: + break + last_loc = loc + return loc + + def preParse(self, instring: str, loc: int) -> int: + if self.ignoreExprs: + loc = self._skipIgnorables(instring, loc) + + if self.skipWhitespace: + instrlen = len(instring) + white_chars = self.whiteChars + while loc < instrlen and instring[loc] in white_chars: + loc += 1 + + return loc + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + return loc, [] + + def postParse(self, instring, loc, tokenlist): + return tokenlist + + # @profile + def _parseNoCache( + self, instring, loc, do_actions=True, callPreParse=True + ) -> tuple[int, ParseResults]: + debugging = self.debug # and do_actions) + len_instring = len(instring) + + if debugging or self.failAction: + # print("Match {} at loc {}({}, {})".format(self, loc, lineno(loc, instring), col(loc, instring))) + try: + if callPreParse and self.callPreparse: + pre_loc = self.preParse(instring, loc) + else: + pre_loc = loc + tokens_start = pre_loc + if self.debugActions.debug_try: + self.debugActions.debug_try(instring, tokens_start, self, False) + if self.mayIndexError or pre_loc >= len_instring: + try: + loc, tokens = self.parseImpl(instring, pre_loc, do_actions) + except IndexError: + raise ParseException(instring, len_instring, self.errmsg, self) + else: + loc, tokens = self.parseImpl(instring, pre_loc, do_actions) + except Exception as err: + # print("Exception raised:", err) + if self.debugActions.debug_fail: + self.debugActions.debug_fail( + instring, tokens_start, self, err, False + ) + if self.failAction: + self.failAction(instring, tokens_start, self, err) + raise + else: + if callPreParse and self.callPreparse: + pre_loc = self.preParse(instring, loc) + else: + pre_loc = loc + tokens_start = pre_loc + if self.mayIndexError or pre_loc >= len_instring: + try: + loc, tokens = self.parseImpl(instring, pre_loc, do_actions) + except IndexError: + raise ParseException(instring, len_instring, self.errmsg, self) + else: + loc, tokens = self.parseImpl(instring, pre_loc, do_actions) + + tokens = self.postParse(instring, loc, tokens) + + ret_tokens = ParseResults( + tokens, self.resultsName, asList=self.saveAsList, modal=self.modalResults + ) + if self.parseAction and (do_actions or self.callDuringTry): + if debugging: + try: + for fn in self.parseAction: + try: + tokens = fn(instring, tokens_start, ret_tokens) # type: ignore [call-arg, arg-type] + except IndexError as parse_action_exc: + exc = ParseException("exception raised in parse action") + raise exc from parse_action_exc + + if tokens is not None and tokens is not ret_tokens: + ret_tokens = ParseResults( + tokens, + self.resultsName, + asList=self.saveAsList + and isinstance(tokens, (ParseResults, list)), + modal=self.modalResults, + ) + except Exception as err: + # print "Exception raised in user parse action:", err + if self.debugActions.debug_fail: + self.debugActions.debug_fail( + instring, tokens_start, self, err, False + ) + raise + else: + for fn in self.parseAction: + try: + tokens = fn(instring, tokens_start, ret_tokens) # type: ignore [call-arg, arg-type] + except IndexError as parse_action_exc: + exc = ParseException("exception raised in parse action") + raise exc from parse_action_exc + + if tokens is not None and tokens is not ret_tokens: + ret_tokens = ParseResults( + tokens, + self.resultsName, + asList=self.saveAsList + and isinstance(tokens, (ParseResults, list)), + modal=self.modalResults, + ) + if debugging: + # print("Matched", self, "->", ret_tokens.as_list()) + if self.debugActions.debug_match: + self.debugActions.debug_match( + instring, tokens_start, loc, self, ret_tokens, False + ) + + return loc, ret_tokens + + def try_parse( + self, + instring: str, + loc: int, + *, + raise_fatal: bool = False, + do_actions: bool = False, + ) -> int: + try: + return self._parse(instring, loc, do_actions=do_actions)[0] + except ParseFatalException: + if raise_fatal: + raise + raise ParseException(instring, loc, self.errmsg, self) + + def can_parse_next(self, instring: str, loc: int, do_actions: bool = False) -> bool: + try: + self.try_parse(instring, loc, do_actions=do_actions) + except (ParseException, IndexError): + return False + else: + return True + + # cache for left-recursion in Forward references + recursion_lock = RLock() + recursion_memos: collections.abc.MutableMapping[ + tuple[int, Forward, bool], tuple[int, Union[ParseResults, Exception]] + ] = {} + + class _CacheType(typing.Protocol): + """ + Class to be used for packrat and left-recursion cacheing of results + and exceptions. + """ + + not_in_cache: bool + + def get(self, *args) -> typing.Any: ... + + def set(self, *args) -> None: ... + + def clear(self) -> None: ... + + class NullCache(dict): + """ + A null cache type for initialization of the packrat_cache class variable. + If/when enable_packrat() is called, this null cache will be replaced by a + proper _CacheType class instance. + """ + + not_in_cache: bool = True + + def get(self, *args) -> typing.Any: ... + + def set(self, *args) -> None: ... + + def clear(self) -> None: ... + + # class-level argument cache for optimizing repeated calls when backtracking + # through recursive expressions + packrat_cache: _CacheType = NullCache() + packrat_cache_lock = RLock() + packrat_cache_stats = [0, 0] + + # this method gets repeatedly called during backtracking with the same arguments - + # we can cache these arguments and save ourselves the trouble of re-parsing the contained expression + def _parseCache( + self, instring, loc, do_actions=True, callPreParse=True + ) -> tuple[int, ParseResults]: + HIT, MISS = 0, 1 + lookup = (self, instring, loc, callPreParse, do_actions) + with ParserElement.packrat_cache_lock: + cache = ParserElement.packrat_cache + value = cache.get(lookup) + if value is cache.not_in_cache: + ParserElement.packrat_cache_stats[MISS] += 1 + try: + value = self._parseNoCache(instring, loc, do_actions, callPreParse) + except ParseBaseException as pe: + # cache a copy of the exception, without the traceback + cache.set(lookup, pe.__class__(*pe.args)) + raise + else: + cache.set(lookup, (value[0], value[1].copy(), loc)) + return value + else: + ParserElement.packrat_cache_stats[HIT] += 1 + if self.debug and self.debugActions.debug_try: + try: + self.debugActions.debug_try(instring, loc, self, cache_hit=True) # type: ignore [call-arg] + except TypeError: + pass + if isinstance(value, Exception): + if self.debug and self.debugActions.debug_fail: + try: + self.debugActions.debug_fail( + instring, loc, self, value, cache_hit=True # type: ignore [call-arg] + ) + except TypeError: + pass + raise value + + value = cast(tuple[int, ParseResults, int], value) + loc_, result, endloc = value[0], value[1].copy(), value[2] + if self.debug and self.debugActions.debug_match: + try: + self.debugActions.debug_match( + instring, loc_, endloc, self, result, cache_hit=True # type: ignore [call-arg] + ) + except TypeError: + pass + + return loc_, result + + _parse = _parseNoCache + + @staticmethod + def reset_cache() -> None: + ParserElement.packrat_cache.clear() + ParserElement.packrat_cache_stats[:] = [0] * len( + ParserElement.packrat_cache_stats + ) + ParserElement.recursion_memos.clear() + + _packratEnabled = False + _left_recursion_enabled = False + + @staticmethod + def disable_memoization() -> None: + """ + Disables active Packrat or Left Recursion parsing and their memoization + + This method also works if neither Packrat nor Left Recursion are enabled. + This makes it safe to call before activating Packrat nor Left Recursion + to clear any previous settings. + """ + ParserElement.reset_cache() + ParserElement._left_recursion_enabled = False + ParserElement._packratEnabled = False + ParserElement._parse = ParserElement._parseNoCache + + @staticmethod + def enable_left_recursion( + cache_size_limit: typing.Optional[int] = None, *, force=False + ) -> None: + """ + Enables "bounded recursion" parsing, which allows for both direct and indirect + left-recursion. During parsing, left-recursive :class:`Forward` elements are + repeatedly matched with a fixed recursion depth that is gradually increased + until finding the longest match. + + Example:: + + import pyparsing as pp + pp.ParserElement.enable_left_recursion() + + E = pp.Forward("E") + num = pp.Word(pp.nums) + # match `num`, or `num '+' num`, or `num '+' num '+' num`, ... + E <<= E + '+' - num | num + + print(E.parse_string("1+2+3")) + + Recursion search naturally memoizes matches of ``Forward`` elements and may + thus skip reevaluation of parse actions during backtracking. This may break + programs with parse actions which rely on strict ordering of side-effects. + + Parameters: + + - ``cache_size_limit`` - (default=``None``) - memoize at most this many + ``Forward`` elements during matching; if ``None`` (the default), + memoize all ``Forward`` elements. + + Bounded Recursion parsing works similar but not identical to Packrat parsing, + thus the two cannot be used together. Use ``force=True`` to disable any + previous, conflicting settings. + """ + if force: + ParserElement.disable_memoization() + elif ParserElement._packratEnabled: + raise RuntimeError("Packrat and Bounded Recursion are not compatible") + if cache_size_limit is None: + ParserElement.recursion_memos = _UnboundedMemo() + elif cache_size_limit > 0: + ParserElement.recursion_memos = _LRUMemo(capacity=cache_size_limit) # type: ignore[assignment] + else: + raise NotImplementedError(f"Memo size of {cache_size_limit}") + ParserElement._left_recursion_enabled = True + + @staticmethod + def enable_packrat( + cache_size_limit: Union[int, None] = 128, *, force: bool = False + ) -> None: + """ + Enables "packrat" parsing, which adds memoizing to the parsing logic. + Repeated parse attempts at the same string location (which happens + often in many complex grammars) can immediately return a cached value, + instead of re-executing parsing/validating code. Memoizing is done of + both valid results and parsing exceptions. + + Parameters: + + - ``cache_size_limit`` - (default= ``128``) - if an integer value is provided + will limit the size of the packrat cache; if None is passed, then + the cache size will be unbounded; if 0 is passed, the cache will + be effectively disabled. + + This speedup may break existing programs that use parse actions that + have side-effects. For this reason, packrat parsing is disabled when + you first import pyparsing. To activate the packrat feature, your + program must call the class method :class:`ParserElement.enable_packrat`. + For best results, call ``enable_packrat()`` immediately after + importing pyparsing. + + Example:: + + import pyparsing + pyparsing.ParserElement.enable_packrat() + + Packrat parsing works similar but not identical to Bounded Recursion parsing, + thus the two cannot be used together. Use ``force=True`` to disable any + previous, conflicting settings. + """ + if force: + ParserElement.disable_memoization() + elif ParserElement._left_recursion_enabled: + raise RuntimeError("Packrat and Bounded Recursion are not compatible") + + if ParserElement._packratEnabled: + return + + ParserElement._packratEnabled = True + if cache_size_limit is None: + ParserElement.packrat_cache = _UnboundedCache() + else: + ParserElement.packrat_cache = _FifoCache(cache_size_limit) + ParserElement._parse = ParserElement._parseCache + + def parse_string( + self, instring: str, parse_all: bool = False, *, parseAll: bool = False + ) -> ParseResults: + """ + Parse a string with respect to the parser definition. This function is intended as the primary interface to the + client code. + + :param instring: The input string to be parsed. + :param parse_all: If set, the entire input string must match the grammar. + :param parseAll: retained for pre-PEP8 compatibility, will be removed in a future release. + :raises ParseException: Raised if ``parse_all`` is set and the input string does not match the whole grammar. + :returns: the parsed data as a :class:`ParseResults` object, which may be accessed as a `list`, a `dict`, or + an object with attributes if the given parser includes results names. + + If the input string is required to match the entire grammar, ``parse_all`` flag must be set to ``True``. This + is also equivalent to ending the grammar with :class:`StringEnd`\\ (). + + To report proper column numbers, ``parse_string`` operates on a copy of the input string where all tabs are + converted to spaces (8 spaces per tab, as per the default in ``string.expandtabs``). If the input string + contains tabs and the grammar uses parse actions that use the ``loc`` argument to index into the string + being parsed, one can ensure a consistent view of the input string by doing one of the following: + + - calling ``parse_with_tabs`` on your grammar before calling ``parse_string`` (see :class:`parse_with_tabs`), + - define your parse action using the full ``(s,loc,toks)`` signature, and reference the input string using the + parse action's ``s`` argument, or + - explicitly expand the tabs in your input string before calling ``parse_string``. + + Examples: + + By default, partial matches are OK. + + >>> res = Word('a').parse_string('aaaaabaaa') + >>> print(res) + ['aaaaa'] + + The parsing behavior varies by the inheriting class of this abstract class. Please refer to the children + directly to see more examples. + + It raises an exception if parse_all flag is set and instring does not match the whole grammar. + + >>> res = Word('a').parse_string('aaaaabaaa', parse_all=True) + Traceback (most recent call last): + ... + pyparsing.ParseException: Expected end of text, found 'b' (at char 5), (line:1, col:6) + """ + parseAll = parse_all or parseAll + + ParserElement.reset_cache() + if not self.streamlined: + self.streamline() + for e in self.ignoreExprs: + e.streamline() + if not self.keepTabs: + instring = instring.expandtabs() + try: + loc, tokens = self._parse(instring, 0) + if parseAll: + loc = self.preParse(instring, loc) + se = Empty() + StringEnd().set_debug(False) + se._parse(instring, loc) + except _ParseActionIndexError as pa_exc: + raise pa_exc.exc + except ParseBaseException as exc: + if ParserElement.verbose_stacktrace: + raise + + # catch and re-raise exception from here, clearing out pyparsing internal stack trace + raise exc.with_traceback(None) + else: + return tokens + + def scan_string( + self, + instring: str, + max_matches: int = _MAX_INT, + overlap: bool = False, + always_skip_whitespace=True, + *, + debug: bool = False, + maxMatches: int = _MAX_INT, + ) -> Generator[tuple[ParseResults, int, int], None, None]: + """ + Scan the input string for expression matches. Each match will return the + matching tokens, start location, and end location. May be called with optional + ``max_matches`` argument, to clip scanning after 'n' matches are found. If + ``overlap`` is specified, then overlapping matches will be reported. + + Note that the start and end locations are reported relative to the string + being parsed. See :class:`parse_string` for more information on parsing + strings with embedded tabs. + + Example:: + + source = "sldjf123lsdjjkf345sldkjf879lkjsfd987" + print(source) + for tokens, start, end in Word(alphas).scan_string(source): + print(' '*start + '^'*(end-start)) + print(' '*start + tokens[0]) + + prints:: + + sldjf123lsdjjkf345sldkjf879lkjsfd987 + ^^^^^ + sldjf + ^^^^^^^ + lsdjjkf + ^^^^^^ + sldkjf + ^^^^^^ + lkjsfd + """ + maxMatches = min(maxMatches, max_matches) + if not self.streamlined: + self.streamline() + for e in self.ignoreExprs: + e.streamline() + + if not self.keepTabs: + instring = str(instring).expandtabs() + instrlen = len(instring) + loc = 0 + if always_skip_whitespace: + preparser = Empty() + preparser.ignoreExprs = self.ignoreExprs + preparser.whiteChars = self.whiteChars + preparseFn = preparser.preParse + else: + preparseFn = self.preParse + parseFn = self._parse + ParserElement.resetCache() + matches = 0 + try: + while loc <= instrlen and matches < maxMatches: + try: + preloc: int = preparseFn(instring, loc) + nextLoc: int + tokens: ParseResults + nextLoc, tokens = parseFn(instring, preloc, callPreParse=False) + except ParseException: + loc = preloc + 1 + else: + if nextLoc > loc: + matches += 1 + if debug: + print( + { + "tokens": tokens.asList(), + "start": preloc, + "end": nextLoc, + } + ) + yield tokens, preloc, nextLoc + if overlap: + nextloc = preparseFn(instring, loc) + if nextloc > loc: + loc = nextLoc + else: + loc += 1 + else: + loc = nextLoc + else: + loc = preloc + 1 + except ParseBaseException as exc: + if ParserElement.verbose_stacktrace: + raise + + # catch and re-raise exception from here, clears out pyparsing internal stack trace + raise exc.with_traceback(None) + + def transform_string(self, instring: str, *, debug: bool = False) -> str: + """ + Extension to :class:`scan_string`, to modify matching text with modified tokens that may + be returned from a parse action. To use ``transform_string``, define a grammar and + attach a parse action to it that modifies the returned token list. + Invoking ``transform_string()`` on a target string will then scan for matches, + and replace the matched text patterns according to the logic in the parse + action. ``transform_string()`` returns the resulting transformed string. + + Example:: + + wd = Word(alphas) + wd.set_parse_action(lambda toks: toks[0].title()) + + print(wd.transform_string("now is the winter of our discontent made glorious summer by this sun of york.")) + + prints:: + + Now Is The Winter Of Our Discontent Made Glorious Summer By This Sun Of York. + """ + out: list[str] = [] + lastE = 0 + # force preservation of s, to minimize unwanted transformation of string, and to + # keep string locs straight between transform_string and scan_string + self.keepTabs = True + try: + for t, s, e in self.scan_string(instring, debug=debug): + if s > lastE: + out.append(instring[lastE:s]) + lastE = e + + if not t: + continue + + if isinstance(t, ParseResults): + out += t.as_list() + elif isinstance(t, Iterable) and not isinstance(t, str_type): + out.extend(t) + else: + out.append(t) + + out.append(instring[lastE:]) + out = [o for o in out if o] + return "".join([str(s) for s in _flatten(out)]) + except ParseBaseException as exc: + if ParserElement.verbose_stacktrace: + raise + + # catch and re-raise exception from here, clears out pyparsing internal stack trace + raise exc.with_traceback(None) + + def search_string( + self, + instring: str, + max_matches: int = _MAX_INT, + *, + debug: bool = False, + maxMatches: int = _MAX_INT, + ) -> ParseResults: + """ + Another extension to :class:`scan_string`, simplifying the access to the tokens found + to match the given parse expression. May be called with optional + ``max_matches`` argument, to clip searching after 'n' matches are found. + + Example:: + + # a capitalized word starts with an uppercase letter, followed by zero or more lowercase letters + cap_word = Word(alphas.upper(), alphas.lower()) + + print(cap_word.search_string("More than Iron, more than Lead, more than Gold I need Electricity")) + + # the sum() builtin can be used to merge results into a single ParseResults object + print(sum(cap_word.search_string("More than Iron, more than Lead, more than Gold I need Electricity"))) + + prints:: + + [['More'], ['Iron'], ['Lead'], ['Gold'], ['I'], ['Electricity']] + ['More', 'Iron', 'Lead', 'Gold', 'I', 'Electricity'] + """ + maxMatches = min(maxMatches, max_matches) + try: + return ParseResults( + [ + t + for t, s, e in self.scan_string( + instring, maxMatches, always_skip_whitespace=False, debug=debug + ) + ] + ) + except ParseBaseException as exc: + if ParserElement.verbose_stacktrace: + raise + + # catch and re-raise exception from here, clears out pyparsing internal stack trace + raise exc.with_traceback(None) + + def split( + self, + instring: str, + maxsplit: int = _MAX_INT, + include_separators: bool = False, + *, + includeSeparators=False, + ) -> Generator[str, None, None]: + """ + Generator method to split a string using the given expression as a separator. + May be called with optional ``maxsplit`` argument, to limit the number of splits; + and the optional ``include_separators`` argument (default= ``False``), if the separating + matching text should be included in the split results. + + Example:: + + punc = one_of(list(".,;:/-!?")) + print(list(punc.split("This, this?, this sentence, is badly punctuated!"))) + + prints:: + + ['This', ' this', '', ' this sentence', ' is badly punctuated', ''] + """ + includeSeparators = includeSeparators or include_separators + last = 0 + for t, s, e in self.scan_string(instring, max_matches=maxsplit): + yield instring[last:s] + if includeSeparators: + yield t[0] + last = e + yield instring[last:] + + def __add__(self, other) -> ParserElement: + """ + Implementation of ``+`` operator - returns :class:`And`. Adding strings to a :class:`ParserElement` + converts them to :class:`Literal`\\ s by default. + + Example:: + + greet = Word(alphas) + "," + Word(alphas) + "!" + hello = "Hello, World!" + print(hello, "->", greet.parse_string(hello)) + + prints:: + + Hello, World! -> ['Hello', ',', 'World', '!'] + + ``...`` may be used as a parse expression as a short form of :class:`SkipTo`:: + + Literal('start') + ... + Literal('end') + + is equivalent to:: + + Literal('start') + SkipTo('end')("_skipped*") + Literal('end') + + Note that the skipped text is returned with '_skipped' as a results name, + and to support having multiple skips in the same parser, the value returned is + a list of all skipped text. + """ + if other is Ellipsis: + return _PendingSkip(self) + + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return And([self, other]) + + def __radd__(self, other) -> ParserElement: + """ + Implementation of ``+`` operator when left operand is not a :class:`ParserElement` + """ + if other is Ellipsis: + return SkipTo(self)("_skipped*") + self + + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return other + self + + def __sub__(self, other) -> ParserElement: + """ + Implementation of ``-`` operator, returns :class:`And` with error stop + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return self + And._ErrorStop() + other + + def __rsub__(self, other) -> ParserElement: + """ + Implementation of ``-`` operator when left operand is not a :class:`ParserElement` + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return other - self + + def __mul__(self, other) -> ParserElement: + """ + Implementation of ``*`` operator, allows use of ``expr * 3`` in place of + ``expr + expr + expr``. Expressions may also be multiplied by a 2-integer + tuple, similar to ``{min, max}`` multipliers in regular expressions. Tuples + may also include ``None`` as in: + + - ``expr*(n, None)`` or ``expr*(n, )`` is equivalent + to ``expr*n + ZeroOrMore(expr)`` + (read as "at least n instances of ``expr``") + - ``expr*(None, n)`` is equivalent to ``expr*(0, n)`` + (read as "0 to n instances of ``expr``") + - ``expr*(None, None)`` is equivalent to ``ZeroOrMore(expr)`` + - ``expr*(1, None)`` is equivalent to ``OneOrMore(expr)`` + + Note that ``expr*(None, n)`` does not raise an exception if + more than n exprs exist in the input stream; that is, + ``expr*(None, n)`` does not enforce a maximum number of expr + occurrences. If this behavior is desired, then write + ``expr*(None, n) + ~expr`` + """ + if other is Ellipsis: + other = (0, None) + elif isinstance(other, tuple) and other[:1] == (Ellipsis,): + other = ((0,) + other[1:] + (None,))[:2] + + if not isinstance(other, (int, tuple)): + return NotImplemented + + if isinstance(other, int): + minElements, optElements = other, 0 + else: + other = tuple(o if o is not Ellipsis else None for o in other) + other = (other + (None, None))[:2] + if other[0] is None: + other = (0, other[1]) + if isinstance(other[0], int) and other[1] is None: + if other[0] == 0: + return ZeroOrMore(self) + if other[0] == 1: + return OneOrMore(self) + else: + return self * other[0] + ZeroOrMore(self) + elif isinstance(other[0], int) and isinstance(other[1], int): + minElements, optElements = other + optElements -= minElements + else: + return NotImplemented + + if minElements < 0: + raise ValueError("cannot multiply ParserElement by negative value") + if optElements < 0: + raise ValueError( + "second tuple value must be greater or equal to first tuple value" + ) + if minElements == optElements == 0: + return And([]) + + if optElements: + + def makeOptionalList(n): + if n > 1: + return Opt(self + makeOptionalList(n - 1)) + else: + return Opt(self) + + if minElements: + if minElements == 1: + ret = self + makeOptionalList(optElements) + else: + ret = And([self] * minElements) + makeOptionalList(optElements) + else: + ret = makeOptionalList(optElements) + else: + if minElements == 1: + ret = self + else: + ret = And([self] * minElements) + return ret + + def __rmul__(self, other) -> ParserElement: + return self.__mul__(other) + + def __or__(self, other) -> ParserElement: + """ + Implementation of ``|`` operator - returns :class:`MatchFirst` + """ + if other is Ellipsis: + return _PendingSkip(self, must_skip=True) + + if isinstance(other, str_type): + # `expr | ""` is equivalent to `Opt(expr)` + if other == "": + return Opt(self) + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return MatchFirst([self, other]) + + def __ror__(self, other) -> ParserElement: + """ + Implementation of ``|`` operator when left operand is not a :class:`ParserElement` + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return other | self + + def __xor__(self, other) -> ParserElement: + """ + Implementation of ``^`` operator - returns :class:`Or` + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return Or([self, other]) + + def __rxor__(self, other) -> ParserElement: + """ + Implementation of ``^`` operator when left operand is not a :class:`ParserElement` + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return other ^ self + + def __and__(self, other) -> ParserElement: + """ + Implementation of ``&`` operator - returns :class:`Each` + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return Each([self, other]) + + def __rand__(self, other) -> ParserElement: + """ + Implementation of ``&`` operator when left operand is not a :class:`ParserElement` + """ + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return other & self + + def __invert__(self) -> ParserElement: + """ + Implementation of ``~`` operator - returns :class:`NotAny` + """ + return NotAny(self) + + # disable __iter__ to override legacy use of sequential access to __getitem__ to + # iterate over a sequence + __iter__ = None + + def __getitem__(self, key): + """ + use ``[]`` indexing notation as a short form for expression repetition: + + - ``expr[n]`` is equivalent to ``expr*n`` + - ``expr[m, n]`` is equivalent to ``expr*(m, n)`` + - ``expr[n, ...]`` or ``expr[n,]`` is equivalent + to ``expr*n + ZeroOrMore(expr)`` + (read as "at least n instances of ``expr``") + - ``expr[..., n]`` is equivalent to ``expr*(0, n)`` + (read as "0 to n instances of ``expr``") + - ``expr[...]`` and ``expr[0, ...]`` are equivalent to ``ZeroOrMore(expr)`` + - ``expr[1, ...]`` is equivalent to ``OneOrMore(expr)`` + + ``None`` may be used in place of ``...``. + + Note that ``expr[..., n]`` and ``expr[m, n]`` do not raise an exception + if more than ``n`` ``expr``\\ s exist in the input stream. If this behavior is + desired, then write ``expr[..., n] + ~expr``. + + For repetition with a stop_on expression, use slice notation: + + - ``expr[...: end_expr]`` and ``expr[0, ...: end_expr]`` are equivalent to ``ZeroOrMore(expr, stop_on=end_expr)`` + - ``expr[1, ...: end_expr]`` is equivalent to ``OneOrMore(expr, stop_on=end_expr)`` + + """ + + stop_on_defined = False + stop_on = NoMatch() + if isinstance(key, slice): + key, stop_on = key.start, key.stop + if key is None: + key = ... + stop_on_defined = True + elif isinstance(key, tuple) and isinstance(key[-1], slice): + key, stop_on = (key[0], key[1].start), key[1].stop + stop_on_defined = True + + # convert single arg keys to tuples + if isinstance(key, str_type): + key = (key,) + try: + iter(key) + except TypeError: + key = (key, key) + + if len(key) > 2: + raise TypeError( + f"only 1 or 2 index arguments supported ({key[:5]}{f'... [{len(key)}]' if len(key) > 5 else ''})" + ) + + # clip to 2 elements + ret = self * tuple(key[:2]) + ret = typing.cast(_MultipleMatch, ret) + + if stop_on_defined: + ret.stopOn(stop_on) + + return ret + + def __call__(self, name: typing.Optional[str] = None) -> ParserElement: + """ + Shortcut for :class:`set_results_name`, with ``list_all_matches=False``. + + If ``name`` is given with a trailing ``'*'`` character, then ``list_all_matches`` will be + passed as ``True``. + + If ``name`` is omitted, same as calling :class:`copy`. + + Example:: + + # these are equivalent + userdata = Word(alphas).set_results_name("name") + Word(nums + "-").set_results_name("socsecno") + userdata = Word(alphas)("name") + Word(nums + "-")("socsecno") + """ + if name is not None: + return self._setResultsName(name) + + return self.copy() + + def suppress(self) -> ParserElement: + """ + Suppresses the output of this :class:`ParserElement`; useful to keep punctuation from + cluttering up returned output. + """ + return Suppress(self) + + def ignore_whitespace(self, recursive: bool = True) -> ParserElement: + """ + Enables the skipping of whitespace before matching the characters in the + :class:`ParserElement`'s defined pattern. + + :param recursive: If ``True`` (the default), also enable whitespace skipping in child elements (if any) + """ + self.skipWhitespace = True + return self + + def leave_whitespace(self, recursive: bool = True) -> ParserElement: + """ + Disables the skipping of whitespace before matching the characters in the + :class:`ParserElement`'s defined pattern. This is normally only used internally by + the pyparsing module, but may be needed in some whitespace-sensitive grammars. + + :param recursive: If true (the default), also disable whitespace skipping in child elements (if any) + """ + self.skipWhitespace = False + return self + + def set_whitespace_chars( + self, chars: Union[set[str], str], copy_defaults: bool = False + ) -> ParserElement: + """ + Overrides the default whitespace chars + """ + self.skipWhitespace = True + self.whiteChars = set(chars) + self.copyDefaultWhiteChars = copy_defaults + return self + + def parse_with_tabs(self) -> ParserElement: + """ + Overrides default behavior to expand ```` s to spaces before parsing the input string. + Must be called before ``parse_string`` when the input grammar contains elements that + match ```` characters. + """ + self.keepTabs = True + return self + + def ignore(self, other: ParserElement) -> ParserElement: + """ + Define expression to be ignored (e.g., comments) while doing pattern + matching; may be called repeatedly, to define multiple comment or other + ignorable patterns. + + Example:: + + patt = Word(alphas)[...] + patt.parse_string('ablaj /* comment */ lskjd') + # -> ['ablaj'] + + patt.ignore(c_style_comment) + patt.parse_string('ablaj /* comment */ lskjd') + # -> ['ablaj', 'lskjd'] + """ + if isinstance(other, str_type): + other = Suppress(other) + + if isinstance(other, Suppress): + if other not in self.ignoreExprs: + self.ignoreExprs.append(other) + else: + self.ignoreExprs.append(Suppress(other.copy())) + return self + + def set_debug_actions( + self, + start_action: DebugStartAction, + success_action: DebugSuccessAction, + exception_action: DebugExceptionAction, + ) -> ParserElement: + """ + Customize display of debugging messages while doing pattern matching: + + - ``start_action`` - method to be called when an expression is about to be parsed; + should have the signature ``fn(input_string: str, location: int, expression: ParserElement, cache_hit: bool)`` + + - ``success_action`` - method to be called when an expression has successfully parsed; + should have the signature ``fn(input_string: str, start_location: int, end_location: int, expression: ParserELement, parsed_tokens: ParseResults, cache_hit: bool)`` + + - ``exception_action`` - method to be called when expression fails to parse; + should have the signature ``fn(input_string: str, location: int, expression: ParserElement, exception: Exception, cache_hit: bool)`` + """ + self.debugActions = self.DebugActions( + start_action or _default_start_debug_action, # type: ignore[truthy-function] + success_action or _default_success_debug_action, # type: ignore[truthy-function] + exception_action or _default_exception_debug_action, # type: ignore[truthy-function] + ) + self.debug = True + return self + + def set_debug(self, flag: bool = True, recurse: bool = False) -> ParserElement: + """ + Enable display of debugging messages while doing pattern matching. + Set ``flag`` to ``True`` to enable, ``False`` to disable. + Set ``recurse`` to ``True`` to set the debug flag on this expression and all sub-expressions. + + Example:: + + wd = Word(alphas).set_name("alphaword") + integer = Word(nums).set_name("numword") + term = wd | integer + + # turn on debugging for wd + wd.set_debug() + + term[1, ...].parse_string("abc 123 xyz 890") + + prints:: + + Match alphaword at loc 0(1,1) + Matched alphaword -> ['abc'] + Match alphaword at loc 3(1,4) + Exception raised:Expected alphaword (at char 4), (line:1, col:5) + Match alphaword at loc 7(1,8) + Matched alphaword -> ['xyz'] + Match alphaword at loc 11(1,12) + Exception raised:Expected alphaword (at char 12), (line:1, col:13) + Match alphaword at loc 15(1,16) + Exception raised:Expected alphaword (at char 15), (line:1, col:16) + + The output shown is that produced by the default debug actions - custom debug actions can be + specified using :class:`set_debug_actions`. Prior to attempting + to match the ``wd`` expression, the debugging message ``"Match at loc (,)"`` + is shown. Then if the parse succeeds, a ``"Matched"`` message is shown, or an ``"Exception raised"`` + message is shown. Also note the use of :class:`set_name` to assign a human-readable name to the expression, + which makes debugging and exception messages easier to understand - for instance, the default + name created for the :class:`Word` expression without calling ``set_name`` is ``"W:(A-Za-z)"``. + """ + if recurse: + for expr in self.visit_all(): + expr.set_debug(flag, recurse=False) + return self + + if flag: + self.set_debug_actions( + _default_start_debug_action, + _default_success_debug_action, + _default_exception_debug_action, + ) + else: + self.debug = False + return self + + @property + def default_name(self) -> str: + if self._defaultName is None: + self._defaultName = self._generateDefaultName() + return self._defaultName + + @abstractmethod + def _generateDefaultName(self) -> str: + """ + Child classes must define this method, which defines how the ``default_name`` is set. + """ + + def set_name(self, name: typing.Optional[str]) -> ParserElement: + """ + Define name for this expression, makes debugging and exception messages clearer. If + `__diag__.enable_debug_on_named_expressions` is set to True, setting a name will also + enable debug for this expression. + + If `name` is None, clears any custom name for this expression, and clears the + debug flag is it was enabled via `__diag__.enable_debug_on_named_expressions`. + + Example:: + + integer = Word(nums) + integer.parse_string("ABC") # -> Exception: Expected W:(0-9) (at char 0), (line:1, col:1) + + integer.set_name("integer") + integer.parse_string("ABC") # -> Exception: Expected integer (at char 0), (line:1, col:1) + """ + self.customName = name # type: ignore[assignment] + self.errmsg = f"Expected {str(self)}" + + if __diag__.enable_debug_on_named_expressions: + self.set_debug(name is not None) + + return self + + @property + def name(self) -> str: + # This will use a user-defined name if available, but otherwise defaults back to the auto-generated name + return self.customName if self.customName is not None else self.default_name + + @name.setter + def name(self, new_name) -> None: + self.set_name(new_name) + + def __str__(self) -> str: + return self.name + + def __repr__(self) -> str: + return str(self) + + def streamline(self) -> ParserElement: + self.streamlined = True + self._defaultName = None + return self + + def recurse(self) -> list[ParserElement]: + return [] + + def _checkRecursion(self, parseElementList): + subRecCheckList = parseElementList[:] + [self] + for e in self.recurse(): + e._checkRecursion(subRecCheckList) + + def validate(self, validateTrace=None) -> None: + """ + Check defined expressions for valid structure, check for infinite recursive definitions. + """ + warnings.warn( + "ParserElement.validate() is deprecated, and should not be used to check for left recursion", + DeprecationWarning, + stacklevel=2, + ) + self._checkRecursion([]) + + def parse_file( + self, + file_or_filename: Union[str, Path, TextIO], + encoding: str = "utf-8", + parse_all: bool = False, + *, + parseAll: bool = False, + ) -> ParseResults: + """ + Execute the parse expression on the given file or filename. + If a filename is specified (instead of a file object), + the entire file is opened, read, and closed before parsing. + """ + parseAll = parseAll or parse_all + try: + file_or_filename = typing.cast(TextIO, file_or_filename) + file_contents = file_or_filename.read() + except AttributeError: + file_or_filename = typing.cast(str, file_or_filename) + with open(file_or_filename, "r", encoding=encoding) as f: + file_contents = f.read() + try: + return self.parse_string(file_contents, parseAll) + except ParseBaseException as exc: + if ParserElement.verbose_stacktrace: + raise + + # catch and re-raise exception from here, clears out pyparsing internal stack trace + raise exc.with_traceback(None) + + def __eq__(self, other): + if self is other: + return True + elif isinstance(other, str_type): + return self.matches(other, parse_all=True) + elif isinstance(other, ParserElement): + return vars(self) == vars(other) + return False + + def __hash__(self): + return id(self) + + def matches( + self, test_string: str, parse_all: bool = True, *, parseAll: bool = True + ) -> bool: + """ + Method for quick testing of a parser against a test string. Good for simple + inline microtests of sub expressions while building up larger parser. + + Parameters: + + - ``test_string`` - to test against this expression for a match + - ``parse_all`` - (default= ``True``) - flag to pass to :class:`parse_string` when running tests + + Example:: + + expr = Word(nums) + assert expr.matches("100") + """ + parseAll = parseAll and parse_all + try: + self.parse_string(str(test_string), parse_all=parseAll) + return True + except ParseBaseException: + return False + + def run_tests( + self, + tests: Union[str, list[str]], + parse_all: bool = True, + comment: typing.Optional[Union[ParserElement, str]] = "#", + full_dump: bool = True, + print_results: bool = True, + failure_tests: bool = False, + post_parse: typing.Optional[ + Callable[[str, ParseResults], typing.Optional[str]] + ] = None, + file: typing.Optional[TextIO] = None, + with_line_numbers: bool = False, + *, + parseAll: bool = True, + fullDump: bool = True, + printResults: bool = True, + failureTests: bool = False, + postParse: typing.Optional[ + Callable[[str, ParseResults], typing.Optional[str]] + ] = None, + ) -> tuple[bool, list[tuple[str, Union[ParseResults, Exception]]]]: + """ + Execute the parse expression on a series of test strings, showing each + test, the parsed results or where the parse failed. Quick and easy way to + run a parse expression against a list of sample strings. + + Parameters: + + - ``tests`` - a list of separate test strings, or a multiline string of test strings + - ``parse_all`` - (default= ``True``) - flag to pass to :class:`parse_string` when running tests + - ``comment`` - (default= ``'#'``) - expression for indicating embedded comments in the test + string; pass None to disable comment filtering + - ``full_dump`` - (default= ``True``) - dump results as list followed by results names in nested outline; + if False, only dump nested list + - ``print_results`` - (default= ``True``) prints test output to stdout + - ``failure_tests`` - (default= ``False``) indicates if these tests are expected to fail parsing + - ``post_parse`` - (default= ``None``) optional callback for successful parse results; called as + `fn(test_string, parse_results)` and returns a string to be added to the test output + - ``file`` - (default= ``None``) optional file-like object to which test output will be written; + if None, will default to ``sys.stdout`` + - ``with_line_numbers`` - default= ``False``) show test strings with line and column numbers + + Returns: a (success, results) tuple, where success indicates that all tests succeeded + (or failed if ``failure_tests`` is True), and the results contain a list of lines of each + test's output + + Example:: + + number_expr = pyparsing_common.number.copy() + + result = number_expr.run_tests(''' + # unsigned integer + 100 + # negative integer + -100 + # float with scientific notation + 6.02e23 + # integer with scientific notation + 1e-12 + ''') + print("Success" if result[0] else "Failed!") + + result = number_expr.run_tests(''' + # stray character + 100Z + # missing leading digit before '.' + -.100 + # too many '.' + 3.14.159 + ''', failure_tests=True) + print("Success" if result[0] else "Failed!") + + prints:: + + # unsigned integer + 100 + [100] + + # negative integer + -100 + [-100] + + # float with scientific notation + 6.02e23 + [6.02e+23] + + # integer with scientific notation + 1e-12 + [1e-12] + + Success + + # stray character + 100Z + ^ + FAIL: Expected end of text (at char 3), (line:1, col:4) + + # missing leading digit before '.' + -.100 + ^ + FAIL: Expected {real number with scientific notation | real number | signed integer} (at char 0), (line:1, col:1) + + # too many '.' + 3.14.159 + ^ + FAIL: Expected end of text (at char 4), (line:1, col:5) + + Success + + Each test string must be on a single line. If you want to test a string that spans multiple + lines, create a test like this:: + + expr.run_tests(r"this is a test\\n of strings that spans \\n 3 lines") + + (Note that this is a raw string literal, you must include the leading ``'r'``.) + """ + from .testing import pyparsing_test + + parseAll = parseAll and parse_all + fullDump = fullDump and full_dump + printResults = printResults and print_results + failureTests = failureTests or failure_tests + postParse = postParse or post_parse + if isinstance(tests, str_type): + tests = typing.cast(str, tests) + line_strip = type(tests).strip + tests = [line_strip(test_line) for test_line in tests.rstrip().splitlines()] + comment_specified = comment is not None + if comment_specified: + if isinstance(comment, str_type): + comment = typing.cast(str, comment) + comment = Literal(comment) + comment = typing.cast(ParserElement, comment) + if file is None: + file = sys.stdout + print_ = file.write + + result: Union[ParseResults, Exception] + allResults: list[tuple[str, Union[ParseResults, Exception]]] = [] + comments: list[str] = [] + success = True + NL = Literal(r"\n").add_parse_action(replace_with("\n")).ignore(quoted_string) + BOM = "\ufeff" + nlstr = "\n" + for t in tests: + if comment_specified and comment.matches(t, False) or comments and not t: + comments.append( + pyparsing_test.with_line_numbers(t) if with_line_numbers else t + ) + continue + if not t: + continue + out = [ + f"{nlstr}{nlstr.join(comments) if comments else ''}", + pyparsing_test.with_line_numbers(t) if with_line_numbers else t, + ] + comments.clear() + try: + # convert newline marks to actual newlines, and strip leading BOM if present + t = NL.transform_string(t.lstrip(BOM)) + result = self.parse_string(t, parse_all=parseAll) + except ParseBaseException as pe: + fatal = "(FATAL) " if isinstance(pe, ParseFatalException) else "" + out.append(pe.explain()) + out.append(f"FAIL: {fatal}{pe}") + if ParserElement.verbose_stacktrace: + out.extend(traceback.format_tb(pe.__traceback__)) + success = success and failureTests + result = pe + except Exception as exc: + tag = "FAIL-EXCEPTION" + + # see if this exception was raised in a parse action + tb = exc.__traceback__ + it = iter(traceback.walk_tb(tb)) + for f, line in it: + if (f.f_code.co_filename, line) == pa_call_line_synth: + next_f = next(it)[0] + tag += f" (raised in parse action {next_f.f_code.co_name!r})" + break + + out.append(f"{tag}: {type(exc).__name__}: {exc}") + if ParserElement.verbose_stacktrace: + out.extend(traceback.format_tb(exc.__traceback__)) + success = success and failureTests + result = exc + else: + success = success and not failureTests + if postParse is not None: + try: + pp_value = postParse(t, result) + if pp_value is not None: + if isinstance(pp_value, ParseResults): + out.append(pp_value.dump()) + else: + out.append(str(pp_value)) + else: + out.append(result.dump()) + except Exception as e: + out.append(result.dump(full=fullDump)) + out.append( + f"{postParse.__name__} failed: {type(e).__name__}: {e}" + ) + else: + out.append(result.dump(full=fullDump)) + out.append("") + + if printResults: + print_("\n".join(out)) + + allResults.append((t, result)) + + return success, allResults + + def create_diagram( + self, + output_html: Union[TextIO, Path, str], + vertical: int = 3, + show_results_names: bool = False, + show_groups: bool = False, + embed: bool = False, + show_hidden: bool = False, + **kwargs, + ) -> None: + """ + Create a railroad diagram for the parser. + + Parameters: + + - ``output_html`` (str or file-like object) - output target for generated + diagram HTML + - ``vertical`` (int) - threshold for formatting multiple alternatives vertically + instead of horizontally (default=3) + - ``show_results_names`` - bool flag whether diagram should show annotations for + defined results names + - ``show_groups`` - bool flag whether groups should be highlighted with an unlabeled surrounding box + - ``show_hidden`` - bool flag to show diagram elements for internal elements that are usually hidden + - ``embed`` - bool flag whether generated HTML should omit , , and tags to embed + the resulting HTML in an enclosing HTML source + - ``head`` - str containing additional HTML to insert into the section of the generated code; + can be used to insert custom CSS styling + - ``body`` - str containing additional HTML to insert at the beginning of the section of the + generated code + + Additional diagram-formatting keyword arguments can also be included; + see railroad.Diagram class. + """ + + try: + from .diagram import to_railroad, railroad_to_html + except ImportError as ie: + raise Exception( + "must ``pip install pyparsing[diagrams]`` to generate parser railroad diagrams" + ) from ie + + self.streamline() + + railroad = to_railroad( + self, + vertical=vertical, + show_results_names=show_results_names, + show_groups=show_groups, + show_hidden=show_hidden, + diagram_kwargs=kwargs, + ) + if not isinstance(output_html, (str, Path)): + # we were passed a file-like object, just write to it + output_html.write(railroad_to_html(railroad, embed=embed, **kwargs)) + return + + with open(output_html, "w", encoding="utf-8") as diag_file: + diag_file.write(railroad_to_html(railroad, embed=embed, **kwargs)) + + # Compatibility synonyms + # fmt: off + inlineLiteralsUsing = staticmethod(replaced_by_pep8("inlineLiteralsUsing", inline_literals_using)) + setDefaultWhitespaceChars = staticmethod(replaced_by_pep8( + "setDefaultWhitespaceChars", set_default_whitespace_chars + )) + disableMemoization = staticmethod(replaced_by_pep8("disableMemoization", disable_memoization)) + enableLeftRecursion = staticmethod(replaced_by_pep8("enableLeftRecursion", enable_left_recursion)) + enablePackrat = staticmethod(replaced_by_pep8("enablePackrat", enable_packrat)) + resetCache = staticmethod(replaced_by_pep8("resetCache", reset_cache)) + + setResultsName = replaced_by_pep8("setResultsName", set_results_name) + setBreak = replaced_by_pep8("setBreak", set_break) + setParseAction = replaced_by_pep8("setParseAction", set_parse_action) + addParseAction = replaced_by_pep8("addParseAction", add_parse_action) + addCondition = replaced_by_pep8("addCondition", add_condition) + setFailAction = replaced_by_pep8("setFailAction", set_fail_action) + tryParse = replaced_by_pep8("tryParse", try_parse) + parseString = replaced_by_pep8("parseString", parse_string) + scanString = replaced_by_pep8("scanString", scan_string) + transformString = replaced_by_pep8("transformString", transform_string) + searchString = replaced_by_pep8("searchString", search_string) + ignoreWhitespace = replaced_by_pep8("ignoreWhitespace", ignore_whitespace) + leaveWhitespace = replaced_by_pep8("leaveWhitespace", leave_whitespace) + setWhitespaceChars = replaced_by_pep8("setWhitespaceChars", set_whitespace_chars) + parseWithTabs = replaced_by_pep8("parseWithTabs", parse_with_tabs) + setDebugActions = replaced_by_pep8("setDebugActions", set_debug_actions) + setDebug = replaced_by_pep8("setDebug", set_debug) + setName = replaced_by_pep8("setName", set_name) + parseFile = replaced_by_pep8("parseFile", parse_file) + runTests = replaced_by_pep8("runTests", run_tests) + canParseNext = replaced_by_pep8("canParseNext", can_parse_next) + defaultName = default_name + # fmt: on + + +class _PendingSkip(ParserElement): + # internal placeholder class to hold a place were '...' is added to a parser element, + # once another ParserElement is added, this placeholder will be replaced with a SkipTo + def __init__(self, expr: ParserElement, must_skip: bool = False) -> None: + super().__init__() + self.anchor = expr + self.must_skip = must_skip + + def _generateDefaultName(self) -> str: + return str(self.anchor + Empty()).replace("Empty", "...") + + def __add__(self, other) -> ParserElement: + skipper = SkipTo(other).set_name("...")("_skipped*") + if self.must_skip: + + def must_skip(t): + if not t._skipped or t._skipped.as_list() == [""]: + del t[0] + t.pop("_skipped", None) + + def show_skip(t): + if t._skipped.as_list()[-1:] == [""]: + t.pop("_skipped") + t["_skipped"] = f"missing <{self.anchor!r}>" + + return ( + self.anchor + skipper().add_parse_action(must_skip) + | skipper().add_parse_action(show_skip) + ) + other + + return self.anchor + skipper + other + + def __repr__(self): + return self.defaultName + + def parseImpl(self, *args) -> ParseImplReturnType: + raise Exception( + "use of `...` expression without following SkipTo target expression" + ) + + +class Token(ParserElement): + """Abstract :class:`ParserElement` subclass, for defining atomic + matching patterns. + """ + + def __init__(self) -> None: + super().__init__(savelist=False) + + def _generateDefaultName(self) -> str: + return type(self).__name__ + + +class NoMatch(Token): + """ + A token that will never match. + """ + + def __init__(self) -> None: + super().__init__() + self._may_return_empty = True + self.mayIndexError = False + self.errmsg = "Unmatchable token" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + raise ParseException(instring, loc, self.errmsg, self) + + +class Literal(Token): + """ + Token to exactly match a specified string. + + Example:: + + Literal('abc').parse_string('abc') # -> ['abc'] + Literal('abc').parse_string('abcdef') # -> ['abc'] + Literal('abc').parse_string('ab') # -> Exception: Expected "abc" + + For case-insensitive matching, use :class:`CaselessLiteral`. + + For keyword matching (force word break before and after the matched string), + use :class:`Keyword` or :class:`CaselessKeyword`. + """ + + def __new__(cls, match_string: str = "", *, matchString: str = ""): + # Performance tuning: select a subclass with optimized parseImpl + if cls is Literal: + match_string = matchString or match_string + if not match_string: + return super().__new__(Empty) + if len(match_string) == 1: + return super().__new__(_SingleCharLiteral) + + # Default behavior + return super().__new__(cls) + + # Needed to make copy.copy() work correctly if we customize __new__ + def __getnewargs__(self): + return (self.match,) + + def __init__(self, match_string: str = "", *, matchString: str = "") -> None: + super().__init__() + match_string = matchString or match_string + self.match = match_string + self.matchLen = len(match_string) + self.firstMatchChar = match_string[:1] + self.errmsg = f"Expected {self.name}" + self._may_return_empty = False + self.mayIndexError = False + + def _generateDefaultName(self) -> str: + return repr(self.match) + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if instring[loc] == self.firstMatchChar and instring.startswith( + self.match, loc + ): + return loc + self.matchLen, self.match + raise ParseException(instring, loc, self.errmsg, self) + + +class Empty(Literal): + """ + An empty token, will always match. + """ + + def __init__(self, match_string="", *, matchString="") -> None: + super().__init__("") + self._may_return_empty = True + self.mayIndexError = False + + def _generateDefaultName(self) -> str: + return "Empty" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + return loc, [] + + +class _SingleCharLiteral(Literal): + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if instring[loc] == self.firstMatchChar: + return loc + 1, self.match + raise ParseException(instring, loc, self.errmsg, self) + + +ParserElement._literalStringClass = Literal + + +class Keyword(Token): + """ + Token to exactly match a specified string as a keyword, that is, + it must be immediately preceded and followed by whitespace or + non-keyword characters. Compare with :class:`Literal`: + + - ``Literal("if")`` will match the leading ``'if'`` in + ``'ifAndOnlyIf'``. + - ``Keyword("if")`` will not; it will only match the leading + ``'if'`` in ``'if x=1'``, or ``'if(y==2)'`` + + Accepts two optional constructor arguments in addition to the + keyword string: + + - ``ident_chars`` is a string of characters that would be valid + identifier characters, defaulting to all alphanumerics + "_" and + "$" + - ``caseless`` allows case-insensitive matching, default is ``False``. + + Example:: + + Keyword("start").parse_string("start") # -> ['start'] + Keyword("start").parse_string("starting") # -> Exception + + For case-insensitive matching, use :class:`CaselessKeyword`. + """ + + DEFAULT_KEYWORD_CHARS = alphanums + "_$" + + def __init__( + self, + match_string: str = "", + ident_chars: typing.Optional[str] = None, + caseless: bool = False, + *, + matchString: str = "", + identChars: typing.Optional[str] = None, + ) -> None: + super().__init__() + identChars = identChars or ident_chars + if identChars is None: + identChars = Keyword.DEFAULT_KEYWORD_CHARS + match_string = matchString or match_string + self.match = match_string + self.matchLen = len(match_string) + self.firstMatchChar = match_string[:1] + if not self.firstMatchChar: + raise ValueError("null string passed to Keyword; use Empty() instead") + self.errmsg = f"Expected {type(self).__name__} {self.name}" + self._may_return_empty = False + self.mayIndexError = False + self.caseless = caseless + if caseless: + self.caselessmatch = match_string.upper() + identChars = identChars.upper() + self.identChars = set(identChars) + + def _generateDefaultName(self) -> str: + return repr(self.match) + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + errmsg = self.errmsg or "" + errloc = loc + if self.caseless: + if instring[loc : loc + self.matchLen].upper() == self.caselessmatch: + if loc == 0 or instring[loc - 1].upper() not in self.identChars: + if ( + loc >= len(instring) - self.matchLen + or instring[loc + self.matchLen].upper() not in self.identChars + ): + return loc + self.matchLen, self.match + + # followed by keyword char + errmsg += ", was immediately followed by keyword character" + errloc = loc + self.matchLen + else: + # preceded by keyword char + errmsg += ", keyword was immediately preceded by keyword character" + errloc = loc - 1 + # else no match just raise plain exception + + elif ( + instring[loc] == self.firstMatchChar + and self.matchLen == 1 + or instring.startswith(self.match, loc) + ): + if loc == 0 or instring[loc - 1] not in self.identChars: + if ( + loc >= len(instring) - self.matchLen + or instring[loc + self.matchLen] not in self.identChars + ): + return loc + self.matchLen, self.match + + # followed by keyword char + errmsg += ", keyword was immediately followed by keyword character" + errloc = loc + self.matchLen + else: + # preceded by keyword char + errmsg += ", keyword was immediately preceded by keyword character" + errloc = loc - 1 + # else no match just raise plain exception + + raise ParseException(instring, errloc, errmsg, self) + + @staticmethod + def set_default_keyword_chars(chars) -> None: + """ + Overrides the default characters used by :class:`Keyword` expressions. + """ + Keyword.DEFAULT_KEYWORD_CHARS = chars + + # Compatibility synonyms + setDefaultKeywordChars = staticmethod( + replaced_by_pep8("setDefaultKeywordChars", set_default_keyword_chars) + ) + + +class CaselessLiteral(Literal): + """ + Token to match a specified string, ignoring case of letters. + Note: the matched results will always be in the case of the given + match string, NOT the case of the input text. + + Example:: + + CaselessLiteral("CMD")[1, ...].parse_string("cmd CMD Cmd10") + # -> ['CMD', 'CMD', 'CMD'] + + (Contrast with example for :class:`CaselessKeyword`.) + """ + + def __init__(self, match_string: str = "", *, matchString: str = "") -> None: + match_string = matchString or match_string + super().__init__(match_string.upper()) + # Preserve the defining literal. + self.returnString = match_string + self.errmsg = f"Expected {self.name}" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if instring[loc : loc + self.matchLen].upper() == self.match: + return loc + self.matchLen, self.returnString + raise ParseException(instring, loc, self.errmsg, self) + + +class CaselessKeyword(Keyword): + """ + Caseless version of :class:`Keyword`. + + Example:: + + CaselessKeyword("CMD")[1, ...].parse_string("cmd CMD Cmd10") + # -> ['CMD', 'CMD'] + + (Contrast with example for :class:`CaselessLiteral`.) + """ + + def __init__( + self, + match_string: str = "", + ident_chars: typing.Optional[str] = None, + *, + matchString: str = "", + identChars: typing.Optional[str] = None, + ) -> None: + identChars = identChars or ident_chars + match_string = matchString or match_string + super().__init__(match_string, identChars, caseless=True) + + +class CloseMatch(Token): + """A variation on :class:`Literal` which matches "close" matches, + that is, strings with at most 'n' mismatching characters. + :class:`CloseMatch` takes parameters: + + - ``match_string`` - string to be matched + - ``caseless`` - a boolean indicating whether to ignore casing when comparing characters + - ``max_mismatches`` - (``default=1``) maximum number of + mismatches allowed to count as a match + + The results from a successful parse will contain the matched text + from the input string and the following named results: + + - ``mismatches`` - a list of the positions within the + match_string where mismatches were found + - ``original`` - the original match_string used to compare + against the input string + + If ``mismatches`` is an empty list, then the match was an exact + match. + + Example:: + + patt = CloseMatch("ATCATCGAATGGA") + patt.parse_string("ATCATCGAAXGGA") # -> (['ATCATCGAAXGGA'], {'mismatches': [[9]], 'original': ['ATCATCGAATGGA']}) + patt.parse_string("ATCAXCGAAXGGA") # -> Exception: Expected 'ATCATCGAATGGA' (with up to 1 mismatches) (at char 0), (line:1, col:1) + + # exact match + patt.parse_string("ATCATCGAATGGA") # -> (['ATCATCGAATGGA'], {'mismatches': [[]], 'original': ['ATCATCGAATGGA']}) + + # close match allowing up to 2 mismatches + patt = CloseMatch("ATCATCGAATGGA", max_mismatches=2) + patt.parse_string("ATCAXCGAAXGGA") # -> (['ATCAXCGAAXGGA'], {'mismatches': [[4, 9]], 'original': ['ATCATCGAATGGA']}) + """ + + def __init__( + self, + match_string: str, + max_mismatches: typing.Optional[int] = None, + *, + maxMismatches: int = 1, + caseless=False, + ) -> None: + maxMismatches = max_mismatches if max_mismatches is not None else maxMismatches + super().__init__() + self.match_string = match_string + self.maxMismatches = maxMismatches + self.errmsg = f"Expected {self.match_string!r} (with up to {self.maxMismatches} mismatches)" + self.caseless = caseless + self.mayIndexError = False + self._may_return_empty = False + + def _generateDefaultName(self) -> str: + return f"{type(self).__name__}:{self.match_string!r}" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + start = loc + instrlen = len(instring) + maxloc = start + len(self.match_string) + + if maxloc <= instrlen: + match_string = self.match_string + match_stringloc = 0 + mismatches = [] + maxMismatches = self.maxMismatches + + for match_stringloc, s_m in enumerate( + zip(instring[loc:maxloc], match_string) + ): + src, mat = s_m + if self.caseless: + src, mat = src.lower(), mat.lower() + + if src != mat: + mismatches.append(match_stringloc) + if len(mismatches) > maxMismatches: + break + else: + loc = start + match_stringloc + 1 + results = ParseResults([instring[start:loc]]) + results["original"] = match_string + results["mismatches"] = mismatches + return loc, results + + raise ParseException(instring, loc, self.errmsg, self) + + +class Word(Token): + """Token for matching words composed of allowed character sets. + + Parameters: + + - ``init_chars`` - string of all characters that should be used to + match as a word; "ABC" will match "AAA", "ABAB", "CBAC", etc.; + if ``body_chars`` is also specified, then this is the string of + initial characters + - ``body_chars`` - string of characters that + can be used for matching after a matched initial character as + given in ``init_chars``; if omitted, same as the initial characters + (default=``None``) + - ``min`` - minimum number of characters to match (default=1) + - ``max`` - maximum number of characters to match (default=0) + - ``exact`` - exact number of characters to match (default=0) + - ``as_keyword`` - match as a keyword (default=``False``) + - ``exclude_chars`` - characters that might be + found in the input ``body_chars`` string but which should not be + accepted for matching ;useful to define a word of all + printables except for one or two characters, for instance + (default=``None``) + + :class:`srange` is useful for defining custom character set strings + for defining :class:`Word` expressions, using range notation from + regular expression character sets. + + A common mistake is to use :class:`Word` to match a specific literal + string, as in ``Word("Address")``. Remember that :class:`Word` + uses the string argument to define *sets* of matchable characters. + This expression would match "Add", "AAA", "dAred", or any other word + made up of the characters 'A', 'd', 'r', 'e', and 's'. To match an + exact literal string, use :class:`Literal` or :class:`Keyword`. + + pyparsing includes helper strings for building Words: + + - :class:`alphas` + - :class:`nums` + - :class:`alphanums` + - :class:`hexnums` + - :class:`alphas8bit` (alphabetic characters in ASCII range 128-255 + - accented, tilded, umlauted, etc.) + - :class:`punc8bit` (non-alphabetic characters in ASCII range + 128-255 - currency, symbols, superscripts, diacriticals, etc.) + - :class:`printables` (any non-whitespace character) + + ``alphas``, ``nums``, and ``printables`` are also defined in several + Unicode sets - see :class:`pyparsing_unicode``. + + Example:: + + # a word composed of digits + integer = Word(nums) # equivalent to Word("0123456789") or Word(srange("0-9")) + + # a word with a leading capital, and zero or more lowercase + capitalized_word = Word(alphas.upper(), alphas.lower()) + + # hostnames are alphanumeric, with leading alpha, and '-' + hostname = Word(alphas, alphanums + '-') + + # roman numeral (not a strict parser, accepts invalid mix of characters) + roman = Word("IVXLCDM") + + # any string of non-whitespace characters, except for ',' + csv_value = Word(printables, exclude_chars=",") + """ + + def __init__( + self, + init_chars: str = "", + body_chars: typing.Optional[str] = None, + min: int = 1, + max: int = 0, + exact: int = 0, + as_keyword: bool = False, + exclude_chars: typing.Optional[str] = None, + *, + initChars: typing.Optional[str] = None, + bodyChars: typing.Optional[str] = None, + asKeyword: bool = False, + excludeChars: typing.Optional[str] = None, + ) -> None: + initChars = initChars or init_chars + bodyChars = bodyChars or body_chars + asKeyword = asKeyword or as_keyword + excludeChars = excludeChars or exclude_chars + super().__init__() + if not initChars: + raise ValueError( + f"invalid {type(self).__name__}, initChars cannot be empty string" + ) + + initChars_set = set(initChars) + if excludeChars: + excludeChars_set = set(excludeChars) + initChars_set -= excludeChars_set + if bodyChars: + bodyChars = "".join(set(bodyChars) - excludeChars_set) + self.initChars = initChars_set + self.initCharsOrig = "".join(sorted(initChars_set)) + + if bodyChars: + self.bodyChars = set(bodyChars) + self.bodyCharsOrig = "".join(sorted(bodyChars)) + else: + self.bodyChars = initChars_set + self.bodyCharsOrig = self.initCharsOrig + + self.maxSpecified = max > 0 + + if min < 1: + raise ValueError( + "cannot specify a minimum length < 1; use Opt(Word()) if zero-length word is permitted" + ) + + if self.maxSpecified and min > max: + raise ValueError( + f"invalid args, if min and max both specified min must be <= max (min={min}, max={max})" + ) + + self.minLen = min + + if max > 0: + self.maxLen = max + else: + self.maxLen = _MAX_INT + + if exact > 0: + min = max = exact + self.maxLen = exact + self.minLen = exact + + self.errmsg = f"Expected {self.name}" + self.mayIndexError = False + self.asKeyword = asKeyword + if self.asKeyword: + self.errmsg += " as a keyword" + + # see if we can make a regex for this Word + if " " not in (self.initChars | self.bodyChars): + if len(self.initChars) == 1: + re_leading_fragment = re.escape(self.initCharsOrig) + else: + re_leading_fragment = f"[{_collapse_string_to_ranges(self.initChars)}]" + + if self.bodyChars == self.initChars: + if max == 0 and self.minLen == 1: + repeat = "+" + elif max == 1: + repeat = "" + else: + if self.minLen != self.maxLen: + repeat = f"{{{self.minLen},{'' if self.maxLen == _MAX_INT else self.maxLen}}}" + else: + repeat = f"{{{self.minLen}}}" + self.reString = f"{re_leading_fragment}{repeat}" + else: + if max == 1: + re_body_fragment = "" + repeat = "" + else: + re_body_fragment = f"[{_collapse_string_to_ranges(self.bodyChars)}]" + if max == 0 and self.minLen == 1: + repeat = "*" + elif max == 2: + repeat = "?" if min <= 1 else "" + else: + if min != max: + repeat = f"{{{min - 1 if min > 0 else ''},{max - 1 if max > 0 else ''}}}" + else: + repeat = f"{{{min - 1 if min > 0 else ''}}}" + + self.reString = f"{re_leading_fragment}{re_body_fragment}{repeat}" + + if self.asKeyword: + self.reString = rf"\b{self.reString}\b" + + try: + self.re = re.compile(self.reString) + except re.error: + self.re = None # type: ignore[assignment] + else: + self.re_match = self.re.match + self.parseImpl = self.parseImpl_regex # type: ignore[method-assign] + + def _generateDefaultName(self) -> str: + def charsAsStr(s): + max_repr_len = 16 + s = _collapse_string_to_ranges(s, re_escape=False) + + if len(s) > max_repr_len: + return s[: max_repr_len - 3] + "..." + + return s + + if self.initChars != self.bodyChars: + base = f"W:({charsAsStr(self.initChars)}, {charsAsStr(self.bodyChars)})" + else: + base = f"W:({charsAsStr(self.initChars)})" + + # add length specification + if self.minLen > 1 or self.maxLen != _MAX_INT: + if self.minLen == self.maxLen: + if self.minLen == 1: + return base[2:] + else: + return base + f"{{{self.minLen}}}" + elif self.maxLen == _MAX_INT: + return base + f"{{{self.minLen},...}}" + else: + return base + f"{{{self.minLen},{self.maxLen}}}" + return base + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if instring[loc] not in self.initChars: + raise ParseException(instring, loc, self.errmsg, self) + + start = loc + loc += 1 + instrlen = len(instring) + body_chars: set[str] = self.bodyChars + maxloc = start + self.maxLen + maxloc = min(maxloc, instrlen) + while loc < maxloc and instring[loc] in body_chars: + loc += 1 + + throw_exception = False + if loc - start < self.minLen: + throw_exception = True + elif self.maxSpecified and loc < instrlen and instring[loc] in body_chars: + throw_exception = True + elif self.asKeyword and ( + (start > 0 and instring[start - 1] in body_chars) + or (loc < instrlen and instring[loc] in body_chars) + ): + throw_exception = True + + if throw_exception: + raise ParseException(instring, loc, self.errmsg, self) + + return loc, instring[start:loc] + + def parseImpl_regex(self, instring, loc, do_actions=True) -> ParseImplReturnType: + result = self.re_match(instring, loc) + if not result: + raise ParseException(instring, loc, self.errmsg, self) + + loc = result.end() + return loc, result.group() + + +class Char(Word): + """A short-cut class for defining :class:`Word` ``(characters, exact=1)``, + when defining a match of any single character in a string of + characters. + """ + + def __init__( + self, + charset: str, + as_keyword: bool = False, + exclude_chars: typing.Optional[str] = None, + *, + asKeyword: bool = False, + excludeChars: typing.Optional[str] = None, + ) -> None: + asKeyword = asKeyword or as_keyword + excludeChars = excludeChars or exclude_chars + super().__init__( + charset, exact=1, as_keyword=asKeyword, exclude_chars=excludeChars + ) + + +class Regex(Token): + r"""Token for matching strings that match a given regular + expression. Defined with string specifying the regular expression in + a form recognized by the stdlib Python `re module `_. + If the given regex contains named groups (defined using ``(?P...)``), + these will be preserved as named :class:`ParseResults`. + + If instead of the Python stdlib ``re`` module you wish to use a different RE module + (such as the ``regex`` module), you can do so by building your ``Regex`` object with + a compiled RE that was compiled using ``regex``. + + Example:: + + realnum = Regex(r"[+-]?\d+\.\d*") + # ref: https://stackoverflow.com/questions/267399/how-do-you-match-only-valid-roman-numerals-with-a-regular-expression + roman = Regex(r"M{0,4}(CM|CD|D?{0,3})(XC|XL|L?X{0,3})(IX|IV|V?I{0,3})") + + # named fields in a regex will be returned as named results + date = Regex(r'(?P\d{4})-(?P\d\d?)-(?P\d\d?)') + + # the Regex class will accept re's compiled using the regex module + import regex + parser = pp.Regex(regex.compile(r'[0-9]')) + """ + + def __init__( + self, + pattern: Any, + flags: Union[re.RegexFlag, int] = 0, + as_group_list: bool = False, + as_match: bool = False, + *, + asGroupList: bool = False, + asMatch: bool = False, + ) -> None: + """The parameters ``pattern`` and ``flags`` are passed + to the ``re.compile()`` function as-is. See the Python + `re module `_ module for an + explanation of the acceptable patterns and flags. + """ + super().__init__() + asGroupList = asGroupList or as_group_list + asMatch = asMatch or as_match + + if isinstance(pattern, str_type): + if not pattern: + raise ValueError("null string passed to Regex; use Empty() instead") + + self._re = None + self._may_return_empty = None # type: ignore [assignment] + self.reString = self.pattern = pattern + + elif hasattr(pattern, "pattern") and hasattr(pattern, "match"): + self._re = pattern + self._may_return_empty = None # type: ignore [assignment] + self.pattern = self.reString = pattern.pattern + + elif callable(pattern): + # defer creating this pattern until we really need it + self.pattern = pattern + self._may_return_empty = None # type: ignore [assignment] + self._re = None + + else: + raise TypeError( + "Regex may only be constructed with a string or a compiled RE object," + " or a callable that takes no arguments and returns a string or a" + " compiled RE object" + ) + + self.flags = flags + self.errmsg = f"Expected {self.name}" + self.mayIndexError = False + self.asGroupList = asGroupList + self.asMatch = asMatch + if self.asGroupList: + self.parseImpl = self.parseImplAsGroupList # type: ignore [method-assign] + if self.asMatch: + self.parseImpl = self.parseImplAsMatch # type: ignore [method-assign] + + @cached_property + def re(self) -> re.Pattern: + if self._re: + return self._re + + if callable(self.pattern): + # replace self.pattern with the string returned by calling self.pattern() + self.pattern = cast(Callable[[], str], self.pattern)() + + # see if we got a compiled RE back instead of a str - if so, we're done + if hasattr(self.pattern, "pattern") and hasattr(self.pattern, "match"): + self._re = cast(re.Pattern[str], self.pattern) + self.pattern = self.reString = self._re.pattern + return self._re + + try: + self._re = re.compile(self.pattern, self.flags) + except re.error: + raise ValueError(f"invalid pattern ({self.pattern!r}) passed to Regex") + else: + self._may_return_empty = self.re.match("", pos=0) is not None + return self._re + + @cached_property + def re_match(self) -> Callable[[str, int], Any]: + return self.re.match + + @property + def mayReturnEmpty(self): + if self._may_return_empty is None: + # force compile of regex pattern, to set may_return_empty flag + self.re # noqa + return self._may_return_empty + + @mayReturnEmpty.setter + def mayReturnEmpty(self, value): + self._may_return_empty = value + + def _generateDefaultName(self) -> str: + unescaped = repr(self.pattern).replace("\\\\", "\\") + return f"Re:({unescaped})" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + # explicit check for matching past the length of the string; + # this is done because the re module will not complain about + # a match with `pos > len(instring)`, it will just return "" + if loc > len(instring) and self.mayReturnEmpty: + raise ParseException(instring, loc, self.errmsg, self) + + result = self.re_match(instring, loc) + if not result: + raise ParseException(instring, loc, self.errmsg, self) + + loc = result.end() + ret = ParseResults(result.group()) + d = result.groupdict() + + for k, v in d.items(): + ret[k] = v + + return loc, ret + + def parseImplAsGroupList(self, instring, loc, do_actions=True): + if loc > len(instring) and self.mayReturnEmpty: + raise ParseException(instring, loc, self.errmsg, self) + + result = self.re_match(instring, loc) + if not result: + raise ParseException(instring, loc, self.errmsg, self) + + loc = result.end() + ret = result.groups() + return loc, ret + + def parseImplAsMatch(self, instring, loc, do_actions=True): + if loc > len(instring) and self.mayReturnEmpty: + raise ParseException(instring, loc, self.errmsg, self) + + result = self.re_match(instring, loc) + if not result: + raise ParseException(instring, loc, self.errmsg, self) + + loc = result.end() + ret = result + return loc, ret + + def sub(self, repl: str) -> ParserElement: + r""" + Return :class:`Regex` with an attached parse action to transform the parsed + result as if called using `re.sub(expr, repl, string) `_. + + Example:: + + make_html = Regex(r"(\w+):(.*?):").sub(r"<\1>\2") + print(make_html.transform_string("h1:main title:")) + # prints "

main title

" + """ + if self.asGroupList: + raise TypeError("cannot use sub() with Regex(as_group_list=True)") + + if self.asMatch and callable(repl): + raise TypeError( + "cannot use sub() with a callable with Regex(as_match=True)" + ) + + if self.asMatch: + + def pa(tokens): + return tokens[0].expand(repl) + + else: + + def pa(tokens): + return self.re.sub(repl, tokens[0]) + + return self.add_parse_action(pa) + + +class QuotedString(Token): + r""" + Token for matching strings that are delimited by quoting characters. + + Defined with the following parameters: + + - ``quote_char`` - string of one or more characters defining the + quote delimiting string + - ``esc_char`` - character to re_escape quotes, typically backslash + (default= ``None``) + - ``esc_quote`` - special quote sequence to re_escape an embedded quote + string (such as SQL's ``""`` to re_escape an embedded ``"``) + (default= ``None``) + - ``multiline`` - boolean indicating whether quotes can span + multiple lines (default= ``False``) + - ``unquote_results`` - boolean indicating whether the matched text + should be unquoted (default= ``True``) + - ``end_quote_char`` - string of one or more characters defining the + end of the quote delimited string (default= ``None`` => same as + quote_char) + - ``convert_whitespace_escapes`` - convert escaped whitespace + (``'\t'``, ``'\n'``, etc.) to actual whitespace + (default= ``True``) + + Example:: + + qs = QuotedString('"') + print(qs.search_string('lsjdf "This is the quote" sldjf')) + complex_qs = QuotedString('{{', end_quote_char='}}') + print(complex_qs.search_string('lsjdf {{This is the "quote"}} sldjf')) + sql_qs = QuotedString('"', esc_quote='""') + print(sql_qs.search_string('lsjdf "This is the quote with ""embedded"" quotes" sldjf')) + + prints:: + + [['This is the quote']] + [['This is the "quote"']] + [['This is the quote with "embedded" quotes']] + """ + + ws_map = dict(((r"\t", "\t"), (r"\n", "\n"), (r"\f", "\f"), (r"\r", "\r"))) + + def __init__( + self, + quote_char: str = "", + esc_char: typing.Optional[str] = None, + esc_quote: typing.Optional[str] = None, + multiline: bool = False, + unquote_results: bool = True, + end_quote_char: typing.Optional[str] = None, + convert_whitespace_escapes: bool = True, + *, + quoteChar: str = "", + escChar: typing.Optional[str] = None, + escQuote: typing.Optional[str] = None, + unquoteResults: bool = True, + endQuoteChar: typing.Optional[str] = None, + convertWhitespaceEscapes: bool = True, + ) -> None: + super().__init__() + esc_char = escChar or esc_char + esc_quote = escQuote or esc_quote + unquote_results = unquoteResults and unquote_results + end_quote_char = endQuoteChar or end_quote_char + convert_whitespace_escapes = ( + convertWhitespaceEscapes and convert_whitespace_escapes + ) + quote_char = quoteChar or quote_char + + # remove white space from quote chars + quote_char = quote_char.strip() + if not quote_char: + raise ValueError("quote_char cannot be the empty string") + + if end_quote_char is None: + end_quote_char = quote_char + else: + end_quote_char = end_quote_char.strip() + if not end_quote_char: + raise ValueError("end_quote_char cannot be the empty string") + + self.quote_char: str = quote_char + self.quote_char_len: int = len(quote_char) + self.first_quote_char: str = quote_char[0] + self.end_quote_char: str = end_quote_char + self.end_quote_char_len: int = len(end_quote_char) + self.esc_char: str = esc_char or "" + self.has_esc_char: bool = esc_char is not None + self.esc_quote: str = esc_quote or "" + self.unquote_results: bool = unquote_results + self.convert_whitespace_escapes: bool = convert_whitespace_escapes + self.multiline = multiline + self.re_flags = re.RegexFlag(0) + + # fmt: off + # build up re pattern for the content between the quote delimiters + inner_pattern: list[str] = [] + + if esc_quote: + inner_pattern.append(rf"(?:{re.escape(esc_quote)})") + + if esc_char: + inner_pattern.append(rf"(?:{re.escape(esc_char)}.)") + + if len(self.end_quote_char) > 1: + inner_pattern.append( + "(?:" + + "|".join( + f"(?:{re.escape(self.end_quote_char[:i])}(?!{re.escape(self.end_quote_char[i:])}))" + for i in range(len(self.end_quote_char) - 1, 0, -1) + ) + + ")" + ) + + if self.multiline: + self.re_flags |= re.MULTILINE | re.DOTALL + inner_pattern.append( + rf"(?:[^{_escape_regex_range_chars(self.end_quote_char[0])}" + rf"{(_escape_regex_range_chars(self.esc_char) if self.has_esc_char else '')}])" + ) + else: + inner_pattern.append( + rf"(?:[^{_escape_regex_range_chars(self.end_quote_char[0])}\n\r" + rf"{(_escape_regex_range_chars(self.esc_char) if self.has_esc_char else '')}])" + ) + + self.pattern = "".join( + [ + re.escape(self.quote_char), + "(?:", + '|'.join(inner_pattern), + ")*", + re.escape(self.end_quote_char), + ] + ) + + if self.unquote_results: + if self.convert_whitespace_escapes: + self.unquote_scan_re = re.compile( + rf"({'|'.join(re.escape(k) for k in self.ws_map)})" + rf"|(\\[0-7]{3}|\\0|\\x[0-9a-fA-F]{2}|\\u[0-9a-fA-F]{4})" + rf"|({re.escape(self.esc_char)}.)" + rf"|(\n|.)", + flags=self.re_flags, + ) + else: + self.unquote_scan_re = re.compile( + rf"({re.escape(self.esc_char)}.)" + rf"|(\n|.)", + flags=self.re_flags + ) + # fmt: on + + try: + self.re = re.compile(self.pattern, self.re_flags) + self.reString = self.pattern + self.re_match = self.re.match + except re.error: + raise ValueError(f"invalid pattern {self.pattern!r} passed to Regex") + + self.errmsg = f"Expected {self.name}" + self.mayIndexError = False + self._may_return_empty = True + + def _generateDefaultName(self) -> str: + if self.quote_char == self.end_quote_char and isinstance( + self.quote_char, str_type + ): + return f"string enclosed in {self.quote_char!r}" + + return f"quoted string, starting with {self.quote_char} ending with {self.end_quote_char}" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + # check first character of opening quote to see if that is a match + # before doing the more complicated regex match + result = ( + instring[loc] == self.first_quote_char + and self.re_match(instring, loc) + or None + ) + if not result: + raise ParseException(instring, loc, self.errmsg, self) + + # get ending loc and matched string from regex matching result + loc = result.end() + ret = result.group() + + def convert_escaped_numerics(s: str) -> str: + if s == "0": + return "\0" + if s.isdigit() and len(s) == 3: + return chr(int(s, base=8)) + elif s.startswith(("u", "x")): + return chr(int(s[1:], base=16)) + else: + return s + + if self.unquote_results: + # strip off quotes + ret = ret[self.quote_char_len : -self.end_quote_char_len] + + if isinstance(ret, str_type): + # fmt: off + if self.convert_whitespace_escapes: + # as we iterate over matches in the input string, + # collect from whichever match group of the unquote_scan_re + # regex matches (only 1 group will match at any given time) + ret = "".join( + # match group 1 matches \t, \n, etc. + self.ws_map[match.group(1)] if match.group(1) + # match group 2 matches escaped octal, null, hex, and Unicode + # sequences + else convert_escaped_numerics(match.group(2)[1:]) if match.group(2) + # match group 3 matches escaped characters + else match.group(3)[-1] if match.group(3) + # match group 4 matches any character + else match.group(4) + for match in self.unquote_scan_re.finditer(ret) + ) + else: + ret = "".join( + # match group 1 matches escaped characters + match.group(1)[-1] if match.group(1) + # match group 2 matches any character + else match.group(2) + for match in self.unquote_scan_re.finditer(ret) + ) + # fmt: on + + # replace escaped quotes + if self.esc_quote: + ret = ret.replace(self.esc_quote, self.end_quote_char) + + return loc, ret + + +class CharsNotIn(Token): + """Token for matching words composed of characters *not* in a given + set (will include whitespace in matched characters if not listed in + the provided exclusion set - see example). Defined with string + containing all disallowed characters, and an optional minimum, + maximum, and/or exact length. The default value for ``min`` is + 1 (a minimum value < 1 is not valid); the default values for + ``max`` and ``exact`` are 0, meaning no maximum or exact + length restriction. + + Example:: + + # define a comma-separated-value as anything that is not a ',' + csv_value = CharsNotIn(',') + print(DelimitedList(csv_value).parse_string("dkls,lsdkjf,s12 34,@!#,213")) + + prints:: + + ['dkls', 'lsdkjf', 's12 34', '@!#', '213'] + """ + + def __init__( + self, + not_chars: str = "", + min: int = 1, + max: int = 0, + exact: int = 0, + *, + notChars: str = "", + ) -> None: + super().__init__() + self.skipWhitespace = False + self.notChars = not_chars or notChars + self.notCharsSet = set(self.notChars) + + if min < 1: + raise ValueError( + "cannot specify a minimum length < 1; use" + " Opt(CharsNotIn()) if zero-length char group is permitted" + ) + + self.minLen = min + + if max > 0: + self.maxLen = max + else: + self.maxLen = _MAX_INT + + if exact > 0: + self.maxLen = exact + self.minLen = exact + + self.errmsg = f"Expected {self.name}" + self._may_return_empty = self.minLen == 0 + self.mayIndexError = False + + def _generateDefaultName(self) -> str: + not_chars_str = _collapse_string_to_ranges(self.notChars) + if len(not_chars_str) > 16: + return f"!W:({self.notChars[: 16 - 3]}...)" + else: + return f"!W:({self.notChars})" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + notchars = self.notCharsSet + if instring[loc] in notchars: + raise ParseException(instring, loc, self.errmsg, self) + + start = loc + loc += 1 + maxlen = min(start + self.maxLen, len(instring)) + while loc < maxlen and instring[loc] not in notchars: + loc += 1 + + if loc - start < self.minLen: + raise ParseException(instring, loc, self.errmsg, self) + + return loc, instring[start:loc] + + +class White(Token): + """Special matching class for matching whitespace. Normally, + whitespace is ignored by pyparsing grammars. This class is included + when some whitespace structures are significant. Define with + a string containing the whitespace characters to be matched; default + is ``" \\t\\r\\n"``. Also takes optional ``min``, + ``max``, and ``exact`` arguments, as defined for the + :class:`Word` class. + """ + + whiteStrs = { + " ": "", + "\t": "", + "\n": "", + "\r": "", + "\f": "", + "\u00A0": "", + "\u1680": "", + "\u180E": "", + "\u2000": "", + "\u2001": "", + "\u2002": "", + "\u2003": "", + "\u2004": "", + "\u2005": "", + "\u2006": "", + "\u2007": "", + "\u2008": "", + "\u2009": "", + "\u200A": "", + "\u200B": "", + "\u202F": "", + "\u205F": "", + "\u3000": "", + } + + def __init__( + self, ws: str = " \t\r\n", min: int = 1, max: int = 0, exact: int = 0 + ) -> None: + super().__init__() + self.matchWhite = ws + self.set_whitespace_chars( + "".join(c for c in self.whiteStrs if c not in self.matchWhite), + copy_defaults=True, + ) + # self.leave_whitespace() + self._may_return_empty = True + self.errmsg = f"Expected {self.name}" + + self.minLen = min + + if max > 0: + self.maxLen = max + else: + self.maxLen = _MAX_INT + + if exact > 0: + self.maxLen = exact + self.minLen = exact + + def _generateDefaultName(self) -> str: + return "".join(White.whiteStrs[c] for c in self.matchWhite) + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if instring[loc] not in self.matchWhite: + raise ParseException(instring, loc, self.errmsg, self) + start = loc + loc += 1 + maxloc = start + self.maxLen + maxloc = min(maxloc, len(instring)) + while loc < maxloc and instring[loc] in self.matchWhite: + loc += 1 + + if loc - start < self.minLen: + raise ParseException(instring, loc, self.errmsg, self) + + return loc, instring[start:loc] + + +class PositionToken(Token): + def __init__(self) -> None: + super().__init__() + self._may_return_empty = True + self.mayIndexError = False + + +class GoToColumn(PositionToken): + """Token to advance to a specific column of input text; useful for + tabular report scraping. + """ + + def __init__(self, colno: int) -> None: + super().__init__() + self.col = colno + + def preParse(self, instring: str, loc: int) -> int: + if col(loc, instring) == self.col: + return loc + + instrlen = len(instring) + if self.ignoreExprs: + loc = self._skipIgnorables(instring, loc) + while ( + loc < instrlen + and instring[loc].isspace() + and col(loc, instring) != self.col + ): + loc += 1 + + return loc + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + thiscol = col(loc, instring) + if thiscol > self.col: + raise ParseException(instring, loc, "Text not in expected column", self) + newloc = loc + self.col - thiscol + ret = instring[loc:newloc] + return newloc, ret + + +class LineStart(PositionToken): + r"""Matches if current position is at the beginning of a line within + the parse string + + Example:: + + test = '''\ + AAA this line + AAA and this line + AAA but not this one + B AAA and definitely not this one + ''' + + for t in (LineStart() + 'AAA' + rest_of_line).search_string(test): + print(t) + + prints:: + + ['AAA', ' this line'] + ['AAA', ' and this line'] + + """ + + def __init__(self) -> None: + super().__init__() + self.leave_whitespace() + self.orig_whiteChars = set() | self.whiteChars + self.whiteChars.discard("\n") + self.skipper = Empty().set_whitespace_chars(self.whiteChars) + self.set_name("start of line") + + def preParse(self, instring: str, loc: int) -> int: + if loc == 0: + return loc + + ret = self.skipper.preParse(instring, loc) + + if "\n" in self.orig_whiteChars: + while instring[ret : ret + 1] == "\n": + ret = self.skipper.preParse(instring, ret + 1) + + return ret + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if col(loc, instring) == 1: + return loc, [] + raise ParseException(instring, loc, self.errmsg, self) + + +class LineEnd(PositionToken): + """Matches if current position is at the end of a line within the + parse string + """ + + def __init__(self) -> None: + super().__init__() + self.whiteChars.discard("\n") + self.set_whitespace_chars(self.whiteChars, copy_defaults=False) + self.set_name("end of line") + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if loc < len(instring): + if instring[loc] == "\n": + return loc + 1, "\n" + else: + raise ParseException(instring, loc, self.errmsg, self) + elif loc == len(instring): + return loc + 1, [] + else: + raise ParseException(instring, loc, self.errmsg, self) + + +class StringStart(PositionToken): + """Matches if current position is at the beginning of the parse + string + """ + + def __init__(self) -> None: + super().__init__() + self.set_name("start of text") + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + # see if entire string up to here is just whitespace and ignoreables + if loc != 0 and loc != self.preParse(instring, 0): + raise ParseException(instring, loc, self.errmsg, self) + + return loc, [] + + +class StringEnd(PositionToken): + """ + Matches if current position is at the end of the parse string + """ + + def __init__(self) -> None: + super().__init__() + self.set_name("end of text") + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if loc < len(instring): + raise ParseException(instring, loc, self.errmsg, self) + if loc == len(instring): + return loc + 1, [] + if loc > len(instring): + return loc, [] + + raise ParseException(instring, loc, self.errmsg, self) + + +class WordStart(PositionToken): + """Matches if the current position is at the beginning of a + :class:`Word`, and is not preceded by any character in a given + set of ``word_chars`` (default= ``printables``). To emulate the + ``\b`` behavior of regular expressions, use + ``WordStart(alphanums)``. ``WordStart`` will also match at + the beginning of the string being parsed, or at the beginning of + a line. + """ + + def __init__( + self, word_chars: str = printables, *, wordChars: str = printables + ) -> None: + wordChars = word_chars if wordChars == printables else wordChars + super().__init__() + self.wordChars = set(wordChars) + self.set_name("start of a word") + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if loc != 0: + if ( + instring[loc - 1] in self.wordChars + or instring[loc] not in self.wordChars + ): + raise ParseException(instring, loc, self.errmsg, self) + return loc, [] + + +class WordEnd(PositionToken): + """Matches if the current position is at the end of a :class:`Word`, + and is not followed by any character in a given set of ``word_chars`` + (default= ``printables``). To emulate the ``\b`` behavior of + regular expressions, use ``WordEnd(alphanums)``. ``WordEnd`` + will also match at the end of the string being parsed, or at the end + of a line. + """ + + def __init__( + self, word_chars: str = printables, *, wordChars: str = printables + ) -> None: + wordChars = word_chars if wordChars == printables else wordChars + super().__init__() + self.wordChars = set(wordChars) + self.skipWhitespace = False + self.set_name("end of a word") + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + instrlen = len(instring) + if instrlen > 0 and loc < instrlen: + if ( + instring[loc] in self.wordChars + or instring[loc - 1] not in self.wordChars + ): + raise ParseException(instring, loc, self.errmsg, self) + return loc, [] + + +class Tag(Token): + """ + A meta-element for inserting a named result into the parsed + tokens that may be checked later in a parse action or while + processing the parsed results. Accepts an optional tag value, + defaulting to `True`. + + Example:: + + end_punc = "." | ("!" + Tag("enthusiastic"))) + greeting = "Hello," + Word(alphas) + end_punc + + result = greeting.parse_string("Hello, World.") + print(result.dump()) + + result = greeting.parse_string("Hello, World!") + print(result.dump()) + + prints:: + + ['Hello,', 'World', '.'] + + ['Hello,', 'World', '!'] + - enthusiastic: True + """ + + def __init__(self, tag_name: str, value: Any = True) -> None: + super().__init__() + self._may_return_empty = True + self.mayIndexError = False + self.leave_whitespace() + self.tag_name = tag_name + self.tag_value = value + self.add_parse_action(self._add_tag) + self.show_in_diagram = False + + def _add_tag(self, tokens: ParseResults): + tokens[self.tag_name] = self.tag_value + + def _generateDefaultName(self) -> str: + return f"{type(self).__name__}:{self.tag_name}={self.tag_value!r}" + + +class ParseExpression(ParserElement): + """Abstract subclass of ParserElement, for combining and + post-processing parsed tokens. + """ + + def __init__( + self, exprs: typing.Iterable[ParserElement], savelist: bool = False + ) -> None: + super().__init__(savelist) + self.exprs: list[ParserElement] + if isinstance(exprs, _generatorType): + exprs = list(exprs) + + if isinstance(exprs, str_type): + self.exprs = [self._literalStringClass(exprs)] + elif isinstance(exprs, ParserElement): + self.exprs = [exprs] + elif isinstance(exprs, Iterable): + exprs = list(exprs) + # if sequence of strings provided, wrap with Literal + if any(isinstance(expr, str_type) for expr in exprs): + exprs = ( + self._literalStringClass(e) if isinstance(e, str_type) else e + for e in exprs + ) + self.exprs = list(exprs) + else: + try: + self.exprs = list(exprs) + except TypeError: + self.exprs = [exprs] + self.callPreparse = False + + def recurse(self) -> list[ParserElement]: + return self.exprs[:] + + def append(self, other) -> ParserElement: + self.exprs.append(other) + self._defaultName = None + return self + + def leave_whitespace(self, recursive: bool = True) -> ParserElement: + """ + Extends ``leave_whitespace`` defined in base class, and also invokes ``leave_whitespace`` on + all contained expressions. + """ + super().leave_whitespace(recursive) + + if recursive: + self.exprs = [e.copy() for e in self.exprs] + for e in self.exprs: + e.leave_whitespace(recursive) + return self + + def ignore_whitespace(self, recursive: bool = True) -> ParserElement: + """ + Extends ``ignore_whitespace`` defined in base class, and also invokes ``leave_whitespace`` on + all contained expressions. + """ + super().ignore_whitespace(recursive) + if recursive: + self.exprs = [e.copy() for e in self.exprs] + for e in self.exprs: + e.ignore_whitespace(recursive) + return self + + def ignore(self, other) -> ParserElement: + if isinstance(other, Suppress): + if other not in self.ignoreExprs: + super().ignore(other) + for e in self.exprs: + e.ignore(self.ignoreExprs[-1]) + else: + super().ignore(other) + for e in self.exprs: + e.ignore(self.ignoreExprs[-1]) + return self + + def _generateDefaultName(self) -> str: + return f"{type(self).__name__}:({self.exprs})" + + def streamline(self) -> ParserElement: + if self.streamlined: + return self + + super().streamline() + + for e in self.exprs: + e.streamline() + + # collapse nested :class:`And`'s of the form ``And(And(And(a, b), c), d)`` to ``And(a, b, c, d)`` + # but only if there are no parse actions or resultsNames on the nested And's + # (likewise for :class:`Or`'s and :class:`MatchFirst`'s) + if len(self.exprs) == 2: + other = self.exprs[0] + if ( + isinstance(other, self.__class__) + and not other.parseAction + and other.resultsName is None + and not other.debug + ): + self.exprs = other.exprs[:] + [self.exprs[1]] + self._defaultName = None + self._may_return_empty |= other.mayReturnEmpty + self.mayIndexError |= other.mayIndexError + + other = self.exprs[-1] + if ( + isinstance(other, self.__class__) + and not other.parseAction + and other.resultsName is None + and not other.debug + ): + self.exprs = self.exprs[:-1] + other.exprs[:] + self._defaultName = None + self._may_return_empty |= other.mayReturnEmpty + self.mayIndexError |= other.mayIndexError + + self.errmsg = f"Expected {self}" + + return self + + def validate(self, validateTrace=None) -> None: + warnings.warn( + "ParserElement.validate() is deprecated, and should not be used to check for left recursion", + DeprecationWarning, + stacklevel=2, + ) + tmp = (validateTrace if validateTrace is not None else [])[:] + [self] + for e in self.exprs: + e.validate(tmp) + self._checkRecursion([]) + + def copy(self) -> ParserElement: + ret = super().copy() + ret = typing.cast(ParseExpression, ret) + ret.exprs = [e.copy() for e in self.exprs] + return ret + + def _setResultsName(self, name, list_all_matches=False) -> ParserElement: + if not ( + __diag__.warn_ungrouped_named_tokens_in_collection + and Diagnostics.warn_ungrouped_named_tokens_in_collection + not in self.suppress_warnings_ + ): + return super()._setResultsName(name, list_all_matches) + + for e in self.exprs: + if ( + isinstance(e, ParserElement) + and e.resultsName + and ( + Diagnostics.warn_ungrouped_named_tokens_in_collection + not in e.suppress_warnings_ + ) + ): + warning = ( + "warn_ungrouped_named_tokens_in_collection:" + f" setting results name {name!r} on {type(self).__name__} expression" + f" collides with {e.resultsName!r} on contained expression" + ) + warnings.warn(warning, stacklevel=3) + break + + return super()._setResultsName(name, list_all_matches) + + # Compatibility synonyms + # fmt: off + leaveWhitespace = replaced_by_pep8("leaveWhitespace", leave_whitespace) + ignoreWhitespace = replaced_by_pep8("ignoreWhitespace", ignore_whitespace) + # fmt: on + + +class And(ParseExpression): + """ + Requires all given :class:`ParserElement` s to be found in the given order. + Expressions may be separated by whitespace. + May be constructed using the ``'+'`` operator. + May also be constructed using the ``'-'`` operator, which will + suppress backtracking. + + Example:: + + integer = Word(nums) + name_expr = Word(alphas)[1, ...] + + expr = And([integer("id"), name_expr("name"), integer("age")]) + # more easily written as: + expr = integer("id") + name_expr("name") + integer("age") + """ + + class _ErrorStop(Empty): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.leave_whitespace() + + def _generateDefaultName(self) -> str: + return "-" + + def __init__( + self, + exprs_arg: typing.Iterable[Union[ParserElement, str]], + savelist: bool = True, + ) -> None: + # instantiate exprs as a list, converting strs to ParserElements + exprs: list[ParserElement] = [ + self._literalStringClass(e) if isinstance(e, str) else e for e in exprs_arg + ] + + # convert any Ellipsis elements to SkipTo + if Ellipsis in exprs: + + # Ellipsis cannot be the last element + if exprs[-1] is Ellipsis: + raise Exception("cannot construct And with sequence ending in ...") + + tmp: list[ParserElement] = [] + for cur_expr, next_expr in zip(exprs, exprs[1:]): + if cur_expr is Ellipsis: + tmp.append(SkipTo(next_expr)("_skipped*")) + else: + tmp.append(cur_expr) + + exprs[:-1] = tmp + + super().__init__(exprs, savelist) + if self.exprs: + self._may_return_empty = all(e.mayReturnEmpty for e in self.exprs) + if not isinstance(self.exprs[0], White): + self.set_whitespace_chars( + self.exprs[0].whiteChars, + copy_defaults=self.exprs[0].copyDefaultWhiteChars, + ) + self.skipWhitespace = self.exprs[0].skipWhitespace + else: + self.skipWhitespace = False + else: + self._may_return_empty = True + self.callPreparse = True + + def streamline(self) -> ParserElement: + # collapse any _PendingSkip's + if self.exprs and any( + isinstance(e, ParseExpression) + and e.exprs + and isinstance(e.exprs[-1], _PendingSkip) + for e in self.exprs[:-1] + ): + deleted_expr_marker = NoMatch() + for i, e in enumerate(self.exprs[:-1]): + if e is deleted_expr_marker: + continue + if ( + isinstance(e, ParseExpression) + and e.exprs + and isinstance(e.exprs[-1], _PendingSkip) + ): + e.exprs[-1] = e.exprs[-1] + self.exprs[i + 1] + self.exprs[i + 1] = deleted_expr_marker + self.exprs = [e for e in self.exprs if e is not deleted_expr_marker] + + super().streamline() + + # link any IndentedBlocks to the prior expression + prev: ParserElement + cur: ParserElement + for prev, cur in zip(self.exprs, self.exprs[1:]): + # traverse cur or any first embedded expr of cur looking for an IndentedBlock + # (but watch out for recursive grammar) + seen = set() + while True: + if id(cur) in seen: + break + seen.add(id(cur)) + if isinstance(cur, IndentedBlock): + prev.add_parse_action( + lambda s, l, t, cur_=cur: setattr( + cur_, "parent_anchor", col(l, s) + ) + ) + break + subs = cur.recurse() + next_first = next(iter(subs), None) + if next_first is None: + break + cur = typing.cast(ParserElement, next_first) + + self._may_return_empty = all(e.mayReturnEmpty for e in self.exprs) + return self + + def parseImpl(self, instring, loc, do_actions=True): + # pass False as callPreParse arg to _parse for first element, since we already + # pre-parsed the string as part of our And pre-parsing + loc, resultlist = self.exprs[0]._parse( + instring, loc, do_actions, callPreParse=False + ) + errorStop = False + for e in self.exprs[1:]: + # if isinstance(e, And._ErrorStop): + if type(e) is And._ErrorStop: + errorStop = True + continue + if errorStop: + try: + loc, exprtokens = e._parse(instring, loc, do_actions) + except ParseSyntaxException: + raise + except ParseBaseException as pe: + pe.__traceback__ = None + raise ParseSyntaxException._from_exception(pe) + except IndexError: + raise ParseSyntaxException( + instring, len(instring), self.errmsg, self + ) + else: + loc, exprtokens = e._parse(instring, loc, do_actions) + resultlist += exprtokens + return loc, resultlist + + def __iadd__(self, other): + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return self.append(other) # And([self, other]) + + def _checkRecursion(self, parseElementList): + subRecCheckList = parseElementList[:] + [self] + for e in self.exprs: + e._checkRecursion(subRecCheckList) + if not e.mayReturnEmpty: + break + + def _generateDefaultName(self) -> str: + inner = " ".join(str(e) for e in self.exprs) + # strip off redundant inner {}'s + while len(inner) > 1 and inner[0 :: len(inner) - 1] == "{}": + inner = inner[1:-1] + return f"{{{inner}}}" + + +class Or(ParseExpression): + """Requires that at least one :class:`ParserElement` is found. If + two expressions match, the expression that matches the longest + string will be used. May be constructed using the ``'^'`` + operator. + + Example:: + + # construct Or using '^' operator + + number = Word(nums) ^ Combine(Word(nums) + '.' + Word(nums)) + print(number.search_string("123 3.1416 789")) + + prints:: + + [['123'], ['3.1416'], ['789']] + """ + + def __init__( + self, exprs: typing.Iterable[ParserElement], savelist: bool = False + ) -> None: + super().__init__(exprs, savelist) + if self.exprs: + self._may_return_empty = any(e.mayReturnEmpty for e in self.exprs) + self.skipWhitespace = all(e.skipWhitespace for e in self.exprs) + else: + self._may_return_empty = True + + def streamline(self) -> ParserElement: + super().streamline() + if self.exprs: + self._may_return_empty = any(e.mayReturnEmpty for e in self.exprs) + self.saveAsList = any(e.saveAsList for e in self.exprs) + self.skipWhitespace = all( + e.skipWhitespace and not isinstance(e, White) for e in self.exprs + ) + else: + self.saveAsList = False + return self + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + maxExcLoc = -1 + maxException = None + matches: list[tuple[int, ParserElement]] = [] + fatals: list[ParseFatalException] = [] + if all(e.callPreparse for e in self.exprs): + loc = self.preParse(instring, loc) + for e in self.exprs: + try: + loc2 = e.try_parse(instring, loc, raise_fatal=True) + except ParseFatalException as pfe: + pfe.__traceback__ = None + pfe.parser_element = e + fatals.append(pfe) + maxException = None + maxExcLoc = -1 + except ParseException as err: + if not fatals: + err.__traceback__ = None + if err.loc > maxExcLoc: + maxException = err + maxExcLoc = err.loc + except IndexError: + if len(instring) > maxExcLoc: + maxException = ParseException( + instring, len(instring), e.errmsg, self + ) + maxExcLoc = len(instring) + else: + # save match among all matches, to retry longest to shortest + matches.append((loc2, e)) + + if matches: + # re-evaluate all matches in descending order of length of match, in case attached actions + # might change whether or how much they match of the input. + matches.sort(key=itemgetter(0), reverse=True) + + if not do_actions: + # no further conditions or parse actions to change the selection of + # alternative, so the first match will be the best match + best_expr = matches[0][1] + return best_expr._parse(instring, loc, do_actions) + + longest: tuple[int, typing.Optional[ParseResults]] = -1, None + for loc1, expr1 in matches: + if loc1 <= longest[0]: + # already have a longer match than this one will deliver, we are done + return longest + + try: + loc2, toks = expr1._parse(instring, loc, do_actions) + except ParseException as err: + err.__traceback__ = None + if err.loc > maxExcLoc: + maxException = err + maxExcLoc = err.loc + else: + if loc2 >= loc1: + return loc2, toks + # didn't match as much as before + elif loc2 > longest[0]: + longest = loc2, toks + + if longest != (-1, None): + return longest + + if fatals: + if len(fatals) > 1: + fatals.sort(key=lambda e: -e.loc) + if fatals[0].loc == fatals[1].loc: + fatals.sort(key=lambda e: (-e.loc, -len(str(e.parser_element)))) + max_fatal = fatals[0] + raise max_fatal + + if maxException is not None: + # infer from this check that all alternatives failed at the current position + # so emit this collective error message instead of any single error message + parse_start_loc = self.preParse(instring, loc) + if maxExcLoc == parse_start_loc: + maxException.msg = self.errmsg or "" + raise maxException + + raise ParseException(instring, loc, "no defined alternatives to match", self) + + def __ixor__(self, other): + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return self.append(other) # Or([self, other]) + + def _generateDefaultName(self) -> str: + return f"{{{' ^ '.join(str(e) for e in self.exprs)}}}" + + def _setResultsName(self, name, list_all_matches=False) -> ParserElement: + if ( + __diag__.warn_multiple_tokens_in_named_alternation + and Diagnostics.warn_multiple_tokens_in_named_alternation + not in self.suppress_warnings_ + ): + if any( + isinstance(e, And) + and Diagnostics.warn_multiple_tokens_in_named_alternation + not in e.suppress_warnings_ + for e in self.exprs + ): + warning = ( + "warn_multiple_tokens_in_named_alternation:" + f" setting results name {name!r} on {type(self).__name__} expression" + " will return a list of all parsed tokens in an And alternative," + " in prior versions only the first token was returned; enclose" + " contained argument in Group" + ) + warnings.warn(warning, stacklevel=3) + + return super()._setResultsName(name, list_all_matches) + + +class MatchFirst(ParseExpression): + """Requires that at least one :class:`ParserElement` is found. If + more than one expression matches, the first one listed is the one that will + match. May be constructed using the ``'|'`` operator. + + Example:: + + # construct MatchFirst using '|' operator + + # watch the order of expressions to match + number = Word(nums) | Combine(Word(nums) + '.' + Word(nums)) + print(number.search_string("123 3.1416 789")) # Fail! -> [['123'], ['3'], ['1416'], ['789']] + + # put more selective expression first + number = Combine(Word(nums) + '.' + Word(nums)) | Word(nums) + print(number.search_string("123 3.1416 789")) # Better -> [['123'], ['3.1416'], ['789']] + """ + + def __init__( + self, exprs: typing.Iterable[ParserElement], savelist: bool = False + ) -> None: + super().__init__(exprs, savelist) + if self.exprs: + self._may_return_empty = any(e.mayReturnEmpty for e in self.exprs) + self.skipWhitespace = all(e.skipWhitespace for e in self.exprs) + else: + self._may_return_empty = True + + def streamline(self) -> ParserElement: + if self.streamlined: + return self + + super().streamline() + if self.exprs: + self.saveAsList = any(e.saveAsList for e in self.exprs) + self._may_return_empty = any(e.mayReturnEmpty for e in self.exprs) + self.skipWhitespace = all( + e.skipWhitespace and not isinstance(e, White) for e in self.exprs + ) + else: + self.saveAsList = False + self._may_return_empty = True + return self + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + maxExcLoc = -1 + maxException = None + + for e in self.exprs: + try: + return e._parse(instring, loc, do_actions) + except ParseFatalException as pfe: + pfe.__traceback__ = None + pfe.parser_element = e + raise + except ParseException as err: + if err.loc > maxExcLoc: + maxException = err + maxExcLoc = err.loc + except IndexError: + if len(instring) > maxExcLoc: + maxException = ParseException( + instring, len(instring), e.errmsg, self + ) + maxExcLoc = len(instring) + + if maxException is not None: + # infer from this check that all alternatives failed at the current position + # so emit this collective error message instead of any individual error message + parse_start_loc = self.preParse(instring, loc) + if maxExcLoc == parse_start_loc: + maxException.msg = self.errmsg or "" + raise maxException + + raise ParseException(instring, loc, "no defined alternatives to match", self) + + def __ior__(self, other): + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return self.append(other) # MatchFirst([self, other]) + + def _generateDefaultName(self) -> str: + return f"{{{' | '.join(str(e) for e in self.exprs)}}}" + + def _setResultsName(self, name, list_all_matches=False) -> ParserElement: + if ( + __diag__.warn_multiple_tokens_in_named_alternation + and Diagnostics.warn_multiple_tokens_in_named_alternation + not in self.suppress_warnings_ + ): + if any( + isinstance(e, And) + and Diagnostics.warn_multiple_tokens_in_named_alternation + not in e.suppress_warnings_ + for e in self.exprs + ): + warning = ( + "warn_multiple_tokens_in_named_alternation:" + f" setting results name {name!r} on {type(self).__name__} expression" + " will return a list of all parsed tokens in an And alternative," + " in prior versions only the first token was returned; enclose" + " contained argument in Group" + ) + warnings.warn(warning, stacklevel=3) + + return super()._setResultsName(name, list_all_matches) + + +class Each(ParseExpression): + """Requires all given :class:`ParserElement` s to be found, but in + any order. Expressions may be separated by whitespace. + + May be constructed using the ``'&'`` operator. + + Example:: + + color = one_of("RED ORANGE YELLOW GREEN BLUE PURPLE BLACK WHITE BROWN") + shape_type = one_of("SQUARE CIRCLE TRIANGLE STAR HEXAGON OCTAGON") + integer = Word(nums) + shape_attr = "shape:" + shape_type("shape") + posn_attr = "posn:" + Group(integer("x") + ',' + integer("y"))("posn") + color_attr = "color:" + color("color") + size_attr = "size:" + integer("size") + + # use Each (using operator '&') to accept attributes in any order + # (shape and posn are required, color and size are optional) + shape_spec = shape_attr & posn_attr & Opt(color_attr) & Opt(size_attr) + + shape_spec.run_tests(''' + shape: SQUARE color: BLACK posn: 100, 120 + shape: CIRCLE size: 50 color: BLUE posn: 50,80 + color:GREEN size:20 shape:TRIANGLE posn:20,40 + ''' + ) + + prints:: + + shape: SQUARE color: BLACK posn: 100, 120 + ['shape:', 'SQUARE', 'color:', 'BLACK', 'posn:', ['100', ',', '120']] + - color: BLACK + - posn: ['100', ',', '120'] + - x: 100 + - y: 120 + - shape: SQUARE + + + shape: CIRCLE size: 50 color: BLUE posn: 50,80 + ['shape:', 'CIRCLE', 'size:', '50', 'color:', 'BLUE', 'posn:', ['50', ',', '80']] + - color: BLUE + - posn: ['50', ',', '80'] + - x: 50 + - y: 80 + - shape: CIRCLE + - size: 50 + + + color: GREEN size: 20 shape: TRIANGLE posn: 20,40 + ['color:', 'GREEN', 'size:', '20', 'shape:', 'TRIANGLE', 'posn:', ['20', ',', '40']] + - color: GREEN + - posn: ['20', ',', '40'] + - x: 20 + - y: 40 + - shape: TRIANGLE + - size: 20 + """ + + def __init__( + self, exprs: typing.Iterable[ParserElement], savelist: bool = True + ) -> None: + super().__init__(exprs, savelist) + if self.exprs: + self._may_return_empty = all(e.mayReturnEmpty for e in self.exprs) + else: + self._may_return_empty = True + self.skipWhitespace = True + self.initExprGroups = True + self.saveAsList = True + + def __iand__(self, other): + if isinstance(other, str_type): + other = self._literalStringClass(other) + if not isinstance(other, ParserElement): + return NotImplemented + return self.append(other) # Each([self, other]) + + def streamline(self) -> ParserElement: + super().streamline() + if self.exprs: + self._may_return_empty = all(e.mayReturnEmpty for e in self.exprs) + else: + self._may_return_empty = True + return self + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if self.initExprGroups: + self.opt1map = dict( + (id(e.expr), e) for e in self.exprs if isinstance(e, Opt) + ) + opt1 = [e.expr for e in self.exprs if isinstance(e, Opt)] + opt2 = [ + e + for e in self.exprs + if e.mayReturnEmpty and not isinstance(e, (Opt, Regex, ZeroOrMore)) + ] + self.optionals = opt1 + opt2 + self.multioptionals = [ + e.expr.set_results_name(e.resultsName, list_all_matches=True) + for e in self.exprs + if isinstance(e, _MultipleMatch) + ] + self.multirequired = [ + e.expr.set_results_name(e.resultsName, list_all_matches=True) + for e in self.exprs + if isinstance(e, OneOrMore) + ] + self.required = [ + e for e in self.exprs if not isinstance(e, (Opt, ZeroOrMore, OneOrMore)) + ] + self.required += self.multirequired + self.initExprGroups = False + + tmpLoc = loc + tmpReqd = self.required[:] + tmpOpt = self.optionals[:] + multis = self.multioptionals[:] + matchOrder: list[ParserElement] = [] + + keepMatching = True + failed: list[ParserElement] = [] + fatals: list[ParseFatalException] = [] + while keepMatching: + tmpExprs = tmpReqd + tmpOpt + multis + failed.clear() + fatals.clear() + for e in tmpExprs: + try: + tmpLoc = e.try_parse(instring, tmpLoc, raise_fatal=True) + except ParseFatalException as pfe: + pfe.__traceback__ = None + pfe.parser_element = e + fatals.append(pfe) + failed.append(e) + except ParseException: + failed.append(e) + else: + matchOrder.append(self.opt1map.get(id(e), e)) + if e in tmpReqd: + tmpReqd.remove(e) + elif e in tmpOpt: + tmpOpt.remove(e) + if len(failed) == len(tmpExprs): + keepMatching = False + + # look for any ParseFatalExceptions + if fatals: + if len(fatals) > 1: + fatals.sort(key=lambda e: -e.loc) + if fatals[0].loc == fatals[1].loc: + fatals.sort(key=lambda e: (-e.loc, -len(str(e.parser_element)))) + max_fatal = fatals[0] + raise max_fatal + + if tmpReqd: + missing = ", ".join([str(e) for e in tmpReqd]) + raise ParseException( + instring, + loc, + f"Missing one or more required elements ({missing})", + ) + + # add any unmatched Opts, in case they have default values defined + matchOrder += [e for e in self.exprs if isinstance(e, Opt) and e.expr in tmpOpt] + + total_results = ParseResults([]) + for e in matchOrder: + loc, results = e._parse(instring, loc, do_actions) + total_results += results + + return loc, total_results + + def _generateDefaultName(self) -> str: + return f"{{{' & '.join(str(e) for e in self.exprs)}}}" + + +class ParseElementEnhance(ParserElement): + """Abstract subclass of :class:`ParserElement`, for combining and + post-processing parsed tokens. + """ + + def __init__(self, expr: Union[ParserElement, str], savelist: bool = False) -> None: + super().__init__(savelist) + if isinstance(expr, str_type): + expr_str = typing.cast(str, expr) + if issubclass(self._literalStringClass, Token): + expr = self._literalStringClass(expr_str) # type: ignore[call-arg] + elif issubclass(type(self), self._literalStringClass): + expr = Literal(expr_str) + else: + expr = self._literalStringClass(Literal(expr_str)) # type: ignore[assignment, call-arg] + expr = typing.cast(ParserElement, expr) + self.expr = expr + if expr is not None: + self.mayIndexError = expr.mayIndexError + self._may_return_empty = expr.mayReturnEmpty + self.set_whitespace_chars( + expr.whiteChars, copy_defaults=expr.copyDefaultWhiteChars + ) + self.skipWhitespace = expr.skipWhitespace + self.saveAsList = expr.saveAsList + self.callPreparse = expr.callPreparse + self.ignoreExprs.extend(expr.ignoreExprs) + + def recurse(self) -> list[ParserElement]: + return [self.expr] if self.expr is not None else [] + + def parseImpl(self, instring, loc, do_actions=True): + if self.expr is None: + raise ParseException(instring, loc, "No expression defined", self) + + try: + return self.expr._parse(instring, loc, do_actions, callPreParse=False) + except ParseSyntaxException: + raise + except ParseBaseException as pbe: + pbe.pstr = pbe.pstr or instring + pbe.loc = pbe.loc or loc + pbe.parser_element = pbe.parser_element or self + if not isinstance(self, Forward) and self.customName is not None: + if self.errmsg: + pbe.msg = self.errmsg + raise + + def leave_whitespace(self, recursive: bool = True) -> ParserElement: + super().leave_whitespace(recursive) + + if recursive: + if self.expr is not None: + self.expr = self.expr.copy() + self.expr.leave_whitespace(recursive) + return self + + def ignore_whitespace(self, recursive: bool = True) -> ParserElement: + super().ignore_whitespace(recursive) + + if recursive: + if self.expr is not None: + self.expr = self.expr.copy() + self.expr.ignore_whitespace(recursive) + return self + + def ignore(self, other) -> ParserElement: + if not isinstance(other, Suppress) or other not in self.ignoreExprs: + super().ignore(other) + if self.expr is not None: + self.expr.ignore(self.ignoreExprs[-1]) + + return self + + def streamline(self) -> ParserElement: + super().streamline() + if self.expr is not None: + self.expr.streamline() + return self + + def _checkRecursion(self, parseElementList): + if self in parseElementList: + raise RecursiveGrammarException(parseElementList + [self]) + subRecCheckList = parseElementList[:] + [self] + if self.expr is not None: + self.expr._checkRecursion(subRecCheckList) + + def validate(self, validateTrace=None) -> None: + warnings.warn( + "ParserElement.validate() is deprecated, and should not be used to check for left recursion", + DeprecationWarning, + stacklevel=2, + ) + if validateTrace is None: + validateTrace = [] + tmp = validateTrace[:] + [self] + if self.expr is not None: + self.expr.validate(tmp) + self._checkRecursion([]) + + def _generateDefaultName(self) -> str: + return f"{type(self).__name__}:({self.expr})" + + # Compatibility synonyms + # fmt: off + leaveWhitespace = replaced_by_pep8("leaveWhitespace", leave_whitespace) + ignoreWhitespace = replaced_by_pep8("ignoreWhitespace", ignore_whitespace) + # fmt: on + + +class IndentedBlock(ParseElementEnhance): + """ + Expression to match one or more expressions at a given indentation level. + Useful for parsing text where structure is implied by indentation (like Python source code). + """ + + class _Indent(Empty): + def __init__(self, ref_col: int) -> None: + super().__init__() + self.errmsg = f"expected indent at column {ref_col}" + self.add_condition(lambda s, l, t: col(l, s) == ref_col) + + class _IndentGreater(Empty): + def __init__(self, ref_col: int) -> None: + super().__init__() + self.errmsg = f"expected indent at column greater than {ref_col}" + self.add_condition(lambda s, l, t: col(l, s) > ref_col) + + def __init__( + self, expr: ParserElement, *, recursive: bool = False, grouped: bool = True + ) -> None: + super().__init__(expr, savelist=True) + # if recursive: + # raise NotImplementedError("IndentedBlock with recursive is not implemented") + self._recursive = recursive + self._grouped = grouped + self.parent_anchor = 1 + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + # advance parse position to non-whitespace by using an Empty() + # this should be the column to be used for all subsequent indented lines + anchor_loc = Empty().preParse(instring, loc) + + # see if self.expr matches at the current location - if not it will raise an exception + # and no further work is necessary + self.expr.try_parse(instring, anchor_loc, do_actions=do_actions) + + indent_col = col(anchor_loc, instring) + peer_detect_expr = self._Indent(indent_col) + + inner_expr = Empty() + peer_detect_expr + self.expr + if self._recursive: + sub_indent = self._IndentGreater(indent_col) + nested_block = IndentedBlock( + self.expr, recursive=self._recursive, grouped=self._grouped + ) + nested_block.set_debug(self.debug) + nested_block.parent_anchor = indent_col + inner_expr += Opt(sub_indent + nested_block) + + inner_expr.set_name(f"inner {hex(id(inner_expr))[-4:].upper()}@{indent_col}") + block = OneOrMore(inner_expr) + + trailing_undent = self._Indent(self.parent_anchor) | StringEnd() + + if self._grouped: + wrapper = Group + else: + wrapper = lambda expr: expr # type: ignore[misc, assignment] + return (wrapper(block) + Optional(trailing_undent)).parseImpl( + instring, anchor_loc, do_actions + ) + + +class AtStringStart(ParseElementEnhance): + """Matches if expression matches at the beginning of the parse + string:: + + AtStringStart(Word(nums)).parse_string("123") + # prints ["123"] + + AtStringStart(Word(nums)).parse_string(" 123") + # raises ParseException + """ + + def __init__(self, expr: Union[ParserElement, str]) -> None: + super().__init__(expr) + self.callPreparse = False + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if loc != 0: + raise ParseException(instring, loc, "not found at string start") + return super().parseImpl(instring, loc, do_actions) + + +class AtLineStart(ParseElementEnhance): + r"""Matches if an expression matches at the beginning of a line within + the parse string + + Example:: + + test = '''\ + AAA this line + AAA and this line + AAA but not this one + B AAA and definitely not this one + ''' + + for t in (AtLineStart('AAA') + rest_of_line).search_string(test): + print(t) + + prints:: + + ['AAA', ' this line'] + ['AAA', ' and this line'] + + """ + + def __init__(self, expr: Union[ParserElement, str]) -> None: + super().__init__(expr) + self.callPreparse = False + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if col(loc, instring) != 1: + raise ParseException(instring, loc, "not found at line start") + return super().parseImpl(instring, loc, do_actions) + + +class FollowedBy(ParseElementEnhance): + """Lookahead matching of the given parse expression. + ``FollowedBy`` does *not* advance the parsing position within + the input string, it only verifies that the specified parse + expression matches at the current position. ``FollowedBy`` + always returns a null token list. If any results names are defined + in the lookahead expression, those *will* be returned for access by + name. + + Example:: + + # use FollowedBy to match a label only if it is followed by a ':' + data_word = Word(alphas) + label = data_word + FollowedBy(':') + attr_expr = Group(label + Suppress(':') + OneOrMore(data_word, stop_on=label).set_parse_action(' '.join)) + + attr_expr[1, ...].parse_string("shape: SQUARE color: BLACK posn: upper left").pprint() + + prints:: + + [['shape', 'SQUARE'], ['color', 'BLACK'], ['posn', 'upper left']] + """ + + def __init__(self, expr: Union[ParserElement, str]) -> None: + super().__init__(expr) + self._may_return_empty = True + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + # by using self._expr.parse and deleting the contents of the returned ParseResults list + # we keep any named results that were defined in the FollowedBy expression + _, ret = self.expr._parse(instring, loc, do_actions=do_actions) + del ret[:] + + return loc, ret + + +class PrecededBy(ParseElementEnhance): + """Lookbehind matching of the given parse expression. + ``PrecededBy`` does not advance the parsing position within the + input string, it only verifies that the specified parse expression + matches prior to the current position. ``PrecededBy`` always + returns a null token list, but if a results name is defined on the + given expression, it is returned. + + Parameters: + + - ``expr`` - expression that must match prior to the current parse + location + - ``retreat`` - (default= ``None``) - (int) maximum number of characters + to lookbehind prior to the current parse location + + If the lookbehind expression is a string, :class:`Literal`, + :class:`Keyword`, or a :class:`Word` or :class:`CharsNotIn` + with a specified exact or maximum length, then the retreat + parameter is not required. Otherwise, retreat must be specified to + give a maximum number of characters to look back from + the current parse position for a lookbehind match. + + Example:: + + # VB-style variable names with type prefixes + int_var = PrecededBy("#") + pyparsing_common.identifier + str_var = PrecededBy("$") + pyparsing_common.identifier + + """ + + def __init__(self, expr: Union[ParserElement, str], retreat: int = 0) -> None: + super().__init__(expr) + self.expr = self.expr().leave_whitespace() + self._may_return_empty = True + self.mayIndexError = False + self.exact = False + if isinstance(expr, str_type): + expr = typing.cast(str, expr) + retreat = len(expr) + self.exact = True + elif isinstance(expr, (Literal, Keyword)): + retreat = expr.matchLen + self.exact = True + elif isinstance(expr, (Word, CharsNotIn)) and expr.maxLen != _MAX_INT: + retreat = expr.maxLen + self.exact = True + elif isinstance(expr, PositionToken): + retreat = 0 + self.exact = True + self.retreat = retreat + self.errmsg = f"not preceded by {expr}" + self.skipWhitespace = False + self.parseAction.append(lambda s, l, t: t.__delitem__(slice(None, None))) + + def parseImpl(self, instring, loc=0, do_actions=True) -> ParseImplReturnType: + if self.exact: + if loc < self.retreat: + raise ParseException(instring, loc, self.errmsg, self) + start = loc - self.retreat + _, ret = self.expr._parse(instring, start) + return loc, ret + + # retreat specified a maximum lookbehind window, iterate + test_expr = self.expr + StringEnd() + instring_slice = instring[max(0, loc - self.retreat) : loc] + last_expr: ParseBaseException = ParseException(instring, loc, self.errmsg, self) + + for offset in range(1, min(loc, self.retreat + 1) + 1): + try: + # print('trying', offset, instring_slice, repr(instring_slice[loc - offset:])) + _, ret = test_expr._parse(instring_slice, len(instring_slice) - offset) + except ParseBaseException as pbe: + last_expr = pbe + else: + break + else: + raise last_expr + + return loc, ret + + +class Located(ParseElementEnhance): + """ + Decorates a returned token with its starting and ending + locations in the input string. + + This helper adds the following results names: + + - ``locn_start`` - location where matched expression begins + - ``locn_end`` - location where matched expression ends + - ``value`` - the actual parsed results + + Be careful if the input text contains ```` characters, you + may want to call :class:`ParserElement.parse_with_tabs` + + Example:: + + wd = Word(alphas) + for match in Located(wd).search_string("ljsdf123lksdjjf123lkkjj1222"): + print(match) + + prints:: + + [0, ['ljsdf'], 5] + [8, ['lksdjjf'], 15] + [18, ['lkkjj'], 23] + + """ + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + start = loc + loc, tokens = self.expr._parse(instring, start, do_actions, callPreParse=False) + ret_tokens = ParseResults([start, tokens, loc]) + ret_tokens["locn_start"] = start + ret_tokens["value"] = tokens + ret_tokens["locn_end"] = loc + if self.resultsName: + # must return as a list, so that the name will be attached to the complete group + return loc, [ret_tokens] + else: + return loc, ret_tokens + + +class NotAny(ParseElementEnhance): + """ + Lookahead to disallow matching with the given parse expression. + ``NotAny`` does *not* advance the parsing position within the + input string, it only verifies that the specified parse expression + does *not* match at the current position. Also, ``NotAny`` does + *not* skip over leading whitespace. ``NotAny`` always returns + a null token list. May be constructed using the ``'~'`` operator. + + Example:: + + AND, OR, NOT = map(CaselessKeyword, "AND OR NOT".split()) + + # take care not to mistake keywords for identifiers + ident = ~(AND | OR | NOT) + Word(alphas) + boolean_term = Opt(NOT) + ident + + # very crude boolean expression - to support parenthesis groups and + # operation hierarchy, use infix_notation + boolean_expr = boolean_term + ((AND | OR) + boolean_term)[...] + + # integers that are followed by "." are actually floats + integer = Word(nums) + ~Char(".") + """ + + def __init__(self, expr: Union[ParserElement, str]) -> None: + super().__init__(expr) + # do NOT use self.leave_whitespace(), don't want to propagate to exprs + # self.leave_whitespace() + self.skipWhitespace = False + + self._may_return_empty = True + self.errmsg = f"Found unwanted token, {self.expr}" + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if self.expr.can_parse_next(instring, loc, do_actions=do_actions): + raise ParseException(instring, loc, self.errmsg, self) + return loc, [] + + def _generateDefaultName(self) -> str: + return f"~{{{self.expr}}}" + + +class _MultipleMatch(ParseElementEnhance): + def __init__( + self, + expr: Union[str, ParserElement], + stop_on: typing.Optional[Union[ParserElement, str]] = None, + *, + stopOn: typing.Optional[Union[ParserElement, str]] = None, + ) -> None: + super().__init__(expr) + stopOn = stopOn or stop_on + self.saveAsList = True + ender = stopOn + if isinstance(ender, str_type): + ender = self._literalStringClass(ender) + self.stopOn(ender) + + def stopOn(self, ender) -> ParserElement: + if isinstance(ender, str_type): + ender = self._literalStringClass(ender) + self.not_ender = ~ender if ender is not None else None + return self + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + self_expr_parse = self.expr._parse + self_skip_ignorables = self._skipIgnorables + check_ender = False + if self.not_ender is not None: + try_not_ender = self.not_ender.try_parse + check_ender = True + + # must be at least one (but first see if we are the stopOn sentinel; + # if so, fail) + if check_ender: + try_not_ender(instring, loc) + loc, tokens = self_expr_parse(instring, loc, do_actions) + try: + hasIgnoreExprs = not not self.ignoreExprs + while 1: + if check_ender: + try_not_ender(instring, loc) + if hasIgnoreExprs: + preloc = self_skip_ignorables(instring, loc) + else: + preloc = loc + loc, tmptokens = self_expr_parse(instring, preloc, do_actions) + tokens += tmptokens + except (ParseException, IndexError): + pass + + return loc, tokens + + def _setResultsName(self, name, list_all_matches=False) -> ParserElement: + if ( + __diag__.warn_ungrouped_named_tokens_in_collection + and Diagnostics.warn_ungrouped_named_tokens_in_collection + not in self.suppress_warnings_ + ): + for e in [self.expr] + self.expr.recurse(): + if ( + isinstance(e, ParserElement) + and e.resultsName + and ( + Diagnostics.warn_ungrouped_named_tokens_in_collection + not in e.suppress_warnings_ + ) + ): + warning = ( + "warn_ungrouped_named_tokens_in_collection:" + f" setting results name {name!r} on {type(self).__name__} expression" + f" collides with {e.resultsName!r} on contained expression" + ) + warnings.warn(warning, stacklevel=3) + break + + return super()._setResultsName(name, list_all_matches) + + +class OneOrMore(_MultipleMatch): + """ + Repetition of one or more of the given expression. + + Parameters: + + - ``expr`` - expression that must match one or more times + - ``stop_on`` - (default= ``None``) - expression for a terminating sentinel + (only required if the sentinel would ordinarily match the repetition + expression) + + Example:: + + data_word = Word(alphas) + label = data_word + FollowedBy(':') + attr_expr = Group(label + Suppress(':') + OneOrMore(data_word).set_parse_action(' '.join)) + + text = "shape: SQUARE posn: upper left color: BLACK" + attr_expr[1, ...].parse_string(text).pprint() # Fail! read 'color' as data instead of next label -> [['shape', 'SQUARE color']] + + # use stop_on attribute for OneOrMore to avoid reading label string as part of the data + attr_expr = Group(label + Suppress(':') + OneOrMore(data_word, stop_on=label).set_parse_action(' '.join)) + OneOrMore(attr_expr).parse_string(text).pprint() # Better -> [['shape', 'SQUARE'], ['posn', 'upper left'], ['color', 'BLACK']] + + # could also be written as + (attr_expr * (1,)).parse_string(text).pprint() + """ + + def _generateDefaultName(self) -> str: + return f"{{{self.expr}}}..." + + +class ZeroOrMore(_MultipleMatch): + """ + Optional repetition of zero or more of the given expression. + + Parameters: + + - ``expr`` - expression that must match zero or more times + - ``stop_on`` - expression for a terminating sentinel + (only required if the sentinel would ordinarily match the repetition + expression) - (default= ``None``) + + Example: similar to :class:`OneOrMore` + """ + + def __init__( + self, + expr: Union[str, ParserElement], + stop_on: typing.Optional[Union[ParserElement, str]] = None, + *, + stopOn: typing.Optional[Union[ParserElement, str]] = None, + ) -> None: + super().__init__(expr, stopOn=stopOn or stop_on) + self._may_return_empty = True + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + try: + return super().parseImpl(instring, loc, do_actions) + except (ParseException, IndexError): + return loc, ParseResults([], name=self.resultsName) + + def _generateDefaultName(self) -> str: + return f"[{self.expr}]..." + + +class DelimitedList(ParseElementEnhance): + def __init__( + self, + expr: Union[str, ParserElement], + delim: Union[str, ParserElement] = ",", + combine: bool = False, + min: typing.Optional[int] = None, + max: typing.Optional[int] = None, + *, + allow_trailing_delim: bool = False, + ) -> None: + """Helper to define a delimited list of expressions - the delimiter + defaults to ','. By default, the list elements and delimiters can + have intervening whitespace, and comments, but this can be + overridden by passing ``combine=True`` in the constructor. If + ``combine`` is set to ``True``, the matching tokens are + returned as a single token string, with the delimiters included; + otherwise, the matching tokens are returned as a list of tokens, + with the delimiters suppressed. + + If ``allow_trailing_delim`` is set to True, then the list may end with + a delimiter. + + Example:: + + DelimitedList(Word(alphas)).parse_string("aa,bb,cc") # -> ['aa', 'bb', 'cc'] + DelimitedList(Word(hexnums), delim=':', combine=True).parse_string("AA:BB:CC:DD:EE") # -> ['AA:BB:CC:DD:EE'] + """ + if isinstance(expr, str_type): + expr = ParserElement._literalStringClass(expr) + expr = typing.cast(ParserElement, expr) + + if min is not None and min < 1: + raise ValueError("min must be greater than 0") + + if max is not None and min is not None and max < min: + raise ValueError("max must be greater than, or equal to min") + + self.content = expr + self.raw_delim = str(delim) + self.delim = delim + self.combine = combine + if not combine: + self.delim = Suppress(delim) + self.min = min or 1 + self.max = max + self.allow_trailing_delim = allow_trailing_delim + + delim_list_expr = self.content + (self.delim + self.content) * ( + self.min - 1, + None if self.max is None else self.max - 1, + ) + if self.allow_trailing_delim: + delim_list_expr += Opt(self.delim) + + if self.combine: + delim_list_expr = Combine(delim_list_expr) + + super().__init__(delim_list_expr, savelist=True) + + def _generateDefaultName(self) -> str: + content_expr = self.content.streamline() + return f"{content_expr} [{self.raw_delim} {content_expr}]..." + + +class _NullToken: + def __bool__(self): + return False + + def __str__(self): + return "" + + +class Opt(ParseElementEnhance): + """ + Optional matching of the given expression. + + Parameters: + + - ``expr`` - expression that must match zero or more times + - ``default`` (optional) - value to be returned if the optional expression is not found. + + Example:: + + # US postal code can be a 5-digit zip, plus optional 4-digit qualifier + zip = Combine(Word(nums, exact=5) + Opt('-' + Word(nums, exact=4))) + zip.run_tests(''' + # traditional ZIP code + 12345 + + # ZIP+4 form + 12101-0001 + + # invalid ZIP + 98765- + ''') + + prints:: + + # traditional ZIP code + 12345 + ['12345'] + + # ZIP+4 form + 12101-0001 + ['12101-0001'] + + # invalid ZIP + 98765- + ^ + FAIL: Expected end of text (at char 5), (line:1, col:6) + """ + + __optionalNotMatched = _NullToken() + + def __init__( + self, expr: Union[ParserElement, str], default: Any = __optionalNotMatched + ) -> None: + super().__init__(expr, savelist=False) + self.saveAsList = self.expr.saveAsList + self.defaultValue = default + self._may_return_empty = True + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + self_expr = self.expr + try: + loc, tokens = self_expr._parse( + instring, loc, do_actions, callPreParse=False + ) + except (ParseException, IndexError): + default_value = self.defaultValue + if default_value is not self.__optionalNotMatched: + if self_expr.resultsName: + tokens = ParseResults([default_value]) + tokens[self_expr.resultsName] = default_value + else: + tokens = [default_value] # type: ignore[assignment] + else: + tokens = [] # type: ignore[assignment] + return loc, tokens + + def _generateDefaultName(self) -> str: + inner = str(self.expr) + # strip off redundant inner {}'s + while len(inner) > 1 and inner[0 :: len(inner) - 1] == "{}": + inner = inner[1:-1] + return f"[{inner}]" + + +Optional = Opt + + +class SkipTo(ParseElementEnhance): + """ + Token for skipping over all undefined text until the matched + expression is found. + + Parameters: + + - ``expr`` - target expression marking the end of the data to be skipped + - ``include`` - if ``True``, the target expression is also parsed + (the skipped text and target expression are returned as a 2-element + list) (default= ``False``). + - ``ignore`` - (default= ``None``) used to define grammars (typically quoted strings and + comments) that might contain false matches to the target expression + - ``fail_on`` - (default= ``None``) define expressions that are not allowed to be + included in the skipped test; if found before the target expression is found, + the :class:`SkipTo` is not a match + + Example:: + + report = ''' + Outstanding Issues Report - 1 Jan 2000 + + # | Severity | Description | Days Open + -----+----------+-------------------------------------------+----------- + 101 | Critical | Intermittent system crash | 6 + 94 | Cosmetic | Spelling error on Login ('log|n') | 14 + 79 | Minor | System slow when running too many reports | 47 + ''' + integer = Word(nums) + SEP = Suppress('|') + # use SkipTo to simply match everything up until the next SEP + # - ignore quoted strings, so that a '|' character inside a quoted string does not match + # - parse action will call token.strip() for each matched token, i.e., the description body + string_data = SkipTo(SEP, ignore=quoted_string) + string_data.set_parse_action(token_map(str.strip)) + ticket_expr = (integer("issue_num") + SEP + + string_data("sev") + SEP + + string_data("desc") + SEP + + integer("days_open")) + + for tkt in ticket_expr.search_string(report): + print tkt.dump() + + prints:: + + ['101', 'Critical', 'Intermittent system crash', '6'] + - days_open: '6' + - desc: 'Intermittent system crash' + - issue_num: '101' + - sev: 'Critical' + ['94', 'Cosmetic', "Spelling error on Login ('log|n')", '14'] + - days_open: '14' + - desc: "Spelling error on Login ('log|n')" + - issue_num: '94' + - sev: 'Cosmetic' + ['79', 'Minor', 'System slow when running too many reports', '47'] + - days_open: '47' + - desc: 'System slow when running too many reports' + - issue_num: '79' + - sev: 'Minor' + """ + + def __init__( + self, + other: Union[ParserElement, str], + include: bool = False, + ignore: typing.Optional[Union[ParserElement, str]] = None, + fail_on: typing.Optional[Union[ParserElement, str]] = None, + *, + failOn: typing.Optional[Union[ParserElement, str]] = None, + ) -> None: + super().__init__(other) + failOn = failOn or fail_on + self.ignoreExpr = ignore + self._may_return_empty = True + self.mayIndexError = False + self.includeMatch = include + self.saveAsList = False + if isinstance(failOn, str_type): + self.failOn = self._literalStringClass(failOn) + else: + self.failOn = failOn + self.errmsg = f"No match found for {self.expr}" + self.ignorer = Empty().leave_whitespace() + self._update_ignorer() + + def _update_ignorer(self): + # rebuild internal ignore expr from current ignore exprs and assigned ignoreExpr + self.ignorer.ignoreExprs.clear() + for e in self.expr.ignoreExprs: + self.ignorer.ignore(e) + if self.ignoreExpr: + self.ignorer.ignore(self.ignoreExpr) + + def ignore(self, expr): + super().ignore(expr) + self._update_ignorer() + + def parseImpl(self, instring, loc, do_actions=True): + startloc = loc + instrlen = len(instring) + self_expr_parse = self.expr._parse + self_failOn_canParseNext = ( + self.failOn.canParseNext if self.failOn is not None else None + ) + ignorer_try_parse = self.ignorer.try_parse if self.ignorer.ignoreExprs else None + + tmploc = loc + while tmploc <= instrlen: + if self_failOn_canParseNext is not None: + # break if failOn expression matches + if self_failOn_canParseNext(instring, tmploc): + break + + if ignorer_try_parse is not None: + # advance past ignore expressions + prev_tmploc = tmploc + while 1: + try: + tmploc = ignorer_try_parse(instring, tmploc) + except ParseBaseException: + break + # see if all ignorers matched, but didn't actually ignore anything + if tmploc == prev_tmploc: + break + prev_tmploc = tmploc + + try: + self_expr_parse(instring, tmploc, do_actions=False, callPreParse=False) + except (ParseException, IndexError): + # no match, advance loc in string + tmploc += 1 + else: + # matched skipto expr, done + break + + else: + # ran off the end of the input string without matching skipto expr, fail + raise ParseException(instring, loc, self.errmsg, self) + + # build up return values + loc = tmploc + skiptext = instring[startloc:loc] + skipresult = ParseResults(skiptext) + + if self.includeMatch: + loc, mat = self_expr_parse(instring, loc, do_actions, callPreParse=False) + skipresult += mat + + return loc, skipresult + + +class Forward(ParseElementEnhance): + """ + Forward declaration of an expression to be defined later - + used for recursive grammars, such as algebraic infix notation. + When the expression is known, it is assigned to the ``Forward`` + variable using the ``'<<'`` operator. + + Note: take care when assigning to ``Forward`` not to overlook + precedence of operators. + + Specifically, ``'|'`` has a lower precedence than ``'<<'``, so that:: + + fwd_expr << a | b | c + + will actually be evaluated as:: + + (fwd_expr << a) | b | c + + thereby leaving b and c out as parseable alternatives. It is recommended that you + explicitly group the values inserted into the ``Forward``:: + + fwd_expr << (a | b | c) + + Converting to use the ``'<<='`` operator instead will avoid this problem. + + See :class:`ParseResults.pprint` for an example of a recursive + parser created using ``Forward``. + """ + + def __init__( + self, other: typing.Optional[Union[ParserElement, str]] = None + ) -> None: + self.caller_frame = traceback.extract_stack(limit=2)[0] + super().__init__(other, savelist=False) # type: ignore[arg-type] + self.lshift_line = None + + def __lshift__(self, other) -> Forward: + if hasattr(self, "caller_frame"): + del self.caller_frame + if isinstance(other, str_type): + other = self._literalStringClass(other) + + if not isinstance(other, ParserElement): + return NotImplemented + + self.expr = other + self.streamlined = other.streamlined + self.mayIndexError = self.expr.mayIndexError + self._may_return_empty = self.expr.mayReturnEmpty + self.set_whitespace_chars( + self.expr.whiteChars, copy_defaults=self.expr.copyDefaultWhiteChars + ) + self.skipWhitespace = self.expr.skipWhitespace + self.saveAsList = self.expr.saveAsList + self.ignoreExprs.extend(self.expr.ignoreExprs) + self.lshift_line = traceback.extract_stack(limit=2)[-2] # type: ignore[assignment] + return self + + def __ilshift__(self, other) -> Forward: + if not isinstance(other, ParserElement): + return NotImplemented + + return self << other + + def __or__(self, other) -> ParserElement: + caller_line = traceback.extract_stack(limit=2)[-2] + if ( + __diag__.warn_on_match_first_with_lshift_operator + and caller_line == self.lshift_line + and Diagnostics.warn_on_match_first_with_lshift_operator + not in self.suppress_warnings_ + ): + warnings.warn( + "warn_on_match_first_with_lshift_operator:" + " using '<<' operator with '|' is probably an error, use '<<='", + stacklevel=2, + ) + ret = super().__or__(other) + return ret + + def __del__(self): + # see if we are getting dropped because of '=' reassignment of var instead of '<<=' or '<<' + if ( + self.expr is None + and __diag__.warn_on_assignment_to_Forward + and Diagnostics.warn_on_assignment_to_Forward not in self.suppress_warnings_ + ): + warnings.warn_explicit( + "warn_on_assignment_to_Forward:" + " Forward defined here but no expression attached later using '<<=' or '<<'", + UserWarning, + filename=self.caller_frame.filename, + lineno=self.caller_frame.lineno, + ) + + def parseImpl(self, instring, loc, do_actions=True) -> ParseImplReturnType: + if ( + self.expr is None + and __diag__.warn_on_parse_using_empty_Forward + and Diagnostics.warn_on_parse_using_empty_Forward + not in self.suppress_warnings_ + ): + # walk stack until parse_string, scan_string, search_string, or transform_string is found + parse_fns = ( + "parse_string", + "scan_string", + "search_string", + "transform_string", + ) + tb = traceback.extract_stack(limit=200) + for i, frm in enumerate(reversed(tb), start=1): + if frm.name in parse_fns: + stacklevel = i + 1 + break + else: + stacklevel = 2 + warnings.warn( + "warn_on_parse_using_empty_Forward:" + " Forward expression was never assigned a value, will not parse any input", + stacklevel=stacklevel, + ) + if not ParserElement._left_recursion_enabled: + return super().parseImpl(instring, loc, do_actions) + # ## Bounded Recursion algorithm ## + # Recursion only needs to be processed at ``Forward`` elements, since they are + # the only ones that can actually refer to themselves. The general idea is + # to handle recursion stepwise: We start at no recursion, then recurse once, + # recurse twice, ..., until more recursion offers no benefit (we hit the bound). + # + # The "trick" here is that each ``Forward`` gets evaluated in two contexts + # - to *match* a specific recursion level, and + # - to *search* the bounded recursion level + # and the two run concurrently. The *search* must *match* each recursion level + # to find the best possible match. This is handled by a memo table, which + # provides the previous match to the next level match attempt. + # + # See also "Left Recursion in Parsing Expression Grammars", Medeiros et al. + # + # There is a complication since we not only *parse* but also *transform* via + # actions: We do not want to run the actions too often while expanding. Thus, + # we expand using `do_actions=False` and only run `do_actions=True` if the next + # recursion level is acceptable. + with ParserElement.recursion_lock: + memo = ParserElement.recursion_memos + try: + # we are parsing at a specific recursion expansion - use it as-is + prev_loc, prev_result = memo[loc, self, do_actions] + if isinstance(prev_result, Exception): + raise prev_result + return prev_loc, prev_result.copy() + except KeyError: + act_key = (loc, self, True) + peek_key = (loc, self, False) + # we are searching for the best recursion expansion - keep on improving + # both `do_actions` cases must be tracked separately here! + prev_loc, prev_peek = memo[peek_key] = ( + loc - 1, + ParseException( + instring, loc, "Forward recursion without base case", self + ), + ) + if do_actions: + memo[act_key] = memo[peek_key] + while True: + try: + new_loc, new_peek = super().parseImpl(instring, loc, False) + except ParseException: + # we failed before getting any match - do not hide the error + if isinstance(prev_peek, Exception): + raise + new_loc, new_peek = prev_loc, prev_peek + # the match did not get better: we are done + if new_loc <= prev_loc: + if do_actions: + # replace the match for do_actions=False as well, + # in case the action did backtrack + prev_loc, prev_result = memo[peek_key] = memo[act_key] + del memo[peek_key], memo[act_key] + return prev_loc, copy.copy(prev_result) + del memo[peek_key] + return prev_loc, copy.copy(prev_peek) + # the match did get better: see if we can improve further + if do_actions: + try: + memo[act_key] = super().parseImpl(instring, loc, True) + except ParseException as e: + memo[peek_key] = memo[act_key] = (new_loc, e) + raise + prev_loc, prev_peek = memo[peek_key] = new_loc, new_peek + + def leave_whitespace(self, recursive: bool = True) -> ParserElement: + self.skipWhitespace = False + return self + + def ignore_whitespace(self, recursive: bool = True) -> ParserElement: + self.skipWhitespace = True + return self + + def streamline(self) -> ParserElement: + if not self.streamlined: + self.streamlined = True + if self.expr is not None: + self.expr.streamline() + return self + + def validate(self, validateTrace=None) -> None: + warnings.warn( + "ParserElement.validate() is deprecated, and should not be used to check for left recursion", + DeprecationWarning, + stacklevel=2, + ) + if validateTrace is None: + validateTrace = [] + + if self not in validateTrace: + tmp = validateTrace[:] + [self] + if self.expr is not None: + self.expr.validate(tmp) + self._checkRecursion([]) + + def _generateDefaultName(self) -> str: + # Avoid infinite recursion by setting a temporary _defaultName + save_default_name = self._defaultName + self._defaultName = ": ..." + + # Use the string representation of main expression. + try: + if self.expr is not None: + ret_string = str(self.expr)[:1000] + else: + ret_string = "None" + except Exception: + ret_string = "..." + + self._defaultName = save_default_name + return f"{type(self).__name__}: {ret_string}" + + def copy(self) -> ParserElement: + if self.expr is not None: + return super().copy() + else: + ret = Forward() + ret <<= self + return ret + + def _setResultsName(self, name, list_all_matches=False) -> ParserElement: + # fmt: off + if ( + __diag__.warn_name_set_on_empty_Forward + and Diagnostics.warn_name_set_on_empty_Forward not in self.suppress_warnings_ + and self.expr is None + ): + warning = ( + "warn_name_set_on_empty_Forward:" + f" setting results name {name!r} on {type(self).__name__} expression" + " that has no contained expression" + ) + warnings.warn(warning, stacklevel=3) + # fmt: on + + return super()._setResultsName(name, list_all_matches) + + # Compatibility synonyms + # fmt: off + leaveWhitespace = replaced_by_pep8("leaveWhitespace", leave_whitespace) + ignoreWhitespace = replaced_by_pep8("ignoreWhitespace", ignore_whitespace) + # fmt: on + + +class TokenConverter(ParseElementEnhance): + """ + Abstract subclass of :class:`ParseElementEnhance`, for converting parsed results. + """ + + def __init__(self, expr: Union[ParserElement, str], savelist=False) -> None: + super().__init__(expr) # , savelist) + self.saveAsList = False + + +class Combine(TokenConverter): + """Converter to concatenate all matching tokens to a single string. + By default, the matching patterns must also be contiguous in the + input string; this can be disabled by specifying + ``'adjacent=False'`` in the constructor. + + Example:: + + real = Word(nums) + '.' + Word(nums) + print(real.parse_string('3.1416')) # -> ['3', '.', '1416'] + # will also erroneously match the following + print(real.parse_string('3. 1416')) # -> ['3', '.', '1416'] + + real = Combine(Word(nums) + '.' + Word(nums)) + print(real.parse_string('3.1416')) # -> ['3.1416'] + # no match when there are internal spaces + print(real.parse_string('3. 1416')) # -> Exception: Expected W:(0123...) + """ + + def __init__( + self, + expr: ParserElement, + join_string: str = "", + adjacent: bool = True, + *, + joinString: typing.Optional[str] = None, + ) -> None: + super().__init__(expr) + joinString = joinString if joinString is not None else join_string + # suppress whitespace-stripping in contained parse expressions, but re-enable it on the Combine itself + if adjacent: + self.leave_whitespace() + self.adjacent = adjacent + self.skipWhitespace = True + self.joinString = joinString + self.callPreparse = True + + def ignore(self, other) -> ParserElement: + if self.adjacent: + ParserElement.ignore(self, other) + else: + super().ignore(other) + return self + + def postParse(self, instring, loc, tokenlist): + retToks = tokenlist.copy() + del retToks[:] + retToks += ParseResults( + ["".join(tokenlist._asStringList(self.joinString))], modal=self.modalResults + ) + + if self.resultsName and retToks.haskeys(): + return [retToks] + else: + return retToks + + +class Group(TokenConverter): + """Converter to return the matched tokens as a list - useful for + returning tokens of :class:`ZeroOrMore` and :class:`OneOrMore` expressions. + + The optional ``aslist`` argument when set to True will return the + parsed tokens as a Python list instead of a pyparsing ParseResults. + + Example:: + + ident = Word(alphas) + num = Word(nums) + term = ident | num + func = ident + Opt(DelimitedList(term)) + print(func.parse_string("fn a, b, 100")) + # -> ['fn', 'a', 'b', '100'] + + func = ident + Group(Opt(DelimitedList(term))) + print(func.parse_string("fn a, b, 100")) + # -> ['fn', ['a', 'b', '100']] + """ + + def __init__(self, expr: ParserElement, aslist: bool = False) -> None: + super().__init__(expr) + self.saveAsList = True + self._asPythonList = aslist + + def postParse(self, instring, loc, tokenlist): + if self._asPythonList: + return ParseResults.List( + tokenlist.asList() + if isinstance(tokenlist, ParseResults) + else list(tokenlist) + ) + + return [tokenlist] + + +class Dict(TokenConverter): + """Converter to return a repetitive expression as a list, but also + as a dictionary. Each element can also be referenced using the first + token in the expression as its key. Useful for tabular report + scraping when the first column can be used as a item key. + + The optional ``asdict`` argument when set to True will return the + parsed tokens as a Python dict instead of a pyparsing ParseResults. + + Example:: + + data_word = Word(alphas) + label = data_word + FollowedBy(':') + + text = "shape: SQUARE posn: upper left color: light blue texture: burlap" + attr_expr = (label + Suppress(':') + OneOrMore(data_word, stop_on=label).set_parse_action(' '.join)) + + # print attributes as plain groups + print(attr_expr[1, ...].parse_string(text).dump()) + + # instead of OneOrMore(expr), parse using Dict(Group(expr)[1, ...]) - Dict will auto-assign names + result = Dict(Group(attr_expr)[1, ...]).parse_string(text) + print(result.dump()) + + # access named fields as dict entries, or output as dict + print(result['shape']) + print(result.as_dict()) + + prints:: + + ['shape', 'SQUARE', 'posn', 'upper left', 'color', 'light blue', 'texture', 'burlap'] + [['shape', 'SQUARE'], ['posn', 'upper left'], ['color', 'light blue'], ['texture', 'burlap']] + - color: 'light blue' + - posn: 'upper left' + - shape: 'SQUARE' + - texture: 'burlap' + SQUARE + {'color': 'light blue', 'posn': 'upper left', 'texture': 'burlap', 'shape': 'SQUARE'} + + See more examples at :class:`ParseResults` of accessing fields by results name. + """ + + def __init__(self, expr: ParserElement, asdict: bool = False) -> None: + super().__init__(expr) + self.saveAsList = True + self._asPythonDict = asdict + + def postParse(self, instring, loc, tokenlist): + for i, tok in enumerate(tokenlist): + if len(tok) == 0: + continue + + ikey = tok[0] + if isinstance(ikey, int): + ikey = str(ikey).strip() + + if len(tok) == 1: + tokenlist[ikey] = _ParseResultsWithOffset("", i) + + elif len(tok) == 2 and not isinstance(tok[1], ParseResults): + tokenlist[ikey] = _ParseResultsWithOffset(tok[1], i) + + else: + try: + dictvalue = tok.copy() # ParseResults(i) + except Exception: + exc = TypeError( + "could not extract dict values from parsed results" + " - Dict expression must contain Grouped expressions" + ) + raise exc from None + + del dictvalue[0] + + if len(dictvalue) != 1 or ( + isinstance(dictvalue, ParseResults) and dictvalue.haskeys() + ): + tokenlist[ikey] = _ParseResultsWithOffset(dictvalue, i) + else: + tokenlist[ikey] = _ParseResultsWithOffset(dictvalue[0], i) + + if self._asPythonDict: + return [tokenlist.as_dict()] if self.resultsName else tokenlist.as_dict() + + return [tokenlist] if self.resultsName else tokenlist + + +class Suppress(TokenConverter): + """Converter for ignoring the results of a parsed expression. + + Example:: + + source = "a, b, c,d" + wd = Word(alphas) + wd_list1 = wd + (',' + wd)[...] + print(wd_list1.parse_string(source)) + + # often, delimiters that are useful during parsing are just in the + # way afterward - use Suppress to keep them out of the parsed output + wd_list2 = wd + (Suppress(',') + wd)[...] + print(wd_list2.parse_string(source)) + + # Skipped text (using '...') can be suppressed as well + source = "lead in START relevant text END trailing text" + start_marker = Keyword("START") + end_marker = Keyword("END") + find_body = Suppress(...) + start_marker + ... + end_marker + print(find_body.parse_string(source) + + prints:: + + ['a', ',', 'b', ',', 'c', ',', 'd'] + ['a', 'b', 'c', 'd'] + ['START', 'relevant text ', 'END'] + + (See also :class:`DelimitedList`.) + """ + + def __init__(self, expr: Union[ParserElement, str], savelist: bool = False) -> None: + if expr is ...: + expr = _PendingSkip(NoMatch()) + super().__init__(expr) + + def __add__(self, other) -> ParserElement: + if isinstance(self.expr, _PendingSkip): + return Suppress(SkipTo(other)) + other + + return super().__add__(other) + + def __sub__(self, other) -> ParserElement: + if isinstance(self.expr, _PendingSkip): + return Suppress(SkipTo(other)) - other + + return super().__sub__(other) + + def postParse(self, instring, loc, tokenlist): + return [] + + def suppress(self) -> ParserElement: + return self + + +def trace_parse_action(f: ParseAction) -> ParseAction: + """Decorator for debugging parse actions. + + When the parse action is called, this decorator will print + ``">> entering method-name(line:, , )"``. + When the parse action completes, the decorator will print + ``"<<"`` followed by the returned value, or any exception that the parse action raised. + + Example:: + + wd = Word(alphas) + + @trace_parse_action + def remove_duplicate_chars(tokens): + return ''.join(sorted(set(''.join(tokens)))) + + wds = wd[1, ...].set_parse_action(remove_duplicate_chars) + print(wds.parse_string("slkdjs sld sldd sdlf sdljf")) + + prints:: + + >>entering remove_duplicate_chars(line: 'slkdjs sld sldd sdlf sdljf', 0, (['slkdjs', 'sld', 'sldd', 'sdlf', 'sdljf'], {})) + < 3: + thisFunc = f"{type(paArgs[0]).__name__}.{thisFunc}" + sys.stderr.write(f">>entering {thisFunc}(line: {line(l, s)!r}, {l}, {t!r})\n") + try: + ret = f(*paArgs) + except Exception as exc: + sys.stderr.write( + f"< str: + r"""Helper to easily define string ranges for use in :class:`Word` + construction. Borrows syntax from regexp ``'[]'`` string range + definitions:: + + srange("[0-9]") -> "0123456789" + srange("[a-z]") -> "abcdefghijklmnopqrstuvwxyz" + srange("[a-z$_]") -> "abcdefghijklmnopqrstuvwxyz$_" + + The input string must be enclosed in []'s, and the returned string + is the expanded character set joined into a single string. The + values enclosed in the []'s may be: + + - a single character + - an escaped character with a leading backslash (such as ``\-`` + or ``\]``) + - an escaped hex character with a leading ``'\x'`` + (``\x21``, which is a ``'!'`` character) (``\0x##`` + is also supported for backwards compatibility) + - an escaped octal character with a leading ``'\0'`` + (``\041``, which is a ``'!'`` character) + - a range of any of the above, separated by a dash (``'a-z'``, + etc.) + - any combination of the above (``'aeiouy'``, + ``'a-zA-Z0-9_$'``, etc.) + """ + + def _expanded(p): + if isinstance(p, ParseResults): + yield from (chr(c) for c in range(ord(p[0]), ord(p[1]) + 1)) + else: + yield p + + try: + return "".join( + [c for part in _reBracketExpr.parse_string(s).body for c in _expanded(part)] + ) + except Exception as e: + return "" + + +def token_map(func, *args) -> ParseAction: + """Helper to define a parse action by mapping a function to all + elements of a :class:`ParseResults` list. If any additional args are passed, + they are forwarded to the given function as additional arguments + after the token, as in + ``hex_integer = Word(hexnums).set_parse_action(token_map(int, 16))``, + which will convert the parsed data to an integer using base 16. + + Example (compare the last to example in :class:`ParserElement.transform_string`:: + + hex_ints = Word(hexnums)[1, ...].set_parse_action(token_map(int, 16)) + hex_ints.run_tests(''' + 00 11 22 aa FF 0a 0d 1a + ''') + + upperword = Word(alphas).set_parse_action(token_map(str.upper)) + upperword[1, ...].run_tests(''' + my kingdom for a horse + ''') + + wd = Word(alphas).set_parse_action(token_map(str.title)) + wd[1, ...].set_parse_action(' '.join).run_tests(''' + now is the winter of our discontent made glorious summer by this sun of york + ''') + + prints:: + + 00 11 22 aa FF 0a 0d 1a + [0, 17, 34, 170, 255, 10, 13, 26] + + my kingdom for a horse + ['MY', 'KINGDOM', 'FOR', 'A', 'HORSE'] + + now is the winter of our discontent made glorious summer by this sun of york + ['Now Is The Winter Of Our Discontent Made Glorious Summer By This Sun Of York'] + """ + + def pa(s, l, t): + return [func(tokn, *args) for tokn in t] + + func_name = getattr(func, "__name__", getattr(func, "__class__").__name__) + pa.__name__ = func_name + + return pa + + +def autoname_elements() -> None: + """ + Utility to simplify mass-naming of parser elements, for + generating railroad diagram with named subdiagrams. + """ + + # guard against _getframe not being implemented in the current Python + getframe_fn = getattr(sys, "_getframe", lambda _: None) + calling_frame = getframe_fn(1) + if calling_frame is None: + return + + # find all locals in the calling frame that are ParserElements + calling_frame = typing.cast(types.FrameType, calling_frame) + for name, var in calling_frame.f_locals.items(): + # if no custom name defined, set the name to the var name + if isinstance(var, ParserElement) and not var.customName: + var.set_name(name) + + +dbl_quoted_string = Combine( + Regex(r'"(?:[^"\n\r\\]|(?:"")|(?:\\(?:[^x]|x[0-9a-fA-F]+)))*') + '"' +).set_name("string enclosed in double quotes") + +sgl_quoted_string = Combine( + Regex(r"'(?:[^'\n\r\\]|(?:'')|(?:\\(?:[^x]|x[0-9a-fA-F]+)))*") + "'" +).set_name("string enclosed in single quotes") + +quoted_string = Combine( + (Regex(r'"(?:[^"\n\r\\]|(?:"")|(?:\\(?:[^x]|x[0-9a-fA-F]+)))*') + '"').set_name( + "double quoted string" + ) + | (Regex(r"'(?:[^'\n\r\\]|(?:'')|(?:\\(?:[^x]|x[0-9a-fA-F]+)))*") + "'").set_name( + "single quoted string" + ) +).set_name("quoted string using single or double quotes") + +python_quoted_string = Combine( + (Regex(r'"""(?:[^"\\]|""(?!")|"(?!"")|\\.)*', flags=re.MULTILINE) + '"""').set_name( + "multiline double quoted string" + ) + ^ ( + Regex(r"'''(?:[^'\\]|''(?!')|'(?!'')|\\.)*", flags=re.MULTILINE) + "'''" + ).set_name("multiline single quoted string") + ^ (Regex(r'"(?:[^"\n\r\\]|(?:\\")|(?:\\(?:[^x]|x[0-9a-fA-F]+)))*') + '"').set_name( + "double quoted string" + ) + ^ (Regex(r"'(?:[^'\n\r\\]|(?:\\')|(?:\\(?:[^x]|x[0-9a-fA-F]+)))*") + "'").set_name( + "single quoted string" + ) +).set_name("Python quoted string") + +unicode_string = Combine("u" + quoted_string.copy()).set_name("unicode string literal") + + +alphas8bit = srange(r"[\0xc0-\0xd6\0xd8-\0xf6\0xf8-\0xff]") +punc8bit = srange(r"[\0xa1-\0xbf\0xd7\0xf7]") + +# build list of built-in expressions, for future reference if a global default value +# gets updated +_builtin_exprs: list[ParserElement] = [ + v for v in vars().values() if isinstance(v, ParserElement) +] + +# Compatibility synonyms +# fmt: off +sglQuotedString = sgl_quoted_string +dblQuotedString = dbl_quoted_string +quotedString = quoted_string +unicodeString = unicode_string +lineStart = line_start +lineEnd = line_end +stringStart = string_start +stringEnd = string_end +nullDebugAction = replaced_by_pep8("nullDebugAction", null_debug_action) +traceParseAction = replaced_by_pep8("traceParseAction", trace_parse_action) +conditionAsParseAction = replaced_by_pep8("conditionAsParseAction", condition_as_parse_action) +tokenMap = replaced_by_pep8("tokenMap", token_map) +# fmt: on diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/diagram/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/diagram/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..526cf3862a471a2073dbc559cca533404a10671e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/diagram/__init__.py @@ -0,0 +1,754 @@ +# mypy: ignore-errors +from __future__ import annotations + +import itertools +import railroad +import pyparsing +import dataclasses +import typing +from typing import ( + Generic, + TypeVar, + Callable, + Iterable, +) +from jinja2 import Template +from io import StringIO +import inspect +import re + + +jinja2_template_source = """\ +{% if not embed %} + + + +{% endif %} + {% if not head %} + + {% else %} + {{ head | safe }} + {% endif %} +{% if not embed %} + + +{% endif %} + +{{ body | safe }} +{% for diagram in diagrams %} +
+

{{ diagram.title }}

+
{{ diagram.text }}
+
+ {{ diagram.svg }} +
+
+{% endfor %} +{% if not embed %} + + +{% endif %} +""" + +template = Template(jinja2_template_source) + + +_bookmark_lookup = {} +_bookmark_ids = itertools.count(start=1) + +def _make_bookmark(s: str) -> str: + """ + Converts a string into a valid HTML bookmark (ID or anchor name). + """ + if s in _bookmark_lookup: + return _bookmark_lookup[s] + + # Replace invalid characters with hyphens and ensure only valid characters + bookmark = re.sub(r'[^a-zA-Z0-9-]+', '-', s) + + # Ensure it starts with a letter by adding 'z' if necessary + if not bookmark[:1].isalpha(): + bookmark = f"z{bookmark}" + + # Convert to lowercase and strip hyphens + bookmark = bookmark.lower().strip('-') + + _bookmark_lookup[s] = bookmark = f"{bookmark}-{next(_bookmark_ids):04d}" + + return bookmark + + +def _collapse_verbose_regex(regex_str: str) -> str: + if "\n" not in regex_str: + return regex_str + collapsed = pyparsing.Regex(r"#.*$").suppress().transform_string(regex_str) + collapsed = re.sub(r"\s*\n\s*", "", collapsed) + return collapsed + + +@dataclasses.dataclass +class NamedDiagram: + """ + A simple structure for associating a name with a railroad diagram + """ + + name: str + index: int + diagram: railroad.DiagramItem = None + + @property + def bookmark(self): + bookmark = _make_bookmark(self.name) + return bookmark + + +T = TypeVar("T") + + +class EachItem(railroad.Group): + """ + Custom railroad item to compose a: + - Group containing a + - OneOrMore containing a + - Choice of the elements in the Each + with the group label indicating that all must be matched + """ + + all_label = "[ALL]" + + def __init__(self, *items) -> None: + choice_item = railroad.Choice(len(items) - 1, *items) + one_or_more_item = railroad.OneOrMore(item=choice_item) + super().__init__(one_or_more_item, label=self.all_label) + + +class AnnotatedItem(railroad.Group): + """ + Simple subclass of Group that creates an annotation label + """ + + def __init__(self, label: str, item) -> None: + super().__init__(item=item, label=f"[{label}]" if label else "") + + +class EditablePartial(Generic[T]): + """ + Acts like a functools.partial, but can be edited. In other words, it represents a type that hasn't yet been + constructed. + """ + + # We need this here because the railroad constructors actually transform the data, so can't be called until the + # entire tree is assembled + + def __init__(self, func: Callable[..., T], args: list, kwargs: dict) -> None: + self.func = func + self.args = args + self.kwargs = kwargs + + @classmethod + def from_call(cls, func: Callable[..., T], *args, **kwargs) -> EditablePartial[T]: + """ + If you call this function in the same way that you would call the constructor, it will store the arguments + as you expect. For example EditablePartial.from_call(Fraction, 1, 3)() == Fraction(1, 3) + """ + return EditablePartial(func=func, args=list(args), kwargs=kwargs) + + @property + def name(self): + return self.kwargs["name"] + + def __call__(self) -> T: + """ + Evaluate the partial and return the result + """ + args = self.args.copy() + kwargs = self.kwargs.copy() + + # This is a helpful hack to allow you to specify varargs parameters (e.g. *args) as keyword args (e.g. + # args=['list', 'of', 'things']) + arg_spec = inspect.getfullargspec(self.func) + if arg_spec.varargs in self.kwargs: + args += kwargs.pop(arg_spec.varargs) + + return self.func(*args, **kwargs) + + +def railroad_to_html(diagrams: list[NamedDiagram], embed=False, **kwargs) -> str: + """ + Given a list of NamedDiagram, produce a single HTML string that visualises those diagrams + :params kwargs: kwargs to be passed in to the template + """ + data = [] + for diagram in diagrams: + if diagram.diagram is None: + continue + io = StringIO() + try: + css = kwargs.get("css") + diagram.diagram.writeStandalone(io.write, css=css) + except AttributeError: + diagram.diagram.writeSvg(io.write) + title = diagram.name + if diagram.index == 0: + title += " (root)" + data.append( + { + "title": title, "text": "", "svg": io.getvalue(), "bookmark": diagram.bookmark + } + ) + + return template.render(diagrams=data, embed=embed, **kwargs) + + +def resolve_partial(partial: EditablePartial[T]) -> T: + """ + Recursively resolves a collection of Partials into whatever type they are + """ + if isinstance(partial, EditablePartial): + partial.args = resolve_partial(partial.args) + partial.kwargs = resolve_partial(partial.kwargs) + return partial() + elif isinstance(partial, list): + return [resolve_partial(x) for x in partial] + elif isinstance(partial, dict): + return {key: resolve_partial(x) for key, x in partial.items()} + else: + return partial + + +def to_railroad( + element: pyparsing.ParserElement, + diagram_kwargs: typing.Optional[dict] = None, + vertical: int = 3, + show_results_names: bool = False, + show_groups: bool = False, + show_hidden: bool = False, +) -> list[NamedDiagram]: + """ + Convert a pyparsing element tree into a list of diagrams. This is the recommended entrypoint to diagram + creation if you want to access the Railroad tree before it is converted to HTML + :param element: base element of the parser being diagrammed + :param diagram_kwargs: kwargs to pass to the Diagram() constructor + :param vertical: (optional) - int - limit at which number of alternatives should be + shown vertically instead of horizontally + :param show_results_names - bool to indicate whether results name annotations should be + included in the diagram + :param show_groups - bool to indicate whether groups should be highlighted with an unlabeled + surrounding box + :param show_hidden - bool to indicate whether internal elements that are typically hidden + should be shown + """ + # Convert the whole tree underneath the root + lookup = ConverterState(diagram_kwargs=diagram_kwargs or {}) + _to_diagram_element( + element, + lookup=lookup, + parent=None, + vertical=vertical, + show_results_names=show_results_names, + show_groups=show_groups, + show_hidden=show_hidden, + ) + + root_id = id(element) + # Convert the root if it hasn't been already + if root_id in lookup: + if not element.customName: + lookup[root_id].name = "" + lookup[root_id].mark_for_extraction(root_id, lookup, force=True) + + # Now that we're finished, we can convert from intermediate structures into Railroad elements + diags = list(lookup.diagrams.values()) + if len(diags) > 1: + # collapse out duplicate diags with the same name + seen = set() + deduped_diags = [] + for d in diags: + # don't extract SkipTo elements, they are uninformative as subdiagrams + if d.name == "...": + continue + if d.name is not None and d.name not in seen: + seen.add(d.name) + deduped_diags.append(d) + resolved = [resolve_partial(partial) for partial in deduped_diags] + else: + # special case - if just one diagram, always display it, even if + # it has no name + resolved = [resolve_partial(partial) for partial in diags] + return sorted(resolved, key=lambda diag: diag.index) + + +def _should_vertical( + specification: int, exprs: Iterable[pyparsing.ParserElement] +) -> bool: + """ + Returns true if we should return a vertical list of elements + """ + if specification is None: + return False + else: + return len(_visible_exprs(exprs)) >= specification + + +@dataclasses.dataclass +class ElementState: + """ + State recorded for an individual pyparsing Element + """ + + #: The pyparsing element that this represents + element: pyparsing.ParserElement + #: The output Railroad element in an unconverted state + converted: EditablePartial + #: The parent Railroad element, which we store so that we can extract this if it's duplicated + parent: EditablePartial + #: The order in which we found this element, used for sorting diagrams if this is extracted into a diagram + number: int + #: The name of the element + name: str = None + #: The index of this inside its parent + parent_index: typing.Optional[int] = None + #: If true, we should extract this out into a subdiagram + extract: bool = False + #: If true, all of this element's children have been filled out + complete: bool = False + + def mark_for_extraction( + self, el_id: int, state: ConverterState, name: str = None, force: bool = False + ): + """ + Called when this instance has been seen twice, and thus should eventually be extracted into a sub-diagram + :param el_id: id of the element + :param state: element/diagram state tracker + :param name: name to use for this element's text + :param force: If true, force extraction now, regardless of the state of this. Only useful for extracting the + root element when we know we're finished + """ + self.extract = True + + # Set the name + if not self.name: + if name: + # Allow forcing a custom name + self.name = name + elif self.element.customName: + self.name = self.element.customName + else: + self.name = "" + + # Just because this is marked for extraction doesn't mean we can do it yet. We may have to wait for children + # to be added + # Also, if this is just a string literal etc, don't bother extracting it + if force or (self.complete and _worth_extracting(self.element)): + state.extract_into_diagram(el_id) + + +class ConverterState: + """ + Stores some state that persists between recursions into the element tree + """ + + def __init__(self, diagram_kwargs: typing.Optional[dict] = None) -> None: + #: A dictionary mapping ParserElements to state relating to them + self._element_diagram_states: dict[int, ElementState] = {} + #: A dictionary mapping ParserElement IDs to subdiagrams generated from them + self.diagrams: dict[int, EditablePartial[NamedDiagram]] = {} + #: The index of the next unnamed element + self.unnamed_index: int = 1 + #: The index of the next element. This is used for sorting + self.index: int = 0 + #: Shared kwargs that are used to customize the construction of diagrams + self.diagram_kwargs: dict = diagram_kwargs or {} + self.extracted_diagram_names: set[str] = set() + + def __setitem__(self, key: int, value: ElementState): + self._element_diagram_states[key] = value + + def __getitem__(self, key: int) -> ElementState: + return self._element_diagram_states[key] + + def __delitem__(self, key: int): + del self._element_diagram_states[key] + + def __contains__(self, key: int): + return key in self._element_diagram_states + + def get(self, key, default=None): + try: + return self[key] + except KeyError: + return default + + def generate_unnamed(self) -> int: + """ + Generate a number used in the name of an otherwise unnamed diagram + """ + self.unnamed_index += 1 + return self.unnamed_index + + def generate_index(self) -> int: + """ + Generate a number used to index a diagram + """ + self.index += 1 + return self.index + + def extract_into_diagram(self, el_id: int): + """ + Used when we encounter the same token twice in the same tree. When this + happens, we replace all instances of that token with a terminal, and + create a new subdiagram for the token + """ + position = self[el_id] + + # Replace the original definition of this element with a regular block + if position.parent: + href = f"#{_make_bookmark(position.name)}" + ret = EditablePartial.from_call(railroad.NonTerminal, text=position.name, href=href) + if "item" in position.parent.kwargs: + position.parent.kwargs["item"] = ret + elif "items" in position.parent.kwargs: + position.parent.kwargs["items"][position.parent_index] = ret + + # If the element we're extracting is a group, skip to its content but keep the title + if position.converted.func == railroad.Group: + content = position.converted.kwargs["item"] + else: + content = position.converted + + self.diagrams[el_id] = EditablePartial.from_call( + NamedDiagram, + name=position.name, + diagram=EditablePartial.from_call( + railroad.Diagram, content, **self.diagram_kwargs + ), + index=position.number, + ) + + del self[el_id] + + +def _worth_extracting(element: pyparsing.ParserElement) -> bool: + """ + Returns true if this element is worth having its own sub-diagram. Simply, if any of its children + themselves have children, then its complex enough to extract + """ + children = element.recurse() + return any(child.recurse() for child in children) + + +def _apply_diagram_item_enhancements(fn): + """ + decorator to ensure enhancements to a diagram item (such as results name annotations) + get applied on return from _to_diagram_element (we do this since there are several + returns in _to_diagram_element) + """ + + def _inner( + element: pyparsing.ParserElement, + parent: typing.Optional[EditablePartial], + lookup: ConverterState = None, + vertical: int = None, + index: int = 0, + name_hint: str = None, + show_results_names: bool = False, + show_groups: bool = False, + show_hidden: bool = False, + ) -> typing.Optional[EditablePartial]: + ret = fn( + element, + parent, + lookup, + vertical, + index, + name_hint, + show_results_names, + show_groups, + show_hidden, + ) + + # apply annotation for results name, if present + if show_results_names and ret is not None: + element_results_name = element.resultsName + if element_results_name: + # add "*" to indicate if this is a "list all results" name + modal_tag = "" if element.modalResults else "*" + ret = EditablePartial.from_call( + railroad.Group, + item=ret, + label=f"{repr(element_results_name)}{modal_tag}", + ) + + return ret + + return _inner + + +def _visible_exprs(exprs: Iterable[pyparsing.ParserElement]): + non_diagramming_exprs = ( + pyparsing.ParseElementEnhance, + pyparsing.PositionToken, + pyparsing.And._ErrorStop, + ) + return [ + e + for e in exprs + if not isinstance(e, non_diagramming_exprs) + ] + + +@_apply_diagram_item_enhancements +def _to_diagram_element( + element: pyparsing.ParserElement, + parent: typing.Optional[EditablePartial], + lookup: ConverterState = None, + vertical: int = None, + index: int = 0, + name_hint: str = None, + show_results_names: bool = False, + show_groups: bool = False, + show_hidden: bool = False, +) -> typing.Optional[EditablePartial]: + """ + Recursively converts a PyParsing Element to a railroad Element + :param lookup: The shared converter state that keeps track of useful things + :param index: The index of this element within the parent + :param parent: The parent of this element in the output tree + :param vertical: Controls at what point we make a list of elements vertical. If this is an integer (the default), + it sets the threshold of the number of items before we go vertical. If True, always go vertical, if False, never + do so + :param name_hint: If provided, this will override the generated name + :param show_results_names: bool flag indicating whether to add annotations for results names + :param show_groups: bool flag indicating whether to show groups using bounding box + :param show_hidden: bool flag indicating whether to show elements that are typically hidden + :returns: The converted version of the input element, but as a Partial that hasn't yet been constructed + """ + exprs = element.recurse() + name = name_hint or element.customName or type(element).__name__ + + # Python's id() is used to provide a unique identifier for elements + el_id = id(element) + + element_results_name = element.resultsName + + # Here we basically bypass processing certain wrapper elements if they contribute nothing to the diagram + if not element.customName: + if isinstance( + element, + ( + # pyparsing.TokenConverter, + pyparsing.Forward, + pyparsing.Located, + ), + ): + # However, if this element has a useful custom name, and its child does not, we can pass it on to the child + if exprs: + if not exprs[0].customName: + propagated_name = name + else: + propagated_name = None + + return _to_diagram_element( + element.expr, + parent=parent, + lookup=lookup, + vertical=vertical, + index=index, + name_hint=propagated_name, + show_results_names=show_results_names, + show_groups=show_groups, + show_hidden=show_hidden, + ) + + # If the element isn't worth extracting, we always treat it as the first time we say it + if _worth_extracting(element): + looked_up = lookup.get(el_id) + if looked_up and looked_up.name is not None: + # If we've seen this element exactly once before, we are only just now finding out that it's a duplicate, + # so we have to extract it into a new diagram. + looked_up.mark_for_extraction(el_id, lookup, name=name_hint) + href = f"#{_make_bookmark(looked_up.name)}" + ret = EditablePartial.from_call(railroad.NonTerminal, text=looked_up.name, href=href) + return ret + + elif el_id in lookup.diagrams: + # If we have seen the element at least twice before, and have already extracted it into a subdiagram, we + # just put in a marker element that refers to the sub-diagram + text = lookup.diagrams[el_id].kwargs["name"] + ret = EditablePartial.from_call( + railroad.NonTerminal, text=text, href=f"#{_make_bookmark(text)}" + ) + return ret + + # Recursively convert child elements + # Here we find the most relevant Railroad element for matching pyparsing Element + # We use ``items=[]`` here to hold the place for where the child elements will go once created + + # see if this element is normally hidden, and whether hidden elements are desired + # if not, just return None + if not element.show_in_diagram and not show_hidden: + return None + + if isinstance(element, pyparsing.And): + # detect And's created with ``expr*N`` notation - for these use a OneOrMore with a repeat + # (all will have the same name, and resultsName) + if not exprs: + return None + if len(set((e.name, e.resultsName) for e in exprs)) == 1 and len(exprs) > 2: + ret = EditablePartial.from_call( + railroad.OneOrMore, item="", repeat=str(len(exprs)) + ) + elif _should_vertical(vertical, exprs): + ret = EditablePartial.from_call(railroad.Stack, items=[]) + else: + ret = EditablePartial.from_call(railroad.Sequence, items=[]) + elif isinstance(element, (pyparsing.Or, pyparsing.MatchFirst)): + if not exprs: + return None + if _should_vertical(vertical, exprs): + ret = EditablePartial.from_call(railroad.Choice, 0, items=[]) + else: + ret = EditablePartial.from_call(railroad.HorizontalChoice, items=[]) + elif isinstance(element, pyparsing.Each): + if not exprs: + return None + ret = EditablePartial.from_call(EachItem, items=[]) + elif isinstance(element, pyparsing.NotAny): + ret = EditablePartial.from_call(AnnotatedItem, label="NOT", item="") + elif isinstance(element, pyparsing.FollowedBy): + ret = EditablePartial.from_call(AnnotatedItem, label="LOOKAHEAD", item="") + elif isinstance(element, pyparsing.PrecededBy): + ret = EditablePartial.from_call(AnnotatedItem, label="LOOKBEHIND", item="") + elif isinstance(element, pyparsing.Group): + if show_groups: + ret = EditablePartial.from_call(AnnotatedItem, label="", item="") + else: + ret = EditablePartial.from_call( + railroad.Group, item=None, label=element_results_name + ) + elif isinstance(element, pyparsing.TokenConverter): + label = type(element).__name__.lower() + if label == "tokenconverter": + ret = EditablePartial.from_call(railroad.Sequence, items=[]) + else: + ret = EditablePartial.from_call(AnnotatedItem, label=label, item="") + elif isinstance(element, pyparsing.Opt): + ret = EditablePartial.from_call(railroad.Optional, item="") + elif isinstance(element, pyparsing.OneOrMore): + if element.not_ender is not None: + args = [ + parent, + lookup, + vertical, + index, + name_hint, + show_results_names, + show_groups, + show_hidden, + ] + return _to_diagram_element( + (~element.not_ender.expr + element.expr)[1, ...].set_name(element.name), + *args, + ) + ret = EditablePartial.from_call(railroad.OneOrMore, item=None) + elif isinstance(element, pyparsing.ZeroOrMore): + if element.not_ender is not None: + args = [ + parent, + lookup, + vertical, + index, + name_hint, + show_results_names, + show_groups, + show_hidden, + ] + return _to_diagram_element( + (~element.not_ender.expr + element.expr)[...].set_name(element.name), + *args, + ) + ret = EditablePartial.from_call(railroad.ZeroOrMore, item="") + elif isinstance(element, pyparsing.Empty) and not element.customName: + # Skip unnamed "Empty" elements + ret = None + elif isinstance(element, pyparsing.ParseElementEnhance): + ret = EditablePartial.from_call(railroad.Sequence, items=[]) + elif len(exprs) > 0 and not element_results_name: + ret = EditablePartial.from_call(railroad.Group, item="", label=name) + elif isinstance(element, pyparsing.Regex): + collapsed_patt = _collapse_verbose_regex(element.pattern) + ret = EditablePartial.from_call(railroad.Terminal, collapsed_patt) + elif len(exprs) > 0: + ret = EditablePartial.from_call(railroad.Sequence, items=[]) + else: + terminal = EditablePartial.from_call(railroad.Terminal, element.defaultName) + ret = terminal + + if ret is None: + return + + # Indicate this element's position in the tree so we can extract it if necessary + lookup[el_id] = ElementState( + element=element, + converted=ret, + parent=parent, + parent_index=index, + number=lookup.generate_index(), + ) + if element.customName: + lookup[el_id].mark_for_extraction(el_id, lookup, element.customName) + + i = 0 + for expr in exprs: + # Add a placeholder index in case we have to extract the child before we even add it to the parent + if "items" in ret.kwargs: + ret.kwargs["items"].insert(i, None) + + item = _to_diagram_element( + expr, + parent=ret, + lookup=lookup, + vertical=vertical, + index=i, + show_results_names=show_results_names, + show_groups=show_groups, + show_hidden=show_hidden, + ) + + # Some elements don't need to be shown in the diagram + if item is not None: + if "item" in ret.kwargs: + ret.kwargs["item"] = item + elif "items" in ret.kwargs: + # If we've already extracted the child, don't touch this index, since it's occupied by a nonterminal + ret.kwargs["items"][i] = item + i += 1 + elif "items" in ret.kwargs: + # If we're supposed to skip this element, remove it from the parent + del ret.kwargs["items"][i] + + # If all this items children are none, skip this item + if ret and ( + ("items" in ret.kwargs and len(ret.kwargs["items"]) == 0) + or ("item" in ret.kwargs and ret.kwargs["item"] is None) + ): + ret = EditablePartial.from_call(railroad.Terminal, name) + + # Mark this element as "complete", ie it has all of its children + if el_id in lookup: + lookup[el_id].complete = True + + if el_id in lookup and lookup[el_id].extract and lookup[el_id].complete: + lookup.extract_into_diagram(el_id) + if ret is not None: + text = lookup.diagrams[el_id].kwargs["name"] + href = f"#{_make_bookmark(text)}" + ret = EditablePartial.from_call( + railroad.NonTerminal, text=text, href=href + ) + + return ret diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/exceptions.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..fe07a85585616964ac0aebf473ec24fa557f5dd4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/exceptions.py @@ -0,0 +1,313 @@ +# exceptions.py +from __future__ import annotations + +import copy +import re +import sys +import typing +from functools import cached_property + +from .unicode import pyparsing_unicode as ppu +from .util import ( + _collapse_string_to_ranges, + col, + line, + lineno, + replaced_by_pep8, +) + + +class _ExceptionWordUnicodeSet( + ppu.Latin1, ppu.LatinA, ppu.LatinB, ppu.Greek, ppu.Cyrillic +): + pass + + +_extract_alphanums = _collapse_string_to_ranges(_ExceptionWordUnicodeSet.alphanums) +_exception_word_extractor = re.compile("([" + _extract_alphanums + "]{1,16})|.") + + +class ParseBaseException(Exception): + """base exception class for all parsing runtime exceptions""" + + loc: int + msg: str + pstr: str + parser_element: typing.Any # "ParserElement" + args: tuple[str, int, typing.Optional[str]] + + __slots__ = ( + "loc", + "msg", + "pstr", + "parser_element", + "args", + ) + + # Performance tuning: we construct a *lot* of these, so keep this + # constructor as small and fast as possible + def __init__( + self, + pstr: str, + loc: int = 0, + msg: typing.Optional[str] = None, + elem=None, + ) -> None: + if msg is None: + msg, pstr = pstr, "" + + self.loc = loc + self.msg = msg + self.pstr = pstr + self.parser_element = elem + self.args = (pstr, loc, msg) + + @staticmethod + def explain_exception(exc: Exception, depth: int = 16) -> str: + """ + Method to take an exception and translate the Python internal traceback into a list + of the pyparsing expressions that caused the exception to be raised. + + Parameters: + + - exc - exception raised during parsing (need not be a ParseException, in support + of Python exceptions that might be raised in a parse action) + - depth (default=16) - number of levels back in the stack trace to list expression + and function names; if None, the full stack trace names will be listed; if 0, only + the failing input line, marker, and exception string will be shown + + Returns a multi-line string listing the ParserElements and/or function names in the + exception's stack trace. + """ + import inspect + from .core import ParserElement + + if depth is None: + depth = sys.getrecursionlimit() + ret: list[str] = [] + if isinstance(exc, ParseBaseException): + ret.append(exc.line) + ret.append(f"{'^':>{exc.column}}") + ret.append(f"{type(exc).__name__}: {exc}") + + if depth <= 0 or exc.__traceback__ is None: + return "\n".join(ret) + + callers = inspect.getinnerframes(exc.__traceback__, context=depth) + seen: set[int] = set() + for ff in callers[-depth:]: + frm = ff[0] + + f_self = frm.f_locals.get("self", None) + if isinstance(f_self, ParserElement): + if not frm.f_code.co_name.startswith(("parseImpl", "_parseNoCache")): + continue + if id(f_self) in seen: + continue + seen.add(id(f_self)) + + self_type = type(f_self) + ret.append(f"{self_type.__module__}.{self_type.__name__} - {f_self}") + + elif f_self is not None: + self_type = type(f_self) + ret.append(f"{self_type.__module__}.{self_type.__name__}") + + else: + code = frm.f_code + if code.co_name in ("wrapper", ""): + continue + + ret.append(code.co_name) + + depth -= 1 + if not depth: + break + + return "\n".join(ret) + + @classmethod + def _from_exception(cls, pe) -> ParseBaseException: + """ + internal factory method to simplify creating one type of ParseException + from another - avoids having __init__ signature conflicts among subclasses + """ + return cls(pe.pstr, pe.loc, pe.msg, pe.parser_element) + + @cached_property + def line(self) -> str: + """ + Return the line of text where the exception occurred. + """ + return line(self.loc, self.pstr) + + @cached_property + def lineno(self) -> int: + """ + Return the 1-based line number of text where the exception occurred. + """ + return lineno(self.loc, self.pstr) + + @cached_property + def col(self) -> int: + """ + Return the 1-based column on the line of text where the exception occurred. + """ + return col(self.loc, self.pstr) + + @cached_property + def column(self) -> int: + """ + Return the 1-based column on the line of text where the exception occurred. + """ + return col(self.loc, self.pstr) + + @cached_property + def found(self) -> str: + if not self.pstr: + return "" + + if self.loc >= len(self.pstr): + return "end of text" + + # pull out next word at error location + found_match = _exception_word_extractor.match(self.pstr, self.loc) + if found_match is not None: + found_text = found_match.group(0) + else: + found_text = self.pstr[self.loc : self.loc + 1] + + return repr(found_text).replace(r"\\", "\\") + + # pre-PEP8 compatibility + @property + def parserElement(self): + return self.parser_element + + @parserElement.setter + def parserElement(self, elem): + self.parser_element = elem + + def copy(self): + return copy.copy(self) + + def formatted_message(self) -> str: + found_phrase = f", found {self.found}" if self.found else "" + return f"{self.msg}{found_phrase} (at char {self.loc}), (line:{self.lineno}, col:{self.column})" + + def __str__(self) -> str: + return self.formatted_message() + + def __repr__(self): + return str(self) + + def mark_input_line( + self, marker_string: typing.Optional[str] = None, *, markerString: str = ">!<" + ) -> str: + """ + Extracts the exception line from the input string, and marks + the location of the exception with a special symbol. + """ + markerString = marker_string if marker_string is not None else markerString + line_str = self.line + line_column = self.column - 1 + if markerString: + line_str = f"{line_str[:line_column]}{markerString}{line_str[line_column:]}" + return line_str.strip() + + def explain(self, depth: int = 16) -> str: + """ + Method to translate the Python internal traceback into a list + of the pyparsing expressions that caused the exception to be raised. + + Parameters: + + - depth (default=16) - number of levels back in the stack trace to list expression + and function names; if None, the full stack trace names will be listed; if 0, only + the failing input line, marker, and exception string will be shown + + Returns a multi-line string listing the ParserElements and/or function names in the + exception's stack trace. + + Example:: + + # an expression to parse 3 integers + expr = pp.Word(pp.nums) * 3 + try: + # a failing parse - the third integer is prefixed with "A" + expr.parse_string("123 456 A789") + except pp.ParseException as pe: + print(pe.explain(depth=0)) + + prints:: + + 123 456 A789 + ^ + ParseException: Expected W:(0-9), found 'A' (at char 8), (line:1, col:9) + + Note: the diagnostic output will include string representations of the expressions + that failed to parse. These representations will be more helpful if you use `set_name` to + give identifiable names to your expressions. Otherwise they will use the default string + forms, which may be cryptic to read. + + Note: pyparsing's default truncation of exception tracebacks may also truncate the + stack of expressions that are displayed in the ``explain`` output. To get the full listing + of parser expressions, you may have to set ``ParserElement.verbose_stacktrace = True`` + """ + return self.explain_exception(self, depth) + + # Compatibility synonyms + # fmt: off + markInputline = replaced_by_pep8("markInputline", mark_input_line) + # fmt: on + + +class ParseException(ParseBaseException): + """ + Exception thrown when a parse expression doesn't match the input string + + Example:: + + integer = Word(nums).set_name("integer") + try: + integer.parse_string("ABC") + except ParseException as pe: + print(pe, f"column: {pe.column}") + + prints:: + + Expected integer, found 'ABC' (at char 0), (line:1, col:1) column: 1 + + """ + + +class ParseFatalException(ParseBaseException): + """ + User-throwable exception thrown when inconsistent parse content + is found; stops all parsing immediately + """ + + +class ParseSyntaxException(ParseFatalException): + """ + Just like :class:`ParseFatalException`, but thrown internally + when an :class:`ErrorStop` ('-' operator) indicates + that parsing is to stop immediately because an unbacktrackable + syntax error has been found. + """ + + +class RecursiveGrammarException(Exception): + """ + Exception thrown by :class:`ParserElement.validate` if the + grammar could be left-recursive; parser may need to enable + left recursion using :class:`ParserElement.enable_left_recursion` + + Deprecated: only used by deprecated method ParserElement.validate. + """ + + def __init__(self, parseElementList) -> None: + self.parseElementTrace = parseElementList + + def __str__(self) -> str: + return f"RecursiveGrammarException: {self.parseElementTrace}" diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/helpers.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..2badd68d64f6c2269889ea035974e15478eb9410 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/helpers.py @@ -0,0 +1,1131 @@ +# helpers.py +import html.entities +import operator +import re +import sys +import typing + +from . import __diag__ +from .core import * +from .util import ( + _bslash, + _flatten, + _escape_regex_range_chars, + make_compressed_re, + replaced_by_pep8, +) + + +# +# global helpers +# +def counted_array( + expr: ParserElement, + int_expr: typing.Optional[ParserElement] = None, + *, + intExpr: typing.Optional[ParserElement] = None, +) -> ParserElement: + """Helper to define a counted list of expressions. + + This helper defines a pattern of the form:: + + integer expr expr expr... + + where the leading integer tells how many expr expressions follow. + The matched tokens returns the array of expr tokens as a list - the + leading count token is suppressed. + + If ``int_expr`` is specified, it should be a pyparsing expression + that produces an integer value. + + Example:: + + counted_array(Word(alphas)).parse_string('2 ab cd ef') # -> ['ab', 'cd'] + + # in this parser, the leading integer value is given in binary, + # '10' indicating that 2 values are in the array + binary_constant = Word('01').set_parse_action(lambda t: int(t[0], 2)) + counted_array(Word(alphas), int_expr=binary_constant).parse_string('10 ab cd ef') # -> ['ab', 'cd'] + + # if other fields must be parsed after the count but before the + # list items, give the fields results names and they will + # be preserved in the returned ParseResults: + count_with_metadata = integer + Word(alphas)("type") + typed_array = counted_array(Word(alphanums), int_expr=count_with_metadata)("items") + result = typed_array.parse_string("3 bool True True False") + print(result.dump()) + + # prints + # ['True', 'True', 'False'] + # - items: ['True', 'True', 'False'] + # - type: 'bool' + """ + intExpr = intExpr or int_expr + array_expr = Forward() + + def count_field_parse_action(s, l, t): + nonlocal array_expr + n = t[0] + array_expr <<= (expr * n) if n else Empty() + # clear list contents, but keep any named results + del t[:] + + if intExpr is None: + intExpr = Word(nums).set_parse_action(lambda t: int(t[0])) + else: + intExpr = intExpr.copy() + intExpr.set_name("arrayLen") + intExpr.add_parse_action(count_field_parse_action, call_during_try=True) + return (intExpr + array_expr).set_name(f"(len) {expr}...") + + +def match_previous_literal(expr: ParserElement) -> ParserElement: + """Helper to define an expression that is indirectly defined from + the tokens matched in a previous expression, that is, it looks for + a 'repeat' of a previous expression. For example:: + + first = Word(nums) + second = match_previous_literal(first) + match_expr = first + ":" + second + + will match ``"1:1"``, but not ``"1:2"``. Because this + matches a previous literal, will also match the leading + ``"1:1"`` in ``"1:10"``. If this is not desired, use + :class:`match_previous_expr`. Do *not* use with packrat parsing + enabled. + """ + rep = Forward() + + def copy_token_to_repeater(s, l, t): + if not t: + rep << Empty() + return + + if len(t) == 1: + rep << t[0] + return + + # flatten t tokens + tflat = _flatten(t.as_list()) + rep << And(Literal(tt) for tt in tflat) + + expr.add_parse_action(copy_token_to_repeater, callDuringTry=True) + rep.set_name("(prev) " + str(expr)) + return rep + + +def match_previous_expr(expr: ParserElement) -> ParserElement: + """Helper to define an expression that is indirectly defined from + the tokens matched in a previous expression, that is, it looks for + a 'repeat' of a previous expression. For example:: + + first = Word(nums) + second = match_previous_expr(first) + match_expr = first + ":" + second + + will match ``"1:1"``, but not ``"1:2"``. Because this + matches by expressions, will *not* match the leading ``"1:1"`` + in ``"1:10"``; the expressions are evaluated first, and then + compared, so ``"1"`` is compared with ``"10"``. Do *not* use + with packrat parsing enabled. + """ + rep = Forward() + e2 = expr.copy() + rep <<= e2 + + def copy_token_to_repeater(s, l, t): + matchTokens = _flatten(t.as_list()) + + def must_match_these_tokens(s, l, t): + theseTokens = _flatten(t.as_list()) + if theseTokens != matchTokens: + raise ParseException( + s, l, f"Expected {matchTokens}, found{theseTokens}" + ) + + rep.set_parse_action(must_match_these_tokens, callDuringTry=True) + + expr.add_parse_action(copy_token_to_repeater, callDuringTry=True) + rep.set_name("(prev) " + str(expr)) + return rep + + +def one_of( + strs: Union[typing.Iterable[str], str], + caseless: bool = False, + use_regex: bool = True, + as_keyword: bool = False, + *, + useRegex: bool = True, + asKeyword: bool = False, +) -> ParserElement: + """Helper to quickly define a set of alternative :class:`Literal` s, + and makes sure to do longest-first testing when there is a conflict, + regardless of the input order, but returns + a :class:`MatchFirst` for best performance. + + Parameters: + + - ``strs`` - a string of space-delimited literals, or a collection of + string literals + - ``caseless`` - treat all literals as caseless - (default= ``False``) + - ``use_regex`` - as an optimization, will + generate a :class:`Regex` object; otherwise, will generate + a :class:`MatchFirst` object (if ``caseless=True`` or ``as_keyword=True``, or if + creating a :class:`Regex` raises an exception) - (default= ``True``) + - ``as_keyword`` - enforce :class:`Keyword`-style matching on the + generated expressions - (default= ``False``) + - ``asKeyword`` and ``useRegex`` are retained for pre-PEP8 compatibility, + but will be removed in a future release + + Example:: + + comp_oper = one_of("< = > <= >= !=") + var = Word(alphas) + number = Word(nums) + term = var | number + comparison_expr = term + comp_oper + term + print(comparison_expr.search_string("B = 12 AA=23 B<=AA AA>12")) + + prints:: + + [['B', '=', '12'], ['AA', '=', '23'], ['B', '<=', 'AA'], ['AA', '>', '12']] + """ + asKeyword = asKeyword or as_keyword + useRegex = useRegex and use_regex + + if ( + isinstance(caseless, str_type) + and __diag__.warn_on_multiple_string_args_to_oneof + ): + warnings.warn( + "warn_on_multiple_string_args_to_oneof:" + " More than one string argument passed to one_of, pass" + " choices as a list or space-delimited string", + stacklevel=2, + ) + + if caseless: + is_equal = lambda a, b: a.upper() == b.upper() + masks = lambda a, b: b.upper().startswith(a.upper()) + else: + is_equal = operator.eq + masks = lambda a, b: b.startswith(a) + + symbols: list[str] + if isinstance(strs, str_type): + strs = typing.cast(str, strs) + symbols = strs.split() + elif isinstance(strs, Iterable): + symbols = list(strs) + else: + raise TypeError("Invalid argument to one_of, expected string or iterable") + if not symbols: + return NoMatch() + + # reorder given symbols to take care to avoid masking longer choices with shorter ones + # (but only if the given symbols are not just single characters) + i = 0 + while i < len(symbols) - 1: + cur = symbols[i] + for j, other in enumerate(symbols[i + 1 :]): + if is_equal(other, cur): + del symbols[i + j + 1] + break + if len(other) > len(cur) and masks(cur, other): + del symbols[i + j + 1] + symbols.insert(i, other) + break + else: + i += 1 + + if useRegex: + re_flags: int = re.IGNORECASE if caseless else 0 + + try: + if all(len(sym) == 1 for sym in symbols): + # symbols are just single characters, create range regex pattern + patt = f"[{''.join(_escape_regex_range_chars(sym) for sym in symbols)}]" + else: + patt = "|".join(re.escape(sym) for sym in symbols) + + # wrap with \b word break markers if defining as keywords + if asKeyword: + patt = rf"\b(?:{patt})\b" + + ret = Regex(patt, flags=re_flags) + ret.set_name(" | ".join(re.escape(s) for s in symbols)) + + if caseless: + # add parse action to return symbols as specified, not in random + # casing as found in input string + symbol_map = {sym.lower(): sym for sym in symbols} + ret.add_parse_action(lambda s, l, t: symbol_map[t[0].lower()]) + + return ret + + except re.error: + warnings.warn( + "Exception creating Regex for one_of, building MatchFirst", stacklevel=2 + ) + + # last resort, just use MatchFirst of Token class corresponding to caseless + # and asKeyword settings + CASELESS = KEYWORD = True + parse_element_class = { + (CASELESS, KEYWORD): CaselessKeyword, + (CASELESS, not KEYWORD): CaselessLiteral, + (not CASELESS, KEYWORD): Keyword, + (not CASELESS, not KEYWORD): Literal, + }[(caseless, asKeyword)] + return MatchFirst(parse_element_class(sym) for sym in symbols).set_name( + " | ".join(symbols) + ) + + +def dict_of(key: ParserElement, value: ParserElement) -> Dict: + """Helper to easily and clearly define a dictionary by specifying + the respective patterns for the key and value. Takes care of + defining the :class:`Dict`, :class:`ZeroOrMore`, and + :class:`Group` tokens in the proper order. The key pattern + can include delimiting markers or punctuation, as long as they are + suppressed, thereby leaving the significant key text. The value + pattern can include named results, so that the :class:`Dict` results + can include named token fields. + + Example:: + + text = "shape: SQUARE posn: upper left color: light blue texture: burlap" + attr_expr = (label + Suppress(':') + OneOrMore(data_word, stop_on=label).set_parse_action(' '.join)) + print(attr_expr[1, ...].parse_string(text).dump()) + + attr_label = label + attr_value = Suppress(':') + OneOrMore(data_word, stop_on=label).set_parse_action(' '.join) + + # similar to Dict, but simpler call format + result = dict_of(attr_label, attr_value).parse_string(text) + print(result.dump()) + print(result['shape']) + print(result.shape) # object attribute access works too + print(result.as_dict()) + + prints:: + + [['shape', 'SQUARE'], ['posn', 'upper left'], ['color', 'light blue'], ['texture', 'burlap']] + - color: 'light blue' + - posn: 'upper left' + - shape: 'SQUARE' + - texture: 'burlap' + SQUARE + SQUARE + {'color': 'light blue', 'shape': 'SQUARE', 'posn': 'upper left', 'texture': 'burlap'} + """ + return Dict(OneOrMore(Group(key + value))) + + +def original_text_for( + expr: ParserElement, as_string: bool = True, *, asString: bool = True +) -> ParserElement: + """Helper to return the original, untokenized text for a given + expression. Useful to restore the parsed fields of an HTML start + tag into the raw tag text itself, or to revert separate tokens with + intervening whitespace back to the original matching input text. By + default, returns a string containing the original parsed text. + + If the optional ``as_string`` argument is passed as + ``False``, then the return value is + a :class:`ParseResults` containing any results names that + were originally matched, and a single token containing the original + matched text from the input string. So if the expression passed to + :class:`original_text_for` contains expressions with defined + results names, you must set ``as_string`` to ``False`` if you + want to preserve those results name values. + + The ``asString`` pre-PEP8 argument is retained for compatibility, + but will be removed in a future release. + + Example:: + + src = "this is test bold text normal text " + for tag in ("b", "i"): + opener, closer = make_html_tags(tag) + patt = original_text_for(opener + ... + closer) + print(patt.search_string(src)[0]) + + prints:: + + [' bold text '] + ['text'] + """ + asString = asString and as_string + + locMarker = Empty().set_parse_action(lambda s, loc, t: loc) + endlocMarker = locMarker.copy() + endlocMarker.callPreparse = False + matchExpr = locMarker("_original_start") + expr + endlocMarker("_original_end") + if asString: + extractText = lambda s, l, t: s[t._original_start : t._original_end] + else: + + def extractText(s, l, t): + t[:] = [s[t.pop("_original_start") : t.pop("_original_end")]] + + matchExpr.set_parse_action(extractText) + matchExpr.ignoreExprs = expr.ignoreExprs + matchExpr.suppress_warning(Diagnostics.warn_ungrouped_named_tokens_in_collection) + return matchExpr + + +def ungroup(expr: ParserElement) -> ParserElement: + """Helper to undo pyparsing's default grouping of And expressions, + even if all but one are non-empty. + """ + return TokenConverter(expr).add_parse_action(lambda t: t[0]) + + +def locatedExpr(expr: ParserElement) -> ParserElement: + """ + (DEPRECATED - future code should use the :class:`Located` class) + Helper to decorate a returned token with its starting and ending + locations in the input string. + + This helper adds the following results names: + + - ``locn_start`` - location where matched expression begins + - ``locn_end`` - location where matched expression ends + - ``value`` - the actual parsed results + + Be careful if the input text contains ```` characters, you + may want to call :class:`ParserElement.parse_with_tabs` + + Example:: + + wd = Word(alphas) + for match in locatedExpr(wd).search_string("ljsdf123lksdjjf123lkkjj1222"): + print(match) + + prints:: + + [[0, 'ljsdf', 5]] + [[8, 'lksdjjf', 15]] + [[18, 'lkkjj', 23]] + """ + locator = Empty().set_parse_action(lambda ss, ll, tt: ll) + return Group( + locator("locn_start") + + expr("value") + + locator.copy().leaveWhitespace()("locn_end") + ) + + +# define special default value to permit None as a significant value for +# ignore_expr +_NO_IGNORE_EXPR_GIVEN = NoMatch() + + +def nested_expr( + opener: Union[str, ParserElement] = "(", + closer: Union[str, ParserElement] = ")", + content: typing.Optional[ParserElement] = None, + ignore_expr: ParserElement = _NO_IGNORE_EXPR_GIVEN, + *, + ignoreExpr: ParserElement = _NO_IGNORE_EXPR_GIVEN, +) -> ParserElement: + """Helper method for defining nested lists enclosed in opening and + closing delimiters (``"("`` and ``")"`` are the default). + + Parameters: + + - ``opener`` - opening character for a nested list + (default= ``"("``); can also be a pyparsing expression + - ``closer`` - closing character for a nested list + (default= ``")"``); can also be a pyparsing expression + - ``content`` - expression for items within the nested lists + (default= ``None``) + - ``ignore_expr`` - expression for ignoring opening and closing delimiters + (default= :class:`quoted_string`) + - ``ignoreExpr`` - this pre-PEP8 argument is retained for compatibility + but will be removed in a future release + + If an expression is not provided for the content argument, the + nested expression will capture all whitespace-delimited content + between delimiters as a list of separate values. + + Use the ``ignore_expr`` argument to define expressions that may + contain opening or closing characters that should not be treated as + opening or closing characters for nesting, such as quoted_string or + a comment expression. Specify multiple expressions using an + :class:`Or` or :class:`MatchFirst`. The default is + :class:`quoted_string`, but if no expressions are to be ignored, then + pass ``None`` for this argument. + + Example:: + + data_type = one_of("void int short long char float double") + decl_data_type = Combine(data_type + Opt(Word('*'))) + ident = Word(alphas+'_', alphanums+'_') + number = pyparsing_common.number + arg = Group(decl_data_type + ident) + LPAR, RPAR = map(Suppress, "()") + + code_body = nested_expr('{', '}', ignore_expr=(quoted_string | c_style_comment)) + + c_function = (decl_data_type("type") + + ident("name") + + LPAR + Opt(DelimitedList(arg), [])("args") + RPAR + + code_body("body")) + c_function.ignore(c_style_comment) + + source_code = ''' + int is_odd(int x) { + return (x%2); + } + + int dec_to_hex(char hchar) { + if (hchar >= '0' && hchar <= '9') { + return (ord(hchar)-ord('0')); + } else { + return (10+ord(hchar)-ord('A')); + } + } + ''' + for func in c_function.search_string(source_code): + print("%(name)s (%(type)s) args: %(args)s" % func) + + + prints:: + + is_odd (int) args: [['int', 'x']] + dec_to_hex (int) args: [['char', 'hchar']] + """ + if ignoreExpr != ignore_expr: + ignoreExpr = ignore_expr if ignoreExpr is _NO_IGNORE_EXPR_GIVEN else ignoreExpr + + if ignoreExpr is _NO_IGNORE_EXPR_GIVEN: + ignoreExpr = quoted_string() + + if opener == closer: + raise ValueError("opening and closing strings cannot be the same") + + if content is None: + if isinstance(opener, str_type) and isinstance(closer, str_type): + opener = typing.cast(str, opener) + closer = typing.cast(str, closer) + if len(opener) == 1 and len(closer) == 1: + if ignoreExpr is not None: + content = Combine( + OneOrMore( + ~ignoreExpr + + CharsNotIn( + opener + closer + ParserElement.DEFAULT_WHITE_CHARS, + exact=1, + ) + ) + ) + else: + content = Combine( + Empty() + + CharsNotIn( + opener + closer + ParserElement.DEFAULT_WHITE_CHARS + ) + ) + else: + if ignoreExpr is not None: + content = Combine( + OneOrMore( + ~ignoreExpr + + ~Literal(opener) + + ~Literal(closer) + + CharsNotIn(ParserElement.DEFAULT_WHITE_CHARS, exact=1) + ) + ) + else: + content = Combine( + OneOrMore( + ~Literal(opener) + + ~Literal(closer) + + CharsNotIn(ParserElement.DEFAULT_WHITE_CHARS, exact=1) + ) + ) + else: + raise ValueError( + "opening and closing arguments must be strings if no content expression is given" + ) + + # for these internally-created context expressions, simulate whitespace-skipping + if ParserElement.DEFAULT_WHITE_CHARS: + content.set_parse_action( + lambda t: t[0].strip(ParserElement.DEFAULT_WHITE_CHARS) + ) + + ret = Forward() + if ignoreExpr is not None: + ret <<= Group( + Suppress(opener) + ZeroOrMore(ignoreExpr | ret | content) + Suppress(closer) + ) + else: + ret <<= Group(Suppress(opener) + ZeroOrMore(ret | content) + Suppress(closer)) + + ret.set_name(f"nested {opener}{closer} expression") + + # don't override error message from content expressions + ret.errmsg = None + return ret + + +def _makeTags(tagStr, xml, suppress_LT=Suppress("<"), suppress_GT=Suppress(">")): + """Internal helper to construct opening and closing tag expressions, given a tag name""" + if isinstance(tagStr, str_type): + resname = tagStr + tagStr = Keyword(tagStr, caseless=not xml) + else: + resname = tagStr.name + + tagAttrName = Word(alphas, alphanums + "_-:") + if xml: + tagAttrValue = dbl_quoted_string.copy().set_parse_action(remove_quotes) + openTag = ( + suppress_LT + + tagStr("tag") + + Dict(ZeroOrMore(Group(tagAttrName + Suppress("=") + tagAttrValue))) + + Opt("/", default=[False])("empty").set_parse_action( + lambda s, l, t: t[0] == "/" + ) + + suppress_GT + ) + else: + tagAttrValue = quoted_string.copy().set_parse_action(remove_quotes) | Word( + printables, exclude_chars=">" + ) + openTag = ( + suppress_LT + + tagStr("tag") + + Dict( + ZeroOrMore( + Group( + tagAttrName.set_parse_action(lambda t: t[0].lower()) + + Opt(Suppress("=") + tagAttrValue) + ) + ) + ) + + Opt("/", default=[False])("empty").set_parse_action( + lambda s, l, t: t[0] == "/" + ) + + suppress_GT + ) + closeTag = Combine(Literal("", adjacent=False) + + openTag.set_name(f"<{resname}>") + # add start results name in parse action now that ungrouped names are not reported at two levels + openTag.add_parse_action( + lambda t: t.__setitem__( + "start" + "".join(resname.replace(":", " ").title().split()), t.copy() + ) + ) + closeTag = closeTag( + "end" + "".join(resname.replace(":", " ").title().split()) + ).set_name(f"") + openTag.tag = resname + closeTag.tag = resname + openTag.tag_body = SkipTo(closeTag()) + return openTag, closeTag + + +def make_html_tags( + tag_str: Union[str, ParserElement], +) -> tuple[ParserElement, ParserElement]: + """Helper to construct opening and closing tag expressions for HTML, + given a tag name. Matches tags in either upper or lower case, + attributes with namespaces and with quoted or unquoted values. + + Example:: + + text = 'More info at the pyparsing wiki page' + # make_html_tags returns pyparsing expressions for the opening and + # closing tags as a 2-tuple + a, a_end = make_html_tags("A") + link_expr = a + SkipTo(a_end)("link_text") + a_end + + for link in link_expr.search_string(text): + # attributes in the tag (like "href" shown here) are + # also accessible as named results + print(link.link_text, '->', link.href) + + prints:: + + pyparsing -> https://github.com/pyparsing/pyparsing/wiki + """ + return _makeTags(tag_str, False) + + +def make_xml_tags( + tag_str: Union[str, ParserElement], +) -> tuple[ParserElement, ParserElement]: + """Helper to construct opening and closing tag expressions for XML, + given a tag name. Matches tags only in the given upper/lower case. + + Example: similar to :class:`make_html_tags` + """ + return _makeTags(tag_str, True) + + +any_open_tag: ParserElement +any_close_tag: ParserElement +any_open_tag, any_close_tag = make_html_tags( + Word(alphas, alphanums + "_:").set_name("any tag") +) + +_htmlEntityMap = {k.rstrip(";"): v for k, v in html.entities.html5.items()} +_most_common_entities = "nbsp lt gt amp quot apos cent pound euro copy".replace( + " ", "|" +) +common_html_entity = Regex( + lambda: f"&(?P{_most_common_entities}|{make_compressed_re(_htmlEntityMap)});" +).set_name("common HTML entity") + + +def replace_html_entity(s, l, t): + """Helper parser action to replace common HTML entities with their special characters""" + return _htmlEntityMap.get(t.entity) + + +class OpAssoc(Enum): + """Enumeration of operator associativity + - used in constructing InfixNotationOperatorSpec for :class:`infix_notation`""" + + LEFT = 1 + RIGHT = 2 + + +InfixNotationOperatorArgType = Union[ + ParserElement, str, tuple[Union[ParserElement, str], Union[ParserElement, str]] +] +InfixNotationOperatorSpec = Union[ + tuple[ + InfixNotationOperatorArgType, + int, + OpAssoc, + typing.Optional[ParseAction], + ], + tuple[ + InfixNotationOperatorArgType, + int, + OpAssoc, + ], +] + + +def infix_notation( + base_expr: ParserElement, + op_list: list[InfixNotationOperatorSpec], + lpar: Union[str, ParserElement] = Suppress("("), + rpar: Union[str, ParserElement] = Suppress(")"), +) -> Forward: + """Helper method for constructing grammars of expressions made up of + operators working in a precedence hierarchy. Operators may be unary + or binary, left- or right-associative. Parse actions can also be + attached to operator expressions. The generated parser will also + recognize the use of parentheses to override operator precedences + (see example below). + + Note: if you define a deep operator list, you may see performance + issues when using infix_notation. See + :class:`ParserElement.enable_packrat` for a mechanism to potentially + improve your parser performance. + + Parameters: + + - ``base_expr`` - expression representing the most basic operand to + be used in the expression + - ``op_list`` - list of tuples, one for each operator precedence level + in the expression grammar; each tuple is of the form ``(op_expr, + num_operands, right_left_assoc, (optional)parse_action)``, where: + + - ``op_expr`` is the pyparsing expression for the operator; may also + be a string, which will be converted to a Literal; if ``num_operands`` + is 3, ``op_expr`` is a tuple of two expressions, for the two + operators separating the 3 terms + - ``num_operands`` is the number of terms for this operator (must be 1, + 2, or 3) + - ``right_left_assoc`` is the indicator whether the operator is right + or left associative, using the pyparsing-defined constants + ``OpAssoc.RIGHT`` and ``OpAssoc.LEFT``. + - ``parse_action`` is the parse action to be associated with + expressions matching this operator expression (the parse action + tuple member may be omitted); if the parse action is passed + a tuple or list of functions, this is equivalent to calling + ``set_parse_action(*fn)`` + (:class:`ParserElement.set_parse_action`) + - ``lpar`` - expression for matching left-parentheses; if passed as a + str, then will be parsed as ``Suppress(lpar)``. If lpar is passed as + an expression (such as ``Literal('(')``), then it will be kept in + the parsed results, and grouped with them. (default= ``Suppress('(')``) + - ``rpar`` - expression for matching right-parentheses; if passed as a + str, then will be parsed as ``Suppress(rpar)``. If rpar is passed as + an expression (such as ``Literal(')')``), then it will be kept in + the parsed results, and grouped with them. (default= ``Suppress(')')``) + + Example:: + + # simple example of four-function arithmetic with ints and + # variable names + integer = pyparsing_common.signed_integer + varname = pyparsing_common.identifier + + arith_expr = infix_notation(integer | varname, + [ + ('-', 1, OpAssoc.RIGHT), + (one_of('* /'), 2, OpAssoc.LEFT), + (one_of('+ -'), 2, OpAssoc.LEFT), + ]) + + arith_expr.run_tests(''' + 5+3*6 + (5+3)*6 + -2--11 + ''', full_dump=False) + + prints:: + + 5+3*6 + [[5, '+', [3, '*', 6]]] + + (5+3)*6 + [[[5, '+', 3], '*', 6]] + + (5+x)*y + [[[5, '+', 'x'], '*', 'y']] + + -2--11 + [[['-', 2], '-', ['-', 11]]] + """ + + # captive version of FollowedBy that does not do parse actions or capture results names + class _FB(FollowedBy): + def parseImpl(self, instring, loc, doActions=True): + self.expr.try_parse(instring, loc) + return loc, [] + + _FB.__name__ = "FollowedBy>" + + ret = Forward() + ret.set_name(f"{base_expr.name}_expression") + if isinstance(lpar, str): + lpar = Suppress(lpar) + if isinstance(rpar, str): + rpar = Suppress(rpar) + + nested_expr = (lpar + ret + rpar).set_name(f"nested_{base_expr.name}") + + # if lpar and rpar are not suppressed, wrap in group + if not (isinstance(lpar, Suppress) and isinstance(rpar, Suppress)): + lastExpr = base_expr | Group(nested_expr) + else: + lastExpr = base_expr | nested_expr + + arity: int + rightLeftAssoc: opAssoc + pa: typing.Optional[ParseAction] + opExpr1: ParserElement + opExpr2: ParserElement + matchExpr: ParserElement + match_lookahead: ParserElement + for operDef in op_list: + opExpr, arity, rightLeftAssoc, pa = (operDef + (None,))[:4] # type: ignore[assignment] + if isinstance(opExpr, str_type): + opExpr = ParserElement._literalStringClass(opExpr) + opExpr = typing.cast(ParserElement, opExpr) + if arity == 3: + if not isinstance(opExpr, (tuple, list)) or len(opExpr) != 2: + raise ValueError( + "if numterms=3, opExpr must be a tuple or list of two expressions" + ) + opExpr1, opExpr2 = opExpr + term_name = f"{opExpr1}{opExpr2} operations" + else: + term_name = f"{opExpr} operations" + + if not 1 <= arity <= 3: + raise ValueError("operator must be unary (1), binary (2), or ternary (3)") + + if rightLeftAssoc not in (OpAssoc.LEFT, OpAssoc.RIGHT): + raise ValueError("operator must indicate right or left associativity") + + thisExpr: ParserElement = Forward().set_name(term_name) + thisExpr = typing.cast(Forward, thisExpr) + match_lookahead = And([]) + if rightLeftAssoc is OpAssoc.LEFT: + if arity == 1: + match_lookahead = _FB(lastExpr + opExpr) + matchExpr = Group(lastExpr + opExpr[1, ...]) + elif arity == 2: + if opExpr is not None: + match_lookahead = _FB(lastExpr + opExpr + lastExpr) + matchExpr = Group(lastExpr + (opExpr + lastExpr)[1, ...]) + else: + match_lookahead = _FB(lastExpr + lastExpr) + matchExpr = Group(lastExpr[2, ...]) + elif arity == 3: + match_lookahead = _FB( + lastExpr + opExpr1 + lastExpr + opExpr2 + lastExpr + ) + matchExpr = Group( + lastExpr + (opExpr1 + lastExpr + opExpr2 + lastExpr)[1, ...] + ) + elif rightLeftAssoc is OpAssoc.RIGHT: + if arity == 1: + # try to avoid LR with this extra test + if not isinstance(opExpr, Opt): + opExpr = Opt(opExpr) + match_lookahead = _FB(opExpr.expr + thisExpr) + matchExpr = Group(opExpr + thisExpr) + elif arity == 2: + if opExpr is not None: + match_lookahead = _FB(lastExpr + opExpr + thisExpr) + matchExpr = Group(lastExpr + (opExpr + thisExpr)[1, ...]) + else: + match_lookahead = _FB(lastExpr + thisExpr) + matchExpr = Group(lastExpr + thisExpr[1, ...]) + elif arity == 3: + match_lookahead = _FB( + lastExpr + opExpr1 + thisExpr + opExpr2 + thisExpr + ) + matchExpr = Group(lastExpr + opExpr1 + thisExpr + opExpr2 + thisExpr) + + # suppress lookahead expr from railroad diagrams + match_lookahead.show_in_diagram = False + + # TODO - determine why this statement can't be included in the following + # if pa block + matchExpr = match_lookahead + matchExpr + + if pa: + if isinstance(pa, (tuple, list)): + matchExpr.set_parse_action(*pa) + else: + matchExpr.set_parse_action(pa) + + thisExpr <<= (matchExpr | lastExpr).setName(term_name) + lastExpr = thisExpr + + ret <<= lastExpr + return ret + + +def indentedBlock(blockStatementExpr, indentStack, indent=True, backup_stacks=[]): + """ + (DEPRECATED - use :class:`IndentedBlock` class instead) + Helper method for defining space-delimited indentation blocks, + such as those used to define block statements in Python source code. + + Parameters: + + - ``blockStatementExpr`` - expression defining syntax of statement that + is repeated within the indented block + - ``indentStack`` - list created by caller to manage indentation stack + (multiple ``statementWithIndentedBlock`` expressions within a single + grammar should share a common ``indentStack``) + - ``indent`` - boolean indicating whether block must be indented beyond + the current level; set to ``False`` for block of left-most statements + (default= ``True``) + + A valid block must contain at least one ``blockStatement``. + + (Note that indentedBlock uses internal parse actions which make it + incompatible with packrat parsing.) + + Example:: + + data = ''' + def A(z): + A1 + B = 100 + G = A2 + A2 + A3 + B + def BB(a,b,c): + BB1 + def BBA(): + bba1 + bba2 + bba3 + C + D + def spam(x,y): + def eggs(z): + pass + ''' + + + indentStack = [1] + stmt = Forward() + + identifier = Word(alphas, alphanums) + funcDecl = ("def" + identifier + Group("(" + Opt(delimitedList(identifier)) + ")") + ":") + func_body = indentedBlock(stmt, indentStack) + funcDef = Group(funcDecl + func_body) + + rvalue = Forward() + funcCall = Group(identifier + "(" + Opt(delimitedList(rvalue)) + ")") + rvalue << (funcCall | identifier | Word(nums)) + assignment = Group(identifier + "=" + rvalue) + stmt << (funcDef | assignment | identifier) + + module_body = stmt[1, ...] + + parseTree = module_body.parseString(data) + parseTree.pprint() + + prints:: + + [['def', + 'A', + ['(', 'z', ')'], + ':', + [['A1'], [['B', '=', '100']], [['G', '=', 'A2']], ['A2'], ['A3']]], + 'B', + ['def', + 'BB', + ['(', 'a', 'b', 'c', ')'], + ':', + [['BB1'], [['def', 'BBA', ['(', ')'], ':', [['bba1'], ['bba2'], ['bba3']]]]]], + 'C', + 'D', + ['def', + 'spam', + ['(', 'x', 'y', ')'], + ':', + [[['def', 'eggs', ['(', 'z', ')'], ':', [['pass']]]]]]] + """ + backup_stacks.append(indentStack[:]) + + def reset_stack(): + indentStack[:] = backup_stacks[-1] + + def checkPeerIndent(s, l, t): + if l >= len(s): + return + curCol = col(l, s) + if curCol != indentStack[-1]: + if curCol > indentStack[-1]: + raise ParseException(s, l, "illegal nesting") + raise ParseException(s, l, "not a peer entry") + + def checkSubIndent(s, l, t): + curCol = col(l, s) + if curCol > indentStack[-1]: + indentStack.append(curCol) + else: + raise ParseException(s, l, "not a subentry") + + def checkUnindent(s, l, t): + if l >= len(s): + return + curCol = col(l, s) + if not (indentStack and curCol in indentStack): + raise ParseException(s, l, "not an unindent") + if curCol < indentStack[-1]: + indentStack.pop() + + NL = OneOrMore(LineEnd().set_whitespace_chars("\t ").suppress()) + INDENT = (Empty() + Empty().set_parse_action(checkSubIndent)).set_name("INDENT") + PEER = Empty().set_parse_action(checkPeerIndent).set_name("") + UNDENT = Empty().set_parse_action(checkUnindent).set_name("UNINDENT") + if indent: + smExpr = Group( + Opt(NL) + + INDENT + + OneOrMore(PEER + Group(blockStatementExpr) + Opt(NL)) + + UNDENT + ) + else: + smExpr = Group( + Opt(NL) + + OneOrMore(PEER + Group(blockStatementExpr) + Opt(NL)) + + Opt(UNDENT) + ) + + # add a parse action to remove backup_stack from list of backups + smExpr.add_parse_action( + lambda: backup_stacks.pop(-1) and None if backup_stacks else None + ) + smExpr.set_fail_action(lambda a, b, c, d: reset_stack()) + blockStatementExpr.ignore(_bslash + LineEnd()) + return smExpr.set_name("indented block") + + +# it's easy to get these comment structures wrong - they're very common, +# so may as well make them available +c_style_comment = Regex(r"/\*(?:[^*]|\*(?!/))*\*\/").set_name("C style comment") +"Comment of the form ``/* ... */``" + +html_comment = Regex(r"").set_name("HTML comment") +"Comment of the form ````" + +rest_of_line = Regex(r".*").leave_whitespace().set_name("rest of line") +dbl_slash_comment = Regex(r"//(?:\\\n|[^\n])*").set_name("// comment") +"Comment of the form ``// ... (to end of line)``" + +cpp_style_comment = Regex( + r"(?:/\*(?:[^*]|\*(?!/))*\*\/)|(?://(?:\\\n|[^\n])*)" +).set_name("C++ style comment") +"Comment of either form :class:`c_style_comment` or :class:`dbl_slash_comment`" + +java_style_comment = cpp_style_comment +"Same as :class:`cpp_style_comment`" + +python_style_comment = Regex(r"#.*").set_name("Python style comment") +"Comment of the form ``# ... (to end of line)``" + + +# build list of built-in expressions, for future reference if a global default value +# gets updated +_builtin_exprs: list[ParserElement] = [ + v for v in vars().values() if isinstance(v, ParserElement) +] + + +# compatibility function, superseded by DelimitedList class +def delimited_list( + expr: Union[str, ParserElement], + delim: Union[str, ParserElement] = ",", + combine: bool = False, + min: typing.Optional[int] = None, + max: typing.Optional[int] = None, + *, + allow_trailing_delim: bool = False, +) -> ParserElement: + """(DEPRECATED - use :class:`DelimitedList` class)""" + return DelimitedList( + expr, delim, combine, min, max, allow_trailing_delim=allow_trailing_delim + ) + + +# Compatibility synonyms +# fmt: off +opAssoc = OpAssoc +anyOpenTag = any_open_tag +anyCloseTag = any_close_tag +commonHTMLEntity = common_html_entity +cStyleComment = c_style_comment +htmlComment = html_comment +restOfLine = rest_of_line +dblSlashComment = dbl_slash_comment +cppStyleComment = cpp_style_comment +javaStyleComment = java_style_comment +pythonStyleComment = python_style_comment +delimitedList = replaced_by_pep8("delimitedList", DelimitedList) +delimited_list = replaced_by_pep8("delimited_list", DelimitedList) +countedArray = replaced_by_pep8("countedArray", counted_array) +matchPreviousLiteral = replaced_by_pep8("matchPreviousLiteral", match_previous_literal) +matchPreviousExpr = replaced_by_pep8("matchPreviousExpr", match_previous_expr) +oneOf = replaced_by_pep8("oneOf", one_of) +dictOf = replaced_by_pep8("dictOf", dict_of) +originalTextFor = replaced_by_pep8("originalTextFor", original_text_for) +nestedExpr = replaced_by_pep8("nestedExpr", nested_expr) +makeHTMLTags = replaced_by_pep8("makeHTMLTags", make_html_tags) +makeXMLTags = replaced_by_pep8("makeXMLTags", make_xml_tags) +replaceHTMLEntity = replaced_by_pep8("replaceHTMLEntity", replace_html_entity) +infixNotation = replaced_by_pep8("infixNotation", infix_notation) +# fmt: on diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/py.typed b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/results.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/results.py new file mode 100644 index 0000000000000000000000000000000000000000..956230352c84617cb6565142f28de874fa2bfacb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/results.py @@ -0,0 +1,816 @@ +# results.py +from __future__ import annotations + +import collections +from collections.abc import ( + MutableMapping, + Mapping, + MutableSequence, + Iterator, + Iterable, +) +import pprint +from typing import Any + +from .util import replaced_by_pep8 + + +str_type: tuple[type, ...] = (str, bytes) +_generator_type = type((_ for _ in ())) + + +class _ParseResultsWithOffset: + tup: tuple[ParseResults, int] + __slots__ = ["tup"] + + def __init__(self, p1: ParseResults, p2: int) -> None: + self.tup: tuple[ParseResults, int] = (p1, p2) + + def __getitem__(self, i): + return self.tup[i] + + def __getstate__(self): + return self.tup + + def __setstate__(self, *args): + self.tup = args[0] + + +class ParseResults: + """Structured parse results, to provide multiple means of access to + the parsed data: + + - as a list (``len(results)``) + - by list index (``results[0], results[1]``, etc.) + - by attribute (``results.`` - see :class:`ParserElement.set_results_name`) + + Example:: + + integer = Word(nums) + date_str = (integer.set_results_name("year") + '/' + + integer.set_results_name("month") + '/' + + integer.set_results_name("day")) + # equivalent form: + # date_str = (integer("year") + '/' + # + integer("month") + '/' + # + integer("day")) + + # parse_string returns a ParseResults object + result = date_str.parse_string("1999/12/31") + + def test(s, fn=repr): + print(f"{s} -> {fn(eval(s))}") + test("list(result)") + test("result[0]") + test("result['month']") + test("result.day") + test("'month' in result") + test("'minutes' in result") + test("result.dump()", str) + + prints:: + + list(result) -> ['1999', '/', '12', '/', '31'] + result[0] -> '1999' + result['month'] -> '12' + result.day -> '31' + 'month' in result -> True + 'minutes' in result -> False + result.dump() -> ['1999', '/', '12', '/', '31'] + - day: '31' + - month: '12' + - year: '1999' + """ + + _null_values: tuple[Any, ...] = (None, [], ()) + + _name: str + _parent: ParseResults + _all_names: set[str] + _modal: bool + _toklist: list[Any] + _tokdict: dict[str, Any] + + __slots__ = ( + "_name", + "_parent", + "_all_names", + "_modal", + "_toklist", + "_tokdict", + ) + + class List(list): + """ + Simple wrapper class to distinguish parsed list results that should be preserved + as actual Python lists, instead of being converted to :class:`ParseResults`:: + + LBRACK, RBRACK = map(pp.Suppress, "[]") + element = pp.Forward() + item = ppc.integer + element_list = LBRACK + pp.DelimitedList(element) + RBRACK + + # add parse actions to convert from ParseResults to actual Python collection types + def as_python_list(t): + return pp.ParseResults.List(t.as_list()) + element_list.add_parse_action(as_python_list) + + element <<= item | element_list + + element.run_tests(''' + 100 + [2,3,4] + [[2, 1],3,4] + [(2, 1),3,4] + (2,3,4) + ''', post_parse=lambda s, r: (r[0], type(r[0]))) + + prints:: + + 100 + (100, ) + + [2,3,4] + ([2, 3, 4], ) + + [[2, 1],3,4] + ([[2, 1], 3, 4], ) + + (Used internally by :class:`Group` when `aslist=True`.) + """ + + def __new__(cls, contained=None): + if contained is None: + contained = [] + + if not isinstance(contained, list): + raise TypeError( + f"{cls.__name__} may only be constructed with a list, not {type(contained).__name__}" + ) + + return list.__new__(cls) + + def __new__(cls, toklist=None, name=None, **kwargs): + if isinstance(toklist, ParseResults): + return toklist + self = object.__new__(cls) + self._name = None + self._parent = None + self._all_names = set() + + if toklist is None: + self._toklist = [] + elif isinstance(toklist, (list, _generator_type)): + self._toklist = ( + [toklist[:]] + if isinstance(toklist, ParseResults.List) + else list(toklist) + ) + else: + self._toklist = [toklist] + self._tokdict = dict() + return self + + # Performance tuning: we construct a *lot* of these, so keep this + # constructor as small and fast as possible + def __init__( + self, toklist=None, name=None, asList=True, modal=True, isinstance=isinstance + ) -> None: + self._tokdict: dict[str, _ParseResultsWithOffset] + self._modal = modal + + if name is None or name == "": + return + + if isinstance(name, int): + name = str(name) + + if not modal: + self._all_names = {name} + + self._name = name + + if toklist in self._null_values: + return + + if isinstance(toklist, (str_type, type)): + toklist = [toklist] + + if asList: + if isinstance(toklist, ParseResults): + self[name] = _ParseResultsWithOffset(ParseResults(toklist._toklist), 0) + else: + self[name] = _ParseResultsWithOffset(ParseResults(toklist[0]), 0) + self[name]._name = name + return + + try: + self[name] = toklist[0] + except (KeyError, TypeError, IndexError): + if toklist is not self: + self[name] = toklist + else: + self._name = name + + def __getitem__(self, i): + if isinstance(i, (int, slice)): + return self._toklist[i] + + if i not in self._all_names: + return self._tokdict[i][-1][0] + + return ParseResults([v[0] for v in self._tokdict[i]]) + + def __setitem__(self, k, v, isinstance=isinstance): + if isinstance(v, _ParseResultsWithOffset): + self._tokdict[k] = self._tokdict.get(k, list()) + [v] + sub = v[0] + elif isinstance(k, (int, slice)): + self._toklist[k] = v + sub = v + else: + self._tokdict[k] = self._tokdict.get(k, []) + [ + _ParseResultsWithOffset(v, 0) + ] + sub = v + if isinstance(sub, ParseResults): + sub._parent = self + + def __delitem__(self, i): + if not isinstance(i, (int, slice)): + del self._tokdict[i] + return + + mylen = len(self._toklist) + del self._toklist[i] + + # convert int to slice + if isinstance(i, int): + if i < 0: + i += mylen + i = slice(i, i + 1) + # get removed indices + removed = list(range(*i.indices(mylen))) + removed.reverse() + # fixup indices in token dictionary + for occurrences in self._tokdict.values(): + for j in removed: + for k, (value, position) in enumerate(occurrences): + occurrences[k] = _ParseResultsWithOffset( + value, position - (position > j) + ) + + def __contains__(self, k) -> bool: + return k in self._tokdict + + def __len__(self) -> int: + return len(self._toklist) + + def __bool__(self) -> bool: + return not not (self._toklist or self._tokdict) + + def __iter__(self) -> Iterator: + return iter(self._toklist) + + def __reversed__(self) -> Iterator: + return iter(self._toklist[::-1]) + + def keys(self): + return iter(self._tokdict) + + def values(self): + return (self[k] for k in self.keys()) + + def items(self): + return ((k, self[k]) for k in self.keys()) + + def haskeys(self) -> bool: + """ + Since ``keys()`` returns an iterator, this method is helpful in bypassing + code that looks for the existence of any defined results names.""" + return not not self._tokdict + + def pop(self, *args, **kwargs): + """ + Removes and returns item at specified index (default= ``last``). + Supports both ``list`` and ``dict`` semantics for ``pop()``. If + passed no argument or an integer argument, it will use ``list`` + semantics and pop tokens from the list of parsed tokens. If passed + a non-integer argument (most likely a string), it will use ``dict`` + semantics and pop the corresponding value from any defined results + names. A second default return value argument is supported, just as in + ``dict.pop()``. + + Example:: + + numlist = Word(nums)[...] + print(numlist.parse_string("0 123 321")) # -> ['0', '123', '321'] + + def remove_first(tokens): + tokens.pop(0) + numlist.add_parse_action(remove_first) + print(numlist.parse_string("0 123 321")) # -> ['123', '321'] + + label = Word(alphas) + patt = label("LABEL") + Word(nums)[1, ...] + print(patt.parse_string("AAB 123 321").dump()) + + # Use pop() in a parse action to remove named result (note that corresponding value is not + # removed from list form of results) + def remove_LABEL(tokens): + tokens.pop("LABEL") + return tokens + patt.add_parse_action(remove_LABEL) + print(patt.parse_string("AAB 123 321").dump()) + + prints:: + + ['AAB', '123', '321'] + - LABEL: 'AAB' + + ['AAB', '123', '321'] + """ + if not args: + args = [-1] + for k, v in kwargs.items(): + if k == "default": + args = (args[0], v) + else: + raise TypeError(f"pop() got an unexpected keyword argument {k!r}") + if isinstance(args[0], int) or len(args) == 1 or args[0] in self: + index = args[0] + ret = self[index] + del self[index] + return ret + else: + defaultvalue = args[1] + return defaultvalue + + def get(self, key, default_value=None): + """ + Returns named result matching the given key, or if there is no + such name, then returns the given ``default_value`` or ``None`` if no + ``default_value`` is specified. + + Similar to ``dict.get()``. + + Example:: + + integer = Word(nums) + date_str = integer("year") + '/' + integer("month") + '/' + integer("day") + + result = date_str.parse_string("1999/12/31") + print(result.get("year")) # -> '1999' + print(result.get("hour", "not specified")) # -> 'not specified' + print(result.get("hour")) # -> None + """ + if key in self: + return self[key] + else: + return default_value + + def insert(self, index, ins_string): + """ + Inserts new element at location index in the list of parsed tokens. + + Similar to ``list.insert()``. + + Example:: + + numlist = Word(nums)[...] + print(numlist.parse_string("0 123 321")) # -> ['0', '123', '321'] + + # use a parse action to insert the parse location in the front of the parsed results + def insert_locn(locn, tokens): + tokens.insert(0, locn) + numlist.add_parse_action(insert_locn) + print(numlist.parse_string("0 123 321")) # -> [0, '0', '123', '321'] + """ + self._toklist.insert(index, ins_string) + # fixup indices in token dictionary + for occurrences in self._tokdict.values(): + for k, (value, position) in enumerate(occurrences): + occurrences[k] = _ParseResultsWithOffset( + value, position + (position > index) + ) + + def append(self, item): + """ + Add single element to end of ``ParseResults`` list of elements. + + Example:: + + numlist = Word(nums)[...] + print(numlist.parse_string("0 123 321")) # -> ['0', '123', '321'] + + # use a parse action to compute the sum of the parsed integers, and add it to the end + def append_sum(tokens): + tokens.append(sum(map(int, tokens))) + numlist.add_parse_action(append_sum) + print(numlist.parse_string("0 123 321")) # -> ['0', '123', '321', 444] + """ + self._toklist.append(item) + + def extend(self, itemseq): + """ + Add sequence of elements to end of ``ParseResults`` list of elements. + + Example:: + + patt = Word(alphas)[1, ...] + + # use a parse action to append the reverse of the matched strings, to make a palindrome + def make_palindrome(tokens): + tokens.extend(reversed([t[::-1] for t in tokens])) + return ''.join(tokens) + patt.add_parse_action(make_palindrome) + print(patt.parse_string("lskdj sdlkjf lksd")) # -> 'lskdjsdlkjflksddsklfjkldsjdksl' + """ + if isinstance(itemseq, ParseResults): + self.__iadd__(itemseq) + else: + self._toklist.extend(itemseq) + + def clear(self): + """ + Clear all elements and results names. + """ + del self._toklist[:] + self._tokdict.clear() + + def __getattr__(self, name): + try: + return self[name] + except KeyError: + if name.startswith("__"): + raise AttributeError(name) + return "" + + def __add__(self, other: ParseResults) -> ParseResults: + ret = self.copy() + ret += other + return ret + + def __iadd__(self, other: ParseResults) -> ParseResults: + if not other: + return self + + if other._tokdict: + offset = len(self._toklist) + addoffset = lambda a: offset if a < 0 else a + offset + otheritems = other._tokdict.items() + otherdictitems = [ + (k, _ParseResultsWithOffset(v[0], addoffset(v[1]))) + for k, vlist in otheritems + for v in vlist + ] + for k, v in otherdictitems: + self[k] = v + if isinstance(v[0], ParseResults): + v[0]._parent = self + + self._toklist += other._toklist + self._all_names |= other._all_names + return self + + def __radd__(self, other) -> ParseResults: + if isinstance(other, int) and other == 0: + # useful for merging many ParseResults using sum() builtin + return self.copy() + else: + # this may raise a TypeError - so be it + return other + self + + def __repr__(self) -> str: + return f"{type(self).__name__}({self._toklist!r}, {self.as_dict()})" + + def __str__(self) -> str: + return ( + "[" + + ", ".join( + [ + str(i) if isinstance(i, ParseResults) else repr(i) + for i in self._toklist + ] + ) + + "]" + ) + + def _asStringList(self, sep=""): + out = [] + for item in self._toklist: + if out and sep: + out.append(sep) + if isinstance(item, ParseResults): + out += item._asStringList() + else: + out.append(str(item)) + return out + + def as_list(self, *, flatten: bool = False) -> list: + """ + Returns the parse results as a nested list of matching tokens, all converted to strings. + If flatten is True, all the nesting levels in the returned list are collapsed. + + Example:: + + patt = Word(alphas)[1, ...] + result = patt.parse_string("sldkj lsdkj sldkj") + # even though the result prints in string-like form, it is actually a pyparsing ParseResults + print(type(result), result) # -> ['sldkj', 'lsdkj', 'sldkj'] + + # Use as_list() to create an actual list + result_list = result.as_list() + print(type(result_list), result_list) # -> ['sldkj', 'lsdkj', 'sldkj'] + """ + + def flattened(pr): + to_visit = collections.deque([*self]) + while to_visit: + to_do = to_visit.popleft() + if isinstance(to_do, ParseResults): + to_visit.extendleft(to_do[::-1]) + else: + yield to_do + + if flatten: + return [*flattened(self)] + else: + return [ + res.as_list() if isinstance(res, ParseResults) else res + for res in self._toklist + ] + + def as_dict(self) -> dict: + """ + Returns the named parse results as a nested dictionary. + + Example:: + + integer = Word(nums) + date_str = integer("year") + '/' + integer("month") + '/' + integer("day") + + result = date_str.parse_string('12/31/1999') + print(type(result), repr(result)) # -> (['12', '/', '31', '/', '1999'], {'day': [('1999', 4)], 'year': [('12', 0)], 'month': [('31', 2)]}) + + result_dict = result.as_dict() + print(type(result_dict), repr(result_dict)) # -> {'day': '1999', 'year': '12', 'month': '31'} + + # even though a ParseResults supports dict-like access, sometime you just need to have a dict + import json + print(json.dumps(result)) # -> Exception: TypeError: ... is not JSON serializable + print(json.dumps(result.as_dict())) # -> {"month": "31", "day": "1999", "year": "12"} + """ + + def to_item(obj): + if isinstance(obj, ParseResults): + return obj.as_dict() if obj.haskeys() else [to_item(v) for v in obj] + else: + return obj + + return dict((k, to_item(v)) for k, v in self.items()) + + def copy(self) -> ParseResults: + """ + Returns a new shallow copy of a :class:`ParseResults` object. `ParseResults` + items contained within the source are shared with the copy. Use + :class:`ParseResults.deepcopy()` to create a copy with its own separate + content values. + """ + ret = ParseResults(self._toklist) + ret._tokdict = self._tokdict.copy() + ret._parent = self._parent + ret._all_names |= self._all_names + ret._name = self._name + return ret + + def deepcopy(self) -> ParseResults: + """ + Returns a new deep copy of a :class:`ParseResults` object. + """ + ret = self.copy() + # replace values with copies if they are of known mutable types + for i, obj in enumerate(self._toklist): + if isinstance(obj, ParseResults): + ret._toklist[i] = obj.deepcopy() + elif isinstance(obj, (str, bytes)): + pass + elif isinstance(obj, MutableMapping): + ret._toklist[i] = dest = type(obj)() + for k, v in obj.items(): + dest[k] = v.deepcopy() if isinstance(v, ParseResults) else v + elif isinstance(obj, Iterable): + ret._toklist[i] = type(obj)( + v.deepcopy() if isinstance(v, ParseResults) else v for v in obj # type: ignore[call-arg] + ) + return ret + + def get_name(self) -> str | None: + r""" + Returns the results name for this token expression. Useful when several + different expressions might match at a particular location. + + Example:: + + integer = Word(nums) + ssn_expr = Regex(r"\d\d\d-\d\d-\d\d\d\d") + house_number_expr = Suppress('#') + Word(nums, alphanums) + user_data = (Group(house_number_expr)("house_number") + | Group(ssn_expr)("ssn") + | Group(integer)("age")) + user_info = user_data[1, ...] + + result = user_info.parse_string("22 111-22-3333 #221B") + for item in result: + print(item.get_name(), ':', item[0]) + + prints:: + + age : 22 + ssn : 111-22-3333 + house_number : 221B + """ + if self._name: + return self._name + elif self._parent: + par: ParseResults = self._parent + parent_tokdict_items = par._tokdict.items() + return next( + ( + k + for k, vlist in parent_tokdict_items + for v, loc in vlist + if v is self + ), + None, + ) + elif ( + len(self) == 1 + and len(self._tokdict) == 1 + and next(iter(self._tokdict.values()))[0][1] in (0, -1) + ): + return next(iter(self._tokdict.keys())) + else: + return None + + def dump(self, indent="", full=True, include_list=True, _depth=0) -> str: + """ + Diagnostic method for listing out the contents of + a :class:`ParseResults`. Accepts an optional ``indent`` argument so + that this string can be embedded in a nested display of other data. + + Example:: + + integer = Word(nums) + date_str = integer("year") + '/' + integer("month") + '/' + integer("day") + + result = date_str.parse_string('1999/12/31') + print(result.dump()) + + prints:: + + ['1999', '/', '12', '/', '31'] + - day: '31' + - month: '12' + - year: '1999' + """ + out = [] + NL = "\n" + out.append(indent + str(self.as_list()) if include_list else "") + + if not full: + return "".join(out) + + if self.haskeys(): + items = sorted((str(k), v) for k, v in self.items()) + for k, v in items: + if out: + out.append(NL) + out.append(f"{indent}{(' ' * _depth)}- {k}: ") + if not isinstance(v, ParseResults): + out.append(repr(v)) + continue + + if not v: + out.append(str(v)) + continue + + out.append( + v.dump( + indent=indent, + full=full, + include_list=include_list, + _depth=_depth + 1, + ) + ) + if not any(isinstance(vv, ParseResults) for vv in self): + return "".join(out) + + v = self + incr = " " + nl = "\n" + for i, vv in enumerate(v): + if isinstance(vv, ParseResults): + vv_dump = vv.dump( + indent=indent, + full=full, + include_list=include_list, + _depth=_depth + 1, + ) + out.append( + f"{nl}{indent}{incr * _depth}[{i}]:{nl}{indent}{incr * (_depth + 1)}{vv_dump}" + ) + else: + out.append( + f"{nl}{indent}{incr * _depth}[{i}]:{nl}{indent}{incr * (_depth + 1)}{vv}" + ) + + return "".join(out) + + def pprint(self, *args, **kwargs): + """ + Pretty-printer for parsed results as a list, using the + `pprint `_ module. + Accepts additional positional or keyword args as defined for + `pprint.pprint `_ . + + Example:: + + ident = Word(alphas, alphanums) + num = Word(nums) + func = Forward() + term = ident | num | Group('(' + func + ')') + func <<= ident + Group(Optional(DelimitedList(term))) + result = func.parse_string("fna a,b,(fnb c,d,200),100") + result.pprint(width=40) + + prints:: + + ['fna', + ['a', + 'b', + ['(', 'fnb', ['c', 'd', '200'], ')'], + '100']] + """ + pprint.pprint(self.as_list(), *args, **kwargs) + + # add support for pickle protocol + def __getstate__(self): + return ( + self._toklist, + ( + self._tokdict.copy(), + None, + self._all_names, + self._name, + ), + ) + + def __setstate__(self, state): + self._toklist, (self._tokdict, par, inAccumNames, self._name) = state + self._all_names = set(inAccumNames) + self._parent = None + + def __getnewargs__(self): + return self._toklist, self._name + + def __dir__(self): + return dir(type(self)) + list(self.keys()) + + @classmethod + def from_dict(cls, other, name=None) -> ParseResults: + """ + Helper classmethod to construct a ``ParseResults`` from a ``dict``, preserving the + name-value relations as results names. If an optional ``name`` argument is + given, a nested ``ParseResults`` will be returned. + """ + + def is_iterable(obj): + try: + iter(obj) + except Exception: + return False + # str's are iterable, but in pyparsing, we don't want to iterate over them + else: + return not isinstance(obj, str_type) + + ret = cls([]) + for k, v in other.items(): + if isinstance(v, Mapping): + ret += cls.from_dict(v, name=k) + else: + ret += cls([v], name=k, asList=is_iterable(v)) + if name is not None: + ret = cls([ret], name=name) + return ret + + asList = as_list + """Deprecated - use :class:`as_list`""" + asDict = as_dict + """Deprecated - use :class:`as_dict`""" + getName = get_name + """Deprecated - use :class:`get_name`""" + + +MutableMapping.register(ParseResults) +MutableSequence.register(ParseResults) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/testing.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..836b2f86fbeb0ad7c12f36996f4099bcf96e492a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/testing.py @@ -0,0 +1,362 @@ +# testing.py + +from contextlib import contextmanager +import re +import typing + + +from .core import ( + ParserElement, + ParseException, + Keyword, + __diag__, + __compat__, +) + + +class pyparsing_test: + """ + namespace class for classes useful in writing unit tests + """ + + class reset_pyparsing_context: + """ + Context manager to be used when writing unit tests that modify pyparsing config values: + - packrat parsing + - bounded recursion parsing + - default whitespace characters. + - default keyword characters + - literal string auto-conversion class + - __diag__ settings + + Example:: + + with reset_pyparsing_context(): + # test that literals used to construct a grammar are automatically suppressed + ParserElement.inlineLiteralsUsing(Suppress) + + term = Word(alphas) | Word(nums) + group = Group('(' + term[...] + ')') + + # assert that the '()' characters are not included in the parsed tokens + self.assertParseAndCheckList(group, "(abc 123 def)", ['abc', '123', 'def']) + + # after exiting context manager, literals are converted to Literal expressions again + """ + + def __init__(self): + self._save_context = {} + + def save(self): + self._save_context["default_whitespace"] = ParserElement.DEFAULT_WHITE_CHARS + self._save_context["default_keyword_chars"] = Keyword.DEFAULT_KEYWORD_CHARS + + self._save_context["literal_string_class"] = ( + ParserElement._literalStringClass + ) + + self._save_context["verbose_stacktrace"] = ParserElement.verbose_stacktrace + + self._save_context["packrat_enabled"] = ParserElement._packratEnabled + if ParserElement._packratEnabled: + self._save_context["packrat_cache_size"] = ( + ParserElement.packrat_cache.size + ) + else: + self._save_context["packrat_cache_size"] = None + self._save_context["packrat_parse"] = ParserElement._parse + self._save_context["recursion_enabled"] = ( + ParserElement._left_recursion_enabled + ) + + self._save_context["__diag__"] = { + name: getattr(__diag__, name) for name in __diag__._all_names + } + + self._save_context["__compat__"] = { + "collect_all_And_tokens": __compat__.collect_all_And_tokens + } + + return self + + def restore(self): + # reset pyparsing global state + if ( + ParserElement.DEFAULT_WHITE_CHARS + != self._save_context["default_whitespace"] + ): + ParserElement.set_default_whitespace_chars( + self._save_context["default_whitespace"] + ) + + ParserElement.verbose_stacktrace = self._save_context["verbose_stacktrace"] + + Keyword.DEFAULT_KEYWORD_CHARS = self._save_context["default_keyword_chars"] + ParserElement.inlineLiteralsUsing( + self._save_context["literal_string_class"] + ) + + for name, value in self._save_context["__diag__"].items(): + (__diag__.enable if value else __diag__.disable)(name) + + ParserElement._packratEnabled = False + if self._save_context["packrat_enabled"]: + ParserElement.enable_packrat(self._save_context["packrat_cache_size"]) + else: + ParserElement._parse = self._save_context["packrat_parse"] + ParserElement._left_recursion_enabled = self._save_context[ + "recursion_enabled" + ] + + __compat__.collect_all_And_tokens = self._save_context["__compat__"] + + return self + + def copy(self): + ret = type(self)() + ret._save_context.update(self._save_context) + return ret + + def __enter__(self): + return self.save() + + def __exit__(self, *args): + self.restore() + + class TestParseResultsAsserts: + """ + A mixin class to add parse results assertion methods to normal unittest.TestCase classes. + """ + + def assertParseResultsEquals( + self, result, expected_list=None, expected_dict=None, msg=None + ): + """ + Unit test assertion to compare a :class:`ParseResults` object with an optional ``expected_list``, + and compare any defined results names with an optional ``expected_dict``. + """ + if expected_list is not None: + self.assertEqual(expected_list, result.as_list(), msg=msg) + if expected_dict is not None: + self.assertEqual(expected_dict, result.as_dict(), msg=msg) + + def assertParseAndCheckList( + self, expr, test_string, expected_list, msg=None, verbose=True + ): + """ + Convenience wrapper assert to test a parser element and input string, and assert that + the resulting ``ParseResults.asList()`` is equal to the ``expected_list``. + """ + result = expr.parse_string(test_string, parse_all=True) + if verbose: + print(result.dump()) + else: + print(result.as_list()) + self.assertParseResultsEquals(result, expected_list=expected_list, msg=msg) + + def assertParseAndCheckDict( + self, expr, test_string, expected_dict, msg=None, verbose=True + ): + """ + Convenience wrapper assert to test a parser element and input string, and assert that + the resulting ``ParseResults.asDict()`` is equal to the ``expected_dict``. + """ + result = expr.parse_string(test_string, parseAll=True) + if verbose: + print(result.dump()) + else: + print(result.as_list()) + self.assertParseResultsEquals(result, expected_dict=expected_dict, msg=msg) + + def assertRunTestResults( + self, run_tests_report, expected_parse_results=None, msg=None + ): + """ + Unit test assertion to evaluate output of ``ParserElement.runTests()``. If a list of + list-dict tuples is given as the ``expected_parse_results`` argument, then these are zipped + with the report tuples returned by ``runTests`` and evaluated using ``assertParseResultsEquals``. + Finally, asserts that the overall ``runTests()`` success value is ``True``. + + :param run_tests_report: tuple(bool, [tuple(str, ParseResults or Exception)]) returned from runTests + :param expected_parse_results (optional): [tuple(str, list, dict, Exception)] + """ + run_test_success, run_test_results = run_tests_report + + if expected_parse_results is None: + self.assertTrue( + run_test_success, msg=msg if msg is not None else "failed runTests" + ) + return + + merged = [ + (*rpt, expected) + for rpt, expected in zip(run_test_results, expected_parse_results) + ] + for test_string, result, expected in merged: + # expected should be a tuple containing a list and/or a dict or an exception, + # and optional failure message string + # an empty tuple will skip any result validation + fail_msg = next((exp for exp in expected if isinstance(exp, str)), None) + expected_exception = next( + ( + exp + for exp in expected + if isinstance(exp, type) and issubclass(exp, Exception) + ), + None, + ) + if expected_exception is not None: + with self.assertRaises( + expected_exception=expected_exception, msg=fail_msg or msg + ): + if isinstance(result, Exception): + raise result + else: + expected_list = next( + (exp for exp in expected if isinstance(exp, list)), None + ) + expected_dict = next( + (exp for exp in expected if isinstance(exp, dict)), None + ) + if (expected_list, expected_dict) != (None, None): + self.assertParseResultsEquals( + result, + expected_list=expected_list, + expected_dict=expected_dict, + msg=fail_msg or msg, + ) + else: + # warning here maybe? + print(f"no validation for {test_string!r}") + + # do this last, in case some specific test results can be reported instead + self.assertTrue( + run_test_success, msg=msg if msg is not None else "failed runTests" + ) + + @contextmanager + def assertRaisesParseException( + self, exc_type=ParseException, expected_msg=None, msg=None + ): + if expected_msg is not None: + if isinstance(expected_msg, str): + expected_msg = re.escape(expected_msg) + with self.assertRaisesRegex(exc_type, expected_msg, msg=msg) as ctx: + yield ctx + + else: + with self.assertRaises(exc_type, msg=msg) as ctx: + yield ctx + + @staticmethod + def with_line_numbers( + s: str, + start_line: typing.Optional[int] = None, + end_line: typing.Optional[int] = None, + expand_tabs: bool = True, + eol_mark: str = "|", + mark_spaces: typing.Optional[str] = None, + mark_control: typing.Optional[str] = None, + *, + indent: typing.Union[str, int] = "", + base_1: bool = True, + ) -> str: + """ + Helpful method for debugging a parser - prints a string with line and column numbers. + (Line and column numbers are 1-based by default - if debugging a parse action, + pass base_1=False, to correspond to the loc value passed to the parse action.) + + :param s: tuple(bool, str - string to be printed with line and column numbers + :param start_line: int - (optional) starting line number in s to print (default=1) + :param end_line: int - (optional) ending line number in s to print (default=len(s)) + :param expand_tabs: bool - (optional) expand tabs to spaces, to match the pyparsing default + :param eol_mark: str - (optional) string to mark the end of lines, helps visualize trailing spaces (default="|") + :param mark_spaces: str - (optional) special character to display in place of spaces + :param mark_control: str - (optional) convert non-printing control characters to a placeholding + character; valid values: + - "unicode" - replaces control chars with Unicode symbols, such as "␍" and "␊" + - any single character string - replace control characters with given string + - None (default) - string is displayed as-is + :param indent: str | int - (optional) string to indent with line and column numbers; if an int + is passed, converted to " " * indent + :param base_1: bool - (optional) whether to label string using base 1; if False, string will be + labeled based at 0 (default=True) + + :return: str - input string with leading line numbers and column number headers + """ + if expand_tabs: + s = s.expandtabs() + if isinstance(indent, int): + indent = " " * indent + indent = indent.expandtabs() + if mark_control is not None: + mark_control = typing.cast(str, mark_control) + if mark_control == "unicode": + transtable_map = { + c: u for c, u in zip(range(0, 33), range(0x2400, 0x2433)) + } + transtable_map[127] = 0x2421 + tbl = str.maketrans(transtable_map) + eol_mark = "" + else: + ord_mark_control = ord(mark_control) + tbl = str.maketrans( + {c: ord_mark_control for c in list(range(0, 32)) + [127]} + ) + s = s.translate(tbl) + if mark_spaces is not None and mark_spaces != " ": + if mark_spaces == "unicode": + tbl = str.maketrans({9: 0x2409, 32: 0x2423}) + s = s.translate(tbl) + else: + s = s.replace(" ", mark_spaces) + if start_line is None: + start_line = 0 + if end_line is None: + end_line = len(s) + end_line = min(end_line, len(s)) + start_line = min(max(0, start_line), end_line) + + if mark_control != "unicode": + s_lines = s.splitlines()[start_line - base_1 : end_line] + else: + s_lines = [ + line + "␊" for line in s.split("␊")[start_line - base_1 : end_line] + ] + if not s_lines: + return "" + + lineno_width = len(str(end_line)) + max_line_len = max(len(line) for line in s_lines) + lead = indent + " " * (lineno_width + 1) + if max_line_len >= 99: + header0 = ( + lead + + ("" if base_1 else " ") + + "".join( + f"{' ' * 99}{(i + 1) % 100}" + for i in range(1 if base_1 else 0, max(max_line_len // 100, 1)) + ) + + "\n" + ) + else: + header0 = "" + header1 = ( + ("" if base_1 else " ") + + lead + + "".join(f" {(i + 1) % 10}" for i in range(-(-max_line_len // 10))) + + "\n" + ) + digits = "1234567890" + header2 = ( + lead + ("" if base_1 else "0") + digits * (-(-max_line_len // 10)) + "\n" + ) + return ( + header1 + + header2 + + "\n".join( + f"{indent}{i:{lineno_width}d}:{line}{eol_mark}" + for i, line in enumerate(s_lines, start=start_line + base_1) + ) + + "\n" + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/tools/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/tools/cvt_pyparsing_pep8_names.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/tools/cvt_pyparsing_pep8_names.py new file mode 100644 index 0000000000000000000000000000000000000000..f4a8bd9f516994ce6b9734ca07c5895bb8276fd4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/tools/cvt_pyparsing_pep8_names.py @@ -0,0 +1,116 @@ +from functools import lru_cache +import pyparsing as pp + + +@lru_cache(maxsize=None) +def camel_to_snake(s: str) -> str: + """ + Convert CamelCase to snake_case. + """ + return "".join("_" + c.lower() if c.isupper() else c for c in s).lstrip("_") + + +pre_pep8_method_names = """ +addCondition addParseAction anyCloseTag anyOpenTag asDict asList cStyleComment canParseNext conditionAsParseAction +convertToDate convertToDatetime convertToFloat convertToInteger countedArray cppStyleComment dblQuotedString +dblSlashComment defaultName dictOf disableMemoization downcaseTokens enableLeftRecursion enablePackrat getName +htmlComment ignoreWhitespace indentedBlock infixNotation inlineLiteralsUsing javaStyleComment leaveWhitespace +lineEnd lineStart locatedExpr matchOnlyAtCol matchPreviousExpr matchPreviousLiteral nestedExpr nullDebugAction oneOf +originalTextFor parseFile parseString parseWithTabs pythonStyleComment quotedString removeQuotes replaceWith +resetCache restOfLine runTests scanString searchString setBreak setDebug setDebugActions setDefaultWhitespaceChars +setFailAction setName setParseAction setResultsName setWhitespaceChars sglQuotedString stringEnd stringStart tokenMap +traceParseAction transformString tryParse unicodeString upcaseTokens withAttribute withClass +""".split() + +special_changes = { + "opAssoc": "OpAssoc", + "delimitedList": "DelimitedList", + "delimited_list": "DelimitedList", + "replaceHTMLEntity": "replace_html_entity", + "makeHTMLTags": "make_html_tags", + "makeXMLTags": "make_xml_tags", + "commonHTMLEntity": "common_html_entity", + "stripHTMLTags": "strip_html_tags", +} + +pre_pep8_arg_names = """parseAll maxMatches listAllMatches callDuringTry includeSeparators fullDump printResults +failureTests postParse matchString identChars maxMismatches initChars bodyChars asKeyword excludeChars asGroupList +asMatch quoteChar escChar escQuote unquoteResults endQuoteChar convertWhitespaceEscapes notChars wordChars stopOn +failOn joinString markerString intExpr useRegex asString ignoreExpr""".split() + +pre_pep8_method_name = pp.one_of(pre_pep8_method_names, as_keyword=True) +pre_pep8_method_name.set_parse_action(lambda t: camel_to_snake(t[0])) +special_pre_pep8_name = pp.one_of(special_changes, as_keyword=True) +special_pre_pep8_name.set_parse_action(lambda t: special_changes[t[0]]) +# only replace arg names if part of an arg list +pre_pep8_arg_name = pp.Regex( + rf"{pp.util.make_compressed_re(pre_pep8_arg_names)}\s*=" +) +pre_pep8_arg_name.set_parse_action(lambda t: camel_to_snake(t[0])) + +pep8_converter = pre_pep8_method_name | special_pre_pep8_name | pre_pep8_arg_name + +if __name__ == "__main__": + import argparse + from pathlib import Path + import sys + + argparser = argparse.ArgumentParser( + description = ( + "Utility to convert Python pyparsing scripts using legacy" + " camelCase names to use PEP8 snake_case names." + "\nBy default, this script will only show whether this script would make any changes." + ) + ) + argparser.add_argument("--verbose", "-v", action="store_true", help="Show unified diff for each source file") + argparser.add_argument("-vv", action="store_true", dest="verbose2", help="Show unified diff for each source file, plus names of scanned files with no changes") + argparser.add_argument("--update", "-u", action="store_true", help="Update source files in-place") + argparser.add_argument("--encoding", type=str, default="utf-8", help="Encoding of source files (default: utf-8)") + argparser.add_argument("--exit-zero-even-if-changed", "-exit0", action="store_true", help="Exit with status code 0 even if changes were made") + argparser.add_argument("source_filename", nargs="+", help="Source filenames or filename patterns of Python files to be converted") + args = argparser.parse_args() + + + def show_diffs(original, modified): + import difflib + + diff = difflib.unified_diff( + original.splitlines(), modified.splitlines(), lineterm="" + ) + sys.stdout.writelines(f"{diff_line}\n" for diff_line in diff) + + exit_status = 0 + + for filename_pattern in args.source_filename: + + for filename in Path().glob(filename_pattern): + if not Path(filename).is_file(): + continue + + try: + original_contents = Path(filename).read_text(encoding=args.encoding) + modified_contents = pep8_converter.transform_string( + original_contents + ) + + if modified_contents != original_contents: + if args.update: + Path(filename).write_text(modified_contents, encoding=args.encoding) + print(f"Converted {filename}") + else: + print(f"Found required changes in {filename}") + + if args.verbose: + show_diffs(original_contents, modified_contents) + print() + + exit_status = 1 + + else: + if args.verbose2: + print(f"No required changes in {filename}") + + except Exception as e: + print(f"Failed to convert {filename}: {type(e).__name__}: {e}") + + sys.exit(exit_status if not args.exit_zero_even_if_changed else 0) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/unicode.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/unicode.py new file mode 100644 index 0000000000000000000000000000000000000000..066486c28eea020d420c2e90fdda76f69f1c9ead --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/unicode.py @@ -0,0 +1,356 @@ +# unicode.py + +import sys +from itertools import filterfalse +from typing import Union + + +class _lazyclassproperty: + def __init__(self, fn): + self.fn = fn + self.__doc__ = fn.__doc__ + self.__name__ = fn.__name__ + + def __get__(self, obj, cls): + if cls is None: + cls = type(obj) + if not hasattr(cls, "_intern") or any( + cls._intern is getattr(superclass, "_intern", []) + for superclass in cls.__mro__[1:] + ): + cls._intern = {} + attrname = self.fn.__name__ + if attrname not in cls._intern: + cls._intern[attrname] = self.fn(cls) + return cls._intern[attrname] + + +UnicodeRangeList = list[Union[tuple[int, int], tuple[int]]] + + +class unicode_set: + """ + A set of Unicode characters, for language-specific strings for + ``alphas``, ``nums``, ``alphanums``, and ``printables``. + A unicode_set is defined by a list of ranges in the Unicode character + set, in a class attribute ``_ranges``. Ranges can be specified using + 2-tuples or a 1-tuple, such as:: + + _ranges = [ + (0x0020, 0x007e), + (0x00a0, 0x00ff), + (0x0100,), + ] + + Ranges are left- and right-inclusive. A 1-tuple of (x,) is treated as (x, x). + + A unicode set can also be defined using multiple inheritance of other unicode sets:: + + class CJK(Chinese, Japanese, Korean): + pass + """ + + _ranges: UnicodeRangeList = [] + + @_lazyclassproperty + def _chars_for_ranges(cls) -> list[str]: + ret: list[int] = [] + for cc in cls.__mro__: # type: ignore[attr-defined] + if cc is unicode_set: + break + for rr in getattr(cc, "_ranges", ()): + ret.extend(range(rr[0], rr[-1] + 1)) + return sorted(chr(c) for c in set(ret)) + + @_lazyclassproperty + def printables(cls) -> str: + """all non-whitespace characters in this range""" + return "".join(filterfalse(str.isspace, cls._chars_for_ranges)) + + @_lazyclassproperty + def alphas(cls) -> str: + """all alphabetic characters in this range""" + return "".join(filter(str.isalpha, cls._chars_for_ranges)) + + @_lazyclassproperty + def nums(cls) -> str: + """all numeric digit characters in this range""" + return "".join(filter(str.isdigit, cls._chars_for_ranges)) + + @_lazyclassproperty + def alphanums(cls) -> str: + """all alphanumeric characters in this range""" + return cls.alphas + cls.nums + + @_lazyclassproperty + def identchars(cls) -> str: + """all characters in this range that are valid identifier characters, plus underscore '_'""" + return "".join( + sorted( + set(filter(str.isidentifier, cls._chars_for_ranges)) + | set( + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyzªµº" + "ÀÁÂÃÄÅÆÇÈÉÊËÌÍÎÏÐÑÒÓÔÕÖØÙÚÛÜÝÞßàáâãäåæçèéêëìíîïðñòóôõöøùúûüýþÿ" + "_" + ) + ) + ) + + @_lazyclassproperty + def identbodychars(cls) -> str: + """ + all characters in this range that are valid identifier body characters, + plus the digits 0-9, and · (Unicode MIDDLE DOT) + """ + identifier_chars = set( + c for c in cls._chars_for_ranges if ("_" + c).isidentifier() + ) + return "".join( + sorted(identifier_chars | set(cls.identchars) | set("0123456789·")) + ) + + @_lazyclassproperty + def identifier(cls): + """ + a pyparsing Word expression for an identifier using this range's definitions for + identchars and identbodychars + """ + from pyparsing import Word + + return Word(cls.identchars, cls.identbodychars) + + +class pyparsing_unicode(unicode_set): + """ + A namespace class for defining common language unicode_sets. + """ + + # fmt: off + + # define ranges in language character sets + _ranges: UnicodeRangeList = [ + (0x0020, sys.maxunicode), + ] + + class BasicMultilingualPlane(unicode_set): + """Unicode set for the Basic Multilingual Plane""" + _ranges: UnicodeRangeList = [ + (0x0020, 0xFFFF), + ] + + class Latin1(unicode_set): + """Unicode set for Latin-1 Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0020, 0x007E), + (0x00A0, 0x00FF), + ] + + class LatinA(unicode_set): + """Unicode set for Latin-A Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0100, 0x017F), + ] + + class LatinB(unicode_set): + """Unicode set for Latin-B Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0180, 0x024F), + ] + + class Greek(unicode_set): + """Unicode set for Greek Unicode Character Ranges""" + _ranges: UnicodeRangeList = [ + (0x0342, 0x0345), + (0x0370, 0x0377), + (0x037A, 0x037F), + (0x0384, 0x038A), + (0x038C,), + (0x038E, 0x03A1), + (0x03A3, 0x03E1), + (0x03F0, 0x03FF), + (0x1D26, 0x1D2A), + (0x1D5E,), + (0x1D60,), + (0x1D66, 0x1D6A), + (0x1F00, 0x1F15), + (0x1F18, 0x1F1D), + (0x1F20, 0x1F45), + (0x1F48, 0x1F4D), + (0x1F50, 0x1F57), + (0x1F59,), + (0x1F5B,), + (0x1F5D,), + (0x1F5F, 0x1F7D), + (0x1F80, 0x1FB4), + (0x1FB6, 0x1FC4), + (0x1FC6, 0x1FD3), + (0x1FD6, 0x1FDB), + (0x1FDD, 0x1FEF), + (0x1FF2, 0x1FF4), + (0x1FF6, 0x1FFE), + (0x2129,), + (0x2719, 0x271A), + (0xAB65,), + (0x10140, 0x1018D), + (0x101A0,), + (0x1D200, 0x1D245), + (0x1F7A1, 0x1F7A7), + ] + + class Cyrillic(unicode_set): + """Unicode set for Cyrillic Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0400, 0x052F), + (0x1C80, 0x1C88), + (0x1D2B,), + (0x1D78,), + (0x2DE0, 0x2DFF), + (0xA640, 0xA672), + (0xA674, 0xA69F), + (0xFE2E, 0xFE2F), + ] + + class Chinese(unicode_set): + """Unicode set for Chinese Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x2E80, 0x2E99), + (0x2E9B, 0x2EF3), + (0x31C0, 0x31E3), + (0x3400, 0x4DB5), + (0x4E00, 0x9FEF), + (0xA700, 0xA707), + (0xF900, 0xFA6D), + (0xFA70, 0xFAD9), + (0x16FE2, 0x16FE3), + (0x1F210, 0x1F212), + (0x1F214, 0x1F23B), + (0x1F240, 0x1F248), + (0x20000, 0x2A6D6), + (0x2A700, 0x2B734), + (0x2B740, 0x2B81D), + (0x2B820, 0x2CEA1), + (0x2CEB0, 0x2EBE0), + (0x2F800, 0x2FA1D), + ] + + class Japanese(unicode_set): + """Unicode set for Japanese Unicode Character Range, combining Kanji, Hiragana, and Katakana ranges""" + + class Kanji(unicode_set): + "Unicode set for Kanji Unicode Character Range" + _ranges: UnicodeRangeList = [ + (0x4E00, 0x9FBF), + (0x3000, 0x303F), + ] + + class Hiragana(unicode_set): + """Unicode set for Hiragana Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x3041, 0x3096), + (0x3099, 0x30A0), + (0x30FC,), + (0xFF70,), + (0x1B001,), + (0x1B150, 0x1B152), + (0x1F200,), + ] + + class Katakana(unicode_set): + """Unicode set for Katakana Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x3099, 0x309C), + (0x30A0, 0x30FF), + (0x31F0, 0x31FF), + (0x32D0, 0x32FE), + (0xFF65, 0xFF9F), + (0x1B000,), + (0x1B164, 0x1B167), + (0x1F201, 0x1F202), + (0x1F213,), + ] + + 漢字 = Kanji + カタカナ = Katakana + ひらがな = Hiragana + + _ranges = ( + Kanji._ranges + + Hiragana._ranges + + Katakana._ranges + ) + + class Hangul(unicode_set): + """Unicode set for Hangul (Korean) Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x1100, 0x11FF), + (0x302E, 0x302F), + (0x3131, 0x318E), + (0x3200, 0x321C), + (0x3260, 0x327B), + (0x327E,), + (0xA960, 0xA97C), + (0xAC00, 0xD7A3), + (0xD7B0, 0xD7C6), + (0xD7CB, 0xD7FB), + (0xFFA0, 0xFFBE), + (0xFFC2, 0xFFC7), + (0xFFCA, 0xFFCF), + (0xFFD2, 0xFFD7), + (0xFFDA, 0xFFDC), + ] + + Korean = Hangul + + class CJK(Chinese, Japanese, Hangul): + """Unicode set for combined Chinese, Japanese, and Korean (CJK) Unicode Character Range""" + + class Thai(unicode_set): + """Unicode set for Thai Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0E01, 0x0E3A), + (0x0E3F, 0x0E5B) + ] + + class Arabic(unicode_set): + """Unicode set for Arabic Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0600, 0x061B), + (0x061E, 0x06FF), + (0x0700, 0x077F), + ] + + class Hebrew(unicode_set): + """Unicode set for Hebrew Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0591, 0x05C7), + (0x05D0, 0x05EA), + (0x05EF, 0x05F4), + (0xFB1D, 0xFB36), + (0xFB38, 0xFB3C), + (0xFB3E,), + (0xFB40, 0xFB41), + (0xFB43, 0xFB44), + (0xFB46, 0xFB4F), + ] + + class Devanagari(unicode_set): + """Unicode set for Devanagari Unicode Character Range""" + _ranges: UnicodeRangeList = [ + (0x0900, 0x097F), + (0xA8E0, 0xA8FF) + ] + + BMP = BasicMultilingualPlane + + # add language identifiers using language Unicode + العربية = Arabic + 中文 = Chinese + кириллица = Cyrillic + Ελληνικά = Greek + עִברִית = Hebrew + 日本語 = Japanese + 한국어 = Korean + ไทย = Thai + देवनागरी = Devanagari + + # fmt: on diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/util.py b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/util.py new file mode 100644 index 0000000000000000000000000000000000000000..1cb16e2e6205bc650dd006a8517a9c155d926c01 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/pyparsing/util.py @@ -0,0 +1,438 @@ +# util.py +import contextlib +import re +from functools import lru_cache, wraps +import inspect +import itertools +import types +from typing import Callable, Union, Iterable, TypeVar, cast +import warnings + +_bslash = chr(92) +C = TypeVar("C", bound=Callable) + + +class __config_flags: + """Internal class for defining compatibility and debugging flags""" + + _all_names: list[str] = [] + _fixed_names: list[str] = [] + _type_desc = "configuration" + + @classmethod + def _set(cls, dname, value): + if dname in cls._fixed_names: + warnings.warn( + f"{cls.__name__}.{dname} {cls._type_desc} is {str(getattr(cls, dname)).upper()}" + f" and cannot be overridden", + stacklevel=3, + ) + return + if dname in cls._all_names: + setattr(cls, dname, value) + else: + raise ValueError(f"no such {cls._type_desc} {dname!r}") + + enable = classmethod(lambda cls, name: cls._set(name, True)) + disable = classmethod(lambda cls, name: cls._set(name, False)) + + +@lru_cache(maxsize=128) +def col(loc: int, strg: str) -> int: + """ + Returns current column within a string, counting newlines as line separators. + The first column is number 1. + + Note: the default parsing behavior is to expand tabs in the input string + before starting the parsing process. See + :class:`ParserElement.parse_string` for more + information on parsing strings containing ```` s, and suggested + methods to maintain a consistent view of the parsed string, the parse + location, and line and column positions within the parsed string. + """ + s = strg + return 1 if 0 < loc < len(s) and s[loc - 1] == "\n" else loc - s.rfind("\n", 0, loc) + + +@lru_cache(maxsize=128) +def lineno(loc: int, strg: str) -> int: + """Returns current line number within a string, counting newlines as line separators. + The first line is number 1. + + Note - the default parsing behavior is to expand tabs in the input string + before starting the parsing process. See :class:`ParserElement.parse_string` + for more information on parsing strings containing ```` s, and + suggested methods to maintain a consistent view of the parsed string, the + parse location, and line and column positions within the parsed string. + """ + return strg.count("\n", 0, loc) + 1 + + +@lru_cache(maxsize=128) +def line(loc: int, strg: str) -> str: + """ + Returns the line of text containing loc within a string, counting newlines as line separators. + """ + last_cr = strg.rfind("\n", 0, loc) + next_cr = strg.find("\n", loc) + return strg[last_cr + 1 : next_cr] if next_cr >= 0 else strg[last_cr + 1 :] + + +class _UnboundedCache: + def __init__(self): + cache = {} + cache_get = cache.get + self.not_in_cache = not_in_cache = object() + + def get(_, key): + return cache_get(key, not_in_cache) + + def set_(_, key, value): + cache[key] = value + + def clear(_): + cache.clear() + + self.size = None + self.get = types.MethodType(get, self) + self.set = types.MethodType(set_, self) + self.clear = types.MethodType(clear, self) + + +class _FifoCache: + def __init__(self, size): + cache = {} + self.size = size + self.not_in_cache = not_in_cache = object() + cache_get = cache.get + cache_pop = cache.pop + + def get(_, key): + return cache_get(key, not_in_cache) + + def set_(_, key, value): + cache[key] = value + while len(cache) > size: + # pop oldest element in cache by getting the first key + cache_pop(next(iter(cache))) + + def clear(_): + cache.clear() + + self.get = types.MethodType(get, self) + self.set = types.MethodType(set_, self) + self.clear = types.MethodType(clear, self) + + +class LRUMemo: + """ + A memoizing mapping that retains `capacity` deleted items + + The memo tracks retained items by their access order; once `capacity` items + are retained, the least recently used item is discarded. + """ + + def __init__(self, capacity): + self._capacity = capacity + self._active = {} + self._memory = {} + + def __getitem__(self, key): + try: + return self._active[key] + except KeyError: + self._memory[key] = self._memory.pop(key) + return self._memory[key] + + def __setitem__(self, key, value): + self._memory.pop(key, None) + self._active[key] = value + + def __delitem__(self, key): + try: + value = self._active.pop(key) + except KeyError: + pass + else: + oldest_keys = list(self._memory)[: -(self._capacity + 1)] + for key_to_delete in oldest_keys: + self._memory.pop(key_to_delete) + self._memory[key] = value + + def clear(self): + self._active.clear() + self._memory.clear() + + +class UnboundedMemo(dict): + """ + A memoizing mapping that retains all deleted items + """ + + def __delitem__(self, key): + pass + + +def _escape_regex_range_chars(s: str) -> str: + # escape these chars: ^-[] + for c in r"\^-[]": + s = s.replace(c, _bslash + c) + s = s.replace("\n", r"\n") + s = s.replace("\t", r"\t") + return str(s) + + +class _GroupConsecutive: + """ + Used as a callable `key` for itertools.groupby to group + characters that are consecutive: + itertools.groupby("abcdejkmpqrs", key=IsConsecutive()) + yields: + (0, iter(['a', 'b', 'c', 'd', 'e'])) + (1, iter(['j', 'k'])) + (2, iter(['m'])) + (3, iter(['p', 'q', 'r', 's'])) + """ + + def __init__(self) -> None: + self.prev = 0 + self.counter = itertools.count() + self.value = -1 + + def __call__(self, char: str) -> int: + c_int = ord(char) + self.prev, prev = c_int, self.prev + if c_int - prev > 1: + self.value = next(self.counter) + return self.value + + +def _collapse_string_to_ranges( + s: Union[str, Iterable[str]], re_escape: bool = True +) -> str: + r""" + Take a string or list of single-character strings, and return + a string of the consecutive characters in that string collapsed + into groups, as might be used in a regular expression '[a-z]' + character set: + 'a' -> 'a' -> '[a]' + 'bc' -> 'bc' -> '[bc]' + 'defgh' -> 'd-h' -> '[d-h]' + 'fdgeh' -> 'd-h' -> '[d-h]' + 'jklnpqrtu' -> 'j-lnp-rtu' -> '[j-lnp-rtu]' + Duplicates get collapsed out: + 'aaa' -> 'a' -> '[a]' + 'bcbccb' -> 'bc' -> '[bc]' + 'defghhgf' -> 'd-h' -> '[d-h]' + 'jklnpqrjjjtu' -> 'j-lnp-rtu' -> '[j-lnp-rtu]' + Spaces are preserved: + 'ab c' -> ' a-c' -> '[ a-c]' + Characters that are significant when defining regex ranges + get escaped: + 'acde[]-' -> r'\-\[\]ac-e' -> r'[\-\[\]ac-e]' + """ + + # Developer notes: + # - Do not optimize this code assuming that the given input string + # or internal lists will be short (such as in loading generators into + # lists to make it easier to find the last element); this method is also + # used to generate regex ranges for character sets in the pyparsing.unicode + # classes, and these can be _very_ long lists of strings + + def escape_re_range_char(c: str) -> str: + return "\\" + c if c in r"\^-][" else c + + def no_escape_re_range_char(c: str) -> str: + return c + + if not re_escape: + escape_re_range_char = no_escape_re_range_char + + ret = [] + + # reduce input string to remove duplicates, and put in sorted order + s_chars: list[str] = sorted(set(s)) + + if len(s_chars) > 2: + # find groups of characters that are consecutive (can be collapsed + # down to "-") + for _, chars in itertools.groupby(s_chars, key=_GroupConsecutive()): + # _ is unimportant, is just used to identify groups + # chars is an iterator of one or more consecutive characters + # that comprise the current group + first = last = next(chars) + with contextlib.suppress(ValueError): + *_, last = chars + + if first == last: + # there was only a single char in this group + ret.append(escape_re_range_char(first)) + + elif last == chr(ord(first) + 1): + # there were only 2 characters in this group + # 'a','b' -> 'ab' + ret.append(f"{escape_re_range_char(first)}{escape_re_range_char(last)}") + + else: + # there were > 2 characters in this group, make into a range + # 'c','d','e' -> 'c-e' + ret.append( + f"{escape_re_range_char(first)}-{escape_re_range_char(last)}" + ) + else: + # only 1 or 2 chars were given to form into groups + # 'a' -> ['a'] + # 'bc' -> ['b', 'c'] + # 'dg' -> ['d', 'g'] + # no need to list them with "-", just return as a list + # (after escaping) + ret = [escape_re_range_char(c) for c in s_chars] + + return "".join(ret) + + +def _flatten(ll: Iterable) -> list: + ret = [] + to_visit = [*ll] + while to_visit: + i = to_visit.pop(0) + if isinstance(i, Iterable) and not isinstance(i, str): + to_visit[:0] = i + else: + ret.append(i) + return ret + + +def make_compressed_re( + word_list: Iterable[str], + max_level: int = 2, + *, + non_capturing_groups: bool = True, + _level: int = 1, +) -> str: + """ + Create a regular expression string from a list of words, collapsing by common + prefixes and optional suffixes. + + Calls itself recursively to build nested sublists for each group of suffixes + that have a shared prefix. + """ + + def get_suffixes_from_common_prefixes(namelist: list[str]): + if len(namelist) > 1: + for prefix, suffixes in itertools.groupby(namelist, key=lambda s: s[:1]): + yield prefix, sorted([s[1:] for s in suffixes], key=len, reverse=True) + else: + yield namelist[0][0], [namelist[0][1:]] + + if _level == 1: + if not word_list: + raise ValueError("no words given to make_compressed_re()") + + if "" in word_list: + raise ValueError("word list cannot contain empty string") + else: + # internal recursive call, just return empty string if no words + if not word_list: + return "" + + # dedupe the word list + word_list = list({}.fromkeys(word_list)) + + if max_level == 0: + if any(len(wd) > 1 for wd in word_list): + return "|".join( + sorted([re.escape(wd) for wd in word_list], key=len, reverse=True) + ) + else: + return f"[{''.join(_escape_regex_range_chars(wd) for wd in word_list)}]" + + ret = [] + sep = "" + ncgroup = "?:" if non_capturing_groups else "" + + for initial, suffixes in get_suffixes_from_common_prefixes(sorted(word_list)): + ret.append(sep) + sep = "|" + + initial = re.escape(initial) + + trailing = "" + if "" in suffixes: + trailing = "?" + suffixes.remove("") + + if len(suffixes) > 1: + if all(len(s) == 1 for s in suffixes): + ret.append( + f"{initial}[{''.join(_escape_regex_range_chars(s) for s in suffixes)}]{trailing}" + ) + else: + if _level < max_level: + suffix_re = make_compressed_re( + sorted(suffixes), + max_level, + non_capturing_groups=non_capturing_groups, + _level=_level + 1, + ) + ret.append(f"{initial}({ncgroup}{suffix_re}){trailing}") + else: + if all(len(s) == 1 for s in suffixes): + ret.append( + f"{initial}[{''.join(_escape_regex_range_chars(s) for s in suffixes)}]{trailing}" + ) + else: + suffixes.sort(key=len, reverse=True) + ret.append( + f"{initial}({ncgroup}{'|'.join(re.escape(s) for s in suffixes)}){trailing}" + ) + else: + if suffixes: + suffix = re.escape(suffixes[0]) + if len(suffix) > 1 and trailing: + ret.append(f"{initial}({ncgroup}{suffix}){trailing}") + else: + ret.append(f"{initial}{suffix}{trailing}") + else: + ret.append(initial) + return "".join(ret) + + +def replaced_by_pep8(compat_name: str, fn: C) -> C: + # In a future version, uncomment the code in the internal _inner() functions + # to begin emitting DeprecationWarnings. + + # Unwrap staticmethod/classmethod + fn = getattr(fn, "__func__", fn) + + # (Presence of 'self' arg in signature is used by explain_exception() methods, so we take + # some extra steps to add it if present in decorated function.) + if ["self"] == list(inspect.signature(fn).parameters)[:1]: + + @wraps(fn) + def _inner(self, *args, **kwargs): + # warnings.warn( + # f"Deprecated - use {fn.__name__}", DeprecationWarning, stacklevel=2 + # ) + return fn(self, *args, **kwargs) + + else: + + @wraps(fn) + def _inner(*args, **kwargs): + # warnings.warn( + # f"Deprecated - use {fn.__name__}", DeprecationWarning, stacklevel=2 + # ) + return fn(*args, **kwargs) + + _inner.__doc__ = f"""Deprecated - use :class:`{fn.__name__}`""" + _inner.__name__ = compat_name + _inner.__annotations__ = fn.__annotations__ + if isinstance(fn, types.FunctionType): + _inner.__kwdefaults__ = fn.__kwdefaults__ # type: ignore [attr-defined] + elif isinstance(fn, type) and hasattr(fn, "__init__"): + _inner.__kwdefaults__ = fn.__init__.__kwdefaults__ # type: ignore [misc,attr-defined] + else: + _inner.__kwdefaults__ = None # type: ignore [attr-defined] + _inner.__qualname__ = fn.__qualname__ + return cast(C, _inner) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/INSTALLER b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/LICENSE b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..1e65815cf0b3132689485874a93034ede7206bf4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/LICENSE @@ -0,0 +1,54 @@ +Copyright 2017- Paul Ganssle +Copyright 2017- dateutil contributors (see AUTHORS file) + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +The above license applies to all contributions after 2017-12-01, as well as +all contributions that have been re-licensed (see AUTHORS file for the list of +contributors who have re-licensed their code). +-------------------------------------------------------------------------------- +dateutil - Extensions to the standard Python datetime module. + +Copyright (c) 2003-2011 - Gustavo Niemeyer +Copyright (c) 2012-2014 - Tomi Pieviläinen +Copyright (c) 2014-2016 - Yaron de Leeuw +Copyright (c) 2015- - Paul Ganssle +Copyright (c) 2015- - dateutil contributors (see AUTHORS file) + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + * Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +The above BSD License Applies to all code, even that also covered by Apache 2.0. \ No newline at end of file diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/METADATA b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..577f2bf2b7749e1b123b8225d610b1b257e430cc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/METADATA @@ -0,0 +1,204 @@ +Metadata-Version: 2.1 +Name: python-dateutil +Version: 2.9.0.post0 +Summary: Extensions to the standard Python datetime module +Home-page: https://github.com/dateutil/dateutil +Author: Gustavo Niemeyer +Author-email: gustavo@niemeyer.net +Maintainer: Paul Ganssle +Maintainer-email: dateutil@python.org +License: Dual License +Project-URL: Documentation, https://dateutil.readthedocs.io/en/stable/ +Project-URL: Source, https://github.com/dateutil/dateutil +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: BSD License +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 2 +Classifier: Programming Language :: Python :: 2.7 +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.3 +Classifier: Programming Language :: Python :: 3.4 +Classifier: Programming Language :: Python :: 3.5 +Classifier: Programming Language :: Python :: 3.6 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Topic :: Software Development :: Libraries +Requires-Python: !=3.0.*,!=3.1.*,!=3.2.*,>=2.7 +Description-Content-Type: text/x-rst +License-File: LICENSE +Requires-Dist: six >=1.5 + +dateutil - powerful extensions to datetime +========================================== + +|pypi| |support| |licence| + +|gitter| |readthedocs| + +|travis| |appveyor| |pipelines| |coverage| + +.. |pypi| image:: https://img.shields.io/pypi/v/python-dateutil.svg?style=flat-square + :target: https://pypi.org/project/python-dateutil/ + :alt: pypi version + +.. |support| image:: https://img.shields.io/pypi/pyversions/python-dateutil.svg?style=flat-square + :target: https://pypi.org/project/python-dateutil/ + :alt: supported Python version + +.. |travis| image:: https://img.shields.io/travis/dateutil/dateutil/master.svg?style=flat-square&label=Travis%20Build + :target: https://travis-ci.org/dateutil/dateutil + :alt: travis build status + +.. |appveyor| image:: https://img.shields.io/appveyor/ci/dateutil/dateutil/master.svg?style=flat-square&logo=appveyor + :target: https://ci.appveyor.com/project/dateutil/dateutil + :alt: appveyor build status + +.. |pipelines| image:: https://dev.azure.com/pythondateutilazure/dateutil/_apis/build/status/dateutil.dateutil?branchName=master + :target: https://dev.azure.com/pythondateutilazure/dateutil/_build/latest?definitionId=1&branchName=master + :alt: azure pipelines build status + +.. |coverage| image:: https://codecov.io/gh/dateutil/dateutil/branch/master/graphs/badge.svg?branch=master + :target: https://codecov.io/gh/dateutil/dateutil?branch=master + :alt: Code coverage + +.. |gitter| image:: https://badges.gitter.im/dateutil/dateutil.svg + :alt: Join the chat at https://gitter.im/dateutil/dateutil + :target: https://gitter.im/dateutil/dateutil + +.. |licence| image:: https://img.shields.io/pypi/l/python-dateutil.svg?style=flat-square + :target: https://pypi.org/project/python-dateutil/ + :alt: licence + +.. |readthedocs| image:: https://img.shields.io/readthedocs/dateutil/latest.svg?style=flat-square&label=Read%20the%20Docs + :alt: Read the documentation at https://dateutil.readthedocs.io/en/latest/ + :target: https://dateutil.readthedocs.io/en/latest/ + +The `dateutil` module provides powerful extensions to +the standard `datetime` module, available in Python. + +Installation +============ +`dateutil` can be installed from PyPI using `pip` (note that the package name is +different from the importable name):: + + pip install python-dateutil + +Download +======== +dateutil is available on PyPI +https://pypi.org/project/python-dateutil/ + +The documentation is hosted at: +https://dateutil.readthedocs.io/en/stable/ + +Code +==== +The code and issue tracker are hosted on GitHub: +https://github.com/dateutil/dateutil/ + +Features +======== + +* Computing of relative deltas (next month, next year, + next Monday, last week of month, etc); +* Computing of relative deltas between two given + date and/or datetime objects; +* Computing of dates based on very flexible recurrence rules, + using a superset of the `iCalendar `_ + specification. Parsing of RFC strings is supported as well. +* Generic parsing of dates in almost any string format; +* Timezone (tzinfo) implementations for tzfile(5) format + files (/etc/localtime, /usr/share/zoneinfo, etc), TZ + environment string (in all known formats), iCalendar + format files, given ranges (with help from relative deltas), + local machine timezone, fixed offset timezone, UTC timezone, + and Windows registry-based time zones. +* Internal up-to-date world timezone information based on + Olson's database. +* Computing of Easter Sunday dates for any given year, + using Western, Orthodox or Julian algorithms; +* A comprehensive test suite. + +Quick example +============= +Here's a snapshot, just to give an idea about the power of the +package. For more examples, look at the documentation. + +Suppose you want to know how much time is left, in +years/months/days/etc, before the next easter happening on a +year with a Friday 13th in August, and you want to get today's +date out of the "date" unix system command. Here is the code: + +.. code-block:: python3 + + >>> from dateutil.relativedelta import * + >>> from dateutil.easter import * + >>> from dateutil.rrule import * + >>> from dateutil.parser import * + >>> from datetime import * + >>> now = parse("Sat Oct 11 17:13:46 UTC 2003") + >>> today = now.date() + >>> year = rrule(YEARLY,dtstart=now,bymonth=8,bymonthday=13,byweekday=FR)[0].year + >>> rdelta = relativedelta(easter(year), today) + >>> print("Today is: %s" % today) + Today is: 2003-10-11 + >>> print("Year with next Aug 13th on a Friday is: %s" % year) + Year with next Aug 13th on a Friday is: 2004 + >>> print("How far is the Easter of that year: %s" % rdelta) + How far is the Easter of that year: relativedelta(months=+6) + >>> print("And the Easter of that year is: %s" % (today+rdelta)) + And the Easter of that year is: 2004-04-11 + +Being exactly 6 months ahead was **really** a coincidence :) + +Contributing +============ + +We welcome many types of contributions - bug reports, pull requests (code, infrastructure or documentation fixes). For more information about how to contribute to the project, see the ``CONTRIBUTING.md`` file in the repository. + + +Author +====== +The dateutil module was written by Gustavo Niemeyer +in 2003. + +It is maintained by: + +* Gustavo Niemeyer 2003-2011 +* Tomi Pieviläinen 2012-2014 +* Yaron de Leeuw 2014-2016 +* Paul Ganssle 2015- + +Starting with version 2.4.1 and running until 2.8.2, all source and binary +distributions will be signed by a PGP key that has, at the very least, been +signed by the key which made the previous release. A table of release signing +keys can be found below: + +=========== ============================ +Releases Signing key fingerprint +=========== ============================ +2.4.1-2.8.2 `6B49 ACBA DCF6 BD1C A206 67AB CD54 FCE3 D964 BEFB`_ +=========== ============================ + +New releases *may* have signed tags, but binary and source distributions +uploaded to PyPI will no longer have GPG signatures attached. + +Contact +======= +Our mailing list is available at `dateutil@python.org `_. As it is hosted by the PSF, it is subject to the `PSF code of +conduct `_. + +License +======= + +All contributions after December 1, 2017 released under dual license - either `Apache 2.0 License `_ or the `BSD 3-Clause License `_. Contributions before December 1, 2017 - except those those explicitly relicensed - are released only under the BSD 3-Clause License. + + +.. _6B49 ACBA DCF6 BD1C A206 67AB CD54 FCE3 D964 BEFB: + https://pgp.mit.edu/pks/lookup?op=vindex&search=0xCD54FCE3D964BEFB diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/RECORD b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..176502e4e04b23de8580d433f4ee9101ac917a93 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/RECORD @@ -0,0 +1,44 @@ +dateutil/__init__.py,sha256=Mqam67WO9IkTmUFyI66vS6IoSXTp9G388DadH2LCMLY,620 +dateutil/__pycache__/__init__.cpython-310.pyc,, +dateutil/__pycache__/_common.cpython-310.pyc,, +dateutil/__pycache__/_version.cpython-310.pyc,, +dateutil/__pycache__/easter.cpython-310.pyc,, +dateutil/__pycache__/relativedelta.cpython-310.pyc,, +dateutil/__pycache__/rrule.cpython-310.pyc,, +dateutil/__pycache__/tzwin.cpython-310.pyc,, +dateutil/__pycache__/utils.cpython-310.pyc,, +dateutil/_common.py,sha256=77w0yytkrxlYbSn--lDVPUMabUXRR9I3lBv_vQRUqUY,932 +dateutil/_version.py,sha256=BV031OxDDAmy58neUg5yyqLkLaqIw7ibK9As3jiMib0,166 +dateutil/easter.py,sha256=dyBi-lKvimH1u_k6p7Z0JJK72QhqVtVBsqByvpEPKvc,2678 +dateutil/parser/__init__.py,sha256=wWk6GFuxTpjoggCGtgkceJoti4pVjl4_fHQXpNOaSYg,1766 +dateutil/parser/__pycache__/__init__.cpython-310.pyc,, +dateutil/parser/__pycache__/_parser.cpython-310.pyc,, +dateutil/parser/__pycache__/isoparser.cpython-310.pyc,, +dateutil/parser/_parser.py,sha256=7klDdyicksQB_Xgl-3UAmBwzCYor1AIZqklIcT6dH_8,58796 +dateutil/parser/isoparser.py,sha256=8Fy999bnCd1frSdOYuOraWfJTtd5W7qQ51NwNuH_hXM,13233 +dateutil/relativedelta.py,sha256=IY_mglMjoZbYfrvloTY2ce02aiVjPIkiZfqgNTZRfuA,24903 +dateutil/rrule.py,sha256=KJzKlaCd1jEbu4A38ZltslaoAUh9nSbdbOFdjp70Kew,66557 +dateutil/tz/__init__.py,sha256=F-Mz13v6jYseklQf9Te9J6nzcLDmq47gORa61K35_FA,444 +dateutil/tz/__pycache__/__init__.cpython-310.pyc,, +dateutil/tz/__pycache__/_common.cpython-310.pyc,, +dateutil/tz/__pycache__/_factories.cpython-310.pyc,, +dateutil/tz/__pycache__/tz.cpython-310.pyc,, +dateutil/tz/__pycache__/win.cpython-310.pyc,, +dateutil/tz/_common.py,sha256=cgzDTANsOXvEc86cYF77EsliuSab8Puwpsl5-bX3_S4,12977 +dateutil/tz/_factories.py,sha256=unb6XQNXrPMveksTCU-Ag8jmVZs4SojoPUcAHpWnrvU,2569 +dateutil/tz/tz.py,sha256=EUnEdMfeThXiY6l4sh9yBabZ63_POzy01zSsh9thn1o,62855 +dateutil/tz/win.py,sha256=xJszWgSwE1xPx_HJj4ZkepyukC_hNy016WMcXhbRaB8,12935 +dateutil/tzwin.py,sha256=7Ar4vdQCnnM0mKR3MUjbIKsZrBVfHgdwsJZc_mGYRew,59 +dateutil/utils.py,sha256=dKCchEw8eObi0loGTx91unBxm_7UGlU3v_FjFMdqwYM,1965 +dateutil/zoneinfo/__init__.py,sha256=KYg0pthCMjcp5MXSEiBJn3nMjZeNZav7rlJw5-tz1S4,5889 +dateutil/zoneinfo/__pycache__/__init__.cpython-310.pyc,, +dateutil/zoneinfo/__pycache__/rebuild.cpython-310.pyc,, +dateutil/zoneinfo/dateutil-zoneinfo.tar.gz,sha256=0-pS57bpaN4NiE3xKIGTWW-pW4A9tPkqGCeac5gARHU,156400 +dateutil/zoneinfo/rebuild.py,sha256=MiqYzCIHvNbMH-LdRYLv-4T0EIA7hDKt5GLR0IRTLdI,2392 +python_dateutil-2.9.0.post0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +python_dateutil-2.9.0.post0.dist-info/LICENSE,sha256=ugD1Gg2SgjtaHN4n2LW50jIeZ-2NqbwWPv-W1eF-V34,2889 +python_dateutil-2.9.0.post0.dist-info/METADATA,sha256=qdQ22jIr6AgzL5jYgyWZjofLaTpniplp_rTPrXKabpM,8354 +python_dateutil-2.9.0.post0.dist-info/RECORD,, +python_dateutil-2.9.0.post0.dist-info/WHEEL,sha256=-G_t0oGuE7UD0DrSpVZnq1hHMBV9DD2XkS5v7XpmTnk,110 +python_dateutil-2.9.0.post0.dist-info/top_level.txt,sha256=4tjdWkhRZvF7LA_BYe_L9gB2w_p2a-z5y6ArjaRkot8,9 +python_dateutil-2.9.0.post0.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/WHEEL b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..4724c45738f6ac125bb3a21787855562e6870440 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: bdist_wheel (0.42.0) +Root-Is-Purelib: true +Tag: py2-none-any +Tag: py3-none-any + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/top_level.txt b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..66501480ba5b63f98ee9a59c1f99e5e6917da6d9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/top_level.txt @@ -0,0 +1 @@ +dateutil diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/zip-safe b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/zip-safe new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/python_dateutil-2.9.0.post0.dist-info/zip-safe @@ -0,0 +1 @@ + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy.libs/libscipy_openblas-f07f5a5d207a3a47104dca54d6d0c86a.dll b/external/alphageometry/.venv-ag/Lib/site-packages/scipy.libs/libscipy_openblas-f07f5a5d207a3a47104dca54d6d0c86a.dll new file mode 100644 index 0000000000000000000000000000000000000000..1fda406d8daf44a45ffcc165b7374478aa48987e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy.libs/libscipy_openblas-f07f5a5d207a3a47104dca54d6d0c86a.dll @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b2103f2ae4d8547998b5d188e9801fba6cb12404ae8e4bfff319e8cc1949000 +size 20155392 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/__config__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/__config__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ce033f8b171623ece455094c092fa786b1d0f06 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/__config__.py @@ -0,0 +1,161 @@ +# This file is generated by SciPy's build process +# It contains system_info results at the time of building this package. +from enum import Enum + +__all__ = ["show"] +_built_with_meson = True + + +class DisplayModes(Enum): + stdout = "stdout" + dicts = "dicts" + + +def _cleanup(d): + """ + Removes empty values in a `dict` recursively + This ensures we remove values that Meson could not provide to CONFIG + """ + if isinstance(d, dict): + return { k: _cleanup(v) for k, v in d.items() if v != '' and _cleanup(v) != '' } + else: + return d + + +CONFIG = _cleanup( + { + "Compilers": { + "c": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.3.0", + "commands": r"cc", + "args": r"", + "linker args": r"", + }, + "cython": { + "name": r"cython", + "linker": r"cython", + "version": r"3.0.12", + "commands": r"cython", + "args": r"", + "linker args": r"", + }, + "c++": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.3.0", + "commands": r"c++", + "args": r"", + "linker args": r"", + }, + "fortran": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.3.0", + "commands": r"gfortran", + "args": r"", + "linker args": r"", + }, + "pythran": { + "version": r"0.17.0", + "include directory": r"C:\Users\runneradmin\AppData\Local\Temp\pip-build-env-1iowlesx\overlay\Lib\site-packages/pythran" + }, + }, + "Machine Information": { + "host": { + "cpu": r"x86_64", + "family": r"x86_64", + "endian": r"little", + "system": r"windows", + }, + "build": { + "cpu": r"x86_64", + "family": r"x86_64", + "endian": r"little", + "system": r"windows", + }, + "cross-compiled": bool("False".lower().replace('false', '')), + }, + "Build Dependencies": { + "blas": { + "name": "scipy-openblas", + "found": bool("True".lower().replace('false', '')), + "version": "0.3.28", + "detection method": "pkgconfig", + "include directory": r"C:/Users/runneradmin/AppData/Local/Temp/cibw-run-6w_ufqtf/cp310-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/include", + "lib directory": r"C:/Users/runneradmin/AppData/Local/Temp/cibw-run-6w_ufqtf/cp310-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/lib", + "openblas configuration": r"OpenBLAS 0.3.28 DYNAMIC_ARCH NO_AFFINITY Haswell MAX_THREADS=24", + "pc file directory": r"D:/a/scipy/scipy", + }, + "lapack": { + "name": "scipy-openblas", + "found": bool("True".lower().replace('false', '')), + "version": "0.3.28", + "detection method": "pkgconfig", + "include directory": r"C:/Users/runneradmin/AppData/Local/Temp/cibw-run-6w_ufqtf/cp310-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/include", + "lib directory": r"C:/Users/runneradmin/AppData/Local/Temp/cibw-run-6w_ufqtf/cp310-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/lib", + "openblas configuration": r"OpenBLAS 0.3.28 DYNAMIC_ARCH NO_AFFINITY Haswell MAX_THREADS=24", + "pc file directory": r"D:/a/scipy/scipy", + }, + "pybind11": { + "name": "pybind11", + "version": "2.13.6", + "detection method": "config-tool", + "include directory": r"unknown", + }, + }, + "Python Information": { + "path": r"C:\Users\runneradmin\AppData\Local\Temp\cibw-run-6w_ufqtf\cp310-win_amd64\build\venv\Scripts\python.exe", + "version": "3.10", + }, + } +) + + +def _check_pyyaml(): + import yaml + + return yaml + + +def show(mode=DisplayModes.stdout.value): + """ + Show libraries and system information on which SciPy was built + and is being used + + Parameters + ---------- + mode : {`'stdout'`, `'dicts'`}, optional. + Indicates how to display the config information. + `'stdout'` prints to console, `'dicts'` returns a dictionary + of the configuration. + + Returns + ------- + out : {`dict`, `None`} + If mode is `'dicts'`, a dict is returned, else None + + Notes + ----- + 1. The `'stdout'` mode will give more readable + output if ``pyyaml`` is installed + + """ + if mode == DisplayModes.stdout.value: + try: # Non-standard library, check import + yaml = _check_pyyaml() + + print(yaml.dump(CONFIG)) + except ModuleNotFoundError: + import warnings + import json + + warnings.warn("Install `pyyaml` for better output", stacklevel=1) + print(json.dumps(CONFIG, indent=2)) + elif mode == DisplayModes.dicts.value: + return CONFIG + else: + raise AttributeError( + f"Invalid `mode`, use one of: {', '.join([e.value for e in DisplayModes])}" + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fbb72626379239daefe13489ec4df6037e6d228 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/__init__.py @@ -0,0 +1,153 @@ +""" +SciPy: A scientific computing package for Python +================================================ + +Documentation is available in the docstrings and +online at https://docs.scipy.org. + +Subpackages +----------- +Using any of these subpackages requires an explicit import. For example, +``import scipy.cluster``. + +:: + + cluster --- Vector Quantization / Kmeans + constants --- Physical and mathematical constants and units + datasets --- Dataset methods + differentiate --- Finite difference differentiation tools + fft --- Discrete Fourier transforms + fftpack --- Legacy discrete Fourier transforms + integrate --- Integration routines + interpolate --- Interpolation Tools + io --- Data input and output + linalg --- Linear algebra routines + ndimage --- N-D image package + odr --- Orthogonal Distance Regression + optimize --- Optimization Tools + signal --- Signal Processing Tools + sparse --- Sparse Matrices + spatial --- Spatial data structures and algorithms + special --- Special functions + stats --- Statistical Functions + +Public API in the main SciPy namespace +-------------------------------------- +:: + + __version__ --- SciPy version string + LowLevelCallable --- Low-level callback function + show_config --- Show scipy build configuration + test --- Run scipy unittests + +""" + + +# start delvewheel patch +def _delvewheel_patch_1_10_1(): + import os + if os.path.isdir(libs_dir := os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'scipy.libs'))): + os.add_dll_directory(libs_dir) + + +_delvewheel_patch_1_10_1() +del _delvewheel_patch_1_10_1 +# end delvewheel patch + +import importlib as _importlib + +from numpy import __version__ as __numpy_version__ + + +try: + from scipy.__config__ import show as show_config +except ImportError as e: + msg = """Error importing SciPy: you cannot import SciPy while + being in scipy source directory; please exit the SciPy source + tree first and relaunch your Python interpreter.""" + raise ImportError(msg) from e + + +from scipy.version import version as __version__ + + +# Allow distributors to run custom init code +from . import _distributor_init +del _distributor_init + + +from scipy._lib import _pep440 +# In maintenance branch, change to np_maxversion N+3 if numpy is at N +np_minversion = '1.23.5' +np_maxversion = '2.5.0' +if (_pep440.parse(__numpy_version__) < _pep440.Version(np_minversion) or + _pep440.parse(__numpy_version__) >= _pep440.Version(np_maxversion)): + import warnings + warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}" + f" is required for this version of SciPy (detected " + f"version {__numpy_version__})", + UserWarning, stacklevel=2) +del _pep440 + + +# This is the first import of an extension module within SciPy. If there's +# a general issue with the install, such that extension modules are missing +# or cannot be imported, this is where we'll get a failure - so give an +# informative error message. +try: + from scipy._lib._ccallback import LowLevelCallable +except ImportError as e: + msg = "The `scipy` install you are using seems to be broken, " + \ + "(extension modules cannot be imported), " + \ + "please try reinstalling." + raise ImportError(msg) from e + + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester + + +submodules = [ + 'cluster', + 'constants', + 'datasets', + 'differentiate', + 'fft', + 'fftpack', + 'integrate', + 'interpolate', + 'io', + 'linalg', + 'ndimage', + 'odr', + 'optimize', + 'signal', + 'sparse', + 'spatial', + 'special', + 'stats' +] + +__all__ = submodules + [ + 'LowLevelCallable', + 'test', + 'show_config', + '__version__', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + if name in submodules: + return _importlib.import_module(f'scipy.{name}') + else: + try: + return globals()[name] + except KeyError: + raise AttributeError( + f"Module 'scipy' has no attribute '{name}'" + ) \ No newline at end of file diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_distributor_init.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_distributor_init.py new file mode 100644 index 0000000000000000000000000000000000000000..0c3acda987d79141950e75e4ffe7927d392c0882 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_distributor_init.py @@ -0,0 +1,18 @@ +""" Distributor init file + +Distributors: you can replace the contents of this file with your own custom +code to support particular distributions of SciPy. + +For example, this is a good place to put any checks for hardware requirements +or BLAS/LAPACK library initialization. + +The SciPy standard source distribution will not put code in this file beyond +the try-except import of `_distributor_init_local` (which is not part of a +standard source distribution), so you can safely replace this file with your +own version. +""" + +try: + from . import _distributor_init_local # noqa: F401 +except ImportError: + pass diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0ce03d7cd467370ebcace81ebe7d00726046e20f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/__init__.py @@ -0,0 +1 @@ +from ._helpers import * # noqa: F403 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/_linalg.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/_linalg.py new file mode 100644 index 0000000000000000000000000000000000000000..937ef6ad6e19ac80c9d29b39ac90a32a01857336 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/_linalg.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple +if TYPE_CHECKING: + from typing import Literal, Optional, Tuple, Union + from ._typing import ndarray + +import math + +import numpy as np +if np.__version__[0] == "2": + from numpy.lib.array_utils import normalize_axis_tuple +else: + from numpy.core.numeric import normalize_axis_tuple + +from ._aliases import matmul, matrix_transpose, tensordot, vecdot, isdtype +from .._internal import get_xp + +# These are in the main NumPy namespace but not in numpy.linalg +def cross(x1: ndarray, x2: ndarray, /, xp, *, axis: int = -1, **kwargs) -> ndarray: + return xp.cross(x1, x2, axis=axis, **kwargs) + +def outer(x1: ndarray, x2: ndarray, /, xp, **kwargs) -> ndarray: + return xp.outer(x1, x2, **kwargs) + +class EighResult(NamedTuple): + eigenvalues: ndarray + eigenvectors: ndarray + +class QRResult(NamedTuple): + Q: ndarray + R: ndarray + +class SlogdetResult(NamedTuple): + sign: ndarray + logabsdet: ndarray + +class SVDResult(NamedTuple): + U: ndarray + S: ndarray + Vh: ndarray + +# These functions are the same as their NumPy counterparts except they return +# a namedtuple. +def eigh(x: ndarray, /, xp, **kwargs) -> EighResult: + return EighResult(*xp.linalg.eigh(x, **kwargs)) + +def qr(x: ndarray, /, xp, *, mode: Literal['reduced', 'complete'] = 'reduced', + **kwargs) -> QRResult: + return QRResult(*xp.linalg.qr(x, mode=mode, **kwargs)) + +def slogdet(x: ndarray, /, xp, **kwargs) -> SlogdetResult: + return SlogdetResult(*xp.linalg.slogdet(x, **kwargs)) + +def svd(x: ndarray, /, xp, *, full_matrices: bool = True, **kwargs) -> SVDResult: + return SVDResult(*xp.linalg.svd(x, full_matrices=full_matrices, **kwargs)) + +# These functions have additional keyword arguments + +# The upper keyword argument is new from NumPy +def cholesky(x: ndarray, /, xp, *, upper: bool = False, **kwargs) -> ndarray: + L = xp.linalg.cholesky(x, **kwargs) + if upper: + U = get_xp(xp)(matrix_transpose)(L) + if get_xp(xp)(isdtype)(U.dtype, 'complex floating'): + U = xp.conj(U) + return U + return L + +# The rtol keyword argument of matrix_rank() and pinv() is new from NumPy. +# Note that it has a different semantic meaning from tol and rcond. +def matrix_rank(x: ndarray, + /, + xp, + *, + rtol: Optional[Union[float, ndarray]] = None, + **kwargs) -> ndarray: + # this is different from xp.linalg.matrix_rank, which supports 1 + # dimensional arrays. + if x.ndim < 2: + raise xp.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional") + S = get_xp(xp)(svdvals)(x, **kwargs) + if rtol is None: + tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * xp.finfo(S.dtype).eps + else: + # this is different from xp.linalg.matrix_rank, which does not + # multiply the tolerance by the largest singular value. + tol = S.max(axis=-1, keepdims=True)*xp.asarray(rtol)[..., xp.newaxis] + return xp.count_nonzero(S > tol, axis=-1) + +def pinv(x: ndarray, /, xp, *, rtol: Optional[Union[float, ndarray]] = None, **kwargs) -> ndarray: + # this is different from xp.linalg.pinv, which does not multiply the + # default tolerance by max(M, N). + if rtol is None: + rtol = max(x.shape[-2:]) * xp.finfo(x.dtype).eps + return xp.linalg.pinv(x, rcond=rtol, **kwargs) + +# These functions are new in the array API spec + +def matrix_norm(x: ndarray, /, xp, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> ndarray: + return xp.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord) + +# svdvals is not in NumPy (but it is in SciPy). It is equivalent to +# xp.linalg.svd(compute_uv=False). +def svdvals(x: ndarray, /, xp) -> Union[ndarray, Tuple[ndarray, ...]]: + return xp.linalg.svd(x, compute_uv=False) + +def vector_norm(x: ndarray, /, xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> ndarray: + # xp.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or + # when axis=None and the input is 2-D, so to force a vector norm, we make + # it so the input is 1-D (for axis=None), or reshape so that norm is done + # on a single dimension. + if axis is None: + # Note: xp.linalg.norm() doesn't handle 0-D arrays + _x = x.ravel() + _axis = 0 + elif isinstance(axis, tuple): + # Note: The axis argument supports any number of axes, whereas + # xp.linalg.norm() only supports a single axis for vector norm. + normalized_axis = normalize_axis_tuple(axis, x.ndim) + rest = tuple(i for i in range(x.ndim) if i not in normalized_axis) + newshape = axis + rest + _x = xp.transpose(x, newshape).reshape( + (math.prod([x.shape[i] for i in axis]), *[x.shape[i] for i in rest])) + _axis = 0 + else: + _x = x + _axis = axis + + res = xp.linalg.norm(_x, axis=_axis, ord=ord) + + if keepdims: + # We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks + # above to avoid matrix norm logic. + shape = list(x.shape) + _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim) + for i in _axis: + shape[i] = 1 + res = xp.reshape(res, tuple(shape)) + + return res + +# xp.diagonal and xp.trace operate on the first two axes whereas these +# operates on the last two + +def diagonal(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray: + return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs) + +def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarray: + return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs)) + +__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult', + 'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet', + 'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm', + 'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal', + 'trace'] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/_typing.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..1a4fec36af33f80f9efb7da6c2e91cb403deb9d1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/common/_typing.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +__all__ = [ + "NestedSequence", + "SupportsBufferProtocol", +] + +from typing import ( + Any, + TypeVar, + Protocol, +) + +_T_co = TypeVar("_T_co", covariant=True) + +class NestedSequence(Protocol[_T_co]): + def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ... + def __len__(self, /) -> int: ... + +SupportsBufferProtocol = Any + +Array = Any +Device = Any diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73179b25003f948d948fbb190535334db28fffe0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/__init__.py @@ -0,0 +1,16 @@ +from cupy import * # noqa: F403 + +# from cupy import * doesn't overwrite these builtin names +from cupy import abs, max, min, round # noqa: F401 + +# These imports may overwrite names from the import * above. +from ._aliases import * # noqa: F403 + +# See the comment in the numpy __init__.py +__import__(__package__ + '.linalg') + +__import__(__package__ + '.fft') + +from ..common._helpers import * # noqa: F401,F403 + +__array_api_version__ = '2023.12' diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_aliases.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..d24858210f7b12d24cb851e45c1d70f622c4a31c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_aliases.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import cupy as cp + +from ..common import _aliases +from .._internal import get_xp + +from ._info import __array_namespace_info__ + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from typing import Optional, Union + from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol + +bool = cp.bool_ + +# Basic renames +acos = cp.arccos +acosh = cp.arccosh +asin = cp.arcsin +asinh = cp.arcsinh +atan = cp.arctan +atan2 = cp.arctan2 +atanh = cp.arctanh +bitwise_left_shift = cp.left_shift +bitwise_invert = cp.invert +bitwise_right_shift = cp.right_shift +concat = cp.concatenate +pow = cp.power + +arange = get_xp(cp)(_aliases.arange) +empty = get_xp(cp)(_aliases.empty) +empty_like = get_xp(cp)(_aliases.empty_like) +eye = get_xp(cp)(_aliases.eye) +full = get_xp(cp)(_aliases.full) +full_like = get_xp(cp)(_aliases.full_like) +linspace = get_xp(cp)(_aliases.linspace) +ones = get_xp(cp)(_aliases.ones) +ones_like = get_xp(cp)(_aliases.ones_like) +zeros = get_xp(cp)(_aliases.zeros) +zeros_like = get_xp(cp)(_aliases.zeros_like) +UniqueAllResult = get_xp(cp)(_aliases.UniqueAllResult) +UniqueCountsResult = get_xp(cp)(_aliases.UniqueCountsResult) +UniqueInverseResult = get_xp(cp)(_aliases.UniqueInverseResult) +unique_all = get_xp(cp)(_aliases.unique_all) +unique_counts = get_xp(cp)(_aliases.unique_counts) +unique_inverse = get_xp(cp)(_aliases.unique_inverse) +unique_values = get_xp(cp)(_aliases.unique_values) +astype = _aliases.astype +std = get_xp(cp)(_aliases.std) +var = get_xp(cp)(_aliases.var) +cumulative_sum = get_xp(cp)(_aliases.cumulative_sum) +clip = get_xp(cp)(_aliases.clip) +permute_dims = get_xp(cp)(_aliases.permute_dims) +reshape = get_xp(cp)(_aliases.reshape) +argsort = get_xp(cp)(_aliases.argsort) +sort = get_xp(cp)(_aliases.sort) +nonzero = get_xp(cp)(_aliases.nonzero) +ceil = get_xp(cp)(_aliases.ceil) +floor = get_xp(cp)(_aliases.floor) +trunc = get_xp(cp)(_aliases.trunc) +matmul = get_xp(cp)(_aliases.matmul) +matrix_transpose = get_xp(cp)(_aliases.matrix_transpose) +tensordot = get_xp(cp)(_aliases.tensordot) +sign = get_xp(cp)(_aliases.sign) + +_copy_default = object() + +# asarray also adds the copy keyword, which is not present in numpy 1.0. +def asarray( + obj: Union[ + ndarray, + bool, + int, + float, + NestedSequence[bool | int | float], + SupportsBufferProtocol, + ], + /, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + copy: Optional[bool] = _copy_default, + **kwargs, +) -> ndarray: + """ + Array API compatibility wrapper for asarray(). + + See the corresponding documentation in the array library and/or the array API + specification for more details. + """ + with cp.cuda.Device(device): + # cupy is like NumPy 1.26 (except without _CopyMode). See the comments + # in asarray in numpy/_aliases.py. + if copy is not _copy_default: + # A future version of CuPy will change the meaning of copy=False + # to mean no-copy. We don't know for certain what version it will + # be yet, so to avoid breaking that version, we use a different + # default value for copy so asarray(obj) with no copy kwarg will + # always do the copy-if-needed behavior. + + # This will still need to be updated to remove the + # NotImplementedError for copy=False, but at least this won't + # break the default or existing behavior. + if copy is None: + copy = False + elif copy is False: + raise NotImplementedError("asarray(copy=False) is not yet supported in cupy") + kwargs['copy'] = copy + + return cp.array(obj, dtype=dtype, **kwargs) + +# These functions are completely new here. If the library already has them +# (i.e., numpy 2.0), use the library version instead of our wrapper. +if hasattr(cp, 'vecdot'): + vecdot = cp.vecdot +else: + vecdot = get_xp(cp)(_aliases.vecdot) + +if hasattr(cp, 'isdtype'): + isdtype = cp.isdtype +else: + isdtype = get_xp(cp)(_aliases.isdtype) + +if hasattr(cp, 'unstack'): + unstack = cp.unstack +else: + unstack = get_xp(cp)(_aliases.unstack) + +__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'bool', + 'acos', 'acosh', 'asin', 'asinh', 'atan', + 'atan2', 'atanh', 'bitwise_left_shift', + 'bitwise_invert', 'bitwise_right_shift', + 'concat', 'pow', 'sign'] + +_all_ignore = ['cp', 'get_xp'] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_info.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_info.py new file mode 100644 index 0000000000000000000000000000000000000000..d794100bbab633f315f7b71dd4031fa8fd499a04 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_info.py @@ -0,0 +1,326 @@ +""" +Array API Inspection namespace + +This is the namespace for inspection functions as defined by the array API +standard. See +https://data-apis.org/array-api/latest/API_specification/inspection.html for +more details. + +""" +from cupy import ( + dtype, + cuda, + bool_ as bool, + intp, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + complex64, + complex128, +) + +class __array_namespace_info__: + """ + Get the array API inspection namespace for CuPy. + + The array API inspection namespace defines the following functions: + + - capabilities() + - default_device() + - default_dtypes() + - dtypes() + - devices() + + See + https://data-apis.org/array-api/latest/API_specification/inspection.html + for more details. + + Returns + ------- + info : ModuleType + The array API inspection namespace for CuPy. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': cupy.float64, + 'complex floating': cupy.complex128, + 'integral': cupy.int64, + 'indexing': cupy.int64} + + """ + + __module__ = 'cupy' + + def capabilities(self): + """ + Return a dictionary of array API library capabilities. + + The resulting dictionary has the following keys: + + - **"boolean indexing"**: boolean indicating whether an array library + supports boolean indexing. Always ``True`` for CuPy. + + - **"data-dependent shapes"**: boolean indicating whether an array + library supports data-dependent output shapes. Always ``True`` for + CuPy. + + See + https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html + for more details. + + See Also + -------- + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + capabilities : dict + A dictionary of array API library capabilities. + + Examples + -------- + >>> info = xp.__array_namespace_info__() + >>> info.capabilities() + {'boolean indexing': True, + 'data-dependent shapes': True} + + """ + return { + "boolean indexing": True, + "data-dependent shapes": True, + # 'max rank' will be part of the 2024.12 standard + # "max rank": 64, + } + + def default_device(self): + """ + The default device used for new CuPy arrays. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + device : str + The default device used for new CuPy arrays. + + Examples + -------- + >>> info = xp.__array_namespace_info__() + >>> info.default_device() + Device(0) + + """ + return cuda.Device(0) + + def default_dtypes(self, *, device=None): + """ + The default data types used for new CuPy arrays. + + For CuPy, this always returns the following dictionary: + + - **"real floating"**: ``cupy.float64`` + - **"complex floating"**: ``cupy.complex128`` + - **"integral"**: ``cupy.intp`` + - **"indexing"**: ``cupy.intp`` + + Parameters + ---------- + device : str, optional + The device to get the default data types for. + + Returns + ------- + dtypes : dict + A dictionary describing the default data types used for new CuPy + arrays. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = xp.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': cupy.float64, + 'complex floating': cupy.complex128, + 'integral': cupy.int64, + 'indexing': cupy.int64} + + """ + # TODO: Does this depend on device? + return { + "real floating": dtype(float64), + "complex floating": dtype(complex128), + "integral": dtype(intp), + "indexing": dtype(intp), + } + + def dtypes(self, *, device=None, kind=None): + """ + The array API data types supported by CuPy. + + Note that this function only returns data types that are defined by + the array API. + + Parameters + ---------- + device : str, optional + The device to get the data types for. + kind : str or tuple of str, optional + The kind of data types to return. If ``None``, all data types are + returned. If a string, only data types of that kind are returned. + If a tuple, a dictionary containing the union of the given kinds + is returned. The following kinds are supported: + + - ``'bool'``: boolean data types (i.e., ``bool``). + - ``'signed integer'``: signed integer data types (i.e., ``int8``, + ``int16``, ``int32``, ``int64``). + - ``'unsigned integer'``: unsigned integer data types (i.e., + ``uint8``, ``uint16``, ``uint32``, ``uint64``). + - ``'integral'``: integer data types. Shorthand for ``('signed + integer', 'unsigned integer')``. + - ``'real floating'``: real-valued floating-point data types + (i.e., ``float32``, ``float64``). + - ``'complex floating'``: complex floating-point data types (i.e., + ``complex64``, ``complex128``). + - ``'numeric'``: numeric data types. Shorthand for ``('integral', + 'real floating', 'complex floating')``. + + Returns + ------- + dtypes : dict + A dictionary mapping the names of data types to the corresponding + CuPy data types. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = xp.__array_namespace_info__() + >>> info.dtypes(kind='signed integer') + {'int8': cupy.int8, + 'int16': cupy.int16, + 'int32': cupy.int32, + 'int64': cupy.int64} + + """ + # TODO: Does this depend on device? + if kind is None: + return { + "bool": dtype(bool), + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + "float32": dtype(float32), + "float64": dtype(float64), + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if kind == "bool": + return {"bool": bool} + if kind == "signed integer": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + } + if kind == "unsigned integer": + return { + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + } + if kind == "integral": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + } + if kind == "real floating": + return { + "float32": dtype(float32), + "float64": dtype(float64), + } + if kind == "complex floating": + return { + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if kind == "numeric": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + "float32": dtype(float32), + "float64": dtype(float64), + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if isinstance(kind, tuple): + res = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + def devices(self): + """ + The devices supported by CuPy. + + Returns + ------- + devices : list of str + The devices supported by CuPy. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes + + """ + return [cuda.Device(i) for i in range(cuda.runtime.getDeviceCount())] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_typing.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..488cc44c0d9ad1246d7ce8ddf9a8ea2a3ec4f3b1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/_typing.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +__all__ = [ + "ndarray", + "Device", + "Dtype", +] + +import sys +from typing import ( + Union, + TYPE_CHECKING, +) + +from cupy import ( + ndarray, + dtype, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, +) + +from cupy.cuda.device import Device + +if TYPE_CHECKING or sys.version_info >= (3, 9): + Dtype = dtype[Union[ + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + ]] +else: + Dtype = dtype diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/fft.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/fft.py new file mode 100644 index 0000000000000000000000000000000000000000..324fd9a8c680ea26c435c88bbea0133bb5ab1d7d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/fft.py @@ -0,0 +1,36 @@ +from cupy.fft import * # noqa: F403 +# cupy.fft doesn't have __all__. If it is added, replace this with +# +# from cupy.fft import __all__ as linalg_all +_n = {} +exec('from cupy.fft import *', _n) +del _n['__builtins__'] +fft_all = list(_n) +del _n + +from ..common import _fft +from .._internal import get_xp + +import cupy as cp + +fft = get_xp(cp)(_fft.fft) +ifft = get_xp(cp)(_fft.ifft) +fftn = get_xp(cp)(_fft.fftn) +ifftn = get_xp(cp)(_fft.ifftn) +rfft = get_xp(cp)(_fft.rfft) +irfft = get_xp(cp)(_fft.irfft) +rfftn = get_xp(cp)(_fft.rfftn) +irfftn = get_xp(cp)(_fft.irfftn) +hfft = get_xp(cp)(_fft.hfft) +ihfft = get_xp(cp)(_fft.ihfft) +fftfreq = get_xp(cp)(_fft.fftfreq) +rfftfreq = get_xp(cp)(_fft.rfftfreq) +fftshift = get_xp(cp)(_fft.fftshift) +ifftshift = get_xp(cp)(_fft.ifftshift) + +__all__ = fft_all + _fft.__all__ + +del get_xp +del cp +del fft_all +del _fft diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/linalg.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/linalg.py new file mode 100644 index 0000000000000000000000000000000000000000..49bda5c42d605aea6abbaf072a6cc688e04ee8a3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/cupy/linalg.py @@ -0,0 +1,49 @@ +from cupy.linalg import * # noqa: F403 +# cupy.linalg doesn't have __all__. If it is added, replace this with +# +# from cupy.linalg import __all__ as linalg_all +_n = {} +exec('from cupy.linalg import *', _n) +del _n['__builtins__'] +linalg_all = list(_n) +del _n + +from ..common import _linalg +from .._internal import get_xp + +import cupy as cp + +# These functions are in both the main and linalg namespaces +from ._aliases import matmul, matrix_transpose, tensordot, vecdot # noqa: F401 + +cross = get_xp(cp)(_linalg.cross) +outer = get_xp(cp)(_linalg.outer) +EighResult = _linalg.EighResult +QRResult = _linalg.QRResult +SlogdetResult = _linalg.SlogdetResult +SVDResult = _linalg.SVDResult +eigh = get_xp(cp)(_linalg.eigh) +qr = get_xp(cp)(_linalg.qr) +slogdet = get_xp(cp)(_linalg.slogdet) +svd = get_xp(cp)(_linalg.svd) +cholesky = get_xp(cp)(_linalg.cholesky) +matrix_rank = get_xp(cp)(_linalg.matrix_rank) +pinv = get_xp(cp)(_linalg.pinv) +matrix_norm = get_xp(cp)(_linalg.matrix_norm) +svdvals = get_xp(cp)(_linalg.svdvals) +diagonal = get_xp(cp)(_linalg.diagonal) +trace = get_xp(cp)(_linalg.trace) + +# These functions are completely new here. If the library already has them +# (i.e., numpy 2.0), use the library version instead of our wrapper. +if hasattr(cp.linalg, 'vector_norm'): + vector_norm = cp.linalg.vector_norm +else: + vector_norm = get_xp(cp)(_linalg.vector_norm) + +__all__ = linalg_all + _linalg.__all__ + +del get_xp +del cp +del linalg_all +del _linalg diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3139722dddde5862abaa1f6fabb52a6fe64bbc69 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/__init__.py @@ -0,0 +1,9 @@ +from dask.array import * # noqa: F403 + +# These imports may overwrite names from the import * above. +from ._aliases import * # noqa: F403 + +__array_api_version__ = '2023.12' + +__import__(__package__ + '.linalg') +__import__(__package__ + '.fft') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/_aliases.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..5b128607a80a8568b400716bbef60351dd025b6c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/_aliases.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +from ...common import _aliases +from ...common._helpers import _check_device + +from ..._internal import get_xp + +from ._info import __array_namespace_info__ + +import numpy as np +from numpy import ( + # Dtypes + iinfo, + finfo, + bool_ as bool, + float32, + float64, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + complex64, + complex128, + can_cast, + result_type, +) + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from typing import Optional, Union + + from ...common._typing import Device, Dtype, Array, NestedSequence, SupportsBufferProtocol + +import dask.array as da + +isdtype = get_xp(np)(_aliases.isdtype) +unstack = get_xp(da)(_aliases.unstack) +astype = _aliases.astype + +# Common aliases + +# This arange func is modified from the common one to +# not pass stop/step as keyword arguments, which will cause +# an error with dask + +# TODO: delete the xp stuff, it shouldn't be necessary +def _dask_arange( + start: Union[int, float], + /, + stop: Optional[Union[int, float]] = None, + step: Union[int, float] = 1, + *, + xp, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs, +) -> Array: + _check_device(xp, device) + args = [start] + if stop is not None: + args.append(stop) + else: + # stop is None, so start is actually stop + # prepend the default value for start which is 0 + args.insert(0, 0) + args.append(step) + return xp.arange(*args, dtype=dtype, **kwargs) + +arange = get_xp(da)(_dask_arange) +eye = get_xp(da)(_aliases.eye) + +linspace = get_xp(da)(_aliases.linspace) +eye = get_xp(da)(_aliases.eye) +UniqueAllResult = get_xp(da)(_aliases.UniqueAllResult) +UniqueCountsResult = get_xp(da)(_aliases.UniqueCountsResult) +UniqueInverseResult = get_xp(da)(_aliases.UniqueInverseResult) +unique_all = get_xp(da)(_aliases.unique_all) +unique_counts = get_xp(da)(_aliases.unique_counts) +unique_inverse = get_xp(da)(_aliases.unique_inverse) +unique_values = get_xp(da)(_aliases.unique_values) +permute_dims = get_xp(da)(_aliases.permute_dims) +std = get_xp(da)(_aliases.std) +var = get_xp(da)(_aliases.var) +cumulative_sum = get_xp(da)(_aliases.cumulative_sum) +empty = get_xp(da)(_aliases.empty) +empty_like = get_xp(da)(_aliases.empty_like) +full = get_xp(da)(_aliases.full) +full_like = get_xp(da)(_aliases.full_like) +ones = get_xp(da)(_aliases.ones) +ones_like = get_xp(da)(_aliases.ones_like) +zeros = get_xp(da)(_aliases.zeros) +zeros_like = get_xp(da)(_aliases.zeros_like) +reshape = get_xp(da)(_aliases.reshape) +matrix_transpose = get_xp(da)(_aliases.matrix_transpose) +vecdot = get_xp(da)(_aliases.vecdot) + +nonzero = get_xp(da)(_aliases.nonzero) +ceil = get_xp(np)(_aliases.ceil) +floor = get_xp(np)(_aliases.floor) +trunc = get_xp(np)(_aliases.trunc) +matmul = get_xp(np)(_aliases.matmul) +tensordot = get_xp(np)(_aliases.tensordot) +sign = get_xp(np)(_aliases.sign) + +# asarray also adds the copy keyword, which is not present in numpy 1.0. +def asarray( + obj: Union[ + Array, + bool, + int, + float, + NestedSequence[bool | int | float], + SupportsBufferProtocol, + ], + /, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + copy: "Optional[Union[bool, np._CopyMode]]" = None, + **kwargs, +) -> Array: + """ + Array API compatibility wrapper for asarray(). + + See the corresponding documentation in the array library and/or the array API + specification for more details. + """ + if copy is False: + # copy=False is not yet implemented in dask + raise NotImplementedError("copy=False is not yet implemented") + elif copy is True: + if isinstance(obj, da.Array) and dtype is None: + return obj.copy() + # Go through numpy, since dask copy is no-op by default + obj = np.array(obj, dtype=dtype, copy=True) + return da.array(obj, dtype=dtype) + else: + if not isinstance(obj, da.Array) or dtype is not None and obj.dtype != dtype: + obj = np.asarray(obj, dtype=dtype) + return da.from_array(obj) + return obj + + return da.asarray(obj, dtype=dtype, **kwargs) + +from dask.array import ( + # Element wise aliases + arccos as acos, + arccosh as acosh, + arcsin as asin, + arcsinh as asinh, + arctan as atan, + arctan2 as atan2, + arctanh as atanh, + left_shift as bitwise_left_shift, + right_shift as bitwise_right_shift, + invert as bitwise_invert, + power as pow, + # Other + concatenate as concat, +) + +# dask.array.clip does not work unless all three arguments are provided. +# Furthermore, the masking workaround in common._aliases.clip cannot work with +# dask (meaning uint64 promoting to float64 is going to just be unfixed for +# now). +@get_xp(da) +def clip( + x: Array, + /, + min: Optional[Union[int, float, Array]] = None, + max: Optional[Union[int, float, Array]] = None, + *, + xp, +) -> Array: + def _isscalar(a): + return isinstance(a, (int, float, type(None))) + min_shape = () if _isscalar(min) else min.shape + max_shape = () if _isscalar(max) else max.shape + + # TODO: This won't handle dask unknown shapes + import numpy as np + result_shape = np.broadcast_shapes(x.shape, min_shape, max_shape) + + if min is not None: + min = xp.broadcast_to(xp.asarray(min), result_shape) + if max is not None: + max = xp.broadcast_to(xp.asarray(max), result_shape) + + if min is None and max is None: + return xp.positive(x) + + if min is None: + return astype(xp.minimum(x, max), x.dtype) + if max is None: + return astype(xp.maximum(x, min), x.dtype) + + return astype(xp.minimum(xp.maximum(x, min), max), x.dtype) + +# exclude these from all since dask.array has no sorting functions +_da_unsupported = ['sort', 'argsort'] + +_common_aliases = [alias for alias in _aliases.__all__ if alias not in _da_unsupported] + +__all__ = _common_aliases + ['__array_namespace_info__', 'asarray', 'acos', + 'acosh', 'asin', 'asinh', 'atan', 'atan2', + 'atanh', 'bitwise_left_shift', 'bitwise_invert', + 'bitwise_right_shift', 'concat', 'pow', 'iinfo', 'finfo', 'can_cast', + 'result_type', 'bool', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64', + 'uint8', 'uint16', 'uint32', 'uint64', + 'complex64', 'complex128', 'iinfo', 'finfo', + 'can_cast', 'result_type'] + +_all_ignore = ["get_xp", "da", "np"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/_info.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/_info.py new file mode 100644 index 0000000000000000000000000000000000000000..4c299f9084bbc14de62a8e26b253fca89294394a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/_info.py @@ -0,0 +1,345 @@ +""" +Array API Inspection namespace + +This is the namespace for inspection functions as defined by the array API +standard. See +https://data-apis.org/array-api/latest/API_specification/inspection.html for +more details. + +""" +from numpy import ( + dtype, + bool_ as bool, + intp, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + complex64, + complex128, +) + +from ...common._helpers import _DASK_DEVICE + +class __array_namespace_info__: + """ + Get the array API inspection namespace for Dask. + + The array API inspection namespace defines the following functions: + + - capabilities() + - default_device() + - default_dtypes() + - dtypes() + - devices() + + See + https://data-apis.org/array-api/latest/API_specification/inspection.html + for more details. + + Returns + ------- + info : ModuleType + The array API inspection namespace for Dask. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': dask.float64, + 'complex floating': dask.complex128, + 'integral': dask.int64, + 'indexing': dask.int64} + + """ + + __module__ = 'dask.array' + + def capabilities(self): + """ + Return a dictionary of array API library capabilities. + + The resulting dictionary has the following keys: + + - **"boolean indexing"**: boolean indicating whether an array library + supports boolean indexing. Always ``False`` for Dask. + + - **"data-dependent shapes"**: boolean indicating whether an array + library supports data-dependent output shapes. Always ``False`` for + Dask. + + See + https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html + for more details. + + See Also + -------- + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + capabilities : dict + A dictionary of array API library capabilities. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.capabilities() + {'boolean indexing': True, + 'data-dependent shapes': True} + + """ + return { + "boolean indexing": False, + "data-dependent shapes": False, + # 'max rank' will be part of the 2024.12 standard + # "max rank": 64, + } + + def default_device(self): + """ + The default device used for new Dask arrays. + + For Dask, this always returns ``'cpu'``. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + device : str + The default device used for new Dask arrays. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_device() + 'cpu' + + """ + return "cpu" + + def default_dtypes(self, *, device=None): + """ + The default data types used for new Dask arrays. + + For Dask, this always returns the following dictionary: + + - **"real floating"**: ``numpy.float64`` + - **"complex floating"**: ``numpy.complex128`` + - **"integral"**: ``numpy.intp`` + - **"indexing"**: ``numpy.intp`` + + Parameters + ---------- + device : str, optional + The device to get the default data types for. + + Returns + ------- + dtypes : dict + A dictionary describing the default data types used for new Dask + arrays. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': dask.float64, + 'complex floating': dask.complex128, + 'integral': dask.int64, + 'indexing': dask.int64} + + """ + if device not in ["cpu", _DASK_DEVICE, None]: + raise ValueError( + 'Device not understood. Only "cpu" or _DASK_DEVICE is allowed, but received:' + f' {device}' + ) + return { + "real floating": dtype(float64), + "complex floating": dtype(complex128), + "integral": dtype(intp), + "indexing": dtype(intp), + } + + def dtypes(self, *, device=None, kind=None): + """ + The array API data types supported by Dask. + + Note that this function only returns data types that are defined by + the array API. + + Parameters + ---------- + device : str, optional + The device to get the data types for. + kind : str or tuple of str, optional + The kind of data types to return. If ``None``, all data types are + returned. If a string, only data types of that kind are returned. + If a tuple, a dictionary containing the union of the given kinds + is returned. The following kinds are supported: + + - ``'bool'``: boolean data types (i.e., ``bool``). + - ``'signed integer'``: signed integer data types (i.e., ``int8``, + ``int16``, ``int32``, ``int64``). + - ``'unsigned integer'``: unsigned integer data types (i.e., + ``uint8``, ``uint16``, ``uint32``, ``uint64``). + - ``'integral'``: integer data types. Shorthand for ``('signed + integer', 'unsigned integer')``. + - ``'real floating'``: real-valued floating-point data types + (i.e., ``float32``, ``float64``). + - ``'complex floating'``: complex floating-point data types (i.e., + ``complex64``, ``complex128``). + - ``'numeric'``: numeric data types. Shorthand for ``('integral', + 'real floating', 'complex floating')``. + + Returns + ------- + dtypes : dict + A dictionary mapping the names of data types to the corresponding + Dask data types. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.dtypes(kind='signed integer') + {'int8': dask.int8, + 'int16': dask.int16, + 'int32': dask.int32, + 'int64': dask.int64} + + """ + if device not in ["cpu", _DASK_DEVICE, None]: + raise ValueError( + 'Device not understood. Only "cpu" or _DASK_DEVICE is allowed, but received:' + f' {device}' + ) + if kind is None: + return { + "bool": dtype(bool), + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + "float32": dtype(float32), + "float64": dtype(float64), + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if kind == "bool": + return {"bool": bool} + if kind == "signed integer": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + } + if kind == "unsigned integer": + return { + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + } + if kind == "integral": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + } + if kind == "real floating": + return { + "float32": dtype(float32), + "float64": dtype(float64), + } + if kind == "complex floating": + return { + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if kind == "numeric": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + "float32": dtype(float32), + "float64": dtype(float64), + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if isinstance(kind, tuple): + res = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + def devices(self): + """ + The devices supported by Dask. + + For Dask, this always returns ``['cpu', DASK_DEVICE]``. + + Returns + ------- + devices : list of str + The devices supported by Dask. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.devices() + ['cpu', DASK_DEVICE] + + """ + return ["cpu", _DASK_DEVICE] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/fft.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/fft.py new file mode 100644 index 0000000000000000000000000000000000000000..e725a6bbceb26fab00a9339006303cfea8dc400a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/fft.py @@ -0,0 +1,24 @@ +from dask.array.fft import * # noqa: F403 +# dask.array.fft doesn't have __all__. If it is added, replace this with +# +# from dask.array.fft import __all__ as linalg_all +_n = {} +exec('from dask.array.fft import *', _n) +del _n['__builtins__'] +fft_all = list(_n) +del _n + +from ...common import _fft +from ..._internal import get_xp + +import dask.array as da + +fftfreq = get_xp(da)(_fft.fftfreq) +rfftfreq = get_xp(da)(_fft.rfftfreq) + +__all__ = [elem for elem in fft_all if elem != "annotations"] + ["fftfreq", "rfftfreq"] + +del get_xp +del da +del fft_all +del _fft diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/linalg.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/linalg.py new file mode 100644 index 0000000000000000000000000000000000000000..0c92b7e3b661ccd100463e3abb5b913c98a9c1c2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/dask/array/linalg.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from ...common import _linalg +from ..._internal import get_xp + +# Exports +from dask.array.linalg import * # noqa: F403 +from dask.array import outer + +# These functions are in both the main and linalg namespaces +from dask.array import matmul, tensordot +from ._aliases import matrix_transpose, vecdot + +import dask.array as da + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from ...common._typing import Array + from typing import Literal + +# dask.array.linalg doesn't have __all__. If it is added, replace this with +# +# from dask.array.linalg import __all__ as linalg_all +_n = {} +exec('from dask.array.linalg import *', _n) +del _n['__builtins__'] +if 'annotations' in _n: + del _n['annotations'] +linalg_all = list(_n) +del _n + +EighResult = _linalg.EighResult +QRResult = _linalg.QRResult +SlogdetResult = _linalg.SlogdetResult +SVDResult = _linalg.SVDResult +# TODO: use the QR wrapper once dask +# supports the mode keyword on QR +# https://github.com/dask/dask/issues/10388 +#qr = get_xp(da)(_linalg.qr) +def qr(x: Array, mode: Literal['reduced', 'complete'] = 'reduced', + **kwargs) -> QRResult: + if mode != "reduced": + raise ValueError("dask arrays only support using mode='reduced'") + return QRResult(*da.linalg.qr(x, **kwargs)) +trace = get_xp(da)(_linalg.trace) +cholesky = get_xp(da)(_linalg.cholesky) +matrix_rank = get_xp(da)(_linalg.matrix_rank) +matrix_norm = get_xp(da)(_linalg.matrix_norm) + + +# Wrap the svd functions to not pass full_matrices to dask +# when full_matrices=False (as that is the default behavior for dask), +# and dask doesn't have the full_matrices keyword +def svd(x: Array, full_matrices: bool = True, **kwargs) -> SVDResult: + if full_matrices: + raise ValueError("full_matrics=True is not supported by dask.") + return da.linalg.svd(x, coerce_signs=False, **kwargs) + +def svdvals(x: Array) -> Array: + # TODO: can't avoid computing U or V for dask + _, s, _ = svd(x) + return s + +vector_norm = get_xp(da)(_linalg.vector_norm) +diagonal = get_xp(da)(_linalg.diagonal) + +__all__ = linalg_all + ["trace", "outer", "matmul", "tensordot", + "matrix_transpose", "vecdot", "EighResult", + "QRResult", "SlogdetResult", "SVDResult", "qr", + "cholesky", "matrix_rank", "matrix_norm", "svdvals", + "vector_norm", "diagonal"] + +_all_ignore = ['get_xp', 'da', 'linalg_all'] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b313ddcc025a3e8ed71440000b96c1f3641dc3eb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/__init__.py @@ -0,0 +1,30 @@ +from numpy import * # noqa: F403 + +# from numpy import * doesn't overwrite these builtin names +from numpy import abs, max, min, round # noqa: F401 + +# These imports may overwrite names from the import * above. +from ._aliases import * # noqa: F403 + +# Don't know why, but we have to do an absolute import to import linalg. If we +# instead do +# +# from . import linalg +# +# It doesn't overwrite np.linalg from above. The import is generated +# dynamically so that the library can be vendored. +__import__(__package__ + '.linalg') + +__import__(__package__ + '.fft') + +from .linalg import matrix_transpose, vecdot # noqa: F401 + +from ..common._helpers import * # noqa: F403 + +try: + # Used in asarray(). Not present in older versions. + from numpy import _CopyMode # noqa: F401 +except ImportError: + pass + +__array_api_version__ = '2023.12' diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_aliases.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..5741e911c684b367fb8a4ce864399c813b2ad84f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_aliases.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from ..common import _aliases + +from .._internal import get_xp + +from ._info import __array_namespace_info__ + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from typing import Optional, Union + from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol + +import numpy as np +bool = np.bool_ + +# Basic renames +acos = np.arccos +acosh = np.arccosh +asin = np.arcsin +asinh = np.arcsinh +atan = np.arctan +atan2 = np.arctan2 +atanh = np.arctanh +bitwise_left_shift = np.left_shift +bitwise_invert = np.invert +bitwise_right_shift = np.right_shift +concat = np.concatenate +pow = np.power + +arange = get_xp(np)(_aliases.arange) +empty = get_xp(np)(_aliases.empty) +empty_like = get_xp(np)(_aliases.empty_like) +eye = get_xp(np)(_aliases.eye) +full = get_xp(np)(_aliases.full) +full_like = get_xp(np)(_aliases.full_like) +linspace = get_xp(np)(_aliases.linspace) +ones = get_xp(np)(_aliases.ones) +ones_like = get_xp(np)(_aliases.ones_like) +zeros = get_xp(np)(_aliases.zeros) +zeros_like = get_xp(np)(_aliases.zeros_like) +UniqueAllResult = get_xp(np)(_aliases.UniqueAllResult) +UniqueCountsResult = get_xp(np)(_aliases.UniqueCountsResult) +UniqueInverseResult = get_xp(np)(_aliases.UniqueInverseResult) +unique_all = get_xp(np)(_aliases.unique_all) +unique_counts = get_xp(np)(_aliases.unique_counts) +unique_inverse = get_xp(np)(_aliases.unique_inverse) +unique_values = get_xp(np)(_aliases.unique_values) +astype = _aliases.astype +std = get_xp(np)(_aliases.std) +var = get_xp(np)(_aliases.var) +cumulative_sum = get_xp(np)(_aliases.cumulative_sum) +clip = get_xp(np)(_aliases.clip) +permute_dims = get_xp(np)(_aliases.permute_dims) +reshape = get_xp(np)(_aliases.reshape) +argsort = get_xp(np)(_aliases.argsort) +sort = get_xp(np)(_aliases.sort) +nonzero = get_xp(np)(_aliases.nonzero) +ceil = get_xp(np)(_aliases.ceil) +floor = get_xp(np)(_aliases.floor) +trunc = get_xp(np)(_aliases.trunc) +matmul = get_xp(np)(_aliases.matmul) +matrix_transpose = get_xp(np)(_aliases.matrix_transpose) +tensordot = get_xp(np)(_aliases.tensordot) +sign = get_xp(np)(_aliases.sign) + +def _supports_buffer_protocol(obj): + try: + memoryview(obj) + except TypeError: + return False + return True + +# asarray also adds the copy keyword, which is not present in numpy 1.0. +# asarray() is different enough between numpy, cupy, and dask, the logic +# complicated enough that it's easier to define it separately for each module +# rather than trying to combine everything into one function in common/ +def asarray( + obj: Union[ + ndarray, + bool, + int, + float, + NestedSequence[bool | int | float], + SupportsBufferProtocol, + ], + /, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + copy: "Optional[Union[bool, np._CopyMode]]" = None, + **kwargs, +) -> ndarray: + """ + Array API compatibility wrapper for asarray(). + + See the corresponding documentation in the array library and/or the array API + specification for more details. + """ + if device not in ["cpu", None]: + raise ValueError(f"Unsupported device for NumPy: {device!r}") + + if hasattr(np, '_CopyMode'): + if copy is None: + copy = np._CopyMode.IF_NEEDED + elif copy is False: + copy = np._CopyMode.NEVER + elif copy is True: + copy = np._CopyMode.ALWAYS + else: + # Not present in older NumPys. In this case, we cannot really support + # copy=False. + if copy is False: + raise NotImplementedError("asarray(copy=False) requires a newer version of NumPy.") + + return np.array(obj, copy=copy, dtype=dtype, **kwargs) + +# These functions are completely new here. If the library already has them +# (i.e., numpy 2.0), use the library version instead of our wrapper. +if hasattr(np, 'vecdot'): + vecdot = np.vecdot +else: + vecdot = get_xp(np)(_aliases.vecdot) + +if hasattr(np, 'isdtype'): + isdtype = np.isdtype +else: + isdtype = get_xp(np)(_aliases.isdtype) + +if hasattr(np, 'unstack'): + unstack = np.unstack +else: + unstack = get_xp(np)(_aliases.unstack) + +__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'bool', + 'acos', 'acosh', 'asin', 'asinh', 'atan', + 'atan2', 'atanh', 'bitwise_left_shift', + 'bitwise_invert', 'bitwise_right_shift', + 'concat', 'pow'] + +_all_ignore = ['np', 'get_xp'] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_info.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_info.py new file mode 100644 index 0000000000000000000000000000000000000000..2a8c759dff034e92be0484d884889c59a2bfc338 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_info.py @@ -0,0 +1,346 @@ +""" +Array API Inspection namespace + +This is the namespace for inspection functions as defined by the array API +standard. See +https://data-apis.org/array-api/latest/API_specification/inspection.html for +more details. + +""" +from numpy import ( + dtype, + bool_ as bool, + intp, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + complex64, + complex128, +) + + +class __array_namespace_info__: + """ + Get the array API inspection namespace for NumPy. + + The array API inspection namespace defines the following functions: + + - capabilities() + - default_device() + - default_dtypes() + - dtypes() + - devices() + + See + https://data-apis.org/array-api/latest/API_specification/inspection.html + for more details. + + Returns + ------- + info : ModuleType + The array API inspection namespace for NumPy. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': numpy.float64, + 'complex floating': numpy.complex128, + 'integral': numpy.int64, + 'indexing': numpy.int64} + + """ + + __module__ = 'numpy' + + def capabilities(self): + """ + Return a dictionary of array API library capabilities. + + The resulting dictionary has the following keys: + + - **"boolean indexing"**: boolean indicating whether an array library + supports boolean indexing. Always ``True`` for NumPy. + + - **"data-dependent shapes"**: boolean indicating whether an array + library supports data-dependent output shapes. Always ``True`` for + NumPy. + + See + https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html + for more details. + + See Also + -------- + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + capabilities : dict + A dictionary of array API library capabilities. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.capabilities() + {'boolean indexing': True, + 'data-dependent shapes': True} + + """ + return { + "boolean indexing": True, + "data-dependent shapes": True, + # 'max rank' will be part of the 2024.12 standard + # "max rank": 64, + } + + def default_device(self): + """ + The default device used for new NumPy arrays. + + For NumPy, this always returns ``'cpu'``. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + device : str + The default device used for new NumPy arrays. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_device() + 'cpu' + + """ + return "cpu" + + def default_dtypes(self, *, device=None): + """ + The default data types used for new NumPy arrays. + + For NumPy, this always returns the following dictionary: + + - **"real floating"**: ``numpy.float64`` + - **"complex floating"**: ``numpy.complex128`` + - **"integral"**: ``numpy.intp`` + - **"indexing"**: ``numpy.intp`` + + Parameters + ---------- + device : str, optional + The device to get the default data types for. For NumPy, only + ``'cpu'`` is allowed. + + Returns + ------- + dtypes : dict + A dictionary describing the default data types used for new NumPy + arrays. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': numpy.float64, + 'complex floating': numpy.complex128, + 'integral': numpy.int64, + 'indexing': numpy.int64} + + """ + if device not in ["cpu", None]: + raise ValueError( + 'Device not understood. Only "cpu" is allowed, but received:' + f' {device}' + ) + return { + "real floating": dtype(float64), + "complex floating": dtype(complex128), + "integral": dtype(intp), + "indexing": dtype(intp), + } + + def dtypes(self, *, device=None, kind=None): + """ + The array API data types supported by NumPy. + + Note that this function only returns data types that are defined by + the array API. + + Parameters + ---------- + device : str, optional + The device to get the data types for. For NumPy, only ``'cpu'`` is + allowed. + kind : str or tuple of str, optional + The kind of data types to return. If ``None``, all data types are + returned. If a string, only data types of that kind are returned. + If a tuple, a dictionary containing the union of the given kinds + is returned. The following kinds are supported: + + - ``'bool'``: boolean data types (i.e., ``bool``). + - ``'signed integer'``: signed integer data types (i.e., ``int8``, + ``int16``, ``int32``, ``int64``). + - ``'unsigned integer'``: unsigned integer data types (i.e., + ``uint8``, ``uint16``, ``uint32``, ``uint64``). + - ``'integral'``: integer data types. Shorthand for ``('signed + integer', 'unsigned integer')``. + - ``'real floating'``: real-valued floating-point data types + (i.e., ``float32``, ``float64``). + - ``'complex floating'``: complex floating-point data types (i.e., + ``complex64``, ``complex128``). + - ``'numeric'``: numeric data types. Shorthand for ``('integral', + 'real floating', 'complex floating')``. + + Returns + ------- + dtypes : dict + A dictionary mapping the names of data types to the corresponding + NumPy data types. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.dtypes(kind='signed integer') + {'int8': numpy.int8, + 'int16': numpy.int16, + 'int32': numpy.int32, + 'int64': numpy.int64} + + """ + if device not in ["cpu", None]: + raise ValueError( + 'Device not understood. Only "cpu" is allowed, but received:' + f' {device}' + ) + if kind is None: + return { + "bool": dtype(bool), + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + "float32": dtype(float32), + "float64": dtype(float64), + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if kind == "bool": + return {"bool": bool} + if kind == "signed integer": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + } + if kind == "unsigned integer": + return { + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + } + if kind == "integral": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + } + if kind == "real floating": + return { + "float32": dtype(float32), + "float64": dtype(float64), + } + if kind == "complex floating": + return { + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if kind == "numeric": + return { + "int8": dtype(int8), + "int16": dtype(int16), + "int32": dtype(int32), + "int64": dtype(int64), + "uint8": dtype(uint8), + "uint16": dtype(uint16), + "uint32": dtype(uint32), + "uint64": dtype(uint64), + "float32": dtype(float32), + "float64": dtype(float64), + "complex64": dtype(complex64), + "complex128": dtype(complex128), + } + if isinstance(kind, tuple): + res = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + def devices(self): + """ + The devices supported by NumPy. + + For NumPy, this always returns ``['cpu']``. + + Returns + ------- + devices : list of str + The devices supported by NumPy. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.devices() + ['cpu'] + + """ + return ["cpu"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_typing.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..5d0ea57b6adf252204890253ab36da9416a57a55 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/_typing.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +__all__ = [ + "ndarray", + "Device", + "Dtype", +] + +import sys +from typing import ( + Literal, + Union, + TYPE_CHECKING, +) + +from numpy import ( + ndarray, + dtype, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, +) + +Device = Literal["cpu"] +if TYPE_CHECKING or sys.version_info >= (3, 9): + Dtype = dtype[Union[ + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + ]] +else: + Dtype = dtype diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/fft.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/fft.py new file mode 100644 index 0000000000000000000000000000000000000000..1e8f0aa7960464212bdff643e4551db2aa0657d0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/fft.py @@ -0,0 +1,29 @@ +from numpy.fft import * # noqa: F403 +from numpy.fft import __all__ as fft_all + +from ..common import _fft +from .._internal import get_xp + +import numpy as np + +fft = get_xp(np)(_fft.fft) +ifft = get_xp(np)(_fft.ifft) +fftn = get_xp(np)(_fft.fftn) +ifftn = get_xp(np)(_fft.ifftn) +rfft = get_xp(np)(_fft.rfft) +irfft = get_xp(np)(_fft.irfft) +rfftn = get_xp(np)(_fft.rfftn) +irfftn = get_xp(np)(_fft.irfftn) +hfft = get_xp(np)(_fft.hfft) +ihfft = get_xp(np)(_fft.ihfft) +fftfreq = get_xp(np)(_fft.fftfreq) +rfftfreq = get_xp(np)(_fft.rfftfreq) +fftshift = get_xp(np)(_fft.fftshift) +ifftshift = get_xp(np)(_fft.ifftshift) + +__all__ = fft_all + _fft.__all__ + +del get_xp +del np +del fft_all +del _fft diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/linalg.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/linalg.py new file mode 100644 index 0000000000000000000000000000000000000000..1775d4d9310deddce2e9f777b68d5ba869718ee3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/numpy/linalg.py @@ -0,0 +1,90 @@ +from numpy.linalg import * # noqa: F403 +from numpy.linalg import __all__ as linalg_all +import numpy as _np + +from ..common import _linalg +from .._internal import get_xp + +# These functions are in both the main and linalg namespaces +from ._aliases import matmul, matrix_transpose, tensordot, vecdot # noqa: F401 + +import numpy as np + +cross = get_xp(np)(_linalg.cross) +outer = get_xp(np)(_linalg.outer) +EighResult = _linalg.EighResult +QRResult = _linalg.QRResult +SlogdetResult = _linalg.SlogdetResult +SVDResult = _linalg.SVDResult +eigh = get_xp(np)(_linalg.eigh) +qr = get_xp(np)(_linalg.qr) +slogdet = get_xp(np)(_linalg.slogdet) +svd = get_xp(np)(_linalg.svd) +cholesky = get_xp(np)(_linalg.cholesky) +matrix_rank = get_xp(np)(_linalg.matrix_rank) +pinv = get_xp(np)(_linalg.pinv) +matrix_norm = get_xp(np)(_linalg.matrix_norm) +svdvals = get_xp(np)(_linalg.svdvals) +diagonal = get_xp(np)(_linalg.diagonal) +trace = get_xp(np)(_linalg.trace) + +# Note: unlike np.linalg.solve, the array API solve() only accepts x2 as a +# vector when it is exactly 1-dimensional. All other cases treat x2 as a stack +# of matrices. The np.linalg.solve behavior of allowing stacks of both +# matrices and vectors is ambiguous c.f. +# https://github.com/numpy/numpy/issues/15349 and +# https://github.com/data-apis/array-api/issues/285. + +# To workaround this, the below is the code from np.linalg.solve except +# only calling solve1 in the exactly 1D case. + +# This code is here instead of in common because it is numpy specific. Also +# note that CuPy's solve() does not currently support broadcasting (see +# https://github.com/cupy/cupy/blob/main/cupy/cublas.py#L43). +def solve(x1: _np.ndarray, x2: _np.ndarray, /) -> _np.ndarray: + try: + from numpy.linalg._linalg import ( + _makearray, _assert_stacked_2d, _assert_stacked_square, + _commonType, isComplexType, _raise_linalgerror_singular + ) + except ImportError: + from numpy.linalg.linalg import ( + _makearray, _assert_stacked_2d, _assert_stacked_square, + _commonType, isComplexType, _raise_linalgerror_singular + ) + from numpy.linalg import _umath_linalg + + x1, _ = _makearray(x1) + _assert_stacked_2d(x1) + _assert_stacked_square(x1) + x2, wrap = _makearray(x2) + t, result_t = _commonType(x1, x2) + + # This part is different from np.linalg.solve + if x2.ndim == 1: + gufunc = _umath_linalg.solve1 + else: + gufunc = _umath_linalg.solve + + # This does nothing currently but is left in because it will be relevant + # when complex dtype support is added to the spec in 2022. + signature = 'DD->D' if isComplexType(t) else 'dd->d' + with _np.errstate(call=_raise_linalgerror_singular, invalid='call', + over='ignore', divide='ignore', under='ignore'): + r = gufunc(x1, x2, signature=signature) + + return wrap(r.astype(result_t, copy=False)) + +# These functions are completely new here. If the library already has them +# (i.e., numpy 2.0), use the library version instead of our wrapper. +if hasattr(np.linalg, 'vector_norm'): + vector_norm = np.linalg.vector_norm +else: + vector_norm = get_xp(np)(_linalg.vector_norm) + +__all__ = linalg_all + _linalg.__all__ + ['solve'] + +del get_xp +del np +del linalg_all +del _linalg diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d494b05eee0f23ad504435bc35fcfd722e3ca323 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/__init__.py @@ -0,0 +1,24 @@ +from torch import * # noqa: F403 + +# Several names are not included in the above import * +import torch +for n in dir(torch): + if (n.startswith('_') + or n.endswith('_') + or 'cuda' in n + or 'cpu' in n + or 'backward' in n): + continue + exec(n + ' = torch.' + n) + +# These imports may overwrite names from the import * above. +from ._aliases import * # noqa: F403 + +# See the comment in the numpy __init__.py +__import__(__package__ + '.linalg') + +__import__(__package__ + '.fft') + +from ..common._helpers import * # noqa: F403 + +__array_api_version__ = '2023.12' diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/_aliases.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..7c1e79a3acdf758cb0bff527ffe22e5d67a7f725 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/_aliases.py @@ -0,0 +1,752 @@ +from __future__ import annotations + +from functools import wraps as _wraps +from builtins import all as _builtin_all, any as _builtin_any + +from ..common._aliases import (matrix_transpose as _aliases_matrix_transpose, + vecdot as _aliases_vecdot, + clip as _aliases_clip, + unstack as _aliases_unstack, + cumulative_sum as _aliases_cumulative_sum, + ) +from .._internal import get_xp + +from ._info import __array_namespace_info__ + +import torch + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from typing import List, Optional, Sequence, Tuple, Union + from ..common._typing import Device + from torch import dtype as Dtype + + array = torch.Tensor + +_int_dtypes = { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, +} + +_array_api_dtypes = { + torch.bool, + *_int_dtypes, + torch.float32, + torch.float64, + torch.complex64, + torch.complex128, +} + +_promotion_table = { + # bool + (torch.bool, torch.bool): torch.bool, + # ints + (torch.int8, torch.int8): torch.int8, + (torch.int8, torch.int16): torch.int16, + (torch.int8, torch.int32): torch.int32, + (torch.int8, torch.int64): torch.int64, + (torch.int16, torch.int8): torch.int16, + (torch.int16, torch.int16): torch.int16, + (torch.int16, torch.int32): torch.int32, + (torch.int16, torch.int64): torch.int64, + (torch.int32, torch.int8): torch.int32, + (torch.int32, torch.int16): torch.int32, + (torch.int32, torch.int32): torch.int32, + (torch.int32, torch.int64): torch.int64, + (torch.int64, torch.int8): torch.int64, + (torch.int64, torch.int16): torch.int64, + (torch.int64, torch.int32): torch.int64, + (torch.int64, torch.int64): torch.int64, + # uints + (torch.uint8, torch.uint8): torch.uint8, + # ints and uints (mixed sign) + (torch.int8, torch.uint8): torch.int16, + (torch.int16, torch.uint8): torch.int16, + (torch.int32, torch.uint8): torch.int32, + (torch.int64, torch.uint8): torch.int64, + (torch.uint8, torch.int8): torch.int16, + (torch.uint8, torch.int16): torch.int16, + (torch.uint8, torch.int32): torch.int32, + (torch.uint8, torch.int64): torch.int64, + # floats + (torch.float32, torch.float32): torch.float32, + (torch.float32, torch.float64): torch.float64, + (torch.float64, torch.float32): torch.float64, + (torch.float64, torch.float64): torch.float64, + # complexes + (torch.complex64, torch.complex64): torch.complex64, + (torch.complex64, torch.complex128): torch.complex128, + (torch.complex128, torch.complex64): torch.complex128, + (torch.complex128, torch.complex128): torch.complex128, + # Mixed float and complex + (torch.float32, torch.complex64): torch.complex64, + (torch.float32, torch.complex128): torch.complex128, + (torch.float64, torch.complex64): torch.complex128, + (torch.float64, torch.complex128): torch.complex128, +} + + +def _two_arg(f): + @_wraps(f) + def _f(x1, x2, /, **kwargs): + x1, x2 = _fix_promotion(x1, x2) + return f(x1, x2, **kwargs) + if _f.__doc__ is None: + _f.__doc__ = f"""\ +Array API compatibility wrapper for torch.{f.__name__}. + +See the corresponding PyTorch documentation and/or the array API specification +for more details. + +""" + return _f + +def _fix_promotion(x1, x2, only_scalar=True): + if not isinstance(x1, torch.Tensor) or not isinstance(x2, torch.Tensor): + return x1, x2 + if x1.dtype not in _array_api_dtypes or x2.dtype not in _array_api_dtypes: + return x1, x2 + # If an argument is 0-D pytorch downcasts the other argument + if not only_scalar or x1.shape == (): + dtype = result_type(x1, x2) + x2 = x2.to(dtype) + if not only_scalar or x2.shape == (): + dtype = result_type(x1, x2) + x1 = x1.to(dtype) + return x1, x2 + +def result_type(*arrays_and_dtypes: Union[array, Dtype]) -> Dtype: + if len(arrays_and_dtypes) == 0: + raise TypeError("At least one array or dtype must be provided") + if len(arrays_and_dtypes) == 1: + x = arrays_and_dtypes[0] + if isinstance(x, torch.dtype): + return x + return x.dtype + if len(arrays_and_dtypes) > 2: + return result_type(arrays_and_dtypes[0], result_type(*arrays_and_dtypes[1:])) + + x, y = arrays_and_dtypes + xdt = x.dtype if not isinstance(x, torch.dtype) else x + ydt = y.dtype if not isinstance(y, torch.dtype) else y + + if (xdt, ydt) in _promotion_table: + return _promotion_table[xdt, ydt] + + # This doesn't result_type(dtype, dtype) for non-array API dtypes + # because torch.result_type only accepts tensors. This does however, allow + # cross-kind promotion. + x = torch.tensor([], dtype=x) if isinstance(x, torch.dtype) else x + y = torch.tensor([], dtype=y) if isinstance(y, torch.dtype) else y + return torch.result_type(x, y) + +def can_cast(from_: Union[Dtype, array], to: Dtype, /) -> bool: + if not isinstance(from_, torch.dtype): + from_ = from_.dtype + return torch.can_cast(from_, to) + +# Basic renames +bitwise_invert = torch.bitwise_not +newaxis = None +# torch.conj sets the conjugation bit, which breaks conversion to other +# libraries. See https://github.com/data-apis/array-api-compat/issues/173 +conj = torch.conj_physical + +# Two-arg elementwise functions +# These require a wrapper to do the correct type promotion on 0-D tensors +add = _two_arg(torch.add) +atan2 = _two_arg(torch.atan2) +bitwise_and = _two_arg(torch.bitwise_and) +bitwise_left_shift = _two_arg(torch.bitwise_left_shift) +bitwise_or = _two_arg(torch.bitwise_or) +bitwise_right_shift = _two_arg(torch.bitwise_right_shift) +bitwise_xor = _two_arg(torch.bitwise_xor) +copysign = _two_arg(torch.copysign) +divide = _two_arg(torch.divide) +# Also a rename. torch.equal does not broadcast +equal = _two_arg(torch.eq) +floor_divide = _two_arg(torch.floor_divide) +greater = _two_arg(torch.greater) +greater_equal = _two_arg(torch.greater_equal) +hypot = _two_arg(torch.hypot) +less = _two_arg(torch.less) +less_equal = _two_arg(torch.less_equal) +logaddexp = _two_arg(torch.logaddexp) +# logical functions are not included here because they only accept bool in the +# spec, so type promotion is irrelevant. +maximum = _two_arg(torch.maximum) +minimum = _two_arg(torch.minimum) +multiply = _two_arg(torch.multiply) +not_equal = _two_arg(torch.not_equal) +pow = _two_arg(torch.pow) +remainder = _two_arg(torch.remainder) +subtract = _two_arg(torch.subtract) + +# These wrappers are mostly based on the fact that pytorch uses 'dim' instead +# of 'axis'. + +# torch.min and torch.max return a tuple and don't support multiple axes https://github.com/pytorch/pytorch/issues/58745 +def max(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array: + # https://github.com/pytorch/pytorch/issues/29137 + if axis == (): + return torch.clone(x) + return torch.amax(x, axis, keepdims=keepdims) + +def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array: + # https://github.com/pytorch/pytorch/issues/29137 + if axis == (): + return torch.clone(x) + return torch.amin(x, axis, keepdims=keepdims) + +clip = get_xp(torch)(_aliases_clip) +unstack = get_xp(torch)(_aliases_unstack) +cumulative_sum = get_xp(torch)(_aliases_cumulative_sum) + +# torch.sort also returns a tuple +# https://github.com/pytorch/pytorch/issues/70921 +def sort(x: array, /, *, axis: int = -1, descending: bool = False, stable: bool = True, **kwargs) -> array: + return torch.sort(x, dim=axis, descending=descending, stable=stable, **kwargs).values + +def _normalize_axes(axis, ndim): + axes = [] + if ndim == 0 and axis: + # Better error message in this case + raise IndexError(f"Dimension out of range: {axis[0]}") + lower, upper = -ndim, ndim - 1 + for a in axis: + if a < lower or a > upper: + # Match torch error message (e.g., from sum()) + raise IndexError(f"Dimension out of range (expected to be in range of [{lower}, {upper}], but got {a}") + if a < 0: + a = a + ndim + if a in axes: + # Use IndexError instead of RuntimeError, and "axis" instead of "dim" + raise IndexError(f"Axis {a} appears multiple times in the list of axes") + axes.append(a) + return sorted(axes) + +def _axis_none_keepdims(x, ndim, keepdims): + # Apply keepdims when axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + # Note that this is only valid for the axis=None case. + if keepdims: + for i in range(ndim): + x = torch.unsqueeze(x, 0) + return x + +def _reduce_multiple_axes(f, x, axis, keepdims=False, **kwargs): + # Some reductions don't support multiple axes + # (https://github.com/pytorch/pytorch/issues/56586). + axes = _normalize_axes(axis, x.ndim) + for a in reversed(axes): + x = torch.movedim(x, a, -1) + x = torch.flatten(x, -len(axes)) + + out = f(x, -1, **kwargs) + + if keepdims: + for a in axes: + out = torch.unsqueeze(out, a) + return out + +def prod(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + dtype: Optional[Dtype] = None, + keepdims: bool = False, + **kwargs) -> array: + x = torch.asarray(x) + ndim = x.ndim + + # https://github.com/pytorch/pytorch/issues/29137. Separate from the logic + # below because it still needs to upcast. + if axis == (): + if dtype is None: + # We can't upcast uint8 according to the spec because there is no + # torch.uint64, so at least upcast to int64 which is what sum does + # when axis=None. + if x.dtype in [torch.int8, torch.int16, torch.int32, torch.uint8]: + return x.to(torch.int64) + return x.clone() + return x.to(dtype) + + # torch.prod doesn't support multiple axes + # (https://github.com/pytorch/pytorch/issues/56586). + if isinstance(axis, tuple): + return _reduce_multiple_axes(torch.prod, x, axis, keepdims=keepdims, dtype=dtype, **kwargs) + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.prod(x, dtype=dtype, **kwargs) + res = _axis_none_keepdims(res, ndim, keepdims) + return res + + return torch.prod(x, axis, dtype=dtype, keepdims=keepdims, **kwargs) + + +def sum(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + dtype: Optional[Dtype] = None, + keepdims: bool = False, + **kwargs) -> array: + x = torch.asarray(x) + ndim = x.ndim + + # https://github.com/pytorch/pytorch/issues/29137. + # Make sure it upcasts. + if axis == (): + if dtype is None: + # We can't upcast uint8 according to the spec because there is no + # torch.uint64, so at least upcast to int64 which is what sum does + # when axis=None. + if x.dtype in [torch.int8, torch.int16, torch.int32, torch.uint8]: + return x.to(torch.int64) + return x.clone() + return x.to(dtype) + + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.sum(x, dtype=dtype, **kwargs) + res = _axis_none_keepdims(res, ndim, keepdims) + return res + + return torch.sum(x, axis, dtype=dtype, keepdims=keepdims, **kwargs) + +def any(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, + **kwargs) -> array: + x = torch.asarray(x) + ndim = x.ndim + if axis == (): + return x.to(torch.bool) + # torch.any doesn't support multiple axes + # (https://github.com/pytorch/pytorch/issues/56586). + if isinstance(axis, tuple): + res = _reduce_multiple_axes(torch.any, x, axis, keepdims=keepdims, **kwargs) + return res.to(torch.bool) + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.any(x, **kwargs) + res = _axis_none_keepdims(res, ndim, keepdims) + return res.to(torch.bool) + + # torch.any doesn't return bool for uint8 + return torch.any(x, axis, keepdims=keepdims).to(torch.bool) + +def all(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, + **kwargs) -> array: + x = torch.asarray(x) + ndim = x.ndim + if axis == (): + return x.to(torch.bool) + # torch.all doesn't support multiple axes + # (https://github.com/pytorch/pytorch/issues/56586). + if isinstance(axis, tuple): + res = _reduce_multiple_axes(torch.all, x, axis, keepdims=keepdims, **kwargs) + return res.to(torch.bool) + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.all(x, **kwargs) + res = _axis_none_keepdims(res, ndim, keepdims) + return res.to(torch.bool) + + # torch.all doesn't return bool for uint8 + return torch.all(x, axis, keepdims=keepdims).to(torch.bool) + +def mean(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, + **kwargs) -> array: + # https://github.com/pytorch/pytorch/issues/29137 + if axis == (): + return torch.clone(x) + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.mean(x, **kwargs) + res = _axis_none_keepdims(res, x.ndim, keepdims) + return res + return torch.mean(x, axis, keepdims=keepdims, **kwargs) + +def std(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, + **kwargs) -> array: + # Note, float correction is not supported + # https://github.com/pytorch/pytorch/issues/61492. We don't try to + # implement it here for now. + + if isinstance(correction, float): + _correction = int(correction) + if correction != _correction: + raise NotImplementedError("float correction in torch std() is not yet supported") + else: + _correction = correction + + # https://github.com/pytorch/pytorch/issues/29137 + if axis == (): + return torch.zeros_like(x) + if isinstance(axis, int): + axis = (axis,) + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.std(x, tuple(range(x.ndim)), correction=_correction, **kwargs) + res = _axis_none_keepdims(res, x.ndim, keepdims) + return res + return torch.std(x, axis, correction=_correction, keepdims=keepdims, **kwargs) + +def var(x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, + **kwargs) -> array: + # Note, float correction is not supported + # https://github.com/pytorch/pytorch/issues/61492. We don't try to + # implement it here for now. + + # if isinstance(correction, float): + # correction = int(correction) + + # https://github.com/pytorch/pytorch/issues/29137 + if axis == (): + return torch.zeros_like(x) + if isinstance(axis, int): + axis = (axis,) + if axis is None: + # torch doesn't support keepdims with axis=None + # (https://github.com/pytorch/pytorch/issues/71209) + res = torch.var(x, tuple(range(x.ndim)), correction=correction, **kwargs) + res = _axis_none_keepdims(res, x.ndim, keepdims) + return res + return torch.var(x, axis, correction=correction, keepdims=keepdims, **kwargs) + +# torch.concat doesn't support dim=None +# https://github.com/pytorch/pytorch/issues/70925 +def concat(arrays: Union[Tuple[array, ...], List[array]], + /, + *, + axis: Optional[int] = 0, + **kwargs) -> array: + if axis is None: + arrays = tuple(ar.flatten() for ar in arrays) + axis = 0 + return torch.concat(arrays, axis, **kwargs) + +# torch.squeeze only accepts int dim and doesn't require it +# https://github.com/pytorch/pytorch/issues/70924. Support for tuple dim was +# added at https://github.com/pytorch/pytorch/pull/89017. +def squeeze(x: array, /, axis: Union[int, Tuple[int, ...]]) -> array: + if isinstance(axis, int): + axis = (axis,) + for a in axis: + if x.shape[a] != 1: + raise ValueError("squeezed dimensions must be equal to 1") + axes = _normalize_axes(axis, x.ndim) + # Remove this once pytorch 1.14 is released with the above PR #89017. + sequence = [a - i for i, a in enumerate(axes)] + for a in sequence: + x = torch.squeeze(x, a) + return x + +# torch.broadcast_to uses size instead of shape +def broadcast_to(x: array, /, shape: Tuple[int, ...], **kwargs) -> array: + return torch.broadcast_to(x, shape, **kwargs) + +# torch.permute uses dims instead of axes +def permute_dims(x: array, /, axes: Tuple[int, ...]) -> array: + return torch.permute(x, axes) + +# The axis parameter doesn't work for flip() and roll() +# https://github.com/pytorch/pytorch/issues/71210. Also torch.flip() doesn't +# accept axis=None +def flip(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, **kwargs) -> array: + if axis is None: + axis = tuple(range(x.ndim)) + # torch.flip doesn't accept dim as an int but the method does + # https://github.com/pytorch/pytorch/issues/18095 + return x.flip(axis, **kwargs) + +def roll(x: array, /, shift: Union[int, Tuple[int, ...]], *, axis: Optional[Union[int, Tuple[int, ...]]] = None, **kwargs) -> array: + return torch.roll(x, shift, axis, **kwargs) + +def nonzero(x: array, /, **kwargs) -> Tuple[array, ...]: + if x.ndim == 0: + raise ValueError("nonzero() does not support zero-dimensional arrays") + return torch.nonzero(x, as_tuple=True, **kwargs) + +def where(condition: array, x1: array, x2: array, /) -> array: + x1, x2 = _fix_promotion(x1, x2) + return torch.where(condition, x1, x2) + +# torch.reshape doesn't have the copy keyword +def reshape(x: array, + /, + shape: Tuple[int, ...], + copy: Optional[bool] = None, + **kwargs) -> array: + if copy is not None: + raise NotImplementedError("torch.reshape doesn't yet support the copy keyword") + return torch.reshape(x, shape, **kwargs) + +# torch.arange doesn't support returning empty arrays +# (https://github.com/pytorch/pytorch/issues/70915), and doesn't support some +# keyword argument combinations +# (https://github.com/pytorch/pytorch/issues/70914) +def arange(start: Union[int, float], + /, + stop: Optional[Union[int, float]] = None, + step: Union[int, float] = 1, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs) -> array: + if stop is None: + start, stop = 0, start + if step > 0 and stop <= start or step < 0 and stop >= start: + if dtype is None: + if _builtin_all(isinstance(i, int) for i in [start, stop, step]): + dtype = torch.int64 + else: + dtype = torch.float32 + return torch.empty(0, dtype=dtype, device=device, **kwargs) + return torch.arange(start, stop, step, dtype=dtype, device=device, **kwargs) + +# torch.eye does not accept None as a default for the second argument and +# doesn't support off-diagonals (https://github.com/pytorch/pytorch/issues/70910) +def eye(n_rows: int, + n_cols: Optional[int] = None, + /, + *, + k: int = 0, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs) -> array: + if n_cols is None: + n_cols = n_rows + z = torch.zeros(n_rows, n_cols, dtype=dtype, device=device, **kwargs) + if abs(k) <= n_rows + n_cols: + z.diagonal(k).fill_(1) + return z + +# torch.linspace doesn't have the endpoint parameter +def linspace(start: Union[int, float], + stop: Union[int, float], + /, + num: int, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + endpoint: bool = True, + **kwargs) -> array: + if not endpoint: + return torch.linspace(start, stop, num+1, dtype=dtype, device=device, **kwargs)[:-1] + return torch.linspace(start, stop, num, dtype=dtype, device=device, **kwargs) + +# torch.full does not accept an int size +# https://github.com/pytorch/pytorch/issues/70906 +def full(shape: Union[int, Tuple[int, ...]], + fill_value: Union[bool, int, float, complex], + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs) -> array: + if isinstance(shape, int): + shape = (shape,) + + return torch.full(shape, fill_value, dtype=dtype, device=device, **kwargs) + +# ones, zeros, and empty do not accept shape as a keyword argument +def ones(shape: Union[int, Tuple[int, ...]], + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs) -> array: + return torch.ones(shape, dtype=dtype, device=device, **kwargs) + +def zeros(shape: Union[int, Tuple[int, ...]], + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs) -> array: + return torch.zeros(shape, dtype=dtype, device=device, **kwargs) + +def empty(shape: Union[int, Tuple[int, ...]], + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + **kwargs) -> array: + return torch.empty(shape, dtype=dtype, device=device, **kwargs) + +# tril and triu do not call the keyword argument k + +def tril(x: array, /, *, k: int = 0) -> array: + return torch.tril(x, k) + +def triu(x: array, /, *, k: int = 0) -> array: + return torch.triu(x, k) + +# Functions that aren't in torch https://github.com/pytorch/pytorch/issues/58742 +def expand_dims(x: array, /, *, axis: int = 0) -> array: + return torch.unsqueeze(x, axis) + +def astype(x: array, dtype: Dtype, /, *, copy: bool = True) -> array: + return x.to(dtype, copy=copy) + +def broadcast_arrays(*arrays: array) -> List[array]: + shape = torch.broadcast_shapes(*[a.shape for a in arrays]) + return [torch.broadcast_to(a, shape) for a in arrays] + +# Note that these named tuples aren't actually part of the standard namespace, +# but I don't see any issue with exporting the names here regardless. +from ..common._aliases import (UniqueAllResult, UniqueCountsResult, + UniqueInverseResult) + +# https://github.com/pytorch/pytorch/issues/70920 +def unique_all(x: array) -> UniqueAllResult: + # torch.unique doesn't support returning indices. + # https://github.com/pytorch/pytorch/issues/36748. The workaround + # suggested in that issue doesn't actually function correctly (it relies + # on non-deterministic behavior of scatter()). + raise NotImplementedError("unique_all() not yet implemented for pytorch (see https://github.com/pytorch/pytorch/issues/36748)") + + # values, inverse_indices, counts = torch.unique(x, return_counts=True, return_inverse=True) + # # torch.unique incorrectly gives a 0 count for nan values. + # # https://github.com/pytorch/pytorch/issues/94106 + # counts[torch.isnan(values)] = 1 + # return UniqueAllResult(values, indices, inverse_indices, counts) + +def unique_counts(x: array) -> UniqueCountsResult: + values, counts = torch.unique(x, return_counts=True) + + # torch.unique incorrectly gives a 0 count for nan values. + # https://github.com/pytorch/pytorch/issues/94106 + counts[torch.isnan(values)] = 1 + return UniqueCountsResult(values, counts) + +def unique_inverse(x: array) -> UniqueInverseResult: + values, inverse = torch.unique(x, return_inverse=True) + return UniqueInverseResult(values, inverse) + +def unique_values(x: array) -> array: + return torch.unique(x) + +def matmul(x1: array, x2: array, /, **kwargs) -> array: + # torch.matmul doesn't type promote (but differently from _fix_promotion) + x1, x2 = _fix_promotion(x1, x2, only_scalar=False) + return torch.matmul(x1, x2, **kwargs) + +matrix_transpose = get_xp(torch)(_aliases_matrix_transpose) +_vecdot = get_xp(torch)(_aliases_vecdot) + +def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array: + x1, x2 = _fix_promotion(x1, x2, only_scalar=False) + return _vecdot(x1, x2, axis=axis) + +# torch.tensordot uses dims instead of axes +def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2, **kwargs) -> array: + # Note: torch.tensordot fails with integer dtypes when there is only 1 + # element in the axis (https://github.com/pytorch/pytorch/issues/84530). + x1, x2 = _fix_promotion(x1, x2, only_scalar=False) + return torch.tensordot(x1, x2, dims=axes, **kwargs) + + +def isdtype( + dtype: Dtype, kind: Union[Dtype, str, Tuple[Union[Dtype, str], ...]], + *, _tuple=True, # Disallow nested tuples +) -> bool: + """ + Returns a boolean indicating whether a provided dtype is of a specified data type ``kind``. + + Note that outside of this function, this compat library does not yet fully + support complex numbers. + + See + https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html + for more details + """ + if isinstance(kind, tuple) and _tuple: + return _builtin_any(isdtype(dtype, k, _tuple=False) for k in kind) + elif isinstance(kind, str): + if kind == 'bool': + return dtype == torch.bool + elif kind == 'signed integer': + return dtype in _int_dtypes and dtype.is_signed + elif kind == 'unsigned integer': + return dtype in _int_dtypes and not dtype.is_signed + elif kind == 'integral': + return dtype in _int_dtypes + elif kind == 'real floating': + return dtype.is_floating_point + elif kind == 'complex floating': + return dtype.is_complex + elif kind == 'numeric': + return isdtype(dtype, ('integral', 'real floating', 'complex floating')) + else: + raise ValueError(f"Unrecognized data type kind: {kind!r}") + else: + return dtype == kind + +def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -> array: + if axis is None: + if x.ndim != 1: + raise ValueError("axis must be specified when ndim > 1") + axis = 0 + return torch.index_select(x, axis, indices, **kwargs) + +def sign(x: array, /) -> array: + # torch sign() does not support complex numbers and does not propagate + # nans. See https://github.com/data-apis/array-api-compat/issues/136 + if x.dtype.is_complex: + out = x/torch.abs(x) + # sign(0) = 0 but the above formula would give nan + out[x == 0+0j] = 0+0j + return out + else: + out = torch.sign(x) + if x.dtype.is_floating_point: + out[torch.isnan(x)] = torch.nan + return out + + +__all__ = ['__array_namespace_info__', 'result_type', 'can_cast', + 'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add', + 'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or', + 'bitwise_right_shift', 'bitwise_xor', 'copysign', 'divide', + 'equal', 'floor_divide', 'greater', 'greater_equal', 'hypot', + 'less', 'less_equal', 'logaddexp', 'maximum', 'minimum', + 'multiply', 'not_equal', 'pow', 'remainder', 'subtract', 'max', + 'min', 'clip', 'unstack', 'cumulative_sum', 'sort', 'prod', 'sum', + 'any', 'all', 'mean', 'std', 'var', 'concat', 'squeeze', + 'broadcast_to', 'flip', 'roll', 'nonzero', 'where', 'reshape', + 'arange', 'eye', 'linspace', 'full', 'ones', 'zeros', 'empty', + 'tril', 'triu', 'expand_dims', 'astype', 'broadcast_arrays', + 'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult', + 'unique_all', 'unique_counts', 'unique_inverse', 'unique_values', + 'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype', + 'take', 'sign'] + +_all_ignore = ['torch', 'get_xp'] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/_info.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/_info.py new file mode 100644 index 0000000000000000000000000000000000000000..5614ef6d5bb6118728f28e165fb8804f332b87e2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/_info.py @@ -0,0 +1,358 @@ +""" +Array API Inspection namespace + +This is the namespace for inspection functions as defined by the array API +standard. See +https://data-apis.org/array-api/latest/API_specification/inspection.html for +more details. + +""" +import torch + +from functools import cache + +class __array_namespace_info__: + """ + Get the array API inspection namespace for PyTorch. + + The array API inspection namespace defines the following functions: + + - capabilities() + - default_device() + - default_dtypes() + - dtypes() + - devices() + + See + https://data-apis.org/array-api/latest/API_specification/inspection.html + for more details. + + Returns + ------- + info : ModuleType + The array API inspection namespace for PyTorch. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': numpy.float64, + 'complex floating': numpy.complex128, + 'integral': numpy.int64, + 'indexing': numpy.int64} + + """ + + __module__ = 'torch' + + def capabilities(self): + """ + Return a dictionary of array API library capabilities. + + The resulting dictionary has the following keys: + + - **"boolean indexing"**: boolean indicating whether an array library + supports boolean indexing. Always ``True`` for PyTorch. + + - **"data-dependent shapes"**: boolean indicating whether an array + library supports data-dependent output shapes. Always ``True`` for + PyTorch. + + See + https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html + for more details. + + See Also + -------- + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + capabilities : dict + A dictionary of array API library capabilities. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.capabilities() + {'boolean indexing': True, + 'data-dependent shapes': True} + + """ + return { + "boolean indexing": True, + "data-dependent shapes": True, + # 'max rank' will be part of the 2024.12 standard + # "max rank": 64, + } + + def default_device(self): + """ + The default device used for new PyTorch arrays. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Returns + ------- + device : str + The default device used for new PyTorch arrays. + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_device() + 'cpu' + + """ + return torch.device("cpu") + + def default_dtypes(self, *, device=None): + """ + The default data types used for new PyTorch arrays. + + Parameters + ---------- + device : str, optional + The device to get the default data types for. For PyTorch, only + ``'cpu'`` is allowed. + + Returns + ------- + dtypes : dict + A dictionary describing the default data types used for new PyTorch + arrays. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.default_dtypes() + {'real floating': torch.float32, + 'complex floating': torch.complex64, + 'integral': torch.int64, + 'indexing': torch.int64} + + """ + # Note: if the default is set to float64, the devices like MPS that + # don't support float64 will error. We still return the default_dtype + # value here because this error doesn't represent a different default + # per-device. + default_floating = torch.get_default_dtype() + default_complex = torch.complex64 if default_floating == torch.float32 else torch.complex128 + default_integral = torch.int64 + return { + "real floating": default_floating, + "complex floating": default_complex, + "integral": default_integral, + "indexing": default_integral, + } + + + def _dtypes(self, kind): + bool = torch.bool + int8 = torch.int8 + int16 = torch.int16 + int32 = torch.int32 + int64 = torch.int64 + uint8 = torch.uint8 + # uint16, uint32, and uint64 are present in newer versions of pytorch, + # but they aren't generally supported by the array API functions, so + # we omit them from this function. + float32 = torch.float32 + float64 = torch.float64 + complex64 = torch.complex64 + complex128 = torch.complex128 + + if kind is None: + return { + "bool": bool, + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + "float32": float32, + "float64": float64, + "complex64": complex64, + "complex128": complex128, + } + if kind == "bool": + return {"bool": bool} + if kind == "signed integer": + return { + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + } + if kind == "unsigned integer": + return { + "uint8": uint8, + } + if kind == "integral": + return { + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + } + if kind == "real floating": + return { + "float32": float32, + "float64": float64, + } + if kind == "complex floating": + return { + "complex64": complex64, + "complex128": complex128, + } + if kind == "numeric": + return { + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + "float32": float32, + "float64": float64, + "complex64": complex64, + "complex128": complex128, + } + if isinstance(kind, tuple): + res = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + @cache + def dtypes(self, *, device=None, kind=None): + """ + The array API data types supported by PyTorch. + + Note that this function only returns data types that are defined by + the array API. + + Parameters + ---------- + device : str, optional + The device to get the data types for. + kind : str or tuple of str, optional + The kind of data types to return. If ``None``, all data types are + returned. If a string, only data types of that kind are returned. + If a tuple, a dictionary containing the union of the given kinds + is returned. The following kinds are supported: + + - ``'bool'``: boolean data types (i.e., ``bool``). + - ``'signed integer'``: signed integer data types (i.e., ``int8``, + ``int16``, ``int32``, ``int64``). + - ``'unsigned integer'``: unsigned integer data types (i.e., + ``uint8``, ``uint16``, ``uint32``, ``uint64``). + - ``'integral'``: integer data types. Shorthand for ``('signed + integer', 'unsigned integer')``. + - ``'real floating'``: real-valued floating-point data types + (i.e., ``float32``, ``float64``). + - ``'complex floating'``: complex floating-point data types (i.e., + ``complex64``, ``complex128``). + - ``'numeric'``: numeric data types. Shorthand for ``('integral', + 'real floating', 'complex floating')``. + + Returns + ------- + dtypes : dict + A dictionary mapping the names of data types to the corresponding + PyTorch data types. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.devices + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.dtypes(kind='signed integer') + {'int8': numpy.int8, + 'int16': numpy.int16, + 'int32': numpy.int32, + 'int64': numpy.int64} + + """ + res = self._dtypes(kind) + for k, v in res.copy().items(): + try: + torch.empty((0,), dtype=v, device=device) + except: + del res[k] + return res + + @cache + def devices(self): + """ + The devices supported by PyTorch. + + Returns + ------- + devices : list of str + The devices supported by PyTorch. + + See Also + -------- + __array_namespace_info__.capabilities, + __array_namespace_info__.default_device, + __array_namespace_info__.default_dtypes, + __array_namespace_info__.dtypes + + Examples + -------- + >>> info = np.__array_namespace_info__() + >>> info.devices() + [device(type='cpu'), device(type='mps', index=0), device(type='meta')] + + """ + # Torch doesn't have a straightforward way to get the list of all + # currently supported devices. To do this, we first parse the error + # message of torch.device to get the list of all possible types of + # device: + try: + torch.device('notadevice') + except RuntimeError as e: + # The error message is something like: + # "Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: notadevice" + devices_names = e.args[0].split('Expected one of ')[1].split(' device type')[0].split(', ') + + # Next we need to check for different indices for different devices. + # device(device_name, index=index) doesn't actually check if the + # device name or index is valid. We have to try to create a tensor + # with it (which is why this function is cached). + devices = [] + for device_name in devices_names: + i = 0 + while True: + try: + a = torch.empty((0,), device=torch.device(device_name, index=i)) + if a.device in devices: + break + devices.append(a.device) + except: + break + i += 1 + + return devices diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/fft.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/fft.py new file mode 100644 index 0000000000000000000000000000000000000000..465f1b677a39b64d0408b962175327a58e0c7692 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/fft.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import torch + array = torch.Tensor + from typing import Union, Sequence, Literal + +from torch.fft import * # noqa: F403 +import torch.fft + +# Several torch fft functions do not map axes to dim + +def fftn( + x: array, + /, + *, + s: Sequence[int] = None, + axes: Sequence[int] = None, + norm: Literal["backward", "ortho", "forward"] = "backward", + **kwargs, +) -> array: + return torch.fft.fftn(x, s=s, dim=axes, norm=norm, **kwargs) + +def ifftn( + x: array, + /, + *, + s: Sequence[int] = None, + axes: Sequence[int] = None, + norm: Literal["backward", "ortho", "forward"] = "backward", + **kwargs, +) -> array: + return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, **kwargs) + +def rfftn( + x: array, + /, + *, + s: Sequence[int] = None, + axes: Sequence[int] = None, + norm: Literal["backward", "ortho", "forward"] = "backward", + **kwargs, +) -> array: + return torch.fft.rfftn(x, s=s, dim=axes, norm=norm, **kwargs) + +def irfftn( + x: array, + /, + *, + s: Sequence[int] = None, + axes: Sequence[int] = None, + norm: Literal["backward", "ortho", "forward"] = "backward", + **kwargs, +) -> array: + return torch.fft.irfftn(x, s=s, dim=axes, norm=norm, **kwargs) + +def fftshift( + x: array, + /, + *, + axes: Union[int, Sequence[int]] = None, + **kwargs, +) -> array: + return torch.fft.fftshift(x, dim=axes, **kwargs) + +def ifftshift( + x: array, + /, + *, + axes: Union[int, Sequence[int]] = None, + **kwargs, +) -> array: + return torch.fft.ifftshift(x, dim=axes, **kwargs) + + +__all__ = torch.fft.__all__ + [ + "fftn", + "ifftn", + "rfftn", + "irfftn", + "fftshift", + "ifftshift", +] + +_all_ignore = ['torch'] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/linalg.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/linalg.py new file mode 100644 index 0000000000000000000000000000000000000000..caacbcaba0bbd371ae17a3324fd5be7a673b0ec6 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_compat/torch/linalg.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import torch + array = torch.Tensor + from torch import dtype as Dtype + from typing import Optional, Union, Tuple, Literal + inf = float('inf') + +from ._aliases import _fix_promotion, sum + +from torch.linalg import * # noqa: F403 + +# torch.linalg doesn't define __all__ +# from torch.linalg import __all__ as linalg_all +from torch import linalg as torch_linalg +linalg_all = [i for i in dir(torch_linalg) if not i.startswith('_')] + +# outer is implemented in torch but aren't in the linalg namespace +from torch import outer +# These functions are in both the main and linalg namespaces +from ._aliases import matmul, matrix_transpose, tensordot + +# Note: torch.linalg.cross does not default to axis=-1 (it defaults to the +# first axis with size 3), see https://github.com/pytorch/pytorch/issues/58743 + +# torch.cross also does not support broadcasting when it would add new +# dimensions https://github.com/pytorch/pytorch/issues/39656 +def cross(x1: array, x2: array, /, *, axis: int = -1) -> array: + x1, x2 = _fix_promotion(x1, x2, only_scalar=False) + if not (-min(x1.ndim, x2.ndim) <= axis < max(x1.ndim, x2.ndim)): + raise ValueError(f"axis {axis} out of bounds for cross product of arrays with shapes {x1.shape} and {x2.shape}") + if not (x1.shape[axis] == x2.shape[axis] == 3): + raise ValueError(f"cross product axis must have size 3, got {x1.shape[axis]} and {x2.shape[axis]}") + x1, x2 = torch.broadcast_tensors(x1, x2) + return torch_linalg.cross(x1, x2, dim=axis) + +def vecdot(x1: array, x2: array, /, *, axis: int = -1, **kwargs) -> array: + from ._aliases import isdtype + + x1, x2 = _fix_promotion(x1, x2, only_scalar=False) + + # torch.linalg.vecdot incorrectly allows broadcasting along the contracted dimension + if x1.shape[axis] != x2.shape[axis]: + raise ValueError("x1 and x2 must have the same size along the given axis") + + # torch.linalg.vecdot doesn't support integer dtypes + if isdtype(x1.dtype, 'integral') or isdtype(x2.dtype, 'integral'): + if kwargs: + raise RuntimeError("vecdot kwargs not supported for integral dtypes") + + x1_ = torch.moveaxis(x1, axis, -1) + x2_ = torch.moveaxis(x2, axis, -1) + x1_, x2_ = torch.broadcast_tensors(x1_, x2_) + + res = x1_[..., None, :] @ x2_[..., None] + return res[..., 0, 0] + return torch.linalg.vecdot(x1, x2, dim=axis, **kwargs) + +def solve(x1: array, x2: array, /, **kwargs) -> array: + x1, x2 = _fix_promotion(x1, x2, only_scalar=False) + # Torch tries to emulate NumPy 1 solve behavior by using batched 1-D solve + # whenever + # 1. x1.ndim - 1 == x2.ndim + # 2. x1.shape[:-1] == x2.shape + # + # See linalg_solve_is_vector_rhs in + # aten/src/ATen/native/LinearAlgebraUtils.h and + # TORCH_META_FUNC(_linalg_solve_ex) in + # aten/src/ATen/native/BatchLinearAlgebra.cpp in the PyTorch source code. + # + # The easiest way to work around this is to prepend a size 1 dimension to + # x2, since x2 is already one dimension less than x1. + # + # See https://github.com/pytorch/pytorch/issues/52915 + if x2.ndim != 1 and x1.ndim - 1 == x2.ndim and x1.shape[:-1] == x2.shape: + x2 = x2[None] + return torch.linalg.solve(x1, x2, **kwargs) + +# torch.trace doesn't support the offset argument and doesn't support stacking +def trace(x: array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> array: + # Use our wrapped sum to make sure it does upcasting correctly + return sum(torch.diagonal(x, offset=offset, dim1=-2, dim2=-1), axis=-1, dtype=dtype) + +def vector_norm( + x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, + ord: Union[int, float, Literal[inf, -inf]] = 2, + **kwargs, +) -> array: + # torch.vector_norm incorrectly treats axis=() the same as axis=None + if axis == (): + out = kwargs.get('out') + if out is None: + dtype = None + if x.dtype == torch.complex64: + dtype = torch.float32 + elif x.dtype == torch.complex128: + dtype = torch.float64 + + out = torch.zeros_like(x, dtype=dtype) + + # The norm of a single scalar works out to abs(x) in every case except + # for ord=0, which is x != 0. + if ord == 0: + out[:] = (x != 0) + else: + out[:] = torch.abs(x) + return out + return torch.linalg.vector_norm(x, ord=ord, axis=axis, keepdim=keepdims, **kwargs) + +__all__ = linalg_all + ['outer', 'matmul', 'matrix_transpose', 'tensordot', + 'cross', 'vecdot', 'solve', 'trace', 'vector_norm'] + +_all_ignore = ['torch_linalg', 'sum'] + +del linalg_all diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..845a081fa8abc8138c57935a1e77cc1cbce46956 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/__init__.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from ._funcs import atleast_nd, cov, create_diagonal, expand_dims, kron, sinc + +__version__ = "0.2.0" + +__all__ = [ + "__version__", + "atleast_nd", + "cov", + "create_diagonal", + "expand_dims", + "kron", + "sinc", +] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/_funcs.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/_funcs.py new file mode 100644 index 0000000000000000000000000000000000000000..172db59bb5518c201f04b2eaa439457ef134f0ce --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/_funcs.py @@ -0,0 +1,484 @@ +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ._typing import Array, ModuleType + +__all__ = ["atleast_nd", "cov", "create_diagonal", "expand_dims", "kron", "sinc"] + + +def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType) -> Array: + """ + Recursively expand the dimension of an array to at least `ndim`. + + Parameters + ---------- + x : array + ndim : int + The minimum number of dimensions for the result. + xp : array_namespace + The standard-compatible namespace for `x`. + + Returns + ------- + res : array + An array with ``res.ndim`` >= `ndim`. + If ``x.ndim`` >= `ndim`, `x` is returned. + If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes + until ``res.ndim`` equals `ndim`. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([1]) + >>> xpx.atleast_nd(x, ndim=3, xp=xp) + Array([[[1]]], dtype=array_api_strict.int64) + + >>> x = xp.asarray([[[1, 2], + ... [3, 4]]]) + >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x + True + + """ + if x.ndim < ndim: + x = xp.expand_dims(x, axis=0) + x = atleast_nd(x, ndim=ndim, xp=xp) + return x + + +def cov(m: Array, /, *, xp: ModuleType) -> Array: + """ + Estimate a covariance matrix. + + Covariance indicates the level to which two variables vary together. + If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`, + then the covariance matrix element :math:`C_{ij}` is the covariance of + :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance + of :math:`x_i`. + + This provides a subset of the functionality of ``numpy.cov``. + + Parameters + ---------- + m : array + A 1-D or 2-D array containing multiple variables and observations. + Each row of `m` represents a variable, and each column a single + observation of all those variables. + xp : array_namespace + The standard-compatible namespace for `m`. + + Returns + ------- + res : array + The covariance matrix of the variables. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + + Consider two variables, :math:`x_0` and :math:`x_1`, which + correlate perfectly, but in opposite directions: + + >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T + >>> x + Array([[0, 1, 2], + [2, 1, 0]], dtype=array_api_strict.int64) + + Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance + matrix shows this clearly: + + >>> xpx.cov(x, xp=xp) + Array([[ 1., -1.], + [-1., 1.]], dtype=array_api_strict.float64) + + + Note that element :math:`C_{0,1}`, which shows the correlation between + :math:`x_0` and :math:`x_1`, is negative. + + Further, note how `x` and `y` are combined: + + >>> x = xp.asarray([-2.1, -1, 4.3]) + >>> y = xp.asarray([3, 1.1, 0.12]) + >>> X = xp.stack((x, y), axis=0) + >>> xpx.cov(X, xp=xp) + Array([[11.71 , -4.286 ], + [-4.286 , 2.14413333]], dtype=array_api_strict.float64) + + >>> xpx.cov(x, xp=xp) + Array(11.71, dtype=array_api_strict.float64) + + >>> xpx.cov(y, xp=xp) + Array(2.14413333, dtype=array_api_strict.float64) + + """ + m = xp.asarray(m, copy=True) + dtype = ( + xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64) + ) + + m = atleast_nd(m, ndim=2, xp=xp) + m = xp.astype(m, dtype) + + avg = _mean(m, axis=1, xp=xp) + fact = m.shape[1] - 1 + + if fact <= 0: + warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2) + fact = 0.0 + + m -= avg[:, None] + m_transpose = m.T + if xp.isdtype(m_transpose.dtype, "complex floating"): + m_transpose = xp.conj(m_transpose) + c = m @ m_transpose + c /= fact + axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1) + return xp.squeeze(c, axis=axes) + + +def create_diagonal(x: Array, /, *, offset: int = 0, xp: ModuleType) -> Array: + """ + Construct a diagonal array. + + Parameters + ---------- + x : array + A 1-D array + offset : int, optional + Offset from the leading diagonal (default is ``0``). + Use positive ints for diagonals above the leading diagonal, + and negative ints for diagonals below the leading diagonal. + xp : array_namespace + The standard-compatible namespace for `x`. + + Returns + ------- + res : array + A 2-D array with `x` on the diagonal (offset by `offset`). + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([2, 4, 8]) + + >>> xpx.create_diagonal(x, xp=xp) + Array([[2, 0, 0], + [0, 4, 0], + [0, 0, 8]], dtype=array_api_strict.int64) + + >>> xpx.create_diagonal(x, offset=-2, xp=xp) + Array([[0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [2, 0, 0, 0, 0], + [0, 4, 0, 0, 0], + [0, 0, 8, 0, 0]], dtype=array_api_strict.int64) + + """ + if x.ndim != 1: + err_msg = "`x` must be 1-dimensional." + raise ValueError(err_msg) + n = x.shape[0] + abs(offset) + diag = xp.zeros(n**2, dtype=x.dtype) + i = offset if offset >= 0 else abs(offset) * n + diag[i : min(n * (n - offset), diag.shape[0]) : n + 1] = x + return xp.reshape(diag, (n, n)) + + +def _mean( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, + xp: ModuleType, +) -> Array: + """ + Complex mean, https://github.com/data-apis/array-api/issues/846. + """ + if xp.isdtype(x.dtype, "complex floating"): + x_real = xp.real(x) + x_imag = xp.imag(x) + mean_real = xp.mean(x_real, axis=axis, keepdims=keepdims) + mean_imag = xp.mean(x_imag, axis=axis, keepdims=keepdims) + return mean_real + (mean_imag * xp.asarray(1j)) + return xp.mean(x, axis=axis, keepdims=keepdims) + + +def expand_dims( + a: Array, /, *, axis: int | tuple[int, ...] = (0,), xp: ModuleType +) -> Array: + """ + Expand the shape of an array. + + Insert (a) new axis/axes that will appear at the position(s) specified by + `axis` in the expanded array shape. + + This is ``xp.expand_dims`` for `axis` an int *or a tuple of ints*. + Roughly equivalent to ``numpy.expand_dims`` for NumPy arrays. + + Parameters + ---------- + a : array + axis : int or tuple of ints, optional + Position(s) in the expanded axes where the new axis (or axes) is/are placed. + If multiple positions are provided, they should be unique (note that a position + given by a positive index could also be referred to by a negative index - + that will also result in an error). + Default: ``(0,)``. + xp : array_namespace + The standard-compatible namespace for `a`. + + Returns + ------- + res : array + `a` with an expanded shape. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([1, 2]) + >>> x.shape + (2,) + + The following is equivalent to ``x[xp.newaxis, :]`` or ``x[xp.newaxis]``: + + >>> y = xpx.expand_dims(x, axis=0, xp=xp) + >>> y + Array([[1, 2]], dtype=array_api_strict.int64) + >>> y.shape + (1, 2) + + The following is equivalent to ``x[:, xp.newaxis]``: + + >>> y = xpx.expand_dims(x, axis=1, xp=xp) + >>> y + Array([[1], + [2]], dtype=array_api_strict.int64) + >>> y.shape + (2, 1) + + ``axis`` may also be a tuple: + + >>> y = xpx.expand_dims(x, axis=(0, 1), xp=xp) + >>> y + Array([[[1, 2]]], dtype=array_api_strict.int64) + + >>> y = xpx.expand_dims(x, axis=(2, 0), xp=xp) + >>> y + Array([[[1], + [2]]], dtype=array_api_strict.int64) + + """ + if not isinstance(axis, tuple): + axis = (axis,) + ndim = a.ndim + len(axis) + if axis != () and (min(axis) < -ndim or max(axis) >= ndim): + err_msg = ( + f"a provided axis position is out of bounds for array of dimension {a.ndim}" + ) + raise IndexError(err_msg) + axis = tuple(dim % ndim for dim in axis) + if len(set(axis)) != len(axis): + err_msg = "Duplicate dimensions specified in `axis`." + raise ValueError(err_msg) + for i in sorted(axis): + a = xp.expand_dims(a, axis=i) + return a + + +def kron(a: Array, b: Array, /, *, xp: ModuleType) -> Array: + """ + Kronecker product of two arrays. + + Computes the Kronecker product, a composite array made of blocks of the + second array scaled by the first. + + Equivalent to ``numpy.kron`` for NumPy arrays. + + Parameters + ---------- + a, b : array + xp : array_namespace + The standard-compatible namespace for `a` and `b`. + + Returns + ------- + res : array + The Kronecker product of `a` and `b`. + + Notes + ----- + The function assumes that the number of dimensions of `a` and `b` + are the same, if necessary prepending the smallest with ones. + If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``, + the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``. + The elements are products of elements from `a` and `b`, organized + explicitly by:: + + kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN] + + where:: + + kt = it * st + jt, t = 0,...,N + + In the common 2-D case (N=1), the block structure can be visualized:: + + [[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ], + [ ... ... ], + [ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]] + + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> xpx.kron(xp.asarray([1, 10, 100]), xp.asarray([5, 6, 7]), xp=xp) + Array([ 5, 6, 7, 50, 60, 70, 500, + 600, 700], dtype=array_api_strict.int64) + + >>> xpx.kron(xp.asarray([5, 6, 7]), xp.asarray([1, 10, 100]), xp=xp) + Array([ 5, 50, 500, 6, 60, 600, 7, + 70, 700], dtype=array_api_strict.int64) + + >>> xpx.kron(xp.eye(2), xp.ones((2, 2)), xp=xp) + Array([[1., 1., 0., 0.], + [1., 1., 0., 0.], + [0., 0., 1., 1.], + [0., 0., 1., 1.]], dtype=array_api_strict.float64) + + + >>> a = xp.reshape(xp.arange(100), (2, 5, 2, 5)) + >>> b = xp.reshape(xp.arange(24), (2, 3, 4)) + >>> c = xpx.kron(a, b, xp=xp) + >>> c.shape + (2, 10, 6, 20) + >>> I = (1, 3, 0, 2) + >>> J = (0, 2, 1) + >>> J1 = (0,) + J # extend to ndim=4 + >>> S1 = (1,) + b.shape + >>> K = tuple(xp.asarray(I) * xp.asarray(S1) + xp.asarray(J1)) + >>> c[K] == a[I]*b[J] + Array(True, dtype=array_api_strict.bool) + + """ + + b = xp.asarray(b) + singletons = (1,) * (b.ndim - a.ndim) + a = xp.broadcast_to(xp.asarray(a), singletons + a.shape) + + nd_b, nd_a = b.ndim, a.ndim + nd_max = max(nd_b, nd_a) + if nd_a == 0 or nd_b == 0: + return xp.multiply(a, b) + + a_shape = a.shape + b_shape = b.shape + + # Equalise the shapes by prepending smaller one with 1s + a_shape = (1,) * max(0, nd_b - nd_a) + a_shape + b_shape = (1,) * max(0, nd_a - nd_b) + b_shape + + # Insert empty dimensions + a_arr = expand_dims(a, axis=tuple(range(nd_b - nd_a)), xp=xp) + b_arr = expand_dims(b, axis=tuple(range(nd_a - nd_b)), xp=xp) + + # Compute the product + a_arr = expand_dims(a_arr, axis=tuple(range(1, nd_max * 2, 2)), xp=xp) + b_arr = expand_dims(b_arr, axis=tuple(range(0, nd_max * 2, 2)), xp=xp) + result = xp.multiply(a_arr, b_arr) + + # Reshape back and return + a_shape = xp.asarray(a_shape) + b_shape = xp.asarray(b_shape) + return xp.reshape(result, tuple(xp.multiply(a_shape, b_shape))) + + +def sinc(x: Array, /, *, xp: ModuleType) -> Array: + r""" + Return the normalized sinc function. + + The sinc function is equal to :math:`\sin(\pi x)/(\pi x)` for any argument + :math:`x\ne 0`. ``sinc(0)`` takes the limit value 1, making ``sinc`` not + only everywhere continuous but also infinitely differentiable. + + .. note:: + + Note the normalization factor of ``pi`` used in the definition. + This is the most commonly used definition in signal processing. + Use ``sinc(x / xp.pi)`` to obtain the unnormalized sinc function + :math:`\sin(x)/x` that is more common in mathematics. + + Parameters + ---------- + x : array + Array (possibly multi-dimensional) of values for which to calculate + ``sinc(x)``. Must have a real floating point dtype. + xp : array_namespace + The standard-compatible namespace for `x`. + + Returns + ------- + res : array + ``sinc(x)`` calculated elementwise, which has the same shape as the input. + + Notes + ----- + The name sinc is short for "sine cardinal" or "sinus cardinalis". + + The sinc function is used in various signal processing applications, + including in anti-aliasing, in the construction of a Lanczos resampling + filter, and in interpolation. + + For bandlimited interpolation of discrete-time signals, the ideal + interpolation kernel is proportional to the sinc function. + + References + ---------- + .. [1] Weisstein, Eric W. "Sinc Function." From MathWorld--A Wolfram Web + Resource. https://mathworld.wolfram.com/SincFunction.html + .. [2] Wikipedia, "Sinc function", + https://en.wikipedia.org/wiki/Sinc_function + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.linspace(-4, 4, 41) + >>> xpx.sinc(x, xp=xp) + Array([-3.89817183e-17, -4.92362781e-02, + -8.40918587e-02, -8.90384387e-02, + -5.84680802e-02, 3.89817183e-17, + 6.68206631e-02, 1.16434881e-01, + 1.26137788e-01, 8.50444803e-02, + -3.89817183e-17, -1.03943254e-01, + -1.89206682e-01, -2.16236208e-01, + -1.55914881e-01, 3.89817183e-17, + 2.33872321e-01, 5.04551152e-01, + 7.56826729e-01, 9.35489284e-01, + 1.00000000e+00, 9.35489284e-01, + 7.56826729e-01, 5.04551152e-01, + 2.33872321e-01, 3.89817183e-17, + -1.55914881e-01, -2.16236208e-01, + -1.89206682e-01, -1.03943254e-01, + -3.89817183e-17, 8.50444803e-02, + 1.26137788e-01, 1.16434881e-01, + 6.68206631e-02, 3.89817183e-17, + -5.84680802e-02, -8.90384387e-02, + -8.40918587e-02, -4.92362781e-02, + -3.89817183e-17], dtype=array_api_strict.float64) + + """ + if not xp.isdtype(x.dtype, "real floating"): + err_msg = "`x` must have a real floating data type." + raise ValueError(err_msg) + # no scalars in `where` - array-api#807 + y = xp.pi * xp.where( + x, x, xp.asarray(xp.finfo(x.dtype).smallest_normal, dtype=x.dtype) + ) + return xp.sin(y) / y diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/_typing.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..e681154ba8beadf8ed1f839c38fd1bb3dfff08e8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/array_api_extra/_typing.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from types import ModuleType +from typing import Any + +Array = Any # To be changed to a Protocol later (see array-api#589) + +__all__ = ["Array", "ModuleType"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8fce9593e615c820101980266822a7e0ab5da7b7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/__init__.py @@ -0,0 +1,20 @@ +from .main import minimize +from .utils import show_versions + +# PEP0440 compatible formatted version, see: +# https://www.python.org/dev/peps/pep-0440/ +# +# Final release markers: +# X.Y.0 # For first release after an increment in Y +# X.Y.Z # For bugfix releases +# +# Admissible pre-release markers: +# X.YaN # Alpha release +# X.YbN # Beta release +# X.YrcN # Release Candidate +# +# Dev branch marker is: 'X.Y.dev' or 'X.Y.devN' where N is an integer. +# 'X.Y.dev0' is the canonical version of 'X.Y.dev'. +__version__ = "1.1.2" + +__all__ = ["minimize", "show_versions"] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/framework.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/framework.py new file mode 100644 index 0000000000000000000000000000000000000000..59b8ea0442c81af1f39ff12381bc993e83dbb0f5 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/framework.py @@ -0,0 +1,1240 @@ +import warnings + +import numpy as np +from scipy.optimize import lsq_linear + +from .models import Models, Quadratic +from .settings import Options, Constants +from .subsolvers import ( + cauchy_geometry, + spider_geometry, + normal_byrd_omojokun, + tangential_byrd_omojokun, + constrained_tangential_byrd_omojokun, +) +from .subsolvers.optim import qr_tangential_byrd_omojokun +from .utils import get_arrays_tol + + +TINY = np.finfo(float).tiny +EPS = np.finfo(float).eps + + +class TrustRegion: + """ + Trust-region framework. + """ + + def __init__(self, pb, options, constants): + """ + Initialize the trust-region framework. + + Parameters + ---------- + pb : `cobyqa.problem.Problem` + Problem to solve. + options : dict + Options of the solver. + constants : dict + Constants of the solver. + + Raises + ------ + `cobyqa.utils.MaxEvalError` + If the maximum number of evaluations is reached. + `cobyqa.utils.TargetSuccess` + If a nearly feasible point has been found with an objective + function value below the target. + `cobyqa.utils.FeasibleSuccess` + If a feasible point has been found for a feasibility problem. + `numpy.linalg.LinAlgError` + If the initial interpolation system is ill-defined. + """ + # Set the initial penalty parameter. + self._penalty = 0.0 + + # Initialize the models. + self._pb = pb + self._models = Models(self._pb, options, self.penalty) + self._constants = constants + + # Set the index of the best interpolation point. + self._best_index = 0 + self.set_best_index() + + # Set the initial Lagrange multipliers. + self._lm_linear_ub = np.zeros(self.m_linear_ub) + self._lm_linear_eq = np.zeros(self.m_linear_eq) + self._lm_nonlinear_ub = np.zeros(self.m_nonlinear_ub) + self._lm_nonlinear_eq = np.zeros(self.m_nonlinear_eq) + self.set_multipliers(self.x_best) + + # Set the initial trust-region radius and the resolution. + self._resolution = options[Options.RHOBEG] + self._radius = self.resolution + + @property + def n(self): + """ + Number of variables. + + Returns + ------- + int + Number of variables. + """ + return self._pb.n + + @property + def m_linear_ub(self): + """ + Number of linear inequality constraints. + + Returns + ------- + int + Number of linear inequality constraints. + """ + return self._pb.m_linear_ub + + @property + def m_linear_eq(self): + """ + Number of linear equality constraints. + + Returns + ------- + int + Number of linear equality constraints. + """ + return self._pb.m_linear_eq + + @property + def m_nonlinear_ub(self): + """ + Number of nonlinear inequality constraints. + + Returns + ------- + int + Number of nonlinear inequality constraints. + """ + return self._pb.m_nonlinear_ub + + @property + def m_nonlinear_eq(self): + """ + Number of nonlinear equality constraints. + + Returns + ------- + int + Number of nonlinear equality constraints. + """ + return self._pb.m_nonlinear_eq + + @property + def radius(self): + """ + Trust-region radius. + + Returns + ------- + float + Trust-region radius. + """ + return self._radius + + @radius.setter + def radius(self, radius): + """ + Set the trust-region radius. + + Parameters + ---------- + radius : float + New trust-region radius. + """ + self._radius = radius + if ( + self.radius + <= self._constants[Constants.DECREASE_RADIUS_THRESHOLD] + * self.resolution + ): + self._radius = self.resolution + + @property + def resolution(self): + """ + Resolution of the trust-region framework. + + The resolution is a lower bound on the trust-region radius. + + Returns + ------- + float + Resolution of the trust-region framework. + """ + return self._resolution + + @resolution.setter + def resolution(self, resolution): + """ + Set the resolution of the trust-region framework. + + Parameters + ---------- + resolution : float + New resolution of the trust-region framework. + """ + self._resolution = resolution + + @property + def penalty(self): + """ + Penalty parameter. + + Returns + ------- + float + Penalty parameter. + """ + return self._penalty + + @property + def models(self): + """ + Models of the objective function and constraints. + + Returns + ------- + `cobyqa.models.Models` + Models of the objective function and constraints. + """ + return self._models + + @property + def best_index(self): + """ + Index of the best interpolation point. + + Returns + ------- + int + Index of the best interpolation point. + """ + return self._best_index + + @property + def x_best(self): + """ + Best interpolation point. + + Its value is interpreted as relative to the origin, not the base point. + + Returns + ------- + `numpy.ndarray` + Best interpolation point. + """ + return self.models.interpolation.point(self.best_index) + + @property + def fun_best(self): + """ + Value of the objective function at `x_best`. + + Returns + ------- + float + Value of the objective function at `x_best`. + """ + return self.models.fun_val[self.best_index] + + @property + def cub_best(self): + """ + Values of the nonlinear inequality constraints at `x_best`. + + Returns + ------- + `numpy.ndarray`, shape (m_nonlinear_ub,) + Values of the nonlinear inequality constraints at `x_best`. + """ + return self.models.cub_val[self.best_index, :] + + @property + def ceq_best(self): + """ + Values of the nonlinear equality constraints at `x_best`. + + Returns + ------- + `numpy.ndarray`, shape (m_nonlinear_eq,) + Values of the nonlinear equality constraints at `x_best`. + """ + return self.models.ceq_val[self.best_index, :] + + def lag_model(self, x): + """ + Evaluate the Lagrangian model at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the Lagrangian model is evaluated. + + Returns + ------- + float + Value of the Lagrangian model at `x`. + """ + return ( + self.models.fun(x) + + self._lm_linear_ub + @ (self._pb.linear.a_ub @ x - self._pb.linear.b_ub) + + self._lm_linear_eq + @ (self._pb.linear.a_eq @ x - self._pb.linear.b_eq) + + self._lm_nonlinear_ub @ self.models.cub(x) + + self._lm_nonlinear_eq @ self.models.ceq(x) + ) + + def lag_model_grad(self, x): + """ + Evaluate the gradient of the Lagrangian model at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the gradient of the Lagrangian model is evaluated. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Gradient of the Lagrangian model at `x`. + """ + return ( + self.models.fun_grad(x) + + self._lm_linear_ub @ self._pb.linear.a_ub + + self._lm_linear_eq @ self._pb.linear.a_eq + + self._lm_nonlinear_ub @ self.models.cub_grad(x) + + self._lm_nonlinear_eq @ self.models.ceq_grad(x) + ) + + def lag_model_hess(self): + """ + Evaluate the Hessian matrix of the Lagrangian model at a given point. + + Returns + ------- + `numpy.ndarray`, shape (n, n) + Hessian matrix of the Lagrangian model at `x`. + """ + hess = self.models.fun_hess() + if self.m_nonlinear_ub > 0: + hess += self._lm_nonlinear_ub @ self.models.cub_hess() + if self.m_nonlinear_eq > 0: + hess += self._lm_nonlinear_eq @ self.models.ceq_hess() + return hess + + def lag_model_hess_prod(self, v): + """ + Evaluate the right product of the Hessian matrix of the Lagrangian + model with a given vector. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Vector with which the Hessian matrix of the Lagrangian model is + multiplied from the right. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Right product of the Hessian matrix of the Lagrangian model with + `v`. + """ + return ( + self.models.fun_hess_prod(v) + + self._lm_nonlinear_ub @ self.models.cub_hess_prod(v) + + self._lm_nonlinear_eq @ self.models.ceq_hess_prod(v) + ) + + def lag_model_curv(self, v): + """ + Evaluate the curvature of the Lagrangian model along a given direction. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Direction along which the curvature of the Lagrangian model is + evaluated. + + Returns + ------- + float + Curvature of the Lagrangian model along `v`. + """ + return ( + self.models.fun_curv(v) + + self._lm_nonlinear_ub @ self.models.cub_curv(v) + + self._lm_nonlinear_eq @ self.models.ceq_curv(v) + ) + + def sqp_fun(self, step): + """ + Evaluate the objective function of the SQP subproblem. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Step along which the objective function of the SQP subproblem is + evaluated. + + Returns + ------- + float + Value of the objective function of the SQP subproblem along `step`. + """ + return step @ ( + self.models.fun_grad(self.x_best) + + 0.5 * self.lag_model_hess_prod(step) + ) + + def sqp_cub(self, step): + """ + Evaluate the linearization of the nonlinear inequality constraints. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Step along which the linearization of the nonlinear inequality + constraints is evaluated. + + Returns + ------- + `numpy.ndarray`, shape (m_nonlinear_ub,) + Value of the linearization of the nonlinear inequality constraints + along `step`. + """ + return ( + self.models.cub(self.x_best) + + self.models.cub_grad(self.x_best) @ step + ) + + def sqp_ceq(self, step): + """ + Evaluate the linearization of the nonlinear equality constraints. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Step along which the linearization of the nonlinear equality + constraints is evaluated. + + Returns + ------- + `numpy.ndarray`, shape (m_nonlinear_ub,) + Value of the linearization of the nonlinear equality constraints + along `step`. + """ + return ( + self.models.ceq(self.x_best) + + self.models.ceq_grad(self.x_best) @ step + ) + + def merit(self, x, fun_val=None, cub_val=None, ceq_val=None): + """ + Evaluate the merit function at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the merit function is evaluated. + fun_val : float, optional + Value of the objective function at `x`. If not provided, the + objective function is evaluated at `x`. + cub_val : `numpy.ndarray`, shape (m_nonlinear_ub,), optional + Values of the nonlinear inequality constraints. If not provided, + the nonlinear inequality constraints are evaluated at `x`. + ceq_val : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Values of the nonlinear equality constraints. If not provided, + the nonlinear equality constraints are evaluated at `x`. + + Returns + ------- + float + Value of the merit function at `x`. + """ + if fun_val is None or cub_val is None or ceq_val is None: + fun_val, cub_val, ceq_val = self._pb(x, self.penalty) + m_val = fun_val + if self._penalty > 0.0: + c_val = self._pb.violation(x, cub_val=cub_val, ceq_val=ceq_val) + if np.count_nonzero(c_val): + m_val += self._penalty * np.linalg.norm(c_val) + return m_val + + def get_constraint_linearizations(self, x): + """ + Get the linearizations of the constraints at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the linearizations of the constraints are evaluated. + + Returns + ------- + `numpy.ndarray`, shape (m_linear_ub + m_nonlinear_ub, n) + Left-hand side matrix of the linearized inequality constraints. + `numpy.ndarray`, shape (m_linear_ub + m_nonlinear_ub,) + Right-hand side vector of the linearized inequality constraints. + `numpy.ndarray`, shape (m_linear_eq + m_nonlinear_eq, n) + Left-hand side matrix of the linearized equality constraints. + `numpy.ndarray`, shape (m_linear_eq + m_nonlinear_eq,) + Right-hand side vector of the linearized equality constraints. + """ + aub = np.block( + [ + [self._pb.linear.a_ub], + [self.models.cub_grad(x)], + ] + ) + bub = np.block( + [ + self._pb.linear.b_ub - self._pb.linear.a_ub @ x, + -self.models.cub(x), + ] + ) + aeq = np.block( + [ + [self._pb.linear.a_eq], + [self.models.ceq_grad(x)], + ] + ) + beq = np.block( + [ + self._pb.linear.b_eq - self._pb.linear.a_eq @ x, + -self.models.ceq(x), + ] + ) + return aub, bub, aeq, beq + + def get_trust_region_step(self, options): + """ + Get the trust-region step. + + The trust-region step is computed by solving the derivative-free + trust-region SQP subproblem using a Byrd-Omojokun composite-step + approach. For more details, see Section 5.2.3 of [1]_. + + Parameters + ---------- + options : dict + Options of the solver. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Normal step. + `numpy.ndarray`, shape (n,) + Tangential step. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization + Methods and Software*. PhD thesis, Department of Applied + Mathematics, The Hong Kong Polytechnic University, Hong Kong, China, + 2022. URL: https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + # Evaluate the linearizations of the constraints. + aub, bub, aeq, beq = self.get_constraint_linearizations(self.x_best) + xl = self._pb.bounds.xl - self.x_best + xu = self._pb.bounds.xu - self.x_best + + # Evaluate the normal step. + radius = self._constants[Constants.BYRD_OMOJOKUN_FACTOR] * self.radius + normal_step = normal_byrd_omojokun( + aub, + bub, + aeq, + beq, + xl, + xu, + radius, + options[Options.DEBUG], + **self._constants, + ) + if options[Options.DEBUG]: + tol = get_arrays_tol(xl, xu) + if (np.any(normal_step + tol < xl) + or np.any(xu < normal_step - tol)): + warnings.warn( + "the normal step does not respect the bound constraint.", + RuntimeWarning, + 2, + ) + if np.linalg.norm(normal_step) > 1.1 * radius: + warnings.warn( + "the normal step does not respect the trust-region " + "constraint.", + RuntimeWarning, + 2, + ) + + # Evaluate the tangential step. + radius = np.sqrt(self.radius**2.0 - normal_step @ normal_step) + xl -= normal_step + xu -= normal_step + bub = np.maximum(bub - aub @ normal_step, 0.0) + g_best = self.models.fun_grad(self.x_best) + self.lag_model_hess_prod( + normal_step + ) + if self._pb.type in ["unconstrained", "bound-constrained"]: + tangential_step = tangential_byrd_omojokun( + g_best, + self.lag_model_hess_prod, + xl, + xu, + radius, + options[Options.DEBUG], + **self._constants, + ) + else: + tangential_step = constrained_tangential_byrd_omojokun( + g_best, + self.lag_model_hess_prod, + xl, + xu, + aub, + bub, + aeq, + radius, + options["debug"], + **self._constants, + ) + if options[Options.DEBUG]: + tol = get_arrays_tol(xl, xu) + if np.any(tangential_step + tol < xl) or np.any( + xu < tangential_step - tol + ): + warnings.warn( + "The tangential step does not respect the bound " + "constraints.", + RuntimeWarning, + 2, + ) + if ( + np.linalg.norm(normal_step + tangential_step) + > 1.1 * np.sqrt(2.0) * self.radius + ): + warnings.warn( + "The trial step does not respect the trust-region " + "constraint.", + RuntimeWarning, + 2, + ) + return normal_step, tangential_step + + def get_geometry_step(self, k_new, options): + """ + Get the geometry-improving step. + + Three different geometry-improving steps are computed and the best one + is returned. For more details, see Section 5.2.7 of [1]_. + + Parameters + ---------- + k_new : int + Index of the interpolation point to be modified. + options : dict + Options of the solver. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Geometry-improving step. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the computation of a determinant fails. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization + Methods and Software*. PhD thesis, Department of Applied + Mathematics, The Hong Kong Polytechnic University, Hong Kong, China, + 2022. URL: https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + if options[Options.DEBUG]: + assert ( + k_new != self.best_index + ), "The index `k_new` must be different from the best index." + + # Build the k_new-th Lagrange polynomial. + coord_vec = np.squeeze(np.eye(1, self.models.npt, k_new)) + lag = Quadratic( + self.models.interpolation, + coord_vec, + options[Options.DEBUG], + ) + g_lag = lag.grad(self.x_best, self.models.interpolation) + + # Compute a simple constrained Cauchy step. + xl = self._pb.bounds.xl - self.x_best + xu = self._pb.bounds.xu - self.x_best + step = cauchy_geometry( + 0.0, + g_lag, + lambda v: lag.curv(v, self.models.interpolation), + xl, + xu, + self.radius, + options[Options.DEBUG], + ) + sigma = self.models.determinants(self.x_best + step, k_new) + + # Compute the solution on the straight lines joining the interpolation + # points to the k-th one, and choose it if it provides a larger value + # of the determinant of the interpolation system in absolute value. + xpt = ( + self.models.interpolation.xpt + - self.models.interpolation.xpt[:, self.best_index, np.newaxis] + ) + xpt[:, [0, self.best_index]] = xpt[:, [self.best_index, 0]] + step_alt = spider_geometry( + 0.0, + g_lag, + lambda v: lag.curv(v, self.models.interpolation), + xpt[:, 1:], + xl, + xu, + self.radius, + options[Options.DEBUG], + ) + sigma_alt = self.models.determinants(self.x_best + step_alt, k_new) + if abs(sigma_alt) > abs(sigma): + step = step_alt + sigma = sigma_alt + + # Compute a Cauchy step on the tangent space of the active constraints. + if self._pb.type in [ + "linearly constrained", + "nonlinearly constrained", + ]: + aub, bub, aeq, beq = ( + self.get_constraint_linearizations(self.x_best)) + tol_bd = get_arrays_tol(xl, xu) + tol_ub = get_arrays_tol(bub) + free_xl = xl <= -tol_bd + free_xu = xu >= tol_bd + free_ub = bub >= tol_ub + + # Compute the Cauchy step. + n_act, q = qr_tangential_byrd_omojokun( + aub, + aeq, + free_xl, + free_xu, + free_ub, + ) + g_lag_proj = q[:, n_act:] @ (q[:, n_act:].T @ g_lag) + norm_g_lag_proj = np.linalg.norm(g_lag_proj) + if 0 < n_act < self._pb.n and norm_g_lag_proj > TINY * self.radius: + step_alt = (self.radius / norm_g_lag_proj) * g_lag_proj + if lag.curv(step_alt, self.models.interpolation) < 0.0: + step_alt = -step_alt + + # Evaluate the constraint violation at the Cauchy step. + cbd = np.block([xl - step_alt, step_alt - xu]) + cub = aub @ step_alt - bub + ceq = aeq @ step_alt - beq + maxcv_val = max( + np.max(array, initial=0.0) + for array in [cbd, cub, np.abs(ceq)] + ) + + # Accept the new step if it is nearly feasible and do not + # drastically worsen the determinant of the interpolation + # system in absolute value. + tol = np.max(np.abs(step_alt[~free_xl]), initial=0.0) + tol = np.max(np.abs(step_alt[~free_xu]), initial=tol) + tol = np.max(np.abs(aub[~free_ub, :] @ step_alt), initial=tol) + tol = min(10.0 * tol, 1e-2 * np.linalg.norm(step_alt)) + if maxcv_val <= tol: + sigma_alt = self.models.determinants( + self.x_best + step_alt, k_new + ) + if abs(sigma_alt) >= 0.1 * abs(sigma): + step = np.clip(step_alt, xl, xu) + + if options[Options.DEBUG]: + tol = get_arrays_tol(xl, xu) + if np.any(step + tol < xl) or np.any(xu < step - tol): + warnings.warn( + "The geometry step does not respect the bound " + "constraints.", + RuntimeWarning, + 2, + ) + if np.linalg.norm(step) > 1.1 * self.radius: + warnings.warn( + "The geometry step does not respect the " + "trust-region constraint.", + RuntimeWarning, + 2, + ) + return step + + def get_second_order_correction_step(self, step, options): + """ + Get the second-order correction step. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Trust-region step. + options : dict + Options of the solver. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Second-order correction step. + """ + # Evaluate the linearizations of the constraints. + aub, bub, aeq, beq = self.get_constraint_linearizations(self.x_best) + xl = self._pb.bounds.xl - self.x_best + xu = self._pb.bounds.xu - self.x_best + radius = np.linalg.norm(step) + soc_step = normal_byrd_omojokun( + aub, + bub, + aeq, + beq, + xl, + xu, + radius, + options[Options.DEBUG], + **self._constants, + ) + if options[Options.DEBUG]: + tol = get_arrays_tol(xl, xu) + if np.any(soc_step + tol < xl) or np.any(xu < soc_step - tol): + warnings.warn( + "The second-order correction step does not " + "respect the bound constraints.", + RuntimeWarning, + 2, + ) + if np.linalg.norm(soc_step) > 1.1 * radius: + warnings.warn( + "The second-order correction step does not " + "respect the trust-region constraint.", + RuntimeWarning, + 2, + ) + return soc_step + + def get_reduction_ratio(self, step, fun_val, cub_val, ceq_val): + """ + Get the reduction ratio. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Trust-region step. + fun_val : float + Objective function value at the trial point. + cub_val : `numpy.ndarray`, shape (m_nonlinear_ub,) + Nonlinear inequality constraint values at the trial point. + ceq_val : `numpy.ndarray`, shape (m_nonlinear_eq,) + Nonlinear equality constraint values at the trial point. + + Returns + ------- + float + Reduction ratio. + """ + merit_old = self.merit( + self.x_best, + self.fun_best, + self.cub_best, + self.ceq_best, + ) + merit_new = self.merit(self.x_best + step, fun_val, cub_val, ceq_val) + merit_model_old = self.merit( + self.x_best, + 0.0, + self.models.cub(self.x_best), + self.models.ceq(self.x_best), + ) + merit_model_new = self.merit( + self.x_best + step, + self.sqp_fun(step), + self.sqp_cub(step), + self.sqp_ceq(step), + ) + if abs(merit_model_old - merit_model_new) > TINY * abs( + merit_old - merit_new + ): + return (merit_old - merit_new) / abs( + merit_model_old - merit_model_new + ) + else: + return -1.0 + + def increase_penalty(self, step): + """ + Increase the penalty parameter. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Trust-region step. + """ + aub, bub, aeq, beq = self.get_constraint_linearizations(self.x_best) + viol_diff = max( + np.linalg.norm( + np.block( + [ + np.maximum(0.0, -bub), + beq, + ] + ) + ) + - np.linalg.norm( + np.block( + [ + np.maximum(0.0, aub @ step - bub), + aeq @ step - beq, + ] + ) + ), + 0.0, + ) + sqp_val = self.sqp_fun(step) + + threshold = np.linalg.norm( + np.block( + [ + self._lm_linear_ub, + self._lm_linear_eq, + self._lm_nonlinear_ub, + self._lm_nonlinear_eq, + ] + ) + ) + if abs(viol_diff) > TINY * abs(sqp_val): + threshold = max(threshold, sqp_val / viol_diff) + best_index_save = self.best_index + if ( + self._penalty + <= self._constants[Constants.PENALTY_INCREASE_THRESHOLD] + * threshold + ): + self._penalty = max( + self._constants[Constants.PENALTY_INCREASE_FACTOR] * threshold, + 1.0, + ) + self.set_best_index() + return best_index_save == self.best_index + + def decrease_penalty(self): + """ + Decrease the penalty parameter. + """ + self._penalty = min(self._penalty, self._get_low_penalty()) + self.set_best_index() + + def set_best_index(self): + """ + Set the index of the best point. + """ + best_index = self.best_index + m_best = self.merit( + self.x_best, + self.models.fun_val[best_index], + self.models.cub_val[best_index, :], + self.models.ceq_val[best_index, :], + ) + r_best = self._pb.maxcv( + self.x_best, + self.models.cub_val[best_index, :], + self.models.ceq_val[best_index, :], + ) + tol = ( + 10.0 + * EPS + * max(self.models.n, self.models.npt) + * max(abs(m_best), 1.0) + ) + for k in range(self.models.npt): + if k != self.best_index: + x_val = self.models.interpolation.point(k) + m_val = self.merit( + x_val, + self.models.fun_val[k], + self.models.cub_val[k, :], + self.models.ceq_val[k, :], + ) + r_val = self._pb.maxcv( + x_val, + self.models.cub_val[k, :], + self.models.ceq_val[k, :], + ) + if m_val < m_best or (m_val < m_best + tol and r_val < r_best): + best_index = k + m_best = m_val + r_best = r_val + self._best_index = best_index + + def get_index_to_remove(self, x_new=None): + """ + Get the index of the interpolation point to remove. + + If `x_new` is not provided, the index returned should be used during + the geometry-improvement phase. Otherwise, the index returned is the + best index for included `x_new` in the interpolation set. + + Parameters + ---------- + x_new : `numpy.ndarray`, shape (n,), optional + New point to be included in the interpolation set. + + Returns + ------- + int + Index of the interpolation point to remove. + float + Distance between `x_best` and the removed point. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the computation of a determinant fails. + """ + dist_sq = np.sum( + ( + self.models.interpolation.xpt + - self.models.interpolation.xpt[:, self.best_index, np.newaxis] + ) + ** 2.0, + axis=0, + ) + if x_new is None: + sigma = 1.0 + weights = dist_sq + else: + sigma = self.models.determinants(x_new) + weights = ( + np.maximum( + 1.0, + dist_sq + / max( + self._constants[Constants.LOW_RADIUS_FACTOR] + * self.radius, + self.resolution, + ) + ** 2.0, + ) + ** 3.0 + ) + weights[self.best_index] = -1.0 # do not remove the best point + k_max = np.argmax(weights * np.abs(sigma)) + return k_max, np.sqrt(dist_sq[k_max]) + + def update_radius(self, step, ratio): + """ + Update the trust-region radius. + + Parameters + ---------- + step : `numpy.ndarray`, shape (n,) + Trust-region step. + ratio : float + Reduction ratio. + """ + s_norm = np.linalg.norm(step) + if ratio <= self._constants[Constants.LOW_RATIO]: + self.radius *= self._constants[Constants.DECREASE_RADIUS_FACTOR] + elif ratio <= self._constants[Constants.HIGH_RATIO]: + self.radius = max( + self._constants[Constants.DECREASE_RADIUS_FACTOR] + * self.radius, + s_norm, + ) + else: + self.radius = min( + self._constants[Constants.INCREASE_RADIUS_FACTOR] + * self.radius, + max( + self._constants[Constants.DECREASE_RADIUS_FACTOR] + * self.radius, + self._constants[Constants.INCREASE_RADIUS_THRESHOLD] + * s_norm, + ), + ) + + def enhance_resolution(self, options): + """ + Enhance the resolution of the trust-region framework. + + Parameters + ---------- + options : dict + Options of the solver. + """ + if ( + self._constants[Constants.LARGE_RESOLUTION_THRESHOLD] + * options[Options.RHOEND] + < self.resolution + ): + self.resolution *= self._constants[ + Constants.DECREASE_RESOLUTION_FACTOR + ] + elif ( + self._constants[Constants.MODERATE_RESOLUTION_THRESHOLD] + * options[Options.RHOEND] + < self.resolution + ): + self.resolution = np.sqrt(self.resolution + * options[Options.RHOEND]) + else: + self.resolution = options[Options.RHOEND] + + # Reduce the trust-region radius. + self._radius = max( + self._constants[Constants.DECREASE_RADIUS_FACTOR] * self._radius, + self.resolution, + ) + + def shift_x_base(self, options): + """ + Shift the base point to `x_best`. + + Parameters + ---------- + options : dict + Options of the solver. + """ + self.models.shift_x_base(np.copy(self.x_best), options) + + def set_multipliers(self, x): + """ + Set the Lagrange multipliers. + + This method computes and set the Lagrange multipliers of the linear and + nonlinear constraints to be the QP multipliers. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the Lagrange multipliers are computed. + """ + # Build the constraints of the least-squares problem. + incl_linear_ub = self._pb.linear.a_ub @ x >= self._pb.linear.b_ub + incl_nonlinear_ub = self.cub_best >= 0.0 + incl_xl = self._pb.bounds.xl >= x + incl_xu = self._pb.bounds.xu <= x + m_linear_ub = np.count_nonzero(incl_linear_ub) + m_nonlinear_ub = np.count_nonzero(incl_nonlinear_ub) + m_xl = np.count_nonzero(incl_xl) + m_xu = np.count_nonzero(incl_xu) + + if ( + m_linear_ub + m_nonlinear_ub + self.m_linear_eq + + self.m_nonlinear_eq > 0 + ): + identity = np.eye(self._pb.n) + c_jac = np.r_[ + -identity[incl_xl, :], + identity[incl_xu, :], + self._pb.linear.a_ub[incl_linear_ub, :], + self.models.cub_grad(x, incl_nonlinear_ub), + self._pb.linear.a_eq, + self.models.ceq_grad(x), + ] + + # Solve the least-squares problem. + g_best = self.models.fun_grad(x) + xl_lm = np.full(c_jac.shape[0], -np.inf) + xl_lm[: m_xl + m_xu + m_linear_ub + m_nonlinear_ub] = 0.0 + res = lsq_linear( + c_jac.T, + -g_best, + bounds=(xl_lm, np.inf), + method="bvls", + ) + + # Extract the Lagrange multipliers. + self._lm_linear_ub[incl_linear_ub] = res.x[ + m_xl + m_xu:m_xl + m_xu + m_linear_ub + ] + self._lm_linear_ub[~incl_linear_ub] = 0.0 + self._lm_nonlinear_ub[incl_nonlinear_ub] = res.x[ + m_xl + + m_xu + + m_linear_ub:m_xl + + m_xu + + m_linear_ub + + m_nonlinear_ub + ] + self._lm_nonlinear_ub[~incl_nonlinear_ub] = 0.0 + self._lm_linear_eq[:] = res.x[ + m_xl + + m_xu + + m_linear_ub + + m_nonlinear_ub:m_xl + + m_xu + + m_linear_ub + + m_nonlinear_ub + + self.m_linear_eq + ] + self._lm_nonlinear_eq[:] = res.x[ + m_xl + m_xu + m_linear_ub + m_nonlinear_ub + self.m_linear_eq: + ] + + def _get_low_penalty(self): + r_val_ub = np.c_[ + ( + self.models.interpolation.x_base[np.newaxis, :] + + self.models.interpolation.xpt.T + ) + @ self._pb.linear.a_ub.T + - self._pb.linear.b_ub[np.newaxis, :], + self.models.cub_val, + ] + r_val_eq = ( + self.models.interpolation.x_base[np.newaxis, :] + + self.models.interpolation.xpt.T + ) @ self._pb.linear.a_eq.T - self._pb.linear.b_eq[np.newaxis, :] + r_val_eq = np.block( + [ + r_val_eq, + -r_val_eq, + self.models.ceq_val, + -self.models.ceq_val, + ] + ) + r_val = np.block([r_val_ub, r_val_eq]) + c_min = np.nanmin(r_val, axis=0) + c_max = np.nanmax(r_val, axis=0) + indices = ( + c_min + < self._constants[Constants.THRESHOLD_RATIO_CONSTRAINTS] * c_max + ) + if np.any(indices): + f_min = np.nanmin(self.models.fun_val) + f_max = np.nanmax(self.models.fun_val) + c_min_neg = np.minimum(0.0, c_min[indices]) + c_diff = np.min(c_max[indices] - c_min_neg) + if c_diff > TINY * (f_max - f_min): + penalty = (f_max - f_min) / c_diff + else: + penalty = np.inf + else: + penalty = 0.0 + return penalty diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/main.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/main.py new file mode 100644 index 0000000000000000000000000000000000000000..1e8fd966e40bb5fc93248f0b62cb64ccfd0f3e90 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/main.py @@ -0,0 +1,1506 @@ +import warnings + +import numpy as np +from scipy.optimize import ( + Bounds, + LinearConstraint, + NonlinearConstraint, + OptimizeResult, +) + +from .framework import TrustRegion +from .problem import ( + ObjectiveFunction, + BoundConstraints, + LinearConstraints, + NonlinearConstraints, + Problem, +) +from .utils import ( + MaxEvalError, + TargetSuccess, + CallbackSuccess, + FeasibleSuccess, + exact_1d_array, +) +from .settings import ( + ExitStatus, + Options, + Constants, + DEFAULT_OPTIONS, + DEFAULT_CONSTANTS, + PRINT_OPTIONS, +) + + +def minimize( + fun, + x0, + args=(), + bounds=None, + constraints=(), + callback=None, + options=None, + **kwargs, +): + r""" + Minimize a scalar function using the COBYQA method. + + The Constrained Optimization BY Quadratic Approximations (COBYQA) method is + a derivative-free optimization method designed to solve general nonlinear + optimization problems. A complete description of COBYQA is given in [3]_. + + Parameters + ---------- + fun : {callable, None} + Objective function to be minimized. + + ``fun(x, *args) -> float`` + + where ``x`` is an array with shape (n,) and `args` is a tuple. If `fun` + is ``None``, the objective function is assumed to be the zero function, + resulting in a feasibility problem. + x0 : array_like, shape (n,) + Initial guess. + args : tuple, optional + Extra arguments passed to the objective function. + bounds : {`scipy.optimize.Bounds`, array_like, shape (n, 2)}, optional + Bound constraints of the problem. It can be one of the cases below. + + #. An instance of `scipy.optimize.Bounds`. For the time being, the + argument ``keep_feasible`` is disregarded, and all the constraints + are considered unrelaxable and will be enforced. + #. An array with shape (n, 2). The bound constraints for ``x[i]`` are + ``bounds[i][0] <= x[i] <= bounds[i][1]``. Set ``bounds[i][0]`` to + :math:`-\infty` if there is no lower bound, and set ``bounds[i][1]`` + to :math:`\infty` if there is no upper bound. + + The COBYQA method always respect the bound constraints. + constraints : {Constraint, list}, optional + General constraints of the problem. It can be one of the cases below. + + #. An instance of `scipy.optimize.LinearConstraint`. The argument + ``keep_feasible`` is disregarded. + #. An instance of `scipy.optimize.NonlinearConstraint`. The arguments + ``jac``, ``hess``, ``keep_feasible``, ``finite_diff_rel_step``, and + ``finite_diff_jac_sparsity`` are disregarded. + + #. A list, each of whose elements are described in the cases above. + + callback : callable, optional + A callback executed at each objective function evaluation. The method + terminates if a ``StopIteration`` exception is raised by the callback + function. Its signature can be one of the following: + + ``callback(intermediate_result)`` + + where ``intermediate_result`` is a keyword parameter that contains an + instance of `scipy.optimize.OptimizeResult`, with attributes ``x`` + and ``fun``, being the point at which the objective function is + evaluated and the value of the objective function, respectively. The + name of the parameter must be ``intermediate_result`` for the callback + to be passed an instance of `scipy.optimize.OptimizeResult`. + + Alternatively, the callback function can have the signature: + + ``callback(xk)`` + + where ``xk`` is the point at which the objective function is evaluated. + Introspection is used to determine which of the signatures to invoke. + options : dict, optional + Options passed to the solver. Accepted keys are: + + disp : bool, optional + Whether to print information about the optimization procedure. + Default is ``False``. + maxfev : int, optional + Maximum number of function evaluations. Default is ``500 * n``. + maxiter : int, optional + Maximum number of iterations. Default is ``1000 * n``. + target : float, optional + Target on the objective function value. The optimization + procedure is terminated when the objective function value of a + feasible point is less than or equal to this target. Default is + ``-numpy.inf``. + feasibility_tol : float, optional + Tolerance on the constraint violation. If the maximum + constraint violation at a point is less than or equal to this + tolerance, the point is considered feasible. Default is + ``numpy.sqrt(numpy.finfo(float).eps)``. + radius_init : float, optional + Initial trust-region radius. Typically, this value should be in + the order of one tenth of the greatest expected change to `x0`. + Default is ``1.0``. + radius_final : float, optional + Final trust-region radius. It should indicate the accuracy + required in the final values of the variables. Default is + ``1e-6``. + nb_points : int, optional + Number of interpolation points used to build the quadratic + models of the objective and constraint functions. Default is + ``2 * n + 1``. + scale : bool, optional + Whether to scale the variables according to the bounds. Default + is ``False``. + filter_size : int, optional + Maximum number of points in the filter. The filter is used to + select the best point returned by the optimization procedure. + Default is ``sys.maxsize``. + store_history : bool, optional + Whether to store the history of the function evaluations. + Default is ``False``. + history_size : int, optional + Maximum number of function evaluations to store in the history. + Default is ``sys.maxsize``. + debug : bool, optional + Whether to perform additional checks during the optimization + procedure. This option should be used only for debugging + purposes and is highly discouraged to general users. Default is + ``False``. + + Other constants (from the keyword arguments) are described below. They + are not intended to be changed by general users. They should only be + changed by users with a deep understanding of the algorithm, who want + to experiment with different settings. + + Returns + ------- + `scipy.optimize.OptimizeResult` + Result of the optimization procedure, with the following fields: + + message : str + Description of the cause of the termination. + success : bool + Whether the optimization procedure terminated successfully. + status : int + Termination status of the optimization procedure. + x : `numpy.ndarray`, shape (n,) + Solution point. + fun : float + Objective function value at the solution point. + maxcv : float + Maximum constraint violation at the solution point. + nfev : int + Number of function evaluations. + nit : int + Number of iterations. + + If ``store_history`` is True, the result also has the following fields: + + fun_history : `numpy.ndarray`, shape (nfev,) + History of the objective function values. + maxcv_history : `numpy.ndarray`, shape (nfev,) + History of the maximum constraint violations. + + A description of the termination statuses is given below. + + .. list-table:: + :widths: 25 75 + :header-rows: 1 + + * - Exit status + - Description + * - 0 + - The lower bound for the trust-region radius has been reached. + * - 1 + - The target objective function value has been reached. + * - 2 + - All variables are fixed by the bound constraints. + * - 3 + - The callback requested to stop the optimization procedure. + * - 4 + - The feasibility problem received has been solved successfully. + * - 5 + - The maximum number of function evaluations has been exceeded. + * - 6 + - The maximum number of iterations has been exceeded. + * - -1 + - The bound constraints are infeasible. + * - -2 + - A linear algebra error occurred. + + Other Parameters + ---------------- + decrease_radius_factor : float, optional + Factor by which the trust-region radius is reduced when the reduction + ratio is low or negative. Default is ``0.5``. + increase_radius_factor : float, optional + Factor by which the trust-region radius is increased when the reduction + ratio is large. Default is ``numpy.sqrt(2.0)``. + increase_radius_threshold : float, optional + Threshold that controls the increase of the trust-region radius when + the reduction ratio is large. Default is ``2.0``. + decrease_radius_threshold : float, optional + Threshold used to determine whether the trust-region radius should be + reduced to the resolution. Default is ``1.4``. + decrease_resolution_factor : float, optional + Factor by which the resolution is reduced when the current value is far + from its final value. Default is ``0.1``. + large_resolution_threshold : float, optional + Threshold used to determine whether the resolution is far from its + final value. Default is ``250.0``. + moderate_resolution_threshold : float, optional + Threshold used to determine whether the resolution is close to its + final value. Default is ``16.0``. + low_ratio : float, optional + Threshold used to determine whether the reduction ratio is low. Default + is ``0.1``. + high_ratio : float, optional + Threshold used to determine whether the reduction ratio is high. + Default is ``0.7``. + very_low_ratio : float, optional + Threshold used to determine whether the reduction ratio is very low. + This is used to determine whether the models should be reset. Default + is ``0.01``. + penalty_increase_threshold : float, optional + Threshold used to determine whether the penalty parameter should be + increased. Default is ``1.5``. + penalty_increase_factor : float, optional + Factor by which the penalty parameter is increased. Default is ``2.0``. + short_step_threshold : float, optional + Factor used to determine whether the trial step is too short. Default + is ``0.5``. + low_radius_factor : float, optional + Factor used to determine which interpolation point should be removed + from the interpolation set at each iteration. Default is ``0.1``. + byrd_omojokun_factor : float, optional + Factor by which the trust-region radius is reduced for the computations + of the normal step in the Byrd-Omojokun composite-step approach. + Default is ``0.8``. + threshold_ratio_constraints : float, optional + Threshold used to determine which constraints should be taken into + account when decreasing the penalty parameter. Default is ``2.0``. + large_shift_factor : float, optional + Factor used to determine whether the point around which the quadratic + models are built should be updated. Default is ``10.0``. + large_gradient_factor : float, optional + Factor used to determine whether the models should be reset. Default is + ``10.0``. + resolution_factor : float, optional + Factor by which the resolution is decreased. Default is ``2.0``. + improve_tcg : bool, optional + Whether to improve the steps computed by the truncated conjugate + gradient method when the trust-region boundary is reached. Default is + ``True``. + + References + ---------- + .. [1] J. Nocedal and S. J. Wright. *Numerical Optimization*. Springer Ser. + Oper. Res. Financ. Eng. Springer, New York, NY, USA, second edition, + 2006. `doi:10.1007/978-0-387-40065-5 + `_. + .. [2] M. J. D. Powell. A direct search optimization method that models the + objective and constraint functions by linear interpolation. In S. Gomez + and J.-P. Hennart, editors, *Advances in Optimization and Numerical + Analysis*, volume 275 of Math. Appl., pages 51--67. Springer, Dordrecht, + Netherlands, 1994. `doi:10.1007/978-94-015-8330-5_4 + `_. + .. [3] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods + and Software*. PhD thesis, Department of Applied Mathematics, The Hong + Kong Polytechnic University, Hong Kong, China, 2022. URL: + https://theses.lib.polyu.edu.hk/handle/200/12294. + + Examples + -------- + To demonstrate how to use `minimize`, we first minimize the Rosenbrock + function implemented in `scipy.optimize` in an unconstrained setting. + + .. testsetup:: + + import numpy as np + np.set_printoptions(precision=3, suppress=True) + + >>> from cobyqa import minimize + >>> from scipy.optimize import rosen + + To solve the problem using COBYQA, run: + + >>> x0 = [1.3, 0.7, 0.8, 1.9, 1.2] + >>> res = minimize(rosen, x0) + >>> res.x + array([1., 1., 1., 1., 1.]) + + To see how bound and constraints are handled using `minimize`, we solve + Example 16.4 of [1]_, defined as + + .. math:: + + \begin{aligned} + \min_{x \in \mathbb{R}^2} & \quad (x_1 - 1)^2 + (x_2 - 2.5)^2\\ + \text{s.t.} & \quad -x_1 + 2x_2 \le 2,\\ + & \quad x_1 + 2x_2 \le 6,\\ + & \quad x_1 - 2x_2 \le 2,\\ + & \quad x_1 \ge 0,\\ + & \quad x_2 \ge 0. + \end{aligned} + + >>> import numpy as np + >>> from scipy.optimize import Bounds, LinearConstraint + + Its objective function can be implemented as: + + >>> def fun(x): + ... return (x[0] - 1.0)**2 + (x[1] - 2.5)**2 + + This problem can be solved using `minimize` as: + + >>> x0 = [2.0, 0.0] + >>> bounds = Bounds([0.0, 0.0], np.inf) + >>> constraints = LinearConstraint([ + ... [-1.0, 2.0], + ... [1.0, 2.0], + ... [1.0, -2.0], + ... ], -np.inf, [2.0, 6.0, 2.0]) + >>> res = minimize(fun, x0, bounds=bounds, constraints=constraints) + >>> res.x + array([1.4, 1.7]) + + To see how nonlinear constraints are handled, we solve Problem (F) of [2]_, + defined as + + .. math:: + + \begin{aligned} + \min_{x \in \mathbb{R}^2} & \quad -x_1 - x_2\\ + \text{s.t.} & \quad x_1^2 - x_2 \le 0,\\ + & \quad x_1^2 + x_2^2 \le 1. + \end{aligned} + + >>> from scipy.optimize import NonlinearConstraint + + Its objective and constraint functions can be implemented as: + + >>> def fun(x): + ... return -x[0] - x[1] + >>> + >>> def cub(x): + ... return [x[0]**2 - x[1], x[0]**2 + x[1]**2] + + This problem can be solved using `minimize` as: + + >>> x0 = [1.0, 1.0] + >>> constraints = NonlinearConstraint(cub, -np.inf, [0.0, 1.0]) + >>> res = minimize(fun, x0, constraints=constraints) + >>> res.x + array([0.707, 0.707]) + + Finally, to see how to supply linear and nonlinear constraints + simultaneously, we solve Problem (G) of [2]_, defined as + + .. math:: + + \begin{aligned} + \min_{x \in \mathbb{R}^3} & \quad x_3\\ + \text{s.t.} & \quad 5x_1 - x_2 + x_3 \ge 0,\\ + & \quad -5x_1 - x_2 + x_3 \ge 0,\\ + & \quad x_1^2 + x_2^2 + 4x_2 \le x_3. + \end{aligned} + + Its objective and nonlinear constraint functions can be implemented as: + + >>> def fun(x): + ... return x[2] + >>> + >>> def cub(x): + ... return x[0]**2 + x[1]**2 + 4.0*x[1] - x[2] + + This problem can be solved using `minimize` as: + + >>> x0 = [1.0, 1.0, 1.0] + >>> constraints = [ + ... LinearConstraint( + ... [[5.0, -1.0, 1.0], [-5.0, -1.0, 1.0]], + ... [0.0, 0.0], + ... np.inf, + ... ), + ... NonlinearConstraint(cub, -np.inf, 0.0), + ... ] + >>> res = minimize(fun, x0, constraints=constraints) + >>> res.x + array([ 0., -3., -3.]) + """ + # Get basic options that are needed for the initialization. + if options is None: + options = {} + else: + options = dict(options) + verbose = options.get(Options.VERBOSE, DEFAULT_OPTIONS[Options.VERBOSE]) + verbose = bool(verbose) + feasibility_tol = options.get( + Options.FEASIBILITY_TOL, + DEFAULT_OPTIONS[Options.FEASIBILITY_TOL], + ) + feasibility_tol = float(feasibility_tol) + scale = options.get(Options.SCALE, DEFAULT_OPTIONS[Options.SCALE]) + scale = bool(scale) + store_history = options.get( + Options.STORE_HISTORY, + DEFAULT_OPTIONS[Options.STORE_HISTORY], + ) + store_history = bool(store_history) + if Options.HISTORY_SIZE in options and options[Options.HISTORY_SIZE] <= 0: + raise ValueError("The size of the history must be positive.") + history_size = options.get( + Options.HISTORY_SIZE, + DEFAULT_OPTIONS[Options.HISTORY_SIZE], + ) + history_size = int(history_size) + if Options.FILTER_SIZE in options and options[Options.FILTER_SIZE] <= 0: + raise ValueError("The size of the filter must be positive.") + filter_size = options.get( + Options.FILTER_SIZE, + DEFAULT_OPTIONS[Options.FILTER_SIZE], + ) + filter_size = int(filter_size) + debug = options.get(Options.DEBUG, DEFAULT_OPTIONS[Options.DEBUG]) + debug = bool(debug) + + # Initialize the objective function. + if not isinstance(args, tuple): + args = (args,) + obj = ObjectiveFunction(fun, verbose, debug, *args) + + # Initialize the bound constraints. + if not hasattr(x0, "__len__"): + x0 = [x0] + n_orig = len(x0) + bounds = BoundConstraints(_get_bounds(bounds, n_orig)) + + # Initialize the constraints. + linear_constraints, nonlinear_constraints = _get_constraints(constraints) + linear = LinearConstraints(linear_constraints, n_orig, debug) + nonlinear = NonlinearConstraints(nonlinear_constraints, verbose, debug) + + # Initialize the problem (and remove the fixed variables). + pb = Problem( + obj, + x0, + bounds, + linear, + nonlinear, + callback, + feasibility_tol, + scale, + store_history, + history_size, + filter_size, + debug, + ) + + # Set the default options. + _set_default_options(options, pb.n) + constants = _set_default_constants(**kwargs) + + # Initialize the models and skip the computations whenever possible. + if not pb.bounds.is_feasible: + # The bound constraints are infeasible. + return _build_result( + pb, + 0.0, + False, + ExitStatus.INFEASIBLE_ERROR, + 0, + options, + ) + elif pb.n == 0: + # All variables are fixed by the bound constraints. + return _build_result( + pb, + 0.0, + True, + ExitStatus.FIXED_SUCCESS, + 0, + options, + ) + if verbose: + print("Starting the optimization procedure.") + print(f"Initial trust-region radius: {options[Options.RHOBEG]}.") + print(f"Final trust-region radius: {options[Options.RHOEND]}.") + print( + f"Maximum number of function evaluations: " + f"{options[Options.MAX_EVAL]}." + ) + print(f"Maximum number of iterations: {options[Options.MAX_ITER]}.") + print() + try: + framework = TrustRegion(pb, options, constants) + except TargetSuccess: + # The target on the objective function value has been reached + return _build_result( + pb, + 0.0, + True, + ExitStatus.TARGET_SUCCESS, + 0, + options, + ) + except CallbackSuccess: + # The callback raised a StopIteration exception. + return _build_result( + pb, + 0.0, + True, + ExitStatus.CALLBACK_SUCCESS, + 0, + options, + ) + except FeasibleSuccess: + # The feasibility problem has been solved successfully. + return _build_result( + pb, + 0.0, + True, + ExitStatus.FEASIBLE_SUCCESS, + 0, + options, + ) + except MaxEvalError: + # The maximum number of function evaluations has been exceeded. + return _build_result( + pb, + 0.0, + False, + ExitStatus.MAX_ITER_WARNING, + 0, + options, + ) + except np.linalg.LinAlgError: + # The construction of the initial interpolation set failed. + return _build_result( + pb, + 0.0, + False, + ExitStatus.LINALG_ERROR, + 0, + options, + ) + + # Start the optimization procedure. + success = False + n_iter = 0 + k_new = None + n_short_steps = 0 + n_very_short_steps = 0 + n_alt_models = 0 + while True: + # Stop the optimization procedure if the maximum number of iterations + # has been exceeded. We do not write the main loop as a for loop + # because we want to access the number of iterations outside the loop. + if n_iter >= options[Options.MAX_ITER]: + status = ExitStatus.MAX_ITER_WARNING + break + n_iter += 1 + + # Update the point around which the quadratic models are built. + if ( + np.linalg.norm( + framework.x_best - framework.models.interpolation.x_base + ) + >= constants[Constants.LARGE_SHIFT_FACTOR] * framework.radius + ): + framework.shift_x_base(options) + + # Evaluate the trial step. + radius_save = framework.radius + normal_step, tangential_step = framework.get_trust_region_step(options) + step = normal_step + tangential_step + s_norm = np.linalg.norm(step) + + # If the trial step is too short, we do not attempt to evaluate the + # objective and constraint functions. Instead, we reduce the + # trust-region radius and check whether the resolution should be + # enhanced and whether the geometry of the interpolation set should be + # improved. Otherwise, we entertain a classical iteration. The + # criterion for performing an exceptional jump is taken from NEWUOA. + if ( + s_norm + <= constants[Constants.SHORT_STEP_THRESHOLD] * framework.resolution + ): + framework.radius *= constants[Constants.DECREASE_RESOLUTION_FACTOR] + if radius_save > framework.resolution: + n_short_steps = 0 + n_very_short_steps = 0 + else: + n_short_steps += 1 + n_very_short_steps += 1 + if s_norm > 0.1 * framework.resolution: + n_very_short_steps = 0 + enhance_resolution = n_short_steps >= 5 or n_very_short_steps >= 3 + if enhance_resolution: + n_short_steps = 0 + n_very_short_steps = 0 + improve_geometry = False + else: + try: + k_new, dist_new = framework.get_index_to_remove() + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + improve_geometry = dist_new > max( + framework.radius, + constants[Constants.RESOLUTION_FACTOR] + * framework.resolution, + ) + else: + # Increase the penalty parameter if necessary. + same_best_point = framework.increase_penalty(step) + if same_best_point: + # Evaluate the objective and constraint functions. + try: + fun_val, cub_val, ceq_val = _eval( + pb, + framework, + step, + options, + ) + except TargetSuccess: + status = ExitStatus.TARGET_SUCCESS + success = True + break + except FeasibleSuccess: + status = ExitStatus.FEASIBLE_SUCCESS + success = True + break + except CallbackSuccess: + status = ExitStatus.CALLBACK_SUCCESS + success = True + break + except MaxEvalError: + status = ExitStatus.MAX_EVAL_WARNING + break + + # Perform a second-order correction step if necessary. + merit_old = framework.merit( + framework.x_best, + framework.fun_best, + framework.cub_best, + framework.ceq_best, + ) + merit_new = framework.merit( + framework.x_best + step, fun_val, cub_val, ceq_val + ) + if ( + pb.type == "nonlinearly constrained" + and merit_new > merit_old + and np.linalg.norm(normal_step) + > constants[Constants.BYRD_OMOJOKUN_FACTOR] ** 2.0 + * framework.radius + ): + soc_step = framework.get_second_order_correction_step( + step, options + ) + if np.linalg.norm(soc_step) > 0.0: + step += soc_step + + # Evaluate the objective and constraint functions. + try: + fun_val, cub_val, ceq_val = _eval( + pb, + framework, + step, + options, + ) + except TargetSuccess: + status = ExitStatus.TARGET_SUCCESS + success = True + break + except FeasibleSuccess: + status = ExitStatus.FEASIBLE_SUCCESS + success = True + break + except CallbackSuccess: + status = ExitStatus.CALLBACK_SUCCESS + success = True + break + except MaxEvalError: + status = ExitStatus.MAX_EVAL_WARNING + break + + # Calculate the reduction ratio. + ratio = framework.get_reduction_ratio( + step, + fun_val, + cub_val, + ceq_val, + ) + + # Choose an interpolation point to remove. + try: + k_new = framework.get_index_to_remove( + framework.x_best + step + )[0] + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + + # Update the interpolation set. + try: + ill_conditioned = framework.models.update_interpolation( + k_new, framework.x_best + step, fun_val, cub_val, + ceq_val + ) + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + framework.set_best_index() + + # Update the trust-region radius. + framework.update_radius(step, ratio) + + # Attempt to replace the models by the alternative ones. + if framework.radius <= framework.resolution: + if ratio >= constants[Constants.VERY_LOW_RATIO]: + n_alt_models = 0 + else: + n_alt_models += 1 + grad = framework.models.fun_grad(framework.x_best) + try: + grad_alt = framework.models.fun_alt_grad( + framework.x_best + ) + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + if np.linalg.norm(grad) < constants[ + Constants.LARGE_GRADIENT_FACTOR + ] * np.linalg.norm(grad_alt): + n_alt_models = 0 + if n_alt_models >= 3: + try: + framework.models.reset_models() + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + n_alt_models = 0 + + # Update the Lagrange multipliers. + framework.set_multipliers(framework.x_best + step) + + # Check whether the resolution should be enhanced. + try: + k_new, dist_new = framework.get_index_to_remove() + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + improve_geometry = ( + ill_conditioned + or ratio <= constants[Constants.LOW_RATIO] + and dist_new + > max( + framework.radius, + constants[Constants.RESOLUTION_FACTOR] + * framework.resolution, + ) + ) + enhance_resolution = ( + radius_save <= framework.resolution + and ratio <= constants[Constants.LOW_RATIO] + and not improve_geometry + ) + else: + # When increasing the penalty parameter, the best point so far + # may change. In this case, we restart the iteration. + enhance_resolution = False + improve_geometry = False + + # Reduce the resolution if necessary. + if enhance_resolution: + if framework.resolution <= options[Options.RHOEND]: + success = True + status = ExitStatus.RADIUS_SUCCESS + break + framework.enhance_resolution(options) + framework.decrease_penalty() + + if verbose: + maxcv_val = pb.maxcv( + framework.x_best, framework.cub_best, framework.ceq_best + ) + _print_step( + f"New trust-region radius: {framework.resolution}", + pb, + pb.build_x(framework.x_best), + framework.fun_best, + maxcv_val, + pb.n_eval, + n_iter, + ) + print() + + # Improve the geometry of the interpolation set if necessary. + if improve_geometry: + try: + step = framework.get_geometry_step(k_new, options) + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + + # Evaluate the objective and constraint functions. + try: + fun_val, cub_val, ceq_val = _eval(pb, framework, step, options) + except TargetSuccess: + status = ExitStatus.TARGET_SUCCESS + success = True + break + except FeasibleSuccess: + status = ExitStatus.FEASIBLE_SUCCESS + success = True + break + except CallbackSuccess: + status = ExitStatus.CALLBACK_SUCCESS + success = True + break + except MaxEvalError: + status = ExitStatus.MAX_EVAL_WARNING + break + + # Update the interpolation set. + try: + framework.models.update_interpolation( + k_new, + framework.x_best + step, + fun_val, + cub_val, + ceq_val, + ) + except np.linalg.LinAlgError: + status = ExitStatus.LINALG_ERROR + break + framework.set_best_index() + + return _build_result( + pb, + framework.penalty, + success, + status, + n_iter, + options, + ) + + +def _get_bounds(bounds, n): + """ + Uniformize the bounds. + """ + if bounds is None: + return Bounds(np.full(n, -np.inf), np.full(n, np.inf)) + elif isinstance(bounds, Bounds): + if bounds.lb.shape != (n,) or bounds.ub.shape != (n,): + raise ValueError(f"The bounds must have {n} elements.") + return Bounds(bounds.lb, bounds.ub) + elif hasattr(bounds, "__len__"): + bounds = np.asarray(bounds) + if bounds.shape != (n, 2): + raise ValueError( + "The shape of the bounds is not compatible with " + "the number of variables." + ) + return Bounds(bounds[:, 0], bounds[:, 1]) + else: + raise TypeError( + "The bounds must be an instance of " + "scipy.optimize.Bounds or an array-like object." + ) + + +def _get_constraints(constraints): + """ + Extract the linear and nonlinear constraints. + """ + if isinstance(constraints, dict) or not hasattr(constraints, "__len__"): + constraints = (constraints,) + + # Extract the linear and nonlinear constraints. + linear_constraints = [] + nonlinear_constraints = [] + for constraint in constraints: + if isinstance(constraint, LinearConstraint): + lb = exact_1d_array( + constraint.lb, + "The lower bound of the linear constraints must be a vector.", + ) + ub = exact_1d_array( + constraint.ub, + "The upper bound of the linear constraints must be a vector.", + ) + linear_constraints.append( + LinearConstraint( + constraint.A, + *np.broadcast_arrays(lb, ub), + ) + ) + elif isinstance(constraint, NonlinearConstraint): + lb = exact_1d_array( + constraint.lb, + "The lower bound of the " + "nonlinear constraints must be a " + "vector.", + ) + ub = exact_1d_array( + constraint.ub, + "The upper bound of the " + "nonlinear constraints must be a " + "vector.", + ) + nonlinear_constraints.append( + NonlinearConstraint( + constraint.fun, + *np.broadcast_arrays(lb, ub), + ) + ) + elif isinstance(constraint, dict): + if "type" not in constraint or constraint["type"] not in ( + "eq", + "ineq", + ): + raise ValueError('The constraint type must be "eq" or "ineq".') + if "fun" not in constraint or not callable(constraint["fun"]): + raise ValueError("The constraint function must be callable.") + nonlinear_constraints.append( + { + "fun": constraint["fun"], + "type": constraint["type"], + "args": constraint.get("args", ()), + } + ) + else: + raise TypeError( + "The constraints must be instances of " + "scipy.optimize.LinearConstraint, " + "scipy.optimize.NonlinearConstraint, or dict." + ) + return linear_constraints, nonlinear_constraints + + +def _set_default_options(options, n): + """ + Set the default options. + """ + if Options.RHOBEG in options and options[Options.RHOBEG] <= 0.0: + raise ValueError("The initial trust-region radius must be positive.") + if Options.RHOEND in options and options[Options.RHOEND] < 0.0: + raise ValueError("The final trust-region radius must be nonnegative.") + if Options.RHOBEG in options and Options.RHOEND in options: + if options[Options.RHOBEG] < options[Options.RHOEND]: + raise ValueError( + "The initial trust-region radius must be greater " + "than or equal to the final trust-region radius." + ) + elif Options.RHOBEG in options: + options[Options.RHOEND.value] = np.min( + [ + DEFAULT_OPTIONS[Options.RHOEND], + options[Options.RHOBEG], + ] + ) + elif Options.RHOEND in options: + options[Options.RHOBEG.value] = np.max( + [ + DEFAULT_OPTIONS[Options.RHOBEG], + options[Options.RHOEND], + ] + ) + else: + options[Options.RHOBEG.value] = DEFAULT_OPTIONS[Options.RHOBEG] + options[Options.RHOEND.value] = DEFAULT_OPTIONS[Options.RHOEND] + options[Options.RHOBEG.value] = float(options[Options.RHOBEG]) + options[Options.RHOEND.value] = float(options[Options.RHOEND]) + if Options.NPT in options and options[Options.NPT] <= 0: + raise ValueError("The number of interpolation points must be " + "positive.") + if ( + Options.NPT in options + and options[Options.NPT] > ((n + 1) * (n + 2)) // 2 + ): + raise ValueError( + f"The number of interpolation points must be at most " + f"{((n + 1) * (n + 2)) // 2}." + ) + options.setdefault(Options.NPT.value, DEFAULT_OPTIONS[Options.NPT](n)) + options[Options.NPT.value] = int(options[Options.NPT]) + if Options.MAX_EVAL in options and options[Options.MAX_EVAL] <= 0: + raise ValueError( + "The maximum number of function evaluations must be positive." + ) + options.setdefault( + Options.MAX_EVAL.value, + np.max( + [ + DEFAULT_OPTIONS[Options.MAX_EVAL](n), + options[Options.NPT] + 1, + ] + ), + ) + options[Options.MAX_EVAL.value] = int(options[Options.MAX_EVAL]) + if Options.MAX_ITER in options and options[Options.MAX_ITER] <= 0: + raise ValueError("The maximum number of iterations must be positive.") + options.setdefault( + Options.MAX_ITER.value, + DEFAULT_OPTIONS[Options.MAX_ITER](n), + ) + options[Options.MAX_ITER.value] = int(options[Options.MAX_ITER]) + options.setdefault(Options.TARGET.value, DEFAULT_OPTIONS[Options.TARGET]) + options[Options.TARGET.value] = float(options[Options.TARGET]) + options.setdefault( + Options.FEASIBILITY_TOL.value, + DEFAULT_OPTIONS[Options.FEASIBILITY_TOL], + ) + options[Options.FEASIBILITY_TOL.value] = float( + options[Options.FEASIBILITY_TOL] + ) + options.setdefault(Options.VERBOSE.value, DEFAULT_OPTIONS[Options.VERBOSE]) + options[Options.VERBOSE.value] = bool(options[Options.VERBOSE]) + options.setdefault(Options.SCALE.value, DEFAULT_OPTIONS[Options.SCALE]) + options[Options.SCALE.value] = bool(options[Options.SCALE]) + options.setdefault( + Options.FILTER_SIZE.value, + DEFAULT_OPTIONS[Options.FILTER_SIZE], + ) + options[Options.FILTER_SIZE.value] = int(options[Options.FILTER_SIZE]) + options.setdefault( + Options.STORE_HISTORY.value, + DEFAULT_OPTIONS[Options.STORE_HISTORY], + ) + options[Options.STORE_HISTORY.value] = bool(options[Options.STORE_HISTORY]) + options.setdefault( + Options.HISTORY_SIZE.value, + DEFAULT_OPTIONS[Options.HISTORY_SIZE], + ) + options[Options.HISTORY_SIZE.value] = int(options[Options.HISTORY_SIZE]) + options.setdefault(Options.DEBUG.value, DEFAULT_OPTIONS[Options.DEBUG]) + options[Options.DEBUG.value] = bool(options[Options.DEBUG]) + + # Check whether they are any unknown options. + for key in options: + if key not in Options.__members__.values(): + warnings.warn(f"Unknown option: {key}.", RuntimeWarning, 3) + + +def _set_default_constants(**kwargs): + """ + Set the default constants. + """ + constants = dict(kwargs) + constants.setdefault( + Constants.DECREASE_RADIUS_FACTOR.value, + DEFAULT_CONSTANTS[Constants.DECREASE_RADIUS_FACTOR], + ) + constants[Constants.DECREASE_RADIUS_FACTOR.value] = float( + constants[Constants.DECREASE_RADIUS_FACTOR] + ) + if ( + constants[Constants.DECREASE_RADIUS_FACTOR] <= 0.0 + or constants[Constants.DECREASE_RADIUS_FACTOR] >= 1.0 + ): + raise ValueError( + "The constant decrease_radius_factor must be in the interval " + "(0, 1)." + ) + constants.setdefault( + Constants.INCREASE_RADIUS_THRESHOLD.value, + DEFAULT_CONSTANTS[Constants.INCREASE_RADIUS_THRESHOLD], + ) + constants[Constants.INCREASE_RADIUS_THRESHOLD.value] = float( + constants[Constants.INCREASE_RADIUS_THRESHOLD] + ) + if constants[Constants.INCREASE_RADIUS_THRESHOLD] <= 1.0: + raise ValueError( + "The constant increase_radius_threshold must be greater than 1." + ) + if ( + Constants.INCREASE_RADIUS_FACTOR in constants + and constants[Constants.INCREASE_RADIUS_FACTOR] <= 1.0 + ): + raise ValueError( + "The constant increase_radius_factor must be greater than 1." + ) + if ( + Constants.DECREASE_RADIUS_THRESHOLD in constants + and constants[Constants.DECREASE_RADIUS_THRESHOLD] <= 1.0 + ): + raise ValueError( + "The constant decrease_radius_threshold must be greater than 1." + ) + if ( + Constants.INCREASE_RADIUS_FACTOR in constants + and Constants.DECREASE_RADIUS_THRESHOLD in constants + ): + if ( + constants[Constants.DECREASE_RADIUS_THRESHOLD] + >= constants[Constants.INCREASE_RADIUS_FACTOR] + ): + raise ValueError( + "The constant decrease_radius_threshold must be " + "less than increase_radius_factor." + ) + elif Constants.INCREASE_RADIUS_FACTOR in constants: + constants[Constants.DECREASE_RADIUS_THRESHOLD.value] = np.min( + [ + DEFAULT_CONSTANTS[Constants.DECREASE_RADIUS_THRESHOLD], + 0.5 * (1.0 + constants[Constants.INCREASE_RADIUS_FACTOR]), + ] + ) + elif Constants.DECREASE_RADIUS_THRESHOLD in constants: + constants[Constants.INCREASE_RADIUS_FACTOR.value] = np.max( + [ + DEFAULT_CONSTANTS[Constants.INCREASE_RADIUS_FACTOR], + 2.0 * constants[Constants.DECREASE_RADIUS_THRESHOLD], + ] + ) + else: + constants[Constants.INCREASE_RADIUS_FACTOR.value] = DEFAULT_CONSTANTS[ + Constants.INCREASE_RADIUS_FACTOR + ] + constants[Constants.DECREASE_RADIUS_THRESHOLD.value] = ( + DEFAULT_CONSTANTS[Constants.DECREASE_RADIUS_THRESHOLD]) + constants.setdefault( + Constants.DECREASE_RESOLUTION_FACTOR.value, + DEFAULT_CONSTANTS[Constants.DECREASE_RESOLUTION_FACTOR], + ) + constants[Constants.DECREASE_RESOLUTION_FACTOR.value] = float( + constants[Constants.DECREASE_RESOLUTION_FACTOR] + ) + if ( + constants[Constants.DECREASE_RESOLUTION_FACTOR] <= 0.0 + or constants[Constants.DECREASE_RESOLUTION_FACTOR] >= 1.0 + ): + raise ValueError( + "The constant decrease_resolution_factor must be in the interval " + "(0, 1)." + ) + if ( + Constants.LARGE_RESOLUTION_THRESHOLD in constants + and constants[Constants.LARGE_RESOLUTION_THRESHOLD] <= 1.0 + ): + raise ValueError( + "The constant large_resolution_threshold must be greater than 1." + ) + if ( + Constants.MODERATE_RESOLUTION_THRESHOLD in constants + and constants[Constants.MODERATE_RESOLUTION_THRESHOLD] <= 1.0 + ): + raise ValueError( + "The constant moderate_resolution_threshold must be greater than " + "1." + ) + if ( + Constants.LARGE_RESOLUTION_THRESHOLD in constants + and Constants.MODERATE_RESOLUTION_THRESHOLD in constants + ): + if ( + constants[Constants.MODERATE_RESOLUTION_THRESHOLD] + > constants[Constants.LARGE_RESOLUTION_THRESHOLD] + ): + raise ValueError( + "The constant moderate_resolution_threshold " + "must be at most large_resolution_threshold." + ) + elif Constants.LARGE_RESOLUTION_THRESHOLD in constants: + constants[Constants.MODERATE_RESOLUTION_THRESHOLD.value] = np.min( + [ + DEFAULT_CONSTANTS[Constants.MODERATE_RESOLUTION_THRESHOLD], + constants[Constants.LARGE_RESOLUTION_THRESHOLD], + ] + ) + elif Constants.MODERATE_RESOLUTION_THRESHOLD in constants: + constants[Constants.LARGE_RESOLUTION_THRESHOLD.value] = np.max( + [ + DEFAULT_CONSTANTS[Constants.LARGE_RESOLUTION_THRESHOLD], + constants[Constants.MODERATE_RESOLUTION_THRESHOLD], + ] + ) + else: + constants[Constants.LARGE_RESOLUTION_THRESHOLD.value] = ( + DEFAULT_CONSTANTS[Constants.LARGE_RESOLUTION_THRESHOLD] + ) + constants[Constants.MODERATE_RESOLUTION_THRESHOLD.value] = ( + DEFAULT_CONSTANTS[Constants.MODERATE_RESOLUTION_THRESHOLD] + ) + if Constants.LOW_RATIO in constants and ( + constants[Constants.LOW_RATIO] <= 0.0 + or constants[Constants.LOW_RATIO] >= 1.0 + ): + raise ValueError( + "The constant low_ratio must be in the interval (0, 1)." + ) + if Constants.HIGH_RATIO in constants and ( + constants[Constants.HIGH_RATIO] <= 0.0 + or constants[Constants.HIGH_RATIO] >= 1.0 + ): + raise ValueError( + "The constant high_ratio must be in the interval (0, 1)." + ) + if Constants.LOW_RATIO in constants and Constants.HIGH_RATIO in constants: + if constants[Constants.LOW_RATIO] > constants[Constants.HIGH_RATIO]: + raise ValueError( + "The constant low_ratio must be at most high_ratio." + ) + elif Constants.LOW_RATIO in constants: + constants[Constants.HIGH_RATIO.value] = np.max( + [ + DEFAULT_CONSTANTS[Constants.HIGH_RATIO], + constants[Constants.LOW_RATIO], + ] + ) + elif Constants.HIGH_RATIO in constants: + constants[Constants.LOW_RATIO.value] = np.min( + [ + DEFAULT_CONSTANTS[Constants.LOW_RATIO], + constants[Constants.HIGH_RATIO], + ] + ) + else: + constants[Constants.LOW_RATIO.value] = DEFAULT_CONSTANTS[ + Constants.LOW_RATIO + ] + constants[Constants.HIGH_RATIO.value] = DEFAULT_CONSTANTS[ + Constants.HIGH_RATIO + ] + constants.setdefault( + Constants.VERY_LOW_RATIO.value, + DEFAULT_CONSTANTS[Constants.VERY_LOW_RATIO], + ) + constants[Constants.VERY_LOW_RATIO.value] = float( + constants[Constants.VERY_LOW_RATIO] + ) + if ( + constants[Constants.VERY_LOW_RATIO] <= 0.0 + or constants[Constants.VERY_LOW_RATIO] >= 1.0 + ): + raise ValueError( + "The constant very_low_ratio must be in the interval (0, 1)." + ) + if ( + Constants.PENALTY_INCREASE_THRESHOLD in constants + and constants[Constants.PENALTY_INCREASE_THRESHOLD] < 1.0 + ): + raise ValueError( + "The constant penalty_increase_threshold must be " + "greater than or equal to 1." + ) + if ( + Constants.PENALTY_INCREASE_FACTOR in constants + and constants[Constants.PENALTY_INCREASE_FACTOR] <= 1.0 + ): + raise ValueError( + "The constant penalty_increase_factor must be greater than 1." + ) + if ( + Constants.PENALTY_INCREASE_THRESHOLD in constants + and Constants.PENALTY_INCREASE_FACTOR in constants + ): + if ( + constants[Constants.PENALTY_INCREASE_FACTOR] + < constants[Constants.PENALTY_INCREASE_THRESHOLD] + ): + raise ValueError( + "The constant penalty_increase_factor must be " + "greater than or equal to " + "penalty_increase_threshold." + ) + elif Constants.PENALTY_INCREASE_THRESHOLD in constants: + constants[Constants.PENALTY_INCREASE_FACTOR.value] = np.max( + [ + DEFAULT_CONSTANTS[Constants.PENALTY_INCREASE_FACTOR], + constants[Constants.PENALTY_INCREASE_THRESHOLD], + ] + ) + elif Constants.PENALTY_INCREASE_FACTOR in constants: + constants[Constants.PENALTY_INCREASE_THRESHOLD.value] = np.min( + [ + DEFAULT_CONSTANTS[Constants.PENALTY_INCREASE_THRESHOLD], + constants[Constants.PENALTY_INCREASE_FACTOR], + ] + ) + else: + constants[Constants.PENALTY_INCREASE_THRESHOLD.value] = ( + DEFAULT_CONSTANTS[Constants.PENALTY_INCREASE_THRESHOLD] + ) + constants[Constants.PENALTY_INCREASE_FACTOR.value] = DEFAULT_CONSTANTS[ + Constants.PENALTY_INCREASE_FACTOR + ] + constants.setdefault( + Constants.SHORT_STEP_THRESHOLD.value, + DEFAULT_CONSTANTS[Constants.SHORT_STEP_THRESHOLD], + ) + constants[Constants.SHORT_STEP_THRESHOLD.value] = float( + constants[Constants.SHORT_STEP_THRESHOLD] + ) + if ( + constants[Constants.SHORT_STEP_THRESHOLD] <= 0.0 + or constants[Constants.SHORT_STEP_THRESHOLD] >= 1.0 + ): + raise ValueError( + "The constant short_step_threshold must be in the interval (0, 1)." + ) + constants.setdefault( + Constants.LOW_RADIUS_FACTOR.value, + DEFAULT_CONSTANTS[Constants.LOW_RADIUS_FACTOR], + ) + constants[Constants.LOW_RADIUS_FACTOR.value] = float( + constants[Constants.LOW_RADIUS_FACTOR] + ) + if ( + constants[Constants.LOW_RADIUS_FACTOR] <= 0.0 + or constants[Constants.LOW_RADIUS_FACTOR] >= 1.0 + ): + raise ValueError( + "The constant low_radius_factor must be in the interval (0, 1)." + ) + constants.setdefault( + Constants.BYRD_OMOJOKUN_FACTOR.value, + DEFAULT_CONSTANTS[Constants.BYRD_OMOJOKUN_FACTOR], + ) + constants[Constants.BYRD_OMOJOKUN_FACTOR.value] = float( + constants[Constants.BYRD_OMOJOKUN_FACTOR] + ) + if ( + constants[Constants.BYRD_OMOJOKUN_FACTOR] <= 0.0 + or constants[Constants.BYRD_OMOJOKUN_FACTOR] >= 1.0 + ): + raise ValueError( + "The constant byrd_omojokun_factor must be in the interval (0, 1)." + ) + constants.setdefault( + Constants.THRESHOLD_RATIO_CONSTRAINTS.value, + DEFAULT_CONSTANTS[Constants.THRESHOLD_RATIO_CONSTRAINTS], + ) + constants[Constants.THRESHOLD_RATIO_CONSTRAINTS.value] = float( + constants[Constants.THRESHOLD_RATIO_CONSTRAINTS] + ) + if constants[Constants.THRESHOLD_RATIO_CONSTRAINTS] <= 1.0: + raise ValueError( + "The constant threshold_ratio_constraints must be greater than 1." + ) + constants.setdefault( + Constants.LARGE_SHIFT_FACTOR.value, + DEFAULT_CONSTANTS[Constants.LARGE_SHIFT_FACTOR], + ) + constants[Constants.LARGE_SHIFT_FACTOR.value] = float( + constants[Constants.LARGE_SHIFT_FACTOR] + ) + if constants[Constants.LARGE_SHIFT_FACTOR] < 0.0: + raise ValueError("The constant large_shift_factor must be " + "nonnegative.") + constants.setdefault( + Constants.LARGE_GRADIENT_FACTOR.value, + DEFAULT_CONSTANTS[Constants.LARGE_GRADIENT_FACTOR], + ) + constants[Constants.LARGE_GRADIENT_FACTOR.value] = float( + constants[Constants.LARGE_GRADIENT_FACTOR] + ) + if constants[Constants.LARGE_GRADIENT_FACTOR] <= 1.0: + raise ValueError( + "The constant large_gradient_factor must be greater than 1." + ) + constants.setdefault( + Constants.RESOLUTION_FACTOR.value, + DEFAULT_CONSTANTS[Constants.RESOLUTION_FACTOR], + ) + constants[Constants.RESOLUTION_FACTOR.value] = float( + constants[Constants.RESOLUTION_FACTOR] + ) + if constants[Constants.RESOLUTION_FACTOR] <= 1.0: + raise ValueError( + "The constant resolution_factor must be greater than 1." + ) + constants.setdefault( + Constants.IMPROVE_TCG.value, + DEFAULT_CONSTANTS[Constants.IMPROVE_TCG], + ) + constants[Constants.IMPROVE_TCG.value] = bool( + constants[Constants.IMPROVE_TCG] + ) + + # Check whether they are any unknown options. + for key in kwargs: + if key not in Constants.__members__.values(): + warnings.warn(f"Unknown constant: {key}.", RuntimeWarning, 3) + return constants + + +def _eval(pb, framework, step, options): + """ + Evaluate the objective and constraint functions. + """ + if pb.n_eval >= options[Options.MAX_EVAL]: + raise MaxEvalError + x_eval = framework.x_best + step + fun_val, cub_val, ceq_val = pb(x_eval, framework.penalty) + r_val = pb.maxcv(x_eval, cub_val, ceq_val) + if ( + fun_val <= options[Options.TARGET] + and r_val <= options[Options.FEASIBILITY_TOL] + ): + raise TargetSuccess + if pb.is_feasibility and r_val <= options[Options.FEASIBILITY_TOL]: + raise FeasibleSuccess + return fun_val, cub_val, ceq_val + + +def _build_result(pb, penalty, success, status, n_iter, options): + """ + Build the result of the optimization process. + """ + # Build the result. + x, fun, maxcv = pb.best_eval(penalty) + success = success and np.isfinite(fun) and np.isfinite(maxcv) + if status not in [ExitStatus.TARGET_SUCCESS, ExitStatus.FEASIBLE_SUCCESS]: + success = success and maxcv <= options[Options.FEASIBILITY_TOL] + result = OptimizeResult() + result.message = { + ExitStatus.RADIUS_SUCCESS: "The lower bound for the trust-region " + "radius has been reached", + ExitStatus.TARGET_SUCCESS: "The target objective function value has " + "been reached", + ExitStatus.FIXED_SUCCESS: "All variables are fixed by the bound " + "constraints", + ExitStatus.CALLBACK_SUCCESS: "The callback requested to stop the " + "optimization procedure", + ExitStatus.FEASIBLE_SUCCESS: "The feasibility problem received has " + "been solved successfully", + ExitStatus.MAX_EVAL_WARNING: "The maximum number of function " + "evaluations has been exceeded", + ExitStatus.MAX_ITER_WARNING: "The maximum number of iterations has " + "been exceeded", + ExitStatus.INFEASIBLE_ERROR: "The bound constraints are infeasible", + ExitStatus.LINALG_ERROR: "A linear algebra error occurred", + }.get(status, "Unknown exit status") + result.success = success + result.status = status.value + result.x = pb.build_x(x) + result.fun = fun + result.maxcv = maxcv + result.nfev = pb.n_eval + result.nit = n_iter + if options[Options.STORE_HISTORY]: + result.fun_history = pb.fun_history + result.maxcv_history = pb.maxcv_history + + # Print the result if requested. + if options[Options.VERBOSE]: + _print_step( + result.message, + pb, + result.x, + result.fun, + result.maxcv, + result.nfev, + result.nit, + ) + return result + + +def _print_step(message, pb, x, fun_val, r_val, n_eval, n_iter): + """ + Print information about the current state of the optimization process. + """ + print() + print(f"{message}.") + print(f"Number of function evaluations: {n_eval}.") + print(f"Number of iterations: {n_iter}.") + if not pb.is_feasibility: + print(f"Least value of {pb.fun_name}: {fun_val}.") + print(f"Maximum constraint violation: {r_val}.") + with np.printoptions(**PRINT_OPTIONS): + print(f"Corresponding point: {x}.") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/models.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/models.py new file mode 100644 index 0000000000000000000000000000000000000000..8917e594e6d13921185f27e9b34167845aa89371 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/models.py @@ -0,0 +1,1529 @@ +import warnings + +import numpy as np +from scipy.linalg import eigh + +from .settings import Options +from .utils import MaxEvalError, TargetSuccess, FeasibleSuccess + + +EPS = np.finfo(float).eps + + +class Interpolation: + """ + Interpolation set. + + This class stores a base point around which the models are expanded and the + interpolation points. The coordinates of the interpolation points are + relative to the base point. + """ + + def __init__(self, pb, options): + """ + Initialize the interpolation set. + + Parameters + ---------- + pb : `cobyqa.problem.Problem` + Problem to be solved. + options : dict + Options of the solver. + """ + # Reduce the initial trust-region radius if necessary. + self._debug = options[Options.DEBUG] + max_radius = 0.5 * np.min(pb.bounds.xu - pb.bounds.xl) + if options[Options.RHOBEG] > max_radius: + options[Options.RHOBEG.value] = max_radius + options[Options.RHOEND.value] = np.min( + [ + options[Options.RHOEND], + max_radius, + ] + ) + + # Set the initial point around which the models are expanded. + self._x_base = np.copy(pb.x0) + very_close_xl_idx = ( + self.x_base <= pb.bounds.xl + 0.5 * options[Options.RHOBEG] + ) + self.x_base[very_close_xl_idx] = pb.bounds.xl[very_close_xl_idx] + close_xl_idx = ( + pb.bounds.xl + 0.5 * options[Options.RHOBEG] < self.x_base + ) & (self.x_base <= pb.bounds.xl + options[Options.RHOBEG]) + self.x_base[close_xl_idx] = np.minimum( + pb.bounds.xl[close_xl_idx] + options[Options.RHOBEG], + pb.bounds.xu[close_xl_idx], + ) + very_close_xu_idx = ( + self.x_base >= pb.bounds.xu - 0.5 * options[Options.RHOBEG] + ) + self.x_base[very_close_xu_idx] = pb.bounds.xu[very_close_xu_idx] + close_xu_idx = ( + self.x_base < pb.bounds.xu - 0.5 * options[Options.RHOBEG] + ) & (pb.bounds.xu - options[Options.RHOBEG] <= self.x_base) + self.x_base[close_xu_idx] = np.maximum( + pb.bounds.xu[close_xu_idx] - options[Options.RHOBEG], + pb.bounds.xl[close_xu_idx], + ) + + # Set the initial interpolation set. + self._xpt = np.zeros((pb.n, options[Options.NPT])) + for k in range(1, options[Options.NPT]): + if k <= pb.n: + if very_close_xu_idx[k - 1]: + self.xpt[k - 1, k] = -options[Options.RHOBEG] + else: + self.xpt[k - 1, k] = options[Options.RHOBEG] + elif k <= 2 * pb.n: + if very_close_xl_idx[k - pb.n - 1]: + self.xpt[k - pb.n - 1, k] = 2.0 * options[Options.RHOBEG] + elif very_close_xu_idx[k - pb.n - 1]: + self.xpt[k - pb.n - 1, k] = -2.0 * options[Options.RHOBEG] + else: + self.xpt[k - pb.n - 1, k] = -options[Options.RHOBEG] + else: + spread = (k - pb.n - 1) // pb.n + k1 = k - (1 + spread) * pb.n - 1 + k2 = (k1 + spread) % pb.n + self.xpt[k1, k] = self.xpt[k1, k1 + 1] + self.xpt[k2, k] = self.xpt[k2, k2 + 1] + + @property + def n(self): + """ + Number of variables. + + Returns + ------- + int + Number of variables. + """ + return self.xpt.shape[0] + + @property + def npt(self): + """ + Number of interpolation points. + + Returns + ------- + int + Number of interpolation points. + """ + return self.xpt.shape[1] + + @property + def xpt(self): + """ + Interpolation points. + + Returns + ------- + `numpy.ndarray`, shape (n, npt) + Interpolation points. + """ + return self._xpt + + @xpt.setter + def xpt(self, xpt): + """ + Set the interpolation points. + + Parameters + ---------- + xpt : `numpy.ndarray`, shape (n, npt) + New interpolation points. + """ + if self._debug: + assert xpt.shape == ( + self.n, + self.npt, + ), "The shape of `xpt` is not valid." + self._xpt = xpt + + @property + def x_base(self): + """ + Base point around which the models are expanded. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Base point around which the models are expanded. + """ + return self._x_base + + @x_base.setter + def x_base(self, x_base): + """ + Set the base point around which the models are expanded. + + Parameters + ---------- + x_base : `numpy.ndarray`, shape (n,) + New base point around which the models are expanded. + """ + if self._debug: + assert x_base.shape == ( + self.n, + ), "The shape of `x_base` is not valid." + self._x_base = x_base + + def point(self, k): + """ + Get the `k`-th interpolation point. + + The return point is relative to the origin. + + Parameters + ---------- + k : int + Index of the interpolation point. + + Returns + ------- + `numpy.ndarray`, shape (n,) + `k`-th interpolation point. + """ + if self._debug: + assert 0 <= k < self.npt, "The index `k` is not valid." + return self.x_base + self.xpt[:, k] + + +_cache = {"xpt": None, "a": None, "right_scaling": None, "eigh": None} + + +def build_system(interpolation): + """ + Build the left-hand side matrix of the interpolation system. The + matrix below stores W * diag(right_scaling), + where W is the theoretical matrix of the interpolation system. The + right scaling matrices is chosen to keep the elements in + the matrix well-balanced. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + """ + + # Compute the scaled directions from the base point to the + # interpolation points. We scale the directions to avoid numerical + # difficulties. + if _cache["xpt"] is not None and np.array_equal( + interpolation.xpt, _cache["xpt"] + ): + return _cache["a"], _cache["right_scaling"], _cache["eigh"] + + scale = np.max(np.linalg.norm(interpolation.xpt, axis=0), initial=EPS) + xpt_scale = interpolation.xpt / scale + + n, npt = xpt_scale.shape + a = np.zeros((npt + n + 1, npt + n + 1)) + a[:npt, :npt] = 0.5 * (xpt_scale.T @ xpt_scale) ** 2.0 + a[:npt, npt] = 1.0 + a[:npt, npt + 1:] = xpt_scale.T + a[npt, :npt] = 1.0 + a[npt + 1:, :npt] = xpt_scale + + # Build the left and right scaling diagonal matrices. + right_scaling = np.empty(npt + n + 1) + right_scaling[:npt] = 1.0 / scale**2.0 + right_scaling[npt] = scale**2.0 + right_scaling[npt + 1:] = scale + + eig_values, eig_vectors = eigh(a, check_finite=False) + + _cache["xpt"] = np.copy(interpolation.xpt) + _cache["a"] = np.copy(a) + _cache["right_scaling"] = np.copy(right_scaling) + _cache["eigh"] = (eig_values, eig_vectors) + + return a, right_scaling, (eig_values, eig_vectors) + + +class Quadratic: + """ + Quadratic model. + + This class stores the Hessian matrix of the quadratic model using the + implicit/explicit representation designed by Powell for NEWUOA [1]_. + + References + ---------- + .. [1] M. J. D. Powell. The NEWUOA software for unconstrained optimization + without derivatives. In G. Di Pillo and M. Roma, editors, *Large-Scale + Nonlinear Optimization*, volume 83 of Nonconvex Optim. Appl., pages + 255--297. Springer, Boston, MA, USA, 2006. `doi:10.1007/0-387-30065-1_16 + `_. + """ + + def __init__(self, interpolation, values, debug): + """ + Initialize the quadratic model. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + values : `numpy.ndarray`, shape (npt,) + Values of the interpolated function at the interpolation points. + debug : bool + Whether to make debugging tests during the execution. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + self._debug = debug + if self._debug: + assert values.shape == ( + interpolation.npt, + ), "The shape of `values` is not valid." + if interpolation.npt < interpolation.n + 1: + raise ValueError( + f"The number of interpolation points must be at least " + f"{interpolation.n + 1}." + ) + self._const, self._grad, self._i_hess, _ = self._get_model( + interpolation, + values, + ) + self._e_hess = np.zeros((self.n, self.n)) + + def __call__(self, x, interpolation): + """ + Evaluate the quadratic model at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the quadratic model is evaluated. + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + + Returns + ------- + float + Value of the quadratic model at `x`. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + x_diff = x - interpolation.x_base + return ( + self._const + + self._grad @ x_diff + + 0.5 + * ( + self._i_hess @ (interpolation.xpt.T @ x_diff) ** 2.0 + + x_diff @ self._e_hess @ x_diff + ) + ) + + @property + def n(self): + """ + Number of variables. + + Returns + ------- + int + Number of variables. + """ + return self._grad.size + + @property + def npt(self): + """ + Number of interpolation points used to define the quadratic model. + + Returns + ------- + int + Number of interpolation points used to define the quadratic model. + """ + return self._i_hess.size + + def grad(self, x, interpolation): + """ + Evaluate the gradient of the quadratic model at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which the gradient of the quadratic model is evaluated. + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Gradient of the quadratic model at `x`. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + x_diff = x - interpolation.x_base + return self._grad + self.hess_prod(x_diff, interpolation) + + def hess(self, interpolation): + """ + Evaluate the Hessian matrix of the quadratic model. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + + Returns + ------- + `numpy.ndarray`, shape (n, n) + Hessian matrix of the quadratic model. + """ + return self._e_hess + interpolation.xpt @ ( + self._i_hess[:, np.newaxis] * interpolation.xpt.T + ) + + def hess_prod(self, v, interpolation): + """ + Evaluate the right product of the Hessian matrix of the quadratic model + with a given vector. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Vector with which the Hessian matrix of the quadratic model is + multiplied from the right. + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Right product of the Hessian matrix of the quadratic model with + `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + return self._e_hess @ v + interpolation.xpt @ ( + self._i_hess * (interpolation.xpt.T @ v) + ) + + def curv(self, v, interpolation): + """ + Evaluate the curvature of the quadratic model along a given direction. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Direction along which the curvature of the quadratic model is + evaluated. + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + + Returns + ------- + float + Curvature of the quadratic model along `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + return ( + v @ self._e_hess @ v + + self._i_hess @ (interpolation.xpt.T @ v) ** 2.0 + ) + + def update(self, interpolation, k_new, dir_old, values_diff): + """ + Update the quadratic model. + + This method applies the derivative-free symmetric Broyden update to the + quadratic model. The `knew`-th interpolation point must be updated + before calling this method. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Updated interpolation set. + k_new : int + Index of the updated interpolation point. + dir_old : `numpy.ndarray`, shape (n,) + Value of ``interpolation.xpt[:, k_new]`` before the update. + values_diff : `numpy.ndarray`, shape (npt,) + Differences between the values of the interpolated nonlinear + function and the previous quadratic model at the updated + interpolation points. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + if self._debug: + assert 0 <= k_new < self.npt, "The index `k_new` is not valid." + assert dir_old.shape == ( + self.n, + ), "The shape of `dir_old` is not valid." + assert values_diff.shape == ( + self.npt, + ), "The shape of `values_diff` is not valid." + + # Forward the k_new-th element of the implicit Hessian matrix to the + # explicit Hessian matrix. This must be done because the implicit + # Hessian matrix is related to the interpolation points, and the + # k_new-th interpolation point is modified. + self._e_hess += self._i_hess[k_new] * np.outer(dir_old, dir_old) + self._i_hess[k_new] = 0.0 + + # Update the quadratic model. + const, grad, i_hess, ill_conditioned = self._get_model( + interpolation, + values_diff, + ) + self._const += const + self._grad += grad + self._i_hess += i_hess + return ill_conditioned + + def shift_x_base(self, interpolation, new_x_base): + """ + Shift the point around which the quadratic model is defined. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Previous interpolation set. + new_x_base : `numpy.ndarray`, shape (n,) + Point that will replace ``interpolation.x_base``. + """ + if self._debug: + assert new_x_base.shape == ( + self.n, + ), "The shape of `new_x_base` is not valid." + self._const = self(new_x_base, interpolation) + self._grad = self.grad(new_x_base, interpolation) + shift = new_x_base - interpolation.x_base + update = np.outer( + shift, + (interpolation.xpt - 0.5 * shift[:, np.newaxis]) @ self._i_hess, + ) + self._e_hess += update + update.T + + @staticmethod + def solve_systems(interpolation, rhs): + """ + Solve the interpolation systems. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + rhs : `numpy.ndarray`, shape (npt + n + 1, m) + Right-hand side vectors of the ``m`` interpolation systems. + + Returns + ------- + `numpy.ndarray`, shape (npt + n + 1, m) + Solutions of the interpolation systems. + `numpy.ndarray`, shape (m, ) + Whether the interpolation systems are ill-conditioned. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation systems are ill-defined. + """ + n, npt = interpolation.xpt.shape + assert ( + rhs.ndim == 2 and rhs.shape[0] == npt + n + 1 + ), "The shape of `rhs` is not valid." + + # Build the left-hand side matrix of the interpolation system. The + # matrix below stores diag(left_scaling) * W * diag(right_scaling), + # where W is the theoretical matrix of the interpolation system. The + # left and right scaling matrices are chosen to keep the elements in + # the matrix well-balanced. + a, right_scaling, eig = build_system(interpolation) + + # Build the solution. After a discussion with Mike Saunders and Alexis + # Montoison during their visit to the Hong Kong Polytechnic University + # in 2024, we decided to use the eigendecomposition of the symmetric + # matrix a. This is more stable than the previously employed LBL + # decomposition, and allows us to directly detect ill-conditioning of + # the system and to build the least-squares solution if necessary. + # Numerical experiments have shown that this strategy improves the + # performance of the solver. + rhs_scaled = rhs * right_scaling[:, np.newaxis] + if not (np.all(np.isfinite(a)) and np.all(np.isfinite(rhs_scaled))): + raise np.linalg.LinAlgError( + "The interpolation system is ill-defined." + ) + + # calculated in build_system + eig_values, eig_vectors = eig + + large_eig_values = np.abs(eig_values) > EPS + eig_vectors = eig_vectors[:, large_eig_values] + inv_eig_values = 1.0 / eig_values[large_eig_values] + ill_conditioned = ~np.all(large_eig_values, 0) + left_scaled_solutions = eig_vectors @ ( + (eig_vectors.T @ rhs_scaled) * inv_eig_values[:, np.newaxis] + ) + return ( + left_scaled_solutions * right_scaling[:, np.newaxis], + ill_conditioned, + ) + + @staticmethod + def _get_model(interpolation, values): + """ + Solve the interpolation system. + + Parameters + ---------- + interpolation : `cobyqa.models.Interpolation` + Interpolation set. + values : `numpy.ndarray`, shape (npt,) + Values of the interpolated function at the interpolation points. + + Returns + ------- + float + Constant term of the quadratic model. + `numpy.ndarray`, shape (n,) + Gradient of the quadratic model at ``interpolation.x_base``. + `numpy.ndarray`, shape (npt,) + Implicit Hessian matrix of the quadratic model. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + assert values.shape == ( + interpolation.npt, + ), "The shape of `values` is not valid." + n, npt = interpolation.xpt.shape + x, ill_conditioned = Quadratic.solve_systems( + interpolation, + np.block( + [ + [ + values, + np.zeros(n + 1), + ] + ] + ).T, + ) + return x[npt, 0], x[npt + 1:, 0], x[:npt, 0], ill_conditioned + + +class Models: + """ + Models for a nonlinear optimization problem. + """ + + def __init__(self, pb, options, penalty): + """ + Initialize the models. + + Parameters + ---------- + pb : `cobyqa.problem.Problem` + Problem to be solved. + options : dict + Options of the solver. + penalty : float + Penalty parameter used to select the point in the filter to forward + to the callback function. + + Raises + ------ + `cobyqa.utils.MaxEvalError` + If the maximum number of evaluations is reached. + `cobyqa.utils.TargetSuccess` + If a nearly feasible point has been found with an objective + function value below the target. + `cobyqa.utils.FeasibleSuccess` + If a feasible point has been found for a feasibility problem. + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + # Set the initial interpolation set. + self._debug = options[Options.DEBUG] + self._interpolation = Interpolation(pb, options) + + # Evaluate the nonlinear functions at the initial interpolation points. + x_eval = self.interpolation.point(0) + fun_init, cub_init, ceq_init = pb(x_eval, penalty) + self._fun_val = np.full(options[Options.NPT], np.nan) + self._cub_val = np.full((options[Options.NPT], cub_init.size), np.nan) + self._ceq_val = np.full((options[Options.NPT], ceq_init.size), np.nan) + for k in range(options[Options.NPT]): + if k >= options[Options.MAX_EVAL]: + raise MaxEvalError + if k == 0: + self.fun_val[k] = fun_init + self.cub_val[k, :] = cub_init + self.ceq_val[k, :] = ceq_init + else: + x_eval = self.interpolation.point(k) + self.fun_val[k], self.cub_val[k, :], self.ceq_val[k, :] = pb( + x_eval, + penalty, + ) + + # Stop the iterations if the problem is a feasibility problem and + # the current interpolation point is feasible. + if ( + pb.is_feasibility + and pb.maxcv( + self.interpolation.point(k), + self.cub_val[k, :], + self.ceq_val[k, :], + ) + <= options[Options.FEASIBILITY_TOL] + ): + raise FeasibleSuccess + + # Stop the iterations if the current interpolation point is nearly + # feasible and has an objective function value below the target. + if ( + self._fun_val[k] <= options[Options.TARGET] + and pb.maxcv( + self.interpolation.point(k), + self.cub_val[k, :], + self.ceq_val[k, :], + ) + <= options[Options.FEASIBILITY_TOL] + ): + raise TargetSuccess + + # Build the initial quadratic models. + self._fun = Quadratic( + self.interpolation, + self._fun_val, + options[Options.DEBUG], + ) + self._cub = np.empty(self.m_nonlinear_ub, dtype=Quadratic) + self._ceq = np.empty(self.m_nonlinear_eq, dtype=Quadratic) + for i in range(self.m_nonlinear_ub): + self._cub[i] = Quadratic( + self.interpolation, + self.cub_val[:, i], + options[Options.DEBUG], + ) + for i in range(self.m_nonlinear_eq): + self._ceq[i] = Quadratic( + self.interpolation, + self.ceq_val[:, i], + options[Options.DEBUG], + ) + if self._debug: + self._check_interpolation_conditions() + + @property + def n(self): + """ + Dimension of the problem. + + Returns + ------- + int + Dimension of the problem. + """ + return self.interpolation.n + + @property + def npt(self): + """ + Number of interpolation points. + + Returns + ------- + int + Number of interpolation points. + """ + return self.interpolation.npt + + @property + def m_nonlinear_ub(self): + """ + Number of nonlinear inequality constraints. + + Returns + ------- + int + Number of nonlinear inequality constraints. + """ + return self.cub_val.shape[1] + + @property + def m_nonlinear_eq(self): + """ + Number of nonlinear equality constraints. + + Returns + ------- + int + Number of nonlinear equality constraints. + """ + return self.ceq_val.shape[1] + + @property + def interpolation(self): + """ + Interpolation set. + + Returns + ------- + `cobyqa.models.Interpolation` + Interpolation set. + """ + return self._interpolation + + @property + def fun_val(self): + """ + Values of the objective function at the interpolation points. + + Returns + ------- + `numpy.ndarray`, shape (npt,) + Values of the objective function at the interpolation points. + """ + return self._fun_val + + @property + def cub_val(self): + """ + Values of the nonlinear inequality constraint functions at the + interpolation points. + + Returns + ------- + `numpy.ndarray`, shape (npt, m_nonlinear_ub) + Values of the nonlinear inequality constraint functions at the + interpolation points. + """ + return self._cub_val + + @property + def ceq_val(self): + """ + Values of the nonlinear equality constraint functions at the + interpolation points. + + Returns + ------- + `numpy.ndarray`, shape (npt, m_nonlinear_eq) + Values of the nonlinear equality constraint functions at the + interpolation points. + """ + return self._ceq_val + + def fun(self, x): + """ + Evaluate the quadratic model of the objective function at a given + point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the quadratic model of the objective + function. + + Returns + ------- + float + Value of the quadratic model of the objective function at `x`. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + return self._fun(x, self.interpolation) + + def fun_grad(self, x): + """ + Evaluate the gradient of the quadratic model of the objective function + at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the gradient of the quadratic model of + the objective function. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Gradient of the quadratic model of the objective function at `x`. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + return self._fun.grad(x, self.interpolation) + + def fun_hess(self): + """ + Evaluate the Hessian matrix of the quadratic model of the objective + function. + + Returns + ------- + `numpy.ndarray`, shape (n, n) + Hessian matrix of the quadratic model of the objective function. + """ + return self._fun.hess(self.interpolation) + + def fun_hess_prod(self, v): + """ + Evaluate the right product of the Hessian matrix of the quadratic model + of the objective function with a given vector. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Vector with which the Hessian matrix of the quadratic model of the + objective function is multiplied from the right. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Right product of the Hessian matrix of the quadratic model of the + objective function with `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + return self._fun.hess_prod(v, self.interpolation) + + def fun_curv(self, v): + """ + Evaluate the curvature of the quadratic model of the objective function + along a given direction. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Direction along which the curvature of the quadratic model of the + objective function is evaluated. + + Returns + ------- + float + Curvature of the quadratic model of the objective function along + `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + return self._fun.curv(v, self.interpolation) + + def fun_alt_grad(self, x): + """ + Evaluate the gradient of the alternative quadratic model of the + objective function at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the gradient of the alternative + quadratic model of the objective function. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Gradient of the alternative quadratic model of the objective + function at `x`. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + model = Quadratic(self.interpolation, self.fun_val, self._debug) + return model.grad(x, self.interpolation) + + def cub(self, x, mask=None): + """ + Evaluate the quadratic models of the nonlinear inequality functions at + a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the quadratic models of the nonlinear + inequality functions. + mask : `numpy.ndarray`, shape (m_nonlinear_ub,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Values of the quadratic model of the nonlinear inequality + functions. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_ub, + ), "The shape of `mask` is not valid." + return np.array( + [model(x, self.interpolation) for model in self._get_cub(mask)] + ) + + def cub_grad(self, x, mask=None): + """ + Evaluate the gradients of the quadratic models of the nonlinear + inequality functions at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the gradients of the quadratic models of + the nonlinear inequality functions. + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Gradients of the quadratic model of the nonlinear inequality + functions. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_ub, + ), "The shape of `mask` is not valid." + return np.reshape( + [model.grad(x, self.interpolation) + for model in self._get_cub(mask)], + (-1, self.n), + ) + + def cub_hess(self, mask=None): + """ + Evaluate the Hessian matrices of the quadratic models of the nonlinear + inequality functions. + + Parameters + ---------- + mask : `numpy.ndarray`, shape (m_nonlinear_ub,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Hessian matrices of the quadratic models of the nonlinear + inequality functions. + """ + if self._debug: + assert mask is None or mask.shape == ( + self.m_nonlinear_ub, + ), "The shape of `mask` is not valid." + return np.reshape( + [model.hess(self.interpolation) for model in self._get_cub(mask)], + (-1, self.n, self.n), + ) + + def cub_hess_prod(self, v, mask=None): + """ + Evaluate the right product of the Hessian matrices of the quadratic + models of the nonlinear inequality functions with a given vector. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Vector with which the Hessian matrices of the quadratic models of + the nonlinear inequality functions are multiplied from the right. + mask : `numpy.ndarray`, shape (m_nonlinear_ub,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Right products of the Hessian matrices of the quadratic models of + the nonlinear inequality functions with `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_ub, + ), "The shape of `mask` is not valid." + return np.reshape( + [ + model.hess_prod(v, self.interpolation) + for model in self._get_cub(mask) + ], + (-1, self.n), + ) + + def cub_curv(self, v, mask=None): + """ + Evaluate the curvature of the quadratic models of the nonlinear + inequality functions along a given direction. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Direction along which the curvature of the quadratic models of the + nonlinear inequality functions is evaluated. + mask : `numpy.ndarray`, shape (m_nonlinear_ub,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Curvature of the quadratic models of the nonlinear inequality + functions along `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_ub, + ), "The shape of `mask` is not valid." + return np.array( + [model.curv(v, self.interpolation) + for model in self._get_cub(mask)] + ) + + def ceq(self, x, mask=None): + """ + Evaluate the quadratic models of the nonlinear equality functions at a + given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the quadratic models of the nonlinear + equality functions. + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Values of the quadratic model of the nonlinear equality functions. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_eq, + ), "The shape of `mask` is not valid." + return np.array( + [model(x, self.interpolation) for model in self._get_ceq(mask)] + ) + + def ceq_grad(self, x, mask=None): + """ + Evaluate the gradients of the quadratic models of the nonlinear + equality functions at a given point. + + Parameters + ---------- + x : `numpy.ndarray`, shape (n,) + Point at which to evaluate the gradients of the quadratic models of + the nonlinear equality functions. + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Gradients of the quadratic model of the nonlinear equality + functions. + """ + if self._debug: + assert x.shape == (self.n,), "The shape of `x` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_eq, + ), "The shape of `mask` is not valid." + return np.reshape( + [model.grad(x, self.interpolation) + for model in self._get_ceq(mask)], + (-1, self.n), + ) + + def ceq_hess(self, mask=None): + """ + Evaluate the Hessian matrices of the quadratic models of the nonlinear + equality functions. + + Parameters + ---------- + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Hessian matrices of the quadratic models of the nonlinear equality + functions. + """ + if self._debug: + assert mask is None or mask.shape == ( + self.m_nonlinear_eq, + ), "The shape of `mask` is not valid." + return np.reshape( + [model.hess(self.interpolation) for model in self._get_ceq(mask)], + (-1, self.n, self.n), + ) + + def ceq_hess_prod(self, v, mask=None): + """ + Evaluate the right product of the Hessian matrices of the quadratic + models of the nonlinear equality functions with a given vector. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Vector with which the Hessian matrices of the quadratic models of + the nonlinear equality functions are multiplied from the right. + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Right products of the Hessian matrices of the quadratic models of + the nonlinear equality functions with `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_eq, + ), "The shape of `mask` is not valid." + return np.reshape( + [ + model.hess_prod(v, self.interpolation) + for model in self._get_ceq(mask) + ], + (-1, self.n), + ) + + def ceq_curv(self, v, mask=None): + """ + Evaluate the curvature of the quadratic models of the nonlinear + equality functions along a given direction. + + Parameters + ---------- + v : `numpy.ndarray`, shape (n,) + Direction along which the curvature of the quadratic models of the + nonlinear equality functions is evaluated. + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to consider. + + Returns + ------- + `numpy.ndarray` + Curvature of the quadratic models of the nonlinear equality + functions along `v`. + """ + if self._debug: + assert v.shape == (self.n,), "The shape of `v` is not valid." + assert mask is None or mask.shape == ( + self.m_nonlinear_eq, + ), "The shape of `mask` is not valid." + return np.array( + [model.curv(v, self.interpolation) + for model in self._get_ceq(mask)] + ) + + def reset_models(self): + """ + Set the quadratic models of the objective function, nonlinear + inequality constraints, and nonlinear equality constraints to the + alternative quadratic models. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + self._fun = Quadratic(self.interpolation, self.fun_val, self._debug) + for i in range(self.m_nonlinear_ub): + self._cub[i] = Quadratic( + self.interpolation, + self.cub_val[:, i], + self._debug, + ) + for i in range(self.m_nonlinear_eq): + self._ceq[i] = Quadratic( + self.interpolation, + self.ceq_val[:, i], + self._debug, + ) + if self._debug: + self._check_interpolation_conditions() + + def update_interpolation(self, k_new, x_new, fun_val, cub_val, ceq_val): + """ + Update the interpolation set. + + This method updates the interpolation set by replacing the `knew`-th + interpolation point with `xnew`. It also updates the function values + and the quadratic models. + + Parameters + ---------- + k_new : int + Index of the updated interpolation point. + x_new : `numpy.ndarray`, shape (n,) + New interpolation point. Its value is interpreted as relative to + the origin, not the base point. + fun_val : float + Value of the objective function at `x_new`. + Objective function value at `x_new`. + cub_val : `numpy.ndarray`, shape (m_nonlinear_ub,) + Values of the nonlinear inequality constraints at `x_new`. + ceq_val : `numpy.ndarray`, shape (m_nonlinear_eq,) + Values of the nonlinear equality constraints at `x_new`. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + """ + if self._debug: + assert 0 <= k_new < self.npt, "The index `k_new` is not valid." + assert x_new.shape == (self.n,), \ + "The shape of `x_new` is not valid." + assert isinstance(fun_val, float), \ + "The function value is not valid." + assert cub_val.shape == ( + self.m_nonlinear_ub, + ), "The shape of `cub_val` is not valid." + assert ceq_val.shape == ( + self.m_nonlinear_eq, + ), "The shape of `ceq_val` is not valid." + + # Compute the updates in the interpolation conditions. + fun_diff = np.zeros(self.npt) + cub_diff = np.zeros(self.cub_val.shape) + ceq_diff = np.zeros(self.ceq_val.shape) + fun_diff[k_new] = fun_val - self.fun(x_new) + cub_diff[k_new, :] = cub_val - self.cub(x_new) + ceq_diff[k_new, :] = ceq_val - self.ceq(x_new) + + # Update the function values. + self.fun_val[k_new] = fun_val + self.cub_val[k_new, :] = cub_val + self.ceq_val[k_new, :] = ceq_val + + # Update the interpolation set. + dir_old = np.copy(self.interpolation.xpt[:, k_new]) + self.interpolation.xpt[:, k_new] = x_new - self.interpolation.x_base + + # Update the quadratic models. + ill_conditioned = self._fun.update( + self.interpolation, + k_new, + dir_old, + fun_diff, + ) + for i in range(self.m_nonlinear_ub): + ill_conditioned = ill_conditioned or self._cub[i].update( + self.interpolation, + k_new, + dir_old, + cub_diff[:, i], + ) + for i in range(self.m_nonlinear_eq): + ill_conditioned = ill_conditioned or self._ceq[i].update( + self.interpolation, + k_new, + dir_old, + ceq_diff[:, i], + ) + if self._debug: + self._check_interpolation_conditions() + return ill_conditioned + + def determinants(self, x_new, k_new=None): + """ + Compute the normalized determinants of the new interpolation systems. + + Parameters + ---------- + x_new : `numpy.ndarray`, shape (n,) + New interpolation point. Its value is interpreted as relative to + the origin, not the base point. + k_new : int, optional + Index of the updated interpolation point. If `k_new` is not + specified, all the possible determinants are computed. + + Returns + ------- + {float, `numpy.ndarray`, shape (npt,)} + Determinant(s) of the new interpolation system. + + Raises + ------ + `numpy.linalg.LinAlgError` + If the interpolation system is ill-defined. + + Notes + ----- + The determinants are normalized by the determinant of the current + interpolation system. For stability reasons, the calculations are done + using the formula (2.12) in [1]_. + + References + ---------- + .. [1] M. J. D. Powell. On updating the inverse of a KKT matrix. + Technical Report DAMTP 2004/NA01, Department of Applied Mathematics + and Theoretical Physics, University of Cambridge, Cambridge, UK, + 2004. + """ + if self._debug: + assert x_new.shape == (self.n,), \ + "The shape of `x_new` is not valid." + assert ( + k_new is None or 0 <= k_new < self.npt + ), "The index `k_new` is not valid." + + # Compute the values independent of k_new. + shift = x_new - self.interpolation.x_base + new_col = np.empty((self.npt + self.n + 1, 1)) + new_col[: self.npt, 0] = ( + 0.5 * (self.interpolation.xpt.T @ shift) ** 2.0) + new_col[self.npt, 0] = 1.0 + new_col[self.npt + 1:, 0] = shift + inv_new_col = Quadratic.solve_systems(self.interpolation, new_col)[0] + beta = 0.5 * (shift @ shift) ** 2.0 - new_col[:, 0] @ inv_new_col[:, 0] + + # Compute the values that depend on k. + if k_new is None: + coord_vec = np.eye(self.npt + self.n + 1, self.npt) + alpha = np.diag( + Quadratic.solve_systems( + self.interpolation, + coord_vec, + )[0] + ) + tau = inv_new_col[: self.npt, 0] + else: + coord_vec = np.eye(self.npt + self.n + 1, 1, -k_new) + alpha = Quadratic.solve_systems( + self.interpolation, + coord_vec, + )[ + 0 + ][k_new, 0] + tau = inv_new_col[k_new, 0] + return alpha * beta + tau**2.0 + + def shift_x_base(self, new_x_base, options): + """ + Shift the base point without changing the interpolation set. + + Parameters + ---------- + new_x_base : `numpy.ndarray`, shape (n,) + New base point. + options : dict + Options of the solver. + """ + if self._debug: + assert new_x_base.shape == ( + self.n, + ), "The shape of `new_x_base` is not valid." + + # Update the models. + self._fun.shift_x_base(self.interpolation, new_x_base) + for model in self._cub: + model.shift_x_base(self.interpolation, new_x_base) + for model in self._ceq: + model.shift_x_base(self.interpolation, new_x_base) + + # Update the base point and the interpolation points. + shift = new_x_base - self.interpolation.x_base + self.interpolation.x_base += shift + self.interpolation.xpt -= shift[:, np.newaxis] + if options[Options.DEBUG]: + self._check_interpolation_conditions() + + def _get_cub(self, mask=None): + """ + Get the quadratic models of the nonlinear inequality constraints. + + Parameters + ---------- + mask : `numpy.ndarray`, shape (m_nonlinear_ub,), optional + Mask of the quadratic models to return. + + Returns + ------- + `numpy.ndarray` + Quadratic models of the nonlinear inequality constraints. + """ + return self._cub if mask is None else self._cub[mask] + + def _get_ceq(self, mask=None): + """ + Get the quadratic models of the nonlinear equality constraints. + + Parameters + ---------- + mask : `numpy.ndarray`, shape (m_nonlinear_eq,), optional + Mask of the quadratic models to return. + + Returns + ------- + `numpy.ndarray` + Quadratic models of the nonlinear equality constraints. + """ + return self._ceq if mask is None else self._ceq[mask] + + def _check_interpolation_conditions(self): + """ + Check the interpolation conditions of all quadratic models. + """ + error_fun = 0.0 + error_cub = 0.0 + error_ceq = 0.0 + for k in range(self.npt): + error_fun = np.max( + [ + error_fun, + np.abs( + self.fun(self.interpolation.point(k)) - self.fun_val[k] + ), + ] + ) + error_cub = np.max( + np.abs( + self.cub(self.interpolation.point(k)) - self.cub_val[k, :] + ), + initial=error_cub, + ) + error_ceq = np.max( + np.abs( + self.ceq(self.interpolation.point(k)) - self.ceq_val[k, :] + ), + initial=error_ceq, + ) + tol = 10.0 * np.sqrt(EPS) * max(self.n, self.npt) + if error_fun > tol * np.max(np.abs(self.fun_val), initial=1.0): + warnings.warn( + "The interpolation conditions for the objective function are " + "not satisfied.", + RuntimeWarning, + 2, + ) + if error_cub > tol * np.max(np.abs(self.cub_val), initial=1.0): + warnings.warn( + "The interpolation conditions for the inequality constraint " + "function are not satisfied.", + RuntimeWarning, + 2, + ) + if error_ceq > tol * np.max(np.abs(self.ceq_val), initial=1.0): + warnings.warn( + "The interpolation conditions for the equality constraint " + "function are not satisfied.", + RuntimeWarning, + 2, + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/problem.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/problem.py new file mode 100644 index 0000000000000000000000000000000000000000..046441a76486eab3d665af26147f6843d8bfaef2 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/problem.py @@ -0,0 +1,1296 @@ +from contextlib import suppress +from inspect import signature +import copy + +import numpy as np +from scipy.optimize import ( + Bounds, + LinearConstraint, + NonlinearConstraint, + OptimizeResult, +) +from scipy.optimize._constraints import PreparedConstraint + + +from .settings import PRINT_OPTIONS, BARRIER +from .utils import CallbackSuccess, get_arrays_tol +from .utils import exact_1d_array + + +class ObjectiveFunction: + """ + Real-valued objective function. + """ + + def __init__(self, fun, verbose, debug, *args): + """ + Initialize the objective function. + + Parameters + ---------- + fun : {callable, None} + Function to evaluate, or None. + + ``fun(x, *args) -> float`` + + where ``x`` is an array with shape (n,) and `args` is a tuple. + verbose : bool + Whether to print the function evaluations. + debug : bool + Whether to make debugging tests during the execution. + *args : tuple + Additional arguments to be passed to the function. + """ + if debug: + assert fun is None or callable(fun) + assert isinstance(verbose, bool) + assert isinstance(debug, bool) + + self._fun = fun + self._verbose = verbose + self._args = args + self._n_eval = 0 + + def __call__(self, x): + """ + Evaluate the objective function. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the objective function is evaluated. + + Returns + ------- + float + Function value at `x`. + """ + x = np.array(x, dtype=float) + if self._fun is None: + f = 0.0 + else: + f = float(np.squeeze(self._fun(x, *self._args))) + self._n_eval += 1 + if self._verbose: + with np.printoptions(**PRINT_OPTIONS): + print(f"{self.name}({x}) = {f}") + return f + + @property + def n_eval(self): + """ + Number of function evaluations. + + Returns + ------- + int + Number of function evaluations. + """ + return self._n_eval + + @property + def name(self): + """ + Name of the objective function. + + Returns + ------- + str + Name of the objective function. + """ + name = "" + if self._fun is not None: + try: + name = self._fun.__name__ + except AttributeError: + name = "fun" + return name + + +class BoundConstraints: + """ + Bound constraints ``xl <= x <= xu``. + """ + + def __init__(self, bounds): + """ + Initialize the bound constraints. + + Parameters + ---------- + bounds : scipy.optimize.Bounds + Bound constraints. + """ + self._xl = np.array(bounds.lb, float) + self._xu = np.array(bounds.ub, float) + + # Remove the ill-defined bounds. + self.xl[np.isnan(self.xl)] = -np.inf + self.xu[np.isnan(self.xu)] = np.inf + + self.is_feasible = ( + np.all(self.xl <= self.xu) + and np.all(self.xl < np.inf) + and np.all(self.xu > -np.inf) + ) + self.m = np.count_nonzero(self.xl > -np.inf) + np.count_nonzero( + self.xu < np.inf + ) + self.pcs = PreparedConstraint(bounds, np.ones(bounds.lb.size)) + + @property + def xl(self): + """ + Lower bound. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Lower bound. + """ + return self._xl + + @property + def xu(self): + """ + Upper bound. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Upper bound. + """ + return self._xu + + def maxcv(self, x): + """ + Evaluate the maximum constraint violation. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the maximum constraint violation is evaluated. + + Returns + ------- + float + Maximum constraint violation at `x`. + """ + x = np.asarray(x, dtype=float) + return self.violation(x) + + def violation(self, x): + # shortcut for no bounds + if self.is_feasible: + return np.array([0]) + else: + return self.pcs.violation(x) + + def project(self, x): + """ + Project a point onto the feasible set. + + Parameters + ---------- + x : array_like, shape (n,) + Point to be projected. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Projection of `x` onto the feasible set. + """ + return np.clip(x, self.xl, self.xu) if self.is_feasible else x + + +class LinearConstraints: + """ + Linear constraints ``a_ub @ x <= b_ub`` and ``a_eq @ x == b_eq``. + """ + + def __init__(self, constraints, n, debug): + """ + Initialize the linear constraints. + + Parameters + ---------- + constraints : list of LinearConstraint + Linear constraints. + n : int + Number of variables. + debug : bool + Whether to make debugging tests during the execution. + """ + if debug: + assert isinstance(constraints, list) + for constraint in constraints: + assert isinstance(constraint, LinearConstraint) + assert isinstance(debug, bool) + + self._a_ub = np.empty((0, n)) + self._b_ub = np.empty(0) + self._a_eq = np.empty((0, n)) + self._b_eq = np.empty(0) + for constraint in constraints: + is_equality = np.abs( + constraint.ub - constraint.lb + ) <= get_arrays_tol(constraint.lb, constraint.ub) + if np.any(is_equality): + self._a_eq = np.vstack((self.a_eq, constraint.A[is_equality])) + self._b_eq = np.concatenate( + ( + self.b_eq, + 0.5 + * ( + constraint.lb[is_equality] + + constraint.ub[is_equality] + ), + ) + ) + if not np.all(is_equality): + self._a_ub = np.vstack( + ( + self.a_ub, + constraint.A[~is_equality], + -constraint.A[~is_equality], + ) + ) + self._b_ub = np.concatenate( + ( + self.b_ub, + constraint.ub[~is_equality], + -constraint.lb[~is_equality], + ) + ) + + # Remove the ill-defined constraints. + self.a_ub[np.isnan(self.a_ub)] = 0.0 + self.a_eq[np.isnan(self.a_eq)] = 0.0 + undef_ub = np.isnan(self.b_ub) | np.isinf(self.b_ub) + undef_eq = np.isnan(self.b_eq) + self._a_ub = self.a_ub[~undef_ub, :] + self._b_ub = self.b_ub[~undef_ub] + self._a_eq = self.a_eq[~undef_eq, :] + self._b_eq = self.b_eq[~undef_eq] + self.pcs = [ + PreparedConstraint(c, np.ones(n)) for c in constraints if c.A.size + ] + + @property + def a_ub(self): + """ + Left-hand side matrix of the linear inequality constraints. + + Returns + ------- + `numpy.ndarray`, shape (m, n) + Left-hand side matrix of the linear inequality constraints. + """ + return self._a_ub + + @property + def b_ub(self): + """ + Right-hand side vector of the linear inequality constraints. + + Returns + ------- + `numpy.ndarray`, shape (m, n) + Right-hand side vector of the linear inequality constraints. + """ + return self._b_ub + + @property + def a_eq(self): + """ + Left-hand side matrix of the linear equality constraints. + + Returns + ------- + `numpy.ndarray`, shape (m, n) + Left-hand side matrix of the linear equality constraints. + """ + return self._a_eq + + @property + def b_eq(self): + """ + Right-hand side vector of the linear equality constraints. + + Returns + ------- + `numpy.ndarray`, shape (m, n) + Right-hand side vector of the linear equality constraints. + """ + return self._b_eq + + @property + def m_ub(self): + """ + Number of linear inequality constraints. + + Returns + ------- + int + Number of linear inequality constraints. + """ + return self.b_ub.size + + @property + def m_eq(self): + """ + Number of linear equality constraints. + + Returns + ------- + int + Number of linear equality constraints. + """ + return self.b_eq.size + + def maxcv(self, x): + """ + Evaluate the maximum constraint violation. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the maximum constraint violation is evaluated. + + Returns + ------- + float + Maximum constraint violation at `x`. + """ + return np.max(self.violation(x), initial=0.0) + + def violation(self, x): + if len(self.pcs): + return np.concatenate([pc.violation(x) for pc in self.pcs]) + return np.array([]) + + +class NonlinearConstraints: + """ + Nonlinear constraints ``c_ub(x) <= 0`` and ``c_eq(x) == b_eq``. + """ + + def __init__(self, constraints, verbose, debug): + """ + Initialize the nonlinear constraints. + + Parameters + ---------- + constraints : list + Nonlinear constraints. + verbose : bool + Whether to print the function evaluations. + debug : bool + Whether to make debugging tests during the execution. + """ + if debug: + assert isinstance(constraints, list) + for constraint in constraints: + assert isinstance(constraint, NonlinearConstraint) + assert isinstance(verbose, bool) + assert isinstance(debug, bool) + + self._constraints = constraints + self.pcs = [] + self._verbose = verbose + + # map of indexes for equality and inequality constraints + self._map_ub = None + self._map_eq = None + self._m_ub = self._m_eq = None + + def __call__(self, x): + """ + Calculates the residual (slack) for the constraints. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the constraints are evaluated. + + Returns + ------- + `numpy.ndarray`, shape (m_nonlinear_ub,) + Nonlinear inequality constraint slack values. + `numpy.ndarray`, shape (m_nonlinear_eq,) + Nonlinear equality constraint slack values. + """ + if not len(self._constraints): + self._m_eq = self._m_ub = 0 + return np.array([]), np.array([]) + + x = np.array(x, dtype=float) + # first time around the constraints haven't been prepared + if not len(self.pcs): + self._map_ub = [] + self._map_eq = [] + self._m_eq = 0 + self._m_ub = 0 + + for constraint in self._constraints: + if not callable(constraint.jac): + # having a callable constraint function prevents + # constraint.fun from being evaluated when preparing + # constraint + c = copy.copy(constraint) + c.jac = lambda x0: x0 + c.hess = lambda x0, v: 0.0 + pc = PreparedConstraint(c, x) + else: + pc = PreparedConstraint(constraint, x) + # we're going to be using the same x value again immediately + # after this initialisation + pc.fun.f_updated = True + + self.pcs.append(pc) + idx = np.arange(pc.fun.m) + + # figure out equality and inequality maps + lb, ub = pc.bounds[0], pc.bounds[1] + arr_tol = get_arrays_tol(lb, ub) + is_equality = np.abs(ub - lb) <= arr_tol + self._map_eq.append(idx[is_equality]) + self._map_ub.append(idx[~is_equality]) + + # these values will be corrected to their proper values later + self._m_eq += np.count_nonzero(is_equality) + self._m_ub += np.count_nonzero(~is_equality) + + c_ub = [] + c_eq = [] + for i, pc in enumerate(self.pcs): + val = pc.fun.fun(x) + if self._verbose: + with np.printoptions(**PRINT_OPTIONS): + with suppress(AttributeError): + fun_name = self._constraints[i].fun.__name__ + print(f"{fun_name}({x}) = {val}") + + # separate violations into c_eq and c_ub + eq_idx = self._map_eq[i] + ub_idx = self._map_ub[i] + + ub_val = val[ub_idx] + if len(ub_idx): + xl = pc.bounds[0][ub_idx] + xu = pc.bounds[1][ub_idx] + + # calculate slack within lower bound + finite_xl = xl > -np.inf + _v = xl[finite_xl] - ub_val[finite_xl] + c_ub.append(_v) + + # calculate slack within lower bound + finite_xu = xu < np.inf + _v = ub_val[finite_xu] - xu[finite_xu] + c_ub.append(_v) + + # equality constraints taken from midpoint between lb and ub + eq_val = val[eq_idx] + if len(eq_idx): + midpoint = 0.5 * (pc.bounds[1][eq_idx] + pc.bounds[0][eq_idx]) + eq_val -= midpoint + c_eq.append(eq_val) + + if self._m_eq: + c_eq = np.concatenate(c_eq) + else: + c_eq = np.array([]) + + if self._m_ub: + c_ub = np.concatenate(c_ub) + else: + c_ub = np.array([]) + + self._m_ub = c_ub.size + self._m_eq = c_eq.size + + return c_ub, c_eq + + @property + def m_ub(self): + """ + Number of nonlinear inequality constraints. + + Returns + ------- + int + Number of nonlinear inequality constraints. + + Raises + ------ + ValueError + If the number of nonlinear inequality constraints is unknown. + """ + if self._m_ub is None: + raise ValueError( + "The number of nonlinear inequality constraints is unknown." + ) + else: + return self._m_ub + + @property + def m_eq(self): + """ + Number of nonlinear equality constraints. + + Returns + ------- + int + Number of nonlinear equality constraints. + + Raises + ------ + ValueError + If the number of nonlinear equality constraints is unknown. + """ + if self._m_eq is None: + raise ValueError( + "The number of nonlinear equality constraints is unknown." + ) + else: + return self._m_eq + + @property + def n_eval(self): + """ + Number of function evaluations. + + Returns + ------- + int + Number of function evaluations. + """ + if len(self.pcs): + return self.pcs[0].fun.nfev + else: + return 0 + + def maxcv(self, x, cub_val=None, ceq_val=None): + """ + Evaluate the maximum constraint violation. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the maximum constraint violation is evaluated. + cub_val : array_like, shape (m_nonlinear_ub,), optional + Values of the nonlinear inequality constraints. If not provided, + the nonlinear inequality constraints are evaluated at `x`. + ceq_val : array_like, shape (m_nonlinear_eq,), optional + Values of the nonlinear equality constraints. If not provided, + the nonlinear equality constraints are evaluated at `x`. + + Returns + ------- + float + Maximum constraint violation at `x`. + """ + return np.max( + self.violation(x, cub_val=cub_val, ceq_val=ceq_val), initial=0.0 + ) + + def violation(self, x, cub_val=None, ceq_val=None): + return np.concatenate([pc.violation(x) for pc in self.pcs]) + + +class Problem: + """ + Optimization problem. + """ + + def __init__( + self, + obj, + x0, + bounds, + linear, + nonlinear, + callback, + feasibility_tol, + scale, + store_history, + history_size, + filter_size, + debug, + ): + """ + Initialize the nonlinear problem. + + The problem is preprocessed to remove all the variables that are fixed + by the bound constraints. + + Parameters + ---------- + obj : ObjectiveFunction + Objective function. + x0 : array_like, shape (n,) + Initial guess. + bounds : BoundConstraints + Bound constraints. + linear : LinearConstraints + Linear constraints. + nonlinear : NonlinearConstraints + Nonlinear constraints. + callback : {callable, None} + Callback function. + feasibility_tol : float + Tolerance on the constraint violation. + scale : bool + Whether to scale the problem according to the bounds. + store_history : bool + Whether to store the function evaluations. + history_size : int + Maximum number of function evaluations to store. + filter_size : int + Maximum number of points in the filter. + debug : bool + Whether to make debugging tests during the execution. + """ + if debug: + assert isinstance(obj, ObjectiveFunction) + assert isinstance(bounds, BoundConstraints) + assert isinstance(linear, LinearConstraints) + assert isinstance(nonlinear, NonlinearConstraints) + assert isinstance(feasibility_tol, float) + assert isinstance(scale, bool) + assert isinstance(store_history, bool) + assert isinstance(history_size, int) + if store_history: + assert history_size > 0 + assert isinstance(filter_size, int) + assert filter_size > 0 + assert isinstance(debug, bool) + + self._obj = obj + self._linear = linear + self._nonlinear = nonlinear + if callback is not None: + if not callable(callback): + raise TypeError("The callback must be a callable function.") + self._callback = callback + + # Check the consistency of the problem. + x0 = exact_1d_array(x0, "The initial guess must be a vector.") + n = x0.size + if bounds.xl.size != n: + raise ValueError(f"The bounds must have {n} elements.") + if linear.a_ub.shape[1] != n: + raise ValueError( + f"The left-hand side matrices of the linear constraints must " + f"have {n} columns." + ) + + # Check which variables are fixed. + tol = get_arrays_tol(bounds.xl, bounds.xu) + self._fixed_idx = (bounds.xl <= bounds.xu) & ( + np.abs(bounds.xl - bounds.xu) < tol + ) + self._fixed_val = 0.5 * ( + bounds.xl[self._fixed_idx] + bounds.xu[self._fixed_idx] + ) + self._fixed_val = np.clip( + self._fixed_val, + bounds.xl[self._fixed_idx], + bounds.xu[self._fixed_idx], + ) + + # Set the bound constraints. + self._orig_bounds = bounds + self._bounds = BoundConstraints( + Bounds(bounds.xl[~self._fixed_idx], bounds.xu[~self._fixed_idx]) + ) + + # Set the initial guess. + self._x0 = self._bounds.project(x0[~self._fixed_idx]) + + # Set the linear constraints. + b_eq = linear.b_eq - linear.a_eq[:, self._fixed_idx] @ self._fixed_val + self._linear = LinearConstraints( + [ + LinearConstraint( + linear.a_ub[:, ~self._fixed_idx], + -np.inf, + linear.b_ub + - linear.a_ub[:, self._fixed_idx] @ self._fixed_val, + ), + LinearConstraint(linear.a_eq[:, ~self._fixed_idx], b_eq, b_eq), + ], + self.n, + debug, + ) + + # Scale the problem if necessary. + scale = ( + scale + and self._bounds.is_feasible + and np.all(np.isfinite(self._bounds.xl)) + and np.all(np.isfinite(self._bounds.xu)) + ) + if scale: + self._scaling_factor = 0.5 * (self._bounds.xu - self._bounds.xl) + self._scaling_shift = 0.5 * (self._bounds.xu + self._bounds.xl) + self._bounds = BoundConstraints( + Bounds(-np.ones(self.n), np.ones(self.n)) + ) + b_eq = self._linear.b_eq - self._linear.a_eq @ self._scaling_shift + self._linear = LinearConstraints( + [ + LinearConstraint( + self._linear.a_ub @ np.diag(self._scaling_factor), + -np.inf, + self._linear.b_ub + - self._linear.a_ub @ self._scaling_shift, + ), + LinearConstraint( + self._linear.a_eq @ np.diag(self._scaling_factor), + b_eq, + b_eq, + ), + ], + self.n, + debug, + ) + self._x0 = (self._x0 - self._scaling_shift) / self._scaling_factor + else: + self._scaling_factor = np.ones(self.n) + self._scaling_shift = np.zeros(self.n) + + # Set the initial filter. + self._feasibility_tol = feasibility_tol + self._filter_size = filter_size + self._fun_filter = [] + self._maxcv_filter = [] + self._x_filter = [] + + # Set the initial history. + self._store_history = store_history + self._history_size = history_size + self._fun_history = [] + self._maxcv_history = [] + self._x_history = [] + + def __call__(self, x, penalty=0.0): + """ + Evaluate the objective and nonlinear constraint functions. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the functions are evaluated. + penalty : float, optional + Penalty parameter used to select the point in the filter to forward + to the callback function. + + Returns + ------- + float + Objective function value. + `numpy.ndarray`, shape (m_nonlinear_ub,) + Nonlinear inequality constraint function values. + `numpy.ndarray`, shape (m_nonlinear_eq,) + Nonlinear equality constraint function values. + + Raises + ------ + `cobyqa.utils.CallbackSuccess` + If the callback function raises a ``StopIteration``. + """ + # Evaluate the objective and nonlinear constraint functions. + x = np.asarray(x, dtype=float) + x_full = self.build_x(x) + fun_val = self._obj(x_full) + cub_val, ceq_val = self._nonlinear(x_full) + maxcv_val = self.maxcv(x, cub_val, ceq_val) + if self._store_history: + self._fun_history.append(fun_val) + self._maxcv_history.append(maxcv_val) + self._x_history.append(x) + if len(self._fun_history) > self._history_size: + self._fun_history.pop(0) + self._maxcv_history.pop(0) + self._x_history.pop(0) + + # Add the point to the filter if it is not dominated by any point. + if np.isnan(fun_val) and np.isnan(maxcv_val): + include_point = len(self._fun_filter) == 0 + elif np.isnan(fun_val): + include_point = all( + np.isnan(fun_filter) + and maxcv_val < maxcv_filter + or np.isnan(maxcv_filter) + for fun_filter, maxcv_filter in zip( + self._fun_filter, + self._maxcv_filter, + ) + ) + elif np.isnan(maxcv_val): + include_point = all( + np.isnan(maxcv_filter) + and fun_val < fun_filter + or np.isnan(fun_filter) + for fun_filter, maxcv_filter in zip( + self._fun_filter, + self._maxcv_filter, + ) + ) + else: + include_point = all( + fun_val < fun_filter or maxcv_val < maxcv_filter + for fun_filter, maxcv_filter in zip( + self._fun_filter, + self._maxcv_filter, + ) + ) + if include_point: + self._fun_filter.append(fun_val) + self._maxcv_filter.append(maxcv_val) + self._x_filter.append(x) + + # Remove the points in the filter that are dominated by the new + # point. We must iterate in reverse order to avoid problems when + # removing elements from the list. + for k in range(len(self._fun_filter) - 2, -1, -1): + if np.isnan(fun_val): + remove_point = np.isnan(self._fun_filter[k]) + elif np.isnan(maxcv_val): + remove_point = np.isnan(self._maxcv_filter[k]) + else: + remove_point = ( + np.isnan(self._fun_filter[k]) + or np.isnan(self._maxcv_filter[k]) + or fun_val <= self._fun_filter[k] + and maxcv_val <= self._maxcv_filter[k] + ) + if remove_point: + self._fun_filter.pop(k) + self._maxcv_filter.pop(k) + self._x_filter.pop(k) + + # Keep only the most recent points in the filter. + if len(self._fun_filter) > self._filter_size: + self._fun_filter.pop(0) + self._maxcv_filter.pop(0) + self._x_filter.pop(0) + + # Evaluate the callback function after updating the filter to ensure + # that the current point can be returned by the method. + if self._callback is not None: + sig = signature(self._callback) + try: + x_best, fun_best, _ = self.best_eval(penalty) + x_best = self.build_x(x_best) + if set(sig.parameters) == {"intermediate_result"}: + intermediate_result = OptimizeResult( + x=x_best, + fun=fun_best, + # maxcv=maxcv_best, + ) + self._callback(intermediate_result=intermediate_result) + else: + self._callback(x_best) + except StopIteration as exc: + raise CallbackSuccess from exc + + # Apply the extreme barriers and return. + if np.isnan(fun_val): + fun_val = BARRIER + cub_val[np.isnan(cub_val)] = BARRIER + ceq_val[np.isnan(ceq_val)] = BARRIER + fun_val = max(min(fun_val, BARRIER), -BARRIER) + cub_val = np.maximum(np.minimum(cub_val, BARRIER), -BARRIER) + ceq_val = np.maximum(np.minimum(ceq_val, BARRIER), -BARRIER) + return fun_val, cub_val, ceq_val + + @property + def n(self): + """ + Number of variables. + + Returns + ------- + int + Number of variables. + """ + return self.x0.size + + @property + def n_orig(self): + """ + Number of variables in the original problem (with fixed variables). + + Returns + ------- + int + Number of variables in the original problem (with fixed variables). + """ + return self._fixed_idx.size + + @property + def x0(self): + """ + Initial guess. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Initial guess. + """ + return self._x0 + + @property + def n_eval(self): + """ + Number of function evaluations. + + Returns + ------- + int + Number of function evaluations. + """ + return self._obj.n_eval + + @property + def fun_name(self): + """ + Name of the objective function. + + Returns + ------- + str + Name of the objective function. + """ + return self._obj.name + + @property + def bounds(self): + """ + Bound constraints. + + Returns + ------- + BoundConstraints + Bound constraints. + """ + return self._bounds + + @property + def linear(self): + """ + Linear constraints. + + Returns + ------- + LinearConstraints + Linear constraints. + """ + return self._linear + + @property + def m_bounds(self): + """ + Number of bound constraints. + + Returns + ------- + int + Number of bound constraints. + """ + return self.bounds.m + + @property + def m_linear_ub(self): + """ + Number of linear inequality constraints. + + Returns + ------- + int + Number of linear inequality constraints. + """ + return self.linear.m_ub + + @property + def m_linear_eq(self): + """ + Number of linear equality constraints. + + Returns + ------- + int + Number of linear equality constraints. + """ + return self.linear.m_eq + + @property + def m_nonlinear_ub(self): + """ + Number of nonlinear inequality constraints. + + Returns + ------- + int + Number of nonlinear inequality constraints. + + Raises + ------ + ValueError + If the number of nonlinear inequality constraints is not known. + """ + return self._nonlinear.m_ub + + @property + def m_nonlinear_eq(self): + """ + Number of nonlinear equality constraints. + + Returns + ------- + int + Number of nonlinear equality constraints. + + Raises + ------ + ValueError + If the number of nonlinear equality constraints is not known. + """ + return self._nonlinear.m_eq + + @property + def fun_history(self): + """ + History of objective function evaluations. + + Returns + ------- + `numpy.ndarray`, shape (n_eval,) + History of objective function evaluations. + """ + return np.array(self._fun_history, dtype=float) + + @property + def maxcv_history(self): + """ + History of maximum constraint violations. + + Returns + ------- + `numpy.ndarray`, shape (n_eval,) + History of maximum constraint violations. + """ + return np.array(self._maxcv_history, dtype=float) + + @property + def type(self): + """ + Type of the problem. + + The problem can be either 'unconstrained', 'bound-constrained', + 'linearly constrained', or 'nonlinearly constrained'. + + Returns + ------- + str + Type of the problem. + """ + try: + if self.m_nonlinear_ub > 0 or self.m_nonlinear_eq > 0: + return "nonlinearly constrained" + elif self.m_linear_ub > 0 or self.m_linear_eq > 0: + return "linearly constrained" + elif self.m_bounds > 0: + return "bound-constrained" + else: + return "unconstrained" + except ValueError: + # The number of nonlinear constraints is not known. It may be zero + # if the user provided a nonlinear inequality and/or equality + # constraint function that returns an empty array. However, as this + # is not known before the first call to the function, we assume + # that the problem is nonlinearly constrained. + return "nonlinearly constrained" + + @property + def is_feasibility(self): + """ + Whether the problem is a feasibility problem. + + Returns + ------- + bool + Whether the problem is a feasibility problem. + """ + return self.fun_name == "" + + def build_x(self, x): + """ + Build the full vector of variables from the reduced vector. + + Parameters + ---------- + x : array_like, shape (n,) + Reduced vector of variables. + + Returns + ------- + `numpy.ndarray`, shape (n_orig,) + Full vector of variables. + """ + x_full = np.empty(self.n_orig) + x_full[self._fixed_idx] = self._fixed_val + x_full[~self._fixed_idx] = (x * self._scaling_factor + + self._scaling_shift) + return self._orig_bounds.project(x_full) + + def maxcv(self, x, cub_val=None, ceq_val=None): + """ + Evaluate the maximum constraint violation. + + Parameters + ---------- + x : array_like, shape (n,) + Point at which the maximum constraint violation is evaluated. + cub_val : array_like, shape (m_nonlinear_ub,), optional + Values of the nonlinear inequality constraints. If not provided, + the nonlinear inequality constraints are evaluated at `x`. + ceq_val : array_like, shape (m_nonlinear_eq,), optional + Values of the nonlinear equality constraints. If not provided, + the nonlinear equality constraints are evaluated at `x`. + + Returns + ------- + float + Maximum constraint violation at `x`. + """ + violation = self.violation(x, cub_val=cub_val, ceq_val=ceq_val) + if np.count_nonzero(violation): + return np.max(violation, initial=0.0) + else: + return 0.0 + + def violation(self, x, cub_val=None, ceq_val=None): + violation = [] + if not self.bounds.is_feasible: + b = self.bounds.violation(x) + violation.append(b) + + if len(self.linear.pcs): + lc = self.linear.violation(x) + violation.append(lc) + if len(self._nonlinear.pcs): + nlc = self._nonlinear.violation(x, cub_val, ceq_val) + violation.append(nlc) + + if len(violation): + return np.concatenate(violation) + + def best_eval(self, penalty): + """ + Return the best point in the filter and the corresponding objective and + nonlinear constraint function evaluations. + + Parameters + ---------- + penalty : float + Penalty parameter + + Returns + ------- + `numpy.ndarray`, shape (n,) + Best point. + float + Corresponding objective function value. + float + Corresponding maximum constraint violation. + """ + # If the filter is empty, i.e., if no function evaluation has been + # performed, we evaluate the objective and nonlinear constraint + # functions at the initial guess. + if len(self._fun_filter) == 0: + self(self.x0) + + # Find the best point in the filter. + fun_filter = np.array(self._fun_filter) + maxcv_filter = np.array(self._maxcv_filter) + x_filter = np.array(self._x_filter) + finite_idx = np.isfinite(maxcv_filter) + if np.any(finite_idx): + # At least one point has a finite maximum constraint violation. + feasible_idx = maxcv_filter <= self._feasibility_tol + if np.any(feasible_idx) and not np.all( + np.isnan(fun_filter[feasible_idx]) + ): + # At least one point is feasible and has a well-defined + # objective function value. We select the point with the least + # objective function value. If there is a tie, we select the + # point with the least maximum constraint violation. If there + # is still a tie, we select the most recent point. + fun_min_idx = feasible_idx & ( + fun_filter <= np.nanmin(fun_filter[feasible_idx]) + ) + if np.count_nonzero(fun_min_idx) > 1: + fun_min_idx &= maxcv_filter <= np.min( + maxcv_filter[fun_min_idx] + ) + i = np.flatnonzero(fun_min_idx)[-1] + elif np.any(feasible_idx): + # At least one point is feasible but no feasible point has a + # well-defined objective function value. We select the most + # recent feasible point. + i = np.flatnonzero(feasible_idx)[-1] + else: + # No point is feasible. We first compute the merit function + # value for each point. + merit_filter = np.full_like(fun_filter, np.nan) + merit_filter[finite_idx] = ( + fun_filter[finite_idx] + penalty * maxcv_filter[finite_idx] + ) + if np.all(np.isnan(merit_filter)): + # No point has a well-defined merit function value. In + # other words, among the points with a well-defined maximum + # constraint violation, none has a well-defined objective + # function value. We select the point with the least + # maximum constraint violation. If there is a tie, we + # select the most recent point. + min_maxcv_idx = maxcv_filter <= np.nanmin(maxcv_filter) + i = np.flatnonzero(min_maxcv_idx)[-1] + else: + # At least one point has a well-defined merit function + # value. We select the point with the least merit function + # value. If there is a tie, we select the point with the + # least maximum constraint violation. If there is still a + # tie, we select the point with the least objective + # function value. If there is still a tie, we select the + # most recent point. + merit_min_idx = merit_filter <= np.nanmin(merit_filter) + if np.count_nonzero(merit_min_idx) > 1: + merit_min_idx &= maxcv_filter <= np.min( + maxcv_filter[merit_min_idx] + ) + + if np.count_nonzero(merit_min_idx) > 1: + merit_min_idx &= fun_filter <= np.min( + fun_filter[merit_min_idx] + ) + i = np.flatnonzero(merit_min_idx)[-1] + elif not np.all(np.isnan(fun_filter)): + # No maximum constraint violation is well-defined but at least one + # point has a well-defined objective function value. We select the + # point with the least objective function value. If there is a tie, + # we select the most recent point. + fun_min_idx = fun_filter <= np.nanmin(fun_filter) + i = np.flatnonzero(fun_min_idx)[-1] + else: + # No point has a well-defined maximum constraint violation or + # objective function value. We select the most recent point. + i = len(fun_filter) - 1 + return ( + self.bounds.project(x_filter[i, :]), + fun_filter[i], + maxcv_filter[i], + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/settings.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/settings.py new file mode 100644 index 0000000000000000000000000000000000000000..fec9a091c40f7da974c8752e2143c67922483c77 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/settings.py @@ -0,0 +1,132 @@ +import sys +from enum import Enum + +import numpy as np + + +# Exit status. +class ExitStatus(Enum): + """ + Exit statuses. + """ + + RADIUS_SUCCESS = 0 + TARGET_SUCCESS = 1 + FIXED_SUCCESS = 2 + CALLBACK_SUCCESS = 3 + FEASIBLE_SUCCESS = 4 + MAX_EVAL_WARNING = 5 + MAX_ITER_WARNING = 6 + INFEASIBLE_ERROR = -1 + LINALG_ERROR = -2 + + +class Options(str, Enum): + """ + Options. + """ + + DEBUG = "debug" + FEASIBILITY_TOL = "feasibility_tol" + FILTER_SIZE = "filter_size" + HISTORY_SIZE = "history_size" + MAX_EVAL = "maxfev" + MAX_ITER = "maxiter" + NPT = "nb_points" + RHOBEG = "radius_init" + RHOEND = "radius_final" + SCALE = "scale" + STORE_HISTORY = "store_history" + TARGET = "target" + VERBOSE = "disp" + + +class Constants(str, Enum): + """ + Constants. + """ + + DECREASE_RADIUS_FACTOR = "decrease_radius_factor" + INCREASE_RADIUS_FACTOR = "increase_radius_factor" + INCREASE_RADIUS_THRESHOLD = "increase_radius_threshold" + DECREASE_RADIUS_THRESHOLD = "decrease_radius_threshold" + DECREASE_RESOLUTION_FACTOR = "decrease_resolution_factor" + LARGE_RESOLUTION_THRESHOLD = "large_resolution_threshold" + MODERATE_RESOLUTION_THRESHOLD = "moderate_resolution_threshold" + LOW_RATIO = "low_ratio" + HIGH_RATIO = "high_ratio" + VERY_LOW_RATIO = "very_low_ratio" + PENALTY_INCREASE_THRESHOLD = "penalty_increase_threshold" + PENALTY_INCREASE_FACTOR = "penalty_increase_factor" + SHORT_STEP_THRESHOLD = "short_step_threshold" + LOW_RADIUS_FACTOR = "low_radius_factor" + BYRD_OMOJOKUN_FACTOR = "byrd_omojokun_factor" + THRESHOLD_RATIO_CONSTRAINTS = "threshold_ratio_constraints" + LARGE_SHIFT_FACTOR = "large_shift_factor" + LARGE_GRADIENT_FACTOR = "large_gradient_factor" + RESOLUTION_FACTOR = "resolution_factor" + IMPROVE_TCG = "improve_tcg" + + +# Default options. +DEFAULT_OPTIONS = { + Options.DEBUG.value: False, + Options.FEASIBILITY_TOL.value: np.sqrt(np.finfo(float).eps), + Options.FILTER_SIZE.value: sys.maxsize, + Options.HISTORY_SIZE.value: sys.maxsize, + Options.MAX_EVAL.value: lambda n: 500 * n, + Options.MAX_ITER.value: lambda n: 1000 * n, + Options.NPT.value: lambda n: 2 * n + 1, + Options.RHOBEG.value: 1.0, + Options.RHOEND.value: 1e-6, + Options.SCALE.value: False, + Options.STORE_HISTORY.value: False, + Options.TARGET.value: -np.inf, + Options.VERBOSE.value: False, +} + +# Default constants. +DEFAULT_CONSTANTS = { + Constants.DECREASE_RADIUS_FACTOR.value: 0.5, + Constants.INCREASE_RADIUS_FACTOR.value: np.sqrt(2.0), + Constants.INCREASE_RADIUS_THRESHOLD.value: 2.0, + Constants.DECREASE_RADIUS_THRESHOLD.value: 1.4, + Constants.DECREASE_RESOLUTION_FACTOR.value: 0.1, + Constants.LARGE_RESOLUTION_THRESHOLD.value: 250.0, + Constants.MODERATE_RESOLUTION_THRESHOLD.value: 16.0, + Constants.LOW_RATIO.value: 0.1, + Constants.HIGH_RATIO.value: 0.7, + Constants.VERY_LOW_RATIO.value: 0.01, + Constants.PENALTY_INCREASE_THRESHOLD.value: 1.5, + Constants.PENALTY_INCREASE_FACTOR.value: 2.0, + Constants.SHORT_STEP_THRESHOLD.value: 0.5, + Constants.LOW_RADIUS_FACTOR.value: 0.1, + Constants.BYRD_OMOJOKUN_FACTOR.value: 0.8, + Constants.THRESHOLD_RATIO_CONSTRAINTS.value: 2.0, + Constants.LARGE_SHIFT_FACTOR.value: 10.0, + Constants.LARGE_GRADIENT_FACTOR.value: 10.0, + Constants.RESOLUTION_FACTOR.value: 2.0, + Constants.IMPROVE_TCG.value: True, +} + +# Printing options. +PRINT_OPTIONS = { + "threshold": 6, + "edgeitems": 2, + "linewidth": sys.maxsize, + "formatter": { + "float_kind": lambda x: np.format_float_scientific( + x, + precision=3, + unique=False, + pad_left=2, + ) + }, +} + +# Constants. +BARRIER = 2.0 ** min( + 100, + np.finfo(float).maxexp // 2, + -np.finfo(float).minexp // 2, +) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..43719b8a45d6b56b2f088fc2466d1b3cee4d951d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/__init__.py @@ -0,0 +1,14 @@ +from .geometry import cauchy_geometry, spider_geometry +from .optim import ( + tangential_byrd_omojokun, + constrained_tangential_byrd_omojokun, + normal_byrd_omojokun, +) + +__all__ = [ + "cauchy_geometry", + "spider_geometry", + "tangential_byrd_omojokun", + "constrained_tangential_byrd_omojokun", + "normal_byrd_omojokun", +] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/geometry.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..65bb0620a21391faa01f102523c6dd220d3d36ca --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/geometry.py @@ -0,0 +1,387 @@ +import inspect + +import numpy as np + +from ..utils import get_arrays_tol + + +TINY = np.finfo(float).tiny + + +def cauchy_geometry(const, grad, curv, xl, xu, delta, debug): + r""" + Maximize approximately the absolute value of a quadratic function subject + to bound constraints in a trust region. + + This function solves approximately + + .. math:: + + \max_{s \in \mathbb{R}^n} \quad \bigg\lvert c + g^{\mathsf{T}} s + + \frac{1}{2} s^{\mathsf{T}} H s \bigg\rvert \quad \text{s.t.} \quad + \left\{ \begin{array}{l} + l \le s \le u,\\ + \lVert s \rVert \le \Delta, + \end{array} \right. + + by maximizing the objective function along the constrained Cauchy + direction. + + Parameters + ---------- + const : float + Constant :math:`c` as shown above. + grad : `numpy.ndarray`, shape (n,) + Gradient :math:`g` as shown above. + curv : callable + Curvature of :math:`H` along any vector. + + ``curv(s) -> float`` + + returns :math:`s^{\mathsf{T}} H s`. + xl : `numpy.ndarray`, shape (n,) + Lower bounds :math:`l` as shown above. + xu : `numpy.ndarray`, shape (n,) + Upper bounds :math:`u` as shown above. + delta : float + Trust-region radius :math:`\Delta` as shown above. + debug : bool + Whether to make debugging tests during the execution. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Approximate solution :math:`s`. + + Notes + ----- + This function is described as the first alternative in Section 6.5 of [1]_. + It is assumed that the origin is feasible with respect to the bound + constraints and that `delta` is finite and positive. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods + and Software*. PhD thesis, Department of Applied Mathematics, The Hong + Kong Polytechnic University, Hong Kong, China, 2022. URL: + https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + if debug: + assert isinstance(const, float) + assert isinstance(grad, np.ndarray) and grad.ndim == 1 + assert inspect.signature(curv).bind(grad) + assert isinstance(xl, np.ndarray) and xl.shape == grad.shape + assert isinstance(xu, np.ndarray) and xu.shape == grad.shape + assert isinstance(delta, float) + assert isinstance(debug, bool) + tol = get_arrays_tol(xl, xu) + assert np.all(xl <= tol) + assert np.all(xu >= -tol) + assert np.isfinite(delta) and delta > 0.0 + xl = np.minimum(xl, 0.0) + xu = np.maximum(xu, 0.0) + + # To maximize the absolute value of a quadratic function, we maximize the + # function itself or its negative, and we choose the solution that provides + # the largest function value. + step1, q_val1 = _cauchy_geom(const, grad, curv, xl, xu, delta, debug) + step2, q_val2 = _cauchy_geom( + -const, + -grad, + lambda x: -curv(x), + xl, + xu, + delta, + debug, + ) + step = step1 if abs(q_val1) >= abs(q_val2) else step2 + + if debug: + assert np.all(xl <= step) + assert np.all(step <= xu) + assert np.linalg.norm(step) < 1.1 * delta + return step + + +def spider_geometry(const, grad, curv, xpt, xl, xu, delta, debug): + r""" + Maximize approximately the absolute value of a quadratic function subject + to bound constraints in a trust region. + + This function solves approximately + + .. math:: + + \max_{s \in \mathbb{R}^n} \quad \bigg\lvert c + g^{\mathsf{T}} s + + \frac{1}{2} s^{\mathsf{T}} H s \bigg\rvert \quad \text{s.t.} \quad + \left\{ \begin{array}{l} + l \le s \le u,\\ + \lVert s \rVert \le \Delta, + \end{array} \right. + + by maximizing the objective function along given straight lines. + + Parameters + ---------- + const : float + Constant :math:`c` as shown above. + grad : `numpy.ndarray`, shape (n,) + Gradient :math:`g` as shown above. + curv : callable + Curvature of :math:`H` along any vector. + + ``curv(s) -> float`` + + returns :math:`s^{\mathsf{T}} H s`. + xpt : `numpy.ndarray`, shape (n, npt) + Points defining the straight lines. The straight lines considered are + the ones passing through the origin and the points in `xpt`. + xl : `numpy.ndarray`, shape (n,) + Lower bounds :math:`l` as shown above. + xu : `numpy.ndarray`, shape (n,) + Upper bounds :math:`u` as shown above. + delta : float + Trust-region radius :math:`\Delta` as shown above. + debug : bool + Whether to make debugging tests during the execution. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Approximate solution :math:`s`. + + Notes + ----- + This function is described as the second alternative in Section 6.5 of + [1]_. It is assumed that the origin is feasible with respect to the bound + constraints and that `delta` is finite and positive. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods + and Software*. PhD thesis, Department of Applied Mathematics, The Hong + Kong Polytechnic University, Hong Kong, China, 2022. URL: + https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + if debug: + assert isinstance(const, float) + assert isinstance(grad, np.ndarray) and grad.ndim == 1 + assert inspect.signature(curv).bind(grad) + assert ( + isinstance(xpt, np.ndarray) + and xpt.ndim == 2 + and xpt.shape[0] == grad.size + ) + assert isinstance(xl, np.ndarray) and xl.shape == grad.shape + assert isinstance(xu, np.ndarray) and xu.shape == grad.shape + assert isinstance(delta, float) + assert isinstance(debug, bool) + tol = get_arrays_tol(xl, xu) + assert np.all(xl <= tol) + assert np.all(xu >= -tol) + assert np.isfinite(delta) and delta > 0.0 + xl = np.minimum(xl, 0.0) + xu = np.maximum(xu, 0.0) + + # Iterate through the straight lines. + step = np.zeros_like(grad) + q_val = const + s_norm = np.linalg.norm(xpt, axis=0) + + # Set alpha_xl to the step size for the lower-bound constraint and + # alpha_xu to the step size for the upper-bound constraint. + + # xl.shape = (N,) + # xpt.shape = (N, M) + # i_xl_pos.shape = (M, N) + i_xl_pos = (xl > -np.inf) & (xpt.T > -TINY * xl) + i_xl_neg = (xl > -np.inf) & (xpt.T < TINY * xl) + i_xu_pos = (xu < np.inf) & (xpt.T > TINY * xu) + i_xu_neg = (xu < np.inf) & (xpt.T < -TINY * xu) + + # (M, N) + alpha_xl_pos = np.atleast_2d( + np.broadcast_to(xl, i_xl_pos.shape)[i_xl_pos] / xpt.T[i_xl_pos] + ) + # (M,) + alpha_xl_pos = np.max(alpha_xl_pos, axis=1, initial=-np.inf) + # make sure it's (M,) + alpha_xl_pos = np.broadcast_to(np.atleast_1d(alpha_xl_pos), xpt.shape[1]) + + alpha_xl_neg = np.atleast_2d( + np.broadcast_to(xl, i_xl_neg.shape)[i_xl_neg] / xpt.T[i_xl_neg] + ) + alpha_xl_neg = np.max(alpha_xl_neg, axis=1, initial=np.inf) + alpha_xl_neg = np.broadcast_to(np.atleast_1d(alpha_xl_neg), xpt.shape[1]) + + alpha_xu_neg = np.atleast_2d( + np.broadcast_to(xu, i_xu_neg.shape)[i_xu_neg] / xpt.T[i_xu_neg] + ) + alpha_xu_neg = np.max(alpha_xu_neg, axis=1, initial=-np.inf) + alpha_xu_neg = np.broadcast_to(np.atleast_1d(alpha_xu_neg), xpt.shape[1]) + + alpha_xu_pos = np.atleast_2d( + np.broadcast_to(xu, i_xu_pos.shape)[i_xu_pos] / xpt.T[i_xu_pos] + ) + alpha_xu_pos = np.max(alpha_xu_pos, axis=1, initial=np.inf) + alpha_xu_pos = np.broadcast_to(np.atleast_1d(alpha_xu_pos), xpt.shape[1]) + + for k in range(xpt.shape[1]): + # Set alpha_tr to the step size for the trust-region constraint. + if s_norm[k] > TINY * delta: + alpha_tr = max(delta / s_norm[k], 0.0) + else: + # The current straight line is basically zero. + continue + + alpha_bd_pos = max(min(alpha_xu_pos[k], alpha_xl_neg[k]), 0.0) + alpha_bd_neg = min(max(alpha_xl_pos[k], alpha_xu_neg[k]), 0.0) + + # Set alpha_quad_pos and alpha_quad_neg to the step size to the extrema + # of the quadratic function along the positive and negative directions. + grad_step = grad @ xpt[:, k] + curv_step = curv(xpt[:, k]) + if ( + grad_step >= 0.0 + and curv_step < -TINY * grad_step + or grad_step <= 0.0 + and curv_step > -TINY * grad_step + ): + alpha_quad_pos = max(-grad_step / curv_step, 0.0) + else: + alpha_quad_pos = np.inf + if ( + grad_step >= 0.0 + and curv_step > TINY * grad_step + or grad_step <= 0.0 + and curv_step < TINY * grad_step + ): + alpha_quad_neg = min(-grad_step / curv_step, 0.0) + else: + alpha_quad_neg = -np.inf + + # Select the step that provides the largest value of the objective + # function if it improves the current best. The best positive step is + # either the one that reaches the constraints or the one that reaches + # the extremum of the objective function along the current direction + # (only possible if the resulting step is feasible). We test both, and + # we perform similar calculations along the negative step. + # N.B.: we select the largest possible step among all the ones that + # maximize the objective function. This is to avoid returning the zero + # step in some extreme cases. + alpha_pos = min(alpha_tr, alpha_bd_pos) + alpha_neg = max(-alpha_tr, alpha_bd_neg) + q_val_pos = ( + const + alpha_pos * grad_step + 0.5 * alpha_pos**2.0 * curv_step + ) + q_val_neg = ( + const + alpha_neg * grad_step + 0.5 * alpha_neg**2.0 * curv_step + ) + if alpha_quad_pos < alpha_pos: + q_val_quad_pos = ( + const + + alpha_quad_pos * grad_step + + 0.5 * alpha_quad_pos**2.0 * curv_step + ) + if abs(q_val_quad_pos) > abs(q_val_pos): + alpha_pos = alpha_quad_pos + q_val_pos = q_val_quad_pos + if alpha_quad_neg > alpha_neg: + q_val_quad_neg = ( + const + + alpha_quad_neg * grad_step + + 0.5 * alpha_quad_neg**2.0 * curv_step + ) + if abs(q_val_quad_neg) > abs(q_val_neg): + alpha_neg = alpha_quad_neg + q_val_neg = q_val_quad_neg + if abs(q_val_pos) >= abs(q_val_neg) and abs(q_val_pos) > abs(q_val): + step = np.clip(alpha_pos * xpt[:, k], xl, xu) + q_val = q_val_pos + elif abs(q_val_neg) > abs(q_val_pos) and abs(q_val_neg) > abs(q_val): + step = np.clip(alpha_neg * xpt[:, k], xl, xu) + q_val = q_val_neg + + if debug: + assert np.all(xl <= step) + assert np.all(step <= xu) + assert np.linalg.norm(step) < 1.1 * delta + return step + + +def _cauchy_geom(const, grad, curv, xl, xu, delta, debug): + """ + Same as `bound_constrained_cauchy_step` without the absolute value. + """ + # Calculate the initial active set. + fixed_xl = (xl < 0.0) & (grad > 0.0) + fixed_xu = (xu > 0.0) & (grad < 0.0) + + # Calculate the Cauchy step. + cauchy_step = np.zeros_like(grad) + cauchy_step[fixed_xl] = xl[fixed_xl] + cauchy_step[fixed_xu] = xu[fixed_xu] + if np.linalg.norm(cauchy_step) > delta: + working = fixed_xl | fixed_xu + while True: + # Calculate the Cauchy step for the directions in the working set. + g_norm = np.linalg.norm(grad[working]) + delta_reduced = np.sqrt( + delta**2.0 - cauchy_step[~working] @ cauchy_step[~working] + ) + if g_norm > TINY * abs(delta_reduced): + mu = max(delta_reduced / g_norm, 0.0) + else: + break + cauchy_step[working] = mu * grad[working] + + # Update the working set. + fixed_xl = working & (cauchy_step < xl) + fixed_xu = working & (cauchy_step > xu) + if not np.any(fixed_xl) and not np.any(fixed_xu): + # Stop the calculations as the Cauchy step is now feasible. + break + cauchy_step[fixed_xl] = xl[fixed_xl] + cauchy_step[fixed_xu] = xu[fixed_xu] + working = working & ~(fixed_xl | fixed_xu) + + # Calculate the step that maximizes the quadratic along the Cauchy step. + grad_step = grad @ cauchy_step + if grad_step >= 0.0: + # Set alpha_tr to the step size for the trust-region constraint. + s_norm = np.linalg.norm(cauchy_step) + if s_norm > TINY * delta: + alpha_tr = max(delta / s_norm, 0.0) + else: + # The Cauchy step is basically zero. + alpha_tr = 0.0 + + # Set alpha_quad to the step size for the maximization problem. + curv_step = curv(cauchy_step) + if curv_step < -TINY * grad_step: + alpha_quad = max(-grad_step / curv_step, 0.0) + else: + alpha_quad = np.inf + + # Set alpha_bd to the step size for the bound constraints. + i_xl = (xl > -np.inf) & (cauchy_step < TINY * xl) + i_xu = (xu < np.inf) & (cauchy_step > TINY * xu) + alpha_xl = np.min(xl[i_xl] / cauchy_step[i_xl], initial=np.inf) + alpha_xu = np.min(xu[i_xu] / cauchy_step[i_xu], initial=np.inf) + alpha_bd = min(alpha_xl, alpha_xu) + + # Calculate the solution and the corresponding function value. + alpha = min(alpha_tr, alpha_quad, alpha_bd) + step = np.clip(alpha * cauchy_step, xl, xu) + q_val = const + alpha * grad_step + 0.5 * alpha**2.0 * curv_step + else: + # This case is never reached in exact arithmetic. It prevents this + # function to return a step that decreases the objective function. + step = np.zeros_like(grad) + q_val = const + + if debug: + assert np.all(xl <= step) + assert np.all(step <= xu) + assert np.linalg.norm(step) < 1.1 * delta + return step, q_val diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/optim.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/optim.py new file mode 100644 index 0000000000000000000000000000000000000000..8638f2b382b01d2b9051a31fcbe6e6945dcf638d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/subsolvers/optim.py @@ -0,0 +1,1203 @@ +import inspect + +import numpy as np +from scipy.linalg import qr + +from ..utils import get_arrays_tol + + +TINY = np.finfo(float).tiny +EPS = np.finfo(float).eps + + +def tangential_byrd_omojokun(grad, hess_prod, xl, xu, delta, debug, **kwargs): + r""" + Minimize approximately a quadratic function subject to bound constraints in + a trust region. + + This function solves approximately + + .. math:: + + \min_{s \in \mathbb{R}^n} \quad g^{\mathsf{T}} s + \frac{1}{2} + s^{\mathsf{T}} H s \quad \text{s.t.} \quad + \left\{ \begin{array}{l} + l \le s \le u\\ + \lVert s \rVert \le \Delta, + \end{array} \right. + + using an active-set variation of the truncated conjugate gradient method. + + Parameters + ---------- + grad : `numpy.ndarray`, shape (n,) + Gradient :math:`g` as shown above. + hess_prod : callable + Product of the Hessian matrix :math:`H` with any vector. + + ``hess_prod(s) -> `numpy.ndarray`, shape (n,)`` + + returns the product :math:`H s`. + xl : `numpy.ndarray`, shape (n,) + Lower bounds :math:`l` as shown above. + xu : `numpy.ndarray`, shape (n,) + Upper bounds :math:`u` as shown above. + delta : float + Trust-region radius :math:`\Delta` as shown above. + debug : bool + Whether to make debugging tests during the execution. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Approximate solution :math:`s`. + + Other Parameters + ---------------- + improve_tcg : bool, optional + If True, a solution generated by the truncated conjugate gradient + method that is on the boundary of the trust region is improved by + moving around the trust-region boundary on the two-dimensional space + spanned by the solution and the gradient of the quadratic function at + the solution (default is True). + + Notes + ----- + This function implements Algorithm 6.2 of [1]_. It is assumed that the + origin is feasible with respect to the bound constraints and that `delta` + is finite and positive. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods + and Software*. PhD thesis, Department of Applied Mathematics, The Hong + Kong Polytechnic University, Hong Kong, China, 2022. URL: + https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + if debug: + assert isinstance(grad, np.ndarray) and grad.ndim == 1 + assert inspect.signature(hess_prod).bind(grad) + assert isinstance(xl, np.ndarray) and xl.shape == grad.shape + assert isinstance(xu, np.ndarray) and xu.shape == grad.shape + assert isinstance(delta, float) + assert isinstance(debug, bool) + tol = get_arrays_tol(xl, xu) + assert np.all(xl <= tol) + assert np.all(xu >= -tol) + assert np.isfinite(delta) and delta > 0.0 + xl = np.minimum(xl, 0.0) + xu = np.maximum(xu, 0.0) + + # Copy the arrays that may be modified by the code below. + n = grad.size + grad = np.copy(grad) + grad_orig = np.copy(grad) + + # Calculate the initial active set. + free_bd = ((xl < 0.0) | (grad < 0.0)) & ((xu > 0.0) | (grad > 0.0)) + + # Set the initial iterate and the initial search direction. + step = np.zeros_like(grad) + sd = np.zeros_like(step) + sd[free_bd] = -grad[free_bd] + + k = 0 + reduct = 0.0 + boundary_reached = False + while k < np.count_nonzero(free_bd): + # Stop the computations if sd is not a descent direction. + grad_sd = grad @ sd + if grad_sd >= -10.0 * EPS * n * max(1.0, np.linalg.norm(grad)): + break + + # Set alpha_tr to the step size for the trust-region constraint. + try: + alpha_tr = _alpha_tr(step, sd, delta) + except ZeroDivisionError: + break + + # Stop the computations if a step along sd is expected to give a + # relatively small reduction in the objective function. + if -alpha_tr * grad_sd <= 1e-8 * reduct: + break + + # Set alpha_quad to the step size for the minimization problem. + hess_sd = hess_prod(sd) + curv_sd = sd @ hess_sd + if curv_sd > TINY * abs(grad_sd): + alpha_quad = max(-grad_sd / curv_sd, 0.0) + else: + alpha_quad = np.inf + + # Stop the computations if the reduction in the objective function + # provided by an unconstrained step is small. + alpha = min(alpha_tr, alpha_quad) + if -alpha * (grad_sd + 0.5 * alpha * curv_sd) <= 1e-8 * reduct: + break + + # Set alpha_bd to the step size for the bound constraints. + i_xl = (xl > -np.inf) & (sd < -TINY * np.abs(xl - step)) + i_xu = (xu < np.inf) & (sd > TINY * np.abs(xu - step)) + all_alpha_xl = np.full_like(step, np.inf) + all_alpha_xu = np.full_like(step, np.inf) + all_alpha_xl[i_xl] = np.maximum( + (xl[i_xl] - step[i_xl]) / sd[i_xl], + 0.0, + ) + all_alpha_xu[i_xu] = np.maximum( + (xu[i_xu] - step[i_xu]) / sd[i_xu], + 0.0, + ) + alpha_xl = np.min(all_alpha_xl) + alpha_xu = np.min(all_alpha_xu) + alpha_bd = min(alpha_xl, alpha_xu) + + # Update the iterate. + alpha = min(alpha, alpha_bd) + if alpha > 0.0: + step[free_bd] = np.clip( + step[free_bd] + alpha * sd[free_bd], + xl[free_bd], + xu[free_bd], + ) + grad += alpha * hess_sd + reduct -= alpha * (grad_sd + 0.5 * alpha * curv_sd) + + if alpha < min(alpha_tr, alpha_bd): + # The current iteration is a conjugate gradient iteration. Update + # the search direction so that it is conjugate (with respect to H) + # to all the previous search directions. + beta = (grad[free_bd] @ hess_sd[free_bd]) / curv_sd + sd[free_bd] = beta * sd[free_bd] - grad[free_bd] + sd[~free_bd] = 0.0 + k += 1 + elif alpha < alpha_tr: + # The iterate is restricted by a bound constraint. Add this bound + # constraint to the active set, and restart the calculations. + if alpha_xl <= alpha: + i_new = np.argmin(all_alpha_xl) + step[i_new] = xl[i_new] + else: + i_new = np.argmin(all_alpha_xu) + step[i_new] = xu[i_new] + free_bd[i_new] = False + sd[free_bd] = -grad[free_bd] + sd[~free_bd] = 0.0 + k = 0 + else: + # The current iterate is on the trust-region boundary. Add all the + # active bounds to the working set to prepare for the improvement + # of the solution, and stop the iterations. + if alpha_xl <= alpha: + i_new = _argmin(all_alpha_xl) + step[i_new] = xl[i_new] + free_bd[i_new] = False + if alpha_xu <= alpha: + i_new = _argmin(all_alpha_xu) + step[i_new] = xu[i_new] + free_bd[i_new] = False + boundary_reached = True + break + + # Attempt to improve the solution on the trust-region boundary. + if kwargs.get("improve_tcg", True) and boundary_reached: + step_base = np.copy(step) + step_comparator = grad_orig @ step_base + 0.5 * step_base @ hess_prod( + step_base + ) + + while np.count_nonzero(free_bd) > 0: + # Check whether a substantial reduction in the objective function + # is possible, and set the search direction. + step_sq = step[free_bd] @ step[free_bd] + grad_sq = grad[free_bd] @ grad[free_bd] + grad_step = grad[free_bd] @ step[free_bd] + grad_sd = -np.sqrt(max(step_sq * grad_sq - grad_step**2.0, 0.0)) + sd[free_bd] = grad_step * step[free_bd] - step_sq * grad[free_bd] + sd[~free_bd] = 0.0 + if grad_sd >= -1e-8 * reduct or np.any( + grad_sd >= -TINY * np.abs(sd[free_bd]) + ): + break + sd[free_bd] /= -grad_sd + + # Calculate an upper bound for the tangent of half the angle theta + # of this alternative iteration. The step will be updated as: + # step = cos(theta) * step + sin(theta) * sd. + temp_xl = np.zeros(n) + temp_xu = np.zeros(n) + temp_xl[free_bd] = ( + step[free_bd] ** 2.0 + sd[free_bd] ** 2.0 - xl[free_bd] ** 2.0 + ) + temp_xu[free_bd] = ( + step[free_bd] ** 2.0 + sd[free_bd] ** 2.0 - xu[free_bd] ** 2.0 + ) + temp_xl[temp_xl > 0.0] = ( + np.sqrt(temp_xl[temp_xl > 0.0]) - sd[temp_xl > 0.0] + ) + temp_xu[temp_xu > 0.0] = ( + np.sqrt(temp_xu[temp_xu > 0.0]) + sd[temp_xu > 0.0] + ) + dist_xl = np.maximum(step - xl, 0.0) + dist_xu = np.maximum(xu - step, 0.0) + i_xl = temp_xl > TINY * dist_xl + i_xu = temp_xu > TINY * dist_xu + all_t_xl = np.ones(n) + all_t_xu = np.ones(n) + all_t_xl[i_xl] = np.minimum( + all_t_xl[i_xl], + dist_xl[i_xl] / temp_xl[i_xl], + ) + all_t_xu[i_xu] = np.minimum( + all_t_xu[i_xu], + dist_xu[i_xu] / temp_xu[i_xu], + ) + t_xl = np.min(all_t_xl) + t_xu = np.min(all_t_xu) + t_bd = min(t_xl, t_xu) + + # Calculate some curvature information. + hess_step = hess_prod(step) + hess_sd = hess_prod(sd) + curv_step = step @ hess_step + curv_sd = sd @ hess_sd + curv_step_sd = step @ hess_sd + + # For a range of equally spaced values of tan(0.5 * theta), + # calculate the reduction in the objective function that would be + # obtained by accepting the corresponding angle. + n_samples = 20 + n_samples = int((n_samples - 3) * t_bd + 3) + t_samples = np.linspace(t_bd / n_samples, t_bd, n_samples) + sin_values = 2.0 * t_samples / (1.0 + t_samples**2.0) + all_reduct = sin_values * ( + grad_step * t_samples + - grad_sd + - t_samples * curv_step + + sin_values + * (t_samples * curv_step_sd - 0.5 * (curv_sd - curv_step)) + ) + if np.all(all_reduct <= 0.0): + # No reduction in the objective function is obtained. + break + + # Accept the angle that provides the largest reduction in the + # objective function, and update the iterate. + i_max = np.argmax(all_reduct) + cos_value = (1.0 - t_samples[i_max] ** 2.0) / ( + 1.0 + t_samples[i_max] ** 2.0 + ) + step[free_bd] = ( + cos_value * step[free_bd] + sin_values[i_max] * sd[free_bd] + ) + grad += (cos_value - 1.0) * hess_step + sin_values[i_max] * hess_sd + reduct += all_reduct[i_max] + + # If the above angle is restricted by bound constraints, add them + # to the working set, and restart the alternative iteration. + # Otherwise, the calculations are terminated. + if t_bd < 1.0 and i_max == n_samples - 1: + if t_xl <= t_bd: + i_new = _argmin(all_t_xl) + step[i_new] = xl[i_new] + free_bd[i_new] = False + if t_xu <= t_bd: + i_new = _argmin(all_t_xu) + step[i_new] = xu[i_new] + free_bd[i_new] = False + else: + break + + # Ensure that the alternative iteration improves the objective + # function. + if grad_orig @ step + 0.5 * step @ hess_prod(step) > step_comparator: + step = step_base + + if debug: + assert np.all(xl <= step) + assert np.all(step <= xu) + assert np.linalg.norm(step) < 1.1 * delta + return step + + +def constrained_tangential_byrd_omojokun( + grad, + hess_prod, + xl, + xu, + aub, + bub, + aeq, + delta, + debug, + **kwargs, +): + r""" + Minimize approximately a quadratic function subject to bound and linear + constraints in a trust region. + + This function solves approximately + + .. math:: + + \min_{s \in \mathbb{R}^n} \quad g^{\mathsf{T}} s + \frac{1}{2} + s^{\mathsf{T}} H s \quad \text{s.t.} \quad + \left\{ \begin{array}{l} + l \le s \le u,\\ + A_{\scriptscriptstyle I} s \le b_{\scriptscriptstyle I},\\ + A_{\scriptscriptstyle E} s = 0,\\ + \lVert s \rVert \le \Delta, + \end{array} \right. + + using an active-set variation of the truncated conjugate gradient method. + + Parameters + ---------- + grad : `numpy.ndarray`, shape (n,) + Gradient :math:`g` as shown above. + hess_prod : callable + Product of the Hessian matrix :math:`H` with any vector. + + ``hess_prod(s) -> `numpy.ndarray`, shape (n,)`` + + returns the product :math:`H s`. + xl : `numpy.ndarray`, shape (n,) + Lower bounds :math:`l` as shown above. + xu : `numpy.ndarray`, shape (n,) + Upper bounds :math:`u` as shown above. + aub : `numpy.ndarray`, shape (m_linear_ub, n) + Coefficient matrix :math:`A_{\scriptscriptstyle I}` as shown above. + bub : `numpy.ndarray`, shape (m_linear_ub,) + Right-hand side :math:`b_{\scriptscriptstyle I}` as shown above. + aeq : `numpy.ndarray`, shape (m_linear_eq, n) + Coefficient matrix :math:`A_{\scriptscriptstyle E}` as shown above. + delta : float + Trust-region radius :math:`\Delta` as shown above. + debug : bool + Whether to make debugging tests during the execution. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Approximate solution :math:`s`. + + Other Parameters + ---------------- + improve_tcg : bool, optional + If True, a solution generated by the truncated conjugate gradient + method that is on the boundary of the trust region is improved by + moving around the trust-region boundary on the two-dimensional space + spanned by the solution and the gradient of the quadratic function at + the solution (default is True). + + Notes + ----- + This function implements Algorithm 6.3 of [1]_. It is assumed that the + origin is feasible with respect to the bound and linear constraints, and + that `delta` is finite and positive. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods + and Software*. PhD thesis, Department of Applied Mathematics, The Hong + Kong Polytechnic University, Hong Kong, China, 2022. URL: + https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + if debug: + assert isinstance(grad, np.ndarray) and grad.ndim == 1 + assert inspect.signature(hess_prod).bind(grad) + assert isinstance(xl, np.ndarray) and xl.shape == grad.shape + assert isinstance(xu, np.ndarray) and xu.shape == grad.shape + assert ( + isinstance(aub, np.ndarray) + and aub.ndim == 2 + and aub.shape[1] == grad.size + ) + assert ( + isinstance(bub, np.ndarray) + and bub.ndim == 1 + and bub.size == aub.shape[0] + ) + assert ( + isinstance(aeq, np.ndarray) + and aeq.ndim == 2 + and aeq.shape[1] == grad.size + ) + assert isinstance(delta, float) + assert isinstance(debug, bool) + tol = get_arrays_tol(xl, xu) + assert np.all(xl <= tol) + assert np.all(xu >= -tol) + assert np.all(bub >= -tol) + assert np.isfinite(delta) and delta > 0.0 + xl = np.minimum(xl, 0.0) + xu = np.maximum(xu, 0.0) + bub = np.maximum(bub, 0.0) + + # Copy the arrays that may be modified by the code below. + n = grad.size + grad = np.copy(grad) + grad_orig = np.copy(grad) + + # Calculate the initial active set. + free_xl = (xl < 0.0) | (grad < 0.0) + free_xu = (xu > 0.0) | (grad > 0.0) + free_ub = (bub > 0.0) | (aub @ grad > 0.0) + n_act, q = qr_tangential_byrd_omojokun(aub, aeq, free_xl, free_xu, free_ub) + + # Set the initial iterate and the initial search direction. + step = np.zeros_like(grad) + sd = -q[:, n_act:] @ (q[:, n_act:].T @ grad) + resid = np.copy(bub) + + k = 0 + reduct = 0.0 + boundary_reached = False + while k < n - n_act: + # Stop the computations if sd is not a descent direction. + grad_sd = grad @ sd + if grad_sd >= -10.0 * EPS * n * max(1.0, np.linalg.norm(grad)): + break + + # Set alpha_tr to the step size for the trust-region constraint. + try: + alpha_tr = _alpha_tr(step, sd, delta) + except ZeroDivisionError: + break + + # Stop the computations if a step along sd is expected to give a + # relatively small reduction in the objective function. + if -alpha_tr * grad_sd <= 1e-8 * reduct: + break + + # Set alpha_quad to the step size for the minimization problem. + hess_sd = hess_prod(sd) + curv_sd = sd @ hess_sd + if curv_sd > TINY * abs(grad_sd): + alpha_quad = max(-grad_sd / curv_sd, 0.0) + else: + alpha_quad = np.inf + + # Stop the computations if the reduction in the objective function + # provided by an unconstrained step is small. + alpha = min(alpha_tr, alpha_quad) + if -alpha * (grad_sd + 0.5 * alpha * curv_sd) <= 1e-8 * reduct: + break + + # Set alpha_bd to the step size for the bound constraints. + i_xl = free_xl & (xl > -np.inf) & (sd < -TINY * np.abs(xl - step)) + i_xu = free_xu & (xu < np.inf) & (sd > TINY * np.abs(xu - step)) + all_alpha_xl = np.full_like(step, np.inf) + all_alpha_xu = np.full_like(step, np.inf) + all_alpha_xl[i_xl] = np.maximum( + (xl[i_xl] - step[i_xl]) / sd[i_xl], + 0.0, + ) + all_alpha_xu[i_xu] = np.maximum( + (xu[i_xu] - step[i_xu]) / sd[i_xu], + 0.0, + ) + alpha_xl = np.min(all_alpha_xl) + alpha_xu = np.min(all_alpha_xu) + alpha_bd = min(alpha_xl, alpha_xu) + + # Set alpha_ub to the step size for the linear constraints. + aub_sd = aub @ sd + i_ub = free_ub & (aub_sd > TINY * np.abs(resid)) + all_alpha_ub = np.full_like(bub, np.inf) + all_alpha_ub[i_ub] = resid[i_ub] / aub_sd[i_ub] + alpha_ub = np.min(all_alpha_ub, initial=np.inf) + + # Update the iterate. + alpha = min(alpha, alpha_bd, alpha_ub) + if alpha > 0.0: + step = np.clip(step + alpha * sd, xl, xu) + grad += alpha * hess_sd + resid = np.maximum(0.0, resid - alpha * aub_sd) + reduct -= alpha * (grad_sd + 0.5 * alpha * curv_sd) + + if alpha < min(alpha_tr, alpha_bd, alpha_ub): + # The current iteration is a conjugate gradient iteration. Update + # the search direction so that it is conjugate (with respect to H) + # to all the previous search directions. + grad_proj = q[:, n_act:] @ (q[:, n_act:].T @ grad) + beta = (grad_proj @ hess_sd) / curv_sd + sd = beta * sd - grad_proj + k += 1 + elif alpha < alpha_tr: + # The iterate is restricted by a bound/linear constraint. Add this + # constraint to the active set, and restart the calculations. + if alpha_xl <= alpha: + i_new = np.argmin(all_alpha_xl) + step[i_new] = xl[i_new] + free_xl[i_new] = False + elif alpha_xu <= alpha: + i_new = np.argmin(all_alpha_xu) + step[i_new] = xu[i_new] + free_xu[i_new] = False + else: + i_new = np.argmin(all_alpha_ub) + free_ub[i_new] = False + n_act, q = qr_tangential_byrd_omojokun( + aub, + aeq, + free_xl, + free_xu, + free_ub, + ) + sd = -q[:, n_act:] @ (q[:, n_act:].T @ grad) + k = 0 + else: + # The current iterate is on the trust-region boundary. Add all the + # active bound/linear constraints to the working set to prepare for + # the improvement of the solution, and stop the iterations. + if alpha_xl <= alpha: + i_new = _argmin(all_alpha_xl) + step[i_new] = xl[i_new] + free_xl[i_new] = False + if alpha_xu <= alpha: + i_new = _argmin(all_alpha_xu) + step[i_new] = xu[i_new] + free_xu[i_new] = False + if alpha_ub <= alpha: + i_new = _argmin(all_alpha_ub) + free_ub[i_new] = False + n_act, q = qr_tangential_byrd_omojokun( + aub, + aeq, + free_xl, + free_xu, + free_ub, + ) + boundary_reached = True + break + + # Attempt to improve the solution on the trust-region boundary. + if kwargs.get("improve_tcg", True) and boundary_reached and n_act < n: + step_base = np.copy(step) + while n_act < n: + # Check whether a substantial reduction in the objective function + # is possible, and set the search direction. + step_proj = q[:, n_act:] @ (q[:, n_act:].T @ step) + grad_proj = q[:, n_act:] @ (q[:, n_act:].T @ grad) + step_sq = step_proj @ step_proj + grad_sq = grad_proj @ grad_proj + grad_step = grad_proj @ step_proj + grad_sd = -np.sqrt(max(step_sq * grad_sq - grad_step**2.0, 0.0)) + sd = q[:, n_act:] @ ( + q[:, n_act:].T @ (grad_step * step - step_sq * grad) + ) + if grad_sd >= -1e-8 * reduct or np.any( + grad_sd >= -TINY * np.abs(sd) + ): + break + sd /= -grad_sd + + # Calculate an upper bound for the tangent of half the angle theta + # of this alternative iteration for the bound constraints. The step + # will be updated as: + # step += (cos(theta) - 1) * step_proj + sin(theta) * sd. + temp_xl = np.zeros(n) + temp_xu = np.zeros(n) + dist_xl = np.maximum(step - xl, 0.0) + dist_xu = np.maximum(xu - step, 0.0) + temp_xl[free_xl] = sd[free_xl] ** 2.0 - dist_xl[free_xl] * ( + dist_xl[free_xl] - 2.0 * step_proj[free_xl] + ) + temp_xu[free_xu] = sd[free_xu] ** 2.0 - dist_xu[free_xu] * ( + dist_xu[free_xu] + 2.0 * step_proj[free_xu] + ) + temp_xl[temp_xl > 0.0] = ( + np.sqrt(temp_xl[temp_xl > 0.0]) - sd[temp_xl > 0.0] + ) + temp_xu[temp_xu > 0.0] = ( + np.sqrt(temp_xu[temp_xu > 0.0]) + sd[temp_xu > 0.0] + ) + i_xl = temp_xl > TINY * dist_xl + i_xu = temp_xu > TINY * dist_xu + all_t_xl = np.ones(n) + all_t_xu = np.ones(n) + all_t_xl[i_xl] = np.minimum( + all_t_xl[i_xl], + dist_xl[i_xl] / temp_xl[i_xl], + ) + all_t_xu[i_xu] = np.minimum( + all_t_xu[i_xu], + dist_xu[i_xu] / temp_xu[i_xu], + ) + t_xl = np.min(all_t_xl) + t_xu = np.min(all_t_xu) + t_bd = min(t_xl, t_xu) + + # Calculate an upper bound for the tangent of half the angle theta + # of this alternative iteration for the linear constraints. + temp_ub = np.zeros_like(resid) + aub_step = aub @ step_proj + aub_sd = aub @ sd + temp_ub[free_ub] = aub_sd[free_ub] ** 2.0 - resid[free_ub] * ( + resid[free_ub] + 2.0 * aub_step[free_ub] + ) + temp_ub[temp_ub > 0.0] = ( + np.sqrt(temp_ub[temp_ub > 0.0]) + aub_sd[temp_ub > 0.0] + ) + i_ub = temp_ub > TINY * resid + all_t_ub = np.ones_like(resid) + all_t_ub[i_ub] = np.minimum( + all_t_ub[i_ub], + resid[i_ub] / temp_ub[i_ub], + ) + t_ub = np.min(all_t_ub, initial=1.0) + t_min = min(t_bd, t_ub) + + # Calculate some curvature information. + hess_step = hess_prod(step_proj) + hess_sd = hess_prod(sd) + curv_step = step_proj @ hess_step + curv_sd = sd @ hess_sd + curv_step_sd = step_proj @ hess_sd + + # For a range of equally spaced values of tan(0.5 * theta), + # calculate the reduction in the objective function that would be + # obtained by accepting the corresponding angle. + n_samples = 20 + n_samples = int((n_samples - 3) * t_min + 3) + t_samples = np.linspace(t_min / n_samples, t_min, n_samples) + sin_values = 2.0 * t_samples / (1.0 + t_samples**2.0) + all_reduct = sin_values * ( + grad_step * t_samples + - grad_sd + - sin_values + * ( + 0.5 * t_samples**2.0 * curv_step + - 2.0 * t_samples * curv_step_sd + + 0.5 * curv_sd + ) + ) + if np.all(all_reduct <= 0.0): + # No reduction in the objective function is obtained. + break + + # Accept the angle that provides the largest reduction in the + # objective function, and update the iterate. + i_max = np.argmax(all_reduct) + cos_value = (1.0 - t_samples[i_max] ** 2.0) / ( + 1.0 + t_samples[i_max] ** 2.0 + ) + step = np.clip( + step + (cos_value - 1.0) * step_proj + sin_values[i_max] * sd, + xl, + xu, + ) + grad += (cos_value - 1.0) * hess_step + sin_values[i_max] * hess_sd + resid = np.maximum( + 0.0, + resid + - (cos_value - 1.0) * aub_step + - sin_values[i_max] * aub_sd, + ) + reduct += all_reduct[i_max] + + # If the above angle is restricted by bound constraints, add them + # to the working set, and restart the alternative iteration. + # Otherwise, the calculations are terminated. + if t_min < 1.0 and i_max == n_samples - 1: + if t_xl <= t_min: + i_new = _argmin(all_t_xl) + step[i_new] = xl[i_new] + free_xl[i_new] = False + if t_xu <= t_min: + i_new = _argmin(all_t_xu) + step[i_new] = xu[i_new] + free_xl[i_new] = False + if t_ub <= t_min: + i_new = _argmin(all_t_ub) + free_ub[i_new] = False + n_act, q = qr_tangential_byrd_omojokun( + aub, + aeq, + free_xl, + free_xu, + free_ub, + ) + else: + break + + # Ensure that the alternative iteration improves the objective + # function. + if grad_orig @ step + 0.5 * step @ hess_prod( + step + ) > grad_orig @ step_base + 0.5 * step_base @ hess_prod(step_base): + step = step_base + + if debug: + tol = get_arrays_tol(xl, xu) + assert np.all(xl <= step) + assert np.all(step <= xu) + assert np.all(aub @ step <= bub + tol) + assert np.all(np.abs(aeq @ step) <= tol) + assert np.linalg.norm(step) < 1.1 * delta + return step + + +def normal_byrd_omojokun(aub, bub, aeq, beq, xl, xu, delta, debug, **kwargs): + r""" + Minimize approximately a linear constraint violation subject to bound + constraints in a trust region. + + This function solves approximately + + .. math:: + + \min_{s \in \mathbb{R}^n} \quad \frac{1}{2} \big( \lVert \max \{ + A_{\scriptscriptstyle I} s - b_{\scriptscriptstyle I}, 0 \} \rVert^2 + + \lVert A_{\scriptscriptstyle E} s - b_{\scriptscriptstyle E} \rVert^2 + \big) \quad \text{s.t.} + \quad + \left\{ \begin{array}{l} + l \le s \le u,\\ + \lVert s \rVert \le \Delta, + \end{array} \right. + + using a variation of the truncated conjugate gradient method. + + Parameters + ---------- + aub : `numpy.ndarray`, shape (m_linear_ub, n) + Matrix :math:`A_{\scriptscriptstyle I}` as shown above. + bub : `numpy.ndarray`, shape (m_linear_ub,) + Vector :math:`b_{\scriptscriptstyle I}` as shown above. + aeq : `numpy.ndarray`, shape (m_linear_eq, n) + Matrix :math:`A_{\scriptscriptstyle E}` as shown above. + beq : `numpy.ndarray`, shape (m_linear_eq,) + Vector :math:`b_{\scriptscriptstyle E}` as shown above. + xl : `numpy.ndarray`, shape (n,) + Lower bounds :math:`l` as shown above. + xu : `numpy.ndarray`, shape (n,) + Upper bounds :math:`u` as shown above. + delta : float + Trust-region radius :math:`\Delta` as shown above. + debug : bool + Whether to make debugging tests during the execution. + + Returns + ------- + `numpy.ndarray`, shape (n,) + Approximate solution :math:`s`. + + Other Parameters + ---------------- + improve_tcg : bool, optional + If True, a solution generated by the truncated conjugate gradient + method that is on the boundary of the trust region is improved by + moving around the trust-region boundary on the two-dimensional space + spanned by the solution and the gradient of the quadratic function at + the solution (default is True). + + Notes + ----- + This function implements Algorithm 6.4 of [1]_. It is assumed that the + origin is feasible with respect to the bound constraints and that `delta` + is finite and positive. + + References + ---------- + .. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods + and Software*. PhD thesis, Department of Applied Mathematics, The Hong + Kong Polytechnic University, Hong Kong, China, 2022. URL: + https://theses.lib.polyu.edu.hk/handle/200/12294. + """ + if debug: + assert isinstance(aub, np.ndarray) and aub.ndim == 2 + assert ( + isinstance(bub, np.ndarray) + and bub.ndim == 1 + and bub.size == aub.shape[0] + ) + assert ( + isinstance(aeq, np.ndarray) + and aeq.ndim == 2 + and aeq.shape[1] == aub.shape[1] + ) + assert ( + isinstance(beq, np.ndarray) + and beq.ndim == 1 + and beq.size == aeq.shape[0] + ) + assert isinstance(xl, np.ndarray) and xl.shape == (aub.shape[1],) + assert isinstance(xu, np.ndarray) and xu.shape == (aub.shape[1],) + assert isinstance(delta, float) + assert isinstance(debug, bool) + tol = get_arrays_tol(xl, xu) + assert np.all(xl <= tol) + assert np.all(xu >= -tol) + assert np.isfinite(delta) and delta > 0.0 + xl = np.minimum(xl, 0.0) + xu = np.maximum(xu, 0.0) + + # Calculate the initial active set. + m_linear_ub, n = aub.shape + grad = np.r_[aeq.T @ -beq, np.maximum(0.0, -bub)] + free_xl = (xl < 0.0) | (grad[:n] < 0.0) + free_xu = (xu > 0.0) | (grad[:n] > 0.0) + free_slack = bub < 0.0 + free_ub = (bub > 0.0) | (aub @ grad[:n] - grad[n:] > 0.0) + n_act, q = qr_normal_byrd_omojokun( + aub, + free_xl, + free_xu, + free_slack, + free_ub, + ) + + # Calculate an upper bound on the norm of the slack variables. It is not + # used in the original algorithm, but it may prevent undesired behaviors + # engendered by computer rounding errors. + delta_slack = np.sqrt(beq @ beq + grad[n:] @ grad[n:]) + + # Set the initial iterate and the initial search direction. + step = np.zeros(n) + sd = -q[:, n_act:] @ (q[:, n_act:].T @ grad) + resid = bub + grad[n:] + + k = 0 + reduct = 0.0 + boundary_reached = False + while k < n + m_linear_ub - n_act: + # Stop the computations if sd is not a descent direction. + grad_sd = grad @ sd + if grad_sd >= -10.0 * EPS * n * max(1.0, np.linalg.norm(grad)): + break + + # Set alpha_tr to the step size for the trust-region constraint. + try: + alpha_tr = _alpha_tr(step, sd[:n], delta) + except ZeroDivisionError: + alpha_tr = np.inf + + # Prevent undesired behaviors engendered by computer rounding errors by + # considering the trust-region constraint on the slack variables. + try: + alpha_tr = min(alpha_tr, _alpha_tr(grad[n:], sd[n:], delta_slack)) + except ZeroDivisionError: + pass + + # Stop the computations if a step along sd is expected to give a + # relatively small reduction in the objective function. + if -alpha_tr * grad_sd <= 1e-8 * reduct: + break + + # Set alpha_quad to the step size for the minimization problem. + hess_sd = np.r_[aeq.T @ (aeq @ sd[:n]), sd[n:]] + curv_sd = sd @ hess_sd + if curv_sd > TINY * abs(grad_sd): + alpha_quad = max(-grad_sd / curv_sd, 0.0) + else: + alpha_quad = np.inf + + # Stop the computations if the reduction in the objective function + # provided by an unconstrained step is small. + alpha = min(alpha_tr, alpha_quad) + if -alpha * (grad_sd + 0.5 * alpha * curv_sd) <= 1e-8 * reduct: + break + + # Set alpha_bd to the step size for the bound constraints. + i_xl = free_xl & (xl > -np.inf) & (sd[:n] < -TINY * np.abs(xl - step)) + i_xu = free_xu & (xu < np.inf) & (sd[:n] > TINY * np.abs(xu - step)) + i_slack = free_slack & (sd[n:] < -TINY * np.abs(grad[n:])) + all_alpha_xl = np.full_like(step, np.inf) + all_alpha_xu = np.full_like(step, np.inf) + all_alpha_slack = np.full_like(bub, np.inf) + all_alpha_xl[i_xl] = np.maximum( + (xl[i_xl] - step[i_xl]) / sd[:n][i_xl], + 0.0, + ) + all_alpha_xu[i_xu] = np.maximum( + (xu[i_xu] - step[i_xu]) / sd[:n][i_xu], + 0.0, + ) + all_alpha_slack[i_slack] = np.maximum( + -grad[n:][i_slack] / sd[n:][i_slack], + 0.0, + ) + alpha_xl = np.min(all_alpha_xl) + alpha_xu = np.min(all_alpha_xu) + alpha_slack = np.min(all_alpha_slack, initial=np.inf) + alpha_bd = min(alpha_xl, alpha_xu, alpha_slack) + + # Set alpha_ub to the step size for the linear constraints. + aub_sd = aub @ sd[:n] - sd[n:] + i_ub = free_ub & (aub_sd > TINY * np.abs(resid)) + all_alpha_ub = np.full_like(bub, np.inf) + all_alpha_ub[i_ub] = resid[i_ub] / aub_sd[i_ub] + alpha_ub = np.min(all_alpha_ub, initial=np.inf) + + # Update the iterate. + alpha = min(alpha, alpha_bd, alpha_ub) + if alpha > 0.0: + step = np.clip(step + alpha * sd[:n], xl, xu) + grad += alpha * hess_sd + resid = np.maximum(0.0, resid - alpha * aub_sd) + reduct -= alpha * (grad_sd + 0.5 * alpha * curv_sd) + + if alpha < min(alpha_tr, alpha_bd, alpha_ub): + # The current iteration is a conjugate gradient iteration. Update + # the search direction so that it is conjugate (with respect to H) + # to all the previous search directions. + grad_proj = q[:, n_act:] @ (q[:, n_act:].T @ grad) + beta = (grad_proj @ hess_sd) / curv_sd + sd = beta * sd - grad_proj + k += 1 + elif alpha < alpha_tr: + # The iterate is restricted by a bound/linear constraint. Add this + # constraint to the active set, and restart the calculations. + if alpha_xl <= alpha: + i_new = np.argmin(all_alpha_xl) + step[i_new] = xl[i_new] + free_xl[i_new] = False + elif alpha_xu <= alpha: + i_new = np.argmin(all_alpha_xu) + step[i_new] = xu[i_new] + free_xu[i_new] = False + elif alpha_slack <= alpha: + i_new = np.argmin(all_alpha_slack) + free_slack[i_new] = False + else: + i_new = np.argmin(all_alpha_ub) + free_ub[i_new] = False + n_act, q = qr_normal_byrd_omojokun( + aub, free_xl, free_xu, free_slack, free_ub + ) + sd = -q[:, n_act:] @ (q[:, n_act:].T @ grad) + k = 0 + else: + # The current iterate is on the trust-region boundary. Add all the + # active bound constraints to the working set to prepare for the + # improvement of the solution, and stop the iterations. + if alpha_xl <= alpha: + i_new = _argmin(all_alpha_xl) + step[i_new] = xl[i_new] + free_xl[i_new] = False + if alpha_xu <= alpha: + i_new = _argmin(all_alpha_xu) + step[i_new] = xu[i_new] + free_xu[i_new] = False + boundary_reached = True + break + + # Attempt to improve the solution on the trust-region boundary. + if kwargs.get("improve_tcg", True) and boundary_reached: + step_base = np.copy(step) + free_bd = free_xl & free_xu + grad = aub.T @ np.maximum(aub @ step - bub, 0.0) + aeq.T @ ( + aeq @ step - beq + ) + sd = np.zeros(n) + while np.count_nonzero(free_bd) > 0: + # Check whether a substantial reduction in the objective function + # is possible, and set the search direction. + step_sq = step[free_bd] @ step[free_bd] + grad_sq = grad[free_bd] @ grad[free_bd] + grad_step = grad[free_bd] @ step[free_bd] + grad_sd = -np.sqrt(max(step_sq * grad_sq - grad_step**2.0, 0.0)) + sd[free_bd] = grad_step * step[free_bd] - step_sq * grad[free_bd] + sd[~free_bd] = 0.0 + if grad_sd >= -1e-8 * reduct or np.any( + grad_sd >= -TINY * np.abs(sd[free_bd]) + ): + break + sd[free_bd] /= -grad_sd + + # Calculate an upper bound for the tangent of half the angle theta + # of this alternative iteration. The step will be updated as: + # step = cos(theta) * step + sin(theta) * sd. + temp_xl = np.zeros(n) + temp_xu = np.zeros(n) + temp_xl[free_bd] = ( + step[free_bd] ** 2.0 + sd[free_bd] ** 2.0 - xl[free_bd] ** 2.0 + ) + temp_xu[free_bd] = ( + step[free_bd] ** 2.0 + sd[free_bd] ** 2.0 - xu[free_bd] ** 2.0 + ) + temp_xl[temp_xl > 0.0] = ( + np.sqrt(temp_xl[temp_xl > 0.0]) - sd[temp_xl > 0.0] + ) + temp_xu[temp_xu > 0.0] = ( + np.sqrt(temp_xu[temp_xu > 0.0]) + sd[temp_xu > 0.0] + ) + dist_xl = np.maximum(step - xl, 0.0) + dist_xu = np.maximum(xu - step, 0.0) + i_xl = temp_xl > TINY * dist_xl + i_xu = temp_xu > TINY * dist_xu + all_t_xl = np.ones(n) + all_t_xu = np.ones(n) + all_t_xl[i_xl] = np.minimum( + all_t_xl[i_xl], + dist_xl[i_xl] / temp_xl[i_xl], + ) + all_t_xu[i_xu] = np.minimum( + all_t_xu[i_xu], + dist_xu[i_xu] / temp_xu[i_xu], + ) + t_xl = np.min(all_t_xl) + t_xu = np.min(all_t_xu) + t_bd = min(t_xl, t_xu) + + # For a range of equally spaced values of tan(0.5 * theta), + # calculate the reduction in the objective function that would be + # obtained by accepting the corresponding angle. + n_samples = 20 + n_samples = int((n_samples - 3) * t_bd + 3) + t_samples = np.linspace(t_bd / n_samples, t_bd, n_samples) + resid_ub = np.maximum(aub @ step - bub, 0.0) + resid_eq = aeq @ step - beq + step_proj = np.copy(step) + step_proj[~free_bd] = 0.0 + all_reduct = np.empty(n_samples) + for i in range(n_samples): + sin_value = 2.0 * t_samples[i] / (1.0 + t_samples[i] ** 2.0) + step_alt = np.clip( + step + sin_value * (sd - t_samples[i] * step_proj), + xl, + xu, + ) + resid_ub_alt = np.maximum(aub @ step_alt - bub, 0.0) + resid_eq_alt = aeq @ step_alt - beq + all_reduct[i] = 0.5 * ( + resid_ub @ resid_ub + + resid_eq @ resid_eq + - resid_ub_alt @ resid_ub_alt + - resid_eq_alt @ resid_eq_alt + ) + if np.all(all_reduct <= 0.0): + # No reduction in the objective function is obtained. + break + + # Accept the angle that provides the largest reduction in the + # objective function, and update the iterate. + i_max = np.argmax(all_reduct) + cos_value = (1.0 - t_samples[i_max] ** 2.0) / ( + 1.0 + t_samples[i_max] ** 2.0 + ) + sin_value = (2.0 * t_samples[i_max] + / (1.0 + t_samples[i_max] ** 2.0)) + step[free_bd] = cos_value * step[free_bd] + sin_value * sd[free_bd] + grad = aub.T @ np.maximum(aub @ step - bub, 0.0) + aeq.T @ ( + aeq @ step - beq + ) + reduct += all_reduct[i_max] + + # If the above angle is restricted by bound constraints, add them + # to the working set, and restart the alternative iteration. + # Otherwise, the calculations are terminated. + if t_bd < 1.0 and i_max == n_samples - 1: + if t_xl <= t_bd: + i_new = _argmin(all_t_xl) + step[i_new] = xl[i_new] + free_bd[i_new] = False + if t_xu <= t_bd: + i_new = _argmin(all_t_xu) + step[i_new] = xu[i_new] + free_bd[i_new] = False + else: + break + + # Ensure that the alternative iteration improves the objective + # function. + resid_ub = np.maximum(aub @ step - bub, 0.0) + resid_ub_base = np.maximum(aub @ step_base - bub, 0.0) + resid_eq = aeq @ step - beq + resid_eq_base = aeq @ step_base - beq + if ( + resid_ub @ resid_ub + resid_eq @ resid_eq + > resid_ub_base @ resid_ub_base + resid_eq_base @ resid_eq_base + ): + step = step_base + + if debug: + assert np.all(xl <= step) + assert np.all(step <= xu) + assert np.linalg.norm(step) < 1.1 * delta + return step + + +def qr_tangential_byrd_omojokun(aub, aeq, free_xl, free_xu, free_ub): + n = free_xl.size + identity = np.eye(n) + q, r, _ = qr( + np.block( + [ + [aeq], + [aub[~free_ub, :]], + [-identity[~free_xl, :]], + [identity[~free_xu, :]], + ] + ).T, + pivoting=True, + ) + n_act = np.count_nonzero( + np.abs(np.diag(r)) + >= 10.0 + * EPS + * n + * np.linalg.norm(r[: np.min(r.shape), : np.min(r.shape)], axis=0) + ) + return n_act, q + + +def qr_normal_byrd_omojokun(aub, free_xl, free_xu, free_slack, free_ub): + m_linear_ub, n = aub.shape + identity_n = np.eye(n) + identity_m = np.eye(m_linear_ub) + q, r, _ = qr( + np.block( + [ + [ + aub[~free_ub, :], + -identity_m[~free_ub, :], + ], + [ + np.zeros((m_linear_ub - np.count_nonzero(free_slack), n)), + -identity_m[~free_slack, :], + ], + [ + -identity_n[~free_xl, :], + np.zeros((n - np.count_nonzero(free_xl), m_linear_ub)), + ], + [ + identity_n[~free_xu, :], + np.zeros((n - np.count_nonzero(free_xu), m_linear_ub)), + ], + ] + ).T, + pivoting=True, + ) + n_act = np.count_nonzero( + np.abs(np.diag(r)) + >= 10.0 + * EPS + * (n + m_linear_ub) + * np.linalg.norm(r[: np.min(r.shape), : np.min(r.shape)], axis=0) + ) + return n_act, q + + +def _alpha_tr(step, sd, delta): + step_sd = step @ sd + sd_sq = sd @ sd + dist_tr_sq = delta**2.0 - step @ step + temp = np.sqrt(max(step_sd**2.0 + sd_sq * dist_tr_sq, 0.0)) + if step_sd <= 0.0 and sd_sq > TINY * abs(temp - step_sd): + alpha_tr = max((temp - step_sd) / sd_sq, 0.0) + elif abs(temp + step_sd) > TINY * dist_tr_sq: + alpha_tr = max(dist_tr_sq / (temp + step_sd), 0.0) + else: + raise ZeroDivisionError + return alpha_tr + + +def _argmax(x): + return np.flatnonzero(x >= np.max(x)) + + +def _argmin(x): + return np.flatnonzero(x <= np.min(x)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ff8d6fb0c047ae54e7e1fa68e96cbd023c228c69 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/__init__.py @@ -0,0 +1,18 @@ +from .exceptions import ( + MaxEvalError, + TargetSuccess, + CallbackSuccess, + FeasibleSuccess, +) +from .math import get_arrays_tol, exact_1d_array +from .versions import show_versions + +__all__ = [ + "MaxEvalError", + "TargetSuccess", + "CallbackSuccess", + "FeasibleSuccess", + "get_arrays_tol", + "exact_1d_array", + "show_versions", +] diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/exceptions.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..5bb1bf42be088e26ca74f6aadcb7bc3199a43c3b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/exceptions.py @@ -0,0 +1,22 @@ +class MaxEvalError(Exception): + """ + Exception raised when the maximum number of evaluations is reached. + """ + + +class TargetSuccess(Exception): + """ + Exception raised when the target value is reached. + """ + + +class CallbackSuccess(StopIteration): + """ + Exception raised when the callback function raises a ``StopIteration``. + """ + + +class FeasibleSuccess(Exception): + """ + Exception raised when a feasible point of a feasible problem is found. + """ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/math.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/math.py new file mode 100644 index 0000000000000000000000000000000000000000..2dd197a890f23107f830fe3dc2742a0ea59313f9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/math.py @@ -0,0 +1,77 @@ +import numpy as np + + +EPS = np.finfo(float).eps + + +def get_arrays_tol(*arrays): + """ + Get a relative tolerance for a set of arrays. + + Parameters + ---------- + *arrays: tuple + Set of `numpy.ndarray` to get the tolerance for. + + Returns + ------- + float + Relative tolerance for the set of arrays. + + Raises + ------ + ValueError + If no array is provided. + """ + if len(arrays) == 0: + raise ValueError("At least one array must be provided.") + size = max(array.size for array in arrays) + weight = max( + np.max(np.abs(array[np.isfinite(array)]), initial=1.0) + for array in arrays + ) + return 10.0 * EPS * max(size, 1.0) * weight + + +def exact_1d_array(x, message): + """ + Preprocess a 1-dimensional array. + + Parameters + ---------- + x : array_like + Array to be preprocessed. + message : str + Error message if `x` cannot be interpreter as a 1-dimensional array. + + Returns + ------- + `numpy.ndarray` + Preprocessed array. + """ + x = np.atleast_1d(np.squeeze(x)).astype(float) + if x.ndim != 1: + raise ValueError(message) + return x + + +def exact_2d_array(x, message): + """ + Preprocess a 2-dimensional array. + + Parameters + ---------- + x : array_like + Array to be preprocessed. + message : str + Error message if `x` cannot be interpreter as a 2-dimensional array. + + Returns + ------- + `numpy.ndarray` + Preprocessed array. + """ + x = np.atleast_2d(x).astype(float) + if x.ndim != 2: + raise ValueError(message) + return x diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/versions.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/versions.py new file mode 100644 index 0000000000000000000000000000000000000000..386d01e691adeaa9ddfc3138ed1fa2254dda8f41 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/cobyqa/utils/versions.py @@ -0,0 +1,67 @@ +import os +import platform +import sys +from importlib.metadata import PackageNotFoundError, version + + +def _get_sys_info(): + """ + Get useful system information. + + Returns + ------- + dict + Useful system information. + """ + return { + "python": sys.version.replace(os.linesep, " "), + "executable": sys.executable, + "machine": platform.platform(), + } + + +def _get_deps_info(): + """ + Get the versions of the dependencies. + + Returns + ------- + dict + Versions of the dependencies. + """ + deps = ["cobyqa", "numpy", "scipy", "setuptools", "pip"] + deps_info = {} + for module in deps: + try: + deps_info[module] = version(module) + except PackageNotFoundError: + deps_info[module] = None + return deps_info + + +def show_versions(): + """ + Display useful system and dependencies information. + + When reporting issues, please include this information. + """ + print("System settings") + print("---------------") + sys_info = _get_sys_info() + print( + "\n".join( + f"{k:>{max(map(len, sys_info.keys())) + 1}}: {v}" + for k, v in sys_info.items() + ) + ) + + print() + print("Python dependencies") + print("-------------------") + deps_info = _get_deps_info() + print( + "\n".join( + f"{k:>{max(map(len, deps_info.keys())) + 1}}: {v}" + for k, v in deps_info.items() + ) + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_array_api.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_array_api.py new file mode 100644 index 0000000000000000000000000000000000000000..25ce6837761fd9db7eefac44204f810d875540c0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_array_api.py @@ -0,0 +1,191 @@ +import numpy as np +import pytest + +from scipy.conftest import array_api_compatible +from scipy._lib._array_api import ( + _GLOBAL_CONFIG, array_namespace, _asarray, xp_copy, xp_assert_equal, is_numpy, + np_compat, xp_default_dtype +) +from scipy._lib._array_api_no_0d import xp_assert_equal as xp_assert_equal_no_0d + +skip_xp_backends = pytest.mark.skip_xp_backends + + +@pytest.mark.skipif(not _GLOBAL_CONFIG["SCIPY_ARRAY_API"], + reason="Array API test; set environment variable SCIPY_ARRAY_API=1 to run it") +class TestArrayAPI: + + def test_array_namespace(self): + x, y = np.array([0, 1, 2]), np.array([0, 1, 2]) + xp = array_namespace(x, y) + assert 'array_api_compat.numpy' in xp.__name__ + + _GLOBAL_CONFIG["SCIPY_ARRAY_API"] = False + xp = array_namespace(x, y) + assert 'array_api_compat.numpy' in xp.__name__ + _GLOBAL_CONFIG["SCIPY_ARRAY_API"] = True + + @array_api_compatible + def test_asarray(self, xp): + x, y = _asarray([0, 1, 2], xp=xp), _asarray(np.arange(3), xp=xp) + ref = xp.asarray([0, 1, 2]) + xp_assert_equal(x, ref) + xp_assert_equal(y, ref) + + @pytest.mark.filterwarnings("ignore: the matrix subclass") + def test_raises(self): + msg = "of type `numpy.ma.MaskedArray` are not supported" + with pytest.raises(TypeError, match=msg): + array_namespace(np.ma.array(1), np.array(1)) + + msg = "of type `numpy.matrix` are not supported" + with pytest.raises(TypeError, match=msg): + array_namespace(np.array(1), np.matrix(1)) + + msg = "only boolean and numerical dtypes are supported" + with pytest.raises(TypeError, match=msg): + array_namespace([object()]) + with pytest.raises(TypeError, match=msg): + array_namespace('abc') + + def test_array_likes(self): + # should be no exceptions + array_namespace([0, 1, 2]) + array_namespace(1, 2, 3) + array_namespace(1) + + @skip_xp_backends('jax.numpy', + reason="JAX arrays do not support item assignment") + @pytest.mark.usefixtures("skip_xp_backends") + @array_api_compatible + def test_copy(self, xp): + for _xp in [xp, None]: + x = xp.asarray([1, 2, 3]) + y = xp_copy(x, xp=_xp) + # with numpy we'd want to use np.shared_memory, but that's not specified + # in the array-api + x[0] = 10 + x[1] = 11 + x[2] = 12 + + assert x[0] != y[0] + assert x[1] != y[1] + assert x[2] != y[2] + assert id(x) != id(y) + + @array_api_compatible + @pytest.mark.parametrize('dtype', ['int32', 'int64', 'float32', 'float64']) + @pytest.mark.parametrize('shape', [(), (3,)]) + def test_strict_checks(self, xp, dtype, shape): + # Check that `_strict_check` behaves as expected + dtype = getattr(xp, dtype) + x = xp.broadcast_to(xp.asarray(1, dtype=dtype), shape) + x = x if shape else x[()] + y = np_compat.asarray(1)[()] + + kwarg_names = ["check_namespace", "check_dtype", "check_shape", "check_0d"] + options = dict(zip(kwarg_names, [True, False, False, False])) + if xp == np: + xp_assert_equal(x, y, **options) + else: + with pytest.raises(AssertionError, match="Namespaces do not match."): + xp_assert_equal(x, y, **options) + + options = dict(zip(kwarg_names, [False, True, False, False])) + if y.dtype.name in str(x.dtype): + xp_assert_equal(x, y, **options) + else: + with pytest.raises(AssertionError, match="dtypes do not match."): + xp_assert_equal(x, y, **options) + + options = dict(zip(kwarg_names, [False, False, True, False])) + if x.shape == y.shape: + xp_assert_equal(x, y, **options) + else: + with pytest.raises(AssertionError, match="Shapes do not match."): + xp_assert_equal(x, xp.asarray(y), **options) + + options = dict(zip(kwarg_names, [False, False, False, True])) + if is_numpy(xp) and x.shape == y.shape: + xp_assert_equal(x, y, **options) + elif is_numpy(xp): + with pytest.raises(AssertionError, match="Array-ness does not match."): + xp_assert_equal(x, y, **options) + + + @array_api_compatible + def test_check_scalar(self, xp): + if not is_numpy(xp): + pytest.skip("Scalars only exist in NumPy") + + # identity always passes + xp_assert_equal(xp.float64(0), xp.float64(0)) + xp_assert_equal(xp.asarray(0.), xp.asarray(0.)) + xp_assert_equal(xp.float64(0), xp.float64(0), check_0d=False) + xp_assert_equal(xp.asarray(0.), xp.asarray(0.), check_0d=False) + + # Check default convention: 0d-arrays are distinguished from scalars + message = "Array-ness does not match:.*" + with pytest.raises(AssertionError, match=message): + xp_assert_equal(xp.asarray(0.), xp.float64(0)) + with pytest.raises(AssertionError, match=message): + xp_assert_equal(xp.float64(0), xp.asarray(0.)) + with pytest.raises(AssertionError, match=message): + xp_assert_equal(xp.asarray(42), xp.int64(42)) + with pytest.raises(AssertionError, match=message): + xp_assert_equal(xp.int64(42), xp.asarray(42)) + + # with `check_0d=False`, scalars-vs-0d passes (if values match) + xp_assert_equal(xp.asarray(0.), xp.float64(0), check_0d=False) + xp_assert_equal(xp.float64(0), xp.asarray(0.), check_0d=False) + # also with regular python objects + xp_assert_equal(xp.asarray(0.), 0., check_0d=False) + xp_assert_equal(0., xp.asarray(0.), check_0d=False) + xp_assert_equal(xp.asarray(42), 42, check_0d=False) + xp_assert_equal(42, xp.asarray(42), check_0d=False) + + # as an alternative to `check_0d=False`, explicitly expect scalar + xp_assert_equal(xp.float64(0), xp.asarray(0.)[()]) + + + @array_api_compatible + def test_check_scalar_no_0d(self, xp): + if not is_numpy(xp): + pytest.skip("Scalars only exist in NumPy") + + # identity passes, if first argument is not 0d (or check_0d=True) + xp_assert_equal_no_0d(xp.float64(0), xp.float64(0)) + xp_assert_equal_no_0d(xp.float64(0), xp.float64(0), check_0d=True) + xp_assert_equal_no_0d(xp.asarray(0.), xp.asarray(0.), check_0d=True) + + # by default, 0d values are forbidden as the first argument + message = "Result is a NumPy 0d-array.*" + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.asarray(0.), xp.asarray(0.)) + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.asarray(0.), xp.float64(0)) + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.asarray(42), xp.int64(42)) + + # Check default convention: 0d-arrays are NOT distinguished from scalars + xp_assert_equal_no_0d(xp.float64(0), xp.asarray(0.)) + xp_assert_equal_no_0d(xp.int64(42), xp.asarray(42)) + + # opt in to 0d-check remains possible + message = "Array-ness does not match:.*" + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.asarray(0.), xp.float64(0), check_0d=True) + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.float64(0), xp.asarray(0.), check_0d=True) + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.asarray(42), xp.int64(0), check_0d=True) + with pytest.raises(AssertionError, match=message): + xp_assert_equal_no_0d(xp.int64(0), xp.asarray(42), check_0d=True) + + # scalars-vs-0d passes (if values match) also with regular python objects + xp_assert_equal_no_0d(0., xp.asarray(0.)) + xp_assert_equal_no_0d(42, xp.asarray(42)) + + @array_api_compatible + def test_default_dtype(self, xp): + assert xp_default_dtype(xp) == xp.asarray(1.).dtype diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_bunch.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_bunch.py new file mode 100644 index 0000000000000000000000000000000000000000..1fbecf081ee9d6b97ffb765af98fd000af8d120b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_bunch.py @@ -0,0 +1,162 @@ +import pytest +import pickle +from numpy.testing import assert_equal +from scipy._lib._bunch import _make_tuple_bunch + + +# `Result` is defined at the top level of the module so it can be +# used to test pickling. +Result = _make_tuple_bunch('Result', ['x', 'y', 'z'], ['w', 'beta']) + + +class TestMakeTupleBunch: + + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + # Tests with Result + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + def setup_method(self): + # Set up an instance of Result. + self.result = Result(x=1, y=2, z=3, w=99, beta=0.5) + + def test_attribute_access(self): + assert_equal(self.result.x, 1) + assert_equal(self.result.y, 2) + assert_equal(self.result.z, 3) + assert_equal(self.result.w, 99) + assert_equal(self.result.beta, 0.5) + + def test_indexing(self): + assert_equal(self.result[0], 1) + assert_equal(self.result[1], 2) + assert_equal(self.result[2], 3) + assert_equal(self.result[-1], 3) + with pytest.raises(IndexError, match='index out of range'): + self.result[3] + + def test_unpacking(self): + x0, y0, z0 = self.result + assert_equal((x0, y0, z0), (1, 2, 3)) + assert_equal(self.result, (1, 2, 3)) + + def test_slice(self): + assert_equal(self.result[1:], (2, 3)) + assert_equal(self.result[::2], (1, 3)) + assert_equal(self.result[::-1], (3, 2, 1)) + + def test_len(self): + assert_equal(len(self.result), 3) + + def test_repr(self): + s = repr(self.result) + assert_equal(s, 'Result(x=1, y=2, z=3, w=99, beta=0.5)') + + def test_hash(self): + assert_equal(hash(self.result), hash((1, 2, 3))) + + def test_pickle(self): + s = pickle.dumps(self.result) + obj = pickle.loads(s) + assert isinstance(obj, Result) + assert_equal(obj.x, self.result.x) + assert_equal(obj.y, self.result.y) + assert_equal(obj.z, self.result.z) + assert_equal(obj.w, self.result.w) + assert_equal(obj.beta, self.result.beta) + + def test_read_only_existing(self): + with pytest.raises(AttributeError, match="can't set attribute"): + self.result.x = -1 + + def test_read_only_new(self): + self.result.plate_of_shrimp = "lattice of coincidence" + assert self.result.plate_of_shrimp == "lattice of coincidence" + + def test_constructor_missing_parameter(self): + with pytest.raises(TypeError, match='missing'): + # `w` is missing. + Result(x=1, y=2, z=3, beta=0.75) + + def test_constructor_incorrect_parameter(self): + with pytest.raises(TypeError, match='unexpected'): + # `foo` is not an existing field. + Result(x=1, y=2, z=3, w=123, beta=0.75, foo=999) + + def test_module(self): + m = 'scipy._lib.tests.test_bunch' + assert_equal(Result.__module__, m) + assert_equal(self.result.__module__, m) + + def test_extra_fields_per_instance(self): + # This test exists to ensure that instances of the same class + # store their own values for the extra fields. That is, the values + # are stored per instance and not in the class. + result1 = Result(x=1, y=2, z=3, w=-1, beta=0.0) + result2 = Result(x=4, y=5, z=6, w=99, beta=1.0) + assert_equal(result1.w, -1) + assert_equal(result1.beta, 0.0) + # The rest of these checks aren't essential, but let's check + # them anyway. + assert_equal(result1[:], (1, 2, 3)) + assert_equal(result2.w, 99) + assert_equal(result2.beta, 1.0) + assert_equal(result2[:], (4, 5, 6)) + + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + # Other tests + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + def test_extra_field_names_is_optional(self): + Square = _make_tuple_bunch('Square', ['width', 'height']) + sq = Square(width=1, height=2) + assert_equal(sq.width, 1) + assert_equal(sq.height, 2) + s = repr(sq) + assert_equal(s, 'Square(width=1, height=2)') + + def test_tuple_like(self): + Tup = _make_tuple_bunch('Tup', ['a', 'b']) + tu = Tup(a=1, b=2) + assert isinstance(tu, tuple) + assert isinstance(tu + (1,), tuple) + + def test_explicit_module(self): + m = 'some.module.name' + Foo = _make_tuple_bunch('Foo', ['x'], ['a', 'b'], module=m) + foo = Foo(x=1, a=355, b=113) + assert_equal(Foo.__module__, m) + assert_equal(foo.__module__, m) + + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + # Argument validation + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + @pytest.mark.parametrize('args', [('123', ['a'], ['b']), + ('Foo', ['-3'], ['x']), + ('Foo', ['a'], ['+-*/'])]) + def test_identifiers_not_allowed(self, args): + with pytest.raises(ValueError, match='identifiers'): + _make_tuple_bunch(*args) + + @pytest.mark.parametrize('args', [('Foo', ['a', 'b', 'a'], ['x']), + ('Foo', ['a', 'b'], ['b', 'x'])]) + def test_repeated_field_names(self, args): + with pytest.raises(ValueError, match='Duplicate'): + _make_tuple_bunch(*args) + + @pytest.mark.parametrize('args', [('Foo', ['_a'], ['x']), + ('Foo', ['a'], ['_x'])]) + def test_leading_underscore_not_allowed(self, args): + with pytest.raises(ValueError, match='underscore'): + _make_tuple_bunch(*args) + + @pytest.mark.parametrize('args', [('Foo', ['def'], ['x']), + ('Foo', ['a'], ['or']), + ('and', ['a'], ['x'])]) + def test_keyword_not_allowed_in_fields(self, args): + with pytest.raises(ValueError, match='keyword'): + _make_tuple_bunch(*args) + + def test_at_least_one_field_name_required(self): + with pytest.raises(ValueError, match='at least one name'): + _make_tuple_bunch('Qwerty', [], ['a', 'b']) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_ccallback.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_ccallback.py new file mode 100644 index 0000000000000000000000000000000000000000..fe047d5f56bc2f13765038f8662c8a25b3d64042 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_ccallback.py @@ -0,0 +1,204 @@ +from numpy.testing import assert_equal, assert_ +from pytest import raises as assert_raises + +import time +import pytest +import ctypes +import threading +from scipy._lib import _ccallback_c as _test_ccallback_cython +from scipy._lib import _test_ccallback +from scipy._lib._ccallback import LowLevelCallable + +try: + import cffi + HAVE_CFFI = True +except ImportError: + HAVE_CFFI = False + + +ERROR_VALUE = 2.0 + + +def callback_python(a, user_data=None): + if a == ERROR_VALUE: + raise ValueError("bad value") + + if user_data is None: + return a + 1 + else: + return a + user_data + +def _get_cffi_func(base, signature): + if not HAVE_CFFI: + pytest.skip("cffi not installed") + + # Get function address + voidp = ctypes.cast(base, ctypes.c_void_p) + address = voidp.value + + # Create corresponding cffi handle + ffi = cffi.FFI() + func = ffi.cast(signature, address) + return func + + +def _get_ctypes_data(): + value = ctypes.c_double(2.0) + return ctypes.cast(ctypes.pointer(value), ctypes.c_voidp) + + +def _get_cffi_data(): + if not HAVE_CFFI: + pytest.skip("cffi not installed") + ffi = cffi.FFI() + return ffi.new('double *', 2.0) + + +CALLERS = { + 'simple': _test_ccallback.test_call_simple, + 'nodata': _test_ccallback.test_call_nodata, + 'nonlocal': _test_ccallback.test_call_nonlocal, + 'cython': _test_ccallback_cython.test_call_cython, +} + +# These functions have signatures known to the callers +FUNCS = { + 'python': lambda: callback_python, + 'capsule': lambda: _test_ccallback.test_get_plus1_capsule(), + 'cython': lambda: LowLevelCallable.from_cython(_test_ccallback_cython, + "plus1_cython"), + 'ctypes': lambda: _test_ccallback_cython.plus1_ctypes, + 'cffi': lambda: _get_cffi_func(_test_ccallback_cython.plus1_ctypes, + 'double (*)(double, int *, void *)'), + 'capsule_b': lambda: _test_ccallback.test_get_plus1b_capsule(), + 'cython_b': lambda: LowLevelCallable.from_cython(_test_ccallback_cython, + "plus1b_cython"), + 'ctypes_b': lambda: _test_ccallback_cython.plus1b_ctypes, + 'cffi_b': lambda: _get_cffi_func(_test_ccallback_cython.plus1b_ctypes, + 'double (*)(double, double, int *, void *)'), +} + +# These functions have signatures the callers don't know +BAD_FUNCS = { + 'capsule_bc': lambda: _test_ccallback.test_get_plus1bc_capsule(), + 'cython_bc': lambda: LowLevelCallable.from_cython(_test_ccallback_cython, + "plus1bc_cython"), + 'ctypes_bc': lambda: _test_ccallback_cython.plus1bc_ctypes, + 'cffi_bc': lambda: _get_cffi_func( + _test_ccallback_cython.plus1bc_ctypes, + 'double (*)(double, double, double, int *, void *)' + ), +} + +USER_DATAS = { + 'ctypes': _get_ctypes_data, + 'cffi': _get_cffi_data, + 'capsule': _test_ccallback.test_get_data_capsule, +} + + +def test_callbacks(): + def check(caller, func, user_data): + caller = CALLERS[caller] + func = FUNCS[func]() + user_data = USER_DATAS[user_data]() + + if func is callback_python: + def func2(x): + return func(x, 2.0) + else: + func2 = LowLevelCallable(func, user_data) + func = LowLevelCallable(func) + + # Test basic call + assert_equal(caller(func, 1.0), 2.0) + + # Test 'bad' value resulting to an error + assert_raises(ValueError, caller, func, ERROR_VALUE) + + # Test passing in user_data + assert_equal(caller(func2, 1.0), 3.0) + + for caller in sorted(CALLERS.keys()): + for func in sorted(FUNCS.keys()): + for user_data in sorted(USER_DATAS.keys()): + check(caller, func, user_data) + + +def test_bad_callbacks(): + def check(caller, func, user_data): + caller = CALLERS[caller] + user_data = USER_DATAS[user_data]() + func = BAD_FUNCS[func]() + + if func is callback_python: + def func2(x): + return func(x, 2.0) + else: + func2 = LowLevelCallable(func, user_data) + func = LowLevelCallable(func) + + # Test that basic call fails + assert_raises(ValueError, caller, LowLevelCallable(func), 1.0) + + # Test that passing in user_data also fails + assert_raises(ValueError, caller, func2, 1.0) + + # Test error message + llfunc = LowLevelCallable(func) + try: + caller(llfunc, 1.0) + except ValueError as err: + msg = str(err) + assert_(llfunc.signature in msg, msg) + assert_('double (double, double, int *, void *)' in msg, msg) + + for caller in sorted(CALLERS.keys()): + for func in sorted(BAD_FUNCS.keys()): + for user_data in sorted(USER_DATAS.keys()): + check(caller, func, user_data) + + +def test_signature_override(): + caller = _test_ccallback.test_call_simple + func = _test_ccallback.test_get_plus1_capsule() + + llcallable = LowLevelCallable(func, signature="bad signature") + assert_equal(llcallable.signature, "bad signature") + assert_raises(ValueError, caller, llcallable, 3) + + llcallable = LowLevelCallable(func, signature="double (double, int *, void *)") + assert_equal(llcallable.signature, "double (double, int *, void *)") + assert_equal(caller(llcallable, 3), 4) + + +def test_threadsafety(): + def callback(a, caller): + if a <= 0: + return 1 + else: + res = caller(lambda x: callback(x, caller), a - 1) + return 2*res + + def check(caller): + caller = CALLERS[caller] + + results = [] + + count = 10 + + def run(): + time.sleep(0.01) + r = caller(lambda x: callback(x, caller), count) + results.append(r) + + threads = [threading.Thread(target=run) for j in range(20)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert_equal(results, [2.0**count]*len(threads)) + + for caller in CALLERS.keys(): + check(caller) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_config.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_config.py new file mode 100644 index 0000000000000000000000000000000000000000..50ecfe3f33b41866fcacd8f4660c5ea67300b287 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_config.py @@ -0,0 +1,45 @@ +""" +Check the SciPy config is valid. +""" +import scipy +import pytest +from unittest.mock import patch + +pytestmark = pytest.mark.skipif( + not hasattr(scipy.__config__, "_built_with_meson"), + reason="Requires Meson builds", +) + + +class TestSciPyConfigs: + REQUIRED_CONFIG_KEYS = [ + "Compilers", + "Machine Information", + "Python Information", + ] + + @pytest.mark.thread_unsafe + @patch("scipy.__config__._check_pyyaml") + def test_pyyaml_not_found(self, mock_yaml_importer): + mock_yaml_importer.side_effect = ModuleNotFoundError() + with pytest.warns(UserWarning): + scipy.show_config() + + def test_dict_mode(self): + config = scipy.show_config(mode="dicts") + + assert isinstance(config, dict) + assert all([key in config for key in self.REQUIRED_CONFIG_KEYS]), ( + "Required key missing," + " see index of `False` with `REQUIRED_CONFIG_KEYS`" + ) + + def test_invalid_mode(self): + with pytest.raises(AttributeError): + scipy.show_config(mode="foo") + + def test_warn_to_add_tests(self): + assert len(scipy.__config__.DisplayModes) == 2, ( + "New mode detected," + " please add UT if applicable and increment this count" + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_deprecation.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_deprecation.py new file mode 100644 index 0000000000000000000000000000000000000000..57baecb4c4aea3e7664eb2fa38e4b3fe9bb61f65 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_deprecation.py @@ -0,0 +1,10 @@ +import pytest + +@pytest.mark.thread_unsafe +def test_cython_api_deprecation(): + match = ("`scipy._lib._test_deprecation_def.foo_deprecated` " + "is deprecated, use `foo` instead!\n" + "Deprecated in Scipy 42.0.0") + with pytest.warns(DeprecationWarning, match=match): + from .. import _test_deprecation_call + assert _test_deprecation_call.call() == (1, 1) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_doccer.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_doccer.py new file mode 100644 index 0000000000000000000000000000000000000000..13d14bcb6d21285cebe833e330cb03767d09567f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_doccer.py @@ -0,0 +1,143 @@ +''' Some tests for the documenting decorator and support functions ''' + +import sys +import pytest +from numpy.testing import assert_equal, suppress_warnings + +from scipy._lib import doccer + +# python -OO strips docstrings +DOCSTRINGS_STRIPPED = sys.flags.optimize > 1 + +docstring = \ +"""Docstring + %(strtest1)s + %(strtest2)s + %(strtest3)s +""" +param_doc1 = \ +"""Another test + with some indent""" + +param_doc2 = \ +"""Another test, one line""" + +param_doc3 = \ +""" Another test + with some indent""" + +doc_dict = {'strtest1':param_doc1, + 'strtest2':param_doc2, + 'strtest3':param_doc3} + +filled_docstring = \ +"""Docstring + Another test + with some indent + Another test, one line + Another test + with some indent +""" + + +def test_unindent(): + with suppress_warnings() as sup: + sup.filter(category=DeprecationWarning) + assert_equal(doccer.unindent_string(param_doc1), param_doc1) + assert_equal(doccer.unindent_string(param_doc2), param_doc2) + assert_equal(doccer.unindent_string(param_doc3), param_doc1) + + +def test_unindent_dict(): + with suppress_warnings() as sup: + sup.filter(category=DeprecationWarning) + d2 = doccer.unindent_dict(doc_dict) + assert_equal(d2['strtest1'], doc_dict['strtest1']) + assert_equal(d2['strtest2'], doc_dict['strtest2']) + assert_equal(d2['strtest3'], doc_dict['strtest1']) + + +def test_docformat(): + with suppress_warnings() as sup: + sup.filter(category=DeprecationWarning) + udd = doccer.unindent_dict(doc_dict) + formatted = doccer.docformat(docstring, udd) + assert_equal(formatted, filled_docstring) + single_doc = 'Single line doc %(strtest1)s' + formatted = doccer.docformat(single_doc, doc_dict) + # Note - initial indent of format string does not + # affect subsequent indent of inserted parameter + assert_equal(formatted, """Single line doc Another test + with some indent""") + + +@pytest.mark.skipif(DOCSTRINGS_STRIPPED, reason="docstrings stripped") +def test_decorator(): + with suppress_warnings() as sup: + sup.filter(category=DeprecationWarning) + # with unindentation of parameters + decorator = doccer.filldoc(doc_dict, True) + + @decorator + def func(): + """ Docstring + %(strtest3)s + """ + + def expected(): + """ Docstring + Another test + with some indent + """ + assert_equal(func.__doc__, expected.__doc__) + + # without unindentation of parameters + + # The docstring should be unindented for Python 3.13+ + # because of https://github.com/python/cpython/issues/81283 + decorator = doccer.filldoc(doc_dict, False if \ + sys.version_info < (3, 13) else True) + + @decorator + def func(): + """ Docstring + %(strtest3)s + """ + def expected(): + """ Docstring + Another test + with some indent + """ + assert_equal(func.__doc__, expected.__doc__) + + +@pytest.mark.skipif(DOCSTRINGS_STRIPPED, reason="docstrings stripped") +def test_inherit_docstring_from(): + + with suppress_warnings() as sup: + sup.filter(category=DeprecationWarning) + + class Foo: + def func(self): + '''Do something useful.''' + return + + def func2(self): + '''Something else.''' + + class Bar(Foo): + @doccer.inherit_docstring_from(Foo) + def func(self): + '''%(super)sABC''' + return + + @doccer.inherit_docstring_from(Foo) + def func2(self): + # No docstring. + return + + assert_equal(Bar.func.__doc__, Foo.func.__doc__ + 'ABC') + assert_equal(Bar.func2.__doc__, Foo.func2.__doc__) + bar = Bar() + assert_equal(bar.func.__doc__, Foo.func.__doc__ + 'ABC') + assert_equal(bar.func2.__doc__, Foo.func2.__doc__) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_import_cycles.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_import_cycles.py new file mode 100644 index 0000000000000000000000000000000000000000..04bb6d3a57110807cf373f2c9b850277884e3bff --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_import_cycles.py @@ -0,0 +1,18 @@ +import pytest +import sys +import subprocess + +from .test_public_api import PUBLIC_MODULES + +# Regression tests for gh-6793. +# Check that all modules are importable in a new Python process. +# This is not necessarily true if there are import cycles present. + +@pytest.mark.fail_slow(40) +@pytest.mark.slow +@pytest.mark.thread_unsafe +def test_public_modules_importable(): + pids = [subprocess.Popen([sys.executable, '-c', f'import {module}']) + for module in PUBLIC_MODULES] + for i, pid in enumerate(pids): + assert pid.wait() == 0, f'Failed to import {PUBLIC_MODULES[i]}' diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_public_api.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_public_api.py new file mode 100644 index 0000000000000000000000000000000000000000..744ac1c5ba59ac56616daf7d00195c00afb5da70 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/_lib/tests/test_public_api.py @@ -0,0 +1,469 @@ +""" +This test script is adopted from: + https://github.com/numpy/numpy/blob/main/numpy/tests/test_public_api.py +""" + +import pkgutil +import types +import importlib +import warnings +from importlib import import_module + +import pytest + +import numpy as np +import scipy + +from scipy.conftest import xp_available_backends + + +def test_dir_testing(): + """Assert that output of dir has only one "testing/tester" + attribute without duplicate""" + assert len(dir(scipy)) == len(set(dir(scipy))) + + +# Historically SciPy has not used leading underscores for private submodules +# much. This has resulted in lots of things that look like public modules +# (i.e. things that can be imported as `import scipy.somesubmodule.somefile`), +# but were never intended to be public. The PUBLIC_MODULES list contains +# modules that are either public because they were meant to be, or because they +# contain public functions/objects that aren't present in any other namespace +# for whatever reason and therefore should be treated as public. +PUBLIC_MODULES = ["scipy." + s for s in [ + "cluster", + "cluster.vq", + "cluster.hierarchy", + "constants", + "datasets", + "differentiate", + "fft", + "fftpack", + "integrate", + "interpolate", + "io", + "io.arff", + "io.matlab", + "io.wavfile", + "linalg", + "linalg.blas", + "linalg.cython_blas", + "linalg.lapack", + "linalg.cython_lapack", + "linalg.interpolative", + "ndimage", + "odr", + "optimize", + "optimize.elementwise", + "signal", + "signal.windows", + "sparse", + "sparse.linalg", + "sparse.csgraph", + "spatial", + "spatial.distance", + "spatial.transform", + "special", + "stats", + "stats.contingency", + "stats.distributions", + "stats.mstats", + "stats.qmc", + "stats.sampling" +]] + +# The PRIVATE_BUT_PRESENT_MODULES list contains modules that lacked underscores +# in their name and hence looked public, but weren't meant to be. All these +# namespace were deprecated in the 1.8.0 release - see "clear split between +# public and private API" in the 1.8.0 release notes. +# These private modules support will be removed in SciPy v2.0.0, as the +# deprecation messages emitted by each of these modules say. +PRIVATE_BUT_PRESENT_MODULES = [ + 'scipy.constants.codata', + 'scipy.constants.constants', + 'scipy.fftpack.basic', + 'scipy.fftpack.convolve', + 'scipy.fftpack.helper', + 'scipy.fftpack.pseudo_diffs', + 'scipy.fftpack.realtransforms', + 'scipy.integrate.dop', + 'scipy.integrate.lsoda', + 'scipy.integrate.odepack', + 'scipy.integrate.quadpack', + 'scipy.integrate.vode', + 'scipy.interpolate.dfitpack', + 'scipy.interpolate.fitpack', + 'scipy.interpolate.fitpack2', + 'scipy.interpolate.interpnd', + 'scipy.interpolate.interpolate', + 'scipy.interpolate.ndgriddata', + 'scipy.interpolate.polyint', + 'scipy.interpolate.rbf', + 'scipy.io.arff.arffread', + 'scipy.io.harwell_boeing', + 'scipy.io.idl', + 'scipy.io.matlab.byteordercodes', + 'scipy.io.matlab.mio', + 'scipy.io.matlab.mio4', + 'scipy.io.matlab.mio5', + 'scipy.io.matlab.mio5_params', + 'scipy.io.matlab.mio5_utils', + 'scipy.io.matlab.mio_utils', + 'scipy.io.matlab.miobase', + 'scipy.io.matlab.streams', + 'scipy.io.mmio', + 'scipy.io.netcdf', + 'scipy.linalg.basic', + 'scipy.linalg.decomp', + 'scipy.linalg.decomp_cholesky', + 'scipy.linalg.decomp_lu', + 'scipy.linalg.decomp_qr', + 'scipy.linalg.decomp_schur', + 'scipy.linalg.decomp_svd', + 'scipy.linalg.matfuncs', + 'scipy.linalg.misc', + 'scipy.linalg.special_matrices', + 'scipy.misc', + 'scipy.misc.common', + 'scipy.misc.doccer', + 'scipy.ndimage.filters', + 'scipy.ndimage.fourier', + 'scipy.ndimage.interpolation', + 'scipy.ndimage.measurements', + 'scipy.ndimage.morphology', + 'scipy.odr.models', + 'scipy.odr.odrpack', + 'scipy.optimize.cobyla', + 'scipy.optimize.cython_optimize', + 'scipy.optimize.lbfgsb', + 'scipy.optimize.linesearch', + 'scipy.optimize.minpack', + 'scipy.optimize.minpack2', + 'scipy.optimize.moduleTNC', + 'scipy.optimize.nonlin', + 'scipy.optimize.optimize', + 'scipy.optimize.slsqp', + 'scipy.optimize.tnc', + 'scipy.optimize.zeros', + 'scipy.signal.bsplines', + 'scipy.signal.filter_design', + 'scipy.signal.fir_filter_design', + 'scipy.signal.lti_conversion', + 'scipy.signal.ltisys', + 'scipy.signal.signaltools', + 'scipy.signal.spectral', + 'scipy.signal.spline', + 'scipy.signal.waveforms', + 'scipy.signal.wavelets', + 'scipy.signal.windows.windows', + 'scipy.sparse.base', + 'scipy.sparse.bsr', + 'scipy.sparse.compressed', + 'scipy.sparse.construct', + 'scipy.sparse.coo', + 'scipy.sparse.csc', + 'scipy.sparse.csr', + 'scipy.sparse.data', + 'scipy.sparse.dia', + 'scipy.sparse.dok', + 'scipy.sparse.extract', + 'scipy.sparse.lil', + 'scipy.sparse.linalg.dsolve', + 'scipy.sparse.linalg.eigen', + 'scipy.sparse.linalg.interface', + 'scipy.sparse.linalg.isolve', + 'scipy.sparse.linalg.matfuncs', + 'scipy.sparse.sparsetools', + 'scipy.sparse.spfuncs', + 'scipy.sparse.sputils', + 'scipy.spatial.ckdtree', + 'scipy.spatial.kdtree', + 'scipy.spatial.qhull', + 'scipy.spatial.transform.rotation', + 'scipy.special.add_newdocs', + 'scipy.special.basic', + 'scipy.special.cython_special', + 'scipy.special.orthogonal', + 'scipy.special.sf_error', + 'scipy.special.specfun', + 'scipy.special.spfun_stats', + 'scipy.stats.biasedurn', + 'scipy.stats.kde', + 'scipy.stats.morestats', + 'scipy.stats.mstats_basic', + 'scipy.stats.mstats_extras', + 'scipy.stats.mvn', + 'scipy.stats.stats', +] + + +def is_unexpected(name): + """Check if this needs to be considered.""" + if '._' in name or '.tests' in name or '.setup' in name: + return False + + if name in PUBLIC_MODULES: + return False + + if name in PRIVATE_BUT_PRESENT_MODULES: + return False + + return True + + +SKIP_LIST = [ + 'scipy.conftest', + 'scipy.version', + 'scipy.special.libsf_error_state' +] + + +# XXX: this test does more than it says on the tin - in using `pkgutil.walk_packages`, +# it will raise if it encounters any exceptions which are not handled by `ignore_errors` +# while attempting to import each discovered package. +# For now, `ignore_errors` only ignores what is necessary, but this could be expanded - +# for example, to all errors from private modules or git subpackages - if desired. +@pytest.mark.thread_unsafe +def test_all_modules_are_expected(): + """ + Test that we don't add anything that looks like a new public module by + accident. Check is based on filenames. + """ + + def ignore_errors(name): + # if versions of other array libraries are installed which are incompatible + # with the installed NumPy version, there can be errors on importing + # `array_api_compat`. This should only raise if SciPy is configured with + # that library as an available backend. + backends = {'cupy', 'torch', 'dask.array'} + for backend in backends: + path = f'array_api_compat.{backend}' + if path in name and backend not in xp_available_backends: + return + raise + + modnames = [] + + with np.testing.suppress_warnings() as sup: + sup.filter(DeprecationWarning,"scipy.misc") + for _, modname, _ in pkgutil.walk_packages(path=scipy.__path__, + prefix=scipy.__name__ + '.', + onerror=ignore_errors): + if is_unexpected(modname) and modname not in SKIP_LIST: + # We have a name that is new. If that's on purpose, add it to + # PUBLIC_MODULES. We don't expect to have to add anything to + # PRIVATE_BUT_PRESENT_MODULES. Use an underscore in the name! + modnames.append(modname) + + if modnames: + raise AssertionError(f'Found unexpected modules: {modnames}') + + +# Stuff that clearly shouldn't be in the API and is detected by the next test +# below +SKIP_LIST_2 = [ + 'scipy.char', + 'scipy.rec', + 'scipy.emath', + 'scipy.math', + 'scipy.random', + 'scipy.ctypeslib', + 'scipy.ma' +] + + +def test_all_modules_are_expected_2(): + """ + Method checking all objects. The pkgutil-based method in + `test_all_modules_are_expected` does not catch imports into a namespace, + only filenames. + """ + + def find_unexpected_members(mod_name): + members = [] + module = importlib.import_module(mod_name) + if hasattr(module, '__all__'): + objnames = module.__all__ + else: + objnames = dir(module) + + for objname in objnames: + if not objname.startswith('_'): + fullobjname = mod_name + '.' + objname + if isinstance(getattr(module, objname), types.ModuleType): + if is_unexpected(fullobjname) and fullobjname not in SKIP_LIST_2: + members.append(fullobjname) + + return members + with np.testing.suppress_warnings() as sup: + sup.filter(DeprecationWarning, "scipy.misc") + unexpected_members = find_unexpected_members("scipy") + + for modname in PUBLIC_MODULES: + unexpected_members.extend(find_unexpected_members(modname)) + + if unexpected_members: + raise AssertionError("Found unexpected object(s) that look like " + f"modules: {unexpected_members}") + + +def test_api_importable(): + """ + Check that all submodules listed higher up in this file can be imported + Note that if a PRIVATE_BUT_PRESENT_MODULES entry goes missing, it may + simply need to be removed from the list (deprecation may or may not be + needed - apply common sense). + """ + def check_importable(module_name): + try: + importlib.import_module(module_name) + except (ImportError, AttributeError): + return False + + return True + + module_names = [] + for module_name in PUBLIC_MODULES: + if not check_importable(module_name): + module_names.append(module_name) + + if module_names: + raise AssertionError("Modules in the public API that cannot be " + f"imported: {module_names}") + + with warnings.catch_warnings(record=True): + warnings.filterwarnings('always', category=DeprecationWarning) + warnings.filterwarnings('always', category=ImportWarning) + for module_name in PRIVATE_BUT_PRESENT_MODULES: + if not check_importable(module_name): + module_names.append(module_name) + + if module_names: + raise AssertionError("Modules that are not really public but looked " + "public and can not be imported: " + f"{module_names}") + + +@pytest.mark.thread_unsafe +@pytest.mark.parametrize(("module_name", "correct_module"), + [('scipy.constants.codata', None), + ('scipy.constants.constants', None), + ('scipy.fftpack.basic', None), + ('scipy.fftpack.helper', None), + ('scipy.fftpack.pseudo_diffs', None), + ('scipy.fftpack.realtransforms', None), + ('scipy.integrate.dop', None), + ('scipy.integrate.lsoda', None), + ('scipy.integrate.odepack', None), + ('scipy.integrate.quadpack', None), + ('scipy.integrate.vode', None), + ('scipy.interpolate.fitpack', None), + ('scipy.interpolate.fitpack2', None), + ('scipy.interpolate.interpolate', None), + ('scipy.interpolate.ndgriddata', None), + ('scipy.interpolate.polyint', None), + ('scipy.interpolate.rbf', None), + ('scipy.io.harwell_boeing', None), + ('scipy.io.idl', None), + ('scipy.io.mmio', None), + ('scipy.io.netcdf', None), + ('scipy.io.arff.arffread', 'arff'), + ('scipy.io.matlab.byteordercodes', 'matlab'), + ('scipy.io.matlab.mio_utils', 'matlab'), + ('scipy.io.matlab.mio', 'matlab'), + ('scipy.io.matlab.mio4', 'matlab'), + ('scipy.io.matlab.mio5_params', 'matlab'), + ('scipy.io.matlab.mio5_utils', 'matlab'), + ('scipy.io.matlab.mio5', 'matlab'), + ('scipy.io.matlab.miobase', 'matlab'), + ('scipy.io.matlab.streams', 'matlab'), + ('scipy.linalg.basic', None), + ('scipy.linalg.decomp', None), + ('scipy.linalg.decomp_cholesky', None), + ('scipy.linalg.decomp_lu', None), + ('scipy.linalg.decomp_qr', None), + ('scipy.linalg.decomp_schur', None), + ('scipy.linalg.decomp_svd', None), + ('scipy.linalg.matfuncs', None), + ('scipy.linalg.misc', None), + ('scipy.linalg.special_matrices', None), + ('scipy.ndimage.filters', None), + ('scipy.ndimage.fourier', None), + ('scipy.ndimage.interpolation', None), + ('scipy.ndimage.measurements', None), + ('scipy.ndimage.morphology', None), + ('scipy.odr.models', None), + ('scipy.odr.odrpack', None), + ('scipy.optimize.cobyla', None), + ('scipy.optimize.lbfgsb', None), + ('scipy.optimize.linesearch', None), + ('scipy.optimize.minpack', None), + ('scipy.optimize.minpack2', None), + ('scipy.optimize.moduleTNC', None), + ('scipy.optimize.nonlin', None), + ('scipy.optimize.optimize', None), + ('scipy.optimize.slsqp', None), + ('scipy.optimize.tnc', None), + ('scipy.optimize.zeros', None), + ('scipy.signal.bsplines', None), + ('scipy.signal.filter_design', None), + ('scipy.signal.fir_filter_design', None), + ('scipy.signal.lti_conversion', None), + ('scipy.signal.ltisys', None), + ('scipy.signal.signaltools', None), + ('scipy.signal.spectral', None), + ('scipy.signal.waveforms', None), + ('scipy.signal.wavelets', None), + ('scipy.signal.windows.windows', 'windows'), + ('scipy.sparse.lil', None), + ('scipy.sparse.linalg.dsolve', 'linalg'), + ('scipy.sparse.linalg.eigen', 'linalg'), + ('scipy.sparse.linalg.interface', 'linalg'), + ('scipy.sparse.linalg.isolve', 'linalg'), + ('scipy.sparse.linalg.matfuncs', 'linalg'), + ('scipy.sparse.sparsetools', None), + ('scipy.sparse.spfuncs', None), + ('scipy.sparse.sputils', None), + ('scipy.spatial.ckdtree', None), + ('scipy.spatial.kdtree', None), + ('scipy.spatial.qhull', None), + ('scipy.spatial.transform.rotation', 'transform'), + ('scipy.special.add_newdocs', None), + ('scipy.special.basic', None), + ('scipy.special.orthogonal', None), + ('scipy.special.sf_error', None), + ('scipy.special.specfun', None), + ('scipy.special.spfun_stats', None), + ('scipy.stats.biasedurn', None), + ('scipy.stats.kde', None), + ('scipy.stats.morestats', None), + ('scipy.stats.mstats_basic', 'mstats'), + ('scipy.stats.mstats_extras', 'mstats'), + ('scipy.stats.mvn', None), + ('scipy.stats.stats', None)]) +def test_private_but_present_deprecation(module_name, correct_module): + # gh-18279, gh-17572, gh-17771 noted that deprecation warnings + # for imports from private modules + # were misleading. Check that this is resolved. + module = import_module(module_name) + if correct_module is None: + import_name = f'scipy.{module_name.split(".")[1]}' + else: + import_name = f'scipy.{module_name.split(".")[1]}.{correct_module}' + + correct_import = import_module(import_name) + + # Attributes that were formerly in `module_name` can still be imported from + # `module_name`, albeit with a deprecation warning. + for attr_name in module.__all__: + # ensure attribute is present where the warning is pointing + assert getattr(correct_import, attr_name, None) is not None + message = f"Please import `{attr_name}` from the `{import_name}`..." + with pytest.deprecated_call(match=message): + getattr(module, attr_name) + + # Attributes that were not in `module_name` get an error notifying the user + # that the attribute is not in `module_name` and that `module_name` is deprecated. + message = f"`{module_name}` is deprecated..." + with pytest.raises(AttributeError, match=message): + getattr(module, "ekki") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..93a68daf892e865a9d949a77d8857b1d1b68c1f3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/__init__.py @@ -0,0 +1,31 @@ +""" +========================================= +Clustering package (:mod:`scipy.cluster`) +========================================= + +.. currentmodule:: scipy.cluster + +Clustering algorithms are useful in information theory, target detection, +communications, compression, and other areas. The `vq` module only +supports vector quantization and the k-means algorithms. + +The `hierarchy` module provides functions for hierarchical and +agglomerative clustering. Its features include generating hierarchical +clusters from distance matrices, +calculating statistics on clusters, cutting linkages +to generate flat clusters, and visualizing clusters with dendrograms. + +.. toctree:: + :maxdepth: 1 + + cluster.vq + cluster.hierarchy + +""" +__all__ = ['vq', 'hierarchy'] + +from . import vq, hierarchy + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_hierarchy.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_hierarchy.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..6deadf38b55f7320fc1a323910b3c993c96dfe93 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_hierarchy.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_optimal_leaf_ordering.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_optimal_leaf_ordering.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..f4e4b7a71a356725f01025e183edb26bab4d06fe Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_optimal_leaf_ordering.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_vq.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_vq.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..1e5a457ef9f678deb08b8d76ed3dd6a5dc713d84 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/_vq.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/hierarchy.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/hierarchy.py new file mode 100644 index 0000000000000000000000000000000000000000..9de6befccdb002c4280789da09c51aa55da7b07f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/hierarchy.py @@ -0,0 +1,4178 @@ +""" +Hierarchical clustering (:mod:`scipy.cluster.hierarchy`) +======================================================== + +.. currentmodule:: scipy.cluster.hierarchy + +These functions cut hierarchical clusterings into flat clusterings +or find the roots of the forest formed by a cut by providing the flat +cluster ids of each observation. + +.. autosummary:: + :toctree: generated/ + + fcluster + fclusterdata + leaders + +These are routines for agglomerative clustering. + +.. autosummary:: + :toctree: generated/ + + linkage + single + complete + average + weighted + centroid + median + ward + +These routines compute statistics on hierarchies. + +.. autosummary:: + :toctree: generated/ + + cophenet + from_mlab_linkage + inconsistent + maxinconsts + maxdists + maxRstat + to_mlab_linkage + +Routines for visualizing flat clusters. + +.. autosummary:: + :toctree: generated/ + + dendrogram + +These are data structures and routines for representing hierarchies as +tree objects. + +.. autosummary:: + :toctree: generated/ + + ClusterNode + leaves_list + to_tree + cut_tree + optimal_leaf_ordering + +These are predicates for checking the validity of linkage and +inconsistency matrices as well as for checking isomorphism of two +flat cluster assignments. + +.. autosummary:: + :toctree: generated/ + + is_valid_im + is_valid_linkage + is_isomorphic + is_monotonic + correspond + num_obs_linkage + +Utility routines for plotting: + +.. autosummary:: + :toctree: generated/ + + set_link_color_palette + +Utility classes: + +.. autosummary:: + :toctree: generated/ + + DisjointSet -- data structure for incremental connectivity queries + +""" +# Copyright (C) Damian Eads, 2007-2008. New BSD License. + +# hierarchy.py (derived from cluster.py, http://scipy-cluster.googlecode.com) +# +# Author: Damian Eads +# Date: September 22, 2007 +# +# Copyright (c) 2007, 2008, Damian Eads +# +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# - Redistributions of source code must retain the above +# copyright notice, this list of conditions and the +# following disclaimer. +# - Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer +# in the documentation and/or other materials provided with the +# distribution. +# - Neither the name of the author nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import warnings +import bisect +from collections import deque + +import numpy as np +from . import _hierarchy, _optimal_leaf_ordering +import scipy.spatial.distance as distance +from scipy._lib._array_api import array_namespace, _asarray, xp_copy, is_jax +from scipy._lib._disjoint_set import DisjointSet + + +_LINKAGE_METHODS = {'single': 0, 'complete': 1, 'average': 2, 'centroid': 3, + 'median': 4, 'ward': 5, 'weighted': 6} +_EUCLIDEAN_METHODS = ('centroid', 'median', 'ward') + +__all__ = ['ClusterNode', 'DisjointSet', 'average', 'centroid', 'complete', + 'cophenet', 'correspond', 'cut_tree', 'dendrogram', 'fcluster', + 'fclusterdata', 'from_mlab_linkage', 'inconsistent', + 'is_isomorphic', 'is_monotonic', 'is_valid_im', 'is_valid_linkage', + 'leaders', 'leaves_list', 'linkage', 'maxRstat', 'maxdists', + 'maxinconsts', 'median', 'num_obs_linkage', 'optimal_leaf_ordering', + 'set_link_color_palette', 'single', 'to_mlab_linkage', 'to_tree', + 'ward', 'weighted'] + + +class ClusterWarning(UserWarning): + pass + + +def _warning(s): + warnings.warn(f'scipy.cluster: {s}', ClusterWarning, stacklevel=3) + + +def int_floor(arr, xp): + # array_api_strict is strict about not allowing `int()` on a float array. + # That's typically not needed, here it is - so explicitly convert + return int(xp.astype(xp.asarray(arr), xp.int64)) + + +def single(y): + """ + Perform single/min/nearest linkage on the condensed distance matrix ``y``. + + Parameters + ---------- + y : ndarray + The upper triangular of the distance matrix. The result of + ``pdist`` is returned in this form. + + Returns + ------- + Z : ndarray + The linkage matrix. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import single, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = single(y) + >>> Z + array([[ 0., 1., 1., 2.], + [ 2., 12., 1., 3.], + [ 3., 4., 1., 2.], + [ 5., 14., 1., 3.], + [ 6., 7., 1., 2.], + [ 8., 16., 1., 3.], + [ 9., 10., 1., 2.], + [11., 18., 1., 3.], + [13., 15., 2., 6.], + [17., 20., 2., 9.], + [19., 21., 2., 12.]]) + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 7, 8, 9, 10, 11, 12, 4, 5, 6, 1, 2, 3], dtype=int32) + >>> fcluster(Z, 1, criterion='distance') + array([3, 3, 3, 4, 4, 4, 2, 2, 2, 1, 1, 1], dtype=int32) + >>> fcluster(Z, 2, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + """ + return linkage(y, method='single', metric='euclidean') + + +def complete(y): + """ + Perform complete/max/farthest point linkage on a condensed distance matrix. + + Parameters + ---------- + y : ndarray + The upper triangular of the distance matrix. The result of + ``pdist`` is returned in this form. + + Returns + ------- + Z : ndarray + A linkage matrix containing the hierarchical clustering. See + the `linkage` function documentation for more information + on its structure. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import complete, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = complete(y) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.41421356, 3. ], + [ 5. , 13. , 1.41421356, 3. ], + [ 8. , 14. , 1.41421356, 3. ], + [11. , 15. , 1.41421356, 3. ], + [16. , 17. , 4.12310563, 6. ], + [18. , 19. , 4.12310563, 6. ], + [20. , 21. , 5.65685425, 12. ]]) + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=int32) + >>> fcluster(Z, 1.5, criterion='distance') + array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], dtype=int32) + >>> fcluster(Z, 4.5, criterion='distance') + array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2], dtype=int32) + >>> fcluster(Z, 6, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + """ + return linkage(y, method='complete', metric='euclidean') + + +def average(y): + """ + Perform average/UPGMA linkage on a condensed distance matrix. + + Parameters + ---------- + y : ndarray + The upper triangular of the distance matrix. The result of + ``pdist`` is returned in this form. + + Returns + ------- + Z : ndarray + A linkage matrix containing the hierarchical clustering. See + `linkage` for more information on its structure. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import average, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = average(y) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.20710678, 3. ], + [ 5. , 13. , 1.20710678, 3. ], + [ 8. , 14. , 1.20710678, 3. ], + [11. , 15. , 1.20710678, 3. ], + [16. , 17. , 3.39675184, 6. ], + [18. , 19. , 3.39675184, 6. ], + [20. , 21. , 4.09206523, 12. ]]) + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=int32) + >>> fcluster(Z, 1.5, criterion='distance') + array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], dtype=int32) + >>> fcluster(Z, 4, criterion='distance') + array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2], dtype=int32) + >>> fcluster(Z, 6, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + + """ + return linkage(y, method='average', metric='euclidean') + + +def weighted(y): + """ + Perform weighted/WPGMA linkage on the condensed distance matrix. + + See `linkage` for more information on the return + structure and algorithm. + + Parameters + ---------- + y : ndarray + The upper triangular of the distance matrix. The result of + ``pdist`` is returned in this form. + + Returns + ------- + Z : ndarray + A linkage matrix containing the hierarchical clustering. See + `linkage` for more information on its structure. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import weighted, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = weighted(y) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 9. , 11. , 1. , 2. ], + [ 2. , 12. , 1.20710678, 3. ], + [ 8. , 13. , 1.20710678, 3. ], + [ 5. , 14. , 1.20710678, 3. ], + [10. , 15. , 1.20710678, 3. ], + [18. , 19. , 3.05595762, 6. ], + [16. , 17. , 3.32379407, 6. ], + [20. , 21. , 4.06357713, 12. ]]) + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 7, 8, 9, 1, 2, 3, 10, 11, 12, 4, 6, 5], dtype=int32) + >>> fcluster(Z, 1.5, criterion='distance') + array([3, 3, 3, 1, 1, 1, 4, 4, 4, 2, 2, 2], dtype=int32) + >>> fcluster(Z, 4, criterion='distance') + array([2, 2, 2, 1, 1, 1, 2, 2, 2, 1, 1, 1], dtype=int32) + >>> fcluster(Z, 6, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + + """ + return linkage(y, method='weighted', metric='euclidean') + + +def centroid(y): + """ + Perform centroid/UPGMC linkage. + + See `linkage` for more information on the input matrix, + return structure, and algorithm. + + The following are common calling conventions: + + 1. ``Z = centroid(y)`` + + Performs centroid/UPGMC linkage on the condensed distance + matrix ``y``. + + 2. ``Z = centroid(X)`` + + Performs centroid/UPGMC linkage on the observation matrix ``X`` + using Euclidean distance as the distance metric. + + Parameters + ---------- + y : ndarray + A condensed distance matrix. A condensed + distance matrix is a flat array containing the upper + triangular of the distance matrix. This is the form that + ``pdist`` returns. Alternatively, a collection of + m observation vectors in n dimensions may be passed as + an m by n array. + + Returns + ------- + Z : ndarray + A linkage matrix containing the hierarchical clustering. See + the `linkage` function documentation for more information + on its structure. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import centroid, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = centroid(y) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 2. , 12. , 1.11803399, 3. ], + [ 5. , 13. , 1.11803399, 3. ], + [ 8. , 15. , 1.11803399, 3. ], + [11. , 14. , 1.11803399, 3. ], + [18. , 19. , 3.33333333, 6. ], + [16. , 17. , 3.33333333, 6. ], + [20. , 21. , 3.33333333, 12. ]]) # may vary + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6], dtype=int32) # may vary + >>> fcluster(Z, 1.1, criterion='distance') + array([5, 5, 6, 7, 7, 8, 1, 1, 2, 3, 3, 4], dtype=int32) # may vary + >>> fcluster(Z, 2, criterion='distance') + array([3, 3, 3, 4, 4, 4, 1, 1, 1, 2, 2, 2], dtype=int32) # may vary + >>> fcluster(Z, 4, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + + """ + return linkage(y, method='centroid', metric='euclidean') + + +def median(y): + """ + Perform median/WPGMC linkage. + + See `linkage` for more information on the return structure + and algorithm. + + The following are common calling conventions: + + 1. ``Z = median(y)`` + + Performs median/WPGMC linkage on the condensed distance matrix + ``y``. See ``linkage`` for more information on the return + structure and algorithm. + + 2. ``Z = median(X)`` + + Performs median/WPGMC linkage on the observation matrix ``X`` + using Euclidean distance as the distance metric. See `linkage` + for more information on the return structure and algorithm. + + Parameters + ---------- + y : ndarray + A condensed distance matrix. A condensed + distance matrix is a flat array containing the upper + triangular of the distance matrix. This is the form that + ``pdist`` returns. Alternatively, a collection of + m observation vectors in n dimensions may be passed as + an m by n array. + + Returns + ------- + Z : ndarray + The hierarchical clustering encoded as a linkage matrix. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import median, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = median(y) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 2. , 12. , 1.11803399, 3. ], + [ 5. , 13. , 1.11803399, 3. ], + [ 8. , 15. , 1.11803399, 3. ], + [11. , 14. , 1.11803399, 3. ], + [18. , 19. , 3. , 6. ], + [16. , 17. , 3.5 , 6. ], + [20. , 21. , 3.25 , 12. ]]) + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6], dtype=int32) + >>> fcluster(Z, 1.1, criterion='distance') + array([5, 5, 6, 7, 7, 8, 1, 1, 2, 3, 3, 4], dtype=int32) + >>> fcluster(Z, 2, criterion='distance') + array([3, 3, 3, 4, 4, 4, 1, 1, 1, 2, 2, 2], dtype=int32) + >>> fcluster(Z, 4, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + + """ + return linkage(y, method='median', metric='euclidean') + + +def ward(y): + """ + Perform Ward's linkage on a condensed distance matrix. + + See `linkage` for more information on the return structure + and algorithm. + + The following are common calling conventions: + + 1. ``Z = ward(y)`` + Performs Ward's linkage on the condensed distance matrix ``y``. + + 2. ``Z = ward(X)`` + Performs Ward's linkage on the observation matrix ``X`` using + Euclidean distance as the distance metric. + + Parameters + ---------- + y : ndarray + A condensed distance matrix. A condensed + distance matrix is a flat array containing the upper + triangular of the distance matrix. This is the form that + ``pdist`` returns. Alternatively, a collection of + m observation vectors in n dimensions may be passed as + an m by n array. + + Returns + ------- + Z : ndarray + The hierarchical clustering encoded as a linkage matrix. See + `linkage` for more information on the return structure and + algorithm. + + See Also + -------- + linkage : for advanced creation of hierarchical clusterings. + scipy.spatial.distance.pdist : pairwise distance metrics + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, fcluster + >>> from scipy.spatial.distance import pdist + + First, we need a toy dataset to play with:: + + x x x x + x x + + x x + x x x x + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + Then, we get a condensed distance matrix from this dataset: + + >>> y = pdist(X) + + Finally, we can perform the clustering: + + >>> Z = ward(y) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + + The linkage matrix ``Z`` represents a dendrogram - see + `scipy.cluster.hierarchy.linkage` for a detailed explanation of its + contents. + + We can use `scipy.cluster.hierarchy.fcluster` to see to which cluster + each initial point would belong given a distance threshold: + + >>> fcluster(Z, 0.9, criterion='distance') + array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=int32) + >>> fcluster(Z, 1.1, criterion='distance') + array([1, 1, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8], dtype=int32) + >>> fcluster(Z, 3, criterion='distance') + array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], dtype=int32) + >>> fcluster(Z, 9, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + Also, `scipy.cluster.hierarchy.dendrogram` can be used to generate a + plot of the dendrogram. + + """ + return linkage(y, method='ward', metric='euclidean') + + +def linkage(y, method='single', metric='euclidean', optimal_ordering=False): + """ + Perform hierarchical/agglomerative clustering. + + The input y may be either a 1-D condensed distance matrix + or a 2-D array of observation vectors. + + If y is a 1-D condensed distance matrix, + then y must be a :math:`\\binom{n}{2}` sized + vector, where n is the number of original observations paired + in the distance matrix. The behavior of this function is very + similar to the MATLAB linkage function. + + A :math:`(n-1)` by 4 matrix ``Z`` is returned. At the + :math:`i`-th iteration, clusters with indices ``Z[i, 0]`` and + ``Z[i, 1]`` are combined to form cluster :math:`n + i`. A + cluster with an index less than :math:`n` corresponds to one of + the :math:`n` original observations. The distance between + clusters ``Z[i, 0]`` and ``Z[i, 1]`` is given by ``Z[i, 2]``. The + fourth value ``Z[i, 3]`` represents the number of original + observations in the newly formed cluster. + + The following linkage methods are used to compute the distance + :math:`d(s, t)` between two clusters :math:`s` and + :math:`t`. The algorithm begins with a forest of clusters that + have yet to be used in the hierarchy being formed. When two + clusters :math:`s` and :math:`t` from this forest are combined + into a single cluster :math:`u`, :math:`s` and :math:`t` are + removed from the forest, and :math:`u` is added to the + forest. When only one cluster remains in the forest, the algorithm + stops, and this cluster becomes the root. + + A distance matrix is maintained at each iteration. The ``d[i,j]`` + entry corresponds to the distance between cluster :math:`i` and + :math:`j` in the original forest. + + At each iteration, the algorithm must update the distance matrix + to reflect the distance of the newly formed cluster u with the + remaining clusters in the forest. + + Suppose there are :math:`|u|` original observations + :math:`u[0], \\ldots, u[|u|-1]` in cluster :math:`u` and + :math:`|v|` original objects :math:`v[0], \\ldots, v[|v|-1]` in + cluster :math:`v`. Recall, :math:`s` and :math:`t` are + combined to form cluster :math:`u`. Let :math:`v` be any + remaining cluster in the forest that is not :math:`u`. + + The following are methods for calculating the distance between the + newly formed cluster :math:`u` and each :math:`v`. + + * method='single' assigns + + .. math:: + d(u,v) = \\min(dist(u[i],v[j])) + + for all points :math:`i` in cluster :math:`u` and + :math:`j` in cluster :math:`v`. This is also known as the + Nearest Point Algorithm. + + * method='complete' assigns + + .. math:: + d(u, v) = \\max(dist(u[i],v[j])) + + for all points :math:`i` in cluster u and :math:`j` in + cluster :math:`v`. This is also known by the Farthest Point + Algorithm or Voor Hees Algorithm. + + * method='average' assigns + + .. math:: + d(u,v) = \\sum_{ij} \\frac{d(u[i], v[j])} + {(|u|*|v|)} + + for all points :math:`i` and :math:`j` where :math:`|u|` + and :math:`|v|` are the cardinalities of clusters :math:`u` + and :math:`v`, respectively. This is also called the UPGMA + algorithm. + + * method='weighted' assigns + + .. math:: + d(u,v) = (dist(s,v) + dist(t,v))/2 + + where cluster u was formed with cluster s and t and v + is a remaining cluster in the forest (also called WPGMA). + + * method='centroid' assigns + + .. math:: + dist(s,t) = ||c_s-c_t||_2 + + where :math:`c_s` and :math:`c_t` are the centroids of + clusters :math:`s` and :math:`t`, respectively. When two + clusters :math:`s` and :math:`t` are combined into a new + cluster :math:`u`, the new centroid is computed over all the + original objects in clusters :math:`s` and :math:`t`. The + distance then becomes the Euclidean distance between the + centroid of :math:`u` and the centroid of a remaining cluster + :math:`v` in the forest. This is also known as the UPGMC + algorithm. + + * method='median' assigns :math:`d(s,t)` like the ``centroid`` + method. When two clusters :math:`s` and :math:`t` are combined + into a new cluster :math:`u`, the average of centroids s and t + give the new centroid :math:`u`. This is also known as the + WPGMC algorithm. + + * method='ward' uses the Ward variance minimization algorithm. + The new entry :math:`d(u,v)` is computed as follows, + + .. math:: + + d(u,v) = \\sqrt{\\frac{|v|+|s|} + {T}d(v,s)^2 + + \\frac{|v|+|t|} + {T}d(v,t)^2 + - \\frac{|v|} + {T}d(s,t)^2} + + where :math:`u` is the newly joined cluster consisting of + clusters :math:`s` and :math:`t`, :math:`v` is an unused + cluster in the forest, :math:`T=|v|+|s|+|t|`, and + :math:`|*|` is the cardinality of its argument. This is also + known as the incremental algorithm. + + Warning: When the minimum distance pair in the forest is chosen, there + may be two or more pairs with the same minimum distance. This + implementation may choose a different minimum than the MATLAB + version. + + Parameters + ---------- + y : ndarray + A condensed distance matrix. A condensed distance matrix + is a flat array containing the upper triangular of the distance matrix. + This is the form that ``pdist`` returns. Alternatively, a collection of + :math:`m` observation vectors in :math:`n` dimensions may be passed as + an :math:`m` by :math:`n` array. All elements of the condensed distance + matrix must be finite, i.e., no NaNs or infs. + method : str, optional + The linkage algorithm to use. See the ``Linkage Methods`` section below + for full descriptions. + metric : str or function, optional + The distance metric to use in the case that y is a collection of + observation vectors; ignored otherwise. See the ``pdist`` + function for a list of valid distance metrics. A custom distance + function can also be used. + optimal_ordering : bool, optional + If True, the linkage matrix will be reordered so that the distance + between successive leaves is minimal. This results in a more intuitive + tree structure when the data are visualized. defaults to False, because + this algorithm can be slow, particularly on large datasets [2]_. See + also the `optimal_leaf_ordering` function. + + .. versionadded:: 1.0.0 + + Returns + ------- + Z : ndarray + The hierarchical clustering encoded as a linkage matrix. + + Notes + ----- + 1. For method 'single', an optimized algorithm based on minimum spanning + tree is implemented. It has time complexity :math:`O(n^2)`. + For methods 'complete', 'average', 'weighted' and 'ward', an algorithm + called nearest-neighbors chain is implemented. It also has time + complexity :math:`O(n^2)`. + For other methods, a naive algorithm is implemented with :math:`O(n^3)` + time complexity. + All algorithms use :math:`O(n^2)` memory. + Refer to [1]_ for details about the algorithms. + 2. Methods 'centroid', 'median', and 'ward' are correctly defined only if + Euclidean pairwise metric is used. If `y` is passed as precomputed + pairwise distances, then it is the user's responsibility to assure that + these distances are in fact Euclidean, otherwise the produced result + will be incorrect. + + See Also + -------- + scipy.spatial.distance.pdist : pairwise distance metrics + + References + ---------- + .. [1] Daniel Mullner, "Modern hierarchical, agglomerative clustering + algorithms", :arXiv:`1109.2378v1`. + .. [2] Ziv Bar-Joseph, David K. Gifford, Tommi S. Jaakkola, "Fast optimal + leaf ordering for hierarchical clustering", 2001. Bioinformatics + :doi:`10.1093/bioinformatics/17.suppl_1.S22` + + Examples + -------- + >>> from scipy.cluster.hierarchy import dendrogram, linkage + >>> from matplotlib import pyplot as plt + >>> X = [[i] for i in [2, 8, 0, 4, 1, 9, 9, 0]] + + >>> Z = linkage(X, 'ward') + >>> fig = plt.figure(figsize=(25, 10)) + >>> dn = dendrogram(Z) + + >>> Z = linkage(X, 'single') + >>> fig = plt.figure(figsize=(25, 10)) + >>> dn = dendrogram(Z) + >>> plt.show() + """ + xp = array_namespace(y) + y = _asarray(y, order='C', dtype=xp.float64, xp=xp) + + if method not in _LINKAGE_METHODS: + raise ValueError(f"Invalid method: {method}") + + if method in _EUCLIDEAN_METHODS and metric != 'euclidean' and y.ndim == 2: + msg = f"`method={method}` requires the distance metric to be Euclidean" + raise ValueError(msg) + + if y.ndim == 1: + distance.is_valid_y(y, throw=True, name='y') + elif y.ndim == 2: + if (y.shape[0] == y.shape[1] and np.allclose(np.diag(y), 0) and + xp.all(y >= 0) and np.allclose(y, y.T)): + warnings.warn('The symmetric non-negative hollow observation ' + 'matrix looks suspiciously like an uncondensed ' + 'distance matrix', + ClusterWarning, stacklevel=2) + y = distance.pdist(y, metric) + y = xp.asarray(y) + else: + raise ValueError("`y` must be 1 or 2 dimensional.") + + if not xp.all(xp.isfinite(y)): + raise ValueError("The condensed distance matrix must contain only " + "finite values.") + + n = int(distance.num_obs_y(y)) + method_code = _LINKAGE_METHODS[method] + + y = np.asarray(y) + if method == 'single': + result = _hierarchy.mst_single_linkage(y, n) + elif method in ['complete', 'average', 'weighted', 'ward']: + result = _hierarchy.nn_chain(y, n, method_code) + else: + result = _hierarchy.fast_linkage(y, n, method_code) + result = xp.asarray(result) + + if optimal_ordering: + y = xp.asarray(y) + return optimal_leaf_ordering(result, y) + else: + return result + + +class ClusterNode: + """ + A tree node class for representing a cluster. + + Leaf nodes correspond to original observations, while non-leaf nodes + correspond to non-singleton clusters. + + The `to_tree` function converts a matrix returned by the linkage + function into an easy-to-use tree representation. + + All parameter names are also attributes. + + Parameters + ---------- + id : int + The node id. + left : ClusterNode instance, optional + The left child tree node. + right : ClusterNode instance, optional + The right child tree node. + dist : float, optional + Distance for this cluster in the linkage matrix. + count : int, optional + The number of samples in this cluster. + + See Also + -------- + to_tree : for converting a linkage matrix ``Z`` into a tree object. + + """ + + def __init__(self, id, left=None, right=None, dist=0.0, count=1): + if id < 0: + raise ValueError('The id must be non-negative.') + if dist < 0: + raise ValueError('The distance must be non-negative.') + if (left is None and right is not None) or \ + (left is not None and right is None): + raise ValueError('Only full or proper binary trees are permitted.' + ' This node has one child.') + if count < 1: + raise ValueError('A cluster must contain at least one original ' + 'observation.') + self.id = id + self.left = left + self.right = right + self.dist = dist + if self.left is None: + self.count = count + else: + self.count = left.count + right.count + + def __lt__(self, node): + if not isinstance(node, ClusterNode): + raise ValueError("Can't compare ClusterNode " + f"to type {type(node)}") + return self.dist < node.dist + + def __gt__(self, node): + if not isinstance(node, ClusterNode): + raise ValueError("Can't compare ClusterNode " + f"to type {type(node)}") + return self.dist > node.dist + + def __eq__(self, node): + if not isinstance(node, ClusterNode): + raise ValueError("Can't compare ClusterNode " + f"to type {type(node)}") + return self.dist == node.dist + + def get_id(self): + """ + The identifier of the target node. + + For ``0 <= i < n``, `i` corresponds to original observation i. + For ``n <= i < 2n-1``, `i` corresponds to non-singleton cluster formed + at iteration ``i-n``. + + Returns + ------- + id : int + The identifier of the target node. + + """ + return self.id + + def get_count(self): + """ + The number of leaf nodes (original observations) belonging to + the cluster node nd. If the target node is a leaf, 1 is + returned. + + Returns + ------- + get_count : int + The number of leaf nodes below the target node. + + """ + return self.count + + def get_left(self): + """ + Return a reference to the left child tree object. + + Returns + ------- + left : ClusterNode + The left child of the target node. If the node is a leaf, + None is returned. + + """ + return self.left + + def get_right(self): + """ + Return a reference to the right child tree object. + + Returns + ------- + right : ClusterNode + The left child of the target node. If the node is a leaf, + None is returned. + + """ + return self.right + + def is_leaf(self): + """ + Return True if the target node is a leaf. + + Returns + ------- + leafness : bool + True if the target node is a leaf node. + + """ + return self.left is None + + def pre_order(self, func=(lambda x: x.id)): + """ + Perform pre-order traversal without recursive function calls. + + When a leaf node is first encountered, ``func`` is called with + the leaf node as its argument, and its result is appended to + the list. + + For example, the statement:: + + ids = root.pre_order(lambda x: x.id) + + returns a list of the node ids corresponding to the leaf nodes + of the tree as they appear from left to right. + + Parameters + ---------- + func : function + Applied to each leaf ClusterNode object in the pre-order traversal. + Given the ``i``-th leaf node in the pre-order traversal ``n[i]``, + the result of ``func(n[i])`` is stored in ``L[i]``. If not + provided, the index of the original observation to which the node + corresponds is used. + + Returns + ------- + L : list + The pre-order traversal. + + """ + # Do a preorder traversal, caching the result. To avoid having to do + # recursion, we'll store the previous index we've visited in a vector. + n = self.count + + curNode = [None] * (2 * n) + lvisited = set() + rvisited = set() + curNode[0] = self + k = 0 + preorder = [] + while k >= 0: + nd = curNode[k] + ndid = nd.id + if nd.is_leaf(): + preorder.append(func(nd)) + k = k - 1 + else: + if ndid not in lvisited: + curNode[k + 1] = nd.left + lvisited.add(ndid) + k = k + 1 + elif ndid not in rvisited: + curNode[k + 1] = nd.right + rvisited.add(ndid) + k = k + 1 + # If we've visited the left and right of this non-leaf + # node already, go up in the tree. + else: + k = k - 1 + + return preorder + + +_cnode_bare = ClusterNode(0) +_cnode_type = type(ClusterNode) + + +def _order_cluster_tree(Z): + """ + Return clustering nodes in bottom-up order by distance. + + Parameters + ---------- + Z : scipy.cluster.linkage array + The linkage matrix. + + Returns + ------- + nodes : list + A list of ClusterNode objects. + """ + q = deque() + tree = to_tree(Z) + q.append(tree) + nodes = [] + + while q: + node = q.popleft() + if not node.is_leaf(): + bisect.insort_left(nodes, node) + q.append(node.get_right()) + q.append(node.get_left()) + return nodes + + +def cut_tree(Z, n_clusters=None, height=None): + """ + Given a linkage matrix Z, return the cut tree. + + Parameters + ---------- + Z : scipy.cluster.linkage array + The linkage matrix. + n_clusters : array_like, optional + Number of clusters in the tree at the cut point. + height : array_like, optional + The height at which to cut the tree. Only possible for ultrametric + trees. + + Returns + ------- + cutree : array + An array indicating group membership at each agglomeration step. I.e., + for a full cut tree, in the first column each data point is in its own + cluster. At the next step, two nodes are merged. Finally, all + singleton and non-singleton clusters are in one group. If `n_clusters` + or `height` are given, the columns correspond to the columns of + `n_clusters` or `height`. + + Examples + -------- + >>> from scipy import cluster + >>> import numpy as np + >>> from numpy.random import default_rng + >>> rng = default_rng() + >>> X = rng.random((50, 4)) + >>> Z = cluster.hierarchy.ward(X) + >>> cutree = cluster.hierarchy.cut_tree(Z, n_clusters=[5, 10]) + >>> cutree[:10] + array([[0, 0], + [1, 1], + [2, 2], + [3, 3], + [3, 4], + [2, 2], + [0, 0], + [1, 5], + [3, 6], + [4, 7]]) # random + + """ + xp = array_namespace(Z) + nobs = num_obs_linkage(Z) + nodes = _order_cluster_tree(Z) + + if height is not None and n_clusters is not None: + raise ValueError("At least one of either height or n_clusters " + "must be None") + elif height is None and n_clusters is None: # return the full cut tree + cols_idx = xp.arange(nobs) + elif height is not None: + height = xp.asarray(height) + heights = xp.asarray([x.dist for x in nodes]) + cols_idx = xp.searchsorted(heights, height) + else: + n_clusters = xp.asarray(n_clusters) + cols_idx = nobs - xp.searchsorted(xp.arange(nobs), n_clusters) + + try: + n_cols = len(cols_idx) + except TypeError: # scalar + n_cols = 1 + cols_idx = xp.asarray([cols_idx]) + + groups = xp.zeros((n_cols, nobs), dtype=xp.int64) + last_group = xp.arange(nobs) + if 0 in cols_idx: + groups[0] = last_group + + for i, node in enumerate(nodes): + idx = node.pre_order() + this_group = xp_copy(last_group, xp=xp) + # TODO ARRAY_API complex indexing not supported + this_group[idx] = xp.min(last_group[idx]) + this_group[this_group > xp.max(last_group[idx])] -= 1 + if i + 1 in cols_idx: + groups[np.nonzero(i + 1 == cols_idx)[0]] = this_group + last_group = this_group + + return groups.T + + +def to_tree(Z, rd=False): + """ + Convert a linkage matrix into an easy-to-use tree object. + + The reference to the root `ClusterNode` object is returned (by default). + + Each `ClusterNode` object has a ``left``, ``right``, ``dist``, ``id``, + and ``count`` attribute. The left and right attributes point to + ClusterNode objects that were combined to generate the cluster. + If both are None then the `ClusterNode` object is a leaf node, its count + must be 1, and its distance is meaningless but set to 0. + + *Note: This function is provided for the convenience of the library + user. ClusterNodes are not used as input to any of the functions in this + library.* + + Parameters + ---------- + Z : ndarray + The linkage matrix in proper form (see the `linkage` + function documentation). + rd : bool, optional + When False (default), a reference to the root `ClusterNode` object is + returned. Otherwise, a tuple ``(r, d)`` is returned. ``r`` is a + reference to the root node while ``d`` is a list of `ClusterNode` + objects - one per original entry in the linkage matrix plus entries + for all clustering steps. If a cluster id is + less than the number of samples ``n`` in the data that the linkage + matrix describes, then it corresponds to a singleton cluster (leaf + node). + See `linkage` for more information on the assignment of cluster ids + to clusters. + + Returns + ------- + tree : ClusterNode or tuple (ClusterNode, list of ClusterNode) + If ``rd`` is False, a `ClusterNode`. + If ``rd`` is True, a list of length ``2*n - 1``, with ``n`` the number + of samples. See the description of `rd` above for more details. + + See Also + -------- + linkage, is_valid_linkage, ClusterNode + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster import hierarchy + >>> rng = np.random.default_rng() + >>> x = rng.random((5, 2)) + >>> Z = hierarchy.linkage(x) + >>> hierarchy.to_tree(Z) + >> rootnode, nodelist = hierarchy.to_tree(Z, rd=True) + >>> rootnode + >> len(nodelist) + 9 + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='c', xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + # Number of original objects is equal to the number of rows plus 1. + n = Z.shape[0] + 1 + + # Create a list full of None's to store the node objects + d = [None] * (n * 2 - 1) + + # Create the nodes corresponding to the n original objects. + for i in range(0, n): + d[i] = ClusterNode(i) + + nd = None + + for i in range(Z.shape[0]): + row = Z[i, :] + + fi = int_floor(row[0], xp) + fj = int_floor(row[1], xp) + if fi > i + n: + raise ValueError(('Corrupt matrix Z. Index to derivative cluster ' + 'is used before it is formed. See row %d, ' + 'column 0') % fi) + if fj > i + n: + raise ValueError(('Corrupt matrix Z. Index to derivative cluster ' + 'is used before it is formed. See row %d, ' + 'column 1') % fj) + + nd = ClusterNode(i + n, d[fi], d[fj], row[2]) + # ^ id ^ left ^ right ^ dist + if row[3] != nd.count: + raise ValueError(('Corrupt matrix Z. The count Z[%d,3] is ' + 'incorrect.') % i) + d[n + i] = nd + + if rd: + return (nd, d) + else: + return nd + + +def optimal_leaf_ordering(Z, y, metric='euclidean'): + """ + Given a linkage matrix Z and distance, reorder the cut tree. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded as a linkage matrix. See + `linkage` for more information on the return structure and + algorithm. + y : ndarray + The condensed distance matrix from which Z was generated. + Alternatively, a collection of m observation vectors in n + dimensions may be passed as an m by n array. + metric : str or function, optional + The distance metric to use in the case that y is a collection of + observation vectors; ignored otherwise. See the ``pdist`` + function for a list of valid distance metrics. A custom distance + function can also be used. + + Returns + ------- + Z_ordered : ndarray + A copy of the linkage matrix Z, reordered to minimize the distance + between adjacent leaves. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster import hierarchy + >>> rng = np.random.default_rng() + >>> X = rng.standard_normal((10, 10)) + >>> Z = hierarchy.ward(X) + >>> hierarchy.leaves_list(Z) + array([0, 3, 1, 9, 2, 5, 7, 4, 6, 8], dtype=int32) + >>> hierarchy.leaves_list(hierarchy.optimal_leaf_ordering(Z, X)) + array([3, 0, 2, 5, 7, 4, 8, 6, 9, 1], dtype=int32) + + """ + xp = array_namespace(Z, y) + Z = _asarray(Z, order='C', xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + y = _asarray(y, order='C', dtype=xp.float64, xp=xp) + + if y.ndim == 1: + distance.is_valid_y(y, throw=True, name='y') + elif y.ndim == 2: + if (y.shape[0] == y.shape[1] and np.allclose(np.diag(y), 0) and + np.all(y >= 0) and np.allclose(y, y.T)): + warnings.warn('The symmetric non-negative hollow observation ' + 'matrix looks suspiciously like an uncondensed ' + 'distance matrix', + ClusterWarning, stacklevel=2) + y = distance.pdist(y, metric) + y = xp.asarray(y) + else: + raise ValueError("`y` must be 1 or 2 dimensional.") + + if not xp.all(xp.isfinite(y)): + raise ValueError("The condensed distance matrix must contain only " + "finite values.") + + Z = np.asarray(Z) + y = np.asarray(y) + return xp.asarray(_optimal_leaf_ordering.optimal_leaf_ordering(Z, y)) + + +def cophenet(Z, Y=None): + """ + Calculate the cophenetic distances between each observation in + the hierarchical clustering defined by the linkage ``Z``. + + Suppose ``p`` and ``q`` are original observations in + disjoint clusters ``s`` and ``t``, respectively and + ``s`` and ``t`` are joined by a direct parent cluster + ``u``. The cophenetic distance between observations + ``i`` and ``j`` is simply the distance between + clusters ``s`` and ``t``. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded as an array + (see `linkage` function). + Y : ndarray (optional) + Calculates the cophenetic correlation coefficient ``c`` of a + hierarchical clustering defined by the linkage matrix `Z` + of a set of :math:`n` observations in :math:`m` + dimensions. `Y` is the condensed distance matrix from which + `Z` was generated. + + Returns + ------- + c : ndarray + The cophentic correlation distance (if ``Y`` is passed). + d : ndarray + The cophenetic distance matrix in condensed form. The + :math:`ij` th entry is the cophenetic distance between + original observations :math:`i` and :math:`j`. + + See Also + -------- + linkage : + for a description of what a linkage matrix is. + scipy.spatial.distance.squareform : + transforming condensed matrices into square ones. + + Examples + -------- + >>> from scipy.cluster.hierarchy import single, cophenet + >>> from scipy.spatial.distance import pdist, squareform + + Given a dataset ``X`` and a linkage matrix ``Z``, the cophenetic distance + between two points of ``X`` is the distance between the largest two + distinct clusters that each of the points: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + ``X`` corresponds to this dataset :: + + x x x x + x x + + x x + x x x x + + >>> Z = single(pdist(X)) + >>> Z + array([[ 0., 1., 1., 2.], + [ 2., 12., 1., 3.], + [ 3., 4., 1., 2.], + [ 5., 14., 1., 3.], + [ 6., 7., 1., 2.], + [ 8., 16., 1., 3.], + [ 9., 10., 1., 2.], + [11., 18., 1., 3.], + [13., 15., 2., 6.], + [17., 20., 2., 9.], + [19., 21., 2., 12.]]) + >>> cophenet(Z) + array([1., 1., 2., 2., 2., 2., 2., 2., 2., 2., 2., 1., 2., 2., 2., 2., 2., + 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 1., 1., 2., 2., + 2., 2., 2., 2., 1., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., + 1., 1., 2., 2., 2., 1., 2., 2., 2., 2., 2., 2., 1., 1., 1.]) + + The output of the `scipy.cluster.hierarchy.cophenet` method is + represented in condensed form. We can use + `scipy.spatial.distance.squareform` to see the output as a + regular matrix (where each element ``ij`` denotes the cophenetic distance + between each ``i``, ``j`` pair of points in ``X``): + + >>> squareform(cophenet(Z)) + array([[0., 1., 1., 2., 2., 2., 2., 2., 2., 2., 2., 2.], + [1., 0., 1., 2., 2., 2., 2., 2., 2., 2., 2., 2.], + [1., 1., 0., 2., 2., 2., 2., 2., 2., 2., 2., 2.], + [2., 2., 2., 0., 1., 1., 2., 2., 2., 2., 2., 2.], + [2., 2., 2., 1., 0., 1., 2., 2., 2., 2., 2., 2.], + [2., 2., 2., 1., 1., 0., 2., 2., 2., 2., 2., 2.], + [2., 2., 2., 2., 2., 2., 0., 1., 1., 2., 2., 2.], + [2., 2., 2., 2., 2., 2., 1., 0., 1., 2., 2., 2.], + [2., 2., 2., 2., 2., 2., 1., 1., 0., 2., 2., 2.], + [2., 2., 2., 2., 2., 2., 2., 2., 2., 0., 1., 1.], + [2., 2., 2., 2., 2., 2., 2., 2., 2., 1., 0., 1.], + [2., 2., 2., 2., 2., 2., 2., 2., 2., 1., 1., 0.]]) + + In this example, the cophenetic distance between points on ``X`` that are + very close (i.e., in the same corner) is 1. For other pairs of points is 2, + because the points will be located in clusters at different + corners - thus, the distance between these clusters will be larger. + + """ + xp = array_namespace(Z, Y) + # Ensure float64 C-contiguous array. Cython code doesn't deal with striding. + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + n = Z.shape[0] + 1 + zz = np.zeros((n * (n-1)) // 2, dtype=np.float64) + + Z = np.asarray(Z) + _hierarchy.cophenetic_distances(Z, zz, int(n)) + zz = xp.asarray(zz) + if Y is None: + return zz + + Y = _asarray(Y, order='C', xp=xp) + distance.is_valid_y(Y, throw=True, name='Y') + + z = xp.mean(zz) + y = xp.mean(Y) + Yy = Y - y + Zz = zz - z + numerator = (Yy * Zz) + denomA = Yy**2 + denomB = Zz**2 + c = xp.sum(numerator) / xp.sqrt(xp.sum(denomA) * xp.sum(denomB)) + return (c, zz) + + +def inconsistent(Z, d=2): + r""" + Calculate inconsistency statistics on a linkage matrix. + + Parameters + ---------- + Z : ndarray + The :math:`(n-1)` by 4 matrix encoding the linkage (hierarchical + clustering). See `linkage` documentation for more information on its + form. + d : int, optional + The number of links up to `d` levels below each non-singleton cluster. + + Returns + ------- + R : ndarray + A :math:`(n-1)` by 4 matrix where the ``i``'th row contains the link + statistics for the non-singleton cluster ``i``. The link statistics are + computed over the link heights for links :math:`d` levels below the + cluster ``i``. ``R[i,0]`` and ``R[i,1]`` are the mean and standard + deviation of the link heights, respectively; ``R[i,2]`` is the number + of links included in the calculation; and ``R[i,3]`` is the + inconsistency coefficient, + + .. math:: \frac{\mathtt{Z[i,2]} - \mathtt{R[i,0]}} {R[i,1]} + + Notes + ----- + This function behaves similarly to the MATLAB(TM) ``inconsistent`` + function. + + Examples + -------- + >>> from scipy.cluster.hierarchy import inconsistent, linkage + >>> from matplotlib import pyplot as plt + >>> X = [[i] for i in [2, 8, 0, 4, 1, 9, 9, 0]] + >>> Z = linkage(X, 'ward') + >>> print(Z) + [[ 5. 6. 0. 2. ] + [ 2. 7. 0. 2. ] + [ 0. 4. 1. 2. ] + [ 1. 8. 1.15470054 3. ] + [ 9. 10. 2.12132034 4. ] + [ 3. 12. 4.11096096 5. ] + [11. 13. 14.07183949 8. ]] + >>> inconsistent(Z) + array([[ 0. , 0. , 1. , 0. ], + [ 0. , 0. , 1. , 0. ], + [ 1. , 0. , 1. , 0. ], + [ 0.57735027, 0.81649658, 2. , 0.70710678], + [ 1.04044011, 1.06123822, 3. , 1.01850858], + [ 3.11614065, 1.40688837, 2. , 0.70710678], + [ 6.44583366, 6.76770586, 3. , 1.12682288]]) + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + if (not d == np.floor(d)) or d < 0: + raise ValueError('The second argument d must be a nonnegative ' + 'integer value.') + + n = Z.shape[0] + 1 + R = np.zeros((n - 1, 4), dtype=np.float64) + + Z = np.asarray(Z) + _hierarchy.inconsistent(Z, R, int(n), int(d)) + R = xp.asarray(R) + return R + + +def from_mlab_linkage(Z): + """ + Convert a linkage matrix generated by MATLAB(TM) to a new + linkage matrix compatible with this module. + + The conversion does two things: + + * the indices are converted from ``1..N`` to ``0..(N-1)`` form, + and + + * a fourth column ``Z[:,3]`` is added where ``Z[i,3]`` represents the + number of original observations (leaves) in the non-singleton + cluster ``i``. + + This function is useful when loading in linkages from legacy data + files generated by MATLAB. + + Parameters + ---------- + Z : ndarray + A linkage matrix generated by MATLAB(TM). + + Returns + ------- + ZS : ndarray + A linkage matrix compatible with ``scipy.cluster.hierarchy``. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + to_mlab_linkage : transform from SciPy to MATLAB format. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster.hierarchy import ward, from_mlab_linkage + + Given a linkage matrix in MATLAB format ``mZ``, we can use + `scipy.cluster.hierarchy.from_mlab_linkage` to import + it into SciPy format: + + >>> mZ = np.array([[1, 2, 1], [4, 5, 1], [7, 8, 1], + ... [10, 11, 1], [3, 13, 1.29099445], + ... [6, 14, 1.29099445], + ... [9, 15, 1.29099445], + ... [12, 16, 1.29099445], + ... [17, 18, 5.77350269], + ... [19, 20, 5.77350269], + ... [21, 22, 8.16496581]]) + + >>> Z = from_mlab_linkage(mZ) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [ 11. , 15. , 1.29099445, 3. ], + [ 16. , 17. , 5.77350269, 6. ], + [ 18. , 19. , 5.77350269, 6. ], + [ 20. , 21. , 8.16496581, 12. ]]) + + As expected, the linkage matrix ``Z`` returned includes an + additional column counting the number of original samples in + each cluster. Also, all cluster indices are reduced by 1 + (MATLAB format uses 1-indexing, whereas SciPy uses 0-indexing). + + """ + xp = array_namespace(Z) + Z = _asarray(Z, dtype=xp.float64, order='C', xp=xp) + Zs = Z.shape + + # If it's empty, return it. + if len(Zs) == 0 or (len(Zs) == 1 and Zs[0] == 0): + return xp_copy(Z, xp=xp) + + if len(Zs) != 2: + raise ValueError("The linkage array must be rectangular.") + + # If it contains no rows, return it. + if Zs[0] == 0: + return xp_copy(Z, xp=xp) + + if xp.min(Z[:, 0:2]) != 1.0 and xp.max(Z[:, 0:2]) != 2 * Zs[0]: + raise ValueError('The format of the indices is not 1..N') + + Zpart = xp.concat((Z[:, 0:2] - 1.0, Z[:, 2:]), axis=1) + CS = np.zeros((Zs[0],), dtype=np.float64) + if is_jax(xp): + # calculate_cluster_sizes doesn't accept read-only arrays + Zpart = np.array(Zpart, copy=True) + else: + Zpart = np.asarray(Zpart) + _hierarchy.calculate_cluster_sizes(Zpart, CS, int(Zs[0]) + 1) + res = np.hstack([Zpart, CS.reshape(Zs[0], 1)]) + return xp.asarray(res) + + +def to_mlab_linkage(Z): + """ + Convert a linkage matrix to a MATLAB(TM) compatible one. + + Converts a linkage matrix ``Z`` generated by the linkage function + of this module to a MATLAB(TM) compatible one. The return linkage + matrix has the last column removed and the cluster indices are + converted to ``1..N`` indexing. + + Parameters + ---------- + Z : ndarray + A linkage matrix generated by ``scipy.cluster.hierarchy``. + + Returns + ------- + to_mlab_linkage : ndarray + A linkage matrix compatible with MATLAB(TM)'s hierarchical + clustering functions. + + The return linkage matrix has the last column removed + and the cluster indices are converted to ``1..N`` indexing. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + from_mlab_linkage : transform from Matlab to SciPy format. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, to_mlab_linkage + >>> from scipy.spatial.distance import pdist + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + + After a linkage matrix ``Z`` has been created, we can use + `scipy.cluster.hierarchy.to_mlab_linkage` to convert it + into MATLAB format: + + >>> mZ = to_mlab_linkage(Z) + >>> mZ + array([[ 1. , 2. , 1. ], + [ 4. , 5. , 1. ], + [ 7. , 8. , 1. ], + [ 10. , 11. , 1. ], + [ 3. , 13. , 1.29099445], + [ 6. , 14. , 1.29099445], + [ 9. , 15. , 1.29099445], + [ 12. , 16. , 1.29099445], + [ 17. , 18. , 5.77350269], + [ 19. , 20. , 5.77350269], + [ 21. , 22. , 8.16496581]]) + + The new linkage matrix ``mZ`` uses 1-indexing for all the + clusters (instead of 0-indexing). Also, the last column of + the original linkage matrix has been dropped. + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + Zs = Z.shape + if len(Zs) == 0 or (len(Zs) == 1 and Zs[0] == 0): + return xp_copy(Z, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + return xp.concat((Z[:, :2] + 1.0, Z[:, 2:3]), axis=1) + + +def is_monotonic(Z): + """ + Return True if the linkage passed is monotonic. + + The linkage is monotonic if for every cluster :math:`s` and :math:`t` + joined, the distance between them is no less than the distance + between any previously joined clusters. + + Parameters + ---------- + Z : ndarray + The linkage matrix to check for monotonicity. + + Returns + ------- + b : bool + A boolean indicating whether the linkage is monotonic. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + + Examples + -------- + >>> from scipy.cluster.hierarchy import median, ward, is_monotonic + >>> from scipy.spatial.distance import pdist + + By definition, some hierarchical clustering algorithms - such as + `scipy.cluster.hierarchy.ward` - produce monotonic assignments of + samples to clusters; however, this is not always true for other + hierarchical methods - e.g. `scipy.cluster.hierarchy.median`. + + Given a linkage matrix ``Z`` (as the result of a hierarchical clustering + method) we can test programmatically whether it has the monotonicity + property or not, using `scipy.cluster.hierarchy.is_monotonic`: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + >>> is_monotonic(Z) + True + + >>> Z = median(pdist(X)) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 2. , 12. , 1.11803399, 3. ], + [ 5. , 13. , 1.11803399, 3. ], + [ 8. , 15. , 1.11803399, 3. ], + [11. , 14. , 1.11803399, 3. ], + [18. , 19. , 3. , 6. ], + [16. , 17. , 3.5 , 6. ], + [20. , 21. , 3.25 , 12. ]]) + >>> is_monotonic(Z) + False + + Note that this method is equivalent to just verifying that the distances + in the third column of the linkage matrix appear in a monotonically + increasing order. + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='c', xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + # We expect the i'th value to be greater than its successor. + return xp.all(Z[1:, 2] >= Z[:-1, 2]) + + +def is_valid_im(R, warning=False, throw=False, name=None): + """Return True if the inconsistency matrix passed is valid. + + It must be a :math:`n` by 4 array of doubles. The standard + deviations ``R[:,1]`` must be nonnegative. The link counts + ``R[:,2]`` must be positive and no greater than :math:`n-1`. + + Parameters + ---------- + R : ndarray + The inconsistency matrix to check for validity. + warning : bool, optional + When True, issues a Python warning if the linkage + matrix passed is invalid. + throw : bool, optional + When True, throws a Python exception if the linkage + matrix passed is invalid. + name : str, optional + This string refers to the variable name of the invalid + linkage matrix. + + Returns + ------- + b : bool + True if the inconsistency matrix is valid. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + inconsistent : for the creation of a inconsistency matrix. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, inconsistent, is_valid_im + >>> from scipy.spatial.distance import pdist + + Given a data set ``X``, we can apply a clustering method to obtain a + linkage matrix ``Z``. `scipy.cluster.hierarchy.inconsistent` can + be also used to obtain the inconsistency matrix ``R`` associated to + this clustering process: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + >>> R = inconsistent(Z) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + >>> R + array([[1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1.14549722, 0.20576415, 2. , 0.70710678], + [1.14549722, 0.20576415, 2. , 0.70710678], + [1.14549722, 0.20576415, 2. , 0.70710678], + [1.14549722, 0.20576415, 2. , 0.70710678], + [2.78516386, 2.58797734, 3. , 1.15470054], + [2.78516386, 2.58797734, 3. , 1.15470054], + [6.57065706, 1.38071187, 3. , 1.15470054]]) + + Now we can use `scipy.cluster.hierarchy.is_valid_im` to verify that + ``R`` is correct: + + >>> is_valid_im(R) + True + + However, if ``R`` is wrongly constructed (e.g., one of the standard + deviations is set to a negative value), then the check will fail: + + >>> R[-1,1] = R[-1,1] * -1 + >>> is_valid_im(R) + False + + """ + xp = array_namespace(R) + R = _asarray(R, order='c', xp=xp) + valid = True + name_str = f"{name!r} " if name else '' + try: + if R.dtype != xp.float64: + raise TypeError(f'Inconsistency matrix {name_str}must contain doubles ' + '(double).') + if len(R.shape) != 2: + raise ValueError(f'Inconsistency matrix {name_str}must have shape=2 (i.e. ' + 'be two-dimensional).') + if R.shape[1] != 4: + raise ValueError(f'Inconsistency matrix {name_str}' + 'must have 4 columns.') + if R.shape[0] < 1: + raise ValueError(f'Inconsistency matrix {name_str}' + 'must have at least one row.') + if xp.any(R[:, 0] < 0): + raise ValueError(f'Inconsistency matrix {name_str}' + 'contains negative link height means.') + if xp.any(R[:, 1] < 0): + raise ValueError(f'Inconsistency matrix {name_str}' + 'contains negative link height standard deviations.') + if xp.any(R[:, 2] < 0): + raise ValueError(f'Inconsistency matrix {name_str}' + 'contains negative link counts.') + except Exception as e: + if throw: + raise + if warning: + _warning(str(e)) + valid = False + + return valid + + +def is_valid_linkage(Z, warning=False, throw=False, name=None): + """ + Check the validity of a linkage matrix. + + A linkage matrix is valid if it is a 2-D array (type double) + with :math:`n` rows and 4 columns. The first two columns must contain + indices between 0 and :math:`2n-1`. For a given row ``i``, the following + two expressions have to hold: + + .. math:: + + 0 \\leq \\mathtt{Z[i,0]} \\leq i+n-1 + 0 \\leq Z[i,1] \\leq i+n-1 + + I.e., a cluster cannot join another cluster unless the cluster being joined + has been generated. + + The fourth column of `Z` represents the number of original observations + in a cluster, so a valid ``Z[i, 3]`` value may not exceed the number of + original observations. + + Parameters + ---------- + Z : array_like + Linkage matrix. + warning : bool, optional + When True, issues a Python warning if the linkage + matrix passed is invalid. + throw : bool, optional + When True, throws a Python exception if the linkage + matrix passed is invalid. + name : str, optional + This string refers to the variable name of the invalid + linkage matrix. + + Returns + ------- + b : bool + True if the inconsistency matrix is valid. + + See Also + -------- + linkage: for a description of what a linkage matrix is. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, is_valid_linkage + >>> from scipy.spatial.distance import pdist + + All linkage matrices generated by the clustering methods in this module + will be valid (i.e., they will have the appropriate dimensions and the two + required expressions will hold for all the rows). + + We can check this using `scipy.cluster.hierarchy.is_valid_linkage`: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + >>> is_valid_linkage(Z) + True + + However, if we create a linkage matrix in a wrong way - or if we modify + a valid one in a way that any of the required expressions don't hold + anymore, then the check will fail: + + >>> Z[3][1] = 20 # the cluster number 20 is not defined at this point + >>> is_valid_linkage(Z) + False + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='c', xp=xp) + valid = True + name_str = f"{name!r} " if name else '' + try: + if Z.dtype != xp.float64: + raise TypeError(f'Linkage matrix {name_str}must contain doubles.') + if len(Z.shape) != 2: + raise ValueError(f'Linkage matrix {name_str}must have shape=2 (i.e. be' + ' two-dimensional).') + if Z.shape[1] != 4: + raise ValueError(f'Linkage matrix {name_str}must have 4 columns.') + if Z.shape[0] == 0: + raise ValueError('Linkage must be computed on at least two ' + 'observations.') + n = Z.shape[0] + if n > 1: + if (xp.any(Z[:, 0] < 0) or xp.any(Z[:, 1] < 0)): + raise ValueError(f'Linkage {name_str}contains negative indices.') + if xp.any(Z[:, 2] < 0): + raise ValueError(f'Linkage {name_str}contains negative distances.') + if xp.any(Z[:, 3] < 0): + raise ValueError(f'Linkage {name_str}contains negative counts.') + if xp.any(Z[:, 3] > (Z.shape[0] + 1)): + raise ValueError('Linkage matrix contains excessive observations' + 'in a cluster') + if _check_hierarchy_uses_cluster_before_formed(Z): + raise ValueError(f'Linkage {name_str}uses non-singleton cluster before' + ' it is formed.') + if _check_hierarchy_uses_cluster_more_than_once(Z): + raise ValueError(f'Linkage {name_str}uses the same cluster more than once.') + except Exception as e: + if throw: + raise + if warning: + _warning(str(e)) + valid = False + + return valid + + +def _check_hierarchy_uses_cluster_before_formed(Z): + n = Z.shape[0] + 1 + for i in range(0, n - 1): + if Z[i, 0] >= n + i or Z[i, 1] >= n + i: + return True + return False + + +def _check_hierarchy_uses_cluster_more_than_once(Z): + n = Z.shape[0] + 1 + chosen = set() + for i in range(0, n - 1): + used_more_than_once = ( + (float(Z[i, 0]) in chosen) + or (float(Z[i, 1]) in chosen) + or Z[i, 0] == Z[i, 1] + ) + if used_more_than_once: + return True + chosen.add(float(Z[i, 0])) + chosen.add(float(Z[i, 1])) + return False + + +def _check_hierarchy_not_all_clusters_used(Z): + n = Z.shape[0] + 1 + chosen = set() + for i in range(0, n - 1): + chosen.add(int(Z[i, 0])) + chosen.add(int(Z[i, 1])) + must_chosen = set(range(0, 2 * n - 2)) + return len(must_chosen.difference(chosen)) > 0 + + +def num_obs_linkage(Z): + """ + Return the number of original observations of the linkage matrix passed. + + Parameters + ---------- + Z : ndarray + The linkage matrix on which to perform the operation. + + Returns + ------- + n : int + The number of original observations in the linkage. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, num_obs_linkage + >>> from scipy.spatial.distance import pdist + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + + ``Z`` is a linkage matrix obtained after using the Ward clustering method + with ``X``, a dataset with 12 data points. + + >>> num_obs_linkage(Z) + 12 + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='c', xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + return (Z.shape[0] + 1) + + +def correspond(Z, Y): + """ + Check for correspondence between linkage and condensed distance matrices. + + They must have the same number of original observations for + the check to succeed. + + This function is useful as a sanity check in algorithms that make + extensive use of linkage and distance matrices that must + correspond to the same set of original observations. + + Parameters + ---------- + Z : array_like + The linkage matrix to check for correspondence. + Y : array_like + The condensed distance matrix to check for correspondence. + + Returns + ------- + b : bool + A boolean indicating whether the linkage matrix and distance + matrix could possibly correspond to one another. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, correspond + >>> from scipy.spatial.distance import pdist + + This method can be used to check if a given linkage matrix ``Z`` has been + obtained from the application of a cluster method over a dataset ``X``: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + >>> X_condensed = pdist(X) + >>> Z = ward(X_condensed) + + Here, we can compare ``Z`` and ``X`` (in condensed form): + + >>> correspond(Z, X_condensed) + True + + """ + is_valid_linkage(Z, throw=True) + distance.is_valid_y(Y, throw=True) + xp = array_namespace(Z, Y) + Z = _asarray(Z, order='c', xp=xp) + Y = _asarray(Y, order='c', xp=xp) + return distance.num_obs_y(Y) == num_obs_linkage(Z) + + +def fcluster(Z, t, criterion='inconsistent', depth=2, R=None, monocrit=None): + """ + Form flat clusters from the hierarchical clustering defined by + the given linkage matrix. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded with the matrix returned + by the `linkage` function. + t : scalar + For criteria 'inconsistent', 'distance' or 'monocrit', + this is the threshold to apply when forming flat clusters. + For 'maxclust' or 'maxclust_monocrit' criteria, + this would be max number of clusters requested. + criterion : str, optional + The criterion to use in forming flat clusters. This can + be any of the following values: + + ``inconsistent`` : + If a cluster node and all its + descendants have an inconsistent value less than or equal + to `t`, then all its leaf descendants belong to the + same flat cluster. When no non-singleton cluster meets + this criterion, every node is assigned to its own + cluster. (Default) + + ``distance`` : + Forms flat clusters so that the original + observations in each flat cluster have no greater a + cophenetic distance than `t`. + + ``maxclust`` : + Finds a minimum threshold ``r`` so that + the cophenetic distance between any two original + observations in the same flat cluster is no more than + ``r`` and no more than `t` flat clusters are formed. + + ``monocrit`` : + Forms a flat cluster from a cluster node c + with index i when ``monocrit[j] <= t``. + + For example, to threshold on the maximum mean distance + as computed in the inconsistency matrix R with a + threshold of 0.8 do:: + + MR = maxRstat(Z, R, 3) + fcluster(Z, t=0.8, criterion='monocrit', monocrit=MR) + + ``maxclust_monocrit`` : + Forms a flat cluster from a + non-singleton cluster node ``c`` when ``monocrit[i] <= + r`` for all cluster indices ``i`` below and including + ``c``. ``r`` is minimized such that no more than ``t`` + flat clusters are formed. monocrit must be + monotonic. For example, to minimize the threshold t on + maximum inconsistency values so that no more than 3 flat + clusters are formed, do:: + + MI = maxinconsts(Z, R) + fcluster(Z, t=3, criterion='maxclust_monocrit', monocrit=MI) + depth : int, optional + The maximum depth to perform the inconsistency calculation. + It has no meaning for the other criteria. Default is 2. + R : ndarray, optional + The inconsistency matrix to use for the ``'inconsistent'`` + criterion. This matrix is computed if not provided. + monocrit : ndarray, optional + An array of length n-1. `monocrit[i]` is the + statistics upon which non-singleton i is thresholded. The + monocrit vector must be monotonic, i.e., given a node c with + index i, for all node indices j corresponding to nodes + below c, ``monocrit[i] >= monocrit[j]``. + + Returns + ------- + fcluster : ndarray + An array of length ``n``. ``T[i]`` is the flat cluster number to + which original observation ``i`` belongs. + + See Also + -------- + linkage : for information about hierarchical clustering methods work. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, fcluster + >>> from scipy.spatial.distance import pdist + + All cluster linkage methods - e.g., `scipy.cluster.hierarchy.ward` + generate a linkage matrix ``Z`` as their output: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + + This matrix represents a dendrogram, where the first and second elements + are the two clusters merged at each step, the third element is the + distance between these clusters, and the fourth element is the size of + the new cluster - the number of original data points included. + + `scipy.cluster.hierarchy.fcluster` can be used to flatten the + dendrogram, obtaining as a result an assignation of the original data + points to single clusters. + + This assignation mostly depends on a distance threshold ``t`` - the maximum + inter-cluster distance allowed: + + >>> fcluster(Z, t=0.9, criterion='distance') + array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=int32) + + >>> fcluster(Z, t=1.1, criterion='distance') + array([1, 1, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8], dtype=int32) + + >>> fcluster(Z, t=3, criterion='distance') + array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], dtype=int32) + + >>> fcluster(Z, t=9, criterion='distance') + array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32) + + In the first case, the threshold ``t`` is too small to allow any two + samples in the data to form a cluster, so 12 different clusters are + returned. + + In the second case, the threshold is large enough to allow the first + 4 points to be merged with their nearest neighbors. So, here, only 8 + clusters are returned. + + The third case, with a much higher threshold, allows for up to 8 data + points to be connected - so 4 clusters are returned here. + + Lastly, the threshold of the fourth case is large enough to allow for + all data points to be merged together - so a single cluster is returned. + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + n = Z.shape[0] + 1 + T = np.zeros((n,), dtype='i') + + if monocrit is not None: + monocrit = np.asarray(monocrit, order='C', dtype=np.float64) + + Z = np.asarray(Z) + monocrit = np.asarray(monocrit) + if criterion == 'inconsistent': + if R is None: + R = inconsistent(Z, depth) + else: + R = _asarray(R, order='C', dtype=xp.float64, xp=xp) + is_valid_im(R, throw=True, name='R') + # Since the C code does not support striding using strides. + # The dimensions are used instead. + R = np.asarray(R) + _hierarchy.cluster_in(Z, R, T, float(t), int(n)) + elif criterion == 'distance': + _hierarchy.cluster_dist(Z, T, float(t), int(n)) + elif criterion == 'maxclust': + _hierarchy.cluster_maxclust_dist(Z, T, int(n), t) + elif criterion == 'monocrit': + _hierarchy.cluster_monocrit(Z, monocrit, T, float(t), int(n)) + elif criterion == 'maxclust_monocrit': + _hierarchy.cluster_maxclust_monocrit(Z, monocrit, T, int(n), int(t)) + else: + raise ValueError(f'Invalid cluster formation criterion: {str(criterion)}') + return xp.asarray(T) + + +def fclusterdata(X, t, criterion='inconsistent', + metric='euclidean', depth=2, method='single', R=None): + """ + Cluster observation data using a given metric. + + Clusters the original observations in the n-by-m data + matrix X (n observations in m dimensions), using the euclidean + distance metric to calculate distances between original observations, + performs hierarchical clustering using the single linkage algorithm, + and forms flat clusters using the inconsistency method with `t` as the + cut-off threshold. + + A 1-D array ``T`` of length ``n`` is returned. ``T[i]`` is + the index of the flat cluster to which the original observation ``i`` + belongs. + + Parameters + ---------- + X : (N, M) ndarray + N by M data matrix with N observations in M dimensions. + t : scalar + For criteria 'inconsistent', 'distance' or 'monocrit', + this is the threshold to apply when forming flat clusters. + For 'maxclust' or 'maxclust_monocrit' criteria, + this would be max number of clusters requested. + criterion : str, optional + Specifies the criterion for forming flat clusters. Valid + values are 'inconsistent' (default), 'distance', or 'maxclust' + cluster formation algorithms. See `fcluster` for descriptions. + metric : str or function, optional + The distance metric for calculating pairwise distances. See + ``distance.pdist`` for descriptions and linkage to verify + compatibility with the linkage method. + depth : int, optional + The maximum depth for the inconsistency calculation. See + `inconsistent` for more information. + method : str, optional + The linkage method to use (single, complete, average, + weighted, median centroid, ward). See `linkage` for more + information. Default is "single". + R : ndarray, optional + The inconsistency matrix. It will be computed if necessary + if it is not passed. + + Returns + ------- + fclusterdata : ndarray + A vector of length n. T[i] is the flat cluster number to + which original observation i belongs. + + See Also + -------- + scipy.spatial.distance.pdist : pairwise distance metrics + + Notes + ----- + This function is similar to the MATLAB function ``clusterdata``. + + Examples + -------- + >>> from scipy.cluster.hierarchy import fclusterdata + + This is a convenience method that abstracts all the steps to perform in a + typical SciPy's hierarchical clustering workflow. + + * Transform the input data into a condensed matrix with + `scipy.spatial.distance.pdist`. + + * Apply a clustering method. + + * Obtain flat clusters at a user defined distance threshold ``t`` using + `scipy.cluster.hierarchy.fcluster`. + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> fclusterdata(X, t=1) + array([3, 3, 3, 4, 4, 4, 2, 2, 2, 1, 1, 1], dtype=int32) + + The output here (for the dataset ``X``, distance threshold ``t``, and the + default settings) is four clusters with three data points each. + + """ + xp = array_namespace(X) + X = _asarray(X, order='C', dtype=xp.float64, xp=xp) + + if X.ndim != 2: + raise TypeError('The observation matrix X must be an n by m ' + 'array.') + + Y = distance.pdist(X, metric=metric) + Y = xp.asarray(Y) + Z = linkage(Y, method=method) + if R is None: + R = inconsistent(Z, d=depth) + else: + R = _asarray(R, order='c', xp=xp) + T = fcluster(Z, criterion=criterion, depth=depth, R=R, t=t) + return T + + +def leaves_list(Z): + """ + Return a list of leaf node ids. + + The return corresponds to the observation vector index as it appears + in the tree from left to right. Z is a linkage matrix. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded as a matrix. `Z` is + a linkage matrix. See `linkage` for more information. + + Returns + ------- + leaves_list : ndarray + The list of leaf node ids. + + See Also + -------- + dendrogram : for information about dendrogram structure. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, dendrogram, leaves_list + >>> from scipy.spatial.distance import pdist + >>> from matplotlib import pyplot as plt + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + + The linkage matrix ``Z`` represents a dendrogram, that is, a tree that + encodes the structure of the clustering performed. + `scipy.cluster.hierarchy.leaves_list` shows the mapping between + indices in the ``X`` dataset and leaves in the dendrogram: + + >>> leaves_list(Z) + array([ 2, 0, 1, 5, 3, 4, 8, 6, 7, 11, 9, 10], dtype=int32) + + >>> fig = plt.figure(figsize=(25, 10)) + >>> dn = dendrogram(Z) + >>> plt.show() + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='C', xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + n = Z.shape[0] + 1 + ML = np.zeros((n,), dtype='i') + Z = np.asarray(Z) + _hierarchy.prelist(Z, ML, n) + return xp.asarray(ML) + + +# Maps number of leaves to text size. +# +# p <= 20, size="12" +# 20 < p <= 30, size="10" +# 30 < p <= 50, size="8" +# 50 < p <= np.inf, size="6" + +_dtextsizes = {20: 12, 30: 10, 50: 8, 85: 6, np.inf: 5} +_drotation = {20: 0, 40: 45, np.inf: 90} +_dtextsortedkeys = list(_dtextsizes.keys()) +_dtextsortedkeys.sort() +_drotationsortedkeys = list(_drotation.keys()) +_drotationsortedkeys.sort() + + +def _remove_dups(L): + """ + Remove duplicates AND preserve the original order of the elements. + + The set class is not guaranteed to do this. + """ + seen_before = set() + L2 = [] + for i in L: + if i not in seen_before: + seen_before.add(i) + L2.append(i) + return L2 + + +def _get_tick_text_size(p): + for k in _dtextsortedkeys: + if p <= k: + return _dtextsizes[k] + + +def _get_tick_rotation(p): + for k in _drotationsortedkeys: + if p <= k: + return _drotation[k] + + +def _plot_dendrogram(icoords, dcoords, ivl, p, n, mh, orientation, + no_labels, color_list, leaf_font_size=None, + leaf_rotation=None, contraction_marks=None, + ax=None, above_threshold_color='C0'): + # Import matplotlib here so that it's not imported unless dendrograms + # are plotted. Raise an informative error if importing fails. + try: + # if an axis is provided, don't use pylab at all + if ax is None: + import matplotlib.pylab + import matplotlib.patches + import matplotlib.collections + except ImportError as e: + raise ImportError("You must install the matplotlib library to plot " + "the dendrogram. Use no_plot=True to calculate the " + "dendrogram without plotting.") from e + + if ax is None: + ax = matplotlib.pylab.gca() + # if we're using pylab, we want to trigger a draw at the end + trigger_redraw = True + else: + trigger_redraw = False + + # Independent variable plot width + ivw = len(ivl) * 10 + # Dependent variable plot height + dvw = mh + mh * 0.05 + + iv_ticks = np.arange(5, len(ivl) * 10 + 5, 10) + if orientation in ('top', 'bottom'): + if orientation == 'top': + ax.set_ylim([0, dvw]) + ax.set_xlim([0, ivw]) + else: + ax.set_ylim([dvw, 0]) + ax.set_xlim([0, ivw]) + + xlines = icoords + ylines = dcoords + if no_labels: + ax.set_xticks([]) + ax.set_xticklabels([]) + else: + ax.set_xticks(iv_ticks) + + if orientation == 'top': + ax.xaxis.set_ticks_position('bottom') + else: + ax.xaxis.set_ticks_position('top') + + # Make the tick marks invisible because they cover up the links + for line in ax.get_xticklines(): + line.set_visible(False) + + leaf_rot = (float(_get_tick_rotation(len(ivl))) + if (leaf_rotation is None) else leaf_rotation) + leaf_font = (float(_get_tick_text_size(len(ivl))) + if (leaf_font_size is None) else leaf_font_size) + ax.set_xticklabels(ivl, rotation=leaf_rot, size=leaf_font) + + elif orientation in ('left', 'right'): + if orientation == 'left': + ax.set_xlim([dvw, 0]) + ax.set_ylim([0, ivw]) + else: + ax.set_xlim([0, dvw]) + ax.set_ylim([0, ivw]) + + xlines = dcoords + ylines = icoords + if no_labels: + ax.set_yticks([]) + ax.set_yticklabels([]) + else: + ax.set_yticks(iv_ticks) + + if orientation == 'left': + ax.yaxis.set_ticks_position('right') + else: + ax.yaxis.set_ticks_position('left') + + # Make the tick marks invisible because they cover up the links + for line in ax.get_yticklines(): + line.set_visible(False) + + leaf_font = (float(_get_tick_text_size(len(ivl))) + if (leaf_font_size is None) else leaf_font_size) + + if leaf_rotation is not None: + ax.set_yticklabels(ivl, rotation=leaf_rotation, size=leaf_font) + else: + ax.set_yticklabels(ivl, size=leaf_font) + + # Let's use collections instead. This way there is a separate legend item + # for each tree grouping, rather than stupidly one for each line segment. + colors_used = _remove_dups(color_list) + color_to_lines = {} + for color in colors_used: + color_to_lines[color] = [] + for (xline, yline, color) in zip(xlines, ylines, color_list): + color_to_lines[color].append(list(zip(xline, yline))) + + colors_to_collections = {} + # Construct the collections. + for color in colors_used: + coll = matplotlib.collections.LineCollection(color_to_lines[color], + colors=(color,)) + colors_to_collections[color] = coll + + # Add all the groupings below the color threshold. + for color in colors_used: + if color != above_threshold_color: + ax.add_collection(colors_to_collections[color]) + # If there's a grouping of links above the color threshold, it goes last. + if above_threshold_color in colors_to_collections: + ax.add_collection(colors_to_collections[above_threshold_color]) + + if contraction_marks is not None: + Ellipse = matplotlib.patches.Ellipse + for (x, y) in contraction_marks: + if orientation in ('left', 'right'): + e = Ellipse((y, x), width=dvw / 100, height=1.0) + else: + e = Ellipse((x, y), width=1.0, height=dvw / 100) + ax.add_artist(e) + e.set_clip_box(ax.bbox) + e.set_alpha(0.5) + e.set_facecolor('k') + + if trigger_redraw: + matplotlib.pylab.draw_if_interactive() + + +# C0 is used for above threshold color +_link_line_colors_default = ('C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9') +_link_line_colors = list(_link_line_colors_default) + + +def set_link_color_palette(palette): + """ + Set list of matplotlib color codes for use by dendrogram. + + Note that this palette is global (i.e., setting it once changes the colors + for all subsequent calls to `dendrogram`) and that it affects only the + the colors below ``color_threshold``. + + Note that `dendrogram` also accepts a custom coloring function through its + ``link_color_func`` keyword, which is more flexible and non-global. + + Parameters + ---------- + palette : list of str or None + A list of matplotlib color codes. The order of the color codes is the + order in which the colors are cycled through when color thresholding in + the dendrogram. + + If ``None``, resets the palette to its default (which are matplotlib + default colors C1 to C9). + + Returns + ------- + None + + See Also + -------- + dendrogram + + Notes + ----- + Ability to reset the palette with ``None`` added in SciPy 0.17.0. + + Thread safety: using this function in a multi-threaded fashion may + result in `dendrogram` producing plots with unexpected colors. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster import hierarchy + >>> ytdist = np.array([662., 877., 255., 412., 996., 295., 468., 268., + ... 400., 754., 564., 138., 219., 869., 669.]) + >>> Z = hierarchy.linkage(ytdist, 'single') + >>> dn = hierarchy.dendrogram(Z, no_plot=True) + >>> dn['color_list'] + ['C1', 'C0', 'C0', 'C0', 'C0'] + >>> hierarchy.set_link_color_palette(['c', 'm', 'y', 'k']) + >>> dn = hierarchy.dendrogram(Z, no_plot=True, above_threshold_color='b') + >>> dn['color_list'] + ['c', 'b', 'b', 'b', 'b'] + >>> dn = hierarchy.dendrogram(Z, no_plot=True, color_threshold=267, + ... above_threshold_color='k') + >>> dn['color_list'] + ['c', 'm', 'm', 'k', 'k'] + + Now, reset the color palette to its default: + + >>> hierarchy.set_link_color_palette(None) + + """ + if palette is None: + # reset to its default + palette = _link_line_colors_default + elif not isinstance(palette, (list, tuple)): + raise TypeError("palette must be a list or tuple") + _ptypes = [isinstance(p, str) for p in palette] + + if False in _ptypes: + raise TypeError("all palette list elements must be color strings") + + global _link_line_colors + _link_line_colors = palette + + +def dendrogram(Z, p=30, truncate_mode=None, color_threshold=None, + get_leaves=True, orientation='top', labels=None, + count_sort=False, distance_sort=False, show_leaf_counts=True, + no_plot=False, no_labels=False, leaf_font_size=None, + leaf_rotation=None, leaf_label_func=None, + show_contracted=False, link_color_func=None, ax=None, + above_threshold_color='C0'): + """ + Plot the hierarchical clustering as a dendrogram. + + The dendrogram illustrates how each cluster is + composed by drawing a U-shaped link between a non-singleton + cluster and its children. The top of the U-link indicates a + cluster merge. The two legs of the U-link indicate which clusters + were merged. The length of the two legs of the U-link represents + the distance between the child clusters. It is also the + cophenetic distance between original observations in the two + children clusters. + + Parameters + ---------- + Z : ndarray + The linkage matrix encoding the hierarchical clustering to + render as a dendrogram. See the ``linkage`` function for more + information on the format of ``Z``. + p : int, optional + The ``p`` parameter for ``truncate_mode``. + truncate_mode : str, optional + The dendrogram can be hard to read when the original + observation matrix from which the linkage is derived is + large. Truncation is used to condense the dendrogram. There + are several modes: + + ``None`` + No truncation is performed (default). + Note: ``'none'`` is an alias for ``None`` that's kept for + backward compatibility. + + ``'lastp'`` + The last ``p`` non-singleton clusters formed in the linkage are the + only non-leaf nodes in the linkage; they correspond to rows + ``Z[n-p-2:end]`` in ``Z``. All other non-singleton clusters are + contracted into leaf nodes. + + ``'level'`` + No more than ``p`` levels of the dendrogram tree are displayed. + A "level" includes all nodes with ``p`` merges from the final merge. + + Note: ``'mtica'`` is an alias for ``'level'`` that's kept for + backward compatibility. + + color_threshold : double, optional + For brevity, let :math:`t` be the ``color_threshold``. + Colors all the descendent links below a cluster node + :math:`k` the same color if :math:`k` is the first node below + the cut threshold :math:`t`. All links connecting nodes with + distances greater than or equal to the threshold are colored + with de default matplotlib color ``'C0'``. If :math:`t` is less + than or equal to zero, all nodes are colored ``'C0'``. + If ``color_threshold`` is None or 'default', + corresponding with MATLAB(TM) behavior, the threshold is set to + ``0.7*max(Z[:,2])``. + + get_leaves : bool, optional + Includes a list ``R['leaves']=H`` in the result + dictionary. For each :math:`i`, ``H[i] == j``, cluster node + ``j`` appears in position ``i`` in the left-to-right traversal + of the leaves, where :math:`j < 2n-1` and :math:`i < n`. + orientation : str, optional + The direction to plot the dendrogram, which can be any + of the following strings: + + ``'top'`` + Plots the root at the top, and plot descendent links going downwards. + (default). + + ``'bottom'`` + Plots the root at the bottom, and plot descendent links going + upwards. + + ``'left'`` + Plots the root at the left, and plot descendent links going right. + + ``'right'`` + Plots the root at the right, and plot descendent links going left. + + labels : ndarray, optional + By default, ``labels`` is None so the index of the original observation + is used to label the leaf nodes. Otherwise, this is an :math:`n`-sized + sequence, with ``n == Z.shape[0] + 1``. The ``labels[i]`` value is the + text to put under the :math:`i` th leaf node only if it corresponds to + an original observation and not a non-singleton cluster. + count_sort : str or bool, optional + For each node n, the order (visually, from left-to-right) n's + two descendent links are plotted is determined by this + parameter, which can be any of the following values: + + ``False`` + Nothing is done. + + ``'ascending'`` or ``True`` + The child with the minimum number of original objects in its cluster + is plotted first. + + ``'descending'`` + The child with the maximum number of original objects in its cluster + is plotted first. + + Note, ``distance_sort`` and ``count_sort`` cannot both be True. + distance_sort : str or bool, optional + For each node n, the order (visually, from left-to-right) n's + two descendent links are plotted is determined by this + parameter, which can be any of the following values: + + ``False`` + Nothing is done. + + ``'ascending'`` or ``True`` + The child with the minimum distance between its direct descendents is + plotted first. + + ``'descending'`` + The child with the maximum distance between its direct descendents is + plotted first. + + Note ``distance_sort`` and ``count_sort`` cannot both be True. + show_leaf_counts : bool, optional + When True, leaf nodes representing :math:`k>1` original + observation are labeled with the number of observations they + contain in parentheses. + no_plot : bool, optional + When True, the final rendering is not performed. This is + useful if only the data structures computed for the rendering + are needed or if matplotlib is not available. + no_labels : bool, optional + When True, no labels appear next to the leaf nodes in the + rendering of the dendrogram. + leaf_rotation : double, optional + Specifies the angle (in degrees) to rotate the leaf + labels. When unspecified, the rotation is based on the number of + nodes in the dendrogram (default is 0). + leaf_font_size : int, optional + Specifies the font size (in points) of the leaf labels. When + unspecified, the size based on the number of nodes in the + dendrogram. + leaf_label_func : lambda or function, optional + When ``leaf_label_func`` is a callable function, for each + leaf with cluster index :math:`k < 2n-1`. The function + is expected to return a string with the label for the + leaf. + + Indices :math:`k < n` correspond to original observations + while indices :math:`k \\geq n` correspond to non-singleton + clusters. + + For example, to label singletons with their node id and + non-singletons with their id, count, and inconsistency + coefficient, simply do:: + + # First define the leaf label function. + def llf(id): + if id < n: + return str(id) + else: + return '[%d %d %1.2f]' % (id, count, R[n-id,3]) + + # The text for the leaf nodes is going to be big so force + # a rotation of 90 degrees. + dendrogram(Z, leaf_label_func=llf, leaf_rotation=90) + + # leaf_label_func can also be used together with ``truncate_mode``, + # in which case you will get your leaves labeled after truncation: + dendrogram(Z, leaf_label_func=llf, leaf_rotation=90, + truncate_mode='level', p=2) + + show_contracted : bool, optional + When True the heights of non-singleton nodes contracted + into a leaf node are plotted as crosses along the link + connecting that leaf node. This really is only useful when + truncation is used (see ``truncate_mode`` parameter). + link_color_func : callable, optional + If given, `link_color_function` is called with each non-singleton id + corresponding to each U-shaped link it will paint. The function is + expected to return the color to paint the link, encoded as a matplotlib + color string code. For example:: + + dendrogram(Z, link_color_func=lambda k: colors[k]) + + colors the direct links below each untruncated non-singleton node + ``k`` using ``colors[k]``. + ax : matplotlib Axes instance, optional + If None and `no_plot` is not True, the dendrogram will be plotted + on the current axes. Otherwise if `no_plot` is not True the + dendrogram will be plotted on the given ``Axes`` instance. This can be + useful if the dendrogram is part of a more complex figure. + above_threshold_color : str, optional + This matplotlib color string sets the color of the links above the + color_threshold. The default is ``'C0'``. + + Returns + ------- + R : dict + A dictionary of data structures computed to render the + dendrogram. Its has the following keys: + + ``'color_list'`` + A list of color names. The k'th element represents the color of the + k'th link. + + ``'icoord'`` and ``'dcoord'`` + Each of them is a list of lists. Let ``icoord = [I1, I2, ..., Ip]`` + where ``Ik = [xk1, xk2, xk3, xk4]`` and ``dcoord = [D1, D2, ..., Dp]`` + where ``Dk = [yk1, yk2, yk3, yk4]``, then the k'th link painted is + ``(xk1, yk1)`` - ``(xk2, yk2)`` - ``(xk3, yk3)`` - ``(xk4, yk4)``. + + ``'ivl'`` + A list of labels corresponding to the leaf nodes. + + ``'leaves'`` + For each i, ``H[i] == j``, cluster node ``j`` appears in position + ``i`` in the left-to-right traversal of the leaves, where + :math:`j < 2n-1` and :math:`i < n`. If ``j`` is less than ``n``, the + ``i``-th leaf node corresponds to an original observation. + Otherwise, it corresponds to a non-singleton cluster. + + ``'leaves_color_list'`` + A list of color names. The k'th element represents the color of the + k'th leaf. + + See Also + -------- + linkage, set_link_color_palette + + Notes + ----- + It is expected that the distances in ``Z[:,2]`` be monotonic, otherwise + crossings appear in the dendrogram. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster import hierarchy + >>> import matplotlib.pyplot as plt + + A very basic example: + + >>> ytdist = np.array([662., 877., 255., 412., 996., 295., 468., 268., + ... 400., 754., 564., 138., 219., 869., 669.]) + >>> Z = hierarchy.linkage(ytdist, 'single') + >>> plt.figure() + >>> dn = hierarchy.dendrogram(Z) + + Now, plot in given axes, improve the color scheme and use both vertical and + horizontal orientations: + + >>> hierarchy.set_link_color_palette(['m', 'c', 'y', 'k']) + >>> fig, axes = plt.subplots(1, 2, figsize=(8, 3)) + >>> dn1 = hierarchy.dendrogram(Z, ax=axes[0], above_threshold_color='y', + ... orientation='top') + >>> dn2 = hierarchy.dendrogram(Z, ax=axes[1], + ... above_threshold_color='#bcbddc', + ... orientation='right') + >>> hierarchy.set_link_color_palette(None) # reset to default after use + >>> plt.show() + + """ + # This feature was thought about but never implemented (still useful?): + # + # ... = dendrogram(..., leaves_order=None) + # + # Plots the leaves in the order specified by a vector of + # original observation indices. If the vector contains duplicates + # or results in a crossing, an exception will be thrown. Passing + # None orders leaf nodes based on the order they appear in the + # pre-order traversal. + xp = array_namespace(Z) + Z = _asarray(Z, order='c', xp=xp) + + if orientation not in ["top", "left", "bottom", "right"]: + raise ValueError("orientation must be one of 'top', 'left', " + "'bottom', or 'right'") + + if labels is not None: + try: + len_labels = len(labels) + except (TypeError, AttributeError): + len_labels = labels.shape[0] + if Z.shape[0] + 1 != len_labels: + raise ValueError("Dimensions of Z and labels must be consistent.") + + is_valid_linkage(Z, throw=True, name='Z') + Zs = Z.shape + n = Zs[0] + 1 + if isinstance(p, (int, float)): + p = int(p) + else: + raise TypeError('The second argument must be a number') + + if truncate_mode not in ('lastp', 'mtica', 'level', 'none', None): + # 'mtica' is kept working for backwards compat. + raise ValueError('Invalid truncation mode.') + + if truncate_mode == 'lastp': + if p > n or p == 0: + p = n + + if truncate_mode == 'mtica': + # 'mtica' is an alias + truncate_mode = 'level' + + if truncate_mode == 'level': + if p <= 0: + p = np.inf + + if get_leaves: + lvs = [] + else: + lvs = None + + icoord_list = [] + dcoord_list = [] + color_list = [] + current_color = [0] + currently_below_threshold = [False] + ivl = [] # list of leaves + + if color_threshold is None or (isinstance(color_threshold, str) and + color_threshold == 'default'): + color_threshold = xp.max(Z[:, 2]) * 0.7 + + R = {'icoord': icoord_list, 'dcoord': dcoord_list, 'ivl': ivl, + 'leaves': lvs, 'color_list': color_list} + + # Empty list will be filled in _dendrogram_calculate_info + contraction_marks = [] if show_contracted else None + + _dendrogram_calculate_info( + Z=Z, p=p, + truncate_mode=truncate_mode, + color_threshold=color_threshold, + get_leaves=get_leaves, + orientation=orientation, + labels=labels, + count_sort=count_sort, + distance_sort=distance_sort, + show_leaf_counts=show_leaf_counts, + i=2*n - 2, + iv=0.0, + ivl=ivl, + n=n, + icoord_list=icoord_list, + dcoord_list=dcoord_list, + lvs=lvs, + current_color=current_color, + color_list=color_list, + currently_below_threshold=currently_below_threshold, + leaf_label_func=leaf_label_func, + contraction_marks=contraction_marks, + link_color_func=link_color_func, + above_threshold_color=above_threshold_color) + + if not no_plot: + mh = xp.max(Z[:, 2]) + _plot_dendrogram(icoord_list, dcoord_list, ivl, p, n, mh, orientation, + no_labels, color_list, + leaf_font_size=leaf_font_size, + leaf_rotation=leaf_rotation, + contraction_marks=contraction_marks, + ax=ax, + above_threshold_color=above_threshold_color) + + R["leaves_color_list"] = _get_leaves_color_list(R) + + return R + + +def _get_leaves_color_list(R): + leaves_color_list = [None] * len(R['leaves']) + for link_x, link_y, link_color in zip(R['icoord'], + R['dcoord'], + R['color_list']): + for (xi, yi) in zip(link_x, link_y): + if yi == 0.0 and (xi % 5 == 0 and xi % 2 == 1): + # if yi is 0.0 and xi is divisible by 5 and odd, + # the point is a leaf + # xi of leaves are 5, 15, 25, 35, ... (see `iv_ticks`) + # index of leaves are 0, 1, 2, 3, ... as below + leaf_index = (int(xi) - 5) // 10 + # each leaf has a same color of its link. + leaves_color_list[leaf_index] = link_color + return leaves_color_list + + +def _append_singleton_leaf_node(Z, p, n, level, lvs, ivl, leaf_label_func, + i, labels): + # If the leaf id structure is not None and is a list then the caller + # to dendrogram has indicated that cluster id's corresponding to the + # leaf nodes should be recorded. + + if lvs is not None: + lvs.append(int(i)) + + # If leaf node labels are to be displayed... + if ivl is not None: + # If a leaf_label_func has been provided, the label comes from the + # string returned from the leaf_label_func, which is a function + # passed to dendrogram. + if leaf_label_func: + ivl.append(leaf_label_func(int(i))) + else: + # Otherwise, if the dendrogram caller has passed a labels list + # for the leaf nodes, use it. + if labels is not None: + ivl.append(labels[int(i - n)]) + else: + # Otherwise, use the id as the label for the leaf.x + ivl.append(str(int(i))) + + +def _append_nonsingleton_leaf_node(Z, p, n, level, lvs, ivl, leaf_label_func, + i, labels, show_leaf_counts): + # If the leaf id structure is not None and is a list then the caller + # to dendrogram has indicated that cluster id's corresponding to the + # leaf nodes should be recorded. + + if lvs is not None: + lvs.append(int(i)) + if ivl is not None: + if leaf_label_func: + ivl.append(leaf_label_func(int(i))) + else: + if show_leaf_counts: + ivl.append("(" + str(np.asarray(Z[i - n, 3], dtype=np.int64)) + ")") + else: + ivl.append("") + + +def _append_contraction_marks(Z, iv, i, n, contraction_marks, xp): + _append_contraction_marks_sub(Z, iv, int_floor(Z[i - n, 0], xp), + n, contraction_marks, xp) + _append_contraction_marks_sub(Z, iv, int_floor(Z[i - n, 1], xp), + n, contraction_marks, xp) + + +def _append_contraction_marks_sub(Z, iv, i, n, contraction_marks, xp): + if i >= n: + contraction_marks.append((iv, Z[i - n, 2])) + _append_contraction_marks_sub(Z, iv, int_floor(Z[i - n, 0], xp), + n, contraction_marks, xp) + _append_contraction_marks_sub(Z, iv, int_floor(Z[i - n, 1], xp), + n, contraction_marks, xp) + + +def _dendrogram_calculate_info(Z, p, truncate_mode, + color_threshold=np.inf, get_leaves=True, + orientation='top', labels=None, + count_sort=False, distance_sort=False, + show_leaf_counts=False, i=-1, iv=0.0, + ivl=None, n=0, icoord_list=None, dcoord_list=None, + lvs=None, mhr=False, + current_color=None, color_list=None, + currently_below_threshold=None, + leaf_label_func=None, level=0, + contraction_marks=None, + link_color_func=None, + above_threshold_color='C0'): + """ + Calculate the endpoints of the links as well as the labels for the + the dendrogram rooted at the node with index i. iv is the independent + variable value to plot the left-most leaf node below the root node i + (if orientation='top', this would be the left-most x value where the + plotting of this root node i and its descendents should begin). + + ivl is a list to store the labels of the leaf nodes. The leaf_label_func + is called whenever ivl != None, labels == None, and + leaf_label_func != None. When ivl != None and labels != None, the + labels list is used only for labeling the leaf nodes. When + ivl == None, no labels are generated for leaf nodes. + + When get_leaves==True, a list of leaves is built as they are visited + in the dendrogram. + + Returns a tuple with l being the independent variable coordinate that + corresponds to the midpoint of cluster to the left of cluster i if + i is non-singleton, otherwise the independent coordinate of the leaf + node if i is a leaf node. + + Returns + ------- + A tuple (left, w, h, md), where: + * left is the independent variable coordinate of the center of the + the U of the subtree + + * w is the amount of space used for the subtree (in independent + variable units) + + * h is the height of the subtree in dependent variable units + + * md is the ``max(Z[*,2]``) for all nodes ``*`` below and including + the target node. + + """ + xp = array_namespace(Z) + if n == 0: + raise ValueError("Invalid singleton cluster count n.") + + if i == -1: + raise ValueError("Invalid root cluster index i.") + + if truncate_mode == 'lastp': + # If the node is a leaf node but corresponds to a non-singleton + # cluster, its label is either the empty string or the number of + # original observations belonging to cluster i. + if 2*n - p > i >= n: + d = Z[i - n, 2] + _append_nonsingleton_leaf_node(Z, p, n, level, lvs, ivl, + leaf_label_func, i, labels, + show_leaf_counts) + if contraction_marks is not None: + _append_contraction_marks(Z, iv + 5.0, i, n, contraction_marks, xp) + return (iv + 5.0, 10.0, 0.0, d) + elif i < n: + _append_singleton_leaf_node(Z, p, n, level, lvs, ivl, + leaf_label_func, i, labels) + return (iv + 5.0, 10.0, 0.0, 0.0) + elif truncate_mode == 'level': + if i > n and level > p: + d = Z[i - n, 2] + _append_nonsingleton_leaf_node(Z, p, n, level, lvs, ivl, + leaf_label_func, i, labels, + show_leaf_counts) + if contraction_marks is not None: + _append_contraction_marks(Z, iv + 5.0, i, n, contraction_marks, xp) + return (iv + 5.0, 10.0, 0.0, d) + elif i < n: + _append_singleton_leaf_node(Z, p, n, level, lvs, ivl, + leaf_label_func, i, labels) + return (iv + 5.0, 10.0, 0.0, 0.0) + + # Otherwise, only truncate if we have a leaf node. + # + # Only place leaves if they correspond to original observations. + if i < n: + _append_singleton_leaf_node(Z, p, n, level, lvs, ivl, + leaf_label_func, i, labels) + return (iv + 5.0, 10.0, 0.0, 0.0) + + # !!! Otherwise, we don't have a leaf node, so work on plotting a + # non-leaf node. + # Actual indices of a and b + aa = int_floor(Z[i - n, 0], xp) + ab = int_floor(Z[i - n, 1], xp) + if aa >= n: + # The number of singletons below cluster a + na = Z[aa - n, 3] + # The distance between a's two direct children. + da = Z[aa - n, 2] + else: + na = 1 + da = 0.0 + if ab >= n: + nb = Z[ab - n, 3] + db = Z[ab - n, 2] + else: + nb = 1 + db = 0.0 + + if count_sort == 'ascending' or count_sort is True: + # If a has a count greater than b, it and its descendents should + # be drawn to the right. Otherwise, to the left. + if na > nb: + # The cluster index to draw to the left (ua) will be ab + # and the one to draw to the right (ub) will be aa + ua = ab + ub = aa + else: + ua = aa + ub = ab + elif count_sort == 'descending': + # If a has a count less than or equal to b, it and its + # descendents should be drawn to the left. Otherwise, to + # the right. + if na > nb: + ua = aa + ub = ab + else: + ua = ab + ub = aa + elif distance_sort == 'ascending' or distance_sort is True: + # If a has a distance greater than b, it and its descendents should + # be drawn to the right. Otherwise, to the left. + if da > db: + ua = ab + ub = aa + else: + ua = aa + ub = ab + elif distance_sort == 'descending': + # If a has a distance less than or equal to b, it and its + # descendents should be drawn to the left. Otherwise, to + # the right. + if da > db: + ua = aa + ub = ab + else: + ua = ab + ub = aa + else: + ua = aa + ub = ab + + # Updated iv variable and the amount of space used. + (uiva, uwa, uah, uamd) = \ + _dendrogram_calculate_info( + Z=Z, p=p, + truncate_mode=truncate_mode, + color_threshold=color_threshold, + get_leaves=get_leaves, + orientation=orientation, + labels=labels, + count_sort=count_sort, + distance_sort=distance_sort, + show_leaf_counts=show_leaf_counts, + i=ua, iv=iv, ivl=ivl, n=n, + icoord_list=icoord_list, + dcoord_list=dcoord_list, lvs=lvs, + current_color=current_color, + color_list=color_list, + currently_below_threshold=currently_below_threshold, + leaf_label_func=leaf_label_func, + level=level + 1, contraction_marks=contraction_marks, + link_color_func=link_color_func, + above_threshold_color=above_threshold_color) + + h = Z[i - n, 2] + if h >= color_threshold or color_threshold <= 0: + c = above_threshold_color + + if currently_below_threshold[0]: + current_color[0] = (current_color[0] + 1) % len(_link_line_colors) + currently_below_threshold[0] = False + else: + currently_below_threshold[0] = True + c = _link_line_colors[current_color[0]] + + (uivb, uwb, ubh, ubmd) = \ + _dendrogram_calculate_info( + Z=Z, p=p, + truncate_mode=truncate_mode, + color_threshold=color_threshold, + get_leaves=get_leaves, + orientation=orientation, + labels=labels, + count_sort=count_sort, + distance_sort=distance_sort, + show_leaf_counts=show_leaf_counts, + i=ub, iv=iv + uwa, ivl=ivl, n=n, + icoord_list=icoord_list, + dcoord_list=dcoord_list, lvs=lvs, + current_color=current_color, + color_list=color_list, + currently_below_threshold=currently_below_threshold, + leaf_label_func=leaf_label_func, + level=level + 1, contraction_marks=contraction_marks, + link_color_func=link_color_func, + above_threshold_color=above_threshold_color) + + max_dist = max(uamd, ubmd, h) + + icoord_list.append([uiva, uiva, uivb, uivb]) + dcoord_list.append([uah, h, h, ubh]) + if link_color_func is not None: + v = link_color_func(int(i)) + if not isinstance(v, str): + raise TypeError("link_color_func must return a matplotlib " + "color string!") + color_list.append(v) + else: + color_list.append(c) + + return (((uiva + uivb) / 2), uwa + uwb, h, max_dist) + + +def is_isomorphic(T1, T2): + """ + Determine if two different cluster assignments are equivalent. + + Parameters + ---------- + T1 : array_like + An assignment of singleton cluster ids to flat cluster ids. + T2 : array_like + An assignment of singleton cluster ids to flat cluster ids. + + Returns + ------- + b : bool + Whether the flat cluster assignments `T1` and `T2` are + equivalent. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + fcluster : for the creation of flat cluster assignments. + + Examples + -------- + >>> from scipy.cluster.hierarchy import fcluster, is_isomorphic + >>> from scipy.cluster.hierarchy import single, complete + >>> from scipy.spatial.distance import pdist + + Two flat cluster assignments can be isomorphic if they represent the same + cluster assignment, with different labels. + + For example, we can use the `scipy.cluster.hierarchy.single`: method + and flatten the output to four clusters: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = single(pdist(X)) + >>> T = fcluster(Z, 1, criterion='distance') + >>> T + array([3, 3, 3, 4, 4, 4, 2, 2, 2, 1, 1, 1], dtype=int32) + + We can then do the same using the + `scipy.cluster.hierarchy.complete`: method: + + >>> Z = complete(pdist(X)) + >>> T_ = fcluster(Z, 1.5, criterion='distance') + >>> T_ + array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], dtype=int32) + + As we can see, in both cases we obtain four clusters and all the data + points are distributed in the same way - the only thing that changes + are the flat cluster labels (3 => 1, 4 =>2, 2 =>3 and 4 =>1), so both + cluster assignments are isomorphic: + + >>> is_isomorphic(T, T_) + True + + """ + T1 = np.asarray(T1, order='c') + T2 = np.asarray(T2, order='c') + + T1S = T1.shape + T2S = T2.shape + + if len(T1S) != 1: + raise ValueError('T1 must be one-dimensional.') + if len(T2S) != 1: + raise ValueError('T2 must be one-dimensional.') + if T1S[0] != T2S[0]: + raise ValueError('T1 and T2 must have the same number of elements.') + n = T1S[0] + d1 = {} + d2 = {} + for i in range(0, n): + if T1[i] in d1: + if T2[i] not in d2: + return False + if d1[T1[i]] != T2[i] or d2[T2[i]] != T1[i]: + return False + elif T2[i] in d2: + return False + else: + d1[T1[i]] = T2[i] + d2[T2[i]] = T1[i] + return True + + +def maxdists(Z): + """ + Return the maximum distance between any non-singleton cluster. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded as a matrix. See + ``linkage`` for more information. + + Returns + ------- + maxdists : ndarray + A ``(n-1)`` sized numpy array of doubles; ``MD[i]`` represents + the maximum distance between any cluster (including + singletons) below and including the node with index i. More + specifically, ``MD[i] = Z[Q(i)-n, 2].max()`` where ``Q(i)`` is the + set of all node indices below and including node i. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + is_monotonic : for testing for monotonicity of a linkage matrix. + + Examples + -------- + >>> from scipy.cluster.hierarchy import median, maxdists + >>> from scipy.spatial.distance import pdist + + Given a linkage matrix ``Z``, `scipy.cluster.hierarchy.maxdists` + computes for each new cluster generated (i.e., for each row of the linkage + matrix) what is the maximum distance between any two child clusters. + + Due to the nature of hierarchical clustering, in many cases this is going + to be just the distance between the two child clusters that were merged + to form the current one - that is, Z[:,2]. + + However, for non-monotonic cluster assignments such as + `scipy.cluster.hierarchy.median` clustering this is not always the + case: There may be cluster formations were the distance between the two + clusters merged is smaller than the distance between their children. + + We can see this in an example: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = median(pdist(X)) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 2. , 12. , 1.11803399, 3. ], + [ 5. , 13. , 1.11803399, 3. ], + [ 8. , 15. , 1.11803399, 3. ], + [11. , 14. , 1.11803399, 3. ], + [18. , 19. , 3. , 6. ], + [16. , 17. , 3.5 , 6. ], + [20. , 21. , 3.25 , 12. ]]) + >>> maxdists(Z) + array([1. , 1. , 1. , 1. , 1.11803399, + 1.11803399, 1.11803399, 1.11803399, 3. , 3.5 , + 3.5 ]) + + Note that while the distance between the two clusters merged when creating the + last cluster is 3.25, there are two children (clusters 16 and 17) whose distance + is larger (3.5). Thus, `scipy.cluster.hierarchy.maxdists` returns 3.5 in + this case. + + """ + xp = array_namespace(Z) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + n = Z.shape[0] + 1 + MD = np.zeros((n - 1,)) + Z = np.asarray(Z) + _hierarchy.get_max_dist_for_each_cluster(Z, MD, int(n)) + MD = xp.asarray(MD) + return MD + + +def maxinconsts(Z, R): + """ + Return the maximum inconsistency coefficient for each + non-singleton cluster and its children. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded as a matrix. See + `linkage` for more information. + R : ndarray + The inconsistency matrix. + + Returns + ------- + MI : ndarray + A monotonic ``(n-1)``-sized numpy array of doubles. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + inconsistent : for the creation of a inconsistency matrix. + + Examples + -------- + >>> from scipy.cluster.hierarchy import median, inconsistent, maxinconsts + >>> from scipy.spatial.distance import pdist + + Given a data set ``X``, we can apply a clustering method to obtain a + linkage matrix ``Z``. `scipy.cluster.hierarchy.inconsistent` can + be also used to obtain the inconsistency matrix ``R`` associated to + this clustering process: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = median(pdist(X)) + >>> R = inconsistent(Z) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 2. , 12. , 1.11803399, 3. ], + [ 5. , 13. , 1.11803399, 3. ], + [ 8. , 15. , 1.11803399, 3. ], + [11. , 14. , 1.11803399, 3. ], + [18. , 19. , 3. , 6. ], + [16. , 17. , 3.5 , 6. ], + [20. , 21. , 3.25 , 12. ]]) + >>> R + array([[1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.74535599, 1.08655358, 3. , 1.15470054], + [1.91202266, 1.37522872, 3. , 1.15470054], + [3.25 , 0.25 , 3. , 0. ]]) + + Here, `scipy.cluster.hierarchy.maxinconsts` can be used to compute + the maximum value of the inconsistency statistic (the last column of + ``R``) for each non-singleton cluster and its children: + + >>> maxinconsts(Z, R) + array([0. , 0. , 0. , 0. , 0.70710678, + 0.70710678, 0.70710678, 0.70710678, 1.15470054, 1.15470054, + 1.15470054]) + + """ + xp = array_namespace(Z, R) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + R = _asarray(R, order='C', dtype=xp.float64, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + is_valid_im(R, throw=True, name='R') + + n = Z.shape[0] + 1 + if Z.shape[0] != R.shape[0]: + raise ValueError("The inconsistency matrix and linkage matrix each " + "have a different number of rows.") + MI = np.zeros((n - 1,)) + Z = np.asarray(Z) + R = np.asarray(R) + _hierarchy.get_max_Rfield_for_each_cluster(Z, R, MI, int(n), 3) + MI = xp.asarray(MI) + return MI + + +def maxRstat(Z, R, i): + """ + Return the maximum statistic for each non-singleton cluster and its + children. + + Parameters + ---------- + Z : array_like + The hierarchical clustering encoded as a matrix. See `linkage` for more + information. + R : array_like + The inconsistency matrix. + i : int + The column of `R` to use as the statistic. + + Returns + ------- + MR : ndarray + Calculates the maximum statistic for the i'th column of the + inconsistency matrix `R` for each non-singleton cluster + node. ``MR[j]`` is the maximum over ``R[Q(j)-n, i]``, where + ``Q(j)`` the set of all node ids corresponding to nodes below + and including ``j``. + + See Also + -------- + linkage : for a description of what a linkage matrix is. + inconsistent : for the creation of a inconsistency matrix. + + Examples + -------- + >>> from scipy.cluster.hierarchy import median, inconsistent, maxRstat + >>> from scipy.spatial.distance import pdist + + Given a data set ``X``, we can apply a clustering method to obtain a + linkage matrix ``Z``. `scipy.cluster.hierarchy.inconsistent` can + be also used to obtain the inconsistency matrix ``R`` associated to + this clustering process: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = median(pdist(X)) + >>> R = inconsistent(Z) + >>> R + array([[1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1. , 0. , 1. , 0. ], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.05901699, 0.08346263, 2. , 0.70710678], + [1.74535599, 1.08655358, 3. , 1.15470054], + [1.91202266, 1.37522872, 3. , 1.15470054], + [3.25 , 0.25 , 3. , 0. ]]) + + `scipy.cluster.hierarchy.maxRstat` can be used to compute + the maximum value of each column of ``R``, for each non-singleton + cluster and its children: + + >>> maxRstat(Z, R, 0) + array([1. , 1. , 1. , 1. , 1.05901699, + 1.05901699, 1.05901699, 1.05901699, 1.74535599, 1.91202266, + 3.25 ]) + >>> maxRstat(Z, R, 1) + array([0. , 0. , 0. , 0. , 0.08346263, + 0.08346263, 0.08346263, 0.08346263, 1.08655358, 1.37522872, + 1.37522872]) + >>> maxRstat(Z, R, 3) + array([0. , 0. , 0. , 0. , 0.70710678, + 0.70710678, 0.70710678, 0.70710678, 1.15470054, 1.15470054, + 1.15470054]) + + """ + xp = array_namespace(Z, R) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + R = _asarray(R, order='C', dtype=xp.float64, xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + is_valid_im(R, throw=True, name='R') + + if not isinstance(i, int): + raise TypeError('The third argument must be an integer.') + + if i < 0 or i > 3: + raise ValueError('i must be an integer between 0 and 3 inclusive.') + + if Z.shape[0] != R.shape[0]: + raise ValueError("The inconsistency matrix and linkage matrix each " + "have a different number of rows.") + + n = Z.shape[0] + 1 + MR = np.zeros((n - 1,)) + Z = np.asarray(Z) + R = np.asarray(R) + _hierarchy.get_max_Rfield_for_each_cluster(Z, R, MR, int(n), i) + MR = xp.asarray(MR) + return MR + + +def leaders(Z, T): + """ + Return the root nodes in a hierarchical clustering. + + Returns the root nodes in a hierarchical clustering corresponding + to a cut defined by a flat cluster assignment vector ``T``. See + the ``fcluster`` function for more information on the format of ``T``. + + For each flat cluster :math:`j` of the :math:`k` flat clusters + represented in the n-sized flat cluster assignment vector ``T``, + this function finds the lowest cluster node :math:`i` in the linkage + tree Z, such that: + + * leaf descendants belong only to flat cluster j + (i.e., ``T[p]==j`` for all :math:`p` in :math:`S(i)`, where + :math:`S(i)` is the set of leaf ids of descendant leaf nodes + with cluster node :math:`i`) + + * there does not exist a leaf that is not a descendant with + :math:`i` that also belongs to cluster :math:`j` + (i.e., ``T[q]!=j`` for all :math:`q` not in :math:`S(i)`). If + this condition is violated, ``T`` is not a valid cluster + assignment vector, and an exception will be thrown. + + Parameters + ---------- + Z : ndarray + The hierarchical clustering encoded as a matrix. See + `linkage` for more information. + T : ndarray + The flat cluster assignment vector. + + Returns + ------- + L : ndarray + The leader linkage node id's stored as a k-element 1-D array, + where ``k`` is the number of flat clusters found in ``T``. + + ``L[j]=i`` is the linkage cluster node id that is the + leader of flat cluster with id M[j]. If ``i < n``, ``i`` + corresponds to an original observation, otherwise it + corresponds to a non-singleton cluster. + M : ndarray + The leader linkage node id's stored as a k-element 1-D array, where + ``k`` is the number of flat clusters found in ``T``. This allows the + set of flat cluster ids to be any arbitrary set of ``k`` integers. + + For example: if ``L[3]=2`` and ``M[3]=8``, the flat cluster with + id 8's leader is linkage node 2. + + See Also + -------- + fcluster : for the creation of flat cluster assignments. + + Examples + -------- + >>> from scipy.cluster.hierarchy import ward, fcluster, leaders + >>> from scipy.spatial.distance import pdist + + Given a linkage matrix ``Z`` - obtained after apply a clustering method + to a dataset ``X`` - and a flat cluster assignment array ``T``: + + >>> X = [[0, 0], [0, 1], [1, 0], + ... [0, 4], [0, 3], [1, 4], + ... [4, 0], [3, 0], [4, 1], + ... [4, 4], [3, 4], [4, 3]] + + >>> Z = ward(pdist(X)) + >>> Z + array([[ 0. , 1. , 1. , 2. ], + [ 3. , 4. , 1. , 2. ], + [ 6. , 7. , 1. , 2. ], + [ 9. , 10. , 1. , 2. ], + [ 2. , 12. , 1.29099445, 3. ], + [ 5. , 13. , 1.29099445, 3. ], + [ 8. , 14. , 1.29099445, 3. ], + [11. , 15. , 1.29099445, 3. ], + [16. , 17. , 5.77350269, 6. ], + [18. , 19. , 5.77350269, 6. ], + [20. , 21. , 8.16496581, 12. ]]) + + >>> T = fcluster(Z, 3, criterion='distance') + >>> T + array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], dtype=int32) + + `scipy.cluster.hierarchy.leaders` returns the indices of the nodes + in the dendrogram that are the leaders of each flat cluster: + + >>> L, M = leaders(Z, T) + >>> L + array([16, 17, 18, 19], dtype=int32) + + (remember that indices 0-11 point to the 12 data points in ``X``, + whereas indices 12-22 point to the 11 rows of ``Z``) + + `scipy.cluster.hierarchy.leaders` also returns the indices of + the flat clusters in ``T``: + + >>> M + array([1, 2, 3, 4], dtype=int32) + + """ + xp = array_namespace(Z, T) + Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp) + T = _asarray(T, order='C', xp=xp) + is_valid_linkage(Z, throw=True, name='Z') + + if T.dtype != xp.int32: + raise TypeError('T must be a 1-D array of dtype int32.') + + if T.shape[0] != Z.shape[0] + 1: + raise ValueError('Mismatch: len(T)!=Z.shape[0] + 1.') + + n_clusters = int(xp.unique_values(T).shape[0]) + n_obs = int(Z.shape[0] + 1) + L = np.zeros(n_clusters, dtype=np.int32) + M = np.zeros(n_clusters, dtype=np.int32) + Z = np.asarray(Z) + T = np.asarray(T, dtype=np.int32) + s = _hierarchy.leaders(Z, T, L, M, n_clusters, n_obs) + if s >= 0: + raise ValueError(('T is not a valid assignment vector. Error found ' + 'when examining linkage node %d (< 2n-1).') % s) + L, M = xp.asarray(L), xp.asarray(M) + return (L, M) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/hierarchy_test_data.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/hierarchy_test_data.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad7da3af7fa3acf8f399c9eedb149a1420d5a17 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/hierarchy_test_data.py @@ -0,0 +1,145 @@ +from numpy import array + + +Q_X = array([[5.26563660e-01, 3.14160190e-01, 8.00656370e-02], + [7.50205180e-01, 4.60299830e-01, 8.98696460e-01], + [6.65461230e-01, 6.94011420e-01, 9.10465700e-01], + [9.64047590e-01, 1.43082200e-03, 7.39874220e-01], + [1.08159060e-01, 5.53028790e-01, 6.63804780e-02], + [9.31359130e-01, 8.25424910e-01, 9.52315440e-01], + [6.78086960e-01, 3.41903970e-01, 5.61481950e-01], + [9.82730940e-01, 7.04605210e-01, 8.70978630e-02], + [6.14691610e-01, 4.69989230e-02, 6.02406450e-01], + [5.80161260e-01, 9.17354970e-01, 5.88163850e-01], + [1.38246310e+00, 1.96358160e+00, 1.94437880e+00], + [2.10675860e+00, 1.67148730e+00, 1.34854480e+00], + [1.39880070e+00, 1.66142050e+00, 1.32224550e+00], + [1.71410460e+00, 1.49176380e+00, 1.45432170e+00], + [1.54102340e+00, 1.84374950e+00, 1.64658950e+00], + [2.08512480e+00, 1.84524350e+00, 2.17340850e+00], + [1.30748740e+00, 1.53801650e+00, 2.16007740e+00], + [1.41447700e+00, 1.99329070e+00, 1.99107420e+00], + [1.61943490e+00, 1.47703280e+00, 1.89788160e+00], + [1.59880600e+00, 1.54988980e+00, 1.57563350e+00], + [3.37247380e+00, 2.69635310e+00, 3.39981700e+00], + [3.13705120e+00, 3.36528090e+00, 3.06089070e+00], + [3.29413250e+00, 3.19619500e+00, 2.90700170e+00], + [2.65510510e+00, 3.06785900e+00, 2.97198540e+00], + [3.30941040e+00, 2.59283970e+00, 2.57714110e+00], + [2.59557220e+00, 3.33477370e+00, 3.08793190e+00], + [2.58206180e+00, 3.41615670e+00, 3.26441990e+00], + [2.71127000e+00, 2.77032450e+00, 2.63466500e+00], + [2.79617850e+00, 3.25473720e+00, 3.41801560e+00], + [2.64741750e+00, 2.54538040e+00, 3.25354110e+00]]) + +ytdist = array([662., 877., 255., 412., 996., 295., 468., 268., 400., 754., + 564., 138., 219., 869., 669.]) + +linkage_ytdist_single = array([[2., 5., 138., 2.], + [3., 4., 219., 2.], + [0., 7., 255., 3.], + [1., 8., 268., 4.], + [6., 9., 295., 6.]]) + +linkage_ytdist_complete = array([[2., 5., 138., 2.], + [3., 4., 219., 2.], + [1., 6., 400., 3.], + [0., 7., 412., 3.], + [8., 9., 996., 6.]]) + +linkage_ytdist_average = array([[2., 5., 138., 2.], + [3., 4., 219., 2.], + [0., 7., 333.5, 3.], + [1., 6., 347.5, 3.], + [8., 9., 680.77777778, 6.]]) + +linkage_ytdist_weighted = array([[2., 5., 138., 2.], + [3., 4., 219., 2.], + [0., 7., 333.5, 3.], + [1., 6., 347.5, 3.], + [8., 9., 670.125, 6.]]) + +# the optimal leaf ordering of linkage_ytdist_single +linkage_ytdist_single_olo = array([[5., 2., 138., 2.], + [4., 3., 219., 2.], + [7., 0., 255., 3.], + [1., 8., 268., 4.], + [6., 9., 295., 6.]]) + +X = array([[1.43054825, -7.5693489], + [6.95887839, 6.82293382], + [2.87137846, -9.68248579], + [7.87974764, -6.05485803], + [8.24018364, -6.09495602], + [7.39020262, 8.54004355]]) + +linkage_X_centroid = array([[3., 4., 0.36265956, 2.], + [1., 5., 1.77045373, 2.], + [0., 2., 2.55760419, 2.], + [6., 8., 6.43614494, 4.], + [7., 9., 15.17363237, 6.]]) + +linkage_X_median = array([[3., 4., 0.36265956, 2.], + [1., 5., 1.77045373, 2.], + [0., 2., 2.55760419, 2.], + [6., 8., 6.43614494, 4.], + [7., 9., 15.17363237, 6.]]) + +linkage_X_ward = array([[3., 4., 0.36265956, 2.], + [1., 5., 1.77045373, 2.], + [0., 2., 2.55760419, 2.], + [6., 8., 9.10208346, 4.], + [7., 9., 24.7784379, 6.]]) + +# the optimal leaf ordering of linkage_X_ward +linkage_X_ward_olo = array([[4., 3., 0.36265956, 2.], + [5., 1., 1.77045373, 2.], + [2., 0., 2.55760419, 2.], + [6., 8., 9.10208346, 4.], + [7., 9., 24.7784379, 6.]]) + +inconsistent_ytdist = { + 1: array([[138., 0., 1., 0.], + [219., 0., 1., 0.], + [255., 0., 1., 0.], + [268., 0., 1., 0.], + [295., 0., 1., 0.]]), + 2: array([[138., 0., 1., 0.], + [219., 0., 1., 0.], + [237., 25.45584412, 2., 0.70710678], + [261.5, 9.19238816, 2., 0.70710678], + [233.66666667, 83.9424406, 3., 0.7306594]]), + 3: array([[138., 0., 1., 0.], + [219., 0., 1., 0.], + [237., 25.45584412, 2., 0.70710678], + [247.33333333, 25.38372182, 3., 0.81417007], + [239., 69.36377537, 4., 0.80733783]]), + 4: array([[138., 0., 1., 0.], + [219., 0., 1., 0.], + [237., 25.45584412, 2., 0.70710678], + [247.33333333, 25.38372182, 3., 0.81417007], + [235., 60.73302232, 5., 0.98793042]])} + +fcluster_inconsistent = { + 0.8: array([6, 2, 2, 4, 6, 2, 3, 7, 3, 5, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1]), + 1.0: array([6, 2, 2, 4, 6, 2, 3, 7, 3, 5, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1]), + 2.0: array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1])} + +fcluster_distance = { + 0.6: array([4, 4, 4, 4, 4, 4, 4, 5, 4, 4, 6, 6, 6, 6, 6, 7, 6, 6, 6, 6, 3, + 1, 1, 1, 2, 1, 1, 1, 1, 1]), + 1.0: array([2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1]), + 2.0: array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1])} + +fcluster_maxclust = { + 8.0: array([5, 5, 5, 5, 5, 5, 5, 6, 5, 5, 7, 7, 7, 7, 7, 8, 7, 7, 7, 7, 4, + 1, 1, 1, 3, 1, 1, 1, 1, 2]), + 4.0: array([3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 2, + 1, 1, 1, 1, 1, 1, 1, 1, 1]), + 1.0: array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1])} diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_disjoint_set.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_disjoint_set.py new file mode 100644 index 0000000000000000000000000000000000000000..c9518ea1af28df04d811f46ff114d5cf58e594d7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_disjoint_set.py @@ -0,0 +1,202 @@ +import pytest +from pytest import raises as assert_raises +import numpy as np +from scipy.cluster.hierarchy import DisjointSet +import string + + +def generate_random_token(): + k = len(string.ascii_letters) + tokens = list(np.arange(k, dtype=int)) + tokens += list(np.arange(k, dtype=float)) + tokens += list(string.ascii_letters) + tokens += [None for i in range(k)] + tokens = np.array(tokens, dtype=object) + rng = np.random.RandomState(seed=0) + + while 1: + size = rng.randint(1, 3) + element = rng.choice(tokens, size) + if size == 1: + yield element[0] + else: + yield tuple(element) + + +def get_elements(n): + # dict is deterministic without difficulty of comparing numpy ints + elements = {} + for element in generate_random_token(): + if element not in elements: + elements[element] = len(elements) + if len(elements) >= n: + break + return list(elements.keys()) + + +def test_init(): + n = 10 + elements = get_elements(n) + dis = DisjointSet(elements) + assert dis.n_subsets == n + assert list(dis) == elements + + +def test_len(): + n = 10 + elements = get_elements(n) + dis = DisjointSet(elements) + assert len(dis) == n + + dis.add("dummy") + assert len(dis) == n + 1 + + +@pytest.mark.parametrize("n", [10, 100]) +def test_contains(n): + elements = get_elements(n) + dis = DisjointSet(elements) + for x in elements: + assert x in dis + + assert "dummy" not in dis + + +@pytest.mark.parametrize("n", [10, 100]) +def test_add(n): + elements = get_elements(n) + dis1 = DisjointSet(elements) + + dis2 = DisjointSet() + for i, x in enumerate(elements): + dis2.add(x) + assert len(dis2) == i + 1 + + # test idempotency by adding element again + dis2.add(x) + assert len(dis2) == i + 1 + + assert list(dis1) == list(dis2) + + +def test_element_not_present(): + elements = get_elements(n=10) + dis = DisjointSet(elements) + + with assert_raises(KeyError): + dis["dummy"] + + with assert_raises(KeyError): + dis.merge(elements[0], "dummy") + + with assert_raises(KeyError): + dis.connected(elements[0], "dummy") + + +@pytest.mark.parametrize("direction", ["forwards", "backwards"]) +@pytest.mark.parametrize("n", [10, 100]) +def test_linear_union_sequence(n, direction): + elements = get_elements(n) + dis = DisjointSet(elements) + assert elements == list(dis) + + indices = list(range(n - 1)) + if direction == "backwards": + indices = indices[::-1] + + for it, i in enumerate(indices): + assert not dis.connected(elements[i], elements[i + 1]) + assert dis.merge(elements[i], elements[i + 1]) + assert dis.connected(elements[i], elements[i + 1]) + assert dis.n_subsets == n - 1 - it + + roots = [dis[i] for i in elements] + if direction == "forwards": + assert all(elements[0] == r for r in roots) + else: + assert all(elements[-2] == r for r in roots) + assert not dis.merge(elements[0], elements[-1]) + + +@pytest.mark.parametrize("n", [10, 100]) +def test_self_unions(n): + elements = get_elements(n) + dis = DisjointSet(elements) + + for x in elements: + assert dis.connected(x, x) + assert not dis.merge(x, x) + assert dis.connected(x, x) + assert dis.n_subsets == len(elements) + + assert elements == list(dis) + roots = [dis[x] for x in elements] + assert elements == roots + + +@pytest.mark.parametrize("order", ["ab", "ba"]) +@pytest.mark.parametrize("n", [10, 100]) +def test_equal_size_ordering(n, order): + elements = get_elements(n) + dis = DisjointSet(elements) + + rng = np.random.RandomState(seed=0) + indices = np.arange(n) + rng.shuffle(indices) + + for i in range(0, len(indices), 2): + a, b = elements[indices[i]], elements[indices[i + 1]] + if order == "ab": + assert dis.merge(a, b) + else: + assert dis.merge(b, a) + + expected = elements[min(indices[i], indices[i + 1])] + assert dis[a] == expected + assert dis[b] == expected + + +@pytest.mark.parametrize("kmax", [5, 10]) +def test_binary_tree(kmax): + n = 2**kmax + elements = get_elements(n) + dis = DisjointSet(elements) + rng = np.random.RandomState(seed=0) + + for k in 2**np.arange(kmax): + for i in range(0, n, 2 * k): + r1, r2 = rng.randint(0, k, size=2) + a, b = elements[i + r1], elements[i + k + r2] + assert not dis.connected(a, b) + assert dis.merge(a, b) + assert dis.connected(a, b) + + assert elements == list(dis) + roots = [dis[i] for i in elements] + expected_indices = np.arange(n) - np.arange(n) % (2 * k) + expected = [elements[i] for i in expected_indices] + assert roots == expected + + +@pytest.mark.parametrize("n", [10, 100]) +def test_subsets(n): + elements = get_elements(n) + dis = DisjointSet(elements) + + rng = np.random.RandomState(seed=0) + for i, j in rng.randint(0, n, (n, 2)): + x = elements[i] + y = elements[j] + + expected = {element for element in dis if {dis[element]} == {dis[x]}} + assert dis.subset_size(x) == len(dis.subset(x)) + assert expected == dis.subset(x) + + expected = {dis[element]: set() for element in dis} + for element in dis: + expected[dis[element]].add(element) + expected = list(expected.values()) + assert expected == dis.subsets() + + dis.merge(x, y) + assert dis.subset(x) == dis.subset(y) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_hierarchy.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_hierarchy.py new file mode 100644 index 0000000000000000000000000000000000000000..03baf9486ab2909b377b2e326d433fcb4c97d67a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_hierarchy.py @@ -0,0 +1,1300 @@ +# +# Author: Damian Eads +# Date: April 17, 2008 +# +# Copyright (C) 2008 Damian Eads +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# +# 1. Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above +# copyright notice, this list of conditions and the following +# disclaimer in the documentation and/or other materials provided +# with the distribution. +# +# 3. The name of the author may not be used to endorse or promote +# products derived from this software without specific prior +# written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS +# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY +# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +# GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import numpy as np +from numpy.testing import (assert_allclose, assert_equal, assert_array_equal, assert_, + assert_warns) +import pytest +from pytest import raises as assert_raises + +import scipy.cluster.hierarchy +from scipy.cluster.hierarchy import ( + ClusterWarning, linkage, from_mlab_linkage, to_mlab_linkage, + num_obs_linkage, inconsistent, cophenet, fclusterdata, fcluster, + is_isomorphic, single, leaders, + correspond, is_monotonic, maxdists, maxinconsts, maxRstat, + is_valid_linkage, is_valid_im, to_tree, leaves_list, dendrogram, + set_link_color_palette, cut_tree, optimal_leaf_ordering, + _order_cluster_tree, _hierarchy, _LINKAGE_METHODS) +from scipy.spatial.distance import pdist +from scipy.cluster._hierarchy import Heap +from scipy.conftest import array_api_compatible +from scipy._lib._array_api import xp_assert_close, xp_assert_equal + +from threading import Lock + +from . import hierarchy_test_data + + +# Matplotlib is not a scipy dependency but is optionally used in dendrogram, so +# check if it's available +try: + import matplotlib + # and set the backend to be Agg (no gui) + matplotlib.use('Agg') + # before importing pyplot + import matplotlib.pyplot as plt + have_matplotlib = True +except Exception: + have_matplotlib = False + + +pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")] +skip_xp_backends = pytest.mark.skip_xp_backends + + +class TestLinkage: + + @skip_xp_backends(cpu_only=True) + def test_linkage_non_finite_elements_in_distance_matrix(self, xp): + # Tests linkage(Y) where Y contains a non-finite element (e.g. NaN or Inf). + # Exception expected. + y = xp.asarray([xp.nan] + [0.0]*5) + assert_raises(ValueError, linkage, y) + + @skip_xp_backends(cpu_only=True) + def test_linkage_empty_distance_matrix(self, xp): + # Tests linkage(Y) where Y is a 0x4 linkage matrix. Exception expected. + y = xp.zeros((0,)) + assert_raises(ValueError, linkage, y) + + @skip_xp_backends(cpu_only=True) + def test_linkage_tdist(self, xp): + for method in ['single', 'complete', 'average', 'weighted']: + self.check_linkage_tdist(method, xp) + + def check_linkage_tdist(self, method, xp): + # Tests linkage(Y, method) on the tdist data set. + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), method) + expectedZ = getattr(hierarchy_test_data, 'linkage_ytdist_' + method) + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-10) + + @skip_xp_backends(cpu_only=True) + def test_linkage_X(self, xp): + for method in ['centroid', 'median', 'ward']: + self.check_linkage_q(method, xp) + + def check_linkage_q(self, method, xp): + # Tests linkage(Y, method) on the Q data set. + Z = linkage(xp.asarray(hierarchy_test_data.X), method) + expectedZ = getattr(hierarchy_test_data, 'linkage_X_' + method) + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-06) + + y = scipy.spatial.distance.pdist(hierarchy_test_data.X, + metric="euclidean") + Z = linkage(xp.asarray(y), method) + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-06) + + @skip_xp_backends(cpu_only=True) + def test_compare_with_trivial(self, xp): + rng = np.random.RandomState(0) + n = 20 + X = rng.rand(n, 2) + d = pdist(X) + + for method, code in _LINKAGE_METHODS.items(): + Z_trivial = _hierarchy.linkage(d, n, code) + Z = linkage(xp.asarray(d), method) + xp_assert_close(Z, xp.asarray(Z_trivial), rtol=1e-14, atol=1e-15) + + @skip_xp_backends(cpu_only=True) + def test_optimal_leaf_ordering(self, xp): + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), optimal_ordering=True) + expectedZ = getattr(hierarchy_test_data, 'linkage_ytdist_single_olo') + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-10) + + +@skip_xp_backends(cpu_only=True) +class TestLinkageTies: + + _expectations = { + 'single': np.array([[0, 1, 1.41421356, 2], + [2, 3, 1.41421356, 3]]), + 'complete': np.array([[0, 1, 1.41421356, 2], + [2, 3, 2.82842712, 3]]), + 'average': np.array([[0, 1, 1.41421356, 2], + [2, 3, 2.12132034, 3]]), + 'weighted': np.array([[0, 1, 1.41421356, 2], + [2, 3, 2.12132034, 3]]), + 'centroid': np.array([[0, 1, 1.41421356, 2], + [2, 3, 2.12132034, 3]]), + 'median': np.array([[0, 1, 1.41421356, 2], + [2, 3, 2.12132034, 3]]), + 'ward': np.array([[0, 1, 1.41421356, 2], + [2, 3, 2.44948974, 3]]), + } + + def test_linkage_ties(self, xp): + for method in ['single', 'complete', 'average', 'weighted', + 'centroid', 'median', 'ward']: + self.check_linkage_ties(method, xp) + + def check_linkage_ties(self, method, xp): + X = xp.asarray([[-1, -1], [0, 0], [1, 1]]) + Z = linkage(X, method=method) + expectedZ = self._expectations[method] + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-06) + + +@skip_xp_backends(cpu_only=True) +class TestInconsistent: + + def test_inconsistent_tdist(self, xp): + for depth in hierarchy_test_data.inconsistent_ytdist: + self.check_inconsistent_tdist(depth, xp) + + def check_inconsistent_tdist(self, depth, xp): + Z = xp.asarray(hierarchy_test_data.linkage_ytdist_single) + xp_assert_close(inconsistent(Z, depth), + xp.asarray(hierarchy_test_data.inconsistent_ytdist[depth])) + + +@skip_xp_backends(cpu_only=True) +class TestCopheneticDistance: + + def test_linkage_cophenet_tdist_Z(self, xp): + # Tests cophenet(Z) on tdist data set. + expectedM = xp.asarray([268, 295, 255, 255, 295, 295, 268, 268, 295, 295, + 295, 138, 219, 295, 295]) + Z = xp.asarray(hierarchy_test_data.linkage_ytdist_single) + M = cophenet(Z) + xp_assert_close(M, xp.asarray(expectedM, dtype=xp.float64), atol=1e-10) + + def test_linkage_cophenet_tdist_Z_Y(self, xp): + # Tests cophenet(Z, Y) on tdist data set. + Z = xp.asarray(hierarchy_test_data.linkage_ytdist_single) + (c, M) = cophenet(Z, xp.asarray(hierarchy_test_data.ytdist)) + expectedM = xp.asarray([268, 295, 255, 255, 295, 295, 268, 268, 295, 295, + 295, 138, 219, 295, 295], dtype=xp.float64) + expectedc = xp.asarray(0.639931296433393415057366837573, dtype=xp.float64)[()] + xp_assert_close(c, expectedc, atol=1e-10) + xp_assert_close(M, expectedM, atol=1e-10) + + def test_gh_22183(self, xp): + # check for lack of segfault + # (out of bounds memory access) + # and correct interception of + # invalid linkage matrix + arr=[[0.0, 1.0, 1.0, 2.0], + [2.0, 12.0, 1.0, 3.0], + [3.0, 4.0, 1.0, 2.0], + [5.0, 14.0, 1.0, 3.0], + [6.0, 7.0, 1.0, 2.0], + [8.0, 16.0, 1.0, 3.0], + [9.0, 10.0, 1.0, 2.0], + [11.0, 18.0, 1.0, 3.0], + [13.0, 15.0, 2.0, 6.0], + [17.0, 20.0, 2.0, 32.0], + [19.0, 21.0, 2.0, 12.0]] + with pytest.raises(ValueError, match="excessive observations"): + cophenet(xp.asarray(arr)) + + +class TestMLabLinkageConversion: + + def test_mlab_linkage_conversion_empty(self, xp): + # Tests from/to_mlab_linkage on empty linkage array. + X = xp.asarray([], dtype=xp.float64) + xp_assert_equal(from_mlab_linkage(X), X) + xp_assert_equal(to_mlab_linkage(X), X) + + @skip_xp_backends(cpu_only=True) + def test_mlab_linkage_conversion_single_row(self, xp): + # Tests from/to_mlab_linkage on linkage array with single row. + Z = xp.asarray([[0., 1., 3., 2.]]) + Zm = xp.asarray([[1, 2, 3]]) + xp_assert_close(from_mlab_linkage(Zm), xp.asarray(Z, dtype=xp.float64), + rtol=1e-15) + xp_assert_close(to_mlab_linkage(Z), xp.asarray(Zm, dtype=xp.float64), + rtol=1e-15) + + @skip_xp_backends(cpu_only=True) + def test_mlab_linkage_conversion_multiple_rows(self, xp): + # Tests from/to_mlab_linkage on linkage array with multiple rows. + Zm = xp.asarray([[3, 6, 138], [4, 5, 219], + [1, 8, 255], [2, 9, 268], [7, 10, 295]]) + Z = xp.asarray([[2., 5., 138., 2.], + [3., 4., 219., 2.], + [0., 7., 255., 3.], + [1., 8., 268., 4.], + [6., 9., 295., 6.]], + dtype=xp.float64) + xp_assert_close(from_mlab_linkage(Zm), Z, rtol=1e-15) + xp_assert_close(to_mlab_linkage(Z), xp.asarray(Zm, dtype=xp.float64), + rtol=1e-15) + + +@skip_xp_backends(cpu_only=True) +class TestFcluster: + + def test_fclusterdata(self, xp): + for t in hierarchy_test_data.fcluster_inconsistent: + self.check_fclusterdata(t, 'inconsistent', xp) + for t in hierarchy_test_data.fcluster_distance: + self.check_fclusterdata(t, 'distance', xp) + for t in hierarchy_test_data.fcluster_maxclust: + self.check_fclusterdata(t, 'maxclust', xp) + + def check_fclusterdata(self, t, criterion, xp): + # Tests fclusterdata(X, criterion=criterion, t=t) on a random 3-cluster data set + expectedT = xp.asarray(getattr(hierarchy_test_data, 'fcluster_' + criterion)[t]) + X = xp.asarray(hierarchy_test_data.Q_X) + T = fclusterdata(X, criterion=criterion, t=t) + assert_(is_isomorphic(T, expectedT)) + + def test_fcluster(self, xp): + for t in hierarchy_test_data.fcluster_inconsistent: + self.check_fcluster(t, 'inconsistent', xp) + for t in hierarchy_test_data.fcluster_distance: + self.check_fcluster(t, 'distance', xp) + for t in hierarchy_test_data.fcluster_maxclust: + self.check_fcluster(t, 'maxclust', xp) + + def check_fcluster(self, t, criterion, xp): + # Tests fcluster(Z, criterion=criterion, t=t) on a random 3-cluster data set. + expectedT = xp.asarray(getattr(hierarchy_test_data, 'fcluster_' + criterion)[t]) + Z = single(xp.asarray(hierarchy_test_data.Q_X)) + T = fcluster(Z, criterion=criterion, t=t) + assert_(is_isomorphic(T, expectedT)) + + def test_fcluster_monocrit(self, xp): + for t in hierarchy_test_data.fcluster_distance: + self.check_fcluster_monocrit(t, xp) + for t in hierarchy_test_data.fcluster_maxclust: + self.check_fcluster_maxclust_monocrit(t, xp) + + def check_fcluster_monocrit(self, t, xp): + expectedT = xp.asarray(hierarchy_test_data.fcluster_distance[t]) + Z = single(xp.asarray(hierarchy_test_data.Q_X)) + T = fcluster(Z, t, criterion='monocrit', monocrit=maxdists(Z)) + assert_(is_isomorphic(T, expectedT)) + + def check_fcluster_maxclust_monocrit(self, t, xp): + expectedT = xp.asarray(hierarchy_test_data.fcluster_maxclust[t]) + Z = single(xp.asarray(hierarchy_test_data.Q_X)) + T = fcluster(Z, t, criterion='maxclust_monocrit', monocrit=maxdists(Z)) + assert_(is_isomorphic(T, expectedT)) + + def test_fcluster_maxclust_gh_12651(self, xp): + y = xp.asarray([[1], [4], [5]]) + Z = single(y) + assert_array_equal(fcluster(Z, t=1, criterion="maxclust"), + xp.asarray([1, 1, 1])) + assert_array_equal(fcluster(Z, t=2, criterion="maxclust"), + xp.asarray([2, 1, 1])) + assert_array_equal(fcluster(Z, t=3, criterion="maxclust"), + xp.asarray([1, 2, 3])) + assert_array_equal(fcluster(Z, t=5, criterion="maxclust"), + xp.asarray([1, 2, 3])) + + +@skip_xp_backends(cpu_only=True) +class TestLeaders: + + def test_leaders_single(self, xp): + # Tests leaders using a flat clustering generated by single linkage. + X = hierarchy_test_data.Q_X + Y = pdist(X) + Y = xp.asarray(Y) + Z = linkage(Y) + T = fcluster(Z, criterion='maxclust', t=3) + Lright = (xp.asarray([53, 55, 56]), xp.asarray([2, 3, 1])) + T = xp.asarray(T, dtype=xp.int32) + L = leaders(Z, T) + assert_allclose(np.concatenate(L), np.concatenate(Lright), rtol=1e-15) + + +@skip_xp_backends(np_only=True, + reason='`is_isomorphic` only supports NumPy backend') +class TestIsIsomorphic: + + @skip_xp_backends(np_only=True, + reason='array-likes only supported for NumPy backend') + def test_array_like(self, xp): + assert is_isomorphic([1, 1, 1], [2, 2, 2]) + assert is_isomorphic([], []) + + def test_is_isomorphic_1(self, xp): + # Tests is_isomorphic on test case #1 (one flat cluster, different labellings) + a = xp.asarray([1, 1, 1]) + b = xp.asarray([2, 2, 2]) + assert is_isomorphic(a, b) + assert is_isomorphic(b, a) + + def test_is_isomorphic_2(self, xp): + # Tests is_isomorphic on test case #2 (two flat clusters, different labelings) + a = xp.asarray([1, 7, 1]) + b = xp.asarray([2, 3, 2]) + assert is_isomorphic(a, b) + assert is_isomorphic(b, a) + + def test_is_isomorphic_3(self, xp): + # Tests is_isomorphic on test case #3 (no flat clusters) + a = xp.asarray([]) + b = xp.asarray([]) + assert is_isomorphic(a, b) + + def test_is_isomorphic_4A(self, xp): + # Tests is_isomorphic on test case #4A + # (3 flat clusters, different labelings, isomorphic) + a = xp.asarray([1, 2, 3]) + b = xp.asarray([1, 3, 2]) + assert is_isomorphic(a, b) + assert is_isomorphic(b, a) + + def test_is_isomorphic_4B(self, xp): + # Tests is_isomorphic on test case #4B + # (3 flat clusters, different labelings, nonisomorphic) + a = xp.asarray([1, 2, 3, 3]) + b = xp.asarray([1, 3, 2, 3]) + assert is_isomorphic(a, b) is False + assert is_isomorphic(b, a) is False + + def test_is_isomorphic_4C(self, xp): + # Tests is_isomorphic on test case #4C + # (3 flat clusters, different labelings, isomorphic) + a = xp.asarray([7, 2, 3]) + b = xp.asarray([6, 3, 2]) + assert is_isomorphic(a, b) + assert is_isomorphic(b, a) + + def test_is_isomorphic_5(self, xp): + # Tests is_isomorphic on test case #5 (1000 observations, 2/3/5 random + # clusters, random permutation of the labeling). + for nc in [2, 3, 5]: + self.help_is_isomorphic_randperm(1000, nc, xp=xp) + + def test_is_isomorphic_6(self, xp): + # Tests is_isomorphic on test case #5A (1000 observations, 2/3/5 random + # clusters, random permutation of the labeling, slightly + # nonisomorphic.) + for nc in [2, 3, 5]: + self.help_is_isomorphic_randperm(1000, nc, True, 5, xp=xp) + + def test_is_isomorphic_7(self, xp): + # Regression test for gh-6271 + a = xp.asarray([1, 2, 3]) + b = xp.asarray([1, 1, 1]) + assert not is_isomorphic(a, b) + + def help_is_isomorphic_randperm(self, nobs, nclusters, noniso=False, nerrors=0, + *, xp): + for k in range(3): + a = (np.random.rand(nobs) * nclusters).astype(int) + b = np.zeros(a.size, dtype=int) + P = np.random.permutation(nclusters) + for i in range(0, a.shape[0]): + b[i] = P[a[i]] + if noniso: + Q = np.random.permutation(nobs) + b[Q[0:nerrors]] += 1 + b[Q[0:nerrors]] %= nclusters + a = xp.asarray(a) + b = xp.asarray(b) + assert is_isomorphic(a, b) == (not noniso) + assert is_isomorphic(b, a) == (not noniso) + + +@skip_xp_backends(cpu_only=True) +class TestIsValidLinkage: + + def test_is_valid_linkage_various_size(self, xp): + for nrow, ncol, valid in [(2, 5, False), (2, 3, False), + (1, 4, True), (2, 4, True)]: + self.check_is_valid_linkage_various_size(nrow, ncol, valid, xp) + + def check_is_valid_linkage_various_size(self, nrow, ncol, valid, xp): + # Tests is_valid_linkage(Z) with linkage matrices of various sizes + Z = xp.asarray([[0, 1, 3.0, 2, 5], + [3, 2, 4.0, 3, 3]], dtype=xp.float64) + Z = Z[:nrow, :ncol] + assert_(is_valid_linkage(Z) == valid) + if not valid: + assert_raises(ValueError, is_valid_linkage, Z, throw=True) + + def test_is_valid_linkage_int_type(self, xp): + # Tests is_valid_linkage(Z) with integer type. + Z = xp.asarray([[0, 1, 3.0, 2], + [3, 2, 4.0, 3]], dtype=xp.int64) + assert_(is_valid_linkage(Z) is False) + assert_raises(TypeError, is_valid_linkage, Z, throw=True) + + def test_is_valid_linkage_empty(self, xp): + # Tests is_valid_linkage(Z) with empty linkage. + Z = xp.zeros((0, 4), dtype=xp.float64) + assert_(is_valid_linkage(Z) is False) + assert_raises(ValueError, is_valid_linkage, Z, throw=True) + + def test_is_valid_linkage_4_and_up(self, xp): + # Tests is_valid_linkage(Z) on linkage on observation sets between + # sizes 4 and 15 (step size 3). + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + assert_(is_valid_linkage(Z) is True) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_is_valid_linkage_4_and_up_neg_index_left(self, xp): + # Tests is_valid_linkage(Z) on linkage on observation sets between + # sizes 4 and 15 (step size 3) with negative indices (left). + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + Z[i//2,0] = -2 + assert_(is_valid_linkage(Z) is False) + assert_raises(ValueError, is_valid_linkage, Z, throw=True) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_is_valid_linkage_4_and_up_neg_index_right(self, xp): + # Tests is_valid_linkage(Z) on linkage on observation sets between + # sizes 4 and 15 (step size 3) with negative indices (right). + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + Z[i//2,1] = -2 + assert_(is_valid_linkage(Z) is False) + assert_raises(ValueError, is_valid_linkage, Z, throw=True) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_is_valid_linkage_4_and_up_neg_dist(self, xp): + # Tests is_valid_linkage(Z) on linkage on observation sets between + # sizes 4 and 15 (step size 3) with negative distances. + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + Z[i//2,2] = -0.5 + assert_(is_valid_linkage(Z) is False) + assert_raises(ValueError, is_valid_linkage, Z, throw=True) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_is_valid_linkage_4_and_up_neg_counts(self, xp): + # Tests is_valid_linkage(Z) on linkage on observation sets between + # sizes 4 and 15 (step size 3) with negative counts. + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + Z[i//2,3] = -2 + assert_(is_valid_linkage(Z) is False) + assert_raises(ValueError, is_valid_linkage, Z, throw=True) + + +@skip_xp_backends(cpu_only=True) +class TestIsValidInconsistent: + + def test_is_valid_im_int_type(self, xp): + # Tests is_valid_im(R) with integer type. + R = xp.asarray([[0, 1, 3.0, 2], + [3, 2, 4.0, 3]], dtype=xp.int64) + assert_(is_valid_im(R) is False) + assert_raises(TypeError, is_valid_im, R, throw=True) + + def test_is_valid_im_various_size(self, xp): + for nrow, ncol, valid in [(2, 5, False), (2, 3, False), + (1, 4, True), (2, 4, True)]: + self.check_is_valid_im_various_size(nrow, ncol, valid, xp) + + def check_is_valid_im_various_size(self, nrow, ncol, valid, xp): + # Tests is_valid_im(R) with linkage matrices of various sizes + R = xp.asarray([[0, 1, 3.0, 2, 5], + [3, 2, 4.0, 3, 3]], dtype=xp.float64) + R = R[:nrow, :ncol] + assert_(is_valid_im(R) == valid) + if not valid: + assert_raises(ValueError, is_valid_im, R, throw=True) + + def test_is_valid_im_empty(self, xp): + # Tests is_valid_im(R) with empty inconsistency matrix. + R = xp.zeros((0, 4), dtype=xp.float64) + assert_(is_valid_im(R) is False) + assert_raises(ValueError, is_valid_im, R, throw=True) + + def test_is_valid_im_4_and_up(self, xp): + # Tests is_valid_im(R) on im on observation sets between sizes 4 and 15 + # (step size 3). + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + R = inconsistent(Z) + assert_(is_valid_im(R) is True) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment') + def test_is_valid_im_4_and_up_neg_index_left(self, xp): + # Tests is_valid_im(R) on im on observation sets between sizes 4 and 15 + # (step size 3) with negative link height means. + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + R = inconsistent(Z) + R[i//2,0] = -2.0 + assert_(is_valid_im(R) is False) + assert_raises(ValueError, is_valid_im, R, throw=True) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment') + def test_is_valid_im_4_and_up_neg_index_right(self, xp): + # Tests is_valid_im(R) on im on observation sets between sizes 4 and 15 + # (step size 3) with negative link height standard deviations. + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + R = inconsistent(Z) + R[i//2,1] = -2.0 + assert_(is_valid_im(R) is False) + assert_raises(ValueError, is_valid_im, R, throw=True) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment') + def test_is_valid_im_4_and_up_neg_dist(self, xp): + # Tests is_valid_im(R) on im on observation sets between sizes 4 and 15 + # (step size 3) with negative link counts. + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + R = inconsistent(Z) + R[i//2,2] = -0.5 + assert_(is_valid_im(R) is False) + assert_raises(ValueError, is_valid_im, R, throw=True) + + +class TestNumObsLinkage: + + @skip_xp_backends(cpu_only=True) + def test_num_obs_linkage_empty(self, xp): + # Tests num_obs_linkage(Z) with empty linkage. + Z = xp.zeros((0, 4), dtype=xp.float64) + assert_raises(ValueError, num_obs_linkage, Z) + + def test_num_obs_linkage_1x4(self, xp): + # Tests num_obs_linkage(Z) on linkage over 2 observations. + Z = xp.asarray([[0, 1, 3.0, 2]], dtype=xp.float64) + assert_equal(num_obs_linkage(Z), 2) + + def test_num_obs_linkage_2x4(self, xp): + # Tests num_obs_linkage(Z) on linkage over 3 observations. + Z = xp.asarray([[0, 1, 3.0, 2], + [3, 2, 4.0, 3]], dtype=xp.float64) + assert_equal(num_obs_linkage(Z), 3) + + @skip_xp_backends(cpu_only=True) + def test_num_obs_linkage_4_and_up(self, xp): + # Tests num_obs_linkage(Z) on linkage on observation sets between sizes + # 4 and 15 (step size 3). + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + assert_equal(num_obs_linkage(Z), i) + + +@skip_xp_backends(cpu_only=True) +class TestLeavesList: + + def test_leaves_list_1x4(self, xp): + # Tests leaves_list(Z) on a 1x4 linkage. + Z = xp.asarray([[0, 1, 3.0, 2]], dtype=xp.float64) + to_tree(Z) + assert_allclose(leaves_list(Z), [0, 1], rtol=1e-15) + + def test_leaves_list_2x4(self, xp): + # Tests leaves_list(Z) on a 2x4 linkage. + Z = xp.asarray([[0, 1, 3.0, 2], + [3, 2, 4.0, 3]], dtype=xp.float64) + to_tree(Z) + assert_allclose(leaves_list(Z), [0, 1, 2], rtol=1e-15) + + def test_leaves_list_Q(self, xp): + for method in ['single', 'complete', 'average', 'weighted', 'centroid', + 'median', 'ward']: + self.check_leaves_list_Q(method, xp) + + def check_leaves_list_Q(self, method, xp): + # Tests leaves_list(Z) on the Q data set + X = xp.asarray(hierarchy_test_data.Q_X) + Z = linkage(X, method) + node = to_tree(Z) + assert_allclose(node.pre_order(), leaves_list(Z), rtol=1e-15) + + def test_Q_subtree_pre_order(self, xp): + # Tests that pre_order() works when called on sub-trees. + X = xp.asarray(hierarchy_test_data.Q_X) + Z = linkage(X, 'single') + node = to_tree(Z) + assert_allclose(node.pre_order(), (node.get_left().pre_order() + + node.get_right().pre_order()), + rtol=1e-15) + + +@skip_xp_backends(cpu_only=True) +class TestCorrespond: + + def test_correspond_empty(self, xp): + # Tests correspond(Z, y) with empty linkage and condensed distance matrix. + y = xp.zeros((0,), dtype=xp.float64) + Z = xp.zeros((0,4), dtype=xp.float64) + assert_raises(ValueError, correspond, Z, y) + + def test_correspond_2_and_up(self, xp): + # Tests correspond(Z, y) on linkage and CDMs over observation sets of + # different sizes. + for i in range(2, 4): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + assert_(correspond(Z, y)) + for i in range(4, 15, 3): + y = np.random.rand(i*(i-1)//2) + y = xp.asarray(y) + Z = linkage(y) + assert_(correspond(Z, y)) + + def test_correspond_4_and_up(self, xp): + # Tests correspond(Z, y) on linkage and CDMs over observation sets of + # different sizes. Correspondence should be false. + for (i, j) in (list(zip(list(range(2, 4)), list(range(3, 5)))) + + list(zip(list(range(3, 5)), list(range(2, 4))))): + y = np.random.rand(i*(i-1)//2) + y2 = np.random.rand(j*(j-1)//2) + y = xp.asarray(y) + y2 = xp.asarray(y2) + Z = linkage(y) + Z2 = linkage(y2) + assert not correspond(Z, y2) + assert not correspond(Z2, y) + + def test_correspond_4_and_up_2(self, xp): + # Tests correspond(Z, y) on linkage and CDMs over observation sets of + # different sizes. Correspondence should be false. + for (i, j) in (list(zip(list(range(2, 7)), list(range(16, 21)))) + + list(zip(list(range(2, 7)), list(range(16, 21))))): + y = np.random.rand(i*(i-1)//2) + y2 = np.random.rand(j*(j-1)//2) + y = xp.asarray(y) + y2 = xp.asarray(y2) + Z = linkage(y) + Z2 = linkage(y2) + assert not correspond(Z, y2) + assert not correspond(Z2, y) + + def test_num_obs_linkage_multi_matrix(self, xp): + # Tests num_obs_linkage with observation matrices of multiple sizes. + for n in range(2, 10): + X = np.random.rand(n, 4) + Y = pdist(X) + Y = xp.asarray(Y) + Z = linkage(Y) + assert_equal(num_obs_linkage(Z), n) + + +@skip_xp_backends(cpu_only=True) +class TestIsMonotonic: + + def test_is_monotonic_empty(self, xp): + # Tests is_monotonic(Z) on an empty linkage. + Z = xp.zeros((0, 4), dtype=xp.float64) + assert_raises(ValueError, is_monotonic, Z) + + def test_is_monotonic_1x4(self, xp): + # Tests is_monotonic(Z) on 1x4 linkage. Expecting True. + Z = xp.asarray([[0, 1, 0.3, 2]], dtype=xp.float64) + assert is_monotonic(Z) + + def test_is_monotonic_2x4_T(self, xp): + # Tests is_monotonic(Z) on 2x4 linkage. Expecting True. + Z = xp.asarray([[0, 1, 0.3, 2], + [2, 3, 0.4, 3]], dtype=xp.float64) + assert is_monotonic(Z) + + def test_is_monotonic_2x4_F(self, xp): + # Tests is_monotonic(Z) on 2x4 linkage. Expecting False. + Z = xp.asarray([[0, 1, 0.4, 2], + [2, 3, 0.3, 3]], dtype=xp.float64) + assert not is_monotonic(Z) + + def test_is_monotonic_3x4_T(self, xp): + # Tests is_monotonic(Z) on 3x4 linkage. Expecting True. + Z = xp.asarray([[0, 1, 0.3, 2], + [2, 3, 0.4, 2], + [4, 5, 0.6, 4]], dtype=xp.float64) + assert is_monotonic(Z) + + def test_is_monotonic_3x4_F1(self, xp): + # Tests is_monotonic(Z) on 3x4 linkage (case 1). Expecting False. + Z = xp.asarray([[0, 1, 0.3, 2], + [2, 3, 0.2, 2], + [4, 5, 0.6, 4]], dtype=xp.float64) + assert not is_monotonic(Z) + + def test_is_monotonic_3x4_F2(self, xp): + # Tests is_monotonic(Z) on 3x4 linkage (case 2). Expecting False. + Z = xp.asarray([[0, 1, 0.8, 2], + [2, 3, 0.4, 2], + [4, 5, 0.6, 4]], dtype=xp.float64) + assert not is_monotonic(Z) + + def test_is_monotonic_3x4_F3(self, xp): + # Tests is_monotonic(Z) on 3x4 linkage (case 3). Expecting False + Z = xp.asarray([[0, 1, 0.3, 2], + [2, 3, 0.4, 2], + [4, 5, 0.2, 4]], dtype=xp.float64) + assert not is_monotonic(Z) + + def test_is_monotonic_tdist_linkage1(self, xp): + # Tests is_monotonic(Z) on clustering generated by single linkage on + # tdist data set. Expecting True. + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + assert is_monotonic(Z) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment') + def test_is_monotonic_tdist_linkage2(self, xp): + # Tests is_monotonic(Z) on clustering generated by single linkage on + # tdist data set. Perturbing. Expecting False. + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + Z[2,2] = 0.0 + assert not is_monotonic(Z) + + def test_is_monotonic_Q_linkage(self, xp): + # Tests is_monotonic(Z) on clustering generated by single linkage on + # Q data set. Expecting True. + X = xp.asarray(hierarchy_test_data.Q_X) + Z = linkage(X, 'single') + assert is_monotonic(Z) + + +@skip_xp_backends(cpu_only=True) +class TestMaxDists: + + def test_maxdists_empty_linkage(self, xp): + # Tests maxdists(Z) on empty linkage. Expecting exception. + Z = xp.zeros((0, 4), dtype=xp.float64) + assert_raises(ValueError, maxdists, Z) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment') + def test_maxdists_one_cluster_linkage(self, xp): + # Tests maxdists(Z) on linkage with one cluster. + Z = xp.asarray([[0, 1, 0.3, 4]], dtype=xp.float64) + MD = maxdists(Z) + expectedMD = calculate_maximum_distances(Z, xp) + xp_assert_close(MD, expectedMD, atol=1e-15) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment') + def test_maxdists_Q_linkage(self, xp): + for method in ['single', 'complete', 'ward', 'centroid', 'median']: + self.check_maxdists_Q_linkage(method, xp) + + def check_maxdists_Q_linkage(self, method, xp): + # Tests maxdists(Z) on the Q data set + X = xp.asarray(hierarchy_test_data.Q_X) + Z = linkage(X, method) + MD = maxdists(Z) + expectedMD = calculate_maximum_distances(Z, xp) + xp_assert_close(MD, expectedMD, atol=1e-15) + + +class TestMaxInconsts: + + @skip_xp_backends(cpu_only=True) + def test_maxinconsts_empty_linkage(self, xp): + # Tests maxinconsts(Z, R) on empty linkage. Expecting exception. + Z = xp.zeros((0, 4), dtype=xp.float64) + R = xp.zeros((0, 4), dtype=xp.float64) + assert_raises(ValueError, maxinconsts, Z, R) + + def test_maxinconsts_difrow_linkage(self, xp): + # Tests maxinconsts(Z, R) on linkage and inconsistency matrices with + # different numbers of clusters. Expecting exception. + Z = xp.asarray([[0, 1, 0.3, 4]], dtype=xp.float64) + R = np.random.rand(2, 4) + R = xp.asarray(R) + assert_raises(ValueError, maxinconsts, Z, R) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment', + cpu_only=True) + def test_maxinconsts_one_cluster_linkage(self, xp): + # Tests maxinconsts(Z, R) on linkage with one cluster. + Z = xp.asarray([[0, 1, 0.3, 4]], dtype=xp.float64) + R = xp.asarray([[0, 0, 0, 0.3]], dtype=xp.float64) + MD = maxinconsts(Z, R) + expectedMD = calculate_maximum_inconsistencies(Z, R, xp=xp) + xp_assert_close(MD, expectedMD, atol=1e-15) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment', + cpu_only=True) + def test_maxinconsts_Q_linkage(self, xp): + for method in ['single', 'complete', 'ward', 'centroid', 'median']: + self.check_maxinconsts_Q_linkage(method, xp) + + def check_maxinconsts_Q_linkage(self, method, xp): + # Tests maxinconsts(Z, R) on the Q data set + X = xp.asarray(hierarchy_test_data.Q_X) + Z = linkage(X, method) + R = inconsistent(Z) + MD = maxinconsts(Z, R) + expectedMD = calculate_maximum_inconsistencies(Z, R, xp=xp) + xp_assert_close(MD, expectedMD, atol=1e-15) + + +class TestMaxRStat: + + def test_maxRstat_invalid_index(self, xp): + for i in [3.3, -1, 4]: + self.check_maxRstat_invalid_index(i, xp) + + def check_maxRstat_invalid_index(self, i, xp): + # Tests maxRstat(Z, R, i). Expecting exception. + Z = xp.asarray([[0, 1, 0.3, 4]], dtype=xp.float64) + R = xp.asarray([[0, 0, 0, 0.3]], dtype=xp.float64) + if isinstance(i, int): + assert_raises(ValueError, maxRstat, Z, R, i) + else: + assert_raises(TypeError, maxRstat, Z, R, i) + + @skip_xp_backends(cpu_only=True) + def test_maxRstat_empty_linkage(self, xp): + for i in range(4): + self.check_maxRstat_empty_linkage(i, xp) + + def check_maxRstat_empty_linkage(self, i, xp): + # Tests maxRstat(Z, R, i) on empty linkage. Expecting exception. + Z = xp.zeros((0, 4), dtype=xp.float64) + R = xp.zeros((0, 4), dtype=xp.float64) + assert_raises(ValueError, maxRstat, Z, R, i) + + def test_maxRstat_difrow_linkage(self, xp): + for i in range(4): + self.check_maxRstat_difrow_linkage(i, xp) + + def check_maxRstat_difrow_linkage(self, i, xp): + # Tests maxRstat(Z, R, i) on linkage and inconsistency matrices with + # different numbers of clusters. Expecting exception. + Z = xp.asarray([[0, 1, 0.3, 4]], dtype=xp.float64) + R = np.random.rand(2, 4) + R = xp.asarray(R) + assert_raises(ValueError, maxRstat, Z, R, i) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment', + cpu_only=True) + def test_maxRstat_one_cluster_linkage(self, xp): + for i in range(4): + self.check_maxRstat_one_cluster_linkage(i, xp) + + def check_maxRstat_one_cluster_linkage(self, i, xp): + # Tests maxRstat(Z, R, i) on linkage with one cluster. + Z = xp.asarray([[0, 1, 0.3, 4]], dtype=xp.float64) + R = xp.asarray([[0, 0, 0, 0.3]], dtype=xp.float64) + MD = maxRstat(Z, R, 1) + expectedMD = calculate_maximum_inconsistencies(Z, R, 1, xp) + xp_assert_close(MD, expectedMD, atol=1e-15) + + @skip_xp_backends('jax.numpy', reason='jax arrays do not support item assignment', + cpu_only=True) + def test_maxRstat_Q_linkage(self, xp): + for method in ['single', 'complete', 'ward', 'centroid', 'median']: + for i in range(4): + self.check_maxRstat_Q_linkage(method, i, xp) + + def check_maxRstat_Q_linkage(self, method, i, xp): + # Tests maxRstat(Z, R, i) on the Q data set + X = xp.asarray(hierarchy_test_data.Q_X) + Z = linkage(X, method) + R = inconsistent(Z) + MD = maxRstat(Z, R, 1) + expectedMD = calculate_maximum_inconsistencies(Z, R, 1, xp) + xp_assert_close(MD, expectedMD, atol=1e-15) + + +@skip_xp_backends(cpu_only=True) +class TestDendrogram: + + def test_dendrogram_single_linkage_tdist(self, xp): + # Tests dendrogram calculation on single linkage of the tdist data set. + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + R = dendrogram(Z, no_plot=True) + leaves = R["leaves"] + assert_equal(leaves, [2, 5, 1, 0, 3, 4]) + + def test_valid_orientation(self, xp): + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + assert_raises(ValueError, dendrogram, Z, orientation="foo") + + def test_labels_as_array_or_list(self, xp): + # test for gh-12418 + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + labels = [1, 3, 2, 6, 4, 5] + result1 = dendrogram(Z, labels=xp.asarray(labels), no_plot=True) + result2 = dendrogram(Z, labels=labels, no_plot=True) + assert result1 == result2 + + @pytest.mark.skipif(not have_matplotlib, reason="no matplotlib") + def test_valid_label_size(self, xp): + link = xp.asarray([ + [0, 1, 1.0, 4], + [2, 3, 1.0, 5], + [4, 5, 2.0, 6], + ]) + plt.figure() + with pytest.raises(ValueError) as exc_info: + dendrogram(link, labels=list(range(100))) + assert "Dimensions of Z and labels must be consistent."\ + in str(exc_info.value) + + with pytest.raises( + ValueError, + match="Dimensions of Z and labels must be consistent."): + dendrogram(link, labels=[]) + + plt.close() + + @skip_xp_backends('torch', + reason='MPL 3.9.2 & torch DeprecationWarning from __array_wrap__' + ' and NumPy 2.0' + ) + @pytest.mark.skipif(not have_matplotlib, reason="no matplotlib") + def test_dendrogram_plot(self, xp): + for orientation in ['top', 'bottom', 'left', 'right']: + self.check_dendrogram_plot(orientation, xp) + + def check_dendrogram_plot(self, orientation, xp): + # Tests dendrogram plotting. + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + expected = {'color_list': ['C1', 'C0', 'C0', 'C0', 'C0'], + 'dcoord': [[0.0, 138.0, 138.0, 0.0], + [0.0, 219.0, 219.0, 0.0], + [0.0, 255.0, 255.0, 219.0], + [0.0, 268.0, 268.0, 255.0], + [138.0, 295.0, 295.0, 268.0]], + 'icoord': [[5.0, 5.0, 15.0, 15.0], + [45.0, 45.0, 55.0, 55.0], + [35.0, 35.0, 50.0, 50.0], + [25.0, 25.0, 42.5, 42.5], + [10.0, 10.0, 33.75, 33.75]], + 'ivl': ['2', '5', '1', '0', '3', '4'], + 'leaves': [2, 5, 1, 0, 3, 4], + 'leaves_color_list': ['C1', 'C1', 'C0', 'C0', 'C0', 'C0'], + } + + fig = plt.figure() + ax = fig.add_subplot(221) + + # test that dendrogram accepts ax keyword + R1 = dendrogram(Z, ax=ax, orientation=orientation) + R1['dcoord'] = np.asarray(R1['dcoord']) + assert_equal(R1, expected) + + # test that dendrogram accepts and handle the leaf_font_size and + # leaf_rotation keywords + dendrogram(Z, ax=ax, orientation=orientation, + leaf_font_size=20, leaf_rotation=90) + testlabel = ( + ax.get_xticklabels()[0] + if orientation in ['top', 'bottom'] + else ax.get_yticklabels()[0] + ) + assert_equal(testlabel.get_rotation(), 90) + assert_equal(testlabel.get_size(), 20) + dendrogram(Z, ax=ax, orientation=orientation, + leaf_rotation=90) + testlabel = ( + ax.get_xticklabels()[0] + if orientation in ['top', 'bottom'] + else ax.get_yticklabels()[0] + ) + assert_equal(testlabel.get_rotation(), 90) + dendrogram(Z, ax=ax, orientation=orientation, + leaf_font_size=20) + testlabel = ( + ax.get_xticklabels()[0] + if orientation in ['top', 'bottom'] + else ax.get_yticklabels()[0] + ) + assert_equal(testlabel.get_size(), 20) + plt.close() + + # test plotting to gca (will import pylab) + R2 = dendrogram(Z, orientation=orientation) + plt.close() + R2['dcoord'] = np.asarray(R2['dcoord']) + assert_equal(R2, expected) + + @skip_xp_backends('torch', + reason='MPL 3.9.2 & torch DeprecationWarning from __array_wrap__' + ' and NumPy 2.0' + ) + @pytest.mark.skipif(not have_matplotlib, reason="no matplotlib") + def test_dendrogram_truncate_mode(self, xp): + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + + R = dendrogram(Z, 2, 'lastp', show_contracted=True) + plt.close() + R['dcoord'] = np.asarray(R['dcoord']) + assert_equal(R, {'color_list': ['C0'], + 'dcoord': [[0.0, 295.0, 295.0, 0.0]], + 'icoord': [[5.0, 5.0, 15.0, 15.0]], + 'ivl': ['(2)', '(4)'], + 'leaves': [6, 9], + 'leaves_color_list': ['C0', 'C0'], + }) + + R = dendrogram(Z, 2, 'mtica', show_contracted=True) + plt.close() + R['dcoord'] = np.asarray(R['dcoord']) + assert_equal(R, {'color_list': ['C1', 'C0', 'C0', 'C0'], + 'dcoord': [[0.0, 138.0, 138.0, 0.0], + [0.0, 255.0, 255.0, 0.0], + [0.0, 268.0, 268.0, 255.0], + [138.0, 295.0, 295.0, 268.0]], + 'icoord': [[5.0, 5.0, 15.0, 15.0], + [35.0, 35.0, 45.0, 45.0], + [25.0, 25.0, 40.0, 40.0], + [10.0, 10.0, 32.5, 32.5]], + 'ivl': ['2', '5', '1', '0', '(2)'], + 'leaves': [2, 5, 1, 0, 7], + 'leaves_color_list': ['C1', 'C1', 'C0', 'C0', 'C0'], + }) + + @pytest.fixture + def dendrogram_lock(self): + return Lock() + + def test_dendrogram_colors(self, xp, dendrogram_lock): + # Tests dendrogram plots with alternate colors + Z = linkage(xp.asarray(hierarchy_test_data.ytdist), 'single') + + with dendrogram_lock: + # Global color palette might be changed concurrently + set_link_color_palette(['c', 'm', 'y', 'k']) + R = dendrogram(Z, no_plot=True, + above_threshold_color='g', color_threshold=250) + set_link_color_palette(['g', 'r', 'c', 'm', 'y', 'k']) + + color_list = R['color_list'] + assert_equal(color_list, ['c', 'm', 'g', 'g', 'g']) + + # reset color palette (global list) + set_link_color_palette(None) + + def test_dendrogram_leaf_colors_zero_dist(self, xp): + # tests that the colors of leafs are correct for tree + # with two identical points + x = xp.asarray([[1, 0, 0], + [0, 0, 1], + [0, 2, 0], + [0, 0, 1], + [0, 1, 0], + [0, 1, 0]]) + z = linkage(x, "single") + d = dendrogram(z, no_plot=True) + exp_colors = ['C0', 'C1', 'C1', 'C0', 'C2', 'C2'] + colors = d["leaves_color_list"] + assert_equal(colors, exp_colors) + + def test_dendrogram_leaf_colors(self, xp): + # tests that the colors are correct for a tree + # with two near points ((0, 0, 1.1) and (0, 0, 1)) + x = xp.asarray([[1, 0, 0], + [0, 0, 1.1], + [0, 2, 0], + [0, 0, 1], + [0, 1, 0], + [0, 1, 0]]) + z = linkage(x, "single") + d = dendrogram(z, no_plot=True) + exp_colors = ['C0', 'C1', 'C1', 'C0', 'C2', 'C2'] + colors = d["leaves_color_list"] + assert_equal(colors, exp_colors) + + +def calculate_maximum_distances(Z, xp): + # Used for testing correctness of maxdists. + n = Z.shape[0] + 1 + B = xp.zeros((n-1,), dtype=Z.dtype) + q = xp.zeros((3,)) + for i in range(0, n - 1): + q[:] = 0.0 + left = Z[i, 0] + right = Z[i, 1] + if left >= n: + q[0] = B[xp.asarray(left, dtype=xp.int64) - n] + if right >= n: + q[1] = B[xp.asarray(right, dtype=xp.int64) - n] + q[2] = Z[i, 2] + B[i] = xp.max(q) + return B + + +def calculate_maximum_inconsistencies(Z, R, k=3, xp=np): + # Used for testing correctness of maxinconsts. + n = Z.shape[0] + 1 + dtype = xp.result_type(Z, R) + B = xp.zeros((n-1,), dtype=dtype) + q = xp.zeros((3,)) + for i in range(0, n - 1): + q[:] = 0.0 + left = Z[i, 0] + right = Z[i, 1] + if left >= n: + q[0] = B[xp.asarray(left, dtype=xp.int64) - n] + if right >= n: + q[1] = B[xp.asarray(right, dtype=xp.int64) - n] + q[2] = R[i, k] + B[i] = xp.max(q) + return B + + +@pytest.mark.thread_unsafe +@skip_xp_backends(cpu_only=True) +def test_unsupported_uncondensed_distance_matrix_linkage_warning(xp): + assert_warns(ClusterWarning, linkage, xp.asarray([[0, 1], [1, 0]])) + + +def test_euclidean_linkage_value_error(xp): + for method in scipy.cluster.hierarchy._EUCLIDEAN_METHODS: + assert_raises(ValueError, linkage, xp.asarray([[1, 1], [1, 1]]), + method=method, metric='cityblock') + + +@skip_xp_backends(cpu_only=True) +def test_2x2_linkage(xp): + Z1 = linkage(xp.asarray([1]), method='single', metric='euclidean') + Z2 = linkage(xp.asarray([[0, 1], [0, 0]]), method='single', metric='euclidean') + xp_assert_close(Z1, Z2, rtol=1e-15) + + +@skip_xp_backends(cpu_only=True) +def test_node_compare(xp): + np.random.seed(23) + nobs = 50 + X = np.random.randn(nobs, 4) + X = xp.asarray(X) + Z = scipy.cluster.hierarchy.ward(X) + tree = to_tree(Z) + assert_(tree > tree.get_left()) + assert_(tree.get_right() > tree.get_left()) + assert_(tree.get_right() == tree.get_right()) + assert_(tree.get_right() != tree.get_left()) + + +@skip_xp_backends(np_only=True, reason='`cut_tree` uses non-standard indexing') +def test_cut_tree(xp): + np.random.seed(23) + nobs = 50 + X = np.random.randn(nobs, 4) + X = xp.asarray(X) + Z = scipy.cluster.hierarchy.ward(X) + cutree = cut_tree(Z) + + # cutree.dtype varies between int32 and int64 over platforms + xp_assert_close(cutree[:, 0], xp.arange(nobs), rtol=1e-15, check_dtype=False) + xp_assert_close(cutree[:, -1], xp.zeros(nobs), rtol=1e-15, check_dtype=False) + assert_equal(np.asarray(cutree).max(0), np.arange(nobs - 1, -1, -1)) + + xp_assert_close(cutree[:, [-5]], cut_tree(Z, n_clusters=5), rtol=1e-15) + xp_assert_close(cutree[:, [-5, -10]], cut_tree(Z, n_clusters=[5, 10]), rtol=1e-15) + xp_assert_close(cutree[:, [-10, -5]], cut_tree(Z, n_clusters=[10, 5]), rtol=1e-15) + + nodes = _order_cluster_tree(Z) + heights = xp.asarray([node.dist for node in nodes]) + + xp_assert_close(cutree[:, np.searchsorted(heights, [5])], + cut_tree(Z, height=5), rtol=1e-15) + xp_assert_close(cutree[:, np.searchsorted(heights, [5, 10])], + cut_tree(Z, height=[5, 10]), rtol=1e-15) + xp_assert_close(cutree[:, np.searchsorted(heights, [10, 5])], + cut_tree(Z, height=[10, 5]), rtol=1e-15) + + +@skip_xp_backends(cpu_only=True) +def test_optimal_leaf_ordering(xp): + # test with the distance vector y + Z = optimal_leaf_ordering(linkage(xp.asarray(hierarchy_test_data.ytdist)), + xp.asarray(hierarchy_test_data.ytdist)) + expectedZ = hierarchy_test_data.linkage_ytdist_single_olo + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-10) + + # test with the observation matrix X + Z = optimal_leaf_ordering(linkage(xp.asarray(hierarchy_test_data.X), 'ward'), + xp.asarray(hierarchy_test_data.X)) + expectedZ = hierarchy_test_data.linkage_X_ward_olo + xp_assert_close(Z, xp.asarray(expectedZ), atol=1e-06) + + +@skip_xp_backends(np_only=True, reason='`Heap` only supports NumPy backend') +def test_Heap(xp): + values = xp.asarray([2, -1, 0, -1.5, 3]) + heap = Heap(values) + + pair = heap.get_min() + assert_equal(pair['key'], 3) + assert_equal(pair['value'], -1.5) + + heap.remove_min() + pair = heap.get_min() + assert_equal(pair['key'], 1) + assert_equal(pair['value'], -1) + + heap.change_value(1, 2.5) + pair = heap.get_min() + assert_equal(pair['key'], 2) + assert_equal(pair['value'], 0) + + heap.remove_min() + heap.remove_min() + + heap.change_value(1, 10) + pair = heap.get_min() + assert_equal(pair['key'], 4) + assert_equal(pair['value'], 3) + + heap.remove_min() + pair = heap.get_min() + assert_equal(pair['key'], 1) + assert_equal(pair['value'], 10) + + +@skip_xp_backends(cpu_only=True) +def test_centroid_neg_distance(xp): + # gh-21011 + values = xp.asarray([0, 0, -1]) + with pytest.raises(ValueError): + # This is just checking that this doesn't crash + linkage(values, method='centroid') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_vq.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_vq.py new file mode 100644 index 0000000000000000000000000000000000000000..1902041fdce6b4f9241356cf45b83e7fbdda389d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/tests/test_vq.py @@ -0,0 +1,450 @@ +import warnings +import sys +from copy import deepcopy +from threading import Lock + +import numpy as np +from numpy.testing import ( + assert_array_equal, assert_equal, assert_, suppress_warnings +) +import pytest +from pytest import raises as assert_raises + +from scipy.cluster.vq import (kmeans, kmeans2, py_vq, vq, whiten, + ClusterError, _krandinit) +from scipy.cluster import _vq +from scipy.conftest import array_api_compatible +from scipy.sparse._sputils import matrix + +from scipy._lib import array_api_extra as xpx +from scipy._lib._array_api import ( + SCIPY_ARRAY_API, array_namespace, xp_copy, xp_assert_close, xp_assert_equal +) + +pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")] +skip_xp_backends = pytest.mark.skip_xp_backends + +TESTDATA_2D = np.array([ + -2.2, 1.17, -1.63, 1.69, -2.04, 4.38, -3.09, 0.95, -1.7, 4.79, -1.68, 0.68, + -2.26, 3.34, -2.29, 2.55, -1.72, -0.72, -1.99, 2.34, -2.75, 3.43, -2.45, + 2.41, -4.26, 3.65, -1.57, 1.87, -1.96, 4.03, -3.01, 3.86, -2.53, 1.28, + -4.0, 3.95, -1.62, 1.25, -3.42, 3.17, -1.17, 0.12, -3.03, -0.27, -2.07, + -0.55, -1.17, 1.34, -2.82, 3.08, -2.44, 0.24, -1.71, 2.48, -5.23, 4.29, + -2.08, 3.69, -1.89, 3.62, -2.09, 0.26, -0.92, 1.07, -2.25, 0.88, -2.25, + 2.02, -4.31, 3.86, -2.03, 3.42, -2.76, 0.3, -2.48, -0.29, -3.42, 3.21, + -2.3, 1.73, -2.84, 0.69, -1.81, 2.48, -5.24, 4.52, -2.8, 1.31, -1.67, + -2.34, -1.18, 2.17, -2.17, 2.82, -1.85, 2.25, -2.45, 1.86, -6.79, 3.94, + -2.33, 1.89, -1.55, 2.08, -1.36, 0.93, -2.51, 2.74, -2.39, 3.92, -3.33, + 2.99, -2.06, -0.9, -2.83, 3.35, -2.59, 3.05, -2.36, 1.85, -1.69, 1.8, + -1.39, 0.66, -2.06, 0.38, -1.47, 0.44, -4.68, 3.77, -5.58, 3.44, -2.29, + 2.24, -1.04, -0.38, -1.85, 4.23, -2.88, 0.73, -2.59, 1.39, -1.34, 1.75, + -1.95, 1.3, -2.45, 3.09, -1.99, 3.41, -5.55, 5.21, -1.73, 2.52, -2.17, + 0.85, -2.06, 0.49, -2.54, 2.07, -2.03, 1.3, -3.23, 3.09, -1.55, 1.44, + -0.81, 1.1, -2.99, 2.92, -1.59, 2.18, -2.45, -0.73, -3.12, -1.3, -2.83, + 0.2, -2.77, 3.24, -1.98, 1.6, -4.59, 3.39, -4.85, 3.75, -2.25, 1.71, -3.28, + 3.38, -1.74, 0.88, -2.41, 1.92, -2.24, 1.19, -2.48, 1.06, -1.68, -0.62, + -1.3, 0.39, -1.78, 2.35, -3.54, 2.44, -1.32, 0.66, -2.38, 2.76, -2.35, + 3.95, -1.86, 4.32, -2.01, -1.23, -1.79, 2.76, -2.13, -0.13, -5.25, 3.84, + -2.24, 1.59, -4.85, 2.96, -2.41, 0.01, -0.43, 0.13, -3.92, 2.91, -1.75, + -0.53, -1.69, 1.69, -1.09, 0.15, -2.11, 2.17, -1.53, 1.22, -2.1, -0.86, + -2.56, 2.28, -3.02, 3.33, -1.12, 3.86, -2.18, -1.19, -3.03, 0.79, -0.83, + 0.97, -3.19, 1.45, -1.34, 1.28, -2.52, 4.22, -4.53, 3.22, -1.97, 1.75, + -2.36, 3.19, -0.83, 1.53, -1.59, 1.86, -2.17, 2.3, -1.63, 2.71, -2.03, + 3.75, -2.57, -0.6, -1.47, 1.33, -1.95, 0.7, -1.65, 1.27, -1.42, 1.09, -3.0, + 3.87, -2.51, 3.06, -2.6, 0.74, -1.08, -0.03, -2.44, 1.31, -2.65, 2.99, + -1.84, 1.65, -4.76, 3.75, -2.07, 3.98, -2.4, 2.67, -2.21, 1.49, -1.21, + 1.22, -5.29, 2.38, -2.85, 2.28, -5.6, 3.78, -2.7, 0.8, -1.81, 3.5, -3.75, + 4.17, -1.29, 2.99, -5.92, 3.43, -1.83, 1.23, -1.24, -1.04, -2.56, 2.37, + -3.26, 0.39, -4.63, 2.51, -4.52, 3.04, -1.7, 0.36, -1.41, 0.04, -2.1, 1.0, + -1.87, 3.78, -4.32, 3.59, -2.24, 1.38, -1.99, -0.22, -1.87, 1.95, -0.84, + 2.17, -5.38, 3.56, -1.27, 2.9, -1.79, 3.31, -5.47, 3.85, -1.44, 3.69, + -2.02, 0.37, -1.29, 0.33, -2.34, 2.56, -1.74, -1.27, -1.97, 1.22, -2.51, + -0.16, -1.64, -0.96, -2.99, 1.4, -1.53, 3.31, -2.24, 0.45, -2.46, 1.71, + -2.88, 1.56, -1.63, 1.46, -1.41, 0.68, -1.96, 2.76, -1.61, + 2.11]).reshape((200, 2)) + + +# Global data +X = np.array([[3.0, 3], [4, 3], [4, 2], + [9, 2], [5, 1], [6, 2], [9, 4], + [5, 2], [5, 4], [7, 4], [6, 5]]) + +CODET1 = np.array([[3.0000, 3.0000], + [6.2000, 4.0000], + [5.8000, 1.8000]]) + +CODET2 = np.array([[11.0/3, 8.0/3], + [6.7500, 4.2500], + [6.2500, 1.7500]]) + +LABEL1 = np.array([0, 1, 2, 2, 2, 2, 1, 2, 1, 1, 1]) + + +class TestWhiten: + + def test_whiten(self, xp): + desired = xp.asarray([[5.08738849, 2.97091878], + [3.19909255, 0.69660580], + [4.51041982, 0.02640918], + [4.38567074, 0.95120889], + [2.32191480, 1.63195503]]) + + obs = xp.asarray([[0.98744510, 0.82766775], + [0.62093317, 0.19406729], + [0.87545741, 0.00735733], + [0.85124403, 0.26499712], + [0.45067590, 0.45464607]]) + xp_assert_close(whiten(obs), desired, rtol=1e-5) + + @pytest.fixture + def whiten_lock(self): + return Lock() + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_whiten_zero_std(self, xp, whiten_lock): + desired = xp.asarray([[0., 1.0, 2.86666544], + [0., 1.0, 1.32460034], + [0., 1.0, 3.74382172]]) + + obs = xp.asarray([[0., 1., 0.74109533], + [0., 1., 0.34243798], + [0., 1., 0.96785929]]) + + with whiten_lock: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + + xp_assert_close(whiten(obs), desired, rtol=1e-5) + + assert_equal(len(w), 1) + assert_(issubclass(w[-1].category, RuntimeWarning)) + + def test_whiten_not_finite(self, xp): + for bad_value in xp.nan, xp.inf, -xp.inf: + obs = xp.asarray([[0.98744510, bad_value], + [0.62093317, 0.19406729], + [0.87545741, 0.00735733], + [0.85124403, 0.26499712], + [0.45067590, 0.45464607]]) + assert_raises(ValueError, whiten, obs) + + @pytest.mark.skipif(SCIPY_ARRAY_API, + reason='`np.matrix` unsupported in array API mode') + def test_whiten_not_finite_matrix(self, xp): + for bad_value in np.nan, np.inf, -np.inf: + obs = matrix([[0.98744510, bad_value], + [0.62093317, 0.19406729], + [0.87545741, 0.00735733], + [0.85124403, 0.26499712], + [0.45067590, 0.45464607]]) + assert_raises(ValueError, whiten, obs) + + +class TestVq: + + @skip_xp_backends(cpu_only=True) + def test_py_vq(self, xp): + initc = np.concatenate([[X[0]], [X[1]], [X[2]]]) + # label1.dtype varies between int32 and int64 over platforms + label1 = py_vq(xp.asarray(X), xp.asarray(initc))[0] + xp_assert_equal(label1, xp.asarray(LABEL1, dtype=xp.int64), + check_dtype=False) + + @pytest.mark.skipif(SCIPY_ARRAY_API, + reason='`np.matrix` unsupported in array API mode') + def test_py_vq_matrix(self, xp): + initc = np.concatenate([[X[0]], [X[1]], [X[2]]]) + # label1.dtype varies between int32 and int64 over platforms + label1 = py_vq(matrix(X), matrix(initc))[0] + assert_array_equal(label1, LABEL1) + + @skip_xp_backends(np_only=True, reason='`_vq` only supports NumPy backend') + def test_vq(self, xp): + initc = np.concatenate([[X[0]], [X[1]], [X[2]]]) + label1, _ = _vq.vq(xp.asarray(X), xp.asarray(initc)) + assert_array_equal(label1, LABEL1) + _, _ = vq(xp.asarray(X), xp.asarray(initc)) + + @pytest.mark.skipif(SCIPY_ARRAY_API, + reason='`np.matrix` unsupported in array API mode') + def test_vq_matrix(self, xp): + initc = np.concatenate([[X[0]], [X[1]], [X[2]]]) + label1, _ = _vq.vq(matrix(X), matrix(initc)) + assert_array_equal(label1, LABEL1) + _, _ = vq(matrix(X), matrix(initc)) + + @skip_xp_backends(cpu_only=True) + def test_vq_1d(self, xp): + # Test special rank 1 vq algo, python implementation. + data = X[:, 0] + initc = data[:3] + a, b = _vq.vq(data, initc) + data = xp.asarray(data) + initc = xp.asarray(initc) + ta, tb = py_vq(data[:, np.newaxis], initc[:, np.newaxis]) + # ta.dtype varies between int32 and int64 over platforms + xp_assert_equal(ta, xp.asarray(a, dtype=xp.int64), check_dtype=False) + xp_assert_equal(tb, xp.asarray(b)) + + @skip_xp_backends(np_only=True, reason='`_vq` only supports NumPy backend') + def test__vq_sametype(self, xp): + a = xp.asarray([1.0, 2.0], dtype=xp.float64) + b = a.astype(xp.float32) + assert_raises(TypeError, _vq.vq, a, b) + + @skip_xp_backends(np_only=True, reason='`_vq` only supports NumPy backend') + def test__vq_invalid_type(self, xp): + a = xp.asarray([1, 2], dtype=int) + assert_raises(TypeError, _vq.vq, a, a) + + @skip_xp_backends(cpu_only=True) + def test_vq_large_nfeat(self, xp): + X = np.random.rand(20, 20) + code_book = np.random.rand(3, 20) + + codes0, dis0 = _vq.vq(X, code_book) + codes1, dis1 = py_vq( + xp.asarray(X), xp.asarray(code_book) + ) + xp_assert_close(dis1, xp.asarray(dis0), rtol=1e-5) + # codes1.dtype varies between int32 and int64 over platforms + xp_assert_equal(codes1, xp.asarray(codes0, dtype=xp.int64), check_dtype=False) + + X = X.astype(np.float32) + code_book = code_book.astype(np.float32) + + codes0, dis0 = _vq.vq(X, code_book) + codes1, dis1 = py_vq( + xp.asarray(X), xp.asarray(code_book) + ) + xp_assert_close(dis1, xp.asarray(dis0, dtype=xp.float64), rtol=1e-5) + # codes1.dtype varies between int32 and int64 over platforms + xp_assert_equal(codes1, xp.asarray(codes0, dtype=xp.int64), check_dtype=False) + + @skip_xp_backends(cpu_only=True) + def test_vq_large_features(self, xp): + X = np.random.rand(10, 5) * 1000000 + code_book = np.random.rand(2, 5) * 1000000 + + codes0, dis0 = _vq.vq(X, code_book) + codes1, dis1 = py_vq( + xp.asarray(X), xp.asarray(code_book) + ) + xp_assert_close(dis1, xp.asarray(dis0), rtol=1e-5) + # codes1.dtype varies between int32 and int64 over platforms + xp_assert_equal(codes1, xp.asarray(codes0, dtype=xp.int64), check_dtype=False) + + +# Whole class skipped on GPU for now; +# once pdist/cdist are hooked up for CuPy, more tests will work +@skip_xp_backends(cpu_only=True) +class TestKMean: + + def test_large_features(self, xp): + # Generate a data set with large values, and run kmeans on it to + # (regression for 1077). + d = 300 + n = 100 + + m1 = np.random.randn(d) + m2 = np.random.randn(d) + x = 10000 * np.random.randn(n, d) - 20000 * m1 + y = 10000 * np.random.randn(n, d) + 20000 * m2 + + data = np.empty((x.shape[0] + y.shape[0], d), np.float64) + data[:x.shape[0]] = x + data[x.shape[0]:] = y + + # use `seed` to ensure backwards compatibility after SPEC7 + kmeans(xp.asarray(data), 2, seed=1) + + def test_kmeans_simple(self, xp): + rng = np.random.default_rng(54321) + initc = np.concatenate([[X[0]], [X[1]], [X[2]]]) + code1 = kmeans(xp.asarray(X), xp.asarray(initc), iter=1, rng=rng)[0] + xp_assert_close(code1, xp.asarray(CODET2)) + + @pytest.mark.skipif(SCIPY_ARRAY_API, + reason='`np.matrix` unsupported in array API mode') + def test_kmeans_simple_matrix(self, xp): + rng = np.random.default_rng(54321) + initc = np.concatenate([[X[0]], [X[1]], [X[2]]]) + code1 = kmeans(matrix(X), matrix(initc), iter=1, rng=rng)[0] + xp_assert_close(code1, CODET2) + + def test_kmeans_lost_cluster(self, xp): + # This will cause kmeans to have a cluster with no points. + data = xp.asarray(TESTDATA_2D) + initk = xp.asarray([[-1.8127404, -0.67128041], + [2.04621601, 0.07401111], + [-2.31149087, -0.05160469]]) + + kmeans(data, initk) + with suppress_warnings() as sup: + sup.filter(UserWarning, + "One of the clusters is empty. Re-run kmeans with a " + "different initialization") + kmeans2(data, initk, missing='warn') + + assert_raises(ClusterError, kmeans2, data, initk, missing='raise') + + def test_kmeans2_simple(self, xp): + rng = np.random.default_rng(12345678) + initc = xp.asarray(np.concatenate([[X[0]], [X[1]], [X[2]]])) + arrays = [xp.asarray] if SCIPY_ARRAY_API else [np.asarray, matrix] + for tp in arrays: + code1 = kmeans2(tp(X), tp(initc), iter=1, rng=rng)[0] + code2 = kmeans2(tp(X), tp(initc), iter=2, rng=rng)[0] + + xp_assert_close(code1, xp.asarray(CODET1)) + xp_assert_close(code2, xp.asarray(CODET2)) + + @pytest.mark.skipif(SCIPY_ARRAY_API, + reason='`np.matrix` unsupported in array API mode') + def test_kmeans2_simple_matrix(self, xp): + rng = np.random.default_rng(12345678) + initc = xp.asarray(np.concatenate([[X[0]], [X[1]], [X[2]]])) + code1 = kmeans2(matrix(X), matrix(initc), iter=1, rng=rng)[0] + code2 = kmeans2(matrix(X), matrix(initc), iter=2, rng=rng)[0] + + xp_assert_close(code1, CODET1) + xp_assert_close(code2, CODET2) + + def test_kmeans2_rank1(self, xp): + data = xp.asarray(TESTDATA_2D) + data1 = data[:, 0] + + initc = data1[:3] + code = xp_copy(initc, xp=xp) + + # use `seed` to ensure backwards compatibility after SPEC7 + kmeans2(data1, code, iter=1, seed=1)[0] + kmeans2(data1, code, iter=2)[0] + + def test_kmeans2_rank1_2(self, xp): + data = xp.asarray(TESTDATA_2D) + data1 = data[:, 0] + kmeans2(data1, 2, iter=1) + + def test_kmeans2_high_dim(self, xp): + # test kmeans2 when the number of dimensions exceeds the number + # of input points + data = xp.asarray(TESTDATA_2D) + data = xp.reshape(data, (20, 20))[:10, :] + kmeans2(data, 2) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_kmeans2_init(self, xp): + rng = np.random.default_rng(12345678) + data = xp.asarray(TESTDATA_2D) + k = 3 + + kmeans2(data, k, minit='points', rng=rng) + kmeans2(data[:, 1], k, minit='points', rng=rng) # special case (1-D) + + kmeans2(data, k, minit='++', rng=rng) + kmeans2(data[:, 1], k, minit='++', rng=rng) # special case (1-D) + + # minit='random' can give warnings, filter those + with suppress_warnings() as sup: + sup.filter(message="One of the clusters is empty. Re-run.") + kmeans2(data, k, minit='random', rng=rng) + kmeans2(data[:, 1], k, minit='random', rng=rng) # special case (1-D) + + @pytest.fixture + def krand_lock(self): + return Lock() + + @pytest.mark.skipif(sys.platform == 'win32', + reason='Fails with MemoryError in Wine.') + def test_krandinit(self, xp, krand_lock): + data = xp.asarray(TESTDATA_2D) + datas = [xp.reshape(data, (200, 2)), + xp.reshape(data, (20, 20))[:10, :]] + k = int(1e6) + xp_test = array_namespace(data) + with krand_lock: + for data in datas: + rng = np.random.default_rng(1234) + init = _krandinit(data, k, rng, xp_test) + orig_cov = xpx.cov(data.T, xp=xp_test) + init_cov = xpx.cov(init.T, xp=xp_test) + xp_assert_close(orig_cov, init_cov, atol=1.1e-2) + + def test_kmeans2_empty(self, xp): + # Regression test for gh-1032. + assert_raises(ValueError, kmeans2, xp.asarray([]), 2) + + def test_kmeans_0k(self, xp): + # Regression test for gh-1073: fail when k arg is 0. + assert_raises(ValueError, kmeans, xp.asarray(X), 0) + assert_raises(ValueError, kmeans2, xp.asarray(X), 0) + assert_raises(ValueError, kmeans2, xp.asarray(X), xp.asarray([])) + + def test_kmeans_large_thres(self, xp): + # Regression test for gh-1774 + x = xp.asarray([1, 2, 3, 4, 10], dtype=xp.float64) + res = kmeans(x, 1, thresh=1e16) + xp_assert_close(res[0], xp.asarray([4.], dtype=xp.float64)) + xp_assert_close(res[1], xp.asarray(2.3999999999999999, dtype=xp.float64)[()]) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_kmeans2_kpp_low_dim(self, xp): + # Regression test for gh-11462 + rng = np.random.default_rng(2358792345678234568) + prev_res = xp.asarray([[-1.95266667, 0.898], + [-3.153375, 3.3945]], dtype=xp.float64) + res, _ = kmeans2(xp.asarray(TESTDATA_2D), 2, minit='++', rng=rng) + xp_assert_close(res, prev_res) + + @pytest.mark.thread_unsafe + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_kmeans2_kpp_high_dim(self, xp): + # Regression test for gh-11462 + rng = np.random.default_rng(23587923456834568) + n_dim = 100 + size = 10 + centers = np.vstack([5 * np.ones(n_dim), + -5 * np.ones(n_dim)]) + + data = np.vstack([ + rng.multivariate_normal(centers[0], np.eye(n_dim), size=size), + rng.multivariate_normal(centers[1], np.eye(n_dim), size=size) + ]) + + data = xp.asarray(data) + res, _ = kmeans2(data, 2, minit='++', rng=rng) + xp_assert_equal(xp.sign(res), xp.sign(xp.asarray(centers))) + + def test_kmeans_diff_convergence(self, xp): + # Regression test for gh-8727 + obs = xp.asarray([-3, -1, 0, 1, 1, 8], dtype=xp.float64) + res = kmeans(obs, xp.asarray([-3., 0.99])) + xp_assert_close(res[0], xp.asarray([-0.4, 8.], dtype=xp.float64)) + xp_assert_close(res[1], xp.asarray(1.0666666666666667, dtype=xp.float64)[()]) + + @skip_xp_backends('jax.numpy', + reason='jax arrays do not support item assignment') + def test_kmeans_and_kmeans2_random_seed(self, xp): + + seed_list = [ + 1234, np.random.RandomState(1234), np.random.default_rng(1234) + ] + + for seed in seed_list: + seed1 = deepcopy(seed) + seed2 = deepcopy(seed) + data = xp.asarray(TESTDATA_2D) + # test for kmeans + res1, _ = kmeans(data, 2, seed=seed1) + res2, _ = kmeans(data, 2, seed=seed2) + xp_assert_close(res1, res2) # should be same results + # test for kmeans2 + for minit in ["random", "points", "++"]: + res1, _ = kmeans2(data, 2, minit=minit, seed=seed1) + res2, _ = kmeans2(data, 2, minit=minit, seed=seed2) + xp_assert_close(res1, res2) # should be same results diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/vq.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/vq.py new file mode 100644 index 0000000000000000000000000000000000000000..3d779c13903649779cf3ce8f72b0665712ab267b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/cluster/vq.py @@ -0,0 +1,828 @@ +""" +K-means clustering and vector quantization (:mod:`scipy.cluster.vq`) +==================================================================== + +Provides routines for k-means clustering, generating code books +from k-means models and quantizing vectors by comparing them with +centroids in a code book. + +.. autosummary:: + :toctree: generated/ + + whiten -- Normalize a group of observations so each feature has unit variance + vq -- Calculate code book membership of a set of observation vectors + kmeans -- Perform k-means on a set of observation vectors forming k clusters + kmeans2 -- A different implementation of k-means with more methods + -- for initializing centroids + +Background information +---------------------- +The k-means algorithm takes as input the number of clusters to +generate, k, and a set of observation vectors to cluster. It +returns a set of centroids, one for each of the k clusters. An +observation vector is classified with the cluster number or +centroid index of the centroid closest to it. + +A vector v belongs to cluster i if it is closer to centroid i than +any other centroid. If v belongs to i, we say centroid i is the +dominating centroid of v. The k-means algorithm tries to +minimize distortion, which is defined as the sum of the squared distances +between each observation vector and its dominating centroid. +The minimization is achieved by iteratively reclassifying +the observations into clusters and recalculating the centroids until +a configuration is reached in which the centroids are stable. One can +also define a maximum number of iterations. + +Since vector quantization is a natural application for k-means, +information theory terminology is often used. The centroid index +or cluster index is also referred to as a "code" and the table +mapping codes to centroids and, vice versa, is often referred to as a +"code book". The result of k-means, a set of centroids, can be +used to quantize vectors. Quantization aims to find an encoding of +vectors that reduces the expected distortion. + +All routines expect obs to be an M by N array, where the rows are +the observation vectors. The codebook is a k by N array, where the +ith row is the centroid of code word i. The observation vectors +and centroids have the same feature dimension. + +As an example, suppose we wish to compress a 24-bit color image +(each pixel is represented by one byte for red, one for blue, and +one for green) before sending it over the web. By using a smaller +8-bit encoding, we can reduce the amount of data by two +thirds. Ideally, the colors for each of the 256 possible 8-bit +encoding values should be chosen to minimize distortion of the +color. Running k-means with k=256 generates a code book of 256 +codes, which fills up all possible 8-bit sequences. Instead of +sending a 3-byte value for each pixel, the 8-bit centroid index +(or code word) of the dominating centroid is transmitted. The code +book is also sent over the wire so each 8-bit code can be +translated back to a 24-bit pixel value representation. If the +image of interest was of an ocean, we would expect many 24-bit +blues to be represented by 8-bit codes. If it was an image of a +human face, more flesh-tone colors would be represented in the +code book. + +""" +import warnings +import numpy as np +from collections import deque +from scipy._lib._array_api import ( + _asarray, array_namespace, xp_size, xp_copy +) +from scipy._lib._util import (check_random_state, rng_integers, + _transition_to_rng) +from scipy._lib import array_api_extra as xpx +from scipy.spatial.distance import cdist + +from . import _vq + +__docformat__ = 'restructuredtext' + +__all__ = ['whiten', 'vq', 'kmeans', 'kmeans2'] + + +class ClusterError(Exception): + pass + + +def whiten(obs, check_finite=True): + """ + Normalize a group of observations on a per feature basis. + + Before running k-means, it is beneficial to rescale each feature + dimension of the observation set by its standard deviation (i.e. "whiten" + it - as in "white noise" where each frequency has equal power). + Each feature is divided by its standard deviation across all observations + to give it unit variance. + + Parameters + ---------- + obs : ndarray + Each row of the array is an observation. The + columns are the features seen during each observation. + + >>> # f0 f1 f2 + >>> obs = [[ 1., 1., 1.], #o0 + ... [ 2., 2., 2.], #o1 + ... [ 3., 3., 3.], #o2 + ... [ 4., 4., 4.]] #o3 + + check_finite : bool, optional + Whether to check that the input matrices contain only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + Default: True + + Returns + ------- + result : ndarray + Contains the values in `obs` scaled by the standard deviation + of each column. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster.vq import whiten + >>> features = np.array([[1.9, 2.3, 1.7], + ... [1.5, 2.5, 2.2], + ... [0.8, 0.6, 1.7,]]) + >>> whiten(features) + array([[ 4.17944278, 2.69811351, 7.21248917], + [ 3.29956009, 2.93273208, 9.33380951], + [ 1.75976538, 0.7038557 , 7.21248917]]) + + """ + xp = array_namespace(obs) + obs = _asarray(obs, check_finite=check_finite, xp=xp) + std_dev = xp.std(obs, axis=0) + zero_std_mask = std_dev == 0 + if xp.any(zero_std_mask): + std_dev[zero_std_mask] = 1.0 + warnings.warn("Some columns have standard deviation zero. " + "The values of these columns will not change.", + RuntimeWarning, stacklevel=2) + return obs / std_dev + + +def vq(obs, code_book, check_finite=True): + """ + Assign codes from a code book to observations. + + Assigns a code from a code book to each observation. Each + observation vector in the 'M' by 'N' `obs` array is compared with the + centroids in the code book and assigned the code of the closest + centroid. + + The features in `obs` should have unit variance, which can be + achieved by passing them through the whiten function. The code + book can be created with the k-means algorithm or a different + encoding algorithm. + + Parameters + ---------- + obs : ndarray + Each row of the 'M' x 'N' array is an observation. The columns are + the "features" seen during each observation. The features must be + whitened first using the whiten function or something equivalent. + code_book : ndarray + The code book is usually generated using the k-means algorithm. + Each row of the array holds a different code, and the columns are + the features of the code. + + >>> # f0 f1 f2 f3 + >>> code_book = [ + ... [ 1., 2., 3., 4.], #c0 + ... [ 1., 2., 3., 4.], #c1 + ... [ 1., 2., 3., 4.]] #c2 + + check_finite : bool, optional + Whether to check that the input matrices contain only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + Default: True + + Returns + ------- + code : ndarray + A length M array holding the code book index for each observation. + dist : ndarray + The distortion (distance) between the observation and its nearest + code. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster.vq import vq + >>> code_book = np.array([[1., 1., 1.], + ... [2., 2., 2.]]) + >>> features = np.array([[1.9, 2.3, 1.7], + ... [1.5, 2.5, 2.2], + ... [0.8, 0.6, 1.7]]) + >>> vq(features, code_book) + (array([1, 1, 0], dtype=int32), array([0.43588989, 0.73484692, 0.83066239])) + + """ + xp = array_namespace(obs, code_book) + obs = _asarray(obs, xp=xp, check_finite=check_finite) + code_book = _asarray(code_book, xp=xp, check_finite=check_finite) + ct = xp.result_type(obs, code_book) + + c_obs = xp.astype(obs, ct, copy=False) + c_code_book = xp.astype(code_book, ct, copy=False) + + if xp.isdtype(ct, kind='real floating'): + c_obs = np.asarray(c_obs) + c_code_book = np.asarray(c_code_book) + result = _vq.vq(c_obs, c_code_book) + return xp.asarray(result[0]), xp.asarray(result[1]) + return py_vq(obs, code_book, check_finite=False) + + +def py_vq(obs, code_book, check_finite=True): + """ Python version of vq algorithm. + + The algorithm computes the Euclidean distance between each + observation and every frame in the code_book. + + Parameters + ---------- + obs : ndarray + Expects a rank 2 array. Each row is one observation. + code_book : ndarray + Code book to use. Same format than obs. Should have same number of + features (e.g., columns) than obs. + check_finite : bool, optional + Whether to check that the input matrices contain only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + Default: True + + Returns + ------- + code : ndarray + code[i] gives the label of the ith obversation; its code is + code_book[code[i]]. + mind_dist : ndarray + min_dist[i] gives the distance between the ith observation and its + corresponding code. + + Notes + ----- + This function is slower than the C version but works for + all input types. If the inputs have the wrong types for the + C versions of the function, this one is called as a last resort. + + It is about 20 times slower than the C version. + + """ + xp = array_namespace(obs, code_book) + obs = _asarray(obs, xp=xp, check_finite=check_finite) + code_book = _asarray(code_book, xp=xp, check_finite=check_finite) + + if obs.ndim != code_book.ndim: + raise ValueError("Observation and code_book should have the same rank") + + if obs.ndim == 1: + obs = obs[:, xp.newaxis] + code_book = code_book[:, xp.newaxis] + + # Once `cdist` has array API support, this `xp.asarray` call can be removed + dist = xp.asarray(cdist(obs, code_book)) + code = xp.argmin(dist, axis=1) + min_dist = xp.min(dist, axis=1) + return code, min_dist + + +def _kmeans(obs, guess, thresh=1e-5, xp=None): + """ "raw" version of k-means. + + Returns + ------- + code_book + The lowest distortion codebook found. + avg_dist + The average distance a observation is from a code in the book. + Lower means the code_book matches the data better. + + See Also + -------- + kmeans : wrapper around k-means + + Examples + -------- + Note: not whitened in this example. + + >>> import numpy as np + >>> from scipy.cluster.vq import _kmeans + >>> features = np.array([[ 1.9,2.3], + ... [ 1.5,2.5], + ... [ 0.8,0.6], + ... [ 0.4,1.8], + ... [ 1.0,1.0]]) + >>> book = np.array((features[0],features[2])) + >>> _kmeans(features,book) + (array([[ 1.7 , 2.4 ], + [ 0.73333333, 1.13333333]]), 0.40563916697728591) + + """ + xp = np if xp is None else xp + code_book = guess + diff = xp.inf + prev_avg_dists = deque([diff], maxlen=2) + while diff > thresh: + # compute membership and distances between obs and code_book + obs_code, distort = vq(obs, code_book, check_finite=False) + prev_avg_dists.append(xp.mean(distort, axis=-1)) + # recalc code_book as centroids of associated obs + obs = np.asarray(obs) + obs_code = np.asarray(obs_code) + code_book, has_members = _vq.update_cluster_means(obs, obs_code, + code_book.shape[0]) + obs = xp.asarray(obs) + obs_code = xp.asarray(obs_code) + code_book = xp.asarray(code_book) + has_members = xp.asarray(has_members) + code_book = code_book[has_members] + diff = xp.abs(prev_avg_dists[0] - prev_avg_dists[1]) + + return code_book, prev_avg_dists[1] + + +@_transition_to_rng("seed") +def kmeans(obs, k_or_guess, iter=20, thresh=1e-5, check_finite=True, + *, rng=None): + """ + Performs k-means on a set of observation vectors forming k clusters. + + The k-means algorithm adjusts the classification of the observations + into clusters and updates the cluster centroids until the position of + the centroids is stable over successive iterations. In this + implementation of the algorithm, the stability of the centroids is + determined by comparing the absolute value of the change in the average + Euclidean distance between the observations and their corresponding + centroids against a threshold. This yields + a code book mapping centroids to codes and vice versa. + + Parameters + ---------- + obs : ndarray + Each row of the M by N array is an observation vector. The + columns are the features seen during each observation. + The features must be whitened first with the `whiten` function. + + k_or_guess : int or ndarray + The number of centroids to generate. A code is assigned to + each centroid, which is also the row index of the centroid + in the code_book matrix generated. + + The initial k centroids are chosen by randomly selecting + observations from the observation matrix. Alternatively, + passing a k by N array specifies the initial k centroids. + + iter : int, optional + The number of times to run k-means, returning the codebook + with the lowest distortion. This argument is ignored if + initial centroids are specified with an array for the + ``k_or_guess`` parameter. This parameter does not represent the + number of iterations of the k-means algorithm. + + thresh : float, optional + Terminates the k-means algorithm if the change in + distortion since the last k-means iteration is less than + or equal to threshold. + + check_finite : bool, optional + Whether to check that the input matrices contain only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + Default: True + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + codebook : ndarray + A k by N array of k centroids. The ith centroid + codebook[i] is represented with the code i. The centroids + and codes generated represent the lowest distortion seen, + not necessarily the globally minimal distortion. + Note that the number of centroids is not necessarily the same as the + ``k_or_guess`` parameter, because centroids assigned to no observations + are removed during iterations. + + distortion : float + The mean (non-squared) Euclidean distance between the observations + passed and the centroids generated. Note the difference to the standard + definition of distortion in the context of the k-means algorithm, which + is the sum of the squared distances. + + See Also + -------- + kmeans2 : a different implementation of k-means clustering + with more methods for generating initial centroids but without + using a distortion change threshold as a stopping criterion. + + whiten : must be called prior to passing an observation matrix + to kmeans. + + Notes + ----- + For more functionalities or optimal performance, you can use + `sklearn.cluster.KMeans `_. + `This `_ + is a benchmark result of several implementations. + + Examples + -------- + >>> import numpy as np + >>> from scipy.cluster.vq import vq, kmeans, whiten + >>> import matplotlib.pyplot as plt + >>> features = np.array([[ 1.9,2.3], + ... [ 1.5,2.5], + ... [ 0.8,0.6], + ... [ 0.4,1.8], + ... [ 0.1,0.1], + ... [ 0.2,1.8], + ... [ 2.0,0.5], + ... [ 0.3,1.5], + ... [ 1.0,1.0]]) + >>> whitened = whiten(features) + >>> book = np.array((whitened[0],whitened[2])) + >>> kmeans(whitened,book) + (array([[ 2.3110306 , 2.86287398], # random + [ 0.93218041, 1.24398691]]), 0.85684700941625547) + + >>> codes = 3 + >>> kmeans(whitened,codes) + (array([[ 2.3110306 , 2.86287398], # random + [ 1.32544402, 0.65607529], + [ 0.40782893, 2.02786907]]), 0.5196582527686241) + + >>> # Create 50 datapoints in two clusters a and b + >>> pts = 50 + >>> rng = np.random.default_rng() + >>> a = rng.multivariate_normal([0, 0], [[4, 1], [1, 4]], size=pts) + >>> b = rng.multivariate_normal([30, 10], + ... [[10, 2], [2, 1]], + ... size=pts) + >>> features = np.concatenate((a, b)) + >>> # Whiten data + >>> whitened = whiten(features) + >>> # Find 2 clusters in the data + >>> codebook, distortion = kmeans(whitened, 2) + >>> # Plot whitened data and cluster centers in red + >>> plt.scatter(whitened[:, 0], whitened[:, 1]) + >>> plt.scatter(codebook[:, 0], codebook[:, 1], c='r') + >>> plt.show() + + """ + if isinstance(k_or_guess, int): + xp = array_namespace(obs) + else: + xp = array_namespace(obs, k_or_guess) + obs = _asarray(obs, xp=xp, check_finite=check_finite) + guess = _asarray(k_or_guess, xp=xp, check_finite=check_finite) + if iter < 1: + raise ValueError(f"iter must be at least 1, got {iter}") + + # Determine whether a count (scalar) or an initial guess (array) was passed. + if xp_size(guess) != 1: + if xp_size(guess) < 1: + raise ValueError(f"Asked for 0 clusters. Initial book was {guess}") + return _kmeans(obs, guess, thresh=thresh, xp=xp) + + # k_or_guess is a scalar, now verify that it's an integer + k = int(guess) + if k != guess: + raise ValueError("If k_or_guess is a scalar, it must be an integer.") + if k < 1: + raise ValueError("Asked for %d clusters." % k) + + rng = check_random_state(rng) + + # initialize best distance value to a large value + best_dist = xp.inf + for i in range(iter): + # the initial code book is randomly selected from observations + guess = _kpoints(obs, k, rng, xp) + book, dist = _kmeans(obs, guess, thresh=thresh, xp=xp) + if dist < best_dist: + best_book = book + best_dist = dist + return best_book, best_dist + + +def _kpoints(data, k, rng, xp): + """Pick k points at random in data (one row = one observation). + + Parameters + ---------- + data : ndarray + Expect a rank 1 or 2 array. Rank 1 are assumed to describe one + dimensional data, rank 2 multidimensional data, in which case one + row is one observation. + k : int + Number of samples to generate. + rng : `numpy.random.Generator` or `numpy.random.RandomState` + Random number generator. + + Returns + ------- + x : ndarray + A 'k' by 'N' containing the initial centroids + + """ + idx = rng.choice(data.shape[0], size=int(k), replace=False) + # convert to array with default integer dtype (avoids numpy#25607) + idx = xp.asarray(idx, dtype=xp.asarray([1]).dtype) + return xp.take(data, idx, axis=0) + + +def _krandinit(data, k, rng, xp): + """Returns k samples of a random variable whose parameters depend on data. + + More precisely, it returns k observations sampled from a Gaussian random + variable whose mean and covariances are the ones estimated from the data. + + Parameters + ---------- + data : ndarray + Expect a rank 1 or 2 array. Rank 1 is assumed to describe 1-D + data, rank 2 multidimensional data, in which case one + row is one observation. + k : int + Number of samples to generate. + rng : `numpy.random.Generator` or `numpy.random.RandomState` + Random number generator. + + Returns + ------- + x : ndarray + A 'k' by 'N' containing the initial centroids + + """ + mu = xp.mean(data, axis=0) + k = np.asarray(k) + + if data.ndim == 1: + _cov = xpx.cov(data, xp=xp) + x = rng.standard_normal(size=k) + x = xp.asarray(x) + x *= xp.sqrt(_cov) + elif data.shape[1] > data.shape[0]: + # initialize when the covariance matrix is rank deficient + _, s, vh = xp.linalg.svd(data - mu, full_matrices=False) + x = rng.standard_normal(size=(k, xp_size(s))) + x = xp.asarray(x) + sVh = s[:, None] * vh / xp.sqrt(data.shape[0] - xp.asarray(1.)) + x = x @ sVh + else: + _cov = xpx.atleast_nd(xpx.cov(data.T, xp=xp), ndim=2, xp=xp) + + # k rows, d cols (one row = one obs) + # Generate k sample of a random variable ~ Gaussian(mu, cov) + x = rng.standard_normal(size=(k, xp_size(mu))) + x = xp.asarray(x) + x = x @ xp.linalg.cholesky(_cov).T + + x += mu + return x + + +def _kpp(data, k, rng, xp): + """ Picks k points in the data based on the kmeans++ method. + + Parameters + ---------- + data : ndarray + Expect a rank 1 or 2 array. Rank 1 is assumed to describe 1-D + data, rank 2 multidimensional data, in which case one + row is one observation. + k : int + Number of samples to generate. + rng : `numpy.random.Generator` or `numpy.random.RandomState` + Random number generator. + + Returns + ------- + init : ndarray + A 'k' by 'N' containing the initial centroids. + + References + ---------- + .. [1] D. Arthur and S. Vassilvitskii, "k-means++: the advantages of + careful seeding", Proceedings of the Eighteenth Annual ACM-SIAM Symposium + on Discrete Algorithms, 2007. + """ + + ndim = len(data.shape) + if ndim == 1: + data = data[:, None] + + dims = data.shape[1] + + init = xp.empty((int(k), dims)) + + for i in range(k): + if i == 0: + init[i, :] = data[rng_integers(rng, data.shape[0]), :] + + else: + D2 = cdist(init[:i,:], data, metric='sqeuclidean').min(axis=0) + probs = D2/D2.sum() + cumprobs = probs.cumsum() + r = rng.uniform() + cumprobs = np.asarray(cumprobs) + init[i, :] = data[np.searchsorted(cumprobs, r), :] + + if ndim == 1: + init = init[:, 0] + return init + + +_valid_init_meth = {'random': _krandinit, 'points': _kpoints, '++': _kpp} + + +def _missing_warn(): + """Print a warning when called.""" + warnings.warn("One of the clusters is empty. " + "Re-run kmeans with a different initialization.", + stacklevel=3) + + +def _missing_raise(): + """Raise a ClusterError when called.""" + raise ClusterError("One of the clusters is empty. " + "Re-run kmeans with a different initialization.") + + +_valid_miss_meth = {'warn': _missing_warn, 'raise': _missing_raise} + + +@_transition_to_rng("seed") +def kmeans2(data, k, iter=10, thresh=1e-5, minit='random', + missing='warn', check_finite=True, *, rng=None): + """ + Classify a set of observations into k clusters using the k-means algorithm. + + The algorithm attempts to minimize the Euclidean distance between + observations and centroids. Several initialization methods are + included. + + Parameters + ---------- + data : ndarray + A 'M' by 'N' array of 'M' observations in 'N' dimensions or a length + 'M' array of 'M' 1-D observations. + k : int or ndarray + The number of clusters to form as well as the number of + centroids to generate. If `minit` initialization string is + 'matrix', or if a ndarray is given instead, it is + interpreted as initial cluster to use instead. + iter : int, optional + Number of iterations of the k-means algorithm to run. Note + that this differs in meaning from the iters parameter to + the kmeans function. + thresh : float, optional + (not used yet) + minit : str, optional + Method for initialization. Available methods are 'random', + 'points', '++' and 'matrix': + + 'random': generate k centroids from a Gaussian with mean and + variance estimated from the data. + + 'points': choose k observations (rows) at random from data for + the initial centroids. + + '++': choose k observations accordingly to the kmeans++ method + (careful seeding) + + 'matrix': interpret the k parameter as a k by M (or length k + array for 1-D data) array of initial centroids. + missing : str, optional + Method to deal with empty clusters. Available methods are + 'warn' and 'raise': + + 'warn': give a warning and continue. + + 'raise': raise an ClusterError and terminate the algorithm. + check_finite : bool, optional + Whether to check that the input matrices contain only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + Default: True + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + centroid : ndarray + A 'k' by 'N' array of centroids found at the last iteration of + k-means. + label : ndarray + label[i] is the code or index of the centroid the + ith observation is closest to. + + See Also + -------- + kmeans + + References + ---------- + .. [1] D. Arthur and S. Vassilvitskii, "k-means++: the advantages of + careful seeding", Proceedings of the Eighteenth Annual ACM-SIAM Symposium + on Discrete Algorithms, 2007. + + Examples + -------- + >>> from scipy.cluster.vq import kmeans2 + >>> import matplotlib.pyplot as plt + >>> import numpy as np + + Create z, an array with shape (100, 2) containing a mixture of samples + from three multivariate normal distributions. + + >>> rng = np.random.default_rng() + >>> a = rng.multivariate_normal([0, 6], [[2, 1], [1, 1.5]], size=45) + >>> b = rng.multivariate_normal([2, 0], [[1, -1], [-1, 3]], size=30) + >>> c = rng.multivariate_normal([6, 4], [[5, 0], [0, 1.2]], size=25) + >>> z = np.concatenate((a, b, c)) + >>> rng.shuffle(z) + + Compute three clusters. + + >>> centroid, label = kmeans2(z, 3, minit='points') + >>> centroid + array([[ 2.22274463, -0.61666946], # may vary + [ 0.54069047, 5.86541444], + [ 6.73846769, 4.01991898]]) + + How many points are in each cluster? + + >>> counts = np.bincount(label) + >>> counts + array([29, 51, 20]) # may vary + + Plot the clusters. + + >>> w0 = z[label == 0] + >>> w1 = z[label == 1] + >>> w2 = z[label == 2] + >>> plt.plot(w0[:, 0], w0[:, 1], 'o', alpha=0.5, label='cluster 0') + >>> plt.plot(w1[:, 0], w1[:, 1], 'd', alpha=0.5, label='cluster 1') + >>> plt.plot(w2[:, 0], w2[:, 1], 's', alpha=0.5, label='cluster 2') + >>> plt.plot(centroid[:, 0], centroid[:, 1], 'k*', label='centroids') + >>> plt.axis('equal') + >>> plt.legend(shadow=True) + >>> plt.show() + + """ + if int(iter) < 1: + raise ValueError(f"Invalid iter ({iter}), must be a positive integer.") + try: + miss_meth = _valid_miss_meth[missing] + except KeyError as e: + raise ValueError(f"Unknown missing method {missing!r}") from e + + if isinstance(k, int): + xp = array_namespace(data) + else: + xp = array_namespace(data, k) + data = _asarray(data, xp=xp, check_finite=check_finite) + code_book = xp_copy(k, xp=xp) + if data.ndim == 1: + d = 1 + elif data.ndim == 2: + d = data.shape[1] + else: + raise ValueError("Input of rank > 2 is not supported.") + + if xp_size(data) < 1 or xp_size(code_book) < 1: + raise ValueError("Empty input is not supported.") + + # If k is not a single value, it should be compatible with data's shape + if minit == 'matrix' or xp_size(code_book) > 1: + if data.ndim != code_book.ndim: + raise ValueError("k array doesn't match data rank") + nc = code_book.shape[0] + if data.ndim > 1 and code_book.shape[1] != d: + raise ValueError("k array doesn't match data dimension") + else: + nc = int(code_book) + + if nc < 1: + raise ValueError("Cannot ask kmeans2 for %d clusters" + " (k was %s)" % (nc, code_book)) + elif nc != code_book: + warnings.warn("k was not an integer, was converted.", stacklevel=2) + + try: + init_meth = _valid_init_meth[minit] + except KeyError as e: + raise ValueError(f"Unknown init method {minit!r}") from e + else: + rng = check_random_state(rng) + code_book = init_meth(data, code_book, rng, xp) + + data = np.asarray(data) + code_book = np.asarray(code_book) + for i in range(iter): + # Compute the nearest neighbor for each obs using the current code book + label = vq(data, code_book, check_finite=check_finite)[0] + # Update the code book by computing centroids + new_code_book, has_members = _vq.update_cluster_means(data, label, nc) + if not has_members.all(): + miss_meth() + # Set the empty clusters to their previous positions + new_code_book[~has_members] = code_book[~has_members] + code_book = new_code_book + + return xp.asarray(code_book), xp.asarray(label) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/conftest.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..894548d3a92c4757583ef69288b2f5495e07f0cf --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/conftest.py @@ -0,0 +1,552 @@ +# Pytest customization +import json +import os +import warnings +import tempfile +from contextlib import contextmanager + +import numpy as np +import numpy.testing as npt +import pytest +import hypothesis + +from scipy._lib._fpumode import get_fpu_mode +from scipy._lib._testutils import FPUModeChangeWarning +from scipy._lib._array_api import SCIPY_ARRAY_API, SCIPY_DEVICE +from scipy._lib import _pep440 + +try: + from scipy_doctest.conftest import dt_config + HAVE_SCPDT = True +except ModuleNotFoundError: + HAVE_SCPDT = False + +try: + import pytest_run_parallel # noqa:F401 + PARALLEL_RUN_AVAILABLE = True +except Exception: + PARALLEL_RUN_AVAILABLE = False + + +def pytest_configure(config): + config.addinivalue_line("markers", + "slow: Tests that are very slow.") + config.addinivalue_line("markers", + "xslow: mark test as extremely slow (not run unless explicitly requested)") + config.addinivalue_line("markers", + "xfail_on_32bit: mark test as failing on 32-bit platforms") + try: + import pytest_timeout # noqa:F401 + except Exception: + config.addinivalue_line( + "markers", 'timeout: mark a test for a non-default timeout') + try: + # This is a more reliable test of whether pytest_fail_slow is installed + # When I uninstalled it, `import pytest_fail_slow` didn't fail! + from pytest_fail_slow import parse_duration # type: ignore[import-not-found] # noqa:F401,E501 + except Exception: + config.addinivalue_line( + "markers", 'fail_slow: mark a test for a non-default timeout failure') + config.addinivalue_line("markers", + "skip_xp_backends(backends, reason=None, np_only=False, cpu_only=False, " + "exceptions=None): " + "mark the desired skip configuration for the `skip_xp_backends` fixture.") + config.addinivalue_line("markers", + "xfail_xp_backends(backends, reason=None, np_only=False, cpu_only=False, " + "exceptions=None): " + "mark the desired xfail configuration for the `xfail_xp_backends` fixture.") + if not PARALLEL_RUN_AVAILABLE: + config.addinivalue_line( + 'markers', + 'parallel_threads(n): run the given test function in parallel ' + 'using `n` threads.') + config.addinivalue_line( + "markers", + "thread_unsafe: mark the test function as single-threaded", + ) + config.addinivalue_line( + "markers", + "iterations(n): run the given test function `n` times in each thread", + ) + + +def pytest_runtest_setup(item): + mark = item.get_closest_marker("xslow") + if mark is not None: + try: + v = int(os.environ.get('SCIPY_XSLOW', '0')) + except ValueError: + v = False + if not v: + pytest.skip("very slow test; " + "set environment variable SCIPY_XSLOW=1 to run it") + mark = item.get_closest_marker("xfail_on_32bit") + if mark is not None and np.intp(0).itemsize < 8: + pytest.xfail(f'Fails on our 32-bit test platform(s): {mark.args[0]}') + + # Older versions of threadpoolctl have an issue that may lead to this + # warning being emitted, see gh-14441 + with npt.suppress_warnings() as sup: + sup.filter(pytest.PytestUnraisableExceptionWarning) + + try: + from threadpoolctl import threadpool_limits + + HAS_THREADPOOLCTL = True + except Exception: # observed in gh-14441: (ImportError, AttributeError) + # Optional dependency only. All exceptions are caught, for robustness + HAS_THREADPOOLCTL = False + + if HAS_THREADPOOLCTL: + # Set the number of openmp threads based on the number of workers + # xdist is using to prevent oversubscription. Simplified version of what + # sklearn does (it can rely on threadpoolctl and its builtin OpenMP helper + # functions) + try: + xdist_worker_count = int(os.environ['PYTEST_XDIST_WORKER_COUNT']) + except KeyError: + # raises when pytest-xdist is not installed + return + + if not os.getenv('OMP_NUM_THREADS'): + max_openmp_threads = os.cpu_count() // 2 # use nr of physical cores + threads_per_worker = max(max_openmp_threads // xdist_worker_count, 1) + try: + threadpool_limits(threads_per_worker, user_api='blas') + except Exception: + # May raise AttributeError for older versions of OpenBLAS. + # Catch any error for robustness. + return + + +@pytest.fixture(scope="function", autouse=True) +def check_fpu_mode(request): + """ + Check FPU mode was not changed during the test. + """ + old_mode = get_fpu_mode() + yield + new_mode = get_fpu_mode() + + if old_mode != new_mode: + warnings.warn(f"FPU mode changed from {old_mode:#x} to {new_mode:#x} during " + "the test", + category=FPUModeChangeWarning, stacklevel=0) + + +if not PARALLEL_RUN_AVAILABLE: + @pytest.fixture + def num_parallel_threads(): + return 1 + + +# Array API backend handling +xp_available_backends = {'numpy': np} + +if SCIPY_ARRAY_API and isinstance(SCIPY_ARRAY_API, str): + # fill the dict of backends with available libraries + try: + import array_api_strict + xp_available_backends.update({'array_api_strict': array_api_strict}) + if _pep440.parse(array_api_strict.__version__) < _pep440.Version('2.0'): + raise ImportError("array-api-strict must be >= version 2.0") + array_api_strict.set_array_api_strict_flags( + api_version='2023.12' + ) + except ImportError: + pass + + try: + import torch # type: ignore[import-not-found] + xp_available_backends.update({'torch': torch}) + # can use `mps` or `cpu` + torch.set_default_device(SCIPY_DEVICE) + except ImportError: + pass + + try: + import cupy # type: ignore[import-not-found] + xp_available_backends.update({'cupy': cupy}) + except ImportError: + pass + + try: + import jax.numpy # type: ignore[import-not-found] + xp_available_backends.update({'jax.numpy': jax.numpy}) + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_default_device", jax.devices(SCIPY_DEVICE)[0]) + except ImportError: + pass + + # by default, use all available backends + if SCIPY_ARRAY_API.lower() not in ("1", "true"): + SCIPY_ARRAY_API_ = json.loads(SCIPY_ARRAY_API) + + if 'all' in SCIPY_ARRAY_API_: + pass # same as True + else: + # only select a subset of backend by filtering out the dict + try: + xp_available_backends = { + backend: xp_available_backends[backend] + for backend in SCIPY_ARRAY_API_ + } + except KeyError: + msg = f"'--array-api-backend' must be in {xp_available_backends.keys()}" + raise ValueError(msg) + +if 'cupy' in xp_available_backends: + SCIPY_DEVICE = 'cuda' + +array_api_compatible = pytest.mark.parametrize("xp", xp_available_backends.values()) + +skip_xp_invalid_arg = pytest.mark.skipif(SCIPY_ARRAY_API, + reason = ('Test involves masked arrays, object arrays, or other types ' + 'that are not valid input when `SCIPY_ARRAY_API` is used.')) + + +def _backends_kwargs_from_request(request, skip_or_xfail): + """A helper for {skip,xfail}_xp_backends""" + # do not allow multiple backends + args_ = request.keywords[f'{skip_or_xfail}_xp_backends'].args + if len(args_) > 1: + # np_only / cpu_only has args=(), otherwise it's ('numpy',) + # and we do not allow ('numpy', 'cupy') + raise ValueError(f"multiple backends: {args_}") + + markers = list(request.node.iter_markers(f'{skip_or_xfail}_xp_backends')) + backends = [] + kwargs = {} + for marker in markers: + if marker.kwargs.get('np_only'): + kwargs['np_only'] = True + kwargs['exceptions'] = marker.kwargs.get('exceptions', []) + elif marker.kwargs.get('cpu_only'): + if not kwargs.get('np_only'): + # if np_only is given, it is certainly cpu only + kwargs['cpu_only'] = True + kwargs['exceptions'] = marker.kwargs.get('exceptions', []) + + # add backends, if any + if len(marker.args) > 0: + backend = marker.args[0] # was a tuple, ('numpy',) etc + backends.append(backend) + kwargs.update(**{backend: marker.kwargs}) + + return backends, kwargs + + +@pytest.fixture +def skip_xp_backends(xp, request): + """skip_xp_backends(backend=None, reason=None, np_only=False, cpu_only=False, exceptions=None) + + Skip a decorated test for the provided backend, or skip a category of backends. + + See ``skip_or_xfail_backends`` docstring for details. Note that, contrary to + ``skip_or_xfail_backends``, the ``backend`` and ``reason`` arguments are optional + single strings: this function only skips a single backend at a time. + To skip multiple backends, provide multiple decorators. + """ # noqa: E501 + if "skip_xp_backends" not in request.keywords: + return + + backends, kwargs = _backends_kwargs_from_request(request, skip_or_xfail='skip') + skip_or_xfail_xp_backends(xp, backends, kwargs, skip_or_xfail='skip') + + +@pytest.fixture +def xfail_xp_backends(xp, request): + """xfail_xp_backends(backend=None, reason=None, np_only=False, cpu_only=False, exceptions=None) + + xfail a decorated test for the provided backend, or xfail a category of backends. + + See ``skip_or_xfail_backends`` docstring for details. Note that, contrary to + ``skip_or_xfail_backends``, the ``backend`` and ``reason`` arguments are optional + single strings: this function only xfails a single backend at a time. + To xfail multiple backends, provide multiple decorators. + """ # noqa: E501 + if "xfail_xp_backends" not in request.keywords: + return + backends, kwargs = _backends_kwargs_from_request(request, skip_or_xfail='xfail') + skip_or_xfail_xp_backends(xp, backends, kwargs, skip_or_xfail='xfail') + + +def skip_or_xfail_xp_backends(xp, backends, kwargs, skip_or_xfail='skip'): + """ + Skip based on the ``skip_xp_backends`` or ``xfail_xp_backends`` marker. + + See the "Support for the array API standard" docs page for usage examples. + + Parameters + ---------- + backends : tuple + Backends to skip/xfail, e.g. ``("array_api_strict", "torch")``. + These are overriden when ``np_only`` is ``True``, and are not + necessary to provide for non-CPU backends when ``cpu_only`` is ``True``. + For a custom reason to apply, you should pass a dict ``{'reason': '...'}`` + to a keyword matching the name of the backend. + reason : str, optional + A reason for the skip/xfail in the case of ``np_only=True``. + If unprovided, a default reason is used. Note that it is not possible + to specify a custom reason with ``cpu_only``. + np_only : bool, optional + When ``True``, the test is skipped/xfailed for all backends other + than the default NumPy backend. There is no need to provide + any ``backends`` in this case. To specify a reason, pass a + value to ``reason``. Default: ``False``. + cpu_only : bool, optional + When ``True``, the test is skipped/xfailed on non-CPU devices. + There is no need to provide any ``backends`` in this case, + but any ``backends`` will also be skipped on the CPU. + Default: ``False``. + exceptions : list, optional + A list of exceptions for use with ``cpu_only`` or ``np_only``. + This should be provided when delegation is implemented for some, + but not all, non-CPU/non-NumPy backends. + skip_or_xfail : str + ``'skip'`` to skip, ``'xfail'`` to xfail. + """ + skip_or_xfail = getattr(pytest, skip_or_xfail) + np_only = kwargs.get("np_only", False) + cpu_only = kwargs.get("cpu_only", False) + exceptions = kwargs.get("exceptions", []) + + if reasons := kwargs.get("reasons"): + raise ValueError(f"provide a single `reason=` kwarg; got {reasons=} instead") + + # input validation + if np_only and cpu_only: + # np_only is a stricter subset of cpu_only + cpu_only = False + if exceptions and not (cpu_only or np_only): + raise ValueError("`exceptions` is only valid alongside `cpu_only` or `np_only`") + + if np_only: + reason = kwargs.get("reason", "do not run with non-NumPy backends.") + if not isinstance(reason, str) and len(reason) > 1: + raise ValueError("please provide a singleton `reason` " + "when using `np_only`") + if xp.__name__ != 'numpy' and xp.__name__ not in exceptions: + skip_or_xfail(reason=reason) + return + if cpu_only: + reason = ("no array-agnostic implementation or delegation available " + "for this backend and device") + exceptions = [] if exceptions is None else exceptions + if SCIPY_ARRAY_API and SCIPY_DEVICE != 'cpu': + if xp.__name__ == 'cupy' and 'cupy' not in exceptions: + skip_or_xfail(reason=reason) + elif xp.__name__ == 'torch' and 'torch' not in exceptions: + if 'cpu' not in xp.empty(0).device.type: + skip_or_xfail(reason=reason) + elif xp.__name__ == 'jax.numpy' and 'jax.numpy' not in exceptions: + for d in xp.empty(0).devices(): + if 'cpu' not in d.device_kind: + skip_or_xfail(reason=reason) + + if backends is not None: + for i, backend in enumerate(backends): + if xp.__name__ == backend: + reason = kwargs[backend].get('reason') + if not reason: + reason = f"do not run with array API backend: {backend}" + + skip_or_xfail(reason=reason) + + +# Following the approach of NumPy's conftest.py... +# Use a known and persistent tmpdir for hypothesis' caches, which +# can be automatically cleared by the OS or user. +hypothesis.configuration.set_hypothesis_home_dir( + os.path.join(tempfile.gettempdir(), ".hypothesis") +) + +# We register two custom profiles for SciPy - for details see +# https://hypothesis.readthedocs.io/en/latest/settings.html +# The first is designed for our own CI runs; the latter also +# forces determinism and is designed for use via scipy.test() +hypothesis.settings.register_profile( + name="nondeterministic", deadline=None, print_blob=True, +) +hypothesis.settings.register_profile( + name="deterministic", + deadline=None, print_blob=True, database=None, derandomize=True, + suppress_health_check=list(hypothesis.HealthCheck), +) + +# Profile is currently set by environment variable `SCIPY_HYPOTHESIS_PROFILE` +# In the future, it would be good to work the choice into dev.py. +SCIPY_HYPOTHESIS_PROFILE = os.environ.get("SCIPY_HYPOTHESIS_PROFILE", + "deterministic") +hypothesis.settings.load_profile(SCIPY_HYPOTHESIS_PROFILE) + + +############################################################################ +# doctesting stuff + +if HAVE_SCPDT: + + # FIXME: populate the dict once + @contextmanager + def warnings_errors_and_rng(test=None): + """Temporarily turn (almost) all warnings to errors. + + Filter out known warnings which we allow. + """ + known_warnings = dict() + + # these functions are known to emit "divide by zero" RuntimeWarnings + divide_by_zero = [ + 'scipy.linalg.norm', 'scipy.ndimage.center_of_mass', + ] + for name in divide_by_zero: + known_warnings[name] = dict(category=RuntimeWarning, + message='divide by zero') + + # Deprecated stuff in scipy.signal and elsewhere + deprecated = [ + 'scipy.signal.cwt', 'scipy.signal.morlet', 'scipy.signal.morlet2', + 'scipy.signal.ricker', + 'scipy.integrate.simpson', + 'scipy.interpolate.interp2d', + 'scipy.linalg.kron', + ] + for name in deprecated: + known_warnings[name] = dict(category=DeprecationWarning) + + from scipy import integrate + # the functions are known to emit IntegrationWarnings + integration_w = ['scipy.special.ellip_normal', + 'scipy.special.ellip_harm_2', + ] + for name in integration_w: + known_warnings[name] = dict(category=integrate.IntegrationWarning, + message='The occurrence of roundoff') + + # scipy.stats deliberately emits UserWarnings sometimes + user_w = ['scipy.stats.anderson_ksamp', 'scipy.stats.kurtosistest', + 'scipy.stats.normaltest', 'scipy.sparse.linalg.norm'] + for name in user_w: + known_warnings[name] = dict(category=UserWarning) + + # additional one-off warnings to filter + dct = { + 'scipy.sparse.linalg.norm': + dict(category=UserWarning, message="Exited at iteration"), + # tutorials + 'linalg.rst': + dict(message='the matrix subclass is not', + category=PendingDeprecationWarning), + 'stats.rst': + dict(message='The maximum number of subdivisions', + category=integrate.IntegrationWarning), + } + known_warnings.update(dct) + + # these legitimately emit warnings in examples + legit = set('scipy.signal.normalize') + + # Now, the meat of the matter: filter warnings, + # also control the random seed for each doctest. + + # XXX: this matches the refguide-check behavior, but is a tad strange: + # makes sure that the seed the old-fashioned np.random* methods is + # *NOT* reproducible but the new-style `default_rng()` *IS* repoducible. + # Should these two be either both repro or both not repro? + + from scipy._lib._util import _fixed_default_rng + import numpy as np + with _fixed_default_rng(): + np.random.seed(None) + with warnings.catch_warnings(): + if test and test.name in known_warnings: + warnings.filterwarnings('ignore', + **known_warnings[test.name]) + yield + elif test and test.name in legit: + yield + else: + warnings.simplefilter('error', Warning) + yield + + dt_config.user_context_mgr = warnings_errors_and_rng + dt_config.skiplist = set([ + 'scipy.linalg.LinAlgError', # comes from numpy + 'scipy.fftpack.fftshift', # fftpack stuff is also from numpy + 'scipy.fftpack.ifftshift', + 'scipy.fftpack.fftfreq', + 'scipy.special.sinc', # sinc is from numpy + 'scipy.optimize.show_options', # does not have much to doctest + 'scipy.signal.normalize', # manipulates warnings (XXX temp skip) + 'scipy.sparse.linalg.norm', # XXX temp skip + # these below test things which inherit from np.ndarray + # cross-ref https://github.com/numpy/numpy/issues/28019 + 'scipy.io.matlab.MatlabObject.strides', + 'scipy.io.matlab.MatlabObject.dtype', + 'scipy.io.matlab.MatlabOpaque.dtype', + 'scipy.io.matlab.MatlabOpaque.strides', + 'scipy.io.matlab.MatlabFunction.strides', + 'scipy.io.matlab.MatlabFunction.dtype' + ]) + + # these are affected by NumPy 2.0 scalar repr: rely on string comparison + if np.__version__ < "2": + dt_config.skiplist.update(set([ + 'scipy.io.hb_read', + 'scipy.io.hb_write', + 'scipy.sparse.csgraph.connected_components', + 'scipy.sparse.csgraph.depth_first_order', + 'scipy.sparse.csgraph.shortest_path', + 'scipy.sparse.csgraph.floyd_warshall', + 'scipy.sparse.csgraph.dijkstra', + 'scipy.sparse.csgraph.bellman_ford', + 'scipy.sparse.csgraph.johnson', + 'scipy.sparse.csgraph.yen', + 'scipy.sparse.csgraph.breadth_first_order', + 'scipy.sparse.csgraph.reverse_cuthill_mckee', + 'scipy.sparse.csgraph.structural_rank', + 'scipy.sparse.csgraph.construct_dist_matrix', + 'scipy.sparse.csgraph.reconstruct_path', + 'scipy.ndimage.value_indices', + 'scipy.stats.mstats.describe', + ])) + + # help pytest collection a bit: these names are either private + # (distributions), or just do not need doctesting. + dt_config.pytest_extra_ignore = [ + "scipy.stats.distributions", + "scipy.optimize.cython_optimize", + "scipy.test", + "scipy.show_config", + # equivalent to "pytest --ignore=path/to/file" + "scipy/special/_precompute", + "scipy/interpolate/_interpnd_info.py", + "scipy/_lib/array_api_compat", + "scipy/_lib/highs", + "scipy/_lib/unuran", + "scipy/_lib/_gcutils.py", + "scipy/_lib/doccer.py", + "scipy/_lib/_uarray", + ] + + dt_config.pytest_extra_xfail = { + # name: reason + "ND_regular_grid.rst": "ReST parser limitation", + "extrapolation_examples.rst": "ReST parser limitation", + "sampling_pinv.rst": "__cinit__ unexpected argument", + "sampling_srou.rst": "nan in scalar_power", + "probability_distributions.rst": "integration warning", + } + + # tutorials + dt_config.pseudocode = set(['integrate.nquad(func,']) + dt_config.local_resources = { + 'io.rst': [ + "octave_a.mat", + "octave_cells.mat", + "octave_struct.mat" + ] + } + + dt_config.strict_check = True +############################################################################ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ec4ed081f9ed1dfc3adb382793d322c9d400f794 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/__init__.py @@ -0,0 +1,358 @@ +r""" +================================== +Constants (:mod:`scipy.constants`) +================================== + +.. currentmodule:: scipy.constants + +Physical and mathematical constants and units. + + +Mathematical constants +====================== + +================ ================================================================= +``pi`` Pi +``golden`` Golden ratio +``golden_ratio`` Golden ratio +================ ================================================================= + + +Physical constants +================== +The following physical constants are available as attributes of `scipy.constants`. +All units are `SI `_. + +=========================== ================================================================ =============== +Attribute Quantity Units +=========================== ================================================================ =============== +``c`` speed of light in vacuum m s^-1 +``speed_of_light`` speed of light in vacuum m s^-1 +``mu_0`` the magnetic constant :math:`\mu_0` N A^-2 +``epsilon_0`` the electric constant (vacuum permittivity), :math:`\epsilon_0` F m^-1 +``h`` the Planck constant :math:`h` J Hz^-1 +``Planck`` the Planck constant :math:`h` J Hz^-1 +``hbar`` the reduced Planck constant, :math:`\hbar = h/(2\pi)` J s +``G`` Newtonian constant of gravitation m^3 kg^-1 s^-2 +``gravitational_constant`` Newtonian constant of gravitation m^3 kg^-1 s^-2 +``g`` standard acceleration of gravity m s^-2 +``e`` elementary charge C +``elementary_charge`` elementary charge C +``R`` molar gas constant J mol^-1 K^-1 +``gas_constant`` molar gas constant J mol^-1 K^-1 +``alpha`` fine-structure constant (unitless) +``fine_structure`` fine-structure constant (unitless) +``N_A`` Avogadro constant mol^-1 +``Avogadro`` Avogadro constant mol^-1 +``k`` Boltzmann constant J K^-1 +``Boltzmann`` Boltzmann constant J K^-1 +``sigma`` Stefan-Boltzmann constant :math:`\sigma` W m^-2 K^-4 +``Stefan_Boltzmann`` Stefan-Boltzmann constant :math:`\sigma` W m^-2 K^-4 +``Wien`` Wien wavelength displacement law constant m K +``Rydberg`` Rydberg constant m^-1 +``m_e`` electron mass kg +``electron_mass`` electron mass kg +``m_p`` proton mass kg +``proton_mass`` proton mass kg +``m_n`` neutron mass kg +``neutron_mass`` neutron mass kg +=========================== ================================================================ =============== + + +Constants database +------------------ + +In addition to the above variables, :mod:`scipy.constants` also contains the +2022 CODATA recommended values [CODATA2022]_ database containing more physical +constants. + +.. autosummary:: + :toctree: generated/ + + value -- Value in physical_constants indexed by key + unit -- Unit in physical_constants indexed by key + precision -- Relative precision in physical_constants indexed by key + find -- Return list of physical_constant keys with a given string + ConstantWarning -- Constant sought not in newest CODATA data set + +.. data:: physical_constants + + Dictionary of physical constants, of the format + ``physical_constants[name] = (value, unit, uncertainty)``. + The CODATA database uses ellipses to indicate that a value is defined + (exactly) in terms of others but cannot be represented exactly with the + allocated number of digits. In these cases, SciPy calculates the derived + value and reports it to the full precision of a Python ``float``. Although + ``physical_constants`` lists the uncertainty as ``0.0`` to indicate that + the CODATA value is exact, the value in ``physical_constants`` is still + subject to the truncation error inherent in double-precision representation. + +Available constants: + +====================================================================== ==== +%(constant_names)s +====================================================================== ==== + + +Units +===== + +SI prefixes +----------- + +============ ================================================================= +``quetta`` :math:`10^{30}` +``ronna`` :math:`10^{27}` +``yotta`` :math:`10^{24}` +``zetta`` :math:`10^{21}` +``exa`` :math:`10^{18}` +``peta`` :math:`10^{15}` +``tera`` :math:`10^{12}` +``giga`` :math:`10^{9}` +``mega`` :math:`10^{6}` +``kilo`` :math:`10^{3}` +``hecto`` :math:`10^{2}` +``deka`` :math:`10^{1}` +``deci`` :math:`10^{-1}` +``centi`` :math:`10^{-2}` +``milli`` :math:`10^{-3}` +``micro`` :math:`10^{-6}` +``nano`` :math:`10^{-9}` +``pico`` :math:`10^{-12}` +``femto`` :math:`10^{-15}` +``atto`` :math:`10^{-18}` +``zepto`` :math:`10^{-21}` +``yocto`` :math:`10^{-24}` +``ronto`` :math:`10^{-27}` +``quecto`` :math:`10^{-30}` +============ ================================================================= + +Binary prefixes +--------------- + +============ ================================================================= +``kibi`` :math:`2^{10}` +``mebi`` :math:`2^{20}` +``gibi`` :math:`2^{30}` +``tebi`` :math:`2^{40}` +``pebi`` :math:`2^{50}` +``exbi`` :math:`2^{60}` +``zebi`` :math:`2^{70}` +``yobi`` :math:`2^{80}` +============ ================================================================= + +Mass +---- + +================= ============================================================ +``gram`` :math:`10^{-3}` kg +``metric_ton`` :math:`10^{3}` kg +``grain`` one grain in kg +``lb`` one pound (avoirdupous) in kg +``pound`` one pound (avoirdupous) in kg +``blob`` one inch version of a slug in kg (added in 1.0.0) +``slinch`` one inch version of a slug in kg (added in 1.0.0) +``slug`` one slug in kg (added in 1.0.0) +``oz`` one ounce in kg +``ounce`` one ounce in kg +``stone`` one stone in kg +``grain`` one grain in kg +``long_ton`` one long ton in kg +``short_ton`` one short ton in kg +``troy_ounce`` one Troy ounce in kg +``troy_pound`` one Troy pound in kg +``carat`` one carat in kg +``m_u`` atomic mass constant (in kg) +``u`` atomic mass constant (in kg) +``atomic_mass`` atomic mass constant (in kg) +================= ============================================================ + +Angle +----- + +================= ============================================================ +``degree`` degree in radians +``arcmin`` arc minute in radians +``arcminute`` arc minute in radians +``arcsec`` arc second in radians +``arcsecond`` arc second in radians +================= ============================================================ + + +Time +---- + +================= ============================================================ +``minute`` one minute in seconds +``hour`` one hour in seconds +``day`` one day in seconds +``week`` one week in seconds +``year`` one year (365 days) in seconds +``Julian_year`` one Julian year (365.25 days) in seconds +================= ============================================================ + + +Length +------ + +===================== ============================================================ +``inch`` one inch in meters +``foot`` one foot in meters +``yard`` one yard in meters +``mile`` one mile in meters +``mil`` one mil in meters +``pt`` one point in meters +``point`` one point in meters +``survey_foot`` one survey foot in meters +``survey_mile`` one survey mile in meters +``nautical_mile`` one nautical mile in meters +``fermi`` one Fermi in meters +``angstrom`` one Angstrom in meters +``micron`` one micron in meters +``au`` one astronomical unit in meters +``astronomical_unit`` one astronomical unit in meters +``light_year`` one light year in meters +``parsec`` one parsec in meters +===================== ============================================================ + +Pressure +-------- + +================= ============================================================ +``atm`` standard atmosphere in pascals +``atmosphere`` standard atmosphere in pascals +``bar`` one bar in pascals +``torr`` one torr (mmHg) in pascals +``mmHg`` one torr (mmHg) in pascals +``psi`` one psi in pascals +================= ============================================================ + +Area +---- + +================= ============================================================ +``hectare`` one hectare in square meters +``acre`` one acre in square meters +================= ============================================================ + + +Volume +------ + +=================== ======================================================== +``liter`` one liter in cubic meters +``litre`` one liter in cubic meters +``gallon`` one gallon (US) in cubic meters +``gallon_US`` one gallon (US) in cubic meters +``gallon_imp`` one gallon (UK) in cubic meters +``fluid_ounce`` one fluid ounce (US) in cubic meters +``fluid_ounce_US`` one fluid ounce (US) in cubic meters +``fluid_ounce_imp`` one fluid ounce (UK) in cubic meters +``bbl`` one barrel in cubic meters +``barrel`` one barrel in cubic meters +=================== ======================================================== + +Speed +----- + +================== ========================================================== +``kmh`` kilometers per hour in meters per second +``mph`` miles per hour in meters per second +``mach`` one Mach (approx., at 15 C, 1 atm) in meters per second +``speed_of_sound`` one Mach (approx., at 15 C, 1 atm) in meters per second +``knot`` one knot in meters per second +================== ========================================================== + + +Temperature +----------- + +===================== ======================================================= +``zero_Celsius`` zero of Celsius scale in Kelvin +``degree_Fahrenheit`` one Fahrenheit (only differences) in Kelvins +===================== ======================================================= + +.. autosummary:: + :toctree: generated/ + + convert_temperature + +Energy +------ + +==================== ======================================================= +``eV`` one electron volt in Joules +``electron_volt`` one electron volt in Joules +``calorie`` one calorie (thermochemical) in Joules +``calorie_th`` one calorie (thermochemical) in Joules +``calorie_IT`` one calorie (International Steam Table calorie, 1956) in Joules +``erg`` one erg in Joules +``Btu`` one British thermal unit (International Steam Table) in Joules +``Btu_IT`` one British thermal unit (International Steam Table) in Joules +``Btu_th`` one British thermal unit (thermochemical) in Joules +``ton_TNT`` one ton of TNT in Joules +==================== ======================================================= + +Power +----- + +==================== ======================================================= +``hp`` one horsepower in watts +``horsepower`` one horsepower in watts +==================== ======================================================= + +Force +----- + +==================== ======================================================= +``dyn`` one dyne in newtons +``dyne`` one dyne in newtons +``lbf`` one pound force in newtons +``pound_force`` one pound force in newtons +``kgf`` one kilogram force in newtons +``kilogram_force`` one kilogram force in newtons +==================== ======================================================= + +Optics +------ + +.. autosummary:: + :toctree: generated/ + + lambda2nu + nu2lambda + +References +========== + +.. [CODATA2022] CODATA Recommended Values of the Fundamental + Physical Constants 2022. + + https://physics.nist.gov/cuu/Constants/ + +""" # noqa: E501 +# Modules contributed by BasSw (wegwerp@gmail.com) +from ._codata import * +from ._constants import * +from ._codata import _obsolete_constants, physical_constants + +# Deprecated namespaces, to be removed in v2.0.0 +from . import codata, constants + +_constant_names_list = [(_k.lower(), _k, _v) + for _k, _v in physical_constants.items() + if _k not in _obsolete_constants] +_constant_names = "\n".join(["``{}``{} {} {}".format(_x[1], " "*(66-len(_x[1])), + _x[2][0], _x[2][1]) + for _x in sorted(_constant_names_list)]) +if __doc__: + __doc__ = __doc__ % dict(constant_names=_constant_names) + +del _constant_names +del _constant_names_list + +__all__ = [s for s in dir() if not s.startswith('_')] + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/_codata.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/_codata.py new file mode 100644 index 0000000000000000000000000000000000000000..03fb5bc0a5715610e67972937a404261c9c481eb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/_codata.py @@ -0,0 +1,2266 @@ +""" +Fundamental Physical Constants +------------------------------ + +These constants are taken from CODATA Recommended Values of the Fundamental +Physical Constants 2022. + +Object +------ +physical_constants : dict + A dictionary containing physical constants. Keys are the names of physical + constants, values are tuples (value, units, precision). + +Functions +--------- +value(key): + Returns the value of the physical constant(key). +unit(key): + Returns the units of the physical constant(key). +precision(key): + Returns the relative precision of the physical constant(key). +find(sub): + Prints or returns list of keys containing the string sub, default is all. + +Source +------ +The values of the constants provided at this site are recommended for +international use by CODATA and are the latest available. Termed the "2018 +CODATA recommended values," they are generally recognized worldwide for use in +all fields of science and technology. The values became available on 20 May +2019 and replaced the 2014 CODATA set. Also available is an introduction to the +constants for non-experts at + +https://physics.nist.gov/cuu/Constants/introduction.html + +References +---------- +Theoretical and experimental publications relevant to the fundamental constants +and closely related precision measurements published since the mid 1980s, but +also including many older papers of particular interest, some of which date +back to the 1800s. To search the bibliography, visit + +https://physics.nist.gov/cuu/Constants/ + +""" + +# Compiled by Charles Harris, dated October 3, 2002 +# updated to 2002 values by BasSw, 2006 +# Updated to 2006 values by Vincent Davis June 2010 +# Updated to 2014 values by Joseph Booker, 2015 +# Updated to 2018 values by Jakob Jakobson, 2019 +# Updated to 2022 values by Jakob Jakobson, 2024 + +import warnings +import math + +from typing import Any +from collections.abc import Callable + +__all__ = ['physical_constants', 'value', 'unit', 'precision', 'find', + 'ConstantWarning'] + +""" +Source: https://physics.nist.gov/cuu/Constants/ + +The values of the constants provided at this site are recommended for +international use by CODATA and are the latest available. Termed the "2018 +CODATA recommended values," they are generally recognized worldwide for use in +all fields of science and technology. The values became available on 20 May +2019 and replaced the 2014 CODATA set. +""" + +# +# Source: https://physics.nist.gov/cuu/Constants/ +# + +# Quantity Value Uncertainty Unit +# ---------------------------------------------------- --------------------- -------------------- ------------- +txt2002 = """\ +Wien displacement law constant 2.897 7685e-3 0.000 0051e-3 m K +atomic unit of 1st hyperpolarizablity 3.206 361 51e-53 0.000 000 28e-53 C^3 m^3 J^-2 +atomic unit of 2nd hyperpolarizablity 6.235 3808e-65 0.000 0011e-65 C^4 m^4 J^-3 +atomic unit of electric dipole moment 8.478 353 09e-30 0.000 000 73e-30 C m +atomic unit of electric polarizablity 1.648 777 274e-41 0.000 000 016e-41 C^2 m^2 J^-1 +atomic unit of electric quadrupole moment 4.486 551 24e-40 0.000 000 39e-40 C m^2 +atomic unit of magn. dipole moment 1.854 801 90e-23 0.000 000 16e-23 J T^-1 +atomic unit of magn. flux density 2.350 517 42e5 0.000 000 20e5 T +deuteron magn. moment 0.433 073 482e-26 0.000 000 038e-26 J T^-1 +deuteron magn. moment to Bohr magneton ratio 0.466 975 4567e-3 0.000 000 0050e-3 +deuteron magn. moment to nuclear magneton ratio 0.857 438 2329 0.000 000 0092 +deuteron-electron magn. moment ratio -4.664 345 548e-4 0.000 000 050e-4 +deuteron-proton magn. moment ratio 0.307 012 2084 0.000 000 0045 +deuteron-neutron magn. moment ratio -0.448 206 52 0.000 000 11 +electron gyromagn. ratio 1.760 859 74e11 0.000 000 15e11 s^-1 T^-1 +electron gyromagn. ratio over 2 pi 28 024.9532 0.0024 MHz T^-1 +electron magn. moment -928.476 412e-26 0.000 080e-26 J T^-1 +electron magn. moment to Bohr magneton ratio -1.001 159 652 1859 0.000 000 000 0038 +electron magn. moment to nuclear magneton ratio -1838.281 971 07 0.000 000 85 +electron magn. moment anomaly 1.159 652 1859e-3 0.000 000 0038e-3 +electron to shielded proton magn. moment ratio -658.227 5956 0.000 0071 +electron to shielded helion magn. moment ratio 864.058 255 0.000 010 +electron-deuteron magn. moment ratio -2143.923 493 0.000 023 +electron-muon magn. moment ratio 206.766 9894 0.000 0054 +electron-neutron magn. moment ratio 960.920 50 0.000 23 +electron-proton magn. moment ratio -658.210 6862 0.000 0066 +magn. constant 12.566 370 614...e-7 (exact) N A^-2 +magn. flux quantum 2.067 833 72e-15 0.000 000 18e-15 Wb +muon magn. moment -4.490 447 99e-26 0.000 000 40e-26 J T^-1 +muon magn. moment to Bohr magneton ratio -4.841 970 45e-3 0.000 000 13e-3 +muon magn. moment to nuclear magneton ratio -8.890 596 98 0.000 000 23 +muon-proton magn. moment ratio -3.183 345 118 0.000 000 089 +neutron gyromagn. ratio 1.832 471 83e8 0.000 000 46e8 s^-1 T^-1 +neutron gyromagn. ratio over 2 pi 29.164 6950 0.000 0073 MHz T^-1 +neutron magn. moment -0.966 236 45e-26 0.000 000 24e-26 J T^-1 +neutron magn. moment to Bohr magneton ratio -1.041 875 63e-3 0.000 000 25e-3 +neutron magn. moment to nuclear magneton ratio -1.913 042 73 0.000 000 45 +neutron to shielded proton magn. moment ratio -0.684 996 94 0.000 000 16 +neutron-electron magn. moment ratio 1.040 668 82e-3 0.000 000 25e-3 +neutron-proton magn. moment ratio -0.684 979 34 0.000 000 16 +proton gyromagn. ratio 2.675 222 05e8 0.000 000 23e8 s^-1 T^-1 +proton gyromagn. ratio over 2 pi 42.577 4813 0.000 0037 MHz T^-1 +proton magn. moment 1.410 606 71e-26 0.000 000 12e-26 J T^-1 +proton magn. moment to Bohr magneton ratio 1.521 032 206e-3 0.000 000 015e-3 +proton magn. moment to nuclear magneton ratio 2.792 847 351 0.000 000 028 +proton magn. shielding correction 25.689e-6 0.015e-6 +proton-neutron magn. moment ratio -1.459 898 05 0.000 000 34 +shielded helion gyromagn. ratio 2.037 894 70e8 0.000 000 18e8 s^-1 T^-1 +shielded helion gyromagn. ratio over 2 pi 32.434 1015 0.000 0028 MHz T^-1 +shielded helion magn. moment -1.074 553 024e-26 0.000 000 093e-26 J T^-1 +shielded helion magn. moment to Bohr magneton ratio -1.158 671 474e-3 0.000 000 014e-3 +shielded helion magn. moment to nuclear magneton ratio -2.127 497 723 0.000 000 025 +shielded helion to proton magn. moment ratio -0.761 766 562 0.000 000 012 +shielded helion to shielded proton magn. moment ratio -0.761 786 1313 0.000 000 0033 +shielded helion gyromagn. ratio 2.037 894 70e8 0.000 000 18e8 s^-1 T^-1 +shielded helion gyromagn. ratio over 2 pi 32.434 1015 0.000 0028 MHz T^-1 +shielded proton magn. moment 1.410 570 47e-26 0.000 000 12e-26 J T^-1 +shielded proton magn. moment to Bohr magneton ratio 1.520 993 132e-3 0.000 000 016e-3 +shielded proton magn. moment to nuclear magneton ratio 2.792 775 604 0.000 000 030 +{220} lattice spacing of silicon 192.015 5965e-12 0.000 0070e-12 m""" + + +def exact2002(exact): + replace = { + 'magn. constant': 4e-7 * math.pi, + } + return replace + + +txt2006 = """\ +lattice spacing of silicon 192.015 5762 e-12 0.000 0050 e-12 m +alpha particle-electron mass ratio 7294.299 5365 0.000 0031 +alpha particle mass 6.644 656 20 e-27 0.000 000 33 e-27 kg +alpha particle mass energy equivalent 5.971 919 17 e-10 0.000 000 30 e-10 J +alpha particle mass energy equivalent in MeV 3727.379 109 0.000 093 MeV +alpha particle mass in u 4.001 506 179 127 0.000 000 000 062 u +alpha particle molar mass 4.001 506 179 127 e-3 0.000 000 000 062 e-3 kg mol^-1 +alpha particle-proton mass ratio 3.972 599 689 51 0.000 000 000 41 +Angstrom star 1.000 014 98 e-10 0.000 000 90 e-10 m +atomic mass constant 1.660 538 782 e-27 0.000 000 083 e-27 kg +atomic mass constant energy equivalent 1.492 417 830 e-10 0.000 000 074 e-10 J +atomic mass constant energy equivalent in MeV 931.494 028 0.000 023 MeV +atomic mass unit-electron volt relationship 931.494 028 e6 0.000 023 e6 eV +atomic mass unit-hartree relationship 3.423 177 7149 e7 0.000 000 0049 e7 E_h +atomic mass unit-hertz relationship 2.252 342 7369 e23 0.000 000 0032 e23 Hz +atomic mass unit-inverse meter relationship 7.513 006 671 e14 0.000 000 011 e14 m^-1 +atomic mass unit-joule relationship 1.492 417 830 e-10 0.000 000 074 e-10 J +atomic mass unit-kelvin relationship 1.080 9527 e13 0.000 0019 e13 K +atomic mass unit-kilogram relationship 1.660 538 782 e-27 0.000 000 083 e-27 kg +atomic unit of 1st hyperpolarizability 3.206 361 533 e-53 0.000 000 081 e-53 C^3 m^3 J^-2 +atomic unit of 2nd hyperpolarizability 6.235 380 95 e-65 0.000 000 31 e-65 C^4 m^4 J^-3 +atomic unit of action 1.054 571 628 e-34 0.000 000 053 e-34 J s +atomic unit of charge 1.602 176 487 e-19 0.000 000 040 e-19 C +atomic unit of charge density 1.081 202 300 e12 0.000 000 027 e12 C m^-3 +atomic unit of current 6.623 617 63 e-3 0.000 000 17 e-3 A +atomic unit of electric dipole mom. 8.478 352 81 e-30 0.000 000 21 e-30 C m +atomic unit of electric field 5.142 206 32 e11 0.000 000 13 e11 V m^-1 +atomic unit of electric field gradient 9.717 361 66 e21 0.000 000 24 e21 V m^-2 +atomic unit of electric polarizability 1.648 777 2536 e-41 0.000 000 0034 e-41 C^2 m^2 J^-1 +atomic unit of electric potential 27.211 383 86 0.000 000 68 V +atomic unit of electric quadrupole mom. 4.486 551 07 e-40 0.000 000 11 e-40 C m^2 +atomic unit of energy 4.359 743 94 e-18 0.000 000 22 e-18 J +atomic unit of force 8.238 722 06 e-8 0.000 000 41 e-8 N +atomic unit of length 0.529 177 208 59 e-10 0.000 000 000 36 e-10 m +atomic unit of mag. dipole mom. 1.854 801 830 e-23 0.000 000 046 e-23 J T^-1 +atomic unit of mag. flux density 2.350 517 382 e5 0.000 000 059 e5 T +atomic unit of magnetizability 7.891 036 433 e-29 0.000 000 027 e-29 J T^-2 +atomic unit of mass 9.109 382 15 e-31 0.000 000 45 e-31 kg +atomic unit of momentum 1.992 851 565 e-24 0.000 000 099 e-24 kg m s^-1 +atomic unit of permittivity 1.112 650 056... e-10 (exact) F m^-1 +atomic unit of time 2.418 884 326 505 e-17 0.000 000 000 016 e-17 s +atomic unit of velocity 2.187 691 2541 e6 0.000 000 0015 e6 m s^-1 +Avogadro constant 6.022 141 79 e23 0.000 000 30 e23 mol^-1 +Bohr magneton 927.400 915 e-26 0.000 023 e-26 J T^-1 +Bohr magneton in eV/T 5.788 381 7555 e-5 0.000 000 0079 e-5 eV T^-1 +Bohr magneton in Hz/T 13.996 246 04 e9 0.000 000 35 e9 Hz T^-1 +Bohr magneton in inverse meters per tesla 46.686 4515 0.000 0012 m^-1 T^-1 +Bohr magneton in K/T 0.671 7131 0.000 0012 K T^-1 +Bohr radius 0.529 177 208 59 e-10 0.000 000 000 36 e-10 m +Boltzmann constant 1.380 6504 e-23 0.000 0024 e-23 J K^-1 +Boltzmann constant in eV/K 8.617 343 e-5 0.000 015 e-5 eV K^-1 +Boltzmann constant in Hz/K 2.083 6644 e10 0.000 0036 e10 Hz K^-1 +Boltzmann constant in inverse meters per kelvin 69.503 56 0.000 12 m^-1 K^-1 +characteristic impedance of vacuum 376.730 313 461... (exact) ohm +classical electron radius 2.817 940 2894 e-15 0.000 000 0058 e-15 m +Compton wavelength 2.426 310 2175 e-12 0.000 000 0033 e-12 m +Compton wavelength over 2 pi 386.159 264 59 e-15 0.000 000 53 e-15 m +conductance quantum 7.748 091 7004 e-5 0.000 000 0053 e-5 S +conventional value of Josephson constant 483 597.9 e9 (exact) Hz V^-1 +conventional value of von Klitzing constant 25 812.807 (exact) ohm +Cu x unit 1.002 076 99 e-13 0.000 000 28 e-13 m +deuteron-electron mag. mom. ratio -4.664 345 537 e-4 0.000 000 039 e-4 +deuteron-electron mass ratio 3670.482 9654 0.000 0016 +deuteron g factor 0.857 438 2308 0.000 000 0072 +deuteron mag. mom. 0.433 073 465 e-26 0.000 000 011 e-26 J T^-1 +deuteron mag. mom. to Bohr magneton ratio 0.466 975 4556 e-3 0.000 000 0039 e-3 +deuteron mag. mom. to nuclear magneton ratio 0.857 438 2308 0.000 000 0072 +deuteron mass 3.343 583 20 e-27 0.000 000 17 e-27 kg +deuteron mass energy equivalent 3.005 062 72 e-10 0.000 000 15 e-10 J +deuteron mass energy equivalent in MeV 1875.612 793 0.000 047 MeV +deuteron mass in u 2.013 553 212 724 0.000 000 000 078 u +deuteron molar mass 2.013 553 212 724 e-3 0.000 000 000 078 e-3 kg mol^-1 +deuteron-neutron mag. mom. ratio -0.448 206 52 0.000 000 11 +deuteron-proton mag. mom. ratio 0.307 012 2070 0.000 000 0024 +deuteron-proton mass ratio 1.999 007 501 08 0.000 000 000 22 +deuteron rms charge radius 2.1402 e-15 0.0028 e-15 m +electric constant 8.854 187 817... e-12 (exact) F m^-1 +electron charge to mass quotient -1.758 820 150 e11 0.000 000 044 e11 C kg^-1 +electron-deuteron mag. mom. ratio -2143.923 498 0.000 018 +electron-deuteron mass ratio 2.724 437 1093 e-4 0.000 000 0012 e-4 +electron g factor -2.002 319 304 3622 0.000 000 000 0015 +electron gyromag. ratio 1.760 859 770 e11 0.000 000 044 e11 s^-1 T^-1 +electron gyromag. ratio over 2 pi 28 024.953 64 0.000 70 MHz T^-1 +electron mag. mom. -928.476 377 e-26 0.000 023 e-26 J T^-1 +electron mag. mom. anomaly 1.159 652 181 11 e-3 0.000 000 000 74 e-3 +electron mag. mom. to Bohr magneton ratio -1.001 159 652 181 11 0.000 000 000 000 74 +electron mag. mom. to nuclear magneton ratio -1838.281 970 92 0.000 000 80 +electron mass 9.109 382 15 e-31 0.000 000 45 e-31 kg +electron mass energy equivalent 8.187 104 38 e-14 0.000 000 41 e-14 J +electron mass energy equivalent in MeV 0.510 998 910 0.000 000 013 MeV +electron mass in u 5.485 799 0943 e-4 0.000 000 0023 e-4 u +electron molar mass 5.485 799 0943 e-7 0.000 000 0023 e-7 kg mol^-1 +electron-muon mag. mom. ratio 206.766 9877 0.000 0052 +electron-muon mass ratio 4.836 331 71 e-3 0.000 000 12 e-3 +electron-neutron mag. mom. ratio 960.920 50 0.000 23 +electron-neutron mass ratio 5.438 673 4459 e-4 0.000 000 0033 e-4 +electron-proton mag. mom. ratio -658.210 6848 0.000 0054 +electron-proton mass ratio 5.446 170 2177 e-4 0.000 000 0024 e-4 +electron-tau mass ratio 2.875 64 e-4 0.000 47 e-4 +electron to alpha particle mass ratio 1.370 933 555 70 e-4 0.000 000 000 58 e-4 +electron to shielded helion mag. mom. ratio 864.058 257 0.000 010 +electron to shielded proton mag. mom. ratio -658.227 5971 0.000 0072 +electron volt 1.602 176 487 e-19 0.000 000 040 e-19 J +electron volt-atomic mass unit relationship 1.073 544 188 e-9 0.000 000 027 e-9 u +electron volt-hartree relationship 3.674 932 540 e-2 0.000 000 092 e-2 E_h +electron volt-hertz relationship 2.417 989 454 e14 0.000 000 060 e14 Hz +electron volt-inverse meter relationship 8.065 544 65 e5 0.000 000 20 e5 m^-1 +electron volt-joule relationship 1.602 176 487 e-19 0.000 000 040 e-19 J +electron volt-kelvin relationship 1.160 4505 e4 0.000 0020 e4 K +electron volt-kilogram relationship 1.782 661 758 e-36 0.000 000 044 e-36 kg +elementary charge 1.602 176 487 e-19 0.000 000 040 e-19 C +elementary charge over h 2.417 989 454 e14 0.000 000 060 e14 A J^-1 +Faraday constant 96 485.3399 0.0024 C mol^-1 +Faraday constant for conventional electric current 96 485.3401 0.0048 C_90 mol^-1 +Fermi coupling constant 1.166 37 e-5 0.000 01 e-5 GeV^-2 +fine-structure constant 7.297 352 5376 e-3 0.000 000 0050 e-3 +first radiation constant 3.741 771 18 e-16 0.000 000 19 e-16 W m^2 +first radiation constant for spectral radiance 1.191 042 759 e-16 0.000 000 059 e-16 W m^2 sr^-1 +hartree-atomic mass unit relationship 2.921 262 2986 e-8 0.000 000 0042 e-8 u +hartree-electron volt relationship 27.211 383 86 0.000 000 68 eV +Hartree energy 4.359 743 94 e-18 0.000 000 22 e-18 J +Hartree energy in eV 27.211 383 86 0.000 000 68 eV +hartree-hertz relationship 6.579 683 920 722 e15 0.000 000 000 044 e15 Hz +hartree-inverse meter relationship 2.194 746 313 705 e7 0.000 000 000 015 e7 m^-1 +hartree-joule relationship 4.359 743 94 e-18 0.000 000 22 e-18 J +hartree-kelvin relationship 3.157 7465 e5 0.000 0055 e5 K +hartree-kilogram relationship 4.850 869 34 e-35 0.000 000 24 e-35 kg +helion-electron mass ratio 5495.885 2765 0.000 0052 +helion mass 5.006 411 92 e-27 0.000 000 25 e-27 kg +helion mass energy equivalent 4.499 538 64 e-10 0.000 000 22 e-10 J +helion mass energy equivalent in MeV 2808.391 383 0.000 070 MeV +helion mass in u 3.014 932 2473 0.000 000 0026 u +helion molar mass 3.014 932 2473 e-3 0.000 000 0026 e-3 kg mol^-1 +helion-proton mass ratio 2.993 152 6713 0.000 000 0026 +hertz-atomic mass unit relationship 4.439 821 6294 e-24 0.000 000 0064 e-24 u +hertz-electron volt relationship 4.135 667 33 e-15 0.000 000 10 e-15 eV +hertz-hartree relationship 1.519 829 846 006 e-16 0.000 000 000010e-16 E_h +hertz-inverse meter relationship 3.335 640 951... e-9 (exact) m^-1 +hertz-joule relationship 6.626 068 96 e-34 0.000 000 33 e-34 J +hertz-kelvin relationship 4.799 2374 e-11 0.000 0084 e-11 K +hertz-kilogram relationship 7.372 496 00 e-51 0.000 000 37 e-51 kg +inverse fine-structure constant 137.035 999 679 0.000 000 094 +inverse meter-atomic mass unit relationship 1.331 025 0394 e-15 0.000 000 0019 e-15 u +inverse meter-electron volt relationship 1.239 841 875 e-6 0.000 000 031 e-6 eV +inverse meter-hartree relationship 4.556 335 252 760 e-8 0.000 000 000 030 e-8 E_h +inverse meter-hertz relationship 299 792 458 (exact) Hz +inverse meter-joule relationship 1.986 445 501 e-25 0.000 000 099 e-25 J +inverse meter-kelvin relationship 1.438 7752 e-2 0.000 0025 e-2 K +inverse meter-kilogram relationship 2.210 218 70 e-42 0.000 000 11 e-42 kg +inverse of conductance quantum 12 906.403 7787 0.000 0088 ohm +Josephson constant 483 597.891 e9 0.012 e9 Hz V^-1 +joule-atomic mass unit relationship 6.700 536 41 e9 0.000 000 33 e9 u +joule-electron volt relationship 6.241 509 65 e18 0.000 000 16 e18 eV +joule-hartree relationship 2.293 712 69 e17 0.000 000 11 e17 E_h +joule-hertz relationship 1.509 190 450 e33 0.000 000 075 e33 Hz +joule-inverse meter relationship 5.034 117 47 e24 0.000 000 25 e24 m^-1 +joule-kelvin relationship 7.242 963 e22 0.000 013 e22 K +joule-kilogram relationship 1.112 650 056... e-17 (exact) kg +kelvin-atomic mass unit relationship 9.251 098 e-14 0.000 016 e-14 u +kelvin-electron volt relationship 8.617 343 e-5 0.000 015 e-5 eV +kelvin-hartree relationship 3.166 8153 e-6 0.000 0055 e-6 E_h +kelvin-hertz relationship 2.083 6644 e10 0.000 0036 e10 Hz +kelvin-inverse meter relationship 69.503 56 0.000 12 m^-1 +kelvin-joule relationship 1.380 6504 e-23 0.000 0024 e-23 J +kelvin-kilogram relationship 1.536 1807 e-40 0.000 0027 e-40 kg +kilogram-atomic mass unit relationship 6.022 141 79 e26 0.000 000 30 e26 u +kilogram-electron volt relationship 5.609 589 12 e35 0.000 000 14 e35 eV +kilogram-hartree relationship 2.061 486 16 e34 0.000 000 10 e34 E_h +kilogram-hertz relationship 1.356 392 733 e50 0.000 000 068 e50 Hz +kilogram-inverse meter relationship 4.524 439 15 e41 0.000 000 23 e41 m^-1 +kilogram-joule relationship 8.987 551 787... e16 (exact) J +kilogram-kelvin relationship 6.509 651 e39 0.000 011 e39 K +lattice parameter of silicon 543.102 064 e-12 0.000 014 e-12 m +Loschmidt constant (273.15 K, 101.325 kPa) 2.686 7774 e25 0.000 0047 e25 m^-3 +mag. constant 12.566 370 614... e-7 (exact) N A^-2 +mag. flux quantum 2.067 833 667 e-15 0.000 000 052 e-15 Wb +molar gas constant 8.314 472 0.000 015 J mol^-1 K^-1 +molar mass constant 1 e-3 (exact) kg mol^-1 +molar mass of carbon-12 12 e-3 (exact) kg mol^-1 +molar Planck constant 3.990 312 6821 e-10 0.000 000 0057 e-10 J s mol^-1 +molar Planck constant times c 0.119 626 564 72 0.000 000 000 17 J m mol^-1 +molar volume of ideal gas (273.15 K, 100 kPa) 22.710 981 e-3 0.000 040 e-3 m^3 mol^-1 +molar volume of ideal gas (273.15 K, 101.325 kPa) 22.413 996 e-3 0.000 039 e-3 m^3 mol^-1 +molar volume of silicon 12.058 8349 e-6 0.000 0011 e-6 m^3 mol^-1 +Mo x unit 1.002 099 55 e-13 0.000 000 53 e-13 m +muon Compton wavelength 11.734 441 04 e-15 0.000 000 30 e-15 m +muon Compton wavelength over 2 pi 1.867 594 295 e-15 0.000 000 047 e-15 m +muon-electron mass ratio 206.768 2823 0.000 0052 +muon g factor -2.002 331 8414 0.000 000 0012 +muon mag. mom. -4.490 447 86 e-26 0.000 000 16 e-26 J T^-1 +muon mag. mom. anomaly 1.165 920 69 e-3 0.000 000 60 e-3 +muon mag. mom. to Bohr magneton ratio -4.841 970 49 e-3 0.000 000 12 e-3 +muon mag. mom. to nuclear magneton ratio -8.890 597 05 0.000 000 23 +muon mass 1.883 531 30 e-28 0.000 000 11 e-28 kg +muon mass energy equivalent 1.692 833 510 e-11 0.000 000 095 e-11 J +muon mass energy equivalent in MeV 105.658 3668 0.000 0038 MeV +muon mass in u 0.113 428 9256 0.000 000 0029 u +muon molar mass 0.113 428 9256 e-3 0.000 000 0029 e-3 kg mol^-1 +muon-neutron mass ratio 0.112 454 5167 0.000 000 0029 +muon-proton mag. mom. ratio -3.183 345 137 0.000 000 085 +muon-proton mass ratio 0.112 609 5261 0.000 000 0029 +muon-tau mass ratio 5.945 92 e-2 0.000 97 e-2 +natural unit of action 1.054 571 628 e-34 0.000 000 053 e-34 J s +natural unit of action in eV s 6.582 118 99 e-16 0.000 000 16 e-16 eV s +natural unit of energy 8.187 104 38 e-14 0.000 000 41 e-14 J +natural unit of energy in MeV 0.510 998 910 0.000 000 013 MeV +natural unit of length 386.159 264 59 e-15 0.000 000 53 e-15 m +natural unit of mass 9.109 382 15 e-31 0.000 000 45 e-31 kg +natural unit of momentum 2.730 924 06 e-22 0.000 000 14 e-22 kg m s^-1 +natural unit of momentum in MeV/c 0.510 998 910 0.000 000 013 MeV/c +natural unit of time 1.288 088 6570 e-21 0.000 000 0018 e-21 s +natural unit of velocity 299 792 458 (exact) m s^-1 +neutron Compton wavelength 1.319 590 8951 e-15 0.000 000 0020 e-15 m +neutron Compton wavelength over 2 pi 0.210 019 413 82 e-15 0.000 000 000 31 e-15 m +neutron-electron mag. mom. ratio 1.040 668 82 e-3 0.000 000 25 e-3 +neutron-electron mass ratio 1838.683 6605 0.000 0011 +neutron g factor -3.826 085 45 0.000 000 90 +neutron gyromag. ratio 1.832 471 85 e8 0.000 000 43 e8 s^-1 T^-1 +neutron gyromag. ratio over 2 pi 29.164 6954 0.000 0069 MHz T^-1 +neutron mag. mom. -0.966 236 41 e-26 0.000 000 23 e-26 J T^-1 +neutron mag. mom. to Bohr magneton ratio -1.041 875 63 e-3 0.000 000 25 e-3 +neutron mag. mom. to nuclear magneton ratio -1.913 042 73 0.000 000 45 +neutron mass 1.674 927 211 e-27 0.000 000 084 e-27 kg +neutron mass energy equivalent 1.505 349 505 e-10 0.000 000 075 e-10 J +neutron mass energy equivalent in MeV 939.565 346 0.000 023 MeV +neutron mass in u 1.008 664 915 97 0.000 000 000 43 u +neutron molar mass 1.008 664 915 97 e-3 0.000 000 000 43 e-3 kg mol^-1 +neutron-muon mass ratio 8.892 484 09 0.000 000 23 +neutron-proton mag. mom. ratio -0.684 979 34 0.000 000 16 +neutron-proton mass ratio 1.001 378 419 18 0.000 000 000 46 +neutron-tau mass ratio 0.528 740 0.000 086 +neutron to shielded proton mag. mom. ratio -0.684 996 94 0.000 000 16 +Newtonian constant of gravitation 6.674 28 e-11 0.000 67 e-11 m^3 kg^-1 s^-2 +Newtonian constant of gravitation over h-bar c 6.708 81 e-39 0.000 67 e-39 (GeV/c^2)^-2 +nuclear magneton 5.050 783 24 e-27 0.000 000 13 e-27 J T^-1 +nuclear magneton in eV/T 3.152 451 2326 e-8 0.000 000 0045 e-8 eV T^-1 +nuclear magneton in inverse meters per tesla 2.542 623 616 e-2 0.000 000 064 e-2 m^-1 T^-1 +nuclear magneton in K/T 3.658 2637 e-4 0.000 0064 e-4 K T^-1 +nuclear magneton in MHz/T 7.622 593 84 0.000 000 19 MHz T^-1 +Planck constant 6.626 068 96 e-34 0.000 000 33 e-34 J s +Planck constant in eV s 4.135 667 33 e-15 0.000 000 10 e-15 eV s +Planck constant over 2 pi 1.054 571 628 e-34 0.000 000 053 e-34 J s +Planck constant over 2 pi in eV s 6.582 118 99 e-16 0.000 000 16 e-16 eV s +Planck constant over 2 pi times c in MeV fm 197.326 9631 0.000 0049 MeV fm +Planck length 1.616 252 e-35 0.000 081 e-35 m +Planck mass 2.176 44 e-8 0.000 11 e-8 kg +Planck mass energy equivalent in GeV 1.220 892 e19 0.000 061 e19 GeV +Planck temperature 1.416 785 e32 0.000 071 e32 K +Planck time 5.391 24 e-44 0.000 27 e-44 s +proton charge to mass quotient 9.578 833 92 e7 0.000 000 24 e7 C kg^-1 +proton Compton wavelength 1.321 409 8446 e-15 0.000 000 0019 e-15 m +proton Compton wavelength over 2 pi 0.210 308 908 61 e-15 0.000 000 000 30 e-15 m +proton-electron mass ratio 1836.152 672 47 0.000 000 80 +proton g factor 5.585 694 713 0.000 000 046 +proton gyromag. ratio 2.675 222 099 e8 0.000 000 070 e8 s^-1 T^-1 +proton gyromag. ratio over 2 pi 42.577 4821 0.000 0011 MHz T^-1 +proton mag. mom. 1.410 606 662 e-26 0.000 000 037 e-26 J T^-1 +proton mag. mom. to Bohr magneton ratio 1.521 032 209 e-3 0.000 000 012 e-3 +proton mag. mom. to nuclear magneton ratio 2.792 847 356 0.000 000 023 +proton mag. shielding correction 25.694 e-6 0.014 e-6 +proton mass 1.672 621 637 e-27 0.000 000 083 e-27 kg +proton mass energy equivalent 1.503 277 359 e-10 0.000 000 075 e-10 J +proton mass energy equivalent in MeV 938.272 013 0.000 023 MeV +proton mass in u 1.007 276 466 77 0.000 000 000 10 u +proton molar mass 1.007 276 466 77 e-3 0.000 000 000 10 e-3 kg mol^-1 +proton-muon mass ratio 8.880 243 39 0.000 000 23 +proton-neutron mag. mom. ratio -1.459 898 06 0.000 000 34 +proton-neutron mass ratio 0.998 623 478 24 0.000 000 000 46 +proton rms charge radius 0.8768 e-15 0.0069 e-15 m +proton-tau mass ratio 0.528 012 0.000 086 +quantum of circulation 3.636 947 5199 e-4 0.000 000 0050 e-4 m^2 s^-1 +quantum of circulation times 2 7.273 895 040 e-4 0.000 000 010 e-4 m^2 s^-1 +Rydberg constant 10 973 731.568 527 0.000 073 m^-1 +Rydberg constant times c in Hz 3.289 841 960 361 e15 0.000 000 000 022 e15 Hz +Rydberg constant times hc in eV 13.605 691 93 0.000 000 34 eV +Rydberg constant times hc in J 2.179 871 97 e-18 0.000 000 11 e-18 J +Sackur-Tetrode constant (1 K, 100 kPa) -1.151 7047 0.000 0044 +Sackur-Tetrode constant (1 K, 101.325 kPa) -1.164 8677 0.000 0044 +second radiation constant 1.438 7752 e-2 0.000 0025 e-2 m K +shielded helion gyromag. ratio 2.037 894 730 e8 0.000 000 056 e8 s^-1 T^-1 +shielded helion gyromag. ratio over 2 pi 32.434 101 98 0.000 000 90 MHz T^-1 +shielded helion mag. mom. -1.074 552 982 e-26 0.000 000 030 e-26 J T^-1 +shielded helion mag. mom. to Bohr magneton ratio -1.158 671 471 e-3 0.000 000 014 e-3 +shielded helion mag. mom. to nuclear magneton ratio -2.127 497 718 0.000 000 025 +shielded helion to proton mag. mom. ratio -0.761 766 558 0.000 000 011 +shielded helion to shielded proton mag. mom. ratio -0.761 786 1313 0.000 000 0033 +shielded proton gyromag. ratio 2.675 153 362 e8 0.000 000 073 e8 s^-1 T^-1 +shielded proton gyromag. ratio over 2 pi 42.576 3881 0.000 0012 MHz T^-1 +shielded proton mag. mom. 1.410 570 419 e-26 0.000 000 038 e-26 J T^-1 +shielded proton mag. mom. to Bohr magneton ratio 1.520 993 128 e-3 0.000 000 017 e-3 +shielded proton mag. mom. to nuclear magneton ratio 2.792 775 598 0.000 000 030 +speed of light in vacuum 299 792 458 (exact) m s^-1 +standard acceleration of gravity 9.806 65 (exact) m s^-2 +standard atmosphere 101 325 (exact) Pa +Stefan-Boltzmann constant 5.670 400 e-8 0.000 040 e-8 W m^-2 K^-4 +tau Compton wavelength 0.697 72 e-15 0.000 11 e-15 m +tau Compton wavelength over 2 pi 0.111 046 e-15 0.000 018 e-15 m +tau-electron mass ratio 3477.48 0.57 +tau mass 3.167 77 e-27 0.000 52 e-27 kg +tau mass energy equivalent 2.847 05 e-10 0.000 46 e-10 J +tau mass energy equivalent in MeV 1776.99 0.29 MeV +tau mass in u 1.907 68 0.000 31 u +tau molar mass 1.907 68 e-3 0.000 31 e-3 kg mol^-1 +tau-muon mass ratio 16.8183 0.0027 +tau-neutron mass ratio 1.891 29 0.000 31 +tau-proton mass ratio 1.893 90 0.000 31 +Thomson cross section 0.665 245 8558 e-28 0.000 000 0027 e-28 m^2 +triton-electron mag. mom. ratio -1.620 514 423 e-3 0.000 000 021 e-3 +triton-electron mass ratio 5496.921 5269 0.000 0051 +triton g factor 5.957 924 896 0.000 000 076 +triton mag. mom. 1.504 609 361 e-26 0.000 000 042 e-26 J T^-1 +triton mag. mom. to Bohr magneton ratio 1.622 393 657 e-3 0.000 000 021 e-3 +triton mag. mom. to nuclear magneton ratio 2.978 962 448 0.000 000 038 +triton mass 5.007 355 88 e-27 0.000 000 25 e-27 kg +triton mass energy equivalent 4.500 387 03 e-10 0.000 000 22 e-10 J +triton mass energy equivalent in MeV 2808.920 906 0.000 070 MeV +triton mass in u 3.015 500 7134 0.000 000 0025 u +triton molar mass 3.015 500 7134 e-3 0.000 000 0025 e-3 kg mol^-1 +triton-neutron mag. mom. ratio -1.557 185 53 0.000 000 37 +triton-proton mag. mom. ratio 1.066 639 908 0.000 000 010 +triton-proton mass ratio 2.993 717 0309 0.000 000 0025 +unified atomic mass unit 1.660 538 782 e-27 0.000 000 083 e-27 kg +von Klitzing constant 25 812.807 557 0.000 018 ohm +weak mixing angle 0.222 55 0.000 56 +Wien frequency displacement law constant 5.878 933 e10 0.000 010 e10 Hz K^-1 +Wien wavelength displacement law constant 2.897 7685 e-3 0.000 0051 e-3 m K""" + + +def exact2006(exact): + mu0 = 4e-7 * math.pi + c = exact['speed of light in vacuum'] + epsilon0 = 1 / (mu0 * c**2) + replace = { + 'mag. constant': mu0, + 'electric constant': epsilon0, + 'atomic unit of permittivity': 4*math.pi*epsilon0, + 'characteristic impedance of vacuum': math.sqrt(mu0 / epsilon0), + 'hertz-inverse meter relationship': 1/c, + 'joule-kilogram relationship': 1/c**2, + 'kilogram-joule relationship': c**2, + } + return replace + + +txt2010 = """\ +{220} lattice spacing of silicon 192.015 5714 e-12 0.000 0032 e-12 m +alpha particle-electron mass ratio 7294.299 5361 0.000 0029 +alpha particle mass 6.644 656 75 e-27 0.000 000 29 e-27 kg +alpha particle mass energy equivalent 5.971 919 67 e-10 0.000 000 26 e-10 J +alpha particle mass energy equivalent in MeV 3727.379 240 0.000 082 MeV +alpha particle mass in u 4.001 506 179 125 0.000 000 000 062 u +alpha particle molar mass 4.001 506 179 125 e-3 0.000 000 000 062 e-3 kg mol^-1 +alpha particle-proton mass ratio 3.972 599 689 33 0.000 000 000 36 +Angstrom star 1.000 014 95 e-10 0.000 000 90 e-10 m +atomic mass constant 1.660 538 921 e-27 0.000 000 073 e-27 kg +atomic mass constant energy equivalent 1.492 417 954 e-10 0.000 000 066 e-10 J +atomic mass constant energy equivalent in MeV 931.494 061 0.000 021 MeV +atomic mass unit-electron volt relationship 931.494 061 e6 0.000 021 e6 eV +atomic mass unit-hartree relationship 3.423 177 6845 e7 0.000 000 0024 e7 E_h +atomic mass unit-hertz relationship 2.252 342 7168 e23 0.000 000 0016 e23 Hz +atomic mass unit-inverse meter relationship 7.513 006 6042 e14 0.000 000 0053 e14 m^-1 +atomic mass unit-joule relationship 1.492 417 954 e-10 0.000 000 066 e-10 J +atomic mass unit-kelvin relationship 1.080 954 08 e13 0.000 000 98 e13 K +atomic mass unit-kilogram relationship 1.660 538 921 e-27 0.000 000 073 e-27 kg +atomic unit of 1st hyperpolarizability 3.206 361 449 e-53 0.000 000 071 e-53 C^3 m^3 J^-2 +atomic unit of 2nd hyperpolarizability 6.235 380 54 e-65 0.000 000 28 e-65 C^4 m^4 J^-3 +atomic unit of action 1.054 571 726 e-34 0.000 000 047 e-34 J s +atomic unit of charge 1.602 176 565 e-19 0.000 000 035 e-19 C +atomic unit of charge density 1.081 202 338 e12 0.000 000 024 e12 C m^-3 +atomic unit of current 6.623 617 95 e-3 0.000 000 15 e-3 A +atomic unit of electric dipole mom. 8.478 353 26 e-30 0.000 000 19 e-30 C m +atomic unit of electric field 5.142 206 52 e11 0.000 000 11 e11 V m^-1 +atomic unit of electric field gradient 9.717 362 00 e21 0.000 000 21 e21 V m^-2 +atomic unit of electric polarizability 1.648 777 2754 e-41 0.000 000 0016 e-41 C^2 m^2 J^-1 +atomic unit of electric potential 27.211 385 05 0.000 000 60 V +atomic unit of electric quadrupole mom. 4.486 551 331 e-40 0.000 000 099 e-40 C m^2 +atomic unit of energy 4.359 744 34 e-18 0.000 000 19 e-18 J +atomic unit of force 8.238 722 78 e-8 0.000 000 36 e-8 N +atomic unit of length 0.529 177 210 92 e-10 0.000 000 000 17 e-10 m +atomic unit of mag. dipole mom. 1.854 801 936 e-23 0.000 000 041 e-23 J T^-1 +atomic unit of mag. flux density 2.350 517 464 e5 0.000 000 052 e5 T +atomic unit of magnetizability 7.891 036 607 e-29 0.000 000 013 e-29 J T^-2 +atomic unit of mass 9.109 382 91 e-31 0.000 000 40 e-31 kg +atomic unit of mom.um 1.992 851 740 e-24 0.000 000 088 e-24 kg m s^-1 +atomic unit of permittivity 1.112 650 056... e-10 (exact) F m^-1 +atomic unit of time 2.418 884 326 502e-17 0.000 000 000 012e-17 s +atomic unit of velocity 2.187 691 263 79 e6 0.000 000 000 71 e6 m s^-1 +Avogadro constant 6.022 141 29 e23 0.000 000 27 e23 mol^-1 +Bohr magneton 927.400 968 e-26 0.000 020 e-26 J T^-1 +Bohr magneton in eV/T 5.788 381 8066 e-5 0.000 000 0038 e-5 eV T^-1 +Bohr magneton in Hz/T 13.996 245 55 e9 0.000 000 31 e9 Hz T^-1 +Bohr magneton in inverse meters per tesla 46.686 4498 0.000 0010 m^-1 T^-1 +Bohr magneton in K/T 0.671 713 88 0.000 000 61 K T^-1 +Bohr radius 0.529 177 210 92 e-10 0.000 000 000 17 e-10 m +Boltzmann constant 1.380 6488 e-23 0.000 0013 e-23 J K^-1 +Boltzmann constant in eV/K 8.617 3324 e-5 0.000 0078 e-5 eV K^-1 +Boltzmann constant in Hz/K 2.083 6618 e10 0.000 0019 e10 Hz K^-1 +Boltzmann constant in inverse meters per kelvin 69.503 476 0.000 063 m^-1 K^-1 +characteristic impedance of vacuum 376.730 313 461... (exact) ohm +classical electron radius 2.817 940 3267 e-15 0.000 000 0027 e-15 m +Compton wavelength 2.426 310 2389 e-12 0.000 000 0016 e-12 m +Compton wavelength over 2 pi 386.159 268 00 e-15 0.000 000 25 e-15 m +conductance quantum 7.748 091 7346 e-5 0.000 000 0025 e-5 S +conventional value of Josephson constant 483 597.9 e9 (exact) Hz V^-1 +conventional value of von Klitzing constant 25 812.807 (exact) ohm +Cu x unit 1.002 076 97 e-13 0.000 000 28 e-13 m +deuteron-electron mag. mom. ratio -4.664 345 537 e-4 0.000 000 039 e-4 +deuteron-electron mass ratio 3670.482 9652 0.000 0015 +deuteron g factor 0.857 438 2308 0.000 000 0072 +deuteron mag. mom. 0.433 073 489 e-26 0.000 000 010 e-26 J T^-1 +deuteron mag. mom. to Bohr magneton ratio 0.466 975 4556 e-3 0.000 000 0039 e-3 +deuteron mag. mom. to nuclear magneton ratio 0.857 438 2308 0.000 000 0072 +deuteron mass 3.343 583 48 e-27 0.000 000 15 e-27 kg +deuteron mass energy equivalent 3.005 062 97 e-10 0.000 000 13 e-10 J +deuteron mass energy equivalent in MeV 1875.612 859 0.000 041 MeV +deuteron mass in u 2.013 553 212 712 0.000 000 000 077 u +deuteron molar mass 2.013 553 212 712 e-3 0.000 000 000 077 e-3 kg mol^-1 +deuteron-neutron mag. mom. ratio -0.448 206 52 0.000 000 11 +deuteron-proton mag. mom. ratio 0.307 012 2070 0.000 000 0024 +deuteron-proton mass ratio 1.999 007 500 97 0.000 000 000 18 +deuteron rms charge radius 2.1424 e-15 0.0021 e-15 m +electric constant 8.854 187 817... e-12 (exact) F m^-1 +electron charge to mass quotient -1.758 820 088 e11 0.000 000 039 e11 C kg^-1 +electron-deuteron mag. mom. ratio -2143.923 498 0.000 018 +electron-deuteron mass ratio 2.724 437 1095 e-4 0.000 000 0011 e-4 +electron g factor -2.002 319 304 361 53 0.000 000 000 000 53 +electron gyromag. ratio 1.760 859 708 e11 0.000 000 039 e11 s^-1 T^-1 +electron gyromag. ratio over 2 pi 28 024.952 66 0.000 62 MHz T^-1 +electron-helion mass ratio 1.819 543 0761 e-4 0.000 000 0017 e-4 +electron mag. mom. -928.476 430 e-26 0.000 021 e-26 J T^-1 +electron mag. mom. anomaly 1.159 652 180 76 e-3 0.000 000 000 27 e-3 +electron mag. mom. to Bohr magneton ratio -1.001 159 652 180 76 0.000 000 000 000 27 +electron mag. mom. to nuclear magneton ratio -1838.281 970 90 0.000 000 75 +electron mass 9.109 382 91 e-31 0.000 000 40 e-31 kg +electron mass energy equivalent 8.187 105 06 e-14 0.000 000 36 e-14 J +electron mass energy equivalent in MeV 0.510 998 928 0.000 000 011 MeV +electron mass in u 5.485 799 0946 e-4 0.000 000 0022 e-4 u +electron molar mass 5.485 799 0946 e-7 0.000 000 0022 e-7 kg mol^-1 +electron-muon mag. mom. ratio 206.766 9896 0.000 0052 +electron-muon mass ratio 4.836 331 66 e-3 0.000 000 12 e-3 +electron-neutron mag. mom. ratio 960.920 50 0.000 23 +electron-neutron mass ratio 5.438 673 4461 e-4 0.000 000 0032 e-4 +electron-proton mag. mom. ratio -658.210 6848 0.000 0054 +electron-proton mass ratio 5.446 170 2178 e-4 0.000 000 0022 e-4 +electron-tau mass ratio 2.875 92 e-4 0.000 26 e-4 +electron to alpha particle mass ratio 1.370 933 555 78 e-4 0.000 000 000 55 e-4 +electron to shielded helion mag. mom. ratio 864.058 257 0.000 010 +electron to shielded proton mag. mom. ratio -658.227 5971 0.000 0072 +electron-triton mass ratio 1.819 200 0653 e-4 0.000 000 0017 e-4 +electron volt 1.602 176 565 e-19 0.000 000 035 e-19 J +electron volt-atomic mass unit relationship 1.073 544 150 e-9 0.000 000 024 e-9 u +electron volt-hartree relationship 3.674 932 379 e-2 0.000 000 081 e-2 E_h +electron volt-hertz relationship 2.417 989 348 e14 0.000 000 053 e14 Hz +electron volt-inverse meter relationship 8.065 544 29 e5 0.000 000 18 e5 m^-1 +electron volt-joule relationship 1.602 176 565 e-19 0.000 000 035 e-19 J +electron volt-kelvin relationship 1.160 4519 e4 0.000 0011 e4 K +electron volt-kilogram relationship 1.782 661 845 e-36 0.000 000 039 e-36 kg +elementary charge 1.602 176 565 e-19 0.000 000 035 e-19 C +elementary charge over h 2.417 989 348 e14 0.000 000 053 e14 A J^-1 +Faraday constant 96 485.3365 0.0021 C mol^-1 +Faraday constant for conventional electric current 96 485.3321 0.0043 C_90 mol^-1 +Fermi coupling constant 1.166 364 e-5 0.000 005 e-5 GeV^-2 +fine-structure constant 7.297 352 5698 e-3 0.000 000 0024 e-3 +first radiation constant 3.741 771 53 e-16 0.000 000 17 e-16 W m^2 +first radiation constant for spectral radiance 1.191 042 869 e-16 0.000 000 053 e-16 W m^2 sr^-1 +hartree-atomic mass unit relationship 2.921 262 3246 e-8 0.000 000 0021 e-8 u +hartree-electron volt relationship 27.211 385 05 0.000 000 60 eV +Hartree energy 4.359 744 34 e-18 0.000 000 19 e-18 J +Hartree energy in eV 27.211 385 05 0.000 000 60 eV +hartree-hertz relationship 6.579 683 920 729 e15 0.000 000 000 033 e15 Hz +hartree-inverse meter relationship 2.194 746 313 708 e7 0.000 000 000 011 e7 m^-1 +hartree-joule relationship 4.359 744 34 e-18 0.000 000 19 e-18 J +hartree-kelvin relationship 3.157 7504 e5 0.000 0029 e5 K +hartree-kilogram relationship 4.850 869 79 e-35 0.000 000 21 e-35 kg +helion-electron mass ratio 5495.885 2754 0.000 0050 +helion g factor -4.255 250 613 0.000 000 050 +helion mag. mom. -1.074 617 486 e-26 0.000 000 027 e-26 J T^-1 +helion mag. mom. to Bohr magneton ratio -1.158 740 958 e-3 0.000 000 014 e-3 +helion mag. mom. to nuclear magneton ratio -2.127 625 306 0.000 000 025 +helion mass 5.006 412 34 e-27 0.000 000 22 e-27 kg +helion mass energy equivalent 4.499 539 02 e-10 0.000 000 20 e-10 J +helion mass energy equivalent in MeV 2808.391 482 0.000 062 MeV +helion mass in u 3.014 932 2468 0.000 000 0025 u +helion molar mass 3.014 932 2468 e-3 0.000 000 0025 e-3 kg mol^-1 +helion-proton mass ratio 2.993 152 6707 0.000 000 0025 +hertz-atomic mass unit relationship 4.439 821 6689 e-24 0.000 000 0031 e-24 u +hertz-electron volt relationship 4.135 667 516 e-15 0.000 000 091 e-15 eV +hertz-hartree relationship 1.519 829 8460045e-16 0.000 000 0000076e-16 E_h +hertz-inverse meter relationship 3.335 640 951... e-9 (exact) m^-1 +hertz-joule relationship 6.626 069 57 e-34 0.000 000 29 e-34 J +hertz-kelvin relationship 4.799 2434 e-11 0.000 0044 e-11 K +hertz-kilogram relationship 7.372 496 68 e-51 0.000 000 33 e-51 kg +inverse fine-structure constant 137.035 999 074 0.000 000 044 +inverse meter-atomic mass unit relationship 1.331 025 051 20 e-15 0.000 000 000 94 e-15 u +inverse meter-electron volt relationship 1.239 841 930 e-6 0.000 000 027 e-6 eV +inverse meter-hartree relationship 4.556 335 252 755 e-8 0.000 000 000 023 e-8 E_h +inverse meter-hertz relationship 299 792 458 (exact) Hz +inverse meter-joule relationship 1.986 445 684 e-25 0.000 000 088 e-25 J +inverse meter-kelvin relationship 1.438 7770 e-2 0.000 0013 e-2 K +inverse meter-kilogram relationship 2.210 218 902 e-42 0.000 000 098 e-42 kg +inverse of conductance quantum 12 906.403 7217 0.000 0042 ohm +Josephson constant 483 597.870 e9 0.011 e9 Hz V^-1 +joule-atomic mass unit relationship 6.700 535 85 e9 0.000 000 30 e9 u +joule-electron volt relationship 6.241 509 34 e18 0.000 000 14 e18 eV +joule-hartree relationship 2.293 712 48 e17 0.000 000 10 e17 E_h +joule-hertz relationship 1.509 190 311 e33 0.000 000 067 e33 Hz +joule-inverse meter relationship 5.034 117 01 e24 0.000 000 22 e24 m^-1 +joule-kelvin relationship 7.242 9716 e22 0.000 0066 e22 K +joule-kilogram relationship 1.112 650 056... e-17 (exact) kg +kelvin-atomic mass unit relationship 9.251 0868 e-14 0.000 0084 e-14 u +kelvin-electron volt relationship 8.617 3324 e-5 0.000 0078 e-5 eV +kelvin-hartree relationship 3.166 8114 e-6 0.000 0029 e-6 E_h +kelvin-hertz relationship 2.083 6618 e10 0.000 0019 e10 Hz +kelvin-inverse meter relationship 69.503 476 0.000 063 m^-1 +kelvin-joule relationship 1.380 6488 e-23 0.000 0013 e-23 J +kelvin-kilogram relationship 1.536 1790 e-40 0.000 0014 e-40 kg +kilogram-atomic mass unit relationship 6.022 141 29 e26 0.000 000 27 e26 u +kilogram-electron volt relationship 5.609 588 85 e35 0.000 000 12 e35 eV +kilogram-hartree relationship 2.061 485 968 e34 0.000 000 091 e34 E_h +kilogram-hertz relationship 1.356 392 608 e50 0.000 000 060 e50 Hz +kilogram-inverse meter relationship 4.524 438 73 e41 0.000 000 20 e41 m^-1 +kilogram-joule relationship 8.987 551 787... e16 (exact) J +kilogram-kelvin relationship 6.509 6582 e39 0.000 0059 e39 K +lattice parameter of silicon 543.102 0504 e-12 0.000 0089 e-12 m +Loschmidt constant (273.15 K, 100 kPa) 2.651 6462 e25 0.000 0024 e25 m^-3 +Loschmidt constant (273.15 K, 101.325 kPa) 2.686 7805 e25 0.000 0024 e25 m^-3 +mag. constant 12.566 370 614... e-7 (exact) N A^-2 +mag. flux quantum 2.067 833 758 e-15 0.000 000 046 e-15 Wb +molar gas constant 8.314 4621 0.000 0075 J mol^-1 K^-1 +molar mass constant 1 e-3 (exact) kg mol^-1 +molar mass of carbon-12 12 e-3 (exact) kg mol^-1 +molar Planck constant 3.990 312 7176 e-10 0.000 000 0028 e-10 J s mol^-1 +molar Planck constant times c 0.119 626 565 779 0.000 000 000 084 J m mol^-1 +molar volume of ideal gas (273.15 K, 100 kPa) 22.710 953 e-3 0.000 021 e-3 m^3 mol^-1 +molar volume of ideal gas (273.15 K, 101.325 kPa) 22.413 968 e-3 0.000 020 e-3 m^3 mol^-1 +molar volume of silicon 12.058 833 01 e-6 0.000 000 80 e-6 m^3 mol^-1 +Mo x unit 1.002 099 52 e-13 0.000 000 53 e-13 m +muon Compton wavelength 11.734 441 03 e-15 0.000 000 30 e-15 m +muon Compton wavelength over 2 pi 1.867 594 294 e-15 0.000 000 047 e-15 m +muon-electron mass ratio 206.768 2843 0.000 0052 +muon g factor -2.002 331 8418 0.000 000 0013 +muon mag. mom. -4.490 448 07 e-26 0.000 000 15 e-26 J T^-1 +muon mag. mom. anomaly 1.165 920 91 e-3 0.000 000 63 e-3 +muon mag. mom. to Bohr magneton ratio -4.841 970 44 e-3 0.000 000 12 e-3 +muon mag. mom. to nuclear magneton ratio -8.890 596 97 0.000 000 22 +muon mass 1.883 531 475 e-28 0.000 000 096 e-28 kg +muon mass energy equivalent 1.692 833 667 e-11 0.000 000 086 e-11 J +muon mass energy equivalent in MeV 105.658 3715 0.000 0035 MeV +muon mass in u 0.113 428 9267 0.000 000 0029 u +muon molar mass 0.113 428 9267 e-3 0.000 000 0029 e-3 kg mol^-1 +muon-neutron mass ratio 0.112 454 5177 0.000 000 0028 +muon-proton mag. mom. ratio -3.183 345 107 0.000 000 084 +muon-proton mass ratio 0.112 609 5272 0.000 000 0028 +muon-tau mass ratio 5.946 49 e-2 0.000 54 e-2 +natural unit of action 1.054 571 726 e-34 0.000 000 047 e-34 J s +natural unit of action in eV s 6.582 119 28 e-16 0.000 000 15 e-16 eV s +natural unit of energy 8.187 105 06 e-14 0.000 000 36 e-14 J +natural unit of energy in MeV 0.510 998 928 0.000 000 011 MeV +natural unit of length 386.159 268 00 e-15 0.000 000 25 e-15 m +natural unit of mass 9.109 382 91 e-31 0.000 000 40 e-31 kg +natural unit of mom.um 2.730 924 29 e-22 0.000 000 12 e-22 kg m s^-1 +natural unit of mom.um in MeV/c 0.510 998 928 0.000 000 011 MeV/c +natural unit of time 1.288 088 668 33 e-21 0.000 000 000 83 e-21 s +natural unit of velocity 299 792 458 (exact) m s^-1 +neutron Compton wavelength 1.319 590 9068 e-15 0.000 000 0011 e-15 m +neutron Compton wavelength over 2 pi 0.210 019 415 68 e-15 0.000 000 000 17 e-15 m +neutron-electron mag. mom. ratio 1.040 668 82 e-3 0.000 000 25 e-3 +neutron-electron mass ratio 1838.683 6605 0.000 0011 +neutron g factor -3.826 085 45 0.000 000 90 +neutron gyromag. ratio 1.832 471 79 e8 0.000 000 43 e8 s^-1 T^-1 +neutron gyromag. ratio over 2 pi 29.164 6943 0.000 0069 MHz T^-1 +neutron mag. mom. -0.966 236 47 e-26 0.000 000 23 e-26 J T^-1 +neutron mag. mom. to Bohr magneton ratio -1.041 875 63 e-3 0.000 000 25 e-3 +neutron mag. mom. to nuclear magneton ratio -1.913 042 72 0.000 000 45 +neutron mass 1.674 927 351 e-27 0.000 000 074 e-27 kg +neutron mass energy equivalent 1.505 349 631 e-10 0.000 000 066 e-10 J +neutron mass energy equivalent in MeV 939.565 379 0.000 021 MeV +neutron mass in u 1.008 664 916 00 0.000 000 000 43 u +neutron molar mass 1.008 664 916 00 e-3 0.000 000 000 43 e-3 kg mol^-1 +neutron-muon mass ratio 8.892 484 00 0.000 000 22 +neutron-proton mag. mom. ratio -0.684 979 34 0.000 000 16 +neutron-proton mass difference 2.305 573 92 e-30 0.000 000 76 e-30 +neutron-proton mass difference energy equivalent 2.072 146 50 e-13 0.000 000 68 e-13 +neutron-proton mass difference energy equivalent in MeV 1.293 332 17 0.000 000 42 +neutron-proton mass difference in u 0.001 388 449 19 0.000 000 000 45 +neutron-proton mass ratio 1.001 378 419 17 0.000 000 000 45 +neutron-tau mass ratio 0.528 790 0.000 048 +neutron to shielded proton mag. mom. ratio -0.684 996 94 0.000 000 16 +Newtonian constant of gravitation 6.673 84 e-11 0.000 80 e-11 m^3 kg^-1 s^-2 +Newtonian constant of gravitation over h-bar c 6.708 37 e-39 0.000 80 e-39 (GeV/c^2)^-2 +nuclear magneton 5.050 783 53 e-27 0.000 000 11 e-27 J T^-1 +nuclear magneton in eV/T 3.152 451 2605 e-8 0.000 000 0022 e-8 eV T^-1 +nuclear magneton in inverse meters per tesla 2.542 623 527 e-2 0.000 000 056 e-2 m^-1 T^-1 +nuclear magneton in K/T 3.658 2682 e-4 0.000 0033 e-4 K T^-1 +nuclear magneton in MHz/T 7.622 593 57 0.000 000 17 MHz T^-1 +Planck constant 6.626 069 57 e-34 0.000 000 29 e-34 J s +Planck constant in eV s 4.135 667 516 e-15 0.000 000 091 e-15 eV s +Planck constant over 2 pi 1.054 571 726 e-34 0.000 000 047 e-34 J s +Planck constant over 2 pi in eV s 6.582 119 28 e-16 0.000 000 15 e-16 eV s +Planck constant over 2 pi times c in MeV fm 197.326 9718 0.000 0044 MeV fm +Planck length 1.616 199 e-35 0.000 097 e-35 m +Planck mass 2.176 51 e-8 0.000 13 e-8 kg +Planck mass energy equivalent in GeV 1.220 932 e19 0.000 073 e19 GeV +Planck temperature 1.416 833 e32 0.000 085 e32 K +Planck time 5.391 06 e-44 0.000 32 e-44 s +proton charge to mass quotient 9.578 833 58 e7 0.000 000 21 e7 C kg^-1 +proton Compton wavelength 1.321 409 856 23 e-15 0.000 000 000 94 e-15 m +proton Compton wavelength over 2 pi 0.210 308 910 47 e-15 0.000 000 000 15 e-15 m +proton-electron mass ratio 1836.152 672 45 0.000 000 75 +proton g factor 5.585 694 713 0.000 000 046 +proton gyromag. ratio 2.675 222 005 e8 0.000 000 063 e8 s^-1 T^-1 +proton gyromag. ratio over 2 pi 42.577 4806 0.000 0010 MHz T^-1 +proton mag. mom. 1.410 606 743 e-26 0.000 000 033 e-26 J T^-1 +proton mag. mom. to Bohr magneton ratio 1.521 032 210 e-3 0.000 000 012 e-3 +proton mag. mom. to nuclear magneton ratio 2.792 847 356 0.000 000 023 +proton mag. shielding correction 25.694 e-6 0.014 e-6 +proton mass 1.672 621 777 e-27 0.000 000 074 e-27 kg +proton mass energy equivalent 1.503 277 484 e-10 0.000 000 066 e-10 J +proton mass energy equivalent in MeV 938.272 046 0.000 021 MeV +proton mass in u 1.007 276 466 812 0.000 000 000 090 u +proton molar mass 1.007 276 466 812 e-3 0.000 000 000 090 e-3 kg mol^-1 +proton-muon mass ratio 8.880 243 31 0.000 000 22 +proton-neutron mag. mom. ratio -1.459 898 06 0.000 000 34 +proton-neutron mass ratio 0.998 623 478 26 0.000 000 000 45 +proton rms charge radius 0.8775 e-15 0.0051 e-15 m +proton-tau mass ratio 0.528 063 0.000 048 +quantum of circulation 3.636 947 5520 e-4 0.000 000 0024 e-4 m^2 s^-1 +quantum of circulation times 2 7.273 895 1040 e-4 0.000 000 0047 e-4 m^2 s^-1 +Rydberg constant 10 973 731.568 539 0.000 055 m^-1 +Rydberg constant times c in Hz 3.289 841 960 364 e15 0.000 000 000 017 e15 Hz +Rydberg constant times hc in eV 13.605 692 53 0.000 000 30 eV +Rydberg constant times hc in J 2.179 872 171 e-18 0.000 000 096 e-18 J +Sackur-Tetrode constant (1 K, 100 kPa) -1.151 7078 0.000 0023 +Sackur-Tetrode constant (1 K, 101.325 kPa) -1.164 8708 0.000 0023 +second radiation constant 1.438 7770 e-2 0.000 0013 e-2 m K +shielded helion gyromag. ratio 2.037 894 659 e8 0.000 000 051 e8 s^-1 T^-1 +shielded helion gyromag. ratio over 2 pi 32.434 100 84 0.000 000 81 MHz T^-1 +shielded helion mag. mom. -1.074 553 044 e-26 0.000 000 027 e-26 J T^-1 +shielded helion mag. mom. to Bohr magneton ratio -1.158 671 471 e-3 0.000 000 014 e-3 +shielded helion mag. mom. to nuclear magneton ratio -2.127 497 718 0.000 000 025 +shielded helion to proton mag. mom. ratio -0.761 766 558 0.000 000 011 +shielded helion to shielded proton mag. mom. ratio -0.761 786 1313 0.000 000 0033 +shielded proton gyromag. ratio 2.675 153 268 e8 0.000 000 066 e8 s^-1 T^-1 +shielded proton gyromag. ratio over 2 pi 42.576 3866 0.000 0010 MHz T^-1 +shielded proton mag. mom. 1.410 570 499 e-26 0.000 000 035 e-26 J T^-1 +shielded proton mag. mom. to Bohr magneton ratio 1.520 993 128 e-3 0.000 000 017 e-3 +shielded proton mag. mom. to nuclear magneton ratio 2.792 775 598 0.000 000 030 +speed of light in vacuum 299 792 458 (exact) m s^-1 +standard acceleration of gravity 9.806 65 (exact) m s^-2 +standard atmosphere 101 325 (exact) Pa +standard-state pressure 100 000 (exact) Pa +Stefan-Boltzmann constant 5.670 373 e-8 0.000 021 e-8 W m^-2 K^-4 +tau Compton wavelength 0.697 787 e-15 0.000 063 e-15 m +tau Compton wavelength over 2 pi 0.111 056 e-15 0.000 010 e-15 m +tau-electron mass ratio 3477.15 0.31 +tau mass 3.167 47 e-27 0.000 29 e-27 kg +tau mass energy equivalent 2.846 78 e-10 0.000 26 e-10 J +tau mass energy equivalent in MeV 1776.82 0.16 MeV +tau mass in u 1.907 49 0.000 17 u +tau molar mass 1.907 49 e-3 0.000 17 e-3 kg mol^-1 +tau-muon mass ratio 16.8167 0.0015 +tau-neutron mass ratio 1.891 11 0.000 17 +tau-proton mass ratio 1.893 72 0.000 17 +Thomson cross section 0.665 245 8734 e-28 0.000 000 0013 e-28 m^2 +triton-electron mass ratio 5496.921 5267 0.000 0050 +triton g factor 5.957 924 896 0.000 000 076 +triton mag. mom. 1.504 609 447 e-26 0.000 000 038 e-26 J T^-1 +triton mag. mom. to Bohr magneton ratio 1.622 393 657 e-3 0.000 000 021 e-3 +triton mag. mom. to nuclear magneton ratio 2.978 962 448 0.000 000 038 +triton mass 5.007 356 30 e-27 0.000 000 22 e-27 kg +triton mass energy equivalent 4.500 387 41 e-10 0.000 000 20 e-10 J +triton mass energy equivalent in MeV 2808.921 005 0.000 062 MeV +triton mass in u 3.015 500 7134 0.000 000 0025 u +triton molar mass 3.015 500 7134 e-3 0.000 000 0025 e-3 kg mol^-1 +triton-proton mass ratio 2.993 717 0308 0.000 000 0025 +unified atomic mass unit 1.660 538 921 e-27 0.000 000 073 e-27 kg +von Klitzing constant 25 812.807 4434 0.000 0084 ohm +weak mixing angle 0.2223 0.0021 +Wien frequency displacement law constant 5.878 9254 e10 0.000 0053 e10 Hz K^-1 +Wien wavelength displacement law constant 2.897 7721 e-3 0.000 0026 e-3 m K""" + + +exact2010 = exact2006 + + +txt2014 = """\ +{220} lattice spacing of silicon 192.015 5714 e-12 0.000 0032 e-12 m +alpha particle-electron mass ratio 7294.299 541 36 0.000 000 24 +alpha particle mass 6.644 657 230 e-27 0.000 000 082 e-27 kg +alpha particle mass energy equivalent 5.971 920 097 e-10 0.000 000 073 e-10 J +alpha particle mass energy equivalent in MeV 3727.379 378 0.000 023 MeV +alpha particle mass in u 4.001 506 179 127 0.000 000 000 063 u +alpha particle molar mass 4.001 506 179 127 e-3 0.000 000 000 063 e-3 kg mol^-1 +alpha particle-proton mass ratio 3.972 599 689 07 0.000 000 000 36 +Angstrom star 1.000 014 95 e-10 0.000 000 90 e-10 m +atomic mass constant 1.660 539 040 e-27 0.000 000 020 e-27 kg +atomic mass constant energy equivalent 1.492 418 062 e-10 0.000 000 018 e-10 J +atomic mass constant energy equivalent in MeV 931.494 0954 0.000 0057 MeV +atomic mass unit-electron volt relationship 931.494 0954 e6 0.000 0057 e6 eV +atomic mass unit-hartree relationship 3.423 177 6902 e7 0.000 000 0016 e7 E_h +atomic mass unit-hertz relationship 2.252 342 7206 e23 0.000 000 0010 e23 Hz +atomic mass unit-inverse meter relationship 7.513 006 6166 e14 0.000 000 0034 e14 m^-1 +atomic mass unit-joule relationship 1.492 418 062 e-10 0.000 000 018 e-10 J +atomic mass unit-kelvin relationship 1.080 954 38 e13 0.000 000 62 e13 K +atomic mass unit-kilogram relationship 1.660 539 040 e-27 0.000 000 020 e-27 kg +atomic unit of 1st hyperpolarizability 3.206 361 329 e-53 0.000 000 020 e-53 C^3 m^3 J^-2 +atomic unit of 2nd hyperpolarizability 6.235 380 085 e-65 0.000 000 077 e-65 C^4 m^4 J^-3 +atomic unit of action 1.054 571 800 e-34 0.000 000 013 e-34 J s +atomic unit of charge 1.602 176 6208 e-19 0.000 000 0098 e-19 C +atomic unit of charge density 1.081 202 3770 e12 0.000 000 0067 e12 C m^-3 +atomic unit of current 6.623 618 183 e-3 0.000 000 041 e-3 A +atomic unit of electric dipole mom. 8.478 353 552 e-30 0.000 000 052 e-30 C m +atomic unit of electric field 5.142 206 707 e11 0.000 000 032 e11 V m^-1 +atomic unit of electric field gradient 9.717 362 356 e21 0.000 000 060 e21 V m^-2 +atomic unit of electric polarizability 1.648 777 2731 e-41 0.000 000 0011 e-41 C^2 m^2 J^-1 +atomic unit of electric potential 27.211 386 02 0.000 000 17 V +atomic unit of electric quadrupole mom. 4.486 551 484 e-40 0.000 000 028 e-40 C m^2 +atomic unit of energy 4.359 744 650 e-18 0.000 000 054 e-18 J +atomic unit of force 8.238 723 36 e-8 0.000 000 10 e-8 N +atomic unit of length 0.529 177 210 67 e-10 0.000 000 000 12 e-10 m +atomic unit of mag. dipole mom. 1.854 801 999 e-23 0.000 000 011 e-23 J T^-1 +atomic unit of mag. flux density 2.350 517 550 e5 0.000 000 014 e5 T +atomic unit of magnetizability 7.891 036 5886 e-29 0.000 000 0090 e-29 J T^-2 +atomic unit of mass 9.109 383 56 e-31 0.000 000 11 e-31 kg +atomic unit of mom.um 1.992 851 882 e-24 0.000 000 024 e-24 kg m s^-1 +atomic unit of permittivity 1.112 650 056... e-10 (exact) F m^-1 +atomic unit of time 2.418 884 326509e-17 0.000 000 000014e-17 s +atomic unit of velocity 2.187 691 262 77 e6 0.000 000 000 50 e6 m s^-1 +Avogadro constant 6.022 140 857 e23 0.000 000 074 e23 mol^-1 +Bohr magneton 927.400 9994 e-26 0.000 0057 e-26 J T^-1 +Bohr magneton in eV/T 5.788 381 8012 e-5 0.000 000 0026 e-5 eV T^-1 +Bohr magneton in Hz/T 13.996 245 042 e9 0.000 000 086 e9 Hz T^-1 +Bohr magneton in inverse meters per tesla 46.686 448 14 0.000 000 29 m^-1 T^-1 +Bohr magneton in K/T 0.671 714 05 0.000 000 39 K T^-1 +Bohr radius 0.529 177 210 67 e-10 0.000 000 000 12 e-10 m +Boltzmann constant 1.380 648 52 e-23 0.000 000 79 e-23 J K^-1 +Boltzmann constant in eV/K 8.617 3303 e-5 0.000 0050 e-5 eV K^-1 +Boltzmann constant in Hz/K 2.083 6612 e10 0.000 0012 e10 Hz K^-1 +Boltzmann constant in inverse meters per kelvin 69.503 457 0.000 040 m^-1 K^-1 +characteristic impedance of vacuum 376.730 313 461... (exact) ohm +classical electron radius 2.817 940 3227 e-15 0.000 000 0019 e-15 m +Compton wavelength 2.426 310 2367 e-12 0.000 000 0011 e-12 m +Compton wavelength over 2 pi 386.159 267 64 e-15 0.000 000 18 e-15 m +conductance quantum 7.748 091 7310 e-5 0.000 000 0018 e-5 S +conventional value of Josephson constant 483 597.9 e9 (exact) Hz V^-1 +conventional value of von Klitzing constant 25 812.807 (exact) ohm +Cu x unit 1.002 076 97 e-13 0.000 000 28 e-13 m +deuteron-electron mag. mom. ratio -4.664 345 535 e-4 0.000 000 026 e-4 +deuteron-electron mass ratio 3670.482 967 85 0.000 000 13 +deuteron g factor 0.857 438 2311 0.000 000 0048 +deuteron mag. mom. 0.433 073 5040 e-26 0.000 000 0036 e-26 J T^-1 +deuteron mag. mom. to Bohr magneton ratio 0.466 975 4554 e-3 0.000 000 0026 e-3 +deuteron mag. mom. to nuclear magneton ratio 0.857 438 2311 0.000 000 0048 +deuteron mass 3.343 583 719 e-27 0.000 000 041 e-27 kg +deuteron mass energy equivalent 3.005 063 183 e-10 0.000 000 037 e-10 J +deuteron mass energy equivalent in MeV 1875.612 928 0.000 012 MeV +deuteron mass in u 2.013 553 212 745 0.000 000 000 040 u +deuteron molar mass 2.013 553 212 745 e-3 0.000 000 000 040 e-3 kg mol^-1 +deuteron-neutron mag. mom. ratio -0.448 206 52 0.000 000 11 +deuteron-proton mag. mom. ratio 0.307 012 2077 0.000 000 0015 +deuteron-proton mass ratio 1.999 007 500 87 0.000 000 000 19 +deuteron rms charge radius 2.1413 e-15 0.0025 e-15 m +electric constant 8.854 187 817... e-12 (exact) F m^-1 +electron charge to mass quotient -1.758 820 024 e11 0.000 000 011 e11 C kg^-1 +electron-deuteron mag. mom. ratio -2143.923 499 0.000 012 +electron-deuteron mass ratio 2.724 437 107 484 e-4 0.000 000 000 096 e-4 +electron g factor -2.002 319 304 361 82 0.000 000 000 000 52 +electron gyromag. ratio 1.760 859 644 e11 0.000 000 011 e11 s^-1 T^-1 +electron gyromag. ratio over 2 pi 28 024.951 64 0.000 17 MHz T^-1 +electron-helion mass ratio 1.819 543 074 854 e-4 0.000 000 000 088 e-4 +electron mag. mom. -928.476 4620 e-26 0.000 0057 e-26 J T^-1 +electron mag. mom. anomaly 1.159 652 180 91 e-3 0.000 000 000 26 e-3 +electron mag. mom. to Bohr magneton ratio -1.001 159 652 180 91 0.000 000 000 000 26 +electron mag. mom. to nuclear magneton ratio -1838.281 972 34 0.000 000 17 +electron mass 9.109 383 56 e-31 0.000 000 11 e-31 kg +electron mass energy equivalent 8.187 105 65 e-14 0.000 000 10 e-14 J +electron mass energy equivalent in MeV 0.510 998 9461 0.000 000 0031 MeV +electron mass in u 5.485 799 090 70 e-4 0.000 000 000 16 e-4 u +electron molar mass 5.485 799 090 70 e-7 0.000 000 000 16 e-7 kg mol^-1 +electron-muon mag. mom. ratio 206.766 9880 0.000 0046 +electron-muon mass ratio 4.836 331 70 e-3 0.000 000 11 e-3 +electron-neutron mag. mom. ratio 960.920 50 0.000 23 +electron-neutron mass ratio 5.438 673 4428 e-4 0.000 000 0027 e-4 +electron-proton mag. mom. ratio -658.210 6866 0.000 0020 +electron-proton mass ratio 5.446 170 213 52 e-4 0.000 000 000 52 e-4 +electron-tau mass ratio 2.875 92 e-4 0.000 26 e-4 +electron to alpha particle mass ratio 1.370 933 554 798 e-4 0.000 000 000 045 e-4 +electron to shielded helion mag. mom. ratio 864.058 257 0.000 010 +electron to shielded proton mag. mom. ratio -658.227 5971 0.000 0072 +electron-triton mass ratio 1.819 200 062 203 e-4 0.000 000 000 084 e-4 +electron volt 1.602 176 6208 e-19 0.000 000 0098 e-19 J +electron volt-atomic mass unit relationship 1.073 544 1105 e-9 0.000 000 0066 e-9 u +electron volt-hartree relationship 3.674 932 248 e-2 0.000 000 023 e-2 E_h +electron volt-hertz relationship 2.417 989 262 e14 0.000 000 015 e14 Hz +electron volt-inverse meter relationship 8.065 544 005 e5 0.000 000 050 e5 m^-1 +electron volt-joule relationship 1.602 176 6208 e-19 0.000 000 0098 e-19 J +electron volt-kelvin relationship 1.160 452 21 e4 0.000 000 67 e4 K +electron volt-kilogram relationship 1.782 661 907 e-36 0.000 000 011 e-36 kg +elementary charge 1.602 176 6208 e-19 0.000 000 0098 e-19 C +elementary charge over h 2.417 989 262 e14 0.000 000 015 e14 A J^-1 +Faraday constant 96 485.332 89 0.000 59 C mol^-1 +Faraday constant for conventional electric current 96 485.3251 0.0012 C_90 mol^-1 +Fermi coupling constant 1.166 3787 e-5 0.000 0006 e-5 GeV^-2 +fine-structure constant 7.297 352 5664 e-3 0.000 000 0017 e-3 +first radiation constant 3.741 771 790 e-16 0.000 000 046 e-16 W m^2 +first radiation constant for spectral radiance 1.191 042 953 e-16 0.000 000 015 e-16 W m^2 sr^-1 +hartree-atomic mass unit relationship 2.921 262 3197 e-8 0.000 000 0013 e-8 u +hartree-electron volt relationship 27.211 386 02 0.000 000 17 eV +Hartree energy 4.359 744 650 e-18 0.000 000 054 e-18 J +Hartree energy in eV 27.211 386 02 0.000 000 17 eV +hartree-hertz relationship 6.579 683 920 711 e15 0.000 000 000 039 e15 Hz +hartree-inverse meter relationship 2.194 746 313 702 e7 0.000 000 000 013 e7 m^-1 +hartree-joule relationship 4.359 744 650 e-18 0.000 000 054 e-18 J +hartree-kelvin relationship 3.157 7513 e5 0.000 0018 e5 K +hartree-kilogram relationship 4.850 870 129 e-35 0.000 000 060 e-35 kg +helion-electron mass ratio 5495.885 279 22 0.000 000 27 +helion g factor -4.255 250 616 0.000 000 050 +helion mag. mom. -1.074 617 522 e-26 0.000 000 014 e-26 J T^-1 +helion mag. mom. to Bohr magneton ratio -1.158 740 958 e-3 0.000 000 014 e-3 +helion mag. mom. to nuclear magneton ratio -2.127 625 308 0.000 000 025 +helion mass 5.006 412 700 e-27 0.000 000 062 e-27 kg +helion mass energy equivalent 4.499 539 341 e-10 0.000 000 055 e-10 J +helion mass energy equivalent in MeV 2808.391 586 0.000 017 MeV +helion mass in u 3.014 932 246 73 0.000 000 000 12 u +helion molar mass 3.014 932 246 73 e-3 0.000 000 000 12 e-3 kg mol^-1 +helion-proton mass ratio 2.993 152 670 46 0.000 000 000 29 +hertz-atomic mass unit relationship 4.439 821 6616 e-24 0.000 000 0020 e-24 u +hertz-electron volt relationship 4.135 667 662 e-15 0.000 000 025 e-15 eV +hertz-hartree relationship 1.5198298460088 e-16 0.0000000000090e-16 E_h +hertz-inverse meter relationship 3.335 640 951... e-9 (exact) m^-1 +hertz-joule relationship 6.626 070 040 e-34 0.000 000 081 e-34 J +hertz-kelvin relationship 4.799 2447 e-11 0.000 0028 e-11 K +hertz-kilogram relationship 7.372 497 201 e-51 0.000 000 091 e-51 kg +inverse fine-structure constant 137.035 999 139 0.000 000 031 +inverse meter-atomic mass unit relationship 1.331 025 049 00 e-15 0.000 000 000 61 e-15 u +inverse meter-electron volt relationship 1.239 841 9739 e-6 0.000 000 0076 e-6 eV +inverse meter-hartree relationship 4.556 335 252 767 e-8 0.000 000 000 027 e-8 E_h +inverse meter-hertz relationship 299 792 458 (exact) Hz +inverse meter-joule relationship 1.986 445 824 e-25 0.000 000 024 e-25 J +inverse meter-kelvin relationship 1.438 777 36 e-2 0.000 000 83 e-2 K +inverse meter-kilogram relationship 2.210 219 057 e-42 0.000 000 027 e-42 kg +inverse of conductance quantum 12 906.403 7278 0.000 0029 ohm +Josephson constant 483 597.8525 e9 0.0030 e9 Hz V^-1 +joule-atomic mass unit relationship 6.700 535 363 e9 0.000 000 082 e9 u +joule-electron volt relationship 6.241 509 126 e18 0.000 000 038 e18 eV +joule-hartree relationship 2.293 712 317 e17 0.000 000 028 e17 E_h +joule-hertz relationship 1.509 190 205 e33 0.000 000 019 e33 Hz +joule-inverse meter relationship 5.034 116 651 e24 0.000 000 062 e24 m^-1 +joule-kelvin relationship 7.242 9731 e22 0.000 0042 e22 K +joule-kilogram relationship 1.112 650 056... e-17 (exact) kg +kelvin-atomic mass unit relationship 9.251 0842 e-14 0.000 0053 e-14 u +kelvin-electron volt relationship 8.617 3303 e-5 0.000 0050 e-5 eV +kelvin-hartree relationship 3.166 8105 e-6 0.000 0018 e-6 E_h +kelvin-hertz relationship 2.083 6612 e10 0.000 0012 e10 Hz +kelvin-inverse meter relationship 69.503 457 0.000 040 m^-1 +kelvin-joule relationship 1.380 648 52 e-23 0.000 000 79 e-23 J +kelvin-kilogram relationship 1.536 178 65 e-40 0.000 000 88 e-40 kg +kilogram-atomic mass unit relationship 6.022 140 857 e26 0.000 000 074 e26 u +kilogram-electron volt relationship 5.609 588 650 e35 0.000 000 034 e35 eV +kilogram-hartree relationship 2.061 485 823 e34 0.000 000 025 e34 E_h +kilogram-hertz relationship 1.356 392 512 e50 0.000 000 017 e50 Hz +kilogram-inverse meter relationship 4.524 438 411 e41 0.000 000 056 e41 m^-1 +kilogram-joule relationship 8.987 551 787... e16 (exact) J +kilogram-kelvin relationship 6.509 6595 e39 0.000 0037 e39 K +lattice parameter of silicon 543.102 0504 e-12 0.000 0089 e-12 m +Loschmidt constant (273.15 K, 100 kPa) 2.651 6467 e25 0.000 0015 e25 m^-3 +Loschmidt constant (273.15 K, 101.325 kPa) 2.686 7811 e25 0.000 0015 e25 m^-3 +mag. constant 12.566 370 614... e-7 (exact) N A^-2 +mag. flux quantum 2.067 833 831 e-15 0.000 000 013 e-15 Wb +molar gas constant 8.314 4598 0.000 0048 J mol^-1 K^-1 +molar mass constant 1 e-3 (exact) kg mol^-1 +molar mass of carbon-12 12 e-3 (exact) kg mol^-1 +molar Planck constant 3.990 312 7110 e-10 0.000 000 0018 e-10 J s mol^-1 +molar Planck constant times c 0.119 626 565 582 0.000 000 000 054 J m mol^-1 +molar volume of ideal gas (273.15 K, 100 kPa) 22.710 947 e-3 0.000 013 e-3 m^3 mol^-1 +molar volume of ideal gas (273.15 K, 101.325 kPa) 22.413 962 e-3 0.000 013 e-3 m^3 mol^-1 +molar volume of silicon 12.058 832 14 e-6 0.000 000 61 e-6 m^3 mol^-1 +Mo x unit 1.002 099 52 e-13 0.000 000 53 e-13 m +muon Compton wavelength 11.734 441 11 e-15 0.000 000 26 e-15 m +muon Compton wavelength over 2 pi 1.867 594 308 e-15 0.000 000 042 e-15 m +muon-electron mass ratio 206.768 2826 0.000 0046 +muon g factor -2.002 331 8418 0.000 000 0013 +muon mag. mom. -4.490 448 26 e-26 0.000 000 10 e-26 J T^-1 +muon mag. mom. anomaly 1.165 920 89 e-3 0.000 000 63 e-3 +muon mag. mom. to Bohr magneton ratio -4.841 970 48 e-3 0.000 000 11 e-3 +muon mag. mom. to nuclear magneton ratio -8.890 597 05 0.000 000 20 +muon mass 1.883 531 594 e-28 0.000 000 048 e-28 kg +muon mass energy equivalent 1.692 833 774 e-11 0.000 000 043 e-11 J +muon mass energy equivalent in MeV 105.658 3745 0.000 0024 MeV +muon mass in u 0.113 428 9257 0.000 000 0025 u +muon molar mass 0.113 428 9257 e-3 0.000 000 0025 e-3 kg mol^-1 +muon-neutron mass ratio 0.112 454 5167 0.000 000 0025 +muon-proton mag. mom. ratio -3.183 345 142 0.000 000 071 +muon-proton mass ratio 0.112 609 5262 0.000 000 0025 +muon-tau mass ratio 5.946 49 e-2 0.000 54 e-2 +natural unit of action 1.054 571 800 e-34 0.000 000 013 e-34 J s +natural unit of action in eV s 6.582 119 514 e-16 0.000 000 040 e-16 eV s +natural unit of energy 8.187 105 65 e-14 0.000 000 10 e-14 J +natural unit of energy in MeV 0.510 998 9461 0.000 000 0031 MeV +natural unit of length 386.159 267 64 e-15 0.000 000 18 e-15 m +natural unit of mass 9.109 383 56 e-31 0.000 000 11 e-31 kg +natural unit of mom.um 2.730 924 488 e-22 0.000 000 034 e-22 kg m s^-1 +natural unit of mom.um in MeV/c 0.510 998 9461 0.000 000 0031 MeV/c +natural unit of time 1.288 088 667 12 e-21 0.000 000 000 58 e-21 s +natural unit of velocity 299 792 458 (exact) m s^-1 +neutron Compton wavelength 1.319 590 904 81 e-15 0.000 000 000 88 e-15 m +neutron Compton wavelength over 2 pi 0.210 019 415 36 e-15 0.000 000 000 14 e-15 m +neutron-electron mag. mom. ratio 1.040 668 82 e-3 0.000 000 25 e-3 +neutron-electron mass ratio 1838.683 661 58 0.000 000 90 +neutron g factor -3.826 085 45 0.000 000 90 +neutron gyromag. ratio 1.832 471 72 e8 0.000 000 43 e8 s^-1 T^-1 +neutron gyromag. ratio over 2 pi 29.164 6933 0.000 0069 MHz T^-1 +neutron mag. mom. -0.966 236 50 e-26 0.000 000 23 e-26 J T^-1 +neutron mag. mom. to Bohr magneton ratio -1.041 875 63 e-3 0.000 000 25 e-3 +neutron mag. mom. to nuclear magneton ratio -1.913 042 73 0.000 000 45 +neutron mass 1.674 927 471 e-27 0.000 000 021 e-27 kg +neutron mass energy equivalent 1.505 349 739 e-10 0.000 000 019 e-10 J +neutron mass energy equivalent in MeV 939.565 4133 0.000 0058 MeV +neutron mass in u 1.008 664 915 88 0.000 000 000 49 u +neutron molar mass 1.008 664 915 88 e-3 0.000 000 000 49 e-3 kg mol^-1 +neutron-muon mass ratio 8.892 484 08 0.000 000 20 +neutron-proton mag. mom. ratio -0.684 979 34 0.000 000 16 +neutron-proton mass difference 2.305 573 77 e-30 0.000 000 85 e-30 +neutron-proton mass difference energy equivalent 2.072 146 37 e-13 0.000 000 76 e-13 +neutron-proton mass difference energy equivalent in MeV 1.293 332 05 0.000 000 48 +neutron-proton mass difference in u 0.001 388 449 00 0.000 000 000 51 +neutron-proton mass ratio 1.001 378 418 98 0.000 000 000 51 +neutron-tau mass ratio 0.528 790 0.000 048 +neutron to shielded proton mag. mom. ratio -0.684 996 94 0.000 000 16 +Newtonian constant of gravitation 6.674 08 e-11 0.000 31 e-11 m^3 kg^-1 s^-2 +Newtonian constant of gravitation over h-bar c 6.708 61 e-39 0.000 31 e-39 (GeV/c^2)^-2 +nuclear magneton 5.050 783 699 e-27 0.000 000 031 e-27 J T^-1 +nuclear magneton in eV/T 3.152 451 2550 e-8 0.000 000 0015 e-8 eV T^-1 +nuclear magneton in inverse meters per tesla 2.542 623 432 e-2 0.000 000 016 e-2 m^-1 T^-1 +nuclear magneton in K/T 3.658 2690 e-4 0.000 0021 e-4 K T^-1 +nuclear magneton in MHz/T 7.622 593 285 0.000 000 047 MHz T^-1 +Planck constant 6.626 070 040 e-34 0.000 000 081 e-34 J s +Planck constant in eV s 4.135 667 662 e-15 0.000 000 025 e-15 eV s +Planck constant over 2 pi 1.054 571 800 e-34 0.000 000 013 e-34 J s +Planck constant over 2 pi in eV s 6.582 119 514 e-16 0.000 000 040 e-16 eV s +Planck constant over 2 pi times c in MeV fm 197.326 9788 0.000 0012 MeV fm +Planck length 1.616 229 e-35 0.000 038 e-35 m +Planck mass 2.176 470 e-8 0.000 051 e-8 kg +Planck mass energy equivalent in GeV 1.220 910 e19 0.000 029 e19 GeV +Planck temperature 1.416 808 e32 0.000 033 e32 K +Planck time 5.391 16 e-44 0.000 13 e-44 s +proton charge to mass quotient 9.578 833 226 e7 0.000 000 059 e7 C kg^-1 +proton Compton wavelength 1.321 409 853 96 e-15 0.000 000 000 61 e-15 m +proton Compton wavelength over 2 pi 0.210 308910109e-15 0.000 000 000097e-15 m +proton-electron mass ratio 1836.152 673 89 0.000 000 17 +proton g factor 5.585 694 702 0.000 000 017 +proton gyromag. ratio 2.675 221 900 e8 0.000 000 018 e8 s^-1 T^-1 +proton gyromag. ratio over 2 pi 42.577 478 92 0.000 000 29 MHz T^-1 +proton mag. mom. 1.410 606 7873 e-26 0.000 000 0097 e-26 J T^-1 +proton mag. mom. to Bohr magneton ratio 1.521 032 2053 e-3 0.000 000 0046 e-3 +proton mag. mom. to nuclear magneton ratio 2.792 847 3508 0.000 000 0085 +proton mag. shielding correction 25.691 e-6 0.011 e-6 +proton mass 1.672 621 898 e-27 0.000 000 021 e-27 kg +proton mass energy equivalent 1.503 277 593 e-10 0.000 000 018 e-10 J +proton mass energy equivalent in MeV 938.272 0813 0.000 0058 MeV +proton mass in u 1.007 276 466 879 0.000 000 000 091 u +proton molar mass 1.007 276 466 879 e-3 0.000 000 000 091 e-3 kg mol^-1 +proton-muon mass ratio 8.880 243 38 0.000 000 20 +proton-neutron mag. mom. ratio -1.459 898 05 0.000 000 34 +proton-neutron mass ratio 0.998 623 478 44 0.000 000 000 51 +proton rms charge radius 0.8751 e-15 0.0061 e-15 m +proton-tau mass ratio 0.528 063 0.000 048 +quantum of circulation 3.636 947 5486 e-4 0.000 000 0017 e-4 m^2 s^-1 +quantum of circulation times 2 7.273 895 0972 e-4 0.000 000 0033 e-4 m^2 s^-1 +Rydberg constant 10 973 731.568 508 0.000 065 m^-1 +Rydberg constant times c in Hz 3.289 841 960 355 e15 0.000 000 000 019 e15 Hz +Rydberg constant times hc in eV 13.605 693 009 0.000 000 084 eV +Rydberg constant times hc in J 2.179 872 325 e-18 0.000 000 027 e-18 J +Sackur-Tetrode constant (1 K, 100 kPa) -1.151 7084 0.000 0014 +Sackur-Tetrode constant (1 K, 101.325 kPa) -1.164 8714 0.000 0014 +second radiation constant 1.438 777 36 e-2 0.000 000 83 e-2 m K +shielded helion gyromag. ratio 2.037 894 585 e8 0.000 000 027 e8 s^-1 T^-1 +shielded helion gyromag. ratio over 2 pi 32.434 099 66 0.000 000 43 MHz T^-1 +shielded helion mag. mom. -1.074 553 080 e-26 0.000 000 014 e-26 J T^-1 +shielded helion mag. mom. to Bohr magneton ratio -1.158 671 471 e-3 0.000 000 014 e-3 +shielded helion mag. mom. to nuclear magneton ratio -2.127 497 720 0.000 000 025 +shielded helion to proton mag. mom. ratio -0.761 766 5603 0.000 000 0092 +shielded helion to shielded proton mag. mom. ratio -0.761 786 1313 0.000 000 0033 +shielded proton gyromag. ratio 2.675 153 171 e8 0.000 000 033 e8 s^-1 T^-1 +shielded proton gyromag. ratio over 2 pi 42.576 385 07 0.000 000 53 MHz T^-1 +shielded proton mag. mom. 1.410 570 547 e-26 0.000 000 018 e-26 J T^-1 +shielded proton mag. mom. to Bohr magneton ratio 1.520 993 128 e-3 0.000 000 017 e-3 +shielded proton mag. mom. to nuclear magneton ratio 2.792 775 600 0.000 000 030 +speed of light in vacuum 299 792 458 (exact) m s^-1 +standard acceleration of gravity 9.806 65 (exact) m s^-2 +standard atmosphere 101 325 (exact) Pa +standard-state pressure 100 000 (exact) Pa +Stefan-Boltzmann constant 5.670 367 e-8 0.000 013 e-8 W m^-2 K^-4 +tau Compton wavelength 0.697 787 e-15 0.000 063 e-15 m +tau Compton wavelength over 2 pi 0.111 056 e-15 0.000 010 e-15 m +tau-electron mass ratio 3477.15 0.31 +tau mass 3.167 47 e-27 0.000 29 e-27 kg +tau mass energy equivalent 2.846 78 e-10 0.000 26 e-10 J +tau mass energy equivalent in MeV 1776.82 0.16 MeV +tau mass in u 1.907 49 0.000 17 u +tau molar mass 1.907 49 e-3 0.000 17 e-3 kg mol^-1 +tau-muon mass ratio 16.8167 0.0015 +tau-neutron mass ratio 1.891 11 0.000 17 +tau-proton mass ratio 1.893 72 0.000 17 +Thomson cross section 0.665 245 871 58 e-28 0.000 000 000 91 e-28 m^2 +triton-electron mass ratio 5496.921 535 88 0.000 000 26 +triton g factor 5.957 924 920 0.000 000 028 +triton mag. mom. 1.504 609 503 e-26 0.000 000 012 e-26 J T^-1 +triton mag. mom. to Bohr magneton ratio 1.622 393 6616 e-3 0.000 000 0076 e-3 +triton mag. mom. to nuclear magneton ratio 2.978 962 460 0.000 000 014 +triton mass 5.007 356 665 e-27 0.000 000 062 e-27 kg +triton mass energy equivalent 4.500 387 735 e-10 0.000 000 055 e-10 J +triton mass energy equivalent in MeV 2808.921 112 0.000 017 MeV +triton mass in u 3.015 500 716 32 0.000 000 000 11 u +triton molar mass 3.015 500 716 32 e-3 0.000 000 000 11 e-3 kg mol^-1 +triton-proton mass ratio 2.993 717 033 48 0.000 000 000 22 +unified atomic mass unit 1.660 539 040 e-27 0.000 000 020 e-27 kg +von Klitzing constant 25 812.807 4555 0.000 0059 ohm +weak mixing angle 0.2223 0.0021 +Wien frequency displacement law constant 5.878 9238 e10 0.000 0034 e10 Hz K^-1 +Wien wavelength displacement law constant 2.897 7729 e-3 0.000 0017 e-3 m K""" + + +exact2014 = exact2010 + + +txt2018 = """\ +alpha particle-electron mass ratio 7294.299 541 42 0.000 000 24 +alpha particle mass 6.644 657 3357 e-27 0.000 000 0020 e-27 kg +alpha particle mass energy equivalent 5.971 920 1914 e-10 0.000 000 0018 e-10 J +alpha particle mass energy equivalent in MeV 3727.379 4066 0.000 0011 MeV +alpha particle mass in u 4.001 506 179 127 0.000 000 000 063 u +alpha particle molar mass 4.001 506 1777 e-3 0.000 000 0012 e-3 kg mol^-1 +alpha particle-proton mass ratio 3.972 599 690 09 0.000 000 000 22 +alpha particle relative atomic mass 4.001 506 179 127 0.000 000 000 063 +Angstrom star 1.000 014 95 e-10 0.000 000 90 e-10 m +atomic mass constant 1.660 539 066 60 e-27 0.000 000 000 50 e-27 kg +atomic mass constant energy equivalent 1.492 418 085 60 e-10 0.000 000 000 45 e-10 J +atomic mass constant energy equivalent in MeV 931.494 102 42 0.000 000 28 MeV +atomic mass unit-electron volt relationship 9.314 941 0242 e8 0.000 000 0028 e8 eV +atomic mass unit-hartree relationship 3.423 177 6874 e7 0.000 000 0010 e7 E_h +atomic mass unit-hertz relationship 2.252 342 718 71 e23 0.000 000 000 68 e23 Hz +atomic mass unit-inverse meter relationship 7.513 006 6104 e14 0.000 000 0023 e14 m^-1 +atomic mass unit-joule relationship 1.492 418 085 60 e-10 0.000 000 000 45 e-10 J +atomic mass unit-kelvin relationship 1.080 954 019 16 e13 0.000 000 000 33 e13 K +atomic mass unit-kilogram relationship 1.660 539 066 60 e-27 0.000 000 000 50 e-27 kg +atomic unit of 1st hyperpolarizability 3.206 361 3061 e-53 0.000 000 0015 e-53 C^3 m^3 J^-2 +atomic unit of 2nd hyperpolarizability 6.235 379 9905 e-65 0.000 000 0038 e-65 C^4 m^4 J^-3 +atomic unit of action 1.054 571 817... e-34 (exact) J s +atomic unit of charge 1.602 176 634 e-19 (exact) C +atomic unit of charge density 1.081 202 384 57 e12 0.000 000 000 49 e12 C m^-3 +atomic unit of current 6.623 618 237 510 e-3 0.000 000 000 013 e-3 A +atomic unit of electric dipole mom. 8.478 353 6255 e-30 0.000 000 0013 e-30 C m +atomic unit of electric field 5.142 206 747 63 e11 0.000 000 000 78 e11 V m^-1 +atomic unit of electric field gradient 9.717 362 4292 e21 0.000 000 0029 e21 V m^-2 +atomic unit of electric polarizability 1.648 777 274 36 e-41 0.000 000 000 50 e-41 C^2 m^2 J^-1 +atomic unit of electric potential 27.211 386 245 988 0.000 000 000 053 V +atomic unit of electric quadrupole mom. 4.486 551 5246 e-40 0.000 000 0014 e-40 C m^2 +atomic unit of energy 4.359 744 722 2071 e-18 0.000 000 000 0085 e-18 J +atomic unit of force 8.238 723 4983 e-8 0.000 000 0012 e-8 N +atomic unit of length 5.291 772 109 03 e-11 0.000 000 000 80 e-11 m +atomic unit of mag. dipole mom. 1.854 802 015 66 e-23 0.000 000 000 56 e-23 J T^-1 +atomic unit of mag. flux density 2.350 517 567 58 e5 0.000 000 000 71 e5 T +atomic unit of magnetizability 7.891 036 6008 e-29 0.000 000 0048 e-29 J T^-2 +atomic unit of mass 9.109 383 7015 e-31 0.000 000 0028 e-31 kg +atomic unit of momentum 1.992 851 914 10 e-24 0.000 000 000 30 e-24 kg m s^-1 +atomic unit of permittivity 1.112 650 055 45 e-10 0.000 000 000 17 e-10 F m^-1 +atomic unit of time 2.418 884 326 5857 e-17 0.000 000 000 0047 e-17 s +atomic unit of velocity 2.187 691 263 64 e6 0.000 000 000 33 e6 m s^-1 +Avogadro constant 6.022 140 76 e23 (exact) mol^-1 +Bohr magneton 9.274 010 0783 e-24 0.000 000 0028 e-24 J T^-1 +Bohr magneton in eV/T 5.788 381 8060 e-5 0.000 000 0017 e-5 eV T^-1 +Bohr magneton in Hz/T 1.399 624 493 61 e10 0.000 000 000 42 e10 Hz T^-1 +Bohr magneton in inverse meter per tesla 46.686 447 783 0.000 000 014 m^-1 T^-1 +Bohr magneton in K/T 0.671 713 815 63 0.000 000 000 20 K T^-1 +Bohr radius 5.291 772 109 03 e-11 0.000 000 000 80 e-11 m +Boltzmann constant 1.380 649 e-23 (exact) J K^-1 +Boltzmann constant in eV/K 8.617 333 262... e-5 (exact) eV K^-1 +Boltzmann constant in Hz/K 2.083 661 912... e10 (exact) Hz K^-1 +Boltzmann constant in inverse meter per kelvin 69.503 480 04... (exact) m^-1 K^-1 +characteristic impedance of vacuum 376.730 313 668 0.000 000 057 ohm +classical electron radius 2.817 940 3262 e-15 0.000 000 0013 e-15 m +Compton wavelength 2.426 310 238 67 e-12 0.000 000 000 73 e-12 m +conductance quantum 7.748 091 729... e-5 (exact) S +conventional value of ampere-90 1.000 000 088 87... (exact) A +conventional value of coulomb-90 1.000 000 088 87... (exact) C +conventional value of farad-90 0.999 999 982 20... (exact) F +conventional value of henry-90 1.000 000 017 79... (exact) H +conventional value of Josephson constant 483 597.9 e9 (exact) Hz V^-1 +conventional value of ohm-90 1.000 000 017 79... (exact) ohm +conventional value of volt-90 1.000 000 106 66... (exact) V +conventional value of von Klitzing constant 25 812.807 (exact) ohm +conventional value of watt-90 1.000 000 195 53... (exact) W +Cu x unit 1.002 076 97 e-13 0.000 000 28 e-13 m +deuteron-electron mag. mom. ratio -4.664 345 551 e-4 0.000 000 012 e-4 +deuteron-electron mass ratio 3670.482 967 88 0.000 000 13 +deuteron g factor 0.857 438 2338 0.000 000 0022 +deuteron mag. mom. 4.330 735 094 e-27 0.000 000 011 e-27 J T^-1 +deuteron mag. mom. to Bohr magneton ratio 4.669 754 570 e-4 0.000 000 012 e-4 +deuteron mag. mom. to nuclear magneton ratio 0.857 438 2338 0.000 000 0022 +deuteron mass 3.343 583 7724 e-27 0.000 000 0010 e-27 kg +deuteron mass energy equivalent 3.005 063 231 02 e-10 0.000 000 000 91 e-10 J +deuteron mass energy equivalent in MeV 1875.612 942 57 0.000 000 57 MeV +deuteron mass in u 2.013 553 212 745 0.000 000 000 040 u +deuteron molar mass 2.013 553 212 05 e-3 0.000 000 000 61 e-3 kg mol^-1 +deuteron-neutron mag. mom. ratio -0.448 206 53 0.000 000 11 +deuteron-proton mag. mom. ratio 0.307 012 209 39 0.000 000 000 79 +deuteron-proton mass ratio 1.999 007 501 39 0.000 000 000 11 +deuteron relative atomic mass 2.013 553 212 745 0.000 000 000 040 +deuteron rms charge radius 2.127 99 e-15 0.000 74 e-15 m +electron charge to mass quotient -1.758 820 010 76 e11 0.000 000 000 53 e11 C kg^-1 +electron-deuteron mag. mom. ratio -2143.923 4915 0.000 0056 +electron-deuteron mass ratio 2.724 437 107 462 e-4 0.000 000 000 096 e-4 +electron g factor -2.002 319 304 362 56 0.000 000 000 000 35 +electron gyromag. ratio 1.760 859 630 23 e11 0.000 000 000 53 e11 s^-1 T^-1 +electron gyromag. ratio in MHz/T 28 024.951 4242 0.000 0085 MHz T^-1 +electron-helion mass ratio 1.819 543 074 573 e-4 0.000 000 000 079 e-4 +electron mag. mom. -9.284 764 7043 e-24 0.000 000 0028 e-24 J T^-1 +electron mag. mom. anomaly 1.159 652 181 28 e-3 0.000 000 000 18 e-3 +electron mag. mom. to Bohr magneton ratio -1.001 159 652 181 28 0.000 000 000 000 18 +electron mag. mom. to nuclear magneton ratio -1838.281 971 88 0.000 000 11 +electron mass 9.109 383 7015 e-31 0.000 000 0028 e-31 kg +electron mass energy equivalent 8.187 105 7769 e-14 0.000 000 0025 e-14 J +electron mass energy equivalent in MeV 0.510 998 950 00 0.000 000 000 15 MeV +electron mass in u 5.485 799 090 65 e-4 0.000 000 000 16 e-4 u +electron molar mass 5.485 799 0888 e-7 0.000 000 0017 e-7 kg mol^-1 +electron-muon mag. mom. ratio 206.766 9883 0.000 0046 +electron-muon mass ratio 4.836 331 69 e-3 0.000 000 11 e-3 +electron-neutron mag. mom. ratio 960.920 50 0.000 23 +electron-neutron mass ratio 5.438 673 4424 e-4 0.000 000 0026 e-4 +electron-proton mag. mom. ratio -658.210 687 89 0.000 000 20 +electron-proton mass ratio 5.446 170 214 87 e-4 0.000 000 000 33 e-4 +electron relative atomic mass 5.485 799 090 65 e-4 0.000 000 000 16 e-4 +electron-tau mass ratio 2.875 85 e-4 0.000 19 e-4 +electron to alpha particle mass ratio 1.370 933 554 787 e-4 0.000 000 000 045 e-4 +electron to shielded helion mag. mom. ratio 864.058 257 0.000 010 +electron to shielded proton mag. mom. ratio -658.227 5971 0.000 0072 +electron-triton mass ratio 1.819 200 062 251 e-4 0.000 000 000 090 e-4 +electron volt 1.602 176 634 e-19 (exact) J +electron volt-atomic mass unit relationship 1.073 544 102 33 e-9 0.000 000 000 32 e-9 u +electron volt-hartree relationship 3.674 932 217 5655 e-2 0.000 000 000 0071 e-2 E_h +electron volt-hertz relationship 2.417 989 242... e14 (exact) Hz +electron volt-inverse meter relationship 8.065 543 937... e5 (exact) m^-1 +electron volt-joule relationship 1.602 176 634 e-19 (exact) J +electron volt-kelvin relationship 1.160 451 812... e4 (exact) K +electron volt-kilogram relationship 1.782 661 921... e-36 (exact) kg +elementary charge 1.602 176 634 e-19 (exact) C +elementary charge over h-bar 1.519 267 447... e15 (exact) A J^-1 +Faraday constant 96 485.332 12... (exact) C mol^-1 +Fermi coupling constant 1.166 3787 e-5 0.000 0006 e-5 GeV^-2 +fine-structure constant 7.297 352 5693 e-3 0.000 000 0011 e-3 +first radiation constant 3.741 771 852... e-16 (exact) W m^2 +first radiation constant for spectral radiance 1.191 042 972... e-16 (exact) W m^2 sr^-1 +hartree-atomic mass unit relationship 2.921 262 322 05 e-8 0.000 000 000 88 e-8 u +hartree-electron volt relationship 27.211 386 245 988 0.000 000 000 053 eV +Hartree energy 4.359 744 722 2071 e-18 0.000 000 000 0085 e-18 J +Hartree energy in eV 27.211 386 245 988 0.000 000 000 053 eV +hartree-hertz relationship 6.579 683 920 502 e15 0.000 000 000 013 e15 Hz +hartree-inverse meter relationship 2.194 746 313 6320 e7 0.000 000 000 0043 e7 m^-1 +hartree-joule relationship 4.359 744 722 2071 e-18 0.000 000 000 0085 e-18 J +hartree-kelvin relationship 3.157 750 248 0407 e5 0.000 000 000 0061 e5 K +hartree-kilogram relationship 4.850 870 209 5432 e-35 0.000 000 000 0094 e-35 kg +helion-electron mass ratio 5495.885 280 07 0.000 000 24 +helion g factor -4.255 250 615 0.000 000 050 +helion mag. mom. -1.074 617 532 e-26 0.000 000 013 e-26 J T^-1 +helion mag. mom. to Bohr magneton ratio -1.158 740 958 e-3 0.000 000 014 e-3 +helion mag. mom. to nuclear magneton ratio -2.127 625 307 0.000 000 025 +helion mass 5.006 412 7796 e-27 0.000 000 0015 e-27 kg +helion mass energy equivalent 4.499 539 4125 e-10 0.000 000 0014 e-10 J +helion mass energy equivalent in MeV 2808.391 607 43 0.000 000 85 MeV +helion mass in u 3.014 932 247 175 0.000 000 000 097 u +helion molar mass 3.014 932 246 13 e-3 0.000 000 000 91 e-3 kg mol^-1 +helion-proton mass ratio 2.993 152 671 67 0.000 000 000 13 +helion relative atomic mass 3.014 932 247 175 0.000 000 000 097 +helion shielding shift 5.996 743 e-5 0.000 010 e-5 +hertz-atomic mass unit relationship 4.439 821 6652 e-24 0.000 000 0013 e-24 u +hertz-electron volt relationship 4.135 667 696... e-15 (exact) eV +hertz-hartree relationship 1.519 829 846 0570 e-16 0.000 000 000 0029 e-16 E_h +hertz-inverse meter relationship 3.335 640 951... e-9 (exact) m^-1 +hertz-joule relationship 6.626 070 15 e-34 (exact) J +hertz-kelvin relationship 4.799 243 073... e-11 (exact) K +hertz-kilogram relationship 7.372 497 323... e-51 (exact) kg +hyperfine transition frequency of Cs-133 9 192 631 770 (exact) Hz +inverse fine-structure constant 137.035 999 084 0.000 000 021 +inverse meter-atomic mass unit relationship 1.331 025 050 10 e-15 0.000 000 000 40 e-15 u +inverse meter-electron volt relationship 1.239 841 984... e-6 (exact) eV +inverse meter-hartree relationship 4.556 335 252 9120 e-8 0.000 000 000 0088 e-8 E_h +inverse meter-hertz relationship 299 792 458 (exact) Hz +inverse meter-joule relationship 1.986 445 857... e-25 (exact) J +inverse meter-kelvin relationship 1.438 776 877... e-2 (exact) K +inverse meter-kilogram relationship 2.210 219 094... e-42 (exact) kg +inverse of conductance quantum 12 906.403 72... (exact) ohm +Josephson constant 483 597.848 4... e9 (exact) Hz V^-1 +joule-atomic mass unit relationship 6.700 535 2565 e9 0.000 000 0020 e9 u +joule-electron volt relationship 6.241 509 074... e18 (exact) eV +joule-hartree relationship 2.293 712 278 3963 e17 0.000 000 000 0045 e17 E_h +joule-hertz relationship 1.509 190 179... e33 (exact) Hz +joule-inverse meter relationship 5.034 116 567... e24 (exact) m^-1 +joule-kelvin relationship 7.242 970 516... e22 (exact) K +joule-kilogram relationship 1.112 650 056... e-17 (exact) kg +kelvin-atomic mass unit relationship 9.251 087 3014 e-14 0.000 000 0028 e-14 u +kelvin-electron volt relationship 8.617 333 262... e-5 (exact) eV +kelvin-hartree relationship 3.166 811 563 4556 e-6 0.000 000 000 0061 e-6 E_h +kelvin-hertz relationship 2.083 661 912... e10 (exact) Hz +kelvin-inverse meter relationship 69.503 480 04... (exact) m^-1 +kelvin-joule relationship 1.380 649 e-23 (exact) J +kelvin-kilogram relationship 1.536 179 187... e-40 (exact) kg +kilogram-atomic mass unit relationship 6.022 140 7621 e26 0.000 000 0018 e26 u +kilogram-electron volt relationship 5.609 588 603... e35 (exact) eV +kilogram-hartree relationship 2.061 485 788 7409 e34 0.000 000 000 0040 e34 E_h +kilogram-hertz relationship 1.356 392 489... e50 (exact) Hz +kilogram-inverse meter relationship 4.524 438 335... e41 (exact) m^-1 +kilogram-joule relationship 8.987 551 787... e16 (exact) J +kilogram-kelvin relationship 6.509 657 260... e39 (exact) K +lattice parameter of silicon 5.431 020 511 e-10 0.000 000 089 e-10 m +lattice spacing of ideal Si (220) 1.920 155 716 e-10 0.000 000 032 e-10 m +Loschmidt constant (273.15 K, 100 kPa) 2.651 645 804... e25 (exact) m^-3 +Loschmidt constant (273.15 K, 101.325 kPa) 2.686 780 111... e25 (exact) m^-3 +luminous efficacy 683 (exact) lm W^-1 +mag. flux quantum 2.067 833 848... e-15 (exact) Wb +molar gas constant 8.314 462 618... (exact) J mol^-1 K^-1 +molar mass constant 0.999 999 999 65 e-3 0.000 000 000 30 e-3 kg mol^-1 +molar mass of carbon-12 11.999 999 9958 e-3 0.000 000 0036 e-3 kg mol^-1 +molar Planck constant 3.990 312 712... e-10 (exact) J Hz^-1 mol^-1 +molar volume of ideal gas (273.15 K, 100 kPa) 22.710 954 64... e-3 (exact) m^3 mol^-1 +molar volume of ideal gas (273.15 K, 101.325 kPa) 22.413 969 54... e-3 (exact) m^3 mol^-1 +molar volume of silicon 1.205 883 199 e-5 0.000 000 060 e-5 m^3 mol^-1 +Mo x unit 1.002 099 52 e-13 0.000 000 53 e-13 m +muon Compton wavelength 1.173 444 110 e-14 0.000 000 026 e-14 m +muon-electron mass ratio 206.768 2830 0.000 0046 +muon g factor -2.002 331 8418 0.000 000 0013 +muon mag. mom. -4.490 448 30 e-26 0.000 000 10 e-26 J T^-1 +muon mag. mom. anomaly 1.165 920 89 e-3 0.000 000 63 e-3 +muon mag. mom. to Bohr magneton ratio -4.841 970 47 e-3 0.000 000 11 e-3 +muon mag. mom. to nuclear magneton ratio -8.890 597 03 0.000 000 20 +muon mass 1.883 531 627 e-28 0.000 000 042 e-28 kg +muon mass energy equivalent 1.692 833 804 e-11 0.000 000 038 e-11 J +muon mass energy equivalent in MeV 105.658 3755 0.000 0023 MeV +muon mass in u 0.113 428 9259 0.000 000 0025 u +muon molar mass 1.134 289 259 e-4 0.000 000 025 e-4 kg mol^-1 +muon-neutron mass ratio 0.112 454 5170 0.000 000 0025 +muon-proton mag. mom. ratio -3.183 345 142 0.000 000 071 +muon-proton mass ratio 0.112 609 5264 0.000 000 0025 +muon-tau mass ratio 5.946 35 e-2 0.000 40 e-2 +natural unit of action 1.054 571 817... e-34 (exact) J s +natural unit of action in eV s 6.582 119 569... e-16 (exact) eV s +natural unit of energy 8.187 105 7769 e-14 0.000 000 0025 e-14 J +natural unit of energy in MeV 0.510 998 950 00 0.000 000 000 15 MeV +natural unit of length 3.861 592 6796 e-13 0.000 000 0012 e-13 m +natural unit of mass 9.109 383 7015 e-31 0.000 000 0028 e-31 kg +natural unit of momentum 2.730 924 530 75 e-22 0.000 000 000 82 e-22 kg m s^-1 +natural unit of momentum in MeV/c 0.510 998 950 00 0.000 000 000 15 MeV/c +natural unit of time 1.288 088 668 19 e-21 0.000 000 000 39 e-21 s +natural unit of velocity 299 792 458 (exact) m s^-1 +neutron Compton wavelength 1.319 590 905 81 e-15 0.000 000 000 75 e-15 m +neutron-electron mag. mom. ratio 1.040 668 82 e-3 0.000 000 25 e-3 +neutron-electron mass ratio 1838.683 661 73 0.000 000 89 +neutron g factor -3.826 085 45 0.000 000 90 +neutron gyromag. ratio 1.832 471 71 e8 0.000 000 43 e8 s^-1 T^-1 +neutron gyromag. ratio in MHz/T 29.164 6931 0.000 0069 MHz T^-1 +neutron mag. mom. -9.662 3651 e-27 0.000 0023 e-27 J T^-1 +neutron mag. mom. to Bohr magneton ratio -1.041 875 63 e-3 0.000 000 25 e-3 +neutron mag. mom. to nuclear magneton ratio -1.913 042 73 0.000 000 45 +neutron mass 1.674 927 498 04 e-27 0.000 000 000 95 e-27 kg +neutron mass energy equivalent 1.505 349 762 87 e-10 0.000 000 000 86 e-10 J +neutron mass energy equivalent in MeV 939.565 420 52 0.000 000 54 MeV +neutron mass in u 1.008 664 915 95 0.000 000 000 49 u +neutron molar mass 1.008 664 915 60 e-3 0.000 000 000 57 e-3 kg mol^-1 +neutron-muon mass ratio 8.892 484 06 0.000 000 20 +neutron-proton mag. mom. ratio -0.684 979 34 0.000 000 16 +neutron-proton mass difference 2.305 574 35 e-30 0.000 000 82 e-30 kg +neutron-proton mass difference energy equivalent 2.072 146 89 e-13 0.000 000 74 e-13 J +neutron-proton mass difference energy equivalent in MeV 1.293 332 36 0.000 000 46 MeV +neutron-proton mass difference in u 1.388 449 33 e-3 0.000 000 49 e-3 u +neutron-proton mass ratio 1.001 378 419 31 0.000 000 000 49 +neutron relative atomic mass 1.008 664 915 95 0.000 000 000 49 +neutron-tau mass ratio 0.528 779 0.000 036 +neutron to shielded proton mag. mom. ratio -0.684 996 94 0.000 000 16 +Newtonian constant of gravitation 6.674 30 e-11 0.000 15 e-11 m^3 kg^-1 s^-2 +Newtonian constant of gravitation over h-bar c 6.708 83 e-39 0.000 15 e-39 (GeV/c^2)^-2 +nuclear magneton 5.050 783 7461 e-27 0.000 000 0015 e-27 J T^-1 +nuclear magneton in eV/T 3.152 451 258 44 e-8 0.000 000 000 96 e-8 eV T^-1 +nuclear magneton in inverse meter per tesla 2.542 623 413 53 e-2 0.000 000 000 78 e-2 m^-1 T^-1 +nuclear magneton in K/T 3.658 267 7756 e-4 0.000 000 0011 e-4 K T^-1 +nuclear magneton in MHz/T 7.622 593 2291 0.000 000 0023 MHz T^-1 +Planck constant 6.626 070 15 e-34 (exact) J Hz^-1 +Planck constant in eV/Hz 4.135 667 696... e-15 (exact) eV Hz^-1 +Planck length 1.616 255 e-35 0.000 018 e-35 m +Planck mass 2.176 434 e-8 0.000 024 e-8 kg +Planck mass energy equivalent in GeV 1.220 890 e19 0.000 014 e19 GeV +Planck temperature 1.416 784 e32 0.000 016 e32 K +Planck time 5.391 247 e-44 0.000 060 e-44 s +proton charge to mass quotient 9.578 833 1560 e7 0.000 000 0029 e7 C kg^-1 +proton Compton wavelength 1.321 409 855 39 e-15 0.000 000 000 40 e-15 m +proton-electron mass ratio 1836.152 673 43 0.000 000 11 +proton g factor 5.585 694 6893 0.000 000 0016 +proton gyromag. ratio 2.675 221 8744 e8 0.000 000 0011 e8 s^-1 T^-1 +proton gyromag. ratio in MHz/T 42.577 478 518 0.000 000 018 MHz T^-1 +proton mag. mom. 1.410 606 797 36 e-26 0.000 000 000 60 e-26 J T^-1 +proton mag. mom. to Bohr magneton ratio 1.521 032 202 30 e-3 0.000 000 000 46 e-3 +proton mag. mom. to nuclear magneton ratio 2.792 847 344 63 0.000 000 000 82 +proton mag. shielding correction 2.5689 e-5 0.0011 e-5 +proton mass 1.672 621 923 69 e-27 0.000 000 000 51 e-27 kg +proton mass energy equivalent 1.503 277 615 98 e-10 0.000 000 000 46 e-10 J +proton mass energy equivalent in MeV 938.272 088 16 0.000 000 29 MeV +proton mass in u 1.007 276 466 621 0.000 000 000 053 u +proton molar mass 1.007 276 466 27 e-3 0.000 000 000 31 e-3 kg mol^-1 +proton-muon mass ratio 8.880 243 37 0.000 000 20 +proton-neutron mag. mom. ratio -1.459 898 05 0.000 000 34 +proton-neutron mass ratio 0.998 623 478 12 0.000 000 000 49 +proton relative atomic mass 1.007 276 466 621 0.000 000 000 053 +proton rms charge radius 8.414 e-16 0.019 e-16 m +proton-tau mass ratio 0.528 051 0.000 036 +quantum of circulation 3.636 947 5516 e-4 0.000 000 0011 e-4 m^2 s^-1 +quantum of circulation times 2 7.273 895 1032 e-4 0.000 000 0022 e-4 m^2 s^-1 +reduced Compton wavelength 3.861 592 6796 e-13 0.000 000 0012 e-13 m +reduced muon Compton wavelength 1.867 594 306 e-15 0.000 000 042 e-15 m +reduced neutron Compton wavelength 2.100 194 1552 e-16 0.000 000 0012 e-16 m +reduced Planck constant 1.054 571 817... e-34 (exact) J s +reduced Planck constant in eV s 6.582 119 569... e-16 (exact) eV s +reduced Planck constant times c in MeV fm 197.326 980 4... (exact) MeV fm +reduced proton Compton wavelength 2.103 089 103 36 e-16 0.000 000 000 64 e-16 m +reduced tau Compton wavelength 1.110 538 e-16 0.000 075 e-16 m +Rydberg constant 10 973 731.568 160 0.000 021 m^-1 +Rydberg constant times c in Hz 3.289 841 960 2508 e15 0.000 000 000 0064 e15 Hz +Rydberg constant times hc in eV 13.605 693 122 994 0.000 000 000 026 eV +Rydberg constant times hc in J 2.179 872 361 1035 e-18 0.000 000 000 0042 e-18 J +Sackur-Tetrode constant (1 K, 100 kPa) -1.151 707 537 06 0.000 000 000 45 +Sackur-Tetrode constant (1 K, 101.325 kPa) -1.164 870 523 58 0.000 000 000 45 +second radiation constant 1.438 776 877... e-2 (exact) m K +shielded helion gyromag. ratio 2.037 894 569 e8 0.000 000 024 e8 s^-1 T^-1 +shielded helion gyromag. ratio in MHz/T 32.434 099 42 0.000 000 38 MHz T^-1 +shielded helion mag. mom. -1.074 553 090 e-26 0.000 000 013 e-26 J T^-1 +shielded helion mag. mom. to Bohr magneton ratio -1.158 671 471 e-3 0.000 000 014 e-3 +shielded helion mag. mom. to nuclear magneton ratio -2.127 497 719 0.000 000 025 +shielded helion to proton mag. mom. ratio -0.761 766 5618 0.000 000 0089 +shielded helion to shielded proton mag. mom. ratio -0.761 786 1313 0.000 000 0033 +shielded proton gyromag. ratio 2.675 153 151 e8 0.000 000 029 e8 s^-1 T^-1 +shielded proton gyromag. ratio in MHz/T 42.576 384 74 0.000 000 46 MHz T^-1 +shielded proton mag. mom. 1.410 570 560 e-26 0.000 000 015 e-26 J T^-1 +shielded proton mag. mom. to Bohr magneton ratio 1.520 993 128 e-3 0.000 000 017 e-3 +shielded proton mag. mom. to nuclear magneton ratio 2.792 775 599 0.000 000 030 +shielding difference of d and p in HD 2.0200 e-8 0.0020 e-8 +shielding difference of t and p in HT 2.4140 e-8 0.0020 e-8 +speed of light in vacuum 299 792 458 (exact) m s^-1 +standard acceleration of gravity 9.806 65 (exact) m s^-2 +standard atmosphere 101 325 (exact) Pa +standard-state pressure 100 000 (exact) Pa +Stefan-Boltzmann constant 5.670 374 419... e-8 (exact) W m^-2 K^-4 +tau Compton wavelength 6.977 71 e-16 0.000 47 e-16 m +tau-electron mass ratio 3477.23 0.23 +tau energy equivalent 1776.86 0.12 MeV +tau mass 3.167 54 e-27 0.000 21 e-27 kg +tau mass energy equivalent 2.846 84 e-10 0.000 19 e-10 J +tau mass in u 1.907 54 0.000 13 u +tau molar mass 1.907 54 e-3 0.000 13 e-3 kg mol^-1 +tau-muon mass ratio 16.8170 0.0011 +tau-neutron mass ratio 1.891 15 0.000 13 +tau-proton mass ratio 1.893 76 0.000 13 +Thomson cross section 6.652 458 7321 e-29 0.000 000 0060 e-29 m^2 +triton-electron mass ratio 5496.921 535 73 0.000 000 27 +triton g factor 5.957 924 931 0.000 000 012 +triton mag. mom. 1.504 609 5202 e-26 0.000 000 0030 e-26 J T^-1 +triton mag. mom. to Bohr magneton ratio 1.622 393 6651 e-3 0.000 000 0032 e-3 +triton mag. mom. to nuclear magneton ratio 2.978 962 4656 0.000 000 0059 +triton mass 5.007 356 7446 e-27 0.000 000 0015 e-27 kg +triton mass energy equivalent 4.500 387 8060 e-10 0.000 000 0014 e-10 J +triton mass energy equivalent in MeV 2808.921 132 98 0.000 000 85 MeV +triton mass in u 3.015 500 716 21 0.000 000 000 12 u +triton molar mass 3.015 500 715 17 e-3 0.000 000 000 92 e-3 kg mol^-1 +triton-proton mass ratio 2.993 717 034 14 0.000 000 000 15 +triton relative atomic mass 3.015 500 716 21 0.000 000 000 12 +triton to proton mag. mom. ratio 1.066 639 9191 0.000 000 0021 +unified atomic mass unit 1.660 539 066 60 e-27 0.000 000 000 50 e-27 kg +vacuum electric permittivity 8.854 187 8128 e-12 0.000 000 0013 e-12 F m^-1 +vacuum mag. permeability 1.256 637 062 12 e-6 0.000 000 000 19 e-6 N A^-2 +von Klitzing constant 25 812.807 45... (exact) ohm +weak mixing angle 0.222 90 0.000 30 +Wien frequency displacement law constant 5.878 925 757... e10 (exact) Hz K^-1 +Wien wavelength displacement law constant 2.897 771 955... e-3 (exact) m K +W to Z mass ratio 0.881 53 0.000 17 """ + + +def exact2018(exact): + # SI base constants + c = exact['speed of light in vacuum'] + h = exact['Planck constant'] + e = exact['elementary charge'] + k = exact['Boltzmann constant'] + N_A = exact['Avogadro constant'] + + # Other useful constants + R = N_A * k + hbar = h / (2*math.pi) + G_0 = 2 * e**2 / h + + # Wien law numerical constants: https://en.wikipedia.org/wiki/Wien%27s_displacement_law + # (alpha - 3)*exp(alpha) + 3 = 0 + # (x - 5)*exp(x) + 5 = 0 + alpha_W = 2.821439372122078893403 # 3 + lambertw(-3 * exp(-3)) + x_W = 4.965114231744276303699 # 5 + lambertw(-5 * exp(-5)) + + # Conventional electrical unit + # See https://en.wikipedia.org/wiki/Conventional_electrical_unit + K_J90 = exact['conventional value of Josephson constant'] + K_J = 2 * e / h + R_K90 = exact['conventional value of von Klitzing constant'] + R_K = h / e**2 + V_90 = K_J90 / K_J + ohm_90 = R_K / R_K90 + A_90 = V_90 / ohm_90 + + replace = { + 'atomic unit of action': hbar, + 'Boltzmann constant in eV/K': k / e, + 'Boltzmann constant in Hz/K': k / h, + 'Boltzmann constant in inverse meter per kelvin': k / (h * c), + 'conductance quantum': G_0, + 'conventional value of ampere-90': A_90, + 'conventional value of coulomb-90': A_90, + 'conventional value of farad-90': 1 / ohm_90, + 'conventional value of henry-90': ohm_90, + 'conventional value of ohm-90': ohm_90, + 'conventional value of volt-90': V_90, + 'conventional value of watt-90': V_90**2 / ohm_90, + 'electron volt-hertz relationship': e / h, + 'electron volt-inverse meter relationship': e / (h * c), + 'electron volt-kelvin relationship': e / k, + 'electron volt-kilogram relationship': e / c**2, + 'elementary charge over h-bar': e / hbar, + 'Faraday constant': e * N_A, + 'first radiation constant': 2 * math.pi * h * c**2, + 'first radiation constant for spectral radiance': 2 * h * c**2, + 'hertz-electron volt relationship': h / e, + 'hertz-inverse meter relationship': 1 / c, + 'hertz-kelvin relationship': h / k, + 'hertz-kilogram relationship': h / c**2, + 'inverse meter-electron volt relationship': (h * c) / e, + 'inverse meter-joule relationship': h * c, + 'inverse meter-kelvin relationship': h * c / k, + 'inverse meter-kilogram relationship': h / c, + 'inverse of conductance quantum': 1 / G_0, + 'Josephson constant': K_J, + 'joule-electron volt relationship': 1 / e, + 'joule-hertz relationship': 1 / h, + 'joule-inverse meter relationship': 1 / (h * c), + 'joule-kelvin relationship': 1 / k, + 'joule-kilogram relationship': 1 / c**2, + 'kelvin-electron volt relationship': k / e, + 'kelvin-hertz relationship': k / h, + 'kelvin-inverse meter relationship': k / (h * c), + 'kelvin-kilogram relationship': k / c**2, + 'kilogram-electron volt relationship': c**2 / e, + 'kilogram-hertz relationship': c**2 / h, + 'kilogram-inverse meter relationship': c / h, + 'kilogram-joule relationship': c**2, + 'kilogram-kelvin relationship': c**2 / k, + 'Loschmidt constant (273.15 K, 100 kPa)': 100e3 / 273.15 / k, + 'Loschmidt constant (273.15 K, 101.325 kPa)': 101.325e3 / 273.15 / k, + 'mag. flux quantum': h / (2 * e), + 'molar gas constant': R, + 'molar Planck constant': h * N_A, + 'molar volume of ideal gas (273.15 K, 100 kPa)': R * 273.15 / 100e3, + 'molar volume of ideal gas (273.15 K, 101.325 kPa)': R * 273.15 / 101.325e3, + 'natural unit of action': hbar, + 'natural unit of action in eV s': hbar / e, + 'Planck constant in eV/Hz': h / e, + 'reduced Planck constant': hbar, + 'reduced Planck constant in eV s': hbar / e, + 'reduced Planck constant times c in MeV fm': hbar * c / (e * 1e6 * 1e-15), + 'second radiation constant': h * c / k, + 'Stefan-Boltzmann constant': 2 * math.pi**5 * k**4 / (15 * h**3 * c**2), + 'von Klitzing constant': R_K, + 'Wien frequency displacement law constant': alpha_W * k / h, + 'Wien wavelength displacement law constant': h * c / (x_W * k), + } + return replace + + +txt2022 = """\ +alpha particle-electron mass ratio 7294.299 541 71 0.000 000 17 +alpha particle mass 6.644 657 3450 e-27 0.000 000 0021 e-27 kg +alpha particle mass energy equivalent 5.971 920 1997 e-10 0.000 000 0019 e-10 J +alpha particle mass energy equivalent in MeV 3727.379 4118 0.000 0012 MeV +alpha particle mass in u 4.001 506 179 129 0.000 000 000 062 u +alpha particle molar mass 4.001 506 1833 e-3 0.000 000 0012 e-3 kg mol^-1 +alpha particle-proton mass ratio 3.972 599 690 252 0.000 000 000 070 +alpha particle relative atomic mass 4.001 506 179 129 0.000 000 000 062 +alpha particle rms charge radius 1.6785 e-15 0.0021 e-15 m +Angstrom star 1.000 014 95 e-10 0.000 000 90 e-10 m +atomic mass constant 1.660 539 068 92 e-27 0.000 000 000 52 e-27 kg +atomic mass constant energy equivalent 1.492 418 087 68 e-10 0.000 000 000 46 e-10 J +atomic mass constant energy equivalent in MeV 931.494 103 72 0.000 000 29 MeV +atomic mass unit-electron volt relationship 9.314 941 0372 e8 0.000 000 0029 e8 eV +atomic mass unit-hartree relationship 3.423 177 6922 e7 0.000 000 0011 e7 E_h +atomic mass unit-hertz relationship 2.252 342 721 85 e23 0.000 000 000 70 e23 Hz +atomic mass unit-inverse meter relationship 7.513 006 6209 e14 0.000 000 0023 e14 m^-1 +atomic mass unit-joule relationship 1.492 418 087 68 e-10 0.000 000 000 46 e-10 J +atomic mass unit-kelvin relationship 1.080 954 020 67 e13 0.000 000 000 34 e13 K +atomic mass unit-kilogram relationship 1.660 539 068 92 e-27 0.000 000 000 52 e-27 kg +atomic unit of 1st hyperpolarizability 3.206 361 2996 e-53 0.000 000 0015 e-53 C^3 m^3 J^-2 +atomic unit of 2nd hyperpolarizability 6.235 379 9735 e-65 0.000 000 0039 e-65 C^4 m^4 J^-3 +atomic unit of action 1.054 571 817... e-34 (exact) J s +atomic unit of charge 1.602 176 634 e-19 (exact) C +atomic unit of charge density 1.081 202 386 77 e12 0.000 000 000 51 e12 C m^-3 +atomic unit of current 6.623 618 237 5082 e-3 0.000 000 000 0072 e-3 A +atomic unit of electric dipole mom. 8.478 353 6198 e-30 0.000 000 0013 e-30 C m +atomic unit of electric field 5.142 206 751 12 e11 0.000 000 000 80 e11 V m^-1 +atomic unit of electric field gradient 9.717 362 4424 e21 0.000 000 0030 e21 V m^-2 +atomic unit of electric polarizability 1.648 777 272 12 e-41 0.000 000 000 51 e-41 C^2 m^2 J^-1 +atomic unit of electric potential 27.211 386 245 981 0.000 000 000 030 V +atomic unit of electric quadrupole mom. 4.486 551 5185 e-40 0.000 000 0014 e-40 C m^2 +atomic unit of energy 4.359 744 722 2060 e-18 0.000 000 000 0048 e-18 J +atomic unit of force 8.238 723 5038 e-8 0.000 000 0013 e-8 N +atomic unit of length 5.291 772 105 44 e-11 0.000 000 000 82 e-11 m +atomic unit of mag. dipole mom. 1.854 802 013 15 e-23 0.000 000 000 58 e-23 J T^-1 +atomic unit of mag. flux density 2.350 517 570 77 e5 0.000 000 000 73 e5 T +atomic unit of magnetizability 7.891 036 5794 e-29 0.000 000 0049 e-29 J T^-2 +atomic unit of mass 9.109 383 7139 e-31 0.000 000 0028 e-31 kg +atomic unit of momentum 1.992 851 915 45 e-24 0.000 000 000 31 e-24 kg m s^-1 +atomic unit of permittivity 1.112 650 056 20 e-10 0.000 000 000 17 e-10 F m^-1 +atomic unit of time 2.418 884 326 5864 e-17 0.000 000 000 0026 e-17 s +atomic unit of velocity 2.187 691 262 16 e6 0.000 000 000 34 e6 m s^-1 +Avogadro constant 6.022 140 76 e23 (exact) mol^-1 +Bohr magneton 9.274 010 0657 e-24 0.000 000 0029 e-24 J T^-1 +Bohr magneton in eV/T 5.788 381 7982 e-5 0.000 000 0018 e-5 eV T^-1 +Bohr magneton in Hz/T 1.399 624 491 71 e10 0.000 000 000 44 e10 Hz T^-1 +Bohr magneton in inverse meter per tesla 46.686 447 719 0.000 000 015 m^-1 T^-1 +Bohr magneton in K/T 0.671 713 814 72 0.000 000 000 21 K T^-1 +Bohr radius 5.291 772 105 44 e-11 0.000 000 000 82 e-11 m +Boltzmann constant 1.380 649 e-23 (exact) J K^-1 +Boltzmann constant in eV/K 8.617 333 262... e-5 (exact) eV K^-1 +Boltzmann constant in Hz/K 2.083 661 912... e10 (exact) Hz K^-1 +Boltzmann constant in inverse meter per kelvin 69.503 480 04... (exact) m^-1 K^-1 +characteristic impedance of vacuum 376.730 313 412 0.000 000 059 ohm +classical electron radius 2.817 940 3205 e-15 0.000 000 0013 e-15 m +Compton wavelength 2.426 310 235 38 e-12 0.000 000 000 76 e-12 m +conductance quantum 7.748 091 729... e-5 (exact) S +conventional value of ampere-90 1.000 000 088 87... (exact) A +conventional value of coulomb-90 1.000 000 088 87... (exact) C +conventional value of farad-90 0.999 999 982 20... (exact) F +conventional value of henry-90 1.000 000 017 79... (exact) H +conventional value of Josephson constant 483 597.9 e9 (exact) Hz V^-1 +conventional value of ohm-90 1.000 000 017 79... (exact) ohm +conventional value of volt-90 1.000 000 106 66... (exact) V +conventional value of von Klitzing constant 25 812.807 (exact) ohm +conventional value of watt-90 1.000 000 195 53... (exact) W +Copper x unit 1.002 076 97 e-13 0.000 000 28 e-13 m +deuteron-electron mag. mom. ratio -4.664 345 550 e-4 0.000 000 012 e-4 +deuteron-electron mass ratio 3670.482 967 655 0.000 000 063 +deuteron g factor 0.857 438 2335 0.000 000 0022 +deuteron mag. mom. 4.330 735 087 e-27 0.000 000 011 e-27 J T^-1 +deuteron mag. mom. to Bohr magneton ratio 4.669 754 568 e-4 0.000 000 012 e-4 +deuteron mag. mom. to nuclear magneton ratio 0.857 438 2335 0.000 000 0022 +deuteron mass 3.343 583 7768 e-27 0.000 000 0010 e-27 kg +deuteron mass energy equivalent 3.005 063 234 91 e-10 0.000 000 000 94 e-10 J +deuteron mass energy equivalent in MeV 1875.612 945 00 0.000 000 58 MeV +deuteron mass in u 2.013 553 212 544 0.000 000 000 015 u +deuteron molar mass 2.013 553 214 66 e-3 0.000 000 000 63 e-3 kg mol^-1 +deuteron-neutron mag. mom. ratio -0.448 206 52 0.000 000 11 +deuteron-proton mag. mom. ratio 0.307 012 209 30 0.000 000 000 79 +deuteron-proton mass ratio 1.999 007 501 2699 0.000 000 000 0084 +deuteron relative atomic mass 2.013 553 212 544 0.000 000 000 015 +deuteron rms charge radius 2.127 78 e-15 0.000 27 e-15 m +electron charge to mass quotient -1.758 820 008 38 e11 0.000 000 000 55 e11 C kg^-1 +electron-deuteron mag. mom. ratio -2143.923 4921 0.000 0056 +electron-deuteron mass ratio 2.724 437 107 629 e-4 0.000 000 000 047 e-4 +electron g factor -2.002 319 304 360 92 0.000 000 000 000 36 +electron gyromag. ratio 1.760 859 627 84 e11 0.000 000 000 55 e11 s^-1 T^-1 +electron gyromag. ratio in MHz/T 28 024.951 3861 0.000 0087 MHz T^-1 +electron-helion mass ratio 1.819 543 074 649 e-4 0.000 000 000 053 e-4 +electron mag. mom. -9.284 764 6917 e-24 0.000 000 0029 e-24 J T^-1 +electron mag. mom. anomaly 1.159 652 180 46 e-3 0.000 000 000 18 e-3 +electron mag. mom. to Bohr magneton ratio -1.001 159 652 180 46 0.000 000 000 000 18 +electron mag. mom. to nuclear magneton ratio -1838.281 971 877 0.000 000 032 +electron mass 9.109 383 7139 e-31 0.000 000 0028 e-31 kg +electron mass energy equivalent 8.187 105 7880 e-14 0.000 000 0026 e-14 J +electron mass energy equivalent in MeV 0.510 998 950 69 0.000 000 000 16 MeV +electron mass in u 5.485 799 090 441 e-4 0.000 000 000 097 e-4 u +electron molar mass 5.485 799 0962 e-7 0.000 000 0017 e-7 kg mol^-1 +electron-muon mag. mom. ratio 206.766 9881 0.000 0046 +electron-muon mass ratio 4.836 331 70 e-3 0.000 000 11 e-3 +electron-neutron mag. mom. ratio 960.920 48 0.000 23 +electron-neutron mass ratio 5.438 673 4416 e-4 0.000 000 0022 e-4 +electron-proton mag. mom. ratio -658.210 687 89 0.000 000 19 +electron-proton mass ratio 5.446 170 214 889 e-4 0.000 000 000 094 e-4 +electron relative atomic mass 5.485 799 090 441 e-4 0.000 000 000 097 e-4 +electron-tau mass ratio 2.875 85 e-4 0.000 19 e-4 +electron to alpha particle mass ratio 1.370 933 554 733 e-4 0.000 000 000 032 e-4 +electron to shielded helion mag. mom. ratio 864.058 239 86 0.000 000 70 +electron to shielded proton mag. mom. ratio -658.227 5856 0.000 0027 +electron-triton mass ratio 1.819 200 062 327 e-4 0.000 000 000 068 e-4 +electron volt 1.602 176 634 e-19 (exact) J +electron volt-atomic mass unit relationship 1.073 544 100 83 e-9 0.000 000 000 33 e-9 u +electron volt-hartree relationship 3.674 932 217 5665 e-2 0.000 000 000 0040 e-2 E_h +electron volt-hertz relationship 2.417 989 242... e14 (exact) Hz +electron volt-inverse meter relationship 8.065 543 937... e5 (exact) m^-1 +electron volt-joule relationship 1.602 176 634 e-19 (exact) J +electron volt-kelvin relationship 1.160 451 812... e4 (exact) K +electron volt-kilogram relationship 1.782 661 921... e-36 (exact) kg +elementary charge 1.602 176 634 e-19 (exact) C +elementary charge over h-bar 1.519 267 447... e15 (exact) A J^-1 +Faraday constant 96 485.332 12... (exact) C mol^-1 +Fermi coupling constant 1.166 3787 e-5 0.000 0006 e-5 GeV^-2 +fine-structure constant 7.297 352 5643 e-3 0.000 000 0011 e-3 +first radiation constant 3.741 771 852... e-16 (exact) W m^2 +first radiation constant for spectral radiance 1.191 042 972... e-16 (exact) W m^2 sr^-1 +hartree-atomic mass unit relationship 2.921 262 317 97 e-8 0.000 000 000 91 e-8 u +hartree-electron volt relationship 27.211 386 245 981 0.000 000 000 030 eV +Hartree energy 4.359 744 722 2060 e-18 0.000 000 000 0048 e-18 J +Hartree energy in eV 27.211 386 245 981 0.000 000 000 030 eV +hartree-hertz relationship 6.579 683 920 4999 e15 0.000 000 000 0072 e15 Hz +hartree-inverse meter relationship 2.194 746 313 6314 e7 0.000 000 000 0024 e7 m^-1 +hartree-joule relationship 4.359 744 722 2060 e-18 0.000 000 000 0048 e-18 J +hartree-kelvin relationship 3.157 750 248 0398 e5 0.000 000 000 0034 e5 K +hartree-kilogram relationship 4.850 870 209 5419 e-35 0.000 000 000 0053 e-35 kg +helion-electron mass ratio 5495.885 279 84 0.000 000 16 +helion g factor -4.255 250 6995 0.000 000 0034 +helion mag. mom. -1.074 617 551 98 e-26 0.000 000 000 93 e-26 J T^-1 +helion mag. mom. to Bohr magneton ratio -1.158 740 980 83 e-3 0.000 000 000 94 e-3 +helion mag. mom. to nuclear magneton ratio -2.127 625 3498 0.000 000 0017 +helion mass 5.006 412 7862 e-27 0.000 000 0016 e-27 kg +helion mass energy equivalent 4.499 539 4185 e-10 0.000 000 0014 e-10 J +helion mass energy equivalent in MeV 2808.391 611 12 0.000 000 88 MeV +helion mass in u 3.014 932 246 932 0.000 000 000 074 u +helion molar mass 3.014 932 250 10 e-3 0.000 000 000 94 e-3 kg mol^-1 +helion-proton mass ratio 2.993 152 671 552 0.000 000 000 070 +helion relative atomic mass 3.014 932 246 932 0.000 000 000 074 +helion shielding shift 5.996 7029 e-5 0.000 0023 e-5 +hertz-atomic mass unit relationship 4.439 821 6590 e-24 0.000 000 0014 e-24 u +hertz-electron volt relationship 4.135 667 696... e-15 (exact) eV +hertz-hartree relationship 1.519 829 846 0574 e-16 0.000 000 000 0017 e-16 E_h +hertz-inverse meter relationship 3.335 640 951... e-9 (exact) m^-1 +hertz-joule relationship 6.626 070 15 e-34 (exact) J +hertz-kelvin relationship 4.799 243 073... e-11 (exact) K +hertz-kilogram relationship 7.372 497 323... e-51 (exact) kg +hyperfine transition frequency of Cs-133 9 192 631 770 (exact) Hz +inverse fine-structure constant 137.035 999 177 0.000 000 021 +inverse meter-atomic mass unit relationship 1.331 025 048 24 e-15 0.000 000 000 41 e-15 u +inverse meter-electron volt relationship 1.239 841 984... e-6 (exact) eV +inverse meter-hartree relationship 4.556 335 252 9132 e-8 0.000 000 000 0050 e-8 E_h +inverse meter-hertz relationship 299 792 458 (exact) Hz +inverse meter-joule relationship 1.986 445 857... e-25 (exact) J +inverse meter-kelvin relationship 1.438 776 877... e-2 (exact) K +inverse meter-kilogram relationship 2.210 219 094... e-42 (exact) kg +inverse of conductance quantum 12 906.403 72... (exact) ohm +Josephson constant 483 597.848 4... e9 (exact) Hz V^-1 +joule-atomic mass unit relationship 6.700 535 2471 e9 0.000 000 0021 e9 u +joule-electron volt relationship 6.241 509 074... e18 (exact) eV +joule-hartree relationship 2.293 712 278 3969 e17 0.000 000 000 0025 e17 E_h +joule-hertz relationship 1.509 190 179... e33 (exact) Hz +joule-inverse meter relationship 5.034 116 567... e24 (exact) m^-1 +joule-kelvin relationship 7.242 970 516... e22 (exact) K +joule-kilogram relationship 1.112 650 056... e-17 (exact) kg +kelvin-atomic mass unit relationship 9.251 087 2884 e-14 0.000 000 0029 e-14 u +kelvin-electron volt relationship 8.617 333 262... e-5 (exact) eV +kelvin-hartree relationship 3.166 811 563 4564 e-6 0.000 000 000 0035 e-6 E_h +kelvin-hertz relationship 2.083 661 912... e10 (exact) Hz +kelvin-inverse meter relationship 69.503 480 04... (exact) m^-1 +kelvin-joule relationship 1.380 649 e-23 (exact) J +kelvin-kilogram relationship 1.536 179 187... e-40 (exact) kg +kilogram-atomic mass unit relationship 6.022 140 7537 e26 0.000 000 0019 e26 u +kilogram-electron volt relationship 5.609 588 603... e35 (exact) eV +kilogram-hartree relationship 2.061 485 788 7415 e34 0.000 000 000 0022 e34 E_h +kilogram-hertz relationship 1.356 392 489... e50 (exact) Hz +kilogram-inverse meter relationship 4.524 438 335... e41 (exact) m^-1 +kilogram-joule relationship 8.987 551 787... e16 (exact) J +kilogram-kelvin relationship 6.509 657 260... e39 (exact) K +lattice parameter of silicon 5.431 020 511 e-10 0.000 000 089 e-10 m +lattice spacing of ideal Si (220) 1.920 155 716 e-10 0.000 000 032 e-10 m +Loschmidt constant (273.15 K, 100 kPa) 2.651 645 804... e25 (exact) m^-3 +Loschmidt constant (273.15 K, 101.325 kPa) 2.686 780 111... e25 (exact) m^-3 +luminous efficacy 683 (exact) lm W^-1 +mag. flux quantum 2.067 833 848... e-15 (exact) Wb +molar gas constant 8.314 462 618... (exact) J mol^-1 K^-1 +molar mass constant 1.000 000 001 05 e-3 0.000 000 000 31 e-3 kg mol^-1 +molar mass of carbon-12 12.000 000 0126 e-3 0.000 000 0037 e-3 kg mol^-1 +molar Planck constant 3.990 312 712... e-10 (exact) J Hz^-1 mol^-1 +molar volume of ideal gas (273.15 K, 100 kPa) 22.710 954 64... e-3 (exact) m^3 mol^-1 +molar volume of ideal gas (273.15 K, 101.325 kPa) 22.413 969 54... e-3 (exact) m^3 mol^-1 +molar volume of silicon 1.205 883 199 e-5 0.000 000 060 e-5 m^3 mol^-1 +Molybdenum x unit 1.002 099 52 e-13 0.000 000 53 e-13 m +muon Compton wavelength 1.173 444 110 e-14 0.000 000 026 e-14 m +muon-electron mass ratio 206.768 2827 0.000 0046 +muon g factor -2.002 331 841 23 0.000 000 000 82 +muon mag. mom. -4.490 448 30 e-26 0.000 000 10 e-26 J T^-1 +muon mag. mom. anomaly 1.165 920 62 e-3 0.000 000 41 e-3 +muon mag. mom. to Bohr magneton ratio -4.841 970 48 e-3 0.000 000 11 e-3 +muon mag. mom. to nuclear magneton ratio -8.890 597 04 0.000 000 20 +muon mass 1.883 531 627 e-28 0.000 000 042 e-28 kg +muon mass energy equivalent 1.692 833 804 e-11 0.000 000 038 e-11 J +muon mass energy equivalent in MeV 105.658 3755 0.000 0023 MeV +muon mass in u 0.113 428 9257 0.000 000 0025 u +muon molar mass 1.134 289 258 e-4 0.000 000 025 e-4 kg mol^-1 +muon-neutron mass ratio 0.112 454 5168 0.000 000 0025 +muon-proton mag. mom. ratio -3.183 345 146 0.000 000 071 +muon-proton mass ratio 0.112 609 5262 0.000 000 0025 +muon-tau mass ratio 5.946 35 e-2 0.000 40 e-2 +natural unit of action 1.054 571 817... e-34 (exact) J s +natural unit of action in eV s 6.582 119 569... e-16 (exact) eV s +natural unit of energy 8.187 105 7880 e-14 0.000 000 0026 e-14 J +natural unit of energy in MeV 0.510 998 950 69 0.000 000 000 16 MeV +natural unit of length 3.861 592 6744 e-13 0.000 000 0012 e-13 m +natural unit of mass 9.109 383 7139 e-31 0.000 000 0028 e-31 kg +natural unit of momentum 2.730 924 534 46 e-22 0.000 000 000 85 e-22 kg m s^-1 +natural unit of momentum in MeV/c 0.510 998 950 69 0.000 000 000 16 MeV/c +natural unit of time 1.288 088 666 44 e-21 0.000 000 000 40 e-21 s +natural unit of velocity 299 792 458 (exact) m s^-1 +neutron Compton wavelength 1.319 590 903 82 e-15 0.000 000 000 67 e-15 m +neutron-electron mag. mom. ratio 1.040 668 84 e-3 0.000 000 24 e-3 +neutron-electron mass ratio 1838.683 662 00 0.000 000 74 +neutron g factor -3.826 085 52 0.000 000 90 +neutron gyromag. ratio 1.832 471 74 e8 0.000 000 43 e8 s^-1 T^-1 +neutron gyromag. ratio in MHz/T 29.164 6935 0.000 0069 MHz T^-1 +neutron mag. mom. -9.662 3653 e-27 0.000 0023 e-27 J T^-1 +neutron mag. mom. to Bohr magneton ratio -1.041 875 65 e-3 0.000 000 25 e-3 +neutron mag. mom. to nuclear magneton ratio -1.913 042 76 0.000 000 45 +neutron mass 1.674 927 500 56 e-27 0.000 000 000 85 e-27 kg +neutron mass energy equivalent 1.505 349 765 14 e-10 0.000 000 000 76 e-10 J +neutron mass energy equivalent in MeV 939.565 421 94 0.000 000 48 MeV +neutron mass in u 1.008 664 916 06 0.000 000 000 40 u +neutron molar mass 1.008 664 917 12 e-3 0.000 000 000 51 e-3 kg mol^-1 +neutron-muon mass ratio 8.892 484 08 0.000 000 20 +neutron-proton mag. mom. ratio -0.684 979 35 0.000 000 16 +neutron-proton mass difference 2.305 574 61 e-30 0.000 000 67 e-30 kg +neutron-proton mass difference energy equivalent 2.072 147 12 e-13 0.000 000 60 e-13 J +neutron-proton mass difference energy equivalent in MeV 1.293 332 51 0.000 000 38 MeV +neutron-proton mass difference in u 1.388 449 48 e-3 0.000 000 40 e-3 u +neutron-proton mass ratio 1.001 378 419 46 0.000 000 000 40 +neutron relative atomic mass 1.008 664 916 06 0.000 000 000 40 +neutron-tau mass ratio 0.528 779 0.000 036 +neutron to shielded proton mag. mom. ratio -0.684 996 94 0.000 000 16 +Newtonian constant of gravitation 6.674 30 e-11 0.000 15 e-11 m^3 kg^-1 s^-2 +Newtonian constant of gravitation over h-bar c 6.708 83 e-39 0.000 15 e-39 (GeV/c^2)^-2 +nuclear magneton 5.050 783 7393 e-27 0.000 000 0016 e-27 J T^-1 +nuclear magneton in eV/T 3.152 451 254 17 e-8 0.000 000 000 98 e-8 eV T^-1 +nuclear magneton in inverse meter per tesla 2.542 623 410 09 e-2 0.000 000 000 79 e-2 m^-1 T^-1 +nuclear magneton in K/T 3.658 267 7706 e-4 0.000 000 0011 e-4 K T^-1 +nuclear magneton in MHz/T 7.622 593 2188 0.000 000 0024 MHz T^-1 +Planck constant 6.626 070 15 e-34 (exact) J Hz^-1 +Planck constant in eV/Hz 4.135 667 696... e-15 (exact) eV Hz^-1 +Planck length 1.616 255 e-35 0.000 018 e-35 m +Planck mass 2.176 434 e-8 0.000 024 e-8 kg +Planck mass energy equivalent in GeV 1.220 890 e19 0.000 014 e19 GeV +Planck temperature 1.416 784 e32 0.000 016 e32 K +Planck time 5.391 247 e-44 0.000 060 e-44 s +proton charge to mass quotient 9.578 833 1430 e7 0.000 000 0030 e7 C kg^-1 +proton Compton wavelength 1.321 409 853 60 e-15 0.000 000 000 41 e-15 m +proton-electron mass ratio 1836.152 673 426 0.000 000 032 +proton g factor 5.585 694 6893 0.000 000 0016 +proton gyromag. ratio 2.675 221 8708 e8 0.000 000 0011 e8 s^-1 T^-1 +proton gyromag. ratio in MHz/T 42.577 478 461 0.000 000 018 MHz T^-1 +proton mag. mom. 1.410 606 795 45 e-26 0.000 000 000 60 e-26 J T^-1 +proton mag. mom. to Bohr magneton ratio 1.521 032 202 30 e-3 0.000 000 000 45 e-3 +proton mag. mom. to nuclear magneton ratio 2.792 847 344 63 0.000 000 000 82 +proton mag. shielding correction 2.567 15 e-5 0.000 41 e-5 +proton mass 1.672 621 925 95 e-27 0.000 000 000 52 e-27 kg +proton mass energy equivalent 1.503 277 618 02 e-10 0.000 000 000 47 e-10 J +proton mass energy equivalent in MeV 938.272 089 43 0.000 000 29 MeV +proton mass in u 1.007 276 466 5789 0.000 000 000 0083 u +proton molar mass 1.007 276 467 64 e-3 0.000 000 000 31 e-3 kg mol^-1 +proton-muon mass ratio 8.880 243 38 0.000 000 20 +proton-neutron mag. mom. ratio -1.459 898 02 0.000 000 34 +proton-neutron mass ratio 0.998 623 477 97 0.000 000 000 40 +proton relative atomic mass 1.007 276 466 5789 0.000 000 000 0083 +proton rms charge radius 8.4075 e-16 0.0064 e-16 m +proton-tau mass ratio 0.528 051 0.000 036 +quantum of circulation 3.636 947 5467 e-4 0.000 000 0011 e-4 m^2 s^-1 +quantum of circulation times 2 7.273 895 0934 e-4 0.000 000 0023 e-4 m^2 s^-1 +reduced Compton wavelength 3.861 592 6744 e-13 0.000 000 0012 e-13 m +reduced muon Compton wavelength 1.867 594 306 e-15 0.000 000 042 e-15 m +reduced neutron Compton wavelength 2.100 194 1520 e-16 0.000 000 0011 e-16 m +reduced Planck constant 1.054 571 817... e-34 (exact) J s +reduced Planck constant in eV s 6.582 119 569... e-16 (exact) eV s +reduced Planck constant times c in MeV fm 197.326 980 4... (exact) MeV fm +reduced proton Compton wavelength 2.103 089 100 51 e-16 0.000 000 000 66 e-16 m +reduced tau Compton wavelength 1.110 538 e-16 0.000 075 e-16 m +Rydberg constant 10 973 731.568 157 0.000 012 m^-1 +Rydberg constant times c in Hz 3.289 841 960 2500 e15 0.000 000 000 0036 e15 Hz +Rydberg constant times hc in eV 13.605 693 122 990 0.000 000 000 015 eV +Rydberg constant times hc in J 2.179 872 361 1030 e-18 0.000 000 000 0024 e-18 J +Sackur-Tetrode constant (1 K, 100 kPa) -1.151 707 534 96 0.000 000 000 47 +Sackur-Tetrode constant (1 K, 101.325 kPa) -1.164 870 521 49 0.000 000 000 47 +second radiation constant 1.438 776 877... e-2 (exact) m K +shielded helion gyromag. ratio 2.037 894 6078 e8 0.000 000 0018 e8 s^-1 T^-1 +shielded helion gyromag. ratio in MHz/T 32.434 100 033 0.000 000 028 MHz T^-1 +shielded helion mag. mom. -1.074 553 110 35 e-26 0.000 000 000 93 e-26 J T^-1 +shielded helion mag. mom. to Bohr magneton ratio -1.158 671 494 57 e-3 0.000 000 000 94 e-3 +shielded helion mag. mom. to nuclear magneton ratio -2.127 497 7624 0.000 000 0017 +shielded helion to proton mag. mom. ratio -0.761 766 577 21 0.000 000 000 66 +shielded helion to shielded proton mag. mom. ratio -0.761 786 1334 0.000 000 0031 +shielded proton gyromag. ratio 2.675 153 194 e8 0.000 000 011 e8 s^-1 T^-1 +shielded proton gyromag. ratio in MHz/T 42.576 385 43 0.000 000 17 MHz T^-1 +shielded proton mag. mom. 1.410 570 5830 e-26 0.000 000 0058 e-26 J T^-1 +shielded proton mag. mom. to Bohr magneton ratio 1.520 993 1551 e-3 0.000 000 0062 e-3 +shielded proton mag. mom. to nuclear magneton ratio 2.792 775 648 0.000 000 011 +shielding difference of d and p in HD 1.987 70 e-8 0.000 10 e-8 +shielding difference of t and p in HT 2.394 50 e-8 0.000 20 e-8 +speed of light in vacuum 299 792 458 (exact) m s^-1 +standard acceleration of gravity 9.806 65 (exact) m s^-2 +standard atmosphere 101 325 (exact) Pa +standard-state pressure 100 000 (exact) Pa +Stefan-Boltzmann constant 5.670 374 419... e-8 (exact) W m^-2 K^-4 +tau Compton wavelength 6.977 71 e-16 0.000 47 e-16 m +tau-electron mass ratio 3477.23 0.23 +tau energy equivalent 1776.86 0.12 MeV +tau mass 3.167 54 e-27 0.000 21 e-27 kg +tau mass energy equivalent 2.846 84 e-10 0.000 19 e-10 J +tau mass in u 1.907 54 0.000 13 u +tau molar mass 1.907 54 e-3 0.000 13 e-3 kg mol^-1 +tau-muon mass ratio 16.8170 0.0011 +tau-neutron mass ratio 1.891 15 0.000 13 +tau-proton mass ratio 1.893 76 0.000 13 +Thomson cross section 6.652 458 7051 e-29 0.000 000 0062 e-29 m^2 +triton-electron mass ratio 5496.921 535 51 0.000 000 21 +triton g factor 5.957 924 930 0.000 000 012 +triton mag. mom. 1.504 609 5178 e-26 0.000 000 0030 e-26 J T^-1 +triton mag. mom. to Bohr magneton ratio 1.622 393 6648 e-3 0.000 000 0032 e-3 +triton mag. mom. to nuclear magneton ratio 2.978 962 4650 0.000 000 0059 +triton mass 5.007 356 7512 e-27 0.000 000 0016 e-27 kg +triton mass energy equivalent 4.500 387 8119 e-10 0.000 000 0014 e-10 J +triton mass energy equivalent in MeV 2808.921 136 68 0.000 000 88 MeV +triton mass in u 3.015 500 715 97 0.000 000 000 10 u +triton molar mass 3.015 500 719 13 e-3 0.000 000 000 94 e-3 kg mol^-1 +triton-proton mass ratio 2.993 717 034 03 0.000 000 000 10 +triton relative atomic mass 3.015 500 715 97 0.000 000 000 10 +triton to proton mag. mom. ratio 1.066 639 9189 0.000 000 0021 +unified atomic mass unit 1.660 539 068 92 e-27 0.000 000 000 52 e-27 kg +vacuum electric permittivity 8.854 187 8188 e-12 0.000 000 0014 e-12 F m^-1 +vacuum mag. permeability 1.256 637 061 27 e-6 0.000 000 000 20 e-6 N A^-2 +von Klitzing constant 25 812.807 45... (exact) ohm +weak mixing angle 0.223 05 0.000 23 +Wien frequency displacement law constant 5.878 925 757... e10 (exact) Hz K^-1 +Wien wavelength displacement law constant 2.897 771 955... e-3 (exact) m K +W to Z mass ratio 0.881 45 0.000 13 """ + + +exact2022 = exact2018 + + +# ----------------------------------------------------------------------------- + + +def parse_constants_2002to2014( + d: str, exact_func: Callable[[Any], Any] +) -> dict[str, tuple[float, str, float]]: + constants: dict[str, tuple[float, str, float]] = {} + exact: dict[str, float] = {} + need_replace = set() + for line in d.split('\n'): + name = line[:55].rstrip() + val = float(line[55:77].replace(' ', '').replace('...', '')) + is_truncated = '...' in line[55:77] + is_exact = '(exact)' in line[77:99] + if is_truncated and is_exact: + # missing decimals, use computed exact value + need_replace.add(name) + elif is_exact: + exact[name] = val + else: + assert not is_truncated + uncert = float(line[77:99].replace(' ', '').replace('(exact)', '0')) + units = line[99:].rstrip() + constants[name] = (val, units, uncert) + replace = exact_func(exact) + replace_exact(constants, need_replace, replace) + return constants + + +def parse_constants_2018toXXXX( + d: str, exact_func: Callable[[Any], Any] +) -> dict[str, tuple[float, str, float]]: + constants: dict[str, tuple[float, str, float]] = {} + exact: dict[str, float] = {} + need_replace = set() + for line in d.split('\n'): + name = line[:60].rstrip() + val = float(line[60:85].replace(' ', '').replace('...', '')) + is_truncated = '...' in line[60:85] + is_exact = '(exact)' in line[85:110] + if is_truncated and is_exact: + # missing decimals, use computed exact value + need_replace.add(name) + elif is_exact: + exact[name] = val + else: + assert not is_truncated + uncert = float(line[85:110].replace(' ', '').replace('(exact)', '0')) + units = line[110:].rstrip() + constants[name] = (val, units, uncert) + replace = exact_func(exact) + replace_exact(constants, need_replace, replace) + return constants + + +def replace_exact(d, to_replace, exact): + for name in to_replace: + assert name in exact, f'Missing exact value: {name}' + assert abs(exact[name]/d[name][0] - 1) <= 1e-9, \ + f'Bad exact value: {name}: { exact[name]}, {d[name][0]}' + d[name] = (exact[name],) + d[name][1:] + assert set(exact.keys()) == set(to_replace) + + +_physical_constants_2002 = parse_constants_2002to2014(txt2002, exact2002) +_physical_constants_2006 = parse_constants_2002to2014(txt2006, exact2006) +_physical_constants_2010 = parse_constants_2002to2014(txt2010, exact2010) +_physical_constants_2014 = parse_constants_2002to2014(txt2014, exact2014) +_physical_constants_2018 = parse_constants_2018toXXXX(txt2018, exact2018) +_physical_constants_2022 = parse_constants_2018toXXXX(txt2022, exact2022) + +physical_constants: dict[str, tuple[float, str, float]] = {} +physical_constants.update(_physical_constants_2002) +physical_constants.update(_physical_constants_2006) +physical_constants.update(_physical_constants_2010) +physical_constants.update(_physical_constants_2014) +physical_constants.update(_physical_constants_2018) +physical_constants.update(_physical_constants_2022) +_current_constants = _physical_constants_2022 +_current_codata = "CODATA 2022" + +# check obsolete values +_obsolete_constants = {} +for k in physical_constants: + if k not in _current_constants: + _obsolete_constants[k] = True + +# generate some additional aliases +_aliases = {} +for k in _physical_constants_2002: + if 'magn.' in k: + _aliases[k] = k.replace('magn.', 'mag.') +for k in _physical_constants_2006: + if 'momentum' in k: + _aliases[k] = k.replace('momentum', 'mom.um') +for k in _physical_constants_2018: + if 'momentum' in k: + _aliases[k] = k.replace('momentum', 'mom.um') +for k in _physical_constants_2022: + if 'momentum' in k: + _aliases[k] = k.replace('momentum', 'mom.um') + +# CODATA 2018 and 2022: renamed and no longer exact; use as aliases +_aliases['mag. constant'] = 'vacuum mag. permeability' +_aliases['electric constant'] = 'vacuum electric permittivity' + + +_extra_alias_keys = ['natural unit of velocity', + 'natural unit of action', + 'natural unit of action in eV s', + 'natural unit of mass', + 'natural unit of energy', + 'natural unit of energy in MeV', + 'natural unit of mom.um', + 'natural unit of mom.um in MeV/c', + 'natural unit of length', + 'natural unit of time'] + +# finally, insert aliases for values +for k, v in list(_aliases.items()): + if v in _current_constants or v in _extra_alias_keys: + physical_constants[k] = physical_constants[v] + else: + del _aliases[k] + + +class ConstantWarning(DeprecationWarning): + """Accessing a constant no longer in current CODATA data set""" + pass + + +def _check_obsolete(key: str) -> None: + if key in _obsolete_constants and key not in _aliases: + warnings.warn(f"Constant '{key}' is not in current {_current_codata} data set", + ConstantWarning, stacklevel=3) + + +def value(key: str) -> float: + """ + Value in physical_constants indexed by key + + Parameters + ---------- + key : Python string + Key in dictionary `physical_constants` + + Returns + ------- + value : float + Value in `physical_constants` corresponding to `key` + + Examples + -------- + >>> from scipy import constants + >>> constants.value('elementary charge') + 1.602176634e-19 + + """ + _check_obsolete(key) + return physical_constants[key][0] + + +def unit(key: str) -> str: + """ + Unit in physical_constants indexed by key + + Parameters + ---------- + key : Python string + Key in dictionary `physical_constants` + + Returns + ------- + unit : Python string + Unit in `physical_constants` corresponding to `key` + + Examples + -------- + >>> from scipy import constants + >>> constants.unit('proton mass') + 'kg' + + """ + _check_obsolete(key) + return physical_constants[key][1] + + +def precision(key: str) -> float: + """ + Relative precision in physical_constants indexed by key + + Parameters + ---------- + key : Python string + Key in dictionary `physical_constants` + + Returns + ------- + prec : float + Relative precision in `physical_constants` corresponding to `key` + + Examples + -------- + >>> from scipy import constants + >>> constants.precision('proton mass') + 5.1e-37 + + """ + _check_obsolete(key) + return physical_constants[key][2] / physical_constants[key][0] + + +def find(sub: str | None = None, disp: bool = False) -> Any: + """ + Return list of physical_constant keys containing a given string. + + Parameters + ---------- + sub : str + Sub-string to search keys for. By default, return all keys. + disp : bool + If True, print the keys that are found and return None. + Otherwise, return the list of keys without printing anything. + + Returns + ------- + keys : list or None + If `disp` is False, the list of keys is returned. + Otherwise, None is returned. + + Examples + -------- + >>> from scipy.constants import find, physical_constants + + Which keys in the ``physical_constants`` dictionary contain 'boltzmann'? + + >>> find('boltzmann') + ['Boltzmann constant', + 'Boltzmann constant in Hz/K', + 'Boltzmann constant in eV/K', + 'Boltzmann constant in inverse meter per kelvin', + 'Stefan-Boltzmann constant'] + + Get the constant called 'Boltzmann constant in Hz/K': + + >>> physical_constants['Boltzmann constant in Hz/K'] + (20836619120.0, 'Hz K^-1', 0.0) + + Find constants with 'radius' in the key: + + >>> find('radius') + ['Bohr radius', + 'alpha particle rms charge radius', + 'classical electron radius', + 'deuteron rms charge radius', + 'proton rms charge radius'] + >>> physical_constants['classical electron radius'] + (2.8179403262e-15, 'm', 1.3e-24) + + """ + if sub is None: + result = list(_current_constants.keys()) + else: + result = [key for key in _current_constants + if sub.lower() in key.lower()] + + result.sort() + if disp: + for key in result: + print(key) + return + else: + return result + +# This is not used here, but it must be defined to pass +# scipy/_lib/tests/test_public_api.py::test_private_but_present_deprecation +c = value('speed of light in vacuum') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/_constants.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..9f508081e946f89261f0678a7c44543574a5ca2c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/_constants.py @@ -0,0 +1,366 @@ +""" +Collection of physical constants and conversion factors. + +Most constants are in SI units, so you can do +print '10 mile per minute is', 10*mile/minute, 'm/s or', 10*mile/(minute*knot), 'knots' + +The list is not meant to be comprehensive, but just convenient for everyday use. +""" + +import math as _math +from typing import TYPE_CHECKING, Any + +from ._codata import value as _cd + +if TYPE_CHECKING: + import numpy.typing as npt + +from scipy._lib._array_api import array_namespace, _asarray + + +""" +BasSw 2006 +physical constants: imported from CODATA +unit conversion: see e.g., NIST special publication 811 +Use at own risk: double-check values before calculating your Mars orbit-insertion burn. +Some constants exist in a few variants, which are marked with suffixes. +The ones without any suffix should be the most common ones. +""" + +__all__ = [ + 'Avogadro', 'Boltzmann', 'Btu', 'Btu_IT', 'Btu_th', 'G', + 'Julian_year', 'N_A', 'Planck', 'R', 'Rydberg', + 'Stefan_Boltzmann', 'Wien', 'acre', 'alpha', + 'angstrom', 'arcmin', 'arcminute', 'arcsec', + 'arcsecond', 'astronomical_unit', 'atm', + 'atmosphere', 'atomic_mass', 'atto', 'au', 'bar', + 'barrel', 'bbl', 'blob', 'c', 'calorie', + 'calorie_IT', 'calorie_th', 'carat', 'centi', + 'convert_temperature', 'day', 'deci', 'degree', + 'degree_Fahrenheit', 'deka', 'dyn', 'dyne', 'e', + 'eV', 'electron_mass', 'electron_volt', + 'elementary_charge', 'epsilon_0', 'erg', + 'exa', 'exbi', 'femto', 'fermi', 'fine_structure', + 'fluid_ounce', 'fluid_ounce_US', 'fluid_ounce_imp', + 'foot', 'g', 'gallon', 'gallon_US', 'gallon_imp', + 'gas_constant', 'gibi', 'giga', 'golden', 'golden_ratio', + 'grain', 'gram', 'gravitational_constant', 'h', 'hbar', + 'hectare', 'hecto', 'horsepower', 'hour', 'hp', + 'inch', 'k', 'kgf', 'kibi', 'kilo', 'kilogram_force', + 'kmh', 'knot', 'lambda2nu', 'lb', 'lbf', + 'light_year', 'liter', 'litre', 'long_ton', 'm_e', + 'm_n', 'm_p', 'm_u', 'mach', 'mebi', 'mega', + 'metric_ton', 'micro', 'micron', 'mil', 'mile', + 'milli', 'minute', 'mmHg', 'mph', 'mu_0', 'nano', + 'nautical_mile', 'neutron_mass', 'nu2lambda', + 'ounce', 'oz', 'parsec', 'pebi', 'peta', + 'pi', 'pico', 'point', 'pound', 'pound_force', + 'proton_mass', 'psi', 'pt', 'quecto', 'quetta', 'ronna', 'ronto', + 'short_ton', 'sigma', 'slinch', 'slug', 'speed_of_light', + 'speed_of_sound', 'stone', 'survey_foot', + 'survey_mile', 'tebi', 'tera', 'ton_TNT', + 'torr', 'troy_ounce', 'troy_pound', 'u', + 'week', 'yard', 'year', 'yobi', 'yocto', + 'yotta', 'zebi', 'zepto', 'zero_Celsius', 'zetta' +] + + +# mathematical constants +pi = _math.pi +golden = golden_ratio = (1 + _math.sqrt(5)) / 2 + +# SI prefixes +quetta = 1e30 +ronna = 1e27 +yotta = 1e24 +zetta = 1e21 +exa = 1e18 +peta = 1e15 +tera = 1e12 +giga = 1e9 +mega = 1e6 +kilo = 1e3 +hecto = 1e2 +deka = 1e1 +deci = 1e-1 +centi = 1e-2 +milli = 1e-3 +micro = 1e-6 +nano = 1e-9 +pico = 1e-12 +femto = 1e-15 +atto = 1e-18 +zepto = 1e-21 +yocto = 1e-24 +ronto = 1e-27 +quecto = 1e-30 + +# binary prefixes +kibi = 2**10 +mebi = 2**20 +gibi = 2**30 +tebi = 2**40 +pebi = 2**50 +exbi = 2**60 +zebi = 2**70 +yobi = 2**80 + +# physical constants +c = speed_of_light = _cd('speed of light in vacuum') +mu_0 = _cd('vacuum mag. permeability') +epsilon_0 = _cd('vacuum electric permittivity') +h = Planck = _cd('Planck constant') +hbar = _cd('reduced Planck constant') +G = gravitational_constant = _cd('Newtonian constant of gravitation') +g = _cd('standard acceleration of gravity') +e = elementary_charge = _cd('elementary charge') +R = gas_constant = _cd('molar gas constant') +alpha = fine_structure = _cd('fine-structure constant') +N_A = Avogadro = _cd('Avogadro constant') +k = Boltzmann = _cd('Boltzmann constant') +sigma = Stefan_Boltzmann = _cd('Stefan-Boltzmann constant') +Wien = _cd('Wien wavelength displacement law constant') +Rydberg = _cd('Rydberg constant') + +# mass in kg +gram = 1e-3 +metric_ton = 1e3 +grain = 64.79891e-6 +lb = pound = 7000 * grain # avoirdupois +blob = slinch = pound * g / 0.0254 # lbf*s**2/in (added in 1.0.0) +slug = blob / 12 # lbf*s**2/foot (added in 1.0.0) +oz = ounce = pound / 16 +stone = 14 * pound +long_ton = 2240 * pound +short_ton = 2000 * pound + +troy_ounce = 480 * grain # only for metals / gems +troy_pound = 12 * troy_ounce +carat = 200e-6 + +m_e = electron_mass = _cd('electron mass') +m_p = proton_mass = _cd('proton mass') +m_n = neutron_mass = _cd('neutron mass') +m_u = u = atomic_mass = _cd('atomic mass constant') + +# angle in rad +degree = pi / 180 +arcmin = arcminute = degree / 60 +arcsec = arcsecond = arcmin / 60 + +# time in second +minute = 60.0 +hour = 60 * minute +day = 24 * hour +week = 7 * day +year = 365 * day +Julian_year = 365.25 * day + +# length in meter +inch = 0.0254 +foot = 12 * inch +yard = 3 * foot +mile = 1760 * yard +mil = inch / 1000 +pt = point = inch / 72 # typography +survey_foot = 1200.0 / 3937 +survey_mile = 5280 * survey_foot +nautical_mile = 1852.0 +fermi = 1e-15 +angstrom = 1e-10 +micron = 1e-6 +au = astronomical_unit = 149597870700.0 +light_year = Julian_year * c +parsec = au / arcsec + +# pressure in pascal +atm = atmosphere = _cd('standard atmosphere') +bar = 1e5 +torr = mmHg = atm / 760 +psi = pound * g / (inch * inch) + +# area in meter**2 +hectare = 1e4 +acre = 43560 * foot**2 + +# volume in meter**3 +litre = liter = 1e-3 +gallon = gallon_US = 231 * inch**3 # US +# pint = gallon_US / 8 +fluid_ounce = fluid_ounce_US = gallon_US / 128 +bbl = barrel = 42 * gallon_US # for oil + +gallon_imp = 4.54609e-3 # UK +fluid_ounce_imp = gallon_imp / 160 + +# speed in meter per second +kmh = 1e3 / hour +mph = mile / hour +# approx value of mach at 15 degrees in 1 atm. Is this a common value? +mach = speed_of_sound = 340.5 +knot = nautical_mile / hour + +# temperature in kelvin +zero_Celsius = 273.15 +degree_Fahrenheit = 1/1.8 # only for differences + +# energy in joule +eV = electron_volt = elementary_charge # * 1 Volt +calorie = calorie_th = 4.184 +calorie_IT = 4.1868 +erg = 1e-7 +Btu_th = pound * degree_Fahrenheit * calorie_th / gram +Btu = Btu_IT = pound * degree_Fahrenheit * calorie_IT / gram +ton_TNT = 1e9 * calorie_th +# Wh = watt_hour + +# power in watt +hp = horsepower = 550 * foot * pound * g + +# force in newton +dyn = dyne = 1e-5 +lbf = pound_force = pound * g +kgf = kilogram_force = g # * 1 kg + +# functions for conversions that are not linear + + +def convert_temperature( + val: "npt.ArrayLike", + old_scale: str, + new_scale: str, +) -> Any: + """ + Convert from a temperature scale to another one among Celsius, Kelvin, + Fahrenheit, and Rankine scales. + + Parameters + ---------- + val : array_like + Value(s) of the temperature(s) to be converted expressed in the + original scale. + old_scale : str + Specifies as a string the original scale from which the temperature + value(s) will be converted. Supported scales are Celsius ('Celsius', + 'celsius', 'C' or 'c'), Kelvin ('Kelvin', 'kelvin', 'K', 'k'), + Fahrenheit ('Fahrenheit', 'fahrenheit', 'F' or 'f'), and Rankine + ('Rankine', 'rankine', 'R', 'r'). + new_scale : str + Specifies as a string the new scale to which the temperature + value(s) will be converted. Supported scales are Celsius ('Celsius', + 'celsius', 'C' or 'c'), Kelvin ('Kelvin', 'kelvin', 'K', 'k'), + Fahrenheit ('Fahrenheit', 'fahrenheit', 'F' or 'f'), and Rankine + ('Rankine', 'rankine', 'R', 'r'). + + Returns + ------- + res : float or array of floats + Value(s) of the converted temperature(s) expressed in the new scale. + + Notes + ----- + .. versionadded:: 0.18.0 + + Examples + -------- + >>> from scipy.constants import convert_temperature + >>> import numpy as np + >>> convert_temperature(np.array([-40, 40]), 'Celsius', 'Kelvin') + array([ 233.15, 313.15]) + + """ + xp = array_namespace(val) + _val = _asarray(val, xp=xp, subok=True) + # Convert from `old_scale` to Kelvin + if old_scale.lower() in ['celsius', 'c']: + tempo = _val + zero_Celsius + elif old_scale.lower() in ['kelvin', 'k']: + tempo = _val + elif old_scale.lower() in ['fahrenheit', 'f']: + tempo = (_val - 32) * 5 / 9 + zero_Celsius + elif old_scale.lower() in ['rankine', 'r']: + tempo = _val * 5 / 9 + else: + raise NotImplementedError(f"{old_scale=} is unsupported: supported scales " + "are Celsius, Kelvin, Fahrenheit, and " + "Rankine") + # and from Kelvin to `new_scale`. + if new_scale.lower() in ['celsius', 'c']: + res = tempo - zero_Celsius + elif new_scale.lower() in ['kelvin', 'k']: + res = tempo + elif new_scale.lower() in ['fahrenheit', 'f']: + res = (tempo - zero_Celsius) * 9 / 5 + 32 + elif new_scale.lower() in ['rankine', 'r']: + res = tempo * 9 / 5 + else: + raise NotImplementedError(f"{new_scale=} is unsupported: supported " + "scales are 'Celsius', 'Kelvin', " + "'Fahrenheit', and 'Rankine'") + + return res + + +# optics + + +def lambda2nu(lambda_: "npt.ArrayLike") -> Any: + """ + Convert wavelength to optical frequency + + Parameters + ---------- + lambda_ : array_like + Wavelength(s) to be converted. + + Returns + ------- + nu : float or array of floats + Equivalent optical frequency. + + Notes + ----- + Computes ``nu = c / lambda`` where c = 299792458.0, i.e., the + (vacuum) speed of light in meters/second. + + Examples + -------- + >>> from scipy.constants import lambda2nu, speed_of_light + >>> import numpy as np + >>> lambda2nu(np.array((1, speed_of_light))) + array([ 2.99792458e+08, 1.00000000e+00]) + + """ + xp = array_namespace(lambda_) + return c / _asarray(lambda_, xp=xp, subok=True) + + +def nu2lambda(nu: "npt.ArrayLike") -> Any: + """ + Convert optical frequency to wavelength. + + Parameters + ---------- + nu : array_like + Optical frequency to be converted. + + Returns + ------- + lambda : float or array of floats + Equivalent wavelength(s). + + Notes + ----- + Computes ``lambda = c / nu`` where c = 299792458.0, i.e., the + (vacuum) speed of light in meters/second. + + Examples + -------- + >>> from scipy.constants import nu2lambda, speed_of_light + >>> import numpy as np + >>> nu2lambda(np.array((1, speed_of_light))) + array([ 2.99792458e+08, 1.00000000e+00]) + + """ + xp = array_namespace(nu) + return c / _asarray(nu, xp=xp, subok=True) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/codata.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/codata.py new file mode 100644 index 0000000000000000000000000000000000000000..a1f946e1bdfa13ffc485157d15b4284e5b14994b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/codata.py @@ -0,0 +1,21 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.constants` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + +__all__ = [ # noqa: F822 + 'physical_constants', 'value', 'unit', 'precision', 'find', + 'ConstantWarning', 'k', 'c', + +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="constants", module="codata", + private_modules=["_codata"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/constants.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..2615486c1e1f6715dc4f44b6a8217794b7640d82 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/constants.py @@ -0,0 +1,53 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.constants` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'Avogadro', 'Boltzmann', 'Btu', 'Btu_IT', 'Btu_th', 'G', + 'Julian_year', 'N_A', 'Planck', 'R', 'Rydberg', + 'Stefan_Boltzmann', 'Wien', 'acre', 'alpha', + 'angstrom', 'arcmin', 'arcminute', 'arcsec', + 'arcsecond', 'astronomical_unit', 'atm', + 'atmosphere', 'atomic_mass', 'atto', 'au', 'bar', + 'barrel', 'bbl', 'blob', 'c', 'calorie', + 'calorie_IT', 'calorie_th', 'carat', 'centi', + 'convert_temperature', 'day', 'deci', 'degree', + 'degree_Fahrenheit', 'deka', 'dyn', 'dyne', 'e', + 'eV', 'electron_mass', 'electron_volt', + 'elementary_charge', 'epsilon_0', 'erg', + 'exa', 'exbi', 'femto', 'fermi', 'fine_structure', + 'fluid_ounce', 'fluid_ounce_US', 'fluid_ounce_imp', + 'foot', 'g', 'gallon', 'gallon_US', 'gallon_imp', + 'gas_constant', 'gibi', 'giga', 'golden', 'golden_ratio', + 'grain', 'gram', 'gravitational_constant', 'h', 'hbar', + 'hectare', 'hecto', 'horsepower', 'hour', 'hp', + 'inch', 'k', 'kgf', 'kibi', 'kilo', 'kilogram_force', + 'kmh', 'knot', 'lambda2nu', 'lb', 'lbf', + 'light_year', 'liter', 'litre', 'long_ton', 'm_e', + 'm_n', 'm_p', 'm_u', 'mach', 'mebi', 'mega', + 'metric_ton', 'micro', 'micron', 'mil', 'mile', + 'milli', 'minute', 'mmHg', 'mph', 'mu_0', 'nano', + 'nautical_mile', 'neutron_mass', 'nu2lambda', + 'ounce', 'oz', 'parsec', 'pebi', 'peta', + 'pi', 'pico', 'point', 'pound', 'pound_force', + 'proton_mass', 'psi', 'pt', 'short_ton', + 'sigma', 'slinch', 'slug', 'speed_of_light', + 'speed_of_sound', 'stone', 'survey_foot', + 'survey_mile', 'tebi', 'tera', 'ton_TNT', + 'torr', 'troy_ounce', 'troy_pound', 'u', + 'week', 'yard', 'year', 'yobi', 'yocto', + 'yotta', 'zebi', 'zepto', 'zero_Celsius', 'zetta' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="constants", module="constants", + private_modules=["_constants"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/test_codata.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/test_codata.py new file mode 100644 index 0000000000000000000000000000000000000000..f5a203762274e66f48bea265680471ceec607de4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/test_codata.py @@ -0,0 +1,78 @@ +from scipy.constants import find, value, c, speed_of_light, precision +from numpy.testing import assert_equal, assert_, assert_almost_equal +import scipy.constants._codata as _cd +from scipy import constants + + +def test_find(): + keys = find('weak mixing', disp=False) + assert_equal(keys, ['weak mixing angle']) + + keys = find('qwertyuiop', disp=False) + assert_equal(keys, []) + + keys = find('natural unit', disp=False) + assert_equal(keys, sorted(['natural unit of velocity', + 'natural unit of action', + 'natural unit of action in eV s', + 'natural unit of mass', + 'natural unit of energy', + 'natural unit of energy in MeV', + 'natural unit of momentum', + 'natural unit of momentum in MeV/c', + 'natural unit of length', + 'natural unit of time'])) + + +def test_basic_table_parse(): + c_s = 'speed of light in vacuum' + assert_equal(value(c_s), c) + assert_equal(value(c_s), speed_of_light) + + +def test_basic_lookup(): + assert_equal('%d %s' % (_cd.value('speed of light in vacuum'), + _cd.unit('speed of light in vacuum')), + '299792458 m s^-1') + + +def test_find_all(): + assert_(len(find(disp=False)) > 300) + + +def test_find_single(): + assert_equal(find('Wien freq', disp=False)[0], + 'Wien frequency displacement law constant') + + +def test_2002_vs_2006(): + assert_almost_equal(value('magn. flux quantum'), + value('mag. flux quantum')) + + +def test_exact_values(): + # Check that updating stored values with exact ones worked. + exact = dict((k, v[0]) for k, v in _cd._physical_constants_2018.items()) + replace = _cd.exact2018(exact) + for key, val in replace.items(): + assert_equal(val, value(key)) + assert precision(key) == 0 + + +def test_gh11341(): + # gh-11341 noted that these three constants should exist (for backward + # compatibility) and should always have the same value: + a = constants.epsilon_0 + b = constants.physical_constants['electric constant'][0] + c = constants.physical_constants['vacuum electric permittivity'][0] + assert a == b == c + + +def test_gh14467(): + # gh-14467 noted that some physical constants in CODATA are rounded + # to only ten significant figures even though they are supposed to be + # exact. Check that (at least) the case mentioned in the issue is resolved. + res = constants.physical_constants['Boltzmann constant in eV/K'][0] + ref = (constants.physical_constants['Boltzmann constant'][0] + / constants.physical_constants['elementary charge'][0]) + assert res == ref diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/test_constants.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/test_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..757f6413c6039c6eb2f01a0e9237b702ed2b222b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/constants/tests/test_constants.py @@ -0,0 +1,90 @@ +import pytest + +import scipy.constants as sc +from scipy.conftest import array_api_compatible +from scipy._lib._array_api_no_0d import xp_assert_equal, xp_assert_close +from numpy.testing import assert_allclose + + +pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")] +skip_xp_backends = pytest.mark.skip_xp_backends + + +class TestConvertTemperature: + def test_convert_temperature(self, xp): + xp_assert_equal(sc.convert_temperature(xp.asarray(32.), 'f', 'Celsius'), + xp.asarray(0.0)) + xp_assert_equal(sc.convert_temperature(xp.asarray([0., 0.]), + 'celsius', 'Kelvin'), + xp.asarray([273.15, 273.15])) + xp_assert_equal(sc.convert_temperature(xp.asarray([0., 0.]), 'kelvin', 'c'), + xp.asarray([-273.15, -273.15])) + xp_assert_equal(sc.convert_temperature(xp.asarray([32., 32.]), 'f', 'k'), + xp.asarray([273.15, 273.15])) + xp_assert_equal(sc.convert_temperature(xp.asarray([273.15, 273.15]), + 'kelvin', 'F'), + xp.asarray([32., 32.])) + xp_assert_equal(sc.convert_temperature(xp.asarray([0., 0.]), 'C', 'fahrenheit'), + xp.asarray([32., 32.])) + xp_assert_close(sc.convert_temperature(xp.asarray([0., 0.], dtype=xp.float64), + 'c', 'r'), + xp.asarray([491.67, 491.67], dtype=xp.float64), + rtol=0., atol=1e-13) + xp_assert_close(sc.convert_temperature(xp.asarray([491.67, 491.67], + dtype=xp.float64), + 'Rankine', 'C'), + xp.asarray([0., 0.], dtype=xp.float64), rtol=0., atol=1e-13) + xp_assert_close(sc.convert_temperature(xp.asarray([491.67, 491.67], + dtype=xp.float64), + 'r', 'F'), + xp.asarray([32., 32.], dtype=xp.float64), rtol=0., atol=1e-13) + xp_assert_close(sc.convert_temperature(xp.asarray([32., 32.], dtype=xp.float64), + 'fahrenheit', 'R'), + xp.asarray([491.67, 491.67], dtype=xp.float64), + rtol=0., atol=1e-13) + xp_assert_close(sc.convert_temperature(xp.asarray([273.15, 273.15], + dtype=xp.float64), + 'K', 'R'), + xp.asarray([491.67, 491.67], dtype=xp.float64), + rtol=0., atol=1e-13) + xp_assert_close(sc.convert_temperature(xp.asarray([491.67, 0.], + dtype=xp.float64), + 'rankine', 'kelvin'), + xp.asarray([273.15, 0.], dtype=xp.float64), rtol=0., atol=1e-13) + + @skip_xp_backends(np_only=True, reason='Python list input uses NumPy backend') + def test_convert_temperature_array_like(self): + assert_allclose(sc.convert_temperature([491.67, 0.], 'rankine', 'kelvin'), + [273.15, 0.], rtol=0., atol=1e-13) + + + @skip_xp_backends(np_only=True, reason='Python int input uses NumPy backend') + def test_convert_temperature_errors(self, xp): + with pytest.raises(NotImplementedError, match="old_scale="): + sc.convert_temperature(1, old_scale="cheddar", new_scale="kelvin") + with pytest.raises(NotImplementedError, match="new_scale="): + sc.convert_temperature(1, old_scale="kelvin", new_scale="brie") + + +class TestLambdaToNu: + def test_lambda_to_nu(self, xp): + xp_assert_equal(sc.lambda2nu(xp.asarray([sc.speed_of_light, 1])), + xp.asarray([1, sc.speed_of_light])) + + + @skip_xp_backends(np_only=True, reason='Python list input uses NumPy backend') + def test_lambda_to_nu_array_like(self, xp): + assert_allclose(sc.lambda2nu([sc.speed_of_light, 1]), + [1, sc.speed_of_light]) + + +class TestNuToLambda: + def test_nu_to_lambda(self, xp): + xp_assert_equal(sc.nu2lambda(xp.asarray([sc.speed_of_light, 1])), + xp.asarray([1, sc.speed_of_light])) + + @skip_xp_backends(np_only=True, reason='Python list input uses NumPy backend') + def test_nu_to_lambda_array_like(self, xp): + assert_allclose(sc.nu2lambda([sc.speed_of_light, 1]), + [1, sc.speed_of_light]) + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eab63e3ae2a680e7e9d767c85bfdcc41669f3822 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/__init__.py @@ -0,0 +1,90 @@ +""" +================================ +Datasets (:mod:`scipy.datasets`) +================================ + +.. currentmodule:: scipy.datasets + +Dataset Methods +=============== + +.. autosummary:: + :toctree: generated/ + + ascent + face + electrocardiogram + +Utility Methods +=============== + +.. autosummary:: + :toctree: generated/ + + download_all -- Download all the dataset files to specified path. + clear_cache -- Clear cached dataset directory. + + +Usage of Datasets +================= + +SciPy dataset methods can be simply called as follows: ``'()'`` +This downloads the dataset files over the network once, and saves the cache, +before returning a `numpy.ndarray` object representing the dataset. + +Note that the return data structure and data type might be different for +different dataset methods. For a more detailed example on usage, please look +into the particular dataset method documentation above. + + +How dataset retrieval and storage works +======================================= + +SciPy dataset files are stored within individual GitHub repositories under the +SciPy GitHub organization, following a naming convention as +``'dataset-'``, for example `scipy.datasets.face` files live at +https://github.com/scipy/dataset-face. The `scipy.datasets` submodule utilizes +and depends on `Pooch `_, a Python +package built to simplify fetching data files. Pooch uses these repos to +retrieve the respective dataset files when calling the dataset function. + +A registry of all the datasets, essentially a mapping of filenames with their +SHA256 hash and repo urls are maintained, which Pooch uses to handle and verify +the downloads on function call. After downloading the dataset once, the files +are saved in the system cache directory under ``'scipy-data'``. + +Dataset cache locations may vary on different platforms. + +For macOS:: + + '~/Library/Caches/scipy-data' + +For Linux and other Unix-like platforms:: + + '~/.cache/scipy-data' # or the value of the XDG_CACHE_HOME env var, if defined + +For Windows:: + + 'C:\\Users\\\\AppData\\Local\\\\scipy-data\\Cache' + + +In environments with constrained network connectivity for various security +reasons or on systems without continuous internet connections, one may manually +load the cache of the datasets by placing the contents of the dataset repo in +the above mentioned cache directory to avoid fetching dataset errors without +the internet connectivity. + +""" + + +from ._fetchers import face, ascent, electrocardiogram +from ._download_all import download_all +from ._utils import clear_cache + +__all__ = ['ascent', 'electrocardiogram', 'face', + 'download_all', 'clear_cache'] + + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_download_all.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_download_all.py new file mode 100644 index 0000000000000000000000000000000000000000..4acc02709090857a8c377c21fc3c26068ec0c3a4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_download_all.py @@ -0,0 +1,57 @@ +""" +Platform independent script to download all the +`scipy.datasets` module data files. +This doesn't require a full scipy build. + +Run: python _download_all.py +""" + +import argparse +try: + import pooch +except ImportError: + pooch = None + + +if __package__ is None or __package__ == '': + # Running as python script, use absolute import + import _registry # type: ignore +else: + # Running as python module, use relative import + from . import _registry + + +def download_all(path=None): + """ + Utility method to download all the dataset files + for `scipy.datasets` module. + + Parameters + ---------- + path : str, optional + Directory path to download all the dataset files. + If None, default to the system cache_dir detected by pooch. + """ + if pooch is None: + raise ImportError("Missing optional dependency 'pooch' required " + "for scipy.datasets module. Please use pip or " + "conda to install 'pooch'.") + if path is None: + path = pooch.os_cache('scipy-data') + for dataset_name, dataset_hash in _registry.registry.items(): + pooch.retrieve(url=_registry.registry_urls[dataset_name], + known_hash=dataset_hash, + fname=dataset_name, path=path) + + +def main(): + parser = argparse.ArgumentParser(description='Download SciPy data files.') + parser.add_argument("path", nargs='?', type=str, + default=pooch.os_cache('scipy-data'), + help="Directory path to download all the data files.") + args = parser.parse_args() + download_all(args.path) + + +if __name__ == "__main__": + main() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_fetchers.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_fetchers.py new file mode 100644 index 0000000000000000000000000000000000000000..2985690468095e345ebc2750df420fdb86cc4e21 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_fetchers.py @@ -0,0 +1,219 @@ +from numpy import array, frombuffer, load +from ._registry import registry, registry_urls + +try: + import pooch +except ImportError: + pooch = None + data_fetcher = None +else: + data_fetcher = pooch.create( + # Use the default cache folder for the operating system + # Pooch uses appdirs (https://github.com/ActiveState/appdirs) to + # select an appropriate directory for the cache on each platform. + path=pooch.os_cache("scipy-data"), + + # The remote data is on Github + # base_url is a required param, even though we override this + # using individual urls in the registry. + base_url="https://github.com/scipy/", + registry=registry, + urls=registry_urls + ) + + +def fetch_data(dataset_name, data_fetcher=data_fetcher): + if data_fetcher is None: + raise ImportError("Missing optional dependency 'pooch' required " + "for scipy.datasets module. Please use pip or " + "conda to install 'pooch'.") + # The "fetch" method returns the full path to the downloaded data file. + return data_fetcher.fetch(dataset_name) + + +def ascent(): + """ + Get an 8-bit grayscale bit-depth, 512 x 512 derived image for easy + use in demos. + + The image is derived from + https://pixnio.com/people/accent-to-the-top + + Parameters + ---------- + None + + Returns + ------- + ascent : ndarray + convenient image to use for testing and demonstration + + Examples + -------- + >>> import scipy.datasets + >>> ascent = scipy.datasets.ascent() + >>> ascent.shape + (512, 512) + >>> ascent.max() + np.uint8(255) + + >>> import matplotlib.pyplot as plt + >>> plt.gray() + >>> plt.imshow(ascent) + >>> plt.show() + + """ + import pickle + + # The file will be downloaded automatically the first time this is run, + # returning the path to the downloaded file. Afterwards, Pooch finds + # it in the local cache and doesn't repeat the download. + fname = fetch_data("ascent.dat") + # Now we just need to load it with our standard Python tools. + with open(fname, 'rb') as f: + ascent = array(pickle.load(f)) + return ascent + + +def electrocardiogram(): + """ + Load an electrocardiogram as an example for a 1-D signal. + + The returned signal is a 5 minute long electrocardiogram (ECG), a medical + recording of the heart's electrical activity, sampled at 360 Hz. + + Returns + ------- + ecg : ndarray + The electrocardiogram in millivolt (mV) sampled at 360 Hz. + + Notes + ----- + The provided signal is an excerpt (19:35 to 24:35) from the `record 208`_ + (lead MLII) provided by the MIT-BIH Arrhythmia Database [1]_ on + PhysioNet [2]_. The excerpt includes noise induced artifacts, typical + heartbeats as well as pathological changes. + + .. _record 208: https://physionet.org/physiobank/database/html/mitdbdir/records.htm#208 + + .. versionadded:: 1.1.0 + + References + ---------- + .. [1] Moody GB, Mark RG. The impact of the MIT-BIH Arrhythmia Database. + IEEE Eng in Med and Biol 20(3):45-50 (May-June 2001). + (PMID: 11446209); :doi:`10.13026/C2F305` + .. [2] Goldberger AL, Amaral LAN, Glass L, Hausdorff JM, Ivanov PCh, + Mark RG, Mietus JE, Moody GB, Peng C-K, Stanley HE. PhysioBank, + PhysioToolkit, and PhysioNet: Components of a New Research Resource + for Complex Physiologic Signals. Circulation 101(23):e215-e220; + :doi:`10.1161/01.CIR.101.23.e215` + + Examples + -------- + >>> from scipy.datasets import electrocardiogram + >>> ecg = electrocardiogram() + >>> ecg + array([-0.245, -0.215, -0.185, ..., -0.405, -0.395, -0.385], shape=(108000,)) + >>> ecg.shape, ecg.mean(), ecg.std() + ((108000,), -0.16510875, 0.5992473991177294) + + As stated the signal features several areas with a different morphology. + E.g., the first few seconds show the electrical activity of a heart in + normal sinus rhythm as seen below. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> fs = 360 + >>> time = np.arange(ecg.size) / fs + >>> plt.plot(time, ecg) + >>> plt.xlabel("time in s") + >>> plt.ylabel("ECG in mV") + >>> plt.xlim(9, 10.2) + >>> plt.ylim(-1, 1.5) + >>> plt.show() + + After second 16, however, the first premature ventricular contractions, + also called extrasystoles, appear. These have a different morphology + compared to typical heartbeats. The difference can easily be observed + in the following plot. + + >>> plt.plot(time, ecg) + >>> plt.xlabel("time in s") + >>> plt.ylabel("ECG in mV") + >>> plt.xlim(46.5, 50) + >>> plt.ylim(-2, 1.5) + >>> plt.show() + + At several points large artifacts disturb the recording, e.g.: + + >>> plt.plot(time, ecg) + >>> plt.xlabel("time in s") + >>> plt.ylabel("ECG in mV") + >>> plt.xlim(207, 215) + >>> plt.ylim(-2, 3.5) + >>> plt.show() + + Finally, examining the power spectrum reveals that most of the biosignal is + made up of lower frequencies. At 60 Hz the noise induced by the mains + electricity can be clearly observed. + + >>> from scipy.signal import welch + >>> f, Pxx = welch(ecg, fs=fs, nperseg=2048, scaling="spectrum") + >>> plt.semilogy(f, Pxx) + >>> plt.xlabel("Frequency in Hz") + >>> plt.ylabel("Power spectrum of the ECG in mV**2") + >>> plt.xlim(f[[0, -1]]) + >>> plt.show() + """ + fname = fetch_data("ecg.dat") + with load(fname) as file: + ecg = file["ecg"].astype(int) # np.uint16 -> int + # Convert raw output of ADC to mV: (ecg - adc_zero) / adc_gain + ecg = (ecg - 1024) / 200.0 + return ecg + + +def face(gray=False): + """ + Get a 1024 x 768, color image of a raccoon face. + + The image is derived from + https://pixnio.com/fauna-animals/raccoons/raccoon-procyon-lotor + + Parameters + ---------- + gray : bool, optional + If True return 8-bit grey-scale image, otherwise return a color image + + Returns + ------- + face : ndarray + image of a raccoon face + + Examples + -------- + >>> import scipy.datasets + >>> face = scipy.datasets.face() + >>> face.shape + (768, 1024, 3) + >>> face.max() + np.uint8(255) + + >>> import matplotlib.pyplot as plt + >>> plt.gray() + >>> plt.imshow(face) + >>> plt.show() + + """ + import bz2 + fname = fetch_data("face.dat") + with open(fname, 'rb') as f: + rawdata = f.read() + face_data = bz2.decompress(rawdata) + face = frombuffer(face_data, dtype='uint8') + face.shape = (768, 1024, 3) + if gray is True: + face = (0.21 * face[:, :, 0] + 0.71 * face[:, :, 1] + + 0.07 * face[:, :, 2]).astype('uint8') + return face diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_registry.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..9df4cf4de4d23ebdd8c036b26daee004663dfe0f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_registry.py @@ -0,0 +1,26 @@ +########################################################################## +# This file serves as the dataset registry for SciPy Datasets SubModule. +########################################################################## + + +# To generate the SHA256 hash, use the command +# openssl sha256 +registry = { + "ascent.dat": "03ce124c1afc880f87b55f6b061110e2e1e939679184f5614e38dacc6c1957e2", + "ecg.dat": "f20ad3365fb9b7f845d0e5c48b6fe67081377ee466c3a220b7f69f35c8958baf", + "face.dat": "9d8b0b4d081313e2b485748c770472e5a95ed1738146883d84c7030493e82886" +} + +registry_urls = { + "ascent.dat": "https://raw.githubusercontent.com/scipy/dataset-ascent/main/ascent.dat", + "ecg.dat": "https://raw.githubusercontent.com/scipy/dataset-ecg/main/ecg.dat", + "face.dat": "https://raw.githubusercontent.com/scipy/dataset-face/main/face.dat" +} + +# dataset method mapping with their associated filenames +# : ["filename1", "filename2", ...] +method_files_map = { + "ascent": ["ascent.dat"], + "electrocardiogram": ["ecg.dat"], + "face": ["face.dat"] +} diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_utils.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2ab1b4037a9f6eca22eea495b02878323be93c35 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/_utils.py @@ -0,0 +1,81 @@ +import os +import shutil +from ._registry import method_files_map + +try: + import platformdirs +except ImportError: + platformdirs = None # type: ignore[assignment] + + +def _clear_cache(datasets, cache_dir=None, method_map=None): + if method_map is None: + # Use SciPy Datasets method map + method_map = method_files_map + if cache_dir is None: + # Use default cache_dir path + if platformdirs is None: + # platformdirs is pooch dependency + raise ImportError("Missing optional dependency 'pooch' required " + "for scipy.datasets module. Please use pip or " + "conda to install 'pooch'.") + cache_dir = platformdirs.user_cache_dir("scipy-data") + + if not os.path.exists(cache_dir): + print(f"Cache Directory {cache_dir} doesn't exist. Nothing to clear.") + return + + if datasets is None: + print(f"Cleaning the cache directory {cache_dir}!") + shutil.rmtree(cache_dir) + else: + if not isinstance(datasets, (list, tuple)): + # single dataset method passed should be converted to list + datasets = [datasets, ] + for dataset in datasets: + assert callable(dataset) + dataset_name = dataset.__name__ # Name of the dataset method + if dataset_name not in method_map: + raise ValueError(f"Dataset method {dataset_name} doesn't " + "exist. Please check if the passed dataset " + "is a subset of the following dataset " + f"methods: {list(method_map.keys())}") + + data_files = method_map[dataset_name] + data_filepaths = [os.path.join(cache_dir, file) + for file in data_files] + for data_filepath in data_filepaths: + if os.path.exists(data_filepath): + print("Cleaning the file " + f"{os.path.split(data_filepath)[1]} " + f"for dataset {dataset_name}") + os.remove(data_filepath) + else: + print(f"Path {data_filepath} doesn't exist. " + "Nothing to clear.") + + +def clear_cache(datasets=None): + """ + Cleans the scipy datasets cache directory. + + If a scipy.datasets method or a list/tuple of the same is + provided, then clear_cache removes all the data files + associated to the passed dataset method callable(s). + + By default, it removes all the cached data files. + + Parameters + ---------- + datasets : callable or list/tuple of callable or None + + Examples + -------- + >>> from scipy import datasets + >>> ascent_array = datasets.ascent() + >>> ascent_array.shape + (512, 512) + >>> datasets.clear_cache([datasets.ascent]) + Cleaning the file ascent.dat for dataset ascent + """ + _clear_cache(datasets) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/tests/test_data.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/tests/test_data.py new file mode 100644 index 0000000000000000000000000000000000000000..ebc27bdae480ef00346f3b9808884fbc48738512 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/datasets/tests/test_data.py @@ -0,0 +1,128 @@ +from scipy.datasets._registry import registry +from scipy.datasets._fetchers import data_fetcher +from scipy.datasets._utils import _clear_cache +from scipy.datasets import ascent, face, electrocardiogram, download_all +from numpy.testing import assert_equal, assert_almost_equal +import os +from threading import get_ident +import pytest + +try: + import pooch +except ImportError: + raise ImportError("Missing optional dependency 'pooch' required " + "for scipy.datasets module. Please use pip or " + "conda to install 'pooch'.") + + +data_dir = data_fetcher.path # type: ignore + + +def _has_hash(path, expected_hash): + """Check if the provided path has the expected hash.""" + if not os.path.exists(path): + return False + return pooch.file_hash(path) == expected_hash + + +class TestDatasets: + + @pytest.fixture(scope='module', autouse=True) + def test_download_all(self): + # This fixture requires INTERNET CONNECTION + + # test_setup phase + download_all() + + yield + + @pytest.mark.fail_slow(10) + def test_existence_all(self): + assert len(os.listdir(data_dir)) >= len(registry) + + def test_ascent(self): + assert_equal(ascent().shape, (512, 512)) + + # hash check + assert _has_hash(os.path.join(data_dir, "ascent.dat"), + registry["ascent.dat"]) + + def test_face(self): + assert_equal(face().shape, (768, 1024, 3)) + + # hash check + assert _has_hash(os.path.join(data_dir, "face.dat"), + registry["face.dat"]) + + def test_electrocardiogram(self): + # Test shape, dtype and stats of signal + ecg = electrocardiogram() + assert_equal(ecg.dtype, float) + assert_equal(ecg.shape, (108000,)) + assert_almost_equal(ecg.mean(), -0.16510875) + assert_almost_equal(ecg.std(), 0.5992473991177294) + + # hash check + assert _has_hash(os.path.join(data_dir, "ecg.dat"), + registry["ecg.dat"]) + + +def test_clear_cache(tmp_path): + # Note: `tmp_path` is a pytest fixture, it handles cleanup + thread_basepath = tmp_path / str(get_ident()) + thread_basepath.mkdir() + + dummy_basepath = thread_basepath / "dummy_cache_dir" + dummy_basepath.mkdir() + + # Create three dummy dataset files for dummy dataset methods + dummy_method_map = {} + for i in range(4): + dummy_method_map[f"data{i}"] = [f"data{i}.dat"] + data_filepath = dummy_basepath / f"data{i}.dat" + data_filepath.write_text("") + + # clear files associated to single dataset method data0 + # also test callable argument instead of list of callables + def data0(): + pass + _clear_cache(datasets=data0, cache_dir=dummy_basepath, + method_map=dummy_method_map) + assert not os.path.exists(dummy_basepath/"data0.dat") + + # clear files associated to multiple dataset methods "data3" and "data4" + def data1(): + pass + + def data2(): + pass + _clear_cache(datasets=[data1, data2], cache_dir=dummy_basepath, + method_map=dummy_method_map) + assert not os.path.exists(dummy_basepath/"data1.dat") + assert not os.path.exists(dummy_basepath/"data2.dat") + + # clear multiple dataset files "data3_0.dat" and "data3_1.dat" + # associated with dataset method "data3" + def data4(): + pass + # create files + (dummy_basepath / "data4_0.dat").write_text("") + (dummy_basepath / "data4_1.dat").write_text("") + + dummy_method_map["data4"] = ["data4_0.dat", "data4_1.dat"] + _clear_cache(datasets=[data4], cache_dir=dummy_basepath, + method_map=dummy_method_map) + assert not os.path.exists(dummy_basepath/"data4_0.dat") + assert not os.path.exists(dummy_basepath/"data4_1.dat") + + # wrong dataset method should raise ValueError since it + # doesn't exist in the dummy_method_map + def data5(): + pass + with pytest.raises(ValueError): + _clear_cache(datasets=[data5], cache_dir=dummy_basepath, + method_map=dummy_method_map) + + # remove all dataset cache + _clear_cache(datasets=None, cache_dir=dummy_basepath) + assert not os.path.exists(dummy_basepath) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..940cd71abc0b7ea954f04433ca261d542e664b25 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/__init__.py @@ -0,0 +1,27 @@ +""" +============================================================== +Finite Difference Differentiation (:mod:`scipy.differentiate`) +============================================================== + +.. currentmodule:: scipy.differentiate + +SciPy ``differentiate`` provides functions for performing finite difference +numerical differentiation of black-box functions. + +.. autosummary:: + :toctree: generated/ + + derivative + jacobian + hessian + +""" + + +from ._differentiate import * + +__all__ = ['derivative', 'jacobian', 'hessian'] + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/_differentiate.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/_differentiate.py new file mode 100644 index 0000000000000000000000000000000000000000..5637f066eab178f77a1ebbe7304e174f6497747f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/_differentiate.py @@ -0,0 +1,1129 @@ +# mypy: disable-error-code="attr-defined" +import warnings +import numpy as np +import scipy._lib._elementwise_iterative_method as eim +from scipy._lib._util import _RichResult +from scipy._lib._array_api import array_namespace, xp_sign, xp_copy, xp_take_along_axis + +_EERRORINCREASE = -1 # used in derivative + +def _derivative_iv(f, x, args, tolerances, maxiter, order, initial_step, + step_factor, step_direction, preserve_shape, callback): + # Input validation for `derivative` + xp = array_namespace(x) + + if not callable(f): + raise ValueError('`f` must be callable.') + + if not np.iterable(args): + args = (args,) + + tolerances = {} if tolerances is None else tolerances + atol = tolerances.get('atol', None) + rtol = tolerances.get('rtol', None) + + # tolerances are floats, not arrays; OK to use NumPy + message = 'Tolerances and step parameters must be non-negative scalars.' + tols = np.asarray([atol if atol is not None else 1, + rtol if rtol is not None else 1, + step_factor]) + if (not np.issubdtype(tols.dtype, np.number) or np.any(tols < 0) + or np.any(np.isnan(tols)) or tols.shape != (3,)): + raise ValueError(message) + step_factor = float(tols[2]) + + maxiter_int = int(maxiter) + if maxiter != maxiter_int or maxiter <= 0: + raise ValueError('`maxiter` must be a positive integer.') + + order_int = int(order) + if order_int != order or order <= 0: + raise ValueError('`order` must be a positive integer.') + + step_direction = xp.asarray(step_direction) + initial_step = xp.asarray(initial_step) + temp = xp.broadcast_arrays(x, step_direction, initial_step) + x, step_direction, initial_step = temp + + message = '`preserve_shape` must be True or False.' + if preserve_shape not in {True, False}: + raise ValueError(message) + + if callback is not None and not callable(callback): + raise ValueError('`callback` must be callable.') + + return (f, x, args, atol, rtol, maxiter_int, order_int, initial_step, + step_factor, step_direction, preserve_shape, callback) + + +def derivative(f, x, *, args=(), tolerances=None, maxiter=10, + order=8, initial_step=0.5, step_factor=2.0, + step_direction=0, preserve_shape=False, callback=None): + """Evaluate the derivative of a elementwise, real scalar function numerically. + + For each element of the output of `f`, `derivative` approximates the first + derivative of `f` at the corresponding element of `x` using finite difference + differentiation. + + This function works elementwise when `x`, `step_direction`, and `args` contain + (broadcastable) arrays. + + Parameters + ---------- + f : callable + The function whose derivative is desired. The signature must be:: + + f(xi: ndarray, *argsi) -> ndarray + + where each element of ``xi`` is a finite real number and ``argsi`` is a tuple, + which may contain an arbitrary number of arrays that are broadcastable with + ``xi``. `f` must be an elementwise function: each scalar element ``f(xi)[j]`` + must equal ``f(xi[j])`` for valid indices ``j``. It must not mutate the array + ``xi`` or the arrays in ``argsi``. + x : float array_like + Abscissae at which to evaluate the derivative. Must be broadcastable with + `args` and `step_direction`. + args : tuple of array_like, optional + Additional positional array arguments to be passed to `f`. Arrays + must be broadcastable with one another and the arrays of `init`. + If the callable for which the root is desired requires arguments that are + not broadcastable with `x`, wrap that callable with `f` such that `f` + accepts only `x` and broadcastable ``*args``. + tolerances : dictionary of floats, optional + Absolute and relative tolerances. Valid keys of the dictionary are: + + - ``atol`` - absolute tolerance on the derivative + - ``rtol`` - relative tolerance on the derivative + + Iteration will stop when ``res.error < atol + rtol * abs(res.df)``. The default + `atol` is the smallest normal number of the appropriate dtype, and + the default `rtol` is the square root of the precision of the + appropriate dtype. + order : int, default: 8 + The (positive integer) order of the finite difference formula to be + used. Odd integers will be rounded up to the next even integer. + initial_step : float array_like, default: 0.5 + The (absolute) initial step size for the finite difference derivative + approximation. + step_factor : float, default: 2.0 + The factor by which the step size is *reduced* in each iteration; i.e. + the step size in iteration 1 is ``initial_step/step_factor``. If + ``step_factor < 1``, subsequent steps will be greater than the initial + step; this may be useful if steps smaller than some threshold are + undesirable (e.g. due to subtractive cancellation error). + maxiter : int, default: 10 + The maximum number of iterations of the algorithm to perform. See + Notes. + step_direction : integer array_like + An array representing the direction of the finite difference steps (for + use when `x` lies near to the boundary of the domain of the function.) + Must be broadcastable with `x` and all `args`. + Where 0 (default), central differences are used; where negative (e.g. + -1), steps are non-positive; and where positive (e.g. 1), all steps are + non-negative. + preserve_shape : bool, default: False + In the following, "arguments of `f`" refers to the array ``xi`` and + any arrays within ``argsi``. Let ``shape`` be the broadcasted shape + of `x` and all elements of `args` (which is conceptually + distinct from ``xi` and ``argsi`` passed into `f`). + + - When ``preserve_shape=False`` (default), `f` must accept arguments + of *any* broadcastable shapes. + + - When ``preserve_shape=True``, `f` must accept arguments of shape + ``shape`` *or* ``shape + (n,)``, where ``(n,)`` is the number of + abscissae at which the function is being evaluated. + + In either case, for each scalar element ``xi[j]`` within ``xi``, the array + returned by `f` must include the scalar ``f(xi[j])`` at the same index. + Consequently, the shape of the output is always the shape of the input + ``xi``. + + See Examples. + callback : callable, optional + An optional user-supplied function to be called before the first + iteration and after each iteration. + Called as ``callback(res)``, where ``res`` is a ``_RichResult`` + similar to that returned by `derivative` (but containing the current + iterate's values of all variables). If `callback` raises a + ``StopIteration``, the algorithm will terminate immediately and + `derivative` will return a result. `callback` must not mutate + `res` or its attributes. + + Returns + ------- + res : _RichResult + An object similar to an instance of `scipy.optimize.OptimizeResult` with the + following attributes. The descriptions are written as though the values will + be scalars; however, if `f` returns an array, the outputs will be + arrays of the same shape. + + success : bool array + ``True`` where the algorithm terminated successfully (status ``0``); + ``False`` otherwise. + status : int array + An integer representing the exit status of the algorithm. + + - ``0`` : The algorithm converged to the specified tolerances. + - ``-1`` : The error estimate increased, so iteration was terminated. + - ``-2`` : The maximum number of iterations was reached. + - ``-3`` : A non-finite value was encountered. + - ``-4`` : Iteration was terminated by `callback`. + - ``1`` : The algorithm is proceeding normally (in `callback` only). + + df : float array + The derivative of `f` at `x`, if the algorithm terminated + successfully. + error : float array + An estimate of the error: the magnitude of the difference between + the current estimate of the derivative and the estimate in the + previous iteration. + nit : int array + The number of iterations of the algorithm that were performed. + nfev : int array + The number of points at which `f` was evaluated. + x : float array + The value at which the derivative of `f` was evaluated + (after broadcasting with `args` and `step_direction`). + + See Also + -------- + jacobian, hessian + + Notes + ----- + The implementation was inspired by jacobi [1]_, numdifftools [2]_, and + DERIVEST [3]_, but the implementation follows the theory of Taylor series + more straightforwardly (and arguably naively so). + In the first iteration, the derivative is estimated using a finite + difference formula of order `order` with maximum step size `initial_step`. + Each subsequent iteration, the maximum step size is reduced by + `step_factor`, and the derivative is estimated again until a termination + condition is reached. The error estimate is the magnitude of the difference + between the current derivative approximation and that of the previous + iteration. + + The stencils of the finite difference formulae are designed such that + abscissae are "nested": after `f` is evaluated at ``order + 1`` + points in the first iteration, `f` is evaluated at only two new points + in each subsequent iteration; ``order - 1`` previously evaluated function + values required by the finite difference formula are reused, and two + function values (evaluations at the points furthest from `x`) are unused. + + Step sizes are absolute. When the step size is small relative to the + magnitude of `x`, precision is lost; for example, if `x` is ``1e20``, the + default initial step size of ``0.5`` cannot be resolved. Accordingly, + consider using larger initial step sizes for large magnitudes of `x`. + + The default tolerances are challenging to satisfy at points where the + true derivative is exactly zero. If the derivative may be exactly zero, + consider specifying an absolute tolerance (e.g. ``atol=1e-12``) to + improve convergence. + + References + ---------- + .. [1] Hans Dembinski (@HDembinski). jacobi. + https://github.com/HDembinski/jacobi + .. [2] Per A. Brodtkorb and John D'Errico. numdifftools. + https://numdifftools.readthedocs.io/en/latest/ + .. [3] John D'Errico. DERIVEST: Adaptive Robust Numerical Differentiation. + https://www.mathworks.com/matlabcentral/fileexchange/13490-adaptive-robust-numerical-differentiation + .. [4] Numerical Differentition. Wikipedia. + https://en.wikipedia.org/wiki/Numerical_differentiation + + Examples + -------- + Evaluate the derivative of ``np.exp`` at several points ``x``. + + >>> import numpy as np + >>> from scipy.differentiate import derivative + >>> f = np.exp + >>> df = np.exp # true derivative + >>> x = np.linspace(1, 2, 5) + >>> res = derivative(f, x) + >>> res.df # approximation of the derivative + array([2.71828183, 3.49034296, 4.48168907, 5.75460268, 7.3890561 ]) + >>> res.error # estimate of the error + array([7.13740178e-12, 9.16600129e-12, 1.17594823e-11, 1.51061386e-11, + 1.94262384e-11]) + >>> abs(res.df - df(x)) # true error + array([2.53130850e-14, 3.55271368e-14, 5.77315973e-14, 5.59552404e-14, + 6.92779167e-14]) + + Show the convergence of the approximation as the step size is reduced. + Each iteration, the step size is reduced by `step_factor`, so for + sufficiently small initial step, each iteration reduces the error by a + factor of ``1/step_factor**order`` until finite precision arithmetic + inhibits further improvement. + + >>> import matplotlib.pyplot as plt + >>> iter = list(range(1, 12)) # maximum iterations + >>> hfac = 2 # step size reduction per iteration + >>> hdir = [-1, 0, 1] # compare left-, central-, and right- steps + >>> order = 4 # order of differentiation formula + >>> x = 1 + >>> ref = df(x) + >>> errors = [] # true error + >>> for i in iter: + ... res = derivative(f, x, maxiter=i, step_factor=hfac, + ... step_direction=hdir, order=order, + ... # prevent early termination + ... tolerances=dict(atol=0, rtol=0)) + ... errors.append(abs(res.df - ref)) + >>> errors = np.array(errors) + >>> plt.semilogy(iter, errors[:, 0], label='left differences') + >>> plt.semilogy(iter, errors[:, 1], label='central differences') + >>> plt.semilogy(iter, errors[:, 2], label='right differences') + >>> plt.xlabel('iteration') + >>> plt.ylabel('error') + >>> plt.legend() + >>> plt.show() + >>> (errors[1, 1] / errors[0, 1], 1 / hfac**order) + (0.06215223140159822, 0.0625) + + The implementation is vectorized over `x`, `step_direction`, and `args`. + The function is evaluated once before the first iteration to perform input + validation and standardization, and once per iteration thereafter. + + >>> def f(x, p): + ... f.nit += 1 + ... return x**p + >>> f.nit = 0 + >>> def df(x, p): + ... return p*x**(p-1) + >>> x = np.arange(1, 5) + >>> p = np.arange(1, 6).reshape((-1, 1)) + >>> hdir = np.arange(-1, 2).reshape((-1, 1, 1)) + >>> res = derivative(f, x, args=(p,), step_direction=hdir, maxiter=1) + >>> np.allclose(res.df, df(x, p)) + True + >>> res.df.shape + (3, 5, 4) + >>> f.nit + 2 + + By default, `preserve_shape` is False, and therefore the callable + `f` may be called with arrays of any broadcastable shapes. + For example: + + >>> shapes = [] + >>> def f(x, c): + ... shape = np.broadcast_shapes(x.shape, c.shape) + ... shapes.append(shape) + ... return np.sin(c*x) + >>> + >>> c = [1, 5, 10, 20] + >>> res = derivative(f, 0, args=(c,)) + >>> shapes + [(4,), (4, 8), (4, 2), (3, 2), (2, 2), (1, 2)] + + To understand where these shapes are coming from - and to better + understand how `derivative` computes accurate results - note that + higher values of ``c`` correspond with higher frequency sinusoids. + The higher frequency sinusoids make the function's derivative change + faster, so more function evaluations are required to achieve the target + accuracy: + + >>> res.nfev + array([11, 13, 15, 17], dtype=int32) + + The initial ``shape``, ``(4,)``, corresponds with evaluating the + function at a single abscissa and all four frequencies; this is used + for input validation and to determine the size and dtype of the arrays + that store results. The next shape corresponds with evaluating the + function at an initial grid of abscissae and all four frequencies. + Successive calls to the function evaluate the function at two more + abscissae, increasing the effective order of the approximation by two. + However, in later function evaluations, the function is evaluated at + fewer frequencies because the corresponding derivative has already + converged to the required tolerance. This saves function evaluations to + improve performance, but it requires the function to accept arguments of + any shape. + + "Vector-valued" functions are unlikely to satisfy this requirement. + For example, consider + + >>> def f(x): + ... return [x, np.sin(3*x), x+np.sin(10*x), np.sin(20*x)*(x-1)**2] + + This integrand is not compatible with `derivative` as written; for instance, + the shape of the output will not be the same as the shape of ``x``. Such a + function *could* be converted to a compatible form with the introduction of + additional parameters, but this would be inconvenient. In such cases, + a simpler solution would be to use `preserve_shape`. + + >>> shapes = [] + >>> def f(x): + ... shapes.append(x.shape) + ... x0, x1, x2, x3 = x + ... return [x0, np.sin(3*x1), x2+np.sin(10*x2), np.sin(20*x3)*(x3-1)**2] + >>> + >>> x = np.zeros(4) + >>> res = derivative(f, x, preserve_shape=True) + >>> shapes + [(4,), (4, 8), (4, 2), (4, 2), (4, 2), (4, 2)] + + Here, the shape of ``x`` is ``(4,)``. With ``preserve_shape=True``, the + function may be called with argument ``x`` of shape ``(4,)`` or ``(4, n)``, + and this is what we observe. + + """ + # TODO (followup): + # - investigate behavior at saddle points + # - multivariate functions? + # - relative steps? + # - show example of `np.vectorize` + + res = _derivative_iv(f, x, args, tolerances, maxiter, order, initial_step, + step_factor, step_direction, preserve_shape, callback) + (func, x, args, atol, rtol, maxiter, order, + h0, fac, hdir, preserve_shape, callback) = res + + # Initialization + # Since f(x) (no step) is not needed for central differences, it may be + # possible to eliminate this function evaluation. However, it's useful for + # input validation and standardization, and everything else is designed to + # reduce function calls, so let's keep it simple. + temp = eim._initialize(func, (x,), args, preserve_shape=preserve_shape) + func, xs, fs, args, shape, dtype, xp = temp + + finfo = xp.finfo(dtype) + atol = finfo.smallest_normal if atol is None else atol + rtol = finfo.eps**0.5 if rtol is None else rtol # keep same as `hessian` + + x, f = xs[0], fs[0] + df = xp.full_like(f, xp.nan) + + # Ideally we'd broadcast the shape of `hdir` in `_elementwise_algo_init`, but + # it's simpler to do it here than to generalize `_elementwise_algo_init` further. + # `hdir` and `x` are already broadcasted in `_derivative_iv`, so we know + # that `hdir` can be broadcasted to the final shape. Same with `h0`. + hdir = xp.broadcast_to(hdir, shape) + hdir = xp.reshape(hdir, (-1,)) + hdir = xp.astype(xp_sign(hdir), dtype) + h0 = xp.broadcast_to(h0, shape) + h0 = xp.reshape(h0, (-1,)) + h0 = xp.astype(h0, dtype) + h0[h0 <= 0] = xp.asarray(xp.nan, dtype=dtype) + + status = xp.full_like(x, eim._EINPROGRESS, dtype=xp.int32) # in progress + nit, nfev = 0, 1 # one function evaluations performed above + # Boolean indices of left, central, right, and (all) one-sided steps + il = hdir < 0 + ic = hdir == 0 + ir = hdir > 0 + io = il | ir + + # Most of these attributes are reasonably obvious, but: + # - `fs` holds all the function values of all active `x`. The zeroth + # axis corresponds with active points `x`, the first axis corresponds + # with the different steps (in the order described in + # `_derivative_weights`). + # - `terms` (which could probably use a better name) is half the `order`, + # which is always even. + work = _RichResult(x=x, df=df, fs=f[:, xp.newaxis], error=xp.nan, h=h0, + df_last=xp.nan, error_last=xp.nan, fac=fac, + atol=atol, rtol=rtol, nit=nit, nfev=nfev, + status=status, dtype=dtype, terms=(order+1)//2, + hdir=hdir, il=il, ic=ic, ir=ir, io=io, + # Store the weights in an object so they can't get compressed + # Using RichResult to allow dot notation, but a dict would work + diff_state=_RichResult(central=[], right=[], fac=None)) + + # This is the correspondence between terms in the `work` object and the + # final result. In this case, the mapping is trivial. Note that `success` + # is prepended automatically. + res_work_pairs = [('status', 'status'), ('df', 'df'), ('error', 'error'), + ('nit', 'nit'), ('nfev', 'nfev'), ('x', 'x')] + + def pre_func_eval(work): + """Determine the abscissae at which the function needs to be evaluated. + + See `_derivative_weights` for a description of the stencil (pattern + of the abscissae). + + In the first iteration, there is only one stored function value in + `work.fs`, `f(x)`, so we need to evaluate at `order` new points. In + subsequent iterations, we evaluate at two new points. Note that + `work.x` is always flattened into a 1D array after broadcasting with + all `args`, so we add a new axis at the end and evaluate all point + in one call to the function. + + For improvement: + - Consider measuring the step size actually taken, since ``(x + h) - x`` + is not identically equal to `h` with floating point arithmetic. + - Adjust the step size automatically if `x` is too big to resolve the + step. + - We could probably save some work if there are no central difference + steps or no one-sided steps. + """ + n = work.terms # half the order + h = work.h[:, xp.newaxis] # step size + c = work.fac # step reduction factor + d = c**0.5 # square root of step reduction factor (one-sided stencil) + # Note - no need to be careful about dtypes until we allocate `x_eval` + + if work.nit == 0: + hc = h / c**xp.arange(n, dtype=work.dtype) + hc = xp.concat((-xp.flip(hc, axis=-1), hc), axis=-1) + else: + hc = xp.concat((-h, h), axis=-1) / c**(n-1) + + if work.nit == 0: + hr = h / d**xp.arange(2*n, dtype=work.dtype) + else: + hr = xp.concat((h, h/d), axis=-1) / c**(n-1) + + n_new = 2*n if work.nit == 0 else 2 # number of new abscissae + x_eval = xp.zeros((work.hdir.shape[0], n_new), dtype=work.dtype) + il, ic, ir = work.il, work.ic, work.ir + x_eval[ir] = work.x[ir][:, xp.newaxis] + hr[ir] + x_eval[ic] = work.x[ic][:, xp.newaxis] + hc[ic] + x_eval[il] = work.x[il][:, xp.newaxis] - hr[il] + return x_eval + + def post_func_eval(x, f, work): + """ Estimate the derivative and error from the function evaluations + + As in `pre_func_eval`: in the first iteration, there is only one stored + function value in `work.fs`, `f(x)`, so we need to add the `order` new + points. In subsequent iterations, we add two new points. The tricky + part is getting the order to match that of the weights, which is + described in `_derivative_weights`. + + For improvement: + - Change the order of the weights (and steps in `pre_func_eval`) to + simplify `work_fc` concatenation and eliminate `fc` concatenation. + - It would be simple to do one-step Richardson extrapolation with `df` + and `df_last` to increase the order of the estimate and/or improve + the error estimate. + - Process the function evaluations in a more numerically favorable + way. For instance, combining the pairs of central difference evals + into a second-order approximation and using Richardson extrapolation + to produce a higher order approximation seemed to retain accuracy up + to very high order. + - Alternatively, we could use `polyfit` like Jacobi. An advantage of + fitting polynomial to more points than necessary is improved noise + tolerance. + """ + n = work.terms + n_new = n if work.nit == 0 else 1 + il, ic, io = work.il, work.ic, work.io + + # Central difference + # `work_fc` is *all* the points at which the function has been evaluated + # `fc` is the points we're using *this iteration* to produce the estimate + work_fc = (f[ic][:, :n_new], work.fs[ic], f[ic][:, -n_new:]) + work_fc = xp.concat(work_fc, axis=-1) + if work.nit == 0: + fc = work_fc + else: + fc = (work_fc[:, :n], work_fc[:, n:n+1], work_fc[:, -n:]) + fc = xp.concat(fc, axis=-1) + + # One-sided difference + work_fo = xp.concat((work.fs[io], f[io]), axis=-1) + if work.nit == 0: + fo = work_fo + else: + fo = xp.concat((work_fo[:, 0:1], work_fo[:, -2*n:]), axis=-1) + + work.fs = xp.zeros((ic.shape[0], work.fs.shape[-1] + 2*n_new), dtype=work.dtype) + work.fs[ic] = work_fc + work.fs[io] = work_fo + + wc, wo = _derivative_weights(work, n, xp) + work.df_last = xp.asarray(work.df, copy=True) + work.df[ic] = fc @ wc / work.h[ic] + work.df[io] = fo @ wo / work.h[io] + work.df[il] *= -1 + + work.h /= work.fac + work.error_last = work.error + # Simple error estimate - the difference in derivative estimates between + # this iteration and the last. This is typically conservative because if + # convergence has begin, the true error is much closer to the difference + # between the current estimate and the *next* error estimate. However, + # we could use Richarson extrapolation to produce an error estimate that + # is one order higher, and take the difference between that and + # `work.df` (which would just be constant factor that depends on `fac`.) + work.error = xp.abs(work.df - work.df_last) + + def check_termination(work): + """Terminate due to convergence, non-finite values, or error increase""" + stop = xp.astype(xp.zeros_like(work.df), xp.bool) + + i = work.error < work.atol + work.rtol*abs(work.df) + work.status[i] = eim._ECONVERGED + stop[i] = True + + if work.nit > 0: + i = ~((xp.isfinite(work.x) & xp.isfinite(work.df)) | stop) + work.df[i], work.status[i] = xp.nan, eim._EVALUEERR + stop[i] = True + + # With infinite precision, there is a step size below which + # all smaller step sizes will reduce the error. But in floating point + # arithmetic, catastrophic cancellation will begin to cause the error + # to increase again. This heuristic tries to avoid step sizes that are + # too small. There may be more theoretically sound approaches for + # detecting a step size that minimizes the total error, but this + # heuristic seems simple and effective. + i = (work.error > work.error_last*10) & ~stop + work.status[i] = _EERRORINCREASE + stop[i] = True + + return stop + + def post_termination_check(work): + return + + def customize_result(res, shape): + return shape + + return eim._loop(work, callback, shape, maxiter, func, args, dtype, + pre_func_eval, post_func_eval, check_termination, + post_termination_check, customize_result, res_work_pairs, + xp, preserve_shape) + + +def _derivative_weights(work, n, xp): + # This produces the weights of the finite difference formula for a given + # stencil. In experiments, use of a second-order central difference formula + # with Richardson extrapolation was more accurate numerically, but it was + # more complicated, and it would have become even more complicated when + # adding support for one-sided differences. However, now that all the + # function evaluation values are stored, they can be processed in whatever + # way is desired to produce the derivative estimate. We leave alternative + # approaches to future work. To be more self-contained, here is the theory + # for deriving the weights below. + # + # Recall that the Taylor expansion of a univariate, scalar-values function + # about a point `x` may be expressed as: + # f(x + h) = f(x) + f'(x)*h + f''(x)/2!*h**2 + O(h**3) + # Suppose we evaluate f(x), f(x+h), and f(x-h). We have: + # f(x) = f(x) + # f(x + h) = f(x) + f'(x)*h + f''(x)/2!*h**2 + O(h**3) + # f(x - h) = f(x) - f'(x)*h + f''(x)/2!*h**2 + O(h**3) + # We can solve for weights `wi` such that: + # w1*f(x) = w1*(f(x)) + # + w2*f(x + h) = w2*(f(x) + f'(x)*h + f''(x)/2!*h**2) + O(h**3) + # + w3*f(x - h) = w3*(f(x) - f'(x)*h + f''(x)/2!*h**2) + O(h**3) + # = 0 + f'(x)*h + 0 + O(h**3) + # Then + # f'(x) ~ (w1*f(x) + w2*f(x+h) + w3*f(x-h))/h + # is a finite difference derivative approximation with error O(h**2), + # and so it is said to be a "second-order" approximation. Under certain + # conditions (e.g. well-behaved function, `h` sufficiently small), the + # error in the approximation will decrease with h**2; that is, if `h` is + # reduced by a factor of 2, the error is reduced by a factor of 4. + # + # By default, we use eighth-order formulae. Our central-difference formula + # uses abscissae: + # x-h/c**3, x-h/c**2, x-h/c, x-h, x, x+h, x+h/c, x+h/c**2, x+h/c**3 + # where `c` is the step factor. (Typically, the step factor is greater than + # one, so the outermost points - as written above - are actually closest to + # `x`.) This "stencil" is chosen so that each iteration, the step can be + # reduced by the factor `c`, and most of the function evaluations can be + # reused with the new step size. For example, in the next iteration, we + # will have: + # x-h/c**4, x-h/c**3, x-h/c**2, x-h/c, x, x+h/c, x+h/c**2, x+h/c**3, x+h/c**4 + # We do not reuse `x-h` and `x+h` for the new derivative estimate. + # While this would increase the order of the formula and thus the + # theoretical convergence rate, it is also less stable numerically. + # (As noted above, there are other ways of processing the values that are + # more stable. Thus, even now we store `f(x-h)` and `f(x+h)` in `work.fs` + # to simplify future development of this sort of improvement.) + # + # The (right) one-sided formula is produced similarly using abscissae + # x, x+h, x+h/d, x+h/d**2, ..., x+h/d**6, x+h/d**7, x+h/d**7 + # where `d` is the square root of `c`. (The left one-sided formula simply + # uses -h.) When the step size is reduced by factor `c = d**2`, we have + # abscissae: + # x, x+h/d**2, x+h/d**3..., x+h/d**8, x+h/d**9, x+h/d**9 + # `d` is chosen as the square root of `c` so that the rate of the step-size + # reduction is the same per iteration as in the central difference case. + # Note that because the central difference formulas are inherently of even + # order, for simplicity, we use only even-order formulas for one-sided + # differences, too. + + # It's possible for the user to specify `fac` in, say, double precision but + # `x` and `args` in single precision. `fac` gets converted to single + # precision, but we should always use double precision for the intermediate + # calculations here to avoid additional error in the weights. + fac = float(work.fac) + + # Note that if the user switches back to floating point precision with + # `x` and `args`, then `fac` will not necessarily equal the (lower + # precision) cached `_derivative_weights.fac`, and the weights will + # need to be recalculated. This could be fixed, but it's late, and of + # low consequence. + + diff_state = work.diff_state + if fac != diff_state.fac: + diff_state.central = [] + diff_state.right = [] + diff_state.fac = fac + + if len(diff_state.central) != 2*n + 1: + # Central difference weights. Consider refactoring this; it could + # probably be more compact. + # Note: Using NumPy here is OK; we convert to xp-type at the end + i = np.arange(-n, n + 1) + p = np.abs(i) - 1. # center point has power `p` -1, but sign `s` is 0 + s = np.sign(i) + + h = s / fac ** p + A = np.vander(h, increasing=True).T + b = np.zeros(2*n + 1) + b[1] = 1 + weights = np.linalg.solve(A, b) + + # Enforce identities to improve accuracy + weights[n] = 0 + for i in range(n): + weights[-i-1] = -weights[i] + + # Cache the weights. We only need to calculate them once unless + # the step factor changes. + diff_state.central = weights + + # One-sided difference weights. The left one-sided weights (with + # negative steps) are simply the negative of the right one-sided + # weights, so no need to compute them separately. + i = np.arange(2*n + 1) + p = i - 1. + s = np.sign(i) + + h = s / np.sqrt(fac) ** p + A = np.vander(h, increasing=True).T + b = np.zeros(2 * n + 1) + b[1] = 1 + weights = np.linalg.solve(A, b) + + diff_state.right = weights + + return (xp.asarray(diff_state.central, dtype=work.dtype), + xp.asarray(diff_state.right, dtype=work.dtype)) + + +def jacobian(f, x, *, tolerances=None, maxiter=10, order=8, initial_step=0.5, + step_factor=2.0, step_direction=0): + r"""Evaluate the Jacobian of a function numerically. + + Parameters + ---------- + f : callable + The function whose Jacobian is desired. The signature must be:: + + f(xi: ndarray) -> ndarray + + where each element of ``xi`` is a finite real. If the function to be + differentiated accepts additional arguments, wrap it (e.g. using + `functools.partial` or ``lambda``) and pass the wrapped callable + into `jacobian`. `f` must not mutate the array ``xi``. See Notes + regarding vectorization and the dimensionality of the input and output. + x : float array_like + Points at which to evaluate the Jacobian. Must have at least one dimension. + See Notes regarding the dimensionality and vectorization. + tolerances : dictionary of floats, optional + Absolute and relative tolerances. Valid keys of the dictionary are: + + - ``atol`` - absolute tolerance on the derivative + - ``rtol`` - relative tolerance on the derivative + + Iteration will stop when ``res.error < atol + rtol * abs(res.df)``. The default + `atol` is the smallest normal number of the appropriate dtype, and + the default `rtol` is the square root of the precision of the + appropriate dtype. + maxiter : int, default: 10 + The maximum number of iterations of the algorithm to perform. See + Notes. + order : int, default: 8 + The (positive integer) order of the finite difference formula to be + used. Odd integers will be rounded up to the next even integer. + initial_step : float array_like, default: 0.5 + The (absolute) initial step size for the finite difference derivative + approximation. Must be broadcastable with `x` and `step_direction`. + step_factor : float, default: 2.0 + The factor by which the step size is *reduced* in each iteration; i.e. + the step size in iteration 1 is ``initial_step/step_factor``. If + ``step_factor < 1``, subsequent steps will be greater than the initial + step; this may be useful if steps smaller than some threshold are + undesirable (e.g. due to subtractive cancellation error). + step_direction : integer array_like + An array representing the direction of the finite difference steps (e.g. + for use when `x` lies near to the boundary of the domain of the function.) + Must be broadcastable with `x` and `initial_step`. + Where 0 (default), central differences are used; where negative (e.g. + -1), steps are non-positive; and where positive (e.g. 1), all steps are + non-negative. + + Returns + ------- + res : _RichResult + An object similar to an instance of `scipy.optimize.OptimizeResult` with the + following attributes. The descriptions are written as though the values will + be scalars; however, if `f` returns an array, the outputs will be + arrays of the same shape. + + success : bool array + ``True`` where the algorithm terminated successfully (status ``0``); + ``False`` otherwise. + status : int array + An integer representing the exit status of the algorithm. + + - ``0`` : The algorithm converged to the specified tolerances. + - ``-1`` : The error estimate increased, so iteration was terminated. + - ``-2`` : The maximum number of iterations was reached. + - ``-3`` : A non-finite value was encountered. + + df : float array + The Jacobian of `f` at `x`, if the algorithm terminated + successfully. + error : float array + An estimate of the error: the magnitude of the difference between + the current estimate of the Jacobian and the estimate in the + previous iteration. + nit : int array + The number of iterations of the algorithm that were performed. + nfev : int array + The number of points at which `f` was evaluated. + + Each element of an attribute is associated with the corresponding + element of `df`. For instance, element ``i`` of `nfev` is the + number of points at which `f` was evaluated for the sake of + computing element ``i`` of `df`. + + See Also + -------- + derivative, hessian + + Notes + ----- + Suppose we wish to evaluate the Jacobian of a function + :math:`f: \mathbf{R}^m \rightarrow \mathbf{R}^n`. Assign to variables + ``m`` and ``n`` the positive integer values of :math:`m` and :math:`n`, + respectively, and let ``...`` represent an arbitrary tuple of integers. + If we wish to evaluate the Jacobian at a single point, then: + + - argument `x` must be an array of shape ``(m,)`` + - argument `f` must be vectorized to accept an array of shape ``(m, ...)``. + The first axis represents the :math:`m` inputs of :math:`f`; the remainder + are for evaluating the function at multiple points in a single call. + - argument `f` must return an array of shape ``(n, ...)``. The first + axis represents the :math:`n` outputs of :math:`f`; the remainder + are for the result of evaluating the function at multiple points. + - attribute ``df`` of the result object will be an array of shape ``(n, m)``, + the Jacobian. + + This function is also vectorized in the sense that the Jacobian can be + evaluated at ``k`` points in a single call. In this case, `x` would be an + array of shape ``(m, k)``, `f` would accept an array of shape + ``(m, k, ...)`` and return an array of shape ``(n, k, ...)``, and the ``df`` + attribute of the result would have shape ``(n, m, k)``. + + Suppose the desired callable ``f_not_vectorized`` is not vectorized; it can + only accept an array of shape ``(m,)``. A simple solution to satisfy the required + interface is to wrap ``f_not_vectorized`` as follows:: + + def f(x): + return np.apply_along_axis(f_not_vectorized, axis=0, arr=x) + + Alternatively, suppose the desired callable ``f_vec_q`` is vectorized, but + only for 2-D arrays of shape ``(m, q)``. To satisfy the required interface, + consider:: + + def f(x): + m, batch = x.shape[0], x.shape[1:] # x.shape is (m, ...) + x = np.reshape(x, (m, -1)) # `-1` is short for q = prod(batch) + res = f_vec_q(x) # pass shape (m, q) to function + n = res.shape[0] + return np.reshape(res, (n,) + batch) # return shape (n, ...) + + Then pass the wrapped callable ``f`` as the first argument of `jacobian`. + + References + ---------- + .. [1] Jacobian matrix and determinant, *Wikipedia*, + https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant + + Examples + -------- + The Rosenbrock function maps from :math:`\mathbf{R}^m \rightarrow \mathbf{R}`; + the SciPy implementation `scipy.optimize.rosen` is vectorized to accept an + array of shape ``(m, p)`` and return an array of shape ``p``. Suppose we wish + to evaluate the Jacobian (AKA the gradient because the function returns a scalar) + at ``[0.5, 0.5, 0.5]``. + + >>> import numpy as np + >>> from scipy.differentiate import jacobian + >>> from scipy.optimize import rosen, rosen_der + >>> m = 3 + >>> x = np.full(m, 0.5) + >>> res = jacobian(rosen, x) + >>> ref = rosen_der(x) # reference value of the gradient + >>> res.df, ref + (array([-51., -1., 50.]), array([-51., -1., 50.])) + + As an example of a function with multiple outputs, consider Example 4 + from [1]_. + + >>> def f(x): + ... x1, x2, x3 = x + ... return [x1, 5*x3, 4*x2**2 - 2*x3, x3*np.sin(x1)] + + The true Jacobian is given by: + + >>> def df(x): + ... x1, x2, x3 = x + ... one = np.ones_like(x1) + ... return [[one, 0*one, 0*one], + ... [0*one, 0*one, 5*one], + ... [0*one, 8*x2, -2*one], + ... [x3*np.cos(x1), 0*one, np.sin(x1)]] + + Evaluate the Jacobian at an arbitrary point. + + >>> rng = np.random.default_rng(389252938452) + >>> x = rng.random(size=3) + >>> res = jacobian(f, x) + >>> ref = df(x) + >>> res.df.shape == (4, 3) + True + >>> np.allclose(res.df, ref) + True + + Evaluate the Jacobian at 10 arbitrary points in a single call. + + >>> x = rng.random(size=(3, 10)) + >>> res = jacobian(f, x) + >>> ref = df(x) + >>> res.df.shape == (4, 3, 10) + True + >>> np.allclose(res.df, ref) + True + + """ + xp = array_namespace(x) + x = xp.asarray(x) + int_dtype = xp.isdtype(x.dtype, 'integral') + x0 = xp.asarray(x, dtype=xp.asarray(1.0).dtype) if int_dtype else x + + if x0.ndim < 1: + message = "Argument `x` must be at least 1-D." + raise ValueError(message) + + m = x0.shape[0] + i = xp.arange(m) + + def wrapped(x): + p = () if x.ndim == x0.ndim else (x.shape[-1],) # number of abscissae + + new_shape = (m, m) + x0.shape[1:] + p + xph = xp.expand_dims(x0, axis=1) + if x.ndim != x0.ndim: + xph = xp.expand_dims(xph, axis=-1) + xph = xp_copy(xp.broadcast_to(xph, new_shape), xp=xp) + xph[i, i] = x + return f(xph) + + res = derivative(wrapped, x, tolerances=tolerances, + maxiter=maxiter, order=order, initial_step=initial_step, + step_factor=step_factor, preserve_shape=True, + step_direction=step_direction) + + del res.x # the user knows `x`, and the way it gets broadcasted is meaningless here + return res + + +def hessian(f, x, *, tolerances=None, maxiter=10, + order=8, initial_step=0.5, step_factor=2.0): + r"""Evaluate the Hessian of a function numerically. + + Parameters + ---------- + f : callable + The function whose Hessian is desired. The signature must be:: + + f(xi: ndarray) -> ndarray + + where each element of ``xi`` is a finite real. If the function to be + differentiated accepts additional arguments, wrap it (e.g. using + `functools.partial` or ``lambda``) and pass the wrapped callable + into `hessian`. `f` must not mutate the array ``xi``. See Notes + regarding vectorization and the dimensionality of the input and output. + x : float array_like + Points at which to evaluate the Hessian. Must have at least one dimension. + See Notes regarding the dimensionality and vectorization. + tolerances : dictionary of floats, optional + Absolute and relative tolerances. Valid keys of the dictionary are: + + - ``atol`` - absolute tolerance on the derivative + - ``rtol`` - relative tolerance on the derivative + + Iteration will stop when ``res.error < atol + rtol * abs(res.df)``. The default + `atol` is the smallest normal number of the appropriate dtype, and + the default `rtol` is the square root of the precision of the + appropriate dtype. + order : int, default: 8 + The (positive integer) order of the finite difference formula to be + used. Odd integers will be rounded up to the next even integer. + initial_step : float, default: 0.5 + The (absolute) initial step size for the finite difference derivative + approximation. + step_factor : float, default: 2.0 + The factor by which the step size is *reduced* in each iteration; i.e. + the step size in iteration 1 is ``initial_step/step_factor``. If + ``step_factor < 1``, subsequent steps will be greater than the initial + step; this may be useful if steps smaller than some threshold are + undesirable (e.g. due to subtractive cancellation error). + maxiter : int, default: 10 + The maximum number of iterations of the algorithm to perform. See + Notes. + + Returns + ------- + res : _RichResult + An object similar to an instance of `scipy.optimize.OptimizeResult` with the + following attributes. The descriptions are written as though the values will + be scalars; however, if `f` returns an array, the outputs will be + arrays of the same shape. + + success : bool array + ``True`` where the algorithm terminated successfully (status ``0``); + ``False`` otherwise. + status : int array + An integer representing the exit status of the algorithm. + + - ``0`` : The algorithm converged to the specified tolerances. + - ``-1`` : The error estimate increased, so iteration was terminated. + - ``-2`` : The maximum number of iterations was reached. + - ``-3`` : A non-finite value was encountered. + + ddf : float array + The Hessian of `f` at `x`, if the algorithm terminated + successfully. + error : float array + An estimate of the error: the magnitude of the difference between + the current estimate of the Hessian and the estimate in the + previous iteration. + nfev : int array + The number of points at which `f` was evaluated. + + Each element of an attribute is associated with the corresponding + element of `ddf`. For instance, element ``[i, j]`` of `nfev` is the + number of points at which `f` was evaluated for the sake of + computing element ``[i, j]`` of `ddf`. + + See Also + -------- + derivative, jacobian + + Notes + ----- + Suppose we wish to evaluate the Hessian of a function + :math:`f: \mathbf{R}^m \rightarrow \mathbf{R}`, and we assign to variable + ``m`` the positive integer value of :math:`m`. If we wish to evaluate + the Hessian at a single point, then: + + - argument `x` must be an array of shape ``(m,)`` + - argument `f` must be vectorized to accept an array of shape + ``(m, ...)``. The first axis represents the :math:`m` inputs of + :math:`f`; the remaining axes indicated by ellipses are for evaluating + the function at several abscissae in a single call. + - argument `f` must return an array of shape ``(...)``. + - attribute ``dff`` of the result object will be an array of shape ``(m, m)``, + the Hessian. + + This function is also vectorized in the sense that the Hessian can be + evaluated at ``k`` points in a single call. In this case, `x` would be an + array of shape ``(m, k)``, `f` would accept an array of shape + ``(m, ...)`` and return an array of shape ``(...)``, and the ``ddf`` + attribute of the result would have shape ``(m, m, k)``. Note that the + axis associated with the ``k`` points is included within the axes + denoted by ``(...)``. + + Currently, `hessian` is implemented by nesting calls to `jacobian`. + All options passed to `hessian` are used for both the inner and outer + calls with one exception: the `rtol` used in the inner `jacobian` call + is tightened by a factor of 100 with the expectation that the inner + error can be ignored. A consequence is that `rtol` should not be set + less than 100 times the precision of the dtype of `x`; a warning is + emitted otherwise. + + References + ---------- + .. [1] Hessian matrix, *Wikipedia*, + https://en.wikipedia.org/wiki/Hessian_matrix + + Examples + -------- + The Rosenbrock function maps from :math:`\mathbf{R}^m \rightarrow \mathbf{R}`; + the SciPy implementation `scipy.optimize.rosen` is vectorized to accept an + array of shape ``(m, ...)`` and return an array of shape ``...``. Suppose we + wish to evaluate the Hessian at ``[0.5, 0.5, 0.5]``. + + >>> import numpy as np + >>> from scipy.differentiate import hessian + >>> from scipy.optimize import rosen, rosen_hess + >>> m = 3 + >>> x = np.full(m, 0.5) + >>> res = hessian(rosen, x) + >>> ref = rosen_hess(x) # reference value of the Hessian + >>> np.allclose(res.ddf, ref) + True + + `hessian` is vectorized to evaluate the Hessian at multiple points + in a single call. + + >>> rng = np.random.default_rng(4589245925010) + >>> x = rng.random((m, 10)) + >>> res = hessian(rosen, x) + >>> ref = [rosen_hess(xi) for xi in x.T] + >>> ref = np.moveaxis(ref, 0, -1) + >>> np.allclose(res.ddf, ref) + True + + """ + # todo: + # - add ability to vectorize over additional parameters (*args?) + # - error estimate stack with inner jacobian (or use legit 2D stencil) + + kwargs = dict(maxiter=maxiter, order=order, initial_step=initial_step, + step_factor=step_factor) + tolerances = {} if tolerances is None else tolerances + atol = tolerances.get('atol', None) + rtol = tolerances.get('rtol', None) + + xp = array_namespace(x) + x = xp.asarray(x) + dtype = x.dtype if not xp.isdtype(x.dtype, 'integral') else xp.asarray(1.).dtype + finfo = xp.finfo(dtype) + rtol = finfo.eps**0.5 if rtol is None else rtol # keep same as `derivative` + + # tighten the inner tolerance to make the inner error negligible + rtol_min = finfo.eps * 100 + message = (f"The specified `{rtol=}`, but error estimates are likely to be " + f"unreliable when `rtol < {rtol_min}`.") + if 0 < rtol < rtol_min: # rtol <= 0 is an error + warnings.warn(message, RuntimeWarning, stacklevel=2) + rtol = rtol_min + + def df(x): + tolerances = dict(rtol=rtol/100, atol=atol) + temp = jacobian(f, x, tolerances=tolerances, **kwargs) + nfev.append(temp.nfev if len(nfev) == 0 else temp.nfev.sum(axis=-1)) + return temp.df + + nfev = [] # track inner function evaluations + res = jacobian(df, x, tolerances=tolerances, **kwargs) # jacobian of jacobian + + nfev = xp.cumulative_sum(xp.stack(nfev), axis=0) + res_nit = xp.astype(res.nit[xp.newaxis, ...], xp.int64) # appease torch + res.nfev = xp_take_along_axis(nfev, res_nit, axis=0)[0] + res.ddf = res.df + del res.df # this is renamed to ddf + del res.nit # this is only the outer-jacobian nit + + return res diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/tests/test_differentiate.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/tests/test_differentiate.py new file mode 100644 index 0000000000000000000000000000000000000000..6e6309d764bb837094f60156955160f599b36284 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/differentiate/tests/test_differentiate.py @@ -0,0 +1,695 @@ +import math +import pytest + +import numpy as np + +from scipy.conftest import array_api_compatible +import scipy._lib._elementwise_iterative_method as eim +from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal, xp_assert_less +from scipy._lib._array_api import is_numpy, is_torch, array_namespace + +from scipy import stats, optimize, special +from scipy.differentiate import derivative, jacobian, hessian +from scipy.differentiate._differentiate import _EERRORINCREASE + + +pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")] + +array_api_strict_skip_reason = 'Array API does not support fancy indexing assignment.' +jax_skip_reason = 'JAX arrays do not support item assignment.' + + +@pytest.mark.skip_xp_backends('array_api_strict', reason=array_api_strict_skip_reason) +@pytest.mark.skip_xp_backends('jax.numpy',reason=jax_skip_reason) +class TestDerivative: + + def f(self, x): + return special.ndtr(x) + + @pytest.mark.parametrize('x', [0.6, np.linspace(-0.05, 1.05, 10)]) + def test_basic(self, x, xp): + # Invert distribution CDF and compare against distribution `ppf` + default_dtype = xp.asarray(1.).dtype + res = derivative(self.f, xp.asarray(x, dtype=default_dtype)) + ref = xp.asarray(stats.norm().pdf(x), dtype=default_dtype) + xp_assert_close(res.df, ref) + # This would be nice, but doesn't always work out. `error` is an + # estimate, not a bound. + if not is_torch(xp): + xp_assert_less(xp.abs(res.df - ref), res.error) + + @pytest.mark.skip_xp_backends(np_only=True) + @pytest.mark.parametrize('case', stats._distr_params.distcont) + def test_accuracy(self, case): + distname, params = case + dist = getattr(stats, distname)(*params) + x = dist.median() + 0.1 + res = derivative(dist.cdf, x) + ref = dist.pdf(x) + xp_assert_close(res.df, ref, atol=1e-10) + + @pytest.mark.parametrize('order', [1, 6]) + @pytest.mark.parametrize('shape', [tuple(), (12,), (3, 4), (3, 2, 2)]) + def test_vectorization(self, order, shape, xp): + # Test for correct functionality, output shapes, and dtypes for various + # input shapes. + x = np.linspace(-0.05, 1.05, 12).reshape(shape) if shape else 0.6 + n = np.size(x) + state = {} + + @np.vectorize + def _derivative_single(x): + return derivative(self.f, x, order=order) + + def f(x, *args, **kwargs): + state['nit'] += 1 + state['feval'] += 1 if (x.size == n or x.ndim <=1) else x.shape[-1] + return self.f(x, *args, **kwargs) + + state['nit'] = -1 + state['feval'] = 0 + + res = derivative(f, xp.asarray(x, dtype=xp.float64), order=order) + refs = _derivative_single(x).ravel() + + ref_x = [ref.x for ref in refs] + xp_assert_close(xp.reshape(res.x, (-1,)), xp.asarray(ref_x)) + + ref_df = [ref.df for ref in refs] + xp_assert_close(xp.reshape(res.df, (-1,)), xp.asarray(ref_df)) + + ref_error = [ref.error for ref in refs] + xp_assert_close(xp.reshape(res.error, (-1,)), xp.asarray(ref_error), + atol=1e-12) + + ref_success = [bool(ref.success) for ref in refs] + xp_assert_equal(xp.reshape(res.success, (-1,)), xp.asarray(ref_success)) + + ref_flag = [np.int32(ref.status) for ref in refs] + xp_assert_equal(xp.reshape(res.status, (-1,)), xp.asarray(ref_flag)) + + ref_nfev = [np.int32(ref.nfev) for ref in refs] + xp_assert_equal(xp.reshape(res.nfev, (-1,)), xp.asarray(ref_nfev)) + if is_numpy(xp): # can't expect other backends to be exactly the same + assert xp.max(res.nfev) == state['feval'] + + ref_nit = [np.int32(ref.nit) for ref in refs] + xp_assert_equal(xp.reshape(res.nit, (-1,)), xp.asarray(ref_nit)) + if is_numpy(xp): # can't expect other backends to be exactly the same + assert xp.max(res.nit) == state['nit'] + + def test_flags(self, xp): + # Test cases that should produce different status flags; show that all + # can be produced simultaneously. + rng = np.random.default_rng(5651219684984213) + def f(xs, js): + f.nit += 1 + funcs = [lambda x: x - 2.5, # converges + lambda x: xp.exp(x)*rng.random(), # error increases + lambda x: xp.exp(x), # reaches maxiter due to order=2 + lambda x: xp.full_like(x, xp.nan)] # stops due to NaN + res = [funcs[int(j)](x) for x, j in zip(xs, xp.reshape(js, (-1,)))] + return xp.stack(res) + f.nit = 0 + + args = (xp.arange(4, dtype=xp.int64),) + res = derivative(f, xp.ones(4, dtype=xp.float64), + tolerances=dict(rtol=1e-14), + order=2, args=args) + + ref_flags = xp.asarray([eim._ECONVERGED, + _EERRORINCREASE, + eim._ECONVERR, + eim._EVALUEERR], dtype=xp.int32) + xp_assert_equal(res.status, ref_flags) + + def test_flags_preserve_shape(self, xp): + # Same test as above but using `preserve_shape` option to simplify. + rng = np.random.default_rng(5651219684984213) + def f(x): + out = [x - 2.5, # converges + xp.exp(x)*rng.random(), # error increases + xp.exp(x), # reaches maxiter due to order=2 + xp.full_like(x, xp.nan)] # stops due to NaN + return xp.stack(out) + + res = derivative(f, xp.asarray(1, dtype=xp.float64), + tolerances=dict(rtol=1e-14), + order=2, preserve_shape=True) + + ref_flags = xp.asarray([eim._ECONVERGED, + _EERRORINCREASE, + eim._ECONVERR, + eim._EVALUEERR], dtype=xp.int32) + xp_assert_equal(res.status, ref_flags) + + def test_preserve_shape(self, xp): + # Test `preserve_shape` option + def f(x): + out = [x, xp.sin(3*x), x+xp.sin(10*x), xp.sin(20*x)*(x-1)**2] + return xp.stack(out) + + x = xp.asarray(0.) + ref = xp.asarray([xp.asarray(1), 3*xp.cos(3*x), 1+10*xp.cos(10*x), + 20*xp.cos(20*x)*(x-1)**2 + 2*xp.sin(20*x)*(x-1)]) + res = derivative(f, x, preserve_shape=True) + xp_assert_close(res.df, ref) + + def test_convergence(self, xp): + # Test that the convergence tolerances behave as expected + x = xp.asarray(1., dtype=xp.float64) + f = special.ndtr + ref = float(stats.norm.pdf(1.)) + tolerances0 = dict(atol=0, rtol=0) + + tolerances = tolerances0.copy() + tolerances['atol'] = 1e-3 + res1 = derivative(f, x, tolerances=tolerances, order=4) + assert abs(res1.df - ref) < 1e-3 + tolerances['atol'] = 1e-6 + res2 = derivative(f, x, tolerances=tolerances, order=4) + assert abs(res2.df - ref) < 1e-6 + assert abs(res2.df - ref) < abs(res1.df - ref) + + tolerances = tolerances0.copy() + tolerances['rtol'] = 1e-3 + res1 = derivative(f, x, tolerances=tolerances, order=4) + assert abs(res1.df - ref) < 1e-3 * ref + tolerances['rtol'] = 1e-6 + res2 = derivative(f, x, tolerances=tolerances, order=4) + assert abs(res2.df - ref) < 1e-6 * ref + assert abs(res2.df - ref) < abs(res1.df - ref) + + def test_step_parameters(self, xp): + # Test that step factors have the expected effect on accuracy + x = xp.asarray(1., dtype=xp.float64) + f = special.ndtr + ref = float(stats.norm.pdf(1.)) + + res1 = derivative(f, x, initial_step=0.5, maxiter=1) + res2 = derivative(f, x, initial_step=0.05, maxiter=1) + assert abs(res2.df - ref) < abs(res1.df - ref) + + res1 = derivative(f, x, step_factor=2, maxiter=1) + res2 = derivative(f, x, step_factor=20, maxiter=1) + assert abs(res2.df - ref) < abs(res1.df - ref) + + # `step_factor` can be less than 1: `initial_step` is the minimum step + kwargs = dict(order=4, maxiter=1, step_direction=0) + res = derivative(f, x, initial_step=0.5, step_factor=0.5, **kwargs) + ref = derivative(f, x, initial_step=1, step_factor=2, **kwargs) + xp_assert_close(res.df, ref.df, rtol=5e-15) + + # This is a similar test for one-sided difference + kwargs = dict(order=2, maxiter=1, step_direction=1) + res = derivative(f, x, initial_step=1, step_factor=2, **kwargs) + ref = derivative(f, x, initial_step=1/np.sqrt(2), step_factor=0.5, **kwargs) + xp_assert_close(res.df, ref.df, rtol=5e-15) + + kwargs['step_direction'] = -1 + res = derivative(f, x, initial_step=1, step_factor=2, **kwargs) + ref = derivative(f, x, initial_step=1/np.sqrt(2), step_factor=0.5, **kwargs) + xp_assert_close(res.df, ref.df, rtol=5e-15) + + def test_step_direction(self, xp): + # test that `step_direction` works as expected + def f(x): + y = xp.exp(x) + y[(x < 0) + (x > 2)] = xp.nan + return y + + x = xp.linspace(0, 2, 10) + step_direction = xp.zeros_like(x) + step_direction[x < 0.6], step_direction[x > 1.4] = 1, -1 + res = derivative(f, x, step_direction=step_direction) + xp_assert_close(res.df, xp.exp(x)) + assert xp.all(res.success) + + def test_vectorized_step_direction_args(self, xp): + # test that `step_direction` and `args` are vectorized properly + def f(x, p): + return x ** p + + def df(x, p): + return p * x ** (p - 1) + + x = xp.reshape(xp.asarray([1, 2, 3, 4]), (-1, 1, 1)) + hdir = xp.reshape(xp.asarray([-1, 0, 1]), (1, -1, 1)) + p = xp.reshape(xp.asarray([2, 3]), (1, 1, -1)) + res = derivative(f, x, step_direction=hdir, args=(p,)) + ref = xp.broadcast_to(df(x, p), res.df.shape) + ref = xp.asarray(ref, dtype=xp.asarray(1.).dtype) + xp_assert_close(res.df, ref) + + def test_initial_step(self, xp): + # Test that `initial_step` works as expected and is vectorized + def f(x): + return xp.exp(x) + + x = xp.asarray(0., dtype=xp.float64) + step_direction = xp.asarray([-1, 0, 1]) + h0 = xp.reshape(xp.logspace(-3, 0, 10), (-1, 1)) + res = derivative(f, x, initial_step=h0, order=2, maxiter=1, + step_direction=step_direction) + err = xp.abs(res.df - f(x)) + + # error should be smaller for smaller step sizes + assert xp.all(err[:-1, ...] < err[1:, ...]) + + # results of vectorized call should match results with + # initial_step taken one at a time + for i in range(h0.shape[0]): + ref = derivative(f, x, initial_step=h0[i, 0], order=2, maxiter=1, + step_direction=step_direction) + xp_assert_close(res.df[i, :], ref.df, rtol=1e-14) + + def test_maxiter_callback(self, xp): + # Test behavior of `maxiter` parameter and `callback` interface + x = xp.asarray(0.612814, dtype=xp.float64) + maxiter = 3 + + def f(x): + res = special.ndtr(x) + return res + + default_order = 8 + res = derivative(f, x, maxiter=maxiter, tolerances=dict(rtol=1e-15)) + assert not xp.any(res.success) + assert xp.all(res.nfev == default_order + 1 + (maxiter - 1)*2) + assert xp.all(res.nit == maxiter) + + def callback(res): + callback.iter += 1 + callback.res = res + assert hasattr(res, 'x') + assert float(res.df) not in callback.dfs + callback.dfs.add(float(res.df)) + assert res.status == eim._EINPROGRESS + if callback.iter == maxiter: + raise StopIteration + callback.iter = -1 # callback called once before first iteration + callback.res = None + callback.dfs = set() + + res2 = derivative(f, x, callback=callback, tolerances=dict(rtol=1e-15)) + # terminating with callback is identical to terminating due to maxiter + # (except for `status`) + for key in res.keys(): + if key == 'status': + assert res[key] == eim._ECONVERR + assert res2[key] == eim._ECALLBACK + else: + assert res2[key] == callback.res[key] == res[key] + + @pytest.mark.parametrize("hdir", (-1, 0, 1)) + @pytest.mark.parametrize("x", (0.65, [0.65, 0.7])) + @pytest.mark.parametrize("dtype", ('float16', 'float32', 'float64')) + def test_dtype(self, hdir, x, dtype, xp): + if dtype == 'float16' and not is_numpy(xp): + pytest.skip('float16 not tested for alternative backends') + + # Test that dtypes are preserved + dtype = getattr(xp, dtype) + x = xp.asarray(x, dtype=dtype) + + def f(x): + assert x.dtype == dtype + return xp.exp(x) + + def callback(res): + assert res.x.dtype == dtype + assert res.df.dtype == dtype + assert res.error.dtype == dtype + + res = derivative(f, x, order=4, step_direction=hdir, callback=callback) + assert res.x.dtype == dtype + assert res.df.dtype == dtype + assert res.error.dtype == dtype + eps = xp.finfo(dtype).eps + # not sure why torch is less accurate here; might be worth investigating + rtol = eps**0.5 * 50 if is_torch(xp) else eps**0.5 + xp_assert_close(res.df, xp.exp(res.x), rtol=rtol) + + def test_input_validation(self, xp): + # Test input validation for appropriate error messages + one = xp.asarray(1) + + message = '`f` must be callable.' + with pytest.raises(ValueError, match=message): + derivative(None, one) + + message = 'Abscissae and function output must be real numbers.' + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, xp.asarray(-4+1j)) + + message = "When `preserve_shape=False`, the shape of the array..." + with pytest.raises(ValueError, match=message): + derivative(lambda x: [1, 2, 3], xp.asarray([-2, -3])) + + message = 'Tolerances and step parameters must be non-negative...' + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, tolerances=dict(atol=-1)) + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, tolerances=dict(rtol='ekki')) + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, step_factor=object()) + + message = '`maxiter` must be a positive integer.' + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, maxiter=1.5) + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, maxiter=0) + + message = '`order` must be a positive integer' + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, order=1.5) + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, order=0) + + message = '`preserve_shape` must be True or False.' + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, preserve_shape='herring') + + message = '`callback` must be callable.' + with pytest.raises(ValueError, match=message): + derivative(lambda x: x, one, callback='shrubbery') + + def test_special_cases(self, xp): + # Test edge cases and other special cases + + # Test that integers are not passed to `f` + # (otherwise this would overflow) + def f(x): + xp_test = array_namespace(x) # needs `isdtype` + assert xp_test.isdtype(x.dtype, 'real floating') + return x ** 99 - 1 + + if not is_torch(xp): # torch defaults to float32 + res = derivative(f, xp.asarray(7), tolerances=dict(rtol=1e-10)) + assert res.success + xp_assert_close(res.df, xp.asarray(99*7.**98)) + + # Test invalid step size and direction + res = derivative(xp.exp, xp.asarray(1), step_direction=xp.nan) + xp_assert_equal(res.df, xp.asarray(xp.nan)) + xp_assert_equal(res.status, xp.asarray(-3, dtype=xp.int32)) + + res = derivative(xp.exp, xp.asarray(1), initial_step=0) + xp_assert_equal(res.df, xp.asarray(xp.nan)) + xp_assert_equal(res.status, xp.asarray(-3, dtype=xp.int32)) + + # Test that if success is achieved in the correct number + # of iterations if function is a polynomial. Ideally, all polynomials + # of order 0-2 would get exact result with 0 refinement iterations, + # all polynomials of order 3-4 would be differentiated exactly after + # 1 iteration, etc. However, it seems that `derivative` needs an + # extra iteration to detect convergence based on the error estimate. + + for n in range(6): + x = xp.asarray(1.5, dtype=xp.float64) + def f(x): + return 2*x**n + + ref = 2*n*x**(n-1) + + res = derivative(f, x, maxiter=1, order=max(1, n)) + xp_assert_close(res.df, ref, rtol=1e-15) + xp_assert_equal(res.error, xp.asarray(xp.nan, dtype=xp.float64)) + + res = derivative(f, x, order=max(1, n)) + assert res.success + assert res.nit == 2 + xp_assert_close(res.df, ref, rtol=1e-15) + + # Test scalar `args` (not in tuple) + def f(x, c): + return c*x - 1 + + res = derivative(f, xp.asarray(2), args=xp.asarray(3)) + xp_assert_close(res.df, xp.asarray(3.)) + + # no need to run a test on multiple backends if it's xfailed + @pytest.mark.skip_xp_backends(np_only=True) + @pytest.mark.xfail + @pytest.mark.parametrize("case", ( # function, evaluation point + (lambda x: (x - 1) ** 3, 1), + (lambda x: np.where(x > 1, (x - 1) ** 5, (x - 1) ** 3), 1) + )) + def test_saddle_gh18811(self, case): + # With default settings, `derivative` will not always converge when + # the true derivative is exactly zero. This tests that specifying a + # (tight) `atol` alleviates the problem. See discussion in gh-18811. + atol = 1e-16 + res = derivative(*case, step_direction=[-1, 0, 1], atol=atol) + assert np.all(res.success) + xp_assert_close(res.df, 0, atol=atol) + + +class JacobianHessianTest: + def test_iv(self, xp): + jh_func = self.jh_func.__func__ + + # Test input validation + message = "Argument `x` must be at least 1-D." + with pytest.raises(ValueError, match=message): + jh_func(xp.sin, 1, tolerances=dict(atol=-1)) + + # Confirm that other parameters are being passed to `derivative`, + # which raises an appropriate error message. + x = xp.ones(3) + func = optimize.rosen + message = 'Tolerances and step parameters must be non-negative scalars.' + with pytest.raises(ValueError, match=message): + jh_func(func, x, tolerances=dict(atol=-1)) + with pytest.raises(ValueError, match=message): + jh_func(func, x, tolerances=dict(rtol=-1)) + with pytest.raises(ValueError, match=message): + jh_func(func, x, step_factor=-1) + + message = '`order` must be a positive integer.' + with pytest.raises(ValueError, match=message): + jh_func(func, x, order=-1) + + message = '`maxiter` must be a positive integer.' + with pytest.raises(ValueError, match=message): + jh_func(func, x, maxiter=-1) + + +@pytest.mark.skip_xp_backends('array_api_strict', reason=array_api_strict_skip_reason) +@pytest.mark.skip_xp_backends('jax.numpy',reason=jax_skip_reason) +class TestJacobian(JacobianHessianTest): + jh_func = jacobian + + # Example functions and Jacobians from Wikipedia: + # https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant#Examples + + def f1(z, xp): + x, y = z + return xp.stack([x ** 2 * y, 5 * x + xp.sin(y)]) + + def df1(z): + x, y = z + return [[2 * x * y, x ** 2], [np.full_like(x, 5), np.cos(y)]] + + f1.mn = 2, 2 # type: ignore[attr-defined] + f1.ref = df1 # type: ignore[attr-defined] + + def f2(z, xp): + r, phi = z + return xp.stack([r * xp.cos(phi), r * xp.sin(phi)]) + + def df2(z): + r, phi = z + return [[np.cos(phi), -r * np.sin(phi)], + [np.sin(phi), r * np.cos(phi)]] + + f2.mn = 2, 2 # type: ignore[attr-defined] + f2.ref = df2 # type: ignore[attr-defined] + + def f3(z, xp): + r, phi, th = z + return xp.stack([r * xp.sin(phi) * xp.cos(th), r * xp.sin(phi) * xp.sin(th), + r * xp.cos(phi)]) + + def df3(z): + r, phi, th = z + return [[np.sin(phi) * np.cos(th), r * np.cos(phi) * np.cos(th), + -r * np.sin(phi) * np.sin(th)], + [np.sin(phi) * np.sin(th), r * np.cos(phi) * np.sin(th), + r * np.sin(phi) * np.cos(th)], + [np.cos(phi), -r * np.sin(phi), np.zeros_like(r)]] + + f3.mn = 3, 3 # type: ignore[attr-defined] + f3.ref = df3 # type: ignore[attr-defined] + + def f4(x, xp): + x1, x2, x3 = x + return xp.stack([x1, 5 * x3, 4 * x2 ** 2 - 2 * x3, x3 * xp.sin(x1)]) + + def df4(x): + x1, x2, x3 = x + one = np.ones_like(x1) + return [[one, 0 * one, 0 * one], + [0 * one, 0 * one, 5 * one], + [0 * one, 8 * x2, -2 * one], + [x3 * np.cos(x1), 0 * one, np.sin(x1)]] + + f4.mn = 3, 4 # type: ignore[attr-defined] + f4.ref = df4 # type: ignore[attr-defined] + + def f5(x, xp): + x1, x2, x3 = x + return xp.stack([5 * x2, 4 * x1 ** 2 - 2 * xp.sin(x2 * x3), x2 * x3]) + + def df5(x): + x1, x2, x3 = x + one = np.ones_like(x1) + return [[0 * one, 5 * one, 0 * one], + [8 * x1, -2 * x3 * np.cos(x2 * x3), -2 * x2 * np.cos(x2 * x3)], + [0 * one, x3, x2]] + + f5.mn = 3, 3 # type: ignore[attr-defined] + f5.ref = df5 # type: ignore[attr-defined] + + def rosen(x, _): return optimize.rosen(x) + rosen.mn = 5, 1 # type: ignore[attr-defined] + rosen.ref = optimize.rosen_der # type: ignore[attr-defined] + + @pytest.mark.parametrize('dtype', ('float32', 'float64')) + @pytest.mark.parametrize('size', [(), (6,), (2, 3)]) + @pytest.mark.parametrize('func', [f1, f2, f3, f4, f5, rosen]) + def test_examples(self, dtype, size, func, xp): + atol = 1e-10 if dtype == 'float64' else 1.99e-3 + dtype = getattr(xp, dtype) + rng = np.random.default_rng(458912319542) + m, n = func.mn + x = rng.random(size=(m,) + size) + res = jacobian(lambda x: func(x , xp), xp.asarray(x, dtype=dtype)) + # convert list of arrays to single array before converting to xp array + ref = xp.asarray(np.asarray(func.ref(x)), dtype=dtype) + xp_assert_close(res.df, ref, atol=atol) + + def test_attrs(self, xp): + # Test attributes of result object + z = xp.asarray([0.5, 0.25]) + + # case in which some elements of the Jacobian are harder + # to calculate than others + def df1(z): + x, y = z + return xp.stack([xp.cos(0.5*x) * xp.cos(y), xp.sin(2*x) * y**2]) + + def df1_0xy(x, y): + return xp.cos(0.5*x) * xp.cos(y) + + def df1_1xy(x, y): + return xp.sin(2*x) * y**2 + + res = jacobian(df1, z, initial_step=10) + if is_numpy(xp): + assert len(np.unique(res.nit)) == 4 + assert len(np.unique(res.nfev)) == 4 + + res00 = jacobian(lambda x: df1_0xy(x, z[1]), z[0:1], initial_step=10) + res01 = jacobian(lambda y: df1_0xy(z[0], y), z[1:2], initial_step=10) + res10 = jacobian(lambda x: df1_1xy(x, z[1]), z[0:1], initial_step=10) + res11 = jacobian(lambda y: df1_1xy(z[0], y), z[1:2], initial_step=10) + ref = optimize.OptimizeResult() + for attr in ['success', 'status', 'df', 'nit', 'nfev']: + ref_attr = xp.asarray([[getattr(res00, attr), getattr(res01, attr)], + [getattr(res10, attr), getattr(res11, attr)]]) + ref[attr] = xp.squeeze(ref_attr) + rtol = 1.5e-5 if res[attr].dtype == xp.float32 else 1.5e-14 + xp_assert_close(res[attr], ref[attr], rtol=rtol) + + def test_step_direction_size(self, xp): + # Check that `step_direction` and `initial_step` can be used to ensure that + # the usable domain of a function is respected. + rng = np.random.default_rng(23892589425245) + b = rng.random(3) + eps = 1e-7 # torch needs wiggle room? + + def f(x): + x[0, x[0] < b[0]] = xp.nan + x[0, x[0] > b[0] + 0.25] = xp.nan + x[1, x[1] > b[1]] = xp.nan + x[1, x[1] < b[1] - 0.1-eps] = xp.nan + return TestJacobian.f5(x, xp) + + dir = [1, -1, 0] + h0 = [0.25, 0.1, 0.5] + atol = {'atol': 1e-8} + res = jacobian(f, xp.asarray(b, dtype=xp.float64), initial_step=h0, + step_direction=dir, tolerances=atol) + ref = xp.asarray(TestJacobian.df5(b), dtype=xp.float64) + xp_assert_close(res.df, ref, atol=1e-8) + assert xp.all(xp.isfinite(ref)) + + +@pytest.mark.skip_xp_backends('array_api_strict', reason=array_api_strict_skip_reason) +@pytest.mark.skip_xp_backends('jax.numpy',reason=jax_skip_reason) +class TestHessian(JacobianHessianTest): + jh_func = hessian + + @pytest.mark.parametrize('shape', [(), (4,), (2, 4)]) + def test_example(self, shape, xp): + rng = np.random.default_rng(458912319542) + m = 3 + x = xp.asarray(rng.random((m,) + shape), dtype=xp.float64) + res = hessian(optimize.rosen, x) + if shape: + x = xp.reshape(x, (m, -1)) + ref = xp.stack([optimize.rosen_hess(xi) for xi in x.T]) + ref = xp.moveaxis(ref, 0, -1) + ref = xp.reshape(ref, (m, m,) + shape) + else: + ref = optimize.rosen_hess(x) + xp_assert_close(res.ddf, ref, atol=1e-8) + + # # Removed symmetry enforcement; consider adding back in as a feature + # # check symmetry + # for key in ['ddf', 'error', 'nfev', 'success', 'status']: + # assert_equal(res[key], np.swapaxes(res[key], 0, 1)) + + def test_float32(self, xp): + rng = np.random.default_rng(458912319542) + x = xp.asarray(rng.random(3), dtype=xp.float32) + res = hessian(optimize.rosen, x) + ref = optimize.rosen_hess(x) + mask = (ref != 0) + xp_assert_close(res.ddf[mask], ref[mask]) + atol = 1e-2 * xp.abs(xp.min(ref[mask])) + xp_assert_close(res.ddf[~mask], ref[~mask], atol=atol) + + def test_nfev(self, xp): + z = xp.asarray([0.5, 0.25]) + xp_test = array_namespace(z) + + def f1(z): + x, y = xp_test.broadcast_arrays(*z) + f1.nfev = f1.nfev + (math.prod(x.shape[2:]) if x.ndim > 2 else 1) + return xp.sin(x) * y ** 3 + f1.nfev = 0 + + + res = hessian(f1, z, initial_step=10) + f1.nfev = 0 + res00 = hessian(lambda x: f1([x[0], z[1]]), z[0:1], initial_step=10) + assert res.nfev[0, 0] == f1.nfev == res00.nfev[0, 0] + + f1.nfev = 0 + res11 = hessian(lambda y: f1([z[0], y[0]]), z[1:2], initial_step=10) + assert res.nfev[1, 1] == f1.nfev == res11.nfev[0, 0] + + # Removed symmetry enforcement; consider adding back in as a feature + # assert_equal(res.nfev, res.nfev.T) # check symmetry + # assert np.unique(res.nfev).size == 3 + + + @pytest.mark.thread_unsafe + @pytest.mark.skip_xp_backends(np_only=True, + reason='Python list input uses NumPy backend') + def test_small_rtol_warning(self, xp): + message = 'The specified `rtol=1e-15`, but...' + with pytest.warns(RuntimeWarning, match=message): + hessian(xp.sin, [1.], tolerances=dict(rtol=1e-15)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_backend.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..94d03a4c825ca3fcf1b1d07cf6da932d4e0664db --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_backend.py @@ -0,0 +1,196 @@ +import scipy._lib.uarray as ua +from . import _basic_backend +from . import _realtransforms_backend +from . import _fftlog_backend + + +class _ScipyBackend: + """The default backend for fft calculations + + Notes + ----- + We use the domain ``numpy.scipy`` rather than ``scipy`` because ``uarray`` + treats the domain as a hierarchy. This means the user can install a single + backend for ``numpy`` and have it implement ``numpy.scipy.fft`` as well. + """ + __ua_domain__ = "numpy.scipy.fft" + + @staticmethod + def __ua_function__(method, args, kwargs): + + fn = getattr(_basic_backend, method.__name__, None) + if fn is None: + fn = getattr(_realtransforms_backend, method.__name__, None) + if fn is None: + fn = getattr(_fftlog_backend, method.__name__, None) + if fn is None: + return NotImplemented + return fn(*args, **kwargs) + + +_named_backends = { + 'scipy': _ScipyBackend, +} + + +def _backend_from_arg(backend): + """Maps strings to known backends and validates the backend""" + + if isinstance(backend, str): + try: + backend = _named_backends[backend] + except KeyError as e: + raise ValueError(f'Unknown backend {backend}') from e + + if backend.__ua_domain__ != 'numpy.scipy.fft': + raise ValueError('Backend does not implement "numpy.scipy.fft"') + + return backend + + +def set_global_backend(backend, coerce=False, only=False, try_last=False): + """Sets the global fft backend + + This utility method replaces the default backend for permanent use. It + will be tried in the list of backends automatically, unless the + ``only`` flag is set on a backend. This will be the first tried + backend outside the :obj:`set_backend` context manager. + + Parameters + ---------- + backend : {object, 'scipy'} + The backend to use. + Can either be a ``str`` containing the name of a known backend + {'scipy'} or an object that implements the uarray protocol. + coerce : bool + Whether to coerce input types when trying this backend. + only : bool + If ``True``, no more backends will be tried if this fails. + Implied by ``coerce=True``. + try_last : bool + If ``True``, the global backend is tried after registered backends. + + Raises + ------ + ValueError: If the backend does not implement ``numpy.scipy.fft``. + + Notes + ----- + This will overwrite the previously set global backend, which, by default, is + the SciPy implementation. + + Examples + -------- + We can set the global fft backend: + + >>> from scipy.fft import fft, set_global_backend + >>> set_global_backend("scipy") # Sets global backend (default is "scipy"). + >>> fft([1]) # Calls the global backend + array([1.+0.j]) + """ + backend = _backend_from_arg(backend) + ua.set_global_backend(backend, coerce=coerce, only=only, try_last=try_last) + + +def register_backend(backend): + """ + Register a backend for permanent use. + + Registered backends have the lowest priority and will be tried after the + global backend. + + Parameters + ---------- + backend : {object, 'scipy'} + The backend to use. + Can either be a ``str`` containing the name of a known backend + {'scipy'} or an object that implements the uarray protocol. + + Raises + ------ + ValueError: If the backend does not implement ``numpy.scipy.fft``. + + Examples + -------- + We can register a new fft backend: + + >>> from scipy.fft import fft, register_backend, set_global_backend + >>> class NoopBackend: # Define an invalid Backend + ... __ua_domain__ = "numpy.scipy.fft" + ... def __ua_function__(self, func, args, kwargs): + ... return NotImplemented + >>> set_global_backend(NoopBackend()) # Set the invalid backend as global + >>> register_backend("scipy") # Register a new backend + # The registered backend is called because + # the global backend returns `NotImplemented` + >>> fft([1]) + array([1.+0.j]) + >>> set_global_backend("scipy") # Restore global backend to default + + """ + backend = _backend_from_arg(backend) + ua.register_backend(backend) + + +def set_backend(backend, coerce=False, only=False): + """Context manager to set the backend within a fixed scope. + + Upon entering the ``with`` statement, the given backend will be added to + the list of available backends with the highest priority. Upon exit, the + backend is reset to the state before entering the scope. + + Parameters + ---------- + backend : {object, 'scipy'} + The backend to use. + Can either be a ``str`` containing the name of a known backend + {'scipy'} or an object that implements the uarray protocol. + coerce : bool, optional + Whether to allow expensive conversions for the ``x`` parameter. e.g., + copying a NumPy array to the GPU for a CuPy backend. Implies ``only``. + only : bool, optional + If only is ``True`` and this backend returns ``NotImplemented``, then a + BackendNotImplemented error will be raised immediately. Ignoring any + lower priority backends. + + Examples + -------- + >>> import scipy.fft as fft + >>> with fft.set_backend('scipy', only=True): + ... fft.fft([1]) # Always calls the scipy implementation + array([1.+0.j]) + """ + backend = _backend_from_arg(backend) + return ua.set_backend(backend, coerce=coerce, only=only) + + +def skip_backend(backend): + """Context manager to skip a backend within a fixed scope. + + Within the context of a ``with`` statement, the given backend will not be + called. This covers backends registered both locally and globally. Upon + exit, the backend will again be considered. + + Parameters + ---------- + backend : {object, 'scipy'} + The backend to skip. + Can either be a ``str`` containing the name of a known backend + {'scipy'} or an object that implements the uarray protocol. + + Examples + -------- + >>> import scipy.fft as fft + >>> fft.fft([1]) # Calls default SciPy backend + array([1.+0.j]) + >>> with fft.skip_backend('scipy'): # We explicitly skip the SciPy backend + ... fft.fft([1]) # leaving no implementation available + Traceback (most recent call last): + ... + BackendNotImplementedError: No selected backends had an implementation ... + """ + backend = _backend_from_arg(backend) + return ua.skip_backend(backend) + + +set_global_backend('scipy', try_last=True) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_basic.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..44c61d0a6557381edaebdcc36288d0625e7ab257 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_basic.py @@ -0,0 +1,1630 @@ +from scipy._lib.uarray import generate_multimethod, Dispatchable +import numpy as np + + +def _x_replacer(args, kwargs, dispatchables): + """ + uarray argument replacer to replace the transform input array (``x``) + """ + if len(args) > 0: + return (dispatchables[0],) + args[1:], kwargs + kw = kwargs.copy() + kw['x'] = dispatchables[0] + return args, kw + + +def _dispatch(func): + """ + Function annotation that creates a uarray multimethod from the function + """ + return generate_multimethod(func, _x_replacer, domain="numpy.scipy.fft") + + +@_dispatch +def fft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 1-D discrete Fourier Transform. + + This function computes the 1-D *n*-point discrete Fourier + Transform (DFT) with the efficient Fast Fourier Transform (FFT) + algorithm [1]_. + + Parameters + ---------- + x : array_like + Input array, can be complex. + n : int, optional + Length of the transformed axis of the output. + If `n` is smaller than the length of the input, the input is cropped. + If it is larger, the input is padded with zeros. If `n` is not given, + the length of the input along the axis specified by `axis` is used. + axis : int, optional + Axis over which to compute the FFT. If not given, the last axis is + used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode. Default is "backward", meaning no normalization on + the forward transforms and scaling by ``1/n`` on the `ifft`. + "forward" instead applies the ``1/n`` factor on the forward transform. + For ``norm="ortho"``, both directions are scaled by ``1/sqrt(n)``. + + .. versionadded:: 1.6.0 + ``norm={"forward", "backward"}`` options were added + + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See the notes below for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. See below for more + details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axis + indicated by `axis`, or the last one if `axis` is not specified. + + Raises + ------ + IndexError + if `axes` is larger than the last axis of `x`. + + See Also + -------- + ifft : The inverse of `fft`. + fft2 : The 2-D FFT. + fftn : The N-D FFT. + rfftn : The N-D FFT of real input. + fftfreq : Frequency bins for given FFT parameters. + next_fast_len : Size to pad input to for most efficient transforms + + Notes + ----- + FFT (Fast Fourier Transform) refers to a way the discrete Fourier Transform + (DFT) can be calculated efficiently, by using symmetries in the calculated + terms. The symmetry is highest when `n` is a power of 2, and the transform + is therefore most efficient for these sizes. For poorly factorizable sizes, + `scipy.fft` uses Bluestein's algorithm [2]_ and so is never worse than + O(`n` log `n`). Further performance improvements may be seen by zero-padding + the input using `next_fast_len`. + + If ``x`` is a 1d array, then the `fft` is equivalent to :: + + y[k] = np.sum(x * np.exp(-2j * np.pi * k * np.arange(n)/n)) + + The frequency term ``f=k/n`` is found at ``y[k]``. At ``y[n/2]`` we reach + the Nyquist frequency and wrap around to the negative-frequency terms. So, + for an 8-point transform, the frequencies of the result are + [0, 1, 2, 3, -4, -3, -2, -1]. To rearrange the fft output so that the + zero-frequency component is centered, like [-4, -3, -2, -1, 0, 1, 2, 3], + use `fftshift`. + + Transforms can be done in single, double, or extended precision (long + double) floating point. Half precision inputs will be converted to single + precision and non-floating-point inputs will be converted to double + precision. + + If the data type of ``x`` is real, a "real FFT" algorithm is automatically + used, which roughly halves the computation time. To increase efficiency + a little further, use `rfft`, which does the same calculation, but only + outputs half of the symmetrical spectrum. If the data are both real and + symmetrical, the `dct` can again double the efficiency, by generating + half of the spectrum from half of the signal. + + When ``overwrite_x=True`` is specified, the memory referenced by ``x`` may + be used by the implementation in any way. This may include reusing the + memory for the result, but this is in no way guaranteed. You should not + rely on the contents of ``x`` after the transform as this may change in + future without warning. + + The ``workers`` argument specifies the maximum number of parallel jobs to + split the FFT computation into. This will execute independent 1-D + FFTs within ``x``. So, ``x`` must be at least 2-D and the + non-transformed axes must be large enough to split into chunks. If ``x`` is + too small, fewer jobs may be used than requested. + + References + ---------- + .. [1] Cooley, James W., and John W. Tukey, 1965, "An algorithm for the + machine calculation of complex Fourier series," *Math. Comput.* + 19: 297-301. + .. [2] Bluestein, L., 1970, "A linear filtering approach to the + computation of discrete Fourier transform". *IEEE Transactions on + Audio and Electroacoustics.* 18 (4): 451-455. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> scipy.fft.fft(np.exp(2j * np.pi * np.arange(8) / 8)) + array([-2.33486982e-16+1.14423775e-17j, 8.00000000e+00-1.25557246e-15j, + 2.33486982e-16+2.33486982e-16j, 0.00000000e+00+1.22464680e-16j, + -1.14423775e-17+2.33486982e-16j, 0.00000000e+00+5.20784380e-16j, + 1.14423775e-17+1.14423775e-17j, 0.00000000e+00+1.22464680e-16j]) + + In this example, real input has an FFT which is Hermitian, i.e., symmetric + in the real part and anti-symmetric in the imaginary part: + + >>> from scipy.fft import fft, fftfreq, fftshift + >>> import matplotlib.pyplot as plt + >>> t = np.arange(256) + >>> sp = fftshift(fft(np.sin(t))) + >>> freq = fftshift(fftfreq(t.shape[-1])) + >>> plt.plot(freq, sp.real, freq, sp.imag) + [, + ] + >>> plt.show() + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def ifft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 1-D inverse discrete Fourier Transform. + + This function computes the inverse of the 1-D *n*-point + discrete Fourier transform computed by `fft`. In other words, + ``ifft(fft(x)) == x`` to within numerical accuracy. + + The input should be ordered in the same way as is returned by `fft`, + i.e., + + * ``x[0]`` should contain the zero frequency term, + * ``x[1:n//2]`` should contain the positive-frequency terms, + * ``x[n//2 + 1:]`` should contain the negative-frequency terms, in + increasing order starting from the most negative frequency. + + For an even number of input points, ``x[n//2]`` represents the sum of + the values at the positive and negative Nyquist frequencies, as the two + are aliased together. See `fft` for details. + + Parameters + ---------- + x : array_like + Input array, can be complex. + n : int, optional + Length of the transformed axis of the output. + If `n` is smaller than the length of the input, the input is cropped. + If it is larger, the input is padded with zeros. If `n` is not given, + the length of the input along the axis specified by `axis` is used. + See notes about padding issues. + axis : int, optional + Axis over which to compute the inverse DFT. If not given, the last + axis is used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axis + indicated by `axis`, or the last one if `axis` is not specified. + + Raises + ------ + IndexError + If `axes` is larger than the last axis of `x`. + + See Also + -------- + fft : The 1-D (forward) FFT, of which `ifft` is the inverse. + ifft2 : The 2-D inverse FFT. + ifftn : The N-D inverse FFT. + + Notes + ----- + If the input parameter `n` is larger than the size of the input, the input + is padded by appending zeros at the end. Even though this is the common + approach, it might lead to surprising results. If a different padding is + desired, it must be performed before calling `ifft`. + + If ``x`` is a 1-D array, then the `ifft` is equivalent to :: + + y[k] = np.sum(x * np.exp(2j * np.pi * k * np.arange(n)/n)) / len(x) + + As with `fft`, `ifft` has support for all floating point types and is + optimized for real input. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> scipy.fft.ifft([0, 4, 0, 0]) + array([ 1.+0.j, 0.+1.j, -1.+0.j, 0.-1.j]) # may vary + + Create and plot a band-limited signal with random phases: + + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng() + >>> t = np.arange(400) + >>> n = np.zeros((400,), dtype=complex) + >>> n[40:60] = np.exp(1j*rng.uniform(0, 2*np.pi, (20,))) + >>> s = scipy.fft.ifft(n) + >>> plt.plot(t, s.real, 'b-', t, s.imag, 'r--') + [, ] + >>> plt.legend(('real', 'imaginary')) + + >>> plt.show() + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def rfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 1-D discrete Fourier Transform for real input. + + This function computes the 1-D *n*-point discrete Fourier + Transform (DFT) of a real-valued array by means of an efficient algorithm + called the Fast Fourier Transform (FFT). + + Parameters + ---------- + x : array_like + Input array + n : int, optional + Number of points along transformation axis in the input to use. + If `n` is smaller than the length of the input, the input is cropped. + If it is larger, the input is padded with zeros. If `n` is not given, + the length of the input along the axis specified by `axis` is used. + axis : int, optional + Axis over which to compute the FFT. If not given, the last axis is + used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axis + indicated by `axis`, or the last one if `axis` is not specified. + If `n` is even, the length of the transformed axis is ``(n/2)+1``. + If `n` is odd, the length is ``(n+1)/2``. + + Raises + ------ + IndexError + If `axis` is larger than the last axis of `a`. + + See Also + -------- + irfft : The inverse of `rfft`. + fft : The 1-D FFT of general (complex) input. + fftn : The N-D FFT. + rfft2 : The 2-D FFT of real input. + rfftn : The N-D FFT of real input. + + Notes + ----- + When the DFT is computed for purely real input, the output is + Hermitian-symmetric, i.e., the negative frequency terms are just the complex + conjugates of the corresponding positive-frequency terms, and the + negative-frequency terms are therefore redundant. This function does not + compute the negative frequency terms, and the length of the transformed + axis of the output is therefore ``n//2 + 1``. + + When ``X = rfft(x)`` and fs is the sampling frequency, ``X[0]`` contains + the zero-frequency term 0*fs, which is real due to Hermitian symmetry. + + If `n` is even, ``A[-1]`` contains the term representing both positive + and negative Nyquist frequency (+fs/2 and -fs/2), and must also be purely + real. If `n` is odd, there is no term at fs/2; ``A[-1]`` contains + the largest positive frequency (fs/2*(n-1)/n), and is complex in the + general case. + + If the input `a` contains an imaginary part, it is silently discarded. + + Examples + -------- + >>> import scipy.fft + >>> scipy.fft.fft([0, 1, 0, 0]) + array([ 1.+0.j, 0.-1.j, -1.+0.j, 0.+1.j]) # may vary + >>> scipy.fft.rfft([0, 1, 0, 0]) + array([ 1.+0.j, 0.-1.j, -1.+0.j]) # may vary + + Notice how the final element of the `fft` output is the complex conjugate + of the second element, for real input. For `rfft`, this symmetry is + exploited to compute only the non-negative frequency terms. + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def irfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Computes the inverse of `rfft`. + + This function computes the inverse of the 1-D *n*-point + discrete Fourier Transform of real input computed by `rfft`. + In other words, ``irfft(rfft(x), len(x)) == x`` to within numerical + accuracy. (See Notes below for why ``len(a)`` is necessary here.) + + The input is expected to be in the form returned by `rfft`, i.e., the + real zero-frequency term followed by the complex positive frequency terms + in order of increasing frequency. Since the discrete Fourier Transform of + real input is Hermitian-symmetric, the negative frequency terms are taken + to be the complex conjugates of the corresponding positive frequency terms. + + Parameters + ---------- + x : array_like + The input array. + n : int, optional + Length of the transformed axis of the output. + For `n` output points, ``n//2+1`` input points are necessary. If the + input is longer than this, it is cropped. If it is shorter than this, + it is padded with zeros. If `n` is not given, it is taken to be + ``2*(m-1)``, where ``m`` is the length of the input along the axis + specified by `axis`. + axis : int, optional + Axis over which to compute the inverse FFT. If not given, the last + axis is used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The truncated or zero-padded input, transformed along the axis + indicated by `axis`, or the last one if `axis` is not specified. + The length of the transformed axis is `n`, or, if `n` is not given, + ``2*(m-1)`` where ``m`` is the length of the transformed axis of the + input. To get an odd number of output points, `n` must be specified. + + Raises + ------ + IndexError + If `axis` is larger than the last axis of `x`. + + See Also + -------- + rfft : The 1-D FFT of real input, of which `irfft` is inverse. + fft : The 1-D FFT. + irfft2 : The inverse of the 2-D FFT of real input. + irfftn : The inverse of the N-D FFT of real input. + + Notes + ----- + Returns the real valued `n`-point inverse discrete Fourier transform + of `x`, where `x` contains the non-negative frequency terms of a + Hermitian-symmetric sequence. `n` is the length of the result, not the + input. + + If you specify an `n` such that `a` must be zero-padded or truncated, the + extra/removed values will be added/removed at high frequencies. One can + thus resample a series to `m` points via Fourier interpolation by: + ``a_resamp = irfft(rfft(a), m)``. + + The default value of `n` assumes an even output length. By the Hermitian + symmetry, the last imaginary component must be 0 and so is ignored. To + avoid losing information, the correct length of the real input *must* be + given. + + Examples + -------- + >>> import scipy.fft + >>> scipy.fft.ifft([1, -1j, -1, 1j]) + array([0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j]) # may vary + >>> scipy.fft.irfft([1, -1j, -1]) + array([0., 1., 0., 0.]) + + Notice how the last term in the input to the ordinary `ifft` is the + complex conjugate of the second term, and the output has zero imaginary + part everywhere. When calling `irfft`, the negative frequencies are not + specified, and the output array is purely real. + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def hfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the FFT of a signal that has Hermitian symmetry, i.e., a real + spectrum. + + Parameters + ---------- + x : array_like + The input array. + n : int, optional + Length of the transformed axis of the output. For `n` output + points, ``n//2 + 1`` input points are necessary. If the input is + longer than this, it is cropped. If it is shorter than this, it is + padded with zeros. If `n` is not given, it is taken to be ``2*(m-1)``, + where ``m`` is the length of the input along the axis specified by + `axis`. + axis : int, optional + Axis over which to compute the FFT. If not given, the last + axis is used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See `fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The truncated or zero-padded input, transformed along the axis + indicated by `axis`, or the last one if `axis` is not specified. + The length of the transformed axis is `n`, or, if `n` is not given, + ``2*m - 2``, where ``m`` is the length of the transformed axis of + the input. To get an odd number of output points, `n` must be + specified, for instance, as ``2*m - 1`` in the typical case, + + Raises + ------ + IndexError + If `axis` is larger than the last axis of `a`. + + See Also + -------- + rfft : Compute the 1-D FFT for real input. + ihfft : The inverse of `hfft`. + hfftn : Compute the N-D FFT of a Hermitian signal. + + Notes + ----- + `hfft`/`ihfft` are a pair analogous to `rfft`/`irfft`, but for the + opposite case: here the signal has Hermitian symmetry in the time + domain and is real in the frequency domain. So, here, it's `hfft`, for + which you must supply the length of the result if it is to be odd. + * even: ``ihfft(hfft(a, 2*len(a) - 2) == a``, within roundoff error, + * odd: ``ihfft(hfft(a, 2*len(a) - 1) == a``, within roundoff error. + + Examples + -------- + >>> from scipy.fft import fft, hfft + >>> import numpy as np + >>> a = 2 * np.pi * np.arange(10) / 10 + >>> signal = np.cos(a) + 3j * np.sin(3 * a) + >>> fft(signal).round(10) + array([ -0.+0.j, 5.+0.j, -0.+0.j, 15.-0.j, 0.+0.j, 0.+0.j, + -0.+0.j, -15.-0.j, 0.+0.j, 5.+0.j]) + >>> hfft(signal[:6]).round(10) # Input first half of signal + array([ 0., 5., 0., 15., -0., 0., 0., -15., -0., 5.]) + >>> hfft(signal, 10) # Input entire signal and truncate + array([ 0., 5., 0., 15., -0., 0., 0., -15., -0., 5.]) + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def ihfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the inverse FFT of a signal that has Hermitian symmetry. + + Parameters + ---------- + x : array_like + Input array. + n : int, optional + Length of the inverse FFT, the number of points along + transformation axis in the input to use. If `n` is smaller than + the length of the input, the input is cropped. If it is larger, + the input is padded with zeros. If `n` is not given, the length of + the input along the axis specified by `axis` is used. + axis : int, optional + Axis over which to compute the inverse FFT. If not given, the last + axis is used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See `fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axis + indicated by `axis`, or the last one if `axis` is not specified. + The length of the transformed axis is ``n//2 + 1``. + + See Also + -------- + hfft, irfft + + Notes + ----- + `hfft`/`ihfft` are a pair analogous to `rfft`/`irfft`, but for the + opposite case: here, the signal has Hermitian symmetry in the time + domain and is real in the frequency domain. So, here, it's `hfft`, for + which you must supply the length of the result if it is to be odd: + * even: ``ihfft(hfft(a, 2*len(a) - 2) == a``, within roundoff error, + * odd: ``ihfft(hfft(a, 2*len(a) - 1) == a``, within roundoff error. + + Examples + -------- + >>> from scipy.fft import ifft, ihfft + >>> import numpy as np + >>> spectrum = np.array([ 15, -4, 0, -1, 0, -4]) + >>> ifft(spectrum) + array([1.+0.j, 2.+0.j, 3.+0.j, 4.+0.j, 3.+0.j, 2.+0.j]) # may vary + >>> ihfft(spectrum) + array([ 1.-0.j, 2.-0.j, 3.-0.j, 4.-0.j]) # may vary + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def fftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the N-D discrete Fourier Transform. + + This function computes the N-D discrete Fourier Transform over + any number of axes in an M-D array by means of the Fast Fourier + Transform (FFT). + + Parameters + ---------- + x : array_like + Input array, can be complex. + s : sequence of ints, optional + Shape (length of each transformed axis) of the output + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). + This corresponds to ``n`` for ``fft(x, n)``. + Along any axis, if the given shape is smaller than that of the input, + the input is cropped. If it is larger, the input is padded with zeros. + if `s` is not given, the shape of the input along the axes specified + by `axes` is used. + axes : sequence of ints, optional + Axes over which to compute the FFT. If not given, the last ``len(s)`` + axes are used, or all axes if `s` is also not specified. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or by a combination of `s` and `x`, + as explained in the parameters section above. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + ifftn : The inverse of `fftn`, the inverse N-D FFT. + fft : The 1-D FFT, with definitions and conventions used. + rfftn : The N-D FFT of real input. + fft2 : The 2-D FFT. + fftshift : Shifts zero-frequency terms to centre of array. + + Notes + ----- + The output, analogously to `fft`, contains the term for zero frequency in + the low-order corner of all axes, the positive frequency terms in the + first half of all axes, the term for the Nyquist frequency in the middle + of all axes and the negative frequency terms in the second half of all + axes, in order of decreasingly negative frequency. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.mgrid[:3, :3, :3][0] + >>> scipy.fft.fftn(x, axes=(1, 2)) + array([[[ 0.+0.j, 0.+0.j, 0.+0.j], # may vary + [ 0.+0.j, 0.+0.j, 0.+0.j], + [ 0.+0.j, 0.+0.j, 0.+0.j]], + [[ 9.+0.j, 0.+0.j, 0.+0.j], + [ 0.+0.j, 0.+0.j, 0.+0.j], + [ 0.+0.j, 0.+0.j, 0.+0.j]], + [[18.+0.j, 0.+0.j, 0.+0.j], + [ 0.+0.j, 0.+0.j, 0.+0.j], + [ 0.+0.j, 0.+0.j, 0.+0.j]]]) + >>> scipy.fft.fftn(x, (2, 2), axes=(0, 1)) + array([[[ 2.+0.j, 2.+0.j, 2.+0.j], # may vary + [ 0.+0.j, 0.+0.j, 0.+0.j]], + [[-2.+0.j, -2.+0.j, -2.+0.j], + [ 0.+0.j, 0.+0.j, 0.+0.j]]]) + + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng() + >>> [X, Y] = np.meshgrid(2 * np.pi * np.arange(200) / 12, + ... 2 * np.pi * np.arange(200) / 34) + >>> S = np.sin(X) + np.cos(Y) + rng.uniform(0, 1, X.shape) + >>> FS = scipy.fft.fftn(S) + >>> plt.imshow(np.log(np.abs(scipy.fft.fftshift(FS))**2)) + + >>> plt.show() + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def ifftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the N-D inverse discrete Fourier Transform. + + This function computes the inverse of the N-D discrete + Fourier Transform over any number of axes in an M-D array by + means of the Fast Fourier Transform (FFT). In other words, + ``ifftn(fftn(x)) == x`` to within numerical accuracy. + + The input, analogously to `ifft`, should be ordered in the same way as is + returned by `fftn`, i.e., it should have the term for zero frequency + in all axes in the low-order corner, the positive frequency terms in the + first half of all axes, the term for the Nyquist frequency in the middle + of all axes and the negative frequency terms in the second half of all + axes, in order of decreasingly negative frequency. + + Parameters + ---------- + x : array_like + Input array, can be complex. + s : sequence of ints, optional + Shape (length of each transformed axis) of the output + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). + This corresponds to ``n`` for ``ifft(x, n)``. + Along any axis, if the given shape is smaller than that of the input, + the input is cropped. If it is larger, the input is padded with zeros. + if `s` is not given, the shape of the input along the axes specified + by `axes` is used. See notes for issue on `ifft` zero padding. + axes : sequence of ints, optional + Axes over which to compute the IFFT. If not given, the last ``len(s)`` + axes are used, or all axes if `s` is also not specified. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or by a combination of `s` or `x`, + as explained in the parameters section above. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + fftn : The forward N-D FFT, of which `ifftn` is the inverse. + ifft : The 1-D inverse FFT. + ifft2 : The 2-D inverse FFT. + ifftshift : Undoes `fftshift`, shifts zero-frequency terms to beginning + of array. + + Notes + ----- + Zero-padding, analogously with `ifft`, is performed by appending zeros to + the input along the specified dimension. Although this is the common + approach, it might lead to surprising results. If another form of zero + padding is desired, it must be performed before `ifftn` is called. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.eye(4) + >>> scipy.fft.ifftn(scipy.fft.fftn(x, axes=(0,)), axes=(1,)) + array([[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], # may vary + [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]]) + + + Create and plot an image with band-limited frequency content: + + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng() + >>> n = np.zeros((200,200), dtype=complex) + >>> n[60:80, 20:40] = np.exp(1j*rng.uniform(0, 2*np.pi, (20, 20))) + >>> im = scipy.fft.ifftn(n).real + >>> plt.imshow(im) + + >>> plt.show() + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def fft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 2-D discrete Fourier Transform + + This function computes the N-D discrete Fourier Transform + over any axes in an M-D array by means of the + Fast Fourier Transform (FFT). By default, the transform is computed over + the last two axes of the input array, i.e., a 2-dimensional FFT. + + Parameters + ---------- + x : array_like + Input array, can be complex + s : sequence of ints, optional + Shape (length of each transformed axis) of the output + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). + This corresponds to ``n`` for ``fft(x, n)``. + Along each axis, if the given shape is smaller than that of the input, + the input is cropped. If it is larger, the input is padded with zeros. + if `s` is not given, the shape of the input along the axes specified + by `axes` is used. + axes : sequence of ints, optional + Axes over which to compute the FFT. If not given, the last two axes are + used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or the last two axes if `axes` is not given. + + Raises + ------ + ValueError + If `s` and `axes` have different length, or `axes` not given and + ``len(s) != 2``. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + ifft2 : The inverse 2-D FFT. + fft : The 1-D FFT. + fftn : The N-D FFT. + fftshift : Shifts zero-frequency terms to the center of the array. + For 2-D input, swaps first and third quadrants, and second + and fourth quadrants. + + Notes + ----- + `fft2` is just `fftn` with a different default for `axes`. + + The output, analogously to `fft`, contains the term for zero frequency in + the low-order corner of the transformed axes, the positive frequency terms + in the first half of these axes, the term for the Nyquist frequency in the + middle of the axes and the negative frequency terms in the second half of + the axes, in order of decreasingly negative frequency. + + See `fftn` for details and a plotting example, and `fft` for + definitions and conventions used. + + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.mgrid[:5, :5][0] + >>> scipy.fft.fft2(x) + array([[ 50. +0.j , 0. +0.j , 0. +0.j , # may vary + 0. +0.j , 0. +0.j ], + [-12.5+17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 +4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5 -4.0614962j , 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ], + [-12.5-17.20477401j, 0. +0.j , 0. +0.j , + 0. +0.j , 0. +0.j ]]) + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def ifft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 2-D inverse discrete Fourier Transform. + + This function computes the inverse of the 2-D discrete Fourier + Transform over any number of axes in an M-D array by means of + the Fast Fourier Transform (FFT). In other words, ``ifft2(fft2(x)) == x`` + to within numerical accuracy. By default, the inverse transform is + computed over the last two axes of the input array. + + The input, analogously to `ifft`, should be ordered in the same way as is + returned by `fft2`, i.e., it should have the term for zero frequency + in the low-order corner of the two axes, the positive frequency terms in + the first half of these axes, the term for the Nyquist frequency in the + middle of the axes and the negative frequency terms in the second half of + both axes, in order of decreasingly negative frequency. + + Parameters + ---------- + x : array_like + Input array, can be complex. + s : sequence of ints, optional + Shape (length of each axis) of the output (``s[0]`` refers to axis 0, + ``s[1]`` to axis 1, etc.). This corresponds to `n` for ``ifft(x, n)``. + Along each axis, if the given shape is smaller than that of the input, + the input is cropped. If it is larger, the input is padded with zeros. + if `s` is not given, the shape of the input along the axes specified + by `axes` is used. See notes for issue on `ifft` zero padding. + axes : sequence of ints, optional + Axes over which to compute the FFT. If not given, the last two + axes are used. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or the last two axes if `axes` is not given. + + Raises + ------ + ValueError + If `s` and `axes` have different length, or `axes` not given and + ``len(s) != 2``. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + fft2 : The forward 2-D FFT, of which `ifft2` is the inverse. + ifftn : The inverse of the N-D FFT. + fft : The 1-D FFT. + ifft : The 1-D inverse FFT. + + Notes + ----- + `ifft2` is just `ifftn` with a different default for `axes`. + + See `ifftn` for details and a plotting example, and `fft` for + definition and conventions used. + + Zero-padding, analogously with `ifft`, is performed by appending zeros to + the input along the specified dimension. Although this is the common + approach, it might lead to surprising results. If another form of zero + padding is desired, it must be performed before `ifft2` is called. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = 4 * np.eye(4) + >>> scipy.fft.ifft2(x) + array([[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], # may vary + [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j], + [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j], + [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j]]) + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def rfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the N-D discrete Fourier Transform for real input. + + This function computes the N-D discrete Fourier Transform over + any number of axes in an M-D real array by means of the Fast + Fourier Transform (FFT). By default, all axes are transformed, with the + real transform performed over the last axis, while the remaining + transforms are complex. + + Parameters + ---------- + x : array_like + Input array, taken to be real. + s : sequence of ints, optional + Shape (length along each transformed axis) to use from the input. + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). + The final element of `s` corresponds to `n` for ``rfft(x, n)``, while + for the remaining axes, it corresponds to `n` for ``fft(x, n)``. + Along any axis, if the given shape is smaller than that of the input, + the input is cropped. If it is larger, the input is padded with zeros. + if `s` is not given, the shape of the input along the axes specified + by `axes` is used. + axes : sequence of ints, optional + Axes over which to compute the FFT. If not given, the last ``len(s)`` + axes are used, or all axes if `s` is also not specified. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or by a combination of `s` and `x`, + as explained in the parameters section above. + The length of the last axis transformed will be ``s[-1]//2+1``, + while the remaining transformed axes will have lengths according to + `s`, or unchanged from the input. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + irfftn : The inverse of `rfftn`, i.e., the inverse of the N-D FFT + of real input. + fft : The 1-D FFT, with definitions and conventions used. + rfft : The 1-D FFT of real input. + fftn : The N-D FFT. + rfft2 : The 2-D FFT of real input. + + Notes + ----- + The transform for real input is performed over the last transformation + axis, as by `rfft`, then the transform over the remaining axes is + performed as by `fftn`. The order of the output is as for `rfft` for the + final transformation axis, and as for `fftn` for the remaining + transformation axes. + + See `fft` for details, definitions and conventions used. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.ones((2, 2, 2)) + >>> scipy.fft.rfftn(x) + array([[[8.+0.j, 0.+0.j], # may vary + [0.+0.j, 0.+0.j]], + [[0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j]]]) + + >>> scipy.fft.rfftn(x, axes=(2, 0)) + array([[[4.+0.j, 0.+0.j], # may vary + [4.+0.j, 0.+0.j]], + [[0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j]]]) + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def rfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 2-D FFT of a real array. + + Parameters + ---------- + x : array + Input array, taken to be real. + s : sequence of ints, optional + Shape of the FFT. + axes : sequence of ints, optional + Axes over which to compute the FFT. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The result of the real 2-D FFT. + + See Also + -------- + irfft2 : The inverse of the 2-D FFT of real input. + rfft : The 1-D FFT of real input. + rfftn : Compute the N-D discrete Fourier Transform for real + input. + + Notes + ----- + This is really just `rfftn` with different default behavior. + For more details see `rfftn`. + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def irfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Computes the inverse of `rfftn` + + This function computes the inverse of the N-D discrete + Fourier Transform for real input over any number of axes in an + M-D array by means of the Fast Fourier Transform (FFT). In + other words, ``irfftn(rfftn(x), x.shape) == x`` to within numerical + accuracy. (The ``a.shape`` is necessary like ``len(a)`` is for `irfft`, + and for the same reason.) + + The input should be ordered in the same way as is returned by `rfftn`, + i.e., as for `irfft` for the final transformation axis, and as for `ifftn` + along all the other axes. + + Parameters + ---------- + x : array_like + Input array. + s : sequence of ints, optional + Shape (length of each transformed axis) of the output + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). `s` is also the + number of input points used along this axis, except for the last axis, + where ``s[-1]//2+1`` points of the input are used. + Along any axis, if the shape indicated by `s` is smaller than that of + the input, the input is cropped. If it is larger, the input is padded + with zeros. If `s` is not given, the shape of the input along the axes + specified by axes is used. Except for the last axis which is taken to be + ``2*(m-1)``, where ``m`` is the length of the input along that axis. + axes : sequence of ints, optional + Axes over which to compute the inverse FFT. If not given, the last + `len(s)` axes are used, or all axes if `s` is also not specified. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or by a combination of `s` or `x`, + as explained in the parameters section above. + The length of each transformed axis is as given by the corresponding + element of `s`, or the length of the input in every axis except for the + last one if `s` is not given. In the final transformed axis the length + of the output when `s` is not given is ``2*(m-1)``, where ``m`` is the + length of the final transformed axis of the input. To get an odd + number of output points in the final axis, `s` must be specified. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + rfftn : The forward N-D FFT of real input, + of which `ifftn` is the inverse. + fft : The 1-D FFT, with definitions and conventions used. + irfft : The inverse of the 1-D FFT of real input. + irfft2 : The inverse of the 2-D FFT of real input. + + Notes + ----- + See `fft` for definitions and conventions used. + + See `rfft` for definitions and conventions used for real input. + + The default value of `s` assumes an even output length in the final + transformation axis. When performing the final complex to real + transformation, the Hermitian symmetry requires that the last imaginary + component along that axis must be 0 and so it is ignored. To avoid losing + information, the correct length of the real input *must* be given. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.zeros((3, 2, 2)) + >>> x[0, 0, 0] = 3 * 2 * 2 + >>> scipy.fft.irfftn(x) + array([[[1., 1.], + [1., 1.]], + [[1., 1.], + [1., 1.]], + [[1., 1.], + [1., 1.]]]) + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def irfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Computes the inverse of `rfft2` + + Parameters + ---------- + x : array_like + The input array + s : sequence of ints, optional + Shape of the real output to the inverse FFT. + axes : sequence of ints, optional + The axes over which to compute the inverse fft. + Default is the last two axes. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The result of the inverse real 2-D FFT. + + See Also + -------- + rfft2 : The 2-D FFT of real input. + irfft : The inverse of the 1-D FFT of real input. + irfftn : The inverse of the N-D FFT of real input. + + Notes + ----- + This is really `irfftn` with different defaults. + For more details see `irfftn`. + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def hfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the N-D FFT of Hermitian symmetric complex input, i.e., a + signal with a real spectrum. + + This function computes the N-D discrete Fourier Transform for a + Hermitian symmetric complex input over any number of axes in an + M-D array by means of the Fast Fourier Transform (FFT). In other + words, ``ihfftn(hfftn(x, s)) == x`` to within numerical accuracy. (``s`` + here is ``x.shape`` with ``s[-1] = x.shape[-1] * 2 - 1``, this is necessary + for the same reason ``x.shape`` would be necessary for `irfft`.) + + Parameters + ---------- + x : array_like + Input array. + s : sequence of ints, optional + Shape (length of each transformed axis) of the output + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). `s` is also the + number of input points used along this axis, except for the last axis, + where ``s[-1]//2+1`` points of the input are used. + Along any axis, if the shape indicated by `s` is smaller than that of + the input, the input is cropped. If it is larger, the input is padded + with zeros. If `s` is not given, the shape of the input along the axes + specified by axes is used. Except for the last axis which is taken to be + ``2*(m-1)`` where ``m`` is the length of the input along that axis. + axes : sequence of ints, optional + Axes over which to compute the inverse FFT. If not given, the last + `len(s)` axes are used, or all axes if `s` is also not specified. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or by a combination of `s` or `x`, + as explained in the parameters section above. + The length of each transformed axis is as given by the corresponding + element of `s`, or the length of the input in every axis except for the + last one if `s` is not given. In the final transformed axis the length + of the output when `s` is not given is ``2*(m-1)`` where ``m`` is the + length of the final transformed axis of the input. To get an odd + number of output points in the final axis, `s` must be specified. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + ihfftn : The inverse N-D FFT with real spectrum. Inverse of `hfftn`. + fft : The 1-D FFT, with definitions and conventions used. + rfft : Forward FFT of real input. + + Notes + ----- + For a 1-D signal ``x`` to have a real spectrum, it must satisfy + the Hermitian property:: + + x[i] == np.conj(x[-i]) for all i + + This generalizes into higher dimensions by reflecting over each axis in + turn:: + + x[i, j, k, ...] == np.conj(x[-i, -j, -k, ...]) for all i, j, k, ... + + This should not be confused with a Hermitian matrix, for which the + transpose is its own conjugate:: + + x[i, j] == np.conj(x[j, i]) for all i, j + + + The default value of `s` assumes an even output length in the final + transformation axis. When performing the final complex to real + transformation, the Hermitian symmetry requires that the last imaginary + component along that axis must be 0 and so it is ignored. To avoid losing + information, the correct length of the real input *must* be given. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.ones((3, 2, 2)) + >>> scipy.fft.hfftn(x) + array([[[12., 0.], + [ 0., 0.]], + [[ 0., 0.], + [ 0., 0.]], + [[ 0., 0.], + [ 0., 0.]]]) + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def hfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 2-D FFT of a Hermitian complex array. + + Parameters + ---------- + x : array + Input array, taken to be Hermitian complex. + s : sequence of ints, optional + Shape of the real output. + axes : sequence of ints, optional + Axes over which to compute the FFT. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See `fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The real result of the 2-D Hermitian complex real FFT. + + See Also + -------- + hfftn : Compute the N-D discrete Fourier Transform for Hermitian + complex input. + + Notes + ----- + This is really just `hfftn` with different default behavior. + For more details see `hfftn`. + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def ihfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the N-D inverse discrete Fourier Transform for a real + spectrum. + + This function computes the N-D inverse discrete Fourier Transform + over any number of axes in an M-D real array by means of the Fast + Fourier Transform (FFT). By default, all axes are transformed, with the + real transform performed over the last axis, while the remaining transforms + are complex. + + Parameters + ---------- + x : array_like + Input array, taken to be real. + s : sequence of ints, optional + Shape (length along each transformed axis) to use from the input. + (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.). + Along any axis, if the given shape is smaller than that of the input, + the input is cropped. If it is larger, the input is padded with zeros. + if `s` is not given, the shape of the input along the axes specified + by `axes` is used. + axes : sequence of ints, optional + Axes over which to compute the FFT. If not given, the last ``len(s)`` + axes are used, or all axes if `s` is also not specified. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : complex ndarray + The truncated or zero-padded input, transformed along the axes + indicated by `axes`, or by a combination of `s` and `x`, + as explained in the parameters section above. + The length of the last axis transformed will be ``s[-1]//2+1``, + while the remaining transformed axes will have lengths according to + `s`, or unchanged from the input. + + Raises + ------ + ValueError + If `s` and `axes` have different length. + IndexError + If an element of `axes` is larger than the number of axes of `x`. + + See Also + -------- + hfftn : The forward N-D FFT of Hermitian input. + hfft : The 1-D FFT of Hermitian input. + fft : The 1-D FFT, with definitions and conventions used. + fftn : The N-D FFT. + hfft2 : The 2-D FFT of Hermitian input. + + Notes + ----- + The transform for real input is performed over the last transformation + axis, as by `ihfft`, then the transform over the remaining axes is + performed as by `ifftn`. The order of the output is the positive part of + the Hermitian output signal, in the same format as `rfft`. + + Examples + -------- + >>> import scipy.fft + >>> import numpy as np + >>> x = np.ones((2, 2, 2)) + >>> scipy.fft.ihfftn(x) + array([[[1.+0.j, 0.+0.j], # may vary + [0.+0.j, 0.+0.j]], + [[0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j]]]) + >>> scipy.fft.ihfftn(x, axes=(2, 0)) + array([[[1.+0.j, 0.+0.j], # may vary + [1.+0.j, 0.+0.j]], + [[0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j]]]) + + """ + return (Dispatchable(x, np.ndarray),) + + +@_dispatch +def ihfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None, *, + plan=None): + """ + Compute the 2-D inverse FFT of a real spectrum. + + Parameters + ---------- + x : array_like + The input array + s : sequence of ints, optional + Shape of the real input to the inverse FFT. + axes : sequence of ints, optional + The axes over which to compute the inverse fft. + Default is the last two axes. + norm : {"backward", "ortho", "forward"}, optional + Normalization mode (see `fft`). Default is "backward". + overwrite_x : bool, optional + If True, the contents of `x` can be destroyed; the default is False. + See :func:`fft` for more details. + workers : int, optional + Maximum number of workers to use for parallel computation. If negative, + the value wraps around from ``os.cpu_count()``. + See :func:`~scipy.fft.fft` for more details. + plan : object, optional + This argument is reserved for passing in a precomputed plan provided + by downstream FFT vendors. It is currently not used in SciPy. + + .. versionadded:: 1.5.0 + + Returns + ------- + out : ndarray + The result of the inverse real 2-D FFT. + + See Also + -------- + ihfftn : Compute the inverse of the N-D FFT of Hermitian input. + + Notes + ----- + This is really `ihfftn` with different defaults. + For more details see `ihfftn`. + + """ + return (Dispatchable(x, np.ndarray),) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_basic_backend.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_basic_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..d01e075b4fff2d837ecf51465967f7ff38e15272 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_basic_backend.py @@ -0,0 +1,197 @@ +from scipy._lib._array_api import ( + array_namespace, is_numpy, xp_unsupported_param_msg, is_complex, xp_float_to_complex +) +from . import _pocketfft +import numpy as np + + +def _validate_fft_args(workers, plan, norm): + if workers is not None: + raise ValueError(xp_unsupported_param_msg("workers")) + if plan is not None: + raise ValueError(xp_unsupported_param_msg("plan")) + if norm is None: + norm = 'backward' + return norm + + +# these functions expect complex input in the fft standard extension +complex_funcs = {'fft', 'ifft', 'fftn', 'ifftn', 'hfft', 'irfft', 'irfftn'} + +# pocketfft is used whenever SCIPY_ARRAY_API is not set, +# or x is a NumPy array or array-like. +# When SCIPY_ARRAY_API is set, we try to use xp.fft for CuPy arrays, +# PyTorch arrays and other array API standard supporting objects. +# If xp.fft does not exist, we attempt to convert to np and back to use pocketfft. + +def _execute_1D(func_str, pocketfft_func, x, n, axis, norm, overwrite_x, workers, plan): + xp = array_namespace(x) + + if is_numpy(xp): + x = np.asarray(x) + return pocketfft_func(x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + norm = _validate_fft_args(workers, plan, norm) + if hasattr(xp, 'fft'): + xp_func = getattr(xp.fft, func_str) + if func_str in complex_funcs: + try: + res = xp_func(x, n=n, axis=axis, norm=norm) + except: # backends may require complex input # noqa: E722 + x = xp_float_to_complex(x, xp) + res = xp_func(x, n=n, axis=axis, norm=norm) + return res + return xp_func(x, n=n, axis=axis, norm=norm) + + x = np.asarray(x) + y = pocketfft_func(x, n=n, axis=axis, norm=norm) + return xp.asarray(y) + + +def _execute_nD(func_str, pocketfft_func, x, s, axes, norm, overwrite_x, workers, plan): + xp = array_namespace(x) + + if is_numpy(xp): + x = np.asarray(x) + return pocketfft_func(x, s=s, axes=axes, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + norm = _validate_fft_args(workers, plan, norm) + if hasattr(xp, 'fft'): + xp_func = getattr(xp.fft, func_str) + if func_str in complex_funcs: + try: + res = xp_func(x, s=s, axes=axes, norm=norm) + except: # backends may require complex input # noqa: E722 + x = xp_float_to_complex(x, xp) + res = xp_func(x, s=s, axes=axes, norm=norm) + return res + return xp_func(x, s=s, axes=axes, norm=norm) + + x = np.asarray(x) + y = pocketfft_func(x, s=s, axes=axes, norm=norm) + return xp.asarray(y) + + +def fft(x, n=None, axis=-1, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_1D('fft', _pocketfft.fft, x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def ifft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, + plan=None): + return _execute_1D('ifft', _pocketfft.ifft, x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def rfft(x, n=None, axis=-1, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_1D('rfft', _pocketfft.rfft, x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def irfft(x, n=None, axis=-1, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_1D('irfft', _pocketfft.irfft, x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def hfft(x, n=None, axis=-1, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_1D('hfft', _pocketfft.hfft, x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def ihfft(x, n=None, axis=-1, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_1D('ihfft', _pocketfft.ihfft, x, n=n, axis=axis, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def fftn(x, s=None, axes=None, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_nD('fftn', _pocketfft.fftn, x, s=s, axes=axes, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + + +def ifftn(x, s=None, axes=None, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_nD('ifftn', _pocketfft.ifftn, x, s=s, axes=axes, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def fft2(x, s=None, axes=(-2, -1), norm=None, + overwrite_x=False, workers=None, *, plan=None): + return fftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + + +def ifft2(x, s=None, axes=(-2, -1), norm=None, + overwrite_x=False, workers=None, *, plan=None): + return ifftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + + +def rfftn(x, s=None, axes=None, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_nD('rfftn', _pocketfft.rfftn, x, s=s, axes=axes, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def rfft2(x, s=None, axes=(-2, -1), norm=None, + overwrite_x=False, workers=None, *, plan=None): + return rfftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + + +def irfftn(x, s=None, axes=None, norm=None, + overwrite_x=False, workers=None, *, plan=None): + return _execute_nD('irfftn', _pocketfft.irfftn, x, s=s, axes=axes, norm=norm, + overwrite_x=overwrite_x, workers=workers, plan=plan) + + +def irfft2(x, s=None, axes=(-2, -1), norm=None, + overwrite_x=False, workers=None, *, plan=None): + return irfftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + + +def _swap_direction(norm): + if norm in (None, 'backward'): + norm = 'forward' + elif norm == 'forward': + norm = 'backward' + elif norm != 'ortho': + raise ValueError(f'Invalid norm value {norm}; should be "backward", ' + '"ortho", or "forward".') + return norm + + +def hfftn(x, s=None, axes=None, norm=None, + overwrite_x=False, workers=None, *, plan=None): + xp = array_namespace(x) + if is_numpy(xp): + x = np.asarray(x) + return _pocketfft.hfftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + if is_complex(x, xp): + x = xp.conj(x) + return irfftn(x, s, axes, _swap_direction(norm), + overwrite_x, workers, plan=plan) + + +def hfft2(x, s=None, axes=(-2, -1), norm=None, + overwrite_x=False, workers=None, *, plan=None): + return hfftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + + +def ihfftn(x, s=None, axes=None, norm=None, + overwrite_x=False, workers=None, *, plan=None): + xp = array_namespace(x) + if is_numpy(xp): + x = np.asarray(x) + return _pocketfft.ihfftn(x, s, axes, norm, overwrite_x, workers, plan=plan) + return xp.conj(rfftn(x, s, axes, _swap_direction(norm), + overwrite_x, workers, plan=plan)) + +def ihfft2(x, s=None, axes=(-2, -1), norm=None, + overwrite_x=False, workers=None, *, plan=None): + return ihfftn(x, s, axes, norm, overwrite_x, workers, plan=plan) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_debug_backends.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_debug_backends.py new file mode 100644 index 0000000000000000000000000000000000000000..1a09866e5b4f647c7de3c5292afece0a7a0d3378 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/fft/_debug_backends.py @@ -0,0 +1,22 @@ +import numpy as np + +class NumPyBackend: + """Backend that uses numpy.fft""" + __ua_domain__ = "numpy.scipy.fft" + + @staticmethod + def __ua_function__(method, args, kwargs): + kwargs.pop("overwrite_x", None) + + fn = getattr(np.fft, method.__name__, None) + return (NotImplemented if fn is None + else fn(*args, **kwargs)) + + +class EchoBackend: + """Backend that just prints the __ua_function__ arguments""" + __ua_domain__ = "numpy.scipy.fft" + + @staticmethod + def __ua_function__(method, args, kwargs): + print(method, args, kwargs, sep='\n') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/_scipy_spectral_test_shim.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/_scipy_spectral_test_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..2a7df8990c5123e9a5bcbd58c236fa3618f17c3b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/_scipy_spectral_test_shim.py @@ -0,0 +1,480 @@ +"""Helpers to utilize existing stft / istft tests for testing `ShortTimeFFT`. + +This module provides the functions stft_compare() and istft_compare(), which, +compares the output between the existing (i)stft() and the shortTimeFFT based +_(i)stft_wrapper() implementations in this module. + +For testing add the following imports to the file ``tests/test_spectral.py``:: + + from ._scipy_spectral_test_shim import stft_compare as stft + from ._scipy_spectral_test_shim import istft_compare as istft + +and remove the existing imports of stft and istft. + +The idea of these wrappers is not to provide a backward-compatible interface +but to demonstrate that the ShortTimeFFT implementation is at least as capable +as the existing one and delivers comparable results. Furthermore, the +wrappers highlight the different philosophies of the implementations, +especially in the border handling. +""" +import platform +from typing import cast, Literal + +import numpy as np +from numpy.testing import assert_allclose + +from scipy.signal import ShortTimeFFT +from scipy.signal import csd, get_window, stft, istft +from scipy.signal._arraytools import const_ext, even_ext, odd_ext, zero_ext +from scipy.signal._short_time_fft import FFT_MODE_TYPE +from scipy.signal._spectral_py import _spectral_helper, _triage_segments, \ + _median_bias + + +def _stft_wrapper(x, fs=1.0, window='hann', nperseg=256, noverlap=None, + nfft=None, detrend=False, return_onesided=True, + boundary='zeros', padded=True, axis=-1, scaling='spectrum'): + """Wrapper for the SciPy `stft()` function based on `ShortTimeFFT` for + unit testing. + + Handling the boundary and padding is where `ShortTimeFFT` and `stft()` + differ in behavior. Parts of `_spectral_helper()` were copied to mimic + the` stft()` behavior. + + This function is meant to be solely used by `stft_compare()`. + """ + if scaling not in ('psd', 'spectrum'): # same errors as in original stft: + raise ValueError(f"Parameter {scaling=} not in ['spectrum', 'psd']!") + + # The following lines are taken from the original _spectral_helper(): + boundary_funcs = {'even': even_ext, + 'odd': odd_ext, + 'constant': const_ext, + 'zeros': zero_ext, + None: None} + + if boundary not in boundary_funcs: + raise ValueError(f"Unknown boundary option '{boundary}', must be one" + + f" of: {list(boundary_funcs.keys())}") + if x.size == 0: + return np.empty(x.shape), np.empty(x.shape), np.empty(x.shape) + + if nperseg is not None: # if specified by user + nperseg = int(nperseg) + if nperseg < 1: + raise ValueError('nperseg must be a positive integer') + + # parse window; if array like, then set nperseg = win.shape + win, nperseg = _triage_segments(window, nperseg, + input_length=x.shape[axis]) + + if nfft is None: + nfft = nperseg + elif nfft < nperseg: + raise ValueError('nfft must be greater than or equal to nperseg.') + else: + nfft = int(nfft) + + if noverlap is None: + noverlap = nperseg//2 + else: + noverlap = int(noverlap) + if noverlap >= nperseg: + raise ValueError('noverlap must be less than nperseg.') + nstep = nperseg - noverlap + n = x.shape[axis] + + # Padding occurs after boundary extension, so that the extended signal ends + # in zeros, instead of introducing an impulse at the end. + # I.e. if x = [..., 3, 2] + # extend then pad -> [..., 3, 2, 2, 3, 0, 0, 0] + # pad then extend -> [..., 3, 2, 0, 0, 0, 2, 3] + + if boundary is not None: + ext_func = boundary_funcs[boundary] + # Extend by nperseg//2 in front and back: + x = ext_func(x, nperseg//2, axis=axis) + + if padded: + # Pad to integer number of windowed segments + # I.e make x.shape[-1] = nperseg + (nseg-1)*nstep, with integer nseg + x = np.moveaxis(x, axis, -1) + + # This is an edge case where shortTimeFFT returns one more time slice + # than the Scipy stft() shorten to remove last time slice: + if n % 2 == 1 and nperseg % 2 == 1 and noverlap % 2 == 1: + x = x[..., : -1] + + nadd = (-(x.shape[-1]-nperseg) % nstep) % nperseg + zeros_shape = list(x.shape[:-1]) + [nadd] + x = np.concatenate((x, np.zeros(zeros_shape)), axis=-1) + x = np.moveaxis(x, -1, axis) + + # ... end original _spectral_helper() code. + scale_to = {'spectrum': 'magnitude', 'psd': 'psd'}[scaling] + + if np.iscomplexobj(x) and return_onesided: + return_onesided = False + # using cast() to make mypy happy: + fft_mode = cast(FFT_MODE_TYPE, 'onesided' if return_onesided else 'twosided') + + ST = ShortTimeFFT(win, nstep, fs, fft_mode=fft_mode, mfft=nfft, + scale_to=scale_to, phase_shift=None) + + k_off = nperseg // 2 + p0 = 0 # ST.lower_border_end[1] + 1 + nn = x.shape[axis] if padded else n+k_off+1 + # number of frames akin to legacy stft computation + p1 = (x.shape[axis] - nperseg) // nstep + 1 + + detr = None if detrend is False else detrend + Sxx = ST.stft_detrend(x, detr, p0, p1, k_offset=k_off, axis=axis) + t = ST.t(nn, 0, p1 - p0, k_offset=0 if boundary is not None else k_off) + if x.dtype in (np.float32, np.complex64): + Sxx = Sxx.astype(np.complex64) + + return ST.f, t, Sxx + + +def _istft_wrapper(Zxx, fs=1.0, window='hann', nperseg=None, noverlap=None, + nfft=None, input_onesided=True, boundary=True, time_axis=-1, + freq_axis=-2, scaling='spectrum') -> \ + tuple[np.ndarray, np.ndarray, tuple[int, int]]: + """Wrapper for the SciPy `istft()` function based on `ShortTimeFFT` for + unit testing. + + Note that only option handling is implemented as far as to handle the unit + tests. E.g., the case ``nperseg=None`` is not handled. + + This function is meant to be solely used by `istft_compare()`. + """ + # *** Lines are taken from _spectral_py.istft() ***: + if Zxx.ndim < 2: + raise ValueError('Input stft must be at least 2d!') + + if freq_axis == time_axis: + raise ValueError('Must specify differing time and frequency axes!') + + nseg = Zxx.shape[time_axis] + + if input_onesided: + # Assume even segment length + n_default = 2*(Zxx.shape[freq_axis] - 1) + else: + n_default = Zxx.shape[freq_axis] + + # Check windowing parameters + if nperseg is None: + nperseg = n_default + else: + nperseg = int(nperseg) + if nperseg < 1: + raise ValueError('nperseg must be a positive integer') + + if nfft is None: + if input_onesided and (nperseg == n_default + 1): + # Odd nperseg, no FFT padding + nfft = nperseg + else: + nfft = n_default + elif nfft < nperseg: + raise ValueError('nfft must be greater than or equal to nperseg.') + else: + nfft = int(nfft) + + if noverlap is None: + noverlap = nperseg//2 + else: + noverlap = int(noverlap) + if noverlap >= nperseg: + raise ValueError('noverlap must be less than nperseg.') + nstep = nperseg - noverlap + + # Get window as array + if isinstance(window, str) or type(window) is tuple: + win = get_window(window, nperseg) + else: + win = np.asarray(window) + if len(win.shape) != 1: + raise ValueError('window must be 1-D') + if win.shape[0] != nperseg: + raise ValueError(f'window must have length of {nperseg}') + + outputlength = nperseg + (nseg-1)*nstep + # *** End block of: Taken from _spectral_py.istft() *** + + # Using cast() to make mypy happy: + fft_mode = cast(FFT_MODE_TYPE, 'onesided' if input_onesided else 'twosided') + scale_to = cast(Literal['magnitude', 'psd'], + {'spectrum': 'magnitude', 'psd': 'psd'}[scaling]) + + ST = ShortTimeFFT(win, nstep, fs, fft_mode=fft_mode, mfft=nfft, + scale_to=scale_to, phase_shift=None) + + if boundary: + j = nperseg if nperseg % 2 == 0 else nperseg - 1 + k0 = ST.k_min + nperseg // 2 + k1 = outputlength - j + k0 + else: + raise NotImplementedError("boundary=False does not make sense with" + + "ShortTimeFFT.istft()!") + + x = ST.istft(Zxx, k0=k0, k1=k1, f_axis=freq_axis, t_axis=time_axis) + t = np.arange(k1 - k0) * ST.T + k_hi = ST.upper_border_begin(k1 - k0)[0] + # using cast() to make mypy happy: + return t, x, (ST.lower_border_end[0], k_hi) + + +def _csd_wrapper(x, y, fs=1.0, window='hann', nperseg=None, noverlap=None, + nfft=None, detrend='constant', return_onesided=True, + scaling='density', axis=-1, average='mean'): + """Wrapper for the `csd()` function based on `ShortTimeFFT` for + unit testing. + """ + freqs, _, Pxy = _csd_test_shim(x, y, fs, window, nperseg, noverlap, nfft, + detrend, return_onesided, scaling, axis) + + # The following code is taken from csd(): + if len(Pxy.shape) >= 2 and Pxy.size > 0: + if Pxy.shape[-1] > 1: + if average == 'median': + # np.median must be passed real arrays for the desired result + bias = _median_bias(Pxy.shape[-1]) + if np.iscomplexobj(Pxy): + Pxy = (np.median(np.real(Pxy), axis=-1) + + 1j * np.median(np.imag(Pxy), axis=-1)) + else: + Pxy = np.median(Pxy, axis=-1) + Pxy /= bias + elif average == 'mean': + Pxy = Pxy.mean(axis=-1) + else: + raise ValueError(f'average must be "median" or "mean", got {average}') + else: + Pxy = np.reshape(Pxy, Pxy.shape[:-1]) + + return freqs, Pxy + + +def _csd_test_shim(x, y, fs=1.0, window='hann', nperseg=None, noverlap=None, + nfft=None, detrend='constant', return_onesided=True, + scaling='density', axis=-1): + """Compare output of _spectral_helper() and ShortTimeFFT, more + precisely _spect_helper_csd() for used in csd_wrapper(). + + The motivation of this function is to test if the ShortTimeFFT-based + wrapper `_spect_helper_csd()` returns the same values as `_spectral_helper`. + This function should only be usd by csd() in (unit) testing. + """ + freqs, t, Pxy = _spectral_helper(x, y, fs, window, nperseg, noverlap, nfft, + detrend, return_onesided, scaling, axis, + mode='psd') + freqs1, Pxy1 = _spect_helper_csd(x, y, fs, window, nperseg, noverlap, nfft, + detrend, return_onesided, scaling, axis) + + np.testing.assert_allclose(freqs1, freqs) + amax_Pxy = max(np.abs(Pxy).max(), 1) if Pxy.size else 1 + atol = np.finfo(Pxy.dtype).resolution * amax_Pxy # needed for large Pxy + # for c_ in range(Pxy.shape[-1]): + # np.testing.assert_allclose(Pxy1[:, c_], Pxy[:, c_], atol=atol) + np.testing.assert_allclose(Pxy1, Pxy, atol=atol) + return freqs, t, Pxy + + +def _spect_helper_csd(x, y, fs=1.0, window='hann', nperseg=None, noverlap=None, + nfft=None, detrend='constant', return_onesided=True, + scaling='density', axis=-1): + """Wrapper for replacing _spectral_helper() by using the ShortTimeFFT + for use by csd(). + + This function should be only used by _csd_test_shim() and is only useful + for testing the ShortTimeFFT implementation. + """ + + # The following lines are taken from the original _spectral_helper(): + same_data = y is x + axis = int(axis) + + # Ensure we have np.arrays, get outdtype + x = np.asarray(x) + if not same_data: + y = np.asarray(y) + # outdtype = np.result_type(x, y, np.complex64) + # else: + # outdtype = np.result_type(x, np.complex64) + + if not same_data: + # Check if we can broadcast the outer axes together + xouter = list(x.shape) + youter = list(y.shape) + xouter.pop(axis) + youter.pop(axis) + try: + outershape = np.broadcast(np.empty(xouter), np.empty(youter)).shape + except ValueError as e: + raise ValueError('x and y cannot be broadcast together.') from e + + if same_data: + if x.size == 0: + return np.empty(x.shape), np.empty(x.shape) + else: + if x.size == 0 or y.size == 0: + outshape = outershape + (min([x.shape[axis], y.shape[axis]]),) + emptyout = np.moveaxis(np.empty(outshape), -1, axis) + return emptyout, emptyout + + if nperseg is not None: # if specified by user + nperseg = int(nperseg) + if nperseg < 1: + raise ValueError('nperseg must be a positive integer') + + # parse window; if array like, then set nperseg = win.shape + n = x.shape[axis] if same_data else max(x.shape[axis], y.shape[axis]) + win, nperseg = _triage_segments(window, nperseg, input_length=n) + + if nfft is None: + nfft = nperseg + elif nfft < nperseg: + raise ValueError('nfft must be greater than or equal to nperseg.') + else: + nfft = int(nfft) + + if noverlap is None: + noverlap = nperseg // 2 + else: + noverlap = int(noverlap) + if noverlap >= nperseg: + raise ValueError('noverlap must be less than nperseg.') + nstep = nperseg - noverlap + + if np.iscomplexobj(x) and return_onesided: + return_onesided = False + + # using cast() to make mypy happy: + fft_mode = cast(FFT_MODE_TYPE, 'onesided' if return_onesided + else 'twosided') + scale = {'spectrum': 'magnitude', 'density': 'psd'}[scaling] + SFT = ShortTimeFFT(win, nstep, fs, fft_mode=fft_mode, mfft=nfft, + scale_to=scale, phase_shift=None) + + # _spectral_helper() calculates X.conj()*Y instead of X*Y.conj(): + Pxy = SFT.spectrogram(y, x, detr=None if detrend is False else detrend, + p0=0, p1=(n-noverlap)//SFT.hop, k_offset=nperseg//2, + axis=axis).conj() + # Note: + # 'onesided2X' scaling of ShortTimeFFT conflicts with the + # scaling='spectrum' parameter, since it doubles the squared magnitude, + # which in the view of the ShortTimeFFT implementation does not make sense. + # Hence, the doubling of the square is implemented here: + if return_onesided: + f_axis = Pxy.ndim - 1 + axis if axis < 0 else axis + Pxy = np.moveaxis(Pxy, f_axis, -1) + Pxy[..., 1:-1 if SFT.mfft % 2 == 0 else None] *= 2 + Pxy = np.moveaxis(Pxy, -1, f_axis) + + return SFT.f, Pxy + + +def stft_compare(x, fs=1.0, window='hann', nperseg=256, noverlap=None, + nfft=None, detrend=False, return_onesided=True, + boundary='zeros', padded=True, axis=-1, scaling='spectrum'): + """Assert that the results from the existing `stft()` and `_stft_wrapper()` + are close to each other. + + For comparing the STFT values an absolute tolerance of the floating point + resolution was added to circumvent problems with the following tests: + * For float32 the tolerances are much higher in + TestSTFT.test_roundtrip_float32()). + * The TestSTFT.test_roundtrip_scaling() has a high relative deviation. + Interestingly this did not appear in Scipy 1.9.1 but only in the current + development version. + """ + kw = dict(x=x, fs=fs, window=window, nperseg=nperseg, noverlap=noverlap, + nfft=nfft, detrend=detrend, return_onesided=return_onesided, + boundary=boundary, padded=padded, axis=axis, scaling=scaling) + f, t, Zxx = stft(**kw) + f_wrapper, t_wrapper, Zxx_wrapper = _stft_wrapper(**kw) + + e_msg_part = " of `stft_wrapper()` differ from `stft()`." + assert_allclose(f_wrapper, f, err_msg=f"Frequencies {e_msg_part}") + assert_allclose(t_wrapper, t, err_msg=f"Time slices {e_msg_part}") + + # Adapted tolerances to account for: + atol = np.finfo(Zxx.dtype).resolution * 2 + assert_allclose(Zxx_wrapper, Zxx, atol=atol, + err_msg=f"STFT values {e_msg_part}") + return f, t, Zxx + + +def istft_compare(Zxx, fs=1.0, window='hann', nperseg=None, noverlap=None, + nfft=None, input_onesided=True, boundary=True, time_axis=-1, + freq_axis=-2, scaling='spectrum'): + """Assert that the results from the existing `istft()` and + `_istft_wrapper()` are close to each other. + + Quirks: + * If ``boundary=False`` the comparison is skipped, since it does not + make sense with ShortTimeFFT.istft(). Only used in test + TestSTFT.test_roundtrip_boundary_extension(). + * If ShortTimeFFT.istft() decides the STFT is not invertible, the + comparison is skipped, since istft() only emits a warning and does not + return a correct result. Only used in + ShortTimeFFT.test_roundtrip_not_nola(). + * For comparing the signals an absolute tolerance of the floating point + resolution was added to account for the low accuracy of float32 (Occurs + only in TestSTFT.test_roundtrip_float32()). + """ + kw = dict(Zxx=Zxx, fs=fs, window=window, nperseg=nperseg, + noverlap=noverlap, nfft=nfft, input_onesided=input_onesided, + boundary=boundary, time_axis=time_axis, freq_axis=freq_axis, + scaling=scaling) + + t, x = istft(**kw) + if not boundary: # skip test_roundtrip_boundary_extension(): + return t, x # _istft_wrapper does() not implement this case + try: # if inversion fails, istft() only emits a warning: + t_wrapper, x_wrapper, (k_lo, k_hi) = _istft_wrapper(**kw) + except ValueError as v: # Do nothing if inversion fails: + if v.args[0] == "Short-time Fourier Transform not invertible!": + return t, x + raise v + + e_msg_part = " of `istft_wrapper()` differ from `istft()`" + assert_allclose(t, t_wrapper, err_msg=f"Sample times {e_msg_part}") + + # Adapted tolerances to account for resolution loss: + atol = np.finfo(x.dtype).resolution*2 # instead of default atol = 0 + rtol = 1e-7 # default for np.allclose() + + # Relax atol on 32-Bit platforms a bit to pass CI tests. + # - Not clear why there are discrepancies (in the FFT maybe?) + # - Not sure what changed on 'i686' since earlier on those test passed + if x.dtype == np.float32 and platform.machine() == 'i686': + # float32 gets only used by TestSTFT.test_roundtrip_float32() so + # we are using the tolerances from there to circumvent CI problems + atol, rtol = 1e-4, 1e-5 + elif platform.machine() in ('aarch64', 'i386', 'i686'): + atol = max(atol, 1e-12) # 2e-15 seems too tight for 32-Bit platforms + + assert_allclose(x_wrapper[k_lo:k_hi], x[k_lo:k_hi], atol=atol, rtol=rtol, + err_msg=f"Signal values {e_msg_part}") + return t, x + + +def csd_compare(x, y, fs=1.0, window='hann', nperseg=None, noverlap=None, + nfft=None, detrend='constant', return_onesided=True, + scaling='density', axis=-1, average='mean'): + """Assert that the results from the existing `csd()` and `_csd_wrapper()` + are close to each other. """ + kw = dict(x=x, y=y, fs=fs, window=window, nperseg=nperseg, + noverlap=noverlap, nfft=nfft, detrend=detrend, + return_onesided=return_onesided, scaling=scaling, axis=axis, + average=average) + freqs0, Pxy0 = csd(**kw) + freqs1, Pxy1 = _csd_wrapper(**kw) + + assert_allclose(freqs1, freqs0) + assert_allclose(Pxy1, Pxy0) + assert_allclose(freqs1, freqs0) + return freqs0, Pxy0 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_dltisys.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_dltisys.py new file mode 100644 index 0000000000000000000000000000000000000000..55257a9c032eb08005bca54500d1529552d9d223 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_dltisys.py @@ -0,0 +1,599 @@ +# Author: Jeffrey Armstrong +# April 4, 2011 + +import numpy as np +from numpy.testing import suppress_warnings +from pytest import raises as assert_raises +from scipy._lib._array_api import ( + assert_array_almost_equal, assert_almost_equal, xp_assert_close, xp_assert_equal, +) + +from scipy.signal import (dlsim, dstep, dimpulse, tf2zpk, lti, dlti, + StateSpace, TransferFunction, ZerosPolesGain, + dfreqresp, dbode, BadCoefficients) + + +class TestDLTI: + + def test_dlsim(self): + + a = np.asarray([[0.9, 0.1], [-0.2, 0.9]]) + b = np.asarray([[0.4, 0.1, -0.1], [0.0, 0.05, 0.0]]) + c = np.asarray([[0.1, 0.3]]) + d = np.asarray([[0.0, -0.1, 0.0]]) + dt = 0.5 + + # Create an input matrix with inputs down the columns (3 cols) and its + # respective time input vector + u = np.hstack((np.linspace(0, 4.0, num=5)[:, np.newaxis], + np.full((5, 1), 0.01), + np.full((5, 1), -0.002))) + t_in = np.linspace(0, 2.0, num=5) + + # Define the known result + yout_truth = np.array([[-0.001, + -0.00073, + 0.039446, + 0.0915387, + 0.13195948]]).T + xout_truth = np.asarray([[0, 0], + [0.0012, 0.0005], + [0.40233, 0.00071], + [1.163368, -0.079327], + [2.2402985, -0.3035679]]) + + tout, yout, xout = dlsim((a, b, c, d, dt), u, t_in) + + assert_array_almost_equal(yout_truth, yout) + assert_array_almost_equal(xout_truth, xout) + assert_array_almost_equal(t_in, tout) + + # Make sure input with single-dimension doesn't raise error + dlsim((1, 2, 3), 4) + + # Interpolated control - inputs should have different time steps + # than the discrete model uses internally + u_sparse = u[[0, 4], :] + t_sparse = np.asarray([0.0, 2.0]) + + tout, yout, xout = dlsim((a, b, c, d, dt), u_sparse, t_sparse) + + assert_array_almost_equal(yout_truth, yout) + assert_array_almost_equal(xout_truth, xout) + assert len(tout) == len(yout) + + # Transfer functions (assume dt = 0.5) + num = np.asarray([1.0, -0.1]) + den = np.asarray([0.3, 1.0, 0.2]) + yout_truth = np.array([[0.0, + 0.0, + 3.33333333333333, + -4.77777777777778, + 23.0370370370370]]).T + + # Assume use of the first column of the control input built earlier + tout, yout = dlsim((num, den, 0.5), u[:, 0], t_in) + + assert_array_almost_equal(yout, yout_truth) + assert_array_almost_equal(t_in, tout) + + # Retest the same with a 1-D input vector + uflat = np.asarray(u[:, 0]) + uflat = uflat.reshape((5,)) + tout, yout = dlsim((num, den, 0.5), uflat, t_in) + + assert_array_almost_equal(yout, yout_truth) + assert_array_almost_equal(t_in, tout) + + # zeros-poles-gain representation + zd = np.array([0.5, -0.5]) + pd = np.array([1.j / np.sqrt(2), -1.j / np.sqrt(2)]) + k = 1.0 + yout_truth = np.array([[0.0, 1.0, 2.0, 2.25, 2.5]]).T + + tout, yout = dlsim((zd, pd, k, 0.5), u[:, 0], t_in) + + assert_array_almost_equal(yout, yout_truth) + assert_array_almost_equal(t_in, tout) + + # Raise an error for continuous-time systems + system = lti([1], [1, 1]) + assert_raises(AttributeError, dlsim, system, u) + + def test_dstep(self): + + a = np.asarray([[0.9, 0.1], [-0.2, 0.9]]) + b = np.asarray([[0.4, 0.1, -0.1], [0.0, 0.05, 0.0]]) + c = np.asarray([[0.1, 0.3]]) + d = np.asarray([[0.0, -0.1, 0.0]]) + dt = 0.5 + + # Because b.shape[1] == 3, dstep should result in a tuple of three + # result vectors + yout_step_truth = (np.asarray([0.0, 0.04, 0.052, 0.0404, 0.00956, + -0.036324, -0.093318, -0.15782348, + -0.226628324, -0.2969374948]), + np.asarray([-0.1, -0.075, -0.058, -0.04815, + -0.04453, -0.0461895, -0.0521812, + -0.061588875, -0.073549579, + -0.08727047595]), + np.asarray([0.0, -0.01, -0.013, -0.0101, -0.00239, + 0.009081, 0.0233295, 0.03945587, + 0.056657081, 0.0742343737])) + + tout, yout = dstep((a, b, c, d, dt), n=10) + + assert len(yout) == 3 + + for i in range(0, len(yout)): + assert yout[i].shape[0] == 10 + assert_array_almost_equal(yout[i].flatten(), yout_step_truth[i]) + + # Check that the other two inputs (tf, zpk) will work as well + tfin = ([1.0], [1.0, 1.0], 0.5) + yout_tfstep = np.asarray([0.0, 1.0, 0.0]) + tout, yout = dstep(tfin, n=3) + assert len(yout) == 1 + assert_array_almost_equal(yout[0].flatten(), yout_tfstep) + + zpkin = tf2zpk(tfin[0], tfin[1]) + (0.5,) + tout, yout = dstep(zpkin, n=3) + assert len(yout) == 1 + assert_array_almost_equal(yout[0].flatten(), yout_tfstep) + + # Raise an error for continuous-time systems + system = lti([1], [1, 1]) + assert_raises(AttributeError, dstep, system) + + def test_dimpulse(self): + + a = np.asarray([[0.9, 0.1], [-0.2, 0.9]]) + b = np.asarray([[0.4, 0.1, -0.1], [0.0, 0.05, 0.0]]) + c = np.asarray([[0.1, 0.3]]) + d = np.asarray([[0.0, -0.1, 0.0]]) + dt = 0.5 + + # Because b.shape[1] == 3, dimpulse should result in a tuple of three + # result vectors + yout_imp_truth = (np.asarray([0.0, 0.04, 0.012, -0.0116, -0.03084, + -0.045884, -0.056994, -0.06450548, + -0.068804844, -0.0703091708]), + np.asarray([-0.1, 0.025, 0.017, 0.00985, 0.00362, + -0.0016595, -0.0059917, -0.009407675, + -0.011960704, -0.01372089695]), + np.asarray([0.0, -0.01, -0.003, 0.0029, 0.00771, + 0.011471, 0.0142485, 0.01612637, + 0.017201211, 0.0175772927])) + + tout, yout = dimpulse((a, b, c, d, dt), n=10) + + assert len(yout) == 3 + + for i in range(0, len(yout)): + assert yout[i].shape[0] == 10 + assert_array_almost_equal(yout[i].flatten(), yout_imp_truth[i]) + + # Check that the other two inputs (tf, zpk) will work as well + tfin = ([1.0], [1.0, 1.0], 0.5) + yout_tfimpulse = np.asarray([0.0, 1.0, -1.0]) + tout, yout = dimpulse(tfin, n=3) + assert len(yout) == 1 + assert_array_almost_equal(yout[0].flatten(), yout_tfimpulse) + + zpkin = tf2zpk(tfin[0], tfin[1]) + (0.5,) + tout, yout = dimpulse(zpkin, n=3) + assert len(yout) == 1 + assert_array_almost_equal(yout[0].flatten(), yout_tfimpulse) + + # Raise an error for continuous-time systems + system = lti([1], [1, 1]) + assert_raises(AttributeError, dimpulse, system) + + def test_dlsim_trivial(self): + a = np.array([[0.0]]) + b = np.array([[0.0]]) + c = np.array([[0.0]]) + d = np.array([[0.0]]) + n = 5 + u = np.zeros(n).reshape(-1, 1) + tout, yout, xout = dlsim((a, b, c, d, 1), u) + xp_assert_equal(tout, np.arange(float(n))) + xp_assert_equal(yout, np.zeros((n, 1))) + xp_assert_equal(xout, np.zeros((n, 1))) + + def test_dlsim_simple1d(self): + a = np.array([[0.5]]) + b = np.array([[0.0]]) + c = np.array([[1.0]]) + d = np.array([[0.0]]) + n = 5 + u = np.zeros(n).reshape(-1, 1) + tout, yout, xout = dlsim((a, b, c, d, 1), u, x0=1) + xp_assert_equal(tout, np.arange(float(n))) + expected = (0.5 ** np.arange(float(n))).reshape(-1, 1) + xp_assert_equal(yout, expected) + xp_assert_equal(xout, expected) + + def test_dlsim_simple2d(self): + lambda1 = 0.5 + lambda2 = 0.25 + a = np.array([[lambda1, 0.0], + [0.0, lambda2]]) + b = np.array([[0.0], + [0.0]]) + c = np.array([[1.0, 0.0], + [0.0, 1.0]]) + d = np.array([[0.0], + [0.0]]) + n = 5 + u = np.zeros(n).reshape(-1, 1) + tout, yout, xout = dlsim((a, b, c, d, 1), u, x0=1) + xp_assert_equal(tout, np.arange(float(n))) + # The analytical solution: + expected = (np.array([lambda1, lambda2]) ** + np.arange(float(n)).reshape(-1, 1)) + xp_assert_equal(yout, expected) + xp_assert_equal(xout, expected) + + def test_more_step_and_impulse(self): + lambda1 = 0.5 + lambda2 = 0.75 + a = np.array([[lambda1, 0.0], + [0.0, lambda2]]) + b = np.array([[1.0, 0.0], + [0.0, 1.0]]) + c = np.array([[1.0, 1.0]]) + d = np.array([[0.0, 0.0]]) + + n = 10 + + # Check a step response. + ts, ys = dstep((a, b, c, d, 1), n=n) + + # Create the exact step response. + stp0 = (1.0 / (1 - lambda1)) * (1.0 - lambda1 ** np.arange(n)) + stp1 = (1.0 / (1 - lambda2)) * (1.0 - lambda2 ** np.arange(n)) + + xp_assert_close(ys[0][:, 0], stp0) + xp_assert_close(ys[1][:, 0], stp1) + + # Check an impulse response with an initial condition. + x0 = np.array([1.0, 1.0]) + ti, yi = dimpulse((a, b, c, d, 1), n=n, x0=x0) + + # Create the exact impulse response. + imp = (np.array([lambda1, lambda2]) ** + np.arange(-1, n + 1).reshape(-1, 1)) + imp[0, :] = 0.0 + # Analytical solution to impulse response + y0 = imp[:n, 0] + np.dot(imp[1:n + 1, :], x0) + y1 = imp[:n, 1] + np.dot(imp[1:n + 1, :], x0) + + xp_assert_close(yi[0][:, 0], y0) + xp_assert_close(yi[1][:, 0], y1) + + # Check that dt=0.1, n=3 gives 3 time values. + system = ([1.0], [1.0, -0.5], 0.1) + t, (y,) = dstep(system, n=3) + xp_assert_close(t, [0, 0.1, 0.2]) + xp_assert_equal(y.T, [[0, 1.0, 1.5]]) + t, (y,) = dimpulse(system, n=3) + xp_assert_close(t, [0, 0.1, 0.2]) + xp_assert_equal(y.T, [[0, 1, 0.5]]) + + +class TestDlti: + def test_dlti_instantiation(self): + # Test that lti can be instantiated. + + dt = 0.05 + # TransferFunction + s = dlti([1], [-1], dt=dt) + assert isinstance(s, TransferFunction) + assert isinstance(s, dlti) + assert not isinstance(s, lti) + assert s.dt == dt + + # ZerosPolesGain + s = dlti(np.array([]), np.array([-1]), 1, dt=dt) + assert isinstance(s, ZerosPolesGain) + assert isinstance(s, dlti) + assert not isinstance(s, lti) + assert s.dt == dt + + # StateSpace + s = dlti([1], [-1], 1, 3, dt=dt) + assert isinstance(s, StateSpace) + assert isinstance(s, dlti) + assert not isinstance(s, lti) + assert s.dt == dt + + # Number of inputs + assert_raises(ValueError, dlti, 1) + assert_raises(ValueError, dlti, 1, 1, 1, 1, 1) + + +class TestStateSpaceDisc: + def test_initialization(self): + # Check that all initializations work + dt = 0.05 + StateSpace(1, 1, 1, 1, dt=dt) + StateSpace([1], [2], [3], [4], dt=dt) + StateSpace(np.array([[1, 2], [3, 4]]), np.array([[1], [2]]), + np.array([[1, 0]]), np.array([[0]]), dt=dt) + StateSpace(1, 1, 1, 1, dt=True) + + def test_conversion(self): + # Check the conversion functions + s = StateSpace(1, 2, 3, 4, dt=0.05) + assert isinstance(s.to_ss(), StateSpace) + assert isinstance(s.to_tf(), TransferFunction) + assert isinstance(s.to_zpk(), ZerosPolesGain) + + # Make sure copies work + assert StateSpace(s) is not s + assert s.to_ss() is not s + + def test_properties(self): + # Test setters/getters for cross class properties. + # This implicitly tests to_tf() and to_zpk() + + # Getters + s = StateSpace(1, 1, 1, 1, dt=0.05) + xp_assert_equal(s.poles, [1.]) + xp_assert_equal(s.zeros, [0.]) + + +class TestTransferFunction: + def test_initialization(self): + # Check that all initializations work + dt = 0.05 + TransferFunction(1, 1, dt=dt) + TransferFunction([1], [2], dt=dt) + TransferFunction(np.array([1]), np.array([2]), dt=dt) + TransferFunction(1, 1, dt=True) + + def test_conversion(self): + # Check the conversion functions + s = TransferFunction([1, 0], [1, -1], dt=0.05) + assert isinstance(s.to_ss(), StateSpace) + assert isinstance(s.to_tf(), TransferFunction) + assert isinstance(s.to_zpk(), ZerosPolesGain) + + # Make sure copies work + assert TransferFunction(s) is not s + assert s.to_tf() is not s + + def test_properties(self): + # Test setters/getters for cross class properties. + # This implicitly tests to_ss() and to_zpk() + + # Getters + s = TransferFunction([1, 0], [1, -1], dt=0.05) + xp_assert_equal(s.poles, [1.]) + xp_assert_equal(s.zeros, [0.]) + + +class TestZerosPolesGain: + def test_initialization(self): + # Check that all initializations work + dt = 0.05 + ZerosPolesGain(1, 1, 1, dt=dt) + ZerosPolesGain([1], [2], 1, dt=dt) + ZerosPolesGain(np.array([1]), np.array([2]), 1, dt=dt) + ZerosPolesGain(1, 1, 1, dt=True) + + def test_conversion(self): + # Check the conversion functions + s = ZerosPolesGain(1, 2, 3, dt=0.05) + assert isinstance(s.to_ss(), StateSpace) + assert isinstance(s.to_tf(), TransferFunction) + assert isinstance(s.to_zpk(), ZerosPolesGain) + + # Make sure copies work + assert ZerosPolesGain(s) is not s + assert s.to_zpk() is not s + + +class Test_dfreqresp: + + def test_manual(self): + # Test dfreqresp() real part calculation (manual sanity check). + # 1st order low-pass filter: H(z) = 1 / (z - 0.2), + system = TransferFunction(1, [1, -0.2], dt=0.1) + w = [0.1, 1, 10] + w, H = dfreqresp(system, w=w) + + # test real + expected_re = [1.2383, 0.4130, -0.7553] + assert_almost_equal(H.real, expected_re, decimal=4) + + # test imag + expected_im = [-0.1555, -1.0214, 0.3955] + assert_almost_equal(H.imag, expected_im, decimal=4) + + def test_auto(self): + # Test dfreqresp() real part calculation. + # 1st order low-pass filter: H(z) = 1 / (z - 0.2), + system = TransferFunction(1, [1, -0.2], dt=0.1) + w = [0.1, 1, 10, 100] + w, H = dfreqresp(system, w=w) + jw = np.exp(w * 1j) + y = np.polyval(system.num, jw) / np.polyval(system.den, jw) + + # test real + expected_re = y.real + assert_almost_equal(H.real, expected_re) + + # test imag + expected_im = y.imag + assert_almost_equal(H.imag, expected_im) + + def test_freq_range(self): + # Test that freqresp() finds a reasonable frequency range. + # 1st order low-pass filter: H(z) = 1 / (z - 0.2), + # Expected range is from 0.01 to 10. + system = TransferFunction(1, [1, -0.2], dt=0.1) + n = 10 + expected_w = np.linspace(0, np.pi, 10, endpoint=False) + w, H = dfreqresp(system, n=n) + assert_almost_equal(w, expected_w) + + def test_pole_one(self): + # Test that freqresp() doesn't fail on a system with a pole at 0. + # integrator, pole at zero: H(s) = 1 / s + system = TransferFunction([1], [1, -1], dt=0.1) + + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, message="divide by zero") + sup.filter(RuntimeWarning, message="invalid value encountered") + w, H = dfreqresp(system, n=2) + assert w[0] == 0. # a fail would give not-a-number + + def test_error(self): + # Raise an error for continuous-time systems + system = lti([1], [1, 1]) + assert_raises(AttributeError, dfreqresp, system) + + def test_from_state_space(self): + # H(z) = 2 / z^3 - 0.5 * z^2 + + system_TF = dlti([2], [1, -0.5, 0, 0]) + + A = np.array([[0.5, 0, 0], + [1, 0, 0], + [0, 1, 0]]) + B = np.array([[1, 0, 0]]).T + C = np.array([[0, 0, 2]]) + D = 0 + + system_SS = dlti(A, B, C, D) + w = 10.0**np.arange(-3,0,.5) + with suppress_warnings() as sup: + sup.filter(BadCoefficients) + w1, H1 = dfreqresp(system_TF, w=w) + w2, H2 = dfreqresp(system_SS, w=w) + + assert_almost_equal(H1, H2) + + def test_from_zpk(self): + # 1st order low-pass filter: H(s) = 0.3 / (z - 0.2), + system_ZPK = dlti([],[0.2],0.3) + system_TF = dlti(0.3, [1, -0.2]) + w = [0.1, 1, 10, 100] + w1, H1 = dfreqresp(system_ZPK, w=w) + w2, H2 = dfreqresp(system_TF, w=w) + assert_almost_equal(H1, H2) + + +class Test_bode: + + def test_manual(self): + # Test bode() magnitude calculation (manual sanity check). + # 1st order low-pass filter: H(s) = 0.3 / (z - 0.2), + dt = 0.1 + system = TransferFunction(0.3, [1, -0.2], dt=dt) + w = [0.1, 0.5, 1, np.pi] + w2, mag, phase = dbode(system, w=w) + + # Test mag + expected_mag = [-8.5329, -8.8396, -9.6162, -12.0412] + assert_almost_equal(mag, expected_mag, decimal=4) + + # Test phase + expected_phase = [-7.1575, -35.2814, -67.9809, -180.0000] + assert_almost_equal(phase, expected_phase, decimal=4) + + # Test frequency + xp_assert_equal(np.array(w) / dt, w2) + + def test_auto(self): + # Test bode() magnitude calculation. + # 1st order low-pass filter: H(s) = 0.3 / (z - 0.2), + system = TransferFunction(0.3, [1, -0.2], dt=0.1) + w = np.array([0.1, 0.5, 1, np.pi]) + w2, mag, phase = dbode(system, w=w) + jw = np.exp(w * 1j) + y = np.polyval(system.num, jw) / np.polyval(system.den, jw) + + # Test mag + expected_mag = 20.0 * np.log10(abs(y)) + assert_almost_equal(mag, expected_mag) + + # Test phase + expected_phase = np.rad2deg(np.angle(y)) + assert_almost_equal(phase, expected_phase) + + def test_range(self): + # Test that bode() finds a reasonable frequency range. + # 1st order low-pass filter: H(s) = 0.3 / (z - 0.2), + dt = 0.1 + system = TransferFunction(0.3, [1, -0.2], dt=0.1) + n = 10 + # Expected range is from 0.01 to 10. + expected_w = np.linspace(0, np.pi, n, endpoint=False) / dt + w, mag, phase = dbode(system, n=n) + assert_almost_equal(w, expected_w) + + def test_pole_one(self): + # Test that freqresp() doesn't fail on a system with a pole at 0. + # integrator, pole at zero: H(s) = 1 / s + system = TransferFunction([1], [1, -1], dt=0.1) + + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, message="divide by zero") + sup.filter(RuntimeWarning, message="invalid value encountered") + w, mag, phase = dbode(system, n=2) + assert w[0] == 0. # a fail would give not-a-number + + def test_imaginary(self): + # bode() should not fail on a system with pure imaginary poles. + # The test passes if bode doesn't raise an exception. + system = TransferFunction([1], [1, 0, 100], dt=0.1) + dbode(system, n=2) + + def test_error(self): + # Raise an error for continuous-time systems + system = lti([1], [1, 1]) + assert_raises(AttributeError, dbode, system) + + +class TestTransferFunctionZConversion: + """Test private conversions between 'z' and 'z**-1' polynomials.""" + + def test_full(self): + # Numerator and denominator same order + num = np.asarray([2.0, 3, 4]) + den = np.asarray([5.0, 6, 7]) + num2, den2 = TransferFunction._z_to_zinv(num, den) + xp_assert_equal(num, num2) + xp_assert_equal(den, den2) + + num2, den2 = TransferFunction._zinv_to_z(num, den) + xp_assert_equal(num, num2) + xp_assert_equal(den, den2) + + def test_numerator(self): + # Numerator lower order than denominator + num = np.asarray([2.0, 3]) + den = np.asarray([50, 6, 7]) + num2, den2 = TransferFunction._z_to_zinv(num, den) + xp_assert_equal([0.0, 2, 3], num2) + xp_assert_equal(den, den2) + + num2, den2 = TransferFunction._zinv_to_z(num, den) + xp_assert_equal([2.0, 3, 0], num2) + xp_assert_equal(den, den2) + + def test_denominator(self): + # Numerator higher order than denominator + num = np.asarray([2., 3, 4]) + den = np.asarray([5.0, 6]) + num2, den2 = TransferFunction._z_to_zinv(num, den) + xp_assert_equal(num, num2) + xp_assert_equal([0.0, 5, 6], den2) + + num2, den2 = TransferFunction._zinv_to_z(num, den) + xp_assert_equal(num, num2) + xp_assert_equal([5.0, 6, 0], den2) + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_filter_design.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_filter_design.py new file mode 100644 index 0000000000000000000000000000000000000000..f51975b73b6ab2a796dd790769c578e096561dd1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_filter_design.py @@ -0,0 +1,4485 @@ +import warnings + +from scipy._lib import _pep440 +import numpy as np +from numpy.testing import ( + assert_array_almost_equal_nulp, assert_warns, suppress_warnings +) +import pytest +from pytest import raises as assert_raises +from scipy._lib._array_api import ( + xp_assert_close, xp_assert_equal, + assert_array_almost_equal, +) + +from numpy import array, spacing, sin, pi, sort, sqrt +from scipy.signal import (argrelextrema, BadCoefficients, bessel, besselap, bilinear, + buttap, butter, buttord, cheb1ap, cheb1ord, cheb2ap, + cheb2ord, cheby1, cheby2, ellip, ellipap, ellipord, + firwin, freqs_zpk, freqs, freqz, freqz_zpk, + gammatone, group_delay, iircomb, iirdesign, iirfilter, + iirnotch, iirpeak, lp2bp, lp2bs, lp2hp, lp2lp, normalize, + medfilt, order_filter, + sos2tf, sos2zpk, sosfreqz, freqz_sos, tf2sos, tf2zpk, zpk2sos, + zpk2tf, bilinear_zpk, lp2lp_zpk, lp2hp_zpk, lp2bp_zpk, + lp2bs_zpk) +from scipy.signal._filter_design import (_cplxreal, _cplxpair, _norm_factor, + _bessel_poly, _bessel_zeros) + +try: + import mpmath +except ImportError: + mpmath = None + + +def mpmath_check(min_ver): + return pytest.mark.skipif( + mpmath is None + or _pep440.parse(mpmath.__version__) < _pep440.Version(min_ver), + reason=f"mpmath version >= {min_ver} required", + ) + + +class TestCplxPair: + + def test_trivial_input(self): + assert _cplxpair([]).size == 0 + assert _cplxpair(1) == 1 + + def test_output_order(self): + xp_assert_close(_cplxpair([1+1j, 1-1j]), [1-1j, 1+1j]) + + a = [1+1j, 1+1j, 1, 1-1j, 1-1j, 2] + b = [1-1j, 1+1j, 1-1j, 1+1j, 1, 2] + xp_assert_close(_cplxpair(a), b) + + # points spaced around the unit circle + z = np.exp(2j*pi*array([4, 3, 5, 2, 6, 1, 0])/7) + z1 = np.copy(z) + np.random.shuffle(z) + xp_assert_close(_cplxpair(z), z1) + np.random.shuffle(z) + xp_assert_close(_cplxpair(z), z1) + np.random.shuffle(z) + xp_assert_close(_cplxpair(z), z1) + + # Should be able to pair up all the conjugates + x = np.random.rand(10000) + 1j * np.random.rand(10000) + y = x.conj() + z = np.random.rand(10000) + x = np.concatenate((x, y, z)) + np.random.shuffle(x) + c = _cplxpair(x) + + # Every other element of head should be conjugates: + xp_assert_close(c[0:20000:2], np.conj(c[1:20000:2])) + # Real parts of head should be in sorted order: + xp_assert_close(c[0:20000:2].real, np.sort(c[0:20000:2].real)) + # Tail should be sorted real numbers: + xp_assert_close(c[20000:], np.sort(c[20000:])) + + def test_real_integer_input(self): + xp_assert_equal(_cplxpair([2, 0, 1]), [0, 1, 2]) + + def test_tolerances(self): + eps = spacing(1) + xp_assert_close(_cplxpair([1j, -1j, 1+1j*eps], tol=2*eps), + [-1j, 1j, 1+1j*eps]) + + # sorting close to 0 + xp_assert_close(_cplxpair([-eps+1j, +eps-1j]), [-1j, +1j]) + xp_assert_close(_cplxpair([+eps+1j, -eps-1j]), [-1j, +1j]) + xp_assert_close(_cplxpair([+1j, -1j]), [-1j, +1j]) + + def test_unmatched_conjugates(self): + # 1+2j is unmatched + assert_raises(ValueError, _cplxpair, [1+3j, 1-3j, 1+2j]) + + # 1+2j and 1-3j are unmatched + assert_raises(ValueError, _cplxpair, [1+3j, 1-3j, 1+2j, 1-3j]) + + # 1+3j is unmatched + assert_raises(ValueError, _cplxpair, [1+3j, 1-3j, 1+3j]) + + # Not conjugates + assert_raises(ValueError, _cplxpair, [4+5j, 4+5j]) + assert_raises(ValueError, _cplxpair, [1-7j, 1-7j]) + + # No pairs + assert_raises(ValueError, _cplxpair, [1+3j]) + assert_raises(ValueError, _cplxpair, [1-3j]) + + +class TestCplxReal: + + def test_trivial_input(self): + assert all(x.size == 0 for x in _cplxreal([])) + + x = _cplxreal(1) + assert x[0].size == 0 + xp_assert_equal(x[1], np.asarray([1])) + + + def test_output_order(self): + zc, zr = _cplxreal(np.roots(array([1, 0, 0, 1]))) + xp_assert_close(np.append(zc, zr), [1/2 + 1j*sin(pi/3), -1]) + + eps = spacing(1) + + a = [0+1j, 0-1j, eps + 1j, eps - 1j, -eps + 1j, -eps - 1j, + 1, 4, 2, 3, 0, 0, + 2+3j, 2-3j, + 1-eps + 1j, 1+2j, 1-2j, 1+eps - 1j, # sorts out of order + 3+1j, 3+1j, 3+1j, 3-1j, 3-1j, 3-1j, + 2-3j, 2+3j] + zc, zr = _cplxreal(a) + xp_assert_close(zc, [1j, 1j, 1j, 1+1j, 1+2j, 2+3j, 2+3j, 3+1j, 3+1j, + 3+1j]) + xp_assert_close(zr, [0.0, 0, 1, 2, 3, 4]) + + z = array([1-eps + 1j, 1+2j, 1-2j, 1+eps - 1j, 1+eps+3j, 1-2*eps-3j, + 0+1j, 0-1j, 2+4j, 2-4j, 2+3j, 2-3j, 3+7j, 3-7j, 4-eps+1j, + 4+eps-2j, 4-1j, 4-eps+2j]) + + zc, zr = _cplxreal(z) + xp_assert_close(zc, [1j, 1+1j, 1+2j, 1+3j, 2+3j, 2+4j, 3+7j, 4+1j, + 4+2j]) + xp_assert_equal(zr, np.asarray([])) + + def test_unmatched_conjugates(self): + # 1+2j is unmatched + assert_raises(ValueError, _cplxreal, [1+3j, 1-3j, 1+2j]) + + # 1+2j and 1-3j are unmatched + assert_raises(ValueError, _cplxreal, [1+3j, 1-3j, 1+2j, 1-3j]) + + # 1+3j is unmatched + assert_raises(ValueError, _cplxreal, [1+3j, 1-3j, 1+3j]) + + # No pairs + assert_raises(ValueError, _cplxreal, [1+3j]) + assert_raises(ValueError, _cplxreal, [1-3j]) + + def test_real_integer_input(self): + zc, zr = _cplxreal([2, 0, 1, 4]) + xp_assert_equal(zc, []) + xp_assert_equal(zr, [0, 1, 2, 4]) + + +class TestTf2zpk: + + @pytest.mark.parametrize('dt', (np.float64, np.complex128)) + def test_simple(self, dt): + z_r = np.array([0.5, -0.5]) + p_r = np.array([1.j / np.sqrt(2), -1.j / np.sqrt(2)]) + # Sort the zeros/poles so that we don't fail the test if the order + # changes + z_r.sort() + p_r.sort() + b = np.poly(z_r).astype(dt) + a = np.poly(p_r).astype(dt) + + z, p, k = tf2zpk(b, a) + z.sort() + # The real part of `p` is ~0.0, so sort by imaginary part + p = p[np.argsort(p.imag)] + + assert_array_almost_equal(z, z_r) + assert_array_almost_equal(p, p_r) + assert_array_almost_equal(k, 1.) + assert k.dtype == dt + + def test_bad_filter(self): + # Regression test for #651: better handling of badly conditioned + # filter coefficients. + with suppress_warnings(): + warnings.simplefilter("error", BadCoefficients) + assert_raises(BadCoefficients, tf2zpk, [1e-15], [1.0, 1.0]) + + +class TestZpk2Tf: + + def test_identity(self): + """Test the identity transfer function.""" + z = [] + p = [] + k = 1. + b, a = zpk2tf(z, p, k) + b_r = np.array([1.]) # desired result + a_r = np.array([1.]) # desired result + # The test for the *type* of the return values is a regression + # test for ticket #1095. In the case p=[], zpk2tf used to + # return the scalar 1.0 instead of array([1.0]). + xp_assert_equal(b, b_r) + assert isinstance(b, np.ndarray) + xp_assert_equal(a, a_r) + assert isinstance(a, np.ndarray) + + +class TestSos2Zpk: + + def test_basic(self): + sos = [[1, 0, 1, 1, 0, -0.81], + [1, 0, 0, 1, 0, +0.49]] + z, p, k = sos2zpk(sos) + z2 = [1j, -1j, 0, 0] + p2 = [0.9, -0.9, 0.7j, -0.7j] + k2 = 1 + assert_array_almost_equal(sort(z), sort(z2), decimal=4) + assert_array_almost_equal(sort(p), sort(p2), decimal=4) + assert_array_almost_equal(k, k2) + + sos = [[1.00000, +0.61803, 1.0000, 1.00000, +0.60515, 0.95873], + [1.00000, -1.61803, 1.0000, 1.00000, -1.58430, 0.95873], + [1.00000, +1.00000, 0.0000, 1.00000, +0.97915, 0.00000]] + z, p, k = sos2zpk(sos) + z2 = [-0.3090 + 0.9511j, -0.3090 - 0.9511j, 0.8090 + 0.5878j, + 0.8090 - 0.5878j, -1.0000 + 0.0000j, 0] + p2 = [-0.3026 + 0.9312j, -0.3026 - 0.9312j, 0.7922 + 0.5755j, + 0.7922 - 0.5755j, -0.9791 + 0.0000j, 0] + k2 = 1 + assert_array_almost_equal(sort(z), sort(z2), decimal=4) + assert_array_almost_equal(sort(p), sort(p2), decimal=4) + + sos = array([[1, 2, 3, 1, 0.2, 0.3], + [4, 5, 6, 1, 0.4, 0.5]]) + z = array([-1 - 1.41421356237310j, -1 + 1.41421356237310j, + -0.625 - 1.05326872164704j, -0.625 + 1.05326872164704j]) + p = array([-0.2 - 0.678232998312527j, -0.2 + 0.678232998312527j, + -0.1 - 0.538516480713450j, -0.1 + 0.538516480713450j]) + k = 4 + z2, p2, k2 = sos2zpk(sos) + xp_assert_close(_cplxpair(z2), z) + xp_assert_close(_cplxpair(p2), p) + assert k2 == k + + @pytest.mark.thread_unsafe + def test_fewer_zeros(self): + """Test not the expected number of p/z (effectively at origin).""" + sos = butter(3, 0.1, output='sos') + z, p, k = sos2zpk(sos) + assert len(z) == 4 + assert len(p) == 4 + + sos = butter(12, [5., 30.], 'bandpass', fs=1200., analog=False, + output='sos') + with pytest.warns(BadCoefficients, match='Badly conditioned'): + z, p, k = sos2zpk(sos) + assert len(z) == 24 + assert len(p) == 24 + + +class TestSos2Tf: + + def test_basic(self): + sos = [[1, 1, 1, 1, 0, -1], + [-2, 3, 1, 1, 10, 1]] + b, a = sos2tf(sos) + assert_array_almost_equal(b, [-2, 1, 2, 4, 1]) + assert_array_almost_equal(a, [1, 10, 0, -10, -1]) + + +class TestTf2Sos: + + def test_basic(self): + num = [2, 16, 44, 56, 32] + den = [3, 3, -15, 18, -12] + sos = tf2sos(num, den) + sos2 = [[0.6667, 4.0000, 5.3333, 1.0000, +2.0000, -4.0000], + [1.0000, 2.0000, 2.0000, 1.0000, -1.0000, +1.0000]] + assert_array_almost_equal(sos, sos2, decimal=4) + + b = [1, -3, 11, -27, 18] + a = [16, 12, 2, -4, -1] + sos = tf2sos(b, a) + sos2 = [[0.0625, -0.1875, 0.1250, 1.0000, -0.2500, -0.1250], + [1.0000, +0.0000, 9.0000, 1.0000, +1.0000, +0.5000]] + # assert_array_almost_equal(sos, sos2, decimal=4) + + @pytest.mark.parametrize('b, a, analog, sos', + [([1], [1], False, [[1., 0., 0., 1., 0., 0.]]), + ([1], [1], True, [[0., 0., 1., 0., 0., 1.]]), + ([1], [1., 0., -1.01, 0, 0.01], False, + [[1., 0., 0., 1., 0., -0.01], + [1., 0., 0., 1., 0., -1]]), + ([1], [1., 0., -1.01, 0, 0.01], True, + [[0., 0., 1., 1., 0., -1], + [0., 0., 1., 1., 0., -0.01]])]) + def test_analog(self, b, a, analog, sos): + sos2 = tf2sos(b, a, analog=analog) + assert_array_almost_equal(sos, sos2, decimal=4) + + +class TestZpk2Sos: + + @pytest.mark.parametrize('dt', 'fdgFDG') + @pytest.mark.parametrize('pairing, analog', + [('nearest', False), + ('keep_odd', False), + ('minimal', False), + ('minimal', True)]) + def test_dtypes(self, dt, pairing, analog): + z = np.array([-1, -1]).astype(dt) + ct = dt.upper() # the poles have to be complex + p = np.array([0.57149 + 0.29360j, 0.57149 - 0.29360j]).astype(ct) + k = np.array(1).astype(dt) + sos = zpk2sos(z, p, k, pairing=pairing, analog=analog) + sos2 = [[1, 2, 1, 1, -1.14298, 0.41280]] # octave & MATLAB + assert_array_almost_equal(sos, sos2, decimal=4) + + def test_basic(self): + for pairing in ('nearest', 'keep_odd'): + # + # Cases that match octave + # + + z = [-1, -1] + p = [0.57149 + 0.29360j, 0.57149 - 0.29360j] + k = 1 + sos = zpk2sos(z, p, k, pairing=pairing) + sos2 = [[1, 2, 1, 1, -1.14298, 0.41280]] # octave & MATLAB + assert_array_almost_equal(sos, sos2, decimal=4) + + z = [1j, -1j] + p = [0.9, -0.9, 0.7j, -0.7j] + k = 1 + sos = zpk2sos(z, p, k, pairing=pairing) + sos2 = [[1, 0, 1, 1, 0, +0.49], + [1, 0, 0, 1, 0, -0.81]] # octave + # sos2 = [[0, 0, 1, 1, -0.9, 0], + # [1, 0, 1, 1, 0.9, 0]] # MATLAB + assert_array_almost_equal(sos, sos2, decimal=4) + + z = [] + p = [0.8, -0.5+0.25j, -0.5-0.25j] + k = 1. + sos = zpk2sos(z, p, k, pairing=pairing) + sos2 = [[1., 0., 0., 1., 1., 0.3125], + [1., 0., 0., 1., -0.8, 0.]] # octave, MATLAB fails + assert_array_almost_equal(sos, sos2, decimal=4) + + z = [1., 1., 0.9j, -0.9j] + p = [0.99+0.01j, 0.99-0.01j, 0.1+0.9j, 0.1-0.9j] + k = 1 + sos = zpk2sos(z, p, k, pairing=pairing) + sos2 = [[1, 0, 0.81, 1, -0.2, 0.82], + [1, -2, 1, 1, -1.98, 0.9802]] # octave + # sos2 = [[1, -2, 1, 1, -0.2, 0.82], + # [1, 0, 0.81, 1, -1.98, 0.9802]] # MATLAB + assert_array_almost_equal(sos, sos2, decimal=4) + + z = [0.9+0.1j, 0.9-0.1j, -0.9] + p = [0.75+0.25j, 0.75-0.25j, 0.9] + k = 1 + sos = zpk2sos(z, p, k, pairing=pairing) + if pairing == 'keep_odd': + sos2 = [[1, -1.8, 0.82, 1, -1.5, 0.625], + [1, 0.9, 0, 1, -0.9, 0]] # octave; MATLAB fails + assert_array_almost_equal(sos, sos2, decimal=4) + else: # pairing == 'nearest' + sos2 = [[1, 0.9, 0, 1, -1.5, 0.625], + [1, -1.8, 0.82, 1, -0.9, 0]] # our algorithm + assert_array_almost_equal(sos, sos2, decimal=4) + + # + # Cases that differ from octave: + # + + z = [-0.3090 + 0.9511j, -0.3090 - 0.9511j, 0.8090 + 0.5878j, + +0.8090 - 0.5878j, -1.0000 + 0.0000j] + p = [-0.3026 + 0.9312j, -0.3026 - 0.9312j, 0.7922 + 0.5755j, + +0.7922 - 0.5755j, -0.9791 + 0.0000j] + k = 1 + sos = zpk2sos(z, p, k, pairing=pairing) + # sos2 = [[1, 0.618, 1, 1, 0.6052, 0.95870], + # [1, -1.618, 1, 1, -1.5844, 0.95878], + # [1, 1, 0, 1, 0.9791, 0]] # octave, MATLAB fails + sos2 = [[1, 1, 0, 1, +0.97915, 0], + [1, 0.61803, 1, 1, +0.60515, 0.95873], + [1, -1.61803, 1, 1, -1.58430, 0.95873]] + assert_array_almost_equal(sos, sos2, decimal=4) + + z = [-1 - 1.4142j, -1 + 1.4142j, + -0.625 - 1.0533j, -0.625 + 1.0533j] + p = [-0.2 - 0.6782j, -0.2 + 0.6782j, + -0.1 - 0.5385j, -0.1 + 0.5385j] + k = 4 + sos = zpk2sos(z, p, k, pairing=pairing) + sos2 = [[4, 8, 12, 1, 0.2, 0.3], + [1, 1.25, 1.5, 1, 0.4, 0.5]] # MATLAB + # sos2 = [[4, 8, 12, 1, 0.4, 0.5], + # [1, 1.25, 1.5, 1, 0.2, 0.3]] # octave + xp_assert_close(sos, sos2, rtol=1e-4, atol=1e-4) + + z = [] + p = [0.2, -0.5+0.25j, -0.5-0.25j] + k = 1. + sos = zpk2sos(z, p, k, pairing=pairing) + sos2 = [[1., 0., 0., 1., -0.2, 0.], + [1., 0., 0., 1., 1., 0.3125]] + # sos2 = [[1., 0., 0., 1., 1., 0.3125], + # [1., 0., 0., 1., -0.2, 0]] # octave, MATLAB fails + assert_array_almost_equal(sos, sos2, decimal=4) + + # The next two examples are adapted from Leland B. Jackson, + # "Digital Filters and Signal Processing (1995) p.400: + # http://books.google.com/books?id=VZ8uabI1pNMC&lpg=PA400&ots=gRD9pi8Jua&dq=Pole%2Fzero%20pairing%20for%20minimum%20roundoff%20noise%20in%20BSF.&pg=PA400#v=onepage&q=Pole%2Fzero%20pairing%20for%20minimum%20roundoff%20noise%20in%20BSF.&f=false + + deg2rad = np.pi / 180. + k = 1. + + # first example + thetas = [22.5, 45, 77.5] + mags = [0.8, 0.6, 0.9] + z = np.array([np.exp(theta * deg2rad * 1j) for theta in thetas]) + z = np.concatenate((z, np.conj(z))) + p = np.array([mag * np.exp(theta * deg2rad * 1j) + for theta, mag in zip(thetas, mags)]) + p = np.concatenate((p, np.conj(p))) + sos = zpk2sos(z, p, k) + # sos2 = [[1, -0.43288, 1, 1, -0.38959, 0.81], # octave, + # [1, -1.41421, 1, 1, -0.84853, 0.36], # MATLAB fails + # [1, -1.84776, 1, 1, -1.47821, 0.64]] + # Note that pole-zero pairing matches, but ordering is different + sos2 = [[1, -1.41421, 1, 1, -0.84853, 0.36], + [1, -1.84776, 1, 1, -1.47821, 0.64], + [1, -0.43288, 1, 1, -0.38959, 0.81]] + assert_array_almost_equal(sos, sos2, decimal=4) + + # second example + z = np.array([np.exp(theta * deg2rad * 1j) + for theta in (85., 10.)]) + z = np.concatenate((z, np.conj(z), [1, -1])) + sos = zpk2sos(z, p, k) + + # sos2 = [[1, -0.17431, 1, 1, -0.38959, 0.81], # octave "wrong", + # [1, -1.96962, 1, 1, -0.84853, 0.36], # MATLAB fails + # [1, 0, -1, 1, -1.47821, 0.64000]] + # Our pole-zero pairing matches the text, Octave does not + sos2 = [[1, 0, -1, 1, -0.84853, 0.36], + [1, -1.96962, 1, 1, -1.47821, 0.64], + [1, -0.17431, 1, 1, -0.38959, 0.81]] + assert_array_almost_equal(sos, sos2, decimal=4) + + # these examples are taken from the doc string, and show the + # effect of the 'pairing' argument + @pytest.mark.parametrize('pairing, sos', + [('nearest', + np.array([[1., 1., 0.5, 1., -0.75, 0.], + [1., 1., 0., 1., -1.6, 0.65]])), + ('keep_odd', + np.array([[1., 1., 0, 1., -0.75, 0.], + [1., 1., 0.5, 1., -1.6, 0.65]])), + ('minimal', + np.array([[0., 1., 1., 0., 1., -0.75], + [1., 1., 0.5, 1., -1.6, 0.65]]))]) + def test_pairing(self, pairing, sos): + z1 = np.array([-1, -0.5-0.5j, -0.5+0.5j]) + p1 = np.array([0.75, 0.8+0.1j, 0.8-0.1j]) + sos2 = zpk2sos(z1, p1, 1, pairing=pairing) + assert_array_almost_equal(sos, sos2, decimal=4) + + @pytest.mark.parametrize('p, sos_dt', + [([-1, 1, -0.1, 0.1], + [[0., 0., 1., 1., 0., -0.01], + [0., 0., 1., 1., 0., -1]]), + ([-0.7071+0.7071j, -0.7071-0.7071j, -0.1j, 0.1j], + [[0., 0., 1., 1., 0., 0.01], + [0., 0., 1., 1., 1.4142, 1.]])]) + def test_analog(self, p, sos_dt): + # test `analog` argument + # for discrete time, poles closest to unit circle should appear last + # for cont. time, poles closest to imaginary axis should appear last + sos2_dt = zpk2sos([], p, 1, pairing='minimal', analog=False) + sos2_ct = zpk2sos([], p, 1, pairing='minimal', analog=True) + assert_array_almost_equal(sos_dt, sos2_dt, decimal=4) + assert_array_almost_equal(sos_dt[::-1], sos2_ct, decimal=4) + + def test_bad_args(self): + with pytest.raises(ValueError, match=r'pairing must be one of'): + zpk2sos([1], [2], 1, pairing='no_such_pairing') + + with pytest.raises(ValueError, match=r'.*pairing must be "minimal"'): + zpk2sos([1], [2], 1, pairing='keep_odd', analog=True) + + with pytest.raises(ValueError, + match=r'.*must have len\(p\)>=len\(z\)'): + zpk2sos([1, 1], [2], 1, analog=True) + + with pytest.raises(ValueError, match=r'k must be real'): + zpk2sos([1], [2], k=1j) + + +class TestFreqs: + + def test_basic(self): + _, h = freqs([1.0], [1.0], worN=8) + assert_array_almost_equal(h, np.ones(8)) + + def test_output(self): + # 1st order low-pass filter: H(s) = 1 / (s + 1) + w = [0.1, 1, 10, 100] + num = [1] + den = [1, 1] + w, H = freqs(num, den, worN=w) + s = w * 1j + expected = 1 / (s + 1) + assert_array_almost_equal(H.real, expected.real) + assert_array_almost_equal(H.imag, expected.imag) + + def test_freq_range(self): + # Test that freqresp() finds a reasonable frequency range. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + # Expected range is from 0.01 to 10. + num = [1] + den = [1, 1] + n = 10 + expected_w = np.logspace(-2, 1, n) + w, H = freqs(num, den, worN=n) + assert_array_almost_equal(w, expected_w) + + def test_plot(self): + + def plot(w, h): + assert_array_almost_equal(h, np.ones(8)) + + assert_raises(ZeroDivisionError, freqs, [1.0], [1.0], worN=8, + plot=lambda w, h: 1 / 0) + freqs([1.0], [1.0], worN=8, plot=plot) + + def test_backward_compat(self): + # For backward compatibility, test if None act as a wrapper for default + w1, h1 = freqs([1.0], [1.0]) + w2, h2 = freqs([1.0], [1.0], None) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(h1, h2) + + def test_w_or_N_types(self): + # Measure at 8 equally-spaced points + for N in (8, np.int8(8), np.int16(8), np.int32(8), np.int64(8), + np.array(8)): + w, h = freqs([1.0], [1.0], worN=N) + assert len(w) == 8 + assert_array_almost_equal(h, np.ones(8)) + + # Measure at frequency 8 rad/sec + for w in (8.0, 8.0+0j): + w_out, h = freqs([1.0], [1.0], worN=w) + assert_array_almost_equal(w_out, [8]) + assert_array_almost_equal(h, [1]) + + +class TestFreqs_zpk: + + def test_basic(self): + _, h = freqs_zpk([1.0], [1.0], [1.0], worN=8) + assert_array_almost_equal(h, np.ones(8)) + + def test_output(self): + # 1st order low-pass filter: H(s) = 1 / (s + 1) + w = [0.1, 1, 10, 100] + z = [] + p = [-1] + k = 1 + w, H = freqs_zpk(z, p, k, worN=w) + s = w * 1j + expected = 1 / (s + 1) + assert_array_almost_equal(H.real, expected.real) + assert_array_almost_equal(H.imag, expected.imag) + + def test_freq_range(self): + # Test that freqresp() finds a reasonable frequency range. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + # Expected range is from 0.01 to 10. + z = [] + p = [-1] + k = 1 + n = 10 + expected_w = np.logspace(-2, 1, n) + w, H = freqs_zpk(z, p, k, worN=n) + assert_array_almost_equal(w, expected_w) + + def test_vs_freqs(self): + b, a = cheby1(4, 5, 100, analog=True, output='ba') + z, p, k = cheby1(4, 5, 100, analog=True, output='zpk') + + w1, h1 = freqs(b, a) + w2, h2 = freqs_zpk(z, p, k) + xp_assert_close(w1, w2) + xp_assert_close(h1, h2, rtol=1e-6) + + def test_backward_compat(self): + # For backward compatibility, test if None act as a wrapper for default + w1, h1 = freqs_zpk([1.0], [1.0], [1.0]) + w2, h2 = freqs_zpk([1.0], [1.0], [1.0], None) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(h1, h2) + + def test_w_or_N_types(self): + # Measure at 8 equally-spaced points + for N in (8, np.int8(8), np.int16(8), np.int32(8), np.int64(8), + np.array(8)): + w, h = freqs_zpk([], [], 1, worN=N) + assert len(w) == 8 + assert_array_almost_equal(h, np.ones(8)) + + # Measure at frequency 8 rad/sec + for w in (8.0, 8.0+0j): + w_out, h = freqs_zpk([], [], 1, worN=w) + assert_array_almost_equal(w_out, [8]) + assert_array_almost_equal(h, [1]) + + +class TestFreqz: + + def test_ticket1441(self): + """Regression test for ticket 1441.""" + # Because freqz previously used arange instead of linspace, + # when N was large, it would return one more point than + # requested. + N = 100000 + w, h = freqz([1.0], worN=N) + assert w.shape == (N,) + + def test_basic(self): + w, h = freqz([1.0], worN=8) + assert_array_almost_equal(w, np.pi * np.arange(8) / 8.) + assert_array_almost_equal(h, np.ones(8)) + w, h = freqz([1.0], worN=9) + assert_array_almost_equal(w, np.pi * np.arange(9) / 9.) + assert_array_almost_equal(h, np.ones(9)) + + for a in [1, np.ones(2)]: + w, h = freqz(np.ones(2), a, worN=0) + assert w.shape == (0,) + assert h.shape == (0,) + assert h.dtype == np.dtype('complex128') + + t = np.linspace(0, 1, 4, endpoint=False) + for b, a, h_whole in zip( + ([1., 0, 0, 0], np.sin(2 * np.pi * t)), + ([1., 0, 0, 0], [0.5, 0, 0, 0]), + ([1., 1., 1., 1.], [0, -4j, 0, 4j])): + w, h = freqz(b, a, worN=4, whole=True) + expected_w = np.linspace(0, 2 * np.pi, 4, endpoint=False) + assert_array_almost_equal(w, expected_w) + assert_array_almost_equal(h, h_whole) + # simultaneously check int-like support + w, h = freqz(b, a, worN=np.int32(4), whole=True) + assert_array_almost_equal(w, expected_w) + assert_array_almost_equal(h, h_whole) + w, h = freqz(b, a, worN=w, whole=True) + assert_array_almost_equal(w, expected_w) + assert_array_almost_equal(h, h_whole) + + def test_basic_whole(self): + w, h = freqz([1.0], worN=8, whole=True) + assert_array_almost_equal(w, 2 * np.pi * np.arange(8.0) / 8) + assert_array_almost_equal(h, np.ones(8)) + + def test_plot(self): + + def plot(w, h): + assert_array_almost_equal(w, np.pi * np.arange(8.0) / 8) + assert_array_almost_equal(h, np.ones(8)) + + assert_raises(ZeroDivisionError, freqz, [1.0], worN=8, + plot=lambda w, h: 1 / 0) + freqz([1.0], worN=8, plot=plot) + + def test_fft_wrapping(self): + # Some simple real FIR filters + bs = list() # filters + as_ = list() + hs_whole = list() + hs_half = list() + # 3 taps + t = np.linspace(0, 1, 3, endpoint=False) + bs.append(np.sin(2 * np.pi * t)) + as_.append(3.) + hs_whole.append([0, -0.5j, 0.5j]) + hs_half.append([0, np.sqrt(1./12.), -0.5j]) + # 4 taps + t = np.linspace(0, 1, 4, endpoint=False) + bs.append(np.sin(2 * np.pi * t)) + as_.append(0.5) + hs_whole.append([0, -4j, 0, 4j]) + hs_half.append([0, np.sqrt(8), -4j, -np.sqrt(8)]) + del t + for ii, b in enumerate(bs): + # whole + a = as_[ii] + expected_w = np.linspace(0, 2 * np.pi, len(b), endpoint=False) + w, h = freqz(b, a, worN=expected_w, whole=True) # polyval + err_msg = f'b = {b}, a={a}' + assert_array_almost_equal(w, expected_w, err_msg=err_msg) + assert_array_almost_equal(h, hs_whole[ii], err_msg=err_msg) + w, h = freqz(b, a, worN=len(b), whole=True) # FFT + assert_array_almost_equal(w, expected_w, err_msg=err_msg) + assert_array_almost_equal(h, hs_whole[ii], err_msg=err_msg) + # non-whole + expected_w = np.linspace(0, np.pi, len(b), endpoint=False) + w, h = freqz(b, a, worN=expected_w, whole=False) # polyval + assert_array_almost_equal(w, expected_w, err_msg=err_msg) + assert_array_almost_equal(h, hs_half[ii], err_msg=err_msg) + w, h = freqz(b, a, worN=len(b), whole=False) # FFT + assert_array_almost_equal(w, expected_w, err_msg=err_msg) + assert_array_almost_equal(h, hs_half[ii], err_msg=err_msg) + + # some random FIR filters (real + complex) + # assume polyval is accurate + rng = np.random.RandomState(0) + for ii in range(2, 10): # number of taps + b = rng.randn(ii) + for kk in range(2): + a = rng.randn(1) if kk == 0 else rng.randn(3) + for jj in range(2): + if jj == 1: + b = b + rng.randn(ii) * 1j + # whole + expected_w = np.linspace(0, 2 * np.pi, ii, endpoint=False) + w, expected_h = freqz(b, a, worN=expected_w, whole=True) + assert_array_almost_equal(w, expected_w) + w, h = freqz(b, a, worN=ii, whole=True) + assert_array_almost_equal(w, expected_w) + assert_array_almost_equal(h, expected_h) + # half + expected_w = np.linspace(0, np.pi, ii, endpoint=False) + w, expected_h = freqz(b, a, worN=expected_w, whole=False) + assert_array_almost_equal(w, expected_w) + w, h = freqz(b, a, worN=ii, whole=False) + assert_array_almost_equal(w, expected_w) + assert_array_almost_equal(h, expected_h) + + def test_broadcasting1(self): + # Test broadcasting with worN an integer or a 1-D array, + # b and a are n-dimensional arrays. + np.random.seed(123) + b = np.random.rand(3, 5, 1) + a = np.random.rand(2, 1) + for whole in [False, True]: + # Test with worN being integers (one fast for FFT and one not), + # a 1-D array, and an empty array. + for worN in [16, 17, np.linspace(0, 1, 10), np.array([])]: + w, h = freqz(b, a, worN=worN, whole=whole) + for k in range(b.shape[1]): + bk = b[:, k, 0] + ak = a[:, 0] + ww, hh = freqz(bk, ak, worN=worN, whole=whole) + xp_assert_close(ww, w) + xp_assert_close(hh, h[k]) + + def test_broadcasting2(self): + # Test broadcasting with worN an integer or a 1-D array, + # b is an n-dimensional array, and a is left at the default value. + np.random.seed(123) + b = np.random.rand(3, 5, 1) + for whole in [False, True]: + for worN in [16, 17, np.linspace(0, 1, 10)]: + w, h = freqz(b, worN=worN, whole=whole) + for k in range(b.shape[1]): + bk = b[:, k, 0] + ww, hh = freqz(bk, worN=worN, whole=whole) + xp_assert_close(ww, w) + xp_assert_close(hh, h[k]) + + def test_broadcasting3(self): + # Test broadcasting where b.shape[-1] is the same length + # as worN, and a is left at the default value. + np.random.seed(123) + N = 16 + b = np.random.rand(3, N) + for whole in [False, True]: + for worN in [N, np.linspace(0, 1, N)]: + w, h = freqz(b, worN=worN, whole=whole) + assert w.size == N + for k in range(N): + bk = b[:, k] + ww, hh = freqz(bk, worN=w[k], whole=whole) + xp_assert_close(ww, np.asarray(w[k])[None]) + xp_assert_close(hh, np.asarray(h[k])[None]) + + def test_broadcasting4(self): + # Test broadcasting with worN a 2-D array. + np.random.seed(123) + b = np.random.rand(4, 2, 1, 1) + a = np.random.rand(5, 2, 1, 1) + for whole in [False, True]: + for worN in [np.random.rand(6, 7), np.empty((6, 0))]: + w, h = freqz(b, a, worN=worN, whole=whole) + xp_assert_close(w, worN, rtol=1e-14) + assert h.shape == (2,) + worN.shape + for k in range(2): + ww, hh = freqz(b[:, k, 0, 0], a[:, k, 0, 0], + worN=worN.ravel(), + whole=whole) + xp_assert_close(ww, worN.ravel(), rtol=1e-14) + xp_assert_close(hh, h[k, :, :].ravel()) + + def test_backward_compat(self): + # For backward compatibility, test if None act as a wrapper for default + w1, h1 = freqz([1.0], 1) + w2, h2 = freqz([1.0], 1, None) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(h1, h2) + + def test_fs_param(self): + fs = 900 + b = [0.039479155677484369, 0.11843746703245311, 0.11843746703245311, + 0.039479155677484369] + a = [1.0, -1.3199152021838287, 0.80341991081938424, + -0.16767146321568049] + + # N = None, whole=False + w1, h1 = freqz(b, a, fs=fs) + w2, h2 = freqz(b, a) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs/2, 512, endpoint=False)) + + # N = None, whole=True + w1, h1 = freqz(b, a, whole=True, fs=fs) + w2, h2 = freqz(b, a, whole=True) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs, 512, endpoint=False)) + + # N = 5, whole=False + w1, h1 = freqz(b, a, 5, fs=fs) + w2, h2 = freqz(b, a, 5) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs/2, 5, endpoint=False)) + + # N = 5, whole=True + w1, h1 = freqz(b, a, 5, whole=True, fs=fs) + w2, h2 = freqz(b, a, 5, whole=True) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs, 5, endpoint=False)) + + # w is an array_like + for w in ([123], (123,), np.array([123]), (50, 123, 230), + np.array([50, 123, 230])): + w1, h1 = freqz(b, a, w, fs=fs) + w2, h2 = freqz(b, a, 2*pi*np.array(w)/fs) + xp_assert_close(h1, h2) + xp_assert_close(w, w1, check_dtype=False) + + def test_w_or_N_types(self): + # Measure at 7 (polyval) or 8 (fft) equally-spaced points + for N in (7, np.int8(7), np.int16(7), np.int32(7), np.int64(7), + np.array(7), + 8, np.int8(8), np.int16(8), np.int32(8), np.int64(8), + np.array(8)): + + w, h = freqz([1.0], worN=N) + assert_array_almost_equal(w, np.pi * np.arange(N) / N) + assert_array_almost_equal(h, np.ones(N)) + + w, h = freqz([1.0], worN=N, fs=100) + assert_array_almost_equal(w, np.linspace(0, 50, N, endpoint=False)) + assert_array_almost_equal(h, np.ones(N)) + + # Measure at frequency 8 Hz + for w in (8.0, 8.0+0j): + # Only makes sense when fs is specified + w_out, h = freqz([1.0], worN=w, fs=100) + assert_array_almost_equal(w_out, [8]) + assert_array_almost_equal(h, [1]) + + def test_nyquist(self): + w, h = freqz([1.0], worN=8, include_nyquist=True) + assert_array_almost_equal(w, np.pi * np.arange(8) / 7.) + assert_array_almost_equal(h, np.ones(8)) + w, h = freqz([1.0], worN=9, include_nyquist=True) + assert_array_almost_equal(w, np.pi * np.arange(9) / 8.) + assert_array_almost_equal(h, np.ones(9)) + + for a in [1, np.ones(2)]: + w, h = freqz(np.ones(2), a, worN=0, include_nyquist=True) + assert w.shape == (0,) + assert h.shape == (0,) + assert h.dtype == np.dtype('complex128') + + w1, h1 = freqz([1.0], worN=8, whole = True, include_nyquist=True) + w2, h2 = freqz([1.0], worN=8, whole = True, include_nyquist=False) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(h1, h2) + + # https://github.com/scipy/scipy/issues/17289 + # https://github.com/scipy/scipy/issues/15273 + @pytest.mark.parametrize('whole,nyquist,worN', + [(False, False, 32), + (False, True, 32), + (True, False, 32), + (True, True, 32), + (False, False, 257), + (False, True, 257), + (True, False, 257), + (True, True, 257)]) + def test_17289(self, whole, nyquist, worN): + d = [0, 1] + w, Drfft = freqz(d, worN=32, whole=whole, include_nyquist=nyquist) + _, Dpoly = freqz(d, worN=w) + xp_assert_close(Drfft, Dpoly) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + freqz([1.0], fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none."): + freqz([1.0], fs=None) + + +class Testfreqz_sos: + + def test_freqz_sos_basic(self): + # Compare the results of freqz and freqz_sos for a low order + # Butterworth filter. + + N = 500 + + b, a = butter(4, 0.2) + sos = butter(4, 0.2, output='sos') + w, h = freqz(b, a, worN=N) + w2, h2 = freqz_sos(sos, worN=N) + xp_assert_equal(w2, w) + xp_assert_close(h2, h, rtol=1e-10, atol=1e-14) + + b, a = ellip(3, 1, 30, (0.2, 0.3), btype='bandpass') + sos = ellip(3, 1, 30, (0.2, 0.3), btype='bandpass', output='sos') + w, h = freqz(b, a, worN=N) + w2, h2 = freqz_sos(sos, worN=N) + xp_assert_equal(w2, w) + xp_assert_close(h2, h, rtol=1e-10, atol=1e-14) + # must have at least one section + assert_raises(ValueError, freqz_sos, sos[:0]) + + def test_backward_compat(self): + # For backward compatibility, test if None act as a wrapper for default + N = 500 + + sos = butter(4, 0.2, output='sos') + w1, h1 = freqz_sos(sos, worN=N) + w2, h2 = sosfreqz(sos, worN=N) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(h1, h2) + + def test_freqz_sos_design(self): + # Compare freqz_sos output against expected values for different + # filter types + + # from cheb2ord + N, Wn = cheb2ord([0.1, 0.6], [0.2, 0.5], 3, 60) + sos = cheby2(N, 60, Wn, 'stop', output='sos') + w, h = freqz_sos(sos) + h = np.abs(h) + w /= np.pi + xp_assert_close(20 * np.log10(h[w <= 0.1]), np.asarray(0.), atol=3.01, + check_shape=False) + xp_assert_close(20 * np.log10(h[w >= 0.6]), np.asarray(0.), atol=3.01, + check_shape=False) + xp_assert_close(h[(w >= 0.2) & (w <= 0.5)], + np.asarray(0.), atol=1e-3, + check_shape=False) # <= -60 dB + + N, Wn = cheb2ord([0.1, 0.6], [0.2, 0.5], 3, 150) + sos = cheby2(N, 150, Wn, 'stop', output='sos') + w, h = freqz_sos(sos) + dB = 20*np.log10(np.abs(h)) + w /= np.pi + xp_assert_close(dB[w <= 0.1], np.asarray(0.0), atol=3.01, check_shape=False) + xp_assert_close(dB[w >= 0.6], np.asarray(0.0), atol=3.01, check_shape=False) + assert np.all(dB[(w >= 0.2) & (w <= 0.5)] < -149.9) + + # from cheb1ord + N, Wn = cheb1ord(0.2, 0.3, 3, 40) + sos = cheby1(N, 3, Wn, 'low', output='sos') + w, h = freqz_sos(sos) + h = np.abs(h) + w /= np.pi + xp_assert_close(20 * np.log10(h[w <= 0.2]), np.asarray(0.0), atol=3.01, + check_shape=False) + xp_assert_close(h[w >= 0.3], np.asarray(0.0), atol=1e-2, + check_shape=False) # <= -40 dB + + N, Wn = cheb1ord(0.2, 0.3, 1, 150) + sos = cheby1(N, 1, Wn, 'low', output='sos') + w, h = freqz_sos(sos) + dB = 20*np.log10(np.abs(h)) + w /= np.pi + xp_assert_close(dB[w <= 0.2], np.asarray(0.0), atol=1.01, + check_shape=False) + assert np.all(dB[w >= 0.3] < -149.9) + + # adapted from ellipord + N, Wn = ellipord(0.3, 0.2, 3, 60) + sos = ellip(N, 0.3, 60, Wn, 'high', output='sos') + w, h = freqz_sos(sos) + h = np.abs(h) + w /= np.pi + xp_assert_close(20 * np.log10(h[w >= 0.3]), np.asarray(0.0), atol=3.01, + check_shape=False) + xp_assert_close(h[w <= 0.1], np.asarray(0.0), atol=1.5e-3, + check_shape=False) # <= -60 dB (approx) + + # adapted from buttord + N, Wn = buttord([0.2, 0.5], [0.14, 0.6], 3, 40) + sos = butter(N, Wn, 'band', output='sos') + w, h = freqz_sos(sos) + h = np.abs(h) + w /= np.pi + + h014 = h[w <= 0.14] + xp_assert_close(h014, np.zeros_like(h014), atol=1e-2) # <= -40 dB + h06 = h[w >= 0.6] + xp_assert_close(h06, np.zeros_like(h06), atol=1e-2) # <= -40 dB + h0205 = 20 * np.log10(h[(w >= 0.2) & (w <= 0.5)]) + xp_assert_close(h0205, np.zeros_like(h0205), atol=3.01) + + N, Wn = buttord([0.2, 0.5], [0.14, 0.6], 3, 100) + sos = butter(N, Wn, 'band', output='sos') + w, h = freqz_sos(sos) + dB = 20*np.log10(np.maximum(np.abs(h), 1e-10)) + w /= np.pi + + assert np.all(dB[(w > 0) & (w <= 0.14)] < -99.9) + assert np.all(dB[w >= 0.6] < -99.9) + db0205 = dB[(w >= 0.2) & (w <= 0.5)] + xp_assert_close(db0205, np.zeros_like(db0205), atol=3.01) + + def test_freqz_sos_design_ellip(self): + N, Wn = ellipord(0.3, 0.1, 3, 60) + sos = ellip(N, 0.3, 60, Wn, 'high', output='sos') + w, h = freqz_sos(sos) + h = np.abs(h) + w /= np.pi + + h03 = 20 * np.log10(h[w >= 0.3]) + xp_assert_close(h03, np.zeros_like(h03), atol=3.01) + h01 = h[w <= 0.1] + xp_assert_close(h01, np.zeros_like(h01), atol=1.5e-3) # <= -60 dB (approx) + + N, Wn = ellipord(0.3, 0.2, .5, 150) + sos = ellip(N, .5, 150, Wn, 'high', output='sos') + w, h = freqz_sos(sos) + dB = 20*np.log10(np.maximum(np.abs(h), 1e-10)) + w /= np.pi + + db03 = dB[w >= 0.3] + xp_assert_close(db03, np.zeros_like(db03), atol=.55) + # Allow some numerical slop in the upper bound -150, so this is + # a check that dB[w <= 0.2] is less than or almost equal to -150. + assert dB[w <= 0.2].max() < -150*(1 - 1e-12) + + @mpmath_check("0.10") + def test_freqz_sos_against_mp(self): + # Compare the result of freqz_sos applied to a high order Butterworth + # filter against the result computed using mpmath. (signal.freqz fails + # miserably with such high order filters.) + from . import mpsig + N = 500 + order = 25 + Wn = 0.15 + with mpmath.workdps(80): + z_mp, p_mp, k_mp = mpsig.butter_lp(order, Wn) + w_mp, h_mp = mpsig.zpkfreqz(z_mp, p_mp, k_mp, N) + w_mp = np.array([float(x) for x in w_mp]) + h_mp = np.array([complex(x) for x in h_mp]) + + sos = butter(order, Wn, output='sos') + w, h = freqz_sos(sos, worN=N) + xp_assert_close(w, w_mp, rtol=1e-12, atol=1e-14) + xp_assert_close(h, h_mp, rtol=1e-12, atol=1e-14) + + def test_fs_param(self): + fs = 900 + sos = [[0.03934683014103762, 0.07869366028207524, 0.03934683014103762, + 1.0, -0.37256600288916636, 0.0], + [1.0, 1.0, 0.0, 1.0, -0.9495739996946778, 0.45125966317124144]] + + # N = None, whole=False + w1, h1 = freqz_sos(sos, fs=fs) + w2, h2 = freqz_sos(sos) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs/2, 512, endpoint=False)) + + # N = None, whole=True + w1, h1 = freqz_sos(sos, whole=True, fs=fs) + w2, h2 = freqz_sos(sos, whole=True) + xp_assert_close(h1, h2, atol=1e-27) + xp_assert_close(w1, np.linspace(0, fs, 512, endpoint=False)) + + # N = 5, whole=False + w1, h1 = freqz_sos(sos, 5, fs=fs) + w2, h2 = freqz_sos(sos, 5) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs/2, 5, endpoint=False)) + + # N = 5, whole=True + w1, h1 = freqz_sos(sos, 5, whole=True, fs=fs) + w2, h2 = freqz_sos(sos, 5, whole=True) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs, 5, endpoint=False)) + + # w is an array_like + for w in ([123], (123,), np.array([123]), (50, 123, 230), + np.array([50, 123, 230])): + w1, h1 = freqz_sos(sos, w, fs=fs) + w2, h2 = freqz_sos(sos, 2*pi*np.array(w)/fs) + xp_assert_close(h1, h2) + xp_assert_close(w, w1, check_dtype=False) + + def test_w_or_N_types(self): + # Measure at 7 (polyval) or 8 (fft) equally-spaced points + for N in (7, np.int8(7), np.int16(7), np.int32(7), np.int64(7), + np.array(7), + 8, np.int8(8), np.int16(8), np.int32(8), np.int64(8), + np.array(8)): + + w, h = freqz_sos([1, 0, 0, 1, 0, 0], worN=N) + assert_array_almost_equal(w, np.pi * np.arange(N) / N) + assert_array_almost_equal(h, np.ones(N)) + + w, h = freqz_sos([1, 0, 0, 1, 0, 0], worN=N, fs=100) + assert_array_almost_equal(w, np.linspace(0, 50, N, endpoint=False)) + assert_array_almost_equal(h, np.ones(N)) + + # Measure at frequency 8 Hz + for w in (8.0, 8.0+0j): + # Only makes sense when fs is specified + w_out, h = freqz_sos([1, 0, 0, 1, 0, 0], worN=w, fs=100) + assert_array_almost_equal(w_out, [8]) + assert_array_almost_equal(h, [1]) + + def test_fs_validation(self): + sos = butter(4, 0.2, output='sos') + with pytest.raises(ValueError, match="Sampling.*single scalar"): + freqz_sos(sos, fs=np.array([10, 20])) + + +class TestFreqz_zpk: + + def test_ticket1441(self): + """Regression test for ticket 1441.""" + # Because freqz previously used arange instead of linspace, + # when N was large, it would return one more point than + # requested. + N = 100000 + w, h = freqz_zpk([0.5], [0.5], 1.0, worN=N) + assert w.shape == (N,) + + def test_basic(self): + w, h = freqz_zpk([0.5], [0.5], 1.0, worN=8) + assert_array_almost_equal(w, np.pi * np.arange(8.0) / 8) + assert_array_almost_equal(h, np.ones(8)) + + def test_basic_whole(self): + w, h = freqz_zpk([0.5], [0.5], 1.0, worN=8, whole=True) + assert_array_almost_equal(w, 2 * np.pi * np.arange(8.0) / 8) + assert_array_almost_equal(h, np.ones(8)) + + def test_vs_freqz(self): + b, a = cheby1(4, 5, 0.5, analog=False, output='ba') + z, p, k = cheby1(4, 5, 0.5, analog=False, output='zpk') + + w1, h1 = freqz(b, a) + w2, h2 = freqz_zpk(z, p, k) + xp_assert_close(w1, w2) + xp_assert_close(h1, h2, rtol=1e-6) + + def test_backward_compat(self): + # For backward compatibility, test if None act as a wrapper for default + w1, h1 = freqz_zpk([0.5], [0.5], 1.0) + w2, h2 = freqz_zpk([0.5], [0.5], 1.0, None) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(h1, h2) + + def test_fs_param(self): + fs = 900 + z = [-1, -1, -1] + p = [0.4747869998473389+0.4752230717749344j, 0.37256600288916636, + 0.4747869998473389-0.4752230717749344j] + k = 0.03934683014103762 + + # N = None, whole=False + w1, h1 = freqz_zpk(z, p, k, whole=False, fs=fs) + w2, h2 = freqz_zpk(z, p, k, whole=False) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs/2, 512, endpoint=False)) + + # N = None, whole=True + w1, h1 = freqz_zpk(z, p, k, whole=True, fs=fs) + w2, h2 = freqz_zpk(z, p, k, whole=True) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs, 512, endpoint=False)) + + # N = 5, whole=False + w1, h1 = freqz_zpk(z, p, k, 5, fs=fs) + w2, h2 = freqz_zpk(z, p, k, 5) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs/2, 5, endpoint=False)) + + # N = 5, whole=True + w1, h1 = freqz_zpk(z, p, k, 5, whole=True, fs=fs) + w2, h2 = freqz_zpk(z, p, k, 5, whole=True) + xp_assert_close(h1, h2) + xp_assert_close(w1, np.linspace(0, fs, 5, endpoint=False)) + + # w is an array_like + for w in ([123], (123,), np.array([123]), (50, 123, 230), + np.array([50, 123, 230])): + w1, h1 = freqz_zpk(z, p, k, w, fs=fs) + w2, h2 = freqz_zpk(z, p, k, 2*pi*np.array(w)/fs) + xp_assert_close(h1, h2) + xp_assert_close(w, w1, check_dtype=False) + + def test_w_or_N_types(self): + # Measure at 8 equally-spaced points + for N in (8, np.int8(8), np.int16(8), np.int32(8), np.int64(8), + np.array(8)): + + w, h = freqz_zpk([], [], 1, worN=N) + assert_array_almost_equal(w, np.pi * np.arange(8) / 8.) + assert_array_almost_equal(h, np.ones(8)) + + w, h = freqz_zpk([], [], 1, worN=N, fs=100) + assert_array_almost_equal(w, np.linspace(0, 50, 8, endpoint=False)) + assert_array_almost_equal(h, np.ones(8)) + + # Measure at frequency 8 Hz + for w in (8.0, 8.0+0j): + # Only makes sense when fs is specified + w_out, h = freqz_zpk([], [], 1, worN=w, fs=100) + assert_array_almost_equal(w_out, [8]) + assert_array_almost_equal(h, [1]) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + freqz_zpk([1.0], [1.0], [1.0], fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none."): + freqz_zpk([1.0], [1.0], [1.0], fs=None) + + +class TestNormalize: + + def test_allclose(self): + """Test for false positive on allclose in normalize() in + filter_design.py""" + # Test to make sure the allclose call within signal.normalize does not + # choose false positives. Then check against a known output from MATLAB + # to make sure the fix doesn't break anything. + + # These are the coefficients returned from + # `[b,a] = cheby1(8, 0.5, 0.048)' + # in MATLAB. There are at least 15 significant figures in each + # coefficient, so it makes sense to test for errors on the order of + # 1e-13 (this can always be relaxed if different platforms have + # different rounding errors) + b_matlab = np.array([2.150733144728282e-11, 1.720586515782626e-10, + 6.022052805239190e-10, 1.204410561047838e-09, + 1.505513201309798e-09, 1.204410561047838e-09, + 6.022052805239190e-10, 1.720586515782626e-10, + 2.150733144728282e-11]) + a_matlab = np.array([1.000000000000000e+00, -7.782402035027959e+00, + 2.654354569747454e+01, -5.182182531666387e+01, + 6.334127355102684e+01, -4.963358186631157e+01, + 2.434862182949389e+01, -6.836925348604676e+00, + 8.412934944449140e-01]) + + # This is the input to signal.normalize after passing through the + # equivalent steps in signal.iirfilter as was done for MATLAB + b_norm_in = np.array([1.5543135865293012e-06, 1.2434508692234413e-05, + 4.3520780422820447e-05, 8.7041560845640893e-05, + 1.0880195105705122e-04, 8.7041560845640975e-05, + 4.3520780422820447e-05, 1.2434508692234413e-05, + 1.5543135865293012e-06]) + a_norm_in = np.array([7.2269025909127173e+04, -5.6242661430467968e+05, + 1.9182761917308895e+06, -3.7451128364682454e+06, + 4.5776121393762771e+06, -3.5869706138592605e+06, + 1.7596511818472347e+06, -4.9409793515707983e+05, + 6.0799461347219651e+04]) + + b_output, a_output = normalize(b_norm_in, a_norm_in) + + # The test on b works for decimal=14 but the one for a does not. For + # the sake of consistency, both of these are decimal=13. If something + # breaks on another platform, it is probably fine to relax this lower. + assert_array_almost_equal(b_matlab, b_output, decimal=13) + assert_array_almost_equal(a_matlab, a_output, decimal=13) + + def test_errors(self): + """Test the error cases.""" + # all zero denominator + assert_raises(ValueError, normalize, [1, 2], 0) + + # denominator not 1 dimensional + assert_raises(ValueError, normalize, [1, 2], [[1]]) + + # numerator too many dimensions + assert_raises(ValueError, normalize, [[[1, 2]]], 1) + + +class TestLp2lp: + + def test_basic(self): + b = [1] + a = [1, np.sqrt(2), 1] + b_lp, a_lp = lp2lp(b, a, 0.38574256627112119) + assert_array_almost_equal(b_lp, [0.1488], decimal=4) + assert_array_almost_equal(a_lp, [1, 0.5455, 0.1488], decimal=4) + + +class TestLp2hp: + + def test_basic(self): + b = [0.25059432325190018] + a = [1, 0.59724041654134863, 0.92834805757524175, 0.25059432325190018] + b_hp, a_hp = lp2hp(b, a, 2*np.pi*5000) + xp_assert_close(b_hp, [1.0, 0, 0, 0]) + xp_assert_close(a_hp, [1, 1.1638e5, 2.3522e9, 1.2373e14], rtol=1e-4) + + +class TestLp2bp: + + def test_basic(self): + b = [1] + a = [1, 2, 2, 1] + b_bp, a_bp = lp2bp(b, a, 2*np.pi*4000, 2*np.pi*2000) + xp_assert_close(b_bp, [1.9844e12, 0, 0, 0], rtol=1e-6) + xp_assert_close(a_bp, [1, 2.5133e4, 2.2108e9, 3.3735e13, + 1.3965e18, 1.0028e22, 2.5202e26], rtol=1e-4) + + +class TestLp2bs: + + def test_basic(self): + b = [1] + a = [1, 1] + b_bs, a_bs = lp2bs(b, a, 0.41722257286366754, 0.18460575326152251) + assert_array_almost_equal(b_bs, [1, 0, 0.17407], decimal=5) + assert_array_almost_equal(a_bs, [1, 0.18461, 0.17407], decimal=5) + + +class TestBilinear: + + def test_basic(self): + b = [0.14879732743343033] + a = [1, 0.54552236880522209, 0.14879732743343033] + b_z, a_z = bilinear(b, a, 0.5) + assert_array_almost_equal(b_z, [0.087821, 0.17564, 0.087821], + decimal=5) + assert_array_almost_equal(a_z, [1, -1.0048, 0.35606], decimal=4) + + b = [1, 0, 0.17407467530697837] + a = [1, 0.18460575326152251, 0.17407467530697837] + b_z, a_z = bilinear(b, a, 0.5) + assert_array_almost_equal(b_z, [0.86413, -1.2158, 0.86413], + decimal=4) + assert_array_almost_equal(a_z, [1, -1.2158, 0.72826], + decimal=4) + + def test_fs_validation(self): + b = [0.14879732743343033] + a = [1, 0.54552236880522209, 0.14879732743343033] + with pytest.raises(ValueError, match="Sampling.*single scalar"): + bilinear(b, a, fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none"): + bilinear(b, a, fs=None) + + +class TestLp2lp_zpk: + + def test_basic(self): + z = [] + p = [(-1+1j)/np.sqrt(2), (-1-1j)/np.sqrt(2)] + k = 1 + z_lp, p_lp, k_lp = lp2lp_zpk(z, p, k, 5) + xp_assert_equal(z_lp, []) + xp_assert_close(sort(p_lp), sort(p)*5) + xp_assert_close(k_lp, 25.) + + # Pseudo-Chebyshev with both poles and zeros + z = [-2j, +2j] + p = [-0.75, -0.5-0.5j, -0.5+0.5j] + k = 3 + z_lp, p_lp, k_lp = lp2lp_zpk(z, p, k, 20) + xp_assert_close(sort(z_lp), sort([-40j, +40j])) + xp_assert_close(sort(p_lp), sort([-15, -10-10j, -10+10j])) + xp_assert_close(k_lp, 60.) + + def test_fs_validation(self): + z = [-2j, +2j] + p = [-0.75, -0.5 - 0.5j, -0.5 + 0.5j] + k = 3 + + with pytest.raises(ValueError, match="Sampling.*single scalar"): + bilinear_zpk(z, p, k, fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none"): + bilinear_zpk(z, p, k, fs=None) + + +class TestLp2hp_zpk: + + def test_basic(self): + z = [] + p = [(-1+1j)/np.sqrt(2), (-1-1j)/np.sqrt(2)] + k = 1 + + z_hp, p_hp, k_hp = lp2hp_zpk(z, p, k, 5) + xp_assert_equal(z_hp, np.asarray([0.0, 0.0])) + xp_assert_close(sort(p_hp), sort(p)*5) + xp_assert_close(k_hp, 1.0) + + z = [-2j, +2j] + p = [-0.75, -0.5-0.5j, -0.5+0.5j] + k = 3 + z_hp, p_hp, k_hp = lp2hp_zpk(z, p, k, 6) + xp_assert_close(sort(z_hp), sort([-3j, 0, +3j])) + xp_assert_close(sort(p_hp), sort([-8, -6-6j, -6+6j])) + xp_assert_close(k_hp, 32.0) + + +class TestLp2bp_zpk: + + def test_basic(self): + z = [-2j, +2j] + p = [-0.75, -0.5-0.5j, -0.5+0.5j] + k = 3 + z_bp, p_bp, k_bp = lp2bp_zpk(z, p, k, 15, 8) + xp_assert_close(sort(z_bp), sort([-25j, -9j, 0, +9j, +25j])) + xp_assert_close(sort(p_bp), sort([-3 + 6j*sqrt(6), + -3 - 6j*sqrt(6), + +2j+sqrt(-8j-225)-2, + -2j+sqrt(+8j-225)-2, + +2j-sqrt(-8j-225)-2, + -2j-sqrt(+8j-225)-2, ])) + xp_assert_close(k_bp, 24.0) + + +class TestLp2bs_zpk: + + def test_basic(self): + z = [-2j, +2j] + p = [-0.75, -0.5-0.5j, -0.5+0.5j] + k = 3 + + z_bs, p_bs, k_bs = lp2bs_zpk(z, p, k, 35, 12) + + xp_assert_close(sort(z_bs), sort([+35j, -35j, + +3j+sqrt(1234)*1j, + -3j+sqrt(1234)*1j, + +3j-sqrt(1234)*1j, + -3j-sqrt(1234)*1j])) + xp_assert_close(sort(p_bs), sort([+3j*sqrt(129) - 8, + -3j*sqrt(129) - 8, + (-6 + 6j) - sqrt(-1225 - 72j), + (-6 - 6j) - sqrt(-1225 + 72j), + (-6 + 6j) + sqrt(-1225 - 72j), + (-6 - 6j) + sqrt(-1225 + 72j), ])) + xp_assert_close(k_bs, 32.0) + + +class TestBilinear_zpk: + + def test_basic(self): + z = [-2j, +2j] + p = [-0.75, -0.5-0.5j, -0.5+0.5j] + k = 3 + + z_d, p_d, k_d = bilinear_zpk(z, p, k, 10) + + xp_assert_close(sort(z_d), sort([(20-2j)/(20+2j), (20+2j)/(20-2j), + -1])) + xp_assert_close(sort(p_d), sort([77/83, + (1j/2 + 39/2) / (41/2 - 1j/2), + (39/2 - 1j/2) / (1j/2 + 41/2), ])) + xp_assert_close(k_d, 9696/69803) + + +class TestPrototypeType: + + def test_output_type(self): + # Prototypes should consistently output arrays, not lists + # https://github.com/scipy/scipy/pull/441 + for func in (buttap, + besselap, + lambda N: cheb1ap(N, 1), + lambda N: cheb2ap(N, 20), + lambda N: ellipap(N, 1, 20)): + for N in range(7): + z, p, k = func(N) + assert isinstance(z, np.ndarray) + assert isinstance(p, np.ndarray) + + +def dB(x): + # Return magnitude in decibels, avoiding divide-by-zero warnings + # (and deal with some "not less-ordered" errors when -inf shows up) + return 20 * np.log10(np.maximum(np.abs(x), np.finfo(np.float64).tiny)) + + +class TestButtord: + + def test_lowpass(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + N, Wn = buttord(wp, ws, rp, rs, False) + b, a = butter(N, Wn, 'lowpass', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs) + + assert N == 16 + xp_assert_close(Wn, + 2.0002776782743284e-01, rtol=1e-15) + + def test_highpass(self): + wp = 0.3 + ws = 0.2 + rp = 3 + rs = 70 + N, Wn = buttord(wp, ws, rp, rs, False) + b, a = butter(N, Wn, 'highpass', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp < dB(h[wp <= w])) + assert np.all(dB(h[w <= ws]) < -rs) + + assert N == 18 + xp_assert_close(Wn, + 2.9996603079132672e-01, rtol=1e-15) + + def test_bandpass(self): + wp = [0.2, 0.5] + ws = [0.1, 0.6] + rp = 3 + rs = 80 + N, Wn = buttord(wp, ws, rp, rs, False) + b, a = butter(N, Wn, 'bandpass', False) + w, h = freqz(b, a) + w /= np.pi + + assert np.all((-rp - 0.1) < dB(h[np.logical_and(wp[0] <= w, w <= wp[1])])) + + assert np.all(dB(h[np.logical_or(w <= ws[0], ws[1] <= w)]) < (-rs + 0.1)) + + assert N == 18 + xp_assert_close(Wn, [1.9998742411409134e-01, 5.0002139595676276e-01], + rtol=1e-15) + + def test_bandstop(self): + wp = [0.1, 0.6] + ws = [0.2, 0.5] + rp = 3 + rs = 90 + N, Wn = buttord(wp, ws, rp, rs, False) + b, a = butter(N, Wn, 'bandstop', False) + w, h = freqz(b, a) + w /= np.pi + + assert np.all(-rp < dB(h[np.logical_or(w <= wp[0], wp[1] <= w)])) + assert np.all(dB(h[np.logical_and(ws[0] <= w, w <= ws[1])]) < -rs) + + assert N == 20 + xp_assert_close(Wn, [1.4759432329294042e-01, 5.9997365985276407e-01], + rtol=1e-6) + + def test_analog(self): + wp = 200 + ws = 600 + rp = 3 + rs = 60 + N, Wn = buttord(wp, ws, rp, rs, True) + b, a = butter(N, Wn, 'lowpass', True) + w, h = freqs(b, a) + assert np.all(-rp < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs) + + assert N == 7 + xp_assert_close(Wn, 2.0006785355671877e+02, rtol=1e-15) + + n, Wn = buttord(1, 550/450, 1, 26, analog=True) + assert n == 19 + xp_assert_close(Wn, 1.0361980524629517, rtol=1e-15) + + xp_assert_equal(buttord(1, 1.2, 1, 80, analog=True)[0], 55) + + def test_fs_param(self): + wp = [4410, 11025] + ws = [2205, 13230] + rp = 3 + rs = 80 + fs = 44100 + N, Wn = buttord(wp, ws, rp, rs, False, fs=fs) + b, a = butter(N, Wn, 'bandpass', False, fs=fs) + w, h = freqz(b, a, fs=fs) + assert np.all(-rp - 0.1 < dB(h[np.logical_and(wp[0] <= w, w <= wp[1])])) + assert np.all(dB(h[np.logical_or(w <= ws[0], ws[1] <= w)]) < -rs + 0.1) + + assert N == 18 + xp_assert_close(Wn, [4409.722701715714, 11025.47178084662], + rtol=1e-15) + + def test_invalid_input(self): + with pytest.raises(ValueError) as exc_info: + buttord([20, 50], [14, 60], 3, 2) + assert "gpass should be smaller than gstop" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + buttord([20, 50], [14, 60], -1, 2) + assert "gpass should be larger than 0.0" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + buttord([20, 50], [14, 60], 1, -2) + assert "gstop should be larger than 0.0" in str(exc_info.value) + + @pytest.mark.thread_unsafe + def test_runtime_warnings(self): + msg = "Order is zero.*|divide by zero encountered" + with pytest.warns(RuntimeWarning, match=msg): + buttord(0.0, 1.0, 3, 60) + + def test_ellip_butter(self): + # The purpose of the test is to compare to some known output from past + # scipy versions. The values to compare to are generated with scipy + # 1.9.1 (there is nothing special about this particular version though) + n, wn = buttord([0.1, 0.6], [0.2, 0.5], 3, 60) + assert n == 14 + + def test_fs_validation(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + + with pytest.raises(ValueError, match="Sampling.*single scalar"): + buttord(wp, ws, rp, rs, False, fs=np.array([10, 20])) + + +class TestCheb1ord: + + def test_lowpass(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + N, Wn = cheb1ord(wp, ws, rp, rs, False) + b, a = cheby1(N, rp, Wn, 'low', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs + 0.1) + + assert N == 8 + xp_assert_close(Wn, 0.2, rtol=1e-15) + + def test_highpass(self): + wp = 0.3 + ws = 0.2 + rp = 3 + rs = 70 + N, Wn = cheb1ord(wp, ws, rp, rs, False) + b, a = cheby1(N, rp, Wn, 'high', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[wp <= w])) + assert np.all(dB(h[w <= ws]) < -rs + 0.1) + + assert N == 9 + xp_assert_close(Wn, 0.3, rtol=1e-15) + + def test_bandpass(self): + wp = [0.2, 0.5] + ws = [0.1, 0.6] + rp = 3 + rs = 80 + N, Wn = cheb1ord(wp, ws, rp, rs, False) + b, a = cheby1(N, rp, Wn, 'band', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[np.logical_and(wp[0] <= w, w <= wp[1])])) + assert np.all(dB(h[np.logical_or(w <= ws[0], ws[1] <= w)]) < -rs + 0.1) + + assert N == 9 + xp_assert_close(Wn, [0.2, 0.5], rtol=1e-15) + + def test_bandstop(self): + wp = [0.1, 0.6] + ws = [0.2, 0.5] + rp = 3 + rs = 90 + N, Wn = cheb1ord(wp, ws, rp, rs, False) + b, a = cheby1(N, rp, Wn, 'stop', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[np.logical_or(w <= wp[0], wp[1] <= w)])) + assert np.all(dB(h[np.logical_and(ws[0] <= w, w <= ws[1])]) < -rs + 0.1) + + assert N == 10 + xp_assert_close(Wn, [0.14758232569947785, 0.6], rtol=1e-5) + + def test_analog(self): + wp = 700 + ws = 100 + rp = 3 + rs = 70 + N, Wn = cheb1ord(wp, ws, rp, rs, True) + b, a = cheby1(N, rp, Wn, 'high', True) + w, h = freqs(b, a) + assert np.all(-rp - 0.1 < dB(h[wp <= w])) + assert np.all(dB(h[w <= ws]) < -rs + 0.1) + + assert N == 4 + xp_assert_close(Wn, 700.0, rtol=1e-15) + + xp_assert_equal(cheb1ord(1, 1.2, 1, 80, analog=True)[0], 17) + + def test_fs_param(self): + wp = 4800 + ws = 7200 + rp = 3 + rs = 60 + fs = 48000 + N, Wn = cheb1ord(wp, ws, rp, rs, False, fs=fs) + b, a = cheby1(N, rp, Wn, 'low', False, fs=fs) + w, h = freqz(b, a, fs=fs) + assert np.all(-rp - 0.1 < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs + 0.1) + + assert N == 8 + xp_assert_close(Wn, 4800.0, rtol=1e-15) + + def test_invalid_input(self): + with pytest.raises(ValueError) as exc_info: + cheb1ord(0.2, 0.3, 3, 2) + assert "gpass should be smaller than gstop" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + cheb1ord(0.2, 0.3, -1, 2) + assert "gpass should be larger than 0.0" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + cheb1ord(0.2, 0.3, 1, -2) + assert "gstop should be larger than 0.0" in str(exc_info.value) + + def test_ellip_cheb1(self): + # The purpose of the test is to compare to some known output from past + # scipy versions. The values to compare to are generated with scipy + # 1.9.1 (there is nothing special about this particular version though) + n, wn = cheb1ord([0.1, 0.6], [0.2, 0.5], 3, 60) + assert n == 7 + + n2, w2 = cheb2ord([0.1, 0.6], [0.2, 0.5], 3, 60) + assert not (wn == w2).all() + + def test_fs_validation(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + + with pytest.raises(ValueError, match="Sampling.*single scalar"): + cheb1ord(wp, ws, rp, rs, False, fs=np.array([10, 20])) + + +class TestCheb2ord: + + def test_lowpass(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + N, Wn = cheb2ord(wp, ws, rp, rs, False) + b, a = cheby2(N, rs, Wn, 'lp', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs + 0.1) + + assert N == 8 + xp_assert_close(Wn, 0.28647639976553163, rtol=1e-15) + + def test_highpass(self): + wp = 0.3 + ws = 0.2 + rp = 3 + rs = 70 + N, Wn = cheb2ord(wp, ws, rp, rs, False) + b, a = cheby2(N, rs, Wn, 'hp', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[wp <= w])) + assert np.all(dB(h[w <= ws]) < -rs + 0.1) + + assert N == 9 + xp_assert_close(Wn, 0.20697492182903282, rtol=1e-15) + + def test_bandpass(self): + wp = [0.2, 0.5] + ws = [0.1, 0.6] + rp = 3 + rs = 80 + N, Wn = cheb2ord(wp, ws, rp, rs, False) + b, a = cheby2(N, rs, Wn, 'bp', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[np.logical_and(wp[0] <= w, w <= wp[1])])) + assert np.all(dB(h[np.logical_or(w <= ws[0], ws[1] <= w)]) < -rs + 0.1) + + assert N == 9 + xp_assert_close(Wn, [0.14876937565923479, 0.59748447842351482], + rtol=1e-15) + + def test_bandstop(self): + wp = [0.1, 0.6] + ws = [0.2, 0.5] + rp = 3 + rs = 90 + N, Wn = cheb2ord(wp, ws, rp, rs, False) + b, a = cheby2(N, rs, Wn, 'bs', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[np.logical_or(w <= wp[0], wp[1] <= w)])) + assert np.all(dB(h[np.logical_and(ws[0] <= w, w <= ws[1])]) < -rs + 0.1) + + assert N == 10 + xp_assert_close(Wn, [0.19926249974781743, 0.50125246585567362], + rtol=1e-6) + + def test_analog(self): + wp = [20, 50] + ws = [10, 60] + rp = 3 + rs = 80 + N, Wn = cheb2ord(wp, ws, rp, rs, True) + b, a = cheby2(N, rs, Wn, 'bp', True) + w, h = freqs(b, a) + assert np.all(-rp - 0.1 < dB(h[np.logical_and(wp[0] <= w, w <= wp[1])])) + assert np.all(dB(h[np.logical_or(w <= ws[0], ws[1] <= w)]) < -rs + 0.1) + + assert N == 11 + xp_assert_close(Wn, [1.673740595370124e+01, 5.974641487254268e+01], + rtol=1e-15) + + def test_fs_param(self): + wp = 150 + ws = 100 + rp = 3 + rs = 70 + fs = 1000 + N, Wn = cheb2ord(wp, ws, rp, rs, False, fs=fs) + b, a = cheby2(N, rs, Wn, 'hp', False, fs=fs) + w, h = freqz(b, a, fs=fs) + assert np.all(-rp - 0.1 < dB(h[wp <= w])) + assert np.all(dB(h[w <= ws]) < -rs + 0.1) + + assert N == 9 + xp_assert_close(Wn, 103.4874609145164, rtol=1e-15) + + def test_invalid_input(self): + with pytest.raises(ValueError) as exc_info: + cheb2ord([0.1, 0.6], [0.2, 0.5], 3, 2) + assert "gpass should be smaller than gstop" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + cheb2ord([0.1, 0.6], [0.2, 0.5], -1, 2) + assert "gpass should be larger than 0.0" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + cheb2ord([0.1, 0.6], [0.2, 0.5], 1, -2) + assert "gstop should be larger than 0.0" in str(exc_info.value) + + def test_ellip_cheb2(self): + # The purpose of the test is to compare to some known output from past + # scipy versions. The values to compare to are generated with scipy + # 1.9.1 (there is nothing special about this particular version though) + n, wn = cheb2ord([0.1, 0.6], [0.2, 0.5], 3, 60) + assert n == 7 + + n1, w1 = cheb1ord([0.1, 0.6], [0.2, 0.5], 3, 60) + assert not (wn == w1).all() + + def test_fs_validation(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + + with pytest.raises(ValueError, match="Sampling.*single scalar"): + cheb2ord(wp, ws, rp, rs, False, fs=np.array([10, 20])) + + +class TestEllipord: + + def test_lowpass(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + N, Wn = ellipord(wp, ws, rp, rs, False) + b, a = ellip(N, rp, rs, Wn, 'lp', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs + 0.1) + + assert N == 5 + xp_assert_close(Wn, 0.2, rtol=1e-15) + + def test_lowpass_1000dB(self): + # failed when ellipkm1 wasn't used in ellipord and ellipap + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 1000 + N, Wn = ellipord(wp, ws, rp, rs, False) + sos = ellip(N, rp, rs, Wn, 'lp', False, output='sos') + w, h = freqz_sos(sos) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[w <= wp])) + assert np.all(dB(h[ws <= w]) < -rs + 0.1) + + def test_highpass(self): + wp = 0.3 + ws = 0.2 + rp = 3 + rs = 70 + N, Wn = ellipord(wp, ws, rp, rs, False) + b, a = ellip(N, rp, rs, Wn, 'hp', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[wp <= w])) + assert np.all(dB(h[w <= ws]) < -rs + 0.1) + + assert N == 6 + xp_assert_close(Wn, 0.3, rtol=1e-15) + + def test_bandpass(self): + wp = [0.2, 0.5] + ws = [0.1, 0.6] + rp = 3 + rs = 80 + N, Wn = ellipord(wp, ws, rp, rs, False) + b, a = ellip(N, rp, rs, Wn, 'bp', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[np.logical_and(wp[0] <= w, w <= wp[1])])) + assert np.all(dB(h[np.logical_or(w <= ws[0], ws[1] <= w)]) < -rs + 0.1) + + assert N == 6 + xp_assert_close(Wn, [0.2, 0.5], rtol=1e-15) + + def test_bandstop(self): + wp = [0.1, 0.6] + ws = [0.2, 0.5] + rp = 3 + rs = 90 + N, Wn = ellipord(wp, ws, rp, rs, False) + b, a = ellip(N, rp, rs, Wn, 'bs', False) + w, h = freqz(b, a) + w /= np.pi + assert np.all(-rp - 0.1 < dB(h[np.logical_or(w <= wp[0], wp[1] <= w)])) + assert np.all(dB(h[np.logical_and(ws[0] <= w, w <= ws[1])]) < -rs + 0.1) + + assert N == 7 + xp_assert_close(Wn, [0.14758232794342988, 0.6], rtol=1e-5) + + def test_analog(self): + wp = [1000, 6000] + ws = [2000, 5000] + rp = 3 + rs = 90 + N, Wn = ellipord(wp, ws, rp, rs, True) + b, a = ellip(N, rp, rs, Wn, 'bs', True) + w, h = freqs(b, a) + assert np.all(-rp - 0.1 < dB(h[np.logical_or(w <= wp[0], wp[1] <= w)])) + assert np.all(dB(h[np.logical_and(ws[0] <= w, w <= ws[1])]) < -rs + 0.1) + + assert N == 8 + xp_assert_close(Wn, [1666.6666, 6000]) + + assert ellipord(1, 1.2, 1, 80, analog=True)[0] == 9 + + def test_fs_param(self): + wp = [400, 2400] + ws = [800, 2000] + rp = 3 + rs = 90 + fs = 8000 + N, Wn = ellipord(wp, ws, rp, rs, False, fs=fs) + b, a = ellip(N, rp, rs, Wn, 'bs', False, fs=fs) + w, h = freqz(b, a, fs=fs) + assert np.all(-rp - 0.1 < dB(h[np.logical_or(w <= wp[0], wp[1] <= w)])) + assert np.all(dB(h[np.logical_and(ws[0] <= w, w <= ws[1])]) < -rs + 0.1) + + assert N == 7 + xp_assert_close(Wn, [590.3293117737195, 2400], rtol=1e-5) + + def test_invalid_input(self): + with pytest.raises(ValueError) as exc_info: + ellipord(0.2, 0.5, 3, 2) + assert "gpass should be smaller than gstop" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + ellipord(0.2, 0.5, -1, 2) + assert "gpass should be larger than 0.0" in str(exc_info.value) + + with pytest.raises(ValueError) as exc_info: + ellipord(0.2, 0.5, 1, -2) + assert "gstop should be larger than 0.0" in str(exc_info.value) + + def test_ellip_butter(self): + # The purpose of the test is to compare to some known output from past + # scipy versions. The values to compare to are generated with scipy + # 1.9.1 (there is nothing special about this particular version though) + n, wn = ellipord([0.1, 0.6], [0.2, 0.5], 3, 60) + assert n == 5 + + def test_fs_validation(self): + wp = 0.2 + ws = 0.3 + rp = 3 + rs = 60 + + with pytest.raises(ValueError, match="Sampling.*single scalar"): + ellipord(wp, ws, rp, rs, False, fs=np.array([10, 20])) + + +class TestBessel: + + def test_degenerate(self): + for norm in ('delay', 'phase', 'mag'): + # 0-order filter is just a passthrough + b, a = bessel(0, 1, analog=True, norm=norm) + xp_assert_equal(b, np.asarray([1.0])) + xp_assert_equal(a, np.asarray([1.0])) + + # 1-order filter is same for all types + b, a = bessel(1, 1, analog=True, norm=norm) + xp_assert_close(b, np.asarray([1.0]), rtol=1e-15) + xp_assert_close(a, np.asarray([1.0, 1]), rtol=1e-15) + + z, p, k = bessel(1, 0.3, analog=True, output='zpk', norm=norm) + xp_assert_equal(z, np.asarray([])) + xp_assert_close(p, np.asarray([-0.3+0j]), rtol=1e-14) + xp_assert_close(k, 0.3, rtol=1e-14) + + def test_high_order(self): + # high even order, 'phase' + z, p, k = bessel(24, 100, analog=True, output='zpk') + z2 = [] + p2 = [ + -9.055312334014323e+01 + 4.844005815403969e+00j, + -8.983105162681878e+01 + 1.454056170018573e+01j, + -8.837357994162065e+01 + 2.426335240122282e+01j, + -8.615278316179575e+01 + 3.403202098404543e+01j, + -8.312326467067703e+01 + 4.386985940217900e+01j, + -7.921695461084202e+01 + 5.380628489700191e+01j, + -7.433392285433246e+01 + 6.388084216250878e+01j, + -6.832565803501586e+01 + 7.415032695116071e+01j, + -6.096221567378025e+01 + 8.470292433074425e+01j, + -5.185914574820616e+01 + 9.569048385258847e+01j, + -4.027853855197555e+01 + 1.074195196518679e+02j, + -2.433481337524861e+01 + 1.207298683731973e+02j, + ] + k2 = 9.999999999999989e+47 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(np.union1d(p2, np.conj(p2)), key=np.imag)) + xp_assert_close(k, k2, rtol=1e-14) + + # high odd order, 'phase' + z, p, k = bessel(23, 1000, analog=True, output='zpk') + z2 = [] + p2 = [ + -2.497697202208956e+02 + 1.202813187870698e+03j, + -4.126986617510172e+02 + 1.065328794475509e+03j, + -5.304922463809596e+02 + 9.439760364018479e+02j, + -9.027564978975828e+02 + 1.010534334242318e+02j, + -8.909283244406079e+02 + 2.023024699647598e+02j, + -8.709469394347836e+02 + 3.039581994804637e+02j, + -8.423805948131370e+02 + 4.062657947488952e+02j, + -8.045561642249877e+02 + 5.095305912401127e+02j, + -7.564660146766259e+02 + 6.141594859516342e+02j, + -6.965966033906477e+02 + 7.207341374730186e+02j, + -6.225903228776276e+02 + 8.301558302815096e+02j, + -9.066732476324988e+02] + k2 = 9.999999999999983e+68 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(np.union1d(p2, np.conj(p2)), key=np.imag)) + xp_assert_close(k, k2, rtol=1e-14) + + # high even order, 'delay' (Orchard 1965 "The Roots of the + # Maximally Flat-Delay Polynomials" Table 1) + z, p, k = bessel(31, 1, analog=True, output='zpk', norm='delay') + p2 = [-20.876706, + -20.826543 + 1.735732j, + -20.675502 + 3.473320j, + -20.421895 + 5.214702j, + -20.062802 + 6.961982j, + -19.593895 + 8.717546j, + -19.009148 + 10.484195j, + -18.300400 + 12.265351j, + -17.456663 + 14.065350j, + -16.463032 + 15.889910j, + -15.298849 + 17.746914j, + -13.934466 + 19.647827j, + -12.324914 + 21.610519j, + -10.395893 + 23.665701j, + - 8.005600 + 25.875019j, + - 4.792045 + 28.406037j, + ] + xp_assert_close(sorted(p, key=np.imag), + sorted(np.union1d(p2, np.conj(p2)), key=np.imag)) + + # high odd order, 'delay' + z, p, k = bessel(30, 1, analog=True, output='zpk', norm='delay') + p2 = [-20.201029 + 0.867750j, + -20.097257 + 2.604235j, + -19.888485 + 4.343721j, + -19.572188 + 6.088363j, + -19.144380 + 7.840570j, + -18.599342 + 9.603147j, + -17.929195 + 11.379494j, + -17.123228 + 13.173901j, + -16.166808 + 14.992008j, + -15.039580 + 16.841580j, + -13.712245 + 18.733902j, + -12.140295 + 20.686563j, + -10.250119 + 22.729808j, + - 7.901170 + 24.924391j, + - 4.734679 + 27.435615j, + ] + xp_assert_close(sorted(p, key=np.imag), + sorted(np.union1d(p2, np.conj(p2)), key=np.imag)) + + def test_refs(self): + # Compare to http://www.crbond.com/papers/bsf2.pdf + # "Delay Normalized Bessel Polynomial Coefficients" + bond_b = np.asarray([10395.0]) + bond_a = np.asarray([1.0, 21, 210, 1260, 4725, 10395, 10395]) + b, a = bessel(6, 1, norm='delay', analog=True) + xp_assert_close(b, bond_b) + xp_assert_close(a, bond_a) + + # "Delay Normalized Bessel Pole Locations" + bond_poles = { + 1: [-1.0000000000], + 2: [-1.5000000000 + 0.8660254038j], + 3: [-1.8389073227 + 1.7543809598j, -2.3221853546], + 4: [-2.1037893972 + 2.6574180419j, -2.8962106028 + 0.8672341289j], + 5: [-2.3246743032 + 3.5710229203j, -3.3519563992 + 1.7426614162j, + -3.6467385953], + 6: [-2.5159322478 + 4.4926729537j, -3.7357083563 + 2.6262723114j, + -4.2483593959 + 0.8675096732j], + 7: [-2.6856768789 + 5.4206941307j, -4.0701391636 + 3.5171740477j, + -4.7582905282 + 1.7392860611j, -4.9717868585], + 8: [-2.8389839489 + 6.3539112986j, -4.3682892172 + 4.4144425005j, + -5.2048407906 + 2.6161751526j, -5.5878860433 + 0.8676144454j], + 9: [-2.9792607982 + 7.2914636883j, -4.6384398872 + 5.3172716754j, + -5.6044218195 + 3.4981569179j, -6.1293679043 + 1.7378483835j, + -6.2970191817], + 10: [-3.1089162336 + 8.2326994591j, -4.8862195669 + 6.2249854825j, + -5.9675283286 + 4.3849471889j, -6.6152909655 + 2.6115679208j, + -6.9220449054 + 0.8676651955j] + } + + for N in range(1, 11): + p1 = np.sort(bond_poles[N]) + p2 = np.sort(np.concatenate(_cplxreal(besselap(N, 'delay')[1]))) + assert_array_almost_equal(p1, p2, decimal=10) + + # "Frequency Normalized Bessel Pole Locations" + bond_poles = { + 1: [-1.0000000000], + 2: [-1.1016013306 + 0.6360098248j], + 3: [-1.0474091610 + 0.9992644363j, -1.3226757999], + 4: [-0.9952087644 + 1.2571057395j, -1.3700678306 + 0.4102497175j], + 5: [-0.9576765486 + 1.4711243207j, -1.3808773259 + 0.7179095876j, + -1.5023162714], + 6: [-0.9306565229 + 1.6618632689j, -1.3818580976 + 0.9714718907j, + -1.5714904036 + 0.3208963742j], + 7: [-0.9098677806 + 1.8364513530j, -1.3789032168 + 1.1915667778j, + -1.6120387662 + 0.5892445069j, -1.6843681793], + 8: [-0.8928697188 + 1.9983258436j, -1.3738412176 + 1.3883565759j, + -1.6369394181 + 0.8227956251j, -1.7574084004 + 0.2728675751j], + 9: [-0.8783992762 + 2.1498005243j, -1.3675883098 + 1.5677337122j, + -1.6523964846 + 1.0313895670j, -1.8071705350 + 0.5123837306j, + -1.8566005012], + 10: [-0.8657569017 + 2.2926048310j, -1.3606922784 + 1.7335057427j, + -1.6618102414 + 1.2211002186j, -1.8421962445 + 0.7272575978j, + -1.9276196914 + 0.2416234710j] + } + + for N in range(1, 11): + p1 = np.sort(bond_poles[N]) + p2 = np.sort(np.concatenate(_cplxreal(besselap(N, 'mag')[1]))) + assert_array_almost_equal(p1, p2, decimal=10) + + # Compare to https://www.ranecommercial.com/legacy/note147.html + # "Table 1 - Bessel Crossovers of Second, Third, and Fourth-Order" + a = np.asarray([1, 1, 1/3]) + b2, a2 = bessel(2, 1, norm='delay', analog=True) + xp_assert_close(a[::-1], a2/b2) + + a = np.asarray([1, 1, 2/5, 1/15]) + b2, a2 = bessel(3, 1, norm='delay', analog=True) + xp_assert_close(a[::-1], a2/b2) + + a = np.asarray([1, 1, 9/21, 2/21, 1/105]) + b2, a2 = bessel(4, 1, norm='delay', analog=True) + xp_assert_close(a[::-1], a2/b2) + + a = np.asarray([1, np.sqrt(3), 1]) + b2, a2 = bessel(2, 1, norm='phase', analog=True) + xp_assert_close(a[::-1], a2/b2) + + # TODO: Why so inaccurate? Is reference flawed? + a = np.asarray([1, 2.481, 2.463, 1.018]) + b2, a2 = bessel(3, 1, norm='phase', analog=True) + assert_array_almost_equal(a[::-1], a2/b2, decimal=1) + + # TODO: Why so inaccurate? Is reference flawed? + a = np.asarray([1, 3.240, 4.5, 3.240, 1.050]) + b2, a2 = bessel(4, 1, norm='phase', analog=True) + assert_array_almost_equal(a[::-1], a2/b2, decimal=1) + + # Table of -3 dB factors: + N, scale = 2, np.asarray([1.272, 1.272], dtype=np.complex128) + scale2 = besselap(N, 'mag')[1] / besselap(N, 'phase')[1] + assert_array_almost_equal(scale2, scale, decimal=3) + + # TODO: Why so inaccurate? Is reference flawed? + N, scale = 3, np.asarray([1.413, 1.413, 1.413], dtype=np.complex128) + scale2 = besselap(N, 'mag')[1] / besselap(N, 'phase')[1] + assert_array_almost_equal(scale2, scale, decimal=2) + + # TODO: Why so inaccurate? Is reference flawed? + N, scale = 4, np.asarray([1.533]*4, dtype=np.complex128) + scale2 = besselap(N, 'mag')[1] / besselap(N, 'phase')[1] + assert_array_almost_equal(scale, scale2, decimal=1) + + def test_hardcoded(self): + # Compare to values from original hardcoded implementation + originals = { + 0: [], + 1: [-1], + 2: [-.8660254037844386467637229 + .4999999999999999999999996j], + 3: [-.9416000265332067855971980, + -.7456403858480766441810907 + .7113666249728352680992154j], + 4: [-.6572111716718829545787788 + .8301614350048733772399715j, + -.9047587967882449459642624 + .2709187330038746636700926j], + 5: [-.9264420773877602247196260, + -.8515536193688395541722677 + .4427174639443327209850002j, + -.5905759446119191779319432 + .9072067564574549539291747j], + 6: [-.9093906830472271808050953 + .1856964396793046769246397j, + -.7996541858328288520243325 + .5621717346937317988594118j, + -.5385526816693109683073792 + .9616876881954277199245657j], + 7: [-.9194871556490290014311619, + -.8800029341523374639772340 + .3216652762307739398381830j, + -.7527355434093214462291616 + .6504696305522550699212995j, + -.4966917256672316755024763 + 1.002508508454420401230220j], + 8: [-.9096831546652910216327629 + .1412437976671422927888150j, + -.8473250802359334320103023 + .4259017538272934994996429j, + -.7111381808485399250796172 + .7186517314108401705762571j, + -.4621740412532122027072175 + 1.034388681126901058116589j], + 9: [-.9154957797499037686769223, + -.8911217017079759323183848 + .2526580934582164192308115j, + -.8148021112269012975514135 + .5085815689631499483745341j, + -.6743622686854761980403401 + .7730546212691183706919682j, + -.4331415561553618854685942 + 1.060073670135929666774323j], + 10: [-.9091347320900502436826431 + .1139583137335511169927714j, + -.8688459641284764527921864 + .3430008233766309973110589j, + -.7837694413101441082655890 + .5759147538499947070009852j, + -.6417513866988316136190854 + .8175836167191017226233947j, + -.4083220732868861566219785 + 1.081274842819124562037210j], + 11: [-.9129067244518981934637318, + -.8963656705721166099815744 + .2080480375071031919692341j, + -.8453044014712962954184557 + .4178696917801248292797448j, + -.7546938934722303128102142 + .6319150050721846494520941j, + -.6126871554915194054182909 + .8547813893314764631518509j, + -.3868149510055090879155425 + 1.099117466763120928733632j], + 12: [-.9084478234140682638817772 + 95506365213450398415258360e-27j, + -.8802534342016826507901575 + .2871779503524226723615457j, + -.8217296939939077285792834 + .4810212115100676440620548j, + -.7276681615395159454547013 + .6792961178764694160048987j, + -.5866369321861477207528215 + .8863772751320727026622149j, + -.3679640085526312839425808 + 1.114373575641546257595657j], + 13: [-.9110914665984182781070663, + -.8991314665475196220910718 + .1768342956161043620980863j, + -.8625094198260548711573628 + .3547413731172988997754038j, + -.7987460692470972510394686 + .5350752120696801938272504j, + -.7026234675721275653944062 + .7199611890171304131266374j, + -.5631559842430199266325818 + .9135900338325109684927731j, + -.3512792323389821669401925 + 1.127591548317705678613239j], + 14: [-.9077932138396487614720659 + 82196399419401501888968130e-27j, + -.8869506674916445312089167 + .2470079178765333183201435j, + -.8441199160909851197897667 + .4131653825102692595237260j, + -.7766591387063623897344648 + .5819170677377608590492434j, + -.6794256425119233117869491 + .7552857305042033418417492j, + -.5418766775112297376541293 + .9373043683516919569183099j, + -.3363868224902037330610040 + 1.139172297839859991370924j], + 15: [-.9097482363849064167228581, + -.9006981694176978324932918 + .1537681197278439351298882j, + -.8731264620834984978337843 + .3082352470564267657715883j, + -.8256631452587146506294553 + .4642348752734325631275134j, + -.7556027168970728127850416 + .6229396358758267198938604j, + -.6579196593110998676999362 + .7862895503722515897065645j, + -.5224954069658330616875186 + .9581787261092526478889345j, + -.3229963059766444287113517 + 1.149416154583629539665297j], + 16: [-.9072099595087001356491337 + 72142113041117326028823950e-27j, + -.8911723070323647674780132 + .2167089659900576449410059j, + -.8584264231521330481755780 + .3621697271802065647661080j, + -.8074790293236003885306146 + .5092933751171800179676218j, + -.7356166304713115980927279 + .6591950877860393745845254j, + -.6379502514039066715773828 + .8137453537108761895522580j, + -.5047606444424766743309967 + .9767137477799090692947061j, + -.3108782755645387813283867 + 1.158552841199330479412225j], + 17: [-.9087141161336397432860029, + -.9016273850787285964692844 + .1360267995173024591237303j, + -.8801100704438627158492165 + .2725347156478803885651973j, + -.8433414495836129204455491 + .4100759282910021624185986j, + -.7897644147799708220288138 + .5493724405281088674296232j, + -.7166893842372349049842743 + .6914936286393609433305754j, + -.6193710717342144521602448 + .8382497252826992979368621j, + -.4884629337672704194973683 + .9932971956316781632345466j, + -.2998489459990082015466971 + 1.166761272925668786676672j], + 18: [-.9067004324162775554189031 + 64279241063930693839360680e-27j, + -.8939764278132455733032155 + .1930374640894758606940586j, + -.8681095503628830078317207 + .3224204925163257604931634j, + -.8281885016242836608829018 + .4529385697815916950149364j, + -.7726285030739558780127746 + .5852778162086640620016316j, + -.6987821445005273020051878 + .7204696509726630531663123j, + -.6020482668090644386627299 + .8602708961893664447167418j, + -.4734268069916151511140032 + 1.008234300314801077034158j, + -.2897592029880489845789953 + 1.174183010600059128532230j], + 19: [-.9078934217899404528985092, + -.9021937639390660668922536 + .1219568381872026517578164j, + -.8849290585034385274001112 + .2442590757549818229026280j, + -.8555768765618421591093993 + .3672925896399872304734923j, + -.8131725551578197705476160 + .4915365035562459055630005j, + -.7561260971541629355231897 + .6176483917970178919174173j, + -.6818424412912442033411634 + .7466272357947761283262338j, + -.5858613321217832644813602 + .8801817131014566284786759j, + -.4595043449730988600785456 + 1.021768776912671221830298j, + -.2804866851439370027628724 + 1.180931628453291873626003j], + 20: [-.9062570115576771146523497 + 57961780277849516990208850e-27j, + -.8959150941925768608568248 + .1740317175918705058595844j, + -.8749560316673332850673214 + .2905559296567908031706902j, + -.8427907479956670633544106 + .4078917326291934082132821j, + -.7984251191290606875799876 + .5264942388817132427317659j, + -.7402780309646768991232610 + .6469975237605228320268752j, + -.6658120544829934193890626 + .7703721701100763015154510j, + -.5707026806915714094398061 + .8982829066468255593407161j, + -.4465700698205149555701841 + 1.034097702560842962315411j, + -.2719299580251652601727704 + 1.187099379810885886139638j], + 21: [-.9072262653142957028884077, + -.9025428073192696303995083 + .1105252572789856480992275j, + -.8883808106664449854431605 + .2213069215084350419975358j, + -.8643915813643204553970169 + .3326258512522187083009453j, + -.8299435470674444100273463 + .4448177739407956609694059j, + -.7840287980408341576100581 + .5583186348022854707564856j, + -.7250839687106612822281339 + .6737426063024382240549898j, + -.6506315378609463397807996 + .7920349342629491368548074j, + -.5564766488918562465935297 + .9148198405846724121600860j, + -.4345168906815271799687308 + 1.045382255856986531461592j, + -.2640041595834031147954813 + 1.192762031948052470183960j], + 22: [-.9058702269930872551848625 + 52774908289999045189007100e-27j, + -.8972983138153530955952835 + .1584351912289865608659759j, + -.8799661455640176154025352 + .2644363039201535049656450j, + -.8534754036851687233084587 + .3710389319482319823405321j, + -.8171682088462720394344996 + .4785619492202780899653575j, + -.7700332930556816872932937 + .5874255426351153211965601j, + -.7105305456418785989070935 + .6982266265924524000098548j, + -.6362427683267827226840153 + .8118875040246347267248508j, + -.5430983056306302779658129 + .9299947824439872998916657j, + -.4232528745642628461715044 + 1.055755605227545931204656j, + -.2566376987939318038016012 + 1.197982433555213008346532j], + 23: [-.9066732476324988168207439, + -.9027564979912504609412993 + .1010534335314045013252480j, + -.8909283242471251458653994 + .2023024699381223418195228j, + -.8709469395587416239596874 + .3039581993950041588888925j, + -.8423805948021127057054288 + .4062657948237602726779246j, + -.8045561642053176205623187 + .5095305912227258268309528j, + -.7564660146829880581478138 + .6141594859476032127216463j, + -.6965966033912705387505040 + .7207341374753046970247055j, + -.6225903228771341778273152 + .8301558302812980678845563j, + -.5304922463810191698502226 + .9439760364018300083750242j, + -.4126986617510148836149955 + 1.065328794475513585531053j, + -.2497697202208956030229911 + 1.202813187870697831365338j], + 24: [-.9055312363372773709269407 + 48440066540478700874836350e-27j, + -.8983105104397872954053307 + .1454056133873610120105857j, + -.8837358034555706623131950 + .2426335234401383076544239j, + -.8615278304016353651120610 + .3403202112618624773397257j, + -.8312326466813240652679563 + .4386985933597305434577492j, + -.7921695462343492518845446 + .5380628490968016700338001j, + -.7433392285088529449175873 + .6388084216222567930378296j, + -.6832565803536521302816011 + .7415032695091650806797753j, + -.6096221567378335562589532 + .8470292433077202380020454j, + -.5185914574820317343536707 + .9569048385259054576937721j, + -.4027853855197518014786978 + 1.074195196518674765143729j, + -.2433481337524869675825448 + 1.207298683731972524975429j], + 25: [-.9062073871811708652496104, + -.9028833390228020537142561 + 93077131185102967450643820e-27j, + -.8928551459883548836774529 + .1863068969804300712287138j, + -.8759497989677857803656239 + .2798521321771408719327250j, + -.8518616886554019782346493 + .3738977875907595009446142j, + -.8201226043936880253962552 + .4686668574656966589020580j, + -.7800496278186497225905443 + .5644441210349710332887354j, + -.7306549271849967721596735 + .6616149647357748681460822j, + -.6704827128029559528610523 + .7607348858167839877987008j, + -.5972898661335557242320528 + .8626676330388028512598538j, + -.5073362861078468845461362 + .9689006305344868494672405j, + -.3934529878191079606023847 + 1.082433927173831581956863j, + -.2373280669322028974199184 + 1.211476658382565356579418j], + } + for N in originals: + p1 = sorted(np.union1d(originals[N], + np.conj(originals[N])), key=np.imag) + p2 = sorted(besselap(N)[1], key=np.imag) + xp_assert_close(p1, + p2, rtol=1e-14, check_dtype=False) + + def test_norm_phase(self): + # Test some orders and frequencies and see that they have the right + # phase at w0 + for N in (1, 2, 3, 4, 5, 51, 72): + for w0 in (1, 100): + b, a = bessel(N, w0, analog=True, norm='phase') + w = np.linspace(0, w0, 100) + w, h = freqs(b, a, w) + phase = np.unwrap(np.angle(h)) + xp_assert_close(phase[[0, -1]], (0, -N*pi/4), rtol=1e-1) + + def test_norm_mag(self): + # Test some orders and frequencies and see that they have the right + # mag at w0 + for N in (1, 2, 3, 4, 5, 51, 72): + for w0 in (1, 100): + b, a = bessel(N, w0, analog=True, norm='mag') + w = (0, w0) + w, h = freqs(b, a, w) + mag = abs(h) + xp_assert_close(mag, (1, 1/np.sqrt(2))) + + def test_norm_delay(self): + # Test some orders and frequencies and see that they have the right + # delay at DC + for N in (1, 2, 3, 4, 5, 51, 72): + for w0 in (1, 100): + b, a = bessel(N, w0, analog=True, norm='delay') + w = np.linspace(0, 10*w0, 1000) + w, h = freqs(b, a, w) + delay = -np.diff(np.unwrap(np.angle(h)))/np.diff(w) + xp_assert_close(delay[0], 1/w0, rtol=1e-4) + + def test_norm_factor(self): + mpmath_values = { + 1: 1.0, 2: 1.361654128716130520, 3: 1.755672368681210649, + 4: 2.113917674904215843, 5: 2.427410702152628137, + 6: 2.703395061202921876, 7: 2.951722147038722771, + 8: 3.179617237510651330, 9: 3.391693138911660101, + 10: 3.590980594569163482, 11: 3.779607416439620092, + 12: 3.959150821144285315, 13: 4.130825499383535980, + 14: 4.295593409533637564, 15: 4.454233021624377494, + 16: 4.607385465472647917, 17: 4.755586548961147727, + 18: 4.899289677284488007, 19: 5.038882681488207605, + 20: 5.174700441742707423, 21: 5.307034531360917274, + 22: 5.436140703250035999, 23: 5.562244783787878196, + 24: 5.685547371295963521, 25: 5.806227623775418541, + 50: 8.268963160013226298, 51: 8.352374541546012058, + } + for N in mpmath_values: + z, p, k = besselap(N, 'delay') + xp_assert_close(mpmath_values[N], _norm_factor(p, k), rtol=1e-13) + + def test_bessel_poly(self): + xp_assert_equal(_bessel_poly(5), [945, 945, 420, 105, 15, 1]) + xp_assert_equal(_bessel_poly(4, True), [1, 10, 45, 105, 105]) + + def test_bessel_zeros(self): + xp_assert_equal(_bessel_zeros(0), []) + + def test_invalid(self): + assert_raises(ValueError, besselap, 5, 'nonsense') + assert_raises(ValueError, besselap, -5) + assert_raises(ValueError, besselap, 3.2) + assert_raises(ValueError, _bessel_poly, -3) + assert_raises(ValueError, _bessel_poly, 3.3) + + @pytest.mark.fail_slow(10) + def test_fs_param(self): + for norm in ('phase', 'mag', 'delay'): + for fs in (900, 900.1, 1234.567): + for N in (0, 1, 2, 3, 10): + for fc in (100, 100.1, 432.12345): + for btype in ('lp', 'hp'): + ba1 = bessel(N, fc, btype, norm=norm, fs=fs) + ba2 = bessel(N, fc/(fs/2), btype, norm=norm) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + for fc in ((100, 200), (100.1, 200.2), (321.123, 432.123)): + for btype in ('bp', 'bs'): + ba1 = bessel(N, fc, btype, norm=norm, fs=fs) + for seq in (list, tuple, array): + fcnorm = seq([f/(fs/2) for f in fc]) + ba2 = bessel(N, fcnorm, btype, norm=norm) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + + +class TestButter: + + def test_degenerate(self): + # 0-order filter is just a passthrough + b, a = butter(0, 1, analog=True) + xp_assert_equal(b, np.asarray([1.0])) + xp_assert_equal(a, np.asarray([1.0])) + + # 1-order filter is same for all types + b, a = butter(1, 1, analog=True) + assert_array_almost_equal(b, [1]) + assert_array_almost_equal(a, [1, 1]) + + z, p, k = butter(1, 0.3, output='zpk') + xp_assert_equal(z, np.asarray([-1.0])) + xp_assert_close(p, [3.249196962329063e-01 + 0j], rtol=1e-14) + xp_assert_close(k, 3.375401518835469e-01, rtol=1e-14) + + def test_basic(self): + # analog s-plane + for N in range(25): + wn = 0.01 + z, p, k = butter(N, wn, 'low', analog=True, output='zpk') + assert_array_almost_equal([], z) + assert len(p) == N + # All poles should be at distance wn from origin + assert_array_almost_equal(abs(p), np.asarray(wn)) + assert all(np.real(p) <= 0) # No poles in right half of S-plane + assert_array_almost_equal(wn**N, k) + + # digital z-plane + for N in range(25): + wn = 0.01 + z, p, k = butter(N, wn, 'high', analog=False, output='zpk') + xp_assert_equal(np.ones(N), z) # All zeros exactly at DC + assert all(np.abs(p) <= 1) # No poles outside unit circle + + b1, a1 = butter(2, 1, analog=True) + assert_array_almost_equal(b1, [1]) + assert_array_almost_equal(a1, [1, np.sqrt(2), 1]) + + b2, a2 = butter(5, 1, analog=True) + assert_array_almost_equal(b2, [1]) + assert_array_almost_equal(a2, [1, 3.2361, 5.2361, + 5.2361, 3.2361, 1], decimal=4) + + b3, a3 = butter(10, 1, analog=True) + assert_array_almost_equal(b3, [1]) + assert_array_almost_equal(a3, [1, 6.3925, 20.4317, 42.8021, 64.8824, + 74.2334, 64.8824, 42.8021, 20.4317, + 6.3925, 1], decimal=4) + + b2, a2 = butter(19, 1.0441379169150726, analog=True) + assert_array_almost_equal(b2, [2.2720], decimal=4) + assert_array_almost_equal(a2, 1.0e+004 * np.array([ + 0.0001, 0.0013, 0.0080, 0.0335, 0.1045, 0.2570, + 0.5164, 0.8669, 1.2338, 1.5010, 1.5672, 1.4044, + 1.0759, 0.6986, 0.3791, 0.1681, 0.0588, 0.0153, + 0.0026, 0.0002]), decimal=0) + + b, a = butter(5, 0.4) + assert_array_almost_equal(b, [0.0219, 0.1097, 0.2194, + 0.2194, 0.1097, 0.0219], decimal=4) + assert_array_almost_equal(a, [1.0000, -0.9853, 0.9738, + -0.3864, 0.1112, -0.0113], decimal=4) + + def test_highpass(self): + # highpass, high even order + z, p, k = butter(28, 0.43, 'high', output='zpk') + z2 = np.ones(28) + p2 = [ + 2.068257195514592e-01 + 9.238294351481734e-01j, + 2.068257195514592e-01 - 9.238294351481734e-01j, + 1.874933103892023e-01 + 8.269455076775277e-01j, + 1.874933103892023e-01 - 8.269455076775277e-01j, + 1.717435567330153e-01 + 7.383078571194629e-01j, + 1.717435567330153e-01 - 7.383078571194629e-01j, + 1.588266870755982e-01 + 6.564623730651094e-01j, + 1.588266870755982e-01 - 6.564623730651094e-01j, + 1.481881532502603e-01 + 5.802343458081779e-01j, + 1.481881532502603e-01 - 5.802343458081779e-01j, + 1.394122576319697e-01 + 5.086609000582009e-01j, + 1.394122576319697e-01 - 5.086609000582009e-01j, + 1.321840881809715e-01 + 4.409411734716436e-01j, + 1.321840881809715e-01 - 4.409411734716436e-01j, + 1.262633413354405e-01 + 3.763990035551881e-01j, + 1.262633413354405e-01 - 3.763990035551881e-01j, + 1.214660449478046e-01 + 3.144545234797277e-01j, + 1.214660449478046e-01 - 3.144545234797277e-01j, + 1.104868766650320e-01 + 2.771505404367791e-02j, + 1.104868766650320e-01 - 2.771505404367791e-02j, + 1.111768629525075e-01 + 8.331369153155753e-02j, + 1.111768629525075e-01 - 8.331369153155753e-02j, + 1.125740630842972e-01 + 1.394219509611784e-01j, + 1.125740630842972e-01 - 1.394219509611784e-01j, + 1.147138487992747e-01 + 1.963932363793666e-01j, + 1.147138487992747e-01 - 1.963932363793666e-01j, + 1.176516491045901e-01 + 2.546021573417188e-01j, + 1.176516491045901e-01 - 2.546021573417188e-01j, + ] + k2 = 1.446671081817286e-06 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-7) + xp_assert_close(k, k2, rtol=1e-10) + + # highpass, high odd order + z, p, k = butter(27, 0.56, 'high', output='zpk') + z2 = np.ones(27) + p2 = [ + -1.772572785680147e-01 + 9.276431102995948e-01j, + -1.772572785680147e-01 - 9.276431102995948e-01j, + -1.600766565322114e-01 + 8.264026279893268e-01j, + -1.600766565322114e-01 - 8.264026279893268e-01j, + -1.461948419016121e-01 + 7.341841939120078e-01j, + -1.461948419016121e-01 - 7.341841939120078e-01j, + -1.348975284762046e-01 + 6.493235066053785e-01j, + -1.348975284762046e-01 - 6.493235066053785e-01j, + -1.256628210712206e-01 + 5.704921366889227e-01j, + -1.256628210712206e-01 - 5.704921366889227e-01j, + -1.181038235962314e-01 + 4.966120551231630e-01j, + -1.181038235962314e-01 - 4.966120551231630e-01j, + -1.119304913239356e-01 + 4.267938916403775e-01j, + -1.119304913239356e-01 - 4.267938916403775e-01j, + -1.069237739782691e-01 + 3.602914879527338e-01j, + -1.069237739782691e-01 - 3.602914879527338e-01j, + -1.029178030691416e-01 + 2.964677964142126e-01j, + -1.029178030691416e-01 - 2.964677964142126e-01j, + -9.978747500816100e-02 + 2.347687643085738e-01j, + -9.978747500816100e-02 - 2.347687643085738e-01j, + -9.743974496324025e-02 + 1.747028739092479e-01j, + -9.743974496324025e-02 - 1.747028739092479e-01j, + -9.580754551625957e-02 + 1.158246860771989e-01j, + -9.580754551625957e-02 - 1.158246860771989e-01j, + -9.484562207782568e-02 + 5.772118357151691e-02j, + -9.484562207782568e-02 - 5.772118357151691e-02j, + -9.452783117928215e-02 + ] + k2 = 9.585686688851069e-09 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-8) + xp_assert_close(k, k2) + + def test_bandpass(self): + z, p, k = butter(8, [0.25, 0.33], 'band', output='zpk') + z2 = [1, 1, 1, 1, 1, 1, 1, 1, + -1, -1, -1, -1, -1, -1, -1, -1] + p2 = [ + 4.979909925436156e-01 + 8.367609424799387e-01j, + 4.979909925436156e-01 - 8.367609424799387e-01j, + 4.913338722555539e-01 + 7.866774509868817e-01j, + 4.913338722555539e-01 - 7.866774509868817e-01j, + 5.035229361778706e-01 + 7.401147376726750e-01j, + 5.035229361778706e-01 - 7.401147376726750e-01j, + 5.307617160406101e-01 + 7.029184459442954e-01j, + 5.307617160406101e-01 - 7.029184459442954e-01j, + 5.680556159453138e-01 + 6.788228792952775e-01j, + 5.680556159453138e-01 - 6.788228792952775e-01j, + 6.100962560818854e-01 + 6.693849403338664e-01j, + 6.100962560818854e-01 - 6.693849403338664e-01j, + 6.904694312740631e-01 + 6.930501690145245e-01j, + 6.904694312740631e-01 - 6.930501690145245e-01j, + 6.521767004237027e-01 + 6.744414640183752e-01j, + 6.521767004237027e-01 - 6.744414640183752e-01j, + ] + k2 = 3.398854055800844e-08 + xp_assert_equal(z, z2, check_dtype=False) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-13) + xp_assert_close(k, k2, rtol=1e-13) + + # bandpass analog + z, p, k = butter(4, [90.5, 110.5], 'bp', analog=True, output='zpk') + z2 = np.zeros(4, dtype=z.dtype) + p2 = [ + -4.179137760733086e+00 + 1.095935899082837e+02j, + -4.179137760733086e+00 - 1.095935899082837e+02j, + -9.593598668443835e+00 + 1.034745398029734e+02j, + -9.593598668443835e+00 - 1.034745398029734e+02j, + -8.883991981781929e+00 + 9.582087115567160e+01j, + -8.883991981781929e+00 - 9.582087115567160e+01j, + -3.474530886568715e+00 + 9.111599925805801e+01j, + -3.474530886568715e+00 - 9.111599925805801e+01j, + ] + k2 = 1.600000000000001e+05 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag)) + xp_assert_close(k, k2, rtol=1e-15) + + def test_bandstop(self): + z, p, k = butter(7, [0.45, 0.56], 'stop', output='zpk') + z2 = [-1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j, + -1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j, + -1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j, + -1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j, + -1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j, + -1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j, + -1.594474531383421e-02 + 9.998728744679880e-01j, + -1.594474531383421e-02 - 9.998728744679880e-01j] + p2 = [-1.766850742887729e-01 + 9.466951258673900e-01j, + -1.766850742887729e-01 - 9.466951258673900e-01j, + 1.467897662432886e-01 + 9.515917126462422e-01j, + 1.467897662432886e-01 - 9.515917126462422e-01j, + -1.370083529426906e-01 + 8.880376681273993e-01j, + -1.370083529426906e-01 - 8.880376681273993e-01j, + 1.086774544701390e-01 + 8.915240810704319e-01j, + 1.086774544701390e-01 - 8.915240810704319e-01j, + -7.982704457700891e-02 + 8.506056315273435e-01j, + -7.982704457700891e-02 - 8.506056315273435e-01j, + 5.238812787110331e-02 + 8.524011102699969e-01j, + 5.238812787110331e-02 - 8.524011102699969e-01j, + -1.357545000491310e-02 + 8.382287744986582e-01j, + -1.357545000491310e-02 - 8.382287744986582e-01j] + k2 = 4.577122512960063e-01 + xp_assert_close(sorted(z, key=np.imag), + sorted(z2, key=np.imag)) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag)) + xp_assert_close(k, k2, rtol=1e-14) + + def test_ba_output(self): + b, a = butter(4, [100, 300], 'bandpass', analog=True) + b2 = [1.6e+09, 0, 0, 0, 0] + a2 = [1.000000000000000e+00, 5.226251859505511e+02, + 2.565685424949238e+05, 6.794127417357160e+07, + 1.519411254969542e+10, 2.038238225207147e+12, + 2.309116882454312e+14, 1.411088002066486e+16, + 8.099999999999991e+17] + xp_assert_close(b, b2, rtol=1e-14) + xp_assert_close(a, a2, rtol=1e-14) + + def test_fs_param(self): + for fs in (900, 900.1, 1234.567): + for N in (0, 1, 2, 3, 10): + for fc in (100, 100.1, 432.12345): + for btype in ('lp', 'hp'): + ba1 = butter(N, fc, btype, fs=fs) + ba2 = butter(N, fc/(fs/2), btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + for fc in ((100, 200), (100.1, 200.2), (321.123, 432.123)): + for btype in ('bp', 'bs'): + ba1 = butter(N, fc, btype, fs=fs) + for seq in (list, tuple, array): + fcnorm = seq([f/(fs/2) for f in fc]) + ba2 = butter(N, fcnorm, btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + + +class TestCheby1: + + def test_degenerate(self): + # 0-order filter is just a passthrough + # Even-order filters have DC gain of -rp dB + b, a = cheby1(0, 10*np.log10(2), 1, analog=True) + assert_array_almost_equal(b, [1/np.sqrt(2)]) + xp_assert_equal(a, np.asarray([1.0])) + + # 1-order filter is same for all types + b, a = cheby1(1, 10*np.log10(2), 1, analog=True) + assert_array_almost_equal(b, [1]) + assert_array_almost_equal(a, [1, 1]) + + z, p, k = cheby1(1, 0.1, 0.3, output='zpk') + xp_assert_equal(z, np.asarray([-1.0])) + xp_assert_close(p, [-5.390126972799615e-01 + 0j], rtol=1e-14) + xp_assert_close(k, 7.695063486399808e-01, rtol=1e-14) + + def test_basic(self): + for N in range(25): + wn = 0.01 + z, p, k = cheby1(N, 1, wn, 'low', analog=True, output='zpk') + assert_array_almost_equal([], z) + assert len(p) == N + assert all(np.real(p) <= 0) # No poles in right half of S-plane + + for N in range(25): + wn = 0.01 + z, p, k = cheby1(N, 1, wn, 'high', analog=False, output='zpk') + xp_assert_equal(np.ones(N), z) # All zeros exactly at DC + assert all(np.abs(p) <= 1) # No poles outside unit circle + + # Same test as TestNormalize + b, a = cheby1(8, 0.5, 0.048) + assert_array_almost_equal(b, [ + 2.150733144728282e-11, 1.720586515782626e-10, + 6.022052805239190e-10, 1.204410561047838e-09, + 1.505513201309798e-09, 1.204410561047838e-09, + 6.022052805239190e-10, 1.720586515782626e-10, + 2.150733144728282e-11], decimal=14) + assert_array_almost_equal(a, [ + 1.000000000000000e+00, -7.782402035027959e+00, + 2.654354569747454e+01, -5.182182531666387e+01, + 6.334127355102684e+01, -4.963358186631157e+01, + 2.434862182949389e+01, -6.836925348604676e+00, + 8.412934944449140e-01], decimal=14) + + b, a = cheby1(4, 1, [0.4, 0.7], btype='band') + assert_array_almost_equal(b, [0.0084, 0, -0.0335, 0, 0.0502, 0, + -0.0335, 0, 0.0084], decimal=4) + assert_array_almost_equal(a, [1.0, 1.1191, 2.862, 2.2986, 3.4137, + 1.8653, 1.8982, 0.5676, 0.4103], + decimal=4) + + b2, a2 = cheby1(5, 3, 1, analog=True) + assert_array_almost_equal(b2, [0.0626], decimal=4) + assert_array_almost_equal(a2, [1, 0.5745, 1.4150, 0.5489, 0.4080, + 0.0626], decimal=4) + + b, a = cheby1(8, 0.5, 0.1) + assert_array_almost_equal(b, 1.0e-006 * np.array([ + 0.00703924326028, 0.05631394608227, 0.19709881128793, + 0.39419762257586, 0.49274702821983, 0.39419762257586, + 0.19709881128793, 0.05631394608227, 0.00703924326028]), + decimal=13) + assert_array_almost_equal(a, [ + 1.00000000000000, -7.44912258934158, 24.46749067762108, + -46.27560200466141, 55.11160187999928, -42.31640010161038, + 20.45543300484147, -5.69110270561444, 0.69770374759022], + decimal=13) + + b, a = cheby1(8, 0.5, 0.25) + assert_array_almost_equal(b, 1.0e-003 * np.array([ + 0.00895261138923, 0.07162089111382, 0.25067311889837, + 0.50134623779673, 0.62668279724591, 0.50134623779673, + 0.25067311889837, 0.07162089111382, 0.00895261138923]), + decimal=13) + assert_array_almost_equal(a, [1.00000000000000, -5.97529229188545, + 16.58122329202101, -27.71423273542923, + 30.39509758355313, -22.34729670426879, + 10.74509800434910, -3.08924633697497, + 0.40707685889802], decimal=13) + + def test_highpass(self): + # high even order + z, p, k = cheby1(24, 0.7, 0.2, 'high', output='zpk') + z2 = np.ones(24) + p2 = [-6.136558509657073e-01 + 2.700091504942893e-01j, + -6.136558509657073e-01 - 2.700091504942893e-01j, + -3.303348340927516e-01 + 6.659400861114254e-01j, + -3.303348340927516e-01 - 6.659400861114254e-01j, + 8.779713780557169e-03 + 8.223108447483040e-01j, + 8.779713780557169e-03 - 8.223108447483040e-01j, + 2.742361123006911e-01 + 8.356666951611864e-01j, + 2.742361123006911e-01 - 8.356666951611864e-01j, + 4.562984557158206e-01 + 7.954276912303594e-01j, + 4.562984557158206e-01 - 7.954276912303594e-01j, + 5.777335494123628e-01 + 7.435821817961783e-01j, + 5.777335494123628e-01 - 7.435821817961783e-01j, + 6.593260977749194e-01 + 6.955390907990932e-01j, + 6.593260977749194e-01 - 6.955390907990932e-01j, + 7.149590948466562e-01 + 6.559437858502012e-01j, + 7.149590948466562e-01 - 6.559437858502012e-01j, + 7.532432388188739e-01 + 6.256158042292060e-01j, + 7.532432388188739e-01 - 6.256158042292060e-01j, + 7.794365244268271e-01 + 6.042099234813333e-01j, + 7.794365244268271e-01 - 6.042099234813333e-01j, + 7.967253874772997e-01 + 5.911966597313203e-01j, + 7.967253874772997e-01 - 5.911966597313203e-01j, + 8.069756417293870e-01 + 5.862214589217275e-01j, + 8.069756417293870e-01 - 5.862214589217275e-01j] + k2 = 6.190427617192018e-04 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-10) + xp_assert_close(k, k2, rtol=1e-10) + + # high odd order + z, p, k = cheby1(23, 0.8, 0.3, 'high', output='zpk') + z2 = np.ones(23) + p2 = [-7.676400532011010e-01, + -6.754621070166477e-01 + 3.970502605619561e-01j, + -6.754621070166477e-01 - 3.970502605619561e-01j, + -4.528880018446727e-01 + 6.844061483786332e-01j, + -4.528880018446727e-01 - 6.844061483786332e-01j, + -1.986009130216447e-01 + 8.382285942941594e-01j, + -1.986009130216447e-01 - 8.382285942941594e-01j, + 2.504673931532608e-02 + 8.958137635794080e-01j, + 2.504673931532608e-02 - 8.958137635794080e-01j, + 2.001089429976469e-01 + 9.010678290791480e-01j, + 2.001089429976469e-01 - 9.010678290791480e-01j, + 3.302410157191755e-01 + 8.835444665962544e-01j, + 3.302410157191755e-01 - 8.835444665962544e-01j, + 4.246662537333661e-01 + 8.594054226449009e-01j, + 4.246662537333661e-01 - 8.594054226449009e-01j, + 4.919620928120296e-01 + 8.366772762965786e-01j, + 4.919620928120296e-01 - 8.366772762965786e-01j, + 5.385746917494749e-01 + 8.191616180796720e-01j, + 5.385746917494749e-01 - 8.191616180796720e-01j, + 5.855636993537203e-01 + 8.060680937701062e-01j, + 5.855636993537203e-01 - 8.060680937701062e-01j, + 5.688812849391721e-01 + 8.086497795114683e-01j, + 5.688812849391721e-01 - 8.086497795114683e-01j] + k2 = 1.941697029206324e-05 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-10) + xp_assert_close(k, k2, rtol=1e-10) + + z, p, k = cheby1(10, 1, 1000, 'high', analog=True, output='zpk') + z2 = np.zeros(10) + p2 = [-3.144743169501551e+03 + 3.511680029092744e+03j, + -3.144743169501551e+03 - 3.511680029092744e+03j, + -5.633065604514602e+02 + 2.023615191183945e+03j, + -5.633065604514602e+02 - 2.023615191183945e+03j, + -1.946412183352025e+02 + 1.372309454274755e+03j, + -1.946412183352025e+02 - 1.372309454274755e+03j, + -7.987162953085479e+01 + 1.105207708045358e+03j, + -7.987162953085479e+01 - 1.105207708045358e+03j, + -2.250315039031946e+01 + 1.001723931471477e+03j, + -2.250315039031946e+01 - 1.001723931471477e+03j] + k2 = 8.912509381337453e-01 + xp_assert_equal(z, z2) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-13) + xp_assert_close(k, k2, rtol=1e-15) + + def test_bandpass(self): + z, p, k = cheby1(8, 1, [0.3, 0.4], 'bp', output='zpk') + z2 = [1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1, -1] + p2 = [3.077784854851463e-01 + 9.453307017592942e-01j, + 3.077784854851463e-01 - 9.453307017592942e-01j, + 3.280567400654425e-01 + 9.272377218689016e-01j, + 3.280567400654425e-01 - 9.272377218689016e-01j, + 3.677912763284301e-01 + 9.038008865279966e-01j, + 3.677912763284301e-01 - 9.038008865279966e-01j, + 4.194425632520948e-01 + 8.769407159656157e-01j, + 4.194425632520948e-01 - 8.769407159656157e-01j, + 4.740921994669189e-01 + 8.496508528630974e-01j, + 4.740921994669189e-01 - 8.496508528630974e-01j, + 5.234866481897429e-01 + 8.259608422808477e-01j, + 5.234866481897429e-01 - 8.259608422808477e-01j, + 5.844717632289875e-01 + 8.052901363500210e-01j, + 5.844717632289875e-01 - 8.052901363500210e-01j, + 5.615189063336070e-01 + 8.100667803850766e-01j, + 5.615189063336070e-01 - 8.100667803850766e-01j] + k2 = 5.007028718074307e-09 + xp_assert_equal(z, z2, check_dtype=False) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-13) + xp_assert_close(k, k2, rtol=1e-13) + + def test_bandstop(self): + z, p, k = cheby1(7, 1, [0.5, 0.6], 'stop', output='zpk') + z2 = [-1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j, + -1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j, + -1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j, + -1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j, + -1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j, + -1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j, + -1.583844403245361e-01 + 9.873775210440450e-01j, + -1.583844403245361e-01 - 9.873775210440450e-01j] + p2 = [-8.942974551472813e-02 + 3.482480481185926e-01j, + -8.942974551472813e-02 - 3.482480481185926e-01j, + 1.293775154041798e-01 + 8.753499858081858e-01j, + 1.293775154041798e-01 - 8.753499858081858e-01j, + 3.399741945062013e-02 + 9.690316022705607e-01j, + 3.399741945062013e-02 - 9.690316022705607e-01j, + 4.167225522796539e-04 + 9.927338161087488e-01j, + 4.167225522796539e-04 - 9.927338161087488e-01j, + -3.912966549550960e-01 + 8.046122859255742e-01j, + -3.912966549550960e-01 - 8.046122859255742e-01j, + -3.307805547127368e-01 + 9.133455018206508e-01j, + -3.307805547127368e-01 - 9.133455018206508e-01j, + -3.072658345097743e-01 + 9.443589759799366e-01j, + -3.072658345097743e-01 - 9.443589759799366e-01j] + k2 = 3.619438310405028e-01 + xp_assert_close(sorted(z, key=np.imag), + sorted(z2, key=np.imag), rtol=1e-13) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-13) + xp_assert_close(k, k2, rtol=0, atol=5e-16) + + def test_ba_output(self): + # with transfer function conversion, without digital conversion + b, a = cheby1(5, 0.9, [210, 310], 'stop', analog=True) + b2 = [1.000000000000006e+00, 0, + 3.255000000000020e+05, 0, + 4.238010000000026e+10, 0, + 2.758944510000017e+15, 0, + 8.980364380050052e+19, 0, + 1.169243442282517e+24 + ] + a2 = [1.000000000000000e+00, 4.630555945694342e+02, + 4.039266454794788e+05, 1.338060988610237e+08, + 5.844333551294591e+10, 1.357346371637638e+13, + 3.804661141892782e+15, 5.670715850340080e+17, + 1.114411200988328e+20, 8.316815934908471e+21, + 1.169243442282517e+24 + ] + xp_assert_close(b, b2, rtol=1e-14) + xp_assert_close(a, a2, rtol=1e-14) + + def test_fs_param(self): + for fs in (900, 900.1, 1234.567): + for N in (0, 1, 2, 3, 10): + for fc in (100, 100.1, 432.12345): + for btype in ('lp', 'hp'): + ba1 = cheby1(N, 1, fc, btype, fs=fs) + ba2 = cheby1(N, 1, fc/(fs/2), btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + for fc in ((100, 200), (100.1, 200.2), (321.123, 432.123)): + for btype in ('bp', 'bs'): + ba1 = cheby1(N, 1, fc, btype, fs=fs) + for seq in (list, tuple, array): + fcnorm = seq([f/(fs/2) for f in fc]) + ba2 = cheby1(N, 1, fcnorm, btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + +class TestCheby2: + + def test_degenerate(self): + # 0-order filter is just a passthrough + # Stopband ripple factor doesn't matter + b, a = cheby2(0, 123.456, 1, analog=True) + xp_assert_equal(b, np.asarray([1.0])) + xp_assert_equal(a, np.asarray([1.0])) + + # 1-order filter is same for all types + b, a = cheby2(1, 10*np.log10(2), 1, analog=True) + assert_array_almost_equal(b, [1]) + assert_array_almost_equal(a, [1, 1]) + + z, p, k = cheby2(1, 50, 0.3, output='zpk') + xp_assert_equal(z, np.asarray([-1], dtype=np.complex128)) + xp_assert_close(p, [9.967826460175649e-01 + 0j], rtol=1e-14) + xp_assert_close(k, 1.608676991217512e-03, rtol=1e-14) + + def test_basic(self): + for N in range(25): + wn = 0.01 + z, p, k = cheby2(N, 40, wn, 'low', analog=True, output='zpk') + assert len(p) == N + assert all(np.real(p) <= 0) # No poles in right half of S-plane + + for N in range(25): + wn = 0.01 + z, p, k = cheby2(N, 40, wn, 'high', analog=False, output='zpk') + assert all(np.abs(p) <= 1) # No poles outside unit circle + + B, A = cheby2(18, 100, 0.5) + assert_array_almost_equal(B, [ + 0.00167583914216, 0.01249479541868, 0.05282702120282, + 0.15939804265706, 0.37690207631117, 0.73227013789108, + 1.20191856962356, 1.69522872823393, 2.07598674519837, + 2.21972389625291, 2.07598674519838, 1.69522872823395, + 1.20191856962359, 0.73227013789110, 0.37690207631118, + 0.15939804265707, 0.05282702120282, 0.01249479541868, + 0.00167583914216], decimal=13) + assert_array_almost_equal(A, [ + 1.00000000000000, -0.27631970006174, 3.19751214254060, + -0.15685969461355, 4.13926117356269, 0.60689917820044, + 2.95082770636540, 0.89016501910416, 1.32135245849798, + 0.51502467236824, 0.38906643866660, 0.15367372690642, + 0.07255803834919, 0.02422454070134, 0.00756108751837, + 0.00179848550988, 0.00033713574499, 0.00004258794833, + 0.00000281030149], decimal=13) + + def test_highpass(self): + # high even order + z, p, k = cheby2(26, 60, 0.3, 'high', output='zpk') + z2 = [9.981088955489852e-01 + 6.147058341984388e-02j, + 9.981088955489852e-01 - 6.147058341984388e-02j, + 9.832702870387426e-01 + 1.821525257215483e-01j, + 9.832702870387426e-01 - 1.821525257215483e-01j, + 9.550760158089112e-01 + 2.963609353922882e-01j, + 9.550760158089112e-01 - 2.963609353922882e-01j, + 9.162054748821922e-01 + 4.007087817803773e-01j, + 9.162054748821922e-01 - 4.007087817803773e-01j, + 8.700619897368064e-01 + 4.929423232136168e-01j, + 8.700619897368064e-01 - 4.929423232136168e-01j, + 5.889791753434985e-01 + 8.081482110427953e-01j, + 5.889791753434985e-01 - 8.081482110427953e-01j, + 5.984900456570295e-01 + 8.011302423760501e-01j, + 5.984900456570295e-01 - 8.011302423760501e-01j, + 6.172880888914629e-01 + 7.867371958365343e-01j, + 6.172880888914629e-01 - 7.867371958365343e-01j, + 6.448899971038180e-01 + 7.642754030030161e-01j, + 6.448899971038180e-01 - 7.642754030030161e-01j, + 6.804845629637927e-01 + 7.327624168637228e-01j, + 6.804845629637927e-01 - 7.327624168637228e-01j, + 8.202619107108660e-01 + 5.719881098737678e-01j, + 8.202619107108660e-01 - 5.719881098737678e-01j, + 7.228410452536148e-01 + 6.910143437705678e-01j, + 7.228410452536148e-01 - 6.910143437705678e-01j, + 7.702121399578629e-01 + 6.377877856007792e-01j, + 7.702121399578629e-01 - 6.377877856007792e-01j] + p2 = [7.365546198286450e-01 + 4.842085129329526e-02j, + 7.365546198286450e-01 - 4.842085129329526e-02j, + 7.292038510962885e-01 + 1.442201672097581e-01j, + 7.292038510962885e-01 - 1.442201672097581e-01j, + 7.151293788040354e-01 + 2.369925800458584e-01j, + 7.151293788040354e-01 - 2.369925800458584e-01j, + 6.955051820787286e-01 + 3.250341363856910e-01j, + 6.955051820787286e-01 - 3.250341363856910e-01j, + 6.719122956045220e-01 + 4.070475750638047e-01j, + 6.719122956045220e-01 - 4.070475750638047e-01j, + 6.461722130611300e-01 + 4.821965916689270e-01j, + 6.461722130611300e-01 - 4.821965916689270e-01j, + 5.528045062872224e-01 + 8.162920513838372e-01j, + 5.528045062872224e-01 - 8.162920513838372e-01j, + 5.464847782492791e-01 + 7.869899955967304e-01j, + 5.464847782492791e-01 - 7.869899955967304e-01j, + 5.488033111260949e-01 + 7.520442354055579e-01j, + 5.488033111260949e-01 - 7.520442354055579e-01j, + 6.201874719022955e-01 + 5.500894392527353e-01j, + 6.201874719022955e-01 - 5.500894392527353e-01j, + 5.586478152536709e-01 + 7.112676877332921e-01j, + 5.586478152536709e-01 - 7.112676877332921e-01j, + 5.958145844148228e-01 + 6.107074340842115e-01j, + 5.958145844148228e-01 - 6.107074340842115e-01j, + 5.747812938519067e-01 + 6.643001536914696e-01j, + 5.747812938519067e-01 - 6.643001536914696e-01j] + k2 = 9.932997786497189e-02 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-13) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-12) + xp_assert_close(k, k2, rtol=1e-11) + + # high odd order + z, p, k = cheby2(25, 80, 0.5, 'high', output='zpk') + z2 = [9.690690376586687e-01 + 2.467897896011971e-01j, + 9.690690376586687e-01 - 2.467897896011971e-01j, + 9.999999999999492e-01, + 8.835111277191199e-01 + 4.684101698261429e-01j, + 8.835111277191199e-01 - 4.684101698261429e-01j, + 7.613142857900539e-01 + 6.483830335935022e-01j, + 7.613142857900539e-01 - 6.483830335935022e-01j, + 6.232625173626231e-01 + 7.820126817709752e-01j, + 6.232625173626231e-01 - 7.820126817709752e-01j, + 4.864456563413621e-01 + 8.737108351316745e-01j, + 4.864456563413621e-01 - 8.737108351316745e-01j, + 3.618368136816749e-01 + 9.322414495530347e-01j, + 3.618368136816749e-01 - 9.322414495530347e-01j, + 2.549486883466794e-01 + 9.669545833752675e-01j, + 2.549486883466794e-01 - 9.669545833752675e-01j, + 1.676175432109457e-01 + 9.858520980390212e-01j, + 1.676175432109457e-01 - 9.858520980390212e-01j, + 1.975218468277521e-03 + 9.999980492540941e-01j, + 1.975218468277521e-03 - 9.999980492540941e-01j, + 1.786959496651858e-02 + 9.998403260399917e-01j, + 1.786959496651858e-02 - 9.998403260399917e-01j, + 9.967933660557139e-02 + 9.950196127985684e-01j, + 9.967933660557139e-02 - 9.950196127985684e-01j, + 5.013970951219547e-02 + 9.987422137518890e-01j, + 5.013970951219547e-02 - 9.987422137518890e-01j] + p2 = [4.218866331906864e-01, + 4.120110200127552e-01 + 1.361290593621978e-01j, + 4.120110200127552e-01 - 1.361290593621978e-01j, + 3.835890113632530e-01 + 2.664910809911026e-01j, + 3.835890113632530e-01 - 2.664910809911026e-01j, + 3.399195570456499e-01 + 3.863983538639875e-01j, + 3.399195570456499e-01 - 3.863983538639875e-01j, + 2.855977834508353e-01 + 4.929444399540688e-01j, + 2.855977834508353e-01 - 4.929444399540688e-01j, + 2.255765441339322e-01 + 5.851631870205766e-01j, + 2.255765441339322e-01 - 5.851631870205766e-01j, + 1.644087535815792e-01 + 6.637356937277153e-01j, + 1.644087535815792e-01 - 6.637356937277153e-01j, + -7.293633845273095e-02 + 9.739218252516307e-01j, + -7.293633845273095e-02 - 9.739218252516307e-01j, + 1.058259206358626e-01 + 7.304739464862978e-01j, + 1.058259206358626e-01 - 7.304739464862978e-01j, + -5.703971947785402e-02 + 9.291057542169088e-01j, + -5.703971947785402e-02 - 9.291057542169088e-01j, + 5.263875132656864e-02 + 7.877974334424453e-01j, + 5.263875132656864e-02 - 7.877974334424453e-01j, + -3.007943405982616e-02 + 8.846331716180016e-01j, + -3.007943405982616e-02 - 8.846331716180016e-01j, + 6.857277464483946e-03 + 8.383275456264492e-01j, + 6.857277464483946e-03 - 8.383275456264492e-01j] + k2 = 6.507068761705037e-03 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-13) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-12) + xp_assert_close(k, k2, rtol=1e-11) + + def test_bandpass(self): + z, p, k = cheby2(9, 40, [0.07, 0.2], 'pass', output='zpk') + z2 = [-9.999999999999999e-01, + 3.676588029658514e-01 + 9.299607543341383e-01j, + 3.676588029658514e-01 - 9.299607543341383e-01j, + 7.009689684982283e-01 + 7.131917730894889e-01j, + 7.009689684982283e-01 - 7.131917730894889e-01j, + 7.815697973765858e-01 + 6.238178033919218e-01j, + 7.815697973765858e-01 - 6.238178033919218e-01j, + 8.063793628819866e-01 + 5.913986160941200e-01j, + 8.063793628819866e-01 - 5.913986160941200e-01j, + 1.000000000000001e+00, + 9.944493019920448e-01 + 1.052168511576739e-01j, + 9.944493019920448e-01 - 1.052168511576739e-01j, + 9.854674703367308e-01 + 1.698642543566085e-01j, + 9.854674703367308e-01 - 1.698642543566085e-01j, + 9.762751735919308e-01 + 2.165335665157851e-01j, + 9.762751735919308e-01 - 2.165335665157851e-01j, + 9.792277171575134e-01 + 2.027636011479496e-01j, + 9.792277171575134e-01 - 2.027636011479496e-01j] + p2 = [8.143803410489621e-01 + 5.411056063397541e-01j, + 8.143803410489621e-01 - 5.411056063397541e-01j, + 7.650769827887418e-01 + 5.195412242095543e-01j, + 7.650769827887418e-01 - 5.195412242095543e-01j, + 6.096241204063443e-01 + 3.568440484659796e-01j, + 6.096241204063443e-01 - 3.568440484659796e-01j, + 6.918192770246239e-01 + 4.770463577106911e-01j, + 6.918192770246239e-01 - 4.770463577106911e-01j, + 6.986241085779207e-01 + 1.146512226180060e-01j, + 6.986241085779207e-01 - 1.146512226180060e-01j, + 8.654645923909734e-01 + 1.604208797063147e-01j, + 8.654645923909734e-01 - 1.604208797063147e-01j, + 9.164831670444591e-01 + 1.969181049384918e-01j, + 9.164831670444591e-01 - 1.969181049384918e-01j, + 9.630425777594550e-01 + 2.317513360702271e-01j, + 9.630425777594550e-01 - 2.317513360702271e-01j, + 9.438104703725529e-01 + 2.193509900269860e-01j, + 9.438104703725529e-01 - 2.193509900269860e-01j] + k2 = 9.345352824659604e-03 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-13) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-13) + xp_assert_close(k, k2, rtol=1e-11) + + def test_bandstop(self): + z, p, k = cheby2(6, 55, [0.1, 0.9], 'stop', output='zpk') + z2 = [6.230544895101009e-01 + 7.821784343111114e-01j, + 6.230544895101009e-01 - 7.821784343111114e-01j, + 9.086608545660115e-01 + 4.175349702471991e-01j, + 9.086608545660115e-01 - 4.175349702471991e-01j, + 9.478129721465802e-01 + 3.188268649763867e-01j, + 9.478129721465802e-01 - 3.188268649763867e-01j, + -6.230544895100982e-01 + 7.821784343111109e-01j, + -6.230544895100982e-01 - 7.821784343111109e-01j, + -9.086608545660116e-01 + 4.175349702472088e-01j, + -9.086608545660116e-01 - 4.175349702472088e-01j, + -9.478129721465784e-01 + 3.188268649763897e-01j, + -9.478129721465784e-01 - 3.188268649763897e-01j] + p2 = [-9.464094036167638e-01 + 1.720048695084344e-01j, + -9.464094036167638e-01 - 1.720048695084344e-01j, + -8.715844103386737e-01 + 1.370665039509297e-01j, + -8.715844103386737e-01 - 1.370665039509297e-01j, + -8.078751204586425e-01 + 5.729329866682983e-02j, + -8.078751204586425e-01 - 5.729329866682983e-02j, + 9.464094036167665e-01 + 1.720048695084332e-01j, + 9.464094036167665e-01 - 1.720048695084332e-01j, + 8.078751204586447e-01 + 5.729329866683007e-02j, + 8.078751204586447e-01 - 5.729329866683007e-02j, + 8.715844103386721e-01 + 1.370665039509331e-01j, + 8.715844103386721e-01 - 1.370665039509331e-01j] + k2 = 2.917823332763358e-03 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-13) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-13) + xp_assert_close(k, k2, rtol=1e-11) + + def test_ba_output(self): + # with transfer function conversion, without digital conversion + b, a = cheby2(5, 20, [2010, 2100], 'stop', True) + b2 = [1.000000000000000e+00, 0, # Matlab: 6.683253076978249e-12, + 2.111512500000000e+07, 0, # Matlab: 1.134325604589552e-04, + 1.782966433781250e+14, 0, # Matlab: 7.216787944356781e+02, + 7.525901316990656e+20, 0, # Matlab: 2.039829265789886e+09, + 1.587960565565748e+27, 0, # Matlab: 2.161236218626134e+15, + 1.339913493808585e+33] + a2 = [1.000000000000000e+00, 1.849550755473371e+02, + 2.113222918998538e+07, 3.125114149732283e+09, + 1.785133457155609e+14, 1.979158697776348e+16, + 7.535048322653831e+20, 5.567966191263037e+22, + 1.589246884221346e+27, 5.871210648525566e+28, + 1.339913493808590e+33] + xp_assert_close(b, b2, rtol=1e-14) + xp_assert_close(a, a2, rtol=1e-14) + + def test_fs_param(self): + for fs in (900, 900.1, 1234.567): + for N in (0, 1, 2, 3, 10): + for fc in (100, 100.1, 432.12345): + for btype in ('lp', 'hp'): + ba1 = cheby2(N, 20, fc, btype, fs=fs) + ba2 = cheby2(N, 20, fc/(fs/2), btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + for fc in ((100, 200), (100.1, 200.2), (321.123, 432.123)): + for btype in ('bp', 'bs'): + ba1 = cheby2(N, 20, fc, btype, fs=fs) + for seq in (list, tuple, array): + fcnorm = seq([f/(fs/2) for f in fc]) + ba2 = cheby2(N, 20, fcnorm, btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + +class TestEllip: + + def test_degenerate(self): + # 0-order filter is just a passthrough + # Even-order filters have DC gain of -rp dB + # Stopband ripple factor doesn't matter + b, a = ellip(0, 10*np.log10(2), 123.456, 1, analog=True) + assert_array_almost_equal(b, [1/np.sqrt(2)]) + xp_assert_equal(a, np.asarray([1.0])) + + # 1-order filter is same for all types + b, a = ellip(1, 10*np.log10(2), 1, 1, analog=True) + assert_array_almost_equal(b, [1]) + assert_array_almost_equal(a, [1, 1]) + + z, p, k = ellip(1, 1, 55, 0.3, output='zpk') + xp_assert_close(z, [-9.999999999999998e-01], rtol=1e-14) + xp_assert_close(p, [-6.660721153525525e-04], rtol=1e-10) + xp_assert_close(k, 5.003330360576763e-01, rtol=1e-14) + + def test_basic(self): + for N in range(25): + wn = 0.01 + z, p, k = ellip(N, 1, 40, wn, 'low', analog=True, output='zpk') + assert len(p) == N + assert all(np.real(p) <= 0) # No poles in right half of S-plane + + for N in range(25): + wn = 0.01 + z, p, k = ellip(N, 1, 40, wn, 'high', analog=False, output='zpk') + assert all(np.abs(p) <= 1) # No poles outside unit circle + + b3, a3 = ellip(5, 3, 26, 1, analog=True) + assert_array_almost_equal(b3, [0.1420, 0, 0.3764, 0, + 0.2409], decimal=4) + assert_array_almost_equal(a3, [1, 0.5686, 1.8061, 0.8017, 0.8012, + 0.2409], decimal=4) + + b, a = ellip(3, 1, 60, [0.4, 0.7], 'stop') + assert_array_almost_equal(b, [0.3310, 0.3469, 1.1042, 0.7044, 1.1042, + 0.3469, 0.3310], decimal=4) + assert_array_almost_equal(a, [1.0000, 0.6973, 1.1441, 0.5878, 0.7323, + 0.1131, -0.0060], decimal=4) + + def test_highpass(self): + # high even order + z, p, k = ellip(24, 1, 80, 0.3, 'high', output='zpk') + z2 = [9.761875332501075e-01 + 2.169283290099910e-01j, + 9.761875332501075e-01 - 2.169283290099910e-01j, + 8.413503353963494e-01 + 5.404901600661900e-01j, + 8.413503353963494e-01 - 5.404901600661900e-01j, + 7.160082576305009e-01 + 6.980918098681732e-01j, + 7.160082576305009e-01 - 6.980918098681732e-01j, + 6.456533638965329e-01 + 7.636306264739803e-01j, + 6.456533638965329e-01 - 7.636306264739803e-01j, + 6.127321820971366e-01 + 7.902906256703928e-01j, + 6.127321820971366e-01 - 7.902906256703928e-01j, + 5.983607817490196e-01 + 8.012267936512676e-01j, + 5.983607817490196e-01 - 8.012267936512676e-01j, + 5.922577552594799e-01 + 8.057485658286990e-01j, + 5.922577552594799e-01 - 8.057485658286990e-01j, + 5.896952092563588e-01 + 8.076258788449631e-01j, + 5.896952092563588e-01 - 8.076258788449631e-01j, + 5.886248765538837e-01 + 8.084063054565607e-01j, + 5.886248765538837e-01 - 8.084063054565607e-01j, + 5.881802711123132e-01 + 8.087298490066037e-01j, + 5.881802711123132e-01 - 8.087298490066037e-01j, + 5.879995719101164e-01 + 8.088612386766461e-01j, + 5.879995719101164e-01 - 8.088612386766461e-01j, + 5.879354086709576e-01 + 8.089078780868164e-01j, + 5.879354086709576e-01 - 8.089078780868164e-01j] + p2 = [-3.184805259081650e-01 + 4.206951906775851e-01j, + -3.184805259081650e-01 - 4.206951906775851e-01j, + 1.417279173459985e-01 + 7.903955262836452e-01j, + 1.417279173459985e-01 - 7.903955262836452e-01j, + 4.042881216964651e-01 + 8.309042239116594e-01j, + 4.042881216964651e-01 - 8.309042239116594e-01j, + 5.128964442789670e-01 + 8.229563236799665e-01j, + 5.128964442789670e-01 - 8.229563236799665e-01j, + 5.569614712822724e-01 + 8.155957702908510e-01j, + 5.569614712822724e-01 - 8.155957702908510e-01j, + 5.750478870161392e-01 + 8.118633973883931e-01j, + 5.750478870161392e-01 - 8.118633973883931e-01j, + 5.825314018170804e-01 + 8.101960910679270e-01j, + 5.825314018170804e-01 - 8.101960910679270e-01j, + 5.856397379751872e-01 + 8.094825218722543e-01j, + 5.856397379751872e-01 - 8.094825218722543e-01j, + 5.869326035251949e-01 + 8.091827531557583e-01j, + 5.869326035251949e-01 - 8.091827531557583e-01j, + 5.874697218855733e-01 + 8.090593298213502e-01j, + 5.874697218855733e-01 - 8.090593298213502e-01j, + 5.876904783532237e-01 + 8.090127161018823e-01j, + 5.876904783532237e-01 - 8.090127161018823e-01j, + 5.877753105317594e-01 + 8.090050577978136e-01j, + 5.877753105317594e-01 - 8.090050577978136e-01j] + k2 = 4.918081266957108e-02 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-4) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-4) + xp_assert_close(k, k2, rtol=1e-3) + + # high odd order + z, p, k = ellip(23, 1, 70, 0.5, 'high', output='zpk') + z2 = [9.999999999998661e-01, + 6.603717261750994e-01 + 7.509388678638675e-01j, + 6.603717261750994e-01 - 7.509388678638675e-01j, + 2.788635267510325e-01 + 9.603307416968041e-01j, + 2.788635267510325e-01 - 9.603307416968041e-01j, + 1.070215532544218e-01 + 9.942567008268131e-01j, + 1.070215532544218e-01 - 9.942567008268131e-01j, + 4.049427369978163e-02 + 9.991797705105507e-01j, + 4.049427369978163e-02 - 9.991797705105507e-01j, + 1.531059368627931e-02 + 9.998827859909265e-01j, + 1.531059368627931e-02 - 9.998827859909265e-01j, + 5.808061438534933e-03 + 9.999831330689181e-01j, + 5.808061438534933e-03 - 9.999831330689181e-01j, + 2.224277847754599e-03 + 9.999975262909676e-01j, + 2.224277847754599e-03 - 9.999975262909676e-01j, + 8.731857107534554e-04 + 9.999996187732845e-01j, + 8.731857107534554e-04 - 9.999996187732845e-01j, + 3.649057346914968e-04 + 9.999999334218996e-01j, + 3.649057346914968e-04 - 9.999999334218996e-01j, + 1.765538109802615e-04 + 9.999999844143768e-01j, + 1.765538109802615e-04 - 9.999999844143768e-01j, + 1.143655290967426e-04 + 9.999999934602630e-01j, + 1.143655290967426e-04 - 9.999999934602630e-01j] + p2 = [-6.322017026545028e-01, + -4.648423756662754e-01 + 5.852407464440732e-01j, + -4.648423756662754e-01 - 5.852407464440732e-01j, + -2.249233374627773e-01 + 8.577853017985717e-01j, + -2.249233374627773e-01 - 8.577853017985717e-01j, + -9.234137570557621e-02 + 9.506548198678851e-01j, + -9.234137570557621e-02 - 9.506548198678851e-01j, + -3.585663561241373e-02 + 9.821494736043981e-01j, + -3.585663561241373e-02 - 9.821494736043981e-01j, + -1.363917242312723e-02 + 9.933844128330656e-01j, + -1.363917242312723e-02 - 9.933844128330656e-01j, + -5.131505238923029e-03 + 9.975221173308673e-01j, + -5.131505238923029e-03 - 9.975221173308673e-01j, + -1.904937999259502e-03 + 9.990680819857982e-01j, + -1.904937999259502e-03 - 9.990680819857982e-01j, + -6.859439885466834e-04 + 9.996492201426826e-01j, + -6.859439885466834e-04 - 9.996492201426826e-01j, + -2.269936267937089e-04 + 9.998686250679161e-01j, + -2.269936267937089e-04 - 9.998686250679161e-01j, + -5.687071588789117e-05 + 9.999527573294513e-01j, + -5.687071588789117e-05 - 9.999527573294513e-01j, + -6.948417068525226e-07 + 9.999882737700173e-01j, + -6.948417068525226e-07 - 9.999882737700173e-01j] + k2 = 1.220910020289434e-02 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-4) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-4) + xp_assert_close(k, k2, rtol=1e-3) + + def test_bandpass(self): + z, p, k = ellip(7, 1, 40, [0.07, 0.2], 'pass', output='zpk') + z2 = [-9.999999999999991e-01, + 6.856610961780020e-01 + 7.279209168501619e-01j, + 6.856610961780020e-01 - 7.279209168501619e-01j, + 7.850346167691289e-01 + 6.194518952058737e-01j, + 7.850346167691289e-01 - 6.194518952058737e-01j, + 7.999038743173071e-01 + 6.001281461922627e-01j, + 7.999038743173071e-01 - 6.001281461922627e-01j, + 9.999999999999999e-01, + 9.862938983554124e-01 + 1.649980183725925e-01j, + 9.862938983554124e-01 - 1.649980183725925e-01j, + 9.788558330548762e-01 + 2.045513580850601e-01j, + 9.788558330548762e-01 - 2.045513580850601e-01j, + 9.771155231720003e-01 + 2.127093189691258e-01j, + 9.771155231720003e-01 - 2.127093189691258e-01j] + p2 = [8.063992755498643e-01 + 5.858071374778874e-01j, + 8.063992755498643e-01 - 5.858071374778874e-01j, + 8.050395347071724e-01 + 5.639097428109795e-01j, + 8.050395347071724e-01 - 5.639097428109795e-01j, + 8.113124936559144e-01 + 4.855241143973142e-01j, + 8.113124936559144e-01 - 4.855241143973142e-01j, + 8.665595314082394e-01 + 3.334049560919331e-01j, + 8.665595314082394e-01 - 3.334049560919331e-01j, + 9.412369011968871e-01 + 2.457616651325908e-01j, + 9.412369011968871e-01 - 2.457616651325908e-01j, + 9.679465190411238e-01 + 2.228772501848216e-01j, + 9.679465190411238e-01 - 2.228772501848216e-01j, + 9.747235066273385e-01 + 2.178937926146544e-01j, + 9.747235066273385e-01 - 2.178937926146544e-01j] + k2 = 8.354782670263239e-03 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-4) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-4) + xp_assert_close(k, k2, rtol=1e-3) + + z, p, k = ellip(5, 1, 75, [90.5, 110.5], 'pass', True, 'zpk') + z2 = [-5.583607317695175e-14 + 1.433755965989225e+02j, + -5.583607317695175e-14 - 1.433755965989225e+02j, + 5.740106416459296e-14 + 1.261678754570291e+02j, + 5.740106416459296e-14 - 1.261678754570291e+02j, + -2.199676239638652e-14 + 6.974861996895196e+01j, + -2.199676239638652e-14 - 6.974861996895196e+01j, + -3.372595657044283e-14 + 7.926145989044531e+01j, + -3.372595657044283e-14 - 7.926145989044531e+01j, + 0] + p2 = [-8.814960004852743e-01 + 1.104124501436066e+02j, + -8.814960004852743e-01 - 1.104124501436066e+02j, + -2.477372459140184e+00 + 1.065638954516534e+02j, + -2.477372459140184e+00 - 1.065638954516534e+02j, + -3.072156842945799e+00 + 9.995404870405324e+01j, + -3.072156842945799e+00 - 9.995404870405324e+01j, + -2.180456023925693e+00 + 9.379206865455268e+01j, + -2.180456023925693e+00 - 9.379206865455268e+01j, + -7.230484977485752e-01 + 9.056598800801140e+01j, + -7.230484977485752e-01 - 9.056598800801140e+01j] + k2 = 3.774571622827070e-02 + xp_assert_close(sorted(z, key=np.imag), + sorted(z2, key=np.imag), rtol=1e-4) + xp_assert_close(sorted(p, key=np.imag), + sorted(p2, key=np.imag), rtol=1e-6) + xp_assert_close(k, k2, rtol=1e-3) + + def test_bandstop(self): + z, p, k = ellip(8, 1, 65, [0.2, 0.4], 'stop', output='zpk') + z2 = [3.528578094286510e-01 + 9.356769561794296e-01j, + 3.528578094286510e-01 - 9.356769561794296e-01j, + 3.769716042264783e-01 + 9.262248159096587e-01j, + 3.769716042264783e-01 - 9.262248159096587e-01j, + 4.406101783111199e-01 + 8.976985411420985e-01j, + 4.406101783111199e-01 - 8.976985411420985e-01j, + 5.539386470258847e-01 + 8.325574907062760e-01j, + 5.539386470258847e-01 - 8.325574907062760e-01j, + 6.748464963023645e-01 + 7.379581332490555e-01j, + 6.748464963023645e-01 - 7.379581332490555e-01j, + 7.489887970285254e-01 + 6.625826604475596e-01j, + 7.489887970285254e-01 - 6.625826604475596e-01j, + 7.913118471618432e-01 + 6.114127579150699e-01j, + 7.913118471618432e-01 - 6.114127579150699e-01j, + 7.806804740916381e-01 + 6.249303940216475e-01j, + 7.806804740916381e-01 - 6.249303940216475e-01j] + + p2 = [-1.025299146693730e-01 + 5.662682444754943e-01j, + -1.025299146693730e-01 - 5.662682444754943e-01j, + 1.698463595163031e-01 + 8.926678667070186e-01j, + 1.698463595163031e-01 - 8.926678667070186e-01j, + 2.750532687820631e-01 + 9.351020170094005e-01j, + 2.750532687820631e-01 - 9.351020170094005e-01j, + 3.070095178909486e-01 + 9.457373499553291e-01j, + 3.070095178909486e-01 - 9.457373499553291e-01j, + 7.695332312152288e-01 + 2.792567212705257e-01j, + 7.695332312152288e-01 - 2.792567212705257e-01j, + 8.083818999225620e-01 + 4.990723496863960e-01j, + 8.083818999225620e-01 - 4.990723496863960e-01j, + 8.066158014414928e-01 + 5.649811440393374e-01j, + 8.066158014414928e-01 - 5.649811440393374e-01j, + 8.062787978834571e-01 + 5.855780880424964e-01j, + 8.062787978834571e-01 - 5.855780880424964e-01j] + k2 = 2.068622545291259e-01 + xp_assert_close(sorted(z, key=np.angle), + sorted(z2, key=np.angle), rtol=1e-6) + xp_assert_close(sorted(p, key=np.angle), + sorted(p2, key=np.angle), rtol=1e-5) + xp_assert_close(k, k2, rtol=1e-5) + + def test_ba_output(self): + # with transfer function conversion, without digital conversion + b, a = ellip(5, 1, 40, [201, 240], 'stop', True) + b2 = [ + 1.000000000000000e+00, 0, # Matlab: 1.743506051190569e-13, + 2.426561778314366e+05, 0, # Matlab: 3.459426536825722e-08, + 2.348218683400168e+10, 0, # Matlab: 2.559179747299313e-03, + 1.132780692872241e+15, 0, # Matlab: 8.363229375535731e+01, + 2.724038554089566e+19, 0, # Matlab: 1.018700994113120e+06, + 2.612380874940186e+23 + ] + a2 = [ + 1.000000000000000e+00, 1.337266601804649e+02, + 2.486725353510667e+05, 2.628059713728125e+07, + 2.436169536928770e+10, 1.913554568577315e+12, + 1.175208184614438e+15, 6.115751452473410e+16, + 2.791577695211466e+19, 7.241811142725384e+20, + 2.612380874940182e+23 + ] + xp_assert_close(b, b2, rtol=1e-6) + xp_assert_close(a, a2, rtol=1e-4) + + def test_fs_param(self): + for fs in (900, 900.1, 1234.567): + for N in (0, 1, 2, 3, 10): + for fc in (100, 100.1, 432.12345): + for btype in ('lp', 'hp'): + ba1 = ellip(N, 1, 20, fc, btype, fs=fs) + ba2 = ellip(N, 1, 20, fc/(fs/2), btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + for fc in ((100, 200), (100.1, 200.2), (321.123, 432.123)): + for btype in ('bp', 'bs'): + ba1 = ellip(N, 1, 20, fc, btype, fs=fs) + for seq in (list, tuple, array): + fcnorm = seq([f/(fs/2) for f in fc]) + ba2 = ellip(N, 1, 20, fcnorm, btype) + for ba1_, ba2_ in zip(ba1, ba2): + xp_assert_close(ba1_, ba2_) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + iirnotch(0.06, 30, fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none"): + iirnotch(0.06, 30, fs=None) + + +def test_sos_consistency(): + # Consistency checks of output='sos' for the specialized IIR filter + # design functions. + design_funcs = [(bessel, (0.1,)), + (butter, (0.1,)), + (cheby1, (45.0, 0.1)), + (cheby2, (0.087, 0.1)), + (ellip, (0.087, 45, 0.1))] + for func, args in design_funcs: + name = func.__name__ + + b, a = func(2, *args, output='ba') + sos = func(2, *args, output='sos') + xp_assert_close(sos, [np.hstack((b, a))], err_msg=f"{name}(2,...)") + + zpk = func(3, *args, output='zpk') + sos = func(3, *args, output='sos') + xp_assert_close(sos, zpk2sos(*zpk), err_msg=f"{name}(3,...)") + + zpk = func(4, *args, output='zpk') + sos = func(4, *args, output='sos') + xp_assert_close(sos, zpk2sos(*zpk), err_msg=f"{name}(4,...)") + + +class TestIIRNotch: + + def test_ba_output(self): + # Compare coefficients with Matlab ones + # for the equivalent input: + b, a = iirnotch(0.06, 30) + b2 = [ + 9.9686824e-01, -1.9584219e+00, + 9.9686824e-01 + ] + a2 = [ + 1.0000000e+00, -1.9584219e+00, + 9.9373647e-01 + ] + + xp_assert_close(b, b2, rtol=1e-8) + xp_assert_close(a, a2, rtol=1e-8) + + def test_frequency_response(self): + # Get filter coefficients + b, a = iirnotch(0.3, 30) + + # Get frequency response + w, h = freqz(b, a, 1000) + + # Pick 5 point + p = [200, # w0 = 0.200 + 295, # w0 = 0.295 + 300, # w0 = 0.300 + 305, # w0 = 0.305 + 400] # w0 = 0.400 + + # Get frequency response correspondent to each of those points + hp = h[p] + + # Check if the frequency response fulfill the specifications: + # hp[0] and hp[4] correspond to frequencies distant from + # w0 = 0.3 and should be close to 1 + xp_assert_close(abs(hp[0]), np.asarray(1.), rtol=1e-2, check_0d=False) + xp_assert_close(abs(hp[4]), np.asarray(1.), rtol=1e-2, check_0d=False) + + # hp[1] and hp[3] correspond to frequencies approximately + # on the edges of the passband and should be close to -3dB + xp_assert_close(abs(hp[1]), 1/np.sqrt(2), rtol=1e-2) + xp_assert_close(abs(hp[3]), 1/np.sqrt(2), rtol=1e-2) + + # hp[2] correspond to the frequency that should be removed + # the frequency response should be very close to 0 + xp_assert_close(abs(hp[2]), np.asarray(0.0), atol=1e-10, check_0d=False) + + def test_errors(self): + # Exception should be raised if w0 > 1 or w0 <0 + assert_raises(ValueError, iirnotch, w0=2, Q=30) + assert_raises(ValueError, iirnotch, w0=-1, Q=30) + + # Exception should be raised if any of the parameters + # are not float (or cannot be converted to one) + assert_raises(ValueError, iirnotch, w0="blabla", Q=30) + assert_raises(TypeError, iirnotch, w0=-1, Q=[1, 2, 3]) + + def test_fs_param(self): + # Get filter coefficients + b, a = iirnotch(1500, 30, fs=10000) + + # Get frequency response + w, h = freqz(b, a, 1000, fs=10000) + + # Pick 5 point + p = [200, # w0 = 1000 + 295, # w0 = 1475 + 300, # w0 = 1500 + 305, # w0 = 1525 + 400] # w0 = 2000 + + # Get frequency response correspondent to each of those points + hp = h[p] + + # Check if the frequency response fulfill the specifications: + # hp[0] and hp[4] correspond to frequencies distant from + # w0 = 1500 and should be close to 1 + xp_assert_close(abs(hp[0]), np.ones_like(abs(hp[0])), rtol=1e-2, + check_0d=False) + xp_assert_close(abs(hp[4]), np.ones_like(abs(hp[4])), rtol=1e-2, + check_0d=False) + + # hp[1] and hp[3] correspond to frequencies approximately + # on the edges of the passband and should be close to -3dB + xp_assert_close(abs(hp[1]), 1/np.sqrt(2), rtol=1e-2) + xp_assert_close(abs(hp[3]), 1/np.sqrt(2), rtol=1e-2) + + # hp[2] correspond to the frequency that should be removed + # the frequency response should be very close to 0 + xp_assert_close(abs(hp[2]), np.asarray(0.0), atol=1e-10, check_0d=False) + + +class TestIIRPeak: + + def test_ba_output(self): + # Compare coefficients with Matlab ones + # for the equivalent input: + b, a = iirpeak(0.06, 30) + b2 = [ + 3.131764229e-03, 0, + -3.131764229e-03 + ] + a2 = [ + 1.0000000e+00, -1.958421917e+00, + 9.9373647e-01 + ] + xp_assert_close(b, b2, rtol=1e-8) + xp_assert_close(a, a2, rtol=1e-8) + + def test_frequency_response(self): + # Get filter coefficients + b, a = iirpeak(0.3, 30) + + # Get frequency response + w, h = freqz(b, a, 1000) + + # Pick 5 point + p = [30, # w0 = 0.030 + 295, # w0 = 0.295 + 300, # w0 = 0.300 + 305, # w0 = 0.305 + 800] # w0 = 0.800 + + # Get frequency response correspondent to each of those points + hp = h[p] + + # Check if the frequency response fulfill the specifications: + # hp[0] and hp[4] correspond to frequencies distant from + # w0 = 0.3 and should be close to 0 + xp_assert_close(abs(hp[0]), + np.zeros_like(abs(hp[0])), atol=1e-2, check_0d=False) + xp_assert_close(abs(hp[4]), + np.zeros_like(abs(hp[4])), atol=1e-2, check_0d=False) + + # hp[1] and hp[3] correspond to frequencies approximately + # on the edges of the passband and should be close to 10**(-3/20) + xp_assert_close(abs(hp[1]), 1/np.sqrt(2), rtol=1e-2) + xp_assert_close(abs(hp[3]), 1/np.sqrt(2), rtol=1e-2) + + # hp[2] correspond to the frequency that should be retained and + # the frequency response should be very close to 1 + xp_assert_close(abs(hp[2]), np.asarray(1.0), rtol=1e-10, check_0d=False) + + def test_errors(self): + # Exception should be raised if w0 > 1 or w0 <0 + assert_raises(ValueError, iirpeak, w0=2, Q=30) + assert_raises(ValueError, iirpeak, w0=-1, Q=30) + + # Exception should be raised if any of the parameters + # are not float (or cannot be converted to one) + assert_raises(ValueError, iirpeak, w0="blabla", Q=30) + assert_raises(TypeError, iirpeak, w0=-1, Q=[1, 2, 3]) + + def test_fs_param(self): + # Get filter coefficients + b, a = iirpeak(1200, 30, fs=8000) + + # Get frequency response + w, h = freqz(b, a, 1000, fs=8000) + + # Pick 5 point + p = [30, # w0 = 120 + 295, # w0 = 1180 + 300, # w0 = 1200 + 305, # w0 = 1220 + 800] # w0 = 3200 + + # Get frequency response correspondent to each of those points + hp = h[p] + + # Check if the frequency response fulfill the specifications: + # hp[0] and hp[4] correspond to frequencies distant from + # w0 = 1200 and should be close to 0 + xp_assert_close(abs(hp[0]), + np.zeros_like(abs(hp[0])), atol=1e-2, check_0d=False) + xp_assert_close(abs(hp[4]), + np.zeros_like(abs(hp[4])), atol=1e-2, check_0d=False) + + # hp[1] and hp[3] correspond to frequencies approximately + # on the edges of the passband and should be close to 10**(-3/20) + xp_assert_close(abs(hp[1]), 1/np.sqrt(2), rtol=1e-2) + xp_assert_close(abs(hp[3]), 1/np.sqrt(2), rtol=1e-2) + + # hp[2] correspond to the frequency that should be retained and + # the frequency response should be very close to 1 + xp_assert_close(abs(hp[2]), + np.ones_like(abs(hp[2])), rtol=1e-10, check_0d=False) + + +class TestIIRComb: + # Test erroneous input cases + def test_invalid_input(self): + # w0 is <= 0 or >= fs / 2 + fs = 1000 + for args in [(-fs, 30), (0, 35), (fs / 2, 40), (fs, 35)]: + with pytest.raises(ValueError, match='w0 must be between '): + iircomb(*args, fs=fs) + + # fs is not divisible by w0 + for args in [(120, 30), (157, 35)]: + with pytest.raises(ValueError, match='fs must be divisible '): + iircomb(*args, fs=fs) + + # https://github.com/scipy/scipy/issues/14043#issuecomment-1107349140 + # Previously, fs=44100, w0=49.999 was rejected, but fs=2, + # w0=49.999/int(44100/2) was accepted. Now it is rejected, too. + with pytest.raises(ValueError, match='fs must be divisible '): + iircomb(w0=49.999/int(44100/2), Q=30) + + with pytest.raises(ValueError, match='fs must be divisible '): + iircomb(w0=49.999, Q=30, fs=44100) + + # Filter type is not notch or peak + for args in [(0.2, 30, 'natch'), (0.5, 35, 'comb')]: + with pytest.raises(ValueError, match='ftype must be '): + iircomb(*args) + + # Verify that the filter's frequency response contains a + # notch at the cutoff frequency + @pytest.mark.parametrize('ftype', ('notch', 'peak')) + def test_frequency_response(self, ftype): + # Create a notching or peaking comb filter at 1000 Hz + b, a = iircomb(1000, 30, ftype=ftype, fs=10000) + + # Compute the frequency response + freqs, response = freqz(b, a, 1000, fs=10000) + + # Find the notch using argrelextrema + comb_points = argrelextrema(abs(response), np.less)[0] + + # Verify that the first notch sits at 1000 Hz + comb1 = comb_points[0] + xp_assert_close(freqs[comb1], np.asarray(1000.), check_0d=False) + + # Verify pass_zero parameter + @pytest.mark.parametrize('ftype,pass_zero,peak,notch', + [('peak', True, 123.45, 61.725), + ('peak', False, 61.725, 123.45), + ('peak', None, 61.725, 123.45), + ('notch', None, 61.725, 123.45), + ('notch', True, 123.45, 61.725), + ('notch', False, 61.725, 123.45)]) + def test_pass_zero(self, ftype, pass_zero, peak, notch): + # Create a notching or peaking comb filter + b, a = iircomb(123.45, 30, ftype=ftype, fs=1234.5, pass_zero=pass_zero) + + # Compute the frequency response + freqs, response = freqz(b, a, [peak, notch], fs=1234.5) + + # Verify that expected notches are notches and peaks are peaks + assert abs(response[0]) > 0.99 + assert abs(response[1]) < 1e-10 + + # All built-in IIR filters are real, so should have perfectly + # symmetrical poles and zeros. Then ba representation (using + # numpy.poly) will be purely real instead of having negligible + # imaginary parts. + def test_iir_symmetry(self): + b, a = iircomb(400, 30, fs=24000) + z, p, k = tf2zpk(b, a) + xp_assert_equal(sorted(z), sorted(z.conj())) + xp_assert_equal(sorted(p), sorted(p.conj())) + xp_assert_equal(k, np.real(k)) + + assert issubclass(b.dtype.type, np.floating) + assert issubclass(a.dtype.type, np.floating) + + # Verify filter coefficients with MATLAB's iircomb function + def test_ba_output(self): + b_notch, a_notch = iircomb(60, 35, ftype='notch', fs=600) + b_notch2 = [0.957020174408697, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, -0.957020174408697] + a_notch2 = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, -0.914040348817395] + xp_assert_close(b_notch, b_notch2) + xp_assert_close(a_notch, a_notch2) + + b_peak, a_peak = iircomb(60, 35, ftype='peak', fs=600) + b_peak2 = [0.0429798255913026, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, -0.0429798255913026] + a_peak2 = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.914040348817395] + xp_assert_close(b_peak, b_peak2) + xp_assert_close(a_peak, a_peak2) + + # Verify that https://github.com/scipy/scipy/issues/14043 is fixed + def test_nearest_divisor(self): + # Create a notching comb filter + b, a = iircomb(50/int(44100/2), 50.0, ftype='notch') + + # Compute the frequency response at an upper harmonic of 50 + freqs, response = freqz(b, a, [22000], fs=44100) + + # Before bug fix, this would produce N = 881, so that 22 kHz was ~0 dB. + # Now N = 882 correctly and 22 kHz should be a notch <-220 dB + assert abs(response[0]) < 1e-10 + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + iircomb(1000, 30, fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none"): + iircomb(1000, 30, fs=None) + + +class TestIIRDesign: + + def test_exceptions(self): + with pytest.raises(ValueError, match="the same shape"): + iirdesign(0.2, [0.1, 0.3], 1, 40) + with pytest.raises(ValueError, match="the same shape"): + iirdesign(np.array([[0.3, 0.6], [0.3, 0.6]]), + np.array([[0.4, 0.5], [0.4, 0.5]]), 1, 40) + + # discrete filter with non-positive frequency + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign(0, 0.5, 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign(-0.1, 0.5, 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign(0.1, 0, 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign(0.1, -0.5, 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0, 0.3], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([-0.1, 0.3], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, -0.3], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0.3], [0, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0.3], [-0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0.3], [0.1, 0], 1, 40) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0.3], [0.1, -0.5], 1, 40) + + # analog filter with negative frequency + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign(-0.1, 0.5, 1, 40, analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign(0.1, -0.5, 1, 40, analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([-0.1, 0.3], [0.1, 0.5], 1, 40, analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, -0.3], [0.1, 0.5], 1, 40, analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0.3], [-0.1, 0.5], 1, 40, analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirdesign([0.1, 0.3], [0.1, -0.5], 1, 40, analog=True) + + # discrete filter with fs=None, freq > 1 + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign(1, 0.5, 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign(1.1, 0.5, 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign(0.1, 1, 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign(0.1, 1.5, 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([1, 0.3], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([1.1, 0.3], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([0.1, 1], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([0.1, 1.1], [0.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([0.1, 0.3], [1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([0.1, 0.3], [1.1, 0.5], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([0.1, 0.3], [0.1, 1], 1, 40) + with pytest.raises(ValueError, match="must be less than 1"): + iirdesign([0.1, 0.3], [0.1, 1.5], 1, 40) + + # discrete filter with fs>2, wp, ws < fs/2 must pass + iirdesign(100, 500, 1, 40, fs=2000) + iirdesign(500, 100, 1, 40, fs=2000) + iirdesign([200, 400], [100, 500], 1, 40, fs=2000) + iirdesign([100, 500], [200, 400], 1, 40, fs=2000) + + # discrete filter with fs>2, freq > fs/2: this must raise + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign(1000, 400, 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign(1100, 500, 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign(100, 1000, 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign(100, 1100, 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([1000, 400], [100, 500], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([1100, 400], [100, 500], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([200, 1000], [100, 500], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([200, 1100], [100, 500], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([200, 400], [1000, 500], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([200, 400], [1100, 500], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([200, 400], [100, 1000], 1, 40, fs=2000) + with pytest.raises(ValueError, match="must be less than fs/2"): + iirdesign([200, 400], [100, 1100], 1, 40, fs=2000) + + with pytest.raises(ValueError, match="strictly inside stopband"): + iirdesign([0.1, 0.4], [0.5, 0.6], 1, 40) + with pytest.raises(ValueError, match="strictly inside stopband"): + iirdesign([0.5, 0.6], [0.1, 0.4], 1, 40) + with pytest.raises(ValueError, match="strictly inside stopband"): + iirdesign([0.3, 0.6], [0.4, 0.7], 1, 40) + with pytest.raises(ValueError, match="strictly inside stopband"): + iirdesign([0.4, 0.7], [0.3, 0.6], 1, 40) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + iirfilter(1, 1, btype="low", fs=np.array([10, 20])) + + +class TestIIRFilter: + + def test_symmetry(self): + # All built-in IIR filters are real, so should have perfectly + # symmetrical poles and zeros. Then ba representation (using + # numpy.poly) will be purely real instead of having negligible + # imaginary parts. + for N in np.arange(1, 26): + for ftype in ('butter', 'bessel', 'cheby1', 'cheby2', 'ellip'): + z, p, k = iirfilter(N, 1.1, 1, 20, 'low', analog=True, + ftype=ftype, output='zpk') + xp_assert_equal(sorted(z), + sorted(z.conj())) + xp_assert_equal(sorted(p), + sorted(p.conj())) + xp_assert_equal(k, np.real(k)) + + b, a = iirfilter(N, 1.1, 1, 20, 'low', analog=True, + ftype=ftype, output='ba') + assert issubclass(b.dtype.type, np.floating) + assert issubclass(a.dtype.type, np.floating) + + def test_int_inputs(self): + # Using integer frequency arguments and large N should not produce + # numpy integers that wraparound to negative numbers + k = iirfilter(24, 100, btype='low', analog=True, ftype='bessel', + output='zpk')[2] + k2 = 9.999999999999989e+47 + xp_assert_close(np.asarray(k), np.asarray(k2)) + # if fs is specified then the normalization of Wn to have + # 0 <= Wn <= 1 should not cause an integer overflow + # the following line should not raise an exception + iirfilter(20, [1000000000, 1100000000], btype='bp', + analog=False, fs=6250000000) + + def test_invalid_wn_size(self): + # low and high have 1 Wn, band and stop have 2 Wn + assert_raises(ValueError, iirfilter, 1, [0.1, 0.9], btype='low') + assert_raises(ValueError, iirfilter, 1, [0.2, 0.5], btype='high') + assert_raises(ValueError, iirfilter, 1, 0.2, btype='bp') + assert_raises(ValueError, iirfilter, 1, 400, btype='bs', analog=True) + + def test_invalid_wn_range(self): + # For digital filters, 0 <= Wn <= 1 + assert_raises(ValueError, iirfilter, 1, 2, btype='low') + assert_raises(ValueError, iirfilter, 1, [0.5, 1], btype='band') + assert_raises(ValueError, iirfilter, 1, [0., 0.5], btype='band') + assert_raises(ValueError, iirfilter, 1, -1, btype='high') + assert_raises(ValueError, iirfilter, 1, [1, 2], btype='band') + assert_raises(ValueError, iirfilter, 1, [10, 20], btype='stop') + + # analog=True with non-positive critical frequencies + with pytest.raises(ValueError, match="must be greater than 0"): + iirfilter(2, 0, btype='low', analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirfilter(2, -1, btype='low', analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirfilter(2, [0, 100], analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirfilter(2, [-1, 100], analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirfilter(2, [10, 0], analog=True) + with pytest.raises(ValueError, match="must be greater than 0"): + iirfilter(2, [10, -1], analog=True) + + def test_analog_sos(self): + # first order Butterworth filter with Wn = 1 has tf 1/(s+1) + sos = [[0., 0., 1., 0., 1., 1.]] + sos2 = iirfilter(N=1, Wn=1, btype='low', analog=True, output='sos') + assert_array_almost_equal(sos, sos2) + + def test_wn1_ge_wn0(self): + # gh-15773: should raise error if Wn[0] >= Wn[1] + with pytest.raises(ValueError, + match=r"Wn\[0\] must be less than Wn\[1\]"): + iirfilter(2, [0.5, 0.5]) + with pytest.raises(ValueError, + match=r"Wn\[0\] must be less than Wn\[1\]"): + iirfilter(2, [0.6, 0.5]) + + +class TestGroupDelay: + def test_identity_filter(self): + w, gd = group_delay((1, 1)) + assert_array_almost_equal(w, pi * np.arange(512) / 512) + assert_array_almost_equal(gd, np.zeros(512)) + w, gd = group_delay((1, 1), whole=True) + assert_array_almost_equal(w, 2 * pi * np.arange(512) / 512) + assert_array_almost_equal(gd, np.zeros(512)) + + def test_fir(self): + # Let's design linear phase FIR and check that the group delay + # is constant. + N = 100 + b = firwin(N + 1, 0.1) + w, gd = group_delay((b, 1)) + xp_assert_close(gd, np.ones_like(gd)*(0.5 * N)) + + def test_iir(self): + # Let's design Butterworth filter and test the group delay at + # some points against MATLAB answer. + b, a = butter(4, 0.1) + w = np.linspace(0, pi, num=10, endpoint=False) + w, gd = group_delay((b, a), w=w) + matlab_gd = np.array([8.249313898506037, 11.958947880907104, + 2.452325615326005, 1.048918665702008, + 0.611382575635897, 0.418293269460578, + 0.317932917836572, 0.261371844762525, + 0.229038045801298, 0.212185774208521]) + assert_array_almost_equal(gd, matlab_gd) + + @pytest.mark.thread_unsafe + def test_singular(self): + # Let's create a filter with zeros and poles on the unit circle and + # check if warnings are raised at those frequencies. + z1 = np.exp(1j * 0.1 * pi) + z2 = np.exp(1j * 0.25 * pi) + p1 = np.exp(1j * 0.5 * pi) + p2 = np.exp(1j * 0.8 * pi) + b = np.convolve([1, -z1], [1, -z2]) + a = np.convolve([1, -p1], [1, -p2]) + w = np.array([0.1 * pi, 0.25 * pi, -0.5 * pi, -0.8 * pi]) + + w, gd = assert_warns(UserWarning, group_delay, (b, a), w=w) + + def test_backward_compat(self): + # For backward compatibility, test if None act as a wrapper for default + w1, gd1 = group_delay((1, 1)) + w2, gd2 = group_delay((1, 1), None) + assert_array_almost_equal(w1, w2) + assert_array_almost_equal(gd1, gd2) + + def test_fs_param(self): + # Let's design Butterworth filter and test the group delay at + # some points against the normalized frequency answer. + b, a = butter(4, 4800, fs=96000) + w = np.linspace(0, 96000/2, num=10, endpoint=False) + w, gd = group_delay((b, a), w=w, fs=96000) + norm_gd = np.array([8.249313898506037, 11.958947880907104, + 2.452325615326005, 1.048918665702008, + 0.611382575635897, 0.418293269460578, + 0.317932917836572, 0.261371844762525, + 0.229038045801298, 0.212185774208521]) + assert_array_almost_equal(gd, norm_gd) + + def test_w_or_N_types(self): + # Measure at 8 equally-spaced points + for N in (8, np.int8(8), np.int16(8), np.int32(8), np.int64(8), + np.array(8)): + w, gd = group_delay((1, 1), N) + assert_array_almost_equal(w, pi * np.arange(8) / 8) + assert_array_almost_equal(gd, np.zeros(8)) + + # Measure at frequency 8 rad/sec + for w in (8.0, 8.0+0j): + w_out, gd = group_delay((1, 1), w) + assert_array_almost_equal(w_out, [8]) + assert_array_almost_equal(gd, [0]) + + def test_complex_coef(self): + # gh-19586: handle complex coef TFs + # + # for g(z) = (alpha*z+1)/(1+conjugate(alpha)), group delay is + # given by function below. + # + # def gd_expr(w, alpha): + # num = 1j*(abs(alpha)**2-1)*np.exp(1j*w) + # den = (alpha*np.exp(1j*w)+1)*(np.exp(1j*w)+np.conj(alpha)) + # return -np.imag(num/den) + + # arbitrary non-real alpha + alpha = -0.6143077933232609+0.3355978770229421j + # 8 points from from -pi to pi + wref = np.array([-3.141592653589793 , + -2.356194490192345 , + -1.5707963267948966, + -0.7853981633974483, + 0. , + 0.7853981633974483, + 1.5707963267948966, + 2.356194490192345 ]) + gdref = array([0.18759548150354619, + 0.17999770352712252, + 0.23598047471879877, + 0.46539443069907194, + 1.9511492420564165 , + 3.478129975138865 , + 0.6228594960517333 , + 0.27067831839471224]) + b = [alpha,1] + a = [1, np.conjugate(alpha)] + gdtest = group_delay((b,a), wref)[1] + # need nulp=14 for macOS arm64 wheel builds; added 2 for some + # robustness on other platforms. + assert_array_almost_equal_nulp(gdtest, gdref, nulp=16) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + group_delay((1, 1), fs=np.array([10, 20])) + + with pytest.raises(ValueError, match="Sampling.*be none"): + group_delay((1, 1), fs=None) + + +class TestGammatone: + # Test erroneous input cases. + def test_invalid_input(self): + # Cutoff frequency is <= 0 or >= fs / 2. + fs = 16000 + for args in [(-fs, 'iir'), (0, 'fir'), (fs / 2, 'iir'), (fs, 'fir')]: + with pytest.raises(ValueError, match='The frequency must be ' + 'between '): + gammatone(*args, fs=fs) + + # Filter type is not fir or iir + for args in [(440, 'fie'), (220, 'it')]: + with pytest.raises(ValueError, match='ftype must be '): + gammatone(*args, fs=fs) + + # Order is <= 0 or > 24 for FIR filter. + for args in [(440, 'fir', -50), (220, 'fir', 0), (110, 'fir', 25), + (55, 'fir', 50)]: + with pytest.raises(ValueError, match='Invalid order: '): + gammatone(*args, numtaps=None, fs=fs) + + # Verify that the filter's frequency response is approximately + # 1 at the cutoff frequency. + def test_frequency_response(self): + fs = 16000 + ftypes = ['fir', 'iir'] + for ftype in ftypes: + # Create a gammatone filter centered at 1000 Hz. + b, a = gammatone(1000, ftype, fs=fs) + + # Calculate the frequency response. + freqs, response = freqz(b, a) + + # Determine peak magnitude of the response + # and corresponding frequency. + response_max = np.max(np.abs(response)) + freq_hz = freqs[np.argmax(np.abs(response))] / ((2 * np.pi) / fs) + + # Check that the peak magnitude is 1 and the frequency is 1000 Hz. + xp_assert_close(response_max, + np.ones_like(response_max), rtol=1e-2, check_0d=False) + xp_assert_close(freq_hz, + 1000*np.ones_like(freq_hz), rtol=1e-2, check_0d=False) + + # All built-in IIR filters are real, so should have perfectly + # symmetrical poles and zeros. Then ba representation (using + # numpy.poly) will be purely real instead of having negligible + # imaginary parts. + def test_iir_symmetry(self): + b, a = gammatone(440, 'iir', fs=24000) + z, p, k = tf2zpk(b, a) + xp_assert_equal(sorted(z), sorted(z.conj())) + xp_assert_equal(sorted(p), sorted(p.conj())) + xp_assert_equal(k, np.real(k)) + + assert issubclass(b.dtype.type, np.floating) + assert issubclass(a.dtype.type, np.floating) + + # Verify FIR filter coefficients with the paper's + # Mathematica implementation + def test_fir_ba_output(self): + b, _ = gammatone(15, 'fir', fs=1000) + b2 = [0.0, 2.2608075649884e-04, + 1.5077903981357e-03, 4.2033687753998e-03, + 8.1508962726503e-03, 1.2890059089154e-02, + 1.7833890391666e-02, 2.2392613558564e-02, + 2.6055195863104e-02, 2.8435872863284e-02, + 2.9293319149544e-02, 2.852976858014e-02, + 2.6176557156294e-02, 2.2371510270395e-02, + 1.7332485267759e-02] + xp_assert_close(b, b2) + + # Verify IIR filter coefficients with the paper's MATLAB implementation + def test_iir_ba_output(self): + b, a = gammatone(440, 'iir', fs=16000) + b2 = [1.31494461367464e-06, -5.03391196645395e-06, + 7.00649426000897e-06, -4.18951968419854e-06, + 9.02614910412011e-07] + a2 = [1.0, -7.65646235454218, + 25.7584699322366, -49.7319214483238, + 60.2667361289181, -46.9399590980486, + 22.9474798808461, -6.43799381299034, + 0.793651554625368] + xp_assert_close(b, b2) + xp_assert_close(a, a2) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + gammatone(440, 'iir', fs=np.array([10, 20])) + + +class TestOrderFilter: + def test_doc_example(self): + x = np.arange(25).reshape(5, 5) + domain = np.identity(3) + + # minimum of elements 1,3,9 (zero-padded) on phone pad + # 7,5,3 on numpad + expected = np.array( + [[0., 0., 0., 0., 0.], + [0., 0., 1., 2., 0.], + [0., 5., 6., 7., 0.], + [0., 10., 11., 12., 0.], + [0., 0., 0., 0., 0.]], + ) + xp_assert_close(order_filter(x, domain, 0), expected, check_dtype=False) + + # maximum of elements 1,3,9 (zero-padded) on phone pad + # 7,5,3 on numpad + expected = np.array( + [[6., 7., 8., 9., 4.], + [11., 12., 13., 14., 9.], + [16., 17., 18., 19., 14.], + [21., 22., 23., 24., 19.], + [20., 21., 22., 23., 24.]], + ) + xp_assert_close(order_filter(x, domain, 2), expected, check_dtype=False) + + # and, just to complete the set, median of zero-padded elements + expected = np.array( + [[0, 1, 2, 3, 0], + [5, 6, 7, 8, 3], + [10, 11, 12, 13, 8], + [15, 16, 17, 18, 13], + [0, 15, 16, 17, 18]], + ) + xp_assert_close(order_filter(x, domain, 1), expected) + + def test_medfilt_order_filter(self): + x = np.arange(25).reshape(5, 5) + + # median of zero-padded elements 1,5,9 on phone pad + # 7,5,3 on numpad + expected = np.array( + [[0, 1, 2, 3, 0], + [1, 6, 7, 8, 4], + [6, 11, 12, 13, 9], + [11, 16, 17, 18, 14], + [0, 16, 17, 18, 0]], + ) + xp_assert_close(medfilt(x, 3), expected) + + xp_assert_close( + order_filter(x, np.ones((3, 3)), 4), + expected + ) + + def test_order_filter_asymmetric(self): + x = np.arange(25).reshape(5, 5) + domain = np.array( + [[1, 1, 0], + [0, 1, 0], + [0, 0, 0]], + ) + + expected = np.array( + [[0, 0, 0, 0, 0], + [0, 0, 1, 2, 3], + [0, 5, 6, 7, 8], + [0, 10, 11, 12, 13], + [0, 15, 16, 17, 18]] + ) + xp_assert_close(order_filter(x, domain, 0), expected) + + expected = np.array( + [[0, 0, 0, 0, 0], + [0, 1, 2, 3, 4], + [5, 6, 7, 8, 9], + [10, 11, 12, 13, 14], + [15, 16, 17, 18, 19]] + ) + xp_assert_close(order_filter(x, domain, 1), expected) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_fir_filter_design.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_fir_filter_design.py new file mode 100644 index 0000000000000000000000000000000000000000..d245fc875784df2e32203b675e3f37f5570c3549 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_fir_filter_design.py @@ -0,0 +1,654 @@ +import numpy as np +from numpy.testing import assert_warns +from scipy._lib._array_api import ( + xp_assert_close, xp_assert_equal, + assert_almost_equal, assert_array_almost_equal, +) +from pytest import raises as assert_raises +import pytest + +from scipy.fft import fft +from scipy.special import sinc +from scipy.signal import (kaiser_beta, kaiser_atten, kaiserord, + firwin, firwin2, freqz, remez, firls, minimum_phase +) + + +def test_kaiser_beta(): + b = kaiser_beta(58.7) + assert_almost_equal(b, 0.1102 * 50.0) + b = kaiser_beta(22.0) + assert_almost_equal(b, 0.5842 + 0.07886) + b = kaiser_beta(21.0) + assert b == 0.0 + b = kaiser_beta(10.0) + assert b == 0.0 + + +def test_kaiser_atten(): + a = kaiser_atten(1, 1.0) + assert a == 7.95 + a = kaiser_atten(2, 1/np.pi) + assert a == 2.285 + 7.95 + + +def test_kaiserord(): + assert_raises(ValueError, kaiserord, 1.0, 1.0) + numtaps, beta = kaiserord(2.285 + 7.95 - 0.001, 1/np.pi) + assert (numtaps, beta) == (2, 0.0) + + +class TestFirwin: + + def check_response(self, h, expected_response, tol=.05): + N = len(h) + alpha = 0.5 * (N-1) + m = np.arange(0,N) - alpha # time indices of taps + for freq, expected in expected_response: + actual = abs(np.sum(h*np.exp(-1.j*np.pi*m*freq))) + mse = abs(actual-expected)**2 + assert mse < tol, f'response not as expected, mse={mse:g} > {tol:g}' + + def test_response(self): + N = 51 + f = .5 + # increase length just to try even/odd + h = firwin(N, f) # low-pass from 0 to f + self.check_response(h, [(.25,1), (.75,0)]) + + h = firwin(N+1, f, window='nuttall') # specific window + self.check_response(h, [(.25,1), (.75,0)]) + + h = firwin(N+2, f, pass_zero=False) # stop from 0 to f --> high-pass + self.check_response(h, [(.25,0), (.75,1)]) + + f1, f2, f3, f4 = .2, .4, .6, .8 + h = firwin(N+3, [f1, f2], pass_zero=False) # band-pass filter + self.check_response(h, [(.1,0), (.3,1), (.5,0)]) + + h = firwin(N+4, [f1, f2]) # band-stop filter + self.check_response(h, [(.1,1), (.3,0), (.5,1)]) + + h = firwin(N+5, [f1, f2, f3, f4], pass_zero=False, scale=False) + self.check_response(h, [(.1,0), (.3,1), (.5,0), (.7,1), (.9,0)]) + + h = firwin(N+6, [f1, f2, f3, f4]) # multiband filter + self.check_response(h, [(.1,1), (.3,0), (.5,1), (.7,0), (.9,1)]) + + h = firwin(N+7, 0.1, width=.03) # low-pass + self.check_response(h, [(.05,1), (.75,0)]) + + h = firwin(N+8, 0.1, pass_zero=False) # high-pass + self.check_response(h, [(.05,0), (.75,1)]) + + def mse(self, h, bands): + """Compute mean squared error versus ideal response across frequency + band. + h -- coefficients + bands -- list of (left, right) tuples relative to 1==Nyquist of + passbands + """ + w, H = freqz(h, worN=1024) + f = w/np.pi + passIndicator = np.zeros(len(w), bool) + for left, right in bands: + passIndicator |= (f >= left) & (f < right) + Hideal = np.where(passIndicator, 1, 0) + mse = np.mean(abs(abs(H)-Hideal)**2) + return mse + + def test_scaling(self): + """ + For one lowpass, bandpass, and highpass example filter, this test + checks two things: + - the mean squared error over the frequency domain of the unscaled + filter is smaller than the scaled filter (true for rectangular + window) + - the response of the scaled filter is exactly unity at the center + of the first passband + """ + N = 11 + cases = [ + ([.5], True, (0, 1)), + ([0.2, .6], False, (.4, 1)), + ([.5], False, (1, 1)), + ] + for cutoff, pass_zero, expected_response in cases: + h = firwin(N, cutoff, scale=False, pass_zero=pass_zero, window='ones') + hs = firwin(N, cutoff, scale=True, pass_zero=pass_zero, window='ones') + if len(cutoff) == 1: + if pass_zero: + cutoff = [0] + cutoff + else: + cutoff = cutoff + [1] + msg = 'least squares violation' + assert self.mse(h, [cutoff]) < self.mse(hs, [cutoff]), msg + self.check_response(hs, [expected_response], 1e-12) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + firwin(51, .5, fs=np.array([10, 20])) + + +class TestFirWinMore: + """Different author, different style, different tests...""" + + def test_lowpass(self): + width = 0.04 + ntaps, beta = kaiserord(120, width) + kwargs = dict(cutoff=0.5, window=('kaiser', beta), scale=False) + taps = firwin(ntaps, **kwargs) + + # Check the symmetry of taps. + assert_array_almost_equal(taps[:ntaps//2], taps[ntaps:ntaps-ntaps//2-1:-1]) + + # Check the gain at a few samples where + # we know it should be approximately 0 or 1. + freq_samples = np.array([0.0, 0.25, 0.5-width/2, 0.5+width/2, 0.75, 1.0]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [1.0, 1.0, 1.0, 0.0, 0.0, 0.0], decimal=5) + + taps_str = firwin(ntaps, pass_zero='lowpass', **kwargs) + xp_assert_close(taps, taps_str) + + def test_highpass(self): + width = 0.04 + ntaps, beta = kaiserord(120, width) + + # Ensure that ntaps is odd. + ntaps |= 1 + + kwargs = dict(cutoff=0.5, window=('kaiser', beta), scale=False) + taps = firwin(ntaps, pass_zero=False, **kwargs) + + # Check the symmetry of taps. + assert_array_almost_equal(taps[:ntaps//2], taps[ntaps:ntaps-ntaps//2-1:-1]) + + # Check the gain at a few samples where + # we know it should be approximately 0 or 1. + freq_samples = np.array([0.0, 0.25, 0.5-width/2, 0.5+width/2, 0.75, 1.0]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], decimal=5) + + taps_str = firwin(ntaps, pass_zero='highpass', **kwargs) + xp_assert_close(taps, taps_str) + + def test_bandpass(self): + width = 0.04 + ntaps, beta = kaiserord(120, width) + kwargs = dict(cutoff=[0.3, 0.7], window=('kaiser', beta), scale=False) + taps = firwin(ntaps, pass_zero=False, **kwargs) + + # Check the symmetry of taps. + assert_array_almost_equal(taps[:ntaps//2], taps[ntaps:ntaps-ntaps//2-1:-1]) + + # Check the gain at a few samples where + # we know it should be approximately 0 or 1. + freq_samples = np.array([0.0, 0.2, 0.3-width/2, 0.3+width/2, 0.5, + 0.7-width/2, 0.7+width/2, 0.8, 1.0]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0], decimal=5) + + taps_str = firwin(ntaps, pass_zero='bandpass', **kwargs) + xp_assert_close(taps, taps_str) + + def test_bandstop_multi(self): + width = 0.04 + ntaps, beta = kaiserord(120, width) + kwargs = dict(cutoff=[0.2, 0.5, 0.8], window=('kaiser', beta), + scale=False) + taps = firwin(ntaps, **kwargs) + + # Check the symmetry of taps. + assert_array_almost_equal(taps[:ntaps//2], taps[ntaps:ntaps-ntaps//2-1:-1]) + + # Check the gain at a few samples where + # we know it should be approximately 0 or 1. + freq_samples = np.array([0.0, 0.1, 0.2-width/2, 0.2+width/2, 0.35, + 0.5-width/2, 0.5+width/2, 0.65, + 0.8-width/2, 0.8+width/2, 0.9, 1.0]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0], + decimal=5) + + taps_str = firwin(ntaps, pass_zero='bandstop', **kwargs) + xp_assert_close(taps, taps_str) + + def test_fs_nyq(self): + """Test the fs and nyq keywords.""" + nyquist = 1000 + width = 40.0 + relative_width = width/nyquist + ntaps, beta = kaiserord(120, relative_width) + taps = firwin(ntaps, cutoff=[300, 700], window=('kaiser', beta), + pass_zero=False, scale=False, fs=2*nyquist) + + # Check the symmetry of taps. + assert_array_almost_equal(taps[:ntaps//2], taps[ntaps:ntaps-ntaps//2-1:-1]) + + # Check the gain at a few samples where + # we know it should be approximately 0 or 1. + freq_samples = np.array([0.0, 200, 300-width/2, 300+width/2, 500, + 700-width/2, 700+width/2, 800, 1000]) + freqs, response = freqz(taps, worN=np.pi*freq_samples/nyquist) + assert_array_almost_equal(np.abs(response), + [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0], decimal=5) + + def test_bad_cutoff(self): + """Test that invalid cutoff argument raises ValueError.""" + # cutoff values must be greater than 0 and less than 1. + assert_raises(ValueError, firwin, 99, -0.5) + assert_raises(ValueError, firwin, 99, 1.5) + # Don't allow 0 or 1 in cutoff. + assert_raises(ValueError, firwin, 99, [0, 0.5]) + assert_raises(ValueError, firwin, 99, [0.5, 1]) + # cutoff values must be strictly increasing. + assert_raises(ValueError, firwin, 99, [0.1, 0.5, 0.2]) + assert_raises(ValueError, firwin, 99, [0.1, 0.5, 0.5]) + # Must have at least one cutoff value. + assert_raises(ValueError, firwin, 99, []) + # 2D array not allowed. + assert_raises(ValueError, firwin, 99, [[0.1, 0.2],[0.3, 0.4]]) + # cutoff values must be less than nyq. + assert_raises(ValueError, firwin, 99, 50.0, fs=80) + assert_raises(ValueError, firwin, 99, [10, 20, 30], fs=50) + + def test_even_highpass_raises_value_error(self): + """Test that attempt to create a highpass filter with an even number + of taps raises a ValueError exception.""" + assert_raises(ValueError, firwin, 40, 0.5, pass_zero=False) + assert_raises(ValueError, firwin, 40, [.25, 0.5]) + + def test_bad_pass_zero(self): + """Test degenerate pass_zero cases.""" + with assert_raises(ValueError, match='pass_zero must be'): + firwin(41, 0.5, pass_zero='foo') + with assert_raises(TypeError, match='cannot be interpreted'): + firwin(41, 0.5, pass_zero=1.) + for pass_zero in ('lowpass', 'highpass'): + with assert_raises(ValueError, match='cutoff must have one'): + firwin(41, [0.5, 0.6], pass_zero=pass_zero) + for pass_zero in ('bandpass', 'bandstop'): + with assert_raises(ValueError, match='must have at least two'): + firwin(41, [0.5], pass_zero=pass_zero) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + firwin2(51, .5, 1, fs=np.array([10, 20])) + + +class TestFirwin2: + + def test_invalid_args(self): + # `freq` and `gain` have different lengths. + with assert_raises(ValueError, match='must be of same length'): + firwin2(50, [0, 0.5, 1], [0.0, 1.0]) + # `nfreqs` is less than `ntaps`. + with assert_raises(ValueError, match='ntaps must be less than nfreqs'): + firwin2(50, [0, 0.5, 1], [0.0, 1.0, 1.0], nfreqs=33) + # Decreasing value in `freq` + with assert_raises(ValueError, match='must be nondecreasing'): + firwin2(50, [0, 0.5, 0.4, 1.0], [0, .25, .5, 1.0]) + # Value in `freq` repeated more than once. + with assert_raises(ValueError, match='must not occur more than twice'): + firwin2(50, [0, .1, .1, .1, 1.0], [0.0, 0.5, 0.75, 1.0, 1.0]) + # `freq` does not start at 0.0. + with assert_raises(ValueError, match='start with 0'): + firwin2(50, [0.5, 1.0], [0.0, 1.0]) + # `freq` does not end at fs/2. + with assert_raises(ValueError, match='end with fs/2'): + firwin2(50, [0.0, 0.5], [0.0, 1.0]) + # Value 0 is repeated in `freq` + with assert_raises(ValueError, match='0 must not be repeated'): + firwin2(50, [0.0, 0.0, 0.5, 1.0], [1.0, 1.0, 0.0, 0.0]) + # Value fs/2 is repeated in `freq` + with assert_raises(ValueError, match='fs/2 must not be repeated'): + firwin2(50, [0.0, 0.5, 1.0, 1.0], [1.0, 1.0, 0.0, 0.0]) + # Value in `freq` that is too close to a repeated number + with assert_raises(ValueError, match='cannot contain numbers ' + 'that are too close'): + firwin2(50, [0.0, 0.5 - np.finfo(float).eps * 0.5, 0.5, 0.5, 1.0], + [1.0, 1.0, 1.0, 0.0, 0.0]) + + # Type II filter, but the gain at nyquist frequency is not zero. + with assert_raises(ValueError, match='Type II filter'): + firwin2(16, [0.0, 0.5, 1.0], [0.0, 1.0, 1.0]) + + # Type III filter, but the gains at nyquist and zero rate are not zero. + with assert_raises(ValueError, match='Type III filter'): + firwin2(17, [0.0, 0.5, 1.0], [0.0, 1.0, 1.0], antisymmetric=True) + with assert_raises(ValueError, match='Type III filter'): + firwin2(17, [0.0, 0.5, 1.0], [1.0, 1.0, 0.0], antisymmetric=True) + with assert_raises(ValueError, match='Type III filter'): + firwin2(17, [0.0, 0.5, 1.0], [1.0, 1.0, 1.0], antisymmetric=True) + + # Type IV filter, but the gain at zero rate is not zero. + with assert_raises(ValueError, match='Type IV filter'): + firwin2(16, [0.0, 0.5, 1.0], [1.0, 1.0, 0.0], antisymmetric=True) + + def test01(self): + width = 0.04 + beta = 12.0 + ntaps = 400 + # Filter is 1 from w=0 to w=0.5, then decreases linearly from 1 to 0 as w + # increases from w=0.5 to w=1 (w=1 is the Nyquist frequency). + freq = [0.0, 0.5, 1.0] + gain = [1.0, 1.0, 0.0] + taps = firwin2(ntaps, freq, gain, window=('kaiser', beta)) + freq_samples = np.array([0.0, 0.25, 0.5-width/2, 0.5+width/2, + 0.75, 1.0-width/2]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [1.0, 1.0, 1.0, 1.0-width, 0.5, width], decimal=5) + + def test02(self): + width = 0.04 + beta = 12.0 + # ntaps must be odd for positive gain at Nyquist. + ntaps = 401 + # An ideal highpass filter. + freq = [0.0, 0.5, 0.5, 1.0] + gain = [0.0, 0.0, 1.0, 1.0] + taps = firwin2(ntaps, freq, gain, window=('kaiser', beta)) + freq_samples = np.array([0.0, 0.25, 0.5-width, 0.5+width, 0.75, 1.0]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], decimal=5) + + def test03(self): + width = 0.02 + ntaps, beta = kaiserord(120, width) + # ntaps must be odd for positive gain at Nyquist. + ntaps = int(ntaps) | 1 + freq = [0.0, 0.4, 0.4, 0.5, 0.5, 1.0] + gain = [1.0, 1.0, 0.0, 0.0, 1.0, 1.0] + taps = firwin2(ntaps, freq, gain, window=('kaiser', beta)) + freq_samples = np.array([0.0, 0.4-width, 0.4+width, 0.45, + 0.5-width, 0.5+width, 0.75, 1.0]) + freqs, response = freqz(taps, worN=np.pi*freq_samples) + assert_array_almost_equal(np.abs(response), + [1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0], decimal=5) + + def test04(self): + """Test firwin2 when window=None.""" + ntaps = 5 + # Ideal lowpass: gain is 1 on [0,0.5], and 0 on [0.5, 1.0] + freq = [0.0, 0.5, 0.5, 1.0] + gain = [1.0, 1.0, 0.0, 0.0] + taps = firwin2(ntaps, freq, gain, window=None, nfreqs=8193) + alpha = 0.5 * (ntaps - 1) + m = np.arange(0, ntaps) - alpha + h = 0.5 * sinc(0.5 * m) + assert_array_almost_equal(h, taps) + + def test05(self): + """Test firwin2 for calculating Type IV filters""" + ntaps = 1500 + + freq = [0.0, 1.0] + gain = [0.0, 1.0] + taps = firwin2(ntaps, freq, gain, window=None, antisymmetric=True) + assert_array_almost_equal(taps[: ntaps // 2], -taps[ntaps // 2:][::-1]) + + freqs, response = freqz(taps, worN=2048) + assert_array_almost_equal(abs(response), freqs / np.pi, decimal=4) + + def test06(self): + """Test firwin2 for calculating Type III filters""" + ntaps = 1501 + + freq = [0.0, 0.5, 0.55, 1.0] + gain = [0.0, 0.5, 0.0, 0.0] + taps = firwin2(ntaps, freq, gain, window=None, antisymmetric=True) + assert taps[ntaps // 2] == 0.0 + assert_array_almost_equal(taps[: ntaps // 2], -taps[ntaps // 2 + 1:][::-1]) + + freqs, response1 = freqz(taps, worN=2048) + response2 = np.interp(freqs / np.pi, freq, gain) + assert_array_almost_equal(abs(response1), response2, decimal=3) + + def test_fs_nyq(self): + taps1 = firwin2(80, [0.0, 0.5, 1.0], [1.0, 1.0, 0.0]) + taps2 = firwin2(80, [0.0, 30.0, 60.0], [1.0, 1.0, 0.0], fs=120.0) + assert_array_almost_equal(taps1, taps2) + + def test_tuple(self): + taps1 = firwin2(150, (0.0, 0.5, 0.5, 1.0), (1.0, 1.0, 0.0, 0.0)) + taps2 = firwin2(150, [0.0, 0.5, 0.5, 1.0], [1.0, 1.0, 0.0, 0.0]) + assert_array_almost_equal(taps1, taps2) + + def test_input_modyfication(self): + freq1 = np.array([0.0, 0.5, 0.5, 1.0]) + freq2 = np.array(freq1) + firwin2(80, freq1, [1.0, 1.0, 0.0, 0.0]) + xp_assert_equal(freq1, freq2) + + +class TestRemez: + + def test_bad_args(self): + assert_raises(ValueError, remez, 11, [0.1, 0.4], [1], type='pooka') + + def test_hilbert(self): + N = 11 # number of taps in the filter + a = 0.1 # width of the transition band + + # design an unity gain hilbert bandpass filter from w to 0.5-w + h = remez(11, [a, 0.5-a], [1], type='hilbert') + + # make sure the filter has correct # of taps + assert len(h) == N, "Number of Taps" + + # make sure it is type III (anti-symmetric tap coefficients) + assert_array_almost_equal(h[:(N-1)//2], -h[:-(N-1)//2-1:-1]) + + # Since the requested response is symmetric, all even coefficients + # should be zero (or in this case really small) + assert (abs(h[1::2]) < 1e-15).all(), "Even Coefficients Equal Zero" + + # now check the frequency response + w, H = freqz(h, 1) + f = w/2/np.pi + Hmag = abs(H) + + # should have a zero at 0 and pi (in this case close to zero) + assert (Hmag[[0, -1]] < 0.02).all(), "Zero at zero and pi" + + # check that the pass band is close to unity + idx = np.logical_and(f > a, f < 0.5-a) + assert (abs(Hmag[idx] - 1) < 0.015).all(), "Pass Band Close To Unity" + + def test_compare(self): + # test comparison to MATLAB + k = [0.024590270518440, -0.041314581814658, -0.075943803756711, + -0.003530911231040, 0.193140296954975, 0.373400753484939, + 0.373400753484939, 0.193140296954975, -0.003530911231040, + -0.075943803756711, -0.041314581814658, 0.024590270518440] + h = remez(12, [0, 0.3, 0.5, 1], [1, 0], fs=2.) + xp_assert_close(h, k) + + h = [-0.038976016082299, 0.018704846485491, -0.014644062687875, + 0.002879152556419, 0.016849978528150, -0.043276706138248, + 0.073641298245579, -0.103908158578635, 0.129770906801075, + -0.147163447297124, 0.153302248456347, -0.147163447297124, + 0.129770906801075, -0.103908158578635, 0.073641298245579, + -0.043276706138248, 0.016849978528150, 0.002879152556419, + -0.014644062687875, 0.018704846485491, -0.038976016082299] + xp_assert_close(remez(21, [0, 0.8, 0.9, 1], [0, 1], fs=2.), h) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + remez(11, .1, 1, fs=np.array([10, 20])) + +class TestFirls: + + def test_bad_args(self): + # even numtaps + assert_raises(ValueError, firls, 10, [0.1, 0.2], [0, 0]) + # odd bands + assert_raises(ValueError, firls, 11, [0.1, 0.2, 0.4], [0, 0, 0]) + # len(bands) != len(desired) + assert_raises(ValueError, firls, 11, [0.1, 0.2, 0.3, 0.4], [0, 0, 0]) + # non-monotonic bands + assert_raises(ValueError, firls, 11, [0.2, 0.1], [0, 0]) + assert_raises(ValueError, firls, 11, [0.1, 0.2, 0.3, 0.3], [0] * 4) + assert_raises(ValueError, firls, 11, [0.3, 0.4, 0.1, 0.2], [0] * 4) + assert_raises(ValueError, firls, 11, [0.1, 0.3, 0.2, 0.4], [0] * 4) + # negative desired + assert_raises(ValueError, firls, 11, [0.1, 0.2], [-1, 1]) + # len(weight) != len(pairs) + assert_raises(ValueError, firls, 11, [0.1, 0.2], [0, 0], weight=[1, 2]) + # negative weight + assert_raises(ValueError, firls, 11, [0.1, 0.2], [0, 0], weight=[-1]) + + def test_firls(self): + N = 11 # number of taps in the filter + a = 0.1 # width of the transition band + + # design a halfband symmetric low-pass filter + h = firls(11, [0, a, 0.5-a, 0.5], [1, 1, 0, 0], fs=1.0) + + # make sure the filter has correct # of taps + assert h.shape[0] == N + + # make sure it is symmetric + midx = (N-1) // 2 + assert_array_almost_equal(h[:midx], h[:-midx-1:-1]) + + # make sure the center tap is 0.5 + assert_almost_equal(h[midx], 0.5) + + # For halfband symmetric, odd coefficients (except the center) + # should be zero (really small) + hodd = np.hstack((h[1:midx:2], h[-midx+1::2])) + assert_array_almost_equal(hodd, np.zeros_like(hodd)) + + # now check the frequency response + w, H = freqz(h, 1) + f = w/2/np.pi + Hmag = np.abs(H) + + # check that the pass band is close to unity + idx = np.logical_and(f > 0, f < a) + assert_array_almost_equal(Hmag[idx], np.ones_like(Hmag[idx]), decimal=3) + + # check that the stop band is close to zero + idx = np.logical_and(f > 0.5-a, f < 0.5) + assert_array_almost_equal(Hmag[idx], np.zeros_like(Hmag[idx]), decimal=3) + + def test_compare(self): + # compare to OCTAVE output + taps = firls(9, [0, 0.5, 0.55, 1], [1, 1, 0, 0], weight=[1, 2]) + # >> taps = firls(8, [0 0.5 0.55 1], [1 1 0 0], [1, 2]); + known_taps = [-6.26930101730182e-04, -1.03354450635036e-01, + -9.81576747564301e-03, 3.17271686090449e-01, + 5.11409425599933e-01, 3.17271686090449e-01, + -9.81576747564301e-03, -1.03354450635036e-01, + -6.26930101730182e-04] + xp_assert_close(taps, known_taps) + + # compare to MATLAB output + taps = firls(11, [0, 0.5, 0.5, 1], [1, 1, 0, 0], weight=[1, 2]) + # >> taps = firls(10, [0 0.5 0.5 1], [1 1 0 0], [1, 2]); + known_taps = [ + 0.058545300496815, -0.014233383714318, -0.104688258464392, + 0.012403323025279, 0.317930861136062, 0.488047220029700, + 0.317930861136062, 0.012403323025279, -0.104688258464392, + -0.014233383714318, 0.058545300496815] + xp_assert_close(taps, known_taps) + + # With linear changes: + taps = firls(7, (0, 1, 2, 3, 4, 5), [1, 0, 0, 1, 1, 0], fs=20) + # >> taps = firls(6, [0, 0.1, 0.2, 0.3, 0.4, 0.5], [1, 0, 0, 1, 1, 0]) + known_taps = [ + 1.156090832768218, -4.1385894727395849, 7.5288619164321826, + -8.5530572592947856, 7.5288619164321826, -4.1385894727395849, + 1.156090832768218] + xp_assert_close(taps, known_taps) + + def test_rank_deficient(self): + # solve() runs but warns (only sometimes, so here we don't use match) + x = firls(21, [0, 0.1, 0.9, 1], [1, 1, 0, 0]) + w, h = freqz(x, fs=2.) + absh2 = np.abs(h[:2]) + xp_assert_close(absh2, np.ones_like(absh2), atol=1e-5) + absh2 = np.abs(h[-2:]) + xp_assert_close(absh2, np.zeros_like(absh2), atol=1e-6, rtol=1e-7) + # switch to pinvh (tolerances could be higher with longer + # filters, but using shorter ones is faster computationally and + # the idea is the same) + x = firls(101, [0, 0.01, 0.99, 1], [1, 1, 0, 0]) + w, h = freqz(x, fs=2.) + mask = w < 0.01 + assert mask.sum() > 3 + habs = np.abs(h[mask]) + xp_assert_close(habs, np.ones_like(habs), atol=1e-4) + mask = w > 0.99 + assert mask.sum() > 3 + habs = np.abs(h[mask]) + xp_assert_close(habs, np.zeros_like(habs), atol=1e-4) + + def test_fs_validation(self): + with pytest.raises(ValueError, match="Sampling.*single scalar"): + firls(11, .1, 1, fs=np.array([10, 20])) + +class TestMinimumPhase: + @pytest.mark.thread_unsafe + def test_bad_args(self): + # not enough taps + assert_raises(ValueError, minimum_phase, [1.]) + assert_raises(ValueError, minimum_phase, [1., 1.]) + assert_raises(ValueError, minimum_phase, np.full(10, 1j)) + assert_raises(ValueError, minimum_phase, 'foo') + assert_raises(ValueError, minimum_phase, np.ones(10), n_fft=8) + assert_raises(ValueError, minimum_phase, np.ones(10), method='foo') + assert_warns(RuntimeWarning, minimum_phase, np.arange(3)) + with pytest.raises(ValueError, match="is only supported when"): + minimum_phase(np.ones(3), method='hilbert', half=False) + + def test_homomorphic(self): + # check that it can recover frequency responses of arbitrary + # linear-phase filters + + # for some cases we can get the actual filter back + h = [1, -1] + h_new = minimum_phase(np.convolve(h, h[::-1])) + xp_assert_close(h_new, np.asarray(h, dtype=np.float64), rtol=0.05) + + # but in general we only guarantee we get the magnitude back + rng = np.random.RandomState(0) + for n in (2, 3, 10, 11, 15, 16, 17, 20, 21, 100, 101): + h = rng.randn(n) + h_linear = np.convolve(h, h[::-1]) + h_new = minimum_phase(h_linear) + xp_assert_close(np.abs(fft(h_new)), np.abs(fft(h)), rtol=1e-4) + h_new = minimum_phase(h_linear, half=False) + assert len(h_linear) == len(h_new) + xp_assert_close(np.abs(fft(h_new)), np.abs(fft(h_linear)), rtol=1e-4) + + def test_hilbert(self): + # compare to MATLAB output of reference implementation + + # f=[0 0.3 0.5 1]; + # a=[1 1 0 0]; + # h=remez(11,f,a); + h = remez(12, [0, 0.3, 0.5, 1], [1, 0], fs=2.) + k = [0.349585548646686, 0.373552164395447, 0.326082685363438, + 0.077152207480935, -0.129943946349364, -0.059355880509749] + m = minimum_phase(h, 'hilbert') + xp_assert_close(m, k, rtol=5e-3) + + # f=[0 0.8 0.9 1]; + # a=[0 0 1 1]; + # h=remez(20,f,a); + h = remez(21, [0, 0.8, 0.9, 1], [0, 1], fs=2.) + k = [0.232486803906329, -0.133551833687071, 0.151871456867244, + -0.157957283165866, 0.151739294892963, -0.129293146705090, + 0.100787844523204, -0.065832656741252, 0.035361328741024, + -0.014977068692269, -0.158416139047557] + m = minimum_phase(h, 'hilbert', n_fft=2**19) + xp_assert_close(m, k, rtol=2e-3) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_ltisys.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_ltisys.py new file mode 100644 index 0000000000000000000000000000000000000000..7672db6a8ef44c208be7769caa874ff6ad446610 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_ltisys.py @@ -0,0 +1,1225 @@ +import warnings + +import numpy as np +from numpy.testing import suppress_warnings +import pytest +from pytest import raises as assert_raises +from scipy._lib._array_api import( + assert_almost_equal, xp_assert_equal, xp_assert_close +) + +from scipy.signal import (ss2tf, tf2ss, lti, + dlti, bode, freqresp, lsim, impulse, step, + abcd_normalize, place_poles, + TransferFunction, StateSpace, ZerosPolesGain) +from scipy.signal._filter_design import BadCoefficients +import scipy.linalg as linalg + + +def _assert_poles_close(P1,P2, rtol=1e-8, atol=1e-8): + """ + Check each pole in P1 is close to a pole in P2 with a 1e-8 + relative tolerance or 1e-8 absolute tolerance (useful for zero poles). + These tolerances are very strict but the systems tested are known to + accept these poles so we should not be far from what is requested. + """ + P2 = P2.copy() + for p1 in P1: + found = False + for p2_idx in range(P2.shape[0]): + if np.allclose([np.real(p1), np.imag(p1)], + [np.real(P2[p2_idx]), np.imag(P2[p2_idx])], + rtol, atol): + found = True + np.delete(P2, p2_idx) + break + if not found: + raise ValueError("Can't find pole " + str(p1) + " in " + str(P2)) + + +class TestPlacePoles: + + def _check(self, A, B, P, **kwargs): + """ + Perform the most common tests on the poles computed by place_poles + and return the Bunch object for further specific tests + """ + fsf = place_poles(A, B, P, **kwargs) + expected, _ = np.linalg.eig(A - np.dot(B, fsf.gain_matrix)) + _assert_poles_close(expected, fsf.requested_poles) + _assert_poles_close(expected, fsf.computed_poles) + _assert_poles_close(P,fsf.requested_poles) + return fsf + + def test_real(self): + # Test real pole placement using KNV and YT0 algorithm and example 1 in + # section 4 of the reference publication (see place_poles docstring) + A = np.array([1.380, -0.2077, 6.715, -5.676, -0.5814, -4.290, 0, + 0.6750, 1.067, 4.273, -6.654, 5.893, 0.0480, 4.273, + 1.343, -2.104]).reshape(4, 4) + B = np.array([0, 5.679, 1.136, 1.136, 0, 0, -3.146,0]).reshape(4, 2) + P = np.array([-0.2, -0.5, -5.0566, -8.6659]) + + # Check that both KNV and YT compute correct K matrix + self._check(A, B, P, method='KNV0') + self._check(A, B, P, method='YT') + + # Try to reach the specific case in _YT_real where two singular + # values are almost equal. This is to improve code coverage but I + # have no way to be sure this code is really reached + + # on some architectures this can lead to a RuntimeWarning invalid + # value in divide (see gh-7590), so suppress it for now + with np.errstate(invalid='ignore'): + self._check(A, B, (2,2,3,3)) + + def test_complex(self): + # Test complex pole placement on a linearized car model, taken from L. + # Jaulin, Automatique pour la robotique, Cours et Exercices, iSTE + # editions p 184/185 + A = np.array([[0, 7, 0, 0], + [0, 0, 0, 7/3.], + [0, 0, 0, 0], + [0, 0, 0, 0]]) + B = np.array([[0, 0], + [0, 0], + [1, 0], + [0, 1]]) + # Test complex poles on YT + P = np.array([-3, -1, -2-1j, -2+1j]) + # on macOS arm64 this can lead to a RuntimeWarning invalid + # value in divide, so suppress it for now + with np.errstate(divide='ignore', invalid='ignore'): + self._check(A, B, P) + + # Try to reach the specific case in _YT_complex where two singular + # values are almost equal. This is to improve code coverage but I + # have no way to be sure this code is really reached + + P = [0-1e-6j,0+1e-6j,-10,10] + with np.errstate(divide='ignore', invalid='ignore'): + self._check(A, B, P, maxiter=1000) + + # Try to reach the specific case in _YT_complex where the rank two + # update yields two null vectors. This test was found via Monte Carlo. + + A = np.array( + [-2148,-2902, -2267, -598, -1722, -1829, -165, -283, -2546, + -167, -754, -2285, -543, -1700, -584, -2978, -925, -1300, + -1583, -984, -386, -2650, -764, -897, -517, -1598, 2, -1709, + -291, -338, -153, -1804, -1106, -1168, -867, -2297] + ).reshape(6,6) + + B = np.array( + [-108, -374, -524, -1285, -1232, -161, -1204, -672, -637, + -15, -483, -23, -931, -780, -1245, -1129, -1290, -1502, + -952, -1374, -62, -964, -930, -939, -792, -756, -1437, + -491, -1543, -686] + ).reshape(6,5) + P = [-25.-29.j, -25.+29.j, 31.-42.j, 31.+42.j, 33.-41.j, 33.+41.j] + self._check(A, B, P) + + # Use a lot of poles to go through all cases for update_order + # in _YT_loop + + big_A = np.ones((11,11))-np.eye(11) + big_B = np.ones((11,10))-np.diag([1]*10,1)[:,1:] + big_A[:6,:6] = A + big_B[:6,:5] = B + + P = [-10,-20,-30,40,50,60,70,-20-5j,-20+5j,5+3j,5-3j] + with np.errstate(divide='ignore', invalid='ignore'): + self._check(big_A, big_B, P) + + #check with only complex poles and only real poles + P = [-10,-20,-30,-40,-50,-60,-70,-80,-90,-100] + self._check(big_A[:-1,:-1], big_B[:-1,:-1], P) + P = [-10+10j,-20+20j,-30+30j,-40+40j,-50+50j, + -10-10j,-20-20j,-30-30j,-40-40j,-50-50j] + self._check(big_A[:-1,:-1], big_B[:-1,:-1], P) + + # need a 5x5 array to ensure YT handles properly when there + # is only one real pole and several complex + A = np.array([0,7,0,0,0,0,0,7/3.,0,0,0,0,0,0,0,0, + 0,0,0,5,0,0,0,0,9]).reshape(5,5) + B = np.array([0,0,0,0,1,0,0,1,2,3]).reshape(5,2) + P = np.array([-2, -3+1j, -3-1j, -1+1j, -1-1j]) + with np.errstate(divide='ignore', invalid='ignore'): + place_poles(A, B, P) + + # same test with an odd number of real poles > 1 + # this is another specific case of YT + P = np.array([-2, -3, -4, -1+1j, -1-1j]) + with np.errstate(divide='ignore', invalid='ignore'): + self._check(A, B, P) + + def test_tricky_B(self): + # check we handle as we should the 1 column B matrices and + # n column B matrices (with n such as shape(A)=(n, n)) + A = np.array([1.380, -0.2077, 6.715, -5.676, -0.5814, -4.290, 0, + 0.6750, 1.067, 4.273, -6.654, 5.893, 0.0480, 4.273, + 1.343, -2.104]).reshape(4, 4) + B = np.array([0, 5.679, 1.136, 1.136, 0, 0, -3.146, 0, 1, 2, 3, 4, + 5, 6, 7, 8]).reshape(4, 4) + + # KNV or YT are not called here, it's a specific case with only + # one unique solution + P = np.array([-0.2, -0.5, -5.0566, -8.6659]) + fsf = self._check(A, B, P) + # rtol and nb_iter should be set to np.nan as the identity can be + # used as transfer matrix + assert np.isnan(fsf.rtol) + assert np.isnan(fsf.nb_iter) + + # check with complex poles too as they trigger a specific case in + # the specific case :-) + P = np.array((-2+1j,-2-1j,-3,-2)) + fsf = self._check(A, B, P) + assert np.isnan(fsf.rtol) + assert np.isnan(fsf.nb_iter) + + #now test with a B matrix with only one column (no optimisation) + B = B[:,0].reshape(4,1) + P = np.array((-2+1j,-2-1j,-3,-2)) + fsf = self._check(A, B, P) + + # we can't optimize anything, check they are set to 0 as expected + assert fsf.rtol == 0 + assert fsf.nb_iter == 0 + + @pytest.mark.thread_unsafe + def test_errors(self): + # Test input mistakes from user + A = np.array([0,7,0,0,0,0,0,7/3.,0,0,0,0,0,0,0,0]).reshape(4,4) + B = np.array([0,0,0,0,1,0,0,1]).reshape(4,2) + + #should fail as the method keyword is invalid + assert_raises(ValueError, place_poles, A, B, (-2.1,-2.2,-2.3,-2.4), + method="foo") + + #should fail as poles are not 1D array + assert_raises(ValueError, place_poles, A, B, + np.array((-2.1,-2.2,-2.3,-2.4)).reshape(4,1)) + + #should fail as A is not a 2D array + assert_raises(ValueError, place_poles, A[:,:,np.newaxis], B, + (-2.1,-2.2,-2.3,-2.4)) + + #should fail as B is not a 2D array + assert_raises(ValueError, place_poles, A, B[:,:,np.newaxis], + (-2.1,-2.2,-2.3,-2.4)) + + #should fail as there are too many poles + assert_raises(ValueError, place_poles, A, B, (-2.1,-2.2,-2.3,-2.4,-3)) + + #should fail as there are not enough poles + assert_raises(ValueError, place_poles, A, B, (-2.1,-2.2,-2.3)) + + #should fail as the rtol is greater than 1 + assert_raises(ValueError, place_poles, A, B, (-2.1,-2.2,-2.3,-2.4), + rtol=42) + + #should fail as maxiter is smaller than 1 + assert_raises(ValueError, place_poles, A, B, (-2.1,-2.2,-2.3,-2.4), + maxiter=-42) + + # should fail as ndim(B) is two + assert_raises(ValueError, place_poles, A, B, (-2,-2,-2,-2)) + + # uncontrollable system + assert_raises(ValueError, place_poles, np.ones((4,4)), + np.ones((4,2)), (1,2,3,4)) + + # Should not raise ValueError as the poles can be placed but should + # raise a warning as the convergence is not reached + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + fsf = place_poles(A, B, (-1,-2,-3,-4), rtol=1e-16, maxiter=42) + assert len(w) == 1 + assert issubclass(w[-1].category, UserWarning) + assert ("Convergence was not reached after maxiter iterations" + in str(w[-1].message)) + assert fsf.nb_iter == 42 + + # should fail as a complex misses its conjugate + assert_raises(ValueError, place_poles, A, B, (-2+1j,-2-1j,-2+3j,-2)) + + # should fail as A is not square + assert_raises(ValueError, place_poles, A[:,:3], B, (-2,-3,-4,-5)) + + # should fail as B has not the same number of lines as A + assert_raises(ValueError, place_poles, A, B[:3,:], (-2,-3,-4,-5)) + + # should fail as KNV0 does not support complex poles + assert_raises(ValueError, place_poles, A, B, + (-2+1j,-2-1j,-2+3j,-2-3j), method="KNV0") + + +class TestSS2TF: + + def check_matrix_shapes(self, p, q, r): + ss2tf(np.zeros((p, p)), + np.zeros((p, q)), + np.zeros((r, p)), + np.zeros((r, q)), 0) + + def test_shapes(self): + # Each tuple holds: + # number of states, number of inputs, number of outputs + for p, q, r in [(3, 3, 3), (1, 3, 3), (1, 1, 1)]: + self.check_matrix_shapes(p, q, r) + + def test_basic(self): + # Test a round trip through tf2ss and ss2tf. + b = np.array([1.0, 3.0, 5.0]) + a = np.array([1.0, 2.0, 3.0]) + + A, B, C, D = tf2ss(b, a) + xp_assert_close(A, [[-2., -3], [1, 0]], rtol=1e-13) + xp_assert_close(B, [[1.], [0]], rtol=1e-13) + xp_assert_close(C, [[1., 2]], rtol=1e-13) + xp_assert_close(D, [[1.]], rtol=1e-14) + + bb, aa = ss2tf(A, B, C, D) + xp_assert_close(bb[0], b, rtol=1e-13) + xp_assert_close(aa, a, rtol=1e-13) + + def test_zero_order_round_trip(self): + # See gh-5760 + tf = (2, 1) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[0.]], rtol=1e-13) + xp_assert_close(B, [[0.]], rtol=1e-13) + xp_assert_close(C, [[0.]], rtol=1e-13) + xp_assert_close(D, [[2.]], rtol=1e-13) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[2., 0]], rtol=1e-13) + xp_assert_close(den, [1., 0], rtol=1e-13) + + tf = ([[5], [2]], 1) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[0.]], rtol=1e-13) + xp_assert_close(B, [[0.]], rtol=1e-13) + xp_assert_close(C, [[0.], [0]], rtol=1e-13) + xp_assert_close(D, [[5.], [2]], rtol=1e-13) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[5., 0], [2, 0]], rtol=1e-13) + xp_assert_close(den, [1., 0], rtol=1e-13) + + def test_simo_round_trip(self): + # See gh-5753 + tf = ([[1, 2], [1, 1]], [1, 2]) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[-2.]], rtol=1e-13) + xp_assert_close(B, [[1.]], rtol=1e-13) + xp_assert_close(C, [[0.], [-1.]], rtol=1e-13) + xp_assert_close(D, [[1.], [1.]], rtol=1e-13) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[1., 2.], [1., 1.]], rtol=1e-13) + xp_assert_close(den, [1., 2.], rtol=1e-13) + + tf = ([[1, 0, 1], [1, 1, 1]], [1, 1, 1]) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[-1., -1.], [1., 0.]], rtol=1e-13) + xp_assert_close(B, [[1.], [0.]], rtol=1e-13) + xp_assert_close(C, [[-1., 0.], [0., 0.]], rtol=1e-13) + xp_assert_close(D, [[1.], [1.]], rtol=1e-13) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[1., 0., 1.], [1., 1., 1.]], rtol=1e-13) + xp_assert_close(den, [1., 1., 1.], rtol=1e-13) + + tf = ([[1, 2, 3], [1, 2, 3]], [1, 2, 3, 4]) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[-2., -3, -4], [1, 0, 0], [0, 1, 0]], rtol=1e-13) + xp_assert_close(B, [[1.], [0], [0]], rtol=1e-13) + xp_assert_close(C, [[1., 2, 3], [1, 2, 3]], rtol=1e-13) + xp_assert_close(D, [[0.], [0]], rtol=1e-13) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[0., 1, 2, 3], [0, 1, 2, 3]], rtol=1e-13) + xp_assert_close(den, [1., 2, 3, 4], rtol=1e-13) + + tf = (np.array([1, [2, 3]], dtype=object), [1, 6]) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[-6.]], rtol=1e-31) + xp_assert_close(B, [[1.]], rtol=1e-31) + xp_assert_close(C, [[1.], [-9]], rtol=1e-31) + xp_assert_close(D, [[0.], [2]], rtol=1e-31) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[0., 1], [2, 3]], rtol=1e-13) + xp_assert_close(den, [1., 6], rtol=1e-13) + + tf = (np.array([[1, -3], [1, 2, 3]], dtype=object), [1, 6, 5]) + A, B, C, D = tf2ss(*tf) + xp_assert_close(A, [[-6., -5], [1, 0]], rtol=1e-13) + xp_assert_close(B, [[1.], [0]], rtol=1e-13) + xp_assert_close(C, [[1., -3], [-4, -2]], rtol=1e-13) + xp_assert_close(D, [[0.], [1]], rtol=1e-13) + + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[0., 1, -3], [1, 2, 3]], rtol=1e-13) + xp_assert_close(den, [1., 6, 5], rtol=1e-13) + + def test_all_int_arrays(self): + A = [[0, 1, 0], [0, 0, 1], [-3, -4, -2]] + B = [[0], [0], [1]] + C = [[5, 1, 0]] + D = [[0]] + num, den = ss2tf(A, B, C, D) + xp_assert_close(num, [[0.0, 0.0, 1.0, 5.0]], rtol=1e-13, atol=1e-14) + xp_assert_close(den, [1.0, 2.0, 4.0, 3.0], rtol=1e-13) + + def test_multioutput(self): + # Regression test for gh-2669. + + # 4 states + A = np.array([[-1.0, 0.0, 1.0, 0.0], + [-1.0, 0.0, 2.0, 0.0], + [-4.0, 0.0, 3.0, 0.0], + [-8.0, 8.0, 0.0, 4.0]]) + + # 1 input + B = np.array([[0.3], + [0.0], + [7.0], + [0.0]]) + + # 3 outputs + C = np.array([[0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + [8.0, 8.0, 0.0, 0.0]]) + + D = np.array([[0.0], + [0.0], + [1.0]]) + + # Get the transfer functions for all the outputs in one call. + b_all, a = ss2tf(A, B, C, D) + + # Get the transfer functions for each output separately. + b0, a0 = ss2tf(A, B, C[0], D[0]) + b1, a1 = ss2tf(A, B, C[1], D[1]) + b2, a2 = ss2tf(A, B, C[2], D[2]) + + # Check that we got the same results. + xp_assert_close(a0, a, rtol=1e-13) + xp_assert_close(a1, a, rtol=1e-13) + xp_assert_close(a2, a, rtol=1e-13) + xp_assert_close(b_all, np.vstack((b0, b1, b2)), rtol=1e-13, atol=1e-14) + + +class TestLsim: + digits_accuracy = 7 + + def lti_nowarn(self, *args): + with suppress_warnings() as sup: + sup.filter(BadCoefficients) + system = lti(*args) + return system + + def test_first_order(self): + # y' = -y + # exact solution is y(t) = exp(-t) + system = self.lti_nowarn(-1.,1.,1.,0.) + t = np.linspace(0,5) + u = np.zeros_like(t) + tout, y, x = lsim(system, u, t, X0=[1.0]) + expected_x = np.exp(-tout) + assert_almost_equal(x, expected_x) + assert_almost_equal(y, expected_x) + + def test_second_order(self): + t = np.linspace(0, 10, 1001) + u = np.zeros_like(t) + # Second order system with a repeated root: x''(t) + 2*x(t) + x(t) = 0. + # With initial conditions x(0)=1.0 and x'(t)=0.0, the exact solution + # is (1-t)*exp(-t). + system = self.lti_nowarn([1.0], [1.0, 2.0, 1.0]) + tout, y, x = lsim(system, u, t, X0=[1.0, 0.0]) + expected_x = (1.0 - tout) * np.exp(-tout) + assert_almost_equal(x[:, 0], expected_x) + + def test_integrator(self): + # integrator: y' = u + system = self.lti_nowarn(0., 1., 1., 0.) + t = np.linspace(0,5) + u = t + tout, y, x = lsim(system, u, t) + expected_x = 0.5 * tout**2 + assert_almost_equal(x, expected_x, decimal=self.digits_accuracy) + assert_almost_equal(y, expected_x, decimal=self.digits_accuracy) + + def test_two_states(self): + # A system with two state variables, two inputs, and one output. + A = np.array([[-1.0, 0.0], [0.0, -2.0]]) + B = np.array([[1.0, 0.0], [0.0, 1.0]]) + C = np.array([1.0, 0.0]) + D = np.zeros((1, 2)) + + system = self.lti_nowarn(A, B, C, D) + + t = np.linspace(0, 10.0, 21) + u = np.zeros((len(t), 2)) + tout, y, x = lsim(system, U=u, T=t, X0=[1.0, 1.0]) + expected_y = np.exp(-tout) + expected_x0 = np.exp(-tout) + expected_x1 = np.exp(-2.0 * tout) + assert_almost_equal(y, expected_y) + assert_almost_equal(x[:, 0], expected_x0) + assert_almost_equal(x[:, 1], expected_x1) + + def test_double_integrator(self): + # double integrator: y'' = 2u + A = np.array([[0., 1.], [0., 0.]]) + B = np.array([[0.], [1.]]) + C = np.array([[2., 0.]]) + system = self.lti_nowarn(A, B, C, 0.) + t = np.linspace(0,5) + u = np.ones_like(t) + tout, y, x = lsim(system, u, t) + expected_x = np.transpose(np.array([0.5 * tout**2, tout])) + expected_y = tout**2 + assert_almost_equal(x, expected_x, decimal=self.digits_accuracy) + assert_almost_equal(y, expected_y, decimal=self.digits_accuracy) + + def test_jordan_block(self): + # Non-diagonalizable A matrix + # x1' + x1 = x2 + # x2' + x2 = u + # y = x1 + # Exact solution with u = 0 is y(t) = t exp(-t) + A = np.array([[-1., 1.], [0., -1.]]) + B = np.array([[0.], [1.]]) + C = np.array([[1., 0.]]) + system = self.lti_nowarn(A, B, C, 0.) + t = np.linspace(0,5) + u = np.zeros_like(t) + tout, y, x = lsim(system, u, t, X0=[0.0, 1.0]) + expected_y = tout * np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_miso(self): + # A system with two state variables, two inputs, and one output. + A = np.array([[-1.0, 0.0], [0.0, -2.0]]) + B = np.array([[1.0, 0.0], [0.0, 1.0]]) + C = np.array([1.0, 0.0]) + D = np.zeros((1,2)) + system = self.lti_nowarn(A, B, C, D) + + t = np.linspace(0, 5.0, 101) + u = np.zeros((len(t), 2)) + tout, y, x = lsim(system, u, t, X0=[1.0, 1.0]) + expected_y = np.exp(-tout) + expected_x0 = np.exp(-tout) + expected_x1 = np.exp(-2.0*tout) + assert_almost_equal(y, expected_y) + assert_almost_equal(x[:,0], expected_x0) + assert_almost_equal(x[:,1], expected_x1) + + def test_nonzero_initial_time(self): + system = self.lti_nowarn(-1.,1.,1.,0.) + t = np.linspace(1,2) + u = np.zeros_like(t) + tout, y, x = lsim(system, u, t, X0=[1.0]) + expected_y = np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_nonequal_timesteps(self): + t = np.array([0.0, 1.0, 1.0, 3.0]) + u = np.array([0.0, 0.0, 1.0, 1.0]) + # Simple integrator: x'(t) = u(t) + system = ([1.0], [1.0, 0.0]) + with assert_raises(ValueError, + match="Time steps are not equally spaced."): + tout, y, x = lsim(system, u, t, X0=[1.0]) + + +class TestImpulse: + def test_first_order(self): + # First order system: x'(t) + x(t) = u(t) + # Exact impulse response is x(t) = exp(-t). + system = ([1.0], [1.0,1.0]) + tout, y = impulse(system) + expected_y = np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_first_order_fixed_time(self): + # Specify the desired time values for the output. + + # First order system: x'(t) + x(t) = u(t) + # Exact impulse response is x(t) = exp(-t). + system = ([1.0], [1.0,1.0]) + n = 21 + t = np.linspace(0, 2.0, n) + tout, y = impulse(system, T=t) + assert tout.shape == (n,) + assert_almost_equal(tout, t) + expected_y = np.exp(-t) + assert_almost_equal(y, expected_y) + + def test_first_order_initial(self): + # Specify an initial condition as a scalar. + + # First order system: x'(t) + x(t) = u(t), x(0)=3.0 + # Exact impulse response is x(t) = 4*exp(-t). + system = ([1.0], [1.0,1.0]) + tout, y = impulse(system, X0=3.0) + expected_y = 4.0 * np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_first_order_initial_list(self): + # Specify an initial condition as a list. + + # First order system: x'(t) + x(t) = u(t), x(0)=3.0 + # Exact impulse response is x(t) = 4*exp(-t). + system = ([1.0], [1.0,1.0]) + tout, y = impulse(system, X0=[3.0]) + expected_y = 4.0 * np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_integrator(self): + # Simple integrator: x'(t) = u(t) + system = ([1.0], [1.0,0.0]) + tout, y = impulse(system) + expected_y = np.ones_like(tout) + assert_almost_equal(y, expected_y) + + def test_second_order(self): + # Second order system with a repeated root: + # x''(t) + 2*x(t) + x(t) = u(t) + # The exact impulse response is t*exp(-t). + system = ([1.0], [1.0, 2.0, 1.0]) + tout, y = impulse(system) + expected_y = tout * np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_array_like(self): + # Test that function can accept sequences, scalars. + system = ([1.0], [1.0, 2.0, 1.0]) + # TODO: add meaningful test where X0 is a list + tout, y = impulse(system, X0=[3], T=[5, 6]) + tout, y = impulse(system, X0=[3], T=[5]) + + def test_array_like2(self): + system = ([1.0], [1.0, 2.0, 1.0]) + tout, y = impulse(system, X0=3, T=5) + + +class TestStep: + def test_first_order(self): + # First order system: x'(t) + x(t) = u(t) + # Exact step response is x(t) = 1 - exp(-t). + system = ([1.0], [1.0,1.0]) + tout, y = step(system) + expected_y = 1.0 - np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_first_order_fixed_time(self): + # Specify the desired time values for the output. + + # First order system: x'(t) + x(t) = u(t) + # Exact step response is x(t) = 1 - exp(-t). + system = ([1.0], [1.0,1.0]) + n = 21 + t = np.linspace(0, 2.0, n) + tout, y = step(system, T=t) + assert tout.shape == (n,) + assert_almost_equal(tout, t) + expected_y = 1 - np.exp(-t) + assert_almost_equal(y, expected_y) + + def test_first_order_initial(self): + # Specify an initial condition as a scalar. + + # First order system: x'(t) + x(t) = u(t), x(0)=3.0 + # Exact step response is x(t) = 1 + 2*exp(-t). + system = ([1.0], [1.0,1.0]) + tout, y = step(system, X0=3.0) + expected_y = 1 + 2.0*np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_first_order_initial_list(self): + # Specify an initial condition as a list. + + # First order system: x'(t) + x(t) = u(t), x(0)=3.0 + # Exact step response is x(t) = 1 + 2*exp(-t). + system = ([1.0], [1.0,1.0]) + tout, y = step(system, X0=[3.0]) + expected_y = 1 + 2.0*np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_integrator(self): + # Simple integrator: x'(t) = u(t) + # Exact step response is x(t) = t. + system = ([1.0],[1.0,0.0]) + tout, y = step(system) + expected_y = tout + assert_almost_equal(y, expected_y) + + def test_second_order(self): + # Second order system with a repeated root: + # x''(t) + 2*x(t) + x(t) = u(t) + # The exact step response is 1 - (1 + t)*exp(-t). + system = ([1.0], [1.0, 2.0, 1.0]) + tout, y = step(system) + expected_y = 1 - (1 + tout) * np.exp(-tout) + assert_almost_equal(y, expected_y) + + def test_array_like(self): + # Test that function can accept sequences, scalars. + system = ([1.0], [1.0, 2.0, 1.0]) + # TODO: add meaningful test where X0 is a list + tout, y = step(system, T=[5, 6]) + + def test_complex_input(self): + # Test that complex input doesn't raise an error. + # `step` doesn't seem to have been designed for complex input, but this + # works and may be used, so add regression test. See gh-2654. + step(([], [-1], 1+0j)) + + +class TestLti: + def test_lti_instantiation(self): + # Test that lti can be instantiated with sequences, scalars. + # See PR-225. + + # TransferFunction + s = lti([1], [-1]) + assert isinstance(s, TransferFunction) + assert isinstance(s, lti) + assert not isinstance(s, dlti) + assert s.dt is None + + # ZerosPolesGain + s = lti(np.array([]), np.array([-1]), 1) + assert isinstance(s, ZerosPolesGain) + assert isinstance(s, lti) + assert not isinstance(s, dlti) + assert s.dt is None + + # StateSpace + s = lti([], [-1], 1) + s = lti([1], [-1], 1, 3) + assert isinstance(s, StateSpace) + assert isinstance(s, lti) + assert not isinstance(s, dlti) + assert s.dt is None + + +class TestStateSpace: + def test_initialization(self): + # Check that all initializations work + StateSpace(1, 1, 1, 1) + StateSpace([1], [2], [3], [4]) + StateSpace(np.array([[1, 2], [3, 4]]), np.array([[1], [2]]), + np.array([[1, 0]]), np.array([[0]])) + + def test_conversion(self): + # Check the conversion functions + s = StateSpace(1, 2, 3, 4) + assert isinstance(s.to_ss(), StateSpace) + assert isinstance(s.to_tf(), TransferFunction) + assert isinstance(s.to_zpk(), ZerosPolesGain) + + # Make sure copies work + assert StateSpace(s) is not s + assert s.to_ss() is not s + + def test_properties(self): + # Test setters/getters for cross class properties. + # This implicitly tests to_tf() and to_zpk() + + # Getters + s = StateSpace(1, 1, 1, 1) + xp_assert_equal(s.poles, [1.]) + xp_assert_equal(s.zeros, [0.]) + assert s.dt is None + + def test_operators(self): + # Test +/-/* operators on systems + + class BadType: + pass + + s1 = StateSpace(np.array([[-0.5, 0.7], [0.3, -0.8]]), + np.array([[1], [0]]), + np.array([[1, 0]]), + np.array([[0]]), + ) + + s2 = StateSpace(np.array([[-0.2, -0.1], [0.4, -0.1]]), + np.array([[1], [0]]), + np.array([[1, 0]]), + np.array([[0]]) + ) + + s_discrete = s1.to_discrete(0.1) + s2_discrete = s2.to_discrete(0.2) + s3_discrete = s2.to_discrete(0.1) + + # Impulse response + t = np.linspace(0, 1, 100) + u = np.zeros_like(t) + u[0] = 1 + + # Test multiplication + for typ in (int, float, complex, np.float32, np.complex128, np.array): + xp_assert_close(lsim(typ(2) * s1, U=u, T=t)[1], + typ(2) * lsim(s1, U=u, T=t)[1]) + + xp_assert_close(lsim(s1 * typ(2), U=u, T=t)[1], + lsim(s1, U=u, T=t)[1] * typ(2)) + + xp_assert_close(lsim(s1 / typ(2), U=u, T=t)[1], + lsim(s1, U=u, T=t)[1] / typ(2)) + + with assert_raises(TypeError): + typ(2) / s1 + + xp_assert_close(lsim(s1 * 2, U=u, T=t)[1], + lsim(s1, U=2 * u, T=t)[1]) + + xp_assert_close(lsim(s1 * s2, U=u, T=t)[1], + lsim(s1, U=lsim(s2, U=u, T=t)[1], T=t)[1], + atol=1e-5) + + with assert_raises(TypeError): + s1 / s1 + + with assert_raises(TypeError): + s1 * s_discrete + + with assert_raises(TypeError): + # Check different discretization constants + s_discrete * s2_discrete + + with assert_raises(TypeError): + s1 * BadType() + + with assert_raises(TypeError): + BadType() * s1 + + with assert_raises(TypeError): + s1 / BadType() + + with assert_raises(TypeError): + BadType() / s1 + + # Test addition + xp_assert_close(lsim(s1 + 2, U=u, T=t)[1], + 2 * u + lsim(s1, U=u, T=t)[1]) + + # Check for dimension mismatch + with assert_raises(ValueError): + s1 + np.array([1, 2]) + + with assert_raises(ValueError): + np.array([1, 2]) + s1 + + with assert_raises(TypeError): + s1 + s_discrete + + with assert_raises(ValueError): + s1 / np.array([[1, 2], [3, 4]]) + + with assert_raises(TypeError): + # Check different discretization constants + s_discrete + s2_discrete + + with assert_raises(TypeError): + s1 + BadType() + + with assert_raises(TypeError): + BadType() + s1 + + xp_assert_close(lsim(s1 + s2, U=u, T=t)[1], + lsim(s1, U=u, T=t)[1] + lsim(s2, U=u, T=t)[1]) + + # Test subtraction + xp_assert_close(lsim(s1 - 2, U=u, T=t)[1], + -2 * u + lsim(s1, U=u, T=t)[1]) + + xp_assert_close(lsim(2 - s1, U=u, T=t)[1], + 2 * u + lsim(-s1, U=u, T=t)[1]) + + xp_assert_close(lsim(s1 - s2, U=u, T=t)[1], + lsim(s1, U=u, T=t)[1] - lsim(s2, U=u, T=t)[1]) + + with assert_raises(TypeError): + s1 - BadType() + + with assert_raises(TypeError): + BadType() - s1 + + s = s_discrete + s3_discrete + assert s.dt == 0.1 + + s = s_discrete * s3_discrete + assert s.dt == 0.1 + + s = 3 * s_discrete + assert s.dt == 0.1 + + s = -s_discrete + assert s.dt == 0.1 + +class TestTransferFunction: + def test_initialization(self): + # Check that all initializations work + TransferFunction(1, 1) + TransferFunction([1], [2]) + TransferFunction(np.array([1]), np.array([2])) + + def test_conversion(self): + # Check the conversion functions + s = TransferFunction([1, 0], [1, -1]) + assert isinstance(s.to_ss(), StateSpace) + assert isinstance(s.to_tf(), TransferFunction) + assert isinstance(s.to_zpk(), ZerosPolesGain) + + # Make sure copies work + assert TransferFunction(s) is not s + assert s.to_tf() is not s + + def test_properties(self): + # Test setters/getters for cross class properties. + # This implicitly tests to_ss() and to_zpk() + + # Getters + s = TransferFunction([1, 0], [1, -1]) + xp_assert_equal(s.poles, [1.]) + xp_assert_equal(s.zeros, [0.]) + + +class TestZerosPolesGain: + def test_initialization(self): + # Check that all initializations work + ZerosPolesGain(1, 1, 1) + ZerosPolesGain([1], [2], 1) + ZerosPolesGain(np.array([1]), np.array([2]), 1) + + def test_conversion(self): + #Check the conversion functions + s = ZerosPolesGain(1, 2, 3) + assert isinstance(s.to_ss(), StateSpace) + assert isinstance(s.to_tf(), TransferFunction) + assert isinstance(s.to_zpk(), ZerosPolesGain) + + # Make sure copies work + assert ZerosPolesGain(s) is not s + assert s.to_zpk() is not s + + +class Test_abcd_normalize: + def setup_method(self): + self.A = np.array([[1.0, 2.0], [3.0, 4.0]]) + self.B = np.array([[-1.0], [5.0]]) + self.C = np.array([[4.0, 5.0]]) + self.D = np.array([[2.5]]) + + def test_no_matrix_fails(self): + assert_raises(ValueError, abcd_normalize) + + def test_A_nosquare_fails(self): + assert_raises(ValueError, abcd_normalize, [1, -1], + self.B, self.C, self.D) + + def test_AB_mismatch_fails(self): + assert_raises(ValueError, abcd_normalize, self.A, [-1, 5], + self.C, self.D) + + def test_AC_mismatch_fails(self): + assert_raises(ValueError, abcd_normalize, self.A, self.B, + [[4.0], [5.0]], self.D) + + def test_CD_mismatch_fails(self): + assert_raises(ValueError, abcd_normalize, self.A, self.B, + self.C, [2.5, 0]) + + def test_BD_mismatch_fails(self): + assert_raises(ValueError, abcd_normalize, self.A, [-1, 5], + self.C, self.D) + + def test_normalized_matrices_unchanged(self): + A, B, C, D = abcd_normalize(self.A, self.B, self.C, self.D) + xp_assert_equal(A, self.A) + xp_assert_equal(B, self.B) + xp_assert_equal(C, self.C) + xp_assert_equal(D, self.D) + + def test_shapes(self): + A, B, C, D = abcd_normalize(self.A, self.B, [1, 0], 0) + xp_assert_equal(A.shape[0], A.shape[1]) + xp_assert_equal(A.shape[0], B.shape[0]) + xp_assert_equal(A.shape[0], C.shape[1]) + xp_assert_equal(C.shape[0], D.shape[0]) + xp_assert_equal(B.shape[1], D.shape[1]) + + def test_zero_dimension_is_not_none1(self): + B_ = np.zeros((2, 0)) + D_ = np.zeros((0, 0)) + A, B, C, D = abcd_normalize(A=self.A, B=B_, D=D_) + xp_assert_equal(A, self.A) + xp_assert_equal(B, B_) + xp_assert_equal(D, D_) + assert C.shape[0] == D_.shape[0] + assert C.shape[1] == self.A.shape[0] + + def test_zero_dimension_is_not_none2(self): + B_ = np.zeros((2, 0)) + C_ = np.zeros((0, 2)) + A, B, C, D = abcd_normalize(A=self.A, B=B_, C=C_) + xp_assert_equal(A, self.A) + xp_assert_equal(B, B_) + xp_assert_equal(C, C_) + assert D.shape[0] == C_.shape[0] + assert D.shape[1] == B_.shape[1] + + def test_missing_A(self): + A, B, C, D = abcd_normalize(B=self.B, C=self.C, D=self.D) + assert A.shape[0] == A.shape[1] + assert A.shape[0] == B.shape[0] + assert A.shape == (self.B.shape[0], self.B.shape[0]) + + def test_missing_B(self): + A, B, C, D = abcd_normalize(A=self.A, C=self.C, D=self.D) + assert B.shape[0] == A.shape[0] + assert B.shape[1] == D.shape[1] + assert B.shape == (self.A.shape[0], self.D.shape[1]) + + def test_missing_C(self): + A, B, C, D = abcd_normalize(A=self.A, B=self.B, D=self.D) + assert C.shape[0] == D.shape[0] + assert C.shape[1] == A.shape[0] + assert C.shape == (self.D.shape[0], self.A.shape[0]) + + def test_missing_D(self): + A, B, C, D = abcd_normalize(A=self.A, B=self.B, C=self.C) + assert D.shape[0] == C.shape[0] + assert D.shape[1] == B.shape[1] + assert D.shape == (self.C.shape[0], self.B.shape[1]) + + def test_missing_AB(self): + A, B, C, D = abcd_normalize(C=self.C, D=self.D) + assert A.shape[0] == A.shape[1] + assert A.shape[0] == B.shape[0] + assert B.shape[1] == D.shape[1] + assert A.shape == (self.C.shape[1], self.C.shape[1]) + assert B.shape == (self.C.shape[1], self.D.shape[1]) + + def test_missing_AC(self): + A, B, C, D = abcd_normalize(B=self.B, D=self.D) + assert A.shape[0] == A.shape[1] + assert A.shape[0] == B.shape[0] + assert C.shape[0] == D.shape[0] + assert C.shape[1] == A.shape[0] + assert A.shape == (self.B.shape[0], self.B.shape[0]) + assert C.shape == (self.D.shape[0], self.B.shape[0]) + + def test_missing_AD(self): + A, B, C, D = abcd_normalize(B=self.B, C=self.C) + assert A.shape[0] == A.shape[1] + assert A.shape[0] == B.shape[0] + assert D.shape[0] == C.shape[0] + assert D.shape[1] == B.shape[1] + assert A.shape == (self.B.shape[0], self.B.shape[0]) + assert D.shape == (self.C.shape[0], self.B.shape[1]) + + def test_missing_BC(self): + A, B, C, D = abcd_normalize(A=self.A, D=self.D) + assert B.shape[0] == A.shape[0] + assert B.shape[1] == D.shape[1] + assert C.shape[0] == D.shape[0] + assert C.shape[1], A.shape[0] + assert B.shape == (self.A.shape[0], self.D.shape[1]) + assert C.shape == (self.D.shape[0], self.A.shape[0]) + + def test_missing_ABC_fails(self): + assert_raises(ValueError, abcd_normalize, D=self.D) + + def test_missing_BD_fails(self): + assert_raises(ValueError, abcd_normalize, A=self.A, C=self.C) + + def test_missing_CD_fails(self): + assert_raises(ValueError, abcd_normalize, A=self.A, B=self.B) + + +class Test_bode: + + def test_01(self): + # Test bode() magnitude calculation (manual sanity check). + # 1st order low-pass filter: H(s) = 1 / (s + 1), + # cutoff: 1 rad/s, slope: -20 dB/decade + # H(s=0.1) ~= 0 dB + # H(s=1) ~= -3 dB + # H(s=10) ~= -20 dB + # H(s=100) ~= -40 dB + system = lti([1], [1, 1]) + w = [0.1, 1, 10, 100] + w, mag, phase = bode(system, w=w) + expected_mag = [0, -3, -20, -40] + assert_almost_equal(mag, expected_mag, decimal=1) + + def test_02(self): + # Test bode() phase calculation (manual sanity check). + # 1st order low-pass filter: H(s) = 1 / (s + 1), + # angle(H(s=0.1)) ~= -5.7 deg + # angle(H(s=1)) ~= -45 deg + # angle(H(s=10)) ~= -84.3 deg + system = lti([1], [1, 1]) + w = [0.1, 1, 10] + w, mag, phase = bode(system, w=w) + expected_phase = [-5.7, -45, -84.3] + assert_almost_equal(phase, expected_phase, decimal=1) + + def test_03(self): + # Test bode() magnitude calculation. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + system = lti([1], [1, 1]) + w = [0.1, 1, 10, 100] + w, mag, phase = bode(system, w=w) + jw = w * 1j + y = np.polyval(system.num, jw) / np.polyval(system.den, jw) + expected_mag = 20.0 * np.log10(abs(y)) + assert_almost_equal(mag, expected_mag) + + def test_04(self): + # Test bode() phase calculation. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + system = lti([1], [1, 1]) + w = [0.1, 1, 10, 100] + w, mag, phase = bode(system, w=w) + jw = w * 1j + y = np.polyval(system.num, jw) / np.polyval(system.den, jw) + expected_phase = np.arctan2(y.imag, y.real) * 180.0 / np.pi + assert_almost_equal(phase, expected_phase) + + def test_05(self): + # Test that bode() finds a reasonable frequency range. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + system = lti([1], [1, 1]) + n = 10 + # Expected range is from 0.01 to 10. + expected_w = np.logspace(-2, 1, n) + w, mag, phase = bode(system, n=n) + assert_almost_equal(w, expected_w) + + def test_06(self): + # Test that bode() doesn't fail on a system with a pole at 0. + # integrator, pole at zero: H(s) = 1 / s + system = lti([1], [1, 0]) + w, mag, phase = bode(system, n=2) + assert w[0] == 0.01 # a fail would give not-a-number + + def test_07(self): + # bode() should not fail on a system with pure imaginary poles. + # The test passes if bode doesn't raise an exception. + system = lti([1], [1, 0, 100]) + w, mag, phase = bode(system, n=2) + + def test_08(self): + # Test that bode() return continuous phase, issues/2331. + system = lti([], [-10, -30, -40, -60, -70], 1) + w, mag, phase = system.bode(w=np.logspace(-3, 40, 100)) + assert_almost_equal(min(phase), -450, decimal=15) + + def test_from_state_space(self): + # Ensure that bode works with a system that was created from the + # state space representation matrices A, B, C, D. In this case, + # system.num will be a 2-D array with shape (1, n+1), where (n,n) + # is the shape of A. + # A Butterworth lowpass filter is used, so we know the exact + # frequency response. + a = np.array([1.0, 2.0, 2.0, 1.0]) + A = linalg.companion(a).T + B = np.array([[0.0], [0.0], [1.0]]) + C = np.array([[1.0, 0.0, 0.0]]) + D = np.array([[0.0]]) + with suppress_warnings() as sup: + sup.filter(BadCoefficients) + system = lti(A, B, C, D) + w, mag, phase = bode(system, n=100) + + expected_magnitude = 20 * np.log10(np.sqrt(1.0 / (1.0 + w**6))) + assert_almost_equal(mag, expected_magnitude) + + +class Test_freqresp: + + def test_output_manual(self): + # Test freqresp() output calculation (manual sanity check). + # 1st order low-pass filter: H(s) = 1 / (s + 1), + # re(H(s=0.1)) ~= 0.99 + # re(H(s=1)) ~= 0.5 + # re(H(s=10)) ~= 0.0099 + system = lti([1], [1, 1]) + w = [0.1, 1, 10] + w, H = freqresp(system, w=w) + expected_re = [0.99, 0.5, 0.0099] + expected_im = [-0.099, -0.5, -0.099] + assert_almost_equal(H.real, expected_re, decimal=1) + assert_almost_equal(H.imag, expected_im, decimal=1) + + def test_output(self): + # Test freqresp() output calculation. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + system = lti([1], [1, 1]) + w = [0.1, 1, 10, 100] + w, H = freqresp(system, w=w) + s = w * 1j + expected = np.polyval(system.num, s) / np.polyval(system.den, s) + assert_almost_equal(H.real, expected.real) + assert_almost_equal(H.imag, expected.imag) + + def test_freq_range(self): + # Test that freqresp() finds a reasonable frequency range. + # 1st order low-pass filter: H(s) = 1 / (s + 1) + # Expected range is from 0.01 to 10. + system = lti([1], [1, 1]) + n = 10 + expected_w = np.logspace(-2, 1, n) + w, H = freqresp(system, n=n) + assert_almost_equal(w, expected_w) + + def test_pole_zero(self): + # Test that freqresp() doesn't fail on a system with a pole at 0. + # integrator, pole at zero: H(s) = 1 / s + system = lti([1], [1, 0]) + w, H = freqresp(system, n=2) + assert w[0] == 0.01 # a fail would give not-a-number + + def test_from_state_space(self): + # Ensure that freqresp works with a system that was created from the + # state space representation matrices A, B, C, D. In this case, + # system.num will be a 2-D array with shape (1, n+1), where (n,n) is + # the shape of A. + # A Butterworth lowpass filter is used, so we know the exact + # frequency response. + a = np.array([1.0, 2.0, 2.0, 1.0]) + A = linalg.companion(a).T + B = np.array([[0.0],[0.0],[1.0]]) + C = np.array([[1.0, 0.0, 0.0]]) + D = np.array([[0.0]]) + with suppress_warnings() as sup: + sup.filter(BadCoefficients) + system = lti(A, B, C, D) + w, H = freqresp(system, n=100) + s = w * 1j + expected = (1.0 / (1.0 + 2*s + 2*s**2 + s**3)) + assert_almost_equal(H.real, expected.real) + assert_almost_equal(H.imag, expected.imag) + + def test_from_zpk(self): + # 4th order low-pass filter: H(s) = 1 / (s + 1) + system = lti([],[-1]*4,[1]) + w = [0.1, 1, 10, 100] + w, H = freqresp(system, w=w) + s = w * 1j + expected = 1 / (s + 1)**4 + assert_almost_equal(H.real, expected.real) + assert_almost_equal(H.imag, expected.imag) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_max_len_seq.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_max_len_seq.py new file mode 100644 index 0000000000000000000000000000000000000000..49522cbf242b77e6267fab93a42d7da765051194 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_max_len_seq.py @@ -0,0 +1,71 @@ +import numpy as np +from pytest import raises as assert_raises +from scipy._lib._array_api import xp_assert_close, xp_assert_equal + +from numpy.fft import fft, ifft + +from scipy.signal import max_len_seq + + +class TestMLS: + + def test_mls_inputs(self): + # can't all be zero state + assert_raises(ValueError, max_len_seq, + 10, state=np.zeros(10)) + # wrong size state + assert_raises(ValueError, max_len_seq, 10, + state=np.ones(3)) + # wrong length + assert_raises(ValueError, max_len_seq, 10, length=-1) + xp_assert_equal(max_len_seq(10, length=0)[0], + np.asarray([], dtype=np.int8) + ) + # unknown taps + assert_raises(ValueError, max_len_seq, 64) + # bad taps + assert_raises(ValueError, max_len_seq, 10, taps=[-1, 1]) + + def test_mls_output(self): + # define some alternate working taps + alt_taps = {2: [1], 3: [2], 4: [3], 5: [4, 3, 2], 6: [5, 4, 1], 7: [4], + 8: [7, 5, 3]} + # assume the other bit levels work, too slow to test higher orders... + for nbits in range(2, 8): + for state in [None, np.round(np.random.rand(nbits))]: + for taps in [None, alt_taps[nbits]]: + if state is not None and np.all(state == 0): + state[0] = 1 # they can't all be zero + orig_m = max_len_seq(nbits, state=state, + taps=taps)[0] + m = 2. * orig_m - 1. # convert to +/- 1 representation + # First, make sure we got all 1's or -1 + err_msg = "mls had non binary terms" + xp_assert_equal(np.abs(m), np.ones_like(m), + err_msg=err_msg) + # Test via circular cross-correlation, which is just mult. + # in the frequency domain with one signal conjugated + tester = np.real(ifft(fft(m) * np.conj(fft(m)))) + out_len = 2**nbits - 1 + # impulse amplitude == test_len + err_msg = "mls impulse has incorrect value" + xp_assert_close(tester[0], + float(out_len), + err_msg=err_msg + ) + # steady-state is -1 + err_msg = "mls steady-state has incorrect value" + xp_assert_close(tester[1:], + np.full(out_len - 1, -1, dtype=tester.dtype), + err_msg=err_msg) + # let's do the split thing using a couple options + for n in (1, 2**(nbits - 1)): + m1, s1 = max_len_seq(nbits, state=state, taps=taps, + length=n) + m2, s2 = max_len_seq(nbits, state=s1, taps=taps, + length=1) + m3, s3 = max_len_seq(nbits, state=s2, taps=taps, + length=out_len - n - 1) + new_m = np.concatenate((m1, m2, m3)) + xp_assert_equal(orig_m, new_m) + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_peak_finding.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_peak_finding.py new file mode 100644 index 0000000000000000000000000000000000000000..3a3611cff0b1f5871c020d7b41d9022e740f1639 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_peak_finding.py @@ -0,0 +1,915 @@ +import copy + +import numpy as np +import pytest +from pytest import raises, warns +from scipy._lib._array_api import xp_assert_close, xp_assert_equal + +from scipy.signal._peak_finding import ( + argrelmax, + argrelmin, + peak_prominences, + peak_widths, + _unpack_condition_args, + find_peaks, + find_peaks_cwt, + _identify_ridge_lines +) +from scipy.signal.windows import gaussian +from scipy.signal._peak_finding_utils import _local_maxima_1d, PeakPropertyWarning + + +def _gen_gaussians(center_locs, sigmas, total_length): + xdata = np.arange(0, total_length).astype(float) + out_data = np.zeros(total_length, dtype=float) + for ind, sigma in enumerate(sigmas): + tmp = (xdata - center_locs[ind]) / sigma + out_data += np.exp(-(tmp**2)) + return out_data + + +def _gen_gaussians_even(sigmas, total_length): + num_peaks = len(sigmas) + delta = total_length / (num_peaks + 1) + center_locs = np.linspace(delta, total_length - delta, num=num_peaks).astype(int) + out_data = _gen_gaussians(center_locs, sigmas, total_length) + return out_data, center_locs + + +def _gen_ridge_line(start_locs, max_locs, length, distances, gaps): + """ + Generate coordinates for a ridge line. + + Will be a series of coordinates, starting a start_loc (length 2). + The maximum distance between any adjacent columns will be + `max_distance`, the max distance between adjacent rows + will be `map_gap'. + + `max_locs` should be the size of the intended matrix. The + ending coordinates are guaranteed to be less than `max_locs`, + although they may not approach `max_locs` at all. + """ + + def keep_bounds(num, max_val): + out = max(num, 0) + out = min(out, max_val) + return out + + gaps = copy.deepcopy(gaps) + distances = copy.deepcopy(distances) + + locs = np.zeros([length, 2], dtype=int) + locs[0, :] = start_locs + total_length = max_locs[0] - start_locs[0] - sum(gaps) + if total_length < length: + raise ValueError('Cannot generate ridge line according to constraints') + dist_int = length / len(distances) - 1 + gap_int = length / len(gaps) - 1 + for ind in range(1, length): + nextcol = locs[ind - 1, 1] + nextrow = locs[ind - 1, 0] + 1 + if (ind % dist_int == 0) and (len(distances) > 0): + nextcol += ((-1)**ind)*distances.pop() + if (ind % gap_int == 0) and (len(gaps) > 0): + nextrow += gaps.pop() + nextrow = keep_bounds(nextrow, max_locs[0]) + nextcol = keep_bounds(nextcol, max_locs[1]) + locs[ind, :] = [nextrow, nextcol] + + return [locs[:, 0], locs[:, 1]] + + +class TestLocalMaxima1d: + + def test_empty(self): + """Test with empty signal.""" + x = np.array([], dtype=np.float64) + for array in _local_maxima_1d(x): + xp_assert_equal(array, np.array([]), check_dtype=False) + assert array.base is None + + def test_linear(self): + """Test with linear signal.""" + x = np.linspace(0, 100) + for array in _local_maxima_1d(x): + xp_assert_equal(array, np.array([], dtype=np.intp)) + assert array.base is None + + def test_simple(self): + """Test with simple signal.""" + x = np.linspace(-10, 10, 50) + x[2::3] += 1 + expected = np.arange(2, 50, 3, dtype=np.intp) + for array in _local_maxima_1d(x): + # For plateaus of size 1, the edges are identical with the + # midpoints + xp_assert_equal(array, expected, check_dtype=False) + assert array.base is None + + def test_flat_maxima(self): + """Test if flat maxima are detected correctly.""" + x = np.array([-1.3, 0, 1, 0, 2, 2, 0, 3, 3, 3, 2.99, 4, 4, 4, 4, -10, + -5, -5, -5, -5, -5, -10]) + midpoints, left_edges, right_edges = _local_maxima_1d(x) + xp_assert_equal(midpoints, np.array([2, 4, 8, 12, 18]), check_dtype=False) + xp_assert_equal(left_edges, np.array([2, 4, 7, 11, 16]), check_dtype=False) + xp_assert_equal(right_edges, np.array([2, 5, 9, 14, 20]), check_dtype=False) + + @pytest.mark.parametrize('x', [ + np.array([1., 0, 2]), + np.array([3., 3, 0, 4, 4]), + np.array([5., 5, 5, 0, 6, 6, 6]), + ]) + def test_signal_edges(self, x): + """Test if behavior on signal edges is correct.""" + for array in _local_maxima_1d(x): + xp_assert_equal(array, np.array([], dtype=np.intp)) + assert array.base is None + + def test_exceptions(self): + """Test input validation and raised exceptions.""" + with raises(ValueError, match="wrong number of dimensions"): + _local_maxima_1d(np.ones((1, 1))) + with raises(ValueError, match="expected 'const float64_t'"): + _local_maxima_1d(np.ones(1, dtype=int)) + with raises(TypeError, match="list"): + _local_maxima_1d([1., 2.]) + with raises(TypeError, match="'x' must not be None"): + _local_maxima_1d(None) + + +class TestRidgeLines: + + def test_empty(self): + test_matr = np.zeros([20, 100]) + lines = _identify_ridge_lines(test_matr, np.full(20, 2), 1) + assert len(lines) == 0 + + def test_minimal(self): + test_matr = np.zeros([20, 100]) + test_matr[0, 10] = 1 + lines = _identify_ridge_lines(test_matr, np.full(20, 2), 1) + assert len(lines) == 1 + + test_matr = np.zeros([20, 100]) + test_matr[0:2, 10] = 1 + lines = _identify_ridge_lines(test_matr, np.full(20, 2), 1) + assert len(lines) == 1 + + def test_single_pass(self): + distances = [0, 1, 2, 5] + gaps = [0, 1, 2, 0, 1] + test_matr = np.zeros([20, 50]) + 1e-12 + length = 12 + line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps) + test_matr[line[0], line[1]] = 1 + max_distances = np.full(20, max(distances)) + identified_lines = _identify_ridge_lines(test_matr, + max_distances, + max(gaps) + 1) + assert len(identified_lines) == 1 + for iline_, line_ in zip(identified_lines[0], line): + xp_assert_equal(iline_, line_, check_dtype=False) + + def test_single_bigdist(self): + distances = [0, 1, 2, 5] + gaps = [0, 1, 2, 4] + test_matr = np.zeros([20, 50]) + length = 12 + line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps) + test_matr[line[0], line[1]] = 1 + max_dist = 3 + max_distances = np.full(20, max_dist) + #This should get 2 lines, since the distance is too large + identified_lines = _identify_ridge_lines(test_matr, + max_distances, + max(gaps) + 1) + assert len(identified_lines) == 2 + + for iline in identified_lines: + adists = np.diff(iline[1]) + np.testing.assert_array_less(np.abs(adists), max_dist) + + agaps = np.diff(iline[0]) + np.testing.assert_array_less(np.abs(agaps), max(gaps) + 0.1) + + def test_single_biggap(self): + distances = [0, 1, 2, 5] + max_gap = 3 + gaps = [0, 4, 2, 1] + test_matr = np.zeros([20, 50]) + length = 12 + line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps) + test_matr[line[0], line[1]] = 1 + max_dist = 6 + max_distances = np.full(20, max_dist) + #This should get 2 lines, since the gap is too large + identified_lines = _identify_ridge_lines(test_matr, max_distances, max_gap) + assert len(identified_lines) == 2 + + for iline in identified_lines: + adists = np.diff(iline[1]) + np.testing.assert_array_less(np.abs(adists), max_dist) + + agaps = np.diff(iline[0]) + np.testing.assert_array_less(np.abs(agaps), max(gaps) + 0.1) + + def test_single_biggaps(self): + distances = [0] + max_gap = 1 + gaps = [3, 6] + test_matr = np.zeros([50, 50]) + length = 30 + line = _gen_ridge_line([0, 25], test_matr.shape, length, distances, gaps) + test_matr[line[0], line[1]] = 1 + max_dist = 1 + max_distances = np.full(50, max_dist) + #This should get 3 lines, since the gaps are too large + identified_lines = _identify_ridge_lines(test_matr, max_distances, max_gap) + assert len(identified_lines) == 3 + + for iline in identified_lines: + adists = np.diff(iline[1]) + np.testing.assert_array_less(np.abs(adists), max_dist) + + agaps = np.diff(iline[0]) + np.testing.assert_array_less(np.abs(agaps), max(gaps) + 0.1) + + +class TestArgrel: + + def test_empty(self): + # Regression test for gh-2832. + # When there are no relative extrema, make sure that + # the number of empty arrays returned matches the + # dimension of the input. + + empty_array = np.array([], dtype=int) + + z1 = np.zeros(5) + + i = argrelmin(z1) + xp_assert_equal(len(i), 1) + xp_assert_equal(i[0], empty_array, check_dtype=False) + + z2 = np.zeros((3, 5)) + + row, col = argrelmin(z2, axis=0) + xp_assert_equal(row, empty_array, check_dtype=False) + xp_assert_equal(col, empty_array, check_dtype=False) + + row, col = argrelmin(z2, axis=1) + xp_assert_equal(row, empty_array, check_dtype=False) + xp_assert_equal(col, empty_array, check_dtype=False) + + def test_basic(self): + # Note: the docstrings for the argrel{min,max,extrema} functions + # do not give a guarantee of the order of the indices, so we'll + # sort them before testing. + + x = np.array([[1, 2, 2, 3, 2], + [2, 1, 2, 2, 3], + [3, 2, 1, 2, 2], + [2, 3, 2, 1, 2], + [1, 2, 3, 2, 1]]) + + row, col = argrelmax(x, axis=0) + order = np.argsort(row) + xp_assert_equal(row[order], [1, 2, 3], check_dtype=False) + xp_assert_equal(col[order], [4, 0, 1], check_dtype=False) + + row, col = argrelmax(x, axis=1) + order = np.argsort(row) + xp_assert_equal(row[order], [0, 3, 4], check_dtype=False) + xp_assert_equal(col[order], [3, 1, 2], check_dtype=False) + + row, col = argrelmin(x, axis=0) + order = np.argsort(row) + xp_assert_equal(row[order], [1, 2, 3], check_dtype=False) + xp_assert_equal(col[order], [1, 2, 3], check_dtype=False) + + row, col = argrelmin(x, axis=1) + order = np.argsort(row) + xp_assert_equal(row[order], [1, 2, 3], check_dtype=False) + xp_assert_equal(col[order], [1, 2, 3], check_dtype=False) + + def test_highorder(self): + order = 2 + sigmas = [1.0, 2.0, 10.0, 5.0, 15.0] + test_data, act_locs = _gen_gaussians_even(sigmas, 500) + test_data[act_locs + order] = test_data[act_locs]*0.99999 + test_data[act_locs - order] = test_data[act_locs]*0.99999 + rel_max_locs = argrelmax(test_data, order=order, mode='clip')[0] + + assert len(rel_max_locs) == len(act_locs) + assert (rel_max_locs == act_locs).all() + + def test_2d_gaussians(self): + sigmas = [1.0, 2.0, 10.0] + test_data, act_locs = _gen_gaussians_even(sigmas, 100) + rot_factor = 20 + rot_range = np.arange(0, len(test_data)) - rot_factor + test_data_2 = np.vstack([test_data, test_data[rot_range]]) + rel_max_rows, rel_max_cols = argrelmax(test_data_2, axis=1, order=1) + + for rw in range(0, test_data_2.shape[0]): + inds = (rel_max_rows == rw) + + assert len(rel_max_cols[inds]) == len(act_locs) + assert (act_locs == (rel_max_cols[inds] - rot_factor*rw)).all() + + +class TestPeakProminences: + + def test_empty(self): + """ + Test if an empty array is returned if no peaks are provided. + """ + out = peak_prominences([1, 2, 3], []) + for arr, dtype in zip(out, [np.float64, np.intp, np.intp]): + assert arr.size == 0 + assert arr.dtype == dtype + + out = peak_prominences([], []) + for arr, dtype in zip(out, [np.float64, np.intp, np.intp]): + assert arr.size == 0 + assert arr.dtype == dtype + + def test_basic(self): + """ + Test if height of prominences is correctly calculated in signal with + rising baseline (peak widths are 1 sample). + """ + # Prepare basic signal + x = np.array([-1, 1.2, 1.2, 1, 3.2, 1.3, 2.88, 2.1]) + peaks = np.array([1, 2, 4, 6]) + lbases = np.array([0, 0, 0, 5]) + rbases = np.array([3, 3, 5, 7]) + proms = x[peaks] - np.max([x[lbases], x[rbases]], axis=0) + # Test if calculation matches handcrafted result + out = peak_prominences(x, peaks) + xp_assert_equal(out[0], proms, check_dtype=False) + xp_assert_equal(out[1], lbases, check_dtype=False) + xp_assert_equal(out[2], rbases, check_dtype=False) + + def test_edge_cases(self): + """ + Test edge cases. + """ + # Peaks have same height, prominence and bases + x = [0, 2, 1, 2, 1, 2, 0] + peaks = [1, 3, 5] + proms, lbases, rbases = peak_prominences(x, peaks) + xp_assert_equal(proms, np.asarray([2.0, 2, 2]), check_dtype=False) + xp_assert_equal(lbases, [0, 0, 0], check_dtype=False) + xp_assert_equal(rbases, [6, 6, 6], check_dtype=False) + + # Peaks have same height & prominence but different bases + x = [0, 1, 0, 1, 0, 1, 0] + peaks = np.array([1, 3, 5]) + proms, lbases, rbases = peak_prominences(x, peaks) + xp_assert_equal(proms, np.asarray([1.0, 1, 1])) + xp_assert_equal(lbases, peaks - 1, check_dtype=False) + xp_assert_equal(rbases, peaks + 1, check_dtype=False) + + def test_non_contiguous(self): + """ + Test with non-C-contiguous input arrays. + """ + x = np.repeat([-9, 9, 9, 0, 3, 1], 2) + peaks = np.repeat([1, 2, 4], 2) + proms, lbases, rbases = peak_prominences(x[::2], peaks[::2]) + xp_assert_equal(proms, np.asarray([9.0, 9, 2])) + xp_assert_equal(lbases, [0, 0, 3], check_dtype=False) + xp_assert_equal(rbases, [3, 3, 5], check_dtype=False) + + def test_wlen(self): + """ + Test if wlen actually shrinks the evaluation range correctly. + """ + x = [0, 1, 2, 3, 1, 0, -1] + peak = [3] + # Test rounding behavior of wlen + proms = peak_prominences(x, peak) + for prom, val in zip(proms, [3.0, 0, 6]): + assert prom == val + + for wlen, i in [(8, 0), (7, 0), (6, 0), (5, 1), (3.2, 1), (3, 2), (1.1, 2)]: + proms = peak_prominences(x, peak, wlen) + for prom, val in zip(proms, [3. - i, 0 + i, 6 - i]): + assert prom == val + + def test_exceptions(self): + """ + Verify that exceptions and warnings are raised. + """ + # x with dimension > 1 + with raises(ValueError, match='1-D array'): + peak_prominences([[0, 1, 1, 0]], [1, 2]) + # peaks with dimension > 1 + with raises(ValueError, match='1-D array'): + peak_prominences([0, 1, 1, 0], [[1, 2]]) + # x with dimension < 1 + with raises(ValueError, match='1-D array'): + peak_prominences(3, [0,]) + + # empty x with supplied + with raises(ValueError, match='not a valid index'): + peak_prominences([], [0]) + # invalid indices with non-empty x + for p in [-100, -1, 3, 1000]: + with raises(ValueError, match='not a valid index'): + peak_prominences([1, 0, 2], [p]) + + # peaks is not cast-able to np.intp + with raises(TypeError, match='cannot safely cast'): + peak_prominences([0, 1, 1, 0], [1.1, 2.3]) + + # wlen < 3 + with raises(ValueError, match='wlen'): + peak_prominences(np.arange(10), [3, 5], wlen=1) + + @pytest.mark.thread_unsafe + def test_warnings(self): + """ + Verify that appropriate warnings are raised. + """ + msg = "some peaks have a prominence of 0" + for p in [0, 1, 2]: + with warns(PeakPropertyWarning, match=msg): + peak_prominences([1, 0, 2], [p,]) + with warns(PeakPropertyWarning, match=msg): + peak_prominences([0, 1, 1, 1, 0], [2], wlen=2) + + +class TestPeakWidths: + + def test_empty(self): + """ + Test if an empty array is returned if no peaks are provided. + """ + widths = peak_widths([], [])[0] + assert isinstance(widths, np.ndarray) + assert widths.size == 0 + widths = peak_widths([1, 2, 3], [])[0] + assert isinstance(widths, np.ndarray) + assert widths.size == 0 + out = peak_widths([], []) + for arr in out: + assert isinstance(arr, np.ndarray) + assert arr.size == 0 + + @pytest.mark.filterwarnings("ignore:some peaks have a width of 0") + def test_basic(self): + """ + Test a simple use case with easy to verify results at different relative + heights. + """ + x = np.array([1, 0, 1, 2, 1, 0, -1]) + prominence = 2 + for rel_height, width_true, lip_true, rip_true in [ + (0., 0., 3., 3.), # raises warning + (0.25, 1., 2.5, 3.5), + (0.5, 2., 2., 4.), + (0.75, 3., 1.5, 4.5), + (1., 4., 1., 5.), + (2., 5., 1., 6.), + (3., 5., 1., 6.) + ]: + width_calc, height, lip_calc, rip_calc = peak_widths( + x, [3], rel_height) + xp_assert_close(width_calc, np.asarray([width_true])) + xp_assert_close(height, np.asarray([2 - rel_height * prominence])) + xp_assert_close(lip_calc, np.asarray([lip_true])) + xp_assert_close(rip_calc, np.asarray([rip_true])) + + def test_non_contiguous(self): + """ + Test with non-C-contiguous input arrays. + """ + x = np.repeat([0, 100, 50], 4) + peaks = np.repeat([1], 3) + result = peak_widths(x[::4], peaks[::3]) + xp_assert_equal(result, + np.asarray([[0.75], [75], [0.75], [1.5]]) + ) + + def test_exceptions(self): + """ + Verify that argument validation works as intended. + """ + with raises(ValueError, match='1-D array'): + # x with dimension > 1 + peak_widths(np.zeros((3, 4)), np.ones(3)) + with raises(ValueError, match='1-D array'): + # x with dimension < 1 + peak_widths(3, [0]) + with raises(ValueError, match='1-D array'): + # peaks with dimension > 1 + peak_widths(np.arange(10), np.ones((3, 2), dtype=np.intp)) + with raises(ValueError, match='1-D array'): + # peaks with dimension < 1 + peak_widths(np.arange(10), 3) + with raises(ValueError, match='not a valid index'): + # peak pos exceeds x.size + peak_widths(np.arange(10), [8, 11]) + with raises(ValueError, match='not a valid index'): + # empty x with peaks supplied + peak_widths([], [1, 2]) + with raises(TypeError, match='cannot safely cast'): + # peak cannot be safely cast to intp + peak_widths(np.arange(10), [1.1, 2.3]) + with raises(ValueError, match='rel_height'): + # rel_height is < 0 + peak_widths([0, 1, 0, 1, 0], [1, 3], rel_height=-1) + with raises(TypeError, match='None'): + # prominence data contains None + peak_widths([1, 2, 1], [1], prominence_data=(None, None, None)) + + @pytest.mark.thread_unsafe + def test_warnings(self): + """ + Verify that appropriate warnings are raised. + """ + msg = "some peaks have a width of 0" + with warns(PeakPropertyWarning, match=msg): + # Case: rel_height is 0 + peak_widths([0, 1, 0], [1], rel_height=0) + with warns(PeakPropertyWarning, match=msg): + # Case: prominence is 0 and bases are identical + peak_widths( + [0, 1, 1, 1, 0], [2], + prominence_data=(np.array([0.], np.float64), + np.array([2], np.intp), + np.array([2], np.intp)) + ) + + def test_mismatching_prominence_data(self): + """Test with mismatching peak and / or prominence data.""" + x = [0, 1, 0] + peak = [1] + for i, (prominences, left_bases, right_bases) in enumerate([ + ((1.,), (-1,), (2,)), # left base not in x + ((1.,), (0,), (3,)), # right base not in x + ((1.,), (2,), (0,)), # swapped bases same as peak + ((1., 1.), (0, 0), (2, 2)), # array shapes don't match peaks + ((1., 1.), (0,), (2,)), # arrays with different shapes + ((1.,), (0, 0), (2,)), # arrays with different shapes + ((1.,), (0,), (2, 2)) # arrays with different shapes + ]): + # Make sure input is matches output of signal.peak_prominences + prominence_data = (np.array(prominences, dtype=np.float64), + np.array(left_bases, dtype=np.intp), + np.array(right_bases, dtype=np.intp)) + # Test for correct exception + if i < 3: + match = "prominence data is invalid for peak" + else: + match = "arrays in `prominence_data` must have the same shape" + with raises(ValueError, match=match): + peak_widths(x, peak, prominence_data=prominence_data) + + @pytest.mark.filterwarnings("ignore:some peaks have a width of 0") + def test_intersection_rules(self): + """Test if x == eval_height counts as an intersection.""" + # Flatt peak with two possible intersection points if evaluated at 1 + x = [0, 1, 2, 1, 3, 3, 3, 1, 2, 1, 0] + # relative height is 0 -> width is 0 as well, raises warning + xp_assert_close(peak_widths(x, peaks=[5], rel_height=0), + [(0.,), (3.,), (5.,), (5.,)]) + # width_height == x counts as intersection -> nearest 1 is chosen + xp_assert_close(peak_widths(x, peaks=[5], rel_height=2/3), + [(4.,), (1.,), (3.,), (7.,)]) + + +def test_unpack_condition_args(): + """ + Verify parsing of condition arguments for `scipy.signal.find_peaks` function. + """ + x = np.arange(10) + amin_true = x + amax_true = amin_true + 10 + peaks = amin_true[1::2] + + # Test unpacking with None or interval + assert (None, None) == _unpack_condition_args((None, None), x, peaks) + assert (1, None) == _unpack_condition_args(1, x, peaks) + assert (1, None) == _unpack_condition_args((1, None), x, peaks) + assert (None, 2) == _unpack_condition_args((None, 2), x, peaks) + assert (3., 4.5) == _unpack_condition_args((3., 4.5), x, peaks) + + # Test if borders are correctly reduced with `peaks` + amin_calc, amax_calc = _unpack_condition_args((amin_true, amax_true), x, peaks) + xp_assert_equal(amin_calc, amin_true[peaks]) + xp_assert_equal(amax_calc, amax_true[peaks]) + + # Test raises if array borders don't match x + with raises(ValueError, match="array size of lower"): + _unpack_condition_args(amin_true, np.arange(11), peaks) + with raises(ValueError, match="array size of upper"): + _unpack_condition_args((None, amin_true), np.arange(11), peaks) + + +class TestFindPeaks: + + # Keys of optionally returned properties + property_keys = {'peak_heights', 'left_thresholds', 'right_thresholds', + 'prominences', 'left_bases', 'right_bases', 'widths', + 'width_heights', 'left_ips', 'right_ips'} + + def test_constant(self): + """ + Test behavior for signal without local maxima. + """ + open_interval = (None, None) + peaks, props = find_peaks(np.ones(10), + height=open_interval, threshold=open_interval, + prominence=open_interval, width=open_interval) + assert peaks.size == 0 + for key in self.property_keys: + assert props[key].size == 0 + + def test_plateau_size(self): + """ + Test plateau size condition for peaks. + """ + # Prepare signal with peaks with peak_height == plateau_size + plateau_sizes = np.array([1, 2, 3, 4, 8, 20, 111]) + x = np.zeros(plateau_sizes.size * 2 + 1) + x[1::2] = plateau_sizes + repeats = np.ones(x.size, dtype=int) + repeats[1::2] = x[1::2] + x = np.repeat(x, repeats) + + # Test full output + peaks, props = find_peaks(x, plateau_size=(None, None)) + xp_assert_equal(peaks, [1, 3, 7, 11, 18, 33, 100], check_dtype=False) + xp_assert_equal(props["plateau_sizes"], plateau_sizes, check_dtype=False) + xp_assert_equal(props["left_edges"], peaks - (plateau_sizes - 1) // 2, + check_dtype=False) + xp_assert_equal(props["right_edges"], peaks + plateau_sizes // 2, + check_dtype=False) + + # Test conditions + xp_assert_equal(find_peaks(x, plateau_size=4)[0], [11, 18, 33, 100], + check_dtype=False) + xp_assert_equal(find_peaks(x, plateau_size=(None, 3.5))[0], [1, 3, 7], + check_dtype=False) + xp_assert_equal(find_peaks(x, plateau_size=(5, 50))[0], [18, 33], + check_dtype=False) + + def test_height_condition(self): + """ + Test height condition for peaks. + """ + x = (0., 1/3, 0., 2.5, 0, 4., 0) + peaks, props = find_peaks(x, height=(None, None)) + xp_assert_equal(peaks, np.array([1, 3, 5]), check_dtype=False) + xp_assert_equal(props['peak_heights'], np.array([1/3, 2.5, 4.]), + check_dtype=False) + xp_assert_equal(find_peaks(x, height=0.5)[0], np.array([3, 5]), + check_dtype=False) + xp_assert_equal(find_peaks(x, height=(None, 3))[0], np.array([1, 3]), + check_dtype=False) + xp_assert_equal(find_peaks(x, height=(2, 3))[0], np.array([3]), + check_dtype=False) + + def test_threshold_condition(self): + """ + Test threshold condition for peaks. + """ + x = (0, 2, 1, 4, -1) + peaks, props = find_peaks(x, threshold=(None, None)) + xp_assert_equal(peaks, np.array([1, 3]), check_dtype=False) + xp_assert_equal(props['left_thresholds'], np.array([2.0, 3.0])) + xp_assert_equal(props['right_thresholds'], np.array([1.0, 5.0])) + xp_assert_equal(find_peaks(x, threshold=2)[0], np.array([3]), + check_dtype=False) + xp_assert_equal(find_peaks(x, threshold=3.5)[0], np.array([], dtype=int), + check_dtype=False) + xp_assert_equal(find_peaks(x, threshold=(None, 5))[0], np.array([1, 3]), + check_dtype=False) + xp_assert_equal(find_peaks(x, threshold=(None, 4))[0], np.array([1]), + check_dtype=False) + xp_assert_equal(find_peaks(x, threshold=(2, 4))[0], np.array([], dtype=int), + check_dtype=False) + + def test_distance_condition(self): + """ + Test distance condition for peaks. + """ + # Peaks of different height with constant distance 3 + peaks_all = np.arange(1, 21, 3) + x = np.zeros(21) + x[peaks_all] += np.linspace(1, 2, peaks_all.size) + + # Test if peaks with "minimal" distance are still selected (distance = 3) + xp_assert_equal(find_peaks(x, distance=3)[0], peaks_all, check_dtype=False) + + # Select every second peak (distance > 3) + peaks_subset = find_peaks(x, distance=3.0001)[0] + # Test if peaks_subset is subset of peaks_all + assert np.setdiff1d(peaks_subset, peaks_all, assume_unique=True).size == 0 + + # Test if every second peak was removed + dfs = np.diff(peaks_subset) + xp_assert_equal(dfs, 6*np.ones_like(dfs)) + + # Test priority of peak removal + x = [-2, 1, -1, 0, -3] + peaks_subset = find_peaks(x, distance=10)[0] # use distance > x size + assert peaks_subset.size == 1 and peaks_subset[0] == 1 + + def test_prominence_condition(self): + """ + Test prominence condition for peaks. + """ + x = np.linspace(0, 10, 100) + peaks_true = np.arange(1, 99, 2) + offset = np.linspace(1, 10, peaks_true.size) + x[peaks_true] += offset + prominences = x[peaks_true] - x[peaks_true + 1] + interval = (3, 9) + keep = np.nonzero( + (interval[0] <= prominences) & (prominences <= interval[1])) + + peaks_calc, properties = find_peaks(x, prominence=interval) + xp_assert_equal(peaks_calc, peaks_true[keep], check_dtype=False) + xp_assert_equal(properties['prominences'], prominences[keep], check_dtype=False) + xp_assert_equal(properties['left_bases'], + np.zeros_like(properties['left_bases'])) + xp_assert_equal(properties['right_bases'], peaks_true[keep] + 1, + check_dtype=False) + + def test_width_condition(self): + """ + Test width condition for peaks. + """ + x = np.array([1, 0, 1, 2, 1, 0, -1, 4, 0]) + peaks, props = find_peaks(x, width=(None, 2), rel_height=0.75) + assert peaks.size == 1 + xp_assert_equal(peaks, 7*np.ones_like(peaks)) + xp_assert_close(props['widths'], np.asarray([1.35])) + xp_assert_close(props['width_heights'], np.asarray([1.])) + xp_assert_close(props['left_ips'], np.asarray([6.4])) + xp_assert_close(props['right_ips'], np.asarray([7.75])) + + def test_properties(self): + """ + Test returned properties. + """ + open_interval = (None, None) + x = [0, 1, 0, 2, 1.5, 0, 3, 0, 5, 9] + peaks, props = find_peaks(x, + height=open_interval, threshold=open_interval, + prominence=open_interval, width=open_interval) + assert len(props) == len(self.property_keys) + for key in self.property_keys: + assert peaks.size == props[key].size + + def test_raises(self): + """ + Test exceptions raised by function. + """ + with raises(ValueError, match="1-D array"): + find_peaks(np.array(1)) + with raises(ValueError, match="1-D array"): + find_peaks(np.ones((2, 2))) + with raises(ValueError, match="distance"): + find_peaks(np.arange(10), distance=-1) + + @pytest.mark.filterwarnings("ignore:some peaks have a prominence of 0", + "ignore:some peaks have a width of 0") + def test_wlen_smaller_plateau(self): + """ + Test behavior of prominence and width calculation if the given window + length is smaller than a peak's plateau size. + + Regression test for gh-9110. + """ + peaks, props = find_peaks([0, 1, 1, 1, 0], prominence=(None, None), + width=(None, None), wlen=2) + xp_assert_equal(peaks, 2 * np.ones_like(peaks)) + xp_assert_equal(props["prominences"], np.zeros_like(props["prominences"])) + xp_assert_equal(props["widths"], np.zeros_like(props["widths"])) + xp_assert_equal(props["width_heights"], np.ones_like(props["width_heights"])) + for key in ("left_bases", "right_bases", "left_ips", "right_ips"): + xp_assert_equal(props[key], peaks, check_dtype=False) + + @pytest.mark.parametrize("kwargs", [ + {}, + {"distance": 3.0}, + {"prominence": (None, None)}, + {"width": (None, 2)}, + + ]) + def test_readonly_array(self, kwargs): + """ + Test readonly arrays are accepted. + """ + x = np.linspace(0, 10, 15) + x_readonly = x.copy() + x_readonly.flags.writeable = False + + peaks, _ = find_peaks(x) + peaks_readonly, _ = find_peaks(x_readonly, **kwargs) + + xp_assert_close(peaks, peaks_readonly) + + +class TestFindPeaksCwt: + + def test_find_peaks_exact(self): + """ + Generate a series of gaussians and attempt to find the peak locations. + """ + sigmas = [5.0, 3.0, 10.0, 20.0, 10.0, 50.0] + num_points = 500 + test_data, act_locs = _gen_gaussians_even(sigmas, num_points) + widths = np.arange(0.1, max(sigmas)) + found_locs = find_peaks_cwt(test_data, widths, gap_thresh=2, min_snr=0, + min_length=None) + xp_assert_equal(found_locs, act_locs, + check_dtype=False, + err_msg="Found maximum locations did not equal those expected" + ) + + def test_find_peaks_withnoise(self): + """ + Verify that peak locations are (approximately) found + for a series of gaussians with added noise. + """ + sigmas = [5.0, 3.0, 10.0, 20.0, 10.0, 50.0] + num_points = 500 + test_data, act_locs = _gen_gaussians_even(sigmas, num_points) + widths = np.arange(0.1, max(sigmas)) + noise_amp = 0.07 + np.random.seed(18181911) + test_data += (np.random.rand(num_points) - 0.5)*(2*noise_amp) + found_locs = find_peaks_cwt(test_data, widths, min_length=15, + gap_thresh=1, min_snr=noise_amp / 5) + + err_msg ='Different number of peaks found than expected' + assert len(found_locs) == len(act_locs), err_msg + diffs = np.abs(found_locs - act_locs) + max_diffs = np.array(sigmas) / 5 + np.testing.assert_array_less(diffs, max_diffs, 'Maximum location differed' + + f'by more than {max_diffs}') + + def test_find_peaks_nopeak(self): + """ + Verify that no peak is found in + data that's just noise. + """ + noise_amp = 1.0 + num_points = 100 + rng = np.random.RandomState(181819141) + test_data = (rng.rand(num_points) - 0.5)*(2*noise_amp) + widths = np.arange(10, 50) + found_locs = find_peaks_cwt(test_data, widths, min_snr=5, noise_perc=30) + assert len(found_locs) == 0 + + def test_find_peaks_with_non_default_wavelets(self): + x = gaussian(200, 2) + widths = np.array([1, 2, 3, 4]) + a = find_peaks_cwt(x, widths, wavelet=gaussian) + + xp_assert_equal(a, np.asarray([100]), check_dtype=False) + + def test_find_peaks_window_size(self): + """ + Verify that window_size is passed correctly to private function and + affects the result. + """ + sigmas = [2.0, 2.0] + num_points = 1000 + test_data, act_locs = _gen_gaussians_even(sigmas, num_points) + widths = np.arange(0.1, max(sigmas), 0.2) + noise_amp = 0.05 + rng = np.random.RandomState(18181911) + test_data += (rng.rand(num_points) - 0.5)*(2*noise_amp) + + # Possibly contrived negative region to throw off peak finding + # when window_size is too large + test_data[250:320] -= 1 + + found_locs = find_peaks_cwt(test_data, widths, gap_thresh=2, min_snr=3, + min_length=None, window_size=None) + with pytest.raises(AssertionError): + assert found_locs.size == act_locs.size + + found_locs = find_peaks_cwt(test_data, widths, gap_thresh=2, min_snr=3, + min_length=None, window_size=20) + assert found_locs.size == act_locs.size + + def test_find_peaks_with_one_width(self): + """ + Verify that the `width` argument + in `find_peaks_cwt` can be a float + """ + xs = np.arange(0, np.pi, 0.05) + test_data = np.sin(xs) + widths = 1 + found_locs = find_peaks_cwt(test_data, widths) + + np.testing.assert_equal(found_locs, 32) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_short_time_fft.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_short_time_fft.py new file mode 100644 index 0000000000000000000000000000000000000000..563fd26138e48f426a0fc3faddb61a94918b19d8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_short_time_fft.py @@ -0,0 +1,880 @@ +"""Unit tests for module `_short_time_fft`. + +This file's structure loosely groups the tests into the following sequential +categories: + +1. Test function `_calc_dual_canonical_window`. +2. Test for invalid parameters and exceptions in `ShortTimeFFT` (until the + `test_from_window` function). +3. Test algorithmic properties of STFT/ISTFT. Some tests were ported from + ``test_spectral.py``. + +Notes +----- +* Mypy 0.990 does interpret the line:: + + from scipy.stats import norm as normal_distribution + + incorrectly (but the code works), hence a ``type: ignore`` was appended. +""" +import math +from itertools import product +from typing import cast, get_args, Literal + +import numpy as np +import pytest +from scipy._lib._array_api import xp_assert_close, xp_assert_equal +from scipy.fft import fftshift +from scipy.stats import norm as normal_distribution # type: ignore +from scipy.signal import get_window, welch, stft, istft, spectrogram + +from scipy.signal._short_time_fft import FFT_MODE_TYPE, \ + _calc_dual_canonical_window, ShortTimeFFT, PAD_TYPE +from scipy.signal.windows import gaussian + + +def test__calc_dual_canonical_window_roundtrip(): + """Test dual window calculation with a round trip to verify duality. + + Note that this works only for canonical window pairs (having minimal + energy) like a Gaussian. + + The window is the same as in the example of `from ShortTimeFFT.from_dual`. + """ + win = gaussian(51, std=10, sym=True) + d_win = _calc_dual_canonical_window(win, 10) + win2 = _calc_dual_canonical_window(d_win, 10) + xp_assert_close(win2, win) + + +def test__calc_dual_canonical_window_exceptions(): + """Raise all exceptions in `_calc_dual_canonical_window`.""" + # Verify that calculation can fail: + with pytest.raises(ValueError, match="hop=5 is larger than window len.*"): + _calc_dual_canonical_window(np.ones(4), 5) + with pytest.raises(ValueError, match=".* Transform not invertible!"): + _calc_dual_canonical_window(np.array([.1, .2, .3, 0]), 4) + + # Verify that parameter `win` may not be integers: + with pytest.raises(ValueError, match="Parameter 'win' cannot be of int.*"): + _calc_dual_canonical_window(np.ones(4, dtype=int), 1) + + +def test_invalid_initializer_parameters(): + """Verify that exceptions get raised on invalid parameters when + instantiating ShortTimeFFT. """ + with pytest.raises(ValueError, match=r"Parameter win must be 1d, " + + r"but win.shape=\(2, 2\)!"): + ShortTimeFFT(np.ones((2, 2)), hop=4, fs=1) + with pytest.raises(ValueError, match="Parameter win must have " + + "finite entries"): + ShortTimeFFT(np.array([1, np.inf, 2, 3]), hop=4, fs=1) + with pytest.raises(ValueError, match="Parameter hop=0 is not " + + "an integer >= 1!"): + ShortTimeFFT(np.ones(4), hop=0, fs=1) + with pytest.raises(ValueError, match="Parameter hop=2.0 is not " + + "an integer >= 1!"): + # noinspection PyTypeChecker + ShortTimeFFT(np.ones(4), hop=2.0, fs=1) + with pytest.raises(ValueError, match=r"dual_win.shape=\(5,\) must equal " + + r"win.shape=\(4,\)!"): + ShortTimeFFT(np.ones(4), hop=2, fs=1, dual_win=np.ones(5)) + with pytest.raises(ValueError, match="Parameter dual_win must be " + + "a finite array!"): + ShortTimeFFT(np.ones(3), hop=2, fs=1, + dual_win=np.array([np.nan, 2, 3])) + + +def test_exceptions_properties_methods(): + """Verify that exceptions get raised when setting properties or calling + method of ShortTimeFFT to/with invalid values.""" + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=1) + with pytest.raises(ValueError, match="Sampling interval T=-1 must be " + + "positive!"): + SFT.T = -1 + with pytest.raises(ValueError, match="Sampling frequency fs=-1 must be " + + "positive!"): + SFT.fs = -1 + with pytest.raises(ValueError, match="fft_mode='invalid_typ' not in " + + r"\('twosided', 'centered', " + + r"'onesided', 'onesided2X'\)!"): + SFT.fft_mode = 'invalid_typ' + with pytest.raises(ValueError, match="For scaling is None, " + + "fft_mode='onesided2X' is invalid.*"): + SFT.fft_mode = 'onesided2X' + with pytest.raises(ValueError, match="Attribute mfft=7 needs to be " + + "at least the window length.*"): + SFT.mfft = 7 + with pytest.raises(ValueError, match="scaling='invalid' not in.*"): + # noinspection PyTypeChecker + SFT.scale_to('invalid') + with pytest.raises(ValueError, match="phase_shift=3.0 has the unit .*"): + SFT.phase_shift = 3.0 + with pytest.raises(ValueError, match="-mfft < phase_shift < mfft " + + "does not hold.*"): + SFT.phase_shift = 2*SFT.mfft + with pytest.raises(ValueError, match="Parameter padding='invalid' not.*"): + # noinspection PyTypeChecker + g = SFT._x_slices(np.zeros(16), k_off=0, p0=0, p1=1, padding='invalid') + next(g) # execute generator + with pytest.raises(ValueError, match="Trend type must be 'linear' " + + "or 'constant'"): + # noinspection PyTypeChecker + SFT.stft_detrend(np.zeros(16), detr='invalid') + with pytest.raises(ValueError, match="Parameter detr=nan is not a str, " + + "function or None!"): + # noinspection PyTypeChecker + SFT.stft_detrend(np.zeros(16), detr=np.nan) + with pytest.raises(ValueError, match="Invalid Parameter p0=0, p1=200.*"): + SFT.p_range(100, 0, 200) + + with pytest.raises(ValueError, match="f_axis=0 may not be equal to " + + "t_axis=0!"): + SFT.istft(np.zeros((SFT.f_pts, 2)), t_axis=0, f_axis=0) + with pytest.raises(ValueError, match=r"S.shape\[f_axis\]=2 must be equal" + + " to self.f_pts=5.*"): + SFT.istft(np.zeros((2, 2))) + with pytest.raises(ValueError, match=r"S.shape\[t_axis\]=1 needs to have" + + " at least 2 slices.*"): + SFT.istft(np.zeros((SFT.f_pts, 1))) + with pytest.raises(ValueError, match=r".*\(k1=100\) <= \(k_max=12\) " + + "is false!$"): + SFT.istft(np.zeros((SFT.f_pts, 3)), k1=100) + with pytest.raises(ValueError, match=r"\(k1=1\) - \(k0=0\) = 1 has to " + + "be at least.* length 4!"): + SFT.istft(np.zeros((SFT.f_pts, 3)), k0=0, k1=1) + + with pytest.raises(ValueError, match=r"Parameter axes_seq='invalid' " + + r"not in \['tf', 'ft'\]!"): + # noinspection PyTypeChecker + SFT.extent(n=100, axes_seq='invalid') + with pytest.raises(ValueError, match="Attribute fft_mode=twosided must.*"): + SFT.fft_mode = 'twosided' + SFT.extent(n=100) + + +@pytest.mark.parametrize('m', ('onesided', 'onesided2X')) +def test_exceptions_fft_mode_complex_win(m: FFT_MODE_TYPE): + """Verify that one-sided spectra are not allowed with complex-valued + windows or with complex-valued signals. + + The reason being, the `rfft` function only accepts real-valued input. + """ + with pytest.raises(ValueError, + match=f"One-sided spectra, i.e., fft_mode='{m}'.*"): + ShortTimeFFT(np.ones(8)*1j, hop=4, fs=1, fft_mode=m) + + SFT = ShortTimeFFT(np.ones(8)*1j, hop=4, fs=1, fft_mode='twosided') + with pytest.raises(ValueError, + match=f"One-sided spectra, i.e., fft_mode='{m}'.*"): + SFT.fft_mode = m + + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=1, scale_to='psd', fft_mode='onesided') + with pytest.raises(ValueError, match="Complex-valued `x` not allowed for self.*"): + SFT.stft(np.ones(8)*1j) + SFT.fft_mode = 'onesided2X' + with pytest.raises(ValueError, match="Complex-valued `x` not allowed for self.*"): + SFT.stft(np.ones(8)*1j) + + +def test_invalid_fft_mode_RuntimeError(): + """Ensure exception gets raised when property `fft_mode` is invalid. """ + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=1) + SFT._fft_mode = 'invalid_typ' + + with pytest.raises(RuntimeError): + _ = SFT.f + with pytest.raises(RuntimeError): + SFT._fft_func(np.ones(8)) + with pytest.raises(RuntimeError): + SFT._ifft_func(np.ones(8)) + + +@pytest.mark.parametrize('win_params, Nx', [(('gaussian', 2.), 9), # in docstr + ('triang', 7), + (('kaiser', 4.0), 9), + (('exponential', None, 1.), 9), + (4.0, 9)]) +def test_from_window(win_params, Nx: int): + """Verify that `from_window()` handles parameters correctly. + + The window parameterizations are documented in the `get_window` docstring. + """ + w_sym, fs = get_window(win_params, Nx, fftbins=False), 16. + w_per = get_window(win_params, Nx, fftbins=True) + SFT0 = ShortTimeFFT(w_sym, hop=3, fs=fs, fft_mode='twosided', + scale_to='psd', phase_shift=1) + nperseg = len(w_sym) + noverlap = nperseg - SFT0.hop + SFT1 = ShortTimeFFT.from_window(win_params, fs, nperseg, noverlap, + symmetric_win=True, fft_mode='twosided', + scale_to='psd', phase_shift=1) + # periodic window: + SFT2 = ShortTimeFFT.from_window(win_params, fs, nperseg, noverlap, + symmetric_win=False, fft_mode='twosided', + scale_to='psd', phase_shift=1) + # Be informative when comparing instances: + xp_assert_equal(SFT1.win, SFT0.win) + xp_assert_close(SFT2.win, w_per / np.sqrt(sum(w_per**2) * fs)) + for n_ in ('hop', 'T', 'fft_mode', 'mfft', 'scaling', 'phase_shift'): + v0, v1, v2 = (getattr(SFT_, n_) for SFT_ in (SFT0, SFT1, SFT2)) + assert v1 == v0, f"SFT1.{n_}={v1} does not equal SFT0.{n_}={v0}" + assert v2 == v0, f"SFT2.{n_}={v2} does not equal SFT0.{n_}={v0}" + + +def test_dual_win_roundtrip(): + """Verify the duality of `win` and `dual_win`. + + Note that this test does not work for arbitrary windows, since dual windows + are not unique. It always works for invertible STFTs if the windows do not + overlap. + """ + # Non-standard values for keyword arguments (except for `scale_to`): + kw = dict(hop=4, fs=1, fft_mode='twosided', mfft=8, scale_to=None, + phase_shift=2) + SFT0 = ShortTimeFFT(np.ones(4), **kw) + SFT1 = ShortTimeFFT.from_dual(SFT0.dual_win, **kw) + xp_assert_close(SFT1.dual_win, SFT0.win) + + +@pytest.mark.parametrize('scale_to, fac_psd, fac_mag', + [(None, 0.25, 0.125), + ('magnitude', 2.0, 1), + ('psd', 1, 0.5)]) +def test_scaling(scale_to: Literal['magnitude', 'psd'], fac_psd, fac_mag): + """Verify scaling calculations. + + * Verify passing `scale_to`parameter to ``__init__(). + * Roundtrip while changing scaling factor. + """ + SFT = ShortTimeFFT(np.ones(4) * 2, hop=4, fs=1, scale_to=scale_to) + assert SFT.fac_psd == fac_psd + assert SFT.fac_magnitude == fac_mag + # increase coverage by accessing properties twice: + assert SFT.fac_psd == fac_psd + assert SFT.fac_magnitude == fac_mag + + x = np.fft.irfft([0, 0, 7, 0, 0, 0, 0]) # periodic signal + Sx = SFT.stft(x) + Sx_mag, Sx_psd = Sx * SFT.fac_magnitude, Sx * SFT.fac_psd + + SFT.scale_to('magnitude') + x_mag = SFT.istft(Sx_mag, k1=len(x)) + xp_assert_close(x_mag, x) + + SFT.scale_to('psd') + x_psd = SFT.istft(Sx_psd, k1=len(x)) + xp_assert_close(x_psd, x) + + +def test_scale_to(): + """Verify `scale_to()` method.""" + SFT = ShortTimeFFT(np.ones(4) * 2, hop=4, fs=1, scale_to=None) + + SFT.scale_to('magnitude') + assert SFT.scaling == 'magnitude' + assert SFT.fac_psd == 2.0 + assert SFT.fac_magnitude == 1 + + SFT.scale_to('psd') + assert SFT.scaling == 'psd' + assert SFT.fac_psd == 1 + assert SFT.fac_magnitude == 0.5 + + SFT.scale_to('psd') # needed for coverage + + for scale, s_fac in zip(('magnitude', 'psd'), (8, 4)): + SFT = ShortTimeFFT(np.ones(4) * 2, hop=4, fs=1, scale_to=None) + dual_win = SFT.dual_win.copy() + + SFT.scale_to(cast(Literal['magnitude', 'psd'], scale)) + xp_assert_close(SFT.dual_win, dual_win * s_fac) + + +def test_x_slices_padding(): + """Verify padding. + + The reference arrays were taken from the docstrings of `zero_ext`, + `const_ext`, `odd_ext()`, and `even_ext()` from the _array_tools module. + """ + SFT = ShortTimeFFT(np.ones(5), hop=4, fs=1) + x = np.array([[1, 2, 3, 4, 5], [0, 1, 4, 9, 16]], dtype=float) + d = {'zeros': [[[0, 0, 1, 2, 3], [0, 0, 0, 1, 4]], + [[3, 4, 5, 0, 0], [4, 9, 16, 0, 0]]], + 'edge': [[[1, 1, 1, 2, 3], [0, 0, 0, 1, 4]], + [[3, 4, 5, 5, 5], [4, 9, 16, 16, 16]]], + 'even': [[[3, 2, 1, 2, 3], [4, 1, 0, 1, 4]], + [[3, 4, 5, 4, 3], [4, 9, 16, 9, 4]]], + 'odd': [[[-1, 0, 1, 2, 3], [-4, -1, 0, 1, 4]], + [[3, 4, 5, 6, 7], [4, 9, 16, 23, 28]]]} + for p_, xx in d.items(): + gen = SFT._x_slices(np.array(x), 0, 0, 2, padding=cast(PAD_TYPE, p_)) + yy = np.array([y_.copy() for y_ in gen]) # due to inplace copying + xx = np.asarray(xx, dtype=np.float64) + xp_assert_equal(yy, xx, err_msg=f"Failed '{p_}' padding.") + + +def test_invertible(): + """Verify `invertible` property. """ + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=1) + assert SFT.invertible + SFT = ShortTimeFFT(np.ones(8), hop=9, fs=1) + assert not SFT.invertible + + +def test_border_values(): + """Ensure that minimum and maximum values of slices are correct.""" + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=1) + assert SFT.p_min == 0 + assert SFT.k_min == -4 + assert SFT.lower_border_end == (4, 1) + assert SFT.lower_border_end == (4, 1) # needed to test caching + assert SFT.p_max(10) == 4 + assert SFT.k_max(10) == 16 + assert SFT.upper_border_begin(10) == (4, 2) + # Raise exceptions: + with pytest.raises(ValueError, match="^Parameter n must be"): + SFT.upper_border_begin(3) + with pytest.raises(ValueError, match="^Parameter n must be"): + SFT._post_padding(3) + +def test_border_values_exotic(): + """Ensure that the border calculations are correct for windows with + zeros. """ + w = np.array([0, 0, 0, 0, 0, 0, 0, 1.]) + SFT = ShortTimeFFT(w, hop=1, fs=1) + assert SFT.lower_border_end == (0, 0) + + SFT = ShortTimeFFT(np.flip(w), hop=20, fs=1) + assert SFT.upper_border_begin(4) == (16, 1) + assert SFT.upper_border_begin(5) == (16, 1) + assert SFT.upper_border_begin(23) == (36, 2) + assert SFT.upper_border_begin(24) == (36, 2) + assert SFT.upper_border_begin(25) == (36, 2) + + SFT._hop = -1 # provoke unreachable line + with pytest.raises(RuntimeError): + _ = SFT.k_max(4) + with pytest.raises(RuntimeError): + _ = SFT.k_min + + +def test_t(): + """Verify that the times of the slices are correct. """ + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=2) + assert SFT.T == 1/2 + assert SFT.fs == 2. + assert SFT.delta_t == 4 * 1/2 + t_stft = np.arange(0, SFT.p_max(10)) * SFT.delta_t + xp_assert_equal(SFT.t(10), t_stft) + xp_assert_equal(SFT.t(10, 1, 3), t_stft[1:3]) + SFT.T = 1/4 + assert SFT.T == 1/4 + assert SFT.fs == 4 + SFT.fs = 1/8 + assert SFT.fs == 1/8 + assert SFT.T == 8 + + +@pytest.mark.parametrize('fft_mode, f', + [('onesided', [0., 1., 2.]), + ('onesided2X', [0., 1., 2.]), + ('twosided', [0., 1., 2., -2., -1.]), + ('centered', [-2., -1., 0., 1., 2.])]) +def test_f(fft_mode: FFT_MODE_TYPE, f): + """Verify the frequency values property `f`.""" + SFT = ShortTimeFFT(np.ones(5), hop=4, fs=5, fft_mode=fft_mode, + scale_to='psd') + xp_assert_equal(SFT.f, f) + + +@pytest.mark.parametrize('n', [20, 21]) +@pytest.mark.parametrize('m', [5, 6]) +@pytest.mark.parametrize('fft_mode', ['onesided', 'centered']) +def test_extent(n, m, fft_mode: FFT_MODE_TYPE): + """Ensure that the `extent()` method is correct. """ + SFT = ShortTimeFFT(np.ones(m), hop=m, fs=m, fft_mode=fft_mode) + + t0 = SFT.t(n)[0] # first timestamp + t1 = SFT.t(n)[-1] + SFT.delta_t # last timestamp + 1 + t0c, t1c = t0 - SFT.delta_t / 2, t1 - SFT.delta_t / 2 # centered timestamps + + f0 = SFT.f[0] # first frequency + f1 = SFT.f[-1] + SFT.delta_f # last frequency + 1 + f0c, f1c = f0 - SFT.delta_f / 2, f1 - SFT.delta_f / 2 # centered frequencies + + assert SFT.extent(n, 'tf', False) == (t0, t1, f0, f1) + assert SFT.extent(n, 'ft', False) == (f0, f1, t0, t1) + assert SFT.extent(n, 'tf', True) == (t0c, t1c, f0c, f1c) + assert SFT.extent(n, 'ft', True) == (f0c, f1c, t0c, t1c) + + +def test_spectrogram(): + """Verify spectrogram and cross-spectrogram methods. """ + SFT = ShortTimeFFT(np.ones(8), hop=4, fs=1) + x, y = np.ones(10), np.arange(10) + X, Y = SFT.stft(x), SFT.stft(y) + xp_assert_close(SFT.spectrogram(x), X.real**2+X.imag**2) + xp_assert_close(SFT.spectrogram(x, y), X * Y.conj()) + + +@pytest.mark.parametrize('n', [8, 9]) +def test_fft_func_roundtrip(n: int): + """Test roundtrip `ifft_func(fft_func(x)) == x` for all permutations of + relevant parameters. """ + np.random.seed(2394795) + x0 = np.random.rand(n) + w, h_n = np.ones(n), 4 + + pp = dict( + fft_mode=get_args(FFT_MODE_TYPE), + mfft=[None, n, n+1, n+2], + scaling=[None, 'magnitude', 'psd'], + phase_shift=[None, -n+1, 0, n // 2, n-1]) + for f_typ, mfft, scaling, phase_shift in product(*pp.values()): + if f_typ == 'onesided2X' and scaling is None: + continue # this combination is forbidden + SFT = ShortTimeFFT(w, h_n, fs=n, fft_mode=f_typ, mfft=mfft, + scale_to=scaling, phase_shift=phase_shift) + X0 = SFT._fft_func(x0) + x1 = SFT._ifft_func(X0) + xp_assert_close(x0.astype(x1.dtype), x1, + err_msg="_fft_func() roundtrip failed for " + + f"{f_typ=}, {mfft=}, {scaling=}, {phase_shift=}") + + SFT = ShortTimeFFT(w, h_n, fs=1) + SFT._fft_mode = 'invalid_fft' # type: ignore + with pytest.raises(RuntimeError): + SFT._fft_func(x0) + with pytest.raises(RuntimeError): + SFT._ifft_func(x0) + + +@pytest.mark.parametrize('i', range(19)) +def test_impulse_roundtrip(i): + """Roundtrip for an impulse being at different positions `i`.""" + n = 19 + w, h_n = np.ones(8), 3 + x = np.zeros(n) + x[i] = 1 + + SFT = ShortTimeFFT(w, hop=h_n, fs=1, scale_to=None, phase_shift=None) + Sx = SFT.stft(x) + # test slicing the input signal into two parts: + n_q = SFT.nearest_k_p(n // 2) + Sx0 = SFT.stft(x[:n_q], padding='zeros') + Sx1 = SFT.stft(x[n_q:], padding='zeros') + q0_ub = SFT.upper_border_begin(n_q)[1] - SFT.p_min + q1_le = SFT.lower_border_end[1] - SFT.p_min + xp_assert_close(Sx0[:, :q0_ub], Sx[:, :q0_ub], err_msg=f"{i=}") + xp_assert_close(Sx1[:, q1_le:], Sx[:, q1_le-Sx1.shape[1]:], + err_msg=f"{i=}") + + Sx01 = np.hstack((Sx0[:, :q0_ub], + Sx0[:, q0_ub:] + Sx1[:, :q1_le], + Sx1[:, q1_le:])) + xp_assert_close(Sx, Sx01, atol=1e-8, err_msg=f"{i=}") + + y = SFT.istft(Sx, 0, n) + xp_assert_close(y, x, atol=1e-8, err_msg=f"{i=}") + y0 = SFT.istft(Sx, 0, n//2) + xp_assert_close(x[:n//2], y0, atol=1e-8, err_msg=f"{i=}") + y1 = SFT.istft(Sx, n // 2, n) + xp_assert_close(x[n // 2:], y1, atol=1e-8, err_msg=f"{i=}") + + +@pytest.mark.parametrize('hop', [1, 7, 8]) +def test_asymmetric_window_roundtrip(hop: int): + """An asymmetric window could uncover indexing problems. """ + np.random.seed(23371) + + w = np.arange(16) / 8 # must be of type float + w[len(w)//2:] = 1 + SFT = ShortTimeFFT(w, hop, fs=1) + + x = 10 * np.random.randn(64) + Sx = SFT.stft(x) + x1 = SFT.istft(Sx, k1=len(x)) + xp_assert_close(x1, x1, err_msg="Roundtrip for asymmetric window with " + + f" {hop=} failed!") + + +@pytest.mark.parametrize('m_num', [6, 7]) +def test_minimal_length_signal(m_num): + """Verify that the shortest allowed signal works. """ + SFT = ShortTimeFFT(np.ones(m_num), m_num//2, fs=1) + n = math.ceil(m_num/2) + x = np.ones(n) + Sx = SFT.stft(x) + x1 = SFT.istft(Sx, k1=n) + xp_assert_close(x1, x, err_msg=f"Roundtrip minimal length signal ({n=})" + + f" for {m_num} sample window failed!") + with pytest.raises(ValueError, match=rf"len\(x\)={n-1} must be >= ceil.*"): + SFT.stft(x[:-1]) + with pytest.raises(ValueError, match=rf"S.shape\[t_axis\]={Sx.shape[1]-1}" + f" needs to have at least {Sx.shape[1]} slices"): + SFT.istft(Sx[:, :-1], k1=n) + + +def test_tutorial_stft_sliding_win(): + """Verify example in "Sliding Windows" subsection from the "User Guide". + + In :ref:`tutorial_stft_sliding_win` (file ``signal.rst``) of the + :ref:`user_guide` the behavior the border behavior of + ``ShortTimeFFT(np.ones(6), 2, fs=1)`` with a 50 sample signal is discussed. + This test verifies the presented indexes. + """ + SFT = ShortTimeFFT(np.ones(6), 2, fs=1) + + # Lower border: + assert SFT.m_num_mid == 3, f"Slice middle is not 3 but {SFT.m_num_mid=}" + assert SFT.p_min == -1, f"Lowest slice {SFT.p_min=} is not -1" + assert SFT.k_min == -5, f"Lowest slice sample {SFT.p_min=} is not -5" + k_lb, p_lb = SFT.lower_border_end + assert p_lb == 2, f"First unaffected slice {p_lb=} is not 2" + assert k_lb == 5, f"First unaffected sample {k_lb=} is not 5" + + n = 50 # upper signal border + assert (p_max := SFT.p_max(n)) == 27, f"Last slice {p_max=} must be 27" + assert (k_max := SFT.k_max(n)) == 55, f"Last sample {k_max=} must be 55" + k_ub, p_ub = SFT.upper_border_begin(n) + assert p_ub == 24, f"First upper border slice {p_ub=} must be 24" + assert k_ub == 45, f"First upper border slice {k_ub=} must be 45" + + +def test_tutorial_stft_legacy_stft(): + """Verify STFT example in "Comparison with Legacy Implementation" from the + "User Guide". + + In :ref:`tutorial_stft_legacy_stft` (file ``signal.rst``) of the + :ref:`user_guide` the legacy and the new implementation are compared. + """ + fs, N = 200, 1001 # # 200 Hz sampling rate for 5 s signal + t_z = np.arange(N) / fs # time indexes for signal + z = np.exp(2j*np.pi * 70 * (t_z - 0.2 * t_z ** 2)) # complex-valued chirp + + nperseg, noverlap = 50, 40 + win = ('gaussian', 1e-2 * fs) # Gaussian with 0.01 s standard deviation + + # Legacy STFT: + f0_u, t0, Sz0_u = stft(z, fs, win, nperseg, noverlap, + return_onesided=False, scaling='spectrum') + Sz0 = fftshift(Sz0_u, axes=0) + + # New STFT: + SFT = ShortTimeFFT.from_window(win, fs, nperseg, noverlap, + fft_mode='centered', + scale_to='magnitude', phase_shift=None) + Sz1 = SFT.stft(z) + + xp_assert_close(Sz0, Sz1[:, 2:-1]) + + xp_assert_close((abs(Sz1[:, 1]).min(), abs(Sz1[:, 1]).max()), + (6.925060911593139e-07, 8.00271269218721e-07)) + + t0_r, z0_r = istft(Sz0_u, fs, win, nperseg, noverlap, input_onesided=False, + scaling='spectrum') + z1_r = SFT.istft(Sz1, k1=N) + assert len(z0_r) == N + 9 + xp_assert_close(z0_r[:N], z) + xp_assert_close(z1_r, z) + + # Spectrogram is just the absolute square of th STFT: + xp_assert_close(SFT.spectrogram(z), abs(Sz1) ** 2) + + +def test_tutorial_stft_legacy_spectrogram(): + """Verify spectrogram example in "Comparison with Legacy Implementation" + from the "User Guide". + + In :ref:`tutorial_stft_legacy_stft` (file ``signal.rst``) of the + :ref:`user_guide` the legacy and the new implementation are compared. + """ + fs, N = 200, 1001 # 200 Hz sampling rate for almost 5 s signal + t_z = np.arange(N) / fs # time indexes for signal + z = np.exp(2j*np.pi*70 * (t_z - 0.2*t_z**2)) # complex-valued sweep + + nperseg, noverlap = 50, 40 + win = ('gaussian', 1e-2 * fs) # Gaussian with 0.01 s standard dev. + + # Legacy spectrogram: + f2_u, t2, Sz2_u = spectrogram(z, fs, win, nperseg, noverlap, detrend=None, + return_onesided=False, scaling='spectrum', + mode='complex') + + f2, Sz2 = fftshift(f2_u), fftshift(Sz2_u, axes=0) + + # New STFT: + SFT = ShortTimeFFT.from_window(win, fs, nperseg, noverlap, + fft_mode='centered', scale_to='magnitude', + phase_shift=None) + Sz3 = SFT.stft(z, p0=0, p1=(N-noverlap) // SFT.hop, k_offset=nperseg // 2) + t3 = SFT.t(N, p0=0, p1=(N-noverlap) // SFT.hop, k_offset=nperseg // 2) + + xp_assert_close(t2, t3) + xp_assert_close(f2, SFT.f) + xp_assert_close(Sz2, Sz3) + + +def test_permute_axes(): + """Verify correctness of four-dimensional signal by permuting its + shape. """ + n = 25 + SFT = ShortTimeFFT(np.ones(8)/8, hop=3, fs=n) + x0 = np.arange(n, dtype=np.float64) + Sx0 = SFT.stft(x0) + Sx0 = Sx0.reshape((Sx0.shape[0], 1, 1, 1, Sx0.shape[-1])) + SxT = np.moveaxis(Sx0, (0, -1), (-1, 0)) + + atol = 2 * np.finfo(SFT.win.dtype).resolution + for i in range(4): + y = np.reshape(x0, np.roll((n, 1, 1, 1), i)) + Sy = SFT.stft(y, axis=i) + xp_assert_close(Sy, np.moveaxis(Sx0, 0, i)) + + yb0 = SFT.istft(Sy, k1=n, f_axis=i) + xp_assert_close(yb0, y, atol=atol) + # explicit t-axis parameter (for coverage): + yb1 = SFT.istft(Sy, k1=n, f_axis=i, t_axis=Sy.ndim-1) + xp_assert_close(yb1, y, atol=atol) + + SyT = np.moveaxis(Sy, (i, -1), (-1, i)) + xp_assert_close(SyT, np.moveaxis(SxT, 0, i)) + + ybT = SFT.istft(SyT, k1=n, t_axis=i, f_axis=-1) + xp_assert_close(ybT, y, atol=atol) + + +@pytest.mark.parametrize("fft_mode", + ('twosided', 'centered', 'onesided', 'onesided2X')) +def test_roundtrip_multidimensional(fft_mode: FFT_MODE_TYPE): + """Test roundtrip of a multidimensional input signal versus its components. + + This test can uncover potential problems with `fftshift()`. + """ + n = 9 + x = np.arange(4*n*2, dtype=np.float64).reshape(4, n, 2) + SFT = ShortTimeFFT(get_window('hann', 4), hop=2, fs=1, + scale_to='magnitude', fft_mode=fft_mode) + Sx = SFT.stft(x, axis=1) + y = SFT.istft(Sx, k1=n, f_axis=1, t_axis=-1) + xp_assert_close(y, x.astype(y.dtype), err_msg='Multidim. roundtrip failed!') + + for i, j in product(range(x.shape[0]), range(x.shape[2])): + y_ = SFT.istft(Sx[i, :, j, :], k1=n) + xp_assert_close(y_, x[i, :, j].astype(y_.dtype), + err_msg="Multidim. roundtrip for component " + + f"x[{i}, :, {j}] and {fft_mode=} failed!") + +@pytest.mark.parametrize("phase_shift", (0, 4, None)) +def test_roundtrip_two_dimensional(phase_shift: int|None): + """Test roundtrip of a 2 channel input signal with `mfft` set with different + values for `phase_shift` + + Tests for Issue https://github.com/scipy/scipy/issues/21671 + """ + n = 21 + SFT = ShortTimeFFT.from_window('hann', fs=1, nperseg=13, noverlap=7, + mfft=16, phase_shift=phase_shift) + x = np.arange(2*n, dtype=float).reshape(2, n) + Sx = SFT.stft(x) + y = SFT.istft(Sx, k1=n) + xp_assert_close(y, x, atol=2 * np.finfo(SFT.win.dtype).resolution, + err_msg='2-dim. roundtrip failed!') + + +@pytest.mark.parametrize('window, n, nperseg, noverlap', + [('boxcar', 100, 10, 0), # Test no overlap + ('boxcar', 100, 10, 9), # Test high overlap + ('bartlett', 101, 51, 26), # Test odd nperseg + ('hann', 1024, 256, 128), # Test defaults + (('tukey', 0.5), 1152, 256, 64), # Test Tukey + ('hann', 1024, 256, 255), # Test overlapped hann + ('boxcar', 100, 10, 3), # NOLA True, COLA False + ('bartlett', 101, 51, 37), # NOLA True, COLA False + ('hann', 1024, 256, 127), # NOLA True, COLA False + # NOLA True, COLA False: + (('tukey', 0.5), 1152, 256, 14), + ('hann', 1024, 256, 5)]) # NOLA True, COLA False +def test_roundtrip_windows(window, n: int, nperseg: int, noverlap: int): + """Roundtrip test adapted from `test_spectral.TestSTFT`. + + The parameters are taken from the methods test_roundtrip_real(), + test_roundtrip_nola_not_cola(), test_roundtrip_float32(), + test_roundtrip_complex(). + """ + np.random.seed(2394655) + + w = get_window(window, nperseg) + SFT = ShortTimeFFT(w, nperseg - noverlap, fs=1, fft_mode='twosided', + phase_shift=None) + + z = 10 * np.random.randn(n) + 10j * np.random.randn(n) + Sz = SFT.stft(z) + z1 = SFT.istft(Sz, k1=len(z)) + xp_assert_close(z, z1, err_msg="Roundtrip for complex values failed") + + x = 10 * np.random.randn(n) + Sx = SFT.stft(x) + x1 = SFT.istft(Sx, k1=len(z)) + xp_assert_close(x.astype(np.complex128), x1, + err_msg="Roundtrip for float values failed") + + x32 = x.astype(np.float32) + Sx32 = SFT.stft(x32) + x32_1 = SFT.istft(Sx32, k1=len(x32)) + x32_1_r = x32_1.real + xp_assert_close(x32, x32_1_r.astype(np.float32), + err_msg="Roundtrip for 32 Bit float values failed") + xp_assert_close(x32.imag, np.zeros_like(x32.imag), + err_msg="Roundtrip for 32 Bit float values failed") + + +@pytest.mark.parametrize('signal_type', ('real', 'complex')) +def test_roundtrip_complex_window(signal_type): + """Test roundtrip for complex-valued window function + + The purpose of this test is to check if the dual window is calculated + correctly for complex-valued windows. + """ + np.random.seed(1354654) + win = np.exp(2j*np.linspace(0, np.pi, 8)) + SFT = ShortTimeFFT(win, 3, fs=1, fft_mode='twosided') + + z = 10 * np.random.randn(11) + if signal_type == 'complex': + z = z + 2j * z + Sz = SFT.stft(z) + z1 = SFT.istft(Sz, k1=len(z)) + xp_assert_close(z.astype(np.complex128), z1, + err_msg="Roundtrip for complex-valued window failed") + + +def test_average_all_segments(): + """Compare `welch` function with stft mean. + + Ported from `TestSpectrogram.test_average_all_segments` from file + ``test__spectral.py``. + """ + x = np.random.randn(1024) + + fs = 1.0 + window = ('tukey', 0.25) + nperseg, noverlap = 16, 2 + fw, Pw = welch(x, fs, window, nperseg, noverlap) + SFT = ShortTimeFFT.from_window(window, fs, nperseg, noverlap, + fft_mode='onesided2X', scale_to='psd', + phase_shift=None) + # `welch` positions the window differently than the STFT: + P = SFT.spectrogram(x, detr='constant', p0=0, + p1=(len(x)-noverlap)//SFT.hop, k_offset=nperseg//2) + + xp_assert_close(SFT.f, fw) + xp_assert_close(np.mean(P, axis=-1), Pw) + + +@pytest.mark.parametrize('window, N, nperseg, noverlap, mfft', + # from test_roundtrip_padded_FFT: + [('hann', 1024, 256, 128, 512), + ('hann', 1024, 256, 128, 501), + ('boxcar', 100, 10, 0, 33), + (('tukey', 0.5), 1152, 256, 64, 1024), + # from test_roundtrip_padded_signal: + ('boxcar', 101, 10, 0, None), + ('hann', 1000, 256, 128, None), + # from test_roundtrip_boundary_extension: + ('boxcar', 100, 10, 0, None), + ('boxcar', 100, 10, 9, None)]) +@pytest.mark.parametrize('padding', get_args(PAD_TYPE)) +def test_stft_padding_roundtrip(window, N: int, nperseg: int, noverlap: int, + mfft: int, padding): + """Test the parameter 'padding' of `stft` with roundtrips. + + The STFT parametrizations were taken from the methods + `test_roundtrip_padded_FFT`, `test_roundtrip_padded_signal` and + `test_roundtrip_boundary_extension` from class `TestSTFT` in file + ``test_spectral.py``. Note that the ShortTimeFFT does not need the + concept of "boundary extension". + """ + x = normal_distribution.rvs(size=N, random_state=2909) # real signal + z = x * np.exp(1j * np.pi / 4) # complex signal + + SFT = ShortTimeFFT.from_window(window, 1, nperseg, noverlap, + fft_mode='twosided', mfft=mfft) + Sx = SFT.stft(x, padding=padding) + x1 = SFT.istft(Sx, k1=N) + xp_assert_close(x1, x.astype(np.complex128), + err_msg=f"Failed real roundtrip with '{padding}' padding") + + Sz = SFT.stft(z, padding=padding) + z1 = SFT.istft(Sz, k1=N) + xp_assert_close(z1, z, err_msg="Failed complex roundtrip with " + + f" '{padding}' padding") + + +@pytest.mark.parametrize('N_x', (128, 129, 255, 256, 1337)) # signal length +@pytest.mark.parametrize('w_size', (128, 256)) # window length +@pytest.mark.parametrize('t_step', (4, 64)) # SFT time hop +@pytest.mark.parametrize('f_c', (7., 23.)) # frequency of input sine +def test_energy_conservation(N_x: int, w_size: int, t_step: int, f_c: float): + """Test if a `psd`-scaled STFT conserves the L2 norm. + + This test is adapted from MNE-Python [1]_. Besides being battle-tested, + this test has the benefit of using non-standard window including + non-positive values and a 2d input signal. + + Since `ShortTimeFFT` requires the signal length `N_x` to be at least the + window length `w_size`, the parameter `N_x` was changed from + ``(127, 128, 255, 256, 1337)`` to ``(128, 129, 255, 256, 1337)`` to be + more useful. + + .. [1] File ``test_stft.py`` of MNE-Python + https://github.com/mne-tools/mne-python/blob/main/mne/time_frequency/tests/test_stft.py + """ + window = np.sin(np.arange(.5, w_size + .5) / w_size * np.pi) + SFT = ShortTimeFFT(window, t_step, fs=1000, fft_mode='onesided2X', + scale_to='psd') + atol = 2*np.finfo(window.dtype).resolution + N_x = max(N_x, w_size) # minimal sing + # Test with low frequency signal + t = np.arange(N_x).astype(np.float64) + x = np.sin(2 * np.pi * f_c * t * SFT.T) + x = np.array([x, x + 1.]) + X = SFT.stft(x) + xp = SFT.istft(X, k1=N_x) + + max_freq = SFT.f[np.argmax(np.sum(np.abs(X[0]) ** 2, axis=1))] + + assert X.shape[1] == SFT.f_pts + assert np.all(SFT.f >= 0.) + assert np.abs(max_freq - f_c) < 1. + xp_assert_close(x, xp, atol=atol) + + # check L2-norm squared (i.e., energy) conservation: + E_x = np.sum(x**2, axis=-1) * SFT.T # numerical integration + aX2 = X.real**2 + X.imag.real**2 + E_X = np.sum(np.sum(aX2, axis=-1) * SFT.delta_t, axis=-1) * SFT.delta_f + xp_assert_close(E_X, E_x, atol=atol) + + # Test with random signal + np.random.seed(2392795) + x = np.random.randn(2, N_x) + X = SFT.stft(x) + xp = SFT.istft(X, k1=N_x) + + assert X.shape[1] == SFT.f_pts + assert np.all(SFT.f >= 0.) + assert np.abs(max_freq - f_c) < 1. + xp_assert_close(x, xp, atol=atol) + + # check L2-norm squared (i.e., energy) conservation: + E_x = np.sum(x**2, axis=-1) * SFT.T # numeric integration + aX2 = X.real ** 2 + X.imag.real ** 2 + E_X = np.sum(np.sum(aX2, axis=-1) * SFT.delta_t, axis=-1) * SFT.delta_f + xp_assert_close(E_X, E_x, atol=atol) + + # Try with empty array + x = np.zeros((0, N_x)) + X = SFT.stft(x) + xp = SFT.istft(X, k1=N_x) + assert xp.shape == x.shape diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_signaltools.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_signaltools.py new file mode 100644 index 0000000000000000000000000000000000000000..b03c223c7fbc6a8b7490bb6f0a67a176d43c6ebb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_signaltools.py @@ -0,0 +1,3933 @@ +import sys + +from concurrent.futures import ThreadPoolExecutor, as_completed +from decimal import Decimal +from itertools import product +from math import gcd + +import pytest +from pytest import raises as assert_raises +from numpy.testing import ( + assert_equal, + assert_almost_equal, assert_array_equal, assert_array_almost_equal, + assert_allclose, assert_, assert_array_less, + suppress_warnings) +from numpy import array, arange +import numpy as np + +from scipy import fft as sp_fft +from scipy.ndimage import correlate1d +from scipy.optimize import fmin, linear_sum_assignment +from scipy import signal +from scipy.signal import ( + correlate, correlate2d, correlation_lags, convolve, convolve2d, + fftconvolve, oaconvolve, choose_conv_method, envelope, + hilbert, hilbert2, lfilter, lfilter_zi, filtfilt, butter, zpk2tf, zpk2sos, + invres, invresz, vectorstrength, lfiltic, tf2sos, sosfilt, sosfiltfilt, + sosfilt_zi, tf2zpk, BadCoefficients, detrend, unique_roots, residue, + residuez) +from scipy.signal.windows import hann +from scipy.signal._signaltools import (_filtfilt_gust, _compute_factors, + _group_poles) +from scipy.signal._upfirdn import _upfirdn_modes +from scipy._lib import _testutils +from scipy._lib._array_api import xp_assert_close +from scipy._lib._util import ComplexWarning, np_long, np_ulong + + +class _TestConvolve: + + def test_basic(self): + a = [3, 4, 5, 6, 5, 4] + b = [1, 2, 3] + c = convolve(a, b) + assert_array_equal(c, array([3, 10, 22, 28, 32, 32, 23, 12])) + + def test_same(self): + a = [3, 4, 5] + b = [1, 2, 3, 4] + c = convolve(a, b, mode="same") + assert_array_equal(c, array([10, 22, 34])) + + def test_same_eq(self): + a = [3, 4, 5] + b = [1, 2, 3] + c = convolve(a, b, mode="same") + assert_array_equal(c, array([10, 22, 22])) + + def test_complex(self): + x = array([1 + 1j, 2 + 1j, 3 + 1j]) + y = array([1 + 1j, 2 + 1j]) + z = convolve(x, y) + assert_array_equal(z, array([2j, 2 + 6j, 5 + 8j, 5 + 5j])) + + def test_zero_rank(self): + a = 1289 + b = 4567 + c = convolve(a, b) + assert_equal(c, a * b) + + def test_broadcastable(self): + a = np.arange(27).reshape(3, 3, 3) + b = np.arange(3) + for i in range(3): + b_shape = [1]*3 + b_shape[i] = 3 + x = convolve(a, b.reshape(b_shape), method='direct') + y = convolve(a, b.reshape(b_shape), method='fft') + assert_allclose(x, y) + + def test_single_element(self): + a = array([4967]) + b = array([3920]) + c = convolve(a, b) + assert_equal(c, a * b) + + def test_2d_arrays(self): + a = [[1, 2, 3], [3, 4, 5]] + b = [[2, 3, 4], [4, 5, 6]] + c = convolve(a, b) + d = array([[2, 7, 16, 17, 12], + [10, 30, 62, 58, 38], + [12, 31, 58, 49, 30]]) + assert_array_equal(c, d) + + def test_input_swapping(self): + small = arange(8).reshape(2, 2, 2) + big = 1j * arange(27).reshape(3, 3, 3) + big += arange(27)[::-1].reshape(3, 3, 3) + + out_array = array( + [[[0 + 0j, 26 + 0j, 25 + 1j, 24 + 2j], + [52 + 0j, 151 + 5j, 145 + 11j, 93 + 11j], + [46 + 6j, 133 + 23j, 127 + 29j, 81 + 23j], + [40 + 12j, 98 + 32j, 93 + 37j, 54 + 24j]], + + [[104 + 0j, 247 + 13j, 237 + 23j, 135 + 21j], + [282 + 30j, 632 + 96j, 604 + 124j, 330 + 86j], + [246 + 66j, 548 + 180j, 520 + 208j, 282 + 134j], + [142 + 66j, 307 + 161j, 289 + 179j, 153 + 107j]], + + [[68 + 36j, 157 + 103j, 147 + 113j, 81 + 75j], + [174 + 138j, 380 + 348j, 352 + 376j, 186 + 230j], + [138 + 174j, 296 + 432j, 268 + 460j, 138 + 278j], + [70 + 138j, 145 + 323j, 127 + 341j, 63 + 197j]], + + [[32 + 72j, 68 + 166j, 59 + 175j, 30 + 100j], + [68 + 192j, 139 + 433j, 117 + 455j, 57 + 255j], + [38 + 222j, 73 + 499j, 51 + 521j, 21 + 291j], + [12 + 144j, 20 + 318j, 7 + 331j, 0 + 182j]]]) + + assert_array_equal(convolve(small, big, 'full'), out_array) + assert_array_equal(convolve(big, small, 'full'), out_array) + assert_array_equal(convolve(small, big, 'same'), + out_array[1:3, 1:3, 1:3]) + assert_array_equal(convolve(big, small, 'same'), + out_array[0:3, 0:3, 0:3]) + assert_array_equal(convolve(small, big, 'valid'), + out_array[1:3, 1:3, 1:3]) + assert_array_equal(convolve(big, small, 'valid'), + out_array[1:3, 1:3, 1:3]) + + def test_invalid_params(self): + a = [3, 4, 5] + b = [1, 2, 3] + assert_raises(ValueError, convolve, a, b, mode='spam') + assert_raises(ValueError, convolve, a, b, mode='eggs', method='fft') + assert_raises(ValueError, convolve, a, b, mode='ham', method='direct') + assert_raises(ValueError, convolve, a, b, mode='full', method='bacon') + assert_raises(ValueError, convolve, a, b, mode='same', method='bacon') + + +class TestConvolve(_TestConvolve): + + def test_valid_mode2(self): + # See gh-5897 + a = [1, 2, 3, 6, 5, 3] + b = [2, 3, 4, 5, 3, 4, 2, 2, 1] + expected = [70, 78, 73, 65] + + out = convolve(a, b, 'valid') + assert_array_equal(out, expected) + + out = convolve(b, a, 'valid') + assert_array_equal(out, expected) + + a = [1 + 5j, 2 - 1j, 3 + 0j] + b = [2 - 3j, 1 + 0j] + expected = [2 - 3j, 8 - 10j] + + out = convolve(a, b, 'valid') + assert_array_equal(out, expected) + + out = convolve(b, a, 'valid') + assert_array_equal(out, expected) + + def test_same_mode(self): + a = [1, 2, 3, 3, 1, 2] + b = [1, 4, 3, 4, 5, 6, 7, 4, 3, 2, 1, 1, 3] + c = convolve(a, b, 'same') + d = array([57, 61, 63, 57, 45, 36]) + assert_array_equal(c, d) + + def test_invalid_shapes(self): + # By "invalid," we mean that no one + # array has dimensions that are all at + # least as large as the corresponding + # dimensions of the other array. This + # setup should throw a ValueError. + a = np.arange(1, 7).reshape((2, 3)) + b = np.arange(-6, 0).reshape((3, 2)) + + assert_raises(ValueError, convolve, *(a, b), **{'mode': 'valid'}) + assert_raises(ValueError, convolve, *(b, a), **{'mode': 'valid'}) + + def test_convolve_method(self, n=100): + # this types data structure was manually encoded instead of + # using custom filters on the soon-to-be-removed np.sctypes + types = {'uint16', 'uint64', 'int64', 'int32', + 'complex128', 'float64', 'float16', + 'complex64', 'float32', 'int16', + 'uint8', 'uint32', 'int8', 'bool'} + args = [(t1, t2, mode) for t1 in types for t2 in types + for mode in ['valid', 'full', 'same']] + + # These are random arrays, which means test is much stronger than + # convolving testing by convolving two np.ones arrays + rng = np.random.RandomState(42) + array_types = {'i': rng.choice([0, 1], size=n), + 'f': rng.randn(n)} + array_types['b'] = array_types['u'] = array_types['i'] + array_types['c'] = array_types['f'] + 0.5j*array_types['f'] + + for t1, t2, mode in args: + x1 = array_types[np.dtype(t1).kind].astype(t1) + x2 = array_types[np.dtype(t2).kind].astype(t2) + + results = {key: convolve(x1, x2, method=key, mode=mode) + for key in ['fft', 'direct']} + + assert_equal(results['fft'].dtype, results['direct'].dtype) + + if 'bool' in t1 and 'bool' in t2: + assert_equal(choose_conv_method(x1, x2), 'direct') + continue + + # Found by experiment. Found approx smallest value for (rtol, atol) + # threshold to have tests pass. + if any([t in {'complex64', 'float32'} for t in [t1, t2]]): + kwargs = {'rtol': 1.0e-4, 'atol': 1e-6} + elif 'float16' in [t1, t2]: + # atol is default for np.allclose + kwargs = {'rtol': 1e-3, 'atol': 1e-3} + else: + # defaults for np.allclose (different from assert_allclose) + kwargs = {'rtol': 1e-5, 'atol': 1e-8} + + assert_allclose(results['fft'], results['direct'], **kwargs) + + def test_convolve_method_large_input(self): + # This is really a test that convolving two large integers goes to the + # direct method even if they're in the fft method. + for n in [10, 20, 50, 51, 52, 53, 54, 60, 62]: + z = np.array([2**n], dtype=np.int64) + fft = convolve(z, z, method='fft') + direct = convolve(z, z, method='direct') + + # this is the case when integer precision gets to us + # issue #6076 has more detail, hopefully more tests after resolved + if n < 50: + assert_equal(fft, direct) + assert_equal(fft, 2**(2*n)) + assert_equal(direct, 2**(2*n)) + + def test_mismatched_dims(self): + # Input arrays should have the same number of dimensions + assert_raises(ValueError, convolve, [1], 2, method='direct') + assert_raises(ValueError, convolve, 1, [2], method='direct') + assert_raises(ValueError, convolve, [1], 2, method='fft') + assert_raises(ValueError, convolve, 1, [2], method='fft') + assert_raises(ValueError, convolve, [1], [[2]]) + assert_raises(ValueError, convolve, [3], 2) + + @pytest.mark.thread_unsafe + def test_dtype_deprecation(self): + # gh-21211 + a = np.asarray([1, 2, 3, 6, 5, 3], dtype=object) + b = np.asarray([2, 3, 4, 5, 3, 4, 2, 2, 1], dtype=object) + with pytest.deprecated_call(match="dtype=object is not supported"): + convolve(a, b) + + +class _TestConvolve2d: + + def test_2d_arrays(self): + a = [[1, 2, 3], [3, 4, 5]] + b = [[2, 3, 4], [4, 5, 6]] + d = array([[2, 7, 16, 17, 12], + [10, 30, 62, 58, 38], + [12, 31, 58, 49, 30]]) + e = convolve2d(a, b) + assert_array_equal(e, d) + + def test_valid_mode(self): + e = [[2, 3, 4, 5, 6, 7, 8], [4, 5, 6, 7, 8, 9, 10]] + f = [[1, 2, 3], [3, 4, 5]] + h = array([[62, 80, 98, 116, 134]]) + + g = convolve2d(e, f, 'valid') + assert_array_equal(g, h) + + # See gh-5897 + g = convolve2d(f, e, 'valid') + assert_array_equal(g, h) + + def test_valid_mode_complx(self): + e = [[2, 3, 4, 5, 6, 7, 8], [4, 5, 6, 7, 8, 9, 10]] + f = np.array([[1, 2, 3], [3, 4, 5]], dtype=complex) + 1j + h = array([[62.+24.j, 80.+30.j, 98.+36.j, 116.+42.j, 134.+48.j]]) + + g = convolve2d(e, f, 'valid') + assert_array_almost_equal(g, h) + + # See gh-5897 + g = convolve2d(f, e, 'valid') + assert_array_equal(g, h) + + def test_fillvalue(self): + a = [[1, 2, 3], [3, 4, 5]] + b = [[2, 3, 4], [4, 5, 6]] + fillval = 1 + c = convolve2d(a, b, 'full', 'fill', fillval) + d = array([[24, 26, 31, 34, 32], + [28, 40, 62, 64, 52], + [32, 46, 67, 62, 48]]) + assert_array_equal(c, d) + + def test_fillvalue_errors(self): + msg = "could not cast `fillvalue` directly to the output " + with np.testing.suppress_warnings() as sup: + sup.filter(ComplexWarning, "Casting complex values") + with assert_raises(ValueError, match=msg): + convolve2d([[1]], [[1, 2]], fillvalue=1j) + + msg = "`fillvalue` must be scalar or an array with " + with assert_raises(ValueError, match=msg): + convolve2d([[1]], [[1, 2]], fillvalue=[1, 2]) + + def test_fillvalue_empty(self): + # Check that fillvalue being empty raises an error: + assert_raises(ValueError, convolve2d, [[1]], [[1, 2]], + fillvalue=[]) + + def test_wrap_boundary(self): + a = [[1, 2, 3], [3, 4, 5]] + b = [[2, 3, 4], [4, 5, 6]] + c = convolve2d(a, b, 'full', 'wrap') + d = array([[80, 80, 74, 80, 80], + [68, 68, 62, 68, 68], + [80, 80, 74, 80, 80]]) + assert_array_equal(c, d) + + def test_sym_boundary(self): + a = [[1, 2, 3], [3, 4, 5]] + b = [[2, 3, 4], [4, 5, 6]] + c = convolve2d(a, b, 'full', 'symm') + d = array([[34, 30, 44, 62, 66], + [52, 48, 62, 80, 84], + [82, 78, 92, 110, 114]]) + assert_array_equal(c, d) + + @pytest.mark.parametrize('func', [convolve2d, correlate2d]) + @pytest.mark.parametrize('boundary, expected', + [('symm', [[37.0, 42.0, 44.0, 45.0]]), + ('wrap', [[43.0, 44.0, 42.0, 39.0]])]) + def test_same_with_boundary(self, func, boundary, expected): + # Test boundary='symm' and boundary='wrap' with a "long" kernel. + # The size of the kernel requires that the values in the "image" + # be extended more than once to handle the requested boundary method. + # This is a regression test for gh-8684 and gh-8814. + image = np.array([[2.0, -1.0, 3.0, 4.0]]) + kernel = np.ones((1, 21)) + result = func(image, kernel, mode='same', boundary=boundary) + # The expected results were calculated "by hand". Because the + # kernel is all ones, the same result is expected for convolve2d + # and correlate2d. + assert_array_equal(result, expected) + + def test_boundary_extension_same(self): + # Regression test for gh-12686. + # Use ndimage.convolve with appropriate arguments to create the + # expected result. + import scipy.ndimage as ndi + a = np.arange(1, 10*3+1, dtype=float).reshape(10, 3) + b = np.arange(1, 10*10+1, dtype=float).reshape(10, 10) + c = convolve2d(a, b, mode='same', boundary='wrap') + assert_array_equal(c, ndi.convolve(a, b, mode='wrap', origin=(-1, -1))) + + def test_boundary_extension_full(self): + # Regression test for gh-12686. + # Use ndimage.convolve with appropriate arguments to create the + # expected result. + import scipy.ndimage as ndi + a = np.arange(1, 3*3+1, dtype=float).reshape(3, 3) + b = np.arange(1, 6*6+1, dtype=float).reshape(6, 6) + c = convolve2d(a, b, mode='full', boundary='wrap') + apad = np.pad(a, ((3, 3), (3, 3)), 'wrap') + assert_array_equal(c, ndi.convolve(apad, b, mode='wrap')[:-1, :-1]) + + def test_invalid_shapes(self): + # By "invalid," we mean that no one + # array has dimensions that are all at + # least as large as the corresponding + # dimensions of the other array. This + # setup should throw a ValueError. + a = np.arange(1, 7).reshape((2, 3)) + b = np.arange(-6, 0).reshape((3, 2)) + + assert_raises(ValueError, convolve2d, *(a, b), **{'mode': 'valid'}) + assert_raises(ValueError, convolve2d, *(b, a), **{'mode': 'valid'}) + + +class TestConvolve2d(_TestConvolve2d): + + def test_same_mode(self): + e = [[1, 2, 3], [3, 4, 5]] + f = [[2, 3, 4, 5, 6, 7, 8], [4, 5, 6, 7, 8, 9, 10]] + g = convolve2d(e, f, 'same') + h = array([[22, 28, 34], + [80, 98, 116]]) + assert_array_equal(g, h) + + def test_valid_mode2(self): + # See gh-5897 + e = [[1, 2, 3], [3, 4, 5]] + f = [[2, 3, 4, 5, 6, 7, 8], [4, 5, 6, 7, 8, 9, 10]] + expected = [[62, 80, 98, 116, 134]] + + out = convolve2d(e, f, 'valid') + assert_array_equal(out, expected) + + out = convolve2d(f, e, 'valid') + assert_array_equal(out, expected) + + e = [[1 + 1j, 2 - 3j], [3 + 1j, 4 + 0j]] + f = [[2 - 1j, 3 + 2j, 4 + 0j], [4 - 0j, 5 + 1j, 6 - 3j]] + expected = [[27 - 1j, 46. + 2j]] + + out = convolve2d(e, f, 'valid') + assert_array_equal(out, expected) + + # See gh-5897 + out = convolve2d(f, e, 'valid') + assert_array_equal(out, expected) + + def test_consistency_convolve_funcs(self): + # Compare np.convolve, signal.convolve, signal.convolve2d + a = np.arange(5) + b = np.array([3.2, 1.4, 3]) + for mode in ['full', 'valid', 'same']: + assert_almost_equal(np.convolve(a, b, mode=mode), + signal.convolve(a, b, mode=mode)) + assert_almost_equal(np.squeeze( + signal.convolve2d([a], [b], mode=mode)), + signal.convolve(a, b, mode=mode)) + + def test_invalid_dims(self): + assert_raises(ValueError, convolve2d, 3, 4) + assert_raises(ValueError, convolve2d, [3], [4]) + assert_raises(ValueError, convolve2d, [[[3]]], [[[4]]]) + + @pytest.mark.slow + @pytest.mark.xfail_on_32bit("Can't create large array for test") + def test_large_array(self): + # Test indexing doesn't overflow an int (gh-10761) + n = 2**31 // (1000 * np.int64().itemsize) + _testutils.check_free_memory(2 * n * 1001 * np.int64().itemsize / 1e6) + + # Create a chequered pattern of 1s and 0s + a = np.zeros(1001 * n, dtype=np.int64) + a[::2] = 1 + a = np.lib.stride_tricks.as_strided(a, shape=(n, 1000), strides=(8008, 8)) + + count = signal.convolve2d(a, [[1, 1]]) + fails = np.where(count > 1) + assert fails[0].size == 0 + + +class TestFFTConvolve: + + @pytest.mark.parametrize('axes', ['', None, 0, [0], -1, [-1]]) + def test_real(self, axes): + a = array([1, 2, 3]) + expected = array([1, 4, 10, 12, 9.]) + + if axes == '': + out = fftconvolve(a, a) + else: + out = fftconvolve(a, a, axes=axes) + + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [1, [1], -1, [-1]]) + def test_real_axes(self, axes): + a = array([1, 2, 3]) + expected = array([1, 4, 10, 12, 9.]) + + a = np.tile(a, [2, 1]) + expected = np.tile(expected, [2, 1]) + + out = fftconvolve(a, a, axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', ['', None, 0, [0], -1, [-1]]) + def test_complex(self, axes): + a = array([1 + 1j, 2 + 2j, 3 + 3j]) + expected = array([0 + 2j, 0 + 8j, 0 + 20j, 0 + 24j, 0 + 18j]) + + if axes == '': + out = fftconvolve(a, a) + else: + out = fftconvolve(a, a, axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [1, [1], -1, [-1]]) + def test_complex_axes(self, axes): + a = array([1 + 1j, 2 + 2j, 3 + 3j]) + expected = array([0 + 2j, 0 + 8j, 0 + 20j, 0 + 24j, 0 + 18j]) + + a = np.tile(a, [2, 1]) + expected = np.tile(expected, [2, 1]) + + out = fftconvolve(a, a, axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', ['', + None, + [0, 1], + [1, 0], + [0, -1], + [-1, 0], + [-2, 1], + [1, -2], + [-2, -1], + [-1, -2]]) + def test_2d_real_same(self, axes): + a = array([[1, 2, 3], + [4, 5, 6]]) + expected = array([[1, 4, 10, 12, 9], + [8, 26, 56, 54, 36], + [16, 40, 73, 60, 36]]) + + if axes == '': + out = fftconvolve(a, a) + else: + out = fftconvolve(a, a, axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [[1, 2], + [2, 1], + [1, -1], + [-1, 1], + [-2, 2], + [2, -2], + [-2, -1], + [-1, -2]]) + def test_2d_real_same_axes(self, axes): + a = array([[1, 2, 3], + [4, 5, 6]]) + expected = array([[1, 4, 10, 12, 9], + [8, 26, 56, 54, 36], + [16, 40, 73, 60, 36]]) + + a = np.tile(a, [2, 1, 1]) + expected = np.tile(expected, [2, 1, 1]) + + out = fftconvolve(a, a, axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', ['', + None, + [0, 1], + [1, 0], + [0, -1], + [-1, 0], + [-2, 1], + [1, -2], + [-2, -1], + [-1, -2]]) + def test_2d_complex_same(self, axes): + a = array([[1 + 2j, 3 + 4j, 5 + 6j], + [2 + 1j, 4 + 3j, 6 + 5j]]) + expected = array([ + [-3 + 4j, -10 + 20j, -21 + 56j, -18 + 76j, -11 + 60j], + [10j, 44j, 118j, 156j, 122j], + [3 + 4j, 10 + 20j, 21 + 56j, 18 + 76j, 11 + 60j] + ]) + + if axes == '': + out = fftconvolve(a, a) + else: + out = fftconvolve(a, a, axes=axes) + + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [[1, 2], + [2, 1], + [1, -1], + [-1, 1], + [-2, 2], + [2, -2], + [-2, -1], + [-1, -2]]) + def test_2d_complex_same_axes(self, axes): + a = array([[1 + 2j, 3 + 4j, 5 + 6j], + [2 + 1j, 4 + 3j, 6 + 5j]]) + expected = array([ + [-3 + 4j, -10 + 20j, -21 + 56j, -18 + 76j, -11 + 60j], + [10j, 44j, 118j, 156j, 122j], + [3 + 4j, 10 + 20j, 21 + 56j, 18 + 76j, 11 + 60j] + ]) + + a = np.tile(a, [2, 1, 1]) + expected = np.tile(expected, [2, 1, 1]) + + out = fftconvolve(a, a, axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', ['', None, 0, [0], -1, [-1]]) + def test_real_same_mode(self, axes): + a = array([1, 2, 3]) + b = array([3, 3, 5, 6, 8, 7, 9, 0, 1]) + expected_1 = array([35., 41., 47.]) + expected_2 = array([9., 20., 25., 35., 41., 47., 39., 28., 2.]) + + if axes == '': + out = fftconvolve(a, b, 'same') + else: + out = fftconvolve(a, b, 'same', axes=axes) + assert_array_almost_equal(out, expected_1) + + if axes == '': + out = fftconvolve(b, a, 'same') + else: + out = fftconvolve(b, a, 'same', axes=axes) + assert_array_almost_equal(out, expected_2) + + @pytest.mark.parametrize('axes', [1, -1, [1], [-1]]) + def test_real_same_mode_axes(self, axes): + a = array([1, 2, 3]) + b = array([3, 3, 5, 6, 8, 7, 9, 0, 1]) + expected_1 = array([35., 41., 47.]) + expected_2 = array([9., 20., 25., 35., 41., 47., 39., 28., 2.]) + + a = np.tile(a, [2, 1]) + b = np.tile(b, [2, 1]) + expected_1 = np.tile(expected_1, [2, 1]) + expected_2 = np.tile(expected_2, [2, 1]) + + out = fftconvolve(a, b, 'same', axes=axes) + assert_array_almost_equal(out, expected_1) + + out = fftconvolve(b, a, 'same', axes=axes) + assert_array_almost_equal(out, expected_2) + + @pytest.mark.parametrize('axes', ['', None, 0, [0], -1, [-1]]) + def test_valid_mode_real(self, axes): + # See gh-5897 + a = array([3, 2, 1]) + b = array([3, 3, 5, 6, 8, 7, 9, 0, 1]) + expected = array([24., 31., 41., 43., 49., 25., 12.]) + + if axes == '': + out = fftconvolve(a, b, 'valid') + else: + out = fftconvolve(a, b, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + if axes == '': + out = fftconvolve(b, a, 'valid') + else: + out = fftconvolve(b, a, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [1, [1]]) + def test_valid_mode_real_axes(self, axes): + # See gh-5897 + a = array([3, 2, 1]) + b = array([3, 3, 5, 6, 8, 7, 9, 0, 1]) + expected = array([24., 31., 41., 43., 49., 25., 12.]) + + a = np.tile(a, [2, 1]) + b = np.tile(b, [2, 1]) + expected = np.tile(expected, [2, 1]) + + out = fftconvolve(a, b, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', ['', None, 0, [0], -1, [-1]]) + def test_valid_mode_complex(self, axes): + a = array([3 - 1j, 2 + 7j, 1 + 0j]) + b = array([3 + 2j, 3 - 3j, 5 + 0j, 6 - 1j, 8 + 0j]) + expected = array([45. + 12.j, 30. + 23.j, 48 + 32.j]) + + if axes == '': + out = fftconvolve(a, b, 'valid') + else: + out = fftconvolve(a, b, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + if axes == '': + out = fftconvolve(b, a, 'valid') + else: + out = fftconvolve(b, a, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [1, [1], -1, [-1]]) + def test_valid_mode_complex_axes(self, axes): + a = array([3 - 1j, 2 + 7j, 1 + 0j]) + b = array([3 + 2j, 3 - 3j, 5 + 0j, 6 - 1j, 8 + 0j]) + expected = array([45. + 12.j, 30. + 23.j, 48 + 32.j]) + + a = np.tile(a, [2, 1]) + b = np.tile(b, [2, 1]) + expected = np.tile(expected, [2, 1]) + + out = fftconvolve(a, b, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + out = fftconvolve(b, a, 'valid', axes=axes) + assert_array_almost_equal(out, expected) + + def test_valid_mode_ignore_nonaxes(self): + # See gh-5897 + a = array([3, 2, 1]) + b = array([3, 3, 5, 6, 8, 7, 9, 0, 1]) + expected = array([24., 31., 41., 43., 49., 25., 12.]) + + a = np.tile(a, [2, 1]) + b = np.tile(b, [1, 1]) + expected = np.tile(expected, [2, 1]) + + out = fftconvolve(a, b, 'valid', axes=1) + assert_array_almost_equal(out, expected) + + def test_empty(self): + # Regression test for #1745: crashes with 0-length input. + assert_(fftconvolve([], []).size == 0) + assert_(fftconvolve([5, 6], []).size == 0) + assert_(fftconvolve([], [7]).size == 0) + + def test_zero_rank(self): + a = array(4967) + b = array(3920) + out = fftconvolve(a, b) + assert_equal(out, a * b) + + def test_single_element(self): + a = array([4967]) + b = array([3920]) + out = fftconvolve(a, b) + assert_equal(out, a * b) + + @pytest.mark.parametrize('axes', ['', None, 0, [0], -1, [-1]]) + def test_random_data(self, axes): + np.random.seed(1234) + a = np.random.rand(1233) + 1j * np.random.rand(1233) + b = np.random.rand(1321) + 1j * np.random.rand(1321) + expected = np.convolve(a, b, 'full') + + if axes == '': + out = fftconvolve(a, b, 'full') + else: + out = fftconvolve(a, b, 'full', axes=axes) + assert_(np.allclose(out, expected, rtol=1e-10)) + + @pytest.mark.parametrize('axes', [1, [1], -1, [-1]]) + def test_random_data_axes(self, axes): + np.random.seed(1234) + a = np.random.rand(1233) + 1j * np.random.rand(1233) + b = np.random.rand(1321) + 1j * np.random.rand(1321) + expected = np.convolve(a, b, 'full') + + a = np.tile(a, [2, 1]) + b = np.tile(b, [2, 1]) + expected = np.tile(expected, [2, 1]) + + out = fftconvolve(a, b, 'full', axes=axes) + assert_(np.allclose(out, expected, rtol=1e-10)) + + @pytest.mark.parametrize('axes', [[1, 4], + [4, 1], + [1, -1], + [-1, 1], + [-4, 4], + [4, -4], + [-4, -1], + [-1, -4]]) + def test_random_data_multidim_axes(self, axes): + a_shape, b_shape = (123, 22), (132, 11) + np.random.seed(1234) + a = np.random.rand(*a_shape) + 1j * np.random.rand(*a_shape) + b = np.random.rand(*b_shape) + 1j * np.random.rand(*b_shape) + expected = convolve2d(a, b, 'full') + + a = a[:, :, None, None, None] + b = b[:, :, None, None, None] + expected = expected[:, :, None, None, None] + + a = np.moveaxis(a.swapaxes(0, 2), 1, 4) + b = np.moveaxis(b.swapaxes(0, 2), 1, 4) + expected = np.moveaxis(expected.swapaxes(0, 2), 1, 4) + + # use 1 for dimension 2 in a and 3 in b to test broadcasting + a = np.tile(a, [2, 1, 3, 1, 1]) + b = np.tile(b, [2, 1, 1, 4, 1]) + expected = np.tile(expected, [2, 1, 3, 4, 1]) + + out = fftconvolve(a, b, 'full', axes=axes) + assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + @pytest.mark.slow + @pytest.mark.parametrize( + 'n', + list(range(1, 100)) + + list(range(1000, 1500)) + + np.random.RandomState(1234).randint(1001, 10000, 5).tolist()) + def test_many_sizes(self, n): + a = np.random.rand(n) + 1j * np.random.rand(n) + b = np.random.rand(n) + 1j * np.random.rand(n) + expected = np.convolve(a, b, 'full') + + out = fftconvolve(a, b, 'full') + assert_allclose(out, expected, atol=1e-10) + + out = fftconvolve(a, b, 'full', axes=[0]) + assert_allclose(out, expected, atol=1e-10) + + @pytest.mark.thread_unsafe + def test_fft_nan(self): + n = 1000 + rng = np.random.default_rng(43876432987) + sig_nan = rng.standard_normal(n) + + for val in [np.nan, np.inf]: + sig_nan[100] = val + coeffs = signal.firwin(200, 0.2) + + msg = "Use of fft convolution.*|invalid value encountered.*" + with pytest.warns(RuntimeWarning, match=msg): + signal.convolve(sig_nan, coeffs, mode='same', method='fft') + +def fftconvolve_err(*args, **kwargs): + raise RuntimeError('Fell back to fftconvolve') + + +def gen_oa_shapes(sizes): + return [(a, b) for a, b in product(sizes, repeat=2) + if abs(a - b) > 3] + + +def gen_oa_shapes_2d(sizes): + shapes0 = gen_oa_shapes(sizes) + shapes1 = gen_oa_shapes(sizes) + shapes = [ishapes0+ishapes1 for ishapes0, ishapes1 in + zip(shapes0, shapes1)] + + modes = ['full', 'valid', 'same'] + return [ishapes+(imode,) for ishapes, imode in product(shapes, modes) + if imode != 'valid' or + (ishapes[0] > ishapes[1] and ishapes[2] > ishapes[3]) or + (ishapes[0] < ishapes[1] and ishapes[2] < ishapes[3])] + + +def gen_oa_shapes_eq(sizes): + return [(a, b) for a, b in product(sizes, repeat=2) + if a >= b] + + +class TestOAConvolve: + @pytest.mark.slow() + @pytest.mark.parametrize('shape_a_0, shape_b_0', + gen_oa_shapes_eq(list(range(100)) + + list(range(100, 1000, 23))) + ) + def test_real_manylens(self, shape_a_0, shape_b_0): + a = np.random.rand(shape_a_0) + b = np.random.rand(shape_b_0) + + expected = fftconvolve(a, b) + out = oaconvolve(a, b) + + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('shape_a_0, shape_b_0', + gen_oa_shapes([50, 47, 6, 4, 1])) + @pytest.mark.parametrize('is_complex', [True, False]) + @pytest.mark.parametrize('mode', ['full', 'valid', 'same']) + def test_1d_noaxes(self, shape_a_0, shape_b_0, + is_complex, mode, monkeypatch): + a = np.random.rand(shape_a_0) + b = np.random.rand(shape_b_0) + if is_complex: + a = a + 1j*np.random.rand(shape_a_0) + b = b + 1j*np.random.rand(shape_b_0) + + expected = fftconvolve(a, b, mode=mode) + + monkeypatch.setattr(signal._signaltools, 'fftconvolve', + fftconvolve_err) + out = oaconvolve(a, b, mode=mode) + + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [0, 1]) + @pytest.mark.parametrize('shape_a_0, shape_b_0', + gen_oa_shapes([50, 47, 6, 4])) + @pytest.mark.parametrize('shape_a_extra', [1, 3]) + @pytest.mark.parametrize('shape_b_extra', [1, 3]) + @pytest.mark.parametrize('is_complex', [True, False]) + @pytest.mark.parametrize('mode', ['full', 'valid', 'same']) + def test_1d_axes(self, axes, shape_a_0, shape_b_0, + shape_a_extra, shape_b_extra, + is_complex, mode, monkeypatch): + ax_a = [shape_a_extra]*2 + ax_b = [shape_b_extra]*2 + ax_a[axes] = shape_a_0 + ax_b[axes] = shape_b_0 + + a = np.random.rand(*ax_a) + b = np.random.rand(*ax_b) + if is_complex: + a = a + 1j*np.random.rand(*ax_a) + b = b + 1j*np.random.rand(*ax_b) + + expected = fftconvolve(a, b, mode=mode, axes=axes) + + monkeypatch.setattr(signal._signaltools, 'fftconvolve', + fftconvolve_err) + out = oaconvolve(a, b, mode=mode, axes=axes) + + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('shape_a_0, shape_b_0, ' + 'shape_a_1, shape_b_1, mode', + gen_oa_shapes_2d([50, 47, 6, 4])) + @pytest.mark.parametrize('is_complex', [True, False]) + def test_2d_noaxes(self, shape_a_0, shape_b_0, + shape_a_1, shape_b_1, mode, + is_complex, monkeypatch): + a = np.random.rand(shape_a_0, shape_a_1) + b = np.random.rand(shape_b_0, shape_b_1) + if is_complex: + a = a + 1j*np.random.rand(shape_a_0, shape_a_1) + b = b + 1j*np.random.rand(shape_b_0, shape_b_1) + + expected = fftconvolve(a, b, mode=mode) + + monkeypatch.setattr(signal._signaltools, 'fftconvolve', + fftconvolve_err) + out = oaconvolve(a, b, mode=mode) + + assert_array_almost_equal(out, expected) + + @pytest.mark.parametrize('axes', [[0, 1], [0, 2], [1, 2]]) + @pytest.mark.parametrize('shape_a_0, shape_b_0, ' + 'shape_a_1, shape_b_1, mode', + gen_oa_shapes_2d([50, 47, 6, 4])) + @pytest.mark.parametrize('shape_a_extra', [1, 3]) + @pytest.mark.parametrize('shape_b_extra', [1, 3]) + @pytest.mark.parametrize('is_complex', [True, False]) + def test_2d_axes(self, axes, shape_a_0, shape_b_0, + shape_a_1, shape_b_1, mode, + shape_a_extra, shape_b_extra, + is_complex, monkeypatch): + ax_a = [shape_a_extra]*3 + ax_b = [shape_b_extra]*3 + ax_a[axes[0]] = shape_a_0 + ax_b[axes[0]] = shape_b_0 + ax_a[axes[1]] = shape_a_1 + ax_b[axes[1]] = shape_b_1 + + a = np.random.rand(*ax_a) + b = np.random.rand(*ax_b) + if is_complex: + a = a + 1j*np.random.rand(*ax_a) + b = b + 1j*np.random.rand(*ax_b) + + expected = fftconvolve(a, b, mode=mode, axes=axes) + + monkeypatch.setattr(signal._signaltools, 'fftconvolve', + fftconvolve_err) + out = oaconvolve(a, b, mode=mode, axes=axes) + + assert_array_almost_equal(out, expected) + + def test_empty(self): + # Regression test for #1745: crashes with 0-length input. + assert_(oaconvolve([], []).size == 0) + assert_(oaconvolve([5, 6], []).size == 0) + assert_(oaconvolve([], [7]).size == 0) + + def test_zero_rank(self): + a = array(4967) + b = array(3920) + out = oaconvolve(a, b) + assert_equal(out, a * b) + + def test_single_element(self): + a = array([4967]) + b = array([3920]) + out = oaconvolve(a, b) + assert_equal(out, a * b) + + +class TestAllFreqConvolves: + + @pytest.mark.parametrize('convapproach', + [fftconvolve, oaconvolve]) + def test_invalid_shapes(self, convapproach): + a = np.arange(1, 7).reshape((2, 3)) + b = np.arange(-6, 0).reshape((3, 2)) + with assert_raises(ValueError, + match="For 'valid' mode, one must be at least " + "as large as the other in every dimension"): + convapproach(a, b, mode='valid') + + @pytest.mark.parametrize('convapproach', + [fftconvolve, oaconvolve]) + def test_invalid_shapes_axes(self, convapproach): + a = np.zeros([5, 6, 2, 1]) + b = np.zeros([5, 6, 3, 1]) + with assert_raises(ValueError, + match=r"incompatible shapes for in1 and in2:" + r" \(5L?, 6L?, 2L?, 1L?\) and" + r" \(5L?, 6L?, 3L?, 1L?\)"): + convapproach(a, b, axes=[0, 1]) + + @pytest.mark.parametrize('a,b', + [([1], 2), + (1, [2]), + ([3], [[2]])]) + @pytest.mark.parametrize('convapproach', + [fftconvolve, oaconvolve]) + def test_mismatched_dims(self, a, b, convapproach): + with assert_raises(ValueError, + match="in1 and in2 should have the same" + " dimensionality"): + convapproach(a, b) + + @pytest.mark.parametrize('convapproach', + [fftconvolve, oaconvolve]) + def test_invalid_flags(self, convapproach): + with assert_raises(ValueError, + match="acceptable mode flags are 'valid'," + " 'same', or 'full'"): + convapproach([1], [2], mode='chips') + + with assert_raises(ValueError, + match="when provided, axes cannot be empty"): + convapproach([1], [2], axes=[]) + + with assert_raises(ValueError, match="axes must be a scalar or " + "iterable of integers"): + convapproach([1], [2], axes=[[1, 2], [3, 4]]) + + with assert_raises(ValueError, match="axes must be a scalar or " + "iterable of integers"): + convapproach([1], [2], axes=[1., 2., 3., 4.]) + + with assert_raises(ValueError, + match="axes exceeds dimensionality of input"): + convapproach([1], [2], axes=[1]) + + with assert_raises(ValueError, + match="axes exceeds dimensionality of input"): + convapproach([1], [2], axes=[-2]) + + with assert_raises(ValueError, + match="all axes must be unique"): + convapproach([1], [2], axes=[0, 0]) + + @pytest.mark.filterwarnings('ignore::DeprecationWarning') + @pytest.mark.parametrize('dtype', [np.longdouble, np.clongdouble]) + def test_longdtype_input(self, dtype): + x = np.random.random((27, 27)).astype(dtype) + y = np.random.random((4, 4)).astype(dtype) + if np.iscomplexobj(dtype()): + x += .1j + y -= .1j + + res = fftconvolve(x, y) + assert_allclose(res, convolve(x, y, method='direct')) + assert res.dtype == dtype + + +class TestMedFilt: + + IN = [[50, 50, 50, 50, 50, 92, 18, 27, 65, 46], + [50, 50, 50, 50, 50, 0, 72, 77, 68, 66], + [50, 50, 50, 50, 50, 46, 47, 19, 64, 77], + [50, 50, 50, 50, 50, 42, 15, 29, 95, 35], + [50, 50, 50, 50, 50, 46, 34, 9, 21, 66], + [70, 97, 28, 68, 78, 77, 61, 58, 71, 42], + [64, 53, 44, 29, 68, 32, 19, 68, 24, 84], + [3, 33, 53, 67, 1, 78, 74, 55, 12, 83], + [7, 11, 46, 70, 60, 47, 24, 43, 61, 26], + [32, 61, 88, 7, 39, 4, 92, 64, 45, 61]] + + OUT = [[0, 50, 50, 50, 42, 15, 15, 18, 27, 0], + [0, 50, 50, 50, 50, 42, 19, 21, 29, 0], + [50, 50, 50, 50, 50, 47, 34, 34, 46, 35], + [50, 50, 50, 50, 50, 50, 42, 47, 64, 42], + [50, 50, 50, 50, 50, 50, 46, 55, 64, 35], + [33, 50, 50, 50, 50, 47, 46, 43, 55, 26], + [32, 50, 50, 50, 50, 47, 46, 45, 55, 26], + [7, 46, 50, 50, 47, 46, 46, 43, 45, 21], + [0, 32, 33, 39, 32, 32, 43, 43, 43, 0], + [0, 7, 11, 7, 4, 4, 19, 19, 24, 0]] + + KERNEL_SIZE = [7,3] + + def test_basic(self): + d = signal.medfilt(self.IN, self.KERNEL_SIZE) + e = signal.medfilt2d(np.array(self.IN, float), self.KERNEL_SIZE) + assert_array_equal(d, self.OUT) + assert_array_equal(d, e) + + @pytest.mark.parametrize('dtype', [np.ubyte, np.byte, np.ushort, np.short, + np_ulong, np_long, np.ulonglong, np.ulonglong, + np.float32, np.float64]) + def test_types(self, dtype): + # volume input and output types match + in_typed = np.array(self.IN, dtype=dtype) + assert_equal(signal.medfilt(in_typed).dtype, dtype) + assert_equal(signal.medfilt2d(in_typed).dtype, dtype) + + @pytest.mark.parametrize('dtype', [np.bool_, np.complex64, np.complex128, + np.clongdouble, np.float16, np.object_, + "float96", "float128"]) + def test_invalid_dtypes(self, dtype): + # We can only test this on platforms that support a native type of float96 or + # float128; comparing to np.longdouble allows us to filter out non-native types + if (dtype in ["float96", "float128"] + and np.finfo(np.longdouble).dtype != dtype): + pytest.skip(f"Platform does not support {dtype}") + + in_typed = np.array(self.IN, dtype=dtype) + with pytest.raises(ValueError, match="not supported"): + signal.medfilt(in_typed) + + with pytest.raises(ValueError, match="not supported"): + signal.medfilt2d(in_typed) + + def test_none(self): + # gh-1651, trac #1124. Ensure this does not segfault. + msg = "dtype=object is not supported by medfilt" + with assert_raises(ValueError, match=msg): + signal.medfilt(None) + + def test_odd_strides(self): + # Avoid a regression with possible contiguous + # numpy arrays that have odd strides. The stride value below gets + # us into wrong memory if used (but it does not need to be used) + dummy = np.arange(10, dtype=np.float64) + a = dummy[5:6] + a.strides = 16 + assert_(signal.medfilt(a, 1) == 5.) + + @pytest.mark.parametrize("dtype", [np.ubyte, np.float32, np.float64]) + def test_medfilt2d_parallel(self, dtype): + in_typed = np.array(self.IN, dtype=dtype) + expected = np.array(self.OUT, dtype=dtype) + + # This is used to simplify the indexing calculations. + assert in_typed.shape == expected.shape + + # We'll do the calculation in four chunks. M1 and N1 are the dimensions + # of the first output chunk. We have to extend the input by half the + # kernel size to be able to calculate the full output chunk. + M1 = expected.shape[0] // 2 + N1 = expected.shape[1] // 2 + offM = self.KERNEL_SIZE[0] // 2 + 1 + offN = self.KERNEL_SIZE[1] // 2 + 1 + + def apply(chunk): + # in = slice of in_typed to use. + # sel = slice of output to crop it to the correct region. + # out = slice of output array to store in. + M, N = chunk + if M == 0: + Min = slice(0, M1 + offM) + Msel = slice(0, -offM) + Mout = slice(0, M1) + else: + Min = slice(M1 - offM, None) + Msel = slice(offM, None) + Mout = slice(M1, None) + if N == 0: + Nin = slice(0, N1 + offN) + Nsel = slice(0, -offN) + Nout = slice(0, N1) + else: + Nin = slice(N1 - offN, None) + Nsel = slice(offN, None) + Nout = slice(N1, None) + + # Do the calculation, but do not write to the output in the threads. + chunk_data = in_typed[Min, Nin] + med = signal.medfilt2d(chunk_data, self.KERNEL_SIZE) + return med[Msel, Nsel], Mout, Nout + + # Give each chunk to a different thread. + output = np.zeros_like(expected) + with ThreadPoolExecutor(max_workers=4) as pool: + chunks = {(0, 0), (0, 1), (1, 0), (1, 1)} + futures = {pool.submit(apply, chunk) for chunk in chunks} + + # Store each result in the output as it arrives. + for future in as_completed(futures): + data, Mslice, Nslice = future.result() + output[Mslice, Nslice] = data + + assert_array_equal(output, expected) + + +class TestWiener: + + def test_basic(self): + g = array([[5, 6, 4, 3], + [3, 5, 6, 2], + [2, 3, 5, 6], + [1, 6, 9, 7]], 'd') + h = array([[2.16374269, 3.2222222222, 2.8888888889, 1.6666666667], + [2.666666667, 4.33333333333, 4.44444444444, 2.8888888888], + [2.222222222, 4.4444444444, 5.4444444444, 4.801066874837], + [1.33333333333, 3.92735042735, 6.0712560386, 5.0404040404]]) + assert_array_almost_equal(signal.wiener(g), h, decimal=6) + assert_array_almost_equal(signal.wiener(g, mysize=3), h, decimal=6) + + +padtype_options = ["mean", "median", "minimum", "maximum", "line"] +padtype_options += _upfirdn_modes + + +class TestResample: + def test_basic(self): + # Some basic tests + + # Regression test for issue #3603. + # window.shape must equal to sig.shape[0] + sig = np.arange(128) + num = 256 + win = signal.get_window(('kaiser', 8.0), 160) + assert_raises(ValueError, signal.resample, sig, num, window=win) + + # Other degenerate conditions + assert_raises(ValueError, signal.resample_poly, sig, 'yo', 1) + assert_raises(ValueError, signal.resample_poly, sig, 1, 0) + assert_raises(ValueError, signal.resample_poly, sig, 2, 1, padtype='') + assert_raises(ValueError, signal.resample_poly, sig, 2, 1, + padtype='mean', cval=10) + + # test for issue #6505 - should not modify window.shape when axis ≠ 0 + sig2 = np.tile(np.arange(160), (2, 1)) + signal.resample(sig2, num, axis=-1, window=win) + assert_(win.shape == (160,)) + + @pytest.mark.parametrize('window', (None, 'hamming')) + @pytest.mark.parametrize('N', (20, 19)) + @pytest.mark.parametrize('num', (100, 101, 10, 11)) + def test_rfft(self, N, num, window): + # Make sure the speed up using rfft gives the same result as the normal + # way using fft + x = np.linspace(0, 10, N, endpoint=False) + y = np.cos(-x**2/6.0) + assert_allclose(signal.resample(y, num, window=window), + signal.resample(y + 0j, num, window=window).real) + + y = np.array([np.cos(-x**2/6.0), np.sin(-x**2/6.0)]) + y_complex = y + 0j + assert_allclose( + signal.resample(y, num, axis=1, window=window), + signal.resample(y_complex, num, axis=1, window=window).real, + atol=1e-9) + + def test_input_domain(self): + # Test if both input domain modes produce the same results. + tsig = np.arange(256) + 0j + fsig = sp_fft.fft(tsig) + num = 256 + assert_allclose( + signal.resample(fsig, num, domain='freq'), + signal.resample(tsig, num, domain='time'), + atol=1e-9) + + @pytest.mark.parametrize('nx', (1, 2, 3, 5, 8)) + @pytest.mark.parametrize('ny', (1, 2, 3, 5, 8)) + @pytest.mark.parametrize('dtype', ('float', 'complex')) + def test_dc(self, nx, ny, dtype): + x = np.array([1] * nx, dtype) + y = signal.resample(x, ny) + assert_allclose(y, [1] * ny) + + @pytest.mark.thread_unsafe # due to Cython fused types, see cython#6506 + @pytest.mark.parametrize('padtype', padtype_options) + def test_mutable_window(self, padtype): + # Test that a mutable window is not modified + impulse = np.zeros(3) + window = np.random.RandomState(0).randn(2) + window_orig = window.copy() + signal.resample_poly(impulse, 5, 1, window=window, padtype=padtype) + assert_array_equal(window, window_orig) + + @pytest.mark.parametrize('padtype', padtype_options) + def test_output_float32(self, padtype): + # Test that float32 inputs yield a float32 output + x = np.arange(10, dtype=np.float32) + h = np.array([1, 1, 1], dtype=np.float32) + y = signal.resample_poly(x, 1, 2, window=h, padtype=padtype) + assert y.dtype == np.float32 + + @pytest.mark.parametrize('padtype', padtype_options) + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_output_match_dtype(self, padtype, dtype): + # Test that the dtype of x is preserved per issue #14733 + x = np.arange(10, dtype=dtype) + y = signal.resample_poly(x, 1, 2, padtype=padtype) + assert y.dtype == x.dtype + + @pytest.mark.parametrize( + "method, ext, padtype", + [("fft", False, None)] + + list( + product( + ["polyphase"], [False, True], padtype_options, + ) + ), + ) + def test_resample_methods(self, method, ext, padtype): + # Test resampling of sinusoids and random noise (1-sec) + rate = 100 + rates_to = [49, 50, 51, 99, 100, 101, 199, 200, 201] + + # Sinusoids, windowed to avoid edge artifacts + t = np.arange(rate) / float(rate) + freqs = np.array((1., 10., 40.))[:, np.newaxis] + x = np.sin(2 * np.pi * freqs * t) * hann(rate) + + for rate_to in rates_to: + t_to = np.arange(rate_to) / float(rate_to) + y_tos = np.sin(2 * np.pi * freqs * t_to) * hann(rate_to) + if method == 'fft': + y_resamps = signal.resample(x, rate_to, axis=-1) + else: + if ext and rate_to != rate: + # Match default window design + g = gcd(rate_to, rate) + up = rate_to // g + down = rate // g + max_rate = max(up, down) + f_c = 1. / max_rate + half_len = 10 * max_rate + window = signal.firwin(2 * half_len + 1, f_c, + window=('kaiser', 5.0)) + polyargs = {'window': window, 'padtype': padtype} + else: + polyargs = {'padtype': padtype} + + y_resamps = signal.resample_poly(x, rate_to, rate, axis=-1, + **polyargs) + + for y_to, y_resamp, freq in zip(y_tos, y_resamps, freqs): + if freq >= 0.5 * rate_to: + y_to.fill(0.) # mostly low-passed away + if padtype in ['minimum', 'maximum']: + assert_allclose(y_resamp, y_to, atol=3e-1) + else: + assert_allclose(y_resamp, y_to, atol=1e-3) + else: + assert_array_equal(y_to.shape, y_resamp.shape) + corr = np.corrcoef(y_to, y_resamp)[0, 1] + assert_(corr > 0.99, msg=(corr, rate, rate_to)) + + # Random data + rng = np.random.RandomState(0) + x = hann(rate) * np.cumsum(rng.randn(rate)) # low-pass, wind + for rate_to in rates_to: + # random data + t_to = np.arange(rate_to) / float(rate_to) + y_to = np.interp(t_to, t, x) + if method == 'fft': + y_resamp = signal.resample(x, rate_to) + else: + y_resamp = signal.resample_poly(x, rate_to, rate, + padtype=padtype) + assert_array_equal(y_to.shape, y_resamp.shape) + corr = np.corrcoef(y_to, y_resamp)[0, 1] + assert_(corr > 0.99, msg=corr) + + # More tests of fft method (Master 0.18.1 fails these) + if method == 'fft': + x1 = np.array([1.+0.j, 0.+0.j]) + y1_test = signal.resample(x1, 4) + # upsampling a complex array + y1_true = np.array([1.+0.j, 0.5+0.j, 0.+0.j, 0.5+0.j]) + assert_allclose(y1_test, y1_true, atol=1e-12) + x2 = np.array([1., 0.5, 0., 0.5]) + y2_test = signal.resample(x2, 2) # downsampling a real array + y2_true = np.array([1., 0.]) + assert_allclose(y2_test, y2_true, atol=1e-12) + + def test_poly_vs_filtfilt(self): + # Check that up=1.0 gives same answer as filtfilt + slicing + random_state = np.random.RandomState(17) + try_types = (int, np.float32, np.complex64, float, complex) + size = 10000 + down_factors = [2, 11, 79] + + for dtype in try_types: + x = random_state.randn(size).astype(dtype) + if dtype in (np.complex64, np.complex128): + x += 1j * random_state.randn(size) + + # resample_poly assumes zeros outside of signl, whereas filtfilt + # can only constant-pad. Make them equivalent: + x[0] = 0 + x[-1] = 0 + + for down in down_factors: + h = signal.firwin(31, 1. / down, window='hamming') + yf = filtfilt(h, 1.0, x, padtype='constant')[::down] + + # Need to pass convolved version of filter to resample_poly, + # since filtfilt does forward and backward, but resample_poly + # only goes forward + hc = convolve(h, h[::-1]) + y = signal.resample_poly(x, 1, down, window=hc) + assert_allclose(yf, y, atol=1e-7, rtol=1e-7) + + def test_correlate1d(self): + for down in [2, 4]: + for nx in range(1, 40, down): + for nweights in (32, 33): + x = np.random.random((nx,)) + weights = np.random.random((nweights,)) + y_g = correlate1d(x, weights[::-1], mode='constant') + y_s = signal.resample_poly( + x, up=1, down=down, window=weights) + assert_allclose(y_g[::down], y_s) + + @pytest.mark.parametrize('dtype', [np.int32, np.float32]) + def test_gh_15620(self, dtype): + data = np.array([0, 1, 2, 3, 2, 1, 0], dtype=dtype) + actual = signal.resample_poly(data, + up=2, + down=1, + padtype='smooth') + assert np.count_nonzero(actual) > 0 + + +class TestCSpline1DEval: + + def test_basic(self): + y = array([1, 2, 3, 4, 3, 2, 1, 2, 3.0]) + x = arange(len(y)) + dx = x[1] - x[0] + cj = signal.cspline1d(y) + + x2 = arange(len(y) * 10.0) / 10.0 + y2 = signal.cspline1d_eval(cj, x2, dx=dx, x0=x[0]) + + # make sure interpolated values are on knot points + assert_array_almost_equal(y2[::10], y, decimal=5) + + def test_complex(self): + # create some smoothly varying complex signal to interpolate + x = np.arange(2) + y = np.zeros(x.shape, dtype=np.complex64) + T = 10.0 + f = 1.0 / T + y = np.exp(2.0J * np.pi * f * x) + + # get the cspline transform + cy = signal.cspline1d(y) + + # determine new test x value and interpolate + xnew = np.array([0.5]) + ynew = signal.cspline1d_eval(cy, xnew) + + assert_equal(ynew.dtype, y.dtype) + +class TestOrderFilt: + + def test_basic(self): + assert_array_equal(signal.order_filter([1, 2, 3], [1, 0, 1], 1), + [2, 3, 2]) + + +class _TestLinearFilter: + + def generate(self, shape): + x = np.linspace(0, np.prod(shape) - 1, np.prod(shape)).reshape(shape) + return self.convert_dtype(x) + + def convert_dtype(self, arr): + if self.dtype == np.dtype('O'): + arr = np.asarray(arr) + out = np.empty(arr.shape, self.dtype) + iter = np.nditer([arr, out], ['refs_ok','zerosize_ok'], + [['readonly'],['writeonly']]) + for x, y in iter: + y[...] = self.type(x[()]) + return out + else: + return np.asarray(arr, dtype=self.dtype) + + def test_rank_1_IIR(self): + x = self.generate((6,)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, -0.5]) + y_r = self.convert_dtype([0, 2, 4, 6, 8, 10.]) + assert_array_almost_equal(lfilter(b, a, x), y_r) + + def test_rank_1_FIR(self): + x = self.generate((6,)) + b = self.convert_dtype([1, 1]) + a = self.convert_dtype([1]) + y_r = self.convert_dtype([0, 1, 3, 5, 7, 9.]) + assert_array_almost_equal(lfilter(b, a, x), y_r) + + def test_rank_1_IIR_init_cond(self): + x = self.generate((6,)) + b = self.convert_dtype([1, 0, -1]) + a = self.convert_dtype([0.5, -0.5]) + zi = self.convert_dtype([1, 2]) + y_r = self.convert_dtype([1, 5, 9, 13, 17, 21]) + zf_r = self.convert_dtype([13, -10]) + y, zf = lfilter(b, a, x, zi=zi) + assert_array_almost_equal(y, y_r) + assert_array_almost_equal(zf, zf_r) + + def test_rank_1_FIR_init_cond(self): + x = self.generate((6,)) + b = self.convert_dtype([1, 1, 1]) + a = self.convert_dtype([1]) + zi = self.convert_dtype([1, 1]) + y_r = self.convert_dtype([1, 2, 3, 6, 9, 12.]) + zf_r = self.convert_dtype([9, 5]) + y, zf = lfilter(b, a, x, zi=zi) + assert_array_almost_equal(y, y_r) + assert_array_almost_equal(zf, zf_r) + + def test_rank_2_IIR_axis_0(self): + x = self.generate((4, 3)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, 0.5]) + y_r2_a0 = self.convert_dtype([[0, 2, 4], [6, 4, 2], [0, 2, 4], + [6, 4, 2]]) + y = lfilter(b, a, x, axis=0) + assert_array_almost_equal(y_r2_a0, y) + + def test_rank_2_IIR_axis_1(self): + x = self.generate((4, 3)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, 0.5]) + y_r2_a1 = self.convert_dtype([[0, 2, 0], [6, -4, 6], [12, -10, 12], + [18, -16, 18]]) + y = lfilter(b, a, x, axis=1) + assert_array_almost_equal(y_r2_a1, y) + + def test_rank_2_IIR_axis_0_init_cond(self): + x = self.generate((4, 3)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, 0.5]) + zi = self.convert_dtype(np.ones((4,1))) + + y_r2_a0_1 = self.convert_dtype([[1, 1, 1], [7, -5, 7], [13, -11, 13], + [19, -17, 19]]) + zf_r = self.convert_dtype([-5, -17, -29, -41])[:, np.newaxis] + y, zf = lfilter(b, a, x, axis=1, zi=zi) + assert_array_almost_equal(y_r2_a0_1, y) + assert_array_almost_equal(zf, zf_r) + + def test_rank_2_IIR_axis_1_init_cond(self): + x = self.generate((4,3)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, 0.5]) + zi = self.convert_dtype(np.ones((1,3))) + + y_r2_a0_0 = self.convert_dtype([[1, 3, 5], [5, 3, 1], + [1, 3, 5], [5, 3, 1]]) + zf_r = self.convert_dtype([[-23, -23, -23]]) + y, zf = lfilter(b, a, x, axis=0, zi=zi) + assert_array_almost_equal(y_r2_a0_0, y) + assert_array_almost_equal(zf, zf_r) + + def test_rank_3_IIR(self): + x = self.generate((4, 3, 2)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, 0.5]) + + for axis in range(x.ndim): + y = lfilter(b, a, x, axis) + y_r = np.apply_along_axis(lambda w: lfilter(b, a, w), axis, x) + assert_array_almost_equal(y, y_r) + + def test_rank_3_IIR_init_cond(self): + x = self.generate((4, 3, 2)) + b = self.convert_dtype([1, -1]) + a = self.convert_dtype([0.5, 0.5]) + + for axis in range(x.ndim): + zi_shape = list(x.shape) + zi_shape[axis] = 1 + zi = self.convert_dtype(np.ones(zi_shape)) + zi1 = self.convert_dtype([1]) + y, zf = lfilter(b, a, x, axis, zi) + def lf0(w): + return lfilter(b, a, w, zi=zi1)[0] + def lf1(w): + return lfilter(b, a, w, zi=zi1)[1] + y_r = np.apply_along_axis(lf0, axis, x) + zf_r = np.apply_along_axis(lf1, axis, x) + assert_array_almost_equal(y, y_r) + assert_array_almost_equal(zf, zf_r) + + def test_rank_3_FIR(self): + x = self.generate((4, 3, 2)) + b = self.convert_dtype([1, 0, -1]) + a = self.convert_dtype([1]) + + for axis in range(x.ndim): + y = lfilter(b, a, x, axis) + y_r = np.apply_along_axis(lambda w: lfilter(b, a, w), axis, x) + assert_array_almost_equal(y, y_r) + + def test_rank_3_FIR_init_cond(self): + x = self.generate((4, 3, 2)) + b = self.convert_dtype([1, 0, -1]) + a = self.convert_dtype([1]) + + for axis in range(x.ndim): + zi_shape = list(x.shape) + zi_shape[axis] = 2 + zi = self.convert_dtype(np.ones(zi_shape)) + zi1 = self.convert_dtype([1, 1]) + y, zf = lfilter(b, a, x, axis, zi) + def lf0(w): + return lfilter(b, a, w, zi=zi1)[0] + def lf1(w): + return lfilter(b, a, w, zi=zi1)[1] + y_r = np.apply_along_axis(lf0, axis, x) + zf_r = np.apply_along_axis(lf1, axis, x) + assert_array_almost_equal(y, y_r) + assert_array_almost_equal(zf, zf_r) + + def test_zi_pseudobroadcast(self): + x = self.generate((4, 5, 20)) + b,a = signal.butter(8, 0.2, output='ba') + b = self.convert_dtype(b) + a = self.convert_dtype(a) + zi_size = b.shape[0] - 1 + + # lfilter requires x.ndim == zi.ndim exactly. However, zi can have + # length 1 dimensions. + zi_full = self.convert_dtype(np.ones((4, 5, zi_size))) + zi_sing = self.convert_dtype(np.ones((1, 1, zi_size))) + + y_full, zf_full = lfilter(b, a, x, zi=zi_full) + y_sing, zf_sing = lfilter(b, a, x, zi=zi_sing) + + assert_array_almost_equal(y_sing, y_full) + assert_array_almost_equal(zf_full, zf_sing) + + # lfilter does not prepend ones + assert_raises(ValueError, lfilter, b, a, x, -1, np.ones(zi_size)) + + def test_scalar_a(self): + # a can be a scalar. + x = self.generate(6) + b = self.convert_dtype([1, 0, -1]) + a = self.convert_dtype([1]) + y_r = self.convert_dtype([0, 1, 2, 2, 2, 2]) + + y = lfilter(b, a[0], x) + assert_array_almost_equal(y, y_r) + + def test_zi_some_singleton_dims(self): + # lfilter doesn't really broadcast (no prepending of 1's). But does + # do singleton expansion if x and zi have the same ndim. This was + # broken only if a subset of the axes were singletons (gh-4681). + x = self.convert_dtype(np.zeros((3,2,5), 'l')) + b = self.convert_dtype(np.ones(5, 'l')) + a = self.convert_dtype(np.array([1,0,0])) + zi = np.ones((3,1,4), 'l') + zi[1,:,:] *= 2 + zi[2,:,:] *= 3 + zi = self.convert_dtype(zi) + + zf_expected = self.convert_dtype(np.zeros((3,2,4), 'l')) + y_expected = np.zeros((3,2,5), 'l') + y_expected[:,:,:4] = [[[1]], [[2]], [[3]]] + y_expected = self.convert_dtype(y_expected) + + # IIR + y_iir, zf_iir = lfilter(b, a, x, -1, zi) + assert_array_almost_equal(y_iir, y_expected) + assert_array_almost_equal(zf_iir, zf_expected) + + # FIR + y_fir, zf_fir = lfilter(b, a[0], x, -1, zi) + assert_array_almost_equal(y_fir, y_expected) + assert_array_almost_equal(zf_fir, zf_expected) + + def base_bad_size_zi(self, b, a, x, axis, zi): + b = self.convert_dtype(b) + a = self.convert_dtype(a) + x = self.convert_dtype(x) + zi = self.convert_dtype(zi) + assert_raises(ValueError, lfilter, b, a, x, axis, zi) + + def test_bad_size_zi(self): + # rank 1 + x1 = np.arange(6) + self.base_bad_size_zi([1], [1], x1, -1, [1]) + self.base_bad_size_zi([1, 1], [1], x1, -1, [0, 1]) + self.base_bad_size_zi([1, 1], [1], x1, -1, [[0]]) + self.base_bad_size_zi([1, 1], [1], x1, -1, [0, 1, 2]) + self.base_bad_size_zi([1, 1, 1], [1], x1, -1, [[0]]) + self.base_bad_size_zi([1, 1, 1], [1], x1, -1, [0, 1, 2]) + self.base_bad_size_zi([1], [1, 1], x1, -1, [0, 1]) + self.base_bad_size_zi([1], [1, 1], x1, -1, [[0]]) + self.base_bad_size_zi([1], [1, 1], x1, -1, [0, 1, 2]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x1, -1, [0]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x1, -1, [[0], [1]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x1, -1, [0, 1, 2]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x1, -1, [0, 1, 2, 3]) + self.base_bad_size_zi([1, 1], [1, 1, 1], x1, -1, [0]) + self.base_bad_size_zi([1, 1], [1, 1, 1], x1, -1, [[0], [1]]) + self.base_bad_size_zi([1, 1], [1, 1, 1], x1, -1, [0, 1, 2]) + self.base_bad_size_zi([1, 1], [1, 1, 1], x1, -1, [0, 1, 2, 3]) + + # rank 2 + x2 = np.arange(12).reshape((4,3)) + # for axis=0 zi.shape should == (max(len(a),len(b))-1, 3) + self.base_bad_size_zi([1], [1], x2, 0, [0]) + + # for each of these there are 5 cases tested (in this order): + # 1. not deep enough, right # elements + # 2. too deep, right # elements + # 3. right depth, right # elements, transposed + # 4. right depth, too few elements + # 5. right depth, too many elements + + self.base_bad_size_zi([1, 1], [1], x2, 0, [0,1,2]) + self.base_bad_size_zi([1, 1], [1], x2, 0, [[[0,1,2]]]) + self.base_bad_size_zi([1, 1], [1], x2, 0, [[0], [1], [2]]) + self.base_bad_size_zi([1, 1], [1], x2, 0, [[0,1]]) + self.base_bad_size_zi([1, 1], [1], x2, 0, [[0,1,2,3]]) + + self.base_bad_size_zi([1, 1, 1], [1], x2, 0, [0,1,2,3,4,5]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 0, [[[0,1,2],[3,4,5]]]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 0, [[0,1],[2,3],[4,5]]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 0, [[0,1],[2,3]]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 0, [[0,1,2,3],[4,5,6,7]]) + + self.base_bad_size_zi([1], [1, 1], x2, 0, [0,1,2]) + self.base_bad_size_zi([1], [1, 1], x2, 0, [[[0,1,2]]]) + self.base_bad_size_zi([1], [1, 1], x2, 0, [[0], [1], [2]]) + self.base_bad_size_zi([1], [1, 1], x2, 0, [[0,1]]) + self.base_bad_size_zi([1], [1, 1], x2, 0, [[0,1,2,3]]) + + self.base_bad_size_zi([1], [1, 1, 1], x2, 0, [0,1,2,3,4,5]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 0, [[[0,1,2],[3,4,5]]]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 0, [[0,1],[2,3],[4,5]]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 0, [[0,1],[2,3]]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 0, [[0,1,2,3],[4,5,6,7]]) + + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 0, [0,1,2,3,4,5]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 0, [[[0,1,2],[3,4,5]]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 0, [[0,1],[2,3],[4,5]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 0, [[0,1],[2,3]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 0, [[0,1,2,3],[4,5,6,7]]) + + # for axis=1 zi.shape should == (4, max(len(a),len(b))-1) + self.base_bad_size_zi([1], [1], x2, 1, [0]) + + self.base_bad_size_zi([1, 1], [1], x2, 1, [0,1,2,3]) + self.base_bad_size_zi([1, 1], [1], x2, 1, [[[0],[1],[2],[3]]]) + self.base_bad_size_zi([1, 1], [1], x2, 1, [[0, 1, 2, 3]]) + self.base_bad_size_zi([1, 1], [1], x2, 1, [[0],[1],[2]]) + self.base_bad_size_zi([1, 1], [1], x2, 1, [[0],[1],[2],[3],[4]]) + + self.base_bad_size_zi([1, 1, 1], [1], x2, 1, [0,1,2,3,4,5,6,7]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 1, [[[0,1],[2,3],[4,5],[6,7]]]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 1, [[0,1,2,3],[4,5,6,7]]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 1, [[0,1],[2,3],[4,5]]) + self.base_bad_size_zi([1, 1, 1], [1], x2, 1, [[0,1],[2,3],[4,5],[6,7],[8,9]]) + + self.base_bad_size_zi([1], [1, 1], x2, 1, [0,1,2,3]) + self.base_bad_size_zi([1], [1, 1], x2, 1, [[[0],[1],[2],[3]]]) + self.base_bad_size_zi([1], [1, 1], x2, 1, [[0, 1, 2, 3]]) + self.base_bad_size_zi([1], [1, 1], x2, 1, [[0],[1],[2]]) + self.base_bad_size_zi([1], [1, 1], x2, 1, [[0],[1],[2],[3],[4]]) + + self.base_bad_size_zi([1], [1, 1, 1], x2, 1, [0,1,2,3,4,5,6,7]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 1, [[[0,1],[2,3],[4,5],[6,7]]]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 1, [[0,1,2,3],[4,5,6,7]]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 1, [[0,1],[2,3],[4,5]]) + self.base_bad_size_zi([1], [1, 1, 1], x2, 1, [[0,1],[2,3],[4,5],[6,7],[8,9]]) + + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 1, [0,1,2,3,4,5,6,7]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 1, [[[0,1],[2,3],[4,5],[6,7]]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 1, [[0,1,2,3],[4,5,6,7]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 1, [[0,1],[2,3],[4,5]]) + self.base_bad_size_zi([1, 1, 1], [1, 1], x2, 1, [[0,1],[2,3],[4,5],[6,7],[8,9]]) + + def test_empty_zi(self): + # Regression test for #880: empty array for zi crashes. + x = self.generate((5,)) + a = self.convert_dtype([1]) + b = self.convert_dtype([1]) + zi = self.convert_dtype([]) + y, zf = lfilter(b, a, x, zi=zi) + assert_array_almost_equal(y, x) + assert_equal(zf.dtype, self.dtype) + assert_equal(zf.size, 0) + + def test_lfiltic_bad_zi(self): + # Regression test for #3699: bad initial conditions + a = self.convert_dtype([1]) + b = self.convert_dtype([1]) + # "y" sets the datatype of zi, so it truncates if int + zi = lfiltic(b, a, [1., 0]) + zi_1 = lfiltic(b, a, [1, 0]) + zi_2 = lfiltic(b, a, [True, False]) + assert_array_equal(zi, zi_1) + assert_array_equal(zi, zi_2) + + def test_short_x_FIR(self): + # regression test for #5116 + # x shorter than b, with non None zi fails + a = self.convert_dtype([1]) + b = self.convert_dtype([1, 0, -1]) + zi = self.convert_dtype([2, 7]) + x = self.convert_dtype([72]) + ye = self.convert_dtype([74]) + zfe = self.convert_dtype([7, -72]) + y, zf = lfilter(b, a, x, zi=zi) + assert_array_almost_equal(y, ye) + assert_array_almost_equal(zf, zfe) + + def test_short_x_IIR(self): + # regression test for #5116 + # x shorter than b, with non None zi fails + a = self.convert_dtype([1, 1]) + b = self.convert_dtype([1, 0, -1]) + zi = self.convert_dtype([2, 7]) + x = self.convert_dtype([72]) + ye = self.convert_dtype([74]) + zfe = self.convert_dtype([-67, -72]) + y, zf = lfilter(b, a, x, zi=zi) + assert_array_almost_equal(y, ye) + assert_array_almost_equal(zf, zfe) + + def test_do_not_modify_a_b_IIR(self): + x = self.generate((6,)) + b = self.convert_dtype([1, -1]) + b0 = b.copy() + a = self.convert_dtype([0.5, -0.5]) + a0 = a.copy() + y_r = self.convert_dtype([0, 2, 4, 6, 8, 10.]) + y_f = lfilter(b, a, x) + assert_array_almost_equal(y_f, y_r) + assert_equal(b, b0) + assert_equal(a, a0) + + def test_do_not_modify_a_b_FIR(self): + x = self.generate((6,)) + b = self.convert_dtype([1, 0, 1]) + b0 = b.copy() + a = self.convert_dtype([2]) + a0 = a.copy() + y_r = self.convert_dtype([0, 0.5, 1, 2, 3, 4.]) + y_f = lfilter(b, a, x) + assert_array_almost_equal(y_f, y_r) + assert_equal(b, b0) + assert_equal(a, a0) + + @pytest.mark.parametrize("a", [1.0, [1.0], np.array(1.0)]) + @pytest.mark.parametrize("b", [1.0, [1.0], np.array(1.0)]) + def test_scalar_input(self, a, b): + data = np.random.randn(10) + assert_allclose( + lfilter(np.array([1.0]), np.array([1.0]), data), + lfilter(b, a, data)) + + @pytest.mark.thread_unsafe + def test_dtype_deprecation(self): + # gh-21211 + a = np.asarray([1, 2, 3, 6, 5, 3], dtype=object) + b = np.asarray([2, 3, 4, 5, 3, 4, 2, 2, 1], dtype=object) + with pytest.deprecated_call(match="dtype=object is not supported"): + lfilter(a, b, [1, 2, 3, 4]) + + +class TestLinearFilterFloat32(_TestLinearFilter): + dtype = np.dtype('f') + + +class TestLinearFilterFloat64(_TestLinearFilter): + dtype = np.dtype('d') + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +class TestLinearFilterFloatExtended(_TestLinearFilter): + dtype = np.dtype('g') + + +class TestLinearFilterComplex64(_TestLinearFilter): + dtype = np.dtype('F') + + +class TestLinearFilterComplex128(_TestLinearFilter): + dtype = np.dtype('D') + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +class TestLinearFilterComplexExtended(_TestLinearFilter): + dtype = np.dtype('G') + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +class TestLinearFilterDecimal(_TestLinearFilter): + dtype = np.dtype('O') + + def type(self, x): + return Decimal(str(x)) + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +class TestLinearFilterObject(_TestLinearFilter): + dtype = np.dtype('O') + type = float + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +def test_lfilter_bad_object(): + # lfilter: object arrays with non-numeric objects raise TypeError. + # Regression test for ticket #1452. + if hasattr(sys, 'abiflags') and 'd' in sys.abiflags: + pytest.skip('test is flaky when run with python3-dbg') + assert_raises(TypeError, lfilter, [1.0], [1.0], [1.0, None, 2.0]) + assert_raises(TypeError, lfilter, [1.0], [None], [1.0, 2.0, 3.0]) + assert_raises(TypeError, lfilter, [None], [1.0], [1.0, 2.0, 3.0]) + + +_pmf = pytest.mark.filterwarnings('ignore::DeprecationWarning') + +def test_lfilter_notimplemented_input(): + # Should not crash, gh-7991 + assert_raises(NotImplementedError, lfilter, [2,3], [4,5], [1,2,3,4,5]) + + +@pytest.mark.parametrize('dt', [np.ubyte, np.byte, np.ushort, np.short, + np_ulong, np_long, np.ulonglong, np.ulonglong, + np.float32, np.float64, + pytest.param(np.longdouble, marks=_pmf), + pytest.param(Decimal, marks=_pmf)] +) +class TestCorrelateReal: + def _setup_rank1(self, dt): + a = np.linspace(0, 3, 4).astype(dt) + b = np.linspace(1, 2, 2).astype(dt) + + y_r = np.array([0, 2, 5, 8, 3]).astype(dt) + return a, b, y_r + + def equal_tolerance(self, res_dt): + # default value of keyword + decimal = 6 + try: + dt_info = np.finfo(res_dt) + if hasattr(dt_info, 'resolution'): + decimal = int(-0.5*np.log10(dt_info.resolution)) + except Exception: + pass + return decimal + + def equal_tolerance_fft(self, res_dt): + # FFT implementations convert longdouble arguments down to + # double so don't expect better precision, see gh-9520 + if res_dt == np.longdouble: + return self.equal_tolerance(np.float64) + else: + return self.equal_tolerance(res_dt) + + def test_method(self, dt): + if dt == Decimal: + method = choose_conv_method([Decimal(4)], [Decimal(3)]) + assert_equal(method, 'direct') + else: + a, b, y_r = self._setup_rank3(dt) + y_fft = correlate(a, b, method='fft') + y_direct = correlate(a, b, method='direct') + + assert_array_almost_equal(y_r, + y_fft, + decimal=self.equal_tolerance_fft(y_fft.dtype),) + assert_array_almost_equal(y_r, + y_direct, + decimal=self.equal_tolerance(y_direct.dtype),) + assert_equal(y_fft.dtype, dt) + assert_equal(y_direct.dtype, dt) + + def test_rank1_valid(self, dt): + a, b, y_r = self._setup_rank1(dt) + y = correlate(a, b, 'valid') + assert_array_almost_equal(y, y_r[1:4]) + assert_equal(y.dtype, dt) + + # See gh-5897 + y = correlate(b, a, 'valid') + assert_array_almost_equal(y, y_r[1:4][::-1]) + assert_equal(y.dtype, dt) + + def test_rank1_same(self, dt): + a, b, y_r = self._setup_rank1(dt) + y = correlate(a, b, 'same') + assert_array_almost_equal(y, y_r[:-1]) + assert_equal(y.dtype, dt) + + def test_rank1_full(self, dt): + a, b, y_r = self._setup_rank1(dt) + y = correlate(a, b, 'full') + assert_array_almost_equal(y, y_r) + assert_equal(y.dtype, dt) + + def _setup_rank3(self, dt): + a = np.linspace(0, 39, 40).reshape((2, 4, 5), order='F').astype( + dt) + b = np.linspace(0, 23, 24).reshape((2, 3, 4), order='F').astype( + dt) + + y_r = array([[[0., 184., 504., 912., 1360., 888., 472., 160.], + [46., 432., 1062., 1840., 2672., 1698., 864., 266.], + [134., 736., 1662., 2768., 3920., 2418., 1168., 314.], + [260., 952., 1932., 3056., 4208., 2580., 1240., 332.], + [202., 664., 1290., 1984., 2688., 1590., 712., 150.], + [114., 344., 642., 960., 1280., 726., 296., 38.]], + + [[23., 400., 1035., 1832., 2696., 1737., 904., 293.], + [134., 920., 2166., 3680., 5280., 3306., 1640., 474.], + [325., 1544., 3369., 5512., 7720., 4683., 2192., 535.], + [571., 1964., 3891., 6064., 8272., 4989., 2324., 565.], + [434., 1360., 2586., 3920., 5264., 3054., 1312., 230.], + [241., 700., 1281., 1888., 2496., 1383., 532., 39.]], + + [[22., 214., 528., 916., 1332., 846., 430., 132.], + [86., 484., 1098., 1832., 2600., 1602., 772., 206.], + [188., 802., 1698., 2732., 3788., 2256., 1018., 218.], + [308., 1006., 1950., 2996., 4052., 2400., 1078., 230.], + [230., 692., 1290., 1928., 2568., 1458., 596., 78.], + [126., 354., 636., 924., 1212., 654., 234., 0.]]], + dtype=np.float64).astype(dt) + + return a, b, y_r + + def test_rank3_valid(self, dt): + a, b, y_r = self._setup_rank3(dt) + y = correlate(a, b, "valid") + assert_array_almost_equal(y, y_r[1:2, 2:4, 3:5]) + assert_equal(y.dtype, dt) + + # See gh-5897 + y = correlate(b, a, "valid") + assert_array_almost_equal(y, y_r[1:2, 2:4, 3:5][::-1, ::-1, ::-1]) + assert_equal(y.dtype, dt) + + def test_rank3_same(self, dt): + a, b, y_r = self._setup_rank3(dt) + y = correlate(a, b, "same") + assert_array_almost_equal(y, y_r[0:-1, 1:-1, 1:-2]) + assert_equal(y.dtype, dt) + + def test_rank3_all(self, dt): + a, b, y_r = self._setup_rank3(dt) + y = correlate(a, b) + assert_array_almost_equal(y, y_r) + assert_equal(y.dtype, dt) + + +class TestCorrelate: + # Tests that don't depend on dtype + + def test_invalid_shapes(self): + # By "invalid," we mean that no one + # array has dimensions that are all at + # least as large as the corresponding + # dimensions of the other array. This + # setup should throw a ValueError. + a = np.arange(1, 7).reshape((2, 3)) + b = np.arange(-6, 0).reshape((3, 2)) + + assert_raises(ValueError, correlate, *(a, b), **{'mode': 'valid'}) + assert_raises(ValueError, correlate, *(b, a), **{'mode': 'valid'}) + + def test_invalid_params(self): + a = [3, 4, 5] + b = [1, 2, 3] + assert_raises(ValueError, correlate, a, b, mode='spam') + assert_raises(ValueError, correlate, a, b, mode='eggs', method='fft') + assert_raises(ValueError, correlate, a, b, mode='ham', method='direct') + assert_raises(ValueError, correlate, a, b, mode='full', method='bacon') + assert_raises(ValueError, correlate, a, b, mode='same', method='bacon') + + def test_mismatched_dims(self): + # Input arrays should have the same number of dimensions + assert_raises(ValueError, correlate, [1], 2, method='direct') + assert_raises(ValueError, correlate, 1, [2], method='direct') + assert_raises(ValueError, correlate, [1], 2, method='fft') + assert_raises(ValueError, correlate, 1, [2], method='fft') + assert_raises(ValueError, correlate, [1], [[2]]) + assert_raises(ValueError, correlate, [3], 2) + + def test_numpy_fastpath(self): + a = [1, 2, 3] + b = [4, 5] + assert_allclose(correlate(a, b, mode='same'), [5, 14, 23]) + + a = [1, 2, 3] + b = [4, 5, 6] + assert_allclose(correlate(a, b, mode='same'), [17, 32, 23]) + assert_allclose(correlate(a, b, mode='full'), [6, 17, 32, 23, 12]) + assert_allclose(correlate(a, b, mode='valid'), [32]) + + @pytest.mark.thread_unsafe + def test_dtype_deprecation(self): + # gh-21211 + a = np.asarray([1, 2, 3, 6, 5, 3], dtype=object) + b = np.asarray([2, 3, 4, 5, 3, 4, 2, 2, 1], dtype=object) + with pytest.deprecated_call(match="dtype=object is not supported"): + correlate(a, b) + + +@pytest.mark.parametrize("mode", ["valid", "same", "full"]) +@pytest.mark.parametrize("behind", [True, False]) +@pytest.mark.parametrize("input_size", [100, 101, 1000, 1001, 10000, 10001]) +def test_correlation_lags(mode, behind, input_size): + # generate random data + rng = np.random.RandomState(0) + in1 = rng.standard_normal(input_size) + offset = int(input_size/10) + # generate offset version of array to correlate with + if behind: + # y is behind x + in2 = np.concatenate([rng.standard_normal(offset), in1]) + expected = -offset + else: + # y is ahead of x + in2 = in1[offset:] + expected = offset + # cross correlate, returning lag information + correlation = correlate(in1, in2, mode=mode) + lags = correlation_lags(in1.size, in2.size, mode=mode) + # identify the peak + lag_index = np.argmax(correlation) + # Check as expected + assert_equal(lags[lag_index], expected) + # Correlation and lags shape should match + assert_equal(lags.shape, correlation.shape) + + +def test_correlation_lags_invalid_mode(): + with pytest.raises(ValueError, match="Mode asdfgh is invalid"): + correlation_lags(100, 100, mode="asdfgh") + + +@pytest.mark.parametrize('dt', [np.csingle, np.cdouble, + pytest.param(np.clongdouble, marks=_pmf)]) +class TestCorrelateComplex: + # The decimal precision to be used for comparing results. + # This value will be passed as the 'decimal' keyword argument of + # assert_array_almost_equal(). + # Since correlate may chose to use FFT method which converts + # longdoubles to doubles internally don't expect better precision + # for longdouble than for double (see gh-9520). + + def decimal(self, dt): + if dt == np.clongdouble: + dt = np.cdouble + return int(2 * np.finfo(dt).precision / 3) + + def _setup_rank1(self, dt, mode): + np.random.seed(9) + a = np.random.randn(10).astype(dt) + a += 1j * np.random.randn(10).astype(dt) + b = np.random.randn(8).astype(dt) + b += 1j * np.random.randn(8).astype(dt) + + y_r = (correlate(a.real, b.real, mode=mode) + + correlate(a.imag, b.imag, mode=mode)).astype(dt) + y_r += 1j * (-correlate(a.real, b.imag, mode=mode) + + correlate(a.imag, b.real, mode=mode)) + return a, b, y_r + + def test_rank1_valid(self, dt): + a, b, y_r = self._setup_rank1(dt, 'valid') + y = correlate(a, b, 'valid') + assert_array_almost_equal(y, y_r, decimal=self.decimal(dt)) + assert_equal(y.dtype, dt) + + # See gh-5897 + y = correlate(b, a, 'valid') + assert_array_almost_equal(y, y_r[::-1].conj(), decimal=self.decimal(dt)) + assert_equal(y.dtype, dt) + + def test_rank1_same(self, dt): + a, b, y_r = self._setup_rank1(dt, 'same') + y = correlate(a, b, 'same') + assert_array_almost_equal(y, y_r, decimal=self.decimal(dt)) + assert_equal(y.dtype, dt) + + def test_rank1_full(self, dt): + a, b, y_r = self._setup_rank1(dt, 'full') + y = correlate(a, b, 'full') + assert_array_almost_equal(y, y_r, decimal=self.decimal(dt)) + assert_equal(y.dtype, dt) + + def test_swap_full(self, dt): + d = np.array([0.+0.j, 1.+1.j, 2.+2.j], dtype=dt) + k = np.array([1.+3.j, 2.+4.j, 3.+5.j, 4.+6.j], dtype=dt) + y = correlate(d, k) + assert_equal(y, [0.+0.j, 10.-2.j, 28.-6.j, 22.-6.j, 16.-6.j, 8.-4.j]) + + def test_swap_same(self, dt): + d = [0.+0.j, 1.+1.j, 2.+2.j] + k = [1.+3.j, 2.+4.j, 3.+5.j, 4.+6.j] + y = correlate(d, k, mode="same") + assert_equal(y, [10.-2.j, 28.-6.j, 22.-6.j]) + + def test_rank3(self, dt): + a = np.random.randn(10, 8, 6).astype(dt) + a += 1j * np.random.randn(10, 8, 6).astype(dt) + b = np.random.randn(8, 6, 4).astype(dt) + b += 1j * np.random.randn(8, 6, 4).astype(dt) + + y_r = (correlate(a.real, b.real) + + correlate(a.imag, b.imag)).astype(dt) + y_r += 1j * (-correlate(a.real, b.imag) + correlate(a.imag, b.real)) + + y = correlate(a, b, 'full') + assert_array_almost_equal(y, y_r, decimal=self.decimal(dt) - 1) + assert_equal(y.dtype, dt) + + def test_rank0(self, dt): + a = np.array(np.random.randn()).astype(dt) + a += 1j * np.array(np.random.randn()).astype(dt) + b = np.array(np.random.randn()).astype(dt) + b += 1j * np.array(np.random.randn()).astype(dt) + + y_r = (correlate(a.real, b.real) + + correlate(a.imag, b.imag)).astype(dt) + y_r += 1j * np.array(-correlate(a.real, b.imag) + + correlate(a.imag, b.real)) + + y = correlate(a, b, 'full') + assert_array_almost_equal(y, y_r, decimal=self.decimal(dt) - 1) + assert_equal(y.dtype, dt) + + assert_equal(correlate([1], [2j]), correlate(1, 2j)) + assert_equal(correlate([2j], [3j]), correlate(2j, 3j)) + assert_equal(correlate([3j], [4]), correlate(3j, 4)) + + +class TestCorrelate2d: + + def test_consistency_correlate_funcs(self): + # Compare np.correlate, signal.correlate, signal.correlate2d + a = np.arange(5) + b = np.array([3.2, 1.4, 3]) + for mode in ['full', 'valid', 'same']: + assert_almost_equal(np.correlate(a, b, mode=mode), + signal.correlate(a, b, mode=mode)) + assert_almost_equal(np.squeeze(signal.correlate2d([a], [b], + mode=mode)), + signal.correlate(a, b, mode=mode)) + + # See gh-5897 + if mode == 'valid': + assert_almost_equal(np.correlate(b, a, mode=mode), + signal.correlate(b, a, mode=mode)) + assert_almost_equal(np.squeeze(signal.correlate2d([b], [a], + mode=mode)), + signal.correlate(b, a, mode=mode)) + + def test_invalid_shapes(self): + # By "invalid," we mean that no one + # array has dimensions that are all at + # least as large as the corresponding + # dimensions of the other array. This + # setup should throw a ValueError. + a = np.arange(1, 7).reshape((2, 3)) + b = np.arange(-6, 0).reshape((3, 2)) + + assert_raises(ValueError, signal.correlate2d, *(a, b), **{'mode': 'valid'}) + assert_raises(ValueError, signal.correlate2d, *(b, a), **{'mode': 'valid'}) + + def test_complex_input(self): + assert_equal(signal.correlate2d([[1]], [[2j]]), -2j) + assert_equal(signal.correlate2d([[2j]], [[3j]]), 6) + assert_equal(signal.correlate2d([[3j]], [[4]]), 12j) + + +class TestLFilterZI: + + def test_basic(self): + a = np.array([1.0, -1.0, 0.5]) + b = np.array([1.0, 0.0, 2.0]) + zi_expected = np.array([5.0, -1.0]) + zi = lfilter_zi(b, a) + assert_array_almost_equal(zi, zi_expected) + + def test_scale_invariance(self): + # Regression test. There was a bug in which b was not correctly + # rescaled when a[0] was nonzero. + b = np.array([2, 8, 5]) + a = np.array([1, 1, 8]) + zi1 = lfilter_zi(b, a) + zi2 = lfilter_zi(2*b, 2*a) + assert_allclose(zi2, zi1, rtol=1e-12) + + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_types(self, dtype): + b = np.zeros((8), dtype=dtype) + a = np.array([1], dtype=dtype) + assert_equal(np.real(signal.lfilter_zi(b, a)).dtype, dtype) + + +class TestFiltFilt: + filtfilt_kind = 'tf' + + def filtfilt(self, zpk, x, axis=-1, padtype='odd', padlen=None, + method='pad', irlen=None): + if self.filtfilt_kind == 'tf': + b, a = zpk2tf(*zpk) + return filtfilt(b, a, x, axis, padtype, padlen, method, irlen) + elif self.filtfilt_kind == 'sos': + sos = zpk2sos(*zpk) + return sosfiltfilt(sos, x, axis, padtype, padlen) + + def test_basic(self): + zpk = tf2zpk([1, 2, 3], [1, 2, 3]) + out = self.filtfilt(zpk, np.arange(12)) + assert_allclose(out, arange(12), atol=5.28e-11) + + def test_sine(self): + rate = 2000 + t = np.linspace(0, 1.0, rate + 1) + # A signal with low frequency and a high frequency. + xlow = np.sin(5 * 2 * np.pi * t) + xhigh = np.sin(250 * 2 * np.pi * t) + x = xlow + xhigh + + zpk = butter(8, 0.125, output='zpk') + # r is the magnitude of the largest pole. + r = np.abs(zpk[1]).max() + eps = 1e-5 + # n estimates the number of steps for the + # transient to decay by a factor of eps. + n = int(np.ceil(np.log(eps) / np.log(r))) + + # High order lowpass filter... + y = self.filtfilt(zpk, x, padlen=n) + # Result should be just xlow. + err = np.abs(y - xlow).max() + assert_(err < 1e-4) + + # A 2D case. + x2d = np.vstack([xlow, xlow + xhigh]) + y2d = self.filtfilt(zpk, x2d, padlen=n, axis=1) + assert_equal(y2d.shape, x2d.shape) + err = np.abs(y2d - xlow).max() + assert_(err < 1e-4) + + # Use the previous result to check the use of the axis keyword. + # (Regression test for ticket #1620) + y2dt = self.filtfilt(zpk, x2d.T, padlen=n, axis=0) + assert_equal(y2d, y2dt.T) + + def test_axis(self): + # Test the 'axis' keyword on a 3D array. + x = np.arange(10.0 * 11.0 * 12.0).reshape(10, 11, 12) + zpk = butter(3, 0.125, output='zpk') + y0 = self.filtfilt(zpk, x, padlen=0, axis=0) + y1 = self.filtfilt(zpk, np.swapaxes(x, 0, 1), padlen=0, axis=1) + assert_array_equal(y0, np.swapaxes(y1, 0, 1)) + y2 = self.filtfilt(zpk, np.swapaxes(x, 0, 2), padlen=0, axis=2) + assert_array_equal(y0, np.swapaxes(y2, 0, 2)) + + def test_acoeff(self): + if self.filtfilt_kind != 'tf': + return # only necessary for TF + # test for 'a' coefficient as single number + out = signal.filtfilt([.5, .5], 1, np.arange(10)) + assert_allclose(out, np.arange(10), rtol=1e-14, atol=1e-14) + + def test_gust_simple(self): + if self.filtfilt_kind != 'tf': + pytest.skip('gust only implemented for TF systems') + # The input array has length 2. The exact solution for this case + # was computed "by hand". + x = np.array([1.0, 2.0]) + b = np.array([0.5]) + a = np.array([1.0, -0.5]) + y, z1, z2 = _filtfilt_gust(b, a, x) + assert_allclose([z1[0], z2[0]], + [0.3*x[0] + 0.2*x[1], 0.2*x[0] + 0.3*x[1]]) + assert_allclose(y, [z1[0] + 0.25*z2[0] + 0.25*x[0] + 0.125*x[1], + 0.25*z1[0] + z2[0] + 0.125*x[0] + 0.25*x[1]]) + + def test_gust_scalars(self): + if self.filtfilt_kind != 'tf': + pytest.skip('gust only implemented for TF systems') + # The filter coefficients are both scalars, so the filter simply + # multiplies its input by b/a. When it is used in filtfilt, the + # factor is (b/a)**2. + x = np.arange(12) + b = 3.0 + a = 2.0 + y = filtfilt(b, a, x, method="gust") + expected = (b/a)**2 * x + assert_allclose(y, expected) + + +class TestSOSFiltFilt(TestFiltFilt): + filtfilt_kind = 'sos' + + def test_equivalence(self): + """Test equivalence between sosfiltfilt and filtfilt""" + x = np.random.RandomState(0).randn(1000) + for order in range(1, 6): + zpk = signal.butter(order, 0.35, output='zpk') + b, a = zpk2tf(*zpk) + sos = zpk2sos(*zpk) + y = filtfilt(b, a, x) + y_sos = sosfiltfilt(sos, x) + assert_allclose(y, y_sos, atol=1e-12, err_msg=f'order={order}') + + +def filtfilt_gust_opt(b, a, x): + """ + An alternative implementation of filtfilt with Gustafsson edges. + + This function computes the same result as + `scipy.signal._signaltools._filtfilt_gust`, but only 1-d arrays + are accepted. The problem is solved using `fmin` from `scipy.optimize`. + `_filtfilt_gust` is significantly faster than this implementation. + """ + def filtfilt_gust_opt_func(ics, b, a, x): + """Objective function used in filtfilt_gust_opt.""" + m = max(len(a), len(b)) - 1 + z0f = ics[:m] + z0b = ics[m:] + y_f = lfilter(b, a, x, zi=z0f)[0] + y_fb = lfilter(b, a, y_f[::-1], zi=z0b)[0][::-1] + + y_b = lfilter(b, a, x[::-1], zi=z0b)[0][::-1] + y_bf = lfilter(b, a, y_b, zi=z0f)[0] + value = np.sum((y_fb - y_bf)**2) + return value + + m = max(len(a), len(b)) - 1 + zi = lfilter_zi(b, a) + ics = np.concatenate((x[:m].mean()*zi, x[-m:].mean()*zi)) + result = fmin(filtfilt_gust_opt_func, ics, args=(b, a, x), + xtol=1e-10, ftol=1e-12, + maxfun=10000, maxiter=10000, + full_output=True, disp=False) + opt, fopt, niter, funcalls, warnflag = result + if warnflag > 0: + raise RuntimeError("minimization failed in filtfilt_gust_opt: " + "warnflag=%d" % warnflag) + z0f = opt[:m] + z0b = opt[m:] + + # Apply the forward-backward filter using the computed initial + # conditions. + y_b = lfilter(b, a, x[::-1], zi=z0b)[0][::-1] + y = lfilter(b, a, y_b, zi=z0f)[0] + + return y, z0f, z0b + + +def check_filtfilt_gust(b, a, shape, axis, irlen=None): + # Generate x, the data to be filtered. + np.random.seed(123) + x = np.random.randn(*shape) + + # Apply filtfilt to x. This is the main calculation to be checked. + y = filtfilt(b, a, x, axis=axis, method="gust", irlen=irlen) + + # Also call the private function so we can test the ICs. + yg, zg1, zg2 = _filtfilt_gust(b, a, x, axis=axis, irlen=irlen) + + # filtfilt_gust_opt is an independent implementation that gives the + # expected result, but it only handles 1-D arrays, so use some looping + # and reshaping shenanigans to create the expected output arrays. + xx = np.swapaxes(x, axis, -1) + out_shape = xx.shape[:-1] + yo = np.empty_like(xx) + m = max(len(a), len(b)) - 1 + zo1 = np.empty(out_shape + (m,)) + zo2 = np.empty(out_shape + (m,)) + for indx in product(*[range(d) for d in out_shape]): + yo[indx], zo1[indx], zo2[indx] = filtfilt_gust_opt(b, a, xx[indx]) + yo = np.swapaxes(yo, -1, axis) + zo1 = np.swapaxes(zo1, -1, axis) + zo2 = np.swapaxes(zo2, -1, axis) + + assert_allclose(y, yo, rtol=1e-8, atol=1e-9) + assert_allclose(yg, yo, rtol=1e-8, atol=1e-9) + assert_allclose(zg1, zo1, rtol=1e-8, atol=1e-9) + assert_allclose(zg2, zo2, rtol=1e-8, atol=1e-9) + + +@pytest.mark.fail_slow(10) +def test_choose_conv_method(): + for mode in ['valid', 'same', 'full']: + for ndim in [1, 2]: + n, k, true_method = 8, 6, 'direct' + x = np.random.randn(*((n,) * ndim)) + h = np.random.randn(*((k,) * ndim)) + + method = choose_conv_method(x, h, mode=mode) + assert_equal(method, true_method) + + method_try, times = choose_conv_method(x, h, mode=mode, measure=True) + assert_(method_try in {'fft', 'direct'}) + assert_(isinstance(times, dict)) + assert_('fft' in times.keys() and 'direct' in times.keys()) + + x = np.array([2**51], dtype=np.int64) + h = x.copy() + assert_equal(choose_conv_method(x, h, mode=mode), 'direct') + + +@pytest.mark.thread_unsafe +def test_choose_conv_dtype_deprecation(): + # gh-21211 + a = np.asarray([1, 2, 3, 6, 5, 3], dtype=object) + b = np.asarray([2, 3, 4, 5, 3, 4, 2, 2, 1], dtype=object) + with pytest.deprecated_call(match="dtype=object is not supported"): + choose_conv_method(a, b) + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +def test_choose_conv_method_2(): + for mode in ['valid', 'same', 'full']: + x = [Decimal(3), Decimal(2)] + h = [Decimal(1), Decimal(4)] + assert_equal(choose_conv_method(x, h, mode=mode), 'direct') + + n = 10 + for not_fft_conv_supp in ["complex256", "complex192"]: + if hasattr(np, not_fft_conv_supp): + x = np.ones(n, dtype=not_fft_conv_supp) + h = x.copy() + assert_equal(choose_conv_method(x, h, mode=mode), 'direct') + + +@pytest.mark.fail_slow(10) +def test_filtfilt_gust(): + # Design a filter. + z, p, k = signal.ellip(3, 0.01, 120, 0.0875, output='zpk') + + # Find the approximate impulse response length of the filter. + eps = 1e-10 + r = np.max(np.abs(p)) + approx_impulse_len = int(np.ceil(np.log(eps) / np.log(r))) + + np.random.seed(123) + + b, a = zpk2tf(z, p, k) + for irlen in [None, approx_impulse_len]: + signal_len = 5 * approx_impulse_len + + # 1-d test case + check_filtfilt_gust(b, a, (signal_len,), 0, irlen) + + # 3-d test case; test each axis. + for axis in range(3): + shape = [2, 2, 2] + shape[axis] = signal_len + check_filtfilt_gust(b, a, shape, axis, irlen) + + # Test case with length less than 2*approx_impulse_len. + # In this case, `filtfilt_gust` should behave the same as if + # `irlen=None` was given. + length = 2*approx_impulse_len - 50 + check_filtfilt_gust(b, a, (length,), 0, approx_impulse_len) + + +class TestDecimate: + def test_bad_args(self): + x = np.arange(12) + assert_raises(TypeError, signal.decimate, x, q=0.5, n=1) + assert_raises(TypeError, signal.decimate, x, q=2, n=0.5) + + def test_basic_IIR(self): + x = np.arange(12) + y = signal.decimate(x, 2, n=1, ftype='iir', zero_phase=False).round() + assert_array_equal(y, x[::2]) + + def test_basic_FIR(self): + x = np.arange(12) + y = signal.decimate(x, 2, n=1, ftype='fir', zero_phase=False).round() + assert_array_equal(y, x[::2]) + + def test_shape(self): + # Regression test for ticket #1480. + z = np.zeros((30, 30)) + d0 = signal.decimate(z, 2, axis=0, zero_phase=False) + assert_equal(d0.shape, (15, 30)) + d1 = signal.decimate(z, 2, axis=1, zero_phase=False) + assert_equal(d1.shape, (30, 15)) + + def test_phaseshift_FIR(self): + with suppress_warnings() as sup: + sup.filter(BadCoefficients, "Badly conditioned filter") + self._test_phaseshift(method='fir', zero_phase=False) + + def test_zero_phase_FIR(self): + with suppress_warnings() as sup: + sup.filter(BadCoefficients, "Badly conditioned filter") + self._test_phaseshift(method='fir', zero_phase=True) + + def test_phaseshift_IIR(self): + self._test_phaseshift(method='iir', zero_phase=False) + + def test_zero_phase_IIR(self): + self._test_phaseshift(method='iir', zero_phase=True) + + def _test_phaseshift(self, method, zero_phase): + rate = 120 + rates_to = [15, 20, 30, 40] # q = 8, 6, 4, 3 + + t_tot = 100 # Need to let antialiasing filters settle + t = np.arange(rate*t_tot+1) / float(rate) + + # Sinusoids at 0.8*nyquist, windowed to avoid edge artifacts + freqs = np.array(rates_to) * 0.8 / 2 + d = (np.exp(1j * 2 * np.pi * freqs[:, np.newaxis] * t) + * signal.windows.tukey(t.size, 0.1)) + + for rate_to in rates_to: + q = rate // rate_to + t_to = np.arange(rate_to*t_tot+1) / float(rate_to) + d_tos = (np.exp(1j * 2 * np.pi * freqs[:, np.newaxis] * t_to) + * signal.windows.tukey(t_to.size, 0.1)) + + # Set up downsampling filters, match v0.17 defaults + if method == 'fir': + n = 30 + system = signal.dlti(signal.firwin(n + 1, 1. / q, + window='hamming'), 1.) + elif method == 'iir': + n = 8 + wc = 0.8*np.pi/q + system = signal.dlti(*signal.cheby1(n, 0.05, wc/np.pi)) + + # Calculate expected phase response, as unit complex vector + if zero_phase is False: + _, h_resps = signal.freqz(system.num, system.den, + freqs/rate*2*np.pi) + h_resps /= np.abs(h_resps) + else: + h_resps = np.ones_like(freqs) + + y_resamps = signal.decimate(d.real, q, n, ftype=system, + zero_phase=zero_phase) + + # Get phase from complex inner product, like CSD + h_resamps = np.sum(d_tos.conj() * y_resamps, axis=-1) + h_resamps /= np.abs(h_resamps) + subnyq = freqs < 0.5*rate_to + + # Complex vectors should be aligned, only compare below nyquist + assert_allclose(np.angle(h_resps.conj()*h_resamps)[subnyq], 0, + atol=1e-3, rtol=1e-3) + + def test_auto_n(self): + # Test that our value of n is a reasonable choice (depends on + # the downsampling factor) + sfreq = 100. + n = 1000 + t = np.arange(n) / sfreq + # will alias for decimations (>= 15) + x = np.sqrt(2. / n) * np.sin(2 * np.pi * (sfreq / 30.) * t) + assert_allclose(np.linalg.norm(x), 1., rtol=1e-3) + x_out = signal.decimate(x, 30, ftype='fir') + assert_array_less(np.linalg.norm(x_out), 0.01) + + def test_long_float32(self): + # regression: gh-15072. With 32-bit float and either lfilter + # or filtfilt, this is numerically unstable + x = signal.decimate(np.ones(10_000, dtype=np.float32), 10) + assert not any(np.isnan(x)) + + def test_float16_upcast(self): + # float16 must be upcast to float64 + x = signal.decimate(np.ones(100, dtype=np.float16), 10) + assert x.dtype.type == np.float64 + + def test_complex_iir_dlti(self): + # regression: gh-17845 + # centre frequency for filter [Hz] + fcentre = 50 + # filter passband width [Hz] + fwidth = 5 + # sample rate [Hz] + fs = 1e3 + + z, p, k = signal.butter(2, 2*np.pi*fwidth/2, output='zpk', fs=fs) + z = z.astype(complex) * np.exp(2j * np.pi * fcentre/fs) + p = p.astype(complex) * np.exp(2j * np.pi * fcentre/fs) + system = signal.dlti(z, p, k) + + t = np.arange(200) / fs + + # input + u = (np.exp(2j * np.pi * fcentre * t) + + 0.5 * np.exp(-2j * np.pi * fcentre * t)) + + ynzp = signal.decimate(u, 2, ftype=system, zero_phase=False) + ynzpref = signal.lfilter(*signal.zpk2tf(z, p, k), + u)[::2] + + assert_equal(ynzp, ynzpref) + + yzp = signal.decimate(u, 2, ftype=system, zero_phase=True) + yzpref = signal.filtfilt(*signal.zpk2tf(z, p, k), + u)[::2] + + assert_allclose(yzp, yzpref, rtol=1e-10, atol=1e-13) + + def test_complex_fir_dlti(self): + # centre frequency for filter [Hz] + fcentre = 50 + # filter passband width [Hz] + fwidth = 5 + # sample rate [Hz] + fs = 1e3 + numtaps = 20 + + # FIR filter about 0Hz + bbase = signal.firwin(numtaps, fwidth/2, fs=fs) + + # rotate these to desired frequency + zbase = np.roots(bbase) + zrot = zbase * np.exp(2j * np.pi * fcentre/fs) + # FIR filter about 50Hz, maintaining passband gain of 0dB + bz = bbase[0] * np.poly(zrot) + + system = signal.dlti(bz, 1) + + t = np.arange(200) / fs + + # input + u = (np.exp(2j * np.pi * fcentre * t) + + 0.5 * np.exp(-2j * np.pi * fcentre * t)) + + ynzp = signal.decimate(u, 2, ftype=system, zero_phase=False) + ynzpref = signal.upfirdn(bz, u, up=1, down=2)[:100] + + assert_equal(ynzp, ynzpref) + + yzp = signal.decimate(u, 2, ftype=system, zero_phase=True) + yzpref = signal.resample_poly(u, 1, 2, window=bz) + + assert_equal(yzp, yzpref) + + +class TestHilbert: + + def test_bad_args(self): + x = np.array([1.0 + 0.0j]) + assert_raises(ValueError, hilbert, x) + x = np.arange(8.0) + assert_raises(ValueError, hilbert, x, N=0) + + def test_hilbert_theoretical(self): + # test cases by Ariel Rokem + decimal = 14 + + pi = np.pi + t = np.arange(0, 2 * pi, pi / 256) + a0 = np.sin(t) + a1 = np.cos(t) + a2 = np.sin(2 * t) + a3 = np.cos(2 * t) + a = np.vstack([a0, a1, a2, a3]) + + h = hilbert(a) + h_abs = np.abs(h) + h_angle = np.angle(h) + h_real = np.real(h) + + # The real part should be equal to the original signals: + assert_almost_equal(h_real, a, decimal) + # The absolute value should be one everywhere, for this input: + assert_almost_equal(h_abs, np.ones(a.shape), decimal) + # For the 'slow' sine - the phase should go from -pi/2 to pi/2 in + # the first 256 bins: + assert_almost_equal(h_angle[0, :256], + np.arange(-pi / 2, pi / 2, pi / 256), + decimal) + # For the 'slow' cosine - the phase should go from 0 to pi in the + # same interval: + assert_almost_equal( + h_angle[1, :256], np.arange(0, pi, pi / 256), decimal) + # The 'fast' sine should make this phase transition in half the time: + assert_almost_equal(h_angle[2, :128], + np.arange(-pi / 2, pi / 2, pi / 128), + decimal) + # Ditto for the 'fast' cosine: + assert_almost_equal( + h_angle[3, :128], np.arange(0, pi, pi / 128), decimal) + + # The imaginary part of hilbert(cos(t)) = sin(t) Wikipedia + assert_almost_equal(h[1].imag, a0, decimal) + + def test_hilbert_axisN(self): + # tests for axis and N arguments + a = np.arange(18).reshape(3, 6) + # test axis + aa = hilbert(a, axis=-1) + assert_equal(hilbert(a.T, axis=0), aa.T) + # test 1d + assert_almost_equal(hilbert(a[0]), aa[0], 14) + + # test N + aan = hilbert(a, N=20, axis=-1) + assert_equal(aan.shape, [3, 20]) + assert_equal(hilbert(a.T, N=20, axis=0).shape, [20, 3]) + # the next test is just a regression test, + # no idea whether numbers make sense + a0hilb = np.array([0.000000000000000e+00 - 1.72015830311905j, + 1.000000000000000e+00 - 2.047794505137069j, + 1.999999999999999e+00 - 2.244055555687583j, + 3.000000000000000e+00 - 1.262750302935009j, + 4.000000000000000e+00 - 1.066489252384493j, + 5.000000000000000e+00 + 2.918022706971047j, + 8.881784197001253e-17 + 3.845658908989067j, + -9.444121133484362e-17 + 0.985044202202061j, + -1.776356839400251e-16 + 1.332257797702019j, + -3.996802888650564e-16 + 0.501905089898885j, + 1.332267629550188e-16 + 0.668696078880782j, + -1.192678053963799e-16 + 0.235487067862679j, + -1.776356839400251e-16 + 0.286439612812121j, + 3.108624468950438e-16 + 0.031676888064907j, + 1.332267629550188e-16 - 0.019275656884536j, + -2.360035624836702e-16 - 0.1652588660287j, + 0.000000000000000e+00 - 0.332049855010597j, + 3.552713678800501e-16 - 0.403810179797771j, + 8.881784197001253e-17 - 0.751023775297729j, + 9.444121133484362e-17 - 0.79252210110103j]) + assert_almost_equal(aan[0], a0hilb, 14, 'N regression') + + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_hilbert_types(self, dtype): + in_typed = np.zeros(8, dtype=dtype) + assert_equal(np.real(signal.hilbert(in_typed)).dtype, dtype) + + +class TestHilbert2: + + def test_bad_args(self): + # x must be real. + x = np.array([[1.0 + 0.0j]]) + assert_raises(ValueError, hilbert2, x) + + # x must be rank 2. + x = np.arange(24).reshape(2, 3, 4) + assert_raises(ValueError, hilbert2, x) + + # Bad value for N. + x = np.arange(16).reshape(4, 4) + assert_raises(ValueError, hilbert2, x, N=0) + assert_raises(ValueError, hilbert2, x, N=(2, 0)) + assert_raises(ValueError, hilbert2, x, N=(2,)) + + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_hilbert2_types(self, dtype): + in_typed = np.zeros((2, 32), dtype=dtype) + assert_equal(np.real(signal.hilbert2(in_typed)).dtype, dtype) + + +class TestEnvelope: + """Unit tests for function `._signaltools.envelope()`. """ + + @staticmethod + def assert_close(actual, desired, msg): + """Little helper to compare to arrays with proper tolerances""" + xp_assert_close(actual, desired, atol=1e-12, rtol=1e-12, err_msg=msg) + + def test_envelope_invalid_parameters(self): + """For `envelope()` Raise all exceptions that are used to verify function + parameters. """ + with pytest.raises(ValueError, + match=r"Invalid parameter axis=2 for z.shape=.*"): + envelope(np.ones(3), axis=2) + with pytest.raises(ValueError, + match=r"z.shape\[axis\] not > 0 for z.shape=.*"): + envelope(np.ones((3, 0)), axis=1) + for bp_in in [(0, 1, 2), (0, 2.), (None, 2.)]: + ts = ', '.join(map(str, bp_in)) + with pytest.raises(ValueError, + match=rf"bp_in=\({ts}\) isn't a 2-tuple of.*"): + # noinspection PyTypeChecker + envelope(np.ones(4), bp_in=bp_in) + with pytest.raises(ValueError, + match="n_out=10.0 is not a positive integer or.*"): + # noinspection PyTypeChecker + envelope(np.ones(4), n_out=10.) + for bp_in in [(-1, 3), (1, 1), (0, 10)]: + with pytest.raises(ValueError, + match=r"`-n//2 <= bp_in\[0\] < bp_in\[1\] <=.*"): + envelope(np.ones(4), bp_in=bp_in) + with pytest.raises(ValueError, match="residual='undefined' not in .*"): + # noinspection PyTypeChecker + envelope(np.ones(4), residual='undefined') + + def test_envelope_verify_parameters(self): + """Ensure that the various parametrizations produce compatible results. """ + Z, Zr_a = [4, 2, 2, 3, 0], [4, 0, 0, 6, 0, 0, 0, 0] + z = sp_fft.irfft(Z) + n = len(z) + + # the reference envelope: + ze2_0, zr_0 = envelope(z, (1, 3), residual='all', squared=True) + self.assert_close(sp_fft.rfft(ze2_0), np.array([4, 2, 0, 0, 0]).astype(complex), + msg="Envelope calculation error") + self.assert_close(sp_fft.rfft(zr_0), np.array([4, 0, 0, 3, 0]).astype(complex), + msg="Residual calculation error") + + ze_1, zr_1 = envelope(z, (1, 3), residual='all', squared=False) + self.assert_close(ze_1**2, ze2_0, + msg="Unsquared versus Squared envelope calculation error") + self.assert_close(zr_1, zr_0, + msg="Unsquared versus Squared residual calculation error") + + ze2_2, zr_2 = envelope(z, (1, 3), residual='all', squared=True, n_out=3*n) + self.assert_close(ze2_2[::3], ze2_0, + msg="3x up-sampled envelope calculation error") + self.assert_close(zr_2[::3], zr_0, + msg="3x up-sampled residual calculation error") + + ze2_3, zr_3 = envelope(z, (1, 3), residual='lowpass', squared=True) + self.assert_close(ze2_3, ze2_0, + msg="`residual='lowpass'` envelope calculation error") + self.assert_close(sp_fft.rfft(zr_3), np.array([4, 0, 0, 0, 0]).astype(complex), + msg="`residual='lowpass'` residual calculation error") + + ze2_4 = envelope(z, (1, 3), residual=None, squared=True) + self.assert_close(ze2_4, ze2_0, + msg="`residual=None` envelope calculation error") + + # compare complex analytic signal to real version + Z_a = np.copy(Z) + Z_a[1:] *= 2 + z_a = sp_fft.ifft(Z_a, n=n) # analytic signal of Z + self.assert_close(z_a.real, z, + msg="Reference analytic signal error") + ze2_a, zr_a = envelope(z_a, (1, 3), residual='all', squared=True) + self.assert_close(ze2_a, ze2_0.astype(complex), # dtypes must match + msg="Complex envelope calculation error") + self.assert_close(sp_fft.fft(zr_a), np.array(Zr_a).astype(complex), + msg="Complex residual calculation error") + + @pytest.mark.parametrize( + " Z, bp_in, Ze2_desired, Zr_desired", + [([1, 0, 2, 2, 0], (1, None), [4, 2, 0, 0, 0], [1, 0, 0, 0, 0]), + ([4, 0, 2, 0, 0], (0, None), [4, 0, 2, 0, 0], [0, 0, 0, 0, 0]), + ([4, 0, 0, 2, 0], (None, None), [4, 0, 0, 2, 0], [0, 0, 0, 0, 0]), + ([0, 0, 2, 2, 0], (1, 3), [2, 0, 0, 0, 0], [0, 0, 0, 2, 0]), + ([4, 0, 2, 2, 0], (-3, 3), [4, 0, 2, 0, 0], [0, 0, 0, 2, 0]), + ([4, 0, 3, 4, 0], (None, 1), [2, 0, 0, 0, 0], [0, 0, 3, 4, 0]), + ([4, 0, 3, 4, 0], (None, 0), [0, 0, 0, 0, 0], [4, 0, 3, 4, 0])]) + def test_envelope_real_signals(self, Z, bp_in, Ze2_desired, Zr_desired): + """Test envelope calculation with real-valued test signals. + + The comparisons are performed in the Fourier space, since it makes evaluating + the bandpass filter behavior straightforward. Note that also the squared + envelope can be easily calculated by hand, if one recalls that coefficients of + a complex-valued Fourier series representing the signal can be directly + determined by an FFT and that the absolute square of a Fourier series is again + a Fourier series. + """ + z = sp_fft.irfft(Z) + ze2, zr = envelope(z, bp_in, residual='all', squared=True) + ze2_lp, zr_lp = envelope(z, bp_in, residual='lowpass', squared=True) + Ze2, Zr, Ze2_lp, Zr_lp = (sp_fft.rfft(z_) for z_ in (ze2, zr, ze2_lp, zr_lp)) + + Ze2_desired = np.array(Ze2_desired).astype(complex) + Zr_desired = np.array(Zr_desired).astype(complex) + self.assert_close(Ze2, Ze2_desired, + msg="Envelope calculation error (residual='all')") + self.assert_close(Zr, Zr_desired, + msg="Residual calculation error (residual='all')") + + if bp_in[1] is not None: + Zr_desired[bp_in[1]:] = 0 + self.assert_close(Ze2_lp, Ze2_desired, + msg="Envelope calculation error (residual='lowpass')") + self.assert_close(Zr_lp, Zr_desired, + msg="Residual calculation error (residual='lowpass')") + + @pytest.mark.parametrize( + " Z, bp_in, Ze2_desired, Zr_desired", + [([0, 5, 0, 5, 0], (None, None), [5, 0, 10, 0, 5], [0, 0, 0, 0, 0]), + ([1, 5, 0, 5, 2], (-1, 2), [5, 0, 10, 0, 5], [1, 0, 0, 0, 2]), + ([1, 2, 6, 0, 6, 3], (-1, 2), [0, 6, 0, 12, 0, 6], [1, 2, 0, 0, 0, 3]) + ]) + def test_envelope_complex_signals(self, Z, bp_in, Ze2_desired, Zr_desired): + """Test envelope calculation with complex-valued test signals. + + We only need to test for the complex envelope here, since the ``Nones``s in the + bandpass filter were already tested in the previous test. + """ + z = sp_fft.ifft(sp_fft.ifftshift(Z)) + ze2, zr = envelope(z, bp_in, residual='all', squared=True) + Ze2, Zr = (sp_fft.fftshift(sp_fft.fft(z_)) for z_ in (ze2, zr)) + + self.assert_close(Ze2, np.array(Ze2_desired).astype(complex), + msg="Envelope calculation error") + self.assert_close(Zr, np.array(Zr_desired).astype(complex), + msg="Residual calculation error") + + def test_envelope_verify_axis_parameter(self): + """Test for multi-channel envelope calculations. """ + z = sp_fft.irfft([[1, 0, 2, 2, 0], [7, 0, 4, 4, 0]]) + Ze2_desired = np.array([[4, 2, 0, 0, 0], [16, 8, 0, 0, 0]], + dtype=complex) + Zr_desired = np.array([[1, 0, 0, 0, 0], [7, 0, 0, 0, 0]], dtype=complex) + + ze2, zr = envelope(z, squared=True, axis=1) + ye2T, yrT = envelope(z.T, squared=True, axis=0) + Ze2, Ye2, Zr, Yr = (sp_fft.rfft(z_) for z_ in (ze2, ye2T.T, zr, yrT.T)) + + self.assert_close(Ze2, Ze2_desired, msg="2d envelope calculation error") + self.assert_close(Zr, Zr_desired, msg="2d residual calculation error") + self.assert_close(Ye2, Ze2_desired, msg="Transposed 2d envelope calc. error") + self.assert_close(Yr, Zr_desired, msg="Transposed 2d residual calc. error") + + def test_envelope_verify_axis_parameter_complex(self): + """Test for multi-channel envelope calculations with complex values. """ + z = sp_fft.ifft(sp_fft.ifftshift([[1, 5, 0, 5, 2], [1, 10, 0, 10, 2]], axes=1)) + Ze2_des = np.array([[5, 0, 10, 0, 5], [20, 0, 40, 0, 20],], + dtype=complex) + Zr_des = np.array([[1, 0, 0, 0, 2], [1, 0, 0, 0, 2]], dtype=complex) + + kw = dict(bp_in=(-1, 2), residual='all', squared=True) + ze2, zr = envelope(z, axis=1, **kw) + ye2T, yrT = envelope(z.T, axis=0, **kw) + Ze2, Ye2, Zr, Yr = (sp_fft.fftshift(sp_fft.fft(z_), axes=1) + for z_ in (ze2, ye2T.T, zr, yrT.T)) + + self.assert_close(Ze2, Ze2_des, msg="2d envelope calculation error") + self.assert_close(Zr, Zr_des, msg="2d residual calculation error") + self.assert_close(Ye2, Ze2_des, msg="Transposed 2d envelope calc. error") + self.assert_close(Yr, Zr_des, msg="Transposed 2d residual calc. error") + + @pytest.mark.parametrize('X', [[4, 0, 0, 1, 2], [4, 0, 0, 2, 1, 2]]) + def test_compare_envelope_hilbert(self, X): + """Compare output of `envelope()` and `hilbert()`. """ + x = sp_fft.irfft(X) + e_hil = np.abs(hilbert(x)) + e_env = envelope(x, (None, None), residual=None) + self.assert_close(e_hil, e_env, msg="Hilbert-Envelope comparison error") + + +class TestPartialFractionExpansion: + @staticmethod + def assert_rp_almost_equal(r, p, r_true, p_true, decimal=7): + r_true = np.asarray(r_true) + p_true = np.asarray(p_true) + + distance = np.hypot(abs(p[:, None] - p_true), + abs(r[:, None] - r_true)) + + rows, cols = linear_sum_assignment(distance) + assert_almost_equal(p[rows], p_true[cols], decimal=decimal) + assert_almost_equal(r[rows], r_true[cols], decimal=decimal) + + def test_compute_factors(self): + factors, poly = _compute_factors([1, 2, 3], [3, 2, 1]) + assert_equal(len(factors), 3) + assert_almost_equal(factors[0], np.poly([2, 2, 3])) + assert_almost_equal(factors[1], np.poly([1, 1, 1, 3])) + assert_almost_equal(factors[2], np.poly([1, 1, 1, 2, 2])) + assert_almost_equal(poly, np.poly([1, 1, 1, 2, 2, 3])) + + factors, poly = _compute_factors([1, 2, 3], [3, 2, 1], + include_powers=True) + assert_equal(len(factors), 6) + assert_almost_equal(factors[0], np.poly([1, 1, 2, 2, 3])) + assert_almost_equal(factors[1], np.poly([1, 2, 2, 3])) + assert_almost_equal(factors[2], np.poly([2, 2, 3])) + assert_almost_equal(factors[3], np.poly([1, 1, 1, 2, 3])) + assert_almost_equal(factors[4], np.poly([1, 1, 1, 3])) + assert_almost_equal(factors[5], np.poly([1, 1, 1, 2, 2])) + assert_almost_equal(poly, np.poly([1, 1, 1, 2, 2, 3])) + + def test_group_poles(self): + unique, multiplicity = _group_poles( + [1.0, 1.001, 1.003, 2.0, 2.003, 3.0], 0.1, 'min') + assert_equal(unique, [1.0, 2.0, 3.0]) + assert_equal(multiplicity, [3, 2, 1]) + + def test_residue_general(self): + # Test are taken from issue #4464, note that poles in scipy are + # in increasing by absolute value order, opposite to MATLAB. + r, p, k = residue([5, 3, -2, 7], [-4, 0, 8, 3]) + assert_almost_equal(r, [1.3320, -0.6653, -1.4167], decimal=4) + assert_almost_equal(p, [-0.4093, -1.1644, 1.5737], decimal=4) + assert_almost_equal(k, [-1.2500], decimal=4) + + r, p, k = residue([-4, 8], [1, 6, 8]) + assert_almost_equal(r, [8, -12]) + assert_almost_equal(p, [-2, -4]) + assert_equal(k.size, 0) + + r, p, k = residue([4, 1], [1, -1, -2]) + assert_almost_equal(r, [1, 3]) + assert_almost_equal(p, [-1, 2]) + assert_equal(k.size, 0) + + r, p, k = residue([4, 3], [2, -3.4, 1.98, -0.406]) + self.assert_rp_almost_equal( + r, p, [-18.125 - 13.125j, -18.125 + 13.125j, 36.25], + [0.5 - 0.2j, 0.5 + 0.2j, 0.7]) + assert_equal(k.size, 0) + + r, p, k = residue([2, 1], [1, 5, 8, 4]) + self.assert_rp_almost_equal(r, p, [-1, 1, 3], [-1, -2, -2]) + assert_equal(k.size, 0) + + r, p, k = residue([3, -1.1, 0.88, -2.396, 1.348], + [1, -0.7, -0.14, 0.048]) + assert_almost_equal(r, [-3, 4, 1]) + assert_almost_equal(p, [0.2, -0.3, 0.8]) + assert_almost_equal(k, [3, 1]) + + r, p, k = residue([1], [1, 2, -3]) + assert_almost_equal(r, [0.25, -0.25]) + assert_almost_equal(p, [1, -3]) + assert_equal(k.size, 0) + + r, p, k = residue([1, 0, -5], [1, 0, 0, 0, -1]) + self.assert_rp_almost_equal(r, p, + [1, 1.5j, -1.5j, -1], [-1, -1j, 1j, 1]) + assert_equal(k.size, 0) + + r, p, k = residue([3, 8, 6], [1, 3, 3, 1]) + self.assert_rp_almost_equal(r, p, [1, 2, 3], [-1, -1, -1]) + assert_equal(k.size, 0) + + r, p, k = residue([3, -1], [1, -3, 2]) + assert_almost_equal(r, [-2, 5]) + assert_almost_equal(p, [1, 2]) + assert_equal(k.size, 0) + + r, p, k = residue([2, 3, -1], [1, -3, 2]) + assert_almost_equal(r, [-4, 13]) + assert_almost_equal(p, [1, 2]) + assert_almost_equal(k, [2]) + + r, p, k = residue([7, 2, 3, -1], [1, -3, 2]) + assert_almost_equal(r, [-11, 69]) + assert_almost_equal(p, [1, 2]) + assert_almost_equal(k, [7, 23]) + + r, p, k = residue([2, 3, -1], [1, -3, 4, -2]) + self.assert_rp_almost_equal(r, p, [4, -1 + 3.5j, -1 - 3.5j], + [1, 1 - 1j, 1 + 1j]) + assert_almost_equal(k.size, 0) + + def test_residue_leading_zeros(self): + # Leading zeros in numerator or denominator must not affect the answer. + r0, p0, k0 = residue([5, 3, -2, 7], [-4, 0, 8, 3]) + r1, p1, k1 = residue([0, 5, 3, -2, 7], [-4, 0, 8, 3]) + r2, p2, k2 = residue([5, 3, -2, 7], [0, -4, 0, 8, 3]) + r3, p3, k3 = residue([0, 0, 5, 3, -2, 7], [0, 0, 0, -4, 0, 8, 3]) + assert_almost_equal(r0, r1) + assert_almost_equal(r0, r2) + assert_almost_equal(r0, r3) + assert_almost_equal(p0, p1) + assert_almost_equal(p0, p2) + assert_almost_equal(p0, p3) + assert_almost_equal(k0, k1) + assert_almost_equal(k0, k2) + assert_almost_equal(k0, k3) + + def test_resiude_degenerate(self): + # Several tests for zero numerator and denominator. + r, p, k = residue([0, 0], [1, 6, 8]) + assert_almost_equal(r, [0, 0]) + assert_almost_equal(p, [-2, -4]) + assert_equal(k.size, 0) + + r, p, k = residue(0, 1) + assert_equal(r.size, 0) + assert_equal(p.size, 0) + assert_equal(k.size, 0) + + with pytest.raises(ValueError, match="Denominator `a` is zero."): + residue(1, 0) + + def test_residuez_general(self): + r, p, k = residuez([1, 6, 6, 2], [1, -(2 + 1j), (1 + 2j), -1j]) + self.assert_rp_almost_equal(r, p, [-2+2.5j, 7.5+7.5j, -4.5-12j], + [1j, 1, 1]) + assert_almost_equal(k, [2j]) + + r, p, k = residuez([1, 2, 1], [1, -1, 0.3561]) + self.assert_rp_almost_equal(r, p, + [-0.9041 - 5.9928j, -0.9041 + 5.9928j], + [0.5 + 0.3257j, 0.5 - 0.3257j], + decimal=4) + assert_almost_equal(k, [2.8082], decimal=4) + + r, p, k = residuez([1, -1], [1, -5, 6]) + assert_almost_equal(r, [-1, 2]) + assert_almost_equal(p, [2, 3]) + assert_equal(k.size, 0) + + r, p, k = residuez([2, 3, 4], [1, 3, 3, 1]) + self.assert_rp_almost_equal(r, p, [4, -5, 3], [-1, -1, -1]) + assert_equal(k.size, 0) + + r, p, k = residuez([1, -10, -4, 4], [2, -2, -4]) + assert_almost_equal(r, [0.5, -1.5]) + assert_almost_equal(p, [-1, 2]) + assert_almost_equal(k, [1.5, -1]) + + r, p, k = residuez([18], [18, 3, -4, -1]) + self.assert_rp_almost_equal(r, p, + [0.36, 0.24, 0.4], [0.5, -1/3, -1/3]) + assert_equal(k.size, 0) + + r, p, k = residuez([2, 3], np.polymul([1, -1/2], [1, 1/4])) + assert_almost_equal(r, [-10/3, 16/3]) + assert_almost_equal(p, [-0.25, 0.5]) + assert_equal(k.size, 0) + + r, p, k = residuez([1, -2, 1], [1, -1]) + assert_almost_equal(r, [0]) + assert_almost_equal(p, [1]) + assert_almost_equal(k, [1, -1]) + + r, p, k = residuez(1, [1, -1j]) + assert_almost_equal(r, [1]) + assert_almost_equal(p, [1j]) + assert_equal(k.size, 0) + + r, p, k = residuez(1, [1, -1, 0.25]) + assert_almost_equal(r, [0, 1]) + assert_almost_equal(p, [0.5, 0.5]) + assert_equal(k.size, 0) + + r, p, k = residuez(1, [1, -0.75, .125]) + assert_almost_equal(r, [-1, 2]) + assert_almost_equal(p, [0.25, 0.5]) + assert_equal(k.size, 0) + + r, p, k = residuez([1, 6, 2], [1, -2, 1]) + assert_almost_equal(r, [-10, 9]) + assert_almost_equal(p, [1, 1]) + assert_almost_equal(k, [2]) + + r, p, k = residuez([6, 2], [1, -2, 1]) + assert_almost_equal(r, [-2, 8]) + assert_almost_equal(p, [1, 1]) + assert_equal(k.size, 0) + + r, p, k = residuez([1, 6, 6, 2], [1, -2, 1]) + assert_almost_equal(r, [-24, 15]) + assert_almost_equal(p, [1, 1]) + assert_almost_equal(k, [10, 2]) + + r, p, k = residuez([1, 0, 1], [1, 0, 0, 0, 0, -1]) + self.assert_rp_almost_equal(r, p, + [0.2618 + 0.1902j, 0.2618 - 0.1902j, + 0.4, 0.0382 - 0.1176j, 0.0382 + 0.1176j], + [-0.8090 + 0.5878j, -0.8090 - 0.5878j, + 1.0, 0.3090 + 0.9511j, 0.3090 - 0.9511j], + decimal=4) + assert_equal(k.size, 0) + + def test_residuez_trailing_zeros(self): + # Trailing zeros in numerator or denominator must not affect the + # answer. + r0, p0, k0 = residuez([5, 3, -2, 7], [-4, 0, 8, 3]) + r1, p1, k1 = residuez([5, 3, -2, 7, 0], [-4, 0, 8, 3]) + r2, p2, k2 = residuez([5, 3, -2, 7], [-4, 0, 8, 3, 0]) + r3, p3, k3 = residuez([5, 3, -2, 7, 0, 0], [-4, 0, 8, 3, 0, 0, 0]) + assert_almost_equal(r0, r1) + assert_almost_equal(r0, r2) + assert_almost_equal(r0, r3) + assert_almost_equal(p0, p1) + assert_almost_equal(p0, p2) + assert_almost_equal(p0, p3) + assert_almost_equal(k0, k1) + assert_almost_equal(k0, k2) + assert_almost_equal(k0, k3) + + def test_residuez_degenerate(self): + r, p, k = residuez([0, 0], [1, 6, 8]) + assert_almost_equal(r, [0, 0]) + assert_almost_equal(p, [-2, -4]) + assert_equal(k.size, 0) + + r, p, k = residuez(0, 1) + assert_equal(r.size, 0) + assert_equal(p.size, 0) + assert_equal(k.size, 0) + + with pytest.raises(ValueError, match="Denominator `a` is zero."): + residuez(1, 0) + + with pytest.raises(ValueError, + match="First coefficient of determinant `a` must " + "be non-zero."): + residuez(1, [0, 1, 2, 3]) + + def test_inverse_unique_roots_different_rtypes(self): + # This test was inspired by GitHub issue 2496. + r = [3 / 10, -1 / 6, -2 / 15] + p = [0, -2, -5] + k = [] + b_expected = [0, 1, 3] + a_expected = [1, 7, 10, 0] + + # With the default tolerance, the rtype does not matter + # for this example. + for rtype in ('avg', 'mean', 'min', 'minimum', 'max', 'maximum'): + b, a = invres(r, p, k, rtype=rtype) + assert_allclose(b, b_expected) + assert_allclose(a, a_expected) + + b, a = invresz(r, p, k, rtype=rtype) + assert_allclose(b, b_expected) + assert_allclose(a, a_expected) + + def test_inverse_repeated_roots_different_rtypes(self): + r = [3 / 20, -7 / 36, -1 / 6, 2 / 45] + p = [0, -2, -2, -5] + k = [] + b_expected = [0, 0, 1, 3] + b_expected_z = [-1/6, -2/3, 11/6, 3] + a_expected = [1, 9, 24, 20, 0] + + for rtype in ('avg', 'mean', 'min', 'minimum', 'max', 'maximum'): + b, a = invres(r, p, k, rtype=rtype) + assert_allclose(b, b_expected, atol=1e-14) + assert_allclose(a, a_expected) + + b, a = invresz(r, p, k, rtype=rtype) + assert_allclose(b, b_expected_z, atol=1e-14) + assert_allclose(a, a_expected) + + def test_inverse_bad_rtype(self): + r = [3 / 20, -7 / 36, -1 / 6, 2 / 45] + p = [0, -2, -2, -5] + k = [] + with pytest.raises(ValueError, match="`rtype` must be one of"): + invres(r, p, k, rtype='median') + with pytest.raises(ValueError, match="`rtype` must be one of"): + invresz(r, p, k, rtype='median') + + def test_invresz_one_coefficient_bug(self): + # Regression test for issue in gh-4646. + r = [1] + p = [2] + k = [0] + b, a = invresz(r, p, k) + assert_allclose(b, [1.0]) + assert_allclose(a, [1.0, -2.0]) + + def test_invres(self): + b, a = invres([1], [1], []) + assert_almost_equal(b, [1]) + assert_almost_equal(a, [1, -1]) + + b, a = invres([1 - 1j, 2, 0.5 - 3j], [1, 0.5j, 1 + 1j], []) + assert_almost_equal(b, [3.5 - 4j, -8.5 + 0.25j, 3.5 + 3.25j]) + assert_almost_equal(a, [1, -2 - 1.5j, 0.5 + 2j, 0.5 - 0.5j]) + + b, a = invres([0.5, 1], [1 - 1j, 2 + 2j], [1, 2, 3]) + assert_almost_equal(b, [1, -1 - 1j, 1 - 2j, 0.5 - 3j, 10]) + assert_almost_equal(a, [1, -3 - 1j, 4]) + + b, a = invres([-1, 2, 1j, 3 - 1j, 4, -2], + [-1, 2 - 1j, 2 - 1j, 3, 3, 3], []) + assert_almost_equal(b, [4 - 1j, -28 + 16j, 40 - 62j, 100 + 24j, + -292 + 219j, 192 - 268j]) + assert_almost_equal(a, [1, -12 + 2j, 53 - 20j, -96 + 68j, 27 - 72j, + 108 - 54j, -81 + 108j]) + + b, a = invres([-1, 1j], [1, 1], [1, 2]) + assert_almost_equal(b, [1, 0, -4, 3 + 1j]) + assert_almost_equal(a, [1, -2, 1]) + + def test_invresz(self): + b, a = invresz([1], [1], []) + assert_almost_equal(b, [1]) + assert_almost_equal(a, [1, -1]) + + b, a = invresz([1 - 1j, 2, 0.5 - 3j], [1, 0.5j, 1 + 1j], []) + assert_almost_equal(b, [3.5 - 4j, -8.5 + 0.25j, 3.5 + 3.25j]) + assert_almost_equal(a, [1, -2 - 1.5j, 0.5 + 2j, 0.5 - 0.5j]) + + b, a = invresz([0.5, 1], [1 - 1j, 2 + 2j], [1, 2, 3]) + assert_almost_equal(b, [2.5, -3 - 1j, 1 - 2j, -1 - 3j, 12]) + assert_almost_equal(a, [1, -3 - 1j, 4]) + + b, a = invresz([-1, 2, 1j, 3 - 1j, 4, -2], + [-1, 2 - 1j, 2 - 1j, 3, 3, 3], []) + assert_almost_equal(b, [6, -50 + 11j, 100 - 72j, 80 + 58j, + -354 + 228j, 234 - 297j]) + assert_almost_equal(a, [1, -12 + 2j, 53 - 20j, -96 + 68j, 27 - 72j, + 108 - 54j, -81 + 108j]) + + b, a = invresz([-1, 1j], [1, 1], [1, 2]) + assert_almost_equal(b, [1j, 1, -3, 2]) + assert_almost_equal(a, [1, -2, 1]) + + def test_inverse_scalar_arguments(self): + b, a = invres(1, 1, 1) + assert_almost_equal(b, [1, 0]) + assert_almost_equal(a, [1, -1]) + + b, a = invresz(1, 1, 1) + assert_almost_equal(b, [2, -1]) + assert_almost_equal(a, [1, -1]) + + +class TestVectorstrength: + + def test_single_1dperiod(self): + events = np.array([.5]) + period = 5. + targ_strength = 1. + targ_phase = .1 + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 0) + assert_equal(phase.ndim, 0) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_single_2dperiod(self): + events = np.array([.5]) + period = [1, 2, 5.] + targ_strength = [1.] * 3 + targ_phase = np.array([.5, .25, .1]) + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 1) + assert_equal(phase.ndim, 1) + assert_array_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_equal_1dperiod(self): + events = np.array([.25, .25, .25, .25, .25, .25]) + period = 2 + targ_strength = 1. + targ_phase = .125 + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 0) + assert_equal(phase.ndim, 0) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_equal_2dperiod(self): + events = np.array([.25, .25, .25, .25, .25, .25]) + period = [1, 2, ] + targ_strength = [1.] * 2 + targ_phase = np.array([.25, .125]) + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 1) + assert_equal(phase.ndim, 1) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_spaced_1dperiod(self): + events = np.array([.1, 1.1, 2.1, 4.1, 10.1]) + period = 1 + targ_strength = 1. + targ_phase = .1 + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 0) + assert_equal(phase.ndim, 0) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_spaced_2dperiod(self): + events = np.array([.1, 1.1, 2.1, 4.1, 10.1]) + period = [1, .5] + targ_strength = [1.] * 2 + targ_phase = np.array([.1, .2]) + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 1) + assert_equal(phase.ndim, 1) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_partial_1dperiod(self): + events = np.array([.25, .5, .75]) + period = 1 + targ_strength = 1. / 3. + targ_phase = .5 + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 0) + assert_equal(phase.ndim, 0) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_partial_2dperiod(self): + events = np.array([.25, .5, .75]) + period = [1., 1., 1., 1.] + targ_strength = [1. / 3.] * 4 + targ_phase = np.array([.5, .5, .5, .5]) + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 1) + assert_equal(phase.ndim, 1) + assert_almost_equal(strength, targ_strength) + assert_almost_equal(phase, 2 * np.pi * targ_phase) + + def test_opposite_1dperiod(self): + events = np.array([0, .25, .5, .75]) + period = 1. + targ_strength = 0 + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 0) + assert_equal(phase.ndim, 0) + assert_almost_equal(strength, targ_strength) + + def test_opposite_2dperiod(self): + events = np.array([0, .25, .5, .75]) + period = [1.] * 10 + targ_strength = [0.] * 10 + + strength, phase = vectorstrength(events, period) + + assert_equal(strength.ndim, 1) + assert_equal(phase.ndim, 1) + assert_almost_equal(strength, targ_strength) + + def test_2d_events_ValueError(self): + events = np.array([[1, 2]]) + period = 1. + assert_raises(ValueError, vectorstrength, events, period) + + def test_2d_period_ValueError(self): + events = 1. + period = np.array([[1]]) + assert_raises(ValueError, vectorstrength, events, period) + + def test_zero_period_ValueError(self): + events = 1. + period = 0 + assert_raises(ValueError, vectorstrength, events, period) + + def test_negative_period_ValueError(self): + events = 1. + period = -1 + assert_raises(ValueError, vectorstrength, events, period) + + +def assert_allclose_cast(actual, desired, rtol=1e-7, atol=0): + """Wrap assert_allclose while casting object arrays.""" + if actual.dtype.kind == 'O': + dtype = np.array(actual.flat[0]).dtype + actual, desired = actual.astype(dtype), desired.astype(dtype) + assert_allclose(actual, desired, rtol, atol) + + +@pytest.mark.filterwarnings('ignore::DeprecationWarning') +@pytest.mark.parametrize('func', (sosfilt, lfilter)) +def test_nonnumeric_dtypes(func): + x = [Decimal(1), Decimal(2), Decimal(3)] + b = [Decimal(1), Decimal(2), Decimal(3)] + a = [Decimal(1), Decimal(2), Decimal(3)] + x = np.array(x) + assert x.dtype.kind == 'O' + desired = lfilter(np.array(b, float), np.array(a, float), x.astype(float)) + if func is sosfilt: + actual = sosfilt([b + a], x) + else: + actual = lfilter(b, a, x) + assert all(isinstance(x, Decimal) for x in actual) + assert_allclose(actual.astype(float), desired.astype(float)) + # Degenerate cases + if func is lfilter: + args = [1., 1.] + else: + args = [tf2sos(1., 1.)] + + with pytest.raises(ValueError, match='must be at least 1-D'): + func(*args, x=1.) + + +@pytest.mark.parametrize('dt', 'fdFD') +class TestSOSFilt: + + # The test_rank* tests are pulled from _TestLinearFilter + def test_rank1(self, dt): + x = np.linspace(0, 5, 6).astype(dt) + b = np.array([1, -1]).astype(dt) + a = np.array([0.5, -0.5]).astype(dt) + + # Test simple IIR + y_r = np.array([0, 2, 4, 6, 8, 10.]).astype(dt) + sos = tf2sos(b, a) + assert_array_almost_equal(sosfilt(tf2sos(b, a), x), y_r) + + # Test simple FIR + b = np.array([1, 1]).astype(dt) + # NOTE: This was changed (rel. to TestLinear...) to add a pole @zero: + a = np.array([1, 0]).astype(dt) + y_r = np.array([0, 1, 3, 5, 7, 9.]).astype(dt) + assert_array_almost_equal(sosfilt(tf2sos(b, a), x), y_r) + + b = [1, 1, 0] + a = [1, 0, 0] + x = np.ones(8) + sos = np.concatenate((b, a)) + sos.shape = (1, 6) + y = sosfilt(sos, x) + assert_allclose(y, [1, 2, 2, 2, 2, 2, 2, 2]) + + def test_rank2(self, dt): + shape = (4, 3) + x = np.linspace(0, np.prod(shape) - 1, np.prod(shape)).reshape(shape) + x = x.astype(dt) + + b = np.array([1, -1]).astype(dt) + a = np.array([0.5, 0.5]).astype(dt) + + y_r2_a0 = np.array([[0, 2, 4], [6, 4, 2], [0, 2, 4], [6, 4, 2]], + dtype=dt) + + y_r2_a1 = np.array([[0, 2, 0], [6, -4, 6], [12, -10, 12], + [18, -16, 18]], dtype=dt) + + y = sosfilt(tf2sos(b, a), x, axis=0) + assert_array_almost_equal(y_r2_a0, y) + + y = sosfilt(tf2sos(b, a), x, axis=1) + assert_array_almost_equal(y_r2_a1, y) + + def test_rank3(self, dt): + shape = (4, 3, 2) + x = np.linspace(0, np.prod(shape) - 1, np.prod(shape)).reshape(shape) + + b = np.array([1, -1]).astype(dt) + a = np.array([0.5, 0.5]).astype(dt) + + # Test last axis + y = sosfilt(tf2sos(b, a), x) + for i in range(x.shape[0]): + for j in range(x.shape[1]): + assert_array_almost_equal(y[i, j], lfilter(b, a, x[i, j])) + + def test_initial_conditions(self, dt): + b1, a1 = signal.butter(2, 0.25, 'low') + b2, a2 = signal.butter(2, 0.75, 'low') + b3, a3 = signal.butter(2, 0.75, 'low') + b = np.convolve(np.convolve(b1, b2), b3) + a = np.convolve(np.convolve(a1, a2), a3) + sos = np.array((np.r_[b1, a1], np.r_[b2, a2], np.r_[b3, a3])) + + x = np.random.rand(50).astype(dt) + + # Stopping filtering and continuing + y_true, zi = lfilter(b, a, x[:20], zi=np.zeros(6)) + y_true = np.r_[y_true, lfilter(b, a, x[20:], zi=zi)[0]] + assert_allclose_cast(y_true, lfilter(b, a, x)) + + y_sos, zi = sosfilt(sos, x[:20], zi=np.zeros((3, 2))) + y_sos = np.r_[y_sos, sosfilt(sos, x[20:], zi=zi)[0]] + assert_allclose_cast(y_true, y_sos) + + # Use a step function + zi = sosfilt_zi(sos) + x = np.ones(8, dt) + y, zf = sosfilt(sos, x, zi=zi) + + assert_allclose_cast(y, np.ones(8)) + assert_allclose_cast(zf, zi) + + # Initial condition shape matching + x.shape = (1, 1) + x.shape # 3D + assert_raises(ValueError, sosfilt, sos, x, zi=zi) + zi_nd = zi.copy() + zi_nd.shape = (zi.shape[0], 1, 1, zi.shape[-1]) + assert_raises(ValueError, sosfilt, sos, x, + zi=zi_nd[:, :, :, [0, 1, 1]]) + y, zf = sosfilt(sos, x, zi=zi_nd) + assert_allclose_cast(y[0, 0], np.ones(8)) + assert_allclose_cast(zf[:, 0, 0, :], zi) + + def test_initial_conditions_3d_axis1(self, dt): + # Test the use of zi when sosfilt is applied to axis 1 of a 3-d input. + + # Input array is x. + x = np.random.RandomState(159).randint(0, 5, size=(2, 15, 3)) + x = x.astype(dt) + + # Design a filter in ZPK format and convert to SOS + zpk = signal.butter(6, 0.35, output='zpk') + sos = zpk2sos(*zpk) + nsections = sos.shape[0] + + # Filter along this axis. + axis = 1 + + # Initial conditions, all zeros. + shp = list(x.shape) + shp[axis] = 2 + shp = [nsections] + shp + z0 = np.zeros(shp) + + # Apply the filter to x. + yf, zf = sosfilt(sos, x, axis=axis, zi=z0) + + # Apply the filter to x in two stages. + y1, z1 = sosfilt(sos, x[:, :5, :], axis=axis, zi=z0) + y2, z2 = sosfilt(sos, x[:, 5:, :], axis=axis, zi=z1) + + # y should equal yf, and z2 should equal zf. + y = np.concatenate((y1, y2), axis=axis) + assert_allclose_cast(y, yf, rtol=1e-10, atol=1e-13) + assert_allclose_cast(z2, zf, rtol=1e-10, atol=1e-13) + + # let's try the "step" initial condition + zi = sosfilt_zi(sos) + zi.shape = [nsections, 1, 2, 1] + zi = zi * x[:, 0:1, :] + y = sosfilt(sos, x, axis=axis, zi=zi)[0] + # check it against the TF form + b, a = zpk2tf(*zpk) + zi = lfilter_zi(b, a) + zi.shape = [1, zi.size, 1] + zi = zi * x[:, 0:1, :] + y_tf = lfilter(b, a, x, axis=axis, zi=zi)[0] + assert_allclose_cast(y, y_tf, rtol=1e-10, atol=1e-13) + + def test_bad_zi_shape(self, dt): + # The shape of zi is checked before using any values in the + # arguments, so np.empty is fine for creating the arguments. + x = np.empty((3, 15, 3), dt) + sos = np.zeros((4, 6)) + zi = np.empty((4, 3, 3, 2)) # Correct shape is (4, 3, 2, 3) + with pytest.raises(ValueError, match='should be all ones'): + sosfilt(sos, x, zi=zi, axis=1) + sos[:, 3] = 1. + with pytest.raises(ValueError, match='Invalid zi shape'): + sosfilt(sos, x, zi=zi, axis=1) + + def test_sosfilt_zi(self, dt): + sos = signal.butter(6, 0.2, output='sos') + zi = sosfilt_zi(sos) + + y, zf = sosfilt(sos, np.ones(40, dt), zi=zi) + assert_allclose_cast(zf, zi, rtol=1e-13) + + # Expected steady state value of the step response of this filter: + ss = np.prod(sos[:, :3].sum(axis=-1) / sos[:, 3:].sum(axis=-1)) + assert_allclose_cast(y, ss, rtol=1e-13) + + # zi as array-like + _, zf = sosfilt(sos, np.ones(40, dt), zi=zi.tolist()) + assert_allclose_cast(zf, zi, rtol=1e-13) + + @pytest.mark.thread_unsafe + def test_dtype_deprecation(self, dt): + # gh-21211 + sos = np.asarray([1, 2, 3, 1, 5, 3], dtype=object).reshape(1, 6) + x = np.asarray([2, 3, 4, 5, 3, 4, 2, 2, 1], dtype=object) + with pytest.deprecated_call(match="dtype=object is not supported"): + sosfilt(sos, x) + + +class TestDeconvolve: + + def test_basic(self): + # From docstring example + original = [0, 1, 0, 0, 1, 1, 0, 0] + impulse_response = [2, 1] + recorded = [0, 2, 1, 0, 2, 3, 1, 0, 0] + recovered, remainder = signal.deconvolve(recorded, impulse_response) + assert_allclose(recovered, original) + + def test_n_dimensional_signal(self): + recorded = [[0, 0], [0, 0]] + impulse_response = [0, 0] + with pytest.raises(ValueError, match="signal must be 1-D."): + quotient, remainder = signal.deconvolve(recorded, impulse_response) + + def test_n_dimensional_divisor(self): + recorded = [0, 0] + impulse_response = [[0, 0], [0, 0]] + with pytest.raises(ValueError, match="divisor must be 1-D."): + quotient, remainder = signal.deconvolve(recorded, impulse_response) + + +class TestDetrend: + + def test_basic(self): + detrended = detrend(array([1, 2, 3])) + detrended_exact = array([0, 0, 0]) + assert_array_almost_equal(detrended, detrended_exact) + + def test_copy(self): + x = array([1, 1.2, 1.5, 1.6, 2.4]) + copy_array = detrend(x, overwrite_data=False) + inplace = detrend(x, overwrite_data=True) + assert_array_almost_equal(copy_array, inplace) + + @pytest.mark.parametrize('kind', ['linear', 'constant']) + @pytest.mark.parametrize('axis', [0, 1, 2]) + def test_axis(self, axis, kind): + data = np.arange(5*6*7).reshape(5, 6, 7) + detrended = detrend(data, type=kind, axis=axis) + assert detrended.shape == data.shape + + def test_bp(self): + data = [0, 1, 2] + [5, 0, -5, -10] + detrended = detrend(data, type='linear', bp=3) + assert_allclose(detrended, 0, atol=1e-14) + + # repeat with ndim > 1 and axis + data = np.asarray(data)[None, :, None] + + detrended = detrend(data, type="linear", bp=3, axis=1) + assert_allclose(detrended, 0, atol=1e-14) + + # breakpoint index > shape[axis]: raises + with assert_raises(ValueError): + detrend(data, type="linear", bp=3) + + @pytest.mark.parametrize('bp', [np.array([0, 2]), [0, 2]]) + def test_detrend_array_bp(self, bp): + # regression test for https://github.com/scipy/scipy/issues/18675 + rng = np.random.RandomState(12345) + x = rng.rand(10) + # bp = np.array([0, 2]) + + res = detrend(x, bp=bp) + res_scipy_191 = np.array([-4.44089210e-16, -2.22044605e-16, + -1.11128506e-01, -1.69470553e-01, 1.14710683e-01, 6.35468419e-02, + 3.53533144e-01, -3.67877935e-02, -2.00417675e-02, -1.94362049e-01]) + + assert_allclose(res, res_scipy_191, atol=1e-14) + + +class TestUniqueRoots: + def test_real_no_repeat(self): + p = [-1.0, -0.5, 0.3, 1.2, 10.0] + unique, multiplicity = unique_roots(p) + assert_almost_equal(unique, p, decimal=15) + assert_equal(multiplicity, np.ones(len(p))) + + def test_real_repeat(self): + p = [-1.0, -0.95, -0.89, -0.8, 0.5, 1.0, 1.05] + + unique, multiplicity = unique_roots(p, tol=1e-1, rtype='min') + assert_almost_equal(unique, [-1.0, -0.89, 0.5, 1.0], decimal=15) + assert_equal(multiplicity, [2, 2, 1, 2]) + + unique, multiplicity = unique_roots(p, tol=1e-1, rtype='max') + assert_almost_equal(unique, [-0.95, -0.8, 0.5, 1.05], decimal=15) + assert_equal(multiplicity, [2, 2, 1, 2]) + + unique, multiplicity = unique_roots(p, tol=1e-1, rtype='avg') + assert_almost_equal(unique, [-0.975, -0.845, 0.5, 1.025], decimal=15) + assert_equal(multiplicity, [2, 2, 1, 2]) + + def test_complex_no_repeat(self): + p = [-1.0, 1.0j, 0.5 + 0.5j, -1.0 - 1.0j, 3.0 + 2.0j] + unique, multiplicity = unique_roots(p) + assert_almost_equal(unique, p, decimal=15) + assert_equal(multiplicity, np.ones(len(p))) + + def test_complex_repeat(self): + p = [-1.0, -1.0 + 0.05j, -0.95 + 0.15j, -0.90 + 0.15j, 0.0, + 0.5 + 0.5j, 0.45 + 0.55j] + + unique, multiplicity = unique_roots(p, tol=1e-1, rtype='min') + assert_almost_equal(unique, [-1.0, -0.95 + 0.15j, 0.0, 0.45 + 0.55j], + decimal=15) + assert_equal(multiplicity, [2, 2, 1, 2]) + + unique, multiplicity = unique_roots(p, tol=1e-1, rtype='max') + assert_almost_equal(unique, + [-1.0 + 0.05j, -0.90 + 0.15j, 0.0, 0.5 + 0.5j], + decimal=15) + assert_equal(multiplicity, [2, 2, 1, 2]) + + unique, multiplicity = unique_roots(p, tol=1e-1, rtype='avg') + assert_almost_equal( + unique, [-1.0 + 0.025j, -0.925 + 0.15j, 0.0, 0.475 + 0.525j], + decimal=15) + assert_equal(multiplicity, [2, 2, 1, 2]) + + def test_gh_4915(self): + p = np.roots(np.convolve(np.ones(5), np.ones(5))) + true_roots = [-(-1)**(1/5), (-1)**(4/5), -(-1)**(3/5), (-1)**(2/5)] + + unique, multiplicity = unique_roots(p) + unique = np.sort(unique) + + assert_almost_equal(np.sort(unique), true_roots, decimal=7) + assert_equal(multiplicity, [2, 2, 2, 2]) + + def test_complex_roots_extra(self): + unique, multiplicity = unique_roots([1.0, 1.0j, 1.0]) + assert_almost_equal(unique, [1.0, 1.0j], decimal=15) + assert_equal(multiplicity, [2, 1]) + + unique, multiplicity = unique_roots([1, 1 + 2e-9, 1e-9 + 1j], tol=0.1) + assert_almost_equal(unique, [1.0, 1e-9 + 1.0j], decimal=15) + assert_equal(multiplicity, [2, 1]) + + def test_single_unique_root(self): + p = np.random.rand(100) + 1j * np.random.rand(100) + unique, multiplicity = unique_roots(p, 2) + assert_almost_equal(unique, [np.min(p)], decimal=15) + assert_equal(multiplicity, [100]) + + +def test_gh_22684(): + actual = signal.resample_poly(np.arange(2000, dtype=np.complex64), 6, 4) + assert actual.dtype == np.complex64 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_spectral.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_spectral.py new file mode 100644 index 0000000000000000000000000000000000000000..444aff19f025a8e3a0a4334f8cfa6ca57a7f3abc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_spectral.py @@ -0,0 +1,2059 @@ +import sys + +import numpy as np +from numpy.testing import (assert_, + assert_allclose, assert_array_equal, assert_equal, + assert_array_almost_equal_nulp, suppress_warnings) +import pytest +from pytest import raises as assert_raises + +from scipy import signal +from scipy.fft import fftfreq, rfftfreq, fft, irfft +from scipy.integrate import trapezoid +from scipy.signal import (periodogram, welch, lombscargle, coherence, + spectrogram, check_COLA, check_NOLA) +from scipy.signal.windows import hann +from scipy.signal._spectral_py import _spectral_helper + +# Compare ShortTimeFFT.stft() / ShortTimeFFT.istft() with stft() / istft(): +from scipy.signal.tests._scipy_spectral_test_shim import stft_compare as stft +from scipy.signal.tests._scipy_spectral_test_shim import istft_compare as istft +from scipy.signal.tests._scipy_spectral_test_shim import csd_compare as csd + + +class TestPeriodogram: + def test_real_onesided_even(self): + x = np.zeros(16) + x[0] = 1 + f, p = periodogram(x) + assert_allclose(f, np.linspace(0, 0.5, 9)) + q = np.ones(9) + q[0] = 0 + q[-1] /= 2.0 + q /= 8 + assert_allclose(p, q) + + def test_real_onesided_odd(self): + x = np.zeros(15) + x[0] = 1 + f, p = periodogram(x) + assert_allclose(f, np.arange(8.0)/15.0) + q = np.ones(8) + q[0] = 0 + q *= 2.0/15.0 + assert_allclose(p, q, atol=1e-15) + + def test_real_twosided(self): + x = np.zeros(16) + x[0] = 1 + f, p = periodogram(x, return_onesided=False) + assert_allclose(f, fftfreq(16, 1.0)) + q = np.full(16, 1/16.0) + q[0] = 0 + assert_allclose(p, q) + + def test_real_spectrum(self): + x = np.zeros(16) + x[0] = 1 + f, p = periodogram(x, scaling='spectrum') + g, q = periodogram(x, scaling='density') + assert_allclose(f, np.linspace(0, 0.5, 9)) + assert_allclose(p, q/16.0) + + def test_integer_even(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + f, p = periodogram(x) + assert_allclose(f, np.linspace(0, 0.5, 9)) + q = np.ones(9) + q[0] = 0 + q[-1] /= 2.0 + q /= 8 + assert_allclose(p, q) + + def test_integer_odd(self): + x = np.zeros(15, dtype=int) + x[0] = 1 + f, p = periodogram(x) + assert_allclose(f, np.arange(8.0)/15.0) + q = np.ones(8) + q[0] = 0 + q *= 2.0/15.0 + assert_allclose(p, q, atol=1e-15) + + def test_integer_twosided(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + f, p = periodogram(x, return_onesided=False) + assert_allclose(f, fftfreq(16, 1.0)) + q = np.full(16, 1/16.0) + q[0] = 0 + assert_allclose(p, q) + + def test_complex(self): + x = np.zeros(16, np.complex128) + x[0] = 1.0 + 2.0j + f, p = periodogram(x, return_onesided=False) + assert_allclose(f, fftfreq(16, 1.0)) + q = np.full(16, 5.0/16.0) + q[0] = 0 + assert_allclose(p, q) + + def test_unk_scaling(self): + assert_raises(ValueError, periodogram, np.zeros(4, np.complex128), + scaling='foo') + + @pytest.mark.skipif( + sys.maxsize <= 2**32, + reason="On some 32-bit tolerance issue" + ) + def test_nd_axis_m1(self): + x = np.zeros(20, dtype=np.float64) + x = x.reshape((2,1,10)) + x[:,:,0] = 1.0 + f, p = periodogram(x) + assert_array_equal(p.shape, (2, 1, 6)) + assert_array_almost_equal_nulp(p[0,0,:], p[1,0,:], 60) + f0, p0 = periodogram(x[0,0,:]) + assert_array_almost_equal_nulp(p0[np.newaxis,:], p[1,:], 60) + + @pytest.mark.skipif( + sys.maxsize <= 2**32, + reason="On some 32-bit tolerance issue" + ) + def test_nd_axis_0(self): + x = np.zeros(20, dtype=np.float64) + x = x.reshape((10,2,1)) + x[0,:,:] = 1.0 + f, p = periodogram(x, axis=0) + assert_array_equal(p.shape, (6,2,1)) + assert_array_almost_equal_nulp(p[:,0,0], p[:,1,0], 60) + f0, p0 = periodogram(x[:,0,0]) + assert_array_almost_equal_nulp(p0, p[:,1,0]) + + def test_window_external(self): + x = np.zeros(16) + x[0] = 1 + f, p = periodogram(x, 10, 'hann') + win = signal.get_window('hann', 16) + fe, pe = periodogram(x, 10, win) + assert_array_almost_equal_nulp(p, pe) + assert_array_almost_equal_nulp(f, fe) + win_err = signal.get_window('hann', 32) + assert_raises(ValueError, periodogram, x, + 10, win_err) # win longer than signal + + def test_padded_fft(self): + x = np.zeros(16) + x[0] = 1 + f, p = periodogram(x) + fp, pp = periodogram(x, nfft=32) + assert_allclose(f, fp[::2]) + assert_allclose(p, pp[::2]) + assert_array_equal(pp.shape, (17,)) + + def test_empty_input(self): + f, p = periodogram([]) + assert_array_equal(f.shape, (0,)) + assert_array_equal(p.shape, (0,)) + for shape in [(0,), (3,0), (0,5,2)]: + f, p = periodogram(np.empty(shape)) + assert_array_equal(f.shape, shape) + assert_array_equal(p.shape, shape) + + def test_empty_input_other_axis(self): + for shape in [(3,0), (0,5,2)]: + f, p = periodogram(np.empty(shape), axis=1) + assert_array_equal(f.shape, shape) + assert_array_equal(p.shape, shape) + + def test_short_nfft(self): + x = np.zeros(18) + x[0] = 1 + f, p = periodogram(x, nfft=16) + assert_allclose(f, np.linspace(0, 0.5, 9)) + q = np.ones(9) + q[0] = 0 + q[-1] /= 2.0 + q /= 8 + assert_allclose(p, q) + + def test_nfft_is_xshape(self): + x = np.zeros(16) + x[0] = 1 + f, p = periodogram(x, nfft=16) + assert_allclose(f, np.linspace(0, 0.5, 9)) + q = np.ones(9) + q[0] = 0 + q[-1] /= 2.0 + q /= 8 + assert_allclose(p, q) + + def test_real_onesided_even_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + f, p = periodogram(x) + assert_allclose(f, np.linspace(0, 0.5, 9)) + q = np.ones(9, 'f') + q[0] = 0 + q[-1] /= 2.0 + q /= 8 + assert_allclose(p, q) + assert_(p.dtype == q.dtype) + + def test_real_onesided_odd_32(self): + x = np.zeros(15, 'f') + x[0] = 1 + f, p = periodogram(x) + assert_allclose(f, np.arange(8.0)/15.0) + q = np.ones(8, 'f') + q[0] = 0 + q *= 2.0/15.0 + assert_allclose(p, q, atol=1e-7) + assert_(p.dtype == q.dtype) + + def test_real_twosided_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + f, p = periodogram(x, return_onesided=False) + assert_allclose(f, fftfreq(16, 1.0)) + q = np.full(16, 1/16.0, 'f') + q[0] = 0 + assert_allclose(p, q) + assert_(p.dtype == q.dtype) + + def test_complex_32(self): + x = np.zeros(16, 'F') + x[0] = 1.0 + 2.0j + f, p = periodogram(x, return_onesided=False) + assert_allclose(f, fftfreq(16, 1.0)) + q = np.full(16, 5.0/16.0, 'f') + q[0] = 0 + assert_allclose(p, q) + assert_(p.dtype == q.dtype) + + def test_shorter_window_error(self): + x = np.zeros(16) + x[0] = 1 + win = signal.get_window('hann', 10) + expected_msg = ('the size of the window must be the same size ' + 'of the input on the specified axis') + with assert_raises(ValueError, match=expected_msg): + periodogram(x, window=win) + + +class TestWelch: + def test_real_onesided_even(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8) + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.08333333, 0.15277778, 0.22222222, 0.22222222, + 0.11111111]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_real_onesided_odd(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=9) + assert_allclose(f, np.arange(5.0)/9.0) + q = np.array([0.12477455, 0.23430933, 0.17072113, 0.17072113, + 0.17072113]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_real_twosided(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.08333333, 0.07638889, 0.11111111, 0.11111111, + 0.11111111, 0.11111111, 0.11111111, 0.07638889]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_real_spectrum(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8, scaling='spectrum') + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.015625, 0.02864583, 0.04166667, 0.04166667, + 0.02083333]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_integer_onesided_even(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8) + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.08333333, 0.15277778, 0.22222222, 0.22222222, + 0.11111111]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_integer_onesided_odd(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=9) + assert_allclose(f, np.arange(5.0)/9.0) + q = np.array([0.12477455, 0.23430933, 0.17072113, 0.17072113, + 0.17072113]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_integer_twosided(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.08333333, 0.07638889, 0.11111111, 0.11111111, + 0.11111111, 0.11111111, 0.11111111, 0.07638889]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_complex(self): + x = np.zeros(16, np.complex128) + x[0] = 1.0 + 2.0j + x[8] = 1.0 + 2.0j + f, p = welch(x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.41666667, 0.38194444, 0.55555556, 0.55555556, + 0.55555556, 0.55555556, 0.55555556, 0.38194444]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_unk_scaling(self): + assert_raises(ValueError, welch, np.zeros(4, np.complex128), + scaling='foo', nperseg=4) + + def test_detrend_linear(self): + x = np.arange(10, dtype=np.float64) + 0.04 + f, p = welch(x, nperseg=10, detrend='linear') + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_no_detrending(self): + x = np.arange(10, dtype=np.float64) + 0.04 + f1, p1 = welch(x, nperseg=10, detrend=False) + f2, p2 = welch(x, nperseg=10, detrend=lambda x: x) + assert_allclose(f1, f2, atol=1e-15) + assert_allclose(p1, p2, atol=1e-15) + + def test_detrend_external(self): + x = np.arange(10, dtype=np.float64) + 0.04 + f, p = welch(x, nperseg=10, + detrend=lambda seg: signal.detrend(seg, type='l')) + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_detrend_external_nd_m1(self): + x = np.arange(40, dtype=np.float64) + 0.04 + x = x.reshape((2,2,10)) + f, p = welch(x, nperseg=10, + detrend=lambda seg: signal.detrend(seg, type='l')) + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_detrend_external_nd_0(self): + x = np.arange(20, dtype=np.float64) + 0.04 + x = x.reshape((2,1,10)) + x = np.moveaxis(x, 2, 0) + f, p = welch(x, nperseg=10, axis=0, + detrend=lambda seg: signal.detrend(seg, axis=0, type='l')) + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_nd_axis_m1(self): + x = np.arange(20, dtype=np.float64) + 0.04 + x = x.reshape((2,1,10)) + f, p = welch(x, nperseg=10) + assert_array_equal(p.shape, (2, 1, 6)) + assert_allclose(p[0,0,:], p[1,0,:], atol=1e-13, rtol=1e-13) + f0, p0 = welch(x[0,0,:], nperseg=10) + assert_allclose(p0[np.newaxis,:], p[1,:], atol=1e-13, rtol=1e-13) + + def test_nd_axis_0(self): + x = np.arange(20, dtype=np.float64) + 0.04 + x = x.reshape((10,2,1)) + f, p = welch(x, nperseg=10, axis=0) + assert_array_equal(p.shape, (6,2,1)) + assert_allclose(p[:,0,0], p[:,1,0], atol=1e-13, rtol=1e-13) + f0, p0 = welch(x[:,0,0], nperseg=10) + assert_allclose(p0, p[:,1,0], atol=1e-13, rtol=1e-13) + + def test_window_external(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = welch(x, 10, 'hann', nperseg=8) + win = signal.get_window('hann', 8) + fe, pe = welch(x, 10, win, nperseg=None) + assert_array_almost_equal_nulp(p, pe) + assert_array_almost_equal_nulp(f, fe) + assert_array_equal(fe.shape, (5,)) # because win length used as nperseg + assert_array_equal(pe.shape, (5,)) + assert_raises(ValueError, welch, x, + 10, win, nperseg=4) # because nperseg != win.shape[-1] + win_err = signal.get_window('hann', 32) + assert_raises(ValueError, welch, x, + 10, win_err, nperseg=None) # win longer than signal + + def test_empty_input(self): + f, p = welch([]) + assert_array_equal(f.shape, (0,)) + assert_array_equal(p.shape, (0,)) + for shape in [(0,), (3,0), (0,5,2)]: + f, p = welch(np.empty(shape)) + assert_array_equal(f.shape, shape) + assert_array_equal(p.shape, shape) + + def test_empty_input_other_axis(self): + for shape in [(3,0), (0,5,2)]: + f, p = welch(np.empty(shape), axis=1) + assert_array_equal(f.shape, shape) + assert_array_equal(p.shape, shape) + + def test_short_data(self): + x = np.zeros(8) + x[0] = 1 + #for string-like window, input signal length < nperseg value gives + #UserWarning, sets nperseg to x.shape[-1] + with suppress_warnings() as sup: + msg = "nperseg = 256 is greater than input length = 8, using nperseg = 8" + sup.filter(UserWarning, msg) + f, p = welch(x,window='hann') # default nperseg + f1, p1 = welch(x,window='hann', nperseg=256) # user-specified nperseg + f2, p2 = welch(x, nperseg=8) # valid nperseg, doesn't give warning + assert_allclose(f, f2) + assert_allclose(p, p2) + assert_allclose(f1, f2) + assert_allclose(p1, p2) + + def test_window_long_or_nd(self): + assert_raises(ValueError, welch, np.zeros(4), 1, np.array([1,1,1,1,1])) + assert_raises(ValueError, welch, np.zeros(4), 1, + np.arange(6).reshape((2,3))) + + def test_nondefault_noverlap(self): + x = np.zeros(64) + x[::8] = 1 + f, p = welch(x, nperseg=16, noverlap=4) + q = np.array([0, 1./12., 1./3., 1./5., 1./3., 1./5., 1./3., 1./5., + 1./6.]) + assert_allclose(p, q, atol=1e-12) + + def test_bad_noverlap(self): + assert_raises(ValueError, welch, np.zeros(4), 1, 'hann', 2, 7) + + def test_nfft_too_short(self): + assert_raises(ValueError, welch, np.ones(12), nfft=3, nperseg=4) + + def test_real_onesided_even_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8) + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.08333333, 0.15277778, 0.22222222, 0.22222222, + 0.11111111], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype) + + def test_real_onesided_odd_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=9) + assert_allclose(f, np.arange(5.0)/9.0) + q = np.array([0.12477458, 0.23430935, 0.17072113, 0.17072116, + 0.17072113], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype) + + def test_real_twosided_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.08333333, 0.07638889, 0.11111111, + 0.11111111, 0.11111111, 0.11111111, 0.11111111, + 0.07638889], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype) + + def test_complex_32(self): + x = np.zeros(16, 'F') + x[0] = 1.0 + 2.0j + x[8] = 1.0 + 2.0j + f, p = welch(x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.41666666, 0.38194442, 0.55555552, 0.55555552, + 0.55555558, 0.55555552, 0.55555552, 0.38194442], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype, + f'dtype mismatch, {p.dtype}, {q.dtype}') + + def test_padded_freqs(self): + x = np.zeros(12) + + nfft = 24 + f = fftfreq(nfft, 1.0)[:nfft//2+1] + f[-1] *= -1 + fodd, _ = welch(x, nperseg=5, nfft=nfft) + feven, _ = welch(x, nperseg=6, nfft=nfft) + assert_allclose(f, fodd) + assert_allclose(f, feven) + + nfft = 25 + f = fftfreq(nfft, 1.0)[:(nfft + 1)//2] + fodd, _ = welch(x, nperseg=5, nfft=nfft) + feven, _ = welch(x, nperseg=6, nfft=nfft) + assert_allclose(f, fodd) + assert_allclose(f, feven) + + def test_window_correction(self): + A = 20 + fs = 1e4 + nperseg = int(fs//10) + fsig = 300 + ii = int(fsig*nperseg//fs) # Freq index of fsig + + tt = np.arange(fs)/fs + x = A*np.sin(2*np.pi*fsig*tt) + + for window in ['hann', 'bartlett', ('tukey', 0.1), 'flattop']: + _, p_spec = welch(x, fs=fs, nperseg=nperseg, window=window, + scaling='spectrum') + freq, p_dens = welch(x, fs=fs, nperseg=nperseg, window=window, + scaling='density') + + # Check peak height at signal frequency for 'spectrum' + assert_allclose(p_spec[ii], A**2/2.0) + # Check integrated spectrum RMS for 'density' + assert_allclose(np.sqrt(trapezoid(p_dens, freq)), A*np.sqrt(2)/2, + rtol=1e-3) + + def test_axis_rolling(self): + np.random.seed(1234) + + x_flat = np.random.randn(1024) + _, p_flat = welch(x_flat) + + for a in range(3): + newshape = [1,]*3 + newshape[a] = -1 + x = x_flat.reshape(newshape) + + _, p_plus = welch(x, axis=a) # Positive axis index + _, p_minus = welch(x, axis=a-x.ndim) # Negative axis index + + assert_equal(p_flat, p_plus.squeeze(), err_msg=a) + assert_equal(p_flat, p_minus.squeeze(), err_msg=a-x.ndim) + + def test_average(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = welch(x, nperseg=8, average='median') + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([.1, .05, 0., 1.54074396e-33, 0.]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + assert_raises(ValueError, welch, x, nperseg=8, + average='unrecognised-average') + + +class TestCSD: + def test_pad_shorter_x(self): + x = np.zeros(8) + y = np.zeros(12) + + f = np.linspace(0, 0.5, 7) + c = np.zeros(7,dtype=np.complex128) + f1, c1 = csd(x, y, nperseg=12) + + assert_allclose(f, f1) + assert_allclose(c, c1) + + def test_pad_shorter_y(self): + x = np.zeros(12) + y = np.zeros(8) + + f = np.linspace(0, 0.5, 7) + c = np.zeros(7,dtype=np.complex128) + f1, c1 = csd(x, y, nperseg=12) + + assert_allclose(f, f1) + assert_allclose(c, c1) + + def test_real_onesided_even(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8) + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.08333333, 0.15277778, 0.22222222, 0.22222222, + 0.11111111]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_real_onesided_odd(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=9) + assert_allclose(f, np.arange(5.0)/9.0) + q = np.array([0.12477455, 0.23430933, 0.17072113, 0.17072113, + 0.17072113]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_real_twosided(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.08333333, 0.07638889, 0.11111111, 0.11111111, + 0.11111111, 0.11111111, 0.11111111, 0.07638889]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_real_spectrum(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8, scaling='spectrum') + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.015625, 0.02864583, 0.04166667, 0.04166667, + 0.02083333]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_integer_onesided_even(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8) + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.08333333, 0.15277778, 0.22222222, 0.22222222, + 0.11111111]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_integer_onesided_odd(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=9) + assert_allclose(f, np.arange(5.0)/9.0) + q = np.array([0.12477455, 0.23430933, 0.17072113, 0.17072113, + 0.17072113]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_integer_twosided(self): + x = np.zeros(16, dtype=int) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.08333333, 0.07638889, 0.11111111, 0.11111111, + 0.11111111, 0.11111111, 0.11111111, 0.07638889]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_complex(self): + x = np.zeros(16, np.complex128) + x[0] = 1.0 + 2.0j + x[8] = 1.0 + 2.0j + f, p = csd(x, x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.41666667, 0.38194444, 0.55555556, 0.55555556, + 0.55555556, 0.55555556, 0.55555556, 0.38194444]) + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + + def test_unk_scaling(self): + assert_raises(ValueError, csd, np.zeros(4, np.complex128), + np.ones(4, np.complex128), scaling='foo', nperseg=4) + + def test_detrend_linear(self): + x = np.arange(10, dtype=np.float64) + 0.04 + f, p = csd(x, x, nperseg=10, detrend='linear') + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_no_detrending(self): + x = np.arange(10, dtype=np.float64) + 0.04 + f1, p1 = csd(x, x, nperseg=10, detrend=False) + f2, p2 = csd(x, x, nperseg=10, detrend=lambda x: x) + assert_allclose(f1, f2, atol=1e-15) + assert_allclose(p1, p2, atol=1e-15) + + def test_detrend_external(self): + x = np.arange(10, dtype=np.float64) + 0.04 + f, p = csd(x, x, nperseg=10, + detrend=lambda seg: signal.detrend(seg, type='l')) + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_detrend_external_nd_m1(self): + x = np.arange(40, dtype=np.float64) + 0.04 + x = x.reshape((2,2,10)) + f, p = csd(x, x, nperseg=10, + detrend=lambda seg: signal.detrend(seg, type='l')) + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_detrend_external_nd_0(self): + x = np.arange(20, dtype=np.float64) + 0.04 + x = x.reshape((2,1,10)) + x = np.moveaxis(x, 2, 0) + f, p = csd(x, x, nperseg=10, axis=0, + detrend=lambda seg: signal.detrend(seg, axis=0, type='l')) + assert_allclose(p, np.zeros_like(p), atol=1e-15) + + def test_nd_axis_m1(self): + x = np.arange(20, dtype=np.float64) + 0.04 + x = x.reshape((2,1,10)) + f, p = csd(x, x, nperseg=10) + assert_array_equal(p.shape, (2, 1, 6)) + assert_allclose(p[0,0,:], p[1,0,:], atol=1e-13, rtol=1e-13) + f0, p0 = csd(x[0,0,:], x[0,0,:], nperseg=10) + assert_allclose(p0[np.newaxis,:], p[1,:], atol=1e-13, rtol=1e-13) + + def test_nd_axis_0(self): + x = np.arange(20, dtype=np.float64) + 0.04 + x = x.reshape((10,2,1)) + f, p = csd(x, x, nperseg=10, axis=0) + assert_array_equal(p.shape, (6,2,1)) + assert_allclose(p[:,0,0], p[:,1,0], atol=1e-13, rtol=1e-13) + f0, p0 = csd(x[:,0,0], x[:,0,0], nperseg=10) + assert_allclose(p0, p[:,1,0], atol=1e-13, rtol=1e-13) + + def test_window_external(self): + x = np.zeros(16) + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, 10, 'hann', 8) + win = signal.get_window('hann', 8) + fe, pe = csd(x, x, 10, win, nperseg=None) + assert_array_almost_equal_nulp(p, pe) + assert_array_almost_equal_nulp(f, fe) + assert_array_equal(fe.shape, (5,)) # because win length used as nperseg + assert_array_equal(pe.shape, (5,)) + assert_raises(ValueError, csd, x, x, + 10, win, nperseg=256) # because nperseg != win.shape[-1] + win_err = signal.get_window('hann', 32) + assert_raises(ValueError, csd, x, x, + 10, win_err, nperseg=None) # because win longer than signal + + def test_empty_input(self): + f, p = csd([],np.zeros(10)) + assert_array_equal(f.shape, (0,)) + assert_array_equal(p.shape, (0,)) + + f, p = csd(np.zeros(10),[]) + assert_array_equal(f.shape, (0,)) + assert_array_equal(p.shape, (0,)) + + for shape in [(0,), (3,0), (0,5,2)]: + f, p = csd(np.empty(shape), np.empty(shape)) + assert_array_equal(f.shape, shape) + assert_array_equal(p.shape, shape) + + f, p = csd(np.ones(10), np.empty((5,0))) + assert_array_equal(f.shape, (5,0)) + assert_array_equal(p.shape, (5,0)) + + f, p = csd(np.empty((5,0)), np.ones(10)) + assert_array_equal(f.shape, (5,0)) + assert_array_equal(p.shape, (5,0)) + + def test_empty_input_other_axis(self): + for shape in [(3,0), (0,5,2)]: + f, p = csd(np.empty(shape), np.empty(shape), axis=1) + assert_array_equal(f.shape, shape) + assert_array_equal(p.shape, shape) + + f, p = csd(np.empty((10,10,3)), np.zeros((10,0,1)), axis=1) + assert_array_equal(f.shape, (10,0,3)) + assert_array_equal(p.shape, (10,0,3)) + + f, p = csd(np.empty((10,0,1)), np.zeros((10,10,3)), axis=1) + assert_array_equal(f.shape, (10,0,3)) + assert_array_equal(p.shape, (10,0,3)) + + def test_short_data(self): + x = np.zeros(8) + x[0] = 1 + + #for string-like window, input signal length < nperseg value gives + #UserWarning, sets nperseg to x.shape[-1] + with suppress_warnings() as sup: + msg = "nperseg = 256 is greater than input length = 8, using nperseg = 8" + sup.filter(UserWarning, msg) + f, p = csd(x, x, window='hann') # default nperseg + f1, p1 = csd(x, x, window='hann', nperseg=256) # user-specified nperseg + f2, p2 = csd(x, x, nperseg=8) # valid nperseg, doesn't give warning + assert_allclose(f, f2) + assert_allclose(p, p2) + assert_allclose(f1, f2) + assert_allclose(p1, p2) + + def test_window_long_or_nd(self): + assert_raises(ValueError, csd, np.zeros(4), np.ones(4), 1, + np.array([1,1,1,1,1])) + assert_raises(ValueError, csd, np.zeros(4), np.ones(4), 1, + np.arange(6).reshape((2,3))) + + def test_nondefault_noverlap(self): + x = np.zeros(64) + x[::8] = 1 + f, p = csd(x, x, nperseg=16, noverlap=4) + q = np.array([0, 1./12., 1./3., 1./5., 1./3., 1./5., 1./3., 1./5., + 1./6.]) + assert_allclose(p, q, atol=1e-12) + + def test_bad_noverlap(self): + assert_raises(ValueError, csd, np.zeros(4), np.ones(4), 1, 'hann', + 2, 7) + + def test_nfft_too_short(self): + assert_raises(ValueError, csd, np.ones(12), np.zeros(12), nfft=3, + nperseg=4) + + def test_real_onesided_even_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8) + assert_allclose(f, np.linspace(0, 0.5, 5)) + q = np.array([0.08333333, 0.15277778, 0.22222222, 0.22222222, + 0.11111111], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype) + + def test_real_onesided_odd_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=9) + assert_allclose(f, np.arange(5.0)/9.0) + q = np.array([0.12477458, 0.23430935, 0.17072113, 0.17072116, + 0.17072113], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype) + + def test_real_twosided_32(self): + x = np.zeros(16, 'f') + x[0] = 1 + x[8] = 1 + f, p = csd(x, x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.08333333, 0.07638889, 0.11111111, + 0.11111111, 0.11111111, 0.11111111, 0.11111111, + 0.07638889], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype) + + def test_complex_32(self): + x = np.zeros(16, 'F') + x[0] = 1.0 + 2.0j + x[8] = 1.0 + 2.0j + f, p = csd(x, x, nperseg=8, return_onesided=False) + assert_allclose(f, fftfreq(8, 1.0)) + q = np.array([0.41666666, 0.38194442, 0.55555552, 0.55555552, + 0.55555558, 0.55555552, 0.55555552, 0.38194442], 'f') + assert_allclose(p, q, atol=1e-7, rtol=1e-7) + assert_(p.dtype == q.dtype, + f'dtype mismatch, {p.dtype}, {q.dtype}') + + def test_padded_freqs(self): + x = np.zeros(12) + y = np.ones(12) + + nfft = 24 + f = fftfreq(nfft, 1.0)[:nfft//2+1] + f[-1] *= -1 + fodd, _ = csd(x, y, nperseg=5, nfft=nfft) + feven, _ = csd(x, y, nperseg=6, nfft=nfft) + assert_allclose(f, fodd) + assert_allclose(f, feven) + + nfft = 25 + f = fftfreq(nfft, 1.0)[:(nfft + 1)//2] + fodd, _ = csd(x, y, nperseg=5, nfft=nfft) + feven, _ = csd(x, y, nperseg=6, nfft=nfft) + assert_allclose(f, fodd) + assert_allclose(f, feven) + + def test_copied_data(self): + x = np.random.randn(64) + y = x.copy() + + _, p_same = csd(x, x, nperseg=8, average='mean', + return_onesided=False) + _, p_copied = csd(x, y, nperseg=8, average='mean', + return_onesided=False) + assert_allclose(p_same, p_copied) + + _, p_same = csd(x, x, nperseg=8, average='median', + return_onesided=False) + _, p_copied = csd(x, y, nperseg=8, average='median', + return_onesided=False) + assert_allclose(p_same, p_copied) + + +class TestCoherence: + def test_identical_input(self): + x = np.random.randn(20) + y = np.copy(x) # So `y is x` -> False + + f = np.linspace(0, 0.5, 6) + C = np.ones(6) + f1, C1 = coherence(x, y, nperseg=10) + + assert_allclose(f, f1) + assert_allclose(C, C1) + + def test_phase_shifted_input(self): + x = np.random.randn(20) + y = -x + + f = np.linspace(0, 0.5, 6) + C = np.ones(6) + f1, C1 = coherence(x, y, nperseg=10) + + assert_allclose(f, f1) + assert_allclose(C, C1) + + +class TestSpectrogram: + def test_average_all_segments(self): + x = np.random.randn(1024) + + fs = 1.0 + window = ('tukey', 0.25) + nperseg = 16 + noverlap = 2 + + f, _, P = spectrogram(x, fs, window, nperseg, noverlap) + fw, Pw = welch(x, fs, window, nperseg, noverlap) + assert_allclose(f, fw) + assert_allclose(np.mean(P, axis=-1), Pw) + + def test_window_external(self): + x = np.random.randn(1024) + + fs = 1.0 + window = ('tukey', 0.25) + nperseg = 16 + noverlap = 2 + f, _, P = spectrogram(x, fs, window, nperseg, noverlap) + + win = signal.get_window(('tukey', 0.25), 16) + fe, _, Pe = spectrogram(x, fs, win, nperseg=None, noverlap=2) + assert_array_equal(fe.shape, (9,)) # because win length used as nperseg + assert_array_equal(Pe.shape, (9,73)) + assert_raises(ValueError, spectrogram, x, + fs, win, nperseg=8) # because nperseg != win.shape[-1] + win_err = signal.get_window(('tukey', 0.25), 2048) + assert_raises(ValueError, spectrogram, x, + fs, win_err, nperseg=None) # win longer than signal + + def test_short_data(self): + x = np.random.randn(1024) + fs = 1.0 + + #for string-like window, input signal length < nperseg value gives + #UserWarning, sets nperseg to x.shape[-1] + f, _, p = spectrogram(x, fs, window=('tukey',0.25)) # default nperseg + with suppress_warnings() as sup: + sup.filter(UserWarning, + "nperseg = 1025 is greater than input length = 1024, " + "using nperseg = 1024",) + f1, _, p1 = spectrogram(x, fs, window=('tukey',0.25), + nperseg=1025) # user-specified nperseg + f2, _, p2 = spectrogram(x, fs, nperseg=256) # to compare w/default + f3, _, p3 = spectrogram(x, fs, nperseg=1024) # compare w/user-spec'd + assert_allclose(f, f2) + assert_allclose(p, p2) + assert_allclose(f1, f3) + assert_allclose(p1, p3) + +class TestLombscargle: + def test_frequency(self): + """Test if frequency location of peak corresponds to frequency of + generated input signal. + """ + + # Input parameters + ampl = 2. + w = 1. + phi = 0.5 * np.pi + nin = 100 + nout = 1000 + p = 0.7 # Fraction of points to select + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.sin(w*t + phi) + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Calculate Lomb-Scargle periodogram + P = lombscargle(t, y, f) + + # Check if difference between found frequency maximum and input + # frequency is less than accuracy + delta = f[1] - f[0] + assert(w - f[np.argmax(P)] < (delta/2.)) + + # also, check that it works with weights + P = lombscargle(t, y, f, weights=np.ones_like(t, dtype=f.dtype)) + + # Check if difference between found frequency maximum and input + # frequency is less than accuracy + delta = f[1] - f[0] + assert(w - f[np.argmax(P)] < (delta/2.)) + + + def test_amplitude(self): + # Test if height of peak in unnormalized Lomb-Scargle periodogram + # corresponds to amplitude of the generated input signal. + + # Input parameters + ampl = 2. + w = 1. + phi = 0.5 * np.pi + nin = 1000 + nout = 1000 + p = 0.7 # Fraction of points to select + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.sin(w*t + phi) + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f) + + # convert to the amplitude + pgram = np.sqrt(4.0 * pgram / t.shape[0]) + + # Check if amplitude is correct (this will not exactly match, due to + # numerical differences when data is removed) + assert_allclose(pgram[f==w], ampl, rtol=5e-2) + + def test_precenter(self): + # Test if precenter gives the same result as manually precentering + # (for a very simple offset) + + # Input parameters + ampl = 2. + w = 1. + phi = 0.5 * np.pi + nin = 100 + nout = 1000 + p = 0.7 # Fraction of points to select + offset = 0.15 # Offset to be subtracted in pre-centering + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.sin(w*t + phi) + offset + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f, precenter=True) + pgram2 = lombscargle(t, y - y.mean(), f, precenter=False) + + # check if centering worked + assert_allclose(pgram, pgram2) + + # do this again, but with floating_mean=True + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f, precenter=True, floating_mean=True) + pgram2 = lombscargle(t, y - y.mean(), f, precenter=False, floating_mean=True) + + # check if centering worked + assert_allclose(pgram, pgram2) + + def test_normalize(self): + # Test normalize option of Lomb-Scarge. + + # Input parameters + ampl = 2. + w = 1. + phi = 0.5 * np.pi + nin = 100 + nout = 1000 + p = 0.7 # Fraction of points to select + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.sin(w*t + phi) + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f) + pgram2 = lombscargle(t, y, f, normalize=True) + + # Calculate the scale to convert from unnormalized to normalized + weights = np.ones_like(t)/float(t.shape[0]) + YY_hat = (weights * y * y).sum() + YY = YY_hat # correct formula for floating_mean=False + scale_to_use = 2/(YY*t.shape[0]) + + # check if normalization works as expected + assert_allclose(pgram * scale_to_use, pgram2) + assert_allclose(np.max(pgram2), 1.0) + + def test_wrong_shape(self): + + # different length t and y + t = np.linspace(0, 1, 1) + y = np.linspace(0, 1, 2) + f = np.linspace(0, 1, 3) + 0.1 + assert_raises(ValueError, lombscargle, t, y, f) + + # t is 2D, with both axes length > 1 + t = np.repeat(np.expand_dims(np.linspace(0, 1, 2), 1), 2, axis=1) + y = np.linspace(0, 1, 2) + f = np.linspace(0, 1, 3) + 0.1 + assert_raises(ValueError, lombscargle, t, y, f) + + # y is 2D, with both axes length > 1 + t = np.linspace(0, 1, 2) + y = np.repeat(np.expand_dims(np.linspace(0, 1, 2), 1), 2, axis=1) + f = np.linspace(0, 1, 3) + 0.1 + assert_raises(ValueError, lombscargle, t, y, f) + + # f is 2D, with both axes length > 1 + t = np.linspace(0, 1, 2) + y = np.linspace(0, 1, 2) + f = np.repeat(np.expand_dims(np.linspace(0, 1, 3), 1) + 0.1, 2, axis=1) + assert_raises(ValueError, lombscargle, t, y, f) + + # weights is 2D, with both axes length > 1 + t = np.linspace(0, 1, 2) + y = np.linspace(0, 1, 2) + f = np.linspace(0, 1, 3) + 0.1 + weights = np.repeat(np.expand_dims(np.linspace(0, 1, 2), 1), 2, axis=1) + assert_raises(ValueError, lombscargle, t, y, f, weights=weights) + + def test_lombscargle_atan_vs_atan2(self): + # https://github.com/scipy/scipy/issues/3787 + # This raised a ZeroDivisionError. + t = np.linspace(0, 10, 1000, endpoint=False) + y = np.sin(4*t) + f = np.linspace(0, 50, 500, endpoint=False) + 0.1 + lombscargle(t, y, f*2*np.pi) + + def test_wrong_shape_weights(self): + # Weights must be the same shape as t + + t = np.linspace(0, 1, 1) + y = np.linspace(0, 1, 1) + f = np.linspace(0, 1, 3) + 0.1 + weights = np.linspace(1, 2, 2) + assert_raises(ValueError, lombscargle, t, y, f, weights=weights) + + def test_zero_division_weights(self): + # Weights cannot sum to 0 + + t = np.zeros(1) + y = np.zeros(1) + f = np.ones(1) + weights = np.zeros(1) + assert_raises(ValueError, lombscargle, t, y, f, weights=weights) + + def test_normalize_parameter(self): + # Test the validity of the normalize parameter input + + # Input parameters + ampl = 2. + w = 1. + phi = 0 + nin = 100 + nout = 1000 + p = 0.7 # Fraction of points to select + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.sin(w*t + phi) + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # check each of the valid inputs + pgram_false = lombscargle(t, y, f, normalize=False) + pgram_true = lombscargle(t, y, f, normalize=True) + pgram_power = lombscargle(t, y, f, normalize='power') + pgram_norm = lombscargle(t, y, f, normalize='normalize') + pgram_amp = lombscargle(t, y, f, normalize='amplitude') + + # validate the results that should be the same + assert_allclose(pgram_false, pgram_power) + assert_allclose(pgram_true, pgram_norm) + + # validate that the power and norm outputs are proper wrt each other + weights = np.ones_like(y)/float(y.shape[0]) + YY_hat = (weights * y * y).sum() + YY = YY_hat # correct formula for floating_mean=False + assert_allclose(pgram_power * 2.0 / (float(t.shape[0]) * YY), pgram_norm) + + # validate that the amp output is correct for the given input + f_i = np.where(f==w)[0][0] + assert_allclose(np.abs(pgram_amp[f_i]), ampl) + + # check invalid inputs + # 1) a string that is not allowed + assert_raises(ValueError, lombscargle, t, y, f, normalize='lomb') + # 2) something besides a bool or str + assert_raises(ValueError, lombscargle, t, y, f, normalize=2) + + def test_offset_removal(self): + # Verify that the amplitude is the same, even with an offset + # must use floating_mean=True, otherwise it will not remove an offset + + # Input parameters + ampl = 2. + w = 1. + phi = 0.5 * np.pi + nin = 100 + nout = 1000 + p = 0.7 # Fraction of points to select + offset = 2.15 # Large offset + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.sin(w*t + phi) + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f, floating_mean=True) + pgram_offset = lombscargle(t, y + offset, f, floating_mean=True) + + # check if offset removal works as expected + assert_allclose(pgram, pgram_offset) + + def test_floating_mean_false(self): + # Verify that when disabling the floating_mean, the calculations are correct + + # Input parameters + ampl = 2. + w = 1. + phi = 0 + nin = 1000 + nout = 1000 + p = 0.7 # Fraction of points to select + offset = 2 # Large offset + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a cos wave for the selected times + y = ampl * np.cos(w*t + phi) + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f, normalize=True, floating_mean=False) + pgram_offset = lombscargle(t, y + offset, f, normalize=True, + floating_mean=False) + + # check if disabling floating_mean works as expected + # nearly-zero for no offset, exact value will change based on seed + assert(pgram[0] < 0.01) + # significant value with offset, exact value will change based on seed + assert(pgram_offset[0] > 0.5) + + def test_amplitude_is_correct(self): + # Verify that the amplitude is correct (when normalize='amplitude') + + # Input parameters + ampl = 2. + w = 1. + phi = 0.12 + nin = 100 + nout = 1000 + p = 0.7 # Fraction of points to select + offset = 2.15 # Large offset + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.cos(w*t + phi) + offset + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0.01, 10., nout) + + # Get the index of where the exact result should be + f_indx = np.where(f==w)[0][0] + + # Calculate Lomb-Scargle periodogram (amplitude + phase) + pgram = lombscargle(t, y, f, normalize='amplitude', floating_mean=True) + + # Check if amplitude is correct + assert_allclose(np.abs(pgram[f_indx]), ampl) + + # Check if phase is correct + # (phase angle is the negative of the phase offset) + assert_allclose(-np.angle(pgram[f_indx]), phi) + + def test_negative_weight(self): + # Test that a negative weight produces an error + + t = np.zeros(1) + y = np.zeros(1) + f = np.ones(1) + weights = -np.ones(1) + assert_raises(ValueError, lombscargle, t, y, f, weights=weights) + + def test_list_input(self): + # Test that input can be passsed in as lists and with a numerical issue + # https://github.com/scipy/scipy/issues/8787 + + t = [1.98201652e+09, 1.98201752e+09, 1.98201852e+09, 1.98201952e+09, + 1.98202052e+09, 1.98202152e+09, 1.98202252e+09, 1.98202352e+09, + 1.98202452e+09, 1.98202552e+09, 1.98202652e+09, 1.98202752e+09, + 1.98202852e+09, 1.98202952e+09, 1.98203052e+09, 1.98203152e+09, + 1.98203252e+09, 1.98203352e+09, 1.98203452e+09, 1.98203552e+09, + 1.98205452e+09, 1.98205552e+09, 1.98205652e+09, 1.98205752e+09, + 1.98205852e+09, 1.98205952e+09, 1.98206052e+09, 1.98206152e+09, + 1.98206252e+09, 1.98206352e+09, 1.98206452e+09, 1.98206552e+09, + 1.98206652e+09, 1.98206752e+09, 1.98206852e+09, 1.98206952e+09, + 1.98207052e+09, 1.98207152e+09, 1.98207252e+09, 1.98207352e+09, + 1.98209652e+09, 1.98209752e+09, 1.98209852e+09, 1.98209952e+09, + 1.98210052e+09, 1.98210152e+09, 1.98210252e+09, 1.98210352e+09, + 1.98210452e+09, 1.98210552e+09, 1.98210652e+09, 1.98210752e+09, + 1.98210852e+09, 1.98210952e+09, 1.98211052e+09, 1.98211152e+09, + 1.98211252e+09, 1.98211352e+09, 1.98211452e+09, 1.98211552e+09, + 1.98217252e+09, 1.98217352e+09, 1.98217452e+09, 1.98217552e+09, + 1.98217652e+09, 1.98217752e+09, 1.98217852e+09, 1.98217952e+09, + 1.98218052e+09, 1.98218152e+09, 1.98218252e+09, 1.98218352e+09, + 1.98218452e+09, 1.98218552e+09, 1.98218652e+09, 1.98218752e+09, + 1.98218852e+09, 1.98218952e+09, 1.98219052e+09, 1.98219152e+09, + 1.98219352e+09, 1.98219452e+09, 1.98219552e+09, 1.98219652e+09, + 1.98219752e+09, 1.98219852e+09, 1.98219952e+09, 1.98220052e+09, + 1.98220152e+09, 1.98220252e+09, 1.98220352e+09, 1.98220452e+09, + 1.98220552e+09, 1.98220652e+09, 1.98220752e+09, 1.98220852e+09, + 1.98220952e+09, 1.98221052e+09, 1.98221152e+09, 1.98221252e+09, + 1.98222752e+09, 1.98222852e+09, 1.98222952e+09, 1.98223052e+09, + 1.98223152e+09, 1.98223252e+09, 1.98223352e+09, 1.98223452e+09, + 1.98223552e+09, 1.98223652e+09, 1.98223752e+09, 1.98223852e+09, + 1.98223952e+09, 1.98224052e+09, 1.98224152e+09, 1.98224252e+09, + 1.98224352e+09, 1.98224452e+09, 1.98224552e+09, 1.98224652e+09, + 1.98224752e+09] + y = [2.97600000e+03, 3.18200000e+03, 3.74900000e+03, 4.53500000e+03, + 5.43300000e+03, 6.38000000e+03, 7.34000000e+03, 8.29200000e+03, + 9.21900000e+03, 1.01120000e+04, 1.09620000e+04, 1.17600000e+04, + 1.25010000e+04, 1.31790000e+04, 1.37900000e+04, 1.43290000e+04, + 1.47940000e+04, 1.51800000e+04, 1.54870000e+04, 1.57110000e+04, + 5.74200000e+03, 4.82300000e+03, 3.99100000e+03, 3.33600000e+03, + 2.99600000e+03, 3.08400000e+03, 3.56700000e+03, 4.30700000e+03, + 5.18200000e+03, 6.11900000e+03, 7.07900000e+03, 8.03400000e+03, + 8.97000000e+03, 9.87300000e+03, 1.07350000e+04, 1.15480000e+04, + 1.23050000e+04, 1.30010000e+04, 1.36300000e+04, 1.41890000e+04, + 6.00000000e+03, 5.06800000e+03, 4.20500000e+03, 3.49000000e+03, + 3.04900000e+03, 3.01600000e+03, 3.40400000e+03, 4.08800000e+03, + 4.93500000e+03, 5.86000000e+03, 6.81700000e+03, 7.77500000e+03, + 8.71800000e+03, 9.63100000e+03, 1.05050000e+04, 1.13320000e+04, + 1.21050000e+04, 1.28170000e+04, 1.34660000e+04, 1.40440000e+04, + 1.32730000e+04, 1.26040000e+04, 1.18720000e+04, 1.10820000e+04, + 1.02400000e+04, 9.35300000e+03, 8.43000000e+03, 7.48100000e+03, + 6.52100000e+03, 5.57000000e+03, 4.66200000e+03, 3.85400000e+03, + 3.24600000e+03, 2.97900000e+03, 3.14700000e+03, 3.68800000e+03, + 4.45900000e+03, 5.35000000e+03, 6.29400000e+03, 7.25400000e+03, + 9.13800000e+03, 1.00340000e+04, 1.08880000e+04, 1.16910000e+04, + 1.24370000e+04, 1.31210000e+04, 1.37380000e+04, 1.42840000e+04, + 1.47550000e+04, 1.51490000e+04, 1.54630000e+04, 1.56950000e+04, + 1.58430000e+04, 1.59070000e+04, 1.58860000e+04, 1.57800000e+04, + 1.55910000e+04, 1.53190000e+04, 1.49650000e+04, 1.45330000e+04, + 3.01000000e+03, 3.05900000e+03, 3.51200000e+03, 4.23400000e+03, + 5.10000000e+03, 6.03400000e+03, 6.99300000e+03, 7.95000000e+03, + 8.88800000e+03, 9.79400000e+03, 1.06600000e+04, 1.14770000e+04, + 1.22400000e+04, 1.29410000e+04, 1.35770000e+04, 1.41430000e+04, + 1.46350000e+04, 1.50500000e+04, 1.53850000e+04, 1.56400000e+04, + 1.58110000e+04] + + periods = np.linspace(400, 120, 1000) + angular_freq = 2 * np.pi / periods + + lombscargle(t, y, angular_freq, precenter=True, normalize=True) + + def test_zero_freq(self): + # Verify that function works when freqs includes 0 + # The value at f=0 will depend on the seed + + # Input parameters + ampl = 2. + w = 1. + phi = 0.12 + nin = 100 + nout = 1001 + p = 0.7 # Fraction of points to select + offset = 0 + + # Randomly select a fraction of an array with timesteps + rng = np.random.RandomState(2353425) + r = rng.rand(nin) + t = np.linspace(0.01*np.pi, 10.*np.pi, nin)[r >= p] + + # Plot a sine wave for the selected times + y = ampl * np.cos(w*t + phi) + offset + + # Define the array of frequencies for which to compute the periodogram + f = np.linspace(0, 10., nout) + + # Calculate Lomb-Scargle periodogram + pgram = lombscargle(t, y, f, normalize=True, floating_mean=True) + + # exact value will change based on seed + # testing to make sure it is very small + assert(pgram[0] < 1e-4) + + def test_simple_div_zero(self): + # these are bare-minimum examples that would, without the eps adjustments, + # cause division-by-zero errors + + # first, test with example that will cause first SS sum to be 0.0 + t = [t + 1 for t in range(0, 32)] + y = np.ones(len(t)) + freqs = [2.0*np.pi] * 2 # must have 2+ elements + lombscargle(t, y, freqs) + + # second, test with example that will cause first CC sum to be 0.0 + t = [t*4 + 1 for t in range(0, 32)] + y = np.ones(len(t)) + freqs = [np.pi/2.0] * 2 # must have 2+ elements + + lombscargle(t, y, freqs) + + +class TestSTFT: + @pytest.mark.thread_unsafe + def test_input_validation(self): + + def chk_VE(match): + """Assert for a ValueError matching regexp `match`. + + This little wrapper allows a more concise code layout. + """ + return pytest.raises(ValueError, match=match) + + # Checks for check_COLA(): + with chk_VE('nperseg must be a positive integer'): + check_COLA('hann', -10, 0) + with chk_VE('noverlap must be less than nperseg.'): + check_COLA('hann', 10, 20) + with chk_VE('window must be 1-D'): + check_COLA(np.ones((2, 2)), 10, 0) + with chk_VE('window must have length of nperseg'): + check_COLA(np.ones(20), 10, 0) + + # Checks for check_NOLA(): + with chk_VE('nperseg must be a positive integer'): + check_NOLA('hann', -10, 0) + with chk_VE('noverlap must be less than nperseg'): + check_NOLA('hann', 10, 20) + with chk_VE('window must be 1-D'): + check_NOLA(np.ones((2, 2)), 10, 0) + with chk_VE('window must have length of nperseg'): + check_NOLA(np.ones(20), 10, 0) + with chk_VE('noverlap must be a nonnegative integer'): + check_NOLA('hann', 64, -32) + + x = np.zeros(1024) + z = stft(x)[2] + + # Checks for stft(): + with chk_VE('window must be 1-D'): + stft(x, window=np.ones((2, 2))) + with chk_VE('value specified for nperseg is different ' + + 'from length of window'): + stft(x, window=np.ones(10), nperseg=256) + with chk_VE('nperseg must be a positive integer'): + stft(x, nperseg=-256) + with chk_VE('noverlap must be less than nperseg.'): + stft(x, nperseg=256, noverlap=1024) + with chk_VE('nfft must be greater than or equal to nperseg.'): + stft(x, nperseg=256, nfft=8) + + # Checks for istft(): + with chk_VE('Input stft must be at least 2d!'): + istft(x) + with chk_VE('window must be 1-D'): + istft(z, window=np.ones((2, 2))) + with chk_VE('window must have length of 256'): + istft(z, window=np.ones(10), nperseg=256) + with chk_VE('nperseg must be a positive integer'): + istft(z, nperseg=-256) + with chk_VE('noverlap must be less than nperseg.'): + istft(z, nperseg=256, noverlap=1024) + with chk_VE('nfft must be greater than or equal to nperseg.'): + istft(z, nperseg=256, nfft=8) + with pytest.warns(UserWarning, match="NOLA condition failed, " + + "STFT may not be invertible"): + istft(z, nperseg=256, noverlap=0, window='hann') + with chk_VE('Must specify differing time and frequency axes!'): + istft(z, time_axis=0, freq_axis=0) + + # Checks for _spectral_helper(): + with chk_VE("Unknown value for mode foo, must be one of: " + + r"\{'psd', 'stft'\}"): + _spectral_helper(x, x, mode='foo') + with chk_VE("x and y must be equal if mode is 'stft'"): + _spectral_helper(x[:512], x[512:], mode='stft') + with chk_VE("Unknown boundary option 'foo', must be one of: " + + r"\['even', 'odd', 'constant', 'zeros', None\]"): + _spectral_helper(x, x, boundary='foo') + + scaling = "not_valid" + with chk_VE(fr"Parameter {scaling=} not in \['spectrum', 'psd'\]!"): + stft(x, scaling=scaling) + with chk_VE(fr"Parameter {scaling=} not in \['spectrum', 'psd'\]!"): + istft(z, scaling=scaling) + + def test_check_COLA(self): + settings = [ + ('boxcar', 10, 0), + ('boxcar', 10, 9), + ('bartlett', 51, 26), + ('hann', 256, 128), + ('hann', 256, 192), + ('blackman', 300, 200), + (('tukey', 0.5), 256, 64), + ('hann', 256, 255), + ] + + for setting in settings: + msg = '{}, {}, {}'.format(*setting) + assert_equal(True, check_COLA(*setting), err_msg=msg) + + def test_check_NOLA(self): + settings_pass = [ + ('boxcar', 10, 0), + ('boxcar', 10, 9), + ('boxcar', 10, 7), + ('bartlett', 51, 26), + ('bartlett', 51, 10), + ('hann', 256, 128), + ('hann', 256, 192), + ('hann', 256, 37), + ('blackman', 300, 200), + ('blackman', 300, 123), + (('tukey', 0.5), 256, 64), + (('tukey', 0.5), 256, 38), + ('hann', 256, 255), + ('hann', 256, 39), + ] + for setting in settings_pass: + msg = '{}, {}, {}'.format(*setting) + assert_equal(True, check_NOLA(*setting), err_msg=msg) + + w_fail = np.ones(16) + w_fail[::2] = 0 + settings_fail = [ + (w_fail, len(w_fail), len(w_fail) // 2), + ('hann', 64, 0), + ] + for setting in settings_fail: + msg = '{}, {}, {}'.format(*setting) + assert_equal(False, check_NOLA(*setting), err_msg=msg) + + def test_average_all_segments(self): + rng = np.random.RandomState(1234) + x = rng.randn(1024) + + fs = 1.0 + window = 'hann' + nperseg = 16 + noverlap = 8 + + # Compare twosided, because onesided welch doubles non-DC terms to + # account for power at negative frequencies. stft doesn't do this, + # because it breaks invertibility. + f, _, Z = stft(x, fs, window, nperseg, noverlap, padded=False, + return_onesided=False, boundary=None) + fw, Pw = welch(x, fs, window, nperseg, noverlap, return_onesided=False, + scaling='spectrum', detrend=False) + + assert_allclose(f, fw) + assert_allclose(np.mean(np.abs(Z)**2, axis=-1), Pw) + + def test_permute_axes(self): + rng = np.random.RandomState(1234) + x = rng.randn(1024) + + fs = 1.0 + window = 'hann' + nperseg = 16 + noverlap = 8 + + f1, t1, Z1 = stft(x, fs, window, nperseg, noverlap) + f2, t2, Z2 = stft(x.reshape((-1, 1, 1)), fs, window, nperseg, noverlap, + axis=0) + + t3, x1 = istft(Z1, fs, window, nperseg, noverlap) + t4, x2 = istft(Z2.T, fs, window, nperseg, noverlap, time_axis=0, + freq_axis=-1) + + assert_allclose(f1, f2) + assert_allclose(t1, t2) + assert_allclose(t3, t4) + assert_allclose(Z1, Z2[:, 0, 0, :]) + assert_allclose(x1, x2[:, 0, 0]) + + @pytest.mark.parametrize('scaling', ['spectrum', 'psd']) + def test_roundtrip_real(self, scaling): + rng = np.random.RandomState(1234) + + settings = [ + ('boxcar', 100, 10, 0), # Test no overlap + ('boxcar', 100, 10, 9), # Test high overlap + ('bartlett', 101, 51, 26), # Test odd nperseg + ('hann', 1024, 256, 128), # Test defaults + (('tukey', 0.5), 1152, 256, 64), # Test Tukey + ('hann', 1024, 256, 255), # Test overlapped hann + ] + + for window, N, nperseg, noverlap in settings: + t = np.arange(N) + x = 10*rng.randn(t.size) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=False, + scaling=scaling) + + tr, xr = istft(zz, nperseg=nperseg, noverlap=noverlap, + window=window, scaling=scaling) + + msg = f'{window}, {noverlap}' + assert_allclose(t, tr, err_msg=msg) + assert_allclose(x, xr, err_msg=msg) + + @pytest.mark.thread_unsafe + def test_roundtrip_not_nola(self): + rng = np.random.RandomState(1234) + + w_fail = np.ones(16) + w_fail[::2] = 0 + settings = [ + (w_fail, 256, len(w_fail), len(w_fail) // 2), + ('hann', 256, 64, 0), + ] + + for window, N, nperseg, noverlap in settings: + msg = f'{window}, {N}, {nperseg}, {noverlap}' + assert not check_NOLA(window, nperseg, noverlap), msg + + t = np.arange(N) + x = 10 * rng.randn(t.size) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=True, + boundary='zeros') + with pytest.warns(UserWarning, match='NOLA'): + tr, xr = istft(zz, nperseg=nperseg, noverlap=noverlap, + window=window, boundary=True) + + assert np.allclose(t, tr[:len(t)]), msg + assert not np.allclose(x, xr[:len(x)]), msg + + def test_roundtrip_nola_not_cola(self): + rng = np.random.RandomState(1234) + + settings = [ + ('boxcar', 100, 10, 3), # NOLA True, COLA False + ('bartlett', 101, 51, 37), # NOLA True, COLA False + ('hann', 1024, 256, 127), # NOLA True, COLA False + (('tukey', 0.5), 1152, 256, 14), # NOLA True, COLA False + ('hann', 1024, 256, 5), # NOLA True, COLA False + ] + + for window, N, nperseg, noverlap in settings: + msg = f'{window}, {nperseg}, {noverlap}' + assert check_NOLA(window, nperseg, noverlap), msg + assert not check_COLA(window, nperseg, noverlap), msg + + t = np.arange(N) + x = 10 * rng.randn(t.size) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=True, + boundary='zeros') + + tr, xr = istft(zz, nperseg=nperseg, noverlap=noverlap, + window=window, boundary=True) + + msg = f'{window}, {noverlap}' + assert_allclose(t, tr[:len(t)], err_msg=msg) + assert_allclose(x, xr[:len(x)], err_msg=msg) + + def test_roundtrip_float32(self): + rng = np.random.RandomState(1234) + + settings = [('hann', 1024, 256, 128)] + + for window, N, nperseg, noverlap in settings: + t = np.arange(N) + x = 10*rng.randn(t.size) + x = x.astype(np.float32) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=False) + + tr, xr = istft(zz, nperseg=nperseg, noverlap=noverlap, + window=window) + + msg = f'{window}, {noverlap}' + assert_allclose(t, t, err_msg=msg) + assert_allclose(x, xr, err_msg=msg, rtol=1e-4, atol=1e-5) + assert_(x.dtype == xr.dtype) + + @pytest.mark.thread_unsafe + @pytest.mark.parametrize('scaling', ['spectrum', 'psd']) + def test_roundtrip_complex(self, scaling): + rng = np.random.RandomState(1234) + + settings = [ + ('boxcar', 100, 10, 0), # Test no overlap + ('boxcar', 100, 10, 9), # Test high overlap + ('bartlett', 101, 51, 26), # Test odd nperseg + ('hann', 1024, 256, 128), # Test defaults + (('tukey', 0.5), 1152, 256, 64), # Test Tukey + ('hann', 1024, 256, 255), # Test overlapped hann + ] + + for window, N, nperseg, noverlap in settings: + t = np.arange(N) + x = 10*rng.randn(t.size) + 10j*rng.randn(t.size) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=False, + return_onesided=False, scaling=scaling) + + tr, xr = istft(zz, nperseg=nperseg, noverlap=noverlap, + window=window, input_onesided=False, + scaling=scaling) + + msg = f'{window}, {nperseg}, {noverlap}' + assert_allclose(t, tr, err_msg=msg) + assert_allclose(x, xr, err_msg=msg) + + # Check that asking for onesided switches to twosided + with suppress_warnings() as sup: + sup.filter(UserWarning, + "Input data is complex, switching to return_onesided=False") + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=False, + return_onesided=True, scaling=scaling) + + tr, xr = istft(zz, nperseg=nperseg, noverlap=noverlap, + window=window, input_onesided=False, scaling=scaling) + + msg = f'{window}, {nperseg}, {noverlap}' + assert_allclose(t, tr, err_msg=msg) + assert_allclose(x, xr, err_msg=msg) + + def test_roundtrip_boundary_extension(self): + rng = np.random.RandomState(1234) + + # Test against boxcar, since window is all ones, and thus can be fully + # recovered with no boundary extension + + settings = [ + ('boxcar', 100, 10, 0), # Test no overlap + ('boxcar', 100, 10, 9), # Test high overlap + ] + + for window, N, nperseg, noverlap in settings: + t = np.arange(N) + x = 10*rng.randn(t.size) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=True, + boundary=None) + + _, xr = istft(zz, noverlap=noverlap, window=window, boundary=False) + + for boundary in ['even', 'odd', 'constant', 'zeros']: + _, _, zz_ext = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=True, + boundary=boundary) + + _, xr_ext = istft(zz_ext, noverlap=noverlap, window=window, + boundary=True) + + msg = f'{window}, {noverlap}, {boundary}' + assert_allclose(x, xr, err_msg=msg) + assert_allclose(x, xr_ext, err_msg=msg) + + def test_roundtrip_padded_signal(self): + rng = np.random.RandomState(1234) + + settings = [ + ('boxcar', 101, 10, 0), + ('hann', 1000, 256, 128), + ] + + for window, N, nperseg, noverlap in settings: + t = np.arange(N) + x = 10*rng.randn(t.size) + + _, _, zz = stft(x, nperseg=nperseg, noverlap=noverlap, + window=window, detrend=None, padded=True) + + tr, xr = istft(zz, noverlap=noverlap, window=window) + + msg = f'{window}, {noverlap}' + # Account for possible zero-padding at the end + assert_allclose(t, tr[:t.size], err_msg=msg) + assert_allclose(x, xr[:x.size], err_msg=msg) + + def test_roundtrip_padded_FFT(self): + rng = np.random.RandomState(1234) + + settings = [ + ('hann', 1024, 256, 128, 512), + ('hann', 1024, 256, 128, 501), + ('boxcar', 100, 10, 0, 33), + (('tukey', 0.5), 1152, 256, 64, 1024), + ] + + for window, N, nperseg, noverlap, nfft in settings: + t = np.arange(N) + x = 10*rng.randn(t.size) + xc = x*np.exp(1j*np.pi/4) + + # real signal + _, _, z = stft(x, nperseg=nperseg, noverlap=noverlap, nfft=nfft, + window=window, detrend=None, padded=True) + + # complex signal + _, _, zc = stft(xc, nperseg=nperseg, noverlap=noverlap, nfft=nfft, + window=window, detrend=None, padded=True, + return_onesided=False) + + tr, xr = istft(z, nperseg=nperseg, noverlap=noverlap, nfft=nfft, + window=window) + + tr, xcr = istft(zc, nperseg=nperseg, noverlap=noverlap, nfft=nfft, + window=window, input_onesided=False) + + msg = f'{window}, {noverlap}' + assert_allclose(t, tr, err_msg=msg) + assert_allclose(x, xr, err_msg=msg) + assert_allclose(xc, xcr, err_msg=msg) + + def test_axis_rolling(self): + rng = np.random.RandomState(1234) + + x_flat = rng.randn(1024) + _, _, z_flat = stft(x_flat) + + for a in range(3): + newshape = [1,]*3 + newshape[a] = -1 + x = x_flat.reshape(newshape) + + _, _, z_plus = stft(x, axis=a) # Positive axis index + _, _, z_minus = stft(x, axis=a-x.ndim) # Negative axis index + + assert_equal(z_flat, z_plus.squeeze(), err_msg=a) + assert_equal(z_flat, z_minus.squeeze(), err_msg=a-x.ndim) + + # z_flat has shape [n_freq, n_time] + + # Test vs. transpose + _, x_transpose_m = istft(z_flat.T, time_axis=-2, freq_axis=-1) + _, x_transpose_p = istft(z_flat.T, time_axis=0, freq_axis=1) + + assert_allclose(x_flat, x_transpose_m, err_msg='istft transpose minus') + assert_allclose(x_flat, x_transpose_p, err_msg='istft transpose plus') + + def test_roundtrip_scaling(self): + """Verify behavior of scaling parameter. """ + # Create 1024 sample cosine signal with amplitude 2: + X = np.zeros(513, dtype=complex) + X[256] = 1024 + x = np.fft.irfft(X) + power_x = sum(x**2) / len(x) # power of signal x is 2 + + # Calculate magnitude-scaled STFT: + Zs = stft(x, boundary='even', scaling='spectrum')[2] + + # Test round trip: + x1 = istft(Zs, boundary=True, scaling='spectrum')[1] + assert_allclose(x1, x) + + # For a Hann-windowed 256 sample length FFT, we expect a peak at + # frequency 64 (since it is 1/4 the length of X) with a height of 1 + # (half the amplitude). A Hann window of a perfectly centered sine has + # the magnitude [..., 0, 0, 0.5, 1, 0.5, 0, 0, ...]. + # Note that in this case the 'even' padding works for the beginning + # but not for the end of the STFT. + assert_allclose(abs(Zs[63, :-1]), 0.5) + assert_allclose(abs(Zs[64, :-1]), 1) + assert_allclose(abs(Zs[65, :-1]), 0.5) + # All other values should be zero: + Zs[63:66, :-1] = 0 + # Note since 'rtol' does not have influence here, atol needs to be set: + assert_allclose(Zs[:, :-1], 0, atol=np.finfo(Zs.dtype).resolution) + + # Calculate two-sided psd-scaled STFT: + # - using 'even' padding since signal is axis symmetric - this ensures + # stationary behavior on the boundaries + # - using the two-sided transform allows determining the spectral + # power by `sum(abs(Zp[:, k])**2) / len(f)` for the k-th time slot. + Zp = stft(x, return_onesided=False, boundary='even', scaling='psd')[2] + + # Calculate spectral power of Zd by summing over the frequency axis: + psd_Zp = np.sum(Zp.real**2 + Zp.imag**2, axis=0) / Zp.shape[0] + # Spectral power of Zp should be equal to the signal's power: + assert_allclose(psd_Zp, power_x) + + # Test round trip: + x1 = istft(Zp, input_onesided=False, boundary=True, scaling='psd')[1] + assert_allclose(x1, x) + + # The power of the one-sided psd-scaled STFT can be determined + # analogously (note that the two sides are not of equal shape): + Zp0 = stft(x, return_onesided=True, boundary='even', scaling='psd')[2] + + # Since x is real, its Fourier transform is conjugate symmetric, i.e., + # the missing 'second side' can be expressed through the 'first side': + Zp1 = np.conj(Zp0[-2:0:-1, :]) # 'second side' is conjugate reversed + assert_allclose(Zp[:129, :], Zp0) + assert_allclose(Zp[129:, :], Zp1) + + # Calculate the spectral power: + s2 = (np.sum(Zp0.real ** 2 + Zp0.imag ** 2, axis=0) + + np.sum(Zp1.real ** 2 + Zp1.imag ** 2, axis=0)) + psd_Zp01 = s2 / (Zp0.shape[0] + Zp1.shape[0]) + assert_allclose(psd_Zp01, power_x) + + # Test round trip: + x1 = istft(Zp0, input_onesided=True, boundary=True, scaling='psd')[1] + assert_allclose(x1, x) + + +class TestSampledSpectralRepresentations: + """Check energy/power relations from `Spectral Analysis` section in the user guide. + + A 32 sample cosine signal is used to compare the numerical to the expected results + stated in :ref:`tutorial_SpectralAnalysis` in + file ``doc/source/tutorial/signal.rst`` + """ + n: int = 32 #: number of samples + T: float = 1/16 #: sampling interval + a_ref: float = 3 #: amplitude of reference + l_a: int = 3 #: index in fft for defining frequency of test signal + + x_ref: np.ndarray #: reference signal + X_ref: np.ndarray #: two-sided FFT of x_ref + E_ref: float #: energy of signal + P_ref: float #: power of signal + + def setup_method(self): + """Create Cosine signal with amplitude a from spectrum. """ + f = rfftfreq(self.n, self.T) + X_ref = np.zeros_like(f) + self.l_a = 3 + X_ref[self.l_a] = self.a_ref/2 * self.n # set amplitude + self.x_ref = irfft(X_ref) + self.X_ref = fft(self.x_ref) + + # Closed form expression for continuous-time signal: + self.E_ref = self.tau * self.a_ref**2 / 2 # energy of signal + self.P_ref = self.a_ref**2 / 2 # power of signal + + @property + def tau(self) -> float: + """Duration of signal. """ + return self.n * self.T + + @property + def delta_f(self) -> float: + """Bin width """ + return 1 / (self.n * self.T) + + def test_reference_signal(self): + """Test energy and power formulas. """ + # Verify that amplitude is a: + assert_allclose(2*self.a_ref, np.ptp(self.x_ref), rtol=0.1) + # Verify that energy expression for sampled signal: + assert_allclose(self.T * sum(self.x_ref ** 2), self.E_ref) + + # Verify that spectral energy and power formulas are correct: + sum_X_ref_squared = sum(self.X_ref.real**2 + self.X_ref.imag**2) + assert_allclose(self.T/self.n * sum_X_ref_squared, self.E_ref) + assert_allclose(1/self.n**2 * sum_X_ref_squared, self.P_ref) + + def test_windowed_DFT(self): + """Verify spectral representations of windowed DFT. + + Furthermore, the scalings of `periodogram` and `welch` are verified. + """ + w = hann(self.n, sym=False) + c_amp, c_rms = abs(sum(w)), np.sqrt(sum(w.real**2 + w.imag**2)) + Xw = fft(self.x_ref*w) # unnormalized windowed DFT + + # Verify that the *spectrum* peak is consistent: + assert_allclose(self.tau * Xw[self.l_a] / c_amp, self.a_ref * self.tau / 2) + # Verify that the *amplitude spectrum* peak is consistent: + assert_allclose(Xw[self.l_a] / c_amp, self.a_ref/2) + + # Verify spectral power/energy equals signal's power/energy: + X_ESD = self.tau * self.T * abs(Xw / c_rms)**2 # Energy Spectral Density + X_PSD = self.T * abs(Xw / c_rms)**2 # Power Spectral Density + assert_allclose(self.delta_f * sum(X_ESD), self.E_ref) + assert_allclose(self.delta_f * sum(X_PSD), self.P_ref) + + # Verify scalings of periodogram: + kw = dict(fs=1/self.T, window=w, detrend=False, return_onesided=False) + _, P_mag = periodogram(self.x_ref, scaling='spectrum', **kw) + _, P_psd = periodogram(self.x_ref, scaling='density', **kw) + + # Verify that periodogram calculates a squared magnitude spectrum: + float_res = np.finfo(P_mag.dtype).resolution + assert_allclose(P_mag, abs(Xw/c_amp)**2, atol=float_res*max(P_mag)) + # Verify that periodogram calculates a PSD: + assert_allclose(P_psd, X_PSD, atol=float_res*max(P_psd)) + + # Ensure that scaling of welch is the same as of periodogram: + kw = dict(nperseg=len(self.x_ref), noverlap=0, **kw) + assert_allclose(welch(self.x_ref, scaling='spectrum', **kw)[1], P_mag, + atol=float_res*max(P_mag)) + assert_allclose(welch(self.x_ref, scaling='density', **kw)[1], P_psd, + atol=float_res*max(P_psd)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_splines.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_splines.py new file mode 100644 index 0000000000000000000000000000000000000000..a31e0eeca83a9f434f96c3460347cf31379091cf --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_splines.py @@ -0,0 +1,362 @@ +# pylint: disable=missing-docstring +import numpy as np +import pytest +from scipy._lib._array_api import xp_assert_close + +from scipy.signal._spline import ( + symiirorder1_ic, symiirorder2_ic_fwd, symiirorder2_ic_bwd) +from scipy.signal import symiirorder1, symiirorder2 + + +def _compute_symiirorder2_bwd_hs(k, cs, rsq, omega): + cssq = cs * cs + k = np.abs(k) + rsupk = np.power(rsq, k / 2.0) + + c0 = (cssq * (1.0 + rsq) / (1.0 - rsq) / + (1 - 2 * rsq * np.cos(2 * omega) + rsq * rsq)) + gamma = (1.0 - rsq) / (1.0 + rsq) / np.tan(omega) + return c0 * rsupk * (np.cos(omega * k) + gamma * np.sin(omega * k)) + + +class TestSymIIR: + @pytest.mark.parametrize( + 'dtype', [np.float32, np.float64, np.complex64, np.complex128]) + @pytest.mark.parametrize('precision', [-1.0, 0.7, 0.5, 0.25, 0.0075]) + def test_symiir1_ic(self, dtype, precision): + c_precision = precision + if precision <= 0.0 or precision > 1.0: + if dtype in {np.float32, np.complex64}: + c_precision = 1e-6 + else: + c_precision = 1e-11 + + # Symmetrical initial conditions for a IIR filter of order 1 are: + # x[0] + z1 * \sum{k = 0}^{n - 1} x[k] * z1^k + + # Check the initial condition for a low-pass filter + # with coefficient b = 0.85 on a step signal. The initial condition is + # a geometric series: 1 + b * \sum_{k = 0}^{n - 1} u[k] b^k. + + # Finding the initial condition corresponds to + # 1. Computing the index n such that b**n < precision, which + # corresponds to ceil(log(precision) / log(b)) + # 2. Computing the geometric series until n, this can be computed + # using the partial sum formula: (1 - b**n) / (1 - b) + # This holds due to the input being a step signal. + b = 0.85 + n_exp = int(np.ceil(np.log(c_precision) / np.log(b))) + expected = np.asarray([[(1 - b ** n_exp) / (1 - b)]], dtype=dtype) + expected = 1 + b * expected + + # Create a step signal of size n + 1 + x = np.ones(n_exp + 1, dtype=dtype) + xp_assert_close(symiirorder1_ic(x, b, precision), expected, + atol=2e-6, rtol=2e-7) + + # Check the conditions for a exponential decreasing signal with base 2. + # Same conditions hold, as the product of 0.5^n * 0.85^n is + # still a geometric series + b_d = np.asarray(b, dtype=dtype) + expected = np.asarray( + [[(1 - (0.5 * b_d) ** n_exp) / (1 - (0.5 * b_d))]], dtype=dtype) + expected = 1 + b_d * expected + + # Create an exponential decreasing signal of size n + 1 + x = 2 ** -np.arange(n_exp + 1, dtype=dtype) + xp_assert_close(symiirorder1_ic(x, b, precision), expected, + atol=2e-6, rtol=2e-7) + + def test_symiir1_ic_fails(self): + # Test that symiirorder1_ic fails whenever \sum_{n = 1}^{n} b^n > eps + b = 0.85 + # Create a step signal of size 100 + x = np.ones(100, dtype=np.float64) + + # Compute the closed form for the geometrical series + precision = 1 / (1 - b) + pytest.raises(ValueError, symiirorder1_ic, x, b, precision) + + # Test that symiirorder1_ic fails when |z1| >= 1 + pytest.raises(ValueError, symiirorder1_ic, x, 1.0, -1) + pytest.raises(ValueError, symiirorder1_ic, x, 2.0, -1) + + @pytest.mark.parametrize( + 'dtype', [np.float32, np.float64, np.complex64, np.complex128]) + @pytest.mark.parametrize('precision', [-1.0, 0.7, 0.5, 0.25, 0.0075]) + def test_symiir1(self, dtype, precision): + c_precision = precision + if precision <= 0.0 or precision > 1.0: + if dtype in {np.float32, np.complex64}: + c_precision = 1e-6 + else: + c_precision = 1e-11 + + # Test for a low-pass filter with c0 = 0.15 and z1 = 0.85 + # using an unit step over 200 samples. + c0 = 0.15 + z1 = 0.85 + n = 200 + signal = np.ones(n, dtype=dtype) + + # Find the initial condition. See test_symiir1_ic for a detailed + # explanation + n_exp = int(np.ceil(np.log(c_precision) / np.log(z1))) + initial = np.asarray((1 - z1 ** n_exp) / (1 - z1), dtype=dtype) + initial = 1 + z1 * initial + + # Forward pass + # The transfer function for the system 1 / (1 - z1 * z^-1) when + # applied to an unit step with initial conditions y0 is + # 1 / (1 - z1 * z^-1) * (z^-1 / (1 - z^-1) + y0) + + # Solving the inverse Z-transform for the given expression yields: + # y[n] = y0 * z1**n * u[n] + + # -z1 / (1 - z1) * z1**(k - 1) * u[k - 1] + + # 1 / (1 - z1) * u[k - 1] + # d is the Kronecker delta function, and u is the unit step + + # y0 * z1**n * u[n] + pos = np.arange(n, dtype=dtype) + comp1 = initial * z1**pos + + # -z1 / (1 - z1) * z1**(k - 1) * u[k - 1] + comp2 = np.zeros(n, dtype=dtype) + comp2[1:] = -z1 / (1 - z1) * z1**pos[:-1] + + # 1 / (1 - z1) * u[k - 1] + comp3 = np.zeros(n, dtype=dtype) + comp3[1:] = 1 / (1 - z1) + + expected_fwd = comp1 + comp2 + comp3 + + # Reverse condition + sym_cond = -c0 / (z1 - 1.0) * expected_fwd[-1] + + # Backward pass + # The transfer function for the forward result is equivalent to + # the forward system times c0 / (1 - z1 * z). + + # Computing a closed form for the complete expression is difficult + # The result will be computed iteratively from the difference equation + exp_out = np.zeros(n, dtype=dtype) + exp_out[0] = sym_cond + + for i in range(1, n): + exp_out[i] = c0 * expected_fwd[n - 1 - i] + z1 * exp_out[i - 1] + + exp_out = exp_out[::-1] + + out = symiirorder1(signal, c0, z1, precision) + xp_assert_close(out, exp_out, atol=4e-6, rtol=6e-7) + + @pytest.mark.parametrize('dtype', ['float32', 'float64']) + def test_symiir1_values(self, dtype): + rng = np.random.RandomState(1234) + dtype = getattr(np, dtype) + s = rng.uniform(size=16).astype(dtype) + res = symiirorder1(s, 0.5, 0.1) + + # values from scipy 1.9.1 + exp_res = np.array([0.14387447, 0.35166047, 0.29735238, 0.46295986, 0.45174927, + 0.19982875, 0.20355805, 0.47378628, 0.57232247, 0.51597393, + 0.25935107, 0.31438554, 0.41096728, 0.4190693 , 0.25812255, + 0.33671467], dtype=res.dtype) + assert res.dtype == dtype + atol = {np.float64: 1e-15, np.float32: 1e-7}[dtype] + xp_assert_close(res, exp_res, atol=atol) + + s = s + 1j*s + res = symiirorder1(s, 0.5, 0.1) + assert res.dtype == np.complex64 if dtype == np.float32 else np.complex128 + xp_assert_close(res, exp_res + 1j*exp_res, atol=atol) + + @pytest.mark.parametrize( + 'dtype', ['float32', 'float64']) + @pytest.mark.parametrize('precision', [-1.0, 0.7, 0.5, 0.25, 0.0075]) + def test_symiir2_initial_fwd(self, dtype, precision): + dtype = getattr(np, dtype) + c_precision = precision + if precision <= 0.0 or precision > 1.0: + if dtype in {np.float32, np.complex64}: + c_precision = 1e-6 + else: + c_precision = 1e-11 + + # Compute the initial conditions for a order-two symmetrical low-pass + # filter with r = 0.5 and omega = pi / 3 for an unit step input. + r = np.asarray(0.5, dtype=dtype) + omega = np.asarray(np.pi / 3.0, dtype=dtype) + cs = 1 - 2 * r * np.cos(omega) + r**2 + + # The index n for the initial condition is bound from 0 to the + # first position where sin(omega * (n + 2)) = 0 => omega * (n + 2) = pi + # For omega = pi / 3, the maximum initial condition occurs when + # sqrt(3) / 2 * r**n < precision. + # => n = log(2 * sqrt(3) / 3 * precision) / log(r) + ub = np.ceil(np.log(c_precision / np.sin(omega)) / np.log(c_precision)) + lb = np.ceil(np.pi / omega) - 2 + n_exp = min(ub, lb) + + # The forward initial condition for a filter of order two is: + # \frac{cs}{\sin(\omega)} \sum_{n = 0}^{N - 1} { + # r^(n + 1) \sin{\omega(n + 2)}} + cs + # The closed expression for this sum is: + # s[n] = 2 * r * np.cos(omega) - + # r**2 - r**(n + 2) * np.sin(omega * (n + 3)) / np.sin(omega) + + # r**(n + 3) * np.sin(omega * (n + 2)) / np.sin(omega) + cs + fwd_initial_1 = ( + cs + + 2 * r * np.cos(omega) - + r**2 - + r**(n_exp + 2) * np.sin(omega * (n_exp + 3)) / np.sin(omega) + + r**(n_exp + 3) * np.sin(omega * (n_exp + 2)) / np.sin(omega)) + + # The second initial condition is given by + # s[n] = 1 / np.sin(omega) * ( + # r**2 * np.sin(3 * omega) - + # r**3 * np.sin(2 * omega) - + # r**(n + 3) * np.sin(omega * (n + 4)) + + # r**(n + 4) * np.sin(omega * (n + 3))) + ub = np.ceil(np.log(c_precision / np.sin(omega)) / np.log(c_precision)) + lb = np.ceil(np.pi / omega) - 3 + n_exp = min(ub, lb) + + fwd_initial_2 = ( + cs + cs * 2 * r * np.cos(omega) + + (r**2 * np.sin(3 * omega) - + r**3 * np.sin(2 * omega) - + r**(n_exp + 3) * np.sin(omega * (n_exp + 4)) + + r**(n_exp + 4) * np.sin(omega * (n_exp + 3))) / np.sin(omega)) + + expected = np.r_[fwd_initial_1, fwd_initial_2][None, :] + expected = expected.astype(dtype) + + n = 100 + signal = np.ones(n, dtype=dtype) + + out = symiirorder2_ic_fwd(signal, r, omega, precision) + xp_assert_close(out, expected, atol=4e-6, rtol=6e-7) + + @pytest.mark.parametrize( + 'dtype', [np.float32, np.float64]) + @pytest.mark.parametrize('precision', [-1.0, 0.7, 0.5, 0.25, 0.0075]) + def test_symiir2_initial_bwd(self, dtype, precision): + c_precision = precision + if precision <= 0.0 or precision > 1.0: + if dtype in {np.float32, np.complex64}: + c_precision = 1e-6 + else: + c_precision = 1e-11 + + r = np.asarray(0.5, dtype=dtype) + omega = np.asarray(np.pi / 3.0, dtype=dtype) + cs = 1 - 2 * r * np.cos(omega) + r * r + a2 = 2 * r * np.cos(omega) + a3 = -r * r + + n = 100 + signal = np.ones(n, dtype=dtype) + + # Compute initial forward conditions + ic = symiirorder2_ic_fwd(signal, r, omega, precision) + out = np.zeros(n + 2, dtype=dtype) + out[:2] = ic[0] + + # Apply the forward system cs / (1 - a2 * z^-1 - a3 * z^-2)) + for i in range(2, n + 2): + out[i] = cs * signal[i - 2] + a2 * out[i - 1] + a3 * out[i - 2] + + # Find the backward initial conditions + ic2 = np.zeros(2, dtype=dtype) + idx = np.arange(n) + + diff = (_compute_symiirorder2_bwd_hs(idx, cs, r * r, omega) + + _compute_symiirorder2_bwd_hs(idx + 1, cs, r * r, omega)) + ic2_0_all = np.cumsum(diff * out[:1:-1]) + pos = np.where(diff ** 2 < c_precision)[0] + ic2[0] = ic2_0_all[pos[0]] + + diff = (_compute_symiirorder2_bwd_hs(idx - 1, cs, r * r, omega) + + _compute_symiirorder2_bwd_hs(idx + 2, cs, r * r, omega)) + ic2_1_all = np.cumsum(diff * out[:1:-1]) + pos = np.where(diff ** 2 < c_precision)[0] + ic2[1] = ic2_1_all[pos[0]] + + out_ic = symiirorder2_ic_bwd(out, r, omega, precision)[0] + xp_assert_close(out_ic, ic2, atol=4e-6, rtol=6e-7) + + @pytest.mark.parametrize( + 'dtype', [np.float32, np.float64]) + @pytest.mark.parametrize('precision', [-1.0, 0.7, 0.5, 0.25, 0.0075]) + def test_symiir2(self, dtype, precision): + r = np.asarray(0.5, dtype=dtype) + omega = np.asarray(np.pi / 3.0, dtype=dtype) + cs = 1 - 2 * r * np.cos(omega) + r * r + a2 = 2 * r * np.cos(omega) + a3 = -r * r + + n = 100 + signal = np.ones(n, dtype=dtype) + + # Compute initial forward conditions + ic = symiirorder2_ic_fwd(signal, r, omega, precision) + out1 = np.zeros(n + 2, dtype=dtype) + out1[:2] = ic[0] + + # Apply the forward system cs / (1 - a2 * z^-1 - a3 * z^-2)) + for i in range(2, n + 2): + out1[i] = cs * signal[i - 2] + a2 * out1[i - 1] + a3 * out1[i - 2] + + # Find the backward initial conditions + ic2 = symiirorder2_ic_bwd(out1, r, omega, precision)[0] + + # Apply the system cs / (1 - a2 * z - a3 * z^2)) in backwards + exp = np.empty(n, dtype=dtype) + exp[-2:] = ic2[::-1] + + for i in range(n - 3, -1, -1): + exp[i] = cs * out1[i] + a2 * exp[i + 1] + a3 * exp[i + 2] + + out = symiirorder2(signal, r, omega, precision) + xp_assert_close(out, exp, atol=4e-6, rtol=6e-7) + + @pytest.mark.parametrize('dtyp', ['float32', 'float64']) + def test_symiir2_values(self, dtyp): + dtyp = getattr(np, dtyp) + rng = np.random.RandomState(1234) + s = rng.uniform(size=16).astype(dtyp) + res = symiirorder2(s, 0.1, 0.1, precision=1e-10) + + # values from scipy 1.9.1 + exp_res = np.array([0.26572609, 0.53408018, 0.51032696, 0.72115829, 0.69486885, + 0.3649055 , 0.37349478, 0.74165032, 0.89718521, 0.80582483, + 0.46758053, 0.51898709, 0.65025605, 0.65394321, 0.45273595, + 0.53539183], dtype=dtyp) + + assert res.dtype == dtyp + # The values in SciPy 1.14 agree with those in SciPy 1.9.1 to this + # accuracy only. Implementation differences are twofold: + # 1. boundary conditions are computed differently + # 2. the filter itself uses sosfilt instead of a hardcoded iteration + # The boundary conditions seem are tested separately (see + # test_symiir2_initial_{fwd,bwd} above, so the difference is likely + # due to a different way roundoff errors accumulate in the filter. + # In that respect, sosfilt is likely doing a better job. + xp_assert_close(res, exp_res, atol=2e-6) + + s = s + 1j*s + with pytest.raises(TypeError): + res = symiirorder2(s, 0.5, 0.1) + + def test_symiir1_integer_input(self): + s = np.where(np.arange(100) % 2, -1, 1) + expected = symiirorder1(s.astype(float), 0.5, 0.5) + out = symiirorder1(s, 0.5, 0.5) + xp_assert_close(out, expected) + + def test_symiir2_integer_input(self): + s = np.where(np.arange(100) % 2, -1, 1) + expected = symiirorder2(s.astype(float), 0.5, np.pi / 3.0) + out = symiirorder2(s, 0.5, np.pi / 3.0) + xp_assert_close(out, expected) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_upfirdn.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_upfirdn.py new file mode 100644 index 0000000000000000000000000000000000000000..f6634116593cd0b184bdcf21cc62a95818f72693 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_upfirdn.py @@ -0,0 +1,288 @@ +# Code adapted from "upfirdn" python library with permission: +# +# Copyright (c) 2009, Motorola, Inc +# +# All Rights Reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of Motorola nor the names of its contributors may be +# used to endorse or promote products derived from this software without +# specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import numpy as np +from itertools import product + +from scipy._lib._array_api import xp_assert_close +from pytest import raises as assert_raises +import pytest + +from scipy.signal import upfirdn, firwin +from scipy.signal._upfirdn import _output_len, _upfirdn_modes +from scipy.signal._upfirdn_apply import _pad_test + + +def upfirdn_naive(x, h, up=1, down=1): + """Naive upfirdn processing in Python. + + Note: arg order (x, h) differs to facilitate apply_along_axis use. + """ + h = np.asarray(h) + out = np.zeros(len(x) * up, x.dtype) + out[::up] = x + out = np.convolve(h, out)[::down][:_output_len(len(h), len(x), up, down)] + return out + + +class UpFIRDnCase: + """Test _UpFIRDn object""" + def __init__(self, up, down, h, x_dtype): + self.up = up + self.down = down + self.h = np.atleast_1d(h) + self.x_dtype = x_dtype + self.rng = np.random.RandomState(17) + + def __call__(self): + # tiny signal + self.scrub(np.ones(1, self.x_dtype)) + # ones + self.scrub(np.ones(10, self.x_dtype)) # ones + # randn + x = self.rng.randn(10).astype(self.x_dtype) + if self.x_dtype in (np.complex64, np.complex128): + x += 1j * self.rng.randn(10) + self.scrub(x) + # ramp + self.scrub(np.arange(10).astype(self.x_dtype)) + # 3D, random + size = (2, 3, 5) + x = self.rng.randn(*size).astype(self.x_dtype) + if self.x_dtype in (np.complex64, np.complex128): + x += 1j * self.rng.randn(*size) + for axis in range(len(size)): + self.scrub(x, axis=axis) + x = x[:, ::2, 1::3].T + for axis in range(len(size)): + self.scrub(x, axis=axis) + + def scrub(self, x, axis=-1): + yr = np.apply_along_axis(upfirdn_naive, axis, x, + self.h, self.up, self.down) + want_len = _output_len(len(self.h), x.shape[axis], self.up, self.down) + assert yr.shape[axis] == want_len + y = upfirdn(self.h, x, self.up, self.down, axis=axis) + assert y.shape[axis] == want_len + assert y.shape == yr.shape + dtypes = (self.h.dtype, x.dtype) + if all(d == np.complex64 for d in dtypes): + assert y.dtype == np.complex64 + elif np.complex64 in dtypes and np.float32 in dtypes: + assert y.dtype == np.complex64 + elif all(d == np.float32 for d in dtypes): + assert y.dtype == np.float32 + elif np.complex128 in dtypes or np.complex64 in dtypes: + assert y.dtype == np.complex128 + else: + assert y.dtype == np.float64 + xp_assert_close(yr.astype(y.dtype), y) + + +_UPFIRDN_TYPES = (int, np.float32, np.complex64, float, complex) + + +class TestUpfirdn: + + def test_valid_input(self): + assert_raises(ValueError, upfirdn, [1], [1], 1, 0) # up or down < 1 + assert_raises(ValueError, upfirdn, [], [1], 1, 1) # h.ndim != 1 + assert_raises(ValueError, upfirdn, [[1]], [1], 1, 1) + + @pytest.mark.parametrize('len_h', [1, 2, 3, 4, 5]) + @pytest.mark.parametrize('len_x', [1, 2, 3, 4, 5]) + def test_singleton(self, len_h, len_x): + # gh-9844: lengths producing expected outputs + h = np.zeros(len_h) + h[len_h // 2] = 1. # make h a delta + x = np.ones(len_x) + y = upfirdn(h, x, 1, 1) + want = np.pad(x, (len_h // 2, (len_h - 1) // 2), 'constant') + xp_assert_close(y, want) + + def test_shift_x(self): + # gh-9844: shifted x can change values? + y = upfirdn([1, 1], [1.], 1, 1) + xp_assert_close(y, np.asarray([1.0, 1.0])) # was [0, 1] in the issue + y = upfirdn([1, 1], [0., 1.], 1, 1) + xp_assert_close(y, np.asarray([0.0, 1.0, 1.0])) + + # A bunch of lengths/factors chosen because they exposed differences + # between the "old way" and new way of computing length, and then + # got `expected` from MATLAB + @pytest.mark.parametrize('len_h, len_x, up, down, expected', [ + (2, 2, 5, 2, [1, 0, 0, 0]), + (2, 3, 6, 3, [1, 0, 1, 0, 1]), + (2, 4, 4, 3, [1, 0, 0, 0, 1]), + (3, 2, 6, 2, [1, 0, 0, 1, 0]), + (4, 11, 3, 5, [1, 0, 0, 1, 0, 0, 1]), + ]) + def test_length_factors(self, len_h, len_x, up, down, expected): + # gh-9844: weird factors + h = np.zeros(len_h) + h[0] = 1. + x = np.ones(len_x) + y = upfirdn(h, x, up, down) + expected = np.asarray(expected, dtype=np.float64) + xp_assert_close(y, expected) + + @pytest.mark.parametrize('down, want_len', [ # lengths from MATLAB + (2, 5015), + (11, 912), + (79, 127), + ]) + def test_vs_convolve(self, down, want_len): + # Check that up=1.0 gives same answer as convolve + slicing + random_state = np.random.RandomState(17) + try_types = (int, np.float32, np.complex64, float, complex) + size = 10000 + + for dtype in try_types: + x = random_state.randn(size).astype(dtype) + if dtype in (np.complex64, np.complex128): + x += 1j * random_state.randn(size) + + h = firwin(31, 1. / down, window='hamming') + yl = upfirdn_naive(x, h, 1, down) + y = upfirdn(h, x, up=1, down=down) + assert y.shape == (want_len,) + assert yl.shape[0] == y.shape[0] + xp_assert_close(yl, y, atol=1e-7, rtol=1e-7) + + @pytest.mark.parametrize('x_dtype', _UPFIRDN_TYPES) + @pytest.mark.parametrize('h', (1., 1j)) + @pytest.mark.parametrize('up, down', [(1, 1), (2, 2), (3, 2), (2, 3)]) + def test_vs_naive_delta(self, x_dtype, h, up, down): + UpFIRDnCase(up, down, h, x_dtype)() + + @pytest.mark.parametrize('x_dtype', _UPFIRDN_TYPES) + @pytest.mark.parametrize('h_dtype', _UPFIRDN_TYPES) + @pytest.mark.parametrize('p_max, q_max', + list(product((10, 100), (10, 100)))) + def test_vs_naive(self, x_dtype, h_dtype, p_max, q_max): + tests = self._random_factors(p_max, q_max, h_dtype, x_dtype) + for test in tests: + test() + + def _random_factors(self, p_max, q_max, h_dtype, x_dtype): + n_rep = 3 + longest_h = 25 + random_state = np.random.RandomState(17) + tests = [] + + for _ in range(n_rep): + # Randomize the up/down factors somewhat + p_add = q_max if p_max > q_max else 1 + q_add = p_max if q_max > p_max else 1 + p = random_state.randint(p_max) + p_add + q = random_state.randint(q_max) + q_add + + # Generate random FIR coefficients + len_h = random_state.randint(longest_h) + 1 + h = np.atleast_1d(random_state.randint(len_h)) + h = h.astype(h_dtype) + if h_dtype is complex: + h += 1j * random_state.randint(len_h) + + tests.append(UpFIRDnCase(p, q, h, x_dtype)) + + return tests + + @pytest.mark.parametrize('mode', _upfirdn_modes) + def test_extensions(self, mode): + """Test vs. manually computed results for modes not in numpy's pad.""" + x = np.array([1, 2, 3, 1], dtype=float) + npre, npost = 6, 6 + y = _pad_test(x, npre=npre, npost=npost, mode=mode) + if mode == 'antisymmetric': + y_expected = np.asarray( + [3.0, 1, -1, -3, -2, -1, 1, 2, 3, 1, -1, -3, -2, -1, 1, 2]) + elif mode == 'antireflect': + y_expected = np.asarray( + [1.0, 2, 3, 1, -1, 0, 1, 2, 3, 1, -1, 0, 1, 2, 3, 1]) + elif mode == 'smooth': + y_expected = np.asarray( + [-5.0, -4, -3, -2, -1, 0, 1, 2, 3, 1, -1, -3, -5, -7, -9, -11]) + elif mode == "line": + lin_slope = (x[-1] - x[0]) / (len(x) - 1) + left = x[0] + np.arange(-npre, 0, 1) * lin_slope + right = x[-1] + np.arange(1, npost + 1) * lin_slope + y_expected = np.concatenate((left, x, right)) + else: + y_expected = np.pad(x, (npre, npost), mode=mode) + xp_assert_close(y, y_expected) + + @pytest.mark.parametrize( + 'size, h_len, mode, dtype', + product( + [8], + [4, 5, 26], # include cases with h_len > 2*size + _upfirdn_modes, + [np.float32, np.float64, np.complex64, np.complex128], + ) + ) + def test_modes(self, size, h_len, mode, dtype): + random_state = np.random.RandomState(5) + x = random_state.randn(size).astype(dtype) + if dtype in (np.complex64, np.complex128): + x += 1j * random_state.randn(size) + h = np.arange(1, 1 + h_len, dtype=x.real.dtype) + + y = upfirdn(h, x, up=1, down=1, mode=mode) + # expected result: pad the input, filter with zero padding, then crop + npad = h_len - 1 + if mode in ['antisymmetric', 'antireflect', 'smooth', 'line']: + # use _pad_test test function for modes not supported by np.pad. + xpad = _pad_test(x, npre=npad, npost=npad, mode=mode) + else: + xpad = np.pad(x, npad, mode=mode) + ypad = upfirdn(h, xpad, up=1, down=1, mode='constant') + y_expected = ypad[npad:-npad] + + atol = rtol = np.finfo(dtype).eps * 1e2 + xp_assert_close(y, y_expected, atol=atol, rtol=rtol) + + +def test_output_len_long_input(): + # Regression test for gh-17375. On Windows, a large enough input + # that should have been well within the capabilities of 64 bit integers + # would result in a 32 bit overflow because of a bug in Cython 0.29.32. + len_h = 1001 + in_len = 10**8 + up = 320 + down = 441 + out_len = _output_len(len_h, in_len, up, down) + # The expected value was computed "by hand" from the formula + # (((in_len - 1) * up + len_h) - 1) // down + 1 + assert out_len == 72562360 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_waveforms.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_waveforms.py new file mode 100644 index 0000000000000000000000000000000000000000..016500679e31570348ef10cc3ab8523e382232eb --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_waveforms.py @@ -0,0 +1,380 @@ +import numpy as np +from pytest import raises as assert_raises +from scipy._lib._array_api import ( + assert_almost_equal, xp_assert_equal, xp_assert_close +) + +import scipy.signal._waveforms as waveforms + + +# These chirp_* functions are the instantaneous frequencies of the signals +# returned by chirp(). + +def chirp_linear(t, f0, f1, t1): + f = f0 + (f1 - f0) * t / t1 + return f + + +def chirp_quadratic(t, f0, f1, t1, vertex_zero=True): + if vertex_zero: + f = f0 + (f1 - f0) * t**2 / t1**2 + else: + f = f1 - (f1 - f0) * (t1 - t)**2 / t1**2 + return f + + +def chirp_geometric(t, f0, f1, t1): + f = f0 * (f1/f0)**(t/t1) + return f + + +def chirp_hyperbolic(t, f0, f1, t1): + f = f0*f1*t1 / ((f0 - f1)*t + f1*t1) + return f + + +def compute_frequency(t, theta): + """ + Compute theta'(t)/(2*pi), where theta'(t) is the derivative of theta(t). + """ + # Assume theta and t are 1-D NumPy arrays. + # Assume that t is uniformly spaced. + dt = t[1] - t[0] + f = np.diff(theta)/(2*np.pi) / dt + tf = 0.5*(t[1:] + t[:-1]) + return tf, f + + +class TestChirp: + + def test_linear_at_zero(self): + w = waveforms.chirp(t=0, f0=1.0, f1=2.0, t1=1.0, method='linear') + assert_almost_equal(w, 1.0) + + def test_linear_freq_01(self): + method = 'linear' + f0 = 1.0 + f1 = 2.0 + t1 = 1.0 + t = np.linspace(0, t1, 100) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_linear(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_linear_freq_02(self): + method = 'linear' + f0 = 200.0 + f1 = 100.0 + t1 = 10.0 + t = np.linspace(0, t1, 100) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_linear(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_linear_complex_power(self): + method = 'linear' + f0 = 1.0 + f1 = 2.0 + t1 = 1.0 + t = np.linspace(0, t1, 100) + w_real = waveforms.chirp(t, f0, t1, f1, method, complex=False) + w_complex = waveforms.chirp(t, f0, t1, f1, method, complex=True) + w_pwr_r = np.var(w_real) + w_pwr_c = np.var(w_complex) + + # Making sure that power of the real part is not affected with + # complex conversion operation + err = w_pwr_r - np.real(w_pwr_c) + + assert(err < 1e-6) + + def test_linear_complex_at_zero(self): + w = waveforms.chirp(t=0, f0=-10.0, f1=1.0, t1=1.0, method='linear', + complex=True) + xp_assert_close(w, 1.0+0.0j) # dtype must match + + def test_quadratic_at_zero(self): + w = waveforms.chirp(t=0, f0=1.0, f1=2.0, t1=1.0, method='quadratic') + assert_almost_equal(w, 1.0) + + def test_quadratic_at_zero2(self): + w = waveforms.chirp(t=0, f0=1.0, f1=2.0, t1=1.0, method='quadratic', + vertex_zero=False) + assert_almost_equal(w, 1.0) + + def test_quadratic_complex_at_zero(self): + w = waveforms.chirp(t=0, f0=-1.0, f1=2.0, t1=1.0, method='quadratic', + complex=True) + xp_assert_close(w, 1.0+0j) + + def test_quadratic_freq_01(self): + method = 'quadratic' + f0 = 1.0 + f1 = 2.0 + t1 = 1.0 + t = np.linspace(0, t1, 2000) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_quadratic(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_quadratic_freq_02(self): + method = 'quadratic' + f0 = 20.0 + f1 = 10.0 + t1 = 10.0 + t = np.linspace(0, t1, 2000) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_quadratic(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_logarithmic_at_zero(self): + w = waveforms.chirp(t=0, f0=1.0, f1=2.0, t1=1.0, method='logarithmic') + assert_almost_equal(w, 1.0) + + def test_logarithmic_freq_01(self): + method = 'logarithmic' + f0 = 1.0 + f1 = 2.0 + t1 = 1.0 + t = np.linspace(0, t1, 10000) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_geometric(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_logarithmic_freq_02(self): + method = 'logarithmic' + f0 = 200.0 + f1 = 100.0 + t1 = 10.0 + t = np.linspace(0, t1, 10000) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_geometric(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_logarithmic_freq_03(self): + method = 'logarithmic' + f0 = 100.0 + f1 = 100.0 + t1 = 10.0 + t = np.linspace(0, t1, 10000) + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + abserr = np.max(np.abs(f - chirp_geometric(tf, f0, f1, t1))) + assert abserr < 1e-6 + + def test_hyperbolic_at_zero(self): + w = waveforms.chirp(t=0, f0=10.0, f1=1.0, t1=1.0, method='hyperbolic') + assert_almost_equal(w, 1.0) + + def test_hyperbolic_freq_01(self): + method = 'hyperbolic' + t1 = 1.0 + t = np.linspace(0, t1, 10000) + # f0 f1 + cases = [[10.0, 1.0], + [1.0, 10.0], + [-10.0, -1.0], + [-1.0, -10.0]] + for f0, f1 in cases: + phase = waveforms._chirp_phase(t, f0, t1, f1, method) + tf, f = compute_frequency(t, phase) + expected = chirp_hyperbolic(tf, f0, f1, t1) + xp_assert_close(f, expected, atol=1e-7) + + def test_hyperbolic_zero_freq(self): + # f0=0 or f1=0 must raise a ValueError. + method = 'hyperbolic' + t1 = 1.0 + t = np.linspace(0, t1, 5) + assert_raises(ValueError, waveforms.chirp, t, 0, t1, 1, method) + assert_raises(ValueError, waveforms.chirp, t, 1, t1, 0, method) + + def test_unknown_method(self): + method = "foo" + f0 = 10.0 + f1 = 20.0 + t1 = 1.0 + t = np.linspace(0, t1, 10) + assert_raises(ValueError, waveforms.chirp, t, f0, t1, f1, method) + + def test_integer_t1(self): + f0 = 10.0 + f1 = 20.0 + t = np.linspace(-1, 1, 11) + t1 = 3.0 + float_result = waveforms.chirp(t, f0, t1, f1) + t1 = 3 + int_result = waveforms.chirp(t, f0, t1, f1) + err_msg = "Integer input 't1=3' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + def test_integer_f0(self): + f1 = 20.0 + t1 = 3.0 + t = np.linspace(-1, 1, 11) + f0 = 10.0 + float_result = waveforms.chirp(t, f0, t1, f1) + f0 = 10 + int_result = waveforms.chirp(t, f0, t1, f1) + err_msg = "Integer input 'f0=10' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + def test_integer_f1(self): + f0 = 10.0 + t1 = 3.0 + t = np.linspace(-1, 1, 11) + f1 = 20.0 + float_result = waveforms.chirp(t, f0, t1, f1) + f1 = 20 + int_result = waveforms.chirp(t, f0, t1, f1) + err_msg = "Integer input 'f1=20' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + def test_integer_all(self): + f0 = 10 + t1 = 3 + f1 = 20 + t = np.linspace(-1, 1, 11) + float_result = waveforms.chirp(t, float(f0), float(t1), float(f1)) + int_result = waveforms.chirp(t, f0, t1, f1) + err_msg = "Integer input 'f0=10, t1=3, f1=20' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + +class TestSweepPoly: + + def test_sweep_poly_quad1(self): + p = np.poly1d([1.0, 0.0, 1.0]) + t = np.linspace(0, 3.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = p(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + def test_sweep_poly_const(self): + p = np.poly1d(2.0) + t = np.linspace(0, 3.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = p(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + def test_sweep_poly_linear(self): + p = np.poly1d([-1.0, 10.0]) + t = np.linspace(0, 3.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = p(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + def test_sweep_poly_quad2(self): + p = np.poly1d([1.0, 0.0, -2.0]) + t = np.linspace(0, 3.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = p(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + def test_sweep_poly_cubic(self): + p = np.poly1d([2.0, 1.0, 0.0, -2.0]) + t = np.linspace(0, 2.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = p(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + def test_sweep_poly_cubic2(self): + """Use an array of coefficients instead of a poly1d.""" + p = np.array([2.0, 1.0, 0.0, -2.0]) + t = np.linspace(0, 2.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = np.poly1d(p)(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + def test_sweep_poly_cubic3(self): + """Use a list of coefficients instead of a poly1d.""" + p = [2.0, 1.0, 0.0, -2.0] + t = np.linspace(0, 2.0, 10000) + phase = waveforms._sweep_poly_phase(t, p) + tf, f = compute_frequency(t, phase) + expected = np.poly1d(p)(tf) + abserr = np.max(np.abs(f - expected)) + assert abserr < 1e-6 + + +class TestGaussPulse: + + def test_integer_fc(self): + float_result = waveforms.gausspulse('cutoff', fc=1000.0) + int_result = waveforms.gausspulse('cutoff', fc=1000) + err_msg = "Integer input 'fc=1000' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + def test_integer_bw(self): + float_result = waveforms.gausspulse('cutoff', bw=1.0) + int_result = waveforms.gausspulse('cutoff', bw=1) + err_msg = "Integer input 'bw=1' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + def test_integer_bwr(self): + float_result = waveforms.gausspulse('cutoff', bwr=-6.0) + int_result = waveforms.gausspulse('cutoff', bwr=-6) + err_msg = "Integer input 'bwr=-6' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + def test_integer_tpr(self): + float_result = waveforms.gausspulse('cutoff', tpr=-60.0) + int_result = waveforms.gausspulse('cutoff', tpr=-60) + err_msg = "Integer input 'tpr=-60' gives wrong result" + xp_assert_equal(int_result, float_result, err_msg=err_msg) + + +class TestUnitImpulse: + + def test_no_index(self): + xp_assert_equal(waveforms.unit_impulse(7), + np.asarray([1.0, 0, 0, 0, 0, 0, 0])) + xp_assert_equal(waveforms.unit_impulse((3, 3)), + np.asarray([[1.0, 0, 0], [0, 0, 0], [0, 0, 0]])) + + def test_index(self): + xp_assert_equal(waveforms.unit_impulse(10, 3), + np.asarray([0.0, 0, 0, 1, 0, 0, 0, 0, 0, 0])) + xp_assert_equal(waveforms.unit_impulse((3, 3), (1, 1)), + np.asarray([[0.0, 0, 0], [0, 1, 0], [0, 0, 0]])) + + # Broadcasting + imp = waveforms.unit_impulse((4, 4), 2) + xp_assert_equal(imp, np.asarray([[0.0, 0, 0, 0], + [0.0, 0, 0, 0], + [0.0, 0, 1, 0], + [0.0, 0, 0, 0]])) + + def test_mid(self): + xp_assert_equal(waveforms.unit_impulse((3, 3), 'mid'), + np.asarray([[0.0, 0, 0], [0, 1, 0], [0, 0, 0]])) + xp_assert_equal(waveforms.unit_impulse(9, 'mid'), + np.asarray([0.0, 0, 0, 0, 1, 0, 0, 0, 0])) + + def test_dtype(self): + imp = waveforms.unit_impulse(7) + assert np.issubdtype(imp.dtype, np.floating) + + imp = waveforms.unit_impulse(5, 3, dtype=int) + assert np.issubdtype(imp.dtype, np.integer) + + imp = waveforms.unit_impulse((5, 2), (3, 1), dtype=complex) + assert np.issubdtype(imp.dtype, np.complexfloating) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_wavelets.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_wavelets.py new file mode 100644 index 0000000000000000000000000000000000000000..ce54a4d1d019993cb939f911d0efbbaf8c9862b1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_wavelets.py @@ -0,0 +1,59 @@ +import numpy as np +from numpy.testing import assert_array_equal, assert_array_almost_equal + +import scipy.signal._wavelets as wavelets + + +class TestWavelets: + def test_ricker(self): + w = wavelets._ricker(1.0, 1) + expected = 2 / (np.sqrt(3 * 1.0) * (np.pi ** 0.25)) + assert_array_equal(w, expected) + + lengths = [5, 11, 15, 51, 101] + for length in lengths: + w = wavelets._ricker(length, 1.0) + assert len(w) == length + max_loc = np.argmax(w) + assert max_loc == (length // 2) + + points = 100 + w = wavelets._ricker(points, 2.0) + half_vec = np.arange(0, points // 2) + # Wavelet should be symmetric + assert_array_almost_equal(w[half_vec], w[-(half_vec + 1)]) + + # Check zeros + aas = [5, 10, 15, 20, 30] + points = 99 + for a in aas: + w = wavelets._ricker(points, a) + vec = np.arange(0, points) - (points - 1.0) / 2 + exp_zero1 = np.argmin(np.abs(vec - a)) + exp_zero2 = np.argmin(np.abs(vec + a)) + assert_array_almost_equal(w[exp_zero1], 0) + assert_array_almost_equal(w[exp_zero2], 0) + + def test_cwt(self): + widths = [1.0] + def delta_wavelet(s, t): + return np.array([1]) + len_data = 100 + test_data = np.sin(np.pi * np.arange(0, len_data) / 10.0) + + # Test delta function input gives same data as output + cwt_dat = wavelets._cwt(test_data, delta_wavelet, widths) + assert cwt_dat.shape == (len(widths), len_data) + assert_array_almost_equal(test_data, cwt_dat.flatten()) + + # Check proper shape on output + widths = [1, 3, 4, 5, 10] + cwt_dat = wavelets._cwt(test_data, wavelets._ricker, widths) + assert cwt_dat.shape == (len(widths), len_data) + + widths = [len_data * 10] + # Note: this wavelet isn't defined quite right, but is fine for this test + def flat_wavelet(l, w): + return np.full(w, 1 / w) + cwt_dat = wavelets._cwt(test_data, flat_wavelet, widths) + assert_array_almost_equal(cwt_dat, np.mean(test_data)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_windows.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_windows.py new file mode 100644 index 0000000000000000000000000000000000000000..21aae294eaa2ee41e9d5114503e6ff8647bd44a0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/tests/test_windows.py @@ -0,0 +1,846 @@ +import numpy as np +from numpy import array +from numpy.testing import (assert_array_almost_equal, assert_array_equal, + assert_allclose, + assert_equal, assert_, assert_array_less, + suppress_warnings) +from pytest import raises as assert_raises + +from scipy.fft import fft +from scipy.signal import windows, get_window, resample + + +window_funcs = [ + ('boxcar', ()), + ('triang', ()), + ('parzen', ()), + ('bohman', ()), + ('blackman', ()), + ('nuttall', ()), + ('blackmanharris', ()), + ('flattop', ()), + ('bartlett', ()), + ('barthann', ()), + ('hamming', ()), + ('kaiser', (1,)), + ('dpss', (2,)), + ('gaussian', (0.5,)), + ('general_gaussian', (1.5, 2)), + ('chebwin', (1,)), + ('cosine', ()), + ('hann', ()), + ('exponential', ()), + ('taylor', ()), + ('tukey', (0.5,)), + ('lanczos', ()), + ] + + +class TestBartHann: + + def test_basic(self): + assert_allclose(windows.barthann(6, sym=True), + [0, 0.35857354213752, 0.8794264578624801, + 0.8794264578624801, 0.3585735421375199, 0], + rtol=1e-15, atol=1e-15) + assert_allclose(windows.barthann(7), + [0, 0.27, 0.73, 1.0, 0.73, 0.27, 0], + rtol=1e-15, atol=1e-15) + assert_allclose(windows.barthann(6, False), + [0, 0.27, 0.73, 1.0, 0.73, 0.27], + rtol=1e-15, atol=1e-15) + + +class TestBartlett: + + def test_basic(self): + assert_allclose(windows.bartlett(6), [0, 0.4, 0.8, 0.8, 0.4, 0]) + assert_allclose(windows.bartlett(7), [0, 1/3, 2/3, 1.0, 2/3, 1/3, 0]) + assert_allclose(windows.bartlett(6, False), + [0, 1/3, 2/3, 1.0, 2/3, 1/3]) + + +class TestBlackman: + + def test_basic(self): + assert_allclose(windows.blackman(6, sym=False), + [0, 0.13, 0.63, 1.0, 0.63, 0.13], atol=1e-14) + assert_allclose(windows.blackman(7, sym=False), + [0, 0.09045342435412804, 0.4591829575459636, + 0.9203636180999081, 0.9203636180999081, + 0.4591829575459636, 0.09045342435412804], atol=1e-8) + assert_allclose(windows.blackman(6), + [0, 0.2007701432625305, 0.8492298567374694, + 0.8492298567374694, 0.2007701432625305, 0], + atol=1e-14) + assert_allclose(windows.blackman(7, True), + [0, 0.13, 0.63, 1.0, 0.63, 0.13, 0], atol=1e-14) + + +class TestBlackmanHarris: + + def test_basic(self): + assert_allclose(windows.blackmanharris(6, False), + [6.0e-05, 0.055645, 0.520575, 1.0, 0.520575, 0.055645]) + assert_allclose(windows.blackmanharris(7, sym=False), + [6.0e-05, 0.03339172347815117, 0.332833504298565, + 0.8893697722232837, 0.8893697722232838, + 0.3328335042985652, 0.03339172347815122]) + assert_allclose(windows.blackmanharris(6), + [6.0e-05, 0.1030114893456638, 0.7938335106543362, + 0.7938335106543364, 0.1030114893456638, 6.0e-05]) + assert_allclose(windows.blackmanharris(7, sym=True), + [6.0e-05, 0.055645, 0.520575, 1.0, 0.520575, 0.055645, + 6.0e-05]) + + +class TestTaylor: + + def test_normalized(self): + """Tests windows of small length that are normalized to 1. See the + documentation for the Taylor window for more information on + normalization. + """ + assert_allclose(windows.taylor(1, 2, 15), 1.0) + assert_allclose( + windows.taylor(5, 2, 15), + np.array([0.75803341, 0.90757699, 1.0, 0.90757699, 0.75803341]) + ) + assert_allclose( + windows.taylor(6, 2, 15), + np.array([ + 0.7504082, 0.86624416, 0.98208011, 0.98208011, 0.86624416, + 0.7504082 + ]) + ) + + def test_non_normalized(self): + """Test windows of small length that are not normalized to 1. See + the documentation for the Taylor window for more information on + normalization. + """ + assert_allclose( + windows.taylor(5, 2, 15, norm=False), + np.array([ + 0.87508054, 1.04771499, 1.15440894, 1.04771499, 0.87508054 + ]) + ) + assert_allclose( + windows.taylor(6, 2, 15, norm=False), + np.array([ + 0.86627793, 1.0, 1.13372207, 1.13372207, 1.0, 0.86627793 + ]) + ) + + def test_correctness(self): + """This test ensures the correctness of the implemented Taylor + Windowing function. A Taylor Window of 1024 points is created, its FFT + is taken, and the Peak Sidelobe Level (PSLL) and 3dB and 18dB bandwidth + are found and checked. + + A publication from Sandia National Laboratories was used as reference + for the correctness values [1]_. + + References + ----- + .. [1] Armin Doerry, "Catalog of Window Taper Functions for + Sidelobe Control", 2017. + https://www.researchgate.net/profile/Armin_Doerry/publication/316281181_Catalog_of_Window_Taper_Functions_for_Sidelobe_Control/links/58f92cb2a6fdccb121c9d54d/Catalog-of-Window-Taper-Functions-for-Sidelobe-Control.pdf + """ + M_win = 1024 + N_fft = 131072 + # Set norm=False for correctness as the values obtained from the + # scientific publication do not normalize the values. Normalizing + # changes the sidelobe level from the desired value. + w = windows.taylor(M_win, nbar=4, sll=35, norm=False, sym=False) + f = fft(w, N_fft) + spec = 20 * np.log10(np.abs(f / np.amax(f))) + + first_zero = np.argmax(np.diff(spec) > 0) + + PSLL = np.amax(spec[first_zero:-first_zero]) + + BW_3dB = 2*np.argmax(spec <= -3.0102999566398121) / N_fft * M_win + BW_18dB = 2*np.argmax(spec <= -18.061799739838872) / N_fft * M_win + + assert_allclose(PSLL, -35.1672, atol=1) + assert_allclose(BW_3dB, 1.1822, atol=0.1) + assert_allclose(BW_18dB, 2.6112, atol=0.1) + + +class TestBohman: + + def test_basic(self): + assert_allclose(windows.bohman(6), + [0, 0.1791238937062839, 0.8343114522576858, + 0.8343114522576858, 0.1791238937062838, 0]) + assert_allclose(windows.bohman(7, sym=True), + [0, 0.1089977810442293, 0.6089977810442293, 1.0, + 0.6089977810442295, 0.1089977810442293, 0]) + assert_allclose(windows.bohman(6, False), + [0, 0.1089977810442293, 0.6089977810442293, 1.0, + 0.6089977810442295, 0.1089977810442293]) + + +class TestBoxcar: + + def test_basic(self): + assert_allclose(windows.boxcar(6), [1, 1, 1, 1, 1, 1]) + assert_allclose(windows.boxcar(7), [1, 1, 1, 1, 1, 1, 1]) + assert_allclose(windows.boxcar(6, False), [1, 1, 1, 1, 1, 1]) + + +cheb_odd_true = array([0.200938, 0.107729, 0.134941, 0.165348, + 0.198891, 0.235450, 0.274846, 0.316836, + 0.361119, 0.407338, 0.455079, 0.503883, + 0.553248, 0.602637, 0.651489, 0.699227, + 0.745266, 0.789028, 0.829947, 0.867485, + 0.901138, 0.930448, 0.955010, 0.974482, + 0.988591, 0.997138, 1.000000, 0.997138, + 0.988591, 0.974482, 0.955010, 0.930448, + 0.901138, 0.867485, 0.829947, 0.789028, + 0.745266, 0.699227, 0.651489, 0.602637, + 0.553248, 0.503883, 0.455079, 0.407338, + 0.361119, 0.316836, 0.274846, 0.235450, + 0.198891, 0.165348, 0.134941, 0.107729, + 0.200938]) + +cheb_even_true = array([0.203894, 0.107279, 0.133904, + 0.163608, 0.196338, 0.231986, + 0.270385, 0.311313, 0.354493, + 0.399594, 0.446233, 0.493983, + 0.542378, 0.590916, 0.639071, + 0.686302, 0.732055, 0.775783, + 0.816944, 0.855021, 0.889525, + 0.920006, 0.946060, 0.967339, + 0.983557, 0.994494, 1.000000, + 1.000000, 0.994494, 0.983557, + 0.967339, 0.946060, 0.920006, + 0.889525, 0.855021, 0.816944, + 0.775783, 0.732055, 0.686302, + 0.639071, 0.590916, 0.542378, + 0.493983, 0.446233, 0.399594, + 0.354493, 0.311313, 0.270385, + 0.231986, 0.196338, 0.163608, + 0.133904, 0.107279, 0.203894]) + + +class TestChebWin: + + def test_basic(self): + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + assert_allclose(windows.chebwin(6, 100), + [0.1046401879356917, 0.5075781475823447, 1.0, 1.0, + 0.5075781475823447, 0.1046401879356917]) + assert_allclose(windows.chebwin(7, 100), + [0.05650405062850233, 0.316608530648474, + 0.7601208123539079, 1.0, 0.7601208123539079, + 0.316608530648474, 0.05650405062850233]) + assert_allclose(windows.chebwin(6, 10), + [1.0, 0.6071201674458373, 0.6808391469897297, + 0.6808391469897297, 0.6071201674458373, 1.0]) + assert_allclose(windows.chebwin(7, 10), + [1.0, 0.5190521247588651, 0.5864059018130382, + 0.6101519801307441, 0.5864059018130382, + 0.5190521247588651, 1.0]) + assert_allclose(windows.chebwin(6, 10, False), + [1.0, 0.5190521247588651, 0.5864059018130382, + 0.6101519801307441, 0.5864059018130382, + 0.5190521247588651]) + + def test_cheb_odd_high_attenuation(self): + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + cheb_odd = windows.chebwin(53, at=-40) + assert_array_almost_equal(cheb_odd, cheb_odd_true, decimal=4) + + def test_cheb_even_high_attenuation(self): + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + cheb_even = windows.chebwin(54, at=40) + assert_array_almost_equal(cheb_even, cheb_even_true, decimal=4) + + def test_cheb_odd_low_attenuation(self): + cheb_odd_low_at_true = array([1.000000, 0.519052, 0.586405, + 0.610151, 0.586405, 0.519052, + 1.000000]) + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + cheb_odd = windows.chebwin(7, at=10) + assert_array_almost_equal(cheb_odd, cheb_odd_low_at_true, decimal=4) + + def test_cheb_even_low_attenuation(self): + cheb_even_low_at_true = array([1.000000, 0.451924, 0.51027, + 0.541338, 0.541338, 0.51027, + 0.451924, 1.000000]) + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + cheb_even = windows.chebwin(8, at=-10) + assert_array_almost_equal(cheb_even, cheb_even_low_at_true, decimal=4) + + +exponential_data = { + (4, None, 0.2, False): + array([4.53999297624848542e-05, + 6.73794699908546700e-03, 1.00000000000000000e+00, + 6.73794699908546700e-03]), + (4, None, 0.2, True): array([0.00055308437014783, 0.0820849986238988, + 0.0820849986238988, 0.00055308437014783]), + (4, None, 1.0, False): array([0.1353352832366127, 0.36787944117144233, 1., + 0.36787944117144233]), + (4, None, 1.0, True): array([0.22313016014842982, 0.60653065971263342, + 0.60653065971263342, 0.22313016014842982]), + (4, 2, 0.2, False): + array([4.53999297624848542e-05, 6.73794699908546700e-03, + 1.00000000000000000e+00, 6.73794699908546700e-03]), + (4, 2, 0.2, True): None, + (4, 2, 1.0, False): array([0.1353352832366127, 0.36787944117144233, 1., + 0.36787944117144233]), + (4, 2, 1.0, True): None, + (5, None, 0.2, True): + array([4.53999297624848542e-05, + 6.73794699908546700e-03, 1.00000000000000000e+00, + 6.73794699908546700e-03, 4.53999297624848542e-05]), + (5, None, 1.0, True): array([0.1353352832366127, 0.36787944117144233, 1., + 0.36787944117144233, 0.1353352832366127]), + (5, 2, 0.2, True): None, + (5, 2, 1.0, True): None +} + + +def test_exponential(): + for k, v in exponential_data.items(): + if v is None: + assert_raises(ValueError, windows.exponential, *k) + else: + win = windows.exponential(*k) + assert_allclose(win, v, rtol=1e-14) + + +class TestFlatTop: + + def test_basic(self): + assert_allclose(windows.flattop(6, sym=False), + [-0.000421051, -0.051263156, 0.19821053, 1.0, + 0.19821053, -0.051263156]) + assert_allclose(windows.flattop(7, sym=False), + [-0.000421051, -0.03684078115492348, + 0.01070371671615342, 0.7808739149387698, + 0.7808739149387698, 0.01070371671615342, + -0.03684078115492348]) + assert_allclose(windows.flattop(6), + [-0.000421051, -0.0677142520762119, 0.6068721525762117, + 0.6068721525762117, -0.0677142520762119, + -0.000421051]) + assert_allclose(windows.flattop(7, True), + [-0.000421051, -0.051263156, 0.19821053, 1.0, + 0.19821053, -0.051263156, -0.000421051]) + + +class TestGaussian: + + def test_basic(self): + assert_allclose(windows.gaussian(6, 1.0), + [0.04393693362340742, 0.3246524673583497, + 0.8824969025845955, 0.8824969025845955, + 0.3246524673583497, 0.04393693362340742]) + assert_allclose(windows.gaussian(7, 1.2), + [0.04393693362340742, 0.2493522087772962, + 0.7066482778577162, 1.0, 0.7066482778577162, + 0.2493522087772962, 0.04393693362340742]) + assert_allclose(windows.gaussian(7, 3), + [0.6065306597126334, 0.8007374029168081, + 0.9459594689067654, 1.0, 0.9459594689067654, + 0.8007374029168081, 0.6065306597126334]) + assert_allclose(windows.gaussian(6, 3, False), + [0.6065306597126334, 0.8007374029168081, + 0.9459594689067654, 1.0, 0.9459594689067654, + 0.8007374029168081]) + + +class TestGeneralCosine: + + def test_basic(self): + assert_allclose(windows.general_cosine(5, [0.5, 0.3, 0.2]), + [0.4, 0.3, 1, 0.3, 0.4]) + assert_allclose(windows.general_cosine(4, [0.5, 0.3, 0.2], sym=False), + [0.4, 0.3, 1, 0.3]) + + +class TestGeneralHamming: + + def test_basic(self): + assert_allclose(windows.general_hamming(5, 0.7), + [0.4, 0.7, 1.0, 0.7, 0.4]) + assert_allclose(windows.general_hamming(5, 0.75, sym=False), + [0.5, 0.6727457514, 0.9522542486, + 0.9522542486, 0.6727457514]) + assert_allclose(windows.general_hamming(6, 0.75, sym=True), + [0.5, 0.6727457514, 0.9522542486, + 0.9522542486, 0.6727457514, 0.5]) + + +class TestHamming: + + def test_basic(self): + assert_allclose(windows.hamming(6, False), + [0.08, 0.31, 0.77, 1.0, 0.77, 0.31]) + assert_allclose(windows.hamming(7, sym=False), + [0.08, 0.2531946911449826, 0.6423596296199047, + 0.9544456792351128, 0.9544456792351128, + 0.6423596296199047, 0.2531946911449826]) + assert_allclose(windows.hamming(6), + [0.08, 0.3978521825875242, 0.9121478174124757, + 0.9121478174124757, 0.3978521825875242, 0.08]) + assert_allclose(windows.hamming(7, sym=True), + [0.08, 0.31, 0.77, 1.0, 0.77, 0.31, 0.08]) + + +class TestHann: + + def test_basic(self): + assert_allclose(windows.hann(6, sym=False), + [0, 0.25, 0.75, 1.0, 0.75, 0.25], + rtol=1e-15, atol=1e-15) + assert_allclose(windows.hann(7, sym=False), + [0, 0.1882550990706332, 0.6112604669781572, + 0.9504844339512095, 0.9504844339512095, + 0.6112604669781572, 0.1882550990706332], + rtol=1e-15, atol=1e-15) + assert_allclose(windows.hann(6, True), + [0, 0.3454915028125263, 0.9045084971874737, + 0.9045084971874737, 0.3454915028125263, 0], + rtol=1e-15, atol=1e-15) + assert_allclose(windows.hann(7), + [0, 0.25, 0.75, 1.0, 0.75, 0.25, 0], + rtol=1e-15, atol=1e-15) + + +class TestKaiser: + + def test_basic(self): + assert_allclose(windows.kaiser(6, 0.5), + [0.9403061933191572, 0.9782962393705389, + 0.9975765035372042, 0.9975765035372042, + 0.9782962393705389, 0.9403061933191572]) + assert_allclose(windows.kaiser(7, 0.5), + [0.9403061933191572, 0.9732402256999829, + 0.9932754654413773, 1.0, 0.9932754654413773, + 0.9732402256999829, 0.9403061933191572]) + assert_allclose(windows.kaiser(6, 2.7), + [0.2603047507678832, 0.6648106293528054, + 0.9582099802511439, 0.9582099802511439, + 0.6648106293528054, 0.2603047507678832]) + assert_allclose(windows.kaiser(7, 2.7), + [0.2603047507678832, 0.5985765418119844, + 0.8868495172060835, 1.0, 0.8868495172060835, + 0.5985765418119844, 0.2603047507678832]) + assert_allclose(windows.kaiser(6, 2.7, False), + [0.2603047507678832, 0.5985765418119844, + 0.8868495172060835, 1.0, 0.8868495172060835, + 0.5985765418119844]) + + +class TestKaiserBesselDerived: + + def test_basic(self): + M = 100 + w = windows.kaiser_bessel_derived(M, beta=4.0) + w2 = windows.get_window(('kaiser bessel derived', 4.0), + M, fftbins=False) + assert_allclose(w, w2) + + # Test for Princen-Bradley condition + assert_allclose(w[:M // 2] ** 2 + w[-M // 2:] ** 2, 1.) + + # Test actual values from other implementations + # M = 2: sqrt(2) / 2 + # M = 4: 0.518562710536, 0.855039598640 + # M = 6: 0.436168993154, 0.707106781187, 0.899864772847 + # Ref:https://github.com/scipy/scipy/pull/4747#issuecomment-172849418 + assert_allclose(windows.kaiser_bessel_derived(2, beta=np.pi / 2)[:1], + np.sqrt(2) / 2) + + assert_allclose(windows.kaiser_bessel_derived(4, beta=np.pi / 2)[:2], + [0.518562710536, 0.855039598640]) + + assert_allclose(windows.kaiser_bessel_derived(6, beta=np.pi / 2)[:3], + [0.436168993154, 0.707106781187, 0.899864772847]) + + def test_exceptions(self): + M = 100 + # Assert ValueError for odd window length + msg = ("Kaiser-Bessel Derived windows are only defined for even " + "number of points") + with assert_raises(ValueError, match=msg): + windows.kaiser_bessel_derived(M + 1, beta=4.) + + # Assert ValueError for non-symmetric setting + msg = ("Kaiser-Bessel Derived windows are only defined for " + "symmetric shapes") + with assert_raises(ValueError, match=msg): + windows.kaiser_bessel_derived(M + 1, beta=4., sym=False) + + +class TestNuttall: + + def test_basic(self): + assert_allclose(windows.nuttall(6, sym=False), + [0.0003628, 0.0613345, 0.5292298, 1.0, 0.5292298, + 0.0613345]) + assert_allclose(windows.nuttall(7, sym=False), + [0.0003628, 0.03777576895352025, 0.3427276199688195, + 0.8918518610776603, 0.8918518610776603, + 0.3427276199688196, 0.0377757689535203]) + assert_allclose(windows.nuttall(6), + [0.0003628, 0.1105152530498718, 0.7982580969501282, + 0.7982580969501283, 0.1105152530498719, 0.0003628]) + assert_allclose(windows.nuttall(7, True), + [0.0003628, 0.0613345, 0.5292298, 1.0, 0.5292298, + 0.0613345, 0.0003628]) + + +class TestParzen: + + def test_basic(self): + assert_allclose(windows.parzen(6), + [0.009259259259259254, 0.25, 0.8611111111111112, + 0.8611111111111112, 0.25, 0.009259259259259254]) + assert_allclose(windows.parzen(7, sym=True), + [0.00583090379008747, 0.1574344023323616, + 0.6501457725947521, 1.0, 0.6501457725947521, + 0.1574344023323616, 0.00583090379008747]) + assert_allclose(windows.parzen(6, False), + [0.00583090379008747, 0.1574344023323616, + 0.6501457725947521, 1.0, 0.6501457725947521, + 0.1574344023323616]) + + +class TestTriang: + + def test_basic(self): + + assert_allclose(windows.triang(6, True), + [1/6, 1/2, 5/6, 5/6, 1/2, 1/6]) + assert_allclose(windows.triang(7), + [1/4, 1/2, 3/4, 1, 3/4, 1/2, 1/4]) + assert_allclose(windows.triang(6, sym=False), + [1/4, 1/2, 3/4, 1, 3/4, 1/2]) + + +tukey_data = { + (4, 0.5, True): array([0.0, 1.0, 1.0, 0.0]), + (4, 0.9, True): array([0.0, 0.84312081893436686, + 0.84312081893436686, 0.0]), + (4, 1.0, True): array([0.0, 0.75, 0.75, 0.0]), + (4, 0.5, False): array([0.0, 1.0, 1.0, 1.0]), + (4, 0.9, False): array([0.0, 0.58682408883346526, + 1.0, 0.58682408883346526]), + (4, 1.0, False): array([0.0, 0.5, 1.0, 0.5]), + (5, 0.0, True): array([1.0, 1.0, 1.0, 1.0, 1.0]), + (5, 0.8, True): array([0.0, 0.69134171618254492, + 1.0, 0.69134171618254492, 0.0]), + (5, 1.0, True): array([0.0, 0.5, 1.0, 0.5, 0.0]), + + (6, 0): [1, 1, 1, 1, 1, 1], + (7, 0): [1, 1, 1, 1, 1, 1, 1], + (6, .25): [0, 1, 1, 1, 1, 0], + (7, .25): [0, 1, 1, 1, 1, 1, 0], + (6,): [0, 0.9045084971874737, 1.0, 1.0, 0.9045084971874735, 0], + (7,): [0, 0.75, 1.0, 1.0, 1.0, 0.75, 0], + (6, .75): [0, 0.5522642316338269, 1.0, 1.0, 0.5522642316338267, 0], + (7, .75): [0, 0.4131759111665348, 0.9698463103929542, 1.0, + 0.9698463103929542, 0.4131759111665347, 0], + (6, 1): [0, 0.3454915028125263, 0.9045084971874737, 0.9045084971874737, + 0.3454915028125263, 0], + (7, 1): [0, 0.25, 0.75, 1.0, 0.75, 0.25, 0], +} + + +class TestTukey: + + def test_basic(self): + # Test against hardcoded data + for k, v in tukey_data.items(): + if v is None: + assert_raises(ValueError, windows.tukey, *k) + else: + win = windows.tukey(*k) + assert_allclose(win, v, rtol=1e-15, atol=1e-15) + + def test_extremes(self): + # Test extremes of alpha correspond to boxcar and hann + tuk0 = windows.tukey(100, 0) + box0 = windows.boxcar(100) + assert_array_almost_equal(tuk0, box0) + + tuk1 = windows.tukey(100, 1) + han1 = windows.hann(100) + assert_array_almost_equal(tuk1, han1) + + +dpss_data = { + # All values from MATLAB: + # * taper[1] of (3, 1.4, 3) sign-flipped + # * taper[3] of (5, 1.5, 5) sign-flipped + (4, 0.1, 2): ([[0.497943898, 0.502047681, 0.502047681, 0.497943898], [0.670487993, 0.224601537, -0.224601537, -0.670487993]], [0.197961815, 0.002035474]), # noqa: E501 + (3, 1.4, 3): ([[0.410233151, 0.814504464, 0.410233151], [0.707106781, 0.0, -0.707106781], [0.575941629, -0.580157287, 0.575941629]], [0.999998093, 0.998067480, 0.801934426]), # noqa: E501 + (5, 1.5, 5): ([[0.1745071052, 0.4956749177, 0.669109327, 0.495674917, 0.174507105], [0.4399493348, 0.553574369, 0.0, -0.553574369, -0.439949334], [0.631452756, 0.073280238, -0.437943884, 0.073280238, 0.631452756], [0.553574369, -0.439949334, 0.0, 0.439949334, -0.553574369], [0.266110290, -0.498935248, 0.600414741, -0.498935248, 0.266110290147157]], [0.999728571, 0.983706916, 0.768457889, 0.234159338, 0.013947282907567]), # noqa: E501 + (100, 2, 4): ([[0.0030914414, 0.0041266922, 0.005315076, 0.006665149, 0.008184854, 0.0098814158, 0.011761239, 0.013829809, 0.016091597, 0.018549973, 0.02120712, 0.02406396, 0.027120092, 0.030373728, 0.033821651, 0.037459181, 0.041280145, 0.045276872, 0.049440192, 0.053759447, 0.058222524, 0.062815894, 0.067524661, 0.072332638, 0.077222418, 0.082175473, 0.087172252, 0.092192299, 0.097214376, 0.1022166, 0.10717657, 0.11207154, 0.11687856, 0.12157463, 0.12613686, 0.13054266, 0.13476986, 0.13879691, 0.14260302, 0.14616832, 0.14947401, 0.1525025, 0.15523755, 0.15766438, 0.15976981, 0.16154233, 0.16297223, 0.16405162, 0.16477455, 0.16513702, 0.16513702, 0.16477455, 0.16405162, 0.16297223, 0.16154233, 0.15976981, 0.15766438, 0.15523755, 0.1525025, 0.14947401, 0.14616832, 0.14260302, 0.13879691, 0.13476986, 0.13054266, 0.12613686, 0.12157463, 0.11687856, 0.11207154, 0.10717657, 0.1022166, 0.097214376, 0.092192299, 0.087172252, 0.082175473, 0.077222418, 0.072332638, 0.067524661, 0.062815894, 0.058222524, 0.053759447, 0.049440192, 0.045276872, 0.041280145, 0.037459181, 0.033821651, 0.030373728, 0.027120092, 0.02406396, 0.02120712, 0.018549973, 0.016091597, 0.013829809, 0.011761239, 0.0098814158, 0.008184854, 0.006665149, 0.005315076, 0.0041266922, 0.0030914414], [0.018064449, 0.022040342, 0.026325013, 0.030905288, 0.035764398, 0.040881982, 0.046234148, 0.051793558, 0.057529559, 0.063408356, 0.069393216, 0.075444716, 0.081521022, 0.087578202, 0.093570567, 0.099451049, 0.10517159, 0.11068356, 0.11593818, 0.12088699, 0.12548227, 0.12967752, 0.1334279, 0.13669069, 0.13942569, 0.1415957, 0.14316686, 0.14410905, 0.14439626, 0.14400686, 0.14292389, 0.1411353, 0.13863416, 0.13541876, 0.13149274, 0.12686516, 0.12155045, 0.1155684, 0.10894403, 0.10170748, 0.093893752, 0.08554251, 0.076697768, 0.067407559, 0.057723559, 0.04770068, 0.037396627, 0.026871428, 0.016186944, 0.0054063557, -0.0054063557, -0.016186944, -0.026871428, -0.037396627, -0.04770068, -0.057723559, -0.067407559, -0.076697768, -0.08554251, -0.093893752, -0.10170748, -0.10894403, -0.1155684, -0.12155045, -0.12686516, -0.13149274, -0.13541876, -0.13863416, -0.1411353, -0.14292389, -0.14400686, -0.14439626, -0.14410905, -0.14316686, -0.1415957, -0.13942569, -0.13669069, -0.1334279, -0.12967752, -0.12548227, -0.12088699, -0.11593818, -0.11068356, -0.10517159, -0.099451049, -0.093570567, -0.087578202, -0.081521022, -0.075444716, -0.069393216, -0.063408356, -0.057529559, -0.051793558, -0.046234148, -0.040881982, -0.035764398, -0.030905288, -0.026325013, -0.022040342, -0.018064449], [0.064817553, 0.072567801, 0.080292992, 0.087918235, 0.095367076, 0.10256232, 0.10942687, 0.1158846, 0.12186124, 0.12728523, 0.13208858, 0.13620771, 0.13958427, 0.14216587, 0.14390678, 0.14476863, 0.1447209, 0.14374148, 0.14181704, 0.13894336, 0.13512554, 0.13037812, 0.1247251, 0.11819984, 0.11084487, 0.10271159, 0.093859853, 0.084357497, 0.074279719, 0.063708406, 0.052731374, 0.041441525, 0.029935953, 0.018314987, 0.0066811877, -0.0048616765, -0.016209689, -0.027259848, -0.037911124, -0.048065512, -0.05762905, -0.066512804, -0.0746338, -0.081915903, -0.088290621, -0.09369783, -0.098086416, -0.10141482, -0.10365146, -0.10477512, -0.10477512, -0.10365146, -0.10141482, -0.098086416, -0.09369783, -0.088290621, -0.081915903, -0.0746338, -0.066512804, -0.05762905, -0.048065512, -0.037911124, -0.027259848, -0.016209689, -0.0048616765, 0.0066811877, 0.018314987, 0.029935953, 0.041441525, 0.052731374, 0.063708406, 0.074279719, 0.084357497, 0.093859853, 0.10271159, 0.11084487, 0.11819984, 0.1247251, 0.13037812, 0.13512554, 0.13894336, 0.14181704, 0.14374148, 0.1447209, 0.14476863, 0.14390678, 0.14216587, 0.13958427, 0.13620771, 0.13208858, 0.12728523, 0.12186124, 0.1158846, 0.10942687, 0.10256232, 0.095367076, 0.087918235, 0.080292992, 0.072567801, 0.064817553], [0.14985551, 0.15512305, 0.15931467, 0.16236806, 0.16423291, 0.16487165, 0.16426009, 0.1623879, 0.1592589, 0.15489114, 0.14931693, 0.14258255, 0.13474785, 0.1258857, 0.11608124, 0.10543095, 0.094041635, 0.082029213, 0.069517411, 0.056636348, 0.043521028, 0.030309756, 0.017142511, 0.0041592774, -0.0085016282, -0.020705223, -0.032321494, -0.043226982, -0.053306291, -0.062453515, -0.070573544, -0.077583253, -0.083412547, -0.088005244, -0.091319802, -0.093329861, -0.094024602, -0.093408915, -0.091503383, -0.08834406, -0.08398207, -0.078483012, -0.071926192, -0.064403681, -0.056019215, -0.046886954, -0.037130106, -0.026879442, -0.016271713, -0.005448, 0.005448, 0.016271713, 0.026879442, 0.037130106, 0.046886954, 0.056019215, 0.064403681, 0.071926192, 0.078483012, 0.08398207, 0.08834406, 0.091503383, 0.093408915, 0.094024602, 0.093329861, 0.091319802, 0.088005244, 0.083412547, 0.077583253, 0.070573544, 0.062453515, 0.053306291, 0.043226982, 0.032321494, 0.020705223, 0.0085016282, -0.0041592774, -0.017142511, -0.030309756, -0.043521028, -0.056636348, -0.069517411, -0.082029213, -0.094041635, -0.10543095, -0.11608124, -0.1258857, -0.13474785, -0.14258255, -0.14931693, -0.15489114, -0.1592589, -0.1623879, -0.16426009, -0.16487165, -0.16423291, -0.16236806, -0.15931467, -0.15512305, -0.14985551]], [0.999943140, 0.997571533, 0.959465463, 0.721862496]), # noqa: E501 +} + + +class TestDPSS: + + def test_basic(self): + # Test against hardcoded data + for k, v in dpss_data.items(): + win, ratios = windows.dpss(*k, return_ratios=True) + assert_allclose(win, v[0], atol=1e-7, err_msg=k) + assert_allclose(ratios, v[1], rtol=1e-5, atol=1e-7, err_msg=k) + + def test_unity(self): + # Test unity value handling (gh-2221) + for M in range(1, 21): + # corrected w/approximation (default) + win = windows.dpss(M, M / 2.1) + expected = M % 2 # one for odd, none for even + assert_equal(np.isclose(win, 1.).sum(), expected, + err_msg=f'{win}') + # corrected w/subsample delay (slower) + win_sub = windows.dpss(M, M / 2.1, norm='subsample') + if M > 2: + # @M=2 the subsample doesn't do anything + assert_equal(np.isclose(win_sub, 1.).sum(), expected, + err_msg=f'{win_sub}') + assert_allclose(win, win_sub, rtol=0.03) # within 3% + # not the same, l2-norm + win_2 = windows.dpss(M, M / 2.1, norm=2) + expected = 1 if M == 1 else 0 + assert_equal(np.isclose(win_2, 1.).sum(), expected, + err_msg=f'{win_2}') + + def test_extremes(self): + # Test extremes of alpha + lam = windows.dpss(31, 6, 4, return_ratios=True)[1] + assert_array_almost_equal(lam, 1.) + lam = windows.dpss(31, 7, 4, return_ratios=True)[1] + assert_array_almost_equal(lam, 1.) + lam = windows.dpss(31, 8, 4, return_ratios=True)[1] + assert_array_almost_equal(lam, 1.) + + def test_degenerate(self): + # Test failures + assert_raises(ValueError, windows.dpss, 4, 1.5, -1) # Bad Kmax + assert_raises(ValueError, windows.dpss, 4, 1.5, -5) + assert_raises(TypeError, windows.dpss, 4, 1.5, 1.1) + assert_raises(ValueError, windows.dpss, 3, 1.5, 3) # NW must be < N/2. + assert_raises(ValueError, windows.dpss, 3, -1, 3) # NW must be pos + assert_raises(ValueError, windows.dpss, 3, 0, 3) + assert_raises(ValueError, windows.dpss, -1, 1, 3) # negative M + + +class TestLanczos: + + def test_basic(self): + # Analytical results: + # sinc(x) = sinc(-x) + # sinc(pi) = 0, sinc(0) = 1 + # Hand computation on WolframAlpha: + # sinc(2 pi / 3) = 0.413496672 + # sinc(pi / 3) = 0.826993343 + # sinc(3 pi / 5) = 0.504551152 + # sinc(pi / 5) = 0.935489284 + assert_allclose(windows.lanczos(6, sym=False), + [0., 0.413496672, + 0.826993343, 1., 0.826993343, + 0.413496672], + atol=1e-9) + assert_allclose(windows.lanczos(6), + [0., 0.504551152, + 0.935489284, 0.935489284, + 0.504551152, 0.], + atol=1e-9) + assert_allclose(windows.lanczos(7, sym=True), + [0., 0.413496672, + 0.826993343, 1., 0.826993343, + 0.413496672, 0.], + atol=1e-9) + + def test_array_size(self): + for n in [0, 10, 11]: + assert_equal(len(windows.lanczos(n, sym=False)), n) + assert_equal(len(windows.lanczos(n, sym=True)), n) + + +class TestGetWindow: + + def test_boxcar(self): + w = windows.get_window('boxcar', 12) + assert_array_equal(w, np.ones_like(w)) + + # window is a tuple of len 1 + w = windows.get_window(('boxcar',), 16) + assert_array_equal(w, np.ones_like(w)) + + def test_cheb_odd(self): + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + w = windows.get_window(('chebwin', -40), 53, fftbins=False) + assert_array_almost_equal(w, cheb_odd_true, decimal=4) + + def test_cheb_even(self): + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + w = windows.get_window(('chebwin', 40), 54, fftbins=False) + assert_array_almost_equal(w, cheb_even_true, decimal=4) + + def test_dpss(self): + win1 = windows.get_window(('dpss', 3), 64, fftbins=False) + win2 = windows.dpss(64, 3) + assert_array_almost_equal(win1, win2, decimal=4) + + def test_kaiser_float(self): + win1 = windows.get_window(7.2, 64) + win2 = windows.kaiser(64, 7.2, False) + assert_allclose(win1, win2) + + def test_invalid_inputs(self): + # Window is not a float, tuple, or string + assert_raises(ValueError, windows.get_window, set('hann'), 8) + + # Unknown window type error + assert_raises(ValueError, windows.get_window, 'broken', 4) + + def test_array_as_window(self): + # GitHub issue 3603 + osfactor = 128 + sig = np.arange(128) + + win = windows.get_window(('kaiser', 8.0), osfactor // 2) + with assert_raises(ValueError, match='must have the same length'): + resample(sig, len(sig) * osfactor, window=win) + + def test_general_cosine(self): + assert_allclose(get_window(('general_cosine', [0.5, 0.3, 0.2]), 4), + [0.4, 0.3, 1, 0.3]) + assert_allclose(get_window(('general_cosine', [0.5, 0.3, 0.2]), 4, + fftbins=False), + [0.4, 0.55, 0.55, 0.4]) + + def test_general_hamming(self): + assert_allclose(get_window(('general_hamming', 0.7), 5), + [0.4, 0.6072949, 0.9427051, 0.9427051, 0.6072949]) + assert_allclose(get_window(('general_hamming', 0.7), 5, fftbins=False), + [0.4, 0.7, 1.0, 0.7, 0.4]) + + def test_lanczos(self): + assert_allclose(get_window('lanczos', 6), + [0., 0.413496672, 0.826993343, 1., 0.826993343, + 0.413496672], atol=1e-9) + assert_allclose(get_window('lanczos', 6, fftbins=False), + [0., 0.504551152, 0.935489284, 0.935489284, + 0.504551152, 0.], atol=1e-9) + assert_allclose(get_window('lanczos', 6), get_window('sinc', 6)) + + +def test_windowfunc_basics(): + for window_name, params in window_funcs: + window = getattr(windows, window_name) + with suppress_warnings() as sup: + sup.filter(UserWarning, "This window is not suitable") + # Check symmetry for odd and even lengths + w1 = window(8, *params, sym=True) + w2 = window(7, *params, sym=False) + assert_array_almost_equal(w1[:-1], w2) + + w1 = window(9, *params, sym=True) + w2 = window(8, *params, sym=False) + assert_array_almost_equal(w1[:-1], w2) + + # Check that functions run and output lengths are correct + assert_equal(len(window(6, *params, sym=True)), 6) + assert_equal(len(window(6, *params, sym=False)), 6) + assert_equal(len(window(7, *params, sym=True)), 7) + assert_equal(len(window(7, *params, sym=False)), 7) + + # Check invalid lengths + assert_raises(ValueError, window, 5.5, *params) + assert_raises(ValueError, window, -7, *params) + + # Check degenerate cases + assert_array_equal(window(0, *params, sym=True), []) + assert_array_equal(window(0, *params, sym=False), []) + assert_array_equal(window(1, *params, sym=True), [1]) + assert_array_equal(window(1, *params, sym=False), [1]) + + # Check dtype + assert_(window(0, *params, sym=True).dtype == 'float') + assert_(window(0, *params, sym=False).dtype == 'float') + assert_(window(1, *params, sym=True).dtype == 'float') + assert_(window(1, *params, sym=False).dtype == 'float') + assert_(window(6, *params, sym=True).dtype == 'float') + assert_(window(6, *params, sym=False).dtype == 'float') + + # Check normalization + assert_array_less(window(10, *params, sym=True), 1.01) + assert_array_less(window(10, *params, sym=False), 1.01) + assert_array_less(window(9, *params, sym=True), 1.01) + assert_array_less(window(9, *params, sym=False), 1.01) + + # Check that DFT-even spectrum is purely real for odd and even + assert_allclose(fft(window(10, *params, sym=False)).imag, + 0, atol=1e-14) + assert_allclose(fft(window(11, *params, sym=False)).imag, + 0, atol=1e-14) + + +def test_needs_params(): + for winstr in ['kaiser', 'ksr', 'kaiser_bessel_derived', 'kbd', + 'gaussian', 'gauss', 'gss', + 'general gaussian', 'general_gaussian', + 'general gauss', 'general_gauss', 'ggs', + 'dss', 'dpss', 'general cosine', 'general_cosine', + 'chebwin', 'cheb', 'general hamming', 'general_hamming', + ]: + assert_raises(ValueError, get_window, winstr, 7) + + +def test_not_needs_params(): + for winstr in ['barthann', + 'bartlett', + 'blackman', + 'blackmanharris', + 'bohman', + 'boxcar', + 'cosine', + 'flattop', + 'hamming', + 'nuttall', + 'parzen', + 'taylor', + 'exponential', + 'poisson', + 'tukey', + 'tuk', + 'triangle', + 'lanczos', + 'sinc', + ]: + win = get_window(winstr, 7) + assert_equal(len(win), 7) + + +def test_symmetric(): + + for win in [windows.lanczos]: + # Even sampling points + w = win(4096) + error = np.max(np.abs(w-np.flip(w))) + assert_equal(error, 0.0) + + # Odd sampling points + w = win(4097) + error = np.max(np.abs(w-np.flip(w))) + assert_equal(error, 0.0) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/windows/_windows.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/windows/_windows.py new file mode 100644 index 0000000000000000000000000000000000000000..4e661a8557ccc3660a32fe3007c74502f7f634dd --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/windows/_windows.py @@ -0,0 +1,2374 @@ +"""The suite of window functions.""" + +import operator +import warnings + +import numpy as np +from scipy import linalg, special, fft as sp_fft + +__all__ = ['boxcar', 'triang', 'parzen', 'bohman', 'blackman', 'nuttall', + 'blackmanharris', 'flattop', 'bartlett', 'barthann', + 'hamming', 'kaiser', 'kaiser_bessel_derived', 'gaussian', + 'general_cosine', 'general_gaussian', 'general_hamming', + 'chebwin', 'cosine', 'hann', 'exponential', 'tukey', 'taylor', + 'dpss', 'get_window', 'lanczos'] + + +def _len_guards(M): + """Handle small or incorrect window lengths""" + if int(M) != M or M < 0: + raise ValueError('Window length M must be a non-negative integer') + return M <= 1 + + +def _extend(M, sym): + """Extend window by 1 sample if needed for DFT-even symmetry""" + if not sym: + return M + 1, True + else: + return M, False + + +def _truncate(w, needed): + """Truncate window by 1 sample if needed for DFT-even symmetry""" + if needed: + return w[:-1] + else: + return w + + +def general_cosine(M, a, sym=True): + r""" + Generic weighted sum of cosine terms window + + Parameters + ---------- + M : int + Number of points in the output window + a : array_like + Sequence of weighting coefficients. This uses the convention of being + centered on the origin, so these will typically all be positive + numbers, not alternating sign. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The array of window values. + + References + ---------- + .. [1] A. Nuttall, "Some windows with very good sidelobe behavior," IEEE + Transactions on Acoustics, Speech, and Signal Processing, vol. 29, + no. 1, pp. 84-91, Feb 1981. :doi:`10.1109/TASSP.1981.1163506`. + .. [2] Heinzel G. et al., "Spectrum and spectral density estimation by the + Discrete Fourier transform (DFT), including a comprehensive list of + window functions and some new flat-top windows", February 15, 2002 + https://holometer.fnal.gov/GH_FFT.pdf + + Examples + -------- + Heinzel describes a flat-top window named "HFT90D" with formula: [2]_ + + .. math:: w_j = 1 - 1.942604 \cos(z) + 1.340318 \cos(2z) + - 0.440811 \cos(3z) + 0.043097 \cos(4z) + + where + + .. math:: z = \frac{2 \pi j}{N}, j = 0...N - 1 + + Since this uses the convention of starting at the origin, to reproduce the + window, we need to convert every other coefficient to a positive number: + + >>> HFT90D = [1, 1.942604, 1.340318, 0.440811, 0.043097] + + The paper states that the highest sidelobe is at -90.2 dB. Reproduce + Figure 42 by plotting the window and its frequency response, and confirm + the sidelobe level in red: + + >>> import numpy as np + >>> from scipy.signal.windows import general_cosine + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = general_cosine(1000, HFT90D, sym=False) + >>> plt.plot(window) + >>> plt.title("HFT90D window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 10000) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = np.abs(fftshift(A / abs(A).max())) + >>> response = 20 * np.log10(np.maximum(response, 1e-10)) + >>> plt.plot(freq, response) + >>> plt.axis([-50/1000, 50/1000, -140, 0]) + >>> plt.title("Frequency response of the HFT90D window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + >>> plt.axhline(-90.2, color='red') + >>> plt.show() + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + fac = np.linspace(-np.pi, np.pi, M) + w = np.zeros(M) + for k in range(len(a)): + w += a[k] * np.cos(k * fac) + + return _truncate(w, needs_trunc) + + +def boxcar(M, sym=True): + """Return a boxcar or rectangular window. + + Also known as a rectangular window or Dirichlet window, this is equivalent + to no window at all. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + Whether the window is symmetric. (Has no effect for boxcar.) + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.boxcar(51) + >>> plt.plot(window) + >>> plt.title("Boxcar window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the boxcar window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + w = np.ones(M, float) + + return _truncate(w, needs_trunc) + + +def triang(M, sym=True): + """Return a triangular window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + See Also + -------- + bartlett : A triangular window that touches zero + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.triang(51) + >>> plt.plot(window) + >>> plt.title("Triangular window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = np.abs(fftshift(A / abs(A).max())) + >>> response = 20 * np.log10(np.maximum(response, 1e-10)) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the triangular window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(1, (M + 1) // 2 + 1) + if M % 2 == 0: + w = (2 * n - 1.0) / M + w = np.r_[w, w[::-1]] + else: + w = 2 * n / (M + 1.0) + w = np.r_[w, w[-2::-1]] + + return _truncate(w, needs_trunc) + + +def parzen(M, sym=True): + """Return a Parzen window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + References + ---------- + .. [1] E. Parzen, "Mathematical Considerations in the Estimation of + Spectra", Technometrics, Vol. 3, No. 2 (May, 1961), pp. 167-190 + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.parzen(51) + >>> plt.plot(window) + >>> plt.title("Parzen window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Parzen window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(-(M - 1) / 2.0, (M - 1) / 2.0 + 0.5, 1.0) + na = np.extract(n < -(M - 1) / 4.0, n) + nb = np.extract(abs(n) <= (M - 1) / 4.0, n) + wa = 2 * (1 - np.abs(na) / (M / 2.0)) ** 3.0 + wb = (1 - 6 * (np.abs(nb) / (M / 2.0)) ** 2.0 + + 6 * (np.abs(nb) / (M / 2.0)) ** 3.0) + w = np.r_[wa, wb, wa[::-1]] + + return _truncate(w, needs_trunc) + + +def bohman(M, sym=True): + """Return a Bohman window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.bohman(51) + >>> plt.plot(window) + >>> plt.title("Bohman window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2047) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Bohman window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + fac = np.abs(np.linspace(-1, 1, M)[1:-1]) + w = (1 - fac) * np.cos(np.pi * fac) + 1.0 / np.pi * np.sin(np.pi * fac) + w = np.r_[0, w, 0] + + return _truncate(w, needs_trunc) + + +def blackman(M, sym=True): + r""" + Return a Blackman window. + + The Blackman window is a taper formed by using the first three terms of + a summation of cosines. It was designed to have close to the minimal + leakage possible. It is close to optimal, only slightly worse than a + Kaiser window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Blackman window is defined as + + .. math:: w(n) = 0.42 - 0.5 \cos(2\pi n/M) + 0.08 \cos(4\pi n/M) + + The "exact Blackman" window was designed to null out the third and fourth + sidelobes, but has discontinuities at the boundaries, resulting in a + 6 dB/oct fall-off. This window is an approximation of the "exact" window, + which does not null the sidelobes as well, but is smooth at the edges, + improving the fall-off rate to 18 dB/oct. [3]_ + + Most references to the Blackman window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. It is known as a + "near optimal" tapering function, almost as good (by some measures) + as the Kaiser window. + + References + ---------- + .. [1] Blackman, R.B. and Tukey, J.W., (1958) The measurement of power + spectra, Dover Publications, New York. + .. [2] Oppenheim, A.V., and R.W. Schafer. Discrete-Time Signal Processing. + Upper Saddle River, NJ: Prentice-Hall, 1999, pp. 468-471. + .. [3] Harris, Fredric J. (Jan 1978). "On the use of Windows for Harmonic + Analysis with the Discrete Fourier Transform". Proceedings of the + IEEE 66 (1): 51-83. :doi:`10.1109/PROC.1978.10837`. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.blackman(51) + >>> plt.plot(window) + >>> plt.title("Blackman window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = np.abs(fftshift(A / abs(A).max())) + >>> response = 20 * np.log10(np.maximum(response, 1e-10)) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Blackman window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + # Docstring adapted from NumPy's blackman function + return general_cosine(M, [0.42, 0.50, 0.08], sym) + + +def nuttall(M, sym=True): + """Return a minimum 4-term Blackman-Harris window according to Nuttall. + + This variation is called "Nuttall4c" by Heinzel. [2]_ + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + References + ---------- + .. [1] A. Nuttall, "Some windows with very good sidelobe behavior," IEEE + Transactions on Acoustics, Speech, and Signal Processing, vol. 29, + no. 1, pp. 84-91, Feb 1981. :doi:`10.1109/TASSP.1981.1163506`. + .. [2] Heinzel G. et al., "Spectrum and spectral density estimation by the + Discrete Fourier transform (DFT), including a comprehensive list of + window functions and some new flat-top windows", February 15, 2002 + https://holometer.fnal.gov/GH_FFT.pdf + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.nuttall(51) + >>> plt.plot(window) + >>> plt.title("Nuttall window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Nuttall window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + return general_cosine(M, [0.3635819, 0.4891775, 0.1365995, 0.0106411], sym) + + +def blackmanharris(M, sym=True): + """Return a minimum 4-term Blackman-Harris window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.blackmanharris(51) + >>> plt.plot(window) + >>> plt.title("Blackman-Harris window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Blackman-Harris window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + return general_cosine(M, [0.35875, 0.48829, 0.14128, 0.01168], sym) + + +def flattop(M, sym=True): + """Return a flat top window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + Flat top windows are used for taking accurate measurements of signal + amplitude in the frequency domain, with minimal scalloping error from the + center of a frequency bin to its edges, compared to others. This is a + 5th-order cosine window, with the 5 terms optimized to make the main lobe + maximally flat. [1]_ + + References + ---------- + .. [1] D'Antona, Gabriele, and A. Ferrero, "Digital Signal Processing for + Measurement Systems", Springer Media, 2006, p. 70 + :doi:`10.1007/0-387-28666-7`. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.flattop(51) + >>> plt.plot(window) + >>> plt.title("Flat top window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the flat top window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + a = [0.21557895, 0.41663158, 0.277263158, 0.083578947, 0.006947368] + return general_cosine(M, a, sym) + + +def bartlett(M, sym=True): + r""" + Return a Bartlett window. + + The Bartlett window is very similar to a triangular window, except + that the end points are at zero. It is often used in signal + processing for tapering a signal, without generating too much + ripple in the frequency domain. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The triangular window, with the first and last samples equal to zero + and the maximum value normalized to 1 (though the value 1 does not + appear if `M` is even and `sym` is True). + + See Also + -------- + triang : A triangular window that does not touch zero at the ends + + Notes + ----- + The Bartlett window is defined as + + .. math:: w(n) = \frac{2}{M-1} \left( + \frac{M-1}{2} - \left|n - \frac{M-1}{2}\right| + \right) + + Most references to the Bartlett window come from the signal + processing literature, where it is used as one of many windowing + functions for smoothing values. Note that convolution with this + window produces linear interpolation. It is also known as an + apodization (which means"removing the foot", i.e. smoothing + discontinuities at the beginning and end of the sampled signal) or + tapering function. The Fourier transform of the Bartlett is the product + of two sinc functions. + Note the excellent discussion in Kanasewich. [2]_ + + References + ---------- + .. [1] M.S. Bartlett, "Periodogram Analysis and Continuous Spectra", + Biometrika 37, 1-16, 1950. + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", + The University of Alberta Press, 1975, pp. 109-110. + .. [3] A.V. Oppenheim and R.W. Schafer, "Discrete-Time Signal + Processing", Prentice-Hall, 1999, pp. 468-471. + .. [4] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [5] W.H. Press, B.P. Flannery, S.A. Teukolsky, and W.T. Vetterling, + "Numerical Recipes", Cambridge University Press, 1986, page 429. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.bartlett(51) + >>> plt.plot(window) + >>> plt.title("Bartlett window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Bartlett window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + # Docstring adapted from NumPy's bartlett function + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(0, M) + w = np.where(np.less_equal(n, (M - 1) / 2.0), + 2.0 * n / (M - 1), 2.0 - 2.0 * n / (M - 1)) + + return _truncate(w, needs_trunc) + + +def hann(M, sym=True): + r""" + Return a Hann window. + + The Hann window is a taper formed by using a raised cosine or sine-squared + with ends that touch zero. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Hann window is defined as + + .. math:: w(n) = 0.5 - 0.5 \cos\left(\frac{2\pi{n}}{M-1}\right) + \qquad 0 \leq n \leq M-1 + + The window was named for Julius von Hann, an Austrian meteorologist. It is + also known as the Cosine Bell. It is sometimes erroneously referred to as + the "Hanning" window, from the use of "hann" as a verb in the original + paper and confusion with the very similar Hamming window. + + Most references to the Hann window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. + + References + ---------- + .. [1] Blackman, R.B. and Tukey, J.W., (1958) The measurement of power + spectra, Dover Publications, New York. + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", + The University of Alberta Press, 1975, pp. 106-108. + .. [3] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [4] W.H. Press, B.P. Flannery, S.A. Teukolsky, and W.T. Vetterling, + "Numerical Recipes", Cambridge University Press, 1986, page 425. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.hann(51) + >>> plt.plot(window) + >>> plt.title("Hann window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = np.abs(fftshift(A / abs(A).max())) + >>> response = 20 * np.log10(np.maximum(response, 1e-10)) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Hann window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + # Docstring adapted from NumPy's hanning function + return general_hamming(M, 0.5, sym) + + +def tukey(M, alpha=0.5, sym=True): + r"""Return a Tukey window, also known as a tapered cosine window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + alpha : float, optional + Shape parameter of the Tukey window, representing the fraction of the + window inside the cosine tapered region. + If zero, the Tukey window is equivalent to a rectangular window. + If one, the Tukey window is equivalent to a Hann window. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + References + ---------- + .. [1] Harris, Fredric J. (Jan 1978). "On the use of Windows for Harmonic + Analysis with the Discrete Fourier Transform". Proceedings of the + IEEE 66 (1): 51-83. :doi:`10.1109/PROC.1978.10837` + .. [2] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function#Tukey_window + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.tukey(51) + >>> plt.plot(window) + >>> plt.title("Tukey window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + >>> plt.ylim([0, 1.1]) + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Tukey window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + + if alpha <= 0: + return np.ones(M, 'd') + elif alpha >= 1.0: + return hann(M, sym=sym) + + M, needs_trunc = _extend(M, sym) + + n = np.arange(0, M) + width = int(np.floor(alpha*(M-1)/2.0)) + n1 = n[0:width+1] + n2 = n[width+1:M-width-1] + n3 = n[M-width-1:] + + w1 = 0.5 * (1 + np.cos(np.pi * (-1 + 2.0*n1/alpha/(M-1)))) + w2 = np.ones(n2.shape) + w3 = 0.5 * (1 + np.cos(np.pi * (-2.0/alpha + 1 + 2.0*n3/alpha/(M-1)))) + + w = np.concatenate((w1, w2, w3)) + + return _truncate(w, needs_trunc) + + +def barthann(M, sym=True): + """Return a modified Bartlett-Hann window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.barthann(51) + >>> plt.plot(window) + >>> plt.title("Bartlett-Hann window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Bartlett-Hann window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(0, M) + fac = np.abs(n / (M - 1.0) - 0.5) + w = 0.62 - 0.48 * fac + 0.38 * np.cos(2 * np.pi * fac) + + return _truncate(w, needs_trunc) + + +def general_hamming(M, alpha, sym=True): + r"""Return a generalized Hamming window. + + The generalized Hamming window is constructed by multiplying a rectangular + window by one period of a cosine function [1]_. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + alpha : float + The window coefficient, :math:`\alpha` + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + See Also + -------- + hamming, hann + + Notes + ----- + The generalized Hamming window is defined as + + .. math:: w(n) = \alpha - \left(1 - \alpha\right) + \cos\left(\frac{2\pi{n}}{M-1}\right) \qquad 0 \leq n \leq M-1 + + Both the common Hamming window and Hann window are special cases of the + generalized Hamming window with :math:`\alpha` = 0.54 and :math:`\alpha` = + 0.5, respectively [2]_. + + References + ---------- + .. [1] DSPRelated, "Generalized Hamming Window Family", + https://www.dsprelated.com/freebooks/sasp/Generalized_Hamming_Window_Family.html + .. [2] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [3] Riccardo Piantanida ESA, "Sentinel-1 Level 1 Detailed Algorithm + Definition", + https://sentinel.esa.int/documents/247904/1877131/Sentinel-1-Level-1-Detailed-Algorithm-Definition + .. [4] Matthieu Bourbigot ESA, "Sentinel-1 Product Definition", + https://sentinel.esa.int/documents/247904/1877131/Sentinel-1-Product-Definition + + Examples + -------- + The Sentinel-1A/B Instrument Processing Facility uses generalized Hamming + windows in the processing of spaceborne Synthetic Aperture Radar (SAR) + data [3]_. The facility uses various values for the :math:`\alpha` + parameter based on operating mode of the SAR instrument. Some common + :math:`\alpha` values include 0.75, 0.7 and 0.52 [4]_. As an example, we + plot these different windows. + + >>> import numpy as np + >>> from scipy.signal.windows import general_hamming + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> fig1, spatial_plot = plt.subplots() + >>> spatial_plot.set_title("Generalized Hamming Windows") + >>> spatial_plot.set_ylabel("Amplitude") + >>> spatial_plot.set_xlabel("Sample") + + >>> fig2, freq_plot = plt.subplots() + >>> freq_plot.set_title("Frequency Responses") + >>> freq_plot.set_ylabel("Normalized magnitude [dB]") + >>> freq_plot.set_xlabel("Normalized frequency [cycles per sample]") + + >>> for alpha in [0.75, 0.7, 0.52]: + ... window = general_hamming(41, alpha) + ... spatial_plot.plot(window, label="{:.2f}".format(alpha)) + ... A = fft(window, 2048) / (len(window)/2.0) + ... freq = np.linspace(-0.5, 0.5, len(A)) + ... response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + ... freq_plot.plot(freq, response, label="{:.2f}".format(alpha)) + >>> freq_plot.legend(loc="upper right") + >>> spatial_plot.legend(loc="upper right") + + """ + return general_cosine(M, [alpha, 1. - alpha], sym) + + +def hamming(M, sym=True): + r"""Return a Hamming window. + + The Hamming window is a taper formed by using a raised cosine with + non-zero endpoints, optimized to minimize the nearest side lobe. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Hamming window is defined as + + .. math:: w(n) = 0.54 - 0.46 \cos\left(\frac{2\pi{n}}{M-1}\right) + \qquad 0 \leq n \leq M-1 + + The Hamming was named for R. W. Hamming, an associate of J. W. Tukey and + is described in Blackman and Tukey. It was recommended for smoothing the + truncated autocovariance function in the time domain. + Most references to the Hamming window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. + + References + ---------- + .. [1] Blackman, R.B. and Tukey, J.W., (1958) The measurement of power + spectra, Dover Publications, New York. + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", The + University of Alberta Press, 1975, pp. 109-110. + .. [3] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [4] W.H. Press, B.P. Flannery, S.A. Teukolsky, and W.T. Vetterling, + "Numerical Recipes", Cambridge University Press, 1986, page 425. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.hamming(51) + >>> plt.plot(window) + >>> plt.title("Hamming window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Hamming window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + # Docstring adapted from NumPy's hamming function + return general_hamming(M, 0.54, sym) + + +def kaiser(M, beta, sym=True): + r"""Return a Kaiser window. + + The Kaiser window is a taper formed by using a Bessel function. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + beta : float + Shape parameter, determines trade-off between main-lobe width and + side lobe level. As beta gets large, the window narrows. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Kaiser window is defined as + + .. math:: w(n) = I_0\left( \beta \sqrt{1-\frac{4n^2}{(M-1)^2}} + \right)/I_0(\beta) + + with + + .. math:: \quad -\frac{M-1}{2} \leq n \leq \frac{M-1}{2}, + + where :math:`I_0` is the modified zeroth-order Bessel function. + + The Kaiser was named for Jim Kaiser, who discovered a simple approximation + to the DPSS window based on Bessel functions. + The Kaiser window is a very good approximation to the discrete prolate + spheroidal sequence, or Slepian window, which is the transform which + maximizes the energy in the main lobe of the window relative to total + energy. + + The Kaiser can approximate other windows by varying the beta parameter. + (Some literature uses alpha = beta/pi.) [4]_ + + ==== ======================= + beta Window shape + ==== ======================= + 0 Rectangular + 5 Similar to a Hamming + 6 Similar to a Hann + 8.6 Similar to a Blackman + ==== ======================= + + A beta value of 14 is probably a good starting point. Note that as beta + gets large, the window narrows, and so the number of samples needs to be + large enough to sample the increasingly narrow spike, otherwise NaNs will + be returned. + + Most references to the Kaiser window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. + + References + ---------- + .. [1] J. F. Kaiser, "Digital Filters" - Ch 7 in "Systems analysis by + digital computer", Editors: F.F. Kuo and J.F. Kaiser, p 218-285. + John Wiley and Sons, New York, (1966). + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", The + University of Alberta Press, 1975, pp. 177-178. + .. [3] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [4] F. J. Harris, "On the use of windows for harmonic analysis with the + discrete Fourier transform," Proceedings of the IEEE, vol. 66, + no. 1, pp. 51-83, Jan. 1978. :doi:`10.1109/PROC.1978.10837`. + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.kaiser(51, beta=14) + >>> plt.plot(window) + >>> plt.title(r"Kaiser window ($\beta$=14)") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title(r"Frequency response of the Kaiser window ($\beta$=14)") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + # Docstring adapted from NumPy's kaiser function + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(0, M) + alpha = (M - 1) / 2.0 + w = (special.i0(beta * np.sqrt(1 - ((n - alpha) / alpha) ** 2.0)) / + special.i0(beta)) + + return _truncate(w, needs_trunc) + + +def kaiser_bessel_derived(M, beta, *, sym=True): + """Return a Kaiser-Bessel derived window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + Note that this window is only defined for an even + number of points. + beta : float + Kaiser window shape parameter. + sym : bool, optional + This parameter only exists to comply with the interface offered by + the other window functions and to be callable by `get_window`. + When True (default), generates a symmetric window, for use in filter + design. + + Returns + ------- + w : ndarray + The window, normalized to fulfil the Princen-Bradley condition. + + See Also + -------- + kaiser + + Notes + ----- + It is designed to be suitable for use with the modified discrete cosine + transform (MDCT) and is mainly used in audio signal processing and + audio coding. + + .. versionadded:: 1.9.0 + + References + ---------- + .. [1] Bosi, Marina, and Richard E. Goldberg. Introduction to Digital + Audio Coding and Standards. Dordrecht: Kluwer, 2003. + .. [2] Wikipedia, "Kaiser window", + https://en.wikipedia.org/wiki/Kaiser_window + + Examples + -------- + Plot the Kaiser-Bessel derived window based on the wikipedia + reference [2]_: + + >>> import numpy as np + >>> from scipy import signal + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots() + >>> N = 50 + >>> for alpha in [0.64, 2.55, 7.64, 31.83]: + ... ax.plot(signal.windows.kaiser_bessel_derived(2*N, np.pi*alpha), + ... label=f"{alpha=}") + >>> ax.grid(True) + >>> ax.set_title("Kaiser-Bessel derived window") + >>> ax.set_ylabel("Amplitude") + >>> ax.set_xlabel("Sample") + >>> ax.set_xticks([0, N, 2*N-1]) + >>> ax.set_xticklabels(["0", "N", "2N+1"]) # doctest: +SKIP + >>> ax.set_yticks([0.0, 0.2, 0.4, 0.6, 0.707, 0.8, 1.0]) + >>> fig.legend(loc="center") + >>> fig.tight_layout() + >>> fig.show() + """ + if not sym: + raise ValueError( + "Kaiser-Bessel Derived windows are only defined for symmetric " + "shapes" + ) + elif M < 1: + return np.array([]) + elif M % 2: + raise ValueError( + "Kaiser-Bessel Derived windows are only defined for even number " + "of points" + ) + + kaiser_window = kaiser(M // 2 + 1, beta) + csum = np.cumsum(kaiser_window) + half_window = np.sqrt(csum[:-1] / csum[-1]) + w = np.concatenate((half_window, half_window[::-1]), axis=0) + return w + + +def gaussian(M, std, sym=True): + r"""Return a Gaussian window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + std : float + The standard deviation, sigma. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Gaussian window is defined as + + .. math:: w(n) = e^{ -\frac{1}{2}\left(\frac{n}{\sigma}\right)^2 } + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.gaussian(51, std=7) + >>> plt.plot(window) + >>> plt.title(r"Gaussian window ($\sigma$=7)") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title(r"Frequency response of the Gaussian window ($\sigma$=7)") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(0, M) - (M - 1.0) / 2.0 + sig2 = 2 * std * std + w = np.exp(-n ** 2 / sig2) + + return _truncate(w, needs_trunc) + + +def general_gaussian(M, p, sig, sym=True): + r"""Return a window with a generalized Gaussian shape. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + p : float + Shape parameter. p = 1 is identical to `gaussian`, p = 0.5 is + the same shape as the Laplace distribution. + sig : float + The standard deviation, sigma. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The generalized Gaussian window is defined as + + .. math:: w(n) = e^{ -\frac{1}{2}\left|\frac{n}{\sigma}\right|^{2p} } + + the half-power point is at + + .. math:: (2 \log(2))^{1/(2 p)} \sigma + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.general_gaussian(51, p=1.5, sig=7) + >>> plt.plot(window) + >>> plt.title(r"Generalized Gaussian window (p=1.5, $\sigma$=7)") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title(r"Freq. resp. of the gen. Gaussian " + ... r"window (p=1.5, $\sigma$=7)") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + n = np.arange(0, M) - (M - 1.0) / 2.0 + w = np.exp(-0.5 * np.abs(n / sig) ** (2 * p)) + + return _truncate(w, needs_trunc) + + +# `chebwin` contributed by Kumar Appaiah. +def chebwin(M, at, sym=True): + r"""Return a Dolph-Chebyshev window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + at : float + Attenuation (in dB). + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value always normalized to 1 + + Notes + ----- + This window optimizes for the narrowest main lobe width for a given order + `M` and sidelobe equiripple attenuation `at`, using Chebyshev + polynomials. It was originally developed by Dolph to optimize the + directionality of radio antenna arrays. + + Unlike most windows, the Dolph-Chebyshev is defined in terms of its + frequency response: + + .. math:: W(k) = \frac + {\cos\{M \cos^{-1}[\beta \cos(\frac{\pi k}{M})]\}} + {\cosh[M \cosh^{-1}(\beta)]} + + where + + .. math:: \beta = \cosh \left [\frac{1}{M} + \cosh^{-1}(10^\frac{A}{20}) \right ] + + and 0 <= abs(k) <= M-1. A is the attenuation in decibels (`at`). + + The time domain window is then generated using the IFFT, so + power-of-two `M` are the fastest to generate, and prime number `M` are + the slowest. + + The equiripple condition in the frequency domain creates impulses in the + time domain, which appear at the ends of the window. + + References + ---------- + .. [1] C. Dolph, "A current distribution for broadside arrays which + optimizes the relationship between beam width and side-lobe level", + Proceedings of the IEEE, Vol. 34, Issue 6 + .. [2] Peter Lynch, "The Dolph-Chebyshev Window: A Simple Optimal Filter", + American Meteorological Society (April 1997) + http://mathsci.ucd.ie/~plynch/Publications/Dolph.pdf + .. [3] F. J. Harris, "On the use of windows for harmonic analysis with the + discrete Fourier transforms", Proceedings of the IEEE, Vol. 66, + No. 1, January 1978 + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.chebwin(51, at=100) + >>> plt.plot(window) + >>> plt.title("Dolph-Chebyshev window (100 dB)") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Dolph-Chebyshev window (100 dB)") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ + if np.abs(at) < 45: + warnings.warn("This window is not suitable for spectral analysis " + "for attenuation values lower than about 45dB because " + "the equivalent noise bandwidth of a Chebyshev window " + "does not grow monotonically with increasing sidelobe " + "attenuation when the attenuation is smaller than " + "about 45 dB.", + stacklevel=2) + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + # compute the parameter beta + order = M - 1.0 + beta = np.cosh(1.0 / order * np.arccosh(10 ** (np.abs(at) / 20.))) + k = np.r_[0:M] * 1.0 + x = beta * np.cos(np.pi * k / M) + # Find the window's DFT coefficients + # Use analytic definition of Chebyshev polynomial instead of expansion + # from scipy.special. Using the expansion in scipy.special leads to errors. + p = np.zeros(x.shape) + p[x > 1] = np.cosh(order * np.arccosh(x[x > 1])) + p[x < -1] = (2 * (M % 2) - 1) * np.cosh(order * np.arccosh(-x[x < -1])) + p[np.abs(x) <= 1] = np.cos(order * np.arccos(x[np.abs(x) <= 1])) + + # Appropriate IDFT and filling up + # depending on even/odd M + if M % 2: + w = np.real(sp_fft.fft(p)) + n = (M + 1) // 2 + w = w[:n] + w = np.concatenate((w[n - 1:0:-1], w)) + else: + p = p * np.exp(1.j * np.pi / M * np.r_[0:M]) + w = np.real(sp_fft.fft(p)) + n = M // 2 + 1 + w = np.concatenate((w[n - 1:0:-1], w[1:n])) + w = w / max(w) + + return _truncate(w, needs_trunc) + + +def cosine(M, sym=True): + """Return a window with a simple cosine shape. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + + .. versionadded:: 0.13.0 + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.cosine(51) + >>> plt.plot(window) + >>> plt.title("Cosine window") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2047) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the cosine window") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + >>> plt.show() + + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + w = np.sin(np.pi / M * (np.arange(0, M) + .5)) + + return _truncate(w, needs_trunc) + + +def exponential(M, center=None, tau=1., sym=True): + r"""Return an exponential (or Poisson) window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + center : float, optional + Parameter defining the center location of the window function. + The default value if not given is ``center = (M-1) / 2``. This + parameter must take its default value for symmetric windows. + tau : float, optional + Parameter defining the decay. For ``center = 0`` use + ``tau = -(M-1) / ln(x)`` if ``x`` is the fraction of the window + remaining at the end. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Exponential window is defined as + + .. math:: w(n) = e^{-|n-center| / \tau} + + References + ---------- + .. [1] S. Gade and H. Herlufsen, "Windows to FFT analysis (Part I)", + Technical Review 3, Bruel & Kjaer, 1987. + + Examples + -------- + Plot the symmetric window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> M = 51 + >>> tau = 3.0 + >>> window = signal.windows.exponential(M, tau=tau) + >>> plt.plot(window) + >>> plt.title("Exponential Window (tau=3.0)") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -35, 0]) + >>> plt.title("Frequency response of the Exponential window (tau=3.0)") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + This function can also generate non-symmetric windows: + + >>> tau2 = -(M-1) / np.log(0.01) + >>> window2 = signal.windows.exponential(M, 0, tau2, False) + >>> plt.figure() + >>> plt.plot(window2) + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + """ + if sym and center is not None: + raise ValueError("If sym==True, center must be None.") + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + if center is None: + center = (M-1) / 2 + + n = np.arange(0, M) + w = np.exp(-np.abs(n-center) / tau) + + return _truncate(w, needs_trunc) + + +def taylor(M, nbar=4, sll=30, norm=True, sym=True): + """ + Return a Taylor window. + + The Taylor window taper function approximates the Dolph-Chebyshev window's + constant sidelobe level for a parameterized number of near-in sidelobes, + but then allows a taper beyond [2]_. + + The SAR (synthetic aperture radar) community commonly uses Taylor + weighting for image formation processing because it provides strong, + selectable sidelobe suppression with minimum broadening of the + mainlobe [1]_. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + nbar : int, optional + Number of nearly constant level sidelobes adjacent to the mainlobe. + sll : float, optional + Desired suppression of sidelobe level in decibels (dB) relative to the + DC gain of the mainlobe. This should be a positive number. + norm : bool, optional + When True (default), divides the window by the largest (middle) value + for odd-length windows or the value that would occur between the two + repeated middle values for even-length windows such that all values + are less than or equal to 1. When False the DC gain will remain at 1 + (0 dB) and the sidelobes will be `sll` dB down. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + out : array + The window. When `norm` is True (default), the maximum value is + normalized to 1 (though the value 1 does not appear if `M` is + even and `sym` is True). + + See Also + -------- + chebwin, kaiser, bartlett, blackman, hamming, hann + + References + ---------- + .. [1] W. Carrara, R. Goodman, and R. Majewski, "Spotlight Synthetic + Aperture Radar: Signal Processing Algorithms" Pages 512-513, + July 1995. + .. [2] Armin Doerry, "Catalog of Window Taper Functions for + Sidelobe Control", 2017. + https://www.researchgate.net/profile/Armin_Doerry/publication/316281181_Catalog_of_Window_Taper_Functions_for_Sidelobe_Control/links/58f92cb2a6fdccb121c9d54d/Catalog-of-Window-Taper-Functions-for-Sidelobe-Control.pdf + + Examples + -------- + Plot the window and its frequency response: + + >>> import numpy as np + >>> from scipy import signal + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + + >>> window = signal.windows.taylor(51, nbar=20, sll=100, norm=False) + >>> plt.plot(window) + >>> plt.title("Taylor window (100 dB)") + >>> plt.ylabel("Amplitude") + >>> plt.xlabel("Sample") + + >>> plt.figure() + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> plt.plot(freq, response) + >>> plt.axis([-0.5, 0.5, -120, 0]) + >>> plt.title("Frequency response of the Taylor window (100 dB)") + >>> plt.ylabel("Normalized magnitude [dB]") + >>> plt.xlabel("Normalized frequency [cycles per sample]") + + """ # noqa: E501 + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + # Original text uses a negative sidelobe level parameter and then negates + # it in the calculation of B. To keep consistent with other methods we + # assume the sidelobe level parameter to be positive. + B = 10**(sll / 20) + A = np.arccosh(B) / np.pi + s2 = nbar**2 / (A**2 + (nbar - 0.5)**2) + ma = np.arange(1, nbar) + + Fm = np.empty(nbar-1) + signs = np.empty_like(ma) + signs[::2] = 1 + signs[1::2] = -1 + m2 = ma*ma + for mi, m in enumerate(ma): + numer = signs[mi] * np.prod(1 - m2[mi]/s2/(A**2 + (ma - 0.5)**2)) + denom = 2 * np.prod(1 - m2[mi]/m2[:mi]) * np.prod(1 - m2[mi]/m2[mi+1:]) + Fm[mi] = numer / denom + + def W(n): + return 1 + 2*np.dot(Fm, np.cos( + 2*np.pi*ma[:, np.newaxis]*(n-M/2.+0.5)/M)) + + w = W(np.arange(M)) + + # normalize (Note that this is not described in the original text [1]) + if norm: + scale = 1.0 / W((M - 1) / 2) + w *= scale + + return _truncate(w, needs_trunc) + + +def dpss(M, NW, Kmax=None, sym=True, norm=None, return_ratios=False): + """ + Compute the Discrete Prolate Spheroidal Sequences (DPSS). + + DPSS (or Slepian sequences) are often used in multitaper power spectral + density estimation (see [1]_). The first window in the sequence can be + used to maximize the energy concentration in the main lobe, and is also + called the Slepian window. + + Parameters + ---------- + M : int + Window length. + NW : float + Standardized half bandwidth corresponding to ``2*NW = BW/f0 = BW*M*dt`` + where ``dt`` is taken as 1. + Kmax : int | None, optional + Number of DPSS windows to return (orders ``0`` through ``Kmax-1``). + If None (default), return only a single window of shape ``(M,)`` + instead of an array of windows of shape ``(Kmax, M)``. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + norm : {2, 'approximate', 'subsample'} | None, optional + If 'approximate' or 'subsample', then the windows are normalized by the + maximum, and a correction scale-factor for even-length windows + is applied either using ``M**2/(M**2+NW)`` ("approximate") or + a FFT-based subsample shift ("subsample"), see Notes for details. + If None, then "approximate" is used when ``Kmax=None`` and 2 otherwise + (which uses the l2 norm). + return_ratios : bool, optional + If True, also return the concentration ratios in addition to the + windows. + + Returns + ------- + v : ndarray, shape (Kmax, M) or (M,) + The DPSS windows. Will be 1D if `Kmax` is None. + r : ndarray, shape (Kmax,) or float, optional + The concentration ratios for the windows. Only returned if + `return_ratios` evaluates to True. Will be 0D if `Kmax` is None. + + Notes + ----- + This computation uses the tridiagonal eigenvector formulation given + in [2]_. + + The default normalization for ``Kmax=None``, i.e. window-generation mode, + simply using the l-infinity norm would create a window with two unity + values, which creates slight normalization differences between even and odd + orders. The approximate correction of ``M**2/float(M**2+NW)`` for even + sample numbers is used to counteract this effect (see Examples below). + + For very long signals (e.g., 1e6 elements), it can be useful to compute + windows orders of magnitude shorter and use interpolation (e.g., + `scipy.interpolate.interp1d`) to obtain tapers of length `M`, + but this in general will not preserve orthogonality between the tapers. + + .. versionadded:: 1.1 + + References + ---------- + .. [1] Percival DB, Walden WT. Spectral Analysis for Physical Applications: + Multitaper and Conventional Univariate Techniques. + Cambridge University Press; 1993. + .. [2] Slepian, D. Prolate spheroidal wave functions, Fourier analysis, and + uncertainty V: The discrete case. Bell System Technical Journal, + Volume 57 (1978), 1371430. + .. [3] Kaiser, JF, Schafer RW. On the Use of the I0-Sinh Window for + Spectrum Analysis. IEEE Transactions on Acoustics, Speech and + Signal Processing. ASSP-28 (1): 105-107; 1980. + + Examples + -------- + We can compare the window to `kaiser`, which was invented as an alternative + that was easier to calculate [3]_ (example adapted from + `here `_): + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.signal import windows, freqz + >>> M = 51 + >>> fig, axes = plt.subplots(3, 2, figsize=(5, 7)) + >>> for ai, alpha in enumerate((1, 3, 5)): + ... win_dpss = windows.dpss(M, alpha) + ... beta = alpha*np.pi + ... win_kaiser = windows.kaiser(M, beta) + ... for win, c in ((win_dpss, 'k'), (win_kaiser, 'r')): + ... win /= win.sum() + ... axes[ai, 0].plot(win, color=c, lw=1.) + ... axes[ai, 0].set(xlim=[0, M-1], title=r'$\\alpha$ = %s' % alpha, + ... ylabel='Amplitude') + ... w, h = freqz(win) + ... axes[ai, 1].plot(w, 20 * np.log10(np.abs(h)), color=c, lw=1.) + ... axes[ai, 1].set(xlim=[0, np.pi], + ... title=r'$\\beta$ = %0.2f' % beta, + ... ylabel='Magnitude (dB)') + >>> for ax in axes.ravel(): + ... ax.grid(True) + >>> axes[2, 1].legend(['DPSS', 'Kaiser']) + >>> fig.tight_layout() + >>> plt.show() + + And here are examples of the first four windows, along with their + concentration ratios: + + >>> M = 512 + >>> NW = 2.5 + >>> win, eigvals = windows.dpss(M, NW, 4, return_ratios=True) + >>> fig, ax = plt.subplots(1) + >>> ax.plot(win.T, linewidth=1.) + >>> ax.set(xlim=[0, M-1], ylim=[-0.1, 0.1], xlabel='Samples', + ... title='DPSS, M=%d, NW=%0.1f' % (M, NW)) + >>> ax.legend(['win[%d] (%0.4f)' % (ii, ratio) + ... for ii, ratio in enumerate(eigvals)]) + >>> fig.tight_layout() + >>> plt.show() + + Using a standard :math:`l_{\\infty}` norm would produce two unity values + for even `M`, but only one unity value for odd `M`. This produces uneven + window power that can be counteracted by the approximate correction + ``M**2/float(M**2+NW)``, which can be selected by using + ``norm='approximate'`` (which is the same as ``norm=None`` when + ``Kmax=None``, as is the case here). Alternatively, the slower + ``norm='subsample'`` can be used, which uses subsample shifting in the + frequency domain (FFT) to compute the correction: + + >>> Ms = np.arange(1, 41) + >>> factors = (50, 20, 10, 5, 2.0001) + >>> energy = np.empty((3, len(Ms), len(factors))) + >>> for mi, M in enumerate(Ms): + ... for fi, factor in enumerate(factors): + ... NW = M / float(factor) + ... # Corrected using empirical approximation (default) + ... win = windows.dpss(M, NW) + ... energy[0, mi, fi] = np.sum(win ** 2) / np.sqrt(M) + ... # Corrected using subsample shifting + ... win = windows.dpss(M, NW, norm='subsample') + ... energy[1, mi, fi] = np.sum(win ** 2) / np.sqrt(M) + ... # Uncorrected (using l-infinity norm) + ... win /= win.max() + ... energy[2, mi, fi] = np.sum(win ** 2) / np.sqrt(M) + >>> fig, ax = plt.subplots(1) + >>> hs = ax.plot(Ms, energy[2], '-o', markersize=4, + ... markeredgecolor='none') + >>> leg = [hs[-1]] + >>> for hi, hh in enumerate(hs): + ... h1 = ax.plot(Ms, energy[0, :, hi], '-o', markersize=4, + ... color=hh.get_color(), markeredgecolor='none', + ... alpha=0.66) + ... h2 = ax.plot(Ms, energy[1, :, hi], '-o', markersize=4, + ... color=hh.get_color(), markeredgecolor='none', + ... alpha=0.33) + ... if hi == len(hs) - 1: + ... leg.insert(0, h1[0]) + ... leg.insert(0, h2[0]) + >>> ax.set(xlabel='M (samples)', ylabel=r'Power / $\\sqrt{M}$') + >>> ax.legend(leg, ['Uncorrected', r'Corrected: $\\frac{M^2}{M^2+NW}$', + ... 'Corrected (subsample)']) + >>> fig.tight_layout() + + """ + if _len_guards(M): + return np.ones(M) + if norm is None: + norm = 'approximate' if Kmax is None else 2 + known_norms = (2, 'approximate', 'subsample') + if norm not in known_norms: + raise ValueError(f'norm must be one of {known_norms}, got {norm}') + if Kmax is None: + singleton = True + Kmax = 1 + else: + singleton = False + Kmax = operator.index(Kmax) + if not 0 < Kmax <= M: + raise ValueError('Kmax must be greater than 0 and less than M') + if NW >= M/2.: + raise ValueError('NW must be less than M/2.') + if NW <= 0: + raise ValueError('NW must be positive') + M, needs_trunc = _extend(M, sym) + W = float(NW) / M + nidx = np.arange(M) + + # Here we want to set up an optimization problem to find a sequence + # whose energy is maximally concentrated within band [-W,W]. + # Thus, the measure lambda(T,W) is the ratio between the energy within + # that band, and the total energy. This leads to the eigen-system + # (A - (l1)I)v = 0, where the eigenvector corresponding to the largest + # eigenvalue is the sequence with maximally concentrated energy. The + # collection of eigenvectors of this system are called Slepian + # sequences, or discrete prolate spheroidal sequences (DPSS). Only the + # first K, K = 2NW/dt orders of DPSS will exhibit good spectral + # concentration + # [see https://en.wikipedia.org/wiki/Spectral_concentration_problem] + + # Here we set up an alternative symmetric tri-diagonal eigenvalue + # problem such that + # (B - (l2)I)v = 0, and v are our DPSS (but eigenvalues l2 != l1) + # the main diagonal = ([M-1-2*t]/2)**2 cos(2PIW), t=[0,1,2,...,M-1] + # and the first off-diagonal = t(M-t)/2, t=[1,2,...,M-1] + # [see Percival and Walden, 1993] + d = ((M - 1 - 2 * nidx) / 2.) ** 2 * np.cos(2 * np.pi * W) + e = nidx[1:] * (M - nidx[1:]) / 2. + + # only calculate the highest Kmax eigenvalues + w, windows = linalg.eigh_tridiagonal( + d, e, select='i', select_range=(M - Kmax, M - 1)) + w = w[::-1] + windows = windows[:, ::-1].T + + # By convention (Percival and Walden, 1993 pg 379) + # * symmetric tapers (k=0,2,4,...) should have a positive average. + fix_even = (windows[::2].sum(axis=1) < 0) + for i, f in enumerate(fix_even): + if f: + windows[2 * i] *= -1 + # * antisymmetric tapers should begin with a positive lobe + # (this depends on the definition of "lobe", here we'll take the first + # point above the numerical noise, which should be good enough for + # sufficiently smooth functions, and more robust than relying on an + # algorithm that uses max(abs(w)), which is susceptible to numerical + # noise problems) + thresh = max(1e-7, 1. / M) + for i, w in enumerate(windows[1::2]): + if w[w * w > thresh][0] < 0: + windows[2 * i + 1] *= -1 + + # Now find the eigenvalues of the original spectral concentration problem + # Use the autocorr sequence technique from Percival and Walden, 1993 pg 390 + if return_ratios: + dpss_rxx = _fftautocorr(windows) + r = 4 * W * np.sinc(2 * W * nidx) + r[0] = 2 * W + ratios = np.dot(dpss_rxx, r) + if singleton: + ratios = ratios[0] + # Deal with sym and Kmax=None + if norm != 2: + windows /= windows.max() + if M % 2 == 0: + if norm == 'approximate': + correction = M**2 / float(M**2 + NW) + else: + s = sp_fft.rfft(windows[0]) + shift = -(1 - 1./M) * np.arange(1, M//2 + 1) + s[1:] *= 2 * np.exp(-1j * np.pi * shift) + correction = M / s.real.sum() + windows *= correction + # else we're already l2 normed, so do nothing + if needs_trunc: + windows = windows[:, :-1] + if singleton: + windows = windows[0] + return (windows, ratios) if return_ratios else windows + + +def lanczos(M, *, sym=True): + r"""Return a Lanczos window also known as a sinc window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero, an empty array + is returned. An exception is thrown when it is negative. + sym : bool, optional + When True (default), generates a symmetric window, for use in filter + design. + When False, generates a periodic window, for use in spectral analysis. + + Returns + ------- + w : ndarray + The window, with the maximum value normalized to 1 (though the value 1 + does not appear if `M` is even and `sym` is True). + + Notes + ----- + The Lanczos window is defined as + + .. math:: w(n) = sinc \left( \frac{2n}{M - 1} - 1 \right) + + where + + .. math:: sinc(x) = \frac{\sin(\pi x)}{\pi x} + + The Lanczos window has reduced Gibbs oscillations and is widely used for + filtering climate timeseries with good properties in the physical and + spectral domains. + + .. versionadded:: 1.10 + + References + ---------- + .. [1] Lanczos, C., and Teichmann, T. (1957). Applied analysis. + Physics Today, 10, 44. + .. [2] Duchon C. E. (1979) Lanczos Filtering in One and Two Dimensions. + Journal of Applied Meteorology, Vol 18, pp 1016-1022. + .. [3] Thomson, R. E. and Emery, W. J. (2014) Data Analysis Methods in + Physical Oceanography (Third Edition), Elsevier, pp 593-637. + .. [4] Wikipedia, "Window function", + http://en.wikipedia.org/wiki/Window_function + + Examples + -------- + Plot the window + + >>> import numpy as np + >>> from scipy.signal.windows import lanczos + >>> from scipy.fft import fft, fftshift + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1) + >>> window = lanczos(51) + >>> ax.plot(window) + >>> ax.set_title("Lanczos window") + >>> ax.set_ylabel("Amplitude") + >>> ax.set_xlabel("Sample") + >>> fig.tight_layout() + >>> plt.show() + + and its frequency response: + + >>> fig, ax = plt.subplots(1) + >>> A = fft(window, 2048) / (len(window)/2.0) + >>> freq = np.linspace(-0.5, 0.5, len(A)) + >>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max()))) + >>> ax.plot(freq, response) + >>> ax.set_xlim(-0.5, 0.5) + >>> ax.set_ylim(-120, 0) + >>> ax.set_title("Frequency response of the lanczos window") + >>> ax.set_ylabel("Normalized magnitude [dB]") + >>> ax.set_xlabel("Normalized frequency [cycles per sample]") + >>> fig.tight_layout() + >>> plt.show() + """ + if _len_guards(M): + return np.ones(M) + M, needs_trunc = _extend(M, sym) + + # To make sure that the window is symmetric, we concatenate the right hand + # half of the window and the flipped one which is the left hand half of + # the window. + def _calc_right_side_lanczos(n, m): + return np.sinc(2. * np.arange(n, m) / (m - 1) - 1.0) + + if M % 2 == 0: + wh = _calc_right_side_lanczos(M/2, M) + w = np.r_[np.flip(wh), wh] + else: + wh = _calc_right_side_lanczos((M+1)/2, M) + w = np.r_[np.flip(wh), 1.0, wh] + + return _truncate(w, needs_trunc) + + +def _fftautocorr(x): + """Compute the autocorrelation of a real array and crop the result.""" + N = x.shape[-1] + use_N = sp_fft.next_fast_len(2*N-1) + x_fft = sp_fft.rfft(x, use_N, axis=-1) + cxy = sp_fft.irfft(x_fft * x_fft.conj(), n=use_N)[:, :N] + # Or equivalently (but in most cases slower): + # cxy = np.array([np.convolve(xx, yy[::-1], mode='full') + # for xx, yy in zip(x, x)])[:, N-1:2*N-1] + return cxy + + +_win_equiv_raw = { + ('barthann', 'brthan', 'bth'): (barthann, False), + ('bartlett', 'bart', 'brt'): (bartlett, False), + ('blackman', 'black', 'blk'): (blackman, False), + ('blackmanharris', 'blackharr', 'bkh'): (blackmanharris, False), + ('bohman', 'bman', 'bmn'): (bohman, False), + ('boxcar', 'box', 'ones', + 'rect', 'rectangular'): (boxcar, False), + ('chebwin', 'cheb'): (chebwin, True), + ('cosine', 'halfcosine'): (cosine, False), + ('dpss',): (dpss, True), + ('exponential', 'poisson'): (exponential, False), + ('flattop', 'flat', 'flt'): (flattop, False), + ('gaussian', 'gauss', 'gss'): (gaussian, True), + ('general cosine', 'general_cosine'): (general_cosine, True), + ('general gaussian', 'general_gaussian', + 'general gauss', 'general_gauss', 'ggs'): (general_gaussian, True), + ('general hamming', 'general_hamming'): (general_hamming, True), + ('hamming', 'hamm', 'ham'): (hamming, False), + ('hann', 'han'): (hann, False), + ('kaiser', 'ksr'): (kaiser, True), + ('kaiser bessel derived', 'kbd'): (kaiser_bessel_derived, True), + ('lanczos', 'sinc'): (lanczos, False), + ('nuttall', 'nutl', 'nut'): (nuttall, False), + ('parzen', 'parz', 'par'): (parzen, False), + ('taylor', 'taylorwin'): (taylor, False), + ('triangle', 'triang', 'tri'): (triang, False), + ('tukey', 'tuk'): (tukey, False), +} + +# Fill dict with all valid window name strings +_win_equiv = {} +for k, v in _win_equiv_raw.items(): + for key in k: + _win_equiv[key] = v[0] + +# Keep track of which windows need additional parameters +_needs_param = set() +for k, v in _win_equiv_raw.items(): + if v[1]: + _needs_param.update(k) + + +def get_window(window, Nx, fftbins=True): + """ + Return a window of a given length and type. + + Parameters + ---------- + window : string, float, or tuple + The type of window to create. See below for more details. + Nx : int + The number of samples in the window. + fftbins : bool, optional + If True (default), create a "periodic" window, ready to use with + `ifftshift` and be multiplied by the result of an FFT (see also + :func:`~scipy.fft.fftfreq`). + If False, create a "symmetric" window, for use in filter design. + + Returns + ------- + get_window : ndarray + Returns a window of length `Nx` and type `window` + + Notes + ----- + Window types: + + - `~scipy.signal.windows.boxcar` + - `~scipy.signal.windows.triang` + - `~scipy.signal.windows.blackman` + - `~scipy.signal.windows.hamming` + - `~scipy.signal.windows.hann` + - `~scipy.signal.windows.bartlett` + - `~scipy.signal.windows.flattop` + - `~scipy.signal.windows.parzen` + - `~scipy.signal.windows.bohman` + - `~scipy.signal.windows.blackmanharris` + - `~scipy.signal.windows.nuttall` + - `~scipy.signal.windows.barthann` + - `~scipy.signal.windows.cosine` + - `~scipy.signal.windows.exponential` + - `~scipy.signal.windows.tukey` + - `~scipy.signal.windows.taylor` + - `~scipy.signal.windows.lanczos` + - `~scipy.signal.windows.kaiser` (needs beta) + - `~scipy.signal.windows.kaiser_bessel_derived` (needs beta) + - `~scipy.signal.windows.gaussian` (needs standard deviation) + - `~scipy.signal.windows.general_cosine` (needs weighting coefficients) + - `~scipy.signal.windows.general_gaussian` (needs power, width) + - `~scipy.signal.windows.general_hamming` (needs window coefficient) + - `~scipy.signal.windows.dpss` (needs normalized half-bandwidth) + - `~scipy.signal.windows.chebwin` (needs attenuation) + + + If the window requires no parameters, then `window` can be a string. + + If the window requires parameters, then `window` must be a tuple + with the first argument the string name of the window, and the next + arguments the needed parameters. + + If `window` is a floating point number, it is interpreted as the beta + parameter of the `~scipy.signal.windows.kaiser` window. + + Each of the window types listed above is also the name of + a function that can be called directly to create a window of + that type. + + Examples + -------- + >>> from scipy import signal + >>> signal.get_window('triang', 7) + array([ 0.125, 0.375, 0.625, 0.875, 0.875, 0.625, 0.375]) + >>> signal.get_window(('kaiser', 4.0), 9) + array([ 0.08848053, 0.29425961, 0.56437221, 0.82160913, 0.97885093, + 0.97885093, 0.82160913, 0.56437221, 0.29425961]) + >>> signal.get_window(('exponential', None, 1.), 9) + array([ 0.011109 , 0.03019738, 0.082085 , 0.22313016, 0.60653066, + 0.60653066, 0.22313016, 0.082085 , 0.03019738]) + >>> signal.get_window(4.0, 9) + array([ 0.08848053, 0.29425961, 0.56437221, 0.82160913, 0.97885093, + 0.97885093, 0.82160913, 0.56437221, 0.29425961]) + + """ + sym = not fftbins + try: + beta = float(window) + except (TypeError, ValueError) as e: + args = () + if isinstance(window, tuple): + winstr = window[0] + if len(window) > 1: + args = window[1:] + elif isinstance(window, str): + if window in _needs_param: + raise ValueError("The '" + window + "' window needs one or " + "more parameters -- pass a tuple.") from e + else: + winstr = window + else: + raise ValueError( + f"{str(type(window))} as window type is not supported.") from e + + try: + winfunc = _win_equiv[winstr] + except KeyError as e: + raise ValueError("Unknown window type.") from e + + if winfunc is dpss: + params = (Nx,) + args + (None,) + else: + params = (Nx,) + args + else: + winfunc = kaiser + params = (Nx, beta) + + return winfunc(*params, sym=sym) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/windows/windows.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/windows/windows.py new file mode 100644 index 0000000000000000000000000000000000000000..e3512b67ae78368c4921f31e7abf3c6dffe67c94 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/signal/windows/windows.py @@ -0,0 +1,23 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.signal.windows` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + +__all__ = [ # noqa: F822 + 'boxcar', 'triang', 'parzen', 'bohman', 'blackman', 'nuttall', + 'blackmanharris', 'flattop', 'bartlett', 'barthann', + 'hamming', 'kaiser', 'gaussian', 'general_cosine', + 'general_gaussian', 'general_hamming', 'chebwin', 'cosine', + 'hann', 'exponential', 'tukey', 'taylor', 'dpss', 'get_window', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="signal.windows", module="windows", + private_modules=["_windows"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..021169fa0c9ea579a45ff413b44f5cb24bf5029a --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/__init__.py @@ -0,0 +1,331 @@ +""" +=================================== +Sparse arrays (:mod:`scipy.sparse`) +=================================== + +.. currentmodule:: scipy.sparse + +.. toctree:: + :hidden: + + sparse.csgraph + sparse.linalg + sparse.migration_to_sparray + +SciPy 2-D sparse array package for numeric data. + +.. note:: + + This package is switching to an array interface, compatible with + NumPy arrays, from the older matrix interface. We recommend that + you use the array objects (`bsr_array`, `coo_array`, etc.) for + all new work. + + When using the array interface, please note that: + + - ``x * y`` no longer performs matrix multiplication, but + element-wise multiplication (just like with NumPy arrays). To + make code work with both arrays and matrices, use ``x @ y`` for + matrix multiplication. + - Operations such as ``sum``, that used to produce dense matrices, now + produce arrays, whose multiplication behavior differs similarly. + - Sparse arrays use array style *slicing* operations, returning scalars, + 1D, or 2D sparse arrays. If you need 2D results, use an appropriate index. + E.g. ``A[:, i, None]`` or ``A[:, [i]]``. + + The construction utilities (`eye`, `kron`, `random`, `diags`, etc.) + have appropriate replacements (see :ref:`sparse-construction-functions`). + + For more information see + :ref:`Migration from spmatrix to sparray `. + + +Submodules +========== + +.. autosummary:: + + csgraph - Compressed sparse graph routines + linalg - Sparse linear algebra routines + + +Sparse array classes +==================== + +.. autosummary:: + :toctree: generated/ + + bsr_array - Block Sparse Row array + coo_array - A sparse array in COOrdinate format + csc_array - Compressed Sparse Column array + csr_array - Compressed Sparse Row array + dia_array - Sparse array with DIAgonal storage + dok_array - Dictionary Of Keys based sparse array + lil_array - Row-based list of lists sparse array + sparray - Sparse array base class + +.. _sparse-construction-functions: + +Building sparse arrays +---------------------- + +.. autosummary:: + :toctree: generated/ + + diags_array - Return a sparse array from diagonals + eye_array - Sparse MxN array whose k-th diagonal is all ones + random_array - Random values in a given shape array + block_array - Build a sparse array from sub-blocks + +.. _combining-arrays: + +Combining arrays +---------------- + +.. autosummary:: + :toctree: generated/ + + kron - Kronecker product of two sparse arrays + kronsum - Kronecker sum of sparse arrays + block_diag - Build a block diagonal sparse array + tril - Lower triangular portion of a sparse array + triu - Upper triangular portion of a sparse array + hstack - Stack sparse arrays horizontally (column wise) + vstack - Stack sparse arrays vertically (row wise) + +Sparse tools +------------ + +.. autosummary:: + :toctree: generated/ + + save_npz - Save a sparse array to a file using ``.npz`` format. + load_npz - Load a sparse array from a file using ``.npz`` format. + find - Return the indices and values of the nonzero elements + get_index_dtype - determine a good dtype for index arrays. + safely_cast_index_arrays - cast index array dtype or raise if shape too big + +Identifying sparse arrays +------------------------- + +.. autosummary:: + :toctree: generated/ + + issparse - Check if the argument is a sparse object (array or matrix). + + +Sparse matrix classes +===================== + +.. autosummary:: + :toctree: generated/ + + bsr_matrix - Block Sparse Row matrix + coo_matrix - A sparse matrix in COOrdinate format + csc_matrix - Compressed Sparse Column matrix + csr_matrix - Compressed Sparse Row matrix + dia_matrix - Sparse matrix with DIAgonal storage + dok_matrix - Dictionary Of Keys based sparse matrix + lil_matrix - Row-based list of lists sparse matrix + spmatrix - Sparse matrix base class + +Building sparse matrices +------------------------ + +.. autosummary:: + :toctree: generated/ + + eye - Sparse MxN matrix whose k-th diagonal is all ones + identity - Identity matrix in sparse matrix format + diags - Return a sparse matrix from diagonals + spdiags - Return a sparse matrix from diagonals + bmat - Build a sparse matrix from sparse sub-blocks + random - Random values in a given shape matrix + rand - Random values in a given shape matrix (old interface) + +**Combining matrices use the same functions as for** :ref:`combining-arrays`. + +Identifying sparse matrices +--------------------------- + +.. autosummary:: + :toctree: generated/ + + issparse + isspmatrix + isspmatrix_csc + isspmatrix_csr + isspmatrix_bsr + isspmatrix_lil + isspmatrix_dok + isspmatrix_coo + isspmatrix_dia + + +Warnings +======== + +.. autosummary:: + :toctree: generated/ + + SparseEfficiencyWarning + SparseWarning + + +Usage information +================= + +There are seven available sparse array types: + + 1. csc_array: Compressed Sparse Column format + 2. csr_array: Compressed Sparse Row format + 3. bsr_array: Block Sparse Row format + 4. lil_array: List of Lists format + 5. dok_array: Dictionary of Keys format + 6. coo_array: COOrdinate format (aka IJV, triplet format) + 7. dia_array: DIAgonal format + +To construct an array efficiently, use any of `coo_array`, +`dok_array` or `lil_array`. `dok_array` and `lil_array` +support basic slicing and fancy indexing with a similar syntax +to NumPy arrays. The COO format does not support indexing (yet) +but can also be used to efficiently construct arrays using coord +and value info. + +Despite their similarity to NumPy arrays, it is **strongly discouraged** +to use NumPy functions directly on these arrays because NumPy typically +treats them as generic Python objects rather than arrays, leading to +unexpected (and incorrect) results. If you do want to apply a NumPy +function to these arrays, first check if SciPy has its own implementation +for the given sparse array class, or **convert the sparse array to +a NumPy array** (e.g., using the `toarray` method of the class) +before applying the method. + +All conversions among the CSR, CSC, and COO formats are efficient, +linear-time operations. + +To perform manipulations such as multiplication or inversion, first +convert the array to either CSC or CSR format. The `lil_array` +format is row-based, so conversion to CSR is efficient, whereas +conversion to CSC is less so. + +Matrix vector product +--------------------- + +To do a vector product between a 2D sparse array and a vector use +the matmul operator (i.e., ``@``) which performs a dot product (like the +``dot`` method): + +>>> import numpy as np +>>> from scipy.sparse import csr_array +>>> A = csr_array([[1, 2, 0], [0, 0, 3], [4, 0, 5]]) +>>> v = np.array([1, 0, -1]) +>>> A @ v +array([ 1, -3, -1], dtype=int64) + +The CSR format is especially suitable for fast matrix vector products. + +Example 1 +--------- + +Construct a 1000x1000 `lil_array` and add some values to it: + +>>> from scipy.sparse import lil_array +>>> from scipy.sparse.linalg import spsolve +>>> from numpy.linalg import solve, norm +>>> from numpy.random import rand + +>>> A = lil_array((1000, 1000)) +>>> A[0, :100] = rand(100) +>>> A.setdiag(rand(1000)) + +Now convert it to CSR format and solve A x = b for x: + +>>> A = A.tocsr() +>>> b = rand(1000) +>>> x = spsolve(A, b) + +Convert it to a dense array and solve, and check that the result +is the same: + +>>> x_ = solve(A.toarray(), b) + +Now we can compute norm of the error with: + +>>> err = norm(x-x_) +>>> err < 1e-10 +True + +It should be small :) + + +Example 2 +--------- + +Construct an array in COO format: + +>>> from scipy import sparse +>>> from numpy import array +>>> I = array([0,3,1,0]) +>>> J = array([0,3,1,2]) +>>> V = array([4,5,7,9]) +>>> A = sparse.coo_array((V,(I,J)),shape=(4,4)) + +Notice that the indices do not need to be sorted. + +Duplicate (i,j) entries are summed when converting to CSR or CSC. + +>>> I = array([0,0,1,3,1,0,0]) +>>> J = array([0,2,1,3,1,0,0]) +>>> V = array([1,1,1,1,1,1,1]) +>>> B = sparse.coo_array((V,(I,J)),shape=(4,4)).tocsr() + +This is useful for constructing finite-element stiffness and mass matrices. + +Further details +--------------- + +CSR column indices are not necessarily sorted. Likewise for CSC row +indices. Use the ``.sorted_indices()`` and ``.sort_indices()`` methods when +sorted indices are required (e.g., when passing data to other libraries). + +""" + +# Original code by Travis Oliphant. +# Modified and extended by Ed Schofield, Robert Cimrman, +# Nathan Bell, and Jake Vanderplas. + +import warnings as _warnings + +from ._base import * +from ._csr import * +from ._csc import * +from ._lil import * +from ._dok import * +from ._coo import * +from ._dia import * +from ._bsr import * +from ._construct import * +from ._extract import * +from ._matrix import spmatrix +from ._matrix_io import * +from ._sputils import get_index_dtype, safely_cast_index_arrays + +# For backward compatibility with v0.19. +from . import csgraph + +# Deprecated namespaces, to be removed in v2.0.0 +from . import ( + base, bsr, compressed, construct, coo, csc, csr, data, dia, dok, extract, + lil, sparsetools, sputils +) + +__all__ = [s for s in dir() if not s.startswith('_')] + +# Filter PendingDeprecationWarning for np.matrix introduced with numpy 1.15 +msg = 'the matrix subclass is not the recommended way' +_warnings.filterwarnings('ignore', message=msg) + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_coo.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_coo.py new file mode 100644 index 0000000000000000000000000000000000000000..d23ad6327ccc6eb5278daf817ed1cd3c1cba1b63 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_coo.py @@ -0,0 +1,1647 @@ +""" A sparse matrix in COOrdinate or 'triplet' format""" + +__docformat__ = "restructuredtext en" + +__all__ = ['coo_array', 'coo_matrix', 'isspmatrix_coo'] + +import math +from warnings import warn + +import numpy as np + +from .._lib._util import copy_if_needed +from ._matrix import spmatrix +from ._sparsetools import (coo_tocsr, coo_todense, coo_todense_nd, + coo_matvec, coo_matvec_nd, coo_matmat_dense, + coo_matmat_dense_nd) +from ._base import issparse, SparseEfficiencyWarning, _spbase, sparray +from ._data import _data_matrix, _minmax_mixin +from ._sputils import (upcast_char, to_native, isshape, getdtype, + getdata, downcast_intp_index, get_index_dtype, + check_shape, check_reshape_kwargs, isscalarlike, isdense) + +import operator + + +class _coo_base(_data_matrix, _minmax_mixin): + _format = 'coo' + _allow_nd = range(1, 65) + + def __init__(self, arg1, shape=None, dtype=None, copy=False, *, maxprint=None): + _data_matrix.__init__(self, arg1, maxprint=maxprint) + if not copy: + copy = copy_if_needed + + if isinstance(arg1, tuple): + if isshape(arg1, allow_nd=self._allow_nd): + self._shape = check_shape(arg1, allow_nd=self._allow_nd) + idx_dtype = self._get_index_dtype(maxval=max(self._shape)) + data_dtype = getdtype(dtype, default=float) + self.coords = tuple(np.array([], dtype=idx_dtype) + for _ in range(len(self._shape))) + self.data = np.array([], dtype=data_dtype) + self.has_canonical_format = True + else: + try: + obj, coords = arg1 + except (TypeError, ValueError) as e: + raise TypeError('invalid input format') from e + + if shape is None: + if any(len(idx) == 0 for idx in coords): + raise ValueError('cannot infer dimensions from zero ' + 'sized index arrays') + shape = tuple(operator.index(np.max(idx)) + 1 + for idx in coords) + self._shape = check_shape(shape, allow_nd=self._allow_nd) + idx_dtype = self._get_index_dtype(coords, + maxval=max(self.shape), + check_contents=True) + self.coords = tuple(np.array(idx, copy=copy, dtype=idx_dtype) + for idx in coords) + self.data = getdata(obj, copy=copy, dtype=dtype) + self.has_canonical_format = False + else: + if issparse(arg1): + if arg1.format == self.format and copy: + self.coords = tuple(idx.copy() for idx in arg1.coords) + self.data = arg1.data.astype(getdtype(dtype, arg1)) # copy=True + self._shape = check_shape(arg1.shape, allow_nd=self._allow_nd) + self.has_canonical_format = arg1.has_canonical_format + else: + coo = arg1.tocoo() + self.coords = tuple(coo.coords) + self.data = coo.data.astype(getdtype(dtype, coo), copy=False) + self._shape = check_shape(coo.shape, allow_nd=self._allow_nd) + self.has_canonical_format = False + else: + # dense argument + M = np.asarray(arg1) + if not isinstance(self, sparray): + M = np.atleast_2d(M) + if M.ndim != 2: + raise TypeError(f'expected 2D array or matrix, not {M.ndim}D') + + self._shape = check_shape(M.shape, allow_nd=self._allow_nd) + if shape is not None: + if check_shape(shape, allow_nd=self._allow_nd) != self._shape: + message = f'inconsistent shapes: {shape} != {self._shape}' + raise ValueError(message) + + index_dtype = self._get_index_dtype(maxval=max(self._shape)) + coords = M.nonzero() + self.coords = tuple(idx.astype(index_dtype, copy=False) + for idx in coords) + self.data = getdata(M[coords], copy=copy, dtype=dtype) + self.has_canonical_format = True + + if len(self._shape) > 2: + self.coords = tuple(idx.astype(np.int64, copy=False) for idx in self.coords) + + self._check() + + @property + def row(self): + if self.ndim > 1: + return self.coords[-2] + result = np.zeros_like(self.col) + result.setflags(write=False) + return result + + + @row.setter + def row(self, new_row): + if self.ndim < 2: + raise ValueError('cannot set row attribute of a 1-dimensional sparse array') + new_row = np.asarray(new_row, dtype=self.coords[-2].dtype) + self.coords = self.coords[:-2] + (new_row,) + self.coords[-1:] + + @property + def col(self): + return self.coords[-1] + + @col.setter + def col(self, new_col): + new_col = np.asarray(new_col, dtype=self.coords[-1].dtype) + self.coords = self.coords[:-1] + (new_col,) + + def reshape(self, *args, **kwargs): + shape = check_shape(args, self.shape, allow_nd=self._allow_nd) + order, copy = check_reshape_kwargs(kwargs) + + # Return early if reshape is not required + if shape == self.shape: + if copy: + return self.copy() + else: + return self + + # When reducing the number of dimensions, we need to be careful about + # index overflow. This is why we can't simply call + # `np.ravel_multi_index()` followed by `np.unravel_index()` here. + flat_coords = _ravel_coords(self.coords, self.shape, order=order) + if len(shape) == 2: + if order == 'C': + new_coords = divmod(flat_coords, shape[1]) + else: + new_coords = divmod(flat_coords, shape[0])[::-1] + else: + new_coords = np.unravel_index(flat_coords, shape, order=order) + + idx_dtype = self._get_index_dtype(self.coords, maxval=max(shape)) + new_coords = tuple(np.asarray(co, dtype=idx_dtype) for co in new_coords) + + # Handle copy here rather than passing on to the constructor so that no + # copy will be made of `new_coords` regardless. + if copy: + new_data = self.data.copy() + else: + new_data = self.data + + return self.__class__((new_data, new_coords), shape=shape, copy=False) + + reshape.__doc__ = _spbase.reshape.__doc__ + + def _getnnz(self, axis=None): + if axis is None or (axis == 0 and self.ndim == 1): + nnz = len(self.data) + if any(len(idx) != nnz for idx in self.coords): + raise ValueError('all index and data arrays must have the ' + 'same length') + + if self.data.ndim != 1 or any(idx.ndim != 1 for idx in self.coords): + raise ValueError('coordinates and data arrays must be 1-D') + + return int(nnz) + + if axis < 0: + axis += self.ndim + if axis >= self.ndim: + raise ValueError('axis out of bounds') + + return np.bincount(downcast_intp_index(self.coords[1 - axis]), + minlength=self.shape[1 - axis]) + + _getnnz.__doc__ = _spbase._getnnz.__doc__ + + def count_nonzero(self, axis=None): + self.sum_duplicates() + if axis is None: + return np.count_nonzero(self.data) + + if axis < 0: + axis += self.ndim + if axis < 0 or axis >= self.ndim: + raise ValueError('axis out of bounds') + mask = self.data != 0 + coord = self.coords[1 - axis][mask] + return np.bincount(downcast_intp_index(coord), minlength=self.shape[1 - axis]) + + count_nonzero.__doc__ = _spbase.count_nonzero.__doc__ + + def _check(self): + """ Checks data structure for consistency """ + if self.ndim != len(self.coords): + raise ValueError('mismatching number of index arrays for shape; ' + f'got {len(self.coords)}, expected {self.ndim}') + + # index arrays should have integer data types + for i, idx in enumerate(self.coords): + if idx.dtype.kind != 'i': + warn(f'index array {i} has non-integer dtype ({idx.dtype.name})', + stacklevel=3) + + idx_dtype = self._get_index_dtype(self.coords, maxval=max(self.shape)) + self.coords = tuple(np.asarray(idx, dtype=idx_dtype) + for idx in self.coords) + self.data = to_native(self.data) + + if self.nnz > 0: + for i, idx in enumerate(self.coords): + if idx.max() >= self.shape[i]: + raise ValueError(f'axis {i} index {idx.max()} exceeds ' + f'matrix dimension {self.shape[i]}') + if idx.min() < 0: + raise ValueError(f'negative axis {i} index: {idx.min()}') + + def transpose(self, axes=None, copy=False): + if axes is None: + axes = range(self.ndim)[::-1] + elif isinstance(self, sparray): + if not hasattr(axes, "__len__") or len(axes) != self.ndim: + raise ValueError("axes don't match matrix dimensions") + if len(set(axes)) != self.ndim: + raise ValueError("repeated axis in transpose") + elif axes != (1, 0): + raise ValueError("Sparse matrices do not support an 'axes' " + "parameter because swapping dimensions is the " + "only logical permutation.") + + permuted_shape = tuple(self._shape[i] for i in axes) + permuted_coords = tuple(self.coords[i] for i in axes) + return self.__class__((self.data, permuted_coords), + shape=permuted_shape, copy=copy) + + transpose.__doc__ = _spbase.transpose.__doc__ + + def resize(self, *shape) -> None: + shape = check_shape(shape, allow_nd=self._allow_nd) + if self.ndim > 2: + raise ValueError("only 1-D or 2-D input accepted") + if len(shape) > 2: + raise ValueError("shape argument must be 1-D or 2-D") + # Check for added dimensions. + if len(shape) > self.ndim: + flat_coords = _ravel_coords(self.coords, self.shape) + max_size = math.prod(shape) + self.coords = np.unravel_index(flat_coords[:max_size], shape) + self.data = self.data[:max_size] + self._shape = shape + return + + # Check for removed dimensions. + if len(shape) < self.ndim: + tmp_shape = ( + self._shape[:len(shape) - 1] # Original shape without last axis + + (-1,) # Last axis is used to flatten the array + + (1,) * (self.ndim - len(shape)) # Pad with ones + ) + tmp = self.reshape(tmp_shape) + self.coords = tmp.coords[:len(shape)] + self._shape = tmp.shape[:len(shape)] + + # Handle truncation of existing dimensions. + is_truncating = any(old > new for old, new in zip(self.shape, shape)) + if is_truncating: + mask = np.logical_and.reduce([ + idx < size for idx, size in zip(self.coords, shape) + ]) + if not mask.all(): + self.coords = tuple(idx[mask] for idx in self.coords) + self.data = self.data[mask] + + self._shape = shape + + resize.__doc__ = _spbase.resize.__doc__ + + def toarray(self, order=None, out=None): + B = self._process_toarray_args(order, out) + fortran = int(B.flags.f_contiguous) + if not fortran and not B.flags.c_contiguous: + raise ValueError("Output array must be C or F contiguous") + # This handles both 0D and 1D cases correctly regardless of the + # original shape. + if self.ndim == 1: + coo_todense_nd(np.array([1]), self.nnz, self.ndim, + self.coords[0], self.data, B.ravel('A'), fortran) + elif self.ndim == 2: + M, N = self.shape + coo_todense(M, N, self.nnz, self.row, self.col, self.data, + B.ravel('A'), fortran) + else: + if fortran: + strides = np.append(1, np.cumprod(self.shape[:-1])) + else: + strides = np.append(np.cumprod(self.shape[1:][::-1])[::-1], 1) + coords = np.concatenate(self.coords) + coo_todense_nd(strides, self.nnz, self.ndim, + coords, self.data, B.ravel('A'), fortran) + # Note: reshape() doesn't copy here, but does return a new array (view). + return B.reshape(self.shape) + + toarray.__doc__ = _spbase.toarray.__doc__ + + def tocsc(self, copy=False): + """Convert this array/matrix to Compressed Sparse Column format + + Duplicate entries will be summed together. + + Examples + -------- + >>> from numpy import array + >>> from scipy.sparse import coo_array + >>> row = array([0, 0, 1, 3, 1, 0, 0]) + >>> col = array([0, 2, 1, 3, 1, 0, 0]) + >>> data = array([1, 1, 1, 1, 1, 1, 1]) + >>> A = coo_array((data, (row, col)), shape=(4, 4)).tocsc() + >>> A.toarray() + array([[3, 0, 1, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 1]]) + + """ + if self.ndim != 2: + raise ValueError(f'Cannot convert. CSC format must be 2D. Got {self.ndim}D') + if self.nnz == 0: + return self._csc_container(self.shape, dtype=self.dtype) + else: + from ._csc import csc_array + indptr, indices, data, shape = self._coo_to_compressed(csc_array._swap) + + x = self._csc_container((data, indices, indptr), shape=shape) + if not self.has_canonical_format: + x.sum_duplicates() + return x + + def tocsr(self, copy=False): + """Convert this array/matrix to Compressed Sparse Row format + + Duplicate entries will be summed together. + + Examples + -------- + >>> from numpy import array + >>> from scipy.sparse import coo_array + >>> row = array([0, 0, 1, 3, 1, 0, 0]) + >>> col = array([0, 2, 1, 3, 1, 0, 0]) + >>> data = array([1, 1, 1, 1, 1, 1, 1]) + >>> A = coo_array((data, (row, col)), shape=(4, 4)).tocsr() + >>> A.toarray() + array([[3, 0, 1, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 1]]) + + """ + if self.ndim > 2: + raise ValueError(f'Cannot convert. CSR must be 1D or 2D. Got {self.ndim}D') + if self.nnz == 0: + return self._csr_container(self.shape, dtype=self.dtype) + else: + from ._csr import csr_array + arrays = self._coo_to_compressed(csr_array._swap, copy=copy) + indptr, indices, data, shape = arrays + + x = self._csr_container((data, indices, indptr), shape=self.shape) + if not self.has_canonical_format: + x.sum_duplicates() + return x + + def _coo_to_compressed(self, swap, copy=False): + """convert (shape, coords, data) to (indptr, indices, data, shape)""" + M, N = swap(self._shape_as_2d) + # convert idx_dtype intc to int32 for pythran. + # tested in scipy/optimize/tests/test__numdiff.py::test_group_columns + idx_dtype = self._get_index_dtype(self.coords, maxval=max(self.nnz, N)) + + if self.ndim == 1: + indices = self.coords[0].copy() if copy else self.coords[0] + nnz = len(indices) + indptr = np.array([0, nnz], dtype=idx_dtype) + data = self.data.copy() if copy else self.data + return indptr, indices, data, self.shape + + # ndim == 2 + major, minor = swap(self.coords) + nnz = len(major) + major = major.astype(idx_dtype, copy=False) + minor = minor.astype(idx_dtype, copy=False) + + indptr = np.empty(M + 1, dtype=idx_dtype) + indices = np.empty_like(minor, dtype=idx_dtype) + data = np.empty_like(self.data, dtype=self.dtype) + + coo_tocsr(M, N, nnz, major, minor, self.data, indptr, indices, data) + return indptr, indices, data, self.shape + + def tocoo(self, copy=False): + if copy: + return self.copy() + else: + return self + + tocoo.__doc__ = _spbase.tocoo.__doc__ + + def todia(self, copy=False): + if self.ndim != 2: + raise ValueError(f'Cannot convert. DIA format must be 2D. Got {self.ndim}D') + self.sum_duplicates() + ks = self.col - self.row # the diagonal for each nonzero + diags, diag_idx = np.unique(ks, return_inverse=True) + + if len(diags) > 100: + # probably undesired, should todia() have a maxdiags parameter? + warn(f"Constructing a DIA matrix with {len(diags)} diagonals " + "is inefficient", + SparseEfficiencyWarning, stacklevel=2) + + #initialize and fill in data array + if self.data.size == 0: + data = np.zeros((0, 0), dtype=self.dtype) + else: + data = np.zeros((len(diags), self.col.max()+1), dtype=self.dtype) + data[diag_idx, self.col] = self.data + + return self._dia_container((data, diags), shape=self.shape) + + todia.__doc__ = _spbase.todia.__doc__ + + def todok(self, copy=False): + if self.ndim > 2: + raise ValueError(f'Cannot convert. DOK must be 1D or 2D. Got {self.ndim}D') + self.sum_duplicates() + dok = self._dok_container(self.shape, dtype=self.dtype) + # ensure that 1d coordinates are not tuples + if self.ndim == 1: + coords = self.coords[0] + else: + coords = zip(*self.coords) + + dok._dict = dict(zip(coords, self.data)) + return dok + + todok.__doc__ = _spbase.todok.__doc__ + + def diagonal(self, k=0): + if self.ndim != 2: + raise ValueError("diagonal requires two dimensions") + rows, cols = self.shape + if k <= -rows or k >= cols: + return np.empty(0, dtype=self.data.dtype) + diag = np.zeros(min(rows + min(k, 0), cols - max(k, 0)), + dtype=self.dtype) + diag_mask = (self.row + k) == self.col + + if self.has_canonical_format: + row = self.row[diag_mask] + data = self.data[diag_mask] + else: + inds = tuple(idx[diag_mask] for idx in self.coords) + (row, _), data = self._sum_duplicates(inds, self.data[diag_mask]) + diag[row + min(k, 0)] = data + + return diag + + diagonal.__doc__ = _data_matrix.diagonal.__doc__ + + def _setdiag(self, values, k): + if self.ndim != 2: + raise ValueError("setting a diagonal requires two dimensions") + M, N = self.shape + if values.ndim and not len(values): + return + idx_dtype = self.row.dtype + + # Determine which triples to keep and where to put the new ones. + full_keep = self.col - self.row != k + if k < 0: + max_index = min(M+k, N) + if values.ndim: + max_index = min(max_index, len(values)) + keep = np.logical_or(full_keep, self.col >= max_index) + new_row = np.arange(-k, -k + max_index, dtype=idx_dtype) + new_col = np.arange(max_index, dtype=idx_dtype) + else: + max_index = min(M, N-k) + if values.ndim: + max_index = min(max_index, len(values)) + keep = np.logical_or(full_keep, self.row >= max_index) + new_row = np.arange(max_index, dtype=idx_dtype) + new_col = np.arange(k, k + max_index, dtype=idx_dtype) + + # Define the array of data consisting of the entries to be added. + if values.ndim: + new_data = values[:max_index] + else: + new_data = np.empty(max_index, dtype=self.dtype) + new_data[:] = values + + # Update the internal structure. + self.coords = (np.concatenate((self.row[keep], new_row)), + np.concatenate((self.col[keep], new_col))) + self.data = np.concatenate((self.data[keep], new_data)) + self.has_canonical_format = False + + # needed by _data_matrix + def _with_data(self, data, copy=True): + """Returns a matrix with the same sparsity structure as self, + but with different data. By default the index arrays are copied. + """ + if copy: + coords = tuple(idx.copy() for idx in self.coords) + else: + coords = self.coords + return self.__class__((data, coords), shape=self.shape, dtype=data.dtype) + + def sum_duplicates(self) -> None: + """Eliminate duplicate entries by adding them together + + This is an *in place* operation + """ + if self.has_canonical_format: + return + summed = self._sum_duplicates(self.coords, self.data) + self.coords, self.data = summed + self.has_canonical_format = True + + def _sum_duplicates(self, coords, data): + # Assumes coords not in canonical format. + if len(data) == 0: + return coords, data + # Sort coords w.r.t. rows, then cols. This corresponds to C-order, + # which we rely on for argmin/argmax to return the first index in the + # same way that numpy does (in the case of ties). + order = np.lexsort(coords[::-1]) + coords = tuple(idx[order] for idx in coords) + data = data[order] + unique_mask = np.logical_or.reduce([ + idx[1:] != idx[:-1] for idx in coords + ]) + unique_mask = np.append(True, unique_mask) + coords = tuple(idx[unique_mask] for idx in coords) + unique_inds, = np.nonzero(unique_mask) + data = np.add.reduceat(data, downcast_intp_index(unique_inds), dtype=self.dtype) + return coords, data + + def eliminate_zeros(self): + """Remove zero entries from the array/matrix + + This is an *in place* operation + """ + mask = self.data != 0 + self.data = self.data[mask] + self.coords = tuple(idx[mask] for idx in self.coords) + + ####################### + # Arithmetic handlers # + ####################### + + def _add_dense(self, other): + if other.shape != self.shape: + raise ValueError(f'Incompatible shapes ({self.shape} and {other.shape})') + dtype = upcast_char(self.dtype.char, other.dtype.char) + result = np.array(other, dtype=dtype, copy=True) + fortran = int(result.flags.f_contiguous) + if self.ndim == 1: + coo_todense_nd(np.array([1]), self.nnz, self.ndim, + self.coords[0], self.data, result.ravel('A'), fortran) + elif self.ndim == 2: + M, N = self._shape_as_2d + coo_todense(M, N, self.nnz, self.row, self.col, self.data, + result.ravel('A'), fortran) + else: + if fortran: + strides = np.append(1, np.cumprod(self.shape[:-1])) + else: + strides = np.append(np.cumprod(self.shape[1:][::-1])[::-1], 1) + coords = np.concatenate(self.coords) + coo_todense_nd(strides, self.nnz, self.ndim, + coords, self.data, result.ravel('A'), fortran) + return self._container(result, copy=False) + + + def _add_sparse(self, other): + if self.ndim < 3: + return self.tocsr()._add_sparse(other) + + if other.shape != self.shape: + raise ValueError(f'Incompatible shapes ({self.shape} and {other.shape})') + other = self.__class__(other) + new_data = np.concatenate((self.data, other.data)) + new_coords = tuple(np.concatenate((self.coords, other.coords), axis=1)) + A = self.__class__((new_data, new_coords), shape=self.shape) + return A + + + def _sub_sparse(self, other): + if self.ndim < 3: + return self.tocsr()._sub_sparse(other) + + if other.shape != self.shape: + raise ValueError(f'Incompatible shapes ({self.shape} and {other.shape})') + other = self.__class__(other) + new_data = np.concatenate((self.data, -other.data)) + new_coords = tuple(np.concatenate((self.coords, other.coords), axis=1)) + A = coo_array((new_data, new_coords), shape=self.shape) + return A + + + def _matmul_vector(self, other): + if self.ndim > 2: + result = np.zeros(math.prod(self.shape[:-1]), + dtype=upcast_char(self.dtype.char, other.dtype.char)) + shape = np.array(self.shape) + strides = np.append(np.cumprod(shape[:-1][::-1])[::-1][1:], 1) + coords = np.concatenate(self.coords) + coo_matvec_nd(self.nnz, len(self.shape), strides, coords, self.data, + other, result) + + result = result.reshape(self.shape[:-1]) + return result + + # self.ndim <= 2 + result_shape = self.shape[0] if self.ndim > 1 else 1 + result = np.zeros(result_shape, + dtype=upcast_char(self.dtype.char, other.dtype.char)) + if self.ndim == 2: + col = self.col + row = self.row + elif self.ndim == 1: + col = self.coords[0] + row = np.zeros_like(col) + else: + raise NotImplementedError( + f"coo_matvec not implemented for ndim={self.ndim}") + + coo_matvec(self.nnz, row, col, self.data, other, result) + # Array semantics return a scalar here, not a single-element array. + if isinstance(self, sparray) and result_shape == 1: + return result[0] + return result + + + def _rmatmul_dispatch(self, other): + if isscalarlike(other): + return self._mul_scalar(other) + else: + # Don't use asarray unless we have to + try: + o_ndim = other.ndim + except AttributeError: + other = np.asarray(other) + o_ndim = other.ndim + perm = tuple(range(o_ndim)[:-2]) + tuple(range(o_ndim)[-2:][::-1]) + tr = other.transpose(perm) + + s_ndim = self.ndim + perm = tuple(range(s_ndim)[:-2]) + tuple(range(s_ndim)[-2:][::-1]) + ret = self.transpose(perm)._matmul_dispatch(tr) + if ret is NotImplemented: + return NotImplemented + + if s_ndim == 1 or o_ndim == 1: + perm = range(ret.ndim) + else: + perm = tuple(range(ret.ndim)[:-2]) + tuple(range(ret.ndim)[-2:][::-1]) + return ret.transpose(perm) + + + def _matmul_dispatch(self, other): + if isscalarlike(other): + return self.multiply(other) + + if not (issparse(other) or isdense(other)): + # If it's a list or whatever, treat it like an array + other_a = np.asanyarray(other) + + if other_a.ndim == 0 and other_a.dtype == np.object_: + # Not interpretable as an array; return NotImplemented so that + # other's __rmatmul__ can kick in if that's implemented. + return NotImplemented + + try: + other.shape + except AttributeError: + other = other_a + + if self.ndim < 3 and other.ndim < 3: + return _spbase._matmul_dispatch(self, other) + + N = self.shape[-1] + err_prefix = "matmul: dimension mismatch with signature" + if other.__class__ is np.ndarray: + if other.shape == (N,): + return self._matmul_vector(other) + if other.shape == (N, 1): + result = self._matmul_vector(other.ravel()) + return result.reshape(*self.shape[:-1], 1) + if other.ndim == 1: + msg = f"{err_prefix} (n,k={N}),(k={other.shape[0]},)->(n,)" + raise ValueError(msg) + if other.shape[-2] == N: + # check for batch dimensions compatibility + batch_shape_A = self.shape[:-2] + batch_shape_B = other.shape[:-2] + if batch_shape_A != batch_shape_B: + try: + # This will raise an error if the shapes are not broadcastable + np.broadcast_shapes(batch_shape_A, batch_shape_B) + except ValueError: + raise ValueError("Batch dimensions are not broadcastable") + + return self._matmul_multivector(other) + else: + raise ValueError( + f"{err_prefix} (n,..,k={N}),(k={other.shape[-2]},..,m)->(n,..,m)" + ) + + + if isscalarlike(other): + # scalar value + return self._mul_scalar(other) + + if issparse(other): + self_is_1d = self.ndim == 1 + other_is_1d = other.ndim == 1 + + # reshape to 2-D if self or other is 1-D + if self_is_1d: + self = self.reshape(self._shape_as_2d) # prepend 1 to shape + + if other_is_1d: + other = other.reshape((other.shape[0], 1)) # append 1 to shape + + # Check if the inner dimensions match for matrix multiplication + if N != other.shape[-2]: + raise ValueError( + f"{err_prefix} (n,..,k={N}),(k={other.shape[-2]},..,m)->(n,..,m)" + ) + + # If A or B has more than 2 dimensions, check for + # batch dimensions compatibility + if self.ndim > 2 or other.ndim > 2: + batch_shape_A = self.shape[:-2] + batch_shape_B = other.shape[:-2] + if batch_shape_A != batch_shape_B: + try: + # This will raise an error if the shapes are not broadcastable + np.broadcast_shapes(batch_shape_A, batch_shape_B) + except ValueError: + raise ValueError("Batch dimensions are not broadcastable") + + result = self._matmul_sparse(other) + + # reshape back if a or b were originally 1-D + if self_is_1d: + # if self was originally 1-D, reshape result accordingly + result = result.reshape(tuple(result.shape[:-2]) + + tuple(result.shape[-1:])) + if other_is_1d: + result = result.reshape(result.shape[:-1]) + return result + + + def _matmul_multivector(self, other): + result_dtype = upcast_char(self.dtype.char, other.dtype.char) + if self.ndim >= 3 or other.ndim >= 3: + # if self has shape (N,), reshape to (1,N) + if self.ndim == 1: + result = self.reshape(1, self.shape[0])._matmul_multivector(other) + return result.reshape(tuple(other.shape[:-2]) + tuple(other.shape[-1:])) + + broadcast_shape = np.broadcast_shapes(self.shape[:-2], other.shape[:-2]) + self_shape = broadcast_shape + self.shape[-2:] + other_shape = broadcast_shape + other.shape[-2:] + + self = self._broadcast_to(self_shape) + other = np.broadcast_to(other, other_shape) + result_shape = broadcast_shape + self.shape[-2:-1] + other.shape[-1:] + result = np.zeros(result_shape, dtype=result_dtype) + coo_matmat_dense_nd(self.nnz, len(self.shape), other.shape[-1], + np.array(other_shape), np.array(result_shape), + np.concatenate(self.coords), + self.data, other.ravel('C'), result) + return result + + if self.ndim == 2: + result_shape = (self.shape[0], other.shape[1]) + col = self.col + row = self.row + elif self.ndim == 1: + result_shape = (other.shape[1],) + col = self.coords[0] + row = np.zeros_like(col) + result = np.zeros(result_shape, dtype=result_dtype) + coo_matmat_dense(self.nnz, other.shape[-1], row, col, + self.data, other.ravel('C'), result) + return result.view(type=type(other)) + + + def dot(self, other): + """Return the dot product of two arrays. + + Strictly speaking a dot product involves two vectors. + But in the sense that an array with ndim >= 1 is a collection + of vectors, the function computes the collection of dot products + between each vector in the first array with each vector in the + second array. The axis upon which the sum of products is performed + is the last axis of the first array and the second to last axis of + the second array. If the second array is 1-D, the last axis is used. + + Thus, if both arrays are 1-D, the inner product is returned. + If both are 2-D, we have matrix multiplication. If `other` is 1-D, + the sum product is taken along the last axis of each array. If + `other` is N-D for N>=2, the sum product is over the last axis of + the first array and the second-to-last axis of the second array. + + Parameters + ---------- + other : array_like (dense or sparse) + Second array + + Returns + ------- + output : array (sparse or dense) + The dot product of this array with `other`. + It will be dense/sparse if `other` is dense/sparse. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import coo_array + >>> A = coo_array([[1, 2, 0], [0, 0, 3], [4, 0, 5]]) + >>> v = np.array([1, 0, -1]) + >>> A.dot(v) + array([ 1, -3, -1], dtype=int64) + + For 2-D arrays it is the matrix product: + + >>> A = coo_array([[1, 0], [0, 1]]) + >>> B = coo_array([[4, 1], [2, 2]]) + >>> A.dot(B).toarray() + array([[4, 1], + [2, 2]]) + + For 3-D arrays the shape extends unused axes by other unused axes. + + >>> A = coo_array(np.arange(3*4*5*6)).reshape((3,4,5,6)) + >>> B = coo_array(np.arange(3*4*5*6)).reshape((5,4,6,3)) + >>> A.dot(B).shape + (3, 4, 5, 5, 4, 3) + """ + if not (issparse(other) or isdense(other) or isscalarlike(other)): + # If it's a list or whatever, treat it like an array + o_array = np.asanyarray(other) + + if o_array.ndim == 0 and o_array.dtype == np.object_: + raise TypeError(f"dot argument not supported type: '{type(other)}'") + try: + other.shape + except AttributeError: + other = o_array + + if self.ndim < 3 and (np.isscalar(other) or other.ndim<3): + return _spbase.dot(self, other) + # Handle scalar multiplication + if np.isscalar(other): + return self * other + if isdense(other): + return self._dense_dot(other) + elif other.format != "coo": + raise TypeError("input must be a COO matrix/array") + elif self.ndim == 1 and other.ndim == 1: + # Handle inner product of vectors (1-D arrays) + if self.shape[0] != other.shape[0]: + raise ValueError(f"shapes {self.shape} and {other.shape}" + " are not aligned for inner product") + return self @ other + elif self.ndim == 2 and other.ndim == 2: + # Handle matrix multiplication (2-D arrays) + if self.shape[1] != other.shape[0]: + raise ValueError(f"shapes {self.shape} and {other.shape}" + " are not aligned for matmul") + return self @ other + else: + return self._sparse_dot(other) + + + def _sparse_dot(self, other): + self_is_1d = self.ndim == 1 + other_is_1d = other.ndim == 1 + + # reshape to 2-D if self or other is 1-D + if self_is_1d: + self = self.reshape(self._shape_as_2d) # prepend 1 to shape + if other_is_1d: + other = other.reshape((other.shape[0], 1)) # append 1 to shape + + if self.shape[-1] != other.shape[-2]: + raise ValueError(f"shapes {self.shape} and {other.shape}" + " are not aligned for n-D dot") + + # Prepare the tensors for dot operation + # Ravel non-reduced axes coordinates + self_raveled_coords = _ravel_non_reduced_axes(self.coords, + self.shape, [self.ndim-1]) + other_raveled_coords = _ravel_non_reduced_axes(other.coords, + other.shape, [other.ndim-2]) + + # Get the shape of the non-reduced axes + self_nonreduced_shape = self.shape[:-1] + other_nonreduced_shape = other.shape[:-2] + other.shape[-1:] + + # Create 2D coords arrays + ravel_coords_shape_self = (math.prod(self_nonreduced_shape), self.shape[-1]) + ravel_coords_shape_other = (other.shape[-2], math.prod(other_nonreduced_shape)) + + self_2d_coords = (self_raveled_coords, self.coords[-1]) + other_2d_coords = (other.coords[-2], other_raveled_coords) + + self_2d = coo_array((self.data, self_2d_coords), ravel_coords_shape_self) + other_2d = coo_array((other.data, other_2d_coords), ravel_coords_shape_other) + + prod = (self_2d @ other_2d).tocoo() # routes via 2-D CSR + + # Combine the shapes of the non-reduced axes + combined_shape = self_nonreduced_shape + other_nonreduced_shape + + # Unravel the 2D coordinates to get multi-dimensional coordinates + shapes = (self_nonreduced_shape, other_nonreduced_shape) + prod_coords = [] + for c, s in zip(prod.coords, shapes): + prod_coords.extend(np.unravel_index(c, s)) + + prod_arr = coo_array((prod.data, prod_coords), combined_shape) + + # reshape back if a or b were originally 1-D + # TODO: Move this logic before computation of prod_coords for efficiency + if self_is_1d: + prod_arr = prod_arr.reshape(combined_shape[1:]) + if other_is_1d: + prod_arr = prod_arr.reshape(combined_shape[:-1]) + + return prod_arr + + def _dense_dot(self, other): + self_is_1d = self.ndim == 1 + other_is_1d = other.ndim == 1 + + # reshape to 2-D if self or other is 1-D + if self_is_1d: + self = self.reshape(self._shape_as_2d) # prepend 1 to shape + if other_is_1d: + other = other.reshape((other.shape[0], 1)) # append 1 to shape + + if self.shape[-1] != other.shape[-2]: + raise ValueError(f"shapes {self.shape} and {other.shape}" + " are not aligned for n-D dot") + + new_shape_self = ( + self.shape[:-1] + (1,) * (len(other.shape) - 1) + self.shape[-1:] + ) + new_shape_other = (1,) * (len(self.shape) - 1) + other.shape + + result_shape = self.shape[:-1] + other.shape[:-2] + other.shape[-1:] + result = self.reshape(new_shape_self) @ other.reshape(new_shape_other) + prod_arr = result.reshape(result_shape) + + # reshape back if a or b were originally 1-D + if self_is_1d: + prod_arr = prod_arr.reshape(result_shape[1:]) + if other_is_1d: + prod_arr = prod_arr.reshape(result_shape[:-1]) + + return prod_arr + + def tensordot(self, other, axes=2): + """Return the tensordot product with another array along the given axes. + + The tensordot differs from dot and matmul in that any axis can be + chosen for each of the first and second array and the sum of the + products is computed just like for matrix multiplication, only not + just for the rows of the first times the columns of the second. It + takes the dot product of the collection of vectors along the specified + axes. Here we can even take the sum of the products along two or even + more axes if desired. So, tensordot is a dot product computation + applied to arrays of any dimension >= 1. It is like matmul but over + arbitrary axes for each matrix. + + Given two tensors, `a` and `b`, and the desired axes specified as a + 2-tuple/list/array containing two sequences of axis numbers, + ``(a_axes, b_axes)``, sum the products of `a`'s and `b`'s elements + (components) over the axes specified by ``a_axes`` and ``b_axes``. + The `axes` input can be a single non-negative integer, ``N``; + if it is, then the last ``N`` dimensions of `a` and the first + ``N`` dimensions of `b` are summed over. + + Parameters + ---------- + a, b : array_like + Tensors to "dot". + + axes : int or (2,) array_like + * integer_like + If an int N, sum over the last N axes of `a` and the first N axes + of `b` in order. The sizes of the corresponding axes must match. + * (2,) array_like + A 2-tuple of sequences of axes to be summed over, the first applying + to `a`, the second to `b`. The sequences must be the same length. + The shape of the corresponding axes must match between `a` and `b`. + + Returns + ------- + output : coo_array + The tensor dot product of this array with `other`. + It will be dense/sparse if `other` is dense/sparse. + + See Also + -------- + dot + + Examples + -------- + >>> import numpy as np + >>> import scipy.sparse + >>> A = scipy.sparse.coo_array([[[2, 3], [0, 0]], [[0, 1], [0, 5]]]) + >>> A.shape + (2, 2, 2) + + Integer axes N are shorthand for (range(-N, 0), range(0, N)): + + >>> A.tensordot(A, axes=1).toarray() + array([[[[ 4, 9], + [ 0, 15]], + + [[ 0, 0], + [ 0, 0]]], + + + [[[ 0, 1], + [ 0, 5]], + + [[ 0, 5], + [ 0, 25]]]]) + >>> A.tensordot(A, axes=2).toarray() + array([[ 4, 6], + [ 0, 25]]) + >>> A.tensordot(A, axes=3) + array(39) + + Using tuple for axes: + + >>> a = scipy.sparse.coo_array(np.arange(60).reshape(3,4,5)) + >>> b = np.arange(24).reshape(4,3,2) + >>> c = a.tensordot(b, axes=([1,0],[0,1])) + >>> c.shape + (5, 2) + >>> c + array([[4400, 4730], + [4532, 4874], + [4664, 5018], + [4796, 5162], + [4928, 5306]]) + + """ + if not isdense(other) and not issparse(other): + # If it's a list or whatever, treat it like an array + other_array = np.asanyarray(other) + + if other_array.ndim == 0 and other_array.dtype == np.object_: + raise TypeError(f"tensordot arg not supported type: '{type(other)}'") + try: + other.shape + except AttributeError: + other = other_array + + axes_self, axes_other = _process_axes(self.ndim, other.ndim, axes) + + # Check for shape compatibility along specified axes + if any(self.shape[ax] != other.shape[bx] + for ax, bx in zip(axes_self, axes_other)): + raise ValueError("sizes of the corresponding axes must match") + + if isdense(other): + return self._dense_tensordot(other, axes_self, axes_other) + else: + return self._sparse_tensordot(other, axes_self, axes_other) + + + def _sparse_tensordot(self, other, axes_self, axes_other): + ndim_self = len(self.shape) + ndim_other = len(other.shape) + + # Prepare the tensors for tensordot operation + # Ravel non-reduced axes coordinates + self_non_red_coords = _ravel_non_reduced_axes(self.coords, self.shape, + axes_self) + self_reduced_coords = np.ravel_multi_index( + [self.coords[ax] for ax in axes_self], [self.shape[ax] for ax in axes_self]) + other_non_red_coords = _ravel_non_reduced_axes(other.coords, other.shape, + axes_other) + other_reduced_coords = np.ravel_multi_index( + [other.coords[a] for a in axes_other], [other.shape[a] for a in axes_other] + ) + # Get the shape of the non-reduced axes + self_nonreduced_shape = tuple(self.shape[ax] for ax in range(ndim_self) + if ax not in axes_self) + other_nonreduced_shape = tuple(other.shape[ax] for ax in range(ndim_other) + if ax not in axes_other) + + # Create 2D coords arrays + ravel_coords_shape_self = (math.prod(self_nonreduced_shape), + math.prod([self.shape[ax] for ax in axes_self])) + ravel_coords_shape_other = (math.prod([other.shape[ax] for ax in axes_other]), + math.prod(other_nonreduced_shape)) + + self_2d_coords = (self_non_red_coords, self_reduced_coords) + other_2d_coords = (other_reduced_coords, other_non_red_coords) + + self_2d = coo_array((self.data, self_2d_coords), ravel_coords_shape_self) + other_2d = coo_array((other.data, other_2d_coords), ravel_coords_shape_other) + + # Perform matrix multiplication (routed via 2-D CSR) + prod = (self_2d @ other_2d).tocoo() + + # Combine the shapes of the non-contracted axes + combined_shape = self_nonreduced_shape + other_nonreduced_shape + + # Unravel the 2D coordinates to get multi-dimensional coordinates + coords = [] + for c, s in zip(prod.coords, (self_nonreduced_shape, other_nonreduced_shape)): + if s: + coords.extend(np.unravel_index(c, s)) + + if coords == []: # if result is scalar + return sum(prod.data) + + # Construct the resulting COO array with combined coordinates and shape + return coo_array((prod.data, coords), shape=combined_shape) + + + def _dense_tensordot(self, other, axes_self, axes_other): + ndim_self = len(self.shape) + ndim_other = len(other.shape) + + non_reduced_axes_self = [ax for ax in range(ndim_self) if ax not in axes_self] + reduced_shape_self = [self.shape[s] for s in axes_self] + non_reduced_shape_self = [self.shape[s] for s in non_reduced_axes_self] + + non_reduced_axes_other = [ax for ax in range(ndim_other) + if ax not in axes_other] + reduced_shape_other = [other.shape[s] for s in axes_other] + non_reduced_shape_other = [other.shape[s] for s in non_reduced_axes_other] + + permute_self = non_reduced_axes_self + axes_self + permute_other = ( + non_reduced_axes_other[:-1] + axes_other + non_reduced_axes_other[-1:] + ) + self = self.transpose(permute_self) + other = np.transpose(other, permute_other) + + reshape_self = (*non_reduced_shape_self, math.prod(reduced_shape_self)) + reshape_other = (*non_reduced_shape_other[:-1], math.prod(reduced_shape_other), + *non_reduced_shape_other[-1:]) + + return self.reshape(reshape_self).dot(other.reshape(reshape_other)) + + + def _matmul_sparse(self, other): + """ + Perform sparse-sparse matrix multiplication for two n-D COO arrays. + The method converts input n-D arrays to 2-D block array format, + uses csr_matmat to multiply them, and then converts the + result back to n-D COO array. + + Parameters: + self (COO): The first n-D sparse array in COO format. + other (COO): The second n-D sparse array in COO format. + + Returns: + prod (COO): The resulting n-D sparse array after multiplication. + """ + if self.ndim < 3 and other.ndim < 3: + return _spbase._matmul_sparse(self, other) + + # Get the shapes of self and other + self_shape = self.shape + other_shape = other.shape + + # Determine the new shape to broadcast self and other + broadcast_shape = np.broadcast_shapes(self_shape[:-2], other_shape[:-2]) + self_new_shape = tuple(broadcast_shape) + self_shape[-2:] + other_new_shape = tuple(broadcast_shape) + other_shape[-2:] + + self_broadcasted = self._broadcast_to(self_new_shape) + other_broadcasted = other._broadcast_to(other_new_shape) + + # Convert n-D COO arrays to 2-D block diagonal arrays + self_block_diag = _block_diag(self_broadcasted) + other_block_diag = _block_diag(other_broadcasted) + + # Use csr_matmat to perform sparse matrix multiplication + prod_block_diag = (self_block_diag @ other_block_diag).tocoo() + + # Convert the 2-D block diagonal array back to n-D + return _extract_block_diag( + prod_block_diag, + shape=(*broadcast_shape, self.shape[-2], other.shape[-1]), + ) + + + def _broadcast_to(self, new_shape, copy=False): + if self.shape == new_shape: + return self.copy() if copy else self + + old_shape = self.shape + + # Check if the new shape is compatible for broadcasting + if len(new_shape) < len(old_shape): + raise ValueError("New shape must have at least as many dimensions" + " as the current shape") + + # Add leading ones to shape to ensure same length as `new_shape` + shape = (1,) * (len(new_shape) - len(old_shape)) + tuple(old_shape) + + # Ensure the old shape can be broadcast to the new shape + if any((o != 1 and o != n) for o, n in zip(shape, new_shape)): + raise ValueError(f"current shape {old_shape} cannot be " + "broadcast to new shape {new_shape}") + + # Reshape the COO array to match the new dimensions + self = self.reshape(shape) + + idx_dtype = get_index_dtype(self.coords, maxval=max(new_shape)) + coords = self.coords + new_data = self.data + new_coords = coords[-1:] # Copy last coordinate to start + cum_repeat = 1 # Cumulative repeat factor for broadcasting + + if shape[-1] != new_shape[-1]: # broadcasting the n-th (col) dimension + repeat_count = new_shape[-1] + cum_repeat *= repeat_count + new_data = np.tile(new_data, repeat_count) + new_dim = np.repeat(np.arange(0, repeat_count, dtype=idx_dtype), self.nnz) + new_coords = (new_dim,) + + for i in range(-2, -(len(shape)+1), -1): + if shape[i] != new_shape[i]: + repeat_count = new_shape[i] # number of times to repeat data, coords + cum_repeat *= repeat_count # update cumulative repeat factor + nnz = len(new_data) # Number of non-zero elements so far + + # Tile data and coordinates to match the new repeat count + new_data = np.tile(new_data, repeat_count) + new_coords = tuple(np.tile(new_coords[i+1:], repeat_count)) + + # Create new dimensions and stack them + new_dim = np.repeat(np.arange(0, repeat_count, dtype=idx_dtype), nnz) + new_coords = (new_dim,) + new_coords + else: + # If no broadcasting needed, tile the coordinates + new_dim = np.tile(coords[i], cum_repeat) + new_coords = (new_dim,) + new_coords + + return coo_array((new_data, new_coords), new_shape) + + +def _block_diag(self): + """ + Converts an N-D COO array into a 2-D COO array in block diagonal form. + + Parameters: + self (coo_array): An N-Dimensional COO sparse array. + + Returns: + coo_array: A 2-Dimensional COO sparse array in block diagonal form. + """ + if self.ndim<2: + raise ValueError("array must have atleast dim=2") + num_blocks = math.prod(self.shape[:-2]) + n_col = self.shape[-1] + n_row = self.shape[-2] + res_arr = self.reshape((num_blocks, n_row, n_col)) + new_coords = ( + res_arr.coords[1] + res_arr.coords[0] * res_arr.shape[1], + res_arr.coords[2] + res_arr.coords[0] * res_arr.shape[2], + ) + + new_shape = (num_blocks * n_row, num_blocks * n_col) + return coo_array((self.data, tuple(new_coords)), shape=new_shape) + + +def _extract_block_diag(self, shape): + n_row, n_col = shape[-2], shape[-1] + + # Extract data and coordinates from the block diagonal COO array + data = self.data + row, col = self.row, self.col + + # Initialize new coordinates array + new_coords = np.empty((len(shape), self.nnz), dtype=int) + + # Calculate within-block indices + new_coords[-2] = row % n_row + new_coords[-1] = col % n_col + + # Calculate coordinates for higher dimensions + temp_block_idx = row // n_row + for i in range(len(shape) - 3, -1, -1): + size = shape[i] + new_coords[i] = temp_block_idx % size + temp_block_idx = temp_block_idx // size + + # Create the new COO array with the original n-D shape + return coo_array((data, tuple(new_coords)), shape=shape) + + +def _process_axes(ndim_a, ndim_b, axes): + if isinstance(axes, int): + if axes < 1 or axes > min(ndim_a, ndim_b): + raise ValueError("axes integer is out of bounds for input arrays") + axes_a = list(range(ndim_a - axes, ndim_a)) + axes_b = list(range(axes)) + elif isinstance(axes, (tuple, list)): + if len(axes) != 2: + raise ValueError("axes must be a tuple/list of length 2") + axes_a, axes_b = axes + if len(axes_a) != len(axes_b): + raise ValueError("axes lists/tuples must be of the same length") + if any(ax >= ndim_a or ax < -ndim_a for ax in axes_a) or \ + any(bx >= ndim_b or bx < -ndim_b for bx in axes_b): + raise ValueError("axes indices are out of bounds for input arrays") + else: + raise TypeError("axes must be an integer or a tuple/list of integers") + + axes_a = [axis + ndim_a if axis < 0 else axis for axis in axes_a] + axes_b = [axis + ndim_b if axis < 0 else axis for axis in axes_b] + return axes_a, axes_b + + +def _ravel_non_reduced_axes(coords, shape, axes): + ndim = len(shape) + non_reduced_axes = [ax for ax in range(ndim) if ax not in axes] + + if not non_reduced_axes: + # Return an array with one row + return np.zeros_like(coords[0]) + + # Extract the shape of the non-reduced axes + non_reduced_shape = [shape[ax] for ax in non_reduced_axes] + + # Extract the coordinates of the non-reduced axes + non_reduced_coords = tuple(coords[idx] for idx in non_reduced_axes) + + # Ravel the coordinates into 1D + return np.ravel_multi_index(non_reduced_coords, non_reduced_shape) + + +def _ravel_coords(coords, shape, order='C'): + """Like np.ravel_multi_index, but avoids some overflow issues.""" + if len(coords) == 1: + return coords[0] + # Handle overflow as in https://github.com/scipy/scipy/pull/9132 + if len(coords) == 2: + nrows, ncols = shape + row, col = coords + if order == 'C': + maxval = (ncols * max(0, nrows - 1) + max(0, ncols - 1)) + idx_dtype = get_index_dtype(maxval=maxval) + return np.multiply(ncols, row, dtype=idx_dtype) + col + elif order == 'F': + maxval = (nrows * max(0, ncols - 1) + max(0, nrows - 1)) + idx_dtype = get_index_dtype(maxval=maxval) + return np.multiply(nrows, col, dtype=idx_dtype) + row + else: + raise ValueError("'order' must be 'C' or 'F'") + return np.ravel_multi_index(coords, shape, order=order) + + +def isspmatrix_coo(x): + """Is `x` of coo_matrix type? + + Parameters + ---------- + x + object to check for being a coo matrix + + Returns + ------- + bool + True if `x` is a coo matrix, False otherwise + + Examples + -------- + >>> from scipy.sparse import coo_array, coo_matrix, csr_matrix, isspmatrix_coo + >>> isspmatrix_coo(coo_matrix([[5]])) + True + >>> isspmatrix_coo(coo_array([[5]])) + False + >>> isspmatrix_coo(csr_matrix([[5]])) + False + """ + return isinstance(x, coo_matrix) + + +# This namespace class separates array from matrix with isinstance +class coo_array(_coo_base, sparray): + """ + A sparse array in COOrdinate format. + + Also known as the 'ijv' or 'triplet' format. + + This can be instantiated in several ways: + coo_array(D) + where D is an ndarray + + coo_array(S) + with another sparse array or matrix S (equivalent to S.tocoo()) + + coo_array(shape, [dtype]) + to construct an empty sparse array with shape `shape` + dtype is optional, defaulting to dtype='d'. + + coo_array((data, coords), [shape]) + to construct from existing data and index arrays: + 1. data[:] the entries of the sparse array, in any order + 2. coords[i][:] the axis-i coordinates of the data entries + + Where ``A[coords] = data``, and coords is a tuple of index arrays. + When shape is not specified, it is inferred from the index arrays. + + Attributes + ---------- + dtype : dtype + Data type of the sparse array + shape : tuple of integers + Shape of the sparse array + ndim : int + Number of dimensions of the sparse array + nnz + size + data + COO format data array of the sparse array + coords + COO format tuple of index arrays + has_canonical_format : bool + Whether the matrix has sorted coordinates and no duplicates + format + T + + Notes + ----- + + Sparse arrays can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the COO format + - facilitates fast conversion among sparse formats + - permits duplicate entries (see example) + - very fast conversion to and from CSR/CSC formats + + Disadvantages of the COO format + - does not directly support: + + arithmetic operations + + slicing + + Intended Usage + - COO is a fast format for constructing sparse arrays + - Once a COO array has been constructed, convert to CSR or + CSC format for fast arithmetic and matrix vector operations + - By default when converting to CSR or CSC format, duplicate (i,j) + entries will be summed together. This facilitates efficient + construction of finite element matrices and the like. (see example) + + Canonical format + - Entries and coordinates sorted by row, then column. + - There are no duplicate entries (i.e. duplicate (i,j) locations) + - Data arrays MAY have explicit zeros. + + Examples + -------- + + >>> # Constructing an empty sparse array + >>> import numpy as np + >>> from scipy.sparse import coo_array + >>> coo_array((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> # Constructing a sparse array using ijv format + >>> row = np.array([0, 3, 1, 0]) + >>> col = np.array([0, 3, 1, 2]) + >>> data = np.array([4, 5, 7, 9]) + >>> coo_array((data, (row, col)), shape=(4, 4)).toarray() + array([[4, 0, 9, 0], + [0, 7, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 5]]) + + >>> # Constructing a sparse array with duplicate coordinates + >>> row = np.array([0, 0, 1, 3, 1, 0, 0]) + >>> col = np.array([0, 2, 1, 3, 1, 0, 0]) + >>> data = np.array([1, 1, 1, 1, 1, 1, 1]) + >>> coo = coo_array((data, (row, col)), shape=(4, 4)) + >>> # Duplicate coordinates are maintained until implicitly or explicitly summed + >>> np.max(coo.data) + 1 + >>> coo.toarray() + array([[3, 0, 1, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 1]]) + + """ + + +class coo_matrix(spmatrix, _coo_base): + """ + A sparse matrix in COOrdinate format. + + Also known as the 'ijv' or 'triplet' format. + + This can be instantiated in several ways: + coo_matrix(D) + where D is a 2-D ndarray + + coo_matrix(S) + with another sparse array or matrix S (equivalent to S.tocoo()) + + coo_matrix((M, N), [dtype]) + to construct an empty matrix with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + coo_matrix((data, (i, j)), [shape=(M, N)]) + to construct from three arrays: + 1. data[:] the entries of the matrix, in any order + 2. i[:] the row indices of the matrix entries + 3. j[:] the column indices of the matrix entries + + Where ``A[i[k], j[k]] = data[k]``. When shape is not + specified, it is inferred from the index arrays + + Attributes + ---------- + dtype : dtype + Data type of the matrix + shape : 2-tuple + Shape of the matrix + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + COO format data array of the matrix + row + COO format row index array of the matrix + col + COO format column index array of the matrix + has_canonical_format : bool + Whether the matrix has sorted indices and no duplicates + format + T + + Notes + ----- + + Sparse matrices can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the COO format + - facilitates fast conversion among sparse formats + - permits duplicate entries (see example) + - very fast conversion to and from CSR/CSC formats + + Disadvantages of the COO format + - does not directly support: + + arithmetic operations + + slicing + + Intended Usage + - COO is a fast format for constructing sparse matrices + - Once a COO matrix has been constructed, convert to CSR or + CSC format for fast arithmetic and matrix vector operations + - By default when converting to CSR or CSC format, duplicate (i,j) + entries will be summed together. This facilitates efficient + construction of finite element matrices and the like. (see example) + + Canonical format + - Entries and coordinates sorted by row, then column. + - There are no duplicate entries (i.e. duplicate (i,j) locations) + - Data arrays MAY have explicit zeros. + + Examples + -------- + + >>> # Constructing an empty matrix + >>> import numpy as np + >>> from scipy.sparse import coo_matrix + >>> coo_matrix((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> # Constructing a matrix using ijv format + >>> row = np.array([0, 3, 1, 0]) + >>> col = np.array([0, 3, 1, 2]) + >>> data = np.array([4, 5, 7, 9]) + >>> coo_matrix((data, (row, col)), shape=(4, 4)).toarray() + array([[4, 0, 9, 0], + [0, 7, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 5]]) + + >>> # Constructing a matrix with duplicate coordinates + >>> row = np.array([0, 0, 1, 3, 1, 0, 0]) + >>> col = np.array([0, 2, 1, 3, 1, 0, 0]) + >>> data = np.array([1, 1, 1, 1, 1, 1, 1]) + >>> coo = coo_matrix((data, (row, col)), shape=(4, 4)) + >>> # Duplicate coordinates are maintained until implicitly or explicitly summed + >>> np.max(coo.data) + 1 + >>> coo.toarray() + array([[3, 0, 1, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 1]]) + + """ + + def __setstate__(self, state): + if 'coords' not in state: + # For retro-compatibility with the previous attributes + # storing nnz coordinates for 2D COO matrix. + state['coords'] = (state.pop('row'), state.pop('col')) + self.__dict__.update(state) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csc.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csc.py new file mode 100644 index 0000000000000000000000000000000000000000..e651136d8269d955f067ad434610bbea6f7339b1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csc.py @@ -0,0 +1,367 @@ +"""Compressed Sparse Column matrix format""" +__docformat__ = "restructuredtext en" + +__all__ = ['csc_array', 'csc_matrix', 'isspmatrix_csc'] + + +import numpy as np + +from ._matrix import spmatrix +from ._base import _spbase, sparray +from ._sparsetools import csc_tocsr, expandptr +from ._sputils import upcast + +from ._compressed import _cs_matrix + + +class _csc_base(_cs_matrix): + _format = 'csc' + + def transpose(self, axes=None, copy=False): + if axes is not None and axes != (1, 0): + raise ValueError("Sparse arrays/matrices do not support " + "an 'axes' parameter because swapping " + "dimensions is the only logical permutation.") + + M, N = self.shape + + return self._csr_container((self.data, self.indices, + self.indptr), (N, M), copy=copy) + + transpose.__doc__ = _spbase.transpose.__doc__ + + def __iter__(self): + yield from self.tocsr() + + def tocsc(self, copy=False): + if copy: + return self.copy() + else: + return self + + tocsc.__doc__ = _spbase.tocsc.__doc__ + + def tocsr(self, copy=False): + M,N = self.shape + idx_dtype = self._get_index_dtype((self.indptr, self.indices), + maxval=max(self.nnz, N)) + indptr = np.empty(M + 1, dtype=idx_dtype) + indices = np.empty(self.nnz, dtype=idx_dtype) + data = np.empty(self.nnz, dtype=upcast(self.dtype)) + + csc_tocsr(M, N, + self.indptr.astype(idx_dtype), + self.indices.astype(idx_dtype), + self.data, + indptr, + indices, + data) + + A = self._csr_container( + (data, indices, indptr), + shape=self.shape, copy=False + ) + A.has_sorted_indices = True + return A + + tocsr.__doc__ = _spbase.tocsr.__doc__ + + def nonzero(self): + # CSC can't use _cs_matrix's .nonzero method because it + # returns the indices sorted for self transposed. + + # Get row and col indices, from _cs_matrix.tocoo + major_dim, minor_dim = self._swap(self.shape) + minor_indices = self.indices + major_indices = np.empty(len(minor_indices), dtype=self.indices.dtype) + expandptr(major_dim, self.indptr, major_indices) + row, col = self._swap((major_indices, minor_indices)) + + # Remove explicit zeros + nz_mask = self.data != 0 + row = row[nz_mask] + col = col[nz_mask] + + # Sort them to be in C-style order + ind = np.argsort(row, kind='mergesort') + row = row[ind] + col = col[ind] + + return row, col + + nonzero.__doc__ = _cs_matrix.nonzero.__doc__ + + def _getrow(self, i): + """Returns a copy of row i of the matrix, as a (1 x n) + CSR matrix (row vector). + """ + M, N = self.shape + i = int(i) + if i < 0: + i += M + if i < 0 or i >= M: + raise IndexError('index (%d) out of range' % i) + return self._get_submatrix(minor=i).tocsr() + + def _getcol(self, i): + """Returns a copy of column i of the matrix, as a (m x 1) + CSC matrix (column vector). + """ + M, N = self.shape + i = int(i) + if i < 0: + i += N + if i < 0 or i >= N: + raise IndexError('index (%d) out of range' % i) + return self._get_submatrix(major=i, copy=True) + + def _get_intXarray(self, row, col): + return self._major_index_fancy(col)._get_submatrix(minor=row) + + def _get_intXslice(self, row, col): + if col.step in (1, None): + return self._get_submatrix(major=col, minor=row, copy=True) + return self._major_slice(col)._get_submatrix(minor=row) + + def _get_sliceXint(self, row, col): + if row.step in (1, None): + return self._get_submatrix(major=col, minor=row, copy=True) + return self._get_submatrix(major=col)._minor_slice(row) + + def _get_sliceXarray(self, row, col): + return self._major_index_fancy(col)._minor_slice(row) + + def _get_arrayXint(self, row, col): + res = self._get_submatrix(major=col)._minor_index_fancy(row) + if row.ndim > 1: + return res.reshape(row.shape) + return res + + def _get_arrayXslice(self, row, col): + return self._major_slice(col)._minor_index_fancy(row) + + # these functions are used by the parent class (_cs_matrix) + # to remove redundancy between csc_array and csr_matrix + @staticmethod + def _swap(x): + """swap the members of x if this is a column-oriented matrix + """ + return x[1], x[0] + + +def isspmatrix_csc(x): + """Is `x` of csc_matrix type? + + Parameters + ---------- + x + object to check for being a csc matrix + + Returns + ------- + bool + True if `x` is a csc matrix, False otherwise + + Examples + -------- + >>> from scipy.sparse import csc_array, csc_matrix, coo_matrix, isspmatrix_csc + >>> isspmatrix_csc(csc_matrix([[5]])) + True + >>> isspmatrix_csc(csc_array([[5]])) + False + >>> isspmatrix_csc(coo_matrix([[5]])) + False + """ + return isinstance(x, csc_matrix) + + +# This namespace class separates array from matrix with isinstance +class csc_array(_csc_base, sparray): + """ + Compressed Sparse Column array. + + This can be instantiated in several ways: + csc_array(D) + where D is a 2-D ndarray + + csc_array(S) + with another sparse array or matrix S (equivalent to S.tocsc()) + + csc_array((M, N), [dtype]) + to construct an empty array with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + csc_array((data, (row_ind, col_ind)), [shape=(M, N)]) + where ``data``, ``row_ind`` and ``col_ind`` satisfy the + relationship ``a[row_ind[k], col_ind[k]] = data[k]``. + + csc_array((data, indices, indptr), [shape=(M, N)]) + is the standard CSC representation where the row indices for + column i are stored in ``indices[indptr[i]:indptr[i+1]]`` + and their corresponding values are stored in + ``data[indptr[i]:indptr[i+1]]``. If the shape parameter is + not supplied, the array dimensions are inferred from + the index arrays. + + Attributes + ---------- + dtype : dtype + Data type of the array + shape : 2-tuple + Shape of the array + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + CSC format data array of the array + indices + CSC format index array of the array + indptr + CSC format index pointer array of the array + has_sorted_indices + has_canonical_format + T + + Notes + ----- + + Sparse arrays can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the CSC format + - efficient arithmetic operations CSC + CSC, CSC * CSC, etc. + - efficient column slicing + - fast matrix vector products (CSR, BSR may be faster) + + Disadvantages of the CSC format + - slow row slicing operations (consider CSR) + - changes to the sparsity structure are expensive (consider LIL or DOK) + + Canonical format + - Within each column, indices are sorted by row. + - There are no duplicate entries. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import csc_array + >>> csc_array((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> row = np.array([0, 2, 2, 0, 1, 2]) + >>> col = np.array([0, 0, 1, 2, 2, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csc_array((data, (row, col)), shape=(3, 3)).toarray() + array([[1, 0, 4], + [0, 0, 5], + [2, 3, 6]]) + + >>> indptr = np.array([0, 2, 3, 6]) + >>> indices = np.array([0, 2, 2, 0, 1, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csc_array((data, indices, indptr), shape=(3, 3)).toarray() + array([[1, 0, 4], + [0, 0, 5], + [2, 3, 6]]) + + """ + + +class csc_matrix(spmatrix, _csc_base): + """ + Compressed Sparse Column matrix. + + This can be instantiated in several ways: + csc_matrix(D) + where D is a 2-D ndarray + + csc_matrix(S) + with another sparse array or matrix S (equivalent to S.tocsc()) + + csc_matrix((M, N), [dtype]) + to construct an empty matrix with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + csc_matrix((data, (row_ind, col_ind)), [shape=(M, N)]) + where ``data``, ``row_ind`` and ``col_ind`` satisfy the + relationship ``a[row_ind[k], col_ind[k]] = data[k]``. + + csc_matrix((data, indices, indptr), [shape=(M, N)]) + is the standard CSC representation where the row indices for + column i are stored in ``indices[indptr[i]:indptr[i+1]]`` + and their corresponding values are stored in + ``data[indptr[i]:indptr[i+1]]``. If the shape parameter is + not supplied, the matrix dimensions are inferred from + the index arrays. + + Attributes + ---------- + dtype : dtype + Data type of the matrix + shape : 2-tuple + Shape of the matrix + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + CSC format data array of the matrix + indices + CSC format index array of the matrix + indptr + CSC format index pointer array of the matrix + has_sorted_indices + has_canonical_format + T + + Notes + ----- + + Sparse matrices can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the CSC format + - efficient arithmetic operations CSC + CSC, CSC * CSC, etc. + - efficient column slicing + - fast matrix vector products (CSR, BSR may be faster) + + Disadvantages of the CSC format + - slow row slicing operations (consider CSR) + - changes to the sparsity structure are expensive (consider LIL or DOK) + + Canonical format + - Within each column, indices are sorted by row. + - There are no duplicate entries. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import csc_matrix + >>> csc_matrix((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> row = np.array([0, 2, 2, 0, 1, 2]) + >>> col = np.array([0, 0, 1, 2, 2, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csc_matrix((data, (row, col)), shape=(3, 3)).toarray() + array([[1, 0, 4], + [0, 0, 5], + [2, 3, 6]]) + + >>> indptr = np.array([0, 2, 3, 6]) + >>> indices = np.array([0, 2, 2, 0, 1, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csc_matrix((data, indices, indptr), shape=(3, 3)).toarray() + array([[1, 0, 4], + [0, 0, 5], + [2, 3, 6]]) + + """ + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csparsetools.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csparsetools.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..5b05a56952b55f0cbc5cf507f2057be84093b42d Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csparsetools.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csr.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csr.py new file mode 100644 index 0000000000000000000000000000000000000000..92ced1ce0c5026534c208c0dd1b141235b85c3d0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_csr.py @@ -0,0 +1,558 @@ +"""Compressed Sparse Row matrix format""" + +__docformat__ = "restructuredtext en" + +__all__ = ['csr_array', 'csr_matrix', 'isspmatrix_csr'] + +import numpy as np + +from ._matrix import spmatrix +from ._base import _spbase, sparray +from ._sparsetools import (csr_tocsc, csr_tobsr, csr_count_blocks, + get_csr_submatrix, csr_sample_values) +from ._sputils import upcast + +from ._compressed import _cs_matrix + + +class _csr_base(_cs_matrix): + _format = 'csr' + _allow_nd = (1, 2) + + def transpose(self, axes=None, copy=False): + if axes is not None and axes != (1, 0): + raise ValueError("Sparse arrays/matrices do not support " + "an 'axes' parameter because swapping " + "dimensions is the only logical permutation.") + + if self.ndim == 1: + return self.copy() if copy else self + M, N = self.shape + return self._csc_container((self.data, self.indices, + self.indptr), shape=(N, M), copy=copy) + + transpose.__doc__ = _spbase.transpose.__doc__ + + def tolil(self, copy=False): + if self.ndim != 2: + raise ValueError("Cannot convert a 1d sparse array to lil format") + lil = self._lil_container(self.shape, dtype=self.dtype) + + self.sum_duplicates() + ptr,ind,dat = self.indptr,self.indices,self.data + rows, data = lil.rows, lil.data + + for n in range(self.shape[0]): + start = ptr[n] + end = ptr[n+1] + rows[n] = ind[start:end].tolist() + data[n] = dat[start:end].tolist() + + return lil + + tolil.__doc__ = _spbase.tolil.__doc__ + + def tocsr(self, copy=False): + if copy: + return self.copy() + else: + return self + + tocsr.__doc__ = _spbase.tocsr.__doc__ + + def tocsc(self, copy=False): + if self.ndim != 2: + raise ValueError("Cannot convert a 1d sparse array to csc format") + M, N = self.shape + idx_dtype = self._get_index_dtype((self.indptr, self.indices), + maxval=max(self.nnz, M)) + indptr = np.empty(N + 1, dtype=idx_dtype) + indices = np.empty(self.nnz, dtype=idx_dtype) + data = np.empty(self.nnz, dtype=upcast(self.dtype)) + + csr_tocsc(M, N, + self.indptr.astype(idx_dtype), + self.indices.astype(idx_dtype), + self.data, + indptr, + indices, + data) + + A = self._csc_container((data, indices, indptr), shape=self.shape) + A.has_sorted_indices = True + return A + + tocsc.__doc__ = _spbase.tocsc.__doc__ + + def tobsr(self, blocksize=None, copy=True): + if self.ndim != 2: + raise ValueError("Cannot convert a 1d sparse array to bsr format") + if blocksize is None: + from ._spfuncs import estimate_blocksize + return self.tobsr(blocksize=estimate_blocksize(self)) + + elif blocksize == (1,1): + arg1 = (self.data.reshape(-1,1,1),self.indices,self.indptr) + return self._bsr_container(arg1, shape=self.shape, copy=copy) + + else: + R,C = blocksize + M,N = self.shape + + if R < 1 or C < 1 or M % R != 0 or N % C != 0: + raise ValueError(f'invalid blocksize {blocksize}') + + blks = csr_count_blocks(M,N,R,C,self.indptr,self.indices) + + idx_dtype = self._get_index_dtype((self.indptr, self.indices), + maxval=max(N//C, blks)) + indptr = np.empty(M//R+1, dtype=idx_dtype) + indices = np.empty(blks, dtype=idx_dtype) + data = np.zeros((blks,R,C), dtype=self.dtype) + + csr_tobsr(M, N, R, C, + self.indptr.astype(idx_dtype), + self.indices.astype(idx_dtype), + self.data, + indptr, indices, data.ravel()) + + return self._bsr_container( + (data, indices, indptr), shape=self.shape + ) + + tobsr.__doc__ = _spbase.tobsr.__doc__ + + # these functions are used by the parent class (_cs_matrix) + # to remove redundancy between csc_matrix and csr_array + @staticmethod + def _swap(x): + """swap the members of x if this is a column-oriented matrix + """ + return x + + def __iter__(self): + if self.ndim == 1: + zero = self.dtype.type(0) + u = 0 + for v, d in zip(self.indices, self.data): + for _ in range(v - u): + yield zero + yield d + u = v + 1 + for _ in range(self.shape[0] - u): + yield zero + return + + indptr = np.zeros(2, dtype=self.indptr.dtype) + # return 1d (sparray) or 2drow (spmatrix) + shape = self.shape[1:] if isinstance(self, sparray) else (1, self.shape[1]) + i0 = 0 + for i1 in self.indptr[1:]: + indptr[1] = i1 - i0 + indices = self.indices[i0:i1] + data = self.data[i0:i1] + yield self.__class__((data, indices, indptr), shape=shape, copy=True) + i0 = i1 + + def _getrow(self, i): + """Returns a copy of row i of the matrix, as a (1 x n) + CSR matrix (row vector). + """ + if self.ndim == 1: + if i not in (0, -1): + raise IndexError(f'index ({i}) out of range') + return self.reshape((1, self.shape[0]), copy=True) + + M, N = self.shape + i = int(i) + if i < 0: + i += M + if i < 0 or i >= M: + raise IndexError('index (%d) out of range' % i) + indptr, indices, data = get_csr_submatrix( + M, N, self.indptr, self.indices, self.data, i, i + 1, 0, N) + return self.__class__((data, indices, indptr), shape=(1, N), + dtype=self.dtype, copy=False) + + def _getcol(self, i): + """Returns a copy of column i. A (m x 1) sparse array (column vector). + """ + if self.ndim == 1: + raise ValueError("getcol not provided for 1d arrays. Use indexing A[j]") + M, N = self.shape + i = int(i) + if i < 0: + i += N + if i < 0 or i >= N: + raise IndexError('index (%d) out of range' % i) + indptr, indices, data = get_csr_submatrix( + M, N, self.indptr, self.indices, self.data, 0, M, i, i + 1) + return self.__class__((data, indices, indptr), shape=(M, 1), + dtype=self.dtype, copy=False) + + def _get_int(self, idx): + spot = np.flatnonzero(self.indices == idx) + if spot.size: + return self.data[spot[0]] + return self.data.dtype.type(0) + + def _get_slice(self, idx): + if idx == slice(None): + return self.copy() + if idx.step in (1, None): + ret = self._get_submatrix(0, idx, copy=True) + return ret.reshape(ret.shape[-1]) + return self._minor_slice(idx) + + def _get_array(self, idx): + idx_dtype = self._get_index_dtype(self.indices) + idx = np.asarray(idx, dtype=idx_dtype) + if idx.size == 0: + return self.__class__([], dtype=self.dtype) + + M, N = 1, self.shape[0] + row = np.zeros_like(idx, dtype=idx_dtype) + col = np.asarray(idx, dtype=idx_dtype) + val = np.empty(row.size, dtype=self.dtype) + csr_sample_values(M, N, self.indptr, self.indices, self.data, + row.size, row, col, val) + + new_shape = col.shape if col.shape[0] > 1 else (col.shape[0],) + return self.__class__(val.reshape(new_shape)) + + def _get_intXarray(self, row, col): + return self._getrow(row)._minor_index_fancy(col) + + def _get_intXslice(self, row, col): + if col.step in (1, None): + return self._get_submatrix(row, col, copy=True) + # TODO: uncomment this once it's faster: + # return self._getrow(row)._minor_slice(col) + + M, N = self.shape + start, stop, stride = col.indices(N) + + ii, jj = self.indptr[row:row+2] + row_indices = self.indices[ii:jj] + row_data = self.data[ii:jj] + + if stride > 0: + ind = (row_indices >= start) & (row_indices < stop) + else: + ind = (row_indices <= start) & (row_indices > stop) + + if abs(stride) > 1: + ind &= (row_indices - start) % stride == 0 + + row_indices = (row_indices[ind] - start) // stride + row_data = row_data[ind] + row_indptr = np.array([0, len(row_indices)]) + + if stride < 0: + row_data = row_data[::-1] + row_indices = abs(row_indices[::-1]) + + shape = (1, max(0, int(np.ceil(float(stop - start) / stride)))) + return self.__class__((row_data, row_indices, row_indptr), shape=shape, + dtype=self.dtype, copy=False) + + def _get_sliceXint(self, row, col): + if row.step in (1, None): + return self._get_submatrix(row, col, copy=True) + return self._major_slice(row)._get_submatrix(minor=col) + + def _get_sliceXarray(self, row, col): + return self._major_slice(row)._minor_index_fancy(col) + + def _get_arrayXint(self, row, col): + res = self._major_index_fancy(row)._get_submatrix(minor=col) + if row.ndim > 1: + return res.reshape(row.shape) + return res + + def _get_arrayXslice(self, row, col): + if col.step not in (1, None): + col = np.arange(*col.indices(self.shape[1])) + return self._get_arrayXarray(row, col) + return self._major_index_fancy(row)._get_submatrix(minor=col) + + def _set_int(self, idx, x): + self._set_many(0, idx, x) + + def _set_array(self, idx, x): + x = np.broadcast_to(x, idx.shape) + self._set_many(np.zeros_like(idx), idx, x) + + +def isspmatrix_csr(x): + """Is `x` of csr_matrix type? + + Parameters + ---------- + x + object to check for being a csr matrix + + Returns + ------- + bool + True if `x` is a csr matrix, False otherwise + + Examples + -------- + >>> from scipy.sparse import csr_array, csr_matrix, coo_matrix, isspmatrix_csr + >>> isspmatrix_csr(csr_matrix([[5]])) + True + >>> isspmatrix_csr(csr_array([[5]])) + False + >>> isspmatrix_csr(coo_matrix([[5]])) + False + """ + return isinstance(x, csr_matrix) + + +# This namespace class separates array from matrix with isinstance +class csr_array(_csr_base, sparray): + """ + Compressed Sparse Row array. + + This can be instantiated in several ways: + csr_array(D) + where D is a 2-D ndarray + + csr_array(S) + with another sparse array or matrix S (equivalent to S.tocsr()) + + csr_array((M, N), [dtype]) + to construct an empty array with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + csr_array((data, (row_ind, col_ind)), [shape=(M, N)]) + where ``data``, ``row_ind`` and ``col_ind`` satisfy the + relationship ``a[row_ind[k], col_ind[k]] = data[k]``. + + csr_array((data, indices, indptr), [shape=(M, N)]) + is the standard CSR representation where the column indices for + row i are stored in ``indices[indptr[i]:indptr[i+1]]`` and their + corresponding values are stored in ``data[indptr[i]:indptr[i+1]]``. + If the shape parameter is not supplied, the array dimensions + are inferred from the index arrays. + + Attributes + ---------- + dtype : dtype + Data type of the array + shape : 2-tuple + Shape of the array + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + CSR format data array of the array + indices + CSR format index array of the array + indptr + CSR format index pointer array of the array + has_sorted_indices + has_canonical_format + T + + Notes + ----- + + Sparse arrays can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the CSR format + - efficient arithmetic operations CSR + CSR, CSR * CSR, etc. + - efficient row slicing + - fast matrix vector products + + Disadvantages of the CSR format + - slow column slicing operations (consider CSC) + - changes to the sparsity structure are expensive (consider LIL or DOK) + + Canonical Format + - Within each row, indices are sorted by column. + - There are no duplicate entries. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import csr_array + >>> csr_array((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> row = np.array([0, 0, 1, 2, 2, 2]) + >>> col = np.array([0, 2, 2, 0, 1, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csr_array((data, (row, col)), shape=(3, 3)).toarray() + array([[1, 0, 2], + [0, 0, 3], + [4, 5, 6]]) + + >>> indptr = np.array([0, 2, 3, 6]) + >>> indices = np.array([0, 2, 2, 0, 1, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csr_array((data, indices, indptr), shape=(3, 3)).toarray() + array([[1, 0, 2], + [0, 0, 3], + [4, 5, 6]]) + + Duplicate entries are summed together: + + >>> row = np.array([0, 1, 2, 0]) + >>> col = np.array([0, 1, 1, 0]) + >>> data = np.array([1, 2, 4, 8]) + >>> csr_array((data, (row, col)), shape=(3, 3)).toarray() + array([[9, 0, 0], + [0, 2, 0], + [0, 4, 0]]) + + As an example of how to construct a CSR array incrementally, + the following snippet builds a term-document array from texts: + + >>> docs = [["hello", "world", "hello"], ["goodbye", "cruel", "world"]] + >>> indptr = [0] + >>> indices = [] + >>> data = [] + >>> vocabulary = {} + >>> for d in docs: + ... for term in d: + ... index = vocabulary.setdefault(term, len(vocabulary)) + ... indices.append(index) + ... data.append(1) + ... indptr.append(len(indices)) + ... + >>> csr_array((data, indices, indptr), dtype=int).toarray() + array([[2, 1, 0, 0], + [0, 1, 1, 1]]) + + """ + + +class csr_matrix(spmatrix, _csr_base): + """ + Compressed Sparse Row matrix. + + This can be instantiated in several ways: + csr_matrix(D) + where D is a 2-D ndarray + + csr_matrix(S) + with another sparse array or matrix S (equivalent to S.tocsr()) + + csr_matrix((M, N), [dtype]) + to construct an empty matrix with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + csr_matrix((data, (row_ind, col_ind)), [shape=(M, N)]) + where ``data``, ``row_ind`` and ``col_ind`` satisfy the + relationship ``a[row_ind[k], col_ind[k]] = data[k]``. + + csr_matrix((data, indices, indptr), [shape=(M, N)]) + is the standard CSR representation where the column indices for + row i are stored in ``indices[indptr[i]:indptr[i+1]]`` and their + corresponding values are stored in ``data[indptr[i]:indptr[i+1]]``. + If the shape parameter is not supplied, the matrix dimensions + are inferred from the index arrays. + + Attributes + ---------- + dtype : dtype + Data type of the matrix + shape : 2-tuple + Shape of the matrix + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + CSR format data array of the matrix + indices + CSR format index array of the matrix + indptr + CSR format index pointer array of the matrix + has_sorted_indices + has_canonical_format + T + + Notes + ----- + + Sparse matrices can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the CSR format + - efficient arithmetic operations CSR + CSR, CSR * CSR, etc. + - efficient row slicing + - fast matrix vector products + + Disadvantages of the CSR format + - slow column slicing operations (consider CSC) + - changes to the sparsity structure are expensive (consider LIL or DOK) + + Canonical Format + - Within each row, indices are sorted by column. + - There are no duplicate entries. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import csr_matrix + >>> csr_matrix((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> row = np.array([0, 0, 1, 2, 2, 2]) + >>> col = np.array([0, 2, 2, 0, 1, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csr_matrix((data, (row, col)), shape=(3, 3)).toarray() + array([[1, 0, 2], + [0, 0, 3], + [4, 5, 6]]) + + >>> indptr = np.array([0, 2, 3, 6]) + >>> indices = np.array([0, 2, 2, 0, 1, 2]) + >>> data = np.array([1, 2, 3, 4, 5, 6]) + >>> csr_matrix((data, indices, indptr), shape=(3, 3)).toarray() + array([[1, 0, 2], + [0, 0, 3], + [4, 5, 6]]) + + Duplicate entries are summed together: + + >>> row = np.array([0, 1, 2, 0]) + >>> col = np.array([0, 1, 1, 0]) + >>> data = np.array([1, 2, 4, 8]) + >>> csr_matrix((data, (row, col)), shape=(3, 3)).toarray() + array([[9, 0, 0], + [0, 2, 0], + [0, 4, 0]]) + + As an example of how to construct a CSR matrix incrementally, + the following snippet builds a term-document matrix from texts: + + >>> docs = [["hello", "world", "hello"], ["goodbye", "cruel", "world"]] + >>> indptr = [0] + >>> indices = [] + >>> data = [] + >>> vocabulary = {} + >>> for d in docs: + ... for term in d: + ... index = vocabulary.setdefault(term, len(vocabulary)) + ... indices.append(index) + ... data.append(1) + ... indptr.append(len(indices)) + ... + >>> csr_matrix((data, indices, indptr), dtype=int).toarray() + array([[2, 1, 0, 0], + [0, 1, 1, 1]]) + + """ + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_data.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_data.py new file mode 100644 index 0000000000000000000000000000000000000000..8345414bee1a96be2b07f1db0bd747fd85a9bfdd --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_data.py @@ -0,0 +1,569 @@ +"""Base class for sparse matrice with a .data attribute + + subclasses must provide a _with_data() method that + creates a new matrix with the same sparsity pattern + as self but with a different data array + +""" + +import math +import numpy as np + +from ._base import _spbase, sparray, _ufuncs_with_fixed_point_at_zero +from ._sputils import isscalarlike, validateaxis + +__all__ = [] + + +# TODO implement all relevant operations +# use .data.__methods__() instead of /=, *=, etc. +class _data_matrix(_spbase): + def __init__(self, arg1, *, maxprint=None): + _spbase.__init__(self, arg1, maxprint=maxprint) + + @property + def dtype(self): + return self.data.dtype + + @dtype.setter + def dtype(self, newtype): + self.data.dtype = newtype + + def _deduped_data(self): + if hasattr(self, 'sum_duplicates'): + self.sum_duplicates() + return self.data + + def __abs__(self): + return self._with_data(abs(self._deduped_data())) + + def __round__(self, ndigits=0): + return self._with_data(np.around(self._deduped_data(), decimals=ndigits)) + + def _real(self): + return self._with_data(self.data.real) + + def _imag(self): + return self._with_data(self.data.imag) + + def __neg__(self): + if self.dtype.kind == 'b': + raise NotImplementedError('negating a boolean sparse array is not ' + 'supported') + return self._with_data(-self.data) + + def __imul__(self, other): # self *= other + if isscalarlike(other): + self.data *= other + return self + return NotImplemented + + def __itruediv__(self, other): # self /= other + if isscalarlike(other): + recip = 1.0 / other + self.data *= recip + return self + else: + return NotImplemented + + def astype(self, dtype, casting='unsafe', copy=True): + dtype = np.dtype(dtype) + if self.dtype != dtype: + matrix = self._with_data( + self.data.astype(dtype, casting=casting, copy=True), + copy=True + ) + return matrix._with_data(matrix._deduped_data(), copy=False) + elif copy: + return self.copy() + else: + return self + + astype.__doc__ = _spbase.astype.__doc__ + + def conjugate(self, copy=True): + if np.issubdtype(self.dtype, np.complexfloating): + return self._with_data(self.data.conjugate(), copy=copy) + elif copy: + return self.copy() + else: + return self + + conjugate.__doc__ = _spbase.conjugate.__doc__ + + def copy(self): + return self._with_data(self.data.copy(), copy=True) + + copy.__doc__ = _spbase.copy.__doc__ + + def power(self, n, dtype=None): + """ + This function performs element-wise power. + + Parameters + ---------- + n : scalar + n is a non-zero scalar (nonzero avoids dense ones creation) + If zero power is desired, special case it to use `np.ones` + + dtype : If dtype is not specified, the current dtype will be preserved. + + Raises + ------ + NotImplementedError : if n is a zero scalar + If zero power is desired, special case it to use + ``np.ones(A.shape, dtype=A.dtype)`` + """ + if not isscalarlike(n): + raise NotImplementedError("input is not scalar") + if not n: + raise NotImplementedError( + "zero power is not supported as it would densify the matrix.\n" + "Use `np.ones(A.shape, dtype=A.dtype)` for this case." + ) + + data = self._deduped_data() + if dtype is not None: + data = data.astype(dtype) + return self._with_data(data ** n) + + ########################### + # Multiplication handlers # + ########################### + + def _mul_scalar(self, other): + return self._with_data(self.data * other) + + +# Add the numpy unary ufuncs for which func(0) = 0 to _data_matrix. +for npfunc in _ufuncs_with_fixed_point_at_zero: + name = npfunc.__name__ + + def _create_method(op): + def method(self): + result = op(self._deduped_data()) + return self._with_data(result, copy=True) + + method.__doc__ = (f"Element-wise {name}.\n\n" + f"See `numpy.{name}` for more information.") + method.__name__ = name + + return method + + setattr(_data_matrix, name, _create_method(npfunc)) + + +def _find_missing_index(ind, n): + for k, a in enumerate(ind): + if k != a: + return k + + k += 1 + if k < n: + return k + else: + return -1 + + +class _minmax_mixin: + """Mixin for min and max methods. + + These are not implemented for dia_matrix, hence the separate class. + """ + + def _min_or_max_axis(self, axis, min_or_max, explicit): + N = self.shape[axis] + if N == 0: + raise ValueError("zero-size array to reduction operation") + M = self.shape[1 - axis] + idx_dtype = self._get_index_dtype(maxval=M) + + mat = self.tocsc() if axis == 0 else self.tocsr() + mat.sum_duplicates() + + major_index, value = mat._minor_reduce(min_or_max) + if not explicit: + not_full = np.diff(mat.indptr)[major_index] < N + value[not_full] = min_or_max(value[not_full], 0) + + mask = value != 0 + major_index = np.compress(mask, major_index).astype(idx_dtype, copy=False) + value = np.compress(mask, value) + + if isinstance(self, sparray): + coords = (major_index,) + shape = (M,) + return self._coo_container((value, coords), shape=shape, dtype=self.dtype) + + if axis == 0: + return self._coo_container( + (value, (np.zeros(len(value), dtype=idx_dtype), major_index)), + dtype=self.dtype, shape=(1, M) + ) + else: + return self._coo_container( + (value, (major_index, np.zeros(len(value), dtype=idx_dtype))), + dtype=self.dtype, shape=(M, 1) + ) + + def _min_or_max(self, axis, out, min_or_max, explicit): + if out is not None: + raise ValueError("Sparse arrays do not support an 'out' parameter.") + + validateaxis(axis) + if self.ndim == 1: + if axis not in (None, 0, -1): + raise ValueError("axis out of range") + axis = None # avoid calling special axis case. no impact on 1d + + if axis is None: + if 0 in self.shape: + raise ValueError("zero-size array to reduction operation") + + zero = self.dtype.type(0) + if self.nnz == 0: + return zero + m = min_or_max.reduce(self._deduped_data().ravel()) + if self.nnz != math.prod(self.shape) and not explicit: + m = min_or_max(zero, m) + return m + + if axis < 0: + axis += 2 + + if (axis == 0) or (axis == 1): + return self._min_or_max_axis(axis, min_or_max, explicit) + else: + raise ValueError("axis out of range") + + def _arg_min_or_max_axis(self, axis, argmin_or_argmax, compare, explicit): + if self.shape[axis] == 0: + raise ValueError("Cannot apply the operation along a zero-sized dimension.") + + if axis < 0: + axis += 2 + + zero = self.dtype.type(0) + + mat = self.tocsc() if axis == 0 else self.tocsr() + mat.sum_duplicates() + + ret_size, line_size = mat._swap(mat.shape) + ret = np.zeros(ret_size, dtype=int) + + nz_lines, = np.nonzero(np.diff(mat.indptr)) + for i in nz_lines: + p, q = mat.indptr[i:i + 2] + data = mat.data[p:q] + indices = mat.indices[p:q] + extreme_index = argmin_or_argmax(data) + extreme_value = data[extreme_index] + if explicit: + if q - p > 0: + ret[i] = indices[extreme_index] + else: + if compare(extreme_value, zero) or q - p == line_size: + ret[i] = indices[extreme_index] + else: + zero_ind = _find_missing_index(indices, line_size) + if extreme_value == zero: + ret[i] = min(extreme_index, zero_ind) + else: + ret[i] = zero_ind + + if isinstance(self, sparray): + return ret + + if axis == 1: + ret = ret.reshape(-1, 1) + + return self._ascontainer(ret) + + def _arg_min_or_max(self, axis, out, argmin_or_argmax, compare, explicit): + if out is not None: + raise ValueError("Sparse types do not support an 'out' parameter.") + + validateaxis(axis) + + if self.ndim == 1: + if axis not in (None, 0, -1): + raise ValueError("axis out of range") + axis = None # avoid calling special axis case. no impact on 1d + + if axis is not None: + return self._arg_min_or_max_axis(axis, argmin_or_argmax, compare, explicit) + + if 0 in self.shape: + raise ValueError("Cannot apply the operation to an empty matrix.") + + if self.nnz == 0: + if explicit: + raise ValueError("Cannot apply the operation to zero matrix " + "when explicit=True.") + return 0 + + zero = self.dtype.type(0) + mat = self.tocoo() + # Convert to canonical form: no duplicates, sorted indices. + mat.sum_duplicates() + extreme_index = argmin_or_argmax(mat.data) + if explicit: + return extreme_index + extreme_value = mat.data[extreme_index] + num_col = mat.shape[-1] + + # If the min value is less than zero, or max is greater than zero, + # then we do not need to worry about implicit zeros. + if compare(extreme_value, zero): + # cast to Python int to avoid overflow and RuntimeError + return int(mat.row[extreme_index]) * num_col + int(mat.col[extreme_index]) + + # Cheap test for the rare case where we have no implicit zeros. + size = math.prod(self.shape) + if size == mat.nnz: + return int(mat.row[extreme_index]) * num_col + int(mat.col[extreme_index]) + + # At this stage, any implicit zero could be the min or max value. + # After sum_duplicates(), the `row` and `col` arrays are guaranteed to + # be sorted in C-order, which means the linearized indices are sorted. + linear_indices = mat.row * num_col + mat.col + first_implicit_zero_index = _find_missing_index(linear_indices, size) + if extreme_value == zero: + return min(first_implicit_zero_index, extreme_index) + return first_implicit_zero_index + + def max(self, axis=None, out=None, *, explicit=False): + """Return the maximum of the array/matrix or maximum along an axis. + + By default, all elements are taken into account, not just the non-zero ones. + But with `explicit` set, only the stored elements are considered. + + Parameters + ---------- + axis : {-2, -1, 0, 1, None} optional + Axis along which the sum is computed. The default is to + compute the maximum over all elements, returning + a scalar (i.e., `axis` = `None`). + + out : None, optional + This argument is in the signature *solely* for NumPy + compatibility reasons. Do not pass in anything except + for the default value, as this argument is not used. + + explicit : {False, True} optional (default: False) + When set to True, only the stored elements will be considered. + If a row/column is empty, the sparse.coo_array returned + has no stored element (i.e. an implicit zero) for that row/column. + + .. versionadded:: 1.15.0 + + Returns + ------- + amax : coo_array or scalar + Maximum of `a`. If `axis` is None, the result is a scalar value. + If `axis` is given, the result is a sparse.coo_array of dimension + ``a.ndim - 1``. + + See Also + -------- + min : The minimum value of a sparse array/matrix along a given axis. + numpy.max : NumPy's implementation of 'max' + + """ + return self._min_or_max(axis, out, np.maximum, explicit) + + def min(self, axis=None, out=None, *, explicit=False): + """Return the minimum of the array/matrix or maximum along an axis. + + By default, all elements are taken into account, not just the non-zero ones. + But with `explicit` set, only the stored elements are considered. + + Parameters + ---------- + axis : {-2, -1, 0, 1, None} optional + Axis along which the sum is computed. The default is to + compute the minimum over all elements, returning + a scalar (i.e., `axis` = `None`). + + out : None, optional + This argument is in the signature *solely* for NumPy + compatibility reasons. Do not pass in anything except for + the default value, as this argument is not used. + + explicit : {False, True} optional (default: False) + When set to True, only the stored elements will be considered. + If a row/column is empty, the sparse.coo_array returned + has no stored element (i.e. an implicit zero) for that row/column. + + .. versionadded:: 1.15.0 + + Returns + ------- + amin : coo_matrix or scalar + Minimum of `a`. If `axis` is None, the result is a scalar value. + If `axis` is given, the result is a sparse.coo_array of dimension + ``a.ndim - 1``. + + See Also + -------- + max : The maximum value of a sparse array/matrix along a given axis. + numpy.min : NumPy's implementation of 'min' + + """ + return self._min_or_max(axis, out, np.minimum, explicit) + + def nanmax(self, axis=None, out=None, *, explicit=False): + """Return the maximum, ignoring any Nans, along an axis. + + Return the maximum, ignoring any Nans, of the array/matrix along an axis. + By default this takes all elements into account, but with `explicit` set, + only stored elements are considered. + + .. versionadded:: 1.11.0 + + Parameters + ---------- + axis : {-2, -1, 0, 1, None} optional + Axis along which the maximum is computed. The default is to + compute the maximum over all elements, returning + a scalar (i.e., `axis` = `None`). + + out : None, optional + This argument is in the signature *solely* for NumPy + compatibility reasons. Do not pass in anything except + for the default value, as this argument is not used. + + explicit : {False, True} optional (default: False) + When set to True, only the stored elements will be considered. + If a row/column is empty, the sparse.coo_array returned + has no stored element (i.e. an implicit zero) for that row/column. + + .. versionadded:: 1.15.0 + + Returns + ------- + amax : coo_array or scalar + Maximum of `a`. If `axis` is None, the result is a scalar value. + If `axis` is given, the result is a sparse.coo_array of dimension + ``a.ndim - 1``. + + See Also + -------- + nanmin : The minimum value of a sparse array/matrix along a given axis, + ignoring NaNs. + max : The maximum value of a sparse array/matrix along a given axis, + propagating NaNs. + numpy.nanmax : NumPy's implementation of 'nanmax'. + + """ + return self._min_or_max(axis, out, np.fmax, explicit) + + def nanmin(self, axis=None, out=None, *, explicit=False): + """Return the minimum, ignoring any Nans, along an axis. + + Return the minimum, ignoring any Nans, of the array/matrix along an axis. + By default this takes all elements into account, but with `explicit` set, + only stored elements are considered. + + .. versionadded:: 1.11.0 + + Parameters + ---------- + axis : {-2, -1, 0, 1, None} optional + Axis along which the minimum is computed. The default is to + compute the minimum over all elements, returning + a scalar (i.e., `axis` = `None`). + + out : None, optional + This argument is in the signature *solely* for NumPy + compatibility reasons. Do not pass in anything except for + the default value, as this argument is not used. + + explicit : {False, True} optional (default: False) + When set to True, only the stored elements will be considered. + If a row/column is empty, the sparse.coo_array returned + has no stored element (i.e. an implicit zero) for that row/column. + + .. versionadded:: 1.15.0 + + Returns + ------- + amin : coo_array or scalar + Minimum of `a`. If `axis` is None, the result is a scalar value. + If `axis` is given, the result is a sparse.coo_array of dimension + ``a.ndim - 1``. + + See Also + -------- + nanmax : The maximum value of a sparse array/matrix along a given axis, + ignoring NaNs. + min : The minimum value of a sparse array/matrix along a given axis, + propagating NaNs. + numpy.nanmin : NumPy's implementation of 'nanmin'. + + """ + return self._min_or_max(axis, out, np.fmin, explicit) + + def argmax(self, axis=None, out=None, *, explicit=False): + """Return indices of maximum elements along an axis. + + By default, implicit zero elements are taken into account. If there are + several minimum values, the index of the first occurrence is returned. + If `explicit` is set, only explicitly stored elements will be considered. + + Parameters + ---------- + axis : {-2, -1, 0, 1, None}, optional + Axis along which the argmax is computed. If None (default), index + of the maximum element in the flatten data is returned. + + out : None, optional + This argument is in the signature *solely* for NumPy + compatibility reasons. Do not pass in anything except for + the default value, as this argument is not used. + + explicit : {False, True} optional (default: False) + When set to True, only explicitly stored elements will be considered. + If axis is not None and a row/column has no stored elements, argmax + is undefined, so the index ``0`` is returned for that row/column. + + .. versionadded:: 1.15.0 + + Returns + ------- + ind : numpy.matrix or int + Indices of maximum elements. If matrix, its size along `axis` is 1. + """ + return self._arg_min_or_max(axis, out, np.argmax, np.greater, explicit) + + def argmin(self, axis=None, out=None, *, explicit=False): + """Return indices of minimum elements along an axis. + + By default, implicit zero elements are taken into account. If there are + several minimum values, the index of the first occurrence is returned. + If `explicit` is set, only explicitly stored elements will be considered. + + Parameters + ---------- + axis : {-2, -1, 0, 1, None}, optional + Axis along which the argmin is computed. If None (default), index + of the minimum element in the flatten data is returned. + + out : None, optional + This argument is in the signature *solely* for NumPy + compatibility reasons. Do not pass in anything except for + the default value, as this argument is not used. + + explicit : {False, True} optional (default: False) + When set to True, only explicitly stored elements will be considered. + If axis is not None and a row/column has no stored elements, argmin + is undefined, so the index ``0`` is returned for that row/column. + + .. versionadded:: 1.15.0 + + Returns + ------- + ind : numpy.matrix or int + Indices of minimum elements. If matrix, its size along `axis` is 1. + """ + return self._arg_min_or_max(axis, out, np.argmin, np.less, explicit) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_dia.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_dia.py new file mode 100644 index 0000000000000000000000000000000000000000..ac355dcac3cf6194f32db54f8010b31038b533e3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_dia.py @@ -0,0 +1,590 @@ +"""Sparse DIAgonal format""" + +__docformat__ = "restructuredtext en" + +__all__ = ['dia_array', 'dia_matrix', 'isspmatrix_dia'] + +import numpy as np + +from .._lib._util import copy_if_needed +from ._matrix import spmatrix +from ._base import issparse, _formats, _spbase, sparray +from ._data import _data_matrix +from ._sputils import ( + isshape, upcast_char, getdtype, get_sum_dtype, validateaxis, check_shape +) +from ._sparsetools import dia_matvec + + +class _dia_base(_data_matrix): + _format = 'dia' + + def __init__(self, arg1, shape=None, dtype=None, copy=False, *, maxprint=None): + _data_matrix.__init__(self, arg1, maxprint=maxprint) + + if issparse(arg1): + if arg1.format == "dia": + if copy: + arg1 = arg1.copy() + self.data = arg1.data + self.offsets = arg1.offsets + self._shape = check_shape(arg1.shape) + else: + if arg1.format == self.format and copy: + A = arg1.copy() + else: + A = arg1.todia() + self.data = A.data + self.offsets = A.offsets + self._shape = check_shape(A.shape) + elif isinstance(arg1, tuple): + if isshape(arg1): + # It's a tuple of matrix dimensions (M, N) + # create empty matrix + self._shape = check_shape(arg1) + self.data = np.zeros((0,0), getdtype(dtype, default=float)) + idx_dtype = self._get_index_dtype(maxval=max(self.shape)) + self.offsets = np.zeros((0), dtype=idx_dtype) + else: + try: + # Try interpreting it as (data, offsets) + data, offsets = arg1 + except Exception as e: + message = 'unrecognized form for dia_array constructor' + raise ValueError(message) from e + else: + if shape is None: + raise ValueError('expected a shape argument') + if not copy: + copy = copy_if_needed + self.data = np.atleast_2d(np.array(arg1[0], dtype=dtype, copy=copy)) + offsets = np.array(arg1[1], + dtype=self._get_index_dtype(maxval=max(shape)), + copy=copy) + self.offsets = np.atleast_1d(offsets) + self._shape = check_shape(shape) + else: + # must be dense, convert to COO first, then to DIA + try: + arg1 = np.asarray(arg1) + except Exception as e: + raise ValueError("unrecognized form for " + f"{self.format}_matrix constructor") from e + if isinstance(self, sparray) and arg1.ndim != 2: + raise ValueError(f"DIA arrays don't support {arg1.ndim}D input. Use 2D") + A = self._coo_container(arg1, dtype=dtype, shape=shape).todia() + self.data = A.data + self.offsets = A.offsets + self._shape = check_shape(A.shape) + + if dtype is not None: + newdtype = getdtype(dtype) + self.data = self.data.astype(newdtype) + + # check format + if self.offsets.ndim != 1: + raise ValueError('offsets array must have rank 1') + + if self.data.ndim != 2: + raise ValueError('data array must have rank 2') + + if self.data.shape[0] != len(self.offsets): + raise ValueError('number of diagonals (%d) ' + 'does not match the number of offsets (%d)' + % (self.data.shape[0], len(self.offsets))) + + if len(np.unique(self.offsets)) != len(self.offsets): + raise ValueError('offset array contains duplicate values') + + def __repr__(self): + _, fmt = _formats[self.format] + sparse_cls = 'array' if isinstance(self, sparray) else 'matrix' + d = self.data.shape[0] + return ( + f"<{fmt} sparse {sparse_cls} of dtype '{self.dtype}'\n" + f"\twith {self.nnz} stored elements ({d} diagonals) and shape {self.shape}>" + ) + + def _data_mask(self): + """Returns a mask of the same shape as self.data, where + mask[i,j] is True when data[i,j] corresponds to a stored element.""" + num_rows, num_cols = self.shape + offset_inds = np.arange(self.data.shape[1]) + row = offset_inds - self.offsets[:,None] + mask = (row >= 0) + mask &= (row < num_rows) + mask &= (offset_inds < num_cols) + return mask + + def count_nonzero(self, axis=None): + if axis is not None: + raise NotImplementedError( + "count_nonzero over an axis is not implemented for DIA format" + ) + mask = self._data_mask() + return np.count_nonzero(self.data[mask]) + + count_nonzero.__doc__ = _spbase.count_nonzero.__doc__ + + def _getnnz(self, axis=None): + if axis is not None: + raise NotImplementedError("_getnnz over an axis is not implemented " + "for DIA format") + M,N = self.shape + nnz = 0 + for k in self.offsets: + if k > 0: + nnz += min(M,N-k) + else: + nnz += min(M+k,N) + return int(nnz) + + _getnnz.__doc__ = _spbase._getnnz.__doc__ + + def sum(self, axis=None, dtype=None, out=None): + validateaxis(axis) + + if axis is not None and axis < 0: + axis += 2 + + res_dtype = get_sum_dtype(self.dtype) + num_rows, num_cols = self.shape + ret = None + + if axis == 0: + mask = self._data_mask() + x = (self.data * mask).sum(axis=0) + if x.shape[0] == num_cols: + res = x + else: + res = np.zeros(num_cols, dtype=x.dtype) + res[:x.shape[0]] = x + ret = self._ascontainer(res, dtype=res_dtype) + + else: + row_sums = np.zeros((num_rows, 1), dtype=res_dtype) + one = np.ones(num_cols, dtype=res_dtype) + dia_matvec(num_rows, num_cols, len(self.offsets), + self.data.shape[1], self.offsets, self.data, one, row_sums) + + row_sums = self._ascontainer(row_sums) + + if axis is None: + return row_sums.sum(dtype=dtype, out=out) + + ret = self._ascontainer(row_sums.sum(axis=axis)) + + return ret.sum(axis=(), dtype=dtype, out=out) + + sum.__doc__ = _spbase.sum.__doc__ + + def _add_sparse(self, other): + # If other is not DIA format, let them handle us instead. + if not isinstance(other, _dia_base): + return other._add_sparse(self) + + # Fast path for exact equality of the sparsity structure. + if np.array_equal(self.offsets, other.offsets): + return self._with_data(self.data + other.data) + + # Find the union of the offsets (which will be sorted and unique). + new_offsets = np.union1d(self.offsets, other.offsets) + self_idx = np.searchsorted(new_offsets, self.offsets) + other_idx = np.searchsorted(new_offsets, other.offsets) + + self_d = self.data.shape[1] + other_d = other.data.shape[1] + # Fast path for a sparsity structure where the final offsets are a + # permutation of the existing offsets and the diagonal lengths match. + if self_d == other_d and len(new_offsets) == len(self.offsets): + new_data = self.data[_invert_index(self_idx)] + new_data[other_idx, :] += other.data + elif self_d == other_d and len(new_offsets) == len(other.offsets): + new_data = other.data[_invert_index(other_idx)] + new_data[self_idx, :] += self.data + else: + # Maximum diagonal length of the result. + d = min(self.shape[0] + new_offsets[-1], self.shape[1]) + + # Add all diagonals to a freshly-allocated data array. + new_data = np.zeros( + (len(new_offsets), d), + dtype=np.result_type(self.data, other.data), + ) + new_data[self_idx, :self_d] += self.data[:, :d] + new_data[other_idx, :other_d] += other.data[:, :d] + return self._dia_container((new_data, new_offsets), shape=self.shape) + + def _mul_scalar(self, other): + return self._with_data(self.data * other) + + def _matmul_vector(self, other): + x = other + + y = np.zeros(self.shape[0], dtype=upcast_char(self.dtype.char, + x.dtype.char)) + + L = self.data.shape[1] + + M,N = self.shape + + dia_matvec(M,N, len(self.offsets), L, self.offsets, self.data, + x.ravel(), y.ravel()) + + return y + + def _setdiag(self, values, k=0): + M, N = self.shape + + if values.ndim == 0: + # broadcast + values_n = np.inf + else: + values_n = len(values) + + if k < 0: + n = min(M + k, N, values_n) + min_index = 0 + max_index = n + else: + n = min(M, N - k, values_n) + min_index = k + max_index = k + n + + if values.ndim != 0: + # allow also longer sequences + values = values[:n] + + data_rows, data_cols = self.data.shape + if k in self.offsets: + if max_index > data_cols: + data = np.zeros((data_rows, max_index), dtype=self.data.dtype) + data[:, :data_cols] = self.data + self.data = data + self.data[self.offsets == k, min_index:max_index] = values + else: + self.offsets = np.append(self.offsets, self.offsets.dtype.type(k)) + m = max(max_index, data_cols) + data = np.zeros((data_rows + 1, m), dtype=self.data.dtype) + data[:-1, :data_cols] = self.data + data[-1, min_index:max_index] = values + self.data = data + + def todia(self, copy=False): + if copy: + return self.copy() + else: + return self + + todia.__doc__ = _spbase.todia.__doc__ + + def transpose(self, axes=None, copy=False): + if axes is not None and axes != (1, 0): + raise ValueError("Sparse arrays/matrices do not support " + "an 'axes' parameter because swapping " + "dimensions is the only logical permutation.") + + num_rows, num_cols = self.shape + max_dim = max(self.shape) + + # flip diagonal offsets + offsets = -self.offsets + + # re-align the data matrix + r = np.arange(len(offsets), dtype=np.intc)[:, None] + c = np.arange(num_rows, dtype=np.intc) - (offsets % max_dim)[:, None] + pad_amount = max(0, max_dim-self.data.shape[1]) + data = np.hstack((self.data, np.zeros((self.data.shape[0], pad_amount), + dtype=self.data.dtype))) + data = data[r, c] + return self._dia_container((data, offsets), shape=( + num_cols, num_rows), copy=copy) + + transpose.__doc__ = _spbase.transpose.__doc__ + + def diagonal(self, k=0): + rows, cols = self.shape + if k <= -rows or k >= cols: + return np.empty(0, dtype=self.data.dtype) + idx, = np.nonzero(self.offsets == k) + first_col = max(0, k) + last_col = min(rows + k, cols) + result_size = last_col - first_col + if idx.size == 0: + return np.zeros(result_size, dtype=self.data.dtype) + result = self.data[idx[0], first_col:last_col] + padding = result_size - len(result) + if padding > 0: + result = np.pad(result, (0, padding), mode='constant') + return result + + diagonal.__doc__ = _spbase.diagonal.__doc__ + + def tocsc(self, copy=False): + if self.nnz == 0: + return self._csc_container(self.shape, dtype=self.dtype) + + num_rows, num_cols = self.shape + num_offsets, offset_len = self.data.shape + offset_inds = np.arange(offset_len) + + row = offset_inds - self.offsets[:,None] + mask = (row >= 0) + mask &= (row < num_rows) + mask &= (offset_inds < num_cols) + mask &= (self.data != 0) + + idx_dtype = self._get_index_dtype(maxval=max(self.shape)) + indptr = np.zeros(num_cols + 1, dtype=idx_dtype) + indptr[1:offset_len+1] = np.cumsum(mask.sum(axis=0)[:num_cols]) + if offset_len < num_cols: + indptr[offset_len+1:] = indptr[offset_len] + indices = row.T[mask.T].astype(idx_dtype, copy=False) + data = self.data.T[mask.T] + return self._csc_container((data, indices, indptr), shape=self.shape, + dtype=self.dtype) + + tocsc.__doc__ = _spbase.tocsc.__doc__ + + def tocoo(self, copy=False): + num_rows, num_cols = self.shape + num_offsets, offset_len = self.data.shape + offset_inds = np.arange(offset_len) + + row = offset_inds - self.offsets[:,None] + mask = (row >= 0) + mask &= (row < num_rows) + mask &= (offset_inds < num_cols) + mask &= (self.data != 0) + row = row[mask] + col = np.tile(offset_inds, num_offsets)[mask.ravel()] + idx_dtype = self._get_index_dtype( + arrays=(self.offsets,), maxval=max(self.shape) + ) + row = row.astype(idx_dtype, copy=False) + col = col.astype(idx_dtype, copy=False) + data = self.data[mask] + # Note: this cannot set has_canonical_format=True, because despite the + # lack of duplicates, we do not generate sorted indices. + return self._coo_container( + (data, (row, col)), shape=self.shape, dtype=self.dtype, copy=False + ) + + tocoo.__doc__ = _spbase.tocoo.__doc__ + + # needed by _data_matrix + def _with_data(self, data, copy=True): + """Returns a matrix with the same sparsity structure as self, + but with different data. By default the structure arrays are copied. + """ + if copy: + return self._dia_container( + (data, self.offsets.copy()), shape=self.shape + ) + else: + return self._dia_container( + (data, self.offsets), shape=self.shape + ) + + def resize(self, *shape): + shape = check_shape(shape) + M, N = shape + # we do not need to handle the case of expanding N + self.data = self.data[:, :N] + + if (M > self.shape[0] and + np.any(self.offsets + self.shape[0] < self.data.shape[1])): + # explicitly clear values that were previously hidden + mask = (self.offsets[:, None] + self.shape[0] <= + np.arange(self.data.shape[1])) + self.data[mask] = 0 + + self._shape = shape + + resize.__doc__ = _spbase.resize.__doc__ + + +def _invert_index(idx): + """Helper function to invert an index array.""" + inv = np.zeros_like(idx) + inv[idx] = np.arange(len(idx)) + return inv + + +def isspmatrix_dia(x): + """Is `x` of dia_matrix type? + + Parameters + ---------- + x + object to check for being a dia matrix + + Returns + ------- + bool + True if `x` is a dia matrix, False otherwise + + Examples + -------- + >>> from scipy.sparse import dia_array, dia_matrix, coo_matrix, isspmatrix_dia + >>> isspmatrix_dia(dia_matrix([[5]])) + True + >>> isspmatrix_dia(dia_array([[5]])) + False + >>> isspmatrix_dia(coo_matrix([[5]])) + False + """ + return isinstance(x, dia_matrix) + + +# This namespace class separates array from matrix with isinstance +class dia_array(_dia_base, sparray): + """ + Sparse array with DIAgonal storage. + + This can be instantiated in several ways: + dia_array(D) + where D is a 2-D ndarray + + dia_array(S) + with another sparse array or matrix S (equivalent to S.todia()) + + dia_array((M, N), [dtype]) + to construct an empty array with shape (M, N), + dtype is optional, defaulting to dtype='d'. + + dia_array((data, offsets), shape=(M, N)) + where the ``data[k,:]`` stores the diagonal entries for + diagonal ``offsets[k]`` (See example below) + + Attributes + ---------- + dtype : dtype + Data type of the array + shape : 2-tuple + Shape of the array + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + DIA format data array of the array + offsets + DIA format offset array of the array + T + + Notes + ----- + + Sparse arrays can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + Sparse arrays with DIAgonal storage do not support slicing. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import dia_array + >>> dia_array((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> data = np.array([[1, 2, 3, 4]]).repeat(3, axis=0) + >>> offsets = np.array([0, -1, 2]) + >>> dia_array((data, offsets), shape=(4, 4)).toarray() + array([[1, 0, 3, 0], + [1, 2, 0, 4], + [0, 2, 3, 0], + [0, 0, 3, 4]]) + + >>> from scipy.sparse import dia_array + >>> n = 10 + >>> ex = np.ones(n) + >>> data = np.array([ex, 2 * ex, ex]) + >>> offsets = np.array([-1, 0, 1]) + >>> dia_array((data, offsets), shape=(n, n)).toarray() + array([[2., 1., 0., ..., 0., 0., 0.], + [1., 2., 1., ..., 0., 0., 0.], + [0., 1., 2., ..., 0., 0., 0.], + ..., + [0., 0., 0., ..., 2., 1., 0.], + [0., 0., 0., ..., 1., 2., 1.], + [0., 0., 0., ..., 0., 1., 2.]]) + """ + + +class dia_matrix(spmatrix, _dia_base): + """ + Sparse matrix with DIAgonal storage. + + This can be instantiated in several ways: + dia_matrix(D) + where D is a 2-D ndarray + + dia_matrix(S) + with another sparse array or matrix S (equivalent to S.todia()) + + dia_matrix((M, N), [dtype]) + to construct an empty matrix with shape (M, N), + dtype is optional, defaulting to dtype='d'. + + dia_matrix((data, offsets), shape=(M, N)) + where the ``data[k,:]`` stores the diagonal entries for + diagonal ``offsets[k]`` (See example below) + + Attributes + ---------- + dtype : dtype + Data type of the matrix + shape : 2-tuple + Shape of the matrix + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + DIA format data array of the matrix + offsets + DIA format offset array of the matrix + T + + Notes + ----- + + Sparse matrices can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + Sparse matrices with DIAgonal storage do not support slicing. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.sparse import dia_matrix + >>> dia_matrix((3, 4), dtype=np.int8).toarray() + array([[0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]], dtype=int8) + + >>> data = np.array([[1, 2, 3, 4]]).repeat(3, axis=0) + >>> offsets = np.array([0, -1, 2]) + >>> dia_matrix((data, offsets), shape=(4, 4)).toarray() + array([[1, 0, 3, 0], + [1, 2, 0, 4], + [0, 2, 3, 0], + [0, 0, 3, 4]]) + + >>> from scipy.sparse import dia_matrix + >>> n = 10 + >>> ex = np.ones(n) + >>> data = np.array([ex, 2 * ex, ex]) + >>> offsets = np.array([-1, 0, 1]) + >>> dia_matrix((data, offsets), shape=(n, n)).toarray() + array([[2., 1., 0., ..., 0., 0., 0.], + [1., 2., 1., ..., 0., 0., 0.], + [0., 1., 2., ..., 0., 0., 0.], + ..., + [0., 0., 0., ..., 2., 1., 0.], + [0., 0., 0., ..., 1., 2., 1.], + [0., 0., 0., ..., 0., 1., 2.]]) + """ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_dok.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_dok.py new file mode 100644 index 0000000000000000000000000000000000000000..7bf395791730e7bcc59b3e670580de18f70e6a74 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_dok.py @@ -0,0 +1,692 @@ +"""Dictionary Of Keys based matrix""" + +__docformat__ = "restructuredtext en" + +__all__ = ['dok_array', 'dok_matrix', 'isspmatrix_dok'] + +import itertools +from warnings import warn +import numpy as np + +from ._matrix import spmatrix +from ._base import _spbase, sparray, issparse +from ._index import IndexMixin +from ._sputils import (isdense, getdtype, isshape, isintlike, isscalarlike, + upcast, upcast_scalar, check_shape) + + +class _dok_base(_spbase, IndexMixin, dict): + _format = 'dok' + _allow_nd = (1, 2) + + def __init__(self, arg1, shape=None, dtype=None, copy=False, *, maxprint=None): + _spbase.__init__(self, arg1, maxprint=maxprint) + + if isinstance(arg1, tuple) and isshape(arg1, allow_nd=self._allow_nd): + self._shape = check_shape(arg1, allow_nd=self._allow_nd) + self._dict = {} + self.dtype = getdtype(dtype, default=float) + elif issparse(arg1): # Sparse ctor + if arg1.format == self.format: + arg1 = arg1.copy() if copy else arg1 + else: + arg1 = arg1.todok() + + if dtype is not None: + arg1 = arg1.astype(dtype, copy=False) + + self._dict = arg1._dict + self._shape = check_shape(arg1.shape, allow_nd=self._allow_nd) + self.dtype = getdtype(arg1.dtype) + else: # Dense ctor + try: + arg1 = np.asarray(arg1) + except Exception as e: + raise TypeError('Invalid input format.') from e + + if arg1.ndim > 2: + raise ValueError(f"DOK arrays don't yet support {arg1.ndim}D input.") + + if arg1.ndim == 1: + if dtype is not None: + arg1 = arg1.astype(dtype) + self._dict = {i: v for i, v in enumerate(arg1) if v != 0} + self.dtype = getdtype(arg1.dtype) + else: + d = self._coo_container(arg1, shape=shape, dtype=dtype).todok() + self._dict = d._dict + self.dtype = getdtype(d.dtype) + self._shape = check_shape(arg1.shape, allow_nd=self._allow_nd) + + def update(self, val): + # Prevent direct usage of update + raise NotImplementedError("Direct update to DOK sparse format is not allowed.") + + def _getnnz(self, axis=None): + if axis is not None: + raise NotImplementedError( + "_getnnz over an axis is not implemented for DOK format." + ) + return len(self._dict) + + def count_nonzero(self, axis=None): + if axis is not None: + raise NotImplementedError( + "count_nonzero over an axis is not implemented for DOK format." + ) + return sum(x != 0 for x in self.values()) + + _getnnz.__doc__ = _spbase._getnnz.__doc__ + count_nonzero.__doc__ = _spbase.count_nonzero.__doc__ + + def __len__(self): + return len(self._dict) + + def __contains__(self, key): + return key in self._dict + + def setdefault(self, key, default=None, /): + return self._dict.setdefault(key, default) + + def __delitem__(self, key, /): + del self._dict[key] + + def clear(self): + return self._dict.clear() + + def pop(self, /, *args): + return self._dict.pop(*args) + + def __reversed__(self): + raise TypeError("reversed is not defined for dok_array type") + + def __or__(self, other): + type_names = f"{type(self).__name__} and {type(other).__name__}" + raise TypeError(f"unsupported operand type for |: {type_names}") + + def __ror__(self, other): + type_names = f"{type(self).__name__} and {type(other).__name__}" + raise TypeError(f"unsupported operand type for |: {type_names}") + + def __ior__(self, other): + type_names = f"{type(self).__name__} and {type(other).__name__}" + raise TypeError(f"unsupported operand type for |: {type_names}") + + def popitem(self): + return self._dict.popitem() + + def items(self): + return self._dict.items() + + def keys(self): + return self._dict.keys() + + def values(self): + return self._dict.values() + + def get(self, key, default=0.0): + """This provides dict.get method functionality with type checking""" + if key in self._dict: + return self._dict[key] + if isintlike(key) and self.ndim == 1: + key = (key,) + if self.ndim != len(key): + raise IndexError(f'Index {key} length needs to match self.shape') + try: + for i in key: + assert isintlike(i) + except (AssertionError, TypeError, ValueError) as e: + raise IndexError('Index must be or consist of integers.') from e + key = tuple(i + M if i < 0 else i for i, M in zip(key, self.shape)) + if any(i < 0 or i >= M for i, M in zip(key, self.shape)): + raise IndexError('Index out of bounds.') + if self.ndim == 1: + key = key[0] + return self._dict.get(key, default) + + # 1D get methods + def _get_int(self, idx): + return self._dict.get(idx, self.dtype.type(0)) + + def _get_slice(self, idx): + i_range = range(*idx.indices(self.shape[0])) + return self._get_array(list(i_range)) + + def _get_array(self, idx): + idx = np.asarray(idx) + if idx.ndim == 0: + val = self._dict.get(int(idx), self.dtype.type(0)) + return np.array(val, stype=self.dtype) + new_dok = self._dok_container(idx.shape, dtype=self.dtype) + dok_vals = [self._dict.get(i, 0) for i in idx.ravel()] + if dok_vals: + if len(idx.shape) == 1: + for i, v in enumerate(dok_vals): + if v: + new_dok._dict[i] = v + else: + new_idx = np.unravel_index(np.arange(len(dok_vals)), idx.shape) + new_idx = new_idx[0] if len(new_idx) == 1 else zip(*new_idx) + for i, v in zip(new_idx, dok_vals, strict=True): + if v: + new_dok._dict[i] = v + return new_dok + + # 2D get methods + def _get_intXint(self, row, col): + return self._dict.get((row, col), self.dtype.type(0)) + + def _get_intXslice(self, row, col): + return self._get_sliceXslice(slice(row, row + 1), col) + + def _get_sliceXint(self, row, col): + return self._get_sliceXslice(row, slice(col, col + 1)) + + def _get_sliceXslice(self, row, col): + row_start, row_stop, row_step = row.indices(self.shape[0]) + col_start, col_stop, col_step = col.indices(self.shape[1]) + row_range = range(row_start, row_stop, row_step) + col_range = range(col_start, col_stop, col_step) + shape = (len(row_range), len(col_range)) + # Switch paths only when advantageous + # (count the iterations in the loops, adjust for complexity) + if len(self) >= 2 * shape[0] * shape[1]: + # O(nr*nc) path: loop over + return self._get_columnXarray(row_range, col_range) + # O(nnz) path: loop over entries of self + newdok = self._dok_container(shape, dtype=self.dtype) + for key in self.keys(): + i, ri = divmod(int(key[0]) - row_start, row_step) + if ri != 0 or i < 0 or i >= shape[0]: + continue + j, rj = divmod(int(key[1]) - col_start, col_step) + if rj != 0 or j < 0 or j >= shape[1]: + continue + newdok._dict[i, j] = self._dict[key] + return newdok + + def _get_intXarray(self, row, col): + return self._get_columnXarray([row], col.ravel()) + + def _get_arrayXint(self, row, col): + res = self._get_columnXarray(row.ravel(), [col]) + if row.ndim > 1: + return res.reshape(row.shape) + return res + + def _get_sliceXarray(self, row, col): + row = list(range(*row.indices(self.shape[0]))) + return self._get_columnXarray(row, col) + + def _get_arrayXslice(self, row, col): + col = list(range(*col.indices(self.shape[1]))) + return self._get_columnXarray(row, col) + + def _get_columnXarray(self, row, col): + # outer indexing + newdok = self._dok_container((len(row), len(col)), dtype=self.dtype) + + for i, r in enumerate(row): + for j, c in enumerate(col): + v = self._dict.get((r, c), 0) + if v: + newdok._dict[i, j] = v + return newdok + + def _get_arrayXarray(self, row, col): + # inner indexing + i, j = map(np.atleast_2d, np.broadcast_arrays(row, col)) + newdok = self._dok_container(i.shape, dtype=self.dtype) + + for key in itertools.product(range(i.shape[0]), range(i.shape[1])): + v = self._dict.get((i[key], j[key]), 0) + if v: + newdok._dict[key] = v + return newdok + + # 1D set methods + def _set_int(self, idx, x): + if x: + self._dict[idx] = x + elif idx in self._dict: + del self._dict[idx] + + def _set_array(self, idx, x): + idx_set = idx.ravel() + x_set = x.ravel() + if len(idx_set) != len(x_set): + if len(x_set) == 1: + x_set = np.full(len(idx_set), x_set[0], dtype=self.dtype) + else: + raise ValueError("Need len(index)==len(data) or len(data)==1") + for i, v in zip(idx_set, x_set): + if v: + self._dict[i] = v + elif i in self._dict: + del self._dict[i] + + # 2D set methods + def _set_intXint(self, row, col, x): + key = (row, col) + if x: + self._dict[key] = x + elif key in self._dict: + del self._dict[key] + + def _set_arrayXarray(self, row, col, x): + row = list(map(int, row.ravel())) + col = list(map(int, col.ravel())) + x = x.ravel() + self._dict.update(zip(zip(row, col), x)) + + for i in np.nonzero(x == 0)[0]: + key = (row[i], col[i]) + if self._dict[key] == 0: + # may have been superseded by later update + del self._dict[key] + + def __add__(self, other): + if isscalarlike(other): + res_dtype = upcast_scalar(self.dtype, other) + new = self._dok_container(self.shape, dtype=res_dtype) + # Add this scalar to each element. + for key in itertools.product(*[range(d) for d in self.shape]): + aij = self._dict.get(key, 0) + other + if aij: + new[key] = aij + elif issparse(other): + if other.shape != self.shape: + raise ValueError("Matrix dimensions are not equal.") + res_dtype = upcast(self.dtype, other.dtype) + new = self._dok_container(self.shape, dtype=res_dtype) + new._dict = self._dict.copy() + if other.format == "dok": + o_items = other.items() + else: + other = other.tocoo() + if self.ndim == 1: + o_items = zip(other.coords[0], other.data) + else: + o_items = zip(zip(*other.coords), other.data) + with np.errstate(over='ignore'): + new._dict.update((k, new[k] + v) for k, v in o_items) + elif isdense(other): + new = self.todense() + other + else: + return NotImplemented + return new + + def __radd__(self, other): + return self + other # addition is commutative + + def __neg__(self): + if self.dtype.kind == 'b': + raise NotImplementedError( + 'Negating a sparse boolean matrix is not supported.' + ) + new = self._dok_container(self.shape, dtype=self.dtype) + new._dict.update((k, -v) for k, v in self.items()) + return new + + def _mul_scalar(self, other): + res_dtype = upcast_scalar(self.dtype, other) + # Multiply this scalar by every element. + new = self._dok_container(self.shape, dtype=res_dtype) + new._dict.update(((k, v * other) for k, v in self.items())) + return new + + def _matmul_vector(self, other): + res_dtype = upcast(self.dtype, other.dtype) + + # vector @ vector + if self.ndim == 1: + if issparse(other): + if other.format == "dok": + keys = self.keys() & other.keys() + else: + keys = self.keys() & other.tocoo().coords[0] + return res_dtype(sum(self._dict[k] * other._dict[k] for k in keys)) + elif isdense(other): + return res_dtype(sum(other[k] * v for k, v in self.items())) + else: + return NotImplemented + + # matrix @ vector + result = np.zeros(self.shape[0], dtype=res_dtype) + for (i, j), v in self.items(): + result[i] += v * other[j] + return result + + def _matmul_multivector(self, other): + result_dtype = upcast(self.dtype, other.dtype) + # vector @ multivector + if self.ndim == 1: + # works for other 1d or 2d + return sum(v * other[j] for j, v in self._dict.items()) + + # matrix @ multivector + M = self.shape[0] + new_shape = (M,) if other.ndim == 1 else (M, other.shape[1]) + result = np.zeros(new_shape, dtype=result_dtype) + for (i, j), v in self.items(): + result[i] += v * other[j] + return result + + def __imul__(self, other): + if isscalarlike(other): + self._dict.update((k, v * other) for k, v in self.items()) + return self + return NotImplemented + + def __truediv__(self, other): + if isscalarlike(other): + res_dtype = upcast_scalar(self.dtype, other) + new = self._dok_container(self.shape, dtype=res_dtype) + new._dict.update(((k, v / other) for k, v in self.items())) + return new + return self.tocsr() / other + + def __itruediv__(self, other): + if isscalarlike(other): + self._dict.update((k, v / other) for k, v in self.items()) + return self + return NotImplemented + + def __reduce__(self): + # this approach is necessary because __setstate__ is called after + # __setitem__ upon unpickling and since __init__ is not called there + # is no shape attribute hence it is not possible to unpickle it. + return dict.__reduce__(self) + + def diagonal(self, k=0): + if self.ndim == 2: + return super().diagonal(k) + raise ValueError("diagonal requires two dimensions") + + def transpose(self, axes=None, copy=False): + if self.ndim == 1: + return self.copy() + + if axes is not None and axes != (1, 0): + raise ValueError( + "Sparse arrays/matrices do not support " + "an 'axes' parameter because swapping " + "dimensions is the only logical permutation." + ) + + M, N = self.shape + new = self._dok_container((N, M), dtype=self.dtype, copy=copy) + new._dict.update((((right, left), val) for (left, right), val in self.items())) + return new + + transpose.__doc__ = _spbase.transpose.__doc__ + + def conjtransp(self): + """DEPRECATED: Return the conjugate transpose. + + .. deprecated:: 1.14.0 + + `conjtransp` is deprecated and will be removed in v1.16.0. + Use ``.T.conj()`` instead. + """ + msg = ("`conjtransp` is deprecated and will be removed in v1.16.0. " + "Use `.T.conj()` instead.") + warn(msg, DeprecationWarning, stacklevel=2) + + if self.ndim == 1: + new = self.tocoo() + new.data = new.data.conjugate() + return new + + M, N = self.shape + new = self._dok_container((N, M), dtype=self.dtype) + new._dict = {(right, left): np.conj(val) for (left, right), val in self.items()} + return new + + def copy(self): + new = self._dok_container(self.shape, dtype=self.dtype) + new._dict.update(self._dict) + return new + + copy.__doc__ = _spbase.copy.__doc__ + + @classmethod + def fromkeys(cls, iterable, value=1, /): + tmp = dict.fromkeys(iterable, value) + if isinstance(next(iter(tmp)), tuple): + shape = tuple(max(idx) + 1 for idx in zip(*tmp)) + else: + shape = (max(tmp) + 1,) + result = cls(shape, dtype=type(value)) + result._dict = tmp + return result + + def tocoo(self, copy=False): + nnz = self.nnz + if nnz == 0: + return self._coo_container(self.shape, dtype=self.dtype) + + idx_dtype = self._get_index_dtype(maxval=max(self.shape)) + data = np.fromiter(self.values(), dtype=self.dtype, count=nnz) + # handle 1d keys specially b/c not a tuple + inds = zip(*self.keys()) if self.ndim > 1 else (self.keys(),) + coords = tuple(np.fromiter(ix, dtype=idx_dtype, count=nnz) for ix in inds) + A = self._coo_container((data, coords), shape=self.shape, dtype=self.dtype) + A.has_canonical_format = True + return A + + tocoo.__doc__ = _spbase.tocoo.__doc__ + + def todok(self, copy=False): + if copy: + return self.copy() + return self + + todok.__doc__ = _spbase.todok.__doc__ + + def tocsc(self, copy=False): + if self.ndim == 1: + raise NotImplementedError("tocsr() not valid for 1d sparse array") + return self.tocoo(copy=False).tocsc(copy=copy) + + tocsc.__doc__ = _spbase.tocsc.__doc__ + + def resize(self, *shape): + shape = check_shape(shape, allow_nd=self._allow_nd) + if len(shape) != len(self.shape): + # TODO implement resize across dimensions + raise NotImplementedError + + if self.ndim == 1: + newN = shape[-1] + for i in list(self._dict): + if i >= newN: + del self._dict[i] + self._shape = shape + return + + newM, newN = shape + M, N = self.shape + if newM < M or newN < N: + # Remove all elements outside new dimensions + for i, j in list(self.keys()): + if i >= newM or j >= newN: + del self._dict[i, j] + self._shape = shape + + resize.__doc__ = _spbase.resize.__doc__ + + # Added for 1d to avoid `tocsr` from _base.py + def astype(self, dtype, casting='unsafe', copy=True): + dtype = np.dtype(dtype) + if self.dtype != dtype: + result = self._dok_container(self.shape, dtype=dtype) + data = np.array(list(self._dict.values()), dtype=dtype) + result._dict = dict(zip(self._dict, data)) + return result + elif copy: + return self.copy() + return self + + +def isspmatrix_dok(x): + """Is `x` of dok_array type? + + Parameters + ---------- + x + object to check for being a dok matrix + + Returns + ------- + bool + True if `x` is a dok matrix, False otherwise + + Examples + -------- + >>> from scipy.sparse import dok_array, dok_matrix, coo_matrix, isspmatrix_dok + >>> isspmatrix_dok(dok_matrix([[5]])) + True + >>> isspmatrix_dok(dok_array([[5]])) + False + >>> isspmatrix_dok(coo_matrix([[5]])) + False + """ + return isinstance(x, dok_matrix) + + +# This namespace class separates array from matrix with isinstance +class dok_array(_dok_base, sparray): + """ + Dictionary Of Keys based sparse array. + + This is an efficient structure for constructing sparse + arrays incrementally. + + This can be instantiated in several ways: + dok_array(D) + where D is a 2-D ndarray + + dok_array(S) + with another sparse array or matrix S (equivalent to S.todok()) + + dok_array((M,N), [dtype]) + create the array with initial shape (M,N) + dtype is optional, defaulting to dtype='d' + + Attributes + ---------- + dtype : dtype + Data type of the array + shape : 2-tuple + Shape of the array + ndim : int + Number of dimensions (this is always 2) + nnz + Number of nonzero elements + size + T + + Notes + ----- + + Sparse arrays can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + - Allows for efficient O(1) access of individual elements. + - Duplicates are not allowed. + - Can be efficiently converted to a coo_array once constructed. + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse import dok_array + >>> S = dok_array((5, 5), dtype=np.float32) + >>> for i in range(5): + ... for j in range(5): + ... S[i, j] = i + j # Update element + + """ + + +class dok_matrix(spmatrix, _dok_base): + """ + Dictionary Of Keys based sparse matrix. + + This is an efficient structure for constructing sparse + matrices incrementally. + + This can be instantiated in several ways: + dok_matrix(D) + where D is a 2-D ndarray + + dok_matrix(S) + with another sparse array or matrix S (equivalent to S.todok()) + + dok_matrix((M,N), [dtype]) + create the matrix with initial shape (M,N) + dtype is optional, defaulting to dtype='d' + + Attributes + ---------- + dtype : dtype + Data type of the matrix + shape : 2-tuple + Shape of the matrix + ndim : int + Number of dimensions (this is always 2) + nnz + Number of nonzero elements + size + T + + Notes + ----- + + Sparse matrices can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + - Allows for efficient O(1) access of individual elements. + - Duplicates are not allowed. + - Can be efficiently converted to a coo_matrix once constructed. + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse import dok_matrix + >>> S = dok_matrix((5, 5), dtype=np.float32) + >>> for i in range(5): + ... for j in range(5): + ... S[i, j] = i + j # Update element + + """ + + def set_shape(self, shape): + new_matrix = self.reshape(shape, copy=False).asformat(self.format) + self.__dict__ = new_matrix.__dict__ + + def get_shape(self): + """Get shape of a sparse matrix.""" + return self._shape + + shape = property(fget=get_shape, fset=set_shape) + + def __reversed__(self): + return self._dict.__reversed__() + + def __or__(self, other): + if isinstance(other, _dok_base): + return self._dict | other._dict + return self._dict | other + + def __ror__(self, other): + if isinstance(other, _dok_base): + return self._dict | other._dict + return self._dict | other + + def __ior__(self, other): + if isinstance(other, _dok_base): + self._dict |= other._dict + else: + self._dict |= other + return self diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_extract.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_extract.py new file mode 100644 index 0000000000000000000000000000000000000000..ff8b4552c4ff31dcae77054e6739ed8921b53854 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_extract.py @@ -0,0 +1,178 @@ +"""Functions to extract parts of sparse matrices +""" + +__docformat__ = "restructuredtext en" + +__all__ = ['find', 'tril', 'triu'] + + +from ._coo import coo_matrix, coo_array +from ._base import sparray + + +def find(A): + """Return the indices and values of the nonzero elements of a matrix + + Parameters + ---------- + A : dense or sparse array or matrix + Matrix whose nonzero elements are desired. + + Returns + ------- + (I,J,V) : tuple of arrays + I,J, and V contain the row indices, column indices, and values + of the nonzero entries. + + + Examples + -------- + >>> from scipy.sparse import csr_array, find + >>> A = csr_array([[7.0, 8.0, 0],[0, 0, 9.0]]) + >>> find(A) + (array([0, 0, 1], dtype=int32), + array([0, 1, 2], dtype=int32), + array([ 7., 8., 9.])) + + """ + + A = coo_array(A, copy=True) + A.sum_duplicates() + # remove explicit zeros + nz_mask = A.data != 0 + return A.row[nz_mask], A.col[nz_mask], A.data[nz_mask] + + +def tril(A, k=0, format=None): + """Return the lower triangular portion of a sparse array or matrix + + Returns the elements on or below the k-th diagonal of A. + - k = 0 corresponds to the main diagonal + - k > 0 is above the main diagonal + - k < 0 is below the main diagonal + + Parameters + ---------- + A : dense or sparse array or matrix + Matrix whose lower trianglar portion is desired. + k : integer : optional + The top-most diagonal of the lower triangle. + format : string + Sparse format of the result, e.g. format="csr", etc. + + Returns + ------- + L : sparse matrix + Lower triangular portion of A in sparse format. + + See Also + -------- + triu : upper triangle in sparse format + + Examples + -------- + >>> from scipy.sparse import csr_array, tril + >>> A = csr_array([[1, 2, 0, 0, 3], [4, 5, 0, 6, 7], [0, 0, 8, 9, 0]], + ... dtype='int32') + >>> A.toarray() + array([[1, 2, 0, 0, 3], + [4, 5, 0, 6, 7], + [0, 0, 8, 9, 0]], dtype=int32) + >>> tril(A).toarray() + array([[1, 0, 0, 0, 0], + [4, 5, 0, 0, 0], + [0, 0, 8, 0, 0]], dtype=int32) + >>> tril(A).nnz + 4 + >>> tril(A, k=1).toarray() + array([[1, 2, 0, 0, 0], + [4, 5, 0, 0, 0], + [0, 0, 8, 9, 0]], dtype=int32) + >>> tril(A, k=-1).toarray() + array([[0, 0, 0, 0, 0], + [4, 0, 0, 0, 0], + [0, 0, 0, 0, 0]], dtype=int32) + >>> tril(A, format='csc') + + + """ + coo_sparse = coo_array if isinstance(A, sparray) else coo_matrix + + # convert to COOrdinate format where things are easy + A = coo_sparse(A, copy=False) + mask = A.row + k >= A.col + + row = A.row[mask] + col = A.col[mask] + data = A.data[mask] + new_coo = coo_sparse((data, (row, col)), shape=A.shape, dtype=A.dtype) + return new_coo.asformat(format) + + +def triu(A, k=0, format=None): + """Return the upper triangular portion of a sparse array or matrix + + Returns the elements on or above the k-th diagonal of A. + - k = 0 corresponds to the main diagonal + - k > 0 is above the main diagonal + - k < 0 is below the main diagonal + + Parameters + ---------- + A : dense or sparse array or matrix + Matrix whose upper trianglar portion is desired. + k : integer : optional + The bottom-most diagonal of the upper triangle. + format : string + Sparse format of the result, e.g. format="csr", etc. + + Returns + ------- + L : sparse array or matrix + Upper triangular portion of A in sparse format. + Sparse array if A is a sparse array, otherwise matrix. + + See Also + -------- + tril : lower triangle in sparse format + + Examples + -------- + >>> from scipy.sparse import csr_array, triu + >>> A = csr_array([[1, 2, 0, 0, 3], [4, 5, 0, 6, 7], [0, 0, 8, 9, 0]], + ... dtype='int32') + >>> A.toarray() + array([[1, 2, 0, 0, 3], + [4, 5, 0, 6, 7], + [0, 0, 8, 9, 0]], dtype=int32) + >>> triu(A).toarray() + array([[1, 2, 0, 0, 3], + [0, 5, 0, 6, 7], + [0, 0, 8, 9, 0]], dtype=int32) + >>> triu(A).nnz + 8 + >>> triu(A, k=1).toarray() + array([[0, 2, 0, 0, 3], + [0, 0, 0, 6, 7], + [0, 0, 0, 9, 0]], dtype=int32) + >>> triu(A, k=-1).toarray() + array([[1, 2, 0, 0, 3], + [4, 5, 0, 6, 7], + [0, 0, 8, 9, 0]], dtype=int32) + >>> triu(A, format='csc') + + + """ + coo_sparse = coo_array if isinstance(A, sparray) else coo_matrix + + # convert to COOrdinate format where things are easy + A = coo_sparse(A, copy=False) + mask = A.row + k <= A.col + + row = A.row[mask] + col = A.col[mask] + data = A.data[mask] + new_coo = coo_sparse((data, (row, col)), shape=A.shape, dtype=A.dtype) + return new_coo.asformat(format) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_index.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_index.py new file mode 100644 index 0000000000000000000000000000000000000000..bf7594f2ad8ee8189f18a383ee728489a35e94c1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_index.py @@ -0,0 +1,444 @@ +"""Indexing mixin for sparse array/matrix classes. +""" +import numpy as np +from ._sputils import isintlike +from ._base import sparray, issparse + +INT_TYPES = (int, np.integer) + + +def _broadcast_arrays(a, b): + """ + Same as np.broadcast_arrays(a, b) but old writeability rules. + + NumPy >= 1.17.0 transitions broadcast_arrays to return + read-only arrays. Set writeability explicitly to avoid warnings. + Retain the old writeability rules, as our Cython code assumes + the old behavior. + """ + x, y = np.broadcast_arrays(a, b) + x.flags.writeable = a.flags.writeable + y.flags.writeable = b.flags.writeable + return x, y + + +class IndexMixin: + """ + This class provides common dispatching and validation logic for indexing. + """ + def __getitem__(self, key): + index, new_shape = self._validate_indices(key) + + # 1D array + if len(index) == 1: + idx = index[0] + if isinstance(idx, np.ndarray): + if idx.shape == (): + idx = idx.item() + if isinstance(idx, INT_TYPES): + res = self._get_int(idx) + elif isinstance(idx, slice): + res = self._get_slice(idx) + else: # assume array idx + res = self._get_array(idx) + + # package the result and return + if not isinstance(self, sparray): + return res + # handle np.newaxis in idx when result would otherwise be a scalar + if res.shape == () and new_shape != (): + if len(new_shape) == 1: + return self.__class__([res], shape=new_shape, dtype=self.dtype) + if len(new_shape) == 2: + return self.__class__([[res]], shape=new_shape, dtype=self.dtype) + return res.reshape(new_shape) + + # 2D array + row, col = index + + # Dispatch to specialized methods. + if isinstance(row, INT_TYPES): + if isinstance(col, INT_TYPES): + res = self._get_intXint(row, col) + elif isinstance(col, slice): + res = self._get_intXslice(row, col) + elif col.ndim == 1: + res = self._get_intXarray(row, col) + elif col.ndim == 2: + res = self._get_intXarray(row, col) + else: + raise IndexError('index results in >2 dimensions') + elif isinstance(row, slice): + if isinstance(col, INT_TYPES): + res = self._get_sliceXint(row, col) + elif isinstance(col, slice): + if row == slice(None) and row == col: + res = self.copy() + else: + res = self._get_sliceXslice(row, col) + elif col.ndim == 1: + res = self._get_sliceXarray(row, col) + else: + raise IndexError('index results in >2 dimensions') + else: + if isinstance(col, INT_TYPES): + res = self._get_arrayXint(row, col) + elif isinstance(col, slice): + res = self._get_arrayXslice(row, col) + # arrayXarray preprocess + elif (row.ndim == 2 and row.shape[1] == 1 + and (col.ndim == 1 or col.shape[0] == 1)): + # outer indexing + res = self._get_columnXarray(row[:, 0], col.ravel()) + else: + # inner indexing + row, col = _broadcast_arrays(row, col) + if row.shape != col.shape: + raise IndexError('number of row and column indices differ') + if row.size == 0: + res = self.__class__(np.atleast_2d(row).shape, dtype=self.dtype) + else: + res = self._get_arrayXarray(row, col) + + # handle spmatrix (must be 2d, dont let 1d new_shape start reshape) + if not isinstance(self, sparray): + if new_shape == () or (len(new_shape) == 1 and res.ndim != 0): + # res handles cases not inflated by None + return res + if len(new_shape) == 1: + # shape inflated to 1D by None in index. Make 2D + new_shape = (1,) + new_shape + # reshape if needed (when None changes shape, e.g. A[1,:,None]) + return res if new_shape == res.shape else res.reshape(new_shape) + + # package the result and return + if res.shape != new_shape: + # handle formats that support indexing but not 1D (lil for now) + if self.format == "lil" and len(new_shape) != 2: + if res.shape == (): + return self._coo_container([res], shape = new_shape) + return res.tocoo().reshape(new_shape) + return res.reshape(new_shape) + return res + + def __setitem__(self, key, x): + index, _ = self._validate_indices(key) + + # 1D array + if len(index) == 1: + idx = index[0] + + if issparse(x): + x = x.toarray() + else: + x = np.asarray(x, dtype=self.dtype) + + if isinstance(idx, INT_TYPES): + if x.size != 1: + raise ValueError('Trying to assign a sequence to an item') + self._set_int(idx, x.flat[0]) + return + + if isinstance(idx, slice): + # check for simple case of slice that gives 1 item + # Note: Python `range` does not use lots of memory + idx_range = range(*idx.indices(self.shape[0])) + N = len(idx_range) + if N == 1 and x.size == 1: + self._set_int(idx_range[0], x.flat[0]) + return + idx = np.arange(*idx.indices(self.shape[0])) + idx_shape = idx.shape + else: + idx_shape = idx.squeeze().shape + # broadcast scalar to full 1d + if x.squeeze().shape != idx_shape: + x = np.broadcast_to(x, idx.shape) + if x.size != 0: + self._set_array(idx, x) + return + + # 2D array + row, col = index + + if isinstance(row, INT_TYPES) and isinstance(col, INT_TYPES): + x = np.asarray(x, dtype=self.dtype) + if x.size != 1: + raise ValueError('Trying to assign a sequence to an item') + self._set_intXint(row, col, x.flat[0]) + return + + if isinstance(row, slice): + row = np.arange(*row.indices(self.shape[0]))[:, None] + else: + row = np.atleast_1d(row) + + if isinstance(col, slice): + col = np.arange(*col.indices(self.shape[1]))[None, :] + if row.ndim == 1: + row = row[:, None] + else: + col = np.atleast_1d(col) + + i, j = _broadcast_arrays(row, col) + if i.shape != j.shape: + raise IndexError('number of row and column indices differ') + + if issparse(x): + if 0 in x.shape: + return + if i.ndim == 1: + # Inner indexing, so treat them like row vectors. + i = i[None] + j = j[None] + x = x.tocoo(copy=False).reshape(x._shape_as_2d, copy=True) + broadcast_row = x.shape[0] == 1 and i.shape[0] != 1 + broadcast_col = x.shape[1] == 1 and i.shape[1] != 1 + if not ((broadcast_row or x.shape[0] == i.shape[0]) and + (broadcast_col or x.shape[1] == i.shape[1])): + raise ValueError('shape mismatch in assignment') + x.sum_duplicates() + self._set_arrayXarray_sparse(i, j, x) + else: + # Make x and i into the same shape + x = np.asarray(x, dtype=self.dtype) + if x.squeeze().shape != i.squeeze().shape: + x = np.broadcast_to(x, i.shape) + if x.size == 0: + return + x = x.reshape(i.shape) + self._set_arrayXarray(i, j, x) + + def _validate_indices(self, key): + """Returns two tuples: (index tuple, requested shape tuple)""" + # single ellipsis + if key is Ellipsis: + return (slice(None),) * self.ndim, self.shape + + if not isinstance(key, tuple): + key = [key] + + ellps_pos = None + index_1st = [] + prelim_ndim = 0 + for i, idx in enumerate(key): + if idx is Ellipsis: + if ellps_pos is not None: + raise IndexError('an index can only have a single ellipsis') + ellps_pos = i + elif idx is None: + index_1st.append(idx) + elif isinstance(idx, slice) or isintlike(idx): + index_1st.append(idx) + prelim_ndim += 1 + elif (ix := _compatible_boolean_index(idx, self.ndim)) is not None: + index_1st.append(ix) + prelim_ndim += ix.ndim + elif issparse(idx): + # TODO: make sparse matrix indexing work for sparray + raise IndexError( + 'Indexing with sparse matrices is not supported ' + 'except boolean indexing where matrix and index ' + 'are equal shapes.') + else: # dense array + index_1st.append(np.asarray(idx)) + prelim_ndim += 1 + ellip_slices = (self.ndim - prelim_ndim) * [slice(None)] + if ellip_slices: + if ellps_pos is None: + index_1st.extend(ellip_slices) + else: + index_1st = index_1st[:ellps_pos] + ellip_slices + index_1st[ellps_pos:] + + # second pass (have processed ellipsis and preprocessed arrays) + idx_shape = [] + index_ndim = 0 + index = [] + array_indices = [] + for i, idx in enumerate(index_1st): + if idx is None: + idx_shape.append(1) + elif isinstance(idx, slice): + index.append(idx) + Ms = self._shape[index_ndim] + len_slice = len(range(*idx.indices(Ms))) + idx_shape.append(len_slice) + index_ndim += 1 + elif isintlike(idx): + N = self._shape[index_ndim] + if not (-N <= idx < N): + raise IndexError(f'index ({idx}) out of range') + idx = int(idx + N if idx < 0 else idx) + index.append(idx) + index_ndim += 1 + # bool array (checked in first pass) + elif idx.dtype.kind == 'b': + ix = idx + tmp_ndim = index_ndim + ix.ndim + mid_shape = self._shape[index_ndim:tmp_ndim] + if ix.shape != mid_shape: + raise IndexError( + f"bool index {i} has shape {mid_shape} instead of {ix.shape}" + ) + index.extend(ix.nonzero()) + array_indices.extend(range(index_ndim, tmp_ndim)) + index_ndim = tmp_ndim + else: # dense array + N = self._shape[index_ndim] + idx = self._asindices(idx, N) + index.append(idx) + array_indices.append(index_ndim) + index_ndim += 1 + if index_ndim > self.ndim: + raise IndexError( + f'invalid index ndim. Array is {self.ndim}D. Index needs {index_ndim}D' + ) + if len(array_indices) > 1: + idx_arrays = _broadcast_arrays(*(index[i] for i in array_indices)) + if any(idx_arrays[0].shape != ix.shape for ix in idx_arrays[1:]): + shapes = " ".join(str(ix.shape) for ix in idx_arrays) + msg = (f'shape mismatch: indexing arrays could not be broadcast ' + f'together with shapes {shapes}') + raise IndexError(msg) + # TODO: handle this for nD (adjacent arrays stay, separated move to start) + idx_shape = list(idx_arrays[0].shape) + idx_shape + elif len(array_indices) == 1: + arr_index = array_indices[0] + arr_shape = list(index[arr_index].shape) + idx_shape = idx_shape[:arr_index] + arr_shape + idx_shape[arr_index:] + if (ndim := len(idx_shape)) > 2: + raise IndexError(f'Only 1D or 2D arrays allowed. Index makes {ndim}D') + return tuple(index), tuple(idx_shape) + + def _asindices(self, idx, length): + """Convert `idx` to a valid index for an axis with a given length. + + Subclasses that need special validation can override this method. + """ + try: + x = np.asarray(idx) + except (ValueError, TypeError, MemoryError) as e: + raise IndexError('invalid index') from e + + if x.ndim not in (1, 2): + raise IndexError('Index dimension must be 1 or 2') + + if x.size == 0: + return x + + # Check bounds + max_indx = x.max() + if max_indx >= length: + raise IndexError('index (%d) out of range' % max_indx) + + min_indx = x.min() + if min_indx < 0: + if min_indx < -length: + raise IndexError('index (%d) out of range' % min_indx) + if x is idx or not x.flags.owndata: + x = x.copy() + x[x < 0] += length + return x + + def _getrow(self, i): + """Return a copy of row i of the matrix, as a (1 x n) row vector. + """ + M, N = self.shape + i = int(i) + if i < -M or i >= M: + raise IndexError('index (%d) out of range' % i) + if i < 0: + i += M + return self._get_intXslice(i, slice(None)) + + def _getcol(self, i): + """Return a copy of column i of the matrix, as a (m x 1) column vector. + """ + M, N = self.shape + i = int(i) + if i < -N or i >= N: + raise IndexError('index (%d) out of range' % i) + if i < 0: + i += N + return self._get_sliceXint(slice(None), i) + + def _get_int(self, idx): + raise NotImplementedError() + + def _get_slice(self, idx): + raise NotImplementedError() + + def _get_array(self, idx): + raise NotImplementedError() + + def _get_intXint(self, row, col): + raise NotImplementedError() + + def _get_intXarray(self, row, col): + raise NotImplementedError() + + def _get_intXslice(self, row, col): + raise NotImplementedError() + + def _get_sliceXint(self, row, col): + raise NotImplementedError() + + def _get_sliceXslice(self, row, col): + raise NotImplementedError() + + def _get_sliceXarray(self, row, col): + raise NotImplementedError() + + def _get_arrayXint(self, row, col): + raise NotImplementedError() + + def _get_arrayXslice(self, row, col): + raise NotImplementedError() + + def _get_columnXarray(self, row, col): + raise NotImplementedError() + + def _get_arrayXarray(self, row, col): + raise NotImplementedError() + + def _set_int(self, idx, x): + raise NotImplementedError() + + def _set_array(self, idx, x): + raise NotImplementedError() + + def _set_intXint(self, row, col, x): + raise NotImplementedError() + + def _set_arrayXarray(self, row, col, x): + raise NotImplementedError() + + def _set_arrayXarray_sparse(self, row, col, x): + # Fall back to densifying x + x = np.asarray(x.toarray(), dtype=self.dtype) + x, _ = _broadcast_arrays(x, row) + self._set_arrayXarray(row, col, x) + + +def _compatible_boolean_index(idx, desired_ndim): + """Check for boolean array or array-like. peek before asarray for array-like""" + # use attribute ndim to indicate a compatible array and check dtype + # if not, look at 1st element as quick rejection of bool, else slower asanyarray + if not hasattr(idx, 'ndim'): + # is first element boolean? + try: + ix = next(iter(idx), None) + for _ in range(desired_ndim): + if isinstance(ix, bool): + break + ix = next(iter(ix), None) + else: + return None + except TypeError: + return None + # since first is boolean, construct array and check all elements + idx = np.asanyarray(idx) + + if idx.dtype.kind == 'b': + return idx + return None diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_lil.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_lil.py new file mode 100644 index 0000000000000000000000000000000000000000..60b0c9a7ed80e32d5ff373ff504c59ce6ca08199 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_lil.py @@ -0,0 +1,632 @@ +"""List of Lists sparse matrix class +""" + +__docformat__ = "restructuredtext en" + +__all__ = ['lil_array', 'lil_matrix', 'isspmatrix_lil'] + +from bisect import bisect_left + +import numpy as np + +from ._matrix import spmatrix +from ._base import _spbase, sparray, issparse +from ._index import IndexMixin, INT_TYPES, _broadcast_arrays +from ._sputils import (getdtype, isshape, isscalarlike, upcast_scalar, + check_shape, check_reshape_kwargs) +from . import _csparsetools + + +class _lil_base(_spbase, IndexMixin): + _format = 'lil' + + def __init__(self, arg1, shape=None, dtype=None, copy=False, *, maxprint=None): + _spbase.__init__(self, arg1, maxprint=maxprint) + self.dtype = getdtype(dtype, arg1, default=float) + + # First get the shape + if issparse(arg1): + if arg1.format == "lil" and copy: + A = arg1.copy() + else: + A = arg1.tolil() + + if dtype is not None: + newdtype = getdtype(dtype) + A = A.astype(newdtype, copy=False) + + self._shape = check_shape(A.shape) + self.dtype = A.dtype + self.rows = A.rows + self.data = A.data + elif isinstance(arg1,tuple): + if isshape(arg1): + if shape is not None: + raise ValueError('invalid use of shape parameter') + M, N = arg1 + self._shape = check_shape((M, N)) + self.rows = np.empty((M,), dtype=object) + self.data = np.empty((M,), dtype=object) + for i in range(M): + self.rows[i] = [] + self.data[i] = [] + else: + raise TypeError('unrecognized lil_array constructor usage') + else: + # assume A is dense + try: + A = self._ascontainer(arg1) + except TypeError as e: + raise TypeError('unsupported matrix type') from e + if isinstance(self, sparray) and A.ndim != 2: + raise ValueError(f"LIL arrays don't support {A.ndim}D input. Use 2D") + A = self._csr_container(A, dtype=dtype).tolil() + + self._shape = check_shape(A.shape) + self.dtype = getdtype(A.dtype) + self.rows = A.rows + self.data = A.data + + def __iadd__(self,other): + self[:,:] = self + other + return self + + def __isub__(self,other): + self[:,:] = self - other + return self + + def __imul__(self,other): + if isscalarlike(other): + self[:,:] = self * other + return self + else: + return NotImplemented + + def __itruediv__(self,other): + if isscalarlike(other): + self[:,:] = self / other + return self + else: + return NotImplemented + + # Whenever the dimensions change, empty lists should be created for each + # row + + def _getnnz(self, axis=None): + if axis is None: + return sum([len(rowvals) for rowvals in self.data]) + if axis < 0: + axis += 2 + if axis == 0: + out = np.zeros(self.shape[1], dtype=np.intp) + for row in self.rows: + out[row] += 1 + return out + elif axis == 1: + return np.array([len(rowvals) for rowvals in self.data], dtype=np.intp) + else: + raise ValueError('axis out of bounds') + + _getnnz.__doc__ = _spbase._getnnz.__doc__ + + def count_nonzero(self, axis=None): + if axis is None: + return sum(np.count_nonzero(rowvals) for rowvals in self.data) + + if axis < 0: + axis += 2 + if axis == 0: + out = np.zeros(self.shape[1], dtype=np.intp) + for row, data in zip(self.rows, self.data): + mask = [c for c, d in zip(row, data) if d != 0] + out[mask] += 1 + return out + elif axis == 1: + return np.array( + [np.count_nonzero(rowvals) for rowvals in self.data], dtype=np.intp, + ) + else: + raise ValueError('axis out of bounds') + + count_nonzero.__doc__ = _spbase.count_nonzero.__doc__ + + def getrowview(self, i): + """Returns a view of the 'i'th row (without copying). + """ + new = self._lil_container((1, self.shape[1]), dtype=self.dtype) + new.rows[0] = self.rows[i] + new.data[0] = self.data[i] + return new + + def getrow(self, i): + """Returns a copy of the 'i'th row. + """ + M, N = self.shape + if i < 0: + i += M + if i < 0 or i >= M: + raise IndexError('row index out of bounds') + new = self._lil_container((1, N), dtype=self.dtype) + new.rows[0] = self.rows[i][:] + new.data[0] = self.data[i][:] + return new + + def __getitem__(self, key): + # Fast path for simple (int, int) indexing. + if (isinstance(key, tuple) and len(key) == 2 and + isinstance(key[0], INT_TYPES) and + isinstance(key[1], INT_TYPES)): + # lil_get1 handles validation for us. + return self._get_intXint(*key) + # Everything else takes the normal path. + return IndexMixin.__getitem__(self, key) + + def _asindices(self, idx, N): + # LIL routines handle bounds-checking for us, so don't do it here. + try: + x = np.asarray(idx) + except (ValueError, TypeError, MemoryError) as e: + raise IndexError('invalid index') from e + if x.ndim not in (1, 2): + raise IndexError('Index dimension must be <= 2') + return x + + def _get_intXint(self, row, col): + v = _csparsetools.lil_get1(self.shape[0], self.shape[1], self.rows, + self.data, row, col) + return self.dtype.type(v) + + def _get_sliceXint(self, row, col): + row = range(*row.indices(self.shape[0])) + return self._get_row_ranges(row, slice(col, col+1)) + + def _get_arrayXint(self, row, col): + res = self._get_row_ranges(row.ravel(), slice(col, col+1)) + if row.ndim > 1: + return res.reshape(row.shape) + return res + + def _get_intXslice(self, row, col): + return self._get_row_ranges((row,), col) + + def _get_sliceXslice(self, row, col): + row = range(*row.indices(self.shape[0])) + return self._get_row_ranges(row, col) + + def _get_arrayXslice(self, row, col): + return self._get_row_ranges(row, col) + + def _get_intXarray(self, row, col): + row = np.array(row, dtype=col.dtype, ndmin=1) + return self._get_columnXarray(row, col) + + def _get_sliceXarray(self, row, col): + row = np.arange(*row.indices(self.shape[0])) + return self._get_columnXarray(row, col) + + def _get_columnXarray(self, row, col): + # outer indexing + row, col = _broadcast_arrays(row[:,None], col) + return self._get_arrayXarray(row, col) + + def _get_arrayXarray(self, row, col): + # inner indexing + i, j = map(np.atleast_2d, _prepare_index_for_memoryview(row, col)) + new = self._lil_container(i.shape, dtype=self.dtype) + _csparsetools.lil_fancy_get(self.shape[0], self.shape[1], + self.rows, self.data, + new.rows, new.data, + i, j) + return new + + def _get_row_ranges(self, rows, col_slice): + """ + Fast path for indexing in the case where column index is slice. + + This gains performance improvement over brute force by more + efficient skipping of zeros, by accessing the elements + column-wise in order. + + Parameters + ---------- + rows : sequence or range + Rows indexed. If range, must be within valid bounds. + col_slice : slice + Columns indexed + + """ + j_start, j_stop, j_stride = col_slice.indices(self.shape[1]) + col_range = range(j_start, j_stop, j_stride) + nj = len(col_range) + new = self._lil_container((len(rows), nj), dtype=self.dtype) + + _csparsetools.lil_get_row_ranges(self.shape[0], self.shape[1], + self.rows, self.data, + new.rows, new.data, + rows, + j_start, j_stop, j_stride, nj) + + return new + + def _set_intXint(self, row, col, x): + _csparsetools.lil_insert(self.shape[0], self.shape[1], self.rows, + self.data, row, col, x) + + def _set_arrayXarray(self, row, col, x): + i, j, x = map(np.atleast_2d, _prepare_index_for_memoryview(row, col, x)) + _csparsetools.lil_fancy_set(self.shape[0], self.shape[1], + self.rows, self.data, + i, j, x) + + def _set_arrayXarray_sparse(self, row, col, x): + # Fall back to densifying x + x = np.asarray(x.toarray(), dtype=self.dtype) + x, _ = _broadcast_arrays(x, row) + self._set_arrayXarray(row, col, x) + + def __setitem__(self, key, x): + if isinstance(key, tuple) and len(key) == 2: + row, col = key + # Fast path for simple (int, int) indexing. + if isinstance(row, INT_TYPES) and isinstance(col, INT_TYPES): + x = self.dtype.type(x) + if x.size > 1: + raise ValueError("Trying to assign a sequence to an item") + return self._set_intXint(row, col, x) + # Fast path for full-matrix sparse assignment. + if (isinstance(row, slice) and isinstance(col, slice) and + row == slice(None) and col == slice(None) and + issparse(x) and x.shape == self.shape): + x = self._lil_container(x, dtype=self.dtype) + self.rows = x.rows + self.data = x.data + return + # Everything else takes the normal path. + IndexMixin.__setitem__(self, key, x) + + def _mul_scalar(self, other): + if other == 0: + # Multiply by zero: return the zero matrix + new = self._lil_container(self.shape, dtype=self.dtype) + else: + res_dtype = upcast_scalar(self.dtype, other) + + new = self.copy() + new = new.astype(res_dtype) + # Multiply this scalar by every element. + for j, rowvals in enumerate(new.data): + new.data[j] = [val*other for val in rowvals] + return new + + def __truediv__(self, other): # self / other + if isscalarlike(other): + new = self.copy() + new.dtype = np.result_type(self, other) + # Divide every element by this scalar + for j, rowvals in enumerate(new.data): + new.data[j] = [val/other for val in rowvals] + return new + else: + return self.tocsr() / other + + def copy(self): + M, N = self.shape + new = self._lil_container(self.shape, dtype=self.dtype) + # This is ~14x faster than calling deepcopy() on rows and data. + _csparsetools.lil_get_row_ranges(M, N, self.rows, self.data, + new.rows, new.data, range(M), + 0, N, 1, N) + return new + + copy.__doc__ = _spbase.copy.__doc__ + + def reshape(self, *args, **kwargs): + shape = check_shape(args, self.shape) + order, copy = check_reshape_kwargs(kwargs) + + # Return early if reshape is not required + if shape == self.shape: + if copy: + return self.copy() + else: + return self + + new = self._lil_container(shape, dtype=self.dtype) + + if order == 'C': + ncols = self.shape[1] + for i, row in enumerate(self.rows): + for col, j in enumerate(row): + new_r, new_c = np.unravel_index(i * ncols + j, shape) + new[new_r, new_c] = self[i, j] + elif order == 'F': + nrows = self.shape[0] + for i, row in enumerate(self.rows): + for col, j in enumerate(row): + new_r, new_c = np.unravel_index(i + j * nrows, shape, order) + new[new_r, new_c] = self[i, j] + else: + raise ValueError("'order' must be 'C' or 'F'") + + return new + + reshape.__doc__ = _spbase.reshape.__doc__ + + def resize(self, *shape): + shape = check_shape(shape) + new_M, new_N = shape + M, N = self.shape + + if new_M < M: + self.rows = self.rows[:new_M] + self.data = self.data[:new_M] + elif new_M > M: + self.rows = np.resize(self.rows, new_M) + self.data = np.resize(self.data, new_M) + for i in range(M, new_M): + self.rows[i] = [] + self.data[i] = [] + + if new_N < N: + for row, data in zip(self.rows, self.data): + trunc = bisect_left(row, new_N) + del row[trunc:] + del data[trunc:] + + self._shape = shape + + resize.__doc__ = _spbase.resize.__doc__ + + def toarray(self, order=None, out=None): + d = self._process_toarray_args(order, out) + for i, row in enumerate(self.rows): + for pos, j in enumerate(row): + d[i, j] = self.data[i][pos] + return d + + toarray.__doc__ = _spbase.toarray.__doc__ + + def transpose(self, axes=None, copy=False): + return self.tocsr(copy=copy).transpose(axes=axes, copy=False).tolil(copy=False) + + transpose.__doc__ = _spbase.transpose.__doc__ + + def tolil(self, copy=False): + if copy: + return self.copy() + else: + return self + + tolil.__doc__ = _spbase.tolil.__doc__ + + def tocsr(self, copy=False): + M, N = self.shape + if M == 0 or N == 0: + return self._csr_container((M, N), dtype=self.dtype) + + # construct indptr array + if M*N <= np.iinfo(np.int32).max: + # fast path: it is known that 64-bit indexing will not be needed. + idx_dtype = np.int32 + indptr = np.empty(M + 1, dtype=idx_dtype) + indptr[0] = 0 + _csparsetools.lil_get_lengths(self.rows, indptr[1:]) + np.cumsum(indptr, out=indptr) + nnz = indptr[-1] + else: + idx_dtype = self._get_index_dtype(maxval=N) + lengths = np.empty(M, dtype=idx_dtype) + _csparsetools.lil_get_lengths(self.rows, lengths) + nnz = lengths.sum(dtype=np.int64) + idx_dtype = self._get_index_dtype(maxval=max(N, nnz)) + indptr = np.empty(M + 1, dtype=idx_dtype) + indptr[0] = 0 + np.cumsum(lengths, dtype=idx_dtype, out=indptr[1:]) + + indices = np.empty(nnz, dtype=idx_dtype) + data = np.empty(nnz, dtype=self.dtype) + _csparsetools.lil_flatten_to_array(self.rows, indices) + _csparsetools.lil_flatten_to_array(self.data, data) + + # init csr matrix + return self._csr_container((data, indices, indptr), shape=self.shape) + + tocsr.__doc__ = _spbase.tocsr.__doc__ + + +def _prepare_index_for_memoryview(i, j, x=None): + """ + Convert index and data arrays to form suitable for passing to the + Cython fancy getset routines. + + The conversions are necessary since to (i) ensure the integer + index arrays are in one of the accepted types, and (ii) to ensure + the arrays are writable so that Cython memoryview support doesn't + choke on them. + + Parameters + ---------- + i, j + Index arrays + x : optional + Data arrays + + Returns + ------- + i, j, x + Re-formatted arrays (x is omitted, if input was None) + + """ + if i.dtype > j.dtype: + j = j.astype(i.dtype) + elif i.dtype < j.dtype: + i = i.astype(j.dtype) + + if not i.flags.writeable or i.dtype not in (np.int32, np.int64): + i = i.astype(np.intp) + if not j.flags.writeable or j.dtype not in (np.int32, np.int64): + j = j.astype(np.intp) + + if x is not None: + if not x.flags.writeable: + x = x.copy() + return i, j, x + else: + return i, j + + +def isspmatrix_lil(x): + """Is `x` of lil_matrix type? + + Parameters + ---------- + x + object to check for being a lil matrix + + Returns + ------- + bool + True if `x` is a lil matrix, False otherwise + + Examples + -------- + >>> from scipy.sparse import lil_array, lil_matrix, coo_matrix, isspmatrix_lil + >>> isspmatrix_lil(lil_matrix([[5]])) + True + >>> isspmatrix_lil(lil_array([[5]])) + False + >>> isspmatrix_lil(coo_matrix([[5]])) + False + """ + return isinstance(x, lil_matrix) + + +# This namespace class separates array from matrix with isinstance +class lil_array(_lil_base, sparray): + """ + Row-based LIst of Lists sparse array. + + This is a structure for constructing sparse arrays incrementally. + Note that inserting a single item can take linear time in the worst case; + to construct the array efficiently, make sure the items are pre-sorted by + index, per row. + + This can be instantiated in several ways: + lil_array(D) + where D is a 2-D ndarray + + lil_array(S) + with another sparse array or matrix S (equivalent to S.tolil()) + + lil_array((M, N), [dtype]) + to construct an empty array with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + Attributes + ---------- + dtype : dtype + Data type of the array + shape : 2-tuple + Shape of the array + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + LIL format data array of the array + rows + LIL format row index array of the array + T + + Notes + ----- + Sparse arrays can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the LIL format + - supports flexible slicing + - changes to the array sparsity structure are efficient + + Disadvantages of the LIL format + - arithmetic operations LIL + LIL are slow (consider CSR or CSC) + - slow column slicing (consider CSC) + - slow matrix vector products (consider CSR or CSC) + + Intended Usage + - LIL is a convenient format for constructing sparse arrays + - once an array has been constructed, convert to CSR or + CSC format for fast arithmetic and matrix vector operations + - consider using the COO format when constructing large arrays + + Data Structure + - An array (``self.rows``) of rows, each of which is a sorted + list of column indices of non-zero elements. + - The corresponding nonzero values are stored in similar + fashion in ``self.data``. + + """ + + +class lil_matrix(spmatrix, _lil_base): + """ + Row-based LIst of Lists sparse matrix. + + This is a structure for constructing sparse matrices incrementally. + Note that inserting a single item can take linear time in the worst case; + to construct the matrix efficiently, make sure the items are pre-sorted by + index, per row. + + This can be instantiated in several ways: + lil_matrix(D) + where D is a 2-D ndarray + + lil_matrix(S) + with another sparse array or matrix S (equivalent to S.tolil()) + + lil_matrix((M, N), [dtype]) + to construct an empty matrix with shape (M, N) + dtype is optional, defaulting to dtype='d'. + + Attributes + ---------- + dtype : dtype + Data type of the matrix + shape : 2-tuple + Shape of the matrix + ndim : int + Number of dimensions (this is always 2) + nnz + size + data + LIL format data array of the matrix + rows + LIL format row index array of the matrix + T + + Notes + ----- + Sparse matrices can be used in arithmetic operations: they support + addition, subtraction, multiplication, division, and matrix power. + + Advantages of the LIL format + - supports flexible slicing + - changes to the matrix sparsity structure are efficient + + Disadvantages of the LIL format + - arithmetic operations LIL + LIL are slow (consider CSR or CSC) + - slow column slicing (consider CSC) + - slow matrix vector products (consider CSR or CSC) + + Intended Usage + - LIL is a convenient format for constructing sparse matrices + - once a matrix has been constructed, convert to CSR or + CSC format for fast arithmetic and matrix vector operations + - consider using the COO format when constructing large matrices + + Data Structure + - An array (``self.rows``) of rows, each of which is a sorted + list of column indices of non-zero elements. + - The corresponding nonzero values are stored in similar + fashion in ``self.data``. + + """ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_matrix.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..16b8093d71a55aa70c7230f7317b96f43d0ce376 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_matrix.py @@ -0,0 +1,146 @@ +class spmatrix: + """This class provides a base class for all sparse matrix classes. + + It cannot be instantiated. Most of the work is provided by subclasses. + """ + _allow_nd = (2,) + + @property + def _bsr_container(self): + from ._bsr import bsr_matrix + return bsr_matrix + + @property + def _coo_container(self): + from ._coo import coo_matrix + return coo_matrix + + @property + def _csc_container(self): + from ._csc import csc_matrix + return csc_matrix + + @property + def _csr_container(self): + from ._csr import csr_matrix + return csr_matrix + + @property + def _dia_container(self): + from ._dia import dia_matrix + return dia_matrix + + @property + def _dok_container(self): + from ._dok import dok_matrix + return dok_matrix + + @property + def _lil_container(self): + from ._lil import lil_matrix + return lil_matrix + + # Restore matrix multiplication + def __mul__(self, other): + return self._matmul_dispatch(other) + + def __rmul__(self, other): + return self._rmatmul_dispatch(other) + + # Restore matrix power + def __pow__(self, power): + from .linalg import matrix_power + + return matrix_power(self, power) + + ## Backward compatibility + + def set_shape(self, shape): + """Set the shape of the matrix in-place""" + # Make sure copy is False since this is in place + # Make sure format is unchanged because we are doing a __dict__ swap + new_self = self.reshape(shape, copy=False).asformat(self.format) + self.__dict__ = new_self.__dict__ + + def get_shape(self): + """Get the shape of the matrix""" + return self._shape + + shape = property(fget=get_shape, fset=set_shape, + doc="Shape of the matrix") + + def asfptype(self): + """Upcast matrix to a floating point format (if necessary)""" + return self._asfptype() + + def getmaxprint(self): + """Maximum number of elements to display when printed.""" + return self._getmaxprint() + + def getformat(self): + """Matrix storage format""" + return self.format + + def getnnz(self, axis=None): + """Number of stored values, including explicit zeros. + + Parameters + ---------- + axis : None, 0, or 1 + Select between the number of values across the whole array, in + each column, or in each row. + """ + return self._getnnz(axis=axis) + + def getH(self): + """Return the Hermitian transpose of this matrix. + + See Also + -------- + numpy.matrix.getH : NumPy's implementation of `getH` for matrices + """ + return self.conjugate().transpose() + + def getcol(self, j): + """Returns a copy of column j of the matrix, as an (m x 1) sparse + matrix (column vector). + """ + return self._getcol(j) + + def getrow(self, i): + """Returns a copy of row i of the matrix, as a (1 x n) sparse + matrix (row vector). + """ + return self._getrow(i) + + def todense(self, order=None, out=None): + """ + Return a dense representation of this sparse matrix. + + Parameters + ---------- + order : {'C', 'F'}, optional + Whether to store multi-dimensional data in C (row-major) + or Fortran (column-major) order in memory. The default + is 'None', which provides no ordering guarantees. + Cannot be specified in conjunction with the `out` + argument. + + out : ndarray, 2-D, optional + If specified, uses this array (or `numpy.matrix`) as the + output buffer instead of allocating a new array to + return. The provided array must have the same shape and + dtype as the sparse matrix on which you are calling the + method. + + Returns + ------- + arr : numpy.matrix, 2-D + A NumPy matrix object with the same shape and containing + the same data represented by the sparse matrix, with the + requested memory order. If `out` was passed and was an + array (rather than a `numpy.matrix`), it will be filled + with the appropriate values and returned wrapped in a + `numpy.matrix` object that shares the same memory. + """ + return super().todense(order, out) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_matrix_io.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_matrix_io.py new file mode 100644 index 0000000000000000000000000000000000000000..6b923f3f430eea247c8e2d548d2fc92aac650e96 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_matrix_io.py @@ -0,0 +1,167 @@ +import numpy as np +import scipy as sp + +__all__ = ['save_npz', 'load_npz'] + + +# Make loading safe vs. malicious input +PICKLE_KWARGS = dict(allow_pickle=False) + + +def save_npz(file, matrix, compressed=True): + """ Save a sparse matrix or array to a file using ``.npz`` format. + + Parameters + ---------- + file : str or file-like object + Either the file name (string) or an open file (file-like object) + where the data will be saved. If file is a string, the ``.npz`` + extension will be appended to the file name if it is not already + there. + matrix: spmatrix or sparray + The sparse matrix or array to save. + Supported formats: ``csc``, ``csr``, ``bsr``, ``dia`` or ``coo``. + compressed : bool, optional + Allow compressing the file. Default: True + + See Also + -------- + scipy.sparse.load_npz: Load a sparse matrix from a file using ``.npz`` format. + numpy.savez: Save several arrays into a ``.npz`` archive. + numpy.savez_compressed : Save several arrays into a compressed ``.npz`` archive. + + Examples + -------- + Store sparse matrix to disk, and load it again: + + >>> import numpy as np + >>> import scipy as sp + >>> sparse_matrix = sp.sparse.csc_matrix([[0, 0, 3], [4, 0, 0]]) + >>> sparse_matrix + + >>> sparse_matrix.toarray() + array([[0, 0, 3], + [4, 0, 0]], dtype=int64) + + >>> sp.sparse.save_npz('/tmp/sparse_matrix.npz', sparse_matrix) + >>> sparse_matrix = sp.sparse.load_npz('/tmp/sparse_matrix.npz') + + >>> sparse_matrix + + >>> sparse_matrix.toarray() + array([[0, 0, 3], + [4, 0, 0]], dtype=int64) + """ + arrays_dict = {} + if matrix.format in ('csc', 'csr', 'bsr'): + arrays_dict.update(indices=matrix.indices, indptr=matrix.indptr) + elif matrix.format == 'dia': + arrays_dict.update(offsets=matrix.offsets) + elif matrix.format == 'coo': + arrays_dict.update(row=matrix.row, col=matrix.col) + else: + msg = f'Save is not implemented for sparse matrix of format {matrix.format}.' + raise NotImplementedError(msg) + arrays_dict.update( + format=matrix.format.encode('ascii'), + shape=matrix.shape, + data=matrix.data + ) + if isinstance(matrix, sp.sparse.sparray): + arrays_dict.update(_is_array=True) + if compressed: + np.savez_compressed(file, **arrays_dict) + else: + np.savez(file, **arrays_dict) + + +def load_npz(file): + """ Load a sparse array/matrix from a file using ``.npz`` format. + + Parameters + ---------- + file : str or file-like object + Either the file name (string) or an open file (file-like object) + where the data will be loaded. + + Returns + ------- + result : csc_array, csr_array, bsr_array, dia_array or coo_array + A sparse array/matrix containing the loaded data. + + Raises + ------ + OSError + If the input file does not exist or cannot be read. + + See Also + -------- + scipy.sparse.save_npz: Save a sparse array/matrix to a file using ``.npz`` format. + numpy.load: Load several arrays from a ``.npz`` archive. + + Examples + -------- + Store sparse array/matrix to disk, and load it again: + + >>> import numpy as np + >>> import scipy as sp + >>> sparse_array = sp.sparse.csc_array([[0, 0, 3], [4, 0, 0]]) + >>> sparse_array + + >>> sparse_array.toarray() + array([[0, 0, 3], + [4, 0, 0]], dtype=int64) + + >>> sp.sparse.save_npz('/tmp/sparse_array.npz', sparse_array) + >>> sparse_array = sp.sparse.load_npz('/tmp/sparse_array.npz') + + >>> sparse_array + + >>> sparse_array.toarray() + array([[0, 0, 3], + [4, 0, 0]], dtype=int64) + + In this example we force the result to be csr_array from csr_matrix + >>> sparse_matrix = sp.sparse.csc_matrix([[0, 0, 3], [4, 0, 0]]) + >>> sp.sparse.save_npz('/tmp/sparse_matrix.npz', sparse_matrix) + >>> tmp = sp.sparse.load_npz('/tmp/sparse_matrix.npz') + >>> sparse_array = sp.sparse.csr_array(tmp) + """ + with np.load(file, **PICKLE_KWARGS) as loaded: + sparse_format = loaded.get('format') + if sparse_format is None: + raise ValueError(f'The file {file} does not contain ' + f'a sparse array or matrix.') + sparse_format = sparse_format.item() + + if not isinstance(sparse_format, str): + # Play safe with Python 2 vs 3 backward compatibility; + # files saved with SciPy < 1.0.0 may contain unicode or bytes. + sparse_format = sparse_format.decode('ascii') + + if loaded.get('_is_array'): + sparse_type = sparse_format + '_array' + else: + sparse_type = sparse_format + '_matrix' + + try: + cls = getattr(sp.sparse, f'{sparse_type}') + except AttributeError as e: + raise ValueError(f'Unknown format "{sparse_type}"') from e + + if sparse_format in ('csc', 'csr', 'bsr'): + return cls((loaded['data'], loaded['indices'], loaded['indptr']), + shape=loaded['shape']) + elif sparse_format == 'dia': + return cls((loaded['data'], loaded['offsets']), + shape=loaded['shape']) + elif sparse_format == 'coo': + return cls((loaded['data'], (loaded['row'], loaded['col'])), + shape=loaded['shape']) + else: + raise NotImplementedError(f'Load is not implemented for ' + f'sparse matrix of format {sparse_format}.') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_sparsetools.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_sparsetools.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..5ca70fbb56072792f2d60f6938b3e4d9ddba5e23 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_sparsetools.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_spfuncs.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_spfuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..86d52a51384cb01fbe64ced3c58696374b068d97 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_spfuncs.py @@ -0,0 +1,76 @@ +""" Functions that operate on sparse matrices +""" + +__all__ = ['count_blocks','estimate_blocksize'] + +from ._base import issparse +from ._csr import csr_array +from ._sparsetools import csr_count_blocks + + +def estimate_blocksize(A,efficiency=0.7): + """Attempt to determine the blocksize of a sparse matrix + + Returns a blocksize=(r,c) such that + - A.nnz / A.tobsr( (r,c) ).nnz > efficiency + """ + if not (issparse(A) and A.format in ("csc", "csr")): + A = csr_array(A) + + if A.nnz == 0: + return (1,1) + + if not 0 < efficiency < 1.0: + raise ValueError('efficiency must satisfy 0.0 < efficiency < 1.0') + + high_efficiency = (1.0 + efficiency) / 2.0 + nnz = float(A.nnz) + M,N = A.shape + + if M % 2 == 0 and N % 2 == 0: + e22 = nnz / (4 * count_blocks(A,(2,2))) + else: + e22 = 0.0 + + if M % 3 == 0 and N % 3 == 0: + e33 = nnz / (9 * count_blocks(A,(3,3))) + else: + e33 = 0.0 + + if e22 > high_efficiency and e33 > high_efficiency: + e66 = nnz / (36 * count_blocks(A,(6,6))) + if e66 > efficiency: + return (6,6) + else: + return (3,3) + else: + if M % 4 == 0 and N % 4 == 0: + e44 = nnz / (16 * count_blocks(A,(4,4))) + else: + e44 = 0.0 + + if e44 > efficiency: + return (4,4) + elif e33 > efficiency: + return (3,3) + elif e22 > efficiency: + return (2,2) + else: + return (1,1) + + +def count_blocks(A,blocksize): + """For a given blocksize=(r,c) count the number of occupied + blocks in a sparse matrix A + """ + r,c = blocksize + if r < 1 or c < 1: + raise ValueError('r and c must be positive') + + if issparse(A): + if A.format == "csr": + M,N = A.shape + return csr_count_blocks(M,N,r,c,A.indptr,A.indices) + elif A.format == "csc": + return count_blocks(A.T,(c,r)) + return count_blocks(csr_array(A),blocksize) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_sputils.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_sputils.py new file mode 100644 index 0000000000000000000000000000000000000000..ee829386aafa7ce4254805e532948c7aefcbf6f9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/_sputils.py @@ -0,0 +1,617 @@ +""" Utility functions for sparse matrix module +""" + +import sys +from typing import Any, Literal, Union +import operator +import numpy as np +from math import prod +import scipy.sparse as sp +from scipy._lib._util import np_long, np_ulong + + +__all__ = ['upcast', 'getdtype', 'getdata', 'isscalarlike', 'isintlike', + 'isshape', 'issequence', 'isdense', 'ismatrix', 'get_sum_dtype', + 'broadcast_shapes'] + +supported_dtypes = [np.bool_, np.byte, np.ubyte, np.short, np.ushort, np.intc, + np.uintc, np_long, np_ulong, np.longlong, np.ulonglong, + np.float32, np.float64, np.longdouble, + np.complex64, np.complex128, np.clongdouble] + +_upcast_memo = {} + + +def upcast(*args): + """Returns the nearest supported sparse dtype for the + combination of one or more types. + + upcast(t0, t1, ..., tn) -> T where T is a supported dtype + + Examples + -------- + >>> from scipy.sparse._sputils import upcast + >>> upcast('int32') + + >>> upcast('bool') + + >>> upcast('int32','float32') + + >>> upcast('bool',complex,float) + + + """ + + t = _upcast_memo.get(hash(args)) + if t is not None: + return t + + upcast = np.result_type(*args) + + for t in supported_dtypes: + if np.can_cast(upcast, t): + _upcast_memo[hash(args)] = t + return t + + raise TypeError(f'no supported conversion for types: {args!r}') + + +def upcast_char(*args): + """Same as `upcast` but taking dtype.char as input (faster).""" + t = _upcast_memo.get(args) + if t is not None: + return t + t = upcast(*map(np.dtype, args)) + _upcast_memo[args] = t + return t + + +def upcast_scalar(dtype, scalar): + """Determine data type for binary operation between an array of + type `dtype` and a scalar. + """ + return (np.array([0], dtype=dtype) * scalar).dtype + + +def downcast_intp_index(arr): + """ + Down-cast index array to np.intp dtype if it is of a larger dtype. + + Raise an error if the array contains a value that is too large for + intp. + """ + if arr.dtype.itemsize > np.dtype(np.intp).itemsize: + if arr.size == 0: + return arr.astype(np.intp) + maxval = arr.max() + minval = arr.min() + if maxval > np.iinfo(np.intp).max or minval < np.iinfo(np.intp).min: + raise ValueError("Cannot deal with arrays with indices larger " + "than the machine maximum address size " + "(e.g. 64-bit indices on 32-bit machine).") + return arr.astype(np.intp) + return arr + + +def to_native(A): + """ + Ensure that the data type of the NumPy array `A` has native byte order. + + `A` must be a NumPy array. If the data type of `A` does not have native + byte order, a copy of `A` with a native byte order is returned. Otherwise + `A` is returned. + """ + dt = A.dtype + if dt.isnative: + # Don't call `asarray()` if A is already native, to avoid unnecessarily + # creating a view of the input array. + return A + return np.asarray(A, dtype=dt.newbyteorder('native')) + + +def getdtype(dtype, a=None, default=None): + """Form a supported numpy dtype based on input arguments. + + Returns a valid ``numpy.dtype`` from `dtype` if not None, + or else ``a.dtype`` if possible, or else the given `default` + if not None, or else raise a ``TypeError``. + + The resulting ``dtype`` must be in ``supported_dtypes``: + bool_, int8, uint8, int16, uint16, int32, uint32, + int64, uint64, longlong, ulonglong, float32, float64, + longdouble, complex64, complex128, clongdouble + """ + if dtype is None: + try: + newdtype = a.dtype + except AttributeError as e: + if default is not None: + newdtype = np.dtype(default) + else: + raise TypeError("could not interpret data type") from e + else: + newdtype = np.dtype(dtype) + + if newdtype not in supported_dtypes: + supported_dtypes_fmt = ", ".join(t.__name__ for t in supported_dtypes) + raise ValueError(f"scipy.sparse does not support dtype {newdtype.name}. " + f"The only supported types are: {supported_dtypes_fmt}.") + return newdtype + + +def getdata(obj, dtype=None, copy=False) -> np.ndarray: + """ + This is a wrapper of `np.array(obj, dtype=dtype, copy=copy)` + that will generate a warning if the result is an object array. + """ + data = np.array(obj, dtype=dtype, copy=copy) + # Defer to getdtype for checking that the dtype is OK. + # This is called for the validation only; we don't need the return value. + getdtype(data.dtype) + return data + + +def safely_cast_index_arrays(A, idx_dtype=np.int32, msg=""): + """Safely cast sparse array indices to `idx_dtype`. + + Check the shape of `A` to determine if it is safe to cast its index + arrays to dtype `idx_dtype`. If any dimension in shape is larger than + fits in the dtype, casting is unsafe so raise ``ValueError``. + If safe, cast the index arrays to `idx_dtype` and return the result + without changing the input `A`. The caller can assign results to `A` + attributes if desired or use the recast index arrays directly. + + Unless downcasting is needed, the original index arrays are returned. + You can test e.g. ``A.indptr is new_indptr`` to see if downcasting occurred. + + .. versionadded:: 1.15.0 + + Parameters + ---------- + A : sparse array or matrix + The array for which index arrays should be downcast. + idx_dtype : dtype + Desired dtype. Should be an integer dtype (default: ``np.int32``). + Most of scipy.sparse uses either int64 or int32. + msg : string, optional + A string to be added to the end of the ValueError message + if the array shape is too big to fit in `idx_dtype`. + The error message is ``f" values too large for {msg}"`` + It should indicate why the downcasting is needed, e.g. "SuperLU", + and defaults to f"dtype {idx_dtype}". + + Returns + ------- + idx_arrays : ndarray or tuple of ndarrays + Based on ``A.format``, index arrays are returned after casting to `idx_dtype`. + For CSC/CSR, returns ``(indices, indptr)``. + For COO, returns ``coords``. + For DIA, returns ``offsets``. + For BSR, returns ``(indices, indptr)``. + + Raises + ------ + ValueError + If the array has shape that would not fit in the new dtype, or if + the sparse format does not use index arrays. + + Examples + -------- + >>> import numpy as np + >>> from scipy import sparse + >>> data = [3] + >>> coords = (np.array([3]), np.array([1])) # Note: int64 arrays + >>> A = sparse.coo_array((data, coords)) + >>> A.coords[0].dtype + dtype('int64') + + >>> # rescast after construction, raising exception if shape too big + >>> coords = sparse.safely_cast_index_arrays(A, np.int32) + >>> A.coords[0] is coords[0] # False if casting is needed + False + >>> A.coords = coords # set the index dtype of A + >>> A.coords[0].dtype + dtype('int32') + """ + if not msg: + msg = f"dtype {idx_dtype}" + # check for safe downcasting + max_value = np.iinfo(idx_dtype).max + + if A.format in ("csc", "csr"): + # indptr[-1] is max b/c indptr always sorted + if A.indptr[-1] > max_value: + raise ValueError(f"indptr values too large for {msg}") + + # check shape vs dtype + if max(*A.shape) > max_value: + if (A.indices > max_value).any(): + raise ValueError(f"indices values too large for {msg}") + + indices = A.indices.astype(idx_dtype, copy=False) + indptr = A.indptr.astype(idx_dtype, copy=False) + return indices, indptr + + elif A.format == "coo": + if max(*A.shape) > max_value: + if any((co > max_value).any() for co in A.coords): + raise ValueError(f"coords values too large for {msg}") + return tuple(co.astype(idx_dtype, copy=False) for co in A.coords) + + elif A.format == "dia": + if max(*A.shape) > max_value: + if (A.offsets > max_value).any(): + raise ValueError(f"offsets values too large for {msg}") + offsets = A.offsets.astype(idx_dtype, copy=False) + return offsets + + elif A.format == 'bsr': + R, C = A.blocksize + if A.indptr[-1] * R > max_value: + raise ValueError("indptr values too large for {msg}") + if max(*A.shape) > max_value: + if (A.indices * C > max_value).any(): + raise ValueError(f"indices values too large for {msg}") + indices = A.indices.astype(idx_dtype, copy=False) + indptr = A.indptr.astype(idx_dtype, copy=False) + return indices, indptr + + else: + raise TypeError(f'Format {A.format} is not associated with index arrays. ' + 'DOK and LIL have dict and list, not array.') + + +def get_index_dtype(arrays=(), maxval=None, check_contents=False): + """ + Based on input (integer) arrays `a`, determine a suitable index data + type that can hold the data in the arrays. + + Parameters + ---------- + arrays : tuple of array_like + Input arrays whose types/contents to check + maxval : float, optional + Maximum value needed + check_contents : bool, optional + Whether to check the values in the arrays and not just their types. + Default: False (check only the types) + + Returns + ------- + dtype : dtype + Suitable index data type (int32 or int64) + + Examples + -------- + >>> import numpy as np + >>> from scipy import sparse + >>> # select index dtype based on shape + >>> shape = (3, 3) + >>> idx_dtype = sparse.get_index_dtype(maxval=max(shape)) + >>> data = [1.1, 3.0, 1.5] + >>> indices = np.array([0, 1, 0], dtype=idx_dtype) + >>> indptr = np.array([0, 2, 3, 3], dtype=idx_dtype) + >>> A = sparse.csr_array((data, indices, indptr), shape=shape) + >>> A.indptr.dtype + dtype('int32') + + >>> # select based on larger of existing arrays and shape + >>> shape = (3, 3) + >>> idx_dtype = sparse.get_index_dtype(A.indptr, maxval=max(shape)) + >>> idx_dtype + + """ + # not using intc directly due to misinteractions with pythran + if np.intc().itemsize != 4: + return np.int64 + + int32min = np.int32(np.iinfo(np.int32).min) + int32max = np.int32(np.iinfo(np.int32).max) + + if maxval is not None: + maxval = np.int64(maxval) + if maxval > int32max: + return np.int64 + + if isinstance(arrays, np.ndarray): + arrays = (arrays,) + + for arr in arrays: + arr = np.asarray(arr) + if not np.can_cast(arr.dtype, np.int32): + if check_contents: + if arr.size == 0: + # a bigger type not needed + continue + elif np.issubdtype(arr.dtype, np.integer): + maxval = arr.max() + minval = arr.min() + if minval >= int32min and maxval <= int32max: + # a bigger type not needed + continue + return np.int64 + return np.int32 + + +def get_sum_dtype(dtype: np.dtype) -> np.dtype: + """Mimic numpy's casting for np.sum""" + if dtype.kind == 'u' and np.can_cast(dtype, np.uint): + return np.uint + if np.can_cast(dtype, np.int_): + return np.int_ + return dtype + + +def isscalarlike(x) -> bool: + """Is x either a scalar, an array scalar, or a 0-dim array?""" + return np.isscalar(x) or (isdense(x) and x.ndim == 0) + + +def isintlike(x) -> bool: + """Is x appropriate as an index into a sparse matrix? Returns True + if it can be cast safely to a machine int. + """ + # Fast-path check to eliminate non-scalar values. operator.index would + # catch this case too, but the exception catching is slow. + if np.ndim(x) != 0: + return False + try: + operator.index(x) + except (TypeError, ValueError): + try: + loose_int = bool(int(x) == x) + except (TypeError, ValueError): + return False + if loose_int: + msg = "Inexact indices into sparse matrices are not allowed" + raise ValueError(msg) + return loose_int + return True + + +def isshape(x, nonneg=False, *, allow_nd=(2,)) -> bool: + """Is x a valid tuple of dimensions? + + If nonneg, also checks that the dimensions are non-negative. + Shapes of length in the tuple allow_nd are allowed. + """ + ndim = len(x) + if ndim not in allow_nd: + return False + + for d in x: + if not isintlike(d): + return False + if nonneg and d < 0: + return False + return True + + +def issequence(t) -> bool: + return ((isinstance(t, list | tuple) and + (len(t) == 0 or np.isscalar(t[0]))) or + (isinstance(t, np.ndarray) and (t.ndim == 1))) + + +def ismatrix(t) -> bool: + return ((isinstance(t, list | tuple) and + len(t) > 0 and issequence(t[0])) or + (isinstance(t, np.ndarray) and t.ndim == 2)) + + +def isdense(x) -> bool: + return isinstance(x, np.ndarray) + + +def validateaxis(axis) -> None: + if axis is None: + return + axis_type = type(axis) + + # In NumPy, you can pass in tuples for 'axis', but they are + # not very useful for sparse matrices given their limited + # dimensions, so let's make it explicit that they are not + # allowed to be passed in + if isinstance(axis, tuple): + raise TypeError("Tuples are not accepted for the 'axis' parameter. " + "Please pass in one of the following: " + "{-2, -1, 0, 1, None}.") + + # If not a tuple, check that the provided axis is actually + # an integer and raise a TypeError similar to NumPy's + if not np.issubdtype(np.dtype(axis_type), np.integer): + raise TypeError(f"axis must be an integer, not {axis_type.__name__}") + + if not (-2 <= axis <= 1): + raise ValueError("axis out of range") + + +def check_shape(args, current_shape=None, *, allow_nd=(2,)) -> tuple[int, ...]: + """Imitate numpy.matrix handling of shape arguments + + Parameters + ---------- + args : array_like + Data structures providing information about the shape of the sparse array. + current_shape : tuple, optional + The current shape of the sparse array or matrix. + If None (default), the current shape will be inferred from args. + allow_nd : tuple of ints, optional default: (2,) + If shape does not have a length in the tuple allow_nd an error is raised. + + Returns + ------- + new_shape: tuple + The new shape after validation. + """ + if len(args) == 0: + raise TypeError("function missing 1 required positional argument: 'shape'") + if len(args) == 1: + try: + shape_iter = iter(args[0]) + except TypeError: + new_shape = (operator.index(args[0]), ) + else: + new_shape = tuple(operator.index(arg) for arg in shape_iter) + else: + new_shape = tuple(operator.index(arg) for arg in args) + + if current_shape is None: + if len(new_shape) not in allow_nd: + raise ValueError(f'shape must have length in {allow_nd}. Got {new_shape=}') + if any(d < 0 for d in new_shape): + raise ValueError("'shape' elements cannot be negative") + else: + # Check the current size only if needed + current_size = prod(current_shape) + + # Check for negatives + negative_indexes = [i for i, x in enumerate(new_shape) if x < 0] + if not negative_indexes: + new_size = prod(new_shape) + if new_size != current_size: + raise ValueError(f'cannot reshape array of size {current_size}' + f' into shape {new_shape}') + elif len(negative_indexes) == 1: + skip = negative_indexes[0] + specified = prod(new_shape[:skip] + new_shape[skip+1:]) + unspecified, remainder = divmod(current_size, specified) + if remainder != 0: + err_shape = tuple('newshape' if x < 0 else x for x in new_shape) + raise ValueError(f'cannot reshape array of size {current_size}' + f' into shape {err_shape}') + new_shape = new_shape[:skip] + (unspecified,) + new_shape[skip+1:] + else: + raise ValueError('can only specify one unknown dimension') + + if len(new_shape) not in allow_nd: + raise ValueError(f'shape must have length in {allow_nd}. Got {new_shape=}') + + return new_shape + + +def broadcast_shapes(*shapes): + """Check if shapes can be broadcast and return resulting shape + + This is similar to the NumPy ``broadcast_shapes`` function but + does not check memory consequences of the resulting dense matrix. + + Parameters + ---------- + *shapes : tuple of shape tuples + The tuple of shapes to be considered for broadcasting. + Shapes should be tuples of non-negative integers. + + Returns + ------- + new_shape : tuple of integers + The shape that results from broadcasting th input shapes. + """ + if not shapes: + return () + shapes = [shp if isinstance(shp, (tuple, list)) else (shp,) for shp in shapes] + big_shp = max(shapes, key=len) + out = list(big_shp) + for shp in shapes: + if shp is big_shp: + continue + for i, x in enumerate(shp, start=-len(shp)): + if x != 1 and x != out[i]: + if out[i] != 1: + raise ValueError("shapes cannot be broadcast to a single shape.") + out[i] = x + return (*out,) + + +def check_reshape_kwargs(kwargs): + """Unpack keyword arguments for reshape function. + + This is useful because keyword arguments after star arguments are not + allowed in Python 2, but star keyword arguments are. This function unpacks + 'order' and 'copy' from the star keyword arguments (with defaults) and + throws an error for any remaining. + """ + + order = kwargs.pop('order', 'C') + copy = kwargs.pop('copy', False) + if kwargs: # Some unused kwargs remain + raise TypeError("reshape() got unexpected keywords arguments: " + f"{', '.join(kwargs.keys())}") + return order, copy + + +def is_pydata_spmatrix(m) -> bool: + """ + Check whether object is pydata/sparse matrix, avoiding importing the module. + """ + base_cls = getattr(sys.modules.get('sparse'), 'SparseArray', None) + return base_cls is not None and isinstance(m, base_cls) + + +def convert_pydata_sparse_to_scipy( + arg: Any, + target_format: None | Literal["csc", "csr"] = None, + accept_fv: Any = None, +) -> Union[Any, "sp.spmatrix"]: + """ + Convert a pydata/sparse array to scipy sparse matrix, + pass through anything else. + """ + if is_pydata_spmatrix(arg): + # The `accept_fv` keyword is new in PyData Sparse 0.15.4 (May 2024), + # remove the `except` once the minimum supported version is >=0.15.4 + try: + arg = arg.to_scipy_sparse(accept_fv=accept_fv) + except TypeError: + arg = arg.to_scipy_sparse() + if target_format is not None: + arg = arg.asformat(target_format) + elif arg.format not in ("csc", "csr"): + arg = arg.tocsc() + return arg + + +############################################################################### +# Wrappers for NumPy types that are deprecated + +# Numpy versions of these functions raise deprecation warnings, the +# ones below do not. + +def matrix(*args, **kwargs): + return np.array(*args, **kwargs).view(np.matrix) + + +def asmatrix(data, dtype=None): + if isinstance(data, np.matrix) and (dtype is None or data.dtype == dtype): + return data + return np.asarray(data, dtype=dtype).view(np.matrix) + +############################################################################### + + +def _todata(s) -> np.ndarray: + """Access nonzero values, possibly after summing duplicates. + + Parameters + ---------- + s : sparse array + Input sparse array. + + Returns + ------- + data: ndarray + Nonzero values of the array, with shape (s.nnz,) + + """ + if isinstance(s, sp._data._data_matrix): + return s._deduped_data() + + if isinstance(s, sp.dok_array): + return np.fromiter(s.values(), dtype=s.dtype, count=s.nnz) + + if isinstance(s, sp.lil_array): + data = np.empty(s.nnz, dtype=s.dtype) + sp._csparsetools.lil_flatten_to_array(s.data, data) + return data + + return s.tocoo()._deduped_data() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..87e482f62d500f3b830f389456dc3c513c6028cf --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/__init__.py @@ -0,0 +1,210 @@ +r""" +Compressed sparse graph routines (:mod:`scipy.sparse.csgraph`) +============================================================== + +.. currentmodule:: scipy.sparse.csgraph + +Fast graph algorithms based on sparse matrix representations. + +Contents +-------- + +.. autosummary:: + :toctree: generated/ + + connected_components -- determine connected components of a graph + laplacian -- compute the laplacian of a graph + shortest_path -- compute the shortest path between points on a positive graph + dijkstra -- use Dijkstra's algorithm for shortest path + floyd_warshall -- use the Floyd-Warshall algorithm for shortest path + bellman_ford -- use the Bellman-Ford algorithm for shortest path + johnson -- use Johnson's algorithm for shortest path + yen -- use Yen's algorithm for K-shortest paths between to nodes. + breadth_first_order -- compute a breadth-first order of nodes + depth_first_order -- compute a depth-first order of nodes + breadth_first_tree -- construct the breadth-first tree from a given node + depth_first_tree -- construct a depth-first tree from a given node + minimum_spanning_tree -- construct the minimum spanning tree of a graph + reverse_cuthill_mckee -- compute permutation for reverse Cuthill-McKee ordering + maximum_flow -- solve the maximum flow problem for a graph + maximum_bipartite_matching -- compute a maximum matching of a bipartite graph + min_weight_full_bipartite_matching - compute a minimum weight full matching of a bipartite graph + structural_rank -- compute the structural rank of a graph + NegativeCycleError + +.. autosummary:: + :toctree: generated/ + + construct_dist_matrix + csgraph_from_dense + csgraph_from_masked + csgraph_masked_from_dense + csgraph_to_dense + csgraph_to_masked + reconstruct_path + +Graph Representations +--------------------- +This module uses graphs which are stored in a matrix format. A +graph with N nodes can be represented by an (N x N) adjacency matrix G. +If there is a connection from node i to node j, then G[i, j] = w, where +w is the weight of the connection. For nodes i and j which are +not connected, the value depends on the representation: + +- for dense array representations, non-edges are represented by + G[i, j] = 0, infinity, or NaN. + +- for dense masked representations (of type np.ma.MaskedArray), non-edges + are represented by masked values. This can be useful when graphs with + zero-weight edges are desired. + +- for sparse array representations, non-edges are represented by + non-entries in the matrix. This sort of sparse representation also + allows for edges with zero weights. + +As a concrete example, imagine that you would like to represent the following +undirected graph:: + + G + + (0) + / \ + 1 2 + / \ + (2) (1) + +This graph has three nodes, where node 0 and 1 are connected by an edge of +weight 2, and nodes 0 and 2 are connected by an edge of weight 1. +We can construct the dense, masked, and sparse representations as follows, +keeping in mind that an undirected graph is represented by a symmetric matrix:: + + >>> import numpy as np + >>> G_dense = np.array([[0, 2, 1], + ... [2, 0, 0], + ... [1, 0, 0]]) + >>> G_masked = np.ma.masked_values(G_dense, 0) + >>> from scipy.sparse import csr_array + >>> G_sparse = csr_array(G_dense) + +This becomes more difficult when zero edges are significant. For example, +consider the situation when we slightly modify the above graph:: + + G2 + + (0) + / \ + 0 2 + / \ + (2) (1) + +This is identical to the previous graph, except nodes 0 and 2 are connected +by an edge of zero weight. In this case, the dense representation above +leads to ambiguities: how can non-edges be represented if zero is a meaningful +value? In this case, either a masked or sparse representation must be used +to eliminate the ambiguity:: + + >>> import numpy as np + >>> G2_data = np.array([[np.inf, 2, 0 ], + ... [2, np.inf, np.inf], + ... [0, np.inf, np.inf]]) + >>> G2_masked = np.ma.masked_invalid(G2_data) + >>> from scipy.sparse.csgraph import csgraph_from_dense + >>> # G2_sparse = csr_array(G2_data) would give the wrong result + >>> G2_sparse = csgraph_from_dense(G2_data, null_value=np.inf) + >>> G2_sparse.data + array([ 2., 0., 2., 0.]) + +Here we have used a utility routine from the csgraph submodule in order to +convert the dense representation to a sparse representation which can be +understood by the algorithms in submodule. By viewing the data array, we +can see that the zero values are explicitly encoded in the graph. + +Directed vs. undirected +^^^^^^^^^^^^^^^^^^^^^^^ +Matrices may represent either directed or undirected graphs. This is +specified throughout the csgraph module by a boolean keyword. Graphs are +assumed to be directed by default. In a directed graph, traversal from node +i to node j can be accomplished over the edge G[i, j], but not the edge +G[j, i]. Consider the following dense graph:: + + >>> import numpy as np + >>> G_dense = np.array([[0, 1, 0], + ... [2, 0, 3], + ... [0, 4, 0]]) + +When ``directed=True`` we get the graph:: + + ---1--> ---3--> + (0) (1) (2) + <--2--- <--4--- + +In a non-directed graph, traversal from node i to node j can be +accomplished over either G[i, j] or G[j, i]. If both edges are not null, +and the two have unequal weights, then the smaller of the two is used. + +So for the same graph, when ``directed=False`` we get the graph:: + + (0)--1--(1)--3--(2) + +Note that a symmetric matrix will represent an undirected graph, regardless +of whether the 'directed' keyword is set to True or False. In this case, +using ``directed=True`` generally leads to more efficient computation. + +The routines in this module accept as input either scipy.sparse representations +(csr, csc, or lil format), masked representations, or dense representations +with non-edges indicated by zeros, infinities, and NaN entries. +""" # noqa: E501 + +__docformat__ = "restructuredtext en" + +__all__ = ['connected_components', + 'laplacian', + 'shortest_path', + 'floyd_warshall', + 'dijkstra', + 'bellman_ford', + 'johnson', + 'yen', + 'breadth_first_order', + 'depth_first_order', + 'breadth_first_tree', + 'depth_first_tree', + 'minimum_spanning_tree', + 'reverse_cuthill_mckee', + 'maximum_flow', + 'maximum_bipartite_matching', + 'min_weight_full_bipartite_matching', + 'structural_rank', + 'construct_dist_matrix', + 'reconstruct_path', + 'csgraph_masked_from_dense', + 'csgraph_from_dense', + 'csgraph_from_masked', + 'csgraph_to_dense', + 'csgraph_to_masked', + 'NegativeCycleError'] + +from ._laplacian import laplacian +from ._shortest_path import ( + shortest_path, floyd_warshall, dijkstra, bellman_ford, johnson, yen, + NegativeCycleError +) +from ._traversal import ( + breadth_first_order, depth_first_order, breadth_first_tree, + depth_first_tree, connected_components +) +from ._min_spanning_tree import minimum_spanning_tree +from ._flow import maximum_flow +from ._matching import ( + maximum_bipartite_matching, min_weight_full_bipartite_matching +) +from ._reordering import reverse_cuthill_mckee, structural_rank +from ._tools import ( + construct_dist_matrix, reconstruct_path, csgraph_from_dense, + csgraph_to_dense, csgraph_masked_from_dense, csgraph_from_masked, + csgraph_to_masked +) + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_flow.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_flow.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..3bd843cc4feb45aab4901c7e05d2150b900dc321 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_flow.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_laplacian.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_laplacian.py new file mode 100644 index 0000000000000000000000000000000000000000..1924c777243ad7fc4d3ba78bf2fc8c231c85bccc --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_laplacian.py @@ -0,0 +1,563 @@ +""" +Laplacian of a compressed-sparse graph +""" + +import numpy as np +from scipy.sparse import issparse +from scipy.sparse.linalg import LinearOperator +from scipy.sparse._sputils import convert_pydata_sparse_to_scipy, is_pydata_spmatrix + + +############################################################################### +# Graph laplacian +def laplacian( + csgraph, + normed=False, + return_diag=False, + use_out_degree=False, + *, + copy=True, + form="array", + dtype=None, + symmetrized=False, +): + """ + Return the Laplacian of a directed graph. + + Parameters + ---------- + csgraph : array_like or sparse array or matrix, 2 dimensions + compressed-sparse graph, with shape (N, N). + normed : bool, optional + If True, then compute symmetrically normalized Laplacian. + Default: False. + return_diag : bool, optional + If True, then also return an array related to vertex degrees. + Default: False. + use_out_degree : bool, optional + If True, then use out-degree instead of in-degree. + This distinction matters only if the graph is asymmetric. + Default: False. + copy: bool, optional + If False, then change `csgraph` in place if possible, + avoiding doubling the memory use. + Default: True, for backward compatibility. + form: 'array', or 'function', or 'lo' + Determines the format of the output Laplacian: + + * 'array' is a numpy array; + * 'function' is a pointer to evaluating the Laplacian-vector + or Laplacian-matrix product; + * 'lo' results in the format of the `LinearOperator`. + + Choosing 'function' or 'lo' always avoids doubling + the memory use, ignoring `copy` value. + Default: 'array', for backward compatibility. + dtype: None or one of numeric numpy dtypes, optional + The dtype of the output. If ``dtype=None``, the dtype of the + output matches the dtype of the input csgraph, except for + the case ``normed=True`` and integer-like csgraph, where + the output dtype is 'float' allowing accurate normalization, + but dramatically increasing the memory use. + Default: None, for backward compatibility. + symmetrized: bool, optional + If True, then the output Laplacian is symmetric/Hermitian. + The symmetrization is done by ``csgraph + csgraph.T.conj`` + without dividing by 2 to preserve integer dtypes if possible + prior to the construction of the Laplacian. + The symmetrization will increase the memory footprint of + sparse matrices unless the sparsity pattern is symmetric or + `form` is 'function' or 'lo'. + Default: False, for backward compatibility. + + Returns + ------- + lap : ndarray, or sparse array or matrix, or `LinearOperator` + The N x N Laplacian of csgraph. It will be a NumPy array (dense) + if the input was dense, or a sparse array otherwise, or + the format of a function or `LinearOperator` if + `form` equals 'function' or 'lo', respectively. + diag : ndarray, optional + The length-N main diagonal of the Laplacian matrix. + For the normalized Laplacian, this is the array of square roots + of vertex degrees or 1 if the degree is zero. + + Notes + ----- + The Laplacian matrix of a graph is sometimes referred to as the + "Kirchhoff matrix" or just the "Laplacian", and is useful in many + parts of spectral graph theory. + In particular, the eigen-decomposition of the Laplacian can give + insight into many properties of the graph, e.g., + is commonly used for spectral data embedding and clustering. + + The constructed Laplacian doubles the memory use if ``copy=True`` and + ``form="array"`` which is the default. + Choosing ``copy=False`` has no effect unless ``form="array"`` + or the matrix is sparse in the ``coo`` format, or dense array, except + for the integer input with ``normed=True`` that forces the float output. + + Sparse input is reformatted into ``coo`` if ``form="array"``, + which is the default. + + If the input adjacency matrix is not symmetric, the Laplacian is + also non-symmetric unless ``symmetrized=True`` is used. + + Diagonal entries of the input adjacency matrix are ignored and + replaced with zeros for the purpose of normalization where ``normed=True``. + The normalization uses the inverse square roots of row-sums of the input + adjacency matrix, and thus may fail if the row-sums contain + negative or complex with a non-zero imaginary part values. + + The normalization is symmetric, making the normalized Laplacian also + symmetric if the input csgraph was symmetric. + + References + ---------- + .. [1] Laplacian matrix. https://en.wikipedia.org/wiki/Laplacian_matrix + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse import csgraph + + Our first illustration is the symmetric graph + + >>> G = np.arange(4) * np.arange(4)[:, np.newaxis] + >>> G + array([[0, 0, 0, 0], + [0, 1, 2, 3], + [0, 2, 4, 6], + [0, 3, 6, 9]]) + + and its symmetric Laplacian matrix + + >>> csgraph.laplacian(G) + array([[ 0, 0, 0, 0], + [ 0, 5, -2, -3], + [ 0, -2, 8, -6], + [ 0, -3, -6, 9]]) + + The non-symmetric graph + + >>> G = np.arange(9).reshape(3, 3) + >>> G + array([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + + has different row- and column sums, resulting in two varieties + of the Laplacian matrix, using an in-degree, which is the default + + >>> L_in_degree = csgraph.laplacian(G) + >>> L_in_degree + array([[ 9, -1, -2], + [-3, 8, -5], + [-6, -7, 7]]) + + or alternatively an out-degree + + >>> L_out_degree = csgraph.laplacian(G, use_out_degree=True) + >>> L_out_degree + array([[ 3, -1, -2], + [-3, 8, -5], + [-6, -7, 13]]) + + Constructing a symmetric Laplacian matrix, one can add the two as + + >>> L_in_degree + L_out_degree.T + array([[ 12, -4, -8], + [ -4, 16, -12], + [ -8, -12, 20]]) + + or use the ``symmetrized=True`` option + + >>> csgraph.laplacian(G, symmetrized=True) + array([[ 12, -4, -8], + [ -4, 16, -12], + [ -8, -12, 20]]) + + that is equivalent to symmetrizing the original graph + + >>> csgraph.laplacian(G + G.T) + array([[ 12, -4, -8], + [ -4, 16, -12], + [ -8, -12, 20]]) + + The goal of normalization is to make the non-zero diagonal entries + of the Laplacian matrix to be all unit, also scaling off-diagonal + entries correspondingly. The normalization can be done manually, e.g., + + >>> G = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]) + >>> L, d = csgraph.laplacian(G, return_diag=True) + >>> L + array([[ 2, -1, -1], + [-1, 2, -1], + [-1, -1, 2]]) + >>> d + array([2, 2, 2]) + >>> scaling = np.sqrt(d) + >>> scaling + array([1.41421356, 1.41421356, 1.41421356]) + >>> (1/scaling)*L*(1/scaling) + array([[ 1. , -0.5, -0.5], + [-0.5, 1. , -0.5], + [-0.5, -0.5, 1. ]]) + + Or using ``normed=True`` option + + >>> L, d = csgraph.laplacian(G, return_diag=True, normed=True) + >>> L + array([[ 1. , -0.5, -0.5], + [-0.5, 1. , -0.5], + [-0.5, -0.5, 1. ]]) + + which now instead of the diagonal returns the scaling coefficients + + >>> d + array([1.41421356, 1.41421356, 1.41421356]) + + Zero scaling coefficients are substituted with 1s, where scaling + has thus no effect, e.g., + + >>> G = np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0]]) + >>> G + array([[0, 0, 0], + [0, 0, 1], + [0, 1, 0]]) + >>> L, d = csgraph.laplacian(G, return_diag=True, normed=True) + >>> L + array([[ 0., -0., -0.], + [-0., 1., -1.], + [-0., -1., 1.]]) + >>> d + array([1., 1., 1.]) + + Only the symmetric normalization is implemented, resulting + in a symmetric Laplacian matrix if and only if its graph is symmetric + and has all non-negative degrees, like in the examples above. + + The output Laplacian matrix is by default a dense array or a sparse + array or matrix inferring its class, shape, format, and dtype from + the input graph matrix: + + >>> G = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]).astype(np.float32) + >>> G + array([[0., 1., 1.], + [1., 0., 1.], + [1., 1., 0.]], dtype=float32) + >>> csgraph.laplacian(G) + array([[ 2., -1., -1.], + [-1., 2., -1.], + [-1., -1., 2.]], dtype=float32) + + but can alternatively be generated matrix-free as a LinearOperator: + + >>> L = csgraph.laplacian(G, form="lo") + >>> L + <3x3 _CustomLinearOperator with dtype=float32> + >>> L(np.eye(3)) + array([[ 2., -1., -1.], + [-1., 2., -1.], + [-1., -1., 2.]]) + + or as a lambda-function: + + >>> L = csgraph.laplacian(G, form="function") + >>> L + . at 0x0000012AE6F5A598> + >>> L(np.eye(3)) + array([[ 2., -1., -1.], + [-1., 2., -1.], + [-1., -1., 2.]]) + + The Laplacian matrix is used for + spectral data clustering and embedding + as well as for spectral graph partitioning. + Our final example illustrates the latter + for a noisy directed linear graph. + + >>> from scipy.sparse import diags_array, random_array + >>> from scipy.sparse.linalg import lobpcg + + Create a directed linear graph with ``N=35`` vertices + using a sparse adjacency matrix ``G``: + + >>> N = 35 + >>> G = diags_array(np.ones(N - 1), offsets=1, format="csr") + + Fix a random seed ``rng`` and add a random sparse noise to the graph ``G``: + + >>> rng = np.random.default_rng() + >>> G += 1e-2 * random_array((N, N), density=0.1, rng=rng) + + Set initial approximations for eigenvectors: + + >>> X = rng.random((N, 2)) + + The constant vector of ones is always a trivial eigenvector + of the non-normalized Laplacian to be filtered out: + + >>> Y = np.ones((N, 1)) + + Alternating (1) the sign of the graph weights allows determining + labels for spectral max- and min- cuts in a single loop. + Since the graph is undirected, the option ``symmetrized=True`` + must be used in the construction of the Laplacian. + The option ``normed=True`` cannot be used in (2) for the negative weights + here as the symmetric normalization evaluates square roots. + The option ``form="lo"`` in (2) is matrix-free, i.e., guarantees + a fixed memory footprint and read-only access to the graph. + Calling the eigenvalue solver ``lobpcg`` (3) computes the Fiedler vector + that determines the labels as the signs of its components in (5). + Since the sign in an eigenvector is not deterministic and can flip, + we fix the sign of the first component to be always +1 in (4). + + >>> for cut in ["max", "min"]: + ... G = -G # 1. + ... L = csgraph.laplacian(G, symmetrized=True, form="lo") # 2. + ... _, eves = lobpcg(L, X, Y=Y, largest=False, tol=1e-2) # 3. + ... eves *= np.sign(eves[0, 0]) # 4. + ... print(cut + "-cut labels:\\n", 1 * (eves[:, 0]>0)) # 5. + max-cut labels: + [1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1] + min-cut labels: + [1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0] + + As anticipated for a (slightly noisy) linear graph, + the max-cut strips all the edges of the graph coloring all + odd vertices into one color and all even vertices into another one, + while the balanced min-cut partitions the graph + in the middle by deleting a single edge. + Both determined partitions are optimal. + """ + is_pydata_sparse = is_pydata_spmatrix(csgraph) + if is_pydata_sparse: + pydata_sparse_cls = csgraph.__class__ + csgraph = convert_pydata_sparse_to_scipy(csgraph) + if csgraph.ndim != 2 or csgraph.shape[0] != csgraph.shape[1]: + raise ValueError('csgraph must be a square matrix or array') + + if normed and ( + np.issubdtype(csgraph.dtype, np.signedinteger) + or np.issubdtype(csgraph.dtype, np.uint) + ): + csgraph = csgraph.astype(np.float64) + + if form == "array": + create_lap = ( + _laplacian_sparse if issparse(csgraph) else _laplacian_dense + ) + else: + create_lap = ( + _laplacian_sparse_flo + if issparse(csgraph) + else _laplacian_dense_flo + ) + + degree_axis = 1 if use_out_degree else 0 + + lap, d = create_lap( + csgraph, + normed=normed, + axis=degree_axis, + copy=copy, + form=form, + dtype=dtype, + symmetrized=symmetrized, + ) + if is_pydata_sparse: + lap = pydata_sparse_cls.from_scipy_sparse(lap) + if return_diag: + return lap, d + return lap + + +def _setdiag_dense(m, d): + step = len(d) + 1 + m.flat[::step] = d + + +def _laplace(m, d): + return lambda v: v * d[:, np.newaxis] - m @ v + + +def _laplace_normed(m, d, nd): + laplace = _laplace(m, d) + return lambda v: nd[:, np.newaxis] * laplace(v * nd[:, np.newaxis]) + + +def _laplace_sym(m, d): + return ( + lambda v: v * d[:, np.newaxis] + - m @ v + - np.transpose(np.conjugate(np.transpose(np.conjugate(v)) @ m)) + ) + + +def _laplace_normed_sym(m, d, nd): + laplace_sym = _laplace_sym(m, d) + return lambda v: nd[:, np.newaxis] * laplace_sym(v * nd[:, np.newaxis]) + + +def _linearoperator(mv, shape, dtype): + return LinearOperator(matvec=mv, matmat=mv, shape=shape, dtype=dtype) + + +def _laplacian_sparse_flo(graph, normed, axis, copy, form, dtype, symmetrized): + # The keyword argument `copy` is unused and has no effect here. + del copy + + if dtype is None: + dtype = graph.dtype + + graph_sum = np.asarray(graph.sum(axis=axis)).ravel() + graph_diagonal = graph.diagonal() + diag = graph_sum - graph_diagonal + if symmetrized: + graph_sum += np.asarray(graph.sum(axis=1 - axis)).ravel() + diag = graph_sum - graph_diagonal - graph_diagonal + + if normed: + isolated_node_mask = diag == 0 + w = np.where(isolated_node_mask, 1, np.sqrt(diag)) + if symmetrized: + md = _laplace_normed_sym(graph, graph_sum, 1.0 / w) + else: + md = _laplace_normed(graph, graph_sum, 1.0 / w) + if form == "function": + return md, w.astype(dtype, copy=False) + elif form == "lo": + m = _linearoperator(md, shape=graph.shape, dtype=dtype) + return m, w.astype(dtype, copy=False) + else: + raise ValueError(f"Invalid form: {form!r}") + else: + if symmetrized: + md = _laplace_sym(graph, graph_sum) + else: + md = _laplace(graph, graph_sum) + if form == "function": + return md, diag.astype(dtype, copy=False) + elif form == "lo": + m = _linearoperator(md, shape=graph.shape, dtype=dtype) + return m, diag.astype(dtype, copy=False) + else: + raise ValueError(f"Invalid form: {form!r}") + + +def _laplacian_sparse(graph, normed, axis, copy, form, dtype, symmetrized): + # The keyword argument `form` is unused and has no effect here. + del form + + if dtype is None: + dtype = graph.dtype + + needs_copy = False + if graph.format in ('lil', 'dok'): + m = graph.tocoo() + else: + m = graph + if copy: + needs_copy = True + + if symmetrized: + m += m.T.conj() + + w = np.asarray(m.sum(axis=axis)).ravel() - m.diagonal() + if normed: + m = m.tocoo(copy=needs_copy) + isolated_node_mask = (w == 0) + w = np.where(isolated_node_mask, 1, np.sqrt(w)) + m.data /= w[m.row] + m.data /= w[m.col] + m.data *= -1 + m.setdiag(1 - isolated_node_mask) + else: + if m.format == 'dia': + m = m.copy() + else: + m = m.tocoo(copy=needs_copy) + m.data *= -1 + m.setdiag(w) + + return m.astype(dtype, copy=False), w.astype(dtype) + + +def _laplacian_dense_flo(graph, normed, axis, copy, form, dtype, symmetrized): + + if copy: + m = np.array(graph) + else: + m = np.asarray(graph) + + if dtype is None: + dtype = m.dtype + + graph_sum = m.sum(axis=axis) + graph_diagonal = m.diagonal() + diag = graph_sum - graph_diagonal + if symmetrized: + graph_sum += m.sum(axis=1 - axis) + diag = graph_sum - graph_diagonal - graph_diagonal + + if normed: + isolated_node_mask = diag == 0 + w = np.where(isolated_node_mask, 1, np.sqrt(diag)) + if symmetrized: + md = _laplace_normed_sym(m, graph_sum, 1.0 / w) + else: + md = _laplace_normed(m, graph_sum, 1.0 / w) + if form == "function": + return md, w.astype(dtype, copy=False) + elif form == "lo": + m = _linearoperator(md, shape=graph.shape, dtype=dtype) + return m, w.astype(dtype, copy=False) + else: + raise ValueError(f"Invalid form: {form!r}") + else: + if symmetrized: + md = _laplace_sym(m, graph_sum) + else: + md = _laplace(m, graph_sum) + if form == "function": + return md, diag.astype(dtype, copy=False) + elif form == "lo": + m = _linearoperator(md, shape=graph.shape, dtype=dtype) + return m, diag.astype(dtype, copy=False) + else: + raise ValueError(f"Invalid form: {form!r}") + + +def _laplacian_dense(graph, normed, axis, copy, form, dtype, symmetrized): + + if form != "array": + raise ValueError(f'{form!r} must be "array"') + + if dtype is None: + dtype = graph.dtype + + if copy: + m = np.array(graph) + else: + m = np.asarray(graph) + + if dtype is None: + dtype = m.dtype + + if symmetrized: + m += m.T.conj() + np.fill_diagonal(m, 0) + w = m.sum(axis=axis) + if normed: + isolated_node_mask = (w == 0) + w = np.where(isolated_node_mask, 1, np.sqrt(w)) + m /= w + m /= w[:, np.newaxis] + m *= -1 + _setdiag_dense(m, 1 - isolated_node_mask) + else: + m *= -1 + _setdiag_dense(m, w) + + return m.astype(dtype, copy=False), w.astype(dtype, copy=False) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_matching.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_matching.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..8633e3cf2b05f4401879c3824f20d5529708bcc1 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_matching.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_min_spanning_tree.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_min_spanning_tree.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..8b375f34645045975ec20124b34ad26548ae7d25 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_min_spanning_tree.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_reordering.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_reordering.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..80b606440938d6bf8fbfb4ea6990326acedf19c6 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_reordering.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_shortest_path.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_shortest_path.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..8e1ac2cc43b2969f10951711a9f843b616a9879b Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_shortest_path.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_tools.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_tools.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..d4a5d8b05777d19f22214e23672f11b98a885e25 Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_tools.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_traversal.cp310-win_amd64.dll.a b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_traversal.cp310-win_amd64.dll.a new file mode 100644 index 0000000000000000000000000000000000000000..4ab436636fc839fcd08ea96ea11710c376babaaa Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_traversal.cp310-win_amd64.dll.a differ diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_validation.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..002d8ff391edd5f8167b439317352d6a47637851 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/_validation.py @@ -0,0 +1,66 @@ +import numpy as np +from scipy.sparse import issparse +from scipy.sparse._sputils import convert_pydata_sparse_to_scipy +from scipy.sparse.csgraph._tools import ( + csgraph_to_dense, csgraph_from_dense, + csgraph_masked_from_dense, csgraph_from_masked +) + +DTYPE = np.float64 + + +def validate_graph(csgraph, directed, dtype=DTYPE, + csr_output=True, dense_output=True, + copy_if_dense=False, copy_if_sparse=False, + null_value_in=0, null_value_out=np.inf, + infinity_null=True, nan_null=True): + """Routine for validation and conversion of csgraph inputs""" + if not (csr_output or dense_output): + raise ValueError("Internal: dense or csr output must be true") + + accept_fv = [null_value_in] + if infinity_null: + accept_fv.append(np.inf) + if nan_null: + accept_fv.append(np.nan) + csgraph = convert_pydata_sparse_to_scipy(csgraph, accept_fv=accept_fv) + + # if undirected and csc storage, then transposing in-place + # is quicker than later converting to csr. + if (not directed) and issparse(csgraph) and csgraph.format == "csc": + csgraph = csgraph.T + + if issparse(csgraph): + if csr_output: + csgraph = csgraph.tocsr(copy=copy_if_sparse).astype(DTYPE, copy=False) + else: + csgraph = csgraph_to_dense(csgraph, null_value=null_value_out) + elif np.ma.isMaskedArray(csgraph): + if dense_output: + mask = csgraph.mask + csgraph = np.array(csgraph.data, dtype=DTYPE, copy=copy_if_dense) + csgraph[mask] = null_value_out + else: + csgraph = csgraph_from_masked(csgraph) + else: + if dense_output: + csgraph = csgraph_masked_from_dense(csgraph, + copy=copy_if_dense, + null_value=null_value_in, + nan_null=nan_null, + infinity_null=infinity_null) + mask = csgraph.mask + csgraph = np.asarray(csgraph.data, dtype=DTYPE) + csgraph[mask] = null_value_out + else: + csgraph = csgraph_from_dense(csgraph, null_value=null_value_in, + infinity_null=infinity_null, + nan_null=nan_null) + + if csgraph.ndim != 2: + raise ValueError("compressed-sparse graph must be 2-D") + + if csgraph.shape[0] != csgraph.shape[1]: + raise ValueError("compressed-sparse graph must be shape (N, N)") + + return csgraph diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_connected_components.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_connected_components.py new file mode 100644 index 0000000000000000000000000000000000000000..6ed7e321a2c5ca8539e19c6fe72447af5c524223 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_connected_components.py @@ -0,0 +1,119 @@ +import numpy as np +from numpy.testing import assert_equal, assert_array_almost_equal +from scipy.sparse import csgraph, csr_array + + +def test_weak_connections(): + Xde = np.array([[0, 1, 0], + [0, 0, 0], + [0, 0, 0]]) + + Xsp = csgraph.csgraph_from_dense(Xde, null_value=0) + + for X in Xsp, Xde: + n_components, labels =\ + csgraph.connected_components(X, directed=True, + connection='weak') + + assert_equal(n_components, 2) + assert_array_almost_equal(labels, [0, 0, 1]) + + +def test_strong_connections(): + X1de = np.array([[0, 1, 0], + [0, 0, 0], + [0, 0, 0]]) + X2de = X1de + X1de.T + + X1sp = csgraph.csgraph_from_dense(X1de, null_value=0) + X2sp = csgraph.csgraph_from_dense(X2de, null_value=0) + + for X in X1sp, X1de: + n_components, labels =\ + csgraph.connected_components(X, directed=True, + connection='strong') + + assert_equal(n_components, 3) + labels.sort() + assert_array_almost_equal(labels, [0, 1, 2]) + + for X in X2sp, X2de: + n_components, labels =\ + csgraph.connected_components(X, directed=True, + connection='strong') + + assert_equal(n_components, 2) + labels.sort() + assert_array_almost_equal(labels, [0, 0, 1]) + + +def test_strong_connections2(): + X = np.array([[0, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 0, 0], + [0, 0, 1, 0, 1, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 0]]) + n_components, labels =\ + csgraph.connected_components(X, directed=True, + connection='strong') + assert_equal(n_components, 5) + labels.sort() + assert_array_almost_equal(labels, [0, 1, 2, 2, 3, 4]) + + +def test_weak_connections2(): + X = np.array([[0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 0, 0], + [0, 0, 1, 0, 1, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 0]]) + n_components, labels =\ + csgraph.connected_components(X, directed=True, + connection='weak') + assert_equal(n_components, 2) + labels.sort() + assert_array_almost_equal(labels, [0, 0, 1, 1, 1, 1]) + + +def test_ticket1876(): + # Regression test: this failed in the original implementation + # There should be two strongly-connected components; previously gave one + g = np.array([[0, 1, 1, 0], + [1, 0, 0, 1], + [0, 0, 0, 1], + [0, 0, 1, 0]]) + n_components, labels = csgraph.connected_components(g, connection='strong') + + assert_equal(n_components, 2) + assert_equal(labels[0], labels[1]) + assert_equal(labels[2], labels[3]) + + +def test_fully_connected_graph(): + # Fully connected dense matrices raised an exception. + # https://github.com/scipy/scipy/issues/3818 + g = np.ones((4, 4)) + n_components, labels = csgraph.connected_components(g) + assert_equal(n_components, 1) + + +def test_int64_indices_undirected(): + # See https://github.com/scipy/scipy/issues/18716 + g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2)) + assert g.indices.dtype == np.int64 + n, labels = csgraph.connected_components(g, directed=False) + assert n == 1 + assert_array_almost_equal(labels, [0, 0]) + + +def test_int64_indices_directed(): + # See https://github.com/scipy/scipy/issues/18716 + g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2)) + assert g.indices.dtype == np.int64 + n, labels = csgraph.connected_components(g, directed=True, + connection='strong') + assert n == 2 + assert_array_almost_equal(labels, [1, 0]) + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_conversions.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_conversions.py new file mode 100644 index 0000000000000000000000000000000000000000..6abf2177429ea7e1323d87dde87542f5662fc20b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_conversions.py @@ -0,0 +1,61 @@ +import numpy as np +from numpy.testing import assert_array_almost_equal +from scipy.sparse import csr_array +from scipy.sparse.csgraph import csgraph_from_dense, csgraph_to_dense + + +def test_csgraph_from_dense(): + np.random.seed(1234) + G = np.random.random((10, 10)) + some_nulls = (G < 0.4) + all_nulls = (G < 0.8) + + for null_value in [0, np.nan, np.inf]: + G[all_nulls] = null_value + with np.errstate(invalid="ignore"): + G_csr = csgraph_from_dense(G, null_value=0) + + G[all_nulls] = 0 + assert_array_almost_equal(G, G_csr.toarray()) + + for null_value in [np.nan, np.inf]: + G[all_nulls] = 0 + G[some_nulls] = null_value + with np.errstate(invalid="ignore"): + G_csr = csgraph_from_dense(G, null_value=0) + + G[all_nulls] = 0 + assert_array_almost_equal(G, G_csr.toarray()) + + +def test_csgraph_to_dense(): + np.random.seed(1234) + G = np.random.random((10, 10)) + nulls = (G < 0.8) + G[nulls] = np.inf + + G_csr = csgraph_from_dense(G) + + for null_value in [0, 10, -np.inf, np.inf]: + G[nulls] = null_value + assert_array_almost_equal(G, csgraph_to_dense(G_csr, null_value)) + + +def test_multiple_edges(): + # create a random square matrix with an even number of elements + np.random.seed(1234) + X = np.random.random((10, 10)) + Xcsr = csr_array(X) + + # now double-up every other column + Xcsr.indices[::2] = Xcsr.indices[1::2] + + # normal sparse toarray() will sum the duplicated edges + Xdense = Xcsr.toarray() + assert_array_almost_equal(Xdense[:, 1::2], + X[:, ::2] + X[:, 1::2]) + + # csgraph_to_dense chooses the minimum of each duplicated edge + Xdense = csgraph_to_dense(Xcsr) + assert_array_almost_equal(Xdense[:, 1::2], + np.minimum(X[:, ::2], X[:, 1::2])) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_flow.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..587ffe22b2296b9c93e5ed299e9455634f2d03af --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_flow.py @@ -0,0 +1,209 @@ +import numpy as np +from numpy.testing import assert_array_equal +import pytest + +from scipy.sparse import csr_array, csc_array, csr_matrix +from scipy.sparse.csgraph import maximum_flow +from scipy.sparse.csgraph._flow import ( + _add_reverse_edges, _make_edge_pointers, _make_tails +) + +methods = ['edmonds_karp', 'dinic'] + +def test_raises_on_dense_input(): + with pytest.raises(TypeError): + graph = np.array([[0, 1], [0, 0]]) + maximum_flow(graph, 0, 1) + maximum_flow(graph, 0, 1, method='edmonds_karp') + + +def test_raises_on_csc_input(): + with pytest.raises(TypeError): + graph = csc_array([[0, 1], [0, 0]]) + maximum_flow(graph, 0, 1) + maximum_flow(graph, 0, 1, method='edmonds_karp') + + +def test_raises_on_floating_point_input(): + with pytest.raises(ValueError): + graph = csr_array([[0, 1.5], [0, 0]], dtype=np.float64) + maximum_flow(graph, 0, 1) + maximum_flow(graph, 0, 1, method='edmonds_karp') + + +def test_raises_on_non_square_input(): + with pytest.raises(ValueError): + graph = csr_array([[0, 1, 2], [2, 1, 0]]) + maximum_flow(graph, 0, 1) + + +def test_raises_when_source_is_sink(): + with pytest.raises(ValueError): + graph = csr_array([[0, 1], [0, 0]]) + maximum_flow(graph, 0, 0) + maximum_flow(graph, 0, 0, method='edmonds_karp') + + +@pytest.mark.parametrize('method', methods) +@pytest.mark.parametrize('source', [-1, 2, 3]) +def test_raises_when_source_is_out_of_bounds(source, method): + with pytest.raises(ValueError): + graph = csr_array([[0, 1], [0, 0]]) + maximum_flow(graph, source, 1, method=method) + + +@pytest.mark.parametrize('method', methods) +@pytest.mark.parametrize('sink', [-1, 2, 3]) +def test_raises_when_sink_is_out_of_bounds(sink, method): + with pytest.raises(ValueError): + graph = csr_array([[0, 1], [0, 0]]) + maximum_flow(graph, 0, sink, method=method) + + +@pytest.mark.parametrize('method', methods) +def test_simple_graph(method): + # This graph looks as follows: + # (0) --5--> (1) + graph = csr_array([[0, 5], [0, 0]]) + res = maximum_flow(graph, 0, 1, method=method) + assert res.flow_value == 5 + expected_flow = np.array([[0, 5], [-5, 0]]) + assert_array_equal(res.flow.toarray(), expected_flow) + + +@pytest.mark.parametrize('method', methods) +def test_return_type(method): + graph = csr_array([[0, 5], [0, 0]]) + assert isinstance(maximum_flow(graph, 0, 1, method=method).flow, csr_array) + graph = csr_matrix([[0, 5], [0, 0]]) + assert isinstance(maximum_flow(graph, 0, 1, method=method).flow, csr_matrix) + + +@pytest.mark.parametrize('method', methods) +def test_bottle_neck_graph(method): + # This graph cannot use the full capacity between 0 and 1: + # (0) --5--> (1) --3--> (2) + graph = csr_array([[0, 5, 0], [0, 0, 3], [0, 0, 0]]) + res = maximum_flow(graph, 0, 2, method=method) + assert res.flow_value == 3 + expected_flow = np.array([[0, 3, 0], [-3, 0, 3], [0, -3, 0]]) + assert_array_equal(res.flow.toarray(), expected_flow) + + +@pytest.mark.parametrize('method', methods) +def test_backwards_flow(method): + # This example causes backwards flow between vertices 3 and 4, + # and so this test ensures that we handle that accordingly. See + # https://stackoverflow.com/q/38843963/5085211 + # for more information. + graph = csr_array([[0, 10, 0, 0, 10, 0, 0, 0], + [0, 0, 10, 0, 0, 0, 0, 0], + [0, 0, 0, 10, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 10], + [0, 0, 0, 10, 0, 10, 0, 0], + [0, 0, 0, 0, 0, 0, 10, 0], + [0, 0, 0, 0, 0, 0, 0, 10], + [0, 0, 0, 0, 0, 0, 0, 0]]) + res = maximum_flow(graph, 0, 7, method=method) + assert res.flow_value == 20 + expected_flow = np.array([[0, 10, 0, 0, 10, 0, 0, 0], + [-10, 0, 10, 0, 0, 0, 0, 0], + [0, -10, 0, 10, 0, 0, 0, 0], + [0, 0, -10, 0, 0, 0, 0, 10], + [-10, 0, 0, 0, 0, 10, 0, 0], + [0, 0, 0, 0, -10, 0, 10, 0], + [0, 0, 0, 0, 0, -10, 0, 10], + [0, 0, 0, -10, 0, 0, -10, 0]]) + assert_array_equal(res.flow.toarray(), expected_flow) + + +@pytest.mark.parametrize('method', methods) +def test_example_from_clrs_chapter_26_1(method): + # See page 659 in CLRS second edition, but note that the maximum flow + # we find is slightly different than the one in CLRS; we push a flow of + # 12 to v_1 instead of v_2. + graph = csr_array([[0, 16, 13, 0, 0, 0], + [0, 0, 10, 12, 0, 0], + [0, 4, 0, 0, 14, 0], + [0, 0, 9, 0, 0, 20], + [0, 0, 0, 7, 0, 4], + [0, 0, 0, 0, 0, 0]]) + res = maximum_flow(graph, 0, 5, method=method) + assert res.flow_value == 23 + expected_flow = np.array([[0, 12, 11, 0, 0, 0], + [-12, 0, 0, 12, 0, 0], + [-11, 0, 0, 0, 11, 0], + [0, -12, 0, 0, -7, 19], + [0, 0, -11, 7, 0, 4], + [0, 0, 0, -19, -4, 0]]) + assert_array_equal(res.flow.toarray(), expected_flow) + + +@pytest.mark.parametrize('method', methods) +def test_disconnected_graph(method): + # This tests the following disconnected graph: + # (0) --5--> (1) (2) --3--> (3) + graph = csr_array([[0, 5, 0, 0], + [0, 0, 0, 0], + [0, 0, 9, 3], + [0, 0, 0, 0]]) + res = maximum_flow(graph, 0, 3, method=method) + assert res.flow_value == 0 + expected_flow = np.zeros((4, 4), dtype=np.int32) + assert_array_equal(res.flow.toarray(), expected_flow) + + +@pytest.mark.parametrize('method', methods) +def test_add_reverse_edges_large_graph(method): + # Regression test for https://github.com/scipy/scipy/issues/14385 + n = 100_000 + indices = np.arange(1, n) + indptr = np.array(list(range(n)) + [n - 1]) + data = np.ones(n - 1, dtype=np.int32) + graph = csr_array((data, indices, indptr), shape=(n, n)) + res = maximum_flow(graph, 0, n - 1, method=method) + assert res.flow_value == 1 + expected_flow = graph - graph.transpose() + assert_array_equal(res.flow.data, expected_flow.data) + assert_array_equal(res.flow.indices, expected_flow.indices) + assert_array_equal(res.flow.indptr, expected_flow.indptr) + + +@pytest.mark.parametrize("a,b_data_expected", [ + ([[]], []), + ([[0], [0]], []), + ([[1, 0, 2], [0, 0, 0], [0, 3, 0]], [1, 2, 0, 0, 3]), + ([[9, 8, 7], [4, 5, 6], [0, 0, 0]], [9, 8, 7, 4, 5, 6, 0, 0])]) +def test_add_reverse_edges(a, b_data_expected): + """Test that the reversal of the edges of the input graph works + as expected. + """ + a = csr_array(a, dtype=np.int32, shape=(len(a), len(a))) + b = _add_reverse_edges(a) + assert_array_equal(b.data, b_data_expected) + + +@pytest.mark.parametrize("a,expected", [ + ([[]], []), + ([[0]], []), + ([[1]], [0]), + ([[0, 1], [10, 0]], [1, 0]), + ([[1, 0, 2], [0, 0, 3], [4, 5, 0]], [0, 3, 4, 1, 2]) +]) +def test_make_edge_pointers(a, expected): + a = csr_array(a, dtype=np.int32) + rev_edge_ptr = _make_edge_pointers(a) + assert_array_equal(rev_edge_ptr, expected) + + +@pytest.mark.parametrize("a,expected", [ + ([[]], []), + ([[0]], []), + ([[1]], [0]), + ([[0, 1], [10, 0]], [0, 1]), + ([[1, 0, 2], [0, 0, 3], [4, 5, 0]], [0, 0, 1, 2, 2]) +]) +def test_make_tails(a, expected): + a = csr_array(a, dtype=np.int32) + tails = _make_tails(a) + assert_array_equal(tails, expected) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_graph_laplacian.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_graph_laplacian.py new file mode 100644 index 0000000000000000000000000000000000000000..220c0a2a4d23847c8cd158a25045dc35c0396297 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_graph_laplacian.py @@ -0,0 +1,368 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose +from pytest import raises as assert_raises +from scipy import sparse + +from scipy.sparse import csgraph +from scipy._lib._util import np_long, np_ulong + + +def check_int_type(mat): + return np.issubdtype(mat.dtype, np.signedinteger) or np.issubdtype( + mat.dtype, np_ulong + ) + + +def test_laplacian_value_error(): + for t in int, float, complex: + for m in ([1, 1], + [[[1]]], + [[1, 2, 3], [4, 5, 6]], + [[1, 2], [3, 4], [5, 5]]): + A = np.array(m, dtype=t) + assert_raises(ValueError, csgraph.laplacian, A) + + +def _explicit_laplacian(x, normed=False): + if sparse.issparse(x): + x = x.toarray() + x = np.asarray(x) + y = -1.0 * x + for j in range(y.shape[0]): + y[j,j] = x[j,j+1:].sum() + x[j,:j].sum() + if normed: + d = np.diag(y).copy() + d[d == 0] = 1.0 + y /= d[:,None]**.5 + y /= d[None,:]**.5 + return y + + +def _check_symmetric_graph_laplacian(mat, normed, copy=True): + if not hasattr(mat, 'shape'): + mat = eval(mat, dict(np=np, sparse=sparse)) + + if sparse.issparse(mat): + sp_mat = mat + mat = sp_mat.toarray() + else: + sp_mat = sparse.csr_array(mat) + + mat_copy = np.copy(mat) + sp_mat_copy = sparse.csr_array(sp_mat, copy=True) + + n_nodes = mat.shape[0] + explicit_laplacian = _explicit_laplacian(mat, normed=normed) + laplacian = csgraph.laplacian(mat, normed=normed, copy=copy) + sp_laplacian = csgraph.laplacian(sp_mat, normed=normed, + copy=copy) + + if copy: + assert_allclose(mat, mat_copy) + _assert_allclose_sparse(sp_mat, sp_mat_copy) + else: + if not (normed and check_int_type(mat)): + assert_allclose(laplacian, mat) + if sp_mat.format == 'coo': + _assert_allclose_sparse(sp_laplacian, sp_mat) + + assert_allclose(laplacian, sp_laplacian.toarray()) + + for tested in [laplacian, sp_laplacian.toarray()]: + if not normed: + assert_allclose(tested.sum(axis=0), np.zeros(n_nodes)) + assert_allclose(tested.T, tested) + assert_allclose(tested, explicit_laplacian) + + +def test_symmetric_graph_laplacian(): + symmetric_mats = ( + 'np.arange(10) * np.arange(10)[:, np.newaxis]', + 'np.ones((7, 7))', + 'np.eye(19)', + 'sparse.diags([1, 1], [-1, 1], shape=(4, 4))', + 'sparse.diags([1, 1], [-1, 1], shape=(4, 4)).toarray()', + 'sparse.diags([1, 1], [-1, 1], shape=(4, 4)).todense()', + 'np.vander(np.arange(4)) + np.vander(np.arange(4)).T' + ) + for mat in symmetric_mats: + for normed in True, False: + for copy in True, False: + _check_symmetric_graph_laplacian(mat, normed, copy) + + +def _assert_allclose_sparse(a, b, **kwargs): + # helper function that can deal with sparse matrices + if sparse.issparse(a): + a = a.toarray() + if sparse.issparse(b): + b = b.toarray() + assert_allclose(a, b, **kwargs) + + +def _check_laplacian_dtype_none( + A, desired_L, desired_d, normed, use_out_degree, copy, dtype, arr_type +): + mat = arr_type(A, dtype=dtype) + L, d = csgraph.laplacian( + mat, + normed=normed, + return_diag=True, + use_out_degree=use_out_degree, + copy=copy, + dtype=None, + ) + if normed and check_int_type(mat): + assert L.dtype == np.float64 + assert d.dtype == np.float64 + _assert_allclose_sparse(L, desired_L, atol=1e-12) + _assert_allclose_sparse(d, desired_d, atol=1e-12) + else: + assert L.dtype == dtype + assert d.dtype == dtype + desired_L = np.asarray(desired_L).astype(dtype) + desired_d = np.asarray(desired_d).astype(dtype) + _assert_allclose_sparse(L, desired_L, atol=1e-12) + _assert_allclose_sparse(d, desired_d, atol=1e-12) + + if not copy: + if not (normed and check_int_type(mat)): + if type(mat) is np.ndarray: + assert_allclose(L, mat) + elif mat.format == "coo": + _assert_allclose_sparse(L, mat) + + +def _check_laplacian_dtype( + A, desired_L, desired_d, normed, use_out_degree, copy, dtype, arr_type +): + mat = arr_type(A, dtype=dtype) + L, d = csgraph.laplacian( + mat, + normed=normed, + return_diag=True, + use_out_degree=use_out_degree, + copy=copy, + dtype=dtype, + ) + assert L.dtype == dtype + assert d.dtype == dtype + desired_L = np.asarray(desired_L).astype(dtype) + desired_d = np.asarray(desired_d).astype(dtype) + _assert_allclose_sparse(L, desired_L, atol=1e-12) + _assert_allclose_sparse(d, desired_d, atol=1e-12) + + if not copy: + if not (normed and check_int_type(mat)): + if type(mat) is np.ndarray: + assert_allclose(L, mat) + elif mat.format == 'coo': + _assert_allclose_sparse(L, mat) + + +INT_DTYPES = (np.intc, np_long, np.longlong) +REAL_DTYPES = (np.float32, np.float64, np.longdouble) +COMPLEX_DTYPES = (np.complex64, np.complex128, np.clongdouble) +DTYPES = INT_DTYPES + REAL_DTYPES + COMPLEX_DTYPES + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("arr_type", [np.array, + sparse.csr_matrix, + sparse.coo_matrix, + sparse.csr_array, + sparse.coo_array]) +@pytest.mark.parametrize("copy", [True, False]) +@pytest.mark.parametrize("normed", [True, False]) +@pytest.mark.parametrize("use_out_degree", [True, False]) +def test_asymmetric_laplacian(use_out_degree, normed, + copy, dtype, arr_type): + # adjacency matrix + A = [[0, 1, 0], + [4, 2, 0], + [0, 0, 0]] + A = arr_type(np.array(A), dtype=dtype) + A_copy = A.copy() + + if not normed and use_out_degree: + # Laplacian matrix using out-degree + L = [[1, -1, 0], + [-4, 4, 0], + [0, 0, 0]] + d = [1, 4, 0] + + if normed and use_out_degree: + # normalized Laplacian matrix using out-degree + L = [[1, -0.5, 0], + [-2, 1, 0], + [0, 0, 0]] + d = [1, 2, 1] + + if not normed and not use_out_degree: + # Laplacian matrix using in-degree + L = [[4, -1, 0], + [-4, 1, 0], + [0, 0, 0]] + d = [4, 1, 0] + + if normed and not use_out_degree: + # normalized Laplacian matrix using in-degree + L = [[1, -0.5, 0], + [-2, 1, 0], + [0, 0, 0]] + d = [2, 1, 1] + + _check_laplacian_dtype_none( + A, + L, + d, + normed=normed, + use_out_degree=use_out_degree, + copy=copy, + dtype=dtype, + arr_type=arr_type, + ) + + _check_laplacian_dtype( + A_copy, + L, + d, + normed=normed, + use_out_degree=use_out_degree, + copy=copy, + dtype=dtype, + arr_type=arr_type, + ) + + +@pytest.mark.parametrize("fmt", ['csr', 'csc', 'coo', 'lil', + 'dok', 'dia', 'bsr']) +@pytest.mark.parametrize("normed", [True, False]) +@pytest.mark.parametrize("copy", [True, False]) +def test_sparse_formats(fmt, normed, copy): + mat = sparse.diags_array([1, 1], offsets=[-1, 1], shape=(4, 4), format=fmt) + _check_symmetric_graph_laplacian(mat, normed, copy) + + +@pytest.mark.parametrize( + "arr_type", [np.asarray, + sparse.csr_matrix, + sparse.coo_matrix, + sparse.csr_array, + sparse.coo_array] +) +@pytest.mark.parametrize("form", ["array", "function", "lo"]) +def test_laplacian_symmetrized(arr_type, form): + # adjacency matrix + n = 3 + mat = arr_type(np.arange(n * n).reshape(n, n)) + L_in, d_in = csgraph.laplacian( + mat, + return_diag=True, + form=form, + ) + L_out, d_out = csgraph.laplacian( + mat, + return_diag=True, + use_out_degree=True, + form=form, + ) + Ls, ds = csgraph.laplacian( + mat, + return_diag=True, + symmetrized=True, + form=form, + ) + Ls_normed, ds_normed = csgraph.laplacian( + mat, + return_diag=True, + symmetrized=True, + normed=True, + form=form, + ) + mat += mat.T + Lss, dss = csgraph.laplacian(mat, return_diag=True, form=form) + Lss_normed, dss_normed = csgraph.laplacian( + mat, + return_diag=True, + normed=True, + form=form, + ) + + assert_allclose(ds, d_in + d_out) + assert_allclose(ds, dss) + assert_allclose(ds_normed, dss_normed) + + d = {} + for L in ["L_in", "L_out", "Ls", "Ls_normed", "Lss", "Lss_normed"]: + if form == "array": + d[L] = eval(L) + else: + d[L] = eval(L)(np.eye(n, dtype=mat.dtype)) + + _assert_allclose_sparse(d["Ls"], d["L_in"] + d["L_out"].T) + _assert_allclose_sparse(d["Ls"], d["Lss"]) + _assert_allclose_sparse(d["Ls_normed"], d["Lss_normed"]) + + +@pytest.mark.parametrize( + "arr_type", [np.asarray, + sparse.csr_matrix, + sparse.coo_matrix, + sparse.csr_array, + sparse.coo_array] +) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("normed", [True, False]) +@pytest.mark.parametrize("symmetrized", [True, False]) +@pytest.mark.parametrize("use_out_degree", [True, False]) +@pytest.mark.parametrize("form", ["function", "lo"]) +def test_format(dtype, arr_type, normed, symmetrized, use_out_degree, form): + n = 3 + mat = [[0, 1, 0], [4, 2, 0], [0, 0, 0]] + mat = arr_type(np.array(mat), dtype=dtype) + Lo, do = csgraph.laplacian( + mat, + return_diag=True, + normed=normed, + symmetrized=symmetrized, + use_out_degree=use_out_degree, + dtype=dtype, + ) + La, da = csgraph.laplacian( + mat, + return_diag=True, + normed=normed, + symmetrized=symmetrized, + use_out_degree=use_out_degree, + dtype=dtype, + form="array", + ) + assert_allclose(do, da) + _assert_allclose_sparse(Lo, La) + + L, d = csgraph.laplacian( + mat, + return_diag=True, + normed=normed, + symmetrized=symmetrized, + use_out_degree=use_out_degree, + dtype=dtype, + form=form, + ) + assert_allclose(d, do) + assert d.dtype == dtype + Lm = L(np.eye(n, dtype=mat.dtype)).astype(dtype) + _assert_allclose_sparse(Lm, Lo, rtol=2e-7, atol=2e-7) + x = np.arange(6).reshape(3, 2) + if not (normed and dtype in INT_DTYPES): + assert_allclose(L(x), Lo @ x) + else: + # Normalized Lo is casted to integer, but L() is not + pass + + +def test_format_error_message(): + with pytest.raises(ValueError, match="Invalid form: 'toto'"): + _ = csgraph.laplacian(np.eye(1), form='toto') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_matching.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..49d9b2a81772a850f10fcc1b04d26b5e74f61933 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_matching.py @@ -0,0 +1,295 @@ +from itertools import product + +import numpy as np +from numpy.testing import assert_array_equal, assert_equal +import pytest + +from scipy.sparse import csr_array, diags_array +from scipy.sparse.csgraph import ( + maximum_bipartite_matching, min_weight_full_bipartite_matching +) + + +def test_maximum_bipartite_matching_raises_on_dense_input(): + with pytest.raises(TypeError): + graph = np.array([[0, 1], [0, 0]]) + maximum_bipartite_matching(graph) + + +def test_maximum_bipartite_matching_empty_graph(): + graph = csr_array((0, 0)) + x = maximum_bipartite_matching(graph, perm_type='row') + y = maximum_bipartite_matching(graph, perm_type='column') + expected_matching = np.array([]) + assert_array_equal(expected_matching, x) + assert_array_equal(expected_matching, y) + + +def test_maximum_bipartite_matching_empty_left_partition(): + graph = csr_array((2, 0)) + x = maximum_bipartite_matching(graph, perm_type='row') + y = maximum_bipartite_matching(graph, perm_type='column') + assert_array_equal(np.array([]), x) + assert_array_equal(np.array([-1, -1]), y) + + +def test_maximum_bipartite_matching_empty_right_partition(): + graph = csr_array((0, 3)) + x = maximum_bipartite_matching(graph, perm_type='row') + y = maximum_bipartite_matching(graph, perm_type='column') + assert_array_equal(np.array([-1, -1, -1]), x) + assert_array_equal(np.array([]), y) + + +def test_maximum_bipartite_matching_graph_with_no_edges(): + graph = csr_array((2, 2)) + x = maximum_bipartite_matching(graph, perm_type='row') + y = maximum_bipartite_matching(graph, perm_type='column') + assert_array_equal(np.array([-1, -1]), x) + assert_array_equal(np.array([-1, -1]), y) + + +def test_maximum_bipartite_matching_graph_that_causes_augmentation(): + # In this graph, column 1 is initially assigned to row 1, but it should be + # reassigned to make room for row 2. + graph = csr_array([[1, 1], [1, 0]]) + x = maximum_bipartite_matching(graph, perm_type='column') + y = maximum_bipartite_matching(graph, perm_type='row') + expected_matching = np.array([1, 0]) + assert_array_equal(expected_matching, x) + assert_array_equal(expected_matching, y) + + +def test_maximum_bipartite_matching_graph_with_more_rows_than_columns(): + graph = csr_array([[1, 1], [1, 0], [0, 1]]) + x = maximum_bipartite_matching(graph, perm_type='column') + y = maximum_bipartite_matching(graph, perm_type='row') + assert_array_equal(np.array([0, -1, 1]), x) + assert_array_equal(np.array([0, 2]), y) + + +def test_maximum_bipartite_matching_graph_with_more_columns_than_rows(): + graph = csr_array([[1, 1, 0], [0, 0, 1]]) + x = maximum_bipartite_matching(graph, perm_type='column') + y = maximum_bipartite_matching(graph, perm_type='row') + assert_array_equal(np.array([0, 2]), x) + assert_array_equal(np.array([0, -1, 1]), y) + + +def test_maximum_bipartite_matching_explicit_zeros_count_as_edges(): + data = [0, 0] + indices = [1, 0] + indptr = [0, 1, 2] + graph = csr_array((data, indices, indptr), shape=(2, 2)) + x = maximum_bipartite_matching(graph, perm_type='row') + y = maximum_bipartite_matching(graph, perm_type='column') + expected_matching = np.array([1, 0]) + assert_array_equal(expected_matching, x) + assert_array_equal(expected_matching, y) + + +def test_maximum_bipartite_matching_feasibility_of_result(): + # This is a regression test for GitHub issue #11458 + data = np.ones(50, dtype=int) + indices = [11, 12, 19, 22, 23, 5, 22, 3, 8, 10, 5, 6, 11, 12, 13, 5, 13, + 14, 20, 22, 3, 15, 3, 13, 14, 11, 12, 19, 22, 23, 5, 22, 3, 8, + 10, 5, 6, 11, 12, 13, 5, 13, 14, 20, 22, 3, 15, 3, 13, 14] + indptr = [0, 5, 7, 10, 10, 15, 20, 22, 22, 23, 25, 30, 32, 35, 35, 40, 45, + 47, 47, 48, 50] + graph = csr_array((data, indices, indptr), shape=(20, 25)) + x = maximum_bipartite_matching(graph, perm_type='row') + y = maximum_bipartite_matching(graph, perm_type='column') + assert (x != -1).sum() == 13 + assert (y != -1).sum() == 13 + # Ensure that each element of the matching is in fact an edge in the graph. + for u, v in zip(range(graph.shape[0]), y): + if v != -1: + assert graph[u, v] + for u, v in zip(x, range(graph.shape[1])): + if u != -1: + assert graph[u, v] + + +def test_matching_large_random_graph_with_one_edge_incident_to_each_vertex(): + np.random.seed(42) + A = diags_array(np.ones(25), offsets=0, format='csr') + rand_perm = np.random.permutation(25) + rand_perm2 = np.random.permutation(25) + + Rrow = np.arange(25) + Rcol = rand_perm + Rdata = np.ones(25, dtype=int) + Rmat = csr_array((Rdata, (Rrow, Rcol))) + + Crow = rand_perm2 + Ccol = np.arange(25) + Cdata = np.ones(25, dtype=int) + Cmat = csr_array((Cdata, (Crow, Ccol))) + # Randomly permute identity matrix + B = Rmat @ A @ Cmat + + # Row permute + perm = maximum_bipartite_matching(B, perm_type='row') + Rrow = np.arange(25) + Rcol = perm + Rdata = np.ones(25, dtype=int) + Rmat = csr_array((Rdata, (Rrow, Rcol))) + C1 = Rmat @ B + + # Column permute + perm2 = maximum_bipartite_matching(B, perm_type='column') + Crow = perm2 + Ccol = np.arange(25) + Cdata = np.ones(25, dtype=int) + Cmat = csr_array((Cdata, (Crow, Ccol))) + C2 = B @ Cmat + + # Should get identity matrix back + assert_equal(any(C1.diagonal() == 0), False) + assert_equal(any(C2.diagonal() == 0), False) + + +@pytest.mark.parametrize('num_rows,num_cols', [(0, 0), (2, 0), (0, 3)]) +def test_min_weight_full_matching_trivial_graph(num_rows, num_cols): + biadjacency = csr_array((num_cols, num_rows)) + row_ind, col_ind = min_weight_full_bipartite_matching(biadjacency) + assert len(row_ind) == 0 + assert len(col_ind) == 0 + + +@pytest.mark.parametrize('biadjacency', + [ + [[1, 1, 1], [1, 0, 0], [1, 0, 0]], + [[1, 1, 1], [0, 0, 1], [0, 0, 1]], + [[1, 0, 0, 1], [1, 1, 0, 1], [0, 0, 0, 0]], + [[1, 0, 0], [2, 0, 0]], + [[0, 1, 0], [0, 2, 0]], + [[1, 0], [2, 0], [5, 0]] + ]) +def test_min_weight_full_matching_infeasible_problems(biadjacency): + with pytest.raises(ValueError): + min_weight_full_bipartite_matching(csr_array(biadjacency)) + + +def test_min_weight_full_matching_large_infeasible(): + # Regression test for GitHub issue #17269 + a = np.asarray([ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.001, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.001, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.001, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.001, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.001, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.001, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.001, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.001, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.001], + [0.0, 0.11687445, 0.0, 0.0, 0.01319788, 0.07509257, 0.0, + 0.0, 0.0, 0.74228317, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.81087935, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.8408466, 0.0, 0.0, 0.0, 0.0, 0.01194389, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.82994211, 0.0, 0.0, 0.0, 0.11468516, 0.0, 0.0, 0.0, + 0.11173505, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0], + [0.18796507, 0.0, 0.04002318, 0.0, 0.0, 0.0, 0.0, 0.0, 0.75883335, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.71545464, 0.0, 0.0, 0.0, 0.0, 0.0, 0.02748488, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.78470564, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.14829198, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.10870609, 0.0, 0.0, 0.0, 0.8918677, 0.0, 0.0, 0.0, 0.06306644, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.63844085, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.7442354, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.09850549, 0.0, 0.0, 0.18638258, + 0.2769244, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.73182464, 0.0, 0.0, 0.46443561, + 0.38589284, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.29510278, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.09666032, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + ]) + with pytest.raises(ValueError, match='no full matching exists'): + min_weight_full_bipartite_matching(csr_array(a)) + + +@pytest.mark.thread_unsafe +def test_explicit_zero_causes_warning(): + with pytest.warns(UserWarning): + biadjacency = csr_array(((2, 0, 3), (0, 1, 1), (0, 2, 3))) + min_weight_full_bipartite_matching(biadjacency) + + +# General test for linear sum assignment solvers to make it possible to rely +# on the same tests for scipy.optimize.linear_sum_assignment. +def linear_sum_assignment_assertions( + solver, array_type, sign, test_case +): + cost_matrix, expected_cost = test_case + maximize = sign == -1 + cost_matrix = sign * array_type(cost_matrix) + expected_cost = sign * np.array(expected_cost) + + row_ind, col_ind = solver(cost_matrix, maximize=maximize) + assert_array_equal(row_ind, np.sort(row_ind)) + assert_array_equal(expected_cost, + np.array(cost_matrix[row_ind, col_ind]).flatten()) + + cost_matrix = cost_matrix.T + row_ind, col_ind = solver(cost_matrix, maximize=maximize) + assert_array_equal(row_ind, np.sort(row_ind)) + assert_array_equal(np.sort(expected_cost), + np.sort(np.array( + cost_matrix[row_ind, col_ind])).flatten()) + + +linear_sum_assignment_test_cases = product( + [-1, 1], + [ + # Square + ([[400, 150, 400], + [400, 450, 600], + [300, 225, 300]], + [150, 400, 300]), + + # Rectangular variant + ([[400, 150, 400, 1], + [400, 450, 600, 2], + [300, 225, 300, 3]], + [150, 2, 300]), + + ([[10, 10, 8], + [9, 8, 1], + [9, 7, 4]], + [10, 1, 7]), + + # Square + ([[10, 10, 8, 11], + [9, 8, 1, 1], + [9, 7, 4, 10]], + [10, 1, 4]), + + # Rectangular variant + ([[10, float("inf"), float("inf")], + [float("inf"), float("inf"), 1], + [float("inf"), 7, float("inf")]], + [10, 1, 7]) + ]) + + +@pytest.mark.parametrize('sign,test_case', linear_sum_assignment_test_cases) +def test_min_weight_full_matching_small_inputs(sign, test_case): + linear_sum_assignment_assertions( + min_weight_full_bipartite_matching, csr_array, sign, test_case) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_pydata_sparse.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_pydata_sparse.py new file mode 100644 index 0000000000000000000000000000000000000000..e1b24cbecff8a48102bd0b8f1a18e23c2c72128f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_pydata_sparse.py @@ -0,0 +1,194 @@ +import pytest + +import numpy as np +import scipy.sparse as sp +import scipy.sparse.csgraph as spgraph +from scipy._lib import _pep440 + +from numpy.testing import assert_equal + +try: + import sparse +except Exception: + sparse = None + +pytestmark = pytest.mark.skipif(sparse is None, + reason="pydata/sparse not installed") + + +msg = "pydata/sparse (0.15.1) does not implement necessary operations" + + +sparse_params = (pytest.param("COO"), + pytest.param("DOK", marks=[pytest.mark.xfail(reason=msg)])) + + +def check_sparse_version(min_ver): + if sparse is None: + return pytest.mark.skip(reason="sparse is not installed") + return pytest.mark.skipif( + _pep440.parse(sparse.__version__) < _pep440.Version(min_ver), + reason=f"sparse version >= {min_ver} required" + ) + + +@pytest.fixture(params=sparse_params) +def sparse_cls(request): + return getattr(sparse, request.param) + + +@pytest.fixture +def graphs(sparse_cls): + graph = [ + [0, 1, 1, 0, 0], + [0, 0, 1, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 0], + ] + A_dense = np.array(graph) + A_sparse = sparse_cls(A_dense) + return A_dense, A_sparse + + +@pytest.mark.parametrize( + "func", + [ + spgraph.shortest_path, + spgraph.dijkstra, + spgraph.floyd_warshall, + spgraph.bellman_ford, + spgraph.johnson, + spgraph.reverse_cuthill_mckee, + spgraph.maximum_bipartite_matching, + spgraph.structural_rank, + ] +) +def test_csgraph_equiv(func, graphs): + A_dense, A_sparse = graphs + actual = func(A_sparse) + desired = func(sp.csc_array(A_dense)) + assert_equal(actual, desired) + + +def test_connected_components(graphs): + A_dense, A_sparse = graphs + func = spgraph.connected_components + + actual_comp, actual_labels = func(A_sparse) + desired_comp, desired_labels, = func(sp.csc_array(A_dense)) + + assert actual_comp == desired_comp + assert_equal(actual_labels, desired_labels) + + +def test_laplacian(graphs): + A_dense, A_sparse = graphs + sparse_cls = type(A_sparse) + func = spgraph.laplacian + + actual = func(A_sparse) + desired = func(sp.csc_array(A_dense)) + + assert isinstance(actual, sparse_cls) + + assert_equal(actual.todense(), desired.todense()) + + +@pytest.mark.parametrize( + "func", [spgraph.breadth_first_order, spgraph.depth_first_order] +) +def test_order_search(graphs, func): + A_dense, A_sparse = graphs + + actual = func(A_sparse, 0) + desired = func(sp.csc_array(A_dense), 0) + + assert_equal(actual, desired) + + +@pytest.mark.parametrize( + "func", [spgraph.breadth_first_tree, spgraph.depth_first_tree] +) +def test_tree_search(graphs, func): + A_dense, A_sparse = graphs + sparse_cls = type(A_sparse) + + actual = func(A_sparse, 0) + desired = func(sp.csc_array(A_dense), 0) + + assert isinstance(actual, sparse_cls) + + assert_equal(actual.todense(), desired.todense()) + + +def test_minimum_spanning_tree(graphs): + A_dense, A_sparse = graphs + sparse_cls = type(A_sparse) + func = spgraph.minimum_spanning_tree + + actual = func(A_sparse) + desired = func(sp.csc_array(A_dense)) + + assert isinstance(actual, sparse_cls) + + assert_equal(actual.todense(), desired.todense()) + + +def test_maximum_flow(graphs): + A_dense, A_sparse = graphs + sparse_cls = type(A_sparse) + func = spgraph.maximum_flow + + actual = func(A_sparse, 0, 2) + desired = func(sp.csr_array(A_dense), 0, 2) + + assert actual.flow_value == desired.flow_value + assert isinstance(actual.flow, sparse_cls) + + assert_equal(actual.flow.todense(), desired.flow.todense()) + + +def test_min_weight_full_bipartite_matching(graphs): + A_dense, A_sparse = graphs + func = spgraph.min_weight_full_bipartite_matching + + actual = func(A_sparse[0:2, 1:3]) + desired = func(sp.csc_array(A_dense)[0:2, 1:3]) + + assert_equal(actual, desired) + + +@check_sparse_version("0.15.4") +@pytest.mark.parametrize( + "func", + [ + spgraph.shortest_path, + spgraph.dijkstra, + spgraph.floyd_warshall, + spgraph.bellman_ford, + spgraph.johnson, + spgraph.minimum_spanning_tree, + ] +) +@pytest.mark.parametrize( + "fill_value, comp_func", + [(np.inf, np.isposinf), (np.nan, np.isnan)], +) +def test_nonzero_fill_value(graphs, func, fill_value, comp_func): + A_dense, A_sparse = graphs + A_sparse = A_sparse.astype(float) + A_sparse.fill_value = fill_value + sparse_cls = type(A_sparse) + + actual = func(A_sparse) + desired = func(sp.csc_array(A_dense)) + + if func == spgraph.minimum_spanning_tree: + assert isinstance(actual, sparse_cls) + assert comp_func(actual.fill_value) + actual = actual.todense() + actual[comp_func(actual)] = 0.0 + assert_equal(actual, desired.todense()) + else: + assert_equal(actual, desired) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_reordering.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_reordering.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e2e3fbf1950adf8559e4806062996720c789f9 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_reordering.py @@ -0,0 +1,70 @@ +import numpy as np +from numpy.testing import assert_equal +from scipy.sparse.csgraph import reverse_cuthill_mckee, structural_rank +from scipy.sparse import csc_array, csr_array, coo_array + + +def test_graph_reverse_cuthill_mckee(): + A = np.array([[1, 0, 0, 0, 1, 0, 0, 0], + [0, 1, 1, 0, 0, 1, 0, 1], + [0, 1, 1, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 0, 0, 1, 0], + [1, 0, 1, 0, 1, 0, 0, 0], + [0, 1, 0, 0, 0, 1, 0, 1], + [0, 0, 0, 1, 0, 0, 1, 0], + [0, 1, 0, 0, 0, 1, 0, 1]], dtype=int) + + graph = csr_array(A) + perm = reverse_cuthill_mckee(graph) + correct_perm = np.array([6, 3, 7, 5, 1, 2, 4, 0]) + assert_equal(perm, correct_perm) + + # Test int64 indices input + graph.indices = graph.indices.astype('int64') + graph.indptr = graph.indptr.astype('int64') + perm = reverse_cuthill_mckee(graph, True) + assert_equal(perm, correct_perm) + + +def test_graph_reverse_cuthill_mckee_ordering(): + data = np.ones(63,dtype=int) + rows = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, + 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, + 6, 6, 6, 7, 7, 7, 7, 8, 8, 8, 8, 9, 9, + 9, 10, 10, 10, 10, 10, 11, 11, 11, 11, + 12, 12, 12, 13, 13, 13, 13, 14, 14, 14, + 14, 15, 15, 15, 15, 15]) + cols = np.array([0, 2, 5, 8, 10, 1, 3, 9, 11, 0, 2, + 7, 10, 1, 3, 11, 4, 6, 12, 14, 0, 7, 13, + 15, 4, 6, 14, 2, 5, 7, 15, 0, 8, 10, 13, + 1, 9, 11, 0, 2, 8, 10, 15, 1, 3, 9, 11, + 4, 12, 14, 5, 8, 13, 15, 4, 6, 12, 14, + 5, 7, 10, 13, 15]) + graph = csr_array((data, (rows,cols))) + perm = reverse_cuthill_mckee(graph) + correct_perm = np.array([12, 14, 4, 6, 10, 8, 2, 15, + 0, 13, 7, 5, 9, 11, 1, 3]) + assert_equal(perm, correct_perm) + + +def test_graph_structural_rank(): + # Test square matrix #1 + A = csc_array([[1, 1, 0], + [1, 0, 1], + [0, 1, 0]]) + assert_equal(structural_rank(A), 3) + + # Test square matrix #2 + rows = np.array([0,0,0,0,0,1,1,2,2,3,3,3,3,3,3,4,4,5,5,6,6,7,7]) + cols = np.array([0,1,2,3,4,2,5,2,6,0,1,3,5,6,7,4,5,5,6,2,6,2,4]) + data = np.ones_like(rows) + B = coo_array((data,(rows,cols)), shape=(8,8)) + assert_equal(structural_rank(B), 6) + + #Test non-square matrix + C = csc_array([[1, 0, 2, 0], + [2, 0, 4, 0]]) + assert_equal(structural_rank(C), 2) + + #Test tall matrix + assert_equal(structural_rank(C.T), 2) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_shortest_path.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_shortest_path.py new file mode 100644 index 0000000000000000000000000000000000000000..bec46357596dc0d0403afc549e52795ff6a8833e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_shortest_path.py @@ -0,0 +1,484 @@ +from io import StringIO +import warnings +import numpy as np +from numpy.testing import assert_array_almost_equal, assert_array_equal, assert_allclose +from pytest import raises as assert_raises +from scipy.sparse.csgraph import (shortest_path, dijkstra, johnson, + bellman_ford, construct_dist_matrix, yen, + NegativeCycleError) +import scipy.sparse +from scipy.io import mmread +import pytest + +directed_G = np.array([[0, 3, 3, 0, 0], + [0, 0, 0, 2, 4], + [0, 0, 0, 0, 0], + [1, 0, 0, 0, 0], + [2, 0, 0, 2, 0]], dtype=float) + +undirected_G = np.array([[0, 3, 3, 1, 2], + [3, 0, 0, 2, 4], + [3, 0, 0, 0, 0], + [1, 2, 0, 0, 2], + [2, 4, 0, 2, 0]], dtype=float) + +unweighted_G = (directed_G > 0).astype(float) + +directed_SP = [[0, 3, 3, 5, 7], + [3, 0, 6, 2, 4], + [np.inf, np.inf, 0, np.inf, np.inf], + [1, 4, 4, 0, 8], + [2, 5, 5, 2, 0]] + +directed_2SP_0_to_3 = [[-9999, 0, -9999, 1, -9999], + [-9999, 0, -9999, 4, 1]] + +directed_sparse_zero_G = scipy.sparse.csr_array( + ( + [0, 1, 2, 3, 1], + ([0, 1, 2, 3, 4], [1, 2, 0, 4, 3]), + ), + shape=(5, 5), +) + +directed_sparse_zero_SP = [[0, 0, 1, np.inf, np.inf], + [3, 0, 1, np.inf, np.inf], + [2, 2, 0, np.inf, np.inf], + [np.inf, np.inf, np.inf, 0, 3], + [np.inf, np.inf, np.inf, 1, 0]] + +undirected_sparse_zero_G = scipy.sparse.csr_array( + ( + [0, 0, 1, 1, 2, 2, 1, 1], + ([0, 1, 1, 2, 2, 0, 3, 4], [1, 0, 2, 1, 0, 2, 4, 3]) + ), + shape=(5, 5), +) + +undirected_sparse_zero_SP = [[0, 0, 1, np.inf, np.inf], + [0, 0, 1, np.inf, np.inf], + [1, 1, 0, np.inf, np.inf], + [np.inf, np.inf, np.inf, 0, 1], + [np.inf, np.inf, np.inf, 1, 0]] + +directed_pred = np.array([[-9999, 0, 0, 1, 1], + [3, -9999, 0, 1, 1], + [-9999, -9999, -9999, -9999, -9999], + [3, 0, 0, -9999, 1], + [4, 0, 0, 4, -9999]], dtype=float) + +undirected_SP = np.array([[0, 3, 3, 1, 2], + [3, 0, 6, 2, 4], + [3, 6, 0, 4, 5], + [1, 2, 4, 0, 2], + [2, 4, 5, 2, 0]], dtype=float) + +undirected_SP_limit_2 = np.array([[0, np.inf, np.inf, 1, 2], + [np.inf, 0, np.inf, 2, np.inf], + [np.inf, np.inf, 0, np.inf, np.inf], + [1, 2, np.inf, 0, 2], + [2, np.inf, np.inf, 2, 0]], dtype=float) + +undirected_SP_limit_0 = np.ones((5, 5), dtype=float) - np.eye(5) +undirected_SP_limit_0[undirected_SP_limit_0 > 0] = np.inf + +undirected_pred = np.array([[-9999, 0, 0, 0, 0], + [1, -9999, 0, 1, 1], + [2, 0, -9999, 0, 0], + [3, 3, 0, -9999, 3], + [4, 4, 0, 4, -9999]], dtype=float) + +directed_negative_weighted_G = np.array([[0, 0, 0], + [-1, 0, 0], + [0, -1, 0]], dtype=float) + +directed_negative_weighted_SP = np.array([[0, np.inf, np.inf], + [-1, 0, np.inf], + [-2, -1, 0]], dtype=float) + +methods = ['auto', 'FW', 'D', 'BF', 'J'] + + +def test_dijkstra_limit(): + limits = [0, 2, np.inf] + results = [undirected_SP_limit_0, + undirected_SP_limit_2, + undirected_SP] + + def check(limit, result): + SP = dijkstra(undirected_G, directed=False, limit=limit) + assert_array_almost_equal(SP, result) + + for limit, result in zip(limits, results): + check(limit, result) + + +def test_directed(): + def check(method): + SP = shortest_path(directed_G, method=method, directed=True, + overwrite=False) + assert_array_almost_equal(SP, directed_SP) + + for method in methods: + check(method) + + +def test_undirected(): + def check(method, directed_in): + if directed_in: + SP1 = shortest_path(directed_G, method=method, directed=False, + overwrite=False) + assert_array_almost_equal(SP1, undirected_SP) + else: + SP2 = shortest_path(undirected_G, method=method, directed=True, + overwrite=False) + assert_array_almost_equal(SP2, undirected_SP) + + for method in methods: + for directed_in in (True, False): + check(method, directed_in) + + +def test_directed_sparse_zero(): + # test directed sparse graph with zero-weight edge and two connected components + def check(method): + SP = shortest_path(directed_sparse_zero_G, method=method, directed=True, + overwrite=False) + assert_array_almost_equal(SP, directed_sparse_zero_SP) + + for method in methods: + check(method) + + +def test_undirected_sparse_zero(): + def check(method, directed_in): + if directed_in: + SP1 = shortest_path(directed_sparse_zero_G, method=method, directed=False, + overwrite=False) + assert_array_almost_equal(SP1, undirected_sparse_zero_SP) + else: + SP2 = shortest_path(undirected_sparse_zero_G, method=method, directed=True, + overwrite=False) + assert_array_almost_equal(SP2, undirected_sparse_zero_SP) + + for method in methods: + for directed_in in (True, False): + check(method, directed_in) + + +@pytest.mark.parametrize('directed, SP_ans', + ((True, directed_SP), + (False, undirected_SP))) +@pytest.mark.parametrize('indices', ([0, 2, 4], [0, 4], [3, 4], [0, 0])) +def test_dijkstra_indices_min_only(directed, SP_ans, indices): + SP_ans = np.array(SP_ans) + indices = np.array(indices, dtype=np.int64) + min_ind_ans = indices[np.argmin(SP_ans[indices, :], axis=0)] + min_d_ans = np.zeros(SP_ans.shape[0], SP_ans.dtype) + for k in range(SP_ans.shape[0]): + min_d_ans[k] = SP_ans[min_ind_ans[k], k] + min_ind_ans[np.isinf(min_d_ans)] = -9999 + + SP, pred, sources = dijkstra(directed_G, + directed=directed, + indices=indices, + min_only=True, + return_predecessors=True) + assert_array_almost_equal(SP, min_d_ans) + assert_array_equal(min_ind_ans, sources) + SP = dijkstra(directed_G, + directed=directed, + indices=indices, + min_only=True, + return_predecessors=False) + assert_array_almost_equal(SP, min_d_ans) + + +@pytest.mark.parametrize('n', (10, 100, 1000)) +def test_dijkstra_min_only_random(n): + rng = np.random.default_rng(7345782358920239234) + data = scipy.sparse.random_array((n, n), density=0.5, format='lil', + rng=rng, dtype=np.float64) + data.setdiag(np.zeros(n, dtype=np.bool_)) + # choose some random vertices + v = np.arange(n) + rng.shuffle(v) + indices = v[:int(n*.1)] + ds, pred, sources = dijkstra(data, + directed=True, + indices=indices, + min_only=True, + return_predecessors=True) + for k in range(n): + p = pred[k] + s = sources[k] + while p != -9999: + assert sources[p] == s + p = pred[p] + + +def test_dijkstra_random(): + # reproduces the hang observed in gh-17782 + n = 10 + indices = [0, 4, 4, 5, 7, 9, 0, 6, 2, 3, 7, 9, 1, 2, 9, 2, 5, 6] + indptr = [0, 0, 2, 5, 6, 7, 8, 12, 15, 18, 18] + data = [0.33629, 0.40458, 0.47493, 0.42757, 0.11497, 0.91653, 0.69084, + 0.64979, 0.62555, 0.743, 0.01724, 0.99945, 0.31095, 0.15557, + 0.02439, 0.65814, 0.23478, 0.24072] + graph = scipy.sparse.csr_array((data, indices, indptr), shape=(n, n)) + dijkstra(graph, directed=True, return_predecessors=True) + + +def test_gh_17782_segfault(): + text = """%%MatrixMarket matrix coordinate real general + 84 84 22 + 2 1 4.699999809265137e+00 + 6 14 1.199999973177910e-01 + 9 6 1.199999973177910e-01 + 10 16 2.012000083923340e+01 + 11 10 1.422000026702881e+01 + 12 1 9.645999908447266e+01 + 13 18 2.012000083923340e+01 + 14 13 4.679999828338623e+00 + 15 11 1.199999973177910e-01 + 16 12 1.199999973177910e-01 + 18 15 1.199999973177910e-01 + 32 2 2.299999952316284e+00 + 33 20 6.000000000000000e+00 + 33 32 5.000000000000000e+00 + 36 9 3.720000028610229e+00 + 36 37 3.720000028610229e+00 + 36 38 3.720000028610229e+00 + 37 44 8.159999847412109e+00 + 38 32 7.903999328613281e+01 + 43 20 2.400000000000000e+01 + 43 33 4.000000000000000e+00 + 44 43 6.028000259399414e+01 + """ + data = mmread(StringIO(text), spmatrix=False) + dijkstra(data, directed=True, return_predecessors=True) + + +def test_shortest_path_indices(): + indices = np.arange(4) + + def check(func, indshape): + outshape = indshape + (5,) + SP = func(directed_G, directed=False, + indices=indices.reshape(indshape)) + assert_array_almost_equal(SP, undirected_SP[indices].reshape(outshape)) + + for indshape in [(4,), (4, 1), (2, 2)]: + for func in (dijkstra, bellman_ford, johnson, shortest_path): + check(func, indshape) + + assert_raises(ValueError, shortest_path, directed_G, method='FW', + indices=indices) + + +def test_predecessors(): + SP_res = {True: directed_SP, + False: undirected_SP} + pred_res = {True: directed_pred, + False: undirected_pred} + + def check(method, directed): + SP, pred = shortest_path(directed_G, method, directed=directed, + overwrite=False, + return_predecessors=True) + assert_array_almost_equal(SP, SP_res[directed]) + assert_array_almost_equal(pred, pred_res[directed]) + + for method in methods: + for directed in (True, False): + check(method, directed) + + +def test_construct_shortest_path(): + def check(method, directed): + SP1, pred = shortest_path(directed_G, + directed=directed, + overwrite=False, + return_predecessors=True) + SP2 = construct_dist_matrix(directed_G, pred, directed=directed) + assert_array_almost_equal(SP1, SP2) + + for method in methods: + for directed in (True, False): + check(method, directed) + + +def test_unweighted_path(): + def check(method, directed): + SP1 = shortest_path(directed_G, + directed=directed, + overwrite=False, + unweighted=True) + SP2 = shortest_path(unweighted_G, + directed=directed, + overwrite=False, + unweighted=False) + assert_array_almost_equal(SP1, SP2) + + for method in methods: + for directed in (True, False): + check(method, directed) + + +def test_negative_cycles(): + # create a small graph with a negative cycle + graph = np.ones([5, 5]) + graph.flat[::6] = 0 + graph[1, 2] = -2 + + def check(method, directed): + assert_raises(NegativeCycleError, shortest_path, graph, method, + directed) + + for directed in (True, False): + for method in ['FW', 'J', 'BF']: + check(method, directed) + + assert_raises(NegativeCycleError, yen, graph, 0, 1, 1, + directed=directed) + + +@pytest.mark.parametrize("method", ['FW', 'J', 'BF']) +def test_negative_weights(method): + SP = shortest_path(directed_negative_weighted_G, method, directed=True) + assert_allclose(SP, directed_negative_weighted_SP, atol=1e-10) + + +def test_masked_input(): + np.ma.masked_equal(directed_G, 0) + + def check(method): + SP = shortest_path(directed_G, method=method, directed=True, + overwrite=False) + assert_array_almost_equal(SP, directed_SP) + + for method in methods: + check(method) + + +def test_overwrite(): + G = np.array([[0, 3, 3, 1, 2], + [3, 0, 0, 2, 4], + [3, 0, 0, 0, 0], + [1, 2, 0, 0, 2], + [2, 4, 0, 2, 0]], dtype=float) + foo = G.copy() + shortest_path(foo, overwrite=False) + assert_array_equal(foo, G) + + +@pytest.mark.parametrize('method', methods) +def test_buffer(method): + # Smoke test that sparse matrices with read-only buffers (e.g., those from + # joblib workers) do not cause:: + # + # ValueError: buffer source array is read-only + # + G = scipy.sparse.csr_array([[1.]]) + G.data.flags['WRITEABLE'] = False + shortest_path(G, method=method) + + +def test_NaN_warnings(): + with warnings.catch_warnings(record=True) as record: + shortest_path(np.array([[0, 1], [np.nan, 0]])) + for r in record: + assert r.category is not RuntimeWarning + + +def test_sparse_matrices(): + # Test that using lil,csr and csc sparse matrix do not cause error + G_dense = np.array([[0, 3, 0, 0, 0], + [0, 0, -1, 0, 0], + [0, 0, 0, 2, 0], + [0, 0, 0, 0, 4], + [0, 0, 0, 0, 0]], dtype=float) + SP = shortest_path(G_dense) + G_csr = scipy.sparse.csr_array(G_dense) + G_csc = scipy.sparse.csc_array(G_dense) + G_lil = scipy.sparse.lil_array(G_dense) + assert_array_almost_equal(SP, shortest_path(G_csr)) + assert_array_almost_equal(SP, shortest_path(G_csc)) + assert_array_almost_equal(SP, shortest_path(G_lil)) + + +def test_yen_directed(): + distances, predecessors = yen( + directed_G, + source=0, + sink=3, + K=2, + return_predecessors=True + ) + assert_allclose(distances, [5., 9.]) + assert_allclose(predecessors, directed_2SP_0_to_3) + + +def test_yen_undirected(): + distances = yen( + undirected_G, + source=0, + sink=3, + K=4, + ) + assert_allclose(distances, [1., 4., 5., 8.]) + +def test_yen_unweighted(): + # Ask for more paths than there are, verify only the available paths are returned + distances, predecessors = yen( + directed_G, + source=0, + sink=3, + K=4, + unweighted=True, + return_predecessors=True, + ) + assert_allclose(distances, [2., 3.]) + assert_allclose(predecessors, directed_2SP_0_to_3) + +def test_yen_no_paths(): + distances = yen( + directed_G, + source=2, + sink=3, + K=1, + ) + assert distances.size == 0 + +def test_yen_negative_weights(): + distances = yen( + directed_negative_weighted_G, + source=2, + sink=0, + K=1, + ) + assert_allclose(distances, [-2.]) + + +@pytest.mark.parametrize("min_only", (True, False)) +@pytest.mark.parametrize("directed", (True, False)) +@pytest.mark.parametrize("return_predecessors", (True, False)) +@pytest.mark.parametrize("index_dtype", (np.int32, np.int64)) +@pytest.mark.parametrize("indices", (None, [1])) +def test_20904(min_only, directed, return_predecessors, index_dtype, indices): + """Test two failures from gh-20904: int32 and indices-as-None.""" + adj_mat = scipy.sparse.eye_array(4, format="csr") + adj_mat = scipy.sparse.csr_array( + ( + adj_mat.data, + adj_mat.indices.astype(index_dtype), + adj_mat.indptr.astype(index_dtype), + ), + ) + dijkstra( + adj_mat, + directed, + indices=indices, + min_only=min_only, + return_predecessors=return_predecessors, + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_spanning_tree.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_spanning_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..0cbf02d2d44a8d30761e64018e0aa4e6dcc7db9f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_spanning_tree.py @@ -0,0 +1,66 @@ +"""Test the minimum spanning tree function""" +import numpy as np +from numpy.testing import assert_ +import numpy.testing as npt +from scipy.sparse import csr_array +from scipy.sparse.csgraph import minimum_spanning_tree + + +def test_minimum_spanning_tree(): + + # Create a graph with two connected components. + graph = [[0,1,0,0,0], + [1,0,0,0,0], + [0,0,0,8,5], + [0,0,8,0,1], + [0,0,5,1,0]] + graph = np.asarray(graph) + + # Create the expected spanning tree. + expected = [[0,1,0,0,0], + [0,0,0,0,0], + [0,0,0,0,5], + [0,0,0,0,1], + [0,0,0,0,0]] + expected = np.asarray(expected) + + # Ensure minimum spanning tree code gives this expected output. + csgraph = csr_array(graph) + mintree = minimum_spanning_tree(csgraph) + mintree_array = mintree.toarray() + npt.assert_array_equal(mintree_array, expected, + 'Incorrect spanning tree found.') + + # Ensure that the original graph was not modified. + npt.assert_array_equal(csgraph.toarray(), graph, + 'Original graph was modified.') + + # Now let the algorithm modify the csgraph in place. + mintree = minimum_spanning_tree(csgraph, overwrite=True) + npt.assert_array_equal(mintree.toarray(), expected, + 'Graph was not properly modified to contain MST.') + + np.random.seed(1234) + for N in (5, 10, 15, 20): + + # Create a random graph. + graph = 3 + np.random.random((N, N)) + csgraph = csr_array(graph) + + # The spanning tree has at most N - 1 edges. + mintree = minimum_spanning_tree(csgraph) + assert_(mintree.nnz < N) + + # Set the sub diagonal to 1 to create a known spanning tree. + idx = np.arange(N-1) + graph[idx,idx+1] = 1 + csgraph = csr_array(graph) + mintree = minimum_spanning_tree(csgraph) + + # We expect to see this pattern in the spanning tree and otherwise + # have this zero. + expected = np.zeros((N, N)) + expected[idx, idx+1] = 1 + + npt.assert_array_equal(mintree.toarray(), expected, + 'Incorrect spanning tree found.') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_traversal.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_traversal.py new file mode 100644 index 0000000000000000000000000000000000000000..7bd20420ef23948ea1e4deb73d5a13eea479ce0e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/csgraph/tests/test_traversal.py @@ -0,0 +1,148 @@ +import numpy as np +import pytest +from numpy.testing import assert_array_almost_equal +from scipy.sparse import csr_array, csr_matrix, coo_array, coo_matrix +from scipy.sparse.csgraph import (breadth_first_tree, depth_first_tree, + csgraph_to_dense, csgraph_from_dense, csgraph_masked_from_dense) + + +def test_graph_breadth_first(): + csgraph = np.array([[0, 1, 2, 0, 0], + [1, 0, 0, 0, 3], + [2, 0, 0, 7, 0], + [0, 0, 7, 0, 1], + [0, 3, 0, 1, 0]]) + csgraph = csgraph_from_dense(csgraph, null_value=0) + + bfirst = np.array([[0, 1, 2, 0, 0], + [0, 0, 0, 0, 3], + [0, 0, 0, 7, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0]]) + + for directed in [True, False]: + bfirst_test = breadth_first_tree(csgraph, 0, directed) + assert_array_almost_equal(csgraph_to_dense(bfirst_test), + bfirst) + + +def test_graph_depth_first(): + csgraph = np.array([[0, 1, 2, 0, 0], + [1, 0, 0, 0, 3], + [2, 0, 0, 7, 0], + [0, 0, 7, 0, 1], + [0, 3, 0, 1, 0]]) + csgraph = csgraph_from_dense(csgraph, null_value=0) + + dfirst = np.array([[0, 1, 0, 0, 0], + [0, 0, 0, 0, 3], + [0, 0, 0, 0, 0], + [0, 0, 7, 0, 0], + [0, 0, 0, 1, 0]]) + + for directed in [True, False]: + dfirst_test = depth_first_tree(csgraph, 0, directed) + assert_array_almost_equal(csgraph_to_dense(dfirst_test), dfirst) + + +def test_return_type(): + from .._laplacian import laplacian + from .._min_spanning_tree import minimum_spanning_tree + + np_csgraph = np.array([[0, 1, 2, 0, 0], + [1, 0, 0, 0, 3], + [2, 0, 0, 7, 0], + [0, 0, 7, 0, 1], + [0, 3, 0, 1, 0]]) + csgraph = csr_array(np_csgraph) + assert isinstance(laplacian(csgraph), coo_array) + assert isinstance(minimum_spanning_tree(csgraph), csr_array) + for directed in [True, False]: + assert isinstance(depth_first_tree(csgraph, 0, directed), csr_array) + assert isinstance(breadth_first_tree(csgraph, 0, directed), csr_array) + + csgraph = csgraph_from_dense(np_csgraph, null_value=0) + assert isinstance(csgraph, csr_array) + assert isinstance(laplacian(csgraph), coo_array) + assert isinstance(minimum_spanning_tree(csgraph), csr_array) + for directed in [True, False]: + assert isinstance(depth_first_tree(csgraph, 0, directed), csr_array) + assert isinstance(breadth_first_tree(csgraph, 0, directed), csr_array) + + csgraph = csgraph_masked_from_dense(np_csgraph, null_value=0) + assert isinstance(csgraph, np.ma.MaskedArray) + assert csgraph._baseclass is np.ndarray + # laplacian doesnt work with masked arrays so not here + assert isinstance(minimum_spanning_tree(csgraph), csr_array) + for directed in [True, False]: + assert isinstance(depth_first_tree(csgraph, 0, directed), csr_array) + assert isinstance(breadth_first_tree(csgraph, 0, directed), csr_array) + + # start of testing with matrix/spmatrix types + with np.testing.suppress_warnings() as sup: + sup.filter(DeprecationWarning, "the matrix subclass.*") + sup.filter(PendingDeprecationWarning, "the matrix subclass.*") + + nm_csgraph = np.matrix([[0, 1, 2, 0, 0], + [1, 0, 0, 0, 3], + [2, 0, 0, 7, 0], + [0, 0, 7, 0, 1], + [0, 3, 0, 1, 0]]) + + csgraph = csr_matrix(nm_csgraph) + assert isinstance(laplacian(csgraph), coo_matrix) + assert isinstance(minimum_spanning_tree(csgraph), csr_matrix) + for directed in [True, False]: + assert isinstance(depth_first_tree(csgraph, 0, directed), csr_matrix) + assert isinstance(breadth_first_tree(csgraph, 0, directed), csr_matrix) + + csgraph = csgraph_from_dense(nm_csgraph, null_value=0) + assert isinstance(csgraph, csr_matrix) + assert isinstance(laplacian(csgraph), coo_matrix) + assert isinstance(minimum_spanning_tree(csgraph), csr_matrix) + for directed in [True, False]: + assert isinstance(depth_first_tree(csgraph, 0, directed), csr_matrix) + assert isinstance(breadth_first_tree(csgraph, 0, directed), csr_matrix) + + mm_csgraph = csgraph_masked_from_dense(nm_csgraph, null_value=0) + assert isinstance(mm_csgraph, np.ma.MaskedArray) + # laplacian doesnt work with masked arrays so not here + assert isinstance(minimum_spanning_tree(csgraph), csr_matrix) + for directed in [True, False]: + assert isinstance(depth_first_tree(csgraph, 0, directed), csr_matrix) + assert isinstance(breadth_first_tree(csgraph, 0, directed), csr_matrix) + # end of testing with matrix/spmatrix types + + +def test_graph_breadth_first_trivial_graph(): + csgraph = np.array([[0]]) + csgraph = csgraph_from_dense(csgraph, null_value=0) + + bfirst = np.array([[0]]) + + for directed in [True, False]: + bfirst_test = breadth_first_tree(csgraph, 0, directed) + assert_array_almost_equal(csgraph_to_dense(bfirst_test), bfirst) + + +def test_graph_depth_first_trivial_graph(): + csgraph = np.array([[0]]) + csgraph = csgraph_from_dense(csgraph, null_value=0) + + bfirst = np.array([[0]]) + + for directed in [True, False]: + bfirst_test = depth_first_tree(csgraph, 0, directed) + assert_array_almost_equal(csgraph_to_dense(bfirst_test), + bfirst) + + +@pytest.mark.parametrize('directed', [True, False]) +@pytest.mark.parametrize('tree_func', [breadth_first_tree, depth_first_tree]) +def test_int64_indices(tree_func, directed): + # See https://github.com/scipy/scipy/issues/18716 + g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2)) + assert g.indices.dtype == np.int64 + tree = tree_func(g, 0, directed=directed) + assert_array_almost_equal(csgraph_to_dense(tree), [[0, 1], [0, 0]]) + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b9500b1111e714b235e0d2da55105d093c97c5db --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/__init__.py @@ -0,0 +1,148 @@ +""" +Sparse linear algebra (:mod:`scipy.sparse.linalg`) +================================================== + +.. currentmodule:: scipy.sparse.linalg + +Abstract linear operators +------------------------- + +.. autosummary:: + :toctree: generated/ + + LinearOperator -- abstract representation of a linear operator + aslinearoperator -- convert an object to an abstract linear operator + +Matrix Operations +----------------- + +.. autosummary:: + :toctree: generated/ + + inv -- compute the sparse matrix inverse + expm -- compute the sparse matrix exponential + expm_multiply -- compute the product of a matrix exponential and a matrix + matrix_power -- compute the matrix power by raising a matrix to an exponent + +Matrix norms +------------ + +.. autosummary:: + :toctree: generated/ + + norm -- Norm of a sparse matrix + onenormest -- Estimate the 1-norm of a sparse matrix + +Solving linear problems +----------------------- + +Direct methods for linear equation systems: + +.. autosummary:: + :toctree: generated/ + + spsolve -- Solve the sparse linear system Ax=b + spsolve_triangular -- Solve sparse linear system Ax=b for a triangular A. + is_sptriangular -- Check if sparse A is triangular. + spbandwidth -- Find the bandwidth of a sparse matrix. + factorized -- Pre-factorize matrix to a function solving a linear system + MatrixRankWarning -- Warning on exactly singular matrices + use_solver -- Select direct solver to use + +Iterative methods for linear equation systems: + +.. autosummary:: + :toctree: generated/ + + bicg -- Use BIConjugate Gradient iteration to solve Ax = b + bicgstab -- Use BIConjugate Gradient STABilized iteration to solve Ax = b + cg -- Use Conjugate Gradient iteration to solve Ax = b + cgs -- Use Conjugate Gradient Squared iteration to solve Ax = b + gmres -- Use Generalized Minimal RESidual iteration to solve Ax = b + lgmres -- Solve a matrix equation using the LGMRES algorithm + minres -- Use MINimum RESidual iteration to solve Ax = b + qmr -- Use Quasi-Minimal Residual iteration to solve Ax = b + gcrotmk -- Solve a matrix equation using the GCROT(m,k) algorithm + tfqmr -- Use Transpose-Free Quasi-Minimal Residual iteration to solve Ax = b + +Iterative methods for least-squares problems: + +.. autosummary:: + :toctree: generated/ + + lsqr -- Find the least-squares solution to a sparse linear equation system + lsmr -- Find the least-squares solution to a sparse linear equation system + +Matrix factorizations +--------------------- + +Eigenvalue problems: + +.. autosummary:: + :toctree: generated/ + + eigs -- Find k eigenvalues and eigenvectors of the square matrix A + eigsh -- Find k eigenvalues and eigenvectors of a symmetric matrix + lobpcg -- Solve symmetric partial eigenproblems with optional preconditioning + +Singular values problems: + +.. autosummary:: + :toctree: generated/ + + svds -- Compute k singular values/vectors for a sparse matrix + +The `svds` function supports the following solvers: + +.. toctree:: + + sparse.linalg.svds-arpack + sparse.linalg.svds-lobpcg + sparse.linalg.svds-propack + +Complete or incomplete LU factorizations + +.. autosummary:: + :toctree: generated/ + + splu -- Compute a LU decomposition for a sparse matrix + spilu -- Compute an incomplete LU decomposition for a sparse matrix + SuperLU -- Object representing an LU factorization + +Sparse arrays with structure +---------------------------- + +.. autosummary:: + :toctree: generated/ + + LaplacianNd -- Laplacian on a uniform rectangular grid in ``N`` dimensions + +Exceptions +---------- + +.. autosummary:: + :toctree: generated/ + + ArpackNoConvergence + ArpackError + +""" + +from ._isolve import * +from ._dsolve import * +from ._interface import * +from ._eigen import * +from ._matfuncs import * +from ._onenormest import * +from ._norm import * +from ._expm_multiply import * +from ._special_sparse_arrays import * + +# Deprecated namespaces, to be removed in v2.0.0 +from . import isolve, dsolve, interface, eigen, matfuncs + +__all__ = [s for s in dir() if not s.startswith('_')] + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_expm_multiply.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_expm_multiply.py new file mode 100644 index 0000000000000000000000000000000000000000..5dff7ec42716d3a2a7159e0e272ae6bc5e4f93f7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_expm_multiply.py @@ -0,0 +1,816 @@ +"""Compute the action of the matrix exponential.""" +from warnings import warn + +import numpy as np + +import scipy.linalg +import scipy.sparse.linalg +from scipy.linalg._decomp_qr import qr +from scipy.sparse._sputils import is_pydata_spmatrix +from scipy.sparse.linalg import aslinearoperator +from scipy.sparse.linalg._interface import IdentityOperator +from scipy.sparse.linalg._onenormest import onenormest + +__all__ = ['expm_multiply'] + + +def _exact_inf_norm(A): + # A compatibility function which should eventually disappear. + if scipy.sparse.issparse(A): + return max(abs(A).sum(axis=1).flat) + elif is_pydata_spmatrix(A): + return max(abs(A).sum(axis=1)) + else: + return np.linalg.norm(A, np.inf) + + +def _exact_1_norm(A): + # A compatibility function which should eventually disappear. + if scipy.sparse.issparse(A): + return max(abs(A).sum(axis=0).flat) + elif is_pydata_spmatrix(A): + return max(abs(A).sum(axis=0)) + else: + return np.linalg.norm(A, 1) + + +def _trace(A): + # A compatibility function which should eventually disappear. + if is_pydata_spmatrix(A): + return A.to_scipy_sparse().trace() + else: + return A.trace() + + +def traceest(A, m3, seed=None): + """Estimate `np.trace(A)` using `3*m3` matrix-vector products. + + The result is not deterministic. + + Parameters + ---------- + A : LinearOperator + Linear operator whose trace will be estimated. Has to be square. + m3 : int + Number of matrix-vector products divided by 3 used to estimate the + trace. + seed : optional + Seed for `numpy.random.default_rng`. + Can be provided to obtain deterministic results. + + Returns + ------- + trace : LinearOperator.dtype + Estimate of the trace + + Notes + ----- + This is the Hutch++ algorithm given in [1]_. + + References + ---------- + .. [1] Meyer, Raphael A., Cameron Musco, Christopher Musco, and David P. + Woodruff. "Hutch++: Optimal Stochastic Trace Estimation." In Symposium + on Simplicity in Algorithms (SOSA), pp. 142-155. Society for Industrial + and Applied Mathematics, 2021 + https://doi.org/10.1137/1.9781611976496.16 + + """ + rng = np.random.default_rng(seed) + if len(A.shape) != 2 or A.shape[-1] != A.shape[-2]: + raise ValueError("Expected A to be like a square matrix.") + n = A.shape[-1] + S = rng.choice([-1.0, +1.0], [n, m3]) + Q, _ = qr(A.matmat(S), overwrite_a=True, mode='economic') + trQAQ = np.trace(Q.conj().T @ A.matmat(Q)) + G = rng.choice([-1, +1], [n, m3]) + right = G - Q@(Q.conj().T @ G) + trGAG = np.trace(right.conj().T @ A.matmat(right)) + return trQAQ + trGAG/m3 + + +def _ident_like(A): + # A compatibility function which should eventually disappear. + if scipy.sparse.issparse(A): + # Creates a sparse matrix in dia format + out = scipy.sparse.eye(A.shape[0], A.shape[1], dtype=A.dtype) + if scipy.sparse.issparse(A): + return out.asformat(A.format) + return scipy.sparse.dia_array(out).asformat(A.format) + elif is_pydata_spmatrix(A): + import sparse + return sparse.eye(A.shape[0], A.shape[1], dtype=A.dtype) + elif isinstance(A, scipy.sparse.linalg.LinearOperator): + return IdentityOperator(A.shape, dtype=A.dtype) + else: + return np.eye(A.shape[0], A.shape[1], dtype=A.dtype) + + +def expm_multiply(A, B, start=None, stop=None, num=None, + endpoint=None, traceA=None): + """ + Compute the action of the matrix exponential of A on B. + + Parameters + ---------- + A : transposable linear operator + The operator whose exponential is of interest. + B : ndarray, sparse array + The matrix or vector to be multiplied by the matrix exponential of A. + start : scalar, optional + The starting time point of the sequence. + stop : scalar, optional + The end time point of the sequence, unless `endpoint` is set to False. + In that case, the sequence consists of all but the last of ``num + 1`` + evenly spaced time points, so that `stop` is excluded. + Note that the step size changes when `endpoint` is False. + num : int, optional + Number of time points to use. + endpoint : bool, optional + If True, `stop` is the last time point. Otherwise, it is not included. + traceA : scalar, optional + Trace of `A`. If not given the trace is estimated for linear operators, + or calculated exactly for sparse matrices. It is used to precondition + `A`, thus an approximate trace is acceptable. + For linear operators, `traceA` should be provided to ensure performance + as the estimation is not guaranteed to be reliable for all cases. + + .. versionadded:: 1.9.0 + + Returns + ------- + expm_A_B : ndarray + The result of the action :math:`e^{t_k A} B`. + + Warns + ----- + UserWarning + If `A` is a linear operator and ``traceA=None`` (default). + + Notes + ----- + The optional arguments defining the sequence of evenly spaced time points + are compatible with the arguments of `numpy.linspace`. + + The output ndarray shape is somewhat complicated so I explain it here. + The ndim of the output could be either 1, 2, or 3. + It would be 1 if you are computing the expm action on a single vector + at a single time point. + It would be 2 if you are computing the expm action on a vector + at multiple time points, or if you are computing the expm action + on a matrix at a single time point. + It would be 3 if you want the action on a matrix with multiple + columns at multiple time points. + If multiple time points are requested, expm_A_B[0] will always + be the action of the expm at the first time point, + regardless of whether the action is on a vector or a matrix. + + References + ---------- + .. [1] Awad H. Al-Mohy and Nicholas J. Higham (2011) + "Computing the Action of the Matrix Exponential, + with an Application to Exponential Integrators." + SIAM Journal on Scientific Computing, + 33 (2). pp. 488-511. ISSN 1064-8275 + http://eprints.ma.man.ac.uk/1591/ + + .. [2] Nicholas J. Higham and Awad H. Al-Mohy (2010) + "Computing Matrix Functions." + Acta Numerica, + 19. 159-208. ISSN 0962-4929 + http://eprints.ma.man.ac.uk/1451/ + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse import csc_array + >>> from scipy.sparse.linalg import expm, expm_multiply + >>> A = csc_array([[1, 0], [0, 1]]) + >>> A.toarray() + array([[1, 0], + [0, 1]], dtype=int64) + >>> B = np.array([np.exp(-1.), np.exp(-2.)]) + >>> B + array([ 0.36787944, 0.13533528]) + >>> expm_multiply(A, B, start=1, stop=2, num=3, endpoint=True) + array([[ 1. , 0.36787944], + [ 1.64872127, 0.60653066], + [ 2.71828183, 1. ]]) + >>> expm(A).dot(B) # Verify 1st timestep + array([ 1. , 0.36787944]) + >>> expm(1.5*A).dot(B) # Verify 2nd timestep + array([ 1.64872127, 0.60653066]) + >>> expm(2*A).dot(B) # Verify 3rd timestep + array([ 2.71828183, 1. ]) + """ + if all(arg is None for arg in (start, stop, num, endpoint)): + X = _expm_multiply_simple(A, B, traceA=traceA) + else: + X, status = _expm_multiply_interval(A, B, start, stop, num, + endpoint, traceA=traceA) + return X + + +def _expm_multiply_simple(A, B, t=1.0, traceA=None, balance=False): + """ + Compute the action of the matrix exponential at a single time point. + + Parameters + ---------- + A : transposable linear operator + The operator whose exponential is of interest. + B : ndarray + The matrix to be multiplied by the matrix exponential of A. + t : float + A time point. + traceA : scalar, optional + Trace of `A`. If not given the trace is estimated for linear operators, + or calculated exactly for sparse matrices. It is used to precondition + `A`, thus an approximate trace is acceptable + balance : bool + Indicates whether or not to apply balancing. + + Returns + ------- + F : ndarray + :math:`e^{t A} B` + + Notes + ----- + This is algorithm (3.2) in Al-Mohy and Higham (2011). + + """ + if balance: + raise NotImplementedError + if len(A.shape) != 2 or A.shape[0] != A.shape[1]: + raise ValueError('expected A to be like a square matrix') + if A.shape[1] != B.shape[0]: + raise ValueError(f'shapes of matrices A {A.shape} and B {B.shape}' + ' are incompatible') + ident = _ident_like(A) + is_linear_operator = isinstance(A, scipy.sparse.linalg.LinearOperator) + n = A.shape[0] + if len(B.shape) == 1: + n0 = 1 + elif len(B.shape) == 2: + n0 = B.shape[1] + else: + raise ValueError('expected B to be like a matrix or a vector') + u_d = 2**-53 + tol = u_d + if traceA is None: + if is_linear_operator: + warn("Trace of LinearOperator not available, it will be estimated." + " Provide `traceA` to ensure performance.", stacklevel=3) + # m3=1 is bit arbitrary choice, a more accurate trace (larger m3) might + # speed up exponential calculation, but trace estimation is more costly + traceA = traceest(A, m3=1) if is_linear_operator else _trace(A) + mu = traceA / float(n) + A = A - mu * ident + A_1_norm = onenormest(A) if is_linear_operator else _exact_1_norm(A) + if t*A_1_norm == 0: + m_star, s = 0, 1 + else: + ell = 2 + norm_info = LazyOperatorNormInfo(t*A, A_1_norm=t*A_1_norm, ell=ell) + m_star, s = _fragment_3_1(norm_info, n0, tol, ell=ell) + return _expm_multiply_simple_core(A, B, t, mu, m_star, s, tol, balance) + + +def _expm_multiply_simple_core(A, B, t, mu, m_star, s, tol=None, balance=False): + """ + A helper function. + """ + if balance: + raise NotImplementedError + if tol is None: + u_d = 2 ** -53 + tol = u_d + F = B + eta = np.exp(t*mu / float(s)) + for i in range(s): + c1 = _exact_inf_norm(B) + for j in range(m_star): + coeff = t / float(s*(j+1)) + B = coeff * A.dot(B) + c2 = _exact_inf_norm(B) + F = F + B + if c1 + c2 <= tol * _exact_inf_norm(F): + break + c1 = c2 + F = eta * F + B = F + return F + + +# This table helps to compute bounds. +# They seem to have been difficult to calculate, involving symbolic +# manipulation of equations, followed by numerical root finding. +_theta = { + # The first 30 values are from table A.3 of Computing Matrix Functions. + 1: 2.29e-16, + 2: 2.58e-8, + 3: 1.39e-5, + 4: 3.40e-4, + 5: 2.40e-3, + 6: 9.07e-3, + 7: 2.38e-2, + 8: 5.00e-2, + 9: 8.96e-2, + 10: 1.44e-1, + # 11 + 11: 2.14e-1, + 12: 3.00e-1, + 13: 4.00e-1, + 14: 5.14e-1, + 15: 6.41e-1, + 16: 7.81e-1, + 17: 9.31e-1, + 18: 1.09, + 19: 1.26, + 20: 1.44, + # 21 + 21: 1.62, + 22: 1.82, + 23: 2.01, + 24: 2.22, + 25: 2.43, + 26: 2.64, + 27: 2.86, + 28: 3.08, + 29: 3.31, + 30: 3.54, + # The rest are from table 3.1 of + # Computing the Action of the Matrix Exponential. + 35: 4.7, + 40: 6.0, + 45: 7.2, + 50: 8.5, + 55: 9.9, + } + + +def _onenormest_matrix_power(A, p, + t=2, itmax=5, compute_v=False, compute_w=False): + """ + Efficiently estimate the 1-norm of A^p. + + Parameters + ---------- + A : ndarray + Matrix whose 1-norm of a power is to be computed. + p : int + Non-negative integer power. + t : int, optional + A positive parameter controlling the tradeoff between + accuracy versus time and memory usage. + Larger values take longer and use more memory + but give more accurate output. + itmax : int, optional + Use at most this many iterations. + compute_v : bool, optional + Request a norm-maximizing linear operator input vector if True. + compute_w : bool, optional + Request a norm-maximizing linear operator output vector if True. + + Returns + ------- + est : float + An underestimate of the 1-norm of the sparse matrix. + v : ndarray, optional + The vector such that ||Av||_1 == est*||v||_1. + It can be thought of as an input to the linear operator + that gives an output with particularly large norm. + w : ndarray, optional + The vector Av which has relatively large 1-norm. + It can be thought of as an output of the linear operator + that is relatively large in norm compared to the input. + + """ + #XXX Eventually turn this into an API function in the _onenormest module, + #XXX and remove its underscore, + #XXX but wait until expm_multiply goes into scipy. + from scipy.sparse.linalg._onenormest import onenormest + return onenormest(aslinearoperator(A) ** p) + +class LazyOperatorNormInfo: + """ + Information about an operator is lazily computed. + + The information includes the exact 1-norm of the operator, + in addition to estimates of 1-norms of powers of the operator. + This uses the notation of Computing the Action (2011). + This class is specialized enough to probably not be of general interest + outside of this module. + + """ + + def __init__(self, A, A_1_norm=None, ell=2, scale=1): + """ + Provide the operator and some norm-related information. + + Parameters + ---------- + A : linear operator + The operator of interest. + A_1_norm : float, optional + The exact 1-norm of A. + ell : int, optional + A technical parameter controlling norm estimation quality. + scale : int, optional + If specified, return the norms of scale*A instead of A. + + """ + self._A = A + self._A_1_norm = A_1_norm + self._ell = ell + self._d = {} + self._scale = scale + + def set_scale(self,scale): + """ + Set the scale parameter. + """ + self._scale = scale + + def onenorm(self): + """ + Compute the exact 1-norm. + """ + if self._A_1_norm is None: + self._A_1_norm = _exact_1_norm(self._A) + return self._scale*self._A_1_norm + + def d(self, p): + """ + Lazily estimate :math:`d_p(A) ~= || A^p ||^(1/p)` + where :math:`||.||` is the 1-norm. + """ + if p not in self._d: + est = _onenormest_matrix_power(self._A, p, self._ell) + self._d[p] = est ** (1.0 / p) + return self._scale*self._d[p] + + def alpha(self, p): + """ + Lazily compute max(d(p), d(p+1)). + """ + return max(self.d(p), self.d(p+1)) + +def _compute_cost_div_m(m, p, norm_info): + """ + A helper function for computing bounds. + + This is equation (3.10). + It measures cost in terms of the number of required matrix products. + + Parameters + ---------- + m : int + A valid key of _theta. + p : int + A matrix power. + norm_info : LazyOperatorNormInfo + Information about 1-norms of related operators. + + Returns + ------- + cost_div_m : int + Required number of matrix products divided by m. + + """ + return int(np.ceil(norm_info.alpha(p) / _theta[m])) + + +def _compute_p_max(m_max): + """ + Compute the largest positive integer p such that p*(p-1) <= m_max + 1. + + Do this in a slightly dumb way, but safe and not too slow. + + Parameters + ---------- + m_max : int + A count related to bounds. + + """ + sqrt_m_max = np.sqrt(m_max) + p_low = int(np.floor(sqrt_m_max)) + p_high = int(np.ceil(sqrt_m_max + 1)) + return max(p for p in range(p_low, p_high+1) if p*(p-1) <= m_max + 1) + + +def _fragment_3_1(norm_info, n0, tol, m_max=55, ell=2): + """ + A helper function for the _expm_multiply_* functions. + + Parameters + ---------- + norm_info : LazyOperatorNormInfo + Information about norms of certain linear operators of interest. + n0 : int + Number of columns in the _expm_multiply_* B matrix. + tol : float + Expected to be + :math:`2^{-24}` for single precision or + :math:`2^{-53}` for double precision. + m_max : int + A value related to a bound. + ell : int + The number of columns used in the 1-norm approximation. + This is usually taken to be small, maybe between 1 and 5. + + Returns + ------- + best_m : int + Related to bounds for error control. + best_s : int + Amount of scaling. + + Notes + ----- + This is code fragment (3.1) in Al-Mohy and Higham (2011). + The discussion of default values for m_max and ell + is given between the definitions of equation (3.11) + and the definition of equation (3.12). + + """ + if ell < 1: + raise ValueError('expected ell to be a positive integer') + best_m = None + best_s = None + if _condition_3_13(norm_info.onenorm(), n0, m_max, ell): + for m, theta in _theta.items(): + s = int(np.ceil(norm_info.onenorm() / theta)) + if best_m is None or m * s < best_m * best_s: + best_m = m + best_s = s + else: + # Equation (3.11). + for p in range(2, _compute_p_max(m_max) + 1): + for m in range(p*(p-1)-1, m_max+1): + if m in _theta: + s = _compute_cost_div_m(m, p, norm_info) + if best_m is None or m * s < best_m * best_s: + best_m = m + best_s = s + best_s = max(best_s, 1) + return best_m, best_s + + +def _condition_3_13(A_1_norm, n0, m_max, ell): + """ + A helper function for the _expm_multiply_* functions. + + Parameters + ---------- + A_1_norm : float + The precomputed 1-norm of A. + n0 : int + Number of columns in the _expm_multiply_* B matrix. + m_max : int + A value related to a bound. + ell : int + The number of columns used in the 1-norm approximation. + This is usually taken to be small, maybe between 1 and 5. + + Returns + ------- + value : bool + Indicates whether or not the condition has been met. + + Notes + ----- + This is condition (3.13) in Al-Mohy and Higham (2011). + + """ + + # This is the rhs of equation (3.12). + p_max = _compute_p_max(m_max) + a = 2 * ell * p_max * (p_max + 3) + + # Evaluate the condition (3.13). + b = _theta[m_max] / float(n0 * m_max) + return A_1_norm <= a * b + + +def _expm_multiply_interval(A, B, start=None, stop=None, num=None, + endpoint=None, traceA=None, balance=False, + status_only=False): + """ + Compute the action of the matrix exponential at multiple time points. + + Parameters + ---------- + A : transposable linear operator + The operator whose exponential is of interest. + B : ndarray + The matrix to be multiplied by the matrix exponential of A. + start : scalar, optional + The starting time point of the sequence. + stop : scalar, optional + The end time point of the sequence, unless `endpoint` is set to False. + In that case, the sequence consists of all but the last of ``num + 1`` + evenly spaced time points, so that `stop` is excluded. + Note that the step size changes when `endpoint` is False. + num : int, optional + Number of time points to use. + traceA : scalar, optional + Trace of `A`. If not given the trace is estimated for linear operators, + or calculated exactly for sparse matrices. It is used to precondition + `A`, thus an approximate trace is acceptable + endpoint : bool, optional + If True, `stop` is the last time point. Otherwise, it is not included. + balance : bool + Indicates whether or not to apply balancing. + status_only : bool + A flag that is set to True for some debugging and testing operations. + + Returns + ------- + F : ndarray + :math:`e^{t_k A} B` + status : int + An integer status for testing and debugging. + + Notes + ----- + This is algorithm (5.2) in Al-Mohy and Higham (2011). + + There seems to be a typo, where line 15 of the algorithm should be + moved to line 6.5 (between lines 6 and 7). + + """ + if balance: + raise NotImplementedError + if len(A.shape) != 2 or A.shape[0] != A.shape[1]: + raise ValueError('expected A to be like a square matrix') + if A.shape[1] != B.shape[0]: + raise ValueError(f'shapes of matrices A {A.shape} and B {B.shape}' + ' are incompatible') + ident = _ident_like(A) + is_linear_operator = isinstance(A, scipy.sparse.linalg.LinearOperator) + n = A.shape[0] + if len(B.shape) == 1: + n0 = 1 + elif len(B.shape) == 2: + n0 = B.shape[1] + else: + raise ValueError('expected B to be like a matrix or a vector') + u_d = 2**-53 + tol = u_d + if traceA is None: + if is_linear_operator: + warn("Trace of LinearOperator not available, it will be estimated." + " Provide `traceA` to ensure performance.", stacklevel=3) + # m3=5 is bit arbitrary choice, a more accurate trace (larger m3) might + # speed up exponential calculation, but trace estimation is also costly + # an educated guess would need to consider the number of time points + traceA = traceest(A, m3=5) if is_linear_operator else _trace(A) + mu = traceA / float(n) + + # Get the linspace samples, attempting to preserve the linspace defaults. + linspace_kwargs = {'retstep': True} + if num is not None: + linspace_kwargs['num'] = num + if endpoint is not None: + linspace_kwargs['endpoint'] = endpoint + samples, step = np.linspace(start, stop, **linspace_kwargs) + + # Convert the linspace output to the notation used by the publication. + nsamples = len(samples) + if nsamples < 2: + raise ValueError('at least two time points are required') + q = nsamples - 1 + h = step + t_0 = samples[0] + t_q = samples[q] + + # Define the output ndarray. + # Use an ndim=3 shape, such that the last two indices + # are the ones that may be involved in level 3 BLAS operations. + X_shape = (nsamples,) + B.shape + X = np.empty(X_shape, dtype=np.result_type(A.dtype, B.dtype, float)) + t = t_q - t_0 + A = A - mu * ident + A_1_norm = onenormest(A) if is_linear_operator else _exact_1_norm(A) + ell = 2 + norm_info = LazyOperatorNormInfo(t*A, A_1_norm=t*A_1_norm, ell=ell) + if t*A_1_norm == 0: + m_star, s = 0, 1 + else: + m_star, s = _fragment_3_1(norm_info, n0, tol, ell=ell) + + # Compute the expm action up to the initial time point. + action_t0 = _expm_multiply_simple_core(A, B, t_0, mu, m_star, s) + if scipy.sparse.issparse(action_t0): + action_t0 = action_t0.toarray() + elif is_pydata_spmatrix(action_t0): + action_t0 = action_t0.todense() + X[0] = action_t0 + + # Compute the expm action at the rest of the time points. + if q <= s: + if status_only: + return 0 + else: + return _expm_multiply_interval_core_0(A, X, + h, mu, q, norm_info, tol, ell,n0) + elif not (q % s): + if status_only: + return 1 + else: + return _expm_multiply_interval_core_1(A, X, + h, mu, m_star, s, q, tol) + elif (q % s): + if status_only: + return 2 + else: + return _expm_multiply_interval_core_2(A, X, + h, mu, m_star, s, q, tol) + else: + raise Exception('internal error') + + +def _expm_multiply_interval_core_0(A, X, h, mu, q, norm_info, tol, ell, n0): + """ + A helper function, for the case q <= s. + """ + + # Compute the new values of m_star and s which should be applied + # over intervals of size t/q + if norm_info.onenorm() == 0: + m_star, s = 0, 1 + else: + norm_info.set_scale(1./q) + m_star, s = _fragment_3_1(norm_info, n0, tol, ell=ell) + norm_info.set_scale(1) + + for k in range(q): + X[k+1] = _expm_multiply_simple_core(A, X[k], h, mu, m_star, s) + return X, 0 + + +def _expm_multiply_interval_core_1(A, X, h, mu, m_star, s, q, tol): + """ + A helper function, for the case q > s and q % s == 0. + """ + d = q // s + input_shape = X.shape[1:] + K_shape = (m_star + 1, ) + input_shape + K = np.empty(K_shape, dtype=X.dtype) + for i in range(s): + Z = X[i*d] + K[0] = Z + high_p = 0 + for k in range(1, d+1): + F = K[0] + c1 = _exact_inf_norm(F) + for p in range(1, m_star+1): + if p > high_p: + K[p] = h * A.dot(K[p-1]) / float(p) + coeff = float(pow(k, p)) + F = F + coeff * K[p] + inf_norm_K_p_1 = _exact_inf_norm(K[p]) + c2 = coeff * inf_norm_K_p_1 + if c1 + c2 <= tol * _exact_inf_norm(F): + break + c1 = c2 + X[k + i*d] = np.exp(k*h*mu) * F + return X, 1 + + +def _expm_multiply_interval_core_2(A, X, h, mu, m_star, s, q, tol): + """ + A helper function, for the case q > s and q % s > 0. + """ + d = q // s + j = q // d + r = q - d * j + input_shape = X.shape[1:] + K_shape = (m_star + 1, ) + input_shape + K = np.empty(K_shape, dtype=X.dtype) + for i in range(j + 1): + Z = X[i*d] + K[0] = Z + high_p = 0 + if i < j: + effective_d = d + else: + effective_d = r + for k in range(1, effective_d+1): + F = K[0] + c1 = _exact_inf_norm(F) + for p in range(1, m_star+1): + if p == high_p + 1: + K[p] = h * A.dot(K[p-1]) / float(p) + high_p = p + coeff = float(pow(k, p)) + F = F + coeff * K[p] + inf_norm_K_p_1 = _exact_inf_norm(K[p]) + c2 = coeff * inf_norm_K_p_1 + if c1 + c2 <= tol * _exact_inf_norm(F): + break + c1 = c2 + X[k + i*d] = np.exp(k*h*mu) * F + return X, 2 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_interface.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..07c6bae191e9578f859286f3141f1b71102a2d56 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_interface.py @@ -0,0 +1,921 @@ +"""Abstract linear algebra library. + +This module defines a class hierarchy that implements a kind of "lazy" +matrix representation, called the ``LinearOperator``. It can be used to do +linear algebra with extremely large sparse or structured matrices, without +representing those explicitly in memory. Such matrices can be added, +multiplied, transposed, etc. + +As a motivating example, suppose you want have a matrix where almost all of +the elements have the value one. The standard sparse matrix representation +skips the storage of zeros, but not ones. By contrast, a LinearOperator is +able to represent such matrices efficiently. First, we need a compact way to +represent an all-ones matrix:: + + >>> import numpy as np + >>> from scipy.sparse.linalg._interface import LinearOperator + >>> class Ones(LinearOperator): + ... def __init__(self, shape): + ... super().__init__(dtype=None, shape=shape) + ... def _matvec(self, x): + ... return np.repeat(x.sum(), self.shape[0]) + +Instances of this class emulate ``np.ones(shape)``, but using a constant +amount of storage, independent of ``shape``. The ``_matvec`` method specifies +how this linear operator multiplies with (operates on) a vector. We can now +add this operator to a sparse matrix that stores only offsets from one:: + + >>> from scipy.sparse.linalg._interface import aslinearoperator + >>> from scipy.sparse import csr_array + >>> offsets = csr_array([[1, 0, 2], [0, -1, 0], [0, 0, 3]]) + >>> A = aslinearoperator(offsets) + Ones(offsets.shape) + >>> A.dot([1, 2, 3]) + array([13, 4, 15]) + +The result is the same as that given by its dense, explicitly-stored +counterpart:: + + >>> (np.ones(A.shape, A.dtype) + offsets.toarray()).dot([1, 2, 3]) + array([13, 4, 15]) + +Several algorithms in the ``scipy.sparse`` library are able to operate on +``LinearOperator`` instances. +""" + +import warnings + +import numpy as np + +from scipy.sparse import issparse +from scipy.sparse._sputils import isshape, isintlike, asmatrix, is_pydata_spmatrix + +__all__ = ['LinearOperator', 'aslinearoperator'] + + +class LinearOperator: + """Common interface for performing matrix vector products + + Many iterative methods (e.g. cg, gmres) do not need to know the + individual entries of a matrix to solve a linear system A@x=b. + Such solvers only require the computation of matrix vector + products, A@v where v is a dense vector. This class serves as + an abstract interface between iterative solvers and matrix-like + objects. + + To construct a concrete LinearOperator, either pass appropriate + callables to the constructor of this class, or subclass it. + + A subclass must implement either one of the methods ``_matvec`` + and ``_matmat``, and the attributes/properties ``shape`` (pair of + integers) and ``dtype`` (may be None). It may call the ``__init__`` + on this class to have these attributes validated. Implementing + ``_matvec`` automatically implements ``_matmat`` (using a naive + algorithm) and vice-versa. + + Optionally, a subclass may implement ``_rmatvec`` or ``_adjoint`` + to implement the Hermitian adjoint (conjugate transpose). As with + ``_matvec`` and ``_matmat``, implementing either ``_rmatvec`` or + ``_adjoint`` implements the other automatically. Implementing + ``_adjoint`` is preferable; ``_rmatvec`` is mostly there for + backwards compatibility. + + Parameters + ---------- + shape : tuple + Matrix dimensions (M, N). + matvec : callable f(v) + Returns returns A @ v. + rmatvec : callable f(v) + Returns A^H @ v, where A^H is the conjugate transpose of A. + matmat : callable f(V) + Returns A @ V, where V is a dense matrix with dimensions (N, K). + dtype : dtype + Data type of the matrix. + rmatmat : callable f(V) + Returns A^H @ V, where V is a dense matrix with dimensions (M, K). + + Attributes + ---------- + args : tuple + For linear operators describing products etc. of other linear + operators, the operands of the binary operation. + ndim : int + Number of dimensions (this is always 2) + + See Also + -------- + aslinearoperator : Construct LinearOperators + + Notes + ----- + The user-defined matvec() function must properly handle the case + where v has shape (N,) as well as the (N,1) case. The shape of + the return type is handled internally by LinearOperator. + + It is highly recommended to explicitly specify the `dtype`, otherwise + it is determined automatically at the cost of a single matvec application + on `int8` zero vector using the promoted `dtype` of the output. + Python `int` could be difficult to automatically cast to numpy integers + in the definition of the `matvec` so the determination may be inaccurate. + It is assumed that `matmat`, `rmatvec`, and `rmatmat` would result in + the same dtype of the output given an `int8` input as `matvec`. + + LinearOperator instances can also be multiplied, added with each + other and exponentiated, all lazily: the result of these operations + is always a new, composite LinearOperator, that defers linear + operations to the original operators and combines the results. + + More details regarding how to subclass a LinearOperator and several + examples of concrete LinearOperator instances can be found in the + external project `PyLops `_. + + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse.linalg import LinearOperator + >>> def mv(v): + ... return np.array([2*v[0], 3*v[1]]) + ... + >>> A = LinearOperator((2,2), matvec=mv) + >>> A + <2x2 _CustomLinearOperator with dtype=int8> + >>> A.matvec(np.ones(2)) + array([ 2., 3.]) + >>> A @ np.ones(2) + array([ 2., 3.]) + + """ + + ndim = 2 + # Necessary for right matmul with numpy arrays. + __array_ufunc__ = None + + def __new__(cls, *args, **kwargs): + if cls is LinearOperator: + # Operate as _CustomLinearOperator factory. + return super().__new__(_CustomLinearOperator) + else: + obj = super().__new__(cls) + + if (type(obj)._matvec == LinearOperator._matvec + and type(obj)._matmat == LinearOperator._matmat): + warnings.warn("LinearOperator subclass should implement" + " at least one of _matvec and _matmat.", + category=RuntimeWarning, stacklevel=2) + + return obj + + def __init__(self, dtype, shape): + """Initialize this LinearOperator. + + To be called by subclasses. ``dtype`` may be None; ``shape`` should + be convertible to a length-2 tuple. + """ + if dtype is not None: + dtype = np.dtype(dtype) + + shape = tuple(shape) + if not isshape(shape): + raise ValueError(f"invalid shape {shape!r} (must be 2-d)") + + self.dtype = dtype + self.shape = shape + + def _init_dtype(self): + """Determine the dtype by executing `matvec` on an `int8` test vector. + + In `np.promote_types` hierarchy, the type `int8` is the smallest, + so we call `matvec` on `int8` and use the promoted dtype of the output + to set the default `dtype` of the `LinearOperator`. + We assume that `matmat`, `rmatvec`, and `rmatmat` would result in + the same dtype of the output given an `int8` input as `matvec`. + + Called from subclasses at the end of the __init__ routine. + """ + if self.dtype is None: + v = np.zeros(self.shape[-1], dtype=np.int8) + try: + matvec_v = np.asarray(self.matvec(v)) + except OverflowError: + # Python large `int` promoted to `np.int64`or `np.int32` + self.dtype = np.dtype(int) + else: + self.dtype = matvec_v.dtype + + def _matmat(self, X): + """Default matrix-matrix multiplication handler. + + Falls back on the user-defined _matvec method, so defining that will + define matrix multiplication (though in a very suboptimal way). + """ + + return np.hstack([self.matvec(col.reshape(-1,1)) for col in X.T]) + + def _matvec(self, x): + """Default matrix-vector multiplication handler. + + If self is a linear operator of shape (M, N), then this method will + be called on a shape (N,) or (N, 1) ndarray, and should return a + shape (M,) or (M, 1) ndarray. + + This default implementation falls back on _matmat, so defining that + will define matrix-vector multiplication as well. + """ + return self.matmat(x.reshape(-1, 1)) + + def matvec(self, x): + """Matrix-vector multiplication. + + Performs the operation y=A@x where A is an MxN linear + operator and x is a column vector or 1-d array. + + Parameters + ---------- + x : {matrix, ndarray} + An array with shape (N,) or (N,1). + + Returns + ------- + y : {matrix, ndarray} + A matrix or ndarray with shape (M,) or (M,1) depending + on the type and shape of the x argument. + + Notes + ----- + This matvec wraps the user-specified matvec routine or overridden + _matvec method to ensure that y has the correct shape and type. + + """ + + x = np.asanyarray(x) + + M,N = self.shape + + if x.shape != (N,) and x.shape != (N,1): + raise ValueError('dimension mismatch') + + y = self._matvec(x) + + if isinstance(x, np.matrix): + y = asmatrix(y) + else: + y = np.asarray(y) + + if x.ndim == 1: + y = y.reshape(M) + elif x.ndim == 2: + y = y.reshape(M,1) + else: + raise ValueError('invalid shape returned by user-defined matvec()') + + return y + + def rmatvec(self, x): + """Adjoint matrix-vector multiplication. + + Performs the operation y = A^H @ x where A is an MxN linear + operator and x is a column vector or 1-d array. + + Parameters + ---------- + x : {matrix, ndarray} + An array with shape (M,) or (M,1). + + Returns + ------- + y : {matrix, ndarray} + A matrix or ndarray with shape (N,) or (N,1) depending + on the type and shape of the x argument. + + Notes + ----- + This rmatvec wraps the user-specified rmatvec routine or overridden + _rmatvec method to ensure that y has the correct shape and type. + + """ + + x = np.asanyarray(x) + + M,N = self.shape + + if x.shape != (M,) and x.shape != (M,1): + raise ValueError('dimension mismatch') + + y = self._rmatvec(x) + + if isinstance(x, np.matrix): + y = asmatrix(y) + else: + y = np.asarray(y) + + if x.ndim == 1: + y = y.reshape(N) + elif x.ndim == 2: + y = y.reshape(N,1) + else: + raise ValueError('invalid shape returned by user-defined rmatvec()') + + return y + + def _rmatvec(self, x): + """Default implementation of _rmatvec; defers to adjoint.""" + if type(self)._adjoint == LinearOperator._adjoint: + # _adjoint not overridden, prevent infinite recursion + if (hasattr(self, "_rmatmat") + and type(self)._rmatmat != LinearOperator._rmatmat): + # Try to use _rmatmat as a fallback + return self._rmatmat(x.reshape(-1, 1)).reshape(-1) + raise NotImplementedError + else: + return self.H.matvec(x) + + def matmat(self, X): + """Matrix-matrix multiplication. + + Performs the operation y=A@X where A is an MxN linear + operator and X dense N*K matrix or ndarray. + + Parameters + ---------- + X : {matrix, ndarray} + An array with shape (N,K). + + Returns + ------- + Y : {matrix, ndarray} + A matrix or ndarray with shape (M,K) depending on + the type of the X argument. + + Notes + ----- + This matmat wraps any user-specified matmat routine or overridden + _matmat method to ensure that y has the correct type. + + """ + if not (issparse(X) or is_pydata_spmatrix(X)): + X = np.asanyarray(X) + + if X.ndim != 2: + raise ValueError(f'expected 2-d ndarray or matrix, not {X.ndim}-d') + + if X.shape[0] != self.shape[1]: + raise ValueError(f'dimension mismatch: {self.shape}, {X.shape}') + + try: + Y = self._matmat(X) + except Exception as e: + if issparse(X) or is_pydata_spmatrix(X): + raise TypeError( + "Unable to multiply a LinearOperator with a sparse matrix." + " Wrap the matrix in aslinearoperator first." + ) from e + raise + + if isinstance(Y, np.matrix): + Y = asmatrix(Y) + + return Y + + def rmatmat(self, X): + """Adjoint matrix-matrix multiplication. + + Performs the operation y = A^H @ x where A is an MxN linear + operator and x is a column vector or 1-d array, or 2-d array. + The default implementation defers to the adjoint. + + Parameters + ---------- + X : {matrix, ndarray} + A matrix or 2D array. + + Returns + ------- + Y : {matrix, ndarray} + A matrix or 2D array depending on the type of the input. + + Notes + ----- + This rmatmat wraps the user-specified rmatmat routine. + + """ + if not (issparse(X) or is_pydata_spmatrix(X)): + X = np.asanyarray(X) + + if X.ndim != 2: + raise ValueError('expected 2-d ndarray or matrix, not %d-d' + % X.ndim) + + if X.shape[0] != self.shape[0]: + raise ValueError(f'dimension mismatch: {self.shape}, {X.shape}') + + try: + Y = self._rmatmat(X) + except Exception as e: + if issparse(X) or is_pydata_spmatrix(X): + raise TypeError( + "Unable to multiply a LinearOperator with a sparse matrix." + " Wrap the matrix in aslinearoperator() first." + ) from e + raise + + if isinstance(Y, np.matrix): + Y = asmatrix(Y) + return Y + + def _rmatmat(self, X): + """Default implementation of _rmatmat defers to rmatvec or adjoint.""" + if type(self)._adjoint == LinearOperator._adjoint: + return np.hstack([self.rmatvec(col.reshape(-1, 1)) for col in X.T]) + else: + return self.H.matmat(X) + + def __call__(self, x): + return self@x + + def __mul__(self, x): + return self.dot(x) + + def __truediv__(self, other): + if not np.isscalar(other): + raise ValueError("Can only divide a linear operator by a scalar.") + + return _ScaledLinearOperator(self, 1.0/other) + + def dot(self, x): + """Matrix-matrix or matrix-vector multiplication. + + Parameters + ---------- + x : array_like + 1-d or 2-d array, representing a vector or matrix. + + Returns + ------- + Ax : array + 1-d or 2-d array (depending on the shape of x) that represents + the result of applying this linear operator on x. + + """ + if isinstance(x, LinearOperator): + return _ProductLinearOperator(self, x) + elif np.isscalar(x): + return _ScaledLinearOperator(self, x) + else: + if not issparse(x) and not is_pydata_spmatrix(x): + # Sparse matrices shouldn't be converted to numpy arrays. + x = np.asarray(x) + + if x.ndim == 1 or x.ndim == 2 and x.shape[1] == 1: + return self.matvec(x) + elif x.ndim == 2: + return self.matmat(x) + else: + raise ValueError(f'expected 1-d or 2-d array or matrix, got {x!r}') + + def __matmul__(self, other): + if np.isscalar(other): + raise ValueError("Scalar operands are not allowed, " + "use '*' instead") + return self.__mul__(other) + + def __rmatmul__(self, other): + if np.isscalar(other): + raise ValueError("Scalar operands are not allowed, " + "use '*' instead") + return self.__rmul__(other) + + def __rmul__(self, x): + if np.isscalar(x): + return _ScaledLinearOperator(self, x) + else: + return self._rdot(x) + + def _rdot(self, x): + """Matrix-matrix or matrix-vector multiplication from the right. + + Parameters + ---------- + x : array_like + 1-d or 2-d array, representing a vector or matrix. + + Returns + ------- + xA : array + 1-d or 2-d array (depending on the shape of x) that represents + the result of applying this linear operator on x from the right. + + Notes + ----- + This is copied from dot to implement right multiplication. + """ + if isinstance(x, LinearOperator): + return _ProductLinearOperator(x, self) + elif np.isscalar(x): + return _ScaledLinearOperator(self, x) + else: + if not issparse(x) and not is_pydata_spmatrix(x): + # Sparse matrices shouldn't be converted to numpy arrays. + x = np.asarray(x) + + # We use transpose instead of rmatvec/rmatmat to avoid + # unnecessary complex conjugation if possible. + if x.ndim == 1 or x.ndim == 2 and x.shape[0] == 1: + return self.T.matvec(x.T).T + elif x.ndim == 2: + return self.T.matmat(x.T).T + else: + raise ValueError(f'expected 1-d or 2-d array or matrix, got {x!r}') + + def __pow__(self, p): + if np.isscalar(p): + return _PowerLinearOperator(self, p) + else: + return NotImplemented + + def __add__(self, x): + if isinstance(x, LinearOperator): + return _SumLinearOperator(self, x) + else: + return NotImplemented + + def __neg__(self): + return _ScaledLinearOperator(self, -1) + + def __sub__(self, x): + return self.__add__(-x) + + def __repr__(self): + M,N = self.shape + if self.dtype is None: + dt = 'unspecified dtype' + else: + dt = 'dtype=' + str(self.dtype) + + return '<%dx%d %s with %s>' % (M, N, self.__class__.__name__, dt) + + def adjoint(self): + """Hermitian adjoint. + + Returns the Hermitian adjoint of self, aka the Hermitian + conjugate or Hermitian transpose. For a complex matrix, the + Hermitian adjoint is equal to the conjugate transpose. + + Can be abbreviated self.H instead of self.adjoint(). + + Returns + ------- + A_H : LinearOperator + Hermitian adjoint of self. + """ + return self._adjoint() + + H = property(adjoint) + + def transpose(self): + """Transpose this linear operator. + + Returns a LinearOperator that represents the transpose of this one. + Can be abbreviated self.T instead of self.transpose(). + """ + return self._transpose() + + T = property(transpose) + + def _adjoint(self): + """Default implementation of _adjoint; defers to rmatvec.""" + return _AdjointLinearOperator(self) + + def _transpose(self): + """ Default implementation of _transpose; defers to rmatvec + conj""" + return _TransposedLinearOperator(self) + + +class _CustomLinearOperator(LinearOperator): + """Linear operator defined in terms of user-specified operations.""" + + def __init__(self, shape, matvec, rmatvec=None, matmat=None, + dtype=None, rmatmat=None): + super().__init__(dtype, shape) + + self.args = () + + self.__matvec_impl = matvec + self.__rmatvec_impl = rmatvec + self.__rmatmat_impl = rmatmat + self.__matmat_impl = matmat + + self._init_dtype() + + def _matmat(self, X): + if self.__matmat_impl is not None: + return self.__matmat_impl(X) + else: + return super()._matmat(X) + + def _matvec(self, x): + return self.__matvec_impl(x) + + def _rmatvec(self, x): + func = self.__rmatvec_impl + if func is None: + raise NotImplementedError("rmatvec is not defined") + return self.__rmatvec_impl(x) + + def _rmatmat(self, X): + if self.__rmatmat_impl is not None: + return self.__rmatmat_impl(X) + else: + return super()._rmatmat(X) + + def _adjoint(self): + return _CustomLinearOperator(shape=(self.shape[1], self.shape[0]), + matvec=self.__rmatvec_impl, + rmatvec=self.__matvec_impl, + matmat=self.__rmatmat_impl, + rmatmat=self.__matmat_impl, + dtype=self.dtype) + + +class _AdjointLinearOperator(LinearOperator): + """Adjoint of arbitrary Linear Operator""" + + def __init__(self, A): + shape = (A.shape[1], A.shape[0]) + super().__init__(dtype=A.dtype, shape=shape) + self.A = A + self.args = (A,) + + def _matvec(self, x): + return self.A._rmatvec(x) + + def _rmatvec(self, x): + return self.A._matvec(x) + + def _matmat(self, x): + return self.A._rmatmat(x) + + def _rmatmat(self, x): + return self.A._matmat(x) + +class _TransposedLinearOperator(LinearOperator): + """Transposition of arbitrary Linear Operator""" + + def __init__(self, A): + shape = (A.shape[1], A.shape[0]) + super().__init__(dtype=A.dtype, shape=shape) + self.A = A + self.args = (A,) + + def _matvec(self, x): + # NB. np.conj works also on sparse matrices + return np.conj(self.A._rmatvec(np.conj(x))) + + def _rmatvec(self, x): + return np.conj(self.A._matvec(np.conj(x))) + + def _matmat(self, x): + # NB. np.conj works also on sparse matrices + return np.conj(self.A._rmatmat(np.conj(x))) + + def _rmatmat(self, x): + return np.conj(self.A._matmat(np.conj(x))) + +def _get_dtype(operators, dtypes=None): + if dtypes is None: + dtypes = [] + for obj in operators: + if obj is not None and hasattr(obj, 'dtype'): + dtypes.append(obj.dtype) + return np.result_type(*dtypes) + + +class _SumLinearOperator(LinearOperator): + def __init__(self, A, B): + if not isinstance(A, LinearOperator) or \ + not isinstance(B, LinearOperator): + raise ValueError('both operands have to be a LinearOperator') + if A.shape != B.shape: + raise ValueError(f'cannot add {A} and {B}: shape mismatch') + self.args = (A, B) + super().__init__(_get_dtype([A, B]), A.shape) + + def _matvec(self, x): + return self.args[0].matvec(x) + self.args[1].matvec(x) + + def _rmatvec(self, x): + return self.args[0].rmatvec(x) + self.args[1].rmatvec(x) + + def _rmatmat(self, x): + return self.args[0].rmatmat(x) + self.args[1].rmatmat(x) + + def _matmat(self, x): + return self.args[0].matmat(x) + self.args[1].matmat(x) + + def _adjoint(self): + A, B = self.args + return A.H + B.H + + +class _ProductLinearOperator(LinearOperator): + def __init__(self, A, B): + if not isinstance(A, LinearOperator) or \ + not isinstance(B, LinearOperator): + raise ValueError('both operands have to be a LinearOperator') + if A.shape[1] != B.shape[0]: + raise ValueError(f'cannot multiply {A} and {B}: shape mismatch') + super().__init__(_get_dtype([A, B]), + (A.shape[0], B.shape[1])) + self.args = (A, B) + + def _matvec(self, x): + return self.args[0].matvec(self.args[1].matvec(x)) + + def _rmatvec(self, x): + return self.args[1].rmatvec(self.args[0].rmatvec(x)) + + def _rmatmat(self, x): + return self.args[1].rmatmat(self.args[0].rmatmat(x)) + + def _matmat(self, x): + return self.args[0].matmat(self.args[1].matmat(x)) + + def _adjoint(self): + A, B = self.args + return B.H @ A.H + + +class _ScaledLinearOperator(LinearOperator): + def __init__(self, A, alpha): + if not isinstance(A, LinearOperator): + raise ValueError('LinearOperator expected as A') + if not np.isscalar(alpha): + raise ValueError('scalar expected as alpha') + if isinstance(A, _ScaledLinearOperator): + A, alpha_original = A.args + # Avoid in-place multiplication so that we don't accidentally mutate + # the original prefactor. + alpha = alpha * alpha_original + + dtype = _get_dtype([A], [type(alpha)]) + super().__init__(dtype, A.shape) + self.args = (A, alpha) + # Note: args[1] is alpha (a scalar), so use `*` below, not `@` + + def _matvec(self, x): + return self.args[1] * self.args[0].matvec(x) + + def _rmatvec(self, x): + return np.conj(self.args[1]) * self.args[0].rmatvec(x) + + def _rmatmat(self, x): + return np.conj(self.args[1]) * self.args[0].rmatmat(x) + + def _matmat(self, x): + return self.args[1] * self.args[0].matmat(x) + + def _adjoint(self): + A, alpha = self.args + return A.H * np.conj(alpha) + + +class _PowerLinearOperator(LinearOperator): + def __init__(self, A, p): + if not isinstance(A, LinearOperator): + raise ValueError('LinearOperator expected as A') + if A.shape[0] != A.shape[1]: + raise ValueError(f'square LinearOperator expected, got {A!r}') + if not isintlike(p) or p < 0: + raise ValueError('non-negative integer expected as p') + + super().__init__(_get_dtype([A]), A.shape) + self.args = (A, p) + + def _power(self, fun, x): + res = np.array(x, copy=True) + for i in range(self.args[1]): + res = fun(res) + return res + + def _matvec(self, x): + return self._power(self.args[0].matvec, x) + + def _rmatvec(self, x): + return self._power(self.args[0].rmatvec, x) + + def _rmatmat(self, x): + return self._power(self.args[0].rmatmat, x) + + def _matmat(self, x): + return self._power(self.args[0].matmat, x) + + def _adjoint(self): + A, p = self.args + return A.H ** p + + +class MatrixLinearOperator(LinearOperator): + def __init__(self, A): + super().__init__(A.dtype, A.shape) + self.A = A + self.__adj = None + self.args = (A,) + + def _matmat(self, X): + return self.A.dot(X) + + def _adjoint(self): + if self.__adj is None: + self.__adj = _AdjointMatrixOperator(self.A) + return self.__adj + + +class _AdjointMatrixOperator(MatrixLinearOperator): + def __init__(self, adjoint_array): + self.A = adjoint_array.T.conj() + self.args = (adjoint_array,) + self.shape = adjoint_array.shape[1], adjoint_array.shape[0] + + @property + def dtype(self): + return self.args[0].dtype + + def _adjoint(self): + return MatrixLinearOperator(self.args[0]) + + +class IdentityOperator(LinearOperator): + def __init__(self, shape, dtype=None): + super().__init__(dtype, shape) + + def _matvec(self, x): + return x + + def _rmatvec(self, x): + return x + + def _rmatmat(self, x): + return x + + def _matmat(self, x): + return x + + def _adjoint(self): + return self + + +def aslinearoperator(A): + """Return A as a LinearOperator. + + 'A' may be any of the following types: + - ndarray + - matrix + - sparse array (e.g. csr_array, lil_array, etc.) + - LinearOperator + - An object with .shape and .matvec attributes + + See the LinearOperator documentation for additional information. + + Notes + ----- + If 'A' has no .dtype attribute, the data type is determined by calling + :func:`LinearOperator.matvec()` - set the .dtype attribute to prevent this + call upon the linear operator creation. + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse.linalg import aslinearoperator + >>> M = np.array([[1,2,3],[4,5,6]], dtype=np.int32) + >>> aslinearoperator(M) + <2x3 MatrixLinearOperator with dtype=int32> + """ + if isinstance(A, LinearOperator): + return A + + elif isinstance(A, np.ndarray) or isinstance(A, np.matrix): + if A.ndim > 2: + raise ValueError('array must have ndim <= 2') + A = np.atleast_2d(np.asarray(A)) + return MatrixLinearOperator(A) + + elif issparse(A) or is_pydata_spmatrix(A): + return MatrixLinearOperator(A) + + else: + if hasattr(A, 'shape') and hasattr(A, 'matvec'): + rmatvec = None + rmatmat = None + dtype = None + + if hasattr(A, 'rmatvec'): + rmatvec = A.rmatvec + if hasattr(A, 'rmatmat'): + rmatmat = A.rmatmat + if hasattr(A, 'dtype'): + dtype = A.dtype + return LinearOperator(A.shape, A.matvec, rmatvec=rmatvec, + rmatmat=rmatmat, dtype=dtype) + + else: + raise TypeError('type not understood') diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_matfuncs.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_matfuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..c638885c2365d5157eb983a3c9ba73e6590bec3e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_matfuncs.py @@ -0,0 +1,940 @@ +""" +Sparse matrix functions +""" + +# +# Authors: Travis Oliphant, March 2002 +# Anthony Scopatz, August 2012 (Sparse Updates) +# Jake Vanderplas, August 2012 (Sparse Updates) +# + +__all__ = ['expm', 'inv', 'matrix_power'] + +import numpy as np +from scipy.linalg._basic import solve, solve_triangular + +from scipy.sparse._base import issparse +from scipy.sparse.linalg import spsolve +from scipy.sparse._sputils import is_pydata_spmatrix, isintlike + +import scipy.sparse +import scipy.sparse.linalg +from scipy.sparse.linalg._interface import LinearOperator +from scipy.sparse._construct import eye_array + +from ._expm_multiply import _ident_like, _exact_1_norm as _onenorm + + +UPPER_TRIANGULAR = 'upper_triangular' + + +def inv(A): + """ + Compute the inverse of a sparse arrays + + Parameters + ---------- + A : (M, M) sparse arrays + square matrix to be inverted + + Returns + ------- + Ainv : (M, M) sparse arrays + inverse of `A` + + Notes + ----- + This computes the sparse inverse of `A`. If the inverse of `A` is expected + to be non-sparse, it will likely be faster to convert `A` to dense and use + `scipy.linalg.inv`. + + Examples + -------- + >>> from scipy.sparse import csc_array + >>> from scipy.sparse.linalg import inv + >>> A = csc_array([[1., 0.], [1., 2.]]) + >>> Ainv = inv(A) + >>> Ainv + + >>> A.dot(Ainv) + + >>> A.dot(Ainv).toarray() + array([[ 1., 0.], + [ 0., 1.]]) + + .. versionadded:: 0.12.0 + + """ + # Check input + if not (issparse(A) or is_pydata_spmatrix(A)): + raise TypeError('Input must be a sparse arrays') + + # Use sparse direct solver to solve "AX = I" accurately + I = _ident_like(A) + Ainv = spsolve(A, I) + return Ainv + + +def _onenorm_matrix_power_nnm(A, p): + """ + Compute the 1-norm of a non-negative integer power of a non-negative matrix. + + Parameters + ---------- + A : a square ndarray or matrix or sparse arrays + Input matrix with non-negative entries. + p : non-negative integer + The power to which the matrix is to be raised. + + Returns + ------- + out : float + The 1-norm of the matrix power p of A. + + """ + # Check input + if int(p) != p or p < 0: + raise ValueError('expected non-negative integer p') + p = int(p) + if len(A.shape) != 2 or A.shape[0] != A.shape[1]: + raise ValueError('expected A to be like a square matrix') + + # Explicitly make a column vector so that this works when A is a + # numpy matrix (in addition to ndarray and sparse arrays). + v = np.ones((A.shape[0], 1), dtype=float) + M = A.T + for i in range(p): + v = M.dot(v) + return np.max(v) + + +def _is_upper_triangular(A): + # This function could possibly be of wider interest. + if issparse(A): + lower_part = scipy.sparse.tril(A, -1) + # Check structural upper triangularity, + # then coincidental upper triangularity if needed. + return lower_part.nnz == 0 or lower_part.count_nonzero() == 0 + elif is_pydata_spmatrix(A): + import sparse + lower_part = sparse.tril(A, -1) + return lower_part.nnz == 0 + else: + return not np.tril(A, -1).any() + + +def _smart_matrix_product(A, B, alpha=None, structure=None): + """ + A matrix product that knows about sparse and structured matrices. + + Parameters + ---------- + A : 2d ndarray + First matrix. + B : 2d ndarray + Second matrix. + alpha : float + The matrix product will be scaled by this constant. + structure : str, optional + A string describing the structure of both matrices `A` and `B`. + Only `upper_triangular` is currently supported. + + Returns + ------- + M : 2d ndarray + Matrix product of A and B. + + """ + if len(A.shape) != 2: + raise ValueError('expected A to be a rectangular matrix') + if len(B.shape) != 2: + raise ValueError('expected B to be a rectangular matrix') + f = None + if structure == UPPER_TRIANGULAR: + if (not issparse(A) and not issparse(B) + and not is_pydata_spmatrix(A) and not is_pydata_spmatrix(B)): + f, = scipy.linalg.get_blas_funcs(('trmm',), (A, B)) + if f is not None: + if alpha is None: + alpha = 1. + out = f(alpha, A, B) + else: + if alpha is None: + out = A.dot(B) + else: + out = alpha * A.dot(B) + return out + + +class MatrixPowerOperator(LinearOperator): + + def __init__(self, A, p, structure=None): + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError('expected A to be like a square matrix') + if p < 0: + raise ValueError('expected p to be a non-negative integer') + self._A = A + self._p = p + self._structure = structure + self.dtype = A.dtype + self.ndim = A.ndim + self.shape = A.shape + + def _matvec(self, x): + for i in range(self._p): + x = self._A.dot(x) + return x + + def _rmatvec(self, x): + A_T = self._A.T + x = x.ravel() + for i in range(self._p): + x = A_T.dot(x) + return x + + def _matmat(self, X): + for i in range(self._p): + X = _smart_matrix_product(self._A, X, structure=self._structure) + return X + + @property + def T(self): + return MatrixPowerOperator(self._A.T, self._p) + + +class ProductOperator(LinearOperator): + """ + For now, this is limited to products of multiple square matrices. + """ + + def __init__(self, *args, **kwargs): + self._structure = kwargs.get('structure', None) + for A in args: + if len(A.shape) != 2 or A.shape[0] != A.shape[1]: + raise ValueError( + 'For now, the ProductOperator implementation is ' + 'limited to the product of multiple square matrices.') + if args: + n = args[0].shape[0] + for A in args: + for d in A.shape: + if d != n: + raise ValueError( + 'The square matrices of the ProductOperator ' + 'must all have the same shape.') + self.shape = (n, n) + self.ndim = len(self.shape) + self.dtype = np.result_type(*[x.dtype for x in args]) + self._operator_sequence = args + + def _matvec(self, x): + for A in reversed(self._operator_sequence): + x = A.dot(x) + return x + + def _rmatvec(self, x): + x = x.ravel() + for A in self._operator_sequence: + x = A.T.dot(x) + return x + + def _matmat(self, X): + for A in reversed(self._operator_sequence): + X = _smart_matrix_product(A, X, structure=self._structure) + return X + + @property + def T(self): + T_args = [A.T for A in reversed(self._operator_sequence)] + return ProductOperator(*T_args) + + +def _onenormest_matrix_power(A, p, + t=2, itmax=5, compute_v=False, compute_w=False, structure=None): + """ + Efficiently estimate the 1-norm of A^p. + + Parameters + ---------- + A : ndarray + Matrix whose 1-norm of a power is to be computed. + p : int + Non-negative integer power. + t : int, optional + A positive parameter controlling the tradeoff between + accuracy versus time and memory usage. + Larger values take longer and use more memory + but give more accurate output. + itmax : int, optional + Use at most this many iterations. + compute_v : bool, optional + Request a norm-maximizing linear operator input vector if True. + compute_w : bool, optional + Request a norm-maximizing linear operator output vector if True. + + Returns + ------- + est : float + An underestimate of the 1-norm of the sparse arrays. + v : ndarray, optional + The vector such that ||Av||_1 == est*||v||_1. + It can be thought of as an input to the linear operator + that gives an output with particularly large norm. + w : ndarray, optional + The vector Av which has relatively large 1-norm. + It can be thought of as an output of the linear operator + that is relatively large in norm compared to the input. + + """ + return scipy.sparse.linalg.onenormest( + MatrixPowerOperator(A, p, structure=structure)) + + +def _onenormest_product(operator_seq, + t=2, itmax=5, compute_v=False, compute_w=False, structure=None): + """ + Efficiently estimate the 1-norm of the matrix product of the args. + + Parameters + ---------- + operator_seq : linear operator sequence + Matrices whose 1-norm of product is to be computed. + t : int, optional + A positive parameter controlling the tradeoff between + accuracy versus time and memory usage. + Larger values take longer and use more memory + but give more accurate output. + itmax : int, optional + Use at most this many iterations. + compute_v : bool, optional + Request a norm-maximizing linear operator input vector if True. + compute_w : bool, optional + Request a norm-maximizing linear operator output vector if True. + structure : str, optional + A string describing the structure of all operators. + Only `upper_triangular` is currently supported. + + Returns + ------- + est : float + An underestimate of the 1-norm of the sparse arrays. + v : ndarray, optional + The vector such that ||Av||_1 == est*||v||_1. + It can be thought of as an input to the linear operator + that gives an output with particularly large norm. + w : ndarray, optional + The vector Av which has relatively large 1-norm. + It can be thought of as an output of the linear operator + that is relatively large in norm compared to the input. + + """ + return scipy.sparse.linalg.onenormest( + ProductOperator(*operator_seq, structure=structure)) + + +class _ExpmPadeHelper: + """ + Help lazily evaluate a matrix exponential. + + The idea is to not do more work than we need for high expm precision, + so we lazily compute matrix powers and store or precompute + other properties of the matrix. + + """ + + def __init__(self, A, structure=None, use_exact_onenorm=False): + """ + Initialize the object. + + Parameters + ---------- + A : a dense or sparse square numpy matrix or ndarray + The matrix to be exponentiated. + structure : str, optional + A string describing the structure of matrix `A`. + Only `upper_triangular` is currently supported. + use_exact_onenorm : bool, optional + If True then only the exact one-norm of matrix powers and products + will be used. Otherwise, the one-norm of powers and products + may initially be estimated. + """ + self.A = A + self._A2 = None + self._A4 = None + self._A6 = None + self._A8 = None + self._A10 = None + self._d4_exact = None + self._d6_exact = None + self._d8_exact = None + self._d10_exact = None + self._d4_approx = None + self._d6_approx = None + self._d8_approx = None + self._d10_approx = None + self.ident = _ident_like(A) + self.structure = structure + self.use_exact_onenorm = use_exact_onenorm + + @property + def A2(self): + if self._A2 is None: + self._A2 = _smart_matrix_product( + self.A, self.A, structure=self.structure) + return self._A2 + + @property + def A4(self): + if self._A4 is None: + self._A4 = _smart_matrix_product( + self.A2, self.A2, structure=self.structure) + return self._A4 + + @property + def A6(self): + if self._A6 is None: + self._A6 = _smart_matrix_product( + self.A4, self.A2, structure=self.structure) + return self._A6 + + @property + def A8(self): + if self._A8 is None: + self._A8 = _smart_matrix_product( + self.A6, self.A2, structure=self.structure) + return self._A8 + + @property + def A10(self): + if self._A10 is None: + self._A10 = _smart_matrix_product( + self.A4, self.A6, structure=self.structure) + return self._A10 + + @property + def d4_tight(self): + if self._d4_exact is None: + self._d4_exact = _onenorm(self.A4)**(1/4.) + return self._d4_exact + + @property + def d6_tight(self): + if self._d6_exact is None: + self._d6_exact = _onenorm(self.A6)**(1/6.) + return self._d6_exact + + @property + def d8_tight(self): + if self._d8_exact is None: + self._d8_exact = _onenorm(self.A8)**(1/8.) + return self._d8_exact + + @property + def d10_tight(self): + if self._d10_exact is None: + self._d10_exact = _onenorm(self.A10)**(1/10.) + return self._d10_exact + + @property + def d4_loose(self): + if self.use_exact_onenorm: + return self.d4_tight + if self._d4_exact is not None: + return self._d4_exact + else: + if self._d4_approx is None: + self._d4_approx = _onenormest_matrix_power(self.A2, 2, + structure=self.structure)**(1/4.) + return self._d4_approx + + @property + def d6_loose(self): + if self.use_exact_onenorm: + return self.d6_tight + if self._d6_exact is not None: + return self._d6_exact + else: + if self._d6_approx is None: + self._d6_approx = _onenormest_matrix_power(self.A2, 3, + structure=self.structure)**(1/6.) + return self._d6_approx + + @property + def d8_loose(self): + if self.use_exact_onenorm: + return self.d8_tight + if self._d8_exact is not None: + return self._d8_exact + else: + if self._d8_approx is None: + self._d8_approx = _onenormest_matrix_power(self.A4, 2, + structure=self.structure)**(1/8.) + return self._d8_approx + + @property + def d10_loose(self): + if self.use_exact_onenorm: + return self.d10_tight + if self._d10_exact is not None: + return self._d10_exact + else: + if self._d10_approx is None: + self._d10_approx = _onenormest_product((self.A4, self.A6), + structure=self.structure)**(1/10.) + return self._d10_approx + + def pade3(self): + b = (120., 60., 12., 1.) + U = _smart_matrix_product(self.A, + b[3]*self.A2 + b[1]*self.ident, + structure=self.structure) + V = b[2]*self.A2 + b[0]*self.ident + return U, V + + def pade5(self): + b = (30240., 15120., 3360., 420., 30., 1.) + U = _smart_matrix_product(self.A, + b[5]*self.A4 + b[3]*self.A2 + b[1]*self.ident, + structure=self.structure) + V = b[4]*self.A4 + b[2]*self.A2 + b[0]*self.ident + return U, V + + def pade7(self): + b = (17297280., 8648640., 1995840., 277200., 25200., 1512., 56., 1.) + U = _smart_matrix_product(self.A, + b[7]*self.A6 + b[5]*self.A4 + b[3]*self.A2 + b[1]*self.ident, + structure=self.structure) + V = b[6]*self.A6 + b[4]*self.A4 + b[2]*self.A2 + b[0]*self.ident + return U, V + + def pade9(self): + b = (17643225600., 8821612800., 2075673600., 302702400., 30270240., + 2162160., 110880., 3960., 90., 1.) + U = _smart_matrix_product(self.A, + (b[9]*self.A8 + b[7]*self.A6 + b[5]*self.A4 + + b[3]*self.A2 + b[1]*self.ident), + structure=self.structure) + V = (b[8]*self.A8 + b[6]*self.A6 + b[4]*self.A4 + + b[2]*self.A2 + b[0]*self.ident) + return U, V + + def pade13_scaled(self, s): + b = (64764752532480000., 32382376266240000., 7771770303897600., + 1187353796428800., 129060195264000., 10559470521600., + 670442572800., 33522128640., 1323241920., 40840800., 960960., + 16380., 182., 1.) + B = self.A * 2**-s + B2 = self.A2 * 2**(-2*s) + B4 = self.A4 * 2**(-4*s) + B6 = self.A6 * 2**(-6*s) + U2 = _smart_matrix_product(B6, + b[13]*B6 + b[11]*B4 + b[9]*B2, + structure=self.structure) + U = _smart_matrix_product(B, + (U2 + b[7]*B6 + b[5]*B4 + + b[3]*B2 + b[1]*self.ident), + structure=self.structure) + V2 = _smart_matrix_product(B6, + b[12]*B6 + b[10]*B4 + b[8]*B2, + structure=self.structure) + V = V2 + b[6]*B6 + b[4]*B4 + b[2]*B2 + b[0]*self.ident + return U, V + + +def expm(A): + """ + Compute the matrix exponential using Pade approximation. + + Parameters + ---------- + A : (M,M) array_like or sparse array + 2D Array or Matrix (sparse or dense) to be exponentiated + + Returns + ------- + expA : (M,M) ndarray + Matrix exponential of `A` + + Notes + ----- + This is algorithm (6.1) which is a simplification of algorithm (5.1). + + .. versionadded:: 0.12.0 + + References + ---------- + .. [1] Awad H. Al-Mohy and Nicholas J. Higham (2009) + "A New Scaling and Squaring Algorithm for the Matrix Exponential." + SIAM Journal on Matrix Analysis and Applications. + 31 (3). pp. 970-989. ISSN 1095-7162 + + Examples + -------- + >>> from scipy.sparse import csc_array + >>> from scipy.sparse.linalg import expm + >>> A = csc_array([[1, 0, 0], [0, 2, 0], [0, 0, 3]]) + >>> A.toarray() + array([[1, 0, 0], + [0, 2, 0], + [0, 0, 3]], dtype=int64) + >>> Aexp = expm(A) + >>> Aexp + + >>> Aexp.toarray() + array([[ 2.71828183, 0. , 0. ], + [ 0. , 7.3890561 , 0. ], + [ 0. , 0. , 20.08553692]]) + """ + return _expm(A, use_exact_onenorm='auto') + + +def _expm(A, use_exact_onenorm): + # Core of expm, separated to allow testing exact and approximate + # algorithms. + + # Avoid indiscriminate asarray() to allow sparse or other strange arrays. + if isinstance(A, (list, tuple, np.matrix)): + A = np.asarray(A) + if len(A.shape) != 2 or A.shape[0] != A.shape[1]: + raise ValueError('expected a square matrix') + + # gracefully handle size-0 input, + # carefully handling sparse scenario + if A.shape == (0, 0): + out = np.zeros([0, 0], dtype=A.dtype) + if issparse(A) or is_pydata_spmatrix(A): + return A.__class__(out) + return out + + # Trivial case + if A.shape == (1, 1): + out = [[np.exp(A[0, 0])]] + + # Avoid indiscriminate casting to ndarray to + # allow for sparse or other strange arrays + if issparse(A) or is_pydata_spmatrix(A): + return A.__class__(out) + + return np.array(out) + + # Ensure input is of float type, to avoid integer overflows etc. + if ((isinstance(A, np.ndarray) or issparse(A) or is_pydata_spmatrix(A)) + and not np.issubdtype(A.dtype, np.inexact)): + A = A.astype(float) + + # Detect upper triangularity. + structure = UPPER_TRIANGULAR if _is_upper_triangular(A) else None + + if use_exact_onenorm == "auto": + # Hardcode a matrix order threshold for exact vs. estimated one-norms. + use_exact_onenorm = A.shape[0] < 200 + + # Track functions of A to help compute the matrix exponential. + h = _ExpmPadeHelper( + A, structure=structure, use_exact_onenorm=use_exact_onenorm) + + # Try Pade order 3. + eta_1 = max(h.d4_loose, h.d6_loose) + if eta_1 < 1.495585217958292e-002 and _ell(h.A, 3) == 0: + U, V = h.pade3() + return _solve_P_Q(U, V, structure=structure) + + # Try Pade order 5. + eta_2 = max(h.d4_tight, h.d6_loose) + if eta_2 < 2.539398330063230e-001 and _ell(h.A, 5) == 0: + U, V = h.pade5() + return _solve_P_Q(U, V, structure=structure) + + # Try Pade orders 7 and 9. + eta_3 = max(h.d6_tight, h.d8_loose) + if eta_3 < 9.504178996162932e-001 and _ell(h.A, 7) == 0: + U, V = h.pade7() + return _solve_P_Q(U, V, structure=structure) + if eta_3 < 2.097847961257068e+000 and _ell(h.A, 9) == 0: + U, V = h.pade9() + return _solve_P_Q(U, V, structure=structure) + + # Use Pade order 13. + eta_4 = max(h.d8_loose, h.d10_loose) + eta_5 = min(eta_3, eta_4) + theta_13 = 4.25 + + # Choose smallest s>=0 such that 2**(-s) eta_5 <= theta_13 + if eta_5 == 0: + # Nilpotent special case + s = 0 + else: + s = max(int(np.ceil(np.log2(eta_5 / theta_13))), 0) + s = s + _ell(2**-s * h.A, 13) + U, V = h.pade13_scaled(s) + X = _solve_P_Q(U, V, structure=structure) + if structure == UPPER_TRIANGULAR: + # Invoke Code Fragment 2.1. + X = _fragment_2_1(X, h.A, s) + else: + # X = r_13(A)^(2^s) by repeated squaring. + for i in range(s): + X = X.dot(X) + return X + + +def _solve_P_Q(U, V, structure=None): + """ + A helper function for expm_2009. + + Parameters + ---------- + U : ndarray + Pade numerator. + V : ndarray + Pade denominator. + structure : str, optional + A string describing the structure of both matrices `U` and `V`. + Only `upper_triangular` is currently supported. + + Notes + ----- + The `structure` argument is inspired by similar args + for theano and cvxopt functions. + + """ + P = U + V + Q = -U + V + if issparse(U) or is_pydata_spmatrix(U): + return spsolve(Q, P) + elif structure is None: + return solve(Q, P) + elif structure == UPPER_TRIANGULAR: + return solve_triangular(Q, P) + else: + raise ValueError('unsupported matrix structure: ' + str(structure)) + + +def _exp_sinch(a, x): + """ + Stably evaluate exp(a)*sinh(x)/x + + Notes + ----- + The strategy of falling back to a sixth order Taylor expansion + was suggested by the Spallation Neutron Source docs + which was found on the internet by google search. + http://www.ornl.gov/~t6p/resources/xal/javadoc/gov/sns/tools/math/ElementaryFunction.html + The details of the cutoff point and the Horner-like evaluation + was picked without reference to anything in particular. + + Note that sinch is not currently implemented in scipy.special, + whereas the "engineer's" definition of sinc is implemented. + The implementation of sinc involves a scaling factor of pi + that distinguishes it from the "mathematician's" version of sinc. + + """ + + # If x is small then use sixth order Taylor expansion. + # How small is small? I am using the point where the relative error + # of the approximation is less than 1e-14. + # If x is large then directly evaluate sinh(x) / x. + if abs(x) < 0.0135: + x2 = x*x + return np.exp(a) * (1 + (x2/6.)*(1 + (x2/20.)*(1 + (x2/42.)))) + else: + return (np.exp(a + x) - np.exp(a - x)) / (2*x) + + +def _eq_10_42(lam_1, lam_2, t_12): + """ + Equation (10.42) of Functions of Matrices: Theory and Computation. + + Notes + ----- + This is a helper function for _fragment_2_1 of expm_2009. + Equation (10.42) is on page 251 in the section on Schur algorithms. + In particular, section 10.4.3 explains the Schur-Parlett algorithm. + expm([[lam_1, t_12], [0, lam_1]) + = + [[exp(lam_1), t_12*exp((lam_1 + lam_2)/2)*sinch((lam_1 - lam_2)/2)], + [0, exp(lam_2)] + """ + + # The plain formula t_12 * (exp(lam_2) - exp(lam_2)) / (lam_2 - lam_1) + # apparently suffers from cancellation, according to Higham's textbook. + # A nice implementation of sinch, defined as sinh(x)/x, + # will apparently work around the cancellation. + a = 0.5 * (lam_1 + lam_2) + b = 0.5 * (lam_1 - lam_2) + return t_12 * _exp_sinch(a, b) + + +def _fragment_2_1(X, T, s): + """ + A helper function for expm_2009. + + Notes + ----- + The argument X is modified in-place, but this modification is not the same + as the returned value of the function. + This function also takes pains to do things in ways that are compatible + with sparse arrays, for example by avoiding fancy indexing + and by using methods of the matrices whenever possible instead of + using functions of the numpy or scipy libraries themselves. + + """ + # Form X = r_m(2^-s T) + # Replace diag(X) by exp(2^-s diag(T)). + n = X.shape[0] + diag_T = np.ravel(T.diagonal().copy()) + + # Replace diag(X) by exp(2^-s diag(T)). + scale = 2 ** -s + exp_diag = np.exp(scale * diag_T) + for k in range(n): + X[k, k] = exp_diag[k] + + for i in range(s-1, -1, -1): + X = X.dot(X) + + # Replace diag(X) by exp(2^-i diag(T)). + scale = 2 ** -i + exp_diag = np.exp(scale * diag_T) + for k in range(n): + X[k, k] = exp_diag[k] + + # Replace (first) superdiagonal of X by explicit formula + # for superdiagonal of exp(2^-i T) from Eq (10.42) of + # the author's 2008 textbook + # Functions of Matrices: Theory and Computation. + for k in range(n-1): + lam_1 = scale * diag_T[k] + lam_2 = scale * diag_T[k+1] + t_12 = scale * T[k, k+1] + value = _eq_10_42(lam_1, lam_2, t_12) + X[k, k+1] = value + + # Return the updated X matrix. + return X + + +def _ell(A, m): + """ + A helper function for expm_2009. + + Parameters + ---------- + A : linear operator + A linear operator whose norm of power we care about. + m : int + The power of the linear operator + + Returns + ------- + value : int + A value related to a bound. + + """ + if len(A.shape) != 2 or A.shape[0] != A.shape[1]: + raise ValueError('expected A to be like a square matrix') + + # The c_i are explained in (2.2) and (2.6) of the 2005 expm paper. + # They are coefficients of terms of a generating function series expansion. + c_i = {3: 100800., + 5: 10059033600., + 7: 4487938430976000., + 9: 5914384781877411840000., + 13: 113250775606021113483283660800000000. + } + abs_c_recip = c_i[m] + + # This is explained after Eq. (1.2) of the 2009 expm paper. + # It is the "unit roundoff" of IEEE double precision arithmetic. + u = 2**-53 + + # Compute the one-norm of matrix power p of abs(A). + A_abs_onenorm = _onenorm_matrix_power_nnm(abs(A), 2*m + 1) + + # Treat zero norm as a special case. + if not A_abs_onenorm: + return 0 + + alpha = A_abs_onenorm / (_onenorm(A) * abs_c_recip) + log2_alpha_div_u = np.log2(alpha/u) + value = int(np.ceil(log2_alpha_div_u / (2 * m))) + return max(value, 0) + +def matrix_power(A, power): + """ + Raise a square matrix to the integer power, `power`. + + For non-negative integers, ``A**power`` is computed using repeated + matrix multiplications. Negative integers are not supported. + + Parameters + ---------- + A : (M, M) square sparse array or matrix + sparse array that will be raised to power `power` + power : int + Exponent used to raise sparse array `A` + + Returns + ------- + A**power : (M, M) sparse array or matrix + The output matrix will be the same shape as A, and will preserve + the class of A, but the format of the output may be changed. + + Notes + ----- + This uses a recursive implementation of the matrix power. For computing + the matrix power using a reasonably large `power`, this may be less efficient + than computing the product directly, using A @ A @ ... @ A. + This is contingent upon the number of nonzero entries in the matrix. + + .. versionadded:: 1.12.0 + + Examples + -------- + >>> from scipy import sparse + >>> A = sparse.csc_array([[0,1,0],[1,0,1],[0,1,0]]) + >>> A.todense() + array([[0, 1, 0], + [1, 0, 1], + [0, 1, 0]]) + >>> (A @ A).todense() + array([[1, 0, 1], + [0, 2, 0], + [1, 0, 1]]) + >>> A2 = sparse.linalg.matrix_power(A, 2) + >>> A2.todense() + array([[1, 0, 1], + [0, 2, 0], + [1, 0, 1]]) + >>> A4 = sparse.linalg.matrix_power(A, 4) + >>> A4.todense() + array([[2, 0, 2], + [0, 4, 0], + [2, 0, 2]]) + + """ + M, N = A.shape + if M != N: + raise TypeError('sparse matrix is not square') + + if isintlike(power): + power = int(power) + if power < 0: + raise ValueError('exponent must be >= 0') + + if power == 0: + return eye_array(M, dtype=A.dtype) + + if power == 1: + return A.copy() + + tmp = matrix_power(A, power // 2) + if power % 2: + return A @ tmp @ tmp + else: + return tmp @ tmp + else: + raise ValueError("exponent must be an integer") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_norm.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..14ca4719f3b450a64b7dcc8bac768a697a5a3dfe --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_norm.py @@ -0,0 +1,195 @@ +"""Sparse matrix norms. + +""" +import numpy as np +from scipy.sparse import issparse +from scipy.sparse.linalg import svds +from scipy.sparse._sputils import convert_pydata_sparse_to_scipy +import scipy.sparse as sp + +from numpy import sqrt, abs + +__all__ = ['norm'] + + +def _sparse_frobenius_norm(x): + data = sp._sputils._todata(x) + return np.linalg.norm(data) + + +def norm(x, ord=None, axis=None): + """ + Norm of a sparse matrix + + This function is able to return one of seven different matrix norms, + depending on the value of the ``ord`` parameter. + + Parameters + ---------- + x : a sparse array + Input sparse array. + ord : {non-zero int, inf, -inf, 'fro'}, optional + Order of the norm (see table under ``Notes``). inf means numpy's + `inf` object. + axis : {int, 2-tuple of ints, None}, optional + If `axis` is an integer, it specifies the axis of `x` along which to + compute the vector norms. If `axis` is a 2-tuple, it specifies the + axes that hold 2-D matrices, and the matrix norms of these matrices + are computed. If `axis` is None then either a vector norm (when `x` + is 1-D) or a matrix norm (when `x` is 2-D) is returned. + + Returns + ------- + n : float or ndarray + + Notes + ----- + Some of the ord are not implemented because some associated functions like, + _multi_svd_norm, are not yet available for sparse array. + + This docstring is modified based on numpy.linalg.norm. + https://github.com/numpy/numpy/blob/main/numpy/linalg/linalg.py + + The following norms can be calculated: + + ===== ============================ + ord norm for sparse arrays + ===== ============================ + None Frobenius norm + 'fro' Frobenius norm + inf max(sum(abs(x), axis=1)) + -inf min(sum(abs(x), axis=1)) + 0 abs(x).sum(axis=axis) + 1 max(sum(abs(x), axis=0)) + -1 min(sum(abs(x), axis=0)) + 2 Spectral norm (the largest singular value) + -2 Not implemented + other Not implemented + ===== ============================ + + The Frobenius norm is given by [1]_: + + :math:`||A||_F = [\\sum_{i,j} abs(a_{i,j})^2]^{1/2}` + + References + ---------- + .. [1] G. H. Golub and C. F. Van Loan, *Matrix Computations*, + Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 + + Examples + -------- + >>> from scipy.sparse import csr_array, diags_array + >>> import numpy as np + >>> from scipy.sparse.linalg import norm + >>> a = np.arange(9) - 4 + >>> a + array([-4, -3, -2, -1, 0, 1, 2, 3, 4]) + >>> b = a.reshape((3, 3)) + >>> b + array([[-4, -3, -2], + [-1, 0, 1], + [ 2, 3, 4]]) + + >>> b = csr_array(b) + >>> norm(b) + 7.745966692414834 + >>> norm(b, 'fro') + 7.745966692414834 + >>> norm(b, np.inf) + 9 + >>> norm(b, -np.inf) + 2 + >>> norm(b, 1) + 7 + >>> norm(b, -1) + 6 + + The matrix 2-norm or the spectral norm is the largest singular + value, computed approximately and with limitations. + + >>> b = diags_array([-1, 1], [0, 1], shape=(9, 10)) + >>> norm(b, 2) + 1.9753... + """ + x = convert_pydata_sparse_to_scipy(x, target_format="csr") + if not issparse(x): + raise TypeError("input is not sparse. use numpy.linalg.norm") + + # Check the default case first and handle it immediately. + if axis is None and ord in (None, 'fro', 'f'): + return _sparse_frobenius_norm(x) + + # Some norms require functions that are not implemented for all types. + x = x.tocsr() + + if axis is None: + axis = tuple(range(x.ndim)) + elif not isinstance(axis, tuple): + msg = "'axis' must be None, an integer or a tuple of integers" + try: + int_axis = int(axis) + except TypeError as e: + raise TypeError(msg) from e + if axis != int_axis: + raise TypeError(msg) + axis = (int_axis,) + + nd = x.ndim + if len(axis) == 2: + row_axis, col_axis = axis + if not (-nd <= row_axis < nd and -nd <= col_axis < nd): + message = f'Invalid axis {axis!r} for an array with shape {x.shape!r}' + raise ValueError(message) + if row_axis % nd == col_axis % nd: + raise ValueError('Duplicate axes given.') + if ord == 2: + # Only solver="lobpcg" supports all numpy dtypes + _, s, _ = svds(x, k=1, solver="lobpcg") + return s[0] + elif ord == -2: + raise NotImplementedError + #return _multi_svd_norm(x, row_axis, col_axis, amin) + elif ord == 1: + return abs(x).sum(axis=row_axis).max() + elif ord == np.inf: + return abs(x).sum(axis=col_axis).max() + elif ord == -1: + return abs(x).sum(axis=row_axis).min() + elif ord == -np.inf: + return abs(x).sum(axis=col_axis).min() + elif ord in (None, 'f', 'fro'): + # The axis order does not matter for this norm. + return _sparse_frobenius_norm(x) + else: + raise ValueError("Invalid norm order for matrices.") + elif len(axis) == 1: + a, = axis + if not (-nd <= a < nd): + message = f'Invalid axis {axis!r} for an array with shape {x.shape!r}' + raise ValueError(message) + if ord == np.inf: + M = abs(x).max(axis=a) + elif ord == -np.inf: + M = abs(x).min(axis=a) + elif ord == 0: + # Zero norm + M = (x != 0).sum(axis=a) + elif ord == 1: + # special case for speedup + M = abs(x).sum(axis=a) + elif ord in (2, None): + M = sqrt(abs(x).power(2).sum(axis=a)) + else: + try: + ord + 1 + except TypeError as e: + raise ValueError('Invalid norm order for vectors.') from e + M = np.power(abs(x).power(ord).sum(axis=a), 1 / ord) + if hasattr(M, 'toarray'): + return M.toarray().ravel() + elif hasattr(M, 'A'): + return M.A.ravel() + else: + return M.ravel() + else: + raise ValueError("Improper number of dimensions to norm.") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_onenormest.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_onenormest.py new file mode 100644 index 0000000000000000000000000000000000000000..e197e69701f1fa309a6e5825d189379d824dfe69 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_onenormest.py @@ -0,0 +1,467 @@ +"""Sparse block 1-norm estimator. +""" + +import numpy as np +from scipy.sparse.linalg import aslinearoperator + + +__all__ = ['onenormest'] + + +def onenormest(A, t=2, itmax=5, compute_v=False, compute_w=False): + """ + Compute a lower bound of the 1-norm of a sparse array. + + Parameters + ---------- + A : ndarray or other linear operator + A linear operator that can be transposed and that can + produce matrix products. + t : int, optional + A positive parameter controlling the tradeoff between + accuracy versus time and memory usage. + Larger values take longer and use more memory + but give more accurate output. + itmax : int, optional + Use at most this many iterations. + compute_v : bool, optional + Request a norm-maximizing linear operator input vector if True. + compute_w : bool, optional + Request a norm-maximizing linear operator output vector if True. + + Returns + ------- + est : float + An underestimate of the 1-norm of the sparse array. + v : ndarray, optional + The vector such that ||Av||_1 == est*||v||_1. + It can be thought of as an input to the linear operator + that gives an output with particularly large norm. + w : ndarray, optional + The vector Av which has relatively large 1-norm. + It can be thought of as an output of the linear operator + that is relatively large in norm compared to the input. + + Notes + ----- + This is algorithm 2.4 of [1]. + + In [2] it is described as follows. + "This algorithm typically requires the evaluation of + about 4t matrix-vector products and almost invariably + produces a norm estimate (which is, in fact, a lower + bound on the norm) correct to within a factor 3." + + .. versionadded:: 0.13.0 + + References + ---------- + .. [1] Nicholas J. Higham and Francoise Tisseur (2000), + "A Block Algorithm for Matrix 1-Norm Estimation, + with an Application to 1-Norm Pseudospectra." + SIAM J. Matrix Anal. Appl. Vol. 21, No. 4, pp. 1185-1201. + + .. [2] Awad H. Al-Mohy and Nicholas J. Higham (2009), + "A new scaling and squaring algorithm for the matrix exponential." + SIAM J. Matrix Anal. Appl. Vol. 31, No. 3, pp. 970-989. + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse import csc_array + >>> from scipy.sparse.linalg import onenormest + >>> A = csc_array([[1., 0., 0.], [5., 8., 2.], [0., -1., 0.]], dtype=float) + >>> A.toarray() + array([[ 1., 0., 0.], + [ 5., 8., 2.], + [ 0., -1., 0.]]) + >>> onenormest(A) + 9.0 + >>> np.linalg.norm(A.toarray(), ord=1) + 9.0 + """ + + # Check the input. + A = aslinearoperator(A) + if A.shape[0] != A.shape[1]: + raise ValueError('expected the operator to act like a square matrix') + + # If the operator size is small compared to t, + # then it is easier to compute the exact norm. + # Otherwise estimate the norm. + n = A.shape[1] + if t >= n: + A_explicit = np.asarray(aslinearoperator(A).matmat(np.identity(n))) + if A_explicit.shape != (n, n): + raise Exception('internal error: ', + 'unexpected shape ' + str(A_explicit.shape)) + col_abs_sums = abs(A_explicit).sum(axis=0) + if col_abs_sums.shape != (n, ): + raise Exception('internal error: ', + 'unexpected shape ' + str(col_abs_sums.shape)) + argmax_j = np.argmax(col_abs_sums) + v = elementary_vector(n, argmax_j) + w = A_explicit[:, argmax_j] + est = col_abs_sums[argmax_j] + else: + est, v, w, nmults, nresamples = _onenormest_core(A, A.H, t, itmax) + + # Report the norm estimate along with some certificates of the estimate. + if compute_v or compute_w: + result = (est,) + if compute_v: + result += (v,) + if compute_w: + result += (w,) + return result + else: + return est + + +def _blocked_elementwise(func): + """ + Decorator for an elementwise function, to apply it blockwise along + first dimension, to avoid excessive memory usage in temporaries. + """ + block_size = 2**20 + + def wrapper(x): + if x.shape[0] < block_size: + return func(x) + else: + y0 = func(x[:block_size]) + y = np.zeros((x.shape[0],) + y0.shape[1:], dtype=y0.dtype) + y[:block_size] = y0 + del y0 + for j in range(block_size, x.shape[0], block_size): + y[j:j+block_size] = func(x[j:j+block_size]) + return y + return wrapper + + +@_blocked_elementwise +def sign_round_up(X): + """ + This should do the right thing for both real and complex matrices. + + From Higham and Tisseur: + "Everything in this section remains valid for complex matrices + provided that sign(A) is redefined as the matrix (aij / |aij|) + (and sign(0) = 1) transposes are replaced by conjugate transposes." + + """ + Y = X.copy() + Y[Y == 0] = 1 + Y /= np.abs(Y) + return Y + + +@_blocked_elementwise +def _max_abs_axis1(X): + return np.max(np.abs(X), axis=1) + + +def _sum_abs_axis0(X): + block_size = 2**20 + r = None + for j in range(0, X.shape[0], block_size): + y = np.sum(np.abs(X[j:j+block_size]), axis=0) + if r is None: + r = y + else: + r += y + return r + + +def elementary_vector(n, i): + v = np.zeros(n, dtype=float) + v[i] = 1 + return v + + +def vectors_are_parallel(v, w): + # Columns are considered parallel when they are equal or negative. + # Entries are required to be in {-1, 1}, + # which guarantees that the magnitudes of the vectors are identical. + if v.ndim != 1 or v.shape != w.shape: + raise ValueError('expected conformant vectors with entries in {-1,1}') + n = v.shape[0] + return np.dot(v, w) == n + + +def every_col_of_X_is_parallel_to_a_col_of_Y(X, Y): + for v in X.T: + if not any(vectors_are_parallel(v, w) for w in Y.T): + return False + return True + + +def column_needs_resampling(i, X, Y=None): + # column i of X needs resampling if either + # it is parallel to a previous column of X or + # it is parallel to a column of Y + n, t = X.shape + v = X[:, i] + if any(vectors_are_parallel(v, X[:, j]) for j in range(i)): + return True + if Y is not None: + if any(vectors_are_parallel(v, w) for w in Y.T): + return True + return False + + +def resample_column(i, X): + X[:, i] = np.random.randint(0, 2, size=X.shape[0])*2 - 1 + + +def less_than_or_close(a, b): + return np.allclose(a, b) or (a < b) + + +def _algorithm_2_2(A, AT, t): + """ + This is Algorithm 2.2. + + Parameters + ---------- + A : ndarray or other linear operator + A linear operator that can produce matrix products. + AT : ndarray or other linear operator + The transpose of A. + t : int, optional + A positive parameter controlling the tradeoff between + accuracy versus time and memory usage. + + Returns + ------- + g : sequence + A non-negative decreasing vector + such that g[j] is a lower bound for the 1-norm + of the column of A of jth largest 1-norm. + The first entry of this vector is therefore a lower bound + on the 1-norm of the linear operator A. + This sequence has length t. + ind : sequence + The ith entry of ind is the index of the column A whose 1-norm + is given by g[i]. + This sequence of indices has length t, and its entries are + chosen from range(n), possibly with repetition, + where n is the order of the operator A. + + Notes + ----- + This algorithm is mainly for testing. + It uses the 'ind' array in a way that is similar to + its usage in algorithm 2.4. This algorithm 2.2 may be easier to test, + so it gives a chance of uncovering bugs related to indexing + which could have propagated less noticeably to algorithm 2.4. + + """ + A_linear_operator = aslinearoperator(A) + AT_linear_operator = aslinearoperator(AT) + n = A_linear_operator.shape[0] + + # Initialize the X block with columns of unit 1-norm. + X = np.ones((n, t)) + if t > 1: + X[:, 1:] = np.random.randint(0, 2, size=(n, t-1))*2 - 1 + X /= float(n) + + # Iteratively improve the lower bounds. + # Track extra things, to assert invariants for debugging. + g_prev = None + h_prev = None + k = 1 + ind = range(t) + while True: + Y = np.asarray(A_linear_operator.matmat(X)) + g = _sum_abs_axis0(Y) + best_j = np.argmax(g) + g.sort() + g = g[::-1] + S = sign_round_up(Y) + Z = np.asarray(AT_linear_operator.matmat(S)) + h = _max_abs_axis1(Z) + + # If this algorithm runs for fewer than two iterations, + # then its return values do not have the properties indicated + # in the description of the algorithm. + # In particular, the entries of g are not 1-norms of any + # column of A until the second iteration. + # Therefore we will require the algorithm to run for at least + # two iterations, even though this requirement is not stated + # in the description of the algorithm. + if k >= 2: + if less_than_or_close(max(h), np.dot(Z[:, best_j], X[:, best_j])): + break + ind = np.argsort(h)[::-1][:t] + h = h[ind] + for j in range(t): + X[:, j] = elementary_vector(n, ind[j]) + + # Check invariant (2.2). + if k >= 2: + if not less_than_or_close(g_prev[0], h_prev[0]): + raise Exception('invariant (2.2) is violated') + if not less_than_or_close(h_prev[0], g[0]): + raise Exception('invariant (2.2) is violated') + + # Check invariant (2.3). + if k >= 3: + for j in range(t): + if not less_than_or_close(g[j], g_prev[j]): + raise Exception('invariant (2.3) is violated') + + # Update for the next iteration. + g_prev = g + h_prev = h + k += 1 + + # Return the lower bounds and the corresponding column indices. + return g, ind + + +def _onenormest_core(A, AT, t, itmax): + """ + Compute a lower bound of the 1-norm of a sparse array. + + Parameters + ---------- + A : ndarray or other linear operator + A linear operator that can produce matrix products. + AT : ndarray or other linear operator + The transpose of A. + t : int, optional + A positive parameter controlling the tradeoff between + accuracy versus time and memory usage. + itmax : int, optional + Use at most this many iterations. + + Returns + ------- + est : float + An underestimate of the 1-norm of the sparse array. + v : ndarray, optional + The vector such that ||Av||_1 == est*||v||_1. + It can be thought of as an input to the linear operator + that gives an output with particularly large norm. + w : ndarray, optional + The vector Av which has relatively large 1-norm. + It can be thought of as an output of the linear operator + that is relatively large in norm compared to the input. + nmults : int, optional + The number of matrix products that were computed. + nresamples : int, optional + The number of times a parallel column was observed, + necessitating a re-randomization of the column. + + Notes + ----- + This is algorithm 2.4. + + """ + # This function is a more or less direct translation + # of Algorithm 2.4 from the Higham and Tisseur (2000) paper. + A_linear_operator = aslinearoperator(A) + AT_linear_operator = aslinearoperator(AT) + if itmax < 2: + raise ValueError('at least two iterations are required') + if t < 1: + raise ValueError('at least one column is required') + n = A.shape[0] + if t >= n: + raise ValueError('t should be smaller than the order of A') + # Track the number of big*small matrix multiplications + # and the number of resamplings. + nmults = 0 + nresamples = 0 + # "We now explain our choice of starting matrix. We take the first + # column of X to be the vector of 1s [...] This has the advantage that + # for a matrix with nonnegative elements the algorithm converges + # with an exact estimate on the second iteration, and such matrices + # arise in applications [...]" + X = np.ones((n, t), dtype=float) + # "The remaining columns are chosen as rand{-1,1}, + # with a check for and correction of parallel columns, + # exactly as for S in the body of the algorithm." + if t > 1: + for i in range(1, t): + # These are technically initial samples, not resamples, + # so the resampling count is not incremented. + resample_column(i, X) + for i in range(t): + while column_needs_resampling(i, X): + resample_column(i, X) + nresamples += 1 + # "Choose starting matrix X with columns of unit 1-norm." + X /= float(n) + # "indices of used unit vectors e_j" + ind_hist = np.zeros(0, dtype=np.intp) + est_old = 0 + S = np.zeros((n, t), dtype=float) + k = 1 + ind = None + while True: + Y = np.asarray(A_linear_operator.matmat(X)) + nmults += 1 + mags = _sum_abs_axis0(Y) + est = np.max(mags) + best_j = np.argmax(mags) + if est > est_old or k == 2: + if k >= 2: + ind_best = ind[best_j] + w = Y[:, best_j] + # (1) + if k >= 2 and est <= est_old: + est = est_old + break + est_old = est + S_old = S + if k > itmax: + break + S = sign_round_up(Y) + del Y + # (2) + if every_col_of_X_is_parallel_to_a_col_of_Y(S, S_old): + break + if t > 1: + # "Ensure that no column of S is parallel to another column of S + # or to a column of S_old by replacing columns of S by rand{-1,1}." + for i in range(t): + while column_needs_resampling(i, S, S_old): + resample_column(i, S) + nresamples += 1 + del S_old + # (3) + Z = np.asarray(AT_linear_operator.matmat(S)) + nmults += 1 + h = _max_abs_axis1(Z) + del Z + # (4) + if k >= 2 and max(h) == h[ind_best]: + break + # "Sort h so that h_first >= ... >= h_last + # and re-order ind correspondingly." + # + # Later on, we will need at most t+len(ind_hist) largest + # entries, so drop the rest + ind = np.argsort(h)[::-1][:t+len(ind_hist)].copy() + del h + if t > 1: + # (5) + # Break if the most promising t vectors have been visited already. + if np.isin(ind[:t], ind_hist).all(): + break + # Put the most promising unvisited vectors at the front of the list + # and put the visited vectors at the end of the list. + # Preserve the order of the indices induced by the ordering of h. + seen = np.isin(ind, ind_hist) + ind = np.concatenate((ind[~seen], ind[seen])) + for j in range(t): + X[:, j] = elementary_vector(n, ind[j]) + + new_ind = ind[:t][~np.isin(ind[:t], ind_hist)] + ind_hist = np.concatenate((ind_hist, new_ind)) + k += 1 + v = elementary_vector(n, ind_best) + return est, v, w, nmults, nresamples diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_special_sparse_arrays.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_special_sparse_arrays.py new file mode 100644 index 0000000000000000000000000000000000000000..1a6383b022c7cd46dfdcf08803f5408e1cc1060e --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_special_sparse_arrays.py @@ -0,0 +1,948 @@ +import numpy as np +from scipy.sparse.linalg import LinearOperator +from scipy.sparse import kron, eye, dia_array + +__all__ = ['LaplacianNd'] +# Sakurai and Mikota classes are intended for tests and benchmarks +# and explicitly not included in the public API of this module. + + +class LaplacianNd(LinearOperator): + """ + The grid Laplacian in ``N`` dimensions and its eigenvalues/eigenvectors. + + Construct Laplacian on a uniform rectangular grid in `N` dimensions + and output its eigenvalues and eigenvectors. + The Laplacian ``L`` is square, negative definite, real symmetric array + with signed integer entries and zeros otherwise. + + Parameters + ---------- + grid_shape : tuple + A tuple of integers of length ``N`` (corresponding to the dimension of + the Lapacian), where each entry gives the size of that dimension. The + Laplacian matrix is square of the size ``np.prod(grid_shape)``. + boundary_conditions : {'neumann', 'dirichlet', 'periodic'}, optional + The type of the boundary conditions on the boundaries of the grid. + Valid values are ``'dirichlet'`` or ``'neumann'``(default) or + ``'periodic'``. + dtype : dtype + Numerical type of the array. Default is ``np.int8``. + + Methods + ------- + toarray() + Construct a dense array from Laplacian data + tosparse() + Construct a sparse array from Laplacian data + eigenvalues(m=None) + Construct a 1D array of `m` largest (smallest in absolute value) + eigenvalues of the Laplacian matrix in ascending order. + eigenvectors(m=None): + Construct the array with columns made of `m` eigenvectors (``float``) + of the ``Nd`` Laplacian corresponding to the `m` ordered eigenvalues. + + .. versionadded:: 1.12.0 + + Notes + ----- + Compared to the MATLAB/Octave implementation [1] of 1-, 2-, and 3-D + Laplacian, this code allows the arbitrary N-D case and the matrix-free + callable option, but is currently limited to pure Dirichlet, Neumann or + Periodic boundary conditions only. + + The Laplacian matrix of a graph (`scipy.sparse.csgraph.laplacian`) of a + rectangular grid corresponds to the negative Laplacian with the Neumann + conditions, i.e., ``boundary_conditions = 'neumann'``. + + All eigenvalues and eigenvectors of the discrete Laplacian operator for + an ``N``-dimensional regular grid of shape `grid_shape` with the grid + step size ``h=1`` are analytically known [2]. + + References + ---------- + .. [1] https://github.com/lobpcg/blopex/blob/master/blopex_\ +tools/matlab/laplacian/laplacian.m + .. [2] "Eigenvalues and eigenvectors of the second derivative", Wikipedia + https://en.wikipedia.org/wiki/Eigenvalues_and_eigenvectors_\ +of_the_second_derivative + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse.linalg import LaplacianNd + >>> from scipy.sparse import diags, csgraph + >>> from scipy.linalg import eigvalsh + + The one-dimensional Laplacian demonstrated below for pure Neumann boundary + conditions on a regular grid with ``n=6`` grid points is exactly the + negative graph Laplacian for the undirected linear graph with ``n`` + vertices using the sparse adjacency matrix ``G`` represented by the + famous tri-diagonal matrix: + + >>> n = 6 + >>> G = diags(np.ones(n - 1), 1, format='csr') + >>> Lf = csgraph.laplacian(G, symmetrized=True, form='function') + >>> grid_shape = (n, ) + >>> lap = LaplacianNd(grid_shape, boundary_conditions='neumann') + >>> np.array_equal(lap.matmat(np.eye(n)), -Lf(np.eye(n))) + True + + Since all matrix entries of the Laplacian are integers, ``'int8'`` is + the default dtype for storing matrix representations. + + >>> lap.tosparse() + + >>> lap.toarray() + array([[-1, 1, 0, 0, 0, 0], + [ 1, -2, 1, 0, 0, 0], + [ 0, 1, -2, 1, 0, 0], + [ 0, 0, 1, -2, 1, 0], + [ 0, 0, 0, 1, -2, 1], + [ 0, 0, 0, 0, 1, -1]], dtype=int8) + >>> np.array_equal(lap.matmat(np.eye(n)), lap.toarray()) + True + >>> np.array_equal(lap.tosparse().toarray(), lap.toarray()) + True + + Any number of extreme eigenvalues and/or eigenvectors can be computed. + + >>> lap = LaplacianNd(grid_shape, boundary_conditions='periodic') + >>> lap.eigenvalues() + array([-4., -3., -3., -1., -1., 0.]) + >>> lap.eigenvalues()[-2:] + array([-1., 0.]) + >>> lap.eigenvalues(2) + array([-1., 0.]) + >>> lap.eigenvectors(1) + array([[0.40824829], + [0.40824829], + [0.40824829], + [0.40824829], + [0.40824829], + [0.40824829]]) + >>> lap.eigenvectors(2) + array([[ 0.5 , 0.40824829], + [ 0. , 0.40824829], + [-0.5 , 0.40824829], + [-0.5 , 0.40824829], + [ 0. , 0.40824829], + [ 0.5 , 0.40824829]]) + >>> lap.eigenvectors() + array([[ 0.40824829, 0.28867513, 0.28867513, 0.5 , 0.5 , + 0.40824829], + [-0.40824829, -0.57735027, -0.57735027, 0. , 0. , + 0.40824829], + [ 0.40824829, 0.28867513, 0.28867513, -0.5 , -0.5 , + 0.40824829], + [-0.40824829, 0.28867513, 0.28867513, -0.5 , -0.5 , + 0.40824829], + [ 0.40824829, -0.57735027, -0.57735027, 0. , 0. , + 0.40824829], + [-0.40824829, 0.28867513, 0.28867513, 0.5 , 0.5 , + 0.40824829]]) + + The two-dimensional Laplacian is illustrated on a regular grid with + ``grid_shape = (2, 3)`` points in each dimension. + + >>> grid_shape = (2, 3) + >>> n = np.prod(grid_shape) + + Numeration of grid points is as follows: + + >>> np.arange(n).reshape(grid_shape + (-1,)) + array([[[0], + [1], + [2]], + + [[3], + [4], + [5]]]) + + Each of the boundary conditions ``'dirichlet'``, ``'periodic'``, and + ``'neumann'`` is illustrated separately; with ``'dirichlet'`` + + >>> lap = LaplacianNd(grid_shape, boundary_conditions='dirichlet') + >>> lap.tosparse() + + >>> lap.toarray() + array([[-4, 1, 0, 1, 0, 0], + [ 1, -4, 1, 0, 1, 0], + [ 0, 1, -4, 0, 0, 1], + [ 1, 0, 0, -4, 1, 0], + [ 0, 1, 0, 1, -4, 1], + [ 0, 0, 1, 0, 1, -4]], dtype=int8) + >>> np.array_equal(lap.matmat(np.eye(n)), lap.toarray()) + True + >>> np.array_equal(lap.tosparse().toarray(), lap.toarray()) + True + >>> lap.eigenvalues() + array([-6.41421356, -5. , -4.41421356, -3.58578644, -3. , + -1.58578644]) + >>> eigvals = eigvalsh(lap.toarray().astype(np.float64)) + >>> np.allclose(lap.eigenvalues(), eigvals) + True + >>> np.allclose(lap.toarray() @ lap.eigenvectors(), + ... lap.eigenvectors() @ np.diag(lap.eigenvalues())) + True + + with ``'periodic'`` + + >>> lap = LaplacianNd(grid_shape, boundary_conditions='periodic') + >>> lap.tosparse() + + >>> lap.toarray() + array([[-4, 1, 1, 2, 0, 0], + [ 1, -4, 1, 0, 2, 0], + [ 1, 1, -4, 0, 0, 2], + [ 2, 0, 0, -4, 1, 1], + [ 0, 2, 0, 1, -4, 1], + [ 0, 0, 2, 1, 1, -4]], dtype=int8) + >>> np.array_equal(lap.matmat(np.eye(n)), lap.toarray()) + True + >>> np.array_equal(lap.tosparse().toarray(), lap.toarray()) + True + >>> lap.eigenvalues() + array([-7., -7., -4., -3., -3., 0.]) + >>> eigvals = eigvalsh(lap.toarray().astype(np.float64)) + >>> np.allclose(lap.eigenvalues(), eigvals) + True + >>> np.allclose(lap.toarray() @ lap.eigenvectors(), + ... lap.eigenvectors() @ np.diag(lap.eigenvalues())) + True + + and with ``'neumann'`` + + >>> lap = LaplacianNd(grid_shape, boundary_conditions='neumann') + >>> lap.tosparse() + + >>> lap.toarray() + array([[-2, 1, 0, 1, 0, 0], + [ 1, -3, 1, 0, 1, 0], + [ 0, 1, -2, 0, 0, 1], + [ 1, 0, 0, -2, 1, 0], + [ 0, 1, 0, 1, -3, 1], + [ 0, 0, 1, 0, 1, -2]], dtype=int8) + >>> np.array_equal(lap.matmat(np.eye(n)), lap.toarray()) + True + >>> np.array_equal(lap.tosparse().toarray(), lap.toarray()) + True + >>> lap.eigenvalues() + array([-5., -3., -3., -2., -1., 0.]) + >>> eigvals = eigvalsh(lap.toarray().astype(np.float64)) + >>> np.allclose(lap.eigenvalues(), eigvals) + True + >>> np.allclose(lap.toarray() @ lap.eigenvectors(), + ... lap.eigenvectors() @ np.diag(lap.eigenvalues())) + True + + """ + + def __init__(self, grid_shape, *, + boundary_conditions='neumann', + dtype=np.int8): + + if boundary_conditions not in ('dirichlet', 'neumann', 'periodic'): + raise ValueError( + f"Unknown value {boundary_conditions!r} is given for " + "'boundary_conditions' parameter. The valid options are " + "'dirichlet', 'periodic', and 'neumann' (default)." + ) + + self.grid_shape = grid_shape + self.boundary_conditions = boundary_conditions + # LaplacianNd folds all dimensions in `grid_shape` into a single one + N = np.prod(grid_shape) + super().__init__(dtype=dtype, shape=(N, N)) + + def _eigenvalue_ordering(self, m): + """Compute `m` largest eigenvalues in each of the ``N`` directions, + i.e., up to ``m * N`` total, order them and return `m` largest. + """ + grid_shape = self.grid_shape + if m is None: + indices = np.indices(grid_shape) + Leig = np.zeros(grid_shape) + else: + grid_shape_min = min(grid_shape, + tuple(np.ones_like(grid_shape) * m)) + indices = np.indices(grid_shape_min) + Leig = np.zeros(grid_shape_min) + + for j, n in zip(indices, grid_shape): + if self.boundary_conditions == 'dirichlet': + Leig += -4 * np.sin(np.pi * (j + 1) / (2 * (n + 1))) ** 2 + elif self.boundary_conditions == 'neumann': + Leig += -4 * np.sin(np.pi * j / (2 * n)) ** 2 + else: # boundary_conditions == 'periodic' + Leig += -4 * np.sin(np.pi * np.floor((j + 1) / 2) / n) ** 2 + + Leig_ravel = Leig.ravel() + ind = np.argsort(Leig_ravel) + eigenvalues = Leig_ravel[ind] + if m is not None: + eigenvalues = eigenvalues[-m:] + ind = ind[-m:] + + return eigenvalues, ind + + def eigenvalues(self, m=None): + """Return the requested number of eigenvalues. + + Parameters + ---------- + m : int, optional + The positive number of smallest eigenvalues to return. + If not provided, then all eigenvalues will be returned. + + Returns + ------- + eigenvalues : float array + The requested `m` smallest or all eigenvalues, in ascending order. + """ + eigenvalues, _ = self._eigenvalue_ordering(m) + return eigenvalues + + def _ev1d(self, j, n): + """Return 1 eigenvector in 1d with index `j` + and number of grid points `n` where ``j < n``. + """ + if self.boundary_conditions == 'dirichlet': + i = np.pi * (np.arange(n) + 1) / (n + 1) + ev = np.sqrt(2. / (n + 1.)) * np.sin(i * (j + 1)) + elif self.boundary_conditions == 'neumann': + i = np.pi * (np.arange(n) + 0.5) / n + ev = np.sqrt((1. if j == 0 else 2.) / n) * np.cos(i * j) + else: # boundary_conditions == 'periodic' + if j == 0: + ev = np.sqrt(1. / n) * np.ones(n) + elif j + 1 == n and n % 2 == 0: + ev = np.sqrt(1. / n) * np.tile([1, -1], n//2) + else: + i = 2. * np.pi * (np.arange(n) + 0.5) / n + ev = np.sqrt(2. / n) * np.cos(i * np.floor((j + 1) / 2)) + # make small values exact zeros correcting round-off errors + # due to symmetry of eigenvectors the exact 0. is correct + ev[np.abs(ev) < np.finfo(np.float64).eps] = 0. + return ev + + def _one_eve(self, k): + """Return 1 eigenvector in Nd with multi-index `j` + as a tensor product of the corresponding 1d eigenvectors. + """ + phi = [self._ev1d(j, n) for j, n in zip(k, self.grid_shape)] + result = phi[0] + for phi in phi[1:]: + result = np.tensordot(result, phi, axes=0) + return np.asarray(result).ravel() + + def eigenvectors(self, m=None): + """Return the requested number of eigenvectors for ordered eigenvalues. + + Parameters + ---------- + m : int, optional + The positive number of eigenvectors to return. If not provided, + then all eigenvectors will be returned. + + Returns + ------- + eigenvectors : float array + An array with columns made of the requested `m` or all eigenvectors. + The columns are ordered according to the `m` ordered eigenvalues. + """ + _, ind = self._eigenvalue_ordering(m) + if m is None: + grid_shape_min = self.grid_shape + else: + grid_shape_min = min(self.grid_shape, + tuple(np.ones_like(self.grid_shape) * m)) + + N_indices = np.unravel_index(ind, grid_shape_min) + N_indices = [tuple(x) for x in zip(*N_indices)] + eigenvectors_list = [self._one_eve(k) for k in N_indices] + return np.column_stack(eigenvectors_list) + + def toarray(self): + """ + Converts the Laplacian data to a dense array. + + Returns + ------- + L : ndarray + The shape is ``(N, N)`` where ``N = np.prod(grid_shape)``. + + """ + grid_shape = self.grid_shape + n = np.prod(grid_shape) + L = np.zeros([n, n], dtype=np.int8) + # Scratch arrays + L_i = np.empty_like(L) + Ltemp = np.empty_like(L) + + for ind, dim in enumerate(grid_shape): + # Start zeroing out L_i + L_i[:] = 0 + # Allocate the top left corner with the kernel of L_i + # Einsum returns writable view of arrays + np.einsum("ii->i", L_i[:dim, :dim])[:] = -2 + np.einsum("ii->i", L_i[: dim - 1, 1:dim])[:] = 1 + np.einsum("ii->i", L_i[1:dim, : dim - 1])[:] = 1 + + if self.boundary_conditions == 'neumann': + L_i[0, 0] = -1 + L_i[dim - 1, dim - 1] = -1 + elif self.boundary_conditions == 'periodic': + if dim > 1: + L_i[0, dim - 1] += 1 + L_i[dim - 1, 0] += 1 + else: + L_i[0, 0] += 1 + + # kron is too slow for large matrices hence the next two tricks + # 1- kron(eye, mat) is block_diag(mat, mat, ...) + # 2- kron(mat, eye) can be performed by 4d stride trick + + # 1- + new_dim = dim + # for block_diag we tile the top left portion on the diagonal + if ind > 0: + tiles = np.prod(grid_shape[:ind]) + for j in range(1, tiles): + L_i[j*dim:(j+1)*dim, j*dim:(j+1)*dim] = L_i[:dim, :dim] + new_dim += dim + # 2- + # we need the keep L_i, but reset the array + Ltemp[:new_dim, :new_dim] = L_i[:new_dim, :new_dim] + tiles = int(np.prod(grid_shape[ind+1:])) + # Zero out the top left, the rest is already 0 + L_i[:new_dim, :new_dim] = 0 + idx = [x for x in range(tiles)] + L_i.reshape( + (new_dim, tiles, + new_dim, tiles) + )[:, idx, :, idx] = Ltemp[:new_dim, :new_dim] + + L += L_i + + return L.astype(self.dtype) + + def tosparse(self): + """ + Constructs a sparse array from the Laplacian data. The returned sparse + array format is dependent on the selected boundary conditions. + + Returns + ------- + L : scipy.sparse.sparray + The shape is ``(N, N)`` where ``N = np.prod(grid_shape)``. + + """ + N = len(self.grid_shape) + p = np.prod(self.grid_shape) + L = dia_array((p, p), dtype=np.int8) + + for i in range(N): + dim = self.grid_shape[i] + data = np.ones([3, dim], dtype=np.int8) + data[1, :] *= -2 + + if self.boundary_conditions == 'neumann': + data[1, 0] = -1 + data[1, -1] = -1 + + L_i = dia_array((data, [-1, 0, 1]), shape=(dim, dim), + dtype=np.int8 + ) + + if self.boundary_conditions == 'periodic': + t = dia_array((dim, dim), dtype=np.int8) + t.setdiag([1], k=-dim+1) + t.setdiag([1], k=dim-1) + L_i += t + + for j in range(i): + L_i = kron(eye(self.grid_shape[j], dtype=np.int8), L_i) + for j in range(i + 1, N): + L_i = kron(L_i, eye(self.grid_shape[j], dtype=np.int8)) + L += L_i + return L.astype(self.dtype) + + def _matvec(self, x): + grid_shape = self.grid_shape + N = len(grid_shape) + X = x.reshape(grid_shape + (-1,)) + Y = -2 * N * X + for i in range(N): + Y += np.roll(X, 1, axis=i) + Y += np.roll(X, -1, axis=i) + if self.boundary_conditions in ('neumann', 'dirichlet'): + Y[(slice(None),)*i + (0,) + (slice(None),)*(N-i-1) + ] -= np.roll(X, 1, axis=i)[ + (slice(None),) * i + (0,) + (slice(None),) * (N-i-1) + ] + Y[ + (slice(None),) * i + (-1,) + (slice(None),) * (N-i-1) + ] -= np.roll(X, -1, axis=i)[ + (slice(None),) * i + (-1,) + (slice(None),) * (N-i-1) + ] + + if self.boundary_conditions == 'neumann': + Y[ + (slice(None),) * i + (0,) + (slice(None),) * (N-i-1) + ] += np.roll(X, 0, axis=i)[ + (slice(None),) * i + (0,) + (slice(None),) * (N-i-1) + ] + Y[ + (slice(None),) * i + (-1,) + (slice(None),) * (N-i-1) + ] += np.roll(X, 0, axis=i)[ + (slice(None),) * i + (-1,) + (slice(None),) * (N-i-1) + ] + + return Y.reshape(-1, X.shape[-1]) + + def _matmat(self, x): + return self._matvec(x) + + def _adjoint(self): + return self + + def _transpose(self): + return self + + +class Sakurai(LinearOperator): + """ + Construct a Sakurai matrix in various formats and its eigenvalues. + + Constructs the "Sakurai" matrix motivated by reference [1]_: + square real symmetric positive definite and 5-diagonal + with the main diagonal ``[5, 6, 6, ..., 6, 6, 5], the ``+1`` and ``-1`` + diagonals filled with ``-4``, and the ``+2`` and ``-2`` diagonals + made of ``1``. Its eigenvalues are analytically known to be + ``16. * np.power(np.cos(0.5 * k * np.pi / (n + 1)), 4)``. + The matrix gets ill-conditioned with its size growing. + It is useful for testing and benchmarking sparse eigenvalue solvers + especially those taking advantage of its banded 5-diagonal structure. + See the notes below for details. + + Parameters + ---------- + n : int + The size of the matrix. + dtype : dtype + Numerical type of the array. Default is ``np.int8``. + + Methods + ------- + toarray() + Construct a dense array from Laplacian data + tosparse() + Construct a sparse array from Laplacian data + tobanded() + The Sakurai matrix in the format for banded symmetric matrices, + i.e., (3, n) ndarray with 3 upper diagonals + placing the main diagonal at the bottom. + eigenvalues + All eigenvalues of the Sakurai matrix ordered ascending. + + Notes + ----- + Reference [1]_ introduces a generalized eigenproblem for the matrix pair + `A` and `B` where `A` is the identity so we turn it into an eigenproblem + just for the matrix `B` that this function outputs in various formats + together with its eigenvalues. + + .. versionadded:: 1.12.0 + + References + ---------- + .. [1] T. Sakurai, H. Tadano, Y. Inadomi, and U. Nagashima, + "A moment-based method for large-scale generalized + eigenvalue problems", + Appl. Num. Anal. Comp. Math. Vol. 1 No. 2 (2004). + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse.linalg._special_sparse_arrays import Sakurai + >>> from scipy.linalg import eig_banded + >>> n = 6 + >>> sak = Sakurai(n) + + Since all matrix entries are small integers, ``'int8'`` is + the default dtype for storing matrix representations. + + >>> sak.toarray() + array([[ 5, -4, 1, 0, 0, 0], + [-4, 6, -4, 1, 0, 0], + [ 1, -4, 6, -4, 1, 0], + [ 0, 1, -4, 6, -4, 1], + [ 0, 0, 1, -4, 6, -4], + [ 0, 0, 0, 1, -4, 5]], dtype=int8) + >>> sak.tobanded() + array([[ 1, 1, 1, 1, 1, 1], + [-4, -4, -4, -4, -4, -4], + [ 5, 6, 6, 6, 6, 5]], dtype=int8) + >>> sak.tosparse() + + >>> np.array_equal(sak.dot(np.eye(n)), sak.tosparse().toarray()) + True + >>> sak.eigenvalues() + array([0.03922866, 0.56703972, 2.41789479, 5.97822974, + 10.54287655, 14.45473055]) + >>> sak.eigenvalues(2) + array([0.03922866, 0.56703972]) + + The banded form can be used in scipy functions for banded matrices, e.g., + + >>> e = eig_banded(sak.tobanded(), eigvals_only=True) + >>> np.allclose(sak.eigenvalues, e, atol= n * n * n * np.finfo(float).eps) + True + + """ + def __init__(self, n, dtype=np.int8): + self.n = n + self.dtype = dtype + shape = (n, n) + super().__init__(dtype, shape) + + def eigenvalues(self, m=None): + """Return the requested number of eigenvalues. + + Parameters + ---------- + m : int, optional + The positive number of smallest eigenvalues to return. + If not provided, then all eigenvalues will be returned. + + Returns + ------- + eigenvalues : `np.float64` array + The requested `m` smallest or all eigenvalues, in ascending order. + """ + if m is None: + m = self.n + k = np.arange(self.n + 1 -m, self.n + 1) + return np.flip(16. * np.power(np.cos(0.5 * k * np.pi / (self.n + 1)), 4)) + + def tobanded(self): + """ + Construct the Sakurai matrix as a banded array. + """ + d0 = np.r_[5, 6 * np.ones(self.n - 2, dtype=self.dtype), 5] + d1 = -4 * np.ones(self.n, dtype=self.dtype) + d2 = np.ones(self.n, dtype=self.dtype) + return np.array([d2, d1, d0]).astype(self.dtype) + + def tosparse(self): + """ + Construct the Sakurai matrix is a sparse format. + """ + from scipy.sparse import spdiags + d = self.tobanded() + # the banded format has the main diagonal at the bottom + # `spdiags` has no `dtype` parameter so inherits dtype from banded + return spdiags([d[0], d[1], d[2], d[1], d[0]], [-2, -1, 0, 1, 2], + self.n, self.n) + + def toarray(self): + return self.tosparse().toarray() + + def _matvec(self, x): + """ + Construct matrix-free callable banded-matrix-vector multiplication by + the Sakurai matrix without constructing or storing the matrix itself + using the knowledge of its entries and the 5-diagonal format. + """ + x = x.reshape(self.n, -1) + result_dtype = np.promote_types(x.dtype, self.dtype) + sx = np.zeros_like(x, dtype=result_dtype) + sx[0, :] = 5 * x[0, :] - 4 * x[1, :] + x[2, :] + sx[-1, :] = 5 * x[-1, :] - 4 * x[-2, :] + x[-3, :] + sx[1: -1, :] = (6 * x[1: -1, :] - 4 * (x[:-2, :] + x[2:, :]) + + np.pad(x[:-3, :], ((1, 0), (0, 0))) + + np.pad(x[3:, :], ((0, 1), (0, 0)))) + return sx + + def _matmat(self, x): + """ + Construct matrix-free callable matrix-matrix multiplication by + the Sakurai matrix without constructing or storing the matrix itself + by reusing the ``_matvec(x)`` that supports both 1D and 2D arrays ``x``. + """ + return self._matvec(x) + + def _adjoint(self): + return self + + def _transpose(self): + return self + + +class MikotaM(LinearOperator): + """ + Construct a mass matrix in various formats of Mikota pair. + + The mass matrix `M` is square real diagonal + positive definite with entries that are reciprocal to integers. + + Parameters + ---------- + shape : tuple of int + The shape of the matrix. + dtype : dtype + Numerical type of the array. Default is ``np.float64``. + + Methods + ------- + toarray() + Construct a dense array from Mikota data + tosparse() + Construct a sparse array from Mikota data + tobanded() + The format for banded symmetric matrices, + i.e., (1, n) ndarray with the main diagonal. + """ + def __init__(self, shape, dtype=np.float64): + self.shape = shape + self.dtype = dtype + super().__init__(dtype, shape) + + def _diag(self): + # The matrix is constructed from its diagonal 1 / [1, ..., N+1]; + # compute in a function to avoid duplicated code & storage footprint + return (1. / np.arange(1, self.shape[0] + 1)).astype(self.dtype) + + def tobanded(self): + return self._diag() + + def tosparse(self): + from scipy.sparse import diags + return diags([self._diag()], [0], shape=self.shape, dtype=self.dtype) + + def toarray(self): + return np.diag(self._diag()).astype(self.dtype) + + def _matvec(self, x): + """ + Construct matrix-free callable banded-matrix-vector multiplication by + the Mikota mass matrix without constructing or storing the matrix itself + using the knowledge of its entries and the diagonal format. + """ + x = x.reshape(self.shape[0], -1) + return self._diag()[:, np.newaxis] * x + + def _matmat(self, x): + """ + Construct matrix-free callable matrix-matrix multiplication by + the Mikota mass matrix without constructing or storing the matrix itself + by reusing the ``_matvec(x)`` that supports both 1D and 2D arrays ``x``. + """ + return self._matvec(x) + + def _adjoint(self): + return self + + def _transpose(self): + return self + + +class MikotaK(LinearOperator): + """ + Construct a stiffness matrix in various formats of Mikota pair. + + The stiffness matrix `K` is square real tri-diagonal symmetric + positive definite with integer entries. + + Parameters + ---------- + shape : tuple of int + The shape of the matrix. + dtype : dtype + Numerical type of the array. Default is ``np.int32``. + + Methods + ------- + toarray() + Construct a dense array from Mikota data + tosparse() + Construct a sparse array from Mikota data + tobanded() + The format for banded symmetric matrices, + i.e., (2, n) ndarray with 2 upper diagonals + placing the main diagonal at the bottom. + """ + def __init__(self, shape, dtype=np.int32): + self.shape = shape + self.dtype = dtype + super().__init__(dtype, shape) + # The matrix is constructed from its diagonals; + # we precompute these to avoid duplicating the computation + n = shape[0] + self._diag0 = np.arange(2 * n - 1, 0, -2, dtype=self.dtype) + self._diag1 = - np.arange(n - 1, 0, -1, dtype=self.dtype) + + def tobanded(self): + return np.array([np.pad(self._diag1, (1, 0), 'constant'), self._diag0]) + + def tosparse(self): + from scipy.sparse import diags + return diags([self._diag1, self._diag0, self._diag1], [-1, 0, 1], + shape=self.shape, dtype=self.dtype) + + def toarray(self): + return self.tosparse().toarray() + + def _matvec(self, x): + """ + Construct matrix-free callable banded-matrix-vector multiplication by + the Mikota stiffness matrix without constructing or storing the matrix + itself using the knowledge of its entries and the 3-diagonal format. + """ + x = x.reshape(self.shape[0], -1) + result_dtype = np.promote_types(x.dtype, self.dtype) + kx = np.zeros_like(x, dtype=result_dtype) + d1 = self._diag1 + d0 = self._diag0 + kx[0, :] = d0[0] * x[0, :] + d1[0] * x[1, :] + kx[-1, :] = d1[-1] * x[-2, :] + d0[-1] * x[-1, :] + kx[1: -1, :] = (d1[:-1, None] * x[: -2, :] + + d0[1: -1, None] * x[1: -1, :] + + d1[1:, None] * x[2:, :]) + return kx + + def _matmat(self, x): + """ + Construct matrix-free callable matrix-matrix multiplication by + the Stiffness mass matrix without constructing or storing the matrix itself + by reusing the ``_matvec(x)`` that supports both 1D and 2D arrays ``x``. + """ + return self._matvec(x) + + def _adjoint(self): + return self + + def _transpose(self): + return self + + +class MikotaPair: + """ + Construct the Mikota pair of matrices in various formats and + eigenvalues of the generalized eigenproblem with them. + + The Mikota pair of matrices [1, 2]_ models a vibration problem + of a linear mass-spring system with the ends attached where + the stiffness of the springs and the masses increase along + the system length such that vibration frequencies are subsequent + integers 1, 2, ..., `n` where `n` is the number of the masses. Thus, + eigenvalues of the generalized eigenvalue problem for + the matrix pair `K` and `M` where `K` is the system stiffness matrix + and `M` is the system mass matrix are the squares of the integers, + i.e., 1, 4, 9, ..., ``n * n``. + + The stiffness matrix `K` is square real tri-diagonal symmetric + positive definite. The mass matrix `M` is diagonal with diagonal + entries 1, 1/2, 1/3, ...., ``1/n``. Both matrices get + ill-conditioned with `n` growing. + + Parameters + ---------- + n : int + The size of the matrices of the Mikota pair. + dtype : dtype + Numerical type of the array. Default is ``np.float64``. + + Attributes + ---------- + eigenvalues : 1D ndarray, ``np.uint64`` + All eigenvalues of the Mikota pair ordered ascending. + + Methods + ------- + MikotaK() + A `LinearOperator` custom object for the stiffness matrix. + MikotaM() + A `LinearOperator` custom object for the mass matrix. + + .. versionadded:: 1.12.0 + + References + ---------- + .. [1] J. Mikota, "Frequency tuning of chain structure multibody oscillators + to place the natural frequencies at omega1 and N-1 integer multiples + omega2,..., omegaN", Z. Angew. Math. Mech. 81 (2001), S2, S201-S202. + Appl. Num. Anal. Comp. Math. Vol. 1 No. 2 (2004). + .. [2] Peter C. Muller and Metin Gurgoze, + "Natural frequencies of a multi-degree-of-freedom vibration system", + Proc. Appl. Math. Mech. 6, 319-320 (2006). + http://dx.doi.org/10.1002/pamm.200610141. + + Examples + -------- + >>> import numpy as np + >>> from scipy.sparse.linalg._special_sparse_arrays import MikotaPair + >>> n = 6 + >>> mik = MikotaPair(n) + >>> mik_k = mik.k + >>> mik_m = mik.m + >>> mik_k.toarray() + array([[11., -5., 0., 0., 0., 0.], + [-5., 9., -4., 0., 0., 0.], + [ 0., -4., 7., -3., 0., 0.], + [ 0., 0., -3., 5., -2., 0.], + [ 0., 0., 0., -2., 3., -1.], + [ 0., 0., 0., 0., -1., 1.]]) + >>> mik_k.tobanded() + array([[ 0., -5., -4., -3., -2., -1.], + [11., 9., 7., 5., 3., 1.]]) + >>> mik_m.tobanded() + array([1. , 0.5 , 0.33333333, 0.25 , 0.2 , + 0.16666667]) + >>> mik_k.tosparse() + + >>> mik_m.tosparse() + + >>> np.array_equal(mik_k(np.eye(n)), mik_k.toarray()) + True + >>> np.array_equal(mik_m(np.eye(n)), mik_m.toarray()) + True + >>> mik.eigenvalues() + array([ 1, 4, 9, 16, 25, 36]) + >>> mik.eigenvalues(2) + array([ 1, 4]) + + """ + def __init__(self, n, dtype=np.float64): + self.n = n + self.dtype = dtype + self.shape = (n, n) + self.m = MikotaM(self.shape, self.dtype) + self.k = MikotaK(self.shape, self.dtype) + + def eigenvalues(self, m=None): + """Return the requested number of eigenvalues. + + Parameters + ---------- + m : int, optional + The positive number of smallest eigenvalues to return. + If not provided, then all eigenvalues will be returned. + + Returns + ------- + eigenvalues : `np.uint64` array + The requested `m` smallest or all eigenvalues, in ascending order. + """ + if m is None: + m = self.n + arange_plus1 = np.arange(1, m + 1, dtype=np.uint64) + return arange_plus1 * arange_plus1 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_svdp.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_svdp.py new file mode 100644 index 0000000000000000000000000000000000000000..265c23215ac54686e584f7ce5c63c4043ad383ef --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/_svdp.py @@ -0,0 +1,309 @@ +""" +Python wrapper for PROPACK +-------------------------- + +PROPACK is a collection of Fortran routines for iterative computation +of partial SVDs of large matrices or linear operators. + +Based on BSD licensed pypropack project: + http://github.com/jakevdp/pypropack + Author: Jake Vanderplas + +PROPACK source is BSD licensed, and available at + http://soi.stanford.edu/~rmunk/PROPACK/ +""" + +__all__ = ['_svdp'] + +import numpy as np + +from scipy.sparse.linalg import aslinearoperator +from scipy.linalg import LinAlgError + +from ._propack import _spropack # type: ignore[attr-defined] +from ._propack import _dpropack # type: ignore[attr-defined] +from ._propack import _cpropack # type: ignore[attr-defined] +from ._propack import _zpropack # type: ignore[attr-defined] + + +_lansvd_dict = { + 'f': _spropack.slansvd, + 'd': _dpropack.dlansvd, + 'F': _cpropack.clansvd, + 'D': _zpropack.zlansvd, +} + + +_lansvd_irl_dict = { + 'f': _spropack.slansvd_irl, + 'd': _dpropack.dlansvd_irl, + 'F': _cpropack.clansvd_irl, + 'D': _zpropack.zlansvd_irl, +} + +_which_converter = { + 'LM': 'L', + 'SM': 'S', +} + + +class _AProd: + """ + Wrapper class for linear operator + + The call signature of the __call__ method matches the callback of + the PROPACK routines. + """ + def __init__(self, A): + try: + self.A = aslinearoperator(A) + except TypeError: + self.A = aslinearoperator(np.asarray(A)) + + def __call__(self, transa, m, n, x, y, sparm, iparm): + if transa == 'n': + y[:] = self.A.matvec(x) + else: + y[:] = self.A.rmatvec(x) + + @property + def shape(self): + return self.A.shape + + @property + def dtype(self): + try: + return self.A.dtype + except AttributeError: + return self.A.matvec(np.zeros(self.A.shape[1])).dtype + + +def _svdp(A, k, which='LM', irl_mode=True, kmax=None, + compute_u=True, compute_v=True, v0=None, full_output=False, tol=0, + delta=None, eta=None, anorm=0, cgs=False, elr=True, + min_relgap=0.002, shifts=None, maxiter=None, rng=None): + """ + Compute the singular value decomposition of a linear operator using PROPACK + + Parameters + ---------- + A : array_like, sparse matrix, or LinearOperator + Operator for which SVD will be computed. If `A` is a LinearOperator + object, it must define both ``matvec`` and ``rmatvec`` methods. + k : int + Number of singular values/vectors to compute + which : {"LM", "SM"} + Which singular triplets to compute: + - 'LM': compute triplets corresponding to the `k` largest singular + values + - 'SM': compute triplets corresponding to the `k` smallest singular + values + `which='SM'` requires `irl_mode=True`. Computes largest singular + values by default. + irl_mode : bool, optional + If `True`, then compute SVD using IRL (implicitly restarted Lanczos) + mode. Default is `True`. + kmax : int, optional + Maximal number of iterations / maximal dimension of the Krylov + subspace. Default is ``10 * k``. + compute_u : bool, optional + If `True` (default) then compute left singular vectors, `u`. + compute_v : bool, optional + If `True` (default) then compute right singular vectors, `v`. + tol : float, optional + The desired relative accuracy for computed singular values. + If not specified, it will be set based on machine precision. + v0 : array_like, optional + Starting vector for iterations: must be of length ``A.shape[0]``. + If not specified, PROPACK will generate a starting vector. + full_output : bool, optional + If `True`, then return sigma_bound. Default is `False`. + delta : float, optional + Level of orthogonality to maintain between Lanczos vectors. + Default is set based on machine precision. + eta : float, optional + Orthogonality cutoff. During reorthogonalization, vectors with + component larger than `eta` along the Lanczos vector will be purged. + Default is set based on machine precision. + anorm : float, optional + Estimate of ``||A||``. Default is ``0``. + cgs : bool, optional + If `True`, reorthogonalization is done using classical Gram-Schmidt. + If `False` (default), it is done using modified Gram-Schmidt. + elr : bool, optional + If `True` (default), then extended local orthogonality is enforced + when obtaining singular vectors. + min_relgap : float, optional + The smallest relative gap allowed between any shift in IRL mode. + Default is ``0.001``. Accessed only if ``irl_mode=True``. + shifts : int, optional + Number of shifts per restart in IRL mode. Default is determined + to satisfy ``k <= min(kmax-shifts, m, n)``. Must be + >= 0, but choosing 0 might lead to performance degradation. + Accessed only if ``irl_mode=True``. + maxiter : int, optional + Maximum number of restarts in IRL mode. Default is ``1000``. + Accessed only if ``irl_mode=True``. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + u : ndarray + The `k` largest (``which="LM"``) or smallest (``which="SM"``) left + singular vectors, ``shape == (A.shape[0], 3)``, returned only if + ``compute_u=True``. + sigma : ndarray + The top `k` singular values, ``shape == (k,)`` + vt : ndarray + The `k` largest (``which="LM"``) or smallest (``which="SM"``) right + singular vectors, ``shape == (3, A.shape[1])``, returned only if + ``compute_v=True``. + sigma_bound : ndarray + the error bounds on the singular values sigma, returned only if + ``full_output=True``. + + """ + if rng is None: + raise ValueError("`rng` must be a normalized numpy.random.Generator instance") + + which = which.upper() + if which not in {'LM', 'SM'}: + raise ValueError("`which` must be either 'LM' or 'SM'") + if not irl_mode and which == 'SM': + raise ValueError("`which`='SM' requires irl_mode=True") + + aprod = _AProd(A) + typ = aprod.dtype.char + + try: + lansvd_irl = _lansvd_irl_dict[typ] + lansvd = _lansvd_dict[typ] + except KeyError: + # work with non-supported types using native system precision + if np.iscomplexobj(np.empty(0, dtype=typ)): + typ = np.dtype(complex).char + else: + typ = np.dtype(float).char + lansvd_irl = _lansvd_irl_dict[typ] + lansvd = _lansvd_dict[typ] + + m, n = aprod.shape + if (k < 1) or (k > min(m, n)): + raise ValueError("k must be positive and not greater than m or n") + + if kmax is None: + kmax = 10*k + if maxiter is None: + maxiter = 1000 + + # guard against unnecessarily large kmax + kmax = min(m + 1, n + 1, kmax) + if kmax < k: + raise ValueError( + "kmax must be greater than or equal to k, " + f"but kmax ({kmax}) < k ({k})") + + # convert python args to fortran args + jobu = 'y' if compute_u else 'n' + jobv = 'y' if compute_v else 'n' + + # these will be the output arrays + u = np.zeros((m, kmax + 1), order='F', dtype=typ) + v = np.zeros((n, kmax), order='F', dtype=typ) + + # Specify the starting vector. if v0 is all zero, PROPACK will generate + # a random starting vector: the random seed cannot be controlled in that + # case, so we'll instead use numpy to generate a random vector + if v0 is None: + u[:, 0] = rng.uniform(size=m) + if np.iscomplexobj(np.empty(0, dtype=typ)): # complex type + u[:, 0] += 1j * rng.uniform(size=m) + else: + try: + u[:, 0] = v0 + except ValueError: + raise ValueError(f"v0 must be of length {m}") + + # process options for the fit + if delta is None: + delta = np.sqrt(np.finfo(typ).eps) + if eta is None: + eta = np.finfo(typ).eps ** 0.75 + + if irl_mode: + doption = np.array((delta, eta, anorm, min_relgap), dtype=typ.lower()) + + # validate or find default shifts + if shifts is None: + shifts = kmax - k + if k > min(kmax - shifts, m, n): + raise ValueError('shifts must satisfy ' + 'k <= min(kmax-shifts, m, n)!') + elif shifts < 0: + raise ValueError('shifts must be >= 0!') + + else: + doption = np.array((delta, eta, anorm), dtype=typ.lower()) + + ioption = np.array((int(bool(cgs)), int(bool(elr))), dtype='i') + + # If computing `u` or `v` (left and right singular vectors, + # respectively), `blocksize` controls how large a fraction of the + # work is done via fast BLAS level 3 operations. A larger blocksize + # may lead to faster computation at the expense of greater memory + # consumption. `blocksize` must be ``>= 1``. Choosing blocksize + # of 16, but docs don't specify; it's almost surely a + # power of 2. + blocksize = 16 + + # Determine lwork & liwork: + # the required lengths are specified in the PROPACK documentation + if compute_u or compute_v: + lwork = m + n + 9*kmax + 5*kmax*kmax + 4 + max( + 3*kmax*kmax + 4*kmax + 4, + blocksize*max(m, n)) + liwork = 8*kmax + else: + lwork = m + n + 9*kmax + 2*kmax*kmax + 4 + max(m + n, 4*kmax + 4) + liwork = 2*kmax + 1 + work = np.empty(lwork, dtype=typ.lower()) + iwork = np.empty(liwork, dtype=np.int32) + + # dummy arguments: these are passed to aprod, and not used in this wrapper + dparm = np.empty(1, dtype=typ.lower()) + iparm = np.empty(1, dtype=np.int32) + + if typ.isupper(): + # PROPACK documentation is unclear on the required length of zwork. + # Use the same length Julia's wrapper uses + # see https://github.com/JuliaSmoothOptimizers/PROPACK.jl/ + zwork = np.empty(m + n + 32*m, dtype=typ) + works = work, zwork, iwork + else: + works = work, iwork + + if irl_mode: + u, sigma, bnd, v, info = lansvd_irl(_which_converter[which], jobu, + jobv, m, n, shifts, k, maxiter, + aprod, u, v, tol, *works, doption, + ioption, dparm, iparm) + else: + u, sigma, bnd, v, info = lansvd(jobu, jobv, m, n, k, aprod, u, v, tol, + *works, doption, ioption, dparm, iparm) + + if info > 0: + raise LinAlgError( + f"An invariant subspace of dimension {info} was found.") + elif info < 0: + raise LinAlgError( + f"k={k} singular triplets did not converge within " + f"kmax={kmax} iterations") + + # info == 0: The K largest (or smallest) singular triplets were computed + # successfully! + + return u[:, :k], sigma, v[:, :k].conj().T, bnd diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/dsolve.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/dsolve.py new file mode 100644 index 0000000000000000000000000000000000000000..b737de449b512000e00113b9df3a575a7df08668 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/dsolve.py @@ -0,0 +1,22 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.sparse.linalg` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'MatrixRankWarning', 'SuperLU', 'factorized', + 'spilu', 'splu', 'spsolve', + 'spsolve_triangular', 'use_solver', 'test' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="sparse.linalg", module="dsolve", + private_modules=["_dsolve"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/eigen.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/eigen.py new file mode 100644 index 0000000000000000000000000000000000000000..fae7e9b3221cc8be89c876a4fd9692ea29b4855d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/eigen.py @@ -0,0 +1,21 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.sparse.linalg` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'ArpackError', 'ArpackNoConvergence', 'ArpackError', + 'eigs', 'eigsh', 'lobpcg', 'svds', 'test' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="sparse.linalg", module="eigen", + private_modules=["_eigen"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/interface.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/interface.py new file mode 100644 index 0000000000000000000000000000000000000000..7e697532c38344d29bcbd003b7851dcf9c7f29ec --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/interface.py @@ -0,0 +1,20 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.sparse.linalg` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'LinearOperator', 'aslinearoperator', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="sparse.linalg", module="interface", + private_modules=["_interface"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/isolve.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/isolve.py new file mode 100644 index 0000000000000000000000000000000000000000..13aa3bda05be7a39f9eafb234c80aea33063916d --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/isolve.py @@ -0,0 +1,22 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.sparse.linalg` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'bicg', 'bicgstab', 'cg', 'cgs', 'gcrotmk', 'gmres', + 'lgmres', 'lsmr', 'lsqr', + 'minres', 'qmr', 'tfqmr', 'test' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="sparse.linalg", module="isolve", + private_modules=["_isolve"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/matfuncs.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/matfuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..c326554c206e0b32351f9fb6259ac7e505abfe53 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/matfuncs.py @@ -0,0 +1,18 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.sparse.linalg` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = ["expm", "inv", "spsolve", "LinearOperator"] # noqa: F822 + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="sparse.linalg", module="matfuncs", + private_modules=["_matfuncs"], all=__all__, + attribute=name) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/propack_test_data.npz b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/propack_test_data.npz new file mode 100644 index 0000000000000000000000000000000000000000..0bf01015610346655c749ead87a47d5575e2b67b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/propack_test_data.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bfe34d9a92353e08f400f3837136e553a8e91d441186913d39b59bf8a627bba3 +size 600350 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_expm_multiply.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_expm_multiply.py new file mode 100644 index 0000000000000000000000000000000000000000..6ea0380b8a6c6d4e6255501e29fcd22a03505314 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_expm_multiply.py @@ -0,0 +1,367 @@ +"""Test functions for the sparse.linalg._expm_multiply module.""" +from functools import partial +from itertools import product + +import numpy as np +import pytest +from numpy.testing import (assert_allclose, assert_, assert_equal, + suppress_warnings) +from scipy.sparse import SparseEfficiencyWarning +import scipy.sparse +from scipy.sparse.linalg import aslinearoperator +import scipy.linalg +from scipy.sparse.linalg import expm as sp_expm +from scipy.sparse.linalg._expm_multiply import (_theta, _compute_p_max, + _onenormest_matrix_power, expm_multiply, _expm_multiply_simple, + _expm_multiply_interval) +from scipy._lib._util import np_long + + +IMPRECISE = {np.single, np.csingle} +REAL_DTYPES = (np.intc, np_long, np.longlong, + np.float32, np.float64, np.longdouble) +COMPLEX_DTYPES = (np.complex64, np.complex128, np.clongdouble) +DTYPES = REAL_DTYPES + COMPLEX_DTYPES + + +def estimated(func): + """If trace is estimated, it should warn. + + We warn that estimation of trace might impact performance. + All result have to be correct nevertheless! + + """ + def wrapped(*args, **kwds): + with pytest.warns(UserWarning, + match="Trace of LinearOperator not available"): + return func(*args, **kwds) + return wrapped + + +def less_than_or_close(a, b): + return np.allclose(a, b) or (a < b) + + +class TestExpmActionSimple: + """ + These tests do not consider the case of multiple time steps in one call. + """ + + def test_theta_monotonicity(self): + pairs = sorted(_theta.items()) + for (m_a, theta_a), (m_b, theta_b) in zip(pairs[:-1], pairs[1:]): + assert_(theta_a < theta_b) + + def test_p_max_default(self): + m_max = 55 + expected_p_max = 8 + observed_p_max = _compute_p_max(m_max) + assert_equal(observed_p_max, expected_p_max) + + def test_p_max_range(self): + for m_max in range(1, 55+1): + p_max = _compute_p_max(m_max) + assert_(p_max*(p_max - 1) <= m_max + 1) + p_too_big = p_max + 1 + assert_(p_too_big*(p_too_big - 1) > m_max + 1) + + def test_onenormest_matrix_power(self): + rng = np.random.RandomState(1234) + n = 40 + nsamples = 10 + for i in range(nsamples): + A = scipy.linalg.inv(rng.randn(n, n)) + for p in range(4): + if not p: + M = np.identity(n) + else: + M = np.dot(M, A) + estimated = _onenormest_matrix_power(A, p) + exact = np.linalg.norm(M, 1) + assert_(less_than_or_close(estimated, exact)) + assert_(less_than_or_close(exact, 3*estimated)) + + @pytest.mark.thread_unsafe + def test_expm_multiply(self): + np.random.seed(1234) + n = 40 + k = 3 + nsamples = 10 + for i in range(nsamples): + A = scipy.linalg.inv(np.random.randn(n, n)) + B = np.random.randn(n, k) + observed = expm_multiply(A, B) + expected = np.dot(sp_expm(A), B) + assert_allclose(observed, expected) + observed = estimated(expm_multiply)(aslinearoperator(A), B) + assert_allclose(observed, expected) + traceA = np.trace(A) + observed = expm_multiply(aslinearoperator(A), B, traceA=traceA) + assert_allclose(observed, expected) + + @pytest.mark.thread_unsafe + def test_matrix_vector_multiply(self): + np.random.seed(1234) + n = 40 + nsamples = 10 + for i in range(nsamples): + A = scipy.linalg.inv(np.random.randn(n, n)) + v = np.random.randn(n) + observed = expm_multiply(A, v) + expected = np.dot(sp_expm(A), v) + assert_allclose(observed, expected) + observed = estimated(expm_multiply)(aslinearoperator(A), v) + assert_allclose(observed, expected) + + @pytest.mark.thread_unsafe + def test_scaled_expm_multiply(self): + np.random.seed(1234) + n = 40 + k = 3 + nsamples = 10 + for i, t in product(range(nsamples), [0.2, 1.0, 1.5]): + with np.errstate(invalid='ignore'): + A = scipy.linalg.inv(np.random.randn(n, n)) + B = np.random.randn(n, k) + observed = _expm_multiply_simple(A, B, t=t) + expected = np.dot(sp_expm(t*A), B) + assert_allclose(observed, expected) + observed = estimated(_expm_multiply_simple)( + aslinearoperator(A), B, t=t + ) + assert_allclose(observed, expected) + + @pytest.mark.thread_unsafe + def test_scaled_expm_multiply_single_timepoint(self): + np.random.seed(1234) + t = 0.1 + n = 5 + k = 2 + A = np.random.randn(n, n) + B = np.random.randn(n, k) + observed = _expm_multiply_simple(A, B, t=t) + expected = sp_expm(t*A).dot(B) + assert_allclose(observed, expected) + observed = estimated(_expm_multiply_simple)( + aslinearoperator(A), B, t=t + ) + assert_allclose(observed, expected) + + @pytest.mark.thread_unsafe + def test_sparse_expm_multiply(self): + rng = np.random.default_rng(1234) + n = 40 + k = 3 + nsamples = 10 + for i in range(nsamples): + A = scipy.sparse.random_array((n, n), density=0.05, rng=rng) + B = rng.standard_normal((n, k)) + observed = expm_multiply(A, B) + with suppress_warnings() as sup: + sup.filter(SparseEfficiencyWarning, + "splu converted its input to CSC format") + sup.filter(SparseEfficiencyWarning, + "spsolve is more efficient when sparse b is in the" + " CSC matrix format") + expected = sp_expm(A).dot(B) + assert_allclose(observed, expected) + observed = estimated(expm_multiply)(aslinearoperator(A), B) + assert_allclose(observed, expected) + + @pytest.mark.thread_unsafe + def test_complex(self): + A = np.array([ + [1j, 1j], + [0, 1j]], dtype=complex) + B = np.array([1j, 1j]) + observed = expm_multiply(A, B) + expected = np.array([ + 1j * np.exp(1j) + 1j * (1j*np.cos(1) - np.sin(1)), + 1j * np.exp(1j)], dtype=complex) + assert_allclose(observed, expected) + observed = estimated(expm_multiply)(aslinearoperator(A), B) + assert_allclose(observed, expected) + + +class TestExpmActionInterval: + + @pytest.mark.fail_slow(20) + def test_sparse_expm_multiply_interval(self): + rng = np.random.default_rng(1234) + start = 0.1 + stop = 3.2 + n = 40 + k = 3 + endpoint = True + for num in (14, 13, 2): + A = scipy.sparse.random_array((n, n), density=0.05, rng=rng) + B = rng.standard_normal((n, k)) + v = rng.standard_normal((n,)) + for target in (B, v): + X = expm_multiply(A, target, start=start, stop=stop, + num=num, endpoint=endpoint) + samples = np.linspace(start=start, stop=stop, + num=num, endpoint=endpoint) + with suppress_warnings() as sup: + sup.filter(SparseEfficiencyWarning, + "splu converted its input to CSC format") + sup.filter(SparseEfficiencyWarning, + "spsolve is more efficient when sparse b is in" + " the CSC matrix format") + for solution, t in zip(X, samples): + assert_allclose(solution, sp_expm(t*A).dot(target)) + + @pytest.mark.thread_unsafe + @pytest.mark.fail_slow(20) + def test_expm_multiply_interval_vector(self): + np.random.seed(1234) + interval = {'start': 0.1, 'stop': 3.2, 'endpoint': True} + for num, n in product([14, 13, 2], [1, 2, 5, 20, 40]): + A = scipy.linalg.inv(np.random.randn(n, n)) + v = np.random.randn(n) + samples = np.linspace(num=num, **interval) + X = expm_multiply(A, v, num=num, **interval) + for solution, t in zip(X, samples): + assert_allclose(solution, sp_expm(t*A).dot(v)) + # test for linear operator with unknown trace -> estimate trace + Xguess = estimated(expm_multiply)(aslinearoperator(A), v, + num=num, **interval) + # test for linear operator with given trace + Xgiven = expm_multiply(aslinearoperator(A), v, num=num, **interval, + traceA=np.trace(A)) + # test robustness for linear operator with wrong trace + Xwrong = expm_multiply(aslinearoperator(A), v, num=num, **interval, + traceA=np.trace(A)*5) + for sol_guess, sol_given, sol_wrong, t in zip(Xguess, Xgiven, + Xwrong, samples): + correct = sp_expm(t*A).dot(v) + assert_allclose(sol_guess, correct) + assert_allclose(sol_given, correct) + assert_allclose(sol_wrong, correct) + + @pytest.mark.thread_unsafe + @pytest.mark.fail_slow(20) + def test_expm_multiply_interval_matrix(self): + np.random.seed(1234) + interval = {'start': 0.1, 'stop': 3.2, 'endpoint': True} + for num, n, k in product([14, 13, 2], [1, 2, 5, 20, 40], [1, 2]): + A = scipy.linalg.inv(np.random.randn(n, n)) + B = np.random.randn(n, k) + samples = np.linspace(num=num, **interval) + X = expm_multiply(A, B, num=num, **interval) + for solution, t in zip(X, samples): + assert_allclose(solution, sp_expm(t*A).dot(B)) + X = estimated(expm_multiply)(aslinearoperator(A), B, num=num, + **interval) + for solution, t in zip(X, samples): + assert_allclose(solution, sp_expm(t*A).dot(B)) + + def test_sparse_expm_multiply_interval_dtypes(self): + # Test A & B int + A = scipy.sparse.diags_array(np.arange(5),format='csr', dtype=int) + B = np.ones(5, dtype=int) + Aexpm = scipy.sparse.diags_array(np.exp(np.arange(5)),format='csr') + BI = np.identity(5, dtype=int) + BI_sparse = scipy.sparse.csr_array(BI) + assert_allclose(expm_multiply(A,B,0,1)[-1], Aexpm.dot(B)) + assert_allclose(np.diag(expm_multiply(A, BI_sparse, 0, 1)[-1]), Aexpm.dot(B)) + + # Test A complex, B int + A = scipy.sparse.diags_array(-1j*np.arange(5),format='csr', dtype=complex) + B = np.ones(5, dtype=int) + Aexpm = scipy.sparse.diags_array(np.exp(-1j*np.arange(5)),format='csr') + assert_allclose(expm_multiply(A,B,0,1)[-1], Aexpm.dot(B)) + assert_allclose(np.diag(expm_multiply(A, BI_sparse, 0, 1)[-1]), Aexpm.dot(B)) + + # Test A int, B complex + A = scipy.sparse.diags_array(np.arange(5),format='csr', dtype=int) + B = np.full(5, 1j, dtype=complex) + Aexpm = scipy.sparse.diags_array(np.exp(np.arange(5)),format='csr') + assert_allclose(expm_multiply(A,B,0,1)[-1], Aexpm.dot(B)) + BI = np.identity(5, dtype=complex)*1j + assert_allclose( + np.diag(expm_multiply(A, scipy.sparse.csr_array(BI), 0, 1)[-1]), + Aexpm.dot(B) + ) + + def test_expm_multiply_interval_status_0(self): + self._help_test_specific_expm_interval_status(0) + + def test_expm_multiply_interval_status_1(self): + self._help_test_specific_expm_interval_status(1) + + def test_expm_multiply_interval_status_2(self): + self._help_test_specific_expm_interval_status(2) + + def _help_test_specific_expm_interval_status(self, target_status): + rng = np.random.RandomState(1234) + start = 0.1 + stop = 3.2 + num = 13 + endpoint = True + n = 5 + k = 2 + nrepeats = 10 + nsuccesses = 0 + for num in [14, 13, 2] * nrepeats: + A = rng.randn(n, n) + B = rng.randn(n, k) + status = _expm_multiply_interval(A, B, + start=start, stop=stop, num=num, endpoint=endpoint, + status_only=True) + if status == target_status: + X, status = _expm_multiply_interval(A, B, + start=start, stop=stop, num=num, endpoint=endpoint, + status_only=False) + assert_equal(X.shape, (num, n, k)) + samples = np.linspace(start=start, stop=stop, + num=num, endpoint=endpoint) + for solution, t in zip(X, samples): + assert_allclose(solution, sp_expm(t*A).dot(B)) + nsuccesses += 1 + if not nsuccesses: + msg = 'failed to find a status-' + str(target_status) + ' interval' + raise Exception(msg) + + +@pytest.mark.thread_unsafe +@pytest.mark.parametrize("dtype_a", DTYPES) +@pytest.mark.parametrize("dtype_b", DTYPES) +@pytest.mark.parametrize("b_is_matrix", [False, True]) +def test_expm_multiply_dtype(dtype_a, dtype_b, b_is_matrix): + """Make sure `expm_multiply` handles all numerical dtypes correctly.""" + assert_allclose_ = (partial(assert_allclose, rtol=1.8e-3, atol=1e-5) + if {dtype_a, dtype_b} & IMPRECISE else assert_allclose) + rng = np.random.default_rng(1234) + # test data + n = 7 + b_shape = (n, 3) if b_is_matrix else (n, ) + if dtype_a in REAL_DTYPES: + A = scipy.linalg.inv(rng.random([n, n])).astype(dtype_a) + else: + A = scipy.linalg.inv( + rng.random([n, n]) + 1j*rng.random([n, n]) + ).astype(dtype_a) + if dtype_b in REAL_DTYPES: + B = (2*rng.random(b_shape)).astype(dtype_b) + else: + B = (rng.random(b_shape) + 1j*rng.random(b_shape)).astype(dtype_b) + + # single application + sol_mat = expm_multiply(A, B) + sol_op = estimated(expm_multiply)(aslinearoperator(A), B) + direct_sol = np.dot(sp_expm(A), B) + assert_allclose_(sol_mat, direct_sol) + assert_allclose_(sol_op, direct_sol) + sol_op = expm_multiply(aslinearoperator(A), B, traceA=np.trace(A)) + assert_allclose_(sol_op, direct_sol) + + # for time points + interval = {'start': 0.1, 'stop': 3.2, 'num': 13, 'endpoint': True} + samples = np.linspace(**interval) + X_mat = expm_multiply(A, B, **interval) + X_op = estimated(expm_multiply)(aslinearoperator(A), B, **interval) + for sol_mat, sol_op, t in zip(X_mat, X_op, samples): + direct_sol = sp_expm(t*A).dot(B) + assert_allclose_(sol_mat, direct_sol) + assert_allclose_(sol_op, direct_sol) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_interface.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..c532fba288be68f687e03be5a22144fd376b45ab --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_interface.py @@ -0,0 +1,561 @@ +"""Test functions for the sparse.linalg._interface module +""" + +from functools import partial +from itertools import product +import operator +import pytest +from pytest import raises as assert_raises, warns +from numpy.testing import assert_, assert_equal + +import numpy as np +import scipy.sparse as sparse + +import scipy.sparse.linalg._interface as interface +from scipy.sparse._sputils import matrix +from scipy._lib._gcutils import assert_deallocated, IS_PYPY + + +class TestLinearOperator: + def setup_method(self): + self.A = np.array([[1,2,3], + [4,5,6]]) + self.B = np.array([[1,2], + [3,4], + [5,6]]) + self.C = np.array([[1,2], + [3,4]]) + + def test_matvec(self): + def get_matvecs(A): + return [{ + 'shape': A.shape, + 'matvec': lambda x: np.dot(A, x).reshape(A.shape[0]), + 'rmatvec': lambda x: np.dot(A.T.conj(), + x).reshape(A.shape[1]) + }, + { + 'shape': A.shape, + 'matvec': lambda x: np.dot(A, x), + 'rmatvec': lambda x: np.dot(A.T.conj(), x), + 'rmatmat': lambda x: np.dot(A.T.conj(), x), + 'matmat': lambda x: np.dot(A, x) + }] + + for matvecs in get_matvecs(self.A): + A = interface.LinearOperator(**matvecs) + + assert_(A.args == ()) + + assert_equal(A.matvec(np.array([1,2,3])), [14,32]) + assert_equal(A.matvec(np.array([[1],[2],[3]])), [[14],[32]]) + assert_equal(A @ np.array([1,2,3]), [14,32]) + assert_equal(A @ np.array([[1],[2],[3]]), [[14],[32]]) + assert_equal(A.dot(np.array([1,2,3])), [14,32]) + assert_equal(A.dot(np.array([[1],[2],[3]])), [[14],[32]]) + + assert_equal(A.matvec(matrix([[1],[2],[3]])), [[14],[32]]) + assert_equal(A @ matrix([[1],[2],[3]]), [[14],[32]]) + assert_equal(A.dot(matrix([[1],[2],[3]])), [[14],[32]]) + + assert_equal((2*A)@[1,1,1], [12,30]) + assert_equal((2 * A).rmatvec([1, 1]), [10, 14, 18]) + assert_equal((2*A).H.matvec([1,1]), [10, 14, 18]) + assert_equal((2*A).adjoint().matvec([1,1]), [10, 14, 18]) + assert_equal((2*A)@[[1],[1],[1]], [[12],[30]]) + assert_equal((2 * A).matmat([[1], [1], [1]]), [[12], [30]]) + assert_equal((A*2)@[1,1,1], [12,30]) + assert_equal((A*2)@[[1],[1],[1]], [[12],[30]]) + assert_equal((2j*A)@[1,1,1], [12j,30j]) + assert_equal((A+A)@[1,1,1], [12, 30]) + assert_equal((A + A).rmatvec([1, 1]), [10, 14, 18]) + assert_equal((A+A).H.matvec([1,1]), [10, 14, 18]) + assert_equal((A+A).adjoint().matvec([1,1]), [10, 14, 18]) + assert_equal((A+A)@[[1],[1],[1]], [[12], [30]]) + assert_equal((A+A).matmat([[1],[1],[1]]), [[12], [30]]) + assert_equal((-A)@[1,1,1], [-6,-15]) + assert_equal((-A)@[[1],[1],[1]], [[-6],[-15]]) + assert_equal((A-A)@[1,1,1], [0,0]) + assert_equal((A - A) @ [[1], [1], [1]], [[0], [0]]) + + X = np.array([[1, 2], [3, 4]]) + # A_asarray = np.array([[1, 2, 3], [4, 5, 6]]) + assert_equal((2 * A).rmatmat(X), np.dot((2 * self.A).T, X)) + assert_equal((A * 2).rmatmat(X), np.dot((self.A * 2).T, X)) + assert_equal((2j * A).rmatmat(X), + np.dot((2j * self.A).T.conj(), X)) + assert_equal((A * 2j).rmatmat(X), + np.dot((self.A * 2j).T.conj(), X)) + assert_equal((A + A).rmatmat(X), + np.dot((self.A + self.A).T, X)) + assert_equal((A + 2j * A).rmatmat(X), + np.dot((self.A + 2j * self.A).T.conj(), X)) + assert_equal((-A).rmatmat(X), np.dot((-self.A).T, X)) + assert_equal((A - A).rmatmat(X), + np.dot((self.A - self.A).T, X)) + assert_equal((2j * A).rmatmat(2j * X), + np.dot((2j * self.A).T.conj(), 2j * X)) + + z = A+A + assert_(len(z.args) == 2 and z.args[0] is A and z.args[1] is A) + z = 2*A + assert_(len(z.args) == 2 and z.args[0] is A and z.args[1] == 2) + + assert_(isinstance(A.matvec([1, 2, 3]), np.ndarray)) + assert_(isinstance(A.matvec(np.array([[1],[2],[3]])), np.ndarray)) + assert_(isinstance(A @ np.array([1,2,3]), np.ndarray)) + assert_(isinstance(A @ np.array([[1],[2],[3]]), np.ndarray)) + assert_(isinstance(A.dot(np.array([1,2,3])), np.ndarray)) + assert_(isinstance(A.dot(np.array([[1],[2],[3]])), np.ndarray)) + + assert_(isinstance(A.matvec(matrix([[1],[2],[3]])), np.ndarray)) + assert_(isinstance(A @ matrix([[1],[2],[3]]), np.ndarray)) + assert_(isinstance(A.dot(matrix([[1],[2],[3]])), np.ndarray)) + + assert_(isinstance(2*A, interface._ScaledLinearOperator)) + assert_(isinstance(2j*A, interface._ScaledLinearOperator)) + assert_(isinstance(A+A, interface._SumLinearOperator)) + assert_(isinstance(-A, interface._ScaledLinearOperator)) + assert_(isinstance(A-A, interface._SumLinearOperator)) + assert_(isinstance(A/2, interface._ScaledLinearOperator)) + assert_(isinstance(A/2j, interface._ScaledLinearOperator)) + assert_(((A * 3) / 3).args[0] is A) # check for simplification + + # Test that prefactor is of _ScaledLinearOperator is not mutated + # when the operator is multiplied by a number + result = A @ np.array([1, 2, 3]) + B = A * 3 + C = A / 5 + assert_equal(A @ np.array([1, 2, 3]), result) + + assert_((2j*A).dtype == np.complex128) + + # Test division by non-scalar + msg = "Can only divide a linear operator by a scalar." + with assert_raises(ValueError, match=msg): + A / np.array([1, 2]) + + assert_raises(ValueError, A.matvec, np.array([1,2])) + assert_raises(ValueError, A.matvec, np.array([1,2,3,4])) + assert_raises(ValueError, A.matvec, np.array([[1],[2]])) + assert_raises(ValueError, A.matvec, np.array([[1],[2],[3],[4]])) + + assert_raises(ValueError, lambda: A@A) + assert_raises(ValueError, lambda: A**2) + + for matvecsA, matvecsB in product(get_matvecs(self.A), + get_matvecs(self.B)): + A = interface.LinearOperator(**matvecsA) + B = interface.LinearOperator(**matvecsB) + # AtimesB = np.array([[22, 28], [49, 64]]) + AtimesB = self.A.dot(self.B) + X = np.array([[1, 2], [3, 4]]) + + assert_equal((A @ B).rmatmat(X), np.dot((AtimesB).T, X)) + assert_equal((2j * A @ B).rmatmat(X), + np.dot((2j * AtimesB).T.conj(), X)) + + assert_equal((A@B)@[1,1], [50,113]) + assert_equal((A@B)@[[1],[1]], [[50],[113]]) + assert_equal((A@B).matmat([[1],[1]]), [[50],[113]]) + + assert_equal((A @ B).rmatvec([1, 1]), [71, 92]) + assert_equal((A @ B).H.matvec([1, 1]), [71, 92]) + assert_equal((A @ B).adjoint().matvec([1, 1]), [71, 92]) + + assert_(isinstance(A@B, interface._ProductLinearOperator)) + + assert_raises(ValueError, lambda: A+B) + assert_raises(ValueError, lambda: A**2) + + z = A@B + assert_(len(z.args) == 2 and z.args[0] is A and z.args[1] is B) + + for matvecsC in get_matvecs(self.C): + C = interface.LinearOperator(**matvecsC) + X = np.array([[1, 2], [3, 4]]) + + assert_equal(C.rmatmat(X), np.dot((self.C).T, X)) + assert_equal((C**2).rmatmat(X), + np.dot((np.dot(self.C, self.C)).T, X)) + + assert_equal((C**2)@[1,1], [17,37]) + assert_equal((C**2).rmatvec([1, 1]), [22, 32]) + assert_equal((C**2).H.matvec([1, 1]), [22, 32]) + assert_equal((C**2).adjoint().matvec([1, 1]), [22, 32]) + assert_equal((C**2).matmat([[1],[1]]), [[17],[37]]) + + assert_(isinstance(C**2, interface._PowerLinearOperator)) + + def test_matmul(self): + D = {'shape': self.A.shape, + 'matvec': lambda x: np.dot(self.A, x).reshape(self.A.shape[0]), + 'rmatvec': lambda x: np.dot(self.A.T.conj(), + x).reshape(self.A.shape[1]), + 'rmatmat': lambda x: np.dot(self.A.T.conj(), x), + 'matmat': lambda x: np.dot(self.A, x)} + A = interface.LinearOperator(**D) + B = np.array([[1 + 1j, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + b = B[0] + + assert_equal(operator.matmul(A, b), A * b) + assert_equal(operator.matmul(A, b.reshape(-1, 1)), A * b.reshape(-1, 1)) + assert_equal(operator.matmul(A, B), A @ B) + assert_equal(operator.matmul(b, A.H), b * A.H) + assert_equal(operator.matmul(b, A.adjoint()), b * A.adjoint()) + assert_equal(operator.matmul(b.reshape(1, -1), A.H), b.reshape(1, -1) * A.H) + assert_equal(operator.matmul(b.reshape(1, -1), A.adjoint()), + b.reshape(1, -1) * A.adjoint()) + assert_equal(operator.matmul(B, A.H), B @ A.H) + assert_equal(operator.matmul(B, A.adjoint()), B @ A.adjoint()) + assert_raises(ValueError, operator.matmul, A, 2) + assert_raises(ValueError, operator.matmul, 2, A) + + +class TestAsLinearOperator: + def setup_method(self): + self.cases = [] + + def make_cases(original, dtype): + cases = [] + + cases.append((matrix(original, dtype=dtype), original)) + cases.append((np.array(original, dtype=dtype), original)) + cases.append((sparse.csr_array(original, dtype=dtype), original)) + + # Test default implementations of _adjoint and _rmatvec, which + # refer to each other. + def mv(x, dtype): + y = original.dot(x) + if len(x.shape) == 2: + y = y.reshape(-1, 1) + return y + + def rmv(x, dtype): + return original.T.conj().dot(x) + + class BaseMatlike(interface.LinearOperator): + args = () + + def __init__(self, dtype): + self.dtype = np.dtype(dtype) + self.shape = original.shape + + def _matvec(self, x): + return mv(x, self.dtype) + + class HasRmatvec(BaseMatlike): + args = () + + def _rmatvec(self,x): + return rmv(x, self.dtype) + + class HasAdjoint(BaseMatlike): + args = () + + def _adjoint(self): + shape = self.shape[1], self.shape[0] + matvec = partial(rmv, dtype=self.dtype) + rmatvec = partial(mv, dtype=self.dtype) + return interface.LinearOperator(matvec=matvec, + rmatvec=rmatvec, + dtype=self.dtype, + shape=shape) + + class HasRmatmat(HasRmatvec): + def _matmat(self, x): + return original.dot(x) + + def _rmatmat(self, x): + return original.T.conj().dot(x) + + cases.append((HasRmatvec(dtype), original)) + cases.append((HasAdjoint(dtype), original)) + cases.append((HasRmatmat(dtype), original)) + return cases + + original = np.array([[1,2,3], [4,5,6]]) + self.cases += make_cases(original, np.int32) + self.cases += make_cases(original, np.float32) + self.cases += make_cases(original, np.float64) + self.cases += [(interface.aslinearoperator(M).T, A.T) + for M, A in make_cases(original.T, np.float64)] + self.cases += [(interface.aslinearoperator(M).H, A.T.conj()) + for M, A in make_cases(original.T, np.float64)] + self.cases += [(interface.aslinearoperator(M).adjoint(), A.T.conj()) + for M, A in make_cases(original.T, np.float64)] + + original = np.array([[1, 2j, 3j], [4j, 5j, 6]]) + self.cases += make_cases(original, np.complex128) + self.cases += [(interface.aslinearoperator(M).T, A.T) + for M, A in make_cases(original.T, np.complex128)] + self.cases += [(interface.aslinearoperator(M).H, A.T.conj()) + for M, A in make_cases(original.T, np.complex128)] + self.cases += [(interface.aslinearoperator(M).adjoint(), A.T.conj()) + for M, A in make_cases(original.T, np.complex128)] + + def test_basic(self): + + for M, A_array in self.cases: + A = interface.aslinearoperator(M) + M,N = A.shape + + xs = [np.array([1, 2, 3]), + np.array([[1], [2], [3]])] + ys = [np.array([1, 2]), np.array([[1], [2]])] + + if A.dtype == np.complex128: + xs += [np.array([1, 2j, 3j]), + np.array([[1], [2j], [3j]])] + ys += [np.array([1, 2j]), np.array([[1], [2j]])] + + x2 = np.array([[1, 4], [2, 5], [3, 6]]) + + for x in xs: + assert_equal(A.matvec(x), A_array.dot(x)) + assert_equal(A @ x, A_array.dot(x)) + + assert_equal(A.matmat(x2), A_array.dot(x2)) + assert_equal(A @ x2, A_array.dot(x2)) + + for y in ys: + assert_equal(A.rmatvec(y), A_array.T.conj().dot(y)) + assert_equal(A.T.matvec(y), A_array.T.dot(y)) + assert_equal(A.H.matvec(y), A_array.T.conj().dot(y)) + assert_equal(A.adjoint().matvec(y), A_array.T.conj().dot(y)) + + for y in ys: + if y.ndim < 2: + continue + assert_equal(A.rmatmat(y), A_array.T.conj().dot(y)) + assert_equal(A.T.matmat(y), A_array.T.dot(y)) + assert_equal(A.H.matmat(y), A_array.T.conj().dot(y)) + assert_equal(A.adjoint().matmat(y), A_array.T.conj().dot(y)) + + if hasattr(M,'dtype'): + assert_equal(A.dtype, M.dtype) + + assert_(hasattr(A, 'args')) + + def test_dot(self): + + for M, A_array in self.cases: + A = interface.aslinearoperator(M) + M,N = A.shape + + x0 = np.array([1, 2, 3]) + x1 = np.array([[1], [2], [3]]) + x2 = np.array([[1, 4], [2, 5], [3, 6]]) + + assert_equal(A.dot(x0), A_array.dot(x0)) + assert_equal(A.dot(x1), A_array.dot(x1)) + assert_equal(A.dot(x2), A_array.dot(x2)) + + +def test_repr(): + A = interface.LinearOperator(shape=(1, 1), matvec=lambda x: 1) + repr_A = repr(A) + assert_('unspecified dtype' not in repr_A, repr_A) + + +def test_identity(): + ident = interface.IdentityOperator((3, 3)) + assert_equal(ident @ [1, 2, 3], [1, 2, 3]) + assert_equal(ident.dot(np.arange(9).reshape(3, 3)).ravel(), np.arange(9)) + + assert_raises(ValueError, ident.matvec, [1, 2, 3, 4]) + + +def test_attributes(): + A = interface.aslinearoperator(np.arange(16).reshape(4, 4)) + + def always_four_ones(x): + x = np.asarray(x) + assert_(x.shape == (3,) or x.shape == (3, 1)) + return np.ones(4) + + B = interface.LinearOperator(shape=(4, 3), matvec=always_four_ones) + + ops = [A, B, A * B, A @ B, A.H, A.adjoint(), A + A, B + B, A**4] + for op in ops: + assert_(hasattr(op, "dtype")) + assert_(hasattr(op, "shape")) + assert_(hasattr(op, "_matvec")) + +def matvec(x): + """ Needed for test_pickle as local functions are not pickleable """ + return np.zeros(3) + +def test_pickle(): + import pickle + + for protocol in range(pickle.HIGHEST_PROTOCOL + 1): + A = interface.LinearOperator((3, 3), matvec) + s = pickle.dumps(A, protocol=protocol) + B = pickle.loads(s) + + for k in A.__dict__: + assert_equal(getattr(A, k), getattr(B, k)) + + +@pytest.mark.thread_unsafe +def test_inheritance(): + class Empty(interface.LinearOperator): + pass + + with warns(RuntimeWarning, match="should implement at least"): + assert_raises(TypeError, Empty) + + class Identity(interface.LinearOperator): + def __init__(self, n): + super().__init__(dtype=None, shape=(n, n)) + + def _matvec(self, x): + return x + + id3 = Identity(3) + assert_equal(id3.matvec([1, 2, 3]), [1, 2, 3]) + assert_raises(NotImplementedError, id3.rmatvec, [4, 5, 6]) + + class MatmatOnly(interface.LinearOperator): + def __init__(self, A): + super().__init__(A.dtype, A.shape) + self.A = A + + def _matmat(self, x): + return self.A.dot(x) + + mm = MatmatOnly(np.random.randn(5, 3)) + assert_equal(mm.matvec(np.random.randn(3)).shape, (5,)) + +def test_dtypes_of_operator_sum(): + # gh-6078 + + mat_complex = np.random.rand(2,2) + 1j * np.random.rand(2,2) + mat_real = np.random.rand(2,2) + + complex_operator = interface.aslinearoperator(mat_complex) + real_operator = interface.aslinearoperator(mat_real) + + sum_complex = complex_operator + complex_operator + sum_real = real_operator + real_operator + + assert_equal(sum_real.dtype, np.float64) + assert_equal(sum_complex.dtype, np.complex128) + +def test_no_double_init(): + call_count = [0] + + def matvec(v): + call_count[0] += 1 + return v + + # It should call matvec exactly once (in order to determine the + # operator dtype) + interface.LinearOperator((2, 2), matvec=matvec) + assert_equal(call_count[0], 1) + +INT_DTYPES = (np.int8, np.int16, np.int32, np.int64) +REAL_DTYPES = (np.float32, np.float64, np.longdouble) +COMPLEX_DTYPES = (np.complex64, np.complex128, np.clongdouble) +INEXACTDTYPES = REAL_DTYPES + COMPLEX_DTYPES +ALLDTYPES = INT_DTYPES + INEXACTDTYPES + + +@pytest.mark.parametrize("test_dtype", ALLDTYPES) +def test_determine_lo_dtype_from_matvec(test_dtype): + # gh-19209 + scalar = np.array(1, dtype=test_dtype) + def mv(v): + return np.array([scalar * v[0], v[1]]) + + lo = interface.LinearOperator((2, 2), matvec=mv) + assert lo.dtype == np.dtype(test_dtype) + +def test_determine_lo_dtype_for_int(): + # gh-19209 + # test Python int larger than int8 max cast to some int + def mv(v): + return np.array([128 * v[0], v[1]]) + + lo = interface.LinearOperator((2, 2), matvec=mv) + assert lo.dtype in INT_DTYPES + +def test_adjoint_conjugate(): + X = np.array([[1j]]) + A = interface.aslinearoperator(X) + + B = 1j * A + Y = 1j * X + + v = np.array([1]) + + assert_equal(B.dot(v), Y.dot(v)) + assert_equal(B.H.dot(v), Y.T.conj().dot(v)) + assert_equal(B.adjoint().dot(v), Y.T.conj().dot(v)) + +def test_ndim(): + X = np.array([[1]]) + A = interface.aslinearoperator(X) + assert_equal(A.ndim, 2) + +def test_transpose_noconjugate(): + X = np.array([[1j]]) + A = interface.aslinearoperator(X) + + B = 1j * A + Y = 1j * X + + v = np.array([1]) + + assert_equal(B.dot(v), Y.dot(v)) + assert_equal(B.T.dot(v), Y.T.dot(v)) + +def test_transpose_multiplication(): + class MyMatrix(interface.LinearOperator): + def __init__(self, A): + super().__init__(A.dtype, A.shape) + self.A = A + def _matmat(self, other): return self.A @ other + def _rmatmat(self, other): return self.A.T @ other + + A = MyMatrix(np.array([[1, 2], [3, 4]])) + X = np.array([1, 2]) + B = np.array([[10, 20], [30, 40]]) + X2 = X.reshape(-1, 1) + Y = np.array([[1, 2], [3, 4]]) + + assert_equal(A @ B, Y @ B) + assert_equal(B.T @ A, B.T @ Y) + assert_equal(A.T @ B, Y.T @ B) + assert_equal(A @ X, Y @ X) + assert_equal(X.T @ A, X.T @ Y) + assert_equal(A.T @ X, Y.T @ X) + assert_equal(A @ X2, Y @ X2) + assert_equal(X2.T @ A, X2.T @ Y) + assert_equal(A.T @ X2, Y.T @ X2) + +def test_sparse_matmat_exception(): + A = interface.LinearOperator((2, 2), matvec=lambda x: x) + B = sparse.eye_array(2) + msg = "Unable to multiply a LinearOperator with a sparse matrix." + with assert_raises(TypeError, match=msg): + A @ B + with assert_raises(TypeError, match=msg): + B @ A + with assert_raises(ValueError): + A @ np.identity(4) + with assert_raises(ValueError): + np.identity(4) @ A + + +@pytest.mark.skipif(IS_PYPY, reason="Test not meaningful on PyPy") +def test_MatrixLinearOperator_refcycle(): + # gh-10634 + # Test that MatrixLinearOperator can be automatically garbage collected + A = np.eye(2) + with assert_deallocated(interface.MatrixLinearOperator, A) as op: + op.adjoint() + del op diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_matfuncs.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_matfuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..f8296080b52e7e7da3b823fc680edf584a034283 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_matfuncs.py @@ -0,0 +1,592 @@ +# +# Created by: Pearu Peterson, March 2002 +# +""" Test functions for scipy.linalg._matfuncs module + +""" +import math + +import numpy as np +from numpy import array, eye, exp, random +from numpy.testing import ( + assert_allclose, assert_, assert_array_almost_equal, assert_equal, + assert_array_almost_equal_nulp, suppress_warnings) + +from scipy.sparse import csc_array, SparseEfficiencyWarning +from scipy.sparse._construct import eye_array +from scipy.sparse.linalg._matfuncs import (expm, _expm, + ProductOperator, MatrixPowerOperator, + _onenorm_matrix_power_nnm, matrix_power) +from scipy.sparse._sputils import matrix +from scipy.linalg import logm +from scipy.special import factorial, binom +import scipy.sparse +import scipy.sparse.linalg + + +def _burkardt_13_power(n, p): + """ + A helper function for testing matrix functions. + + Parameters + ---------- + n : integer greater than 1 + Order of the square matrix to be returned. + p : non-negative integer + Power of the matrix. + + Returns + ------- + out : ndarray representing a square matrix + A Forsythe matrix of order n, raised to the power p. + + """ + # Input validation. + if n != int(n) or n < 2: + raise ValueError('n must be an integer greater than 1') + n = int(n) + if p != int(p) or p < 0: + raise ValueError('p must be a non-negative integer') + p = int(p) + + # Construct the matrix explicitly. + a, b = divmod(p, n) + large = np.power(10.0, -n*a) + small = large * np.power(10.0, -n) + return np.diag([large]*(n-b), b) + np.diag([small]*b, b-n) + + +def test_onenorm_matrix_power_nnm(): + np.random.seed(1234) + for n in range(1, 5): + for p in range(5): + M = np.random.random((n, n)) + Mp = np.linalg.matrix_power(M, p) + observed = _onenorm_matrix_power_nnm(M, p) + expected = np.linalg.norm(Mp, 1) + assert_allclose(observed, expected) + +def test_matrix_power(): + np.random.seed(1234) + row, col = np.random.randint(0, 4, size=(2, 6)) + data = np.random.random(size=(6,)) + Amat = csc_array((data, (row, col)), shape=(4, 4)) + A = csc_array((data, (row, col)), shape=(4, 4)) + Adense = A.toarray() + for power in (2, 5, 6): + Apow = matrix_power(A, power).toarray() + Amat_pow = matrix_power(Amat, power).toarray() + Adense_pow = np.linalg.matrix_power(Adense, power) + assert_allclose(Apow, Adense_pow) + assert_allclose(Apow, Amat_pow) + + +class TestExpM: + def test_zero_ndarray(self): + a = array([[0.,0],[0,0]]) + assert_array_almost_equal(expm(a),[[1,0],[0,1]]) + + def test_zero_sparse(self): + a = csc_array([[0.,0],[0,0]]) + assert_array_almost_equal(expm(a).toarray(),[[1,0],[0,1]]) + + def test_zero_matrix(self): + a = matrix([[0.,0],[0,0]]) + assert_array_almost_equal(expm(a),[[1,0],[0,1]]) + + def test_misc_types(self): + A = expm(np.array([[1]])) + assert_allclose(expm(((1,),)), A) + assert_allclose(expm([[1]]), A) + assert_allclose(expm(matrix([[1]])), A) + assert_allclose(expm(np.array([[1]])), A) + assert_allclose(expm(csc_array([[1]])).toarray(), A) + B = expm(np.array([[1j]])) + assert_allclose(expm(((1j,),)), B) + assert_allclose(expm([[1j]]), B) + assert_allclose(expm(matrix([[1j]])), B) + assert_allclose(expm(csc_array([[1j]])).toarray(), B) + + def test_bidiagonal_sparse(self): + A = csc_array([ + [1, 3, 0], + [0, 1, 5], + [0, 0, 2]], dtype=float) + e1 = math.exp(1) + e2 = math.exp(2) + expected = np.array([ + [e1, 3*e1, 15*(e2 - 2*e1)], + [0, e1, 5*(e2 - e1)], + [0, 0, e2]], dtype=float) + observed = expm(A).toarray() + assert_array_almost_equal(observed, expected) + + def test_padecases_dtype_float(self): + for dtype in [np.float32, np.float64]: + for scale in [1e-2, 1e-1, 5e-1, 1, 10]: + A = scale * eye(3, dtype=dtype) + observed = expm(A) + expected = exp(scale, dtype=dtype) * eye(3, dtype=dtype) + assert_array_almost_equal_nulp(observed, expected, nulp=100) + + def test_padecases_dtype_complex(self): + for dtype in [np.complex64, np.complex128]: + for scale in [1e-2, 1e-1, 5e-1, 1, 10]: + A = scale * eye(3, dtype=dtype) + observed = expm(A) + expected = exp(scale, dtype=dtype) * eye(3, dtype=dtype) + assert_array_almost_equal_nulp(observed, expected, nulp=100) + + def test_padecases_dtype_sparse_float(self): + # float32 and complex64 lead to errors in spsolve/UMFpack + dtype = np.float64 + for scale in [1e-2, 1e-1, 5e-1, 1, 10]: + a = scale * eye_array(3, 3, dtype=dtype, format='csc') + e = exp(scale, dtype=dtype) * eye(3, dtype=dtype) + with suppress_warnings() as sup: + sup.filter(SparseEfficiencyWarning, "Changing the sparsity structure") + exact_onenorm = _expm(a, use_exact_onenorm=True).toarray() + inexact_onenorm = _expm(a, use_exact_onenorm=False).toarray() + assert_array_almost_equal_nulp(exact_onenorm, e, nulp=100) + assert_array_almost_equal_nulp(inexact_onenorm, e, nulp=100) + + def test_padecases_dtype_sparse_complex(self): + # float32 and complex64 lead to errors in spsolve/UMFpack + dtype = np.complex128 + for scale in [1e-2, 1e-1, 5e-1, 1, 10]: + a = scale * eye_array(3, 3, dtype=dtype, format='csc') + e = exp(scale) * eye(3, dtype=dtype) + with suppress_warnings() as sup: + sup.filter(SparseEfficiencyWarning, "Changing the sparsity structure") + assert_array_almost_equal_nulp(expm(a).toarray(), e, nulp=100) + + def test_logm_consistency(self): + random.seed(1234) + for dtype in [np.float64, np.complex128]: + for n in range(1, 10): + for scale in [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2]: + # make logm(A) be of a given scale + A = (eye(n) + random.rand(n, n) * scale).astype(dtype) + if np.iscomplexobj(A): + A = A + 1j * random.rand(n, n) * scale + assert_array_almost_equal(expm(logm(A)), A) + + def test_integer_matrix(self): + Q = np.array([ + [-3, 1, 1, 1], + [1, -3, 1, 1], + [1, 1, -3, 1], + [1, 1, 1, -3]]) + assert_allclose(expm(Q), expm(1.0 * Q)) + + def test_integer_matrix_2(self): + # Check for integer overflows + Q = np.array([[-500, 500, 0, 0], + [0, -550, 360, 190], + [0, 630, -630, 0], + [0, 0, 0, 0]], dtype=np.int16) + assert_allclose(expm(Q), expm(1.0 * Q)) + + Q = csc_array(Q) + assert_allclose(expm(Q).toarray(), expm(1.0 * Q).toarray()) + + def test_triangularity_perturbation(self): + # Experiment (1) of + # Awad H. Al-Mohy and Nicholas J. Higham (2012) + # Improved Inverse Scaling and Squaring Algorithms + # for the Matrix Logarithm. + A = np.array([ + [3.2346e-1, 3e4, 3e4, 3e4], + [0, 3.0089e-1, 3e4, 3e4], + [0, 0, 3.221e-1, 3e4], + [0, 0, 0, 3.0744e-1]], + dtype=float) + A_logm = np.array([ + [-1.12867982029050462e+00, 9.61418377142025565e+04, + -4.52485573953179264e+09, 2.92496941103871812e+14], + [0.00000000000000000e+00, -1.20101052953082288e+00, + 9.63469687211303099e+04, -4.68104828911105442e+09], + [0.00000000000000000e+00, 0.00000000000000000e+00, + -1.13289322264498393e+00, 9.53249183094775653e+04], + [0.00000000000000000e+00, 0.00000000000000000e+00, + 0.00000000000000000e+00, -1.17947533272554850e+00]], + dtype=float) + assert_allclose(expm(A_logm), A, rtol=1e-4) + + # Perturb the upper triangular matrix by tiny amounts, + # so that it becomes technically not upper triangular. + random.seed(1234) + tiny = 1e-17 + A_logm_perturbed = A_logm.copy() + A_logm_perturbed[1, 0] = tiny + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "Ill-conditioned.*") + A_expm_logm_perturbed = expm(A_logm_perturbed) + rtol = 1e-4 + atol = 100 * tiny + assert_(not np.allclose(A_expm_logm_perturbed, A, rtol=rtol, atol=atol)) + + def test_burkardt_1(self): + # This matrix is diagonal. + # The calculation of the matrix exponential is simple. + # + # This is the first of a series of matrix exponential tests + # collected by John Burkardt from the following sources. + # + # Alan Laub, + # Review of "Linear System Theory" by Joao Hespanha, + # SIAM Review, + # Volume 52, Number 4, December 2010, pages 779--781. + # + # Cleve Moler and Charles Van Loan, + # Nineteen Dubious Ways to Compute the Exponential of a Matrix, + # Twenty-Five Years Later, + # SIAM Review, + # Volume 45, Number 1, March 2003, pages 3--49. + # + # Cleve Moler, + # Cleve's Corner: A Balancing Act for the Matrix Exponential, + # 23 July 2012. + # + # Robert Ward, + # Numerical computation of the matrix exponential + # with accuracy estimate, + # SIAM Journal on Numerical Analysis, + # Volume 14, Number 4, September 1977, pages 600--610. + exp1 = np.exp(1) + exp2 = np.exp(2) + A = np.array([ + [1, 0], + [0, 2], + ], dtype=float) + desired = np.array([ + [exp1, 0], + [0, exp2], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_2(self): + # This matrix is symmetric. + # The calculation of the matrix exponential is straightforward. + A = np.array([ + [1, 3], + [3, 2], + ], dtype=float) + desired = np.array([ + [39.322809708033859, 46.166301438885753], + [46.166301438885768, 54.711576854329110], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_3(self): + # This example is due to Laub. + # This matrix is ill-suited for the Taylor series approach. + # As powers of A are computed, the entries blow up too quickly. + exp1 = np.exp(1) + exp39 = np.exp(39) + A = np.array([ + [0, 1], + [-39, -40], + ], dtype=float) + desired = np.array([ + [ + 39/(38*exp1) - 1/(38*exp39), + -np.expm1(-38) / (38*exp1)], + [ + 39*np.expm1(-38) / (38*exp1), + -1/(38*exp1) + 39/(38*exp39)], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_4(self): + # This example is due to Moler and Van Loan. + # The example will cause problems for the series summation approach, + # as well as for diagonal Pade approximations. + A = np.array([ + [-49, 24], + [-64, 31], + ], dtype=float) + U = np.array([[3, 1], [4, 2]], dtype=float) + V = np.array([[1, -1/2], [-2, 3/2]], dtype=float) + w = np.array([-17, -1], dtype=float) + desired = np.dot(U * np.exp(w), V) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_5(self): + # This example is due to Moler and Van Loan. + # This matrix is strictly upper triangular + # All powers of A are zero beyond some (low) limit. + # This example will cause problems for Pade approximations. + A = np.array([ + [0, 6, 0, 0], + [0, 0, 6, 0], + [0, 0, 0, 6], + [0, 0, 0, 0], + ], dtype=float) + desired = np.array([ + [1, 6, 18, 36], + [0, 1, 6, 18], + [0, 0, 1, 6], + [0, 0, 0, 1], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_6(self): + # This example is due to Moler and Van Loan. + # This matrix does not have a complete set of eigenvectors. + # That means the eigenvector approach will fail. + exp1 = np.exp(1) + A = np.array([ + [1, 1], + [0, 1], + ], dtype=float) + desired = np.array([ + [exp1, exp1], + [0, exp1], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_7(self): + # This example is due to Moler and Van Loan. + # This matrix is very close to example 5. + # Mathematically, it has a complete set of eigenvectors. + # Numerically, however, the calculation will be suspect. + exp1 = np.exp(1) + eps = np.spacing(1) + A = np.array([ + [1 + eps, 1], + [0, 1 - eps], + ], dtype=float) + desired = np.array([ + [exp1, exp1], + [0, exp1], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_8(self): + # This matrix was an example in Wikipedia. + exp4 = np.exp(4) + exp16 = np.exp(16) + A = np.array([ + [21, 17, 6], + [-5, -1, -6], + [4, 4, 16], + ], dtype=float) + desired = np.array([ + [13*exp16 - exp4, 13*exp16 - 5*exp4, 2*exp16 - 2*exp4], + [-9*exp16 + exp4, -9*exp16 + 5*exp4, -2*exp16 + 2*exp4], + [16*exp16, 16*exp16, 4*exp16], + ], dtype=float) * 0.25 + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_9(self): + # This matrix is due to the NAG Library. + # It is an example for function F01ECF. + A = np.array([ + [1, 2, 2, 2], + [3, 1, 1, 2], + [3, 2, 1, 2], + [3, 3, 3, 1], + ], dtype=float) + desired = np.array([ + [740.7038, 610.8500, 542.2743, 549.1753], + [731.2510, 603.5524, 535.0884, 542.2743], + [823.7630, 679.4257, 603.5524, 610.8500], + [998.4355, 823.7630, 731.2510, 740.7038], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_10(self): + # This is Ward's example #1. + # It is defective and nonderogatory. + A = np.array([ + [4, 2, 0], + [1, 4, 1], + [1, 1, 4], + ], dtype=float) + assert_allclose(sorted(scipy.linalg.eigvals(A)), (3, 3, 6)) + desired = np.array([ + [147.8666224463699, 183.7651386463682, 71.79703239999647], + [127.7810855231823, 183.7651386463682, 91.88256932318415], + [127.7810855231824, 163.6796017231806, 111.9681062463718], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_11(self): + # This is Ward's example #2. + # It is a symmetric matrix. + A = np.array([ + [29.87942128909879, 0.7815750847907159, -2.289519314033932], + [0.7815750847907159, 25.72656945571064, 8.680737820540137], + [-2.289519314033932, 8.680737820540137, 34.39400925519054], + ], dtype=float) + assert_allclose(scipy.linalg.eigvalsh(A), (20, 30, 40)) + desired = np.array([ + [ + 5.496313853692378E+15, + -1.823188097200898E+16, + -3.047577080858001E+16], + [ + -1.823188097200899E+16, + 6.060522870222108E+16, + 1.012918429302482E+17], + [ + -3.047577080858001E+16, + 1.012918429302482E+17, + 1.692944112408493E+17], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_12(self): + # This is Ward's example #3. + # Ward's algorithm has difficulty estimating the accuracy + # of its results. + A = np.array([ + [-131, 19, 18], + [-390, 56, 54], + [-387, 57, 52], + ], dtype=float) + assert_allclose(sorted(scipy.linalg.eigvals(A)), (-20, -2, -1)) + desired = np.array([ + [-1.509644158793135, 0.3678794391096522, 0.1353352811751005], + [-5.632570799891469, 1.471517758499875, 0.4060058435250609], + [-4.934938326088363, 1.103638317328798, 0.5413411267617766], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_burkardt_13(self): + # This is Ward's example #4. + # This is a version of the Forsythe matrix. + # The eigenvector problem is badly conditioned. + # Ward's algorithm has difficulty estimating the accuracy + # of its results for this problem. + # + # Check the construction of one instance of this family of matrices. + A4_actual = _burkardt_13_power(4, 1) + A4_desired = [[0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + [1e-4, 0, 0, 0]] + assert_allclose(A4_actual, A4_desired) + # Check the expm for a few instances. + for n in (2, 3, 4, 10): + # Approximate expm using Taylor series. + # This works well for this matrix family + # because each matrix in the summation, + # even before dividing by the factorial, + # is entrywise positive with max entry 10**(-floor(p/n)*n). + k = max(1, int(np.ceil(16/n))) + desired = np.zeros((n, n), dtype=float) + for p in range(n*k): + Ap = _burkardt_13_power(n, p) + assert_equal(np.min(Ap), 0) + assert_allclose(np.max(Ap), np.power(10, -np.floor(p/n)*n)) + desired += Ap / factorial(p) + actual = expm(_burkardt_13_power(n, 1)) + assert_allclose(actual, desired) + + def test_burkardt_14(self): + # This is Moler's example. + # This badly scaled matrix caused problems for MATLAB's expm(). + A = np.array([ + [0, 1e-8, 0], + [-(2e10 + 4e8/6.), -3, 2e10], + [200./3., 0, -200./3.], + ], dtype=float) + desired = np.array([ + [0.446849468283175, 1.54044157383952e-09, 0.462811453558774], + [-5743067.77947947, -0.0152830038686819, -4526542.71278401], + [0.447722977849494, 1.54270484519591e-09, 0.463480648837651], + ], dtype=float) + actual = expm(A) + assert_allclose(actual, desired) + + def test_pascal(self): + # Test pascal triangle. + # Nilpotent exponential, used to trigger a failure (gh-8029) + + for scale in [1.0, 1e-3, 1e-6]: + for n in range(0, 80, 3): + sc = scale ** np.arange(n, -1, -1) + if np.any(sc < 1e-300): + break + + A = np.diag(np.arange(1, n + 1), -1) * scale + B = expm(A) + + got = B + expected = binom(np.arange(n + 1)[:,None], + np.arange(n + 1)[None,:]) * sc[None,:] / sc[:,None] + atol = 1e-13 * abs(expected).max() + assert_allclose(got, expected, atol=atol) + + def test_matrix_input(self): + # Large np.matrix inputs should work, gh-5546 + A = np.zeros((200, 200)) + A[-1,0] = 1 + B0 = expm(A) + with suppress_warnings() as sup: + sup.filter(DeprecationWarning, "the matrix subclass.*") + sup.filter(PendingDeprecationWarning, "the matrix subclass.*") + B = expm(np.matrix(A)) + assert_allclose(B, B0) + + def test_exp_sinch_overflow(self): + # Check overflow in intermediate steps is fixed (gh-11839) + L = np.array([[1.0, -0.5, -0.5, 0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, -0.5, -0.5, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0, 0.0, -0.5, -0.5], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]) + + E0 = expm(-L) + E1 = expm(-2**11 * L) + E2 = E0 + for j in range(11): + E2 = E2 @ E2 + + assert_allclose(E1, E2) + + +class TestOperators: + + def test_product_operator(self): + random.seed(1234) + n = 5 + k = 2 + nsamples = 10 + for i in range(nsamples): + A = np.random.randn(n, n) + B = np.random.randn(n, n) + C = np.random.randn(n, n) + D = np.random.randn(n, k) + op = ProductOperator(A, B, C) + assert_allclose(op.matmat(D), A.dot(B).dot(C).dot(D)) + assert_allclose(op.T.matmat(D), (A.dot(B).dot(C)).T.dot(D)) + + def test_matrix_power_operator(self): + random.seed(1234) + n = 5 + k = 2 + p = 3 + nsamples = 10 + for i in range(nsamples): + A = np.random.randn(n, n) + B = np.random.randn(n, k) + op = MatrixPowerOperator(A, p) + assert_allclose(op.matmat(B), np.linalg.matrix_power(A, p).dot(B)) + assert_allclose(op.T.matmat(B), np.linalg.matrix_power(A, p).T.dot(B)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_norm.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..3aa5e6d2cad3e57a928f368b26a3da430a122620 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_norm.py @@ -0,0 +1,154 @@ +"""Test functions for the sparse.linalg.norm module +""" + +import pytest +import numpy as np +from numpy.linalg import norm as npnorm +from numpy.testing import assert_allclose, assert_equal +from pytest import raises as assert_raises + +import scipy.sparse +from scipy.sparse.linalg import norm as spnorm + + +# https://github.com/scipy/scipy/issues/16031 +# https://github.com/scipy/scipy/issues/21690 +def test_sparray_norm(): + row = np.array([0, 0, 1, 1]) + col = np.array([0, 1, 2, 3]) + data = np.array([4, 5, 7, 9]) + test_arr = scipy.sparse.coo_array((data, (row, col)), shape=(2, 4)) + test_mat = scipy.sparse.coo_matrix((data, (row, col)), shape=(2, 4)) + for ord in (1, np.inf, None): + for ax in [0, 1, None, (0, 1), (1, 0)]: + for A in (test_arr, test_mat): + expected = npnorm(A.toarray(), ord=ord, axis=ax) + actual = spnorm(A, ord=ord, axis=ax) + assert hasattr(actual, "dtype") + assert_equal(actual, expected) + # test 1d array and 1d-like (column) matrix + test_arr_1d = scipy.sparse.coo_array((data, (col,)), shape=(4,)) + test_mat_col = scipy.sparse.coo_matrix((data, (col, [0, 0, 0, 0])), shape=(4, 1)) + for ord in (1, np.inf, None): + for ax in [0, None]: + for A in (test_arr_1d, test_mat_col): + expected = npnorm(A.toarray(), ord=ord, axis=ax) + assert_equal(spnorm(A, ord=ord, axis=ax), expected) + + +class TestNorm: + def setup_method(self): + a = np.arange(9) - 4 + b = a.reshape((3, 3)) + self.b = scipy.sparse.csr_array(b) + + @pytest.mark.thread_unsafe + def test_matrix_norm(self): + + # Frobenius norm is the default + assert_allclose(spnorm(self.b), 7.745966692414834) + assert_allclose(spnorm(self.b, 'fro'), 7.745966692414834) + + assert_allclose(spnorm(self.b, np.inf), 9) + assert_allclose(spnorm(self.b, -np.inf), 2) + assert_allclose(spnorm(self.b, 1), 7) + assert_allclose(spnorm(self.b, -1), 6) + # Only floating or complex floating dtype supported by svds. + with pytest.warns(UserWarning, match="The problem size"): + assert_allclose(spnorm(self.b.astype(np.float64), 2), + 7.348469228349534) + + # _multi_svd_norm is not implemented for sparse array + assert_raises(NotImplementedError, spnorm, self.b, -2) + + def test_matrix_norm_axis(self): + for m, axis in ((self.b, None), (self.b, (0, 1)), (self.b.T, (1, 0))): + assert_allclose(spnorm(m, axis=axis), 7.745966692414834) + assert_allclose(spnorm(m, 'fro', axis=axis), 7.745966692414834) + assert_allclose(spnorm(m, np.inf, axis=axis), 9) + assert_allclose(spnorm(m, -np.inf, axis=axis), 2) + assert_allclose(spnorm(m, 1, axis=axis), 7) + assert_allclose(spnorm(m, -1, axis=axis), 6) + + def test_vector_norm(self): + v = [4.5825756949558398, 4.2426406871192848, 4.5825756949558398] + for m, a in (self.b, 0), (self.b.T, 1): + for axis in a, (a, ), a-2, (a-2, ): + assert_allclose(spnorm(m, 1, axis=axis), [7, 6, 7]) + assert_allclose(spnorm(m, np.inf, axis=axis), [4, 3, 4]) + assert_allclose(spnorm(m, axis=axis), v) + assert_allclose(spnorm(m, ord=2, axis=axis), v) + assert_allclose(spnorm(m, ord=None, axis=axis), v) + + def test_norm_exceptions(self): + m = self.b + assert_raises(TypeError, spnorm, m, None, 1.5) + assert_raises(TypeError, spnorm, m, None, [2]) + assert_raises(ValueError, spnorm, m, None, ()) + assert_raises(ValueError, spnorm, m, None, (0, 1, 2)) + assert_raises(ValueError, spnorm, m, None, (0, 0)) + assert_raises(ValueError, spnorm, m, None, (0, 2)) + assert_raises(ValueError, spnorm, m, None, (-3, 0)) + assert_raises(ValueError, spnorm, m, None, 2) + assert_raises(ValueError, spnorm, m, None, -3) + assert_raises(ValueError, spnorm, m, 'plate_of_shrimp', 0) + assert_raises(ValueError, spnorm, m, 'plate_of_shrimp', (0, 1)) + + +class TestVsNumpyNorm: + _sparse_types = ( + scipy.sparse.bsr_array, + scipy.sparse.coo_array, + scipy.sparse.csc_array, + scipy.sparse.csr_array, + scipy.sparse.dia_array, + scipy.sparse.dok_array, + scipy.sparse.lil_array, + ) + _test_matrices = ( + (np.arange(9) - 4).reshape((3, 3)), + [ + [1, 2, 3], + [-1, 1, 4]], + [ + [1, 0, 3], + [-1, 1, 4j]], + ) + + def test_sparse_matrix_norms(self): + for sparse_type in self._sparse_types: + for M in self._test_matrices: + S = sparse_type(M) + assert_allclose(spnorm(S), npnorm(M)) + assert_allclose(spnorm(S, 'fro'), npnorm(M, 'fro')) + assert_allclose(spnorm(S, np.inf), npnorm(M, np.inf)) + assert_allclose(spnorm(S, -np.inf), npnorm(M, -np.inf)) + assert_allclose(spnorm(S, 1), npnorm(M, 1)) + assert_allclose(spnorm(S, -1), npnorm(M, -1)) + + def test_sparse_matrix_norms_with_axis(self): + for sparse_type in self._sparse_types: + for M in self._test_matrices: + S = sparse_type(M) + for axis in None, (0, 1), (1, 0): + assert_allclose(spnorm(S, axis=axis), npnorm(M, axis=axis)) + for ord in 'fro', np.inf, -np.inf, 1, -1: + assert_allclose(spnorm(S, ord, axis=axis), + npnorm(M, ord, axis=axis)) + # Some numpy matrix norms are allergic to negative axes. + for axis in (-2, -1), (-1, -2), (1, -2): + assert_allclose(spnorm(S, axis=axis), npnorm(M, axis=axis)) + assert_allclose(spnorm(S, 'f', axis=axis), + npnorm(M, 'f', axis=axis)) + assert_allclose(spnorm(S, 'fro', axis=axis), + npnorm(M, 'fro', axis=axis)) + + def test_sparse_vector_norms(self): + for sparse_type in self._sparse_types: + for M in self._test_matrices: + S = sparse_type(M) + for axis in (0, 1, -1, -2, (0, ), (1, ), (-1, ), (-2, )): + assert_allclose(spnorm(S, axis=axis), npnorm(M, axis=axis)) + for ord in None, 2, np.inf, -np.inf, 1, 0.5, 0.42: + assert_allclose(spnorm(S, ord, axis=axis), + npnorm(M, ord, axis=axis)) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_onenormest.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_onenormest.py new file mode 100644 index 0000000000000000000000000000000000000000..4e5cccb9b5114f1f16d13d70b8d296b9bb8a422b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_onenormest.py @@ -0,0 +1,252 @@ +"""Test functions for the sparse.linalg._onenormest module +""" + +import numpy as np +from numpy.testing import assert_allclose, assert_equal, assert_ +import pytest +import scipy.linalg +import scipy.sparse.linalg +from scipy.sparse.linalg._onenormest import _onenormest_core, _algorithm_2_2 + + +class MatrixProductOperator(scipy.sparse.linalg.LinearOperator): + """ + This is purely for onenormest testing. + """ + + def __init__(self, A, B): + if A.ndim != 2 or B.ndim != 2: + raise ValueError('expected ndarrays representing matrices') + if A.shape[1] != B.shape[0]: + raise ValueError('incompatible shapes') + self.A = A + self.B = B + self.ndim = 2 + self.shape = (A.shape[0], B.shape[1]) + + def _matvec(self, x): + return np.dot(self.A, np.dot(self.B, x)) + + def _rmatvec(self, x): + return np.dot(np.dot(x, self.A), self.B) + + def _matmat(self, X): + return np.dot(self.A, np.dot(self.B, X)) + + @property + def T(self): + return MatrixProductOperator(self.B.T, self.A.T) + + +class TestOnenormest: + + @pytest.mark.xslow + def test_onenormest_table_3_t_2(self): + # This will take multiple seconds if your computer is slow like mine. + # It is stochastic, so the tolerance could be too strict. + np.random.seed(1234) + t = 2 + n = 100 + itmax = 5 + nsamples = 5000 + observed = [] + expected = [] + nmult_list = [] + nresample_list = [] + for i in range(nsamples): + A = scipy.linalg.inv(np.random.randn(n, n)) + est, v, w, nmults, nresamples = _onenormest_core(A, A.T, t, itmax) + observed.append(est) + expected.append(scipy.linalg.norm(A, 1)) + nmult_list.append(nmults) + nresample_list.append(nresamples) + observed = np.array(observed, dtype=float) + expected = np.array(expected, dtype=float) + relative_errors = np.abs(observed - expected) / expected + + # check the mean underestimation ratio + underestimation_ratio = observed / expected + assert_(0.99 < np.mean(underestimation_ratio) < 1.0) + + # check the max and mean required column resamples + assert_equal(np.max(nresample_list), 2) + assert_(0.05 < np.mean(nresample_list) < 0.2) + + # check the proportion of norms computed exactly correctly + nexact = np.count_nonzero(relative_errors < 1e-14) + proportion_exact = nexact / float(nsamples) + assert_(0.9 < proportion_exact < 0.95) + + # check the average number of matrix*vector multiplications + assert_(3.5 < np.mean(nmult_list) < 4.5) + + @pytest.mark.xslow + def test_onenormest_table_4_t_7(self): + # This will take multiple seconds if your computer is slow like mine. + # It is stochastic, so the tolerance could be too strict. + np.random.seed(1234) + t = 7 + n = 100 + itmax = 5 + nsamples = 5000 + observed = [] + expected = [] + nmult_list = [] + nresample_list = [] + for i in range(nsamples): + A = np.random.randint(-1, 2, size=(n, n)) + est, v, w, nmults, nresamples = _onenormest_core(A, A.T, t, itmax) + observed.append(est) + expected.append(scipy.linalg.norm(A, 1)) + nmult_list.append(nmults) + nresample_list.append(nresamples) + observed = np.array(observed, dtype=float) + expected = np.array(expected, dtype=float) + relative_errors = np.abs(observed - expected) / expected + + # check the mean underestimation ratio + underestimation_ratio = observed / expected + assert_(0.90 < np.mean(underestimation_ratio) < 0.99) + + # check the required column resamples + assert_equal(np.max(nresample_list), 0) + + # check the proportion of norms computed exactly correctly + nexact = np.count_nonzero(relative_errors < 1e-14) + proportion_exact = nexact / float(nsamples) + assert_(0.15 < proportion_exact < 0.25) + + # check the average number of matrix*vector multiplications + assert_(3.5 < np.mean(nmult_list) < 4.5) + + def test_onenormest_table_5_t_1(self): + # "note that there is no randomness and hence only one estimate for t=1" + t = 1 + n = 100 + itmax = 5 + alpha = 1 - 1e-6 + A = -scipy.linalg.inv(np.identity(n) + alpha*np.eye(n, k=1)) + first_col = np.array([1] + [0]*(n-1)) + first_row = np.array([(-alpha)**i for i in range(n)]) + B = -scipy.linalg.toeplitz(first_col, first_row) + assert_allclose(A, B) + est, v, w, nmults, nresamples = _onenormest_core(B, B.T, t, itmax) + exact_value = scipy.linalg.norm(B, 1) + underest_ratio = est / exact_value + assert_allclose(underest_ratio, 0.05, rtol=1e-4) + assert_equal(nmults, 11) + assert_equal(nresamples, 0) + # check the non-underscored version of onenormest + est_plain = scipy.sparse.linalg.onenormest(B, t=t, itmax=itmax) + assert_allclose(est, est_plain) + + @pytest.mark.xslow + def test_onenormest_table_6_t_2(self): + #TODO this test seems to give estimates that match the table, + #TODO even though no attempt has been made to deal with + #TODO complex numbers in the one-norm estimation. + # This will take multiple seconds if your computer is slow like mine. + # It is stochastic, so the tolerance could be too strict. + np.random.seed(1234) + t = 2 + n = 100 + itmax = 5 + nsamples = 5000 + observed = [] + expected = [] + nmult_list = [] + nresample_list = [] + for i in range(nsamples): + A_inv = np.random.rand(n, n) + 1j * np.random.rand(n, n) + A = scipy.linalg.inv(A_inv) + est, v, w, nmults, nresamples = _onenormest_core(A, A.T, t, itmax) + observed.append(est) + expected.append(scipy.linalg.norm(A, 1)) + nmult_list.append(nmults) + nresample_list.append(nresamples) + observed = np.array(observed, dtype=float) + expected = np.array(expected, dtype=float) + relative_errors = np.abs(observed - expected) / expected + + # check the mean underestimation ratio + underestimation_ratio = observed / expected + underestimation_ratio_mean = np.mean(underestimation_ratio) + assert_(0.90 < underestimation_ratio_mean < 0.99) + + # check the required column resamples + max_nresamples = np.max(nresample_list) + assert_equal(max_nresamples, 0) + + # check the proportion of norms computed exactly correctly + nexact = np.count_nonzero(relative_errors < 1e-14) + proportion_exact = nexact / float(nsamples) + assert_(0.7 < proportion_exact < 0.8) + + # check the average number of matrix*vector multiplications + mean_nmult = np.mean(nmult_list) + assert_(4 < mean_nmult < 5) + + def _help_product_norm_slow(self, A, B): + # for profiling + C = np.dot(A, B) + return scipy.linalg.norm(C, 1) + + def _help_product_norm_fast(self, A, B): + # for profiling + t = 2 + itmax = 5 + D = MatrixProductOperator(A, B) + est, v, w, nmults, nresamples = _onenormest_core(D, D.T, t, itmax) + return est + + @pytest.mark.slow + def test_onenormest_linear_operator(self): + # Define a matrix through its product A B. + # Depending on the shapes of A and B, + # it could be easy to multiply this product by a small matrix, + # but it could be annoying to look at all of + # the entries of the product explicitly. + np.random.seed(1234) + n = 6000 + k = 3 + A = np.random.randn(n, k) + B = np.random.randn(k, n) + fast_estimate = self._help_product_norm_fast(A, B) + exact_value = self._help_product_norm_slow(A, B) + assert_(fast_estimate <= exact_value <= 3*fast_estimate, + f'fast: {fast_estimate:g}\nexact:{exact_value:g}') + + def test_returns(self): + np.random.seed(1234) + A = scipy.sparse.rand(50, 50, 0.1) + + s0 = scipy.linalg.norm(A.toarray(), 1) + s1, v = scipy.sparse.linalg.onenormest(A, compute_v=True) + s2, w = scipy.sparse.linalg.onenormest(A, compute_w=True) + s3, v2, w2 = scipy.sparse.linalg.onenormest(A, compute_w=True, compute_v=True) + + assert_allclose(s1, s0, rtol=1e-9) + assert_allclose(np.linalg.norm(A.dot(v), 1), s0*np.linalg.norm(v, 1), rtol=1e-9) + assert_allclose(A.dot(v), w, rtol=1e-9) + + +class TestAlgorithm_2_2: + + @pytest.mark.thread_unsafe + def test_randn_inv(self): + rng = np.random.RandomState(1234) + n = 20 + nsamples = 100 + for i in range(nsamples): + + # Choose integer t uniformly between 1 and 3 inclusive. + t = rng.randint(1, 4) + + # Choose n uniformly between 10 and 40 inclusive. + n = rng.randint(10, 41) + + # Sample the inverse of a matrix with random normal entries. + A = scipy.linalg.inv(rng.randn(n, n)) + + # Compute the 1-norm bounds. + g, ind = _algorithm_2_2(A, A.T, t) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_propack.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_propack.py new file mode 100644 index 0000000000000000000000000000000000000000..a206a1ee6ebd5ddf2335008bc823a20e7bac13a7 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_propack.py @@ -0,0 +1,165 @@ +import os +import pytest + +import numpy as np +from numpy.testing import assert_allclose +from pytest import raises as assert_raises +from scipy.sparse.linalg._svdp import _svdp +from scipy.sparse import csr_array, csc_array + + +# dtype_flavour to tolerance +TOLS = { + np.float32: 1e-4, + np.float64: 1e-8, + np.complex64: 1e-4, + np.complex128: 1e-8, +} + + +def is_complex_type(dtype): + return np.dtype(dtype).kind == "c" + + +_dtypes = [] +for dtype_flavour in TOLS.keys(): + marks = [] + if is_complex_type(dtype_flavour): + marks = [pytest.mark.slow] + _dtypes.append(pytest.param(dtype_flavour, marks=marks, + id=dtype_flavour.__name__)) +_dtypes = tuple(_dtypes) # type: ignore[assignment] + + +def generate_matrix(constructor, n, m, f, + dtype=float, rseed=0, **kwargs): + """Generate a random sparse array""" + rng = np.random.RandomState(rseed) + if is_complex_type(dtype): + M = (- 5 + 10 * rng.rand(n, m) + - 5j + 10j * rng.rand(n, m)).astype(dtype) + else: + M = (-5 + 10 * rng.rand(n, m)).astype(dtype) + M[M.real > 10 * f - 5] = 0 + return constructor(M, **kwargs) + + +def assert_orthogonal(u1, u2, rtol, atol): + """Check that the first k rows of u1 and u2 are orthogonal""" + A = abs(np.dot(u1.conj().T, u2)) + assert_allclose(A, np.eye(u1.shape[1], u2.shape[1]), rtol=rtol, atol=atol) + + +def check_svdp(n, m, constructor, dtype, k, irl_mode, which, f=0.8): + tol = TOLS[dtype] + + M = generate_matrix(np.asarray, n, m, f, dtype) + Msp = constructor(M) + + u1, sigma1, vt1 = np.linalg.svd(M, full_matrices=False) + u2, sigma2, vt2, _ = _svdp(Msp, k=k, which=which, irl_mode=irl_mode, + tol=tol, rng=np.random.default_rng(0)) + + # check the which + if which.upper() == 'SM': + u1 = np.roll(u1, k, 1) + vt1 = np.roll(vt1, k, 0) + sigma1 = np.roll(sigma1, k) + + # check that singular values agree + assert_allclose(sigma1[:k], sigma2, rtol=tol, atol=tol) + + # check that singular vectors are orthogonal + assert_orthogonal(u1, u2, rtol=tol, atol=tol) + assert_orthogonal(vt1.T, vt2.T, rtol=tol, atol=tol) + + +@pytest.mark.parametrize('ctor', (np.array, csr_array, csc_array)) +@pytest.mark.parametrize('dtype', _dtypes) +@pytest.mark.parametrize('irl', (True, False)) +@pytest.mark.parametrize('which', ('LM', 'SM')) +def test_svdp(ctor, dtype, irl, which): + np.random.seed(0) + n, m, k = 10, 20, 3 + if which == 'SM' and not irl: + message = "`which`='SM' requires irl_mode=True" + with assert_raises(ValueError, match=message): + check_svdp(n, m, ctor, dtype, k, irl, which) + else: + check_svdp(n, m, ctor, dtype, k, irl, which) + + +@pytest.mark.xslow +@pytest.mark.parametrize('dtype', _dtypes) +@pytest.mark.parametrize('irl', (False, True)) +def test_examples(dtype, irl): + # Note: atol for complex64 bumped from 1e-4 to 1e-3 due to test failures + # with BLIS, Netlib, and MKL+AVX512 - see + # https://github.com/conda-forge/scipy-feedstock/pull/198#issuecomment-999180432 + atol = { + np.float32: 1.3e-4, + np.float64: 1e-9, + np.complex64: 1e-3, + np.complex128: 1e-9, + }[dtype] + + path_prefix = os.path.dirname(__file__) + # Test matrices from `illc1850.coord` and `mhd1280b.cua` distributed with + # PROPACK 2.1: http://sun.stanford.edu/~rmunk/PROPACK/ + relative_path = "propack_test_data.npz" + filename = os.path.join(path_prefix, relative_path) + with np.load(filename, allow_pickle=True) as data: + if is_complex_type(dtype): + A = data['A_complex'].item().astype(dtype) + else: + A = data['A_real'].item().astype(dtype) + + k = 200 + u, s, vh, _ = _svdp(A, k, irl_mode=irl, rng=np.random.default_rng(0)) + + # complex example matrix has many repeated singular values, so check only + # beginning non-repeated singular vectors to avoid permutations + sv_check = 27 if is_complex_type(dtype) else k + u = u[:, :sv_check] + vh = vh[:sv_check, :] + s = s[:sv_check] + + # Check orthogonality of singular vectors + assert_allclose(np.eye(u.shape[1]), u.conj().T @ u, atol=atol) + assert_allclose(np.eye(vh.shape[0]), vh @ vh.conj().T, atol=atol) + + # Ensure the norm of the difference between the np.linalg.svd and + # PROPACK reconstructed matrices is small + u3, s3, vh3 = np.linalg.svd(A.todense()) + u3 = u3[:, :sv_check] + s3 = s3[:sv_check] + vh3 = vh3[:sv_check, :] + A3 = u3 @ np.diag(s3) @ vh3 + recon = u @ np.diag(s) @ vh + assert_allclose(np.linalg.norm(A3 - recon), 0, atol=atol) + + +@pytest.mark.parametrize('shifts', (None, -10, 0, 1, 10, 70)) +@pytest.mark.parametrize('dtype', _dtypes[:2]) +def test_shifts(shifts, dtype): + rng = np.random.default_rng(0) + n, k = 70, 10 + A = rng.random((n, n)) + if shifts is not None and ((shifts < 0) or (k > min(n-1-shifts, n))): + with pytest.raises(ValueError): + _svdp(A, k, shifts=shifts, kmax=5*k, irl_mode=True, rng=rng) + else: + _svdp(A, k, shifts=shifts, kmax=5*k, irl_mode=True, rng=rng) + + +@pytest.mark.slow +@pytest.mark.xfail() +def test_shifts_accuracy(): + rng = np.random.default_rng(0) + n, k = 70, 10 + A = rng.random((n, n)).astype(np.float64) + u1, s1, vt1, _ = _svdp(A, k, shifts=None, which='SM', irl_mode=True, rng=rng) + u2, s2, vt2, _ = _svdp(A, k, shifts=32, which='SM', irl_mode=True, rng=rng) + # shifts <= 32 doesn't agree with shifts > 32 + # Does agree when which='LM' instead of 'SM' + assert_allclose(s1, s2) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_pydata_sparse.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_pydata_sparse.py new file mode 100644 index 0000000000000000000000000000000000000000..6bccb71339b65c288b1764c34f66086c670f53da --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_pydata_sparse.py @@ -0,0 +1,258 @@ +import pytest + +import numpy as np +import scipy.sparse as sp +import scipy.sparse.linalg as splin + +from numpy.testing import assert_allclose, assert_equal + +try: + import sparse +except Exception: + sparse = None + +pytestmark = pytest.mark.skipif(sparse is None, + reason="pydata/sparse not installed") + + +msg = "pydata/sparse (0.15.1) does not implement necessary operations" + + +sparse_params = (pytest.param("COO"), + pytest.param("DOK", marks=[pytest.mark.xfail(reason=msg)])) + +scipy_sparse_classes = [ + sp.bsr_array, + sp.csr_array, + sp.coo_array, + sp.csc_array, + sp.dia_array, + sp.dok_array +] + + +@pytest.fixture(params=sparse_params) +def sparse_cls(request): + return getattr(sparse, request.param) + + +@pytest.fixture(params=scipy_sparse_classes) +def sp_sparse_cls(request): + return request.param + + +@pytest.fixture +def same_matrix(sparse_cls, sp_sparse_cls): + np.random.seed(1234) + A_dense = np.random.rand(9, 9) + return sp_sparse_cls(A_dense), sparse_cls(A_dense) + + +@pytest.fixture +def matrices(sparse_cls): + np.random.seed(1234) + A_dense = np.random.rand(9, 9) + A_dense = A_dense @ A_dense.T + A_sparse = sparse_cls(A_dense) + b = np.random.rand(9) + return A_dense, A_sparse, b + + +def test_isolve_gmres(matrices): + # Several of the iterative solvers use the same + # isolve.utils.make_system wrapper code, so test just one of them. + A_dense, A_sparse, b = matrices + x, info = splin.gmres(A_sparse, b, atol=1e-15) + assert info == 0 + assert isinstance(x, np.ndarray) + assert_allclose(A_sparse @ x, b) + + +def test_lsmr(matrices): + A_dense, A_sparse, b = matrices + res0 = splin.lsmr(A_dense, b) + res = splin.lsmr(A_sparse, b) + assert_allclose(res[0], res0[0], atol=1e-3) + + +# test issue 17012 +def test_lsmr_output_shape(): + x = splin.lsmr(A=np.ones((10, 1)), b=np.zeros(10), x0=np.ones(1))[0] + assert_equal(x.shape, (1,)) + + +def test_lsqr(matrices): + A_dense, A_sparse, b = matrices + res0 = splin.lsqr(A_dense, b) + res = splin.lsqr(A_sparse, b) + assert_allclose(res[0], res0[0], atol=1e-5) + + +def test_eigs(matrices): + A_dense, A_sparse, v0 = matrices + + M_dense = np.diag(v0**2) + M_sparse = A_sparse.__class__(M_dense) + + w_dense, v_dense = splin.eigs(A_dense, k=3, v0=v0) + w, v = splin.eigs(A_sparse, k=3, v0=v0) + + assert_allclose(w, w_dense) + assert_allclose(v, v_dense) + + for M in [M_sparse, M_dense]: + w_dense, v_dense = splin.eigs(A_dense, M=M_dense, k=3, v0=v0) + w, v = splin.eigs(A_sparse, M=M, k=3, v0=v0) + + assert_allclose(w, w_dense) + assert_allclose(v, v_dense) + + w_dense, v_dense = splin.eigsh(A_dense, M=M_dense, k=3, v0=v0) + w, v = splin.eigsh(A_sparse, M=M, k=3, v0=v0) + + assert_allclose(w, w_dense) + assert_allclose(v, v_dense) + + +def test_svds(matrices): + A_dense, A_sparse, v0 = matrices + + u0, s0, vt0 = splin.svds(A_dense, k=2, v0=v0) + u, s, vt = splin.svds(A_sparse, k=2, v0=v0) + + assert_allclose(s, s0) + assert_allclose(np.abs(u), np.abs(u0)) + assert_allclose(np.abs(vt), np.abs(vt0)) + + +def test_lobpcg(matrices): + A_dense, A_sparse, x = matrices + X = x[:,None] + + w_dense, v_dense = splin.lobpcg(A_dense, X) + w, v = splin.lobpcg(A_sparse, X) + + assert_allclose(w, w_dense) + assert_allclose(v, v_dense) + + +def test_spsolve(matrices): + A_dense, A_sparse, b = matrices + b2 = np.random.rand(len(b), 3) + + x0 = splin.spsolve(sp.csc_array(A_dense), b) + x = splin.spsolve(A_sparse, b) + assert isinstance(x, np.ndarray) + assert_allclose(x, x0) + + x0 = splin.spsolve(sp.csc_array(A_dense), b) + x = splin.spsolve(A_sparse, b, use_umfpack=True) + assert isinstance(x, np.ndarray) + assert_allclose(x, x0) + + x0 = splin.spsolve(sp.csc_array(A_dense), b2) + x = splin.spsolve(A_sparse, b2) + assert isinstance(x, np.ndarray) + assert_allclose(x, x0) + + x0 = splin.spsolve(sp.csc_array(A_dense), + sp.csc_array(A_dense)) + x = splin.spsolve(A_sparse, A_sparse) + assert isinstance(x, type(A_sparse)) + assert_allclose(x.todense(), x0.todense()) + + +def test_splu(matrices): + A_dense, A_sparse, b = matrices + n = len(b) + sparse_cls = type(A_sparse) + + lu = splin.splu(A_sparse) + + assert isinstance(lu.L, sparse_cls) + assert isinstance(lu.U, sparse_cls) + + _Pr_scipy = sp.csc_array((np.ones(n), (lu.perm_r, np.arange(n)))) + _Pc_scipy = sp.csc_array((np.ones(n), (np.arange(n), lu.perm_c))) + Pr = sparse_cls.from_scipy_sparse(_Pr_scipy) + Pc = sparse_cls.from_scipy_sparse(_Pc_scipy) + A2 = Pr.T @ lu.L @ lu.U @ Pc.T + + assert_allclose(A2.todense(), A_sparse.todense()) + + z = lu.solve(A_sparse.todense()) + assert_allclose(z, np.eye(n), atol=1e-10) + + +def test_spilu(matrices): + A_dense, A_sparse, b = matrices + sparse_cls = type(A_sparse) + + lu = splin.spilu(A_sparse) + + assert isinstance(lu.L, sparse_cls) + assert isinstance(lu.U, sparse_cls) + + z = lu.solve(A_sparse.todense()) + assert_allclose(z, np.eye(len(b)), atol=1e-3) + + +def test_spsolve_triangular(matrices): + A_dense, A_sparse, b = matrices + A_sparse = sparse.tril(A_sparse) + + x = splin.spsolve_triangular(A_sparse, b) + assert_allclose(A_sparse @ x, b) + + +def test_onenormest(matrices): + A_dense, A_sparse, b = matrices + est0 = splin.onenormest(A_dense) + est = splin.onenormest(A_sparse) + assert_allclose(est, est0) + + +def test_norm(matrices): + A_dense, A_sparse, b = matrices + norm0 = splin.norm(sp.csr_array(A_dense)) + norm = splin.norm(A_sparse) + assert_allclose(norm, norm0) + + +def test_inv(matrices): + A_dense, A_sparse, b = matrices + x0 = splin.inv(sp.csc_array(A_dense)) + x = splin.inv(A_sparse) + assert_allclose(x.todense(), x0.todense()) + + +def test_expm(matrices): + A_dense, A_sparse, b = matrices + x0 = splin.expm(sp.csc_array(A_dense)) + x = splin.expm(A_sparse) + assert_allclose(x.todense(), x0.todense()) + + +def test_expm_multiply(matrices): + A_dense, A_sparse, b = matrices + x0 = splin.expm_multiply(A_dense, b) + x = splin.expm_multiply(A_sparse, b) + assert_allclose(x, x0) + + x0 = splin.expm_multiply(A_dense, A_dense) + x = splin.expm_multiply(A_sparse, A_sparse) + assert_allclose(x.todense(), x0) + + +def test_eq(same_matrix): + sp_sparse, pd_sparse = same_matrix + # temporary splint until pydata sparse support sparray equality + sp_sparse = sp.coo_matrix(sp_sparse).asformat(sp_sparse.format) + assert (sp_sparse == pd_sparse).all() + + +def test_ne(same_matrix): + sp_sparse, pd_sparse = same_matrix + # temporary splint until pydata sparse support sparray equality + sp_sparse = sp.coo_matrix(sp_sparse).asformat(sp_sparse.format) + assert not (sp_sparse != pd_sparse).any() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_special_sparse_arrays.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_special_sparse_arrays.py new file mode 100644 index 0000000000000000000000000000000000000000..5a51c93edb80e96d6b96305f35bbbe6f47c91f81 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/sparse/linalg/tests/test_special_sparse_arrays.py @@ -0,0 +1,337 @@ +import pytest +import numpy as np +from numpy.testing import assert_array_equal, assert_allclose + +from scipy.sparse import diags, csgraph +from scipy.linalg import eigh + +from scipy.sparse.linalg import LaplacianNd +from scipy.sparse.linalg._special_sparse_arrays import Sakurai +from scipy.sparse.linalg._special_sparse_arrays import MikotaPair + +INT_DTYPES = [np.int8, np.int16, np.int32, np.int64] +REAL_DTYPES = [np.float32, np.float64] +COMPLEX_DTYPES = [np.complex64, np.complex128] +ALLDTYPES = INT_DTYPES + REAL_DTYPES + COMPLEX_DTYPES + + +class TestLaplacianNd: + """ + LaplacianNd tests + """ + + @pytest.mark.parametrize('bc', ['neumann', 'dirichlet', 'periodic']) + def test_1d_specific_shape(self, bc): + lap = LaplacianNd(grid_shape=(6, ), boundary_conditions=bc) + lapa = lap.toarray() + if bc == 'neumann': + a = np.array( + [ + [-1, 1, 0, 0, 0, 0], + [1, -2, 1, 0, 0, 0], + [0, 1, -2, 1, 0, 0], + [0, 0, 1, -2, 1, 0], + [0, 0, 0, 1, -2, 1], + [0, 0, 0, 0, 1, -1], + ] + ) + elif bc == 'dirichlet': + a = np.array( + [ + [-2, 1, 0, 0, 0, 0], + [1, -2, 1, 0, 0, 0], + [0, 1, -2, 1, 0, 0], + [0, 0, 1, -2, 1, 0], + [0, 0, 0, 1, -2, 1], + [0, 0, 0, 0, 1, -2], + ] + ) + else: + a = np.array( + [ + [-2, 1, 0, 0, 0, 1], + [1, -2, 1, 0, 0, 0], + [0, 1, -2, 1, 0, 0], + [0, 0, 1, -2, 1, 0], + [0, 0, 0, 1, -2, 1], + [1, 0, 0, 0, 1, -2], + ] + ) + assert_array_equal(a, lapa) + + def test_1d_with_graph_laplacian(self): + n = 6 + G = diags(np.ones(n - 1), 1, format='dia') + Lf = csgraph.laplacian(G, symmetrized=True, form='function') + La = csgraph.laplacian(G, symmetrized=True, form='array') + grid_shape = (n,) + bc = 'neumann' + lap = LaplacianNd(grid_shape, boundary_conditions=bc) + assert_array_equal(lap(np.eye(n)), -Lf(np.eye(n))) + assert_array_equal(lap.toarray(), -La.toarray()) + # https://github.com/numpy/numpy/issues/24351 + assert_array_equal(lap.tosparse().toarray(), -La.toarray()) + + @pytest.mark.parametrize('grid_shape', [(6, ), (2, 3), (2, 3, 4)]) + @pytest.mark.parametrize('bc', ['neumann', 'dirichlet', 'periodic']) + def test_eigenvalues(self, grid_shape, bc): + lap = LaplacianNd(grid_shape, boundary_conditions=bc, dtype=np.float64) + L = lap.toarray() + eigvals = eigh(L, eigvals_only=True) + n = np.prod(grid_shape) + eigenvalues = lap.eigenvalues() + dtype = eigenvalues.dtype + atol = n * n * np.finfo(dtype).eps + # test the default ``m = None`` + assert_allclose(eigenvalues, eigvals, atol=atol) + # test every ``m > 0`` + for m in np.arange(1, n + 1): + assert_array_equal(lap.eigenvalues(m), eigenvalues[-m:]) + + @pytest.mark.parametrize('grid_shape', [(6, ), (2, 3), (2, 3, 4)]) + @pytest.mark.parametrize('bc', ['neumann', 'dirichlet', 'periodic']) + def test_eigenvectors(self, grid_shape, bc): + lap = LaplacianNd(grid_shape, boundary_conditions=bc, dtype=np.float64) + n = np.prod(grid_shape) + eigenvalues = lap.eigenvalues() + eigenvectors = lap.eigenvectors() + dtype = eigenvectors.dtype + atol = n * n * max(np.finfo(dtype).eps, np.finfo(np.double).eps) + # test the default ``m = None`` every individual eigenvector + for i in np.arange(n): + r = lap.toarray() @ eigenvectors[:, i] - eigenvectors[:, i] * eigenvalues[i] + assert_allclose(r, np.zeros_like(r), atol=atol) + # test every ``m > 0`` + for m in np.arange(1, n + 1): + e = lap.eigenvalues(m) + ev = lap.eigenvectors(m) + r = lap.toarray() @ ev - ev @ np.diag(e) + assert_allclose(r, np.zeros_like(r), atol=atol) + + @pytest.mark.parametrize('grid_shape', [(6, ), (2, 3), (2, 3, 4)]) + @pytest.mark.parametrize('bc', ['neumann', 'dirichlet', 'periodic']) + def test_toarray_tosparse_consistency(self, grid_shape, bc): + lap = LaplacianNd(grid_shape, boundary_conditions=bc) + n = np.prod(grid_shape) + assert_array_equal(lap.toarray(), lap(np.eye(n))) + assert_array_equal(lap.tosparse().toarray(), lap.toarray()) + + @pytest.mark.parametrize('dtype', ALLDTYPES) + @pytest.mark.parametrize('grid_shape', [(6, ), (2, 3), (2, 3, 4)]) + @pytest.mark.parametrize('bc', ['neumann', 'dirichlet', 'periodic']) + def test_linearoperator_shape_dtype(self, grid_shape, bc, dtype): + lap = LaplacianNd(grid_shape, boundary_conditions=bc, dtype=dtype) + n = np.prod(grid_shape) + assert lap.shape == (n, n) + assert lap.dtype == dtype + assert_array_equal( + LaplacianNd( + grid_shape, boundary_conditions=bc, dtype=dtype + ).toarray(), + LaplacianNd(grid_shape, boundary_conditions=bc) + .toarray() + .astype(dtype), + ) + assert_array_equal( + LaplacianNd(grid_shape, boundary_conditions=bc, dtype=dtype) + .tosparse() + .toarray(), + LaplacianNd(grid_shape, boundary_conditions=bc) + .tosparse() + .toarray() + .astype(dtype), + ) + + @pytest.mark.parametrize('dtype', ALLDTYPES) + @pytest.mark.parametrize('grid_shape', [(6, ), (2, 3), (2, 3, 4)]) + @pytest.mark.parametrize('bc', ['neumann', 'dirichlet', 'periodic']) + def test_dot(self, grid_shape, bc, dtype): + """ Test the dot-product for type preservation and consistency. + """ + lap = LaplacianNd(grid_shape, boundary_conditions=bc) + n = np.prod(grid_shape) + x0 = np.arange(n) + x1 = x0.reshape((-1, 1)) + x2 = np.arange(2 * n).reshape((n, 2)) + input_set = [x0, x1, x2] + for x in input_set: + y = lap.dot(x.astype(dtype)) + assert x.shape == y.shape + assert y.dtype == dtype + if x.ndim == 2: + yy = lap.toarray() @ x.astype(dtype) + assert yy.dtype == dtype + np.array_equal(y, yy) + + def test_boundary_conditions_value_error(self): + with pytest.raises(ValueError, match="Unknown value 'robin'"): + LaplacianNd(grid_shape=(6, ), boundary_conditions='robin') + + +class TestSakurai: + """ + Sakurai tests + """ + + def test_specific_shape(self): + sak = Sakurai(6) + assert_array_equal(sak.toarray(), sak(np.eye(6))) + a = np.array( + [ + [ 5, -4, 1, 0, 0, 0], + [-4, 6, -4, 1, 0, 0], + [ 1, -4, 6, -4, 1, 0], + [ 0, 1, -4, 6, -4, 1], + [ 0, 0, 1, -4, 6, -4], + [ 0, 0, 0, 1, -4, 5] + ] + ) + + np.array_equal(a, sak.toarray()) + np.array_equal(sak.tosparse().toarray(), sak.toarray()) + ab = np.array( + [ + [ 1, 1, 1, 1, 1, 1], + [-4, -4, -4, -4, -4, -4], + [ 5, 6, 6, 6, 6, 5] + ] + ) + np.array_equal(ab, sak.tobanded()) + e = np.array( + [0.03922866, 0.56703972, 2.41789479, 5.97822974, + 10.54287655, 14.45473055] + ) + np.array_equal(e, sak.eigenvalues()) + np.array_equal(e[:2], sak.eigenvalues(2)) + + # `Sakurai` default `dtype` is `np.int8` as its entries are small integers + @pytest.mark.parametrize('dtype', ALLDTYPES) + def test_linearoperator_shape_dtype(self, dtype): + n = 7 + sak = Sakurai(n, dtype=dtype) + assert sak.shape == (n, n) + assert sak.dtype == dtype + assert_array_equal(sak.toarray(), Sakurai(n).toarray().astype(dtype)) + assert_array_equal(sak.tosparse().toarray(), + Sakurai(n).tosparse().toarray().astype(dtype)) + + @pytest.mark.parametrize('dtype', ALLDTYPES) + @pytest.mark.parametrize('argument_dtype', ALLDTYPES) + def test_dot(self, dtype, argument_dtype): + """ Test the dot-product for type preservation and consistency. + """ + result_dtype = np.promote_types(argument_dtype, dtype) + n = 5 + sak = Sakurai(n) + x0 = np.arange(n) + x1 = x0.reshape((-1, 1)) + x2 = np.arange(2 * n).reshape((n, 2)) + input_set = [x0, x1, x2] + for x in input_set: + y = sak.dot(x.astype(argument_dtype)) + assert x.shape == y.shape + assert np.can_cast(y.dtype, result_dtype) + if x.ndim == 2: + ya = sak.toarray() @ x.astype(argument_dtype) + np.array_equal(y, ya) + assert np.can_cast(ya.dtype, result_dtype) + ys = sak.tosparse() @ x.astype(argument_dtype) + np.array_equal(y, ys) + assert np.can_cast(ys.dtype, result_dtype) + +class TestMikotaPair: + """ + MikotaPair tests + """ + # both MikotaPair `LinearOperator`s share the same dtype + # while `MikotaK` `dtype` can be as small as its default `np.int32` + # since its entries are integers, the `MikotaM` involves inverses + # so its smallest still accurate `dtype` is `np.float32` + tested_types = REAL_DTYPES + COMPLEX_DTYPES + + def test_specific_shape(self): + n = 6 + mik = MikotaPair(n) + mik_k = mik.k + mik_m = mik.m + assert_array_equal(mik_k.toarray(), mik_k(np.eye(n))) + assert_array_equal(mik_m.toarray(), mik_m(np.eye(n))) + + k = np.array( + [ + [11, -5, 0, 0, 0, 0], + [-5, 9, -4, 0, 0, 0], + [ 0, -4, 7, -3, 0, 0], + [ 0, 0, -3, 5, -2, 0], + [ 0, 0, 0, -2, 3, -1], + [ 0, 0, 0, 0, -1, 1] + ] + ) + np.array_equal(k, mik_k.toarray()) + np.array_equal(mik_k.tosparse().toarray(), k) + kb = np.array( + [ + [ 0, -5, -4, -3, -2, -1], + [11, 9, 7, 5, 3, 1] + ] + ) + np.array_equal(kb, mik_k.tobanded()) + + minv = np.arange(1, n + 1) + np.array_equal(np.diag(1. / minv), mik_m.toarray()) + np.array_equal(mik_m.tosparse().toarray(), mik_m.toarray()) + np.array_equal(1. / minv, mik_m.tobanded()) + + e = np.array([ 1, 4, 9, 16, 25, 36]) + np.array_equal(e, mik.eigenvalues()) + np.array_equal(e[:2], mik.eigenvalues(2)) + + @pytest.mark.parametrize('dtype', tested_types) + def test_linearoperator_shape_dtype(self, dtype): + n = 7 + mik = MikotaPair(n, dtype=dtype) + mik_k = mik.k + mik_m = mik.m + assert mik_k.shape == (n, n) + assert mik_k.dtype == dtype + assert mik_m.shape == (n, n) + assert mik_m.dtype == dtype + mik_default_dtype = MikotaPair(n) + mikd_k = mik_default_dtype.k + mikd_m = mik_default_dtype.m + assert mikd_k.shape == (n, n) + assert mikd_k.dtype == np.float64 + assert mikd_m.shape == (n, n) + assert mikd_m.dtype == np.float64 + assert_array_equal(mik_k.toarray(), + mikd_k.toarray().astype(dtype)) + assert_array_equal(mik_k.tosparse().toarray(), + mikd_k.tosparse().toarray().astype(dtype)) + + @pytest.mark.parametrize('dtype', tested_types) + @pytest.mark.parametrize('argument_dtype', ALLDTYPES) + def test_dot(self, dtype, argument_dtype): + """ Test the dot-product for type preservation and consistency. + """ + result_dtype = np.promote_types(argument_dtype, dtype) + n = 5 + mik = MikotaPair(n, dtype=dtype) + mik_k = mik.k + mik_m = mik.m + x0 = np.arange(n) + x1 = x0.reshape((-1, 1)) + x2 = np.arange(2 * n).reshape((n, 2)) + lo_set = [mik_k, mik_m] + input_set = [x0, x1, x2] + for lo in lo_set: + for x in input_set: + y = lo.dot(x.astype(argument_dtype)) + assert x.shape == y.shape + assert np.can_cast(y.dtype, result_dtype) + if x.ndim == 2: + ya = lo.toarray() @ x.astype(argument_dtype) + np.array_equal(y, ya) + assert np.can_cast(ya.dtype, result_dtype) + ys = lo.tosparse() @ x.astype(argument_dtype) + np.array_equal(y, ys) + assert np.can_cast(ys.dtype, result_dtype) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/scipy/version.py b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/version.py new file mode 100644 index 0000000000000000000000000000000000000000..d942e3b38948ccc6bc6f27c44ba38d3c67fb07f1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/scipy/version.py @@ -0,0 +1,12 @@ + +""" +Module to expose more detailed version info for the installed `scipy` +""" +version = "1.15.3" +full_version = version +short_version = version.split('.dev')[0] +git_revision = "e29dcb65a2040f04819b426a04b60d44a8f69c04" +release = 'dev' not in version and '+' not in version + +if not release: + version = full_version diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/typeguard/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/typeguard/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6781cad094b29db0959df4fb0c29c0fa9d480c0c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/typeguard/__init__.py @@ -0,0 +1,48 @@ +import os +from typing import Any + +from ._checkers import TypeCheckerCallable as TypeCheckerCallable +from ._checkers import TypeCheckLookupCallback as TypeCheckLookupCallback +from ._checkers import check_type_internal as check_type_internal +from ._checkers import checker_lookup_functions as checker_lookup_functions +from ._checkers import load_plugins as load_plugins +from ._config import CollectionCheckStrategy as CollectionCheckStrategy +from ._config import ForwardRefPolicy as ForwardRefPolicy +from ._config import TypeCheckConfiguration as TypeCheckConfiguration +from ._decorators import typechecked as typechecked +from ._decorators import typeguard_ignore as typeguard_ignore +from ._exceptions import InstrumentationWarning as InstrumentationWarning +from ._exceptions import TypeCheckError as TypeCheckError +from ._exceptions import TypeCheckWarning as TypeCheckWarning +from ._exceptions import TypeHintWarning as TypeHintWarning +from ._functions import TypeCheckFailCallback as TypeCheckFailCallback +from ._functions import check_type as check_type +from ._functions import warn_on_error as warn_on_error +from ._importhook import ImportHookManager as ImportHookManager +from ._importhook import TypeguardFinder as TypeguardFinder +from ._importhook import install_import_hook as install_import_hook +from ._memo import TypeCheckMemo as TypeCheckMemo +from ._suppression import suppress_type_checks as suppress_type_checks +from ._utils import Unset as Unset + +# Re-export imports so they look like they live directly in this package +for value in list(locals().values()): + if getattr(value, "__module__", "").startswith(f"{__name__}."): + value.__module__ = __name__ + + +config: TypeCheckConfiguration + + +def __getattr__(name: str) -> Any: + if name == "config": + from ._config import global_config + + return global_config + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Automatically load checker lookup functions unless explicitly disabled +if "TYPEGUARD_DISABLE_PLUGIN_AUTOLOAD" not in os.environ: + load_plugins() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/typeguard/_utils.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/typeguard/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9bcc8417f8773115de078d307682b457af848052 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/typeguard/_utils.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import inspect +import sys +from importlib import import_module +from inspect import currentframe +from types import CodeType, FrameType, FunctionType +from typing import TYPE_CHECKING, Any, Callable, ForwardRef, Union, cast, final +from weakref import WeakValueDictionary + +if TYPE_CHECKING: + from ._memo import TypeCheckMemo + +if sys.version_info >= (3, 13): + from typing import get_args, get_origin + + def evaluate_forwardref(forwardref: ForwardRef, memo: TypeCheckMemo) -> Any: + return forwardref._evaluate( + memo.globals, memo.locals, type_params=(), recursive_guard=frozenset() + ) + +elif sys.version_info >= (3, 10): + from typing import get_args, get_origin + + def evaluate_forwardref(forwardref: ForwardRef, memo: TypeCheckMemo) -> Any: + return forwardref._evaluate( + memo.globals, memo.locals, recursive_guard=frozenset() + ) + +else: + from typing_extensions import get_args, get_origin + + evaluate_extra_args: tuple[frozenset[Any], ...] = ( + (frozenset(),) if sys.version_info >= (3, 9) else () + ) + + def evaluate_forwardref(forwardref: ForwardRef, memo: TypeCheckMemo) -> Any: + from ._union_transformer import compile_type_hint, type_substitutions + + if not forwardref.__forward_evaluated__: + forwardref.__forward_code__ = compile_type_hint(forwardref.__forward_arg__) + + try: + return forwardref._evaluate(memo.globals, memo.locals, *evaluate_extra_args) + except NameError: + if sys.version_info < (3, 10): + # Try again, with the type substitutions (list -> List etc.) in place + new_globals = memo.globals.copy() + new_globals.setdefault("Union", Union) + if sys.version_info < (3, 9): + new_globals.update(type_substitutions) + + return forwardref._evaluate( + new_globals, memo.locals or new_globals, *evaluate_extra_args + ) + + raise + + +_functions_map: WeakValueDictionary[CodeType, FunctionType] = WeakValueDictionary() + + +def get_type_name(type_: Any) -> str: + name: str + for attrname in "__name__", "_name", "__forward_arg__": + candidate = getattr(type_, attrname, None) + if isinstance(candidate, str): + name = candidate + break + else: + origin = get_origin(type_) + candidate = getattr(origin, "_name", None) + if candidate is None: + candidate = type_.__class__.__name__.strip("_") + + if isinstance(candidate, str): + name = candidate + else: + return "(unknown)" + + args = get_args(type_) + if args: + if name == "Literal": + formatted_args = ", ".join(repr(arg) for arg in args) + else: + formatted_args = ", ".join(get_type_name(arg) for arg in args) + + name += f"[{formatted_args}]" + + module = getattr(type_, "__module__", None) + if module and module not in (None, "typing", "typing_extensions", "builtins"): + name = module + "." + name + + return name + + +def qualified_name(obj: Any, *, add_class_prefix: bool = False) -> str: + """ + Return the qualified name (e.g. package.module.Type) for the given object. + + Builtins and types from the :mod:`typing` package get special treatment by having + the module name stripped from the generated name. + + """ + if obj is None: + return "None" + elif inspect.isclass(obj): + prefix = "class " if add_class_prefix else "" + type_ = obj + else: + prefix = "" + type_ = type(obj) + + module = type_.__module__ + qualname = type_.__qualname__ + name = qualname if module in ("typing", "builtins") else f"{module}.{qualname}" + return prefix + name + + +def function_name(func: Callable[..., Any]) -> str: + """ + Return the qualified name of the given function. + + Builtins and types from the :mod:`typing` package get special treatment by having + the module name stripped from the generated name. + + """ + # For partial functions and objects with __call__ defined, __qualname__ does not + # exist + module = getattr(func, "__module__", "") + qualname = (module + ".") if module not in ("builtins", "") else "" + return qualname + getattr(func, "__qualname__", repr(func)) + + +def resolve_reference(reference: str) -> Any: + modulename, varname = reference.partition(":")[::2] + if not modulename or not varname: + raise ValueError(f"{reference!r} is not a module:varname reference") + + obj = import_module(modulename) + for attr in varname.split("."): + obj = getattr(obj, attr) + + return obj + + +def is_method_of(obj: object, cls: type) -> bool: + return ( + inspect.isfunction(obj) + and obj.__module__ == cls.__module__ + and obj.__qualname__.startswith(cls.__qualname__ + ".") + ) + + +def get_stacklevel() -> int: + level = 1 + frame = cast(FrameType, currentframe()).f_back + while frame and frame.f_globals.get("__name__", "").startswith("typeguard."): + level += 1 + frame = frame.f_back + + return level + + +@final +class Unset: + __slots__ = () + + def __repr__(self) -> str: + return "" + + +unset = Unset() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/INSTALLER b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/LICENSE.txt b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..a31470f14c5978d5fcc3bc173b8399b6c9a6443f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2012 Daniel Holth and contributors + +Permission is hereby granted, free of charge, to any person obtaining a +copy of this software and associated documentation files (the "Software"), +to deal in the Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, sublicense, +and/or sell copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR +OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, +ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR +OTHER DEALINGS IN THE SOFTWARE. diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/METADATA b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..f645dcb673aaa7e3e7e054d57af5e8c0992b70e8 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/METADATA @@ -0,0 +1,66 @@ +Metadata-Version: 2.3 +Name: wheel +Version: 0.45.1 +Summary: A built-package format for Python +Keywords: wheel,packaging +Author-email: Daniel Holth +Maintainer-email: Alex Grönholm +Requires-Python: >=3.8 +Description-Content-Type: text/x-rst +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Topic :: System :: Archiving :: Packaging +Classifier: License :: OSI Approved :: MIT License +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Requires-Dist: pytest >= 6.0.0 ; extra == "test" +Requires-Dist: setuptools >= 65 ; extra == "test" +Project-URL: Changelog, https://wheel.readthedocs.io/en/stable/news.html +Project-URL: Documentation, https://wheel.readthedocs.io/ +Project-URL: Issue Tracker, https://github.com/pypa/wheel/issues +Project-URL: Source, https://github.com/pypa/wheel +Provides-Extra: test + +wheel +===== + +This is a command line tool for manipulating Python wheel files, as defined in +`PEP 427`_. It contains the following functionality: + +* Convert ``.egg`` archives into ``.whl`` +* Unpack wheel archives +* Repack wheel archives +* Add or remove tags in existing wheel archives + +.. _PEP 427: https://www.python.org/dev/peps/pep-0427/ + +Historical note +--------------- + +This project used to contain the implementation of the setuptools_ ``bdist_wheel`` +command, but as of setuptools v70.1, it no longer needs ``wheel`` installed for that to +work. Thus, you should install this **only** if you intend to use the ``wheel`` command +line tool! + +.. _setuptools: https://pypi.org/project/setuptools/ + +Documentation +------------- + +The documentation_ can be found on Read The Docs. + +.. _documentation: https://wheel.readthedocs.io/ + +Code of Conduct +--------------- + +Everyone interacting in the wheel project's codebases, issue trackers, chat +rooms, and mailing lists is expected to follow the `PSF Code of Conduct`_. + +.. _PSF Code of Conduct: https://github.com/pypa/.github/blob/main/CODE_OF_CONDUCT.md + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/RECORD b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..c1535b697f8dd60212541f637c69c9dfb4f1c8f1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/RECORD @@ -0,0 +1,68 @@ +../../bin/wheel,sha256=pBhV19bQIgjS-r541fG3kLU6QtcyKaKdQ2RE9YIzeiU,249 +wheel-0.45.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +wheel-0.45.1.dist-info/LICENSE.txt,sha256=MMI2GGeRCPPo6h0qZYx8pBe9_IkcmO8aifpP8MmChlQ,1107 +wheel-0.45.1.dist-info/METADATA,sha256=mKz84H7m7jsxJyzeIcTVORiTb0NPMV39KvOIYhGgmjA,2313 +wheel-0.45.1.dist-info/RECORD,, +wheel-0.45.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +wheel-0.45.1.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82 +wheel-0.45.1.dist-info/entry_points.txt,sha256=rTY1BbkPHhkGMm4Q3F0pIzJBzW2kMxoG1oriffvGdA0,104 +wheel/__init__.py,sha256=mrxMnvdXACur_LWegbUfh5g5ysWZrd63UJn890wvGNk,59 +wheel/__main__.py,sha256=NkMUnuTCGcOkgY0IBLgBCVC_BGGcWORx2K8jYGS12UE,455 +wheel/__pycache__/__init__.cpython-311.pyc,, +wheel/__pycache__/__main__.cpython-311.pyc,, +wheel/__pycache__/_bdist_wheel.cpython-311.pyc,, +wheel/__pycache__/_setuptools_logging.cpython-311.pyc,, +wheel/__pycache__/bdist_wheel.cpython-311.pyc,, +wheel/__pycache__/macosx_libfile.cpython-311.pyc,, +wheel/__pycache__/metadata.cpython-311.pyc,, +wheel/__pycache__/util.cpython-311.pyc,, +wheel/__pycache__/wheelfile.cpython-311.pyc,, +wheel/_bdist_wheel.py,sha256=UghCQjSH_pVfcZh6oRjzSw_TQhcf3anSx1OkiLSL82M,21694 +wheel/_setuptools_logging.py,sha256=-5KC-lne0ilOUWIDfOkqapUWGMFZhuKYDIavIZiB5kM,781 +wheel/bdist_wheel.py,sha256=tpf9WufiSO1RuEMg5oPhIfSG8DMziCZ_4muCKF69Cqo,1107 +wheel/cli/__init__.py,sha256=Npq6_jKi03dhIcRnmbuFhwviVJxwO0tYEnEhWMv9cJo,4402 +wheel/cli/__pycache__/__init__.cpython-311.pyc,, +wheel/cli/__pycache__/convert.cpython-311.pyc,, +wheel/cli/__pycache__/pack.cpython-311.pyc,, +wheel/cli/__pycache__/tags.cpython-311.pyc,, +wheel/cli/__pycache__/unpack.cpython-311.pyc,, +wheel/cli/convert.py,sha256=Bi0ntEXb9nTllCxWeTRQ4j-nPs3szWSEKipG_GgnMkQ,12634 +wheel/cli/pack.py,sha256=CAFcHdBVulvsHYJlndKVO7KMI9JqBTZz5ii0PKxxCOs,3103 +wheel/cli/tags.py,sha256=lHw-LaWrkS5Jy_qWcw-6pSjeNM6yAjDnqKI3E5JTTCU,4760 +wheel/cli/unpack.py,sha256=Y_J7ynxPSoFFTT7H0fMgbBlVErwyDGcObgme5MBuz58,1021 +wheel/macosx_libfile.py,sha256=k1x7CE3LPtOVGqj6NXQ1nTGYVPaeRrhVzUG_KPq3zDs,16572 +wheel/metadata.py,sha256=JC4p7jlQZu2bUTAQ2fevkqLjg_X6gnNyRhLn6OUO1tc,6171 +wheel/util.py,sha256=aL7aibHwYUgfc8WlolL5tXdkV4DatbJxZHb1kwHFJAU,423 +wheel/vendored/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +wheel/vendored/__pycache__/__init__.cpython-311.pyc,, +wheel/vendored/packaging/LICENSE,sha256=ytHvW9NA1z4HS6YU0m996spceUDD2MNIUuZcSQlobEg,197 +wheel/vendored/packaging/LICENSE.APACHE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174 +wheel/vendored/packaging/LICENSE.BSD,sha256=tw5-m3QvHMb5SLNMFqo5_-zpQZY2S8iP8NIYDwAo-sU,1344 +wheel/vendored/packaging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +wheel/vendored/packaging/__pycache__/__init__.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/_elffile.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/_manylinux.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/_musllinux.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/_parser.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/_structures.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/_tokenizer.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/markers.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/requirements.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/specifiers.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/tags.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/utils.cpython-311.pyc,, +wheel/vendored/packaging/__pycache__/version.cpython-311.pyc,, +wheel/vendored/packaging/_elffile.py,sha256=hbmK8OD6Z7fY6hwinHEUcD1by7czkGiNYu7ShnFEk2k,3266 +wheel/vendored/packaging/_manylinux.py,sha256=P7sdR5_7XBY09LVYYPhHmydMJIIwPXWsh4olk74Uuj4,9588 +wheel/vendored/packaging/_musllinux.py,sha256=z1s8To2hQ0vpn_d-O2i5qxGwEK8WmGlLt3d_26V7NeY,2674 +wheel/vendored/packaging/_parser.py,sha256=4tT4emSl2qTaU7VTQE1Xa9o1jMPCsBezsYBxyNMUN-s,10347 +wheel/vendored/packaging/_structures.py,sha256=q3eVNmbWJGG_S0Dit_S3Ao8qQqz_5PYTXFAKBZe5yr4,1431 +wheel/vendored/packaging/_tokenizer.py,sha256=alCtbwXhOFAmFGZ6BQ-wCTSFoRAJ2z-ysIf7__MTJ_k,5292 +wheel/vendored/packaging/markers.py,sha256=_TSPI1BhJYO7Bp9AzTmHQxIqHEVXaTjmDh9G-w8qzPA,8232 +wheel/vendored/packaging/requirements.py,sha256=dgoBeVprPu2YE6Q8nGfwOPTjATHbRa_ZGLyXhFEln6Q,2933 +wheel/vendored/packaging/specifiers.py,sha256=IWSt0SrLSP72heWhAC8UL0eGvas7XIQHjqiViVfmPKE,39778 +wheel/vendored/packaging/tags.py,sha256=fedHXiOHkBxNZTXotXv8uXPmMFU9ae-TKBujgYHigcA,18950 +wheel/vendored/packaging/utils.py,sha256=XgdmP3yx9-wQEFjO7OvMj9RjEf5JlR5HFFR69v7SQ9E,5268 +wheel/vendored/packaging/version.py,sha256=PFJaYZDxBgyxkfYhH3SQw4qfE9ICCWrTmitvq14y3bs,16234 +wheel/vendored/vendor.txt,sha256=Z2ENjB1i5prfez8CdM1Sdr3c6Zxv2rRRolMpLmBncAE,16 +wheel/wheelfile.py,sha256=USCttNlJwafxt51YYFFKG7jnxz8dfhbyqAZL6jMTA9s,8411 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/REQUESTED b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/WHEEL b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..e3c6feefa22927866e3fd5575379ea972b432aaf --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: flit 3.10.1 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/entry_points.txt b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..06c9f69debd05e53c3b6053a905805e616c61fad --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel-0.45.1.dist-info/entry_points.txt @@ -0,0 +1,6 @@ +[console_scripts] +wheel=wheel.cli:main + +[distutils.commands] +bdist_wheel=wheel.bdist_wheel:bdist_wheel + diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/__main__.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..0be74537494dc2cf18c2e3b318ffd22b886aef6b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/__main__.py @@ -0,0 +1,23 @@ +""" +Wheel command line tool (enable python -m wheel syntax) +""" + +from __future__ import annotations + +import sys + + +def main(): # needed for console script + if __package__ == "": + # To be able to run 'python wheel-0.9.whl/wheel': + import os.path + + path = os.path.dirname(os.path.dirname(__file__)) + sys.path[0:0] = [path] + import wheel.cli + + sys.exit(wheel.cli.main()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ba1217f5bdb6106c37ecc2be74d53ef2237b717 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/__init__.py @@ -0,0 +1,155 @@ +""" +Wheel command-line utility. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from argparse import ArgumentTypeError + + +class WheelError(Exception): + pass + + +def unpack_f(args: argparse.Namespace) -> None: + from .unpack import unpack + + unpack(args.wheelfile, args.dest) + + +def pack_f(args: argparse.Namespace) -> None: + from .pack import pack + + pack(args.directory, args.dest_dir, args.build_number) + + +def convert_f(args: argparse.Namespace) -> None: + from .convert import convert + + convert(args.files, args.dest_dir, args.verbose) + + +def tags_f(args: argparse.Namespace) -> None: + from .tags import tags + + names = ( + tags( + wheel, + args.python_tag, + args.abi_tag, + args.platform_tag, + args.build, + args.remove, + ) + for wheel in args.wheel + ) + + for name in names: + print(name) + + +def version_f(args: argparse.Namespace) -> None: + from .. import __version__ + + print(f"wheel {__version__}") + + +def parse_build_tag(build_tag: str) -> str: + if build_tag and not build_tag[0].isdigit(): + raise ArgumentTypeError("build tag must begin with a digit") + elif "-" in build_tag: + raise ArgumentTypeError("invalid character ('-') in build tag") + + return build_tag + + +TAGS_HELP = """\ +Make a new wheel with given tags. Any tags unspecified will remain the same. +Starting the tags with a "+" will append to the existing tags. Starting with a +"-" will remove a tag (use --option=-TAG syntax). Multiple tags can be +separated by ".". The original file will remain unless --remove is given. The +output filename(s) will be displayed on stdout for further processing. +""" + + +def parser(): + p = argparse.ArgumentParser() + s = p.add_subparsers(help="commands") + + unpack_parser = s.add_parser("unpack", help="Unpack wheel") + unpack_parser.add_argument( + "--dest", "-d", help="Destination directory", default="." + ) + unpack_parser.add_argument("wheelfile", help="Wheel file") + unpack_parser.set_defaults(func=unpack_f) + + repack_parser = s.add_parser("pack", help="Repack wheel") + repack_parser.add_argument("directory", help="Root directory of the unpacked wheel") + repack_parser.add_argument( + "--dest-dir", + "-d", + default=os.path.curdir, + help="Directory to store the wheel (default %(default)s)", + ) + repack_parser.add_argument( + "--build-number", help="Build tag to use in the wheel name" + ) + repack_parser.set_defaults(func=pack_f) + + convert_parser = s.add_parser("convert", help="Convert egg or wininst to wheel") + convert_parser.add_argument("files", nargs="*", help="Files to convert") + convert_parser.add_argument( + "--dest-dir", + "-d", + default=os.path.curdir, + help="Directory to store wheels (default %(default)s)", + ) + convert_parser.add_argument("--verbose", "-v", action="store_true") + convert_parser.set_defaults(func=convert_f) + + tags_parser = s.add_parser( + "tags", help="Add or replace the tags on a wheel", description=TAGS_HELP + ) + tags_parser.add_argument("wheel", nargs="*", help="Existing wheel(s) to retag") + tags_parser.add_argument( + "--remove", + action="store_true", + help="Remove the original files, keeping only the renamed ones", + ) + tags_parser.add_argument( + "--python-tag", metavar="TAG", help="Specify an interpreter tag(s)" + ) + tags_parser.add_argument("--abi-tag", metavar="TAG", help="Specify an ABI tag(s)") + tags_parser.add_argument( + "--platform-tag", metavar="TAG", help="Specify a platform tag(s)" + ) + tags_parser.add_argument( + "--build", type=parse_build_tag, metavar="BUILD", help="Specify a build tag" + ) + tags_parser.set_defaults(func=tags_f) + + version_parser = s.add_parser("version", help="Print version and exit") + version_parser.set_defaults(func=version_f) + + help_parser = s.add_parser("help", help="Show this help") + help_parser.set_defaults(func=lambda args: p.print_help()) + + return p + + +def main(): + p = parser() + args = p.parse_args() + if not hasattr(args, "func"): + p.print_help() + else: + try: + args.func(args) + return 0 + except WheelError as e: + print(e, file=sys.stderr) + + return 1 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/convert.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..61d4775c58596ec12a0c4d4a548de430dee579bf --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/convert.py @@ -0,0 +1,332 @@ +from __future__ import annotations + +import os.path +import re +from abc import ABCMeta, abstractmethod +from collections import defaultdict +from collections.abc import Iterator +from email.message import Message +from email.parser import Parser +from email.policy import EmailPolicy +from glob import iglob +from pathlib import Path +from textwrap import dedent +from zipfile import ZipFile + +from .. import __version__ +from ..metadata import generate_requirements +from ..vendored.packaging.tags import parse_tag +from ..wheelfile import WheelFile + +egg_filename_re = re.compile( + r""" + (?P.+?)-(?P.+?) + (-(?Ppy\d\.\d+) + (-(?P.+?))? + )?.egg$""", + re.VERBOSE, +) +egg_info_re = re.compile( + r""" + ^(?P.+?)-(?P.+?) + (-(?Ppy\d\.\d+) + )?.egg-info/""", + re.VERBOSE, +) +wininst_re = re.compile( + r"\.(?Pwin32|win-amd64)(?:-(?Ppy\d\.\d))?\.exe$" +) +pyd_re = re.compile(r"\.(?P[a-z0-9]+)-(?Pwin32|win_amd64)\.pyd$") +serialization_policy = EmailPolicy( + utf8=True, + mangle_from_=False, + max_line_length=0, +) +GENERATOR = f"wheel {__version__}" + + +def convert_requires(requires: str, metadata: Message) -> None: + extra: str | None = None + requirements: dict[str | None, list[str]] = defaultdict(list) + for line in requires.splitlines(): + line = line.strip() + if not line: + continue + + if line.startswith("[") and line.endswith("]"): + extra = line[1:-1] + continue + + requirements[extra].append(line) + + for key, value in generate_requirements(requirements): + metadata.add_header(key, value) + + +def convert_pkg_info(pkginfo: str, metadata: Message): + parsed_message = Parser().parsestr(pkginfo) + for key, value in parsed_message.items(): + key_lower = key.lower() + if value == "UNKNOWN": + continue + + if key_lower == "description": + description_lines = value.splitlines() + value = "\n".join( + ( + description_lines[0].lstrip(), + dedent("\n".join(description_lines[1:])), + "\n", + ) + ) + metadata.set_payload(value) + elif key_lower == "home-page": + metadata.add_header("Project-URL", f"Homepage, {value}") + elif key_lower == "download-url": + metadata.add_header("Project-URL", f"Download, {value}") + else: + metadata.add_header(key, value) + + metadata.replace_header("Metadata-Version", "2.4") + + +def normalize(name: str) -> str: + return re.sub(r"[-_.]+", "-", name).lower().replace("-", "_") + + +class ConvertSource(metaclass=ABCMeta): + name: str + version: str + pyver: str = "py2.py3" + abi: str = "none" + platform: str = "any" + metadata: Message + + @property + def dist_info_dir(self) -> str: + return f"{self.name}-{self.version}.dist-info" + + @abstractmethod + def generate_contents(self) -> Iterator[tuple[str, bytes]]: + pass + + +class EggFileSource(ConvertSource): + def __init__(self, path: Path): + if not (match := egg_filename_re.match(path.name)): + raise ValueError(f"Invalid egg file name: {path.name}") + + # Binary wheels are assumed to be for CPython + self.path = path + self.name = normalize(match.group("name")) + self.version = match.group("ver") + if pyver := match.group("pyver"): + self.pyver = pyver.replace(".", "") + if arch := match.group("arch"): + self.abi = self.pyver.replace("py", "cp") + self.platform = normalize(arch) + + self.metadata = Message() + + def generate_contents(self) -> Iterator[tuple[str, bytes]]: + with ZipFile(self.path, "r") as zip_file: + for filename in sorted(zip_file.namelist()): + # Skip pure directory entries + if filename.endswith("/"): + continue + + # Handle files in the egg-info directory specially, selectively moving + # them to the dist-info directory while converting as needed + if filename.startswith("EGG-INFO/"): + if filename == "EGG-INFO/requires.txt": + requires = zip_file.read(filename).decode("utf-8") + convert_requires(requires, self.metadata) + elif filename == "EGG-INFO/PKG-INFO": + pkginfo = zip_file.read(filename).decode("utf-8") + convert_pkg_info(pkginfo, self.metadata) + elif filename == "EGG-INFO/entry_points.txt": + yield ( + f"{self.dist_info_dir}/entry_points.txt", + zip_file.read(filename), + ) + + continue + + # For any other file, just pass it through + yield filename, zip_file.read(filename) + + +class EggDirectorySource(EggFileSource): + def generate_contents(self) -> Iterator[tuple[str, bytes]]: + for dirpath, _, filenames in os.walk(self.path): + for filename in sorted(filenames): + path = Path(dirpath, filename) + if path.parent.name == "EGG-INFO": + if path.name == "requires.txt": + requires = path.read_text("utf-8") + convert_requires(requires, self.metadata) + elif path.name == "PKG-INFO": + pkginfo = path.read_text("utf-8") + convert_pkg_info(pkginfo, self.metadata) + if name := self.metadata.get("Name"): + self.name = normalize(name) + + if version := self.metadata.get("Version"): + self.version = version + elif path.name == "entry_points.txt": + yield ( + f"{self.dist_info_dir}/entry_points.txt", + path.read_bytes(), + ) + + continue + + # For any other file, just pass it through + yield str(path.relative_to(self.path)), path.read_bytes() + + +class WininstFileSource(ConvertSource): + """ + Handles distributions created with ``bdist_wininst``. + + The egginfo filename has the format:: + + name-ver(-pyver)(-arch).egg-info + + The installer filename has the format:: + + name-ver.arch(-pyver).exe + + Some things to note: + + 1. The installer filename is not definitive. An installer can be renamed + and work perfectly well as an installer. So more reliable data should + be used whenever possible. + 2. The egg-info data should be preferred for the name and version, because + these come straight from the distutils metadata, and are mandatory. + 3. The pyver from the egg-info data should be ignored, as it is + constructed from the version of Python used to build the installer, + which is irrelevant - the installer filename is correct here (even to + the point that when it's not there, any version is implied). + 4. The architecture must be taken from the installer filename, as it is + not included in the egg-info data. + 5. Architecture-neutral installers still have an architecture because the + installer format itself (being executable) is architecture-specific. We + should therefore ignore the architecture if the content is pure-python. + """ + + def __init__(self, path: Path): + self.path = path + self.metadata = Message() + + # Determine the initial architecture and Python version from the file name + # (if possible) + if match := wininst_re.search(path.name): + self.platform = normalize(match.group("platform")) + if pyver := match.group("pyver"): + self.pyver = pyver.replace(".", "") + + # Look for an .egg-info directory and any .pyd files for more precise info + egg_info_found = pyd_found = False + with ZipFile(self.path) as zip_file: + for filename in zip_file.namelist(): + prefix, filename = filename.split("/", 1) + if not egg_info_found and (match := egg_info_re.match(filename)): + egg_info_found = True + self.name = normalize(match.group("name")) + self.version = match.group("ver") + if pyver := match.group("pyver"): + self.pyver = pyver.replace(".", "") + elif not pyd_found and (match := pyd_re.search(filename)): + pyd_found = True + self.abi = match.group("abi") + self.platform = match.group("platform") + + if egg_info_found and pyd_found: + break + + def generate_contents(self) -> Iterator[tuple[str, bytes]]: + dist_info_dir = f"{self.name}-{self.version}.dist-info" + data_dir = f"{self.name}-{self.version}.data" + with ZipFile(self.path, "r") as zip_file: + for filename in sorted(zip_file.namelist()): + # Skip pure directory entries + if filename.endswith("/"): + continue + + # Handle files in the egg-info directory specially, selectively moving + # them to the dist-info directory while converting as needed + prefix, target_filename = filename.split("/", 1) + if egg_info_re.search(target_filename): + basename = target_filename.rsplit("/", 1)[-1] + if basename == "requires.txt": + requires = zip_file.read(filename).decode("utf-8") + convert_requires(requires, self.metadata) + elif basename == "PKG-INFO": + pkginfo = zip_file.read(filename).decode("utf-8") + convert_pkg_info(pkginfo, self.metadata) + elif basename == "entry_points.txt": + yield ( + f"{dist_info_dir}/entry_points.txt", + zip_file.read(filename), + ) + + continue + elif prefix == "SCRIPTS": + target_filename = f"{data_dir}/scripts/{target_filename}" + + # For any other file, just pass it through + yield target_filename, zip_file.read(filename) + + +def convert(files: list[str], dest_dir: str, verbose: bool) -> None: + for pat in files: + for archive in iglob(pat): + path = Path(archive) + if path.suffix == ".egg": + if path.is_dir(): + source: ConvertSource = EggDirectorySource(path) + else: + source = EggFileSource(path) + else: + source = WininstFileSource(path) + + if verbose: + print(f"{archive}...", flush=True, end="") + + dest_path = Path(dest_dir) / ( + f"{source.name}-{source.version}-{source.pyver}-{source.abi}" + f"-{source.platform}.whl" + ) + with WheelFile(dest_path, "w") as wheelfile: + for name_or_zinfo, contents in source.generate_contents(): + wheelfile.writestr(name_or_zinfo, contents) + + # Write the METADATA file + wheelfile.writestr( + f"{source.dist_info_dir}/METADATA", + source.metadata.as_string(policy=serialization_policy).encode( + "utf-8" + ), + ) + + # Write the WHEEL file + wheel_message = Message() + wheel_message.add_header("Wheel-Version", "1.0") + wheel_message.add_header("Generator", GENERATOR) + wheel_message.add_header( + "Root-Is-Purelib", str(source.platform == "any").lower() + ) + tags = parse_tag(f"{source.pyver}-{source.abi}-{source.platform}") + for tag in sorted(tags, key=lambda tag: tag.interpreter): + wheel_message.add_header("Tag", str(tag)) + + wheelfile.writestr( + f"{source.dist_info_dir}/WHEEL", + wheel_message.as_string(policy=serialization_policy).encode( + "utf-8" + ), + ) + + if verbose: + print("OK") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/pack.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/pack.py new file mode 100644 index 0000000000000000000000000000000000000000..64469c0c730c24449e77c23f46d0e3f68d647d67 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/pack.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import email.policy +import os.path +import re +from email.generator import BytesGenerator +from email.parser import BytesParser + +from wheel.cli import WheelError +from wheel.wheelfile import WheelFile + +DIST_INFO_RE = re.compile(r"^(?P(?P.+?)-(?P\d.*?))\.dist-info$") + + +def pack(directory: str, dest_dir: str, build_number: str | None) -> None: + """Repack a previously unpacked wheel directory into a new wheel file. + + The .dist-info/WHEEL file must contain one or more tags so that the target + wheel file name can be determined. + + :param directory: The unpacked wheel directory + :param dest_dir: Destination directory (defaults to the current directory) + """ + # Find the .dist-info directory + dist_info_dirs = [ + fn + for fn in os.listdir(directory) + if os.path.isdir(os.path.join(directory, fn)) and DIST_INFO_RE.match(fn) + ] + if len(dist_info_dirs) > 1: + raise WheelError(f"Multiple .dist-info directories found in {directory}") + elif not dist_info_dirs: + raise WheelError(f"No .dist-info directories found in {directory}") + + # Determine the target wheel filename + dist_info_dir = dist_info_dirs[0] + name_version = DIST_INFO_RE.match(dist_info_dir).group("namever") + + # Read the tags and the existing build number from .dist-info/WHEEL + wheel_file_path = os.path.join(directory, dist_info_dir, "WHEEL") + with open(wheel_file_path, "rb") as f: + info = BytesParser(policy=email.policy.compat32).parse(f) + tags: list[str] = info.get_all("Tag", []) + existing_build_number = info.get("Build") + + if not tags: + raise WheelError( + f"No tags present in {dist_info_dir}/WHEEL; cannot determine target " + f"wheel filename" + ) + + # Set the wheel file name and add/replace/remove the Build tag in .dist-info/WHEEL + build_number = build_number if build_number is not None else existing_build_number + if build_number is not None: + del info["Build"] + if build_number: + info["Build"] = build_number + name_version += "-" + build_number + + if build_number != existing_build_number: + with open(wheel_file_path, "wb") as f: + BytesGenerator(f, maxheaderlen=0).flatten(info) + + # Reassemble the tags for the wheel file + tagline = compute_tagline(tags) + + # Repack the wheel + wheel_path = os.path.join(dest_dir, f"{name_version}-{tagline}.whl") + with WheelFile(wheel_path, "w") as wf: + print(f"Repacking wheel as {wheel_path}...", end="", flush=True) + wf.write_files(directory) + + print("OK") + + +def compute_tagline(tags: list[str]) -> str: + """Compute a tagline from a list of tags. + + :param tags: A list of tags + :return: A tagline + """ + impls = sorted({tag.split("-")[0] for tag in tags}) + abivers = sorted({tag.split("-")[1] for tag in tags}) + platforms = sorted({tag.split("-")[2] for tag in tags}) + return "-".join([".".join(impls), ".".join(abivers), ".".join(platforms)]) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/tags.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/tags.py new file mode 100644 index 0000000000000000000000000000000000000000..88da72e9ec4a31e2427bdb2bcf2b3e0ce3c0beb0 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/tags.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import email.policy +import itertools +import os +from collections.abc import Iterable +from email.parser import BytesParser + +from ..wheelfile import WheelFile + + +def _compute_tags(original_tags: Iterable[str], new_tags: str | None) -> set[str]: + """Add or replace tags. Supports dot-separated tags""" + if new_tags is None: + return set(original_tags) + + if new_tags.startswith("+"): + return {*original_tags, *new_tags[1:].split(".")} + + if new_tags.startswith("-"): + return set(original_tags) - set(new_tags[1:].split(".")) + + return set(new_tags.split(".")) + + +def tags( + wheel: str, + python_tags: str | None = None, + abi_tags: str | None = None, + platform_tags: str | None = None, + build_tag: str | None = None, + remove: bool = False, +) -> str: + """Change the tags on a wheel file. + + The tags are left unchanged if they are not specified. To specify "none", + use ["none"]. To append to the previous tags, a tag should start with a + "+". If a tag starts with "-", it will be removed from existing tags. + Processing is done left to right. + + :param wheel: The paths to the wheels + :param python_tags: The Python tags to set + :param abi_tags: The ABI tags to set + :param platform_tags: The platform tags to set + :param build_tag: The build tag to set + :param remove: Remove the original wheel + """ + with WheelFile(wheel, "r") as f: + assert f.filename, f"{f.filename} must be available" + + wheel_info = f.read(f.dist_info_path + "/WHEEL") + info = BytesParser(policy=email.policy.compat32).parsebytes(wheel_info) + + original_wheel_name = os.path.basename(f.filename) + namever = f.parsed_filename.group("namever") + build = f.parsed_filename.group("build") + original_python_tags = f.parsed_filename.group("pyver").split(".") + original_abi_tags = f.parsed_filename.group("abi").split(".") + original_plat_tags = f.parsed_filename.group("plat").split(".") + + tags: list[str] = info.get_all("Tag", []) + existing_build_tag = info.get("Build") + + impls = {tag.split("-")[0] for tag in tags} + abivers = {tag.split("-")[1] for tag in tags} + platforms = {tag.split("-")[2] for tag in tags} + + if impls != set(original_python_tags): + msg = f"Wheel internal tags {impls!r} != filename tags {original_python_tags!r}" + raise AssertionError(msg) + + if abivers != set(original_abi_tags): + msg = f"Wheel internal tags {abivers!r} != filename tags {original_abi_tags!r}" + raise AssertionError(msg) + + if platforms != set(original_plat_tags): + msg = ( + f"Wheel internal tags {platforms!r} != filename tags {original_plat_tags!r}" + ) + raise AssertionError(msg) + + if existing_build_tag != build: + msg = ( + f"Incorrect filename '{build}' " + f"& *.dist-info/WHEEL '{existing_build_tag}' build numbers" + ) + raise AssertionError(msg) + + # Start changing as needed + if build_tag is not None: + build = build_tag + + final_python_tags = sorted(_compute_tags(original_python_tags, python_tags)) + final_abi_tags = sorted(_compute_tags(original_abi_tags, abi_tags)) + final_plat_tags = sorted(_compute_tags(original_plat_tags, platform_tags)) + + final_tags = [ + namever, + ".".join(final_python_tags), + ".".join(final_abi_tags), + ".".join(final_plat_tags), + ] + if build: + final_tags.insert(1, build) + + final_wheel_name = "-".join(final_tags) + ".whl" + + if original_wheel_name != final_wheel_name: + del info["Tag"], info["Build"] + for a, b, c in itertools.product( + final_python_tags, final_abi_tags, final_plat_tags + ): + info["Tag"] = f"{a}-{b}-{c}" + if build: + info["Build"] = build + + original_wheel_path = os.path.join( + os.path.dirname(f.filename), original_wheel_name + ) + final_wheel_path = os.path.join(os.path.dirname(f.filename), final_wheel_name) + + with WheelFile(original_wheel_path, "r") as fin, WheelFile( + final_wheel_path, "w" + ) as fout: + fout.comment = fin.comment # preserve the comment + for item in fin.infolist(): + if item.is_dir(): + continue + if item.filename == f.dist_info_path + "/RECORD": + continue + if item.filename == f.dist_info_path + "/WHEEL": + fout.writestr(item, info.as_bytes()) + else: + fout.writestr(item, fin.read(item)) + + if remove: + os.remove(original_wheel_path) + + return final_wheel_name diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/unpack.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/unpack.py new file mode 100644 index 0000000000000000000000000000000000000000..d48840e6ec0512225233bf02d1d7ce203415b04c --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/cli/unpack.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from pathlib import Path + +from ..wheelfile import WheelFile + + +def unpack(path: str, dest: str = ".") -> None: + """Unpack a wheel. + + Wheel content will be unpacked to {dest}/{name}-{ver}, where {name} + is the package name and {ver} its version. + + :param path: The path to the wheel. + :param dest: Destination directory (default to current directory). + """ + with WheelFile(path) as wf: + namever = wf.parsed_filename.group("namever") + destination = Path(dest) / namever + print(f"Unpacking to: {destination}...", end="", flush=True) + for zinfo in wf.filelist: + wf.extract(zinfo, destination) + + # Set permissions to the same values as they were set in the archive + # We have to do this manually due to + # https://github.com/python/cpython/issues/59999 + permissions = zinfo.external_attr >> 16 & 0o777 + destination.joinpath(zinfo.filename).chmod(permissions) + + print("OK") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..6f62d44e4ef733c0e713afcd2371fed7f2b3de67 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE @@ -0,0 +1,3 @@ +This software is made available under the terms of *either* of the licenses +found in LICENSE.APACHE or LICENSE.BSD. Contributions to this software is made +under the terms of *both* these licenses. diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE.APACHE b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE.APACHE new file mode 100644 index 0000000000000000000000000000000000000000..f433b1a53f5b830a205fd2df78e2b34974656c7b --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE.APACHE @@ -0,0 +1,177 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE.BSD b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE.BSD new file mode 100644 index 0000000000000000000000000000000000000000..42ce7b75c92fb01a3f6ed17eea363f756b7da582 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/LICENSE.BSD @@ -0,0 +1,23 @@ +Copyright (c) Donald Stufft and individual contributors. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + 1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/__init__.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_elffile.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_elffile.py new file mode 100644 index 0000000000000000000000000000000000000000..6fb19b30bb53c18f38a9ef02dd7c4478670fb962 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_elffile.py @@ -0,0 +1,108 @@ +""" +ELF file parser. + +This provides a class ``ELFFile`` that parses an ELF executable in a similar +interface to ``ZipFile``. Only the read interface is implemented. + +Based on: https://gist.github.com/lyssdod/f51579ae8d93c8657a5564aefc2ffbca +ELF header: https://refspecs.linuxfoundation.org/elf/gabi4+/ch4.eheader.html +""" + +import enum +import os +import struct +from typing import IO, Optional, Tuple + + +class ELFInvalid(ValueError): + pass + + +class EIClass(enum.IntEnum): + C32 = 1 + C64 = 2 + + +class EIData(enum.IntEnum): + Lsb = 1 + Msb = 2 + + +class EMachine(enum.IntEnum): + I386 = 3 + S390 = 22 + Arm = 40 + X8664 = 62 + AArc64 = 183 + + +class ELFFile: + """ + Representation of an ELF executable. + """ + + def __init__(self, f: IO[bytes]) -> None: + self._f = f + + try: + ident = self._read("16B") + except struct.error: + raise ELFInvalid("unable to parse identification") + magic = bytes(ident[:4]) + if magic != b"\x7fELF": + raise ELFInvalid(f"invalid magic: {magic!r}") + + self.capacity = ident[4] # Format for program header (bitness). + self.encoding = ident[5] # Data structure encoding (endianness). + + try: + # e_fmt: Format for program header. + # p_fmt: Format for section header. + # p_idx: Indexes to find p_type, p_offset, and p_filesz. + e_fmt, self._p_fmt, self._p_idx = { + (1, 1): ("HHIIIIIHHH", ">IIIIIIII", (0, 1, 4)), # 32-bit MSB. + (2, 1): ("HHIQQQIHHH", ">IIQQQQQQ", (0, 2, 5)), # 64-bit MSB. + }[(self.capacity, self.encoding)] + except KeyError: + raise ELFInvalid( + f"unrecognized capacity ({self.capacity}) or " + f"encoding ({self.encoding})" + ) + + try: + ( + _, + self.machine, # Architecture type. + _, + _, + self._e_phoff, # Offset of program header. + _, + self.flags, # Processor-specific flags. + _, + self._e_phentsize, # Size of section. + self._e_phnum, # Number of sections. + ) = self._read(e_fmt) + except struct.error as e: + raise ELFInvalid("unable to parse machine and section information") from e + + def _read(self, fmt: str) -> Tuple[int, ...]: + return struct.unpack(fmt, self._f.read(struct.calcsize(fmt))) + + @property + def interpreter(self) -> Optional[str]: + """ + The path recorded in the ``PT_INTERP`` section header. + """ + for index in range(self._e_phnum): + self._f.seek(self._e_phoff + self._e_phentsize * index) + try: + data = self._read(self._p_fmt) + except struct.error: + continue + if data[self._p_idx[0]] != 3: # Not PT_INTERP. + continue + self._f.seek(data[self._p_idx[1]]) + return os.fsdecode(self._f.read(data[self._p_idx[2]])).strip("\0") + return None diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_manylinux.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_manylinux.py new file mode 100644 index 0000000000000000000000000000000000000000..1f5f4ab3e514d1846d3bd189bf081fdb1528ba08 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_manylinux.py @@ -0,0 +1,260 @@ +import collections +import contextlib +import functools +import os +import re +import sys +import warnings +from typing import Dict, Generator, Iterator, NamedTuple, Optional, Sequence, Tuple + +from ._elffile import EIClass, EIData, ELFFile, EMachine + +EF_ARM_ABIMASK = 0xFF000000 +EF_ARM_ABI_VER5 = 0x05000000 +EF_ARM_ABI_FLOAT_HARD = 0x00000400 + + +# `os.PathLike` not a generic type until Python 3.9, so sticking with `str` +# as the type for `path` until then. +@contextlib.contextmanager +def _parse_elf(path: str) -> Generator[Optional[ELFFile], None, None]: + try: + with open(path, "rb") as f: + yield ELFFile(f) + except (OSError, TypeError, ValueError): + yield None + + +def _is_linux_armhf(executable: str) -> bool: + # hard-float ABI can be detected from the ELF header of the running + # process + # https://static.docs.arm.com/ihi0044/g/aaelf32.pdf + with _parse_elf(executable) as f: + return ( + f is not None + and f.capacity == EIClass.C32 + and f.encoding == EIData.Lsb + and f.machine == EMachine.Arm + and f.flags & EF_ARM_ABIMASK == EF_ARM_ABI_VER5 + and f.flags & EF_ARM_ABI_FLOAT_HARD == EF_ARM_ABI_FLOAT_HARD + ) + + +def _is_linux_i686(executable: str) -> bool: + with _parse_elf(executable) as f: + return ( + f is not None + and f.capacity == EIClass.C32 + and f.encoding == EIData.Lsb + and f.machine == EMachine.I386 + ) + + +def _have_compatible_abi(executable: str, archs: Sequence[str]) -> bool: + if "armv7l" in archs: + return _is_linux_armhf(executable) + if "i686" in archs: + return _is_linux_i686(executable) + allowed_archs = { + "x86_64", + "aarch64", + "ppc64", + "ppc64le", + "s390x", + "loongarch64", + "riscv64", + } + return any(arch in allowed_archs for arch in archs) + + +# If glibc ever changes its major version, we need to know what the last +# minor version was, so we can build the complete list of all versions. +# For now, guess what the highest minor version might be, assume it will +# be 50 for testing. Once this actually happens, update the dictionary +# with the actual value. +_LAST_GLIBC_MINOR: Dict[int, int] = collections.defaultdict(lambda: 50) + + +class _GLibCVersion(NamedTuple): + major: int + minor: int + + +def _glibc_version_string_confstr() -> Optional[str]: + """ + Primary implementation of glibc_version_string using os.confstr. + """ + # os.confstr is quite a bit faster than ctypes.DLL. It's also less likely + # to be broken or missing. This strategy is used in the standard library + # platform module. + # https://github.com/python/cpython/blob/fcf1d003bf4f0100c/Lib/platform.py#L175-L183 + try: + # Should be a string like "glibc 2.17". + version_string: Optional[str] = os.confstr("CS_GNU_LIBC_VERSION") + assert version_string is not None + _, version = version_string.rsplit() + except (AssertionError, AttributeError, OSError, ValueError): + # os.confstr() or CS_GNU_LIBC_VERSION not available (or a bad value)... + return None + return version + + +def _glibc_version_string_ctypes() -> Optional[str]: + """ + Fallback implementation of glibc_version_string using ctypes. + """ + try: + import ctypes + except ImportError: + return None + + # ctypes.CDLL(None) internally calls dlopen(NULL), and as the dlopen + # manpage says, "If filename is NULL, then the returned handle is for the + # main program". This way we can let the linker do the work to figure out + # which libc our process is actually using. + # + # We must also handle the special case where the executable is not a + # dynamically linked executable. This can occur when using musl libc, + # for example. In this situation, dlopen() will error, leading to an + # OSError. Interestingly, at least in the case of musl, there is no + # errno set on the OSError. The single string argument used to construct + # OSError comes from libc itself and is therefore not portable to + # hard code here. In any case, failure to call dlopen() means we + # can proceed, so we bail on our attempt. + try: + process_namespace = ctypes.CDLL(None) + except OSError: + return None + + try: + gnu_get_libc_version = process_namespace.gnu_get_libc_version + except AttributeError: + # Symbol doesn't exist -> therefore, we are not linked to + # glibc. + return None + + # Call gnu_get_libc_version, which returns a string like "2.5" + gnu_get_libc_version.restype = ctypes.c_char_p + version_str: str = gnu_get_libc_version() + # py2 / py3 compatibility: + if not isinstance(version_str, str): + version_str = version_str.decode("ascii") + + return version_str + + +def _glibc_version_string() -> Optional[str]: + """Returns glibc version string, or None if not using glibc.""" + return _glibc_version_string_confstr() or _glibc_version_string_ctypes() + + +def _parse_glibc_version(version_str: str) -> Tuple[int, int]: + """Parse glibc version. + + We use a regexp instead of str.split because we want to discard any + random junk that might come after the minor version -- this might happen + in patched/forked versions of glibc (e.g. Linaro's version of glibc + uses version strings like "2.20-2014.11"). See gh-3588. + """ + m = re.match(r"(?P[0-9]+)\.(?P[0-9]+)", version_str) + if not m: + warnings.warn( + f"Expected glibc version with 2 components major.minor," + f" got: {version_str}", + RuntimeWarning, + ) + return -1, -1 + return int(m.group("major")), int(m.group("minor")) + + +@functools.lru_cache +def _get_glibc_version() -> Tuple[int, int]: + version_str = _glibc_version_string() + if version_str is None: + return (-1, -1) + return _parse_glibc_version(version_str) + + +# From PEP 513, PEP 600 +def _is_compatible(arch: str, version: _GLibCVersion) -> bool: + sys_glibc = _get_glibc_version() + if sys_glibc < version: + return False + # Check for presence of _manylinux module. + try: + import _manylinux + except ImportError: + return True + if hasattr(_manylinux, "manylinux_compatible"): + result = _manylinux.manylinux_compatible(version[0], version[1], arch) + if result is not None: + return bool(result) + return True + if version == _GLibCVersion(2, 5): + if hasattr(_manylinux, "manylinux1_compatible"): + return bool(_manylinux.manylinux1_compatible) + if version == _GLibCVersion(2, 12): + if hasattr(_manylinux, "manylinux2010_compatible"): + return bool(_manylinux.manylinux2010_compatible) + if version == _GLibCVersion(2, 17): + if hasattr(_manylinux, "manylinux2014_compatible"): + return bool(_manylinux.manylinux2014_compatible) + return True + + +_LEGACY_MANYLINUX_MAP = { + # CentOS 7 w/ glibc 2.17 (PEP 599) + (2, 17): "manylinux2014", + # CentOS 6 w/ glibc 2.12 (PEP 571) + (2, 12): "manylinux2010", + # CentOS 5 w/ glibc 2.5 (PEP 513) + (2, 5): "manylinux1", +} + + +def platform_tags(archs: Sequence[str]) -> Iterator[str]: + """Generate manylinux tags compatible to the current platform. + + :param archs: Sequence of compatible architectures. + The first one shall be the closest to the actual architecture and be the part of + platform tag after the ``linux_`` prefix, e.g. ``x86_64``. + The ``linux_`` prefix is assumed as a prerequisite for the current platform to + be manylinux-compatible. + + :returns: An iterator of compatible manylinux tags. + """ + if not _have_compatible_abi(sys.executable, archs): + return + # Oldest glibc to be supported regardless of architecture is (2, 17). + too_old_glibc2 = _GLibCVersion(2, 16) + if set(archs) & {"x86_64", "i686"}: + # On x86/i686 also oldest glibc to be supported is (2, 5). + too_old_glibc2 = _GLibCVersion(2, 4) + current_glibc = _GLibCVersion(*_get_glibc_version()) + glibc_max_list = [current_glibc] + # We can assume compatibility across glibc major versions. + # https://sourceware.org/bugzilla/show_bug.cgi?id=24636 + # + # Build a list of maximum glibc versions so that we can + # output the canonical list of all glibc from current_glibc + # down to too_old_glibc2, including all intermediary versions. + for glibc_major in range(current_glibc.major - 1, 1, -1): + glibc_minor = _LAST_GLIBC_MINOR[glibc_major] + glibc_max_list.append(_GLibCVersion(glibc_major, glibc_minor)) + for arch in archs: + for glibc_max in glibc_max_list: + if glibc_max.major == too_old_glibc2.major: + min_minor = too_old_glibc2.minor + else: + # For other glibc major versions oldest supported is (x, 0). + min_minor = -1 + for glibc_minor in range(glibc_max.minor, min_minor, -1): + glibc_version = _GLibCVersion(glibc_max.major, glibc_minor) + tag = "manylinux_{}_{}".format(*glibc_version) + if _is_compatible(arch, glibc_version): + yield f"{tag}_{arch}" + # Handle the legacy manylinux1, manylinux2010, manylinux2014 tags. + if glibc_version in _LEGACY_MANYLINUX_MAP: + legacy_tag = _LEGACY_MANYLINUX_MAP[glibc_version] + if _is_compatible(arch, glibc_version): + yield f"{legacy_tag}_{arch}" diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_musllinux.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_musllinux.py new file mode 100644 index 0000000000000000000000000000000000000000..eb4251b5c1e82772b2b0ea539943da5141fd55ec --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_musllinux.py @@ -0,0 +1,83 @@ +"""PEP 656 support. + +This module implements logic to detect if the currently running Python is +linked against musl, and what musl version is used. +""" + +import functools +import re +import subprocess +import sys +from typing import Iterator, NamedTuple, Optional, Sequence + +from ._elffile import ELFFile + + +class _MuslVersion(NamedTuple): + major: int + minor: int + + +def _parse_musl_version(output: str) -> Optional[_MuslVersion]: + lines = [n for n in (n.strip() for n in output.splitlines()) if n] + if len(lines) < 2 or lines[0][:4] != "musl": + return None + m = re.match(r"Version (\d+)\.(\d+)", lines[1]) + if not m: + return None + return _MuslVersion(major=int(m.group(1)), minor=int(m.group(2))) + + +@functools.lru_cache +def _get_musl_version(executable: str) -> Optional[_MuslVersion]: + """Detect currently-running musl runtime version. + + This is done by checking the specified executable's dynamic linking + information, and invoking the loader to parse its output for a version + string. If the loader is musl, the output would be something like:: + + musl libc (x86_64) + Version 1.2.2 + Dynamic Program Loader + """ + try: + with open(executable, "rb") as f: + ld = ELFFile(f).interpreter + except (OSError, TypeError, ValueError): + return None + if ld is None or "musl" not in ld: + return None + proc = subprocess.run([ld], stderr=subprocess.PIPE, text=True) + return _parse_musl_version(proc.stderr) + + +def platform_tags(archs: Sequence[str]) -> Iterator[str]: + """Generate musllinux tags compatible to the current platform. + + :param archs: Sequence of compatible architectures. + The first one shall be the closest to the actual architecture and be the part of + platform tag after the ``linux_`` prefix, e.g. ``x86_64``. + The ``linux_`` prefix is assumed as a prerequisite for the current platform to + be musllinux-compatible. + + :returns: An iterator of compatible musllinux tags. + """ + sys_musl = _get_musl_version(sys.executable) + if sys_musl is None: # Python not dynamically linked against musl. + return + for arch in archs: + for minor in range(sys_musl.minor, -1, -1): + yield f"musllinux_{sys_musl.major}_{minor}_{arch}" + + +if __name__ == "__main__": # pragma: no cover + import sysconfig + + plat = sysconfig.get_platform() + assert plat.startswith("linux-"), "not linux" + + print("plat:", plat) + print("musl:", _get_musl_version(sys.executable)) + print("tags:", end=" ") + for t in platform_tags(re.sub(r"[.-]", "_", plat.split("-", 1)[-1])): + print(t, end="\n ") diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_parser.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..513686a2190f9911c08ca1ca263f37b799f44702 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_parser.py @@ -0,0 +1,356 @@ +"""Handwritten parser of dependency specifiers. + +The docstring for each __parse_* function contains EBNF-inspired grammar representing +the implementation. +""" + +import ast +from typing import Any, List, NamedTuple, Optional, Tuple, Union + +from ._tokenizer import DEFAULT_RULES, Tokenizer + + +class Node: + def __init__(self, value: str) -> None: + self.value = value + + def __str__(self) -> str: + return self.value + + def __repr__(self) -> str: + return f"<{self.__class__.__name__}('{self}')>" + + def serialize(self) -> str: + raise NotImplementedError + + +class Variable(Node): + def serialize(self) -> str: + return str(self) + + +class Value(Node): + def serialize(self) -> str: + return f'"{self}"' + + +class Op(Node): + def serialize(self) -> str: + return str(self) + + +MarkerVar = Union[Variable, Value] +MarkerItem = Tuple[MarkerVar, Op, MarkerVar] +# MarkerAtom = Union[MarkerItem, List["MarkerAtom"]] +# MarkerList = List[Union["MarkerList", MarkerAtom, str]] +# mypy does not support recursive type definition +# https://github.com/python/mypy/issues/731 +MarkerAtom = Any +MarkerList = List[Any] + + +class ParsedRequirement(NamedTuple): + name: str + url: str + extras: List[str] + specifier: str + marker: Optional[MarkerList] + + +# -------------------------------------------------------------------------------------- +# Recursive descent parser for dependency specifier +# -------------------------------------------------------------------------------------- +def parse_requirement(source: str) -> ParsedRequirement: + return _parse_requirement(Tokenizer(source, rules=DEFAULT_RULES)) + + +def _parse_requirement(tokenizer: Tokenizer) -> ParsedRequirement: + """ + requirement = WS? IDENTIFIER WS? extras WS? requirement_details + """ + tokenizer.consume("WS") + + name_token = tokenizer.expect( + "IDENTIFIER", expected="package name at the start of dependency specifier" + ) + name = name_token.text + tokenizer.consume("WS") + + extras = _parse_extras(tokenizer) + tokenizer.consume("WS") + + url, specifier, marker = _parse_requirement_details(tokenizer) + tokenizer.expect("END", expected="end of dependency specifier") + + return ParsedRequirement(name, url, extras, specifier, marker) + + +def _parse_requirement_details( + tokenizer: Tokenizer, +) -> Tuple[str, str, Optional[MarkerList]]: + """ + requirement_details = AT URL (WS requirement_marker?)? + | specifier WS? (requirement_marker)? + """ + + specifier = "" + url = "" + marker = None + + if tokenizer.check("AT"): + tokenizer.read() + tokenizer.consume("WS") + + url_start = tokenizer.position + url = tokenizer.expect("URL", expected="URL after @").text + if tokenizer.check("END", peek=True): + return (url, specifier, marker) + + tokenizer.expect("WS", expected="whitespace after URL") + + # The input might end after whitespace. + if tokenizer.check("END", peek=True): + return (url, specifier, marker) + + marker = _parse_requirement_marker( + tokenizer, span_start=url_start, after="URL and whitespace" + ) + else: + specifier_start = tokenizer.position + specifier = _parse_specifier(tokenizer) + tokenizer.consume("WS") + + if tokenizer.check("END", peek=True): + return (url, specifier, marker) + + marker = _parse_requirement_marker( + tokenizer, + span_start=specifier_start, + after=( + "version specifier" + if specifier + else "name and no valid version specifier" + ), + ) + + return (url, specifier, marker) + + +def _parse_requirement_marker( + tokenizer: Tokenizer, *, span_start: int, after: str +) -> MarkerList: + """ + requirement_marker = SEMICOLON marker WS? + """ + + if not tokenizer.check("SEMICOLON"): + tokenizer.raise_syntax_error( + f"Expected end or semicolon (after {after})", + span_start=span_start, + ) + tokenizer.read() + + marker = _parse_marker(tokenizer) + tokenizer.consume("WS") + + return marker + + +def _parse_extras(tokenizer: Tokenizer) -> List[str]: + """ + extras = (LEFT_BRACKET wsp* extras_list? wsp* RIGHT_BRACKET)? + """ + if not tokenizer.check("LEFT_BRACKET", peek=True): + return [] + + with tokenizer.enclosing_tokens( + "LEFT_BRACKET", + "RIGHT_BRACKET", + around="extras", + ): + tokenizer.consume("WS") + extras = _parse_extras_list(tokenizer) + tokenizer.consume("WS") + + return extras + + +def _parse_extras_list(tokenizer: Tokenizer) -> List[str]: + """ + extras_list = identifier (wsp* ',' wsp* identifier)* + """ + extras: List[str] = [] + + if not tokenizer.check("IDENTIFIER"): + return extras + + extras.append(tokenizer.read().text) + + while True: + tokenizer.consume("WS") + if tokenizer.check("IDENTIFIER", peek=True): + tokenizer.raise_syntax_error("Expected comma between extra names") + elif not tokenizer.check("COMMA"): + break + + tokenizer.read() + tokenizer.consume("WS") + + extra_token = tokenizer.expect("IDENTIFIER", expected="extra name after comma") + extras.append(extra_token.text) + + return extras + + +def _parse_specifier(tokenizer: Tokenizer) -> str: + """ + specifier = LEFT_PARENTHESIS WS? version_many WS? RIGHT_PARENTHESIS + | WS? version_many WS? + """ + with tokenizer.enclosing_tokens( + "LEFT_PARENTHESIS", + "RIGHT_PARENTHESIS", + around="version specifier", + ): + tokenizer.consume("WS") + parsed_specifiers = _parse_version_many(tokenizer) + tokenizer.consume("WS") + + return parsed_specifiers + + +def _parse_version_many(tokenizer: Tokenizer) -> str: + """ + version_many = (SPECIFIER (WS? COMMA WS? SPECIFIER)*)? + """ + parsed_specifiers = "" + while tokenizer.check("SPECIFIER"): + span_start = tokenizer.position + parsed_specifiers += tokenizer.read().text + if tokenizer.check("VERSION_PREFIX_TRAIL", peek=True): + tokenizer.raise_syntax_error( + ".* suffix can only be used with `==` or `!=` operators", + span_start=span_start, + span_end=tokenizer.position + 1, + ) + if tokenizer.check("VERSION_LOCAL_LABEL_TRAIL", peek=True): + tokenizer.raise_syntax_error( + "Local version label can only be used with `==` or `!=` operators", + span_start=span_start, + span_end=tokenizer.position, + ) + tokenizer.consume("WS") + if not tokenizer.check("COMMA"): + break + parsed_specifiers += tokenizer.read().text + tokenizer.consume("WS") + + return parsed_specifiers + + +# -------------------------------------------------------------------------------------- +# Recursive descent parser for marker expression +# -------------------------------------------------------------------------------------- +def parse_marker(source: str) -> MarkerList: + return _parse_full_marker(Tokenizer(source, rules=DEFAULT_RULES)) + + +def _parse_full_marker(tokenizer: Tokenizer) -> MarkerList: + retval = _parse_marker(tokenizer) + tokenizer.expect("END", expected="end of marker expression") + return retval + + +def _parse_marker(tokenizer: Tokenizer) -> MarkerList: + """ + marker = marker_atom (BOOLOP marker_atom)+ + """ + expression = [_parse_marker_atom(tokenizer)] + while tokenizer.check("BOOLOP"): + token = tokenizer.read() + expr_right = _parse_marker_atom(tokenizer) + expression.extend((token.text, expr_right)) + return expression + + +def _parse_marker_atom(tokenizer: Tokenizer) -> MarkerAtom: + """ + marker_atom = WS? LEFT_PARENTHESIS WS? marker WS? RIGHT_PARENTHESIS WS? + | WS? marker_item WS? + """ + + tokenizer.consume("WS") + if tokenizer.check("LEFT_PARENTHESIS", peek=True): + with tokenizer.enclosing_tokens( + "LEFT_PARENTHESIS", + "RIGHT_PARENTHESIS", + around="marker expression", + ): + tokenizer.consume("WS") + marker: MarkerAtom = _parse_marker(tokenizer) + tokenizer.consume("WS") + else: + marker = _parse_marker_item(tokenizer) + tokenizer.consume("WS") + return marker + + +def _parse_marker_item(tokenizer: Tokenizer) -> MarkerItem: + """ + marker_item = WS? marker_var WS? marker_op WS? marker_var WS? + """ + tokenizer.consume("WS") + marker_var_left = _parse_marker_var(tokenizer) + tokenizer.consume("WS") + marker_op = _parse_marker_op(tokenizer) + tokenizer.consume("WS") + marker_var_right = _parse_marker_var(tokenizer) + tokenizer.consume("WS") + return (marker_var_left, marker_op, marker_var_right) + + +def _parse_marker_var(tokenizer: Tokenizer) -> MarkerVar: + """ + marker_var = VARIABLE | QUOTED_STRING + """ + if tokenizer.check("VARIABLE"): + return process_env_var(tokenizer.read().text.replace(".", "_")) + elif tokenizer.check("QUOTED_STRING"): + return process_python_str(tokenizer.read().text) + else: + tokenizer.raise_syntax_error( + message="Expected a marker variable or quoted string" + ) + + +def process_env_var(env_var: str) -> Variable: + if env_var in ("platform_python_implementation", "python_implementation"): + return Variable("platform_python_implementation") + else: + return Variable(env_var) + + +def process_python_str(python_str: str) -> Value: + value = ast.literal_eval(python_str) + return Value(str(value)) + + +def _parse_marker_op(tokenizer: Tokenizer) -> Op: + """ + marker_op = IN | NOT IN | OP + """ + if tokenizer.check("IN"): + tokenizer.read() + return Op("in") + elif tokenizer.check("NOT"): + tokenizer.read() + tokenizer.expect("WS", expected="whitespace after 'not'") + tokenizer.expect("IN", expected="'in' after 'not'") + return Op("not in") + elif tokenizer.check("OP"): + return Op(tokenizer.read().text) + else: + return tokenizer.raise_syntax_error( + "Expected marker operator, one of " + "<=, <, !=, ==, >=, >, ~=, ===, in, not in" + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_structures.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_structures.py new file mode 100644 index 0000000000000000000000000000000000000000..90a6465f9682c886363eea5327dac64bf623a6ff --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_structures.py @@ -0,0 +1,61 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. + + +class InfinityType: + def __repr__(self) -> str: + return "Infinity" + + def __hash__(self) -> int: + return hash(repr(self)) + + def __lt__(self, other: object) -> bool: + return False + + def __le__(self, other: object) -> bool: + return False + + def __eq__(self, other: object) -> bool: + return isinstance(other, self.__class__) + + def __gt__(self, other: object) -> bool: + return True + + def __ge__(self, other: object) -> bool: + return True + + def __neg__(self: object) -> "NegativeInfinityType": + return NegativeInfinity + + +Infinity = InfinityType() + + +class NegativeInfinityType: + def __repr__(self) -> str: + return "-Infinity" + + def __hash__(self) -> int: + return hash(repr(self)) + + def __lt__(self, other: object) -> bool: + return True + + def __le__(self, other: object) -> bool: + return True + + def __eq__(self, other: object) -> bool: + return isinstance(other, self.__class__) + + def __gt__(self, other: object) -> bool: + return False + + def __ge__(self, other: object) -> bool: + return False + + def __neg__(self: object) -> InfinityType: + return Infinity + + +NegativeInfinity = NegativeInfinityType() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_tokenizer.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..dd0d648d49a7c1a62d25ce5c9107aa448a8a22d1 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/_tokenizer.py @@ -0,0 +1,192 @@ +import contextlib +import re +from dataclasses import dataclass +from typing import Dict, Iterator, NoReturn, Optional, Tuple, Union + +from .specifiers import Specifier + + +@dataclass +class Token: + name: str + text: str + position: int + + +class ParserSyntaxError(Exception): + """The provided source text could not be parsed correctly.""" + + def __init__( + self, + message: str, + *, + source: str, + span: Tuple[int, int], + ) -> None: + self.span = span + self.message = message + self.source = source + + super().__init__() + + def __str__(self) -> str: + marker = " " * self.span[0] + "~" * (self.span[1] - self.span[0]) + "^" + return "\n ".join([self.message, self.source, marker]) + + +DEFAULT_RULES: "Dict[str, Union[str, re.Pattern[str]]]" = { + "LEFT_PARENTHESIS": r"\(", + "RIGHT_PARENTHESIS": r"\)", + "LEFT_BRACKET": r"\[", + "RIGHT_BRACKET": r"\]", + "SEMICOLON": r";", + "COMMA": r",", + "QUOTED_STRING": re.compile( + r""" + ( + ('[^']*') + | + ("[^"]*") + ) + """, + re.VERBOSE, + ), + "OP": r"(===|==|~=|!=|<=|>=|<|>)", + "BOOLOP": r"\b(or|and)\b", + "IN": r"\bin\b", + "NOT": r"\bnot\b", + "VARIABLE": re.compile( + r""" + \b( + python_version + |python_full_version + |os[._]name + |sys[._]platform + |platform_(release|system) + |platform[._](version|machine|python_implementation) + |python_implementation + |implementation_(name|version) + |extra + )\b + """, + re.VERBOSE, + ), + "SPECIFIER": re.compile( + Specifier._operator_regex_str + Specifier._version_regex_str, + re.VERBOSE | re.IGNORECASE, + ), + "AT": r"\@", + "URL": r"[^ \t]+", + "IDENTIFIER": r"\b[a-zA-Z0-9][a-zA-Z0-9._-]*\b", + "VERSION_PREFIX_TRAIL": r"\.\*", + "VERSION_LOCAL_LABEL_TRAIL": r"\+[a-z0-9]+(?:[-_\.][a-z0-9]+)*", + "WS": r"[ \t]+", + "END": r"$", +} + + +class Tokenizer: + """Context-sensitive token parsing. + + Provides methods to examine the input stream to check whether the next token + matches. + """ + + def __init__( + self, + source: str, + *, + rules: "Dict[str, Union[str, re.Pattern[str]]]", + ) -> None: + self.source = source + self.rules: Dict[str, re.Pattern[str]] = { + name: re.compile(pattern) for name, pattern in rules.items() + } + self.next_token: Optional[Token] = None + self.position = 0 + + def consume(self, name: str) -> None: + """Move beyond provided token name, if at current position.""" + if self.check(name): + self.read() + + def check(self, name: str, *, peek: bool = False) -> bool: + """Check whether the next token has the provided name. + + By default, if the check succeeds, the token *must* be read before + another check. If `peek` is set to `True`, the token is not loaded and + would need to be checked again. + """ + assert ( + self.next_token is None + ), f"Cannot check for {name!r}, already have {self.next_token!r}" + assert name in self.rules, f"Unknown token name: {name!r}" + + expression = self.rules[name] + + match = expression.match(self.source, self.position) + if match is None: + return False + if not peek: + self.next_token = Token(name, match[0], self.position) + return True + + def expect(self, name: str, *, expected: str) -> Token: + """Expect a certain token name next, failing with a syntax error otherwise. + + The token is *not* read. + """ + if not self.check(name): + raise self.raise_syntax_error(f"Expected {expected}") + return self.read() + + def read(self) -> Token: + """Consume the next token and return it.""" + token = self.next_token + assert token is not None + + self.position += len(token.text) + self.next_token = None + + return token + + def raise_syntax_error( + self, + message: str, + *, + span_start: Optional[int] = None, + span_end: Optional[int] = None, + ) -> NoReturn: + """Raise ParserSyntaxError at the given position.""" + span = ( + self.position if span_start is None else span_start, + self.position if span_end is None else span_end, + ) + raise ParserSyntaxError( + message, + source=self.source, + span=span, + ) + + @contextlib.contextmanager + def enclosing_tokens( + self, open_token: str, close_token: str, *, around: str + ) -> Iterator[None]: + if self.check(open_token): + open_position = self.position + self.read() + else: + open_position = None + + yield + + if open_position is None: + return + + if not self.check(close_token): + self.raise_syntax_error( + f"Expected matching {close_token} for {open_token}, after {around}", + span_start=open_position, + ) + + self.read() diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/markers.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/markers.py new file mode 100644 index 0000000000000000000000000000000000000000..c96d22a5a445e7353cd2454dec4255d3785c07b3 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/markers.py @@ -0,0 +1,253 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. + +import operator +import os +import platform +import sys +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from ._parser import ( + MarkerAtom, + MarkerList, + Op, + Value, + Variable, +) +from ._parser import ( + parse_marker as _parse_marker, +) +from ._tokenizer import ParserSyntaxError +from .specifiers import InvalidSpecifier, Specifier +from .utils import canonicalize_name + +__all__ = [ + "InvalidMarker", + "UndefinedComparison", + "UndefinedEnvironmentName", + "Marker", + "default_environment", +] + +Operator = Callable[[str, str], bool] + + +class InvalidMarker(ValueError): + """ + An invalid marker was found, users should refer to PEP 508. + """ + + +class UndefinedComparison(ValueError): + """ + An invalid operation was attempted on a value that doesn't support it. + """ + + +class UndefinedEnvironmentName(ValueError): + """ + A name was attempted to be used that does not exist inside of the + environment. + """ + + +def _normalize_extra_values(results: Any) -> Any: + """ + Normalize extra values. + """ + if isinstance(results[0], tuple): + lhs, op, rhs = results[0] + if isinstance(lhs, Variable) and lhs.value == "extra": + normalized_extra = canonicalize_name(rhs.value) + rhs = Value(normalized_extra) + elif isinstance(rhs, Variable) and rhs.value == "extra": + normalized_extra = canonicalize_name(lhs.value) + lhs = Value(normalized_extra) + results[0] = lhs, op, rhs + return results + + +def _format_marker( + marker: Union[List[str], MarkerAtom, str], first: Optional[bool] = True +) -> str: + assert isinstance(marker, (list, tuple, str)) + + # Sometimes we have a structure like [[...]] which is a single item list + # where the single item is itself it's own list. In that case we want skip + # the rest of this function so that we don't get extraneous () on the + # outside. + if ( + isinstance(marker, list) + and len(marker) == 1 + and isinstance(marker[0], (list, tuple)) + ): + return _format_marker(marker[0]) + + if isinstance(marker, list): + inner = (_format_marker(m, first=False) for m in marker) + if first: + return " ".join(inner) + else: + return "(" + " ".join(inner) + ")" + elif isinstance(marker, tuple): + return " ".join([m.serialize() for m in marker]) + else: + return marker + + +_operators: Dict[str, Operator] = { + "in": lambda lhs, rhs: lhs in rhs, + "not in": lambda lhs, rhs: lhs not in rhs, + "<": operator.lt, + "<=": operator.le, + "==": operator.eq, + "!=": operator.ne, + ">=": operator.ge, + ">": operator.gt, +} + + +def _eval_op(lhs: str, op: Op, rhs: str) -> bool: + try: + spec = Specifier("".join([op.serialize(), rhs])) + except InvalidSpecifier: + pass + else: + return spec.contains(lhs, prereleases=True) + + oper: Optional[Operator] = _operators.get(op.serialize()) + if oper is None: + raise UndefinedComparison(f"Undefined {op!r} on {lhs!r} and {rhs!r}.") + + return oper(lhs, rhs) + + +def _normalize(*values: str, key: str) -> Tuple[str, ...]: + # PEP 685 – Comparison of extra names for optional distribution dependencies + # https://peps.python.org/pep-0685/ + # > When comparing extra names, tools MUST normalize the names being + # > compared using the semantics outlined in PEP 503 for names + if key == "extra": + return tuple(canonicalize_name(v) for v in values) + + # other environment markers don't have such standards + return values + + +def _evaluate_markers(markers: MarkerList, environment: Dict[str, str]) -> bool: + groups: List[List[bool]] = [[]] + + for marker in markers: + assert isinstance(marker, (list, tuple, str)) + + if isinstance(marker, list): + groups[-1].append(_evaluate_markers(marker, environment)) + elif isinstance(marker, tuple): + lhs, op, rhs = marker + + if isinstance(lhs, Variable): + environment_key = lhs.value + lhs_value = environment[environment_key] + rhs_value = rhs.value + else: + lhs_value = lhs.value + environment_key = rhs.value + rhs_value = environment[environment_key] + + lhs_value, rhs_value = _normalize(lhs_value, rhs_value, key=environment_key) + groups[-1].append(_eval_op(lhs_value, op, rhs_value)) + else: + assert marker in ["and", "or"] + if marker == "or": + groups.append([]) + + return any(all(item) for item in groups) + + +def format_full_version(info: "sys._version_info") -> str: + version = "{0.major}.{0.minor}.{0.micro}".format(info) + kind = info.releaselevel + if kind != "final": + version += kind[0] + str(info.serial) + return version + + +def default_environment() -> Dict[str, str]: + iver = format_full_version(sys.implementation.version) + implementation_name = sys.implementation.name + return { + "implementation_name": implementation_name, + "implementation_version": iver, + "os_name": os.name, + "platform_machine": platform.machine(), + "platform_release": platform.release(), + "platform_system": platform.system(), + "platform_version": platform.version(), + "python_full_version": platform.python_version(), + "platform_python_implementation": platform.python_implementation(), + "python_version": ".".join(platform.python_version_tuple()[:2]), + "sys_platform": sys.platform, + } + + +class Marker: + def __init__(self, marker: str) -> None: + # Note: We create a Marker object without calling this constructor in + # packaging.requirements.Requirement. If any additional logic is + # added here, make sure to mirror/adapt Requirement. + try: + self._markers = _normalize_extra_values(_parse_marker(marker)) + # The attribute `_markers` can be described in terms of a recursive type: + # MarkerList = List[Union[Tuple[Node, ...], str, MarkerList]] + # + # For example, the following expression: + # python_version > "3.6" or (python_version == "3.6" and os_name == "unix") + # + # is parsed into: + # [ + # (, ')>, ), + # 'and', + # [ + # (, , ), + # 'or', + # (, , ) + # ] + # ] + except ParserSyntaxError as e: + raise InvalidMarker(str(e)) from e + + def __str__(self) -> str: + return _format_marker(self._markers) + + def __repr__(self) -> str: + return f"" + + def __hash__(self) -> int: + return hash((self.__class__.__name__, str(self))) + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, Marker): + return NotImplemented + + return str(self) == str(other) + + def evaluate(self, environment: Optional[Dict[str, str]] = None) -> bool: + """Evaluate a marker. + + Return the boolean from evaluating the given marker against the + environment. environment is an optional argument to override all or + part of the determined environment. + + The environment is determined from the current Python process. + """ + current_environment = default_environment() + current_environment["extra"] = "" + if environment is not None: + current_environment.update(environment) + # The API used to allow setting extra to None. We need to handle this + # case for backwards compatibility. + if current_environment["extra"] is None: + current_environment["extra"] = "" + + return _evaluate_markers(self._markers, current_environment) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/requirements.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/requirements.py new file mode 100644 index 0000000000000000000000000000000000000000..bdc43a7e98d87dba0c2069bfb4554f71d228cad4 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/requirements.py @@ -0,0 +1,90 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. + +from typing import Any, Iterator, Optional, Set + +from ._parser import parse_requirement as _parse_requirement +from ._tokenizer import ParserSyntaxError +from .markers import Marker, _normalize_extra_values +from .specifiers import SpecifierSet +from .utils import canonicalize_name + + +class InvalidRequirement(ValueError): + """ + An invalid requirement was found, users should refer to PEP 508. + """ + + +class Requirement: + """Parse a requirement. + + Parse a given requirement string into its parts, such as name, specifier, + URL, and extras. Raises InvalidRequirement on a badly-formed requirement + string. + """ + + # TODO: Can we test whether something is contained within a requirement? + # If so how do we do that? Do we need to test against the _name_ of + # the thing as well as the version? What about the markers? + # TODO: Can we normalize the name and extra name? + + def __init__(self, requirement_string: str) -> None: + try: + parsed = _parse_requirement(requirement_string) + except ParserSyntaxError as e: + raise InvalidRequirement(str(e)) from e + + self.name: str = parsed.name + self.url: Optional[str] = parsed.url or None + self.extras: Set[str] = set(parsed.extras or []) + self.specifier: SpecifierSet = SpecifierSet(parsed.specifier) + self.marker: Optional[Marker] = None + if parsed.marker is not None: + self.marker = Marker.__new__(Marker) + self.marker._markers = _normalize_extra_values(parsed.marker) + + def _iter_parts(self, name: str) -> Iterator[str]: + yield name + + if self.extras: + formatted_extras = ",".join(sorted(self.extras)) + yield f"[{formatted_extras}]" + + if self.specifier: + yield str(self.specifier) + + if self.url: + yield f"@ {self.url}" + if self.marker: + yield " " + + if self.marker: + yield f"; {self.marker}" + + def __str__(self) -> str: + return "".join(self._iter_parts(self.name)) + + def __repr__(self) -> str: + return f"" + + def __hash__(self) -> int: + return hash( + ( + self.__class__.__name__, + *self._iter_parts(canonicalize_name(self.name)), + ) + ) + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, Requirement): + return NotImplemented + + return ( + canonicalize_name(self.name) == canonicalize_name(other.name) + and self.extras == other.extras + and self.specifier == other.specifier + and self.url == other.url + and self.marker == other.marker + ) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/specifiers.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/specifiers.py new file mode 100644 index 0000000000000000000000000000000000000000..6d4066ae2770a3112f37d0e30d9a98fe59c4861f --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/specifiers.py @@ -0,0 +1,1011 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. +""" +.. testsetup:: + + from packaging.specifiers import Specifier, SpecifierSet, InvalidSpecifier + from packaging.version import Version +""" + +import abc +import itertools +import re +from typing import Callable, Iterable, Iterator, List, Optional, Tuple, TypeVar, Union + +from .utils import canonicalize_version +from .version import Version + +UnparsedVersion = Union[Version, str] +UnparsedVersionVar = TypeVar("UnparsedVersionVar", bound=UnparsedVersion) +CallableOperator = Callable[[Version, str], bool] + + +def _coerce_version(version: UnparsedVersion) -> Version: + if not isinstance(version, Version): + version = Version(version) + return version + + +class InvalidSpecifier(ValueError): + """ + Raised when attempting to create a :class:`Specifier` with a specifier + string that is invalid. + + >>> Specifier("lolwat") + Traceback (most recent call last): + ... + packaging.specifiers.InvalidSpecifier: Invalid specifier: 'lolwat' + """ + + +class BaseSpecifier(metaclass=abc.ABCMeta): + @abc.abstractmethod + def __str__(self) -> str: + """ + Returns the str representation of this Specifier-like object. This + should be representative of the Specifier itself. + """ + + @abc.abstractmethod + def __hash__(self) -> int: + """ + Returns a hash value for this Specifier-like object. + """ + + @abc.abstractmethod + def __eq__(self, other: object) -> bool: + """ + Returns a boolean representing whether or not the two Specifier-like + objects are equal. + + :param other: The other object to check against. + """ + + @property + @abc.abstractmethod + def prereleases(self) -> Optional[bool]: + """Whether or not pre-releases as a whole are allowed. + + This can be set to either ``True`` or ``False`` to explicitly enable or disable + prereleases or it can be set to ``None`` (the default) to use default semantics. + """ + + @prereleases.setter + def prereleases(self, value: bool) -> None: + """Setter for :attr:`prereleases`. + + :param value: The value to set. + """ + + @abc.abstractmethod + def contains(self, item: str, prereleases: Optional[bool] = None) -> bool: + """ + Determines if the given item is contained within this specifier. + """ + + @abc.abstractmethod + def filter( + self, iterable: Iterable[UnparsedVersionVar], prereleases: Optional[bool] = None + ) -> Iterator[UnparsedVersionVar]: + """ + Takes an iterable of items and filters them so that only items which + are contained within this specifier are allowed in it. + """ + + +class Specifier(BaseSpecifier): + """This class abstracts handling of version specifiers. + + .. tip:: + + It is generally not required to instantiate this manually. You should instead + prefer to work with :class:`SpecifierSet` instead, which can parse + comma-separated version specifiers (which is what package metadata contains). + """ + + _operator_regex_str = r""" + (?P(~=|==|!=|<=|>=|<|>|===)) + """ + _version_regex_str = r""" + (?P + (?: + # The identity operators allow for an escape hatch that will + # do an exact string match of the version you wish to install. + # This will not be parsed by PEP 440 and we cannot determine + # any semantic meaning from it. This operator is discouraged + # but included entirely as an escape hatch. + (?<====) # Only match for the identity operator + \s* + [^\s;)]* # The arbitrary version can be just about anything, + # we match everything except for whitespace, a + # semi-colon for marker support, and a closing paren + # since versions can be enclosed in them. + ) + | + (?: + # The (non)equality operators allow for wild card and local + # versions to be specified so we have to define these two + # operators separately to enable that. + (?<===|!=) # Only match for equals and not equals + + \s* + v? + (?:[0-9]+!)? # epoch + [0-9]+(?:\.[0-9]+)* # release + + # You cannot use a wild card and a pre-release, post-release, a dev or + # local version together so group them with a | and make them optional. + (?: + \.\* # Wild card syntax of .* + | + (?: # pre release + [-_\.]? + (alpha|beta|preview|pre|a|b|c|rc) + [-_\.]? + [0-9]* + )? + (?: # post release + (?:-[0-9]+)|(?:[-_\.]?(post|rev|r)[-_\.]?[0-9]*) + )? + (?:[-_\.]?dev[-_\.]?[0-9]*)? # dev release + (?:\+[a-z0-9]+(?:[-_\.][a-z0-9]+)*)? # local + )? + ) + | + (?: + # The compatible operator requires at least two digits in the + # release segment. + (?<=~=) # Only match for the compatible operator + + \s* + v? + (?:[0-9]+!)? # epoch + [0-9]+(?:\.[0-9]+)+ # release (We have a + instead of a *) + (?: # pre release + [-_\.]? + (alpha|beta|preview|pre|a|b|c|rc) + [-_\.]? + [0-9]* + )? + (?: # post release + (?:-[0-9]+)|(?:[-_\.]?(post|rev|r)[-_\.]?[0-9]*) + )? + (?:[-_\.]?dev[-_\.]?[0-9]*)? # dev release + ) + | + (?: + # All other operators only allow a sub set of what the + # (non)equality operators do. Specifically they do not allow + # local versions to be specified nor do they allow the prefix + # matching wild cards. + (?=": "greater_than_equal", + "<": "less_than", + ">": "greater_than", + "===": "arbitrary", + } + + def __init__(self, spec: str = "", prereleases: Optional[bool] = None) -> None: + """Initialize a Specifier instance. + + :param spec: + The string representation of a specifier which will be parsed and + normalized before use. + :param prereleases: + This tells the specifier if it should accept prerelease versions if + applicable or not. The default of ``None`` will autodetect it from the + given specifiers. + :raises InvalidSpecifier: + If the given specifier is invalid (i.e. bad syntax). + """ + match = self._regex.search(spec) + if not match: + raise InvalidSpecifier(f"Invalid specifier: '{spec}'") + + self._spec: Tuple[str, str] = ( + match.group("operator").strip(), + match.group("version").strip(), + ) + + # Store whether or not this Specifier should accept prereleases + self._prereleases = prereleases + + # https://github.com/python/mypy/pull/13475#pullrequestreview-1079784515 + @property # type: ignore[override] + def prereleases(self) -> bool: + # If there is an explicit prereleases set for this, then we'll just + # blindly use that. + if self._prereleases is not None: + return self._prereleases + + # Look at all of our specifiers and determine if they are inclusive + # operators, and if they are if they are including an explicit + # prerelease. + operator, version = self._spec + if operator in ["==", ">=", "<=", "~=", "==="]: + # The == specifier can include a trailing .*, if it does we + # want to remove before parsing. + if operator == "==" and version.endswith(".*"): + version = version[:-2] + + # Parse the version, and if it is a pre-release than this + # specifier allows pre-releases. + if Version(version).is_prerelease: + return True + + return False + + @prereleases.setter + def prereleases(self, value: bool) -> None: + self._prereleases = value + + @property + def operator(self) -> str: + """The operator of this specifier. + + >>> Specifier("==1.2.3").operator + '==' + """ + return self._spec[0] + + @property + def version(self) -> str: + """The version of this specifier. + + >>> Specifier("==1.2.3").version + '1.2.3' + """ + return self._spec[1] + + def __repr__(self) -> str: + """A representation of the Specifier that shows all internal state. + + >>> Specifier('>=1.0.0') + =1.0.0')> + >>> Specifier('>=1.0.0', prereleases=False) + =1.0.0', prereleases=False)> + >>> Specifier('>=1.0.0', prereleases=True) + =1.0.0', prereleases=True)> + """ + pre = ( + f", prereleases={self.prereleases!r}" + if self._prereleases is not None + else "" + ) + + return f"<{self.__class__.__name__}({str(self)!r}{pre})>" + + def __str__(self) -> str: + """A string representation of the Specifier that can be round-tripped. + + >>> str(Specifier('>=1.0.0')) + '>=1.0.0' + >>> str(Specifier('>=1.0.0', prereleases=False)) + '>=1.0.0' + """ + return "{}{}".format(*self._spec) + + @property + def _canonical_spec(self) -> Tuple[str, str]: + canonical_version = canonicalize_version( + self._spec[1], + strip_trailing_zero=(self._spec[0] != "~="), + ) + return self._spec[0], canonical_version + + def __hash__(self) -> int: + return hash(self._canonical_spec) + + def __eq__(self, other: object) -> bool: + """Whether or not the two Specifier-like objects are equal. + + :param other: The other object to check against. + + The value of :attr:`prereleases` is ignored. + + >>> Specifier("==1.2.3") == Specifier("== 1.2.3.0") + True + >>> (Specifier("==1.2.3", prereleases=False) == + ... Specifier("==1.2.3", prereleases=True)) + True + >>> Specifier("==1.2.3") == "==1.2.3" + True + >>> Specifier("==1.2.3") == Specifier("==1.2.4") + False + >>> Specifier("==1.2.3") == Specifier("~=1.2.3") + False + """ + if isinstance(other, str): + try: + other = self.__class__(str(other)) + except InvalidSpecifier: + return NotImplemented + elif not isinstance(other, self.__class__): + return NotImplemented + + return self._canonical_spec == other._canonical_spec + + def _get_operator(self, op: str) -> CallableOperator: + operator_callable: CallableOperator = getattr( + self, f"_compare_{self._operators[op]}" + ) + return operator_callable + + def _compare_compatible(self, prospective: Version, spec: str) -> bool: + # Compatible releases have an equivalent combination of >= and ==. That + # is that ~=2.2 is equivalent to >=2.2,==2.*. This allows us to + # implement this in terms of the other specifiers instead of + # implementing it ourselves. The only thing we need to do is construct + # the other specifiers. + + # We want everything but the last item in the version, but we want to + # ignore suffix segments. + prefix = _version_join( + list(itertools.takewhile(_is_not_suffix, _version_split(spec)))[:-1] + ) + + # Add the prefix notation to the end of our string + prefix += ".*" + + return self._get_operator(">=")(prospective, spec) and self._get_operator("==")( + prospective, prefix + ) + + def _compare_equal(self, prospective: Version, spec: str) -> bool: + # We need special logic to handle prefix matching + if spec.endswith(".*"): + # In the case of prefix matching we want to ignore local segment. + normalized_prospective = canonicalize_version( + prospective.public, strip_trailing_zero=False + ) + # Get the normalized version string ignoring the trailing .* + normalized_spec = canonicalize_version(spec[:-2], strip_trailing_zero=False) + # Split the spec out by bangs and dots, and pretend that there is + # an implicit dot in between a release segment and a pre-release segment. + split_spec = _version_split(normalized_spec) + + # Split the prospective version out by bangs and dots, and pretend + # that there is an implicit dot in between a release segment and + # a pre-release segment. + split_prospective = _version_split(normalized_prospective) + + # 0-pad the prospective version before shortening it to get the correct + # shortened version. + padded_prospective, _ = _pad_version(split_prospective, split_spec) + + # Shorten the prospective version to be the same length as the spec + # so that we can determine if the specifier is a prefix of the + # prospective version or not. + shortened_prospective = padded_prospective[: len(split_spec)] + + return shortened_prospective == split_spec + else: + # Convert our spec string into a Version + spec_version = Version(spec) + + # If the specifier does not have a local segment, then we want to + # act as if the prospective version also does not have a local + # segment. + if not spec_version.local: + prospective = Version(prospective.public) + + return prospective == spec_version + + def _compare_not_equal(self, prospective: Version, spec: str) -> bool: + return not self._compare_equal(prospective, spec) + + def _compare_less_than_equal(self, prospective: Version, spec: str) -> bool: + # NB: Local version identifiers are NOT permitted in the version + # specifier, so local version labels can be universally removed from + # the prospective version. + return Version(prospective.public) <= Version(spec) + + def _compare_greater_than_equal(self, prospective: Version, spec: str) -> bool: + # NB: Local version identifiers are NOT permitted in the version + # specifier, so local version labels can be universally removed from + # the prospective version. + return Version(prospective.public) >= Version(spec) + + def _compare_less_than(self, prospective: Version, spec_str: str) -> bool: + # Convert our spec to a Version instance, since we'll want to work with + # it as a version. + spec = Version(spec_str) + + # Check to see if the prospective version is less than the spec + # version. If it's not we can short circuit and just return False now + # instead of doing extra unneeded work. + if not prospective < spec: + return False + + # This special case is here so that, unless the specifier itself + # includes is a pre-release version, that we do not accept pre-release + # versions for the version mentioned in the specifier (e.g. <3.1 should + # not match 3.1.dev0, but should match 3.0.dev0). + if not spec.is_prerelease and prospective.is_prerelease: + if Version(prospective.base_version) == Version(spec.base_version): + return False + + # If we've gotten to here, it means that prospective version is both + # less than the spec version *and* it's not a pre-release of the same + # version in the spec. + return True + + def _compare_greater_than(self, prospective: Version, spec_str: str) -> bool: + # Convert our spec to a Version instance, since we'll want to work with + # it as a version. + spec = Version(spec_str) + + # Check to see if the prospective version is greater than the spec + # version. If it's not we can short circuit and just return False now + # instead of doing extra unneeded work. + if not prospective > spec: + return False + + # This special case is here so that, unless the specifier itself + # includes is a post-release version, that we do not accept + # post-release versions for the version mentioned in the specifier + # (e.g. >3.1 should not match 3.0.post0, but should match 3.2.post0). + if not spec.is_postrelease and prospective.is_postrelease: + if Version(prospective.base_version) == Version(spec.base_version): + return False + + # Ensure that we do not allow a local version of the version mentioned + # in the specifier, which is technically greater than, to match. + if prospective.local is not None: + if Version(prospective.base_version) == Version(spec.base_version): + return False + + # If we've gotten to here, it means that prospective version is both + # greater than the spec version *and* it's not a pre-release of the + # same version in the spec. + return True + + def _compare_arbitrary(self, prospective: Version, spec: str) -> bool: + return str(prospective).lower() == str(spec).lower() + + def __contains__(self, item: Union[str, Version]) -> bool: + """Return whether or not the item is contained in this specifier. + + :param item: The item to check for. + + This is used for the ``in`` operator and behaves the same as + :meth:`contains` with no ``prereleases`` argument passed. + + >>> "1.2.3" in Specifier(">=1.2.3") + True + >>> Version("1.2.3") in Specifier(">=1.2.3") + True + >>> "1.0.0" in Specifier(">=1.2.3") + False + >>> "1.3.0a1" in Specifier(">=1.2.3") + False + >>> "1.3.0a1" in Specifier(">=1.2.3", prereleases=True) + True + """ + return self.contains(item) + + def contains( + self, item: UnparsedVersion, prereleases: Optional[bool] = None + ) -> bool: + """Return whether or not the item is contained in this specifier. + + :param item: + The item to check for, which can be a version string or a + :class:`Version` instance. + :param prereleases: + Whether or not to match prereleases with this Specifier. If set to + ``None`` (the default), it uses :attr:`prereleases` to determine + whether or not prereleases are allowed. + + >>> Specifier(">=1.2.3").contains("1.2.3") + True + >>> Specifier(">=1.2.3").contains(Version("1.2.3")) + True + >>> Specifier(">=1.2.3").contains("1.0.0") + False + >>> Specifier(">=1.2.3").contains("1.3.0a1") + False + >>> Specifier(">=1.2.3", prereleases=True).contains("1.3.0a1") + True + >>> Specifier(">=1.2.3").contains("1.3.0a1", prereleases=True) + True + """ + + # Determine if prereleases are to be allowed or not. + if prereleases is None: + prereleases = self.prereleases + + # Normalize item to a Version, this allows us to have a shortcut for + # "2.0" in Specifier(">=2") + normalized_item = _coerce_version(item) + + # Determine if we should be supporting prereleases in this specifier + # or not, if we do not support prereleases than we can short circuit + # logic if this version is a prereleases. + if normalized_item.is_prerelease and not prereleases: + return False + + # Actually do the comparison to determine if this item is contained + # within this Specifier or not. + operator_callable: CallableOperator = self._get_operator(self.operator) + return operator_callable(normalized_item, self.version) + + def filter( + self, iterable: Iterable[UnparsedVersionVar], prereleases: Optional[bool] = None + ) -> Iterator[UnparsedVersionVar]: + """Filter items in the given iterable, that match the specifier. + + :param iterable: + An iterable that can contain version strings and :class:`Version` instances. + The items in the iterable will be filtered according to the specifier. + :param prereleases: + Whether or not to allow prereleases in the returned iterator. If set to + ``None`` (the default), it will be intelligently decide whether to allow + prereleases or not (based on the :attr:`prereleases` attribute, and + whether the only versions matching are prereleases). + + This method is smarter than just ``filter(Specifier().contains, [...])`` + because it implements the rule from :pep:`440` that a prerelease item + SHOULD be accepted if no other versions match the given specifier. + + >>> list(Specifier(">=1.2.3").filter(["1.2", "1.3", "1.5a1"])) + ['1.3'] + >>> list(Specifier(">=1.2.3").filter(["1.2", "1.2.3", "1.3", Version("1.4")])) + ['1.2.3', '1.3', ] + >>> list(Specifier(">=1.2.3").filter(["1.2", "1.5a1"])) + ['1.5a1'] + >>> list(Specifier(">=1.2.3").filter(["1.3", "1.5a1"], prereleases=True)) + ['1.3', '1.5a1'] + >>> list(Specifier(">=1.2.3", prereleases=True).filter(["1.3", "1.5a1"])) + ['1.3', '1.5a1'] + """ + + yielded = False + found_prereleases = [] + + kw = {"prereleases": prereleases if prereleases is not None else True} + + # Attempt to iterate over all the values in the iterable and if any of + # them match, yield them. + for version in iterable: + parsed_version = _coerce_version(version) + + if self.contains(parsed_version, **kw): + # If our version is a prerelease, and we were not set to allow + # prereleases, then we'll store it for later in case nothing + # else matches this specifier. + if parsed_version.is_prerelease and not ( + prereleases or self.prereleases + ): + found_prereleases.append(version) + # Either this is not a prerelease, or we should have been + # accepting prereleases from the beginning. + else: + yielded = True + yield version + + # Now that we've iterated over everything, determine if we've yielded + # any values, and if we have not and we have any prereleases stored up + # then we will go ahead and yield the prereleases. + if not yielded and found_prereleases: + for version in found_prereleases: + yield version + + +_prefix_regex = re.compile(r"^([0-9]+)((?:a|b|c|rc)[0-9]+)$") + + +def _version_split(version: str) -> List[str]: + """Split version into components. + + The split components are intended for version comparison. The logic does + not attempt to retain the original version string, so joining the + components back with :func:`_version_join` may not produce the original + version string. + """ + result: List[str] = [] + + epoch, _, rest = version.rpartition("!") + result.append(epoch or "0") + + for item in rest.split("."): + match = _prefix_regex.search(item) + if match: + result.extend(match.groups()) + else: + result.append(item) + return result + + +def _version_join(components: List[str]) -> str: + """Join split version components into a version string. + + This function assumes the input came from :func:`_version_split`, where the + first component must be the epoch (either empty or numeric), and all other + components numeric. + """ + epoch, *rest = components + return f"{epoch}!{'.'.join(rest)}" + + +def _is_not_suffix(segment: str) -> bool: + return not any( + segment.startswith(prefix) for prefix in ("dev", "a", "b", "rc", "post") + ) + + +def _pad_version(left: List[str], right: List[str]) -> Tuple[List[str], List[str]]: + left_split, right_split = [], [] + + # Get the release segment of our versions + left_split.append(list(itertools.takewhile(lambda x: x.isdigit(), left))) + right_split.append(list(itertools.takewhile(lambda x: x.isdigit(), right))) + + # Get the rest of our versions + left_split.append(left[len(left_split[0]) :]) + right_split.append(right[len(right_split[0]) :]) + + # Insert our padding + left_split.insert(1, ["0"] * max(0, len(right_split[0]) - len(left_split[0]))) + right_split.insert(1, ["0"] * max(0, len(left_split[0]) - len(right_split[0]))) + + return ( + list(itertools.chain.from_iterable(left_split)), + list(itertools.chain.from_iterable(right_split)), + ) + + +class SpecifierSet(BaseSpecifier): + """This class abstracts handling of a set of version specifiers. + + It can be passed a single specifier (``>=3.0``), a comma-separated list of + specifiers (``>=3.0,!=3.1``), or no specifier at all. + """ + + def __init__( + self, specifiers: str = "", prereleases: Optional[bool] = None + ) -> None: + """Initialize a SpecifierSet instance. + + :param specifiers: + The string representation of a specifier or a comma-separated list of + specifiers which will be parsed and normalized before use. + :param prereleases: + This tells the SpecifierSet if it should accept prerelease versions if + applicable or not. The default of ``None`` will autodetect it from the + given specifiers. + + :raises InvalidSpecifier: + If the given ``specifiers`` are not parseable than this exception will be + raised. + """ + + # Split on `,` to break each individual specifier into it's own item, and + # strip each item to remove leading/trailing whitespace. + split_specifiers = [s.strip() for s in specifiers.split(",") if s.strip()] + + # Make each individual specifier a Specifier and save in a frozen set for later. + self._specs = frozenset(map(Specifier, split_specifiers)) + + # Store our prereleases value so we can use it later to determine if + # we accept prereleases or not. + self._prereleases = prereleases + + @property + def prereleases(self) -> Optional[bool]: + # If we have been given an explicit prerelease modifier, then we'll + # pass that through here. + if self._prereleases is not None: + return self._prereleases + + # If we don't have any specifiers, and we don't have a forced value, + # then we'll just return None since we don't know if this should have + # pre-releases or not. + if not self._specs: + return None + + # Otherwise we'll see if any of the given specifiers accept + # prereleases, if any of them do we'll return True, otherwise False. + return any(s.prereleases for s in self._specs) + + @prereleases.setter + def prereleases(self, value: bool) -> None: + self._prereleases = value + + def __repr__(self) -> str: + """A representation of the specifier set that shows all internal state. + + Note that the ordering of the individual specifiers within the set may not + match the input string. + + >>> SpecifierSet('>=1.0.0,!=2.0.0') + =1.0.0')> + >>> SpecifierSet('>=1.0.0,!=2.0.0', prereleases=False) + =1.0.0', prereleases=False)> + >>> SpecifierSet('>=1.0.0,!=2.0.0', prereleases=True) + =1.0.0', prereleases=True)> + """ + pre = ( + f", prereleases={self.prereleases!r}" + if self._prereleases is not None + else "" + ) + + return f"" + + def __str__(self) -> str: + """A string representation of the specifier set that can be round-tripped. + + Note that the ordering of the individual specifiers within the set may not + match the input string. + + >>> str(SpecifierSet(">=1.0.0,!=1.0.1")) + '!=1.0.1,>=1.0.0' + >>> str(SpecifierSet(">=1.0.0,!=1.0.1", prereleases=False)) + '!=1.0.1,>=1.0.0' + """ + return ",".join(sorted(str(s) for s in self._specs)) + + def __hash__(self) -> int: + return hash(self._specs) + + def __and__(self, other: Union["SpecifierSet", str]) -> "SpecifierSet": + """Return a SpecifierSet which is a combination of the two sets. + + :param other: The other object to combine with. + + >>> SpecifierSet(">=1.0.0,!=1.0.1") & '<=2.0.0,!=2.0.1' + =1.0.0')> + >>> SpecifierSet(">=1.0.0,!=1.0.1") & SpecifierSet('<=2.0.0,!=2.0.1') + =1.0.0')> + """ + if isinstance(other, str): + other = SpecifierSet(other) + elif not isinstance(other, SpecifierSet): + return NotImplemented + + specifier = SpecifierSet() + specifier._specs = frozenset(self._specs | other._specs) + + if self._prereleases is None and other._prereleases is not None: + specifier._prereleases = other._prereleases + elif self._prereleases is not None and other._prereleases is None: + specifier._prereleases = self._prereleases + elif self._prereleases == other._prereleases: + specifier._prereleases = self._prereleases + else: + raise ValueError( + "Cannot combine SpecifierSets with True and False prerelease " + "overrides." + ) + + return specifier + + def __eq__(self, other: object) -> bool: + """Whether or not the two SpecifierSet-like objects are equal. + + :param other: The other object to check against. + + The value of :attr:`prereleases` is ignored. + + >>> SpecifierSet(">=1.0.0,!=1.0.1") == SpecifierSet(">=1.0.0,!=1.0.1") + True + >>> (SpecifierSet(">=1.0.0,!=1.0.1", prereleases=False) == + ... SpecifierSet(">=1.0.0,!=1.0.1", prereleases=True)) + True + >>> SpecifierSet(">=1.0.0,!=1.0.1") == ">=1.0.0,!=1.0.1" + True + >>> SpecifierSet(">=1.0.0,!=1.0.1") == SpecifierSet(">=1.0.0") + False + >>> SpecifierSet(">=1.0.0,!=1.0.1") == SpecifierSet(">=1.0.0,!=1.0.2") + False + """ + if isinstance(other, (str, Specifier)): + other = SpecifierSet(str(other)) + elif not isinstance(other, SpecifierSet): + return NotImplemented + + return self._specs == other._specs + + def __len__(self) -> int: + """Returns the number of specifiers in this specifier set.""" + return len(self._specs) + + def __iter__(self) -> Iterator[Specifier]: + """ + Returns an iterator over all the underlying :class:`Specifier` instances + in this specifier set. + + >>> sorted(SpecifierSet(">=1.0.0,!=1.0.1"), key=str) + [, =1.0.0')>] + """ + return iter(self._specs) + + def __contains__(self, item: UnparsedVersion) -> bool: + """Return whether or not the item is contained in this specifier. + + :param item: The item to check for. + + This is used for the ``in`` operator and behaves the same as + :meth:`contains` with no ``prereleases`` argument passed. + + >>> "1.2.3" in SpecifierSet(">=1.0.0,!=1.0.1") + True + >>> Version("1.2.3") in SpecifierSet(">=1.0.0,!=1.0.1") + True + >>> "1.0.1" in SpecifierSet(">=1.0.0,!=1.0.1") + False + >>> "1.3.0a1" in SpecifierSet(">=1.0.0,!=1.0.1") + False + >>> "1.3.0a1" in SpecifierSet(">=1.0.0,!=1.0.1", prereleases=True) + True + """ + return self.contains(item) + + def contains( + self, + item: UnparsedVersion, + prereleases: Optional[bool] = None, + installed: Optional[bool] = None, + ) -> bool: + """Return whether or not the item is contained in this SpecifierSet. + + :param item: + The item to check for, which can be a version string or a + :class:`Version` instance. + :param prereleases: + Whether or not to match prereleases with this SpecifierSet. If set to + ``None`` (the default), it uses :attr:`prereleases` to determine + whether or not prereleases are allowed. + + >>> SpecifierSet(">=1.0.0,!=1.0.1").contains("1.2.3") + True + >>> SpecifierSet(">=1.0.0,!=1.0.1").contains(Version("1.2.3")) + True + >>> SpecifierSet(">=1.0.0,!=1.0.1").contains("1.0.1") + False + >>> SpecifierSet(">=1.0.0,!=1.0.1").contains("1.3.0a1") + False + >>> SpecifierSet(">=1.0.0,!=1.0.1", prereleases=True).contains("1.3.0a1") + True + >>> SpecifierSet(">=1.0.0,!=1.0.1").contains("1.3.0a1", prereleases=True) + True + """ + # Ensure that our item is a Version instance. + if not isinstance(item, Version): + item = Version(item) + + # Determine if we're forcing a prerelease or not, if we're not forcing + # one for this particular filter call, then we'll use whatever the + # SpecifierSet thinks for whether or not we should support prereleases. + if prereleases is None: + prereleases = self.prereleases + + # We can determine if we're going to allow pre-releases by looking to + # see if any of the underlying items supports them. If none of them do + # and this item is a pre-release then we do not allow it and we can + # short circuit that here. + # Note: This means that 1.0.dev1 would not be contained in something + # like >=1.0.devabc however it would be in >=1.0.debabc,>0.0.dev0 + if not prereleases and item.is_prerelease: + return False + + if installed and item.is_prerelease: + item = Version(item.base_version) + + # We simply dispatch to the underlying specs here to make sure that the + # given version is contained within all of them. + # Note: This use of all() here means that an empty set of specifiers + # will always return True, this is an explicit design decision. + return all(s.contains(item, prereleases=prereleases) for s in self._specs) + + def filter( + self, iterable: Iterable[UnparsedVersionVar], prereleases: Optional[bool] = None + ) -> Iterator[UnparsedVersionVar]: + """Filter items in the given iterable, that match the specifiers in this set. + + :param iterable: + An iterable that can contain version strings and :class:`Version` instances. + The items in the iterable will be filtered according to the specifier. + :param prereleases: + Whether or not to allow prereleases in the returned iterator. If set to + ``None`` (the default), it will be intelligently decide whether to allow + prereleases or not (based on the :attr:`prereleases` attribute, and + whether the only versions matching are prereleases). + + This method is smarter than just ``filter(SpecifierSet(...).contains, [...])`` + because it implements the rule from :pep:`440` that a prerelease item + SHOULD be accepted if no other versions match the given specifier. + + >>> list(SpecifierSet(">=1.2.3").filter(["1.2", "1.3", "1.5a1"])) + ['1.3'] + >>> list(SpecifierSet(">=1.2.3").filter(["1.2", "1.3", Version("1.4")])) + ['1.3', ] + >>> list(SpecifierSet(">=1.2.3").filter(["1.2", "1.5a1"])) + [] + >>> list(SpecifierSet(">=1.2.3").filter(["1.3", "1.5a1"], prereleases=True)) + ['1.3', '1.5a1'] + >>> list(SpecifierSet(">=1.2.3", prereleases=True).filter(["1.3", "1.5a1"])) + ['1.3', '1.5a1'] + + An "empty" SpecifierSet will filter items based on the presence of prerelease + versions in the set. + + >>> list(SpecifierSet("").filter(["1.3", "1.5a1"])) + ['1.3'] + >>> list(SpecifierSet("").filter(["1.5a1"])) + ['1.5a1'] + >>> list(SpecifierSet("", prereleases=True).filter(["1.3", "1.5a1"])) + ['1.3', '1.5a1'] + >>> list(SpecifierSet("").filter(["1.3", "1.5a1"], prereleases=True)) + ['1.3', '1.5a1'] + """ + # Determine if we're forcing a prerelease or not, if we're not forcing + # one for this particular filter call, then we'll use whatever the + # SpecifierSet thinks for whether or not we should support prereleases. + if prereleases is None: + prereleases = self.prereleases + + # If we have any specifiers, then we want to wrap our iterable in the + # filter method for each one, this will act as a logical AND amongst + # each specifier. + if self._specs: + for spec in self._specs: + iterable = spec.filter(iterable, prereleases=bool(prereleases)) + return iter(iterable) + # If we do not have any specifiers, then we need to have a rough filter + # which will filter out any pre-releases, unless there are no final + # releases. + else: + filtered: List[UnparsedVersionVar] = [] + found_prereleases: List[UnparsedVersionVar] = [] + + for item in iterable: + parsed_version = _coerce_version(item) + + # Store any item which is a pre-release for later unless we've + # already found a final version or we are accepting prereleases + if parsed_version.is_prerelease and not prereleases: + if not filtered: + found_prereleases.append(item) + else: + filtered.append(item) + + # If we've found no items except for pre-releases, then we'll go + # ahead and use the pre-releases + if not filtered and found_prereleases and prereleases is None: + return iter(found_prereleases) + + return iter(filtered) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/tags.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/tags.py new file mode 100644 index 0000000000000000000000000000000000000000..89f1926137dd2d2a6bd63616bf5b9f722fc8d584 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/tags.py @@ -0,0 +1,571 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. + +import logging +import platform +import re +import struct +import subprocess +import sys +import sysconfig +from importlib.machinery import EXTENSION_SUFFIXES +from typing import ( + Dict, + FrozenSet, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, + Union, + cast, +) + +from . import _manylinux, _musllinux + +logger = logging.getLogger(__name__) + +PythonVersion = Sequence[int] +MacVersion = Tuple[int, int] + +INTERPRETER_SHORT_NAMES: Dict[str, str] = { + "python": "py", # Generic. + "cpython": "cp", + "pypy": "pp", + "ironpython": "ip", + "jython": "jy", +} + + +_32_BIT_INTERPRETER = struct.calcsize("P") == 4 + + +class Tag: + """ + A representation of the tag triple for a wheel. + + Instances are considered immutable and thus are hashable. Equality checking + is also supported. + """ + + __slots__ = ["_interpreter", "_abi", "_platform", "_hash"] + + def __init__(self, interpreter: str, abi: str, platform: str) -> None: + self._interpreter = interpreter.lower() + self._abi = abi.lower() + self._platform = platform.lower() + # The __hash__ of every single element in a Set[Tag] will be evaluated each time + # that a set calls its `.disjoint()` method, which may be called hundreds of + # times when scanning a page of links for packages with tags matching that + # Set[Tag]. Pre-computing the value here produces significant speedups for + # downstream consumers. + self._hash = hash((self._interpreter, self._abi, self._platform)) + + @property + def interpreter(self) -> str: + return self._interpreter + + @property + def abi(self) -> str: + return self._abi + + @property + def platform(self) -> str: + return self._platform + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Tag): + return NotImplemented + + return ( + (self._hash == other._hash) # Short-circuit ASAP for perf reasons. + and (self._platform == other._platform) + and (self._abi == other._abi) + and (self._interpreter == other._interpreter) + ) + + def __hash__(self) -> int: + return self._hash + + def __str__(self) -> str: + return f"{self._interpreter}-{self._abi}-{self._platform}" + + def __repr__(self) -> str: + return f"<{self} @ {id(self)}>" + + +def parse_tag(tag: str) -> FrozenSet[Tag]: + """ + Parses the provided tag (e.g. `py3-none-any`) into a frozenset of Tag instances. + + Returning a set is required due to the possibility that the tag is a + compressed tag set. + """ + tags = set() + interpreters, abis, platforms = tag.split("-") + for interpreter in interpreters.split("."): + for abi in abis.split("."): + for platform_ in platforms.split("."): + tags.add(Tag(interpreter, abi, platform_)) + return frozenset(tags) + + +def _get_config_var(name: str, warn: bool = False) -> Union[int, str, None]: + value: Union[int, str, None] = sysconfig.get_config_var(name) + if value is None and warn: + logger.debug( + "Config variable '%s' is unset, Python ABI tag may be incorrect", name + ) + return value + + +def _normalize_string(string: str) -> str: + return string.replace(".", "_").replace("-", "_").replace(" ", "_") + + +def _is_threaded_cpython(abis: List[str]) -> bool: + """ + Determine if the ABI corresponds to a threaded (`--disable-gil`) build. + + The threaded builds are indicated by a "t" in the abiflags. + """ + if len(abis) == 0: + return False + # expect e.g., cp313 + m = re.match(r"cp\d+(.*)", abis[0]) + if not m: + return False + abiflags = m.group(1) + return "t" in abiflags + + +def _abi3_applies(python_version: PythonVersion, threading: bool) -> bool: + """ + Determine if the Python version supports abi3. + + PEP 384 was first implemented in Python 3.2. The threaded (`--disable-gil`) + builds do not support abi3. + """ + return len(python_version) > 1 and tuple(python_version) >= (3, 2) and not threading + + +def _cpython_abis(py_version: PythonVersion, warn: bool = False) -> List[str]: + py_version = tuple(py_version) # To allow for version comparison. + abis = [] + version = _version_nodot(py_version[:2]) + threading = debug = pymalloc = ucs4 = "" + with_debug = _get_config_var("Py_DEBUG", warn) + has_refcount = hasattr(sys, "gettotalrefcount") + # Windows doesn't set Py_DEBUG, so checking for support of debug-compiled + # extension modules is the best option. + # https://github.com/pypa/pip/issues/3383#issuecomment-173267692 + has_ext = "_d.pyd" in EXTENSION_SUFFIXES + if with_debug or (with_debug is None and (has_refcount or has_ext)): + debug = "d" + if py_version >= (3, 13) and _get_config_var("Py_GIL_DISABLED", warn): + threading = "t" + if py_version < (3, 8): + with_pymalloc = _get_config_var("WITH_PYMALLOC", warn) + if with_pymalloc or with_pymalloc is None: + pymalloc = "m" + if py_version < (3, 3): + unicode_size = _get_config_var("Py_UNICODE_SIZE", warn) + if unicode_size == 4 or ( + unicode_size is None and sys.maxunicode == 0x10FFFF + ): + ucs4 = "u" + elif debug: + # Debug builds can also load "normal" extension modules. + # We can also assume no UCS-4 or pymalloc requirement. + abis.append(f"cp{version}{threading}") + abis.insert(0, f"cp{version}{threading}{debug}{pymalloc}{ucs4}") + return abis + + +def cpython_tags( + python_version: Optional[PythonVersion] = None, + abis: Optional[Iterable[str]] = None, + platforms: Optional[Iterable[str]] = None, + *, + warn: bool = False, +) -> Iterator[Tag]: + """ + Yields the tags for a CPython interpreter. + + The tags consist of: + - cp-- + - cp-abi3- + - cp-none- + - cp-abi3- # Older Python versions down to 3.2. + + If python_version only specifies a major version then user-provided ABIs and + the 'none' ABItag will be used. + + If 'abi3' or 'none' are specified in 'abis' then they will be yielded at + their normal position and not at the beginning. + """ + if not python_version: + python_version = sys.version_info[:2] + + interpreter = f"cp{_version_nodot(python_version[:2])}" + + if abis is None: + if len(python_version) > 1: + abis = _cpython_abis(python_version, warn) + else: + abis = [] + abis = list(abis) + # 'abi3' and 'none' are explicitly handled later. + for explicit_abi in ("abi3", "none"): + try: + abis.remove(explicit_abi) + except ValueError: + pass + + platforms = list(platforms or platform_tags()) + for abi in abis: + for platform_ in platforms: + yield Tag(interpreter, abi, platform_) + + threading = _is_threaded_cpython(abis) + use_abi3 = _abi3_applies(python_version, threading) + if use_abi3: + yield from (Tag(interpreter, "abi3", platform_) for platform_ in platforms) + yield from (Tag(interpreter, "none", platform_) for platform_ in platforms) + + if use_abi3: + for minor_version in range(python_version[1] - 1, 1, -1): + for platform_ in platforms: + interpreter = "cp{version}".format( + version=_version_nodot((python_version[0], minor_version)) + ) + yield Tag(interpreter, "abi3", platform_) + + +def _generic_abi() -> List[str]: + """ + Return the ABI tag based on EXT_SUFFIX. + """ + # The following are examples of `EXT_SUFFIX`. + # We want to keep the parts which are related to the ABI and remove the + # parts which are related to the platform: + # - linux: '.cpython-310-x86_64-linux-gnu.so' => cp310 + # - mac: '.cpython-310-darwin.so' => cp310 + # - win: '.cp310-win_amd64.pyd' => cp310 + # - win: '.pyd' => cp37 (uses _cpython_abis()) + # - pypy: '.pypy38-pp73-x86_64-linux-gnu.so' => pypy38_pp73 + # - graalpy: '.graalpy-38-native-x86_64-darwin.dylib' + # => graalpy_38_native + + ext_suffix = _get_config_var("EXT_SUFFIX", warn=True) + if not isinstance(ext_suffix, str) or ext_suffix[0] != ".": + raise SystemError("invalid sysconfig.get_config_var('EXT_SUFFIX')") + parts = ext_suffix.split(".") + if len(parts) < 3: + # CPython3.7 and earlier uses ".pyd" on Windows. + return _cpython_abis(sys.version_info[:2]) + soabi = parts[1] + if soabi.startswith("cpython"): + # non-windows + abi = "cp" + soabi.split("-")[1] + elif soabi.startswith("cp"): + # windows + abi = soabi.split("-")[0] + elif soabi.startswith("pypy"): + abi = "-".join(soabi.split("-")[:2]) + elif soabi.startswith("graalpy"): + abi = "-".join(soabi.split("-")[:3]) + elif soabi: + # pyston, ironpython, others? + abi = soabi + else: + return [] + return [_normalize_string(abi)] + + +def generic_tags( + interpreter: Optional[str] = None, + abis: Optional[Iterable[str]] = None, + platforms: Optional[Iterable[str]] = None, + *, + warn: bool = False, +) -> Iterator[Tag]: + """ + Yields the tags for a generic interpreter. + + The tags consist of: + - -- + + The "none" ABI will be added if it was not explicitly provided. + """ + if not interpreter: + interp_name = interpreter_name() + interp_version = interpreter_version(warn=warn) + interpreter = "".join([interp_name, interp_version]) + if abis is None: + abis = _generic_abi() + else: + abis = list(abis) + platforms = list(platforms or platform_tags()) + if "none" not in abis: + abis.append("none") + for abi in abis: + for platform_ in platforms: + yield Tag(interpreter, abi, platform_) + + +def _py_interpreter_range(py_version: PythonVersion) -> Iterator[str]: + """ + Yields Python versions in descending order. + + After the latest version, the major-only version will be yielded, and then + all previous versions of that major version. + """ + if len(py_version) > 1: + yield f"py{_version_nodot(py_version[:2])}" + yield f"py{py_version[0]}" + if len(py_version) > 1: + for minor in range(py_version[1] - 1, -1, -1): + yield f"py{_version_nodot((py_version[0], minor))}" + + +def compatible_tags( + python_version: Optional[PythonVersion] = None, + interpreter: Optional[str] = None, + platforms: Optional[Iterable[str]] = None, +) -> Iterator[Tag]: + """ + Yields the sequence of tags that are compatible with a specific version of Python. + + The tags consist of: + - py*-none- + - -none-any # ... if `interpreter` is provided. + - py*-none-any + """ + if not python_version: + python_version = sys.version_info[:2] + platforms = list(platforms or platform_tags()) + for version in _py_interpreter_range(python_version): + for platform_ in platforms: + yield Tag(version, "none", platform_) + if interpreter: + yield Tag(interpreter, "none", "any") + for version in _py_interpreter_range(python_version): + yield Tag(version, "none", "any") + + +def _mac_arch(arch: str, is_32bit: bool = _32_BIT_INTERPRETER) -> str: + if not is_32bit: + return arch + + if arch.startswith("ppc"): + return "ppc" + + return "i386" + + +def _mac_binary_formats(version: MacVersion, cpu_arch: str) -> List[str]: + formats = [cpu_arch] + if cpu_arch == "x86_64": + if version < (10, 4): + return [] + formats.extend(["intel", "fat64", "fat32"]) + + elif cpu_arch == "i386": + if version < (10, 4): + return [] + formats.extend(["intel", "fat32", "fat"]) + + elif cpu_arch == "ppc64": + # TODO: Need to care about 32-bit PPC for ppc64 through 10.2? + if version > (10, 5) or version < (10, 4): + return [] + formats.append("fat64") + + elif cpu_arch == "ppc": + if version > (10, 6): + return [] + formats.extend(["fat32", "fat"]) + + if cpu_arch in {"arm64", "x86_64"}: + formats.append("universal2") + + if cpu_arch in {"x86_64", "i386", "ppc64", "ppc", "intel"}: + formats.append("universal") + + return formats + + +def mac_platforms( + version: Optional[MacVersion] = None, arch: Optional[str] = None +) -> Iterator[str]: + """ + Yields the platform tags for a macOS system. + + The `version` parameter is a two-item tuple specifying the macOS version to + generate platform tags for. The `arch` parameter is the CPU architecture to + generate platform tags for. Both parameters default to the appropriate value + for the current system. + """ + version_str, _, cpu_arch = platform.mac_ver() + if version is None: + version = cast("MacVersion", tuple(map(int, version_str.split(".")[:2]))) + if version == (10, 16): + # When built against an older macOS SDK, Python will report macOS 10.16 + # instead of the real version. + version_str = subprocess.run( + [ + sys.executable, + "-sS", + "-c", + "import platform; print(platform.mac_ver()[0])", + ], + check=True, + env={"SYSTEM_VERSION_COMPAT": "0"}, + stdout=subprocess.PIPE, + text=True, + ).stdout + version = cast("MacVersion", tuple(map(int, version_str.split(".")[:2]))) + else: + version = version + if arch is None: + arch = _mac_arch(cpu_arch) + else: + arch = arch + + if (10, 0) <= version and version < (11, 0): + # Prior to Mac OS 11, each yearly release of Mac OS bumped the + # "minor" version number. The major version was always 10. + for minor_version in range(version[1], -1, -1): + compat_version = 10, minor_version + binary_formats = _mac_binary_formats(compat_version, arch) + for binary_format in binary_formats: + yield "macosx_{major}_{minor}_{binary_format}".format( + major=10, minor=minor_version, binary_format=binary_format + ) + + if version >= (11, 0): + # Starting with Mac OS 11, each yearly release bumps the major version + # number. The minor versions are now the midyear updates. + for major_version in range(version[0], 10, -1): + compat_version = major_version, 0 + binary_formats = _mac_binary_formats(compat_version, arch) + for binary_format in binary_formats: + yield "macosx_{major}_{minor}_{binary_format}".format( + major=major_version, minor=0, binary_format=binary_format + ) + + if version >= (11, 0): + # Mac OS 11 on x86_64 is compatible with binaries from previous releases. + # Arm64 support was introduced in 11.0, so no Arm binaries from previous + # releases exist. + # + # However, the "universal2" binary format can have a + # macOS version earlier than 11.0 when the x86_64 part of the binary supports + # that version of macOS. + if arch == "x86_64": + for minor_version in range(16, 3, -1): + compat_version = 10, minor_version + binary_formats = _mac_binary_formats(compat_version, arch) + for binary_format in binary_formats: + yield "macosx_{major}_{minor}_{binary_format}".format( + major=compat_version[0], + minor=compat_version[1], + binary_format=binary_format, + ) + else: + for minor_version in range(16, 3, -1): + compat_version = 10, minor_version + binary_format = "universal2" + yield "macosx_{major}_{minor}_{binary_format}".format( + major=compat_version[0], + minor=compat_version[1], + binary_format=binary_format, + ) + + +def _linux_platforms(is_32bit: bool = _32_BIT_INTERPRETER) -> Iterator[str]: + linux = _normalize_string(sysconfig.get_platform()) + if not linux.startswith("linux_"): + # we should never be here, just yield the sysconfig one and return + yield linux + return + if is_32bit: + if linux == "linux_x86_64": + linux = "linux_i686" + elif linux == "linux_aarch64": + linux = "linux_armv8l" + _, arch = linux.split("_", 1) + archs = {"armv8l": ["armv8l", "armv7l"]}.get(arch, [arch]) + yield from _manylinux.platform_tags(archs) + yield from _musllinux.platform_tags(archs) + for arch in archs: + yield f"linux_{arch}" + + +def _generic_platforms() -> Iterator[str]: + yield _normalize_string(sysconfig.get_platform()) + + +def platform_tags() -> Iterator[str]: + """ + Provides the platform tags for this installation. + """ + if platform.system() == "Darwin": + return mac_platforms() + elif platform.system() == "Linux": + return _linux_platforms() + else: + return _generic_platforms() + + +def interpreter_name() -> str: + """ + Returns the name of the running interpreter. + + Some implementations have a reserved, two-letter abbreviation which will + be returned when appropriate. + """ + name = sys.implementation.name + return INTERPRETER_SHORT_NAMES.get(name) or name + + +def interpreter_version(*, warn: bool = False) -> str: + """ + Returns the version of the running interpreter. + """ + version = _get_config_var("py_version_nodot", warn=warn) + if version: + version = str(version) + else: + version = _version_nodot(sys.version_info[:2]) + return version + + +def _version_nodot(version: PythonVersion) -> str: + return "".join(map(str, version)) + + +def sys_tags(*, warn: bool = False) -> Iterator[Tag]: + """ + Returns the sequence of tag triples for the running interpreter. + + The order of the sequence corresponds to priority order for the + interpreter, from most to least important. + """ + + interp_name = interpreter_name() + if interp_name == "cp": + yield from cpython_tags(warn=warn) + else: + yield from generic_tags() + + if interp_name == "pp": + interp = "pp3" + elif interp_name == "cp": + interp = "cp" + interpreter_version(warn=warn) + else: + interp = None + yield from compatible_tags(interpreter=interp) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/utils.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c2c2f75aa806282d322c76c2117c0f0fdfb09d25 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/utils.py @@ -0,0 +1,172 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. + +import re +from typing import FrozenSet, NewType, Tuple, Union, cast + +from .tags import Tag, parse_tag +from .version import InvalidVersion, Version + +BuildTag = Union[Tuple[()], Tuple[int, str]] +NormalizedName = NewType("NormalizedName", str) + + +class InvalidName(ValueError): + """ + An invalid distribution name; users should refer to the packaging user guide. + """ + + +class InvalidWheelFilename(ValueError): + """ + An invalid wheel filename was found, users should refer to PEP 427. + """ + + +class InvalidSdistFilename(ValueError): + """ + An invalid sdist filename was found, users should refer to the packaging user guide. + """ + + +# Core metadata spec for `Name` +_validate_regex = re.compile( + r"^([A-Z0-9]|[A-Z0-9][A-Z0-9._-]*[A-Z0-9])$", re.IGNORECASE +) +_canonicalize_regex = re.compile(r"[-_.]+") +_normalized_regex = re.compile(r"^([a-z0-9]|[a-z0-9]([a-z0-9-](?!--))*[a-z0-9])$") +# PEP 427: The build number must start with a digit. +_build_tag_regex = re.compile(r"(\d+)(.*)") + + +def canonicalize_name(name: str, *, validate: bool = False) -> NormalizedName: + if validate and not _validate_regex.match(name): + raise InvalidName(f"name is invalid: {name!r}") + # This is taken from PEP 503. + value = _canonicalize_regex.sub("-", name).lower() + return cast(NormalizedName, value) + + +def is_normalized_name(name: str) -> bool: + return _normalized_regex.match(name) is not None + + +def canonicalize_version( + version: Union[Version, str], *, strip_trailing_zero: bool = True +) -> str: + """ + This is very similar to Version.__str__, but has one subtle difference + with the way it handles the release segment. + """ + if isinstance(version, str): + try: + parsed = Version(version) + except InvalidVersion: + # Legacy versions cannot be normalized + return version + else: + parsed = version + + parts = [] + + # Epoch + if parsed.epoch != 0: + parts.append(f"{parsed.epoch}!") + + # Release segment + release_segment = ".".join(str(x) for x in parsed.release) + if strip_trailing_zero: + # NB: This strips trailing '.0's to normalize + release_segment = re.sub(r"(\.0)+$", "", release_segment) + parts.append(release_segment) + + # Pre-release + if parsed.pre is not None: + parts.append("".join(str(x) for x in parsed.pre)) + + # Post-release + if parsed.post is not None: + parts.append(f".post{parsed.post}") + + # Development release + if parsed.dev is not None: + parts.append(f".dev{parsed.dev}") + + # Local version segment + if parsed.local is not None: + parts.append(f"+{parsed.local}") + + return "".join(parts) + + +def parse_wheel_filename( + filename: str, +) -> Tuple[NormalizedName, Version, BuildTag, FrozenSet[Tag]]: + if not filename.endswith(".whl"): + raise InvalidWheelFilename( + f"Invalid wheel filename (extension must be '.whl'): {filename}" + ) + + filename = filename[:-4] + dashes = filename.count("-") + if dashes not in (4, 5): + raise InvalidWheelFilename( + f"Invalid wheel filename (wrong number of parts): {filename}" + ) + + parts = filename.split("-", dashes - 2) + name_part = parts[0] + # See PEP 427 for the rules on escaping the project name. + if "__" in name_part or re.match(r"^[\w\d._]*$", name_part, re.UNICODE) is None: + raise InvalidWheelFilename(f"Invalid project name: {filename}") + name = canonicalize_name(name_part) + + try: + version = Version(parts[1]) + except InvalidVersion as e: + raise InvalidWheelFilename( + f"Invalid wheel filename (invalid version): {filename}" + ) from e + + if dashes == 5: + build_part = parts[2] + build_match = _build_tag_regex.match(build_part) + if build_match is None: + raise InvalidWheelFilename( + f"Invalid build number: {build_part} in '{filename}'" + ) + build = cast(BuildTag, (int(build_match.group(1)), build_match.group(2))) + else: + build = () + tags = parse_tag(parts[-1]) + return (name, version, build, tags) + + +def parse_sdist_filename(filename: str) -> Tuple[NormalizedName, Version]: + if filename.endswith(".tar.gz"): + file_stem = filename[: -len(".tar.gz")] + elif filename.endswith(".zip"): + file_stem = filename[: -len(".zip")] + else: + raise InvalidSdistFilename( + f"Invalid sdist filename (extension must be '.tar.gz' or '.zip'):" + f" {filename}" + ) + + # We are requiring a PEP 440 version, which cannot contain dashes, + # so we split on the last dash. + name_part, sep, version_part = file_stem.rpartition("-") + if not sep: + raise InvalidSdistFilename(f"Invalid sdist filename: {filename}") + + name = canonicalize_name(name_part) + + try: + version = Version(version_part) + except InvalidVersion as e: + raise InvalidSdistFilename( + f"Invalid sdist filename (invalid version): {filename}" + ) from e + + return (name, version) diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/version.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/version.py new file mode 100644 index 0000000000000000000000000000000000000000..cda8e99935c8d92010b84437ae83d75031245d61 --- /dev/null +++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/packaging/version.py @@ -0,0 +1,561 @@ +# This file is dual licensed under the terms of the Apache License, Version +# 2.0, and the BSD License. See the LICENSE file in the root of this repository +# for complete details. +""" +.. testsetup:: + + from packaging.version import parse, Version +""" + +import itertools +import re +from typing import Any, Callable, NamedTuple, Optional, SupportsInt, Tuple, Union + +from ._structures import Infinity, InfinityType, NegativeInfinity, NegativeInfinityType + +__all__ = ["VERSION_PATTERN", "parse", "Version", "InvalidVersion"] + +LocalType = Tuple[Union[int, str], ...] + +CmpPrePostDevType = Union[InfinityType, NegativeInfinityType, Tuple[str, int]] +CmpLocalType = Union[ + NegativeInfinityType, + Tuple[Union[Tuple[int, str], Tuple[NegativeInfinityType, Union[int, str]]], ...], +] +CmpKey = Tuple[ + int, + Tuple[int, ...], + CmpPrePostDevType, + CmpPrePostDevType, + CmpPrePostDevType, + CmpLocalType, +] +VersionComparisonMethod = Callable[[CmpKey, CmpKey], bool] + + +class _Version(NamedTuple): + epoch: int + release: Tuple[int, ...] + dev: Optional[Tuple[str, int]] + pre: Optional[Tuple[str, int]] + post: Optional[Tuple[str, int]] + local: Optional[LocalType] + + +def parse(version: str) -> "Version": + """Parse the given version string. + + >>> parse('1.0.dev1') + + + :param version: The version string to parse. + :raises InvalidVersion: When the version string is not a valid version. + """ + return Version(version) + + +class InvalidVersion(ValueError): + """Raised when a version string is not a valid version. + + >>> Version("invalid") + Traceback (most recent call last): + ... + packaging.version.InvalidVersion: Invalid version: 'invalid' + """ + + +class _BaseVersion: + _key: Tuple[Any, ...] + + def __hash__(self) -> int: + return hash(self._key) + + # Please keep the duplicated `isinstance` check + # in the six comparisons hereunder + # unless you find a way to avoid adding overhead function calls. + def __lt__(self, other: "_BaseVersion") -> bool: + if not isinstance(other, _BaseVersion): + return NotImplemented + + return self._key < other._key + + def __le__(self, other: "_BaseVersion") -> bool: + if not isinstance(other, _BaseVersion): + return NotImplemented + + return self._key <= other._key + + def __eq__(self, other: object) -> bool: + if not isinstance(other, _BaseVersion): + return NotImplemented + + return self._key == other._key + + def __ge__(self, other: "_BaseVersion") -> bool: + if not isinstance(other, _BaseVersion): + return NotImplemented + + return self._key >= other._key + + def __gt__(self, other: "_BaseVersion") -> bool: + if not isinstance(other, _BaseVersion): + return NotImplemented + + return self._key > other._key + + def __ne__(self, other: object) -> bool: + if not isinstance(other, _BaseVersion): + return NotImplemented + + return self._key != other._key + + +# Deliberately not anchored to the start and end of the string, to make it +# easier for 3rd party code to reuse +_VERSION_PATTERN = r""" + v? + (?: + (?:(?P[0-9]+)!)? # epoch + (?P[0-9]+(?:\.[0-9]+)*) # release segment + (?P
                                          # pre-release
+            [-_\.]?
+            (?Palpha|a|beta|b|preview|pre|c|rc)
+            [-_\.]?
+            (?P[0-9]+)?
+        )?
+        (?P                                         # post release
+            (?:-(?P[0-9]+))
+            |
+            (?:
+                [-_\.]?
+                (?Ppost|rev|r)
+                [-_\.]?
+                (?P[0-9]+)?
+            )
+        )?
+        (?P                                          # dev release
+            [-_\.]?
+            (?Pdev)
+            [-_\.]?
+            (?P[0-9]+)?
+        )?
+    )
+    (?:\+(?P[a-z0-9]+(?:[-_\.][a-z0-9]+)*))?       # local version
+"""
+
+VERSION_PATTERN = _VERSION_PATTERN
+"""
+A string containing the regular expression used to match a valid version.
+
+The pattern is not anchored at either end, and is intended for embedding in larger
+expressions (for example, matching a version number as part of a file name). The
+regular expression should be compiled with the ``re.VERBOSE`` and ``re.IGNORECASE``
+flags set.
+
+:meta hide-value:
+"""
+
+
+class Version(_BaseVersion):
+    """This class abstracts handling of a project's versions.
+
+    A :class:`Version` instance is comparison aware and can be compared and
+    sorted using the standard Python interfaces.
+
+    >>> v1 = Version("1.0a5")
+    >>> v2 = Version("1.0")
+    >>> v1
+    
+    >>> v2
+    
+    >>> v1 < v2
+    True
+    >>> v1 == v2
+    False
+    >>> v1 > v2
+    False
+    >>> v1 >= v2
+    False
+    >>> v1 <= v2
+    True
+    """
+
+    _regex = re.compile(r"^\s*" + VERSION_PATTERN + r"\s*$", re.VERBOSE | re.IGNORECASE)
+    _key: CmpKey
+
+    def __init__(self, version: str) -> None:
+        """Initialize a Version object.
+
+        :param version:
+            The string representation of a version which will be parsed and normalized
+            before use.
+        :raises InvalidVersion:
+            If the ``version`` does not conform to PEP 440 in any way then this
+            exception will be raised.
+        """
+
+        # Validate the version and parse it into pieces
+        match = self._regex.search(version)
+        if not match:
+            raise InvalidVersion(f"Invalid version: '{version}'")
+
+        # Store the parsed out pieces of the version
+        self._version = _Version(
+            epoch=int(match.group("epoch")) if match.group("epoch") else 0,
+            release=tuple(int(i) for i in match.group("release").split(".")),
+            pre=_parse_letter_version(match.group("pre_l"), match.group("pre_n")),
+            post=_parse_letter_version(
+                match.group("post_l"), match.group("post_n1") or match.group("post_n2")
+            ),
+            dev=_parse_letter_version(match.group("dev_l"), match.group("dev_n")),
+            local=_parse_local_version(match.group("local")),
+        )
+
+        # Generate a key which will be used for sorting
+        self._key = _cmpkey(
+            self._version.epoch,
+            self._version.release,
+            self._version.pre,
+            self._version.post,
+            self._version.dev,
+            self._version.local,
+        )
+
+    def __repr__(self) -> str:
+        """A representation of the Version that shows all internal state.
+
+        >>> Version('1.0.0')
+        
+        """
+        return f""
+
+    def __str__(self) -> str:
+        """A string representation of the version that can be rounded-tripped.
+
+        >>> str(Version("1.0a5"))
+        '1.0a5'
+        """
+        parts = []
+
+        # Epoch
+        if self.epoch != 0:
+            parts.append(f"{self.epoch}!")
+
+        # Release segment
+        parts.append(".".join(str(x) for x in self.release))
+
+        # Pre-release
+        if self.pre is not None:
+            parts.append("".join(str(x) for x in self.pre))
+
+        # Post-release
+        if self.post is not None:
+            parts.append(f".post{self.post}")
+
+        # Development release
+        if self.dev is not None:
+            parts.append(f".dev{self.dev}")
+
+        # Local version segment
+        if self.local is not None:
+            parts.append(f"+{self.local}")
+
+        return "".join(parts)
+
+    @property
+    def epoch(self) -> int:
+        """The epoch of the version.
+
+        >>> Version("2.0.0").epoch
+        0
+        >>> Version("1!2.0.0").epoch
+        1
+        """
+        return self._version.epoch
+
+    @property
+    def release(self) -> Tuple[int, ...]:
+        """The components of the "release" segment of the version.
+
+        >>> Version("1.2.3").release
+        (1, 2, 3)
+        >>> Version("2.0.0").release
+        (2, 0, 0)
+        >>> Version("1!2.0.0.post0").release
+        (2, 0, 0)
+
+        Includes trailing zeroes but not the epoch or any pre-release / development /
+        post-release suffixes.
+        """
+        return self._version.release
+
+    @property
+    def pre(self) -> Optional[Tuple[str, int]]:
+        """The pre-release segment of the version.
+
+        >>> print(Version("1.2.3").pre)
+        None
+        >>> Version("1.2.3a1").pre
+        ('a', 1)
+        >>> Version("1.2.3b1").pre
+        ('b', 1)
+        >>> Version("1.2.3rc1").pre
+        ('rc', 1)
+        """
+        return self._version.pre
+
+    @property
+    def post(self) -> Optional[int]:
+        """The post-release number of the version.
+
+        >>> print(Version("1.2.3").post)
+        None
+        >>> Version("1.2.3.post1").post
+        1
+        """
+        return self._version.post[1] if self._version.post else None
+
+    @property
+    def dev(self) -> Optional[int]:
+        """The development number of the version.
+
+        >>> print(Version("1.2.3").dev)
+        None
+        >>> Version("1.2.3.dev1").dev
+        1
+        """
+        return self._version.dev[1] if self._version.dev else None
+
+    @property
+    def local(self) -> Optional[str]:
+        """The local version segment of the version.
+
+        >>> print(Version("1.2.3").local)
+        None
+        >>> Version("1.2.3+abc").local
+        'abc'
+        """
+        if self._version.local:
+            return ".".join(str(x) for x in self._version.local)
+        else:
+            return None
+
+    @property
+    def public(self) -> str:
+        """The public portion of the version.
+
+        >>> Version("1.2.3").public
+        '1.2.3'
+        >>> Version("1.2.3+abc").public
+        '1.2.3'
+        >>> Version("1.2.3+abc.dev1").public
+        '1.2.3'
+        """
+        return str(self).split("+", 1)[0]
+
+    @property
+    def base_version(self) -> str:
+        """The "base version" of the version.
+
+        >>> Version("1.2.3").base_version
+        '1.2.3'
+        >>> Version("1.2.3+abc").base_version
+        '1.2.3'
+        >>> Version("1!1.2.3+abc.dev1").base_version
+        '1!1.2.3'
+
+        The "base version" is the public version of the project without any pre or post
+        release markers.
+        """
+        parts = []
+
+        # Epoch
+        if self.epoch != 0:
+            parts.append(f"{self.epoch}!")
+
+        # Release segment
+        parts.append(".".join(str(x) for x in self.release))
+
+        return "".join(parts)
+
+    @property
+    def is_prerelease(self) -> bool:
+        """Whether this version is a pre-release.
+
+        >>> Version("1.2.3").is_prerelease
+        False
+        >>> Version("1.2.3a1").is_prerelease
+        True
+        >>> Version("1.2.3b1").is_prerelease
+        True
+        >>> Version("1.2.3rc1").is_prerelease
+        True
+        >>> Version("1.2.3dev1").is_prerelease
+        True
+        """
+        return self.dev is not None or self.pre is not None
+
+    @property
+    def is_postrelease(self) -> bool:
+        """Whether this version is a post-release.
+
+        >>> Version("1.2.3").is_postrelease
+        False
+        >>> Version("1.2.3.post1").is_postrelease
+        True
+        """
+        return self.post is not None
+
+    @property
+    def is_devrelease(self) -> bool:
+        """Whether this version is a development release.
+
+        >>> Version("1.2.3").is_devrelease
+        False
+        >>> Version("1.2.3.dev1").is_devrelease
+        True
+        """
+        return self.dev is not None
+
+    @property
+    def major(self) -> int:
+        """The first item of :attr:`release` or ``0`` if unavailable.
+
+        >>> Version("1.2.3").major
+        1
+        """
+        return self.release[0] if len(self.release) >= 1 else 0
+
+    @property
+    def minor(self) -> int:
+        """The second item of :attr:`release` or ``0`` if unavailable.
+
+        >>> Version("1.2.3").minor
+        2
+        >>> Version("1").minor
+        0
+        """
+        return self.release[1] if len(self.release) >= 2 else 0
+
+    @property
+    def micro(self) -> int:
+        """The third item of :attr:`release` or ``0`` if unavailable.
+
+        >>> Version("1.2.3").micro
+        3
+        >>> Version("1").micro
+        0
+        """
+        return self.release[2] if len(self.release) >= 3 else 0
+
+
+def _parse_letter_version(
+    letter: Optional[str], number: Union[str, bytes, SupportsInt, None]
+) -> Optional[Tuple[str, int]]:
+    if letter:
+        # We consider there to be an implicit 0 in a pre-release if there is
+        # not a numeral associated with it.
+        if number is None:
+            number = 0
+
+        # We normalize any letters to their lower case form
+        letter = letter.lower()
+
+        # We consider some words to be alternate spellings of other words and
+        # in those cases we want to normalize the spellings to our preferred
+        # spelling.
+        if letter == "alpha":
+            letter = "a"
+        elif letter == "beta":
+            letter = "b"
+        elif letter in ["c", "pre", "preview"]:
+            letter = "rc"
+        elif letter in ["rev", "r"]:
+            letter = "post"
+
+        return letter, int(number)
+    if not letter and number:
+        # We assume if we are given a number, but we are not given a letter
+        # then this is using the implicit post release syntax (e.g. 1.0-1)
+        letter = "post"
+
+        return letter, int(number)
+
+    return None
+
+
+_local_version_separators = re.compile(r"[\._-]")
+
+
+def _parse_local_version(local: Optional[str]) -> Optional[LocalType]:
+    """
+    Takes a string like abc.1.twelve and turns it into ("abc", 1, "twelve").
+    """
+    if local is not None:
+        return tuple(
+            part.lower() if not part.isdigit() else int(part)
+            for part in _local_version_separators.split(local)
+        )
+    return None
+
+
+def _cmpkey(
+    epoch: int,
+    release: Tuple[int, ...],
+    pre: Optional[Tuple[str, int]],
+    post: Optional[Tuple[str, int]],
+    dev: Optional[Tuple[str, int]],
+    local: Optional[LocalType],
+) -> CmpKey:
+    # When we compare a release version, we want to compare it with all of the
+    # trailing zeros removed. So we'll use a reverse the list, drop all the now
+    # leading zeros until we come to something non zero, then take the rest
+    # re-reverse it back into the correct order and make it a tuple and use
+    # that for our sorting key.
+    _release = tuple(
+        reversed(list(itertools.dropwhile(lambda x: x == 0, reversed(release))))
+    )
+
+    # We need to "trick" the sorting algorithm to put 1.0.dev0 before 1.0a0.
+    # We'll do this by abusing the pre segment, but we _only_ want to do this
+    # if there is not a pre or a post segment. If we have one of those then
+    # the normal sorting rules will handle this case correctly.
+    if pre is None and post is None and dev is not None:
+        _pre: CmpPrePostDevType = NegativeInfinity
+    # Versions without a pre-release (except as noted above) should sort after
+    # those with one.
+    elif pre is None:
+        _pre = Infinity
+    else:
+        _pre = pre
+
+    # Versions without a post segment should sort before those with one.
+    if post is None:
+        _post: CmpPrePostDevType = NegativeInfinity
+
+    else:
+        _post = post
+
+    # Versions without a development segment should sort after those with one.
+    if dev is None:
+        _dev: CmpPrePostDevType = Infinity
+
+    else:
+        _dev = dev
+
+    if local is None:
+        # Versions without a local segment should sort before those with one.
+        _local: CmpLocalType = NegativeInfinity
+    else:
+        # Versions with a local segment need that segment parsed to implement
+        # the sorting rules in PEP440.
+        # - Alpha numeric segments sort before numeric segments
+        # - Alpha numeric segments sort lexicographically
+        # - Numeric segments sort numerically
+        # - Shorter versions sort before longer versions when the prefixes
+        #   match exactly
+        _local = tuple(
+            (i, "") if isinstance(i, int) else (NegativeInfinity, i) for i in local
+        )
+
+    return epoch, _release, _pre, _post, _dev, _local
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/vendor.txt b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/vendor.txt
new file mode 100644
index 0000000000000000000000000000000000000000..14666103a82ea7fced3d8cffb8a9b2a9e03fb492
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/_vendor/wheel/vendored/vendor.txt
@@ -0,0 +1 @@
+packaging==24.0
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/build_meta.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/build_meta.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f2e930c73439933b1e15506699ce25ec926c46f
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/build_meta.py
@@ -0,0 +1,548 @@
+"""A PEP 517 interface to setuptools
+
+Previously, when a user or a command line tool (let's call it a "frontend")
+needed to make a request of setuptools to take a certain action, for
+example, generating a list of installation requirements, the frontend
+would call "setup.py egg_info" or "setup.py bdist_wheel" on the command line.
+
+PEP 517 defines a different method of interfacing with setuptools. Rather
+than calling "setup.py" directly, the frontend should:
+
+  1. Set the current directory to the directory with a setup.py file
+  2. Import this module into a safe python interpreter (one in which
+     setuptools can potentially set global variables or crash hard).
+  3. Call one of the functions defined in PEP 517.
+
+What each function does is defined in PEP 517. However, here is a "casual"
+definition of the functions (this definition should not be relied on for
+bug reports or API stability):
+
+  - `build_wheel`: build a wheel in the folder and return the basename
+  - `get_requires_for_build_wheel`: get the `setup_requires` to build
+  - `prepare_metadata_for_build_wheel`: get the `install_requires`
+  - `build_sdist`: build an sdist in the folder and return the basename
+  - `get_requires_for_build_sdist`: get the `setup_requires` to build
+
+Again, this is not a formal definition! Just a "taste" of the module.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import os
+import shlex
+import shutil
+import sys
+import tempfile
+import tokenize
+import warnings
+from collections.abc import Iterable, Iterator, Mapping
+from pathlib import Path
+from typing import TYPE_CHECKING, Union
+
+import setuptools
+
+from . import errors
+from ._path import StrPath, same_path
+from ._reqs import parse_strings
+from .warnings import SetuptoolsDeprecationWarning
+
+import distutils
+from distutils.util import strtobool
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeAlias
+
+__all__ = [
+    'get_requires_for_build_sdist',
+    'get_requires_for_build_wheel',
+    'prepare_metadata_for_build_wheel',
+    'build_wheel',
+    'build_sdist',
+    'get_requires_for_build_editable',
+    'prepare_metadata_for_build_editable',
+    'build_editable',
+    '__legacy__',
+    'SetupRequirementsError',
+]
+
+
+class SetupRequirementsError(BaseException):
+    def __init__(self, specifiers) -> None:
+        self.specifiers = specifiers
+
+
+class Distribution(setuptools.dist.Distribution):
+    def fetch_build_eggs(self, specifiers):
+        specifier_list = list(parse_strings(specifiers))
+
+        raise SetupRequirementsError(specifier_list)
+
+    @classmethod
+    @contextlib.contextmanager
+    def patch(cls):
+        """
+        Replace
+        distutils.dist.Distribution with this class
+        for the duration of this context.
+        """
+        orig = distutils.core.Distribution
+        distutils.core.Distribution = cls  # type: ignore[misc] # monkeypatching
+        try:
+            yield
+        finally:
+            distutils.core.Distribution = orig  # type: ignore[misc] # monkeypatching
+
+
+@contextlib.contextmanager
+def no_install_setup_requires():
+    """Temporarily disable installing setup_requires
+
+    Under PEP 517, the backend reports build dependencies to the frontend,
+    and the frontend is responsible for ensuring they're installed.
+    So setuptools (acting as a backend) should not try to install them.
+    """
+    orig = setuptools._install_setup_requires
+    setuptools._install_setup_requires = lambda attrs: None
+    try:
+        yield
+    finally:
+        setuptools._install_setup_requires = orig
+
+
+def _get_immediate_subdirectories(a_dir):
+    return [
+        name for name in os.listdir(a_dir) if os.path.isdir(os.path.join(a_dir, name))
+    ]
+
+
+def _file_with_extension(directory: StrPath, extension: str | tuple[str, ...]):
+    matching = (f for f in os.listdir(directory) if f.endswith(extension))
+    try:
+        (file,) = matching
+    except ValueError:
+        raise ValueError(
+            'No distribution was found. Ensure that `setup.py` '
+            'is not empty and that it calls `setup()`.'
+        ) from None
+    return file
+
+
+def _open_setup_script(setup_script):
+    if not os.path.exists(setup_script):
+        # Supply a default setup.py
+        return io.StringIO("from setuptools import setup; setup()")
+
+    return tokenize.open(setup_script)
+
+
+@contextlib.contextmanager
+def suppress_known_deprecation():
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', 'setup.py install is deprecated')
+        yield
+
+
+_ConfigSettings: TypeAlias = Union[Mapping[str, Union[str, list[str], None]], None]
+"""
+Currently the user can run::
+
+    pip install -e . --config-settings key=value
+    python -m build -C--key=value -C key=value
+
+- pip will pass both key and value as strings and overwriting repeated keys
+  (pypa/pip#11059).
+- build will accumulate values associated with repeated keys in a list.
+  It will also accept keys with no associated value.
+  This means that an option passed by build can be ``str | list[str] | None``.
+- PEP 517 specifies that ``config_settings`` is an optional dict.
+"""
+
+
+class _ConfigSettingsTranslator:
+    """Translate ``config_settings`` into distutils-style command arguments.
+    Only a limited number of options is currently supported.
+    """
+
+    # See pypa/setuptools#1928 pypa/setuptools#2491
+
+    def _get_config(self, key: str, config_settings: _ConfigSettings) -> list[str]:
+        """
+        Get the value of a specific key in ``config_settings`` as a list of strings.
+
+        >>> fn = _ConfigSettingsTranslator()._get_config
+        >>> fn("--global-option", None)
+        []
+        >>> fn("--global-option", {})
+        []
+        >>> fn("--global-option", {'--global-option': 'foo'})
+        ['foo']
+        >>> fn("--global-option", {'--global-option': ['foo']})
+        ['foo']
+        >>> fn("--global-option", {'--global-option': 'foo'})
+        ['foo']
+        >>> fn("--global-option", {'--global-option': 'foo bar'})
+        ['foo', 'bar']
+        """
+        cfg = config_settings or {}
+        opts = cfg.get(key) or []
+        return shlex.split(opts) if isinstance(opts, str) else opts
+
+    def _global_args(self, config_settings: _ConfigSettings) -> Iterator[str]:
+        """
+        Let the user specify ``verbose`` or ``quiet`` + escape hatch via
+        ``--global-option``.
+        Note: ``-v``, ``-vv``, ``-vvv`` have similar effects in setuptools,
+        so we just have to cover the basic scenario ``-v``.
+
+        >>> fn = _ConfigSettingsTranslator()._global_args
+        >>> list(fn(None))
+        []
+        >>> list(fn({"verbose": "False"}))
+        ['-q']
+        >>> list(fn({"verbose": "1"}))
+        ['-v']
+        >>> list(fn({"--verbose": None}))
+        ['-v']
+        >>> list(fn({"verbose": "true", "--global-option": "-q --no-user-cfg"}))
+        ['-v', '-q', '--no-user-cfg']
+        >>> list(fn({"--quiet": None}))
+        ['-q']
+        """
+        cfg = config_settings or {}
+        falsey = {"false", "no", "0", "off"}
+        if "verbose" in cfg or "--verbose" in cfg:
+            level = str(cfg.get("verbose") or cfg.get("--verbose") or "1")
+            yield ("-q" if level.lower() in falsey else "-v")
+        if "quiet" in cfg or "--quiet" in cfg:
+            level = str(cfg.get("quiet") or cfg.get("--quiet") or "1")
+            yield ("-v" if level.lower() in falsey else "-q")
+
+        yield from self._get_config("--global-option", config_settings)
+
+    def __dist_info_args(self, config_settings: _ConfigSettings) -> Iterator[str]:
+        """
+        The ``dist_info`` command accepts ``tag-date`` and ``tag-build``.
+
+        .. warning::
+           We cannot use this yet as it requires the ``sdist`` and ``bdist_wheel``
+           commands run in ``build_sdist`` and ``build_wheel`` to reuse the egg-info
+           directory created in ``prepare_metadata_for_build_wheel``.
+
+        >>> fn = _ConfigSettingsTranslator()._ConfigSettingsTranslator__dist_info_args
+        >>> list(fn(None))
+        []
+        >>> list(fn({"tag-date": "False"}))
+        ['--no-date']
+        >>> list(fn({"tag-date": None}))
+        ['--no-date']
+        >>> list(fn({"tag-date": "true", "tag-build": ".a"}))
+        ['--tag-date', '--tag-build', '.a']
+        """
+        cfg = config_settings or {}
+        if "tag-date" in cfg:
+            val = strtobool(str(cfg["tag-date"] or "false"))
+            yield ("--tag-date" if val else "--no-date")
+        if "tag-build" in cfg:
+            yield from ["--tag-build", str(cfg["tag-build"])]
+
+    def _editable_args(self, config_settings: _ConfigSettings) -> Iterator[str]:
+        """
+        The ``editable_wheel`` command accepts ``editable-mode=strict``.
+
+        >>> fn = _ConfigSettingsTranslator()._editable_args
+        >>> list(fn(None))
+        []
+        >>> list(fn({"editable-mode": "strict"}))
+        ['--mode', 'strict']
+        """
+        cfg = config_settings or {}
+        mode = cfg.get("editable-mode") or cfg.get("editable_mode")
+        if not mode:
+            return
+        yield from ["--mode", str(mode)]
+
+    def _arbitrary_args(self, config_settings: _ConfigSettings) -> Iterator[str]:
+        """
+        Users may expect to pass arbitrary lists of arguments to a command
+        via "--global-option" (example provided in PEP 517 of a "escape hatch").
+
+        >>> fn = _ConfigSettingsTranslator()._arbitrary_args
+        >>> list(fn(None))
+        []
+        >>> list(fn({}))
+        []
+        >>> list(fn({'--build-option': 'foo'}))
+        ['foo']
+        >>> list(fn({'--build-option': ['foo']}))
+        ['foo']
+        >>> list(fn({'--build-option': 'foo'}))
+        ['foo']
+        >>> list(fn({'--build-option': 'foo bar'}))
+        ['foo', 'bar']
+        >>> list(fn({'--global-option': 'foo'}))
+        []
+        """
+        yield from self._get_config("--build-option", config_settings)
+
+
+class _BuildMetaBackend(_ConfigSettingsTranslator):
+    def _get_build_requires(
+        self, config_settings: _ConfigSettings, requirements: list[str]
+    ):
+        sys.argv = [
+            *sys.argv[:1],
+            *self._global_args(config_settings),
+            "egg_info",
+        ]
+        try:
+            with Distribution.patch():
+                self.run_setup()
+        except SetupRequirementsError as e:
+            requirements += e.specifiers
+
+        return requirements
+
+    def run_setup(self, setup_script: str = 'setup.py'):
+        # Note that we can reuse our build directory between calls
+        # Correctness comes first, then optimization later
+        __file__ = os.path.abspath(setup_script)
+        __name__ = '__main__'
+
+        with _open_setup_script(__file__) as f:
+            code = f.read().replace(r'\r\n', r'\n')
+
+        try:
+            exec(code, locals())
+        except SystemExit as e:
+            if e.code:
+                raise
+            # We ignore exit code indicating success
+            SetuptoolsDeprecationWarning.emit(
+                "Running `setup.py` directly as CLI tool is deprecated.",
+                "Please avoid using `sys.exit(0)` or similar statements "
+                "that don't fit in the paradigm of a configuration file.",
+                see_url="https://blog.ganssle.io/articles/2021/10/"
+                "setup-py-deprecated.html",
+            )
+
+    def get_requires_for_build_wheel(self, config_settings: _ConfigSettings = None):
+        return self._get_build_requires(config_settings, requirements=[])
+
+    def get_requires_for_build_sdist(self, config_settings: _ConfigSettings = None):
+        return self._get_build_requires(config_settings, requirements=[])
+
+    def _bubble_up_info_directory(
+        self, metadata_directory: StrPath, suffix: str
+    ) -> str:
+        """
+        PEP 517 requires that the .dist-info directory be placed in the
+        metadata_directory. To comply, we MUST copy the directory to the root.
+
+        Returns the basename of the info directory, e.g. `proj-0.0.0.dist-info`.
+        """
+        info_dir = self._find_info_directory(metadata_directory, suffix)
+        if not same_path(info_dir.parent, metadata_directory):
+            shutil.move(str(info_dir), metadata_directory)
+            # PEP 517 allow other files and dirs to exist in metadata_directory
+        return info_dir.name
+
+    def _find_info_directory(self, metadata_directory: StrPath, suffix: str) -> Path:
+        for parent, dirs, _ in os.walk(metadata_directory):
+            candidates = [f for f in dirs if f.endswith(suffix)]
+
+            if len(candidates) != 0 or len(dirs) != 1:
+                assert len(candidates) == 1, f"Multiple {suffix} directories found"
+                return Path(parent, candidates[0])
+
+        msg = f"No {suffix} directory found in {metadata_directory}"
+        raise errors.InternalError(msg)
+
+    def prepare_metadata_for_build_wheel(
+        self, metadata_directory: StrPath, config_settings: _ConfigSettings = None
+    ):
+        sys.argv = [
+            *sys.argv[:1],
+            *self._global_args(config_settings),
+            "dist_info",
+            "--output-dir",
+            str(metadata_directory),
+            "--keep-egg-info",
+        ]
+        with no_install_setup_requires():
+            self.run_setup()
+
+        self._bubble_up_info_directory(metadata_directory, ".egg-info")
+        return self._bubble_up_info_directory(metadata_directory, ".dist-info")
+
+    def _build_with_temp_dir(
+        self,
+        setup_command: Iterable[str],
+        result_extension: str | tuple[str, ...],
+        result_directory: StrPath,
+        config_settings: _ConfigSettings,
+        arbitrary_args: Iterable[str] = (),
+    ):
+        result_directory = os.path.abspath(result_directory)
+
+        # Build in a temporary directory, then copy to the target.
+        os.makedirs(result_directory, exist_ok=True)
+
+        with tempfile.TemporaryDirectory(
+            prefix=".tmp-", dir=result_directory
+        ) as tmp_dist_dir:
+            sys.argv = [
+                *sys.argv[:1],
+                *self._global_args(config_settings),
+                *setup_command,
+                "--dist-dir",
+                tmp_dist_dir,
+                *arbitrary_args,
+            ]
+            with no_install_setup_requires():
+                self.run_setup()
+
+            result_basename = _file_with_extension(tmp_dist_dir, result_extension)
+            result_path = os.path.join(result_directory, result_basename)
+            if os.path.exists(result_path):
+                # os.rename will fail overwriting on non-Unix.
+                os.remove(result_path)
+            os.rename(os.path.join(tmp_dist_dir, result_basename), result_path)
+
+        return result_basename
+
+    def build_wheel(
+        self,
+        wheel_directory: StrPath,
+        config_settings: _ConfigSettings = None,
+        metadata_directory: StrPath | None = None,
+    ):
+        def _build(cmd: list[str]):
+            with suppress_known_deprecation():
+                return self._build_with_temp_dir(
+                    cmd,
+                    '.whl',
+                    wheel_directory,
+                    config_settings,
+                    self._arbitrary_args(config_settings),
+                )
+
+        if metadata_directory is None:
+            return _build(['bdist_wheel'])
+
+        try:
+            return _build(['bdist_wheel', '--dist-info-dir', str(metadata_directory)])
+        except SystemExit as ex:  # pragma: nocover
+            # pypa/setuptools#4683
+            if "--dist-info-dir not recognized" not in str(ex):
+                raise
+            _IncompatibleBdistWheel.emit()
+            return _build(['bdist_wheel'])
+
+    def build_sdist(
+        self, sdist_directory: StrPath, config_settings: _ConfigSettings = None
+    ):
+        return self._build_with_temp_dir(
+            ['sdist', '--formats', 'gztar'], '.tar.gz', sdist_directory, config_settings
+        )
+
+    def _get_dist_info_dir(self, metadata_directory: StrPath | None) -> str | None:
+        if not metadata_directory:
+            return None
+        dist_info_candidates = list(Path(metadata_directory).glob("*.dist-info"))
+        assert len(dist_info_candidates) <= 1
+        return str(dist_info_candidates[0]) if dist_info_candidates else None
+
+    def build_editable(
+        self,
+        wheel_directory: StrPath,
+        config_settings: _ConfigSettings = None,
+        metadata_directory: StrPath | None = None,
+    ):
+        # XXX can or should we hide our editable_wheel command normally?
+        info_dir = self._get_dist_info_dir(metadata_directory)
+        opts = ["--dist-info-dir", info_dir] if info_dir else []
+        cmd = ["editable_wheel", *opts, *self._editable_args(config_settings)]
+        with suppress_known_deprecation():
+            return self._build_with_temp_dir(
+                cmd, ".whl", wheel_directory, config_settings
+            )
+
+    def get_requires_for_build_editable(self, config_settings: _ConfigSettings = None):
+        return self.get_requires_for_build_wheel(config_settings)
+
+    def prepare_metadata_for_build_editable(
+        self, metadata_directory: StrPath, config_settings: _ConfigSettings = None
+    ):
+        return self.prepare_metadata_for_build_wheel(
+            metadata_directory, config_settings
+        )
+
+
+class _BuildMetaLegacyBackend(_BuildMetaBackend):
+    """Compatibility backend for setuptools
+
+    This is a version of setuptools.build_meta that endeavors
+    to maintain backwards
+    compatibility with pre-PEP 517 modes of invocation. It
+    exists as a temporary
+    bridge between the old packaging mechanism and the new
+    packaging mechanism,
+    and will eventually be removed.
+    """
+
+    def run_setup(self, setup_script: str = 'setup.py'):
+        # In order to maintain compatibility with scripts assuming that
+        # the setup.py script is in a directory on the PYTHONPATH, inject
+        # '' into sys.path. (pypa/setuptools#1642)
+        sys_path = list(sys.path)  # Save the original path
+
+        script_dir = os.path.dirname(os.path.abspath(setup_script))
+        if script_dir not in sys.path:
+            sys.path.insert(0, script_dir)
+
+        # Some setup.py scripts (e.g. in pygame and numpy) use sys.argv[0] to
+        # get the directory of the source code. They expect it to refer to the
+        # setup.py script.
+        sys_argv_0 = sys.argv[0]
+        sys.argv[0] = setup_script
+
+        try:
+            super().run_setup(setup_script=setup_script)
+        finally:
+            # While PEP 517 frontends should be calling each hook in a fresh
+            # subprocess according to the standard (and thus it should not be
+            # strictly necessary to restore the old sys.path), we'll restore
+            # the original path so that the path manipulation does not persist
+            # within the hook after run_setup is called.
+            sys.path[:] = sys_path
+            sys.argv[0] = sys_argv_0
+
+
+class _IncompatibleBdistWheel(SetuptoolsDeprecationWarning):
+    _SUMMARY = "wheel.bdist_wheel is deprecated, please import it from setuptools"
+    _DETAILS = """
+    Ensure that any custom bdist_wheel implementation is a subclass of
+    setuptools.command.bdist_wheel.bdist_wheel.
+    """
+    _DUE_DATE = (2025, 10, 15)
+    # Initially introduced in 2024/10/15, but maybe too disruptive to be enforced?
+    _SEE_URL = "https://github.com/pypa/wheel/pull/631"
+
+
+# The primary backend
+_BACKEND = _BuildMetaBackend()
+
+get_requires_for_build_wheel = _BACKEND.get_requires_for_build_wheel
+get_requires_for_build_sdist = _BACKEND.get_requires_for_build_sdist
+prepare_metadata_for_build_wheel = _BACKEND.prepare_metadata_for_build_wheel
+build_wheel = _BACKEND.build_wheel
+build_sdist = _BACKEND.build_sdist
+get_requires_for_build_editable = _BACKEND.get_requires_for_build_editable
+prepare_metadata_for_build_editable = _BACKEND.prepare_metadata_for_build_editable
+build_editable = _BACKEND.build_editable
+
+
+# The legacy backend
+__legacy__ = _BuildMetaLegacyBackend()
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-32.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-32.exe
new file mode 100644
index 0000000000000000000000000000000000000000..65c3cd99cc7433f271a5b9387abdd1ddb949d1a6
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-32.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-64.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-64.exe
new file mode 100644
index 0000000000000000000000000000000000000000..3ea50eebfe3f0113b231a318cc1ad6e238afd60d
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-64.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-arm64.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-arm64.exe
new file mode 100644
index 0000000000000000000000000000000000000000..da96455a07a0bad4cde5dc5626544325f82c722b
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli-arm64.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli.exe
new file mode 100644
index 0000000000000000000000000000000000000000..65c3cd99cc7433f271a5b9387abdd1ddb949d1a6
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/cli.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/depends.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/depends.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5223b79561c36d9b6c45ead78288098e1cb0f1d
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/depends.py
@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import contextlib
+import dis
+import marshal
+import sys
+from types import CodeType
+from typing import Any, Literal, TypeVar
+
+from packaging.version import Version
+
+from . import _imp
+from ._imp import PY_COMPILED, PY_FROZEN, PY_SOURCE, find_module
+
+_T = TypeVar("_T")
+
+__all__ = ['Require', 'find_module']
+
+
+class Require:
+    """A prerequisite to building or installing a distribution"""
+
+    def __init__(
+        self,
+        name,
+        requested_version,
+        module,
+        homepage: str = '',
+        attribute=None,
+        format=None,
+    ) -> None:
+        if format is None and requested_version is not None:
+            format = Version
+
+        if format is not None:
+            requested_version = format(requested_version)
+            if attribute is None:
+                attribute = '__version__'
+
+        self.__dict__.update(locals())
+        del self.self
+
+    def full_name(self):
+        """Return full package/distribution name, w/version"""
+        if self.requested_version is not None:
+            return f'{self.name}-{self.requested_version}'
+        return self.name
+
+    def version_ok(self, version):
+        """Is 'version' sufficiently up-to-date?"""
+        return (
+            self.attribute is None
+            or self.format is None
+            or str(version) != "unknown"
+            and self.format(version) >= self.requested_version
+        )
+
+    def get_version(
+        self, paths=None, default: _T | Literal["unknown"] = "unknown"
+    ) -> _T | Literal["unknown"] | None | Any:
+        """Get version number of installed module, 'None', or 'default'
+
+        Search 'paths' for module.  If not found, return 'None'.  If found,
+        return the extracted version attribute, or 'default' if no version
+        attribute was specified, or the value cannot be determined without
+        importing the module.  The version is formatted according to the
+        requirement's version format (if any), unless it is 'None' or the
+        supplied 'default'.
+        """
+
+        if self.attribute is None:
+            try:
+                f, _p, _i = find_module(self.module, paths)
+            except ImportError:
+                return None
+            if f:
+                f.close()
+            return default
+
+        v = get_module_constant(self.module, self.attribute, default, paths)
+
+        if v is not None and v is not default and self.format is not None:
+            return self.format(v)
+
+        return v
+
+    def is_present(self, paths=None):
+        """Return true if dependency is present on 'paths'"""
+        return self.get_version(paths) is not None
+
+    def is_current(self, paths=None):
+        """Return true if dependency is present and up-to-date on 'paths'"""
+        version = self.get_version(paths)
+        if version is None:
+            return False
+        return self.version_ok(str(version))
+
+
+def maybe_close(f):
+    @contextlib.contextmanager
+    def empty():
+        yield
+        return
+
+    if not f:
+        return empty()
+
+    return contextlib.closing(f)
+
+
+# Some objects are not available on some platforms.
+# XXX it'd be better to test assertions about bytecode instead.
+if not sys.platform.startswith('java') and sys.platform != 'cli':
+
+    def get_module_constant(
+        module, symbol, default: _T | int = -1, paths=None
+    ) -> _T | int | None | Any:
+        """Find 'module' by searching 'paths', and extract 'symbol'
+
+        Return 'None' if 'module' does not exist on 'paths', or it does not define
+        'symbol'.  If the module defines 'symbol' as a constant, return the
+        constant.  Otherwise, return 'default'."""
+
+        try:
+            f, path, (_suffix, _mode, kind) = info = find_module(module, paths)
+        except ImportError:
+            # Module doesn't exist
+            return None
+
+        with maybe_close(f):
+            if kind == PY_COMPILED:
+                f.read(8)  # skip magic & date
+                code = marshal.load(f)
+            elif kind == PY_FROZEN:
+                code = _imp.get_frozen_object(module, paths)
+            elif kind == PY_SOURCE:
+                code = compile(f.read(), path, 'exec')
+            else:
+                # Not something we can parse; we'll have to import it.  :(
+                imported = _imp.get_module(module, paths, info)
+                return getattr(imported, symbol, None)
+
+        return extract_constant(code, symbol, default)
+
+    def extract_constant(
+        code: CodeType, symbol: str, default: _T | int = -1
+    ) -> _T | int | None | Any:
+        """Extract the constant value of 'symbol' from 'code'
+
+        If the name 'symbol' is bound to a constant value by the Python code
+        object 'code', return that value.  If 'symbol' is bound to an expression,
+        return 'default'.  Otherwise, return 'None'.
+
+        Return value is based on the first assignment to 'symbol'.  'symbol' must
+        be a global, or at least a non-"fast" local in the code block.  That is,
+        only 'STORE_NAME' and 'STORE_GLOBAL' opcodes are checked, and 'symbol'
+        must be present in 'code.co_names'.
+        """
+        if symbol not in code.co_names:
+            # name's not there, can't possibly be an assignment
+            return None
+
+        name_idx = list(code.co_names).index(symbol)
+
+        STORE_NAME = dis.opmap['STORE_NAME']
+        STORE_GLOBAL = dis.opmap['STORE_GLOBAL']
+        LOAD_CONST = dis.opmap['LOAD_CONST']
+
+        const = default
+
+        for byte_code in dis.Bytecode(code):
+            op = byte_code.opcode
+            arg = byte_code.arg
+
+            if op == LOAD_CONST:
+                assert arg is not None
+                const = code.co_consts[arg]
+            elif arg == name_idx and (op == STORE_NAME or op == STORE_GLOBAL):
+                return const
+            else:
+                const = default
+
+        return None
+
+    __all__ += ['get_module_constant', 'extract_constant']
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/discovery.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/discovery.py
new file mode 100644
index 0000000000000000000000000000000000000000..c88839918562bad12f1a2e72309f1eacfe23349c
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/discovery.py
@@ -0,0 +1,614 @@
+"""Automatic discovery of Python modules and packages (for inclusion in the
+distribution) and other config values.
+
+For the purposes of this module, the following nomenclature is used:
+
+- "src-layout": a directory representing a Python project that contains a "src"
+  folder. Everything under the "src" folder is meant to be included in the
+  distribution when packaging the project. Example::
+
+    .
+    ├── tox.ini
+    ├── pyproject.toml
+    └── src/
+        └── mypkg/
+            ├── __init__.py
+            ├── mymodule.py
+            └── my_data_file.txt
+
+- "flat-layout": a Python project that does not use "src-layout" but instead
+  have a directory under the project root for each package::
+
+    .
+    ├── tox.ini
+    ├── pyproject.toml
+    └── mypkg/
+        ├── __init__.py
+        ├── mymodule.py
+        └── my_data_file.txt
+
+- "single-module": a project that contains a single Python script direct under
+  the project root (no directory used)::
+
+    .
+    ├── tox.ini
+    ├── pyproject.toml
+    └── mymodule.py
+
+"""
+
+from __future__ import annotations
+
+import itertools
+import os
+from collections.abc import Iterable, Iterator, Mapping
+from fnmatch import fnmatchcase
+from glob import glob
+from pathlib import Path
+from typing import TYPE_CHECKING, ClassVar
+
+import _distutils_hack.override  # noqa: F401
+
+from ._path import StrPath
+
+from distutils import log
+from distutils.util import convert_path
+
+if TYPE_CHECKING:
+    from setuptools import Distribution
+
+chain_iter = itertools.chain.from_iterable
+
+
+def _valid_name(path: StrPath) -> bool:
+    # Ignore invalid names that cannot be imported directly
+    return os.path.basename(path).isidentifier()
+
+
+class _Filter:
+    """
+    Given a list of patterns, create a callable that will be true only if
+    the input matches at least one of the patterns.
+    """
+
+    def __init__(self, *patterns: str) -> None:
+        self._patterns = dict.fromkeys(patterns)
+
+    def __call__(self, item: str) -> bool:
+        return any(fnmatchcase(item, pat) for pat in self._patterns)
+
+    def __contains__(self, item: str) -> bool:
+        return item in self._patterns
+
+
+class _Finder:
+    """Base class that exposes functionality for module/package finders"""
+
+    ALWAYS_EXCLUDE: ClassVar[tuple[str, ...]] = ()
+    DEFAULT_EXCLUDE: ClassVar[tuple[str, ...]] = ()
+
+    @classmethod
+    def find(
+        cls,
+        where: StrPath = '.',
+        exclude: Iterable[str] = (),
+        include: Iterable[str] = ('*',),
+    ) -> list[str]:
+        """Return a list of all Python items (packages or modules, depending on
+        the finder implementation) found within directory 'where'.
+
+        'where' is the root directory which will be searched.
+        It should be supplied as a "cross-platform" (i.e. URL-style) path;
+        it will be converted to the appropriate local path syntax.
+
+        'exclude' is a sequence of names to exclude; '*' can be used
+        as a wildcard in the names.
+        When finding packages, 'foo.*' will exclude all subpackages of 'foo'
+        (but not 'foo' itself).
+
+        'include' is a sequence of names to include.
+        If it's specified, only the named items will be included.
+        If it's not specified, all found items will be included.
+        'include' can contain shell style wildcard patterns just like
+        'exclude'.
+        """
+
+        exclude = exclude or cls.DEFAULT_EXCLUDE
+        return list(
+            cls._find_iter(
+                convert_path(str(where)),
+                _Filter(*cls.ALWAYS_EXCLUDE, *exclude),
+                _Filter(*include),
+            )
+        )
+
+    @classmethod
+    def _find_iter(
+        cls, where: StrPath, exclude: _Filter, include: _Filter
+    ) -> Iterator[str]:
+        raise NotImplementedError
+
+
+class PackageFinder(_Finder):
+    """
+    Generate a list of all Python packages found within a directory
+    """
+
+    ALWAYS_EXCLUDE = ("ez_setup", "*__pycache__")
+
+    @classmethod
+    def _find_iter(
+        cls, where: StrPath, exclude: _Filter, include: _Filter
+    ) -> Iterator[str]:
+        """
+        All the packages found in 'where' that pass the 'include' filter, but
+        not the 'exclude' filter.
+        """
+        for root, dirs, files in os.walk(str(where), followlinks=True):
+            # Copy dirs to iterate over it, then empty dirs.
+            all_dirs = dirs[:]
+            dirs[:] = []
+
+            for dir in all_dirs:
+                full_path = os.path.join(root, dir)
+                rel_path = os.path.relpath(full_path, where)
+                package = rel_path.replace(os.path.sep, '.')
+
+                # Skip directory trees that are not valid packages
+                if '.' in dir or not cls._looks_like_package(full_path, package):
+                    continue
+
+                # Should this package be included?
+                if include(package) and not exclude(package):
+                    yield package
+
+                # Early pruning if there is nothing else to be scanned
+                if f"{package}*" in exclude or f"{package}.*" in exclude:
+                    continue
+
+                # Keep searching subdirectories, as there may be more packages
+                # down there, even if the parent was excluded.
+                dirs.append(dir)
+
+    @staticmethod
+    def _looks_like_package(path: StrPath, _package_name: str) -> bool:
+        """Does a directory look like a package?"""
+        return os.path.isfile(os.path.join(path, '__init__.py'))
+
+
+class PEP420PackageFinder(PackageFinder):
+    @staticmethod
+    def _looks_like_package(_path: StrPath, _package_name: str) -> bool:
+        return True
+
+
+class ModuleFinder(_Finder):
+    """Find isolated Python modules.
+    This function will **not** recurse subdirectories.
+    """
+
+    @classmethod
+    def _find_iter(
+        cls, where: StrPath, exclude: _Filter, include: _Filter
+    ) -> Iterator[str]:
+        for file in glob(os.path.join(where, "*.py")):
+            module, _ext = os.path.splitext(os.path.basename(file))
+
+            if not cls._looks_like_module(module):
+                continue
+
+            if include(module) and not exclude(module):
+                yield module
+
+    _looks_like_module = staticmethod(_valid_name)
+
+
+# We have to be extra careful in the case of flat layout to not include files
+# and directories not meant for distribution (e.g. tool-related)
+
+
+class FlatLayoutPackageFinder(PEP420PackageFinder):
+    _EXCLUDE = (
+        "ci",
+        "bin",
+        "debian",
+        "doc",
+        "docs",
+        "documentation",
+        "manpages",
+        "news",
+        "newsfragments",
+        "changelog",
+        "test",
+        "tests",
+        "unit_test",
+        "unit_tests",
+        "example",
+        "examples",
+        "scripts",
+        "tools",
+        "util",
+        "utils",
+        "python",
+        "build",
+        "dist",
+        "venv",
+        "env",
+        "requirements",
+        # ---- Task runners / Build tools ----
+        "tasks",  # invoke
+        "fabfile",  # fabric
+        "site_scons",  # SCons
+        # ---- Other tools ----
+        "benchmark",
+        "benchmarks",
+        "exercise",
+        "exercises",
+        "htmlcov",  # Coverage.py
+        # ---- Hidden directories/Private packages ----
+        "[._]*",
+    )
+
+    DEFAULT_EXCLUDE = tuple(chain_iter((p, f"{p}.*") for p in _EXCLUDE))
+    """Reserved package names"""
+
+    @staticmethod
+    def _looks_like_package(_path: StrPath, package_name: str) -> bool:
+        names = package_name.split('.')
+        # Consider PEP 561
+        root_pkg_is_valid = names[0].isidentifier() or names[0].endswith("-stubs")
+        return root_pkg_is_valid and all(name.isidentifier() for name in names[1:])
+
+
+class FlatLayoutModuleFinder(ModuleFinder):
+    DEFAULT_EXCLUDE = (
+        "setup",
+        "conftest",
+        "test",
+        "tests",
+        "example",
+        "examples",
+        "build",
+        # ---- Task runners ----
+        "toxfile",
+        "noxfile",
+        "pavement",
+        "dodo",
+        "tasks",
+        "fabfile",
+        # ---- Other tools ----
+        "[Ss][Cc]onstruct",  # SCons
+        "conanfile",  # Connan: C/C++ build tool
+        "manage",  # Django
+        "benchmark",
+        "benchmarks",
+        "exercise",
+        "exercises",
+        # ---- Hidden files/Private modules ----
+        "[._]*",
+    )
+    """Reserved top-level module names"""
+
+
+def _find_packages_within(root_pkg: str, pkg_dir: StrPath) -> list[str]:
+    nested = PEP420PackageFinder.find(pkg_dir)
+    return [root_pkg] + [".".join((root_pkg, n)) for n in nested]
+
+
+class ConfigDiscovery:
+    """Fill-in metadata and options that can be automatically derived
+    (from other metadata/options, the file system or conventions)
+    """
+
+    def __init__(self, distribution: Distribution) -> None:
+        self.dist = distribution
+        self._called = False
+        self._disabled = False
+        self._skip_ext_modules = False
+
+    def _disable(self):
+        """Internal API to disable automatic discovery"""
+        self._disabled = True
+
+    def _ignore_ext_modules(self):
+        """Internal API to disregard ext_modules.
+
+        Normally auto-discovery would not be triggered if ``ext_modules`` are set
+        (this is done for backward compatibility with existing packages relying on
+        ``setup.py`` or ``setup.cfg``). However, ``setuptools`` can call this function
+        to ignore given ``ext_modules`` and proceed with the auto-discovery if
+        ``packages`` and ``py_modules`` are not given (e.g. when using pyproject.toml
+        metadata).
+        """
+        self._skip_ext_modules = True
+
+    @property
+    def _root_dir(self) -> StrPath:
+        # The best is to wait until `src_root` is set in dist, before using _root_dir.
+        return self.dist.src_root or os.curdir
+
+    @property
+    def _package_dir(self) -> dict[str, str]:
+        if self.dist.package_dir is None:
+            return {}
+        return self.dist.package_dir
+
+    def __call__(
+        self, force: bool = False, name: bool = True, ignore_ext_modules: bool = False
+    ):
+        """Automatically discover missing configuration fields
+        and modifies the given ``distribution`` object in-place.
+
+        Note that by default this will only have an effect the first time the
+        ``ConfigDiscovery`` object is called.
+
+        To repeatedly invoke automatic discovery (e.g. when the project
+        directory changes), please use ``force=True`` (or create a new
+        ``ConfigDiscovery`` instance).
+        """
+        if force is False and (self._called or self._disabled):
+            # Avoid overhead of multiple calls
+            return
+
+        self._analyse_package_layout(ignore_ext_modules)
+        if name:
+            self.analyse_name()  # depends on ``packages`` and ``py_modules``
+
+        self._called = True
+
+    def _explicitly_specified(self, ignore_ext_modules: bool) -> bool:
+        """``True`` if the user has specified some form of package/module listing"""
+        ignore_ext_modules = ignore_ext_modules or self._skip_ext_modules
+        ext_modules = not (self.dist.ext_modules is None or ignore_ext_modules)
+        return (
+            self.dist.packages is not None
+            or self.dist.py_modules is not None
+            or ext_modules
+            or hasattr(self.dist, "configuration")
+            and self.dist.configuration
+            # ^ Some projects use numpy.distutils.misc_util.Configuration
+        )
+
+    def _analyse_package_layout(self, ignore_ext_modules: bool) -> bool:
+        if self._explicitly_specified(ignore_ext_modules):
+            # For backward compatibility, just try to find modules/packages
+            # when nothing is given
+            return True
+
+        log.debug(
+            "No `packages` or `py_modules` configuration, performing "
+            "automatic discovery."
+        )
+
+        return (
+            self._analyse_explicit_layout()
+            or self._analyse_src_layout()
+            # flat-layout is the trickiest for discovery so it should be last
+            or self._analyse_flat_layout()
+        )
+
+    def _analyse_explicit_layout(self) -> bool:
+        """The user can explicitly give a package layout via ``package_dir``"""
+        package_dir = self._package_dir.copy()  # don't modify directly
+        package_dir.pop("", None)  # This falls under the "src-layout" umbrella
+        root_dir = self._root_dir
+
+        if not package_dir:
+            return False
+
+        log.debug(f"`explicit-layout` detected -- analysing {package_dir}")
+        pkgs = chain_iter(
+            _find_packages_within(pkg, os.path.join(root_dir, parent_dir))
+            for pkg, parent_dir in package_dir.items()
+        )
+        self.dist.packages = list(pkgs)
+        log.debug(f"discovered packages -- {self.dist.packages}")
+        return True
+
+    def _analyse_src_layout(self) -> bool:
+        """Try to find all packages or modules under the ``src`` directory
+        (or anything pointed by ``package_dir[""]``).
+
+        The "src-layout" is relatively safe for automatic discovery.
+        We assume that everything within is meant to be included in the
+        distribution.
+
+        If ``package_dir[""]`` is not given, but the ``src`` directory exists,
+        this function will set ``package_dir[""] = "src"``.
+        """
+        package_dir = self._package_dir
+        src_dir = os.path.join(self._root_dir, package_dir.get("", "src"))
+        if not os.path.isdir(src_dir):
+            return False
+
+        log.debug(f"`src-layout` detected -- analysing {src_dir}")
+        package_dir.setdefault("", os.path.basename(src_dir))
+        self.dist.package_dir = package_dir  # persist eventual modifications
+        self.dist.packages = PEP420PackageFinder.find(src_dir)
+        self.dist.py_modules = ModuleFinder.find(src_dir)
+        log.debug(f"discovered packages -- {self.dist.packages}")
+        log.debug(f"discovered py_modules -- {self.dist.py_modules}")
+        return True
+
+    def _analyse_flat_layout(self) -> bool:
+        """Try to find all packages and modules under the project root.
+
+        Since the ``flat-layout`` is more dangerous in terms of accidentally including
+        extra files/directories, this function is more conservative and will raise an
+        error if multiple packages or modules are found.
+
+        This assumes that multi-package dists are uncommon and refuse to support that
+        use case in order to be able to prevent unintended errors.
+        """
+        log.debug(f"`flat-layout` detected -- analysing {self._root_dir}")
+        return self._analyse_flat_packages() or self._analyse_flat_modules()
+
+    def _analyse_flat_packages(self) -> bool:
+        self.dist.packages = FlatLayoutPackageFinder.find(self._root_dir)
+        top_level = remove_nested_packages(remove_stubs(self.dist.packages))
+        log.debug(f"discovered packages -- {self.dist.packages}")
+        self._ensure_no_accidental_inclusion(top_level, "packages")
+        return bool(top_level)
+
+    def _analyse_flat_modules(self) -> bool:
+        self.dist.py_modules = FlatLayoutModuleFinder.find(self._root_dir)
+        log.debug(f"discovered py_modules -- {self.dist.py_modules}")
+        self._ensure_no_accidental_inclusion(self.dist.py_modules, "modules")
+        return bool(self.dist.py_modules)
+
+    def _ensure_no_accidental_inclusion(self, detected: list[str], kind: str):
+        if len(detected) > 1:
+            from inspect import cleandoc
+
+            from setuptools.errors import PackageDiscoveryError
+
+            msg = f"""Multiple top-level {kind} discovered in a flat-layout: {detected}.
+
+            To avoid accidental inclusion of unwanted files or directories,
+            setuptools will not proceed with this build.
+
+            If you are trying to create a single distribution with multiple {kind}
+            on purpose, you should not rely on automatic discovery.
+            Instead, consider the following options:
+
+            1. set up custom discovery (`find` directive with `include` or `exclude`)
+            2. use a `src-layout`
+            3. explicitly set `py_modules` or `packages` with a list of names
+
+            To find more information, look for "package discovery" on setuptools docs.
+            """
+            raise PackageDiscoveryError(cleandoc(msg))
+
+    def analyse_name(self) -> None:
+        """The packages/modules are the essential contribution of the author.
+        Therefore the name of the distribution can be derived from them.
+        """
+        if self.dist.metadata.name or self.dist.name:
+            # get_name() is not reliable (can return "UNKNOWN")
+            return
+
+        log.debug("No `name` configuration, performing automatic discovery")
+
+        name = (
+            self._find_name_single_package_or_module()
+            or self._find_name_from_packages()
+        )
+        if name:
+            self.dist.metadata.name = name
+
+    def _find_name_single_package_or_module(self) -> str | None:
+        """Exactly one module or package"""
+        for field in ('packages', 'py_modules'):
+            items = getattr(self.dist, field, None) or []
+            if items and len(items) == 1:
+                log.debug(f"Single module/package detected, name: {items[0]}")
+                return items[0]
+
+        return None
+
+    def _find_name_from_packages(self) -> str | None:
+        """Try to find the root package that is not a PEP 420 namespace"""
+        if not self.dist.packages:
+            return None
+
+        packages = remove_stubs(sorted(self.dist.packages, key=len))
+        package_dir = self.dist.package_dir or {}
+
+        parent_pkg = find_parent_package(packages, package_dir, self._root_dir)
+        if parent_pkg:
+            log.debug(f"Common parent package detected, name: {parent_pkg}")
+            return parent_pkg
+
+        log.warn("No parent package detected, impossible to derive `name`")
+        return None
+
+
+def remove_nested_packages(packages: list[str]) -> list[str]:
+    """Remove nested packages from a list of packages.
+
+    >>> remove_nested_packages(["a", "a.b1", "a.b2", "a.b1.c1"])
+    ['a']
+    >>> remove_nested_packages(["a", "b", "c.d", "c.d.e.f", "g.h", "a.a1"])
+    ['a', 'b', 'c.d', 'g.h']
+    """
+    pkgs = sorted(packages, key=len)
+    top_level = pkgs[:]
+    size = len(pkgs)
+    for i, name in enumerate(reversed(pkgs)):
+        if any(name.startswith(f"{other}.") for other in top_level):
+            top_level.pop(size - i - 1)
+
+    return top_level
+
+
+def remove_stubs(packages: list[str]) -> list[str]:
+    """Remove type stubs (:pep:`561`) from a list of packages.
+
+    >>> remove_stubs(["a", "a.b", "a-stubs", "a-stubs.b.c", "b", "c-stubs"])
+    ['a', 'a.b', 'b']
+    """
+    return [pkg for pkg in packages if not pkg.split(".")[0].endswith("-stubs")]
+
+
+def find_parent_package(
+    packages: list[str], package_dir: Mapping[str, str], root_dir: StrPath
+) -> str | None:
+    """Find the parent package that is not a namespace."""
+    packages = sorted(packages, key=len)
+    common_ancestors = []
+    for i, name in enumerate(packages):
+        if not all(n.startswith(f"{name}.") for n in packages[i + 1 :]):
+            # Since packages are sorted by length, this condition is able
+            # to find a list of all common ancestors.
+            # When there is divergence (e.g. multiple root packages)
+            # the list will be empty
+            break
+        common_ancestors.append(name)
+
+    for name in common_ancestors:
+        pkg_path = find_package_path(name, package_dir, root_dir)
+        init = os.path.join(pkg_path, "__init__.py")
+        if os.path.isfile(init):
+            return name
+
+    return None
+
+
+def find_package_path(
+    name: str, package_dir: Mapping[str, str], root_dir: StrPath
+) -> str:
+    """Given a package name, return the path where it should be found on
+    disk, considering the ``package_dir`` option.
+
+    >>> path = find_package_path("my.pkg", {"": "root/is/nested"}, ".")
+    >>> path.replace(os.sep, "/")
+    './root/is/nested/my/pkg'
+
+    >>> path = find_package_path("my.pkg", {"my": "root/is/nested"}, ".")
+    >>> path.replace(os.sep, "/")
+    './root/is/nested/pkg'
+
+    >>> path = find_package_path("my.pkg", {"my.pkg": "root/is/nested"}, ".")
+    >>> path.replace(os.sep, "/")
+    './root/is/nested'
+
+    >>> path = find_package_path("other.pkg", {"my.pkg": "root/is/nested"}, ".")
+    >>> path.replace(os.sep, "/")
+    './other/pkg'
+    """
+    parts = name.split(".")
+    for i in range(len(parts), 0, -1):
+        # Look backwards, the most specific package_dir first
+        partial_name = ".".join(parts[:i])
+        if partial_name in package_dir:
+            parent = package_dir[partial_name]
+            return os.path.join(root_dir, parent, *parts[i:])
+
+    parent = package_dir.get("") or ""
+    return os.path.join(root_dir, *parent.split("/"), *parts)
+
+
+def construct_package_dir(packages: list[str], package_path: StrPath) -> dict[str, str]:
+    parent_pkgs = remove_nested_packages(packages)
+    prefix = Path(package_path).parts
+    return {pkg: "/".join([*prefix, *pkg.split(".")]) for pkg in parent_pkgs}
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/dist.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/dist.py
new file mode 100644
index 0000000000000000000000000000000000000000..f06298c8681a637176c7def931f06ddbdb9854dc
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/dist.py
@@ -0,0 +1,1119 @@
+from __future__ import annotations
+
+import functools
+import io
+import itertools
+import numbers
+import os
+import re
+import sys
+from collections.abc import Iterable, Iterator, MutableMapping, Sequence
+from glob import glob
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Union
+
+from more_itertools import partition, unique_everseen
+from packaging.markers import InvalidMarker, Marker
+from packaging.specifiers import InvalidSpecifier, SpecifierSet
+from packaging.version import Version
+
+from . import (
+    _entry_points,
+    _reqs,
+    _static,
+    command as _,  # noqa: F401 # imported for side-effects
+)
+from ._importlib import metadata
+from ._normalization import _canonicalize_license_expression
+from ._path import StrPath
+from ._reqs import _StrOrIter
+from .config import pyprojecttoml, setupcfg
+from .discovery import ConfigDiscovery
+from .errors import InvalidConfigError
+from .monkey import get_unpatched
+from .warnings import InformationOnly, SetuptoolsDeprecationWarning
+
+import distutils.cmd
+import distutils.command
+import distutils.core
+import distutils.dist
+import distutils.log
+from distutils.debug import DEBUG
+from distutils.errors import DistutilsOptionError, DistutilsSetupError
+from distutils.fancy_getopt import translate_longopt
+from distutils.util import strtobool
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeAlias
+
+
+__all__ = ['Distribution']
+
+_sequence = tuple, list
+"""
+:meta private:
+
+Supported iterable types that are known to be:
+- ordered (which `set` isn't)
+- not match a str (which `Sequence[str]` does)
+- not imply a nested type (like `dict`)
+for use with `isinstance`.
+"""
+_Sequence: TypeAlias = Union[tuple[str, ...], list[str]]
+# This is how stringifying _Sequence would look in Python 3.10
+_sequence_type_repr = "tuple[str, ...] | list[str]"
+_OrderedStrSequence: TypeAlias = Union[str, dict[str, Any], Sequence[str]]
+"""
+:meta private:
+Avoid single-use iterable. Disallow sets.
+A poor approximation of an OrderedSequence (dict doesn't match a Sequence).
+"""
+
+
+def __getattr__(name: str) -> Any:  # pragma: no cover
+    if name == "sequence":
+        SetuptoolsDeprecationWarning.emit(
+            "`setuptools.dist.sequence` is an internal implementation detail.",
+            "Please define your own `sequence = tuple, list` instead.",
+            due_date=(2025, 8, 28),  # Originally added on 2024-08-27
+        )
+        return _sequence
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def check_importable(dist, attr, value):
+    try:
+        ep = metadata.EntryPoint(value=value, name=None, group=None)
+        assert not ep.extras
+    except (TypeError, ValueError, AttributeError, AssertionError) as e:
+        raise DistutilsSetupError(
+            f"{attr!r} must be importable 'module:attrs' string (got {value!r})"
+        ) from e
+
+
+def assert_string_list(dist, attr: str, value: _Sequence) -> None:
+    """Verify that value is a string list"""
+    try:
+        # verify that value is a list or tuple to exclude unordered
+        # or single-use iterables
+        assert isinstance(value, _sequence)
+        # verify that elements of value are strings
+        assert ''.join(value) != value
+    except (TypeError, ValueError, AttributeError, AssertionError) as e:
+        raise DistutilsSetupError(
+            f"{attr!r} must be of type <{_sequence_type_repr}> (got {value!r})"
+        ) from e
+
+
+def check_nsp(dist, attr, value):
+    """Verify that namespace packages are valid"""
+    ns_packages = value
+    assert_string_list(dist, attr, ns_packages)
+    for nsp in ns_packages:
+        if not dist.has_contents_for(nsp):
+            raise DistutilsSetupError(
+                f"Distribution contains no modules or packages for namespace package {nsp!r}"
+            )
+        parent, _sep, _child = nsp.rpartition('.')
+        if parent and parent not in ns_packages:
+            distutils.log.warn(
+                "WARNING: %r is declared as a package namespace, but %r"
+                " is not: please correct this in setup.py",
+                nsp,
+                parent,
+            )
+        SetuptoolsDeprecationWarning.emit(
+            "The namespace_packages parameter is deprecated.",
+            "Please replace its usage with implicit namespaces (PEP 420).",
+            see_docs="references/keywords.html#keyword-namespace-packages",
+            # TODO: define due_date, it may break old packages that are no longer
+            # maintained (e.g. sphinxcontrib extensions) when installed from source.
+            # Warning officially introduced in May 2022, however the deprecation
+            # was mentioned much earlier in the docs (May 2020, see #2149).
+        )
+
+
+def check_extras(dist, attr, value):
+    """Verify that extras_require mapping is valid"""
+    try:
+        list(itertools.starmap(_check_extra, value.items()))
+    except (TypeError, ValueError, AttributeError) as e:
+        raise DistutilsSetupError(
+            "'extras_require' must be a dictionary whose values are "
+            "strings or lists of strings containing valid project/version "
+            "requirement specifiers."
+        ) from e
+
+
+def _check_extra(extra, reqs):
+    _name, _sep, marker = extra.partition(':')
+    try:
+        _check_marker(marker)
+    except InvalidMarker:
+        msg = f"Invalid environment marker: {marker} ({extra!r})"
+        raise DistutilsSetupError(msg) from None
+    list(_reqs.parse(reqs))
+
+
+def _check_marker(marker):
+    if not marker:
+        return
+    m = Marker(marker)
+    m.evaluate()
+
+
+def assert_bool(dist, attr, value):
+    """Verify that value is True, False, 0, or 1"""
+    if bool(value) != value:
+        raise DistutilsSetupError(f"{attr!r} must be a boolean value (got {value!r})")
+
+
+def invalid_unless_false(dist, attr, value):
+    if not value:
+        DistDeprecationWarning.emit(f"{attr} is ignored.")
+        # TODO: should there be a `due_date` here?
+        return
+    raise DistutilsSetupError(f"{attr} is invalid.")
+
+
+def check_requirements(dist, attr: str, value: _OrderedStrSequence) -> None:
+    """Verify that install_requires is a valid requirements list"""
+    try:
+        list(_reqs.parse(value))
+        if isinstance(value, set):
+            raise TypeError("Unordered types are not allowed")
+    except (TypeError, ValueError) as error:
+        msg = (
+            f"{attr!r} must be a string or iterable of strings "
+            f"containing valid project/version requirement specifiers; {error}"
+        )
+        raise DistutilsSetupError(msg) from error
+
+
+def check_specifier(dist, attr, value):
+    """Verify that value is a valid version specifier"""
+    try:
+        SpecifierSet(value)
+    except (InvalidSpecifier, AttributeError) as error:
+        msg = f"{attr!r} must be a string containing valid version specifiers; {error}"
+        raise DistutilsSetupError(msg) from error
+
+
+def check_entry_points(dist, attr, value):
+    """Verify that entry_points map is parseable"""
+    try:
+        _entry_points.load(value)
+    except Exception as e:
+        raise DistutilsSetupError(e) from e
+
+
+def check_package_data(dist, attr, value):
+    """Verify that value is a dictionary of package names to glob lists"""
+    if not isinstance(value, dict):
+        raise DistutilsSetupError(
+            f"{attr!r} must be a dictionary mapping package names to lists of "
+            "string wildcard patterns"
+        )
+    for k, v in value.items():
+        if not isinstance(k, str):
+            raise DistutilsSetupError(
+                f"keys of {attr!r} dict must be strings (got {k!r})"
+            )
+        assert_string_list(dist, f'values of {attr!r} dict', v)
+
+
+def check_packages(dist, attr, value):
+    for pkgname in value:
+        if not re.match(r'\w+(\.\w+)*', pkgname):
+            distutils.log.warn(
+                "WARNING: %r not a valid package name; please use only "
+                ".-separated package names in setup.py",
+                pkgname,
+            )
+
+
+if TYPE_CHECKING:
+    # Work around a mypy issue where type[T] can't be used as a base: https://github.com/python/mypy/issues/10962
+    from distutils.core import Distribution as _Distribution
+else:
+    _Distribution = get_unpatched(distutils.core.Distribution)
+
+
+class Distribution(_Distribution):
+    """Distribution with support for tests and package data
+
+    This is an enhanced version of 'distutils.dist.Distribution' that
+    effectively adds the following new optional keyword arguments to 'setup()':
+
+     'install_requires' -- a string or sequence of strings specifying project
+        versions that the distribution requires when installed, in the format
+        used by 'pkg_resources.require()'.  They will be installed
+        automatically when the package is installed.  If you wish to use
+        packages that are not available in PyPI, or want to give your users an
+        alternate download location, you can add a 'find_links' option to the
+        '[easy_install]' section of your project's 'setup.cfg' file, and then
+        setuptools will scan the listed web pages for links that satisfy the
+        requirements.
+
+     'extras_require' -- a dictionary mapping names of optional "extras" to the
+        additional requirement(s) that using those extras incurs. For example,
+        this::
+
+            extras_require = dict(reST = ["docutils>=0.3", "reSTedit"])
+
+        indicates that the distribution can optionally provide an extra
+        capability called "reST", but it can only be used if docutils and
+        reSTedit are installed.  If the user installs your package using
+        EasyInstall and requests one of your extras, the corresponding
+        additional requirements will be installed if needed.
+
+     'package_data' -- a dictionary mapping package names to lists of filenames
+        or globs to use to find data files contained in the named packages.
+        If the dictionary has filenames or globs listed under '""' (the empty
+        string), those names will be searched for in every package, in addition
+        to any names for the specific package.  Data files found using these
+        names/globs will be installed along with the package, in the same
+        location as the package.  Note that globs are allowed to reference
+        the contents of non-package subdirectories, as long as you use '/' as
+        a path separator.  (Globs are automatically converted to
+        platform-specific paths at runtime.)
+
+    In addition to these new keywords, this class also has several new methods
+    for manipulating the distribution's contents.  For example, the 'include()'
+    and 'exclude()' methods can be thought of as in-place add and subtract
+    commands that add or remove packages, modules, extensions, and so on from
+    the distribution.
+    """
+
+    _DISTUTILS_UNSUPPORTED_METADATA = {
+        'long_description_content_type': lambda: None,
+        'project_urls': dict,
+        'provides_extras': dict,  # behaves like an ordered set
+        'license_expression': lambda: None,
+        'license_file': lambda: None,
+        'license_files': lambda: None,
+        'install_requires': list,
+        'extras_require': dict,
+    }
+
+    # Used by build_py, editable_wheel and install_lib commands for legacy namespaces
+    namespace_packages: list[str]  #: :meta private: DEPRECATED
+
+    # Any: Dynamic assignment results in Incompatible types in assignment
+    def __init__(self, attrs: MutableMapping[str, Any] | None = None) -> None:
+        have_package_data = hasattr(self, "package_data")
+        if not have_package_data:
+            self.package_data: dict[str, list[str]] = {}
+        attrs = attrs or {}
+        self.dist_files: list[tuple[str, str, str]] = []
+        self.include_package_data: bool | None = None
+        self.exclude_package_data: dict[str, list[str]] | None = None
+        # Filter-out setuptools' specific options.
+        self.src_root: str | None = attrs.pop("src_root", None)
+        self.dependency_links: list[str] = attrs.pop('dependency_links', [])
+        self.setup_requires: list[str] = attrs.pop('setup_requires', [])
+        for ep in metadata.entry_points(group='distutils.setup_keywords'):
+            vars(self).setdefault(ep.name, None)
+
+        metadata_only = set(self._DISTUTILS_UNSUPPORTED_METADATA)
+        metadata_only -= {"install_requires", "extras_require"}
+        dist_attrs = {k: v for k, v in attrs.items() if k not in metadata_only}
+        _Distribution.__init__(self, dist_attrs)
+
+        # Private API (setuptools-use only, not restricted to Distribution)
+        # Stores files that are referenced by the configuration and need to be in the
+        # sdist (e.g. `version = file: VERSION.txt`)
+        self._referenced_files = set[str]()
+
+        self.set_defaults = ConfigDiscovery(self)
+
+        self._set_metadata_defaults(attrs)
+
+        self.metadata.version = self._normalize_version(self.metadata.version)
+        self._finalize_requires()
+
+    def _validate_metadata(self):
+        required = {"name"}
+        provided = {
+            key
+            for key in vars(self.metadata)
+            if getattr(self.metadata, key, None) is not None
+        }
+        missing = required - provided
+
+        if missing:
+            msg = f"Required package metadata is missing: {missing}"
+            raise DistutilsSetupError(msg)
+
+    def _set_metadata_defaults(self, attrs):
+        """
+        Fill-in missing metadata fields not supported by distutils.
+        Some fields may have been set by other tools (e.g. pbr).
+        Those fields (vars(self.metadata)) take precedence to
+        supplied attrs.
+        """
+        for option, default in self._DISTUTILS_UNSUPPORTED_METADATA.items():
+            vars(self.metadata).setdefault(option, attrs.get(option, default()))
+
+    @staticmethod
+    def _normalize_version(version):
+        from . import sic
+
+        if isinstance(version, numbers.Number):
+            # Some people apparently take "version number" too literally :)
+            version = str(version)
+        elif isinstance(version, sic) or version is None:
+            return version
+
+        normalized = str(Version(version))
+        if version != normalized:
+            InformationOnly.emit(f"Normalizing '{version}' to '{normalized}'")
+            return normalized
+        return version
+
+    def _finalize_requires(self):
+        """
+        Set `metadata.python_requires` and fix environment markers
+        in `install_requires` and `extras_require`.
+        """
+        if getattr(self, 'python_requires', None):
+            self.metadata.python_requires = self.python_requires
+
+        self._normalize_requires()
+        self.metadata.install_requires = self.install_requires
+        self.metadata.extras_require = self.extras_require
+
+        if self.extras_require:
+            for extra in self.extras_require.keys():
+                # Setuptools allows a weird ": syntax for extras
+                extra = extra.split(':')[0]
+                if extra:
+                    self.metadata.provides_extras.setdefault(extra)
+
+    def _normalize_requires(self):
+        """Make sure requirement-related attributes exist and are normalized"""
+        install_requires = getattr(self, "install_requires", None) or []
+        extras_require = getattr(self, "extras_require", None) or {}
+
+        # Preserve the "static"-ness of values parsed from config files
+        list_ = _static.List if _static.is_static(install_requires) else list
+        self.install_requires = list_(map(str, _reqs.parse(install_requires)))
+
+        dict_ = _static.Dict if _static.is_static(extras_require) else dict
+        self.extras_require = dict_(
+            (k, list(map(str, _reqs.parse(v or [])))) for k, v in extras_require.items()
+        )
+
+    def _finalize_license_expression(self) -> None:
+        """
+        Normalize license and license_expression.
+        >>> dist = Distribution({"license_expression": _static.Str("mit aNd  gpl-3.0-OR-later")})
+        >>> _static.is_static(dist.metadata.license_expression)
+        True
+        >>> dist._finalize_license_expression()
+        >>> _static.is_static(dist.metadata.license_expression)  # preserve "static-ness"
+        True
+        >>> print(dist.metadata.license_expression)
+        MIT AND GPL-3.0-or-later
+        """
+        classifiers = self.metadata.get_classifiers()
+        license_classifiers = [cl for cl in classifiers if cl.startswith("License :: ")]
+
+        license_expr = self.metadata.license_expression
+        if license_expr:
+            str_ = _static.Str if _static.is_static(license_expr) else str
+            normalized = str_(_canonicalize_license_expression(license_expr))
+            if license_expr != normalized:
+                InformationOnly.emit(f"Normalizing '{license_expr}' to '{normalized}'")
+                self.metadata.license_expression = normalized
+            if license_classifiers:
+                raise InvalidConfigError(
+                    "License classifiers have been superseded by license expressions "
+                    "(see https://peps.python.org/pep-0639/). Please remove:\n\n"
+                    + "\n".join(license_classifiers),
+                )
+        elif license_classifiers:
+            pypa_guides = "guides/writing-pyproject-toml/#license"
+            SetuptoolsDeprecationWarning.emit(
+                "License classifiers are deprecated.",
+                "Please consider removing the following classifiers in favor of a "
+                "SPDX license expression:\n\n" + "\n".join(license_classifiers),
+                see_url=f"https://packaging.python.org/en/latest/{pypa_guides}",
+                # Warning introduced on 2025-02-17
+                # TODO: Should we add a due date? It may affect old/unmaintained
+                #       packages in the ecosystem and cause problems...
+            )
+
+    def _finalize_license_files(self) -> None:
+        """Compute names of all license files which should be included."""
+        license_files: list[str] | None = self.metadata.license_files
+        patterns = license_files or []
+
+        license_file: str | None = self.metadata.license_file
+        if license_file and license_file not in patterns:
+            patterns.append(license_file)
+
+        if license_files is None and license_file is None:
+            # Default patterns match the ones wheel uses
+            # See https://wheel.readthedocs.io/en/stable/user_guide.html
+            # -> 'Including license files in the generated wheel file'
+            patterns = ['LICEN[CS]E*', 'COPYING*', 'NOTICE*', 'AUTHORS*']
+            files = self._expand_patterns(patterns, enforce_match=False)
+        else:  # Patterns explicitly given by the user
+            files = self._expand_patterns(patterns, enforce_match=True)
+
+        self.metadata.license_files = list(unique_everseen(files))
+
+    @classmethod
+    def _expand_patterns(
+        cls, patterns: list[str], enforce_match: bool = True
+    ) -> Iterator[str]:
+        """
+        >>> getfixture('sample_project_cwd')
+        >>> list(Distribution._expand_patterns(['LICENSE.txt']))
+        ['LICENSE.txt']
+        >>> list(Distribution._expand_patterns(['pyproject.toml', 'LIC*']))
+        ['pyproject.toml', 'LICENSE.txt']
+        >>> list(Distribution._expand_patterns(['src/**/*.dat']))
+        ['src/sample/package_data.dat']
+        """
+        return (
+            path.replace(os.sep, "/")
+            for pattern in patterns
+            for path in sorted(cls._find_pattern(pattern, enforce_match))
+            if not path.endswith('~') and os.path.isfile(path)
+        )
+
+    @staticmethod
+    def _find_pattern(pattern: str, enforce_match: bool = True) -> list[str]:
+        r"""
+        >>> getfixture('sample_project_cwd')
+        >>> Distribution._find_pattern("LICENSE.txt")
+        ['LICENSE.txt']
+        >>> Distribution._find_pattern("/LICENSE.MIT")
+        Traceback (most recent call last):
+        ...
+        setuptools.errors.InvalidConfigError: Pattern '/LICENSE.MIT' should be relative...
+        >>> Distribution._find_pattern("../LICENSE.MIT")
+        Traceback (most recent call last):
+        ...
+        setuptools.warnings.SetuptoolsDeprecationWarning: ...Pattern '../LICENSE.MIT' cannot contain '..'...
+        >>> Distribution._find_pattern("LICEN{CSE*")
+        Traceback (most recent call last):
+        ...
+        setuptools.warnings.SetuptoolsDeprecationWarning: ...Pattern 'LICEN{CSE*' contains invalid characters...
+        """
+        pypa_guides = "specifications/glob-patterns/"
+        if ".." in pattern:
+            SetuptoolsDeprecationWarning.emit(
+                f"Pattern {pattern!r} cannot contain '..'",
+                """
+                Please ensure the files specified are contained by the root
+                of the Python package (normally marked by `pyproject.toml`).
+                """,
+                see_url=f"https://packaging.python.org/en/latest/{pypa_guides}",
+                due_date=(2026, 3, 20),  # Introduced in 2025-03-20
+                # Replace with InvalidConfigError after deprecation
+            )
+        if pattern.startswith((os.sep, "/")) or ":\\" in pattern:
+            raise InvalidConfigError(
+                f"Pattern {pattern!r} should be relative and must not start with '/'"
+            )
+        if re.match(r'^[\w\-\.\/\*\?\[\]]+$', pattern) is None:
+            SetuptoolsDeprecationWarning.emit(
+                "Please provide a valid glob pattern.",
+                "Pattern {pattern!r} contains invalid characters.",
+                pattern=pattern,
+                see_url=f"https://packaging.python.org/en/latest/{pypa_guides}",
+                due_date=(2026, 3, 20),  # Introduced in 2025-02-20
+            )
+
+        found = glob(pattern, recursive=True)
+
+        if enforce_match and not found:
+            SetuptoolsDeprecationWarning.emit(
+                "Cannot find any files for the given pattern.",
+                "Pattern {pattern!r} did not match any files.",
+                pattern=pattern,
+                due_date=(2026, 3, 20),  # Introduced in 2025-02-20
+                # PEP 639 requires us to error, but as a transition period
+                # we will only issue a warning to give people time to prepare.
+                # After the transition, this should raise an InvalidConfigError.
+            )
+        return found
+
+    # FIXME: 'Distribution._parse_config_files' is too complex (14)
+    def _parse_config_files(self, filenames=None):  # noqa: C901
+        """
+        Adapted from distutils.dist.Distribution.parse_config_files,
+        this method provides the same functionality in subtly-improved
+        ways.
+        """
+        from configparser import ConfigParser
+
+        # Ignore install directory options if we have a venv
+        ignore_options = (
+            []
+            if sys.prefix == sys.base_prefix
+            else [
+                'install-base',
+                'install-platbase',
+                'install-lib',
+                'install-platlib',
+                'install-purelib',
+                'install-headers',
+                'install-scripts',
+                'install-data',
+                'prefix',
+                'exec-prefix',
+                'home',
+                'user',
+                'root',
+            ]
+        )
+
+        ignore_options = frozenset(ignore_options)
+
+        if filenames is None:
+            filenames = self.find_config_files()
+
+        if DEBUG:
+            self.announce("Distribution.parse_config_files():")
+
+        parser = ConfigParser()
+        parser.optionxform = str
+        for filename in filenames:
+            with open(filename, encoding='utf-8') as reader:
+                if DEBUG:
+                    self.announce("  reading {filename}".format(**locals()))
+                parser.read_file(reader)
+            for section in parser.sections():
+                options = parser.options(section)
+                opt_dict = self.get_option_dict(section)
+
+                for opt in options:
+                    if opt == '__name__' or opt in ignore_options:
+                        continue
+
+                    val = parser.get(section, opt)
+                    opt = self._enforce_underscore(opt, section)
+                    opt = self._enforce_option_lowercase(opt, section)
+                    opt_dict[opt] = (filename, val)
+
+            # Make the ConfigParser forget everything (so we retain
+            # the original filenames that options come from)
+            parser.__init__()
+
+        if 'global' not in self.command_options:
+            return
+
+        # If there was a "global" section in the config file, use it
+        # to set Distribution options.
+
+        for opt, (src, val) in self.command_options['global'].items():
+            alias = self.negative_opt.get(opt)
+            if alias:
+                val = not strtobool(val)
+            elif opt in ('verbose', 'dry_run'):  # ugh!
+                val = strtobool(val)
+
+            try:
+                setattr(self, alias or opt, val)
+            except ValueError as e:
+                raise DistutilsOptionError(e) from e
+
+    def _enforce_underscore(self, opt: str, section: str) -> str:
+        if "-" not in opt or self._skip_setupcfg_normalization(section):
+            return opt
+
+        underscore_opt = opt.replace('-', '_')
+        affected = f"(Affected: {self.metadata.name})." if self.metadata.name else ""
+        SetuptoolsDeprecationWarning.emit(
+            f"Invalid dash-separated key {opt!r} in {section!r} (setup.cfg), "
+            f"please use the underscore name {underscore_opt!r} instead.",
+            f"""
+            Usage of dash-separated {opt!r} will not be supported in future
+            versions. Please use the underscore name {underscore_opt!r} instead.
+            {affected}
+            """,
+            see_docs="userguide/declarative_config.html",
+            due_date=(2026, 3, 3),
+            # Warning initially introduced in 3 Mar 2021
+        )
+        return underscore_opt
+
+    def _enforce_option_lowercase(self, opt: str, section: str) -> str:
+        if opt.islower() or self._skip_setupcfg_normalization(section):
+            return opt
+
+        lowercase_opt = opt.lower()
+        affected = f"(Affected: {self.metadata.name})." if self.metadata.name else ""
+        SetuptoolsDeprecationWarning.emit(
+            f"Invalid uppercase key {opt!r} in {section!r} (setup.cfg), "
+            f"please use lowercase {lowercase_opt!r} instead.",
+            f"""
+            Usage of uppercase key {opt!r} in {section!r} will not be supported in
+            future versions. Please use lowercase {lowercase_opt!r} instead.
+            {affected}
+            """,
+            see_docs="userguide/declarative_config.html",
+            due_date=(2026, 3, 3),
+            # Warning initially introduced in 6 Mar 2021
+        )
+        return lowercase_opt
+
+    def _skip_setupcfg_normalization(self, section: str) -> bool:
+        skip = (
+            'options.extras_require',
+            'options.data_files',
+            'options.entry_points',
+            'options.package_data',
+            'options.exclude_package_data',
+        )
+        return section in skip or not self._is_setuptools_section(section)
+
+    def _is_setuptools_section(self, section: str) -> bool:
+        return (
+            section == "metadata"
+            or section.startswith("options")
+            or section in _setuptools_commands()
+        )
+
+    # FIXME: 'Distribution._set_command_options' is too complex (14)
+    def _set_command_options(self, command_obj, option_dict=None):  # noqa: C901
+        """
+        Set the options for 'command_obj' from 'option_dict'.  Basically
+        this means copying elements of a dictionary ('option_dict') to
+        attributes of an instance ('command').
+
+        'command_obj' must be a Command instance.  If 'option_dict' is not
+        supplied, uses the standard option dictionary for this command
+        (from 'self.command_options').
+
+        (Adopted from distutils.dist.Distribution._set_command_options)
+        """
+        command_name = command_obj.get_command_name()
+        if option_dict is None:
+            option_dict = self.get_option_dict(command_name)
+
+        if DEBUG:
+            self.announce(f"  setting options for '{command_name}' command:")
+        for option, (source, value) in option_dict.items():
+            if DEBUG:
+                self.announce(f"    {option} = {value} (from {source})")
+            try:
+                bool_opts = [translate_longopt(o) for o in command_obj.boolean_options]
+            except AttributeError:
+                bool_opts = []
+            try:
+                neg_opt = command_obj.negative_opt
+            except AttributeError:
+                neg_opt = {}
+
+            try:
+                is_string = isinstance(value, str)
+                if option in neg_opt and is_string:
+                    setattr(command_obj, neg_opt[option], not strtobool(value))
+                elif option in bool_opts and is_string:
+                    setattr(command_obj, option, strtobool(value))
+                elif hasattr(command_obj, option):
+                    setattr(command_obj, option, value)
+                else:
+                    raise DistutilsOptionError(
+                        f"error in {source}: command '{command_name}' has no such option '{option}'"
+                    )
+            except ValueError as e:
+                raise DistutilsOptionError(e) from e
+
+    def _get_project_config_files(self, filenames: Iterable[StrPath] | None):
+        """Add default file and split between INI and TOML"""
+        tomlfiles = []
+        standard_project_metadata = Path(self.src_root or os.curdir, "pyproject.toml")
+        if filenames is not None:
+            parts = partition(lambda f: Path(f).suffix == ".toml", filenames)
+            filenames = list(parts[0])  # 1st element => predicate is False
+            tomlfiles = list(parts[1])  # 2nd element => predicate is True
+        elif standard_project_metadata.exists():
+            tomlfiles = [standard_project_metadata]
+        return filenames, tomlfiles
+
+    def parse_config_files(
+        self,
+        filenames: Iterable[StrPath] | None = None,
+        ignore_option_errors: bool = False,
+    ) -> None:
+        """Parses configuration files from various levels
+        and loads configuration.
+        """
+        inifiles, tomlfiles = self._get_project_config_files(filenames)
+
+        self._parse_config_files(filenames=inifiles)
+
+        setupcfg.parse_configuration(
+            self, self.command_options, ignore_option_errors=ignore_option_errors
+        )
+        for filename in tomlfiles:
+            pyprojecttoml.apply_configuration(self, filename, ignore_option_errors)
+
+        self._finalize_requires()
+        self._finalize_license_expression()
+        self._finalize_license_files()
+
+    def fetch_build_eggs(self, requires: _StrOrIter) -> list[metadata.Distribution]:
+        """Resolve pre-setup requirements"""
+        from .installer import _fetch_build_eggs
+
+        return _fetch_build_eggs(self, requires)
+
+    def finalize_options(self) -> None:
+        """
+        Allow plugins to apply arbitrary operations to the
+        distribution. Each hook may optionally define a 'order'
+        to influence the order of execution. Smaller numbers
+        go first and the default is 0.
+        """
+        group = 'setuptools.finalize_distribution_options'
+
+        def by_order(hook):
+            return getattr(hook, 'order', 0)
+
+        defined = metadata.entry_points(group=group)
+        filtered = itertools.filterfalse(self._removed, defined)
+        loaded = map(lambda e: e.load(), filtered)
+        for ep in sorted(loaded, key=by_order):
+            ep(self)
+
+    @staticmethod
+    def _removed(ep):
+        """
+        When removing an entry point, if metadata is loaded
+        from an older version of Setuptools, that removed
+        entry point will attempt to be loaded and will fail.
+        See #2765 for more details.
+        """
+        removed = {
+            # removed 2021-09-05
+            '2to3_doctests',
+        }
+        return ep.name in removed
+
+    def _finalize_setup_keywords(self):
+        for ep in metadata.entry_points(group='distutils.setup_keywords'):
+            value = getattr(self, ep.name, None)
+            if value is not None:
+                ep.load()(self, ep.name, value)
+
+    def get_egg_cache_dir(self):
+        from . import windows_support
+
+        egg_cache_dir = os.path.join(os.curdir, '.eggs')
+        if not os.path.exists(egg_cache_dir):
+            os.mkdir(egg_cache_dir)
+            windows_support.hide_file(egg_cache_dir)
+            readme_txt_filename = os.path.join(egg_cache_dir, 'README.txt')
+            with open(readme_txt_filename, 'w', encoding="utf-8") as f:
+                f.write(
+                    'This directory contains eggs that were downloaded '
+                    'by setuptools to build, test, and run plug-ins.\n\n'
+                )
+                f.write(
+                    'This directory caches those eggs to prevent '
+                    'repeated downloads.\n\n'
+                )
+                f.write('However, it is safe to delete this directory.\n\n')
+
+        return egg_cache_dir
+
+    def fetch_build_egg(self, req):
+        """Fetch an egg needed for building"""
+        from .installer import fetch_build_egg
+
+        return fetch_build_egg(self, req)
+
+    def get_command_class(self, command: str) -> type[distutils.cmd.Command]:  # type: ignore[override] # Not doing complex overrides yet
+        """Pluggable version of get_command_class()"""
+        if command in self.cmdclass:
+            return self.cmdclass[command]
+
+        # Special case bdist_wheel so it's never loaded from "wheel"
+        if command == 'bdist_wheel':
+            from .command.bdist_wheel import bdist_wheel
+
+            return bdist_wheel
+
+        eps = metadata.entry_points(group='distutils.commands', name=command)
+        for ep in eps:
+            self.cmdclass[command] = cmdclass = ep.load()
+            return cmdclass
+        else:
+            return _Distribution.get_command_class(self, command)
+
+    def print_commands(self):
+        for ep in metadata.entry_points(group='distutils.commands'):
+            if ep.name not in self.cmdclass:
+                cmdclass = ep.load()
+                self.cmdclass[ep.name] = cmdclass
+        return _Distribution.print_commands(self)
+
+    def get_command_list(self):
+        for ep in metadata.entry_points(group='distutils.commands'):
+            if ep.name not in self.cmdclass:
+                cmdclass = ep.load()
+                self.cmdclass[ep.name] = cmdclass
+        return _Distribution.get_command_list(self)
+
+    def include(self, **attrs) -> None:
+        """Add items to distribution that are named in keyword arguments
+
+        For example, 'dist.include(py_modules=["x"])' would add 'x' to
+        the distribution's 'py_modules' attribute, if it was not already
+        there.
+
+        Currently, this method only supports inclusion for attributes that are
+        lists or tuples.  If you need to add support for adding to other
+        attributes in this or a subclass, you can add an '_include_X' method,
+        where 'X' is the name of the attribute.  The method will be called with
+        the value passed to 'include()'.  So, 'dist.include(foo={"bar":"baz"})'
+        will try to call 'dist._include_foo({"bar":"baz"})', which can then
+        handle whatever special inclusion logic is needed.
+        """
+        for k, v in attrs.items():
+            include = getattr(self, '_include_' + k, None)
+            if include:
+                include(v)
+            else:
+                self._include_misc(k, v)
+
+    def exclude_package(self, package: str) -> None:
+        """Remove packages, modules, and extensions in named package"""
+
+        pfx = package + '.'
+        if self.packages:
+            self.packages = [
+                p for p in self.packages if p != package and not p.startswith(pfx)
+            ]
+
+        if self.py_modules:
+            self.py_modules = [
+                p for p in self.py_modules if p != package and not p.startswith(pfx)
+            ]
+
+        if self.ext_modules:
+            self.ext_modules = [
+                p
+                for p in self.ext_modules
+                if p.name != package and not p.name.startswith(pfx)
+            ]
+
+    def has_contents_for(self, package: str) -> bool:
+        """Return true if 'exclude_package(package)' would do something"""
+
+        pfx = package + '.'
+
+        for p in self.iter_distribution_names():
+            if p == package or p.startswith(pfx):
+                return True
+
+        return False
+
+    def _exclude_misc(self, name: str, value: _Sequence) -> None:
+        """Handle 'exclude()' for list/tuple attrs without a special handler"""
+        if not isinstance(value, _sequence):
+            raise DistutilsSetupError(
+                f"{name}: setting must be of type <{_sequence_type_repr}> (got {value!r})"
+            )
+        try:
+            old = getattr(self, name)
+        except AttributeError as e:
+            raise DistutilsSetupError(f"{name}: No such distribution setting") from e
+        if old is not None and not isinstance(old, _sequence):
+            raise DistutilsSetupError(
+                name + ": this setting cannot be changed via include/exclude"
+            )
+        elif old:
+            setattr(self, name, [item for item in old if item not in value])
+
+    def _include_misc(self, name: str, value: _Sequence) -> None:
+        """Handle 'include()' for list/tuple attrs without a special handler"""
+
+        if not isinstance(value, _sequence):
+            raise DistutilsSetupError(
+                f"{name}: setting must be of type <{_sequence_type_repr}> (got {value!r})"
+            )
+        try:
+            old = getattr(self, name)
+        except AttributeError as e:
+            raise DistutilsSetupError(f"{name}: No such distribution setting") from e
+        if old is None:
+            setattr(self, name, value)
+        elif not isinstance(old, _sequence):
+            raise DistutilsSetupError(
+                name + ": this setting cannot be changed via include/exclude"
+            )
+        else:
+            new = [item for item in value if item not in old]
+            setattr(self, name, list(old) + new)
+
+    def exclude(self, **attrs) -> None:
+        """Remove items from distribution that are named in keyword arguments
+
+        For example, 'dist.exclude(py_modules=["x"])' would remove 'x' from
+        the distribution's 'py_modules' attribute.  Excluding packages uses
+        the 'exclude_package()' method, so all of the package's contained
+        packages, modules, and extensions are also excluded.
+
+        Currently, this method only supports exclusion from attributes that are
+        lists or tuples.  If you need to add support for excluding from other
+        attributes in this or a subclass, you can add an '_exclude_X' method,
+        where 'X' is the name of the attribute.  The method will be called with
+        the value passed to 'exclude()'.  So, 'dist.exclude(foo={"bar":"baz"})'
+        will try to call 'dist._exclude_foo({"bar":"baz"})', which can then
+        handle whatever special exclusion logic is needed.
+        """
+        for k, v in attrs.items():
+            exclude = getattr(self, '_exclude_' + k, None)
+            if exclude:
+                exclude(v)
+            else:
+                self._exclude_misc(k, v)
+
+    def _exclude_packages(self, packages: _Sequence) -> None:
+        if not isinstance(packages, _sequence):
+            raise DistutilsSetupError(
+                f"packages: setting must be of type <{_sequence_type_repr}> (got {packages!r})"
+            )
+        list(map(self.exclude_package, packages))
+
+    def _parse_command_opts(self, parser, args):
+        # Remove --with-X/--without-X options when processing command args
+        self.global_options = self.__class__.global_options
+        self.negative_opt = self.__class__.negative_opt
+
+        # First, expand any aliases
+        command = args[0]
+        aliases = self.get_option_dict('aliases')
+        while command in aliases:
+            _src, alias = aliases[command]
+            del aliases[command]  # ensure each alias can expand only once!
+            import shlex
+
+            args[:1] = shlex.split(alias, True)
+            command = args[0]
+
+        nargs = _Distribution._parse_command_opts(self, parser, args)
+
+        # Handle commands that want to consume all remaining arguments
+        cmd_class = self.get_command_class(command)
+        if getattr(cmd_class, 'command_consumes_arguments', None):
+            self.get_option_dict(command)['args'] = ("command line", nargs)
+            if nargs is not None:
+                return []
+
+        return nargs
+
+    def get_cmdline_options(self) -> dict[str, dict[str, str | None]]:
+        """Return a '{cmd: {opt:val}}' map of all command-line options
+
+        Option names are all long, but do not include the leading '--', and
+        contain dashes rather than underscores.  If the option doesn't take
+        an argument (e.g. '--quiet'), the 'val' is 'None'.
+
+        Note that options provided by config files are intentionally excluded.
+        """
+
+        d: dict[str, dict[str, str | None]] = {}
+
+        for cmd, opts in self.command_options.items():
+            val: str | None
+            for opt, (src, val) in opts.items():
+                if src != "command line":
+                    continue
+
+                opt = opt.replace('_', '-')
+
+                if val == 0:
+                    cmdobj = self.get_command_obj(cmd)
+                    neg_opt = self.negative_opt.copy()
+                    neg_opt.update(getattr(cmdobj, 'negative_opt', {}))
+                    for neg, pos in neg_opt.items():
+                        if pos == opt:
+                            opt = neg
+                            val = None
+                            break
+                    else:
+                        raise AssertionError("Shouldn't be able to get here")
+
+                elif val == 1:
+                    val = None
+
+                d.setdefault(cmd, {})[opt] = val
+
+        return d
+
+    def iter_distribution_names(self):
+        """Yield all packages, modules, and extension names in distribution"""
+
+        yield from self.packages or ()
+
+        yield from self.py_modules or ()
+
+        for ext in self.ext_modules or ():
+            if isinstance(ext, tuple):
+                name, _buildinfo = ext
+            else:
+                name = ext.name
+            if name.endswith('module'):
+                name = name[:-6]
+            yield name
+
+    def handle_display_options(self, option_order):
+        """If there were any non-global "display-only" options
+        (--help-commands or the metadata display options) on the command
+        line, display the requested info and return true; else return
+        false.
+        """
+        import sys
+
+        if self.help_commands:
+            return _Distribution.handle_display_options(self, option_order)
+
+        # Stdout may be StringIO (e.g. in tests)
+        if not isinstance(sys.stdout, io.TextIOWrapper):
+            return _Distribution.handle_display_options(self, option_order)
+
+        # Don't wrap stdout if utf-8 is already the encoding. Provides
+        #  workaround for #334.
+        if sys.stdout.encoding.lower() in ('utf-8', 'utf8'):
+            return _Distribution.handle_display_options(self, option_order)
+
+        # Print metadata in UTF-8 no matter the platform
+        encoding = sys.stdout.encoding
+        sys.stdout.reconfigure(encoding='utf-8')
+        try:
+            return _Distribution.handle_display_options(self, option_order)
+        finally:
+            sys.stdout.reconfigure(encoding=encoding)
+
+    def run_command(self, command) -> None:
+        self.set_defaults()
+        # Postpone defaults until all explicit configuration is considered
+        # (setup() args, config files, command line and plugins)
+
+        super().run_command(command)
+
+
+@functools.cache
+def _setuptools_commands() -> set[str]:
+    try:
+        # Use older API for importlib.metadata compatibility
+        entry_points = metadata.distribution('setuptools').entry_points
+        eps: Iterable[str] = (ep.name for ep in entry_points)
+    except metadata.PackageNotFoundError:
+        # during bootstrapping, distribution doesn't exist
+        eps = []
+    return {*distutils.command.__all__, *eps}
+
+
+class DistDeprecationWarning(SetuptoolsDeprecationWarning):
+    """Class for warning about deprecations in dist in
+    setuptools. Not ignored by default, unlike DeprecationWarning."""
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/errors.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/errors.py
new file mode 100644
index 0000000000000000000000000000000000000000..990ecbf4e2f18eb188addc9e0466152a20193a90
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/errors.py
@@ -0,0 +1,67 @@
+"""setuptools.errors
+
+Provides exceptions used by setuptools modules.
+"""
+
+from __future__ import annotations
+
+from distutils import errors as _distutils_errors
+
+# Re-export errors from distutils to facilitate the migration to PEP632
+
+ByteCompileError = _distutils_errors.DistutilsByteCompileError
+CCompilerError = _distutils_errors.CCompilerError
+ClassError = _distutils_errors.DistutilsClassError
+CompileError = _distutils_errors.CompileError
+ExecError = _distutils_errors.DistutilsExecError
+FileError = _distutils_errors.DistutilsFileError
+InternalError = _distutils_errors.DistutilsInternalError
+LibError = _distutils_errors.LibError
+LinkError = _distutils_errors.LinkError
+ModuleError = _distutils_errors.DistutilsModuleError
+OptionError = _distutils_errors.DistutilsOptionError
+PlatformError = _distutils_errors.DistutilsPlatformError
+PreprocessError = _distutils_errors.PreprocessError
+SetupError = _distutils_errors.DistutilsSetupError
+TemplateError = _distutils_errors.DistutilsTemplateError
+UnknownFileError = _distutils_errors.UnknownFileError
+
+# The root error class in the hierarchy
+BaseError = _distutils_errors.DistutilsError
+
+
+class InvalidConfigError(OptionError):  # type: ignore[valid-type, misc] # distutils imports are `Any` on python 3.12+
+    """Error used for invalid configurations."""
+
+
+class RemovedConfigError(OptionError):  # type: ignore[valid-type, misc] # distutils imports are `Any` on python 3.12+
+    """Error used for configurations that were deprecated and removed."""
+
+
+class RemovedCommandError(BaseError, RuntimeError):  # type: ignore[valid-type, misc] # distutils imports are `Any` on python 3.12+
+    """Error used for commands that have been removed in setuptools.
+
+    Since ``setuptools`` is built on ``distutils``, simply removing a command
+    from ``setuptools`` will make the behavior fall back to ``distutils``; this
+    error is raised if a command exists in ``distutils`` but has been actively
+    removed in ``setuptools``.
+    """
+
+
+class PackageDiscoveryError(BaseError, RuntimeError):  # type: ignore[valid-type, misc] # distutils imports are `Any` on python 3.12+
+    """Impossible to perform automatic discovery of packages and/or modules.
+
+    The current project layout or given discovery options can lead to problems when
+    scanning the project directory.
+
+    Setuptools might also refuse to complete auto-discovery if an error prone condition
+    is detected (e.g. when a project is organised as a flat-layout but contains
+    multiple directories that can be taken as top-level packages inside a single
+    distribution [*]_). In these situations the users are encouraged to be explicit
+    about which packages to include or to make the discovery parameters more specific.
+
+    .. [*] Since multi-package distributions are uncommon it is very likely that the
+       developers did not intend for all the directories to be packaged, and are just
+       leaving auxiliary code in the repository top-level, such as maintenance-related
+       scripts.
+    """
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/extension.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/extension.py
new file mode 100644
index 0000000000000000000000000000000000000000..76e03d9d6bdcdbd72e443e90c85d34429c22c261
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/extension.py
@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import functools
+import re
+from typing import TYPE_CHECKING
+
+from setuptools._path import StrPath
+
+from .monkey import get_unpatched
+
+import distutils.core
+import distutils.errors
+import distutils.extension
+
+
+def _have_cython():
+    """
+    Return True if Cython can be imported.
+    """
+    cython_impl = 'Cython.Distutils.build_ext'
+    try:
+        # from (cython_impl) import build_ext
+        __import__(cython_impl, fromlist=['build_ext']).build_ext
+    except Exception:
+        return False
+    return True
+
+
+# for compatibility
+have_pyrex = _have_cython
+if TYPE_CHECKING:
+    # Work around a mypy issue where type[T] can't be used as a base: https://github.com/python/mypy/issues/10962
+    from distutils.core import Extension as _Extension
+else:
+    _Extension = get_unpatched(distutils.core.Extension)
+
+
+class Extension(_Extension):
+    """
+    Describes a single extension module.
+
+    This means that all source files will be compiled into a single binary file
+    ``.`` (with ```` derived from ``name`` and
+    ```` defined by one of the values in
+    ``importlib.machinery.EXTENSION_SUFFIXES``).
+
+    In the case ``.pyx`` files are passed as ``sources and`` ``Cython`` is **not**
+    installed in the build environment, ``setuptools`` may also try to look for the
+    equivalent ``.cpp`` or ``.c`` files.
+
+    :arg str name:
+      the full name of the extension, including any packages -- ie.
+      *not* a filename or pathname, but Python dotted name
+
+    :arg list[str|os.PathLike[str]] sources:
+      list of source filenames, relative to the distribution root
+      (where the setup script lives), in Unix form (slash-separated)
+      for portability.  Source files may be C, C++, SWIG (.i),
+      platform-specific resource files, or whatever else is recognized
+      by the "build_ext" command as source for a Python extension.
+
+    :keyword list[str] include_dirs:
+      list of directories to search for C/C++ header files (in Unix
+      form for portability)
+
+    :keyword list[tuple[str, str|None]] define_macros:
+      list of macros to define; each macro is defined using a 2-tuple:
+      the first item corresponding to the name of the macro and the second
+      item either a string with its value or None to
+      define it without a particular value (equivalent of "#define
+      FOO" in source or -DFOO on Unix C compiler command line)
+
+    :keyword list[str] undef_macros:
+      list of macros to undefine explicitly
+
+    :keyword list[str] library_dirs:
+      list of directories to search for C/C++ libraries at link time
+
+    :keyword list[str] libraries:
+      list of library names (not filenames or paths) to link against
+
+    :keyword list[str] runtime_library_dirs:
+      list of directories to search for C/C++ libraries at run time
+      (for shared extensions, this is when the extension is loaded).
+      Setting this will cause an exception during build on Windows
+      platforms.
+
+    :keyword list[str] extra_objects:
+      list of extra files to link with (eg. object files not implied
+      by 'sources', static library that must be explicitly specified,
+      binary resource files, etc.)
+
+    :keyword list[str] extra_compile_args:
+      any extra platform- and compiler-specific information to use
+      when compiling the source files in 'sources'.  For platforms and
+      compilers where "command line" makes sense, this is typically a
+      list of command-line arguments, but for other platforms it could
+      be anything.
+
+    :keyword list[str] extra_link_args:
+      any extra platform- and compiler-specific information to use
+      when linking object files together to create the extension (or
+      to create a new static Python interpreter).  Similar
+      interpretation as for 'extra_compile_args'.
+
+    :keyword list[str] export_symbols:
+      list of symbols to be exported from a shared extension.  Not
+      used on all platforms, and not generally necessary for Python
+      extensions, which typically export exactly one symbol: "init" +
+      extension_name.
+
+    :keyword list[str] swig_opts:
+      any extra options to pass to SWIG if a source file has the .i
+      extension.
+
+    :keyword list[str] depends:
+      list of files that the extension depends on
+
+    :keyword str language:
+      extension language (i.e. "c", "c++", "objc"). Will be detected
+      from the source extensions if not provided.
+
+    :keyword bool optional:
+      specifies that a build failure in the extension should not abort the
+      build process, but simply not install the failing extension.
+
+    :keyword bool py_limited_api:
+      opt-in flag for the usage of :doc:`Python's limited API `.
+
+    :raises setuptools.errors.PlatformError: if ``runtime_library_dirs`` is
+      specified on Windows. (since v63)
+    """
+
+    # These 4 are set and used in setuptools/command/build_ext.py
+    # The lack of a default value and risk of `AttributeError` is purposeful
+    # to avoid people forgetting to call finalize_options if they modify the extension list.
+    # See example/rationale in https://github.com/pypa/setuptools/issues/4529.
+    _full_name: str  #: Private API, internal use only.
+    _links_to_dynamic: bool  #: Private API, internal use only.
+    _needs_stub: bool  #: Private API, internal use only.
+    _file_name: str  #: Private API, internal use only.
+
+    def __init__(
+        self,
+        name: str,
+        sources: list[StrPath],
+        *args,
+        py_limited_api: bool = False,
+        **kw,
+    ) -> None:
+        # The *args is needed for compatibility as calls may use positional
+        # arguments. py_limited_api may be set only via keyword.
+        self.py_limited_api = py_limited_api
+        super().__init__(
+            name,
+            sources,  # type: ignore[arg-type] # Vendored version of setuptools supports PathLike
+            *args,
+            **kw,
+        )
+
+    def _convert_pyx_sources_to_lang(self):
+        """
+        Replace sources with .pyx extensions to sources with the target
+        language extension. This mechanism allows language authors to supply
+        pre-converted sources but to prefer the .pyx sources.
+        """
+        if _have_cython():
+            # the build has Cython, so allow it to compile the .pyx files
+            return
+        lang = self.language or ''
+        target_ext = '.cpp' if lang.lower() == 'c++' else '.c'
+        sub = functools.partial(re.sub, '.pyx$', target_ext)
+        self.sources = list(map(sub, self.sources))
+
+
+class Library(Extension):
+    """Just like a regular Extension, but built as a library instead"""
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/glob.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/glob.py
new file mode 100644
index 0000000000000000000000000000000000000000..1dfff2cd50ff87b8cef9d936f1fc9d4a2478b136
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/glob.py
@@ -0,0 +1,185 @@
+"""
+Filename globbing utility. Mostly a copy of `glob` from Python 3.5.
+
+Changes include:
+ * `yield from` and PEP3102 `*` removed.
+ * Hidden files are not ignored.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import os
+import re
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, AnyStr, overload
+
+if TYPE_CHECKING:
+    from _typeshed import BytesPath, StrOrBytesPath, StrPath
+
+__all__ = ["glob", "iglob", "escape"]
+
+
+def glob(pathname: AnyStr, recursive: bool = False) -> list[AnyStr]:
+    """Return a list of paths matching a pathname pattern.
+
+    The pattern may contain simple shell-style wildcards a la
+    fnmatch. However, unlike fnmatch, filenames starting with a
+    dot are special cases that are not matched by '*' and '?'
+    patterns.
+
+    If recursive is true, the pattern '**' will match any files and
+    zero or more directories and subdirectories.
+    """
+    return list(iglob(pathname, recursive=recursive))
+
+
+def iglob(pathname: AnyStr, recursive: bool = False) -> Iterator[AnyStr]:
+    """Return an iterator which yields the paths matching a pathname pattern.
+
+    The pattern may contain simple shell-style wildcards a la
+    fnmatch. However, unlike fnmatch, filenames starting with a
+    dot are special cases that are not matched by '*' and '?'
+    patterns.
+
+    If recursive is true, the pattern '**' will match any files and
+    zero or more directories and subdirectories.
+    """
+    it = _iglob(pathname, recursive)
+    if recursive and _isrecursive(pathname):
+        s = next(it)  # skip empty string
+        assert not s
+    return it
+
+
+def _iglob(pathname: AnyStr, recursive: bool) -> Iterator[AnyStr]:
+    dirname, basename = os.path.split(pathname)
+    glob_in_dir = glob2 if recursive and _isrecursive(basename) else glob1
+
+    if not has_magic(pathname):
+        if basename:
+            if os.path.lexists(pathname):
+                yield pathname
+        else:
+            # Patterns ending with a slash should match only directories
+            if os.path.isdir(dirname):
+                yield pathname
+        return
+
+    if not dirname:
+        yield from glob_in_dir(dirname, basename)
+        return
+    # `os.path.split()` returns the argument itself as a dirname if it is a
+    # drive or UNC path.  Prevent an infinite recursion if a drive or UNC path
+    # contains magic characters (i.e. r'\\?\C:').
+    if dirname != pathname and has_magic(dirname):
+        dirs: Iterable[AnyStr] = _iglob(dirname, recursive)
+    else:
+        dirs = [dirname]
+    if not has_magic(basename):
+        glob_in_dir = glob0
+    for dirname in dirs:
+        for name in glob_in_dir(dirname, basename):
+            yield os.path.join(dirname, name)
+
+
+# These 2 helper functions non-recursively glob inside a literal directory.
+# They return a list of basenames. `glob1` accepts a pattern while `glob0`
+# takes a literal basename (so it only has to check for its existence).
+
+
+@overload
+def glob1(dirname: StrPath, pattern: str) -> list[str]: ...
+@overload
+def glob1(dirname: BytesPath, pattern: bytes) -> list[bytes]: ...
+def glob1(dirname: StrOrBytesPath, pattern: str | bytes) -> list[str] | list[bytes]:
+    if not dirname:
+        if isinstance(pattern, bytes):
+            dirname = os.curdir.encode('ASCII')
+        else:
+            dirname = os.curdir
+    try:
+        names = os.listdir(dirname)
+    except OSError:
+        return []
+    # mypy false-positives: str or bytes type possibility is always kept in sync
+    return fnmatch.filter(names, pattern)  # type: ignore[type-var, return-value]
+
+
+def glob0(dirname, basename):
+    if not basename:
+        # `os.path.split()` returns an empty basename for paths ending with a
+        # directory separator.  'q*x/' should match only directories.
+        if os.path.isdir(dirname):
+            return [basename]
+    else:
+        if os.path.lexists(os.path.join(dirname, basename)):
+            return [basename]
+    return []
+
+
+# This helper function recursively yields relative pathnames inside a literal
+# directory.
+
+
+@overload
+def glob2(dirname: StrPath, pattern: str) -> Iterator[str]: ...
+@overload
+def glob2(dirname: BytesPath, pattern: bytes) -> Iterator[bytes]: ...
+def glob2(dirname: StrOrBytesPath, pattern: str | bytes) -> Iterator[str | bytes]:
+    assert _isrecursive(pattern)
+    yield pattern[:0]
+    yield from _rlistdir(dirname)
+
+
+# Recursively yields relative pathnames inside a literal directory.
+@overload
+def _rlistdir(dirname: StrPath) -> Iterator[str]: ...
+@overload
+def _rlistdir(dirname: BytesPath) -> Iterator[bytes]: ...
+def _rlistdir(dirname: StrOrBytesPath) -> Iterator[str | bytes]:
+    if not dirname:
+        if isinstance(dirname, bytes):
+            dirname = os.curdir.encode('ASCII')
+        else:
+            dirname = os.curdir
+    try:
+        names = os.listdir(dirname)
+    except OSError:
+        return
+    for x in names:
+        yield x
+        # mypy false-positives: str or bytes type possibility is always kept in sync
+        path = os.path.join(dirname, x) if dirname else x  # type: ignore[arg-type]
+        for y in _rlistdir(path):
+            yield os.path.join(x, y)  # type: ignore[arg-type]
+
+
+magic_check = re.compile('([*?[])')
+magic_check_bytes = re.compile(b'([*?[])')
+
+
+def has_magic(s: str | bytes) -> bool:
+    if isinstance(s, bytes):
+        return magic_check_bytes.search(s) is not None
+    else:
+        return magic_check.search(s) is not None
+
+
+def _isrecursive(pattern: str | bytes) -> bool:
+    if isinstance(pattern, bytes):
+        return pattern == b'**'
+    else:
+        return pattern == '**'
+
+
+def escape(pathname):
+    """Escape all special characters."""
+    # Escaping is done by wrapping any of "*?[" between square brackets.
+    # Metacharacters do not work in the drive part and shouldn't be escaped.
+    drive, pathname = os.path.splitdrive(pathname)
+    if isinstance(pathname, bytes):
+        pathname = magic_check_bytes.sub(rb'[\1]', pathname)
+    else:
+        pathname = magic_check.sub(r'[\1]', pathname)
+    return drive + pathname
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-32.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-32.exe
new file mode 100644
index 0000000000000000000000000000000000000000..1eb430c6d614a5daea4139badc09c222a4b0e72a
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-32.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-64.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-64.exe
new file mode 100644
index 0000000000000000000000000000000000000000..031cb77c17ba8d8a983448268851d612e05e80d1
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-64.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-arm64.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-arm64.exe
new file mode 100644
index 0000000000000000000000000000000000000000..1e00ffacb182c2af206e5dd9d9fbc41d236da0d1
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui-arm64.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui.exe b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui.exe
new file mode 100644
index 0000000000000000000000000000000000000000..1eb430c6d614a5daea4139badc09c222a4b0e72a
Binary files /dev/null and b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/gui.exe differ
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/installer.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/installer.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c26e3a1f4ec6a271198fffb7489dc6e6e612c06
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/installer.py
@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import glob
+import itertools
+import os
+import subprocess
+import sys
+import tempfile
+
+import packaging.requirements
+import packaging.utils
+
+from . import _reqs
+from ._importlib import metadata
+from .warnings import SetuptoolsDeprecationWarning
+from .wheel import Wheel
+
+from distutils import log
+from distutils.errors import DistutilsError
+
+
+def _fixup_find_links(find_links):
+    """Ensure find-links option end-up being a list of strings."""
+    if isinstance(find_links, str):
+        return find_links.split()
+    assert isinstance(find_links, (tuple, list))
+    return find_links
+
+
+def fetch_build_egg(dist, req):
+    """Fetch an egg needed for building.
+
+    Use pip/wheel to fetch/build a wheel."""
+    _DeprecatedInstaller.emit()
+    _warn_wheel_not_available(dist)
+    return _fetch_build_egg_no_warn(dist, req)
+
+
+def _present(req):
+    return any(_dist_matches_req(dist, req) for dist in metadata.distributions())
+
+
+def _fetch_build_eggs(dist, requires: _reqs._StrOrIter) -> list[metadata.Distribution]:
+    _DeprecatedInstaller.emit(stacklevel=3)
+    _warn_wheel_not_available(dist)
+
+    parsed_reqs = _reqs.parse(requires)
+
+    missing_reqs = itertools.filterfalse(_present, parsed_reqs)
+
+    needed_reqs = (
+        req for req in missing_reqs if not req.marker or req.marker.evaluate()
+    )
+    resolved_dists = [_fetch_build_egg_no_warn(dist, req) for req in needed_reqs]
+    for dist in resolved_dists:
+        # dist.locate_file('') is the directory containing EGG-INFO, where the importabl
+        # contents can be found.
+        sys.path.insert(0, str(dist.locate_file('')))
+    return resolved_dists
+
+
+def _dist_matches_req(egg_dist, req):
+    return (
+        packaging.utils.canonicalize_name(egg_dist.name)
+        == packaging.utils.canonicalize_name(req.name)
+        and egg_dist.version in req.specifier
+    )
+
+
+def _fetch_build_egg_no_warn(dist, req):  # noqa: C901  # is too complex (16)  # FIXME
+    # Ignore environment markers; if supplied, it is required.
+    req = strip_marker(req)
+    # Take easy_install options into account, but do not override relevant
+    # pip environment variables (like PIP_INDEX_URL or PIP_QUIET); they'll
+    # take precedence.
+    opts = dist.get_option_dict('easy_install')
+    if 'allow_hosts' in opts:
+        raise DistutilsError(
+            'the `allow-hosts` option is not supported '
+            'when using pip to install requirements.'
+        )
+    quiet = 'PIP_QUIET' not in os.environ and 'PIP_VERBOSE' not in os.environ
+    if 'PIP_INDEX_URL' in os.environ:
+        index_url = None
+    elif 'index_url' in opts:
+        index_url = opts['index_url'][1]
+    else:
+        index_url = None
+    find_links = (
+        _fixup_find_links(opts['find_links'][1])[:] if 'find_links' in opts else []
+    )
+    if dist.dependency_links:
+        find_links.extend(dist.dependency_links)
+    eggs_dir = os.path.realpath(dist.get_egg_cache_dir())
+    cached_dists = metadata.Distribution.discover(path=glob.glob(f'{eggs_dir}/*.egg'))
+    for egg_dist in cached_dists:
+        if _dist_matches_req(egg_dist, req):
+            return egg_dist
+    with tempfile.TemporaryDirectory() as tmpdir:
+        cmd = [
+            sys.executable,
+            '-m',
+            'pip',
+            '--disable-pip-version-check',
+            'wheel',
+            '--no-deps',
+            '-w',
+            tmpdir,
+        ]
+        if quiet:
+            cmd.append('--quiet')
+        if index_url is not None:
+            cmd.extend(('--index-url', index_url))
+        for link in find_links or []:
+            cmd.extend(('--find-links', link))
+        # If requirement is a PEP 508 direct URL, directly pass
+        # the URL to pip, as `req @ url` does not work on the
+        # command line.
+        cmd.append(req.url or str(req))
+        try:
+            subprocess.check_call(cmd)
+        except subprocess.CalledProcessError as e:
+            raise DistutilsError(str(e)) from e
+        wheel = Wheel(glob.glob(os.path.join(tmpdir, '*.whl'))[0])
+        dist_location = os.path.join(eggs_dir, wheel.egg_name())
+        wheel.install_as_egg(dist_location)
+        return metadata.Distribution.at(dist_location + '/EGG-INFO')
+
+
+def strip_marker(req):
+    """
+    Return a new requirement without the environment marker to avoid
+    calling pip with something like `babel; extra == "i18n"`, which
+    would always be ignored.
+    """
+    # create a copy to avoid mutating the input
+    req = packaging.requirements.Requirement(str(req))
+    req.marker = None
+    return req
+
+
+def _warn_wheel_not_available(dist):
+    try:
+        metadata.distribution('wheel')
+    except metadata.PackageNotFoundError:
+        dist.announce('WARNING: The wheel package is not available.', log.WARN)
+
+
+class _DeprecatedInstaller(SetuptoolsDeprecationWarning):
+    _SUMMARY = "setuptools.installer and fetch_build_eggs are deprecated."
+    _DETAILS = """
+    Requirements should be satisfied by a PEP 517 installer.
+    If you are using pip, you can try `pip install --use-pep517`.
+    """
+    _DUE_DATE = 2025, 10, 31
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/launch.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/launch.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d162647d55777d7afa1bf1e44a6c200a3f82419
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/launch.py
@@ -0,0 +1,36 @@
+"""
+Launch the Python script on the command line after
+setuptools is bootstrapped via import.
+"""
+
+# Note that setuptools gets imported implicitly by the
+# invocation of this script using python -m setuptools.launch
+
+import sys
+import tokenize
+
+
+def run() -> None:
+    """
+    Run the script in sys.argv[1] as if it had
+    been invoked naturally.
+    """
+    __builtins__
+    script_name = sys.argv[1]
+    namespace = dict(
+        __file__=script_name,
+        __name__='__main__',
+        __doc__=None,
+    )
+    sys.argv[:] = sys.argv[1:]
+
+    open_ = getattr(tokenize, 'open', open)
+    with open_(script_name) as fid:
+        script = fid.read()
+    norm_script = script.replace('\\r\\n', '\\n')
+    code = compile(norm_script, script_name, 'exec')
+    exec(code, namespace)
+
+
+if __name__ == '__main__':
+    run()
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/logging.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..532da899f7dc02f9fea9a44c429086b98fe043d8
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/logging.py
@@ -0,0 +1,40 @@
+import inspect
+import logging
+import sys
+
+from . import monkey
+
+import distutils.log
+
+
+def _not_warning(record):
+    return record.levelno < logging.WARNING
+
+
+def configure() -> None:
+    """
+    Configure logging to emit warning and above to stderr
+    and everything else to stdout. This behavior is provided
+    for compatibility with distutils.log but may change in
+    the future.
+    """
+    err_handler = logging.StreamHandler()
+    err_handler.setLevel(logging.WARNING)
+    out_handler = logging.StreamHandler(sys.stdout)
+    out_handler.addFilter(_not_warning)
+    handlers = err_handler, out_handler
+    logging.basicConfig(
+        format="{message}", style='{', handlers=handlers, level=logging.DEBUG
+    )
+    if inspect.ismodule(distutils.dist.log):
+        monkey.patch_func(set_threshold, distutils.log, 'set_threshold')
+        # For some reason `distutils.log` module is getting cached in `distutils.dist`
+        # and then loaded again when patched,
+        # implying: id(distutils.log) != id(distutils.dist.log).
+        # Make sure the same module object is used everywhere:
+        distutils.dist.log = distutils.log
+
+
+def set_threshold(level: int) -> int:
+    logging.root.setLevel(level * 10)
+    return set_threshold.unpatched(level)
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/modified.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/modified.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ba02fab68734e1e96fd50d7c4b6ffb1442717fb
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/modified.py
@@ -0,0 +1,18 @@
+try:
+    # Ensure a DistutilsError raised by these methods is the same as distutils.errors.DistutilsError
+    from distutils._modified import (
+        newer,
+        newer_group,
+        newer_pairwise,
+        newer_pairwise_group,
+    )
+except ImportError:
+    # fallback for SETUPTOOLS_USE_DISTUTILS=stdlib, because _modified never existed in stdlib
+    from ._distutils._modified import (
+        newer,
+        newer_group,
+        newer_pairwise,
+        newer_pairwise_group,
+    )
+
+__all__ = ['newer', 'newer_pairwise', 'newer_group', 'newer_pairwise_group']
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/tests/test_namespaces.py b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/tests/test_namespaces.py
new file mode 100644
index 0000000000000000000000000000000000000000..a0f4120bf7900b2118cc066034e036ab7af1798b
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/setuptools/tests/test_namespaces.py
@@ -0,0 +1,138 @@
+import subprocess
+import sys
+
+from setuptools._path import paths_on_pythonpath
+
+from . import namespaces
+
+
+class TestNamespaces:
+    def test_mixed_site_and_non_site(self, tmpdir):
+        """
+        Installing two packages sharing the same namespace, one installed
+        to a site dir and the other installed just to a path on PYTHONPATH
+        should leave the namespace in tact and both packages reachable by
+        import.
+        """
+        pkg_A = namespaces.build_namespace_package(tmpdir, 'myns.pkgA')
+        pkg_B = namespaces.build_namespace_package(tmpdir, 'myns.pkgB')
+        site_packages = tmpdir / 'site-packages'
+        path_packages = tmpdir / 'path-packages'
+        targets = site_packages, path_packages
+        # use pip to install to the target directory
+        install_cmd = [
+            sys.executable,
+            '-m',
+            'pip.__main__',
+            'install',
+            str(pkg_A),
+            '-t',
+            str(site_packages),
+        ]
+        subprocess.check_call(install_cmd)
+        namespaces.make_site_dir(site_packages)
+        install_cmd = [
+            sys.executable,
+            '-m',
+            'pip.__main__',
+            'install',
+            str(pkg_B),
+            '-t',
+            str(path_packages),
+        ]
+        subprocess.check_call(install_cmd)
+        try_import = [
+            sys.executable,
+            '-c',
+            'import myns.pkgA; import myns.pkgB',
+        ]
+        with paths_on_pythonpath(map(str, targets)):
+            subprocess.check_call(try_import)
+
+    def test_pkg_resources_import(self, tmpdir):
+        """
+        Ensure that a namespace package doesn't break on import
+        of pkg_resources.
+        """
+        pkg = namespaces.build_namespace_package(tmpdir, 'myns.pkgA')
+        target = tmpdir / 'packages'
+        target.mkdir()
+        install_cmd = [
+            sys.executable,
+            '-m',
+            'pip',
+            'install',
+            '-t',
+            str(target),
+            str(pkg),
+        ]
+        with paths_on_pythonpath([str(target)]):
+            subprocess.check_call(install_cmd)
+        namespaces.make_site_dir(target)
+        try_import = [
+            sys.executable,
+            '-c',
+            'import pkg_resources',
+        ]
+        with paths_on_pythonpath([str(target)]):
+            subprocess.check_call(try_import)
+
+    def test_namespace_package_installed_and_cwd(self, tmpdir):
+        """
+        Installing a namespace packages but also having it in the current
+        working directory, only one version should take precedence.
+        """
+        pkg_A = namespaces.build_namespace_package(tmpdir, 'myns.pkgA')
+        target = tmpdir / 'packages'
+        # use pip to install to the target directory
+        install_cmd = [
+            sys.executable,
+            '-m',
+            'pip.__main__',
+            'install',
+            str(pkg_A),
+            '-t',
+            str(target),
+        ]
+        subprocess.check_call(install_cmd)
+        namespaces.make_site_dir(target)
+
+        # ensure that package imports and pkg_resources imports
+        pkg_resources_imp = [
+            sys.executable,
+            '-c',
+            'import pkg_resources; import myns.pkgA',
+        ]
+        with paths_on_pythonpath([str(target)]):
+            subprocess.check_call(pkg_resources_imp, cwd=str(pkg_A))
+
+    def test_packages_in_the_same_namespace_installed_and_cwd(self, tmpdir):
+        """
+        Installing one namespace package and also have another in the same
+        namespace in the current working directory, both of them must be
+        importable.
+        """
+        pkg_A = namespaces.build_namespace_package(tmpdir, 'myns.pkgA')
+        pkg_B = namespaces.build_namespace_package(tmpdir, 'myns.pkgB')
+        target = tmpdir / 'packages'
+        # use pip to install to the target directory
+        install_cmd = [
+            sys.executable,
+            '-m',
+            'pip.__main__',
+            'install',
+            str(pkg_A),
+            '-t',
+            str(target),
+        ]
+        subprocess.check_call(install_cmd)
+        namespaces.make_site_dir(target)
+
+        # ensure that all packages import and pkg_resources imports
+        pkg_resources_imp = [
+            sys.executable,
+            '-c',
+            'import pkg_resources; import myns.pkgA; import myns.pkgB',
+        ]
+        with paths_on_pythonpath([str(target)]):
+            subprocess.check_call(pkg_resources_imp, cwd=str(pkg_B))
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/LICENSE b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..1cc22a5aa7679ebaa10934212f356823931bdc3e
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/LICENSE
@@ -0,0 +1,18 @@
+Copyright (c) 2010-2024 Benjamin Peterson
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software is furnished to do so,
+subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/METADATA b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..cfde03c2631c5c1d5cdc0949d0ee3379e7110f0e
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/METADATA
@@ -0,0 +1,43 @@
+Metadata-Version: 2.1
+Name: six
+Version: 1.17.0
+Summary: Python 2 and 3 compatibility utilities
+Home-page: https://github.com/benjaminp/six
+Author: Benjamin Peterson
+Author-email: benjamin@python.org
+License: MIT
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Programming Language :: Python :: 2
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Utilities
+Requires-Python: >=2.7, !=3.0.*, !=3.1.*, !=3.2.*
+License-File: LICENSE
+
+.. image:: https://img.shields.io/pypi/v/six.svg
+   :target: https://pypi.org/project/six/
+   :alt: six on PyPI
+
+.. image:: https://readthedocs.org/projects/six/badge/?version=latest
+   :target: https://six.readthedocs.io/
+   :alt: six's documentation on Read the Docs
+
+.. image:: https://img.shields.io/badge/license-MIT-green.svg
+   :target: https://github.com/benjaminp/six/blob/master/LICENSE
+   :alt: MIT License badge
+
+Six is a Python 2 and 3 compatibility library.  It provides utility functions
+for smoothing over the differences between the Python versions with the goal of
+writing Python code that is compatible on both Python versions.  See the
+documentation for more information on what is provided.
+
+Six supports Python 2.7 and 3.3+.  It is contained in only one Python
+file, so it can be easily copied into your project. (The copyright and license
+notice must be retained.)
+
+Online documentation is at https://six.readthedocs.io/.
+
+Bugs can be reported to https://github.com/benjaminp/six.  The code can also
+be found there.
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/RECORD b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/RECORD
new file mode 100644
index 0000000000000000000000000000000000000000..888d9850e5c78ed5c5e19c3974b6311297fb3bae
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/RECORD
@@ -0,0 +1,8 @@
+__pycache__/six.cpython-310.pyc,,
+six-1.17.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+six-1.17.0.dist-info/LICENSE,sha256=Q3W6IOK5xsTnytKUCmKP2Q6VzD1Q7pKq51VxXYuh-9A,1066
+six-1.17.0.dist-info/METADATA,sha256=ViBCB4wnUlSfbYp8htvF3XCAiKe-bYBnLsewcQC3JGg,1658
+six-1.17.0.dist-info/RECORD,,
+six-1.17.0.dist-info/WHEEL,sha256=pxeNX5JdtCe58PUSYP9upmc7jdRPgvT0Gm9kb1SHlVw,109
+six-1.17.0.dist-info/top_level.txt,sha256=_iVH_iYEtEXnD8nYGQYpYFUvkUW9sEO1GYbkeKSAais,4
+six.py,sha256=xRyR9wPT1LNpbJI8tf7CE-BeddkhU5O--sfy-mo5BN8,34703
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/WHEEL b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/WHEEL
new file mode 100644
index 0000000000000000000000000000000000000000..104f3874635f24f0d2918dfeaf6a59652274460c
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/WHEEL
@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.6.0)
+Root-Is-Purelib: true
+Tag: py2-none-any
+Tag: py3-none-any
+
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/top_level.txt b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/top_level.txt
new file mode 100644
index 0000000000000000000000000000000000000000..ffe2fce498955b628014618b28c6bcf152466a4a
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/six-1.17.0.dist-info/top_level.txt
@@ -0,0 +1 @@
+six
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/wheel/bdist_wheel.py b/external/alphageometry/.venv-ag/Lib/site-packages/wheel/bdist_wheel.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd7b8629e5087fe35a2ea7ab6d3cbbb4cacf0031
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/wheel/bdist_wheel.py
@@ -0,0 +1,26 @@
+from typing import TYPE_CHECKING
+from warnings import warn
+
+warn(
+    "The 'wheel' package is no longer the canonical location of the 'bdist_wheel' "
+    "command, and will be removed in a future release. Please update to setuptools "
+    "v70.1 or later which contains an integrated version of this command.",
+    DeprecationWarning,
+    stacklevel=1,
+)
+
+if TYPE_CHECKING:
+    from ._bdist_wheel import bdist_wheel as bdist_wheel
+else:
+    try:
+        # Better integration/compatibility with setuptools:
+        # in the case new fixes or PEPs are implemented in setuptools
+        # there is no need to backport them to the deprecated code base.
+        # This is useful in the case of old packages in the ecosystem
+        # that are still used but have low maintenance.
+        from setuptools.command.bdist_wheel import bdist_wheel
+    except ImportError:
+        # Only used in the case of old setuptools versions.
+        # If the user wants to get the latest fixes/PEPs,
+        # they are encouraged to address the deprecation warning.
+        from ._bdist_wheel import bdist_wheel as bdist_wheel
diff --git a/external/alphageometry/.venv-ag/Lib/site-packages/wheel/macosx_libfile.py b/external/alphageometry/.venv-ag/Lib/site-packages/wheel/macosx_libfile.py
new file mode 100644
index 0000000000000000000000000000000000000000..abdfc9eda1d4c1f33270c155e1610fe73bd54263
--- /dev/null
+++ b/external/alphageometry/.venv-ag/Lib/site-packages/wheel/macosx_libfile.py
@@ -0,0 +1,482 @@
+"""
+This module contains function to analyse dynamic library
+headers to extract system information
+
+Currently only for MacOSX
+
+Library file on macosx system starts with Mach-O or Fat field.
+This can be distinguish by first 32 bites and it is called magic number.
+Proper value of magic number is with suffix _MAGIC. Suffix _CIGAM means
+reversed bytes order.
+Both fields can occur in two types: 32 and 64 bytes.
+
+FAT field inform that this library contains few version of library
+(typically for different types version). It contains
+information where Mach-O headers starts.
+
+Each section started with Mach-O header contains one library
+(So if file starts with this field it contains only one version).
+
+After filed Mach-O there are section fields.
+Each of them starts with two fields:
+cmd - magic number for this command
+cmdsize - total size occupied by this section information.
+
+In this case only sections LC_VERSION_MIN_MACOSX (for macosx 10.13 and earlier)
+and LC_BUILD_VERSION (for macosx 10.14 and newer) are interesting,
+because them contains information about minimal system version.
+
+Important remarks:
+- For fat files this implementation looks for maximum number version.
+  It not check if it is 32 or 64 and do not compare it with currently built package.
+  So it is possible to false report higher version that needed.
+- All structures signatures are taken form macosx header files.
+- I think that binary format will be more stable than `otool` output.
+  and if apple introduce some changes both implementation will need to be updated.
+- The system compile will set the deployment target no lower than
+  11.0 for arm64 builds. For "Universal 2" builds use the x86_64 deployment
+  target when the arm64 target is 11.0.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import os
+import sys
+from io import BufferedIOBase
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Union
+
+    StrPath = Union[str, os.PathLike[str]]
+
+"""here the needed const and struct from mach-o header files"""
+
+FAT_MAGIC = 0xCAFEBABE
+FAT_CIGAM = 0xBEBAFECA
+FAT_MAGIC_64 = 0xCAFEBABF
+FAT_CIGAM_64 = 0xBFBAFECA
+MH_MAGIC = 0xFEEDFACE
+MH_CIGAM = 0xCEFAEDFE
+MH_MAGIC_64 = 0xFEEDFACF
+MH_CIGAM_64 = 0xCFFAEDFE
+
+LC_VERSION_MIN_MACOSX = 0x24
+LC_BUILD_VERSION = 0x32
+
+CPU_TYPE_ARM64 = 0x0100000C
+
+mach_header_fields = [
+    ("magic", ctypes.c_uint32),
+    ("cputype", ctypes.c_int),
+    ("cpusubtype", ctypes.c_int),
+    ("filetype", ctypes.c_uint32),
+    ("ncmds", ctypes.c_uint32),
+    ("sizeofcmds", ctypes.c_uint32),
+    ("flags", ctypes.c_uint32),
+]
+"""
+struct mach_header {
+    uint32_t	magic;		/* mach magic number identifier */
+    cpu_type_t	cputype;	/* cpu specifier */
+    cpu_subtype_t	cpusubtype;	/* machine specifier */
+    uint32_t	filetype;	/* type of file */
+    uint32_t	ncmds;		/* number of load commands */
+    uint32_t	sizeofcmds;	/* the size of all the load commands */
+    uint32_t	flags;		/* flags */
+};
+typedef integer_t cpu_type_t;
+typedef integer_t cpu_subtype_t;
+"""
+
+mach_header_fields_64 = mach_header_fields + [("reserved", ctypes.c_uint32)]
+"""
+struct mach_header_64 {
+    uint32_t	magic;		/* mach magic number identifier */
+    cpu_type_t	cputype;	/* cpu specifier */
+    cpu_subtype_t	cpusubtype;	/* machine specifier */
+    uint32_t	filetype;	/* type of file */
+    uint32_t	ncmds;		/* number of load commands */
+    uint32_t	sizeofcmds;	/* the size of all the load commands */
+    uint32_t	flags;		/* flags */
+    uint32_t	reserved;	/* reserved */
+};
+"""
+
+fat_header_fields = [("magic", ctypes.c_uint32), ("nfat_arch", ctypes.c_uint32)]
+"""
+struct fat_header {
+    uint32_t	magic;		/* FAT_MAGIC or FAT_MAGIC_64 */
+    uint32_t	nfat_arch;	/* number of structs that follow */
+};
+"""
+
+fat_arch_fields = [
+    ("cputype", ctypes.c_int),
+    ("cpusubtype", ctypes.c_int),
+    ("offset", ctypes.c_uint32),
+    ("size", ctypes.c_uint32),
+    ("align", ctypes.c_uint32),
+]
+"""
+struct fat_arch {
+    cpu_type_t	cputype;	/* cpu specifier (int) */
+    cpu_subtype_t	cpusubtype;	/* machine specifier (int) */
+    uint32_t	offset;		/* file offset to this object file */
+    uint32_t	size;		/* size of this object file */
+    uint32_t	align;		/* alignment as a power of 2 */
+};
+"""
+
+fat_arch_64_fields = [
+    ("cputype", ctypes.c_int),
+    ("cpusubtype", ctypes.c_int),
+    ("offset", ctypes.c_uint64),
+    ("size", ctypes.c_uint64),
+    ("align", ctypes.c_uint32),
+    ("reserved", ctypes.c_uint32),
+]
+"""
+struct fat_arch_64 {
+    cpu_type_t	cputype;	/* cpu specifier (int) */
+    cpu_subtype_t	cpusubtype;	/* machine specifier (int) */
+    uint64_t	offset;		/* file offset to this object file */
+    uint64_t	size;		/* size of this object file */
+    uint32_t	align;		/* alignment as a power of 2 */
+    uint32_t	reserved;	/* reserved */
+};
+"""
+
+segment_base_fields = [("cmd", ctypes.c_uint32), ("cmdsize", ctypes.c_uint32)]
+"""base for reading segment info"""
+
+segment_command_fields = [
+    ("cmd", ctypes.c_uint32),
+    ("cmdsize", ctypes.c_uint32),
+    ("segname", ctypes.c_char * 16),
+    ("vmaddr", ctypes.c_uint32),
+    ("vmsize", ctypes.c_uint32),
+    ("fileoff", ctypes.c_uint32),
+    ("filesize", ctypes.c_uint32),
+    ("maxprot", ctypes.c_int),
+    ("initprot", ctypes.c_int),
+    ("nsects", ctypes.c_uint32),
+    ("flags", ctypes.c_uint32),
+]
+"""
+struct segment_command { /* for 32-bit architectures */
+    uint32_t	cmd;		/* LC_SEGMENT */
+    uint32_t	cmdsize;	/* includes sizeof section structs */
+    char		segname[16];	/* segment name */
+    uint32_t	vmaddr;		/* memory address of this segment */
+    uint32_t	vmsize;		/* memory size of this segment */
+    uint32_t	fileoff;	/* file offset of this segment */
+    uint32_t	filesize;	/* amount to map from the file */
+    vm_prot_t	maxprot;	/* maximum VM protection */
+    vm_prot_t	initprot;	/* initial VM protection */
+    uint32_t	nsects;		/* number of sections in segment */
+    uint32_t	flags;		/* flags */
+};
+typedef int vm_prot_t;
+"""
+
+segment_command_fields_64 = [
+    ("cmd", ctypes.c_uint32),
+    ("cmdsize", ctypes.c_uint32),
+    ("segname", ctypes.c_char * 16),
+    ("vmaddr", ctypes.c_uint64),
+    ("vmsize", ctypes.c_uint64),
+    ("fileoff", ctypes.c_uint64),
+    ("filesize", ctypes.c_uint64),
+    ("maxprot", ctypes.c_int),
+    ("initprot", ctypes.c_int),
+    ("nsects", ctypes.c_uint32),
+    ("flags", ctypes.c_uint32),
+]
+"""
+struct segment_command_64 { /* for 64-bit architectures */
+    uint32_t	cmd;		/* LC_SEGMENT_64 */
+    uint32_t	cmdsize;	/* includes sizeof section_64 structs */
+    char		segname[16];	/* segment name */
+    uint64_t	vmaddr;		/* memory address of this segment */
+    uint64_t	vmsize;		/* memory size of this segment */
+    uint64_t	fileoff;	/* file offset of this segment */
+    uint64_t	filesize;	/* amount to map from the file */
+    vm_prot_t	maxprot;	/* maximum VM protection */
+    vm_prot_t	initprot;	/* initial VM protection */
+    uint32_t	nsects;		/* number of sections in segment */
+    uint32_t	flags;		/* flags */
+};
+"""
+
+version_min_command_fields = segment_base_fields + [
+    ("version", ctypes.c_uint32),
+    ("sdk", ctypes.c_uint32),
+]
+"""
+struct version_min_command {
+    uint32_t	cmd;		/* LC_VERSION_MIN_MACOSX or
+                               LC_VERSION_MIN_IPHONEOS or
+                               LC_VERSION_MIN_WATCHOS or
+                               LC_VERSION_MIN_TVOS */
+    uint32_t	cmdsize;	/* sizeof(struct min_version_command) */
+    uint32_t	version;	/* X.Y.Z is encoded in nibbles xxxx.yy.zz */
+    uint32_t	sdk;		/* X.Y.Z is encoded in nibbles xxxx.yy.zz */
+};
+"""
+
+build_version_command_fields = segment_base_fields + [
+    ("platform", ctypes.c_uint32),
+    ("minos", ctypes.c_uint32),
+    ("sdk", ctypes.c_uint32),
+    ("ntools", ctypes.c_uint32),
+]
+"""
+struct build_version_command {
+    uint32_t	cmd;		/* LC_BUILD_VERSION */
+    uint32_t	cmdsize;	/* sizeof(struct build_version_command) plus */
+                                /* ntools * sizeof(struct build_tool_version) */
+    uint32_t	platform;	/* platform */
+    uint32_t	minos;		/* X.Y.Z is encoded in nibbles xxxx.yy.zz */
+    uint32_t	sdk;		/* X.Y.Z is encoded in nibbles xxxx.yy.zz */
+    uint32_t	ntools;		/* number of tool entries following this */
+};
+"""
+
+
+def swap32(x: int) -> int:
+    return (
+        ((x << 24) & 0xFF000000)
+        | ((x << 8) & 0x00FF0000)
+        | ((x >> 8) & 0x0000FF00)
+        | ((x >> 24) & 0x000000FF)
+    )
+
+
+def get_base_class_and_magic_number(
+    lib_file: BufferedIOBase,
+    seek: int | None = None,
+) -> tuple[type[ctypes.Structure], int]:
+    if seek is None:
+        seek = lib_file.tell()
+    else:
+        lib_file.seek(seek)
+    magic_number = ctypes.c_uint32.from_buffer_copy(
+        lib_file.read(ctypes.sizeof(ctypes.c_uint32))
+    ).value
+
+    # Handle wrong byte order
+    if magic_number in [FAT_CIGAM, FAT_CIGAM_64, MH_CIGAM, MH_CIGAM_64]:
+        if sys.byteorder == "little":
+            BaseClass = ctypes.BigEndianStructure
+        else:
+            BaseClass = ctypes.LittleEndianStructure
+
+        magic_number = swap32(magic_number)
+    else:
+        BaseClass = ctypes.Structure
+
+    lib_file.seek(seek)
+    return BaseClass, magic_number
+
+
+def read_data(struct_class: type[ctypes.Structure], lib_file: BufferedIOBase):
+    return struct_class.from_buffer_copy(lib_file.read(ctypes.sizeof(struct_class)))
+
+
+def extract_macosx_min_system_version(path_to_lib: str):
+    with open(path_to_lib, "rb") as lib_file:
+        BaseClass, magic_number = get_base_class_and_magic_number(lib_file, 0)
+        if magic_number not in [FAT_MAGIC, FAT_MAGIC_64, MH_MAGIC, MH_MAGIC_64]:
+            return
+
+        if magic_number in [FAT_MAGIC, FAT_CIGAM_64]:
+
+            class FatHeader(BaseClass):
+                _fields_ = fat_header_fields
+
+            fat_header = read_data(FatHeader, lib_file)
+            if magic_number == FAT_MAGIC:
+
+                class FatArch(BaseClass):
+                    _fields_ = fat_arch_fields
+
+            else:
+
+                class FatArch(BaseClass):
+                    _fields_ = fat_arch_64_fields
+
+            fat_arch_list = [
+                read_data(FatArch, lib_file) for _ in range(fat_header.nfat_arch)
+            ]
+
+            versions_list: list[tuple[int, int, int]] = []
+            for el in fat_arch_list:
+                try:
+                    version = read_mach_header(lib_file, el.offset)
+                    if version is not None:
+                        if el.cputype == CPU_TYPE_ARM64 and len(fat_arch_list) != 1:
+                            # Xcode will not set the deployment target below 11.0.0
+                            # for the arm64 architecture. Ignore the arm64 deployment
+                            # in fat binaries when the target is 11.0.0, that way
+                            # the other architectures can select a lower deployment
+                            # target.
+                            # This is safe because there is no arm64 variant for
+                            # macOS 10.15 or earlier.
+                            if version == (11, 0, 0):
+                                continue
+                        versions_list.append(version)
+                except ValueError:
+                    pass
+
+            if len(versions_list) > 0:
+                return max(versions_list)
+            else:
+                return None
+
+        else:
+            try:
+                return read_mach_header(lib_file, 0)
+            except ValueError:
+                """when some error during read library files"""
+                return None
+
+
+def read_mach_header(
+    lib_file: BufferedIOBase,
+    seek: int | None = None,
+) -> tuple[int, int, int] | None:
+    """
+    This function parses a Mach-O header and extracts
+    information about the minimal macOS version.
+
+    :param lib_file: reference to opened library file with pointer
+    """
+    base_class, magic_number = get_base_class_and_magic_number(lib_file, seek)
+    arch = "32" if magic_number == MH_MAGIC else "64"
+
+    class SegmentBase(base_class):
+        _fields_ = segment_base_fields
+
+    if arch == "32":
+
+        class MachHeader(base_class):
+            _fields_ = mach_header_fields
+
+    else:
+
+        class MachHeader(base_class):
+            _fields_ = mach_header_fields_64
+
+    mach_header = read_data(MachHeader, lib_file)
+    for _i in range(mach_header.ncmds):
+        pos = lib_file.tell()
+        segment_base = read_data(SegmentBase, lib_file)
+        lib_file.seek(pos)
+        if segment_base.cmd == LC_VERSION_MIN_MACOSX:
+
+            class VersionMinCommand(base_class):
+                _fields_ = version_min_command_fields
+
+            version_info = read_data(VersionMinCommand, lib_file)
+            return parse_version(version_info.version)
+        elif segment_base.cmd == LC_BUILD_VERSION:
+
+            class VersionBuild(base_class):
+                _fields_ = build_version_command_fields
+
+            version_info = read_data(VersionBuild, lib_file)
+            return parse_version(version_info.minos)
+        else:
+            lib_file.seek(pos + segment_base.cmdsize)
+            continue
+
+
+def parse_version(version: int) -> tuple[int, int, int]:
+    x = (version & 0xFFFF0000) >> 16
+    y = (version & 0x0000FF00) >> 8
+    z = version & 0x000000FF
+    return x, y, z
+
+
+def calculate_macosx_platform_tag(archive_root: StrPath, platform_tag: str) -> str:
+    """
+    Calculate proper macosx platform tag basing on files which are included to wheel
+
+    Example platform tag `macosx-10.14-x86_64`
+    """
+    prefix, base_version, suffix = platform_tag.split("-")
+    base_version = tuple(int(x) for x in base_version.split("."))
+    base_version = base_version[:2]
+    if base_version[0] > 10:
+        base_version = (base_version[0], 0)
+    assert len(base_version) == 2
+    if "MACOSX_DEPLOYMENT_TARGET" in os.environ:
+        deploy_target = tuple(
+            int(x) for x in os.environ["MACOSX_DEPLOYMENT_TARGET"].split(".")
+        )
+        deploy_target = deploy_target[:2]
+        if deploy_target[0] > 10:
+            deploy_target = (deploy_target[0], 0)
+        if deploy_target < base_version:
+            sys.stderr.write(
+                "[WARNING] MACOSX_DEPLOYMENT_TARGET is set to a lower value ({}) than "
+                "the version on which the Python interpreter was compiled ({}), and "
+                "will be ignored.\n".format(
+                    ".".join(str(x) for x in deploy_target),
+                    ".".join(str(x) for x in base_version),
+                )
+            )
+        else:
+            base_version = deploy_target
+
+    assert len(base_version) == 2
+    start_version = base_version
+    versions_dict: dict[str, tuple[int, int]] = {}
+    for dirpath, _dirnames, filenames in os.walk(archive_root):
+        for filename in filenames:
+            if filename.endswith(".dylib") or filename.endswith(".so"):
+                lib_path = os.path.join(dirpath, filename)
+                min_ver = extract_macosx_min_system_version(lib_path)
+                if min_ver is not None:
+                    min_ver = min_ver[0:2]
+                    if min_ver[0] > 10:
+                        min_ver = (min_ver[0], 0)
+                    versions_dict[lib_path] = min_ver
+
+    if len(versions_dict) > 0:
+        base_version = max(base_version, max(versions_dict.values()))
+
+    # macosx platform tag do not support minor bugfix release
+    fin_base_version = "_".join([str(x) for x in base_version])
+    if start_version < base_version:
+        problematic_files = [k for k, v in versions_dict.items() if v > start_version]
+        problematic_files = "\n".join(problematic_files)
+        if len(problematic_files) == 1:
+            files_form = "this file"
+        else:
+            files_form = "these files"
+        error_message = (
+            "[WARNING] This wheel needs a higher macOS version than {}  "
+            "To silence this warning, set MACOSX_DEPLOYMENT_TARGET to at least "
+            + fin_base_version
+            + " or recreate "
+            + files_form
+            + " with lower "
+            "MACOSX_DEPLOYMENT_TARGET:  \n" + problematic_files
+        )
+
+        if "MACOSX_DEPLOYMENT_TARGET" in os.environ:
+            error_message = error_message.format(
+                "is set in MACOSX_DEPLOYMENT_TARGET variable."
+            )
+        else:
+            error_message = error_message.format(
+                "the version your Python interpreter is compiled against."
+            )
+
+        sys.stderr.write(error_message)
+
+    platform_tag = prefix + "_" + fin_base_version + "_" + suffix
+    return platform_tag