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