mpb_data.py

from __future__ import division

import math

import numpy as np
import meep as mp
from . import map_data
from . import MPBArray


class MPBData(object):

    TWOPI = 6.2831853071795864769252867665590057683943388

    def __init__(self,
                 lattice=None,
                 kpoint=None,
                 rectify=False,
                 x=0,
                 y=0,
                 z=0,
                 periods=0,
                 resolution=0,
                 phase_angle=0,
                 pick_nearest=False,
                 ve=None,
                 verbose=False):

        self.lattice = lattice
        self.kpoint = kpoint
        self.rectify = rectify

        if periods:
            self.multiply_size = [periods, periods, periods]
        else:
            self.multiply_size = [
                x if x else 1,
                y if y else 1,
                z if z else 1
            ]

        self.resolution = resolution
        self.phase_angle = phase_angle
        self.pick_nearest = pick_nearest
        self.ve = ve

        if self.ve:
            self.have_ve = True
            self.rectify = True
        else:
            self.have_ve = False
            self.ve = mp.Vector3(1, 0, 0)

        self.verbose = verbose
        self.scaleby = complex(1, 0)

        self.phase = complex(math.cos(self.TWOPI * self.phase_angle / 360.0),
                             math.sin(self.TWOPI * self.phase_angle / 360.0))
        self.scaleby *= self.phase

    def handle_dataset(self, in_arr):

        out_dims = [1, 1, 1]
        rank = len(in_arr.shape)
        num_ones = 3 - rank
        in_dims = [x for x in in_arr.shape] + [1] * num_ones

        if np.iscomplexobj(in_arr):
            in_arr_re = np.real(in_arr)
            in_arr_im = np.imag(in_arr)
        else:
            in_arr_re = in_arr
            in_arr_im = None

        if self.verbose:
            fmt = "Input data is rank {}, size {}x{}x{}."
            print(fmt.format(rank, in_dims[0], in_dims[1], in_dims[2]))

        if self.resolution > 0:
            out_dims[0] = math.floor(self.Rout.c1.norm() * self.resolution + 0.5)
            out_dims[1] = math.floor(self.Rout.c2.norm() * self.resolution + 0.5)
            out_dims[2] = math.floor(self.Rout.c3.norm() * self.resolution + 0.5)
        else:
            for i in range(3):
                out_dims[i] = in_dims[i] * self.multiply_size[i]

        for i in range(rank, 3):
            out_dims[i] = 1

        N = 1
        for i in range(3):
            out_dims[i] = int(max(out_dims[i], 1))
            N *= out_dims[i]

        if self.verbose:
            print("Output data {}x{}x{}".format(out_dims[0], out_dims[1], out_dims[2]))

        out_arr_re = np.zeros(int(N))

        if isinstance(in_arr_im, np.ndarray):
            out_arr_im = np.zeros(int(N))
        else:
            out_arr_im = np.array([])

        flat_in_arr_re = in_arr_re.ravel()
        flat_in_arr_im = in_arr_im.ravel() if isinstance(in_arr_im, np.ndarray) else np.array([])

        if self.kpoint:
            kvector = [self.kpoint.x, self.kpoint.y, self.kpoint.z]
        else:
            kvector = []

        map_data(flat_in_arr_re, flat_in_arr_im, np.array(in_dims, dtype=np.intc),
                 out_arr_re, out_arr_im, np.array(out_dims, dtype=np.intc), self.coord_map,
                 kvector, self.pick_nearest, self.verbose, False)

        if np.iscomplexobj(in_arr):
            # multiply * scaleby for complex data
            complex_out = np.vectorize(complex)(out_arr_re, out_arr_im)
            complex_out *= self.scaleby

            return np.reshape(complex_out, out_dims[:rank])

        return np.reshape(out_arr_re, out_dims[:rank])

    def handle_cvector_dataset(self, in_arr, multiply_bloch_phase):
        in_x_re = np.real(in_arr[:, :, :, 0]).ravel()
        in_x_im = np.imag(in_arr[:, :, :, 0]).ravel()
        in_y_re = np.real(in_arr[:, :, :, 1]).ravel()
        in_y_im = np.imag(in_arr[:, :, :, 1]).ravel()
        in_z_re = np.real(in_arr[:, :, :, 2]).ravel()
        in_z_im = np.imag(in_arr[:, :, :, 2]).ravel()

        d_in = [[in_x_re, in_x_im], [in_y_re, in_y_im], [in_z_re, in_z_im]]
        in_dims = [in_arr.shape[0], in_arr.shape[1], 1]
        rank = 2

        if self.verbose:
            print("Found complex vector dataset...")

        if self.verbose:
            fmt = "Input data is rank {}, size {}x{}x{}."
            print(fmt.format(rank, in_dims[0], in_dims[1], in_dims[2]))

        # rotate vector field according to cart_map
        if self.verbose:
            fmt1 = "Rotating vectors by matrix [ {:.10g}, {:.10g}, {:.10g}"
            fmt2 = "                             {:.10g}, {:.10g}, {:.10g}"
            fmt3 = "                             {:.10g}, {:.10g}, {:.10g} ]"
            print(fmt1.format(self.cart_map.c1.x, self.cart_map.c2.x, self.cart_map.c3.x))
            print(fmt2.format(self.cart_map.c1.y, self.cart_map.c2.y, self.cart_map.c3.y))
            print(fmt3.format(self.cart_map.c1.z, self.cart_map.c2.z, self.cart_map.c3.z))

        N = in_dims[0] * in_dims[1]
        for ri in range(2):
            for i in range(N):
                v = mp.Vector3(d_in[0][ri][i], d_in[1][ri][i], d_in[2][ri][i])
                v = self.cart_map * v
                d_in[0][ri][i] = v.x
                d_in[1][ri][i] = v.y
                d_in[2][ri][i] = v.z

        out_dims = [1, 1, 1]

        if self.resolution > 0:
            out_dims[0] = self.Rout.c1.norm() * self.resolution + 0.5
            out_dims[1] = self.Rout.c2.norm() * self.resolution + 0.5
            out_dims[2] = self.Rout.c3.norm() * self.resolution + 0.5
        else:
            for i in range(3):
                out_dims[i] = in_dims[i] * self.multiply_size[i]

        out_dims[2] = 1

        N = 1
        for i in range(3):
            out_dims[i] = int(max(out_dims[i], 1))
            N *= out_dims[i]

        if self.verbose:
            fmt = "Output data {}x{}x{}."
            print(fmt.format(out_dims[0], out_dims[1], out_dims[2]))

        if self.kpoint:
            kvector = [self.kpoint.x, self.kpoint.y, self.kpoint.z]
        else:
            kvector = []

        converted = []
        for dim in range(3):
            out_arr_re = np.zeros(int(N))
            out_arr_im = np.zeros(int(N))

            map_data(d_in[dim][0].ravel(), d_in[dim][1].ravel(), np.array(in_dims, dtype=np.intc),
                     out_arr_re, out_arr_im, np.array(out_dims, dtype=np.intc), self.coord_map,
                     kvector, self.pick_nearest, self.verbose, multiply_bloch_phase)

            # multiply * scaleby
            complex_out = np.vectorize(complex)(out_arr_re, out_arr_im)
            complex_out *= self.scaleby
            converted.append(complex_out)

        result = np.zeros(np.prod(out_dims) * 3, np.complex128)
        result[0::3] = converted[0]
        result[1::3] = converted[1]
        result[2::3] = converted[2]

        return np.reshape(result, (out_dims[0], out_dims[1], 3))

    def init_output_lattice(self):

        cart_map = mp.Matrix(
            mp.Vector3(1, 0, 0),
            mp.Vector3(0, 1, 0),
            mp.Vector3(0, 0, 1)
        )

        Rin = mp.Matrix(
            mp.Vector3(*self.lattice[0]),
            mp.Vector3(*self.lattice[1]),
            mp.Vector3(*self.lattice[2])
        )

        if self.verbose:
            print("Read lattice vectors")
            if self.kpoint:
                fmt = "Read Bloch wavevector ({:.6g}, {:.6g}, {:.6g})"
                print(fmt.format(self.kpoint.x, self.kpoint.y, self.kpoint.z))
            fmt = "Input lattice = ({:.6g}, {:.6g}, {:.6g}), ({:.6g}, {:.6g}, {:.6g}), ({:.6g}, {:.6g}, {:.6g})"
            print(fmt.format(Rin.c1.x, Rin.c1.y, Rin.c1.z,
                             Rin.c2.x, Rin.c2.y, Rin.c2.z,
                             Rin.c3.x, Rin.c3.y, Rin.c3.z))

        Rout = mp.Matrix(Rin.c1, Rin.c2, Rin.c3)

        if self.rectify:
            # Orthogonalize the output lattice vectors.  If have_ve is true,
            # then the first new lattice vector should be in the direction
            # of the ve unit vector; otherwise, the first new lattice vector
            # is the first original lattice vector. Note that we do this in
            # such a way as to preserve the volume of the unit cell, and so
            # that our first vector (in the direction of ve) smoothly
            # interpolates between the original lattice vectors.

            if self.have_ve:
                ve = self.ve.unit()
            else:
                ve = Rout.c1.unit()

            # First, compute c1 in the direction of ve by smoothly
            # interpolating the old c1/c2/c3 (formula is slightly tricky)
            V = Rout.c1.cross(Rout.c2).dot(Rout.c3)
            Rout.c2 = Rout.c2 - Rout.c1
            Rout.c3 = Rout.c3 - Rout.c1
            Rout.c1 = ve.scale(V / Rout.c2.cross(Rout.c3).dot(ve))

            # Now, orthogonalize c2 and c3
            Rout.c2 = Rout.c2 - ve.scale(ve.dot(Rout.c2))
            Rout.c3 = Rout.c3 - ve.scale(ve.dot(Rout.c3))
            Rout.c3 = Rout.c3 - Rout.c2.scale(Rout.c2.dot(Rout.c3) /
                                              Rout.c2.dot(Rout.c2))

            cart_map.c1 = Rout.c1.unit()
            cart_map.c2 = Rout.c2.unit()
            cart_map.c3 = Rout.c3.unit()
            cart_map = cart_map.inverse()

        Rout.c1 = Rout.c1.scale(self.multiply_size[0])
        Rout.c2 = Rout.c2.scale(self.multiply_size[1])
        Rout.c3 = Rout.c3.scale(self.multiply_size[2])

        if self.verbose:
            fmt = "Output lattice = ({:.6g}, {:.6g}, {:.6g}), ({:.6g}, {:.6g}, {:.6g}), ({:.6g}, {:.6g}, {:.6g})"
            print(fmt.format(Rout.c1.x, Rout.c1.y, Rout.c1.z,
                             Rout.c2.x, Rout.c2.y, Rout.c2.z,
                             Rout.c3.x, Rout.c3.y, Rout.c3.z))

        self.coord_map = Rin.inverse() * Rout
        self.Rout = Rout
        self.cart_map = cart_map

    def convert(self, arr, kpoint=None):
        if isinstance(arr, MPBArray):
            self.lattice = arr.lattice
            self.kpoint = arr.kpoint

        if self.lattice is None:
            err = ("Couldn't find 'lattice.' You must do one of the following:\n" +
                   "  1. Pass the ModeSolver lattice to the MPBData constructor\n" +
                   "     i.e., MPBData(lattice=ms.get_lattice())\n" +
                   "  2. Create an MPBArray to pass to MPBData.convert()\n" +
                   "     i.e., mpb_arr = MPBArray(arr, ms.get_lattice(), ... ); mpb_data.convert(mpb_arr))")
            raise ValueError(err)

        if kpoint:
            self.kpoint = kpoint

        self.init_output_lattice()

        if len(arr.shape) == 4:
            return self.handle_cvector_dataset(arr, not arr.bloch_phase)
        else:
            return self.handle_dataset(arr)