From 5dcfeede299a8003214f11f9283c09e1baac7aec Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Sat, 20 Sep 2025 03:07:24 -0400
Subject: [PATCH 01/21] download complete

---
 src/mlglobal/cli/ic_download_main.py |  46 ++-
 src/mlglobal/ic_downloader.py        | 431 ++++++++++++++-------------
 src/mlglobal/logger.py               |  39 +++
 3 files changed, 282 insertions(+), 234 deletions(-)
 create mode 100644 src/mlglobal/logger.py

diff --git a/src/mlglobal/cli/ic_download_main.py b/src/mlglobal/cli/ic_download_main.py
index 4df6886..b5684da 100644
--- a/src/mlglobal/cli/ic_download_main.py
+++ b/src/mlglobal/cli/ic_download_main.py
@@ -1,23 +1,31 @@
 import argparse
+import logging
 import os
 from datetime import datetime
 
 from mlglobal.ic_downloader import ICDownloader
+from mlglobal.logger import setup_logging
+
 
 # Default bucket and root directory for each mode  # TODO: store this in a yaml or some config file
 DEFAULTS = {
     "gfs": {
         "bucket_name": "noaa-gfs-bdp-pds",
-        "root_directory": "gdas",
+        "bucket_root_directory": "",
+        "comroot": "/lfs/h1/ops/prod/com/gfs/v16.3",
     },
     "gefs": {
         "bucket_name": "noaa-ncepdev-none-ca-ufs-cpldcld",
-        "root_directory": "gefs",
+        "bucket_root_directory": "Linlin.Cui/gefs_wcoss2",
+        "comroot": "/lfs/h1/ops/prod/com/gefs/v12.3",
     },
 }
 
 
 def main():
+
+    setup_logging()
+
     parser = argparse.ArgumentParser(description="Download IC data for GFS or GEFS")
 
     subparsers = parser.add_subparsers(
@@ -26,27 +34,12 @@ def main():
 
     def _common_args(inparser, dict_in):
         inparser.add_argument(
-            "--start_date",
-            help="Start datetime",
+            "--current_cycle",
+            help="Datetime to download data for in YYYYMMDDHH format",
             type=str,
             metavar="YYYYMMDDHH",
             required=True,
         )
-        inparser.add_argument(
-            "--end_date",
-            help="End datetime",
-            type=str,
-            metavar="YYYYMMDDHH",
-            required=True,
-        )
-        inparser.add_argument(
-            "--levels",
-            help="number of pressure levels",
-            type=int,
-            choices=[13, 37],
-            default=13,
-            required=False,
-        )
         inparser.add_argument(
             "--source",
             help="Data source",
@@ -73,7 +66,7 @@ def _common_args(inparser, dict_in):
             "--root-directory",
             help="Root directory",
             type=str,
-            default=dict_in["root_directory"],
+            default=dict_in["bucket_root_directory"],
             required=False,
         )
         return inparser
@@ -85,27 +78,26 @@ def _common_args(inparser, dict_in):
     # GEFS subparser
     gefs_parser = subparsers.add_parser("gefs", help="Download GEFS ensemble data")
     gefs_parser = _common_args(gefs_parser, DEFAULTS["gefs"])
+    gefs_members = ["c00"] + [f"p{str(i).zfill(2)}" for i in range(1, 31)]
     gefs_parser.add_argument(
         "--member",
         help="Ensemble member",
-        type=int,
-        choices=list(range(0, 31)),
+        type=str,
+        choices=gefs_members,
         default=0,
     )
 
     args = parser.parse_args()
 
     downloader = ICDownloader(
-        mode=args.mode,
-        start_datetime=datetime.strptime(args.start_date, "%Y%m%d%H"),
-        end_datetime=datetime.strptime(args.end_date, "%Y%m%d%H"),
+        current_cycle=datetime.strptime(args.current_cycle, "%Y%m%d%H"),
         member=None if args.mode == "gfs" else args.member,
         download_source=args.source,
-        download_directory=args.target,
+        local_directory=args.target,
         bucket_name=args.bucket_name,
         root_directory=args.root_directory,
     )
-    downloader.download()
+    downloader.get_data()
 
 
 if __name__ == "__main__":
diff --git a/src/mlglobal/ic_downloader.py b/src/mlglobal/ic_downloader.py
index cfdcef3..0ed0be8 100644
--- a/src/mlglobal/ic_downloader.py
+++ b/src/mlglobal/ic_downloader.py
@@ -1,78 +1,154 @@
-import glob
 import os
 import shutil
-from datetime import datetime, timedelta
+from datetime import timedelta
+from logging import getLogger
 
 
-class FileFormats:
-    def __init__(self, mode, num_levels=13):
+logger = getLogger(__name__)
 
-        FILE_FORMATS = {"gfs": self.gfs_file_formats, "gefs": self.gefs_file_formats}
-        self.num_levels = num_levels
-        self.file_formats = FILE_FORMATS[mode]()
 
-    def gfs_file_formats(self):
-        # List of file formats to download
-        if self.num_levels == 13:
-            file_formats = ["pgrb2.0p25.f000", "pgrb2.0p25.f006"]  # , '0p25.f001'
-        else:
-            file_formats = [
-                "pgrb2.0p25.f000",
-                "pgrb2b.0p25.f000",
-                "pgrb2.0p25.f006",
-            ]  # , '0p25.f001'
+class FileLookup:
+    def __init__(self, current_cycle, member=None):
 
-        return file_formats
+        self.current_cycle = current_cycle
+        self.member = member  # GEFS member values are c00, p01, p02, ..., p30
 
-    def gefs_file_formats(self):
+        # Look back 6 and 12 hours for precip files
+        self.current_cycle_m6h = self.current_cycle - timedelta(hours=6)
+        self.current_cycle_m12h = self.current_cycle - timedelta(hours=12)
 
-        # List of file formats to download
-        if self.num_levels == 13:
-            file_formats = ["pgrb2.0p25.f000", "pgrb2s.0p25.f000"]  # , '0p25.f001'
+        if self.member is not None:
+            self.template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{member}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
+            self.get_file_info = self._gefs_file_info
         else:
-            file_formats = [
-                "pgrb2.0p25.f000",
-                "pgrb2b.0p25.f000",
-                "pgrb2.0p25.f006",
-            ]  # , '0p25.f001'
-
-        return file_formats
+            self.template = f"gfs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/gfs.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
+            self.get_file_info = self._gfs_file_info
+
+    def _gfs_file_info(self):
+
+        file_formats =[
+            "pgrb2.0p25.f000",
+            "pgrb2b.0p25.f000",
+            "pgrb2.0p25.f006"
+        ]
+
+        # From current cycle
+        pgrb2_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec="pgrb2.0p25", fhour=0)
+        pgrb2b_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec="pgrb2b.0p25", fhour=0)
+
+        # From current cycle - 6 hours
+        pgrb2_0p25_f000_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=0)
+        pgrb2_0p25_f006_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=6)
+
+        # From current cycle - 12 hours
+        pgrb2_0p25_f006_m12 = self.template.format(cycle=self.current_cycle_m12h, fspec="pgrb2.0p25", fhour=6)
+
+        file_dict = {}
+        file_dict[self.current_cycle] = {"pgrb2.0p25.f000": pgrb2_0p25_f000,
+                                         "pgrb2b.0p25.f000": pgrb2b_0p25_f000}
+        file_dict[self.current_cycle_m6h] = {"pgrb2.0p25.f000": pgrb2_0p25_f000_m6,
+                                             "pgrb2.0p25.f006": pgrb2_0p25_f006_m6}
+        file_dict[self.current_cycle_m12h] = {"pgrb2.0p25.f006": pgrb2_0p25_f006_m12}
+
+        file_list = [
+            pgrb2_0p25_f000,
+            pgrb2b_0p25_f000,
+            pgrb2_0p25_f000_m6,
+            pgrb2_0p25_f006_m6,
+            pgrb2_0p25_f006_m12
+        ]
+
+        return file_dict, file_list, file_formats
+
+    def _gefs_file_info(self):
+
+        file_formats = [
+            "pgrb2.0p25.f000",
+            "pgrb2s.0p25.f000",
+            "pgrb2s.0p25.f006"
+        ]
+
+        # From current cycle
+        pgrb2_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
+        pgrb2s_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
+
+        # From current cycle - 6 hours
+        pgrb2_0p25_f000_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
+        pgrb2s_0p25_f000_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
+        pgrb2s_0p25_f006_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=6)
+
+        # From current cycle - 12 hours
+        pgrb2s_0p25_f006_m12 = self.template.format(cycle=self.current_cycle_m12h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=6)
+
+        file_dict = {}
+        file_dict[self.current_cycle] = {"pgrb2.0p25.f000": pgrb2_0p25_f000,
+                                         "pgrb2s.0p25.f000": pgrb2s_0p25_f000}
+        file_dict[self.current_cycle_m6h] = {"pgrb2.0p25.f000": pgrb2_0p25_f000_m6,
+                                             "pgrb2s.0p25.f000": pgrb2s_0p25_f000_m6,
+                                             "pgrb2s.0p25.f006": pgrb2s_0p25_f006_m6}
+        file_dict[self.current_cycle_m12h] = {"pgrb2s.0p25.f006": pgrb2s_0p25_f006_m12}
+
+        file_list = [
+            pgrb2_0p25_f000,
+            pgrb2s_0p25_f000,
+            pgrb2_0p25_f000_m6,
+            pgrb2s_0p25_f000_m6,
+            pgrb2s_0p25_f006_m6,
+            pgrb2s_0p25_f006_m12
+        ]
+
+        return file_dict, file_list, file_formats
 
 
 class ICDownloader:
 
     def __init__(
         self,
-        mode,
-        start_datetime,
-        end_datetime,
+        current_cycle,
         member=None,
-        num_pressure_levels=13,
-        download_source="s3",
-        download_directory=None,
+        download_source="local",
+        local_directory="./data",
         bucket_name=None,
-        root_directory=None,
+        root_directory=None
     ):
-        self.mode = mode
-        self.start_datetime = start_datetime
-        self.end_datetime = end_datetime
+        self.current_cycle = current_cycle
         self.member = member
-        self.num_pressure_levels = num_pressure_levels
         self.download_source = download_source
-        self.download_directory = download_directory
+        self.local_directory = local_directory
         self.bucket_name = bucket_name
         self.root_directory = root_directory
 
-        ff = FileFormats(mode, num_levels=self.num_pressure_levels)
-        self.file_formats = ff.file_formats
+        # Generate the lookup dictionary
+        lookup = FileLookup(self.current_cycle, member=self.member)
+        self.file_dict, self.file_list, self.file_formats = lookup.get_file_info()
+
+        if self.download_source in ["s3"]:
+            aws_profile = os.environ.get("AWS_PROFILE", "default")
+            self.s3 = self.get_s3_client_by_bucket_type(self.bucket_name, profile_name=aws_profile)
 
-        self.s3 = self.init_s3_client() if self.download_source == "s3" else None
+        os.makedirs(self.local_directory, exist_ok=True)
 
     @staticmethod
-    def init_s3_client():
+    def get_s3_client_by_bucket_type(bucket_name, profile_name='default'):
+        """
+        Initializes and returns a boto3 S3 client for a given bucket.
+
+        The function first attempts to get a public client. If that fails, it
+        assumes the bucket is private and creates a client using the specified
+        AWS profile.
+
+        Args:
+            bucket_name (str): The name of the S3 bucket.
+            profile_name (str): The AWS profile to use for private buckets.
+                                Defaults to 'default'.
+
+        Returns:
+            boto3.client: A configured S3 client.
+        """
 
         try:
             import boto3
+            from botocore.exceptions import ClientError
             from botocore import UNSIGNED
             from botocore.config import Config
         except ImportError as ee:
@@ -80,75 +156,71 @@ def init_s3_client():
                 "boto3 and botocore are required for S3 operations."
             ) from ee
 
+        # 1. Try to get a client configured for a public (unsigned) bucket
+        s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED))
         try:
-            # Try to create S3 client using profile method
-            profile_name = os.environ.get("AWS_PROFILE", "default")
-            session = boto3.Session(profile_name=profile_name)
-            current_credentials = session.get_credentials().get_frozen_credentials()
-            s3 = session.client(
-                "s3",
-                aws_access_key_id=current_credentials.access_key,
-                aws_secret_access_key=current_credentials.secret_key,
-            )
-        except Exception as e1:
-            print(f"Failed to create S3 client with profile method: {e1}")
-            try:
-                # Try to create S3 client using unsigned method
-                s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED))
-            except Exception as e2:
-                print(f"Failed to create S3 client with unsigned method: {e2}")
-                raise RuntimeError(
-                    "Failed to create S3 client with unsigned method."
-                ) from e2
-
-        return s3
-
-    def get_s3_specs(self, ymd, hh, file_format):
-
-        if self.mode == "gefs":
-
-            s3_prefix = (
-                f"Linlin.Cui/gefs_wcoss2/{self.root_directory}.{ymd}/{hh}/atmos/"
-            )
-            s3_file_format = f"{self.member:02d}.t{hh}z.{file_format}"
-
-        elif self.mode == "gfs":
+            # Check if the bucket can be accessed publicly without authentication
+            s3.head_bucket(Bucket=bucket_name)
+            logger.info(f"Bucket '{bucket_name}' is public. Returning an unsigned client.")
+            return s3
+        except ClientError as ee:
+            error_code = ee.response['Error']['Code']
+            # 2. If it's a 403 Forbidden, the bucket is likely private.
+            if error_code in ('403', '404'):
+                logger.warning(f"Bucket '{bucket_name}' is not public. Returning client with profile '{profile_name}'.")
+
+                # Create a session with the specified profile
+                session = boto3.Session(profile_name=profile_name)
+                current_credentials = session.get_credentials().get_frozen_credentials()
+                s3 = session.client(
+                    "s3",
+                    aws_access_key_id=current_credentials.access_key,
+                    aws_secret_access_key=current_credentials.secret_key,
+                )
+                return s3
+            else:
+                # Handle other errors, such as a non-existent bucket
+                logger.error(f"Error accessing bucket '{bucket_name}': {ee}")
+                return None
 
-            if file_format == "pgrb2.0p25.f006":
-                # get prefix for precip from the previous cycle
-                # Convert ymd and hh to datetime object
-                datetime_obj = datetime.strptime(ymd + hh, "%Y%m%d%H")
+    @staticmethod
+    def get_s3_objects(s3, bucket_name: str, prefix: str) -> list:
 
-                # Get the datetime 6 hours before
-                datetime_before = datetime_obj - timedelta(hours=6)
+        objects = []
+        continuation_token = None
 
-                # Get the date string and time string from datetime objects
-                ymd_precip = datetime_before.strftime("%Y%m%d")
-                hh_precip = datetime_before.strftime("%H")
+        while True:
+            list_kwargs = {
+                'Bucket': bucket_name,
+                'Prefix': prefix,
+                'MaxKeys': 1000  # Explicitly set MaxKeys, though it's default
+            }
+            if continuation_token:
+                list_kwargs['ContinuationToken'] = continuation_token
 
-                # Construct the S3 prefix for the directory
-                s3_prefix = f"{self.root_directory}.{ymd_precip}/{hh_precip}/"
+            response = s3.list_objects_v2(**list_kwargs)
 
-            else:
+            if 'Contents' in response:
+                objects.extend(response['Contents'])
 
-                s3_prefix = f"{self.root_directory}.{ymd}/{hh}/"
+            if not response.get('IsTruncated'):
+                # No more objects to retrieve, or less than 1000 objects in total
+                break
 
-            s3_file_format = file_format
+            continuation_token = response.get('NextContinuationToken')
+            if not continuation_token:
+                # Should not happen if 'IsTruncated' is True, but as a safeguard
+                break
 
-        return s3_prefix, s3_file_format
+        return objects
 
-    def get_data_from_s3(
-        self, path_prefix: str, file_format: str, local_directory: str
-    ) -> None:
+    def get_data_from_s3(self, file_list: list, local_directory: str) -> None:
         """
-        Downloads files with a specific format from an S3 bucket to a local directory.
-
+        Download files from S3 bucket to a local directory.
         Parameters
         ----------
-        prefix : str
-            The prefix (folder path) in the S3 bucket to filter objects.
-        file_format : str
-            The file extension or format to filter files (e.g., '.csv', '.json').
+        file_list : list
+            A list of files that need to be downloaded
         local_directory : str
             The local directory path where the downloaded files will be saved.
 
@@ -157,126 +229,71 @@ def get_data_from_s3(
         None
             This function does not return anything. Files are downloaded as a side effect.
 
-        Notes
-        -----
-        Only files ending with the specified `file_format` will be downloaded.
+        Raises
+        ------
+        Exception
+            If the file download operation fails.
         """
 
-        objects = self.s3.list_objects_v2(Bucket=self.bucket_name, Prefix=path_prefix)
-        for obj in objects.get("Contents", []):
-            obj_key = obj["Key"]
-            if obj_key.endswith(f"{file_format}"):
-                local_file_path = os.path.join(
-                    local_directory, os.path.basename(obj_key)
-                )
-                if not os.path.exists(local_file_path):
-                    self.s3.download_file(self.bucket_name, obj_key, local_file_path)
-                    print(f"Downloaded {obj_key} to {local_file_path}")
-                else:
-                    print(f"File {local_file_path} already exists, skipping download.")
+        logger.info(f"Downloading files from S3 bucket: {self.bucket_name}")
+        logger.info(f"Downloading files to {local_directory}")
 
-    def get_local_specs(self, ymd: str, hh: str, file_format: str) -> tuple[str, str]:
-        """
-        Get the local specifications for a given date, time, and file format.
+        for file_name in file_list:
+            local_file_path = os.path.join(local_directory, os.path.basename(file_name))
+            if os.path.exists(local_file_path):
+                logger.warning(f"File already exists, skipping: {file_name}")
+                continue
+            file_name_in_bucket = self.root_directory + "/" + file_name if self.root_directory else file_name
+            try:
+                self.s3.download_file(self.bucket_name, file_name_in_bucket, local_file_path)
+                logger.info(f"Downloaded:  {file_name} -> {local_directory}")
+            except Exception as ee:
+                logger.error(f"Error downloading {file_name}: {ee}")
+
+    def get_data_from_local(self, file_list: list, local_directory: str) -> None:
+        """Copy files from a local directory to another local directory.
 
         Parameters
         ----------
-        ymd : str
-            The date string in the format 'YYYYMMDD'.
-        hh : str
-            The time string in the format 'HH'.
-        file_format : str
-            The file format string (e.g., 'pgrb2.0p25.f006').
+        file_list : list
+            A list of files that need to be copied.
+        local_directory : str
+            The local directory path where the copied files will be saved.
 
         Returns
         -------
