dwind 0.3__py3-none-any.whl → 0.3.2__py3-none-any.whl

dwind/mp.py

@@ -1,12 +1,19 @@
+"""Provides the :py:class:`MultiProcess` class for running a model on `NREL's Kestrel HPC system`_.
+
+.. NREL's Kestrel HPC system: https://nrel.github.io/HPC/Documentation/Systems/Kestrel/
+"""
+
 from __future__ import annotations
 
 import time
 from pathlib import Path
 
 import pandas as pd
+from rich.live import Live
 from rex.utilities.hpc import SLURM
 
-from dwind.helper import split_by_index
+from dwind.utils import hpc
+from dwind.utils.array import split_by_index
 
 
 class MultiProcess:
@@ -116,7 +123,7 @@ class MultiProcess:
 
         # Create the output directory if it doesn't already exist
         self.dir_out = Path.cwd() if dir_out is None else Path(self.dir_out).resolve()
-        self.out_path = self.dir_out / f"chunk_files_{self.run_name}"
+        self.out_path = self.dir_out / "chunk_files"
         if not self.out_path.exists():
             self.out_path.mkdir()
 
@@ -127,78 +134,70 @@ class MultiProcess:
                 log_dir.mkdir()
             self.stdout_path = log_dir
 
-    def check_status(self, job_ids: list[int]):
+    def check_status(self, job_ids: list[int], start_time: float):
         """Prints the status of all :py:attr:`jobs` submitted.
 
         Parameters
         ----------
         job_ids : list[int]
             The list of HPC ``job_id``s to check on.
+        start_time : float
+            The results of initial ``time.perf_counter()``.
         """
-        hpc = SLURM()
-        print(f"{len(job_ids)} job(s) started")
-
-        jobs_status = {j: hpc.check_status(job_id=j) for j in job_ids}
-        n_remaining = len([s for s in jobs_status.values() if s in ("PD", "R")])
-        print(f"{n_remaining} job(s) remaining: {jobs_status}")
-        time.sleep(30)
-
-        while n_remaining > 0:
-            hpc = SLURM()
-            for job, status in jobs_status.items():
-                if status in ("GC", "None"):
-                    continue
-                jobs_status.update({job: hpc.check_status(job_id=job)})
-
-            n_remaining = len([s for s in jobs_status.values() if s in ("PD", "R")])
-            print(f"{n_remaining} job(s) remaining: {jobs_status}")
-            if n_remaining > 0:
-                time.sleep(30)
-
-    def aggregate_outputs(self):
-        """Collect the chunked results files, combine them into a single output parquet file, and
-        delete the chunked results files.
-        """
-        result_files = [f for f in self.out_path.iterdir() if f.suffix in (".pickle", ".pkl")]
-
-        if len(result_files) > 0:
-            result_agents = pd.concat([pd.read_pickle(f) for f in result_files])
-            f_out = self.dir_out / f"run_{self.run_name}.pqt"
-            result_agents.to_parquet(f_out)
-
-        for f in result_files:
-            f.unlink()
-
-    def run_jobs(self, agent_df: pd.DataFrame) -> None:
+        slurm = SLURM()
+        job_status = {
+            j: {
+                "status": slurm.check_status(job_id=j),
+                "start": start_time,
+                "wait": time.perf_counter() - start_time,
+                "run": 0,
+            }
+            for j in job_ids
+        }
+        table, complete = hpc.generate_run_status_table(job_status)
+        with Live(table, refresh_per_second=1) as live:
+            while not complete:
+                time.sleep(5)
+                job_status |= hpc.update_status(job_status)
+                table, complete = hpc.generate_run_status_table(job_status)
+                live.update(table)
+
+    def run_jobs(self, agent_df: pd.DataFrame) -> dict[str, int]:
         """Run :py:attr:`n_jobs` number of jobs for the :py:attr:`agent_df`.
 
-        Parameters
-        ----------
-        agent_df : pandas.DataFrame
-            The agent DataFrame to be chunked and analyzed.
+        Args:
+            agent_df (pandas.DataFrame): The agent DataFrame to be chunked and analyzed.
+
+        Returns:
+            dict[str, int]: Dictionary mapping of each SLURM job id to the chunk run in that job.
         """
         agent_df = agent_df.reset_index(drop=True)
         # chunks = np.array_split(agent_df, self.n_nodes)
         starts, ends = split_by_index(agent_df, self.n_nodes)
