metadata_utils.py
import numpy as np
from ast import literal_eval
import fnmatch
import lenstronomy.Util.param_util as param_util
__all__ = ['add_qphi_columns', 'add_g1g2_columns', 'add_gamma_psi_ext_columns', 'add_relative_src_offset', 'get_kwargs_src_light', 'get_kwargs_lens_light', 'get_kwargs_ps', 'get_kwargs_lens_mass', 'get_kwargs_ps_lensed']
def get_kwargs_ps_lensed(metadata_row):
x_image = np.array(literal_eval(metadata_row['x_image']))
y_image = np.array(literal_eval(metadata_row['y_image']))
magnification = np.array(literal_eval(metadata_row['magnification']))
kwargs = dict(
ra_image=x_image,
dec_image=y_image,
point_amp=magnification,
)
return kwargs
def get_nested_ps(metadata_row):
nested = {'agn_light': get_kwargs_ps(metadata_row)[0]}
return nested
def get_kwargs_src_light(metadata_row):
#TODO: use keys if Series, columns.values if DataFrame
src_light_keys = [col.split('src_light_')[1] for col in metadata_row.keys() if col.startswith('src_light_')]
src_light_cols = ['src_light_' + col for col in src_light_keys]
kwargs = dict(zip(src_light_keys, metadata_row[src_light_cols]))
return kwargs
def get_kwargs_lens_light(metadata_row):
lens_light_keys = [col.split('lens_light_')[1] for col in metadata_row.keys() if col.startswith('lens_light_')]
lens_light_cols = ['lens_light_' + col for col in lens_light_keys]
kwargs = dict(zip(lens_light_keys, metadata_row[lens_light_cols]))
return kwargs
def get_kwargs_ps(metadata_row):
kwargs = {'magnitude': metadata_row['agn_light_magnitude'], 'ra_source': metadata_row['src_light_center_x'], 'dec_source': metadata_row['src_light_center_y']}
return [kwargs]
def get_kwargs_lens_mass(metadata_row):
# Profile
profile_keys = [col.split('lens_mass_')[1] for col in metadata_row.keys() if col.startswith('lens_mass')]
profile_cols = ['lens_mass_' + col for col in profile_keys]
kwargs_profile = dict(zip(profile_keys, metadata_row[profile_cols]))
# External shear
ext_shear_keys = [col.split('external_shear_')[1] for col in metadata_row.keys() if col.startswith('external_shear_')]
ext_shear_cols = ['external_shear_' + col for col in ext_shear_keys]
kwargs_ext_shear = dict(zip(ext_shear_keys, metadata_row[ext_shear_cols]))
return [kwargs_profile, kwargs_ext_shear]
def add_qphi_columns(metadata):
"""Add alternate ellipticity definitions (axis ratio and angle) for each component for which ellipticity was defined in terms of e1, e2
Parameters
----------
metadata : pd.DataFrame
the metadatadata generated by Baobab
Returns
-------
pd.DataFrame
metadata augmented with e1, e2 for each relevant component
"""
e1_col_names = sorted(fnmatch.filter(metadata.columns.values, '*_e1'))
e2_col_names = sorted(fnmatch.filter(metadata.columns.values, '*_e2'))
for i, e1_col_name in enumerate(e1_col_names):
e2_col_name = e2_col_names[i]
comp_name = e1_col_name.split('_e1')[0] # component name, e.g. 'lens_light'
e1 = metadata[e1_col_name].values
e2 = metadata[e2_col_name].values
phi, q = param_util.ellipticity2phi_q(e1, e2)
metadata['{:s}_q'.format(comp_name)] = q
metadata['{:s}_phi'.format(comp_name)] = phi
return metadata
def add_g1g2_columns(metadata):
"""Add alternate shear definitions (gamma1, gamma2) for external shear defined in terms of shear modulus and angle (gamma_ext, psi_ext)
Parameters
----------
metadata : pd.DataFrame
the metadatadata generated by Baobab
Returns
-------
pd.DataFrame
metadata augmented with gamma1, gamma2 for the external shear component
"""
gamma_ext = metadata['external_shear_gamma_ext'].values
psi_ext = metadata['external_shear_psi_ext'].values
gamma1, gamma2 = param_util.shear_polar2cartesian(phi=psi_ext, gamma=gamma_ext)
metadata['external_shear_gamma1'] = gamma1
metadata['external_shear_gamma2'] = gamma2
return metadata
def add_gamma_psi_ext_columns(metadata):
"""Add alternate shear definitions (gamma1, gamma2) for external shear defined in terms of shear modulus and angle (gamma_ext, psi_ext)
Parameters
----------
metadata : pd.DataFrame
the metadatadata generated by Baobab
Returns
-------
pd.DataFrame
metadata augmented with gamma1, gamma2 for the external shear component
"""
gamma1 = metadata['external_shear_gamma1'].values
gamma2 = metadata['external_shear_gamma2'].values
psi_ext, gamma_ext = param_util.shear_cartesian2polar(gamma1=gamma1, gamma2=gamma2)
metadata['external_shear_gamma_ext'] = gamma_ext
metadata['external_shear_psi_ext'] = psi_ext
return metadata
def add_relative_src_offset(metadata):
"""Get the source offset relative to the lens center
Parameters
----------
metadata : pd.DataFrame
the metadata generated by Baobab
Returns
-------
pd.DataFrame
metadata augmented with relative source offset columns added
"""
metadata['src_light_pos_offset_x'] = metadata['src_light_center_x'] - metadata['lens_mass_center_x']
metadata['src_light_pos_offset_y'] = metadata['src_light_center_y'] - metadata['lens_mass_center_y']
return metadata