Merge pull request #87 from NREL/feature/variable-paths

Feature/variable paths

python/altrios/__init__.py

@@ -2,8 +2,14 @@
 __version__ = get_distribution("altrios").version
 
 from pathlib import Path
+import re
 import numpy as np
 import logging
+import inspect
+from typing import List, Union, Dict, Optional
+from typing_extensions import Self
+import pandas as pd
+import polars as pl
 
 from altrios.loaders.powertrain_components import _res_from_excel
 from altrios.utilities import set_param_from_path  # noqa: F401
@@ -12,6 +18,7 @@
 from altrios.utilities import package_root, resources_root
 # make everything in altrios_pyo3 available here
 from altrios.altrios_pyo3 import *
+from altrios import *
 
 DEFAULT_LOGGING_CONFIG = dict(
     format="%(asctime)s.%(msecs)03d | %(filename)s:%(lineno)s | %(levelname)s: %(message)s",
@@ -25,6 +32,143 @@
 def __array__(self):
     return np.array(self.tolist())
 
+# creates a list of all python classes from rust structs that need variable_path_list() and
+# history_path_list() added as methods
+ACCEPTED_RUST_STRUCTS = [
+    attr for attr in altrios_pyo3.__dir__() if not attr.startswith("__") and isinstance(getattr(altrios_pyo3, attr), type) and
+    attr[0].isupper() 
+]
+
+def variable_path_list(self, element_as_list:bool=False) -> List[str]:
+    """
+    Returns list of key paths to all variables and sub-variables within
+    dict version of `self`. See example usage in `altrios/demos/
+    demo_variable_paths.py`.
+
+    # Arguments:  
+    - `element_as_list`: if True, each element is itself a list of the path elements
+    """
+    return variable_path_list_from_py_objs(self.to_pydict(), element_as_list=element_as_list)
+
+def variable_path_list_from_py_objs(
+    obj: Union[Dict, List], 
+    pre_path:Optional[str]=None,
+    element_as_list:bool=False,
+) -> List[str]:
+    """
+    Returns list of key paths to all variables and sub-variables within
+    dict version of class. See example usage in `altrios/demos/
+    demo_variable_paths.py`.
+
+    # Arguments:  
+    - `obj`: altrios object in dictionary form from `to_pydict()`
+    - `pre_path`: This is used to call the method recursively and should not be
+        specified by user.  Specifies a path to be added in front of all paths
+        returned by the method.
+    - `element_as_list`: if True, each element is itself a list of the path elements
+    """
+    key_paths = []
+    if isinstance(obj, dict):
+        for key, val in obj.items():
+            # check for nested dicts and call recursively
+            if isinstance(val, dict):
+                key_path = f"['{key}']" if pre_path is None else pre_path + f"['{key}']"
+                key_paths.extend(variable_path_list_from_py_objs(val, key_path))
+            # check for lists or other iterables that do not contain numeric data
+            elif "__iter__" in dir(val) and (len(val) > 0) and not(isinstance(val[0], float) or isinstance(val[0], int)):
+                key_path = f"['{key}']" if pre_path is None else pre_path + f"['{key}']"
+                key_paths.extend(variable_path_list_from_py_objs(val, key_path))
+            else:
+                key_path = f"['{key}']" if pre_path is None else pre_path + f"['{key}']"
+                key_paths.append(key_path)
+
+    elif isinstance(obj, list):
+        for key, val in enumerate(obj):
+            # check for nested dicts and call recursively
+            if isinstance(val, dict):
+                key_path = f"[{key}]" if pre_path is None else pre_path + f"[{key}]"
+                key_paths.extend(variable_path_list_from_py_objs(val, key_path))
+            # check for lists or other iterables that do not contain numeric data
+            elif "__iter__" in dir(val) and (len(val) > 0) and not(isinstance(val[0], float) or isinstance(val[0], int)):
+                key_path = f"[{key}]" if pre_path is None else pre_path + f"[{key}]"
+                key_paths.extend(variable_path_list_from_py_objs(val, key_path))
+            else:
+                key_path = f"[{key}]" if pre_path is None else pre_path + f"[{key}]"
+                key_paths.append(key_path)
+    if element_as_list:
+        re_for_elems = re.compile("\\[('(\\w+)'|(\\w+))\\]")
+        for i, kp in enumerate(key_paths):
+            kp: str
+            groups = re_for_elems.findall(kp)
+            selected = [g[1] if len(g[1]) > 0 else g[2] for g in groups]
+            key_paths[i] = selected
+
+    return key_paths
+
+def history_path_list(self, element_as_list:bool=False) -> List[str]:
+    """
+    Returns a list of relative paths to all history variables (all variables
+    that contain history as a subpath). 
+    See example usage in `altrios/demos/demo_variable_paths.py`.
+
+    # Arguments
+    - `element_as_list`: if True, each element is itself a list of the path elements
+    """
+    item_str = lambda item: item if not element_as_list else ".".join(item)
+    history_path_list = [
+        item for item in self.variable_path_list(
+            element_as_list=element_as_list) if "history" in item_str(item)
+    ]
+    return history_path_list
+
+def to_pydict(self) -> Dict:
+    """
+    Returns self converted to pure python dictionary with no nested Rust objects
+    """
+    import json
+    return json.loads(self.to_json())
+
+@classmethod
+def from_pydict(cls, pydict: Dict) -> Self:
+    """
+    Instantiates Self from pure python dictionary 
+    """
+    import json
+    return cls.from_json(json.dumps(pydict))
+
+def to_dataframe(self, pandas:bool=False) -> [pd.DataFrame, pl.DataFrame, pl.LazyFrame]:
+    """
+    Returns time series results from altrios object as a Polars or Pandas dataframe.
+
+    # Arguments
+    - `pandas`: returns pandas dataframe if True; otherwise, returns polars dataframe by default
+    """
+    obj_dict = self.to_pydict()
+    history_paths = self.history_path_list(element_as_list=True)   
+    cols = [".".join(hp) for hp in history_paths]
+    vals = []
+    for hp in history_paths:
+        obj:Union[dict|list] = obj_dict
+        for elem in hp:
+            try: 
+                obj = obj[elem]
+            except:
+                obj = obj[int(elem)]
+        vals.append(obj)
+    if not pandas:
+        df = pl.DataFrame({col: val for col, val in zip(cols, vals)})
+    else:
+        df = pd.DataFrame({col: val for col, val in zip(cols, vals)})
+    return df
+
+# adds variable_path_list() and history_path_list() as methods to all classes in
+# ACCEPTED_RUST_STRUCTS
+for item in ACCEPTED_RUST_STRUCTS:
+    setattr(getattr(altrios_pyo3, item), "variable_path_list", variable_path_list)
+    setattr(getattr(altrios_pyo3, item), "history_path_list", history_path_list)
+    setattr(getattr(altrios_pyo3, item), "to_pydict", to_pydict)
+    setattr(getattr(altrios_pyo3, item), "from_pydict", from_pydict)
+    setattr(getattr(altrios_pyo3, item), "to_dataframe", to_dataframe)
 
 setattr(ReversibleEnergyStorage, "from_excel", classmethod(_res_from_excel))  # noqa: F405
 setattr(Pyo3VecWrapper, "__array__", __array__)  # noqa: F405

python/altrios/demos/demo_variable_paths.py

@@ -0,0 +1,134 @@
+"""
+Script demonstrating how to use variable_path_list() and history_path_list()
+demos to find the paths to variables within altrios classes.
+"""
+import os
+import polars as pl
+import time
+
+import altrios as alt
+SAVE_INTERVAL = 100
+
+# https://docs.rs/altrios-core/latest/altrios_core/train/struct.TrainConfig.html
+train_config = alt.TrainConfig(
+    cars_empty=50,
+    cars_loaded=50,
+    rail_vehicle_type="Manifest",
+    train_type=alt.TrainType.Freight,
+    train_length_meters=None,
+    train_mass_kilograms=None,
+)
+
+# instantiate battery model
+# https://docs.rs/altrios-core/latest/altrios_core/consist/locomotive/powertrain/reversible_energy_storage/struct.ReversibleEnergyStorage.html#
+res = alt.ReversibleEnergyStorage.from_file(
+    alt.resources_root() / "powertrains/reversible_energy_storages/Kokam_NMC_75Ah_flx_drive.yaml"
+)
+# instantiate electric drivetrain (motors and any gearboxes)
+# https://docs.rs/altrios-core/latest/altrios_core/consist/locomotive/powertrain/electric_drivetrain/struct.ElectricDrivetrain.html
+edrv = alt.ElectricDrivetrain(
+    pwr_out_frac_interp=[0., 1.],
+    eta_interp=[0.98, 0.98],
+    pwr_out_max_watts=5e9,
+    save_interval=SAVE_INTERVAL,
+)
+
+bel = alt.Locomotive.build_battery_electric_loco(
+    reversible_energy_storage=res,
+    drivetrain=edrv,
+    loco_params=alt.LocoParams.from_dict(dict(
+        pwr_aux_offset_watts=8.55e3,
+        pwr_aux_traction_coeff_ratio=540.e-6,
+        force_max_newtons=667.2e3,
+)))
+
+# construct a vector of one BEL and several conventional locomotives
+loco_vec = [bel.clone()] + [alt.Locomotive.default()] * 7
+# instantiate consist
+loco_con = alt.Consist(
+    loco_vec,
+    SAVE_INTERVAL,
+)
+
+tsb = alt.TrainSimBuilder(
+    train_id="0",
+    origin_id="Minneapolis",
+    destination_id="Superior",
+    train_config=train_config,
+    loco_con=loco_con,
+)
+
+rail_vehicle_file = "rolling_stock/" + train_config.rail_vehicle_type + ".yaml"
+rail_vehicle = alt.RailVehicle.from_file(
+    alt.resources_root() / rail_vehicle_file)
+
+network = alt.Network.from_file(
+    alt.resources_root() / "networks/Taconite-NoBalloon.yaml")
+
+location_map = alt.import_locations(
+    alt.resources_root() / "networks/default_locations.csv")
+
+train_sim: alt.SpeedLimitTrainSim = tsb.make_speed_limit_train_sim(
+    rail_vehicle=rail_vehicle,
+    location_map=location_map,
+    save_interval=1,
+)
+train_sim.set_save_interval(SAVE_INTERVAL)
+
+est_time_net, _consist = alt.make_est_times(train_sim, network)
+
+timed_link_path = alt.run_dispatch(
+    network,
+    alt.SpeedLimitTrainSimVec([train_sim]),
+    [est_time_net],
+    False,
+    False,
+)[0]
+
+# uncomment this line to see example of logging functionality
+# alt.utils.set_log_level("DEBUG")
+
+t0 = time.perf_counter()
+train_sim.walk_timed_path(
+    network=network,
+    timed_path=timed_link_path,
+)
+t1 = time.perf_counter()
+print(f'Time to simulate: {t1 - t0:.5g}')
+assert len(train_sim.history) > 1
+
+# whether to run assertions, enabled by default
+ENABLE_ASSERTS = os.environ.get("ENABLE_ASSERTS", "true").lower() == "true"
+# whether to override reference files used in assertions, disabled by default
+ENABLE_REF_OVERRIDE = os.environ.get("ENABLE_REF_OVERRIDE", "false").lower() == "true"
+# directory for reference files for checking sim results against expected results
+ref_dir = alt.resources_root() / "demos/demo_variable_paths/"
+
+# print out all subpaths for variables in SimDrive
+print("List of variable paths for SimDrive:" + "\n".join(train_sim.variable_path_list()))
+if ENABLE_REF_OVERRIDE:
+    ref_dir.mkdir(exist_ok=True, parents=True)
+    with open(ref_dir / "variable_path_list_expected.txt", 'w') as f:
+        for line in train_sim.variable_path_list():
+            f.write(line + "\n")
+if ENABLE_ASSERTS:
+    print("Checking output of `variable_path_list()`")
+    with open(ref_dir / "variable_path_list_expected.txt", 'r') as f:
+        variable_path_list_expected = [line.strip() for line in f.readlines()]
+    assert variable_path_list_expected == train_sim.variable_path_list()
+    print("\n")
+
+# print out all subpaths for history variables in SimDrive
+print("List of history variable paths for SimDrive:" +  "\n".join(train_sim.history_path_list()))
+print("\n")
+
+# print results as dataframe
+print("Results as dataframe:\n", train_sim.to_dataframe(), sep="")
+if ENABLE_REF_OVERRIDE:
+    df:pl.DataFrame = train_sim.to_dataframe().lazy().collect()
+    df.write_csv(ref_dir / "to_dataframe_expected.csv")
+if ENABLE_ASSERTS:
+    print("Checking output of `to_dataframe`")
+    to_dataframe_expected = pl.scan_csv(ref_dir / "to_dataframe_expected.csv").collect()
+    assert to_dataframe_expected.equals(train_sim.to_dataframe())
+    print("Success!")