Function prints correct result, returns wrong value-CodePudding

I know there are millions of similar questions, but I couldn't solve my case. I have a data set consisting of 4 columns (energy versus imaginary part of dielectric function in 3 polarizations). I want to apply Kramers-Kronig relations to find the real part for every energy point, so the result will be in the same dimensions as input data. For the sake of mwe, the data file looks like this:

 1.92000     0.45314     0.45774     0.44148
 1.92100     0.45387     0.45846     0.44223
 1.92200     0.45460     0.45918     0.44297
 1.92300     0.45533     0.45990     0.44372
 1.92400     0.45605     0.46062     0.44446
 1.92500     0.45677     0.46134     0.44520

So far my code is:

import numpy as np
import math

data = np.loadtxt('opt.mgsin2gan', skiprows=1)
dw = data[1,0] - data[0,0]

## This is the expression being integrated.
def frac(x,pol):                
    arg = len(data)
    mid = 0
    for i in range(len(data)):
        if data[i,0] == x:
            pass
        else:
            mid  = data[i,0]*data[i,pol]/(data[i,0]**2-x**2)
#            return 1 (2/math.pi)*mid*dw
            result = 1 (2/math.pi)*mid*dw
    print(result)

## Evaluates the expression for every energy point
def grid(a):
    for i in data[:,0]:
        frac(i,a)
#    return frac(i,a)

## Evaluates the expression for each polarization
def polarization():
    for i in [1,2,3]:
        grid(i)
#    return grid(i)
if __name__ == "__main__":
    polarization()

It prints the result on the screen, but I can't save the data in a text file since it doesn't return anything. However, when I uncomment return lines in each function, I get the wrong result. How can I save the result in a numpy array? Also, how can I make this code more elegant? Making it simpler probably will make it run in shorter time.

CodePudding user response：

I am guessing that you are looking for something like this:

def frac(x,pol):                
    arg = len(data)
    mid = 0
    for i in range(len(data)):
        if data[i,0] == x:
            pass
        else:
            mid  = data[i,0]*data[i,pol]/(data[i,0]**2-x**2)
#            return 1 (2/math.pi)*mid*dw
            result = 1 (2/math.pi)*mid*dw
    print(result)
    return result

## Evaluates the expression for every energy point
def grid(a):
    # return list here
    return [frac(i, a) for i in data[:,0]]

## Evaluates the expression for each polarization
def polarization():
    # yield values for each i
    for i in [1,2,3]:
        yield from grid(i)

result = list(polarization())

I've changed the grid function to return a list, and polarization function to yield values for each i.

As an aside, line if data[i,0] == x can be dangerous when you are dealing with floating point numbers. Instead, you may want to check if two numbers are numerically close, e.g. using np.isclose.

It is likely that you can speed things up rewriting numpy code or using numba in this case.