Coverage for src / hallmd / models / thruster.py: 78%
213 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-19 21:56 +0000
1"""Module for Hall thruster models.
3!!! Note
4 Only current implementation is for the 1d fluid [Hallthruster.jl code](https://github.com/UM-PEPL/HallThruster.jl).
5 Other thruster codes can be implemented similarly here.
7Includes:
9- `run_hallthruster_jl` - General wrapper to run HallThruster.jl for a single set of inputs
10- `hallthruster_jl()` - PEM wrapper to run HallThruster.jl for a set of PEM inputs
11- `get_jl_env` - Get the path of the julia environment created for HallThruster.jl for a specific git ref
12- `PEM_TO_JULIA` - Mapping of PEM variable names to a path in the HallThruster.jl input/output structure (defaults)
13"""
15import copy
16import json
17import os
18import platform
19import random
20import string
21import subprocess
22import tempfile
23import time
24import typing
25import warnings
26from importlib import resources
27from pathlib import Path
28from typing import Callable, Optional
30import numpy as np
31from pem_core.constants import AVOGADRO_CONSTANT, FUNDAMENTAL_CHARGE, MOLECULAR_WEIGHTS
32from pem_core.types import Dataset
34from hallmd.utils import load_thruster
36__all__ = ["run_hallthruster_jl", "hallthruster_jl", "get_jl_env", "PEM_TO_JULIA", "JL_BINARY"]
38HALLTHRUSTER_VERSION_DEFAULT = "0.18.7"
40# Maps PEM variable names to a path in the HallThruster.jl input/output structure (default values here)
41with resources.files("hallmd.models").joinpath("pem_to_julia.json").open("r") as fd:
42 PEM_TO_JULIA = json.load(fd)
44assert isinstance(PEM_TO_JULIA, dict)
47def get_jl_binary():
48 """use the juliaup config file to find out where the default julia binary is stored.
49 if juliaup is not installed, returns 'julia'.
50 this lets us avoid running juliaup's wrapper executable, which has led to concurrency issues.
51 This is designed for use on the cluster, so on windows it just returns `julia`."""
53 if platform.system() == "Windows":
54 return "julia"
56 home = os.environ["HOME"]
57 juliaup_dir = Path(f"{home}/.julia/juliaup")
58 juliaup_json = juliaup_dir / "juliaup.json"
59 if not os.path.exists(juliaup_json):
60 alt_launcher = Path(f"{home}/.juliaup/bin/julia")
61 if alt_launcher.exists():
62 return str(alt_launcher)
63 return "julia"
65 with open(juliaup_json, "r") as fd:
66 jl_config = json.load(fd)
68 default = jl_config["Default"]
69 versions = jl_config["InstalledVersions"]
70 channels = jl_config["InstalledChannels"]
72 version = versions[channels[default]["Version"]]
73 if binary_path := version.get("BinaryPath"):
74 return juliaup_dir / Path(binary_path)
76 path = juliaup_dir / version["Path"] / "bin" / "julia"
77 return path
80JL_BINARY = get_jl_binary()
83def get_jl_env(git_ref: str) -> Path:
84 """Get the path of the julia environment created for HallThruster.jl for a specific git ref.
86 :param git_ref: The git ref (i.e. commit hash, version tag, branch, etc.) of HallThruster.jl to use.
87 """
88 global_env_dir = Path("~/.julia/environments/").expanduser()
89 env_path = global_env_dir / f"hallthruster_{git_ref}"
90 return env_path
93def _convert_to_julia(pem_data: Dataset, julia_data: dict, pem_to_julia: dict):
94 """Replace all values in the mutable `julia_data` dict with corresponding values in `pem_data`, using the
95 conversion map provided in `pem_to_julia`. Will "blaze" a path into `julia_data` dict if it does not exist.
97 :param pem_data: thruster inputs of the form `{'pem_variable_name': value}`
98 :param julia_data: a `HallThruster.jl` data structure of the form `{'config': {...}, 'simulation': {...}, etc.}`
99 :param pem_to_julia: a conversion map from a pem variable name to a path in `julia_data`. Use strings for dict keys
100 and integers for list indices. For example, `{'P_b': ['config', 'background_pressure', 2]}`
101 will set `julia_data['config']['background_pressure'][2] = pem_data['P_b']`.
102 """
103 for pem_key, value in pem_data.items():
104 if pem_key not in pem_to_julia:
105 raise KeyError(f"Cannot convert PEM data variable {pem_key} since it is not in the provided conversion map")
107 # Blaze a trail into the julia data structure (set dict defaults for str keys and list defaults for int keys)
108 julia_path = pem_to_julia[pem_key]
109 pointer = julia_data
110 for i, key in enumerate(julia_path[:-1]):
111 if isinstance(pointer, dict) and not pointer.get(key):
112 pointer.setdefault(key, {} if isinstance(julia_path[i + 1], str) else [])
113 if isinstance(pointer, list) and len(pointer) <= key:
114 pointer.extend(
115 [{} if isinstance(julia_path[i + 1], str) else [] for _ in range(key - len(pointer) + 1)]
116 )
117 pointer = pointer[key]
118 pointer[julia_path[-1]] = value
121def _convert_to_pem(julia_data: dict, pem_to_julia: dict):
122 """Return a `dict` of PEM outputs from a `HallThruster.jl` output structure, using the conversion map provided."""
123 pem_data = {}
124 for pem_key, julia_path in pem_to_julia.items():
125 pointer = julia_data
126 if julia_path[0] == "output":
127 found_value = True
128 for key in julia_path:
129 try:
130 pointer = pointer[key]
131 except (KeyError, IndexError):
132 found_value = False
133 break
135 if found_value:
136 pem_data[pem_key] = pointer
137 return pem_data
140def _default_model_fidelity(model_fidelity: tuple, json_config: dict, cfl: float = 0.2) -> dict:
141 """Built-in (default) method to convert model fidelity tuple to `ncells` and `ncharge` via:
143 ```python
144 ncells = 50 * (model_fidelity[0] + 2)
145 ncharge = model_fidelity[1] + 1
146 ```
148 Also adjusts the time step `dt` to maintain the CFL condition (based on grid spacing and ion velocity).
150 :param model_fidelity: tuple of integers that determine the number of cells and the number of charge states to use
151 :param json_config: the current set of configurations for HallThruster.jl
152 :param cfl: a conservative estimate for CFL condition to determine time step
153 :returns: a dictionary of simulation parameters that can be converted to Julia via the `pem_to_julia` mapping,
154 namely `{'num_cells': int, 'ncharge': int, 'dt': float}`
155 """
156 if model_fidelity == ():
157 model_fidelity = (2, 2) # default to high-fidelity model
159 num_cells = 50 * (model_fidelity[0] + 2)
160 ncharge = model_fidelity[1] + 1
162 # Estimate conservative time step based on CFL and grid spacing
163 config = json_config.get("config", {})
164 domain = config.get("domain", [0, 0.08])
165 anode_pot = config.get("discharge_voltage", 300)
166 cathode_pot = config.get("cathode_coupling_voltage", 0)
167 propellant = config.get("propellant", "Xenon")
169 if propellant not in MOLECULAR_WEIGHTS:
170 warnings.warn(
171 f"Could not find propellant {propellant} in `hallmd.utils`. "
172 f"Will default to Xenon for estimating uniform time step from CFL..."
173 )
174 propellant = "Xenon"
176 mi = MOLECULAR_WEIGHTS[propellant] / AVOGADRO_CONSTANT / 1000 # kg
177 dx = float(domain[1]) / (num_cells + 1)
178 u = np.sqrt(2 * ncharge * FUNDAMENTAL_CHARGE * (anode_pot - cathode_pot) / mi)
179 dt_s = cfl * dx / u
181 return {"num_cells": num_cells, "ncharge": ncharge, "dt": float(dt_s)}
184def _format_hallthruster_jl_input(
185 thruster_inputs: Dataset,
186 pem_to_julia: dict,
187 thruster: Path | str | dict = "SPT-100",
188 config: dict | None = None,
189 simulation: dict | None = None,
190 postprocess: dict | None = None,
191 model_fidelity: tuple = (2, 2),
192 output_path: str | Path | None = None,
193 fidelity_function: Callable | None = None,
194) -> dict:
195 """Helper function to format PEM inputs for `Hallthruster.jl` as a `dict` writeable to `json`. See the call
196 signature of `hallthruster_jl` for more details on arguments.
198 The return `dict` has the format:
199 ```json
200 {
201 "config": {
202 "thruster": {...},
203 "discharge_voltage": 300,
204 "domain": [0, 0.08],
205 "anode_mass_flow_rate": 1e-6,
206 "background_pressure": 1e-5,
207 etc.
208 },
209 "simulation": {
210 "ncells": 202,
211 "dt": 8.4e-9
212 "duration": 1e-3,
213 etc.
214 }
215 "postprocess": {...}
216 }
217 ```
219 :returns: a json `dict` in the format that `HallThruster.run_simulation()` expects to be called
220 """
222 json_config = {
223 "config": {} if config is None else copy.deepcopy(config),
224 "simulation": {} if simulation is None else copy.deepcopy(simulation),
225 "postprocess": {} if postprocess is None else copy.deepcopy(postprocess),
226 }
228 # Necessary to load thruster specs separately from the config (to protect sensitive data)
229 if isinstance(thruster, str) or isinstance(thruster, Path):
230 thruster = load_thruster(thruster)
232 if thruster is not None:
233 json_config["config"]["thruster"] = thruster # override
235 # Make sure we request time-averaged
236 duration = json_config["simulation"].get("duration", 1e-3)
237 avg_start_time = json_config["postprocess"].get("average_start_time", 0.5 * duration)
238 json_config["postprocess"]["average_start_time"] = avg_start_time
240 # Update/override config with PEM thruster inputs (modify in place)
241 _convert_to_julia(thruster_inputs, json_config, pem_to_julia)
243 # Override model fidelity quantities
244 if model_fidelity is not None:
245 if fidelity_function is None:
246 fidelity_function = _default_model_fidelity
248 assert isinstance(fidelity_function, Callable)
249 fidelity_overrides = fidelity_function(model_fidelity, json_config)
250 _convert_to_julia(fidelity_overrides, json_config, pem_to_julia)
252 if output_path is not None:
253 fname = "hallthruster_jl"
254 if name := json_config["config"].get("thruster", {}).get("name"):
255 fname += f"_{name}"
256 if vd := json_config["config"].get("discharge_voltage"):
257 fname += f"_{round(vd)}V"
258 if mdot := json_config["config"].get("anode_mass_flow_rate"):
259 fname += f"_{mdot:.1e}kg_s"
261 fname += "_" + "".join(random.choices(string.ascii_uppercase + string.digits, k=4)) + ".json"
262 output_file = str((Path(output_path) / fname).resolve())
263 json_config["postprocess"]["output_file"] = output_file
265 # Handle special conversions for anomalous transport models
266 if anom_model := json_config["config"].get("anom_model"):
267 if anom_model.get("type") in ["LogisticPressureShift", "SimpleLogisticShift"]:
268 anom_model = anom_model.get("model", {})
270 match anom_model.get("type", "TwoZoneBohm"):
271 case "TwoZoneBohm":
272 if thruster_inputs.get("a_2") is not None: # Only when the PEM a_2 is used
273 anom_model["c2"] = anom_model["c2"] * anom_model.get("c1", 0.00625)
274 case "GaussianBohm":
275 if thruster_inputs.get("anom_max") is not None: # Only when the PEM anom_max is used
276 anom_model["hall_max"] = anom_model["hall_max"] * anom_model.get("hall_min", 0.00625)
278 return json_config
281def run_hallthruster_jl(
282 json_input: dict | str | Path, jl_env: Optional[str | Path] = None, jl_script: Optional[str | Path] = None, **kwargs
283) -> dict:
284 """Python wrapper for `HallThruster.run_simulation(json_input)` in Julia.
286 :param json_input: either a dictionary containing `config`, `simulation`, and `postprocess` options for
287 HallThruster.jl, or a string/Path containing a path to a JSON file with those inputs.
288 :param jl_env: The julia environment containing HallThruster.jl. Defaults to global Julia environment.
289 :param jl_script: path to a custom Julia script to run. The script should accept the input json file path as
290 a command line argument. Defaults to just calling `HallThruster.run_simulation(input_file)`.
291 :param kwargs: additional keyword arguments to pass to `subprocess.run` when calling the Julia script.
293 :returns: `dict` of `Hallthruster.jl` outputs. The specific outputs depend on the settings
294 provided in the `postprocess` dict in the input. If `postprocess['output_file']` is present,
295 this function will also write the requested outputs and restart information to that file.
296 """
297 # Read JSON input from file if path provided
298 if isinstance(json_input, dict):
299 _json_dict = json_input
300 else:
301 with open(json_input, "r") as fp:
302 _json_dict = json.load(fp)
304 tempfile_args = {"suffix": ".json", "prefix": "hallthruster_jl_", "mode": "w", "delete": False, "encoding": "utf-8"}
306 # Get output file path. If one not provided, create a temporary
307 temp_out = False
308 if "output_file" in _json_dict.get("postprocess", {}):
309 output_file = Path(_json_dict["postprocess"].get("output_file"))
310 elif "output_file" in _json_dict.get("input", {}).get("postprocess", {}):
311 output_file = Path(_json_dict["input"]["postprocess"].get("output_file"))
312 else:
313 temp_out = True
314 fd_out = tempfile.NamedTemporaryFile(**tempfile_args)
315 output_file = Path(fd_out.name)
316 fd_out.close()
318 if _json_dict.get("input"):
319 _json_dict["input"].setdefault("postprocess", {})
320 _json_dict["input"]["postprocess"]["output_file"] = str(output_file.resolve())
321 else:
322 _json_dict.setdefault("postprocess", {})
323 _json_dict["postprocess"]["output_file"] = str(output_file.resolve())
325 # Dump input to temporary file
326 fd = tempfile.NamedTemporaryFile(**tempfile_args)
327 input_file = fd.name
329 json.dump(_json_dict, fd, ensure_ascii=False, indent=4)
330 fd.close()
332 # Run HallThruster.jl on input file
333 if jl_script is None:
334 cmd = [
335 JL_BINARY,
336 "--startup-file=no",
337 "-e",
338 f'using HallThruster; HallThruster.run_simulation(raw"{input_file}")',
339 ]
340 else:
341 cmd = [JL_BINARY, "--startup-file=no", "--", str(Path(jl_script).resolve()), input_file]
343 if jl_env is not None:
344 if Path(jl_env).exists():
345 cmd.insert(1, f"--project={Path(jl_env).resolve()}")
346 else:
347 raise ValueError(
348 f"Could not find Julia environment {jl_env}. Please create it first. "
349 f"See https://github.com/JANUS-Institute/HallThrusterPEM/blob/main/scripts/install_hallthruster.py"
350 )
352 try:
353 output = subprocess.run(cmd, stderr=subprocess.STDOUT, **kwargs)
354 finally:
355 # Delete temporary input file
356 os.unlink(input_file)
358 if output.returncode != 0:
359 raise ChildProcessError(f"HallThruster.jl exited with nonzero error code:\n{output.stdout}")
361 # Load output data
362 with open(output_file, "r") as fp:
363 output_data = json.load(fp)
365 if temp_out:
366 os.unlink(output_file)
367 if d := output_data.get("postprocess"):
368 if "output_file" in d:
369 del d["output_file"]
370 if d := output_data.get("input"):
371 if d2 := d.get("postprocess"):
372 if "output_file" in d2:
373 del d2["output_file"]
375 return output_data
378def hallthruster_jl(
379 thruster_inputs: Dataset | dict | None = None,
380 thruster: Path | str | dict = "SPT-100",
381 config: Optional[dict] = None,
382 simulation: Optional[dict] = None,
383 postprocess: Optional[dict] = None,
384 model_fidelity: tuple = (2, 2),
385 output_path: Optional[str | Path] = None,
386 version: str = HALLTHRUSTER_VERSION_DEFAULT,
387 pem_to_julia: dict | None = None,
388 fidelity_function: Callable[[tuple[int, ...]], dict] | None = None,
389 julia_script: Optional[str | Path] = None,
390 run_kwargs: dict | None = None,
391 shock_threshold: float | None = None,
392) -> Dataset:
393 """Run a single `HallThruster.jl` simulation for a given set of inputs. This function will write a temporary
394 input file to disk, call `HallThruster.run_simulation()` in Julia, and read the output file back into Python. Will
395 return time-averaged performance metrics and ion velocity for use with the PEM.
397 Note that the specific inputs and outputs described here can be configured using the `pem_to_julia` dict.
399 !!! Warning "Required configuration"
400 You must specify a thruster, a domain, a mass flow rate, and a discharge voltage to run the simulation. The
401 thruster must be defined in the `hallmd.devices` directory or as a dictionary with the required fields.
402 The mass flow rate and discharge voltage are specified in `thruster_inputs` as `mdot_a` (kg/s) and
403 `V_a` (V), respectively. The domain is specified as a list `[left_bound, right_bound]` in the
404 `config` dictionary. See the
405 [HallThruster.jl docs](https://um-pepl.github.io/HallThruster.jl/dev/reference/config/) for more details.
407 :param thruster_inputs: named key-value pairs of thruster inputs: `P_b`, `V_a`, `mdot_a`, `T_e`, `u_n`, `l_t`,
408 `a_1`, `a_2`, `delta_z`, `z0`, `p0`, and `V_cc` for background pressure (Torr), anode
409 voltage, anode mass flow rate (kg/s), electron temperature (eV), neutral velocity (m/s),
410 transition length (m), anomalous transport coefficients, and cathode coupling voltage. Will
411 override the corresponding values in `config` if provided.
412 :param thruster: the name of the thruster to simulate (must be importable from `hallmd.devices`, see
413 [`load_device`][hallmd.utils.load_device]), or a dictionary that provides geometry and
414 magnetic field information of the thruster to simulate; see the
415 [Hallthruster.jl docs](https://um-pepl.github.io/HallThruster.jl/dev/tutorials/simulation//run/).
416 Will override `thruster` in `config` if provided. If None, will defer to `config`.
417 Defaults to the SPT-100.
418 :param config: dictionary of configs for `HallThruster.jl`, see the
419 [Hallthruster.jl docs](https://um-pepl.github.io/HallThruster.jl/dev/reference/config/) for
420 options and formatting.
421 :param simulation: dictionary of simulation parameters for `HallThruster.jl`
422 :param postprocess: dictionary of post-processing parameters for `Hallthruster.jl`
423 :param model_fidelity: tuple of integers that determine the number of cells and the number of charge states to use
424 via `ncells = model_fidelity[0] * 50 + 100` and `ncharge = model_fidelity[1] + 1`.
425 Will override `ncells` and `ncharge` in `simulation` and `config` if provided.
426 :param output_path: base path to save output files, will write to current directory if not specified
427 :param version: version of HallThruster.jl to use; will
428 search for a global `hallthruster_{version}` environment in the `~/.julia/environments/` directory.
429 Can also specify a specific git ref (i.e. branch, commit hash, etc.) to use from GitHub. If the
430 `hallthruster_{version}` environment does not exist, an error will be raised -- you should create
431 this environment first before using it.
432 :param pem_to_julia: a `dict` mapping of PEM shorthand variable names to a list of keys that maps into the
433 `HallThruster.jl` input/output data structure. Defaults to the provided PEM_TO_JULIA dict
434 defined in [`hallmd.models.thruster`][hallmd.models.thruster]. For example,
435 `{'P_b': ['config', 'background_pressure']}` will set `config['background_pressure'] = P_b`.
436 If specified, will override and extend the default mapping.
437 :param fidelity_function: a callable that takes a tuple of integers and returns a dictionary of simulation
438 parameters. Defaults to `_default_model_fidelity` which sets `ncells` and `ncharge` based
439 on the input tuple. The returned simulation parameters must be convertable to Julia via
440 the `pem_to_julia` mapping. The callable should also take in the current json config dict.
441 :param julia_script: path to a custom Julia script to run. The script should accept the input json file path as
442 a command line argument. Defaults to just calling `HallThruster.run_simulation(input_file)`.
443 :param run_kwargs: additional keyword arguments to pass to `subprocess.run` when calling the Julia script.
444 Defaults to `check=True`.
445 :param shock_threshold: if provided, an error will be raised if the ion velocity reaches a maximum before this
446 threshold axial location (in m) - used to detect and filter unwanted "shock-like" behavior,
447 for example by providing a threshold of half the domain length. If not provided, then no
448 filtering is performed (default).
449 :returns: `dict` of `Hallthruster.jl` outputs: `I_B0`, `I_d`, `T`, `eta_c`, `eta_m`, `eta_v`, and `u_ion` for ion
450 beam current (A), discharge current (A), thrust (N), current efficiency, mass efficiency, voltage
451 efficiency, and singly-charged ion velocity profile (m/s), all time-averaged.
452 """
453 if pem_to_julia is None:
454 _pem_to_julia = copy.deepcopy(PEM_TO_JULIA)
455 else:
456 tmp = copy.deepcopy(PEM_TO_JULIA)
457 tmp.update(pem_to_julia)
458 _pem_to_julia = tmp
460 thruster_inputs = typing.cast(Dataset, {} if thruster_inputs is None else thruster_inputs)
462 # Format PEM inputs for HallThruster.jl
463 json_data = _format_hallthruster_jl_input(
464 thruster_inputs,
465 thruster=thruster,
466 config=config,
467 simulation=simulation,
468 postprocess=postprocess,
469 model_fidelity=model_fidelity,
470 output_path=output_path,
471 pem_to_julia=_pem_to_julia,
472 fidelity_function=fidelity_function,
473 )
474 # Get julia environment
475 jl_environment = get_jl_env(version) if version is not None else None
477 if run_kwargs is None:
478 _run_kwargs = {"check": True}
479 else:
480 _run_kwargs = run_kwargs
482 # Run Julia
483 t1 = time.time()
484 sim_results = run_hallthruster_jl(json_data, jl_env=jl_environment, jl_script=julia_script, **_run_kwargs)
485 t2 = time.time()
487 # Format QOIs for PEM
488 thruster_outputs = _convert_to_pem(sim_results, _pem_to_julia)
490 # Raise an exception if thrust or beam current are negative (non-physical cases)
491 thrust = thruster_outputs.get("T", 0)
492 beam_current = thruster_outputs.get("I_B0", 0)
493 if thrust < 0 or beam_current < 0:
494 raise ValueError(f"Exception due to non-physical case: thrust={thrust} N, beam current={beam_current} A")
496 # Raise an exception for ion velocity that exhibits "shock" behavior
497 if shock_threshold is not None:
498 z_coords = thruster_outputs.get("u_ion_coords")
499 ion_velocity = thruster_outputs.get("u_ion")
500 if z_coords is not None and ion_velocity is not None:
501 if (z_max := z_coords[np.argmax(ion_velocity)]) < shock_threshold:
502 raise ValueError(f"Exception due to shock-like behavior: max ion velocity occurs at z={z_max:.3f} m")
504 thruster_outputs["model_cost"] = t2 - t1 # seconds
506 if output_path is not None:
507 output_file = Path(json_data["postprocess"].get("output_file"))
508 thruster_outputs["output_path"] = output_file.relative_to(Path(output_path).resolve()).as_posix()
510 thruster_outputs["thruster_output"] = sim_results
512 return typing.cast(Dataset, thruster_outputs)