Python lat_lon.lon_lat_to_cartesian函数代码示例

def plot_difference(basemap,
                    lons1, lats1, data1, label1,
                    lons2, lats2, data2, label2,
                    base_folder="/skynet3_rech1/huziy/veg_fractions/"
                    ):
    xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons2.flatten(), lats2.flatten())
    xt, yt, zt = lat_lon.lon_lat_to_cartesian(lons1.flatten(), lats1.flatten())

    ktree = cKDTree(list(zip(xs, ys, zs)))
    dists, inds = ktree.query(list(zip(xt, yt, zt)))

    # Calculate differences
    diff_dict = {}
    for key, the_field in data2.items():
        diff_dict[key] = the_field.flatten()[inds].reshape(data1[key].shape) - data1[key]

    x, y = basemap(lons1, lats1)
    imname = "sand_clay_diff_{0}-{1}.jpeg".format(label2, label1)
    impath = os.path.join(base_folder, imname)
    plot_sand_and_clay_diff(x, y, basemap, diff_dict["SAND"], diff_dict["CLAY"],
                            out_image=impath)

    del diff_dict["SAND"], diff_dict["CLAY"]
    imname = "veg_fract_diff_{0}-{1}.jpeg".format(label2, label1)
    impath = os.path.join(base_folder, imname)
    plot_veg_fractions_diff(x, y, basemap, diff_dict,
                            out_image=impath)

    def __init__(self, file_path = "", var_name = "",
                 bathymetry_path = "/skynet3_rech1/huziy/NEMO_OFFICIAL/dev_v3_4_STABLE_2012/NEMOGCM/CONFIG/GLK_LIM3_Michigan/EXP00/bathy_meter.nc"):
        """
        :param file_path:
        :param var_name:
        :param bathymetry_path: used to mask land points
        """
        self.current_time_frame = -1
        self.var_name = var_name

        self.cube = iris.load_cube(file_path, constraint=iris.Constraint(cube_func=lambda c: c.var_name == var_name))
        self.lons, self.lats = cartography.get_xy_grids(self.cube)

        lons2d_gl, lats2d_gl = nemo_commons.get_2d_lons_lats_from_nemo(path=bathymetry_path)
        mask_gl = nemo_commons.get_mask(path=bathymetry_path)

        xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons2d_gl.flatten(), lats2d_gl.flatten())
        xt, yt, zt = lat_lon.lon_lat_to_cartesian(self.lons.flatten(), self.lats.flatten())

        tree = cKDTree(list(zip(xs, ys, zs)))
        dists, indices = tree.query(list(zip(xt, yt, zt)))

        self.mask = mask_gl.flatten()[indices].reshape(self.lons.shape)


        self.nt = self.cube.shape[0]
        assert isinstance(self.cube, Cube)
        print(self.nt)

    def read_and_interpolate_homa_data(self, path="", start_year=None, end_year=None, season_to_months=None):
        """
        :param path:
        :param target_cube:
        """
        import pandas as pd

        ds = Dataset(path)
        sst = ds.variables["sst"][:]

        # read longitudes and latitudes from a file
        lons_source = ds.variables["lon"][:]
        lats_source = ds.variables["lat"][:]



        month_to_season = defaultdict(lambda: "no-season")

        for seas, mths in season_to_months.items():
            for m in mths:
                month_to_season[m] = seas


        # time variable
        time_var = ds.variables["time"]
        dates = num2date(time_var[:], time_var.units)

        if hasattr(sst, "mask"):
            sst[sst.mask] = np.nan

        panel = pd.Panel(data=sst, items=dates, major_axis=range(sst.shape[1]), minor_axis=range(sst.shape[2]))



        seasonal_sst = panel.groupby(
            lambda d: (d.year, month_to_season[d.month]), axis="items").mean()


        # source grid
        xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons_source.flatten(), lats_source.flatten())
        kdtree = cKDTree(data=list(zip(xs, ys, zs)))

        # target grid
        xt, yt, zt = lat_lon.lon_lat_to_cartesian(self.lons.flatten(), self.lats.flatten())

        dists, inds = kdtree.query(list(zip(xt, yt, zt)))



        assert isinstance(seasonal_sst, pd.Panel)

        result = {}
        for the_year in range(start_year, end_year + 1):
            result[the_year] = {}
            for the_season in list(season_to_months.keys()):
                the_mean = seasonal_sst.select(lambda item: item == (the_year, the_season), axis="items")
                result[the_year][the_season] = the_mean.values.flatten()[inds].reshape(self.lons.shape)

        return result

