source: sasview/src/sas/sascalc/calculator/sas_gen.py @ 0765c00

Last change on this file since 0765c00 was 9a5097c, checked in by andyfaff, 8 years ago

MAINT: import numpy as np

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