From a1a20072713b397bb006321f364bacc25583a7b3 Mon Sep 17 00:00:00 2001 From: bcarreres Date: Sun, 18 Feb 2024 10:13:23 -0500 Subject: [PATCH] add geout --- snsim/geo_utils.py | 62 ++++++++++++++++++++++++++++++++++++++++++++ snsim/simu.py | 14 +++++++--- snsim/survey_host.py | 52 +++++++++++++++++++------------------ snsim/utils.py | 58 ----------------------------------------- 4 files changed, 99 insertions(+), 87 deletions(-) create mode 100644 snsim/geo_utils.py diff --git a/snsim/geo_utils.py b/snsim/geo_utils.py new file mode 100644 index 0000000..10c3215 --- /dev/null +++ b/snsim/geo_utils.py @@ -0,0 +1,62 @@ +"""This module contains usefull function for the survey and field geometry.""" +import numpy as np +import geopandas as gpd +from shapely import geometry as shp_geo +from shapely import ops as shp_ops +from .constants import _SPHERE_LIMIT_ + + +def _format_corner(corner, RA): + # -- Replace corners that cross sphere edges + # + # 0 ---- 1 + # | | + # 3 ---- 2 + # + # conditions : + # - RA_0 < RA_1 + # - RA_3 < RA_2 + # - RA_0 and RA_3 on the same side of the field center + + sign = (corner[:, 3, :, 0] - RA[:, None]) * (corner[:, 0, :, 0] - RA[:, None]) < 0 + comp = corner[:, 0, :, 0] < corner[:, 3, :, 0] + + corner[:, 1, :, 0][corner[:, 1, :, 0] < corner[:, 0, :, 0]] += 2 * np.pi + corner[:, 2, :, 0][corner[:, 2, :, 0] < corner[:, 3, :, 0]] += 2 * np.pi + + corner[:, 0, :, 0][sign & comp] += 2 * np.pi + corner[:, 1, :, 0][sign & comp] += 2 * np.pi + + corner[:, 2, :, 0][sign & ~comp] += 2 * np.pi + corner[:, 3, :, 0][sign & ~comp] += 2 * np.pi + return corner + + +def _compute_area(polygon): + """Compute survey total area.""" + # It's an integration by dec strip + area = 0 + strip_dec = np.linspace(-np.pi/2, np.pi/2, 10_000) + for da, db in zip(strip_dec[:-1], strip_dec[1:]): + line = shp_geo.LineString([[0, (da + db) * 0.5], [2 * np.pi, (da + db) * 0.5]]) + if line.intersects(polygon): + dRA = line.intersection(polygon).length + area += dRA * (np.sin(db) - np.sin(da)) + return area + + +def _compute_polygon(corners): + """Create polygon on a sphere, check for edges conditions.""" + + # Create polygons + polygons = gpd.GeoSeries([shp_geo.Polygon(corners[:, j, :]) for j in range(corners.shape[1])]) + + # Check if they intersect the 2pi edge line + int_mask = polygons.intersects(_SPHERE_LIMIT_) + + # If they do cut divide them in 2 and translate the one that is beyond the edge at -2pi + polydiv = gpd.GeoSeries(shp_ops.polygonize(polygons[int_mask].boundary.union(_SPHERE_LIMIT_))) + transl_mask = polydiv.boundary.bounds['maxx'] > 2 * np.pi + polydiv[transl_mask] = polydiv[transl_mask].translate(-2*np.pi) + + return shp_geo.MultiPolygon([*polygons[~int_mask].values, *polydiv.values]) diff --git a/snsim/simu.py b/snsim/simu.py index 4b577a0..0e676bb 100644 --- a/snsim/simu.py +++ b/snsim/simu.py @@ -453,8 +453,11 @@ def _cadence_sim(self, seed, generator, Obj_ID=0): else: duration = generator.time_range[1] - generator.time_range[0] - n_obj = self._gen_n_sn(generator._z_time_rate[1], - duration,seed=seeds[0], area=self.survey._envelope_area) + n_obj = self._gen_n_sn( + generator._z_time_rate[1], + duration, + seed=seeds[0], + area=self.survey._envelope_area) lcs = self._sim_lcs(seeds[1], generator, n_obj, Obj_ID=Obj_ID) @@ -482,8 +485,11 @@ def _fix_nsn_sim(self, seed, generator, Obj_ID=0): raise_trigger = 0 n_to_sim = generator._params['force_n'] while len(lcs) < generator._params['force_n']: - lcs += self._sim_lcs(seed, generator, n_to_sim, - Obj_ID=len(lcs)) + lcs += self._sim_lcs( + seed, + generator, + n_to_sim, + Obj_ID=len(lcs)) # -- Arbitrary cut to stop the simulation if no SN are geenrated if n_to_sim == generator._params['force_n'] - len(lcs): diff --git a/snsim/survey_host.py b/snsim/survey_host.py index d9afb4c..a42686c 100644 --- a/snsim/survey_host.py +++ b/snsim/survey_host.py @@ -13,9 +13,9 @@ from shapely import ops as shp_ops import dask.dataframe as daskdf from . import utils as ut +from . import geo_utils as geo_ut from . import io_utils as io_ut from . import nb_fun as nbf -from .constants import C_LIGHT_KMS class SurveyObs: @@ -91,10 +91,10 @@ def _compute_envelope(self): corners = np.stack([nbf.new_coord_on_fields(sub_fields_corners[:, :, i, :], np.stack([f_RA, f_Dec])) for i in range(4)], axis=1) - corners = ut._format_corner(corners, f_RA) + corners = geo_ut._format_corner(corners, f_RA) - envelope = shp_ops.unary_union([ut._compute_polygon(corners[i]) for i in range(4)]).envelope - envelope_area = ut._compute_area(envelope) + envelope = shp_ops.unary_union([geo_ut._compute_polygon(corners[i]) for i in range(4)]).envelope + envelope_area = geo_ut._compute_area(envelope) return envelope, envelope_area def __str__(self): @@ -160,10 +160,11 @@ def _read_start_end_days(self, obs_dic): end_day = ut.init_astropy_time(end_day) if end_day.mjd > max_mjd or start_day.mjd < min_mjd: - warnings.warn(f'Starting day {start_day.mjd:.3f} MJD or' - f'Ending day {end_day.mjd:.3f} MJD is outer of' - f'the survey range : {min_mjd:.3f} - {max_mjd:.3f}', - UserWarning) + warnings.warn( + f'Starting day {start_day.mjd:.3f} MJD or' + f'Ending day {end_day.mjd:.3f} MJD is outer of' + f'the survey range : {min_mjd:.3f} - {max_mjd:.3f}', + UserWarning) if end_day.mjd < start_day.mjd: raise ValueError("The ending day is before the starting day !") @@ -341,7 +342,6 @@ def _init_data(self): end_day = ut.init_astropy_time(maxMJDinObs) return obs_dic, (start_day, end_day) - def _init_fields_map(self, field_config): """Init the subfield map parameters. @@ -357,10 +357,11 @@ def _init_fields_map(self, field_config): """ if field_config == 'rectangle': - sub_fields_corners = {0: np.array([[[-self.field_size_rad[0] / 2, self.field_size_rad[1] / 2], - [ self.field_size_rad[0] / 2, self.field_size_rad[1] / 2], - [ self.field_size_rad[0] / 2, -self.field_size_rad[1] / 2], - [-self.field_size_rad[0] / 2, -self.field_size_rad[1] / 2]]])} + sub_fields_corners = {0: np.array( + [[[-self.field_size_rad[0] / 2, self.field_size_rad[1] / 2], + [ self.field_size_rad[0] / 2, self.field_size_rad[1] / 2], + [ self.field_size_rad[0] / 2, -self.field_size_rad[1] / 2], + [-self.field_size_rad[0] / 2, -self.field_size_rad[1] / 2]]])} else: sub_fields_corners = io_ut._read_sub_field_map(self.field_size_rad, field_config) @@ -404,10 +405,10 @@ def _match_radec_to_obs(df, ObjPoints, config, sub_fields_corners): np.array([df.fieldRA.values, df.fieldDec.values])) for i in range(4)], axis=1) - corners = ut._format_corner(corners, df.fieldRA.values) + corners = geo_ut._format_corner(corners, df.fieldRA.values) # -- Create shapely polygon - fgeo = np.vectorize(lambda i: ut._compute_polygon(corners[i])) + fgeo = np.vectorize(lambda i: geo_ut._compute_polygon(corners[i])) GeoS = gpd.GeoDataFrame(data=df, geometry=fgeo(np.arange(df.shape[0]))) @@ -419,7 +420,7 @@ def _match_radec_to_obs(df, ObjPoints, config, sub_fields_corners): return join.drop(columns=['geometry', 'index_right', 'min_t', 'max_t', '1_zobs', 't0']) def get_observations(self, params, phase_cut=None, nep_cut=None, IDmin=0, - use_dask=False, npartitions=None): + use_dask=False, npartitions=None): """Give the epochs of observations of a given SN. Parameters @@ -447,8 +448,8 @@ def get_observations(self, params, phase_cut=None, nep_cut=None, IDmin=0, """ params = params.copy() ObjPoints = gpd.GeoDataFrame(data=params[['t0', 'min_t', 'max_t', '1_zobs']], - geometry=gpd.points_from_xy(params['ra'], params['dec']), - index=params.index) + geometry=gpd.points_from_xy(params['ra'], params['dec']), + index=params.index) if use_dask: if npartitions is None: @@ -456,7 +457,7 @@ def get_observations(self, params, phase_cut=None, nep_cut=None, IDmin=0, npartitions = len(self.obs_table) // 10 ddf = daskdf.from_pandas(self.obs_table, npartitions=npartitions) meta = daskdf.utils.make_meta({**{k: t for k, t in zip(ddf.columns, ddf.dtypes)}, - 'phase': 'float64'}) + 'phase': 'float64'}) ObsObj = ddf.map_partitions(self._match_radec_to_obs, ObjPoints, self.config, self._sub_field_corners, @@ -482,7 +483,7 @@ def get_observations(self, params, phase_cut=None, nep_cut=None, IDmin=0, params = params.loc[ObsObj.index.unique()] # -- Reset index - new_idx = {k:IDmin + i for i, k in enumerate(ObsObj.index.unique())} + new_idx = {k: IDmin + i for i, k in enumerate(ObsObj.index.unique())} ObsObj['ID'] = ObsObj.index.map(new_idx) params['ID'] = params.index.map(new_idx) @@ -557,8 +558,8 @@ def show_map(self, ax=None): ax.set_xlabel('RA [deg]') ax.set_ylabel('Dec [deg]') ax.set_xlim(-self.config['ra_size'] / 2 - 0.5, - self.config['ra_size'] / 2 + 0.5) - ax.set_ylim(-self.config['dec_size'] / 2 - 0.5, + self.config['ra_size'] / 2 + 0.5) + ax.set_ylim(-self.config['dec_size'] / 2 - 0.5, self.config['dec_size'] / 2 + 0.5) ax.set_aspect('equal') if ax is None: @@ -654,6 +655,7 @@ def __init__(self, config, z_range=None, geometry=None): elif self.config['distrib'].lower() not in self._dist_options: raise ValueError(f"{self.config['distrib']} is not an available option," f"distributions are {self._dist_options}") + @property def config(self): """Get the configuration dic of host.""" @@ -745,9 +747,9 @@ def compute_weights(self, rate=None): weights /= weights.sum() return weights - #elif self.config['distrib'].lower() == 'gal_prop': - #weights that depends on galaxy properties, it will depend on the SN type, to figure out implementation - #see vincenzi et al, and ask alex kim fo his model + # elif self.config['distrib'].lower() == 'gal_prop': + # weights that depends on galaxy properties, it will depend on the SN type, to figure out implementation + # see vincenzi et al, and ask alex kim fo his model def random_choice(self, n, seed=None, rate=None): """Randomly select hosts. diff --git a/snsim/utils.py b/snsim/utils.py index da74c9f..8303431 100644 --- a/snsim/utils.py +++ b/snsim/utils.py @@ -435,64 +435,6 @@ def print_dic(dic, prefix=''): else: print(prefix + f'{K}: {dic[K]}') - -def _format_corner(corner, RA): - # -- Replace corners that cross sphere edges - # - # 0 ---- 1 - # | | - # 3 ---- 2 - # - # conditions : - # - RA_0 < RA_1 - # - RA_3 < RA_2 - # - RA_0 and RA_3 on the same side of the field center - - sign = (corner[:, 3, :, 0] - RA[:, None]) * (corner[:, 0, :, 0] - RA[:, None]) < 0 - comp = corner[:, 0, :, 0] < corner[:, 3, :, 0] - - corner[:, 1, :, 0][corner[:, 1, :, 0] < corner[:, 0, :, 0]] += 2 * np.pi - corner[:, 2, :, 0][corner[:, 2, :, 0] < corner[:, 3, :, 0]] += 2 * np.pi - - - corner[:, 0, :, 0][sign & comp] += 2 * np.pi - corner[:, 1, :, 0][sign & comp] += 2 * np.pi - - corner[:, 2, :, 0][sign & ~comp] += 2 * np.pi - corner[:, 3, :, 0][sign & ~comp] += 2 * np.pi - return corner - - -def _compute_area(polygon): - """Compute survey total area.""" - # It's an integration by dec strip - area = 0 - strip_dec = np.linspace(-np.pi/2, np.pi/2, 10_000) - for da, db in zip(strip_dec[:-1], strip_dec[1:]): - line = shp_geo.LineString([[0, (da + db) * 0.5], [2 * np.pi, (da + db) * 0.5]]) - if line.intersects(polygon): - dRA = line.intersection(polygon).length - area += dRA * (np.sin(db) - np.sin(da)) - return area - - -def _compute_polygon(corners): - """Create polygon on a sphere, check for edges conditions.""" - - # Create polygons - polygons = gpd.GeoSeries([shp_geo.Polygon(corners[:, j, :]) for j in range(corners.shape[1])]) - - # Check if they intersect the 2pi edge line - int_mask = polygons.intersects(_SPHERE_LIMIT_) - - # If they do cut divide them in 2 and translate the one that is beyond the edge at -2pi - polydiv = gpd.GeoSeries(shp_ops.polygonize(polygons[int_mask].boundary.union(_SPHERE_LIMIT_))) - transl_mask = polydiv.boundary.bounds['maxx'] > 2 * np.pi - polydiv[transl_mask] = polydiv[transl_mask].translate(-2*np.pi) - - return shp_geo.MultiPolygon([*polygons[~int_mask].values, *polydiv.values]) - - def Templatelist_fromsncosmo(source_type=None): """ list names of templates in sncosmo built-in sources catalogue Parameters