split main in two parts, one for preproc, one for tests

mvcm/main.py

@@ -168,19 +168,7 @@ def main(
                 try:
                     vnp02[b]
                 except KeyError:
-                    logger.info(f"Band {b} not found in file. No output will be written.")
-                    with Dataset(file_names["MOD03"]) as f:
-                        attrs = {attr: f.getncattr(attr) for attr in f.ncattrs()}
-                        shape = f["geolocation_data/latitude"][:].shape
-                    save_output(
-                        None,
-                        attrs,
-                        out_file,
-                        compression=5,
-                        debug=debug,
-                        shape=shape,
-                    )
-                    return
+                    logger.info(f"Band {b} not found in file. Fill values will be used instead.")
             logger.info(f"All bands found in file {file_names['MOD02']}. The code will run.")
 
     with Dataset(file_names["MOD03"]) as f:

mvcm/main_prep_only.py

+"""Main function for MVCM."""
+
+import argparse
+import logging
+import os
+from datetime import datetime as dt
+
+import numpy as np
+import psutil
+import xarray as xr
+from netCDF4 import Dataset  # type: ignore
+from pkg_resources import get_distribution  # type: ignore
+from rich.logging import RichHandler
+from ruamel.yaml import YAML
+
+import mvcm.read_data as rd
+from mvcm.constants import SensorConstants
+from mvcm.write_output import run_mvcm, save_output
+
+_LOG_FORMAT = "%(message)s"
+logging.basicConfig(level="NOTSET", datefmt="[%X]", format=_LOG_FORMAT, handlers=[RichHandler()])
+logger = logging.getLogger(__name__)
+LOG_LEVELS = ["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"]
+
+proc = psutil.Process(os.getpid())
+
+
+def main(
+    satellite: str = "snpp",
+    sensor: str = "viirs",
+    data_path: str = "",
+    mod02: str = "",
+    mod03: str = "",
+    img02: str | None = "",
+    img03: str | None = "",
+    threshold_file: str = "",
+    geos_atm_1: str = "",
+    geos_atm_2: str = "",
+    geos_land: str = "",
+    geos_ocean: str = "",
+    geos_constants: str = "",
+    ndvi_file: str = "",
+    sst_file: str = "",
+    eco_file: str = "",
+    out_file: str = "",
+    exe_path: str = "/home/pveglio/mvcm/mvcm",
+    verbose: int = 0,
+    *,
+    debug: bool = False,
+    hires_only: bool = False,
+    mvcm_threshold_file: str | None,
+) -> None:
+    """MVCM main function.
+
+    Parameters
+    ----------
+    satellite: str
+        satellite name, not case-sensitive. Available options: [snpp, ]
+    sensor: str
+        sensor name, not case-sensitive. Available options: [viirs, ]
+    data_path: str
+        path where the data is stored
+    mod02: str
+        VIIRS MOD02 file name
+    mod03: str
+        VIIRS MOD03 file name
+    img02: [optional] str
+        VIIRS IMG02 file name
+    img03: [optional] str
+        VIIRS IMG03 file name
+    threshold_file: str
+        thresholds file name
+    geos_atm_1: str
+        file name of first GEOS5 file for atmospheric parameters
+    goes_atm_2: str
+        file name of second GEOS5 file for atmospheric parameters
+    geos_land: str
+        file name of GEOS5 land parameters
+    geos_ocean: str
+        file name of GEOS5 ocean parameters
+    geos_constants: str
+        file name of GEOS5 constants
+    ndvi_file: str
+        NDVI file name
+    sst_file: str
+        Sea surface temperature file name
+    eco_file: str
+        Ecosystem File
+    out_file: str
+        Output file name
+    verbose: int
+        Verbosity level
+    exe_path: str
+        Path to the MVCM executable
+    debug: bool
+        Debug mode
+    hires_only: bool
+        Only run MVCM with high resolution
+
+    Returns
+    -------
+    None
+    """
+    verbose_level = np.minimum(verbose + 1, 4)
+    logger.setLevel(LOG_LEVELS[verbose_level])
+
+    if img02 is None or img03 is None:
+        use_hires = False
+    else:
+        use_hires = True
+
+    file_names = {
+        "MOD02": f"{mod02}",
+        "MOD03": f"{mod03}",
+        "IMG02": f"{img02}",
+        "IMG03": f"{img03}",
+        "GEOS_atm_1": f"{geos_atm_1}",
+        "GEOS_atm_2": f"{geos_atm_2}",
+        "GEOS_land": f"{geos_land}",
+        "GEOS_ocean": f"{geos_ocean}",
+        "GEOS_constants": f"{geos_constants}",
+        "NDVI": f"{ndvi_file}",
+        "SST": f"{sst_file}",
+        "ECO": f"{eco_file}",
+        "ANC_DIR": f"{data_path}",
+    }
+    logger.info(f"Memory usage #1: {proc.memory_info().rss / 1e6} MB")
+
+    if hires_only is False:
+        logger.info("Running regular MVCM before high resolution")
+        run_mvcm(file_names, mvcm_threshold_file, out_file, exe_path)
+
+    with open(threshold_file) as f:
+        text = f.read()
+    thresholds = YAML(typ="safe").load(text)
+
+    viirs_m_bands = SensorConstants.VIIRS_VIS_BANDS + SensorConstants.VIIRS_IR_BANDS
+
+    with Dataset(file_names["MOD03"]) as f:
+        attrs = {attr: f.getncattr(attr) for attr in f.ncattrs()}
+        lat_shape = f["geolocation_data/latitude"][:].shape
+
+    fmt_in = "%Y-%m-%dT%H:%M:%S.%fZ"
+    fmt_out = "%Y-%m-%dT%H%M"
+
+    # We are not processing night granules
+    if use_hires is True:
+        with xr.open_dataset(file_names["IMG02"], group="observation_data") as vnp02:
+            for b in SensorConstants.VIIRS_IMG_BANDS:
+                try:
+                    vnp02[b]
+                except KeyError:
+                    logger.info(f"Band {b} not found in file. Hires data will be fill values.")
+                    attrs = {"version": _VERSION, "satellite": "VIIRS", "sensor": "VIIRS"}
+                    save_output(None, attrs, out_file, compression=5, debug=debug, shape=lat_shape)
+                    return
+            logger.info(f"All bands found in file {file_names['IMG02']}. The code will run.")
+    else:
+        with xr.open_dataset(file_names["MOD02"], group="observation_data") as vnp02:
+            for b in viirs_m_bands:
+                try:
+                    vnp02[b]
+                except KeyError:
+                    logger.info(f"Band {b} not found in file. Hires data will be fill values.")
+                    attrs = {"version": _VERSION, "satellite": "VIIRS", "sensor": "VIIRS"}
+                    save_output(None, attrs, out_file, compression=5, debug=debug, shape=lat_shape)
+                    return
+            logger.info(f"All bands found in file {file_names['MOD02']}. The code will run.")
+
+    logger.info(f"Memory usage #2: {proc.memory_info().rss / 1e6} MB")
+    # rd.get_data(satellite, sensor, file_names, thresholds, hires=use_hires)
+    viirs_data, pixel_type = rd.get_data(satellite, sensor, file_names, thresholds, hires=use_hires)
+    # viirs_data = rd.get_data(satellite, sensor, file_names, thresholds, hires=use_hires)
+
+    # viirs_data = xr.open_dataset("input_data.nc")
+    logger.info(f"Memory usage #3: {proc.memory_info().rss / 1e6} MB")
+
+    with xr.open_dataset(file_names["MOD02"]) as vnp02:
+        time_str = dt.strftime(dt.strptime(vnp02.time_coverage_start, fmt_in), fmt_out)
+
+    viirs_data.to_netcdf(
+        f"{os.path.dirname(out_file)}/input_data_{time_str}.nc", mode="w", format="NETCDF4"
+    )
+    pixel_type.to_netcdf(
+        f"{os.path.dirname(out_file)}/pixel_type_{time_str}.nc", mode="w", format="NETCDF4"
+    )
+
+    logger.info(f"Memory usage #4: {proc.memory_info().rss / 1e6} MB")
+    ###
+
+
+if __name__ == "__main__":
+    _VERSION = get_distribution("mvcm").version
+    parser = argparse.ArgumentParser(prog="MVCM", description="")
+
+    parser.add_argument(
+        "--satellite",
+        help="satellite name, not case-sensitive. Available options: [snpp, ]",
+        choices=[
+            "snpp",
+        ],
+    )
+    parser.add_argument(
+        "--sensor",
+        help="sensor name, not case-sensitive. Available options: [viirs, ]",
+        choices=[
+            "viirs",
+        ],
+    )
+    parser.add_argument("--path", help="path where the data is stored")
+    parser.add_argument("--l1b", help="level 1b file name")
+    parser.add_argument("--geolocation", help="geolocation file name")
+    parser.add_argument("--hires-l1b", help="VIIRS IMG02 file name")
+    parser.add_argument("--hires-geo", help="VIIRS IMG03 file name")
+    parser.add_argument("-t", "--threshold", help="thresholds file name")
+    parser.add_argument(
+        "--atmos-1",
+        help="file name of the first GEOS-IT file for atmospheric parameters",
+    )
+    parser.add_argument(
+        "--atmos-2",
+        help="file name of the second GEOS-IT file for atmospheric parameters",
+    )
+    parser.add_argument("--land", help="file name of GEOS-IT land parameters")
+    parser.add_argument("--ocean", help="file name of GEOS-IT ocean parameters")
+    parser.add_argument("--constants", help="file name of GEOS-IT constants")
+    parser.add_argument("--ndvi", help="NDVI file name")
+    parser.add_argument("--sst", help="Sea surface temperature file name")
+    parser.add_argument("--eco", help="Ecosystem file")
+    parser.add_argument("-o", "--out", help="output file name")
+    parser.add_argument(
+        "-V",
+        "--version",
+        action="version",
+        version=_VERSION,
+        help="print version and exit",
+    )
+    parser.add_argument(
+        "-v", "--verbose", action="count", default=1, help="print verbose information"
+    )
+    parser.add_argument("-d", "--debug", action="store_true", help="activate debug mode")
+    parser.add_argument("--hires-only", action="store_true", help="run only high resolution code")
+    parser.add_argument("-x", "--mvcm-exe-path", help="path to mvcm executable")
+    parser.add_argument(
+        "--mvcm-thresholds", help="thresholds file name for operational MVCM", default=None
+    )
+
+    args = parser.parse_args()
+
+    satellite = args.satellite or "snpp"
+    sensor = args.sensor or "viirs"
+    data_path = args.path
+    mod02 = args.l1b
+    mod03 = args.geolocation
+    img02 = args.hires_l1b or None
+    img03 = args.hires_geo or None
+    threshold_file = args.threshold
+    geos_atm_1 = args.atmos_1
+    geos_atm_2 = args.atmos_2
+    geos_land = args.land
+    geos_ocean = args.ocean
+    constants = args.constants
+    ndvi_file = args.ndvi
+    sst_file = args.sst
+    eco_file = args.eco
+    out_file = args.out
+    verbose = args.verbose or False
+    debug = args.debug or False
+
+    main(
+        satellite=satellite,
+        sensor=sensor,
+        data_path=data_path,
+        mod02=mod02,
+        mod03=mod03,
+        img02=img02,
+        img03=img03,
+        threshold_file=threshold_file,
+        geos_atm_1=geos_atm_1,
+        geos_atm_2=geos_atm_2,
+        geos_land=geos_land,
+        geos_ocean=geos_ocean,
+        geos_constants=constants,
+        ndvi_file=ndvi_file,
+        sst_file=sst_file,
+        eco_file=eco_file,
+        out_file=out_file,
+        verbose=verbose,
+        debug=debug,
+        exe_path=args.mvcm_exe_path,
+        hires_only=args.hires_only,
+        mvcm_threshold_file=args.mvcm_thresholds,
+    )
+
+# tracemalloc.stop()

pyproject.toml

@@ -23,13 +23,14 @@ dependencies = [
     'wheel',
     'setuptools',
     'numpy==1.26',
-    'xarray',
+    'xarray==2023.6.0',
     'attrs',
     'netCDF4',
     'rich',
     'ruamel.yaml',
     'pre-commit',
     'dask',
+    'psutil',
     ]
 dynamic = ['version']