Read & update a standard JSON template#

In this notebook we cover two ways to read and update an existing COOLEST template file:

  • Option 1: using standard JSON utilities;

  • Option 2: using the dedicated COOLEST Python interface.

author: @aymgal

last update: 02/08/23

import os
from pprint import pprint
TEMPLATE_NAME = 'coolest_template'
TEMPLATE_DIR = 'template_dir'

Option 1: Load the JSON as a nested dictionnary#

This option only requires to the standard python module json. This can be useful for example when querying the template online for database purposes.

It allows to update ‘by hand’ the content of the JSON file, following the syntax of python containers (essentially nested dicts and lists).

import json

Load the JSON as standard dictionnary#

with open(os.path.join(os.getcwd(), TEMPLATE_DIR, TEMPLATE_NAME + '.json'), 'r') as f:
    content = json.load(f)

Explore its content#

# print the lensing entities
pprint(content['lensing_entities'])
[{'mass_model': [{'id': '0-massfield-mass-0-ExternalShear',
                  'parameters': {'gamma_ext': {'fixed': False,
                                               'id': '0-massfield-mass-0-ExternalShear-gamma_ext',
                                               'point_estimate': {'value': 0.07},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                 'phi_ext': {'fixed': False,
                                             'id': '0-massfield-mass-0-ExternalShear-phi_ext',
                                             'point_estimate': {'value': None},
                                             'posterior_stats': {'mean': None,
                                                                 'median': None,
                                                                 'percentile_16th': None,
                                                                 'percentile_84th': None},
                                             'prior': {'type': None}}},
                  'type': 'ExternalShear'}],
  'name': 'my lovely external shear',
  'redshift': 0.5,
  'type': 'MassField'},
 {'light_model': [{'id': '1-galaxy-light-0-Sersic',
                   'parameters': {'I_eff': {'fixed': False,
                                            'id': '1-galaxy-light-0-Sersic-I_eff',
                                            'point_estimate': {'value': None},
                                            'posterior_stats': {'mean': None,
                                                                'median': None,
                                                                'percentile_16th': None,
                                                                'percentile_84th': None},
                                            'prior': {'type': None}},
                                  'center_x': {'fixed': False,
                                               'id': '1-galaxy-light-0-Sersic-center_x',
                                               'point_estimate': {'value': None},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                  'center_y': {'fixed': False,
                                               'id': '1-galaxy-light-0-Sersic-center_y',
                                               'point_estimate': {'value': None},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                  'n': {'fixed': False,
                                        'id': '1-galaxy-light-0-Sersic-n',
                                        'point_estimate': {'value': None},
                                        'posterior_stats': {'mean': None,
                                                            'median': None,
                                                            'percentile_16th': None,
                                                            'percentile_84th': None},
                                        'prior': {'type': None}},
                                  'phi': {'fixed': False,
                                          'id': '1-galaxy-light-0-Sersic-phi',
                                          'point_estimate': {'value': None},
                                          'posterior_stats': {'mean': None,
                                                              'median': None,
                                                              'percentile_16th': None,
                                                              'percentile_84th': None},
                                          'prior': {'type': None}},
                                  'q': {'fixed': False,
                                        'id': '1-galaxy-light-0-Sersic-q',
                                        'point_estimate': {'value': None},
                                        'posterior_stats': {'mean': None,
                                                            'median': None,
                                                            'percentile_16th': None,
                                                            'percentile_84th': None},
                                        'prior': {'type': None}},
                                  'theta_eff': {'fixed': False,
                                                'id': '1-galaxy-light-0-Sersic-theta_eff',
                                                'point_estimate': {'value': None},
                                                'posterior_stats': {'mean': None,
                                                                    'median': None,
                                                                    'percentile_16th': None,
                                                                    'percentile_84th': None},
                                                'prior': {'type': None}}},
                   'type': 'Sersic'},
                  {'id': '1-galaxy-light-1-Sersic',
                   'parameters': {'I_eff': {'fixed': False,
                                            'id': '1-galaxy-light-1-Sersic-I_eff',
                                            'point_estimate': {'value': None},
                                            'posterior_stats': {'mean': None,
                                                                'median': None,
                                                                'percentile_16th': None,
                                                                'percentile_84th': None},
                                            'prior': {'type': None}},
                                  'center_x': {'fixed': False,
                                               'id': '1-galaxy-light-1-Sersic-center_x',
                                               'point_estimate': {'value': None},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                  'center_y': {'fixed': False,
                                               'id': '1-galaxy-light-1-Sersic-center_y',
                                               'point_estimate': {'value': None},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                  'n': {'fixed': False,
                                        'id': '1-galaxy-light-1-Sersic-n',
                                        'point_estimate': {'value': None},
                                        'posterior_stats': {'mean': None,
                                                            'median': None,
                                                            'percentile_16th': None,
                                                            'percentile_84th': None},
                                        'prior': {'type': None}},
                                  'phi': {'fixed': False,
                                          'id': '1-galaxy-light-1-Sersic-phi',
                                          'point_estimate': {'value': None},
                                          'posterior_stats': {'mean': None,
                                                              'median': None,
                                                              'percentile_16th': None,
                                                              'percentile_84th': None},
                                          'prior': {'type': None}},
                                  'q': {'fixed': False,
                                        'id': '1-galaxy-light-1-Sersic-q',
                                        'point_estimate': {'value': 0.89},
                                        'posterior_stats': {'mean': None,
                                                            'median': None,
                                                            'percentile_16th': None,
                                                            'percentile_84th': None},
                                        'prior': {'type': None}},
                                  'theta_eff': {'fixed': False,
                                                'id': '1-galaxy-light-1-Sersic-theta_eff',
                                                'point_estimate': {'value': None},
                                                'posterior_stats': {'mean': None,
                                                                    'median': None,
                                                                    'percentile_16th': None,
                                                                    'percentile_84th': None},
                                                'prior': {'type': None}}},
                   'type': 'Sersic'}],
  'mass_model': [{'id': '1-galaxy-mass-0-PEMD',
                  'parameters': {'center_x': {'fixed': False,
                                              'id': '1-galaxy-mass-0-PEMD-center_x',
                                              'point_estimate': {'value': None},
                                              'posterior_stats': {'mean': None,
                                                                  'median': None,
                                                                  'percentile_16th': None,
                                                                  'percentile_84th': None},
                                              'prior': {'type': None}},
                                 'center_y': {'fixed': False,
                                              'id': '1-galaxy-mass-0-PEMD-center_y',
                                              'point_estimate': {'value': None},
                                              'posterior_stats': {'mean': None,
                                                                  'median': None,
                                                                  'percentile_16th': None,
                                                                  'percentile_84th': None},
                                              'prior': {'type': None}},
                                 'gamma': {'fixed': False,
                                           'id': '1-galaxy-mass-0-PEMD-gamma',
                                           'point_estimate': {'value': None},
                                           'posterior_stats': {'mean': None,
                                                               'median': None,
                                                               'percentile_16th': None,
                                                               'percentile_84th': None},
                                           'prior': {'mean': 2.0,
                                                     'type': 'GaussianPrior',
                                                     'width': 0.2}},
                                 'phi': {'fixed': False,
                                         'id': '1-galaxy-mass-0-PEMD-phi',
                                         'point_estimate': {'value': None},
                                         'posterior_stats': {'mean': None,
                                                             'median': None,
                                                             'percentile_16th': None,
                                                             'percentile_84th': None},
                                         'prior': {'type': None}},
                                 'q': {'fixed': False,
                                       'id': '1-galaxy-mass-0-PEMD-q',
                                       'point_estimate': {'value': None},
                                       'posterior_stats': {'mean': None,
                                                           'median': None,
                                                           'percentile_16th': None,
                                                           'percentile_84th': None},
                                       'prior': {'type': None}},
                                 'theta_E': {'fixed': False,
                                             'id': '1-galaxy-mass-0-PEMD-theta_E',
                                             'point_estimate': {'value': None},
                                             'posterior_stats': {'mean': None,
                                                                 'median': None,
                                                                 'percentile_16th': None,
                                                                 'percentile_84th': None},
                                             'prior': {'type': None}}},
                  'type': 'PEMD'},
                 {'id': '1-galaxy-mass-1-PixelatedRegularGridPotential',
                  'parameters': {'pixels': {'field_of_view_x': [-3.0, 1.0],
                                            'field_of_view_y': [-2.0, 2.0],
                                            'fits_file': {'path': 'regul_grid_image.fits'},
                                            'id': '1-galaxy-mass-1-PixelatedRegularGridPotential-pixels',
                                            'num_pix_x': 0,
                                            'num_pix_y': 0}},
                  'type': 'PixelatedRegularGridPotential'}],
  'name': 'a lens galaxy',
  'redshift': 0.5,
  'type': 'Galaxy'},
 {'light_model': [{'id': '2-galaxy-light-0-Sersic',
                   'parameters': {'I_eff': {'fixed': False,
                                            'id': '2-galaxy-light-0-Sersic-I_eff',
                                            'point_estimate': {'value': None},
                                            'posterior_stats': {'mean': None,
                                                                'median': None,
                                                                'percentile_16th': None,
                                                                'percentile_84th': None},
                                            'prior': {'type': None}},
                                  'center_x': {'fixed': False,
                                               'id': '2-galaxy-light-0-Sersic-center_x',
                                               'point_estimate': {'value': None},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                  'center_y': {'fixed': False,
                                               'id': '2-galaxy-light-0-Sersic-center_y',
                                               'point_estimate': {'value': None},
                                               'posterior_stats': {'mean': None,
                                                                   'median': None,
                                                                   'percentile_16th': None,
                                                                   'percentile_84th': None},
                                               'prior': {'type': None}},
                                  'n': {'fixed': False,
                                        'id': '2-galaxy-light-0-Sersic-n',
                                        'point_estimate': {'value': None},
                                        'posterior_stats': {'mean': None,
                                                            'median': None,
                                                            'percentile_16th': None,
                                                            'percentile_84th': None},
                                        'prior': {'type': None}},
                                  'phi': {'fixed': False,
                                          'id': '2-galaxy-light-0-Sersic-phi',
                                          'point_estimate': {'value': None},
                                          'posterior_stats': {'mean': None,
                                                              'median': None,
                                                              'percentile_16th': None,
                                                              'percentile_84th': None},
                                          'prior': {'type': None}},
                                  'q': {'fixed': False,
                                        'id': '2-galaxy-light-0-Sersic-q',
                                        'point_estimate': {'value': None},
                                        'posterior_stats': {'mean': None,
                                                            'median': None,
                                                            'percentile_16th': None,
                                                            'percentile_84th': None},
                                        'prior': {'type': None}},
                                  'theta_eff': {'fixed': False,
                                                'id': '2-galaxy-light-0-Sersic-theta_eff',
                                                'point_estimate': {'value': None},
                                                'posterior_stats': {'mean': 0.11,
                                                                    'median': 0.15,
                                                                    'percentile_16th': 0.03,
                                                                    'percentile_84th': 0.05},
                                                'prior': {'type': None}}},
                   'type': 'Sersic'}],
  'mass_model': [],
  'name': 'a source galaxy',
  'redshift': 2.0,
  'type': 'Galaxy'},
 {'light_model': [{'id': '3-galaxy-light-0-PixelatedRegularGrid',
                   'parameters': {'pixels': {'field_of_view_x': [-3.0, 1.0],
                                             'field_of_view_y': [-2.0, 2.0],
                                             'fits_file': {'path': 'regul_grid_image.fits'},
                                             'id': '3-galaxy-light-0-PixelatedRegularGrid-pixels',
                                             'num_pix_x': 0,
                                             'num_pix_y': 0}},
                   'type': 'PixelatedRegularGrid'}],
  'mass_model': [],
  'name': 'another source',
  'redshift': 1.5,
  'type': 'Galaxy'},
 {'light_model': [{'id': '4-galaxy-light-0-IrregularGrid',
                   'parameters': {'pixels': {'field_of_view_x': [0, 0],
                                             'field_of_view_y': [0, 0],
                                             'fits_file': {'path': 'irreg_grid_pixels.fits'},
                                             'id': '4-galaxy-light-0-IrregularGrid-pixels',
                                             'num_pix': 0}},
                   'type': 'IrregularGrid'}],
  'mass_model': [],
  'name': 'a VKL source',
  'redshift': 1.2,
  'type': 'Galaxy'}]
# example for accessing some of the content
pprint(content['lensing_entities'][1]['mass_model'][0]['parameters']['theta_E'])
{'fixed': False,
 'id': '1-galaxy-mass-0-PEMD-theta_E',
 'point_estimate': {'value': None},
 'posterior_stats': {'mean': None,
                     'median': None,
                     'percentile_16th': None,
                     'percentile_84th': None},
 'prior': {'type': None}}

Add a prior on a given parameter#

# type of prior
content['lensing_entities'][1]\
    ['mass_model'][0]['parameters']['theta_E']\
    ['prior']['type'] = 'gaussian'

# mean
content['lensing_entities'][1]\
    ['mass_model'][0]['parameters']['theta_E']\
    ['prior']['mean'] = 1.0

# width
content['lensing_entities'][1]\
    ['mass_model'][0]['parameters']['theta_E']\
    ['prior']['width'] = 0.01

Assign a point estimate to a given parameter#

# here we add a gaussian prior on the source effective radius
content['lensing_entities'][1]\
    ['light_model'][0]['parameters']['theta_eff']\
    ['value'] = 0.8

Assign statistics about the posterior distribution of a parameter#

# here we add posterior statistics on the external shear strength
content['lensing_entities'][0]\
    ['mass_model'][0]['parameters']['gamma_ext']\
    ['posterior_stats']['mean'] = 0.03

content['lensing_entities'][0]\
    ['mass_model'][0]['parameters']['gamma_ext']\
    ['posterior_stats']['median'] = 0.031

content['lensing_entities'][0]\
    ['mass_model'][0]['parameters']['gamma_ext']\
    ['posterior_stats']['percentile_16th'] = 0.02

content['lensing_entities'][0]\
    ['mass_model'][0]['parameters']['gamma_ext']\
    ['posterior_stats']['percentile_84th'] = 0.03

Save the current content as an updated JSON file#

# uncomment for saving the updated JSON file
with open(os.path.join(os.getcwd(), TEMPLATE_DIR, 'updated_via_json.json'), 'w') as f:
    f.write(json.dumps(content, sort_keys=True, indent=2))
    
# Note: you can also update the original JSON template if you wish

Option 2: Use the COOLEST Python interface#

The coolest.template interface allows to manipulate the content of the JSON file as a hierachy of python objects, and to updated fields with new instances of objects.

from coolest.template.json import JSONSerializer
from coolest.template.classes.probabilities import GaussianPrior
from coolest.template.classes.parameter import PointEstimate
from coolest.template.classes.probabilities import PosteriorStatistics

Decode the JSON file#

decoder = JSONSerializer(os.path.join(os.getcwd(), TEMPLATE_DIR, TEMPLATE_NAME), indent=2)
coolest_object = decoder.load()

Explore the content, print some attributes#

print("Content:")
pprint(vars(coolest_object))
Content:
{'coordinates_origin': <coolest.template.classes.coordinates.CoordinatesOrigin object at 0x123d16b80>,
 'cosmology': <coolest.template.classes.cosmology.Cosmology object at 0x123d16f40>,
 'instrument': <coolest.template.classes.instrument.Instrument object at 0x123d16ee0>,
 'lensing_entities': [<coolest.template.classes.mass_field.MassField object at 0x123d16c10>,
                      <coolest.template.classes.galaxy.Galaxy object at 0x123d16880>,
                      <coolest.template.classes.galaxy.Galaxy object at 0x123d16820>,
                      <coolest.template.classes.galaxy.Galaxy object at 0x123d169d0>,
                      <coolest.template.classes.galaxy.Galaxy object at 0x123d168e0>],
 'meta': {},
 'mode': 'MAP',
 'observation': <coolest.template.classes.observation.Observation object at 0x123d54e50>,
 'standard': 'COOLEST'}
lensing_entities = coolest_object.lensing_entities
print("Attributes of the lens model:")
pprint(lensing_entities)
Attributes of the lens model:
[<coolest.template.classes.mass_field.MassField object at 0x123d16c10>,
 <coolest.template.classes.galaxy.Galaxy object at 0x123d16880>,
 <coolest.template.classes.galaxy.Galaxy object at 0x123d16820>,
 <coolest.template.classes.galaxy.Galaxy object at 0x123d169d0>,
 <coolest.template.classes.galaxy.Galaxy object at 0x123d168e0>]
lens_galaxy = lensing_entities[1]
print("Attributes of the first galaxy in the lens model:")
pprint(vars(lens_galaxy))
Attributes of the first galaxy in the lens model:
{'documentation': 'Class that defines a Galaxy entity (lens galaxy and/or '
                  'source galaxy).\n'
                  '\n'
                  '    Parameters\n'
                  '    ----------\n'
                  '    name : str\n'
                  '        Name associated to the galaxy.\n'
                  '    redshift : float\n'
                  '        Redshift of the galaxy, if any.\n'
                  '    light_model : LightModel, optional\n'
                  '        Light model instance describing the surface '
                  'brightness of the galaxy, by default None\n'
                  '    mass_model : MassModel, optional\n'
                  '        Mass model instance describing the mass '
                  'distribution of the galaxy, by default None',
 'light_model': [<coolest.template.classes.profiles.light.Sersic object at 0x11092ea90>,
                 <coolest.template.classes.profiles.light.Sersic object at 0x11092e940>],
 'mass_model': [<coolest.template.classes.profiles.mass.PEMD object at 0x123d16580>,
                <coolest.template.classes.profiles.mass.PixelatedRegularGridPotential object at 0x123d16130>],
 'name': 'a lens galaxy',
 'redshift': 0.5,
 'type': 'Galaxy'}
lens_profile = lens_galaxy.mass_model[0]
print("Attributes of its first mass profile:")
pprint(vars(lens_profile))
Attributes of its first mass profile:
{'documentation': 'Power-law elliptical mass density profile, \n'
                  '    also called the Elliptical Power-law profile (EPL).\n'
                  '\n'
                  '    This profile is described by the following parameters:\n'
                  '\n'
                  "    - 'gamma': logarithmic radial slope\n"
                  "    - 'theta_E': Einstein radius\n"
                  "    - 'q': axis ratio (semi-major axis / semi-minor axis)\n"
                  "    - 'phi': position angle\n"
                  "    - 'center_x': position along the x coordinate\n"
                  "    - 'center_y': position along the y coordinate",
 'id': '1-galaxy-mass-0-PEMD',
 'parameters': {'center_x': <coolest.template.classes.parameter.NonLinearParameter object at 0x123d1c4f0>,
                'center_y': <coolest.template.classes.parameter.NonLinearParameter object at 0x123d1c310>,
                'gamma': <coolest.template.classes.parameter.NonLinearParameter object at 0x123d16eb0>,
                'phi': <coolest.template.classes.parameter.NonLinearParameter object at 0x123d1c7c0>,
                'q': <coolest.template.classes.parameter.NonLinearParameter object at 0x123d1ce50>,
                'theta_E': <coolest.template.classes.parameter.NonLinearParameter object at 0x123d1ca60>},
 'type': 'PEMD'}

Add a prior on a given parameter#

# here we add a gaussian prior on the Einstein radius
prior = GaussianPrior(mean=1.0, width=0.01)

coolest_object.lensing_entities[1].mass_model[0].parameters['theta_E'].set_prior(prior)

Assign a point estimate to a given parameter#

# here we add a gaussian prior on the source effective radius
estimate = PointEstimate(value=0.8)

coolest_object.lensing_entities[2].light_model[0].parameters['theta_eff'].set_point_estimate(estimate)

Assign statistics about the posterior distribution of a parameter#

# here we add posterior statistics on the shear strength
posterior \
    = PosteriorStatistics(mean=0.03, median=0.031, 
                          percentile_16th=0.02, percentile_84th=0.03)

coolest_object.lensing_entities[0].mass_model[0].parameters['gamma_ext'].set_posterior(posterior)

Encode the updated classes as a new template#

encoder = JSONSerializer(os.path.join(os.getcwd(), TEMPLATE_DIR, 'updated_via_coolest'), 
                         obj=coolest_object, indent=2)

# uncomment for dumping the content to the JSON file
coolest_object_encoded = encoder.dump_simple()

# Note: you can also update the original JSON template if you wish