Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Model selection #6

Merged
merged 2 commits into from
Mar 6, 2023
Merged

Model selection #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80349f0
add a bit more sampling padding to be safe
Mar 3, 2023
6e2be40
added notebook with example on training the model
swagnercarena Mar 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion jaxstronomy/input_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def get_config():
'num_z_bins': 1000,
'los_pad_length': 10,
'subhalos_pad_length': 750,
'sampling_pad_length': 100000,
'sampling_pad_length': 200000,
}

config['rng'] = jax.random.PRNGKey(0)
Expand Down
241 changes: 241 additions & 0 deletions notebooks/SimpleTrainings.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,241 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "29c9b893",
"metadata": {},
"outputs": [],
"source": [
"import functools\n",
"\n",
"from flax.training import checkpoints\n",
"import jax\n",
"import jax.numpy as jnp\n",
"import matplotlib.pyplot as plt\n",
"from scipy.stats import linregress\n",
"\n",
"from jaxstronomy import train\n",
"from jaxstronomy import input_pipeline"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7ef2bbfe",
"metadata": {},
"outputs": [],
"source": [
"# Pull our input configurations\n",
"from jaxstronomy import train_config\n",
"from jaxstronomy import input_config as ic\n",
"\n",
"config = train_config.get_config()\n",
"input_config = ic.get_config()\n",
"image_size = input_config['kwargs_detector']['n_x']\n",
"rng = jax.random.PRNGKey(0)\n",
"\n",
"# A few parameters we'll need later on to generate our images for comparison.\n",
"cosmology_params = input_pipeline.intialize_cosmology_params(input_config, rng)\n",
"grid_x, grid_y = input_pipeline.generate_grids(input_config)\n",
"\n",
"# A function for drawing large batches of images without running out of memory.\n",
"def generate_images_and_truth(input_config, cosmology_params, grid_x, grid_y, rng, n_batches):\n",
" draw_image_and_truth_vmap = jax.jit(jax.vmap(functools.partial(\n",
" input_pipeline.draw_image_and_truth, all_models=input_config['all_models'],\n",
" principal_md_index=input_config['principal_md_index'],\n",
" principal_source_index=input_config['principal_source_index'],\n",
" kwargs_simulation=input_config['kwargs_simulation'], kwargs_detector=input_config['kwargs_detector'],\n",
" kwargs_psf=input_config['kwargs_psf'], truth_parameters=input_config['truth_parameters']),\n",
" in_axes=(None, None, None, None, 0)))\n",
" rng_batch, _ = jax.random.split(rng)\n",
" image = []\n",
" truth = []\n",
" for _ in range(n_batches):\n",
" rng_batch, _ = jax.random.split(rng_batch)\n",
" rng_images = jax.random.split(rng_batch, config.batch_size)\n",
" image_draw, truth_draw = draw_image_and_truth_vmap(input_config['lensing_config'], cosmology_params, grid_x, \n",
" grid_y, rng_images)\n",
" image.append(image_draw)\n",
" truth.append(truth_draw)\n",
" \n",
" return jnp.concatenate(image, axis=0), jnp.concatenate(truth, axis=0)\n",
"\n",
"# A jitted function for getting the outputs out of the current model.\n",
"@jax.jit\n",
"def get_outputs(state, image):\n",
" return state.apply_fn({'params': state.params, 'batch_stats': state.batch_stats}, \n",
" jnp.expand_dims(image, axis=-1), mutable=['batch_stats'])\n",
"\n",
"# We're just going for short trainings on these examples. The problems will be kept very easy, so the\n",
"# signal should be strong.\n",
"config.steps_per_epoch = 10\n",
"config.num_train_steps = config.steps_per_epoch * 50\n",
"config.keep_every_n_steps = config.steps_per_epoch * 1\n",
"config.warmup_steps = config.steps_per_epoch * 5\n",
"config.batch_size = 32\n",
"config"
]
},
{
"cell_type": "markdown",
"id": "f953c9ae",
"metadata": {},
"source": [
"### Training a model that has no variation in any parameter except substructure."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "05203986",
"metadata": {},
"outputs": [],
"source": [
"# Change every parameter but the substructure to static and make the model only predict the substructure\n",
"# parameter.\n",
"input_config['lensing_config']['main_deflector_params']['theta_e'] = input_pipeline.encode_constant(1.1)\n",
"input_config['lensing_config']['main_deflector_params']['slope'] = input_pipeline.encode_constant(2.0)\n",
"input_config['lensing_config']['main_deflector_params']['center_x'] = input_pipeline.encode_constant(0.08)\n",
"input_config['lensing_config']['main_deflector_params']['center_y'] = input_pipeline.encode_constant(-0.16)\n",
"input_config['lensing_config']['main_deflector_params']['axis_ratio'] = input_pipeline.encode_constant(0.9)\n",
"input_config['lensing_config']['main_deflector_params']['angle'] = input_pipeline.encode_constant(0.0)\n",
"input_config['lensing_config']['main_deflector_params']['gamma_ext'] = input_pipeline.encode_constant(0.0)\n",
"input_config['lensing_config']['source_params']['amp'] = input_pipeline.encode_constant(5.0)\n",
"input_config['lensing_config']['source_params']['sersic_radius'] = input_pipeline.encode_constant(1.5)\n",
"input_config['lensing_config']['source_params']['n_sersic'] = input_pipeline.encode_constant(1.5)\n",
"input_config['lensing_config']['source_params']['axis_ratio'] = input_pipeline.encode_constant(0.9)\n",
"input_config['lensing_config']['source_params']['angle'] = input_pipeline.encode_constant(0.0)\n",
"input_config['lensing_config']['source_params']['center_x'] = input_pipeline.encode_constant(-0.08)\n",
"input_config['lensing_config']['source_params']['center_y'] = input_pipeline.encode_constant(0.04)\n",
"input_config['truth_parameters'] = (['subhalo_params'],['sigma_sub'])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "79b64b0c",
"metadata": {},
"outputs": [],
"source": [
"workdir = '/scratch/users/swagnerc/notebook_outputs/no_variation'\n",
"learning_rate = 1e-2\n",
"state = train.train_and_evaluate(config, input_config, workdir, rng, image_size, learning_rate)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4925a46b",
"metadata": {},
"outputs": [],
"source": [
"# Let's take a peak at how we did.\n",
"n_batches = 8\n",
"image, truth = generate_images_and_truth(input_config, cosmology_params, grid_x, grid_y, rng, n_batches)\n",
"state = checkpoints.restore_checkpoint(workdir, state)\n",
"outputs = get_outputs(state, image)\n",
"n_params = len(input_config['truth_parameters'][0])\n",
"x = jnp.linspace(-2, 2, 10)\n",
"for i in range(n_params):\n",
" plt.xlabel('True Normalized ' + input_config['truth_parameters'][1][i])\n",
" plt.ylabel(r'Predicted $\\mu$' + input_config['truth_parameters'][1][i])\n",
" rho = linregress(truth[:,i], outputs[0][:,i]).rvalue\n",
" print('rho for ' + input_config['truth_parameters'][1][i] + ': ' + str(rho))\n",
" plt.plot(truth[:,i], outputs[0][:,i], '.')\n",
" plt.plot(x,x, c='k')\n",
" plt.xlim((-2,2))\n",
" plt.ylim((-2,2))\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "2a79dd8c",
"metadata": {},
"source": [
"### Training a model that has variation in theta_e."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a3f082ff",
"metadata": {},
"outputs": [],
"source": [
"# Varying theta_e should make the problem a bit harder since, to first order, the sigma_sub signal comes from\n",
"# increasing and decreasing the Eisntein radius.\n",
"input_config['lensing_config']['main_deflector_params']['theta_e'] = input_pipeline.encode_normal(mean=1.1, std=0.15)\n",
"input_config['truth_parameters'] = (['main_deflector_params', 'subhalo_params'],['theta_e', 'sigma_sub'])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "62a84344",
"metadata": {},
"outputs": [],
"source": [
"workdir = '/scratch/users/swagnerc/notebook_outputs/theta_e_variation'\n",
"learning_rate = 1e-2\n",
"state = train.train_and_evaluate(config, input_config, workdir, rng, image_size, learning_rate)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "18cb9e65",
"metadata": {},
"outputs": [],
"source": [
"# Let's take a peak at how we did.\n",
"n_batches = 8\n",
"image, truth = generate_images_and_truth(input_config, cosmology_params, grid_x, grid_y, rng, n_batches)\n",
"state = checkpoints.restore_checkpoint(workdir, state)\n",
"outputs = get_outputs(state, image)\n",
"n_params = len(input_config['truth_parameters'][0])\n",
"x = jnp.linspace(-2, 2, 10)\n",
"for i in range(n_params):\n",
" plt.xlabel('True Normalized ' + input_config['truth_parameters'][1][i])\n",
" plt.ylabel(r'Predicted $\\mu$' + input_config['truth_parameters'][1][i])\n",
" rho = linregress(truth[:,i], outputs[0][:,i]).rvalue\n",
" print('rho for ' + input_config['truth_parameters'][1][i] + ': ' + str(rho))\n",
" plt.plot(truth[:,i], outputs[0][:,i], '.')\n",
" plt.plot(x,x, c='k')\n",
" plt.xlim((-2,2))\n",
" plt.ylim((-2,2))\n",
" plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fe26c0da",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}