Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions ndsl/filesystem.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,37 @@
import warnings

import fsspec


def get_fs(path: str) -> fsspec.AbstractFileSystem:
"""Return the fsspec filesystem required to handle a given path."""
warnings.warn(
"Usage of `get_fs()` is discouraged if favor `os.path` and `pathlib` "
"modules. The function will be removed in the next version of NDSL.",
DeprecationWarning,
stacklevel=2,
)
fs, _, _ = fsspec.get_fs_token_paths(path)
return fs


def is_file(filename):
warnings.warn(
"Usage of `is_file()` is discouraged if favor of plain `os.path.isfile()`. "
"The function will be removed in the next version of NDSL.",
DeprecationWarning,
stacklevel=2,
)
return get_fs(filename).isfile(filename)


def open(filename, *args, **kwargs):
warnings.warn(
"Usage of `open()` is discouraged if favor the python built-in file "
"open context manager. The function will be removed in the next version "
"of NDSL.",
DeprecationWarning,
stacklevel=2,
)
fs = get_fs(filename)
return fs.open(filename, *args, **kwargs)
23 changes: 9 additions & 14 deletions ndsl/grid/helper.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import dataclasses
import os
import pathlib

import xarray as xr
Expand All @@ -15,7 +16,6 @@
import ndsl.constants as constants
from ndsl.constants import Z_DIM, Z_INTERFACE_DIM
from ndsl.dsl.typing import Float
from ndsl.filesystem import get_fs
from ndsl.grid.generation import MetricTerms
from ndsl.initialization.allocator import QuantityFactory
from ndsl.quantity import Quantity
Expand Down Expand Up @@ -130,15 +130,15 @@ class VerticalGridData:
"""
Terms defining the vertical grid.

Eulerian vertical grid is defined by p = ak + bk * p_ref
Eulerian vertical grid is defined by p = ak + bk * p_ref.
"""

# TODO: make these non-optional, make FloatFieldK a true type and use it
ak: Quantity
bk: Quantity
"""
reference pressure (Pa) used to define pressure at vertical interfaces,
where p = ak + bk * p_ref
Reference pressure (Pa) used to define pressure at vertical interfaces,
where p = ak + bk * p_ref.
"""

def __post_init__(self):
Expand All @@ -155,22 +155,21 @@ def new_from_metric_terms(cls, metric_terms: MetricTerms) -> VerticalGridData:

@classmethod
def from_restart(cls, restart_path: str, quantity_factory: QuantityFactory):
fs = get_fs(restart_path)
restart_files = fs.ls(restart_path)
restart_files = os.listdir(restart_path)
data_file = restart_files[
[fname.endswith("fv_core.res.nc") for fname in restart_files].index(True)
]

ak_bk_data_file = pathlib.Path(restart_path) / data_file
if not fs.isfile(ak_bk_data_file):
if not ak_bk_data_file.is_file():
raise ValueError(
"""vertical_grid_from_restart is true,
but no fv_core.res.nc in restart data file."""
)

ak = quantity_factory.zeros([Z_INTERFACE_DIM], units="Pa")
bk = quantity_factory.zeros([Z_INTERFACE_DIM], units="")
with fs.open(ak_bk_data_file, "rb") as f:
with open(ak_bk_data_file, "rb") as f:
ds = xr.open_dataset(f).isel(Time=0).drop_vars("Time")
ak.view[:] = ds["ak"].values
bk.view[:] = ds["bk"].values
Expand All @@ -179,9 +178,7 @@ def from_restart(cls, restart_path: str, quantity_factory: QuantityFactory):

@property
def p_ref(self) -> float:
"""
reference pressure (Pa)
"""
"""Reference pressure (Pa)"""
return 1e5

@property
Expand Down Expand Up @@ -230,9 +227,7 @@ def dp(self) -> Quantity:

@property
def ptop(self) -> Float:
"""
top of atmosphere pressure (Pa)
"""
"""Top of atmosphere pressure (Pa)"""
if self.bk.view[0] != 0:
raise ValueError("ptop is not well-defined when top-of-atmosphere bk != 0")
if is_gpu_backend(self.ak.gt4py_backend):
Expand Down
5 changes: 2 additions & 3 deletions ndsl/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import cftime
import xarray as xr

import ndsl.filesystem as filesystem
from ndsl.quantity import Quantity


Expand Down Expand Up @@ -39,7 +38,7 @@ def write_state(state: dict, filename: str) -> None:
if "time" not in state:
raise ValueError('state must include a value for "time"')
ds = to_xarray_dataset(state)
with filesystem.open(filename, "wb") as f:
with open(filename, "wb") as f:
ds.to_netcdf(f)


