Fix saving lists of arrays with recent versions of numpy
[qpms.git] / misc / rectlat_simple_modes.py
blobb39b5e430d00f7b2d19de786c5d23252e77ef9bf
1 #!/usr/bin/env python3
3 import math
4 from qpms.argproc import ArgParser, annotate_pdf_metadata
6 ap = ArgParser(['rectlattice2d', 'const_real_background', 'single_particle', 'single_lMax']) # const_real_background needed for calculation of the diffracted orders
7 ap.add_argument("-k", nargs=2, type=float, required=True, help='k vector', metavar=('K_X', 'K_Y'))
8 ap.add_argument("--kpi", action='store_true', help="Indicates that the k vector is given in natural units instead of SI, i.e. the arguments given by -k shall be automatically multiplied by pi / period (given by -p argument)")
9 ap.add_argument("--rank-tol", type=float, required=False)
10 ap.add_argument("-o", "--output", type=str, required=False, help='output path (if not provided, will be generated automatically)')
11 ap.add_argument("-t", "--rank-tolerance", type=float, default=1e11)
12 ap.add_argument("-c", "--min-candidates", type=int, default=1, help='always try at least this many eigenvalue candidates, even if their SVs in the rank tests are lower than rank_tolerance')
13 ap.add_argument("-T", "--residual-tolerance", type=float, default=2.)
14 ap.add_argument("-b", "--band-index", type=int, required=True, help="Argument's absolute value determines the empty lattice band order (counted from 1), -/+ determines whether the eigenvalues are searched below/above that empty lattice band.")
15 ap.add_argument("-f", "--interval-factor", type=float, default=0.1)
16 ap.add_argument("-N", type=int, default="150", help="Integration contour discretisation size")
17 ap.add_argument("-i", "--imaginary-aspect-ratio", type=float, default=1, help="Aspect ratio of the integration contour (Im/Re)")
18 ap.add_argument("-O", "--plot-out", type=str, required=False, help="path to plot output (optional)")
19 ap.add_argument("-P", "--plot", action='store_true', help="if -p not given, plot to a default path")
21 a=ap.parse_args()
23 px, py = a.period
25 if a.kpi:
26 a.k[0] *= math.pi/px
27 a.k[1] *= math.pi/py
29 particlestr = ("sph" if a.height is None else "cyl") + ("_r%gnm" % (a.radius*1e9))
30 if a.height is not None: particlestr += "_h%gnm" % (a.height * 1e9)
31 defaultprefix = "%s_p%gnmx%gnm_m%s_n%g_b%+d_k(%g_%g)um-1_L%d_cn%d" % (
32 particlestr, px*1e9, py*1e9, str(a.material), a.refractive_index, a.band_index, a.k[0]*1e-6, a.k[1]*1e-6, a.lMax, a.N)
34 import logging
35 logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
37 import numpy as np
38 import qpms
39 from qpms.cybspec import BaseSpec
40 from qpms.cytmatrices import CTMatrix
41 from qpms.qpms_c import Particle, ScatteringSystem, empty_lattice_modes_xy
42 from qpms.cymaterials import EpsMu, EpsMuGenerator, LorentzDrudeModel, lorentz_drude
43 from qpms.constants import eV, hbar
44 eh = eV/hbar
46 def inside_ellipse(point_xy, centre_xy, halfaxes_xy):
47 x = point_xy[0] - centre_xy[0]
48 y = point_xy[1] - centre_xy[1]
49 ax = halfaxes_xy[0]
50 ay = halfaxes_xy[1]
51 return ((x/ax)**2 + (y/ay)**2) <= 1
53 a1 = ap.direct_basis[0]
54 a2 = ap.direct_basis[1]
56 #Particle positions
57 orig_x = [0]
58 orig_y = [0]
59 orig_xy = np.stack(np.meshgrid(orig_x,orig_y),axis=-1)
61 if a.material in lorentz_drude:
62 emg = EpsMuGenerator(lorentz_drude[a.material])
63 else: # constant refractive index
64 emg = EpsMuGenerator(EpsMu(a.material**2))
66 beta = np.array(a.k)
68 empty_freqs = empty_lattice_modes_xy(ap.background_epsmu, ap.reciprocal_basis2pi, np.array([0,0]), 1)
69 empty_freqs = empty_lattice_modes_xy(ap.background_epsmu, ap.reciprocal_basis2pi, beta, (1+abs(a.band_index)) * empty_freqs[1])
71 # make the frequencies in the list unique
72 empty_freqs = list(empty_freqs)
73 i = 0
74 while i < len(empty_freqs) - 1:
75 if math.isclose(empty_freqs[i], empty_freqs[i+1], rel_tol=1e-13):
76 del empty_freqs[i+1]
77 else:
78 i += 1
80 logging.info("Empty freqs: %s", str(empty_freqs))
81 if a.band_index > 0:
82 top = empty_freqs[a.band_index]
83 bottom = empty_freqs[a.band_index - 1]
84 lebeta_om = bottom
85 else: # a.band_index < 0
86 top = empty_freqs[abs(a.band_index) - 1]
87 bottom = empty_freqs[abs(a.band_index) - 2] if abs(a.band_index) > 1 else 0.
88 lebeta_om = top
89 #print(top,bottom,lebeta_om)
90 freqradius = .5 * (top - bottom) * a.interval_factor
92 centfreq = bottom + freqradius if a.band_index > 0 else top - freqradius
94 bspec = BaseSpec(lMax = a.lMax)
95 pp = Particle(orig_xy[0][0], t = ap.tmgen, bspec=bspec)
97 ss, ssw = ScatteringSystem.create([pp], ap.background_emg, centfreq, latticebasis = ap.direct_basis)
99 if freqradius == 0:
100 raise ValueError("Integration contour radius is set to zero. Are you trying to look below the lowest empty lattice band at the gamma point?")
102 freqradius *= (1-1e-13) # to not totally touch the singularities
104 with qpms.pgsl_ignore_error(15):
105 res = ss.find_modes(centfreq, freqradius, freqradius * a.imaginary_aspect_ratio,
106 blochvector = a.k, contour_points = a.N, rank_tol = a.rank_tolerance,
107 res_tol = a.residual_tolerance, rank_min_sel = a.min_candidates)
109 logging.info("Eigenfrequencies found: %s" % str(res['eigval']))
111 res['inside_contour'] = inside_ellipse((res['eigval'].real, res['eigval'].imag),
112 (centfreq.real, centfreq.imag), (freqradius, freqradius * a.imaginary_aspect_ratio))
114 res['refractive_index_internal'] = [emg(om).n for om in res['eigval']]
116 #del res['omega'] If contour points are not needed...
117 #del res['ImTW'] # not if dbg=false anyway
118 outfile = defaultprefix + ".npz" if a.output is None else a.output
119 np.savez(outfile, meta={**vars(a), 'qpms_version' : qpms.__version__()}, empty_freqs=np.array(empty_freqs),
120 ss_positions=ss.positions, ss_fullvec_poffsets=ss.fullvec_poffsets,
121 ss_fullvec_psizes=ss.fullvec_psizes,
122 ss_bspecs_flat = np.concatenate(ss.bspecs),
123 ss_lattice_basis=ss.lattice_basis, ss_reciprocal_basis = ss.reciprocal_basis,
124 **res)
125 logging.info("Saved to %s" % outfile)
128 if a.plot or (a.plot_out is not None):
129 imcut = np.linspace(0, -freqradius)
130 recut1 = np.sqrt(lebeta_om**2+imcut**2) # incomplete Gamma-related cut
131 recut2 = np.sqrt((lebeta_om/2)**2-imcut**2) + lebeta_om/2 # odd-power-lilgamma-related cut
133 import matplotlib
134 matplotlib.use('pdf')
135 from matplotlib import pyplot as plt
136 from matplotlib.backends.backend_pdf import PdfPages
138 fig = plt.figure()
139 ax = fig.add_subplot(111,)
140 #ax.plot(res['omega'].real/eh, res['omega'].imag/eh*1e3, ':') #res['omega'] not implemented in ScatteringSystem
141 ax.add_artist(matplotlib.patches.Ellipse((centfreq.real/eh, centfreq.imag/eh*1e3),
142 2*freqradius/eh, 2*freqradius*a.imaginary_aspect_ratio/eh*1e3, fill=False,
143 ls=':'))
144 ax.scatter(x=res['eigval'].real/eh, y=res['eigval'].imag/eh*1e3 , c = res['inside_contour']
146 ax.plot(recut1/eh, imcut/eh*1e3)
147 ax.plot(recut2/eh, imcut/eh*1e3)
148 for i,om in enumerate(res['eigval']):
149 ax.annotate(str(i), (om.real/eh, om.imag/eh*1e3))
150 xmin = np.amin(res['eigval'].real)/eh
151 xmax = np.amax(res['eigval'].real)/eh
152 xspan = xmax-xmin
153 ymin = np.amin(res['eigval'].imag)/eh*1e3
154 ymax = np.amax(res['eigval'].imag)/eh*1e3
155 yspan = ymax-ymin
156 ax.set_xlim([xmin-.1*xspan, xmax+.1*xspan])
157 ax.set_ylim([ymin-.1*yspan, ymax+.1*yspan])
158 ax.set_xlabel('$\hbar \Re \omega / \mathrm{eV}$')
159 ax.set_ylabel('$\hbar \Im \omega / \mathrm{meV}$')
161 plotfile = defaultprefix + ".pdf" if a.plot_out is None else a.plot_out
162 with PdfPages(plotfile) as pdf:
163 pdf.savefig(fig)
164 annotate_pdf_metadata(pdf, scriptname="rectlat_simple_modes.py")
166 exit(0)