1use 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
29pub struct FemmFunctionLib;
35
36impl FemmFunctionLib {
37 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}