diff --git a/src/wifi/model/he/he-phy.cc b/src/wifi/model/he/he-phy.cc
index 47cff923c..d9326459d 100644
--- a/src/wifi/model/he/he-phy.cc
+++ b/src/wifi/model/he/he-phy.cc
@@ -22,7 +22,6 @@
 #include "he-phy.h"
 
 #include "he-configuration.h"
-#include "he-ppdu.h"
 #include "obss-pd-algorithm.h"
 
 #include "ns3/ap-wifi-mac.h"
@@ -1251,10 +1250,22 @@ Ptr<SpectrumValue>
 HePhy::GetTxPowerSpectralDensity(double txPowerW,
                                  Ptr<const WifiPpdu> ppdu,
                                  const WifiTxVector& txVector) const
+{
+    auto hePpdu = DynamicCast<const HePpdu>(ppdu);
+    NS_ASSERT(hePpdu);
+    HePpdu::TxPsdFlag flag = hePpdu->GetTxPsdFlag();
+    return GetTxPowerSpectralDensity(txPowerW, ppdu, txVector, flag);
+}
+
+Ptr<SpectrumValue>
+HePhy::GetTxPowerSpectralDensity(double txPowerW,
+                                 Ptr<const WifiPpdu> ppdu,
+                                 const WifiTxVector& txVector,
+                                 HePpdu::TxPsdFlag flag) const
 {
     uint16_t centerFrequency = GetCenterFrequencyForChannelWidth(txVector);
     uint16_t channelWidth = txVector.GetChannelWidth();
-    NS_LOG_FUNCTION(this << centerFrequency << channelWidth << txPowerW);
+    NS_LOG_FUNCTION(this << centerFrequency << channelWidth << txPowerW << txVector);
     const auto& puncturedSubchannels = txVector.GetInactiveSubchannels();
     if (!puncturedSubchannels.empty())
     {
@@ -1265,9 +1276,6 @@ HePhy::GetTxPowerSpectralDensity(double txPowerW,
         NS_ASSERT(
             !puncturedSubchannels.at(p20IndexInBitmap)); // the primary channel cannot be punctured
     }
-    auto hePpdu = DynamicCast<const HePpdu>(ppdu);
-    NS_ASSERT(hePpdu);
-    HePpdu::TxPsdFlag flag = hePpdu->GetTxPsdFlag();
     const auto& txMaskRejectionParams = GetTxMaskRejectionParams();
     switch (ppdu->GetType())
     {
@@ -1275,7 +1283,7 @@ HePhy::GetTxPowerSpectralDensity(double txPowerW,
         if (flag == HePpdu::PSD_NON_HE_PORTION)
         {
             // non-OFDMA portion is sent only on the 20 MHz channels covering the RU
-            const uint16_t staId = GetStaId(hePpdu);
+            const uint16_t staId = GetStaId(ppdu);
             centerFrequency = GetCenterFrequencyForNonOfdmaPart(txVector, staId);
             const uint16_t ruWidth = HeRu::GetBandwidth(txVector.GetRu(staId).GetRuType());
             channelWidth = (ruWidth < 20) ? 20 : ruWidth;
@@ -1291,8 +1299,7 @@ HePhy::GetTxPowerSpectralDensity(double txPowerW,
         }
         else
         {
-            const auto band =
-                GetRuBandForTx(ppdu->GetTxVector(), GetStaId(hePpdu));
+            const auto band = GetRuBandForTx(txVector, GetStaId(ppdu));
             return WifiSpectrumValueHelper::CreateHeMuOfdmTxPowerSpectralDensity(
                 centerFrequency,
                 channelWidth,
@@ -1376,19 +1383,45 @@ HePhy::StartTx(Ptr<const WifiPpdu> ppdu, const WifiTxVector& txVector)
     NS_LOG_FUNCTION(this << ppdu);
     if (ppdu->GetType() == WIFI_PPDU_TYPE_UL_MU || ppdu->GetType() == WIFI_PPDU_TYPE_DL_MU)
     {
+        auto nonOfdmaTxPowerDbm =
+            m_wifiPhy->GetTxPowerForTransmission(ppdu) + m_wifiPhy->GetTxGain();
+
+        // temporarily set WifiPpdu flag to PSD_HE_PORTION for correct calculation of the OFDMA TX
+        // power
+        auto hePpdu = DynamicCast<const HePpdu>(ppdu);
+        NS_ASSERT(hePpdu);
+        hePpdu->SetTxPsdFlag(HePpdu::PSD_HE_PORTION);
+        auto ofdmaTxPowerDbm = m_wifiPhy->GetTxPowerForTransmission(ppdu) + m_wifiPhy->GetTxGain();
+        hePpdu->SetTxPsdFlag(HePpdu::PSD_NON_HE_PORTION);
+
         // non-OFDMA part
-        Time nonOfdmaDuration = ppdu->GetType() == WIFI_PPDU_TYPE_UL_MU
+        auto nonOfdmaDuration = ppdu->GetType() == WIFI_PPDU_TYPE_UL_MU
                                     ? CalculateNonOfdmaDurationForHeTb(txVector)
                                     : CalculateNonOfdmaDurationForHeMu(txVector);
-        Transmit(nonOfdmaDuration, ppdu, txVector, "non-OFDMA transmission");
+        auto nonOfdmaTxPowerSpectrum = GetTxPowerSpectralDensity(DbmToW(nonOfdmaTxPowerDbm),
+                                                                 ppdu,
+                                                                 txVector,
+                                                                 HePpdu::PSD_NON_HE_PORTION);
+        Transmit(nonOfdmaDuration,
+                 ppdu,
+                 txVector,
+                 nonOfdmaTxPowerDbm,
+                 nonOfdmaTxPowerSpectrum,
+                 "non-OFDMA transmission");
 
         // OFDMA part
-        Time ofdmaDuration = ppdu->GetTxDuration() - nonOfdmaDuration;
+        auto ofdmaDuration = ppdu->GetTxDuration() - nonOfdmaDuration;
+        auto ofdmaTxPowerSpectrum = GetTxPowerSpectralDensity(DbmToW(ofdmaTxPowerDbm),
+                                                              ppdu,
+                                                              txVector,
+                                                              HePpdu::PSD_HE_PORTION);
         Simulator::Schedule(nonOfdmaDuration,
                             &HePhy::StartTxOfdma,
                             this,
                             ppdu,
                             txVector,
+                            ofdmaTxPowerDbm,
+                            ofdmaTxPowerSpectrum,
                             ofdmaDuration);
     }
     else
@@ -1398,12 +1431,17 @@ HePhy::StartTx(Ptr<const WifiPpdu> ppdu, const WifiTxVector& txVector)
 }
 
 void
-HePhy::StartTxOfdma(Ptr<const WifiPpdu> ppdu, const WifiTxVector& txVector, Time ofdmaDuration)
+HePhy::StartTxOfdma(Ptr<const WifiPpdu> ppdu,
+                    const WifiTxVector& txVector,
+                    double txPowerDbm,
+                    Ptr<SpectrumValue> txPowerSpectrum,
+                    Time ofdmaDuration)
 {
+    NS_LOG_FUNCTION(this << ppdu << txPowerDbm << ofdmaDuration);
     auto hePpdu = DynamicCast<const HePpdu>(ppdu);
     NS_ASSERT(hePpdu);
     hePpdu->SetTxPsdFlag(HePpdu::PSD_HE_PORTION);
-    Transmit(ofdmaDuration, ppdu, txVector, "OFDMA transmission");
+    Transmit(ofdmaDuration, ppdu, txVector, txPowerDbm, txPowerSpectrum, "OFDMA transmission");
 }
 
 Time
diff --git a/src/wifi/model/he/he-phy.h b/src/wifi/model/he/he-phy.h
index 5c38bc927..07ebcc978 100644
--- a/src/wifi/model/he/he-phy.h
+++ b/src/wifi/model/he/he-phy.h
@@ -22,6 +22,8 @@
 #ifndef HE_PHY_H
 #define HE_PHY_H
 
+#include "he-ppdu.h"
+
 #include "ns3/callback.h"
 #include "ns3/vht-phy.h"
 #include "ns3/wifi-phy-band.h"
@@ -527,14 +529,34 @@ class HePhy : public VhtPhy
     uint8_t GetNumberBccEncoders(const WifiTxVector& txVector) const override;
     Time GetSymbolDuration(const WifiTxVector& txVector) const override;
 
+    /**
+     * This is a helper function to create the TX PSD of the non-HE and HE portions.
+     *
+     * \param txPowerW power in W to spread across the bands
+     * \param ppdu the PPDU that will be transmitted
+     * \param txVector the transmission parameters
+     * \param flag flag indicating whether the PSD is for non-HE portion or HE portion
+     * \return Pointer to SpectrumValue
+     */
+    Ptr<SpectrumValue> GetTxPowerSpectralDensity(double txPowerW,
+                                                 Ptr<const WifiPpdu> ppdu,
+                                                 const WifiTxVector& txVector,
+                                                 HePpdu::TxPsdFlag flag) const;
+
     /**
      * Start the transmission of the OFDMA part of the MU PPDU.
      *
      * \param ppdu the PPDU
-     * \param txVector the TXVECTOR
+     * \param txVector the TXVECTOR used for the transmission of the PPDU
+     * \param txPowerDbm the total TX power in dBm
+     * \param txPowerSpectrum the TX PSD
      * \param ofdmaDuration the duration of the OFDMA part
      */
-    void StartTxOfdma(Ptr<const WifiPpdu> ppdu, const WifiTxVector& txVector, Time ofdmaDuration);
+    void StartTxOfdma(Ptr<const WifiPpdu> ppdu,
+                      const WifiTxVector& txVector,
+                      double txPowerDbm,
+                      Ptr<SpectrumValue> txPowerSpectrum,
+                      Time ofdmaDuration);
 
     /**
      * Notify PHY state helper to switch to CCA busy state,
diff --git a/src/wifi/model/phy-entity.cc b/src/wifi/model/phy-entity.cc
index 3cde20667..6aae2c6ee 100644
--- a/src/wifi/model/phy-entity.cc
+++ b/src/wifi/model/phy-entity.cc
@@ -1282,26 +1282,26 @@ void
 PhyEntity::StartTx(Ptr<const WifiPpdu> ppdu, const WifiTxVector& txVector)
 {
     NS_LOG_FUNCTION(this << ppdu << txVector);
-    Transmit(ppdu->GetTxDuration(), ppdu, txVector, "transmission");
+    auto txPowerDbm = m_wifiPhy->GetTxPowerForTransmission(ppdu) + m_wifiPhy->GetTxGain();
+    auto txPowerSpectrum = GetTxPowerSpectralDensity(DbmToW(txPowerDbm), ppdu, txVector);
+    Transmit(ppdu->GetTxDuration(), ppdu, txVector, txPowerDbm, txPowerSpectrum, "transmission");
 }
 
 void
 PhyEntity::Transmit(Time txDuration,
                     Ptr<const WifiPpdu> ppdu,
                     const WifiTxVector& txVector,
-                    std::string type)
+                    double txPowerDbm,
+                    Ptr<SpectrumValue> txPowerSpectrum,
+                    const std::string& type)
 {
-    NS_LOG_FUNCTION(this << txDuration << ppdu << txVector << type);
-    double txPowerWatts =
-        DbmToW(m_wifiPhy->GetTxPowerForTransmission(ppdu) + m_wifiPhy->GetTxGain());
-    NS_LOG_DEBUG("Start " << type << ": signal power before antenna gain=" << WToDbm(txPowerWatts)
-                          << "dBm");
-    Ptr<SpectrumValue> txPowerSpectrum = GetTxPowerSpectralDensity(txPowerWatts, ppdu, txVector);
-    Ptr<WifiSpectrumSignalParameters> txParams = Create<WifiSpectrumSignalParameters>();
+    NS_LOG_FUNCTION(this << txDuration << ppdu << txVector << txPowerDbm << type);
+    NS_LOG_DEBUG("Start " << type << ": signal power before antenna gain=" << txPowerDbm << "dBm");
+    auto txParams = Create<WifiSpectrumSignalParameters>();
     txParams->duration = txDuration;
     txParams->psd = txPowerSpectrum;
     txParams->ppdu = ppdu;
-    NS_LOG_DEBUG("Starting " << type << " with power " << WToDbm(txPowerWatts) << " dBm on channel "
+    NS_LOG_DEBUG("Starting " << type << " with power " << txPowerDbm << " dBm on channel "
                              << +m_wifiPhy->GetChannelNumber() << " for "
                              << txParams->duration.As(Time::MS));
     NS_LOG_DEBUG("Starting " << type << " with integrated spectrum power "
diff --git a/src/wifi/model/phy-entity.h b/src/wifi/model/phy-entity.h
index 4ee8979f5..9031329de 100644
--- a/src/wifi/model/phy-entity.h
+++ b/src/wifi/model/phy-entity.h
@@ -488,12 +488,16 @@ class PhyEntity : public SimpleRefCount<PhyEntity>
      * \param txDuration the duration of the transmission
      * \param ppdu the PPDU to send
      * \param txVector the TXVECTOR used for the transmission of the PPDU
+     * \param txPowerDbm the total TX power in dBm
+     * \param txPowerSpectrum the TX PSD
      * \param type the type of transmission (for logging)
      */
     void Transmit(Time txDuration,
                   Ptr<const WifiPpdu> ppdu,
                   const WifiTxVector& txVector,
-                  std::string type);
+                  double txPowerDbm,
+                  Ptr<SpectrumValue> txPowerSpectrum,
+                  const std::string& type);
 
     /**
      * \param psduMap the PSDU(s) to transmit indexed by STA-ID