Skip to content

Commit

Permalink
Merge pull request #33 from ergodicio/epw_iaw_simul
Browse files Browse the repository at this point in the history
Epw iaw simul
  • Loading branch information
almilder authored Oct 24, 2023
2 parents 04fb769 + b4624e7 commit 48f7ed4
Show file tree
Hide file tree
Showing 16 changed files with 362 additions and 222 deletions.
31 changes: 26 additions & 5 deletions configs/1d/defaults.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,12 @@ parameters:
lb: 0.01
ub: 3.
gradient: 0.0
Te_gradient:
val: 0.0
active: False
lb: 0.
ub: 10.
num_grad_points: 1
Ti:
val: .2
active: False
Expand All @@ -47,6 +53,12 @@ parameters:
lb: .001
ub: 10.
gradient: 0.0
ne_gradient:
val: 0.
active: False
lb: 0.
ub: 15.
num_grad_points: 1
ud:
val: 0
active: False
Expand All @@ -67,6 +79,7 @@ parameters:
lb: -100.
ub: -0.5
fe_decrease_strict: True
symmetric: False
blur:
val: []
active: False
Expand All @@ -83,6 +96,10 @@ parameters:

other:
crop_window: 1
ion_window_start: 525
ion_window_end: 528
ele_window_start: 425
ele_window_end: 625
BinWidth: 10
NumBinInRng: 0
TotalNumBin: 1023
Expand All @@ -107,29 +124,33 @@ other:
lam_res_unit: 5
refit: True
refit_thresh: 0.25
calc_sigmas: True
calc_sigmas: False

plotting:
n_sigmas: 3
rolling_std_width: 5
data_cbar_u: 10
data_cbar_u: data
data_cbar_l: 0


data:
launch_data_visualizer: False
launch_data_visualizer: True
shotnum: 101675
shotDay: False
fit_rng:
blue_min: 450
blue_max: 510
red_min: 545
red_max: 650
iaw_min: 350
iaw_max: 352
iaw_min: 525
iaw_max: 528
bgscaleE: 1.0
bgscaleI: 0.1
bgshotmult: 1
ion_loss_scale: 1.0
ele_t0: 0
ion_t0_shift: 0.0
ele_lam_shift: -2.0
probe_beam: P9
dpixel: 2
lineouts:
Expand Down
16 changes: 8 additions & 8 deletions configs/1d/inputs.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,11 @@ parameters:
ub: 3.75
amp3:
val: 1.
active: False
active: True
lb: 0.01
ub: 3.75
lam:
val: 524.5
val: 526.5
active: True
lb: 523.0
ub: 528.0
Expand All @@ -37,7 +37,7 @@ parameters:
ub: 3.0
Z:
val: 6.
active: False
active: True
lb: 1.0
ub: 25.0
A:
Expand All @@ -46,12 +46,12 @@ parameters:

ud:
val: 0.0
active: False
active: True
lb: -2.0
ub: 2.0
Va:
val: 0.0
active: False
active: True
lb: -2.5
ub: 2.5
ne:
Expand Down Expand Up @@ -96,9 +96,9 @@ data:

other:
extraoptions:
load_ion_spec: False
load_ion_spec: True
load_ele_spec: True
fit_IAW: False
fit_IAW: True
fit_EPWb: True
fit_EPWr: True
refit: False
Expand All @@ -107,4 +107,4 @@ other:

mlflow:
experiment: inverse-thomson-scattering
run: simultaneous_fit_test_1
run: simultaneous_fit_test_3
6 changes: 3 additions & 3 deletions env.yml
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ dependencies:
- python=3.9
- pip==22.3.1
- pip:
- jaxlib==0.4.1
- jax==0.4.1
- jaxlib==0.4.14
- jax==0.4.14
- jaxopt
- numpy
- scipy
Expand All @@ -17,7 +17,7 @@ dependencies:
- mlflow
- boto3
- flatten-dict
- dm-haiku==0.0.9
- dm-haiku==0.0.10
- typing-extensions
- optax
- tqdm
Expand Down
13 changes: 5 additions & 8 deletions env_gpu.yml
Original file line number Diff line number Diff line change
@@ -1,20 +1,16 @@
name: ts
channels:
- conda-forge
- "nvidia/label/cuda-11.7.0"
- nvidia
dependencies:
- python=3.10 # match version used by rayproject/ray:nightly-py38-gpu docker image
- ipython
- ipykernel
- cuda-nvcc
- cudatoolkit==11.7.0
- pip
- pip:
# works for regular pip packages
- --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
- jaxlib==0.4.1+cuda11.cudnn82
- jax==0.4.1
- jaxlib==0.4.14+cuda12.cudnn89
- jax==0.4.14
- jaxopt
- numpy
- scipy
Expand All @@ -25,8 +21,9 @@ dependencies:
- mlflow
- boto3
- flatten-dict
- dm-haiku==0.0.9
- dm-haiku==0.0.10
- typing-extensions
- optax
- tqdm
- xarray
- xarray
- mlflow_export_import
4 changes: 2 additions & 2 deletions init.sh
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@

sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended

cp ~/private/init_confs/zshrc ~/.zshrc
cp ~/init_confs/zshrc ~/.zshrc
mkdir ~/.ssh
cp ~/private/init_confs/config ~/.ssh/config
cp ~/init_confs/config ~/.ssh/config

cat << 'EOF' >> ~/.zshrc
export "PATH=$PATH:/opt/conda/bin"
Expand Down
8 changes: 4 additions & 4 deletions inverse_thomson_scattering/misc/calibration.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ def get_calibrations(shotNum, tstype, CCDsize):
stddev["spect_FWHM_ele"] = 0.9 # nominally this is ~.8 or .9 for h2
stddev["spect_stddev_ele"] = stddev["spect_FWHM_ele"] / 2.3548 # dummy
stddev["ang_FWHM_ele"] = 1 # see Joe's FDR slides ~1-1.2
IAWtime = 0 # means nothing here just kept to allow one code to be used for both
#IAWtime = 0 # means nothing here just kept to allow one code to be used for both

elif tstype == "temporal":
if shotNum < 105000:
Expand Down Expand Up @@ -84,7 +84,7 @@ def get_calibrations(shotNum, tstype, CCDsize):
magI = 5 # (ps / px) this is just a rough guess
magE = 5 # (ps / px) this is just a rough guess

IAWtime = 0 # temporal offset between EPW ross and IAW ross (varies shot to shot, can potentially add a fix based off the fiducials)
# IAWtime = 0 # temporal offset between EPW ross and IAW ross (varies shot to shot, can potentially add a fix based off the fiducials)

else:
if shotNum < 105000:
Expand Down Expand Up @@ -118,7 +118,7 @@ def get_calibrations(shotNum, tstype, CCDsize):
EPWtcc = 1024 - 456.1 # 562;
IAWtcc = 1024 - 519 # 469;

IAWtime = 0 # means nothing here just kept to allow one code to be used for both
#IAWtime = 0 # means nothing here just kept to allow one code to be used for both

## Apply calibrations
axisy = np.arange(1, CCDsize[0] + 1)
Expand All @@ -139,7 +139,7 @@ def get_calibrations(shotNum, tstype, CCDsize):
# axisxE = np.vstack(np.loadtxt("files/angsFRED.txt"))
axisxI = np.arange(1, CCDsize[1] + 1)

return axisxE, axisxI, axisyE, axisyI, magE, IAWtime, stddev
return axisxE, axisxI, axisyE, axisyI, magE, stddev


def get_scattering_angles(config):
Expand Down
46 changes: 46 additions & 0 deletions inverse_thomson_scattering/misc/data_visualizer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import numpy as np
import matplotlib.pyplot as plt
import tempfile, mlflow, os


def launch_data_visualizer(elecData, ionData, all_axes, config):
with tempfile.TemporaryDirectory() as td:
#until this can be made interactive this plots all the data regions
if config["other"]["extraoptions"]["load_ion_spec"]:
X, Y = np.meshgrid(all_axes["iaw_x"], all_axes["iaw_y"])

fig, ax = plt.subplots()
ax.pcolormesh(X, Y, ionData,
cmap="gist_ncar",
vmin=np.amin(ionData),
vmax=np.amax(ionData),)
sline, = ax.plot([all_axes["iaw_x"][config["data"]["lineouts"]["start"]], all_axes["iaw_x"][config["data"]["lineouts"]["start"]]], [all_axes["iaw_y"][0], all_axes["iaw_y"][-1]], lw=2, color = 'w')
eline, = ax.plot([all_axes["iaw_x"][config["data"]["lineouts"]["end"]], all_axes["iaw_x"][config["data"]["lineouts"]["end"]]], [all_axes["iaw_y"][0], all_axes["iaw_y"][-1]], lw=2, color = 'w')

