Last change on this file since 574adc7 was 574adc7, checked in by Paul Kienzle <pkienzle@…>, 6 years ago

convert sascalc to python 2/3 syntax

• Property mode set to `100644`
File size: 39.2 KB
Line
1# pylint: disable=invalid-name
2"""
3SAS generic computation and sld file readers
4"""
5from __future__ import print_function
6
7import sas.sascalc.calculator.core.sld2i as mod
8from sas.sascalc.calculator.BaseComponent import BaseComponent
9from periodictable import formula
10from periodictable import nsf
11import numpy as np
12import os
13import copy
14import sys
15import logging
16
17logger = logging.getLogger(__name__)
18
19MFACTOR_AM = 2.853E-12
20MFACTOR_MT = 2.3164E-9
21METER2ANG = 1.0E+10
23NA = 6.02214129e+23
24
25def mag2sld(mag, v_unit=None):
26    """
27    Convert magnetization to magnatic SLD
28    sldm = Dm * mag where Dm = gamma * classical elec. radius/(2*Bohr magneton)
29    Dm ~ 2.853E-12 [A^(-2)] ==> Shouldn't be 2.90636E-12 [A^(-2)]???
30    """
31    if v_unit == "A/m":
32        factor = MFACTOR_AM
33    elif v_unit == "mT":
34        factor = MFACTOR_MT
35    else:
36        raise ValueError("Invalid valueunit")
37    sld_m = factor * mag
38    return sld_m
39
40def transform_center(pos_x, pos_y, pos_z):
41    """
42    re-center
43    :return: posx, posy, posz   [arrays]
44    """
45    posx = pos_x - (min(pos_x) + max(pos_x)) / 2.0
46    posy = pos_y - (min(pos_y) + max(pos_y)) / 2.0
47    posz = pos_z - (min(pos_z) + max(pos_z)) / 2.0
48    return posx, posy, posz
49
50class GenSAS(BaseComponent):
51    """
52    Generic SAS computation Model based on sld (n & m) arrays
53    """
54    def __init__(self):
55        """
56        Init
57        :Params sld_data: MagSLD object
58        """
59        # Initialize BaseComponent
60        BaseComponent.__init__(self)
61        self.sld_data = None
62        self.data_pos_unit = None
63        self.data_x = None
64        self.data_y = None
65        self.data_z = None
66        self.data_sldn = None
67        self.data_mx = None
68        self.data_my = None
69        self.data_mz = None
70        self.data_vol = None #[A^3]
71        self.is_avg = False
72        ## Name of the model
73        self.name = "GenSAS"
74        ## Define parameters
75        self.params = {}
76        self.params['scale'] = 1.0
77        self.params['background'] = 0.0
78        self.params['solvent_SLD'] = 0.0
79        self.params['total_volume'] = 1.0
80        self.params['Up_frac_in'] = 1.0
81        self.params['Up_frac_out'] = 1.0
82        self.params['Up_theta'] = 0.0
83        self.description = 'GenSAS'
84        ## Parameter details [units, min, max]
85        self.details = {}
86        self.details['scale'] = ['', 0.0, np.inf]
87        self.details['background'] = ['[1/cm]', 0.0, np.inf]
88        self.details['solvent_SLD'] = ['1/A^(2)', -np.inf, np.inf]
89        self.details['total_volume'] = ['A^(3)', 0.0, np.inf]
90        self.details['Up_frac_in'] = ['[u/(u+d)]', 0.0, 1.0]
91        self.details['Up_frac_out'] = ['[u/(u+d)]', 0.0, 1.0]
92        self.details['Up_theta'] = ['[deg]', -np.inf, np.inf]
93        # fixed parameters
94        self.fixed = []
95
96    def set_pixel_volumes(self, volume):
97        """
98        Set the volume of a pixel in (A^3) unit
99        :Param volume: pixel volume [float]
100        """
101        if self.data_vol is None:
102            raise TypeError("data_vol is missing")
103        self.data_vol = volume
104
105    def set_is_avg(self, is_avg=False):
106        """
107        Sets is_avg: [bool]
108        """
109        self.is_avg = is_avg
110
111    def _gen(self, x, y, i):
112        """
113        Evaluate the function
114        :Param x: array of x-values
115        :Param y: array of y-values
116        :Param i: array of initial i-value
117        :return: function value
118        """
119        pos_x = self.data_x
120        pos_y = self.data_y
121        pos_z = self.data_z
122        len_x = len(pos_x)
123        if self.is_avg is None:
124            len_x *= -1
125            pos_x, pos_y, pos_z = transform_center(pos_x, pos_y, pos_z)
126        len_q = len(x)
127        sldn = copy.deepcopy(self.data_sldn)
128        sldn -= self.params['solvent_SLD']
129        model = mod.new_GenI(len_x, pos_x, pos_y, pos_z,
130                             sldn, self.data_mx, self.data_my,
131                             self.data_mz, self.data_vol,
132                             self.params['Up_frac_in'],
133                             self.params['Up_frac_out'],
134                             self.params['Up_theta'])
135        if y == []:
136            mod.genicom(model, len_q, x, i)
137        else:
138            mod.genicomXY(model, len_q, x, y, i)
139        vol_correction = self.data_total_volume / self.params['total_volume']
140        return  self.params['scale'] * vol_correction * i + \
141                        self.params['background']
142
143    def set_sld_data(self, sld_data=None):
144        """
145        Sets sld_data
146        """
147        self.sld_data = sld_data
148        self.data_pos_unit = sld_data.pos_unit
149        self.data_x = sld_data.pos_x
150        self.data_y = sld_data.pos_y
151        self.data_z = sld_data.pos_z
152        self.data_sldn = sld_data.sld_n
153        self.data_mx = sld_data.sld_mx
154        self.data_my = sld_data.sld_my
155        self.data_mz = sld_data.sld_mz
156        self.data_vol = sld_data.vol_pix
157        self.data_total_volume = sum(sld_data.vol_pix)
158        self.params['total_volume'] = sum(sld_data.vol_pix)
159
160    def getProfile(self):
161        """
162        Get SLD profile
163        : return: sld_data
164        """
165        return self.sld_data
166
167    def run(self, x=0.0):
168        """
169        Evaluate the model
170        :param x: simple value
171        :return: (I value)
172        """
173        if x.__class__.__name__ == 'list':
174            if len(x[1]) > 0:
175                msg = "Not a 1D."
176                raise ValueError(msg)
177            i_out = np.zeros_like(x[0])
178            # 1D I is found at y =0 in the 2D pattern
179            out = self._gen(x[0], [], i_out)
180            return out
181        else:
182            msg = "Q must be given as list of qx's and qy's"
183            raise ValueError(msg)
184
185    def runXY(self, x=0.0):
186        """
187        Evaluate the model
188        :param x: simple value
189        :return: I value
190        :Use this runXY() for the computation
191        """
192        if x.__class__.__name__ == 'list':
193            i_out = np.zeros_like(x[0])
194            out = self._gen(x[0], x[1], i_out)
195            return out
196        else:
197            msg = "Q must be given as list of qx's and qy's"
198            raise ValueError(msg)
199
200    def evalDistribution(self, qdist):
201        """
202        Evaluate a distribution of q-values.
203
204        :param qdist: ndarray of scalar q-values (for 1D) or list [qx,qy]
205                      where qx,qy are 1D ndarrays (for 2D).
206        """
207        if qdist.__class__.__name__ == 'list':
208            if len(qdist[1]) < 1:
209                out = self.run(qdist)
210            else:
211                out = self.runXY(qdist)
212            return out
213        else:
214            mesg = "evalDistribution is expecting an ndarray of "
215            mesg += "a list [qx,qy] where qx,qy are arrays."
216            raise RuntimeError(mesg)
217
218class OMF2SLD(object):
219    """
220    Convert OMFData to MAgData
221    """
222    def __init__(self):
223        """
224        Init
225        """
226        self.pos_x = None
227        self.pos_y = None
228        self.pos_z = None
229        self.mx = None
230        self.my = None
231        self.mz = None
232        self.sld_n = None
233        self.vol_pix = None
234        self.output = None
235        self.omfdata = None
236
237    def set_data(self, omfdata, shape='rectangular'):
238        """
239        Set all data
240        """
241        self.omfdata = omfdata
242        length = int(omfdata.xnodes * omfdata.ynodes * omfdata.znodes)
243        pos_x = np.arange(omfdata.xmin,
244                             omfdata.xnodes*omfdata.xstepsize + omfdata.xmin,
245                             omfdata.xstepsize)
246        pos_y = np.arange(omfdata.ymin,
247                             omfdata.ynodes*omfdata.ystepsize + omfdata.ymin,
248                             omfdata.ystepsize)
249        pos_z = np.arange(omfdata.zmin,
250                             omfdata.znodes*omfdata.zstepsize + omfdata.zmin,
251                             omfdata.zstepsize)
252        self.pos_x = np.tile(pos_x, int(omfdata.ynodes * omfdata.znodes))
253        self.pos_y = pos_y.repeat(int(omfdata.xnodes))
254        self.pos_y = np.tile(self.pos_y, int(omfdata.znodes))
255        self.pos_z = pos_z.repeat(int(omfdata.xnodes * omfdata.ynodes))
256        self.mx = omfdata.mx
257        self.my = omfdata.my
258        self.mz = omfdata.mz
259        self.sld_n = np.zeros(length)
260
261        if omfdata.mx is None:
262            self.mx = np.zeros(length)
263        if omfdata.my is None:
264            self.my = np.zeros(length)
265        if omfdata.mz is None:
266            self.mz = np.zeros(length)
267
268        self._check_data_length(length)
269        self.remove_null_points(False, False)
271        if shape.lower() == 'ellipsoid':
272            try:
273                # Pixel (step) size included
274                x_c = max(self.pos_x) + min(self.pos_x)
275                y_c = max(self.pos_y) + min(self.pos_y)
276                z_c = max(self.pos_z) + min(self.pos_z)
277                x_d = max(self.pos_x) - min(self.pos_x)
278                y_d = max(self.pos_y) - min(self.pos_y)
279                z_d = max(self.pos_z) - min(self.pos_z)
280                x_r = (x_d + omfdata.xstepsize) / 2.0
281                y_r = (y_d + omfdata.ystepsize) / 2.0
282                z_r = (z_d + omfdata.zstepsize) / 2.0
283                x_dir2 = ((self.pos_x - x_c / 2.0) / x_r)
284                x_dir2 *= x_dir2
285                y_dir2 = ((self.pos_y - y_c / 2.0) / y_r)
286                y_dir2 *= y_dir2
287                z_dir2 = ((self.pos_z - z_c / 2.0) / z_r)
288                z_dir2 *= z_dir2
289                mask = (x_dir2 + y_dir2 + z_dir2) <= 1.0
290            except:
291                logger.error(sys.exc_value)
295        self.output.set_pix_type('pixel')
296        self.output.set_pixel_symbols('pixel')
297
298    def get_omfdata(self):
299        """
300        Return all data
301        """
302        return self.omfdata
303
304    def get_output(self):
305        """
306        Return output
307        """
308        return self.output
309
310    def _check_data_length(self, length):
311        """
312        Check if the data lengths are consistent
313        :Params length: data length
314        """
315        parts = (self.pos_x, self.pos_y, self.pos_z, self.mx, self.my, self.mz)
316        if any(len(v) != length for v in parts):
317            raise ValueError("Error: Inconsistent data length.")
318
319    def remove_null_points(self, remove=False, recenter=False):
320        """
321        Removes any mx, my, and mz = 0 points
322        """
323        if remove:
324            is_nonzero = (np.fabs(self.mx) + np.fabs(self.my) +
325                          np.fabs(self.mz)).nonzero()
326            if len(is_nonzero[0]) > 0:
327                self.pos_x = self.pos_x[is_nonzero]
328                self.pos_y = self.pos_y[is_nonzero]
329                self.pos_z = self.pos_z[is_nonzero]
330                self.sld_n = self.sld_n[is_nonzero]
331                self.mx = self.mx[is_nonzero]
332                self.my = self.my[is_nonzero]
333                self.mz = self.mz[is_nonzero]
334        if recenter:
335            self.pos_x -= (min(self.pos_x) + max(self.pos_x)) / 2.0
336            self.pos_y -= (min(self.pos_y) + max(self.pos_y)) / 2.0
337            self.pos_z -= (min(self.pos_z) + max(self.pos_z)) / 2.0
338
339    def get_magsld(self):
340        """
341        return MagSLD
342        """
343        return self.output
344
345
347    """
349    """
350    ## File type
351    type_name = "OMF ASCII"
352
353    ## Wildcards
354    type = ["OMF files (*.OMF, *.omf)|*.omf"]
355    ## List of allowed extensions
356    ext = ['.omf', '.OMF']
357
359        """
361        :param path: file path
362        :return: x, y, z, sld_n, sld_mx, sld_my, sld_mz
363        """
364        desc = ""
365        mx = np.zeros(0)
366        my = np.zeros(0)
367        mz = np.zeros(0)
368        try:
369            input_f = open(path, 'rb')
371            lines = buff.split('\n')
372            input_f.close()
373            output = OMFData()
374            valueunit = None
375            for line in lines:
376                toks = line.split()
378                try:
379                    _mx = float(toks[0])
380                    _my = float(toks[1])
381                    _mz = float(toks[2])
382                    _mx = mag2sld(_mx, valueunit)
383                    _my = mag2sld(_my, valueunit)
384                    _mz = mag2sld(_mz, valueunit)
385                    mx = np.append(mx, _mx)
386                    my = np.append(my, _my)
387                    mz = np.append(mz, _mz)
388                except:
389                    # Skip non-data lines
390                    logger.error(sys.exc_value)
392                s_line = line.split(":", 1)
393                if s_line[0].lower().count("oommf") > 0:
394                    oommf = s_line[1].lstrip()
395                if s_line[0].lower().count("title") > 0:
396                    title = s_line[1].lstrip()
397                if s_line[0].lower().count("desc") > 0:
398                    desc += s_line[1].lstrip()
399                    desc += '\n'
400                if s_line[0].lower().count("meshtype") > 0:
401                    meshtype = s_line[1].lstrip()
402                if s_line[0].lower().count("meshunit") > 0:
403                    meshunit = s_line[1].lstrip()
404                    if meshunit.count("m") < 1:
405                        msg = "Error: \n"
406                        msg += "We accept only m as meshunit"
407                        raise ValueError(msg)
408                if s_line[0].lower().count("xbase") > 0:
409                    xbase = s_line[1].lstrip()
410                if s_line[0].lower().count("ybase") > 0:
411                    ybase = s_line[1].lstrip()
412                if s_line[0].lower().count("zbase") > 0:
413                    zbase = s_line[1].lstrip()
414                if s_line[0].lower().count("xstepsize") > 0:
415                    xstepsize = s_line[1].lstrip()
416                if s_line[0].lower().count("ystepsize") > 0:
417                    ystepsize = s_line[1].lstrip()
418                if s_line[0].lower().count("zstepsize") > 0:
419                    zstepsize = s_line[1].lstrip()
420                if s_line[0].lower().count("xnodes") > 0:
421                    xnodes = s_line[1].lstrip()
422                if s_line[0].lower().count("ynodes") > 0:
423                    ynodes = s_line[1].lstrip()
424                if s_line[0].lower().count("znodes") > 0:
425                    znodes = s_line[1].lstrip()
426                if s_line[0].lower().count("xmin") > 0:
427                    xmin = s_line[1].lstrip()
428                if s_line[0].lower().count("ymin") > 0:
429                    ymin = s_line[1].lstrip()
430                if s_line[0].lower().count("zmin") > 0:
431                    zmin = s_line[1].lstrip()
432                if s_line[0].lower().count("xmax") > 0:
433                    xmax = s_line[1].lstrip()
434                if s_line[0].lower().count("ymax") > 0:
435                    ymax = s_line[1].lstrip()
436                if s_line[0].lower().count("zmax") > 0:
437                    zmax = s_line[1].lstrip()
438                if s_line[0].lower().count("valueunit") > 0:
439                    valueunit = s_line[1].lstrip().rstrip()
440                if s_line[0].lower().count("valuemultiplier") > 0:
441                    valuemultiplier = s_line[1].lstrip()
442                if s_line[0].lower().count("valuerangeminmag") > 0:
443                    valuerangeminmag = s_line[1].lstrip()
444                if s_line[0].lower().count("valuerangemaxmag") > 0:
445                    valuerangemaxmag = s_line[1].lstrip()
446                if s_line[0].lower().count("end") > 0:
447                    output.filename = os.path.basename(path)
448                    output.oommf = oommf
449                    output.title = title
450                    output.desc = desc
451                    output.meshtype = meshtype
452                    output.xbase = float(xbase) * METER2ANG
453                    output.ybase = float(ybase) * METER2ANG
454                    output.zbase = float(zbase) * METER2ANG
455                    output.xstepsize = float(xstepsize) * METER2ANG
456                    output.ystepsize = float(ystepsize) * METER2ANG
457                    output.zstepsize = float(zstepsize) * METER2ANG
458                    output.xnodes = float(xnodes)
459                    output.ynodes = float(ynodes)
460                    output.znodes = float(znodes)
461                    output.xmin = float(xmin) * METER2ANG
462                    output.ymin = float(ymin) * METER2ANG
463                    output.zmin = float(zmin) * METER2ANG
464                    output.xmax = float(xmax) * METER2ANG
465                    output.ymax = float(ymax) * METER2ANG
466                    output.zmax = float(zmax) * METER2ANG
467                    output.valuemultiplier = valuemultiplier
468                    output.valuerangeminmag = mag2sld(float(valuerangeminmag), \
469                                                      valueunit)
470                    output.valuerangemaxmag = mag2sld(float(valuerangemaxmag), \
471                                                      valueunit)
472            output.set_m(mx, my, mz)
473            return output
474        except:
475            msg = "%s is not supported: \n" % path
476            msg += "We accept only Text format OMF file."
477            raise RuntimeError(msg)
478
480    """
481    PDB reader class: limited for reading the lines starting with 'ATOM'
482    """
483    type_name = "PDB"
484    ## Wildcards
485    type = ["pdb files (*.PDB, *.pdb)|*.pdb"]
486    ## List of allowed extensions
487    ext = ['.pdb', '.PDB']
488
490        """
492
493        :param path: file path
494        :return: MagSLD
495        :raise RuntimeError: when the file can't be opened
496        """
497        pos_x = np.zeros(0)
498        pos_y = np.zeros(0)
499        pos_z = np.zeros(0)
500        sld_n = np.zeros(0)
501        sld_mx = np.zeros(0)
502        sld_my = np.zeros(0)
503        sld_mz = np.zeros(0)
504        vol_pix = np.zeros(0)
505        pix_symbol = np.zeros(0)
506        x_line = []
507        y_line = []
508        z_line = []
509        x_lines = []
510        y_lines = []
511        z_lines = []
512        try:
513            input_f = open(path, 'rb')
515            lines = buff.split('\n')
516            input_f.close()
517            num = 0
518            for line in lines:
519                try:
520                    # check if line starts with "ATOM"
521                    if line[0:6].strip().count('ATM') > 0 or \
522                                line[0:6].strip() == 'ATOM':
523                        # define fields of interest
524                        atom_name = line[12:16].strip()
525                        try:
526                            float(line[12])
527                            atom_name = atom_name[1].upper()
528                        except:
529                            if len(atom_name) == 4:
530                                atom_name = atom_name[0].upper()
531                            elif line[12] != ' ':
532                                atom_name = atom_name[0].upper() + \
533                                        atom_name[1].lower()
534                            else:
535                                atom_name = atom_name[0].upper()
536                        _pos_x = float(line[30:38].strip())
537                        _pos_y = float(line[38:46].strip())
538                        _pos_z = float(line[46:54].strip())
539                        pos_x = np.append(pos_x, _pos_x)
540                        pos_y = np.append(pos_y, _pos_y)
541                        pos_z = np.append(pos_z, _pos_z)
542                        try:
543                            val = nsf.neutron_sld(atom_name)[0]
544                            # sld in Ang^-2 unit
545                            val *= 1.0e-6
546                            sld_n = np.append(sld_n, val)
547                            atom = formula(atom_name)
548                            # cm to A units
549                            vol = 1.0e+24 * atom.mass / atom.density / NA
550                            vol_pix = np.append(vol_pix, vol)
551                        except:
552                            print("Error: set the sld of %s to zero"% atom_name)
553                            sld_n = np.append(sld_n, 0.0)
554                        sld_mx = np.append(sld_mx, 0)
555                        sld_my = np.append(sld_my, 0)
556                        sld_mz = np.append(sld_mz, 0)
557                        pix_symbol = np.append(pix_symbol, atom_name)
558                    elif line[0:6].strip().count('CONECT') > 0:
559                        toks = line.split()
560                        num = int(toks[1]) - 1
561                        val_list = []
562                        for val in toks[2:]:
563                            try:
564                                int_val = int(val)
565                            except:
566                                break
567                            if int_val == 0:
568                                break
569                            val_list.append(int_val)
570                        #need val_list ordered
571                        for val in val_list:
572                            index = val - 1
573                            if (pos_x[index], pos_x[num]) in x_line and \
574                               (pos_y[index], pos_y[num]) in y_line and \
575                               (pos_z[index], pos_z[num]) in z_line:
576                                continue
577                            x_line.append((pos_x[num], pos_x[index]))
578                            y_line.append((pos_y[num], pos_y[index]))
579                            z_line.append((pos_z[num], pos_z[index]))
580                    if len(x_line) > 0:
581                        x_lines.append(x_line)
582                        y_lines.append(y_line)
583                        z_lines.append(z_line)
584                except:
585                    logger.error(sys.exc_value)
586
587            output = MagSLD(pos_x, pos_y, pos_z, sld_n, sld_mx, sld_my, sld_mz)
588            output.set_conect_lines(x_line, y_line, z_line)
589            output.filename = os.path.basename(path)
590            output.set_pix_type('atom')
591            output.set_pixel_symbols(pix_symbol)
592            output.set_nodes()
593            output.set_pixel_volumes(vol_pix)
594            output.sld_unit = '1/A^(2)'
595            return output
596        except:
597            raise RuntimeError("%s is not a sld file" % path)
598
599    def write(self, path, data):
600        """
601        Write
602        """
603        print("Not implemented... ")
604
606    """
607    Class to load ascii files (7 columns).
608    """
609    ## File type
610    type_name = "SLD ASCII"
611    ## Wildcards
612    type = ["sld files (*.SLD, *.sld)|*.sld",
613            "txt files (*.TXT, *.txt)|*.txt",
614            "all files (*.*)|*.*"]
615    ## List of allowed extensions
616    ext = ['.sld', '.SLD', '.txt', '.TXT', '.*']
618        """
620        :param path: file path
621        :return MagSLD: x, y, z, sld_n, sld_mx, sld_my, sld_mz
622        :raise RuntimeError: when the file can't be opened
623        :raise ValueError: when the length of the data vectors are inconsistent
624        """
625        try:
626            pos_x = np.zeros(0)
627            pos_y = np.zeros(0)
628            pos_z = np.zeros(0)
629            sld_n = np.zeros(0)
630            sld_mx = np.zeros(0)
631            sld_my = np.zeros(0)
632            sld_mz = np.zeros(0)
633            try:
635                input_f = np.loadtxt(path, dtype='float', skiprows=1,
636                                        ndmin=1, unpack=True)
637                pos_x = np.array(input_f[0])
638                pos_y = np.array(input_f[1])
639                pos_z = np.array(input_f[2])
640                sld_n = np.array(input_f[3])
641                sld_mx = np.array(input_f[4])
642                sld_my = np.array(input_f[5])
643                sld_mz = np.array(input_f[6])
644                ncols = len(input_f)
645                if ncols == 8:
646                    vol_pix = np.array(input_f[7])
647                elif ncols == 7:
648                    vol_pix = None
649            except:
650                # For older version of numpy
651                input_f = open(path, 'rb')
653                lines = buff.split('\n')
654                input_f.close()
655                for line in lines:
656                    toks = line.split()
657                    try:
658                        _pos_x = float(toks[0])
659                        _pos_y = float(toks[1])
660                        _pos_z = float(toks[2])
661                        _sld_n = float(toks[3])
662                        _sld_mx = float(toks[4])
663                        _sld_my = float(toks[5])
664                        _sld_mz = float(toks[6])
665                        pos_x = np.append(pos_x, _pos_x)
666                        pos_y = np.append(pos_y, _pos_y)
667                        pos_z = np.append(pos_z, _pos_z)
668                        sld_n = np.append(sld_n, _sld_n)
669                        sld_mx = np.append(sld_mx, _sld_mx)
670                        sld_my = np.append(sld_my, _sld_my)
671                        sld_mz = np.append(sld_mz, _sld_mz)
672                        try:
673                            _vol_pix = float(toks[7])
674                            vol_pix = np.append(vol_pix, _vol_pix)
675                        except:
676                            vol_pix = None
677                    except:
678                        # Skip non-data lines
679                        logger.error(sys.exc_value)
680            output = MagSLD(pos_x, pos_y, pos_z, sld_n,
681                            sld_mx, sld_my, sld_mz)
682            output.filename = os.path.basename(path)
683            output.set_pix_type('pixel')
684            output.set_pixel_symbols('pixel')
685            if vol_pix is not None:
686                output.set_pixel_volumes(vol_pix)
687            return output
688        except:
689            raise RuntimeError("%s is not a sld file" % path)
690
691    def write(self, path, data):
692        """
693        Write sld file
694        :Param path: file path
695        :Param data: MagSLD data object
696        """
697        if path is None:
698            raise ValueError("Missing the file path.")
699        if data is None:
700            raise ValueError("Missing the data to save.")
701        x_val = data.pos_x
702        y_val = data.pos_y
703        z_val = data.pos_z
704        vol_pix = data.vol_pix
705        length = len(x_val)
706        sld_n = data.sld_n
707        if sld_n is None:
708            sld_n = np.zeros(length)
709        sld_mx = data.sld_mx
710        if sld_mx is None:
711            sld_mx = np.zeros(length)
712            sld_my = np.zeros(length)
713            sld_mz = np.zeros(length)
714        else:
715            sld_my = data.sld_my
716            sld_mz = data.sld_mz
717        out = open(path, 'w')
718        # First Line: Column names
719        out.write("X  Y  Z  SLDN SLDMx  SLDMy  SLDMz VOLUMEpix")
720        for ind in range(length):
721            out.write("\n%g  %g  %g  %g  %g  %g  %g %g" % \
722                      (x_val[ind], y_val[ind], z_val[ind], sld_n[ind],
723                       sld_mx[ind], sld_my[ind], sld_mz[ind], vol_pix[ind]))
724        out.close()
725
726
727class OMFData(object):
728    """
729    OMF Data.
730    """
731    _meshunit = "A"
732    _valueunit = "A^(-2)"
733    def __init__(self):
734        """
735        Init for mag SLD
736        """
737        self.filename = 'default'
738        self.oommf = ''
739        self.title = ''
740        self.desc = ''
741        self.meshtype = ''
742        self.meshunit = self._meshunit
743        self.valueunit = self._valueunit
744        self.xbase = 0.0
745        self.ybase = 0.0
746        self.zbase = 0.0
747        self.xstepsize = 6.0
748        self.ystepsize = 6.0
749        self.zstepsize = 6.0
750        self.xnodes = 10.0
751        self.ynodes = 10.0
752        self.znodes = 10.0
753        self.xmin = 0.0
754        self.ymin = 0.0
755        self.zmin = 0.0
756        self.xmax = 60.0
757        self.ymax = 60.0
758        self.zmax = 60.0
759        self.mx = None
760        self.my = None
761        self.mz = None
762        self.valuemultiplier = 1.
763        self.valuerangeminmag = 0
764        self.valuerangemaxmag = 0
765
766    def __str__(self):
767        """
768        doc strings
769        """
770        _str = "Type:            %s\n" % self.__class__.__name__
771        _str += "File:            %s\n" % self.filename
772        _str += "OOMMF:           %s\n" % self.oommf
773        _str += "Title:           %s\n" % self.title
774        _str += "Desc:            %s\n" % self.desc
775        _str += "meshtype:        %s\n" % self.meshtype
776        _str += "meshunit:        %s\n" % str(self.meshunit)
777        _str += "xbase:           %s [%s]\n" % (str(self.xbase), self.meshunit)
778        _str += "ybase:           %s [%s]\n" % (str(self.ybase), self.meshunit)
779        _str += "zbase:           %s [%s]\n" % (str(self.zbase), self.meshunit)
780        _str += "xstepsize:       %s [%s]\n" % (str(self.xstepsize),
781                                                self.meshunit)
782        _str += "ystepsize:       %s [%s]\n" % (str(self.ystepsize),
783                                                self.meshunit)
784        _str += "zstepsize:       %s [%s]\n" % (str(self.zstepsize),
785                                                self.meshunit)
786        _str += "xnodes:          %s\n" % str(self.xnodes)
787        _str += "ynodes:          %s\n" % str(self.ynodes)
788        _str += "znodes:          %s\n" % str(self.znodes)
789        _str += "xmin:            %s [%s]\n" % (str(self.xmin), self.meshunit)
790        _str += "ymin:            %s [%s]\n" % (str(self.ymin), self.meshunit)
791        _str += "zmin:            %s [%s]\n" % (str(self.zmin), self.meshunit)
792        _str += "xmax:            %s [%s]\n" % (str(self.xmax), self.meshunit)
793        _str += "ymax:            %s [%s]\n" % (str(self.ymax), self.meshunit)
794        _str += "zmax:            %s [%s]\n" % (str(self.zmax), self.meshunit)
795        _str += "valueunit:       %s\n" % self.valueunit
796        _str += "valuemultiplier: %s\n" % str(self.valuemultiplier)
797        _str += "ValueRangeMinMag:%s [%s]\n" % (str(self.valuerangeminmag),
798                                                self.valueunit)
799        _str += "ValueRangeMaxMag:%s [%s]\n" % (str(self.valuerangemaxmag),
800                                                self.valueunit)
801        return _str
802
803    def set_m(self, mx, my, mz):
804        """
805        Set the Mx, My, Mz values
806        """
807        self.mx = mx
808        self.my = my
809        self.mz = mz
810
811class MagSLD(object):
812    """
813    Magnetic SLD.
814    """
815    pos_x = None
816    pos_y = None
817    pos_z = None
818    sld_n = None
819    sld_mx = None
820    sld_my = None
821    sld_mz = None
822    # Units
823    _pos_unit = 'A'
824    _sld_unit = '1/A^(2)'
825    _pix_type = 'pixel'
826
827    def __init__(self, pos_x, pos_y, pos_z, sld_n=None,
828                 sld_mx=None, sld_my=None, sld_mz=None, vol_pix=None):
829        """
830        Init for mag SLD
831        :params : All should be numpy 1D array
832        """
833        self.is_data = True
834        self.filename = ''
835        self.xstepsize = 6.0
836        self.ystepsize = 6.0
837        self.zstepsize = 6.0
838        self.xnodes = 10.0
839        self.ynodes = 10.0
840        self.znodes = 10.0
841        self.has_stepsize = False
842        self.has_conect = False
843        self.pos_unit = self._pos_unit
844        self.sld_unit = self._sld_unit
845        self.pix_type = 'pixel'
846        self.pos_x = pos_x
847        self.pos_y = pos_y
848        self.pos_z = pos_z
849        self.sld_n = sld_n
850        self.line_x = None
851        self.line_y = None
852        self.line_z = None
853        self.sld_mx = sld_mx
854        self.sld_my = sld_my
855        self.sld_mz = sld_mz
856        self.vol_pix = vol_pix
857        self.sld_m = None
858        self.sld_phi = None
859        self.sld_theta = None
860        self.pix_symbol = None
861        if sld_mx is not None and sld_my is not None and sld_mz is not None:
862            self.set_sldms(sld_mx, sld_my, sld_mz)
863        self.set_nodes()
864
865    def __str__(self):
866        """
867        doc strings
868        """
869        _str = "Type:       %s\n" % self.__class__.__name__
870        _str += "File:       %s\n" % self.filename
871        _str += "Axis_unit:  %s\n" % self.pos_unit
872        _str += "SLD_unit:   %s\n" % self.sld_unit
873        return _str
874
875    def set_pix_type(self, pix_type):
876        """
877        Set pixel type
878        :Param pix_type: string, 'pixel' or 'atom'
879        """
880        self.pix_type = pix_type
881
882    def set_sldn(self, sld_n):
883        """
884        Sets neutron SLD
885        """
886        if sld_n.__class__.__name__ == 'float':
887            if self.is_data:
888                # For data, put the value to only the pixels w non-zero M
889                is_nonzero = (np.fabs(self.sld_mx) +
890                              np.fabs(self.sld_my) +
891                              np.fabs(self.sld_mz)).nonzero()
892                self.sld_n = np.zeros(len(self.pos_x))
893                if len(self.sld_n[is_nonzero]) > 0:
894                    self.sld_n[is_nonzero] = sld_n
895                else:
896                    self.sld_n.fill(sld_n)
897            else:
898                # For non-data, put the value to all the pixels
899                self.sld_n = np.ones(len(self.pos_x)) * sld_n
900        else:
901            self.sld_n = sld_n
902
903    def set_sldms(self, sld_mx, sld_my, sld_mz):
904        r"""
905        Sets mx, my, mz and abs(m).
906        """ # Note: escaping
907        if sld_mx.__class__.__name__ == 'float':
908            self.sld_mx = np.ones(len(self.pos_x)) * sld_mx
909        else:
910            self.sld_mx = sld_mx
911        if sld_my.__class__.__name__ == 'float':
912            self.sld_my = np.ones(len(self.pos_x)) * sld_my
913        else:
914            self.sld_my = sld_my
915        if sld_mz.__class__.__name__ == 'float':
916            self.sld_mz = np.ones(len(self.pos_x)) * sld_mz
917        else:
918            self.sld_mz = sld_mz
919
920        sld_m = np.sqrt(sld_mx * sld_mx + sld_my * sld_my + \
921                                sld_mz * sld_mz)
922        self.sld_m = sld_m
923
924    def set_pixel_symbols(self, symbol='pixel'):
925        """
926        Set pixel
927        :Params pixel: str; pixel or atomic symbol, or array of strings
928        """
929        if self.sld_n is None:
930            return
931        if symbol.__class__.__name__ == 'str':
932            self.pix_symbol = np.repeat(symbol, len(self.sld_n))
933        else:
934            self.pix_symbol = symbol
935
936    def set_pixel_volumes(self, vol):
937        """
938        Set pixel volumes
939        :Params pixel: str; pixel or atomic symbol, or array of strings
940        """
941        if self.sld_n is None:
942            return
943        if vol.__class__.__name__ == 'ndarray':
944            self.vol_pix = vol
945        elif vol.__class__.__name__.count('float') > 0:
946            self.vol_pix = np.repeat(vol, len(self.sld_n))
947        else:
948            self.vol_pix = None
949
950    def get_sldn(self):
951        """
952        Returns nuclear sld
953        """
954        return self.sld_n
955
956    def set_nodes(self):
957        """
958        Set xnodes, ynodes, and znodes
959        """
960        self.set_stepsize()
961        if self.pix_type == 'pixel':
962            try:
963                xdist = (max(self.pos_x) - min(self.pos_x)) / self.xstepsize
964                ydist = (max(self.pos_y) - min(self.pos_y)) / self.ystepsize
965                zdist = (max(self.pos_z) - min(self.pos_z)) / self.zstepsize
966                self.xnodes = int(xdist) + 1
967                self.ynodes = int(ydist) + 1
968                self.znodes = int(zdist) + 1
969            except:
970                self.xnodes = None
971                self.ynodes = None
972                self.znodes = None
973        else:
974            self.xnodes = None
975            self.ynodes = None
976            self.znodes = None
977
978    def set_stepsize(self):
979        """
980        Set xtepsize, ystepsize, and zstepsize
981        """
982        if self.pix_type == 'pixel':
983            try:
984                xpos_pre = self.pos_x[0]
985                ypos_pre = self.pos_y[0]
986                zpos_pre = self.pos_z[0]
987                for x_pos in self.pos_x:
988                    if xpos_pre != x_pos:
989                        self.xstepsize = np.fabs(x_pos - xpos_pre)
990                        break
991                for y_pos in self.pos_y:
992                    if ypos_pre != y_pos:
993                        self.ystepsize = np.fabs(y_pos - ypos_pre)
994                        break
995                for z_pos in self.pos_z:
996                    if zpos_pre != z_pos:
997                        self.zstepsize = np.fabs(z_pos - zpos_pre)
998                        break
999                #default pix volume
1000                self.vol_pix = np.ones(len(self.pos_x))
1001                vol = self.xstepsize * self.ystepsize * self.zstepsize
1002                self.set_pixel_volumes(vol)
1003                self.has_stepsize = True
1004            except:
1005                self.xstepsize = None
1006                self.ystepsize = None
1007                self.zstepsize = None
1008                self.vol_pix = None
1009                self.has_stepsize = False
1010        else:
1011            self.xstepsize = None
1012            self.ystepsize = None
1013            self.zstepsize = None
1014            self.has_stepsize = True
1015        return self.xstepsize, self.ystepsize, self.zstepsize
1016
1017    def set_conect_lines(self, line_x, line_y, line_z):
1018        """
1019        Set bonding line data if taken from pdb
1020        """
1021        if line_x.__class__.__name__ != 'list' or len(line_x) < 1:
1022            return
1023        if line_y.__class__.__name__ != 'list' or len(line_y) < 1:
1024            return
1025        if line_z.__class__.__name__ != 'list' or len(line_z) < 1:
1026            return
1027        self.has_conect = True
1028        self.line_x = line_x
1029        self.line_y = line_y
1030        self.line_z = line_z
1031
1033    """
1034        Test code
1035    """
1036    from mpl_toolkits.mplot3d import Axes3D
1037    current_dir = os.path.abspath(os.path.curdir)
1038    print(current_dir)
1039    for i in range(6):
1040        current_dir, _ = os.path.split(current_dir)
1041        tfile = os.path.join(current_dir, "test", "CoreXY_ShellZ.txt")
1042        ofile = os.path.join(current_dir, "test", "A_Raw_Example-1.omf")
1043        if os.path.isfile(tfile):
1044            tfpath = tfile
1045            ofpath = ofile
1046            break
1051    foutput = OMF2SLD()
1052    foutput.set_data(ooutput)
1053
1054    import matplotlib.pyplot as plt
1055    fig = plt.figure()
1056    ax = Axes3D(fig)
1057    ax.plot(output.pos_x, output.pos_y, output.pos_z, '.', c="g",
1058            alpha=0.7, markeredgecolor='gray', rasterized=True)
1059    gap = 7
1060    max_mx = max(output.sld_mx)
1061    max_my = max(output.sld_my)
1062    max_mz = max(output.sld_mz)
1063    max_m = max(max_mx, max_my, max_mz)
1064    x2 = output.pos_x+output.sld_mx/max_m * gap
1065    y2 = output.pos_y+output.sld_my/max_m * gap
1066    z2 = output.pos_z+output.sld_mz/max_m * gap
1067    x_arrow = np.column_stack((output.pos_x, x2))
1068    y_arrow = np.column_stack((output.pos_y, y2))
1069    z_arrow = np.column_stack((output.pos_z, z2))
1070    unit_x2 = output.sld_mx / max_m
1071    unit_y2 = output.sld_my / max_m
1072    unit_z2 = output.sld_mz / max_m
1073    color_x = np.fabs(unit_x2 * 0.8)
1074    color_y = np.fabs(unit_y2 * 0.8)
1075    color_z = np.fabs(unit_z2 * 0.8)
1076    colors = np.column_stack((color_x, color_y, color_z))
1077    plt.show()
1078
1079def test():
1080    """
1081        Test code
1082    """
1083    current_dir = os.path.abspath(os.path.curdir)
1084    for i in range(3):
1085        current_dir, _ = os.path.split(current_dir)
1086        ofile = os.path.join(current_dir, "test", "A_Raw_Example-1.omf")
1087        if os.path.isfile(ofile):
1088            ofpath = ofile
1089            break
1092    foutput = OMF2SLD()
1093    foutput.set_data(ooutput)