-        jobs = []
-
-        base_cmd_str = f"module load conda; conda activate {self.env}; "
-        base_cmd_str += "dwind run-chunk "
-
-        base_args = f"--location {self.location} "
-        base_args += f"--sector {self.sector} "
-        base_args += f"--scenario {self.scenario} "
-        base_args += f"--year {self.year} "
-        base_args += f"--repository {self.repository} "
-        base_args += f"--model-config {self.model_config} "
-        base_args += f"--out-path {self.out_path}"
-
-        for i, (start, end) in enumerate(zip(starts, ends, strict=True)):
-            fn = self.out_path / f"agents_{i}.pqt"
+        job_chunk_map = {}
+
+        base_cmd_str = f"module load conda; conda activate {self.env};"
+        base_cmd_str += " dwind run chunk"
+
+        base_args = f" {self.location}"
+        base_args += f" {self.sector}"
+        base_args += f" {self.scenario}"
+        base_args += f" {self.year}"
+        base_args += f" {self.out_path}"
+        base_args += f" {self.repository}"
+        base_args += f" {self.model_config}"
+
+        if not (agent_path := self.out_path / "agent_chunks").is_dir():
+            agent_path.mkdir()
+
+        start_time = time.perf_counter()
+        # for i, (start, end) in enumerate(zip(starts, ends, strict=True)):
+        for i, (start, end) in enumerate(zip(starts, ends)):  # noqa: B905
+            fn = self.out_path / "agent_chunks" / f"agents_{i}.pqt"
             agent_df.iloc[start:end].to_parquet(fn)
 
             job_name = f"{self.run_name}_{i}"
-            cmd_str = f"{base_cmd_str} --chunk-ix {i} {base_args}"
+            cmd_str = f"{base_cmd_str} {i} {base_args}"
             print("cmd:", cmd_str)
 
             slurm_manager = SLURM()
@@ -213,7 +212,7 @@ class MultiProcess:
             )
 
             if job_id:
-                jobs.append(job_id)
+                job_chunk_map[job_id] = i
                 print(f"Kicked off job: {job_name}, with SLURM {job_id=} on Eagle.")
             else:
                 print(
@@ -221,5 +220,6 @@ class MultiProcess:
                 )
 
         # Check on the job statuses until they're complete, then aggregate the results
-        self.check_status(jobs)
-        self.aggregate_outputs()
+        jobs = [*job_chunk_map]
+        self.check_status(jobs, start_time)
+        return job_chunk_map

dwind/resource.py

@@ -1,23 +1,75 @@
+"""Provides the :py:class:`ResourcePotential` class for gathering pre-calculated reV generation
+data.
+"""
+
 import h5py as h5
 import pandas as pd
 
-from dwind import Configuration
+from dwind.config import Sector, Technology, Configuration
 
 
 class ResourcePotential:
+    """Helper class designed to retrieve pre-calculated energy generation data from reV."""
+
     def __init__(
-        self, parcels, model_config: Configuration, tech="wind", application="fom", year="2018"
+        self,
+        parcels: pd.DataFrame,
+        model_config: Configuration,
+        sector: Sector,
+        tech: str = "wind",
+        year: int = 2018,
     ):
+        """Initializes the :py:class:`ResourcePotential` instance.
+
+        Args:
+            parcels (pd.DataFrame): The agent DataFrame containing at least the following columns:
+                "gid", "rev_gid_{tech}", "solar_az_tilt" (solar only), "azimuth_{sector}"
+                (solar only), "tilt_{tech}" (solar only), "turbine_class" (wind only),
+                "wind_turbine_kw" (wind only), and "turbine_height_m" (wind only).
+            model_config (Configuration): The pre-loaded model configuration data object containing
+                the requisite SQL, file, and configuration data.
+            sector (dwind.config.Sector): A valid sector instance.
+            tech (str, optional): One of "solar" or "wind". Defaults to "wind".
+            year (int, optional): Resource year for the reV lookup. Defaults to 2018.
+
+        Raises:
+            ValueError: Raised if :py:attr:`parcels:` is missing any of the required columns.
+        """
         self.df = parcels
