scatter_plan (type)
Stable type · Evaluate the Jacobian
Tier: stable · Kind: type · Topic: Evaluate the Jacobian
| Julia | Python | Rust |
|---|---|---|
ScatterPlan | ScatterPlan | ScatterPlan |
Julia
abstract type ScatterPlan
ScatterPlan
Band fields → storage positions for one (band list, layout). Built by
scatter_plan
, serialized by serialize_scatter_plan,
checked against a source by bind_plan
. Every plan carries the
canonical states (name, shape) and the bands (row var, column var, field
shape) it was built from, which is what binding checks.
Example (evaluate)
# EarthSciASTDiff evaluates nothing: it emits ESM documents that EarthSciAST
# evaluates. The band fields J_k, as a derived model whose right-hand side,
# run once, leaves every band value in `du`. `fields[k]` names entry k's field.
doc, fields = evaluation_document(file)
@assert length(fields) == length(entries)
prob = EarthSciAST.esm_problem(
JSON3.read(stable_json(doc)), (0.0, 1.0)
)
du = zero(prob.u0)
prob.f!(du, prob.u0, prob.p, 0.0) # at the document's initial state
k = findfirst(en -> (en.u, en.v) == ("A", "A"), entries)
@assert du[prob.var_map["JacobianEval.$(fields[k])[1]"]] == -0.04 # -k1
# J·v as a derived model: inputs `_v_<C>`, outputs `_jvp_<R>`.
jdoc, (model, v, jvp) = jvp_document(file)
@assert model == "JvpEval"
J = jac_prototype(file) # sparse matrix on the fixed pattern
@assert size(J) == (12, 12) && nnz(J) == 24
# A scatter plan places the band fields in your own storage (here CSR),
# checked against the source before it is used.
bound = bind_plan(file, scatter_plan(file, Csr()))
@assert bound.len == 24
@assert all(1 <= k <= length(fields) for (k, _, _) in bound.pairs)
Python
ScatterPlan = LaneBlocksPlan | LaneSparseBlocksPlan | ColumnBandedPlan | CsrPlan
No docstring.
Rust
pub struct ScatterPlan {
pub states: Vec<(String, Vec<i64>)>,
pub bands: Vec<(String, String, Vec<i64>)>,
pub kind: PlanKind,
}
impl ScatterPlan { pub fn from_file(file: &EsmFile, model_name: Option<&str>, layout: &Layout) -> Result<ScatterPlan, BandError> }
impl ScatterPlan { pub fn build(sv: &SysView<'_>, entries: &[JacEntry], layout: &Layout) -> Result<ScatterPlan, BandError> }
impl ScatterPlan { pub fn to_json(&self) -> Value }
impl ScatterPlan { pub fn lanes(&self) -> usize }
impl ScatterPlan { pub fn out_len(&self, l: usize) -> usize }
impl ScatterPlan { pub fn bind(&self, file: &EsmFile, model_name: Option<&str>, l: usize) -> Result<BoundScatter, BandError> }
impl ScatterPlan { pub fn bind_entries(&self, sv: &SysView<'_>, entries: &[JacEntry], l: usize) -> Result<BoundScatter, BandError> }
impl ScatterPlan { pub fn targets(&self, l: usize) -> Vec<Option<(String, String)>> }
Band fields → storage positions for one (band list, layout). Plain data:
Send + Sync, built once, shared by every leaf using the operator.