def main(dfs_var_name="t2", cru_var_name="tmp",
         dfs_folder="/home/huziy/skynet3_rech1/NEMO_OFFICIAL/DFS5.2_interpolated",
         cru_file = "data/cru_data/CRUTS3.1/cru_ts_3_10.1901.2009.tmp.dat.nc"):

    if not os.path.isdir(NEMO_IMAGES_DIR):
        os.mkdir(NEMO_IMAGES_DIR)

    #year range is inclusive [start_year, end_year]
    start_year = 1981
    end_year = 2009

    season_name_to_months = OrderedDict([
        ("Winter", (1, 2, 12)),
        ("Spring", list(range(3, 6))),
        ("Summer", list(range(6, 9))),
        ("Fall", list(range(9, 12)))])

    cru_t_manager = CRUDataManager(var_name=cru_var_name, path=cru_file)
    cru_lons, cru_lats = cru_t_manager.lons2d, cru_t_manager.lats2d
    #get seasonal means (CRU)
    season_to_mean_cru = cru_t_manager.get_seasonal_means(season_name_to_months=season_name_to_months,
                                                          start_year=start_year,
                                                          end_year=end_year)
    #get seasonal means Drakkar
    dfs_manager = DFSDataManager(folder_path=dfs_folder, var_name=dfs_var_name)
    season_to_mean_dfs = dfs_manager.get_seasonal_means(season_name_to_months=season_name_to_months,
                                                        start_year=start_year,
                                                        end_year=end_year)

    dfs_lons, dfs_lats = dfs_manager.get_lons_and_lats_2d()
    xt, yt, zt = lat_lon.lon_lat_to_cartesian(dfs_lons.flatten(), dfs_lats.flatten())
    xs, ys, zs = lat_lon.lon_lat_to_cartesian(cru_lons.flatten(), cru_lats.flatten())
    ktree = cKDTree(data=list(zip(xs, ys, zs)))
    dists, inds = ktree.query(list(zip(xt, yt, zt)))

    season_to_err = OrderedDict()
    for k in season_to_mean_dfs:
        interpolated_cru = season_to_mean_cru[k].flatten()[inds].reshape(dfs_lons.shape)
        if dfs_var_name.lower() == "t2":
            #interpolated_cru += 273.15
            season_to_mean_dfs[k] -= 273.15
        elif dfs_var_name.lower() == "precip":  # precipitation in mm/day
            season_to_mean_dfs[k] *= 24 * 60 * 60

        season_to_err[k] = season_to_mean_dfs[k] #- interpolated_cru

    season_indicator = "-".join(sorted(season_to_err.keys()))
    fig_path = os.path.join(NEMO_IMAGES_DIR, "{3}_errors_{0}-{1}_{2}_dfs.jpeg".format(start_year,
                                                                                  end_year,
                                                                                  season_indicator,
                                                                                  dfs_var_name))

    basemap = nemo_commons.get_default_basemap_for_glk(dfs_lons, dfs_lats, resolution="l")
    x, y = basemap(dfs_lons, dfs_lats)
    coords_and_basemap = {
        "basemap": basemap, "x": x, "y": y
    }

    plot_errors_in_one_figure(season_to_err, fig_path=fig_path, **coords_and_basemap)

