qpms/lattices2d.py

   1 """
   2 These functions are mostly deprecated by the C counterparts from lattices.h.
   3 This file is still kept for reference, but might be removed in the future.
   4 """
   5 import numpy as np
   6 from enum import Enum
   7 from math import floor
   8
   9 nx = None
  10
  11 class LatticeType(Enum):
  12     """
  13     All the five Bravais lattices in 2D
  14     """
  15     OBLIQUE=1
  16     RECTANGULAR=2
  17     SQUARE=4
  18     RHOMBIC=5
  19     EQUILATERAL_TRIANGULAR=3
  20     RIGHT_ISOSCELES=SQUARE
  21     PARALLELOGRAMMIC=OBLIQUE
  22     CENTERED_RHOMBIC=RECTANGULAR
  23     RIGHT_TRIANGULAR=RECTANGULAR
  24     CENTERED_RECTANGULAR=RHOMBIC
  25     ISOSCELE_TRIANGULAR=RHOMBIC
  26     RIGHT_ISOSCELE_TRIANGULAR=SQUARE
  27     HEXAGONAL=EQUILATERAL_TRIANGULAR
  28
  29 def reduceBasisSingle(b1, b2):
  30     """
  31     Lagrange-Gauss reduction of a 2D basis.
  32     cf. https://www.math.auckland.ac.nz/~sgal018/crypto-book/ch17.pdf
  33     inputs and outputs are (2,)-shaped numpy arrays
  34     The output shall satisfy |b1| <= |b2| <= |b2 - b1|
  35     TODO doc
  36
  37     TODO perhaps have the (on-demand?) guarantee of obtuse angle between b1, b2?
  38     TODO possibility of returning the (in-order, no-obtuse angles) b as well?
  39     """
  40     b1 = np.array(b1)
  41     b2 = np.array(b2)
  42     if b1.shape != (2,) or b2.shape != (2,):
  43         raise ValueError('Shape of b1 and b2 must be (2,)')
  44     B1 = np.sum(b1 * b1, axis=-1, keepdims=True)
  45     mu = np.sum(b1 * b2, axis=-1, keepdims=True) / B1
  46     b2 = b2 - np.rint(mu) * b1
  47     B2 = np.sum(b2 * b2, axis=-1, keepdims=True)
  48     while(np.any(B2 < B1)):
  49         b2t = b1
  50         b1 = b2
  51         b2 = b2t
  52         B1 = B2
  53         mu = np.sum(b1 * b2, axis=-1, keepdims=True) / B1
  54         b2 = b2 - np.rint(mu) * b1
  55         B2 = np.sum(b2*b2, axis=-1, keepdims=True)
  56     return np.array((b1,b2))
  57
  58 def shortestBase3(b1, b2):
  59     '''
  60     returns the "ordered shortest triple" of base vectors (each pair from
  61     the triple is a base) and there may not be obtuse angle between b1, b2
  62     and between b2, b3
  63     '''
  64     b1, b2 = reduceBasisSingle(b1,b2)
  65     if is_obtuse(b1, b2, tolerance=0):
  66         b3 = b2
  67         b2 = b2 + b1
  68     else:
  69         b3 = b2 - b1
  70     return (b1, b2, b3)
  71
  72 def shortestBase46(b1, b2, tolerance=1e-13):
  73     b1, b2 = reduceBasisSingle(b1,b2)
  74     b1s = np.sum(b1 ** 2)
  75     b2s = np.sum(b2 ** 2)
  76     b3 = b2 - b1
  77     b3s = np.sum(b3 ** 2)
  78     eps = tolerance * (b2s + b1s)
  79     if abs(b3s - b2s - b1s) < eps:
  80         return(b1, b2, -b1, -b2)
  81     else:
  82         if b3s - b2s - b1s > eps: #obtuse
  83             b3 = b2
  84             b2 = b2 + b1
  85         return (b1, b2, b3, -b1, -b2, -b3)
  86
  87
  88 def is_obtuse(b1, b2, tolerance=1e-13):
  89     b1s = np.sum(b1 ** 2)
  90     b2s = np.sum(b2 ** 2)
  91     b3 = b2 - b1
  92     b3s = np.sum(b3 ** 2)
  93     eps = tolerance * (b2s + b1s)
  94     return (b3s - b2s - b1s > eps)
  95
  96 def classifyLatticeSingle(b1, b2, tolerance=1e-13):
  97     """
  98     Given two basis vectors, returns 2D Bravais lattice type.
  99     Tolerance is relative.
 100     TODO doc
 101     """
 102     b1, b2 = reduceBasisSingle(b1, b2)
 103     b1s = np.sum(b1 ** 2)
 104     b2s = np.sum(b2 ** 2)
 105     b3 = b2 - b1
 106     b3s = np.sum(b3 ** 2)
 107     eps = tolerance * (b2s + b1s)
 108     # Avoid obtuse angle between b1 and b2. TODO This should be yet thoroughly tested.
 109     # TODO use is_obtuse here?
 110     if b3s - b2s - b1s > eps:
 111         b3 = b2
 112         b2 = b2 + b1
 113         # N. B. now the assumption |b3| >= |b2| is no longer valid
 114         #b3 = b2 - b1
 115         b2s = np.sum(b2 ** 2)
 116         b3s = np.sum(b3 ** 2)
 117     if abs(b2s - b1s) < eps or abs(b2s - b3s) < eps: # isoscele
 118         if abs(b3s - b1s) < eps:
 119             return LatticeType.EQUILATERAL_TRIANGULAR
 120         elif abs(b3s - 2 * b1s) < eps:
 121             return LatticeType.SQUARE
 122         else:
 123             return LatticeType.RHOMBIC
 124     elif abs(b3s - b2s - b1s) < eps:
 125         return LatticeType.RECTANGULAR
 126     else:
 127         return LatticeType.OBLIQUE
 128
 129 def range2D(maxN, mini=1, minj=0, minN = 0):
 130     """
 131     "Triangle indices"
 132     Generates pairs of non-negative integer indices (i, j) such that
 133     minN ≤ i + j ≤ maxN, i ≥ mini, j ≥ minj.
 134     TODO doc and possibly different orderings
 135     """
 136     for maxn in range(min(mini, minj, minN), floor(maxN+1)): # i + j == maxn
 137         for i in range(mini, maxn + 1):
 138             yield (i, maxn - i)
 139
 140
 141 def generateLattice(b1, b2, maxlayer=5, include_origin=False, order='leaves'):
 142     bvs = shortestBase46(b1, b2)
 143     cc = len(bvs) # "corner count"
 144
 145     if order == 'leaves':
 146         indices = np.array(list(range2D(maxlayer)))
 147         ia = indices[:,0]
 148         ib = indices[:,1]
 149         cc = len(bvs) # 4 for square/rec,
 150         leaves = list()
 151         if include_origin: leaves.append(np.array([[0,0]]))
 152         for c in range(cc):
 153             ba = bvs[c]
 154             bb = bvs[(c+1)%cc]
 155             leaves.append(ia[:,nx]*ba + ib[:,nx]*bb)
 156         return np.concatenate(leaves)
 157     else:
 158         raise ValueError('Lattice point order not implemented: ', order)
 159
 160 def generateLatticeDisk(b1, b2, r, include_origin=False, order='leaves'):
 161     b1, b2 = reduceBasisSingle(b1,b2)
 162     blen = np.linalg.norm(b1, ord=2)
 163     maxlayer = 2*r/blen # FIXME kanon na vrabce? Nestačí odmocnina ze 2?
 164     points = generateLattice(b1,b2, maxlayer=maxlayer, include_origin=include_origin, order=order)
 165     mask = (np.linalg.norm(points, axis=-1, ord=2) <= r)
 166     return points[mask]
 167
 168 def cellCornersWS(b1, b2,):
 169     """
 170     Given basis vectors, returns the corners of the Wigner-Seitz unit cell
 171     (w1, w2, -w1, w2) for rectangular and square lattice or
 172     (w1, w2, w3, -w1, -w2, -w3) otherwise
 173     """
 174     def solveWS(v1, v2):
 175         v1x = v1[0]
 176         v1y = v1[1]
 177         v2x = v2[0]
 178         v2y = v2[1]
 179         lsm = ((-v1y, v2y), (v1x, -v2x))
 180         rs = ((v1x-v2x)/2, (v1y - v2y)/2)
 181         t = np.linalg.solve(lsm, rs)
 182         return np.array(v1)/2 + t[0]*np.array((v1y, -v1x))
 183     b1, b2 = reduceBasisSingle(b1, b2)
 184     latticeType = classifyLatticeSingle(b1, b2)
 185     if latticeType is LatticeType.RECTANGULAR or latticeType is LatticeType.SQUARE:
 186         return np.array( (
 187             (+b1+b2),
 188             (+b2-b1),
 189             (-b1-b2),
 190             (-b2+b1),
 191         )) / 2
 192     else:
 193         bvs = shortestBase46(b1,b2,tolerance=0)
 194         return np.array([solveWS(bvs[i], bvs[(i+1)%6]) for i in range(6)])
 195
 196 def cutWS(points, b1, b2, scale=1., tolerance=1e-13):
 197     """
 198     From given points, return only those that are inside (or on the edge of)
 199     the Wigner-Seitz cell of a (scale*b1, scale*b2)-based lattice.
 200     """
 201     # TODO check input dimensions?
 202     bvs = shortestBase46(b1, b2)
 203     points = np.array(points)
 204     for b in bvs:
 205         mask = (np.tensordot(points, b, axes=(-1, 0)) <=  (scale * (1+tolerance) / 2) *np.linalg.norm(b, ord=2)**2 )
 206         points = points[mask]
 207     return points
 208
 209 def filledWS(b1, b2, density=10, scale=1.):
 210     """
 211     TODO doc
 212     TODO more intelligent generation, anisotropy balancing etc.
 213     """
 214     points = generateLattice(b1,b2,maxlayer=density*scale, include_origin=True)
 215     points = cutWS(points/density, np.array(b1)*scale, np.array(b2)*scale)
 216     return points
 217
 218
 219 def reciprocalBasis1(*pargs):
 220     a = reduceBasisSingle(*pargs)
 221     return np.linalg.inv(a).T
 222
 223 def reciprocalBasis2pi(*pargs):
 224     return 2*np.pi*reciprocalBasis1(*pargs)
 225
 226 # TODO fill it with "points from reciprocal space" instead
 227 def filledWS2(b1,b2, density=10, scale=1.):
 228     b1, b2 = reduceBasisSingle(b1,b2)
 229     b1r, b2r = reciprocalBasis2pi(b1,b2)
 230     b1l = np.linalg.norm(b1, ord=2)
 231     b2l = np.linalg.norm(b2, ord=2)
 232     b1rl = np.linalg.norm(b1r, ord=2)
 233     b2rl = np.linalg.norm(b2r, ord=2)
 234     # Black magick. Think later.™ Really. FIXME
 235     sicher_ratio = np.maximum(b1rl/b2rl, b2rl/b1rl) * np.maximum(b1l/b2l, b2l/b1l) # This really has to be adjusted
 236     points = generateLattice(b1r,b2r,maxlayer=density*scale*sicher_ratio, include_origin=True)
 237     points = cutWS(points*b1l/b1rl/density, b1*scale, b2*scale)
 238     return points
 239
 240
 241 def change_basis(srcbasis, destbasis, srccoords, srccoordsaxis=-1, lattice=True):
 242     srcbasis = np.array(srcbasis)
 243     destbasis = np.array(destbasis)
 244     trmatrix = np.dot(np.linalg.inv(np.transpose(destbasis)), np.transpose(srcbasis))
 245     if lattice: # if srcbasis and destbasis are two bases of the same lattice, its elements are ints
 246         otrmatrix = trmatrix
 247         trmatrix = np.round(trmatrix)
 248         if not np.all(np.isclose(trmatrix, otrmatrix)):
 249             raise ValueError("Given srcbasis and destbasis are not bases"
 250                 "of the same lattice", srcbasis, destbasis, trmatrix-otrmatrix)
 251     destcoords = np.tensordot(srccoords, trmatrix, axes=(srccoordsaxis, -1))
 252     return destcoords
 253
 254 """
 255 TODO
 256 ====
 257
 258 - DOC!!!!!
 259 - (nehoří) výhledově pořešit problém „hodně anisotropních“ mřížek (tj. kompensovat
 260 rozdílné délky základních vektorů).
 261
 262 """