diff --git a/vlasov/beamloading.py b/vlasov/beamloading.py
index a91b5b22b2405942463d7d0a191af2e978584731..45383e5d93ba08a5c16f153012c3b2290a8f6b29 100644
--- a/vlasov/beamloading.py
+++ b/vlasov/beamloading.py
@@ -33,8 +33,8 @@ class BeamLoadingVlasov():
     ring : Synchrotron object
     cavity_list : list of CavityResonator objects
     I0 : beam current in [A].
-    auto_set_MC_theta : if True, allow class to change generator phase for
-        CavityResonator objetcs with m = 1
+    auto_set_MC_theta : if True, allow class to change cavity phase for
+        CavityResonator objetcs with m = 1 (i.e. main cavities)
     F : list of form factor amplitude
     PHI : list of form factor phase
     B1 : lower intergration boundary
@@ -68,6 +68,10 @@ class BeamLoadingVlasov():
             * self.ring.E0 * self.ring.L)
         self.ug = np.zeros((self.n_cavity,))
         self.ub = np.zeros((self.n_cavity,))
+        self.update_potentials()
+            
+    def update_potentials(self):
+        """Update potentials with cavity and ring data."""
         for i in range(self.n_cavity):
             cavity = self.cavity_list[i]
             self.ug[i] = cavity.Vg / (
@@ -128,20 +132,24 @@ class BeamLoadingVlasov():
         B = quad(f, self.B1, self.B2)
         return A[0] / B[0]
 
-    def to_solve(self, x):
+    def to_solve(self, x, CM=True):
         """System of non-linear equation to solve to find the form factors F
         and PHI at equilibrum.
         The system is composed of Eq. (B6) and (B7) of [1] for each cavity.
-        If auto_set_MC_theta == True, the system also find the generator phase
-        to impose energy balance.
+        If auto_set_MC_theta == True, the system also find the main cavity 
+        phase to impose energy balance or cancel center of mass offset.
+        If CM is True, the system imposes zero center of mass offset,
+        if False, the system imposes energy balance.
         """
-        # Update values of F, PHI and theta_g
+        # Update values of F, PHI and theta
         if self.auto_set_MC_theta:
             self.F = x[:-1:2]
             for i in range(self.n_cavity):
                 cavity = self.cavity_list[i]
                 if cavity.m == 1:
-                    cavity.theta_g = x[-1]
+                    cavity.theta = x[-1]
+                    cavity.set_generator(0.5)
+                    self.update_potentials()
         else:
             self.F = x[::2]
         self.PHI = x[1::2]
@@ -155,8 +163,11 @@ class BeamLoadingVlasov():
                     lambda y: self.uexp(y), lambda y: np.cos(self.ring.k1 * cavity.m * y))
                 res[2 * i + 1] = self.F[i] * np.sin(self.PHI[i]) - self.integrate_func(
                     lambda y: self.uexp(y), lambda y: np.sin(self.ring.k1 * cavity.m * y))
-            # Factor 1e-8 for better convergence
-            res[self.n_cavity * 2] = self.energy_balance() * 1e-8
+            # Factor 1e-8 or 1e12 for better convergence
+            if CM is True:
+                res[self.n_cavity * 2] = self.center_of_mass() * 1e12
+            else:
+                res[self.n_cavity * 2] = self.energy_balance() * 1e-8
         else:
             res = np.zeros((self.n_cavity * 2,))
             for i in range(self.n_cavity):
@@ -249,7 +260,8 @@ class BeamLoadingVlasov():
         variance = np.average((z0 - average)**2, weights=values)
         return np.sqrt(variance)
 
-    def beam_equilibrium(self, x0=None, tol=1e-4, method='hybr', options=None, plot = False):
+    def beam_equilibrium(self, x0=None, tol=1e-4, method='hybr', options=None, 
+                         plot = False, CM=True):
         """Solve system of non-linear equation to find the form factors F
         and PHI at equilibrum. Can be used to compute the equilibrium bunch
         profile.
@@ -261,6 +273,8 @@ class BeamLoadingVlasov():
         method : method used by scipy.optimize.root to solve the system
         options : options given to scipy.optimize.root
         plot : if True, plot the equilibrium bunch profile
+        CM : if True, the system imposes zero center of mass offset,
+        if False, the system imposes energy balance.
         
         Returns
         -------
@@ -269,12 +283,17 @@ class BeamLoadingVlasov():
         if x0 is None:
             x0 = [1, 0] * self.n_cavity
             if self.auto_set_MC_theta:
-                x0 = x0 + [self.cavity_list[0].theta_g]
+                x0 = x0 + [self.cavity_list[0].theta]
 
-        print("The initial energy balance is " +
-              str(self.energy_balance()) + " eV")
+        if CM:
+            print("The initial center of mass offset is " +
+                  str(self.center_of_mass()*1e12) + " ps")
+        else:
+            print("The initial energy balance is " +
+                  str(self.energy_balance()) + " eV")
 
-        sol = root(self.to_solve, x0, tol=tol, method=method, options=options)
+        sol = root(lambda x : self.to_solve(x, CM), x0, tol=tol, method=method, 
+                   options=options)
 
         # Update values of F, PHI and theta_g
         if self.auto_set_MC_theta:
@@ -282,13 +301,17 @@ class BeamLoadingVlasov():
             for i in range(self.n_cavity):
                 cavity = self.cavity_list[i]
                 if cavity.m == 1:
-                    cavity.theta_g = sol.x[-1]
+                    cavity.theta = sol.x[-1]
         else:
             self.F = sol.x[::2]
         self.PHI = sol.x[1::2]
 
-        print("The final energy balance is " +
-              str(self.energy_balance()) + " eV")
+        if CM:
+            print("The final center of mass offset is " +
+                  str(self.center_of_mass()*1e12) + " ps")
+        else:
+            print("The final energy balance is " +
+                  str(self.energy_balance()) + " eV")
         print("The algorithm has converged: " + str(sol.success))
         
         if plot: