Skip to content

Commit e0fbc6a

Browse files
jokasimrjokasimr
committed
docs: polarization calibration prototype
1 parent 413ed1a commit e0fbc6a

File tree

1 file changed

+316
-0
lines changed

1 file changed

+316
-0
lines changed

‎docs/user-guide/estia/polcal.ipynb‎

Lines changed: 316 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,316 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "0",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"# flake8: noqa\n",
11+
"import numpy as np\n",
12+
"import scipp as sc\n",
13+
"\n",
14+
"\n",
15+
"def solve_for_calibration_parameters(Io, I):\n",
16+
" Iopp, Iopa, Ioap, Ioaa = Io\n",
17+
" Ipp, Ipa, Iap, Iaa = I\n",
18+
"\n",
19+
" I0 = 2 * (Iopp * Ioaa - Iopa * Ioap) / (Iopp + Ioaa - Iopa - Ioap)\n",
20+
"\n",
21+
" rho = (Ioaa - Ioap) / (Iopp - Iopa)\n",
22+
" alp = (Ioaa - Iopa) / (Iopp - Ioap)\n",
23+
"\n",
24+
" Rspp_plus_Rsaa = 4 * (rho * alp * Ipp + Iaa + rho * Ipa + alp * Iap) / ((1 + rho) * (1 + alp) * I0)\n",
25+
"\n",
26+
" Pp = sc.sqrt(\n",
27+
" (Ipp + Iaa - Ipa - Iap) * (alp * (Ipp - Iap) - Iaa + Ipa) /\n",
28+
" ((rho * alp * Ipp + Iaa + rho * Ipa + alp * Iap) * (rho * (Ipp - Ipa) - Iaa + Iap))\n",
29+
" )\n",
30+
"\n",
31+
" Ap = sc.sqrt(\n",
32+
" (Ipp + Iaa - Ipa - Iap) * (rho * (Ipp - Ipa) - Iaa + Iap) /\n",
33+
" ((rho * alp * Ipp + Iaa + rho * Ipa + alp * Iap) * (alp * (Ipp - Iap) - Iaa + Ipa))\n",
34+
" )\n",
35+
"\n",
36+
" Rs = sc.sqrt(((alp * (Ipp - Iap) - Iaa + Ipa)) * (rho * (Ipp - Ipa) - Iaa + Iap) / ((rho * alp * Ipp + Iaa + rho * Ipa + alp * Iap) * (Ipp + Iaa - Ipa - Iap)))\n",
37+
"\n",
38+
" Pa = -rho * Pp\n",
39+
" Aa = -alp * Ap\n",
40+
" \n",
41+
" Rspp_minus_Rsaa = Rs * Rspp_plus_Rsaa\n",
42+
" Rspp = (Rspp_plus_Rsaa + Rspp_minus_Rsaa) / 2\n",
43+
" Rsaa = Rspp_plus_Rsaa - Rspp\n",
44+
"\n",
45+
" return I0 / 4, Pp, Pa, Ap, Aa, Rspp, Rsaa\n",
46+
"\n",
47+
"\n",
48+
"def generate_valid_calibration_parameters():\n",
49+
" I0 = np.random.random()\n",
50+
" Pp = np.random.random()\n",
51+
" Pa = -np.random.random()\n",
52+
" Ap = np.random.random()\n",
53+
" Aa = -np.random.random()\n",
54+
" Rspp = np.random.random()\n",
55+
" Rsaa = Rspp * np.random.random()\n",
56+
" return tuple(map(sc.scalar, (I0, Pp, Pa, Ap, Aa, Rspp, Rsaa)))\n",
57+
"\n",
58+
"\n",
59+
"def intensity_from_parameters(I0, Pp, Pa, Ap, Aa, Rpp, Rpa, Rap, Raa):\n",
60+
" return (\n",
61+
" I0 * (Rpp * (1 + Ap) * (1 + Pp) + Rpa * (1 - Ap) * (1 + Pp) + Rap * (1 + Ap) * (1 - Pp) + Raa * (1 - Ap) * (1 - Pp)),\n",
62+
" I0 * (Rpp * (1 + Aa) * (1 + Pp) + Rpa * (1 - Aa) * (1 + Pp) + Rap * (1 + Aa) * (1 - Pp) + Raa * (1 - Aa) * (1 - Pp)),\n",
63+
" I0 * (Rpp * (1 + Ap) * (1 + Pa) + Rpa * (1 - Ap) * (1 + Pa) + Rap * (1 + Ap) * (1 - Pa) + Raa * (1 - Ap) * (1 - Pa)),\n",
64+
" I0 * (Rpp * (1 + Aa) * (1 + Pa) + Rpa * (1 - Aa) * (1 + Pa) + Rap * (1 + Aa) * (1 - Pa) + Raa * (1 - Aa) * (1 - Pa)),\n",
65+
" )\n"
66+
]
67+
},
68+
{
69+
"cell_type": "markdown",
70+
"id": "1",
71+
"metadata": {},
72+
"source": [
73+
"## Sanity check: Can we find the polarization parameters given the measured intensities?"
74+
]
75+
},
76+
{
77+
"cell_type": "code",
78+
"execution_count": null,
79+
"id": "2",
80+
"metadata": {},
81+
"outputs": [],
82+
"source": [
83+
"i0, Pp, Pa, Ap, Aa, Rpp, Raa = generate_valid_calibration_parameters()\n",
84+
"\n",
85+
"I0 = intensity_from_parameters(i0, Pp, Pa, Ap, Aa, sc.scalar(1), sc.scalar(0), sc.scalar(0), sc.scalar(1))\n",
86+
"I = intensity_from_parameters(i0, Pp, Pa, Ap, Aa, Rpp, sc.scalar(0), sc.scalar(0), Raa)\n",
87+
"\n",
88+
"solve_for_calibration_parameters(I0, I)"
89+
]
90+
},
91+
{
92+
"cell_type": "code",
93+
"execution_count": null,
94+
"id": "3",
95+
"metadata": {},
96+
"outputs": [],
97+
"source": [
98+
"# Ground truth\n",
99+
"i0, Pp, Pa, Ap, Aa, Rpp, Raa"
100+
]
101+
},
102+
{
103+
"cell_type": "code",
104+
"execution_count": null,
105+
"id": "4",
106+
"metadata": {},
107+
"outputs": [],
108+
"source": [
109+
"def correction_matrix(Pp, Pa, Ap, Aa):\n",
110+
" return [\n",
111+
" [(1 + Pp) * (1 + Ap), (1 + Pp) * (1 - Ap), (1 - Pp) * (1 + Ap), (1 - Pp) * (1 - Ap)],\n",
112+
" [(1 + Pp) * (1 + Aa), (1 + Pp) * (1 - Aa), (1 - Pp) * (1 + Aa), (1 - Pp) * (1 - Aa)],\n",
113+
" [(1 + Pa) * (1 + Ap), (1 + Pa) * (1 - Ap), (1 - Pa) * (1 + Ap), (1 - Pa) * (1 - Ap)],\n",
114+
" [(1 + Pa) * (1 + Aa), (1 + Pa) * (1 - Aa), (1 - Pa) * (1 + Aa), (1 - Pa) * (1 - Aa)],\n",
115+
" ]\n",
116+
" \n",
117+
"\n",
118+
"def reduce_to_q(I, q, qbins):\n",
119+
" if isinstance(I, list | tuple):\n",
120+
" return [reduce_to_q(i, q, qbins) for i in I]\n",
121+
" return I.assign_coords(Q=q(I)).flatten(I.dims, to='_').hist(Q=qbins)\n",
122+
"\n",
123+
"\n",
124+
"def compute_calibration_factors(Io, I):\n",
125+
" I0, Pp, Pa, Ap, Aa, _, _ = solve_for_calibration_parameters(Io, I)\n",
126+
" return I0, correction_matrix(Pp, Pa, Ap, Aa)\n",
127+
"\n",
128+
"\n",
129+
"def linsolve(A, b):\n",
130+
" An = np.stack([[a.values for a in row] for row in A])\n",
131+
" An = An.transpose((*range(2, An.ndim), 0, 1))\n",
132+
" bn = np.stack([bi.values for bi in b], axis=-1)\n",
133+
" xn = np.linalg.solve(An, bn)\n",
134+
" xn = xn.transpose((xn.ndim-1, *range(xn.ndim-1)))\n",
135+
" out = []\n",
136+
" for bi, v in zip(b, xn):\n",
137+
" out.append(bi.copy())\n",
138+
" out[-1].values[:] = v\n",
139+
" return out\n",
140+
"\n",
141+
"\n",
142+
"def compute_reflectivity_alt1(Io, Is, I, q, qbins):\n",
143+
" 'Compute factors on lz, reduce to q, apply correction on q'\n",
144+
" I0, C = compute_calibration_factors(Io, Is)\n",
145+
" return linsolve(\n",
146+
" reduce_to_q(\n",
147+
" [[I0 * c for c in row] for row in C],\n",
148+
" q, qbins,\n",
149+
" ),\n",
150+
" reduce_to_q(I, q, qbins)\n",
151+
" )\n",
152+
"\n",
153+
"\n",
154+
"def compute_reflectivity_alt2(Io, Is, I, q, qbins):\n",
155+
" 'Compute factors on lz, apply correction on lz, reduce to q'\n",
156+
" I0, C = compute_calibration_factors(Io, Is)\n",
157+
" i = (\n",
158+
" reduce_to_q(\n",
159+
" linsolve(C, I),\n",
160+
" q, qbins\n",
161+
" )\n",
162+
" )\n",
163+
" r = reduce_to_q(I0, q, qbins)\n",
164+
" return [ii / r for ii in i]\n",
165+
"\n",
166+
"\n",
167+
"def compute_reflectivity_alt3(Io, Is, I, q, qbins):\n",
168+
" 'Compute factors on q, reduce to q, apply correction on q'\n",
169+
" I0, C = compute_calibration_factors(reduce_to_q(Io, q, qbins), reduce_to_q(Is, q, qbins))\n",
170+
" i = linsolve(C, reduce_to_q(I, q, qbins))\n",
171+
" r = I0\n",
172+
" return [ii / r for ii in i]\n",
173+
"\n"
174+
]
175+
},
176+
{
177+
"cell_type": "code",
178+
"execution_count": null,
179+
"id": "5",
180+
"metadata": {},
181+
"outputs": [],
182+
"source": [
183+
"\n",
184+
"def run_experiment(M, N, K, noise_level, smooth_polarizer_reflectivity, title, path):\n",
185+
" \n",
186+
" wavelength = sc.linspace('wavelength', 0.1, 1, M)\n",
187+
" theta = sc.linspace('z', 0.2, 0.5, N)\n",
188+
" q = theta / wavelength\n",
189+
" \n",
190+
" def Rs(q):\n",
191+
" return 0.5 * (1 + sc.cos(sc.scalar(31., unit='rad') * q / 5)) / (1 + q**2)\n",
192+
" \n",
193+
" def sharp_drop(l):\n",
194+
" return sc.where(l > 0.2, sc.linspace('wavelength', 0.5, 1, M), sc.scalar(0.1))\n",
195+
" \n",
196+
" def noise():\n",
197+
" return sc.array(dims=('wavelength', 'z'), values=noise_level * np.random.randn(M, N))\n",
198+
" \n",
199+
" i0, Pp, Pa, Ap, Aa, Rpp, Raa = (\n",
200+
" sc.array(dims=('wavelength', 'z'), values=np.ones((M, N))),\n",
201+
" sc.linspace('wavelength', 0.5, 1, M) if smooth_polarizer_reflectivity else sharp_drop(wavelength),\n",
202+
" -(sc.linspace('wavelength', 0.5, 1, M) if smooth_polarizer_reflectivity else sharp_drop(wavelength)),\n",
203+
" sc.linspace('wavelength', 0.5, 1, M) if smooth_polarizer_reflectivity else sharp_drop(wavelength),\n",
204+
" -(sc.linspace('wavelength', 0.5, 1, M) if smooth_polarizer_reflectivity else sharp_drop(wavelength)),\n",
205+
" Rs(q),\n",
206+
" 0.5 * Rs(q),\n",
207+
" )\n",
208+
" \n",
209+
" *I0, = map(lambda i: sc.DataArray(i * (1 + noise())), intensity_from_parameters(\n",
210+
" i0, Pp, Pa, Ap, Aa,\n",
211+
" i0, 0*i0,\n",
212+
" 0*i0, i0,\n",
213+
" ))\n",
214+
" *Is, = map(lambda i: sc.DataArray(i * (1 + noise())), intensity_from_parameters(\n",
215+
" i0, Pp, Pa, Ap, Aa,\n",
216+
" Rpp, Rpp*0,\n",
217+
" Raa*0, Raa\n",
218+
" ))\n",
219+
" *I, = map(lambda i: sc.DataArray(i * (1 + noise())), intensity_from_parameters(\n",
220+
" i0, Pp, Pa, Ap, Aa,\n",
221+
" 0.5 * Rpp, 0.2 * Rpp,\n",
222+
" 0.3 * Raa, 0.4 * Raa\n",
223+
" ))\n",
224+
" \n",
225+
" qf = lambda I: q\n",
226+
" qbins = sc.geomspace('Q', 0.2, 5, K)\n",
227+
"\n",
228+
" r1 = compute_reflectivity_alt1(I0, Is, I, qf, qbins)\n",
229+
" r2 = compute_reflectivity_alt2(I0, Is, I, qf, qbins)\n",
230+
" r3 = compute_reflectivity_alt3(I0, Is, I, qf, qbins)\n",
231+
"\n",
232+
" i = reduce_to_q(tuple(map(sc.DataArray, [0.5 * Rpp, 0.2 * Rpp, 0.3 * Raa, 0.4 * Raa])), qf, qbins)\n",
233+
" r = reduce_to_q(sc.DataArray(i0), qf, qbins)\n",
234+
" gt = [ii/r for ii in i]\n",
235+
"\n",
236+
" polarizer_reflectivity_features = \"P_p is sharp\" if not smooth_polarizer_reflectivity else \"P_p is smooth\"\n",
237+
"\n",
238+
" suffix = f'l{M}_z{N}_q{K}_eps{noise_level}_{\"smooth\" if smooth_polarizer_reflectivity else \"sharp\"}'\n",
239+
"\n",
240+
" methods = {1: 'Mix', 2: 'On $\\lambda z$', 3: 'On $Q$'}\n",
241+
"\n",
242+
" p = sc.plot({\n",
243+
" #methods[1]: sc.abs(r1[0] - gt[0]),\n",
244+
" methods[2]: sc.abs(r2[0] - gt[0]),\n",
245+
" methods[3]: sc.abs(r3[0] - gt[0]),\n",
246+
" }, title=f'{title} Error, method comparison, {polarizer_reflectivity_features}')\n",
247+
" p.save(f'{path}/error_{suffix}.png')\n",
248+
" \n",
249+
"\n",
250+
" for m, i in (\n",
251+
" #(1, r1),\n",
252+
" (r2, 2),\n",
253+
" (r3, 3)\n",
254+
" ):\n",
255+
" p = sc.plot({'ground truth': gt[0], 'computed reflectivity': m[0]}, title=f'{title} $R(Q)$, method: {methods[i]}, {polarizer_reflectivity_features}')\n",
256+
" p.save(f'{path}/reflectivity_{i}_{suffix}.png')\n",
257+
"\n",
258+
" if smooth_polarizer_reflectivity:\n",
259+
" p = sc.DataArray(Pp).plot(title='Polarizer reflectivity - smooth')\n",
260+
" p.save(f'{path}/polarizer_reflectivity_smooth.png')\n",
261+
" else:\n",
262+
" p = sc.DataArray(Pp).plot(title='Polarizer reflectivity - sharp drop')\n",
263+
" p.save(f'{path}/polarizer_reflectivity_sharp_drop.png')"
264+
]
265+
},
266+
{
267+
"cell_type": "code",
268+
"execution_count": null,
269+
"id": "6",
270+
"metadata": {},
271+
"outputs": [],
272+
"source": [
273+
"run_experiment(400, 500, 400, 0, True, '', 'docs/user-guide/estia/figures')\n",
274+
"run_experiment(400, 500, 400, 0, False, '', 'docs/user-guide/estia/figures')\n",
275+
"run_experiment(400, 500, 400, 0.01, True, 'Realistic counts -', 'docs/user-guide/estia/figures')\n",
276+
"run_experiment(400, 500, 400, 0.1, True, 'Low counts -', 'docs/user-guide/estia/figures')"
277+
]
278+
},
279+
{
280+
"cell_type": "markdown",
281+
"id": "7",
282+
"metadata": {},
283+
"source": [
284+
"| Property | On $Q$ | On $\\lambda z$ | Calibrate on $\\lambda z$, correct on $Q$ |\n",
285+
"|--------------|---------------|---------------|---------------|\n",
286+
"| Can pre-reduce reference | Yes | Yes | Yes |\n",
287+
"| Can avoid reducing reference on $\\lambda z$ grid | Yes | No | No |\n",
288+
"| Can avoid reducing sample on $\\lambda z$ grid | Yes | No | Yes |\n",
289+
"| Runtime | Least | Longest | Medium |\n",
290+
"| Unbiased in the high counts limit | No | Yes | Yes |\n",
291+
"| Robust to low counts | Yes | No | No |"
292+
]
293+
}
294+
],
295+
"metadata": {
296+
"kernelspec": {
297+
"display_name": "Python 3 (ipykernel)",
298+
"language": "python",
299+
"name": "python3"
300+
},
301+
"language_info": {
302+
"codemirror_mode": {
303+
"name": "ipython",
304+
"version": 3
305+
},
306+
"file_extension": ".py",
307+
"mimetype": "text/x-python",
308+
"name": "python",
309+
"nbconvert_exporter": "python",
310+
"pygments_lexer": "ipython3",
311+
"version": "3.10.14"
312+
}
313+
},
314+
"nbformat": 4,
315+
"nbformat_minor": 5
316+
}

0 commit comments

Comments
 (0)