def main(in_file: Path, target_grid_file: Path, out_dir: Path=None):

    if out_dir is not None:
        out_dir.mkdir(exist_ok=True)
        out_file = out_dir / (in_file.name + "_interpolated")
    else:
        out_file = in_file.parent / (in_file.name + "_interpolated")

    if out_file.exists():
        print(f"Skipping {in_file}, output already exists ({out_file})")
        return

    with xarray.open_dataset(target_grid_file) as ds_grid:
        lons, lats = ds_grid["lon"][:].values, ds_grid["lat"][:].values

        xt, yt, zt = lat_lon.lon_lat_to_cartesian(lons.flatten(), lats.flatten())

        with xarray.open_dataset(in_file) as ds_in:

            lons_s, lats_s = ds_in["lon"][:].values, ds_in["lat"][:].values
            xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons_s.flatten(), lats_s.flatten())

            ktree = KDTree(list(zip(xs, ys, zs)))
            dists, inds = ktree.query(list(zip(xt, yt, zt)), k=1)




            # resample to daily
            ds_in_r = ds_in.resample(t="1D", keep_attrs=True).mean()


            ds_out = xarray.Dataset()
            for vname, var in ds_grid.variables.items():
                ds_out[vname] = var[:]

            ds_out["t"] = ds_in_r["t"][:]


            for vname, var in ds_in_r.variables.items():

                assert isinstance(var, xarray.Variable)

                var = var.squeeze()

                # only interested in (t, x, y) fields
                if var.ndim != 3:
                    print(f"skipping {vname}")
                    continue

                if vname.lower() not in ["t", "time", "lon", "lat"]:
                    print(f"Processing {vname}")
                    var_interpolated = [var[ti].values.flatten()[inds].reshape(lons.shape) for ti in range(var.shape[0])]
                    ds_out[vname] = xarray.DataArray(
                        var_interpolated, dims=("t", "x", "y"),
                        attrs=var.attrs,
                    )

            ds_out.to_netcdf(out_file)

    def interpolate_data_to_model_grid(self, model_lons_2d, model_lats_2d, data_obs):
        x0, y0, z0 = lat_lon.lon_lat_to_cartesian(model_lons_2d.flatten(), model_lats_2d.flatten())
        x, y, z = lat_lon.lon_lat_to_cartesian(self.lons2d.flatten(), self.lats2d.flatten())

        if self.ktree is None:
            self.ktree = KDTree(list(zip(x, y, z)))

        d, i = self.ktree.query(list(zip(x0, y0, z0)))

        return data_obs.flatten()[i].reshape(model_lons_2d.shape)

def get_dataless_model_points_for_stations(station_list, accumulation_area_km2_2d,
                                           model_lons2d, model_lats2d,
                                           i_array, j_array):
    """
    returns a map {station => modelpoint} for comparison modeled streamflows with observed

    this uses exactly the same method for searching model points as one in diagnose_point (nc-version)

    """
    lons = model_lons2d[i_array, j_array]
    lats = model_lats2d[i_array, j_array]
    model_acc_area_1d = accumulation_area_km2_2d[i_array, j_array]
    npoints = 1
    result = {}

    x0, y0, z0 = lat_lon.lon_lat_to_cartesian(lons, lats)
    kdtree = cKDTree(list(zip(x0, y0, z0)))

    for s in station_list:
        # list of model points which could represent the station

        assert isinstance(s, Station)
        x, y, z = lat_lon.lon_lat_to_cartesian(s.longitude, s.latitude)
        dists, inds = kdtree.query((x, y, z), k=5)

        if npoints == 1:

            deltaDaMin = np.min(np.abs(model_acc_area_1d[inds] - s.drainage_km2))

            # this returns a  list of numpy arrays
            imin = np.where(np.abs(model_acc_area_1d[inds] - s.drainage_km2) == deltaDaMin)[0][0]
            selected_cell_index = inds[imin]
            # check if difference in drainage areas is not too big less than 10 %

            print(s.river_name, deltaDaMin / s.drainage_km2)
            # if deltaDaMin / s.drainage_km2 > 0.2:
            #    continue

            mp = ModelPoint()
            mp.accumulation_area = model_acc_area_1d[selected_cell_index]
            mp.longitude = lons[selected_cell_index]
            mp.latitude = lats[selected_cell_index]
            mp.cell_index = selected_cell_index
            mp.distance_to_station = dists[imin]

            print("Distance to station: ", dists[imin])
            print("Model accumulation area: ", mp.accumulation_area)
            print("Obs accumulation area: ", s.drainage_km2)

            result[s] = mp
        else:
            raise Exception("npoints = {0}, is not yet implemented ...")
    return result

