Skip to content

Commit

Permalink
cleaned up tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Bing Li committed Aug 8, 2024
1 parent a286b9a commit 737c1f6
Show file tree
Hide file tree
Showing 8 changed files with 246 additions and 243 deletions.
21 changes: 0 additions & 21 deletions LICENSE

This file was deleted.

52 changes: 33 additions & 19 deletions src/tavi/instrument/resolution/cooper_nathans.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import numpy as np
import numpy.linalg as la
from tavi.utilities import *
from tavi.instrument.tas import TAS

from tavi.instrument.resolution.reso_ellipses import ResoEllipsoid
from tavi.instrument.tas import TAS
from tavi.utilities import *


class CN(TAS):
Expand Down Expand Up @@ -67,31 +68,44 @@ def cooper_nathans(
rez.projection = projection

if self._mat_f is None:
mono_mosaic = np.deg2rad(self.monochromator.mosaic / 60)
mono_mosaic_v = np.deg2rad(self.monochromator.mosaic_v / 60)
ana_mosaic = np.deg2rad(self.analyzer.mosaic / 60)
ana_mosaic_v = np.deg2rad(self.analyzer.mosaic_v / 60)

# matrix F, divergence of monochromator and analyzer, [pop75] Appendix 1
mat_f = np.zeros(((CN.NUM_MONOS + CN.NUM_ANAS) * 2, (CN.NUM_MONOS + CN.NUM_ANAS) * 2))
mat_f[CN.IDX_MONO0_H, CN.IDX_MONO0_H] = 1.0 / self.monochromator.mosaic**2
mat_f[CN.IDX_MONO0_V, CN.IDX_MONO0_V] = 1.0 / self.monochromator.mosaic_v**2
mat_f[CN.IDX_ANA0_H, CN.IDX_ANA0_H] = 1.0 / self.analyzer.mosaic**2
mat_f[CN.IDX_ANA0_V, CN.IDX_ANA0_V] = 1.0 / self.analyzer.mosaic_v**2
mat_f[CN.IDX_MONO0_H, CN.IDX_MONO0_H] = 1.0 / mono_mosaic**2
mat_f[CN.IDX_MONO0_V, CN.IDX_MONO0_V] = 1.0 / mono_mosaic_v**2
mat_f[CN.IDX_ANA0_H, CN.IDX_ANA0_H] = 1.0 / ana_mosaic**2
mat_f[CN.IDX_ANA0_V, CN.IDX_ANA0_V] = 1.0 / ana_mosaic_v**2
self._mat_f = mat_f

if self._mat_g is None:
coll_h_pre_mono = np.deg2rad(self.collimators.h_pre_mono / 60)
coll_v_pre_mono = np.deg2rad(self.collimators.v_pre_mono / 60)
coll_h_pre_sample = np.deg2rad(self.collimators.h_pre_sample / 60)
coll_v_pre_sample = np.deg2rad(self.collimators.v_pre_sample / 60)
coll_h_post_sample = np.deg2rad(self.collimators.h_post_sample / 60)
coll_v_post_sample = np.deg2rad(self.collimators.v_post_sample / 60)
coll_h_post_ana = np.deg2rad(self.collimators.h_post_ana / 60)
coll_v_post_ana = np.deg2rad(self.collimators.v_post_ana / 60)

# matrix G, divergence of collimators, [pop75] Appendix 1
mat_g = np.zeros((CN.NUM_COLLS * 2, CN.NUM_COLLS * 2))
mat_g[CN.IDX_COLL0_H, CN.IDX_COLL0_H] = 1.0 / self.collimators.h_pre_mono**2
mat_g[CN.IDX_COLL0_V, CN.IDX_COLL0_V] = 1.0 / self.collimators.v_pre_mono**2
mat_g[CN.IDX_COLL1_H, CN.IDX_COLL1_H] = 1.0 / self.collimators.h_pre_sample**2
mat_g[CN.IDX_COLL1_V, CN.IDX_COLL1_V] = 1.0 / self.collimators.v_pre_sample**2
mat_g[CN.IDX_COLL2_H, CN.IDX_COLL2_H] = 1.0 / self.collimators.h_post_sample**2
mat_g[CN.IDX_COLL2_V, CN.IDX_COLL2_V] = 1.0 / self.collimators.v_post_sample**2
mat_g[CN.IDX_COLL3_H, CN.IDX_COLL3_H] = 1.0 / self.collimators.h_post_ana**2
mat_g[CN.IDX_COLL3_V, CN.IDX_COLL3_V] = 1.0 / self.collimators.v_post_ana**2
mat_g[CN.IDX_COLL0_H, CN.IDX_COLL0_H] = 1.0 / coll_h_pre_mono**2
mat_g[CN.IDX_COLL0_V, CN.IDX_COLL0_V] = 1.0 / coll_v_pre_mono**2
mat_g[CN.IDX_COLL1_H, CN.IDX_COLL1_H] = 1.0 / coll_h_pre_sample**2
mat_g[CN.IDX_COLL1_V, CN.IDX_COLL1_V] = 1.0 / coll_v_pre_sample**2
mat_g[CN.IDX_COLL2_H, CN.IDX_COLL2_H] = 1.0 / coll_h_post_sample**2
mat_g[CN.IDX_COLL2_V, CN.IDX_COLL2_V] = 1.0 / coll_v_post_sample**2
mat_g[CN.IDX_COLL3_H, CN.IDX_COLL3_H] = 1.0 / coll_h_post_ana**2
mat_g[CN.IDX_COLL3_V, CN.IDX_COLL3_V] = 1.0 / coll_v_post_ana**2

self._mat_g = mat_g

# determine frame
if isinstance(hkl, tuple | list) and len(hkl) == 3:

if projection is None: # Local Q frame
rez.frame = "q"
rez.angles = (90, 90, 90)
Expand Down Expand Up @@ -218,8 +232,10 @@ def cooper_nathans(
mat_cov = mat_ba @ mat_h_inv @ mat_ba.T

# TODO smaple mosaic????
mat_cov[1, 1] += q_mod**2 * self.sample.mosaic**2
mat_cov[2, 2] += q_mod**2 * self.sample.mosaic_v**2
sample_mosaic = np.deg2rad(self.sample.mosaic / 60)
sample_mosaic_v = np.deg2rad(self.sample.mosaic_v / 60)
mat_cov[1, 1] += q_mod**2 * sample_mosaic**2
mat_cov[2, 2] += q_mod**2 * sample_mosaic_v**2

mat_reso = la.inv(mat_cov) * sig2fwhm**2

Expand All @@ -228,7 +244,6 @@ def cooper_nathans(
rez.q = (q_mod, 0, 0)

elif rez.frame == "hkl":

conv_mat_4d = np.eye(4)
conv_mat_4d[0:3, 0:3] = (
np.array(
Expand All @@ -243,7 +258,6 @@ def cooper_nathans(
rez.mat = conv_mat_4d.T @ mat_reso @ conv_mat_4d

elif rez.frame == "proj":

conv_mat_4d = np.eye(4)
conv_mat_4d[0:3, 0:3] = (
np.array(
Expand Down
60 changes: 20 additions & 40 deletions src/tavi/instrument/tas.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import math
import json

import numpy as np
from tavi.utilities import *
from tavi.instrument.tas_cmponents import *
from tavi.sample.xtal import Xtal
from tavi.sample.powder import Powder
from tavi.sample.sample import Sample
from tavi.sample.xtal import Xtal
from tavi.utilities import *


class TAS(object):
Expand Down Expand Up @@ -59,7 +60,9 @@ def load_instrument_from_json(self, path_to_json):

# def save_instrument(self):
# """Save configuration into a dictionary"""
# pass
# # convert python dictionary to json file
# with open("./src/tavi/instrument/instrument_params/takin_test.json", "w") as file:
# json.dump(instrument_config, file)

def load_sample(self, sample):
"""Load sample info"""
Expand All @@ -70,36 +73,13 @@ def load_sample_from_json(self, path_to_json):

with open(path_to_json, "r", encoding="utf-8") as file:
sample_params = json.load(file)
lattice_params = (
sample_params["a"],
sample_params["b"],
sample_params["c"],
sample_params["alpha"],
sample_params["beta"],
sample_params["gamma"],
)

if sample_params["type"] == "xtal":
sample = Xtal(lattice_params=lattice_params)
sample.ub_matrix = np.array(sample_params["ub_matrix"]).reshape(3, 3)
sample.plane_normal = np.array(sample_params["plane_normal"])
elif sample_params["type"] == "powder":
sample = Powder(lattice_params=lattice_params)
else:
print("Sample type needs to be either xtal or powder.")

param_dict = ("shape", "width", "height", "depth", "mosaic", "mosaic_v")

for key, val in sample_params.items():

match key:
case "height" | "width" | "depth":
setattr(sample, key, val * cm2angstrom)
case "mosaic" | "mosaic_v":
setattr(sample, key, val * min2rad)
case _:
if key in param_dict:
setattr(sample, key, val)

if sample_params["type"] == "xtal":
sample = Xtal.from_json(sample_params)
elif sample_params["type"] == "powder":
sample = Powder.from_json(sample_params)
else:
sample = Sample.from_json(sample_params)

self.load_sample(sample)

Expand All @@ -114,9 +94,9 @@ def q_lab(two_theta, ki, kf):
"""
return np.array(
[
-kf * np.sin(two_theta / rad2deg),
-kf * np.sin(np.deg2rad(two_theta)),
0,
ki - kf * np.cos(two_theta / rad2deg),
ki - kf * np.cos(np.deg2rad(two_theta)),
]
)

Expand Down Expand Up @@ -259,11 +239,11 @@ def find_two_theta(self, peak, ei=13.5, ef=None):
if two_theta is None:
print(f"Triangle cannot be closed at q={hkl}, en={ei-ef} meV.")
return None
elif two_theta * rad2deg < S2_MIN_DEG:
elif np.rad2deg(two_theta) < S2_MIN_DEG:
print(f"s2 is smaller than {S2_MIN_DEG} deg at q={hkl}.")
return None
else:
return two_theta * rad2deg * self.goniometer.sense
return np.rad2deg(two_theta) * self.goniometer.sense

def find_angles(self, peak, ei=13.5, ef=None):
"""calculate motor positions for a given peak if UB matrix has been determined
Expand Down Expand Up @@ -294,11 +274,11 @@ def find_angles(self, peak, ei=13.5, ef=None):
if two_theta is None:
print(f"Triangle cannot be closed at q={hkl}, en={ei-ef} meV.")
angles = None
elif two_theta * rad2deg < S2_MIN_DEG:
elif np.rad2deg(two_theta) < S2_MIN_DEG:
print(f"s2 is smaller than {S2_MIN_DEG} deg at q={hkl}.")
angles = None
else:
two_theta = two_theta * rad2deg * self.goniometer.sense
two_theta = np.rad2deg(two_theta) * self.goniometer.sense
q = self.sample.ub_matrix @ hkl
t1 = q / np.linalg.norm(q)

Expand Down
47 changes: 18 additions & 29 deletions src/tavi/instrument/tas_cmponents.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import numpy as np

from tavi.utilities import *


class Source(object):
"""Neutron source"""

def __init__(self, param_dict):

self.name = None
self.shape = "rectangular"
self.width = None
Expand All @@ -31,7 +31,6 @@ class Guide(object):
"""Guide"""

def __init__(self, param_dict):

self.in_use = False
self.div_h = None
self.div_v = None
Expand All @@ -48,7 +47,6 @@ class Monochromator(object):
"""Monochromator"""

def __init__(self, param_dict):

self.type = "PG002"
self.d_spacing = mono_ana_xtal["PG002"]
self.mosaic = 45 # horizontal mosaic, min of arc
Expand All @@ -72,8 +70,8 @@ def __init__(self, param_dict):

for key, val in param_dict.items():
match key:
case "mosaic" | "mosaic_v" | "curvh" | "curvv":
setattr(self, key, val * min2rad)
# case "mosaic" | "mosaic_v" | "curvh" | "curvv":
# setattr(self, key, val * min2rad)
case "width" | "height" | "depth":
# divide by np.sqrt(12) if rectangular
# Diameter D/4 if spherical
Expand Down Expand Up @@ -102,32 +100,29 @@ def __init__(self, param_dict):
self.v_post_ana = 30

for key, val in param_dict.items():
if key in (
"h_pre_mono",
"h_pre_sample",
"h_post_sample",
"h_post_ana",
"v_pre_mono",
"v_pre_sample",
"v_post_sample",
"v_post_ana",
):
setattr(self, key, val * min2rad)
else:
setattr(self, key, val)
# if key in (
# "h_pre_mono",
# "h_pre_sample",
# "h_post_sample",
# "h_post_ana",
# "v_pre_mono",
# "v_pre_sample",
# "v_post_sample",
# "v_post_ana",
# ):
# setattr(self, key, val * min2rad)
# else:
setattr(self, key, val)


# TODO
class Monitor(object):

def __init__(self, param_dict):

self.shape = "rectangular" # or spherical
self.width = 5 * cm2angstrom
self.height = 12 * cm2angstrom

for key, val in param_dict.items():

match key:
case "width" | "height":
# divide by np.sqrt(12) if rectangular
Expand All @@ -143,9 +138,7 @@ def __init__(self, param_dict):


class Analyzer(object):

def __init__(self, param_dict):

self.type = "Pg002"
self.d_spacing = mono_ana_xtal["Pg002"]
self.mosaic = 45 # horizontal mosaic, min of arc
Expand All @@ -168,8 +161,8 @@ def __init__(self, param_dict):

for key, val in param_dict.items():
match key:
case "mosaic" | "mosaic_v" | "curvh" | "curvv":
setattr(self, key, val * min2rad)
# case "mosaic" | "mosaic_v" | "curvh" | "curvv":
# setattr(self, key, val * min2rad)
case "width" | "height" | "depth":
# divide by np.sqrt(12) if rectangular
# Diameter D/4 if spherical
Expand All @@ -188,7 +181,6 @@ def __init__(self, param_dict):

# TODO
class Detector(object):

def __init__(self, param_dict):
self.shape = "rectangular"
self.width = 1.5 # in cm
Expand Down Expand Up @@ -239,7 +231,6 @@ def r_mat(self, angles):

omega, sgl, sgu = angles # s2, s1, sgl, sgu
match self.type:

case "Y-ZX": # HB3
r_mat = rot_y(omega) @ rot_z(-1 * sgl) @ rot_x(sgu)
case "YZ-X": # CG4C
Expand Down Expand Up @@ -267,9 +258,7 @@ def angles_from_r_mat(self, r_mat):
"""

match self.type:

case "Y-ZX" | "YZ-X": # Y-mZ-X (s1, sgl, sgu) for HB1A and HB3, Y-Z-mX (s1, sgl, sgu) for CG4C

# sgl1 = np.arcsin(r_mat[1, 0]) * rad2deg
# sgl2 = np.arccos(np.sqrt(r_mat[0, 0] ** 2 + r_mat[2, 0] ** 2)) * rad2deg
sgl = np.arctan2(r_mat[1, 0], np.sqrt(r_mat[0, 0] ** 2 + r_mat[2, 0] ** 2)) * rad2deg
Expand Down
23 changes: 16 additions & 7 deletions src/tavi/sample/powder.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import numpy as np
from tavi.sample.sample import Sample


# TODO
class Powder(Sample):
"""Powder sample
Expand All @@ -14,9 +14,18 @@ def __init__(self, lattice_params):
super().__init__(lattice_params)
self.type = "powder"


if __name__ == "__main__":
powder = Powder(lattice_params=(1, 1, 1, 90, 90, 120))

print(powder.b_mat() / 2 / np.pi)
print(powder.b_mat() @ np.array([0, 1, 0]) / 2 / np.pi)
@classmethod
def from_json(cls, sample_params):
"""Alternate constructor from json"""
lattice_params = (
sample_params["a"],
sample_params["b"],
sample_params["c"],
sample_params["alpha"],
sample_params["beta"],
sample_params["gamma"],
)

sample = cls(lattice_params=lattice_params)

return sample
Loading

0 comments on commit 737c1f6

Please sign in to comment.