diff --git a/mbtrack2/instability/__init__.py b/mbtrack2/instability/__init__.py
index 0fa6890eaf6414376c93057ec156b3a513fc360b..5c265bf664b149ef7e7681fbab68bc1c7f617d21 100644
--- a/mbtrack2/instability/__init__.py
+++ b/mbtrack2/instability/__init__.py
@@ -7,13 +7,11 @@ from mbtrack2.instability.instabilities import (
mbi_threshold,
rwmbi_growth_rate,
rwmbi_threshold,
+ transverse_space_charge_tune_shift,
)
from mbtrack2.instability.ions import (
fast_beam_ion,
ion_cross_section,
ion_frequency,
plot_critical_mass,
-)
-from mbtrack2.instability.space_charge import (
- transverse_space_charge_tune_shift,
-)
+)
\ No newline at end of file
diff --git a/mbtrack2/instability/instabilities.py b/mbtrack2/instability/instabilities.py
index 0a198eb98c90e9f55f16dc5a0805faa948889837..275faa687e92b7076ed78209ca60435ab31a59f8 100644
--- a/mbtrack2/instability/instabilities.py
+++ b/mbtrack2/instability/instabilities.py
@@ -454,3 +454,79 @@ def rwmbi_threshold(ring, beff, rho_material, plane='x'):
(1-frac_tune) * omega0) / (2*c*Z0*rho_material))**0.5
return Ith
+
+
+def transverse_space_charge_tune_shift(ring, bunch_current, **kwargs):
+ """
+ Return the (maximum) transverse space charge tune shift for a Gaussian
+ bunch in the linear approximation, see Eq.(1) of [1].
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ Ring parameters.
+ bunch_current : float
+ Bunch current in [A].
+ sigma_s : float, optional
+ RMS bunch length in [s].
+ Default is ring.sigma_0.
+ emit_x : float, optional
+ Horizontal emittance in [m.rad].
+ Default is ring.emit[0].
+ emit_y : float, optional
+ Vertical emittance in [m.rad].
+ Default is ring.emit[1].
+ use_lattice : bool, optional
+ If True, use beta fonctions along the lattice.
+ If False, local values of beta fonctions are used.
+ Default is ring.optics.use_local_values.
+ sigma_delta : float, optional
+ Relative energy spread.
+ Default is ring.sigma_delta.
+ gamma : float, optional
+ Relativistic Lorentz gamma.
+ Default is ring.gamma.
+
+ Returns
+ -------
+ tune_shift : array of shape (2,)
+ Horizontal and vertical space charge tune shift.
+
+ Reference
+ ---------
+ [1] : Antipov, S. A., Gubaidulin, V., Agapov, I., Garcia, E. C., &
+ Gamelin, A. (2024). Space Charge and Future Light Sources.
+ arXiv preprint arXiv:2409.08637.
+
+ """
+ sigma_s = kwargs.get('sigma_s', ring.sigma_0)
+ emit_x = kwargs.get('emit_x', ring.emit[0])
+ emit_y = kwargs.get('emit_y', ring.emit[1])
+ use_lattice = kwargs.get('use_lattice', not ring.optics.use_local_values)
+ sigma_delta = kwargs.get('sigma_delta', ring.sigma_delta)
+ gamma = kwargs.get('gamma', ring.gamma)
+
+ q = np.abs(ring.particle.charge)
+ m = ring.particle.mass
+ r_0 = 1 / (4*pi*epsilon_0) * q**2 / (m * c**2)
+ N = bunch_current / ring.f0 / q
+ sigma_z = sigma_s * c
+
+ if use_lattice:
+ s = np.linspace(0, ring.L, 1000)
+ beta = ring.optics.beta(s)
+ sig_x = (emit_x * beta[0] +
+ ring.optics.dispersion(s)[0]**2 * sigma_delta**2)**0.5
+ sig_y = (emit_y * beta[1] +
+ ring.optics.dispersion(s)[2]**2 * sigma_delta**2)**0.5
+ sig_xy = np.array([sig_x, sig_y])
+ return -r_0 * N / ((2 * pi)**1.5 * gamma**3 * sigma_z) * np.trapz(
+ beta / (sig_xy * (sig_y+sig_x)), s)
+ else:
+ beta = ring.optics.local_beta
+ sig_x = np.sqrt(beta[0] * emit_x)
+ sig_y = np.sqrt(beta[1] * emit_y)
+ sig_xy = np.array([sig_x, sig_y])
+ return -r_0 * N * ring.L / (
+ (2 * pi)**1.5 * gamma**3 * sigma_z) * beta / (sig_xy *
+ (sig_y+sig_x))
diff --git a/mbtrack2/instability/space_charge.py b/mbtrack2/instability/space_charge.py
deleted file mode 100644
index bd763e319da49195876d22d001dabac54974e185..0000000000000000000000000000000000000000
--- a/mbtrack2/instability/space_charge.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-General calculations about space charge such as tune shift.
-"""
-
-import numpy as np
-from scipy.constants import c, epsilon_0, pi
-
-
-def transverse_space_charge_tune_shift(ring, bunch_current, **kwargs):
- """
- Return the (maximum) transverse space charge tune shift for a Gaussian
- bunch in the linear approximation, see Eq.(1) of [1].
-
- Parameters
- ----------
- ring : Synchrotron object
- Ring parameters.
- bunch_current : float
- Bunch current in [A].
- sigma_s : float, optional
- RMS bunch length in [s].
- Default is ring.sigma_0.
- emit_x : float, optional
- Horizontal emittance in [m.rad].
- Default is ring.emit[0].
- emit_y : float, optional
- Vertical emittance in [m.rad].
- Default is ring.emit[1].
- use_lattice : bool, optional
- If True, use beta fonctions along the lattice.
- If False, local values of beta fonctions are used.
- Default is ring.optics.use_local_values.
- sigma_delta : float, optional
- Relative energy spread.
- Default is ring.sigma_delta.
- gamma : float, optional
- Relativistic Lorentz gamma.
- Default is ring.gamma.
-
- Returns
- -------
- tune_shift : array of shape (2,)
- Horizontal and vertical space charge tune shift.
-
- Reference
- ---------
- [1] : Antipov, S. A., Gubaidulin, V., Agapov, I., Garcia, E. C., &
- Gamelin, A. (2024). Space Charge and Future Light Sources.
- arXiv preprint arXiv:2409.08637.
-
- """
- sigma_s = kwargs.get('sigma_s', ring.sigma_0)
- emit_x = kwargs.get('emit_x', ring.emit[0])
- emit_y = kwargs.get('emit_y', ring.emit[1])
- use_lattice = kwargs.get('use_lattice', not ring.optics.use_local_values)
- sigma_delta = kwargs.get('sigma_delta', ring.sigma_delta)
- gamma = kwargs.get('gamma', ring.gamma)
-
- q = np.abs(ring.particle.charge)
- m = ring.particle.mass
- r_0 = 1 / (4*pi*epsilon_0) * q**2 / (m * c**2)
- N = bunch_current / ring.f0 / q
- sigma_z = sigma_s * c
-
- if use_lattice:
- s = np.linspace(0, ring.L, 1000)
- beta = ring.optics.beta(s)
- sig_x = (emit_x * beta[0] +
- ring.optics.dispersion(s)[0]**2 * sigma_delta**2)**0.5
- sig_y = (emit_y * beta[1] +
- ring.optics.dispersion(s)[2]**2 * sigma_delta**2)**0.5
- sig_xy = np.array([sig_x, sig_y])
- return -r_0 * N / ((2 * pi)**1.5 * gamma**3 * sigma_z) * np.trapz(
- beta / (sig_xy * (sig_y+sig_x)), s)
- else:
- beta = ring.optics.local_beta
- sig_x = np.sqrt(beta[0] * emit_x)
- sig_y = np.sqrt(beta[1] * emit_y)
- sig_xy = np.array([sig_x, sig_y])
- return -r_0 * N * ring.L / (
- (2 * pi)**1.5 * gamma**3 * sigma_z) * beta / (sig_xy *
- (sig_y+sig_x))