def get_timeseries_for_points(lons, lats, data_path="", varname="LD"):
    """
    return the list of timeseries for the points with the given coordinates,
    using the nearest neighbor interpolation
    :param varname: variable name
    :param data_path: path to the hdf5 file
    :param lons:
    :param lats:
    :return: pd.DataFrame with axes (time, point_index)
    """

    assert len(lons) == len(lats)

    df = None
    indices = None
    lk_fraction = None
    with tb.open_file(data_path) as h:
        var_table = h.get_node("/{}".format(varname))
        for i, row in enumerate(var_table):
            if df is None:
                df = pd.DataFrame(index=range(len(var_table)), columns=["date", ] + list(range(len(lons))))

                # calculate indices of the grid corresponding to the points
                bmp_info = get_basemap_info_from_hdf(file_path=data_path)
                """
                :type bmp_info: BasemapInfo
                """
                grid_lons, grid_lats = bmp_info.lons, bmp_info.lats

                x, y, z = lat_lon.lon_lat_to_cartesian(grid_lons.flatten(), grid_lats.flatten())

                ktree = KDTree(list(zip(x, y, z)))

                x1, y1, z1 = lat_lon.lon_lat_to_cartesian(lons, lats)

                dists, indices = ktree.query(list(zip(x1, y1, z1)))

                lk_fraction = get_array_from_file(path=data_path, var_name="lake_fraction")

            df.loc[i, :] = [datetime(row["year"], row["month"], row["day"], row["hour"]), ] + list(row["field"].flatten()[indices])

    # print lake fractions
    print(lk_fraction.flatten()[indices])
    print(sum(lk_fraction.flatten()[indices] > 0.05))

    df.set_index("date", inplace=True)
    df.sort_index(inplace=True)
    return df

    def get_daily_timeseries_using_mask(self, mask, lons2d_target, lats2d_target, multipliers_2d, start_date=None,
                                        end_date=None):
        """
        multipliers_2d used to multiply the values when aggregating into a single timeseries
        sum(mi * vi) - in space
        """

        bool_vect = np.array([start_date <= t <= end_date for t in self.times])

        new_times = list(filter(lambda t: start_date <= t <= end_date, self.times))
        new_vals = self.var_data[bool_vect, :, :]
        x_out, y_out, z_out = lat_lon.lon_lat_to_cartesian(lons2d_target.flatten(), lats2d_target.flatten())

        print(len(new_times))

        flat_mask = mask.flatten()
        x_out = x_out[flat_mask == 1]
        y_out = y_out[flat_mask == 1]
        z_out = z_out[flat_mask == 1]
        mults = multipliers_2d.flatten()[flat_mask == 1]
        data_interp = [self._interp_and_sum(new_vals[t, :, :].flatten(), flat_mask, x_out, y_out, z_out) for t in
                       range(len(new_times))]

        print("Interpolated all")
        return TimeSeries(time=new_times, data=data_interp).get_ts_of_daily_means()

    def _get_spatial_integrals_over_points_in_time(self, lons2d_target, lats2d_target, mask,
                                                   areas2d, start_date=None, end_date=None, var_name=""):
        """
        i)  Interpolate to the grid (lons2d_target, lats2d_target)
        ii) Apply the mask to the interpoated fields and sum with coefficients from areas2d

        Note: the interpolation is done using nearest neighbor approach

        returns a timeseries of {t -> sum(Ai[mask]*xi[mask])(t)}
        """


        #interpolation
        x1, y1, z1 = lat_lon.lon_lat_to_cartesian(lons2d_target.flatten(), lats2d_target.flatten())

        dists, indices = self.kdtree.query(list(zip(x1, y1, z1)))

        mask1d = mask.flatten().astype(int)
        areas1d = areas2d.flatten()

        result = {}
        for the_date in list(self.date_to_path.keys()):
            if start_date is not None:
                if start_date > the_date: continue

            if end_date is not None:
                if end_date < the_date: continue

            data = self.get_field_for_date(the_date, var_name=var_name)
            result[the_date] = np.sum(data.flatten()[indices][mask1d == 1] * areas1d[mask1d == 1])

        times = list(sorted(result.keys()))
        values = [result[x] for x in times]
        print("nvals, min, max", len(values), min(values), max(values))
        return TimeSeries(time=times, data=values)

