TB2J 0.9.0.1__py3-none-any.whl → 0.9.1rc0__py3-none-any.whl

TB2J/cut_cell.py

@@ -0,0 +1,82 @@
+import numpy as np
+from TB2J.supercell import find_primitive_cell, map_to_primitive
+from TB2J.io_exchange.io_exchange import SpinIO, gen_distance_dict
+
+
+def cut_cell(path, output_path, sc_matrix, origin_atom_id, thr=1e-5):
+    """
+    Cut the exchange parameters
+    :param path: the original TB2J_results path
+    :param output_path: the output path.
+    :param sc_matrix: the matrix which maps the primitive cell to supercell.
+    :param origin: the origin of the primitive cell.
+    :param thr: the atoms which the reduced position is within -thr to 1.0+thr are considered as inside the primitive atoms
+    :returns:
+    """
+    sc_matrix = np.asarray(sc_matrix, dtype=int)
+    sc_excparams = SpinIO.load_pickle(path=path, fname='TB2J.pickle')
+    sc_atoms = sc_excparams.atoms
+    uc_atoms, ids = find_primitive_cell(sc_atoms,
+                                        sc_matrix,
+                                        origin_atom_id=origin_atom_id,
+                                        thr=thr,
+                                        perfect=False)
+    uc_charges = sc_excparams.charges[ids]
+    uc_spinat = sc_excparams.spinat[ids]
+    indmap, Rmap = map_to_primitive(sc_atoms, uc_atoms)
+
+    # TODO index_spin: {iatom: ispin}
+
+    # list of iatom for each spin index.
+    uc_index_spin = [sc_excparams.index_spin[i] for i in ids]
+
+    uc_ind_atoms = {}
+    for iatom, ispin in enumerate(uc_index_spin):
+        if ispin >= 0:
+            uc_ind_atoms[ispin] = iatom
+
+    uc_Jdict = {}
+    uc_Rset = set()
+    for key, val in sc_excparams.exchange_Jdict.items():
+        R, ispin, jspin = key
+        iatom = sc_excparams.ind_atoms[ispin]
+        jatom = sc_excparams.ind_atoms[jspin]
+        uc_ispin = uc_index_spin[indmap[iatom]]
+        uc_jspin = uc_index_spin[indmap[jatom]]
+        uc_R = R @ sc_matrix + Rmap[jatom] - Rmap[iatom]
+        uc_R = tuple(uc_R)
+        if iatom in ids:
+            #print(f"{iatom=}, {indmap[iatom]=},{uc_ispin=}, {uc_jspin=}, {uc_R=} ")
+            uc_Jdict[(uc_R, uc_ispin, uc_jspin)] = val
+            uc_Rset.add(uc_R)
+
+    uc_distance_dict = gen_distance_dict(uc_ind_atoms, uc_atoms, list(uc_Rset))
+
+    assert sc_excparams.colinear, "Cut supercell for non-collinear spin is not yet implemented."
+
+    uc_exc = SpinIO(uc_atoms,
+                    spinat=uc_spinat,
+                    charges=uc_charges,
+                    index_spin=uc_index_spin,
+                    colinear=sc_excparams.colinear,
+                    distance_dict=uc_distance_dict,
+                    exchange_Jdict=uc_Jdict,
+                    description="Cutted")
+
+    uc_exc.write_all(output_path)
+
+
+def run_cut_cell(
+    path="TB2J_results",
+    output_path="./TB2J_cutted",
+    sc_matrix=np.array([[1, 1, 0], [-1, 1, 0], [0, 0, 2]]),
+):
+    cut_cell(path,
+             output_path,
+             np.array(sc_matrix),
+             origin_atom_id=0,
+             thr=1e-19)
+
+
+if __name__ == "__main__":
+    run_cut_cell()

TB2J/green.py

@@ -7,6 +7,8 @@ import tempfile
 from pathos.multiprocessing import ProcessPool
 import sys
 import pickle
