<commit_before>"""
Grid and Forcing for LADIM for the Regional Ocean Model System (ROMS)

"""

# -----------------------------------
# <NAME>, <<EMAIL>>
# Institute of Marine Research
# Bergen, Norway
# 2017-03-01
# -----------------------------------

# import sys
import glob
import logging
import numpy as np
from netCDF4 import Dataset, num2date
from ladim.sample import sample2D


class Grid:

    """Simple ROMS grid object

    Possible grid arguments:
      subgrid = [i0, i1, j0, j1]
        Ordinary python style, start points included, not end points
        Each of the elements can be replaced with None, for no limitation
      Vinfo: dictionary with N, hc, theta_s and theta_b

    """

    # Lagrer en del unødige attributter

    def __init__(self, config):

        logging.info("Initializing zROMS grid object")
        try:
            ncid = Dataset(config.grid_file)
        except OSError:
            logging.error('Grid file {} not found'.format(config.grid_file))
            raise SystemExit(1)

        # Subgrid, only considers internal grid cells
        # 1 <= i0 < i1 <= imax-1, default=end points
        # 1 <= j0 < j1 <= jmax-1, default=end points
        # Here, imax, jmax refers to whole grid
        jmax, imax = ncid.variables['h'].shape
        whole_grid = [1, imax-1, 1, jmax-1]
        if 'subgrid' in config.grid_args:
            limits = list(config.grid_args['subgrid'])
        else:
            limits = whole_grid
        # Allow None if no imposed limitation
        for ind, val in enumerate(limits):
            if val is None:
                limits[ind] = whole_grid[ind]

        self.i0, self.i1, self.j0, self.j1 = limits
        self.imax = self.i1 - self.i0
        self.jmax = self.j1 - self.j0
        # print('Grid : imax, jmax, size = ',
        #      self.imax, self.jmax, self.imax*self.jmax)

        # Slices
        #   rho-points
        self.I = slice(self.i0, self.i1)
        self.J = slice(self.j0, self.j1)
        #   U and V-points
        self.Iu = slice(self.i0-1, self.i1)
        self.Ju = self.J
        self.Iv = self.I
        self.Jv = slice(self.j0-1, self.j1)

        # Vertical z-levels
        # self.z_levels = ncid.variables['zlevels'][:]
        self.z_levels = np.array([0, 1, 2, 3, 5, 10, 15, 20])

        # Read some variables
        self.H = ncid.variables['h'][self.J, self.I]
        self.M = ncid.variables['mask_rho'][self.J, self.I].astype(int)
        # self.Mu = ncid.variables['mask_u'][self.Ju, self.Iu]
        # self.Mv = ncid.variables['mask_v'][self.Jv, self.Iv]
        self.dx = 1./ncid.variables['pm'][self.J, self.I]
        self.dy = 1./ncid.variables['pn'][self.J, self.I]
        self.lon = ncid.variables['lon_rho'][self.J, self.I]
        self.lat = ncid.variables['lat_rho'][self.J, self.I]

        # self.z_r = sdepth(self.H, self.hc, self.Cs_r,
        #                  stagger='rho', Vtransform=self.Vtransform)
        # self.z_w = sdepth(self.H, self.hc, self.Cs_w,
        #                  stagger='w', Vtransform=self.Vtransform)

        # Land masks at u- and v-points
        M = self.M
        Mu = np.zeros((self.jmax, self.imax+1), dtype=int)
        Mu[:, 1:-1] = M[:, :-1] * M[:, 1:]
        Mu[:, 0] = M[:, 0]
        Mu[:, -1] = M[:, -1]
        self.Mu = Mu

        Mv = np.zeros((self.jmax+1, self.imax), dtype=int)
        Mv[1:-1, :] = M[:-1, :] * M[1:, :]
        Mv[0, :] = M[0, :]
        Mv[-1, :] = M[-1, :]
        self.Mv = Mv

        # Close the file(s)
        ncid.close()

    def sample_metric(self, X, Y):
        """Sample the metric coefficients

        Changes slowly, so using nearest neighbour
        """
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0

        # Metric is conform for PolarStereographic
        A = self.dx[J, I]
        return A, A

    def sample_depth(self, X, Y):
        """Return the depth of grid cells"""
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0
        return self.H[J, I]

    def lonlat(self, X, Y):
        """Return the longitude and latitude from grid coordinates"""
        return (sample2D(self.lon, X-self.i0, Y-self.j0),
                sample2D(self.lat, X-self.i0, Y-self.j0))

    def ingrid(self, X, Y):
        """Returns True for points inside the subgrid"""
        # return ((self.i0-0.5 <= X) & (X <= self.i1-0.5) &
        #         (self.j0-0.5 <= Y) & (Y <= self.j1-0.5))
        return ((self.i0 < X) & (X < self.i1-1) &
                (self.j0 < Y) & (Y < self.j1-1))

    def onland(self, X, Y):
        """Returns True for points on land"""
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0
        return self.M[J, I] < 1

    # Error if point outside
    def atsea(self, X, Y):
        """Returns True for points at sea"""
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0
        return self.M[J, I] > 0

