Version: v0.2.0

scatter_plan (type)

Stable type · Evaluate the Jacobian

Tier: stable · Kind: type · Topic: Evaluate the Jacobian

JuliaPythonRust
ScatterPlanScatterPlanScatterPlan

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.

Source

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)

From docs/api/examples/julia/jacobians.jl

Python

ScatterPlan = LaneBlocksPlan | LaneSparseBlocksPlan | ColumnBandedPlan | CsrPlan

No docstring.

Source

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.

Source · Rust docs