Doing many iterations of curve_fit in one go for piecewise function

I'm trying to perform what are many iterations of Scipy's curve_fit at once in order to avoid loops and therefore increase speed.

This is very similar to this problem, which was solved. However, the fact that the functions are piece-wise (discontinuous) makes so that that solution isn't applicable here.

Consider this example:

import numpy as np
from numpy import random as rng
from scipy.optimize import curve_fit
rng.seed(0)
N=20
X=np.logspace(-1,1,N)
Y = np.zeros((4, N))
for i in range(0,4):
    b = i+1
    a = b
    print(a,b)
    Y[i] = (X/b)**(-a) #+ 0.01 * rng.randn(6)
    Y[i, X>b] = 1

This yields these arrays:

Which as you can see are discontinuous at X==b. I can retrieve the original values of a and b by using curve_fit iteratively:

def plaw(r, a, b):
    """ Theoretical power law for the shape of the normalized conditional density """
    import numpy as np
    return np.piecewise(r, [r < b, r >= b], [lambda x: (x/b)**-a, lambda x: 1])


coeffs=[]
for ix in range(Y.shape[0]):
    print(ix)
    c0, pcov = curve_fit(plaw, X, Y[ix])
    coeffs.append(c0)

But this process can be very slow depending of the size of X, Y and the loop, so I'm trying to speed things up by trying to get coeffs without the need for a loop. So far I haven't had any luck.

Things that might be important:

• X and Y only contain positive values
• a and b are always positive
• Although the data to fit in this example is smooth (for the sake of simplicity), the real data has noise

EDIT

This is as far as I've gotten:

y=np.ma.masked_where(Y<1.01, Y)

lX = np.log(X)
lY = np.log(y)
A = np.vstack([lX, np.ones(len(lX))]).T
m,c=np.linalg.lstsq(A, lY.T)[0]

print('a=',-m)
print('b=',np.exp(-c/m))

But even without any noise the output is:

a= [0.18978965578339158 1.1353633705997466 2.220234483915197 3.3324502660995714]
b= [339.4090881838179 7.95073481873057 6.296592007396107 6.402567167503574]

Which is way worse than I was hoping to get.