<commit_before>from netCDF4 import Dataset
import scipy.interpolate
import scipy.spatial

from climate.pyom import PyOM, pyom_method, cyclic

class NorthAtlanticHighResolution(PyOM):
    """ Based on
    http://journals.ametsoc.org/doi/10.1175/1520-0485%282000%29030%3C1532%3ANSOTNA%3E2.0.CO%3B2
    """
    interpolator_initialized = False

    @pyom_method
    def set_parameter(self):
        """set main parameter
        """
        self.nx, self.ny, self.nz = (150,160,45)#(87,89,45)
        self.x_origin = -98.
        self.y_origin = -20.
        self._x_boundary = 17.2
        self._y_boundary = 72.6

        self.dt_mom = 3600.0
        self.dt_tracer = 3600.0

        self.runlen = 86400 * 365.0
        self.enable_diag_snapshots = True

        self.coord_degree = True

        self.congr_epsilon = 1e-10
        self.congr_max_iterations = 20000
        self.enable_streamfunction = True

        self.enable_neutral_diffusion = True
        self.enable_skew_diffusion = True
        self.K_iso_0 = 1000.0
        self.K_iso_steep = 200.0
        self.iso_dslope = 1./1000.0
        self.iso_slopec = 4./1000.0

        self.enable_hor_friction = True
        self.A_h = 5e4
        self.enable_hor_friction_cos_scaling = True
        self.hor_friction_cosPower = 3
        self.enable_tempsalt_sources = True

        self.enable_implicit_vert_friction = True
        self.enable_tke = False
        self.c_k = 0.1
        self.c_eps = 0.7
        self.alpha_tke = 30.0
        self.mxl_min = 1e-8
        self.tke_mxl_choice = 2

        self.K_gm_0 = 1000.0

        self.enable_idemix = False
        self.enable_idemix_hor_diffusion = True

        self.eq_of_state_type = 5

    @pyom_method
    def set_grid(self):
        with Dataset("forcing.cdf", "r") as forcing_file:
            self.dxt[2:-2] = (self._x_boundary - self.x_origin) / self.nx
            self.dyt[2:-2] = (self._y_boundary - self.y_origin) / self.ny
            self.dzt[:] = forcing_file.variables["dzt"][::-1] / 100.0

    @pyom_method
    def set_coriolis(self):
        self.coriolis_t[:,:] = 2 * self.omega * np.sin(self.yt[np.newaxis, :] / 180. * self.pi)

    @pyom_method
    def _init_interpolator(self, coords, mask):
        data_coords = tuple(c[mask] for c in np.meshgrid(*coords, indexing="ij"))
        delaunay = scipy.spatial.Delaunay(np.vstack([i.flatten() for i in data_coords]).T)
        self.linear_interpolator = lambda val: scipy.interpolate.LinearNDInterpolator(delaunay, val)


    @pyom_method
    def _interpolate(self, coords, var, grid=None, fill_value="nan"):
        if grid is None:
            if len(coords) == 2:
                interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], self.zt[-1:], indexing="ij")
            elif len(coords) == 3:
                interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], self.zt, indexing="ij")
        else:
            interp_coords = np.meshgrid(*grid, indexing="ij")
        interp_coords = np.rollaxis(np.asarray(interp_coords),0,len(interp_coords)+1)
        return self.linear_interpolator(var[var > -1e19])(interp_coords)
        # if fill_value == "nan":
        #     fill_value = np.nan
        # if grid is None:
        #     if len(coords) == 2:
        #         interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], indexing="ij")
        #     elif len(coords) == 3:
        #         interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], self.zt, indexing="ij")
        #else:
        #    interp_coords = np.meshgrid(*grid, indexing="ij")
        #interp_coords = np.rollaxis(np.asarray(interp_coords),0,len(coords)+1)
        #
        # var = np.array(var)
        # invalid_mask = var < -1e19
        # if np.any(invalid_mask):
        #     var[invalid_mask] = np.nan
        # interp_values = scipy.interpolate.interpn(coords, var, interp_coords,
        #                 method="linear", bounds_error=False, fill_value=fill_value)
        # mask = np.isnan(interp_values)
        # if np.any(mask):
        #     ccoords = np.meshgrid(*coords, indexing="ij")
        #     interp_values[mask] = scipy.interpolate.griddata(tuple(c[~invalid_mask] for c in ccoords), var[~invalid_mask], interp_coords,
        #                           method="nearest")[mask]
        # return interp_values

    @pyom_method
    def set_topography(self):
        with Dataset("forcing.cdf","r") as forcing_file:
            topo_x, topo_y, topo_kmt, topo_depth = (forcing_file.variables[k][...].T for k in ("xt","yt","kmt","depth_t"))
            topo_bottom_depth = np.where(topo_kmt > 0, -topo_depth[topo_kmt - 1], 100)
            interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], indexing="ij")
            interp_coords = np.rollaxis(np.asarray(interp_coords),0,3)
            z_interp = scipy.interpolate.interpn((topo_x - 360, topo_y), topo_bottom_depth, interp_coords,
                                                  method="linear", bounds_error=False, fill_value=None)
        self.kbot[2:-2, 2:-2] = np.where(z_interp < 0., 1 + np.argmin(np.abs(z_interp[:, :, np.newaxis] - self.zt[np.newaxis, np.newaxis, :]), axis=2), 0)
        self.kbot *= self.kbot < self.nz


    @pyom_method
    def set_initial_conditions(self):
        with Dataset("forcing.cdf", "r") as forcing_file:
            forc_coords = [forcing_file.variables[k][...].T for k in ("xt","yt","zt")]
            forc_coords[0][...] += -360
            forc_coords[2][...] = -.01 * forc_coords[2][::-1]
            self._init_interpolator(forc_coords, forcing_file.variables["temp_ic"][::-1, :, :].T > -1e19)
            temp = self._interpolate(forc_coords, forcing_file.variables["temp_ic"][::-1, :, :].T)
            self.temp[2:-2, 2:-2, :, self.tau] = self.maskT[2:-2, 2:-2, :] * temp
            salt = 35. + 1000 * self._interpolate(forc_coords, forcing_file.variables["salt_ic"][::-1, :, :].T)
            self.salt[2:-2, 2:-2, :, self.tau] = self.maskT[2:-2, 2:-2, :] * salt

            self._taux = np.zeros((self.nx+4, self.ny+4, 12))
            self._tauy = np.zeros((self.nx+4, self.ny+4, 12))
            for k in xrange(12):
                self._taux[2:-2, 2:-2, k] = self._interpolate(forc_coords[:-1], forcing_file.variables["taux"][k,...].T, grid=(self.xu[2:-2], self.yt[2:-2])) / 10. / self.rho_0
                self._tauy[2:-2, 2:-2, k] = self._interpolate(forc_coords[:-1], forcing_file.variables["tauy"][k,...].T, grid=(self.xt[2:-2], self.yu[2:-2])) / 10. / self.rho_0

            for t in (self._taux, self._tauy):
                if self.enable_cyclic_x:
                    cyclic.setcyclic_x(t)
