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

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.2.2ticket-1009ticket-1094-headlessticket-1242-2d-resolutionticket-1243ticket-1249ticket885unittest-saveload
Last change on this file since ad92c5a was 7432acb, checked in by andyfaff, 8 years ago

MAINT: search+replace '!= None' by 'is not None'

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