Skip to content

ufs

UFS-AQM Reader

UFSReader

Bases: GriddedReader

Reader for UFS-AQM model output files.

Source code in monetio/readers/ufs.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@register_reader("ufs")
class UFSReader(GriddedReader):
    """
    Reader for UFS-AQM model output files.
    """

    def open_dataset(
        self,
        files: str | list[str],
        convert_to_ppb: bool = True,
        mech: str = "cb6r3_ae6_aq",
        var_list: list[str] | None = None,
        fname_pm25: str | list[str] | None = None,
        surf_only: bool = False,
        **kwargs: Any,
    ) -> xr.Dataset:
        """
        Reads UFS-AQM netCDF files.

        Parameters
        ----------
        files : Union[str, List[str]]
            File path, list of paths, or glob pattern.
        convert_to_ppb : bool, optional
            Convert gas species from ppmV to ppbV, by default True.
        mech : str, optional
            Mechanism name for species sums, by default "cb6r3_ae6_aq".
        var_list : List[str], optional
            List of variables to include, by default None.
        fname_pm25 : Union[str, List[str]], optional
            Optional separate PM2.5 files to merge, by default None.
        surf_only : bool, optional
            Whether to only keep surface data, by default False.
        **kwargs : Any
            Additional arguments passed to the driver.

        Returns
        -------
        xr.Dataset
            The processed UFS-AQM dataset.

        Examples
        --------
        >>> reader = UFSReader()
        >>> ds = reader.open_dataset("aqm.t12z.dyn.f*.nc", surf_only=True)
        """
        # Prepare kwargs
        if "concat_dim" not in kwargs:
            kwargs["concat_dim"] = "time"
        if "combine" not in kwargs:
            kwargs["combine"] = "nested"

        # Open dataset
        ds = self.driver.open(files, **kwargs)

        # Merge PM25 file if present
        if fname_pm25 is not None:
            ds_pm25 = self.driver.open(fname_pm25, **kwargs)
            ds_pm25 = ds_pm25.drop_vars(["lat", "lon", "pfull"], errors="ignore")
            ds_pm25.attrs = {}
            from monetio.util import _try_merge_exact

            ds = _try_merge_exact(ds, ds_pm25, right_name="PM2.5")

        # Standardize
        rename_dict = {
            "grid_yt": "y",
            "grid_xt": "x",
            "pfull": "z",
            "phalf": "z_i",
            "lon": "longitude",
            "lat": "latitude",
            "tmp": "temperature_k",
            "pressfc": "surfpres_pa",
            "dpres": "dp_pa",
            "hgtsfc": "surfalt_m",
            "delz": "dz_m",
        }
        # Only rename what exists
        actual_rename = {k: v for k, v in rename_dict.items() if k in ds.variables or k in ds.dims}
        if actual_rename:
            ds = ds.rename(actual_rename)

        # Calculations
        if "surfpres_pa" in ds and "ak" in ds and "bk" in ds:
            ds["pres_pa_mid"] = _calc_pressure(ds)

        # Resort z (Lazy)
        if "z" in ds.coords:
            if ds.z.size > 1:
                is_ascending = ds.z[0] < ds.z[-1]
                if is_ascending:
                    ds = ds.isel(z=slice(None, None, -1))
                    if "dz_m" in ds:
                        ds["dz_m"] = ds["dz_m"] * -1.0
        if "z_i" in ds.coords:
            if ds.z_i.size > 1:
                is_ascending = ds.z_i[0] < ds.z_i[-1]
                if is_ascending:
                    ds = ds.isel(z_i=slice(None, None, -1))

        if not surf_only and "dz_m" in ds and "surfalt_m" in ds:
            ds["alt_msl_m_full"] = _calc_hgt(ds)

        if "latitude" in ds.data_vars and "time" in ds["latitude"].dims:
            ds["latitude"] = ds["latitude"].isel(time=0)
        if "longitude" in ds.data_vars and "time" in ds["longitude"].dims:
            ds["longitude"] = ds["longitude"].isel(time=0)

        coords = [c for c in ["latitude", "longitude", "time"] if c in ds.variables]
        ds = ds.set_coords(coords)

        if surf_only and "z" in ds.dims:
            ds = ds.isel(z=0).expand_dims("z", axis=1)

        # Time fix (Avoid eager .indexes)
        if "time" in ds.coords:
            # Check for cftime lazily if possible, or use .indexes only if necessary
            # For now, following standard conventions to avoid eager indexes if we can
            if ds.indexes["time"].__class__.__name__ == "CFTimeIndex":
                ds["time"] = ds.indexes["time"].to_datetimeindex()

        # Unit conversion (Lazy)
        if convert_to_ppb:
            ds = _convert_to_ppb(ds)

        ds = _convert_ugkg_to_ugm3(ds)

        # Add lazy diagnostic variables
        for name, spec in DIAGNOSTICS.items():
            ds = add_lazy_diagnostic(ds, name, spec)

        # Subset if var_list
        if var_list is not None:
            # We must keep coordinates and some essentials
            essentials = [
                "latitude",
                "longitude",
                "time",
                "z",
                "z_i",
                "z_stag",
                "z_soil",
                "pres_pa_mid",
                "temperature_k",
            ]
            to_keep = set(var_list) | set(essentials)
            # Add those that were added as diagnostics
            to_keep |= {name for name in DIAGNOSTICS if name in ds.variables}
            available = [v for v in ds.variables if v in to_keep]
            ds = ds[available]

        # Scientific Hygiene
        ds = _scientific_hygiene(ds)

        # Update history
        ds = update_history(ds, "Read UFS-AQM data.")

        return ds

