Python高效求解多组参数对应非线性方程的优化咨询

First, let's note a couple of small issues in your original code:

• x0 should likely be xs (since you define xs but iterate over x0)
• The variable w isn't defined—I'll assume it's a scalar constant for the examples below.

Your core bottleneck is the manual loop over each element, which doesn't leverage vectorization or parallel processing. Here are two efficient approaches to fix this:

Approach 1: Vectorize with numpy.vectorize (Simpler, Syntax Cleanup) #

While numpy.vectorize is technically a wrapper around loops under the hood, it makes your code cleaner and can reduce Python loop overhead slightly. We'll wrap your solver function to handle batch inputs directly:

import numpy as np
import tensorflow as tf
import scipy.optimize

kon = 0.01
mu = 1.5
fi = 0.5
kappa = 22
w = 1.0  # Assumed constant
n = 100
xs = tf.random.normal(shape=(n,), stddev=0.2)
eps = tf.random.normal(shape=(n,), stddev=0.17)
z = tf.sigmoid(tf.random.normal(shape=(n,), stddev=0.22))

# Convert TF tensors to numpy arrays (scipy works natively with numpy)
z_np = z.numpy()
eps_np = eps.numpy()
xs_np = xs.numpy()

def solve_single(ze, ei, xs_val):
    ei_exp = np.exp(ei)
    xs_exp = np.exp(xs_val)
    # Define the equation for this specific sample
    def F(hi):
        return (mu/fi)*np.log(hi) -(1-mu)*kappa*(hi)**(1+(1/fi)) - mu*(np.log(w*ei_exp*xs_exp)-np.log(kon)) - np.log(ze)
    return scipy.optimize.newton_krylov(F, 0.5)

# Vectorize the solver to handle batch inputs
vectorized_solve = np.vectorize(solve_single)
hvec = vectorized_solve(z_np, eps_np, xs_np)

Why this helps:

• Eliminates manual for loop boilerplate
• Handles batch inputs in a more readable way
• Retains your trusted newton_krylov solver for nonlinear equations

Approach 2: TensorFlow Batch Optimization (True Parallelism, Better for Large n) #

Since your inputs are already TensorFlow tensors, using TF's built-in optimizers (which leverage GPU/TPU acceleration if available) is the best way to get real parallel speedups. We'll frame the problem as minimizing the squared error of your equation (since we want F(hi) = 0), and optimize all hi values in parallel:

import numpy as np
import tensorflow as tf

kon = 0.01
mu = 1.5
fi = 0.5
kappa = 22
w = 1.0  # Assumed constant
n = 100
xs = tf.random.normal(shape=(n,), stddev=0.2)
eps = tf.random.normal(shape=(n,), stddev=0.17)
z = tf.sigmoid(tf.random.normal(shape=(n,), stddev=0.22))

# Precompute constant terms for all samples upfront
ei_exp = tf.exp(eps)
xs_exp = tf.exp(xs)
log_term = mu*(tf.math.log(w*ei_exp*xs_exp) - tf.math.log(kon)) + tf.math.log(z)

# Initialize all hi values (batch of n elements)
h = tf.Variable(tf.ones(n) * 0.5, dtype=tf.float32)

# Define loss as squared error of F(hi) = 0
def loss_fn():
    term1 = (mu / fi) * tf.math.log(h)
    term2 = (1 - mu) * kappa * tf.pow(h, 1 + (1/fi))
    F_val = term1 - term2 - log_term
    return tf.reduce_mean(tf.square(F_val))

# Use Adam optimizer to minimize the loss (adjust iterations/learning rate as needed)
optimizer = tf.optimizers.Adam(learning_rate=0.01)
for _ in range(1000):
    optimizer.minimize(loss_fn, var_list=[h])

# Extract the final solution
hvec = h.numpy()

Why this helps:

• Fully leverages TensorFlow's vectorized operations and hardware acceleration
• Solves all hi values in parallel, no per-sample loop
• Uses automatic differentiation, so you don't need to compute derivatives manually
• Far more scalable for large n (e.g., n=1000 or higher)

Notes on Choosing Between Approaches #

• Use Approach 1 if you need to stick with scipy.optimize solvers (e.g., for their robustness on specific nonlinear equations)
• Use Approach 2 if you want maximum performance, especially with GPU support, and are comfortable with gradient-based optimization

内容的提问来源于stack exchange，提问作者msh855