Merge pull request #28 from openmethane/allow-directory

Allow custom input directory

.env.example

@@ -2,7 +2,6 @@
 INPUTS=inputs
 OUTPUTS=outputs
 INTERMEDIATES=intermediates
-CMAQ_EXAMPLE=cmaq_example
 CLIMATETRACE=climateTrace
 FOSSIL=fossil_fuel_operations
 
@@ -34,6 +33,6 @@ TERMITES=termite_emissions_2010-2016.nc
 WETLANDS=DLEM_totflux_CRU_diagnostic.nc
 
 # These files are used to create the domain info file and ship in the cmaq_example folder
-CROFILE=GRIDCRO2D_1
-DOTFILE=GRIDDOT2D_1
-GEO_EM=geo_em.d01.nc
+CROFILE=tests/test-data/cmaq_example/GRIDCRO2D_1
+DOTFILE=tests/test-data/cmaq_example/GRIDDOT2D_1
+GEO_EM=tests/test-data/cmaq_example/geo_em.d01.nc

Makefile

@@ -26,7 +26,7 @@ download: ## Download the data for the project
 .PHONY: run
 run:  download ## Run the project for an example period
 	poetry run python scripts/omCreateDomainInfo.py
-	poetry run python scripts/omPrior.py 2022-07-01 2022-07-02
+	poetry run python scripts/omPrior.py 2022-07-01 2022-07-01
 
 .PHONY: ruff-fixes
 ruff-fixes:  # Run ruff on the project

scripts/omCreateDomainInfo.py

@@ -25,34 +25,66 @@
 
 from openmethane_prior.omInputs import croFilePath, domainPath, dotFilePath, geomFilePath
 
-root_path = Path(__file__).parent.parent
+root_path = Path(__file__).parents[1]
+# Root directory to use if relative paths are provided
 
-domainXr = xr.Dataset()
 
-with xr.open_dataset(os.path.join(root_path, geomFilePath)) as geomXr:
-    for attr in ["DX", "DY", "TRUELAT1", "TRUELAT2", "MOAD_CEN_LAT", "STAND_LON"]:
-        domainXr.attrs[attr] = geomXr.attrs[attr]
+def create_domain_info(
+    geometry_file: str,
+    cross_file: str,
+    dot_file: str,
+) -> xr.Dataset:
+    """
+    Create a new domain from the input WRF domain and subsets it to match the CMAQ domain
 
-with xr.open_dataset(os.path.join(root_path, croFilePath)) as croXr:
-    for var in ["LAT", "LON"]:
-        domainXr[var] = croXr[var]
-        domainXr[var] = croXr[var].squeeze(
+    Parameters
+    ----------
+    geometry_file
+        Path to the WRF geometry file
+    cross_file
+        Path to the MCIP cross file
+    dot_file
+        Path to the MCIP dot file
+
+    Returns
+    -------
+        The regridded domain information as an xarray dataset
+    """
+    domainXr = xr.Dataset()
+
+    with xr.open_dataset(geometry_file) as geomXr:
+        for attr in ["DX", "DY", "TRUELAT1", "TRUELAT2", "MOAD_CEN_LAT", "STAND_LON"]:
+            domainXr.attrs[attr] = geomXr.attrs[attr]
+
+    with xr.open_dataset(cross_file) as croXr:
+        for var in ["LAT", "LON"]:
+            domainXr[var] = croXr[var]
+            domainXr[var] = croXr[var].squeeze(
+                dim="LAY", drop=True
+            )  # copy but remove the 'LAY' dimension
+
+        domainXr["LANDMASK"] = croXr["LWMASK"].squeeze(
             dim="LAY", drop=True
         )  # copy but remove the 'LAY' dimension
 
-    domainXr["LANDMASK"] = croXr["LWMASK"].squeeze(
-        dim="LAY", drop=True
-    )  # copy but remove the 'LAY' dimension
-
-with xr.open_dataset(os.path.join(root_path, dotFilePath)) as dotXr:
-    # some repetition between the geom and grid files here, XCELL=DX and YCELL=DY
-    # - XCELL, YCELL: size of a single cell in m
-    # - XCENT, YCENT: lat/long of grid centre point
-    # - XORIG, YORIG: position of 0,0 cell in grid coordinates (in m)
-    for attr in ["XCELL", "YCELL", "XCENT", "YCENT", "XORIG", "YORIG"]:
-        domainXr.attrs[attr] = croXr.attrs[attr]
-    for var in ["LATD", "LOND"]:
-        domainXr[var] = dotXr[var].rename({"COL": "COL_D", "ROW": "ROW_D"})
-
-print(os.path.join(root_path, domainPath))
-domainXr.to_netcdf(os.path.join(root_path, domainPath))
+    with xr.open_dataset(dot_file) as dotXr:
+        # some repetition between the geom and grid files here, XCELL=DX and YCELL=DY
+        # - XCELL, YCELL: size of a single cell in m
+        # - XCENT, YCENT: lat/long of grid centre point
+        # - XORIG, YORIG: position of 0,0 cell in grid coordinates (in m)
+        for attr in ["XCELL", "YCELL", "XCENT", "YCENT", "XORIG", "YORIG"]:
+            domainXr.attrs[attr] = croXr.attrs[attr]
+        for var in ["LATD", "LOND"]:
+            domainXr[var] = dotXr[var].rename({"COL": "COL_D", "ROW": "ROW_D"})
+
+    return domainXr
+
+
+if __name__ == "__main__":
+    domain = create_domain_info(
+        geometry_file=os.path.join(root_path, geomFilePath),
+        cross_file=os.path.join(root_path, croFilePath),
+        dot_file=os.path.join(root_path, dotFilePath),
+    )
+    print(f"Writing domain to {os.path.join(root_path, domainPath)}")
+    domain.to_netcdf(os.path.join(root_path, domainPath))

