File: merge_ddb_nc.py

package info (click to toggle)
abinit 9.10.4-3
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 518,712 kB
  • sloc: xml: 877,568; f90: 577,240; python: 80,760; perl: 7,019; ansic: 4,585; sh: 1,925; javascript: 601; fortran: 557; cpp: 454; objc: 323; makefile: 77; csh: 42; pascal: 31
file content (153 lines) | stat: -rwxr-xr-x 4,948 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
#! /usr/bin/python

#
#    Copyright (C) 2003-2022 ABINIT group
#
#    Written by Gabriel Antonius in python (compatible v2.7).
#    This is free software, and you are welcome to redistribute it
#    under certain conditions (GNU General Public License,
#    see ~abinit/COPYING or http://www.gnu.org/copyleft/gpl.txt).
#
#    ABINIT is a project of the Universite Catholique de Louvain,
#    Corning Inc. and other collaborators, see ~abinit/doc/developers/contributors.txt.
#    Please read ~abinit/doc/biblio/generated_files/bib_acknow.html for suggested
#    acknowledgments of the ABINIT effort.
#
#    For more information, see https://www.abinit.org .

"""
This script can be run interactively,
but it is recommended to import it as a module:

    >>> from merge_ddb_nc import merge_ddb_nc
    >>> merge_ddb_nc(out_fname, fnames)

"""

from __future__ import print_function
import numpy as np
import netCDF4 as nc

__version__ = '1.0.0'

def merge_ddb_nc(out_fname, fnames):
    """
    Merge a list of DDB.nc files containing different elements of the same qpoint.

    Arguments
    ---------

    out_fname: Name for the merged file (will overwrite any existing file).
    fnames: List of DDB.nc files.

    """
    if not fnames:
        raise Exception('Empty list of files given for merge')

    fname0 = fnames.pop(0)

    with nc.Dataset(out_fname, 'w') as dsout:
        with nc.Dataset(fname0, 'r') as dsin:
            nc_copy(dsin, dsout)
            q0 = dsin.variables[u'q_point_reduced_coord'][...]

        for fname in fnames:
            with nc.Dataset(fname, 'r') as dsin:

                # Check that the qpoints are the same
                q = dsin.variables[u'q_point_reduced_coord'][...]
                if not all(np.isclose(q0, q)):
                    raise Exception('Cannot merge DDB.nc at different q-points.')

                # Merge dynamical matrix
                dynmat = dsin.variables[u'second_derivative_of_energy'][...]
                dynmat_mask = dsin.variables[u'second_derivative_of_energy_mask'][...]

                out_dynmat = dsin.variables[u'second_derivative_of_energy']
                out_dynmat_mask = dsin.variables[u'second_derivative_of_energy_mask']

                ni,nj,nk,nl = dynmat_mask.shape

                for i in range(ni):
                 for j in range(nj):
                  for k in range(nk):
                   for l in range(nl):
                    if dynmat_mask[i,j,k,l]:
                        dsout.variables[u'second_derivative_of_energy'][i,j,k,l,:] = (
                            dynmat[i,j,k,l,:])

                        dsout.variables[u'second_derivative_of_energy_mask'][i,j,k,l] = (
                            dynmat_mask[i,j,k,l])

                # Born effective charge tensor
                BECT = dsin.variables[u'born_effective_charge_tensor'][...]
                BECT_mask = dsin.variables[u'born_effective_charge_tensor_mask'][...]

                ni,nj,nk = BECT_mask.shape

                for i in range(ni):
                 for j in range(nj):
                  for k in range(nk):
                    if BECT_mask[i,j,k]:
                        dsout.variables[u'born_effective_charge_tensor'][i,j,k] = (
                            BECT[i,j,k])

                        dsout.variables[u'born_effective_charge_tensor_mask'][i,j,k] = (
                            BECT_mask[i,j,k])


def nc_copy(dsin, dsout):
    """
    Copy all dimensions and variable of one nc.Dataset instance into another.
    """

    #Copy dimensions
    for dname, dim in dsin.dimensions.iteritems():
        dsout.createDimension(dname, len(dim))

    #Copy variables
    for vname, varin in dsin.variables.iteritems():
        outVar = dsout.createVariable(vname, varin.datatype, varin.dimensions)
        outVar[...] = varin[...]


def interactive_merge_ddb_nc():
    """Get inputs from the user and run merge_ddb_nc."""

    program_name = 'merge_ddb_nc'
    description = """Merge several DDB.nc files, belonging to the same q-point."""

    def get_user(s):
        return raw_input(s.rstrip() + '\n').split('#')[0]


    print(program_name)
    print(len(program_name) * '-')
    print(description + '\n')

    ui = get_user('Enter a name for the output file in which to merge (will overwrite any existing file):')
    out_fname = str(ui)

    ui = get_user('Enter the number of files to merge:')
    nfiles = int(ui)

    fnames = list()
    for i in range(nfiles):
        ui = get_user('Enter the name of file {}:'.format(i+1))
        fname = str(ui)
        fnames.append(fname)

    # Main execution
    print('Executing...')
    merge_ddb_nc(out_fname, fnames)

    print('All done.')


# =========================================================================== #
# Run interactive program
# =========================================================================== #

if __name__ == '__main__':
    interactive_merge_ddb_nc()