-        tuple[str, str]
-            A tuple containing the local path prefix and the local file format.
-        """
-        # TODO: elevate the hard-coded paths to class constructor (or above)
-        if self.mode == "gefs":
-            gefs_com_dir = "/lfs/h2/emc/da/noscrub/rahul.mahajan/mldata"  # For testing on Acorn
-            gefs_com_dir = "/lfs/h2/emc/ptmp/jun.wang"  # NCO does not mirror all of GEFS data to dev in RT
-            gefs_com_dir = "/lfs/h1/ops/prod/com/gefs/v12.3"
-            fprefix = file_format.split('.')[0]
-            local_prefix = f"{gefs_com_dir}/gefs.{ymd}/{hh}/atmos/{fprefix}p25"
-            local_file_format = f"{self.member:02d}.t{hh}z.{file_format}"
-
-        elif self.mode == "gfs":
-
-            # TODO: elevate the hard-coded paths to constructor
-            gfs_com_dir = "/lfs/h2/emc/da/noscrub/rahul.mahajan/mldata"  # For testing on Acorn
-            gfs_com_dir = "/lfs/h1/ops/prod/com/gfs/v16.3"
-            if file_format == "pgrb2.0p25.f006":
-                # get prefix for precip from the previous cycle
-                # Convert ymd and hh to datetime object
-                datetime_obj = datetime.strptime(ymd + hh, "%Y%m%d%H")
-
-                # Get the datetime 6 hours before
-                datetime_before = datetime_obj - timedelta(hours=6)
-
-                # Get the date string and time string from datetime objects
-                ymd_precip = datetime_before.strftime("%Y%m%d")
-                hh_precip = datetime_before.strftime("%H")
-
-                # Construct the S3 prefix for the directory
-                local_prefix = f"{gfs_com_dir}/gfs.{ymd_precip}/{hh_precip}/atmos"
-                local_file_format = f"gfs.t{hh_precip}z.{file_format}"
-
-            else:
-
-                local_prefix = f"{gfs_com_dir}/gfs.{ymd}/{hh}/atmos"
-                local_file_format = f"gfs.t{hh}z.{file_format}"
-
-        return local_prefix, local_file_format
-
-    def get_data_from_local(
-        self, path_prefix: str, file_format: str, local_directory: str
-    ) -> None:
+        None
+            This function does not return anything. Files are copied as a side effect.
 
-        file_objects = glob.glob(f"{path_prefix}/*")
-        for obj_key in file_objects:
-            if obj_key.endswith(f"{file_format}"):
+        Raises
+        ------
+        OSError
+            If the file copy operation fails.
+        """
 
-                # Define the local file path
-                local_file_path = os.path.join(
-                    local_directory, os.path.basename(obj_key)
-                )
+        logger.info(f"Copying files from directory: {self.root_directory}")
+        logger.info(f"Copying files to {local_directory}")
 
-                # Copy data to the local path
-                try:
-                    shutil.copy2(obj_key, local_file_path)
-                    print(f"Copied:  {obj_key} -> {local_directory}")
-                except OSError:
-                    raise OSError(f"Unable to copy {obj_key} to {local_directory}")
+        for file_name in file_list:
+            local_file_path = os.path.join(local_directory, os.path.basename(file_name))
+            if os.path.exists(local_file_path):
+                logger.warning(f"File already exists, skipping: {file_name}")
+                continue
+            file_name_in_remote = self.root_directory + "/" + file_name if self.root_directory else file_name
+            try:
+                shutil.copy2(file_name_in_remote, local_file_path)
+                logger.info(f"Copied:  {file_name} -> {local_directory}")
+            except OSError:
+                logger.error(f"Unable to copy {file_name} to {local_directory}")
+                raise OSError(f"Unable to copy {file_name} to {local_directory}")
 
         return
 
-    def download(self, loop_interval=6):
-
-        _SPECS_MAP = {"s3": self.get_s3_specs, "local": self.get_local_specs}
-        _GET_DATA_MAP = {"s3": self.get_data_from_s3, "local": self.get_data_from_local}
-
-        interval_dt = timedelta(hours=loop_interval)
-
-        # Loop through the intervals
-        current_datetime = self.start_datetime
-        while current_datetime <= self.end_datetime:
-            ymd = current_datetime.strftime("%Y%m%d")
-            hh = current_datetime.strftime("%H")
-
-            # Define the local directory path where the file will be saved
-            local_directory = os.path.join(self.download_directory, ymd, hh)
-
-            # Create the local directory if it doesn't exist
-            os.makedirs(local_directory, exist_ok=True)
-
-            # Loop over file formats and download data
-            for file_format in self.file_formats:
-                prefix, fformat = _SPECS_MAP[self.download_source](ymd, hh, file_format)
-                _GET_DATA_MAP[self.download_source](prefix, fformat, local_directory)
+    def get_data(self) -> None:
 
-            # Move to the next interval
-            current_datetime += interval_dt
+        GET_DATA_MAP = {"s3": self.get_data_from_s3,
+                        "local": self.get_data_from_local}
 
-        print("Download completed.")
+        logger.info("Starting download...")
+        GET_DATA_MAP[self.download_source](self.file_list, self.local_directory)
+        logger.info("Download completed.")
diff --git a/src/mlglobal/logger.py b/src/mlglobal/logger.py
new file mode 100644
index 0000000..c2cb966
--- /dev/null
+++ b/src/mlglobal/logger.py
@@ -0,0 +1,39 @@
+import logging
+import os
+"""
+Logger setup module for the mlglobal project.
+
+This module configures the logging settings for the entire project using Python's built-in
+logging module. It sets up a default logging format, date format, and log level, and provides
+a logger instance named "mlglobal" for use throughout the project.
+
+Functions
+---------
+setup_logger(level=logging.INFO, fmt="[%(asctime)s] %(levelname)8s - %(message)s", datefmt="%Y-%m-%dT%H:%M:%S")
+    Configures the logging system with the specified log level, format, and date format.
+
+Attributes
+----------
+logger : logging.Logger
+"""
+
+
+log = logging.getLogger()
+
+def setup_logging(
+    level=os.environ.get("LOGGING_LEVEL", logging.INFO),
+    fmt="[%(asctime)s] %(levelname)8s - %(message)s",
+    datefmt="%Y-%m-%dT%H:%M:%S"
+):
+
+    logger = logging.getLogger()
+    for handler in logger.handlers:
+        logger.removeHandler(handler)
+
+    kwargs: dict = {
+        "datefmt": datefmt,
+        "format": fmt,
+        "level": level
+    }
+
+    logging.basicConfig(**kwargs)

From 05317198e1d512a1259287f438a1718d869bb08f Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Tue, 23 Sep 2025 00:47:50 -0400
Subject: [PATCH 02/21] wip

---
 src/mlglobal/__init__.py |   2 +
 src/mlglobal/gen_ic.py   | 281 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 283 insertions(+)
 create mode 100644 src/mlglobal/gen_ic.py

diff --git a/src/mlglobal/__init__.py b/src/mlglobal/__init__.py
index 6e5f509..5127c24 100644
--- a/src/mlglobal/__init__.py
+++ b/src/mlglobal/__init__.py
@@ -14,3 +14,5 @@
 __license__ = "CC-0"
 
 __all__ = ["__version__", "__author__", "__email__", "__license__"]
+
+from .logger import setup_logging
diff --git a/src/mlglobal/gen_ic.py b/src/mlglobal/gen_ic.py
new file mode 100644
index 0000000..4ffd0da
--- /dev/null
+++ b/src/mlglobal/gen_ic.py
@@ -0,0 +1,281 @@
+import yaml
+import os
+import logging
+import grib2io
+import xarray as xr
+import numpy as np
+from datetime import datetime
+from pprint import pprint
+import fnmatch
+
+logger = logging.getLogger(__name__)
+
+class PrepareIC:
+    """Class to prepare initial conditions."""
+
+    def __init__(self,
+                 current_cycle: datetime,
+                 varinfo_file: str,
+                 filelist: list = [],
+                 data_dir: str = "./data",
+                 output_netcdf: str = "ic.nc") -> None:
+
+        self.current_cycle = current_cycle
+        self.varinfo = self.get_var_info(varinfo_file)
+        self.input_file_list = filelist
+        self.data_dir = data_dir
+        self.output_netcdf = output_netcdf
+
+    def get_var_info(self, yaml_file) -> dict:
+
+        with open(yaml_file, "r") as fh:
+            variables_to_extract = yaml.safe_load(fh)
+
+        return variables_to_extract
+
+    def process_files(self):
+
+        file_patterns = list(self.varinfo.keys())
+
+        mergeDSs = []
+        mergeDAs = []
+
+        # Files are all valid at the same time, so we can merge them along the time dimension
+        for file in self.input_file_list:
+            matched = False
+            for pattern in file_patterns:
+                if fnmatch.fnmatch(file, os.path.join(self.data_dir, '*' + pattern)):
+                    logger.debug(f"Matched pattern: {pattern} in file: {file}")
+                    matched = True
+                    break
+
+            if not matched:
+                logger.warning(f"No pattern matched for file: {file}")  # TODO: should this raise an error?
+
+            logger.info(f"Processing file: {os.path.basename(file)}")
+
+            gribfh = grib2io.open(file)
+
+            for var_pattern, details in self.varinfo[pattern].items():
+
+                variable_names = var_pattern.split(", ")
+                levels = details.get("level", None)
+
+                for var_name in variable_names:
+                    logger.info(f"Extracting variable: {var_name} at levels: {levels}")
+                    if len(levels) > 1:
+                        da = self.get_dataarray_3d(gribfh, var_name, levels)
+                    else:
+                        da = self.get_dataarray_2d(gribfh, var_name, levels[0])
+                    mergeDAs.append(da)
+
+            gribfh.close()
+
+        ds = xr.merge(mergeDAs, compat="no_conflicts")
+        mergeDSs.append(ds)
+        ds.close()
+
+        ds = xr.concat(mergeDSs, dim="time")
+
+        # Get 2D static data from the f000 file
+        # From the file list find the first file that ends with .f000
+        f000_file = next((f for f in self.input_file_list if f.endswith(".f000")), None)
+        if f000_file is None:
+            logger.warning("No 00Z f000 file found.")  # TODO: should this raise an error?
+            return
+
+        gribfh = grib2io.open(os.path.join(self.data_dir, f000_file))
+        static_vars = ["LAND", "HGT"]
+        for var_name in static_vars:
+            da = self.get_dataarray_2d(gribfh, var_name, "surface")
+            ds = xr.merge([ds, da], compat="no_conflicts")
+        gribfh.close()
+
+        ds = ds.rename({
+            "LAND_surface": "land_sea_mask",
+            "HGT_surface": "geopotential_at_surface",
+            "PRMSL_meansealevel": "mean_sea_level_pressure",
+            "TMP_2maboveground": "2m_temperature",
+            "UGRD_10maboveground": "10m_u_component_of_wind",
+            "VGRD_10maboveground": "10m_v_component_of_wind",
+            "APCP_surface": "total_precipitation_6hr",
+            "HGT": "geopotential",
+            "TMP": "temperature",
+            "SPFH": "specific_humidity",
+            "VVEL": "vertical_velocity",
+            "UGRD": "u_component_of_wind",
+            "VGRD": "v_component_of_wind",
+        })
+
+        ds = ds.assign_coords(datetime=ds.time)
+
+        # Adjust time values relative to the first time step
+        ds["time"] = ds["time"] - ds["time"][0]
+
+        # Expand dimensions
+        ds = ds.expand_dims(dim="batch")
+        ds["datetime"] = ds["datetime"].expand_dims(dim="batch")
+
+        # Squeeze dimensions
+        ds["geopotential_at_surface"] = ds["geopotential_at_surface"].squeeze("batch")
+        ds["land_sea_mask"] = ds["land_sea_mask"].squeeze("batch")
+
+        # Update geopotential unit to m2/s2 by multiplying 9.80665
+        ds["geopotential_at_surface"] = ds["geopotential_at_surface"] * 9.80665
+        ds["geopotential"] = ds["geopotential"] * 9.80665
+
+        # Update total_precipitation_6hr unit to (m) from (kg/m^2) by dividing it by 1000kg/m³
+        ds["total_precipitation_6hr"] = ds["total_precipitation_6hr"] / 1000.0
+
+        ds.to_netcdf(self.output_netcdf)
+        ds.close()
+
+        return
+
+    @staticmethod
+    def get_dataarray_2d(grbfile, var_name, desired_level):
+
+        logger.info(f"Getting 2D data for variable: {var_name} at level: {desired_level}")
+
+        msg = grbfile.select(shortName=var_name, level=desired_level)[0]
+
+        # create a netcdf dataset using the matching grib message
+        lats, lons = msg.latlons()
+        lats = lats[:,0]
+        lons = lons[0,:]
+
+        #check latitude range
+        reverse_lat = False
+        if lats[0] > 0:
+            reverse_lat = True
+
+        steps = msg.validDate
+        data = msg.data
+        if reverse_lat:
+            data = data[::-1, :]
+            lats = lats[::-1]
+
+        var_name2 = f'{var_name}_{"".join(desired_level.split())}'
+        if len(data.shape) == 2:
+            da = xr.Dataset(
+                data_vars={
+                    var_name2: (["lat", "lon"], data.astype("float32"))
+                },
+                coords={
+                    "lon": lons.astype("float32"),
+                    "lat": lats.astype("float32"),
+                    "time": steps,
+                }
+            )
+        elif len(data.shape) == 3:
+            print("3D data found in 2D function")
+            da = xr.Dataset(
+                data_vars={
+                    var_name: (["level", "lat", "lon"], data.astype("float32"))
+                },
+                coords={
+                    "lon": lons.astype("float32"),
+                    "lat": lats.astype("float32"),
+                    "level": np.array(desired_level).astype("int32"),
+                    "time": steps,
+                }
+            )
+
+        return da
+
+    @staticmethod
+    def get_dataarray_3d(grbfile, var_name, desired_level):
+
+        logger.info(f"Getting 3D data for variable: {var_name} at levels: {desired_level}")
+        data, levels = [], []
+        for ii, level in enumerate(desired_level):
+            msg = grbfile.select(shortName=var_name, level=level)[0]
+
+            if ii == 0:
+                lats, lons = msg.latlons()
+                lats = lats[:,0]
+                lons = lons[0,:]
+
+                #check latitude range, graphcast needs [-90, 90]
+                reverse_lat = False
+                if lats[0] > 0:
+                    reverse_lat = True
+                steps = msg.validDate
+            data.append(msg.data)
+            levels.append(int(level.split(' ')[0]))
+
+        data = np.array(data)
+        if reverse_lat:
+            data = data[:, ::-1, :]
+            lats = lats[::-1]
+
+        da = xr.Dataset(
+            data_vars={
+                var_name: (["level", "lat", "lon"], data.astype("float32"))
+            },
+            coords={
+                "lon": lons.astype("float32"),
+                "lat": lats.astype("float32"),
+                "level": np.array(levels).astype("int32"),
+                "time": steps,
+            }
+        )
+
+        return da
+
+if __name__ == "__main__":
+
+    import argparse
+    from datetime import datetime
+    from mlglobal.logger import setup_logging
+
+    setup_logging()
+
+    parser = argparse.ArgumentParser(
+        description="Process data to generate initial conditions for graphcast"
+        )
+
+    parser.add_argument(
+            "--current-cycle",
+            help="Current cycle in YYYYMMDDHH format",
+            type=str,
+            metavar="YYYYMMDDHH",
+            required=True
+            )
+    parser.add_argument(
+            "--yaml",
+            help="YAML file containing variable information",
+            type=str,
+            required=True
+            )
+    parser.add_argument(
+            "--data-dir",
+            help="Directory containing input GRIB2 files",
+            type=str,
+            default="./data",
+            required=False
+            )
+    parser.add_argument(
+            "--output",
+            help="Output NetCDF file",
+            type=str,
+            default="ic.nc",
+            required=False
+            )
+
+    args = parser.parse_args()
+
+    input_file_list = [os.path.join(args.data_dir, "gfs.t00z.pgrb2.0p25.f000"),
+                       os.path.join(args.data_dir, "gfs.t18z.pgrb2.0p25.f006")]
+    #input_file_list = glob.glob(os.path.join(args.data_dir, "*"))
+
+    prep = PrepareIC(
+        current_cycle=datetime.strptime(args.current_cycle, "%Y%m%d%H"),
+        varinfo_file=args.yaml,
+        filelist=input_file_list,
+        output_netcdf=args.output,
+        data_dir=args.data_dir
+    )
+
+    logger.info("Starting IC preparation...")
+    prep.process_files()

From ca0e2bb8c3e95ed212a72b64d4d211215b14dc78 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Tue, 23 Sep 2025 00:55:49 -0400
Subject: [PATCH 03/21] wip

---
 src/mlglobal/gen_ic.py        |  7 ++++---
 src/mlglobal/gfs_varinfo.yaml | 21 +++++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)
 create mode 100644 src/mlglobal/gfs_varinfo.yaml

diff --git a/src/mlglobal/gen_ic.py b/src/mlglobal/gen_ic.py
index 4ffd0da..4b75366 100644
--- a/src/mlglobal/gen_ic.py
+++ b/src/mlglobal/gen_ic.py
@@ -56,10 +56,11 @@ def process_files(self):
 
             gribfh = grib2io.open(file)
 
-            for var_pattern, details in self.varinfo[pattern].items():
+            #for var_pattern, details in self.varinfo[pattern].items():
+            for var_dict in self.varinfo[pattern]:
 
-                variable_names = var_pattern.split(", ")
-                levels = details.get("level", None)
+                variable_names = var_dict["variables"]
+                levels = var_dict["levels"]
 
                 for var_name in variable_names:
                     logger.info(f"Extracting variable: {var_name} at levels: {levels}")
diff --git a/src/mlglobal/gfs_varinfo.yaml b/src/mlglobal/gfs_varinfo.yaml
new file mode 100644
index 0000000..fba7c7f
--- /dev/null
+++ b/src/mlglobal/gfs_varinfo.yaml
@@ -0,0 +1,21 @@
+pgrb2.0p25.f000:
+  - variables: ["HGT"]
+    levels: ["surface"]
+    first_time_step_only: True
+  - variables: ["TMP"]
+    levels: ["2 m above ground"]
+  - variables: ["PRMSL"]
+    levels: ["mean sea level"]
+  - variables: ["UGRD", "VGRD"]
+    levels: ["10 m above ground"]
+  - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
+    levels: ["50 mb", "100 mb", "150 mb", "200 mb", "250 mb", "300 mb", "400 mb", "500 mb", "600 mb", "700 mb", "850 mb", "925 mb", "1000 mb"]
+pgrb2.0p25.f006:
+  - variables: ["LAND"]
+    levels: ["surface"]
+    first_time_step_only: True
+  - variables: ["APCP"]
+    levels: ["surface"]
+pgrb2b.0p25.f000:
+  - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
+    levels: ["125 mb", "175 mb", "225 mb", "775 mb", "825 mb", "875 mb"]
\ No newline at end of file