Expand Down Expand Up @@ -68,7 +67,7 @@ def read_state(filename: str) -> dict:
state: a model state dictionary
"""
out_dict = {}
with filesystem.open(filename, "rb") as f:
with open(filename, "rb") as f:
time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
ds = xr.open_dataset(f, decode_times=time_coder)
for name, value in ds.data_vars.items():
Expand Down
24 changes: 8 additions & 16 deletions ndsl/monitor/netcdf_monitor.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,12 @@
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Set
from warnings import warn

import fsspec
import numpy as np
import xarray as xr

from ndsl.comm.communicator import Communicator
from ndsl.dsl.typing import Float, get_precision
from ndsl.filesystem import get_fs
from ndsl.logging import ndsl_log
from ndsl.monitor.convert import to_numpy
from ndsl.quantity import Quantity
Expand Down Expand Up @@ -46,17 +43,14 @@ def data(self) -> Quantity:
class _ChunkedNetCDFWriter:
FILENAME_FORMAT = "state_{chunk:04d}_tile{tile}.nc"

def __init__(
self, path: str, tile: int, fs: fsspec.AbstractFileSystem, time_chunk_size: int
):
def __init__(self, path: str, tile: int, time_chunk_size: int) -> None:
self._path = path
self._tile = tile
self._fs = fs
self._time_chunk_size = time_chunk_size
self._i_time = 0
self._chunked: Optional[Dict[str, _TimeChunkedVariable]] = None
self._times: List[Any] = []
self._time_units: Optional[str] = None
self._chunked: dict[str, _TimeChunkedVariable] | None = None
self._times: list = []
self._time_units: str | None = None

def append(self, state):
ndsl_log.debug("appending at time %d", self._i_time)
Expand Down Expand Up @@ -128,11 +122,10 @@ def __init__(
rank = communicator.rank
self._tile_index = communicator.partitioner.tile_index(rank)
self._path = path
self._fs = get_fs(path)
self._communicator = communicator
self._time_chunk_size = time_chunk_size
self.__writer: Optional[_ChunkedNetCDFWriter] = None
self._expected_vars: Optional[Set[str]] = None
self.__writer: _ChunkedNetCDFWriter | None = None
self._expected_vars: set[str] | None = None
self._transfer_type = precision
if self._transfer_type == np.float32 and get_precision() > 32:
warn(
Expand All @@ -145,7 +138,6 @@ def _writer(self):
self.__writer = _ChunkedNetCDFWriter(
path=self._path,
tile=self._tile_index,
fs=self._fs,
time_chunk_size=self._time_chunk_size,
)
return self.__writer
Expand Down Expand Up @@ -178,7 +170,7 @@ def store(self, state: dict) -> None:
if state is not None: # we are on root rank
self._writer.append(state)

def store_constant(self, state: Dict[str, Quantity]) -> None:
def store_constant(self, state: dict[str, Quantity]) -> None:
state = self._communicator.gather_state(
state, transfer_type=self._transfer_type
)
Expand All @@ -189,7 +181,7 @@ def store_constant(self, state: Dict[str, Quantity]) -> None:
for name, quantity in state.items():
path_for_grid = constants_filename + "_" + name + ".nc"

if self._fs.exists(path_for_grid):
if os.path.exists(path_for_grid):
ds = xr.open_dataset(path_for_grid)
ds = ds.load()
ds[name] = xr.DataArray(
Expand Down
11 changes: 5 additions & 6 deletions ndsl/restart/_legacy_restart.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import xarray as xr

import ndsl.constants as constants
import ndsl.filesystem as filesystem
import ndsl.io as io
from ndsl.comm.communicator import Communicator
from ndsl.comm.partitioner import get_tile_index
Expand Down Expand Up @@ -54,16 +53,16 @@ def open_restart(
raise ValueError("no restart files found at {}".format(dirname))

for filename in filenames:
with filesystem.open(filename, "rb") as file:
with open(filename, "rb") as file:
state.update(
load_partial_state_from_restart_file(
file, restart_properties, only_names=only_names
)
)
coupler_res_filename = get_coupler_res_filename(dirname, label)
if filesystem.is_file(coupler_res_filename):
if os.path.isfile(coupler_res_filename):
if only_names is None or "time" in only_names:
with filesystem.open(coupler_res_filename, "r") as f:
with open(coupler_res_filename, "r") as f:
state["time"] = io.get_current_date_from_coupler_res(f)
if to_state is None:
state = communicator.tile.scatter_state(state)
Expand All @@ -78,7 +77,7 @@ def get_coupler_res_filename(dirname, label):

def restart_files(dirname, tile_index, label) -> Generator[BinaryIO, None, None]:
for filename in restart_filenames(dirname, tile_index, label):
with filesystem.open(filename, "rb") as f:
with open(filename, "rb") as f:
yield f


Expand All @@ -89,7 +88,7 @@ def restart_filenames(dirname, tile_index, label):
filename = os.path.join(dirname, prepend_label(name, label) + suffix)
if (
(name in RESTART_NAMES)
or filesystem.is_file(filename)
or os.path.isfile(filename)
or os.path.exists(filename)
):
return_list.append(filename)
Expand Down
19 changes: 19 additions & 0 deletions tests/test_filesystem.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
import pytest

import ndsl.filesystem as fs


def test_is_file_is_deprecated() -> None:
with pytest.deprecated_call():
fs.is_file("path/to/my_file.txt")


def test_open_is_deprecated() -> None:
with pytest.deprecated_call():
with fs.open("README.md", "r"):
pass


def test_get_fs_is_deprecated() -> None:
with pytest.deprecated_call():
fs.get_fs("path/to/my/file.txt")
Loading