src/openmethane_prior/__init__.py

@@ -2,4 +2,10 @@
 
 import importlib.metadata
 
+import dotenv
+
 __version__ = importlib.metadata.version("openmethane_prior")
+
+
+# Load environment variables from a local .env file
+dotenv.load_dotenv()

src/openmethane_prior/omInputs.py

@@ -28,7 +28,6 @@
 from openmethane_prior.omUtils import getenv
 
 inputsPath = getenv("INPUTS")
-cmaqExamplePath = getenv("CMAQ_EXAMPLE")
 climateTracePath = os.path.join(inputsPath, getenv("CLIMATETRACE"))
 fossilPath = os.path.join(climateTracePath, getenv("FOSSIL"))
 
@@ -52,9 +51,9 @@
 termitePath = os.path.join(inputsPath, getenv("TERMITES"))
 wetlandPath = os.path.join(inputsPath, getenv("WETLANDS"))
 
-croFilePath = os.path.join(cmaqExamplePath, getenv("CROFILE"))
-dotFilePath = os.path.join(cmaqExamplePath, getenv("DOTFILE"))
-geomFilePath = os.path.join(cmaqExamplePath, getenv("GEO_EM"))
+croFilePath = getenv("CROFILE")
+dotFilePath = getenv("DOTFILE")
+geomFilePath = getenv("GEO_EM")
 
 
 # list of layers that will be in the output file

src/openmethane_prior/omUtils.py

@@ -24,13 +24,11 @@
 import pickle
 import typing
 
-import dotenv
 import numpy as np
 from numpy.typing import ArrayLike
 
 T = typing.TypeVar("T", bound=ArrayLike | float)
 
-dotenv.load_dotenv()
 getenv = os.environ.get
 
 secsPerYear = 365 * 24 * 60 * 60

tests/conftest.py

@@ -1,8 +1,38 @@
+import os
 from pathlib import Path
 
+import dotenv
 import pytest
+import xarray as xr
 
 
 @pytest.fixture(scope="session")
 def root_dir():
     return Path(__file__).parent.parent
+
+
+# This fixture will be automatically used by all tests to setup the required env variables
+@pytest.fixture(autouse=True)
+def env(monkeypatch, root_dir):
+    initial_env = dict(os.environ)
+
+    # Use the example .env file to drive the tests
+    dotenv.load_dotenv(dotenv_path=root_dir / ".env.example", override=True)
+
+    yield
+
+    # Reset environment to initial state
+    os.environ.clear()
+    os.environ.update(initial_env)
+
+
+@pytest.fixture()
+def cro_xr(root_dir, env):
+    cro_file_path = os.path.join(root_dir, os.environ["CROFILE"])
+    return xr.open_dataset(cro_file_path)
+
+
+@pytest.fixture()
+def dot_xr(root_dir, env):
+    dot_file_path = os.path.join(root_dir, os.environ["DOTFILE"])
+    return xr.open_dataset(dot_file_path)

tests/test_create_domain_info.py

@@ -0,0 +1,42 @@
+import os
+
+import pytest
+import xarray as xr
+from openmethane_prior.omUtils import getenv
+from scripts.omCreateDomainInfo import create_domain_info
+
+
+@pytest.fixture()
+def geom_xr(root_dir):
+    geom_file_path = os.path.join(root_dir, getenv("GEO_EM"))
+    return xr.open_dataset(geom_file_path)
+
+
+@pytest.fixture()
+def input_domain_xr(root_dir):
+    return create_domain_info(
+        geometry_file=os.path.join(root_dir, getenv("GEO_EM")),
+        cross_file=os.path.join(root_dir, getenv("CROFILE")),
+        dot_file=os.path.join(root_dir, getenv("DOTFILE")),
+    )
+
+
+def test_grid_size_for_geo_files(cro_xr, geom_xr, dot_xr):
+    expected_cell_size = 10000
+
+    assert cro_xr.XCELL == expected_cell_size
+    assert cro_xr.YCELL == expected_cell_size
+
+    assert geom_xr.DX == expected_cell_size
+    assert geom_xr.DY == expected_cell_size
+
+    assert dot_xr.XCELL == expected_cell_size
+    assert dot_xr.YCELL == expected_cell_size
+
+
+def test_compare_in_domain_with_cro_dot_files(input_domain_xr, cro_xr, dot_xr):
+    assert dot_xr.NCOLS == input_domain_xr.COL_D.size
+    assert dot_xr.NROWS == input_domain_xr.ROW_D.size
+
+    assert cro_xr.NCOLS == input_domain_xr.COL.size
+    assert cro_xr.NROWS == input_domain_xr.ROW.size