-        self.tech = tech
-        self.application = application
+        self.tech = Technology(tech)
+        self.sector = sector
         self.year = year
         self.config = model_config
 
-        if self.tech not in ("wind", "solar"):
-            raise ValueError("`tech` must be one of 'solar' or 'wind'.")
-
-    def create_rev_gid_to_summary_lkup(self, configs, save_csv=True):
+        solar_cols = ("solar_az_tilt", f"azimuth_{self.sector.value}", f"tilt_{self.tech.value}")
+        # wind_cols = ("turbine_class", "wind_turbine_kw", "turbine_height_m")
+        wind_cols = ("wind_turbine_kw", "turbine_height_m")
+
+        if self.tech is Technology.WIND:
+            cols = wind_cols
+        elif self.tech is Technology.SOLAR:
+            cols = solar_cols
+
+        missing = set(cols).difference(self.df.columns.tolist())
+        if missing:
+            raise ValueError(f"`parcels` is missing the following columns: {', '.join(missing)}")
+
+    def create_rev_gid_to_summary_lkup(
+        self, configs: list[str], *, save_csv: bool = True
+    ) -> pd.DataFrame:
+        """Creates the reV summary tables based on the "gid" mappings in :py:attr:`parcels`.
+
+        Args:
+            configs (list[str]): The list of technology-specific configurations where the generation
+                data should be retrieved.
+            save_csv (bool, optional): If True, save the resulting lookup calculated from reV to the
+                reV folder definied in ``Configuration.rev.generation.{tech}_DIR``. Defaults to
+                True.
+
+        Returns:
+            pd.DataFrame: reV generation lookup table for the technology-specific configurations in
+                :py:attr:`configs`.
+        """
         config_dfs = []
         for c in configs:
             file_str = self.config.rev.DIR / f"rev_{c}_generation_{self.year}.h5"
@@ -30,10 +82,10 @@ class ResourcePotential:
 
             config_df = pd.concat([rev_index, gids, annual_energy, cf_mean], axis=1)
             config_df.columns = [
-                f"rev_index_{self.tech}",
-                f"rev_gid_{self.tech}",
-                f"{self.tech}_naep",
-                f"{self.tech}_cf",
+                f"rev_index_{self.tech.value}",
+                f"rev_gid_{self.tech.value}",
+                f"{self.tech.value}_naep",
+                f"{self.tech.value}_cf",
             ]
 
             config_df["config"] = c
@@ -43,84 +95,107 @@ class ResourcePotential:
 
         if save_csv:
             save_name = (
-                self.config.rev.generation[f"{self.tech}_DIR"]
-                / f"lkup_rev_gid_to_summary_{self.tech}_{self.year}.csv"
+                self.config.rev.generation[f"{self.tech.value}_DIR"]
+                / f"lkup_rev_gid_to_summary_{self.tech.value}_{self.year}.csv"
             )
             summary_df.to_csv(save_name, index=False)
 
         return summary_df
 
     def find_rev_summary_table(self):
-        if self.tech == "solar":
+        """Creates the generation summary data for each of the :py:attr:`tech`-specific
+        configurations specified in :py:attr:`config.rev.settings.{tech}`, then maps it to the
+        agent data (:py:attr:`parcels`), overwriting any previously computed data.
+        """
+        if self.tech is Technology.SOLAR:
             configs = self.config.rev.settings.solar
             config_col = "solar_az_tilt"