+import warnings
+from TB2J.mathutils.fermi import fermi
 
 MAX_EXP_ARGUMENT = np.log(sys.float_info.max)
 
@@ -26,19 +28,6 @@ def eigen_to_G(evals, evecs, efermi, energy):
     )
 
 
-def fermi(e, mu, width=0.01):
-    """
-    the fermi function.
-     .. math::
-        f=\\frac{1}{\exp((e-\mu)/width)+1}
-
-    :param e,mu,width: e,\mu,width
-    """
-
-    x = (e - mu) / width
-    return np.where(x < MAX_EXP_ARGUMENT, 1 / (1.0 + np.exp(x)), 0.0)
-
-
 def find_energy_ingap(evals, rbound, gap=4.0):
     """
     find a energy inside a gap below rbound (right bound),

TB2J/mathutils/__init__.py

@@ -0,0 +1 @@
+from .lowdin import Lowdin

TB2J/mathutils/fermi.py

@@ -0,0 +1,22 @@
+import numpy as np
+import warnings
+import sys
+
+MAX_EXP_ARGUMENT = np.log(sys.float_info.max)
+
+
+def fermi(e, mu, width=0.01):
+    """
+    the fermi function.
+     .. math::
+        f=\\frac{1}{\exp((e-\mu)/width)+1}
+
+    :param e,mu,width: e,\mu,width
+    """
+    x = (e - mu) / width
+    # disable overflow warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        ret = np.where(x < MAX_EXP_ARGUMENT, 1 / (1.0 + np.exp(x)), 0.0)
+
+    return ret

TB2J/mathutils/kR_convert.py

@@ -0,0 +1,90 @@
+import numpy as np
+
+
+def HR_to_k(HR, Rlist, kpts):
+    # Hk[k,:,:] = sum_R (H[R] exp(i2pi k.R))
+    phase = np.exp(2.0j * np.pi * np.tensordot(kpts, Rlist, axes=([1], [1])))
+    Hk = np.einsum("rlm, kr -> klm", HR, phase)
+    return Hk
+
+
+def Hk_to_R(Hk, Rlist, kpts, kweights):
+    phase = np.exp(-2.0j * np.pi * np.tensordot(kpts, Rlist, axes=([1], [1])))
+    HR = np.einsum("klm, kr, k->rlm", Hk, phase, kweights)
+    return HR
+
+
+def k_to_R(kpts, Rlist, Mk, kweights=None):
+    """
+    Transform k-space wavefunctions to real space.
+    params:
+        kpts: k-points
+        Rlist: list of R vectors
+        Mk: matrix of shape [nkpt, n1, n2] in k-space.
+
+    return:
+        MR: matrix of shape [nR, n1, n2], the matrix in R-space.
+
+    """
+    nkpt, n1, n2 = Mk.shape
+    if kweights is None:
+        kweights = np.ones(nkpt, dtype=float) / nkpt
+    phase = np.exp(-2.0j * np.pi * np.tensordot(kpts, Rlist, axes=([1], [1])))
+    MR = np.einsum("klm, kr, k -> rlm", Mk, phase, kweights)
+    return MR
+
+    # nkpt, n1, n2 = Mk.shape
+    # nR = Rlist.shape[0]
+    # MR = np.zeros((nR, n1, n2), dtype=complex)
+    # if kweights is None:
+    #    kweights = np.ones(nkpt, dtype=float)/nkpt
+    # for iR, R in enumerate(Rlist):
+    #    for ik in range(nkpt):
+    #        MR[iR] += Mk[ik] * np.exp(-2.0j*np.pi * np.dot(kpts[ik], R)) * kweights[ik]
+    # return MR
+
+
+def R_to_k(kpts, Rlist, MR):
+    """
+    Transform real-space wavefunctions to k-space.
+    params:
+        kpts: k-points
+        Rlist: list of R vectors
+        MR: matrix of shape [nR, n1, n2] in R-space.
+
+    return:
+        Mk: matrix of shape [nkpt, n1, n2], the matrix in k-space.
+
+    """
+    phase = np.exp(2.0 * np.pi * 1j * np.tensordot(kpts, Rlist, axes=([1], [1])))
+    Mk = np.einsum("rlm, kr -> klm", MR, phase)
+
+    # nkpt, n1, n2 = Mk.shape
+    # nR = Rlist.shape[0]
+    # Mk = np.zeros((nkpt, n1, n2), dtype=complex)
+    # for iR, R in enumerate(Rlist):
+    #    for ik in range(nkpt):
+    #        Mk[ik] += MR[iR] * np.exp(2.0 * np.pi * 1j * np.dot(kpts[ik], R))
+    return Mk
+
+
+def R_to_onek(kpt, Rlist, MR):
+    """
+    Transform real-space wavefunctions to k-space.
+    params:
+        kpt: k-point
+        Rlist: list of R vectors
+        MR: matrix of shape [nR, n1, n2] in R-space.
+
+    return:
+        Mk: matrix of shape [n1, n2], the matrix in k-space.
+
+    """
+    phase = np.exp(2.0j * np.pi * np.dot(Rlist, kpt))
+    Mk = np.einsum("rlm, r -> lm", MR, phase)
+    return Mk
+    # n1, n2 = MR.shape[1:]
+    #    Mk = np.zeros((n1, n2), dtype=complex)
+    #    for iR, R in enumerate(Rlist):
+    #        Mk += MR[iR] * np.exp(2.0j*np.pi * np.dot(kpt, R))
+    #    return Mk

TB2J/mathutils/lowdin.py

@@ -0,0 +1,12 @@
+import numpy as np
+from scipy.linalg import inv, eigh
+
+
+def Lowdin(S):
+    """
+    Calculate S^(-1/2).
+    Which is used in lowind's symmetric orthonormalization.
+    psi_prime = S^(-1/2) psi
+    """
+    eigval, eigvec = eigh(S)
+    return eigvec @ np.diag(np.sqrt(1.0 / eigval)) @ (eigvec.T.conj())

TB2J/mathutils/rotate_spin.py

@@ -0,0 +1,35 @@
+import numpy as np
+from TB2J.pauli import pauli_block_all, s0, s1, s2, s3, gather_pauli_blocks
+
+
+def rotate_Matrix_from_z_to_axis(M, axis, normalize=True):
+    """
+    Given a spinor matrix M, rotate it from z-axis to axis.
+    The spinor matrix M is a 2x2 matrix, which can be decomposed as I, x, y, z components using Pauli matrices.
+    """
+    MI, Mx, My, Mz = pauli_block_all(M)
+    axis = axis / np.linalg.norm(axis)
+    # M_new = s0* MI +  Mz * (axis[0] * s1 + axis[1] * s2 + axis[2] * s3) *2
+    M_new = gather_pauli_blocks(MI, Mz * axis[0], Mz * axis[1], Mz * axis[2])
+    return M_new
+
+
+def test_rotate_Matrix_from_z_to_axis():
+    M = np.array([[1.1, 0], [0, 0.9]])
+    print(pauli_block_all(M))
+    Mnew = rotate_Matrix_from_z_to_axis(M, [1, 1, 1])
+    print(pauli_block_all(Mnew))
+    print(Mnew)
+
+    M = np.array(
+        [
+            [-9.90532976e-06 + 0.0j, 0.00000000e00 + 0.0j],
+            [0.00000000e00 + 0.0j, -9.88431291e-06 + 0.0j],
+        ]
+    )
+    print(M)
+    print(rotate_Matrix_from_z_to_axis(M, [0, 0, 1]))
+
+
+if __name__ == "__main__":
+    test_rotate_Matrix_from_z_to_axis()

TB2J/pauli.py

@@ -143,7 +143,24 @@ def pauli_block_all(M):
     return MI, Mx, My, Mz
 
 
+def gather_pauli_blocks(MI, Mx, My, Mz):
+    """
+    Gather the I, x, y, z component of a matrix.
+    """
+    return np.kron(MI, s0) + np.kron(Mx, s1) + np.kron(My, s2) + np.kron(Mz, s3)
+
+
+def test_gather_pauli_blocks():
+    M = np.random.rand(4, 4)
+    MI, Mx, My, Mz = pauli_block_all(M)
+    M2 = gather_pauli_blocks(MI, Mx, My, Mz)
+    assert np.allclose(M, M2)
+
+
 def op_norm(M):
+    """
+    Return the operator norm of a matrix.
+    """
     return max(svd(M)[1])
 
 

TB2J/spinham/h_matrix.py

@@ -0,0 +1,68 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import aiida
+from aiida_tb2j.data import ExchangeData
+
+def plot_dispersion(bands, kpoint_labels, color='blue', title=None):
+
+    fig, axs = plt.subplots(1, 1, constrained_layout=True)
+    fig.set_size_inches(6, 6/1.618)
+
+    '''
+    Plot the bands
+    '''
+    kpoints = np.arange(len(bands))
+    axs.plot(kpoints, bands, color=color, linewidth=1.5)
+
+    '''
+    Plot the symmetry points
+    '''
+    bmin = bands.min(); bmax = bands.max()
+    ymin = bmin - 0.05*np.abs(bmin-bmax); ymax = bmax + 0.05*np.abs(bmax-bmin);
+    axs.set_xticks(kpoint_labels[0], kpoint_labels[1], fontsize=10)
+    axs.vlines(x=kpoint_labels[0], ymin=ymin, ymax=ymax, color='black', linewidth=0.5)
+    axs.set_xlim([0, len(kpoints)])
+    axs.set_ylim([ymin, ymax])
+
+    if title is not None:
+        plt.title(title, fontsize=10)
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument('-f', '--pickle_filename', type=str, help="Path of the 'TB2J.pickle' file.", required=True)
+    args = parser.parse_args()
+    
+    '''
+    Right now the implementation depends on AiiDA and so we must create and load an AiiDA profile,
+    even if we do not store any information on a data base.
+    '''
+    aiida.load_profile()
+    '''
+    Create an ExchangeData object with the informations from the TB2J.pickle file
+    '''
+    exchange = ExchangeData.load_tb2j(pickle_file=args.pickle_filename, isotropic=False, pbc=(True, True, True))
+    '''
+    Compute the magnon band structure along a high symmetry path generated with
+    the ASE package. The informations is stored in an AiiDA BandsData object.
+    Here tol is the symmetry tolerance to determine the space group of the system.
+    They are in units of eV
+    '''
+    magnon_data = exchange.get_magnon_bands(npoints=300, tol=1e-1, with_DMI=True, with_Jani=True)
+    magnon_bands = 1000*magnon_data.get_bands() # Convert to meV
+    raw_labels = [(k, '$\Gamma$') if s == 'GAMMA' else (k, s) for k, s in magnon_data.labels]
+    kpoint_labels = list( zip( *raw_labels ) )
+    plot_dispersion(magnon_bands, kpoint_labels, color='blue', title='Magnon Bands')
+    '''
+    We can also obtain the dynamical matrix h instead of the actual magnon bands. The result
+    is stored in a numpy array with shape (number of kpoints, 2*natoms, 2*natoms)
+    '''
+    kpoints = magnon_data.get_kpoints() # The shape of the kpoints must be (nkpoints, 3)
+    h_matrix = 1000*exchange._H_matrix(kpoints, with_DMI=True, with_Jani=True) # Convert to meV
+    h_dispersion = np.linalg.eigvalsh(h_matrix) # We can also get the eigenvectors with np.linalg.eigh
+    plot_dispersion(h_dispersion, kpoint_labels, color='red', title='h matrix dispersion')

TB2J/spinham/obtain_J.py

@@ -0,0 +1,79 @@
+import numpy as np
+from aiida_tb2j.data import ExchangeData
+from aiida_tb2j.data.exchange import get_rotation_arrays
+from itertools import combinations_with_replacement
+
+ux, uy, uz = np.eye(3).reshape((3, 1, 3))
+
+def combine_arrays(u, v):
+
+    return np.concatenate([u*v, np.roll(u, -1, axis=-1)*v, np.roll(v, -1, axis=-1)*u], axis=-1)
+
+def get_coefficients(magmoms, indices):
+
+    i, j = indices
+
+    U, V = zip(*[get_rotation_arrays(magmoms, u=u) for u in [ux, uy, uz]])
+    U = np.stack(U).swapaxes(0, 1)
+    V = np.stack(V).swapaxes(0, 1)
+
+    uc = combine_arrays(U[i], U[j].conj())
+    ur = combine_arrays(U[i], U[j])
+    uc2 = combine_arrays(U[i].conj(), U[j])
+    u = np.concatenate([uc, ur, uc2], axis=1)
+
+    return u, V
+
+def get_C(H0, u, V):
+
+    n = int(H0.shape[-1]/2)
+    upi = np.triu_indices(n)
+    dig = np.diag_indices(n)
+    
+    i, j = upi
+    AB0 = H0[:, [i, i, i+n], [j, j+n, j+n]]
+    AB0 = np.swapaxes(AB0, 0, 2).reshape(len(i), 9)
+    J0_flat = np.linalg.solve(u, AB0)
+
+    J0 = np.empty((n, n, 3, 3), dtype=complex)
+    J0[*upi] = J0_flat[:, [0, 6, 5, 3, 1, 7, 8, 4, 2]].reshape(-1, 3, 3)
+    J0 += J0.swapaxes(0, 1)
+    J0[*dig] = 0.0
+
+    C = np.array([np.diag(a) for a in np.einsum('imx,ijxy,jmy->mi', V, 2*J0, V)]) 
+
+    return C
+
+def get_J(H, kpoints, exchange):
+
+    n = int(H.shape[-1]/2)
+    upi = np.triu_indices(n)
+    dig = np.diag_indices(n)
+    i, j = upi
+
+    magmoms = exchange.magmoms()[np.unique(exchange.pairs)]
+    magmoms /= np.linalg.norm(magmoms, axis=-1).reshape(-1, 1)
+    u, V = get_coefficients(magmoms, indices=upi)
+
+    H0 = np.stack([1000*exchange._H_matrix(kpoints=np.zeros((1, 3)), with_DMI=True, with_Jani=True, u=u) for u in [ux, uy, uz]])[:, 0, :, :]
+    C = get_C(H0, u, V)
+    H[:, :, :n, :n] += C.reshape(3, 1, n, n)
+    H[:, :, n:, n:] += C.reshape(3, 1, n, n)
+
+    AB = H[:, :, [i, i, i+n], [j, j+n, j+n]]
+    AB[:, :, 2, :] = AB[:, ::-1, 2, :]
+    AB = np.moveaxis(AB, [2, 3], [1, 0]).reshape(len(i), 9, -1)
+
+    vectors = exchange.get_vectors()
+    exp_summand = np.exp( -2j*np.pi*vectors @ kpoints.T )
+    nAB = np.einsum('nik,ndk->nid', AB, exp_summand) / len(kpoints)
+    
+    ii = np.where(i == j)
+    i0 = np.where(np.linalg.norm(vectors, axis=-1) == 0.0)
+    J = np.linalg.solve(u, nAB).swapaxes(1, 2)
+    J = J[:, :, [0, 6, 5, 3, 1, 7, 8, 4, 2]].reshape(len(i), -1, 3, 3)
+    J *= -1
+    J[ii] *= 2
+    J[i0] *= 0
+
+    return J