def ask_numeric_eigenstate_basis(parameters, *, spec_type):
numeric_basis_q = si.ask_for_bool("Use numeric eigenstate basis?",
default=True)
if numeric_basis_q:
parameters.append(
si.Parameter(name="use_numeric_eigenstates", value=True))
max_energy = u.eV * si.ask_for_input(
"Numeric Eigenstate Max Energy (in eV)?",
default=20,
callback=float)
parameters.append(
si.Parameter(name="numeric_eigenstate_max_energy",
value=max_energy))
if spec_type == ion.mesh.SphericalHarmonicSpecification:
max_angular_momentum = si.ask_for_input(
"Numeric Eigenstate Maximum l?", default=20, callback=int)
parameters.append(
si.Parameter(
name="numeric_eigenstate_max_angular_momentum",
value=max_angular_momentum,
))
return max_energy, max_angular_momentum
else:
return max_energy, None
def test_expandable_and_fixed_parameter():
a = si.Parameter("a", [0, 1])
b = si.Parameter("b", [True, False], expandable=True)
parameters = [a, b]
d = si.expand_parameters(parameters)
assert d == [{"a": [0, 1], "b": True}, {"a": [0, 1], "b": False}]
def test_two_fixed_parameters():
a = si.Parameter("a", [0, 1])
b = si.Parameter("b", [True, False])
parameters = [a, b]
d = si.expand_parameters(parameters)
assert d == [{"a": [0, 1], "b": [True, False]}]
def construct_pulses(parameters, *, time_initial_in_pw, time_final_in_pw):
pulse_parameters = []
pulse_type = PULSE_NAMES_TO_TYPES[si.ask_for_input(
"Pulse Type? [sinc | gaussian | sech | cos2]", default="sinc")]
constructor_names = (name.replace("from_", "")
for name in pulse_type.__dict__
if name.startswith("from_"))
constructor_name = si.ask_for_input(
f'Pulse Constructor? [{" | ".join(constructor_names)}]',
default="omega_min")
constructor = getattr(pulse_type, f"from_{constructor_name}")
constructor_argspec = inspect.getfullargspec(constructor)
if (constructor_argspec.varargs
is not None): # alias for default constructor, super implicit....
constructor_args = inspect.getfullargspec(pulse_type.__init__).args
else:
constructor_args = constructor_argspec.args
asks = (CONSTRUCTOR_ARG_TO_ASK[arg] for arg in CONSTRUCTOR_ARG_TO_ASK
if arg in constructor_args)
for ask in asks:
ask(pulse_parameters)
window_time_in_pw, window_width_in_pw = ask_pulse_window(
pulse_type=pulse_type,
time_initial_in_pw=time_initial_in_pw,
time_final_in_pw=time_final_in_pw,
)
print("Generating pulses...")
pulses = tuple(
constructor(
**d,
window=ion.potentials.LogisticWindow(
window_time=d["pulse_width"] * window_time_in_pw,
window_width=d["pulse_width"] * window_width_in_pw,
),
) for d in tqdm(si.expand_parameters(pulse_parameters), ascii=True))
parameters.append(
si.Parameter(name="electric_potential", value=pulses, expandable=True))
parameters.append(
si.Parameter(
name="electric_potential_dc_correction",
value=si.ask_for_bool("Perform Electric Field DC Correction?",
default=True),
))
parameters.append(
si.Parameter(
name="electric_potential_fluence_correction",
value=si.ask_for_bool("Perform Electric Field Fluence Correction?",
default=False),
))
return pulse_parameters
def test_two_expandable_parameters_with_sized_values():
a = si.Parameter("a", [0, 1], expandable=True)
b = si.Parameter("b", [True, False], expandable=True)
parameters = [a, b]
d = si.expand_parameters(parameters)
assert d == [
{"a": 0, "b": True},
{"a": 0, "b": False},
{"a": 1, "b": True},
{"a": 1, "b": False},
]
def test_expandable_parameters_with_iterator():
a = si.Parameter("a", range(2), expandable=True)
b = si.Parameter("b", [True, False], expandable=True)
parameters = [a, b]
d = si.expand_parameters(parameters)
assert d == [
{"a": 0, "b": True},
{"a": 0, "b": False},
{"a": 1, "b": True},
{"a": 1, "b": False},
]
def ask_ide_tunneling(parameters):
choices = {
cls.__name__.replace("Rate", ""): cls
for cls in ion.tunneling.TUNNELING_MODEL_TYPES
}
model_key = si.ask_for_input(f'Tunneling Model? [{" | ".join(choices)}]',
default=tuple(choices.keys())[0])
try:
cls = choices[model_key]
except KeyError:
raise ion.exceptions.InvalidChoice(
f"{model_key} is not one of {choices.keys()}")
argspec = inspect.getfullargspec(cls.__init__)
arg_names = argspec.args[1:]
arg_defaults = argspec.defaults
if len(arg_names) > 0:
args = {
name: si.ask_for_eval(f"Value for {name}?", default=repr(default))
for name, default in reversed(
tuple(
itertools.zip_longest(reversed(arg_names),
reversed(arg_defaults))))
}
model = cls(**args)
else:
model = cls()
parameters.append(si.Parameter(name="tunneling_model", value=model))
return model
def ask_mesh_operators(parameters, *, spec_type):
choices_by_spec_type = {
ion.mesh.LineSpecification: {
ion.Gauge.LENGTH.value: ion.mesh.LineLengthGaugeOperators,
ion.Gauge.VELOCITY.value: ion.mesh.LineVelocityGaugeOperators,
},
ion.mesh.CylindricalSliceSpecification: {
ion.Gauge.LENGTH.value:
ion.mesh.CylindricalSliceLengthGaugeOperators
},
ion.mesh.SphericalSliceSpecification: {
ion.Gauge.LENGTH.value: ion.mesh.SphericalSliceLengthGaugeOperators
},
ion.mesh.SphericalHarmonicSpecification: {
ion.Gauge.LENGTH.value:
ion.mesh.SphericalHarmonicLengthGaugeOperators,
ion.Gauge.VELOCITY.value:
ion.mesh.SphericalHarmonicVelocityGaugeOperators,
},
}
choices = choices_by_spec_type[spec_type]
choice = si.ask_for_input(
f'Mesh Operators? [{" | ".join(choices.keys())}]',
default=ion.Gauge.LENGTH.value,
)
try:
method = choices[choice]()
except KeyError:
raise ion.exceptions.InvalidChoice(f"{choice} is not one of {choices}")
parameters.append(si.Parameter(name="evolution_method", value=method))
return method
def ask_time_step(parameters):
parameters.append(
si.Parameter(
name="time_step",
value=u.asec *
si.ask_for_input("Time Step (in as)?", default=1, callback=float),
))
def ask_checkpoints(parameters):
do_checkpoints = si.ask_for_bool("Checkpoints?", default=True)
parameters.append(si.Parameter(name="checkpoints", value=do_checkpoints))
if do_checkpoints:
time_between_checkpoints = si.ask_for_input(
"How long between checkpoints (in minutes)?",
default=60,
callback=int)
parameters.append(
si.Parameter(
name="checkpoint_every",
value=datetime.timedelta(minutes=time_between_checkpoints),
))
return do_checkpoints
def ask_pulse_number_of_cycles(pulse_parameters):
number_of_cycles = si.Parameter(
name="number_of_cycles",
value=np.array(
si.ask_for_eval("Number of Cycles?", default="[2, 3, 4]")),
expandable=True,
)
pulse_parameters.append(number_of_cycles)
def ask_data_storage_ide(parameters, *, spec_type):
parameters.append(
si.Parameter(
name="store_data_every",
value=si.ask_for_input("Store Data Every n Time Steps",
default=-1,
callback=int),
))
def ask_pulse_number_of_pulse_widths(pulse_parameters):
number_of_pulse_widths = si.Parameter(
name="number_of_pulse_widths",
value=np.array(
si.ask_for_eval("Number of Pulse Widths to count Cycles over?",
default="[3]")),
expandable=True,
)
pulse_parameters.append(number_of_pulse_widths)
def ask_pulse_amplitudes(pulse_parameters):
amplitude_prefactors = si.Parameter(
name="amplitude",
value=u.atomic_electric_field * np.array(
si.ask_for_eval("Pulse Amplitudes? (in AEF)",
default="[.01, .05, .1, .5, 1, 2]")),
expandable=True,
)
pulse_parameters.append(amplitude_prefactors)
def ask_pulse_omega_mins(pulse_parameters):
omega_mins = si.Parameter(
name="omega_min",
value=u.twopi * u.THz * np.array(
si.ask_for_eval("Pulse Frequency Minimum? (in THz)",
default="[30]")),
expandable=True,
)
pulse_parameters.append(omega_mins)
def ask_pulse_widths(pulse_parameters):
pulse_width = si.Parameter(
name="pulse_width",
value=u.asec * np.array(
si.ask_for_eval("Pulse Widths (in as)?",
default="[50, 100, 200, 400, 800]")),
expandable=True,
)
pulse_parameters.append(pulse_width)
def ask_pulse_fluences(pulse_parameters):
fluence = si.Parameter(
name="fluence",
value=u.Jcm2 * np.array(
si.ask_for_eval("Pulse Fluence (in J/cm^2)?",
default="[.01, .1, 1, 10, 20]")),
expandable=True,
)
pulse_parameters.append(fluence)
def ask_pulse_phases(pulse_parameters):
phases = si.Parameter(
name="phase",
value=np.array(
si.ask_for_eval("Pulse CEP (in rad)?",
default="[0, u.pi / 4, u.pi / 2]")),
expandable=True,
)
pulse_parameters.append(phases)
def ask_initial_state_for_hydrogen_sim(parameters):
initial_state = si.Parameter(
name="initial_state",
value=ion.states.HydrogenBoundState(
n=si.ask_for_input("Initial State n?", default=1, callback=int),
l=si.ask_for_input("Initial State l?", default=0, callback=int),
),
)
parameters.append(initial_state)
return initial_state
def ask_evolution_method_tdse(parameters):
choices = {
"ADI": ion.mesh.AlternatingDirectionImplicit,
"SO": ion.mesh.SplitInteractionOperator,
}
key = si.ask_for_input(f'Evolution Method? [{" | ".join(choices)}]',
default="SO")
try:
method = choices[key]()
except KeyError:
raise ion.exceptions.InvalidChoice(f"{method} is not one of {choices}")
parameters.append(si.Parameter(name="evolution_method", value=method))
return method
def ask_evolution_method_ide(parameters):
choices = {
"FE": ion.ide.ForwardEulerMethod,
"BE": ion.ide.BackwardEulerMethod,
"TRAP": ion.ide.TrapezoidMethod,
"RK4": ion.ide.RungeKuttaFourMethod,
}
key = si.ask_for_input(f'Evolution Method? [{" | ".join(choices.keys())}]',
default="RK4")
try:
method = choices[key]()
except KeyError:
raise ion.exceptions.InvalidChoice(f"{method} is not one of {choices}")
parameters.append(si.Parameter(name="evolution_method", value=method))
return method
def ask_ide_kernel(parameters):
choices = {
"hydrogen":
ion.ide.LengthGaugeHydrogenKernel,
"hydrogen_with_cc":
ion.ide.LengthGaugeHydrogenKernelWithContinuumContinuumInteraction,
}
kernel_key = si.ask_for_input(f'IDE Kernel? [{" | ".join(choices)}]',
default="hydrogen")
try:
kernel = choices[kernel_key]()
except KeyError:
raise ion.exceptions.InvalidChoice(
f"{kernel_key} is not one of {choices.keys()}")
parameters.append(si.Parameter(name="kernel", value=kernel))
return kernel
def ask_mask__radial_cosine(parameters, mesh_kwargs):
outer_radius_default = mesh_kwargs["outer_radius"] / u.bohr_radius
inner = u.bohr_radius * si.ask_for_input(
"Mask Inner Radius (in Bohr radii)?",
default=np.ceil(outer_radius_default * 0.8),
callback=float,
)
outer = u.bohr_radius * si.ask_for_input(
"Mask Outer Radius (in Bohr radii)?",
default=np.ceil(outer_radius_default),
callback=float,
)
smoothness = si.ask_for_input("Mask Smoothness?", default=8, callback=int)
mask = si.Parameter(
name="mask",
value=ion.potentials.RadialCosineMask(inner_radius=inner,
outer_radius=outer,
smoothness=smoothness),
)
parameters.append(mask)
def create_scan(tag):
parameters = []
spec_type = ion.ide.IntegroDifferentialEquationSpecification
evolution_method = utils.ask_evolution_method_ide(parameters)
test_charge = u.electron_charge
test_mass = u.electron_mass
test_energy = ion.states.HydrogenBoundState(1, 0).energy
parameters.append(si.Parameter(name="test_charge", value=test_charge))
parameters.append(si.Parameter(name="test_mass", value=test_mass))
parameters.append(si.Parameter(name="test_energy", value=test_energy))
parameters.append(
si.Parameter(name="integral_prefactor",
value=-((u.electron_charge / u.hbar)**2)))
kernel = utils.ask_ide_kernel(parameters)
utils.ask_time_step(parameters)
time_initial_in_pw, time_final_in_pw, extra_time = utils.ask_time_evolution(
)
# PULSE PARAMETERS
pulse_parameters = utils.construct_pulses(
parameters,
time_initial_in_pw=time_initial_in_pw,
time_final_in_pw=time_final_in_pw,
)
# MISCELLANEOUS
utils.ask_data_storage_ide(parameters, spec_type=spec_type)
# CREATE SPECS
expanded_parameters = si.expand_parameters(parameters)
extra_parameters = dict(checkpoints=True,
checkpoint_every=datetime.timedelta(minutes=20))
specs = []
print("Generating specifications...")
for component, spec_kwargs in enumerate(
tqdm(expanded_parameters, ascii=True)):
electric_potential = spec_kwargs["electric_potential"]
time_initial = time_initial_in_pw * electric_potential.pulse_width
time_final = (time_final_in_pw *
electric_potential.pulse_width) + extra_time
spec = spec_type(
name=component,
component=component,
time_initial=time_initial,
time_final=time_final,
**spec_kwargs,
**extra_parameters,
)
utils.transfer_potential_attrs_to_spec(electric_potential, spec)
spec.time_initial_in_pw = time_initial_in_pw
spec.time_final_in_pw = time_final_in_pw
specs.append(spec)
# CREATE MAP
opts, custom = utils.ask_map_options()
map = utils.run.map(specs,
map_options=htmap.MapOptions(**opts,
custom_options=custom),
tag=tag)
print(f"Created map {map.tag}")
return map
def ask_data_storage_tdse(parameters, *, spec_type):
parameters.append(
si.Parameter(
name="store_data_every",
value=si.ask_for_input("Store Data Every n Time Steps",
default=-1,
callback=int),
))
datastores_questions_defaults = [
(ion.mesh.Fields,
"Store Electric Field and Vector Potential vs. Time?", True),
(ion.mesh.Norm, "Store Wavefunction Norm vs. Time?", True),
(ion.mesh.InnerProducts, "Store Wavefunction Inner Products vs. Time?",
True),
(
ion.mesh.InternalEnergyExpectationValue,
"Store Internal Energy Expectation Value vs. Time?",
False,
),
(
ion.mesh.TotalEnergyExpectationValue,
"Store Total Energy Expectation Value vs. Time?",
False,
),
(ion.mesh.ZExpectationValue, "Store Z Expectation Value vs. Time?",
False),
(ion.mesh.RExpectationValue, "Store R Expectation Value vs. Time?",
False),
(ion.mesh.NormWithinRadius, "Store Norm Within Radius vs. Time?",
False),
]
if spec_type == ion.mesh.SphericalHarmonicSpecification:
datastores_questions_defaults += [
(
ion.mesh.DirectionalRadialProbabilityCurrent,
"Store Radial Probability Current vs. Time?",
False,
),
(ion.mesh.NormBySphericalHarmonic, "Store Norm-by-L?", False),
]
datastores = []
for cls, question, default in datastores_questions_defaults:
if not si.ask_for_bool(question, default=default):
continue
argspec = inspect.getfullargspec(cls.__init__)
arg_names = argspec.args[1:]
arg_defaults = argspec.defaults
if len(arg_names) > 0:
args = {
name: si.ask_for_eval(f"Value for {name}?", default=default)
for name, default in reversed(
tuple(
itertools.zip_longest(reversed(arg_names),
reversed(arg_defaults))))
}
datastore = cls(**args)
else:
datastore = cls()
datastores.append(datastore)
parameters.append(si.Parameter(name="datastores", value=datastores))