<commit_msg>[wip] work on north atlantic model<commit_after>from netCDF4 import Dataset
import scipy.interpolate
import scipy.spatial
import scipy.ndimage

from climate.pyom import PyOM, pyom_method, cyclic

class NorthAtlanticHighResolution(PyOM):
    """ Based on
    http://journals.ametsoc.org/doi/10.1175/1520-0485%282000%29030%3C1532%3ANSOTNA%3E2.0.CO%3B2
    """

    @pyom_method
    def set_parameter(self):
        """set main parameter
        """
        self.nx, self.ny, self.nz = (150,160,45)#(87,89,45)
        self.x_origin = -98.
        self.y_origin = -20.
        self._x_boundary = 17.2
        self._y_boundary = 72.6

        self.dt_mom = 3600.0
        self.dt_tracer = 3600.0

        self.runlen = 86400 * 365.0
        self.enable_diag_snapshots = True

        self.coord_degree = True

        self.congr_epsilon = 1e-10
        self.congr_max_iterations = 20000
        self.enable_streamfunction = True

        self.enable_neutral_diffusion = True
        self.enable_skew_diffusion = True
        self.K_iso_0 = 1000.0
        self.K_iso_steep = 200.0
        self.iso_dslope = 1./1000.0
        self.iso_slopec = 4./1000.0

        self.enable_hor_friction = True
        self.A_h = 5e4
        self.enable_hor_friction_cos_scaling = True
        self.hor_friction_cosPower = 3
        self.enable_tempsalt_sources = True

        self.enable_implicit_vert_friction = True
        self.enable_tke = False
        self.c_k = 0.1
        self.c_eps = 0.7
        self.alpha_tke = 30.0
        self.mxl_min = 1e-8
        self.tke_mxl_choice = 2

        self.K_gm_0 = 1000.0

        self.enable_idemix = False
        self.enable_idemix_hor_diffusion = True

        self.eq_of_state_type = 5

    @pyom_method
    def set_grid(self):
        with Dataset("forcing.cdf", "r") as forcing_file:
            self.dxt[2:-2] = (self._x_boundary - self.x_origin) / self.nx
            self.dyt[2:-2] = (self._y_boundary - self.y_origin) / self.ny
            self.dzt[:] = forcing_file.variables["dzt"][::-1] / 100.0

    @pyom_method
    def set_coriolis(self):
        self.coriolis_t[:,:] = 2 * self.omega * np.sin(self.yt[np.newaxis, :] / 180. * self.pi)

    @pyom_method
    def _interpolate(self, coords, var, grid=None):
        if grid is None:
            if len(coords) == 2:
                interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], self.zt[-1:], indexing="ij")
            elif len(coords) == 3:
                interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], self.zt, indexing="ij")
        else:
            interp_coords = np.meshgrid(*grid, indexing="ij")
        interp_coords = np.rollaxis(np.asarray(interp_coords),0,len(interp_coords)+1)

        if grid is None:
             if len(coords) == 2:
                 interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], indexing="ij")
             elif len(coords) == 3:
                 interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], self.zt, indexing="ij")
        else:
            interp_coords = np.meshgrid(*grid, indexing="ij")
        interp_coords = np.rollaxis(np.asarray(interp_coords),0,len(coords)+1)

        var = np.array(var)
        invalid_mask = var < -1e19
        if np.any(invalid_mask):
            var[invalid_mask] = np.nan
        interp_values = scipy.interpolate.interpn(coords, var, interp_coords,
                        method="linear", bounds_error=False, fill_value=np.nan)
        mask = np.isnan(interp_values)
        if np.any(mask):
            ccoords = np.meshgrid(*coords, indexing="ij")
            interp_values[mask] = scipy.interpolate.griddata(tuple(c[~invalid_mask] for c in ccoords), var[~invalid_mask], interp_coords,
                                                             method="nearest")[mask]
            #interp_values[mask] = scipy.interpolate.interpn(coords, var, interp_coords,
            #                method="nearest", bounds_error=False, fill_value=None)[mask]
        #mask = np.isnan(interp_values)
        #if np.any(mask):
        #    interp_values[mask] = 0
        return interp_values

    @pyom_method
    def set_topography(self):
        with Dataset("forcing.cdf","r") as forcing_file:
            topo_x, topo_y, topo_kmt, topo_depth = (forcing_file.variables[k][...].T for k in ("xt","yt","kmt","depth_t"))
            topo_x -= 360
            topo_bottom_depth = np.where(topo_kmt > 0, -topo_depth[topo_kmt - 1], 100)
        #with Dataset("ETOPO1_Bed_g_gmt4_NA.nc","r") as topography_file:
        #    topo_x, topo_y, topo_bottom_depth = (topography_file.variables[k][...].T for k in ("x","y","z"))
        #    topo_bottom_depth = scipy.ndimage.gaussian_filter(topo_bottom_depth, sigma=(len(topo_x) / self.nx, len(topo_y) / self.ny))
            interp_coords = np.meshgrid(self.xt[2:-2], self.yt[2:-2], indexing="ij")
            interp_coords = np.rollaxis(np.asarray(interp_coords),0,3)
            z_interp = scipy.interpolate.interpn((topo_x, topo_y), topo_bottom_depth, interp_coords,
                                                  method="linear", bounds_error=False, fill_value=None)
        self.kbot[2:-2, 2:-2] = np.where(z_interp < 0., 1 + np.argmin(np.abs(z_interp[:, :, np.newaxis] - self.zt[np.newaxis, np.newaxis, :]), axis=2), 0)
        self.kbot *= self.kbot < self.nz

        import matplotlib.pyplot as plt
        plt.imshow(topo_bottom_depth)
        plt.figure()
        plt.imshow(z_interp)
        plt.show()

    @pyom_method
    def set_initial_conditions(self):
        with Dataset("forcing.cdf", "r") as forcing_file:
            forc_coords = [forcing_file.variables[k][...].T for k in ("xt","yt","zt")]
            forc_coords[0][...] += -360
            forc_coords[2][...] = -.01 * forc_coords[2][::-1]
            #self._init_interpolator(forc_coords, forcing_file.variables["temp_ic"][::-1, :, :].T > -1e19)
            temp = self._interpolate(forc_coords, forcing_file.variables["temp_ic"][::-1, :, :].T)
            self.temp[2:-2, 2:-2, :, self.tau] = self.maskT[2:-2, 2:-2, :] * temp
            salt = 35. + 1000 * self._interpolate(forc_coords, forcing_file.variables["salt_ic"][::-1, :, :].T)
            self.salt[2:-2, 2:-2, :, self.tau] = self.maskT[2:-2, 2:-2, :] * salt

            self._taux = np.zeros((self.nx+4, self.ny+4, 12))
            self._tauy = np.zeros((self.nx+4, self.ny+4, 12))
            forc_u_coords_hor = [forcing_file.variables[k][...].T for k in ("xu","yu")]
            forc_u_coords_hor[0][...] += -360
            for k in xrange(12):
                self._taux[2:-2, 2:-2, k] = self._interpolate(forc_u_coords_hor, forcing_file.variables["taux"][k,...].T, grid=(self.xu[2:-2], self.yt[2:-2])) / 10. / self.rho_0
                self._tauy[2:-2, 2:-2, k] = self._interpolate(forc_u_coords_hor, forcing_file.variables["tauy"][k,...].T, grid=(self.xt[2:-2], self.yu[2:-2])) / 10. / self.rho_0

            for t in (self._taux, self._tauy):
                if self.enable_cyclic_x:
                    cyclic.setcyclic_x(t)

{'fraction_non_alphanumeric': 0.15106227106227108, 'fraction_numerical': 0.03926739926739927, 'mean_word_length': 2.2822793940851165, 'pattern_counts': {'":': 1, '<': 9, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2653280', 'n_tokens_mistral': 5413, 'n_tokens_neox': 5007, 'n_words': 791}