def test_slices(self):
gmu = constants.MU_GAMMA
muexp = MuonExperiment(['mu'])
muexp.spin_system.add_zeeman_term(0, 1.0 / gmu)
muexp.set_powder_average(100)
times = np.linspace(0, 1.0)
cos = np.cos(2 * np.pi * times)
signal = muexp.run_experiment(times)['e'][:, 0]
self.assertTrue(
np.all(
np.isclose(signal,
0.5 * (2.0 / 3.0 + 1.0 / 3.0 * cos),
atol=1e-3)))
# Now slice
signal_0 = muexp.run_experiment(times,
orient_slice=slice(0, 1))['e'][:, 0]
self.assertTrue(
np.all(
np.isclose(signal_0, 0.5 * cos * muexp.weights[0], atol=1e-3)))
def test_powder(self):
muexp = MuonExperiment()
N = 20
muexp.set_powder_average(N)
self.assertTrue(muexp.orientations.shape[0] == muexp.weights.shape[0])
self.assertTrue(muexp.weights.shape[0] >= N)
# Are these correct? Basic test
muexp.set_powder_average(1000)
o = muexp.orientations
w = muexp.weights
f = 3 * np.cos(o[:, 0])**2 - 1
self.assertAlmostEqual(np.sum(f * w), 0.0, 5)
import cProfile
import numpy as np
from muspinsim.experiment import MuonExperiment
exp = MuonExperiment(['e', 'mu', 'H'])
exp.spin_system.add_hyperfine_term(1, np.diag([1,1,4])*100)
exp.spin_system.add_hyperfine_term(2, np.diag([1,1,4])*100/3)
exp.set_powder_average(100)
exp.set_magnetic_field(5.0)
exp.set_dissipation_coupling(0, 1.0)
exp.set_muon_polarization('z')
times = np.linspace(0, 1, 1000)
op = exp.spin_system.operator({1: 'z'})
cProfile.run('exp.run_experiment(times, [op], "i")', sort='cumulative')
class MuonExperimentalSetup(object):
@mpi.execute_on_root
def __init__(self, params, logfile=None):
self._log = logfile
if logfile:
logging.basicConfig(filename=logfile,
level=logging.INFO,
format='[%(levelname)s] [%(threadName)s] [%(asctime)s] %(message)s',
datefmt='%Y-%m-%d %H:%M:%S')
# Create
self.experiment = MuonExperiment(params.spins)
self.log('Hamiltonian created with spins:')
self.log(', '.join(map(str, params.spins)))
self.log('Adding Hamiltonian terms:')
for i, B in params.zeeman.items():
self.experiment.spin_system.add_zeeman_term(i, B)
self.log('\tAdded zeeman term to spin {0}'.format(i+1))
for (i, j), A in params.hyperfine.items():
self.experiment.spin_system.add_hyperfine_term(i, np.array(A), j)
self.log('\tAdded hyperfine term to spin {0}'.format(i+1))
for (i, j), r in params.dipolar.items():
self.experiment.spin_system.add_dipolar_term(i, j, r)
self.log('\tAdded dipolar term to spins {0}, {1}'.format(i+1, j+1))
for i, EFG in params.quadrupolar.items():
self.experiment.spin_system.add_quadrupolar_term(i, EFG)
self.log('\tAdded quadrupolar term to spin {0}'.format(i+1))
for i, d in params.dissipation.items():
self.experiment.spin_system.set_dissipation(i, d)
self.log('\tSet dissipation parameter for spin '
'{0} to {1} MHz'.format(i+1, d))
self.log('')
# Ranges
trange = _make_range(params.time)
self.log('Using time range: '
'{0} to {1} us in {2} steps'.format(*trange))
self.time_axis = np.linspace(*trange)
brange = _make_range(params.field)
self.log('Using field range: '
'{0} to {1} T in {2} steps'.format(*brange))
self.field_axis = np.linspace(*brange)
# Powder averaging
if params.powder is None:
self.experiment.set_single_crystal(0, 0)
else:
scheme = params.powder[0]
N = params.powder[1]
self.experiment.set_powder_average(N, scheme)
self.log('Using powder averaging scheme '
'{0}'.format(scheme.upper()))
self.log(
'{0} orientations generated'.format(
len(self.experiment.weights)))
if params.polarization == 'longitudinal':
self.muon_axis = 'z'
elif params.polarization == 'transverse':
self.muon_axis = 'x'
else:
raise RuntimeError(
'Invalid polarization {0}'.format(params.polarization))
self.experiment.set_muon_polarization(self.muon_axis)
self.log('Muon beam polarized along axis {0}'.format(self.muon_axis))
# Temperature
self.temperature = params.temperature
self.experiment.set_temperature(self.temperature)
self.log('Using temperature of {0} K'.format(self.temperature))
# What to save
ssys = self.experiment.spin_system
self.observable = ssys.operator({ssys.muon_index: self.muon_axis})
self.save = [p[0] for p in params.save]
self.log('*'*20 + '\n')
@mpi.execute_on_root
def log(self, message):
if self._log:
logging.info(message)
def broadcast(self):
# Broadcast the contents of this object to other MPI threads
mpi.broadcast_object(self, only=[
'experiment',
'field_axis',
'time_axis',
'muon_axis',
'temperature',
'observable',
'save'
])
def run(self):
self.broadcast()
exp = self.experiment
ssys = self.experiment.spin_system
if ssys.is_dissipative:
self.log('Spin system is dissipative; using Lindbladian')
# Now slicing the values
N_f = len(self.field_axis) # Fields
N_o = len(exp.weights) # Orientations
results = {
'e': None,
'i': None
}
if 'e' in self.save:
results['e'] = np.zeros((N_f, len(self.time_axis)))
if 'i' in self.save:
results['i'] = np.zeros(N_f)
# Split the tasks
tasks = mpi.split_2D(range(N_f), range(N_o))
field_scan, orient_slice = tasks[mpi.rank]
self.log(f"MPI: {mpi.rank} will run fields: {field_scan}")
if len(tasks) > 1:
tsizes = [len(t[0])*len(t[1]) for t in tasks]
self.log('Splitting jobs over {0} cores'.format(mpi.size))
self.log('Job sizes:\n\t' + ' '.join(map(str, tsizes)))
# Loop over fields
for i in field_scan:
B = self.field_axis[i]
if i%10 == 0: # Reduces logfile spam
self.log('Performing calculations for B = {0} T'.format(B))
exp.set_magnetic_field(B)
field_results = exp.run_experiment(self.time_axis,
operators=[self.observable],
acquire=self.save,
orient_slice=orient_slice)
if 'e' in self.save:
results['e'][i] = field_results['e'][:, 0]
if 'i' in self.save:
results['i'][i] = field_results['i'][0]
# Reduce results
for k, v in results.items():
if v is None:
continue
results[k] = mpi.sum_data(v)
self.log('*'*20)
return results