From 5480e5d710c2c924dbfb8a2177934ca0fb1c368b Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Tue, 23 Sep 2025 09:08:01 -0400
Subject: [PATCH 04/21] ignore netcdf and grib2 files from being included in
 the repo

---
 .gitignore | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/.gitignore b/.gitignore
index 5e46d37..eadc0e2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -99,3 +99,9 @@ dmypy.json
 
 # VSCode files
 *vscode*
+
+# netCDF and grib files
+*.nc
+*.grib2
+*.grib2.idx
+*pgrb2*

From 6f0b2bc255ca6ab6e626c39d6579d062517a9350 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Wed, 24 Sep 2025 14:17:52 -0400
Subject: [PATCH 05/21] wip

---
 src/mlglobal/cli/gefs_download_main.py   |   0
 src/mlglobal/cli/gefs_preprocess_main.py |   0
 src/mlglobal/cli/gen_gefs_ics_main.py    |   0
 src/mlglobal/cli/graphcast_main.py       |   0
 src/mlglobal/cli/ic_download_main.py     |   9 +-
 src/mlglobal/cli/ic_preprocess_main.py   |   0
 src/mlglobal/cli/nc2grib2_main.py        |   0
 src/mlglobal/gen_ic.py                   | 149 ++++++++++++++---------
 src/mlglobal/ic_downloader.py            |  92 ++++++--------
 9 files changed, 131 insertions(+), 119 deletions(-)
 create mode 100644 src/mlglobal/cli/gefs_download_main.py
 create mode 100644 src/mlglobal/cli/gefs_preprocess_main.py
 create mode 100644 src/mlglobal/cli/gen_gefs_ics_main.py
 create mode 100644 src/mlglobal/cli/graphcast_main.py
 create mode 100644 src/mlglobal/cli/ic_preprocess_main.py
 create mode 100644 src/mlglobal/cli/nc2grib2_main.py

diff --git a/src/mlglobal/cli/gefs_download_main.py b/src/mlglobal/cli/gefs_download_main.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mlglobal/cli/gefs_preprocess_main.py b/src/mlglobal/cli/gefs_preprocess_main.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mlglobal/cli/gen_gefs_ics_main.py b/src/mlglobal/cli/gen_gefs_ics_main.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mlglobal/cli/graphcast_main.py b/src/mlglobal/cli/graphcast_main.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mlglobal/cli/ic_download_main.py b/src/mlglobal/cli/ic_download_main.py
index b5684da..0ec20d8 100644
--- a/src/mlglobal/cli/ic_download_main.py
+++ b/src/mlglobal/cli/ic_download_main.py
@@ -1,5 +1,4 @@
 import argparse
-import logging
 import os
 from datetime import datetime
 
@@ -44,8 +43,8 @@ def _common_args(inparser, dict_in):
             "--source",
             help="Data source",
             type=str,
-            choices=["s3", "local"],
-            default="s3",
+            choices=["local", "s3"],
+            default="local",
             required=False,
         )
         inparser.add_argument(
@@ -66,7 +65,7 @@ def _common_args(inparser, dict_in):
             "--root-directory",
             help="Root directory",
             type=str,
-            default=dict_in["bucket_root_directory"],
+            default=dict_in["comroot"],
             required=False,
         )
         return inparser
@@ -84,7 +83,7 @@ def _common_args(inparser, dict_in):
         help="Ensemble member",
         type=str,
         choices=gefs_members,
-        default=0,
+        default="c00",
     )
 
     args = parser.parse_args()
diff --git a/src/mlglobal/cli/ic_preprocess_main.py b/src/mlglobal/cli/ic_preprocess_main.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mlglobal/cli/nc2grib2_main.py b/src/mlglobal/cli/nc2grib2_main.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/mlglobal/gen_ic.py b/src/mlglobal/gen_ic.py
index 4b75366..ff78d38 100644
--- a/src/mlglobal/gen_ic.py
+++ b/src/mlglobal/gen_ic.py
@@ -7,6 +7,8 @@
 from datetime import datetime
 from pprint import pprint
 import fnmatch
+from mlglobal.ic_downloader import FileLookup
+
 
 logger = logging.getLogger(__name__)
 
@@ -16,16 +18,22 @@ class PrepareIC:
     def __init__(self,
                  current_cycle: datetime,
                  varinfo_file: str,
-                 filelist: list = [],
+                 member: str | None = None,
+                 num_levels: int = 13,
                  data_dir: str = "./data",
                  output_netcdf: str = "ic.nc") -> None:
 
         self.current_cycle = current_cycle
         self.varinfo = self.get_var_info(varinfo_file)
-        self.input_file_list = filelist
+        self.member = member
+        self.num_levels = num_levels
         self.data_dir = data_dir
         self.output_netcdf = output_netcdf
 
+        # Generate the lookup dictionary
+        lookup = FileLookup(self.current_cycle, member=self.member, num_levels=self.num_levels)
+        self.file_dict = lookup.get_file_info()
+
     def get_var_info(self, yaml_file) -> dict:
 
         with open(yaml_file, "r") as fh:
@@ -38,76 +46,72 @@ def process_files(self):
         file_patterns = list(self.varinfo.keys())
 
         mergeDSs = []
-        mergeDAs = []
+        for cycle in self.file_dict.keys():
 
-        # Files are all valid at the same time, so we can merge them along the time dimension
-        for file in self.input_file_list:
-            matched = False
-            for pattern in file_patterns:
-                if fnmatch.fnmatch(file, os.path.join(self.data_dir, '*' + pattern)):
-                    logger.debug(f"Matched pattern: {pattern} in file: {file}")
-                    matched = True
-                    break
+            mergeDAs = []
 
-            if not matched:
-                logger.warning(f"No pattern matched for file: {file}")  # TODO: should this raise an error?
+            # Files valid at the same time, so we can merge them along the time dimension
+            for file in self.file_dict[cycle]:
 
-            logger.info(f"Processing file: {os.path.basename(file)}")
+                filename = os.path.basename(file)
+                file = os.path.join(self.data_dir, filename)
 
-            gribfh = grib2io.open(file)
+                matched = False
+                for pattern in file_patterns:
+                    if fnmatch.fnmatch(file, os.path.join(self.data_dir, '*' + pattern)):
+                        logger.debug(f"Matched pattern: {pattern} in file: {filename}")
+                        matched = True
+                        break
 
-            #for var_pattern, details in self.varinfo[pattern].items():
-            for var_dict in self.varinfo[pattern]:
+                if not matched:
+                    logger.error(f"No pattern matched for file: {filename}")  # TODO: should this raise an error?
+                    raise FileNotFoundError(f"No pattern matched for file: {filename}")
 
-                variable_names = var_dict["variables"]
-                levels = var_dict["levels"]
+                logger.info(f"Processing {filename=}")
 
-                for var_name in variable_names:
-                    logger.info(f"Extracting variable: {var_name} at levels: {levels}")
-                    if len(levels) > 1:
-                        da = self.get_dataarray_3d(gribfh, var_name, levels)
-                    else:
-                        da = self.get_dataarray_2d(gribfh, var_name, levels[0])
-                    mergeDAs.append(da)
+                gribfh = grib2io.open(file, "r")
 
-            gribfh.close()
+                for var_dict in self.varinfo[pattern]:
 
-        ds = xr.merge(mergeDAs, compat="no_conflicts")
-        mergeDSs.append(ds)
-        ds.close()
+                    variable_names = var_dict["variables"]
+                    levels = var_dict["levels"]
+
+                    for var_name in variable_names:
+                        logger.info(f"Extracting variable: {var_name} at levels: {levels}")
+                        if len(levels) > 1:
+                            da = self.get_dataarray_3d(gribfh, var_name, levels)
+                        else:
+                            da = self.get_dataarray_2d(gribfh, var_name, levels[0])
+                        mergeDAs.append(da)
+
+                gribfh.close()
+
+            ds = xr.merge(mergeDAs, compat="no_conflicts")
+            mergeDSs.append(ds)
+            ds.close()
 
+        # Concatenate along the time dimension
         ds = xr.concat(mergeDSs, dim="time")
 
         # Get 2D static data from the f000 file
         # From the file list find the first file that ends with .f000
-        f000_file = next((f for f in self.input_file_list if f.endswith(".f000")), None)
+        f000_file = next((f for f in self.file_dict[self.current_cycle] if f.endswith(".f000")), None)
         if f000_file is None:
-            logger.warning("No 00Z f000 file found.")  # TODO: should this raise an error?
-            return
+            logger.error("No 00Z f000 file found.")  # TODO: should this raise an error?
+            raise FileNotFoundError("No f000 file found.")
 
-        gribfh = grib2io.open(os.path.join(self.data_dir, f000_file))
+        f000_file = os.path.join(self.data_dir, os.path.basename(f000_file))
+        gribfh = grib2io.open(f000_file, "r")
         static_vars = ["LAND", "HGT"]
         for var_name in static_vars:
             da = self.get_dataarray_2d(gribfh, var_name, "surface")
             ds = xr.merge([ds, da], compat="no_conflicts")
         gribfh.close()
 
-        ds = ds.rename({
-            "LAND_surface": "land_sea_mask",
-            "HGT_surface": "geopotential_at_surface",
-            "PRMSL_meansealevel": "mean_sea_level_pressure",
-            "TMP_2maboveground": "2m_temperature",
-            "UGRD_10maboveground": "10m_u_component_of_wind",
-            "VGRD_10maboveground": "10m_v_component_of_wind",
-            "APCP_surface": "total_precipitation_6hr",
-            "HGT": "geopotential",
-            "TMP": "temperature",
-            "SPFH": "specific_humidity",
-            "VVEL": "vertical_velocity",
-            "UGRD": "u_component_of_wind",
-            "VGRD": "v_component_of_wind",
-        })
+        # Rename variables to match graphcast naming conventions
+        ds = self.rename_dsvars(ds)
 
+        # Add datetime coordinate
         ds = ds.assign_coords(datetime=ds.time)
 
         # Adjust time values relative to the first time step
@@ -125,14 +129,48 @@ def process_files(self):
         ds["geopotential_at_surface"] = ds["geopotential_at_surface"] * 9.80665
         ds["geopotential"] = ds["geopotential"] * 9.80665
 
-        # Update total_precipitation_6hr unit to (m) from (kg/m^2) by dividing it by 1000kg/m³
-        ds["total_precipitation_6hr"] = ds["total_precipitation_6hr"] / 1000.0
+        # For GEFS, if total_precipitation_6hr is missing, create it with zeros
+        if "total_precipitation_6hr" not in ds:
+            logger.warning("total_precipitation_6hr variable not found. Creating it with zeros.")
+            ds["total_precipitation_6hr"] = xr.zeros_like(ds["2m_temperature"])
+        else:
+            # Update total_precipitation_6hr unit to (m) from (kg/m^2) by dividing it by 1000kg/m³
+            ds["total_precipitation_6hr"] = ds["total_precipitation_6hr"] / 1000.0
 
         ds.to_netcdf(self.output_netcdf)
         ds.close()
 
         return
 
+    @staticmethod
+    def rename_dsvars(ds: xr.Dataset, rename_dict: dict = None) -> xr.Dataset:
+        """Rename dataset variables to match graphcast naming conventions."""
+        logger.info("Renaming dataset variables to match graphcast naming conventions.")
+
+        # Default rename dictionary  #TODO: move to yaml?
+        if rename_dict is None:
+            rename_dict = {
+                "LAND_surface": "land_sea_mask",
+                "HGT_surface": "geopotential_at_surface",
+                "PRMSL_meansealevel": "mean_sea_level_pressure",
+                "TMP_2maboveground": "2m_temperature",
+                "UGRD_10maboveground": "10m_u_component_of_wind",
+                "VGRD_10maboveground": "10m_v_component_of_wind",
+                "HGT": "geopotential",
+                "TMP": "temperature",
+                "SPFH": "specific_humidity",
+                "VVEL": "vertical_velocity",
+                "UGRD": "u_component_of_wind",
+                "VGRD": "v_component_of_wind",
+            }
+            if "APCP_surface" in ds:
+                rename_dict["APCP_surface"] = "total_precipitation_6hr"
+            logger.debug(f"Using rename_dict: {rename_dict}")
+
+        ds = ds.rename(rename_dict)
+
+        return ds
+
     @staticmethod
     def get_dataarray_2d(grbfile, var_name, desired_level):
 
@@ -169,7 +207,7 @@ def get_dataarray_2d(grbfile, var_name, desired_level):
                 }
             )
         elif len(data.shape) == 3:
-            print("3D data found in 2D function")
+            logger.warning("3D data found in 2D function")
             da = xr.Dataset(
                 data_vars={
                     var_name: (["level", "lat", "lon"], data.astype("float32"))
@@ -266,14 +304,11 @@ def get_dataarray_3d(grbfile, var_name, desired_level):
 
     args = parser.parse_args()
 
-    input_file_list = [os.path.join(args.data_dir, "gfs.t00z.pgrb2.0p25.f000"),
-                       os.path.join(args.data_dir, "gfs.t18z.pgrb2.0p25.f006")]
-    #input_file_list = glob.glob(os.path.join(args.data_dir, "*"))
-
     prep = PrepareIC(
         current_cycle=datetime.strptime(args.current_cycle, "%Y%m%d%H"),
         varinfo_file=args.yaml,
-        filelist=input_file_list,
+        member=None,
+        num_levels=13,
         output_netcdf=args.output,
         data_dir=args.data_dir
     )
diff --git a/src/mlglobal/ic_downloader.py b/src/mlglobal/ic_downloader.py
index 0ed0be8..44dad12 100644
--- a/src/mlglobal/ic_downloader.py
+++ b/src/mlglobal/ic_downloader.py
@@ -2,83 +2,63 @@
 import shutil
 from datetime import timedelta
 from logging import getLogger
+from pprint import pprint
 
 
 logger = getLogger(__name__)
 
 
 class FileLookup:
-    def __init__(self, current_cycle, member=None):
+    def __init__(self, current_cycle, num_levels=13, member=None):
 
         self.current_cycle = current_cycle
+        self.num_levels = num_levels
         self.member = member  # GEFS member values are c00, p01, p02, ..., p30
 
-        # Look back 6 and 12 hours for precip files
+        # Look back 6 hours for precip and 2 time-level data
         self.current_cycle_m6h = self.current_cycle - timedelta(hours=6)
-        self.current_cycle_m12h = self.current_cycle - timedelta(hours=12)
 
         if self.member is not None:
-            self.template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{member}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
             self.get_file_info = self._gefs_file_info
         else:
-            self.template = f"gfs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/gfs.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
             self.get_file_info = self._gfs_file_info
 
     def _gfs_file_info(self):
 
-        file_formats =[
-            "pgrb2.0p25.f000",
-            "pgrb2b.0p25.f000",
-            "pgrb2.0p25.f006"
-        ]
+        # Template for GFS files
+        template = f"gfs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/gfs.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
 
         # From current cycle
-        pgrb2_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec="pgrb2.0p25", fhour=0)
-        pgrb2b_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec="pgrb2b.0p25", fhour=0)
+        pgrb2_0p25_f000 = template.format(cycle=self.current_cycle, fspec="pgrb2.0p25", fhour=0)
+        if self.num_levels == 37:  # Need additional pgrb2b file for 37 level data
+            pgrb2b_0p25_f000 = template.format(cycle=self.current_cycle, fspec="pgrb2b.0p25", fhour=0)
 
         # From current cycle - 6 hours
-        pgrb2_0p25_f000_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=0)
-        pgrb2_0p25_f006_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=6)
-
-        # From current cycle - 12 hours
-        pgrb2_0p25_f006_m12 = self.template.format(cycle=self.current_cycle_m12h, fspec="pgrb2.0p25", fhour=6)
+        pgrb2_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=0)
+        pgrb2_0p25_f006_m6 = template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=6)
 
+        # Create a flat list of files valid at current and current-6h cycles
         file_dict = {}
-        file_dict[self.current_cycle] = {"pgrb2.0p25.f000": pgrb2_0p25_f000,
-                                         "pgrb2b.0p25.f000": pgrb2b_0p25_f000}
-        file_dict[self.current_cycle_m6h] = {"pgrb2.0p25.f000": pgrb2_0p25_f000_m6,
-                                             "pgrb2.0p25.f006": pgrb2_0p25_f006_m6}
-        file_dict[self.current_cycle_m12h] = {"pgrb2.0p25.f006": pgrb2_0p25_f006_m12}
+        file_dict[self.current_cycle] = [pgrb2_0p25_f000, pgrb2_0p25_f006_m6]
+        if self.num_levels == 37:
+            file_dict[self.current_cycle].append(pgrb2b_0p25_f000)
+        file_dict[self.current_cycle_m6h] = [pgrb2_0p25_f000_m6]
 
-        file_list = [
-            pgrb2_0p25_f000,
-            pgrb2b_0p25_f000,
-            pgrb2_0p25_f000_m6,
-            pgrb2_0p25_f006_m6,
-            pgrb2_0p25_f006_m12
-        ]
-
-        return file_dict, file_list, file_formats
+        return file_dict
 
     def _gefs_file_info(self):
 
-        file_formats = [
-            "pgrb2.0p25.f000",
-            "pgrb2s.0p25.f000",
-            "pgrb2s.0p25.f006"
-        ]
+        # Template for GEFS files
+        template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{self.member}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
 
         # From current cycle