# -----------------------------------------------
# The Forcing class from the old forcing module
# -----------------------------------------------


class Forcing:
    """
    Class for ROMS forcing

    """

    def __init__(self, config, grid):

        logging.basicConfig(level=logging.INFO)
        logging.info("Initiating forcing")

        self._grid = grid  # Get the grid object, make private?

        self.ibm_forcing = config.ibm_forcing

        # Test for glob, use MFDataset if needed
        files = glob.glob(config.input_file)
        numfiles = len(files)
        if numfiles == 0:
            logging.error("No input file: {}".format(config.input_file))
            raise SystemExit(3)
        logging.info('Number of forcing files = {}'.format(numfiles))
<commit_msg>Fixed a bug in zROMS.py<commit_after>"""
Grid and Forcing for LADIM for the Regional Ocean Model System (ROMS)

"""

# -----------------------------------
# <NAME>, <<EMAIL>>
# Institute of Marine Research
# Bergen, Norway
# 2017-03-01
# -----------------------------------

# import sys
import glob
import logging
import numpy as np
from netCDF4 import Dataset, num2date
from ladim.sample import sample2D


class Grid:

    """Simple ROMS grid object

    Possible grid arguments:
      subgrid = [i0, i1, j0, j1]
        Ordinary python style, start points included, not end points
        Each of the elements can be replaced with None, for no limitation
      Vinfo: dictionary with N, hc, theta_s and theta_b

    """

    # Lagrer en del unødige attributter

    def __init__(self, config):

        logging.info("Initializing zROMS grid object")
        try:
            ncid = Dataset(config.grid_file)
        except OSError:
            logging.error('Grid file {} not found'.format(config.grid_file))
            raise SystemExit(1)

        # Subgrid, only considers internal grid cells
        # 1 <= i0 < i1 <= imax-1, default=end points
        # 1 <= j0 < j1 <= jmax-1, default=end points
        # Here, imax, jmax refers to whole grid
        jmax, imax = ncid.variables['h'].shape
        whole_grid = [1, imax-1, 1, jmax-1]
        if 'subgrid' in config.grid_args:
            limits = list(config.grid_args['subgrid'])
        else:
            limits = whole_grid
        # Allow None if no imposed limitation
        for ind, val in enumerate(limits):
            if val is None:
                limits[ind] = whole_grid[ind]

        self.i0, self.i1, self.j0, self.j1 = limits
        self.imax = self.i1 - self.i0
        self.jmax = self.j1 - self.j0
        # print('Grid : imax, jmax, size = ',
        #      self.imax, self.jmax, self.imax*self.jmax)

        # Slices
        #   rho-points
        self.I = slice(self.i0, self.i1)
        self.J = slice(self.j0, self.j1)
        #   U and V-points
        self.Iu = slice(self.i0-1, self.i1)
        self.Ju = self.J
        self.Iv = self.I
        self.Jv = slice(self.j0-1, self.j1)

        # Vertical z-levels
        # self.z_levels = ncid.variables['zlevels'][:]
        self.z_levels = np.array([0, 1, 2, 3, 5, 10, 15, 20])

        # Read some variables
        self.H = ncid.variables['h'][self.J, self.I]
        self.M = ncid.variables['mask_rho'][self.J, self.I].astype(int)
        # self.Mu = ncid.variables['mask_u'][self.Ju, self.Iu]
        # self.Mv = ncid.variables['mask_v'][self.Jv, self.Iv]
        self.dx = 1./ncid.variables['pm'][self.J, self.I]
        self.dy = 1./ncid.variables['pn'][self.J, self.I]
        self.lon = ncid.variables['lon_rho'][self.J, self.I]
        self.lat = ncid.variables['lat_rho'][self.J, self.I]

        # self.z_r = sdepth(self.H, self.hc, self.Cs_r,
        #                  stagger='rho', Vtransform=self.Vtransform)
        # self.z_w = sdepth(self.H, self.hc, self.Cs_w,
        #                  stagger='w', Vtransform=self.Vtransform)

        # Land masks at u- and v-points
        M = self.M
        Mu = np.zeros((self.jmax, self.imax+1), dtype=int)
        Mu[:, 1:-1] = M[:, :-1] * M[:, 1:]
        Mu[:, 0] = M[:, 0]
        Mu[:, -1] = M[:, -1]
        self.Mu = Mu

        Mv = np.zeros((self.jmax+1, self.imax), dtype=int)
        Mv[1:-1, :] = M[:-1, :] * M[1:, :]
        Mv[0, :] = M[0, :]
        Mv[-1, :] = M[-1, :]
        self.Mv = Mv

        # Close the file(s)
        ncid.close()

    def sample_metric(self, X, Y):
        """Sample the metric coefficients

        Changes slowly, so using nearest neighbour
        """
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0

        # Metric is conform for PolarStereographic
        A = self.dx[J, I]
        return A, A

    def sample_depth(self, X, Y):
        """Return the depth of grid cells"""
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0
        return self.H[J, I]

    def lonlat(self, X, Y):
        """Return the longitude and latitude from grid coordinates"""
        return (sample2D(self.lon, X-self.i0, Y-self.j0),
                sample2D(self.lat, X-self.i0, Y-self.j0))

    def ingrid(self, X, Y):
        """Returns True for points inside the subgrid"""
        # return ((self.i0-0.5 <= X) & (X <= self.i1-0.5) &
        #         (self.j0-0.5 <= Y) & (Y <= self.j1-0.5))
        return ((self.i0 < X) & (X < self.i1-1) &
                (self.j0 < Y) & (Y < self.j1-1))

    def onland(self, X, Y):
        """Returns True for points on land"""
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0
        return self.M[J, I] < 1

    # Error if point outside
    def atsea(self, X, Y):
        """Returns True for points at sea"""
        I = X.round().astype(int) - self.i0
        J = Y.round().astype(int) - self.j0
        return self.M[J, I] > 0

# -----------------------------------------------
# The Forcing class from the old forcing module
# -----------------------------------------------


class Forcing:
    """
    Class for ROMS forcing

    """

    def __init__(self, config, grid):

        logging.basicConfig(level=logging.INFO)
        logging.info("Initiating forcing")

        self._grid = grid  # Get the grid object, make private?

        self.ibm_forcing = config.ibm_forcing

        # Test for glob, use MFDataset if needed
        files = glob.glob(config.input_file)
        files.sort()
        numfiles = len(files)
        if numfiles == 0:
            logging.error("No input file: {}".format(config.input_file))
            raise SystemExit(3)
        logging.info('Number of forcing files = {}'.format(numfiles))

{'fraction_non_alphanumeric': 0.16943893820563646, 'fraction_numerical': 0.019477721278979575, 'mean_word_length': 1.789423076923077, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3592041', 'n_tokens_mistral': 4057, 'n_tokens_neox': 3934, 'n_words': 1005}