source: sasview/src/sas/calculator/sas_gen.py @ b9dbd6b

ESS_GUIESS_GUI_DocsESS_GUI_batch_fittingESS_GUI_bumps_abstractionESS_GUI_iss1116ESS_GUI_iss879ESS_GUI_iss959ESS_GUI_openclESS_GUI_orderingESS_GUI_sync_sascalccostrafo411magnetic_scattrelease-4.1.1release-4.1.2release-4.2.2release_4.0.1ticket-1009ticket-1094-headlessticket-1242-2d-resolutionticket-1243ticket-1249ticket885unittest-saveload
Last change on this file since b9dbd6b was b6627d9, checked in by Doucet, Mathieu <doucetm@…>, 10 years ago

pylint fixes

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