Getting started

Differentiate a model, read its sparsity pattern and structure, evaluate the Jacobian, and serialize it as a document block.

EarthSciASTDiff turns an EarthSciAST (.esm) model into its analytical sparse Jacobian by symbolic differentiation. This guide takes one small model through each step. Each step shows the same task in every binding; pick your language on any tab and the rest of the site follows. Every snippet is cut from a program that runs, and checks its own result, in CI.

The examples use bd_chem.esm from tests/valid/ : per-cell chemistry among three species A, B and C on a four-cell column, so the Jacobian is block-diagonal: four blocks of three species.

Differentiate

The Jacobian is a list of bands: each is a rectangular region of the equation’s index space, a symbolic column-index expression, and a coefficient expression. One entry is produced for each (row variable, column variable) pair that occurs in the model, so the derivative is exact and does not depend on the point where you later evaluate it.

Julia

file = EarthSciAST.load_path(joinpath(FIXTURES, "bd_chem.esm"))
view = sysview(file, "Chem")
entries = jacobian_bands(view)
# One entry per (row variable, column variable) pair that occurs in the model.
@assert sort(unique([(en.u, en.v) for en in entries])) ==
        [("A", "A"), ("A", "B"), ("A", "C"), ("B", "A"), ("B", "B"), ("C", "B")]

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

Python

file = earthsci_ast.load_path(FIXTURES / "bd_chem.esm")
entries = earthsci_astdiff.jacobian_bands(file)
# One entry per (row variable, column variable) pair that occurs in the model.
assert sorted({(en.u, en.v) for en in entries}) == [
    ("A", "A"), ("A", "B"), ("A", "C"), ("B", "A"), ("B", "B"), ("C", "B"),
]
view = earthsci_astdiff.sysview(file)

From docs/api/examples/python/jacobians.py

Rust

let file = load_path(fixture("bd_chem.esm"))?;
let entries = jacobian_bands(&file, None, Wrt::States)?;
// One entry per (row variable, column variable) pair that occurs in the model.
let mut pairs: Vec<(&str, &str)> =
    entries.iter().map(|en| (en.u.as_str(), en.v.as_str())).collect();
pairs.sort();
pairs.dedup();
assert_eq!(
    pairs,
    [("A", "A"), ("A", "B"), ("A", "C"), ("B", "A"), ("B", "B"), ("C", "B")]
);
let view = sysview(&file, None)?;

From docs/api/examples/rust/jacobians.rs

See jacobian_bands and sysview .

Sparsity pattern

The pattern is read off structural occurrence, never off numeric probing, so it is valid at every point. Branches of ifelse, min and max are unioned, and entries that happen to be zero at one point stay in the pattern.

Julia

pat = jacobian_pattern(file)
# 4 cells x 6 variable pairs, one stored entry each, even where a value is 0.
@assert nnz(pat.pattern) == 24
@assert pat.pattern[1, 1]                    # A[1] depends on A[1]
@assert !pat.pattern[1, 2]                   # but not on A[2]
@assert pat.rownames[1] == "A[1]"

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

Python

pairs, _ = earthsci_astdiff.jacobian_pattern(file)
# 4 cells x 6 variable pairs, one stored entry each, even where a value is 0.
assert len(pairs) == 24
assert ("A[1]", "B[1]") in pairs
assert ("A[1]", "A[2]") not in pairs

From docs/api/examples/python/jacobians.py

See jacobian_pattern .

Structure and factorization

The bands are classified as block-diagonal, banded or general, which tells a solver what matrix type to use. For block-diagonal systems a symbolic per-cell block LU plan (Markowitz ordering and fill-in) is available too.

Julia

@assert detect_structure(entries, view) == :block_diagonal
plan = plan_factorization(entries, view)
@assert sort(plan.block_vars) == ["A", "B", "C"]

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

Python

assert earthsci_astdiff.detect_structure(entries, view) == "block_diagonal"
plan = earthsci_astdiff.plan_factorization(entries, view)
assert sorted(plan.block_vars) == ["A", "B", "C"]

From docs/api/examples/python/jacobians.py

Rust

let ctx = ctx_of(&view)?;
assert_eq!(detect_structure(&entries, &ctx), "block_diagonal");
let plan = block_diagonal_lu(&entries);
let mut vars = plan.block_vars.clone();
vars.sort();
assert_eq!(vars, ["A", "B", "C"]);

From docs/api/examples/rust/jacobians.rs

See detect_structure and plan_factorization .

Evaluate

The one-shot call evaluates the Jacobian at a point, using the document’s initial state and parameters unless you pass others.

Julia

res = assemble_jacobian(file)                # at the document's initial state
@assert res.structure == :block_diagonal
@assert res.J[1, 1] == -0.04                 # -k1
@assert res.J[5, 5] == -3.0                  # -2 k2 B, with B = 0.5 by default

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

Python

values = earthsci_astdiff.assemble_jacobian(file, states={"A": [2.0] * 4})
assert values[("A[1]", "A[1]")] == -0.04  # -k1
assert values[("B[1]", "B[1]")] == -3.0   # -2 k2 B, with B = 0.5 by default

From docs/api/examples/python/jacobians.py

Rust

let env = Env {
    states: HashMap::from([
        ("A".to_string(), json!([2.0, 2.0, 2.0, 2.0])),
        ("B".to_string(), json!([0.5, 0.5, 0.5, 0.5])),
        ("C".to_string(), json!([0.1, 0.1, 0.1, 0.1])),
    ]),
    params: HashMap::from([
        ("k1".to_string(), json!(0.04)),
        ("k2".to_string(), json!(3.0)),
        ("k3".to_string(), json!(1.0)),
    ]),
    t: 0.0,
};
let values = assemble_jacobian(&file, None, Wrt::States, &env)?;
assert_eq!(values[&("A[1]".to_string(), "A[1]".to_string())], -0.04); // -k1
assert_eq!(values[&("B[1]".to_string(), "B[1]".to_string())], -3.0); // -2 k2 B

From docs/api/examples/rust/jacobians.rs

See assemble_jacobian .

Prepare once, evaluate many times

To use the Jacobian inside a solver, differentiate and compile once, then fill a fixed sparse prototype in place at each step.

Julia

jac = prepare_jacobian(file)                 # differentiate and compile once
J = jac_prototype(jac)                       # sparse matrix on the fixed pattern
u = fill(2.0, 12)
p = jac.p_default                            # the document's parameter values
jac(J, u, p, 0.0)                            # fill J in place, at any (u, p, t)
@assert size(J) == (12, 12)
@assert nnz(J) == 24
@assert J[1, 1] == -0.04

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

See prepare_jacobian .

Serialize

The bands can be written into the document itself as a jacobians block, so any binding that can evaluate .esm expressions can evaluate the Jacobian. The block carries the entries, the structure classification and the factorization plan.

Julia

doc = jacobian_document(file, "Chem")
block = doc["jacobians"]["Chem"]
@assert block["structure"] == "block_diagonal"
@assert block["factorization"]["type"] == "block_diagonal_lu"

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

Python

block = earthsci_astdiff.jacobian_block(file)
assert block["structure"] == "block_diagonal"
assert block["factorization"]["type"] == "block_diagonal_lu"
json.dumps(block)  # a plain, serializable dict

From docs/api/examples/python/jacobians.py

Rust

let block = jacobian_block(&file, None, Wrt::States, true, 12)?;
assert_eq!(block["structure"], "block_diagonal");
assert_eq!(block["factorization"]["type"], "block_diagonal_lu");

From docs/api/examples/rust/jacobians.rs

See jacobian_block and jacobian_document .