diff --git a/simpegMT/Utils/dataUtils.py b/simpegMT/Utils/dataUtils.py index f401bc59..99c238dc 100644 --- a/simpegMT/Utils/dataUtils.py +++ b/simpegMT/Utils/dataUtils.py @@ -4,6 +4,7 @@ import SimPEG as simpeg import simpegMT as simpegmt import numpy.lib.recfunctions as recFunc from scipy.constants import mu_0 +from scipy import interpolate as sciint def getAppRes(MTdata): # Make impedance @@ -19,6 +20,28 @@ def getAppRes(MTdata): zList.append(zc) return [appResPhs(zList[i][0],np.sum(zList[i][1:3])) for i in np.arange(len(zList))] +def rotateData(MTdata,rotAngle): + ''' + Function that rotates clockwist by rotAngle (- negative for a counter-clockwise rotation) + ''' + recData = MTdata.toRecArray('Complex') + impData = rec2ndarr(recData[['zxx','zxy','zyx','zyy']],complex) + # Make the rotation matrix + # c,s,zxx,zxy,zyx,zyy = sympy.symbols('c,s,zxx,zxy,zyx,zyy') + # rotM = sympy.Matrix([[c,-s],[s, c]]) + # zM = sympy.Matrix([[zxx,zxy],[zyx,zyy]]) + # rotM*zM*rotM.T + # [c*(c*zxx - s*zyx) - s*(c*zxy - s*zyy), c*(c*zxy - s*zyy) + s*(c*zxx - s*zyx)], + # [c*(c*zyx + s*zxx) - s*(c*zyy + s*zxy), c*(c*zyy + s*zxy) + s*(c*zyx + s*zxx)]]) + s = np.sin(-np.deg2rad(rotAngle)) + c = np.cos(-np.deg2rad(rotAngle)) + rotMat = np.array([[c,-s],[s,c]]) + rotData = (rotMat.dot(impData.reshape(-1,2,2).dot(rotMat.T))).transpose(1,0,2).reshape(-1,4) + outRec = recData.copy() + for nr,comp in enumerate(['zxx','zxy','zyx','zyy']): + outRec[comp] = rotData[:,nr] + return simpegmt.DataMT.DataMT.fromRecArray(outRec) + def appResPhs(freq,z): app_res = ((1./(8e-7*np.pi**2))/freq)*np.abs(z)**2 @@ -182,4 +205,37 @@ def convert3Dto1Dobject(MTdata,rxType3D='zyx'): mtData1DList.append(dat1D) # Return the the list of data. - return mtData1DList \ No newline at end of file + return mtData1DList + +def resampleMTdataAtFreq(MTdata,freqs): + ''' + Function to resample MTdata at set of frequencies + + ''' + # Make a rec array + MTrec = MTdata.toRecArray().data + + # Find unique locations + uniLoc = np.unique(MTrec[['x','y','z']]) + uniFreq = MTdata.survey.freqs + # Get the comps + dNames = KraEDIrec.dtype + + # Loop over all the locations and interpolate + for loc in uniLoc: + # Find the index of the station + ind = np.sqrt(np.sum((rec2ndarr(MTrec[['x','y','z']]) - rec2ndarr(loc))**2,axis=1)) < 1. # Find dist of 1 m accuracy + # Make a temporary recArray and interpolate all the components + tArrRec = np.concatenate((simpeg.mkvc(freqs,2),np.ones((len(freqs),1))*rec2ndarr(loc),np.nan*np.ones((len(freqs),12))),axis=1).view(dNames) + for comp in ['zxxr','zxxi','zxyr','zxyi','zyxr','zyxi','zyyr','zyyi','tzxr','tzxi','tzyr','tzyi']: + int1d = sciint.interp1d(MTrec[ind]['freq'],MTrec[ind][comp],bounds_error=False) + tArrRec[comp] = simpeg.mkvc(int1d(freqs),2) + + # Join together + try: + outRecArr = recFunc.stack_arrays((outRecArr,tArrRec)) + except NameError as e: + outRecArr = tArrRec + + # Make the MTdata and return + return simpegmt.DataMT.DataMT.fromRecArray(outRecArr) \ No newline at end of file