diff --git a/configs/1d/defaults.yaml b/configs/1d/defaults.yaml index 6a7f6c7..299904d 100644 --- a/configs/1d/defaults.yaml +++ b/configs/1d/defaults.yaml @@ -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 @@ -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 @@ -67,6 +79,7 @@ parameters: lb: -100. ub: -0.5 fe_decrease_strict: True + symmetric: False blur: val: [] active: False @@ -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 @@ -107,17 +124,17 @@ 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: @@ -125,11 +142,15 @@ data: 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: diff --git a/configs/1d/inputs.yaml b/configs/1d/inputs.yaml index bbb8fba..810bf08 100644 --- a/configs/1d/inputs.yaml +++ b/configs/1d/inputs.yaml @@ -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 @@ -37,7 +37,7 @@ parameters: ub: 3.0 Z: val: 6. - active: False + active: True lb: 1.0 ub: 25.0 A: @@ -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: @@ -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 @@ -107,4 +107,4 @@ other: mlflow: experiment: inverse-thomson-scattering - run: simultaneous_fit_test_1 \ No newline at end of file + run: simultaneous_fit_test_3 \ No newline at end of file diff --git a/env.yml b/env.yml index 228021b..8a0e7b7 100644 --- a/env.yml +++ b/env.yml @@ -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 @@ -17,7 +17,7 @@ dependencies: - mlflow - boto3 - flatten-dict - - dm-haiku==0.0.9 + - dm-haiku==0.0.10 - typing-extensions - optax - tqdm diff --git a/env_gpu.yml b/env_gpu.yml index c24b7a4..83c1125 100644 --- a/env_gpu.yml +++ b/env_gpu.yml @@ -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 @@ -25,8 +21,9 @@ dependencies: - mlflow - boto3 - flatten-dict - - dm-haiku==0.0.9 + - dm-haiku==0.0.10 - typing-extensions - optax - tqdm - - xarray \ No newline at end of file + - xarray + - mlflow_export_import \ No newline at end of file diff --git a/init.sh b/init.sh index cda0b2a..4ebfa63 100644 --- a/init.sh +++ b/init.sh @@ -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" diff --git a/inverse_thomson_scattering/misc/calibration.py b/inverse_thomson_scattering/misc/calibration.py index 6e35373..555b1ef 100644 --- a/inverse_thomson_scattering/misc/calibration.py +++ b/inverse_thomson_scattering/misc/calibration.py @@ -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: @@ -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: @@ -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) @@ -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): diff --git a/inverse_thomson_scattering/misc/data_visualizer.py b/inverse_thomson_scattering/misc/data_visualizer.py new file mode 100644 index 0000000..83439ad --- /dev/null +++ b/inverse_thomson_scattering/misc/data_visualizer.py @@ -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) \ No newline at end of file diff --git a/inverse_thomson_scattering/misc/lineout_plot.py b/inverse_thomson_scattering/misc/lineout_plot.py new file mode 100644 index 0000000..2df7fcb --- /dev/null +++ b/inverse_thomson_scattering/misc/lineout_plot.py @@ -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) \ No newline at end of file diff --git a/inverse_thomson_scattering/misc/plotters.py b/inverse_thomson_scattering/misc/plotters.py index 62c2caa..02ab780 100644 --- a/inverse_thomson_scattering/misc/plotters.py +++ b/inverse_thomson_scattering/misc/plotters.py @@ -6,6 +6,7 @@ from inverse_thomson_scattering.misc.num_dist_func import get_num_dist_func from inverse_thomson_scattering.model.physics.generate_spectra import get_fit_model +from inverse_thomson_scattering.misc.lineout_plot import lineout_plot def plotinput(config, sa): @@ -65,68 +66,44 @@ def plotinput(config, sa): def model_v_actual( sorted_losses, sorted_data, sorted_fits, num_plots, td, config, loss_inds, yaxis, sorted_sqdev, sorted_red_losses ): - # make plots + + if config["other"]["extraoptions"]["load_ion_spec"] and config["other"]["extraoptions"]["load_ele_spec"]: + ele_s_ind = np.argmin(np.abs(yaxis[0]-config["other"]["ele_window_start"])) + ele_e_ind = np.argmin(np.abs(yaxis[0]-config["other"]["ele_window_end"])) + ion_s_ind = np.argmin(np.abs(yaxis[1]-config["other"]["ion_window_start"])) + ion_e_ind = np.argmin(np.abs(yaxis[1]-config["other"]["ion_window_end"])) + s_ind = [ele_s_ind, ion_s_ind] + e_ind = [ele_e_ind, ion_e_ind] + elif config["other"]["extraoptions"]["load_ion_spec"]: + s_ind = [np.argmin(np.abs(yaxis-config["other"]["ion_window_start"]))] + e_ind = [np.argmin(np.abs(yaxis-config["other"]["ion_window_end"]))] + sorted_data = [sorted_data]#np.array(sorted_data).reshape((1,)+np.shape(sorted_data)) + sorted_fits = [sorted_fits]#np.array(sorted_fits).reshape((1,)+np.shape(sorted_fits)) + sorted_sqdev = [sorted_sqdev]#np.array(sorted_sqdev).reshape((1,)+np.shape(sorted_sqdev)) + + elif config["other"]["extraoptions"]["load_ele_spec"]: + s_ind = [np.argmin(np.abs(yaxis-config["other"]["ele_window_start"]))] + e_ind = [np.argmin(np.abs(yaxis-config["other"]["ele_window_end"]))] + sorted_data = [sorted_data]#np.array(sorted_data).reshape((1,)+np.shape(sorted_data)) + sorted_fits = [sorted_fits]#np.array(sorted_fits).reshape((1,)+np.shape(sorted_fits)) + sorted_sqdev = [sorted_sqdev]#np.array(sorted_sqdev).reshape((1,)+np.shape(sorted_sqdev)) + for i in range(num_plots): + # plot model vs actual titlestr = ( r"|Error|$^2$" + f" = {sorted_losses[i]:.2e}, line out # {config['data']['lineouts']['val'][loss_inds[i]]}" ) - filename = f"loss={sorted_losses[i]:.2e}-reduced_loss={sorted_red_losses[i]:.2e}-lineout={config['data']['lineouts']['val'][loss_inds[-1 - i]]}.png" - fig, ax = plt.subplots(2, 1, figsize=(10, 8), tight_layout=True, sharex=True) - ax[0].plot( - yaxis[config["other"]["crop_window"] : -config["other"]["crop_window"]], - np.squeeze(sorted_data[i, config["other"]["crop_window"] : -config["other"]["crop_window"]]), - label="Data", - ) - ax[0].plot( - yaxis[config["other"]["crop_window"] : -config["other"]["crop_window"]], - np.squeeze(sorted_fits[i, config["other"]["crop_window"] : -config["other"]["crop_window"]]), - label="Fit", - ) - ax[0].set_title(titlestr, fontsize=14) - ax[0].set_ylabel("Amp (arb. units)") - ax[0].legend(fontsize=14) - ax[0].grid() - ax[1].plot( - yaxis[config["other"]["crop_window"] : -config["other"]["crop_window"]], - np.squeeze(sorted_sqdev[i, config["other"]["crop_window"] : -config["other"]["crop_window"]]), - label="Residual", - ) - ax[1].set_xlabel("Wavelength (nm)") - ax[1].set_ylabel("$\chi_i^2$") - fig.savefig(os.path.join(td, "worst", filename), bbox_inches="tight") - plt.close(fig) + filename = f"loss={sorted_losses[i]:.2e}-reduced_loss={sorted_red_losses[i]:.2e}-lineout={config['data']['lineouts']['val'][loss_inds[i]]}.png" + + lineout_plot(np.array(sorted_data)[:,i,:], np.array(sorted_fits)[:,i,:], np.array(sorted_sqdev)[:,i,:], yaxis, s_ind, e_ind, titlestr, filename, td, "worst") titlestr = ( - r"|Error|$^2$" - + f" = {sorted_losses[-1 - i]:.2e}, line out # {config['data']['lineouts']['val'][loss_inds[-1 - i]]}" - ) + r"|Error|$^2$" + f" = {sorted_losses[-1 - i]:.2e}, line out # {config['data']['lineouts']['val'][loss_inds[-1 - i]]}" + ) filename = f"loss={sorted_losses[-1 - i]:.2e}-reduced_loss={sorted_red_losses[-1 - i]:.2e}-lineout={config['data']['lineouts']['val'][loss_inds[-1 - i]]}.png" - fig, ax = plt.subplots(2, 1, figsize=(10, 8), tight_layout=True, sharex=True) - ax[0].plot( - yaxis[config["other"]["crop_window"] : -config["other"]["crop_window"]], - np.squeeze(sorted_data[-1 - i, config["other"]["crop_window"] : -config["other"]["crop_window"]]), - label="Data", - ) - ax[0].plot( - yaxis[config["other"]["crop_window"] : -config["other"]["crop_window"]], - np.squeeze(sorted_fits[-1 - i, config["other"]["crop_window"] : -config["other"]["crop_window"]]), - label="Fit", - ) - ax[0].set_title(titlestr, fontsize=14) - ax[0].set_ylabel("Amp (arb. units)") - ax[0].legend(fontsize=14) - ax[0].grid() - ax[1].plot( - yaxis[config["other"]["crop_window"] : -config["other"]["crop_window"]], - np.squeeze(sorted_sqdev[-1 - i, config["other"]["crop_window"] : -config["other"]["crop_window"]]), - label="Residual", - ) - ax[1].set_xlabel("Wavelength (nm)") - ax[1].set_ylabel("$\chi_i^2$") - fig.savefig(os.path.join(td, "best", filename), bbox_inches="tight") - plt.close(fig) - + + lineout_plot(np.array(sorted_data)[:,-1-i,:], np.array(sorted_fits)[:,-1-i,:], np.array(sorted_sqdev)[:,-1-i,:], yaxis, s_ind, e_ind, titlestr, filename, td, "best") def LinePlots( x, @@ -256,7 +233,7 @@ def ColorPlots( cmap = TScmap() if logplot: - C = log(C) + C = np.log(C) im = ax0.imshow( C, diff --git a/inverse_thomson_scattering/model/loss_function.py b/inverse_thomson_scattering/model/loss_function.py index 5809c28..83f48dd 100644 --- a/inverse_thomson_scattering/model/loss_function.py +++ b/inverse_thomson_scattering/model/loss_function.py @@ -249,7 +249,7 @@ def __loss__(self, weights, batch): # other_losses = calc_other_losses(params) normed_batch = self.get_normed_batch(batch) normed_e_data = normed_batch["e_data"] - return i_error + e_error + density_loss + temperature_loss + momentum_loss, [ThryE, normed_e_data, params] + return self.cfg["data"]["ion_loss_scale"]*i_error + e_error + density_loss + temperature_loss + momentum_loss, [ThryE, normed_e_data, params] def initialize_rest_of_params(self, config): # print("in initialize_rest_of_params") diff --git a/inverse_thomson_scattering/model/physics/generate_spectra.py b/inverse_thomson_scattering/model/physics/generate_spectra.py index 380a3f6..a88d677 100644 --- a/inverse_thomson_scattering/model/physics/generate_spectra.py +++ b/inverse_thomson_scattering/model/physics/generate_spectra.py @@ -80,7 +80,7 @@ def fit_model(fitted_params): parameters["ud"]["val"], sa["sa"], (fecur, vcur), - lam, + lam+config["data"]["ele_lam_shift"], ) # if parameters.fe['Type']=='MYDLM': diff --git a/inverse_thomson_scattering/process/postprocess.py b/inverse_thomson_scattering/process/postprocess.py index 53c9473..64a1fec 100644 --- a/inverse_thomson_scattering/process/postprocess.py +++ b/inverse_thomson_scattering/process/postprocess.py @@ -393,46 +393,8 @@ def plot_regular(config, losses, all_params, used_points, all_axes, fits, all_da mlflow.log_metrics( {"number of fits above threshold after refit": int(np.sum(red_losses > config["other"]["refit_thresh"]))} ) - - # this wont work for ion+electron fitting (just electrons will be plotted) - if config["other"]["extraoptions"]["load_ion_spec"]: - coords = (all_axes["x_label"], np.array(all_axes["iaw_x"][config["data"]["lineouts"]["val"]])), ( - "Wavelength", - all_axes["iaw_y"], - ) - # print(coords) - # print(all_axes["x_label"]) - dat = {"fit": fits["ion"], "data": all_data["i_data"]} - sorted_fits = fits["ion"][loss_inds] - sorted_data = all_data["i_data"][loss_inds] - sorted_sqdev = sqdevs["ion"][loss_inds] - y_axis = all_axes["iaw_y"] - if config["other"]["extraoptions"]["load_ele_spec"]: - coords = (all_axes["x_label"], np.array(all_axes["epw_x"][config["data"]["lineouts"]["val"]])), ( - "Wavelength", - all_axes["epw_y"], - ) - dat = {"fit": fits["ele"], "data": all_data["e_data"]} - sorted_fits = fits["ele"][loss_inds] - sorted_data = all_data["e_data"][loss_inds] - sorted_sqdev = sqdevs["ele"][loss_inds] - y_axis = all_axes["epw_y"] - - # fit vs data storage and plot - savedata = xr.Dataset({k: xr.DataArray(v, coords=coords) for k, v in dat.items()}) - savedata.to_netcdf(os.path.join(td, "fit_and_data.nc")) - - fig, ax = plt.subplots(1, 2, figsize=(12, 5), tight_layout=True) - clevs = np.linspace( - np.amin(savedata["data"]) if config["plotting"]["data_cbar_l"] == "data" else config["plotting"]["data_cbar_l"], - np.amax(savedata["data"]) if config["plotting"]["data_cbar_u"] == "data" else config["plotting"]["data_cbar_u"], - 11, - ) - # clevs = np.linspace(0, 300, 11) - savedata["fit"].T.plot(ax=ax[0], cmap="gist_ncar", levels=clevs) - savedata["data"].T.plot(ax=ax[1], cmap="gist_ncar", levels=clevs) - fig.savefig(os.path.join(td, "fit_and_data.png"), bbox_inches="tight") - + + #make histogram fig, ax = plt.subplots(1, 2, figsize=(12, 4), tight_layout=True) if "red_losses_init" not in locals(): red_losses_init = red_losses @@ -465,19 +427,118 @@ def plot_regular(config, losses, all_params, used_points, all_axes, fits, all_da os.makedirs(os.path.join(td, "worst")) os.makedirs(os.path.join(td, "best")) + + if config["other"]["extraoptions"]["load_ion_spec"]: + coords = (all_axes["x_label"], np.array(all_axes["iaw_x"][config["data"]["lineouts"]["val"]])), ( + "Wavelength", + all_axes["iaw_y"], + ) + # print(coords) + # print(all_axes["x_label"]) + ion_dat = {"fit": fits["ion"], "data": all_data["i_data"]} + ion_sorted_fits = fits["ion"][loss_inds] + ion_sorted_data = all_data["i_data"][loss_inds] + ion_sorted_sqdev = sqdevs["ion"][loss_inds] + ion_y_axis = all_axes["iaw_y"] + # fit vs data storage and plot + ion_savedata = xr.Dataset({k: xr.DataArray(v, coords=coords) for k, v in ion_dat.items()}) + ion_savedata.to_netcdf(os.path.join(td, "ion_fit_and_data.nc")) + if config["other"]["extraoptions"]["load_ele_spec"]: + coords = (all_axes["x_label"], np.array(all_axes["epw_x"][config["data"]["lineouts"]["val"]])), ( + "Wavelength", + all_axes["epw_y"], + ) + ele_dat = {"fit": fits["ele"], "data": all_data["e_data"]} + ele_sorted_fits = fits["ele"][loss_inds] + ele_sorted_data = all_data["e_data"][loss_inds] + ele_sorted_sqdev = sqdevs["ele"][loss_inds] + ele_y_axis = all_axes["epw_y"] + # fit vs data storage and plot + ele_savedata = xr.Dataset({k: xr.DataArray(v, coords=coords) for k, v in ele_dat.items()}) + ele_savedata.to_netcdf(os.path.join(td, "ele_fit_and_data.nc")) + + if config["other"]["extraoptions"]["load_ion_spec"] and config["other"]["extraoptions"]["load_ele_spec"]: + fig, ax = plt.subplots(2, 2, figsize=(12, 12), tight_layout=True) + ion_clevs = np.linspace( + np.amin(ion_dat["data"]) if config["plotting"]["data_cbar_l"] == "data" else config["plotting"]["data_cbar_l"], + np.amax(ion_dat["data"]) if config["plotting"]["data_cbar_u"] == "data" else config["plotting"]["data_cbar_u"], + 11, + ) + ele_clevs = np.linspace( + np.amin(ele_dat["data"]) if config["plotting"]["data_cbar_l"] == "data" else config["plotting"]["data_cbar_l"], + np.amax(ele_dat["data"]) if config["plotting"]["data_cbar_u"] == "data" else config["plotting"]["data_cbar_u"], + 11, + ) + # clevs = np.linspace(0, 300, 11) + ele_savedata["fit"].T.plot(ax=ax[0][0], cmap="gist_ncar", levels=ele_clevs) + ele_savedata["data"].T.plot(ax=ax[0][1], cmap="gist_ncar", levels=ele_clevs) + ion_savedata["fit"].T.plot(ax=ax[1][0], cmap="gist_ncar", levels=ion_clevs) + ion_savedata["data"].T.plot(ax=ax[1][1], cmap="gist_ncar", levels=ion_clevs) + fig.savefig(os.path.join(td, "fit_and_data.png"), bbox_inches="tight") + + plotters.model_v_actual( + sorted_losses, + [ele_sorted_data, ion_sorted_data], + [ele_sorted_fits, ion_sorted_fits], + num_plots, + td, + config, + loss_inds, + [ele_y_axis, ion_y_axis], + [ele_sorted_sqdev, ion_sorted_sqdev], + sorted_redchi, + ) + + elif config["other"]["extraoptions"]["load_ion_spec"]: + fig, ax = plt.subplots(1, 2, figsize=(12, 5), tight_layout=True) + clevs = np.linspace( + np.amin(ion_savedata["data"]) if config["plotting"]["data_cbar_l"] == "data" else config["plotting"]["data_cbar_l"], + np.amax(ion_savedata["data"]) if config["plotting"]["data_cbar_u"] == "data" else config["plotting"]["data_cbar_u"], + 11, + ) + # clevs = np.linspace(0, 300, 11) + ion_savedata["fit"].T.plot(ax=ax[0], cmap="gist_ncar", levels=clevs) + ion_savedata["data"].T.plot(ax=ax[1], cmap="gist_ncar", levels=clevs) + fig.savefig(os.path.join(td, "fit_and_data.png"), bbox_inches="tight") + + plotters.model_v_actual( + sorted_losses, + ion_sorted_data, + ion_sorted_fits, + num_plots, + td, + config, + loss_inds, + np.array([ion_y_axis]), + ion_sorted_sqdev, + sorted_redchi, + ) + + elif config["other"]["extraoptions"]["load_ele_spec"]: + fig, ax = plt.subplots(1, 2, figsize=(12, 5), tight_layout=True) + clevs = np.linspace( + np.amin(ele_savedata["data"]) if config["plotting"]["data_cbar_l"] == "data" else config["plotting"]["data_cbar_l"], + np.amax(ele_savedata["data"]) if config["plotting"]["data_cbar_u"] == "data" else config["plotting"]["data_cbar_u"], + 11, + ) + # clevs = np.linspace(0, 300, 11) + ele_savedata["fit"].T.plot(ax=ax[0], cmap="gist_ncar", levels=clevs) + ele_savedata["data"].T.plot(ax=ax[1], cmap="gist_ncar", levels=clevs) + fig.savefig(os.path.join(td, "fit_and_data.png"), bbox_inches="tight") + + plotters.model_v_actual( + sorted_losses, + ele_sorted_data, + ele_sorted_fits, + num_plots, + td, + config, + loss_inds, + np.array([ele_y_axis]), + ele_sorted_sqdev, + sorted_redchi, + ) - plotters.model_v_actual( - sorted_losses, - sorted_data, - sorted_fits, - num_plots, - td, - config, - loss_inds, - y_axis, - sorted_sqdev, - sorted_redchi, - ) sigmas_ds = xr.Dataset( {k: xr.DataArray(sigmas[:, i], coords=(coords[0],)) for i, k in enumerate(all_params.keys())} @@ -509,78 +570,5 @@ def plot_regular(config, losses, all_params, used_points, all_axes, fits, all_da ax.set_ylim(0.8 * np.min(final_params[param]), 1.2 * np.max(final_params[param])) ax.set_ylabel(param, fontsize=14) fig.savefig(os.path.join(td, "learned_" + param + ".png"), bbox_inches="tight") - # ne_vals = pandas.Series(final_params["ne"]) - # ne_std = ne_vals.rolling(5, min_periods=1, center=True).std() - # Te_vals = pandas.Series(final_params["Te"]) - # Te_std = Te_vals.rolling(5, min_periods=1, center=True).std() - # m_vals = pandas.Series(final_params["m"]) - # m_std = m_vals.rolling(5, min_periods=1, center=True).std() - - # print(final_params) - # print(sigmas_ds) - # ne, Te, m plots with errorbars - # fig, ax = plt.subplots(1, 3, figsize=(15, 4)) - # lineouts = np.array(config["data"]["lineouts"]["val"]) - # ax[0].plot(lineouts, final_params["ne"]) - # ax[0].fill_between( - # lineouts, - # (final_params["ne"] - 3 * sigmas_ds.ne), - # (final_params["ne"] + 3 * sigmas_ds.ne), - # color="b", - # alpha=0.1, - # ) - # ax[0].fill_between( - # lineouts, - # (final_params["ne"] - 3 * ne_std.values), - # (final_params["ne"] + 3 * ne_std.values), - # color="r", - # alpha=0.1, - # ) - # ax[0].set_xlabel("lineout", fontsize=14) - # ax[0].grid() - # ax[0].set_ylim(0.8 * np.min(final_params["ne"]), 1.2 * np.max(final_params["ne"])) - # ax[0].set_title("$n_e$(t)", fontsize=14) - - # ax[1].plot(lineouts, final_params["Te"]) - # ax[1].fill_between( - # lineouts, - # (final_params["Te"] - 3 * sigmas_ds.Te), - # (final_params["Te"] + 3 * sigmas_ds.Te), - # color="b", - # alpha=0.1, - # ) - # ax[1].fill_between( - # lineouts, - # (final_params["Te"] - 3 * Te_std.values), - # (final_params["Te"] + 3 * Te_std.values), - # color="r", - # alpha=0.1, - # ) - # ax[1].set_xlabel("lineout", fontsize=14) - # ax[1].grid() - # ax[1].set_ylim(0.8 * np.min(final_params["Te"]), 1.2 * np.max(final_params["Te"])) - # ax[1].set_title("$T_e$(t)", fontsize=14) - - # ax[2].plot(lineouts, final_params["m"]) - # ax[2].fill_between( - # lineouts, - # (final_params["m"] - 3 * sigmas_ds.m), - # (final_params["m"] + 3 * sigmas_ds.m), - # color="b", - # alpha=0.1, - # ) - # ax[2].fill_between( - # lineouts, - # (final_params["m"] - 3 * m_std.values), - # (final_params["m"] + 3 * m_std.values), - # color="r", - # alpha=0.1, - # ) - # ax[2].set_xlabel("lineout", fontsize=14) - # ax[2].grid() - # ax[2].set_ylim(0.8 * np.min(final_params["m"]), 1.2 * np.max(final_params["m"])) - # ax[2].set_title("$m$(t)", fontsize=14) - - # fig.savefig(os.path.join(td, "learned_parameters.png"), bbox_inches="tight") return final_params diff --git a/inverse_thomson_scattering/process/prepare.py b/inverse_thomson_scattering/process/prepare.py index 21483b0..8a95a03 100644 --- a/inverse_thomson_scattering/process/prepare.py +++ b/inverse_thomson_scattering/process/prepare.py @@ -8,6 +8,7 @@ from inverse_thomson_scattering.process.correct_throughput import correctThroughput from inverse_thomson_scattering.misc.calibration import get_calibrations, get_scattering_angles from inverse_thomson_scattering.process.lineouts import get_lineouts +from inverse_thomson_scattering.misc.data_visualizer import launch_data_visualizer def prepare_data(config: Dict) -> Dict: @@ -29,7 +30,7 @@ def prepare_data(config: Dict) -> Dict: sa = get_scattering_angles(config) # Calibrate axes - [axisxE, axisxI, axisyE, axisyI, magE, IAWtime, stddev] = get_calibrations( + [axisxE, axisxI, axisyE, axisyI, magE, stddev] = get_calibrations( config["data"]["shotnum"], config["other"]["extraoptions"]["spectype"], config["other"]["CCDsize"] ) all_axes = {"epw_x": axisxE, "epw_y": axisyE, "iaw_x": axisxI, "iaw_y": axisyI, "x_label": xlab} @@ -49,10 +50,9 @@ def prepare_data(config: Dict) -> Dict: elecData, config["other"]["extraoptions"]["spectype"], axisyE, config["data"]["shotnum"] ) - # Not possible on JupyterLab looking for another solution - # Lauch the data visualizer for manual linout selection - # if config["data"]["launch_data_visualizer"]: - # launch_data_visualizer(elecData, ionData, all_axes, config) + # Lauch the data visualizer to show linout selection, not currently interactable + if config["data"]["launch_data_visualizer"]: + launch_data_visualizer(elecData, ionData, all_axes, config) # load and correct background [BGele, BGion] = get_shot_bg(config, axisyE, elecData) @@ -115,7 +115,8 @@ def prepare_data(config: Dict) -> Dict: else: all_data = get_lineouts( - elecData, ionData, BGele, BGion, axisxE, axisxI, axisyE, axisyI, 0, IAWtime, xlab, sa, config + elecData, ionData, BGele, BGion, axisxE, axisxI, axisyE, axisyI, config["data"]["ele_t0"], + config["data"]["ion_t0_shift"], xlab, sa, config ) config["other"]["PhysParams"]["widIRF"] = stddev diff --git a/queue_job.sh b/queue_job.sh new file mode 100644 index 0000000..fffb587 --- /dev/null +++ b/queue_job.sh @@ -0,0 +1,24 @@ +#!/bin/bash +#SBATCH -A m4434_g +#SBATCH -C gpu +#SBATCH -q shared +#SBATCH -t 2:00:00 +#SBATCH -n 1 +#SBATCH -c 32 +#SBATCH --gpus-per-task=1 + +export SLURM_CPU_BIND="cores" +export AWS_ACCESS_KEY_ID="AKIARRO7HXV5GIK7G6F3" +export AWS_SECRET_ACCESS_KEY="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" +export BASE_TEMPDIR="$PSCRATCH/tmp/" +export MLFLOW_TRACKING_URI="$PSCRATCH/mlflow" + +# copy job stuff over + +module load cudnn/8.9.3_cuda12.lua +module load cudatoolkit/12.0.lua +module load conda +conda activate ts + +cd /global/homes/a/amilder/inverse-thomson-scattering +srun python3 run.py diff --git a/run.py b/run.py index ed06a62..1050736 100644 --- a/run.py +++ b/run.py @@ -3,6 +3,9 @@ import mlflow from flatten_dict import flatten, unflatten from jax import config +from mlflow_export_import.run.export_run import RunExporter +from tqdm import tqdm +import boto3 config.update("jax_enable_x64", True) # config.update("jax_disable_jit", True) @@ -22,10 +25,29 @@ def one_run(config): return loss +def upload_dir_to_s3(local_directory: str, bucket: str, destination: str, run_id: str): + client = boto3.client("s3") + + # enumerate local files recursively + for root, dirs, files in os.walk(local_directory): + for filename in tqdm(files): + # construct the full local path + local_path = os.path.join(root, filename) + + # construct the full Dropbox path + relative_path = os.path.relpath(local_path, local_directory) + s3_path = os.path.join(destination, relative_path) + client.upload_file(local_path, bucket, s3_path) + + with open(os.path.join(local_directory, f"ingest-{run_id}.txt"), "w") as fi: + fi.write("ready") + + client.upload_file(os.path.join(local_directory, f"ingest-{run_id}.txt"), bucket, f"ingest-{run_id}.txt") + if __name__ == "__main__": all_configs = {} - basedir = os.path.join(os.getcwd(), "configs", "arts") + basedir = os.path.join(os.getcwd(), "configs", "1d") for k in ["defaults", "inputs"]: with open(f"{os.path.join(basedir, k)}.yaml", "r") as fi: all_configs[k] = yaml.safe_load(fi) @@ -53,3 +75,12 @@ def one_run(config): config = unflatten(defaults) one_run(config) + + t0 = time.time() + run_exp = RunExporter(mlflow_client=mlflow.MlflowClient()) + with tempfile.TemporaryDirectory() as td2: + run_exp.export_run(run.info.run_id, td2) + print(f"Export took {round(time.time() - t0, 2)} s") + t0 = time.time() + upload_dir_to_s3(td2, "remote-mlflow-staging", f"artifacts/{run.info.run_id}", run.info.run_id) + print(f"Uploading took {round(time.time() - t0, 2)} s") \ No newline at end of file diff --git a/tests/configs/defaults.yaml b/tests/configs/defaults.yaml index d40ebb4..3d65523 100644 --- a/tests/configs/defaults.yaml +++ b/tests/configs/defaults.yaml @@ -96,13 +96,15 @@ parameters: other: crop_window: 256 - active: False + ion_window_start: 525 + ion_window_end: 528 + ele_window_start: 425 + ele_window_end: 625 BinWidth: 10 NumBinInRng: 0 TotalNumBin: 1023 expandedions: False extraoptions: - spectype: 2 load_ion_spec: False load_ele_spec: True fit_IAW: False @@ -117,10 +119,16 @@ other: flatbg: 0 gain: 1 points_per_pixel: 5 + ang_res_unit: 10 + lam_res_unit: 5 + refit: True + refit_thresh: 0.25 + calc_sigmas: False data: shotnum: 101675 shotDay: False + launch_data_visualizer: False fit_rng: blue_min: 450 blue_max: 510 @@ -129,6 +137,10 @@ data: 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 dpixel: 2 lineouts: type: @@ -161,13 +173,19 @@ optimizer: num_epochs: 200 learning_rate: 1.0e-4 parameter_norm: True + refine_factor: 0 + num_mins: 1 nn: use: false conv_filters: 32|32|16 linear_widths: 16|8 - +dist_fit: + window: + len: 0.2 #should be even + type: hamming # one of [hamming, hann, bartlett] + mlflow: experiment: inverse-thomson-scattering run: base \ No newline at end of file