-        pgrb2_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
-        pgrb2s_0p25_f000 = self.template.format(cycle=self.current_cycle, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
+        pgrb2_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
+        pgrb2s_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
 
         # From current cycle - 6 hours
-        pgrb2_0p25_f000_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
-        pgrb2s_0p25_f000_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
-        pgrb2s_0p25_f006_m6 = self.template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=6)
-
-        # From current cycle - 12 hours
-        pgrb2s_0p25_f006_m12 = self.template.format(cycle=self.current_cycle_m12h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=6)
+        pgrb2_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
+        pgrb2s_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
+        pgrb2s_0p25_f006_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=6)
 
         file_dict = {}
         file_dict[self.current_cycle] = {"pgrb2.0p25.f000": pgrb2_0p25_f000,
@@ -86,18 +66,8 @@ def _gefs_file_info(self):
         file_dict[self.current_cycle_m6h] = {"pgrb2.0p25.f000": pgrb2_0p25_f000_m6,
                                              "pgrb2s.0p25.f000": pgrb2s_0p25_f000_m6,
                                              "pgrb2s.0p25.f006": pgrb2s_0p25_f006_m6}
-        file_dict[self.current_cycle_m12h] = {"pgrb2s.0p25.f006": pgrb2s_0p25_f006_m12}
-
-        file_list = [
-            pgrb2_0p25_f000,
-            pgrb2s_0p25_f000,
-            pgrb2_0p25_f000_m6,
-            pgrb2s_0p25_f000_m6,
-            pgrb2s_0p25_f006_m6,
-            pgrb2s_0p25_f006_m12
-        ]
 
-        return file_dict, file_list, file_formats
+        return file_dict
 
 
 class ICDownloader:
@@ -105,6 +75,7 @@ class ICDownloader:
     def __init__(
         self,
         current_cycle,
+        num_levels=13,
         member=None,
         download_source="local",
         local_directory="./data",
@@ -112,6 +83,7 @@ def __init__(
         root_directory=None
     ):
         self.current_cycle = current_cycle
+        self.num_levels = num_levels
         self.member = member
         self.download_source = download_source
         self.local_directory = local_directory
@@ -119,8 +91,14 @@ def __init__(
         self.root_directory = root_directory
 
         # Generate the lookup dictionary
-        lookup = FileLookup(self.current_cycle, member=self.member)
-        self.file_dict, self.file_list, self.file_formats = lookup.get_file_info()
+        lookup = FileLookup(self.current_cycle, member=self.member, num_levels=self.num_levels)
+        self.file_dict = lookup.get_file_info()
+
+        # Flatten the file_dict to a simple list of files
+        file_list = []
+        for files in self.file_dict.values():
+            file_list.extend(files)
+        self.file_list = file_list
 
         if self.download_source in ["s3"]:
             aws_profile = os.environ.get("AWS_PROFILE", "default")

From cd8c4a460f07151030e0e8622eb75ee49a335942 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Wed, 24 Sep 2025 14:19:55 -0400
Subject: [PATCH 06/21] remove undesired files

---
 src/mlglobal/cli/gefs_download_main.py   | 0
 src/mlglobal/cli/gefs_preprocess_main.py | 0
 src/mlglobal/cli/gen_gefs_ics_main.py    | 0
 src/mlglobal/cli/graphcast_main.py       | 0
 src/mlglobal/cli/ic_preprocess_main.py   | 0
 src/mlglobal/cli/nc2grib2_main.py        | 0
 6 files changed, 0 insertions(+), 0 deletions(-)
 delete mode 100644 src/mlglobal/cli/gefs_download_main.py
 delete mode 100644 src/mlglobal/cli/gefs_preprocess_main.py
 delete mode 100644 src/mlglobal/cli/gen_gefs_ics_main.py
 delete mode 100644 src/mlglobal/cli/graphcast_main.py
 delete mode 100644 src/mlglobal/cli/ic_preprocess_main.py
 delete mode 100644 src/mlglobal/cli/nc2grib2_main.py

diff --git a/src/mlglobal/cli/gefs_download_main.py b/src/mlglobal/cli/gefs_download_main.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/mlglobal/cli/gefs_preprocess_main.py b/src/mlglobal/cli/gefs_preprocess_main.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/mlglobal/cli/gen_gefs_ics_main.py b/src/mlglobal/cli/gen_gefs_ics_main.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/mlglobal/cli/graphcast_main.py b/src/mlglobal/cli/graphcast_main.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/mlglobal/cli/ic_preprocess_main.py b/src/mlglobal/cli/ic_preprocess_main.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/mlglobal/cli/nc2grib2_main.py b/src/mlglobal/cli/nc2grib2_main.py
deleted file mode 100644
index e69de29..0000000

From ae8f3273aeaf0578d6ead4a4cbf857b781346323 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Wed, 24 Sep 2025 21:04:49 -0400
Subject: [PATCH 07/21] wip

---
 pyproject.toml                 |  3 ++-
 src/mlglobal/gefs_varinfo.yaml | 21 ++++++++++++++++
 src/mlglobal/gen_ic.py         | 31 +++++++++++++++--------
 src/mlglobal/gfs_varinfo.yaml  | 46 ++++++++++++++++++----------------
 src/mlglobal/ic_downloader.py  | 18 ++++++-------
 5 files changed, 75 insertions(+), 44 deletions(-)
 create mode 100644 src/mlglobal/gefs_varinfo.yaml

diff --git a/pyproject.toml b/pyproject.toml
index 316db50..8cdf584 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,7 +40,8 @@ dependencies = [
     "scipy>=1.9.0",
     "trimesh>=3.15.0",
     "typing-extensions>=4.4.0",
-    "graphcast@git+https://github.com/google-deepmind/graphcast@6819a0f19796ca9c34a079855339b37134b8d930"
+    "grib2io>=2.5.4",
+    "graphcast@git+https://github.com/noaa-emc/graphcast@aea0678cfbd7e866e4a1d364b456861d0a03954b"
 ]
 
 [project.optional-dependencies]
diff --git a/src/mlglobal/gefs_varinfo.yaml b/src/mlglobal/gefs_varinfo.yaml
new file mode 100644
index 0000000..6cfe60c
--- /dev/null
+++ b/src/mlglobal/gefs_varinfo.yaml
@@ -0,0 +1,21 @@
+time_variant:
+  pgrb2s.0p25.f000:
+    - variables: ["TMP"]
+      levels: ["2 m above ground"]
+    - variables: ["PRMSL"]
+      levels: ["mean sea level"]
+    - variables: ["UGRD", "VGRD"]
+      levels: ["10 m above ground"]
+  pgrb2.0p25.f000:
+    - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
+      levels: ["50 mb", "100 mb", "150 mb", "200 mb", "250 mb", "300 mb", "400 mb", "500 mb", "600 mb", "700 mb", "850 mb", "925 mb", "1000 mb"]
+time_invariant:
+  pgrb2.0p25.f000:
+    - variables: ["LAND"]
+      levels: ["surface"]
+      first_time_step_only: True
+  pgrb2s.0p25.f000:
+    - variables: ["HGT"]
+      levels: ["surface"]
+      first_time_step_only: True
+
diff --git a/src/mlglobal/gen_ic.py b/src/mlglobal/gen_ic.py
index ff78d38..1c03d87 100644
--- a/src/mlglobal/gen_ic.py
+++ b/src/mlglobal/gen_ic.py
@@ -41,9 +41,27 @@ def get_var_info(self, yaml_file) -> dict:
 
         return variables_to_extract
 
+
+    def get_matching_pattern(self, file, patterns, fatal=True):
+
+        matched = False
+        for pattern in file_patterns:
+            if fnmatch.fnmatch(file, os.path.join(self.data_dir, '*' + pattern)):
+                logger.debug(f"Matched pattern: {pattern} in file: {filename}")
+                matched = True
+                break
+
+        if not matched:
+            logger.error(f"No pattern matched for file: {filename}")
+            pattern = None
+            if fatal:
+                raise FileNotFoundError(f"No pattern matched for file: {filename}")
+
+        return pattern
+
     def process_files(self):
 
-        file_patterns = list(self.varinfo.keys())
+        file_patterns = list(self.varinfo['time_variant'].keys())
 
         mergeDSs = []
         for cycle in self.file_dict.keys():
@@ -56,16 +74,7 @@ def process_files(self):
                 filename = os.path.basename(file)
                 file = os.path.join(self.data_dir, filename)
 
-                matched = False
-                for pattern in file_patterns:
-                    if fnmatch.fnmatch(file, os.path.join(self.data_dir, '*' + pattern)):
-                        logger.debug(f"Matched pattern: {pattern} in file: {filename}")
-                        matched = True
-                        break
-
-                if not matched:
-                    logger.error(f"No pattern matched for file: {filename}")  # TODO: should this raise an error?
-                    raise FileNotFoundError(f"No pattern matched for file: {filename}")
+                pattern = get_pattern(file, file_patterns)
 
                 logger.info(f"Processing {filename=}")
 
diff --git a/src/mlglobal/gfs_varinfo.yaml b/src/mlglobal/gfs_varinfo.yaml
index fba7c7f..47753be 100644
--- a/src/mlglobal/gfs_varinfo.yaml
+++ b/src/mlglobal/gfs_varinfo.yaml
@@ -1,21 +1,25 @@
-pgrb2.0p25.f000:
-  - variables: ["HGT"]
-    levels: ["surface"]
-    first_time_step_only: True
-  - variables: ["TMP"]
-    levels: ["2 m above ground"]
-  - variables: ["PRMSL"]
-    levels: ["mean sea level"]
-  - variables: ["UGRD", "VGRD"]
-    levels: ["10 m above ground"]
-  - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
-    levels: ["50 mb", "100 mb", "150 mb", "200 mb", "250 mb", "300 mb", "400 mb", "500 mb", "600 mb", "700 mb", "850 mb", "925 mb", "1000 mb"]
-pgrb2.0p25.f006:
-  - variables: ["LAND"]
-    levels: ["surface"]
-    first_time_step_only: True
-  - variables: ["APCP"]
-    levels: ["surface"]
-pgrb2b.0p25.f000:
-  - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
-    levels: ["125 mb", "175 mb", "225 mb", "775 mb", "825 mb", "875 mb"]
\ No newline at end of file
+time_variant:
+  pgrb2.0p25.f000:
+    - variables: ["TMP"]
+      levels: ["2 m above ground"]
+    - variables: ["PRMSL"]
+      levels: ["mean sea level"]
+    - variables: ["UGRD", "VGRD"]
+      levels: ["10 m above ground"]
+    - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
+      levels: ["50 mb", "100 mb", "150 mb", "200 mb", "250 mb", "300 mb", "400 mb", "500 mb", "600 mb", "700 mb", "850 mb", "925 mb", "1000 mb"]
+  pgrb2.0p25.f006:
+    - variables: ["APCP"]
+      levels: ["surface"]
+  pgrb2b.0p25.f000:
+    - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"]
+      levels: ["125 mb", "175 mb", "225 mb", "775 mb", "825 mb", "875 mb"]
+time_invariant:
+  pgrb2.0p25.f000:
+    - variables: ["HGT"]
+      levels: ["surface"]
+      first_time_step_only: True
+    - variables: ["LAND"]
+      levels: ["surface"]
+      first_time_step_only: True
+
diff --git a/src/mlglobal/ic_downloader.py b/src/mlglobal/ic_downloader.py
index 44dad12..26f249d 100644
--- a/src/mlglobal/ic_downloader.py
+++ b/src/mlglobal/ic_downloader.py
@@ -49,23 +49,19 @@ def _gfs_file_info(self):
     def _gefs_file_info(self):
 
         # Template for GEFS files
-        template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{self.member}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
+        template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{{member}}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
 
         # From current cycle
-        pgrb2_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
-        pgrb2s_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
+        pgrb2_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0, member=self.member)
+        pgrb2s_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0, member=self.member)
 
         # From current cycle - 6 hours
-        pgrb2_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0)
-        pgrb2s_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0)
-        pgrb2s_0p25_f006_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=6)
+        pgrb2_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0, member=self.member)
+        pgrb2s_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0, member=self.member)
 
         file_dict = {}
-        file_dict[self.current_cycle] = {"pgrb2.0p25.f000": pgrb2_0p25_f000,
-                                         "pgrb2s.0p25.f000": pgrb2s_0p25_f000}
-        file_dict[self.current_cycle_m6h] = {"pgrb2.0p25.f000": pgrb2_0p25_f000_m6,
-                                             "pgrb2s.0p25.f000": pgrb2s_0p25_f000_m6,
-                                             "pgrb2s.0p25.f006": pgrb2s_0p25_f006_m6}
+        file_dict[self.current_cycle] = [pgrb2_0p25_f000, pgrb2s_0p25_f000]
+        file_dict[self.current_cycle_m6h] = [pgrb2_0p25_f000_m6, pgrb2s_0p25_f000_m6]
 
         return file_dict
 

From 361b2170a211c37446870bc70229196037f6c559 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 00:00:59 -0400
Subject: [PATCH 08/21] tested with s3

---
 .gitmodules                                |   1 +
 {src/mlglobal => config}/gefs_varinfo.yaml |   1 -
 {src/mlglobal => config}/gfs_varinfo.yaml  |   1 -
 pyproject.toml                             |   6 +-
 src/mlglobal/cli/gen_ics_main.py           | 194 ++++++++++++
 src/mlglobal/cli/ic_download_main.py       | 103 ------
 src/mlglobal/gen_ic.py                     | 326 -------------------
 src/mlglobal/ic_downloader.py              | 136 ++++++--
 src/mlglobal/ic_processor.py               | 351 +++++++++++++++++++++
 src/mlglobal/logger.py                     |  16 +-
 10 files changed, 656 insertions(+), 479 deletions(-)
 rename {src/mlglobal => config}/gefs_varinfo.yaml (99%)
 rename {src/mlglobal => config}/gfs_varinfo.yaml (99%)
 create mode 100644 src/mlglobal/cli/gen_ics_main.py
 delete mode 100644 src/mlglobal/cli/ic_download_main.py
 delete mode 100644 src/mlglobal/gen_ic.py
 create mode 100644 src/mlglobal/ic_processor.py

diff --git a/.gitmodules b/.gitmodules
index fd7d84c..90582fd 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,3 +1,4 @@
 [submodule "graphcast"]
 	path = graphcast
 	url = https://github.com/google-deepmind/graphcast.git
+  ignore = dirty
diff --git a/src/mlglobal/gefs_varinfo.yaml b/config/gefs_varinfo.yaml
similarity index 99%
rename from src/mlglobal/gefs_varinfo.yaml
rename to config/gefs_varinfo.yaml
index 6cfe60c..954d9c4 100644
--- a/src/mlglobal/gefs_varinfo.yaml
+++ b/config/gefs_varinfo.yaml
@@ -18,4 +18,3 @@ time_invariant:
     - variables: ["HGT"]
       levels: ["surface"]
       first_time_step_only: True
-
diff --git a/src/mlglobal/gfs_varinfo.yaml b/config/gfs_varinfo.yaml
similarity index 99%
rename from src/mlglobal/gfs_varinfo.yaml
rename to config/gfs_varinfo.yaml
index 47753be..4b5c23a 100644
--- a/src/mlglobal/gfs_varinfo.yaml
+++ b/config/gfs_varinfo.yaml
@@ -22,4 +22,3 @@ time_invariant:
     - variables: ["LAND"]
       levels: ["surface"]
       first_time_step_only: True