open_dataset(files, convert_to_ppb=True, mech='cb6r3_ae6_aq', var_list=None, fname_pm25=None, surf_only=False, **kwargs)

Reads UFS-AQM netCDF files.

Parameters:

Name Type Description Default
files Union[str, List[str]]

File path, list of paths, or glob pattern.

 required
convert_to_ppb bool

Convert gas species from ppmV to ppbV, by default True.

 True
mech str

Mechanism name for species sums, by default "cb6r3_ae6_aq".

 'cb6r3_ae6_aq'
var_list List[str]

List of variables to include, by default None.

 None
fname_pm25 Union[str, List[str]]

Optional separate PM2.5 files to merge, by default None.

 None
surf_only bool

Whether to only keep surface data, by default False.

 False
**kwargs Any

Additional arguments passed to the driver.

 {}

Returns:

Type Description
Dataset

The processed UFS-AQM dataset.

Examples:

>>> reader = UFSReader()
>>> ds = reader.open_dataset("aqm.t12z.dyn.f*.nc", surf_only=True)
Source code in monetio/readers/ufs.py 
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def open_dataset(
    self,
    files: str | list[str],
    convert_to_ppb: bool = True,
    mech: str = "cb6r3_ae6_aq",
    var_list: list[str] | None = None,
    fname_pm25: str | list[str] | None = None,
    surf_only: bool = False,
    **kwargs: Any,
) -> xr.Dataset:
    """
    Reads UFS-AQM netCDF files.

    Parameters
    ----------
    files : Union[str, List[str]]
        File path, list of paths, or glob pattern.
    convert_to_ppb : bool, optional
        Convert gas species from ppmV to ppbV, by default True.
    mech : str, optional
        Mechanism name for species sums, by default "cb6r3_ae6_aq".
    var_list : List[str], optional
        List of variables to include, by default None.
    fname_pm25 : Union[str, List[str]], optional
        Optional separate PM2.5 files to merge, by default None.
    surf_only : bool, optional
        Whether to only keep surface data, by default False.
    **kwargs : Any
        Additional arguments passed to the driver.

    Returns
    -------
    xr.Dataset
        The processed UFS-AQM dataset.

    Examples
    --------
    >>> reader = UFSReader()
    >>> ds = reader.open_dataset("aqm.t12z.dyn.f*.nc", surf_only=True)
    """
    # Prepare kwargs
    if "concat_dim" not in kwargs:
        kwargs["concat_dim"] = "time"
    if "combine" not in kwargs:
        kwargs["combine"] = "nested"

    # Open dataset
    ds = self.driver.open(files, **kwargs)

    # Merge PM25 file if present
    if fname_pm25 is not None:
        ds_pm25 = self.driver.open(fname_pm25, **kwargs)
        ds_pm25 = ds_pm25.drop_vars(["lat", "lon", "pfull"], errors="ignore")
        ds_pm25.attrs = {}
        from monetio.util import _try_merge_exact

        ds = _try_merge_exact(ds, ds_pm25, right_name="PM2.5")

    # Standardize
    rename_dict = {
        "grid_yt": "y",
        "grid_xt": "x",
        "pfull": "z",
        "phalf": "z_i",
        "lon": "longitude",
        "lat": "latitude",
        "tmp": "temperature_k",
        "pressfc": "surfpres_pa",
        "dpres": "dp_pa",
        "hgtsfc": "surfalt_m",
        "delz": "dz_m",
    }
    # Only rename what exists
    actual_rename = {k: v for k, v in rename_dict.items() if k in ds.variables or k in ds.dims}
    if actual_rename:
        ds = ds.rename(actual_rename)

    # Calculations
    if "surfpres_pa" in ds and "ak" in ds and "bk" in ds:
        ds["pres_pa_mid"] = _calc_pressure(ds)

    # Resort z (Lazy)
    if "z" in ds.coords:
        if ds.z.size > 1:
            is_ascending = ds.z[0] < ds.z[-1]
            if is_ascending:
                ds = ds.isel(z=slice(None, None, -1))
                if "dz_m" in ds:
                    ds["dz_m"] = ds["dz_m"] * -1.0
    if "z_i" in ds.coords:
        if ds.z_i.size > 1:
            is_ascending = ds.z_i[0] < ds.z_i[-1]
            if is_ascending:
                ds = ds.isel(z_i=slice(None, None, -1))

    if not surf_only and "dz_m" in ds and "surfalt_m" in ds:
        ds["alt_msl_m_full"] = _calc_hgt(ds)

    if "latitude" in ds.data_vars and "time" in ds["latitude"].dims:
        ds["latitude"] = ds["latitude"].isel(time=0)
    if "longitude" in ds.data_vars and "time" in ds["longitude"].dims:
        ds["longitude"] = ds["longitude"].isel(time=0)

    coords = [c for c in ["latitude", "longitude", "time"] if c in ds.variables]
    ds = ds.set_coords(coords)

    if surf_only and "z" in ds.dims:
        ds = ds.isel(z=0).expand_dims("z", axis=1)

    # Time fix (Avoid eager .indexes)
    if "time" in ds.coords:
        # Check for cftime lazily if possible, or use .indexes only if necessary
        # For now, following standard conventions to avoid eager indexes if we can
        if ds.indexes["time"].__class__.__name__ == "CFTimeIndex":
            ds["time"] = ds.indexes["time"].to_datetimeindex()

    # Unit conversion (Lazy)
    if convert_to_ppb:
        ds = _convert_to_ppb(ds)

    ds = _convert_ugkg_to_ugm3(ds)

    # Add lazy diagnostic variables
    for name, spec in DIAGNOSTICS.items():
        ds = add_lazy_diagnostic(ds, name, spec)

    # Subset if var_list
    if var_list is not None:
        # We must keep coordinates and some essentials
        essentials = [
            "latitude",
            "longitude",
            "time",
            "z",
            "z_i",
            "z_stag",
            "z_soil",
            "pres_pa_mid",
            "temperature_k",
        ]
        to_keep = set(var_list) | set(essentials)
        # Add those that were added as diagnostics
        to_keep |= {name for name in DIAGNOSTICS if name in ds.variables}
        available = [v for v in ds.variables if v in to_keep]
        ds = ds[available]

    # Scientific Hygiene
    ds = _scientific_hygiene(ds)

    # Update history
    ds = update_history(ds, "Read UFS-AQM data.")

    return ds