tests/test_om_prior.py

@@ -7,34 +7,12 @@
 import pytest
 import requests
 import xarray as xr
-
 from openmethane_prior.layers.omGFASEmis import downloadGFAS
 from openmethane_prior.omInputs import livestockDataPath, sectoralEmissionsPath
-from openmethane_prior.omUtils import getenv, secsPerYear
+from openmethane_prior.omUtils import secsPerYear
 from scripts.omDownloadInputs import download_input_files, downloads, remote
 
 
-@pytest.fixture(scope="session")
-def cro_xr(root_dir):
-    cmaqExamplePath = getenv("CMAQ_EXAMPLE")
-    croFilePath = os.path.join(root_dir, cmaqExamplePath, getenv("CROFILE"))
-    return xr.open_dataset(croFilePath)
-
-
-@pytest.fixture(scope="session")
-def dot_xr(root_dir):
-    cmaqExamplePath = getenv("CMAQ_EXAMPLE")
-    dotFilePath = os.path.join(root_dir, cmaqExamplePath, getenv("DOTFILE"))
-    return xr.open_dataset(dotFilePath)
-
-
-@pytest.fixture(scope="session")
-def geom_xr(root_dir):
-    cmaqExamplePath = getenv("CMAQ_EXAMPLE")
-    geomFilePath = os.path.join(root_dir, cmaqExamplePath, getenv("GEO_EM"))
-    return xr.open_dataset(geomFilePath)
-
-
 # Fixture to download and later remove all input files
 @pytest.fixture(scope="session")
 def input_files(root_dir):
@@ -57,6 +35,7 @@ def input_domain_xr(root_dir, input_files):
 
     # Generated by scripts/omCreateDomainInfo.py
     filepath_in_domain = os.path.join(root_dir, "inputs/om-domain-info.nc")
+    assert os.path.exists(filepath_in_domain)
 
     yield xr.load_dataset(filepath_in_domain)
 
@@ -140,33 +119,6 @@ def test_005_agriculture_emissions(root_dir, input_files):
 
 # TODO Update this test when file structure is clear.
 # This test ensures that the grid size for all input files is 10 km.
-def test_006_grid_size_for_geo_files(cro_xr, geom_xr, dot_xr):
-    expected_cell_size = 10000
-
-    assert cro_xr.XCELL == expected_cell_size
-    assert cro_xr.YCELL == expected_cell_size
-
-    assert geom_xr.DX == expected_cell_size
-    assert geom_xr.DY == expected_cell_size
-
-    assert dot_xr.XCELL == expected_cell_size
-    assert dot_xr.YCELL == expected_cell_size
-
-
-def test_007_compare_in_domain_with_cro_dot_files(input_domain_xr, cro_xr, dot_xr):
-    assert dot_xr.NCOLS == input_domain_xr.COL_D.size
-    assert dot_xr.NROWS == input_domain_xr.ROW_D.size
-
-    assert cro_xr.NCOLS == input_domain_xr.COL.size
-    assert cro_xr.NROWS == input_domain_xr.ROW.size
-
-
-def test_008_compare_out_domain_with_cro_dot_files(output_domain_xr, cro_xr, dot_xr):
-    assert dot_xr.NCOLS == output_domain_xr.COL_D.size
-    assert dot_xr.NROWS == output_domain_xr.ROW_D.size
-
-    assert cro_xr.NCOLS == output_domain_xr.COL.size
-    assert cro_xr.NROWS == output_domain_xr.ROW.size
 
 
 def test_009_output_domain_xr(output_domain_xr, num_regression):
@@ -200,3 +152,11 @@ def test_010_emission_discrepancy(root_dir, output_domain_xr, input_files):
             assert (
                 abs(perectenageDifference) < 0.1
             ), f"Discrepency of {perectenageDifference}% in {sector} emissions"
+
+
+def test_compare_out_domain_with_cro_dot_files(output_domain_xr, cro_xr, dot_xr):
+    assert dot_xr.NCOLS == output_domain_xr.COL_D.size
+    assert dot_xr.NROWS == output_domain_xr.ROW_D.size
+
+    assert cro_xr.NCOLS == output_domain_xr.COL.size
+    assert cro_xr.NROWS == output_domain_xr.ROW.size