-
diff --git a/pyproject.toml b/pyproject.toml
index 8cdf584..801f0bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -71,13 +71,9 @@ notebooks = [
 requires = ["setuptools>=64.0.0", "setuptools-scm"]
 build-backend = "setuptools.build_meta"
 
-
 [tool.setuptools]
 package-dir = {"" = "src"}
 packages = ["mlglobal"]
 
 [project.scripts]
-ic-download = "mlglobal.cli.ic_download_main:main"
-#ic-preprocess = "mlglobal.cli.gefs_preprocess_main:main"
-#graphcast-run = "mlglobal.cli.graphcast_main:main"
-#nc2grib2 = "mlglobal.cli.nc2grib2_main:main"
+gen_ics = "mlglobal.cli.gen_ics_main:main"
diff --git a/src/mlglobal/cli/gen_ics_main.py b/src/mlglobal/cli/gen_ics_main.py
new file mode 100644
index 0000000..ba73c3a
--- /dev/null
+++ b/src/mlglobal/cli/gen_ics_main.py
@@ -0,0 +1,194 @@
+import argparse
+import logging
+import os
+from datetime import datetime
+
+from mlglobal.ic_downloader import ICDownloader
+from mlglobal.ic_processor import ICProcessor
+from mlglobal.logger import setup_logging
+
+logger = logging.getLogger(__name__)
+
+_here = os.path.abspath(os.path.dirname(__file__))
+_top = os.path.abspath(os.path.join(os.path.abspath(_here), "../../../"))
+
+# Default bucket and root directory for each model  # TODO: store this in a yaml or some config file
+DEFAULTS = {
+    "gfs": {
+        "bucket_name": "noaa-gfs-bdp-pds",
+        "bucket_root_directory": "",
+        "comroot": "/lfs/h1/ops/prod/com/gfs/v16.3",
+    },
+    "gefs": {
+        "bucket_name": "noaa-ncepdev-none-ca-ufs-cpldcld",
+        "bucket_root_directory": "Linlin.Cui/gefs_wcoss2",
+        "comroot": "/lfs/h1/ops/prod/com/gefs/v12.3",
+    },
+}
+
+# --- Helper strings for dynamic help text ---
+BUCKET_HELP = (
+    f"S3 bucket name. [default: {DEFAULTS['gfs']['bucket_name']} (for GFS), "
+    f"{DEFAULTS['gefs']['bucket_name']} (for GEFS)]"
+)
+
+BUCKET_ROOT_DIR_HELP = (
+    f"S3 bucket root directory. [default: {DEFAULTS['gfs']['bucket_root_directory']} (for GFS), "
+    f"{DEFAULTS['gefs']['bucket_root_directory']} (for GEFS)]"
+)
+
+COMROOT_DIR_HELP = (
+    f"Root directory. [default: {DEFAULTS['gfs']['comroot']} (for GFS), "
+    f"{DEFAULTS['gefs']['comroot']} (for GEFS)]"
+)
+
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        description="Download IC data for GFS or GEFS",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    subparsers = parser.add_subparsers(
+        dest="model",
+        help="Model to download and process initial conditions for [GFS | GEFS]",
+        required=True,
+    )
+
+    def _common_args(inparser, model: str):
+        dict_in = DEFAULTS[model]
+        inparser.add_argument(
+            "--current-cycle",
+            help="Datetime to download and process initial conditions for in YYYYMMDDHH format",
+            type=str,
+            metavar="YYYYMMDDHH",
+            required=True,
+        )
+        inparser.add_argument(
+            "--source",
+            help="Data source for getting model grib2 data",
+            type=str,
+            choices=["local", "s3"],
+            default="local",
+            required=False,
+        )
+        inparser.add_argument(
+            "--target",
+            help="Target directory to store grib2 model data into",
+            type=str,
+            default="./input_data",
+            required=False,
+        )
+        inparser.add_argument(
+            "--bucket-name",
+            help=BUCKET_HELP,
+            type=str,
+            default=dict_in["bucket_name"],
+            required=False,
+        )
+        inparser.add_argument(
+            "--bucket-root-directory",
+            help=BUCKET_ROOT_DIR_HELP,
+            type=str,
+            default=dict_in["bucket_root_directory"],
+            required=False,
+        )
+        inparser.add_argument(
+            "--comroot",
+            help=COMROOT_DIR_HELP,
+            type=str,
+            default=dict_in["comroot"],
+            required=False,
+        )
+        inparser.add_argument(
+            "--num-levels",
+            help="Number of vertical levels to download from the model data",
+            type=int,
+            default=13,
+            required=False,
+        )
+        inparser.add_argument(
+            "--varinfo-yaml",
+            help="Path to the varinfo YAML file",
+            type=str,
+            default=os.path.join(_top, "config", f"{model}_varinfo.yaml"),
+            required=False,
+        )
+        inparser.add_argument(
+            "--output",
+            help="Name of the output NetCDF file",
+            type=str,
+            default="graphcast_ic.nc",
+            required=False,
+        )
+        inparser.add_argument(
+            "--debug",
+            help="Set logging level to DEBUG",
+            action="store_true",
+            required=False,
+        )
+        inparser.add_argument(
+            "--download-only",
+            help="Only download the data, do not process",
+            action="store_true",
+            required=False,
+        )
+
+        return inparser
+
+    # GFS subparser
+    gfs_parser = subparsers.add_parser("gfs", help="Download GFS data")
+    gfs_parser = _common_args(gfs_parser, "gfs")
+
+    # GEFS subparser
+    gefs_parser = subparsers.add_parser("gefs", help="Download GEFS ensemble data")
+    gefs_parser = _common_args(gefs_parser, "gefs")
+    gefs_members = ["c00"] + [f"p{str(i).zfill(2)}" for i in range(1, 31)]
+    gefs_parser.add_argument(
+        "--member",
+        help="Ensemble member",
+        type=str,
+        choices=gefs_members,
+        default="c00",
+    )
+
+    args = parser.parse_args()
+
+    setup_logging(debug=args.debug)
+
+    current_cycle = datetime.strptime(args.current_cycle, "%Y%m%d%H")
+
+    logger.info(f"Downloading {args.model.upper()} data for cycle {args.current_cycle}")
+    downloader = ICDownloader(
+        current_cycle=current_cycle,
+        num_levels=args.num_levels,
+        member=None if args.model == "gfs" else args.member,
+        download_source=args.source,
+        local_directory=args.target,
+        bucket_name=args.bucket_name,
+        bucket_root_directory=args.bucket_root_directory,
+        root_directory=args.comroot,
+    )
+    downloader.get_data()
+
+    # TODO: validate the downloaded files before processing
+
+    if args.download_only:
+        logger.info("Download-only flag set, skipping processing step.")
+        return
+
+    logger.info(f"Processing {args.model.upper()} data for cycle {args.current_cycle}")
+    processor = ICProcessor(
+        current_cycle=current_cycle,
+        num_levels=args.num_levels,
+        member=None if args.model == "gfs" else args.member,
+        varinfo_file=args.varinfo_yaml,
+        data_dir=args.target,
+        output_netcdf=args.output,
+    )
+    processor.process_files()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/mlglobal/cli/ic_download_main.py b/src/mlglobal/cli/ic_download_main.py
deleted file mode 100644
index 0ec20d8..0000000
--- a/src/mlglobal/cli/ic_download_main.py
+++ /dev/null
@@ -1,103 +0,0 @@
-import argparse
-import os
-from datetime import datetime
-
-from mlglobal.ic_downloader import ICDownloader
-from mlglobal.logger import setup_logging
-
-
-# Default bucket and root directory for each mode  # TODO: store this in a yaml or some config file
-DEFAULTS = {
-    "gfs": {
-        "bucket_name": "noaa-gfs-bdp-pds",
-        "bucket_root_directory": "",
-        "comroot": "/lfs/h1/ops/prod/com/gfs/v16.3",
-    },
-    "gefs": {
-        "bucket_name": "noaa-ncepdev-none-ca-ufs-cpldcld",
-        "bucket_root_directory": "Linlin.Cui/gefs_wcoss2",
-        "comroot": "/lfs/h1/ops/prod/com/gefs/v12.3",
-    },
-}
-
-
-def main():
-
-    setup_logging()
-
-    parser = argparse.ArgumentParser(description="Download IC data for GFS or GEFS")
-
-    subparsers = parser.add_subparsers(
-        dest="mode", help="System to download IC data for GFS or GEFS", required=True
-    )
-
-    def _common_args(inparser, dict_in):
-        inparser.add_argument(
-            "--current_cycle",
-            help="Datetime to download data for in YYYYMMDDHH format",
-            type=str,
-            metavar="YYYYMMDDHH",
-            required=True,
-        )
-        inparser.add_argument(
-            "--source",
-            help="Data source",
-            type=str,
-            choices=["local", "s3"],
-            default="local",
-            required=False,
-        )
-        inparser.add_argument(
-            "--target",
-            help="Target directory to store raw data into",
-            type=str,
-            default=os.getcwd(),
-            required=False,
-        )
-        inparser.add_argument(
-            "--bucket-name",
-            help="S3 bucket name",
-            type=str,
-            default=dict_in["bucket_name"],
-            required=False,
-        )
-        inparser.add_argument(
-            "--root-directory",
-            help="Root directory",
-            type=str,
-            default=dict_in["comroot"],
-            required=False,
-        )
-        return inparser
-
-    # GFS subparser
-    gfs_parser = subparsers.add_parser("gfs", help="Download GFS ensemble data")
-    gfs_parser = _common_args(gfs_parser, DEFAULTS["gfs"])
-
-    # GEFS subparser
-    gefs_parser = subparsers.add_parser("gefs", help="Download GEFS ensemble data")
-    gefs_parser = _common_args(gefs_parser, DEFAULTS["gefs"])
-    gefs_members = ["c00"] + [f"p{str(i).zfill(2)}" for i in range(1, 31)]
-    gefs_parser.add_argument(
-        "--member",
-        help="Ensemble member",
-        type=str,
-        choices=gefs_members,
-        default="c00",
-    )
-
-    args = parser.parse_args()
-
-    downloader = ICDownloader(
-        current_cycle=datetime.strptime(args.current_cycle, "%Y%m%d%H"),
-        member=None if args.mode == "gfs" else args.member,
-        download_source=args.source,
-        local_directory=args.target,
-        bucket_name=args.bucket_name,
-        root_directory=args.root_directory,
-    )
-    downloader.get_data()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/mlglobal/gen_ic.py b/src/mlglobal/gen_ic.py
deleted file mode 100644
index 1c03d87..0000000
--- a/src/mlglobal/gen_ic.py
+++ /dev/null
@@ -1,326 +0,0 @@
-import yaml
-import os
-import logging
-import grib2io
-import xarray as xr
-import numpy as np
-from datetime import datetime
-from pprint import pprint
-import fnmatch
-from mlglobal.ic_downloader import FileLookup
-
-
-logger = logging.getLogger(__name__)
-
-class PrepareIC:
-    """Class to prepare initial conditions."""
-
-    def __init__(self,
-                 current_cycle: datetime,
-                 varinfo_file: str,
-                 member: str | None = None,
-                 num_levels: int = 13,
-                 data_dir: str = "./data",
-                 output_netcdf: str = "ic.nc") -> None:
-
-        self.current_cycle = current_cycle
-        self.varinfo = self.get_var_info(varinfo_file)
-        self.member = member
-        self.num_levels = num_levels
-        self.data_dir = data_dir
-        self.output_netcdf = output_netcdf
-
-        # Generate the lookup dictionary
-        lookup = FileLookup(self.current_cycle, member=self.member, num_levels=self.num_levels)
-        self.file_dict = lookup.get_file_info()
-
-    def get_var_info(self, yaml_file) -> dict:
-
-        with open(yaml_file, "r") as fh:
-            variables_to_extract = yaml.safe_load(fh)
-
-        return variables_to_extract
-
-
-    def get_matching_pattern(self, file, patterns, fatal=True):
-
-        matched = False
-        for pattern in file_patterns:
-            if fnmatch.fnmatch(file, os.path.join(self.data_dir, '*' + pattern)):
-                logger.debug(f"Matched pattern: {pattern} in file: {filename}")
-                matched = True
-                break
-
-        if not matched:
-            logger.error(f"No pattern matched for file: {filename}")
-            pattern = None
-            if fatal:
-                raise FileNotFoundError(f"No pattern matched for file: {filename}")
-
-        return pattern
-
-    def process_files(self):
-
-        file_patterns = list(self.varinfo['time_variant'].keys())
-
-        mergeDSs = []
-        for cycle in self.file_dict.keys():
-
-            mergeDAs = []
-
-            # Files valid at the same time, so we can merge them along the time dimension
-            for file in self.file_dict[cycle]:
-
-                filename = os.path.basename(file)
-                file = os.path.join(self.data_dir, filename)
-
-                pattern = get_pattern(file, file_patterns)
-
-                logger.info(f"Processing {filename=}")
-
-                gribfh = grib2io.open(file, "r")
-
-                for var_dict in self.varinfo[pattern]:
-
-                    variable_names = var_dict["variables"]
-                    levels = var_dict["levels"]
-
-                    for var_name in variable_names:
-                        logger.info(f"Extracting variable: {var_name} at levels: {levels}")
-                        if len(levels) > 1:
-                            da = self.get_dataarray_3d(gribfh, var_name, levels)
-                        else:
-                            da = self.get_dataarray_2d(gribfh, var_name, levels[0])
-                        mergeDAs.append(da)
-
-                gribfh.close()
-
-            ds = xr.merge(mergeDAs, compat="no_conflicts")
-            mergeDSs.append(ds)
-            ds.close()
-
-        # Concatenate along the time dimension
-        ds = xr.concat(mergeDSs, dim="time")
-
-        # Get 2D static data from the f000 file
-        # From the file list find the first file that ends with .f000
-        f000_file = next((f for f in self.file_dict[self.current_cycle] if f.endswith(".f000")), None)
-        if f000_file is None:
-            logger.error("No 00Z f000 file found.")  # TODO: should this raise an error?
-            raise FileNotFoundError("No f000 file found.")
-
-        f000_file = os.path.join(self.data_dir, os.path.basename(f000_file))
-        gribfh = grib2io.open(f000_file, "r")
-        static_vars = ["LAND", "HGT"]
-        for var_name in static_vars:
-            da = self.get_dataarray_2d(gribfh, var_name, "surface")
-            ds = xr.merge([ds, da], compat="no_conflicts")
-        gribfh.close()
-
-        # Rename variables to match graphcast naming conventions
-        ds = self.rename_dsvars(ds)
-
-        # Add datetime coordinate
-        ds = ds.assign_coords(datetime=ds.time)
-
-        # Adjust time values relative to the first time step
-        ds["time"] = ds["time"] - ds["time"][0]
-
-        # Expand dimensions
-        ds = ds.expand_dims(dim="batch")
-        ds["datetime"] = ds["datetime"].expand_dims(dim="batch")
-
-        # Squeeze dimensions
-        ds["geopotential_at_surface"] = ds["geopotential_at_surface"].squeeze("batch")
-        ds["land_sea_mask"] = ds["land_sea_mask"].squeeze("batch")
-
-        # Update geopotential unit to m2/s2 by multiplying 9.80665
-        ds["geopotential_at_surface"] = ds["geopotential_at_surface"] * 9.80665
-        ds["geopotential"] = ds["geopotential"] * 9.80665
-
-        # For GEFS, if total_precipitation_6hr is missing, create it with zeros
-        if "total_precipitation_6hr" not in ds:
-            logger.warning("total_precipitation_6hr variable not found. Creating it with zeros.")
-            ds["total_precipitation_6hr"] = xr.zeros_like(ds["2m_temperature"])
-        else:
-            # Update total_precipitation_6hr unit to (m) from (kg/m^2) by dividing it by 1000kg/m³
-            ds["total_precipitation_6hr"] = ds["total_precipitation_6hr"] / 1000.0
-
-        ds.to_netcdf(self.output_netcdf)
-        ds.close()
-
-        return
-
-    @staticmethod
-    def rename_dsvars(ds: xr.Dataset, rename_dict: dict = None) -> xr.Dataset:
-        """Rename dataset variables to match graphcast naming conventions."""
-        logger.info("Renaming dataset variables to match graphcast naming conventions.")
-
-        # Default rename dictionary  #TODO: move to yaml?
-        if rename_dict is None:
-            rename_dict = {
-                "LAND_surface": "land_sea_mask",
-                "HGT_surface": "geopotential_at_surface",
-                "PRMSL_meansealevel": "mean_sea_level_pressure",
-                "TMP_2maboveground": "2m_temperature",
-                "UGRD_10maboveground": "10m_u_component_of_wind",
-                "VGRD_10maboveground": "10m_v_component_of_wind",
-                "HGT": "geopotential",
-                "TMP": "temperature",
-                "SPFH": "specific_humidity",
-                "VVEL": "vertical_velocity",
-                "UGRD": "u_component_of_wind",
-                "VGRD": "v_component_of_wind",
-            }
-            if "APCP_surface" in ds:
-                rename_dict["APCP_surface"] = "total_precipitation_6hr"
-            logger.debug(f"Using rename_dict: {rename_dict}")
-
-        ds = ds.rename(rename_dict)
-
-        return ds
-
-    @staticmethod
-    def get_dataarray_2d(grbfile, var_name, desired_level):
-
-        logger.info(f"Getting 2D data for variable: {var_name} at level: {desired_level}")
-
-        msg = grbfile.select(shortName=var_name, level=desired_level)[0]
-
-        # create a netcdf dataset using the matching grib message
-        lats, lons = msg.latlons()
-        lats = lats[:,0]
-        lons = lons[0,:]
-
-        #check latitude range
-        reverse_lat = False
-        if lats[0] > 0:
-            reverse_lat = True
-
-        steps = msg.validDate
-        data = msg.data
-        if reverse_lat:
-            data = data[::-1, :]
-            lats = lats[::-1]
-
-        var_name2 = f'{var_name}_{"".join(desired_level.split())}'
-        if len(data.shape) == 2:
-            da = xr.Dataset(
-                data_vars={
-                    var_name2: (["lat", "lon"], data.astype("float32"))
-                },
-                coords={
-                    "lon": lons.astype("float32"),
-                    "lat": lats.astype("float32"),
-                    "time": steps,
-                }
-            )
-        elif len(data.shape) == 3:
-            logger.warning("3D data found in 2D function")
-            da = xr.Dataset(
-                data_vars={
-                    var_name: (["level", "lat", "lon"], data.astype("float32"))
-                },
-                coords={
-                    "lon": lons.astype("float32"),
-                    "lat": lats.astype("float32"),
-                    "level": np.array(desired_level).astype("int32"),
-                    "time": steps,
-                }
-            )
-
-        return da
-
-    @staticmethod
-    def get_dataarray_3d(grbfile, var_name, desired_level):
-
-        logger.info(f"Getting 3D data for variable: {var_name} at levels: {desired_level}")
-        data, levels = [], []
-        for ii, level in enumerate(desired_level):
-            msg = grbfile.select(shortName=var_name, level=level)[0]
-
-            if ii == 0:
-                lats, lons = msg.latlons()
-                lats = lats[:,0]
-                lons = lons[0,:]
-
-                #check latitude range, graphcast needs [-90, 90]
-                reverse_lat = False
-                if lats[0] > 0:
-                    reverse_lat = True
-                steps = msg.validDate
-            data.append(msg.data)
-            levels.append(int(level.split(' ')[0]))
-
-        data = np.array(data)
-        if reverse_lat:
-            data = data[:, ::-1, :]
-            lats = lats[::-1]
-
-        da = xr.Dataset(
-            data_vars={
-                var_name: (["level", "lat", "lon"], data.astype("float32"))
-            },
-            coords={
-                "lon": lons.astype("float32"),
-                "lat": lats.astype("float32"),
-                "level": np.array(levels).astype("int32"),
-                "time": steps,
-            }
-        )
-
-        return da
-
-if __name__ == "__main__":
-
-    import argparse
-    from datetime import datetime
-    from mlglobal.logger import setup_logging
-
-    setup_logging()
-
-    parser = argparse.ArgumentParser(
-        description="Process data to generate initial conditions for graphcast"
-        )
-
-    parser.add_argument(
-            "--current-cycle",
-            help="Current cycle in YYYYMMDDHH format",
-            type=str,
-            metavar="YYYYMMDDHH",
-            required=True
-            )
-    parser.add_argument(
-            "--yaml",
-            help="YAML file containing variable information",
-            type=str,
-            required=True
-            )
-    parser.add_argument(
-            "--data-dir",
-            help="Directory containing input GRIB2 files",
-            type=str,
-            default="./data",
-            required=False
-            )
-    parser.add_argument(
-            "--output",
-            help="Output NetCDF file",
-            type=str,
-            default="ic.nc",
-            required=False
-            )
-
-    args = parser.parse_args()
-
-    prep = PrepareIC(
-        current_cycle=datetime.strptime(args.current_cycle, "%Y%m%d%H"),
-        varinfo_file=args.yaml,
-        member=None,
-        num_levels=13,
-        output_netcdf=args.output,
-        data_dir=args.data_dir
-    )
-
-    logger.info("Starting IC preparation...")
-    prep.process_files()
diff --git a/src/mlglobal/ic_downloader.py b/src/mlglobal/ic_downloader.py
index 26f249d..982725c 100644
--- a/src/mlglobal/ic_downloader.py
+++ b/src/mlglobal/ic_downloader.py
@@ -2,8 +2,6 @@
 import shutil
 from datetime import timedelta
 from logging import getLogger
-from pprint import pprint
-
 
 logger = getLogger(__name__)
 
@@ -26,16 +24,24 @@ def __init__(self, current_cycle, num_levels=13, member=None):
     def _gfs_file_info(self):
 
         # Template for GFS files
-        template = f"gfs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/gfs.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
+        template = f"gfs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/gfs.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"  # noqa: F541
 
         # From current cycle
-        pgrb2_0p25_f000 = template.format(cycle=self.current_cycle, fspec="pgrb2.0p25", fhour=0)
+        pgrb2_0p25_f000 = template.format(
+            cycle=self.current_cycle, fspec="pgrb2.0p25", fhour=0
+        )
         if self.num_levels == 37:  # Need additional pgrb2b file for 37 level data
-            pgrb2b_0p25_f000 = template.format(cycle=self.current_cycle, fspec="pgrb2b.0p25", fhour=0)
+            pgrb2b_0p25_f000 = template.format(
+                cycle=self.current_cycle, fspec="pgrb2b.0p25", fhour=0
+            )
 
         # From current cycle - 6 hours
-        pgrb2_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=0)
-        pgrb2_0p25_f006_m6 = template.format(cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=6)
+        pgrb2_0p25_f000_m6 = template.format(
+            cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=0
+        )
+        pgrb2_0p25_f006_m6 = template.format(
+            cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=6
+        )
 
         # Create a flat list of files valid at current and current-6h cycles
         file_dict = {}
