Models

LUMEBaseModel

Bases: BaseModel, ABC

Abstract base class for models using lume-model variables.

Inheriting classes must define the evaluate method and variable names must be unique (respectively). Models build using this framework will be compatible with the lume-epics EPICS server and associated tools.

Attributes:

Name	Type	Description
input_variables	list[ScalarVariable]	List defining the input variables and their order.
output_variables	list[ScalarVariable]	List defining the output variables and their order.
input_validation_config	Optional[dict[str, ConfigEnum]]	Determines the behavior during input validation by specifying the validation config for each input variable: {var_name: value}. Value can be "warn", "error", or None.
output_validation_config	Optional[dict[str, ConfigEnum]]	Determines the behavior during output validation by specifying the validation config for each output variable: {var_name: value}. Value can be "warn", "error", or None.

Source code in lume_model/base.py

class LUMEBaseModel(BaseModel, ABC):
    """Abstract base class for models using lume-model variables.

    Inheriting classes must define the evaluate method and variable names must be unique (respectively).
    Models build using this framework will be compatible with the lume-epics EPICS server and associated tools.

    Attributes:
        input_variables: List defining the input variables and their order.
        output_variables: List defining the output variables and their order.
        input_validation_config: Determines the behavior during input validation by specifying the validation
          config for each input variable: {var_name: value}. Value can be "warn", "error", or None.
        output_validation_config: Determines the behavior during output validation by specifying the validation
          config for each output variable: {var_name: value}. Value can be "warn", "error", or None.
    """

    input_variables: list[ScalarVariable]
    output_variables: list[ScalarVariable]
    input_validation_config: Optional[dict[str, ConfigEnum]] = None
    output_validation_config: Optional[dict[str, ConfigEnum]] = None

    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)

    @field_validator("input_variables", "output_variables", mode="before")
    def validate_input_variables(cls, value):
        new_value = []
        if isinstance(value, dict):
            for name, val in value.items():
                if isinstance(val, dict):
                    variable_class = get_variable(val["variable_class"])
                    new_value.append(variable_class(name=name, **val))
                elif isinstance(val, ScalarVariable):
                    new_value.append(val)
                else:
                    raise TypeError(f"type {type(val)} not supported")
        elif isinstance(value, list):
            new_value = value
        return new_value

    def __init__(self, *args, **kwargs):
        """Initializes LUMEBaseModel.

        Args:
            *args: Accepts a single argument which is the model configuration as dictionary, YAML or JSON
              formatted string or file path.
            **kwargs: See class attributes.
        """
        if len(args) == 1:
            if len(kwargs) > 0:
                raise ValueError(
                    "Cannot specify YAML string and keyword arguments for LUMEBaseModel init."
                )
            super().__init__(**parse_config(args[0], self.model_fields))
        elif len(args) > 1:
            raise ValueError(
                "Arguments to LUMEBaseModel must be either a single YAML string "
                "or keyword arguments passed directly to pydantic."
            )
        else:
            super().__init__(**kwargs)

    @field_validator("input_variables", "output_variables")
    def unique_variable_names(cls, value):
        verify_unique_variable_names(value)
        return value

    @field_validator("input_variables")
    def verify_input_default_value(cls, value):
        """Verifies that input variables have the required default values."""
        for var in value:
            if var.default_value is None:
                raise ValueError(
                    f"Input variable {var.name} must have a default value."
                )
        return value

    @property
    def input_names(self) -> list[str]:
        return [var.name for var in self.input_variables]

    @property
    def output_names(self) -> list[str]:
        return [var.name for var in self.output_variables]

    @property
    def default_input_validation_config(self) -> dict[str, ConfigEnum]:
        """Determines default behavior during input validation (if input_validation_config is None)."""
        return {var.name: var.default_validation_config for var in self.input_variables}

    @property
    def default_output_validation_config(self) -> dict[str, ConfigEnum]:
        """Determines default behavior during output validation (if output_validation_config is None)."""
        return {
            var.name: var.default_validation_config for var in self.output_variables
        }

    def evaluate(self, input_dict: dict[str, Any]) -> dict[str, Any]:
        """Main evaluation function, child classes must implement the _evaluate method."""
        validated_input_dict = self.input_validation(input_dict)
        output_dict = self._evaluate(validated_input_dict)
        self.output_validation(output_dict)
        return output_dict

    @abstractmethod
    def _evaluate(self, input_dict: dict[str, Any]) -> dict[str, Any]:
        pass

    def input_validation(self, input_dict: dict[str, Any]) -> dict[str, Any]:
        for name, value in input_dict.items():
            _config = (
                None
                if self.input_validation_config is None
                else self.input_validation_config.get(name)
            )
            var = self.input_variables[self.input_names.index(name)]
            var.validate_value(value, config=_config)
        return input_dict

    def output_validation(self, output_dict: dict[str, Any]) -> dict[str, Any]:
        for name, value in output_dict.items():
            _config = (
                None
                if self.output_validation_config is None
                else self.output_validation_config.get(name)
            )
            var = self.output_variables[self.output_names.index(name)]
            var.validate_value(value, config=_config)
        return output_dict

    def to_json(self, **kwargs) -> str:
        return json_dumps(self, **kwargs)

    def model_dump(self, **kwargs) -> dict[str, Any]:
        config = super().model_dump(**kwargs)
        config["input_variables"] = [var.model_dump() for var in self.input_variables]
        config["output_variables"] = [var.model_dump() for var in self.output_variables]
        return {"model_class": self.__class__.__name__} | config

    def json(self, **kwargs) -> str:
        result = self.to_json(**kwargs)
        config = json.loads(result)
        return json.dumps(config)

    def yaml(
        self,
        base_key: str = "",
        file_prefix: str = "",
        save_models: bool = False,
        save_jit: bool = False,
    ) -> str:
        """Serializes the object and returns a YAML formatted string defining the model.

        Args:
            base_key: Base key for serialization.
            file_prefix: Prefix for generated filenames.
            save_models: Determines whether models are saved to file.
            save_jit: Determines whether the model is saved as TorchScript
        Returns:
            YAML formatted string defining the model.
        """
        output = json.loads(
            self.to_json(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )
        s = yaml.dump(output, default_flow_style=None, sort_keys=False)
        return s

    def dump(
        self,
        file: Union[str, os.PathLike],
        base_key: str = "",
        save_models: bool = True,
        save_jit: bool = False,
    ):
        """Returns and optionally saves YAML formatted string defining the model.

        Args:
            file: File path to which the YAML formatted string and corresponding files are saved.
            base_key: Base key for serialization.
            save_models: Determines whether models are saved to file.
            save_jit: Determines whether the model is saved as TorchScript.
        """
        file_prefix = os.path.splitext(os.path.abspath(file))[0]
        with open(file, "w") as f:
            f.write(
                self.yaml(
                    base_key=base_key,
                    file_prefix=file_prefix,
                    save_models=save_models,
                    save_jit=save_jit,
                )
            )

    @classmethod
    def from_file(cls, filename: str):
        if not os.path.exists(filename):
            raise OSError(f"File {filename} is not found.")
        with open(filename, "r") as file:
            return cls.from_yaml(file)

    @classmethod
    def from_yaml(cls, yaml_obj: [str, TextIOWrapper]):
        return cls.model_validate(parse_config(yaml_obj, cls.model_fields))

default_input_validation_config property

Determines default behavior during input validation (if input_validation_config is None).

default_output_validation_config property

Determines default behavior during output validation (if output_validation_config is None).

__init__(*args, **kwargs)

Initializes LUMEBaseModel.

Parameters:

Name	Type	Description	Default
*args		Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path.	()
**kwargs		See class attributes.	{}

Source code in lume_model/base.py

def __init__(self, *args, **kwargs):
    """Initializes LUMEBaseModel.

    Args:
        *args: Accepts a single argument which is the model configuration as dictionary, YAML or JSON
          formatted string or file path.
        **kwargs: See class attributes.
    """
    if len(args) == 1:
        if len(kwargs) > 0:
            raise ValueError(
                "Cannot specify YAML string and keyword arguments for LUMEBaseModel init."
            )
        super().__init__(**parse_config(args[0], self.model_fields))
    elif len(args) > 1:
        raise ValueError(
            "Arguments to LUMEBaseModel must be either a single YAML string "
            "or keyword arguments passed directly to pydantic."
        )
    else:
        super().__init__(**kwargs)

dump(file, base_key='', save_models=True, save_jit=False)

Returns and optionally saves YAML formatted string defining the model.

Parameters:

Name	Type	Description	Default
file	Union[str, PathLike]	File path to which the YAML formatted string and corresponding files are saved.	required
base_key	str	Base key for serialization.	''
save_models	bool	Determines whether models are saved to file.	True
save_jit	bool	Determines whether the model is saved as TorchScript.	False

Source code in lume_model/base.py

def dump(
    self,
    file: Union[str, os.PathLike],
    base_key: str = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Returns and optionally saves YAML formatted string defining the model.

    Args:
        file: File path to which the YAML formatted string and corresponding files are saved.
        base_key: Base key for serialization.
        save_models: Determines whether models are saved to file.
        save_jit: Determines whether the model is saved as TorchScript.
    """
    file_prefix = os.path.splitext(os.path.abspath(file))[0]
    with open(file, "w") as f:
        f.write(
            self.yaml(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )

evaluate(input_dict)

Main evaluation function, child classes must implement the _evaluate method.

Source code in lume_model/base.py

def evaluate(self, input_dict: dict[str, Any]) -> dict[str, Any]:
    """Main evaluation function, child classes must implement the _evaluate method."""
    validated_input_dict = self.input_validation(input_dict)
    output_dict = self._evaluate(validated_input_dict)
    self.output_validation(output_dict)
    return output_dict

verify_input_default_value(value)

Verifies that input variables have the required default values.

Source code in lume_model/base.py

@field_validator("input_variables")
def verify_input_default_value(cls, value):
    """Verifies that input variables have the required default values."""
    for var in value:
        if var.default_value is None:
            raise ValueError(
                f"Input variable {var.name} must have a default value."
            )
    return value

yaml(base_key='', file_prefix='', save_models=False, save_jit=False)

Serializes the object and returns a YAML formatted string defining the model.

Parameters:

Name	Type	Description	Default
base_key	str	Base key for serialization.	''
file_prefix	str	Prefix for generated filenames.	''
save_models	bool	Determines whether models are saved to file.	False
save_jit	bool	Determines whether the model is saved as TorchScript	False

Returns: YAML formatted string defining the model.

Source code in lume_model/base.py

def yaml(
    self,
    base_key: str = "",
    file_prefix: str = "",
    save_models: bool = False,
    save_jit: bool = False,
) -> str:
    """Serializes the object and returns a YAML formatted string defining the model.

    Args:
        base_key: Base key for serialization.
        file_prefix: Prefix for generated filenames.
        save_models: Determines whether models are saved to file.
        save_jit: Determines whether the model is saved as TorchScript
    Returns:
        YAML formatted string defining the model.
    """
    output = json.loads(
        self.to_json(
            base_key=base_key,
            file_prefix=file_prefix,
            save_models=save_models,
            save_jit=save_jit,
        )
    )
    s = yaml.dump(output, default_flow_style=None, sort_keys=False)
    return s

process_torch_module(module, base_key='', key='', file_prefix='', save_modules=True, save_jit=False)

Optionally saves the given torch module to file and returns the filename.

Parameters:

Name	Type	Description	Default
base_key	str	Base key at this stage of serialization.	''
key	str	Key corresponding to the torch module.	''
module		The torch module to process.	required
file_prefix	Union[str, PathLike]	Prefix for generated filenames.	''
save_modules	bool	Determines whether torch modules are saved to file.	True
save_jit	bool	Determines whether the model gets saved as TorchScript.	False

Returns:

Type	Description
	Filename under which the torch module is (or would be) saved.

Source code in lume_model/base.py

def process_torch_module(
    module,
    base_key: str = "",
    key: str = "",
    file_prefix: Union[str, os.PathLike] = "",
    save_modules: bool = True,
    save_jit: bool = False,
):
    """Optionally saves the given torch module to file and returns the filename.

    Args:
        base_key: Base key at this stage of serialization.
        key: Key corresponding to the torch module.
        module: The torch module to process.
        file_prefix: Prefix for generated filenames.
        save_modules: Determines whether torch modules are saved to file.
        save_jit: Determines whether the model gets saved as TorchScript.

    Returns:
        Filename under which the torch module is (or would be) saved.
    """
    torch = try_import_module("torch")
    filepath_prefix, filename_prefix = os.path.split(file_prefix)
    prefixes = [ele for ele in [filename_prefix, base_key] if not ele == ""]
    filename = "{}.pt".format(key)
    jit_filename = "{}.jit".format(key)
    if prefixes:
        filename = "_".join((*prefixes, filename))
        jit_filename = "_".join((*prefixes, jit_filename))
    if save_modules:
        filepath = os.path.join(filepath_prefix, filename)
        torch.save(module, filepath)
    if save_jit:
        filepath = os.path.join(filepath_prefix, jit_filename)
        try:
            scripted_model = torch.jit.script(module)
            torch.jit.save(scripted_model, filepath)
        except Exception as e:
            logger.warning(
                "Saving as JIT through scripting has only been evaluated "
                "for NN models that don't depend on BoTorch modules."
            )
            logger.error(f"Failed to script the model: {e}")
            raise e
    return jit_filename if save_jit else filename

model_kwargs_from_dict(config)

Processes model configuration and returns the corresponding keyword arguments for model constructor.

Parameters:

Name	Type	Description	Default
config	dict	Model configuration.	required

Returns:

Type	Description
dict	Configuration as keyword arguments for model constructor.

Source code in lume_model/base.py

def model_kwargs_from_dict(config: dict) -> dict:
    """Processes model configuration and returns the corresponding keyword arguments for model constructor.

    Args:
        config: Model configuration.

    Returns:
        Configuration as keyword arguments for model constructor.
    """
    config = deserialize_variables(config)
    if all(key in config.keys() for key in ["input_variables", "output_variables"]):
        config["input_variables"], config["output_variables"] = variables_from_dict(
            config
        )
    config.pop("model_class", None)
    return config

parse_config(config, model_fields=None)

Parses model configuration and returns keyword arguments for model constructor.

Parameters:

Name	Type	Description	Default
config	Union[dict, str, TextIOWrapper, PathLike]	Model configuration as dictionary, YAML or JSON formatted string, file or file path.	required
model_fields	dict	Fields expected by the model (required for replacing relative paths).	None

Returns:

Type	Description
dict	Configuration as keyword arguments for model constructor.

Source code in lume_model/base.py

def parse_config(
    config: Union[dict, str, TextIOWrapper, os.PathLike],
    model_fields: dict = None,
) -> dict:
    """Parses model configuration and returns keyword arguments for model constructor.

    Args:
        config: Model configuration as dictionary, YAML or JSON formatted string, file or file path.
        model_fields: Fields expected by the model (required for replacing relative paths).

    Returns:
        Configuration as keyword arguments for model constructor.
    """
    config_file = None
    if isinstance(config, dict):
        d = config
    else:
        if isinstance(config, TextIOWrapper):
            yaml_str = config.read()
            config_file = os.path.abspath(config.name)
        elif isinstance(config, (str, os.PathLike)) and os.path.exists(config):
            with open(config) as f:
                yaml_str = f.read()
            config_file = os.path.abspath(config)
        else:
            yaml_str = config
        d = recursive_deserialize(yaml.safe_load(yaml_str))
    if config_file is not None:
        config_dir = os.path.dirname(os.path.realpath(config_file))
        d = replace_relative_paths(d, model_fields, config_dir)
    return model_kwargs_from_dict(d)

json_dumps(v, *, base_key='', file_prefix='', save_models=True, save_jit=False)

Serializes variables before dumping with json.

Parameters:

Name	Type	Description	Default
v		Object to dump.	required
base_key		Base key for serialization.	''
file_prefix	Union[str, PathLike]	Prefix for generated filenames.	''
save_models	bool	Determines whether models are saved to file.	True
save_jit	bool	Determines whether the model is saved as TorchScript.	False

Returns:

Type	Description
	JSON formatted string.

Source code in lume_model/base.py

def json_dumps(
    v,
    *,
    base_key="",
    file_prefix: Union[str, os.PathLike] = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Serializes variables before dumping with json.

    Args:
        v: Object to dump.
        base_key: Base key for serialization.
        file_prefix: Prefix for generated filenames.
        save_models: Determines whether models are saved to file.
        save_jit: Determines whether the model is saved as TorchScript.

    Returns:
        JSON formatted string.
    """
    v = recursive_serialize(
        v.model_dump(), base_key, file_prefix, save_models, save_jit
    )
    v = json.dumps(v)
    return v

json_loads(v)

Loads JSON formatted string and recursively deserializes the result.

Parameters:

Name	Type	Description	Default
v		JSON formatted string to load.	required

Returns:

Type	Description
	Deserialized object.

Source code in lume_model/base.py

def json_loads(v):
    """Loads JSON formatted string and recursively deserializes the result.

    Args:
        v: JSON formatted string to load.

    Returns:
        Deserialized object.
    """
    v = json.loads(v)
    v = recursive_deserialize(v)
    return v

recursive_serialize(v, base_key='', file_prefix='', save_models=True, save_jit=False)

Recursively performs custom serialization for the given object.

Parameters:

Name	Type	Description	Default
v	dict[str, Any]	Object to serialize.	required
base_key	str	Base key at this stage of serialization.	''
file_prefix	Union[str, PathLike]	Prefix for generated filenames.	''
save_models	bool	Determines whether models are saved to file.	True
save_jit	bool	Determines whether the model is saved as TorchScript.	False

Returns:

Type	Description
	Serialized object.

Source code in lume_model/base.py

def recursive_serialize(
    v: dict[str, Any],
    base_key: str = "",
    file_prefix: Union[str, os.PathLike] = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Recursively performs custom serialization for the given object.

    Args:
        v: Object to serialize.
        base_key: Base key at this stage of serialization.
        file_prefix: Prefix for generated filenames.
        save_models: Determines whether models are saved to file.
        save_jit: Determines whether the model is saved as TorchScript.

    Returns:
        Serialized object.
    """
    # try to import modules for LUMEBaseModel child classes
    torch = try_import_module("torch")
    # serialize
    v = serialize_variables(v)
    for key, value in v.items():
        if isinstance(value, dict):
            v[key] = recursive_serialize(value, key)
        elif torch is not None and isinstance(value, torch.nn.Module):
            v[key] = process_torch_module(
                value, base_key, key, file_prefix, save_models, save_jit
            )
        elif (
            isinstance(value, list)
            and torch is not None
            and any(isinstance(ele, torch.nn.Module) for ele in value)
        ):
            v[key] = [
                process_torch_module(
                    value[i], base_key, f"{key}_{i}", file_prefix, save_models, False
                )
                for i in range(len(value))
            ]
        else:
            for _type, func in JSON_ENCODERS.items():
                if isinstance(value, _type):
                    v[key] = func(value)
        # check to make sure object has been serialized, if not use a generic serializer
        try:
            json.dumps(v[key])
        except (TypeError, OverflowError):
            v[key] = f"{v[key].__module__}.{v[key].__class__.__qualname__}"

    return v

recursive_deserialize(v)

Recursively performs custom deserialization for the given object.

Parameters:

Name	Type	Description	Default
v		Object to deserialize.	required

Returns:

Type	Description
	Deserialized object.

Source code in lume_model/base.py

def recursive_deserialize(v):
    """Recursively performs custom deserialization for the given object.

    Args:
        v: Object to deserialize.

    Returns:
        Deserialized object.
    """
    # deserialize
    v = deserialize_variables(v)
    for key, value in v.items():
        if isinstance(value, dict):
            v[key] = recursive_deserialize(value)
    return v

TorchModel

Bases: LUMEBaseModel

LUME-model class for torch models.

By default, the models are assumed to be fixed, so all gradient computation is deactivated and the model and transformers are put in evaluation mode.

Attributes:

Name	Type	Description
model	Module	The torch base model.
input_variables	list[ScalarVariable]	List defining the input variables and their order.
output_variables	list[ScalarVariable]	List defining the output variables and their order.
input_transformers	list[Union[ReversibleInputTransform, Linear]]	List of transformer objects to apply to input before passing to model.
output_transformers	list[Union[ReversibleInputTransform, Linear]]	List of transformer objects to apply to output of model.
output_format	str	Determines format of outputs: "tensor" or "raw".
device	Union[device, str]	Device on which the model will be evaluated. Defaults to "cpu".
fixed_model	bool	If true, the model and transformers are put in evaluation mode and all gradient computation is deactivated.
precision	str	Precision of the model, either "double" or "single".

Source code in lume_model/models/torch_model.py

class TorchModel(LUMEBaseModel):
    """LUME-model class for torch models.

    By default, the models are assumed to be fixed, so all gradient computation is deactivated and the model and
    transformers are put in evaluation mode.

    Attributes:
        model: The torch base model.
        input_variables: List defining the input variables and their order.
        output_variables: List defining the output variables and their order.
        input_transformers: List of transformer objects to apply to input before passing to model.
        output_transformers: List of transformer objects to apply to output of model.
        output_format: Determines format of outputs: "tensor" or "raw".
        device: Device on which the model will be evaluated. Defaults to "cpu".
        fixed_model: If true, the model and transformers are put in evaluation mode and all gradient
          computation is deactivated.
        precision: Precision of the model, either "double" or "single".
    """

    model: torch.nn.Module
    input_transformers: list[Union[ReversibleInputTransform, torch.nn.Linear]] = None
    output_transformers: list[Union[ReversibleInputTransform, torch.nn.Linear]] = None
    output_format: str = "tensor"
    device: Union[torch.device, str] = "cpu"
    fixed_model: bool = True
    precision: str = "double"

    def __init__(self, *args, **kwargs):
        """Initializes TorchModel.

        Args:
            *args: Accepts a single argument which is the model configuration as dictionary, YAML or JSON
              formatted string or file path.
            **kwargs: See class attributes.
        """
        super().__init__(*args, **kwargs)
        self.input_transformers = (
            [] if self.input_transformers is None else self.input_transformers
        )
        self.output_transformers = (
            [] if self.output_transformers is None else self.output_transformers
        )

        # dtype property sets precision across model and transformers
        self.dtype

        # fixed model: set full model in eval mode and deactivate all gradients
        if self.fixed_model:
            is_scripted = isinstance(self.model, torch.jit.ScriptModule)
            self.model.eval().requires_grad_(False) if not is_scripted else None
            for t in self.input_transformers + self.output_transformers:
                if isinstance(t, torch.nn.Module):
                    t.eval().requires_grad_(False)

        # ensure consistent device
        self.to(self.device)

    @property
    def dtype(self):
        if self.precision == "double":
            self._dtype = torch.double
        elif self.precision == "single":
            self._dtype = torch.float
        else:
            raise ValueError(
                f"Unknown precision {self.precision}, "
                f"expected one of ['double', 'single']."
            )
        self._set_precision(self._dtype)
        return self._dtype

    @property
    def _tkwargs(self):
        return {"device": self.device, "dtype": self.dtype}

    @field_validator("model", mode="before")
    def validate_torch_model(cls, v):
        if isinstance(v, (str, os.PathLike)):
            if os.path.exists(v):
                fname = v
                try:
                    v = torch.jit.load(v)
                    print(f"Loaded TorchScript (JIT) model from file: {fname}")
                except RuntimeError:
                    v = torch.load(v, weights_only=False)
                    print(f"Loaded PyTorch model from file: {fname}")
            else:
                raise OSError(f"File {v} is not found.")
        return v

    @field_validator("input_transformers", "output_transformers", mode="before")
    def validate_transformers(cls, v):
        if not isinstance(v, list):
            raise ValueError("Transformers must be passed as list.")
        loaded_transformers = []
        for t in v:
            if isinstance(t, (str, os.PathLike)):
                if os.path.exists(t):
                    t = torch.load(t, weights_only=False)
                else:
                    raise OSError(f"File {t} is not found.")
            loaded_transformers.append(t)
        v = loaded_transformers
        return v

    @field_validator("output_format")
    def validate_output_format(cls, v):
        supported_formats = ["tensor", "variable", "raw"]
        if v not in supported_formats:
            raise ValueError(
                f"Unknown output format {v}, expected one of {supported_formats}."
            )
        return v

    def _set_precision(self, value: torch.dtype):
        """Sets the precision of the model."""
        self.model.to(dtype=value)
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.to(dtype=value)

    def _evaluate(
        self,
        input_dict: dict[str, Union[float, torch.Tensor]],
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Evaluates model on the given input dictionary.

        Args:
            input_dict: Input dictionary on which to evaluate the model.

        Returns:
            Dictionary of output variable names to values.
        """
        formatted_inputs = self._format_inputs(input_dict)
        input_tensor = self._arrange_inputs(formatted_inputs)
        input_tensor = self._transform_inputs(input_tensor)
        output_tensor = self.model(input_tensor)
        output_tensor = self._transform_outputs(output_tensor)
        parsed_outputs = self._parse_outputs(output_tensor)
        output_dict = self._prepare_outputs(parsed_outputs)
        return output_dict

    def input_validation(self, input_dict: dict[str, Union[float, torch.Tensor]]):
        """Validates input dictionary before evaluation.

        Args:
            input_dict: Input dictionary to validate.

        Returns:
            Validated input dictionary.
        """
        # validate input type (ints only are cast to floats for scalars)
        validated_input = InputDictModel(input_dict=input_dict).input_dict
        # format inputs as tensors w/o changing the dtype
        formatted_inputs = self._format_inputs(validated_input)
        # check default values for missing inputs
        filled_inputs = self._fill_default_inputs(formatted_inputs)
        # itemize inputs for validation
        itemized_inputs = self._itemize_dict(filled_inputs)

        for ele in itemized_inputs:
            # validate values that were in the torch tensor
            # any ints in the torch tensor will be cast to floats by Pydantic
            # but others will be caught, e.g. booleans
            ele = InputDictModel(input_dict=ele).input_dict
            # validate each value based on its var class and config
            super().input_validation(ele)

        # return the validated input dict for consistency w/ casting ints to floats
        if any([isinstance(value, torch.Tensor) for value in validated_input.values()]):
            validated_input = {
                k: v.to(**self._tkwargs) for k, v in validated_input.items()
            }

        return validated_input

    def output_validation(self, output_dict: dict[str, Union[float, torch.Tensor]]):
        """Itemizes tensors before performing output validation."""
        itemized_outputs = self._itemize_dict(output_dict)
        for ele in itemized_outputs:
            super().output_validation(ele)

    def random_input(self, n_samples: int = 1) -> dict[str, torch.Tensor]:
        """Generates random input(s) for the model.

        Args:
            n_samples: Number of random samples to generate.

        Returns:
            Dictionary of input variable names to tensors.
        """
        input_dict = {}
        for var in self.input_variables:
            if isinstance(var, ScalarVariable):
                input_dict[var.name] = var.value_range[0] + torch.rand(
                    size=(n_samples,)
                ) * (var.value_range[1] - var.value_range[0])
            else:
                torch.tensor(var.default_value, **self._tkwargs).repeat((n_samples, 1))
        return input_dict

    def random_evaluate(
        self, n_samples: int = 1
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Returns random evaluation(s) of the model.

        Args:
            n_samples: Number of random samples to evaluate.

        Returns:
            Dictionary of variable names to outputs.
        """
        random_input = self.random_input(n_samples)
        return self.evaluate(random_input)

    def to(self, device: Union[torch.device, str]):
        """Updates the device for the model, transformers and default values.

        Args:
            device: Device on which the model will be evaluated.
        """
        self.model.to(device)
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.to(device)
        self.device = device

    def insert_input_transformer(
        self, new_transformer: ReversibleInputTransform, loc: int
    ):
        """Inserts an additional input transformer at the given location.

        Args:
            new_transformer: New transformer to add.
            loc: Location where the new transformer shall be added to the transformer list.
        """
        self.input_transformers = (
            self.input_transformers[:loc]
            + [new_transformer]
            + self.input_transformers[loc:]
        )

    def insert_output_transformer(
        self, new_transformer: ReversibleInputTransform, loc: int
    ):
        """Inserts an additional output transformer at the given location.

        Args:
            new_transformer: New transformer to add.
            loc: Location where the new transformer shall be added to the transformer list.
        """
        self.output_transformers = (
            self.output_transformers[:loc]
            + [new_transformer]
            + self.output_transformers[loc:]
        )

    def update_input_variables_to_transformer(
        self, transformer_loc: int
    ) -> list[ScalarVariable]:
        """Returns input variables updated to the transformer at the given location.

        Updated are the value ranges and default of the input variables. This allows, e.g., to add a
        calibration transformer and to update the input variable specification accordingly.

        Args:
            transformer_loc: The location of the input transformer to adjust for.

        Returns:
            The updated input variables.
        """
        x_old = {
            "min": torch.tensor(
                [var.value_range[0] for var in self.input_variables], dtype=self.dtype
            ),
            "max": torch.tensor(
                [var.value_range[1] for var in self.input_variables], dtype=self.dtype
            ),
            "default": torch.tensor(
                [var.default_value for var in self.input_variables], dtype=self.dtype
            ),
        }
        x_new = {}
        for key in x_old.keys():
            x = x_old[key]

            # Make at least 2D
            if x.ndim == 0:
                x = x.unsqueeze(0)
            if x.ndim == 1:
                x = x.unsqueeze(0)

            # compute previous limits at transformer location
            for i in range(transformer_loc):
                if isinstance(self.input_transformers[i], ReversibleInputTransform):
                    x = self.input_transformers[i].transform(x)
                else:
                    x = self.input_transformers[i](x)
            # untransform of transformer to adjust for
            if isinstance(
                self.input_transformers[transformer_loc], ReversibleInputTransform
            ):
                x = self.input_transformers[transformer_loc].untransform(x)
            else:
                w = self.input_transformers[transformer_loc].weight
                b = self.input_transformers[transformer_loc].bias
                x = torch.matmul((x - b), torch.linalg.inv(w.T))
            # backtrack through transformers
            for transformer in self.input_transformers[:transformer_loc][::-1]:
                if isinstance(
                    self.input_transformers[transformer_loc], ReversibleInputTransform
                ):
                    x = transformer.untransform(x)
                else:
                    w, b = transformer.weight, transformer.bias
                    x = torch.matmul((x - b), torch.linalg.inv(w.T))
            x_new[key] = x
        updated_variables = deepcopy(self.input_variables)
        for i, var in enumerate(updated_variables):
            var.value_range = [x_new["min"][0][i].item(), x_new["max"][0][i].item()]
            var.default_value = x_new["default"][0][i].item()
        return updated_variables

    def _format_inputs(
        self,
        input_dict: dict[str, Union[float, torch.Tensor]],
    ) -> dict[str, torch.Tensor]:
        """Formats values of the input dictionary as tensors.

        Args:
            input_dict: Dictionary of input variable names to values.

        Returns:
            Dictionary of input variable names to tensors.
        """
        formatted_inputs = {}
        for var_name, value in input_dict.items():
            v = value if isinstance(value, torch.Tensor) else torch.tensor(value)
            formatted_inputs[var_name] = v.squeeze()
        return formatted_inputs

    def _fill_default_inputs(
        self, input_dict: dict[str, torch.Tensor]
    ) -> dict[str, torch.Tensor]:
        """Fills missing input variables with default values.

        Args:
            input_dict: Dictionary of input variable names to tensors.

        Returns:
            Dictionary of input variable names to tensors with default values for missing inputs.
        """
        for var in self.input_variables:
            if var.name not in input_dict.keys():
                input_dict[var.name] = torch.tensor(var.default_value, **self._tkwargs)
        return input_dict

    def _arrange_inputs(
        self, formatted_inputs: dict[str, torch.Tensor]
    ) -> torch.Tensor:
        """Enforces order of input variables.

        Enforces the order of the input variables to be passed to the transformers and model and updates the
        returned tensor with default values for any inputs that are missing.

        Args:
            formatted_inputs: Dictionary of input variable names to tensors.

        Returns:
            Ordered input tensor to be passed to the transformers.
        """
        default_tensor = torch.tensor(
            [var.default_value for var in self.input_variables], **self._tkwargs
        )

        # determine input shape
        input_shapes = [formatted_inputs[k].shape for k in formatted_inputs.keys()]
        if not all(ele == input_shapes[0] for ele in input_shapes):
            raise ValueError("Inputs have inconsistent shapes.")

        input_tensor = torch.tile(default_tensor, dims=(*input_shapes[0], 1))
        for key, value in formatted_inputs.items():
            input_tensor[..., self.input_names.index(key)] = value

        if input_tensor.shape[-1] != len(self.input_names):
            raise ValueError(
                f"""
                Last dimension of input tensor doesn't match the expected number of inputs\n
                received: {default_tensor.shape}, expected {len(self.input_names)} as the last dimension
                """
            )
        return input_tensor

    def _transform_inputs(self, input_tensor: torch.Tensor) -> torch.Tensor:
        """Applies transformations to the inputs.

        Args:
            input_tensor: Ordered input tensor to be passed to the transformers.

        Returns:
            Tensor of transformed inputs to be passed to the model.
        """
        # Make at least 2D
        if input_tensor.ndim == 0:
            input_tensor = input_tensor.unsqueeze(0)
        if input_tensor.ndim == 1:
            input_tensor = input_tensor.unsqueeze(0)

        for transformer in self.input_transformers:
            if isinstance(transformer, ReversibleInputTransform):
                input_tensor = transformer.transform(input_tensor)
            else:
                input_tensor = transformer(input_tensor)
        return input_tensor

    def _transform_outputs(self, output_tensor: torch.Tensor) -> torch.Tensor:
        """(Un-)Transforms the model output tensor.

        Args:
            output_tensor: Output tensor from the model.

        Returns:
            (Un-)Transformed output tensor.
        """
        for transformer in self.output_transformers:
            if isinstance(transformer, ReversibleInputTransform):
                output_tensor = transformer.untransform(output_tensor)
            else:
                w, b = transformer.weight, transformer.bias
                output_tensor = torch.matmul((output_tensor - b), torch.linalg.inv(w.T))
        return output_tensor

    def _parse_outputs(self, output_tensor: torch.Tensor) -> dict[str, torch.Tensor]:
        """Constructs dictionary from model output tensor.

        Args:
            output_tensor: (Un-)transformed output tensor from the model.

        Returns:
            Dictionary of output variable names to (un-)transformed tensors.
        """
        parsed_outputs = {}
        if output_tensor.dim() in [0, 1]:
            output_tensor = output_tensor.unsqueeze(0)
        if len(self.output_names) == 1:
            parsed_outputs[self.output_names[0]] = output_tensor.squeeze()
        else:
            for idx, output_name in enumerate(self.output_names):
                parsed_outputs[output_name] = output_tensor[..., idx].squeeze()
        return parsed_outputs

    def _prepare_outputs(
        self,
        parsed_outputs: dict[str, torch.Tensor],
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Updates and returns outputs according to output_format.

        Updates the output variables within the model to reflect the new values.

        Args:
            parsed_outputs: Dictionary of output variable names to transformed tensors.

        Returns:
            Dictionary of output variable names to values depending on output_format.
        """
        if self.output_format.lower() == "tensor":
            return parsed_outputs
        else:
            return {
                key: value.item() if value.squeeze().dim() == 0 else value
                for key, value in parsed_outputs.items()
            }

    @staticmethod
    def _itemize_dict(
        d: dict[str, Union[float, torch.Tensor]],
    ) -> list[dict[str, Union[float, torch.Tensor]]]:
        """Itemizes the given in-/output dictionary.

        Args:
            d: Dictionary to itemize.

        Returns:
            List of in-/output dictionaries, each containing only a single value per in-/output.
        """
        has_tensors = any([isinstance(value, torch.Tensor) for value in d.values()])
        itemized_dicts = []
        if has_tensors:
            for k, v in d.items():
                for i, ele in enumerate(v.flatten()):
                    if i >= len(itemized_dicts):
                        itemized_dicts.append({k: ele.item()})
                    else:
                        itemized_dicts[i][k] = ele.item()
        else:
            itemized_dicts = [d]
        return itemized_dicts

__init__(*args, **kwargs)

Initializes TorchModel.

Parameters:

Name	Type	Description	Default
*args		Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path.	()
**kwargs		See class attributes.	{}

Source code in lume_model/models/torch_model.py

def __init__(self, *args, **kwargs):
    """Initializes TorchModel.

    Args:
        *args: Accepts a single argument which is the model configuration as dictionary, YAML or JSON
          formatted string or file path.
        **kwargs: See class attributes.
    """
    super().__init__(*args, **kwargs)
    self.input_transformers = (
        [] if self.input_transformers is None else self.input_transformers
    )
    self.output_transformers = (
        [] if self.output_transformers is None else self.output_transformers
    )

    # dtype property sets precision across model and transformers
    self.dtype

    # fixed model: set full model in eval mode and deactivate all gradients
    if self.fixed_model:
        is_scripted = isinstance(self.model, torch.jit.ScriptModule)
        self.model.eval().requires_grad_(False) if not is_scripted else None
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.eval().requires_grad_(False)

    # ensure consistent device
    self.to(self.device)

input_validation(input_dict)

Validates input dictionary before evaluation.

Parameters:

Name	Type	Description	Default
input_dict	dict[str, Union[float, Tensor]]	Input dictionary to validate.	required

Returns:

Type	Description
	Validated input dictionary.

Source code in lume_model/models/torch_model.py

def input_validation(self, input_dict: dict[str, Union[float, torch.Tensor]]):
    """Validates input dictionary before evaluation.

    Args:
        input_dict: Input dictionary to validate.

    Returns:
        Validated input dictionary.
    """
    # validate input type (ints only are cast to floats for scalars)
    validated_input = InputDictModel(input_dict=input_dict).input_dict
    # format inputs as tensors w/o changing the dtype
    formatted_inputs = self._format_inputs(validated_input)
    # check default values for missing inputs
    filled_inputs = self._fill_default_inputs(formatted_inputs)
    # itemize inputs for validation
    itemized_inputs = self._itemize_dict(filled_inputs)

    for ele in itemized_inputs:
        # validate values that were in the torch tensor
        # any ints in the torch tensor will be cast to floats by Pydantic
        # but others will be caught, e.g. booleans
        ele = InputDictModel(input_dict=ele).input_dict
        # validate each value based on its var class and config
        super().input_validation(ele)

    # return the validated input dict for consistency w/ casting ints to floats
    if any([isinstance(value, torch.Tensor) for value in validated_input.values()]):
        validated_input = {
            k: v.to(**self._tkwargs) for k, v in validated_input.items()
        }

    return validated_input

insert_input_transformer(new_transformer, loc)

Inserts an additional input transformer at the given location.

Parameters:

Name	Type	Description	Default
new_transformer	ReversibleInputTransform	New transformer to add.	required
loc	int	Location where the new transformer shall be added to the transformer list.	required

Source code in lume_model/models/torch_model.py

def insert_input_transformer(
    self, new_transformer: ReversibleInputTransform, loc: int
):
    """Inserts an additional input transformer at the given location.

    Args:
        new_transformer: New transformer to add.
        loc: Location where the new transformer shall be added to the transformer list.
    """
    self.input_transformers = (
        self.input_transformers[:loc]
        + [new_transformer]
        + self.input_transformers[loc:]
    )

insert_output_transformer(new_transformer, loc)

Inserts an additional output transformer at the given location.

Parameters:

Name	Type	Description	Default
new_transformer	ReversibleInputTransform	New transformer to add.	required
loc	int	Location where the new transformer shall be added to the transformer list.	required

Source code in lume_model/models/torch_model.py

def insert_output_transformer(
    self, new_transformer: ReversibleInputTransform, loc: int
):
    """Inserts an additional output transformer at the given location.

    Args:
        new_transformer: New transformer to add.
        loc: Location where the new transformer shall be added to the transformer list.
    """
    self.output_transformers = (
        self.output_transformers[:loc]
        + [new_transformer]
        + self.output_transformers[loc:]
    )

output_validation(output_dict)

Itemizes tensors before performing output validation.

Source code in lume_model/models/torch_model.py

def output_validation(self, output_dict: dict[str, Union[float, torch.Tensor]]):
    """Itemizes tensors before performing output validation."""
    itemized_outputs = self._itemize_dict(output_dict)
    for ele in itemized_outputs:
        super().output_validation(ele)

random_evaluate(n_samples=1)

Returns random evaluation(s) of the model.

Parameters:

Name	Type	Description	Default
n_samples	int	Number of random samples to evaluate.	1

Returns:

Type	Description
dict[str, Union[float, Tensor]]	Dictionary of variable names to outputs.

Source code in lume_model/models/torch_model.py

def random_evaluate(
    self, n_samples: int = 1
) -> dict[str, Union[float, torch.Tensor]]:
    """Returns random evaluation(s) of the model.

    Args:
        n_samples: Number of random samples to evaluate.

    Returns:
        Dictionary of variable names to outputs.
    """
    random_input = self.random_input(n_samples)
    return self.evaluate(random_input)

random_input(n_samples=1)

Generates random input(s) for the model.

Parameters:

Name	Type	Description	Default
n_samples	int	Number of random samples to generate.	1

Returns:

Type	Description
dict[str, Tensor]	Dictionary of input variable names to tensors.

Source code in lume_model/models/torch_model.py

def random_input(self, n_samples: int = 1) -> dict[str, torch.Tensor]:
    """Generates random input(s) for the model.

    Args:
        n_samples: Number of random samples to generate.

    Returns:
        Dictionary of input variable names to tensors.
    """
    input_dict = {}
    for var in self.input_variables:
        if isinstance(var, ScalarVariable):
            input_dict[var.name] = var.value_range[0] + torch.rand(
                size=(n_samples,)
            ) * (var.value_range[1] - var.value_range[0])
        else:
            torch.tensor(var.default_value, **self._tkwargs).repeat((n_samples, 1))
    return input_dict

to(device)

Updates the device for the model, transformers and default values.

Parameters:

Name	Type	Description	Default
device	Union[device, str]	Device on which the model will be evaluated.	required

Source code in lume_model/models/torch_model.py

def to(self, device: Union[torch.device, str]):
    """Updates the device for the model, transformers and default values.

    Args:
        device: Device on which the model will be evaluated.
    """
    self.model.to(device)
    for t in self.input_transformers + self.output_transformers:
        if isinstance(t, torch.nn.Module):
            t.to(device)
    self.device = device

update_input_variables_to_transformer(transformer_loc)

Returns input variables updated to the transformer at the given location.

Updated are the value ranges and default of the input variables. This allows, e.g., to add a calibration transformer and to update the input variable specification accordingly.

Parameters:

Name	Type	Description	Default
transformer_loc	int	The location of the input transformer to adjust for.	required

Returns:

Type	Description
list[ScalarVariable]	The updated input variables.

Source code in lume_model/models/torch_model.py

def update_input_variables_to_transformer(
    self, transformer_loc: int
) -> list[ScalarVariable]:
    """Returns input variables updated to the transformer at the given location.

    Updated are the value ranges and default of the input variables. This allows, e.g., to add a
    calibration transformer and to update the input variable specification accordingly.

    Args:
        transformer_loc: The location of the input transformer to adjust for.

    Returns:
        The updated input variables.
    """
    x_old = {
        "min": torch.tensor(
            [var.value_range[0] for var in self.input_variables], dtype=self.dtype
        ),
        "max": torch.tensor(
            [var.value_range[1] for var in self.input_variables], dtype=self.dtype
        ),
        "default": torch.tensor(
            [var.default_value for var in self.input_variables], dtype=self.dtype
        ),
    }
    x_new = {}
    for key in x_old.keys():
        x = x_old[key]

        # Make at least 2D
        if x.ndim == 0:
            x = x.unsqueeze(0)
        if x.ndim == 1:
            x = x.unsqueeze(0)

        # compute previous limits at transformer location
        for i in range(transformer_loc):
            if isinstance(self.input_transformers[i], ReversibleInputTransform):
                x = self.input_transformers[i].transform(x)
            else:
                x = self.input_transformers[i](x)
        # untransform of transformer to adjust for
        if isinstance(
            self.input_transformers[transformer_loc], ReversibleInputTransform
        ):
            x = self.input_transformers[transformer_loc].untransform(x)
        else:
            w = self.input_transformers[transformer_loc].weight
            b = self.input_transformers[transformer_loc].bias
            x = torch.matmul((x - b), torch.linalg.inv(w.T))
        # backtrack through transformers
        for transformer in self.input_transformers[:transformer_loc][::-1]:
            if isinstance(
                self.input_transformers[transformer_loc], ReversibleInputTransform
            ):
                x = transformer.untransform(x)
            else:
                w, b = transformer.weight, transformer.bias
                x = torch.matmul((x - b), torch.linalg.inv(w.T))
        x_new[key] = x
    updated_variables = deepcopy(self.input_variables)
    for i, var in enumerate(updated_variables):
        var.value_range = [x_new["min"][0][i].item(), x_new["max"][0][i].item()]
        var.default_value = x_new["default"][0][i].item()
    return updated_variables

InputDictModel

Bases: BaseModel

Pydantic model for input dictionary validation.

Attributes:

Name	Type	Description
input_dict	Dict[str, Union[Tensor, float]]	Input dictionary to validate.

Source code in lume_model/models/torch_model.py

class InputDictModel(BaseModel):
    """Pydantic model for input dictionary validation.

    Attributes:
        input_dict: Input dictionary to validate.
    """

    input_dict: Dict[str, Union[torch.Tensor, float]]

    model_config = ConfigDict(arbitrary_types_allowed=True, strict=True)

TorchModule

Bases: Module

Wrapper to allow a LUME TorchModel to be used like a torch.nn.Module.

As the base model within the TorchModel is assumed to be fixed during instantiation, so is the TorchModule.

Source code in lume_model/models/torch_module.py

class TorchModule(torch.nn.Module):
    """Wrapper to allow a LUME TorchModel to be used like a torch.nn.Module.

    As the base model within the TorchModel is assumed to be fixed during instantiation,
    so is the TorchModule.
    """

    def __init__(
        self,
        *args,
        model: TorchModel = None,
        input_order: list[str] = None,
        output_order: list[str] = None,
    ):
        """Initializes TorchModule.

        Args:
            *args: Accepts a single argument which is the model configuration as dictionary, YAML or JSON
              formatted string or file path.

        Keyword Args:
            model: The TorchModel instance to wrap around. If config is None, this has to be defined.
            input_order: Input names in the order they are passed to the model. If None, the input order of the
              TorchModel is used.
            output_order: Output names in the order they are returned by the model. If None, the output order of
              the TorchModel is used.
        """
        if all(arg is None for arg in [*args, model]):
            raise ValueError(
                "Either a YAML string has to be given or model has to be defined."
            )
        super().__init__()
        if len(args) == 1:
            if not all(v is None for v in [model, input_order, output_order]):
                raise ValueError(
                    "Cannot specify YAML string and keyword arguments for TorchModule init."
                )
            model_fields = {f"model.{k}": v for k, v in TorchModel.model_fields.items()}
            kwargs = parse_config(args[0], model_fields)
            kwargs["model"] = TorchModel(kwargs["model"])
            self.__init__(**kwargs)
        elif len(args) > 1:
            raise ValueError(
                "Arguments to TorchModule must be either a single YAML string or keyword arguments."
            )
        else:
            self._model = model
            self._input_order = input_order
            self._output_order = output_order
            self.register_module("base_model", self._model.model)
            for i, input_transformer in enumerate(self._model.input_transformers):
                self.register_module(f"input_transformers_{i}", input_transformer)
            for i, output_transformer in enumerate(self._model.output_transformers):
                self.register_module(f"output_transformers_{i}", output_transformer)
            if not model.model.training:  # TorchModel defines train/eval mode
                self.eval()

    @property
    def model(self):
        return self._model

    @property
    def input_order(self):
        if self._input_order is None:
            return self._model.input_names
        else:
            return self._input_order

    @property
    def output_order(self):
        if self._output_order is None:
            return self._model.output_names
        else:
            return self._output_order

    def forward(self, x: torch.Tensor):
        # input shape: [n_batch, n_samples, n_dim]
        x = self._validate_input(x)
        model_input = self._tensor_to_dictionary(x)
        y_model = self.evaluate_model(model_input)
        y_model = self.manipulate_output(y_model)
        # squeeze for use as prior mean in botorch GPs
        y = self._dictionary_to_tensor(y_model).squeeze()
        return y

    def yaml(
        self,
        base_key: str = "",
        file_prefix: str = "",
        save_models: bool = False,
        save_jit: bool = False,
    ) -> str:
        """Serializes the object and returns a YAML formatted string defining the TorchModule instance.

        Args:
            base_key: Base key for serialization.
            file_prefix: Prefix for generated filenames.
            save_models: Determines whether models are saved to file.
            save_jit: Determines whether the structure of the model is saved as TorchScript

        Returns:
            YAML formatted string defining the TorchModule instance.
        """
        d = {}
        for k, v in inspect.signature(TorchModule.__init__).parameters.items():
            if k not in ["self", "args", "model"]:
                d[k] = getattr(self, k)
        output = json.loads(
            json.dumps(
                recursive_serialize(d, base_key, file_prefix, save_models, save_jit)
            )
        )
        model_output = json.loads(
            self._model.to_json(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )
        output["model"] = model_output
        # create YAML formatted string
        s = yaml.dump(
            {"model_class": self.__class__.__name__} | output,
            default_flow_style=None,
            sort_keys=False,
        )
        return s

    def dump(
        self,
        file: Union[str, os.PathLike],
        save_models: bool = True,
        base_key: str = "",
        save_jit: bool = False,
    ):
        """Returns and optionally saves YAML formatted string defining the model.

        Args:
            file: File path to which the YAML formatted string and corresponding files are saved.
            base_key: Base key for serialization.
            save_models: Determines whether models are saved to file.
            save_jit : Whether the model is saved using just in time pytorch method
        """
        file_prefix = os.path.splitext(file)[0]
        with open(file, "w") as f:
            f.write(
                self.yaml(
                    save_models=save_models,
                    base_key=base_key,
                    file_prefix=file_prefix,
                    save_jit=save_jit,
                )
            )

    def evaluate_model(self, x: dict[str, torch.Tensor]):
        """Placeholder method to modify model calls."""
        return self._model.evaluate(x)

    def manipulate_output(self, y_model: dict[str, torch.Tensor]):
        """Placeholder method to modify the model output."""
        return y_model

    def _tensor_to_dictionary(self, x: torch.Tensor):
        input_dict = {}
        for idx, input_name in enumerate(self.input_order):
            input_dict[input_name] = x[..., idx].unsqueeze(-1)
        return input_dict

    def _dictionary_to_tensor(self, y_model: dict[str, torch.Tensor]):
        output_tensor = torch.stack(
            [y_model[output_name].unsqueeze(-1) for output_name in self.output_order],
            dim=-1,
        )
        return output_tensor

    @staticmethod
    def _validate_input(x: torch.Tensor) -> torch.Tensor:
        if x.dim() <= 1:
            raise ValueError(
                f"Expected input dim to be at least 2 ([n_samples, n_features]), received: {tuple(x.shape)}"
            )
        else:
            return x

__init__(*args, model=None, input_order=None, output_order=None)

Initializes TorchModule.

Parameters:

Name	Type	Description	Default
*args		Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path.	()

Other Parameters:

Name	Type	Description
model	TorchModel	The TorchModel instance to wrap around. If config is None, this has to be defined.
input_order	list[str]	Input names in the order they are passed to the model. If None, the input order of the TorchModel is used.
output_order	list[str]	Output names in the order they are returned by the model. If None, the output order of the TorchModel is used.

Source code in lume_model/models/torch_module.py

def __init__(
    self,
    *args,
    model: TorchModel = None,
    input_order: list[str] = None,
    output_order: list[str] = None,
):
    """Initializes TorchModule.

    Args:
        *args: Accepts a single argument which is the model configuration as dictionary, YAML or JSON
          formatted string or file path.

    Keyword Args:
        model: The TorchModel instance to wrap around. If config is None, this has to be defined.
        input_order: Input names in the order they are passed to the model. If None, the input order of the
          TorchModel is used.
        output_order: Output names in the order they are returned by the model. If None, the output order of
          the TorchModel is used.
    """
    if all(arg is None for arg in [*args, model]):
        raise ValueError(
            "Either a YAML string has to be given or model has to be defined."
        )
    super().__init__()
    if len(args) == 1:
        if not all(v is None for v in [model, input_order, output_order]):
            raise ValueError(
                "Cannot specify YAML string and keyword arguments for TorchModule init."
            )
        model_fields = {f"model.{k}": v for k, v in TorchModel.model_fields.items()}
        kwargs = parse_config(args[0], model_fields)
        kwargs["model"] = TorchModel(kwargs["model"])
        self.__init__(**kwargs)
    elif len(args) > 1:
        raise ValueError(
            "Arguments to TorchModule must be either a single YAML string or keyword arguments."
        )
    else:
        self._model = model
        self._input_order = input_order
        self._output_order = output_order
        self.register_module("base_model", self._model.model)
        for i, input_transformer in enumerate(self._model.input_transformers):
            self.register_module(f"input_transformers_{i}", input_transformer)
        for i, output_transformer in enumerate(self._model.output_transformers):
            self.register_module(f"output_transformers_{i}", output_transformer)
        if not model.model.training:  # TorchModel defines train/eval mode
            self.eval()

dump(file, save_models=True, base_key='', save_jit=False)

Returns and optionally saves YAML formatted string defining the model.

Parameters:

Name	Type	Description	Default
file	Union[str, PathLike]	File path to which the YAML formatted string and corresponding files are saved.	required
base_key	str	Base key for serialization.	''
save_models	bool	Determines whether models are saved to file.	True
save_jit		Whether the model is saved using just in time pytorch method	False

Source code in lume_model/models/torch_module.py

def dump(
    self,
    file: Union[str, os.PathLike],
    save_models: bool = True,
    base_key: str = "",
    save_jit: bool = False,
):
    """Returns and optionally saves YAML formatted string defining the model.

    Args:
        file: File path to which the YAML formatted string and corresponding files are saved.
        base_key: Base key for serialization.
        save_models: Determines whether models are saved to file.
        save_jit : Whether the model is saved using just in time pytorch method
    """
    file_prefix = os.path.splitext(file)[0]
    with open(file, "w") as f:
        f.write(
            self.yaml(
                save_models=save_models,
                base_key=base_key,
                file_prefix=file_prefix,
                save_jit=save_jit,
            )
        )

evaluate_model(x)

Placeholder method to modify model calls.

Source code in lume_model/models/torch_module.py

def evaluate_model(self, x: dict[str, torch.Tensor]):
    """Placeholder method to modify model calls."""
    return self._model.evaluate(x)

manipulate_output(y_model)

Placeholder method to modify the model output.

Source code in lume_model/models/torch_module.py

def manipulate_output(self, y_model: dict[str, torch.Tensor]):
    """Placeholder method to modify the model output."""
    return y_model

yaml(base_key='', file_prefix='', save_models=False, save_jit=False)

Serializes the object and returns a YAML formatted string defining the TorchModule instance.

Parameters:

Name	Type	Description	Default
base_key	str	Base key for serialization.	''
file_prefix	str	Prefix for generated filenames.	''
save_models	bool	Determines whether models are saved to file.	False
save_jit	bool	Determines whether the structure of the model is saved as TorchScript	False

Returns:

Type	Description
str	YAML formatted string defining the TorchModule instance.

Source code in lume_model/models/torch_module.py

def yaml(
    self,
    base_key: str = "",
    file_prefix: str = "",
    save_models: bool = False,
    save_jit: bool = False,
) -> str:
    """Serializes the object and returns a YAML formatted string defining the TorchModule instance.

    Args:
        base_key: Base key for serialization.
        file_prefix: Prefix for generated filenames.
        save_models: Determines whether models are saved to file.
        save_jit: Determines whether the structure of the model is saved as TorchScript

    Returns:
        YAML formatted string defining the TorchModule instance.
    """
    d = {}
    for k, v in inspect.signature(TorchModule.__init__).parameters.items():
        if k not in ["self", "args", "model"]:
            d[k] = getattr(self, k)
    output = json.loads(
        json.dumps(
            recursive_serialize(d, base_key, file_prefix, save_models, save_jit)
        )
    )
    model_output = json.loads(
        self._model.to_json(
            base_key=base_key,
            file_prefix=file_prefix,
            save_models=save_models,
            save_jit=save_jit,
        )
    )
    output["model"] = model_output
    # create YAML formatted string
    s = yaml.dump(
        {"model_class": self.__class__.__name__} | output,
        default_flow_style=None,
        sort_keys=False,
    )
    return s