-            col_list = ["gid", f"rev_gid_{self.tech}", config_col]
-            self.df[config_col] = self.df[f"azimuth_{self.application}"].map(
+            col_list = ["gid", f"rev_gid_{self.tech.value}", config_col]
+            self.df[config_col] = self.df[f"azimuth_{self.sector.value}"].map(
                 self.config.rev.settings.azimuth_direction_to_degree
             )
             self.df[config_col] = (
-                self.df[config_col].astype(str) + "_" + self.df[f"tilt_{self.tech}"].astype(str)
+                self.df[config_col].astype(str)
+                + "_"
+                + self.df[f"tilt_{self.tech.value}"].astype(str)
             )
-        elif self.tech == "wind":
-            configs = self.rev.settings.wind
+        elif self.tech is Technology.WIND:
+            configs = self.config.rev.settings.wind
             config_col = "turbine_class"
             col_list = [
                 "gid",
-                f"rev_gid_{self.tech}",
+                f"rev_gid_{self.tech.value}",
                 config_col,
                 "turbine_height_m",
                 "wind_turbine_kw",
             ]
             self.df[config_col] = self.df["wind_turbine_kw"].map(self.config.rev.turbine_class_dict)
 
-        out_cols = [*col_list, f"rev_index_{self.tech}", f"{self.tech}_naep", f"{self.tech}_cf"]
+        out_cols = [
+            *col_list,
+            f"rev_index_{self.tech.value}",
+            f"{self.tech.value}_naep",
+            f"{self.tech.value}_cf",
+        ]
+
+        drop_cols = [
+            f"rev_gid_{self.tech.value}",
+            f"{self.tech.value}_naep",
+            f"{self.tech.value}_cf",
+        ]
+        self.df = self.df.drop(columns=[c for c in drop_cols if c in self.df])
 
         f_gen = (
-            self.config.rev.generation[f"{self.tech}_DIR"]
-            / f"lkup_rev_gid_to_summary_{self.tech}_{self.year}.csv"
+            self.config.rev.generation[f"{self.tech.value}_DIR"]
+            / f"lkup_rev_gid_to_summary_{self.tech.value}_{self.year}.csv"
         )
 
         if f_gen.exists():
-            generation_summary = pd.read_csv(f_gen)
+            generation_summary = pd.read_csv(f_gen, dtype_backend="pyarrow")
         else:
             generation_summary = self.create_rev_gid_to_summary_lkup(configs)
 
         generation_summary = (
             generation_summary.reset_index(drop=True)
-            .drop_duplicates(subset=[f"rev_index_{self.tech}", "config"])
+            .drop_duplicates(subset=[f"rev_index_{self.tech.value}", "config"])
             .rename(columns={"config": config_col})
         )
         agents = self.df.merge(
-            generation_summary, how="left", on=[f"rev_index_{self.tech}", config_col]
+            generation_summary, how="left", on=[f"rev_index_{self.tech.value}", config_col]
         )
         return agents[out_cols]
 
     def prepare_agents_for_gen(self):
-        # create lookup column based on each tech
-        if self.tech == "wind":
+        """Create lookup column based on each technology."""
+        if self.tech is Technology.WIND:
             # drop wind turbine size duplicates
             # SINCE WE ASSUME ANY TURBINE IN A GIVEN CLASS HAS THE SAME POWER CURVE
             self.df.drop_duplicates(subset=["gid", "wind_size_kw"], keep="last", inplace=True)
-            # if running FOM application, only consider a single (largest) turbine size
-            if self.application == "fom":
+            # if running FOM sector, only consider a single (largest) turbine size
+            if self.sector is Sector.FOM:
                 self.df = self.df.loc[self.df["wind_size_kw"] == self.df["wind_size_kw_fom"]]
 
             self.df["turbine_class"] = self.df["wind_turbine_kw"].map(
                 self.config.rev.turbine_class_dict
             )
 
-        if self.tech == "solar":
-            # NOTE: tilt and azimuth are application-specific
-            self.df["solar_az_tilt"] = self.df[f"azimuth_{self.application}"].map(
+        if self.tech is Technology.SOLAR:
+            # NOTE: tilt and azimuth are sector-specific
+            self.df["solar_az_tilt"] = self.df[f"azimuth_{self.sector.value}"].map(
                 self.config.rev.settings.azimuth_direction_to_degree
             )
             self.df["solar_az_tilt"] = self.df["solar_az_tilt"].astype(str)
             self.df["solar_az_tilt"] = (
-                self.df["solar_az_tilt"] + "_" + self.df[f"tilt_{self.application}"].astype(str)
+                self.df["solar_az_tilt"] + "_" + self.df[f"tilt_{self.sector.value}"].astype(str)
             )
 
-    def merge_gen_to_agents(self, tech_agents):
-        if self.tech == "wind":
+    def merge_gen_to_agents(self, tech_agents: pd.DataFrame):
+        """Merges :py:attr:`tech_agents` to the parcel data :py:attr:`df`.
+
+        Args:
+            tech_agents (pd.DataFrame): The technology-specific energy generation data.
+        """
+        if self.tech is Technology.WIND:
             cols = ["turbine_height_m", "wind_turbine_kw", "turbine_class"]
         else:
             # NOTE: need to drop duplicates in solar agents
@@ -130,16 +205,23 @@ class ResourcePotential:
             )
             cols = ["solar_az_tilt"]
 
-        cols.extend(["gid", f"rev_index_{self.tech}"])
+        cols.extend(["gid", f"rev_index_{self.tech.value}"])
 
         self.df = self.df.merge(tech_agents, how="left", on=cols)
 
     def match_rev_summary_to_agents(self):
+        """Runs the energy generation gathering and merging steps, and retursns back the updated
+        :py:attr:`df` agent/parcel data.
+
+        Returns:
+            pd.DataFrame: Updated agent/parcel data with rec/alculated "wind_aep" or "solar_aep"
+                information for each agent.
+        """
         self.prepare_agents_for_gen()
         tech_agents = self.find_rev_summary_table()
         self.merge_gen_to_agents(tech_agents)
 
-        if self.tech == "wind":
+        if self.tech is Technology.WIND:
             # fill nan generation values
             self.df = self.df.loc[
                 ~((self.df["wind_naep"].isnull()) & (self.df["turbine_class"] != "none"))
@@ -150,7 +232,7 @@ class ResourcePotential:
             # calculate annual energy production (aep)
             self.df["wind_aep"] = self.df["wind_naep"] * self.df["wind_turbine_kw"]
             # self.df = self.df.drop(columns="turbine_class")
-        else:
+        elif self.tech is Technology.SOLAR:
             # fill nan generation values
             self.df = self.df.loc[~(self.df["solar_naep"].isnull())]
             # size groundmount system to equal wind aep

dwind/run.py

@@ -3,19 +3,26 @@
 from __future__ import annotations
 
 import sys
-import tomllib
 from enum import Enum
-from typing import Annotated
+from typing import Optional, Annotated
 from pathlib import Path
 
 import typer
 import pandas as pd
+from rich.style import Style
+from rich.pretty import pprint
+from rich.console import Console
 
 
-# from memory_profiler import profile
-
+# fmt: off
+if sys.version_info >= (3, 11):  # noqa
+    import tomllib
+else:
+    import tomli as tomllib
+# fmt: on
 
 app = typer.Typer()
+console = Console()
 
 DWIND = Path("/projects/dwind/agents")
 
@@ -71,20 +78,25 @@ def year_callback(ctx: typer.Context, param: typer.CallbackParam, value: int):
 
 
 def load_agents(
-    file_name: str | Path | None = None,
+    file_name: Path | None = None,
     location: str | None = None,
     sector: str | None = None,
+    model_config: str | Path | None = None,
+    *,
+    prepare: bool = False,
 ) -> pd.DataFrame:
     """Load the agent file based on a filename or the location and sector to a Pandas DataFrame,
     and return the data frame.
 
     Args:
-        file_name (str | Path | None, optional): Name of the agent file, if not auto-generating from
+        file_name (Path | None, optional): Name of the agent file, if not auto-generating from
             the :py:attr:`location` and :py:attr:`sector` inputs. Defaults to None.
         location (str | None, optional): The name of the location or grouping, such as
             "colorado_larimer" or "priority1". Defaults to None.
         sector (str | None, optional): The name of the section. Must be one of "btm" or "fom".
             Defaults to None.
+        prepare (bool, optional): True if loading pre-chunked and prepared agent data, which should
+            bypass the standard column checking for additional joins, by default False.
 
     Returns:
         pd.DataFrame: The agent DataFrame.
@@ -97,10 +109,28 @@ def load_agents(
     f_agents = (
         file_name if file_name is not None else DWIND / f"{location}/agents_dwind_{sector}.parquet"
     )
+    if not isinstance(f_agents, Path):
+        f_agents = Path(f_agents).resolve()
+
+    alternative_suffix = (".pqt", ".parquet", ".pkl", ".pickle", ".csv")
+    base_style = Style.parse("cyan")
     if not f_agents.exists():
-        f_agents = f_agents.with_suffix(".pickle")
-    agents = Agents(agent_file=f_agents).agents
-    return agents
+        for suffix in alternative_suffix:
+            if (new_fn := f_agents.with_suffix(suffix)).exists():
+                if new_fn != f_agents:
+                    msg = (
+                        f"Using alternative agent file: {new_fn}\n\t"
+                        f"Requested agent file: {f_agents}"
+                    )
+                    console.print(msg, style=base_style)
+                    f_agents = new_fn
+                    break
+
+    if prepare:
+        return Agents.load_and_prepare_agents(
+            agent_file=f_agents, sector=sector, model_config=model_config
+        )
+    return Agents.load_agents(agent_file=f_agents)
 
 
 @app.command()
@@ -155,7 +185,9 @@ def run_hpc(
     dir_out: Annotated[
         str, typer.Option(help="Path to where the chunked outputs should be saved.")
     ],
-    stdout_path: Annotated[str | None, typer.Option(help="The path to write stdout logs.")] = None,
+    stdout_path: Annotated[
+        Optional[str], typer.Option(help="The path to write stdout logs.")  # noqa
+    ] = None,
 ):
     """Run dwind via the HPC multiprocessing interface."""
     sys.path.append(repository)
@@ -180,7 +212,9 @@ def run_hpc(
         stdout_path=stdout_path,
     )
 
-    agent_df = load_agents(location=location, sector=sector)
+    agent_df = load_agents(
+        location=location, sector=sector, model_config=model_config, prepare=True
+    )
     mp.run_jobs(agent_df)
 
 
@@ -194,15 +228,15 @@ def run_hpc_from_config(
     config_path = Path(config_path).resolve()
     with config_path.open("rb") as f:
         config = tomllib.load(f)
-    print(config)
+    print("Running the following configuration:")
+    pprint(config)
 
     run_hpc(**config)
 
 
 @app.command()
 def run_chunk(
-    # start: Annotated[int, typer.Option(help="chunk start index")],
-    # end: Annotated[int, typer.Option(help="chunk end index")],
+    chunk_ix: Annotated[int, typer.Option(help="Chunk number/index. Used for logging.")],
     location: Annotated[
         str, typer.Option(help="The state, state_county, or priority region to run.")
     ],
@@ -213,7 +247,6 @@ def run_chunk(
     year: Annotated[
         int, typer.Option(callback=year_callback, help="The year basis of the scenario.")
     ],
-    chunk_ix: Annotated[int, typer.Option(help="Chunk number/index. Used for logging.")],
     out_path: Annotated[str, typer.Option(help="save path")],
     repository: Annotated[
         str, typer.Option(help="Path to the dwind repository to use when running the model.")
@@ -230,7 +263,7 @@ def run_chunk(
     from dwind.model import Model
 
     agent_file = Path(out_path).resolve() / f"agents_{chunk_ix}.pqt"
-    agents = load_agents(file_name=agent_file)
+    agents = load_agents(file_name=agent_file, prepare=False)
     agent_file.unlink()
 
     model = Model(
@@ -271,7 +304,7 @@ def run(
     sys.path.append(repository)
     from dwind.model import Model
 
-    agents = load_agents(location=location, sector=sector)
+    agents = load_agents(location=location, sector=sector, model_config=model_config, prepare=True)
     model = Model(
         agents=agents,
         location=location,