@@ -49,15 +55,39 @@ def _gfs_file_info(self):
     def _gefs_file_info(self):
 
         # Template for GEFS files
-        template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{{member}}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"
+        template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{{member}}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}"  # noqa: F541
 
         # From current cycle
-        pgrb2_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0, member=self.member)
-        pgrb2s_0p25_f000 = template.format(cycle=self.current_cycle, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0, member=self.member)
+        pgrb2_0p25_f000 = template.format(
+            cycle=self.current_cycle,
+            fspec_dir="pgrb2p25",
+            fspec="pgrb2.0p25",
+            fhour=0,
+            member=self.member,
+        )
+        pgrb2s_0p25_f000 = template.format(
+            cycle=self.current_cycle,
+            fspec_dir="pgrb2sp25",
+            fspec="pgrb2s.0p25",
+            fhour=0,
+            member=self.member,
+        )
 
         # From current cycle - 6 hours
-        pgrb2_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2p25", fspec="pgrb2.0p25", fhour=0, member=self.member)
-        pgrb2s_0p25_f000_m6 = template.format(cycle=self.current_cycle_m6h, fspec_dir="pgrb2sp25", fspec="pgrb2s.0p25", fhour=0, member=self.member)
+        pgrb2_0p25_f000_m6 = template.format(
+            cycle=self.current_cycle_m6h,
+            fspec_dir="pgrb2p25",
+            fspec="pgrb2.0p25",
+            fhour=0,
+            member=self.member,
+        )
+        pgrb2s_0p25_f000_m6 = template.format(
+            cycle=self.current_cycle_m6h,
+            fspec_dir="pgrb2sp25",
+            fspec="pgrb2s.0p25",
+            fhour=0,
+            member=self.member,
+        )
 
         file_dict = {}
         file_dict[self.current_cycle] = [pgrb2_0p25_f000, pgrb2s_0p25_f000]
@@ -76,7 +106,8 @@ def __init__(
         download_source="local",
         local_directory="./data",
         bucket_name=None,
-        root_directory=None
+        bucket_root_directory=None,
+        root_directory=None,
     ):
         self.current_cycle = current_cycle
         self.num_levels = num_levels