def get_zone_around_lakes_mask(lons, lats, lake_mask, ktree=None, dist_km=100):
    """
    Returns the mask of a zone around lakes (excluding lakes) of a given width
    :type ktree: cKDTree
    :param ktree:
    :param lons:
    :param lats:
    :param lake_mask:
    :param dist_km:
    """

    x, y, z = lat_lon.lon_lat_to_cartesian(lons[lake_mask], lats[lake_mask])

    near_lake_zone = np.zeros_like(lons, dtype=bool)

    nlons = lons.shape[0] * lons.shape[1]
    near_lake_zone.shape = (nlons,)

    for xi, yi, zi in zip(x, y, z):
        dists, inds = ktree.query(np.array([[xi, yi, zi], ]), k=nlons, distance_upper_bound=dist_km * 1000)
        near_lake_zone[inds[inds < nlons]] = True

    near_lake_zone.shape = (lons.shape[0], lons.shape[1])

    # Remove lake points from the mask
    near_lake_zone &= ~lake_mask

    return near_lake_zone

    def get_kdtree(self, lons=None, lats=None, cache=True):
        """

        :param lons:
        :param lats:
        :param cache: if True then reuse the kdtree
        :return:
        """
        if lons is None:
            lons = self.lons
            lats = self.lats

        if lons is None:
            raise Exception(
                "The coordinates (lons and lats) are not yet set for the manager, please read some data first")

        if cache and self.kdtree is not None:
            return self.kdtree

        xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons.flatten(), lats.flatten())
        kdtree = KDTree(list(zip(xs, ys, zs)))

        if cache:
            self.kdtree = kdtree

        return kdtree

    def get_daily_clim_fields_aggregated_and_interpolated_to(self, start_year=None, end_year=None,
                                              lons_target=None, lats_target=None, n_agg_x=1, n_agg_y=1):
        """
        Aggregate fields to the desired resolution prior to interpolation
        :param start_year:
        :param end_year:
        :param lons_target:
        :param lats_target:
        """
        # Return 365 fields
        df = self.get_daily_climatology_fields(start_year=start_year, end_year=end_year)

        assert isinstance(df, pd.Panel)

        lons1d, lats1d = lons_target.flatten(), lats_target.flatten()
        xt, yt, zt = lat_lon.lon_lat_to_cartesian(lons1d, lats1d)

        dists, indices = self.kdtree.query(list(zip(xt, yt, zt)), k=n_agg_x * n_agg_y)

        clim_fields = [
            df.loc[:, day, :].values.flatten()[indices].mean(axis=1).reshape(lons_target.shape) for day in df.major_axis
        ]
        clim_fields = np.asarray(clim_fields)
        clim_fields = np.ma.masked_where(np.isnan(clim_fields), clim_fields)
        return df.major_axis, clim_fields

    def interpolate_data_to(self, data_in, lons2d, lats2d, nneighbours=4):
        """
        Interpolates data_in to the grid defined by (lons2d, lats2d)
        assuming that the data_in field is on the initial CRU grid

        interpolate using 4 nearest neighbors and inverse of squared distance
        """

        x_out, y_out, z_out = lat_lon.lon_lat_to_cartesian(lons2d.flatten(), lats2d.flatten())
        dst, ind = self.kdtree.query(list(zip(x_out, y_out, z_out)), k=nneighbours)

        data_in_flat = data_in.flatten()

        inverse_square = 1.0 / dst ** 2
        if len(dst.shape) > 1:
            norm = np.sum(inverse_square, axis=1)
            norm = np.array([norm] * dst.shape[1]).transpose()
            coefs = inverse_square / norm

            data_out_flat = np.sum(coefs * data_in_flat[ind], axis=1)
        elif len(dst.shape) == 1:
            data_out_flat = data_in_flat[ind]
        else:
            raise Exception("Could not find neighbor points")
        return np.reshape(data_out_flat, lons2d.shape)

    def _init_fields(self, nc_dataset):
        nc_vars = nc_dataset.variables
        lons = nc_vars["lon"][:]
        lats = nc_vars["lat"][:]


        if lons.ndim == 1:
            lats2d, lons2d = np.meshgrid(lats, lons)
        elif lons.ndim == 2:
            lats2d, lons2d = lats, lons
        else:
            raise NotImplementedError("Cannot handle {}-dimensional coordinates".format(lons.ndim))


        self.lons2d, self.lats2d = lons2d, lats2d

        self.times_var = nc_vars["time"]
        self.times_num = nc_vars["time"][:]

        if hasattr(self.times_var, "calendar"):
            self.times = num2date(self.times_num, self.times_var.units, self.times_var.calendar)
        else:
            self.times = num2date(self.times_num, self.times_var.units)


        if not self.lazy:

            self.var_data = nc_vars[self.var_name][:]
            if nc_vars[self.var_name].shape[1:] != self.lons2d.shape:
                print("nc_vars[self.var_name].shape = {}".format(nc_vars[self.var_name].shape))
                self.var_data = np.transpose(self.var_data, axes=[0, 2, 1])


        x_in, y_in, z_in = lat_lon.lon_lat_to_cartesian(self.lons2d.flatten(), self.lats2d.flatten())
        self.kdtree = cKDTree(list(zip(x_in, y_in, z_in)))