source: sasview/src/sas/sascalc/calculator/sas_gen.py @ 9e0dd49

magnetic_scattrelease-4.2.2ticket-1009ticket-1094-headlessticket-1242-2d-resolutionticket-1243ticket-1249unittest-saveload
Last change on this file since 9e0dd49 was a1b8fee, checked in by andyfaff, 8 years ago

MAINT: from future import print_function

  • Property mode set to 100644
File size: 39.4 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
22#Avogadro constant [1/mol]
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
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)
270        mask = np.ones(len(self.sld_n), dtype=bool)
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)
292        self.output = MagSLD(self.pos_x[mask], self.pos_y[mask],
293                             self.pos_z[mask], self.sld_n[mask],
294                             self.mx[mask], self.my[mask], self.mz[mask])
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        msg = "Error: Inconsistent data length."
316        if len(self.pos_x) != length:
317            raise ValueError, msg
318        if len(self.pos_y) != length:
319            raise ValueError, msg
320        if len(self.pos_z) != length:
321            raise ValueError, msg
322        if len(self.mx) != length:
323            raise ValueError, msg
324        if len(self.my) != length:
325            raise ValueError, msg
326        if len(self.mz) != length:
327            raise ValueError, msg
328
329    def remove_null_points(self, remove=False, recenter=False):
330        """
331        Removes any mx, my, and mz = 0 points
332        """
333        if remove:
334            is_nonzero = (np.fabs(self.mx) + np.fabs(self.my) +
335                          np.fabs(self.mz)).nonzero()
336            if len(is_nonzero[0]) > 0:
337                self.pos_x = self.pos_x[is_nonzero]
338                self.pos_y = self.pos_y[is_nonzero]
339                self.pos_z = self.pos_z[is_nonzero]
340                self.sld_n = self.sld_n[is_nonzero]
341                self.mx = self.mx[is_nonzero]
342                self.my = self.my[is_nonzero]
343                self.mz = self.mz[is_nonzero]
344        if recenter:
345            self.pos_x -= (min(self.pos_x) + max(self.pos_x)) / 2.0
346            self.pos_y -= (min(self.pos_y) + max(self.pos_y)) / 2.0
347            self.pos_z -= (min(self.pos_z) + max(self.pos_z)) / 2.0
348
349    def get_magsld(self):
350        """
351        return MagSLD
352        """
353        return self.output
354
355
356class OMFReader(object):
357    """
358    Class to load omf/ascii files (3 columns w/header).
359    """
360    ## File type
361    type_name = "OMF ASCII"
362
363    ## Wildcards
364    type = ["OMF files (*.OMF, *.omf)|*.omf"]
365    ## List of allowed extensions
366    ext = ['.omf', '.OMF']
367
368    def read(self, path):
369        """
370        Load data file
371        :param path: file path
372        :return: x, y, z, sld_n, sld_mx, sld_my, sld_mz
373        """
374        desc = ""
375        mx = np.zeros(0)
376        my = np.zeros(0)
377        mz = np.zeros(0)
378        try:
379            input_f = open(path, 'rb')
380            buff = input_f.read()
381            lines = buff.split('\n')
382            input_f.close()
383            output = OMFData()
384            valueunit = None
385            for line in lines:
386                toks = line.split()
387                # Read data
388                try:
389                    _mx = float(toks[0])
390                    _my = float(toks[1])
391                    _mz = float(toks[2])
392                    _mx = mag2sld(_mx, valueunit)
393                    _my = mag2sld(_my, valueunit)
394                    _mz = mag2sld(_mz, valueunit)
395                    mx = np.append(mx, _mx)
396                    my = np.append(my, _my)
397                    mz = np.append(mz, _mz)
398                except:
399                    # Skip non-data lines
400                    logger.error(sys.exc_value)
401                #Reading Header; Segment count ignored
402                s_line = line.split(":", 1)
403                if s_line[0].lower().count("oommf") > 0:
404                    oommf = s_line[1].lstrip()
405                if s_line[0].lower().count("title") > 0:
406                    title = s_line[1].lstrip()
407                if s_line[0].lower().count("desc") > 0:
408                    desc += s_line[1].lstrip()
409                    desc += '\n'
410                if s_line[0].lower().count("meshtype") > 0:
411                    meshtype = s_line[1].lstrip()
412                if s_line[0].lower().count("meshunit") > 0:
413                    meshunit = s_line[1].lstrip()
414                    if meshunit.count("m") < 1:
415                        msg = "Error: \n"
416                        msg += "We accept only m as meshunit"
417                        raise ValueError, msg
418                if s_line[0].lower().count("xbase") > 0:
419                    xbase = s_line[1].lstrip()
420                if s_line[0].lower().count("ybase") > 0:
421                    ybase = s_line[1].lstrip()
422                if s_line[0].lower().count("zbase") > 0:
423                    zbase = s_line[1].lstrip()
424                if s_line[0].lower().count("xstepsize") > 0:
425                    xstepsize = s_line[1].lstrip()
426                if s_line[0].lower().count("ystepsize") > 0:
427                    ystepsize = s_line[1].lstrip()
428                if s_line[0].lower().count("zstepsize") > 0:
429                    zstepsize = s_line[1].lstrip()
430                if s_line[0].lower().count("xnodes") > 0:
431                    xnodes = s_line[1].lstrip()
432                if s_line[0].lower().count("ynodes") > 0:
433                    ynodes = s_line[1].lstrip()
434                if s_line[0].lower().count("znodes") > 0:
435                    znodes = s_line[1].lstrip()
436                if s_line[0].lower().count("xmin") > 0:
437                    xmin = s_line[1].lstrip()
438                if s_line[0].lower().count("ymin") > 0:
439                    ymin = s_line[1].lstrip()
440                if s_line[0].lower().count("zmin") > 0:
441                    zmin = s_line[1].lstrip()
442                if s_line[0].lower().count("xmax") > 0:
443                    xmax = s_line[1].lstrip()
444                if s_line[0].lower().count("ymax") > 0:
445                    ymax = s_line[1].lstrip()
446                if s_line[0].lower().count("zmax") > 0:
447                    zmax = s_line[1].lstrip()
448                if s_line[0].lower().count("valueunit") > 0:
449                    valueunit = s_line[1].lstrip().rstrip()
450                if s_line[0].lower().count("valuemultiplier") > 0:
451                    valuemultiplier = s_line[1].lstrip()
452                if s_line[0].lower().count("valuerangeminmag") > 0:
453                    valuerangeminmag = s_line[1].lstrip()
454                if s_line[0].lower().count("valuerangemaxmag") > 0:
455                    valuerangemaxmag = s_line[1].lstrip()
456                if s_line[0].lower().count("end") > 0:
457                    output.filename = os.path.basename(path)
458                    output.oommf = oommf
459                    output.title = title
460                    output.desc = desc
461                    output.meshtype = meshtype
462                    output.xbase = float(xbase) * METER2ANG
463                    output.ybase = float(ybase) * METER2ANG
464                    output.zbase = float(zbase) * METER2ANG
465                    output.xstepsize = float(xstepsize) * METER2ANG
466                    output.ystepsize = float(ystepsize) * METER2ANG
467                    output.zstepsize = float(zstepsize) * METER2ANG
468                    output.xnodes = float(xnodes)
469                    output.ynodes = float(ynodes)
470                    output.znodes = float(znodes)
471                    output.xmin = float(xmin) * METER2ANG
472                    output.ymin = float(ymin) * METER2ANG
473                    output.zmin = float(zmin) * METER2ANG
474                    output.xmax = float(xmax) * METER2ANG
475                    output.ymax = float(ymax) * METER2ANG
476                    output.zmax = float(zmax) * METER2ANG
477                    output.valuemultiplier = valuemultiplier
478                    output.valuerangeminmag = mag2sld(float(valuerangeminmag), \
479                                                      valueunit)
480                    output.valuerangemaxmag = mag2sld(float(valuerangemaxmag), \
481                                                      valueunit)
482            output.set_m(mx, my, mz)
483            return output
484        except:
485            msg = "%s is not supported: \n" % path
486            msg += "We accept only Text format OMF file."
487            raise RuntimeError, msg
488
489class PDBReader(object):
490    """
491    PDB reader class: limited for reading the lines starting with 'ATOM'
492    """
493    type_name = "PDB"
494    ## Wildcards
495    type = ["pdb files (*.PDB, *.pdb)|*.pdb"]
496    ## List of allowed extensions
497    ext = ['.pdb', '.PDB']
498
499    def read(self, path):
500        """
501        Load data file
502
503        :param path: file path
504        :return: MagSLD
505        :raise RuntimeError: when the file can't be opened
506        """
507        pos_x = np.zeros(0)
508        pos_y = np.zeros(0)
509        pos_z = np.zeros(0)
510        sld_n = np.zeros(0)
511        sld_mx = np.zeros(0)
512        sld_my = np.zeros(0)
513        sld_mz = np.zeros(0)
514        vol_pix = np.zeros(0)
515        pix_symbol = np.zeros(0)
516        x_line = []
517        y_line = []
518        z_line = []
519        x_lines = []
520        y_lines = []
521        z_lines = []
522        try:
523            input_f = open(path, 'rb')
524            buff = input_f.read()
525            lines = buff.split('\n')
526            input_f.close()
527            num = 0
528            for line in lines:
529                try:
530                    # check if line starts with "ATOM"
531                    if line[0:6].strip().count('ATM') > 0 or \
532                                line[0:6].strip() == 'ATOM':
533                        # define fields of interest
534                        atom_name = line[12:16].strip()
535                        try:
536                            float(line[12])
537                            atom_name = atom_name[1].upper()
538                        except:
539                            if len(atom_name) == 4:
540                                atom_name = atom_name[0].upper()
541                            elif line[12] != ' ':
542                                atom_name = atom_name[0].upper() + \
543                                        atom_name[1].lower()
544                            else:
545                                atom_name = atom_name[0].upper()
546                        _pos_x = float(line[30:38].strip())
547                        _pos_y = float(line[38:46].strip())
548                        _pos_z = float(line[46:54].strip())
549                        pos_x = np.append(pos_x, _pos_x)
550                        pos_y = np.append(pos_y, _pos_y)
551                        pos_z = np.append(pos_z, _pos_z)
552                        try:
553                            val = nsf.neutron_sld(atom_name)[0]
554                            # sld in Ang^-2 unit
555                            val *= 1.0e-6
556                            sld_n = np.append(sld_n, val)
557                            atom = formula(atom_name)
558                            # cm to A units
559                            vol = 1.0e+24 * atom.mass / atom.density / NA
560                            vol_pix = np.append(vol_pix, vol)
561                        except:
562                            print("Error: set the sld of %s to zero"% atom_name)
563                            sld_n = np.append(sld_n, 0.0)
564                        sld_mx = np.append(sld_mx, 0)
565                        sld_my = np.append(sld_my, 0)
566                        sld_mz = np.append(sld_mz, 0)
567                        pix_symbol = np.append(pix_symbol, atom_name)
568                    elif line[0:6].strip().count('CONECT') > 0:
569                        toks = line.split()
570                        num = int(toks[1]) - 1
571                        val_list = []
572                        for val in toks[2:]:
573                            try:
574                                int_val = int(val)
575                            except:
576                                break
577                            if int_val == 0:
578                                break
579                            val_list.append(int_val)
580                        #need val_list ordered
581                        for val in val_list:
582                            index = val - 1
583                            if (pos_x[index], pos_x[num]) in x_line and \
584                               (pos_y[index], pos_y[num]) in y_line and \
585                               (pos_z[index], pos_z[num]) in z_line:
586                                continue
587                            x_line.append((pos_x[num], pos_x[index]))
588                            y_line.append((pos_y[num], pos_y[index]))
589                            z_line.append((pos_z[num], pos_z[index]))
590                    if len(x_line) > 0:
591                        x_lines.append(x_line)
592                        y_lines.append(y_line)
593                        z_lines.append(z_line)
594                except:
595                    logger.error(sys.exc_value)
596
597            output = MagSLD(pos_x, pos_y, pos_z, sld_n, sld_mx, sld_my, sld_mz)
598            output.set_conect_lines(x_line, y_line, z_line)
599            output.filename = os.path.basename(path)
600            output.set_pix_type('atom')
601            output.set_pixel_symbols(pix_symbol)
602            output.set_nodes()
603            output.set_pixel_volumes(vol_pix)
604            output.sld_unit = '1/A^(2)'
605            return output
606        except:
607            raise RuntimeError, "%s is not a sld file" % path
608
609    def write(self, path, data):
610        """
611        Write
612        """
613        print("Not implemented... ")
614
615class SLDReader(object):
616    """
617    Class to load ascii files (7 columns).
618    """
619    ## File type
620    type_name = "SLD ASCII"
621    ## Wildcards
622    type = ["sld files (*.SLD, *.sld)|*.sld",
623            "txt files (*.TXT, *.txt)|*.txt",
624            "all files (*.*)|*.*"]
625    ## List of allowed extensions
626    ext = ['.sld', '.SLD', '.txt', '.TXT', '.*']
627    def read(self, path):
628        """
629        Load data file
630        :param path: file path
631        :return MagSLD: x, y, z, sld_n, sld_mx, sld_my, sld_mz
632        :raise RuntimeError: when the file can't be opened
633        :raise ValueError: when the length of the data vectors are inconsistent
634        """
635        try:
636            pos_x = np.zeros(0)
637            pos_y = np.zeros(0)
638            pos_z = np.zeros(0)
639            sld_n = np.zeros(0)
640            sld_mx = np.zeros(0)
641            sld_my = np.zeros(0)
642            sld_mz = np.zeros(0)
643            try:
644                # Use numpy to speed up loading
645                input_f = np.loadtxt(path, dtype='float', skiprows=1,
646                                        ndmin=1, unpack=True)
647                pos_x = np.array(input_f[0])
648                pos_y = np.array(input_f[1])
649                pos_z = np.array(input_f[2])
650                sld_n = np.array(input_f[3])
651                sld_mx = np.array(input_f[4])
652                sld_my = np.array(input_f[5])
653                sld_mz = np.array(input_f[6])
654                ncols = len(input_f)
655                if ncols == 8:
656                    vol_pix = np.array(input_f[7])
657                elif ncols == 7:
658                    vol_pix = None
659            except:
660                # For older version of numpy
661                input_f = open(path, 'rb')
662                buff = input_f.read()
663                lines = buff.split('\n')
664                input_f.close()
665                for line in lines:
666                    toks = line.split()
667                    try:
668                        _pos_x = float(toks[0])
669                        _pos_y = float(toks[1])
670                        _pos_z = float(toks[2])
671                        _sld_n = float(toks[3])
672                        _sld_mx = float(toks[4])
673                        _sld_my = float(toks[5])
674                        _sld_mz = float(toks[6])
675                        pos_x = np.append(pos_x, _pos_x)
676                        pos_y = np.append(pos_y, _pos_y)
677                        pos_z = np.append(pos_z, _pos_z)
678                        sld_n = np.append(sld_n, _sld_n)
679                        sld_mx = np.append(sld_mx, _sld_mx)
680                        sld_my = np.append(sld_my, _sld_my)
681                        sld_mz = np.append(sld_mz, _sld_mz)
682                        try:
683                            _vol_pix = float(toks[7])
684                            vol_pix = np.append(vol_pix, _vol_pix)
685                        except:
686                            vol_pix = None
687                    except:
688                        # Skip non-data lines
689                        logger.error(sys.exc_value)
690            output = MagSLD(pos_x, pos_y, pos_z, sld_n,
691                            sld_mx, sld_my, sld_mz)
692            output.filename = os.path.basename(path)
693            output.set_pix_type('pixel')
694            output.set_pixel_symbols('pixel')
695            if vol_pix is not None:
696                output.set_pixel_volumes(vol_pix)
697            return output
698        except:
699            raise RuntimeError, "%s is not a sld file" % path
700
701    def write(self, path, data):
702        """
703        Write sld file
704        :Param path: file path
705        :Param data: MagSLD data object
706        """
707        if path is None:
708            raise ValueError, "Missing the file path."
709        if data is None:
710            raise ValueError, "Missing the data to save."
711        x_val = data.pos_x
712        y_val = data.pos_y
713        z_val = data.pos_z
714        vol_pix = data.vol_pix
715        length = len(x_val)
716        sld_n = data.sld_n
717        if sld_n is None:
718            sld_n = np.zeros(length)
719        sld_mx = data.sld_mx
720        if sld_mx is None:
721            sld_mx = np.zeros(length)
722            sld_my = np.zeros(length)
723            sld_mz = np.zeros(length)
724        else:
725            sld_my = data.sld_my
726            sld_mz = data.sld_mz
727        out = open(path, 'w')
728        # First Line: Column names
729        out.write("X  Y  Z  SLDN SLDMx  SLDMy  SLDMz VOLUMEpix")
730        for ind in range(length):
731            out.write("\n%g  %g  %g  %g  %g  %g  %g %g" % \
732                      (x_val[ind], y_val[ind], z_val[ind], sld_n[ind],
733                       sld_mx[ind], sld_my[ind], sld_mz[ind], vol_pix[ind]))
734        out.close()
735
736
737class OMFData(object):
738    """
739    OMF Data.
740    """
741    _meshunit = "A"
742    _valueunit = "A^(-2)"
743    def __init__(self):
744        """
745        Init for mag SLD
746        """
747        self.filename = 'default'
748        self.oommf = ''
749        self.title = ''
750        self.desc = ''
751        self.meshtype = ''
752        self.meshunit = self._meshunit
753        self.valueunit = self._valueunit
754        self.xbase = 0.0
755        self.ybase = 0.0
756        self.zbase = 0.0
757        self.xstepsize = 6.0
758        self.ystepsize = 6.0
759        self.zstepsize = 6.0
760        self.xnodes = 10.0
761        self.ynodes = 10.0
762        self.znodes = 10.0
763        self.xmin = 0.0
764        self.ymin = 0.0
765        self.zmin = 0.0
766        self.xmax = 60.0
767        self.ymax = 60.0
768        self.zmax = 60.0
769        self.mx = None
770        self.my = None
771        self.mz = None
772        self.valuemultiplier = 1.
773        self.valuerangeminmag = 0
774        self.valuerangemaxmag = 0
775
776    def __str__(self):
777        """
778        doc strings
779        """
780        _str = "Type:            %s\n" % self.__class__.__name__
781        _str += "File:            %s\n" % self.filename
782        _str += "OOMMF:           %s\n" % self.oommf
783        _str += "Title:           %s\n" % self.title
784        _str += "Desc:            %s\n" % self.desc
785        _str += "meshtype:        %s\n" % self.meshtype
786        _str += "meshunit:        %s\n" % str(self.meshunit)
787        _str += "xbase:           %s [%s]\n" % (str(self.xbase), self.meshunit)
788        _str += "ybase:           %s [%s]\n" % (str(self.ybase), self.meshunit)
789        _str += "zbase:           %s [%s]\n" % (str(self.zbase), self.meshunit)
790        _str += "xstepsize:       %s [%s]\n" % (str(self.xstepsize),
791                                                self.meshunit)
792        _str += "ystepsize:       %s [%s]\n" % (str(self.ystepsize),
793                                                self.meshunit)
794        _str += "zstepsize:       %s [%s]\n" % (str(self.zstepsize),
795                                                self.meshunit)
796        _str += "xnodes:          %s\n" % str(self.xnodes)
797        _str += "ynodes:          %s\n" % str(self.ynodes)
798        _str += "znodes:          %s\n" % str(self.znodes)
799        _str += "xmin:            %s [%s]\n" % (str(self.xmin), self.meshunit)
800        _str += "ymin:            %s [%s]\n" % (str(self.ymin), self.meshunit)
801        _str += "zmin:            %s [%s]\n" % (str(self.zmin), self.meshunit)
802        _str += "xmax:            %s [%s]\n" % (str(self.xmax), self.meshunit)
803        _str += "ymax:            %s [%s]\n" % (str(self.ymax), self.meshunit)
804        _str += "zmax:            %s [%s]\n" % (str(self.zmax), self.meshunit)
805        _str += "valueunit:       %s\n" % self.valueunit
806        _str += "valuemultiplier: %s\n" % str(self.valuemultiplier)
807        _str += "ValueRangeMinMag:%s [%s]\n" % (str(self.valuerangeminmag),
808                                                self.valueunit)
809        _str += "ValueRangeMaxMag:%s [%s]\n" % (str(self.valuerangemaxmag),
810                                                self.valueunit)
811        return _str
812
813    def set_m(self, mx, my, mz):
814        """
815        Set the Mx, My, Mz values
816        """
817        self.mx = mx
818        self.my = my
819        self.mz = mz
820
821class MagSLD(object):
822    """
823    Magnetic SLD.
824    """
825    pos_x = None
826    pos_y = None
827    pos_z = None
828    sld_n = None
829    sld_mx = None
830    sld_my = None
831    sld_mz = None
832    # Units
833    _pos_unit = 'A'
834    _sld_unit = '1/A^(2)'
835    _pix_type = 'pixel'
836
837    def __init__(self, pos_x, pos_y, pos_z, sld_n=None,
838                 sld_mx=None, sld_my=None, sld_mz=None, vol_pix=None):
839        """
840        Init for mag SLD
841        :params : All should be numpy 1D array
842        """
843        self.is_data = True
844        self.filename = ''
845        self.xstepsize = 6.0
846        self.ystepsize = 6.0
847        self.zstepsize = 6.0
848        self.xnodes = 10.0
849        self.ynodes = 10.0
850        self.znodes = 10.0
851        self.has_stepsize = False
852        self.has_conect = False
853        self.pos_unit = self._pos_unit
854        self.sld_unit = self._sld_unit
855        self.pix_type = 'pixel'
856        self.pos_x = pos_x
857        self.pos_y = pos_y
858        self.pos_z = pos_z
859        self.sld_n = sld_n
860        self.line_x = None
861        self.line_y = None
862        self.line_z = None
863        self.sld_mx = sld_mx
864        self.sld_my = sld_my
865        self.sld_mz = sld_mz
866        self.vol_pix = vol_pix
867        self.sld_m = None
868        self.sld_phi = None
869        self.sld_theta = None
870        self.pix_symbol = None
871        if sld_mx is not None and sld_my is not None and sld_mz is not None:
872            self.set_sldms(sld_mx, sld_my, sld_mz)
873        self.set_nodes()
874
875    def __str__(self):
876        """
877        doc strings
878        """
879        _str = "Type:       %s\n" % self.__class__.__name__
880        _str += "File:       %s\n" % self.filename
881        _str += "Axis_unit:  %s\n" % self.pos_unit
882        _str += "SLD_unit:   %s\n" % self.sld_unit
883        return _str
884
885    def set_pix_type(self, pix_type):
886        """
887        Set pixel type
888        :Param pix_type: string, 'pixel' or 'atom'
889        """
890        self.pix_type = pix_type
891
892    def set_sldn(self, sld_n):
893        """
894        Sets neutron SLD
895        """
896        if sld_n.__class__.__name__ == 'float':
897            if self.is_data:
898                # For data, put the value to only the pixels w non-zero M
899                is_nonzero = (np.fabs(self.sld_mx) +
900                              np.fabs(self.sld_my) +
901                              np.fabs(self.sld_mz)).nonzero()
902                self.sld_n = np.zeros(len(self.pos_x))
903                if len(self.sld_n[is_nonzero]) > 0:
904                    self.sld_n[is_nonzero] = sld_n
905                else:
906                    self.sld_n.fill(sld_n)
907            else:
908                # For non-data, put the value to all the pixels
909                self.sld_n = np.ones(len(self.pos_x)) * sld_n
910        else:
911            self.sld_n = sld_n
912
913    def set_sldms(self, sld_mx, sld_my, sld_mz):
914        r"""
915        Sets (\|m\|, m_theta, m_phi)
916        """
917        if sld_mx.__class__.__name__ == 'float':
918            self.sld_mx = np.ones(len(self.pos_x)) * sld_mx
919        else:
920            self.sld_mx = sld_mx
921        if sld_my.__class__.__name__ == 'float':
922            self.sld_my = np.ones(len(self.pos_x)) * sld_my
923        else:
924            self.sld_my = sld_my
925        if sld_mz.__class__.__name__ == 'float':
926            self.sld_mz = np.ones(len(self.pos_x)) * sld_mz
927        else:
928            self.sld_mz = sld_mz
929
930        sld_m = np.sqrt(sld_mx * sld_mx + sld_my * sld_my + \
931                                sld_mz * sld_mz)
932        self.sld_m = sld_m
933
934    def set_pixel_symbols(self, symbol='pixel'):
935        """
936        Set pixel
937        :Params pixel: str; pixel or atomic symbol, or array of strings
938        """
939        if self.sld_n is None:
940            return
941        if symbol.__class__.__name__ == 'str':
942            self.pix_symbol = np.repeat(symbol, len(self.sld_n))
943        else:
944            self.pix_symbol = symbol
945
946    def set_pixel_volumes(self, vol):
947        """
948        Set pixel volumes
949        :Params pixel: str; pixel or atomic symbol, or array of strings
950        """
951        if self.sld_n is None:
952            return
953        if vol.__class__.__name__ == 'ndarray':
954            self.vol_pix = vol
955        elif vol.__class__.__name__.count('float') > 0:
956            self.vol_pix = np.repeat(vol, len(self.sld_n))
957        else:
958            self.vol_pix = None
959
960    def get_sldn(self):
961        """
962        Returns nuclear sld
963        """
964        return self.sld_n
965
966    def set_nodes(self):
967        """
968        Set xnodes, ynodes, and znodes
969        """
970        self.set_stepsize()
971        if self.pix_type == 'pixel':
972            try:
973                xdist = (max(self.pos_x) - min(self.pos_x)) / self.xstepsize
974                ydist = (max(self.pos_y) - min(self.pos_y)) / self.ystepsize
975                zdist = (max(self.pos_z) - min(self.pos_z)) / self.zstepsize
976                self.xnodes = int(xdist) + 1
977                self.ynodes = int(ydist) + 1
978                self.znodes = int(zdist) + 1
979            except:
980                self.xnodes = None
981                self.ynodes = None
982                self.znodes = None
983        else:
984            self.xnodes = None
985            self.ynodes = None
986            self.znodes = None
987
988    def set_stepsize(self):
989        """
990        Set xtepsize, ystepsize, and zstepsize
991        """
992        if self.pix_type == 'pixel':
993            try:
994                xpos_pre = self.pos_x[0]
995                ypos_pre = self.pos_y[0]
996                zpos_pre = self.pos_z[0]
997                for x_pos in self.pos_x:
998                    if xpos_pre != x_pos:
999                        self.xstepsize = np.fabs(x_pos - xpos_pre)
1000                        break
1001                for y_pos in self.pos_y:
1002                    if ypos_pre != y_pos:
1003                        self.ystepsize = np.fabs(y_pos - ypos_pre)
1004                        break
1005                for z_pos in self.pos_z:
1006                    if zpos_pre != z_pos:
1007                        self.zstepsize = np.fabs(z_pos - zpos_pre)
1008                        break
1009                #default pix volume
1010                self.vol_pix = np.ones(len(self.pos_x))
1011                vol = self.xstepsize * self.ystepsize * self.zstepsize
1012                self.set_pixel_volumes(vol)
1013                self.has_stepsize = True
1014            except:
1015                self.xstepsize = None
1016                self.ystepsize = None
1017                self.zstepsize = None
1018                self.vol_pix = None
1019                self.has_stepsize = False
1020        else:
1021            self.xstepsize = None
1022            self.ystepsize = None
1023            self.zstepsize = None
1024            self.has_stepsize = True
1025        return self.xstepsize, self.ystepsize, self.zstepsize
1026
1027    def set_conect_lines(self, line_x, line_y, line_z):
1028        """
1029        Set bonding line data if taken from pdb
1030        """
1031        if line_x.__class__.__name__ != 'list' or len(line_x) < 1:
1032            return
1033        if line_y.__class__.__name__ != 'list' or len(line_y) < 1:
1034            return
1035        if line_z.__class__.__name__ != 'list' or len(line_z) < 1:
1036            return
1037        self.has_conect = True
1038        self.line_x = line_x
1039        self.line_y = line_y
1040        self.line_z = line_z
1041
1042def test_load():
1043    """
1044        Test code
1045    """
1046    from mpl_toolkits.mplot3d import Axes3D
1047    current_dir = os.path.abspath(os.path.curdir)
1048    print(current_dir)
1049    for i in range(6):
1050        current_dir, _ = os.path.split(current_dir)
1051        tfile = os.path.join(current_dir, "test", "CoreXY_ShellZ.txt")
1052        ofile = os.path.join(current_dir, "test", "A_Raw_Example-1.omf")
1053        if os.path.isfile(tfile):
1054            tfpath = tfile
1055            ofpath = ofile
1056            break
1057    reader = SLDReader()
1058    oreader = OMFReader()
1059    output = reader.read(tfpath)
1060    ooutput = oreader.read(ofpath)
1061    foutput = OMF2SLD()
1062    foutput.set_data(ooutput)
1063
1064    import matplotlib.pyplot as plt
1065    fig = plt.figure()
1066    ax = Axes3D(fig)
1067    ax.plot(output.pos_x, output.pos_y, output.pos_z, '.', c="g",
1068            alpha=0.7, markeredgecolor='gray', rasterized=True)
1069    gap = 7
1070    max_mx = max(output.sld_mx)
1071    max_my = max(output.sld_my)
1072    max_mz = max(output.sld_mz)
1073    max_m = max(max_mx, max_my, max_mz)
1074    x2 = output.pos_x+output.sld_mx/max_m * gap
1075    y2 = output.pos_y+output.sld_my/max_m * gap
1076    z2 = output.pos_z+output.sld_mz/max_m * gap
1077    x_arrow = np.column_stack((output.pos_x, x2))
1078    y_arrow = np.column_stack((output.pos_y, y2))
1079    z_arrow = np.column_stack((output.pos_z, z2))
1080    unit_x2 = output.sld_mx / max_m
1081    unit_y2 = output.sld_my / max_m
1082    unit_z2 = output.sld_mz / max_m
1083    color_x = np.fabs(unit_x2 * 0.8)
1084    color_y = np.fabs(unit_y2 * 0.8)
1085    color_z = np.fabs(unit_z2 * 0.8)
1086    colors = np.column_stack((color_x, color_y, color_z))
1087    plt.show()
1088
1089def test():
1090    """
1091        Test code
1092    """
1093    current_dir = os.path.abspath(os.path.curdir)
1094    for i in range(3):
1095        current_dir, _ = os.path.split(current_dir)
1096        ofile = os.path.join(current_dir, "test", "A_Raw_Example-1.omf")
1097        if os.path.isfile(ofile):
1098            ofpath = ofile
1099            break
1100    oreader = OMFReader()
1101    ooutput = oreader.read(ofpath)
1102    foutput = OMF2SLD()
1103    foutput.set_data(ooutput)
1104    writer = SLDReader()
1105    writer.write(os.path.join(os.path.dirname(ofpath), "out.txt"),
1106                 foutput.output)
1107    model = GenSAS()
1108    model.set_sld_data(foutput.output)
1109    x = np.arange(1000)/10000. + 1e-5
1110    y = np.arange(1000)/10000. + 1e-5
1111    i = np.zeros(1000)
1112    model.runXY([x, y, i])
1113
1114if __name__ == "__main__":
1115    test()
1116    test_load()
Note: See TracBrowser for help on using the repository browser.