@@ -84,10 +115,13 @@ def __init__(
         self.download_source = download_source
         self.local_directory = local_directory
         self.bucket_name = bucket_name
+        self.bucket_root_directory = bucket_root_directory
         self.root_directory = root_directory
 
         # Generate the lookup dictionary
-        lookup = FileLookup(self.current_cycle, member=self.member, num_levels=self.num_levels)
+        lookup = FileLookup(
+            self.current_cycle, member=self.member, num_levels=self.num_levels
+        )
         self.file_dict = lookup.get_file_info()
 
         # Flatten the file_dict to a simple list of files
@@ -98,12 +132,14 @@ def __init__(
 
         if self.download_source in ["s3"]:
             aws_profile = os.environ.get("AWS_PROFILE", "default")
-            self.s3 = self.get_s3_client_by_bucket_type(self.bucket_name, profile_name=aws_profile)
+            self.s3 = self.get_s3_client_by_bucket_type(
+                self.bucket_name, profile_name=aws_profile
+            )
 
         os.makedirs(self.local_directory, exist_ok=True)
 
     @staticmethod
-    def get_s3_client_by_bucket_type(bucket_name, profile_name='default'):
+    def get_s3_client_by_bucket_type(bucket_name, profile_name="default"):
         """
         Initializes and returns a boto3 S3 client for a given bucket.
 
@@ -122,9 +158,9 @@ def get_s3_client_by_bucket_type(bucket_name, profile_name='default'):
 
         try:
             import boto3
-            from botocore.exceptions import ClientError
             from botocore import UNSIGNED
             from botocore.config import Config
+            from botocore.exceptions import ClientError
         except ImportError as ee:
             raise ImportError(
                 "boto3 and botocore are required for S3 operations."
@@ -135,13 +171,17 @@ def get_s3_client_by_bucket_type(bucket_name, profile_name='default'):
         try:
             # Check if the bucket can be accessed publicly without authentication
             s3.head_bucket(Bucket=bucket_name)
-            logger.info(f"Bucket '{bucket_name}' is public. Returning an unsigned client.")
+            logger.info(
+                f"Bucket '{bucket_name}' is public. Returning an unsigned client."
+            )
             return s3
         except ClientError as ee:
-            error_code = ee.response['Error']['Code']
+            error_code = ee.response["Error"]["Code"]
             # 2. If it's a 403 Forbidden, the bucket is likely private.
-            if error_code in ('403', '404'):
-                logger.warning(f"Bucket '{bucket_name}' is not public. Returning client with profile '{profile_name}'.")
+            if error_code in ("403", "404"):
+                logger.warning(
+                    f"Bucket '{bucket_name}' is not public. Returning client with profile '{profile_name}'."
+                )
 
                 # Create a session with the specified profile
                 session = boto3.Session(profile_name=profile_name)
@@ -159,29 +199,47 @@ def get_s3_client_by_bucket_type(bucket_name, profile_name='default'):
 
     @staticmethod
     def get_s3_objects(s3, bucket_name: str, prefix: str) -> list:
+        """Retrieve a list of objects from an S3 bucket with a specific prefix.
+        This method handles pagination to ensure all objects are retrieved.
+        By default, S3 returns up to 1000 objects per request.
+
+        Parameters
+        ----------
+        s3 : boto3.client
+            The S3 client to use for the operation.
+        bucket_name : str
+            The name of the S3 bucket.
+        prefix : str
+            The prefix to filter the objects.
+
+        Returns
+        -------
+        list
+            A list of S3 objects that match the prefix.
+        """
 
         objects = []
         continuation_token = None
 
         while True:
             list_kwargs = {
-                'Bucket': bucket_name,
-                'Prefix': prefix,
-                'MaxKeys': 1000  # Explicitly set MaxKeys, though it's default
+                "Bucket": bucket_name,
+                "Prefix": prefix,
+                "MaxKeys": 1000,  # Explicitly set MaxKeys, though it's default
             }
             if continuation_token:
-                list_kwargs['ContinuationToken'] = continuation_token
+                list_kwargs["ContinuationToken"] = continuation_token
 
             response = s3.list_objects_v2(**list_kwargs)
 
-            if 'Contents' in response:
-                objects.extend(response['Contents'])
+            if "Contents" in response:
+                objects.extend(response["Contents"])
 
-            if not response.get('IsTruncated'):
+            if not response.get("IsTruncated"):
                 # No more objects to retrieve, or less than 1000 objects in total
                 break
 
-            continuation_token = response.get('NextContinuationToken')
+            continuation_token = response.get("NextContinuationToken")
             if not continuation_token:
                 # Should not happen if 'IsTruncated' is True, but as a safeguard
                 break
@@ -217,9 +275,16 @@ def get_data_from_s3(self, file_list: list, local_directory: str) -> None:
             if os.path.exists(local_file_path):
                 logger.warning(f"File already exists, skipping: {file_name}")
                 continue
-            file_name_in_bucket = self.root_directory + "/" + file_name if self.root_directory else file_name
+            file_name_in_bucket = (
+                self.bucket_root_directory + "/" + file_name
+                if self.bucket_root_directory
+                else file_name
+            )
+            print(f"Downloading {file_name_in_bucket} to {local_file_path}")
             try:
-                self.s3.download_file(self.bucket_name, file_name_in_bucket, local_file_path)
+                self.s3.download_file(
+                    self.bucket_name, file_name_in_bucket, local_file_path
+                )
                 logger.info(f"Downloaded:  {file_name} -> {local_directory}")
             except Exception as ee:
                 logger.error(f"Error downloading {file_name}: {ee}")
@@ -253,7 +318,11 @@ def get_data_from_local(self, file_list: list, local_directory: str) -> None:
             if os.path.exists(local_file_path):
                 logger.warning(f"File already exists, skipping: {file_name}")
                 continue
-            file_name_in_remote = self.root_directory + "/" + file_name if self.root_directory else file_name
+            file_name_in_remote = (
+                self.root_directory + "/" + file_name
+                if self.root_directory
+                else file_name
+            )
             try:
                 shutil.copy2(file_name_in_remote, local_file_path)
                 logger.info(f"Copied:  {file_name} -> {local_directory}")
@@ -265,8 +334,7 @@ def get_data_from_local(self, file_list: list, local_directory: str) -> None:
 
     def get_data(self) -> None:
 
-        GET_DATA_MAP = {"s3": self.get_data_from_s3,
-                        "local": self.get_data_from_local}
+        GET_DATA_MAP = {"s3": self.get_data_from_s3, "local": self.get_data_from_local}
 
         logger.info("Starting download...")
         GET_DATA_MAP[self.download_source](self.file_list, self.local_directory)
diff --git a/src/mlglobal/ic_processor.py b/src/mlglobal/ic_processor.py
new file mode 100644
index 0000000..25f3e78
--- /dev/null
+++ b/src/mlglobal/ic_processor.py
@@ -0,0 +1,351 @@
+import fnmatch
+import logging
+import os
+from datetime import datetime
+
+import grib2io
+import numpy as np
+import xarray as xr
+import yaml
+
+from mlglobal.ic_downloader import FileLookup
+
+logger = logging.getLogger(__name__)
+
+
+class ICProcessor:
+    """Class to prepare initial conditions for graphcast from downloaded data files."""
+
+    def __init__(
+        self,
+        current_cycle: datetime,
+        varinfo_file: str,
+        member: str | None = None,
+        num_levels: int = 13,
+        data_dir: str = "./data",
+        output_netcdf: str = "./ic.nc",
+    ) -> None:
+
+        self.current_cycle = current_cycle
+        self.varinfo = self.load_yaml(varinfo_file)
+        self.member = member
+        self.num_levels = num_levels
+        self.data_dir = data_dir
+        self.output_netcdf = output_netcdf
+
+        # Generate the lookup dictionary
+        lookup = FileLookup(
+            self.current_cycle, member=self.member, num_levels=self.num_levels
+        )
+        self.file_dict = lookup.get_file_info()
+
+    @staticmethod
+    def load_yaml(yaml_file: str) -> dict:
+        """Load a YAML file and return the contents as a dictionary.
+        Parameters
+        ----------
+        yaml_file : str
+            Path to the YAML file.
+        Returns
+        -------
+        dict
+            Dictionary containing the contents of the YAML file.
+        Raises
+        ------
+        FileNotFoundError
+            If the YAML file does not exist.
+        yaml.YAMLError
+            If there is an error parsing the YAML file.
+        """
+
+        try:
+            with open(yaml_file, "r") as fh:
+                yaml_dict = yaml.safe_load(fh)
+        except FileNotFoundError:
+            logger.error(f"YAML file {yaml_file} not found.")
+            raise FileNotFoundError(f"YAML file {yaml_file} not found.")
+        except yaml.YAMLError as e:
+            logger.error(f"Error parsing YAML file {yaml_file}: {e}")
+            raise yaml.YAMLError(f"Error parsing YAML file {yaml_file}: {e}")
+
+        return yaml_dict
+
+    def get_matching_pattern(self, file, patterns, fatal=True):
+
+        filename = os.path.basename(file)
+
+        matched = False
+        pattern = None
+        # Check each pattern to see if it matches the filename
+        for pp in patterns:
+            if fnmatch.fnmatch(file, os.path.join(self.data_dir, "*" + pp)):
+                pattern = pp
+                logger.debug(f"Matched pattern: {pattern} in file: {filename}")
+                matched = True
+                break
+
+        if not matched and fatal:
+            logger.error(f"No pattern matched for file: {filename}")
+            raise FileNotFoundError(f"No pattern matched for file: {filename}")
+
+        return pattern
+
+    def process_files(self):
+        """
+        Process input files and extract relevant variables.
+        This function reads GRIB2 files, extracts specified variables,
+        and saves them into a NetCDF file.
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        FileNotFoundError
+            If any of the input files do not exist.
+        ValueError
+            If there is an error processing the files.
+        IOError
+            If there is an error saving the NetCDF file.
+        """
+
+        logger.info("Starting to process files...")
+
+        file_patterns = list(self.varinfo["time_variant"].keys())
+        mergeDSs = []
+        for cycle in self.file_dict.keys():
+
+            mergeDAs = []
+
+            # Files valid at the same time, so we can merge them along the time dimension
+            for file in self.file_dict[cycle]:
+
+                logger.info(f"Processing {file=}")
+
+                file = os.path.join(self.data_dir, os.path.basename(file))
+
+                pattern = self.get_matching_pattern(file, file_patterns)
+
+                gribfh = grib2io.open(file, "r")
+                for var_dict in self.varinfo["time_variant"][pattern]:
+                    for var_name in var_dict["variables"]:
+                        logger.info(
+                            f"Extracting variable: {var_name} at levels: {var_dict['levels']}"
+                        )
+                        da = self.get_dataarray(gribfh, var_name, var_dict["levels"])
+                        mergeDAs.append(da)
+                gribfh.close()
+
+            ds = xr.merge(mergeDAs, compat="no_conflicts")
+            mergeDSs.append(ds)
+            ds.close()
+
+        # Concatenate along the time dimension
+        ds = xr.concat(mergeDSs, dim="time")
+
+        logger.info("Processing static variables...")
+
+        # Now handle the static variables
+        file_patterns = list(self.varinfo["time_invariant"].keys())
+        mergeDAs = []
+        for file in self.file_dict[self.current_cycle]:
+
+            logger.info(f"Processing {file=}")
+
+            file = os.path.join(self.data_dir, os.path.basename(file))
+
+            pattern = self.get_matching_pattern(file, file_patterns, fatal=False)
+            if pattern is None:  # No pattern matched, so skip this file
+                continue
+
+            gribfh = grib2io.open(file, "r")
+            for var_dict in self.varinfo["time_invariant"][pattern]:
+                for var_name in var_dict["variables"]:
+                    logger.info(
+                        f"Extracting variable: {var_name} at levels: {var_dict['levels']}"
+                    )
+                    da = self.get_dataarray(gribfh, var_name, var_dict["levels"])
+                    mergeDAs.append(da)
+            gribfh.close()
+
+        ds_static = xr.merge(mergeDAs, compat="no_conflicts")
+        ds = xr.merge([ds, ds_static], compat="no_conflicts")
+        ds_static.close()
+
+        # Rename variables to match graphcast naming conventions
+        ds = self.rename_dsvars(ds)
+
+        # Add datetime coordinate
+        ds = ds.assign_coords(datetime=ds.time)
+
+        # Adjust time values relative to the first time step
+        ds["time"] = ds["time"] - ds["time"][0]
+
+        # Expand dimensions
+        ds = ds.expand_dims(dim="batch")
+        ds["datetime"] = ds["datetime"].expand_dims(dim="batch")
+
+        # Squeeze dimensions
+        ds["geopotential_at_surface"] = ds["geopotential_at_surface"].squeeze("batch")
+        ds["land_sea_mask"] = ds["land_sea_mask"].squeeze("batch")
+
+        # Update geopotential unit to m2/s2 by multiplying 9.80665
+        ds["geopotential_at_surface"] = ds["geopotential_at_surface"] * 9.80665
+        ds["geopotential"] = ds["geopotential"] * 9.80665
+
+        # For GEFS, if total_precipitation_6hr is missing, create it with zeros
+        if "total_precipitation_6hr" not in ds:
+            logger.warning(
+                "total_precipitation_6hr variable not found. Creating it with zeros."
+            )
+            ds["total_precipitation_6hr"] = xr.zeros_like(ds["2m_temperature"])
+        else:
+            # Update total_precipitation_6hr unit to (m) from (kg/m^2) by dividing it by 1000kg/m³
+            ds["total_precipitation_6hr"] = ds["total_precipitation_6hr"] / 1000.0
+
+        # Save to NetCDF
+        logger.info(f"Saving processed data to {self.output_netcdf}")
+        os.makedirs(os.path.dirname(self.output_netcdf), exist_ok=True)
+        ds.to_netcdf(self.output_netcdf)
+        ds.close()
+
+        return
+
+    @staticmethod
+    def rename_dsvars(ds: xr.Dataset, rename_dict: dict = None) -> xr.Dataset:
+        """Rename dataset variables to match graphcast naming conventions."""
+        logger.info("Renaming dataset variables to match graphcast naming conventions.")
+
+        # Default rename dictionary  # TODO: move to yaml?
+        if rename_dict is None:
+            rename_dict = {
+                "LAND_surface": "land_sea_mask",
+                "HGT_surface": "geopotential_at_surface",
+                "PRMSL_meansealevel": "mean_sea_level_pressure",
+                "TMP_2maboveground": "2m_temperature",
+                "UGRD_10maboveground": "10m_u_component_of_wind",
+                "VGRD_10maboveground": "10m_v_component_of_wind",
+                "HGT": "geopotential",
+                "TMP": "temperature",
+                "SPFH": "specific_humidity",
+                "VVEL": "vertical_velocity",
+                "UGRD": "u_component_of_wind",
+                "VGRD": "v_component_of_wind",
+            }
+            if "APCP_surface" in ds:
+                rename_dict["APCP_surface"] = "total_precipitation_6hr"
+            logger.debug(f"Using rename_dict: {rename_dict}")
+
+        ds = ds.rename(rename_dict)
+
+        return ds
+
+    @staticmethod
+    def get_dataarray(gh, var_name: str, levels: list | int | str) -> xr.Dataset:
+        """Get a data array from the GRIB file.
+
+        Parameters
+        ----------
+        gh : Handle of opened grib2io file
+        var_name : str
+            Name of the variable to extract
+        levels : list | int | str
+            Levels to extract (can be a single level or a list of levels)
+
+        Returns
+        -------
+        xr.Dataset
+            A dataset containing the extracted data
+        """
+
+        if isinstance(levels, list) and len(levels) > 1:
+            return ICProcessor._get_dataarray_3d(gh, var_name, levels)
+        else:
+            return ICProcessor._get_dataarray_2d(
+                gh, var_name, levels[0] if isinstance(levels, list) else levels
+            )
+
+    @staticmethod
+    def _get_dataarray_2d(gh, var_name, level):
+
+        logger.info(f"Getting 2D data for variable: {var_name} at level: {level}")
+
+        msg = gh.select(shortName=var_name, level=level)[0]
+
+        # create a netcdf dataset using the matching grib message
+        lats, lons = msg.latlons()
+        lats = lats[:, 0]
+        lons = lons[0, :]
+
+        # check latitude range, graphcast needs [-90, 90]
+        reverse_lat = False
+        if lats[0] > 0:
+            reverse_lat = True
+
+        steps = msg.validDate
+        data = msg.data
+        if reverse_lat:
+            data = data[::-1, :]
+            lats = lats[::-1]
+
+        if len(data.shape) == 2:
+            var_name2 = f'{var_name}_{"".join(level.split())}'
+            da = xr.Dataset(
+                data_vars={var_name2: (["lat", "lon"], data.astype("float32"))},
+                coords={
+                    "lon": lons.astype("float32"),
+                    "lat": lats.astype("float32"),
+                    "time": steps,
+                },
+            )
+        elif len(data.shape) == 3:
+            logger.warning("3D data found in 2D function")
+            da = xr.Dataset(
+                data_vars={var_name: (["level", "lat", "lon"], data.astype("float32"))},
+                coords={
+                    "lon": lons.astype("float32"),
+                    "lat": lats.astype("float32"),
+                    "level": np.array(level).astype("int32"),
+                    "time": steps,
+                },
+            )
+
+        return da
+
+    @staticmethod
+    def _get_dataarray_3d(gh, var_name, levels):
+
+        logger.info(f"Getting 3D data for variable: {var_name} at levels: {levels}")
+        data, vlevs = [], []
+        for ii, level in enumerate(levels):
+            msg = gh.select(shortName=var_name, level=level)[0]
+
+            if ii == 0:
+                lats, lons = msg.latlons()
+                lats = lats[:, 0]
+                lons = lons[0, :]
+
+                # check latitude range, graphcast needs [-90, 90]
+                reverse_lat = False
+                if lats[0] > 0:
+                    reverse_lat = True
+                steps = msg.validDate
+            data.append(msg.data)
+            vlevs.append(int(level.split(" ")[0]))
+
+        data = np.array(data)
+        if reverse_lat:
+            data = data[:, ::-1, :]
+            lats = lats[::-1]
+
+        da = xr.Dataset(
+            data_vars={var_name: (["level", "lat", "lon"], data.astype("float32"))},
+            coords={
+                "lon": lons.astype("float32"),
+                "lat": lats.astype("float32"),
+                "level": np.array(vlevs).astype("int32"),
+                "time": steps,
+            },
+        )
+
+        return da
diff --git a/src/mlglobal/logger.py b/src/mlglobal/logger.py
index c2cb966..9554c95 100644
--- a/src/mlglobal/logger.py
+++ b/src/mlglobal/logger.py
@@ -1,5 +1,5 @@
 import logging
-import os
+
 """
 Logger setup module for the mlglobal project.
 
@@ -20,20 +20,18 @@
 
 log = logging.getLogger()
 
+
 def setup_logging(
-    level=os.environ.get("LOGGING_LEVEL", logging.INFO),
-    fmt="[%(asctime)s] %(levelname)8s - %(message)s",
-    datefmt="%Y-%m-%dT%H:%M:%S"
+    fmt: str = "[%(asctime)s] %(levelname)8s - %(message)s",
+    datefmt: str = "%Y-%m-%dT%H:%M:%S",
+    debug: bool = False,
 ):
 
     logger = logging.getLogger()
     for handler in logger.handlers:
         logger.removeHandler(handler)
 
-    kwargs: dict = {
-        "datefmt": datefmt,
-        "format": fmt,
-        "level": level
-    }
+    level = logging.DEBUG if debug else logging.INFO
+    kwargs: dict = {"datefmt": datefmt, "format": fmt, "level": level}
 
     logging.basicConfig(**kwargs)

From c5c67054559d67f133be95df5af83f6f6f2c9906 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 00:12:51 -0400
Subject: [PATCH 09/21] fix error in GA yaml;

---
 .github/workflows/pynorms.yaml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.github/workflows/pynorms.yaml b/.github/workflows/pynorms.yaml
index d4b9b0e..1b771a8 100644
--- a/.github/workflows/pynorms.yaml
+++ b/.github/workflows/pynorms.yaml
@@ -5,7 +5,6 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-      - name: Check python coding norms with pre-commit
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:

From 913e602e8c263d6a48337b30922f04b092dd3019 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 00:18:58 -0400
Subject: [PATCH 10/21] exclude some dirs in pre-commmit

---
 .pre-commit-config.yaml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f7ec5b1..9c65b4f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,7 +16,7 @@ repos:
       - id: black
         language_version: python3.12 # Specify your Python version
         files: ^(src/|tests/)
-        exclude: ^(oper/|graphcast/|training/)
+        exclude: ^(oper/|graphcast/|training/)$
 
   # Flake8 linter for style guide enforcement and error checking
   - repo: https://github.com/PyCQA/flake8
@@ -26,7 +26,7 @@ repos:
         name: flake8
         entry: flake8
         files: ^(src/|tests/)
-        exclude: ^(oper/|graphcast/|training/)
+        exclude: ^(oper/|graphcast/|training/)$
 
   # isort for sorting imports alphabetically and separating them into sections
   - repo: https://github.com/PyCQA/isort
@@ -36,4 +36,4 @@ repos:
         name: isort (python)
         args: ["--profile", "black"] # Use the black profile for compatibility
         files: ^(src/|tests/)
-        exclude: ^(oper/|graphcast/|training/)
+        exclude: ^(oper/|graphcast/|training/)$

From a5fb0173650843c0c88d2d4406d72198fc110dc0 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 00:24:41 -0400
Subject: [PATCH 11/21] exclude some dirs in pre-commmit

---
 .pre-commit-config.yaml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9c65b4f..fba6bfe 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -15,8 +15,8 @@ repos:
     hooks:
       - id: black
         language_version: python3.12 # Specify your Python version
-        files: ^(src/|tests/)
-        exclude: ^(oper/|graphcast/|training/)$
+        files: ^(src/|tests/)$
+        exclude: ^(oper/|graphcast/|training/)
 
   # Flake8 linter for style guide enforcement and error checking
   - repo: https://github.com/PyCQA/flake8
@@ -26,7 +26,7 @@ repos:
         name: flake8
         entry: flake8
         files: ^(src/|tests/)
-        exclude: ^(oper/|graphcast/|training/)$
+        exclude: ^(oper/|graphcast/|training/)
 
   # isort for sorting imports alphabetically and separating them into sections
   - repo: https://github.com/PyCQA/isort
@@ -36,4 +36,4 @@ repos:
         name: isort (python)
         args: ["--profile", "black"] # Use the black profile for compatibility
         files: ^(src/|tests/)
-        exclude: ^(oper/|graphcast/|training/)$
+        exclude: ^(oper/|graphcast/|training/)

From 7806b1f53648f4874dbfc2ffaef5d9936d39aebe Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 00:57:12 -0400
Subject: [PATCH 12/21] add tests and automation

---
 .github/workflows/pytests.yaml |  33 +++
 pyproject.toml                 |   1 +
 tests/test_file_lookup.py      | 142 +++++++++++
 tests/test_ic_downloader.py    | 428 +++++++++++++++++++++++++++++++++
 4 files changed, 604 insertions(+)
 create mode 100644 .github/workflows/pytests.yaml
 create mode 100644 tests/test_file_lookup.py
 create mode 100644 tests/test_ic_downloader.py

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
new file mode 100644
index 0000000..522ceb2
--- /dev/null
+++ b/.github/workflows/pytests.yaml
@@ -0,0 +1,33 @@
+name: pytests
+on: [push, pull_request]
+
+jobs:
+  run_pytests:
+    runs-on: ubuntu-latest
+    name: Install wxflow and run tests with pytests
+    strategy:
+      matrix:
+        python: ["3.11", "3.12", "3.13"]
+
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+
+      - name: Install (upgrade) python dependencies
+        run: |
+          pip install --upgrade pip
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Install
+        run: |
+          cd $GITHUB_WORKSPACE
+          pip install .[dev]
+
+      - name: Run pytests
+        run: |
+          cd $GITHUB_WORKSPACE
+          pytest -v
diff --git a/pyproject.toml b/pyproject.toml
index 801f0bf..c0b395c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,6 +49,7 @@ gpu = ["jax[cuda12]>=0.4.1"]
 dev = [
     "pytest>=7.0.0",
     "pytest-cov>=4.0.0",
+    "pytest-mock>=3.10.0",
     "pre-commit>=4.3.0",
     "black>=24.4.2",
     "isort>=6.0.1",
diff --git a/tests/test_file_lookup.py b/tests/test_file_lookup.py
new file mode 100644
index 0000000..26e729e
--- /dev/null
+++ b/tests/test_file_lookup.py
@@ -0,0 +1,142 @@
+import pytest
+from datetime import datetime, timedelta
+from mlglobal.ic_downloader import FileLookup
+
+
+class TestFileLookup:
+    """Test cases for the FileLookup class."""
+
+    def test_init_gfs(self):
+        """Test FileLookup initialization for GFS mode."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member=None)
+
+        assert lookup.current_cycle == current_cycle
+        assert lookup.num_levels == 13
+        assert lookup.member is None
+        assert lookup.current_cycle_m6h == current_cycle - timedelta(hours=6)
+        assert lookup.get_file_info == lookup._gfs_file_info
+
+    def test_init_gefs(self):
+        """Test FileLookup initialization for GEFS mode."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member="c00")
+
+        assert lookup.current_cycle == current_cycle
+        assert lookup.num_levels == 13
+        assert lookup.member == "c00"
+        assert lookup.current_cycle_m6h == current_cycle - timedelta(hours=6)
+        assert lookup.get_file_info == lookup._gefs_file_info
+
+    def test_gfs_file_info_13_levels(self):
+        """Test GFS file info generation for 13 levels."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member=None)
+
+        file_dict = lookup._gfs_file_info()
+
+        # Check structure
+        assert len(file_dict) == 2
+        assert current_cycle in file_dict
+        assert lookup.current_cycle_m6h in file_dict
+
+        # Check current cycle files
+        current_files = file_dict[current_cycle]
+        assert len(current_files) == 2
+        assert "gfs.20230101/12/atmos/gfs.t12z.pgrb2.0p25.f000" in current_files[0]
+        assert "gfs.20230101/06/atmos/gfs.t06z.pgrb2.0p25.f006" in current_files[1]
+
+        # Check previous cycle files
+        prev_files = file_dict[lookup.current_cycle_m6h]
+        assert len(prev_files) == 1
+        assert "gfs.20230101/06/atmos/gfs.t06z.pgrb2.0p25.f000" in prev_files[0]
+
+    def test_gfs_file_info_37_levels(self):
+        """Test GFS file info generation for 37 levels."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=37, member=None)
+
+        file_dict = lookup._gfs_file_info()
+
+        # Check current cycle has additional pgrb2b file
+        current_files = file_dict[current_cycle]
+        assert len(current_files) == 3
+        assert any("pgrb2b.0p25" in f for f in current_files)
+
+    def test_gefs_file_info(self):
+        """Test GEFS file info generation."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member="c00")
+
+        file_dict = lookup._gefs_file_info()
+
+        # Check structure
+        assert len(file_dict) == 2
+        assert current_cycle in file_dict
+        assert lookup.current_cycle_m6h in file_dict
+
+        # Check current cycle files
+        current_files = file_dict[current_cycle]
+        assert len(current_files) == 2
+        assert "gec00.t12z.pgrb2.0p25.f000" in current_files[0]
+        assert "gec00.t12z.pgrb2s.0p25.f000" in current_files[1]
+
+        # Check previous cycle files
+        prev_files = file_dict[lookup.current_cycle_m6h]
+        assert len(prev_files) == 2
+        assert "gec00.t06z.pgrb2.0p25.f000" in prev_files[0]
+        assert "gec00.t06z.pgrb2s.0p25.f000" in prev_files[1]
+
+    @pytest.mark.parametrize("num_levels,expected_current_files", [
+        (13, 2),  # For 13 levels: pgrb2.0p25 files only
+        (37, 3),  # For 37 levels: includes pgrb2b.0p25 file
+    ])
+    def test_gfs_file_count_by_levels(self, num_levels, expected_current_files):
+        """Test that GFS mode generates correct number of files based on levels."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=num_levels, member=None)
+        file_dict = lookup._gfs_file_info()
+
+        # Check current cycle files
+        current_files = file_dict[current_cycle]
+        assert len(current_files) == expected_current_files
+
+    @pytest.mark.parametrize("member", ["c00", "p01", "p15", "p30"])
+    def test_gefs_members(self, member):
+        """Test GEFS mode with different member values."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member=member)
+        file_dict = lookup._gefs_file_info()
+
+        # Verify member appears in file paths
+        for files in file_dict.values():
+            for file_path in files:
+                assert f"ge{member}" in file_path
+
+    def test_current_cycle_m6h_calculation(self):
+        """Test that current_cycle_m6h is calculated correctly."""
+        current_cycle = datetime(2023, 1, 1, 18, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member=None)
+
+        expected_m6h = datetime(2023, 1, 1, 12, 0, 0)
+        assert lookup.current_cycle_m6h == expected_m6h
+
+    def test_file_info_delegation_gfs(self):
+        """Test that get_file_info correctly delegates to GFS method."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member=None)
+
+        gfs_result = lookup._gfs_file_info()
+        delegated_result = lookup.get_file_info()
+
+        assert gfs_result == delegated_result
+
+    def test_file_info_delegation_gefs(self):
+        """Test that get_file_info correctly delegates to GEFS method."""
+        current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+        lookup = FileLookup(current_cycle, num_levels=13, member="c00")
+
+        gefs_result = lookup._gefs_file_info()
+        delegated_result = lookup.get_file_info()
+
+        assert gefs_result == delegated_result
\ No newline at end of file
diff --git a/tests/test_ic_downloader.py b/tests/test_ic_downloader.py
new file mode 100644
index 0000000..1552e5a
--- /dev/null
+++ b/tests/test_ic_downloader.py
@@ -0,0 +1,428 @@
+import pytest
+import os
+import tempfile
+import shutil
+from datetime import datetime
+from botocore.exceptions import ClientError
+
+from mlglobal.ic_downloader import ICDownloader
+
+
+class TestICDownloader:
+    """Test cases for the ICDownloader class."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.test_dir = tempfile.mkdtemp()
+        self.current_cycle = datetime(2023, 1, 1, 12, 0, 0)
+
+    def teardown_method(self):
+        """Clean up test fixtures."""
+        if os.path.exists(self.test_dir):
+            shutil.rmtree(self.test_dir)
+
+    def test_init_gfs(self):
+        """Test ICDownloader initialization for GFS mode."""
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            num_levels=13,
+            member=None,
+            download_source="local",
+            local_directory=self.test_dir,
+            root_directory="/test/root"
+        )
+
+        assert downloader.current_cycle == self.current_cycle
+        assert downloader.num_levels == 13
+        assert downloader.member is None
+        assert downloader.download_source == "local"
+        assert downloader.local_directory == self.test_dir
+        assert downloader.root_directory == "/test/root"
+        assert len(downloader.file_list) > 0
+        assert os.path.exists(self.test_dir)
+
+    def test_init_gefs(self):
+        """Test ICDownloader initialization for GEFS mode."""
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            num_levels=13,
+            member="c00",
+            download_source="local",
+            local_directory=self.test_dir
+        )
+
+        assert downloader.member == "c00"
+        assert len(downloader.file_list) > 0
+
+    def test_init_s3(self, mocker):
+        """Test ICDownloader initialization for S3 source."""
+        mock_s3_client = mocker.MagicMock()
+
+        # Mock the get_s3_client_by_bucket_type method since boto3 is imported inside it
+        mock_get_s3_client = mocker.patch.object(
+            ICDownloader, 'get_s3_client_by_bucket_type',
+            return_value=mock_s3_client
+        )
+
+        # Mock the environment variable
+        mocker.patch.dict(os.environ, {'AWS_PROFILE': 'test-profile'})
+
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="s3",
+            bucket_name="test-bucket"
+        )
+
+        assert downloader.s3 == mock_s3_client
+        mock_get_s3_client.assert_called_once_with("test-bucket", profile_name="test-profile")
+
+    def test_get_s3_client_by_bucket_type_public(self, mocker):
+        """Test S3 client creation for public bucket."""
+        mock_client = mocker.MagicMock()
+        mock_boto3_client = mocker.patch('boto3.client')
+        mock_boto3_client.return_value = mock_client
+        mock_client.head_bucket.return_value = None  # Success
+
+        result = ICDownloader.get_s3_client_by_bucket_type("public-bucket")
+
+        assert result == mock_client
+        mock_client.head_bucket.assert_called_once_with(Bucket="public-bucket")
+
+    def test_get_s3_client_by_bucket_type_private(self, mocker):
+        """Test S3 client creation for private bucket."""
+        # Mock public client failing
+        mock_public_client = mocker.MagicMock()
+        mock_boto3_client = mocker.patch('boto3.client')
+        mock_boto3_client.return_value = mock_public_client
+        mock_public_client.head_bucket.side_effect = ClientError(
+            {"Error": {"Code": "403"}}, "head_bucket"
+        )
+
+        # Mock private client
+        mock_private_client = mocker.MagicMock()
+        mock_session_instance = mocker.MagicMock()
+        mock_session = mocker.patch('boto3.Session')
+        mock_session.return_value = mock_session_instance
+        mock_session_instance.client.return_value = mock_private_client
+        mock_session_instance.get_credentials.return_value.get_frozen_credentials.return_value = mocker.MagicMock(
+            access_key="test-key", secret_key="test-secret"
+        )
+
+        result = ICDownloader.get_s3_client_by_bucket_type("private-bucket", "test-profile")
+
+        assert result == mock_private_client
+        mock_session.assert_called_once_with(profile_name="test-profile")
+
+    def test_get_s3_client_by_bucket_type_error(self, mocker):
+        """Test S3 client creation with error."""
+        mock_client = mocker.MagicMock()
+        mock_boto3_client = mocker.patch('boto3.client')
+        mock_boto3_client.return_value = mock_client
+        mock_client.head_bucket.side_effect = ClientError(
+            {"Error": {"Code": "NoSuchBucket"}}, "head_bucket"
+        )
+
+        result = ICDownloader.get_s3_client_by_bucket_type("nonexistent-bucket")
+
+        assert result is None
+
+    def test_get_s3_objects(self, mocker):
+        """Test S3 objects retrieval with pagination."""
+        mock_s3 = mocker.MagicMock()
+
+        # Mock paginated response
+        mock_s3.list_objects_v2.side_effect = [
+            {
+                "Contents": [{"Key": "file1"}, {"Key": "file2"}],
+                "IsTruncated": True,
+                "NextContinuationToken": "token123"
+            },
+            {
+                "Contents": [{"Key": "file3"}],
+                "IsTruncated": False
+            }
+        ]
+
+        result = ICDownloader.get_s3_objects(mock_s3, "test-bucket", "prefix/")
+
+        assert len(result) == 3
+        assert result[0]["Key"] == "file1"
+        assert result[1]["Key"] == "file2"
+        assert result[2]["Key"] == "file3"
+
+        # Verify pagination calls
+        assert mock_s3.list_objects_v2.call_count == 2
+
+    def test_get_s3_objects_no_contents(self, mocker):
+        """Test S3 objects retrieval with no contents."""
+        mock_s3 = mocker.MagicMock()
+        mock_s3.list_objects_v2.return_value = {"IsTruncated": False}
+
+        result = ICDownloader.get_s3_objects(mock_s3, "test-bucket", "prefix/")
+
+        assert len(result) == 0
+
+    def test_get_data_from_s3(self, mocker):
+        """Test downloading data from S3."""
+        mock_s3 = mocker.MagicMock()
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="s3",
+            local_directory=self.test_dir,
+            bucket_name="test-bucket"
+        )
+        downloader.s3 = mock_s3
+
+        file_list = ["path/to/file1.grib2", "path/to/file2.grib2"]
+
+        downloader.get_data_from_s3(file_list, self.test_dir)
+
+        # Verify download calls
+        assert mock_s3.download_file.call_count == 2
+        mock_s3.download_file.assert_any_call(
+            "test-bucket", "path/to/file1.grib2",
+            os.path.join(self.test_dir, "file1.grib2")
+        )
+        mock_s3.download_file.assert_any_call(
+            "test-bucket", "path/to/file2.grib2",
+            os.path.join(self.test_dir, "file2.grib2")
+        )
+
+    def test_get_data_from_s3_with_bucket_root(self, mocker):
+        """Test downloading data from S3 with bucket root directory."""
+        mock_s3 = mocker.MagicMock()
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="s3",
+            local_directory=self.test_dir,
+            bucket_name="test-bucket",
+            bucket_root_directory="root/dir"
+        )
+        downloader.s3 = mock_s3
+
+        file_list = ["file1.grib2"]
+
+        downloader.get_data_from_s3(file_list, self.test_dir)
+
+        mock_s3.download_file.assert_called_once_with(
+            "test-bucket", "root/dir/file1.grib2",
+            os.path.join(self.test_dir, "file1.grib2")
+        )
+
+    def test_get_data_from_s3_file_exists(self, mocker):
+        """Test S3 download when local file already exists."""
+        mock_s3 = mocker.MagicMock()
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="s3",
+            local_directory=self.test_dir,
+            bucket_name="test-bucket"
+        )
+        downloader.s3 = mock_s3
+
+        # Create existing file
+        existing_file = os.path.join(self.test_dir, "file1.grib2")
+        with open(existing_file, 'w') as f:
+            f.write("test")
+
+        file_list = ["path/to/file1.grib2"]
+
+        downloader.get_data_from_s3(file_list, self.test_dir)
+
+        # Should not attempt download
+        mock_s3.download_file.assert_not_called()
+
+    def test_get_data_from_s3_download_error(self, mocker):
+        """Test S3 download with error."""
+        mock_s3 = mocker.MagicMock()
+        mock_s3.download_file.side_effect = Exception("Download failed")
+
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="s3",
+            local_directory=self.test_dir,
+            bucket_name="test-bucket"
+        )
+        downloader.s3 = mock_s3
+
+        file_list = ["path/to/file1.grib2"]
+
+        # Should not raise exception, just log error
+        downloader.get_data_from_s3(file_list, self.test_dir)
+
+        mock_s3.download_file.assert_called_once()
+
+    def test_get_data_from_local(self):
+        """Test copying data from local directory."""
+        # Create source directory and files
+        source_dir = tempfile.mkdtemp()
+        try:
+            source_file1 = os.path.join(source_dir, "file1.grib2")
+            source_file2 = os.path.join(source_dir, "file2.grib2")
+
+            with open(source_file1, 'w') as f:
+                f.write("test data 1")
+            with open(source_file2, 'w') as f:
+                f.write("test data 2")
+
+            downloader = ICDownloader(
+                current_cycle=self.current_cycle,
+                download_source="local",
+                local_directory=self.test_dir,
+                root_directory=source_dir
+            )
+
+            file_list = ["file1.grib2", "file2.grib2"]
+
+            downloader.get_data_from_local(file_list, self.test_dir)
+
+            # Verify files were copied
+            assert os.path.exists(os.path.join(self.test_dir, "file1.grib2"))
+            assert os.path.exists(os.path.join(self.test_dir, "file2.grib2"))
+
+            # Verify content
+            with open(os.path.join(self.test_dir, "file1.grib2"), 'r') as f:
+                assert f.read() == "test data 1"
+
+        finally:
+            shutil.rmtree(source_dir)
+
+    def test_get_data_from_local_file_exists(self):
+        """Test local copy when destination file already exists."""
+        source_dir = tempfile.mkdtemp()
+        try:
+            source_file = os.path.join(source_dir, "file1.grib2")
+            with open(source_file, 'w') as f:
+                f.write("source data")
+
+            # Create existing destination file
+            dest_file = os.path.join(self.test_dir, "file1.grib2")
+            with open(dest_file, 'w') as f:
+                f.write("existing data")
+
+            downloader = ICDownloader(
+                current_cycle=self.current_cycle,
+                download_source="local",
+                local_directory=self.test_dir,
+                root_directory=source_dir
+            )
+
+            file_list = ["file1.grib2"]
+
+            downloader.get_data_from_local(file_list, self.test_dir)
+
+            # File should not be overwritten
+            with open(dest_file, 'r') as f:
+                assert f.read() == "existing data"
+
+        finally:
+            shutil.rmtree(source_dir)
+
+    def test_get_data_from_local_copy_error(self):
+        """Test local copy with error."""
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="local",
+            local_directory=self.test_dir,
+            root_directory="/nonexistent/path"
+        )
+
+        file_list = ["file1.grib2"]
+
+        with pytest.raises(OSError, match="Unable to copy"):
+            downloader.get_data_from_local(file_list, self.test_dir)
+
+    def test_get_data_s3(self, mocker):
+        """Test get_data method with S3 source."""
+        mock_s3 = mocker.MagicMock()
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="s3",
+            local_directory=self.test_dir,
+            bucket_name="test-bucket"
+        )
+        downloader.s3 = mock_s3
+
+        mock_get_s3 = mocker.patch.object(downloader, 'get_data_from_s3')
+        downloader.get_data()
+        mock_get_s3.assert_called_once_with(downloader.file_list, self.test_dir)
+
+    def test_get_data_local(self, mocker):
+        """Test get_data method with local source."""
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="local",
+            local_directory=self.test_dir,
+            root_directory="/test/root"
+        )
+
+        mock_get_local = mocker.patch.object(downloader, 'get_data_from_local')
+        downloader.get_data()
+        mock_get_local.assert_called_once_with(downloader.file_list, self.test_dir)
+
+    def test_get_data_invalid_source(self):
+        """Test get_data method with invalid source."""
+        downloader = ICDownloader(
+            current_cycle=self.current_cycle,
+            download_source="invalid",
+            local_directory=self.test_dir
+        )
+
+        with pytest.raises(KeyError):
+            downloader.get_data()
+
+
+@pytest.fixture
+def sample_datetime():
+    """Fixture providing a sample datetime for testing."""
+    return datetime(2023, 1, 1, 12, 0, 0)
+
+
+@pytest.fixture
+def temp_directory():
+    """Fixture providing a temporary directory for testing."""
+    temp_dir = tempfile.mkdtemp()
+    yield temp_dir
+    shutil.rmtree(temp_dir)
+
+
+class TestIntegration:
+    """Integration tests for ICDownloader."""
+
+    def test_full_workflow_local(self, sample_datetime, temp_directory, mocker):
+        """Test full workflow with local data source."""
+        # Create mock source data
+        source_dir = os.path.join(temp_directory, "source")
+        os.makedirs(source_dir, exist_ok=True)
+
+        # Create some mock files that match the expected patterns
+        test_files = [
+            "gfs.20230101/12/atmos/gfs.t12z.pgrb2.0p25.f000",
+            "gfs.20230101/12/atmos/gfs.t12z.pgrb2.0p25.f006",
+        ]
+
+        for file_path in test_files:
+            full_path = os.path.join(source_dir, file_path)
+            os.makedirs(os.path.dirname(full_path), exist_ok=True)
+            with open(full_path, 'w') as f:
+                f.write(f"mock data for {file_path}")
+
+        dest_dir = os.path.join(temp_directory, "dest")
+
+        downloader = ICDownloader(
+            current_cycle=sample_datetime,
+            num_levels=13,
+            member=None,
+            download_source="local",
+            local_directory=dest_dir,
+            root_directory=source_dir
+        )
+
+        # Mock the file lookup to return our test files
+        downloader.file_list = test_files
+
+        # This should work without errors
+        downloader.get_data()
+
+        # Verify destination directory was created
+        assert os.path.exists(dest_dir)
\ No newline at end of file

From e0dee8cca18380e573d482d1702037699f31021c Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 00:59:08 -0400
Subject: [PATCH 13/21] fix pynorms on tests

---
 tests/test_file_lookup.py   | 6 ++++--
 tests/test_ic_downloader.py | 7 ++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/tests/test_file_lookup.py b/tests/test_file_lookup.py
index 26e729e..1ad71b0 100644
--- a/tests/test_file_lookup.py
+++ b/tests/test_file_lookup.py
@@ -1,5 +1,7 @@
-import pytest
 from datetime import datetime, timedelta
+
+import pytest
+
 from mlglobal.ic_downloader import FileLookup
 
 
@@ -139,4 +141,4 @@ def test_file_info_delegation_gefs(self):
         gefs_result = lookup._gefs_file_info()
         delegated_result = lookup.get_file_info()
 
-        assert gefs_result == delegated_result
\ No newline at end of file
+        assert gefs_result == delegated_result
diff --git a/tests/test_ic_downloader.py b/tests/test_ic_downloader.py
index 1552e5a..3290ec5 100644
--- a/tests/test_ic_downloader.py
+++ b/tests/test_ic_downloader.py
@@ -1,8 +1,9 @@
-import pytest
 import os
-import tempfile
 import shutil
+import tempfile
 from datetime import datetime
+
+import pytest
 from botocore.exceptions import ClientError
 
 from mlglobal.ic_downloader import ICDownloader
@@ -425,4 +426,4 @@ def test_full_workflow_local(self, sample_datetime, temp_directory, mocker):
         downloader.get_data()
 
         # Verify destination directory was created
-        assert os.path.exists(dest_dir)
\ No newline at end of file
+        assert os.path.exists(dest_dir)

From 2fa5bc621b50a785e023b2e5c66262257db839a5 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:09:18 -0400
Subject: [PATCH 14/21] build g2c and cache it

---
 .github/workflows/pytests.yaml | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index 522ceb2..9ecf9a2 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -1,10 +1,39 @@
 name: pytests
 on: [push, pull_request]
 
+env:
+  g2c_VERSION: 2.2.0
+  G2C_DIR: ${{ github.workspace }}/g2c-${{ env.g2c_VERSION }}
+  G2C_STATIC: True
+  CC: gcc-14
+
 jobs:
+  cache-g2c:
+    runs-on: ubuntu-latest
+    name: Build and cache g2c library
+    steps:
+      - name: Cache g2c installation
+        id: cache-g2c
+        uses: actions/cache@v4
+        with:
+          path: |
+            g2c-${{ env.g2c_VERSION }}
+          key: g2c-${{ env.g2c_VERSION }}
+
+      - name: Checkout, build and install g2c
+        if: steps.cache-g2c.outputs.cache-hit != 'true'
+        run: |
+          wget https://github.com/NOAA-EMC/NCEPLIBS-g2c/archive/refs/tags/v${{ env.g2c_VERSION }}.tar.gz
+          tar -xzvf v${{ env.g2c_VERSION }}.tar.gz
+          cd NCEPLIBS-g2c-${{ env.g2c_VERSION }}
+          mkdir build && cd build
+          cmake -DUSE_Jasper=OFF -DUSE_OpenJPEG=ON -DBUILD_PNG=ON -DBUILD_AEC=ON -DBUILD_SHARED_LIBS=OFF -DCMAKE_INSTALL_PREFIX=../../g2c-${{ env.g2c_VERSION }} ..
+          make -j2
+          make install
+
   run_pytests:
     runs-on: ubuntu-latest
-    name: Install wxflow and run tests with pytests
+    name: Install and run tests with pytests
     strategy:
       matrix:
         python: ["3.11", "3.12", "3.13"]

From a20334947a77b8226a21a42dd9dd03ffe92a74a5 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:10:02 -0400
Subject: [PATCH 15/21] build g2c and cache it

---
 .github/workflows/pytests.yaml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index 9ecf9a2..c5fb851 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -36,7 +36,8 @@ jobs:
     name: Install and run tests with pytests
     strategy:
       matrix:
-        python: ["3.11", "3.12", "3.13"]
+        #python: ["3.11", "3.12", "3.13"]
+        python: ["3.12"]  # Temporarily only test with 3.12 for faster CI
 
     steps:
       - name: Setup Python

From c764a8cccd5505392ac21f09d932ada625351f4b Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:13:04 -0400
Subject: [PATCH 16/21] fix ga yaml

---
 .github/workflows/pytests.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index c5fb851..a8311d5 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -3,7 +3,7 @@ on: [push, pull_request]
 
 env:
   g2c_VERSION: 2.2.0
-  G2C_DIR: ${{ github.workspace }}/g2c-${{ env.g2c_VERSION }}
+  G2C_DIR: ${{ github.workspace }}/g2c-2.2.0
   G2C_STATIC: True
   CC: gcc-14
 

From ebcaaa8453197908e93391d561f3884119a30aa4 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:17:49 -0400
Subject: [PATCH 17/21] need g2c deps, apt-get them

---
 .github/workflows/pytests.yaml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index a8311d5..a1692ad 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -10,8 +10,15 @@ env:
 jobs:
   cache-g2c:
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: true
     name: Build and cache g2c library
     steps:
+      - name: apt-get and install dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libaec-dev libpng-dev zlib1g-dev libjpeg-dev libopenjp2-7-dev
+
       - name: Cache g2c installation
         id: cache-g2c
         uses: actions/cache@v4

From 953039af467280010a5d0ce07639f07c65a4442a Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:25:54 -0400
Subject: [PATCH 18/21] grib2io in the action needs a little more TLC

---
 .github/workflows/pytests.yaml | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index a1692ad..4df9b64 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -3,8 +3,6 @@ on: [push, pull_request]
 
 env:
   g2c_VERSION: 2.2.0
-  G2C_DIR: ${{ github.workspace }}/g2c-2.2.0
-  G2C_STATIC: True
   CC: gcc-14
 
 jobs:
@@ -34,7 +32,8 @@ jobs:
           tar -xzvf v${{ env.g2c_VERSION }}.tar.gz
           cd NCEPLIBS-g2c-${{ env.g2c_VERSION }}
           mkdir build && cd build
-          cmake -DUSE_Jasper=OFF -DUSE_OpenJPEG=ON -DBUILD_PNG=ON -DBUILD_AEC=ON -DBUILD_SHARED_LIBS=OFF -DCMAKE_INSTALL_PREFIX=../../g2c-${{ env.g2c_VERSION }} ..
+          cmake -DUSE_Jasper=OFF -DUSE_OpenJPEG=ON -DBUILD_PNG=ON -DBUILD_AEC=ON -DBUILD_SHARED_LIBS=OFF \
+                -DCMAKE_INSTALL_PREFIX=${{ github.workspace }}/g2c-${{ env.g2c_VERSION }} ..
           make -j2
           make install
 
@@ -61,6 +60,8 @@ jobs:
 
       - name: Install
         run: |
+          export G2C_STATIC=True
+          export G2C_DIR=${{ github.workspace }}/g2c-${{ env.g2c_VERSION }}
           cd $GITHUB_WORKSPACE
           pip install .[dev]
 

From 267373ff548f68226a069f940eddee9d8a9db7b3 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:27:48 -0400
Subject: [PATCH 19/21] duh, run_pytests needs g2c cache, else they both start
 together

---
 .github/workflows/pytests.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index 4df9b64..5351077 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -40,6 +40,7 @@ jobs:
   run_pytests:
     runs-on: ubuntu-latest
     name: Install and run tests with pytests
+    needs: cache-g2c
     strategy:
       matrix:
         #python: ["3.11", "3.12", "3.13"]

From df5d161136066547249c672d92deeee0d1163c24 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 09:35:07 -0400
Subject: [PATCH 20/21] G2C_DIR needs the path to libg2c.a

---
 .github/workflows/pytests.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml
index 5351077..f9cd6b9 100644
--- a/.github/workflows/pytests.yaml
+++ b/.github/workflows/pytests.yaml
@@ -62,7 +62,7 @@ jobs:
       - name: Install
         run: |
           export G2C_STATIC=True
-          export G2C_DIR=${{ github.workspace }}/g2c-${{ env.g2c_VERSION }}
+          export G2C_DIR=${{ github.workspace }}/g2c-${{ env.g2c_VERSION }}/lib
           cd $GITHUB_WORKSPACE
           pip install .[dev]
 

From 151c131bc90b9cbc59bbe8243905a3683dafa490 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Thu, 25 Sep 2025 15:59:17 -0400
Subject: [PATCH 21/21] update environment.yaml to include grib2io

---
 environment.yml | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/environment.yml b/environment.yml
index 07d06cf..f1805e9 100644
--- a/environment.yml
+++ b/environment.yml
@@ -1,4 +1,4 @@
-name: graphcast
+name: condaenv
 channels:
   - conda-forge
 dependencies:
@@ -9,7 +9,8 @@ dependencies:
   - cartopy
   - jupyterlab
   - boto3
+  - nceplibs-g2c
+  - grib2io
   - pip
   - pip:
-    - https://github.com/deepmind/graphcast/archive/master.zip
-    - flax
+    - git+https://github.com/noaa-emc/graphcast@aea0678cfbd7e866e4a1d364b456861d0a03954b