dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/config/common.py

@@ -0,0 +1,142 @@
+from typing import Any, Iterable
+
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.dimension.time import MeasurementType, TimeDimensionType, TimeIntervalType
+from dsgrid.dimension.standard import (
+    EnergyEfficiency,
+    EnergyEndUse,
+    EnergyIntensity,
+    EnergyIntensityRegression,
+    EnergyServiceDemand,
+    EnergyServiceDemandRegression,
+    FractionalIndex,
+    PeggedIndex,
+    Population,
+    Stock,
+    StockShare,
+    StockRegression,
+    WeatherVariable,
+)
+from dsgrid.exceptions import DSGInvalidParameter
+
+
+SUPPORTED_METRIC_TYPES = {
+    x.__name__
+    for x in (
+        EnergyEfficiency,
+        EnergyEndUse,
+        EnergyIntensity,
+        EnergyIntensityRegression,
+        EnergyServiceDemand,
+        EnergyServiceDemandRegression,
+        FractionalIndex,
+        PeggedIndex,
+        Population,
+        Stock,
+        StockShare,
+        StockRegression,
+        WeatherVariable,
+    )
+}
+
+DIMENSION_CLASS_MAP = {
+    DimensionType.GEOGRAPHY: "Geography",
+    DimensionType.MODEL_YEAR: "ModelYear",
+    DimensionType.SCENARIO: "Scenario",
+    DimensionType.SECTOR: "Sector",
+    DimensionType.SUBSECTOR: "Subsector",
+    DimensionType.TIME: "Time",
+    DimensionType.WEATHER_YEAR: "WeatherYear",
+}
+
+
+def make_base_dimension_template(
+    metric_types: Iterable[str],
+    exclude_dimension_types: set[DimensionType] | None = None,
+    time_type: TimeDimensionType | None = None,
+) -> list[dict[str, Any]]:
+    exclude: set[DimensionType] = exclude_dimension_types or set()
+    exclude.update({DimensionType.METRIC, DimensionType.TIME})
+
+    dimensions: list[dict[str, Any]] = []
+    for metric_type in metric_types:
+        if metric_type not in SUPPORTED_METRIC_TYPES:
+            msg = f"{metric_type=} is not one of the {SUPPORTED_METRIC_TYPES=}"
+            raise DSGInvalidParameter(msg)
+        dim = {
+            "type": DimensionType.METRIC.value,
+            "class": metric_type,
+            "name": DimensionType.METRIC.value,
+            "description": DimensionType.METRIC.value,
+            "file": f"dimensions/{metric_type}.csv",
+            "module": "dsgrid.dimension.standard",
+        }
+        dimensions.append(dim)
+
+    dimensions += [
+        {
+            "type": x.value,
+            "class": DIMENSION_CLASS_MAP[x],
+            "name": x.value,
+            "description": x.value,
+            "file": f"dimensions/{x.value}.csv",
+            "module": "dsgrid.dimension.standard",
+        }
+        for x in DimensionType
+        if x not in exclude
+    ]
+    if time_type is not None:
+        time_dim = make_base_time_dimension_template(time_type)
+        dimensions.append(time_dim)
+
+    return dimensions
+
+
+def make_base_time_dimension_template(time_type: TimeDimensionType) -> dict[str, Any]:
+    time_dim = {
+        "type": DimensionType.TIME.value,
+        "time_type": time_type.value,
+        "time_interval_type": TimeIntervalType.PERIOD_BEGINNING.value,
+        "name": time_type.value,
+        "description": time_type.value,
+        "module": "dsgrid.dimension.standard",
+    }
+    match time_type:
+        case TimeDimensionType.DATETIME:
+            time_dim["class"] = "Time"
+            time_dim["time_zone_format"] = {
+                "format_type": "aligned_in_absolute_time",
+                "time_zone": "Etc/GMT+5",
+            }
+            time_dim["measurement_type"] = MeasurementType.TOTAL.value
+            time_dim["ranges"] = [
+                {
+                    "start": "2018-01-01 00:00:00",
+                    "end": "2018-12-31 23:00:00",
+                    "frequency": "P0DT1H",
+                    "str_format": "%Y-%m-%d %H:%M:%S",
+                },
+            ]
+        case TimeDimensionType.ANNUAL:
+            time_dim["class"] = "AnnualTime"
+            time_dim["include_leap_day"] = True
+            time_dim["ranges"] = [
+                {
+                    "start": "2010",
+                    "end": "2024",
+                    "frequency": 1,
+                    "str_format": "%Y",
+                },
+            ]
+        case TimeDimensionType.INDEX:
+            time_dim["class"] = "IndexTime"
+            time_dim["ranges"] = [
+                {
+                    "start": 0,
+                    "end": 8759,
+                    "frequency": "P0DT1H",
+                    "starting_timestamp": "2018-01-01 00:00:00",
+                    "str_format": "%Y-%m-%d %H:%M:%S",
+                },
+            ]
+    return time_dim

