JANUS-Institute / HallThrusterPEM

Prototype of a predictive engineering model (PEM) of a Hall thruster, integrating physics-based sub-models with uncertainty quantification.
https://janus-institute.github.io/HallThrusterPEM/
GNU General Public License v3.0
2 stars 1 forks source link

Interface with HallThruster.jl #11

Open archermarx opened 2 days ago

archermarx commented 2 days ago

We're looking to improve interop with HallThruster.jl, both for usability and efficiency and to reduce the amount of special handling the PEM itself needs to do.

Usability

We currently use JuliaCall for calling HallThruster.jl from python, but this is a pain. Instead, we would like to use the code from a script like we do the other modules, i.e.. julia run_hallthruster.jl.

To accomplish this, some amount of startup latency needs to be ameliorated. I am working on this on the HallThruster.jl side. To improve package load time, we can reduce the number of dependencies (for instance, Unitful and JLD2 impose a decent startup penalty) as well as reducing the amount of code in HallThruster.jl itself. We also need to reduce the "time to first X" (TTFX), which is the time between loading the package and generating a simulation result. This can be reduced by adding more code to precompile statements in HallThruster.jl, eliminating some parameterizations of the Config struct so we don't need to specialize as much, and simplifying program logic and better containing specialization.

Interop via JSON

The current data-exchange format between these models is a JSON, but this interface is pretty ad-hoc and requires a lot of special handling to map onto HallThruster.jl concepts. We'd like to update the JSON format to map cleanly onto the Config struct in HallThruster.jl with minimal, if any, parameter conversions needed.

Additionally, some of the HallThruster.jl interfaces should be cleaned up to map better onto JSON files. These include:

The JSON file would have the following form (key names subject to change)

{
    "config": { ... },
    "simulation" : {...},
    "postprocess": {...}, // optional
    "output": {...} // optional, generated by write_to_json
}

Alternatively, we could use this form:

{
    "input": {
        "config": {...},
        "simulation": {...},
        "postprocess": {...},
    }
    "output": {...}
}

which may be cleaner when reading json files written by HallThruster.jl

Postprocessing

Postprocessing options should be specifiable (optionally) in the JSON input. These options should include things like the following:

"postprocess": {
    "metrics": ["thrust", "discharge_current", "mass_efficiency"],
    "average_start_time": 2e-4,
    "output_file": "myfile.json",
    "save_time_resolved": true
}

Time-averaged output will be computed starting at time "average_start_time" and saved to "time_average_file"

Running simulations

The arguments to run_simulation should be separated into a Config (geometry and operating conditions), an object describing parameters for this specific run (number of cells, timestepping options, checkpointing/saving options), and an object describing postprocessing. This should be reflected in the JSON. The names of these fields should be something like config, simulation, and postprocess, i.e.

run_simulation(config::Config; simulation, postprocess, kwargs...)

where postprocess is optional, and kwargs... specify additional keyword arguments. These arguments will override any equivalent args in the simulation argument, which allows for

run_simulation(config; simulation, postprocess, adaptive = false, verbose = true) # override `adaptive` and `verbose` in `simulation`

The simulation struct can possibly be made optional as well, to be filled in with keyword args. This would match the current interface.

run_simulation should have another method with the signature

run_simulation(json_file; kwargs...)

In this form, the JSON file is expected to specify all arguments, but the kwargs act as above and can override things in the simulation section of the JSON file. This method would simply unpack the JSON file into a config, simulation, and postprocess structs and pass these to the first version.

Reproducability and restarts

In addition to the above, we would like simulation outputs written to JSON to contain the inputs used to generate them to enable straightforward restarts. The current JSON format produced by write_to_json should be sufficient but the inputs used to run the simulation should be included as well, so the format of the input JSONs and output JSONs is identical, i.e.

"input": {...}
"output": {[
    {"time": 0.0, "thrust": ..., "ui_1_1": [...], ...},
    {"time": 1e-8, ...}
]}

If "restart"=true is present in the simulation section, then HallThruster.jl will initialize itself using the last frame in the output section.

Output structure:

The output should contain a return code and error string. If

"output": {
    "retcode": "success",
    "error": "",
    "frames": [...]
}
"output": {
    "retcode": "failure",
    "error": "NaN detected in variable ...",
    "frames": [...]
}
"output": {
    "retcode": "error",
    "error": "ArgumentError:(...)",
    "frames": [] no data saved on error cases
}

This issue will be updated as work on this is done. Parallel issues in the HallThruster.jl repo will be linked here as well.

eckelsjd commented 2 days ago

With this new json format, on the PEM side, the wrapper Python function will look roughly like:

def hallthruster_jl_wrapper(thruster_inputs, config, ...):
    json_data = _format_for_hallthruster_jl(thruster_inputs, config, ...)  # will return {"config": {...}, "simulation": {...}, "postprocess": {...}}
   write_to_file(json_data, 'thruster_inputs.json')

   subprocess.run(['julia', 'run_hallthruster.jl', '--input', 'thruster_inputs.json'])

   output_data = read_from_file('thruster_outputs.json')  # written by the Julia script
   thruster_outputs = _format_for_pem(output_data)  # extract QoIs for the PEM like thrust, discharge current, etc.

   return thruster_outputs  

And the Julia script:

# run_hallthruster.jl

using HallThruster
infile = ARGS[1]
outfile_timeresolved = ARGS[2]
outfile_timeaveraged = ARGS[3]

sol = HallThruster.run_simulation(infile)
avg = HallThruster.time_average(sol)
HallThruster.write_to_json(outfile_timeresolved, sol)
HallThruster.write_to_json(outfile_timeaveraged, sol)

I am adding the dev-amisc branch, as this is closely associated with new amisc configurations. For example, with this new json format, an amisc config file will now be streamlined into one yaml file:

name: HallThrusterPEMv0
components:
    - name: Cathode
      model: !!python/name:hallmd.models.cathode.cathode_coupling
    - name: Thruster
      model: !!python/name:hallmd.models.thruster.hallthruster_jl_wrapper
      config: 
          thruster: SPT-100
          discharge_voltage: 300  # etc.
      simulation:
          dt: 1e-9
          ncells: 200
      postprocess:
          metrics: ["thrust", "discharge_current"]
          average_start_time: 1e-4
    - name: Plume
      model: !!python/name:hallmd.models.plume.current_density

During surrogate training, all the Hallthruster.jl configs will be passed along from this file. Also, the new wrapper function will now be easily parallelizable with the amisc.Component.call_model(executor=...) interface.