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