dsgrid/config/config_base.py

@@ -0,0 +1,148 @@
+import abc
+import logging
+from pathlib import Path
+from typing import Type
+
+import json5
+
+from dsgrid.exceptions import DSGInvalidOperation
+from dsgrid.spark.types import (
+    DataFrame,
+)
+from dsgrid.utils.spark import models_to_dataframe
+
+
+logger = logging.getLogger(__name__)
+
+
+class ConfigBase(abc.ABC):
+    """Base class for all config classes"""
+
+    def __init__(self, model):
+        self._model = model
+
+    @classmethod
+    def load(cls, config_file, *args, **kwargs):
+        """Load the config from a file.
+
+        Parameters
+        ---------
+        config_file : str
+
+        Returns
+        -------
+        ConfigBase
+
+        """
+        # Subclasses can reimplement this method if they need more arguments.
+        return cls._load(config_file, *args, **kwargs)
+
+    @classmethod
+    def load_from_model(cls, model):
+        """Load the config from a model.
+
+        Parameters
+        ---------
+        model : DSGBaseModel
+
+        Returns
+        -------
+        ConfigBase
+
+        """
+        return cls(model)
+
+    @classmethod
+    def _load(cls, config_file):
+        model = cls.model_class().load(config_file)
+        return cls(model)
+
+    @staticmethod
+    @abc.abstractmethod
+    def config_filename() -> str:
+        """Return the config filename.
+
+        Returns
+        -------
+        str
+
+        """
+
+    @property
+    @abc.abstractmethod
+    def config_id(self) -> str:
+        """Return the configuration ID.
+
+        Returns
+        -------
+        str
+
+        """
+
+    @property
+    def model(self):
+        """Return the data model for the config.
+
+        Returns
+        -------
+        DSGBaseModel
+
+        """
+        return self._model
+
+    @staticmethod
+    @abc.abstractmethod
+    def model_class() -> Type:
+        """Return the data model class backing the config"""
+
+    def serialize(self, path, force=False):
+        """Serialize the configuration to a path.
+
+        path : str
+            Directory
+        force : bool
+            If True, overwrite files.
+
+        """
+        filename = Path(path) / self.config_filename()
+        if filename.exists() and not force:
+            msg = f"{filename} exists. Set force=True to overwrite."
+            raise DSGInvalidOperation(msg)
+        filename.write_text(self.model.model_dump_json(indent=2))
+        return filename
+
+
+class ConfigWithRecordFileBase(ConfigBase, abc.ABC):
+    """Intermediate-level base class to provide serialization of record files."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_records_dataframe(self) -> DataFrame:
+        """Return the records in a spark dataframe. Cached on first call."""
+        # id provides uniqueness and the config_id could help inspect what's in cache in case we
+        # ever need that.
+        # Spark doesn't allow dashes in the table name.
+        table_name = f"{self.config_id}__{id(self)}".replace("-", "_")
+        df = models_to_dataframe(self.model.records, table_name=table_name)
+        logger.debug("Loaded %s records dataframe", self.config_id)
+        return df
+
+    @classmethod
+    def load(cls, config_file):
+        config = super().load(config_file)
+        return config
+
+    def serialize(self, path, force=False):
+        dst_config_file = path / self.config_filename()
+        records_file = path / "records.csv"
+        for filename in (dst_config_file, records_file):
+            if filename.exists() and not force:
+                msg = f"{filename} exists. Set force=True to overwrite."
+                raise DSGInvalidOperation(msg)
+
+        self.get_records_dataframe().toPandas().to_csv(records_file, index=False)
+        model_data = self.model.serialize()
+        model_data["file"] = records_file.name
+        dst_config_file.write_text(json5.dumps(model_data, indent=2))
+        return dst_config_file