using numpy broadcasting / vectorization to build new array from other arrays

Someone gave me a more succinct way to accomplish this, finally eliminating the extra FOR loop. It basically hides the loop in a Pandas DataFrame groupby, but it more succinctly describes what the desired steps are:

def GainPctInd2(offset=0, nbars=2):    class GainPctIndFact2(CustomFactor):        window_length = nbars + offset        inputs = [USEquityPricing.close, ms.asset_classification.morningstar_industry_code]        def compute(self, today, assets, out, close, industries):            df = pd.DataFrame(index=assets, data={                    "gain": ((close[-1 - offset] / close[(-1 - offset) - (nbars - 1)]) - 1) * 100,                    "industry_codes": industries[-1]                 })            out[:] = df.groupby("industry_codes").transform(np.mean).values.flatten()    return GainPctIndFact2()

It does not improve the efficiency at all, according to my benchmarks, but it's probably easier to verify correctness. The one problem with their example is that it uses np.mean instead of np.nanmean, and np.nanmean drops the NaN values resulting in a shape mismatch if you try to use it. To fix the NaN issue, someone else suggested this:

def GainPctInd2(offset=0, nbars=2):    class GainPctIndFact2(CustomFactor):        window_length = nbars + offset        inputs = [USEquityPricing.close, ms.asset_classification.morningstar_industry_code]        def compute(self, today, assets, out, close, industries):            df = pd.DataFrame(index=assets, data={                    "gain": ((close[-1 - offset] / close[(-1 - offset) - (nbars - 1)]) - 1) * 100,                    "industry_codes": industries[-1]                 })            nans = isnan(df['industry_codes'])            notnan = ~nans            out[notnan] = df[df['industry_codes'].notnull()].groupby("industry_codes").transform(np.nanmean).values.flatten()            out[nans] = nan    return GainPctIndFact2()

– user36048