diff --git a/configs/1d/defaults.yaml b/configs/1d/defaults.yaml index 8dec5025..ffde634a 100644 --- a/configs/1d/defaults.yaml +++ b/configs/1d/defaults.yaml @@ -69,12 +69,6 @@ parameters: active: False lb: -100.0 ub: 100.0 - Va: - val: 0.0 - angle: 0.0 - active: False - lb: -20.5 - ub: 20.5 other: expandedions: False diff --git a/configs/1d/inputs.yaml b/configs/1d/inputs.yaml index 5e8ea34b..9f7aa190 100644 --- a/configs/1d/inputs.yaml +++ b/configs/1d/inputs.yaml @@ -58,12 +58,6 @@ parameters: num_grad_points: 1 ub: 15.0 val: 0.0 - Va: - active: false - angle: 0.0 - lb: -20.5 - ub: 20.5 - val: 0.0 amp1: active: true lb: 0.01 @@ -96,21 +90,84 @@ parameters: lb: -20.0 ub: 20.0 val: 0.0 + ion-1: A: active: false val: 40.0 Ti: + active: True + lb: 0.01 + same: false + ub: 1.0 + val: 0.1 + Z: + active: True + lb: 1.0 + ub: 18.0 + val: 8.0 + Va: + active: True + angle: 0.0 + lb: 0.0 + ub: 40.5 + val: 15.0 + fract: + active: false + val: 1.0 + type: active: false + ion: null + + ion-2: + A: + active: false + val: 40.0 + Ti: + active: True lb: 0.01 same: false ub: 3.0 val: 0.2 Z: + active: True + lb: 1.0 + ub: 18.0 + val: 8.0 + Va: + active: True + angle: 0.0 + lb: 0.0 + ub: 40.5 + val: 20.0 + fract: + active: false + val: 1.0 + type: active: false + ion: null + + ion-3: + A: + active: false + val: 40.0 + Ti: + active: True + lb: 0.01 + same: false + ub: 1.0 + val: 0.1 + Z: + active: True lb: 1.0 ub: 18.0 val: 8.0 + Va: + active: True + angle: 0.0 + lb: 0.0 + ub: 40.5 + val: 10.0 fract: active: false val: 1.0 \ No newline at end of file diff --git a/tests/configs/1d-defaults.yaml b/tests/configs/1d-defaults.yaml index e5972ea4..649cab5d 100644 --- a/tests/configs/1d-defaults.yaml +++ b/tests/configs/1d-defaults.yaml @@ -29,6 +29,12 @@ parameters: fract: val: 0.1 active: False + Va: + val: 0.0 + angle: 0.0 + active: False + lb: -20.5 + ub: 20.5 general: amp1: @@ -69,12 +75,6 @@ parameters: active: False lb: -100.0 ub: 100.0 - Va: - val: 0.0 - angle: 0.0 - active: False - lb: -20.5 - ub: 20.5 other: expandedions: False diff --git a/tests/configs/1d-inputs.yaml b/tests/configs/1d-inputs.yaml index aa341be8..8a26ca38 100644 --- a/tests/configs/1d-inputs.yaml +++ b/tests/configs/1d-inputs.yaml @@ -38,6 +38,13 @@ parameters: fract: val: 1. active: False + Va: + val: 0.0 + angle: 0.0 + active: False + lb: -20.5 + ub: 20.5 + general: amp1: val: 1. @@ -77,12 +84,6 @@ parameters: active: False lb: -10.0 ub: 10.0 - Va: - val: 0.0 - angle: 0.0 - active: False - lb: -20.5 - ub: 20.5 other: diff --git a/tests/configs/epw_defaults.yaml b/tests/configs/epw_defaults.yaml index c4f64e8e..a6e98582 100644 --- a/tests/configs/epw_defaults.yaml +++ b/tests/configs/epw_defaults.yaml @@ -37,6 +37,12 @@ parameters: fract: val: 0.1 active: False + Va: + val: 0.0 + angle: 0.0 + active: False + lb: -20.5 + ub: 20.5 general: amp1: @@ -77,12 +83,6 @@ parameters: active: False lb: -100.0 ub: 100.0 - Va: - val: 0.0 - angle: 0.0 - active: False - lb: -20.5 - ub: 20.5 other: expandedions: False diff --git a/tests/configs/epw_inputs.yaml b/tests/configs/epw_inputs.yaml index be2eeae4..053e702c 100644 --- a/tests/configs/epw_inputs.yaml +++ b/tests/configs/epw_inputs.yaml @@ -39,6 +39,13 @@ parameters: val: 1.0 active: False + Va: + val: 0.0 + angle: 0.0 + active: True + lb: -20.5 + ub: 20.5 + general: amp1: val: 1. @@ -78,12 +85,6 @@ parameters: active: False lb: -10.0 ub: 10.0 - Va: - val: 0.0 - angle: 0.0 - active: True - lb: -20.5 - ub: 20.5 data: shotnum: 101675 diff --git a/tests/configs/time_test_defaults.yaml b/tests/configs/time_test_defaults.yaml index 686d5b95..9ed1ee76 100644 --- a/tests/configs/time_test_defaults.yaml +++ b/tests/configs/time_test_defaults.yaml @@ -35,6 +35,12 @@ parameters: fract: val: 0.1 active: False + Va: + val: 0.0 + angle: 0.0 + active: False + lb: -20.5 + ub: 20.5 general: amp1: @@ -75,12 +81,6 @@ parameters: active: False lb: -100.0 ub: 100.0 - Va: - val: 0.0 - angle: 0.0 - active: False - lb: -20.5 - ub: 20.5 other: expandedions: False diff --git a/tests/configs/time_test_inputs.yaml b/tests/configs/time_test_inputs.yaml index 4b7b28ff..f3ae1743 100644 --- a/tests/configs/time_test_inputs.yaml +++ b/tests/configs/time_test_inputs.yaml @@ -37,6 +37,12 @@ parameters: fract: val: 1.0 active: False + Va: + val: 0.0 + angle: 0.0 + active: False + lb: -20.5 + ub: 20.5 general: amp1: val: 1. @@ -76,12 +82,6 @@ parameters: active: False lb: -10.0 ub: 10.0 - Va: - val: 0.0 - angle: 0.0 - active: False - lb: -20.5 - ub: 20.5 other: diff --git a/tsadar/core/modules/ts_params.py b/tsadar/core/modules/ts_params.py index 174c0fdd..f8ec8e5f 100644 --- a/tsadar/core/modules/ts_params.py +++ b/tsadar/core/modules/ts_params.py @@ -133,11 +133,14 @@ def __call__(self): class IonParams(eqx.Module): normed_Ti: Array normed_Z: Array + normed_Va: Array #SB fract: Array Ti_scale: float Ti_shift: float Z_scale: float Z_shift: float + Va_scale: float # SB + Va_shift: float # SB A: int act_funs: Dict[str, Callable] inv_act_funs: Dict[str, Callable] @@ -145,7 +148,7 @@ class IonParams(eqx.Module): def __init__(self, cfg, batch_size, batch=True, activate=False): super().__init__() self.act_funs, self.inv_act_funs = {}, {} - for param in ["Ti", "Z"]: + for param in ["Ti", "Z", "Va"]: #SB setattr(self, param + "_scale", cfg[param]["ub"] - cfg[param]["lb"]) setattr(self, param + "_shift", cfg[param]["lb"]) self.act_funs[param], self.inv_act_funs[param] = get_act_and_inv_act(cfg[param], activate) @@ -159,11 +162,15 @@ def __init__(self, cfg, batch_size, batch=True, activate=False): self.normed_Z = self.inv_act_funs["Z"]( jnp.full(batch_size, (cfg["Z"]["val"] - self.Z_shift) / self.Z_scale) ) + self.normed_Va = self.inv_act_funs["Va"]( + jnp.full(batch_size, (cfg["Va"]["val"] - self.Va_shift) / self.Va_scale) # SB + ) self.A = jnp.full(batch_size, cfg["A"]["val"]) self.fract = self.inv_act_funs["fract"](jnp.full(batch_size, cfg["fract"]["val"])) else: self.normed_Ti = self.inv_act_funs["Ti"]((cfg["Ti"]["val"] - self.Ti_shift) / self.Ti_scale) self.normed_Z = self.inv_act_funs["Z"]((cfg["Z"]["val"] - self.Z_shift) / self.Z_scale) + self.normed_Va = self.inv_act_funs["Va"]((cfg["Va"]["val"] - self.Va_shift) / self.Va_scale) # SB self.A = cfg["A"]["val"] self.fract = float(self.inv_act_funs["fract"](cfg["fract"]["val"])) @@ -177,6 +184,7 @@ def __call__(self): "fract": self.act_funs["fract"](self.fract), "Ti": self.act_funs["Ti"](self.normed_Ti) * self.Ti_scale + self.Ti_shift, "Z": self.act_funs["Z"](self.normed_Z) * self.Z_scale + self.Z_shift, + "Va": self.act_funs["Va"](self.normed_Va) * self.Va_scale + self.Va_shift, } @@ -213,7 +221,7 @@ class GeneralParams(eqx.Module): normed_ne_gradient: Array normed_Te_gradient: Array normed_ud: Array - normed_Va: Array + #normed_Va: Array # SB lam_scale: float lam_shift: float amp1_scale: float @@ -228,8 +236,8 @@ class GeneralParams(eqx.Module): Te_gradient_shift: float ud_scale: float ud_shift: float - Va_scale: float - Va_shift: float + #Va_scale: float # SB + #Va_shift: float # SB act_funs: Dict[str, Callable] def __init__(self, cfg, batch_size: int, batch=True, activate=False): @@ -237,7 +245,7 @@ def __init__(self, cfg, batch_size: int, batch=True, activate=False): # this is all a bit ugly but we use setattr instead of = to be able to use the for loop self.act_funs, inv_act_funs = {}, {} - for param in ["lam", "amp1", "amp2", "amp3", "ne_gradient", "Te_gradient", "ud", "Va"]: + for param in ["lam", "amp1", "amp2", "amp3", "ne_gradient", "Te_gradient", "ud"]: # SB removed Va self.act_funs[param], inv_act_funs[param] = get_act_and_inv_act(cfg[param], activate) setattr(self, param + "_scale", cfg[param]["ub"] - cfg[param]["lb"]) setattr(self, param + "_shift", cfg[param]["lb"]) @@ -245,7 +253,7 @@ def __init__(self, cfg, batch_size: int, batch=True, activate=False): # this is where the linear and nonlinear transformations are applied i.e. # the rescaling and the activation function if batch: - for param in ["lam", "amp1", "amp2", "amp3", "ne_gradient", "Te_gradient", "ud", "Va"]: + for param in ["lam", "amp1", "amp2", "amp3", "ne_gradient", "Te_gradient", "ud"]: # SB removed Va setattr( self, "normed_" + param, @@ -257,7 +265,7 @@ def __init__(self, cfg, batch_size: int, batch=True, activate=False): ), ) else: - for param in ["lam", "amp1", "amp2", "amp3", "ne_gradient", "Te_gradient", "ud", "Va"]: + for param in ["lam", "amp1", "amp2", "amp3", "ne_gradient", "Te_gradient", "ud"]: # SB removed Va setattr( self, "normed_" + param, @@ -281,7 +289,7 @@ def __call__(self): self.act_funs["Te_gradient"](self.normed_Te_gradient) * self.Te_gradient_scale + self.Te_gradient_shift ) unnormed_ud = self.act_funs["ud"](self.normed_ud) * self.ud_scale + self.ud_shift - unnormed_Va = self.act_funs["Va"](self.normed_Va) * self.Va_scale + self.Va_shift + # unnormed_Va = self.act_funs["Va"](self.normed_Va) * self.Va_scale + self.Va_shift # SB return { "lam": unnormed_lam, @@ -291,7 +299,7 @@ def __call__(self): "ne_gradient": unnormed_ne_gradient, "Te_gradient": unnormed_Te_gradient, "ud": unnormed_ud, - "Va": unnormed_Va, + #"Va": unnormed_Va, # SB } @@ -375,7 +383,7 @@ def get_filter_spec(cfg_params: Dict, ts_params: ThomsonParams) -> Dict: if key == "fe": filter_spec = get_distribution_filter_spec(filter_spec, dist_params=_params) else: - nkey = f"normed_{key}" + nkey = f"normed_{key}"# if key!="fract" else f"{key}" # SB treat fractions differently if "ion" in species: filter_spec = eqx.tree_at( lambda tree: getattr(getattr(tree, "ions")[ion_num - 1], nkey), diff --git a/tsadar/core/physics/form_factor.py b/tsadar/core/physics/form_factor.py index c0dc8485..eb015bdb 100644 --- a/tsadar/core/physics/form_factor.py +++ b/tsadar/core/physics/form_factor.py @@ -148,9 +148,11 @@ def __call__(self, params): A = [params[species]["A"] for species in params.keys() if "ion" in species] Z = [params[species]["Z"] for species in params.keys() if "ion" in species] Ti = [params[species]["Ti"] for species in params.keys() if "ion" in species] + Va = [params[species]["Va"] for species in params.keys() if "ion" in species] # SB fract = [params[species]["fract"] for species in params.keys() if "ion" in species] - Va = params["general"]["Va"] * 1e6 # flow velocity in 1e6 cm/s - ud = params["general"]["ud"] * 1e6 # drift velocity in 1e6 cm/s + #Va = params["general"]["Va"] * 1e6 # flow velocity in 1e6 cm/s # commented out by SB + Va = jnp.reshape(jnp.array(Va),[1,1,1,-1]) * 1.0e6 + ud = params["general"]["ud"] * 1.0e6 # drift velocity in 1e6 cm/s fe = params["electron"]["fe"] vx = params["electron"]["v"] @@ -165,10 +167,12 @@ def __call__(self, params): # calculate k and omega vectors omgpe = constants * jnp.sqrt(ne[..., jnp.newaxis, jnp.newaxis]) # plasma frequency Rad/cm omg = self.omgs - omgL + omg = omg[..., jnp.newaxis] ks = jnp.sqrt(self.omgs**2 - omgpe**2) / self.C kL = jnp.sqrt(omgL**2 - omgpe**2) / self.C k = jnp.sqrt(ks**2 + kL**2 - 2 * ks * kL * jnp.cos(sarad)) + k = k[...,jnp.newaxis] # SB kdotv = k * Va omgdop = omg - kdotv @@ -185,29 +189,32 @@ def __call__(self, params): Zbar = jnp.sum(Z * fract) ni = fract * ne[..., jnp.newaxis, jnp.newaxis, jnp.newaxis] / Zbar omgpi = constants * Z * jnp.sqrt(ni * self.Me / Mi) + num_species = fract.shape[3] vTi = jnp.sqrt(jnp.array(Ti) / Mi) # ion thermal velocity - kldi = (vTi / omgpi) * (k[..., jnp.newaxis]) + kldi = (vTi / omgpi) * (k) # SB # ion susceptibilities # finding derivative of plasma dispersion function along xii array - xii = 1.0 / jnp.transpose((jnp.sqrt(2.0) * vTi), [1, 0, 2, 3]) * ((omgdop / k)[..., jnp.newaxis]) + xii = 1.0 / jnp.transpose((jnp.sqrt(2.0) * vTi), [1, 0, 2, 3]) * ((omgdop / k)) #SB # num_ion_pts = jnp.shape(xii) # chiI = jnp.zeros(num_ion_pts) ZpiR = jnp.interp(xii, self.xi2, self.Zpi[0, :], left=xii**-2, right=xii**-2) ZpiI = jnp.interp(xii, self.xi2, self.Zpi[1, :], left=0, right=0) - chiI = jnp.sum(-0.5 / (kldi**2) * (ZpiR + 1j * ZpiI), 3) + #chiI = jnp.sum(-0.5 / (kldi**2) * (ZpiR + 1j * ZpiI), 3) + chiI = -0.5 / (kldi**2) * (ZpiR + 1j * ZpiI) # SB # electron susceptibility # calculating normilized phase velcoity(xi's) for electrons - xie = omgdop / (k * vTe) - ud / vTe + udr = ud - Va # drift velocity w.r.t different ions # SB + xie = omgdop / (k * vTe) - udr / vTe #fe_vphi = jnp.exp(jnp.interp(xie, vx, jnp.log(fe))) fe_vphi=jnp.exp(jnp.apply_along_axis(interp1d,0,jnp.squeeze(xie),vx,jnp.log(jnp.squeeze(fe)),extrap=[-50, -50])).reshape(jnp.shape(xie)) df = jnp.diff(fe_vphi, 1, 1) / jnp.diff(xie, 1, 1) - df = jnp.append(df, jnp.zeros((len(ne), 1, len(self.scattering_angles["sa"]))), 1) + df = jnp.append(df, jnp.zeros((len(ne), 1, len(self.scattering_angles["sa"]),num_species)), 1) chiEI = jnp.pi / (klde**2) * 1j * df @@ -227,19 +234,20 @@ def __call__(self, params): # This line needs to be changed if ion distribution is changed!!! ion_comp_fact = jnp.transpose(fract * Z**2 / Zbar / vTi, [1, 0, 2, 3]) ion_comp = ion_comp_fact * ( - (jnp.abs(chiE[..., jnp.newaxis])) ** 2.0 * jnp.exp(-(xii**2)) / jnp.sqrt(2 * jnp.pi) + (jnp.abs(chiE)) ** 2.0 * jnp.exp(-(xii**2)) / jnp.sqrt(2 * jnp.pi) ) ele_comp = (jnp.abs(1.0 + chiI)) ** 2.0 * fe_vphi / vTe # ele_compE = fe_vphi / vTe # commented because unused - SKW_ion_omg = 1.0 / k[..., jnp.newaxis] * ion_comp / ((jnp.abs(epsilon[..., jnp.newaxis])) ** 2) + SKW_ion_omg = 1.0 / k * ion_comp / ((jnp.abs(epsilon)) ** 2) - SKW_ion_omg = jnp.sum(SKW_ion_omg, 3) + #SKW_ion_omg = jnp.sum(SKW_ion_omg, 3) # commented by SB SKW_ele_omg = 1.0 / k * (ele_comp) / ((jnp.abs(epsilon)) ** 2) # SKW_ele_omgE = 2 * jnp.pi * 1.0 / klde * (ele_compE) / ((jnp.abs(1 + (chiE))) ** 2) * vTe / omgpe # commented because unused PsOmg = (SKW_ion_omg + SKW_ele_omg) * (1 + 2 * omgdop / omgL) * re**2.0 * ne[:, None, None] + PsOmg = jnp.sum(PsOmg,3) # PsOmgE = (SKW_ele_omg) * (1 + 2 * omgdop / omgL) * re**2.0 * jnp.transpose(ne) # commented because unused lams = 2 * jnp.pi * self.C / self.omgs PsLam = PsOmg * 2 * jnp.pi * self.C / lams**2 diff --git a/tsadar/utils/data_handling/calibration.py b/tsadar/utils/data_handling/calibration.py index 64159bcc..d7bd8d6d 100644 --- a/tsadar/utils/data_handling/calibration.py +++ b/tsadar/utils/data_handling/calibration.py @@ -360,7 +360,37 @@ def get_calibrations(shotNum, tstype, t0, CCDsize): # 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 < 104000: + if shotNum == 92522: + EPWDisp = 0.27093 + IAWDisp = 0.0057 + EPWoff = 381.75 # only this shot seems to have shifted + IAWoff = 523.74 + + stddev["spect_stddev_ion"] = 0.028 # needs to be checked + stddev["spect_stddev_ele"] = 1.4365 # needs to be checked + + magI = 2.87 # um / px + magE = 5.13 # um / px + + EPWtcc = 1024 - 456.1 # 562; + IAWtcc = 519#1024 - 519 # 469; + + elif 92522 < shotNum <= 92538: + EPWDisp = 0.27093 + IAWDisp = 0.0057 + EPWoff = 385.256 # needs to be checked + IAWoff = 523.74 + + stddev["spect_stddev_ion"] = 0.028 # needs to be checked + stddev["spect_stddev_ele"] = 1.4365 # needs to be checked + + magI = 2.87 # um / px + magE = 5.13 # um / px + + EPWtcc = 1024 - 456.1 # 562; + IAWtcc = 519#1024 - 519 # 469; + + elif shotNum < 104000: EPWDisp = 0.27093 IAWDisp = 0.00438 EPWoff = 396.256 # needs to be checked