Source code for gyroid.util

# -*- coding: utf-8 -*-
"""
gyroid.util
===========

"""

import numpy as np
import scipy.io
import matplotlib.pyplot as plt
from matplotlib import colors
from mayavi import mlab

__all__ = ["render_structure_1d","render_structure_2d","render_structure_3d"]

[docs]def render_structure_1d(basis,grid,Na,c, save_data=True,data_name="struct1d.mat", save_img=True,show_img=True, img_name="struct1d.png", **kwargs): ''' Calculate and render 1D structure for given SABF and unit cell. :param basis: a set of SABFs :type basis: :class:`Basis` :param Na: number of grids in **a** of the unit cell. :type Na: integer :param c: coefficients for each SABF :type c: 1D `numpy.array` :param save_data: if True, save data in file with Matlab mat format :type save_data: bool :param data_name: the file name of the data file :type data_name: string :param save_img: if True, save image in file, the format is determined by the extension of the image file name :type save_img: bool :param img_name: the file name of the image file :type img_name: string :param show_img: if True, show image on the screen :type show_img: bool :param kwargs: any extra key words arguments will be passed to plot functions ''' #struct = basis.generate_structure(Na,c) struct = basis.generate_structure_by_fft((Na,),c,grid) # For debug only #print basis.fft2sabf(np.fft.fftn(struct),grid) a = 1.0 * basis.shape.h[0,0] rx = np.array([a*i/Na for i in np.arange(Na)]) if save_data: scipy.io.savemat(data_name,{"rx":rx,"struct":struct}) if save_img or show_img: plt.plot(rx,struct,**kwargs) if save_img: plt.savefig(img_name) if show_img: plt.show() return rx,struct
[docs]def render_structure_2d(basis,grid,Na,Nb,c, save_data=True,data_name="struct2d.mat", save_img=True,show_img=True, img_name="struct2d.png", levels=None,cmap=None, **kwargs): ''' Calculate and render 2D structure for given SABF and unit cell. :param basis: a set of SABFs :type basis: :class:`Basis` :param Na: number of grids in **a** of the unit cell. :type Na: integer :param Nb: number of grids in **b** of the unit cell. :type Nb: integer :param c: coefficients for each SABF :type c: 1D `numpy.array` :param save_data: if True, save data in file with Matlab mat format :type save_data: bool :param data_name: the file name of the data file :type data_name: string :param save_img: if True, save image in file, the format is determined by the extension of the image file name :type save_img: bool :param img_name: the file name of the image file :type img_name: string :param show_img: if True, show image on the screen :type show_img: bool :param kwargs: any extra key words arguments will be passed to plot functions ''' #struct = basis.generate_structure((Na,Nb),c) struct = basis.generate_structure_by_fft((Na,Nb),c,grid) # For debug only #print basis.fft2sabf(np.fft.fftn(struct),grid) rx = np.zeros((Na,Nb)) ry = np.zeros((Na,Nb)) for (i,j) in np.ndindex(Na,Nb): x = (1.0*np.array([i,j])) / (Na,Nb) rx[i,j],ry[i,j] = np.dot(x,basis.shape.h) if save_data: scipy.io.savemat(data_name,{"rx":rx,"ry":ry,"struct":struct}) if save_img or show_img: dx = rx.max() - rx.min() dy = ry.max() - ry.min() w,h = plt.figaspect(float(dy/dx)) # float is must # No frame, white background, w/h aspect ratio figure fig = plt.figure(figsize=(w,h),frameon=False, dpi=80,facecolor='w') # full figure subplot, no border, no axes ax = fig.add_axes([0,0,1,1],frameon=False,axisbg='w') # no ticks ax.get_xaxis().set_visible(False) ax.get_yaxis().set_visible(False) # Default: there are 256 contour levels if levels is None: step = (struct.max() - struct.min()) / 256 levels = np.arange(struct.min(),struct.max()+step,step) # Default: colormap is monochromatic red if cmap is None: clr = np.zeros((256,3)) for i in np.arange(256): clr[i,0] = i / 255.0 cmap = colors.ListedColormap(clr) # actual plot ax.contourf(rx,ry,struct,levels=levels, cmap=cmap,antialiased=False,**kwargs) #ax.contourf(rx,ry,struct) if save_img: plt.savefig(img_name) if show_img: plt.show() return rx,ry,struct
[docs]def render_structure_3d(basis,grid,Na,Nb,Nc,c, save_data=True,data_name="struct3d.mat", save_img=True,show_img=True, img_name="struct3d.png", levels=None,cmap=None, **kwargs): ''' Calculate and render 3D structure for given SABF and unit cell. :param basis: a set of SABFs :type basis: :class:`Basis` :param Na: number of grids in **a** of the unit cell. :type Na: integer :param Nb: number of grids in **b** of the unit cell. :type Nb: integer :param Nc: number of grids in **c** of the unit cell. :type Nc: integer :param c: coefficients for each SABF :type c: 1D `numpy.array` :param save_data: if True, save data in file with Matlab mat format :type save_data: bool :param data_name: the file name of the data file :type data_name: string :param save_img: if True, save image in file, the format is determined by the extension of the image file name :type save_img: bool :param img_name: the file name of the image file :type img_name: string :param show_img: if True, show image on the screen :type show_img: bool :param kwargs: any extra key words arguments will be passed to plot functions ''' #struct = basis.generate_structure((Na,Nb,Nc),c) struct = basis.generate_structure_by_fft((Na,Nb,Nc),c,grid) # For debug only #print basis.fft2sabf(np.fft.fftn(struct),grid) rx = np.zeros((Na,Nb,Nc)) ry = np.zeros((Na,Nb,Nc)) rz = np.zeros((Na,Nb,Nc)) for (i,j,k) in np.ndindex(Na,Nb,Nc): x = (1.0*np.array([i,j,k])) / (Na,Nb,Nc) rx[i,j,k],ry[i,j,k],rz[i,j,k] = np.dot(x,basis.shape.h) if save_data: scipy.io.savemat(data_name, {"rx":rx,"ry":ry,"rz":rz,"struct":struct}) if save_img or show_img: mlab.contour3d(rx,ry,rz,struct,**kwargs) if save_img: mlab.savefig(img_name) if show_img: plt.show() return rx,ry,rz,struct

Related Topics