lamsline, = ax.plot([all_axes["iaw_x"][0], all_axes["iaw_x"][-1]], [config["data"]["fit_rng"]["iaw_min"], config["data"]["fit_rng"]["iaw_min"]], lw=2, color = 'w', linestyle = '--')
lameline, = ax.plot([all_axes["iaw_x"][0], all_axes["iaw_x"][-1]], [config["data"]["fit_rng"]["iaw_max"], config["data"]["fit_rng"]["iaw_max"]], lw=2, color = 'w', linestyle = '--')
ax.set_xlabel(all_axes["x_label"])
ax.set_ylabel("Wavelength")
fig.savefig(os.path.join(td, "ion_fit_ranges.png"), bbox_inches="tight")


if config["other"]["extraoptions"]["load_ele_spec"]:
X, Y = np.meshgrid(all_axes["epw_x"], all_axes["epw_y"])

fig, ax = plt.subplots()
ax.pcolormesh(X, Y, elecData,
cmap="gist_ncar",
vmin=np.amin(elecData),
vmax=np.amax(elecData),)
sline, = ax.plot([all_axes["epw_x"][config["data"]["lineouts"]["start"]], all_axes["epw_x"][config["data"]["lineouts"]["start"]]], [all_axes["epw_y"][0], all_axes["epw_y"][-1]], lw=2, color = 'w')
eline, = ax.plot([all_axes["epw_x"][config["data"]["lineouts"]["end"]], all_axes["epw_x"][config["data"]["lineouts"]["end"]]], [all_axes["epw_y"][0], all_axes["epw_y"][-1]], lw=2, color = 'w')

lamsline, = ax.plot([all_axes["epw_x"][0], all_axes["epw_x"][-1]], [config["data"]["fit_rng"]["blue_min"], config["data"]["fit_rng"]["blue_min"]], lw=2, color = 'w', linestyle = '--')
lameline, = ax.plot([all_axes["epw_x"][0], all_axes["epw_x"][-1]], [config["data"]["fit_rng"]["blue_max"], config["data"]["fit_rng"]["blue_max"]], lw=2, color = 'w', linestyle = '--')
lamsline, = ax.plot([all_axes["epw_x"][0], all_axes["epw_x"][-1]], [config["data"]["fit_rng"]["red_min"], config["data"]["fit_rng"]["red_min"]], lw=2, color = 'w', linestyle = '--')
lameline, = ax.plot([all_axes["epw_x"][0], all_axes["epw_x"][-1]], [config["data"]["fit_rng"]["red_max"], config["data"]["fit_rng"]["red_max"]], lw=2, color = 'w', linestyle = '--')
ax.set_xlabel(all_axes["x_label"])
ax.set_ylabel("Wavelength")
fig.savefig(os.path.join(td, "electron_fit_ranges.png"), bbox_inches="tight")

mlflow.log_artifacts(td)
37 changes: 37 additions & 0 deletions inverse_thomson_scattering/misc/lineout_plot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import matplotlib.pyplot as plt
import numpy as np
import os

def lineout_plot(sorted_data, sorted_fits, sorted_sqdev, yaxis, s_ind, e_ind, titlestr, filename, td, tag):
if len(sorted_data) == 2:
num_col = 2
else:
num_col = 1

fig, ax = plt.subplots(2, num_col, figsize=(12, 8), squeeze=False, tight_layout=True, sharex=False)
for col in range(num_col):
ax[0][col].plot(
yaxis[col][s_ind[col] : e_ind[col]],
np.squeeze(sorted_data[col][s_ind[col] : e_ind[col]]),
label="Data")
ax[0][col].plot(
yaxis[col][s_ind[col] : e_ind[col]],
np.squeeze(sorted_fits[col][s_ind[col] : e_ind[col]]),
label="Fit")

ax[0][col].set_title(titlestr, fontsize=14)
ax[0][col].set_ylabel("Amp (arb. units)")
ax[0][col].legend(fontsize=14)
ax[0][col].grid()

ax[1][col].plot(
yaxis[col][s_ind[col] : e_ind[col]],
np.squeeze(sorted_sqdev[col][s_ind[col] : e_ind[col]]),
label="Residual",
)
ax[1][col].set_xlabel("Wavelength (nm)")
ax[1][col].set_ylabel("$\chi_i^2$")


fig.savefig(os.path.join(td, tag, filename), bbox_inches="tight")
plt.close(fig)
Loading

0 comments on commit 48f7ed4

Please sign in to comment.