Merge pull request #82 from NREL/car_counting

Enable car-km counting

Author: calbaker

.gitignore

@@ -1217,3 +1217,10 @@ code\ meetings
 altrios/resources/networks/Scenic - ALTRIOS Confidential.yaml
 altrios/dispatch/*.json
 activate.sh
+*.xlsx
+altrios/*.txt
+truck_work_log.txt
+vehicle_average_times.txt
+avg_time_per_train.txt
+crane_work_log.txt
+hostler_work_log.txt

pyproject.toml

@@ -43,10 +43,11 @@ dependencies = [
     "plotly",
     "typing_extensions",
     "pyyaml",
-    "polars==0.20.25",
+    "polars >= 1.20.0",
     "pyarrow",
     "requests",
     "PyYAML==6.0.2",
+    "simpy",
     "msgpack==1.1.0",
 ]
 
@@ -55,7 +56,17 @@ homepage = "https://www.nrel.gov/transportation/altrios.html"
 source = "https://github.com/NREL/altrios"
 
 [project.optional-dependencies]
-dev = ["black", "pytest", "maturin", "ipykernel"]
+dev = [
+    "pytest~=8.3",
+    "maturin~=1.8",
+    "ipykernel~=6.29",
+    "python-lsp-server~=1.10",
+    "ruff~=0.9",
+    "pylsp-mypy~=0.7.0",
+    "memory_profiler~=0.61.0",
+    "pymoo~=0.6",
+    "plotly~=6.0",
+]
 
 [tool.maturin]
 profile = "release"

python/altrios/__init__.py

@@ -1,24 +1,22 @@
 from importlib.metadata import version
 __version__ = version("altrios")
 
-from pathlib import Path
-import re
 import numpy as np
+import re
 import logging
-import inspect
-from typing import List, Union, Dict, Optional
+from typing import List, Union, Dict, Optional, Any
 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
 from altrios.utilities import copy_demo_files  # noqa: F401
-from altrios import utilities as utils
-from altrios.utilities import package_root, resources_root
+from altrios import utilities as utils  # noqa: F401
+from altrios.utilities import package_root, resources_root  # noqa: F401
 # make everything in altrios_pyo3 available here
-from altrios.altrios_pyo3 import *
-from altrios import *
+from altrios.altrios_pyo3 import *  # noqa: F403
+from altrios import *  # noqa: F403
 
 DEFAULT_LOGGING_CONFIG = dict(
     format="%(asctime)s.%(msecs)03d | %(filename)s:%(lineno)s | %(levelname)s: %(message)s",
@@ -29,98 +27,17 @@
 logging.basicConfig(**DEFAULT_LOGGING_CONFIG)
 logger = logging.getLogger(__name__)
 
+
 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
+
+# creates a list of all python classes from rust structs that need python-side serde helpers
 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() 
+    attr for attr in altrios_pyo3.__dir__() if not attr.startswith("__") and isinstance(getattr(altrios_pyo3, attr), type) and  # noqa: F405
+    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/demo_data/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            
-
 # TODO connect to crate features
 data_formats = [
     'yaml',
@@ -129,6 +46,7 @@ def history_path_list(self, element_as_list:bool=False) -> List[str]:
     'json',
 ]
 
+
 def to_pydict(self, data_fmt: str = "msg_pack", flatten: bool = False) -> Dict:
     """
     Returns self converted to pure python dictionary with no nested Rust objects
@@ -140,10 +58,10 @@ def to_pydict(self, data_fmt: str = "msg_pack", flatten: bool = False) -> Dict:
     assert data_fmt in data_formats, f"`data_fmt` must be one of {data_formats}"
     match data_fmt:
         case "msg_pack":
-            import msgpack
+            import msgpack  # type: ignore[import-untyped]
             pydict = msgpack.loads(self.to_msg_pack())
         case "yaml":
-            from yaml import load
+            from yaml import load  # type: ignore[import-untyped]
             try:
                 from yaml import CLoader as Loader
             except ImportError:
@@ -156,10 +74,14 @@ def to_pydict(self, data_fmt: str = "msg_pack", flatten: bool = False) -> Dict:
     if not flatten:
         return pydict
     else:
-        return next(iter(pd.json_normalize(pydict, sep=".").to_dict(orient='records')))
+        hist_len = get_hist_len(pydict)
+        assert hist_len is not None, "Cannot be flattened"
+        flat_dict = get_flattened(pydict, hist_len)
+        return flat_dict
+
 
 @classmethod
-def from_pydict(cls, pydict: Dict, data_fmt: str = "msg_pack", skip_init: bool = True) -> Self:
+def from_pydict(cls, pydict: Dict, data_fmt: str = "msg_pack", skip_init: bool = False) -> Self:
     """
     Instantiates Self from pure python dictionary 
     # Arguments
@@ -176,56 +98,123 @@ def from_pydict(cls, pydict: Dict, data_fmt: str = "msg_pack", skip_init: bool =
             obj = cls.from_yaml(yaml.dump(pydict), skip_init=skip_init)
         case "msg_pack":
             import msgpack
-            try:
-                obj = cls.from_msg_pack(
-                    msgpack.packb(pydict), skip_init=skip_init)
-            except Exception as err:
-                print(
-                    f"{err}\nFalling back to YAML.")
-                obj = cls.from_pydict(
-                    pydict, data_fmt="yaml", skip_init=skip_init)
+            obj = cls.from_msg_pack(
+                msgpack.packb(pydict), skip_init=skip_init)
         case "json":
             from json import dumps
             obj = cls.from_json(dumps(pydict), skip_init=skip_init)
 
     return obj
 
-def to_dataframe(self, pandas:bool=False) -> [pd.DataFrame, pl.DataFrame, pl.LazyFrame]:
+
+def get_flattened(obj: Dict | List, hist_len: int, prepend_str: str = "") -> Dict:
+    """
+    Flattens and returns dictionary, separating keys and indices with a `"."`
+    # Arguments
+    # - `obj`: object to flatten
+    # -  hist_len: length of any lists storing history data
+    # - `prepend_str`: prepend this to all keys in the returned `flat` dict
     """
-    Returns time series results from altrios object as a Polars or Pandas dataframe.
+    flat: Dict = {}
+    if isinstance(obj, dict):
+        for (k, v) in obj.items():
+            new_key = k if (prepend_str == "") else prepend_str + "." + k
+            if isinstance(v, dict) or (isinstance(v, list) and len(v) != hist_len):
+                flat.update(get_flattened(v, hist_len, prepend_str=new_key))
+            else:
+                flat[new_key] = v
+    elif isinstance(obj, list):
+        for (i, v) in enumerate(obj):
+            new_key = i if (prepend_str == "") else prepend_str + "." + f"[{i}]"
+            if isinstance(v, dict) or (isinstance(v, list) and len(v) != hist_len):
+                flat.update(get_flattened(v, hist_len, prepend_str=new_key))
+            else:
+                flat[new_key] = v
+    else:
+        raise TypeError("`obj` should be `dict` or `list`")
+
+    return flat
+
+
+def get_hist_len(obj: Dict) -> Optional[int]:
+    """
+    Finds nested `history` and gets lenth of first element
+    """
+    if 'history' in obj.keys():
+        return len(next(iter(obj['history'].values())))
+
+    elif next(iter(k for k in obj.keys() if re.search("(history\\.\\w+)$", k) is not None), None) is not None:
+        return len(next((v for (k, v) in obj.items() if re.search("(history\\.\\w+)$", k) is not None)))
+
+    for (k, v) in obj.items():
+        if isinstance(v, dict):
+            hist_len = get_hist_len(v)
+            if hist_len is not None:
+                return hist_len
+    return None
+
+
+def to_dataframe(self, pandas: bool = False, allow_partial: bool = False) -> Union[pd.DataFrame, pl.DataFrame]:
+    """
+    Returns time series results from fastsim object as a Polars or Pandas dataframe.
 
     # Arguments
     - `pandas`: returns pandas dataframe if True; otherwise, returns polars dataframe by default
+    - `allow_partial`: tries to return dataframe of length equal to solved time steps if simulation fails early
     """
-    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)})
+    obj_dict = self.to_pydict(flatten=True)
+    history_keys = ['history.', 'speed_trace.', 'power_trace.']
+    hist_len = get_hist_len(obj_dict)
+    assert hist_len is not None
+
+    history_dict: Dict[str, Any] = {}
+    for k, v in obj_dict.items():
+        hk_in_k = any(hk in k for hk in history_keys)
+        if hk_in_k and ("__len__" in dir(v)):
+            if (len(v) == hist_len) or allow_partial:
+                history_dict[k] = v
+
+    if allow_partial:
+        cutoff = min([len(val) for val in history_dict.values()])
+
+        if not pandas:
+            try:
+                df = pl.DataFrame({col: val[:cutoff]
+                                   for col, val in history_dict.items()})
+            except Exception as err:
+                raise Exception(f"{err}\n`save_interval` may not be uniform")
+        else:
+            try:
+                df = pd.DataFrame({col: val[:cutoff]
+                                   for col, val in history_dict.items()})
+            except Exception as err:
+                raise Exception(f"{err}\n`save_interval` may not be uniform")
+
     else:
-        df = pd.DataFrame({col: val for col, val in zip(cols, vals)})
+        if not pandas:
+            try:
+                df = pl.DataFrame(history_dict)
+            except Exception as err:
+                raise Exception(
+                    f"{err}\nTry passing `allow_partial=True` to `to_dataframe` or checking for consistent save intervals")
+        else:
+            try:
+                df = pd.DataFrame(history_dict)
+            except Exception as err:
+                raise Exception(
+                    f"{err}\nTry passing `allow_partial=True` to `to_dataframe` or checking for consistent save intervals")
     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(getattr(altrios_pyo3, item), "to_pydict", to_pydict)  # noqa: F405
+    setattr(getattr(altrios_pyo3, item), "from_pydict", from_pydict)  # noqa: F405
+    setattr(getattr(altrios_pyo3, item), "to_dataframe", to_dataframe)  # noqa: F405
 
 setattr(ReversibleEnergyStorage, "from_excel", classmethod(_res_from_excel))  # noqa: F405
 setattr(Pyo3VecWrapper, "__array__", __array__)  # noqa: F405
-setattr(Pyo3Vec2Wrapper, "__array__", __array__)
-setattr(Pyo3Vec3Wrapper, "__array__", __array__)
-setattr(Pyo3VecBoolWrapper, "__array__", __array__)
+setattr(Pyo3Vec2Wrapper, "__array__", __array__)  # noqa: F405
+setattr(Pyo3Vec3Wrapper, "__array__", __array__)  # noqa: F405
+setattr(Pyo3VecBoolWrapper, "__array__", __array__)  # noqa: F405