From c4345616989307c1ac93535eacc2847ca600729a Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Fri, 14 Jun 2024 12:09:58 -0400 Subject: [PATCH] allow user to specify config values to overwrite in run_stardis function --- stardis/base.py | 12 ++++++++++-- stardis/io/base.py | 25 +++++++++++++++++++++++-- stardis/io/model/__init__.py | 2 +- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/stardis/base.py b/stardis/base.py index 5b3e0729..7016768c 100644 --- a/stardis/base.py +++ b/stardis/base.py @@ -8,7 +8,9 @@ import logging -def run_stardis(config_fname, tracing_lambdas_or_nus): +def run_stardis( + config_fname, tracing_lambdas_or_nus, add_config_keys=None, add_config_vals=None +): """ Runs a STARDIS simulation. @@ -20,6 +22,10 @@ def run_stardis(config_fname, tracing_lambdas_or_nus): Numpy array of the frequencies or wavelengths to calculate the spectrum for. Must have units attached to it, with dimensions of either length or inverse time. + add_config_keys : list, optional + List of additional keys to add or overwrite for the configuration file. + add_config_vals : list, optional + List of corresponding additional values to add to the configuration file. Returns ------- @@ -29,7 +35,9 @@ def run_stardis(config_fname, tracing_lambdas_or_nus): tracing_nus = tracing_lambdas_or_nus.to(u.Hz, u.spectral()) - config, adata, stellar_model = parse_config_to_model(config_fname) + config, adata, stellar_model = parse_config_to_model( + config_fname, add_config_keys, add_config_vals + ) set_num_threads(config.n_threads) stellar_plasma = create_stellar_plasma(stellar_model, adata, config) stellar_radiation_field = create_stellar_radiation_field( diff --git a/stardis/io/base.py b/stardis/io/base.py index 09afd3f2..5f0b4f95 100644 --- a/stardis/io/base.py +++ b/stardis/io/base.py @@ -3,7 +3,7 @@ import numpy as np from tardis.io.atom_data import AtomData -from tardis.io.configuration.config_validator import validate_yaml +from tardis.io.configuration.config_validator import validate_yaml, validate_dict from tardis.io.configuration.config_reader import Configuration from stardis.io.model.marcs import read_marcs_model @@ -15,7 +15,7 @@ SCHEMA_PATH = BASE_DIR / "config_schema.yml" -def parse_config_to_model(config_fname): +def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=None): """ Parses the config and model files and outputs python objects to be passed into run stardis so they can be individually modified in python. @@ -23,6 +23,10 @@ def parse_config_to_model(config_fname): ---------- config_fname : str Filepath to the STARDIS configuration. Must be a YAML file. + add_config_keys : list, optional + List of additional keys to add or overwrite for the configuration file. + add_config_vals : list, optional + List of corresponding additional values to add to the configuration file. Returns ------- @@ -40,6 +44,23 @@ def parse_config_to_model(config_fname): except: raise ValueError("Config failed to validate. Check the config file.") + if ( + not add_config_keys + ): # If a dictionary was passed, update the config with the dictionary + pass + else: + print("Updating config with additional keys and values") + try: + for key, val in zip(add_config_keys, add_config_vals): + config.set_config_item(key, val) + except: + config.set_config_item(add_config_keys, add_config_vals) + + try: + config_dict = validate_dict(config, schemapath=SCHEMA_PATH) + except: + raise ValueError("Additional config keys and values failed to validate.") + adata = AtomData.from_hdf(config.atom_data) # model diff --git a/stardis/io/model/__init__.py b/stardis/io/model/__init__.py index 6212f01d..6ce3ae45 100644 --- a/stardis/io/model/__init__.py +++ b/stardis/io/model/__init__.py @@ -1 +1 @@ -from stardis.io.model.marcs import read_marcs_model +from stardis.io.model.marcs import *