Skip to main content

sim_lib_femm_function/
exports.rs

1//! Library registration that exposes FEMM callables to the runtime.
2//!
3//! Defines the `Lib` that installs the FEMM function exports and the built-in
4//! fixture models so the runtime can call them by name.
5
6use std::{any::Any, sync::Arc};
7
8use sim_kernel::{
9    AbiVersion, Args, Callable, ClassRef, Cx, DefaultFactory, Dependency, Error, Export, Expr,
10    Factory, Lib, LibManifest, LibTarget, Linker, Object, RawArgs, Result as KernelResult, Symbol,
11    Value, Version,
12};
13use sim_lib_femm_core::{FemmLimits, ParamSet};
14use sim_lib_femm_field::Projection;
15use sim_lib_femm_fixtures::{
16    air_core_solenoid, field_as_number_line_integration, gapped_ei_core_inductor,
17    parallel_plate_capacitor, plunger_actuator_ode, slab_heat_conductor,
18    uniform_conductor_resistance,
19};
20use sim_lib_femm_mesh::FemmModel;
21use sim_lib_femm_post::QuantitySpec;
22use sim_lib_femm_sensitiv::{
23    GradientAnswer, gradient_answer, gradient_trust_label, register_femm_adjoint,
24};
25
26use crate::model_value::{ModelValue, model_value};
27use crate::{FemmCall, FemmCallable, ModelCallable, OutputQuery, femm_as_func};
28
29/// The runtime library that installs the FEMM function exports.
30///
31/// Registers `femm/model`, `femm/eval`, `femm/as-func`, `femm/field`, and
32/// `femm/grad` as callables so the runtime can build, evaluate, and
33/// differentiate models by name. See the [crate README](index.html).
34pub struct FemmFunctionLib;
35
36impl FemmFunctionLib {
37    /// Creates the library installer.
38    pub fn new() -> Self {
39        Self
40    }
41}
42
43impl Default for FemmFunctionLib {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl Lib for FemmFunctionLib {
50    fn manifest(&self) -> LibManifest {
51        LibManifest {
52            id: Symbol::qualified("femm", "function"),
53            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
54            abi: AbiVersion { major: 0, minor: 1 },
55            target: LibTarget::HostRegistered,
56            requires: vec![Dependency {
57                id: Symbol::qualified("femm", "field"),
58                minimum_version: None,
59            }],
60            capabilities: Vec::new(),
61            exports: function_symbols()
62                .into_iter()
63                .map(|symbol| Export::Function {
64                    symbol,
65                    function_id: None,
66                })
67                .collect(),
68        }
69    }
70
71    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> KernelResult<()> {
72        register_femm_adjoint()?;
73        for symbol in function_symbols() {
74            linker.function_value(
75                symbol.clone(),
76                DefaultFactory.opaque(Arc::new(FemmFunctionValue { symbol }))?,
77            )?;
78        }
79        Ok(())
80    }
81}
82
83fn function_symbols() -> Vec<Symbol> {
84    vec![
85        Symbol::qualified("femm", "model"),
86        Symbol::qualified("femm", "eval"),
87        Symbol::qualified("femm", "as-func"),
88        Symbol::qualified("femm", "field"),
89        Symbol::qualified("femm", "grad"),
90    ]
91}
92
93#[derive(Clone)]
94struct FemmFunctionValue {
95    symbol: Symbol,
96}
97
98impl Object for FemmFunctionValue {
99    fn display(&self, _cx: &mut Cx) -> KernelResult<String> {
100        Ok(format!("#<function {}>", self.symbol))
101    }
102
103    fn as_any(&self) -> &dyn Any {
104        self
105    }
106}
107
108impl sim_kernel::ObjectCompat for FemmFunctionValue {
109    fn class(&self, cx: &mut Cx) -> KernelResult<ClassRef> {
110        if let Some(class) = cx
111            .registry()
112            .class_by_symbol(&Symbol::qualified("core", "Function"))
113        {
114            return Ok(class.clone());
115        }
116        DefaultFactory.class_stub(
117            sim_kernel::CORE_FUNCTION_CLASS_ID,
118            Symbol::qualified("core", "Function"),
119        )
120    }
121    fn as_expr(&self, _cx: &mut Cx) -> KernelResult<Expr> {
122        Ok(Expr::Symbol(self.symbol.clone()))
123    }
124    fn as_callable(&self) -> Option<&dyn Callable> {
125        Some(self)
126    }
127}
128
129impl Callable for FemmFunctionValue {
130    fn call(&self, cx: &mut Cx, args: Args) -> KernelResult<Value> {
131        match self.symbol.to_string().as_str() {
132            "femm/model" => call_model(cx, args.into_vec()),
133            "femm/eval" => call_eval(cx, args.into_vec()),
134            "femm/as-func" => call_as_func(cx, args.into_vec()),
135            "femm/field" => call_field(cx, args.into_vec()),
136            "femm/grad" => call_grad(cx, args.into_vec()),
137            _ => Err(Error::Eval(format!(
138                "Unknown FEMM function {}",
139                self.symbol
140            ))),
141        }
142    }
143
144    fn call_exprs(&self, cx: &mut Cx, args: RawArgs) -> KernelResult<Value> {
145        let values = args
146            .into_exprs()
147            .into_iter()
148            .map(|expr| cx.eval_expr(expr))
149            .collect::<KernelResult<Vec<_>>>()?;
150        self.call(cx, Args::new(values))
151    }
152}
153
154fn call_model(cx: &mut Cx, args: Vec<Value>) -> KernelResult<Value> {
155    let model = match args.as_slice() {
156        [] => parallel_plate_capacitor(),
157        [name] => example_model(symbolish_or_string(cx, name)?.as_str())
158            .ok_or_else(|| Error::Eval("unknown FEMM example model".to_owned()))?,
159        _ => {
160            return Err(Error::Eval(
161                "femm/model expects zero or one example name".to_owned(),
162            ));
163        }
164    };
165    cx.factory().opaque(Arc::new(model_value(model)))
166}
167
168fn call_eval(cx: &mut Cx, args: Vec<Value>) -> KernelResult<Value> {
169    let [model, query, params] = args.as_slice() else {
170        return Err(Error::Eval(
171            "femm/eval expects model, query, params".to_owned(),
172        ));
173    };
174    let model = model_from_value(model)?;
175    let query = scalar_query_from_value(cx, query)?;
176    let params = params_from_value(cx, params)?;
177    ModelCallable { model }
178        .eval(
179            cx,
180            FemmCall {
181                params,
182                query,
183                want_grad: None,
184                limits: FemmLimits::default(),
185            },
186        )
187        .map(|out| out.value)
188        .map_err(Error::from)
189}
190
191fn call_as_func(cx: &mut Cx, args: Vec<Value>) -> KernelResult<Value> {
192    let [model, vars, query] = args.as_slice() else {
193        return Err(Error::Eval(
194            "femm/as-func expects model, vars, query".to_owned(),
195        ));
196    };
197    let model = model_from_value(model)?;
198    let vars = symbol_list_from_value(cx, vars)?;
199    let query = scalar_query_from_value(cx, query)?;
200    cx.factory()
201        .opaque(Arc::new(femm_as_func(model, vars, query)))
202}
203
204fn call_field(cx: &mut Cx, args: Vec<Value>) -> KernelResult<Value> {
205    let [model, projection, params] = args.as_slice() else {
206        return Err(Error::Eval(
207            "femm/field expects model, projection, params".to_owned(),
208        ));
209    };
210    let model = model_from_value(model)?;
211    let projection = projection_from_value(cx, projection)?;
212    let params = params_from_value(cx, params)?;
213    ModelCallable { model }
214        .eval(
215            cx,
216            FemmCall {
217                params,
218                query: OutputQuery::Field(projection),
219                want_grad: None,
220                limits: FemmLimits::default(),
221            },
222        )
223        .map(|out| out.value)
224        .map_err(Error::from)
225}
226
227fn call_grad(cx: &mut Cx, args: Vec<Value>) -> KernelResult<Value> {
228    let [model, query, wrt, params] = args.as_slice() else {
229        return Err(Error::Eval(
230            "femm/grad expects model, query, wrt, params".to_owned(),
231        ));
232    };
233    let model = model_from_value(model)?;
234    let query = scalar_query_from_value(cx, query)?;
235    let wrt = symbol_list_from_value(cx, wrt)?;
236    let params = params_from_value(cx, params)?;
237    let answer =
238        gradient_answer(cx, &ModelCallable { model }, query, params, &wrt).map_err(Error::from)?;
239    gradient_answer_value(cx, answer)
240}
241
242fn gradient_answer_value(cx: &mut Cx, answer: GradientAnswer) -> KernelResult<Value> {
243    let gradient = cx.factory().list(
244        answer
245            .values
246            .into_iter()
247            .map(|(symbol, value)| {
248                cx.factory().list(vec![
249                    cx.factory().symbol(symbol)?,
250                    cx.factory()
251                        .number_literal(Symbol::qualified("numbers", "f64"), value.to_string())?,
252                ])
253            })
254            .collect::<KernelResult<Vec<_>>>()?,
255    )?;
256    let trust = cx
257        .factory()
258        .string(gradient_trust_label(&answer.trust).to_owned())?;
259    cx.factory().table(vec![
260        (Symbol::new("gradient"), gradient),
261        (Symbol::new("trust"), trust),
262    ])
263}
264
265fn model_from_value(value: &Value) -> KernelResult<FemmModel> {
266    value
267        .object()
268        .downcast_ref::<ModelValue>()
269        .map(|model| model.model.clone())
270        .ok_or_else(|| Error::Eval("expected FEMM model value".to_owned()))
271}
272
273fn example_model(name: &str) -> Option<FemmModel> {
274    Some(match name {
275        "parallel-plate-capacitor" => parallel_plate_capacitor(),
276        "slab-heat-conductor" => slab_heat_conductor(),
277        "uniform-conductor-resistance" => uniform_conductor_resistance(),
278        "air-core-solenoid" => air_core_solenoid(),
279        "gapped-ei-core-inductor" => gapped_ei_core_inductor(),
280        "plunger-actuator-ode" => plunger_actuator_ode(),
281        "field-as-number-line-integration" => field_as_number_line_integration(),
282        _ => return None,
283    })
284}
285
286fn symbolish_or_string(cx: &mut Cx, value: &Value) -> KernelResult<String> {
287    match value.object().as_expr(cx)? {
288        Expr::Symbol(symbol) => Ok(symbol.to_string()),
289        Expr::String(text) => Ok(text),
290        Expr::Quote { expr, .. } => match *expr {
291            Expr::Symbol(symbol) => Ok(symbol.to_string()),
292            _ => Err(Error::Eval("expected symbol or string".to_owned())),
293        },
294        _ => Err(Error::Eval("expected symbol or string".to_owned())),
295    }
296}
297
298fn symbol_list_from_value(cx: &mut Cx, value: &Value) -> KernelResult<Vec<Symbol>> {
299    match value.object().as_expr(cx)? {
300        Expr::List(items) | Expr::Vector(items) => items
301            .into_iter()
302            .map(expr_to_symbol)
303            .collect::<KernelResult<Vec<_>>>(),
304        _ => Err(Error::Eval("expected symbol list".to_owned())),
305    }
306}
307
308fn expr_to_symbol(expr: Expr) -> KernelResult<Symbol> {
309    match expr {
310        Expr::Symbol(symbol) => Ok(symbol),
311        Expr::Quote { expr, .. } => match *expr {
312            Expr::Symbol(symbol) => Ok(symbol),
313            _ => Err(Error::Eval("expected quoted symbol".to_owned())),
314        },
315        _ => Err(Error::Eval("expected symbol".to_owned())),
316    }
317}
318
319fn params_from_value(cx: &mut Cx, value: &Value) -> KernelResult<ParamSet> {
320    match value.object().as_expr(cx)? {
321        Expr::Map(entries) => Ok(ParamSet::new(
322            entries
323                .into_iter()
324                .map(|(key, value_expr)| Ok((expr_to_symbol(key)?, cx.eval_expr(value_expr)?)))
325                .collect::<KernelResult<Vec<_>>>()?,
326        )),
327        Expr::List(items) | Expr::Vector(items) => Ok(ParamSet::new(
328            items
329                .into_iter()
330                .map(|item| match item {
331                    Expr::List(pair) | Expr::Vector(pair) if pair.len() == 2 => Ok((
332                        expr_to_symbol(pair[0].clone())?,
333                        cx.eval_expr(pair[1].clone())?,
334                    )),
335                    _ => Err(Error::Eval(
336                        "expected [symbol value] param entry".to_owned(),
337                    )),
338                })
339                .collect::<KernelResult<Vec<_>>>()?,
340        )),
341        Expr::Nil => Ok(ParamSet::default()),
342        _ => Err(Error::Eval(
343            "expected parameter table or pair list".to_owned(),
344        )),
345    }
346}
347
348fn scalar_query_from_value(cx: &mut Cx, value: &Value) -> KernelResult<OutputQuery> {
349    Ok(OutputQuery::Quantity(QuantitySpec::Custom {
350        name: Symbol::new("q"),
351        expr: value.object().as_expr(cx)?,
352    }))
353}
354
355fn projection_from_value(cx: &mut Cx, value: &Value) -> KernelResult<Projection> {
356    match symbolish_or_string(cx, value)?.as_str() {
357        "potential" => Ok(Projection::Potential),
358        "bx" => Ok(Projection::Bx),
359        "by" => Ok(Projection::By),
360        "bmag" => Ok(Projection::Bmag),
361        "ex" => Ok(Projection::Ex),
362        "ey" => Ok(Projection::Ey),
363        "emag" => Ok(Projection::Emag),
364        "heat-flux-mag" => Ok(Projection::HeatFluxMag),
365        other => Err(Error::Eval(format!("unknown FEMM projection {other}"))),
366    }
367}